diff --git a/.axolotl-complete.bash b/.axolotl-complete.bash
new file mode 100644
index 0000000000..9a51399e65
--- /dev/null
+++ b/.axolotl-complete.bash
@@ -0,0 +1,41 @@
+#!/bin/bash
+
+_axolotl_completions() {
+    local cur prev
+    COMPREPLY=()
+    cur="${COMP_WORDS[COMP_CWORD]}"
+    prev="${COMP_WORDS[COMP_CWORD-1]}"
+
+    # If we're completing the first argument (the command)
+    if [[ $COMP_CWORD -eq 1 ]]; then
+        mapfile -t COMPREPLY < <(compgen -W "delinearize-llama4 fetch lm-eval merge-sharded-fsdp-weights quantize vllm-serve evaluate inference merge-lora preprocess train" -- "$cur")
+        return 0
+    fi
+
+    # Commands that should complete with directories and YAML files
+    local -a yaml_commands=("merge-sharded-fsdp-weights" "quantize" "vllm-serve" "evaluate" "inference" "merge-lora" "preprocess" "train")
+
+    # Check if previous word is in our list
+    if [[ " ${yaml_commands[*]} " =~ (^|[[:space:]])$prev($|[[:space:]]) ]]; then
+        # Use filename completion which handles directories properly
+        compopt -o filenames
+        mapfile -t COMPREPLY < <(compgen -f -- "$cur")
+
+        # Filter to only include directories and YAML files
+        local -a filtered=()
+        for item in "${COMPREPLY[@]}"; do
+            if [[ -d "$item" ]] || [[ "$item" == *.yaml ]] || [[ "$item" == *.yml ]]; then
+                filtered+=("$item")
+            fi
+        done
+        COMPREPLY=("${filtered[@]}")
+
+        return 0
+    fi
+
+    # Default: no completion
+    return 0
+}
+
+# Remove the -o nospace option - let filenames handle it
+complete -F _axolotl_completions axolotl
diff --git a/.bandit b/.bandit
index 2d81286aee..b814287519 100644
--- a/.bandit
+++ b/.bandit
@@ -1,3 +1,3 @@
 [bandit]
 exclude = tests
-skips = B101
+skips = B101,B615,B102,B110
diff --git a/.coderabbit.yaml b/.coderabbit.yaml
new file mode 100644
index 0000000000..821d6bd5b1
--- /dev/null
+++ b/.coderabbit.yaml
@@ -0,0 +1,17 @@
+# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json
+language: "en-US"
+early_access: false
+reviews:
+  profile: "chill"
+  request_changes_workflow: false
+  high_level_summary: true
+  review_status: true
+  collapse_walkthrough: true
+  poem: false
+  sequence_diagrams: false
+  auto_review:
+    enabled: true
+    drafts: false
+    auto_incremental_review: false
+chat:
+  auto_reply: true
diff --git a/.coveragerc b/.coveragerc
new file mode 100644
index 0000000000..c5670410c7
--- /dev/null
+++ b/.coveragerc
@@ -0,0 +1,14 @@
+[run]
+source = axolotl
+omit =
+    */tests/*
+    setup.py
+
+[report]
+exclude_lines =
+    pragma: no cover
+    def __repr__
+    raise NotImplementedError
+    if __name__ == .__main__.:
+    pass
+    raise ImportError
diff --git a/.flake8 b/.flake8
deleted file mode 100644
index fd69af7756..0000000000
--- a/.flake8
+++ /dev/null
@@ -1,5 +0,0 @@
-[flake8]
-max-line-length = 88
-
-select = C,E,F,W,B,B950
-extend-ignore = E203, E501, W503
diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md
index 9eec23e1a3..a18492ac52 100644
--- a/.github/CONTRIBUTING.md
+++ b/.github/CONTRIBUTING.md
@@ -15,38 +15,67 @@ First of all, thank you for your interest in contributing to axolotl! We appreci
   - [Commit Messages](#commit-messages)
 - [Additional Resources](#additional-resources)
 
-## Code of Conductcode
+## Code of Conduct
 
 All contributors are expected to adhere to our [Code of Conduct](CODE_OF_CONDUCT.md). Please read it before participating in the axolotl community.
 
 ## Getting Started
 
-Bugs? Please check for open issue else create a new [Issue](https://github.com/OpenAccess-AI-Collective/axolotl/issues/new).
+Bugs? Please check for open issue else create a new [Issue](https://github.com/axolotl-ai-cloud/axolotl/issues/new).
 
 PRs are **greatly welcome**!
 
 1. Fork the repository and clone it to your local machine.
-2. Set up the development environment by following the instructions in the [README.md](https://github.com/OpenAccess-AI-Collective/axolotl/tree/main/README.md) file.
+2. Set up the development environment by following the instructions in the [README.md](https://github.com/axolotl-ai-cloud/axolotl/tree/main/README.md) file.
 3. Explore the codebase, run tests, and verify that everything works as expected.
 
 Please run below to setup env
 ```bash
-pip3 install -r requirements-dev.txt -r requirements-tests.txt
+# Install axolotl + dev and test dependencies
+export UV_TORCH_BACKEND=cu128  # or cu130
+uv venv --no-project --relocatable
+source .venv/bin/activate
+uv pip install --no-build-isolation -e '.[deepspeed]' --group dev --group test
 pre-commit install
 
 # test
 pytest tests/
 ```
 
+CI tests across a matrix of Python and PyTorch versions — see [tests.yml](workflows/tests.yml) for the current one. Tests default to `-m 'not slow'`. Run the CPU suite locally (GPU e2e runs in separate jobs — see below):
+
+```bash
+pytest -m "not slow" -n4 --dist loadfile --ignore=tests/e2e tests/
+```
+
+### Running e2e (GPU) tests locally
+
+Recommended for larger changes before opening a PR. Needs an NVIDIA GPU. Run in the public Docker image with your checkout mounted ([docs/docker.qmd](../docs/docker.qmd) lists the available tags):
+
+```bash
+docker run --gpus all --rm -it --ipc=host -v "$PWD:/workspace/axolotl" -w /workspace/axolotl \
+  axolotlai/axolotl-uv:main-latest
+```
+
+The runtime image omits test deps, so install them, then run a test:
+
+```bash
+uv pip install --group test                  # tbparse, etc.
+pytest tests/e2e/test_lora_llama.py          # LoRA smoke test
+pytest tests/e2e/multigpu/                    # needs >= 2 GPUs
+```
+
+Some tests require flash-attn (`uv pip install flash-attn --no-build-isolation`). `cicd/cicd.sh` and `cicd/multigpu.sh` list CI's exact run order.
+
 ## How to Contribute
 
 ### Reporting Bugs
 
-If you encounter a bug or issue while using axolotl, please open a new issue on the [GitHub Issues](https://github.com/OpenAccess-AI-Collective/axolotl/issues) page. Provide a clear and concise description of the problem, steps to reproduce it, and any relevant error messages or logs.
+If you encounter a bug or issue while using axolotl, please open a new issue on the [GitHub Issues](https://github.com/axolotl-ai-cloud/axolotl/issues) page. Provide a clear and concise description of the problem, steps to reproduce it, and any relevant error messages or logs.
 
 ### Suggesting Enhancements
 
-We welcome ideas for improvements and new features. To suggest an enhancement, open a new issue on the [GitHub Issues](https://github.com/OpenAccess-AI-Collective/axolotl/issues) page. Describe the enhancement in detail, explain the use case, and outline the benefits it would bring to the project.
+We welcome ideas for improvements and new features. To suggest an enhancement, open a new issue on the [GitHub Issues](https://github.com/axolotl-ai-cloud/axolotl/issues) page. Describe the enhancement in detail, explain the use case, and outline the benefits it would bring to the project.
 
 ### Submitting Pull Requests
 
@@ -55,13 +84,30 @@ We welcome ideas for improvements and new features. To suggest an enhancement, o
 3. Test your changes and ensure that they don't introduce new issues or break existing functionality.
 4. Commit your changes, following the [commit message guidelines](#commit-messages).
 5. Push your branch to your fork on GitHub.
-6. Open a new pull request against the `main` branch of the axolotl repository. Include a clear and concise description of your changes, referencing any related issues.
+6. Open a new pull request against the `main` branch of the axolotl repository. PR formatting is prescribed in the [PR template](PULL_REQUEST_TEMPLATE.md); reference any related issues.
+
+#### Skipping CI Checks
+
+You can skip certain CI checks by including specific keywords in your commit messages:
+
+- `[skip ci]` or `skip ci` - Skips all CI checks for that commit
+- `[skip-e2e]` or `skip-e2e` - Skips only end-to-end tests while running other CI checks. You may also include this in the title of your PR to disable end-to-end tests for the entire PR.
 
 ## Style Guidelines
 
 ### Code Style
 
-axolotl uses [{codestyle}]({URLofCodestyle}) as its code style guide. Please ensure that your code follows these guidelines.
+axolotl uses [Ruff](https://docs.astral.sh/ruff/) as its code style guide. Please ensure that your code follows these guidelines.
+
+Use the pre-commit linter to ensure that your code is formatted consistently. It installs and runs the **exact versions CI uses**, so don't rely on a system-installed `ruff`/`mypy`:
+```bash
+pre-commit install        # one-time
+pre-commit run --all-files
+```
+
+The exact ruff/mypy/bandit versions are pinned in [`.pre-commit-config.yaml`](../.pre-commit-config.yaml) — the same file CI's pre-commit job runs from, so local and CI never drift.
+
+To run ruff outside pre-commit, pin it to the `ruff-pre-commit` rev in that file so output matches CI, e.g. `uvx ruff@<rev> check` / `uvx ruff@<rev> format`.
 
 ### Commit Messages
 
@@ -71,6 +117,6 @@ Write clear and concise commit messages that briefly describe the changes made i
 
 - [GitHub Help](https://help.github.com/)
 - [GitHub Pull Request Documentation](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests)
-- [{codestyle}]({URLofCodestyle})
+- [Ruff](https://docs.astral.sh/ruff/)
 
 Thank you once again for your interest in contributing to axolotl. We look forward to collaborating with you and creating an even better project together!
diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml
index 4f6ea8de73..cd443a1971 100644
--- a/.github/FUNDING.yml
+++ b/.github/FUNDING.yml
@@ -1,13 +1,13 @@
 # These are supported funding model platforms
 
-github: [winglian, OpenAccess-AI-Collective] # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
+github: # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
 patreon: # Replace with a single Patreon username
 open_collective: # Replace with a single Open Collective username
-ko_fi: axolotl_ai # Replace with a single Ko-fi username
+ko_fi: # Replace with a single Ko-fi username
 tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
 community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
 liberapay: # Replace with a single Liberapay username
 issuehunt: # Replace with a single IssueHunt username
 otechie: # Replace with a single Otechie username
 lfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry
-custom: ['https://quickchart.io/qr?text=bitcoin%3Abc1qxlgwlqwfea5s2cxm42xqsfmwjct0rj8w8ea5np&size=480&centerImageUrl=https%3A%2F%2Fupload.wikimedia.org%2Fwikipedia%2Fcommons%2Fthumb%2F4%2F46%2FBitcoin.svg%2F64px-Bitcoin.svg.png'] # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']
+custom: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']
diff --git a/.github/ISSUE_TEMPLATE/bug-report.yaml b/.github/ISSUE_TEMPLATE/bug-report.yaml
index 1ed703d4e8..0bfe067ab6 100644
--- a/.github/ISSUE_TEMPLATE/bug-report.yaml
+++ b/.github/ISSUE_TEMPLATE/bug-report.yaml
@@ -15,7 +15,7 @@ body:
       label: "Please check that this issue hasn't been reported before."
       description: "The **Label filters** may help make your search more focussed."
       options:
-        - label: "I searched previous [Bug Reports](https://github.com/OpenAccess-AI-Collective/axolotl/labels/bug) didn't find any similar reports."
+        - label: "I searched previous [Bug Reports](https://github.com/axolotl-ai-cloud/axolotl/labels/bug) didn't find any similar reports."
           required: true
 
   - type: textarea
diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml
index e0c5754fb9..973b2d7fb4 100644
--- a/.github/ISSUE_TEMPLATE/config.yml
+++ b/.github/ISSUE_TEMPLATE/config.yml
@@ -1,7 +1,7 @@
 blank_issues_enabled: false
 contact_links:
   - name: Ask a question
-    url: https://github.com/OpenAccess-AI-Collective/axolotl/discussions/categories/q-a
+    url: https://github.com/axolotl-ai-cloud/axolotl/discussions/categories/q-a
     about: Ask questions and discuss with other community members
   - name: Discuss the Project in Discord
     url: https://discord.gg/HhrNrHJPRb
diff --git a/.github/ISSUE_TEMPLATE/docs.yml b/.github/ISSUE_TEMPLATE/docs.yml
index 2c27af1aa0..9a926c384f 100644
--- a/.github/ISSUE_TEMPLATE/docs.yml
+++ b/.github/ISSUE_TEMPLATE/docs.yml
@@ -10,7 +10,7 @@ body:
       value: |
         * Ask questions in [Discord](https://discord.gg/HhrNrHJPRb).
         * Before you file an issue read the [Contributing guide](./CONTRIBUTING.md).
-        * Check to make sure someone hasn't already opened a [similar issue](https://github.com/OpenAccess-AI-Collective/axolotl/issues).
+        * Check to make sure someone hasn't already opened a [similar issue](https://github.com/axolotl-ai-cloud/axolotl/issues).
   - type: textarea
     attributes:
       label: What piece of documentation is affected?
diff --git a/.github/ISSUE_TEMPLATE/feature-request.yaml b/.github/ISSUE_TEMPLATE/feature-request.yaml
index 39b6cb74e1..d90ebd274a 100644
--- a/.github/ISSUE_TEMPLATE/feature-request.yaml
+++ b/.github/ISSUE_TEMPLATE/feature-request.yaml
@@ -8,9 +8,9 @@ body:
       label: "⚠️ Please check that this feature request hasn't been suggested before."
       description: "There are two locations for previous feature requests. Please search in both. Thank you. The **Label filters** may help make your search more focussed."
       options:
-        - label: "I searched previous [Ideas in Discussions](https://github.com/OpenAccess-AI-Collective/axolotl/discussions/categories/ideas) didn't find any similar feature requests."
+        - label: "I searched previous [Ideas in Discussions](https://github.com/axolotl-ai-cloud/axolotl/discussions/categories/ideas) didn't find any similar feature requests."
           required: true
-        - label: "I searched previous [Issues](https://github.com/OpenAccess-AI-Collective/axolotl/labels/enhancement) didn't find any similar feature requests."
+        - label: "I searched previous [Issues](https://github.com/axolotl-ai-cloud/axolotl/labels/enhancement) didn't find any similar feature requests."
           required: true
 
   - type: textarea
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index 562806287a..cc0ade6653 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -15,6 +15,11 @@
 <!--- Include details of your testing environment, tests ran to see how -->
 <!--- your change affects other areas of the code, etc. -->
 
+## AI Usage Disclaimer
+
+<!--- Was AI (e.g., ChatGPT, Claude, Copilot) used to generate or assist with this PR? -->
+<!--- Please indicate: No / Yes (specify which tool and to what extent) -->
+
 ## Screenshots (if appropriate)
 
 ## Types of changes
diff --git a/.github/SECURITY.md b/.github/SECURITY.md
index aceb0d1a2e..e8c8ceff85 100644
--- a/.github/SECURITY.md
+++ b/.github/SECURITY.md
@@ -6,4 +6,4 @@ Due to the nature of the fast development that is happening in this project, onl
 
 ## Reporting a Vulnerability
 
-If you find a vulnerability, please contact us on  [Discord](https://discord.gg/xcu3ECkH9a) rather than creating a GitHub issue to allow us some time to fix it before it is a known vulnerability to others.
+If you find a vulnerability, please contact us by email `wing@axolotl.ai` rather than creating a GitHub issue to allow us some time to fix it before it is a known vulnerability to others.
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000000..5ace4600a1
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
index 2e8db61cc7..8b9a5d47bc 100644
--- a/.github/workflows/base.yml
+++ b/.github/workflows/base.yml
@@ -1,68 +1,115 @@
 name: ci-cd-base
 
 on:
+  push:
+    branches:
+      - "main"
+    paths:
+      - 'docker/Dockerfile-uv-base'
+      - '.github/workflows/base.yml'
+  pull_request:
+    paths:
+      - 'docker/Dockerfile-uv-base'
+      - '.github/workflows/base.yml'
   workflow_dispatch:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
 jobs:
-  build-base:
-    if: github.repository_owner == 'OpenAccess-AI-Collective'
-    # this job needs to be run on self-hosted GPU runners...
-    runs-on: axolotl-gpu-runner
+  build-base-uv:
+    if: ${{ github.repository_owner == 'axolotl-ai-cloud' && (github.event_name != 'pull_request' || !github.event.pull_request.draft) }}
+    timeout-minutes: 480
+    runs-on: ubuntu-latest-m
+    env:
+      HAS_DOCKERHUB_CREDS: ${{ secrets.DOCKERHUB_USERNAME != '' && secrets.DOCKERHUB_TOKEN != '' }}
     strategy:
       fail-fast: false
       matrix:
         include:
-          - cuda: "118"
-            cuda_version: 11.8.0
-            python_version: "3.10"
-            pytorch: 2.1.2
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "121"
-            cuda_version: 12.1.0
-            python_version: "3.10"
-            pytorch: 2.1.2
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "121"
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.1.2
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "121"
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "121"
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.3.0
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+          - cuda: "130"
+            cuda_version: 13.0.0
+            cudnn_version: ""
+            python_version: "3.12"
+            pytorch: 2.11.0
+            torch_cuda_arch_list: "9.0 10.0 10.3 12.0+PTX"
+            dockerfile: "Dockerfile-uv-base"
+            platforms: "linux/amd64,linux/arm64"
+          - cuda: "130"
+            cuda_version: 13.0.0
+            cudnn_version: ""
+            python_version: "3.12"
+            pytorch: 2.12.0
+            torch_cuda_arch_list: "9.0 10.0 10.3 12.0+PTX"
+            dockerfile: "Dockerfile-uv-base"
+            platforms: "linux/amd64,linux/arm64"
+          - cuda: "130"
+            cuda_version: 13.0.0
+            cudnn_version: ""
+            python_version: "3.12"
+            pytorch: 2.12.1
+            torch_cuda_arch_list: "9.0 10.0 10.3 12.0+PTX"
+            dockerfile: "Dockerfile-uv-base"
+            platforms: "linux/amd64,linux/arm64"
+          - cuda: "130"
+            cuda_version: 13.0.0
+            cudnn_version: ""
+            python_version: "3.12"
+            pytorch: 2.13.0
+            torch_cuda_arch_list: "9.0 10.0 10.3 12.0+PTX"
+            dockerfile: "Dockerfile-uv-base"
+            platforms: "linux/amd64,linux/arm64"
+          - cuda: "132"
+            cuda_version: 13.2.1
+            cudnn_version: ""
+            python_version: "3.12"
+            pytorch: 2.13.0
+            torch_index_url: "https://download.pytorch.org/whl/cu132"
+            torch_cuda_arch_list: "9.0 10.0 10.3 12.0+PTX"
+            dockerfile: "Dockerfile-uv-base"
+            platforms: "linux/amd64,linux/arm64"
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
       - name: Docker metadata
         id: metadata
-        uses: docker/metadata-action@v3
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
         with:
-          images: winglian/axolotl-base
+          images: |
+            axolotlai/axolotl-base-uv
       - name: Login to Docker Hub
-        uses: docker/login-action@v2
+        uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.7.0
+        if: ${{ github.event_name != 'pull_request' && env.HAS_DOCKERHUB_CREDS == 'true' }}
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v2
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
       - name: Build
-        uses: docker/build-push-action@v4
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
         with:
           context: .
-          file: ./docker/Dockerfile-base
+          file: ./docker/${{ matrix.dockerfile }}
+          platforms: ${{ matrix.platforms }}
           push: ${{ github.event_name != 'pull_request' }}
-          tags: ${{ steps.metadata.outputs.tags }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+          cache-from: type=registry,ref=axolotlai/axolotl-base-uv:${{ steps.metadata.outputs.version }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+          cache-to: type=inline
+          tags: |
+            axolotlai/axolotl-base-uv:${{ steps.metadata.outputs.version }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            axolotlai/axolotl-base:${{ steps.metadata.outputs.version }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
           build-args: |
             CUDA_VERSION=${{ matrix.cuda_version }}
+            CUDNN_VERSION=${{ matrix.cudnn_version }}
             CUDA=${{ matrix.cuda }}
+            TORCH_BACKEND=${{ matrix.torch_backend || format('cu{0}', matrix.cuda) }}
+            TORCH_INDEX_URL=${{ matrix.torch_index_url }}
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
             TORCH_CUDA_ARCH_LIST=${{ matrix.torch_cuda_arch_list }}
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 16cf774d71..b7b85c8aa2 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -3,6 +3,13 @@ on:
   push:
     branches:
       - main
+    paths:
+      - '**/*.md'
+      - '**/*.qmd'
+      - '_quarto.yml'
+      - 'docs/**'
+      - 'src/axolotl/utils/schemas/**.py'
+      - '.github/workflows/docs.yml'
 
 permissions:
     contents: write
@@ -12,19 +19,27 @@ jobs:
     build-deploy:
         runs-on: ubuntu-latest
         steps:
+        - name: cleanup node
+          run: |
+            sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL
         - name: Check out repository
-          uses: actions/checkout@v4
+          uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+          with:
+            persist-credentials: false
         - name: Set up Quarto
-          uses: quarto-dev/quarto-actions/setup@v2
+          uses: quarto-dev/quarto-actions/setup@8a96df13519ee81fd526f2dfca5962811136661b # v2.2.0
         - name: Setup Python
-          uses: actions/setup-python@v3
+          uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
           with:
-            python-version: '3.10'
-        - name: install dependencies
+            python-version: '3.11'
+        - name: Install dependencies
           run: |
-            python3 -m pip install jupyter
+            python3 -m pip install jupyter quartodoc
+            python3 -m pip install -e .
+        - name: Build autodoc
+          run: quartodoc build
         - name: Publish to GitHub Pages (and render)
-          uses: quarto-dev/quarto-actions/publish@v2
+          uses: quarto-dev/quarto-actions/publish@8a96df13519ee81fd526f2dfca5962811136661b # v2.2.0
           with:
             target: gh-pages
           env:
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 671be4b652..9dfc09dc45 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -1,23 +1,39 @@
 name: lint
 on:
   # check on PRs, and manual triggers
+  merge_group:
   pull_request:
+      types: [opened, synchronize, reopened, ready_for_review]
+      # code/pyproject/workflow linting is covered by the pre-commit job in tests.yml;
+      # this workflow handles the doc/example paths that tests.yml does not trigger on
       paths:
-       - '**.py'
-       - 'requirements.txt'
-       - '.github/workflows/*.yml'
-       - "*.md"
+       - "*.[q]md"
        - "examples/**/*.y[a]?ml"
+       - ".pre-commit-config.yaml"
   workflow_dispatch:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+permissions:
+  contents: read
+
 jobs:
   pre-commit:
     name: pre-commit
     runs-on: ubuntu-latest
+    if: ${{ !github.event.pull_request.draft }}
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
-          python-version: "3.10"
+          python-version: "3.11"
           cache: 'pip' # caching pip dependencies
-      - uses: pre-commit/action@v3.0.0
+      - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
+        env:
+          # check-cli-config-options needs an installed axolotl; the dedicated
+          # tests.yml step covers it in CI
+          SKIP: check-cli-config-options
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index f51c42a823..e87b890563 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -4,124 +4,207 @@ on:
   push:
     branches:
       - "main"
+    tags:
+      - "v*"
   workflow_dispatch:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
 jobs:
-  build-axolotl:
-    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]]') && github.repository_owner == 'OpenAccess-AI-Collective' }}
+  build-axolotl-uv:
+    if: ${{ ! contains(github.event.head_commit.message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
     strategy:
       fail-fast: false
       matrix:
         include:
-          - cuda: 118
-            cuda_version: 11.8.0
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_extras:
-            axolotl_args: "--extra-index-url https://download.pytorch.org/whl/cu118"
-            is_latest: true
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.1
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
             axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.3.0
+            platforms: "linux/amd64,linux/arm64"
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
             axolotl_extras:
+            platforms: "linux/amd64,linux/arm64"
+            is_latest: true
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
       - name: Docker metadata
         id: metadata
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
         with:
-          images: winglian/axolotl
+          images: |
+            axolotlai/axolotl-uv
+          tags: |
+            type=ref,event=branch
+            type=pep440,pattern={{version}}
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
       - name: Login to Docker Hub
-        uses: docker/login-action@v3
+        uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.7.0
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       # guidance for testing before pushing: https://docs.docker.com/build/ci/github-actions/test-before-push/
       - name: Build and export to Docker
-        uses: docker/build-push-action@v5
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
         with:
           context: .
+          platforms: ${{ matrix.platforms }}
           build-args: |
-            BASE_TAG=${{ github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
             CUDA=${{ matrix.cuda }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
             AXOLOTL_ARGS=${{ matrix.axolotl_args }}
-          file: ./docker/Dockerfile
+            AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}
+          file: ./docker/Dockerfile-uv
           push: ${{ github.event_name != 'pull_request' }}
+          cache-from: type=registry,ref=axolotlai/axolotl-uv:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+          cache-to: type=inline
           tags: |
-            ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
-            ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
+            axolotlai/axolotl-uv:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            axolotlai/axolotl-uv:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+            ${{ (matrix.is_latest) && format('axolotlai/axolotl-uv:{0}-latest', steps.metadata.outputs.version) || '' }}
+            ${{ (github.ref_type == 'tag' && matrix.is_latest) && format('axolotlai/axolotl-uv:{0}', steps.metadata.outputs.version) || '' }}
+            axolotlai/axolotl:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            axolotlai/axolotl:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+            ${{ (matrix.is_latest) && format('axolotlai/axolotl:{0}-latest', steps.metadata.outputs.version) || '' }}
+            ${{ (github.ref_type == 'tag' && matrix.is_latest) && format('axolotlai/axolotl:{0}', steps.metadata.outputs.version) || '' }}
           labels: ${{ steps.metadata.outputs.labels }}
 
-  build-axolotl-cloud:
-    needs: build-axolotl
-    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]]') && github.repository_owner == 'OpenAccess-AI-Collective' }}
+  build-axolotl-cloud-uv:
+    needs: build-axolotl-uv
+    if: ${{ ! contains(github.event.head_commit.message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
     # this job needs to be run on self-hosted GPU runners...
     strategy:
+      fail-fast: false
       matrix:
         include:
-          - cuda: 118
-            cuda_version: 11.8.0
-            python_version: "3.10"
-            pytorch: 2.1.2
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
             axolotl_extras:
-            is_latest: true
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.10"
-            pytorch: 2.1.2
+            platforms: "linux/amd64,linux/arm64"
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
             axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.1
+            platforms: "linux/amd64,linux/arm64"
+            is_latest: true
+    runs-on: axolotl-gpu-runner
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - name: Docker metadata
+        id: metadata
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
+        with:
+          images: |
+            axolotlai/axolotl-cloud-uv
+          tags: |
+            type=ref,event=branch
+            type=pep440,pattern={{version}}
+      - name: Login to Docker Hub
+        uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.7.0
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
+      - name: Build
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
+        with:
+          context: .
+          platforms: ${{ matrix.platforms }}
+          build-args: |
+            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            CUDA=${{ matrix.cuda }}
+          file: ./docker/Dockerfile-cloud-uv
+          push: ${{ github.event_name != 'pull_request' }}
+          cache-from: type=registry,ref=axolotlai/axolotl-cloud-uv:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+          cache-to: type=inline
+          tags: |
+             axolotlai/axolotl-cloud-uv:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+             ${{ (matrix.is_latest) && format('axolotlai/axolotl-cloud-uv:{0}-latest', steps.metadata.outputs.version) || '' }}
+             ${{ (github.ref_type == 'tag' && matrix.is_latest) && format('axolotlai/axolotl-cloud-uv:{0}', steps.metadata.outputs.version) || '' }}
+             axolotlai/axolotl-cloud:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+             ${{ (matrix.is_latest) && format('axolotlai/axolotl-cloud:{0}-latest', steps.metadata.outputs.version) || '' }}
+             ${{ (github.ref_type == 'tag' && matrix.is_latest) && format('axolotlai/axolotl-cloud:{0}', steps.metadata.outputs.version) || '' }}
+          labels: ${{ steps.metadata.outputs.labels }}
+
+  build-axolotl-cloud-no-tmux-uv:
+    needs: build-axolotl-uv
+    if: ${{ ! contains(github.event.head_commit.message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
+    # this job needs to be run on self-hosted GPU runners...
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
             axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.3.0
+            platforms: "linux/amd64,linux/arm64"
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
             axolotl_extras:
+            platforms: "linux/amd64,linux/arm64"
+            is_latest: true
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
       - name: Docker metadata
         id: metadata
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
         with:
-          images: winglian/axolotl-cloud
+          images: |
+            axolotlai/axolotl-cloud-term
+          tags: |
+            type=ref,event=branch
+            type=pep440,pattern={{version}}
       - name: Login to Docker Hub
-        uses: docker/login-action@v3
+        uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.7.0
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v2
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
       - name: Build
-        uses: docker/build-push-action@v5
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
         with:
           context: .
+          platforms: linux/amd64,linux/arm64
           build-args: |
-            BASE_TAG=${{ github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
             CUDA=${{ matrix.cuda }}
-          file: ./docker/Dockerfile-cloud
+          file: ./docker/Dockerfile-cloud-no-tmux-uv
           push: ${{ github.event_name != 'pull_request' }}
+          cache-from: type=registry,ref=axolotlai/axolotl-cloud-term:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+          cache-to: type=inline
           tags: |
-             ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
-             ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
+             axolotlai/axolotl-cloud-term:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+             ${{ (matrix.is_latest) && format('axolotlai/axolotl-cloud-term:{0}-latest', steps.metadata.outputs.version) || '' }}
+             ${{ (github.ref_type == 'tag' && matrix.is_latest) && format('axolotlai/axolotl-cloud-term:{0}', steps.metadata.outputs.version) || '' }}
           labels: ${{ steps.metadata.outputs.labels }}
diff --git a/.github/workflows/multi-gpu-e2e.yml b/.github/workflows/multi-gpu-e2e.yml
new file mode 100644
index 0000000000..61c26a5bfd
--- /dev/null
+++ b/.github/workflows/multi-gpu-e2e.yml
@@ -0,0 +1,69 @@
+name: docker-multigpu-tests-biweekly
+
+on:
+  pull_request:
+    paths:
+      - "tests/e2e/multigpu/**.py"
+      - "pyproject.toml"
+      - ".github/workflows/multi-gpu-e2e.yml"
+      - "scripts/cutcrossentropy_install.py"
+      - "src/axolotl/core/trainers/mixins/sequence_parallel.py"
+      - "src/axolotl/utils/distributed.py"
+  workflow_dispatch:
+  schedule:
+    - cron: "0 0 * * 1,4" # Runs at 00:00 UTC every monday & thursday
+
+# Cancel jobs on the same ref if a new one is triggered
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+permissions:
+  contents: read
+
+env:
+  MODAL_IMAGE_BUILDER_VERSION: "2025.06"
+
+jobs:
+  test-axolotl-multigpu:
+    if: ${{ ! contains(github.event.head_commit.message, '[skip e2e]') && github.repository_owner == 'axolotl-ai-cloud' && (github.event_name != 'pull_request' || !github.event.pull_request.draft) }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
+            axolotl_extras:
+            # axolotl_extras: fbgemm-gpu
+            num_gpus: 2
+    runs-on: [self-hosted, modal]
+    timeout-minutes: 120
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - name: Install Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.11"
+      - name: Install Modal
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal==1.3.0.post1 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 "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+      - name: Run tests job on Modal
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+        run: |
+          modal run -m cicd.multigpu
diff --git a/.github/workflows/nightlies.yml b/.github/workflows/nightlies.yml
index f4172fa406..3861628866 100644
--- a/.github/workflows/nightlies.yml
+++ b/.github/workflows/nightlies.yml
@@ -5,56 +5,54 @@ on:
   schedule:
     - cron: '0 0 * * *'  # Runs at 00:00 UTC every day
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
 jobs:
   build-axolotl:
-    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]]') && github.repository_owner == 'OpenAccess-AI-Collective' }}
+    if: ${{ ! contains(github.event.head_commit.message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
     strategy:
       fail-fast: false
       matrix:
         include:
-          - cuda: 118
-            cuda_version: 11.8.0
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_extras:
-            axolotl_args: "--extra-index-url https://download.pytorch.org/whl/cu118"
-            is_latest: true
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.1
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
             axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.3.0
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
             axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
       - name: Docker metadata
         id: metadata
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
         with:
-          images: winglian/axolotl
+          images: |
+            axolotlai/axolotl
           tags: |
             type=raw,value={{ branch }}-{{ date 'YYYYMMDD' }}
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
       - name: Login to Docker Hub
-        uses: docker/login-action@v3
+        uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.7.0
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       # guidance for testing before pushing: https://docs.docker.com/build/ci/github-actions/test-before-push/
       - name: Build and export to Docker
-        uses: docker/build-push-action@v5
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
         with:
           context: .
           build-args: |
@@ -62,67 +60,61 @@ jobs:
             CUDA=${{ matrix.cuda }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
             AXOLOTL_ARGS=${{ matrix.axolotl_args }}
-          file: ./docker/Dockerfile
+          file: ./docker/Dockerfile-uv
           push: ${{ github.event_name != 'pull_request' }}
           tags: |
-            ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            axolotlai/axolotl:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            axolotlai/axolotl-uv:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
 
   build-axolotl-cloud:
     needs: build-axolotl
-    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]]') && github.repository_owner == 'OpenAccess-AI-Collective' }}
+    if: ${{ ! contains(github.event.head_commit.message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
     # this job needs to be run on self-hosted GPU runners...
     strategy:
       matrix:
         include:
-          - cuda: 118
-            cuda_version: 11.8.0
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_extras:
-            is_latest: true
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.10"
-            pytorch: 2.1.2
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
             axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.1
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.3.0
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
             axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
       - name: Docker metadata
         id: metadata
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
         with:
-          images: winglian/axolotl-cloud
+          images: |
+            axolotlai/axolotl-cloud
           tags: |
             type=raw,value={{ branch }}-{{ date 'YYYYMMDD' }}
       - name: Login to Docker Hub
-        uses: docker/login-action@v3
+        uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.7.0
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v2
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
       - name: Build
-        uses: docker/build-push-action@v5
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
         with:
           context: .
           build-args: |
             BASE_TAG=${{ github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
             CUDA=${{ matrix.cuda }}
-          file: ./docker/Dockerfile-cloud
+          file: ./docker/Dockerfile-cloud-uv
           push: ${{ github.event_name != 'pull_request' }}
           tags: |
-             ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+             axolotlai/axolotl-cloud:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+             axolotlai/axolotl-cloud-uv:${{ steps.metadata.outputs.version }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
diff --git a/.github/workflows/precommit-autoupdate.yml b/.github/workflows/precommit-autoupdate.yml
new file mode 100644
index 0000000000..ac24bcbff5
--- /dev/null
+++ b/.github/workflows/precommit-autoupdate.yml
@@ -0,0 +1,45 @@
+name: Pre-commit auto-update
+
+on:
+  schedule:
+    - cron: '0 0 1 * *'  # Run monthly
+  workflow_dispatch:  # Manual kickoff
+
+permissions: {}
+
+jobs:
+  auto-update:
+    runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'axolotl-ai-cloud' }}
+    permissions:
+      contents: write
+      pull-requests: write
+    steps:
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        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
+          fi
+
+      - name: Create Pull Request
+        if: steps.update.outputs.changes == 'true'
+        uses: peter-evans/create-pull-request@c5a7806660adbe173f04e3e038b0ccdcd758773c # v6.1.0
+        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.
diff --git a/.github/workflows/preview-docs.yml b/.github/workflows/preview-docs.yml
new file mode 100644
index 0000000000..73b57b3603
--- /dev/null
+++ b/.github/workflows/preview-docs.yml
@@ -0,0 +1,82 @@
+name: Preview
+on:
+  workflow_dispatch:
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
+
+    # Run the workflow only when one of these files changes
+    paths:
+      - '**/*.md'      # any Markdown file
+      - '**/*.qmd'     # any Quarto file
+      - '_quarto.yml'
+      - docs/scripts/generate_config_docs.py
+      - src/axolotl/utils/schemas/**.py
+      - .github/workflows/preview-docs.yml
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+permissions:
+  contents: read
+  pull-requests: write
+
+jobs:
+  preview:
+    runs-on: ubuntu-latest
+    if: ${{ !github.event.pull_request.draft }}
+    steps:
+      - name: cleanup node
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL
+
+      - name: Check out repository
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+          ref: ${{ github.event.pull_request.head.sha }}
+
+      - name: Set up Quarto
+        uses: quarto-dev/quarto-actions/setup@8a96df13519ee81fd526f2dfca5962811136661b # v2.2.0
+
+      - name: Setup Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: '3.11'
+
+      - name: Install dependencies
+        run: |
+          python3 -m pip install jupyter quartodoc
+          python3 -m pip install -e .
+
+      - name: Build autodoc
+        run: quartodoc build
+
+      - name: Quarto render
+        run: quarto render
+
+      - name: Netlify Publish
+        uses: nwtgck/actions-netlify@4cbaf4c08f1a7bfa537d6113472ef4424e4eb654 # v3.0.0
+        if: ${{ github.event.pull_request.head.repo.full_name == github.repository }}
+        id: netlify
+        with:
+          publish-dir: './_site'
+          enable-pull-request-comment: false
+          enable-github-deployment: false
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          deploy-message: "Deployed On Netlify"
+          github-deployment-environment: 'preview'
+          github-deployment-description: 'Preview Deployment'
+        env:
+          NETLIFY_AUTH_TOKEN: ${{ secrets.NETLIFY_AUTH_TOKEN }}
+          NETLIFY_SITE_ID: ${{ secrets.NETLIFY_SITE_ID }}
+
+      - name: Update PR with preview link
+        if: ${{ steps.netlify.outcome == 'success' }}
+        uses: marocchino/sticky-pull-request-comment@773744901bac0e8cbb5a0dc842800d45e9b2b405 # v2.9.4
+        with:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          message: |
+            📖 **Documentation Preview**: ${{ steps.netlify.outputs.deploy-url }}
+
+            Deployed on Netlify from commit ${{ github.event.pull_request.head.sha }}
diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml
index 885239d185..a3eac95e3f 100644
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -3,43 +3,68 @@ name: publish pypi
 on:
   push:
     tags:
-      - '*'
+      - "v*"
+  workflow_dispatch:
+
+permissions: {}
 
 jobs:
+  setup_release:
+    name: Create Release
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+
+      - name: Create release
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        # idempotent: don't fail a re-run if the release already exists
+        run: |
+          gh release view "$GITHUB_REF_NAME" >/dev/null 2>&1 \
+            || gh release create "$GITHUB_REF_NAME" --generate-notes
   pypi-publish:
     name: Upload release to PyPI
     runs-on: ubuntu-latest
+    needs: [setup_release]
     environment:
       name: pypi
       url: https://pypi.org/p/axolotl
     permissions:
-      id-token: write  # IMPORTANT: this permission is mandatory for trusted publishing
+      contents: read
+      id-token: write # IMPORTANT: this permission is mandatory for trusted publishing
     steps:
       - name: Check out repository code
-        uses: actions/checkout@v3
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
 
       - name: Setup Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
-          python-version: "3.10"
+          python-version: "3.11"
 
-      - name: Install dependencies
-        run: |
-          pip3 install wheel packaging
-          pip3 install -e .
-          pip3 install -r requirements-tests.txt
+      - name: Install uv
+        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        with:
+          enable-cache: false
 
       - name: Extract tag name
         id: tag
-        run: echo ::set-output name=TAG_NAME::$(echo $GITHUB_REF | cut -d / -f 3)
+        run: echo "TAG_NAME=$(echo $GITHUB_REF | cut -d / -f 3)" >> "$GITHUB_OUTPUT"
 
-      - name: Update version in setup.py
+      - name: Update version in VERSION file
         run: |
-          sed -i -E 's/version="([0-9.]+)",/version="${{ steps.tag.outputs.TAG_NAME }}",/g' setup.py
+          echo "${{ steps.tag.outputs.TAG_NAME }}" | sed 's/^v//' > VERSION
 
-      - name: Build a binary wheel
-        run: |
-          python setup.py sdist bdist_wheel
+      - name: Build sdist and wheel
+        # PEP 517 build via uv (setuptools backend reads the version from VERSION);
+        # replaces the removed `python setup.py sdist` after the pyproject migration.
+        run: uv build
 
       - name: Publish package distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
+        uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0
diff --git a/.github/workflows/tests-nightly.yml b/.github/workflows/tests-nightly.yml
new file mode 100644
index 0000000000..1c04384d76
--- /dev/null
+++ b/.github/workflows/tests-nightly.yml
@@ -0,0 +1,295 @@
+name: Tests Nightly against upstream main
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 0 * * *" # Runs at 00:00 UTC every day
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
+    paths:
+      - ".github/workflows/tests-nightly.yml"
+      - "cicd/cicd.sh"
+      - "cicd/cicd_cuda_kernels.sh"
+      - "cicd/e2e_tests.py"
+      - "cicd/e2e_cuda_kernels.py"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+permissions:
+  contents: read
+
+env:
+  UV_SYSTEM_PYTHON: "1"
+
+jobs:
+  pre-commit:
+    name: pre-commit
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+          cache: "pip" # caching pip dependencies
+      - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
+        env:
+          # check-cli-config-options needs an installed axolotl; the dedicated
+          # tests.yml step covers it in CI
+          SKIP: no-commit-to-branch,check-cli-config-options
+
+  prime-cdn-s3-cache:
+    name: Prefetch S3 once to prime the CDN cache
+    runs-on: ubuntu-latest
+    if: ${{ !github.event.pull_request.draft }}
+    timeout-minutes: 10
+    steps:
+      - name: Restore Cache from S3
+        id: hf-cache-restore-s3
+        run: |
+          curl -v -H "Range: bytes=0-1023" -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst > /dev/null
+
+  pytest:
+    name: PyTest
+    runs-on: ubuntu-latest
+    needs: [prime-cdn-s3-cache]
+    strategy:
+      fail-fast: false
+      matrix:
+        python_version: ["3.12"] # TODO include py3.14 once https://github.com/mistralai/mistral-common/pull/194 is merged
+        pytorch_version: ["2.11.0", "2.12.0"]
+    timeout-minutes: 20
+
+    steps:
+      - name: Check out repository code
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+
+      - name: Restore Cache from S3
+        id: hf-cache-restore-s3
+        run: |
+          mkdir -p /home/runner/.cache/huggingface/hub
+          curl -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst | tar -xf - -C /home/runner/.cache/huggingface/hub/  --use-compress-program unzstd
+          ls -ltr /home/runner/.cache/huggingface/hub/
+
+      - name: Setup Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: ${{ matrix.python_version }}
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        with:
+          enable-cache: true
+          cache-dependency-glob: "pyproject.toml"
+
+      - name: Install PyTorch
+        run: |
+          uv pip install torch==${{ matrix.pytorch_version }} torchvision
+          uv pip freeze | grep -E "^(torch|torchvision)==" > /tmp/torch-pin.txt
+
+      - name: Install dependencies
+        run: |
+          uv pip install --no-build-isolation -e . --override /tmp/torch-pin.txt
+          python scripts/cutcrossentropy_install.py --uv | sh
+          uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+            codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
+
+      - name: Override with nightly HF packages
+        run: |
+          uv pip install "kernels>=0.15.2,<0.16"
+          uv pip install --no-deps \
+            "transformers @ git+https://github.com/huggingface/transformers.git@main" \
+            "peft @ git+https://github.com/huggingface/peft.git@main" \
+            "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
+            "trl @ git+https://github.com/huggingface/trl.git@main" \
+            "datasets @ git+https://github.com/huggingface/datasets.git@main"
+
+      - name: Make sure PyTorch version wasn't clobbered
+        run: |
+          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__, f'Expected torch ${{ matrix.pytorch_version }} but got {torch.__version__}'"
+
+      - name: Ensure axolotl CLI was installed
+        run: |
+          axolotl --help
+
+      - name: Pre-Download dataset fixture
+        run: |
+          hf download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures
+
+      - name: Show HF cache
+        run: hf cache ls
+
+      - name: Run tests
+        run: |
+          pytest -v --durations=10 -n8 --dist loadfile \
+            --ignore=tests/e2e/ \
+            --ignore=tests/integrations/kernels/ \
+            --ignore=tests/integrations/monkeypatch/test_tiled_mlp_moe.py \
+            --ignore=tests/integrations/test_gemma4_moe.py \
+            --ignore=tests/integrations/test_scattermoe_lora.py \
+            --ignore=tests/integrations/test_scattermoe_lora_kernels.py \
+            --ignore=tests/integrations/test_scattermoe_multi_lora.py \
+            --ignore=tests/integrations/test_sonicmoe_multi_lora.py \
+            --ignore=tests/patched/ \
+            --ignore=tests/cli/ \
+            tests/
+          pytest -v --durations=10 tests/patched/
+          pytest -v --durations=10 tests/cli/
+
+
+  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: 120
+    needs: [pre-commit, pytest]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
+            num_gpus: 1
+            axolotl_extras:
+            nightly_build: "true"
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
+            num_gpus: 1
+            axolotl_extras:
+            nightly_build: "true"
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - name: Install Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.11"
+      - name: Install Modal
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal==1.3.0.post1 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 "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
+      - name: Run tests job on Modal
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+        run: |
+          modal run -m cicd.e2e_tests
+  docker-e2e-multigpu-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: 120
+    needs: [pre-commit, pytest, docker-e2e-tests, docker-e2e-kernel-tests]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
+            num_gpus: 2
+            axolotl_extras:
+            nightly_build: "true"
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - name: Install Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.11"
+      - name: Install Modal
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal==1.3.0.post1 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 "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+      - name: Run tests job on Modal
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+        run: |
+          modal run -m cicd.multigpu
+
+  docker-e2e-kernel-tests:
+    if: github.repository_owner == 'axolotl-ai-cloud'
+    runs-on: [self-hosted, modal]
+    timeout-minutes: 90
+    needs: [pre-commit, pytest]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
+            num_gpus: 1
+            axolotl_extras:
+            nightly_build: "true"
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
+            num_gpus: 1
+            axolotl_extras:
+            nightly_build: "true"
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - name: Install Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.11"
+      - name: Install Modal
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal==1.3.0.post1 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 "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
+      - name: Run tests job on Modal
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+        run: |
+          modal run -m cicd.e2e_cuda_kernels
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index a53640e0b0..8056b3f3c5 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -1,107 +1,495 @@
 name: Tests
 on:
   # check on push/merge to main, PRs, and manual triggers
+  merge_group:
   push:
     branches:
       - "main"
     paths:
-      - '**.py'
-      - 'requirements.txt'
-      - '.github/workflows/*.yml'
+      - "**.py"
+      - "pyproject.toml"
+      - ".github/workflows/*.yml"
+      - "cicd/cicd.sh"
+      - "cicd/cicd_cuda_kernels.sh"
+      - "cicd/Dockerfile-uv.jinja"
   pull_request:
-      paths:
-       - '**.py'
-       - 'requirements.txt'
-       - '.github/workflows/*.yml'
+    types: [opened, synchronize, reopened, ready_for_review]
+    paths:
+      - "**.py"
+      - "pyproject.toml"
+      - ".github/workflows/*.yml"
+      - "cicd/cicd.sh"
+      - "cicd/cicd_cuda_kernels.sh"
+      - "cicd/Dockerfile-uv.jinja"
   workflow_dispatch:
 
+# Cancel jobs on the same ref if a new one is triggered
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+permissions:
+  contents: read
+
+env:
+  TRANSFORMERS_IS_CI: "yes"
+  UV_SYSTEM_PYTHON: "1"
+
 jobs:
   pre-commit:
     name: pre-commit
     runs-on: ubuntu-latest
+    if: ${{ !github.event.pull_request.draft }}
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
-          python-version: "3.10"
-          cache: 'pip' # caching pip dependencies
-      - uses: pre-commit/action@v3.0.0
+          python-version: "3.11"
+          cache: "pip" # caching pip dependencies
+      - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
+        env:
+          # check-cli-config-options needs an installed axolotl; the dedicated
+          # tests.yml step covers it in CI
+          SKIP: no-commit-to-branch,check-cli-config-options
+
+  prime-cdn-s3-cache:
+    name: Prefetch S3 once to prime the CDN cache
+    runs-on: ubuntu-latest
+    if: ${{ !github.event.pull_request.draft }}
+    timeout-minutes: 10
+    steps:
+      - name: Restore Cache from S3
+        id: hf-cache-restore-s3
+        run: |
+          curl -v -H "Range: bytes=0-1023" -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst > /dev/null
 
   pytest:
     name: PyTest
     runs-on: ubuntu-latest
+    if: ${{ !github.event.pull_request.draft }}
+    needs: [prime-cdn-s3-cache]
+    strategy:
+      max-parallel: 2
+      fail-fast: false
+      matrix:
+        python_version: ["3.12", "3.14"]
+        pytorch_version: ["2.11.0", "2.12.0"]
+    timeout-minutes: 30
+
+    steps:
+      - name: cleanup node
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL
+
+      - name: Check out repository code
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+
+      - name: Restore Cache from S3
+        id: hf-cache-restore-s3
+        run: |
+          mkdir -p ~/.cache/huggingface/hub
+          curl -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst | tar -xpf - -C ~/.cache/huggingface/hub/  --use-compress-program unzstd --strip-components=1
+          ls -ltr ~/.cache/huggingface/hub/
+
+      - name: Setup Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: ${{ matrix.python_version }}
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        with:
+          enable-cache: true
+          cache-dependency-glob: "pyproject.toml"
+
+      - name: Install PyTorch
+        run: |
+          uv pip install torch==${{ matrix.pytorch_version }} torchvision
+          uv pip freeze | grep -E "^(torch|torchvision)==" > /tmp/torch-pin.txt
+
+      - name: Install dependencies
+        run: |
+          uv pip install --no-build-isolation -e . --override /tmp/torch-pin.txt
+          python scripts/cutcrossentropy_install.py --uv | sh
+          uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+            codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
+
+      - name: Make sure PyTorch version wasn't clobbered
+        run: |
+          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__, f'Expected torch ${{ matrix.pytorch_version }} but got {torch.__version__}'"
+
+      - name: Ensure axolotl CLI was installed
+        run: |
+          axolotl --help
+
+      - name: Check generated CLI config options
+        if: ${{ matrix.python_version == '3.12' && matrix.pytorch_version == '2.12.0' }}
+        run: |
+          axolotl generate-cli-config-options --check
+
+      - name: Pre-Download dataset fixture
+        run: |
+          hf download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures
+
+      - name: Show HF cache
+        run: hf cache ls
+
+      - name: Run tests
+        run: |
+          pytest -v --durations=10 -n4 --dist loadfile --ignore=tests/utils/ --ignore=tests/integrations/ --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ --ignore=tests/monkeypatch/ tests/ --cov=axolotl --cov-report=xml
+          pytest -v --durations=10 tests/monkeypatch/ --cov=axolotl --cov-append --cov-report=xml
+          pytest -v --durations=10 tests/patched/ --cov=axolotl --cov-append --cov-report=xml
+          pytest -v --durations=10 tests/cli/ --cov=axolotl --cov-append --cov-report=xml
+          pytest -v --durations=10 tests/utils/ --cov=axolotl --cov-append --cov-report=xml
+          pytest -v --durations=10 \
+            --ignore=tests/integrations/kernels/ \
+            --ignore=tests/integrations/monkeypatch/test_tiled_mlp_moe.py \
+            --ignore=tests/integrations/test_gemma4_moe.py \
+            --ignore=tests/integrations/test_scattermoe_lora.py \
+            --ignore=tests/integrations/test_scattermoe_lora_kernels.py \
+            --ignore=tests/integrations/test_scattermoe_multi_lora.py \
+            --ignore=tests/integrations/test_sonicmoe_multi_lora.py \
+            tests/integrations/ --cov=axolotl --cov-append --cov-report=xml
+
+      - name: Show HF cache
+        run: hf cache ls
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: ./coverage.xml
+          flags: unittests,pytorch-${{ matrix.pytorch_version }}
+          fail_ci_if_error: false
+
+  pytest-sdist:
+    name: PyTest from Source Dist
+    runs-on: ubuntu-latest
+    if: ${{ !github.event.pull_request.draft }}
+    needs: [prime-cdn-s3-cache]
     strategy:
+      max-parallel: 2
       fail-fast: false
       matrix:
-        python_version: ["3.10", "3.11"]
-    timeout-minutes: 20
+        python_version: ["3.12", "3.14"]
+        pytorch_version: ["2.11.0", "2.12.0"]
+    timeout-minutes: 30
 
     steps:
+      - name: cleanup node
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL
+
       - name: Check out repository code
-        uses: actions/checkout@v3
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+
+      - name: Restore Cache from S3
+        id: hf-cache-restore-s3
+        run: |
+          mkdir -p ~/.cache/huggingface/hub
+          curl -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst | tar -xpf - -C ~/.cache/huggingface/hub/  --use-compress-program unzstd --strip-components=1
+          ls -ltr ~/.cache/huggingface/hub/
 
       - name: Setup Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
           python-version: ${{ matrix.python_version }}
-          cache: 'pip' # caching pip dependencies
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        with:
+          enable-cache: true
+          cache-dependency-glob: "pyproject.toml"
+
+      - name: Install PyTorch
+        run: |
+          uv pip install torch==${{ matrix.pytorch_version }} torchvision
+          uv pip freeze | grep -E "^(torch|torchvision)==" > /tmp/torch-pin.txt
 
       - name: Install dependencies
         run: |
-          pip3 install --upgrade pip
-          pip3 install --upgrade packaging
-          pip3 install -U -e .
-          pip3 install -r requirements-tests.txt
+          uv pip install packaging setuptools_scm build wheel psutil
+          python -m build --no-isolation --sdist
+          uv pip install --no-build-isolation dist/axolotl*.tar.gz --override /tmp/torch-pin.txt
+          python scripts/cutcrossentropy_install.py --uv | sh
+          uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+            codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
+
+      - name: Make sure PyTorch version wasn't clobbered
+        run: |
+          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__, f'Expected torch ${{ matrix.pytorch_version }} but got {torch.__version__}'"
+
+      - name: Ensure axolotl CLI was installed
+        run: |
+          axolotl --help
+
+      - name: Check generated CLI config options
+        if: ${{ matrix.python_version == '3.12' && matrix.pytorch_version == '2.12.0' }}
+        run: |
+          axolotl generate-cli-config-options --check
+
+      - name: Verify agent docs are discoverable
+        run: |
+          # Agent docs live in docs/agents/ (source of truth) and are resolved
+          # at runtime from the repo checkout or via `axolotl fetch docs`
+          axolotl agent-docs --list
+          axolotl agent-docs | grep -q "Fine-tuning framework"
+          axolotl agent-docs grpo | grep -q "GRPO"
+          axolotl agent-docs sft | grep -q "SFT"
+          python -c "from axolotl.cli.agent_docs import get_doc, list_topics; assert len(list_topics()) >= 5; assert 'GRPO' in get_doc('grpo')"
+
+      - name: Show HF cache
+        run: hf cache ls
 
       - name: Run tests
         run: |
-          pytest --ignore=tests/e2e/ tests/
+          pytest -v --durations=10 -n4 --dist loadfile --ignore=tests/utils/ --ignore=tests/integrations/ --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ --ignore=tests/monkeypatch/ tests/ --cov=axolotl --cov-report=xml
+          pytest -v --durations=10 tests/monkeypatch/ --cov=axolotl --cov-append --cov-report=xml
+          pytest -v --durations=10 tests/patched/ --cov=axolotl --cov-append --cov-report=xml
+          pytest -v --durations=10 tests/cli/ --cov=axolotl --cov-append --cov-report=xml
+          pytest -v --durations=10 tests/utils/ --cov=axolotl --cov-append --cov-report=xml
+          pytest -v --durations=10 \
+            --ignore=tests/integrations/kernels/ \
+            --ignore=tests/integrations/monkeypatch/test_tiled_mlp_moe.py \
+            --ignore=tests/integrations/test_gemma4_moe.py \
+            --ignore=tests/integrations/test_scattermoe_lora.py \
+            --ignore=tests/integrations/test_scattermoe_lora_kernels.py \
+            --ignore=tests/integrations/test_scattermoe_multi_lora.py \
+            --ignore=tests/integrations/test_sonicmoe_multi_lora.py \
+            tests/integrations/ --cov=axolotl --cov-append --cov-report=xml
+
+      - name: Show HF cache
+        run: hf cache ls
+
+  gate-skip-e2e:
+    needs: [pre-commit]
+    runs-on: ubuntu-latest
+    outputs:
+      skip: ${{ steps.compute.outputs.skip }}
+    steps:
+      - uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
+        id: compute
+        with:
+          script: |
+            const token = /\[skip-e2e\]/i;
+            let msg = '';
+            if (context.eventName === 'push') {
+              msg = context.payload.head_commit?.message || '';
+            } else if (context.eventName === 'pull_request') {
+              const { owner, repo } = context.repo;
+              const prNumber = context.payload.pull_request.number;
+              const commits = await github.paginate(
+                github.rest.pulls.listCommits,
+                { owner, repo, pull_number: prNumber, per_page: 100 }
+              );
+              msg = commits.at(-1)?.commit?.message || '';
+            }
+            const title = context.payload.pull_request?.title || '';
+            const body  = context.payload.pull_request?.body  || '';
+            const skip = token.test(msg) || token.test(title) || token.test(body);
+            core.setOutput('skip', String(skip));
+
+  docker-e2e-tests-1st:
+    # Run this job first as a gate for running the remainder of the test matrix
+    if: >
+      github.repository_owner == 'axolotl-ai-cloud' &&
+      (github.event_name != 'pull_request' || !github.event.pull_request.draft) &&
+      needs.gate-skip-e2e.outputs.skip != 'true'
+    # this job needs to be run on self-hosted GPU runners...
+    runs-on: [self-hosted, modal]
+    timeout-minutes: 120
+    needs: [pre-commit, pytest, gate-skip-e2e]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
+            num_gpus: 1
+            axolotl_extras:
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - name: Install Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.11"
+      - name: Install Modal
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal==1.3.0.post1 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
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+      - name: Run tests job on Modal
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+        run: |
+          modal run -m cicd.e2e_tests
 
   docker-e2e-tests:
-    if: github.repository_owner == 'OpenAccess-AI-Collective'
+    if: >
+      github.repository_owner == 'axolotl-ai-cloud' &&
+      (github.event_name != 'pull_request' || !github.event.pull_request.draft) &&
+      needs.gate-skip-e2e.outputs.skip != 'true'
     # this job needs to be run on self-hosted GPU runners...
     runs-on: [self-hosted, modal]
-    timeout-minutes: 60
-    needs: [pre-commit, pytest]
+    timeout-minutes: 120
+    # Only run the remainder of the matrix if the first e2e check passed;
+    # this is to save on wasted compute costs for known failures that get caught in the first run
+    needs: [pre-commit, pytest, gate-skip-e2e, docker-e2e-tests-1st]
 
     strategy:
       fail-fast: false
       matrix:
         include:
-          - cuda: 118
-            cuda_version: 11.8.0
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_args: "--extra-index-url https://download.pytorch.org/whl/cu118"
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
             num_gpus: 1
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.10"
-            pytorch: 2.1.2
+            axolotl_extras:
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - name: Install Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.11"
+      - name: Install Modal
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal==1.3.0.post1 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
+          echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S'}}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+      - name: Run tests job on Modal
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+        run: |
+          modal run -m cicd.e2e_tests
+
+  docker-e2e-kernel-tests:
+    if: >
+      github.repository_owner == 'axolotl-ai-cloud' &&
+      (github.event_name != 'pull_request' || !github.event.pull_request.draft) &&
+      needs.gate-skip-e2e.outputs.skip != 'true'
+    runs-on: [self-hosted, modal]
+    timeout-minutes: 90
+    needs: [pre-commit, pytest, gate-skip-e2e, docker-e2e-tests-1st]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
+            num_gpus: 1
+            axolotl_extras:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
             num_gpus: 1
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.1
+            axolotl_extras:
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+      - name: Install Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.11"
+      - name: Install Modal
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal==1.3.0.post1 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
+          echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S'}}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+      - name: Run tests job on Modal
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+        run: |
+          modal run -m cicd.e2e_cuda_kernels
+
+  docker-e2e-cleanup:
+    runs-on: [self-hosted, modal]
+    timeout-minutes: 90
+    needs: [docker-e2e-tests, docker-e2e-kernel-tests]
+    if: ${{ !github.event.pull_request.draft }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.0
             num_gpus: 1
+            axolotl_extras:
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
       - name: Install Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
-          python-version: "3.10"
+          python-version: "3.11"
       - name: Install Modal
         run: |
           python -m pip install --upgrade pip
-          pip install modal jinja2
+          pip install modal==1.3.0.post1 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
+          modal run -m cicd.cleanup
diff --git a/.gitignore b/.gitignore
index e6dfee67db..8365917bf3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
 **/axolotl.egg-info
 configs
 last_run_prepared/
+outputs
 .vscode
 _site/
 
@@ -31,6 +32,7 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+uv.lock
 
 # PyInstaller
 #  Usually these files are written by a python script from a template
@@ -176,3 +178,19 @@ qlora-out/*
 mlruns/*
 
 /.quarto/
+prepared-datasets/
+submit.sh
+*.out*
+
+# Quartodoc generated files
+objects.json
+site_libs/
+
+typings/
+out/
+
+# vim
+*.swp
+
+# scm auto-versioning
+src/axolotl/_version.py
diff --git a/.isort.cfg b/.isort.cfg
deleted file mode 100644
index 79067a7c91..0000000000
--- a/.isort.cfg
+++ /dev/null
@@ -1,3 +0,0 @@
-[settings]
-profile=black
-known_third_party=wandb
diff --git a/.mypy.ini b/.mypy.ini
index ede9fef887..8e9b0fe43a 100644
--- a/.mypy.ini
+++ b/.mypy.ini
@@ -8,9 +8,16 @@ ignore_missing_imports = True
 [mypy-axolotl.monkeypatch.*]
 ignore_errors = True
 
+[mypy-axolotl.integrations.kernels.libs.scattermoe_lora.cutlass_fp4.*]
+# CuTe-DSL (cutlass-python): @cute.jit, cute.range(unroll=), DSL tensor types not analyzable.
+ignore_errors = True
+
 [mypy-axolotl.models.mixtral.*]
 ignore_errors = True
 
+[mypy-axolotl.integrations.liger.models.*]
+ignore_errors = True
+
 [mypy-axolotl.models.phi.*]
 ignore_errors = True
 
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6c5f205897..a3486461ca 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -3,29 +3,21 @@ default_language_version:
 
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v6.0.0
     hooks:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
--   repo: https://github.com/psf/black
-    rev: 23.3.0
+    -   id: no-commit-to-branch
+        args: ['--branch', 'main']
+-   repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.8
     hooks:
-    -   id: black
--   repo: https://github.com/pycqa/isort
-    rev: 5.12.0
-    hooks:
-      - id: isort
--   repo: https://github.com/PyCQA/flake8
-    rev: 6.0.0
-    hooks:
-    - id: flake8
--   repo: https://github.com/PyCQA/pylint
-    rev: v2.17.4
-    hooks:
-    - id: pylint
+    -   id: ruff
+        args: [--fix]
+    -   id: ruff-format
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.3.0
+    rev: v1.19.1
     hooks:
     - id: mypy
       additional_dependencies:
@@ -34,10 +26,18 @@ repos:
             'pydantic>=2.5.3',
         ]
 -   repo: https://github.com/PyCQA/bandit
-    rev: 1.7.5
+    rev: 1.9.4
     hooks:
     -   id: bandit
         args: [
             '--ini',
             '.bandit',
         ]
+-   repo: local
+    hooks:
+    -   id: check-cli-config-options
+        name: check generated CLI config options are up to date
+        entry: axolotl generate-cli-config-options --check
+        language: system
+        files: ^src/axolotl/(utils/schemas/.*\.py|cli/config_options\.py)$
+        pass_filenames: false
diff --git a/.pylintrc b/.pylintrc
deleted file mode 100644
index ed973d2859..0000000000
--- a/.pylintrc
+++ /dev/null
@@ -1,14 +0,0 @@
-[MASTER]
-init-hook="from pylint.config import find_pylintrc; import os, sys; sys.path.append(os.path.dirname(find_pylintrc()))"
-
-[TYPECHECK]
-
-# List of members which are set dynamically and missed by Pylint inference
-# system, and so shouldn't trigger E1101 when accessed.
-generated-members=numpy.*, torch.*
-
-
-[pylint.messages_control]
-disable=missing-function-docstring, line-too-long, import-error,
-    too-many-arguments, too-many-locals, too-many-statements, too-many-branches, too-few-public-methods,
-    too-many-instance-attributes, fixme, import-outside-toplevel, logging-fstring-interpolation,
diff --git a/.runpod/.gitignore b/.runpod/.gitignore
new file mode 100644
index 0000000000..383570cfc6
--- /dev/null
+++ b/.runpod/.gitignore
@@ -0,0 +1,161 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+pod/scripts/config.yaml
diff --git a/.runpod/Dockerfile b/.runpod/Dockerfile
new file mode 100644
index 0000000000..948d3f78e9
--- /dev/null
+++ b/.runpod/Dockerfile
@@ -0,0 +1,19 @@
+FROM axolotlai/axolotl-cloud:main-py3.11-cu124-2.6.0
+
+COPY .runpod/requirements.txt /requirements.txt
+RUN --mount=type=cache,target=/root/.cache/pip \
+    python3 -m pip install --upgrade pip && \
+    python3 -m pip install --upgrade -r /requirements.txt
+
+# Environment settings
+ARG BASE_VOLUME="/runpod-volume"
+ENV BASE_VOLUME=$BASE_VOLUME
+ENV HF_DATASETS_CACHE="${BASE_VOLUME}/huggingface-cache/datasets"
+ENV HUGGINGFACE_HUB_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
+ENV HF_HUB_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
+ENV TRANSFORMERS_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
+
+COPY .runpod/src /src
+
+WORKDIR /src
+CMD ["python3", "/src/handler.py"]
diff --git a/.runpod/README.md b/.runpod/README.md
new file mode 100644
index 0000000000..2cebaa5e7d
--- /dev/null
+++ b/.runpod/README.md
@@ -0,0 +1,335 @@
+<h1>LLM Post Training- Full fine-tune, LoRA, QLoRa etc. Llama/Mistral/Gemma and more</h1>
+
+# Configuration Options
+
+This document outlines all available configuration options for training models. The configuration can be provided as a JSON request.
+
+## Usage
+
+You can use these configuration Options:
+
+1. As a JSON request body:
+
+```json
+{
+  "input": {
+    "user_id": "user",
+    "model_id": "model-name",
+    "run_id": "run-id",
+    "credentials": {
+      "wandb_api_key": "", # add your Weights & biases key. TODO:  you will be able to set this in Enviornment variables.
+      "hf_token": "", # add your HF_token. TODO:  you will be able to set this in Enviornment variables.
+    },
+    "args": {
+      "base_model": "NousResearch/Llama-3.2-1B",
+      // ... other options
+    }
+  }
+}
+```
+
+## Configuration Options
+
+### Model Configuration
+
+| Option              | Description                                                                                   | Default              |
+| ------------------- | --------------------------------------------------------------------------------------------- | -------------------- |
+| `base_model`        | Path to the base model (local or HuggingFace)                                                 | Required             |
+| `base_model_config` | Configuration path for the base model                                                         | Same as base_model   |
+| `revision_of_model` | Specific model revision from HuggingFace hub                                                  | Latest               |
+| `tokenizer_config`  | Custom tokenizer configuration path                                                           | Optional             |
+| `model_type`        | Type of model to load                                                                         | AutoModelForCausalLM |
+| `tokenizer_type`    | Type of tokenizer to use                                                                      | AutoTokenizer        |
+| `hub_model_id`      | Repository ID where the model will be pushed on Hugging Face Hub (format: username/repo-name) | Optional             |
+
+## Model Family Identification
+
+| Option                     | Default | Description                    |
+| -------------------------- | ------- | ------------------------------ |
+| `is_falcon_derived_model`  | `false` | Whether model is Falcon-based  |
+| `is_llama_derived_model`   | `false` | Whether model is LLaMA-based   |
+| `is_qwen_derived_model`    | `false` | Whether model is Qwen-based    |
+| `is_mistral_derived_model` | `false` | Whether model is Mistral-based |
+
+## Model Configuration Overrides
+
+| Option                                          | Default    | Description                        |
+| ----------------------------------------------- | ---------- | ---------------------------------- |
+| `overrides_of_model_config.rope_scaling.type`   | `"linear"` | RoPE scaling type (linear/dynamic) |
+| `overrides_of_model_config.rope_scaling.factor` | `1.0`      | RoPE scaling factor                |
+
+### Model Loading Options
+
+| Option         | Description                   | Default |
+| -------------- | ----------------------------- | ------- |
+| `load_in_8bit` | Load model in 8-bit precision | false   |
+| `load_in_4bit` | Load model in 4-bit precision | false   |
+| `bf16`         | Use bfloat16 precision        | false   |
+| `fp16`         | Use float16 precision         | false   |
+| `tf32`         | Use tensor float 32 precision | false   |
+
+## Memory and Device Settings
+
+| Option             | Default   | Description             |
+| ------------------ | --------- | ----------------------- |
+| `gpu_memory_limit` | `"20GiB"` | GPU memory limit        |
+| `lora_on_cpu`      | `false`   | Load LoRA on CPU        |
+| `device_map`       | `"auto"`  | Device mapping strategy |
+| `max_memory`       | `null`    | Max memory per device   |
+
+## Training Hyperparameters
+
+| Option                        | Default   | Description                 |
+| ----------------------------- | --------- | --------------------------- |
+| `gradient_accumulation_steps` | `1`       | Gradient accumulation steps |
+| `micro_batch_size`            | `2`       | Batch size per GPU          |
+| `eval_batch_size`             | `null`    | Evaluation batch size       |
+| `num_epochs`                  | `4`       | Number of training epochs   |
+| `warmup_steps`                | `100`     | Warmup steps                |
+| `warmup_ratio`                | `0.05`    | Warmup ratio                |
+| `learning_rate`               | `0.00003` | Learning rate               |
+| `lr_quadratic_warmup`         | `false`   | Quadratic warmup            |
+| `logging_steps`               | `null`    | Logging frequency           |
+| `eval_steps`                  | `null`    | Evaluation frequency        |
+| `evals_per_epoch`             | `null`    | Evaluations per epoch       |
+| `save_strategy`               | `"epoch"` | Checkpoint saving strategy  |
+| `save_steps`                  | `null`    | Saving frequency            |
+| `saves_per_epoch`             | `null`    | Saves per epoch             |
+| `save_total_limit`            | `null`    | Maximum checkpoints to keep |
+| `max_steps`                   | `null`    | Maximum training steps      |
+
+### Dataset Configuration
+
+```yaml
+datasets:
+  - path: vicgalle/alpaca-gpt4 # HuggingFace dataset or TODO: You will be able to add the local path.
+    type: alpaca # Format type (alpaca, gpteacher, oasst, etc.)
+    ds_type: json # Dataset type
+    data_files: path/to/data # Source data files
+    train_on_split: train # Dataset split to use
+```
+
+## Chat Template Settings
+
+| Option                   | Default                          | Description            |
+| ------------------------ | -------------------------------- | ---------------------- |
+| `chat_template`          | `"tokenizer_default"`            | Chat template type     |
+| `chat_template_jinja`    | `null`                           | Custom Jinja template  |
+| `default_system_message` | `"You are a helpful assistant."` | Default system message |
+
+## Dataset Processing
+
+| Option                            | Default                    | Description                         |
+| --------------------------------- | -------------------------- | ----------------------------------- |
+| `dataset_prepared_path`           | `"data/last_run_prepared"` | Path for prepared dataset           |
+| `push_dataset_to_hub`             | `""`                       | Push dataset to HF hub              |
+| `dataset_num_proc`                | `4`                        | Number of preprocessing processes   |
+| `dataset_keep_in_memory`          | `false`                    | Keep dataset in memory              |
+| `shuffle_merged_datasets`         | `true`                     | Shuffle merged datasets             |
+| `shuffle_before_merging_datasets` | `false`                    | Shuffle each dataset before merging |
+| `dataset_exact_deduplication`     | `true`                     | Deduplicate datasets                |
+
+## LoRA Configuration
+
+| Option                     | Default                | Description                    |
+| -------------------------- | ---------------------- | ------------------------------ |
+| `adapter`                  | `"lora"`               | Adapter type (lora/qlora)      |
+| `lora_model_dir`           | `""`                   | Directory with pretrained LoRA |
+| `lora_r`                   | `8`                    | LoRA attention dimension       |
+| `lora_alpha`               | `16`                   | LoRA alpha parameter           |
+| `lora_dropout`             | `0.05`                 | LoRA dropout                   |
+| `lora_target_modules`      | `["q_proj", "v_proj"]` | Modules to apply LoRA          |
+| `lora_target_linear`       | `false`                | Target all linear modules      |
+| `peft_layers_to_transform` | `[]`                   | Layers to transform            |
+| `lora_modules_to_save`     | `[]`                   | Modules to save                |
+| `lora_fan_in_fan_out`      | `false`                | Fan in/out structure           |
+
+## Optimization Settings
+
+| Option                    | Default | Description                |
+| ------------------------- | ------- | -------------------------- |
+| `train_on_inputs`         | `false` | Train on input prompts     |
+| `group_by_length`         | `false` | Group by sequence length   |
+| `gradient_checkpointing`  | `false` | Use gradient checkpointing |
+| `early_stopping_patience` | `3`     | Early stopping patience    |
+
+## Learning Rate Scheduling
+
+| Option                     | Default    | Description          |
+| -------------------------- | ---------- | -------------------- |
+| `lr_scheduler`             | `"cosine"` | Scheduler type       |
+| `lr_scheduler_kwargs`      | `{}`       | Scheduler parameters |
+| `cosine_min_lr_ratio`      | `null`     | Minimum LR ratio     |
+| `cosine_constant_lr_ratio` | `null`     | Constant LR ratio    |
+| `lr_div_factor`            | `null`     | LR division factor   |
+
+## Optimizer Settings
+
+| Option                 | Default      | Description         |
+| ---------------------- | ------------ | ------------------- |
+| `optimizer`            | `"adamw_hf"` | Optimizer choice    |
+| `optim_args`           | `{}`         | Optimizer arguments |
+| `optim_target_modules` | `[]`         | Target modules      |
+| `weight_decay`         | `null`       | Weight decay        |
+| `adam_beta1`           | `null`       | Adam beta1          |
+| `adam_beta2`           | `null`       | Adam beta2          |
+| `adam_epsilon`         | `null`       | Adam epsilon        |
+| `max_grad_norm`        | `null`       | Gradient clipping   |
+
+## Attention Implementations
+
+| Option                     | Default | Description                   |
+| -------------------------- | ------- | ----------------------------- |
+| `flash_optimum`            | `false` | Use better transformers       |
+| `xformers_attention`       | `false` | Use xformers                  |
+| `flash_attention`          | `false` | Use flash attention           |
+| `flash_attn_cross_entropy` | `false` | Flash attention cross entropy |
+| `flash_attn_rms_norm`      | `false` | Flash attention RMS norm      |
+| `flash_attn_fuse_mlp`      | `false` | Fuse MLP operations           |
+| `sdp_attention`            | `false` | Use scaled dot product        |
+| `s2_attention`             | `false` | Use shifted sparse attention  |
+
+## Tokenizer Modifications
+
+| Option           | Default | Description                  |
+| ---------------- | ------- | ---------------------------- |
+| `special_tokens` | -       | Special tokens to add/modify |
+| `tokens`         | `[]`    | Additional tokens            |
+
+## Distributed Training
+
+| Option                  | Default | Description           |
+| ----------------------- | ------- | --------------------- |
+| `fsdp`                  | `null`  | FSDP configuration    |
+| `fsdp_config`           | `null`  | FSDP config options   |
+| `deepspeed`             | `null`  | Deepspeed config path |
+| `ddp_timeout`           | `null`  | DDP timeout           |
+| `ddp_bucket_cap_mb`     | `null`  | DDP bucket capacity   |
+| `ddp_broadcast_buffers` | `null`  | DDP broadcast buffers |
+
+<details>
+<summary><h3>Example Configuration Request:</h3></summary>
+
+Here's a complete example for fine-tuning a LLaMA model using LoRA:
+
+```json
+{
+  "input": {
+    "user_id": "user",
+    "model_id": "llama-test",
+    "run_id": "test-run",
+    "credentials": {
+      "wandb_api_key": "",
+      "hf_token": ""
+    },
+    "args": {
+      "base_model": "NousResearch/Llama-3.2-1B",
+      "load_in_8bit": false,
+      "load_in_4bit": false,
+      "strict": false,
+      "datasets": [
+        {
+          "path": "teknium/GPT4-LLM-Cleaned",
+          "type": "alpaca"
+        }
+      ],
+      "dataset_prepared_path": "last_run_prepared",
+      "val_set_size": 0.1,
+      "output_dir": "./outputs/lora-out",
+      "adapter": "lora",
+      "sequence_len": 2048,
+      "sample_packing": true,
+      "eval_sample_packing": true,
+      "pad_to_sequence_len": true,
+      "lora_r": 16,
+      "lora_alpha": 32,
+      "lora_dropout": 0.05,
+      "lora_target_modules": [
+        "gate_proj",
+        "down_proj",
+        "up_proj",
+        "q_proj",
+        "v_proj",
+        "k_proj",
+        "o_proj"
+      ],
+      "gradient_accumulation_steps": 2,
+      "micro_batch_size": 2,
+      "num_epochs": 1,
+      "optimizer": "adamw_8bit",
+      "lr_scheduler": "cosine",
+      "learning_rate": 0.0002,
+      "train_on_inputs": false,
+      "group_by_length": false,
+      "bf16": "auto",
+      "tf32": false,
+      "gradient_checkpointing": true,
+      "logging_steps": 1,
+      "flash_attention": true,
+      "loss_watchdog_threshold": 5,
+      "loss_watchdog_patience": 3,
+      "warmup_steps": 10,
+      "evals_per_epoch": 4,
+      "saves_per_epoch": 1,
+      "weight_decay": 0,
+      "hub_model_id": "runpod/llama-fr-lora",
+      "wandb_name": "test-run-1",
+      "wandb_project": "test-run-1",
+      "wandb_entity": "axo-test",
+      "special_tokens": {
+        "pad_token": "<|end_of_text|>"
+      }
+    }
+  }
+}
+```
+
+</details>
+
+### Advanced Features
+
+#### Wandb Integration
+
+- `wandb_project`: Project name for Weights & Biases
+- `wandb_entity`: Team name in W&B
+- `wandb_watch`: Monitor model with W&B
+- `wandb_name`: Name of the W&B run
+- `wandb_run_id`: ID for the W&B run
+
+#### Performance Optimization
+
+- `sample_packing`: Enable efficient sequence packing
+- `eval_sample_packing`: Use sequence packing during evaluation
+- `torch_compile`: Enable PyTorch 2.0 compilation
+- `flash_attention`: Use Flash Attention implementation
+- `xformers_attention`: Use xFormers attention implementation
+
+### Available Optimizers
+
+The following optimizers are supported:
+
+- `adamw_hf`: HuggingFace's AdamW implementation
+- `adamw_torch`: PyTorch's AdamW
+- `adamw_torch_fused`: Fused AdamW implementation
+- `adamw_torch_xla`: XLA-optimized AdamW
+- `adamw_apex_fused`: NVIDIA Apex fused AdamW
+- `adafactor`: Adafactor optimizer
+- `adamw_anyprecision`: Anyprecision AdamW
+- `adamw_bnb_8bit`: 8-bit AdamW from bitsandbytes
+- `lion_8bit`: 8-bit Lion optimizer
+- `lion_32bit`: 32-bit Lion optimizer
+- `sgd`: Stochastic Gradient Descent
+- `adagrad`: Adagrad optimizer
+
+## Notes
+
+- Set `load_in_8bit: true` or `load_in_4bit: true` for memory-efficient training
+- Enable `flash_attention: true` for faster training on modern GPUs
+- Use `gradient_checkpointing: true` to reduce memory usage
+- Adjust `micro_batch_size` and `gradient_accumulation_steps` based on your GPU memory
+
+For more detailed information, please refer to the [documentation](https://axolotl-ai-cloud.github.io/axolotl/docs/config-reference.html).
+
+### Errors:
+
+- if you face any issues with the Flash Attention-2, Delete yoor worker and Re-start.
diff --git a/.runpod/hub.json b/.runpod/hub.json
new file mode 100644
index 0000000000..a243a27d80
--- /dev/null
+++ b/.runpod/hub.json
@@ -0,0 +1,93 @@
+{
+  "title": "Axolotl Fine-Tuning",
+  "description": "Serverless fine-tuning of open-source LLMs with Axolotl. Supports LoRA, QLoRA, DPO, and more using Hugging Face models and datasets.",
+  "type": "serverless",
+  "category": "language",
+  "iconUrl": "https://avatars.githubusercontent.com/u/167502477",
+  "config": {
+    "runsOn": "GPU",
+    "containerDiskInGb": 200,
+    "gpuCount": 1,
+    "allowedCudaVersions": [
+      "12.8",
+      "12.7",
+      "12.6",
+      "12.5",
+      "12.4"
+    ],
+    "presets": [],
+    "env": [
+      {
+        "key": "TOKENIZER",
+        "input": {
+          "name": "Tokenizer",
+          "type": "string",
+          "description": "Name or path of the Hugging Face tokenizer to use.",
+          "default": "",
+          "advanced": true
+        }
+      },
+      {
+        "key": "MAX_NUM_SEQS",
+        "input": {
+          "name": "Max Num Seqs",
+          "type": "number",
+          "description": "Maximum number of sequences per iteration.",
+          "default": 256,
+          "advanced": true
+        }
+      },
+      {
+        "key": "DISABLE_LOG_STATS",
+        "input": {
+          "name": "Disable Log Stats",
+          "type": "boolean",
+          "description": "Disable logging statistics.",
+          "default": false,
+          "trueValue": "true",
+          "falseValue": "false"
+        }
+      },
+      {
+        "key": "LOAD_FORMAT",
+        "input": {
+          "name": "Load Format",
+          "type": "string",
+          "description": "The format of the model weights to load.",
+          "default": "auto",
+          "options": [
+            {
+              "label": "auto",
+              "value": "auto"
+            },
+            {
+              "label": "pt",
+              "value": "pt"
+            },
+            {
+              "label": "safetensors",
+              "value": "safetensors"
+            },
+            {
+              "label": "npcache",
+              "value": "npcache"
+            },
+            {
+              "label": "dummy",
+              "value": "dummy"
+            },
+            {
+              "label": "tensorizer",
+              "value": "tensorizer"
+            },
+            {
+              "label": "bitsandbytes",
+              "value": "bitsandbytes"
+            }
+          ],
+          "advanced": true
+        }
+      }
+    ]
+  }
+}
diff --git a/.runpod/requirements.txt b/.runpod/requirements.txt
new file mode 100644
index 0000000000..345bdda350
--- /dev/null
+++ b/.runpod/requirements.txt
@@ -0,0 +1,7 @@
+# Required Python packages get listed here, one per line.
+# Reccomended to lock the version number to avoid unexpected changes.
+
+# You can also install packages from a git repository, e.g.:
+# git+https://github.com/runpod/runpod-python.git
+# To learn more, see https://pip.pypa.io/en/stable/reference/requirements-file-format/
+runpod~=1.7.0
diff --git a/.runpod/src/config/config.yaml b/.runpod/src/config/config.yaml
new file mode 100644
index 0000000000..b43c83dfec
--- /dev/null
+++ b/.runpod/src/config/config.yaml
@@ -0,0 +1,564 @@
+# # This is the huggingface model that contains *.pt, *.safetensors, or *.bin files
+# # This can also be a relative path to a model on disk
+# base_model: ./llama-7b-hf
+# # You can specify an ignore pattern if the model repo contains more than 1 model type (*.pt, etc)
+# base_model_ignore_patterns:
+# # If the base_model repo on hf hub doesn't include configuration .json files,
+# # You can set that here, or leave this empty to default to base_model
+# base_model_config: ./llama-7b-hf
+# # You can specify to choose a specific model revision from huggingface hub
+# model_revision:
+# # Optional tokenizer configuration override in case you want to use a different tokenizer
+# # than the one defined in the base model
+# tokenizer_config:
+# # If you want to specify the type of model to load, AutoModelForCausalLM is a good choice too
+# model_type: AutoModelForCausalLM
+# # Corresponding tokenizer for the model AutoTokenizer is a good choice
+# tokenizer_type: AutoTokenizer
+# # Trust remote code for untrusted source
+# trust_remote_code:
+# # use_fast option for tokenizer loading from_pretrained, default to True
+# tokenizer_use_fast:
+# # Whether to use the legacy tokenizer setting, defaults to True
+# 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:
+
+# # Used to identify which the model is based on
+# is_falcon_derived_model:
+# is_llama_derived_model:
+# # Please note that if you set this to true, `padding_side` will be set to "left" by default
+# is_mistral_derived_model:
+# is_qwen_derived_model:
+
+# # optional overrides to the base model configuration
+# model_config:
+#   # RoPE Scaling https://github.com/huggingface/transformers/pull/24653
+#   rope_scaling:
+#     type: # linear | dynamic
+#     factor: # float
+
+# # Whether you are training a 4-bit GPTQ quantized model
+# gptq: true
+# gptq_groupsize: 128 # group size
+# gptq_model_v1: false # v1 or v2
+
+# # This will attempt to quantize the model down to 8 bits and use adam 8 bit optimizer
+# load_in_8bit: true
+# # Use bitsandbytes 4 bit
+# load_in_4bit:
+
+# # Use CUDA bf16
+# bf16: true # bool or 'full' for `bf16_full_eval`. require >=ampere
+# # Use CUDA fp16
+# fp16: true
+# # Use CUDA tf32
+# tf32: true # require >=ampere
+
+# # No AMP (automatic mixed precision)
+# bfloat16: true # require >=ampere
+# float16: true
+
+# # 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
+#   - path: vicgalle/alpaca-gpt4
+#   # The type of prompt to use for training. [alpaca, sharegpt, gpteacher, oasst, reflection]
+#     type: alpaca # format | format:<prompt_style> (chat/instruct) | <prompt_strategies>.load_<load_fn>
+#     ds_type: # Optional[str] (json|arrow|parquet|text|csv) defines the datatype when path is a file
+#     data_files: # Optional[str] path to source data files
+#     shards: # Optional[int] number of shards to split data into
+#     name: # Optional[str] name of dataset configuration to load
+#     train_on_split: train # Optional[str] name of dataset split to load from
+
+#     # Optional[str] fastchat conversation type, only used with type: sharegpt
+#     conversation:  # Options (see Conversation 'name'): https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
+#     field_human: # Optional[str]. Human key to use for conversation.
+#     field_model: # Optional[str]. Assistant key to use for conversation.
+
+#   # Custom user prompt
+#   - path: repo
+#     type:
+#       # The below are defaults. only set what's needed.
+#       system_prompt: ""
+#       system_format: "{system}"
+#       field_system: system
+#       field_instruction: instruction
+#       field_input: input
+#       field_output: output
+
+#       # Customizable to be single line or multi-line
+#       # 'format' can include {input}
+#       format: |-
+#         User: {instruction} {input}
+#         Assistant:
+#       # 'no_input_format' cannot include {input}
+#       no_input_format: "{instruction} "
+
+#       # For `completion` datasets only, uses the provided field instead of `text` column
+#       field:
+
+# # Axolotl attempts to save the dataset as an arrow after packing the data together so
+# # subsequent training attempts load faster, relative path
+# dataset_prepared_path: data/last_run_prepared
+# # Push prepared dataset to hub
+# push_dataset_to_hub: # repo path
+# # The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()`
+# # if not set.
+# dataset_num_proc: # defaults to os.cpu_count() if not set
+# # push checkpoints to hub
+# hub_model_id: # repo path to push finetuned model
+# # how to push checkpoints to hub
+# # https://huggingface.co/docs/transformers/v4.31.0/en/main_classes/trainer#transformers.TrainingArguments.hub_strategy
+# hub_strategy:
+# # Whether to use hf `use_auth_token` for loading datasets. Useful for fetching private datasets
+# # Required to be true when used in combination with `push_dataset_to_hub`
+# hf_use_auth_token: # boolean
+# # How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc. 0 for no eval.
+# val_set_size: 0.04
+# # Num shards for whole dataset
+# dataset_shard_num:
+# # Index of shard to use for whole dataset
+# dataset_shard_idx:
+
+# # The maximum length of an input to train with, this should typically be less than 2048
+# # as most models have a token/context limit of 2048
+# sequence_len: 2048
+# # Pad inputs so each step uses constant sized buffers
+# # This will reduce memory fragmentation and may prevent OOMs, by re-using memory more efficiently
+# pad_to_sequence_len:
+# # Max sequence length to concatenate training samples together up to
+# # Inspired by StackLLaMA. see https://huggingface.co/blog/stackllama#supervised-fine-tuning
+# # FutureWarning: This will soon be DEPRECATED
+# max_packed_sequence_len: 1024
+# # Use efficient multi-packing with block diagonal attention and per sequence position_ids. Recommend set to 'true'
+# sample_packing:
+# # Set to 'false' if getting errors during eval with sample_packing on.
+# eval_sample_packing:
+# # You can set these packing optimizations AFTER starting a training at least once.
+# # The trainer will provide recommended values for these values.
+# sample_packing_eff_est:
+# total_num_tokens:
+
+# # If you want to use 'lora' or 'qlora' or leave blank to train all parameters in original model
+# adapter: lora
+# # If you already have a lora model trained that you want to load, put that here.
+# # This means after training, if you want to test the model, you should set this to the value of `lora_out_dir`.
+# lora_model_dir:
+
+# # LoRA hyperparameters
+# # For more details about the following options, see:
+# # https://www.anyscale.com/blog/fine-tuning-llms-lora-or-full-parameter-an-in-depth-analysis-with-llama-2
+# lora_r: 8
+# lora_alpha: 16
+# lora_dropout: 0.05
+# lora_target_modules:
+#   - q_proj
+#   - v_proj
+# #  - k_proj
+# #  - o_proj
+# #  - gate_proj
+# #  - down_proj
+# #  - up_proj
+# lora_target_linear: # If true, will target all linear 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.
+# # `embed_tokens` converts tokens to embeddings, and `lm_head` converts embeddings to token probabilities.
+# # https://github.com/huggingface/peft/issues/334#issuecomment-1561727994
+# lora_modules_to_save:
+# #  - embed_tokens
+# #  - lm_head
+
+# # Once you complete training, the model will be saved to the following directory.
+# # If you merge the adapter to the base model, a subdirectory `merged` will be created under this directory.
+# # Make sure `lora_model_dir` points to this directory if you want to use the trained model.
+# lora_out_dir:
+# lora_fan_in_fan_out: false
+
+# # ReLoRA configuration
+# # Must use either 'lora' or 'qlora' adapter, and does not support fsdp or deepspeed
+# relora_steps: # Number of steps per ReLoRA restart
+# relora_warmup_steps: # Number of per-restart warmup steps
+# relora_cpu_offload: # True to perform lora weight merges on cpu during restarts, for modest gpu memory savings
+
+# # wandb configuration if you're using it
+# wandb_mode: # "offline" to save run metadata locally and not sync to the server, "disabled" to turn off wandb
+# wandb_project: # Your wandb project name
+# wandb_entity: # A wandb Team name if using a Team
+# wandb_watch:
+# wandb_run_id: # Set the name of your wandb run
+# wandb_log_model: # "checkpoint" to log model to wandb Artifacts every `save_steps` or "end" to log only at the end of training
+
+# # Where to save the full-finetuned model to
+# output_dir: ./completed-model
+
+# # Whether to use torch.compile and which backend to use
+# torch_compile:  # bool
+# torch_compile_backend:  # Optional[str]
+
+# # Training hyperparameters
+
+# # If greater than 1, backpropagation will be skipped and the gradients will be accumulated for the given number of steps.
+# gradient_accumulation_steps: 1
+# # The number of samples to include in each batch. This is the number of samples sent to each GPU.
+# micro_batch_size: 2
+# eval_batch_size:
+# num_epochs: 4
+# warmup_steps: 100  # cannot use with warmup_ratio
+# warmup_ratio: 0.05  # cannot use with warmup_steps
+# learning_rate: 0.00003
+# lr_quadratic_warmup:
+# logging_steps:
+# save_strategy: # Set to `no` to skip checkpoint saves
+# save_steps: # Leave empty to save at each epoch
+# eval_steps: # Leave empty to eval at each epoch, integers for every N steps. decimal for fraction of total steps
+# save_total_limit: # Checkpoints saved at a time
+# # Maximum number of iterations to train for. It precedes num_epochs which means that
+# # if both are set, num_epochs will not be guaranteed.
+# # e.g., when 1 epoch is 1000 steps => `num_epochs: 2` and `max_steps: 100` will train for 100 steps
+# max_steps:
+
+# eval_table_size: # Approximate number of predictions sent to wandb depending on batch size. Enabled above 0. Default is 0
+# eval_table_max_new_tokens: # Total number of tokens generated for predictions sent to wandb. Default is 128
+
+# # Whether to mask out or include the human's prompt from the training labels
+# train_on_inputs: false
+# # Group similarly sized data to minimize padding.
+# # May be slower to start, as it must download and sort the entire dataset.
+# # Note that training loss may have an oscillating pattern with this enabled.
+# group_by_length: false
+
+# # Whether to use gradient checkpointing https://huggingface.co/docs/transformers/v4.18.0/en/performance#gradient-checkpointing
+# gradient_checkpointing: false
+
+# # Stop training after this many evaluation losses have increased in a row
+# # https://huggingface.co/transformers/v4.2.2/_modules/transformers/trainer_callback.html#EarlyStoppingCallback
+# early_stopping_patience: 3
+
+# # Specify a scheduler and kwargs to use with the optimizer
+# lr_scheduler: # 'one_cycle' | empty for cosine
+# lr_scheduler_kwargs:
+
+# # For one_cycle optim
+# 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/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_apex_fused
+# # - adafactor
+# # - adamw_anyprecision
+# # - sgd
+# # - adagrad
+# # - adamw_bnb_8bit
+# # - lion_8bit
+# # - lion_32bit
+# # - paged_adamw_32bit
+# # - paged_adamw_8bit
+# # - paged_lion_32bit
+# # - paged_lion_8bit
+# optimizer:
+# # Specify weight decay
+# weight_decay:
+# # adamw hyperparams
+# adam_beta1:
+# adam_beta2:
+# adam_epsilon:
+# # Gradient clipping max norm
+# max_grad_norm:
+
+# # Augmentation techniques
+# # NEFT https://arxiv.org/abs/2310.05914, set this to a number (paper default is 5) to add noise to embeddings
+# # currently only supported on Llama and Mistral
+# noisy_embedding_alpha:
+
+# # Whether to bettertransformers
+# flash_optimum:
+# # Whether to use xformers attention patch https://github.com/facebookresearch/xformers:
+# xformers_attention:
+# # Whether to use flash attention patch https://github.com/Dao-AILab/flash-attention:
+# flash_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_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:
+# # Landmark attention (only llama)
+# landmark_attention:
+# # xpos RoPE see https://github.com/kaiokendev/cutoff-len-is-context-len/blob/main/util/xpos_rope_llama_monkey_patch.py
+# # LLaMA only
+# xpos_rope:
+
+# # Resume from a specific checkpoint dir
+# resume_from_checkpoint:
+# # If resume_from_checkpoint isn't set and you simply want it to start where it left off.
+# # Be careful with this being turned on between different models.
+# auto_resume_from_checkpoints: false
+
+# # Don't mess with this, it's here for accelerate and torchrun
+# local_rank:
+
+# # Add or change special tokens.
+# # If you add tokens here, you don't need to add them to the `tokens` list.
+# special_tokens:
+#   # bos_token: "<s>"
+#   # eos_token: "</s>"
+#   # unk_token: "<unk>"
+
+# # Add extra tokens.
+# tokens:
+
+# # FSDP
+# fsdp:
+# fsdp_config:
+
+# # Deepspeed config path. e.g., deepspeed/zero3.json
+# deepspeed:
+
+# # Advanced DDP Arguments
+# ddp_timeout:
+# ddp_bucket_cap_mb:
+# ddp_broadcast_buffers:
+
+# # Path to torch distx for optim 'adamw_anyprecision'
+# torchdistx_path:
+
+# # Set to HF dataset for type: 'completion' for streaming instead of pre-tokenize
+# pretraining_dataset:
+
+# # Debug mode
+# debug:
+
+# # Seed
+# seed:
+
+# # Allow overwrite yml config using from cli
+# strict:
+
+base_model: ${BASE_MODEL}
+base_model_ignore_patterns: ${BASE_MODEL_IGNORE_PATTERNS}
+base_model_config: ${BASE_MODEL_CONFIG}
+revision_of_model: ${REVISION_OF_MODEL}
+tokenizer_config: ${TOKENIZER_CONFIG}
+model_type: ${MODEL_TYPE}
+tokenizer_type: ${TOKENIZER_TYPE}
+trust_remote_code: ${TRUST_REMOTE_CODE}
+tokenizer_use_fast: ${TOKENIZER_USE_FAST}
+tokenizer_legacy: ${TOKENIZER_LEGACY}
+resize_token_embeddings_to_32x: ${RESIZE_TOKEN_EMBEDDINGS_TO_32X}
+
+is_falcon_derived_model: ${IS_FALCON_DERIVED_MODEL}
+is_llama_derived_model: ${IS_LLAMA_DERIVED_MODEL}
+is_qwen_derived_model: ${IS_QWEN_DERIVED_MODEL}
+is_mistral_derived_model: ${IS_MISTRAL_DERIVED_MODEL}
+
+overrides_of_model_config:
+  rope_scaling:
+    type: ${ROPE_SCALING_TYPE}
+    factor: ${ROPE_SCALING_FACTOR}
+
+bnb_config_kwargs:
+  llm_int8_has_fp16_weight: ${BNB_LLM_INT8_HAS_FP16_WEIGHT}
+  bnb_4bit_quant_type: ${BNB_4BIT_QUANT_TYPE}
+  bnb_4bit_use_double_quant: ${BNB_4BIT_USE_DOUBLE_QUANT}
+
+gptq: ${GPTQ}
+load_in_8bit: ${LOAD_IN_8BIT}
+load_in_4bit: ${LOAD_IN_4BIT}
+bf16: ${BF16}
+fp16: ${FP16}
+tf32: ${TF32}
+bfloat16: ${BFLOAT16}
+float16: ${FLOAT16}
+
+gpu_memory_limit: ${GPU_MEMORY_LIMIT}
+lora_on_cpu: ${LORA_ON_CPU}
+
+datasets:
+  - path: ${DATASET_PATH}
+    type: ${DATASET_TYPE}
+    ds_type: ${DATASET_DS_TYPE}
+    data_files: ${DATASET_DATA_FILES}
+    shards: ${DATASET_SHARDS}
+    name: ${DATASET_NAME}
+    train_on_split: ${DATASET_TRAIN_ON_SPLIT}
+    revision: ${DATASET_REVISION}
+    trust_remote_code: ${DATASET_TRUST_REMOTE_CODE}
+
+rl: ${RL}
+dpo_use_weighting: ${DPO_USE_WEIGHTING}
+
+chat_template: ${CHAT_TEMPLATE}
+chat_template_jinja: ${CHAT_TEMPLATE_JINJA}
+default_system_message: ${DEFAULT_SYSTEM_MESSAGE}
+dataset_prepared_path: ${DATASET_PREPARED_PATH}
+push_dataset_to_hub: ${PUSH_DATASET_TO_HUB}
+dataset_num_proc: ${DATASET_NUM_PROC}
+dataset_keep_in_memory: ${DATASET_KEEP_IN_MEMORY}
+hub_model_id: ${HUB_MODEL_ID}
+hub_strategy: ${HUB_STRATEGY}
+hf_use_auth_token: ${HF_USE_AUTH_TOKEN}
+val_set_size: ${VAL_SET_SIZE}
+dataset_shard_num: ${DATASET_SHARD_NUM}
+dataset_shard_idx: ${DATASET_SHARD_IDX}
+
+sequence_len: ${SEQUENCE_LEN}
+pad_to_sequence_len: ${PAD_TO_SEQUENCE_LEN}
+sample_packing: ${SAMPLE_PACKING}
+eval_sample_packing: ${EVAL_SAMPLE_PACKING}
+sample_packing_eff_est: ${SAMPLE_PACKING_EFF_EST}
+total_num_tokens: ${TOTAL_NUM_TOKENS}
+sample_packing_group_size: ${SAMPLE_PACKING_GROUP_SIZE}
+sample_packing_bin_size: ${SAMPLE_PACKING_BIN_SIZE}
+
+batch_flattening: ${BATCH_FLATTENING}
+device_map: ${DEVICE_MAP}
+max_memory: ${MAX_MEMORY}
+
+adapter: ${ADAPTER}
+lora_model_dir: ${LORA_MODEL_DIR}
+
+lora_r: ${LORA_R}
+lora_alpha: ${LORA_ALPHA}
+lora_dropout: ${LORA_DROPOUT}
+lora_target_modules:
+  - ${LORA_TARGET_MODULES}
+lora_target_linear: ${LORA_TARGET_LINEAR}
+peft_layers_to_transform: ${PEFT_LAYERS_TO_TRANSFORM}
+lora_modules_to_save: ${LORA_MODULES_TO_SAVE}
+lora_fan_in_fan_out: ${LORA_FAN_IN_FAN_OUT}
+
+loraplus_lr_ratio: ${LORAPLUS_LR_RATIO}
+loraplus_lr_embedding: ${LORAPLUS_LR_EMBEDDING}
+
+peft:
+  loftq_config:
+    loftq_bits: ${LOFTQ_BITS}
+
+relora_steps: ${RELORA_STEPS}
+relora_warmup_steps: ${RELORA_WARMUP_STEPS}
+relora_anneal_steps: ${RELORA_ANNEAL_STEPS}
+relora_prune_ratio: ${RELORA_PRUNE_RATIO}
+relora_cpu_offload: ${RELORA_CPU_OFFLOAD}
+
+wandb_mode: ${WANDB_MODE}
+wandb_project: ${WANDB_PROJECT}
+wandb_entity: ${WANDB_ENTITY}
+wandb_watch: ${WANDB_WATCH}
+wandb_name: ${WANDB_NAME}
+wandb_run_id: ${WANDB_RUN_ID}
+wandb_log_model: ${WANDB_LOG_MODEL}
+
+mlflow_tracking_uri: ${MLFLOW_TRACKING_URI}
+mlflow_experiment_name: ${MLFLOW_EXPERIMENT_NAME}
+mlflow_run_name: ${MLFLOW_RUN_NAME}
+hf_mlflow_log_artifacts: ${HF_MLFLOW_LOG_ARTIFACTS}
+
+use_comet: ${USE_COMET}
+comet_api_key: ${COMET_API_KEY}
+comet_workspace: ${COMET_WORKSPACE}
+comet_project_name: ${COMET_PROJECT_NAME}
+comet_experiment_key: ${COMET_EXPERIMENT_KEY}
+comet_mode: ${COMET_MODE}
+comet_online: ${COMET_ONLINE}
+comet_experiment_config: ${COMET_EXPERIMENT_CONFIG}
+
+output_dir: ${OUTPUT_DIR}
+
+torch_compile: ${TORCH_COMPILE}
+torch_compile_backend: ${TORCH_COMPILE_BACKEND}
+
+gradient_accumulation_steps: ${GRADIENT_ACCUMULATION_STEPS}
+micro_batch_size: ${MICRO_BATCH_SIZE}
+eval_batch_size: ${EVAL_BATCH_SIZE}
+num_epochs: ${NUM_EPOCHS}
+warmup_steps: ${WARMUP_STEPS}
+warmup_ratio: ${WARMUP_RATIO}
+learning_rate: ${LEARNING_RATE}
+lr_quadratic_warmup: ${LR_QUADRATIC_WARMUP}
+logging_steps: ${LOGGING_STEPS}
+eval_steps: ${EVAL_STEPS}
+evals_per_epoch: ${EVALS_PER_EPOCH}
+save_strategy: ${SAVE_STRATEGY}
+save_steps: ${SAVE_STEPS}
+saves_per_epoch: ${SAVES_PER_EPOCH}
+save_total_limit: ${SAVE_TOTAL_LIMIT}
+max_steps: ${MAX_STEPS}
+
+eval_table_size: ${EVAL_TABLE_SIZE}
+eval_max_new_tokens: ${EVAL_MAX_NEW_TOKENS}
+eval_causal_lm_metrics: ${EVAL_CAUSAL_LM_METRICS}
+
+profiler_steps: ${PROFILER_STEPS}
+loss_watchdog_threshold: ${LOSS_WATCHDOG_THRESHOLD}
+loss_watchdog_patience: ${LOSS_WATCHDOG_PATIENCE}
+
+train_on_inputs: ${TRAIN_ON_INPUTS}
+group_by_length: ${GROUP_BY_LENGTH}
+gradient_checkpointing: ${GRADIENT_CHECKPOINTING}
+early_stopping_patience: ${EARLY_STOPPING_PATIENCE}
+
+lr_scheduler: ${LR_SCHEDULER}
+lr_scheduler_kwargs: ${LR_SCHEDULER_KWARGS}
+cosine_min_lr_ratio: ${COSINE_MIN_LR_RATIO}
+cosine_constant_lr_ratio: ${COSINE_CONSTANT_LR_RATIO}
+lr_div_factor: ${LR_DIV_FACTOR}
+
+optimizer: ${OPTIMIZER}
+optim_args: ${OPTIM_ARGS}
+optim_target_modules: ${OPTIM_TARGET_MODULES}
+weight_decay: ${WEIGHT_DECAY}
+adam_beta1: ${ADAM_BETA1}
+adam_beta2: ${ADAM_BETA2}
+adam_epsilon: ${ADAM_EPSILON}
+max_grad_norm: ${MAX_GRAD_NORM}
+
+neftune_noise_alpha: ${NEFTUNE_NOISE_ALPHA}
+
+flash_optimum: ${FLASH_OPTIMUM}
+xformers_attention: ${XFORMERS_ATTENTION}
+flash_attention: ${FLASH_ATTENTION}
+flash_attn_cross_entropy: ${FLASH_ATTN_CROSS_ENTROPY}
+flash_attn_rms_norm: ${FLASH_ATTN_RMS_NORM}
+flash_attn_fuse_mlp: ${FLASH_ATTN_FUSE_MLP}
+sdp_attention: ${SDP_ATTENTION}
+s2_attention: ${S2_ATTENTION}
+resume_from_checkpoint: ${RESUME_FROM_CHECKPOINT}
+auto_resume_from_checkpoints: ${AUTO_RESUME_FROM_CHECKPOINTS}
+
+local_rank: ${LOCAL_RANK}
+
+special_tokens:
+  bos_token: ${SPECIAL_TOKEN_BOS}
+  eos_token: ${SPECIAL_TOKEN_EOS}
+  unk_token: ${SPECIAL_TOKEN_UNK}
+  pad_token: ${SPECIAL_TOKEN_PAD}
+
+tokens: ${TOKENS}
+
+fsdp: ${FSDP}
+fsdp_config: ${FSDP_CONFIG}
+deepspeed: ${DEEPSPEED}
+
+ddp_timeout: ${DDP_TIMEOUT}
+ddp_bucket_cap_mb: ${DDP_BUCKET_CAP_MB}
+ddp_broadcast_buffers: ${DDP_BROADCAST_BUFFERS}
+
+torchdistx_path: ${TORCHDISTX_PATH}
+pretraining_dataset: ${PRETRAINING_DATASET}
+debug: ${DEBUG}
+seed: ${SEED}
+strict: ${STRICT}
diff --git a/.runpod/src/handler.py b/.runpod/src/handler.py
new file mode 100644
index 0000000000..740c1ed1fb
--- /dev/null
+++ b/.runpod/src/handler.py
@@ -0,0 +1,66 @@
+"""
+Runpod serverless entrypoint handler
+"""
+
+import os
+
+import runpod
+import yaml
+from huggingface_hub._login import login
+from train import train
+from utils import get_output_dir
+
+BASE_VOLUME = os.environ.get("BASE_VOLUME", "/runpod-volume")
+if not os.path.exists(BASE_VOLUME):
+    os.makedirs(BASE_VOLUME)
+
+logger = runpod.RunPodLogger()
+
+
+async def handler(job):
+    runpod_job_id = job["id"]
+    inputs = job["input"]
+    run_id = inputs.get("run_id", "default_run_id")
+    args = inputs.get("args", {})
+
+    # Set output directory
+    output_dir = os.path.join(BASE_VOLUME, get_output_dir(run_id))
+    args["output_dir"] = output_dir
+
+    # First save args to a temporary config file
+    config_path = "/workspace/test_config.yaml"
+
+    # Add run_name and job_id to args before saving
+    args["run_name"] = run_id
+    args["runpod_job_id"] = runpod_job_id
+
+    yaml_data = yaml.dump(args, default_flow_style=False)
+    with open(config_path, "w", encoding="utf-8") as file:
+        file.write(yaml_data)
+
+    # Handle credentials
+    credentials = inputs.get("credentials", {})
+
+    if "wandb_api_key" in credentials:
+        os.environ["WANDB_API_KEY"] = credentials["wandb_api_key"]
+    if "hf_token" in credentials:
+        os.environ["HF_TOKEN"] = credentials["hf_token"]
+
+    if os.environ.get("HF_TOKEN"):
+        login(token=os.environ["HF_TOKEN"])
+    else:
+        logger.info("No HF_TOKEN provided. Skipping login.")
+
+    logger.info("Starting Training.")
+    async for result in train(config_path):  # Pass the config path instead of args
+        logger.info(result)
+    logger.info("Training Complete.")
+
+    # Cleanup
+    if "WANDB_API_KEY" in os.environ:
+        del os.environ["WANDB_API_KEY"]
+    if "HF_TOKEN" in os.environ:
+        del os.environ["HF_TOKEN"]
+
+
+runpod.serverless.start({"handler": handler, "return_aggregate_stream": True})
diff --git a/.runpod/src/test_input.json b/.runpod/src/test_input.json
new file mode 100644
index 0000000000..889e8ee273
--- /dev/null
+++ b/.runpod/src/test_input.json
@@ -0,0 +1,61 @@
+{
+  "input": {
+    "user_id": "user",
+    "model_id": "llama-test",
+    "run_id": "llama-test",
+    "credentials": {
+      "wandb_api_key": "",
+      "hf_token": ""
+    },
+    "args": {
+      "base_model": "NousResearch/Meta-Llama-3-8B",
+      "model_type": "LlamaForCausalLM",
+      "tokenizer_type": "AutoTokenizer",
+      "load_in_8bit": true,
+      "load_in_4bit": false,
+      "strict": false,
+      "datasets": [
+        {
+          "path": "mhenrichsen/alpaca_2k_test",
+          "type": "alpaca"
+        }
+      ],
+      "val_set_size": 0.05,
+      "output_dir": "./outputs/lora-out",
+      "sequence_len": 4096,
+      "sample_packing": true,
+      "eval_sample_packing": false,
+      "pad_to_sequence_len": true,
+      "adapter": "lora",
+      "lora_r": 32,
+      "lora_alpha": 16,
+      "lora_dropout": 0.05,
+      "lora_target_linear": true,
+      "lora_modules_to_save": [
+        "embed_tokens",
+        "lm_head"
+      ],
+      "gradient_accumulation_steps": 4,
+      "micro_batch_size": 2,
+      "num_epochs": 1,
+      "optimizer": "adamw_bnb_8bit",
+      "lr_scheduler": "cosine",
+      "learning_rate": 0.0002,
+      "train_on_inputs": false,
+      "group_by_length": false,
+      "bf16": "auto",
+      "tf32": false,
+      "gradient_checkpointing": true,
+      "logging_steps": 1,
+      "flash_attention": true,
+      "warmup_steps": 1,
+      "evals_per_epoch": 1,
+      "eval_max_new_tokens": 128,
+      "saves_per_epoch": 1,
+      "weight_decay": 0.0,
+      "special_tokens": {
+        "pad_token": "<|end_of_text|>"
+      }
+    }
+  }
+}
diff --git a/.runpod/src/train.py b/.runpod/src/train.py
new file mode 100644
index 0000000000..72edda9405
--- /dev/null
+++ b/.runpod/src/train.py
@@ -0,0 +1,45 @@
+"""
+Runpod train entrypoint
+"""
+
+import asyncio
+
+
+async def train(config_path: str, gpu_id: str = "0", preprocess: bool = True):
+    """
+    Run preprocessing (if enabled) and training with the given config file
+    :param config_path: Path to the YAML config file
+    :param gpu_id: GPU ID to use (default: "0")
+    :param preprocess: Whether to run preprocessing (default: True)
+
+    """
+    # First check if preprocessing is needed
+    if preprocess:
+        # Preprocess command
+        preprocess_cmd = (
+            f"CUDA_VISIBLE_DEVICES={gpu_id} axolotl preprocess {config_path}"
+        )
+        process = await asyncio.create_subprocess_shell(
+            preprocess_cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+
+        if process.stdout is not None:
+            async for line in process.stdout:
+                yield f"Preprocessing: {line.decode().strip()}"
+        await process.wait()
+        yield "Preprocessing completed."
+    else:
+        yield "Skipping preprocessing step."
+
+    # Training command
+    train_cmd = f"axolotl train {config_path}"
+    process = await asyncio.create_subprocess_shell(
+        train_cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT
+    )
+
+    if process.stdout is not None:
+        async for line in process.stdout:
+            yield f"Training: {line.decode().strip()}"
+    await process.wait()
diff --git a/.runpod/src/utils.py b/.runpod/src/utils.py
new file mode 100644
index 0000000000..8245aecf43
--- /dev/null
+++ b/.runpod/src/utils.py
@@ -0,0 +1,89 @@
+"""
+Runpod launcher utils
+"""
+
+import os
+
+import yaml
+
+
+def get_output_dir(run_id):
+    path = f"fine-tuning/{run_id}"
+    return path
+
+
+def make_valid_config(input_args):
+    """
+    Creates and saves updated config file, returns the path to the new config
+    :param input_args: dict of input args
+    :return: str, path to the updated config file
+    """
+    # Load default config
+    with open("config/config.yaml", "r", encoding="utf-8") as fin:
+        all_args = yaml.safe_load(fin)
+
+    if not input_args:
+        print("No args provided, using defaults")
+    else:
+        all_args.update(input_args)
+
+    # Create updated config path
+    updated_config_path = "config/updated_config.yaml"
+
+    # Save updated config to new file
+    with open(updated_config_path, "w", encoding="utf-8") as f:
+        yaml.dump(all_args, f)
+
+    return updated_config_path
+
+
+def set_config_env_vars(args: dict):
+    """
+    Convert API arguments into environment variables.
+    Handles nested dictionaries, lists, and special values.
+
+    Args:
+        args (dict): The arguments dictionary from the API request
+    """
+
+    def process_value(value):
+        """Convert Python values to string format for environment variables"""
+        if value is None:
+            return ""
+        if isinstance(value, bool):
+            return str(value).lower()
+        if isinstance(value, (list, dict)):
+            return str(value)
+        return str(value)
+
+    def set_env_vars(data, prefix=""):
+        """Recursively set environment variables from nested dictionary"""
+        for key, value in data.items():
+            env_key = prefix + key.upper()
+
+            # Handle special cases
+            if isinstance(value, dict):
+                # For nested dictionaries (like special_tokens)
+                set_env_vars(value, f"{env_key}_")
+            elif isinstance(value, list):
+                # Handle list of dictionaries (like datasets)
+                if value and isinstance(value[0], dict):
+                    for i, item in enumerate(value):
+                        set_env_vars(item, f"{env_key}_{i}_")
+                else:
+                    # For simple lists (like lora_target_modules)
+                    os.environ[env_key] = process_value(value)
+            else:
+                # Handle all other cases
+                os.environ[env_key] = process_value(value)
+
+    # Clear any existing related environment variables
+    # This prevents old values from persisting
+    for key in list(os.environ.keys()):
+        if key.startswith(
+            ("BASE_MODEL", "MODEL_TYPE", "TOKENIZER_TYPE", "DATASET", "LORA_", "WANDB_")
+        ):
+            del os.environ[key]
+
+    # Set new environment variables
+    set_env_vars(args)
diff --git a/.runpod/test-input.json b/.runpod/test-input.json
new file mode 100644
index 0000000000..52bc905e3e
--- /dev/null
+++ b/.runpod/test-input.json
@@ -0,0 +1,86 @@
+{
+  "input": {
+    "name": "quick_smoke_test_sft",
+    "user_id": "user",
+    "model_id": "llama-test",
+    "run_id": "llama-test",
+    "credentials": {
+      "wandb_api_key": "",
+      "hf_token": ""
+    },
+    "args": {
+      "base_model": "HuggingFaceTB/SmolLM2-135M",
+      "model_type": "AutoModelForCausalLM",
+      "tokenizer_type": "AutoTokenizer",
+      "load_in_4bit": true,
+      "strict": false,
+      "datasets": [
+        {
+          "path": "mhenrichsen/alpaca_2k_test",
+          "type": "alpaca",
+          "split": "train[:10%]"
+        }
+      ],
+      "val_set_size": 0.02,
+      "output_dir": "./outputs/lora-out",
+      "sequence_len": 4096,
+      "sample_packing": true,
+      "eval_sample_packing": false,
+      "pad_to_sequence_len": true,
+      "adapter": "qlora",
+      "lora_r": 32,
+      "lora_alpha": 64,
+      "lora_dropout": 0.05,
+      "lora_target_linear": true,
+      "lora_modules_to_save": [
+        "embed_tokens",
+        "lm_head"
+      ],
+      "gradient_accumulation_steps": 2,
+      "micro_batch_size": 1,
+      "num_epochs": 1,
+      "optimizer": "adamw_torch_fused",
+      "lr_scheduler": "cosine",
+      "learning_rate": 0.0002,
+      "train_on_inputs": false,
+      "group_by_length": false,
+      "bf16": "auto",
+      "tf32": true,
+      "gradient_checkpointing": true,
+      "logging_steps": 1,
+      "flash_attention": true,
+      "warmup_steps": 1,
+      "evals_per_epoch": 1,
+      "eval_max_new_tokens": 128,
+      "saves_per_epoch": 1,
+      "weight_decay": 0.0,
+      "special_tokens": {
+        "pad_token": "<|endoftext|>"
+      },
+      "max_steps": 20
+    },
+    "timeout": 100000
+  },
+  "config": {
+    "gpuTypeId": "NVIDIA GeForce RTX 4090",
+    "gpuCount": 1,
+    "containerDiskInGb": 200,
+    "env": [
+      {
+        "key": "TOKENIZER",
+        "value": ""
+      },
+      {
+        "key": "DISABLE_LOG_STATS",
+        "value": "true"
+      }
+    ],
+    "allowedCudaVersions": [
+      "12.8",
+      "12.7",
+      "12.6",
+      "12.5",
+      "12.4"
+    ]
+  }
+}
diff --git a/.runpod/tests.json b/.runpod/tests.json
new file mode 100644
index 0000000000..6cc18daec9
--- /dev/null
+++ b/.runpod/tests.json
@@ -0,0 +1,90 @@
+{
+  "tests": [
+    {
+      "name": "quick_smoke_test_sft",
+      "input": {
+        "user_id": "user",
+        "model_id": "llama-test",
+        "run_id": "llama-test",
+        "credentials": {
+          "wandb_api_key": "",
+          "hf_token": ""
+        },
+        "args": {
+          "base_model": "HuggingFaceTB/SmolLM2-135M",
+          "model_type": "AutoModelForCausalLM",
+          "tokenizer_type": "AutoTokenizer",
+          "load_in_4bit": true,
+          "strict": false,
+          "datasets": [
+            {
+              "path": "mhenrichsen/alpaca_2k_test",
+              "type": "alpaca",
+              "split": "train[:10%]"
+            }
+          ],
+          "val_set_size": 0.02,
+          "output_dir": "./outputs/lora-out",
+          "sequence_len": 4096,
+          "sample_packing": true,
+          "eval_sample_packing": false,
+          "pad_to_sequence_len": true,
+          "adapter": "qlora",
+          "lora_r": 32,
+          "lora_alpha": 64,
+          "lora_dropout": 0.05,
+          "lora_target_linear": true,
+          "lora_modules_to_save": [
+            "embed_tokens",
+            "lm_head"
+          ],
+          "gradient_accumulation_steps": 2,
+          "micro_batch_size": 1,
+          "num_epochs": 1,
+          "optimizer": "adamw_torch_fused",
+          "lr_scheduler": "cosine",
+          "learning_rate": 0.0002,
+          "train_on_inputs": false,
+          "group_by_length": false,
+          "bf16": "auto",
+          "tf32": true,
+          "gradient_checkpointing": true,
+          "logging_steps": 1,
+          "flash_attention": true,
+          "warmup_steps": 1,
+          "evals_per_epoch": 1,
+          "eval_max_new_tokens": 128,
+          "saves_per_epoch": 1,
+          "weight_decay": 0.0,
+          "special_tokens": {
+            "pad_token": "<|endoftext|>"
+          },
+          "max_steps": 20
+        }
+      },
+      "timeout": 100000
+    }
+  ],
+  "config": {
+    "gpuTypeId": "NVIDIA GeForce RTX 4090",
+    "gpuCount": 1,
+    "containerDiskInGb": 200,
+    "env": [
+      {
+        "key": "TOKENIZER",
+        "value": ""
+      },
+      {
+        "key": "DISABLE_LOG_STATS",
+        "value": "true"
+      }
+    ],
+    "allowedCudaVersions": [
+      "12.8",
+      "12.7",
+      "12.6",
+      "12.5",
+      "12.4"
+    ]
+  }
+}
diff --git a/AGENTS.md b/AGENTS.md
new file mode 100644
index 0000000000..37237b33ed
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1,118 @@
+# Axolotl
+
+Fine-tuning framework for LLMs. Config-driven: every training run is defined by a single YAML file.
+
+## Tech Stack
+
+Python, PyTorch, HuggingFace Transformers, TRL, PEFT (LoRA/QLoRA), DeepSpeed, FSDP, vLLM (for GRPO generation).
+
+## Commands
+
+```bash
+axolotl train config.yaml              # Train (single or multi-GPU, auto-detected)
+axolotl preprocess config.yaml         # Tokenize dataset and validate config
+axolotl preprocess config.yaml --debug # Inspect tokenized samples and label masking
+axolotl inference config.yaml          # Interactive inference
+axolotl merge-lora config.yaml         # Merge LoRA adapter into base model
+axolotl vllm-serve config.yaml         # Start vLLM server for GRPO/EBFT training
+axolotl fetch examples                 # Download example configs
+axolotl agent-docs                     # Show agent-optimized docs (bundled with pip package)
+axolotl agent-docs grpo                # Topic-specific agent reference
+axolotl config-schema                  # Dump config JSON schema
+```
+
+## Training Methods
+
+| Method | Config Key | When to Use |
+|--------|-----------|-------------|
+| SFT | *(default)* | Input-output pairs, instruction tuning |
+| DPO/IPO | `rl: dpo` / `rl: dpo, dpo_loss_type: ["ipo"]` | Paired preference data (chosen vs rejected) |
+| KTO | `rl: kto` | Unpaired binary preference labels |
+| ORPO | `rl: orpo` | Single-stage alignment, no ref model |
+| GRPO | `rl: grpo` | RL with verifiable reward functions (math, code) |
+| EBFT | `rl: ebft` | Feature-matching rewards from internal representations |
+
+Agent-specific references:
+- [docs/agents/sft.md](docs/agents/sft.md) — supervised fine-tuning
+- [docs/agents/preference_tuning.md](docs/agents/preference_tuning.md) — DPO, IPO, KTO, ORPO, SimPO
+- [docs/agents/grpo.md](docs/agents/grpo.md) — GRPO online RL with reward functions
+- [docs/agents/reward_modelling.md](docs/agents/reward_modelling.md) — outcome and process reward models
+- [docs/agents/pretraining.md](docs/agents/pretraining.md) — continual pretraining
+- [docs/agents/model_architectures.md](docs/agents/model_architectures.md) — model-specific quirks (Gemma4, Qwen3.5 MoE, etc.)
+- [docs/agents/new_model_support.md](docs/agents/new_model_support.md) — debugging and adding support for new model architectures
+
+## Config Pattern
+
+All training is config-driven. A YAML file specifies model, adapter, dataset(s), and hyperparameters:
+
+```yaml
+base_model: meta-llama/Llama-3.1-8B-Instruct
+adapter: lora                    # or qlora, or omit for full fine-tune
+datasets:
+  - path: my_dataset
+    type: chat_template          # prompt strategy (see docs/dataset-formats/)
+output_dir: ./outputs/lora-out
+```
+
+Config schema: `src/axolotl/utils/schemas/config.py` (AxolotlInputConfig).
+
+## Project Structure
+
+```
+src/axolotl/
+  cli/                           # CLI entry points (train, preprocess, inference, merge_lora, vllm_serve)
+  core/
+    builders/                    # TrainerBuilder classes (causal.py for SFT, rl.py for RLHF)
+    trainers/                    # Trainer classes, mixins (optimizer, scheduler, packing)
+      dpo/                       # DPO trainer and config
+      grpo/                      # GRPO trainer and sampler
+  loaders/                       # Model, tokenizer, adapter, processor loading
+  prompt_strategies/             # Dataset format handlers (chat_template, alpaca, dpo/, kto/, orpo/)
+  utils/schemas/                 # Pydantic config schemas (config, model, training, peft, trl, fsdp)
+  integrations/                  # Plugins (liger, cut_cross_entropy, swanlab, nemo_gym)
+  monkeypatch/                   # Runtime patches for HF transformers
+
+examples/                        # Example YAML configs by model (llama-3/, qwen2/, mistral/, ebft/)
+deepspeed_configs/               # DeepSpeed JSON configs (zero2, zero3)
+docs/                            # Quarto documentation site
+```
+
+## Linting & Tests
+
+The repo pins CI tool versions in `.pre-commit-config.yaml` — never run system `ruff`/`mypy`.
+
+- `pre-commit run --all-files` — ruff, ruff-format, mypy, bandit at the CI-pinned versions
+- `uvx ruff@<rev> check --fix && uvx ruff@<rev> format` — auto-fix with the pinned ruff (`<rev>` = the `ruff-pre-commit` rev in `.pre-commit-config.yaml`)
+- `pytest -m 'not slow' --ignore=tests/e2e tests/` — CPU suite
+
+Setup, CI matrix, GPU e2e, skip-CI keywords: [.github/CONTRIBUTING.md](.github/CONTRIBUTING.md).
+
+## Code Conventions
+
+- Config-driven: features are toggled via YAML, not code changes
+- Prompt strategies: `src/axolotl/prompt_strategies/` — each `type:` value maps to a function
+- Plugin system: `plugins:` list in config loads integration modules
+- Trainer mixins: `core/trainers/mixins/` for composable trainer behaviors
+- Schemas: all config validation via Pydantic in `utils/schemas/`
+
+## Comment Style
+
+- Default to no comment. Only add one when the WHY is non-obvious (hidden constraint, subtle invariant, workaround for a specific bug).
+- Don't explain WHAT the code does — names and types already do that.
+- Don't reference the current task, PR, or callers (e.g. "added for X", "used by Y", "fixes #123"). Those belong in commit messages / PR descriptions and rot fast.
+- Prefer one short line max.
+- Don't add planning/decision/analysis markdown files unless explicitly requested.
+
+## Key Documentation
+
+- [Getting Started](docs/getting-started.qmd) — quickstart tutorial
+- [Choosing a Method](docs/choosing_method.qmd) — SFT vs DPO vs GRPO decision guide
+- [Support Matrix](docs/support-matrix.qmd) — what Axolotl supports, feature couplings, and known gaps
+- [Config Reference](docs/config-reference.qmd) — all config options
+- [Dataset Formats](docs/dataset-formats/) — chat_template, alpaca, input_output, completion
+- [RLHF](docs/rlhf.qmd) — DPO, KTO, ORPO, GRPO, EBFT configs and dataset formats
+- [GRPO Deep Dive](docs/grpo.qmd) — async training, custom rewards, scaling
+- [vLLM Serving](docs/vllm_serving.qmd) — vLLM setup for GRPO/EBFT
+- [Multi-GPU](docs/multi-gpu.qmd) — FSDP and DeepSpeed
+- [Training Stability](docs/training_stability.qmd) — debugging loss, NaN, OOM
+- [Debugging](docs/debugging.qmd) — VSCode setup, Docker debugging
diff --git a/CITATION.cff b/CITATION.cff
new file mode 100644
index 0000000000..7bbfeec64d
--- /dev/null
+++ b/CITATION.cff
@@ -0,0 +1,10 @@
+cff-version: 1.2.0
+type: software
+title: "Axolotl: Open Source LLM Post-Training"
+message: "If you use this software, please cite it as below."
+authors:
+  - name: "Axolotl maintainers and contributors"
+repository-code: "https://github.com/axolotl-ai-cloud/axolotl"
+url: "https://axolotl.ai/"
+license: Apache-2.0
+date-released: "2023-05-30"
diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 120000
index 0000000000..47dc3e3d86
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1 @@
+AGENTS.md
\ No newline at end of file
diff --git a/CNAME b/CNAME
new file mode 100644
index 0000000000..153ba56c3e
--- /dev/null
+++ b/CNAME
@@ -0,0 +1 @@
+docs.axolotl.ai
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 0000000000..1e51d07d67
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,9 @@
+include README.md
+include LICENSE
+include VERSION
+include src/axolotl/utils/chat_templates/templates/*.jinja
+include src/axolotl/integrations/kernels/libs/scattermoe_lora/metadata.json
+recursive-include src/axolotl *.yaml
+include AGENTS.md
+recursive-include docs/agents *.md
+recursive-include axolotl *.py
diff --git a/README.md b/README.md
index 0b45bb78b4..0d36a35e5f 100644
--- a/README.md
+++ b/README.md
@@ -1,702 +1,238 @@
-# Axolotl
-
-Axolotl is a tool designed to streamline the fine-tuning of various AI models, offering support for multiple configurations and architectures.
-
-Features:
-- Train various Huggingface models such as llama, pythia, falcon, mpt
-- Supports fullfinetune, lora, qlora, relora, and gptq
-- Customize configurations using a simple yaml file or CLI overwrite
-- Load different dataset formats, use custom formats, or bring your own tokenized datasets
-- Integrated with xformer, flash attention, rope scaling, and multipacking
-- Works with single GPU or multiple GPUs via FSDP or Deepspeed
-- Easily run with Docker locally or on the cloud
-- Log results and optionally checkpoints to wandb or mlflow
-- And more!
-
-<a href="https://www.phorm.ai/query?projectId=e315ba4a-4e14-421f-ab05-38a1f9076f25">
-  <img alt="phorm.ai" src="https://img.shields.io/badge/Phorm-Ask_AI-%23F2777A.svg?&logo=data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iNSIgaGVpZ2h0PSI0IiBmaWxsPSJub25lIiB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciPgogIDxwYXRoIGQ9Ik00LjQzIDEuODgyYTEuNDQgMS40NCAwIDAgMS0uMDk4LjQyNmMtLjA1LjEyMy0uMTE1LjIzLS4xOTIuMzIyLS4wNzUuMDktLjE2LjE2NS0uMjU1LjIyNmExLjM1MyAxLjM1MyAwIDAgMS0uNTk1LjIxMmMtLjA5OS4wMTItLjE5Mi4wMTQtLjI3OS4wMDZsLTEuNTkzLS4xNHYtLjQwNmgxLjY1OGMuMDkuMDAxLjE3LS4xNjkuMjQ2LS4xOTFhLjYwMy42MDMgMCAwIDAgLjItLjEwNi41MjkuNTI5IDAgMCAwIC4xMzgtLjE3LjY1NC42NTQgMCAwIDAgLjA2NS0uMjRsLjAyOC0uMzJhLjkzLjkzIDAgMCAwLS4wMzYtLjI0OS41NjcuNTY3IDAgMCAwLS4xMDMtLjIuNTAyLjUwMiAwIDAgMC0uMTY4LS4xMzguNjA4LjYwOCAwIDAgMC0uMjQtLjA2N0wyLjQzNy43MjkgMS42MjUuNjcxYS4zMjIuMzIyIDAgMCAwLS4yMzIuMDU4LjM3NS4zNzUgMCAwIDAtLjExNi4yMzJsLS4xMTYgMS40NS0uMDU4LjY5Ny0uMDU4Ljc1NEwuNzA1IDRsLS4zNTctLjA3OUwuNjAyLjkwNkMuNjE3LjcyNi42NjMuNTc0LjczOS40NTRhLjk1OC45NTggMCAwIDEgLjI3NC0uMjg1Ljk3MS45NzEgMCAwIDEgLjMzNy0uMTRjLjExOS0uMDI2LjIyNy0uMDM0LjMyNS0uMDI2TDMuMjMyLjE2Yy4xNTkuMDE0LjMzNi4wMy40NTkuMDgyYTEuMTczIDEuMTczIDAgMCAxIC41NDUuNDQ3Yy4wNi4wOTQuMTA5LjE5Mi4xNDQuMjkzYTEuMzkyIDEuMzkyIDAgMCAxIC4wNzguNThsLS4wMjkuMzJaIiBmaWxsPSIjRjI3NzdBIi8+CiAgPHBhdGggZD0iTTQuMDgyIDIuMDA3YTEuNDU1IDEuNDU1IDAgMCAxLS4wOTguNDI3Yy0uMDUuMTI0LS4xMTQuMjMyLS4xOTIuMzI0YTEuMTMgMS4xMyAwIDAgMS0uMjU0LjIyNyAxLjM1MyAxLjM1MyAwIDAgMS0uNTk1LjIxNGMtLjEuMDEyLS4xOTMuMDE0LS4yOC4wMDZsLTEuNTYtLjEwOC4wMzQtLjQwNi4wMy0uMzQ4IDEuNTU5LjE1NGMuMDkgMCAuMTczLS4wMS4yNDgtLjAzM2EuNjAzLjYwMyAwIDAgMCAuMi0uMTA2LjUzMi41MzIgMCAwIDAgLjEzOS0uMTcyLjY2LjY2IDAgMCAwIC4wNjQtLjI0MWwuMDI5LS4zMjFhLjk0Ljk0IDAgMCAwLS4wMzYtLjI1LjU3LjU3IDAgMCAwLS4xMDMtLjIwMi41MDIuNTAyIDAgMCAwLS4xNjgtLjEzOC42MDUuNjA1IDAgMCAwLS4yNC0uMDY3TDEuMjczLjgyN2MtLjA5NC0uMDA4LS4xNjguMDEtLjIyMS4wNTUtLjA1My4wNDUtLjA4NC4xMTQtLjA5Mi4yMDZMLjcwNSA0IDAgMy45MzhsLjI1NS0yLjkxMUExLjAxIDEuMDEgMCAwIDEgLjM5My41NzIuOTYyLjk2MiAwIDAgMSAuNjY2LjI4NmEuOTcuOTcgMCAwIDEgLjMzOC0uMTRDMS4xMjIuMTIgMS4yMy4xMSAxLjMyOC4xMTlsMS41OTMuMTRjLjE2LjAxNC4zLjA0Ny40MjMuMWExLjE3IDEuMTcgMCAwIDEgLjU0NS40NDhjLjA2MS4wOTUuMTA5LjE5My4xNDQuMjk1YTEuNDA2IDEuNDA2IDAgMCAxIC4wNzcuNTgzbC0uMDI4LjMyMloiIGZpbGw9IndoaXRlIi8+CiAgPHBhdGggZD0iTTQuMDgyIDIuMDA3YTEuNDU1IDEuNDU1IDAgMCAxLS4wOTguNDI3Yy0uMDUuMTI0LS4xMTQuMjMyLS4xOTIuMzI0YTEuMTMgMS4xMyAwIDAgMS0uMjU0LjIyNyAxLjM1MyAxLjM1MyAwIDAgMS0uNTk1LjIxNGMtLjEuMDEyLS4xOTMuMDE0LS4yOC4wMDZsLTEuNTYtLjEwOC4wMzQtLjQwNi4wMy0uMzQ4IDEuNTU5LjE1NGMuMDkgMCAuMTczLS4wMS4yNDgtLjAzM2EuNjAzLjYwMyAwIDAgMCAuMi0uMTA2LjUzMi41MzIgMCAwIDAgLjEzOS0uMTcyLjY2LjY2IDAgMCAwIC4wNjQtLjI0MWwuMDI5LS4zMjFhLjk0Ljk0IDAgMCAwLS4wMzYtLjI1LjU3LjU3IDAgMCAwLS4xMDMtLjIwMi41MDIuNTAyIDAgMCAwLS4xNjgtLjEzOC42MDUuNjA1IDAgMCAwLS4yNC0uMDY3TDEuMjczLjgyN2MtLjA5NC0uMDA4LS4xNjguMDEtLjIyMS4wNTUtLjA1My4wNDUtLjA4NC4xMTQtLjA5Mi4yMDZMLjcwNSA0IDAgMy45MzhsLjI1NS0yLjkxMUExLjAxIDEuMDEgMCAwIDEgLjM5My41NzIuOTYyLjk2MiAwIDAgMSAuNjY2LjI4NmEuOTcuOTcgMCAwIDEgLjMzOC0uMTRDMS4xMjIuMTIgMS4yMy4xMSAxLjMyOC4xMTlsMS41OTMuMTRjLjE2LjAxNC4zLjA0Ny40MjMuMWExLjE3IDEuMTcgMCAwIDEgLjU0NS40NDhjLjA2MS4wOTUuMTA5LjE5My4xNDQuMjk1YTEuNDA2IDEuNDA2IDAgMCAxIC4wNzcuNTgzbC0uMDI4LjMyMloiIGZpbGw9IndoaXRlIi8+Cjwvc3ZnPgo=">
-</a>
-
-<table>
-<tr>
-<td>
-
-## Table of Contents
-- [Introduction](#axolotl)
-- [Supported Features](#axolotl-supports)
-- [Quickstart](#quickstart-)
-- [Environment](#environment)
-  - [Docker](#docker)
-  - [Conda/Pip venv](#condapip-venv)
-  - [Cloud GPU](#cloud-gpu) - Latitude.sh, JarvisLabs, RunPod
-  - [Bare Metal Cloud GPU](#bare-metal-cloud-gpu)
-  - [Windows](#windows)
-  - [Mac](#mac)
-  - [Google Colab](#google-colab)
-  - [Launching on public clouds via SkyPilot](#launching-on-public-clouds-via-skypilot)
-  - [Launching on public clouds via dstack](#launching-on-public-clouds-via-dstack)
-- [Dataset](#dataset)
-- [Config](#config)
-  - [Train](#train)
-  - [Inference](#inference-playground)
-  - [Merge LORA to Base](#merge-lora-to-base)
-  - [Special Tokens](#special-tokens)
-  - [All Config Options](#all-config-options)
-- Advanced Topics
-  - [Multipack](./docs/multipack.qmd)<svg width="24" height="24" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M17 13.5v6H5v-12h6m3-3h6v6m0-6-9 9" class="icon_svg-stroke" stroke="#666" stroke-width="1.5" fill="none" fill-rule="evenodd" stroke-linecap="round" stroke-linejoin="round"></path></svg>
-  - [RLHF & DPO](./docs/rlhf.qmd)<svg width="24" height="24" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M17 13.5v6H5v-12h6m3-3h6v6m0-6-9 9" class="icon_svg-stroke" stroke="#666" stroke-width="1.5" fill="none" fill-rule="evenodd" stroke-linecap="round" stroke-linejoin="round"></path></svg>
-  - [Dataset Pre-Processing](./docs/dataset_preprocessing.qmd)<svg width="24" height="24" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M17 13.5v6H5v-12h6m3-3h6v6m0-6-9 9" class="icon_svg-stroke" stroke="#666" stroke-width="1.5" fill="none" fill-rule="evenodd" stroke-linecap="round" stroke-linejoin="round"></path></svg>
-- [Common Errors](#common-errors-)
-  - [Tokenization Mismatch b/w Training & Inference](#tokenization-mismatch-bw-inference--training)
-- [Debugging Axolotl](#debugging-axolotl)
-- [Need Help?](#need-help-)
-- [Badge](#badge-)
-- [Community Showcase](#community-showcase)
-- [Contributing](#contributing-)
-- [Sponsors](#sponsors-)
-
-</td>
-<td>
-
-<div align="center">
-  <img src="image/axolotl.png" alt="axolotl" width="160">
-  <div>
-    <p>
-      <b>Axolotl provides a unified repository for fine-tuning <br />a variety of AI models with ease</b>
-    </p>
-    <p>
-      Go ahead and Axolotl questions!!
-    </p>
-    <img src="https://github.com/OpenAccess-AI-Collective/axolotl/actions/workflows/pre-commit.yml/badge.svg?branch=main" alt="pre-commit">
-    <img alt="PyTest Status" src="https://github.com/OpenAccess-AI-Collective/axolotl/actions/workflows/tests.yml/badge.svg?branch=main">
-  </div>
-</div>
-
-</td>
-</tr>
-</table>
-
-## Axolotl supports
-
-|             | fp16/fp32 | lora | qlora | gptq | gptq w/flash attn | flash attn | xformers attn |
-|-------------|:----------|:-----|-------|------|-------------------|------------|--------------|
-| llama       | ✅         | ✅    | ✅     | ✅             | ✅                 | ✅          | ✅            |
-| Mistral     | ✅         | ✅    | ✅     | ✅             | ✅                 | ✅          | ✅            |
-| Mixtral-MoE | ✅         | ✅    | ✅     | ❓             | ❓                 | ❓          | ❓            |
-| Mixtral8X22 | ✅         | ✅    | ✅     | ❓             | ❓                 | ❓          | ❓            |
-| Pythia      | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
-| cerebras    | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
-| btlm        | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
-| mpt         | ✅         | ❌    | ❓     | ❌             | ❌                 | ❌          | ❓            |
-| falcon      | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓            |
-| gpt-j       | ✅         | ✅    | ✅     | ❌             | ❌                 | ❓          | ❓            |
-| XGen        | ✅         | ❓    | ✅     | ❓             | ❓                 | ❓          | ✅            |
-| phi         | ✅         | ✅    | ✅     | ❓             | ❓                 | ❓          | ❓            |
-| RWKV        | ✅         | ❓    | ❓     | ❓             | ❓                 | ❓          | ❓            |
-| Qwen        | ✅         | ✅    | ✅     | ❓             | ❓                 | ❓          | ❓            |
-| Gemma       | ✅         | ✅    | ✅     | ❓             | ❓                 | ✅          | ❓            |
-
-✅: supported
-❌: not supported
-❓: untested
-
-## Quickstart ⚡
-
-Get started with Axolotl in just a few steps! This quickstart guide will walk you through setting up and running a basic fine-tuning task.
-
-**Requirements**: Python >=3.10 and Pytorch >=2.1.1.
-
-```bash
-git clone https://github.com/OpenAccess-AI-Collective/axolotl
-cd axolotl
-
-pip3 install packaging ninja
-pip3 install -e '.[flash-attn,deepspeed]'
-```
-
-### Usage
-```bash
-# preprocess datasets - optional but recommended
-CUDA_VISIBLE_DEVICES="" python -m axolotl.cli.preprocess examples/openllama-3b/lora.yml
-
-# finetune lora
-accelerate launch -m axolotl.cli.train examples/openllama-3b/lora.yml
-
-# inference
-accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
-    --lora_model_dir="./lora-out"
-
-# gradio
-accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
-    --lora_model_dir="./lora-out" --gradio
-
-# remote yaml files - the yaml config can be hosted on a public URL
-# Note: the yaml config must directly link to the **raw** yaml
-accelerate launch -m axolotl.cli.train https://raw.githubusercontent.com/OpenAccess-AI-Collective/axolotl/main/examples/openllama-3b/lora.yml
-```
-
-## Advanced Setup
-
-### Environment
-
-#### Docker
-
-  ```bash
-  docker run --gpus '"all"' --rm -it winglian/axolotl:main-latest
-  ```
-
-  Or run on the current files for development:
-
-  ```sh
-  docker compose up -d
-  ```
-
->[!Tip]
-> If you want to debug axolotl or prefer to use Docker as your development environment, see the [debugging guide's section on Docker](docs/debugging.qmd#debugging-with-docker).
-
-  <details>
-
-  <summary>Docker advanced</summary>
-
-  A more powerful Docker command to run would be this:
-
-  ```bash
-docker run --privileged --gpus '"all"' --shm-size 10g --rm -it --name axolotl --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --mount type=bind,src="${PWD}",target=/workspace/axolotl -v ${HOME}/.cache/huggingface:/root/.cache/huggingface winglian/axolotl:main-latest
-  ```
-
-  It additionally:
-  * Prevents memory issues when running e.g. deepspeed (e.g. you could hit SIGBUS/signal 7 error) through `--ipc` and `--ulimit` args.
-  * Persists the downloaded HF data (models etc.) and your modifications to axolotl code through `--mount`/`-v` args.
-  * The `--name` argument simply makes it easier to refer to the container in vscode (`Dev Containers: Attach to Running Container...`) or in your terminal.
-  * The `--privileged` flag gives all capabilities to the container.
-  * The `--shm-size 10g` argument increases the shared memory size. Use this if you see `exitcode: -7` errors using deepspeed.
-
-  [More information on nvidia website](https://docs.nvidia.com/deeplearning/frameworks/user-guide/index.html#setincshmem)
-
-  </details>
-
-#### Conda/Pip venv
-  1. Install python >=**3.10**
-
-  2. Install pytorch stable https://pytorch.org/get-started/locally/
-
-  3. Install Axolotl along with python dependencies
-        ```bash
-        pip3 install packaging
-        pip3 install -e '.[flash-attn,deepspeed]'
-        ```
-  4. (Optional) Login to Huggingface to use gated models/datasets.
-        ```bash
-        huggingface-cli login
-        ```
-        Get the token at huggingface.co/settings/tokens
-
-#### Cloud GPU
-
-For cloud GPU providers that support docker images, use [`winglian/axolotl-cloud:main-latest`](https://hub.docker.com/r/winglian/axolotl-cloud/tags)
-
-- on Latitude.sh use this [direct link](https://latitude.sh/blueprint/989e0e79-3bf6-41ea-a46b-1f246e309d5c)
-- on JarvisLabs.ai use this [direct link](https://jarvislabs.ai/templates/axolotl)
-- on RunPod use this [direct link](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
-
-#### Bare Metal Cloud GPU
-
-##### LambdaLabs
-
-  <details>
-
-  <summary>Click to Expand</summary>
-
-  1. Install python
-  ```bash
-  sudo apt update
-  sudo apt install -y python3.10
-
-  sudo update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1
-  sudo update-alternatives --config python # pick 3.10 if given option
-  python -V # should be 3.10
-
-  ```
-
-  2. Install pip
-  ```bash
-  wget https://bootstrap.pypa.io/get-pip.py
-  python get-pip.py
-  ```
-
-  3. Install Pytorch https://pytorch.org/get-started/locally/
-
-  4. Follow instructions on quickstart.
-
-  5. Run
-  ```bash
-  pip3 install protobuf==3.20.3
-  pip3 install -U --ignore-installed requests Pillow psutil scipy
-  ```
-
-  6. Set path
-  ```bash
-  export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH
-  ```
-  </details>
-
-##### GCP
+<p align="center">
+    <picture>
+        <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/887513285d98132142bf5db2a74eb5e0928787f1/image/axolotl_logo_digital_white.svg">
+        <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/887513285d98132142bf5db2a74eb5e0928787f1/image/axolotl_logo_digital_black.svg">
+        <img alt="Axolotl" src="https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/887513285d98132142bf5db2a74eb5e0928787f1/image/axolotl_logo_digital_black.svg" width="400" height="104" style="max-width: 100%;">
+    </picture>
+</p>
+  <p align="center">
+      <strong>A Free and Open Source LLM Fine-tuning Framework</strong><br>
+  </p>
+
+<p align="center">
+    <img src="https://img.shields.io/github/license/axolotl-ai-cloud/axolotl.svg?color=blue" alt="GitHub License">
+    <img src="https://github.com/axolotl-ai-cloud/axolotl/actions/workflows/tests.yml/badge.svg" alt="tests">
+    <a href="https://codecov.io/gh/axolotl-ai-cloud/axolotl"><img src="https://codecov.io/gh/axolotl-ai-cloud/axolotl/branch/main/graph/badge.svg" alt="codecov"></a>
+    <a href="https://github.com/axolotl-ai-cloud/axolotl/releases"><img src="https://img.shields.io/github/release/axolotl-ai-cloud/axolotl.svg" alt="Releases"></a>
+    <br/>
+    <a href="https://github.com/axolotl-ai-cloud/axolotl/graphs/contributors"><img src="https://img.shields.io/github/contributors-anon/axolotl-ai-cloud/axolotl?color=yellow&style=flat-square" alt="contributors" style="height: 20px;"></a>
+    <img src="https://img.shields.io/github/stars/axolotl-ai-cloud/axolotl" alt="GitHub Repo stars">
+    <br/>
+    <a href="https://discord.com/invite/HhrNrHJPRb"><img src="https://img.shields.io/badge/discord-7289da.svg?style=flat-square&logo=discord" alt="discord" style="height: 20px;"></a>
+    <a href="https://twitter.com/axolotl_ai"><img src="https://img.shields.io/twitter/follow/axolotl_ai?style=social" alt="twitter" style="height: 20px;"></a>
+    <a href="https://colab.research.google.com/github/axolotl-ai-cloud/axolotl/blob/main/examples/colab-notebooks/colab-axolotl-example.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google-colab" style="height: 20px;"></a>
+    <br/>
+    <img src="https://github.com/axolotl-ai-cloud/axolotl/actions/workflows/tests-nightly.yml/badge.svg" alt="tests-nightly">
+    <img src="https://github.com/axolotl-ai-cloud/axolotl/actions/workflows/multi-gpu-e2e.yml/badge.svg" alt="multigpu-semi-weekly tests">
+</p>
+
+
+## 🎉 Latest Updates
+
+- 2026/07:
+  - NVFP4 (4-bit) MoE LoRA training is now supported via [ScatterMoE](https://docs.axolotl.ai/docs/custom_integrations.html#scattermoe-nvfp4-w4a16-lora) and [SonicMoE](https://docs.axolotl.ai/docs/custom_integrations.html#sonicmoe-nvfp4-w4a4-lora).
+- 2026/06:
+  - [Expert Parallelism (EP)](https://docs.axolotl.ai/docs/nd_parallelism.html) for distributed MoE training via DeepEP, remote training through [Tinker-compatible APIs](https://github.com/axolotl-ai-cloud/axolotl/pull/3614), [Context Parallelism for hybrid SSM models](https://github.com/axolotl-ai-cloud/axolotl/pull/3572) (Nemotron-H, Falcon-H1, Bamba), [BitNet 1.58-bit](https://github.com/axolotl-ai-cloud/axolotl/pull/3634) fine-tuning, and a [multimodal assistant-only loss-masking fix](https://github.com/axolotl-ai-cloud/axolotl/pull/3625).
+- 2026/04:
+  - New model support has been added in Axolotl for [Mistral Medium 3.5](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/mistral-medium-3_5) and [Gemma 4](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/gemma4).
+  - New RL and kernels: [Async GRPO](https://github.com/axolotl-ai-cloud/axolotl/pull/3486) (up to 58% faster steps), [Flash Attention 4](https://docs.axolotl.ai/docs/attention.html#flash-attention), [NeMo Gym](https://github.com/axolotl-ai-cloud/axolotl/pull/3516), and [EBFT](https://github.com/axolotl-ai-cloud/axolotl/pull/3527).
+  - Axolotl is now [uv-first](https://github.com/axolotl-ai-cloud/axolotl/pull/3545) and has [SonicMoE fused LoRA](https://github.com/axolotl-ai-cloud/axolotl/pull/3519) support.
+- 2026/03:
+  - New model support has been added in Axolotl for [Mistral Small 4](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/mistral4), [Qwen3.5, Qwen3.5 MoE](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/qwen3.5), [GLM-4.7-Flash](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/glm47-flash), [GLM-4.6V](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/glm46v), and [GLM-4.5-Air](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/glm45).
+  - [MoE expert quantization](https://docs.axolotl.ai/docs/expert_quantization.html) support (via `quantize_moe_experts: true`) greatly reduces VRAM when training MoE models (FSDP2 compat).
 
 <details>
 
-<summary>Click to Expand</summary>
-
-Use a Deeplearning linux OS with cuda and pytorch installed. Then follow instructions on quickstart.
-
-Make sure to run the below to uninstall xla.
-```bash
-pip uninstall -y torch_xla[tpu]
-```
+<summary>Expand older updates</summary>
+
+- 2026/02:
+  - [ScatterMoE LoRA](https://github.com/axolotl-ai-cloud/axolotl/pull/3410) support. LoRA fine-tuning directly on MoE expert weights using custom Triton kernels.
+  - Axolotl now has support for [SageAttention](https://github.com/axolotl-ai-cloud/axolotl/pull/2823) and [GDPO](https://github.com/axolotl-ai-cloud/axolotl/pull/3353) (Generalized DPO).
+- 2026/01:
+  - New integration for [EAFT](https://github.com/axolotl-ai-cloud/axolotl/pull/3366) (Entropy-Aware Focal Training), weights loss by entropy of the top-k logit distribution, and [Scalable Softmax](https://github.com/axolotl-ai-cloud/axolotl/pull/3338), improves long context in attention.
+- 2025/12:
+  - Axolotl now includes support for [Kimi-Linear](https://docs.axolotl.ai/docs/models/kimi-linear.html), [Plano-Orchestrator](https://docs.axolotl.ai/docs/models/plano.html), [MiMo](https://docs.axolotl.ai/docs/models/mimo.html), [InternVL 3.5](https://docs.axolotl.ai/docs/models/internvl3_5.html), [Olmo3](https://docs.axolotl.ai/docs/models/olmo3.html), [Trinity](https://docs.axolotl.ai/docs/models/trinity.html), and [Ministral3](https://docs.axolotl.ai/docs/models/ministral3.html).
+  - [Distributed Muon Optimizer](https://github.com/axolotl-ai-cloud/axolotl/pull/3264) support has been added for FSDP2 pretraining.
+- 2025/10: New model support has been added in Axolotl for: [Qwen3 Next](https://docs.axolotl.ai/docs/models/qwen3-next.html), [Qwen2.5-vl, Qwen3-vl](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/qwen2_5-vl), [Qwen3, Qwen3MoE](https://docs.axolotl.ai/docs/models/qwen3.html), [Granite 4](https://docs.axolotl.ai/docs/models/granite4.html), [HunYuan](https://docs.axolotl.ai/docs/models/hunyuan.html), [Magistral 2509](https://docs.axolotl.ai/docs/models/magistral/vision.html), [Apertus](https://docs.axolotl.ai/docs/models/apertus.html), and [Seed-OSS](https://docs.axolotl.ai/docs/models/seed-oss.html).
+- 2025/09: Axolotl now has text diffusion training. Read more [here](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/diffusion).
+- 2025/08: QAT has been updated to include NVFP4 support. See [PR](https://github.com/axolotl-ai-cloud/axolotl/pull/3107).
+- 2025/07:
+  - ND Parallelism support has been added into Axolotl. Compose Context Parallelism (CP), Tensor Parallelism (TP), and Fully Sharded Data Parallelism (FSDP) within a single node and across multiple nodes. Check out the [blog post](https://huggingface.co/blog/accelerate-nd-parallel) for more info.
+  - Axolotl adds more models: [GPT-OSS](https://docs.axolotl.ai/docs/models/gpt-oss.html), [Gemma 3n](https://docs.axolotl.ai/docs/models/gemma3n.html), [Liquid Foundation Model 2 (LFM2)](https://docs.axolotl.ai/docs/models/LiquidAI.html), and [Arcee Foundation Models (AFM)](https://docs.axolotl.ai/docs/models/arcee.html).
+  - FP8 finetuning with fp8 gather op is now possible in Axolotl via `torchao`. Get started [here](https://docs.axolotl.ai/docs/mixed_precision.html#sec-fp8)!
+  - [Voxtral](https://docs.axolotl.ai/docs/models/voxtral.html), [Magistral 1.1](https://docs.axolotl.ai/docs/models/magistral.html), and [Devstral](https://docs.axolotl.ai/docs/models/devstral.html) with mistral-common tokenizer support has been integrated in Axolotl!
+  - TiledMLP support for single-GPU to multi-GPU training with DDP, DeepSpeed and FSDP support has been added to support Arctic Long Sequence Training. (ALST). See [examples](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/alst) for using ALST with Axolotl!
+- 2025/06: Magistral with mistral-common tokenizer support has been added to Axolotl. See [docs](https://docs.axolotl.ai/docs/models/magistral.html) to start training your own Magistral models with Axolotl!
+- 2025/05: Quantization Aware Training (QAT) support has been added to Axolotl. Explore the [docs](https://docs.axolotl.ai/docs/qat.html) to learn more!
+- 2025/04: Llama 4 support has been added in Axolotl. See [docs](https://docs.axolotl.ai/docs/models/llama-4.html) to start training your own Llama 4 models with Axolotl's linearized version!
+- 2025/03: Axolotl has implemented Sequence Parallelism (SP) support. Read the [blog](https://huggingface.co/blog/axolotl-ai-co/long-context-with-sequence-parallelism-in-axolotl) and [docs](https://docs.axolotl.ai/docs/sequence_parallelism.html) to learn how to scale your context length when fine-tuning.
+- 2025/03: (Beta) Fine-tuning Multimodal models is now supported in Axolotl. Check out the [docs](https://docs.axolotl.ai/docs/multimodal.html) to fine-tune your own!
+- 2025/02: Axolotl has added LoRA optimizations to reduce memory usage and improve training speed for LoRA and QLoRA in single GPU and multi-GPU training (DDP and DeepSpeed). Jump into the [docs](https://docs.axolotl.ai/docs/lora_optims.html) to give it a try.
+- 2025/02: Axolotl has added GRPO support. Dive into our [blog](https://huggingface.co/blog/axolotl-ai-co/training-llms-w-interpreter-feedback-wasm) and [GRPO example](https://github.com/axolotl-ai-cloud/grpo_code) and have some fun!
+- 2025/01: Axolotl has added Reward Modelling / Process Reward Modelling fine-tuning support. See [docs](https://docs.axolotl.ai/docs/reward_modelling.html).
 
 </details>
 
-#### Windows
-Please use WSL or Docker!
-
-#### Mac
-
-Use the below instead of the install method in QuickStart.
-```
-pip3 install -e '.'
-```
-More info: [mac.md](/docs/mac.qmd)
-
-#### Google Colab
-
-Please use this example [notebook](examples/colab-notebooks/colab-axolotl-example.ipynb).
-
-#### Launching on public clouds via SkyPilot
-To launch on GPU instances (both on-demand and spot instances) on 7+ clouds (GCP, AWS, Azure, OCI, and more), you can use [SkyPilot](https://skypilot.readthedocs.io/en/latest/index.html):
-
-```bash
-pip install "skypilot-nightly[gcp,aws,azure,oci,lambda,kubernetes,ibm,scp]"  # choose your clouds
-sky check
-```
-
-Get the [example YAMLs](https://github.com/skypilot-org/skypilot/tree/master/llm/axolotl) of using Axolotl to finetune `mistralai/Mistral-7B-v0.1`:
-```
-git clone https://github.com/skypilot-org/skypilot.git
-cd skypilot/llm/axolotl
-```
-
-Use one command to launch:
-```bash
-# On-demand
-HF_TOKEN=xx sky launch axolotl.yaml --env HF_TOKEN
-
-# Managed spot (auto-recovery on preemption)
-HF_TOKEN=xx BUCKET=<unique-name> sky spot launch axolotl-spot.yaml --env HF_TOKEN --env BUCKET
-```
-
-#### Launching on public clouds via dstack
-To launch on GPU instance (both on-demand and spot instances) on public clouds (GCP, AWS, Azure, Lambda Labs, TensorDock, Vast.ai, and CUDO), you can use [dstack](https://dstack.ai/).
+## ✨ Overview
 
-Write a job description in YAML as below:
+Axolotl is a free and open-source tool designed to streamline post-training and fine-tuning for the latest large language models (LLMs).
 
-```yaml
-# dstack.yaml
-type: task
+Features:
 
-image: winglian/axolotl-cloud:main-20240429-py3.11-cu121-2.2.1
+- **Multiple Model Support**: Train various models like GPT-OSS, LLaMA, Mistral, Mixtral, Pythia, and many more models available on the Hugging Face Hub.
+- **Multimodal Training**: Fine-tune vision-language models (VLMs) including LLaMA-Vision, Qwen2-VL, Pixtral, LLaVA, SmolVLM2, GLM-4.6V, InternVL 3.5, Gemma 3n, and audio models like Voxtral with image, video, and audio support.
+- **Training Methods**: Full fine-tuning, LoRA, QLoRA, GPTQ, QAT (int8/int4/FP8/NVFP4/MXFP4), FP8 mixed-precision training, NVFP4/MXFP4 MoE LoRA, Preference Tuning (DPO, IPO, KTO, ORPO), RL (GRPO, GDPO), and Reward Modelling (RM) / Process Reward Modelling (PRM).
+- **Easy Configuration**: Re-use a single YAML configuration file across the full fine-tuning pipeline: dataset preprocessing, training, evaluation, quantization, and inference.
+- **Performance Optimizations**: [Multipacking](https://docs.axolotl.ai/docs/multipack.html), [Flash Attention 2/3/4](https://docs.axolotl.ai/docs/attention.html#flash-attention), [Xformers](https://docs.axolotl.ai/docs/attention.html#xformers), [Flex Attention](https://docs.axolotl.ai/docs/attention.html#flex-attention), [SageAttention](https://docs.axolotl.ai/docs/attention.html#sageattention), [Liger Kernel](https://docs.axolotl.ai/docs/custom_integrations.html#liger-kernels), [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy), [ScatterMoE](https://docs.axolotl.ai/docs/custom_integrations.html#kernels-integration), [Sequence Parallelism (SP)](https://docs.axolotl.ai/docs/sequence_parallelism.html), [LoRA optimizations](https://docs.axolotl.ai/docs/lora_optims.html), [Multi-GPU training (FSDP1, FSDP2, DeepSpeed)](https://docs.axolotl.ai/docs/multi-gpu.html), [Multi-node training (Torchrun, Ray)](https://docs.axolotl.ai/docs/multi-node.html), and many more!
+- **Flexible Dataset Handling**: Load from local, HuggingFace, and cloud (S3, Azure, GCP, OCI) datasets.
+- **Cloud Ready**: We ship [Docker images](https://hub.docker.com/u/axolotlai) and also [PyPI packages](https://pypi.org/project/axolotl/) for use on cloud platforms and local hardware.
 
-env:
-  - HUGGING_FACE_HUB_TOKEN
-  - WANDB_API_KEY
 
-commands:
-  - accelerate launch -m axolotl.cli.train config.yaml
 
-ports:
-  - 6006
+## 🚀 Quick Start - LLM Fine-tuning in Minutes
 
-resources:
-  gpu:
-    memory: 24GB..
-    count: 2
-```
+**Requirements**:
 
-then, simply run the job with `dstack run` command. Append `--spot` option if you want spot instance. `dstack run` command will show you the instance with cheapest price across multi cloud services:
-
-```bash
-pip install dstack
-HUGGING_FACE_HUB_TOKEN=xxx WANDB_API_KEY=xxx dstack run . -f dstack.yaml # --spot
-```
-
-For further and fine-grained use cases, please refer to the official [dstack documents](https://dstack.ai/docs/) and the detailed description of [axolotl example](https://github.com/dstackai/dstack/tree/master/examples/fine-tuning/axolotl) on the official repository.
-
-### Dataset
-
-Axolotl supports a variety of dataset formats.  It is recommended to use a JSONL.  The schema of the JSONL depends upon the task and the prompt template you wish to use.  Instead of a JSONL, you can also use a HuggingFace dataset with columns for each JSONL field.
-
-See [these docs](https://openaccess-ai-collective.github.io/axolotl/docs/dataset-formats/) for more information on how to use different dataset formats.
-
-### Config
-
-See [examples](examples) for quick start. It is recommended to duplicate and modify to your needs. The most important options are:
-
-- model
-  ```yaml
-  base_model: ./llama-7b-hf # local or huggingface repo
-  ```
-  Note: The code will load the right architecture.
-
-- dataset
-  ```yaml
-  datasets:
-      # huggingface repo
-    - path: vicgalle/alpaca-gpt4
-      type: alpaca
-
-      # huggingface repo with specific configuration/subset
-    - path: EleutherAI/pile
-      name: enron_emails
-      type: completion # format from earlier
-      field: text # Optional[str] default: text, field to use for completion data
-
-      # huggingface repo with multiple named configurations/subsets
-    - path: bigcode/commitpackft
-      name:
-        - ruby
-        - python
-        - typescript
-      type: ... # unimplemented custom format
-
-      # fastchat conversation
-      # See 'conversation' options: https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
-    - path: ...
-      type: sharegpt
-      conversation: chatml # default: vicuna_v1.1
-
-      # local
-    - path: data.jsonl # or json
-      ds_type: json # see other options below
-      type: alpaca
-
-      # dataset with splits, but no train split
-    - path: knowrohit07/know_sql
-      type: context_qa.load_v2
-      train_on_split: validation
-
-      # loading from s3 or gcs
-      # s3 creds will be loaded from the system default and gcs only supports public access
-    - path: s3://path_to_ds # Accepts folder with arrow/parquet or file path like above. Supports s3, gcs.
-      ...
-
-      # Loading Data From a Public URL
-      # - The file format is `json` (which includes `jsonl`) by default. For different formats, adjust the `ds_type` option accordingly.
-    - path: https://some.url.com/yourdata.jsonl # The URL should be a direct link to the file you wish to load. URLs must use HTTPS protocol, not HTTP.
-      ds_type: json # this is the default, see other options below.
-  ```
-
-- loading
-  ```yaml
-  load_in_4bit: true
-  load_in_8bit: true
-
-  bf16: auto # require >=ampere, auto will detect if your GPU supports this and choose automatically.
-  fp16: # leave empty to use fp16 when bf16 is 'auto'. set to false if you want to fallback to fp32
-  tf32: true # require >=ampere
-
-  bfloat16: true # require >=ampere, use instead of bf16 when you don't want AMP (automatic mixed precision)
-  float16: true # use instead of fp16 when you don't want AMP
-  ```
-  Note: Repo does not do 4-bit quantization.
-
-- lora
-  ```yaml
-  adapter: lora # 'qlora' or leave blank for full finetune
-  lora_r: 8
-  lora_alpha: 16
-  lora_dropout: 0.05
-  lora_target_modules:
-    - q_proj
-    - v_proj
-  ```
-
-#### All Config Options
-
-See [these docs](docs/config.qmd) for all config options.
-
-### Train
-
-Run
-```bash
-accelerate launch -m axolotl.cli.train your_config.yml
-```
+- NVIDIA GPU (Ampere or newer for `bf16` and Flash Attention) or AMD GPU
+- Python >=3.11 (3.12 recommended)
+- PyTorch ≥2.11.0
 
-> [!TIP]
-> You can also reference a config file that is hosted on a public URL, for example `accelerate launch -m axolotl.cli.train https://yourdomain.com/your_config.yml`
+### Google Colab
 
-#### Preprocess dataset
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/axolotl-ai-cloud/axolotl/blob/main/examples/colab-notebooks/colab-axolotl-example.ipynb#scrollTo=msOCO4NRmRLa)
 
-You can optionally pre-tokenize dataset with the following before finetuning.
-This is recommended for large datasets.
-
-- Set `dataset_prepared_path:` to a local folder for saving and loading pre-tokenized dataset.
-- (Optional): Set `push_dataset_to_hub: hf_user/repo` to push it to Huggingface.
-- (Optional): Use `--debug` to see preprocessed examples.
+### Installation
 
 ```bash
-python -m axolotl.cli.preprocess your_config.yml
-```
-
-#### Multi-GPU
-
-Below are the options available in axolotl for training with multiple GPUs. Note that DeepSpeed
-is the recommended multi-GPU option currently because FSDP may experience
-[loss instability](https://github.com/huggingface/transformers/issues/26498).
-
-##### DeepSpeed
-
-Deepspeed is an optimization suite for multi-gpu systems allowing you to train much larger models than you
-might typically be able to fit into your GPU's VRAM. More information about the various optimization types
-for deepspeed is available at https://huggingface.co/docs/accelerate/main/en/usage_guides/deepspeed#what-is-integrated
-
-We provide several default deepspeed JSON configurations for ZeRO stage 1, 2, and 3.
-
-```yaml
-deepspeed: deepspeed_configs/zero1.json
-```
-
-```shell
-accelerate launch -m axolotl.cli.train examples/llama-2/config.yml --deepspeed deepspeed_configs/zero1.json
-```
-
-##### FSDP
-
-- llama FSDP
-```yaml
-fsdp:
-  - full_shard
-  - auto_wrap
-fsdp_config:
-  fsdp_offload_params: true
-  fsdp_state_dict_type: FULL_STATE_DICT
-  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
-```
-
-##### FSDP + QLoRA
-
-Axolotl supports training with FSDP and QLoRA, see [these docs](docs/fsdp_qlora.qmd) for more information.
-
-##### Weights & Biases Logging
-
-Make sure your `WANDB_API_KEY` environment variable is set (recommended) or you login to wandb with `wandb login`.
-
-- wandb options
-```yaml
-wandb_mode:
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-```
-
-##### Special Tokens
-
-It is important to have special tokens like delimiters, end-of-sequence, beginning-of-sequence in your tokenizer's vocabulary.  This will help you avoid tokenization issues and help your model train better.  You can do this in axolotl like this:
-
-```yml
-special_tokens:
-  bos_token: "<s>"
-  eos_token: "</s>"
-  unk_token: "<unk>"
-tokens: # these are delimiters
-  - "<|im_start|>"
-  - "<|im_end|>"
-```
-
-When you include these tokens in your axolotl config, axolotl adds these tokens to the tokenizer's vocabulary.
-
-### Inference Playground
+# install uv if you don't already have it installed (restart shell after)
+curl -LsSf https://astral.sh/uv/install.sh | sh
 
-Axolotl allows you to load your model in an interactive terminal playground for quick experimentation.
-The config file is the same config file used for training.
+# change depending on system
+export UV_TORCH_BACKEND=cu130
 
-Pass the appropriate flag to the inference command, depending upon what kind of model was trained:
+# create a new virtual environment
+uv venv --python 3.12
+source .venv/bin/activate
 
-- Pretrained LORA:
-  ```bash
-  python -m axolotl.cli.inference examples/your_config.yml --lora_model_dir="./lora-output-dir"
-  ```
-- Full weights finetune:
-  ```bash
-  python -m axolotl.cli.inference examples/your_config.yml --base_model="./completed-model"
-  ```
-- Full weights finetune w/ a prompt from a text file:
-  ```bash
-  cat /tmp/prompt.txt | python -m axolotl.cli.inference examples/your_config.yml \
-    --base_model="./completed-model" --prompter=None --load_in_8bit=True
-  ```
--- With gradio hosting
-  ```bash
-  python -m axolotl.cli.inference examples/your_config.yml --gradio
-  ```
+uv pip install torch==2.12.0 torchvision
+uv pip install --no-build-isolation axolotl[deepspeed]
 
-Please use `--sample_packing False` if you have it on and receive the error similar to below:
-
-> RuntimeError: stack expects each tensor to be equal size, but got [1, 32, 1, 128] at entry 0 and [1, 32, 8, 128] at entry 1
-
-### Merge LORA to base
-
-The following command will merge your LORA adapater with your base model. You can optionally pass the argument `--lora_model_dir` to specify the directory where your LORA adapter was saved, otherwhise, this will be inferred from `output_dir` in your axolotl config file.  The merged model is saved in the sub-directory `{lora_model_dir}/merged`.
-
-```bash
-python3 -m axolotl.cli.merge_lora your_config.yml --lora_model_dir="./completed-model"
+# Download example axolotl configs, deepspeed configs
+axolotl fetch examples
+axolotl fetch deepspeed_configs  # OPTIONAL
 ```
 
-You may need to use the `gpu_memory_limit` and/or `lora_on_cpu` config options to avoid running out of memory. If you still run out of CUDA memory, you can try to merge in system RAM with
+#### Using Docker
 
+Installing with Docker can be less error prone than installing in your own environment.
 ```bash
-CUDA_VISIBLE_DEVICES="" python3 -m axolotl.cli.merge_lora ...
+docker run --gpus '"all"' --ipc=host --rm -it axolotlai/axolotl:main-latest
 ```
 
-although this will be very slow, and using the config options above are recommended instead.
-
-## Common Errors 🧰
-
-See also the [FAQ's](./docs/faq.qmd) and [debugging guide](docs/debugging.qmd).
-
-> If you encounter a 'Cuda out of memory' error, it means your GPU ran out of memory during the training process. Here's how to resolve it:
-
-Please reduce any below
-  - `micro_batch_size`
-  - `eval_batch_size`
-  - `gradient_accumulation_steps`
-  - `sequence_len`
-
-If it does not help, try running without deepspeed and without accelerate (replace "accelerate launch" with "python") in the command.
-
-Using adamw_bnb_8bit might also save you some memory.
-
-> `failed (exitcode: -9)`
-
-Usually means your system has run out of system memory.
-Similarly, you should consider reducing the same settings as when you run out of VRAM.
-Additionally, look into upgrading your system RAM which should be simpler than GPU upgrades.
-
-> RuntimeError: expected scalar type Float but found Half
-
-Try set `fp16: true`
-
-> NotImplementedError: No operator found for `memory_efficient_attention_forward` ...
-
-Try to turn off xformers.
+Other installation approaches are described [here](https://docs.axolotl.ai/docs/installation.html).
 
-> accelerate config missing
+#### Cloud Providers
 
-It's safe to ignore it.
-
-> NCCL Timeouts during training
-
-See the [NCCL](docs/nccl.qmd) guide.
-
-
-### Tokenization Mismatch b/w Inference & Training
-
-For many formats, Axolotl constructs prompts by concatenating token ids _after_ tokenizing strings.  The reason for concatenating token ids rather than operating on strings is to maintain precise accounting for attention masks.
-
-If you decode a prompt constructed by axolotl, you might see spaces between tokens (or lack thereof) that you do not expect, especially around delimiters and special tokens.  When you are starting out with a new format, you should always do the following:
-
-1. Materialize some data using `python -m axolotl.cli.preprocess your_config.yml --debug`, and then decode the first few rows with your model's tokenizer.
-2. During inference, right before you pass a tensor of token ids to your model, decode these tokens back into a string.
-3. Make sure the inference string from #2 looks **exactly** like the data you fine tuned on from #1, including spaces and new lines.  If they aren't the same, adjust your inference server accordingly.
-4. As an additional troubleshooting step, you can look at the token ids between 1 and 2 to make sure they are identical.
-
-Having misalignment between your prompts during training and inference can cause models to perform very poorly, so it is worth checking this.  See [this blog post](https://hamel.dev/notes/llm/05_tokenizer_gotchas.html) for a concrete example.
-
-## Debugging Axolotl
-
-See [this debugging guide](docs/debugging.qmd) for tips on debugging Axolotl, along with an example configuration for debugging with VSCode.
+<details>
 
-## Need help? 🙋
+- [RunPod](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
+- [Vast.ai](https://cloud.vast.ai?ref_id=62897&template_id=bdd4a49fa8bce926defc99471864cace&utm_source=github&utm_medium=developer_community&utm_campaign=template_launch_axolotl&utm_content=readme)
+- [PRIME Intellect](https://app.primeintellect.ai/dashboard/create-cluster?image=axolotl&location=Cheapest&security=Cheapest&show_spot=true)
+- [Modal](https://www.modal.com?utm_source=github&utm_medium=github&utm_campaign=axolotl)
+- [Novita](https://novita.ai/gpus-console?templateId=311)
+- [JarvisLabs.ai](https://jarvislabs.ai/templates/axolotl)
+- [Latitude.sh](https://latitude.sh/blueprint/989e0e79-3bf6-41ea-a46b-1f246e309d5c)
 
-Join our [Discord server](https://discord.gg/HhrNrHJPRb) where we our community members can help you.
+</details>
 
-Need dedicated support? Please contact us at [✉️wing@openaccessaicollective.org](mailto:wing@openaccessaicollective.org) for dedicated support options.
+### Your First Fine-tune
 
-## Badge ❤🏷️
+```bash
+# Fetch axolotl examples
+axolotl fetch examples
 
-Building something cool with Axolotl? Consider adding a badge to your model card.
+# Or, specify a custom path
+axolotl fetch examples --dest path/to/folder
 
-```markdown
-[<img src="https://raw.githubusercontent.com/OpenAccess-AI-Collective/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/OpenAccess-AI-Collective/axolotl)
+# Train a model using LoRA
+axolotl train examples/llama-3/lora-1b.yml
 ```
 
-[<img src="https://raw.githubusercontent.com/OpenAccess-AI-Collective/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/OpenAccess-AI-Collective/axolotl)
+That's it! Check out our [Getting Started Guide](https://docs.axolotl.ai/docs/getting-started.html) for a more detailed walkthrough.
 
-## Community Showcase
 
-Check out some of the projects and models that have been built using Axolotl! Have a model you'd like to add to our Community Showcase? Open a PR with your model.
+## 📚 Documentation
 
-Open Access AI Collective
-- [Minotaur 13b](https://huggingface.co/openaccess-ai-collective/minotaur-13b-fixed)
-- [Manticore 13b](https://huggingface.co/openaccess-ai-collective/manticore-13b)
-- [Hippogriff 30b](https://huggingface.co/openaccess-ai-collective/hippogriff-30b-chat)
+- [Installation Options](https://docs.axolotl.ai/docs/installation.html) - Detailed setup instructions for different environments
+- [Support Matrix](https://docs.axolotl.ai/docs/support-matrix.html) - Feature support, compatibility, and known gaps
+- [Configuration Guide](https://docs.axolotl.ai/docs/config-reference.html) - Full configuration options and examples
+- [Dataset Loading](https://docs.axolotl.ai/docs/dataset_loading.html) - Loading datasets from various sources
+- [Dataset Guide](https://docs.axolotl.ai/docs/dataset-formats/) - Supported formats and how to use them
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [Multipacking](https://docs.axolotl.ai/docs/multipack.html)
+- [API Reference](https://docs.axolotl.ai/docs/api/) - Auto-generated code documentation
+- [FAQ](https://docs.axolotl.ai/docs/faq.html) - Frequently asked questions
 
-PocketDoc Labs
-- [Dan's PersonalityEngine 13b LoRA](https://huggingface.co/PocketDoc/Dans-PersonalityEngine-13b-LoRA)
+## AI Agent Support
 
-## Contributing 🤝
+Axolotl ships with built-in documentation optimized for AI coding agents (Claude Code, Cursor, Copilot, etc.). These docs are bundled with the pip package, no repo clone needed.
 
-Please read the [contributing guide](./.github/CONTRIBUTING.md)
-
-Bugs? Please check the [open issues](https://github.com/OpenAccess-AI-Collective/axolotl/issues/bug) else create a new Issue.
-
-PRs are **greatly welcome**!
-
-Please run the quickstart instructions followed by the below to setup env:
 ```bash
-pip3 install -r requirements-dev.txt -r requirements-tests.txt
-pre-commit install
+# Show overview and available training methods
+axolotl agent-docs
 
-# test
-pytest tests/
+# Topic-specific references
+axolotl agent-docs sft                 # supervised fine-tuning
+axolotl agent-docs grpo                # GRPO online RL
+axolotl agent-docs preference_tuning   # DPO, KTO, ORPO, SimPO
+axolotl agent-docs reward_modelling    # outcome and process reward models
+axolotl agent-docs pretraining         # continual pretraining
+axolotl agent-docs --list              # list all topics
 
-# optional: run against all files
-pre-commit run --all-files
+# Dump config schema for programmatic use
+axolotl config-schema
+axolotl config-schema --field adapter
 ```
 
-Thanks to all of our contributors to date. Help drive open source AI progress forward by contributing to Axolotl.
+If you're working with the source repo, agent docs are also available at `docs/agents/` and the project overview is in `AGENTS.md`.
 
-<a href="https://github.com/openaccess-ai-collective/axolotl/graphs/contributors">
-  <img src="https://contrib.rocks/image?repo=openaccess-ai-collective/axolotl" alt="contributor chart by https://contrib.rocks"/>
-</a>
+## 🤝 Getting Help
 
-## Sponsors 🤝❤
+- Join our [Discord community](https://discord.gg/HhrNrHJPRb) for support
+- Check out our [Examples](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/) directory
+- Read our [Debugging Guide](https://docs.axolotl.ai/docs/debugging.html)
+- Need dedicated support? Please contact [✉️wing@axolotl.ai](mailto:wing@axolotl.ai) for options
 
-OpenAccess AI Collective is run by volunteer contributors such as [winglian](https://github.com/winglian),
-[NanoCode012](https://github.com/NanoCode012), [tmm1](https://github.com/tmm1),
-[mhenrichsen](https://github.com/mhenrichsen), [casper-hansen](https://github.com/casper-hansen),
-[hamelsmu](https://github.com/hamelsmu) and many more who help us accelerate forward by fixing bugs, answering
-community questions and implementing new features. Axolotl needs donations from sponsors for the compute needed to
-run our unit & integration tests, troubleshooting community issues, and providing bounties. If you love axolotl,
-consider sponsoring the project via [GitHub Sponsors](https://github.com/sponsors/OpenAccess-AI-Collective),
-[Ko-fi](https://ko-fi.com/axolotl_ai) or reach out directly to
-[wing@openaccessaicollective.org](mailto:wing@openaccessaicollective.org).
+## 🌟 Contributing
 
----
+Contributions are welcome! Please see our [Contributing Guide](https://github.com/axolotl-ai-cloud/axolotl/blob/main/.github/CONTRIBUTING.md) for details.
 
-#### 💎 Diamond Sponsors - [Contact directly](mailto:wing@openaccessaicollective.org)
+## 📈 Telemetry
 
----
+Axolotl has opt-out telemetry that helps us understand how the project is being used
+and prioritize improvements. We collect basic system information, model types, and
+error rates, never personal data or file paths. Telemetry is enabled by default. To
+disable it, set AXOLOTL_DO_NOT_TRACK=1. For more details, see our [telemetry documentation](https://docs.axolotl.ai/docs/telemetry.html).
 
-#### 🥇 Gold Sponsors - $5000/mo
+## ❤️ Sponsors
 
----
+Interested in sponsoring? Contact us at [wing@axolotl.ai](mailto:wing@axolotl.ai)
 
-#### 🥈 Silver Sponsors - $1000/mo
+## 📝 Citing Axolotl
 
----
+If you use Axolotl in your research or projects, please cite it as follows:
 
-#### 🥉 Bronze Sponsors - $500/mo
+```bibtex
+@software{axolotl,
+  title = {Axolotl: Open Source LLM Post-Training},
+  author = {{Axolotl maintainers and contributors}},
+  url = {https://github.com/axolotl-ai-cloud/axolotl},
+  license = {Apache-2.0},
+  year = {2023}
+}
+```
 
- - [JarvisLabs.ai](https://jarvislabs.ai)
+## 📜 License
 
----
+This project is licensed under the Apache 2.0 License - see the [LICENSE](LICENSE) file for details.
diff --git a/TODO.md b/TODO.md
deleted file mode 100644
index 2002bbbaf1..0000000000
--- a/TODO.md
+++ /dev/null
@@ -1,10 +0,0 @@
-# todo list
-
-- [] Validation of parameters for combinations that won't work
-
-
-
-## things that are known not to work
-
-- FSDP offload and gradient_checkpointing - https://github.com/pytorch/pytorch/issues/82203
-- adamw_bnb_8bit doesn't play well with FSDP offload
diff --git a/VERSION b/VERSION
new file mode 100644
index 0000000000..52fc5b872a
--- /dev/null
+++ b/VERSION
@@ -0,0 +1 @@
+0.17.0.dev
diff --git a/_quarto.yml b/_quarto.yml
index 749f68cce6..4a39091b58 100644
--- a/_quarto.yml
+++ b/_quarto.yml
@@ -1,12 +1,423 @@
 project:
   type: website
+  pre-render:
+   - docs/scripts/generate_config_docs.py
+   - docs/scripts/generate_examples_docs.py
 
+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
+        - prompters
+        - processing_strategies
+        - logging_config
+        - core.builders.base
+        - core.builders.causal
+        - core.builders.rl
+        - core.training_args
+        - core.training_args_base
+        - 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.art
+        - cli.checks
+        - cli.config
+        - cli.delinearize_llama4
+        - cli.inference
+        - cli.merge_lora
+        - cli.merge_sharded_fsdp_weights
+        - cli.preprocess
+        - cli.quantize
+        - cli.vllm_serve
+        - cli.agent_docs
+        - cli.cloud
+        - cli.cloud.base
+        - cli.cloud.baseten
+        - cli.cloud.modal_
+        - cli.utils
+        - cli.utils.args
+        - cli.utils.diffusion
+        - cli.utils.fetch
+        - cli.utils.load
+        - cli.utils.lora_merge
+        - cli.utils.sweeps
+        - cli.utils.train
+    - title: Trainers
+      desc: Training implementations
+      contents:
+        - core.trainers.base
+        - core.trainers.constants
+        - core.trainers.trl
+        - core.trainers.mamba
+        - core.trainers.dpo.args
+        - core.trainers.dpo.trainer
+        - core.trainers.ebft
+        - core.trainers.ebft.args
+        - core.trainers.ebft.kernels
+        - core.trainers.ebft.rewards
+        - core.trainers.ebft.strided
+        - core.trainers.ebft.trainer
+        - core.trainers.grpo
+        - core.trainers.grpo.args
+        - core.trainers.grpo.trainer
+        - core.trainers.grpo.async_trainer
+        - core.trainers.grpo.fast_async_trainer
+        - core.trainers.grpo.replay_buffer
+        - core.trainers.grpo.sampler
+        - core.trainers.utils
+    - title: Model Loading
+      desc: Functionality for loading and patching models, tokenizers, etc.
+      contents:
+        - loaders.model
+        - loaders.tokenizer
+        - loaders.processor
+        - loaders.adapter
+        - loaders.patch_manager
+        - loaders.constants
+        - loaders.utils
+    - title: Mixins
+      desc: Mixin classes for augmenting trainers
+      contents:
+        - core.trainers.mixins.activation_checkpointing
+        - core.trainers.mixins.checkpoints
+        - core.trainers.mixins.distributed_parallel
+        - core.trainers.mixins.layer_offloading
+        - core.trainers.mixins.optimizer
+        - core.trainers.mixins.packing
+        - core.trainers.mixins.rng_state_loader
+        - core.trainers.mixins.scheduler
+    - title: Context Managers
+      desc: Context managers for altering trainer behaviors
+      contents:
+        - utils.ctx_managers.sequence_parallel
+    - 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.context_qa
+        - prompt_strategies.creative_acr
+        - prompt_strategies.input_output
+        - prompt_strategies.pretrain
+        - prompt_strategies.stepwise_supervised
+        - prompt_strategies.metharme
+        - prompt_strategies.orcamini
+        - prompt_strategies.pygmalion
+        - prompt_strategies.messages.chat
+        - prompt_strategies.ebft.ebft_chat_multiturn
+        - prompt_strategies.ebft.ebft_opencode
+        - prompt_strategies.ebft.ebft_reasoning
+        - prompt_strategies.ebft.ebft_strided_chat
+        - prompt_strategies.ebft.ebft_strided_structured
+        - 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.chat_template
+        - prompt_strategies.bradley_terry.llama3
+    - title: Kernels
+      desc: Low-level performance optimizations
+      contents:
+        - kernels.lora
+        - kernels.dora
+        - kernels.geglu
+        - kernels.swiglu
+        - kernels.quantize
+        - kernels.autotune_telemetry
+        - kernels.gemma4_fused_rope
+        - kernels.rms_norm_gated
+        - kernels.utils
+    - title: Monkey Patches
+      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.lora_kernels
+        - monkeypatch.utils
+        - monkeypatch.btlm_attn_hijack_flash
+        - monkeypatch.stablelm_attn_hijack_flash
+        - monkeypatch.transformers_fa_utils
+        - monkeypatch.data.batch_dataset_fetcher
+        - monkeypatch.mixtral
+        - monkeypatch.gradient_checkpointing.offload_cpu
+        - monkeypatch.gradient_checkpointing.offload_disk
+        - monkeypatch.deepspeed_utils
+        - monkeypatch.fsdp2_qlora
+        - monkeypatch.gemma4_hybrid_mask
+        - monkeypatch.gemma4_loss_kwargs
+        - monkeypatch.kernelize_fixes
+        - monkeypatch.moe_quant
+        - monkeypatch.scaled_softmax_attn
+        - monkeypatch.torchao_optim
+        - monkeypatch.trainer_accelerator_args
+        - monkeypatch.accelerate.fsdp2
+        - monkeypatch.accelerate.parallelism_config
+        - monkeypatch.attention.flash_attn_4
+        - monkeypatch.attention.flex_attn
+        - monkeypatch.attention.fp8_attn
+        - monkeypatch.attention.sage_attn
+        - monkeypatch.attention.xformers
+        - monkeypatch.loss.chunked
+        - monkeypatch.loss.eaft
+        - monkeypatch.models.apertus.activation
+        - monkeypatch.models.falcon_h1.modeling
+        - monkeypatch.models.gemma4_unified.fused_attn
+        - monkeypatch.models.granitemoehybrid.modeling
+        - monkeypatch.models.kimi_linear.patch_kimi_linear
+        - monkeypatch.models.llama4.modeling
+        - monkeypatch.models.mamba_utils
+        - monkeypatch.models.mistral3.mistral_common_tokenizer
+        - monkeypatch.models.nemotron_h.modeling
+        - monkeypatch.models.pixtral.modeling_flash_attention_utils
+        - monkeypatch.models.qwen3.fused_attn
+        - monkeypatch.models.qwen3_5.fused_attn
+        - monkeypatch.models.qwen3_5.modeling
+        - monkeypatch.models.qwen3_5_moe.fused_attn
+        - monkeypatch.models.qwen3_moe.fused_attn
+        - monkeypatch.models.qwen3_next.modeling
+        - monkeypatch.models.qwen3_vl.fused_attn
+        - monkeypatch.models.voxtral.modeling
+        - monkeypatch.peft.utils
+        - monkeypatch.ring_attn.adapters.batch
+        - monkeypatch.ring_attn.patch
+        - monkeypatch.tiled_mlp.base
+        - monkeypatch.tiled_mlp.patch
+        - monkeypatch.trainer.lr
+        - monkeypatch.trainer.trl
+        - monkeypatch.trainer.trl_vllm
+        - monkeypatch.trainer.utils
+        - monkeypatch.transformers.trainer_loss_calc
+        - monkeypatch.xformers_
+    - title: Utils
+      desc: Utility functions
+      contents:
+        - utils.tokenization
+        - utils.chat_templates
+        - utils.chat_templates.base
+        - utils.lora
+        - utils.model_shard_quant
+        - utils.bench
+        - utils.comet_
+        - utils.config
+        - utils.cuda13
+        - utils.datasets
+        - utils.environment
+        - utils.fp32_norms
+        - utils.freeze
+        - utils.import_helper
+        - utils.logging
+        - utils.mlflow_
+        - utils.tee
+        - utils.trackio_
+        - utils.train
+        - utils.trainer
+        - utils.wandb_
+        - utils.weight_serde
+        - utils.schedulers
+        - utils.distributed
+        - utils.dict
+        - utils.generation.sft
+        - utils.mistral.mistral3_processor
+        - utils.mistral.mistral_tokenizer
+        - utils.optimizers.adopt
+        - utils.optimizers.qgalore
+        - utils.data.streaming
+        - utils.data.sft
+        - utils.data.rl
+        - utils.data.lock
+        - utils.data.utils
+        - utils.data.wrappers
+        - utils.quantization
+    - 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.deprecated
+        - utils.schemas.dynamic_checkpoint
+        - utils.schemas.fsdp
+        - utils.schemas.quantization
+        - utils.schemas.validation
+        - utils.schemas.vllm
+        - utils.schemas.enums
+        - utils.schemas.utils
+    - title: Integrations
+      desc: Third-party integrations and extensions
+      contents:
+        - integrations.base
+        - integrations.config
+        - integrations.cut_cross_entropy
+        - integrations.cut_cross_entropy.args
+        - integrations.densemixer.args
+        - integrations.densemixer.plugin
+        - integrations.diffusion.args
+        - integrations.diffusion.callbacks
+        - integrations.diffusion.generation
+        - integrations.diffusion.plugin
+        - integrations.diffusion.trainer
+        - integrations.diffusion.utils
+        - integrations.expert_parallel.args
+        - integrations.expert_parallel.buffer
+        - integrations.expert_parallel.experts_fn
+        - integrations.expert_parallel.plugin
+        - integrations.expert_parallel.shard
+        - integrations.grokfast.args
+        - integrations.grokfast.optimizer
+        - integrations.hatchery.args
+        - integrations.hatchery.data
+        - integrations.hatchery.plugin
+        - integrations.hatchery.rewards.math_reward
+        - integrations.hatchery.rl_trainer
+        - integrations.hatchery.trainer
+        - integrations.kd
+        - integrations.kd.args
+        - integrations.kd.callbacks
+        - integrations.kd.chat_template
+        - integrations.kd.collator
+        - integrations.kd.collator_online_teacher
+        - integrations.kd.kernels.liger
+        - integrations.kd.topk_logprob.forward_kl
+        - integrations.kd.trainer
+        - integrations.kd.utils
+        - integrations.kernels.args
+        - integrations.kernels.autotune_callback
+        - integrations.kernels.autotune_collector
+        - integrations.kernels.constants
+        - integrations.kernels.plugin
+        - integrations.liger.args
+        - integrations.liger.plugin
+        - integrations.liger.utils
+        - integrations.liger.models.base
+        - integrations.liger.models.deepseekv2
+        - integrations.liger.models.jamba
+        - integrations.liger.models.llama4
+        - integrations.liger.models.qwen3
+        - integrations.liger.models.qwen3_5
+        - integrations.liger.models.qwen3_5_moe
+        - integrations.liger.models.qwen3_moe
+        - integrations.llm_compressor.args
+        - integrations.llm_compressor.plugin
+        - integrations.llm_compressor.utils
+        - integrations.lm_eval.args
+        - integrations.lm_eval.cli
+        - integrations.mora.args
+        - integrations.mora.plugin
+        - integrations.nemo_gym.args
+        - integrations.nemo_gym.data_producer
+        - integrations.nemo_gym.dataset
+        - integrations.nemo_gym.multi_turn
+        - integrations.nemo_gym.plugin
+        - integrations.nemo_gym.rewards
+        - integrations.nemo_gym.server
+        - integrations.spectrum
+        - integrations.spectrum.args
+        - integrations.swanlab.args
+        - integrations.swanlab.callbacks
+        - integrations.swanlab.completion_logger
+        - integrations.swanlab.plugins
+    - title: Common
+      desc: Common utilities and shared functionality
+      contents:
+        - common.architectures
+        - common.const
+        - common.datasets
+    - title: Models
+      desc: Custom model implementations
+      contents:
+        - models.mamba.configuration_mamba
+        - models.mamba.modeling_mamba
+    - title: Data Processing
+      desc: Data processing utilities
+      contents:
+        - utils.collators.core
+        - utils.collators.batching
+        - utils.collators.dpo
+        - utils.collators.mamba
+        - utils.collators.mm_chat
+        - utils.samplers.multipack
+        - utils.samplers.utils
+    - title: Callbacks
+      desc: Training callbacks
+      contents:
+        - utils.callbacks
+        - utils.callbacks.perplexity
+        - utils.callbacks.profiler
+        - utils.callbacks.lisa
+        - utils.callbacks.mlflow_
+        - utils.callbacks.comet_
+        - utils.callbacks.qat
+        - utils.callbacks.dynamic_checkpoint
+        - utils.callbacks.generation
+        - utils.callbacks.models
+        - utils.callbacks.opentelemetry
+        - utils.callbacks.swanlab
+        - utils.callbacks.tokens_per_second
+        - utils.callbacks.trackio_
+    - title: Scripts
+      desc: Standalone helper scripts
+      contents:
+        - scripts.process_cleanup
+        - scripts.vllm_serve_lora
+        - scripts.vllm_worker_ext
+    - title: Telemetry
+      desc: Usage telemetry
+      contents:
+        - telemetry.callbacks
+        - telemetry.errors
+        - telemetry.manager
+        - telemetry.runtime_metrics
 website:
   title: "Axolotl"
-  description: "Fine-tuning"
+  description: "We make fine-tuning accessible, scalable, and fun"
   favicon: favicon.jpg
+
+  google-analytics: "G-9KYCVJBNMQ"
+
   navbar:
-    title: Axolotl
+    logo: image/axolotl_logo_digital_white.svg
+    title: false
     background: dark
     pinned: false
     collapse: false
@@ -14,7 +425,7 @@ website:
     - icon: twitter
       href: https://twitter.com/axolotl_ai
     - icon: github
-      href: https://github.com/OpenAccess-AI-Collective/axolotl/
+      href: https://github.com/axolotl-ai-cloud/axolotl/
     - icon: discord
       href: https://discord.gg/7m9sfhzaf3
 
@@ -25,27 +436,137 @@ website:
       contents:
         - text: Home
           href: index.qmd
-        - section: "How-To Guides"
+
+        - section: "Getting Started"
           contents:
-          # TODO Edit folder structure after we have more docs.
-            - docs/debugging.qmd
-            - docs/multipack.qmd
-            - docs/fsdp_qlora.qmd
-            - docs/input_output.qmd
-            - docs/rlhf.qmd
-            - docs/nccl.qmd
-            - docs/mac.qmd
-            - docs/multi-node.qmd
+            - docs/getting-started.qmd
+            - docs/choosing_method.qmd
+            - docs/installation.qmd
+            - docs/inference.qmd
+            - docs/support-matrix.qmd
+            - section: "Model Guides"
+              contents:
+                - docs/models/kimi-linear.qmd
+                - docs/models/plano.qmd
+                - docs/models/mimo.qmd
+                - docs/models/internvl3_5.qmd
+                - docs/models/olmo3.qmd
+                - docs/models/trinity.qmd
+                - docs/models/arcee.qmd
+                - section: "Ministral3"
+                  contents:
+                    - docs/models/ministral3.qmd
+                    - docs/models/ministral3/think.qmd
+                    - docs/models/ministral3/vision.qmd
+                - section: "Magistral"
+                  contents:
+                    - docs/models/magistral.qmd
+                    - docs/models/magistral/think.qmd
+                    - docs/models/magistral/vision.qmd
+                - docs/models/ministral.qmd
+                - docs/models/mistral-small.qmd
+                - docs/models/voxtral.qmd
+                - docs/models/devstral.qmd
+                - docs/models/mistral.qmd
+                - docs/models/llama-4.qmd
+                - docs/models/llama-2.qmd
+                - docs/models/qwen3-next.qmd
+                - docs/models/qwen3.qmd
+                - docs/models/gemma3n.qmd
+                - docs/models/apertus.qmd
+                - docs/models/gpt-oss.qmd
+                - docs/models/seed-oss.qmd
+                - docs/models/phi.qmd
+                - docs/models/smolvlm2.qmd
+                - docs/models/granite4.qmd
+                - docs/models/LiquidAI.qmd
+                - docs/models/hunyuan.qmd
+                - docs/models/jamba.qmd
+                - docs/models/orpheus.qmd
+
+            - docs/cli.qmd
+            - docs/telemetry.qmd
+            - docs/config-reference.qmd
+            - text: "API Reference"
+              href: docs/api
+
         - section: "Dataset Formats"
           contents: docs/dataset-formats/*
-        - section: "Reference"
+
+        - section: "Deployments"
           contents:
-            - docs/config.qmd
-        - docs/faq.qmd
+            - docs/docker.qmd
+            - docs/multi-gpu.qmd
+            - docs/multi-node.qmd
+            - docs/ray-integration.qmd
+            - docs/amd_hpc.qmd
+            - docs/mac.qmd
 
+        - section: "How To Guides"
+          contents:
+            - docs/multimodal.qmd
+            - docs/rlhf.qmd
+            - docs/grpo.qmd
+            - docs/ebft.qmd
+            - docs/vllm_serving.qmd
+            - docs/reward_modelling.qmd
+            - docs/lr_groups.qmd
+            - docs/lora.qmd
+            - docs/lora_optims.qmd
+            - docs/dataset_loading.qmd
+            - docs/qat.qmd
+            - docs/quantize.qmd
+            - docs/1_58bit_finetuning.qmd
+            - docs/optimizations.qmd
+
+        - section: "Core Concepts"
+          contents:
+            - docs/batch_vs_grad.qmd
+            - docs/dataset_preprocessing.qmd
+            - docs/streaming.qmd
+            - docs/multipack.qmd
+            - docs/mixed_precision.qmd
+            - docs/optimizers.qmd
+            - docs/attention.qmd
+
+        - section: "Advanced Features"
+          contents:
+            - docs/fsdp_qlora.qmd
+            - docs/torchao.qmd
+            - docs/custom_integrations.qmd
+            - docs/sequence_parallelism.qmd
+            - docs/gradient_checkpointing.qmd
+            - docs/nd_parallelism.qmd
+            - docs/expert_quantization.qmd
+
+        - section: "Troubleshooting"
+          contents:
+            - docs/faq.qmd
+            - docs/training_stability.qmd
+            - docs/debugging.qmd
+            - docs/nccl.qmd
 
 format:
   html:
-    theme: materia
+    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
diff --git a/cicd/Dockerfile-uv.jinja b/cicd/Dockerfile-uv.jinja
new file mode 100644
index 0000000000..728970b008
--- /dev/null
+++ b/cicd/Dockerfile-uv.jinja
@@ -0,0 +1,53 @@
+FROM axolotlai/axolotl-base-uv:{{ BASE_TAG }}
+
+ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
+ENV AXOLOTL_EXTRAS="{{ AXOLOTL_EXTRAS }}"
+ENV AXOLOTL_ARGS="{{ AXOLOTL_ARGS }}"
+ENV CUDA="{{ CUDA }}"
+ENV PYTORCH_VERSION="{{ PYTORCH_VERSION }}"
+ENV GITHUB_REF="{{ GITHUB_REF }}"
+ENV GITHUB_SHA="{{ GITHUB_SHA }}"
+ENV NIGHTLY_BUILD="{{ NIGHTLY_BUILD }}"
+ENV HF_HOME="{{ HF_HOME }}"
+
+WORKDIR /workspace
+
+RUN git clone --depth=1 https://github.com/axolotl-ai-cloud/axolotl.git
+
+WORKDIR /workspace/axolotl
+
+RUN git fetch origin +$GITHUB_REF && \
+    git checkout FETCH_HEAD
+
+RUN uv pip install packaging==26.0 setuptools==78.1.1
+RUN uv pip uninstall causal_conv1d
+RUN uv pip install "huggingface_hub>=1.5.0" "httpx<1"
+RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
+        uv pip install --no-build-isolation -e .[deepspeed,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+    else \
+        uv pip install --no-build-isolation -e .[deepspeed,optimizers,ray] $AXOLOTL_ARGS; \
+    fi
+
+# Override with nightly HF packages for nightly builds
+RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
+        uv pip install "kernels>=0.15.2,<0.16"; \
+        uv pip install --no-deps \
+            "transformers @ git+https://github.com/huggingface/transformers.git@main" \
+            "peft @ git+https://github.com/huggingface/peft.git@main" \
+            "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
+            "trl @ git+https://github.com/huggingface/trl.git@main" \
+            "datasets @ git+https://github.com/huggingface/datasets.git@main"; \
+    fi
+
+RUN python scripts/cutcrossentropy_install.py --uv | sh
+
+# So we can test the Docker image
+RUN uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+    codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
+
+# fix so that git fetch/pull from remote works
+RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
+    git config --get remote.origin.fetch
+
+# helper for huggingface-login cli
+RUN git config --global credential.helper store
diff --git a/cicd/Dockerfile.jinja b/cicd/Dockerfile.jinja
deleted file mode 100644
index ce03e08c2f..0000000000
--- a/cicd/Dockerfile.jinja
+++ /dev/null
@@ -1,40 +0,0 @@
-FROM winglian/axolotl-base:{{ BASE_TAG }}
-
-ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
-ENV AXOLOTL_EXTRAS="{{ AXOLOTL_EXTRAS }}"
-ENV AXOLOTL_ARGS="{{ AXOLOTL_ARGS }}"
-ENV CUDA="{{ CUDA }}"
-ENV BNB_CUDA_VERSION="{{ CUDA }}"
-ENV PYTORCH_VERSION="{{ PYTORCH_VERSION }}"
-ENV GITHUB_REF="{{ GITHUB_REF }}"
-ENV GITHUB_SHA="{{ GITHUB_SHA }}"
-
-RUN apt-get update && \
-    apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev
-
-WORKDIR /workspace
-
-RUN git clone --depth=1 https://github.com/OpenAccess-AI-Collective/axolotl.git
-
-WORKDIR /workspace/axolotl
-
-RUN git fetch origin +$GITHUB_REF && \
-    git checkout FETCH_HEAD
-
-# If AXOLOTL_EXTRAS is set, append it in brackets
-RUN pip install causal_conv1d
-RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,galore,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
-    else \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,galore] $AXOLOTL_ARGS; \
-    fi
-
-# So we can test the Docker image
-RUN pip install pytest
-
-# fix so that git fetch/pull from remote works
-RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
-    git config --get remote.origin.fetch
-
-# helper for huggingface-login cli
-RUN git config --global credential.helper store
diff --git a/src/axolotl/utils/config/models/input/__init__.py b/cicd/__init__.py
similarity index 100%
rename from src/axolotl/utils/config/models/input/__init__.py
rename to cicd/__init__.py
diff --git a/cicd/cicd.sh b/cicd/cicd.sh
index fa2049b6bd..ce3450cb1a 100755
--- a/cicd/cicd.sh
+++ b/cicd/cicd.sh
@@ -1,5 +1,95 @@
 #!/bin/bash
+set -e
 
-pytest --ignore=tests/e2e/ /workspace/axolotl/tests/
-pytest /workspace/axolotl/tests/e2e/patched/
-pytest --ignore=tests/e2e/patched/ /workspace/axolotl/tests/e2e/
+python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__, f'Expected torch $PYTORCH_VERSION but got {torch.__version__}'"
+
+set -o pipefail
+for i in 1 2 3; do
+  if curl --silent --show-error --fail -L \
+    https://axolotl-ci.b-cdn.net/hf-cache.tar.zst \
+    | tar -xpf - -C "${HF_HOME}/hub/" --use-compress-program unzstd --strip-components=1; then
+    echo "HF cache extracted successfully"
+    break
+  fi
+  echo "Attempt $i failed, cleaning up and retrying in 15s..."
+  rm -rf "${HF_HOME}/hub/"*
+  sleep 15
+done
+# hf download "NousResearch/Meta-Llama-3-8B"
+# hf download "NousResearch/Meta-Llama-3-8B-Instruct"
+# hf download "microsoft/Phi-4-reasoning"
+# hf download "microsoft/Phi-3.5-mini-instruct"
+# hf download "microsoft/Phi-3-medium-128k-instruct"
+
+# Run unit tests with initial coverage report
+pytest -v --durations=10 -n8 \
+  --ignore=tests/e2e/ \
+  --ignore=tests/integrations/ \
+  --ignore=tests/patched/ \
+  --ignore=tests/cli \
+  /workspace/axolotl/tests/ \
+  --cov=axolotl
+
+# Run lora kernels tests with coverage append
+pytest -v --durations=10 \
+  /workspace/axolotl/tests/e2e/patched/lora_kernels \
+  --cov=axolotl \
+  --cov-append
+
+# Run patched tests excluding lora kernels with coverage append
+pytest --full-trace -vvv --durations=10 \
+  --ignore=tests/e2e/patched/lora_kernels \
+  /workspace/axolotl/tests/e2e/patched \
+  --cov=axolotl \
+  --cov-append
+
+# Run solo tests with coverage append
+# test_rm_lora is run in its own process below (it fails on py3.11 when sharing
+# the solo process with other tests; isolating it avoids cross-test state).
+pytest -v --durations=10 -n1 \
+  --ignore=tests/e2e/solo/test_reward_model_smollm2.py \
+  /workspace/axolotl/tests/e2e/solo/ \
+  --cov=axolotl \
+  --cov-append
+
+# Run reward-model test isolated in its own process
+pytest -v --durations=10 -s \
+  /workspace/axolotl/tests/e2e/solo/test_reward_model_smollm2.py \
+  --cov=axolotl \
+  --cov-append
+
+# Run E2E integration tests with coverage append
+pytest -v --durations=10 \
+  /workspace/axolotl/tests/e2e/integrations/ \
+  --cov=axolotl \
+  --cov-append
+
+pytest -v --durations=10 -n8 --dist loadfile \
+  --ignore=tests/integrations/kernels/ \
+  --ignore=tests/integrations/monkeypatch/test_tiled_mlp_moe.py \
+  --ignore=tests/integrations/test_gemma4_moe.py \
+  --ignore=tests/integrations/test_scattermoe_lora.py \
+  --ignore=tests/integrations/test_scattermoe_lora_kernels.py \
+  --ignore=tests/integrations/test_scattermoe_multi_lora.py \
+  --ignore=tests/integrations/test_sonicmoe_multi_lora.py \
+  /workspace/axolotl/tests/integrations/ \
+  --cov=axolotl \
+  --cov-append
+
+pytest -v --durations=10 /workspace/axolotl/tests/cli \
+  --cov=axolotl \
+  --cov-append
+
+# Run remaining e2e tests with coverage append and final report
+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/ \
+  --cov=axolotl \
+  --cov-append \
+  --cov-report=xml:e2e-coverage.xml
+
+codecov upload-process -t $CODECOV_TOKEN -f e2e-coverage.xml -F e2e,pytorch-${PYTORCH_VERSION} || true
diff --git a/cicd/cicd_cuda_kernels.sh b/cicd/cicd_cuda_kernels.sh
new file mode 100755
index 0000000000..ac1c444e52
--- /dev/null
+++ b/cicd/cicd_cuda_kernels.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+set -e
+
+python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__, f'Expected torch $PYTORCH_VERSION but got {torch.__version__}'"
+
+set -o pipefail
+for i in 1 2 3; do
+  if curl --silent --show-error --fail -L \
+    https://axolotl-ci.b-cdn.net/hf-cache.tar.zst \
+    | tar -xpf - -C "${HF_HOME}/hub/" --use-compress-program unzstd --strip-components=1; then
+    echo "HF cache extracted successfully"
+    break
+  fi
+  echo "Attempt $i failed, cleaning up and retrying in 15s..."
+  rm -rf "${HF_HOME}/hub/"*
+  sleep 15
+done
+
+pytest -v --durations=10 \
+  /workspace/axolotl/tests/integrations/kernels/ \
+  /workspace/axolotl/tests/integrations/monkeypatch/test_tiled_mlp_moe.py \
+  /workspace/axolotl/tests/integrations/test_gemma4_moe.py \
+  /workspace/axolotl/tests/integrations/test_scattermoe_lora.py \
+  /workspace/axolotl/tests/integrations/test_scattermoe_lora_kernels.py \
+  /workspace/axolotl/tests/integrations/test_scattermoe_multi_lora.py \
+  /workspace/axolotl/tests/integrations/test_sonicmoe_multi_lora.py \
+  --cov=axolotl \
+  --cov-report=xml:e2e-kernel-coverage.xml
+
+codecov upload-process -t "$CODECOV_TOKEN" -f e2e-kernel-coverage.xml -F e2e,kernels,pytorch-${PYTORCH_VERSION} || true
diff --git a/cicd/cleanup.py b/cicd/cleanup.py
new file mode 100644
index 0000000000..0074899938
--- /dev/null
+++ b/cicd/cleanup.py
@@ -0,0 +1,19 @@
+"""Modal app to run axolotl GPU cleanup"""
+
+from .single_gpu import VOLUME_CONFIG, app, cicd_image, run_cmd
+
+
+@app.function(
+    image=cicd_image,
+    timeout=60 * 60,
+    cpu=8.0,
+    memory=131072,
+    volumes=VOLUME_CONFIG,
+)
+def cleanup():
+    run_cmd("./cicd/cleanup.sh", "/workspace/axolotl")
+
+
+@app.local_entrypoint()
+def main():
+    cleanup.remote()
diff --git a/cicd/cleanup.sh b/cicd/cleanup.sh
new file mode 100755
index 0000000000..4ea851bb4e
--- /dev/null
+++ b/cicd/cleanup.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+set -e
+
+# cleanup old cache files for datasets processing and intermediate mappings
+find /workspace/data/huggingface-cache/hub/datasets -name "cache-*" -type f -mtime +1 -exec rm {} \;
+find /workspace/data/huggingface-cache/hub/datasets -name "*.lock" -type f -mtime +1 -exec rm {} \;
diff --git a/cicd/e2e_cuda_kernels.py b/cicd/e2e_cuda_kernels.py
new file mode 100644
index 0000000000..92c0af6241
--- /dev/null
+++ b/cicd/e2e_cuda_kernels.py
@@ -0,0 +1,20 @@
+"""Modal app to run CUDA-heavy single-GPU tests."""
+
+from .single_gpu import GPU_CONFIG, VOLUME_CONFIG, app, cicd_image, run_cmd
+
+
+@app.function(
+    image=cicd_image,
+    gpu=GPU_CONFIG,
+    timeout=90 * 60,
+    cpu=8.0,
+    memory=131072,
+    volumes=VOLUME_CONFIG,
+)
+def cicd_cuda_kernels():
+    run_cmd("./cicd/cicd_cuda_kernels.sh", "/workspace/axolotl")
+
+
+@app.local_entrypoint()
+def main():
+    cicd_cuda_kernels.remote()
diff --git a/cicd/e2e_tests.py b/cicd/e2e_tests.py
new file mode 100644
index 0000000000..5d2b6fed17
--- /dev/null
+++ b/cicd/e2e_tests.py
@@ -0,0 +1,20 @@
+"""Modal app to run axolotl GPU tests"""
+
+from .single_gpu import GPU_CONFIG, VOLUME_CONFIG, app, cicd_image, run_cmd
+
+
+@app.function(
+    image=cicd_image,
+    gpu=GPU_CONFIG,
+    timeout=120 * 60,  # 90 min
+    cpu=8.0,
+    memory=131072,
+    volumes=VOLUME_CONFIG,
+)
+def cicd_pytest():
+    run_cmd("./cicd/cicd.sh", "/workspace/axolotl")
+
+
+@app.local_entrypoint()
+def main():
+    cicd_pytest.remote()
diff --git a/cicd/multigpu.py b/cicd/multigpu.py
new file mode 100644
index 0000000000..de1cdf8937
--- /dev/null
+++ b/cicd/multigpu.py
@@ -0,0 +1,84 @@
+"""
+modal application to run axolotl gpu tests in Modal
+"""
+
+import os
+import pathlib
+import tempfile
+
+import jinja2
+import modal
+from jinja2 import select_autoescape
+from modal import App, Image
+
+cicd_path = pathlib.Path(__file__).parent.resolve()
+
+template_loader = jinja2.FileSystemLoader(searchpath=cicd_path)
+template_env = jinja2.Environment(
+    loader=template_loader, autoescape=select_autoescape()
+)
+dockerfile = os.environ.get("E2E_DOCKERFILE", "Dockerfile-uv.jinja")
+df_template = template_env.get_template(dockerfile)
+
+df_args = {
+    "AXOLOTL_EXTRAS": os.environ.get("AXOLOTL_EXTRAS", ""),
+    "AXOLOTL_ARGS": os.environ.get("AXOLOTL_ARGS", ""),
+    "PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.6.0"),
+    "BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.11-cu126-2.6.0"),
+    "CUDA": os.environ.get("CUDA", "126"),
+    "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
+    "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
+    "NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""),
+    "CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""),
+    "HF_HOME": "/workspace/data/huggingface-cache/hub",
+    "PYTHONUNBUFFERED": os.environ.get("PYTHONUNBUFFERED", "1"),
+    "DEEPSPEED_LOG_LEVEL": os.environ.get("DEEPSPEED_LOG_LEVEL", "WARNING"),
+}
+
+dockerfile_contents = df_template.render(**df_args)
+
+temp_dir = tempfile.mkdtemp()
+with open(pathlib.Path(temp_dir) / "Dockerfile", "w", encoding="utf-8") as f:
+    f.write(dockerfile_contents)
+
+cicd_image = Image.from_dockerfile(
+    pathlib.Path(temp_dir) / "Dockerfile",
+    gpu="A10G",
+).env(df_args)
+
+app = App("Axolotl CI/CD", secrets=[])
+
+hf_cache_volume = modal.Volume.from_name(
+    "axolotl-ci-hf-hub-cache", create_if_missing=True
+)
+VOLUME_CONFIG = {
+    "/workspace/data/huggingface-cache/hub": hf_cache_volume,
+}
+
+N_GPUS = int(os.environ.get("N_GPUS", 2))
+GPU_CONFIG = f"H100:{N_GPUS}"
+
+
+def run_cmd(cmd: str, run_folder: str):
+    import subprocess  # nosec
+
+    # Propagate errors from subprocess.
+    if exit_code := subprocess.call(cmd.split(), cwd=run_folder):  # nosec
+        exit(exit_code)
+
+
+@app.function(
+    image=cicd_image,
+    gpu=GPU_CONFIG,
+    timeout=120 * 60,
+    cpu=16.0,
+    memory=131072 * N_GPUS,
+    volumes=VOLUME_CONFIG,
+)
+def cicd_pytest():
+    run_cmd("./cicd/multigpu.sh", "/workspace/axolotl")
+
+
+@app.local_entrypoint()
+def main():
+    cicd_pytest.remote()
diff --git a/cicd/multigpu.sh b/cicd/multigpu.sh
new file mode 100755
index 0000000000..e00fe909ee
--- /dev/null
+++ b/cicd/multigpu.sh
@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e
+
+# Only run two tests at a time to avoid OOM on GPU (with coverage collection)
+pytest -v --durations=10 -n2 --maxfail=3 \
+  --ignore=/workspace/axolotl/tests/e2e/multigpu/solo/ \
+  --ignore=/workspace/axolotl/tests/e2e/multigpu/patched/ \
+  /workspace/axolotl/tests/e2e/multigpu/ \
+  --cov=axolotl
+
+# Run solo tests with coverage append
+pytest -v --durations=10 -n1 \
+  /workspace/axolotl/tests/e2e/multigpu/solo/ \
+  --cov=axolotl \
+  --cov-append
+
+pytest -v  --durations=10 -n1 /workspace/axolotl/tests/e2e/multigpu/patched/ \
+  --cov=axolotl \
+  --cov-append \
+  --cov-report=xml:multigpu-coverage.xml
+
+# Upload coverage to Codecov if CODECOV_TOKEN is available
+if [ -n "$CODECOV_TOKEN" ]; then
+  codecov upload-process -t "${CODECOV_TOKEN}" -f multigpu-coverage.xml -F multigpu,docker-tests,pytorch-${PYTORCH_VERSION} || true
+fi
diff --git a/cicd/single_gpu.py b/cicd/single_gpu.py
new file mode 100644
index 0000000000..9850ed2be1
--- /dev/null
+++ b/cicd/single_gpu.py
@@ -0,0 +1,72 @@
+"""Modal app to run axolotl GPU tests"""
+
+import os
+import pathlib
+import tempfile
+
+import jinja2
+import modal
+import modal.experimental
+from jinja2 import select_autoescape
+from modal import App
+
+cicd_path = pathlib.Path(__file__).parent.resolve()
+
+template_loader = jinja2.FileSystemLoader(searchpath=cicd_path)
+template_env = jinja2.Environment(
+    loader=template_loader, autoescape=select_autoescape()
+)
+dockerfile = os.environ.get("E2E_DOCKERFILE", "Dockerfile-uv.jinja")
+df_template = template_env.get_template(dockerfile)
+
+df_args = {
+    "AXOLOTL_EXTRAS": os.environ.get("AXOLOTL_EXTRAS", ""),
+    "AXOLOTL_ARGS": os.environ.get("AXOLOTL_ARGS", ""),
+    "PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.6.0"),
+    "BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.11-cu126-2.6.0"),
+    "CUDA": os.environ.get("CUDA", "126"),
+    "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
+    "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
+    "NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""),
+    "CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""),
+    "HF_HOME": "/workspace/data/huggingface-cache/hub",
+    "PYTHONUNBUFFERED": os.environ.get("PYTHONUNBUFFERED", "1"),
+    "DEEPSPEED_LOG_LEVEL": os.environ.get("DEEPSPEED_LOG_LEVEL", "WARNING"),
+}
+
+dockerfile_contents = df_template.render(**df_args)
+
+temp_dir = tempfile.mkdtemp()
+with open(pathlib.Path(temp_dir) / "Dockerfile", "w", encoding="utf-8") as f:
+    f.write(dockerfile_contents)
+
+cicd_image = modal.experimental.raw_dockerfile_image(
+    pathlib.Path(temp_dir) / "Dockerfile",
+    # context_mount=None,
+    # gpu="A10G",
+).env(df_args)
+
+app = App("Axolotl CI/CD", secrets=[])
+
+hf_cache_volume = modal.Volume.from_name(
+    "axolotl-ci-hf-hub-cache", create_if_missing=True
+)
+VOLUME_CONFIG = {
+    "/workspace/data/huggingface-cache/hub": hf_cache_volume,
+}
+
+N_GPUS = int(os.environ.get("N_GPUS", 1))
+GPU_TYPE = os.environ.get("GPU_TYPE", "L40S")
+GPU_CONFIG = f"{GPU_TYPE}:{N_GPUS}"
+
+
+def run_cmd(cmd: str, run_folder: str):
+    import subprocess  # nosec
+
+    sp_env = os.environ.copy()
+    sp_env["AXOLOTL_DATASET_NUM_PROC"] = "8"
+
+    # Propagate errors from subprocess.
+    exit_code = subprocess.call(cmd.split(), cwd=run_folder, env=sp_env)  # nosec
+    if exit_code:
+        raise RuntimeError(f"Command '{cmd}' failed with exit code {exit_code}")
diff --git a/cicd/tests.py b/cicd/tests.py
deleted file mode 100644
index bfbdb7b90a..0000000000
--- a/cicd/tests.py
+++ /dev/null
@@ -1,75 +0,0 @@
-"""
- modal application to run axolotl gpu tests in Modal
- """
-import os
-import pathlib
-import tempfile
-
-import jinja2
-import modal
-from jinja2 import select_autoescape
-from modal import Image, Stub
-
-cicd_path = pathlib.Path(__file__).parent.resolve()
-
-template_loader = jinja2.FileSystemLoader(searchpath=cicd_path)
-template_env = jinja2.Environment(
-    loader=template_loader, autoescape=select_autoescape()
-)
-df_template = template_env.get_template("Dockerfile.jinja")
-
-df_args = {
-    "AXOLOTL_EXTRAS": os.environ.get("AXOLOTL_EXTRAS", ""),
-    "AXOLOTL_ARGS": os.environ.get("AXOLOTL_ARGS", ""),
-    "PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.0.1"),
-    "BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.10-cu118-2.0.1"),
-    "CUDA": os.environ.get("CUDA", "118"),
-    "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
-    "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
-}
-
-dockerfile_contents = df_template.render(**df_args)
-
-temp_dir = tempfile.mkdtemp()
-with open(pathlib.Path(temp_dir) / "Dockerfile", "w", encoding="utf-8") as f:
-    f.write(dockerfile_contents)
-
-cicd_image = (
-    Image.from_dockerfile(
-        pathlib.Path(temp_dir) / "Dockerfile",
-        force_build=True,
-        gpu="A10G",
-    )
-    .env(df_args)
-    .pip_install("fastapi==0.110.0", "pydantic==2.6.3")
-)
-
-stub = Stub("Axolotl CI/CD", secrets=[])
-
-
-N_GPUS = int(os.environ.get("N_GPUS", 1))
-GPU_CONFIG = modal.gpu.A10G(count=N_GPUS)
-
-
-def run_cmd(cmd: str, run_folder: str):
-    import subprocess  # nosec
-
-    # Propagate errors from subprocess.
-    if exit_code := subprocess.call(cmd.split(), cwd=run_folder):  # nosec
-        exit(exit_code)  # pylint: disable=consider-using-sys-exit
-
-
-@stub.function(
-    image=cicd_image,
-    gpu=GPU_CONFIG,
-    timeout=45 * 60,
-    cpu=8.0,
-    memory=131072,
-)
-def cicd_pytest():
-    run_cmd("./cicd/cicd.sh", "/workspace/axolotl")
-
-
-@stub.local_entrypoint()
-def main():
-    cicd_pytest.remote()
diff --git a/codecov.yml b/codecov.yml
new file mode 100644
index 0000000000..a876dd59d0
--- /dev/null
+++ b/codecov.yml
@@ -0,0 +1,58 @@
+codecov:
+  require_ci_to_pass: yes
+  notify:
+    wait_for_ci: true
+
+coverage:
+  precision: 2
+  round: down
+  range: "70...100"
+  status:
+    project:
+      default:
+        # basic
+        target: auto
+        threshold: 1%
+        base: auto
+        # advanced
+        branches: null
+        if_no_uploads: error
+        if_not_found: success
+        if_ci_failed: error
+        only_pulls: true
+        flags: null
+        paths: null
+        informational: true
+    patch:
+      default:
+        # basic
+        target: auto
+        threshold: 1%
+        base: auto
+        # advanced
+        branches: null
+        if_no_uploads: error
+        if_not_found: success
+        if_ci_failed: error
+        only_pulls: false
+        flags: null
+        paths: null
+        informational: true
+
+parsers:
+  gcov:
+    branch_detection:
+      conditional: yes
+      loop: yes
+      method: no
+      macro: no
+
+comment:
+  layout: "reach,diff,flags,files,footer"
+  behavior: default
+  require_changes: no
+  require_base: no
+  require_head: yes
+
+github_checks:
+  annotations: false
diff --git a/deepspeed_configs/zero1_torch_compile.json b/deepspeed_configs/zero1_torch_compile.json
new file mode 100644
index 0000000000..b884513923
--- /dev/null
+++ b/deepspeed_configs/zero1_torch_compile.json
@@ -0,0 +1,27 @@
+{
+  "zero_optimization": {
+    "stage": 1,
+    "overlap_comm": true
+  },
+  "bf16": {
+    "enabled": "auto"
+  },
+  "fp16": {
+    "enabled": "auto",
+    "auto_cast": false,
+    "loss_scale": 0,
+    "initial_scale_power": 32,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "compile": {
+    "disable": false,
+    "backend": "inductor"
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}
diff --git a/deepspeed_configs/zero2_torch_compile.json b/deepspeed_configs/zero2_torch_compile.json
new file mode 100644
index 0000000000..c3bcf98cf5
--- /dev/null
+++ b/deepspeed_configs/zero2_torch_compile.json
@@ -0,0 +1,31 @@
+{
+  "compile": {
+    "disable": false,
+    "backend": "inductor"
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "offload_optimizer": {
+      "device": "cpu"
+    },
+    "contiguous_gradients": true,
+    "overlap_comm": true
+  },
+  "bf16": {
+    "enabled": "auto"
+  },
+  "fp16": {
+    "enabled": "auto",
+    "auto_cast": false,
+    "loss_scale": 0,
+    "initial_scale_power": 32,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}
diff --git a/deepspeed_configs/zero3.json b/deepspeed_configs/zero3.json
index 90ec3677ea..f8c9cdfe0d 100644
--- a/deepspeed_configs/zero3.json
+++ b/deepspeed_configs/zero3.json
@@ -7,9 +7,9 @@
     "reduce_bucket_size": "auto",
     "stage3_prefetch_bucket_size": "auto",
     "stage3_param_persistence_threshold": "auto",
-    "stage3_max_live_parameters": 0,
-    "stage3_max_reuse_distance": 0,
-    "stage3_gather_16bit_weights_on_model_save": true
+    "max_live_parameters": 0,
+    "max_reuse_distance": 0,
+    "gather_16bit_weights_on_model_save": true
   },
   "bf16": {
     "enabled": "auto"
diff --git a/deepspeed_configs/zero3_bf16.json b/deepspeed_configs/zero3_bf16.json
index 16e64d76b4..a69e13cf7c 100644
--- a/deepspeed_configs/zero3_bf16.json
+++ b/deepspeed_configs/zero3_bf16.json
@@ -7,22 +7,13 @@
     "reduce_bucket_size": "auto",
     "stage3_prefetch_bucket_size": "auto",
     "stage3_param_persistence_threshold": "auto",
-    "stage3_max_live_parameters": 0,
-    "stage3_max_reuse_distance": 0,
-    "stage3_gather_16bit_weights_on_model_save": true
+    "max_live_parameters": 0,
+    "max_reuse_distance": 0,
+    "gather_16bit_weights_on_model_save": true
   },
   "bf16": {
     "enabled": true
   },
-  "fp16": {
-    "enabled": "auto",
-    "auto_cast": false,
-    "loss_scale": 0,
-    "initial_scale_power": 32,
-    "loss_scale_window": 1000,
-    "hysteresis": 2,
-    "min_loss_scale": 1
-  },
   "gradient_accumulation_steps": "auto",
   "gradient_clipping": "auto",
   "train_batch_size": "auto",
diff --git a/deepspeed_configs/zero3_bf16_cpuoffload_all.json b/deepspeed_configs/zero3_bf16_cpuoffload_all.json
index 09ca6785b2..5112c570b0 100644
--- a/deepspeed_configs/zero3_bf16_cpuoffload_all.json
+++ b/deepspeed_configs/zero3_bf16_cpuoffload_all.json
@@ -17,22 +17,13 @@
     "reduce_bucket_size": "auto",
     "stage3_prefetch_bucket_size": "auto",
     "stage3_param_persistence_threshold": "auto",
-    "stage3_max_live_parameters": 0,
-    "stage3_max_reuse_distance": 0,
-    "stage3_gather_16bit_weights_on_model_save": true
+    "max_live_parameters": 0,
+    "max_reuse_distance": 0,
+    "gather_16bit_weights_on_model_save": true
   },
   "bf16": {
     "enabled": true
   },
-  "fp16": {
-    "enabled": "auto",
-    "auto_cast": false,
-    "loss_scale": 0,
-    "initial_scale_power": 32,
-    "loss_scale_window": 1000,
-    "hysteresis": 2,
-    "min_loss_scale": 1
-  },
   "gradient_accumulation_steps": "auto",
   "gradient_clipping": "auto",
   "train_batch_size": "auto",
diff --git a/deepspeed_configs/zero3_bf16_cpuoffload_params.json b/deepspeed_configs/zero3_bf16_cpuoffload_params.json
index 41d4a21323..a2ac823410 100644
--- a/deepspeed_configs/zero3_bf16_cpuoffload_params.json
+++ b/deepspeed_configs/zero3_bf16_cpuoffload_params.json
@@ -13,22 +13,13 @@
     "reduce_bucket_size": "auto",
     "stage3_prefetch_bucket_size": "auto",
     "stage3_param_persistence_threshold": "auto",
-    "stage3_max_live_parameters": 0,
-    "stage3_max_reuse_distance": 0,
-    "stage3_gather_16bit_weights_on_model_save": true
+    "max_live_parameters": 0,
+    "max_reuse_distance": 0,
+    "gather_16bit_weights_on_model_save": true
   },
   "bf16": {
     "enabled": true
   },
-  "fp16": {
-    "enabled": "auto",
-    "auto_cast": false,
-    "loss_scale": 0,
-    "initial_scale_power": 32,
-    "loss_scale_window": 1000,
-    "hysteresis": 2,
-    "min_loss_scale": 1
-  },
   "gradient_accumulation_steps": "auto",
   "gradient_clipping": "auto",
   "train_batch_size": "auto",
diff --git a/devtools/dev_sharegpt.yml b/devtools/dev_chat_template.yml
similarity index 83%
rename from devtools/dev_sharegpt.yml
rename to devtools/dev_chat_template.yml
index 9c65b49dcd..32d5e56a0c 100644
--- a/devtools/dev_sharegpt.yml
+++ b/devtools/dev_chat_template.yml
@@ -1,4 +1,4 @@
-# Example config for debugging the sharegpt prompt format
+# Example config for debugging the chat_template prompt format
 base_model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
@@ -7,13 +7,13 @@ load_in_8bit: true
 load_in_4bit: false
 
 datasets:
-  - path: philschmid/guanaco-sharegpt-style
-    type: sharegpt
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
     shards: 10
 val_set_size: 0
 output_dir: temp_debug/axolotl_outputs/model
 dataset_prepared_path: temp_debug/axolotl_outputs/data
-dataset_processes: 1
+dataset_num_proc: 1
 
 sequence_len: 4096
 sample_packing: false
diff --git a/docker-compose.yaml b/docker-compose.yaml
index a16be726cf..4fbbef29e2 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -3,7 +3,7 @@ services:
   axolotl:
     build:
       context: .
-      dockerfile: ./docker/Dockerfile
+      dockerfile: ./docker/Dockerfile-uv
     volumes:
       - .:/workspace/axolotl
       - ~/.cache/huggingface/:/root/.cache/huggingface/
diff --git a/docker/Dockerfile b/docker/Dockerfile
deleted file mode 100644
index ae3fe89ae4..0000000000
--- a/docker/Dockerfile
+++ /dev/null
@@ -1,38 +0,0 @@
-ARG BASE_TAG=main-base
-FROM winglian/axolotl-base:$BASE_TAG
-
-ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
-ARG AXOLOTL_EXTRAS=""
-ARG AXOLOTL_ARGS=""
-ARG CUDA="118"
-ENV BNB_CUDA_VERSION=$CUDA
-ARG PYTORCH_VERSION="2.1.2"
-
-ENV PYTORCH_VERSION=$PYTORCH_VERSION
-
-RUN apt-get update && \
-    apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev
-
-WORKDIR /workspace
-
-RUN git clone --depth=1 https://github.com/OpenAccess-AI-Collective/axolotl.git
-
-WORKDIR /workspace/axolotl
-
-# If AXOLOTL_EXTRAS is set, append it in brackets
-RUN pip install causal_conv1d
-RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,galore,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
-    else \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,galore] $AXOLOTL_ARGS; \
-    fi
-
-# So we can test the Docker image
-RUN pip install pytest
-
-# fix so that git fetch/pull from remote works
-RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
-    git config --get remote.origin.fetch
-
-# helper for huggingface-login cli
-RUN git config --global credential.helper store
diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
deleted file mode 100644
index 1de5537dac..0000000000
--- a/docker/Dockerfile-base
+++ /dev/null
@@ -1,37 +0,0 @@
-ARG CUDA_VERSION="11.8.0"
-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.10"
-ARG PYTORCH_VERSION="2.1.2"
-ARG CUDA="118"
-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 && 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==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA
-
-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
diff --git a/docker/Dockerfile-cloud b/docker/Dockerfile-cloud
deleted file mode 100644
index 69ce143bb2..0000000000
--- a/docker/Dockerfile-cloud
+++ /dev/null
@@ -1,27 +0,0 @@
-ARG BASE_TAG=main
-FROM winglian/axolotl:$BASE_TAG
-
-ENV HF_DATASETS_CACHE="/workspace/data/huggingface-cache/datasets"
-ENV HUGGINGFACE_HUB_CACHE="/workspace/data/huggingface-cache/hub"
-ENV TRANSFORMERS_CACHE="/workspace/data/huggingface-cache/hub"
-ENV HF_HOME="/workspace/data/huggingface-cache/hub"
-ENV HF_HUB_ENABLE_HF_TRANSFER="1"
-
-EXPOSE 8888
-EXPOSE 22
-
-COPY scripts/cloud-entrypoint.sh /root/cloud-entrypoint.sh
-COPY scripts/motd /etc/motd
-
-RUN pip install jupyterlab notebook ipywidgets && \
-    jupyter lab clean
-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 && \
-    printf "[ ! -z \"\$TERM\" -a -r /etc/motd ] && cat /etc/motd\n" >> ~/.bashrc && \
-    chmod +x /workspace/axolotl/scripts/cloud-entrypoint.sh && \
-    chmod +x /root/cloud-entrypoint.sh
-
-ENTRYPOINT ["/root/cloud-entrypoint.sh"]
-CMD ["sleep", "infinity"]
diff --git a/docker/Dockerfile-cloud-no-tmux-uv b/docker/Dockerfile-cloud-no-tmux-uv
new file mode 100644
index 0000000000..0c096e751d
--- /dev/null
+++ b/docker/Dockerfile-cloud-no-tmux-uv
@@ -0,0 +1,31 @@
+# syntax=docker/dockerfile:1
+ARG BASE_TAG=main
+FROM axolotlai/axolotl-uv:$BASE_TAG
+
+ENV HF_DATASETS_CACHE="/workspace/data/huggingface-cache/datasets"
+ENV HF_HUB_CACHE="/workspace/data/huggingface-cache/hub"
+ENV HF_HOME="/workspace/data/huggingface-cache/hub"
+ENV HF_XET_HIGH_PERFORMANCE="1"
+
+EXPOSE 8888
+EXPOSE 22
+
+COPY scripts/cloud-entrypoint.sh /root/cloud-entrypoint.sh
+COPY scripts/motd /etc/motd
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install jupyterlab notebook ipywidgets && \
+    jupyter lab clean
+RUN apt update && \
+    apt install --yes --no-install-recommends openssh-server tmux iproute2 nvtop ibverbs-providers ibverbs-utils infiniband-diags librdmacm-dev librdmacm1 rdmacm-utils slurm-wlm && \
+    rm -rf /var/cache/apt/archives && \
+    rm -rf /var/lib/apt/lists/* && \
+    mkdir -p ~/.ssh && \
+    chmod 700 ~/.ssh && \
+    printf "[ ! -z \"\$TERM\" -a -r /etc/motd ] && cat /etc/motd\n" >> ~/.bashrc && \
+    printf "source /workspace/axolotl-venv/bin/activate\n" >> ~/.bashrc && \
+    chmod +x /workspace/axolotl/scripts/cloud-entrypoint.sh && \
+    chmod +x /root/cloud-entrypoint.sh
+
+ENTRYPOINT ["/root/cloud-entrypoint.sh"]
+CMD ["sleep", "infinity"]
diff --git a/docker/Dockerfile-cloud-uv b/docker/Dockerfile-cloud-uv
new file mode 100644
index 0000000000..26c324a910
--- /dev/null
+++ b/docker/Dockerfile-cloud-uv
@@ -0,0 +1,33 @@
+# syntax=docker/dockerfile:1
+ARG BASE_TAG=main
+FROM axolotlai/axolotl-uv:$BASE_TAG
+
+ENV HF_DATASETS_CACHE="/workspace/data/huggingface-cache/datasets"
+ENV HF_HUB_CACHE="/workspace/data/huggingface-cache/hub"
+ENV HF_HOME="/workspace/data/huggingface-cache"
+ENV HF_XET_HIGH_PERFORMANCE="1"
+
+EXPOSE 8888
+EXPOSE 22
+
+COPY scripts/cloud-entrypoint.sh /root/cloud-entrypoint.sh
+COPY scripts/motd /etc/motd
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install jupyterlab notebook ipywidgets && \
+    jupyter lab clean
+RUN apt update && \
+    apt install --yes --no-install-recommends openssh-server tmux iproute2 nvtop && \
+    rm -rf /var/cache/apt/archives && \
+    rm -rf /var/lib/apt/lists/* && \
+    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 && \
+    printf "[ ! -z \"\$TERM\" -a -r /etc/motd ] && cat /etc/motd\n" >> ~/.bashrc && \
+    printf "source /workspace/axolotl-venv/bin/activate\n" >> ~/.bashrc && \
+    chmod +x /workspace/axolotl/scripts/cloud-entrypoint.sh && \
+    chmod +x /root/cloud-entrypoint.sh && \
+    echo 'set-option -g history-limit 5000' >> ~/.tmux.conf
+
+ENTRYPOINT ["/root/cloud-entrypoint.sh"]
+CMD ["sleep", "infinity"]
diff --git a/docker/Dockerfile-tests b/docker/Dockerfile-tests
deleted file mode 100644
index d5c0595a87..0000000000
--- a/docker/Dockerfile-tests
+++ /dev/null
@@ -1,41 +0,0 @@
-ARG BASE_TAG=main-base
-FROM winglian/axolotl-base:$BASE_TAG
-
-ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
-ARG AXOLOTL_EXTRAS=""
-ARG AXOLOTL_ARGS=""
-ARG CUDA="118"
-ENV BNB_CUDA_VERSION=$CUDA
-ARG PYTORCH_VERSION="2.1.2"
-ARG GITHUB_REF="main"
-
-ENV PYTORCH_VERSION=$PYTORCH_VERSION
-
-RUN apt-get update && \
-    apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev
-
-WORKDIR /workspace
-
-RUN git clone --depth=1 https://github.com/OpenAccess-AI-Collective/axolotl.git
-
-WORKDIR /workspace/axolotl
-
-RUN git fetch origin +$GITHUB_REF && \
-    git checkout FETCH_HEAD
-
-# If AXOLOTL_EXTRAS is set, append it in brackets
-RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
-    else \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm] $AXOLOTL_ARGS; \
-    fi
-
-# So we can test the Docker image
-RUN pip install pytest
-
-# fix so that git fetch/pull from remote works
-RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
-    git config --get remote.origin.fetch
-
-# helper for huggingface-login cli
-RUN git config --global credential.helper store
diff --git a/docker/Dockerfile-uv b/docker/Dockerfile-uv
new file mode 100644
index 0000000000..f7957a6cb1
--- /dev/null
+++ b/docker/Dockerfile-uv
@@ -0,0 +1,45 @@
+# syntax=docker/dockerfile:1
+ARG BASE_TAG=main-base
+FROM axolotlai/axolotl-base-uv:$BASE_TAG
+
+ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
+ARG AXOLOTL_EXTRAS=""
+ARG AXOLOTL_ARGS=""
+ARG CUDA="118"
+ARG PYTORCH_VERSION="2.1.2"
+ARG TARGETARCH
+
+ENV PYTORCH_VERSION=$PYTORCH_VERSION
+
+WORKDIR /workspace/axolotl
+
+COPY . .
+
+# If AXOLOTL_EXTRAS is set, append it in brackets; don't install deepspeed with arm64
+RUN uv pip uninstall causal_conv1d
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install "huggingface_hub>=1.5.0" "httpx<1" && \
+    if [ "$TARGETARCH" = "arm64" ]; then \
+        BASE_EXTRAS="optimizers,ray"; \
+    else \
+        BASE_EXTRAS="deepspeed,optimizers,ray"; \
+    fi && \
+    if [ "$AXOLOTL_EXTRAS" != "" ]; then \
+        uv pip install --no-build-isolation -e .[$BASE_EXTRAS,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+    else \
+        uv pip install --no-build-isolation -e .[$BASE_EXTRAS] $AXOLOTL_ARGS; \
+    fi && \
+    python scripts/cutcrossentropy_install.py --uv | sh && \
+    uv pip install pytest
+
+# fix so that git fetch/pull from remote works with shallow clone
+RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
+    git config --get remote.origin.fetch && \
+    git config --global credential.helper store
+
+COPY .axolotl-complete.bash /root/.axolotl-complete.bash
+RUN chmod +x /root/.axolotl-complete.bash && \
+    chmod +x /workspace/axolotl/scripts/uv-entrypoint.sh && \
+    echo 'source /root/.axolotl-complete.bash' >> ~/.bashrc
+
+ENTRYPOINT ["/workspace/axolotl/scripts/uv-entrypoint.sh"]
diff --git a/docker/Dockerfile-uv-base b/docker/Dockerfile-uv-base
new file mode 100644
index 0000000000..95d4d7540b
--- /dev/null
+++ b/docker/Dockerfile-uv-base
@@ -0,0 +1,51 @@
+# syntax=docker/dockerfile:1
+ARG CUDA_VERSION="12.6.3"
+ARG CUDNN_VERSION=""
+ARG UBUNTU_VERSION="22.04"
+ARG MAX_JOBS=4
+ARG TARGETARCH
+
+FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
+
+ARG TARGETARCH
+ARG PYTHON_VERSION="3.11"
+ARG PYTORCH_VERSION="2.6.0"
+ARG CUDA="126"
+ARG TORCH_BACKEND="cu126"
+ARG TORCH_INDEX_URL=""
+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 curl \
+    && apt-get install -y --allow-change-held-packages vim curl nano zstd libnccl2 libnccl-dev ibverbs-providers ibverbs-utils infiniband-diags librdmacm-dev librdmacm1 rdmacm-utils slurm-wlm rsync s3fs \
+    && rm -rf /var/cache/apt/archives \
+    && rm -rf /var/lib/apt/lists/* \
+    && git lfs install --skip-repo \
+    && curl -LsSf https://astral.sh/uv/install.sh | sh
+
+ENV PATH="/root/.local/bin:${PATH}"
+
+RUN uv python install ${PYTHON_VERSION}
+
+WORKDIR /workspace
+
+RUN uv venv --no-project --relocatable axolotl-venv
+
+ENV PATH="/workspace/axolotl-venv/bin:${PATH}"
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install packaging setuptools wheel psutil \
+    && if [ -n "$TORCH_INDEX_URL" ]; then \
+        uv pip install --index-url "$TORCH_INDEX_URL" torch==${PYTORCH_VERSION} torchvision; \
+    else \
+        UV_TORCH_BACKEND=$TORCH_BACKEND uv pip install torch==${PYTORCH_VERSION} torchvision; \
+    fi \
+    && uv pip install awscli pydantic
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ "$TARGETARCH" = "amd64" ]; then \
+        MAMBA_SKIP_CUDA_BUILD=TRUE CAUSAL_CONV1D_SKIP_CUDA_BUILD=TRUE uv pip install --no-build-isolation mamba_ssm causal_conv1d; \
+    fi
diff --git a/docs/.gitignore b/docs/.gitignore
index 4c23a061fa..94e5b88fb6 100644
--- a/docs/.gitignore
+++ b/docs/.gitignore
@@ -1,2 +1,7 @@
 /.quarto/
 _site/
+/api/*.qmd
+/api/*.html
+config-reference.qmd
+models/**/*.qmd
+models/**/*.html
diff --git a/docs/1_58bit_finetuning.qmd b/docs/1_58bit_finetuning.qmd
new file mode 100644
index 0000000000..02bc3a6f18
--- /dev/null
+++ b/docs/1_58bit_finetuning.qmd
@@ -0,0 +1,70 @@
+---
+title: "1.58-bit Finetuning"
+back-to-top-navigation: true
+toc: true
+toc-expand: 2
+toc-depth: 4
+---
+
+## Overview
+
+1.58-bit finetuning allows you to finetune BitNet models when their prequantized weights are provided. In theory, it will be possible to fine-tune any LLM in 1.58bit format but the performance degradation will be dramatic.
+
+Axolotl supports 1.58-bit finetuning via the [`onebitllms`](https://github.com/tiiuae/onebitllms) library, which replaces standard linear layers with BitNet-compatible counterparts ready to use for training.
+
+::: {.callout-note}
+LoRA is not supported for BitNet models
+:::
+
+## Installation
+
+Install the `onebitllms` package before using this feature:
+
+```bash
+uv pip install onebitllms
+```
+
+Or from source:
+
+```bash
+uv pip install git+https://github.com/tiiuae/onebitllms
+```
+
+## Supported models
+
+For now, only `Falcon-E` series of models are supported. Make sure to use their `-prequantized` version:
+
+```bash
+tiiuae/Falcon-E-3B-Base-prequantized
+tiiuae/Falcon-E-1B-Base-prequantized
+```
+
+In theory, any other model would 'work' but the performance degradation will be huge. This remains an area of exploration.
+
+## Configuration
+
+To enable 1.58-bit finetuning, set the following in your configuration file:
+
+```yaml
+base_model: tiiuae/Falcon-E-3B-Base-prequantized  # A BitNet-compatible model
+
+use_onebitllms: true
+```
+
+::: {.callout-note}
+For BitNet models, it is recommended to use a higher learning rate than classic models (usually in the order of magnitude of 10x).
+:::
+
+## Considerations after training
+
+Once your model has been trained with 1.58bit fine-tuning, you can convert the trained model in ternary format using the `onebitllms` CLI:
+
+```bash
+onebitllms quantize_to_1bit INPUT_PATH OUTPUT_PATH
+```
+
+After that, you can use supported packages such as `llama.cpp` or Apple MLX package to run the trained model.
+
+## Example Configuration
+
+You can find example configurations in `examples/falcon-e` which contain one configuration for SFT and one configuration for DPO.
diff --git a/docs/agents/grpo.md b/docs/agents/grpo.md
new file mode 100644
index 0000000000..0d7ba3f41e
--- /dev/null
+++ b/docs/agents/grpo.md
@@ -0,0 +1,71 @@
+# GRPO — Agent Reference
+
+Online RL with verifiable reward functions. For full config reference, async features, and scaling, see [grpo.qmd](../grpo.qmd). For vLLM setup, see [vllm_serving.qmd](../vllm_serving.qmd).
+
+## Architecture
+
+```
+Terminal 1 (GPU 0)                    Terminal 2 (GPU 1)
+┌──────────────────────┐              ┌──────────────────────────────────┐
+│  vLLM Server         │   HTTP       │  Trainer                         │
+│  Serves base model   │◄────────────►│  1. Send prompts to vLLM         │
+│  + LoRA adapter      │  /generate   │  2. Score completions (rewards)  │
+│                      │  /set_lora   │  3. Compute advantages           │
+│  Punica kernels for  │              │  4. PPO-clip gradient update     │
+│  LoRA inference      │              │  5. Sync LoRA weights to vLLM    │
+└──────────────────────┘              └──────────────────────────────────┘
+```
+
+## Components Required
+
+1. A YAML config with `rl: grpo`
+2. A reward module (Python file with reward functions)
+3. A running vLLM server (`axolotl vllm-serve config.yaml`)
+
+## Reward Function Signature
+
+```python
+def my_reward(completions, **kwargs) -> list[float]:
+    # completions[i][0]["content"] = text of i-th completion
+    # **kwargs contains dataset columns not removed by transform
+    return [score_for_each_completion]
+```
+
+Multiple rewards: `reward_funcs: [r1, r2]` with `reward_weights: [1.0, 0.5]`.
+
+## Key Async Features
+
+| Feature | Config | Purpose |
+|---------|--------|---------|
+| Async prefetch | `async_prefetch: true` | Overlap generation with training |
+| LoRA sync | `vllm_lora_sync: true` | Fast adapter sync via filesystem |
+| Streaming scoring | `streaming_partial_batch: true` | Score one group at a time |
+| Zero-adv skip | `skip_zero_advantage_batches: true` | Skip batches with no learning signal |
+| Replay buffer | `replay_buffer_size: 100` | Cache high-signal groups |
+| IS correction | `vllm_importance_sampling_correction: true` | Fix off-policy distribution shift |
+
+## Health Checks
+
+- `rewards/*/mean` > 0.15 within 20 steps (else: test reward function standalone)
+- `reward_std` > 0 on most steps (else: no learning signal)
+- `entropy` 0.05-0.5 (< 0.01 = mode collapse)
+- `grad_norm` 0.001-1.0 (> 10 = unstable, 0.0 = zero-advantage skip)
+
+See [training_stability.qmd](../training_stability.qmd) for detailed diagnostics.
+
+## File Map
+
+```
+src/axolotl/
+  cli/train.py                     # Entry point
+  cli/vllm_serve.py                # Entry point for vLLM server
+  core/trainers/grpo/
+    trainer.py                     # AxolotlGRPOTrainer
+    sampler.py                     # Sampling utilities
+  core/builders/rl.py              # HFRLTrainerBuilder — routes rl type → trainer
+  scripts/vllm_serve_lora.py       # vLLM serve script with LoRA sync support
+  utils/schemas/trl.py             # TRL config schema (all trl: options)
+
+docs/grpo.qmd                     # Full user docs: async, rewards, scaling, config reference
+docs/vllm_serving.qmd             # vLLM server modes, LoRA sync, weight sync
+```
diff --git a/docs/agents/model_architectures.md b/docs/agents/model_architectures.md
new file mode 100644
index 0000000000..500664ec89
--- /dev/null
+++ b/docs/agents/model_architectures.md
@@ -0,0 +1,200 @@
+# Model Architectures — Agent Reference
+
+Model-specific quirks, required settings, and known issues. Check this before debugging training failures on specific model families.
+
+## VLM (Vision Language Model) Quick Start
+
+All VLM configs require these four lines:
+```yaml
+processor_type: AutoProcessor
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+```
+
+Decision tree for VLM config:
+```text
+Is the model multimodal (has vision/audio encoder)?
+  ├─ YES: Add `freeze_mm_modules: true` if training text only
+  │       Add `chat_template: <model_template>` (e.g. gemma4, qwen3_5, gemma3)
+  │       LoRA: use regex `lora_target_modules` to restrict to language model
+  └─ NO: Train as a regular text model
+
+Is the model MoE (e.g. Gemma4 26B-A4B, Qwen3.5 35B-A3B)?
+  ├─ YES: Add `lora_target_parameters` for expert LoRA
+  │       Consider ScatterMoE kernels (see Plugins section)
+  └─ NO: Standard LoRA config
+```
+
+## Plugins & Optimizations
+
+### Cut Cross Entropy (CCE)
+
+Computes loss from hidden states + lm_head weight without materializing the full logits tensor, saving significant VRAM. Install if not already present:
+
+```bash
+python scripts/cutcrossentropy_install.py | sh
+```
+
+See [Cut Cross Entropy](https://docs.axolotl.ai/docs/optimizations.html#cut-cross-entropy-cce) for the pinned install command and the full list of supported models.
+
+```yaml
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+```
+
+### ScatterMoE Kernels
+
+Fuses expert + LoRA computation into a single kernel for MoE models. Significant speedup for models with many experts.
+
+```yaml
+plugins:
+  - axolotl.integrations.kernels.KernelsPlugin
+use_kernels: true
+use_scattermoe: true
+experts_implementation: scattermoe
+
+# Expert LoRA targets (3D parameter tensors, not nn.Linear):
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+```
+
+Supported: Gemma4 (`gemma4_text`), Mixtral, Qwen MoE variants. The plugin auto-detects model type and routing function. Without ScatterMoE, expert LoRA still works but runs base expert matmul and LoRA as separate operations.
+
+## Gemma 4
+
+**Models**: `google/gemma-4-26B-A4B` (MoE), `google/gemma-4-31B` (dense), `google/gemma-4-E2B`, `google/gemma-4-E4B`
+
+**Architecture**: Multimodal wrapper (`Gemma4ForConditionalGeneration`) over a text backbone (`Gemma4TextModel`), with optional vision/audio encoders. All Gemma4 HF repos have `model_type: "gemma4"` — even text-only variants load as multimodal with a vision tower.
+
+### Required settings
+
+```yaml
+# Always needed for Gemma4:
+freeze_mm_modules: true          # Freeze vision/audio encoders for text-only training
+gradient_checkpointing_kwargs:
+  use_reentrant: false           # Shared per-layer norms cause "marked ready twice" with reentrant
+
+# LoRA target — restrict to language model only (DO NOT use lora_target_linear: true):
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+```
+
+### Auto-detection
+
+Axolotl auto-detects Gemma4 and applies:
+- `use_reentrant: false` for gradient checkpointing
+- `ddp_find_unused_parameters: true` for DDP (skipped when `activation_offloading: true`)
+
+### Multi-GPU
+
+| Strategy | Works? | Notes |
+|----------|--------|-------|
+| DDP | Yes | Auto-sets `ddp_find_unused_parameters=True` |
+| DDP + activation_offloading | Yes | `find_unused_parameters` is skipped (conflicts with checkpoint wrappers) |
+| FSDP1 | No | OOM during dequantization/sharding with QLoRA |
+| FSDP2 | Yes | Use `Gemma4TextDecoderLayer` (not `Gemma4DecoderLayer`) as wrap class |
+| FSDP2 + activation_offloading | Yes | Lowest VRAM (~26 GiB/GPU for 26B-A4B) |
+
+FSDP2 config:
+```yaml
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_version: 2
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer
+```
+
+### MoE (26B-A4B)
+
+- `enable_moe_block: true`, 256 experts, top-k routing
+- No separate `SparseMoeBlock` — MoE is embedded in each decoder layer
+- Expert LoRA targets 3D parameter tensors:
+  ```yaml
+  lora_target_parameters:
+    - experts.gate_up_proj
+    - experts.down_proj
+  ```
+- ScatterMoE kernel acceleration:
+  ```yaml
+  plugins:
+    - axolotl.integrations.kernels.KernelsPlugin
+  use_kernels: true
+  use_scattermoe: true
+  experts_implementation: scattermoe
+  ```
+
+### VLM (Vision) Training
+
+All Gemma4 models load as `Gemma4ForConditionalGeneration` with a vision tower. No custom `ProcessingStrategy` needed — the base class auto-detects the image token.
+
+```yaml
+base_model: google/gemma-4-E2B-it   # or E4B-it, 26B-A4B
+processor_type: AutoProcessor
+freeze_mm_modules: true
+chat_template: gemma4
+
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+```
+
+A starting VLM loss of ~8-15 is typical. In most runs, loss converges below 1.0 within ~30-50 steps, though results may vary across configurations.
+
+For the 26B-A4B MoE variant with ScatterMoE + expert LoRA + CCE, add:
+```yaml
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+use_kernels: true
+use_scattermoe: true
+experts_implementation: scattermoe
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+```
+
+### Common issues
+
+| Symptom | Cause | Fix |
+|---------|-------|-----|
+| `mm_token_type_ids is required` in DDP | `model.config` not accessible through DDP wrapper | Already fixed — `unwrap_model()` in `compute_loss` and `prediction_step` |
+| `marked a variable ready twice` in DDP | `ddp_find_unused_parameters=True` + activation_offloading checkpoint wrappers | Auto-handled — `find_unused_parameters` is skipped when `activation_offloading: true` |
+| Loss ~12 instead of ~0.5 | Using `lora_target_linear: true` (applies LoRA to vision/audio modules) | Use the regex `lora_target_modules` pattern instead |
+| FSDP2 `Could not find Gemma4AudioLayer` | Auto-wrap detects `_no_split_modules` including audio layers that don't exist | Explicitly set `fsdp_transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer` |
+| `Gemma4ClippableLinear not supported` by PEFT | Vision tower uses a non-standard linear wrapper | Axolotl patches this automatically via `_patch_peft_clippable_linear()` |
+
+### E2B/E4B dense models
+
+These have `hidden_size_per_layer_input: 256` (per-layer input embeddings) and `attention_k_eq_v: False`. Known issue: loss starts higher than expected (~12 vs ~0.5 for 26B). Root cause under investigation — may be related to the per-layer input mechanism or the `Gemma4ForConditionalGeneration` loss computation.
+
+## Gemma 3
+
+**Models**: `google/gemma-3-*`
+
+- `ddp_find_unused_parameters: true` needed (multimodal unused params)
+- `use_reentrant: false` recommended
+- Attention mask must be dropped for sample packing (handled automatically)
+- Multi-GPU test currently skipped (`tests/e2e/multigpu/test_gemma3.py`)
+
+## Qwen 3.5 MoE
+
+**Models**: `Qwen/Qwen3.5-35B-A3B`
+
+- Hybrid architecture: DeltaNet linear attention (30 layers) + full attention (10 layers)
+- 256 experts, 8 active per token
+- Known weight scale drift in late DeltaNet layers (36-38) due to AdamW + rare expert interaction
+- Fix: `normalize_weight_scales` config to detect and rescale outliers:
+  ```yaml
+  normalize_weight_scales:
+    - name_pattern: 'linear_attn\.conv1d\.weight'
+      threshold: 1.3
+  ```
+
+## General MoE Notes
+
+- `lora_target_linear: true` with multimodal MoE models will apply LoRA to ALL linear modules including vision/audio encoders — use regex `lora_target_modules` to restrict to language model only
+- Rare experts get larger effective learning rate from AdamW (small second-moment estimates) — can cause weight drift in recurrent/SSM components. Use `normalize_weight_scales` with `dry_run: true` to detect.
+- For ScatterMoE kernel support, set `experts_implementation: scattermoe` and add the KernelsPlugin
diff --git a/docs/agents/new_model_support.md b/docs/agents/new_model_support.md
new file mode 100644
index 0000000000..bc42ada861
--- /dev/null
+++ b/docs/agents/new_model_support.md
@@ -0,0 +1,181 @@
+# New Model Support — Agent Reference
+
+Guide for debugging and adding support for new model architectures in axolotl. Based on lessons learned from Gemma4, Gemma3, Qwen2-VL, and other multimodal/MoE models.
+
+## Quick Validation Checklist
+
+When testing a new model, run through these checks in order:
+
+1. **Does the model load?** `axolotl preprocess config.yaml` — catches config schema errors
+2. **Does LoRA apply?** Check for "Unsupported layer type" warnings from PEFT
+3. **Is the initial loss sane?** First-step loss for a pretrained model should be 0.5–2.0 for SFT
+4. **Does sample packing work?** Compare loss with `sample_packing: true` vs `false` — should be similar
+5. **Is CCE active?** Check for "Applying Cut Cross Entropy" log and verify peak VRAM is lower
+
+## Loss Debugging
+
+### Expected initial loss
+A pretrained model doing SFT should start with loss roughly in the 0.5–2.0 range. If loss starts above 3.0, something is wrong. If it's near `log(vocab_size)` (≈ 12 for 262K vocab), the model is predicting at random — attention masking or model weights are broken.
+
+### Direct comparison technique
+The fastest way to isolate a loss issue — bypass the trainer entirely:
+
+```python
+# Load model via axolotl's pipeline (applies all patches)
+from axolotl.cli.config import load_cfg
+from axolotl.utils.config import normalize_config, prepare_plugins
+from axolotl.loaders.tokenizer import load_tokenizer
+from axolotl.loaders.model import ModelLoader
+
+cfg = load_cfg("your_config.yaml")
+normalize_config(cfg)
+prepare_plugins(cfg)
+tokenizer = load_tokenizer(cfg)
+model, _ = ModelLoader(cfg, tokenizer).load()
+
+# Forward pass on preprocessed data
+model.train()
+out = model(input_ids, labels=labels)
+print(f"Direct loss: {out.loss.item()}")  # Compare to trainer's reported loss
+```
+
+If direct loss is correct (~1.0) but trainer reports 3–4x higher, check `model_accepts_loss_kwargs` (see below).
+
+### `model_accepts_loss_kwargs` inflation
+HF Trainer checks if the model's `forward()` has `**kwargs` and sets `model_accepts_loss_kwargs=True`. This changes loss normalization: the trainer does NOT divide loss by `gradient_accumulation_steps` before logging. The gradient is correct — only the logged loss is inflated.
+
+**Symptom**: Logged loss ≈ actual_loss × gradient_accumulation_steps.
+
+**Which models are affected**: Any model with `**kwargs` in forward (common in multimodal models for extra inputs like `mm_token_type_ids`, `pixel_values`, etc.).
+
+**Fix location**: `src/axolotl/core/trainers/base.py` `__init__()` — after `super().__init__()`, check if the unwrapped model actually has `num_items_in_batch` in its forward signature. If not, set `self.model_accepts_loss_kwargs = False`.
+
+## Multimodal Models (ForConditionalGeneration)
+
+Many recent models use `ForConditionalGeneration` as the top-level class, not `ForCausalLM`:
+- Gemma3 → `Gemma3ForConditionalGeneration`
+- Gemma4 → `Gemma4ForConditionalGeneration`
+- Qwen2-VL → `Qwen2VLForConditionalGeneration`
+- LLaVA → `LlavaForConditionalGeneration`
+
+### Why this matters
+
+| Component | Targets `ForCausalLM` | Needs `ForConditionalGeneration` |
+|-----------|----------------------|--------------------------------|
+| CCE patches | ✅ (default) | ❌ silently inactive if not patched |
+| PEFT LoRA | ✅ | May fail on custom layer types |
+| HF Trainer label handling | ✅ | May need extra inputs |
+
+### Required extra inputs
+Multimodal models require special inputs during training even for text-only data:
+
+| Model | Required Input | Value for Text-Only |
+|-------|---------------|-------------------|
+| Gemma4 | `mm_token_type_ids` | `torch.zeros_like(input_ids)` |
+| Gemma3 | `token_type_ids` | `torch.zeros_like(input_ids)` |
+
+Auto-inject in `compute_loss()` when not provided by the data collator. See `core/trainers/base.py`.
+
+### Custom layer types and PEFT
+Vision towers often use custom module wrappers that PEFT doesn't support:
+
+| Model | Custom Layer | Wraps | Fix |
+|-------|-------------|-------|-----|
+| Gemma4 | `Gemma4ClippableLinear` | `nn.Linear` | Redirect to `.linear` child |
+
+Fix location: `src/axolotl/loaders/adapter.py` `_patch_peft_clippable_linear()`.
+
+## Sample Packing
+
+### How packed sequence detection works (transformers ≥ 5.x)
+`transformers.masking_utils._preprocess_mask_arguments()` detects packed sequences from `position_ids` resets. But **only when `attention_mask is None`**:
+
+```python
+# From masking_utils.py:
+if position_ids is not None and attention_mask is None and past_key_values is None:
+    packed_sequence_mask = find_packed_sequence_indices(position_ids)
+```
+
+If the collator provides an all-ones `attention_mask`, packing detection is **skipped** and the model builds a single causal mask spanning all packed sequences → cross-sequence attention leakage → very high loss.
+
+### Fix for models using `create_causal_mask_mapping`
+For Gemma3, Gemma4, and similar models that use the new transformers masking system, remove `attention_mask` from inputs when sample packing is active:
+
+```python
+# In compute_loss():
+if (
+    self.args.sample_packing
+    and model_type in ("gemma4", "gemma3")
+    and "attention_mask" in inputs
+    and "position_ids" in inputs
+):
+    del inputs["attention_mask"]
+```
+
+Fix location: `src/axolotl/core/trainers/base.py` `compute_loss()`.
+
+### Models that DON'T need this fix
+Older models that use `_prepare_4d_causal_attention_mask` (Llama, Mistral, Qwen2, etc.) handle sample packing via axolotl's multipack attention monkeypatch instead. Only models using the new `create_causal_mask_mapping` / `create_causal_mask` masking system need the `attention_mask` removal.
+
+## Attention Backend Selection
+
+| Backend | Config | head_dim limit | torch_compile | Notes |
+|---------|--------|---------------|---------------|-------|
+| FA2 | `attn_implementation: flash_attention_2` | 256 | ✅ | Fastest when supported |
+| FA4 | auto with `attn_implementation: flash_attention_2` | 256 (SM90+) | ✅ | Auto-detected on H100+ |
+| SDPA | `attn_implementation: sdpa` | None | ✅ | Universal fallback |
+| flex | `attn_implementation: flex_attention` | None | ⚠️ Triton OOM for large head_dim | Good for variable head dims |
+| eager | `attn_implementation: eager` | None | ✅ | Slowest, always works |
+
+**Check model support**: Look at `_supports_flash_attn_2`, `_supports_flex_attn`, `_supports_sdpa` attributes on the model class.
+
+**head_dim gotcha**: The 256 limit is specific to flash-attn CUDA kernels, NOT PyTorch-level. SDPA and flex_attention both handle arbitrary head_dim. Models with `global_head_dim > 256` (Gemma4: 512) must use SDPA or flex.
+
+**flex + compile gotcha**: `torch_compile` with flex_attention can hit Triton shared memory OOM for large head_dim. Falls back to eager per-function (not a crash, but slower). Unsloth disables flex for Gemma4 for this reason.
+
+## Cut Cross Entropy (CCE)
+
+### How CCE patches work
+CCE replaces the model's `forward()` with a fused version that computes loss from hidden states + lm_head weight without materializing the full logits tensor. This saves ~`batch × seq_len × vocab_size × dtype_bytes` of VRAM.
+
+### Adding CCE for a new model
+1. Check if the model type is in `cut_cross_entropy.transformers.patch.PATCH_FNS`
+2. If not, axolotl's generic fallback (`integrations/cut_cross_entropy/__init__.py` `patch_llama_like()`) patches `{Prefix}ForCausalLM.forward` with `cce_forward`
+3. For multimodal models (`ForConditionalGeneration`), a model-specific patch is needed in `ml-cross-entropy` repo
+4. The multimodal `cce_forward` must accept all extra kwargs (pixel_values, mm_token_type_ids, etc.) and pop any that would conflict before calling `self.model()`
+
+### Common CCE pitfall
+If CCE appears active (log says "Applying Cut Cross Entropy") but peak VRAM doesn't decrease, check which class was patched. If the model loads as `ForConditionalGeneration` but CCE patched `ForCausalLM`, the patch is silently inactive.
+
+## MoE Models
+
+### Dense MLP vs MoE experts
+Some MoE models (e.g., Gemma4) have BOTH dense MLP layers and MoE expert layers at every decoder layer:
+- `gate_proj/up_proj/down_proj` → targets the **dense MLP** (`Gemma4TextMLP`)
+- `experts.gate_up_proj/experts.down_proj` → targets the **MoE experts** (`Gemma4TextExperts`)
+
+LoRA on the dense MLP works normally. Expert LoRA via `lora_target_parameters` requires PEFT support for the specific expert module type (may warn "Unsupported layer type").
+
+### ScatterMoE kernels
+`use_scattermoe: true` with `experts_implementation: scattermoe` registers fused expert kernels via transformers' `ExpertsInterface`. Significant speedup for MoE models. Requires the kernels plugin:
+```yaml
+plugins:
+  - axolotl.integrations.kernels.KernelsPlugin
+use_kernels: true
+use_scattermoe: true
+experts_implementation: scattermoe
+```
+
+## Where to Add Model-Specific Fixes
+
+| What | Where | Example |
+|------|-------|---------|
+| Missing forward inputs | `core/trainers/base.py` `compute_loss()` | mm_token_type_ids injection |
+| Attention mask fixes | `core/trainers/base.py` `compute_loss()` | Sample packing mask removal |
+| Loss logging fixes | `core/trainers/base.py` `__init__()` | model_accepts_loss_kwargs override |
+| PEFT/LoRA patches | `loaders/adapter.py` | ClippableLinear redirect |
+| Attention patches | `monkeypatch/attention/` | FA4 tuple fix |
+| Model-specific patches | `loaders/patch_manager.py` `_apply_model_specific_patches()` | Llama4, Kimi, NemotronH |
+| CCE patches | `ml-cross-entropy` repo `transformers/` | Per-model cce_forward |
+| Example configs | `examples/<model>/` | Validated YAML |
+| Config validation | `utils/schemas/validation.py` | Compatibility checks |
diff --git a/docs/agents/preference_tuning.md b/docs/agents/preference_tuning.md
new file mode 100644
index 0000000000..3414d22ceb
--- /dev/null
+++ b/docs/agents/preference_tuning.md
@@ -0,0 +1,121 @@
+# Preference Learning (RLHF) — Agent Reference
+
+Reference for DPO, IPO, KTO, ORPO, and SimPO. For config templates and dataset format examples, see [rlhf.qmd](../rlhf.qmd). For GRPO, see [grpo.qmd](../grpo.qmd). For EBFT, see [ebft.qmd](../ebft.qmd).
+
+## Method Overview
+
+| Method | Data Requirement | Key Idea | Best For |
+|--------|-----------------|----------|----------|
+| **DPO** | Paired (chosen + rejected) | Implicit reward via preference pairs | General alignment, most common |
+| **IPO** | Paired (chosen + rejected) | DPO with different loss (avoids overfitting) | When DPO overfits |
+| **KTO** | Unpaired (completion + binary label) | Kahneman-Tversky loss, no pairs needed | When you only have thumbs-up/down |
+| **ORPO** | Paired (chosen + rejected) | Combined SFT + preference, no ref model | Single-stage alignment, saves VRAM |
+| **SimPO** | Paired (chosen + rejected) | Length-normalized, no ref model | Simple setup, length-robust |
+
+Default: start with DPO. All methods require `sample_packing: false`.
+
+## Architecture
+
+```
+┌──────────────┐   ┌───────────────┐   ┌───────────────┐
+│ Policy Model │   │ Reference     │   │ Preference    │
+│ (trainable)  │   │ Model (frozen)│   │ Dataset       │
+└──────┬───────┘   └──────┬────────┘   └──────┬────────┘
+       └──────────┬───────┘                    │
+                  v                            │
+       Forward pass on chosen + rejected <─────┘
+                  │
+       Preference Loss (DPO/IPO/KTO/...)
+                  │
+       Backprop + Update
+
+Exception: ORPO and SimPO do NOT use a reference model (~50% less VRAM).
+```
+
+No vLLM server needed (unlike GRPO). Offline RL with pre-collected preference data.
+
+## Method Selection
+
+1. Paired preference data (chosen + rejected)?
+   - Default → `rl: dpo`
+   - Overfitting → `rl: dpo, dpo_loss_type: ["ipo"]`
+   - VRAM-limited → `rl: orpo` (no ref model)
+   - Length-sensitive → `rl: simpo` (no ref model)
+2. Only binary labels (good/bad)? → `rl: kto`
+3. Single-stage training (no separate SFT)? → `rl: orpo`
+
+| | DPO | IPO | KTO | ORPO | SimPO |
+|---|---|---|---|---|---|
+| **Reference model** | Yes | Yes | Yes | No | No |
+| **VRAM overhead** | ~2x model | ~2x model | ~2x model | ~1x model | ~1x model |
+| **TRL trainer class** | DPOTrainer | DPOTrainer | KTOTrainer | ORPOTrainer | CPOTrainer |
+
+## Prompt Strategy Resolution
+
+The `type` field resolves to a Python function:
+
+```
+type: "chatml.intel"
+  → axolotl.prompt_strategies.dpo.chatml.intel(cfg, **kwargs)
+  → returns transform_fn(sample) → {"prompt", "chosen", "rejected"}
+
+type: "chat_template.default"
+  → axolotl.prompt_strategies.dpo.chat_template.default(cfg, dataset_idx, **kwargs)
+
+type: {"field_prompt": "prompt", ...}   (dict)
+  → axolotl.prompt_strategies.dpo.user_defined.default(...)
+```
+
+Module base: `axolotl.prompt_strategies.{rl_method}` — replace `dpo` with `kto` or `orpo`.
+
+## Healthy Training Indicators
+
+| Metric | Healthy Range | Problem |
+|--------|--------------|---------|
+| `train/loss` | Decreasing, 0.3-0.7 | Flat or increasing = broken data or too high LR |
+| `rewards/chosen` | Increasing | Flat = model not learning preferences |
+| `rewards/rejected` | Decreasing | Increasing = model prefers wrong responses |
+| `rewards/margins` | Positive and increasing | Negative = prefers rejected over chosen |
+| `rewards/accuracies` | > 0.5, toward 0.7+ | < 0.5 = worse than random |
+| `logps/rejected` | Decreasing | Increasing = reward hacking |
+| `grad_norm` | 0.01 - 10.0 | > 100 = exploding gradients |
+
+Method-specific: DPO/IPO watch `rewards/margins`; KTO loss is noisier; ORPO monitor SFT + odds ratio components; SimPO check length-normalized reward separation.
+
+## Known Issues
+
+| Issue | Fix |
+|-------|-----|
+| Sample packing crash | Set `sample_packing: false` (required for all preference methods) |
+| KTO `KeyError: 'label'` | Ensure dataset has boolean `label` column |
+| ORPO/KTO `KeyError` during tokenization | Add `remove_unused_columns: false` |
+| ORPO template not applied | ORPO requires explicit `chat_template` setting |
+| OOM with ref model (DPO/IPO/KTO) | Use LoRA/QLoRA, or switch to ORPO/SimPO (no ref model) |
+| IPO + label_smoothing | Do not set `dpo_label_smoothing` when `rl: ipo` |
+
+Full troubleshooting: [training_stability.qmd](../training_stability.qmd)
+
+## File Map
+
+```
+src/axolotl/
+  core/trainers/dpo/              # DPO trainer, args, strategy
+  core/builders/rl.py             # HFRLTrainerBuilder — routes rl type → trainer class
+  core/training_args.py           # AxolotlKTOConfig, AxolotlORPOConfig, AxolotlCPOConfig
+  prompt_strategies/
+    dpo/                          # DPO/IPO/SimPO dataset strategies
+      chat_template.py            # chat_template.default, chat_template.argilla_chat
+      chatml.py                   # chatml.default/intel/icr/argilla_chat/prompt_pairs/ultra
+      llama3.py                   # llama3 variants (same subtypes as chatml)
+      user_defined.py             # Custom field mapping
+      passthrough.py              # No transform
+    kto/                          # KTO dataset strategies (chatml, llama3, user_defined)
+    orpo/                         # ORPO dataset strategies (chat_template.argilla)
+  utils/schemas/enums.py          # RLType enum (dpo, ipo, kto, orpo, simpo, grpo, gdpo, ebft)
+  utils/schemas/config.py         # All rl/dpo/kto/orpo/simpo config fields
+
+docs/rlhf.qmd                    # Full user docs: all dataset formats, config templates
+docs/choosing_method.qmd          # SFT vs DPO vs GRPO decision guide
+examples/qwen2/dpo.yaml           # DPO example
+examples/llama-3/qlora-1b-kto.yaml  # KTO example
+```
diff --git a/docs/agents/pretraining.md b/docs/agents/pretraining.md
new file mode 100644
index 0000000000..e6291e2ed8
--- /dev/null
+++ b/docs/agents/pretraining.md
@@ -0,0 +1,75 @@
+# Pretraining / Continual Pretraining — Agent Reference
+
+Train on raw text with no input masking. Two approaches depending on dataset size.
+
+## When to Use
+
+- Continual pretraining on domain-specific corpora
+- Adapting a base model to a new language or domain before fine-tuning
+- Pretraining-style data where the entire text is the training signal
+
+## Choosing an Approach
+
+| | Non-streaming (`type: completion`) | Streaming (`pretraining_dataset`) |
+|---|---|---|
+| **Dataset size** | Fits in memory | Too large to fit in memory |
+| **Tokenization** | Pre-tokenized before training | On-demand during training |
+| **Config key** | `datasets:` | `pretraining_dataset:` |
+| **Long text handling** | Splits texts exceeding `sequence_len` | Concatenates into fixed-length sequences |
+| **Benefit** | Can preprocess on CPU, transfer to GPU | Start training immediately, no preprocessing |
+
+## Non-Streaming: `type: completion`
+
+For smaller datasets that fit in memory. Pre-tokenizes the entire dataset.
+
+```yaml
+datasets:
+  - path: my_corpus
+    type: completion
+    # field: text              # Column name (default: "text")
+```
+
+## Streaming: `pretraining_dataset`
+
+For large corpora. Streams data on-demand without loading everything into memory.
+
+```yaml
+pretraining_dataset:
+  - path: HuggingFaceFW/fineweb-edu
+    type: pretrain
+    text_column: text
+    split: train
+
+max_steps: 1000                          # Required — axolotl can't infer dataset size
+streaming_multipack_buffer_size: 10000   # Buffer for sample packing
+pretrain_multipack_attn: true            # Prevent cross-attention between packed samples
+```
+
+`max_steps` is required for streaming — one step = `sequence_len * micro_batch_size * gradient_accumulation_steps * num_gpus` tokens.
+
+Full streaming docs: [streaming.qmd](../streaming.qmd)
+
+## Dataset Format
+
+```json
+{"text": "The complete document text goes here."}
+```
+
+## Key Settings
+
+- `sample_packing: true` + `pad_to_sequence_len: true` — pack documents into fixed-length sequences
+- `flash_attention: true` — required for sample packing
+- No adapter — typically full fine-tune for pretraining
+- `train_on_inputs: true` — default for completion (all tokens trained on)
+
+## File Map
+
+```
+src/axolotl/
+  prompt_strategies/completion.py    # Non-streaming: completion prompt strategy (no masking)
+  utils/data/sft.py                  # Non-streaming: dataset loading and processing
+  utils/data/streaming.py            # Streaming: encode_streaming(), wrap_streaming_dataset()
+  utils/schemas/config.py            # Config fields: pretraining_dataset, pretrain_multipack_attn, etc.
+
+examples/streaming/pretrain.yaml     # Full streaming pretraining example config
+```
diff --git a/docs/agents/reward_modelling.md b/docs/agents/reward_modelling.md
new file mode 100644
index 0000000000..055dc6b36c
--- /dev/null
+++ b/docs/agents/reward_modelling.md
@@ -0,0 +1,48 @@
+# Reward Modelling — Agent Reference
+
+Train models to score responses for use as reward signals in RL. For full docs, see [reward_modelling.qmd](../reward_modelling.qmd).
+
+## Types
+
+### Outcome Reward Models (ORM)
+
+Train a classifier to predict preference over entire interactions. Uses `AutoModelForSequenceClassification`.
+
+```yaml
+base_model: google/gemma-2-2b
+model_type: AutoModelForSequenceClassification
+num_labels: 1
+reward_model: true
+chat_template: gemma
+datasets:
+  - path: argilla/distilabel-intel-orca-dpo-pairs
+    type: bradley_terry.chat_template
+```
+
+Dataset format: `{"system": "...", "input": "...", "chosen": "...", "rejected": "..."}`
+
+### Process Reward Models (PRM)
+
+Train a token classifier to score each reasoning step. Uses `AutoModelForTokenClassification`.
+
+```yaml
+base_model: Qwen/Qwen2.5-3B
+model_type: AutoModelForTokenClassification
+num_labels: 2
+process_reward_model: true
+datasets:
+  - path: trl-lib/math_shepherd
+    type: stepwise_supervised
+```
+
+Dataset format: see [stepwise_supervised.qmd](../dataset-formats/stepwise_supervised.qmd).
+
+## File Map
+
+```
+src/axolotl/
+  core/builders/causal.py                    # Handles reward_model flag in trainer builder
+  prompt_strategies/bradley_terry/           # Bradley-Terry prompt strategies
+  prompt_strategies/stepwise_supervised.py   # PRM dataset strategy
+  utils/schemas/config.py                    # reward_model, process_reward_model config fields
+```
diff --git a/docs/agents/sft.md b/docs/agents/sft.md
new file mode 100644
index 0000000000..f601cb0f53
--- /dev/null
+++ b/docs/agents/sft.md
@@ -0,0 +1,139 @@
+# SFT — Agent Reference
+
+Supervised fine-tuning pipeline reference. For config templates and dataset format examples, see [getting-started.qmd](../getting-started.qmd) and [dataset-formats/](../dataset-formats/).
+
+## Architecture
+
+```
+YAML Config → axolotl train config.yaml
+
+  1. Load base model (+ quantization if QLoRA/8-bit)
+  2. Apply adapter layers (LoRA/QLoRA) if configured
+  3. Load + tokenize dataset(s)
+     - Apply prompt template (chat_template / alpaca / custom)
+     - Mask inputs (train_on_inputs: false)
+     - Pack samples into sequences (sample_packing: true)
+  4. Training loop (HuggingFace Trainer)
+     - forward → loss → backward → optimizer step → lr scheduler step
+  5. Save model / adapter weights + tokenizer
+
+Multi-GPU: FSDP or DeepSpeed shards model across GPUs automatically.
+```
+
+## Components Required
+
+1. A YAML config — model, dataset(s), adapter settings, hyperparameters
+2. A dataset — HuggingFace Hub, local JSONL/JSON/Parquet, or S3/GCS path
+3. (Optional) A custom prompt strategy — for non-standard dataset formats
+
+No external server processes needed (unlike GRPO which requires vLLM).
+
+## Dataset Format Decision Tree
+
+```
+Is your data in chat/message format?
+  ├─ YES: OpenAI message format (role/content)?
+  │   ├─ YES ──────────────────────> type: chat_template  (recommended)
+  │   └─ NO (custom field names) ──> type: chat_template + message_property_mappings
+  └─ NO: Instruction/response pairs?
+      ├─ YES ──> type: alpaca       (instruction, input, output)
+      └─ NO: Raw text?
+          ├─ YES with segments ─────> type: input_output  (template-free masking)
+          └─ YES continuous ────────> type: completion     (pretraining-style)
+```
+
+Full format specs: [dataset-formats/](../dataset-formats/)
+
+## Model Size to Adapter Choice
+
+| Model Size | LoRA | QLoRA (4-bit) | Full Fine-Tune | VRAM (approx) |
+|-----------|------|---------------|----------------|---------------|
+| 1-3B | Preferred | Low-budget option | Single GPU OK | 8-16 GB (LoRA) |
+| 7-8B | Preferred | Good balance | Needs multi-GPU | 16-24 GB (LoRA) |
+| 13-14B | Preferred | Good balance | Multi-GPU required | 24-40 GB (LoRA) |
+| 30-70B | LoRA or QLoRA | Preferred for single GPU | Multi-node | 40-80 GB (QLoRA) |
+
+## Hyperparameter Ranges
+
+| Parameter | LoRA | QLoRA | Full FT |
+|-----------|------|-------|---------|
+| `learning_rate` | 1e-4 to 3e-4 | 1e-4 to 3e-4 | 1e-5 to 5e-5 |
+| `lora_r` | 16-64 | 16-64 | N/A |
+| `lora_alpha` | 1-2x `lora_r` | 1-2x `lora_r` | N/A |
+| `micro_batch_size` | 2-8 | 2-4 | 1-2 |
+| `gradient_accumulation_steps` | 2-8 | 4-16 | 4-16 |
+| `num_epochs` | 1-3 | 1-3 | 1-3 |
+| `optimizer` | `adamw_8bit` | `adamw_bnb_8bit` | `adamw_torch_fused` |
+
+Effective batch = micro_batch * grad_accum * num_gpus. Lower LR for larger models.
+
+## Healthy Training Indicators
+
+| Metric | Healthy | Problem |
+|--------|---------|---------|
+| `train_loss` | Decreasing, starting ~2-4 for chat models | Flat or increasing from step 1 — data or LR issue |
+| `eval_loss` | Decreasing, tracks train_loss | Increasing while train_loss decreases — overfitting |
+| `grad_norm` | 0.1-10, relatively stable | Spikes >100 — instability. 0.0 — frozen weights |
+| `learning_rate` | Follows scheduler curve | Flat or NaN — config issue |
+
+Watch for: loss never decreasing (check `train_on_inputs`, dataset, LR), loss goes to 0 quickly (overfitting), eval_loss diverging (reduce epochs, add regularization). See [training_stability.qmd](../training_stability.qmd).
+
+## Known Issues
+
+| Issue | Fix |
+|-------|-----|
+| OOM during training | Reduce `micro_batch_size`, enable `gradient_checkpointing`, reduce `sequence_len` |
+| `sample_packing` + SDPA + bf16 = 0.0 loss | Use `attn_implementation: flash_attention_2` or disable `sample_packing` |
+| Missing chat template error | Set `chat_template: chatml` explicitly |
+| Label masking wrong | Run `axolotl preprocess config.yaml --debug` and inspect labels |
+| Loss NaN | Use `bf16: auto`, lower LR, check data for empty samples |
+| Tokenizer pad token / infinite loss | Set `special_tokens: pad_token: "<\|end_of_text\|>"` |
+| FSDP save hangs | Use `fsdp_state_dict_type: FULL_STATE_DICT` |
+| DeepSpeed CheckpointError | Set `use_reentrant: true` in `gradient_checkpointing_kwargs` |
+
+## Profiling
+
+To profile training and identify optimization opportunities:
+
+```yaml
+# Profile steps 3-7 (after warmup/autotuning settles)
+profiler_steps_start: 3
+profiler_steps: 5
+```
+
+This produces `profiler_trace.json` (Chrome trace) and `snapshot.pickle` (memory snapshot) in `output_dir`.
+View the Chrome trace at `chrome://tracing`.
+
+To programmatically inspect the trace:
+```bash
+python scripts/analyze_profile.py output_dir/
+```
+
+The trace shows per-kernel CUDA times, memory allocations, and operator-level breakdown. Look for:
+- **Large matmul kernels**: candidates for fusion or quantization
+- **Memory copies (H2D/D2H)**: unnecessary data movement
+- **Small frequent kernels**: candidates for kernel fusion
+- **Gaps between kernels**: pipeline bubbles from CPU overhead
+
+Full troubleshooting: [training_stability.qmd](../training_stability.qmd), [debugging.qmd](../debugging.qmd)
+
+## File Map
+
+```
+src/axolotl/
+  cli/train.py                     # Entry point for `axolotl train`
+  cli/preprocess.py                # Entry point for `axolotl preprocess`
+  core/builders/causal.py          # HFCausalTrainerBuilder — wires config → SFT trainer
+  core/trainers/base.py            # AxolotlTrainer — base trainer class
+  core/trainers/mixins/            # Packing, optimizer, scheduler, checkpoints
+  prompt_strategies/               # Format handlers: chat_template, alpaca, completion, input_output
+  utils/schemas/config.py          # AxolotlInputConfig — main config schema
+  utils/schemas/datasets.py        # SFTDataset, DatasetConfig
+  utils/schemas/peft.py            # LoraConfig — LoRA parameters
+  integrations/liger/              # Liger kernel plugin
+
+examples/llama-3/                  # LoRA, QLoRA, full FT example configs
+docs/getting-started.qmd           # Quickstart with config templates
+docs/optimizations.qmd             # Flash attention, gradient checkpointing, sample packing
+docs/multi-gpu.qmd                 # FSDP and DeepSpeed setup
+```
diff --git a/docs/amd_hpc.qmd b/docs/amd_hpc.qmd
new file mode 100644
index 0000000000..259b01ab54
--- /dev/null
+++ b/docs/amd_hpc.qmd
@@ -0,0 +1,108 @@
+---
+title: AMD GPUs on HPC Systems
+description: A comprehensive guide for using Axolotl on distributed systems with AMD GPUs
+---
+
+This guide provides step-by-step instructions for installing and configuring Axolotl on a High-Performance Computing (HPC) environment equipped with AMD GPUs.
+
+## Setup
+
+### 1. Install Python
+
+We recommend using Miniforge, a minimal conda-based Python distribution:
+
+```bash
+curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
+bash Miniforge3-$(uname)-$(uname -m).sh
+```
+
+### 2. Configure Python Environment
+Add Python to your PATH and ensure it's available at login:
+
+```bash
+echo 'export PATH=~/miniforge3/bin:$PATH' >> ~/.bashrc
+echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' >> ~/.bash_profile
+```
+
+### 3. Load AMD GPU Software
+
+Load the ROCm module:
+
+```bash
+module load rocm/5.7.1
+```
+
+Note: The specific module name and version may vary depending on your HPC system. Consult your system documentation for the correct module name.
+
+### 4. Install PyTorch
+
+Install PyTorch with ROCm support:
+
+```bash
+pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.7 --force-reinstall
+```
+
+### 5. Install Flash Attention
+
+Clone and install the Flash Attention repository:
+
+```bash
+git clone --recursive https://github.com/ROCmSoftwarePlatform/flash-attention.git
+export GPU_ARCHS="gfx90a"
+cd flash-attention
+export PYTHON_SITE_PACKAGES=$(python -c 'import site; print(site.getsitepackages()[0])')
+patch "${PYTHON_SITE_PACKAGES}/torch/utils/hipify/hipify_python.py" hipify_patch.patch
+pip install --no-build-isolation .
+```
+
+### 6. Install Axolotl
+
+Clone and install Axolotl:
+
+```bash
+git clone https://github.com/axolotl-ai-cloud/axolotl
+cd axolotl
+pip install packaging ninja
+pip install --no-build-isolation -e .
+```
+
+### 7. Apply xformers Workaround
+
+xformers appears to be incompatible with ROCm. Apply the following workarounds:
+ - Edit $HOME/packages/axolotl/src/axolotl/monkeypatch/llama_attn_hijack_flash.py modifying the code to always return `False` for SwiGLU availability from xformers.
+ - Edit $HOME/miniforge3/lib/python3.10/site-packages/xformers/ops/swiglu_op.py replacing the "SwiGLU" function with a pass statement.
+
+### 8. Prepare Job Submission Script
+
+Create a script for job submission using your HPC's particular software (e.g. Slurm, PBS). Include necessary environment setup and the command to run Axolotl training. If the compute node(s) do(es) not have internet access, it is recommended to include
+
+```bash
+export TRANSFORMERS_OFFLINE=1
+export HF_DATASETS_OFFLINE=1
+```
+
+### 9. Download Base Model
+
+Download a base model using the Hugging Face CLI:
+
+```bash
+hf download meta-llama/Meta-Llama-3.1-8B --local-dir ~/hfdata/llama3.1-8B
+```
+
+### 10. Create Axolotl Configuration
+
+Create an Axolotl configuration file (YAML format) tailored to your specific training requirements and dataset. Use FSDP for multi-node training.
+
+Note: Deepspeed did not work at the time of testing. However, if anyone managed to get it working, please let us know.
+
+### 11. Preprocess Data
+
+Run preprocessing on the login node:
+
+```bash
+CUDA_VISIBLE_DEVICES="" python -m axolotl.cli.preprocess /path/to/your/config.yaml
+```
+
+### 12. Train
+
+You are now ready to submit your previously prepared job script. 🚂
diff --git a/docs/attention.qmd b/docs/attention.qmd
new file mode 100644
index 0000000000..64820fb6ef
--- /dev/null
+++ b/docs/attention.qmd
@@ -0,0 +1,226 @@
+---
+title: Attention
+description: Supported attention modules in Axolotl
+---
+
+Axolotl routes attention via a single config field:
+
+```yaml
+attn_implementation: <backend>
+```
+
+`attn_implementation` is passed through to `transformers` verbatim (via
+`model.config._attn_implementation`). Accepted values are the HF-native
+backends, axolotl-registered backends, or a hub-kernel path.
+
+## Backends
+
+| `attn_implementation` | Description |
+|---|---|
+| `eager` | Plain PyTorch attention. No packing support. |
+| `sdpa` | PyTorch `scaled_dot_product_attention`. No packing support. |
+| `flash_attention_2` | Dao-AILab Flash Attention 2. |
+| `flash_attention_3` | Dao-AILab Flash Attention 3 (Hopper+). |
+| `flex_attention` | Torch Flex Attention (requires torch ≥ 2.6). |
+| `xformers` | xFormers memory-efficient attention. |
+| `sage` | SageAttention (QK int8 / PV fp16). |
+| `s2` | Shifted-Sparse Attention (LLaMA only, FA2 under the hood). |
+| `fp8` | torchao FP8 low-precision attention (requires SM90+, torch ≥ 2.11). Loaded as SDPA and patched post-load. |
+| `kernels-community/flash-attn3` | HF hub FA3 kernel. |
+| `kernels-community/sage-attention` | HF hub SageAttention kernel. |
+| Other `<org>/<name>` path | Any hub-kernel path supported by `transformers`. |
+
+Short-form aliases (`flash`, `fa2`, `flex`, `sdp`, etc.) are **not accepted** —
+set the canonical name above.
+
+### Capability flags
+
+Axolotl derives three boolean capability flags from `attn_implementation` and
+exposes them on the validated config:
+
+- `cfg.attn_supports_packing` — backend supports varlen sample packing via
+  `position_ids`. Gates multipack patches and `sample_packing_drop_attention_mask`.
+- `cfg.attn_uses_flash_lib` — backend needs the `flash_attn` (Dao-AILab)
+  monkeypatches (FA4 auto, LLaMA flash hijack, ring-FA).
+- `cfg.attn_needs_dtype_cast` — backend requires fp16/bf16 embeddings
+  (everything except `eager` and `sdpa`).
+
+These are **computed** — they cannot be overridden from YAML.
+
+## Per-backend notes
+
+### SDPA
+
+Default PyTorch attention. See
+[PyTorch docs](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html).
+
+```yaml
+attn_implementation: sdpa
+```
+
+### Flash Attention
+
+Axolotl supports FA2, FA3, and FA4. The best available version is used
+automatically based on your installed packages and GPU.
+
+```yaml
+attn_implementation: flash_attention_2  # or flash_attention_3
+```
+
+#### Flash Attention 2
+
+Requirements: Ampere, Ada, or Hopper GPUs (Turing or lower not supported)
+
+```bash
+pip install flash-attn --no-build-isolation
+```
+
+::: {.callout-tip}
+
+If you get `undefined symbol` while training, ensure you installed PyTorch prior to Axolotl.
+Alternatively, try reinstall or downgrade a version.
+
+:::
+
+#### Flash Attention 3
+
+Requirements: Hopper only and CUDA 12.8 (recommended)
+
+```bash
+git clone https://github.com/Dao-AILab/flash-attention.git
+cd flash-attention/hopper
+python setup.py install
+```
+
+#### Flash Attention 4
+
+Requirements: Hopper or Blackwell GPUs. Auto-applied when `attn_uses_flash_lib`
+is true and FA4 is importable.
+
+FA4 is still a pre-release on PyPI, so `--pre` is required:
+
+```bash
+pip install --pre flash-attn-4
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/Dao-AILab/flash-attention.git
+cd flash-attention/flash_attn/cute
+pip install -e .
+
+# FA2's flash_attn package includes a cute/ stub that shadows FA4.
+# Remove it so Python can find the real FA4 module:
+rm -r $(python -c "import flash_attn; print(flash_attn.__path__[0])")/cute
+```
+
+::: {.callout-note}
+
+**Hopper (SM90) users**: The backward kernel is not yet included in the pip package. To use FA4
+for training on Hopper, install from source using the instructions above.
+
+:::
+
+::: {.callout-warning}
+
+FA4 only supports head dimensions up to 128 (`d ≤ 128`). The DeepSeek shape `(192, 128)` is
+also supported but only on Blackwell. Axolotl automatically detects incompatible head dimensions
+and falls back to FA2/3.
+
+:::
+
+### AMD
+
+Requirements: ROCm 6.0 and above. See
+[Flash Attention AMD docs](https://github.com/Dao-AILab/flash-attention/tree/main?tab=readme-ov-file#amd-rocm-support).
+
+### Flex Attention
+
+```yaml
+attn_implementation: flex_attention
+torch_compile: true  # recommended
+```
+
+Requires torch ≥ 2.6. See [PyTorch docs](https://pytorch.org/blog/flexattention/).
+
+### SageAttention
+
+Requirements: Ampere, Ada, or Hopper GPUs.
+
+```yaml
+attn_implementation: sage
+```
+
+```bash
+pip install sageattention==2.2.0 --no-build-isolation
+```
+
+::: {.callout-warning}
+
+Only LoRA/QLoRA recommended. Full finetuning has been observed to drop loss to 0. See
+[GitHub Issue](https://github.com/thu-ml/SageAttention/issues/198).
+
+:::
+
+For more details: [Sage Attention](https://github.com/thu-ml/SageAttention).
+
+### xFormers
+
+```yaml
+attn_implementation: xformers
+```
+
+::: {.callout-tip}
+
+Recommended for Turing GPUs or below (e.g. Colab T4).
+
+:::
+
+### FP8
+
+torchao low-precision attention. Loaded as SDPA and patched post-load.
+
+Requirements: SM90+ (Hopper/Blackwell), PyTorch ≥ 2.11, torchao ≥ 0.17,
+flash-attn with FA3. KV caching must be disabled.
+
+```yaml
+attn_implementation: fp8
+```
+
+### Hub kernels
+
+```yaml
+attn_implementation: kernels-community/flash-attn3
+```
+
+Passed through to `transformers`; axolotl does not install the kernel itself.
+For recognized hub paths the capability flags are set automatically; for
+arbitrary paths axolotl uses conservative defaults (`attn_supports_packing=False`,
+`attn_uses_flash_lib=False`).
+
+## Migrating from legacy boolean flags
+
+The following legacy config fields are **deprecated** and will be removed in a
+future release. Each emits a `DeprecationWarning` when set and is stripped from
+the validated config.
+
+| Legacy | Canonical |
+|---|---|
+| `flash_attention: true` | `attn_implementation: flash_attention_2` |
+| `sdp_attention: true` | `attn_implementation: sdpa` |
+| `xformers_attention: true` | `attn_implementation: xformers` |
+| `flex_attention: true` | `attn_implementation: flex_attention` |
+| `sage_attention: true` | `attn_implementation: sage` |
+| `eager_attention: true` | `attn_implementation: eager` |
+
+Combining `attn_implementation` with a legacy flag (e.g. `attn_implementation:
+flash_attention_2` **and** `flash_attention: true`) raises — pick one.
+
+::: {.callout-note}
+
+Existing example configs under `examples/` still use the legacy flags. They
+continue to work with a deprecation warning; they will be migrated in a
+follow-up pass.
+
+:::
diff --git a/docs/checkpoint_saving.qmd b/docs/checkpoint_saving.qmd
new file mode 100644
index 0000000000..5f6b155e7b
--- /dev/null
+++ b/docs/checkpoint_saving.qmd
@@ -0,0 +1,86 @@
+---
+title: "Checkpoint Saving"
+format:
+  html:
+    toc: true
+    toc-depth: 2
+    number-sections: true
+execute:
+  enabled: false
+---
+
+## Overview
+
+Axolotl supports on-demand checkpoint saving during training. You can trigger checkpoints via file-based triggers (for programmatic control) or Control+C (for interactive use).
+
+## File-Based Checkpoint Trigger
+
+### Configuration
+
+Enable in your config:
+
+```yaml
+dynamic_checkpoint:
+  enabled: true
+  check_interval: 100  # Optional: check every N steps (default: 100)
+  trigger_file_path: "axolotl_checkpoint.save"  # Optional: custom filename
+```
+
+**Options:**
+- `enabled`: `true` to enable (required)
+- `check_interval`: Steps between file checks. Default: 100. Lower = faster response, higher I/O overhead.
+- `trigger_file_path`: Custom trigger filename. Default: `axolotl_checkpoint.save`
+
+### How It Works
+
+1. Rank 0 checks for trigger file every `check_interval` steps in `output_dir`
+2. When detected, file is deleted and checkpoint is saved
+3. In distributed training, rank 0 broadcasts to synchronize all ranks
+
+### Usage
+
+**Command line:**
+```bash
+touch /path/to/output_dir/axolotl_checkpoint.save
+```
+
+**Programmatic:**
+```python
+from pathlib import Path
+Path("/path/to/output_dir/axolotl_checkpoint.save").touch()
+```
+
+Checkpoint saves within the next `check_interval` steps. The trigger file is auto-deleted after detection, so you can create it multiple times.
+
+**Custom filename:**
+```yaml
+dynamic_checkpoint:
+  enabled: true
+  trigger_file_path: "my_trigger.save"
+```
+```bash
+touch /path/to/output_dir/my_trigger.save
+```
+
+## Control+C (SIGINT) Checkpoint
+
+Pressing `Ctrl+C` during training saves the model state and exits gracefully. **Note:** This saves only the model weights, not optimizer state. For resumable checkpoints, use the file-based trigger.
+
+## Best Practices
+
+- **Check interval**: Lower values (10-50) for fast training, default 100 for slower training
+- **Distributed training**: Create trigger file once; rank 0 handles synchronization
+- **Resume**: Dynamic checkpoints can be resumed like regular checkpoints via `resume_from_checkpoint`
+
+## Example
+
+```yaml
+output_dir: ./outputs/lora-out
+save_steps: 500  # Scheduled checkpoints
+
+dynamic_checkpoint:
+  enabled: true
+  check_interval: 50
+```
+
+This enables scheduled checkpoints every 500 steps plus on-demand saves via file trigger (checked every 50 steps).
diff --git a/docs/choosing_method.qmd b/docs/choosing_method.qmd
new file mode 100644
index 0000000000..942c52a97d
--- /dev/null
+++ b/docs/choosing_method.qmd
@@ -0,0 +1,206 @@
+---
+title: "Which Fine-Tuning Method Should I Use?"
+description: "A decision guide for choosing the right fine-tuning method, adapter, and hardware configuration in Axolotl."
+format:
+  html:
+    toc: true
+    toc-depth: 3
+    number-sections: true
+execute:
+  enabled: false
+---
+
+## Overview {#sec-overview}
+
+Axolotl supports four broad categories of fine-tuning, each suited to different data types, objectives, and resource constraints.
+
+| Method | What It Does | Data You Need |
+|--------|-------------|---------------|
+| **Supervised Fine-Tuning (SFT)** | Teaches the model to produce specific outputs given inputs | Input-output pairs (instructions, conversations, completions) |
+| **Preference Learning (DPO/KTO/ORPO)** | Steers the model toward preferred outputs and away from dispreferred ones | Chosen/rejected response pairs (DPO, ORPO) or binary labels (KTO) |
+| **Reinforcement Learning (GRPO)** | Optimizes the model against a reward signal through online generation | A reward function (code or model-based) and a prompt dataset |
+| **Reward Modeling** | Trains a model to score responses, for use as a reward signal in RL | Preference pairs ranked by quality |
+
+Each method is configured through a YAML file with `rl: <method>` (or omitted for SFT). All methods support LoRA, QLoRA, and full fine-tuning unless otherwise noted.
+
+## Decision Tree {#sec-decision-tree}
+
+Use the following flowchart to choose your method. Start at the top and follow the path that matches your situation.
+
+```
+Do you have a reward function (code-based or model-based)?
+├── YES
+│   └── Use GRPO (rl: grpo)
+│       The model generates its own completions and learns from reward scores.
+│       Best for: math, code, reasoning, tasks with verifiable answers.
+│       See: rlhf.qmd#grpo
+│
+└── NO
+    │
+    Do you have preference pairs (chosen vs. rejected responses)?
+    ├── YES
+    │   │
+    │   Are they paired (same prompt, one chosen, one rejected)?
+    │   ├── YES → Use DPO (rl: dpo)
+    │   │         Direct optimization without a separate reward model.
+    │   │         See: rlhf.qmd#dpo
+    │   │
+    │   └── NO (only binary good/bad labels)
+    │       └── Use KTO (rl: kto)
+    │           Works with unpaired preference data.
+    │           See: rlhf.qmd#kto
+    │
+    └── NO
+        │
+        Do you have input-output examples?
+        ├── YES → Use SFT
+        │         The simplest and most common method.
+        │         See: getting-started.qmd
+        │
+        └── NO
+            └── You need to create training data first.
+                Consider generating preference pairs with an LLM judge,
+                or writing a reward function for GRPO.
+```
+
+::: {.callout-tip}
+**When in doubt, start with SFT.** It is the most straightforward method and works well for most tasks. You can always move to preference learning or RL later to further refine behavior.
+:::
+
+### Method Comparison at a Glance
+
+| Criterion | SFT | DPO | KTO | GRPO |
+|-----------|-----|-----|-----|------|
+| Data complexity | Low (input-output pairs) | Medium (preference pairs) | Medium (binary labels) | Low (prompts + reward code) |
+| Compute cost | Low | Medium | Medium | High (requires vLLM server) |
+| Learning signal | Supervised | Contrastive | Contrastive | Online reward |
+| Online generation | No | No | No | Yes |
+| Reward model needed | No | No | No | No (uses reward functions) |
+| Best for | Task adaptation, instruction following | Safety, style alignment | Unpaired preference data | Reasoning, math, code |
+
+::: {.callout-note}
+**ORPO** is an alternative to DPO that combines SFT and preference optimization in a single training stage, removing the need for a separate SFT step. Configure with `rl: orpo`. See [rlhf.qmd](rlhf.qmd) for details.
+:::
+
+## Adapter Selection {#sec-adapter-selection}
+
+Once you have chosen a method, decide how to apply the parameter updates. The three main options trade off VRAM usage against model quality.
+
+### QLoRA
+
+- **How it works**: The base model is loaded in 4-bit (NF4) quantization. Small low-rank adapter matrices are trained in higher precision on top.
+- **VRAM savings**: Roughly 4x reduction in model memory compared to full fine-tuning.
+- **Quality**: Slight degradation due to quantization noise, but often negligible for task-specific fine-tuning.
+- **When to use**: When your GPU cannot fit the model in full precision, or when you want fast experimentation.
+
+```yaml
+adapter: qlora
+load_in_4bit: true
+lora_r: 32
+lora_alpha: 64
+lora_target_linear: true
+```
+
+### LoRA
+
+- **How it works**: The base model is loaded at full precision (or 8-bit). Low-rank adapter matrices are trained alongside.
+- **VRAM savings**: Roughly 2-3x reduction compared to full fine-tuning (model weights are frozen, only adapters + optimizer states for adapters are stored).
+- **Quality**: Very close to full fine-tuning for most tasks, especially with higher rank values.
+- **When to use**: When you have enough VRAM for the base model but not for full optimizer states.
+
+```yaml
+adapter: lora
+lora_r: 32
+lora_alpha: 64
+lora_target_linear: true
+```
+
+::: {.callout-tip}
+For GRPO training, LoRA is strongly recommended. The vLLM server needs to sync weights from the trainer, and LoRA sync (`trl.vllm_lora_sync: true`) is far more efficient than syncing full merged weights. See [vLLM Serving](vllm_serving.qmd) for details.
+:::
+
+### Full Fine-Tuning
+
+- **How it works**: All model parameters are updated during training. No adapters.
+- **VRAM savings**: None. Requires memory for model weights, gradients, and optimizer states (roughly 4x model size in bf16 with AdamW).
+- **Quality**: Highest potential quality, especially for large distribution shifts.
+- **When to use**: When you have ample GPU memory or multi-GPU setups, and need maximum performance. Also required for pre-training.
+
+```yaml
+# No adapter or load_in_* lines needed
+micro_batch_size: 1
+gradient_accumulation_steps: 16
+```
+
+### Quick Comparison
+
+| | QLoRA | LoRA | Full |
+|---|---|---|---|
+| Trainable params | ~0.1-1% | ~0.1-1% | 100% |
+| Model memory | ~25% of full | ~50-100% of full | 100% |
+| Optimizer memory | Tiny (adapters only) | Tiny (adapters only) | 2x model size (AdamW) |
+| Training speed | Slower (dequantization overhead) | Baseline | Faster per-step (no adapter overhead) |
+| Inference | Merge or serve with adapter | Merge or serve with adapter | Direct |
+| Multi-GPU required? | Rarely | For 13B+ models | For 7B+ models |
+
+## Hardware Mapping {#sec-hardware-mapping}
+
+The tables below provide approximate GPU memory requirements. Actual usage depends on context length, batch size, and optimizer choice.
+
+### SFT / Preference Learning
+
+| Model Size | QLoRA (4-bit) | LoRA (bf16) | Full (bf16 + AdamW) |
+|------------|--------------|-------------|---------------------|
+| 1-3B | 6-8 GB | 8-12 GB | 24-32 GB |
+| 7-8B | 10-14 GB | 16-24 GB | 60-80 GB |
+| 13-14B | 16-20 GB | 28-40 GB | 120+ GB |
+| 30-34B | 24-32 GB | 64-80 GB | 2-4x 80 GB |
+| 70-72B | 40-48 GB | 2x 80 GB | 4-8x 80 GB |
+
+::: {.callout-important}
+These estimates assume a short context length (512-2048 tokens) and micro_batch_size of 1-2. Longer sequences and larger batches increase memory significantly due to activations. Use [gradient checkpointing](gradient_checkpointing.qmd) to reduce activation memory at the cost of ~30% slower training.
+:::
+
+### GRPO (RL Training)
+
+GRPO requires additional GPU(s) for the vLLM generation server. Plan for at least two GPUs: one for training, one for vLLM.
+
+| Model Size | Training GPU (LoRA, bf16) | vLLM GPU | Total GPUs |
+|------------|--------------------------|----------|------------|
+| 0.5-3B | 1x 24 GB | 1x 24 GB | 2x 24 GB |
+| 7-8B | 1x 80 GB | 1x 80 GB | 2x 80 GB |
+| 13-14B | 1-2x 80 GB | 1-2x 80 GB | 2-4x 80 GB |
+| 30-72B | 2-4x 80 GB (FSDP/DeepSpeed) | 2-4x 80 GB (tensor parallel) | 4-8x 80 GB |
+
+::: {.callout-tip}
+For single-GPU GRPO, use `vllm_mode: colocate` with `vllm_enable_sleep_mode: true`. The vLLM engine shares the GPU and offloads VRAM when not generating. This works for smaller models (up to ~3B on a 24 GB GPU) but is slower than the two-GPU server mode.
+:::
+
+### Multi-GPU Threshold
+
+You need multi-GPU training when:
+
+- **Full fine-tuning** of models 7B+ (use FSDP or DeepSpeed ZeRO)
+- **LoRA** of models 30B+ (or 13B+ with long contexts)
+- **GRPO** almost always (separate vLLM server), unless using colocate mode
+
+See [Multi-GPU Training](multi-gpu.qmd) for FSDP and DeepSpeed configuration.
+
+## Quick Links {#sec-quick-links}
+
+| Method | Config Key | Documentation | Example Config |
+|--------|-----------|---------------|----------------|
+| SFT | *(default, no `rl:` key)* | [Getting Started](getting-started.qmd) | `examples/llama-3/lora-1b.yml` |
+| DPO | `rl: dpo` | [RLHF - DPO](rlhf.qmd#dpo) | See rlhf.qmd |
+| KTO | `rl: kto` | [RLHF - KTO](rlhf.qmd#kto) | See rlhf.qmd |
+| ORPO | `rl: orpo` | [RLHF - ORPO](rlhf.qmd#orpo) | See rlhf.qmd |
+| GRPO | `rl: grpo` | [RLHF - GRPO](rlhf.qmd#grpo), [vLLM Serving](vllm_serving.qmd) | See rlhf.qmd |
+| Reward Modeling | `rl: reward_trainer` | [Reward Modelling](reward_modelling.qmd) | See reward_modelling.qmd |
+
+### Related Guides
+
+- [Configuration Reference](config-reference.qmd) -- Full list of all config options
+- [Dataset Formats](dataset-formats) -- How to structure your training data
+- [Optimizations](optimizations.qmd) -- Flash attention, gradient checkpointing, mixed precision
+- [Multi-GPU Training](multi-gpu.qmd) -- FSDP and DeepSpeed setup
+- [vLLM Serving](vllm_serving.qmd) -- Setting up vLLM for GRPO training
diff --git a/docs/cli.qmd b/docs/cli.qmd
new file mode 100644
index 0000000000..b44c5db23b
--- /dev/null
+++ b/docs/cli.qmd
@@ -0,0 +1,348 @@
+---
+title: "Command Line Interface (CLI)"
+format:
+  html:
+    toc: true
+    toc-expand: 2
+    toc-depth: 3
+execute:
+  enabled: false
+---
+
+The Axolotl CLI provides a streamlined interface for training and fine-tuning large language models. This guide covers
+the CLI commands, their usage, and common examples.
+
+
+## Basic Commands
+
+All Axolotl commands follow this general structure:
+
+```bash
+axolotl <command> [config.yml] [options]
+```
+
+The config file can be local or a URL to a raw YAML file.
+
+### Launcher Arguments
+
+For commands that support multi-GPU (`train`, `evaluate`, ...), you can pass launcher-specific arguments using the `--` separator:
+
+```bash
+# Pass torchrun arguments
+axolotl train config.yml --launcher torchrun -- --nproc_per_node=2 --nnodes=1
+
+# Pass accelerate arguments
+axolotl train config.yml --launcher accelerate -- --config_file=accelerate_config.yml --num_processes=4
+```
+
+Arguments after `--` are passed directly to the launcher (torchrun, accelerate launch, etc.).
+
+## Command Reference
+
+### fetch
+
+Downloads example configurations and deepspeed configs to your local machine.
+
+```bash
+# Get example YAML files
+axolotl fetch examples
+
+# Get deepspeed config files
+axolotl fetch deepspeed_configs
+
+# Specify custom destination
+axolotl fetch examples --dest path/to/folder
+```
+
+### preprocess
+
+Preprocesses and tokenizes your dataset before training. This is recommended for large datasets.
+
+```bash
+# Basic preprocessing
+axolotl preprocess config.yml
+
+# Preprocessing with one GPU
+CUDA_VISIBLE_DEVICES="0" axolotl preprocess config.yml
+
+# Debug mode to see processed examples
+axolotl preprocess config.yml --debug
+
+# Debug with limited examples
+axolotl preprocess config.yml --debug --debug-num-examples 5
+```
+
+Configuration options:
+
+```yaml
+dataset_prepared_path: Local folder for saving preprocessed data
+push_dataset_to_hub: HuggingFace repo to push preprocessed data (optional)
+```
+
+### train
+
+Trains or fine-tunes a model using the configuration specified in your YAML file.
+
+```bash
+# Basic training
+axolotl train config.yml
+
+# Train and set/override specific options
+axolotl train config.yml \
+    --learning-rate 1e-4 \
+    --micro-batch-size 2 \
+    --num-epochs 3
+
+# Training without accelerate
+axolotl train config.yml --launcher python
+
+# Pass launcher-specific arguments using -- separator
+axolotl train config.yml --launcher torchrun -- --nproc_per_node=2 --nnodes=1
+axolotl train config.yml --launcher accelerate -- --config_file=accelerate_config.yml
+
+# Resume training from checkpoint
+axolotl train config.yml --resume-from-checkpoint path/to/checkpoint
+```
+
+It is possible to run sweeps over multiple hyperparameters by passing in a sweeps config.
+
+```bash
+# Basic training with sweeps
+axolotl train config.yml --sweep path/to/sweep.yaml
+```
+
+Example sweep config:
+```yaml
+_:
+  # This section is for dependent variables we need to fix
+  - load_in_8bit: false
+    load_in_4bit: false
+    adapter: lora
+  - load_in_8bit: true
+    load_in_4bit: false
+    adapter: lora
+
+# These are independent variables
+learning_rate: [0.0003, 0.0006]
+lora_r:
+  - 16
+  - 32
+lora_alpha:
+  - 16
+  - 32
+  - 64
+```
+
+
+
+### inference
+
+Runs inference using your trained model in CLI, interactive chat, or Gradio
+interface mode.
+
+```bash
+# CLI inference with LoRA
+axolotl inference config.yml --lora-model-dir="./outputs/lora-out"
+
+# CLI inference with full model
+axolotl inference config.yml --base-model="./completed-model"
+
+# Interactive multi-turn chat (see the inference guide for commands)
+axolotl inference config.yml --chat \
+    --lora-model-dir="./outputs/lora-out"
+
+# Gradio web interface
+axolotl inference config.yml --gradio \
+    --lora-model-dir="./outputs/lora-out"
+
+# Inference with input from file
+cat prompt.txt | axolotl inference config.yml \
+    --base-model="./completed-model"
+```
+
+### merge-lora
+
+Merges trained LoRA adapters into the base model.
+
+```bash
+# Basic merge
+axolotl merge-lora config.yml
+
+# Specify LoRA directory (usually used with checkpoints)
+axolotl merge-lora config.yml --lora-model-dir="./lora-output/checkpoint-100"
+
+# Merge using CPU (if out of GPU memory)
+CUDA_VISIBLE_DEVICES="" axolotl merge-lora config.yml
+```
+
+Configuration options:
+
+```yaml
+gpu_memory_limit: Limit GPU memory usage
+lora_on_cpu: Load LoRA weights on CPU
+```
+
+### merge-sharded-fsdp-weights
+
+Merges sharded FSDP model checkpoints into a single combined checkpoint.
+
+```bash
+# Basic merge
+axolotl merge-sharded-fsdp-weights config.yml
+```
+
+### evaluate
+
+Evaluates a model's performance (loss etc) on the train and eval datasets.
+
+```bash
+# Basic evaluation
+axolotl evaluate config.yml
+
+# Evaluation with launcher arguments
+axolotl evaluate config.yml --launcher torchrun -- --nproc_per_node=2
+```
+
+### lm-eval
+
+Runs LM Evaluation Harness on your model.
+
+```bash
+# Basic evaluation
+axolotl lm-eval config.yml
+```
+
+Configuration options:
+
+```yaml
+lm_eval_model: # model to evaluate (local or hf path)
+
+# List of tasks to evaluate
+lm_eval_tasks:
+  - arc_challenge
+  - hellaswag
+lm_eval_batch_size: # Batch size for evaluation
+output_dir: # Directory to save evaluation results
+```
+
+See [LM Eval Harness integration docs](https://docs.axolotl.ai/docs/custom_integrations.html#language-model-evaluation-harness-lm-eval) for full configuration details.
+
+### delinearize-llama4
+
+Delinearizes a Llama 4 linearized model into a regular HuggingFace Llama 4 model. This only works with the non-quantized linearized model.
+
+```bash
+axolotl delinearize-llama4 --model path/to/model_dir --output path/to/output_dir
+```
+
+This would be necessary to use with other frameworks. If you have an adapter, merge it with the non-quantized linearized model before delinearizing.
+
+### quantize
+
+Quantizes a model using the quantization configuration specified in your YAML file.
+
+```bash
+axolotl quantize config.yml
+```
+
+See [Quantization](./quantize.qmd) for more details.
+
+
+## Legacy CLI Usage
+
+While the new Click-based CLI is preferred, Axolotl still supports the legacy module-based CLI:
+
+```bash
+# Preprocess
+python -m axolotl.cli.preprocess config.yml
+
+# Train
+accelerate launch -m axolotl.cli.train config.yml
+
+# Inference
+accelerate launch -m axolotl.cli.inference config.yml \
+    --lora_model_dir="./outputs/lora-out"
+
+# Gradio interface
+accelerate launch -m axolotl.cli.inference config.yml \
+    --lora_model_dir="./outputs/lora-out" --gradio
+```
+
+::: {.callout-important}
+When overriding CLI parameters in the legacy CLI, use same notation as in yaml file (e.g., `--lora_model_dir`).
+
+**Note:** This differs from the new Click-based CLI, which uses dash notation (e.g., `--lora-model-dir`). Keep this in mind if you're referencing newer documentation or switching between CLI versions.
+:::
+
+## Remote Compute with Modal Cloud
+
+Axolotl supports running training and inference workloads on Modal cloud infrastructure. This is configured using a
+cloud YAML file alongside your regular Axolotl config.
+
+### Cloud Configuration
+
+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_count: 1    # Number of GPUs to use
+timeout: 86400  # Maximum runtime in seconds (24 hours)
+branch: main    # Git branch to use (optional)
+
+volumes:        # Persistent storage volumes
+  - name: axolotl-cache
+    mount: /workspace/cache
+  - name: axolotl-data
+    mount: /workspace/data
+  - name: axolotl-artifacts
+    mount: /workspace/artifacts
+
+secrets:        # Secrets to inject
+  - WANDB_API_KEY
+  - HF_TOKEN
+```
+
+### Running on Modal Cloud
+
+Commands that support the --cloud flag:
+
+```bash
+# Preprocess on cloud
+axolotl preprocess config.yml --cloud cloud_config.yml
+
+# Train on cloud
+axolotl train config.yml --cloud cloud_config.yml
+
+# Run lm-eval on cloud
+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
+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
+```
diff --git a/docs/config.qmd b/docs/config.qmd
deleted file mode 100644
index 570a173f9a..0000000000
--- a/docs/config.qmd
+++ /dev/null
@@ -1,453 +0,0 @@
----
-title: Config options
-description: A complete list of all configuration options.
----
-
-```yaml
-# This is the huggingface model that contains *.pt, *.safetensors, or *.bin files
-# This can also be a relative path to a model on disk
-base_model: ./llama-7b-hf
-# You can specify an ignore pattern if the model repo contains more than 1 model type (*.pt, etc)
-base_model_ignore_patterns:
-# If the base_model repo on hf hub doesn't include configuration .json files,
-# You can set that here, or leave this empty to default to base_model
-base_model_config: ./llama-7b-hf
-# You can specify to choose a specific model revision from huggingface hub
-revision_of_model:
-# Optional tokenizer configuration path in case you want to use a different tokenizer
-# than the one defined in the base model
-tokenizer_config:
-# If you want to specify the type of model to load, AutoModelForCausalLM is a good choice too
-model_type: AutoModelForCausalLM
-# Corresponding tokenizer for the model AutoTokenizer is a good choice
-tokenizer_type: AutoTokenizer
-# Trust remote code for untrusted source
-trust_remote_code:
-# use_fast option for tokenizer loading from_pretrained, default to True
-tokenizer_use_fast:
-# Whether to use the legacy tokenizer setting, defaults to True
-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:
-
-# (Internal use only)
-# Used to identify which the model is based on
-is_falcon_derived_model:
-is_llama_derived_model:
-is_qwen_derived_model:
-# Please note that if you set this to true, `padding_side` will be set to "left" by default
-is_mistral_derived_model:
-
-# optional overrides to the base model configuration
-overrides_of_model_config:
-  # RoPE Scaling https://github.com/huggingface/transformers/pull/24653
-  rope_scaling:
-    type: # linear | dynamic
-    factor: # float
-
-# optional overrides to the bnb 4bit quantization configuration
-# https://huggingface.co/docs/transformers/main/main_classes/quantization#transformers.BitsAndBytesConfig
-bnb_config_kwargs:
-  # These are default values
-  llm_int8_has_fp16_weight: false
-  bnb_4bit_quant_type: nf4
-  bnb_4bit_use_double_quant: true
-
-
-# Whether you are training a 4-bit GPTQ quantized model
-gptq: true
-
-# This will attempt to quantize the model down to 8 bits and use adam 8 bit optimizer
-load_in_8bit: true
-# Use bitsandbytes 4 bit
-load_in_4bit:
-
-# Use CUDA bf16
-bf16: true # bool or 'full' for `bf16_full_eval`. require >=ampere
-# Use CUDA fp16
-fp16: true
-# Use CUDA tf32
-tf32: true # require >=ampere
-
-# No AMP (automatic mixed precision)
-bfloat16: true # require >=ampere
-float16: true
-
-# Limit the memory for all available GPUs to this amount (if an integer, expressed in gigabytes); default: unset
-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
-
-# 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
-  - path: vicgalle/alpaca-gpt4
-  # The type of prompt to use for training. [alpaca, sharegpt, gpteacher, oasst, reflection]
-    type: alpaca # format | format:<prompt_style> (chat/instruct) | <prompt_strategies>.load_<load_fn>
-    ds_type: # Optional[str] (json|arrow|parquet|text|csv) defines the datatype when path is a file
-    data_files: # Optional[str] path to source data files
-    shards: # Optional[int] number of shards to split data into
-    name: # Optional[str] name of dataset configuration to load
-    train_on_split: train # Optional[str] name of dataset split to load from
-
-    # Optional[str] fastchat conversation type, only used with type: sharegpt
-    conversation: # Options (see Conversation 'name'): https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
-    field_human: # Optional[str]. Human key to use for conversation.
-    field_model: # Optional[str]. Assistant key to use for conversation.
-    # Add additional keys from your dataset as input or output roles
-    roles:
-      input: # Optional[List[str]]. These will be masked based on train_on_input
-      output: # Optional[List[str]].
-
-  # Custom user instruction prompt
-  - path: repo
-    type:
-      # The below are defaults. only set what's needed if you use a different column name.
-      system_prompt: ""
-      system_format: "{system}"
-      field_system: system
-      field_instruction: instruction
-      field_input: input
-      field_output: output
-
-      # Customizable to be single line or multi-line
-      # Use {instruction}/{input} as key to be replaced
-      # 'format' can include {input}
-      format: |-
-        User: {instruction} {input}
-        Assistant:
-      # 'no_input_format' cannot include {input}
-      no_input_format: "{instruction} "
-
-      # For `completion` datsets only, uses the provided field instead of `text` column
-      field:
-
-# If false, the datasets will not be shuffled and will keep their original order in `datasets`.
-# The same applies to the `test_datasets` option and the `pretraining_dataset` option. Default is true.
-shuffle_merged_datasets: true
-
-# A list of one or more datasets to eval the model with.
-# You can use either test_datasets, or val_set_size, but not both.
-test_datasets:
-  - path: /workspace/data/eval.jsonl
-    ds_type: json
-    # You need to specify a split. For "json" datasets the default split is called "train".
-    split: train
-    type: completion
-    data_files:
-      - /workspace/data/eval.jsonl
-
-# use RL training: 'dpo', 'ipo', 'kto_pair'
-rl:
-
-# Saves the desired chat template to the tokenizer_config.json for easier inferencing
-# Currently supports chatml and inst (mistral/mixtral)
-chat_template: chatml
-# Changes the default system message
-default_system_message: You are a helpful assistant. Please give a long and detailed answer. # Currently only supports chatml.
-# Axolotl attempts to save the dataset as an arrow after packing the data together so
-# subsequent training attempts load faster, relative path
-dataset_prepared_path: data/last_run_prepared
-# Push prepared dataset to hub
-push_dataset_to_hub: # 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
-# Keep dataset in memory while preprocessing
-# Only needed if cached dataset is taking too much storage
-dataset_keep_in_memory:
-# push checkpoints to hub
-hub_model_id: # private repo path to push finetuned model
-# how to push checkpoints to hub
-# https://huggingface.co/docs/transformers/v4.31.0/en/main_classes/trainer#transformers.TrainingArguments.hub_strategy
-hub_strategy:
-# Whether to use hf `use_auth_token` for loading datasets. Useful for fetching private datasets
-# Required to be true when used in combination with `push_dataset_to_hub`
-hf_use_auth_token: # boolean
-# How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc. 0 for no eval.
-val_set_size: 0.04
-# Num shards for whole dataset
-dataset_shard_num:
-# Index of shard to use for whole dataset
-dataset_shard_idx:
-
-# The maximum length of an input to train with, this should typically be less than 2048
-# as most models have a token/context limit of 2048
-sequence_len: 2048
-# Pad inputs so each step uses constant sized buffers
-# This will reduce memory fragmentation and may prevent OOMs, by re-using memory more efficiently
-pad_to_sequence_len:
-# Use efficient multi-packing with block diagonal attention and per sequence position_ids. Recommend set to 'true'
-sample_packing:
-# Set to 'false' if getting errors during eval with sample_packing on.
-eval_sample_packing:
-# You can set these packing optimizations AFTER starting a training at least once.
-# The trainer will provide recommended values for these values.
-sample_packing_eff_est:
-total_num_tokens:
-
-# Passed through to transformers when loading the model when launched without accelerate
-# Use `sequential` when training w/ model parallelism to limit memory
-device_map:
-# Defines the max memory usage per gpu on the system. Passed through to transformers when loading the model.
-max_memory:
-
-# If you want to use 'lora' or 'qlora' or leave blank to train all parameters in original model
-adapter: lora
-# If you already have a lora model trained that you want to load, put that here.
-# This means after training, if you want to test the model, you should set this to the value of `output_dir`.
-# Note that if you merge an adapter to the base model, a new subdirectory `merged` will be created under the `output_dir`.
-lora_model_dir:
-
-# LoRA hyperparameters
-# For more details about the following options, see:
-# https://www.anyscale.com/blog/fine-tuning-llms-lora-or-full-parameter-an-in-depth-analysis-with-llama-2
-lora_r: 8
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules:
-  - q_proj
-  - v_proj
-#  - k_proj
-#  - o_proj
-#  - gate_proj
-#  - down_proj
-#  - up_proj
-lora_target_linear: # If true, will target all linear modules
-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.
-# `embed_tokens` converts tokens to embeddings, and `lm_head` converts embeddings to token probabilities.
-# https://github.com/huggingface/peft/issues/334#issuecomment-1561727994
-lora_modules_to_save:
-#  - embed_tokens
-#  - lm_head
-
-lora_fan_in_fan_out: false
-
-# LoRA+ hyperparameters
-# For more details about the following options, see:
-# https://arxiv.org/abs/2402.12354  and `src/axolotl/core/train_builder.py`
-loraplus_lr_ratio: # loraplus learning rate ratio lr_B / lr_A. Recommended value is 2^4.
-loraplus_lr_embedding: #  loraplus learning rate for lora embedding layers. Default value is 1e-6.
-
-peft:
-  # Configuration options for loftq initialization for LoRA
-  # https://huggingface.co/docs/peft/developer_guides/quantization#loftq-initialization
-  loftq_config:
-    loftq_bits:  # typically 4 bits
-
-# ReLoRA configuration
-# Must use either 'lora' or 'qlora' adapter, and does not support fsdp or deepspeed
-relora_steps: # Number of steps per ReLoRA restart
-relora_warmup_steps: # Number of per-restart warmup steps
-relora_anneal_steps: # Number of anneal steps for each relora cycle
-relora_prune_ratio: # threshold for optimizer magnitude when pruning
-relora_cpu_offload: # True to perform lora weight merges on cpu during restarts, for modest gpu memory savings
-
-# wandb configuration if you're using it
-# Make sure your `WANDB_API_KEY` environment variable is set (recommended) or you login to wandb with `wandb login`.
-wandb_mode: # "offline" to save run metadata locally and not sync to the server, "disabled" to turn off wandb
-wandb_project: # Your wandb project name
-wandb_entity: # A wandb Team name if using a Team
-wandb_watch:
-wandb_name: # Set the name of your wandb run
-wandb_run_id: # Set the ID of your wandb run
-wandb_log_model: # "checkpoint" to log model to wandb Artifacts every `save_steps` or "end" to log only at the end of training
-
-# mlflow configuration if you're using it
-mlflow_tracking_uri: # URI to mlflow
-mlflow_experiment_name: # Your experiment name
-hf_mlflow_log_artifacts:  # set to true to copy each saved checkpoint on each save to mlflow artifact registry
-
-# Where to save the full-finetuned model to
-output_dir: ./completed-model
-
-# Whether to use torch.compile and which backend to use
-torch_compile:  # bool
-torch_compile_backend:  # Optional[str]
-
-# Training hyperparameters
-
-# If greater than 1, backpropagation will be skipped and the gradients will be accumulated for the given number of steps.
-gradient_accumulation_steps: 1
-# The number of samples to include in each batch. This is the number of samples sent to each GPU.
-# Batch size per gpu = micro_batch_size * gradient_accumulation_steps
-micro_batch_size: 2
-eval_batch_size:
-num_epochs: 4
-warmup_steps: 100  # cannot use with warmup_ratio
-warmup_ratio: 0.05  # cannot use with warmup_steps
-learning_rate: 0.00003
-lr_quadratic_warmup:
-logging_steps:
-eval_steps: # Leave empty to eval at each epoch, integers for every N steps. decimal for fraction of total steps
-evals_per_epoch: # number of times per epoch to run evals, mutually exclusive with eval_steps
-save_strategy: # Set to `no` to skip checkpoint saves
-save_steps: # Leave empty to save at each epoch
-saves_per_epoch: # number of times per epoch to save a checkpoint, mutually exclusive with save_steps
-save_total_limit: # Checkpoints saved at a time
-# Maximum number of iterations to train for. It precedes num_epochs which means that
-# if both are set, num_epochs will not be guaranteed.
-# e.g., when 1 epoch is 1000 steps => `num_epochs: 2` and `max_steps: 100` will train for 100 steps
-max_steps:
-
-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
-eval_causal_lm_metrics: # HF evaluate metrics used during evaluation. Default is ["sacrebleu", "comet", "ter", chrf]
-
-loss_watchdog_threshold: # High loss value, indicating the learning has broken down (a good estimate is ~2 times the loss at the start of training)
-loss_watchdog_patience: # Number of high-loss steps in a row before the trainer aborts (default: 3)
-
-# Save model as safetensors (require safetensors package)
-save_safetensors:
-
-# Whether to mask out or include the human's prompt from the training labels
-train_on_inputs: false
-# Group similarly sized data to minimize padding.
-# May be slower to start, as it must download and sort the entire dataset.
-# Note that training loss may have an oscillating pattern with this enabled.
-group_by_length: false
-
-# 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:
-#   use_reentrant: true
-
-# Stop training after this many evaluation losses have increased in a row
-# https://huggingface.co/transformers/v4.2.2/_modules/transformers/trainer_callback.html#EarlyStoppingCallback
-early_stopping_patience: 3
-
-# Specify a scheduler and kwargs to use with the optimizer
-lr_scheduler: # 'one_cycle' | 'log_sweep' | empty for cosine
-lr_scheduler_kwargs:
-cosine_min_lr_ratio: # decay lr to some percentage of the peak lr, e.g. cosine_min_lr_ratio=0.1 for 10% of peak lr
-cosine_constant_lr_ratio: # freeze lr at some percentage of the step, e.g. cosine_constant_lr_ratio=0.8 means start cosine_min_lr at 80% of training step (https://arxiv.org/pdf/2308.04014.pdf)
-
-# For one_cycle optim
-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/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_apex_fused
-# - adafactor
-# - adamw_anyprecision
-# - sgd
-# - adagrad
-# - adamw_bnb_8bit
-# - lion_8bit
-# - lion_32bit
-# - paged_adamw_32bit
-# - paged_adamw_8bit
-# - paged_lion_32bit
-# - paged_lion_8bit
-# - galore_adamw
-# - galore_adamw_8bit
-# - galore_adafactor
-# - galore_adamw_layerwise
-# - galore_adamw_8bit_layerwise
-# - galore_adafactor_layerwise
-optimizer:
-# Dictionary of arguments to pass to the optimizer
-optim_args:
-# For Galore Optimizers the following optim_args are available
-# rank:  # type: int
-# update_proj_gap  # type: int
-# scale  # type: float
-# proj_type:  # type: str, default = std
-
-# The target modules to optimize, i.e. the module names that you would like to train, right now this is used only for GaLore algorithm
-optim_target_modules:
-# - self_attn  # for llama
-# - mlp
-
-# Specify weight decay
-weight_decay:
-# adamw hyperparams
-adam_beta1:
-adam_beta2:
-adam_epsilon:
-# Gradient clipping max norm
-max_grad_norm:
-
-# Augmentation techniques
-# NEFT https://arxiv.org/abs/2310.05914, set this to a number (paper default is 5) to add noise to embeddings
-# currently only supported on Llama and Mistral
-neftune_noise_alpha:
-
-# Whether to bettertransformers
-flash_optimum:
-# Whether to use xformers attention patch https://github.com/facebookresearch/xformers:
-xformers_attention:
-# Whether to use flash attention patch https://github.com/Dao-AILab/flash-attention:
-flash_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:
-# Shifted-sparse attention (only llama) - https://arxiv.org/pdf/2309.12307.pdf
-s2_attention:
-# Resume from a specific checkpoint dir
-resume_from_checkpoint:
-# If resume_from_checkpoint isn't set and you simply want it to start where it left off.
-# Be careful with this being turned on between different models.
-auto_resume_from_checkpoints: false
-
-# Don't mess with this, it's here for accelerate and torchrun
-local_rank:
-
-# Add or change special tokens.
-# If you add tokens here, you don't need to add them to the `tokens` list.
-special_tokens:
-  # bos_token: "<s>"
-  # eos_token: "</s>"
-  # unk_token: "<unk>"
-  # pad_token: "[PAD]"
-
-# Add extra tokens.
-tokens:
-
-# FSDP
-fsdp:
-fsdp_config:
-
-# Deepspeed config path. e.g., deepspeed_configs/zero3.json
-deepspeed:
-
-# Advanced DDP Arguments
-ddp_timeout:
-ddp_bucket_cap_mb:
-ddp_broadcast_buffers:
-
-# Path to torch distx for optim 'adamw_anyprecision'
-torchdistx_path:
-
-# Set to HF dataset for type: 'completion' for streaming instead of pre-tokenize
-pretraining_dataset:
-
-# Debug mode
-debug:
-
-# Seed
-seed:
-
-# Allow overwrite yml config using from cli
-strict:
-```
diff --git a/docs/custom_integrations.qmd b/docs/custom_integrations.qmd
new file mode 100644
index 0000000000..8e1fdaa2e8
--- /dev/null
+++ b/docs/custom_integrations.qmd
@@ -0,0 +1,121 @@
+---
+title: Custom Integrations
+toc: true
+toc-depth: 3
+---
+
+```{python}
+#| echo: false
+
+import os
+import re
+
+def process_readme(integration_name):
+    try:
+        path = f'../src/axolotl/integrations/{integration_name}/README.md'
+        with open(path, 'r') as f:
+            txt = f.read()
+            # Remove h1 headings
+            txt = re.sub(r'^# .*\n?', '', txt, flags=re.MULTILINE)
+            # Convert h2 to h3
+            txt = re.sub(r'^## ', '### ', txt, flags=re.MULTILINE)
+            return txt
+    except FileNotFoundError:
+        return None
+
+def print_section(name, folder_name):
+    output = f"\n## {name}\n"
+    content = process_readme(folder_name)
+    if content:
+        output += content
+    output += f"\nPlease see reference [here](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/{folder_name})\n"
+    return output
+```
+
+```{python}
+#| output: asis
+#| echo: false
+
+# Introduction text
+print("""
+Axolotl adds custom features through `integrations`. They are located within the `src/axolotl/integrations` directory.
+
+To enable them, please check the respective documentations.
+""")
+
+# Sections
+sections = [
+    ("Cut Cross Entropy", "cut_cross_entropy"),
+    ("Grokfast", "grokfast"),
+    ("Knowledge Distillation (KD)", "kd"),
+    ("Liger Kernels", "liger"),
+    ("Language Model Evaluation Harness (LM Eval)", "lm_eval"),
+    ("Spectrum", "spectrum"),
+    ("LLMCompressor", "llm_compressor")
+]
+
+for folder_name in os.listdir("../src/axolotl/integrations/"):
+    if folder_name in [path for name, path in sections]:
+        # skip if already in sections
+        continue
+    if os.path.exists(f"../src/axolotl/integrations/{folder_name}/README.md"):
+        # grab the first heading in README.md as the section name
+        with open(f"../src/axolotl/integrations/{folder_name}/README.md", "r") as f:
+            txt = f.read()
+            matches = re.search(r'^# (.*)\n?', txt, flags=re.MULTILINE)
+            if matches:
+                name = matches.group(1)
+            else:
+                continue
+            sections.append((name, folder_name))
+
+# sort sections by name
+sections = sorted(sections, key=lambda x: x[0])
+
+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)
+
+:::
diff --git a/docs/dataset-formats/conversation.qmd b/docs/dataset-formats/conversation.qmd
index f7d0cac826..adbb88645d 100644
--- a/docs/dataset-formats/conversation.qmd
+++ b/docs/dataset-formats/conversation.qmd
@@ -4,60 +4,449 @@ description: Conversation format for supervised fine-tuning.
 order: 3
 ---
 
-## sharegpt
+## chat_template
 
-conversations where `from` is `human`/`gpt`. (optional: first row with role `system` to override default system prompt)
+Chat Template strategy uses a jinja2 template that converts a list of messages into a prompt. Support using tokenizer's template, a supported template, or custom jinja2.
 
 ```{.json filename="data.jsonl"}
-{"conversations": [{"from": "...", "value": "..."}]}
+{"messages": [{"role": "...", "content": "..."}, {"role": "...", "content": "..."}, ...]}
 ```
 
-Note: `type: sharegpt` opens special configs:
-- `conversation`: enables conversions to many Conversation types. Refer to the 'name' [here](https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py) for options.
-- `roles`: allows you to specify the roles for input and output. This is useful for datasets with custom roles such as `tool` etc to support masking.
-- `field_human`: specify the key to use instead of `human` in the conversation.
-- `field_model`: specify the key to use instead of `gpt` in the conversation.
+See [configs](../config-reference.qmd) for full configs and supported templates.
+
+### Migrating from sharegpt
+
+Most configs can be adapted as follows:
 
 ```yaml
+# old
+chat_template: chatml
 datasets:
-    path: ...
+  - path: ...
     type: sharegpt
+    conversation: chatml
+
+# new (if using tokenizer's chat_template)
+datasets:
+  - path: ...
+    type: chat_template
 
-    conversation: # Options (see Conversation 'name'): https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
-    field_human: # Optional[str]. Human key to use for conversation.
-    field_model: # Optional[str]. Assistant key to use for conversation.
-    # Add additional keys from your dataset as input or output roles
-    roles:
-      input: # Optional[List[str]]. These will be masked based on train_on_input
-      output: # Optional[List[str]].
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+# new (if setting a new chat_template like chatml, gemma, etc)
+chat_template: chatml
+datasets:
+  - path: ...
+    type: chat_template
+
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
 ```
 
-## pygmalion
+We recommend checking the below examples for other usecases.
+
+### Examples
+
+#### Training on last message
+
+(Legacy) Using the default chat template in the tokenizer_config.json on OpenAI messages format, training on only last message.
+
+```yaml
+datasets:
+  - path: ...
+    type: chat_template
+    roles_to_train:
+    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`.
+:::
+
+#### Overriding default chat template
+
+Using the `gemma` chat template to override the tokenizer_config.json's chat template on OpenAI messages format, training on all assistant messages.
+
+```yaml
+chat_template: gemma # this overwrites the tokenizer's chat_template
+datasets:
+  - path: ...
+    type: chat_template
+    roles_to_train: ["assistant"]  # default value
+```
+
+::: {.callout-note}
+If you want to use built-in chat_template, use `chat_template: tokenizer_default` (this is set by default).
+:::
+
+#### Using default chat template with fallback
+
+Using the tokenizer_config.json's chat template or `chatml` as fallback if the former's chat template does not exist, on OpenAI messages format, training on all assistant messages.
+
+```yaml
+chat_template: tokenizer_default_fallback_chatml # this overwrites the tokenizer's chat_template
+datasets:
+  - path: ...
+    type: chat_template
+```
+
+#### Custom Jinja template
+
+Using a custom jinja template on OpenAI messages format, training on all assistant messages.
+
+```yaml
+# chat_template: jinja # `jinja` will be implied if the `chat_template_jinja` is set and this field is empty
+chat_template_jinja: "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'system') %}{{'<|system|>' + '\n' + message['content'] + '<|end|>' + '\n'}}{% elif (message['role'] == 'user') %}{{'<|user|>' + '\n' + message['content'] + '<|end|>' + '\n' + '<|assistant|>' + '\n'}}{% elif message['role'] == 'assistant' %}{{message['content'] + '<|end|>' + '\n'}}{% endif %}{% endfor %}"
+
+datasets:
+  - path: ...
+    type: chat_template
+```
+
+::: {.callout-tip}
+`chat_template_jinja` also accepts a file path to a `.jinja2` file instead of an inline string:
+
+```yaml
+chat_template_jinja: ./path/to/my_template.jinja2
+```
+:::
+
+::: {.callout-important}
+Please make sure that your `tokenizer.eos_token` is same as EOS (End-of-Sequence) token in template. Otherwise, set `eos_token` under `special_tokens: `.
+:::
+
+#### Using template with different token for EOT and EOS
+
+- If you are using a template that has a different EOT (End-of-Turn) token from EOS token or multiple EOT tokens (like Mistral V7 Tekken), set the `eot_tokens: ` config. The handling of EOT tokens follows `train_on_eos: ` which defaults to turn.
+
+```yaml
+eot_tokens:
+  - "[/INST]"
+  # - "[/SYSTEM_PROMPT]"
+
+datasets:
+  - path: ...
+    type: chat_template
+
+    # optional
+    train_on_eot: turn  # defaults read from train_on_eos (which defaults to turn)
+```
+
+::: {.callout-tip}
+See [config documentation](../config-reference.qmd) for detailed explanations of "turn", "last", and "all" options for training on tokens.
+:::
+
+::: {.callout-note}
+Using `eot_tokens` requires each token that exists in `chat_template` to be a single token in the tokenizer. Otherwise, the tokenizer will split the token and cause unexpected behavior.
+
+You can add those tokens as new tokens under `tokens: ` or (recommended) override unused added_tokens via `added_tokens_overrides: `. See [config](../config-reference.qmd) for more details.
+:::
+
+- Continuing from the previous example, if you want to train on all EOT token trainable turns but only last EOS token, set `train_on_eos: last`.
+
+```yaml
+eot_tokens:
+  - "[/INST]"
+  # ...
+
+datasets:
+  - path: ...
+    type: chat_template
+
+    train_on_eos: last
+    train_on_eot: turn
+```
+
+::: {.callout-tip}
+If EOS token only appears at the end of a prompt, `train_on_eos: last` is equivalent to `train_on_eos: turn`. Therefore, generally, you can leave them to their defaults and omit them.
+:::
+
+
+#### Using tool use
+
+Instead of passing `tools` via the system prompt, an alternative method would be to have the `tools` in a separate column and loaded via `chat_template` to let the template dynamically build it.
+
+```json
+{
+    "tools": [
+        {
+            "type": "...",
+            "function": {
+                "name": "...",
+                "description": "...",
+                "parameters": {
+                    "type": "...",
+                    "properties": {
+                        // ...
+                    },
+                    "required": ["..."],
+                },
+            },
+        },
+    ],
+    "messages": [
+        // ...
+        {
+            "role": "assistant", // call the function via assistant
+            "tool_calls": [
+                {
+                    "id": "...",  // required only for mistral
+                    "type": "function",
+                    "function": {
+                        "name": "...",
+                        "arguments": {
+                            "...": "...",
+                        }
+                    }
+                }
+            ]
+        },
+        {
+            "role": "tool",
+            "tool_call_id": "...",  // required only for mistral
+            "name": "...",
+            "content": "..."
+        },
+    ],
+}
+```
+
+::: {.callout-note}
+Tools need to follow [JSON schema](https://json-schema.org/learn/getting-started-step-by-step).
+:::
+
+::: {.callout-warning}
+If you have tool arguments with same name but different dtypes (like `"time": string` and `"time": number`), please save `arguments: ` as JSON string to prevent `datasets` from having casting issues.
+
+```
+"arguments": "{\"...\": \"...\"}"
+```
+
+The same is applicable for tool parameters.
+
+```
+"parameters": "{\"...\": \"...\"}"
+```
+
+:::
+
+Example config for Llama4:
+```yaml
+chat_template: llama4
+datasets:
+  - path: Nanobit/text-tools-2k-test
+    type: chat_template
+    # field_tools: tools # default is `tools`
+```
+
+::: {.callout-tip}
+Look into the `chat_template` you are using to see if it supports `tools` and what the expected role is for the tool answer. In the example above, the tool answer is expected to be in the `tool` or `ipython` role for `llama4` template.
+:::
+
+
+#### Using fine-grained control over token masking
+
+(Advanced) Using fine-grained control over tokens and turns to train in a conversation
+
+For a data sample that looks like:
 
 ```{.json filename="data.jsonl"}
-{"conversations": [{"role": "...", "value": "..."}]}
+{
+  "conversations": [
+    {"from": "system", "value": "You are an AI assistant.", "train": false},
+    {"from": "human", "value": "Hello", "train": false},
+    {"from": "assistant", "value": "Hello", "train": true},
+    {"from": "human", "value": "How are you?", "train": true},
+    {
+      "from": "assistant",
+      "value": "I'm doing very well, thank you!",
+      "train_detail": [
+        {"begin_offset": 0, "end_offset": 8, "train": false},
+        {"begin_offset": 9, "end_offset": 18, "train": true},
+        {"begin_offset": 19, "end_offset": 30, "train": false},
+      ],
+    },
+    {
+        "from": "human",
+        "value": "I'm doing very well, thank you!",
+        "train": true,
+    },
+    {"from": "assistant", "value": "Hi there!", "train": true}
+  ]
+}
 ```
 
-## sharegpt.load_role
+The configuration would look like:
+
+```yaml
+datasets:
+  - path: ...
+    type: chat_template
+    chat_template: tokenizer_default
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+    roles_to_train: []
+    train_on_eos: turn
+    message_field_training: train
+    message_field_training_detail: train_detail
+```
 
-conversations where `role` is used instead of `from`
+::: {.callout-tip}
+It is not necessary to set both `message_field_training` and `message_field_training_detail` at once.
+:::
+
+#### Content parts with per-part training control
+
+Instead of using character offsets with `train_detail`, you can split a message's content into a list of parts, each with its own training flag. This is useful when you want to mask specific sections of a response (e.g., mask reasoning but train on the answer).
 
 ```{.json filename="data.jsonl"}
-{"conversations": [{"role": "...", "value": "..."}]}
+{
+  "messages": [
+    {"role": "user", "content": [{"type": "text", "text": "What is 2+2?"}]},
+    {
+      "role": "assistant",
+      "content": [
+        {"type": "text", "text": "Let me think step by step...", "train": false},
+        {"type": "text", "text": " The answer is 4.", "train": true}
+      ]
+    }
+  ]
+}
 ```
 
-## sharegpt.load_guanaco
+The configuration is the same as standard `chat_template` — no extra fields needed:
 
-conversations where `from` is `prompter` `assistant` instead of default sharegpt
+```yaml
+datasets:
+  - path: ...
+    type: chat_template
+    roles_to_train: ["assistant"]
+```
+
+Each content part supports:
+
+- `type`: `"text"` (required)
+- `text`: the text value (also accepts `content` or `value` as the key)
+- `train`: `true`/`false` (optional) — whether to train on this part
+- `weight`: `0`/`1` (optional) — alternative to `train`
+
+If a part has no `train` or `weight` flag, it inherits the turn-level training decision (from `roles_to_train`, `message_field_training`, or `train_on_inputs`).
+
+::: {.callout-warning title="Whitespace at part boundaries"}
+BPE tokenizers (used by Llama, Qwen, Mistral, GPT, etc.) prepend spaces to word tokens. For example, `" answer"` is a single token — the space is part of it. This means **where you place whitespace between content parts matters**:
+
+**Split BEFORE spaces** (space goes with the next part):
+
+```json
+[
+  {"type": "text", "text": "Let me think...", "train": false},
+  {"type": "text", "text": " The answer is 4.", "train": true}
+]
+```
+
+**DON'T put trailing spaces** on a part (the space merges with the next word into one token that straddles the boundary, and straddling tokens are masked):
+
+```json
+[
+  {"type": "text", "text": "Let me think... ", "train": false},
+  {"type": "text", "text": "The answer is 4.", "train": true}
+]
+```
+
+In the bad example, `" The"` becomes a single token that spans both parts. Because it straddles the boundary, it is conservatively **masked** (not trained) — even though the second part has `train: true`.
+
+**Newlines** typically merge with preceding punctuation (e.g., `":\n"` is one token). Keep newlines with the preceding part:
+
+```json
+[
+  {"type": "text", "text": "Thinking:\n", "train": false},
+  {"type": "text", "text": "The answer is 4.", "train": true}
+]
+```
+
+Axolotl will log a warning if it detects trailing whitespace at a boundary between parts with different training flags.
+:::
+
+::: {.callout-note}
+When all content parts in a message are strings, they are concatenated before being passed to the chat template. This means content parts work with **any** Jinja template — the template sees a plain string, and the per-part training flags are applied during tokenization.
+:::
+
+##### Per-part training on reasoning_content
+
+For templates that support a separate `reasoning_content` field (e.g., `qwen3`), the same content-parts format works on `reasoning_content`. This is useful for masking incorrect reasoning steps while training on self-corrections:
 
 ```{.json filename="data.jsonl"}
-{"conversations": [{"from": "...", "value": "..."}]}
+{
+  "messages": [
+    {"role": "user", "content": [{"type": "text", "text": "What is 2+2?"}]},
+    {
+      "role": "assistant",
+      "reasoning_content": [
+        {"type": "text", "text": "Hmm maybe 2+2=5.", "train": false},
+        {"type": "text", "text": " Wait no, 2+2=4.", "train": true}
+      ],
+      "content": [
+        {"type": "text", "text": "The answer is 4.", "train": true}
+      ]
+    }
+  ]
+}
 ```
 
-## sharegpt_jokes
+The `reasoning_content` and `content` fields are handled independently — each has its own token boundaries and per-part masking. No additional configuration is needed beyond what the template already requires.
+
+::: {.callout-tip}
+When `reasoning_content` is provided as a separate field, `split_thinking` is not needed — the reasoning is already separated from the content in the data.
+:::
 
-creates a chat where bot is asked to tell a joke, then explain why the joke is funny
+The same whitespace rules apply to `reasoning_content` parts as to `content` parts — split before spaces, keep newlines with the preceding part.
+
+
+#### Reasoning split
+
+(For Qwen3 template only) Enable reasoning split, where the reasoning is split from the content and passed as a separate field into the template.
+
+```yaml
+datasets:
+  - path: ...
+    type: chat_template
+    chat_template: qwen3
+    split_thinking: true
+```
+
+For example, a content can look like:
+
+```json
+{
+  "content": "<think>Some thinking outputs</think>Output after thinking."
+}
+```
+
+After split, it will look like:
+
+```json
+{
+  "reasoning_content": "Some thinking outputs",
+  "content": "Output after thinking..."
+}
+```
+
+
+## sharegpt
+
+::: {.callout-important}
+ShareGPT is deprecated!. Please see [chat_template](#chat_template) section.
+:::
+
+## pygmalion
 
 ```{.json filename="data.jsonl"}
-{"conversations": [{"title": "...", "text": "...", "explanation": "..."}]}
+{"conversations": [{"role": "...", "value": "..."}]}
 ```
diff --git a/docs/dataset-formats/index.qmd b/docs/dataset-formats/index.qmd
index 91873a4c19..7018e8c25f 100644
--- a/docs/dataset-formats/index.qmd
+++ b/docs/dataset-formats/index.qmd
@@ -1,14 +1,452 @@
 ---
 title: Dataset Formats
-description: Supported dataset formats.
-listing:
-  fields: [title, description]
-  type: table
-  sort-ui: false
-  filter-ui: false
-  max-description-length: 250
+description: Guide to Dataset Formats in Axolotl
+back-to-top-navigation: true
+toc: true
+toc-depth: 5
 ---
 
-Axolotl supports a variety of dataset formats.  It is recommended to use a JSONL format.  The schema of the JSONL depends upon the task and the prompt template you wish to use. Instead of a JSONL, you can also use a HuggingFace dataset with columns for each JSONL field.
 
-Below are these various formats organized by task:
+Axolotl is a training framework that aims to make the process convenient yet flexible to users by simply passing a config yaml file.
+
+As there are a lot of available options in Axolotl, this guide aims to provide an simplify the user experience to choosing the proper choice.
+
+Axolotl supports 3 kinds of training methods: pre-training, supervised fine-tuning, and preference-based post-training (e.g. DPO, ORPO, PRMs). Each method has their own dataset format which are described below.
+
+::: {.callout-tip}
+
+This guide will mainly use JSONL as an introduction. Please refer to the [dataset loading docs](../dataset_loading.qmd) to understand how to load datasets from other sources.
+
+For `pretraining_dataset:` specifically, please refer to the [Pre-training section](#pre-training).
+:::
+
+## Pre-training
+
+Pre-training trains on raw text corpora with no input masking. The dataset format is simple:
+
+```json
+{"text": "first row"}
+{"text": "second row"}
+```
+
+Axolotl supports two approaches:
+
+### Streaming (large datasets)
+
+For large corpora that don't fit in memory, use `pretraining_dataset` with [streaming](../streaming.qmd). Data is tokenized on-demand during training.
+
+```yaml
+pretraining_dataset:
+  - path: HuggingFaceFW/fineweb-edu
+    type: pretrain
+    text_column: text
+    split: train
+```
+
+::: {.callout-important}
+Streaming requires `max_steps` in your config — Axolotl cannot infer the dataset size. One step = `sequence_len * micro_batch_size * gradient_accumulation_steps * num_gpus` tokens.
+:::
+
+See [Streaming Datasets](../streaming.qmd) for full configuration details.
+
+### Non-streaming (smaller datasets)
+
+For datasets that fit in memory, use `type: completion` under `datasets:`. The entire dataset is pre-tokenized before training, which can be done on a CPU-only machine.
+
+```yaml
+datasets:
+  - path: my_corpus
+    type: completion
+```
+
+::: {.callout-note}
+With `completion`, texts exceeding `sequence_len` are split into multiple samples automatically.
+:::
+
+## Supervised fine-tuning (SFT)
+
+Supervised fine-tuning is the process of training models to respond to an instruction or chat input.
+
+As there are a wide variety of dataset formats, Axolotl tries to support a majority of the formats available in public datasets.
+
+Axolotl provides four approaches for loading datasets, however, it's easier to work backwards from the dataset you have available to figure out which approach to use.
+
+A flow chart is as follows:
+
+1. Do you already have the dataset tokenized? If yes, check [Pre-Tokenized Dataset](#pre-tokenized-dataset).
+
+2. Do you want to format the dataset yourself and manually choose each section to mask? If yes, check [Template Free Dataset](#template-free-dataset)
+
+3. Is your dataset in a "conversation" format, containing a `list[messages]`? If yes, check [Conversation Dataset](#conversation-dataset)
+
+4. Is your dataset in an "instruct" format, containing `{ instruction, response }`? If yes, check [Instruction Dataset](#instruction-dataset)
+
+If you went through the flow chart and did not find one that matches, it is recommended to preprocess your dataset into one of the above or create a thread on Github Discussion.
+
+::: {.callout-tip}
+You can mix and match within each approach or across approaches to train a model on a variety of datasets.
+:::
+
+### Pre-Tokenized Dataset
+
+We suggest this approach when you want to bring your own tokenized dataset.
+
+Axolotl expects the dataset to have three keys:
+
+- `input_ids`: from tokenizing formatted prompt
+- `attention_mask`: for masking padding. If you don't add padding, it would be equal to `len(input_ids) * [1]`
+- `labels`: this is the same as `input_ids`, however, if you want to mask certain tokens, you would set those indices to `-100`.
+
+::: {.callout-tip}
+Make sure to add BOS/EOS tokens to your prompt and mask it appropriately.
+:::
+
+A config for this would look like:
+
+```yaml
+datasets:
+  - path: A.jsonl
+    type:
+```
+
+::: {.callout-note}
+`type: ` is empty!
+:::
+
+Reference: [Pre-Tokenized Dataset Documentation](tokenized.qmd).
+
+### Template Free Dataset
+
+We reccomend this approach when you want granular control over the prompt formatting, special tokens, and masking, whilst letting Axolotl handle the tokenization. This is very useful if your dataset has unique prompts that differ across samples and where one single general template wouldn't suffice.
+
+In the example below, you could see that there is no proper structure. At the same time, it's very flexible as there are no constraints on how your prompt can look.
+
+```json
+{
+    "segments": [
+        {
+            "label": true,
+            "text": "<s>Hello\n"
+        },
+        {
+            "label": true,
+            "text": "hi there!. "
+        },
+        {
+            "label": false,
+            "text": "goodbye "
+        },
+        {
+            "label": true,
+            "text": "farewell</s>"
+        }
+    ]
+}
+```
+
+Each prompt must be have a key called `segments` which is a list of `{ text, label }`.
+
+```yaml
+datasets:
+  - path: A.jsonl
+    type: input_output
+```
+
+Reference: [Template Free Documentation](template_free.qmd).
+
+### Conversation Dataset
+
+`conversation` messages are a list of messages which usually contain a `role` and `content` key.
+
+::: {.callout-tip}
+Fun fact: Axolotl synonymously refers to "chat" messages as `conversation` messages due to how FastChat initially used this term to build a widely used [fastchat conversation](https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py) method for formatting chat messages prior to the creation of `chat_templates`.
+:::
+
+#### What are `chat_templates`?
+
+The current most popular and convenient method for inference is to use `chat_templates` for formatting prompts. Axolotl supports using `chat_templates` for training to ensure that the model performs in the same environment as in inference.
+
+Here's a quick rundown on `chat_template`: A `chat_template` is a Jinja2 template which formats a list of messages into a prompt.
+
+An example of a prompt formatted into a popular template called ChatML can be seen below:
+
+Single prompt (pretty-printed):
+```json
+{
+    "messages": [
+        {
+            "role": "user",
+            "content": "Hi"
+        },
+        {
+            "role": "assistant",
+            "content": "How can I help you?"
+        },
+        {
+            "role": "user",
+            "content": "Can you add 3+5?"
+        },
+        {
+            "role": "assistant",
+            "content": "The answer is 8."
+        }
+    ]
+}
+```
+
+The ChatML template is as follows:
+```jinja2
+{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}
+```
+
+The above prompt formatted into this template will result in:
+
+```
+<|im_start|>user
+Hi<|im_end|>
+<|im_start|>assistant
+How can I help you?<|im_end|>
+<|im_start|>user
+Can you add 3+5?<|im_end|>
+<|im_start|>assistant
+The answer is 8.<|im_end|>
+```
+
+By using delimiters (`<|im_start|>` and `<|im_end|>`), a prompt separates different speakers which helps the model identify which portion belongs to whom.
+
+#### Common Conversation Dataset formats
+
+Older conversation datasets with the following format are colloquially called `sharegpt` datasets.
+
+```json
+{"conversations": [{"from": "...", "value": "..."}]}
+```
+
+Newer conversation datasets usually follow the OpenAI format.
+
+```json
+{"messages": [{"role": "...", "content": "..."}]}
+```
+
+Axolotl supports both as well as allowing customization of any kind of key.
+
+#### Chat Template Usage
+
+To properly use this method, it is important to identify three things:
+
+1. Which `chat_template` would you use?
+
+2. What are the keys in your dataset, and what are the possible roles? For example, in OpenAI format, the keys would be `messages`, `role`, and `content`, respectively, whereas the possible roles are `system`, `user`, and `assistant`.
+
+3. What do you want to mask? For instance, only assistant messages, only last message, or nothing.
+
+##### Choosing a `chat_template`
+
+There are a lot of `chat_templates` out there. Axolotl supports the common ones: [supported chat templates](https://github.com/axolotl-ai-cloud/axolotl/blob/860609392184cf62a7e0ca676658b170e059ce6c/src/axolotl/utils/chat_templates.py#L17). For example, to use ChatML, it would be `chat_template: chatml`.
+
+However, it is also possible to use the already configured template within the tokenizer by specifying `chat_template: tokenizer_default`. If you want a fallback (in case some tokenizer does not have it pre-configured), you can do `chat_template: tokenizer_default_fallback_chatml` to fallback to the ChatML template if a tokenizer template was not found.
+
+One last but powerful approach is to bring your own template. This can be set via:
+
+```yaml
+chat_template_jinja: # your template
+```
+
+##### Setting `chat_template` dataset keys
+
+We currently default to OpenAI format for dataset keys, so if that's your current dataset format, there's nothing to do here.
+
+If your dataset format is different, here are the keys you should check (with their defaults):
+
+```yaml
+datasets:
+    ...
+    field_messages: messages  # this should point to the key containing the list of conversations
+    message_property_mappings:  # this is a mapping from keys in your dataset to keys in chat_template
+      role: role
+      content: content
+```
+
+In some `chat_templates` (e.g. [Gemma](https://huggingface.co/google/gemma-2b-it/blob/main/tokenizer_config.json#L1507)), the roles are hardcoded to `user` and `assistant`. Consequently, you may find it necessary to map the roles in your dataset to these above. We currently have some defaults that should work for common datasets, but if you get a `KeyError`, it would be necessary to add mapping for your roles. Here is an example of how it would look like:
+
+```yaml
+datasets:
+    ...
+    roles:
+      assistant:
+        - gpt
+        - model
+      user:
+        - human
+```
+
+In the example above, all `gpt` and `model` values are converted to `assistant`. All `human` values are converted to `user.`
+
+##### Handling masking
+
+The common use case for `chat_template` is for chat messages, therefore, it is common to mask all non-assistant messages. Assistant messages refer to the bot messages that you want the model to learn on.
+
+To train on all `assistant` messages, you would set the following configs.
+
+```yaml
+datasets:
+    ...
+    roles_to_train: ["assistant"]
+    train_on_eos: "turn"
+```
+
+The `train_on_eos` config means that it would mask all EOS tokens for turns that aren't assistant-turns. The other options are: `all` and `last` to choose which EOS to train on.
+
+Perhaps, you want to train on `assistant` and `narrator` roles, you can simply add `narrator` to the list of `roles_to_train`. You would also need to add it to the mapping of `roles` above.
+
+```yaml
+datasets:
+    ...
+    roles_to_train: ["assistant", "narrator"]
+    roles:
+      assistant:
+        - gpt
+        - model
+      user:
+        - human
+      narrator: ["narrator"]
+```
+
+::: {.callout-tip}
+As chat_templates may use hardcoded EOS/EOT tokens that are different from the tokenizer's EOS, it is highly recommended to set them. For example, `ChatML` uses `<|im_end|>` to end turns.
+
+```yaml
+special_tokens:
+  eos_token: <|im_end|>
+```
+
+:::
+
+##### Applying `chat_template`
+
+Once all the above steps are completed, you could combine all these configs together to form a bespoke configuration for your custom dataset.
+
+```yaml
+datasets:
+  - path: A.jsonl
+    type: chat_template
+
+    # step 1
+    chat_template: chatml
+
+    # step 2
+    field_messages: messages
+    message_property_mappings:
+      role: role
+      content: content
+
+    roles:
+      assistant:
+        - gpt
+        - model
+        - assistant
+      user:
+        - human
+        - user
+
+    # step 3
+    roles_to_train: ["assistant"]
+    train_on_eos: "turn"
+
+special_tokens:
+  eos_token: <|im_end|>
+```
+
+If this config were to be applied to the sample dataset above, the output would look as such (which can be retrieved via `axolotl preprocess config.yaml --debug`):
+
+```
+<|im_start|>(-100, 128256) user(-100, 882)
+(-100, 198) Hi(-100, 13347) <|im_end|>(-100, 128257)
+(-100, 198) <|im_start|>(-100, 128256) assistant(-100, 78191)
+(-100, 198) How(4438, 4438)  can(649, 649)  I(358, 358)  help(1520, 1520)  you(499, 499) ?(30, 30) <|im_end|>(128257, 128257)
+(-100, 198) <|im_start|>(-100, 128256) user(-100, 882)
+(-100, 198) Can(-100, 6854)  you(-100, 499)  add(-100, 923)  (-100, 220) 3(-100, 18) +(-100, 10) 5(-100, 20) ?(-100, 30) <|im_end|>(-100, 128257)
+(-100, 198) <|im_start|>(-100, 128256) assistant(-100, 78191)
+(-100, 198) The(791, 791)  answer(4320, 4320)  is(374, 374)  (220, 220) 8(23, 23) .(13, 13) <|im_end|>(128257, 128257)
+(-100, 198)
+```
+
+The first number refers to the label, the second refers to the `token_id`. For example, `-100` labels appear on non-assistant portions, meaning that they are masked during. For assistant portions, the label is the same as the `token_id`.
+
+::: {.callout-note}
+
+If during `preprocess`, there are a lot of warnings of `Could not find content __ boundary`, please check the FAQ section for [chat_templates](../faq.qmd#chat-templates).
+
+:::
+
+#### Reference
+
+Please see docs [here](conversation.qmd).
+
+### Instruction Dataset
+
+Instruction datasets are used to train instruction-following models and comprise a prompt, containing an instruction, and a single response. In contrast to chat datasets which may be multi-turn, instruct datasets are typically single-turn.
+
+An example is of a common format called Alpaca:
+```json
+{"instruction": "...", "input": "...", "output": "..."}
+```
+
+Using those keys, a prompt can be built based on it.
+```
+Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+
+### Instruction:
+{instruction}
+
+### Input:
+{input}
+
+### Response:
+{output}
+```
+
+This can be configured as such:
+```yaml
+datasets:
+  - path: A.jsonl
+    type: alpaca
+```
+
+Axolotl supports many kinds of instruction dataset. All of them can be found in the [Instruction Dataset Documentation](inst_tune.qmd) with their respective type and sample row format.
+
+#### Custom Instruct Prompt Format
+
+Due to the myriad possibilities of instruction formats, Axolotl allows customizing your own instruction format without having to dive into the code directly.
+
+In the example below, a sample row is used to output in `mistral_v1` format.
+```json
+{"input": "...", "output": "..."}
+```
+
+```yaml
+datasets:
+  - path: repo
+    type:
+      system_prompt: ""
+
+      field_system:
+      field_instruction: input
+      field_input:
+      field_output: output
+
+      # multi-line example with input
+      format: |-
+        [INST] {instruction} {input} [/INST]
+
+      # single-line example without input
+      no_input_format: "[INST] {instruction} [/INST]"
+```
+
+The config sets that the `field_instruction` is actually named `input`, and the `field_input` is empty as we don't have an `input` in this sample. Generally, `instruction` can be thought as the question to the model, and `input` as the additional information with `output` being the response. It is not necessary to have an `input` nor `system`. In the end, the most important part is to understand what format you want it to look like and how you can customize this to your use case.
+
+Reference: [Custom Instruct Prompt Format Documentation](inst_tune.qmd#how-to-add-custom-prompt-format).
+
+## Reinforcement Learning from Human Feedback (RLHF)
+
+As there are multiple RLHF methods with their own dataset requirements. Please see [RLHF documentation](../rlhf.qmd) for more detail.
diff --git a/docs/dataset-formats/inst_tune.qmd b/docs/dataset-formats/inst_tune.qmd
index d89c6adaf5..f5bd7ab8f1 100644
--- a/docs/dataset-formats/inst_tune.qmd
+++ b/docs/dataset-formats/inst_tune.qmd
@@ -186,4 +186,4 @@ datasets:
       no_input_format: "[INST] {instruction} [/INST]"
 ```
 
-See full config options under [here](../config.qmd).
+See full config options under [here](../config-reference.qmd).
diff --git a/docs/dataset-formats/pretraining.qmd b/docs/dataset-formats/pretraining.qmd
index bb591328e2..9803704713 100644
--- a/docs/dataset-formats/pretraining.qmd
+++ b/docs/dataset-formats/pretraining.qmd
@@ -4,23 +4,9 @@ description: Data format for a pre-training completion task.
 order: 1
 ---
 
-For pretraining, there is no prompt template or roles.  The only required field is `text`:
-
-```{.json filename="data.jsonl"}
-{"text": "first row"}
-{"text": "second row"}
-...
-```
-
-:::{.callout-note}
-
-### Streaming is recommended for large datasets
-
-Axolotl usually loads the entire dataset into memory. This will be challenging for large datasets. Use the following config to enable streaming:
-
-```{.yaml filename="config.yaml"}
-pretraining_dataset: # hf path only
-...
-```
+::: {.callout-note}
+Pre-training documentation has been consolidated:
 
+- **Streaming pretraining** (large datasets): See [Streaming Datasets](../streaming.qmd#pretraining-with-streaming)
+- **Non-streaming pretraining** (`type: completion`): See [Dataset Formats](index.qmd#pre-training)
 :::
diff --git a/docs/dataset-formats/stepwise_supervised.qmd b/docs/dataset-formats/stepwise_supervised.qmd
new file mode 100644
index 0000000000..2cec8e1bda
--- /dev/null
+++ b/docs/dataset-formats/stepwise_supervised.qmd
@@ -0,0 +1,26 @@
+---
+title: Stepwise Supervised Format
+description: Format for datasets with stepwise completions and labels
+order: 3
+---
+
+## Stepwise Supervised
+
+The stepwise supervised format is designed for chain-of-thought (COT) reasoning
+datasets where each example contains multiple completion steps and a preference label
+for each step.
+
+### Example
+
+Here's a simple example of a stepwise supervised dataset entry:
+
+```json
+{
+  "prompt": "Which number is larger, 9.8 or 9.11?",
+  "completions": [
+    "The fractional part of 9.8 is 0.8, while the fractional part of 9.11 is 0.11.",
+    "Since 0.11 is greater than 0.8, the number 9.11 is larger than 9.8."
+  ],
+  "labels": [true, false]
+}
+```
diff --git a/docs/dataset-formats/template_free.qmd b/docs/dataset-formats/template_free.qmd
index 5087d6a013..c75c5931e8 100644
--- a/docs/dataset-formats/template_free.qmd
+++ b/docs/dataset-formats/template_free.qmd
@@ -1,7 +1,239 @@
 ---
 title: Template-Free
 description: Construct prompts without a template.
+toc: true
+toc-depth: 3
 order: 4
 ---
 
-See [these docs](../input_output.qmd).
+## Background {#sec-background}
+
+### Masking Inputs {#masking-inputs}
+
+One of the most popular features of
+[axolotl](https://github.com/axolotl-ai-cloud/axolotl) is
+setting the following configuration value:
+
+
+```yaml
+train_on_inputs: false
+```
+
+If you declare a [dataset formats](https://github.com/axolotl-ai-cloud/axolotl?tab=readme-ov-file#dataset)
+such as `alpaca` or `chatml`, axolotl knows what is an input
+(i.e. human) vs. an output (i.e. the assistant) and masks the input
+labels so that your model can focus on predicting the outputs only.
+
+### You may not want prompt templates {#sec-you-may-not-want-prompt-templates}
+
+However, there are many situations where you don't want to use one of
+these formats or templates. This is because they can:
+
+-   Add unnecessary boilerplate to your prompts.
+-   Create artifacts like special delimiters `<|im_start|>` that can
+    quickly become footguns if you don't include them correctly at
+    inference time.
+-   Enforce a *chat* interface when you do not want one. Sometimes you
+    just want to fine-tune a model to a very specific task and do NOT
+    want multi-turn conversations, roles, etc.
+-   Limit you to only certain roles that the template allows.
+
+### The `input_output` format {#sec-the-inputoutput-format}
+
+You can construct your prompts without a template by using the
+`input_output` format, by setting `type: input_output` in your
+configuration file like this:
+
+**config.yml**
+
+```yaml
+train_on_inputs: false # Mask segments of your data
+datasets:
+  - path: output.jsonl
+    type: input_output  # use template free prompt construction
+```
+
+Unlike `type: completion`, which is also template-free,
+`type: input_output` allows you to mask segments of your text. More
+details on how this works are described below.
+
+## Usage {#sec-usage}
+
+This is how you can use the `input_output` format:
+
+### 1. Prepare Data {#sec-1-prepare-data}
+
+To use the `input_output` format, collect your data in the following
+format into a jsonl file (below is the first row from the file
+`output`.jsonl` pretty printed):
+
+```bash
+$ head -n1 output.jsonl | python -m json.tool
+```
+
+:::{.cell-output .cell-output-stdout}
+    {
+        "segments": [
+            {
+                "label": true,
+                "text": "<s>Hello\n"
+            },
+            {
+                "label": true,
+                "text": "hi there!. "
+            },
+            {
+                "label": false,
+                "text": "goodbye "
+            },
+            {
+                "label": true,
+                "text": "farewell</s>"
+            }
+        ]
+    }
+:::
+
+Set `label:false` when you want to mask a segment of text so that the
+model isn't trained on it. Some things to keep in mind:
+
+> [!IMPORTANT]
+> 1.  **EOS, BOS, spaces, newlines etc. are entirely up to you. Axolotl
+    concatenates all the segments as-is.** The tokenizer doesn't add
+    anything additional. Notice how I added spaces, newlines, `<s>`
+    (BOS), and `</s>` (EOS) myself.
+> 2.  Make sure you check the materialized output to validate that the
+    prompt is getting assembled how you like.
+
+### 2. Use `type: input_output` {#sec-2-use-type-inputoutput}
+
+Let's materialize data with our `output.jsonl` file by setting
+`type: input_output` in our axolotl config:
+
+```yaml
+# training_config.yaml
+base_model: mistralai/Mistral-7B-v0.1
+data_seed: 49
+seed: 49
+
+datasets:
+  - path: output.jsonl
+    type: input_output
+val_set_size: 0.1
+
+sequence_len: 896
+sample_packing: false
+
+micro_batch_size: 2
+gradient_accumulation_steps: 3
+eval_batch_size: 2
+num_epochs: 1
+learning_rate: 0.0002
+
+train_on_inputs: false
+special_tokens:
+  bos_token: "<s>"
+  eos_token: "</s>"
+  unk_token: "<unk>"
+```
+
+You can use the following command to materialize your data. The
+`--debug` flag will print the tokens, along with the labels so you can
+verify that the correct items are being ignored:
+
+```bash
+axolotl preprocess training_config.yaml --debug
+
+...
+[2024-03-05 23:36:46,969] [INFO] [axolotl.check_example_labels:35] [PID:607731] [RANK:0] <s>(1, 1) Hello(22557, 22557)
+(13, 13) hi(12014, 12014) there(736, 736) !(28808, 28808) .(28723, 28723) (28705, 28705) good(-100, 1179) bye(-100, 17664) (-100, 28705) fare(19111, 19111) well(5458, 5458) </s>(2, 2)
+
+```
+
+The format is `decoded_token`(`label`, `token_id`), for example,
+`<s>(1, 1)` means that the token is `<s>`, the label is `1` and the
+token_id is `1`. When the label is `-100` then that token is ignored for
+training.
+
+### 3. Check the prompts {#sec-3-check-the-prompts}
+
+Here is another way to check the materialized output:
+
+```python
+from transformers import AutoTokenizer
+from datasets import load_from_disk
+import yaml
+
+directory = !ls last_run_prepared/
+with open('training_config.yaml', 'r') as f:
+    cfg = yaml.safe_load(f)
+model_id = cfg['base_model']
+tok = AutoTokenizer.from_pretrained(model_id)
+ds = load_from_disk(f'last_run_prepared/{directory[0]}/')
+```
+
+```python
+>>> row = ds[0]
+>>> print(tok.decode(row['input_ids']))
+<s> Hello
+    hi there!.  goodbye  farewell</s>
+```
+
+We can check that the right tokens are ignored by comparing the labels
+to each token:
+
+```python
+import pandas as pd
+pd.DataFrame([{'token': tok.decode(i), 'label': l, 'id':i} for i,l in
+              zip(row['input_ids'], row['labels'])])
+```
+
+| token | label | id    |
+|-------|-------|-------|
+| 0     | \<s\> | 1     |
+| 1     | Hello | 22557 |
+| 2     | \\n   | 13    |
+| 3     | hi    | 12014 |
+| 4     | there | 736   |
+| 5     | !     | 28808 |
+| 6     | .     | 28723 |
+| 7     |       | 28705 |
+| 8     | good  | -100  |
+| 9     | bye   | -100  |
+| 10    |       | -100  |
+| 11    | fare  | 19111 |
+| 12    | well  | 5458  |
+| 13    | \</s\>| 2     |
+
+
+
+If we look at the input data, the above table seems correct! (The jsonl
+version is repeated below for reference):
+
+
+```bash
+$ head -n1 output.jsonl | python -m json.tool
+```
+
+:::{.cell-output .cell-output-stdout}
+    {
+        "segments": [
+            {
+                "label": true,
+                "text": "<s>Hello\n"
+            },
+            {
+                "label": true,
+                "text": "hi there!. "
+            },
+            {
+                "label": false,
+                "text": "goodbye "
+            },
+            {
+                "label": true,
+                "text": "farewell</s>"
+            }
+        ]
+    }
+:::
diff --git a/docs/dataset-formats/tokenized.qmd b/docs/dataset-formats/tokenized.qmd
index 8991a21109..61028cae7f 100644
--- a/docs/dataset-formats/tokenized.qmd
+++ b/docs/dataset-formats/tokenized.qmd
@@ -4,9 +4,25 @@ description: How to use a custom pre-tokenized dataset.
 order: 5
 ---
 
-- Do not pass a `type:` in your axolotl config.
+- Pass an empty `type:` in your axolotl config.
 - Columns in Dataset must be exactly `input_ids`, `attention_mask`, `labels`
+- To indicate that a token should be ignored during training, set its corresponding label to `-100`.
+- You must add BOS and EOS, and make sure that you are training on EOS by not setting its label to -100.
+- For pretraining, do not truncate/pad documents to the context window length.
+- For instruction training, documents must be truncated/padded as desired.
+
+Sample config:
 
 ```{.yaml filename="config.yml"}
-- path: ...
+datasets:
+  - path: /path/to/your/file.jsonl
+    ds_type: json
+    type:
+```
+
+Sample jsonl:
+
+```jsonl
+{"input_ids":[271,299,99],"attention_mask":[1,1,1],"labels":[271,-100,99]}
+{"input_ids":[87,227,8383,12],"attention_mask":[1,1,1,1],"labels":[87,227,8383,12]}
 ```
diff --git a/docs/dataset_loading.qmd b/docs/dataset_loading.qmd
new file mode 100644
index 0000000000..bcffe7f0f3
--- /dev/null
+++ b/docs/dataset_loading.qmd
@@ -0,0 +1,268 @@
+---
+title: Dataset Loading
+description: Understanding how to load datasets from different sources
+back-to-top-navigation: true
+toc: true
+toc-depth: 5
+---
+
+## Overview
+
+Datasets can be loaded in a number of different ways depending on the how it is saved (the extension of the file) and where it is stored.
+
+## Loading Datasets
+
+We use the `datasets` library to load datasets and a mix of `load_dataset` and `load_from_disk` to load them.
+
+You may recognize the similar named configs between `load_dataset` and the `datasets` section of the config file.
+
+```yaml
+datasets:
+  - path:
+    name:
+    data_files:
+    split:
+    revision:
+    trust_remote_code:
+```
+
+::: {.callout-tip}
+
+Do not feel overwhelmed by the number of options here. A lot of them are optional. In fact, the most common config to use would be `path` and sometimes `data_files`.
+
+:::
+
+This matches the API of [`datasets.load_dataset`](https://github.com/huggingface/datasets/blob/0b5998ac62f08e358f8dcc17ec6e2f2a5e9450b6/src/datasets/load.py#L1838-L1858), so if you're familiar with that, you will feel right at home.
+
+For HuggingFace's guide to load different dataset types, see [here](https://huggingface.co/docs/datasets/loading).
+
+For full details on the config, see [config-reference.qmd](config-reference.qmd).
+
+::: {.callout-note}
+
+You can set multiple datasets in the config file by more than one entry under `datasets`.
+
+```yaml
+datasets:
+  - path: /path/to/your/dataset
+  - path: /path/to/your/other/dataset
+```
+
+:::
+
+### Local dataset
+
+#### Files
+
+To load a JSON file, you would do something like this:
+
+```python
+from datasets import load_dataset
+
+dataset = load_dataset("json", data_files="data.json")
+```
+
+Which translates to the following config:
+
+```yaml
+datasets:
+  - path: data.json
+    ds_type: json
+```
+
+In the example above, it can be seen that we can just point the `path` to the file or directory along with the `ds_type` to load the dataset.
+
+This works for CSV, JSON, Parquet, and Arrow files.
+
+::: {.callout-tip}
+
+If `path` points to a file and `ds_type` is not specified, we will automatically infer the dataset type from the file extension, so you could omit `ds_type` if you'd like.
+
+:::
+
+#### Directory
+
+If you're loading a directory, you can point the `path` to the directory.
+
+Then, you have two options:
+
+##### Loading entire directory
+
+You do not need any additional configs.
+
+We will attempt to load in the following order:
+- datasets saved with `datasets.save_to_disk`
+- loading entire directory of files (such as with parquet/arrow files)
+
+```yaml
+datasets:
+  - path: /path/to/your/directory
+```
+
+##### Loading specific files in directory
+
+Provide `data_files` with a list of files to load.
+
+```yaml
+datasets:
+    # single file
+  - path: /path/to/your/directory
+    ds_type: csv
+    data_files: file1.csv
+
+    # multiple files
+  - path: /path/to/your/directory
+    ds_type: json
+    data_files:
+      - file1.jsonl
+      - file2.jsonl
+
+    # multiple files for parquet
+  - path: /path/to/your/directory
+    ds_type: parquet
+    data_files:
+      - file1.parquet
+      - file2.parquet
+
+```
+
+### HuggingFace Hub
+
+The method you use to load the dataset depends on how the dataset was created, whether a folder was uploaded directly or a HuggingFace Dataset was pushed.
+
+::: {.callout-note}
+
+If you're using a private dataset, you will need to enable the `hf_use_auth_token` flag in the root-level of the config file.
+
+:::
+
+#### Folder uploaded
+
+This would mean that the dataset is a single file or file(s) uploaded to the Hub.
+
+```yaml
+datasets:
+  - path: org/dataset-name
+    data_files:
+      - file1.jsonl
+      - file2.jsonl
+```
+
+#### HuggingFace Dataset
+
+This means that the dataset is created as a HuggingFace Dataset and pushed to the Hub via `datasets.push_to_hub`.
+
+```yaml
+datasets:
+  - path: org/dataset-name
+```
+
+::: {.callout-note}
+
+There are some other configs which may be required like `name`, `split`, `revision`, `trust_remote_code`, etc depending on the dataset.
+
+:::
+
+### Remote Filesystems
+
+Via the `storage_options` config under `load_dataset`, you can load datasets from remote filesystems like S3, GCS, Azure, and OCI.
+
+::: {.callout-warning}
+
+This is currently experimental. Please let us know if you run into any issues!
+
+:::
+
+The only difference between the providers is that you need to prepend the path with the respective protocols.
+
+```yaml
+datasets:
+    # Single file
+  - path: s3://bucket-name/path/to/your/file.jsonl
+
+    # Directory
+  - path: s3://bucket-name/path/to/your/directory
+```
+
+For directory, we load via `load_from_disk`.
+
+#### S3
+
+Prepend the path with `s3://`.
+
+The credentials are pulled in the following order:
+
+- `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_SESSION_TOKEN` environment variables
+- from the `~/.aws/credentials` file
+- for nodes on EC2, the IAM metadata provider
+
+::: {.callout-note}
+
+We assume you have credentials setup and not using anonymous access. If you want to use anonymous access, let us know! We may have to open a config option for this.
+
+:::
+
+Other environment variables that can be set can be found in [boto3 docs](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#using-environment-variables)
+
+#### GCS
+
+Prepend the path with `gs://` or `gcs://`.
+
+The credentials are loaded in the following order:
+
+- gcloud credentials
+- for nodes on GCP, the google metadata service
+- anonymous access
+
+#### Azure
+
+##### Gen 1
+
+Prepend the path with `adl://`.
+
+Ensure you have the following environment variables set:
+
+- `AZURE_STORAGE_TENANT_ID`
+- `AZURE_STORAGE_CLIENT_ID`
+- `AZURE_STORAGE_CLIENT_SECRET`
+
+##### Gen 2
+
+Prepend the path with `abfs://` or `az://`.
+
+Ensure you have the following environment variables set:
+
+- `AZURE_STORAGE_ACCOUNT_NAME`
+- `AZURE_STORAGE_ACCOUNT_KEY`
+
+Other environment variables that can be set can be found in [adlfs docs](https://github.com/fsspec/adlfs?tab=readme-ov-file#setting-credentials)
+
+#### OCI
+
+Prepend the path with `oci://`.
+
+It would attempt to read in the following order:
+
+- `OCIFS_IAM_TYPE`, `OCIFS_CONFIG_LOCATION`, and `OCIFS_CONFIG_PROFILE` environment variables
+- when on OCI resource, resource principal
+
+Other environment variables:
+
+- `OCI_REGION_METADATA`
+
+Please see the [ocifs docs](https://ocifs.readthedocs.io/en/latest/getting-connected.html#Using-Environment-Variables).
+
+### HTTPS
+
+The path should start with `https://`.
+
+```yaml
+datasets:
+  - path: https://path/to/your/dataset/file.jsonl
+```
+
+This must be publically accessible.
+
+## Next steps
+
+Now that you know how to load datasets, you can learn more on how to load your specific dataset format into your target output format [dataset formats docs](dataset-formats).
diff --git a/docs/dataset_preprocessing.qmd b/docs/dataset_preprocessing.qmd
index c99fce444e..245723e677 100644
--- a/docs/dataset_preprocessing.qmd
+++ b/docs/dataset_preprocessing.qmd
@@ -3,8 +3,11 @@ title: Dataset Preprocessing
 description: How datasets are processed
 ---
 
+## Overview
+
 Dataset pre-processing is the step where Axolotl takes each dataset you've configured alongside
-the (dataset format)[../dataset-formats/] and prompt strategies to:
+the [dataset format](dataset-formats) and prompt strategies to:
+
  - parse the dataset based on the *dataset format*
  - transform the dataset to how you would interact with the model based on the *prompt strategy*
  - tokenize the dataset based on the configured model & tokenizer
@@ -12,10 +15,12 @@ the (dataset format)[../dataset-formats/] and prompt strategies to:
 
 The processing of the datasets can happen one of two ways:
 
-1. Before kicking off training by calling `python -m axolotl.cli.preprocess /path/to/your.yaml --debug`
+1. Before kicking off training by calling `axolotl preprocess config.yaml --debug`
 2. When training is started
 
-What are the benefits of pre-processing? When training interactively or for sweeps
+### What are the benefits of pre-processing?
+
+When training interactively or for sweeps
 (e.g. you are restarting the trainer often), processing the datasets can oftentimes be frustratingly
 slow. Pre-processing will cache the tokenized/formatted datasets according to a hash of dependent
 training parameters so that it will intelligently pull from its cache when possible.
@@ -28,8 +33,12 @@ default path of `./last_run_prepared/`, but will ignore anything already cached
 setting `dataset_prepared_path: ./last_run_prepared`, the trainer will use whatever pre-processed
 data is in the cache.
 
-What are the edge cases? Let's say you are writing a custom prompt strategy or using a user-defined
+### What are the edge cases?
+
+Let's say you are writing a custom prompt strategy or using a user-defined
 prompt template. Because the trainer cannot readily detect these changes, we cannot change the
-calculated hash value for the pre-processed dataset. If you have `dataset_prepared_path: ...` set
+calculated hash value for the pre-processed dataset.
+
+If you have `dataset_prepared_path: ...` set
 and change your prompt templating logic, it may not pick up the changes you made and you will be
 training over the old prompt.
diff --git a/docs/debugging.qmd b/docs/debugging.qmd
index 7237fbd6f2..5977e7a223 100644
--- a/docs/debugging.qmd
+++ b/docs/debugging.qmd
@@ -6,6 +6,10 @@ description: How to debug Axolotl
 
 This document provides some tips and tricks for debugging Axolotl.  It also provides an example configuration for debugging with VSCode.  A good debugging setup is essential to understanding how Axolotl code works behind the scenes.
 
+::: {.callout-tip}
+For training-specific debugging (loss spikes, NaN gradients, OOM errors, RL training stability), see [Training Stability & Debugging](training_stability.qmd).
+:::
+
 ## Table of Contents
 
 - [General Tips](#general-tips)
@@ -29,13 +33,15 @@ While debugging it's helpful to simplify your test scenario as much as possible.
 1. **Make sure you are using the latest version of axolotl**:  This project changes often and bugs get fixed fast.  Check your git branch and make sure you have pulled the latest changes from `main`.
 1. **Eliminate concurrency**: Restrict the number of processes to 1 for both training and data preprocessing:
     - Set `CUDA_VISIBLE_DEVICES` to a single GPU, ex: `export CUDA_VISIBLE_DEVICES=0`.
-    - Set `dataset_processes: 1` in your axolotl config or run the training command with `--dataset_processes=1`.
+    - Set `dataset_num_proc: 1` in your axolotl config or run the training command with `--dataset_num_proc=1`.
 2. **Use a small dataset**: Construct or use a small dataset from HF Hub. When using a small dataset, you will often have to make sure `sample_packing: False` and `eval_sample_packing: False` to avoid errors.  If you are in a pinch and don't have time to construct a small dataset but want to use from the HF Hub, you can shard the data (this will still tokenize the entire dataset, but will only use a fraction of the data for training.  For example, to shard the dataset into 20 pieces, add the following to your axolotl config):
+
     ```yaml
-    dataset:
+    datasets:
         ...
         shards: 20
     ```
+
 3. **Use a small model**: A good example of a small model is [TinyLlama/TinyLlama-1.1B-Chat-v1.0](https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0).
 4. **Minimize iteration time**: Make sure the training loop finishes as fast as possible, with these settings.
     - `micro_batch_size: 1`
@@ -51,12 +57,12 @@ While debugging it's helpful to simplify your test scenario as much as possible.
 
 ### Background
 
-The below example shows how to configure VSCode to debug data preprocessing of the `sharegpt` format.  This is the format used when you have the following in your axolotl config:
+The below example shows how to configure VSCode to debug data preprocessing of the `chat_template` format.  This is the format used when you have the following in your axolotl config:
 
 ```yaml
 datasets:
-  - path: <path to your sharegpt formatted dataset> # example on HF Hub: philschmid/guanaco-sharegpt-style
-    type: sharegpt
+  - path: <path to your chat_template formatted dataset> # example on HF Hub: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
 ```
 
 >[!Important]
@@ -70,8 +76,10 @@ datasets:
 Make sure you have an [editable install](https://setuptools.pypa.io/en/latest/userguide/development_mode.html) of Axolotl, which ensures that changes you make to the code are reflected at runtime.  Run the following commands from the root of this project:
 
 ```bash
-pip3 install packaging
-pip3 install -e '.[flash-attn,deepspeed]'
+export UV_TORCH_BACKEND=cu130  # or cu128
+uv venv --no-project --relocatable
+source .venv/bin/activate
+uv pip install --no-build-isolation -e '.[deepspeed]' --group dev --group test
 ```
 
 #### Remote Hosts
@@ -83,23 +91,23 @@ If you developing on a remote host, you can easily use VSCode to debug remotely.
 
 The easiest way to get started is to modify the [.vscode/launch.json](../.vscode/launch.json) file in this project.  This is just an example configuration, so you may need to modify or copy it to suit your needs.
 
-For example, to mimic the command `cd devtools && CUDA_VISIBLE_DEVICES=0 accelerate launch -m axolotl.cli.train dev_sharegpt.yml`, you would use the below configuration[^1].  Note that we add additional flags that override the axolotl config and incorporate the tips above (see the comments). We also set the working directory to `devtools` and set the `env` variable `HF_HOME` to a temporary folder that is later partially deleted.  This is because we want to delete the HF dataset cache before each run in order to ensure that the data preprocessing code is run from scratch.
+For example, to mimic the command `cd devtools && CUDA_VISIBLE_DEVICES=0 axolotl train dev_chat_template.yml`, you would use the below configuration[^1].  Note that we add additional flags that override the axolotl config and incorporate the tips above (see the comments). We also set the working directory to `devtools` and set the `env` variable `HF_HOME` to a temporary folder that is later partially deleted.  This is because we want to delete the HF dataset cache before each run in order to ensure that the data preprocessing code is run from scratch.
 
-```jsonc
+```json
 // .vscode/launch.json
 {
     "version": "0.2.0",
     "configurations": [
         {
-            "name": "Debug axolotl prompt - sharegpt",
+            "name": "Debug axolotl prompt - chat_template",
             "type": "python",
             "module": "accelerate.commands.launch",
             "request": "launch",
             "args": [
-                "-m", "axolotl.cli.train", "dev_sharegpt.yml",
+                "-m", "axolotl.cli.train", "dev_chat_template.yml",
                 // The flags below simplify debugging by overriding the axolotl config
                 // with the debugging tips above.  Modify as needed.
-                "--dataset_processes=1",      // limits data preprocessing to one process
+                "--dataset_num_proc=1",      // limits data preprocessing to one process
                 "--max_steps=1",              // limits training to just one step
                 "--batch_size=1",             // minimizes batch size
                 "--micro_batch_size=1",       // minimizes batch size
@@ -132,7 +140,7 @@ For example, to mimic the command `cd devtools && CUDA_VISIBLE_DEVICES=0 acceler
 
 Below is the [./vscode/tasks.json](../.vscode/tasks.json) file that defines the `cleanup-for-dataprep` task.  This task is run before each debugging session when you use the above configuration.  Note how there are two tasks that delete the two folders mentioned above.  The third task `cleanup-for-dataprep` is a composite task that combines the two tasks.  A composite task is necessary because VSCode does not allow you to specify multiple tasks in the `preLaunchTask` argument of the `launch.json` file.
 
-```jsonc
+```json
 // .vscode/tasks.json
 // this file is used by launch.json
 {
@@ -185,14 +193,14 @@ style="border-radius: 10px; display: block; margin: auto;" width="560" height="3
 
 ## Debugging With Docker
 
-Using [official Axolotl Docker images](https://hub.docker.com/r/winglian/axolotl/tags) is a great way to debug your code, and is a very popular way to use Axolotl.  Attaching VSCode to Docker takes a few more steps.
+Using [official Axolotl Docker images](https://hub.docker.com/r/axolotlai/axolotl/tags) is a great way to debug your code, and is a very popular way to use Axolotl.  Attaching VSCode to Docker takes a few more steps.
 
 ### Setup
 
 On the host that is running axolotl (ex: if you are using a remote host), clone the axolotl repo and change your current directory to the root:
 
 ```bash
-git clone https://github.com/OpenAccess-AI-Collective/axolotl
+git clone https://github.com/axolotl-ai-cloud/axolotl
 cd axolotl
 ```
 
@@ -202,17 +210,18 @@ cd axolotl
 Next, run the desired docker image and mount the current directory. Below is a docker command you can run to do this:[^2]
 
 ```bash
-docker run --privileged --gpus '"all"' --shm-size 10g --rm -it --name axolotl --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --mount type=bind,src="${PWD}",target=/workspace/axolotl -v ${HOME}/.cache/huggingface:/root/.cache/huggingface winglian/axolotl:main-py3.10-cu118-2.0.1
+docker run --privileged --gpus '"all"' --shm-size 10g --rm -it --name axolotl --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --mount type=bind,src="${PWD}",target=/workspace/axolotl -v ${HOME}/.cache/huggingface:/root/.cache/huggingface axolotlai/axolotl-uv:main-latest
 ```
 
 >[!Tip]
-> To understand which containers are available, see the [Docker section of the README](../README.md#docker) and the [DockerHub repo](https://hub.docker.com/r/winglian/axolotl/tags).  For details of how the Docker containers are built, see axolotl's [Docker CI builds](../.github/workflows/main.yml).
+> To understand which containers are available, see the [Docker section of the README](../README.md#docker) and the [DockerHub repo](https://hub.docker.com/r/axolotlai/axolotl/tags).  For details of how the Docker containers are built, see axolotl's [Docker CI builds](../.github/workflows/main.yml).
 
-You will now be in the container.  Next, perform an editable install of Axolotl:
+You will now be in the container.  Next, install Axolotl with dev dependencies:
 
 ```bash
-pip3 install packaging
-pip3 install -e '.[flash-attn,deepspeed]'
+uv venv --no-project --relocatable
+source .venv/bin/activate
+uv pip install --no-build-isolation -e '.[deepspeed]' --group dev --group test
 ```
 
 ### Attach To Container
@@ -240,6 +249,6 @@ style="border-radius: 10px; display: block; margin: auto;" width="560" height="3
 </div>
 <br>
 
-[^1]: The config actually mimics the command `CUDA_VISIBLE_DEVICES=0 python -m accelerate.commands.launch -m axolotl.cli.train devtools/sharegpt.yml`, but this is the same thing.
+[^1]: The VSCode config uses `accelerate.commands.launch` as the Python module entry point, which is what `axolotl train` invokes under the hood.
 
 [^2]: Many of the below flags are recommended best practices by Nvidia when using nvidia-container-toolkit.  You can read more about these flags [here](https://docs.nvidia.com/deeplearning/frameworks/user-guide/index.html).
diff --git a/docs/docker.qmd b/docs/docker.qmd
new file mode 100644
index 0000000000..3663f4a26a
--- /dev/null
+++ b/docs/docker.qmd
@@ -0,0 +1,167 @@
+---
+title: "Docker"
+format:
+  html:
+    toc: true
+    toc-depth: 4
+---
+
+This section describes the different Docker images that are released by AxolotlAI at
+[Docker Hub](https://hub.docker.com/u/axolotlai).
+
+::: {.callout-important}
+### uv is now the default
+
+All images now use [uv](https://docs.astral.sh/uv/) with a relocatable venv (`/workspace/axolotl-venv`)
+instead of Miniconda + pip. The plain image names are the defaults, and the `-uv` names are kept as
+aliases where they exist. Tags follow the same format for both names.
+:::
+
+## Supported variants
+
+Current Axolotl runtime images target Python 3.12 on CUDA 13.0 with PyTorch
+2.11.0 and 2.12.0. The floating `main-latest` tags point at the Python 3.12 /
+CUDA 13.0 / PyTorch 2.12.0 build. Bare release tags, such as `0.16.1`, point
+at the same latest matrix entry for that release.
+
+Base images may include additional uv-based tags for build compatibility, and
+older tags may remain available on Docker Hub for reproducibility. New runtime
+builds are centered on the CUDA 13.0 matrix.
+
+## Base
+
+The base image is the most minimal image that can install Axolotl. It is based on the `nvidia/cuda` image.
+It includes python, torch, git, git-lfs, awscli, pydantic, and more.
+
+#### Image
+
+| Image | Notes |
+|-------|-------|
+| `axolotlai/axolotl-base` | Default |
+| `axolotlai/axolotl-base-uv` | Alias |
+
+Links: [base](https://hub.docker.com/r/axolotlai/axolotl-base), [base-uv](https://hub.docker.com/r/axolotlai/axolotl-base-uv)
+
+#### Tags format
+
+```bash
+main-base-py{python_version}-cu{cuda_version}-{pytorch_version}
+```
+
+Tags examples:
+
+- `main-base-py3.12-cu130-2.11.0`
+- `main-base-py3.12-cu130-2.12.0`
+
+## Main
+
+The main image is the image that is used to run Axolotl. It is based on the `axolotlai/axolotl-base` image and includes the Axolotl codebase, dependencies, and more.
+
+#### Image
+
+| Image | Notes |
+|-------|-------|
+| `axolotlai/axolotl` | Default |
+| `axolotlai/axolotl-uv` | Alias |
+
+Links: [axolotl](https://hub.docker.com/r/axolotlai/axolotl), [axolotl-uv](https://hub.docker.com/r/axolotlai/axolotl-uv)
+
+#### Tags format {#sec-main-tags}
+
+```bash
+# on push to main
+main-py{python_version}-cu{cuda_version}-{pytorch_version}
+
+# latest stable build from main
+main-latest
+
+# nightly build
+{branch}-{date_in_YYYYMMDD}-py{python_version}-cu{cuda_version}-{pytorch_version}
+
+# tagged release
+{version}
+{version}-py{python_version}-cu{cuda_version}-{pytorch_version}
+{version}-latest
+```
+
+:::{.callout-tip}
+
+There may be some extra tags appended to the image, like `-vllm` which installs those packages.
+
+:::
+
+Tags examples:
+
+- `main-py3.12-cu130-2.11.0`
+- `main-py3.12-cu130-2.12.0`
+- `main-latest`
+- `main-20260315-py3.12-cu130-2.12.0`
+- `0.16.1`
+- `0.16.1-py3.12-cu130-2.12.0`
+- `0.16.1-latest`
+
+## Cloud
+
+The cloud image is the image that is used to run Axolotl in the cloud. It is based on the `axolotlai/axolotl` image and sets ENV variables like HuggingFace cache directories for volume mounts, tmux, and more for different cloud providers.
+
+:::{.callout-tip}
+
+Jupyter lab is run by default. Set `JUPYTER_DISABLE=1` in the environment variables to disable it.
+
+:::
+
+#### Image
+
+| Image | Notes |
+|-------|-------|
+| `axolotlai/axolotl-cloud` | Default |
+| `axolotlai/axolotl-cloud-uv` | Alias |
+
+Links: [cloud](https://hub.docker.com/r/axolotlai/axolotl-cloud), [cloud-uv](https://hub.docker.com/r/axolotlai/axolotl-cloud-uv)
+
+#### Tags format {#sec-cloud-tags}
+
+This uses the same tags as the [`main` image](#sec-main-tags).
+
+#### Environment variables
+
+- `JUPYTER_DISABLE`: Disable Jupyter lab.
+- `JUPYTER_PASSWORD`: Set a password for the Jupyter lab.
+- `PUBLIC_KEY` / `SSH_KEY`: Add a public key for the SSH service.
+
+#### Volume mounts
+
+:::{.callout-tip}
+
+We recommend mounting volumes to `/workspace/data` for data persistence. The
+cloud images set `HF_HOME` to `/workspace/data/huggingface-cache` so HuggingFace
+Hub and dataset caches are kept under the same persistent mount. `/workspace/axolotl`
+contains the source code and is ephemeral.
+
+:::
+
+- `/workspace/data/axolotl-artifacts`: Directory to store Axolotl artifacts.
+- `/workspace/data/huggingface-cache`: Directory to store HuggingFace cache.
+
+## Cloud-no-tmux
+
+This is the same as the [`cloud` image](#sec-cloud) but without tmux.
+
+#### Image
+
+```
+axolotlai/axolotl-cloud-term
+```
+
+Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl-cloud-term)
+
+:::{.callout-note}
+
+The naming may be a bit confusing as it has `-term` appended to the end. This
+image is uv-based and does not have a separate `-uv` alias.
+
+:::
+
+#### Tags format
+
+This uses the same tags as the [`cloud` image](#sec-cloud-tags).
diff --git a/docs/ebft.qmd b/docs/ebft.qmd
new file mode 100644
index 0000000000..d9afc3307c
--- /dev/null
+++ b/docs/ebft.qmd
@@ -0,0 +1,556 @@
+---
+title: "EBFT Training"
+description: "Energy-Based Fine-Tuning uses feature-matching rewards from internal representations to train language models without external reward functions."
+order: 9
+back-to-top-navigation: true
+toc: true
+toc-expand: 2
+toc-depth: 4
+---
+
+## Overview
+
+Energy-Based Fine-Tuning (EBFT) is a training method that optimizes language models by matching the **internal feature representations** of generated text to those of ground-truth completions. Instead of relying on external reward models or hand-crafted reward functions, EBFT extracts hidden states from intermediate layers of a frozen copy of the model and uses cosine similarity between generated and reference features as the reward signal.
+
+Paper: ["Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models"](https://arxiv.org/abs/2603.12248) (Jelassi et al., 2026)
+
+### How EBFT Differs from Other RL Methods
+
+| Method | Reward Signal | Requires | Best For |
+|--------|--------------|----------|----------|
+| **GRPO** | External reward function(s) | Custom reward code or reward model | Tasks with verifiable answers (math, code) |
+| **DPO** | Preference pairs (chosen vs rejected) | Paired preference data | Alignment with human preferences |
+| **EBFT** | Feature similarity to ground truth | Ground-truth completions | Any task with reference outputs |
+
+EBFT's key advantage is that it needs only ground-truth completions -- no reward engineering, no preference annotation, and no reward model training. The model's own internal representations serve as the reward signal. This makes it particularly effective for:
+
+- Code generation (match features of known-good solutions)
+- Instruction following with reference outputs
+- Continual pretraining on unstructured text (strided mode)
+- Multi-turn dialogue with reference conversations
+
+### Reward Formulation
+
+The EBFT reward for each generated completion is:
+
+```
+reward = alignment_coef * cosine_similarity(gen_features, gt_features)
+       - diversity_coef * mean_pairwise_similarity(gen_features)
+```
+
+- **Alignment**: How closely the generated output's internal representations match the ground truth. Higher is better.
+- **Diversity**: Penalizes generated samples that are too similar to each other (prevents mode collapse). Lower is better.
+- **CFM loss** (Cross-Feature Matching): Tracks `||mean(gen_features) - gt_features||^2` as a diagnostic. This is the quantity that EBFT ultimately minimizes.
+
+## Modes
+
+EBFT supports three operational modes, each suited to different use cases.
+
+### Structured Mode (Sync)
+
+Uses vLLM on a separate GPU for generation, with sequential generate-score-train steps. This is the simplest mode and recommended for getting started.
+
+```
+GPU 0: vLLM Server (generates completions, receives weight syncs)
+GPU 1: Trainer (feature extraction, reward computation, GRPO training)
+```
+
+**When to use**: Standard instruction-following or QA datasets where you have prompt/completion pairs. Requires 2 GPUs.
+
+### Structured Mode (Async)
+
+Same architecture as sync, but overlaps generation of the next batch with training on the current batch. Faster throughput at the cost of slightly stale weights during generation.
+
+**When to use**: Same data as sync mode, but when you want faster training and can tolerate weight staleness (controlled by `vllm_sync_interval`).
+
+### Strided Mode
+
+Runs entirely on a single GPU with no vLLM dependency. Places anchor points throughout a document and generates short rollouts at each anchor using block-parallel attention patterns.
+
+```
+Single GPU: Base model + LoRA adapter
+  - Strided block-parallel generation (flex_attention)
+  - Feature extraction via disable_adapter()
+  - No vLLM needed
+```
+
+**When to use**: Unstructured text data (raw code, prose, documents) where there is no natural prompt/completion split. Also works with structured data that includes prompt boundaries. Requires only 1 GPU.
+
+## Quick Start
+
+### Structured Mode
+
+This minimal example fine-tunes Qwen2-0.5B on code data using EBFT with vLLM generation.
+
+**Step 1**: Create a config file `ebft_quickstart.yaml`:
+
+```yaml
+base_model: Qwen/Qwen2-0.5B-Instruct
+
+rl: ebft
+
+ebft:
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+
+trl:
+  num_generations: 4
+  max_completion_length: 256
+  temperature: 0.7
+  use_vllm: true
+  vllm_server_host: 0.0.0.0
+  vllm_server_port: 8000
+  vllm_lora_sync: true
+  vllm_sync_interval: 3
+  use_data_producer: true
+  async_prefetch: false
+  scale_rewards: true
+  loss_type: grpo
+
+vllm:
+  gpu_memory_utilization: 0.5
+  max_model_len: 1024
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+    split: train[:500]
+
+# Standard training settings (see getting-started.qmd for details)
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_target_linear: true
+sequence_len: 1024
+micro_batch_size: 2
+gradient_accumulation_steps: 4
+max_steps: 20
+learning_rate: 5.0e-6
+bf16: auto
+attn_implementation: flash_attention_2
+gradient_checkpointing: true
+output_dir: ./outputs/ebft-quickstart
+```
+
+**Step 2**: Start vLLM on GPU 0:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve ebft_quickstart.yaml
+```
+
+**Step 3**: Wait approximately 30 seconds for vLLM to initialize, then start training on GPU 1:
+
+```bash
+CUDA_VISIBLE_DEVICES=1 axolotl train ebft_quickstart.yaml
+```
+
+::: {.callout-important}
+The `micro_batch_size` must be divisible by `num_generations`. For example, with `num_generations: 4`, valid values are 4, 8, 12, etc.
+:::
+
+### Dataset Format
+
+Structured mode datasets must produce two fields after the transform:
+
+- `prompt`: Either a string or a list of chat messages (`[{"role": "user", "content": "..."}]`)
+- `ground_truth`: A string containing the reference completion
+
+Example raw dataset row:
+
+```json
+{
+  "input": "Write a function to compute fibonacci numbers.",
+  "output": "def fibonacci(n):\n    if n <= 1:\n        return n\n    return fibonacci(n-1) + fibonacci(n-2)"
+}
+```
+
+The `ebft_opencode.transform` converts this to the required `{prompt, ground_truth}` format automatically.
+
+## Feature Extraction
+
+EBFT extracts hidden states from intermediate transformer layers and pools them into per-sequence embeddings. These embeddings are compared between generated and ground-truth completions to compute rewards.
+
+### Feature Layers
+
+The `feature_layers` parameter specifies which layers to extract, as fractions of total model depth:
+
+```yaml
+ebft:
+  feature_layers: [0.25, 0.5, 0.75]  # Quarter, middle, three-quarter depth
+```
+
+For a 32-layer model, this extracts layers 8, 16, and 24. The hidden states from all selected layers are concatenated along the feature dimension, producing embeddings of size `num_layers * hidden_dim`.
+
+::: {.callout-tip}
+Using multiple layers captures both low-level syntactic features (early layers) and high-level semantic features (later layers). The default `[0.25, 0.5, 0.75]` works well across model sizes.
+:::
+
+### Embed Methods
+
+The `embed_method` controls how per-token hidden states are pooled into a single vector per sequence:
+
+| Method | Description | Output Shape | Notes |
+|--------|-------------|-------------|-------|
+| `last_token` | Hidden state at the last non-padding token | `(B, D)` | Default. Good for autoregressive models where the last token summarizes the sequence. |
+| `mean_pooling` | Mean of all non-padding token states | `(B, D)` | Considers the entire sequence equally. |
+| `completion_mean` | Mean over completion tokens only (excludes prompt) | `(B, D)` | Focuses reward signal on generated content. Requires prompt length information. |
+| `concat` | Concatenation of states at 25%, 50%, 75% positions | `(B, 3*D)` | Captures positional structure. Higher dimensional. |
+
+```yaml
+ebft:
+  embed_method: completion_mean  # Focus on completion features
+```
+
+### SVD Whitening
+
+Whitening decorrelates the feature dimensions so that no single direction dominates the feature-matching loss. This is computed via SVD on the generated embeddings, with the same transform applied to the ground-truth embeddings.
+
+```yaml
+ebft:
+  use_whitening: true
+```
+
+When whitening is enabled, the reward computation applies a whitening matrix `W = U @ diag(1/S) @ U^T` derived from the SVD of generated embeddings. This ensures all feature dimensions contribute equally to the alignment reward.
+
+::: {.callout-note}
+Singular values scale with `sqrt(batch_size)`, so reward magnitudes are batch-size dependent. This is acceptable because the number of samples per prompt (`n_samples_per_prompt` or `num_generations`) is fixed during training.
+:::
+
+### Alignment and Diversity Coefficients
+
+The two reward components are weighted by coefficients:
+
+```yaml
+ebft:
+  alignment_coef: 1.0   # Weight for cosine similarity with ground truth
+  diversity_coef: 1.0   # Weight for pairwise similarity penalty
+```
+
+Both values are scaled by 2 internally (per paper equation 7). The final reward per sample is:
+
+```
+reward_j = 2 * alignment_coef * cos(gen_j, gt)
+         - 2 * diversity_coef * (1/(n-1)) * sum_{j' != j} dot(gen_j, gen_j')
+```
+
+Setting `diversity_coef: 0.0` disables the diversity penalty entirely, which may be appropriate when `num_generations` is small (e.g., 2).
+
+## Strided Mode
+
+Strided mode is designed for training on unstructured text data where there is no natural prompt/completion boundary. Instead of generating full completions with vLLM, it places **anchor points** at regular intervals throughout each document and generates short rollouts at each anchor using block-parallel attention.
+
+### How Block-Parallel Generation Works
+
+Given a document of length `S` tokens:
+
+1. **Anchor placement**: Starting at position `anchor_offset`, place anchors every `stride` tokens. Each anchor defines a block.
+2. **Context window**: Each block sees `context_length` tokens of preceding context from the original document.
+3. **Generation**: At each anchor, generate `generate_max_len` tokens autoregressively, conditioned only on the context window.
+4. **Parallelism**: All blocks are processed in a single forward pass using a specialized attention mask that prevents information leakage between blocks.
+
+```
+Document:   [tok0, tok1, ..., tok_S]
+                    |         |         |
+                 anchor_0   anchor_1  anchor_2
+                    |         |         |
+             [ctx][gen]  [ctx][gen]  [ctx][gen]
+```
+
+The attention mask ensures:
+
+- Prompt tokens use standard causal attention
+- Each generated block attends to its own context window and its own preceding generated tokens
+- Blocks do not attend to each other's generated tokens
+
+When `flex_attention` is available (PyTorch >= 2.5), the mask is compiled into efficient fused kernels. Otherwise, a dense 4D attention mask is used as a fallback.
+
+### Strided Mode Configuration
+
+```yaml
+base_model: meta-llama/Llama-3.2-1B
+rl: ebft
+
+ebft:
+  mode: strided
+  stride: 8                    # Tokens between anchor points
+  context_length: 8            # Context window per block
+  generate_max_len: 8          # Tokens to generate per block
+  n_samples_per_prompt: 4      # Independent rollouts per document
+  temperature: 0.6
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: true
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  rl_coef: 1.0                # RL policy gradient loss weight
+  ce_coef: 0.03               # Cross-entropy loss on GT tokens
+  advantage_estimator: rloo    # rloo, group_norm, or reinforce
+  min_completion_prefix: 8     # Skip anchors in prompt region
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_strided_structured.transform
+    split: train[:1%]
+
+sequence_len: 2048
+micro_batch_size: 1
+gradient_accumulation_steps: 2
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_target_linear: true
+
+bf16: auto
+attn_implementation: flex_attention
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true          # Required with flex_attention
+```
+
+Run with a single command (no vLLM needed):
+
+```bash
+CUDA_VISIBLE_DEVICES=0 axolotl train config.yaml
+```
+
+### Advantage Estimators
+
+Strided mode supports three advantage estimation methods:
+
+| Estimator | Formula | Requirements |
+|-----------|---------|-------------|
+| `rloo` | Leave-one-out baseline: `reward_j - mean(rewards_{-j})` | `n_samples_per_prompt >= 2` |
+| `group_norm` | Group normalization: `(reward_j - mean) / std` | `n_samples_per_prompt >= 2` |
+| `reinforce` | Raw reward as advantage (no baseline) | Works with `n_samples_per_prompt = 1` |
+
+::: {.callout-warning}
+When `n_samples_per_prompt: 1`, the trainer automatically falls back to `reinforce` and disables the diversity penalty (which requires multiple samples).
+:::
+
+### Strided Mode Constraints
+
+- **`flex_attention: true`** is strongly recommended. Without it, dense 4D masks consume significantly more memory.
+- **`torch_compile: true`** must NOT be set. `flex_attention` compiles its own kernels internally; adding `torch_compile` causes conflicts and OOM.
+- **Gradient checkpointing** must use `use_reentrant: true`. Non-reentrant checkpointing causes `CheckpointError` with `flex_attention` block masks.
+- **`activation_offloading`** is incompatible with `flex_attention`.
+
+### Cross-Entropy Loss
+
+Strided mode supports an optional cross-entropy loss term on ground-truth tokens. This acts as a regularizer to prevent the model from drifting too far from the original distribution:
+
+```yaml
+ebft:
+  ce_coef: 0.03    # Small CE coefficient
+  rl_coef: 1.0     # RL loss coefficient
+```
+
+The total loss is `rl_coef * rl_loss + ce_coef * ce_loss`. For structured mode, `ce_coef` is typically `0.0` since vLLM generation provides sufficient learning signal.
+
+## Dataset Formats
+
+EBFT provides several built-in dataset transforms in `src/axolotl/prompt_strategies/ebft/`.
+
+### Built-In Transforms
+
+| Transform | Input Format | Output Fields | Use Case |
+|-----------|-------------|---------------|----------|
+| `ebft_opencode.transform` | `{input, output}` | `{prompt, ground_truth}` | OpenCodeInstruct, structured QA |
+| `ebft_strided_structured.transform` | `{input, output}` | `{input_ids, labels, prompt_length}` | Strided mode with structured data |
+| `ebft_strided_chat.transform` | `{messages: [...]}` | `{input_ids, labels, prompt_length}` | Strided mode with chat data |
+| `ebft_chat_multiturn.transform` | `{messages: [...]}` | `{prompt, ground_truth, remaining_turns}` | Multi-turn: first-turn target |
+| `ebft_chat_multiturn.transform_last_turn` | `{messages: [...]}` | `{prompt, ground_truth}` | Multi-turn: last-turn target |
+| `ebft_chat_multiturn.transform_all_turns` | `{messages: [...]}` | `{prompt[], ground_truth[]}` | Multi-turn: one example per turn |
+| `ebft_reasoning.transform` | `{messages: [...]}` (with `<think>`) | `{prompt, ground_truth}` | Reasoning/thinking datasets |
+
+### Structured Mode Datasets
+
+For structured (sync/async) mode, the transform must produce `prompt` and `ground_truth` fields:
+
+```yaml
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+    split: train[:500]
+```
+
+### Multi-Turn Datasets
+
+Multi-turn transforms extract conversation data for sequential rollout. The `transform` variant targets the first assistant turn, while `transform_last_turn` targets the final turn:
+
+```yaml
+datasets:
+  - path: your/multiturn-dataset
+    type: ebft_chat_multiturn.transform
+```
+
+When `remaining_turns` is present in the dataset output, the trainer performs sequential rollouts: it generates the first assistant turn with vLLM, then continues generating subsequent turns by building up the conversation history.
+
+### Strided Mode Datasets
+
+Strided transforms tokenize the full document and produce `input_ids`, `labels`, and `prompt_length`:
+
+```yaml
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_strided_structured.transform
+    split: train[:1%]
+```
+
+### Custom Transforms
+
+To use your own dataset format, write a transform function:
+
+```python
+def transform(cfg, **kwargs):
+    def transform_fn(example, tokenizer=None):
+        return {
+            "prompt": [{"role": "user", "content": example["question"]}],
+            "ground_truth": example["answer"],
+        }
+    return transform_fn, {"remove_columns": "__all__"}
+```
+
+The `"__all__"` sentinel removes all original dataset columns after the mapping step. Reference this transform in your config:
+
+```yaml
+datasets:
+  - path: your/dataset
+    type: your_module.transform
+```
+
+## Configuration Reference
+
+### Common Parameters (All Modes)
+
+These parameters are set under the `ebft:` key in the YAML config.
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `mode` | `"structured"` or `"strided"` | `"structured"` | EBFT operating mode |
+| `feature_layers` | `list[float]` | `[0.25, 0.5, 0.75]` | Fractional layer depths for feature extraction |
+| `embed_method` | `string` | `"last_token"` | Pooling method: `last_token`, `mean_pooling`, `completion_mean`, or `concat` |
+| `use_whitening` | `bool` | `false` | Apply SVD whitening to feature embeddings before reward computation |
+| `alignment_coef` | `float` | `1.0` | Weight for alignment reward (cosine similarity with ground truth) |
+| `diversity_coef` | `float` | `1.0` | Weight for diversity penalty (pairwise dot product between samples) |
+| `ce_coef` | `float` | `0.0` | Cross-entropy loss coefficient on ground-truth tokens |
+| `adaptive_max_tokens` | `bool` | `true` | Dynamically set vLLM `max_tokens` based on ground-truth length (structured mode) |
+| `gt_length_multiplier` | `float` | `1.5` | Multiplier for ground-truth token count when computing adaptive max tokens (min 0.1) |
+
+### Strided Mode Parameters
+
+These additional parameters apply only when `mode: strided`.
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `stride` | `int` | `8` | Number of tokens between anchor points (must be >= 1) |
+| `context_length` | `int` | `8` | Context window size for each generated block (must be >= 1) |
+| `generate_max_len` | `int` | `8` | Number of tokens to generate per block (must be >= 1) |
+| `n_samples_per_prompt` | `int` | `4` | Number of independent rollouts per document (must be >= 1) |
+| `temperature` | `float` | `0.6` | Sampling temperature for strided generation |
+| `top_p` | `float` | `1.0` | Top-p nucleus sampling threshold |
+| `rl_coef` | `float` | `1.0` | RL policy gradient loss coefficient |
+| `advantage_estimator` | `string` | `"rloo"` | Advantage estimation method: `rloo`, `group_norm`, or `reinforce` |
+| `min_completion_prefix` | `int` | `0` | Minimum tokens into the completion span before placing anchors |
+
+### Structured Mode TRL Parameters
+
+These are set under the `trl:` key and control the GRPO training loop.
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `num_generations` | `int` | -- | Number of completions generated per prompt |
+| `max_completion_length` | `int` | -- | Maximum tokens per generated completion |
+| `temperature` | `float` | `0.7` | Sampling temperature for vLLM generation |
+| `use_vllm` | `bool` | -- | Enable vLLM generation backend |
+| `vllm_lora_sync` | `bool` | `false` | Sync LoRA adapters via filesystem (recommended) |
+| `vllm_sync_interval` | `int` | `1` | Steps between weight syncs to vLLM |
+| `use_data_producer` | `bool` | -- | Required for sync mode with LoRA sync |
+| `async_prefetch` | `bool` | `false` | Enable async generation (overlaps with training) |
+| `streaming_partial_batch` | `bool` | `false` | Score groups incrementally (async mode) |
+| `skip_zero_advantage_batches` | `bool` | `false` | Skip micro-batches where all advantages are zero |
+| `scale_rewards` | `bool` | -- | Normalize rewards within each prompt group |
+| `loss_type` | `string` | `"grpo"` | Loss type for policy optimization |
+| `epsilon` | `float` | `0.2` | Clipping parameter for importance sampling |
+
+### Stop Tokens
+
+vLLM needs explicit stop token IDs for generation. Common configurations:
+
+```yaml
+trl:
+  generation_kwargs:
+    stop_token_ids: [151645, 151643]   # Qwen: <|im_end|>, <|endoftext|>
+```
+
+### Multi-Turn Chat Settings
+
+For multi-turn conversations with Qwen3.5, disable thinking mode to prevent `<think>` tags in completions:
+
+```yaml
+trl:
+  chat_template_kwargs:
+    enable_thinking: false
+```
+
+## Monitoring
+
+### Key Metrics
+
+EBFT logs several custom metrics to wandb and the training console. Here is what to watch for:
+
+| Metric | Healthy Range | Interpretation |
+|--------|--------------|----------------|
+| `ebft/alignment` | 0.3 -- 0.9, trending upward | Cosine similarity between generated and ground-truth features. Higher means the model is learning to produce representations that match the reference. |
+| `ebft/diversity` | 0.01 -- 0.1 | Mean pairwise similarity between different generations for the same prompt. Values above 1.0 indicate mode collapse. |
+| `ebft/cfm_loss` | Below 10, trending downward | Cross-Feature Matching loss. This is the core quantity being minimized. Consistently above 100 indicates instability. |
+| `ebft/reward` | Trending upward (may start negative) | Combined reward signal. If stuck at -1.0, the diversity penalty is dominating alignment. |
+| `grad_norm` | 0.1 -- 3.0 | Gradient magnitude. Values of 0.0 indicate zero-advantage skip (normal). Values above 10 suggest instability. |
+| `entropy` | 0.05 -- 0.5 | Policy entropy. Values below 0.01 suggest mode collapse. |
+| `IS ratio min` | Above 0.1 | Importance sampling ratio minimum. Near-zero values mean the policy is too far off-policy; increase `vllm_sync_interval`. |
+
+### Console Log Example
+
+During training, you will see periodic EBFT reward logs:
+
+```
+ebft reward | align +0.412 ^ | divers +0.023 v | cfm 4.231 v | reward +0.389 ^
+```
+
+The arrows indicate the desired direction: alignment and reward should trend upward, while diversity and CFM loss should trend downward.
+
+### Troubleshooting
+
+| Symptom | Likely Cause | Fix |
+|---------|-------------|-----|
+| `alignment` stays below 0.1 | Feature layers not capturing useful information | Try different `feature_layers` or `embed_method` |
+| `diversity` exceeds 1.0 | Mode collapse -- generations are too similar | Increase `diversity_coef` or `temperature` |
+| `reward` stuck at -1.0 | Diversity penalty dominates alignment | Reduce `diversity_coef` or increase `alignment_coef` |
+| `grad_norm` consistently 0.0 | All micro-batches have zero advantage | Increase `num_generations` or check data quality |
+| `CheckpointError` in strided mode | Incompatible gradient checkpointing settings | Set `use_reentrant: true` in `gradient_checkpointing_kwargs` |
+| OOM during training | Logits tensor too large | Reduce `sequence_len` or `micro_batch_size`; strided mode uses chunked lm_head to mitigate this |
+| vLLM 500 errors | `truncate_prompt_tokens` not supported | Ensure you are using `axolotl vllm-serve` (not `trl vllm-serve`) |
+
+### Feature Network Memory
+
+In PEFT (LoRA) mode, the feature network shares base weights with the actor model by using the `disable_adapter()` context manager. This saves an entire model copy in VRAM (approximately 1--16 GB depending on model size). For non-PEFT training, a separate frozen deepcopy is created.
+
+::: {.callout-note}
+The `disable_adapter()` approach relies on an invariant: `merge_adapter()` is never called on the base weights. All weight sync paths (LoRA sync, HTTP, NCCL) compute merged weights as new tensors or save the adapter to the filesystem, leaving base weights unmodified.
+:::
+
+## Examples
+
+Complete example configurations are available in `examples/ebft/`:
+
+| Config | Model | Mode | Description |
+|--------|-------|------|-------------|
+| `llama-1b-ebft-strided-structured.yaml` | Llama 3.2 1B | Strided | Single-GPU strided training on code data |
+| `qwen3-4b-ebft-structured.yaml` | Qwen3 4B | Structured (sync) | Two-GPU structured training |
+| `qwen3-4b-ebft-structured-async.yaml` | Qwen3 4B | Structured (async) | Two-GPU async training with prefetch |
+| `qwen3-8b-ebft-structured.yaml` | Qwen3 8B | Structured (sync) | Two-GPU structured training for larger model |
+| `qwen35-4b-ebft-structured.yaml` | Qwen3.5 4B | Structured (sync) | Two-GPU with Qwen3.5 |
+| `qwen35-4b-ebft-structured-async.yaml` | Qwen3.5 4B | Structured (async) | Two-GPU async with Qwen3.5 |
+| `qwen35-9b-ebft-structured.yaml` | Qwen3.5 9B | Structured (sync) | Two-GPU structured for 9B model |
diff --git a/docs/expert_quantization.qmd b/docs/expert_quantization.qmd
new file mode 100644
index 0000000000..7eabed1cfe
--- /dev/null
+++ b/docs/expert_quantization.qmd
@@ -0,0 +1,67 @@
+---
+title: "MoE Expert Quantization"
+description: "Reduce VRAM usage when training MoE model adapters by quantizing expert weights on load"
+---
+
+Transformers v5 changed MoE expert layers from `nn.Linear` to fused `nn.Parameter` (3D+ tensors).
+This means `bitsandbytes` can no longer quantize them during model loading, resulting in all expert
+weights being loaded in full bf16 precision and causing massive VRAM usage.
+
+`quantize_moe_experts` solves this by quantizing expert weights during model loading.
+It intercepts the weight loading process, quantizes each expert tensor on the fly, and
+immediately frees the original bf16 tensor from VRAM. This dramatically reduces peak memory.
+For example, GLM-4.7-Flash QLoRA drops from ~127GiB to ~23GiB reserved memory.
+
+## Usage
+
+Enable expert quantization in your Axolotl config:
+
+```yaml
+quantize_moe_experts: true
+```
+
+This works with both 4-bit (QLoRA) and 8-bit (LoRA) quantization.
+
+### Expert LoRA targeting
+
+You can optionally apply LoRA adapters directly to expert weights using `lora_target_parameters`:
+
+```yaml
+lora_target_parameters:
+  - mlp.experts.gate_up_proj
+  - mlp.experts.down_proj
+  # - mlp.gate.weight  # router
+```
+
+::: {.callout-note}
+`lora_dropout` must be `0` when using `lora_target_parameters`.
+:::
+
+## Requirements
+
+- Requires (`adapter: lora` and `load_in_8bit: true`) or (`adapter: qlora` and `load_in_4bit: true`)
+- CUDA GPUs only (not tested with ROCm or other backends)
+- FSDP2 compatible for distributed training
+
+## Limitations
+
+- `lora_target_linear` is not compatible with `quantize_moe_experts`. See [Expert LoRA targeting](#expert-lora-targeting) instead.
+- `cpu_ram_efficient_loading` hangs / takes long time with FSDP2 + QLoRA.
+- Total model parameter count may display incorrectly (trainable param count is correct).
+- FSDP LoRA (8-bit) may have a large initial VRAM spike at the first 1-2 steps, which then drops. QLoRA does not exhibit this.
+- FSDP2 may use more VRAM per GPU than single GPU training due to not all layers being properly sharded across ranks.
+- Model loading takes longer due to on-demand quantization, even on consecutive runs.
+- DeepSpeed has not been tested.
+
+## Implementation details
+
+The quantization is applied by patching transformers to intercept weight loading.
+When a 3D+ CUDA tensor with "expert" in its name is detected:
+
+- **4-bit mode:** Uses bitsandbytes NF4 parametrization (configurable via `bnb_4bit_quant_type`).
+- **8-bit mode:** Uses a custom row-wise int8 parametrization with bitsandbytes dequantization.
+
+The original bf16 tensor is freed immediately after quantization. Multiple sub-patches are applied to
+transformers, PEFT and accelerate FSDP2 to support these parametrized expert modules.
+
+For full implementation details, see [PR #3439](https://github.com/axolotl-ai-cloud/axolotl/pull/3439).
diff --git a/docs/faq.qmd b/docs/faq.qmd
index 91413d24e9..834835150a 100644
--- a/docs/faq.qmd
+++ b/docs/faq.qmd
@@ -3,19 +3,152 @@ title: FAQ
 description: Frequently asked questions
 ---
 
+### General
 
 **Q: The trainer stopped and hasn't progressed in several minutes.**
 
 > A: Usually an issue with the GPUs communicating with each other. See the [NCCL doc](nccl.qmd)
 
-**Q: Exitcode -9**
+**Q: exitcode: -9**
 
 > A: This usually happens when you run out of system RAM.
 
-**Q: Exitcode -7 while using deepspeed**
+**Q: exitcode: -7 while using deepspeed**
 
 > A: Try upgrading deepspeed w: `pip install -U deepspeed`
 
 **Q: AttributeError: 'DummyOptim' object has no attribute 'step'**
 
-> A: You may be using deepspeed with single gpu. Please don't set `deepspeed:` in yaml or cli.
+**Q: ModuleNotFoundError: No module named 'mpi4py' using single GPU with deepspeed**
+
+> A: You may be using deepspeed with single gpu. Please remove the `deepspeed:` section in the yaml file or `--deepspeed` CLI flag.
+
+**Q: The codes is stuck on saving preprocessed datasets.**
+
+> A: This is usually an issue with the GPU. This can be resolved through setting the os environment variable `CUDA_VISIBLE_DEVICES=0`. If you are on runpod, this is usually a pod issue. Starting a new pod should take care of it.
+
+**Q: Received mismatch error on merge adapters / loading adapters between torch.Size of checkpoint and model.**
+
+> A: This is likely due to vocab size mismatch. By default, Axolotl expands the model's embeddings if the tokenizer has more tokens than the model. Please use the `axolotl merge-lora` command to merge the adapters instead of using your own scripts.
+
+> On the other hand, if the model has more tokens than the tokenizer, Axolotl does not shrink the model's embeddings unless `shrink_embeddings: true` is set in the config.
+
+**Q: How to call Axolotl via custom python scripts?**
+
+> A: 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: "..."
+> ```
+
+**Q: `IterableDataset error` or `KeyError: 'input_ids'` when using `preprocess` CLI**
+
+> A: This is because you may be using `preprocess` CLI with `pretraining_dataset:` or `skip_prepare_dataset: true` respectively. Please use `axolotl train` CLI directly instead as these datasets are prepared on demand.
+
+**Q: vLLM is not working with Axolotl**
+
+> A: We currently recommend torch 2.12 for use with `vllm`. Please ensure you use the right version. For Docker, please use a `-vllm` tagged image such as `main-py3.12-cu130-2.12.0-vllm`.
+
+**Q: FA2 2.8.0 `undefined symbol` runtime error on CUDA 12.4**
+
+> A: There seems to be a wheel issue with FA2 2.8.0 on CUDA 12.4. Try CUDA 12.6 instead or downgrade to FA2 2.7.4. Please refer to the upstream issue: https://github.com/Dao-AILab/flash-attention/issues/1717.
+
+**Q: Can we mix text and text+image datasets for VLM training?**
+
+> A: Yes, you can for newer VLM arch. The ones that would not work are LLaVA / Pixtral arch. If you notice one not working, please let us know!
+
+**Q: Why is `memory/max_*` different from `nvidia-smi`?**
+
+> A: We use `torch` APIs to retrieve this information. You can see https://docs.pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management for more information.
+
+### Chat templates
+
+**Q: `jinja2.exceptions.UndefinedError: 'dict object' has no attribute 'content' / 'role' / ____`**
+
+> A: This means that the property mapping for the stated attribute does not exist when building `chat_template` prompt. For example, if `no attribute 'content'`, please check you have added the correct mapping for `content` under `message_property_mappings`.
+
+**Q: `Empty template generated for turn ___`**
+
+> A: The `content` is empty for that turn.
+
+**Q: `Could not find content start/end boundary for turn __`**
+
+> A: The specific turn's start/end could not be detected. Please ensure you have set the `eos_token` following your `chat_template`. Otherwise, this could be a `chat_template` which doesn't use proper boundaries for each turn (like system). On the rare occurrence, make sure your content is not `[[dummy_message]]`. Please let us know about this.
+
+**Q: `Content end boundary is before start boundary for turn ___`**
+
+> A: This is an edge case which should not occur. Please create an Issue if this happens.
+
+**Q: `Content end boundary is the same as start boundary for turn ___. This is likely an empty turn.`**
+
+> A: This is likely an empty turn.
+
+**Q: The EOS token is incorrectly being masked or not being masked / `EOS token __ not found in chat template`.**
+
+> A: There can be two reasons:
+
+> 1. This is because of the mismatch between `tokenizer.eos_token` and EOS token in template. Please make sure to set `eos_token: ` under `special_tokens: ` to the same EOS token as in template.
+
+> 2. The EOS token is not in the template. Please check if your template is correct. As an example, `phi_35` template does not use its dedicated EOS token `<|endoftext|>` at the end.
+
+**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.
+
+**Q: The EOT token(s) are incorrectly being masked or not being masked / `EOT token __ not found in chat template`.**
+
+> A: There can be two reasons:
+
+> 1. The EOT token is different from the EOS token and was not specified under `eot_tokens: `. Please set `eot_tokens: ` to the same EOT token(s) as in template.
+
+> 2. There is more than one EOT token per turn in the template. Please raise an issue with examples as we recognize this as an edge case.
+
+**Q: `EOT token encoding failed. Please check if the token is valid and can be encoded.`**
+
+> A: There could be some issue with the tokenizer or unicode encoding. Please raise an issue with examples with the EOT token & tokenizer causing the issue.
+
+**Q: `EOT token __ is encoded as multiple tokens.`**
+
+> A: This is because the EOT token is encoded as multiple tokens which can cause unexpected behavior. Please add it under `tokens: ` or (recommended) override unused added_tokens via `added_tokens_overrides: `.
+
+**Q: `Conflict between train_on_eos and train_on_eot. eos_token is in eot_tokens and train_on_eos != train_on_eot`**
+
+> A: This is because the EOS token is in the `eot_tokens: ` while mismatch between `train_on_eos: ` and `train_on_eot: `. This will cause one to override the other. Please ensure that `train_on_eos: ` and `train_on_eot: ` are the same or remove the EOS token from `eot_tokens: `.
+
+**Q: If `eot_tokens: ` is not provided, what happens?**
+
+> A: If `eot_tokens: ` is not provided, the default behavior is the same as before. EOS tokens used to delimit turns are masked/unmasked depending on whether the turn is trainable.
+
+> Internally, `eot_tokens: tokenizer.eos_token` and `train_on_eot: train_on_eos` (which defaults to `turn`). This transition helps clarify the naming and behavior of EOT/EOS tokens.
+
+**Q: `Data processing error: CAS service error`**
+
+> A: Try disabling XET with `export HF_HUB_DISABLE_XET=1`
+
+**Q: `torch._inductor.exc.LoweringException: NoValidChoicesError: No choices to select, please consider adding ATEN into max_autotune_gemm_backends config (defined in torch/_inductor/config.py) to allow at least one choice. `**
+
+> A: Depending on the version of torch, you may need to include this in your YAML:
+
+> ```yaml
+> flex_attn_compile_kwargs:
+>   dynamic: false
+>   mode: max-autotune-no-cudagraphs
+> ```
+
+**Q: `ValueError("Backward pass should have cleared tracker of all tensors")`
+
+> A: This may happen due to edge cases in using the modern OffloadActivations context manager for CUDA streams. If you encounter this error, you may have success using the naive implementation with `offload_activations: legacy` in your YAML.
+
+**Q: `Error parsing tool_calls arguments as JSON.`
+
+> A: There is an error parsing string arguments to a dict. Please check your dataset and the error message for more details.
diff --git a/docs/fsdp_qlora.qmd b/docs/fsdp_qlora.qmd
index 7f12d44935..01f57e627b 100644
--- a/docs/fsdp_qlora.qmd
+++ b/docs/fsdp_qlora.qmd
@@ -1,5 +1,5 @@
 ---
-title: "FDSP + QLoRA"
+title: "FSDP + QLoRA"
 description: Use FSDP with QLoRA to fine-tune large LLMs on consumer GPUs.
 format:
   html:
@@ -20,16 +20,22 @@ To enable `QLoRA` with `FSDP`, you need to perform the following steps:
 > See the [example config](#example-config) file in addition to reading these instructions.
 
 1. Set `adapter: qlora` in your axolotl config file.
-2. Enable FSDP in your axolotl config, as [described here](https://github.com/OpenAccess-AI-Collective/axolotl?tab=readme-ov-file#fsdp).
+2. Enable FSDP in your axolotl config, as [described here](multi-gpu.qmd#sec-fsdp).
 3. Use one of the supported model types: `llama`, `mistral` or `mixtral`.
 
+## Enabling Swap for FSDP2
+
+If available memory is insufficient even after FSDP's CPU offloading, you can enable swap memory usage by setting `cpu_offload_pin_memory: false` alongside `offload_params: true` in FSDP config.
+
+This disables memory pinning, allowing FSDP to use disk swap space as fallback. Disabling memory pinning itself incurs performance overhead, and actually having to use swap adds more, but it may enable training larger models that would otherwise cause OOM errors on resource constrained systems.
+
 ## Example Config
 
 [examples/llama-2/qlora-fsdp.yml](../examples/llama-2/qlora-fsdp.yml) contains an example of how to enable QLoRA + FSDP in axolotl.
 
 ## References
 
-- [PR #1378](https://github.com/OpenAccess-AI-Collective/axolotl/pull/1378) enabling QLoRA in FSDP in Axolotl.
+- [PR #1378](https://github.com/axolotl-ai-cloud/axolotl/pull/1378) enabling QLoRA in FSDP in Axolotl.
 - [Blog Post](https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html) from the [Answer.AI](https://www.answer.ai/) team describing the work that enabled QLoRA in FSDP.
 - Related HuggingFace PRs Enabling FDSP + QLoRA:
     - Accelerate [PR#2544](https://github.com/huggingface/accelerate/pull/2544 )
diff --git a/docs/getting-started.qmd b/docs/getting-started.qmd
new file mode 100644
index 0000000000..d7eac064fa
--- /dev/null
+++ b/docs/getting-started.qmd
@@ -0,0 +1,195 @@
+---
+title: "Quickstart"
+format:
+  html:
+    toc: true
+    toc-depth: 3
+    number-sections: true
+execute:
+  enabled: false
+---
+
+This guide will walk you through your first model fine-tuning project with Axolotl.
+
+## Quick Example {#sec-quick-example}
+
+Let's start by fine-tuning a small language model using LoRA. This example uses a 1B parameter model to ensure it runs on most GPUs.
+Assuming `axolotl` is installed (if not, see our [Installation Guide](installation.qmd))
+
+1. Download example configs:
+```bash
+axolotl fetch examples
+```
+
+2. Run the training:
+```bash
+axolotl train examples/llama-3/lora-1b.yml
+```
+
+That's it! Let's understand what just happened.
+
+## Understanding the Process {#sec-understanding}
+
+### The Configuration File {#sec-config}
+
+The YAML configuration file controls everything about your training. Here's what (part of) our example config looks like:
+
+```yaml
+base_model: NousResearch/Llama-3.2-1B
+
+load_in_8bit: true
+adapter: lora
+
+datasets:
+  - path: teknium/GPT4-LLM-Cleaned
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+```
+
+::: {.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-reference.qmd) for more details.
+
+### Training {#sec-training}
+
+When you run `axolotl train`, Axolotl:
+
+1. Downloads the base model
+2. (If specified) applies QLoRA/LoRA adapter layers
+3. Loads and processes the dataset
+4. Runs the training loop
+5. Saves the trained model and / or LoRA weights
+
+## Your First Custom Training {#sec-custom}
+
+Let's modify the example for your own data:
+
+1. Create a new config file `my_training.yml`:
+
+```yaml
+base_model: NousResearch/Nous-Hermes-llama-1b-v1
+
+load_in_8bit: true
+adapter: lora
+
+# Training settings
+micro_batch_size: 2
+num_epochs: 3
+learning_rate: 0.0003
+
+# Your dataset
+datasets:
+  - path: my_data.jsonl        # Your local data file
+    type: alpaca               # Or other format
+```
+
+This specific config is for LoRA fine-tuning a model with instruction tuning data using
+the `alpaca` dataset format, which has the following format:
+
+```json
+{
+    "instruction": "Write a description of alpacas.",
+    "input": "",
+    "output": "Alpacas are domesticated South American camelids..."
+}
+```
+
+Please see our [Dataset Formats](dataset-formats) for more dataset formats and how to
+format them.
+
+2. Prepare your JSONL data in the specified format (in this case, the expected `alpaca`
+format):
+
+```json
+{"instruction": "Classify this text", "input": "I love this!", "output": "positive"}
+{"instruction": "Classify this text", "input": "Not good at all", "output": "negative"}
+```
+
+3. Run the training:
+
+```bash
+axolotl train my_training.yml
+```
+
+## Common Tasks {#sec-common-tasks}
+
+::: {.callout-tip}
+
+The same yaml file is used for training, inference, and merging.
+
+:::
+
+### Testing Your Model {#sec-testing}
+
+After training, test your model:
+
+```bash
+axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out"
+```
+
+More details can be found in [Inference](inference.qmd).
+
+### Using a UI {#sec-ui}
+
+Launch a Gradio interface:
+
+```bash
+axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out" --gradio
+```
+
+### Preprocessing Data {#sec-preprocessing}
+
+For large datasets, preprocess first:
+
+```bash
+axolotl preprocess my_training.yml
+```
+
+Please make sure to set `dataset_prepared_path: ` in your config to set the path to save the prepared dataset.
+
+More details can be found in [Dataset Preprocessing](dataset_preprocessing.qmd).
+
+### Merging LoRA weights {#sec-merging-lora}
+
+To merge the LoRA weights back into the base model, run:
+
+```bash
+axolotl merge-lora my_training.yml --lora-model-dir="./outputs/lora-out"
+```
+
+The merged model will be saved in the `{output_dir}/merged` directory.
+
+More details can be found in [Merging LoRA weights](inference.qmd#sec-merging).
+
+## Next Steps {#sec-next-steps}
+
+Now that you have the basics, explore these guides based on what you want to do:
+
+**Choose your path:**
+
+- [Choosing a Fine-Tuning Method](choosing_method.qmd) — SFT vs LoRA vs QLoRA vs GRPO vs DPO, with hardware recommendations
+
+**Core guides:**
+
+- [Dataset Loading](dataset_loading.qmd) — Loading datasets from various sources
+- [Dataset Formats](dataset-formats) — Working with different data formats
+- [Optimizations](optimizations.qmd) — Flash attention, gradient checkpointing, sample packing
+- [Training Stability & Debugging](training_stability.qmd) — Monitoring metrics, fixing NaN, OOM debugging
+
+**Advanced training methods:**
+
+- [RLHF / Preference Learning](rlhf.qmd) — DPO, KTO, GRPO, EBFT
+- [GRPO Training](grpo.qmd) — RL with custom rewards and vLLM generation
+- [vLLM Serving](vllm_serving.qmd) — Setting up vLLM for GRPO
+
+**Scaling up:**
+
+- [Multi-GPU Training](multi-gpu.qmd) — DeepSpeed, FSDP, DDP
+- [Multi-Node Training](multi-node.qmd) — Distributed training across machines
diff --git a/docs/gradient_checkpointing.qmd b/docs/gradient_checkpointing.qmd
new file mode 100644
index 0000000000..21f11b7efd
--- /dev/null
+++ b/docs/gradient_checkpointing.qmd
@@ -0,0 +1,136 @@
+---
+title: Gradient Checkpointing, Activation Offloading, and Layer Offloading
+---
+
+Gradient checkpointing and activation offloading are techniques used to optimize the performance of deep learning
+models by reducing the memory footprint and improving computational efficiency.
+
+### Enabling Gradient Checkpointing
+
+```yaml
+gradient_checkpointing: true
+```
+
+### Enabling Activation Offloading
+
+```yaml
+gradient_checkpointing: true  # required for activation offloading
+activation_offloading: true
+```
+
+Activation offloading variants:
+
+The default `activation_offloading: true` offloads activations to CPU and uses CUDA streams
+to overlap the communications and computations when offloading.
+
+The `activation_offloading: legacy` naively offloads activations to CPU and without additional optimizations.
+
+For resource constrained environments with limited CPU memory, `activation_offloading: disk` offloads
+activations to disk instead of CPU RAM so that much larger context lengths can be trained with minimal memory.
+
+The `activation_offloading: hidden_states` mode keeps gradient checkpointing enabled and moves only
+the decoder-layer checkpoint input (`hidden_states`) to CPU. With the default
+`gradient_checkpointing_kwargs.use_reentrant: false`, Axolotl uses Transformers' non-reentrant
+checkpointing plus a saved-tensor hook that only offloads marked hidden-state tensors. Other saved
+tensors stay on GPU, which avoids the CPU memory and bandwidth cost of full activation offload.
+
+Set `gradient_checkpointing_kwargs.use_reentrant: true` to use the older ALST-style implementation.
+That path patches torch's reentrant `CheckpointFunction` and is DTensor-aware for sequence/context
+parallelism.
+
+### Choosing a mode
+
+| Mode | What moves | Best for |
+|------|-----------|----------|
+| `true` | All activations → CPU, stream-overlapped | General use; LoRA/QLoRA (offloads instead of recomputing) |
+| `legacy` | All activations → CPU, synchronous | Lowest resident GPU memory; when the streamed path's in-flight buffering inflates peak |
+| `disk` | Activations → disk | Severely CPU-RAM-constrained hosts |
+| `hidden_states` | Per-layer checkpoint input only → CPU, stream-overlapped | Long-context training where full activation offload is CPU-memory or bandwidth-bound |
+
+`hidden_states` uses gradient checkpointing and recomputes each layer during backward. With LoRA/QLoRA,
+`activation_offloading: true` can still be faster because adapter activations are smaller and pure offload
+avoids the recompute cost.
+
+### Selective Activation Checkpointing (SAC)
+
+```yaml
+gradient_checkpointing: true
+selective_checkpointing: true
+```
+
+Plain gradient checkpointing recomputes every op in a decoder layer during backward. Selective
+checkpointing keeps the recompute for cheap ops but *saves* the outputs of expensive ones — by
+default the attention op — so backward skips re-running them. At long context, attention dominates
+the recompute cost, so this recovers most of the checkpointing slowdown for one extra
+hidden-state-sized tensor per layer.
+
+This runs eagerly (no `torch.compile`) via torch's selective-checkpoint policy API. Ops are matched
+at the dispatcher level: SDPA and flash-attention are matched out of the box, and custom kernels
+that are not registered as torch ops are simply recomputed as usual — nothing breaks.
+
+The explicit form lets you save additional ops, matched as substrings of the qualified op name:
+
+```yaml
+selective_checkpointing:
+  save:
+    - attention
+    - aten::mm        # example: also save all matmuls
+```
+
+For hybrid models mixing full attention and sliding-window attention (SWA), only full-attention
+calls are saved by default — SWA is cheap to recompute and not worth the memory. Hybrid layers are
+detected two ways: the flash-attention window arguments, and the model's `layer_types` (published
+per-layer via hooks, which also covers SDPA where the window lives inside the attention mask).
+Set `selective_checkpointing.save_sliding_window: true` to save SWA calls too, or override which
+layer types recompute with `selective_checkpointing.recompute_layer_types` (default
+`[sliding_attention, chunked_attention]`; linear-attention layers never dispatch a matchable
+attention op, so they need no entry).
+
+To offload the saved tensors to pinned CPU memory instead of keeping them on GPU — asynchronous
+side-stream copies during forward, prefetched back one region ahead during backward — set:
+
+```yaml
+selective_checkpointing:
+  save: [attention]
+  offload: true
+```
+
+This makes the saved attention outputs cost ~zero GPU memory at the price of CPU RAM and PCIe
+traffic, and stacks with `activation_offloading: hidden_states`.
+
+Requires non-reentrant checkpointing (the default). Composes with
+`activation_offloading: hidden_states`: layer inputs go to CPU while attention outputs are saved
+(on GPU, or on CPU with `offload: true`). It is incompatible with the TRL-offloader modes
+(`true`/`legacy`/`disk`), which replace gradient checkpointing instead of running inside it.
+Saving matmul ops (e.g. `aten::mm`) is rejected for LoRA/QLoRA runs: PEFT mutates the base linear
+output in-place, which invalidates cached tensors.
+
+### Enabling Layer Offloading
+
+```yaml
+layer_offloading: true
+```
+
+Layer offloading reduces GPU memory usage by moving frozen (non-trainable) decoder layer parameters to CPU
+and streaming them back to GPU one layer at a time during the forward and backward passes. This is
+particularly useful for LoRA/QLoRA training where most of the model's parameters are frozen — only the
+trainable adapter weights stay on GPU permanently.
+
+During training, forward and backward hooks on each decoder layer handle the transfer automatically:
+
+- **Forward pass:** Before a layer executes, its frozen params are loaded to GPU. The next layer is
+  prefetched asynchronously on a separate CUDA stream for overlap.
+- **Backward pass:** Same pattern in reverse — the current layer's frozen params are loaded and the
+  previous layer is prefetched.
+
+After each layer finishes, its frozen params are offloaded back to CPU pinned memory.
+
+This approach trades some CPU-GPU transfer overhead for significant GPU memory savings — the freed memory
+is roughly equal to the size of all frozen parameters across all decoder layers, minus one layer's worth
+that is kept on GPU at any given time.
+
+**Requirements:**
+
+- CUDA GPU (CPU-only training is not supported for this feature)
+- Works with any HuggingFace model architecture that uses decoder layers (Llama, Mistral, Qwen, etc.)
+- Best combined with LoRA/QLoRA where most parameters are frozen
diff --git a/docs/grpo.qmd b/docs/grpo.qmd
new file mode 100644
index 0000000000..a98dbe11d4
--- /dev/null
+++ b/docs/grpo.qmd
@@ -0,0 +1,611 @@
+---
+title: "GRPO Training"
+description: "Group Relative Policy Optimization — a reinforcement learning method for training language models with verifiable reward functions."
+order: 8
+---
+
+## Overview
+
+Group Relative Policy Optimization (GRPO) is a reinforcement learning method that improves language models by generating multiple completions per prompt, scoring them with reward functions, and using the relative ranking within each group to compute advantage estimates. Unlike DPO, which requires pre-collected preference pairs, GRPO generates its own training data online and can work with any programmatic reward signal (math correctness, format compliance, code execution results, etc.).
+
+Use GRPO when you have a task with a verifiable reward signal and want the model to discover solution strategies on its own. Use DPO when you already have human preference data. Use SFT when you have gold-standard completions to imitate directly.
+
+Axolotl's GRPO implementation builds on TRL and adds async generation, streaming scoring, importance sampling correction, replay buffers, and multi-GPU scaling via FSDP and DeepSpeed.
+
+
+## Architecture
+
+GRPO training uses a two-process architecture: a vLLM server for fast generation and a trainer process for scoring and gradient updates.
+
+```
+Terminal 1 (GPU 0)                    Terminal 2 (GPU 1)
+┌──────────────────────┐              ┌──────────────────────────────────┐
+│  vLLM Server         │              │  Trainer                         │
+│                      │   HTTP       │                                  │
+│  Serves base model   │◄────────────►│  Background thread:              │
+│  + LoRA adapter      │  /generate   │    Send prompts to vLLM          │
+│                      │  /set_lora   │    Pad & collate completions     │
+│  Punica kernels for  │              │                                  │
+│  LoRA inference      │              │  Main thread:                    │
+│                      │              │    Score completions (rewards)   │
+└──────────────────────┘              │    Compute policy log-probs      │
+                                      │    Calculate advantages          │
+                                      │    PPO-clip gradient update      │
+                                      │    Sync LoRA weights to vLLM     │
+                                      └──────────────────────────────────┘
+```
+
+**Data flow for each training step:**
+
+1. The background thread sends prompts to vLLM, which generates `num_generations` completions per prompt.
+2. The main thread scores completions using your reward functions.
+3. Advantages are computed within each prompt group (group-relative normalization).
+4. Policy log-probabilities are computed by running a forward pass on the training model.
+5. The PPO-clip loss is computed and gradients are applied.
+6. Periodically, LoRA adapter weights are synced back to vLLM so future generations reflect the updated policy.
+
+With async prefetch enabled, step 1 for the *next* batch runs concurrently with steps 2-6 for the *current* batch.
+
+
+## Quick Start
+
+A GRPO training run requires three components: a YAML config, a reward module (Python file), and a running vLLM server.
+
+### 1. Write a reward module
+
+Create a file called `rewards.py` in your working directory:
+
+```python
+# rewards.py
+import re
+
+
+def accuracy_reward(completions, answer, **kwargs) -> list[float]:
+    """Check if the completion contains the correct numerical answer."""
+    rewards = []
+    for completion, correct in zip(completions, answer):
+        text = completion[0]["content"]
+        # Extract the last number from the completion
+        numbers = re.findall(r"-?\d+(?:\.\d+)?", text)
+        predicted = numbers[-1] if numbers else ""
+        rewards.append(1.0 if predicted == str(correct) else 0.0)
+    return rewards
+
+
+def format_reward(completions, **kwargs) -> list[float]:
+    """Reward completions that use a structured thinking format."""
+    rewards = []
+    for completion in completions:
+        text = completion[0]["content"]
+        has_think = "<think>" in text and "</think>" in text
+        has_answer = "<answer>" in text and "</answer>" in text
+        rewards.append(1.0 if has_think and has_answer else 0.0)
+    return rewards
+
+
+def prompt_transform(cfg, *args, **kwargs):
+    """Convert GSM8K dataset rows into chat prompts."""
+    def transform_fn(example, tokenizer=None):
+        label = example["answer"].split("####")[-1].strip().replace(",", "")
+        return {
+            "prompt": [
+                {"role": "system", "content": "Solve the math problem. Show your reasoning in <think> tags and your final numerical answer in <answer> tags."},
+                {"role": "user", "content": example["question"]},
+            ],
+            "answer": label,
+        }
+    return transform_fn, {"remove_columns": ["question"]}
+```
+
+### 2. Write the config
+
+Create `config.yaml`:
+
+```yaml
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+
+rl: grpo
+chat_template: tokenizer_default
+
+vllm:
+  host: 0.0.0.0
+  port: 8000
+  gpu_memory_utilization: 0.85
+  dtype: auto
+  max_model_len: 2048
+
+adapter: lora
+lora_r: 32
+lora_alpha: 64
+lora_target_linear: true
+
+trl:
+  use_vllm: true
+  use_data_producer: true
+  vllm_server_host: 0.0.0.0
+  vllm_server_port: 8000
+  vllm_server_timeout: 300
+  vllm_lora_sync: true
+  num_generations: 8
+  max_completion_length: 512
+  temperature: 0.7
+  reward_funcs:
+    - rewards.accuracy_reward
+    - rewards.format_reward
+  reward_weights:
+    - 1.0
+    - 0.5
+
+datasets:
+  - path: openai/gsm8k
+    name: main
+    type: rewards.prompt_transform
+    split: train
+
+skip_prepare_dataset: true
+val_set_size: 0.0
+sequence_len: 512
+micro_batch_size: 2
+gradient_accumulation_steps: 4
+max_steps: 200
+learning_rate: 5.0e-6
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 10
+
+bf16: true
+attn_implementation: flash_attention_2
+gradient_checkpointing: true
+
+special_tokens:
+  pad_token: "<|endoftext|>"
+
+output_dir: ./grpo-output
+logging_steps: 1
+```
+
+### 3. Start vLLM and train
+
+```bash
+# Terminal 1: Start vLLM server on GPU 0
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
+
+# Wait 30-90 seconds for model loading and CUDA graph capture
+
+# Terminal 2: Train on GPU 1
+CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml
+```
+
+:::{.callout-tip}
+Use `tmux` or separate terminal sessions to manage the two processes. The vLLM server must remain running for the entire training duration.
+:::
+
+
+## Custom Reward Functions
+
+### Function signature
+
+TRL calls reward functions with this signature:
+
+```python
+def my_reward(completions, **kwargs) -> list[float]:
+```
+
+- `completions` is a list of single-element lists, where each element is a dict `{"role": "assistant", "content": "..."}`. So `completions[i][0]["content"]` gives you the text of the i-th completion.
+- `**kwargs` contains all dataset columns that were *not* removed by the dataset transform. This is how you pass ground truth answers, metadata, or any other information to your reward function.
+- Return a `list[float]` with the same length as `completions`. You may return `None` for individual elements to exclude them from aggregation.
+
+### Example: accuracy reward with answer extraction
+
+```python
+def accuracy_reward(completions, answer, **kwargs) -> list[float]:
+    rewards = []
+    for completion, correct_answer in zip(completions, answer):
+        text = completion[0]["content"]
+        # Extract answer from <answer>...</answer> tags
+        match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
+        predicted = match.group(1).strip() if match else ""
+        rewards.append(1.0 if predicted == str(correct_answer) else 0.0)
+    return rewards
+```
+
+### Example: length penalty
+
+```python
+def length_penalty(completions, **kwargs) -> list[float]:
+    """Penalize very short or very long completions."""
+    rewards = []
+    for completion in completions:
+        length = len(completion[0]["content"])
+        if length < 50:
+            rewards.append(-0.5)
+        elif length > 2000:
+            rewards.append(-0.2)
+        else:
+            rewards.append(0.0)
+    return rewards
+```
+
+### Multiple rewards and weighting
+
+You can combine multiple reward functions with different weights:
+
+```yaml
+trl:
+  reward_funcs:
+    - rewards.accuracy_reward
+    - rewards.format_reward
+    - rewards.length_penalty
+  reward_weights:
+    - 1.0    # accuracy is most important
+    - 0.5    # format compliance
+    - 0.1    # mild length preference
+```
+
+Rewards are combined by the `multi_objective_aggregation` strategy:
+
+- `sum_then_normalize` (default): weights and sums all rewards first, then normalizes across the group.
+- `normalize_then_sum` (GDPO): normalizes each reward independently, then sums. This prevents one reward from dominating and is recommended when using multiple reward functions with different scales.
+
+```yaml
+trl:
+  multi_objective_aggregation: normalize_then_sum
+```
+
+### Dataset transforms
+
+The dataset transform converts raw HuggingFace dataset rows into chat-format prompts:
+
+```python
+def prompt_transform(cfg, *args, **kwargs):
+    def map_fn(example, tokenizer=None):
+        return {
+            "prompt": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": example["question"]},
+            ],
+            # Keep 'answer' column for the reward function
+            "answer": example["answer"],
+        }
+    # Remove columns consumed by the transform; keep columns needed by rewards
+    return map_fn, {"remove_columns": ["question"]}
+```
+
+The transform returns a tuple of `(map_function, kwargs_dict)`. The `remove_columns` in the kwargs dict removes columns that are no longer needed. Columns that your reward functions reference via `**kwargs` (like `answer`) must *not* be removed.
+
+:::{.callout-warning}
+The reward module must be importable from the directory where you run `axolotl train`. If your reward file is `rewards.py`, the import path is `rewards.accuracy_reward`. If it is inside a package `my_rewards/scoring.py`, use `my_rewards.scoring.accuracy_reward`.
+:::
+
+### Reward models (neural network rewards)
+
+Instead of a Python function, you can pass a HuggingFace model path as a reward function. TRL will load it as a reward model and use its scalar output as the reward:
+
+```yaml
+trl:
+  reward_funcs:
+    - OpenAssistant/reward-model-deberta-v3-large-v2
+    - rewards.format_reward
+  reward_weights:
+    - 1.0
+    - 0.3
+```
+
+### Using math_verify
+
+The `math_verify` library provides robust mathematical answer verification but uses `signal.alarm()` internally, which only works in the main thread. If you use `math_verify` in a reward function, set `reward_num_workers` to use subprocess workers:
+
+```yaml
+trl:
+  reward_num_workers: 4
+```
+
+Each worker runs in its own subprocess with its own main thread, so `signal.alarm()` works correctly.
+
+
+## vLLM Setup
+
+GRPO requires a running vLLM server for generation. For a complete guide on server modes, LoRA sync, weight synchronization, and restart procedures, see [vLLM Serving](vllm_serving.qmd).
+
+The minimal setup:
+
+```yaml
+vllm:
+  host: 0.0.0.0
+  port: 8000
+  gpu_memory_utilization: 0.85
+
+trl:
+  use_vllm: true
+  vllm_lora_sync: true         # Recommended with LoRA — faster sync, no NCCL contention
+  vllm_sync_interval: 5        # Sync weights every 5 steps
+```
+
+```bash
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml   # GPU 0: vLLM
+CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml         # GPU 1: training
+```
+
+:::{.callout-warning}
+vLLM must be restarted between experiments — stale weight syncs corrupt server state. See [Restart Requirements](vllm_serving.qmd#sec-restart).
+:::
+
+
+## Async Training Features
+
+Async GRPO overlaps generation and training to reduce wall-clock time. While the model trains on the current batch, the next batch is already being generated by vLLM.
+
+### Enabling async prefetch
+
+```yaml
+trl:
+  use_data_producer: true
+  async_prefetch: true
+  prefetch_depth: 1
+  vllm_sync_interval: 2
+```
+
+- `use_data_producer: true` enables the data producer protocol (required for all async features).
+- `async_prefetch: true` runs generation in a background thread.
+- `prefetch_depth` controls how many batches to prefetch ahead (1 is usually sufficient).
+- `vllm_sync_interval` controls how often LoRA weights are synced to vLLM (every N optimizer steps). Lower values mean fresher generations but more sync overhead.
+
+:::{.callout-tip}
+Because the background thread generates with slightly stale model weights, async mode benefits from importance sampling correction (see next section). Enable `vllm_importance_sampling_correction: true` when using `async_prefetch: true`.
+:::
+
+### Streaming partial batch
+
+Instead of scoring the entire batch at once, streaming mode scores one prompt group at a time. This reduces peak memory during scoring and enables finer-grained zero-advantage skipping.
+
+```yaml
+trl:
+  streaming_partial_batch: true
+  streaming_min_groups: 1
+```
+
+`streaming_min_groups` controls the minimum number of prompt groups scored per chunk. Setting it to 1 gives maximum granularity.
+
+### Zero-advantage batch skipping
+
+When all advantages in a micro-batch are zero (every completion in the group got the same reward), there is no learning signal. This feature skips the forward/backward pass entirely for such micro-batches.
+
+```yaml
+trl:
+  skip_zero_advantage_batches: true   # default
+```
+
+This is enabled by default and logged as `skipped_zero_adv_batches` in training metrics. It is a safety net, not a major optimization -- it only saves significant time when the model cannot solve any prompts in the batch.
+
+### Replay buffer
+
+The replay buffer caches rollout groups that had learning signal (non-zero reward variance) and replaces zero-signal groups in later batches. This improves data utilization when many prompts yield no reward variance.
+
+```yaml
+trl:
+  replay_buffer_size: 100
+  replay_recompute_logps: true
+```
+
+:::{.callout-warning}
+When `replay_recompute_logps: false`, replayed data uses stale log-probabilities which creates an IS mismatch. Keep the default `true` unless you have a specific reason to disable it.
+:::
+
+### Deferred re-rolling
+
+Prompts where the model gets zero reward for all generations are buffered and re-injected into later batches, when the model may have improved enough to produce useful completions.
+
+```yaml
+trl:
+  reroll_start_fraction: 0.5   # Start re-rolling after 50% of training
+  reroll_max_groups: 1          # Max groups to replace per batch
+```
+
+Set `reroll_start_fraction: 1.0` to disable. This is most useful for tasks where the model starts weak but steadily improves.
+
+### Parallel reward workers
+
+Reward functions that use `signal.alarm()` (like `math_verify`) only work in the main thread. Parallel reward workers run each function in its own subprocess:
+
+```yaml
+trl:
+  reward_num_workers: 4
+```
+
+Work is sharded across workers by prompt group. For simple reward functions, a single worker is usually sufficient -- the overhead of IPC can exceed the computation time.
+
+
+## Importance Sampling and Off-Policy Correction
+
+When using async prefetch, completions are generated from a slightly older policy. IS correction adjusts the gradient to account for this mismatch.
+
+```yaml
+trl:
+  vllm_importance_sampling_correction: true
+  importance_sampling_level: token     # 'token' recommended (especially with Liger kernel)
+  off_policy_mask_threshold: 0.5       # KL threshold — masks sequences that are too off-policy
+```
+
+Use `token` level IS. Sequence-level has numerical issues with Liger's chunked computation. The `off_policy_mask_threshold` (OPSM) is a safety net that drops sequences where KL divergence exceeds the threshold — 0.5 is a reasonable starting point.
+
+For detailed coverage of IS modes (`token_mask`, `token_truncate`, etc.), capping, and bias-corrected KL, see [vLLM Serving — IS Correction](vllm_serving.qmd#sec-weight-sync).
+
+
+## Scaling
+
+### FP8 training
+
+FP8 quantization halves model VRAM usage with minimal impact on training quality. It does not significantly speed up computation for small models but allows larger models to fit in memory.
+
+```yaml
+fp8: true
+torch_compile: true
+```
+
+:::{.callout-warning}
+FP8 requires patching for zero-padding edge cases. The `act_quant_kernel` can produce NaN when input is all zeros (padding positions). If you see NaN in grad norms, check whether your padding token embedding is non-zero.
+:::
+
+### FSDP (Fully Sharded Data Parallel)
+
+FSDP distributes model parameters across multiple GPUs for training while vLLM runs on a separate GPU:
+
+```yaml
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+```
+
+Launch with:
+
+```bash
+# GPU 0: vLLM
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
+
+# GPUs 0,1: Training (FSDP will use both visible GPUs)
+CUDA_VISIBLE_DEVICES=0,1 axolotl train config.yaml
+```
+
+:::{.callout-warning}
+`async_prefetch: true` can deadlock with FSDP because background threads perform unsynchronized FSDP collectives across ranks. With multi-GPU FSDP, only rank 0 generates in the background thread and results are broadcast to all ranks. If you still see hangs, set `async_prefetch: false`.
+:::
+
+### DeepSpeed ZeRO-3
+
+```yaml
+deepspeed: deepspeed_configs/zero3_bf16.json
+gradient_checkpointing_kwargs:
+  use_reentrant: true   # Required -- non-reentrant causes CheckpointError with ZeRO-3
+```
+
+:::{.callout-note}
+DeepSpeed ZeRO-3 requires `use_reentrant: true` for gradient checkpointing. This is the opposite of the FSDP recommendation. Non-reentrant checkpointing causes tensor metadata mismatches during recomputation with ZeRO-3's parameter partitioning.
+:::
+
+### Multi-GPU considerations
+
+| Concern | Recommendation |
+|---------|---------------|
+| vLLM GPU allocation | Dedicate one or more GPUs to vLLM; do not share with trainer GPUs |
+| Weight sync contention | Use `vllm_lora_sync: true` to avoid NCCL contention between training and vLLM |
+| FSDP + async | Use `async_prefetch: false` or rely on rank-0-only background generation |
+| DeepSpeed + gradient checkpoint | Must use `use_reentrant: true` |
+| OOM during scoring | Reduce `micro_batch_size` or `num_generations`. The logits tensor scales with `batch_size * vocab_size` |
+
+
+## Monitoring and Debugging
+
+For detailed metric ranges, failure diagnosis, and OOM debugging, see [Training Stability & Debugging](training_stability.qmd).
+
+Quick health checks during GRPO training:
+
+- `rewards/*/mean` should be > 0.15 within 20 steps — if it stays at 0, test your reward function standalone
+- `reward_std` should be > 0 on most steps — all-zero means no learning signal
+- `entropy` in 0.05-0.5 — below 0.01 suggests mode collapse
+- `grad_norm` in 0.001-1.0 — > 10 is unstable, 0.0 is expected when zero-advantage skip fires
+
+:::{.callout-tip}
+Pipe training output to a log file: `axolotl train config.yaml 2>&1 | tee /tmp/training.log`
+:::
+
+
+## Configuration Reference
+
+All GRPO-specific options live under the `trl:` key in your config. Standard training options (`learning_rate`, `micro_batch_size`, etc.) are set at the top level as usual.
+
+### Core GRPO
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `use_vllm` | bool | `false` | Enable vLLM for generation |
+| `vllm_mode` | `"server"` or `"colocate"` | `null` | vLLM deployment mode |
+| `vllm_server_host` | str | `"0.0.0.0"` | vLLM server hostname |
+| `vllm_server_port` | int | `8000` | vLLM server port |
+| `vllm_server_timeout` | int | `null` | Timeout (seconds) for vLLM responses |
+| `num_generations` | int | `null` | Completions generated per prompt |
+| `generation_batch_size` | int | `null` | Number of unique prompts per generation step |
+| `max_completion_length` | int | `null` | Maximum tokens per completion |
+| `beta` | float | `null` | KL penalty coefficient |
+| `num_iterations` | int | `null` | Iterations per batch (mu in the GRPO paper) |
+| `epsilon` | float | `null` | PPO clipping lower bound |
+| `epsilon_high` | float | `null` | PPO clipping upper bound |
+| `loss_type` | str | `null` | Loss formulation: `grpo`, `bnpo`, or `dr_grpo` |
+| `scale_rewards` | bool | `true` | Normalize rewards by standard deviation |
+| `mask_truncated_completions` | bool | `false` | Exclude truncated completions from loss |
+
+### Reward functions
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `reward_funcs` | list[str] | `null` | Import paths to reward functions or HF model IDs |
+| `reward_weights` | list[float] | `null` | Relative weights for each reward function |
+| `multi_objective_aggregation` | str | `null` | `"sum_then_normalize"` (GRPO) or `"normalize_then_sum"` (GDPO) |
+| `rollout_func` | str | `null` | Import path to custom rollout function for OpenEnv-style tasks |
+
+### Generation parameters
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `temperature` | float | `null` | Sampling temperature |
+| `top_p` | float | `null` | Nucleus sampling probability |
+| `top_k` | int | `null` | Top-k sampling |
+| `min_p` | float | `null` | Minimum probability threshold |
+| `repetition_penalty` | float | `null` | Penalty for repeated tokens |
+| `generation_kwargs` | dict | `null` | Additional vLLM SamplingParams (e.g., `stop_token_ids`) |
+| `chat_template_kwargs` | dict | `null` | Chat template kwargs (e.g., `{enable_thinking: false}`) |
+| `vllm_guided_decoding_regex` | str | `null` | Regex constraint for guided decoding |
+
+### Async pipeline
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `use_data_producer` | bool | `false` | Enable data producer protocol (required for async features) |
+| `async_prefetch` | bool | `false` | Generate next batch in background thread |
+| `prefetch_depth` | int | `null` | Number of batches to prefetch ahead |
+| `vllm_sync_interval` | int | `null` | Sync LoRA weights to vLLM every N steps |
+| `vllm_lora_sync` | bool | `false` | Use filesystem LoRA sync instead of NCCL merge |
+| `streaming_partial_batch` | bool | `null` | Score prompt groups incrementally |
+| `streaming_min_groups` | int | `null` | Minimum groups per streaming chunk |
+| `skip_zero_advantage_batches` | bool | `true` | Skip micro-batches with zero learning signal |
+| `reward_num_workers` | int | `1` | Subprocess workers for reward computation |
+| `vllm_enable_sleep_mode` | bool | `null` | Offload vLLM weights when idle (colocate mode) |
+
+### Importance sampling
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `vllm_importance_sampling_correction` | bool | `null` | Enable IS correction for async distribution shift |
+| `importance_sampling_level` | `"token"` or `"sequence"` | `null` | Granularity of IS ratios. Use `token` with Liger |
+| `vllm_importance_sampling_mode` | str | `null` | `token_mask`, `token_truncate`, `sequence_mask`, or `sequence_truncate` |
+| `vllm_importance_sampling_cap` | float | `null` | Cap C for IS ratio clipping/masking |
+| `off_policy_mask_threshold` | float | `null` | KL threshold for off-policy sequence masking (OPSM) |
+| `use_bias_correction_kl` | bool | `null` | Apply IS correction to KL divergence term |
+
+### Replay and re-roll
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `replay_buffer_size` | int | `0` | Max cached high-signal groups. 0 = disabled |
+| `replay_recompute_logps` | bool | `true` | Recompute log-probs for replayed data with current model |
+| `reroll_start_fraction` | float | `1.0` | Start re-rolling failed prompts after this fraction of training. 1.0 = disabled |
+| `reroll_max_groups` | int | `1` | Max prompt groups to replace with re-rolls per batch |
+
+### Reference model
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `sync_ref_model` | bool | `false` | Periodically sync reference model with training model |
+| `ref_model_mixup_alpha` | float | `0.9` | EMA coefficient for reference model sync |
+| `ref_model_sync_steps` | int | `64` | Sync reference model every N steps |
+
+### Logging
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `log_completions` | bool | `false` | Log sample completions to W&B |
+| `num_completions_to_print` | int | `null` | Number of completions to print per step |
+| `use_liger_loss` | bool | `null` | Use Liger fused kernel for GRPO loss (reduces VRAM) |
diff --git a/docs/images/ray-cluster-dashboard.png b/docs/images/ray-cluster-dashboard.png
new file mode 100644
index 0000000000..f0b4beb564
Binary files /dev/null and b/docs/images/ray-cluster-dashboard.png differ
diff --git a/docs/inference.qmd b/docs/inference.qmd
new file mode 100644
index 0000000000..e977e3c3c5
--- /dev/null
+++ b/docs/inference.qmd
@@ -0,0 +1,221 @@
+---
+title: "Inference and Merging"
+format:
+  html:
+    toc: true
+    toc-depth: 3
+    number-sections: true
+execute:
+  enabled: false
+---
+
+This guide covers how to use your trained models for inference, including model loading, interactive testing, merging adapters, and common troubleshooting steps.
+
+## Quick Start {#sec-quickstart}
+
+::: {.callout-tip}
+Use the same config used for training on inference/merging.
+:::
+
+### Basic Inference {#sec-basic}
+
+::: {.panel-tabset}
+
+## LoRA Models
+
+```{.bash}
+axolotl inference your_config.yml --lora-model-dir="./lora-output-dir"
+```
+
+## Full Fine-tuned Models
+
+```{.bash}
+axolotl inference your_config.yml --base-model="./completed-model"
+```
+
+:::
+
+### Interactive Chat {#sec-chat}
+
+For multi-turn testing of conversational models, use chat mode. The chat template
+is resolved exactly as it was during training and re-applied to the full
+conversation each turn:
+
+```{.bash}
+axolotl inference your_config.yml --chat
+```
+
+Type a message to chat. End a line with `\` to continue typing on the next line.
+Slash commands control the session:
+
+| Command | Aliases | Description |
+|---------|---------|-------------|
+| `/help` | `/?` | Show all commands |
+| `/new` | `/clear`, `/reset` | Clear the conversation (keeps system prompt and parameters) |
+| `/system [text\|clear]` | | Show, set, or clear the system prompt |
+| `/set <param> <value>` | | Set a generation parameter |
+| `/status` | `/params` | Show model info and current settings |
+| `/history` | | Show the conversation so far |
+| `/retry` | `/regen` | Regenerate the last assistant reply |
+| `/undo` | | Remove the last exchange |
+| `/save [path]` | | Append the conversation as a `chat_template`-format JSONL sample |
+| `/quit` | `/exit`, `/q` | Exit |
+
+Generation parameters can also be set directly, e.g. `/temperature 0.7` (or
+`/temp 0.7`), `/top_p 0.9`, `/top_k 50`, `/max_tokens 512`, `/rep 1.05`,
+`/seed 42`. Setting `temperature` to `0` switches to greedy decoding.
+
+Press `Ctrl+C` during generation to stop the current reply; the partial response
+is kept in the conversation (diffusion replies denoise in one piece, so an
+interrupted diffusion turn is discarded instead).
+
+#### Thinking Models {#sec-chat-thinking}
+
+Thinking blocks (e.g. `<think>...</think>`) stream live in a small dim window,
+then collapse to a one-line summary — `/expand` shows the full reasoning of the
+last reply, and `/collapse off` switches to raw verbatim output. The per-turn
+stats split thinking from reply tokens. If the chat template supports a
+render-time thinking toggle (e.g. Qwen's `enable_thinking`), `/think off`
+disables thinking entirely from the next turn; `/think default` restores the
+template default.
+
+::: {.callout-note}
+Assistant turns are stored the way `transformers` recommends: special tokens
+are stripped and thinking is kept on a separate `reasoning_content` key (via
+the tokenizer's `parse_response` schema when it ships one, marker-splitting
+otherwise), so the chat template decides how prior-turn reasoning is
+re-rendered — matching what the model saw during training. The KV cache is
+re-used across turns whenever the rendered conversation extends the previous
+one, so long chats stay responsive.
+:::
+
+`/save` writes conversations in the `messages` format accepted by
+`type: chat_template` datasets, so a good interactive session can be turned
+directly into training data.
+
+#### Diffusion Models {#sec-chat-diffusion}
+
+With the diffusion plugin enabled, chat mode generates each reply by appending
+a masked block to the conversation and denoising it. Replies arrive in one
+piece (no token streaming), and the parameter set changes accordingly:
+`/tokens N` sets the completion block size, `/steps N` the number of denoising
+steps, and `/temperature` the denoising temperature. Defaults come from the
+`diffusion:` section of your config.
+
+Chat mode is not supported with `--prompter`; use the default inference mode
+for legacy prompters.
+
+## Advanced Usage {#sec-advanced}
+
+### Gradio Interface {#sec-gradio}
+
+Launch an interactive web interface:
+
+```{.bash}
+axolotl inference your_config.yml --gradio
+```
+
+### File-based Prompts {#sec-file-prompts}
+
+Process prompts from a text file:
+
+```{.bash}
+cat /tmp/prompt.txt | axolotl inference your_config.yml \
+  --base-model="./completed-model" --prompter=None
+```
+
+### Memory Optimization {#sec-memory}
+
+For large models or limited memory:
+
+```{.bash}
+axolotl inference your_config.yml --load-in-8bit=True
+```
+
+## Merging LoRA Weights {#sec-merging}
+
+Merge LoRA adapters with the base model:
+
+```{.bash}
+axolotl merge-lora your_config.yml --lora-model-dir="./completed-model"
+```
+
+### Memory Management for Merging {#sec-memory-management}
+
+::: {.panel-tabset}
+
+## Configuration Options
+
+```{.yaml}
+gpu_memory_limit: 20GiB  # Adjust based on your GPU
+lora_on_cpu: true        # Process on CPU if needed
+```
+
+## Force CPU Merging
+
+```{.bash}
+CUDA_VISIBLE_DEVICES="" axolotl merge-lora ...
+```
+
+:::
+
+## Tokenization {#sec-tokenization}
+
+### Common Issues {#sec-tokenization-issues}
+
+::: {.callout-warning}
+Tokenization mismatches between training and inference are a common source of problems.
+:::
+
+To debug:
+
+1. Check training tokenization:
+```{.bash}
+axolotl preprocess your_config.yml --debug
+```
+
+2. Verify inference tokenization by decoding tokens before model input
+
+3. Compare token IDs between training and inference
+
+### Special Tokens {#sec-special-tokens}
+
+Configure special tokens in your YAML:
+
+```{.yaml}
+special_tokens:
+  bos_token: "<s>"
+  eos_token: "</s>"
+  unk_token: "<unk>"
+tokens:
+  - "<|im_start|>"
+  - "<|im_end|>"
+```
+
+## Troubleshooting {#sec-troubleshooting}
+
+### Common Problems {#sec-common-problems}
+
+::: {.panel-tabset}
+
+## Memory Issues
+
+- Use 8-bit loading
+- Reduce batch sizes
+- Try CPU offloading
+
+## Token Issues
+
+- Verify special tokens
+- Check tokenizer settings
+- Compare training and inference preprocessing
+
+## Performance Issues
+
+- Verify model loading
+- Check prompt formatting
+- Ensure temperature/sampling settings
+
+:::
+
+For more details, see our [debugging guide](debugging.qmd).
diff --git a/docs/input_output.qmd b/docs/input_output.qmd
index 3762901b31..f9d2df2336 100644
--- a/docs/input_output.qmd
+++ b/docs/input_output.qmd
@@ -3,263 +3,4 @@ title: Template-free prompt construction
 description: "Template-free prompt construction with the `input_output` format"
 ---
 
-<!-- TOC -->
-
-- [Background](#background)
-    - [Masking Inputs](#masking-inputs)
-    - [You may not want prompt templates](#you-may-not-want-prompt-templates)
-    - [The `input_output` format](#the-input_output-format)
-- [Usage](#usage)
-    - [1. Prepare Data](#1-prepare-data)
-    - [2. Use `type: input_output`](#2-use-type-input_output)
-    - [3. Check the prompts](#3-check-the-prompts)
-
-<!-- /TOC -->
-
-<a id="markdown-background" name="background"></a>
-
-## Background
-
-<a id="markdown-masking-inputs" name="masking-inputs"></a>
-
-### Masking Inputs
-
-One of the most popular features of
-[axolotl](https://github.com/OpenAccess-AI-Collective/axolotl) is
-setting the following configuration value:
-
-
-```yaml
-train_on_inputs: false
-```
-
-If you declare a [dataset formats](https://github.com/OpenAccess-AI-Collective/axolotl?tab=readme-ov-file#dataset)
-such as `alpaca` or `chatml`, axolotl knows what is an input
-(i.e. human) vs. an output (i.e. the assistant) and masks the input
-labels so that your model can focus on predicting the outputs only.
-
-<a id="markdown-you-may-not-want-prompt-templates" name="you-may-not-want-prompt-templates"></a>
-
-### You may not want prompt templates
-
-However, there are many situations where you don't want to use one of
-these formats or templates. This is because they can:
-
--   Add unnecessary boilerplate to your prompts.
--   Create artifacts like special delimiters `<|im_start|>` that can
-    quickly become footguns if you don't include them correctly at
-    inference time.
--   Enforce a *chat* interface when you do not want one. Sometimes you
-    just want to fine-tune a model to a very specific task and do NOT
-    want multi-turn conversations, roles, etc.
--   Limit you to only certain roles that the template allows.
-
-<a id="markdown-the-inputoutput-format" name="the-inputoutput-format"></a>
-
-### The `input_output` format
-
-You can construct your prompts without a template by using the
-`input_output` format, by setting `type: input_output` in your
-configuration file like this:
-
-**config.yml**
-
-```yaml
-train_on_inputs: false # Mask segments of your data
-datasets:
-  - path: output.jsonl
-    type: input_output  # use template free prompt construction
-```
-
-Unlike `type: completion`, which is also template-free,
-`type: input_output` allows you to mask segments of your text. More
-details on how this works are described below.
-
-<a id="markdown-usage" name="usage"></a>
-
-## Usage
-
-This is how you can use the `input_output` format:
-
-<a id="markdown-1-prepare-data" name="1-prepare-data"></a>
-
-### 1. Prepare Data
-
-To use the `input_output` format, collect your data in the following
-format into a jsonl file (below is the first row from the file
-`output`.jsonl` pretty printed):
-
-```bash
-$ head -n1 output.jsonl | python -m json.tool
-```
-
-:::{.cell-output .cell-output-stdout}
-    {
-        "segments": [
-            {
-                "label": true,
-                "text": "<s>Hello\n"
-            },
-            {
-                "label": true,
-                "text": "hi there!. "
-            },
-            {
-                "label": false,
-                "text": "goodbye "
-            },
-            {
-                "label": true,
-                "text": "farewell</s>"
-            }
-        ]
-    }
-:::
-
-Set `label:false` when you want to mask a segment of text so that the
-model isn't trained on it. Some things to keep in mind:
-
-> [!IMPORTANT]
-> 1.  **EOS, BOS, spaces, newlines etc. are entirely up to you. Axolotl
-    concatenates all the segments as-is.** The tokenizer doesn't add
-    anything additional. Notice how I added spaces, newlines, `<s>`
-    (BOS), and `</s>` (EOS) myself.
-> 2.  Make sure you check the materialized output to validate that the
-    prompt is getting assembled how you like.
-
-<a id="markdown-2-use-type-inputoutput" name="2-use-type-inputoutput"></a>
-
-### 2. Use `type: input_output`
-
-Let's materialize data with our `output.jsonl` file by setting
-`type: input_output` in our axolotl config:
-
-```yaml
-# training_config.yaml
-base_model: mistralai/Mistral-7B-v0.1
-data_seed: 49
-seed: 49
-
-datasets:
-  - path: output.jsonl
-    type: input_output
-val_set_size: 0.1
-
-sequence_len: 896
-sample_packing: false
-
-micro_batch_size: 2
-gradient_accumulation_steps: 3
-eval_batch_size: 2
-num_epochs: 1
-learning_rate: 0.0002
-
-train_on_inputs: false
-special_tokens:
-  bos_token: "<s>"
-  eos_token: "</s>"
-  unk_token: "<unk>"
-```
-
-You can use the following command to materialize your data. The
-`--debug` flag will print the tokens, along with the labels so you can
-verify that the correct items are being ignored:
-
-```bash
-$ python -m axolotl.cli.preprocess training_config.yaml --debug
-
-...
-[2024-03-05 23:36:46,969] [INFO] [axolotl.check_example_labels:35] [PID:607731] [RANK:0] <s>(1, 1) Hello(22557, 22557)
-(13, 13) hi(12014, 12014) there(736, 736) !(28808, 28808) .(28723, 28723) (28705, 28705) good(-100, 1179) bye(-100, 17664) (-100, 28705) fare(19111, 19111) well(5458, 5458) </s>(2, 2)
-
-```
-
-The format is `decoded_token`(`label`, `token_id`), for example,
-`<s>(1, 1)` means that the token is `<s>`, the label is `1` and the
-token_id is `1`. When the label is `-100` then that token is ignored for
-training.
-
-<a id="markdown-3-check-the-prompts" name="3-check-the-prompts"></a>
-
-### 3. Check the prompts
-
-Here is another way to check the materialized output:
-
-```python
-from transformers import AutoTokenizer
-from datasets import load_from_disk
-import yaml
-
-directory = !ls last_run_prepared/
-with open('training_config.yaml', 'r') as f:
-    cfg = yaml.safe_load(f)
-model_id = cfg['base_model']
-tok = AutoTokenizer.from_pretrained(model_id)
-ds = load_from_disk(f'last_run_prepared/{directory[0]}/')
-```
-
-```python
->>> row = ds[0]
->>> print(tok.decode(row['input_ids']))
-<s> Hello
-    hi there!.  goodbye  farewell</s>
-```
-
-We can check that the right tokens are ingored by comparing the labels
-to each token:
-
-```python
-import pandas as pd
-pd.DataFrame([{'token': tok.decode(i), 'label': l, 'id':i} for i,l in
-              zip(row['input_ids'], row['labels'])])
-```
-
-| token | label | id    |
-|-------|-------|-------|
-| 0     | \<s\> | 1     |
-| 1     | Hello | 22557 |
-| 2     | \\n   | 13    |
-| 3     | hi    | 12014 |
-| 4     | there | 736   |
-| 5     | !     | 28808 |
-| 6     | .     | 28723 |
-| 7     |       | 28705 |
-| 8     | good  | -100  |
-| 9     | bye   | -100  |
-| 10    |       | -100  |
-| 11    | fare  | 19111 |
-| 12    | well  | 5458  |
-| 13    | \</s\>| 2     |
-
-
-
-If we look at the input data, the above table seems correct! (The jsonl
-version is repeated below for reference):
-
-
-```bash
-$ head -n1 output.jsonl | python -m json.tool
-```
-
-:::{.cell-output .cell-output-stdout}
-    {
-        "segments": [
-            {
-                "label": true,
-                "text": "<s>Hello\n"
-            },
-            {
-                "label": true,
-                "text": "hi there!. "
-            },
-            {
-                "label": false,
-                "text": "goodbye "
-            },
-            {
-                "label": true,
-                "text": "farewell</s>"
-            }
-        ]
-    }
-:::
+The documentation moved to [here](dataset-formats/template_free.qmd).
diff --git a/docs/installation.qmd b/docs/installation.qmd
new file mode 100644
index 0000000000..830d243455
--- /dev/null
+++ b/docs/installation.qmd
@@ -0,0 +1,166 @@
+---
+title: "Installation"
+format:
+  html:
+    toc: true
+    toc-depth: 3
+    number-sections: true
+execute:
+  enabled: false
+---
+
+This guide covers all the ways you can install and set up Axolotl for your environment.
+
+## Requirements {#sec-requirements}
+
+- NVIDIA GPU (Ampere architecture or newer for `bf16` and Flash Attention) or AMD GPU
+- Python ≥3.11
+- PyTorch ≥2.11.0
+
+## Installation {#sec-installation}
+
+::: {.callout-important}
+For Blackwell GPUs, please use PyTorch ≥2.11.0 and CUDA 13.0. CUDA 12.8 cannot
+compile for `sm_103a` (B300), and using CUDA 13.0 also avoids first-launch JIT
+overhead on B100 / B200 / RTX 50-series.
+:::
+
+### Quick Install {#sec-uv}
+
+Axolotl uses [uv](https://docs.astral.sh/uv/) as its package manager. uv is a fast, reliable Python package installer and resolver built in Rust.
+
+Install uv if not already installed:
+```{.bash}
+curl -LsSf https://astral.sh/uv/install.sh | sh
+source $HOME/.local/bin/env
+```
+
+Choose your CUDA version (e.g. `cu128`, `cu130`), create a venv, and install:
+```{.bash}
+export UV_TORCH_BACKEND=cu130  # or cu128
+uv venv
+source .venv/bin/activate
+uv pip install --no-build-isolation axolotl[deepspeed]
+```
+
+### Edge/Development Build {#sec-edge-build}
+
+For the latest features between releases:
+
+```{.bash}
+git clone https://github.com/axolotl-ai-cloud/axolotl.git
+cd axolotl
+export UV_TORCH_BACKEND=cu130  # or cu128
+uv venv
+source .venv/bin/activate
+uv pip install --no-build-isolation -e '.[deepspeed]'
+```
+
+### Docker {#sec-docker}
+
+```{.bash}
+docker run --gpus '"all"' --rm -it --ipc=host axolotlai/axolotl-uv:main-latest
+```
+
+For development with Docker:
+
+```{.bash}
+docker compose up -d
+```
+
+::: {.callout-tip}
+### Advanced Docker Configuration
+```{.bash}
+docker run --privileged --gpus '"all"' --shm-size 10g --rm -it \
+  --name axolotl --ipc=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  --mount type=bind,src="${PWD}",target=/workspace/axolotl \
+  -v ${HOME}/.cache/huggingface:/root/.cache/huggingface \
+  axolotlai/axolotl-uv:main-latest
+```
+:::
+
+::: {.callout-important}
+For Blackwell GPUs, please use `axolotlai/axolotl-uv:main-py3.12-cu130-2.12.0` or the cloud variant `axolotlai/axolotl-cloud-uv:main-py3.12-cu130-2.12.0`.
+:::
+
+Please refer to the [Docker documentation](docker.qmd) for more information on the different Docker images that are available.
+
+## Cloud Environments {#sec-cloud}
+
+### Cloud GPU Providers {#sec-cloud-gpu}
+
+For providers supporting Docker:
+
+- Use `axolotlai/axolotl-cloud-uv:main-latest`
+- Available on:
+    - [RunPod](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
+    - [Vast.ai](https://cloud.vast.ai?ref_id=62897&template_id=bdd4a49fa8bce926defc99471864cace&utm_source=axolotl&utm_medium=partner&utm_campaign=template_launch_july2025&utm_content=docs_link)
+    - [PRIME Intellect](https://app.primeintellect.ai/dashboard/create-cluster?image=axolotl&location=Cheapest&security=Cheapest&show_spot=true)
+    - [Modal](https://www.modal.com?utm_source=github&utm_medium=github&utm_campaign=axolotl)
+    - [Novita](https://novita.ai/gpus-console?templateId=311)
+    - [JarvisLabs.ai](https://jarvislabs.ai/templates/axolotl)
+    - [Latitude.sh](https://latitude.sh/blueprint/989e0e79-3bf6-41ea-a46b-1f246e309d5c)
+
+### Google Colab {#sec-colab}
+
+[![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/axolotl-ai-cloud/axolotl/blob/main/examples/colab-notebooks/colab-axolotl-example.ipynb#scrollTo=msOCO4NRmRLa)
+
+## Platform-Specific Instructions {#sec-platform-specific}
+
+### macOS {#sec-macos}
+
+```{.bash}
+uv pip install --no-build-isolation -e '.'
+```
+
+See @sec-troubleshooting for Mac-specific issues.
+
+### Windows {#sec-windows}
+
+::: {.callout-important}
+We recommend using WSL2 (Windows Subsystem for Linux) or Docker.
+:::
+
+## Migrating from pip to uv {#sec-migrating}
+
+If you have an existing pip-based Axolotl installation, you can migrate to uv:
+
+```{.bash}
+# Install uv
+curl -LsSf https://astral.sh/uv/install.sh | sh
+source $HOME/.local/bin/env
+
+# Create a fresh venv (recommended for a clean start)
+export UV_TORCH_BACKEND=cu130  # or cu128
+uv venv
+source .venv/bin/activate
+
+# Reinstall axolotl
+uv pip install --no-build-isolation axolotl[deepspeed]
+```
+
+## Using pip (Alternative) {#sec-pip}
+
+If you are unable to install uv, you can still use pip directly.
+
+::: {.callout-important}
+Please make sure to have PyTorch installed before installing Axolotl with pip.
+
+Follow the instructions at: [https://pytorch.org/get-started/locally/](https://pytorch.org/get-started/locally/)
+:::
+
+```{.bash}
+pip3 install -U packaging setuptools wheel ninja
+pip3 install --no-build-isolation axolotl[deepspeed]
+```
+
+For editable/development installs:
+```{.bash}
+pip3 install -U packaging setuptools wheel ninja
+pip3 install --no-build-isolation -e '.[deepspeed]'
+```
+
+## Troubleshooting {#sec-troubleshooting}
+
+If you encounter installation issues, see our [FAQ](faq.qmd) and [Debugging Guide](debugging.qmd).
diff --git a/docs/lora.qmd b/docs/lora.qmd
new file mode 100644
index 0000000000..576bf48277
--- /dev/null
+++ b/docs/lora.qmd
@@ -0,0 +1,57 @@
+---
+title: "LoRA"
+description: "Per-module LoRA rank and alpha overrides"
+---
+
+## Per-module rank and alpha (`lora_rank_pattern` / `lora_alpha_pattern`)
+
+By default `lora_r` and `lora_alpha` apply uniformly to every targeted module.
+Two optional config fields let you override them on a per-module basis:
+
+- `lora_rank_pattern` — `dict[str, int]` mapping a module-name regex to a rank
+  that overrides `lora_r` for the matched modules.
+- `lora_alpha_pattern` — `dict[str, int]` mapping the same kind of regex to an
+  alpha that overrides `lora_alpha`.
+
+Both fields are forwarded directly to PEFT (`LoraConfig.rank_pattern` /
+`LoraConfig.alpha_pattern`). Pattern keys are matched against module paths using
+the same suffix-anchored rule PEFT uses internally, so a key like
+`"layers.0.self_attn.q_proj"` matches any module path ending in that suffix.
+
+### When to use it
+
+- **Asymmetric attention vs MLP ranks.** Raise the rank on attention projections
+  while keeping MLPs at a smaller rank, or vice versa.
+- **Multimodal towers.** Use a larger rank on a vision tower than on the language
+  tower when one modality benefits from more adapter capacity than the other.
+
+### Example
+
+```yaml
+adapter: lora
+lora_r: 8
+lora_alpha: 16
+lora_target_modules:
+  - q_proj
+  - v_proj
+  - up_proj
+  - down_proj
+
+# Bump attention projections to rank 32 / alpha 64, leave MLPs at the global 8/16.
+lora_rank_pattern:
+  ".*\\.(q_proj|v_proj)$": 32
+lora_alpha_pattern:
+  ".*\\.(q_proj|v_proj)$": 64
+```
+
+### Resuming and continuing training
+
+- `resume_from_checkpoint` rebuilds the model from your YAML, so keep the patterns identical to the original run; a changed `lora_rank_pattern` fails at load with a shape mismatch.
+- `lora_model_dir` (continue-training from a saved adapter) takes per-module rank/alpha from the adapter's saved `adapter_config.json`; YAML patterns cannot override it, and axolotl warns if they differ.
+
+### Merging adapters trained with patterns
+
+`axolotl merge-lora` (default `merge_method: memory_efficient`) honors both
+fields when merging: each module is merged with its own rank (inferred from
+the adapter weight shapes) and its own alpha (resolved against
+`alpha_pattern`). No extra flags are needed.
diff --git a/docs/lora_optims.qmd b/docs/lora_optims.qmd
new file mode 100644
index 0000000000..7006b5d19b
--- /dev/null
+++ b/docs/lora_optims.qmd
@@ -0,0 +1,140 @@
+---
+title: "LoRA Optimizations"
+description: "Custom autograd functions and Triton kernels in Axolotl for optimized LoRA fine-tuning"
+---
+
+Inspired by [Unsloth](https://github.com/unslothai/unsloth), we've implemented two
+optimizations for LoRA and QLoRA fine-tuning, supporting both single GPU and multi-GPU
+(including the DDP, DeepSpeed, and FSDP2 settings) training. These include (1) SwiGLU
+and GEGLU activation function Triton kernels, and (2) LoRA MLP and attention custom
+autograd functions. Our goal was to leverage operator fusion and tensor re-use in order
+to improve speed and reduce memory usage during the forward and backward passes of
+these calculations.
+
+We currently support several common model architectures, including (but not limited to):
+
+- `llama`
+- `mistral`
+- `qwen2`
+- `gemma`
+- `gemma2`
+- `gemma3`
+
+<details>
+
+The set of models we support is currently limited by our attention patching strategy,
+which assumes (and replaces) specific code blocks for query / key / value and output
+projections:
+
+```python
+ORIGINAL_QKV_CODE = """
+    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+    key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+""".lstrip(
+    "\n"
+)
+
+ORIGINAL_O_CODE = """
+    attn_output = self.o_proj(attn_output)
+""".lstrip(
+    "\n"
+)
+```
+
+Is replaced with:
+
+```python
+PATCHED_QKV_CODE = """
+    query_states, key_states, value_states = self.apply_qkv(hidden_states)
+    query_states = query_states.view(hidden_shape).transpose(1, 2)
+    key_states = key_states.view(hidden_shape).transpose(1, 2)
+    value_states = value_states.view(hidden_shape).transpose(1, 2)
+""".lstrip(
+    "\n"
+)
+
+PATCHED_O_CODE = """
+    attn_output = self.apply_o(attn_output)
+""".lstrip(
+    "\n"
+)
+```
+
+Where `apply_qkv` and `apply_o` are defined in the `axolotl.kernels.lora` module.
+
+We welcome testing of other model architectures and / or PRs to expand our patching
+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
+`lora_mlp_kernel` option enables the optimized MLP path, while `lora_qkv_kernel` and
+`lora_o_kernel` enable the fused query-key-value projection and optimized output
+projection, respectively.
+
+```yaml
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+```
+
+::: {.callout-note}
+Currently, LoRA kernels are not supported for RLHF training, only SFT.
+:::
+
+::: {.callout-warning}
+LoRA kernels do not support remote modeling code.
+:::
+
+## Requirements
+
+- One or more NVIDIA or AMD GPUs (in order to use the Triton kernels)
+    - Note: Set `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` to enable [memory-efficient attention on AMD GPUs](https://github.com/ROCm/aotriton/issues/16#issuecomment-2346675491)
+- Targeted LoRA adapters cannot use Dropout
+    - This may limit model expressivity / cause overfitting
+- Targeted LoRA adapters cannot have bias terms
+    - This may limit model expressivity
+
+Models with pre-existing LoRA adapters that use Dropout or have bias terms may need to
+be re-finetuned without these features in order to be useful.
+
+## Implementation details
+
+### Custom autograd functions
+
+The LoRA MLP autograd function optimizes the entire MLP computation path. It fuses the
+LoRA and base weight computations together and provides a single, efficient backward
+pass for the entire MLP block.
+
+For attention components, similar optimizations are provided through a function that
+handles the query, key, and value projections, and a function that handles the output
+projection. They are designed to work with the existing `transformers` attention
+implementation via some monkey-patching logic.
+
+### Triton kernels
+
+Two activation functions (SwiGLU and GeGLU) are implemented with Triton kernels for
+improved speed and memory performance. These kernels handle both the forward and
+backward passes.
+
+### Integration
+
+The custom autograd functions and Triton kernels are designed to work together. The
+autograd function manages the high-level computation flow and gradient tracking, while
+calling the Triton kernels for the activation function computation. During the backward
+pass, the kernel computes both the activation output and the required gradients, which
+the autograd function then uses to compute the final gradients for the entire
+computation path.
+
+## Future Work
+
+- Support for additional model architectures
+- Support for dropout and bias
+- Additional operator fusions
diff --git a/docs/lr_groups.qmd b/docs/lr_groups.qmd
new file mode 100644
index 0000000000..ce53507223
--- /dev/null
+++ b/docs/lr_groups.qmd
@@ -0,0 +1,35 @@
+---
+title: Learning Rate Groups
+description: "Setting different learning rates by module name"
+---
+
+## Background
+
+Inspired by LoRA+, Axolotl allows practitioners to specify separate learning rates for each module or groups of
+modules in a model.
+
+## Example
+
+```yaml
+lr_groups:
+  - name: o_proj
+    modules:
+      - self_attn.o_proj.weight
+    lr: 1e-6
+  - name: q_proj
+    modules:
+      - model.layers.2.self_attn.q_proj.weight
+    lr: 1e-5
+
+learning_rate: 2e-5
+```
+
+In this example, we have a default learning rate of 2e-5 across the entire model, but we have a separate learning rate
+of 1e-6 for all the self attention `o_proj` modules across all layers, and a learning are of 1e-5 to the 3rd layer's
+self attention `q_proj` module.
+
+::: {.callout-note}
+
+We currently only support varying `lr` for now. If you're interested in adding support for others (`weight_decay`), we welcome PRs. See https://github.com/axolotl-ai-cloud/axolotl/blob/613bcf90e58f3ab81d3827e7fc572319908db9fb/src/axolotl/core/trainers/mixins/optimizer.py#L17
+
+:::
diff --git a/docs/mac.qmd b/docs/mac.qmd
index 2a83035381..2e6c5c429e 100644
--- a/docs/mac.qmd
+++ b/docs/mac.qmd
@@ -19,4 +19,5 @@ Current support:
 - [ ] DeepSpeed
 
 Untested:
+
 - FSDP
diff --git a/docs/mixed_precision.qmd b/docs/mixed_precision.qmd
new file mode 100644
index 0000000000..ac0f668029
--- /dev/null
+++ b/docs/mixed_precision.qmd
@@ -0,0 +1,169 @@
+---
+title: "Mixed Precision Training"
+format:
+  html:
+    toc: true
+    toc-depth: 3
+    number-sections: true
+    code-tools: true
+execute:
+  enabled: false
+---
+
+Mixed precision training uses lower precision data types to reduce memory usage and increase training speed while maintaining model quality. Axolotl supports several mixed precision formats:
+
+- **FP16** - Half precision 16-bit (Pascal generation+)
+- **BF16** - Brain Float 16-bit (Ampere generation+)
+- **FP8** - 8-bit floating point (Hopper generation+)
+
+## FP16 Mixed Precision {#sec-fp16}
+
+### Overview {#sec-fp16-overview}
+
+FP16 is the traditional half-precision format, supported on older GPUs but can be less numerically stable than BF16.
+
+### Configuration {#sec-fp16-config}
+
+```{.yaml}
+fp16: true
+```
+
+### FP16 Considerations {#sec-fp16-considerations}
+
+- May require gradient scaling to prevent underflow
+- Less numerically stable than BF16
+- Can cause training instability with some model architectures
+- Consider using BF16 if your hardware supports it
+
+## BF16 Mixed Precision {#sec-bf16}
+
+### Overview {#sec-bf16-overview}
+
+BF16 (Brain Float 16) offers better numerical stability than FP16 and is the recommended mixed precision format for modern GPUs. It provides the same dynamic range as FP32 while using half the memory.
+
+### Configuration {#sec-bf16-config}
+
+```{.yaml}
+# Automatic BF16 detection (recommended)
+bf16: auto
+
+# Or explicitly enable
+bf16: true
+
+# For evaluation with BF16
+bf16: full  # Equivalent to bf16_full_eval in the HF trainer
+```
+
+### Keeping norms in fp32 (FSDP2) {#sec-fp32-norms}
+
+Some models declare RMSNorm/LayerNorm layers as fp32 for training
+stability — the variance computation in RMSNorm is numerically poor in
+bf16, and the learned gain γ quantizes harshly. With FSDP1 this fights
+the flat-param dtype uniformity constraint; with FSDP2 each norm can have
+its own `MixedPrecisionPolicy`. Enable with:
+
+```{.yaml}
+fsdp_version: 2
+fp32_norms: true
+# fp32_norm_classes:        # optional override
+#   - RMSNorm
+#   - LayerNorm
+```
+
+Defaults match any class whose name ends in `RMSNorm` or `LayerNorm`. Use
+fully qualified names (`module.path.ClassName`) to pin a specific
+implementation.
+
+## FP8 Mixed Precision {#sec-fp8}
+
+::: {.callout-note}
+FP8 support is experimental and requires compatible hardware (H100, H200) and recent PyTorch versions with TorchAO.
+:::
+
+### What is FP8? {#sec-fp8-overview}
+
+FP8 (8-bit floating point) can provide significant time savings compared to FP16/BF16 while maintaining training stability. Axolotl's implementation uses PyTorch's TorchAO library with "tensorwise" scaling strategy.
+
+### Requirements {#sec-fp8-software}
+
+- Hopper+ GPUs (H100/H200)
+- PyTorch 2.7+ (+ compatible TorchAO version)
+- CUDA 12.4+
+
+### Configuration {#sec-fp8-config}
+
+Add to your YAML config:
+
+```{.yaml}
+# Enable FP8 mixed precision
+fp8: true
+
+# Optional: Enable FP8 for FSDP all-gather operations
+fp8_enable_fsdp_float8_all_gather: true
+
+# Enable torch.compile (almost always necessary for FP8 speedups)
+torch_compile: true
+```
+
+::: {.callout-important}
+**torch.compile is critical for FP8 performance**
+
+FP8 training requires `torch_compile: true` to see meaningful speedups. Without compilation, FP8 may actually be slower and use more memory than FP16/BF16.
+:::
+
+### Advanced FP8 Configs {#sec-fp8-advanced}
+
+For [FSDP](multi-gpu.qmd#sec-fsdp) (Fully Sharded Data Parallel) training:
+
+```{.yaml}
+fp8: true
+fp8_enable_fsdp_float8_all_gather: true
+
+torch_compile: true
+
+# FSDP configuration
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  reshard_after_forward: true
+```
+
+## Best Practices {#sec-best-practices}
+
+### Choosing Precision Format {#sec-choosing-format}
+
+- **Start with automatic detection**: `bf16: auto`
+- **For Hopper+ (H100/H200)**: Try FP8 + torch.compile for maximum speed
+- **For Ampere (A100/RTX 30/40)**: Use BF16
+- **For older Pascal/Turing GPUs**: Use FP16 with caution
+- **For very old or unsupported GPUs**: Use FP32
+
+### Validation and Testing {#sec-validation}
+
+Always validate your mixed precision setup:
+
+- **Start with a small dataset** to verify stability
+- **Monitor loss curves** for irregularities
+- **Compare with FP32 baseline** when possible
+- **Test evaluation metrics** match expectations
+
+### FP8 Particulars {#sec-fp8-details}
+
+- Use cases
+  - Single GPU training
+  - Multi GPU training with FSDP2 or Deepspeed
+- Speedups
+  - Please refer to the [TorchAO FP8 training benchmarks](https://github.com/pytorch/ao/tree/main/torchao/float8#rowwise-scaling) for expected matmul speedups for different (M, K, N) settings
+  - Concrete number for LLaMA 3 8B training can be found [here](https://github.com/pytorch/ao/tree/main/torchao/float8#training-benchmarks)
+- Known issues:
+  - FP8 + DDP + `torch.compile` (causes [error](https://gist.github.com/djsaunde/0c1664c32e44a64d31b5e01b4aafe5c4))
+  - FP8 + FSDP2 + `torch.compile` + FSDP2 activation checkpointing tends to be _slower_ than the BF16 equivalent training
+  - Flash Attention 2 does not play nicely with `torch.compile`
+
+See `examples/llama-3/3b-fp8-fsdp2.yaml` for an optimized example config. Enabling FP8 mixed precision + FP8 all-gather training results in ~10% faster iterations per second vs. BF16 for a relatively small (3B param) model
+
+For more information on multi-GPU training, see our [Multi-GPU guide](multi-gpu.qmd).
diff --git a/docs/multi-gpu.qmd b/docs/multi-gpu.qmd
new file mode 100644
index 0000000000..c949a9ebfa
--- /dev/null
+++ b/docs/multi-gpu.qmd
@@ -0,0 +1,193 @@
+---
+title: "Multi-GPU"
+format:
+  html:
+    toc: true
+    toc-depth: 3
+    # number-sections: true
+    code-tools: true
+execute:
+  enabled: false
+---
+
+This guide covers advanced training configurations for multi-GPU setups using Axolotl.
+
+## Overview {#sec-overview}
+
+When training on multiple GPUs, Axolotl supports 3 sharding/parallelism strategies. Additionally, you can layer specific optimization features on top of that strategy.
+
+You generally cannot combine these strategies; they are mutually exclusive.
+
+1.  **DeepSpeed**: Powerful optimization library, supports ZeRO stages 1-3.
+2.  **FSDP (Fully Sharded Data Parallel)**: PyTorch's native sharding implementation (Recommended).
+3.  **DDP (Distributed Data Parallel)**: PyTorch's native parallelism implementation (Default if neither of the above are selected).
+
+These features can often be combined with the strategies above:
+
+*   **Sequence Parallelism**: Splits long sequences across GPUs (Compatible with DDP, DeepSpeed, and FSDP).
+*   **FSDP + QLoRA**: Combines 4-bit quantization with FSDP (Specific to FSDP).
+
+## DeepSpeed {#sec-deepspeed}
+
+### Configuration {#sec-deepspeed-config}
+
+Add to your YAML config:
+
+```{.yaml}
+deepspeed: deepspeed_configs/zero1.json
+```
+### Usage {#sec-deepspeed-usage}
+
+```{.bash}
+# Fetch deepspeed configs (if not already present)
+axolotl fetch deepspeed_configs
+
+# Passing arg via config
+axolotl train config.yml
+
+# Passing arg via cli
+axolotl train config.yml --deepspeed deepspeed_configs/zero1.json
+```
+
+### ZeRO Stages {#sec-zero-stages}
+
+We provide default configurations for:
+
+- ZeRO Stage 1 (`zero1.json`)
+- ZeRO Stage 1 with torch compile (`zero1_torch_compile.json`)
+- ZeRO Stage 2 (`zero2.json`)
+- ZeRO Stage 3 (`zero3.json`)
+- ZeRO Stage 3 with bf16 (`zero3_bf16.json`)
+- ZeRO Stage 3 with bf16 and CPU offload params(`zero3_bf16_cpuoffload_params.json`)
+- ZeRO Stage 3 with bf16 and CPU offload params and optimizer (`zero3_bf16_cpuoffload_all.json`)
+
+::: {.callout-tip}
+
+Choose the configuration that offloads the least amount to memory while still being able to fit on VRAM for best performance.
+
+Start from Stage 1 -> Stage 2 -> Stage 3.
+
+:::
+
+## Fully Sharded Data Parallel (FSDP) {#sec-fsdp}
+
+FSDP allows you to shard model parameters, gradients, and optimizer states across data parallel workers.
+
+::: {.callout-note}
+
+FSDP2 is recommended for new users. FSDP1 is deprecated and will be removed in an upcoming release of Axolotl.
+
+:::
+
+### FSDP + QLoRA {#sec-fsdp-qlora}
+
+For combining FSDP with QLoRA, see our [dedicated guide](fsdp_qlora.qmd).
+
+### Migrating from FSDP1 to FSDP2 {#sec-migrate-fsdp1-fsdp2}
+
+To migrate your config from FSDP1 to FSDP2, you must use the `fsdp_version` top-level config field to specify the FSDP version, and
+also follow the config field mapping below to update field names.
+
+#### Config mapping
+
+FSDP1 | FSDP2
+-------- | --------
+fsdp_sharding_strategy | reshard_after_forward
+fsdp_backward_prefetch_policy | **REMOVED**
+fsdp_backward_prefetch | **REMOVED**
+fsdp_forward_prefetch | **REMOVED**
+fsdp_sync_module_states | **REMOVED**
+fsdp_cpu_ram_efficient_loading | cpu_ram_efficient_loading
+fsdp_state_dict_type | state_dict_type
+fsdp_use_orig_params | **REMOVED**
+fsdp_activation_checkpointing | activation_checkpointing
+
+For more details, please see the migration guide in the [torchtitan repo](https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md). In Axolotl,
+if you were using the following FSDP1 config:
+
+```{.yaml}
+fsdp_version: 1
+fsdp_config:
+  fsdp_offload_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: Qwen3DecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+```
+
+You can migrate to the following FSDP2 config:
+
+```{.yaml}
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen3DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  reshard_after_forward: true
+```
+
+When using `auto_wrap_policy: SIZE_BASED_WRAP`, set `min_num_params` to the minimum parameter count a module must have to be wrapped (e.g. `100000000`).
+
+### FSDP1 (deprecated) {#sec-fsdp-config}
+
+::: {.callout-note}
+
+Using `fsdp` to configure FSDP is deprecated and will be removed in an upcoming release of Axolotl. Please use `fsdp_config` as above instead.
+
+:::
+
+```{.yaml}
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_offload_params: true
+  fsdp_state_dict_type: FULL_STATE_DICT
+  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.
+
+See our [dedicated guide](sequence_parallelism.qmd) for more information.
+
+## Performance Optimization {#sec-performance}
+
+### Liger Kernel Integration {#sec-liger}
+
+Please see [docs](custom_integrations.qmd#liger) for more info.
+
+## Troubleshooting {#sec-troubleshooting}
+
+### NCCL Issues {#sec-nccl}
+
+For NCCL-related problems, see our [NCCL troubleshooting guide](nccl.qmd).
+
+### Common Problems {#sec-common-problems}
+
+::: {.panel-tabset}
+
+## Memory Issues
+
+- Reduce `micro_batch_size`
+- Reduce `eval_batch_size`
+- Adjust `gradient_accumulation_steps`
+- Consider using a higher ZeRO stage
+
+## Training Instability
+
+- Start with DeepSpeed ZeRO-2
+- Monitor loss values
+- Check learning rates
+
+:::
+
+For more detailed troubleshooting, see our [debugging guide](debugging.qmd).
diff --git a/docs/multi-node.qmd b/docs/multi-node.qmd
index 5c6fa976b9..16196a2d70 100644
--- a/docs/multi-node.qmd
+++ b/docs/multi-node.qmd
@@ -3,6 +3,18 @@ title: Multi Node
 description: How to use Axolotl on multiple machines
 ---
 
+The below are three ways to train multi-node in Axolotl.
+
+::: {.callout-important}
+Each machine needs a copy of Axolotl, we suggest using the same commit to ensure compatibility.
+
+You will also need to have the same configuration file for your model on each machine.
+
+Make sure the main machine is reachable by other machines.
+:::
+
+## Accelerate
+
 You will need to create a configuration for accelerate, either by using `accelerate config` and follow the instructions or you can use one of the preset below:
 
 ~/.cache/huggingface/accelerate/default_config.yaml
@@ -26,23 +38,57 @@ tpu_use_sudo: false
 use_cpu: false
 ```
 
-Configure your model to use FSDP with for example:
+Configure your model to use FSDP in the Axolotl yaml. For example:
 ```yaml
-fsdp:
-  - full_shard
-  - auto_wrap
+fsdp_version: 2
 fsdp_config:
-  fsdp_offload_params: true
-  fsdp_state_dict_type: FULL_STATE_DICT
-  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  offload_params: true
+  state_dict_type: FULL_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  reshard_after_forward: true
 ```
 
-## Machine configuration
+All you have to do now is launch using accelerate as you would usually do on each machine and voila, the processes will start once you have launched accelerate on every machine.
 
-On each machine you need a copy of Axolotl, we suggest using the same commit to ensure compatibility.
+## Raytrain
 
-You will also need to have the same configuration file for your model on each machine.
+Please see ray train doc [here](ray-integration.qmd).
 
-On the main machine only, make sure the port you set as `main_process_port` is open in TCP and reachable by other machines.
+## Torchrun
 
-All you have to do now is launch using accelerate as you would usually do on each machine and voila, the processes will start once you have launched accelerate on every machine.
+If you are using Infiniband, we recommend torchrun to utilize the full bandwidth.
+
+Set the following env (change buffersize/socketname depending on your system):
+
+```bash
+export NCCL_IB_DISABLE=0
+export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
+export NCCL_BUFFSIZE=2097152
+```
+
+Run the following on each node:
+
+### Option 1: New Axolotl CLI with launcher args (Recommended)
+
+```bash
+axolotl train config.yaml --launcher torchrun -- --nnodes $num_nodes --nproc_per_node $gpu_per_node --rdzv_id $rdzv_id --rdzv_backend c10d --rdzv_endpoint "$head_node_ip:$head_node_port"
+```
+
+### Option 2: Direct torchrun (Legacy)
+
+```bash
+torchrun --nnodes $num_nodes --nproc_per_node $gpu_per_node --rdzv_id $rdzv_id --rdzv_backend c10d --rdzv_endpoint "$head_node_ip:$head_node_port" -m axolotl.cli.train config.yaml
+```
+
+Please make sure to substitute the placeholder variables:
+
+- `num_nodes`: Number of nodes (containing GPUs)
+- `gpu_per_node`: Number of gpus per node
+- `head_node_ip`: IP of the head node (make sure other machines can connect to this)
+- `head_node_port`: Port of the head node (make sure other machines can connect to this. Default 29400)
+- `rdzv_id`: A unique job ID that is used by the job across nodes.
+
+The new CLI approach (Option 1) is recommended as it provides consistent argument handling and works seamlessly with other Axolotl CLI features.
+
+More info on the available configs can be found on the Pytorch docs [here](https://pytorch.org/docs/stable/elastic/run.html)
diff --git a/docs/multimodal.qmd b/docs/multimodal.qmd
new file mode 100644
index 0000000000..5197f48b10
--- /dev/null
+++ b/docs/multimodal.qmd
@@ -0,0 +1,371 @@
+---
+title: MultiModal / Vision Language Models (BETA)
+format:
+  html:
+    toc: true
+    toc-depth: 3
+---
+
+## Supported Models
+
+- [Gemma-4](#sec-gemma-4) *(NEW)*
+- [Mllama](#sec-mllama)
+- [Llama4](#sec-llama4)
+- [Pixtral](#sec-pixtral)
+- [Llava-1.5](#sec-llava-15)
+- [Mistral-Small-3.1](#sec-mistral-small-31)
+- [Mistral-Small-4](#sec-mistral-small-4)
+- [Magistral-Small-2509](#sec-magistral-small-2509)
+- [Voxtral](#sec-voxtral)
+- [Gemma-3](#sec-gemma-3)
+- [Gemma-3n](#sec-gemma-3n)
+- [Qwen2-VL](#sec-qwen2-vl)
+- [Qwen2.5-VL](#sec-qwen25-vl)
+- [Qwen3.5](#sec-qwen3-5)
+- [GLM-4.6V](#sec-glm-4-6v)
+- [SmolVLM2](#sec-smolvlm2)
+- [LFM2-VL](#sec-lfm2-vl)
+- [Intern-VL](#sec-intern-vl)
+
+## 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.
+
+```yaml
+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 if specified
+
+# example dataset
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+# (optional) if doing lora, only finetune the Language model,
+# leave the vision model and vision tower frozen
+# load_in_8bit: true
+adapter: lora
+lora_target_modules: 'model.language_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-tip}
+Some of our chat_templates have been extended to support broader dataset types. This should not break any existing configs.
+:::
+
+::: {.callout-note}
+As of now, we do not truncate nor drop samples based on `sequence_len` as each arch has different ways to process non-text tokens. We are looking for help on this.
+:::
+
+### Mllama {#sec-mllama}
+
+```yaml
+base_model: meta-llama/Llama-3.2-11B-Vision-Instruct
+
+chat_template: llama3_2_vision
+```
+
+### Llama4 {#sec-llama4}
+
+```yaml
+base_model: meta-llama/Llama-4-Scout-17B-16E-Instruct
+
+chat_template: llama4
+```
+
+### 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}
+
+::: {.callout-tip}
+Please make sure to install vision lib via `pip install 'mistral-common[opencv]==1.8.5'`
+:::
+
+```yaml
+base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
+```
+
+### Mistral-Small-4 {#sec-mistral-small-4}
+
+```yaml
+base_model: mistralai/Mistral-Small-4-119B-2603
+```
+
+### Magistral-Small-2509 {#sec-magistral-small-2509}
+
+::: {.callout-tip}
+Please make sure to install vision lib via `pip install 'mistral-common[opencv]==1.8.5'`
+:::
+
+```yaml
+base_model: mistralai/Magistral-Small-2509
+```
+
+### Voxtral {#sec-voxtral}
+
+::: {.callout-tip}
+Please make sure to install audio lib via `pip3 install librosa==0.11.0 'mistral_common[audio]==1.8.3'`
+:::
+
+```yaml
+base_model: mistralai/Voxtral-Mini-3B-2507
+
+processor_type: VoxtralProcessor
+```
+
+### Gemma-4 {#sec-gemma-4}
+
+All Gemma 4 variants (E2B, E4B, 26B-A4B, 31B) load as multimodal models even for text-only training.
+
+```yaml
+base_model: google/gemma-4-E2B-it  # or E4B-it, 26B-A4B, 31B
+
+chat_template: gemma4
+freeze_mm_modules: true  # freeze vision/audio encoders for text-only or vision LoRA
+
+# For the 26B-A4B MoE model, enable ScatterMoE and expert LoRA:
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+use_kernels: true
+use_scattermoe: true
+experts_implementation: scattermoe
+
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+# MoE expert LoRA (3D tensors, not nn.Linear) — only for 26B-A4B:
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+```
+
+::: {.callout-warning}
+Gemma 4 VLM training starts with high loss (~8-15). This is expected — see the [training stability guide](training_stability.qmd) for details.
+:::
+
+::: {.callout-tip}
+For DDP training, axolotl auto-detects Gemma4 and sets `use_reentrant=False` and `ddp_find_unused_parameters=True`. However, when `activation_offloading: true`, `ddp_find_unused_parameters` is skipped (checkpoint wrappers conflict with it); use `freeze_mm_modules: true` instead to handle unused vision/audio params. For FSDP2, use `fsdp_transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer`.
+:::
+
+### 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
+```
+
+### Gemma-3n {#sec-gemma-3n}
+
+::: {.callout-warning}
+The model's initial loss and grad norm will be very high. We suspect this to be due to the Conv in the vision layers.
+:::
+
+::: {.callout-tip}
+Please make sure to install `timm` via `pip3 install timm==1.0.17`
+:::
+
+```yaml
+base_model: google/gemma-3n-E2B-it
+
+chat_template: gemma3n
+```
+
+### 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
+```
+
+### Qwen3-VL {#sec-qwen3-vl}
+
+```yaml
+base_model: Qwen/Qwen3-VL-4B-Instruct
+
+chat_template: qwen2_vl  # same as qwen2-vl
+```
+
+### Qwen3.5 {#sec-qwen3-5}
+
+```yaml
+base_model: Qwen/Qwen3.5-9B
+
+chat_template: qwen3_5
+```
+
+### GLM-4.6V {#sec-glm-4-6v}
+
+Both GLM-4.6V (106B MoE) and GLM-4.6V-Flash (9B) are supported.
+
+```yaml
+# GLM-4.6V (106B MoE version)
+base_model: zai-org/GLM-4.6V
+
+# OR GLM-4.6V-Flash (9B version)
+base_model: zai-org/GLM-4.6V-Flash
+```
+
+### SmolVLM2 {#sec-smolvlm2}
+
+::: {.callout-tip}
+Please make sure to install `num2words` via `pip3 install num2words==0.5.14`
+:::
+
+```yaml
+base_model: HuggingFaceTB/SmolVLM2-500M-Video-Instruct
+```
+
+### LFM2-VL {#sec-lfm2-vl}
+
+::: {.callout-warning}
+Please uninstall `causal-conv1d` via `pip3 uninstall -y causal-conv1d`
+:::
+
+```yaml
+base_model: LiquidAI/LFM2-VL-450M
+```
+
+### Intern-VL {#sec-intern-vl}
+
+::: {.callout-tip}
+Please make sure to install `timm` via `pip3 install timm==1.0.19`
+:::
+
+```yaml
+base_model: OpenGVLab/InternVL3_5-8B
+```
+
+## 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`, `image`, `path`, `url`, `base64`, or `audio`).
+
+::: {.callout-note}
+The recommended column name is `messages` with the OpenAI-style `role`/`content` schema shown above. If your upstream data already lives under a different column name (e.g., `conversations`, `dialogue`), set `field_messages: <your_column>` in the dataset config; this is a rename-only escape hatch — the inner schema must still match one of the two formats above (`role`/`content` or ShareGPT's `from`/`value`).
+:::
+
+### Image
+
+::: {.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`.
+:::
+
+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`
+
+### Audio
+
+For audio loading, you can use the following keys within `content` alongside `"type": "audio"`:
+
+- `"path": "/path/to/audio.mp3"`
+- `"url": "https://example.com/audio.mp3"`
+- `"audio": np.ndarray`
+
+::: {.callout-tip}
+
+You may need to install `librosa` via `pip3 install librosa==0.11.0`.
+
+:::
+
+### Video
+
+::: {.callout-warning}
+
+This is not well tested at the moment. We welcome contributors!
+
+:::
+
+For video loading, you can use the following keys within `content` alongside `"type": "video"`:
+
+- `"path": "/path/to/video.mp4"`
+- `"url": "https://example.com/video.mp4"`
+- `"video": np.ndarray | list[PIL.Image.Image] | torch.Tensor` (or list of the aforementioned)
+
+### Example
+
+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", "url": "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."}
+            ]
+        }
+    ]
+  }
+]
+```
+
+## FAQ
+
+1. `PIL.UnidentifiedImageError: cannot identify image file ...`
+
+`PIL` could not retrieve the file at `url` using `requests`. Please check for typo. One alternative reason is that the request is blocked by the server.
diff --git a/docs/multimodal_assistant_mask.md b/docs/multimodal_assistant_mask.md
new file mode 100644
index 0000000000..339ab420f8
--- /dev/null
+++ b/docs/multimodal_assistant_mask.md
@@ -0,0 +1,84 @@
+# Multimodal assistant-only loss masking
+
+## Correct placement
+
+```yaml
+# Top-level: only train_on_inputs lives here.
+train_on_inputs: false
+
+datasets:
+  - path: data/train.jsonl
+    type: chat_template
+    roles_to_train:          # per-dataset — this is what the MM scanner reads
+      - assistant
+    train_on_eos: turn       # per-dataset — same
+
+test_datasets:
+  - path: data/val.jsonl
+    type: chat_template
+    split: train
+    roles_to_train:
+      - assistant
+    train_on_eos: turn
+```
+
+## How to verify at runtime
+
+`build_collator` logs the resolved knobs at INFO:
+
+```text
+MM collator: train_on_inputs=False roles_to_train=['assistant'] train_on_eos=turn role_boundaries_override=none
+```
+
+If `roles_to_train` logs as `None`, the YAML knobs are not reaching the
+scanner — check that they are under `datasets[0]`, not at the root.
+
+Each verified strategy additionally logs its resolved boundary token ids at
+strategy init (e.g. `<|turn>model` → `[105, 4368]`, `<turn|>` → `[106]` for
+Gemma 4). If a strategy emits the "has no built-in role boundaries ... only
+pad and media tokens are masked" one-shot warning instead, it is on the
+fallback path — declare per-role markers in YAML via `cfg.role_boundaries`
+(below) to activate masking. The strategies currently on this path are
+listed in the audit table above under `fallback + warn`.
+
+## Config-based override: `cfg.role_boundaries`
+
+For the "unverified" strategies above, or for custom chat templates that
+don't match a built-in strategy's markers, users can declare role boundaries
+directly in YAML without subclassing:
+
+```yaml
+role_boundaries:
+  - role: assistant
+    start: "<|turn>model"
+    end: "<turn|>"
+  - role: user
+    start: "<|turn>user"
+    end: "<turn|>"
+  # Optional keys:
+  # include_start: false   # default False
+  # include_end: true      # default True, respects cfg.train_on_eos
+  # end: eos_token         # sentinel: resolves to tokenizer.eos_token_id
+  # end: null              # span runs to end of sequence
+```
+
+Semantics:
+
+- `start` and `end` are literal strings; axolotl encodes them at strategy
+  init via `tokenizer.encode(..., add_special_tokens=False)` and logs the
+  resolved token-id sequences at INFO level.
+- The special value `end: eos_token` is the portable way to express
+  "Pixtral-style assistant turns end at EOS" without hard-coding an id.
+- `role_boundaries` is an **opt-in override**. A non-empty list **replaces**
+  the strategy's built-in declarations wholesale (partial overlays are
+  intentionally unsupported — they're hard to reason about at review time).
+  Leaving the field unset *or* setting it to an empty list (`[]`) both mean
+  "use the strategy's built-ins." Writing `role_boundaries: []` is almost
+  always a typo or leftover — honoring it literally would produce all-masked
+  labels and zero gradient, so it is treated the same as unset.
+- `cfg.roles_to_train` still governs which declared roles contribute to
+  loss. You can declare `user` and `assistant` boundaries and set
+  `roles_to_train: ["assistant"]` to have the scanner correctly identify
+  user spans as masking boundaries without training on their content.
+- Invalid specs fail loudly at strategy init (missing `role`/`start`,
+  unencodable markers), not silently at loss-compute time.
diff --git a/docs/nccl.qmd b/docs/nccl.qmd
index 3b616aa665..bd2a744124 100644
--- a/docs/nccl.qmd
+++ b/docs/nccl.qmd
@@ -13,13 +13,13 @@ Often, this timeout will happen after 30 minutes (the default setting) and is ac
 
 Forcing cross-GPU communication via [NVLink](https://en.wikipedia.org/wiki/NVLink) may help without increasing timeouts. To verify that your configuration is leveraging NVLink run the following command:
 
-```shell
+```bash
 nvidia-smi nvlink --status
 ```
 
 To force NCCL to use NVLink, simply set this in the environment:
 
-```shell
+```bash
 export NCCL_P2P_LEVEL=NVL
 ```
 
@@ -33,13 +33,13 @@ If NVLink is not available in your environment there are other options for ``NCC
 
 To validate that acceptable data transfer speeds exist for your training job, running [NCCL Tests](https://github.com/NVIDIA/nccl-tests/blob/master/README.md) can help pinpoint bottlenecks, for example:
 
-```shell
+```bash
 ./build/all_reduce_perf -b 8 -e 128M -f 2 -g 3
 ```
 
 It can be useful when debugging NCCL communication timeouts to activate additional logging in both PyTorch and NCCL:
 
-```shell
+```bash
 export NCCL_DEBUG=INFO
 export NCCL_DEBUG_SUBSYS=ALL
 export TORCH_DISTRIBUTED_DEBUG=INFO
diff --git a/docs/nd_parallelism.qmd b/docs/nd_parallelism.qmd
new file mode 100644
index 0000000000..729eeb9bea
--- /dev/null
+++ b/docs/nd_parallelism.qmd
@@ -0,0 +1,129 @@
+---
+title: "N-D Parallelism (Beta)"
+---
+
+Axolotl enables training models at scale by composing different parallelism techniques. This is essential when:
+
+- A model's weights are too large to fit on a single GPU's memory.
+- A model's activations, especially with very long contexts, are too large for a single GPU.
+- You want to accelerate training by using multiple GPUs or nodes.
+
+or combinations of the above!
+
+## Core Concepts
+
+Parallelism strategies can be combined. The key is understanding how each one divides the workload. PyTorch's `DeviceMesh` is the modern way to manage these combinations, creating a logical grid of your GPUs and assigning different parallel strategies to different dimensions of the grid.
+
+### Data Parallelism {#sec-dp}
+
+Data Parallelism focuses on splitting the global data batch across GPUs.
+
+- Distributed Data Parallel (DDP): The classic approach. The full model is replicated on every GPU. Each GPU processes a different slice of the data batch. Gradients are then averaged across all GPUs after the backward pass to keep the models synchronized. This can substantially improve data throughput compared to single-device training, but requires that each GPU is able to hold the entire model, its gradients, and optimizer states.
+
+- [Fully Sharded Data Parallel (FSDP)](multi-gpu.qmd#fully-sharded-data-parallel-(fsdp)): A highly memory-efficient form of data parallelism (inspired by DeepSpeed's ZeRO). Instead of replicating the model, FSDP shards the model's *parameters, gradients, and optimizer states* across the GPUs in the data-parallel group. During computation, each GPU receives the specific parameters it needs via an `all_gather` operation just before they are used, and they can be discarded immediately after (`reshard-after-forward`).
+    - FSDP maps to ZeRO stages:
+        - ZeRO-2 (`reshard_after_forward=False`): Shards gradients and optimizer states. Model weights are replicated on each GPU.
+        - ZeRO-3 (`reshard_after_forward=True`): Shards gradients, optimizer states, AND model parameters. This provides the most memory savings at the cost of more communication (re-gathering parameters for both forward and backward passes).
+
+### [Experimental] Tensor Parallelism (TP) {#sec-tp}
+
+Also known as "horizontal model parallelism," as described in the [Megatron-LM paper](https://arxiv.org/pdf/1909.08053.pdf). Instead of splitting the batch, TP splits the model's layers themselves across GPUs.
+
+- How it works: For a linear layer `Y = XA`, the weight matrix `A` is split column-wise (`A = [A_1, A_2]`). The computation becomes `Y_1 = XA_1` and `Y_2 = XA_2`, which can happen in parallel on different GPUs. The final output `Y` is simply the concatenation of `Y_1` and `Y_2`. Check [this comment](https://github.com/huggingface/transformers/issues/10321#issuecomment-783543530) for more detailed info.
+- Requirement: TP involves frequent, small communications within a forward/backward pass. It requires a very fast interconnect between GPUs (e.g., NVLink) and is typically not recommended across different nodes.
+
+### Context Parallelism (CP) {#sec-cp}
+
+Context Parallelism, also called [Sequence Parallelism](sequence_parallelism.qmd), addresses the memory bottleneck from long sequences. The input sequence itself is split along the sequence length dimension and distributed across GPUs.
+
+- How it works: If you have a sequence of 8192 tokens and a `context_parallel_size` of 4, each GPU will only handle a chunk of 2048 tokens.
+- The Challenge: Attention is not local; every token needs to "attend to" every other token. Splitting the sequence breaks this.
+- The Solution (`ring-flash-attention`): An efficient communication protocol is used. To compute attention for its local sequence chunk, each GPU passes its Key-Value (KV) cache to its neighbor in a "ring." After `N-1` steps, every GPU has seen the KV-cache from all other GPUs, allowing it to compute the correct attention values for its chunk. This is implemented using the highly optimized `flash-attention` kernel at each step.
+- **Mamba/SSM Hybrid Models**: For hybrid architectures (Nemotron-H, Falcon-H1, Granite MoE Hybrid), attention layers use ring attention as above, while SSM (Mamba2) layers use P2P state passing across ranks with an additive output correction. See [Sequence Parallelism](sequence_parallelism.qmd) for details.
+
+### Expert Parallelism (EP) {#sec-ep}
+
+Expert Parallelism shards the experts of a Mixture-of-Experts (MoE) model across GPUs. Each rank holds `num_experts / ep_size` experts; tokens are routed to the rank that owns the chosen expert via DeepEP's fused dispatch/combine kernels.
+
+- How it works: the EP plugin slices each `Experts` module's `gate_up_proj` / `down_proj` along the experts dimension and replaces the dispatch/combine path with DeepEP. Each rank sees a different batch (so EP is data-parallel for *data*) but holds different expert shards (model-parallel for *expert weights*).
+- When to use: MoE training where the experts dominate the parameter count and don't fit under FSDP alone, or where you want to compose with FSDP on orthogonal mesh axes.
+- Composes with FSDP: `expert_parallel_size × dp_shard_size = world_size`. EP shards experts on the `ep` axis; FSDP shards non-experts (and the per-EP-slice expert weights) on the `dp_shard` axis. Process groups are disjoint.
+- Requires the [DeepEP](https://github.com/deepseek-ai/DeepEP) library (Ampere or Hopper, NVLink). See [Expert Parallelism Integration](custom_integrations.qmd#expert-parallelism-integration) for install + integration details.
+
+### Hybrid Sharding Data Parallel (HSDP) {#sec-hsdp}
+
+HSDP is a 2D strategy that intelligently combines FSDP and DDP, typically for multi-node training.
+
+- Intra-Node (within a machine): Use FSDP. This is efficient because GPUs on the same node have fast interconnects (NVLink), making the `all_gather` operations for sharded parameters fast.
+- Inter-Node (across machines): Use DDP. The gradient synchronization between nodes is less frequent than FSDP's parameter gathering, making it a better fit for the slower node-to-node network (e.g., Ethernet/Infiniband).
+- Example: With 2 nodes of 8 GPUs each (16 total), you could have `dp_shard_size=8` (FSDP within each node) and `dp_replicate_size=2` (DDP across the two nodes).
+
+## Usage
+
+```yaml
+# FSDP config. See https://docs.axolotl.ai/docs/multi-gpu.html#sec-fsdp
+fsdp_version: 2
+fsdp_config:
+  # ...
+
+# The number of GPUs to shard the model parameters across (FSDP dimension).
+dp_shard_size: 4
+
+# The number of times to replicate the sharded model (DDP dimension).
+dp_replicate_size: 2
+
+# Number of GPUs for Tensor Parallelism.
+tensor_parallel_size: 1  # (default is 1, no TP)
+
+# Number of GPUs for Context/Sequence Parallelism.
+context_parallel_size: 1 # (default is 1, no CP)
+
+# Number of GPUs to shard MoE experts across (Expert Parallel).
+expert_parallel_size: 1  # (default is 1, no EP)
+```
+
+Note: We recommend FSDP. DeepSpeed is only compatible with `tensor_parallel_size`.
+
+## Examples
+
+::: {.callout-tip}
+See our example configs [here](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/distributed-parallel).
+:::
+
+1.  HSDP on 2 nodes with 4 GPUs each (8 GPUs total):
+    - You want FSDP within each node and DDP across nodes.
+    - Set `dp_shard_size: 4` and `dp_replicate_size: 2`.
+
+2.  FSDP + TP on a single 8-GPU node:
+    - You want to split the model across 4 GPUs using FSDP, and further split each layer across 2 GPUs with TP.
+    - Set `dp_shard_size: 4` and `tensor_parallel_size: 2`.
+
+3.  FSDP + CP on a single 8-GPU node for long context:
+    - You want to shard the model across all 8 GPUs and also split the sequence length across all 8 GPUs.
+    - Set `dp_shard_size: 8` and `context_parallel_size: 8`. Note: this means the data parallel group and context parallel group are the same. A more common setup might be to shard across a smaller group.
+
+4.  FSDP + EP on a 4-GPU MoE training run:
+    - You want EP to shard the experts and FSDP to shard non-expert params on orthogonal mesh axes.
+    - Set `expert_parallel_size: 2` and `dp_shard_size: 2` (`ep × dp_shard == world_size`). Add the `ExpertParallelPlugin` to `plugins:`.
+
+## Support Matrix
+
+This matrix describes how different parallelism methods can be combined in Axolotl.
+
+| Combination | `dp_replicate_size` | `dp_shard_size` | `tp_size` | `cp_size` | `ep_size` | Status & Notes |
+| --- | :---: | :---: |:---:|:---:|:---:|---|
+| **FSDP** (ZeRO-3) | 1 | >1 | 1 | 1 | 1 | ✅ Fully supported. Shards model across all GPUs. |
+| **HSDP** | >1 | >1 | 1 | 1 | 1 | ✅ Fully supported. FSDP intra-node, DDP inter-node. |
+| **FSDP + TP** | 1 | >1 | >1 | 1 | 1 | ✅ **2D Parallelism**. Shards the model across a `dp_shard` group, and TP-splits layers within the `tp` group. |
+| **HSDP + TP** | >1 | >1 | >1 | 1 | 1 | ✅ **3D Parallelism**. A powerful but complex combination. |
+| **FSDP + CP** | 1 | >1 | 1 | >1 | 1 | ✅ **2D Parallelism**. Combines FSDP with context parallelism. |
+| **FSDP + TP + CP**| 1 | >1 | >1| >1| 1 | ✅ **3D Parallelism**. Another advanced combination. |
+| **EP** | 1 | 1 | 1 | 1 | >1 | ✅ Supported (MoE only). Shards experts across all ranks; non-experts replicated. |
+| **FSDP + EP** | 1 | >1 | 1 | 1 | >1 | ✅ **2D Parallelism** (MoE). EP shards experts on the `ep` axis; FSDP shards non-experts on `dp_shard`. |
+| EP + TP/CP | 1 | * | >1 | >1 | >1 | ❌ Not yet supported in v1; raises `NotImplementedError`. |
+| DDP + TP/CP | >1 | 1 | >1 | >1 | 1 | ❌ **Not Supported**. The `ParallelismConfig` explicitly prevents this, as composing pure DDP with TP or CP is currently not supported. You should use FSDP + TP/CP instead (`dp_shard_size > 1`). |
+| Just TP / CP | 1 | 1 | >1 | >1 | 1 | ✅ Supported. Useful for inference or when the model fits on one GPU but context is too long. |
+
+- `tp_size` refers to `tensor_parallel_size`
+- `cp_size` refers to `context_parallel_size`
+- `ep_size` refers to `expert_parallel_size`
diff --git a/docs/optimizations.qmd b/docs/optimizations.qmd
new file mode 100644
index 0000000000..6776485d95
--- /dev/null
+++ b/docs/optimizations.qmd
@@ -0,0 +1,177 @@
+---
+title: Optimizations Guide
+description: A guide to the performance and memory optimizations available in Axolotl.
+---
+
+Axolotl includes numerous optimizations to speed up training, reduce memory usage, and handle large models.
+
+This guide provides a high-level overview and directs you to the detailed documentation for each feature.
+
+## Speed Optimizations
+
+These optimizations focus on increasing training throughput and reducing total training time.
+
+### Sample Packing
+
+Improves GPU utilization by combining multiple short sequences into a single packed sequence for training. This requires enabling one of the [attention](#attention-implementations) implementations below.
+
+- **Config:** `sample_packing: true`
+- **Learn more:** [Sample Packing](multipack.qmd)
+
+### Attention Implementations
+
+Using an optimized attention implementation is critical for training speed.
+
+- **[Flash Attention 2](https://github.com/Dao-AILab/flash-attention)**: `attn_implementation: flash_attention_2`. **(Recommended)** The industry standard for fast attention on modern GPUs. Requires Ampere or higher. For AMD, check [AMD Support](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#amd-rocm-support).
+- **[Flex Attention](https://pytorch.org/blog/flexattention/)**: `attn_implementation: flex_attention`.
+- **[SDP Attention](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)**: `attn_implementation: sdpa`. PyTorch's native implementation.
+- **[Xformers](https://github.com/facebookresearch/xformers)**: `attn_implementation: xformers`. Works with FP16.
+
+See [Attention](attention.qmd) for the full list of backends and the canonical values.
+
+### LoRA Optimizations
+
+Leverages optimized kernels to accelerate LoRA training and reduce memory usage.
+
+- **Learn more:** [LoRA Optimizations Documentation](lora_optims.qmd)
+
+## Memory Optimizations
+
+These techniques help you fit larger models or use bigger batch sizes on your existing hardware.
+
+### Parameter Efficient Finetuning (LoRA & QLoRA)
+
+Drastically reduces memory by training a small set of "adapter" parameters instead of the full model. This is the most common and effective memory-saving technique.
+
+- Examples: Find configs with `lora` or `qlora` in the [examples directory](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/llama-3).
+- Config Reference: See `adapter`, `load_in_4bit`, and `load_in_8bit` in the [Configuration Reference](config-reference.qmd).
+
+### Gradient Checkpointing & Activation Offloading
+
+These techniques save VRAM by changing how activations are handled.
+
+- Gradient Checkpointing: re-computes activations during the backward pass, trading compute time for VRAM.
+- Activation Offloading: moves activations to CPU RAM or disk, trading I/O overhead for VRAM.
+- Learn more: [Gradient Checkpointing and Offloading Docs](gradient_checkpointing.qmd)
+
+### Layer Offloading
+
+Offloads frozen (non-trainable) decoder layer parameters to CPU and streams them back to GPU one layer at a time during forward/backward passes using CUDA stream prefetching. Especially effective for LoRA/QLoRA where most parameters are frozen.
+
+- **Config:** `layer_offloading: true`
+- **Learn more:** [Layer Offloading Docs](gradient_checkpointing.qmd#enabling-layer-offloading)
+
+### Cut Cross Entropy (CCE)
+
+Reduces VRAM usage by using an optimized cross-entropy loss calculation.
+
+- **Learn more:** [Custom Integrations - CCE](custom_integrations.qmd#cut-cross-entropy)
+
+### Liger Kernels
+
+Provides efficient Triton kernels to improve training speed and reduce memory usage.
+
+- **Learn more:** [Custom Integrations - Liger Kernels](custom_integrations.qmd#liger-kernels)
+
+### Fused RMSNorm + RoPE (Qwen3 / Qwen3-MoE / Qwen3.5 / Qwen3.5-MoE / Qwen3.6 dense / Qwen3.6-MoE)
+
+Replaces the per-layer `q_norm + apply_rotary_pos_emb` (and matching K path) with a single Triton kernel launch on the full-attention layers. Opt-in. The kernel computes in fp32 and rounds once, so it matches an fp32 reference to within bf16 rounding — i.e. it is *more* accurate than the eager bf16 path, which rounds at several intermediate steps. Gemma 4 always uses the fused path (no flag needed). Qwen3.6 checkpoints are loaded by transformers under the `qwen3_5` / `qwen3_5_moe` model_types, so the same flag covers both generations.
+
+```yaml
+fused_attn_kernel: true
+```
+
+- **Compile-safe:** the kernel is wrapped as a `torch.library.triton_op` and traces under `torch.compile(fullgraph=True)`.
+- **Hardware note:** on sm_120 (Blackwell) combining with `torch_compile: true` is a net win; on sm_86 (Ampere consumer) `torch_compile: true` currently regresses the surrounding Inductor-generated kernels — keep compile off there.
+
+### Expert Kernels
+
+Optimized per-expert grouped-GEMM kernels for MoE training, with LoRA support.
+
+- **ScatterMoE**: Triton, any CUDA GPU.
+- **SonicMoE**: CUTLASS / cute-DSL, Hopper (H100/H200) or Blackwell (B200/GB200).
+
+- **Config:** `use_scattermoe: true` or `use_sonicmoe: true`
+- **Learn more:** [Custom Integrations - Kernels Integration](custom_integrations.qmd#kernels-integration)
+
+## Long Context Models
+
+Techniques to train models on sequences longer than their original context window.
+
+### RoPE Scaling
+
+Extends a model's context window by interpolating its Rotary Position Embeddings.
+
+- **Config:** Pass the `rope_scaling` config under the `overrides_of_model_config: `. To learn how to set RoPE, check the respective model config.
+
+### Sequence Parallelism
+
+Splits long sequences across multiple GPUs, enabling training with sequence lengths that would not fit on a single device.
+
+- **Learn more:** [Sequence Parallelism Documentation](sequence_parallelism.qmd)
+
+### Artic Long Sequence Training (ALST)
+
+ALST is a recipe that combines several techniques to train long-context models efficiently. It typically involves:
+
+- TiledMLP to reduce memory usage in MLP layers.
+- Tiled Loss functions (like [CCE](#cut-cross-entropy-(cce) or [Liger](#liger-kernels)).
+- Activation Offloading to CPU.
+
+- Example: [ALST Example Configuration](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/alst)
+
+## Large Models (Distributed Training)
+
+To train models that don't fit on a single GPU, you'll need to use a distributed training strategy like FSDP or DeepSpeed. These frameworks shard the model weights, gradients, and optimizer states across multiple GPUs and nodes.
+
+- **Learn more:** [Multi-GPU Guide](multi-gpu.qmd)
+- **Learn more:** [Multi-Node Guide](multi-node.qmd)
+
+### N-D Parallelism (Beta)
+
+For advanced scaling, Axolotl allows you to compose different parallelism techniques (e.g., Data, Tensor, Sequence, Expert Parallelism). This is a powerful approach to train an extremely large model by overcoming multiple bottlenecks at once.
+
+- **Learn more:** [N-D Parallelism Guide](nd_parallelism.qmd)
+
+
+## Quantization
+
+Techniques to reduce the precision of model weights for memory savings.
+
+### 4-bit Training (QLoRA)
+
+The recommended approach for quantization-based training. It loads the base model in 4-bit using `bitsandbytes` and then trains QLoRA adapters. See [Adapter Finetuning](#adapter-finetuning-lora-qlora) for details.
+
+### FP8 Training
+
+Enables training with 8-bit floating point precision on supported hardware (e.g., NVIDIA Hopper series GPUs) for significant speed and memory gains.
+
+- **Example:** [Llama 3 FP8 FSDP Example](https://github.com/axolotl-ai-cloud/axolotl/blob/main/examples/llama-3/3b-fp8-fsdp2.yaml)
+
+### NVFP4 (W4A4) LoRA
+
+Train LoRA adapters on a ModelOpt NVFP4 MoE checkpoint (the experts stay 4-bit-packed; LoRA `A` / `B` train in bf16).
+Two kernels support it: **SonicMoE** (`use_sonicmoe: true`, W4A4 native on Blackwell SM100+ or W4A16 elsewhere) and **ScatterMoE** (`use_scattermoe: true`, W4A16 on any CUDA GPU).
+
+- **Config:** `use_sonicmoe: true` or `use_scattermoe: true` with a ModelOpt NVFP4 `base_model` (e.g. `nvidia/Qwen3-30B-A3B-NVFP4`)
+- **Examples:** [Qwen3-30B-A3B (SonicMoE)](https://github.com/axolotl-ai-cloud/axolotl/blob/main/examples/qwen3/30b-a3b-nvfp4-lora.yaml), [GLM-5.2 (ScatterMoE)](https://github.com/axolotl-ai-cloud/axolotl/blob/main/examples/glm_moe_dsa/glm-5.2-nvfp4-lora.yaml)
+- **Learn more:** [ScatterMoE NVFP4](custom_integrations.qmd#scattermoe-nvfp4-w4a16-lora) / [SonicMoE NVFP4](custom_integrations.qmd#sonicmoe-nvfp4-w4a4-lora)
+
+### Quantization Aware Training (QAT)
+
+Simulates quantization effects during training, helping the model adapt and potentially improving the final accuracy of the quantized model.
+
+- **Learn more:** [QAT Documentation](qat.qmd)
+
+### GPTQ
+
+Allows you to finetune LoRA adapters on top of a model that has already been quantized using the GPTQ method.
+
+- **Example:** [GPTQ LoRA Example](https://github.com/axolotl-ai-cloud/axolotl/blob/main/examples/llama-2/gptq-lora.yml)
+
+### MoE Expert Quantization
+
+Quantizes MoE expert weights on load to reduce VRAM when training MoE models with adapters. Required for Transformers v5+ MoE models where experts use fused `nn.Parameter` tensors.
+
+- **Config:** `quantize_moe_experts: true`
+- **Learn more:** [MoE Expert Quantization](expert_quantization.qmd)
diff --git a/docs/optimizers.qmd b/docs/optimizers.qmd
new file mode 100644
index 0000000000..97348c161a
--- /dev/null
+++ b/docs/optimizers.qmd
@@ -0,0 +1,191 @@
+---
+title: Optimizers
+description: Configuring optimizers
+---
+
+## Overview
+
+Axolotl supports all optimizers supported by [transformers OptimizerNames](https://github.com/huggingface/transformers/blob/51f94ea06d19a6308c61bbb4dc97c40aabd12bad/src/transformers/training_args.py#L142-L187)
+
+Here is a list of optimizers supported by transformers as of `v4.54.0`:
+
+- `adamw_torch`
+- `adamw_torch_fused`
+- `adamw_torch_xla`
+- `adamw_torch_npu_fused`
+- `adamw_apex_fused`
+- `adafactor`
+- `adamw_anyprecision`
+- `adamw_torch_4bit`
+- `adamw_torch_8bit`
+- `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_radam`
+- `schedule_free_adamw`
+- `schedule_free_sgd`
+- `apollo_adamw`
+- `apollo_adamw_layerwise`
+- `stable_adamw`
+
+
+## Custom Optimizers
+
+Enable custom optimizers by passing a string to the `optimizer` argument. Each optimizer will receive beta and epsilon args, however, some may accept additional args which are detailed below.
+
+### optimi_adamw
+
+```yaml
+optimizer: optimi_adamw
+```
+
+### ao_adamw_4bit
+
+Deprecated: Please use `adamw_torch_4bit`.
+
+### ao_adamw_8bit
+
+Deprecated: Please use `adamw_torch_8bit`.
+
+### ao_adamw_fp8
+
+
+```yaml
+optimizer: ao_adamw_fp8
+```
+
+### adopt_adamw
+
+GitHub: [https://github.com/iShohei220/adopt](https://github.com/iShohei220/adopt)
+Paper: [https://arxiv.org/abs/2411.02853](https://arxiv.org/abs/2411.02853)
+
+```yaml
+optimizer: adopt_adamw
+```
+
+### came_pytorch
+
+GitHub: [https://github.com/yangluo7/CAME/tree/master](https://github.com/yangluo7/CAME/tree/master)
+Paper: [https://arxiv.org/abs/2307.02047](https://arxiv.org/abs/2307.02047)
+
+```yaml
+optimizer: came_pytorch
+
+# optional args (defaults below)
+adam_beta1: 0.9
+adam_beta2: 0.999
+adam_beta3: 0.9999
+adam_epsilon: 1e-30
+adam_epsilon2: 1e-16
+```
+
+### muon
+
+Blog: [https://kellerjordan.github.io/posts/muon/](https://kellerjordan.github.io/posts/muon/)
+Paper: [https://arxiv.org/abs/2502.16982v1](https://arxiv.org/abs/2502.16982v1)
+
+```yaml
+optimizer: muon
+```
+
+### dion
+
+Microsoft's Dion (DIstributed OrthoNormalization) optimizer is a scalable and communication-efficient
+orthonormalizing optimizer that uses low-rank approximations to reduce gradient communication.
+
+GitHub: [https://github.com/microsoft/dion](https://github.com/microsoft/dion)
+Paper: [https://arxiv.org/pdf/2504.05295](https://arxiv.org/pdf/2504.05295)
+Note: Implementation written for PyTorch 2.7+ for DTensor
+
+```yaml
+optimizer: dion
+dion_lr: 0.01
+dion_momentum: 0.95
+lr: 0.00001  # learning rate for embeddings and parameters that fallback to AdamW
+```
+
+### sinkgd
+
+SinkGD (Gradient Multi-Normalization) is a **stateless** optimizer: 2D linear weight
+matrices are updated via the SR-Sinkhorn procedure (alternating row/column L2
+normalization of the raw gradient, no momentum or variance state), while embeddings,
+the LM head, and 1D params (norms/biases) fall back to AdamW. The AdamW fallback uses
+torchao's 8-bit optimizer base, so optimizer-state memory is very small (~87% less than
+8-bit AdamW on an 8B full finetune, since the ~87% of params that are 2D linear carry
+zero optimizer state).
+
+Requires PyTorch >= 2.5.1 (relies on torchao's low-bit optimizer base and `torch.compile`).
+
+Paper: [https://arxiv.org/abs/2502.06742](https://arxiv.org/abs/2502.06742)
+
+```yaml
+optimizer: sinkgd
+learning_rate: 0.001
+optim_args:
+  sinkhorn_iters: 5      # number of SR-Sinkhorn row/column normalization iterations
+  sinkgd_lr_scale: 0.05  # α scale applied to linear-layer updates
+```
+
+### q_galore_adamw8bit
+
+Q-GaLore extends [GaLore](https://arxiv.org/abs/2403.03507) with two extra ideas:
+an INT4-quantized projection matrix and an adaptive SVD scheduler that skips
+re-projection when a layer's gradient subspace stabilizes. Both are wired up in
+axolotl. The third Q-GaLore trick — INT8 weight wrapping — is not yet
+implemented and is tracked as a follow-up.
+
+GitHub: [https://github.com/VITA-Group/Q-GaLore](https://github.com/VITA-Group/Q-GaLore)
+Paper: [https://arxiv.org/abs/2407.08296](https://arxiv.org/abs/2407.08296)
+
+Install: `pip install axolotl[qgalore]`
+
+This optimizer is for **full fine-tuning**. It is incompatible with `adapter`
+(LoRA/QLoRA), `load_in_8bit`, and `load_in_4bit`. DeepSpeed is currently gated
+off; FSDP requires `fsdp_version: 2` with `use_orig_params: true`.
+
+```yaml
+optimizer: q_galore_adamw8bit
+bf16: true
+
+# which parameter substrings get the low-rank projection
+# (defaults to ["attn", "mlp"] if unset — matches the reference impl)
+optim_target_modules:
+  - attn
+  - mlp
+
+# Q-GaLore hyperparameters (defaults shown)
+qgalore_rank: 256
+qgalore_update_proj_gap: 200       # max steps between SVD refreshes
+qgalore_scale: 0.25
+qgalore_proj_type: std
+qgalore_proj_quant: true           # INT-quantize the projection matrix P
+qgalore_proj_bits: 4               # bitwidth for P
+qgalore_proj_group_size: 256       # must divide P's last dim evenly
+qgalore_cos_threshold: 0.4         # skip SVD if P_t is this similar to P_{t-1}
+qgalore_gamma_proj: 2              # grow update_proj_gap by this factor when stable
+qgalore_queue_size: 5
+```
diff --git a/docs/qat.qmd b/docs/qat.qmd
new file mode 100644
index 0000000000..91fe5180cb
--- /dev/null
+++ b/docs/qat.qmd
@@ -0,0 +1,40 @@
+---
+title: "Quantization Aware Training (QAT)"
+back-to-top-navigation: true
+toc: true
+toc-expand: 2
+toc-depth: 4
+---
+
+## Overview
+
+[Quantization Aware Training](https://pytorch.org/blog/introduction-to-quantization-on-pytorch/#quantization-aware-training) (QAT) is a technique for improving the accuracy of models which are quantized
+by applying "fake" quantizations to the model's weights (and optionally, activations) during training. This fake
+quantization allows for the model to adjust for noise introduced by the quantization, so when the model is eventually
+quantized, the accuracy loss is minimized. We use the quantization techniques implemented in [torchao](https://github.com/pytorch/ao) to provide
+support for QAT and post-training quantization (PTQ) in axolotl.
+
+We recommend reviewing the excellent QAT tutorial in the [torchtune library](https://pytorch.org/torchtune/main/tutorials/qat_finetune.html#quantizing-the-qat-model),
+and the QAT documentation in the [torchao library](https://github.com/pytorch/ao/tree/main/torchao/quantization/qat), for more details.
+
+## Configuring QAT in Axolotl
+
+To enable QAT in axolotl, add the following to your configuration file:
+
+```yaml
+qat:
+  activation_dtype: # Optional[str] = "int8". Fake quantization layout to use for activation quantization. Valid options are "int4", "int8", "float8"
+  weight_dtype: # Optional[str] = "int8". Fake quantization layout to use for weight quantization. Valid options are "int4", "fp8", and "nvfp4".
+  group_size: # Optional[int] = 32. The number of elements in each group for per-group fake quantization
+  fake_quant_after_n_steps: # Optional[int] = None. The number of steps to apply fake quantization after
+```
+
+We support the following quantization schemas:
+
+- `Int4WeightOnly` (requires the `fbgemm-gpu` extra when installing Axolotl)
+- `Int8DynamicActivationInt4Weight`
+- `Float8DynamicActivationFloat8Weight`
+- `Float8DynamicActivationInt4Weight`
+- `NVFP4`
+
+Once you have finished training, you must quantize your model by using the same quantization configuration which you used to train the model with. You can use the [`quantize`](./quantize.qmd) command to do this.
diff --git a/docs/quantize.qmd b/docs/quantize.qmd
new file mode 100644
index 0000000000..9c3de1ef1f
--- /dev/null
+++ b/docs/quantize.qmd
@@ -0,0 +1,60 @@
+---
+title: "Quantization with torchao"
+back-to-top-navigation: true
+toc: true
+toc-expand: 2
+toc-depth: 4
+---
+
+Quantization is a technique to lower the memory footprint of your model, potentially at the cost of accuracy or model performance. We support quantizing your model using the [torchao](https://github.com/pytorch/ao) library. Quantization is supported for both post-training quantization (PTQ) and quantization-aware training (QAT).
+
+
+::: {.callout-note}
+
+We do not currently support quantization techniques such as GGUF/GPTQ,EXL2 at the moment.
+
+:::
+
+## Configuring Quantization in Axolotl
+
+Quantization is configured using the `quantization` key in your configuration file.
+
+```yaml
+base_model: # The path to the model to quantize.
+quantization:
+  activation_dtype: # Optional[str] = "int8". Fake quantization layout to use for activation quantization. Valid options are "int4", "int8", "float8"
+  weight_dtype: # Optional[str] = "int8". Fake quantization layout to use for weight quantization. Valid options are "int4", "fp8", and "nvfp4".
+  group_size: # Optional[int] = 32. The number of elements in each group for per-group fake quantization
+  quantize_embedding: # Optional[bool] = False. Whether to quantize the embedding layer.
+
+output_dir:  # The path to the output directory.
+```
+
+Once quantization is complete, your quantized model will be saved in the `{output_dir}/quantized` directory.
+
+You may also use the `quantize` command to quantize a model which has been trained with [QAT](./qat.qmd) - you can do this by using the existing QAT configuration file which
+you used to train the model:
+
+```yaml
+# qat.yml
+qat:
+  activation_dtype: int8
+  weight_dtype: int4
+  group_size: 256
+
+output_dir: # The path to the output directory used during training where the final checkpoint has been saved.
+```
+
+```bash
+axolotl quantize qat.yml
+```
+
+This ensures that an identical quantization configuration is used to quantize the model as was used to train it.
+
+
+::: {.callout-note}
+
+If you have configured pushing to hub with `hub_model_id`, your model hub name will have the quantization schema appended to it,
+e.g. `axolotl-ai-cloud/qat-nvfp4-llama3B` will become `axolotl-ai-cloud/qat-nvfp4-llama3B-nvfp4w`
+
+:::
diff --git a/docs/ray-integration.qmd b/docs/ray-integration.qmd
new file mode 100644
index 0000000000..edf9e2dafd
--- /dev/null
+++ b/docs/ray-integration.qmd
@@ -0,0 +1,91 @@
+---
+title: Ray Train
+description: How to use Axolotl with Ray Train
+---
+
+Axolotl supports using Ray as an alternative to `accelerate` for orchestrating training. This is especially useful for multi-node training since you only have to setup code and dependencies in a single node and launch training as if you were using a single node.
+
+With the `--use-ray` CLI flag, Axolotl will use Ray Train's [`TorchTrainer`](https://docs.ray.io/en/latest/train/api/doc/ray.train.torch.TorchTrainer.html#ray.train.torch.TorchTrainer) to run training.
+
+## Ray cluster setup
+
+A prerequisite using the Ray Train integration is to setup a Ray cluster on your desired node(s). For a detailed guide on how you can get started with ray clusters, check the official Ray docs [here](https://docs.ray.io/en/latest/cluster/getting-started.html).
+
+Every Ray cluster has one _head_ node and a set of worker nodes. The head node is just like any other worker node, but it also runs certain special processes related to scheduling and orchestration. Ray-enabled scripts are run on the head node and depending on the resources (number of CPUs, GPUs, etc) they request, will be scheduled to run certain tasks on the worker nodes. For more on key concepts behind a Ray cluster, you can refer this [doc](https://docs.ray.io/en/latest/cluster/key-concepts.html#cluster-key-concepts).
+
+## Sanity check
+
+To run a sanity check on whether your ray cluster is setup properly, execute the following on the head node:
+
+```bash
+ray status
+```
+
+The output should have a summary of your Ray cluster - list of all the nodes in your cluster, the number of CPUs and GPUs in your cluster, etc. For example, if you have a cluster with 1 CPU-only head node and 2 4xL40S worker nodes, the output can look like this:
+
+
+```
+Node status
+---------------------------------------------------------------
+Active:
+ 1 head
+Idle:
+ 2 4xL40S:48CPU-384GB
+Pending:
+ (no pending nodes)
+Recent failures:
+ (no failures)
+
+Resources
+---------------------------------------------------------------
+Usage:
+ 0.0/96.0 CPU
+ 0.0/8.0 GPU
+ 0B/800.00GiB memory
+ 0B/229.57GiB object_store_memory
+
+Demands:
+ (no resource demands)
+```
+
+You should also be able to see the same on the [Ray dashboard](https://docs.ray.io/en/latest/ray-observability/getting-started.html).
+
+
+## Configuring training with Ray Train
+
+You can find an example configuration at `configs/llama-3/lora-1b-ray.yaml`.
+
+The key parameters to note here are:
+
+```yaml
+use_ray: true
+ray_num_workers: 4
+# optional
+resources_per_worker:
+    GPU: 1
+```
+
+- `use_ray`: This is the flag that enables the Ray Train integration. You can either use the corresponding `--use-ray` flag in the CLI or set `use_ray` in the config file.
+- `ray_num_workers`: This is the number of workers/GPUs to use for training.
+- `resources_per_worker`: This is the Ray [resource request](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html) for each worker. This can be used to request a specific GPU type or a custom resource for each worker. For example, if your ray cluster has GPUs of different types, and you only want to use NVIDIA L40S GPUs, you can do
+
+```yaml
+resources_per_worker:
+    accelerator_type:L40S: 0.001
+```
+
+## Launching training
+
+You can simply run the following command on the head node:
+
+```bash
+axolotl train examples/llama-3/lora-1b-ray.yml --use-ray
+```
+
+This will launch training on the head node and workers will be scheduled automatically by Ray Train to run on the appropriate head or worker nodes.
+
+You can also monitor training progress on the Ray dashboard.
+
+Coming back to the example on a Ray cluster with 1 head node and 2 4xL40S worker nodes, let's say you want to make use of all 8 GPUs. You would be able to just set `ray_num_workers: 8` and run the previous command. The Cluster tab will show the following:
+
+![Ray dashboard](./images/ray-cluster-dashboard.png)
diff --git a/docs/reward_modelling.qmd b/docs/reward_modelling.qmd
new file mode 100644
index 0000000000..b5cf3010dc
--- /dev/null
+++ b/docs/reward_modelling.qmd
@@ -0,0 +1,65 @@
+---
+title: "Reward Modelling"
+description: "Reward models are used to guide models towards behaviors which is preferred by humans, by training over large datasets annotated with human preferences. "
+---
+
+### Overview
+
+Reward modelling is a technique used to train models to predict the reward or value of a given input. This is particularly useful in reinforcement learning scenarios where the model needs to evaluate the quality of its actions or predictions.
+We support the reward modelling techniques supported by `trl`.
+
+### (Outcome) Reward Models
+
+Outcome reward models are trained using data which contains preference annotations for an entire interaction between the user and model (e.g. rather than per-turn or per-step).
+For improved training stability, you can use the `center_rewards_coefficient` parameter to encourage mean-zero reward outputs ([see TRL docs](https://huggingface.co/docs/trl/v0.10.1/en/reward_trainer#centering-rewards)).
+
+```yaml
+base_model: google/gemma-2-2b
+model_type: AutoModelForSequenceClassification
+num_labels: 1
+tokenizer_type: AutoTokenizer
+
+reward_model: true
+chat_template: gemma
+datasets:
+  - path: argilla/distilabel-intel-orca-dpo-pairs
+    type: bradley_terry.chat_template
+
+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
+model_type: AutoModelForTokenClassification
+num_labels: 2
+
+process_reward_model: true
+datasets:
+  - path: trl-lib/math_shepherd
+    type: stepwise_supervised
+    split: train
+
+val_set_size: 0.1
+eval_steps: 100
+```
+
+Please see [stepwise_supervised](dataset-formats/stepwise_supervised.qmd) for more details on the dataset format.
diff --git a/docs/rlhf.qmd b/docs/rlhf.qmd
index b8b2bded09..a27bb2966d 100644
--- a/docs/rlhf.qmd
+++ b/docs/rlhf.qmd
@@ -1,26 +1,45 @@
 ---
 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
 ---
 
-### Overview
+## Overview
 
 Reinforcement Learning from Human Feedback is a method whereby a language model is optimized from data using human
 feedback. Various methods include, but not limited to:
 
-- Proximal Policy Optimization (PPO) (not yet supported in axolotl)
-- Direct Preference Optimization (DPO)
-- Identity Preference Optimization (IPO)
+- [Direct Preference Optimization (DPO)](#dpo)
+- [Identity Preference Optimization (IPO)](#ipo)
+- [Kahneman-Tversky Optimization (KTO)](#kto)
+- [Odds Ratio Preference Optimization (ORPO)](#orpo)
+- [Group Relative Policy Optimization (GRPO)](#grpo) — see also the [GRPO deep dive](grpo.qmd) for async features, custom rewards, and scaling
+- [Group Reward-Decoupled Policy Optimization (GDPO)](#gdpo)
+- [Energy-Based Fine-Tuning (EBFT)](#ebft) — see also the [EBFT guide](ebft.qmd) for detailed mode comparisons and configuration
+- [NeMo Gym Integration](#nemo-gym-integration)
 
+For help choosing between these methods, see [Choosing a Fine-Tuning Method](choosing_method.qmd).
 
-### RLHF using Axolotl
 
->[!IMPORTANT]
->This is a BETA feature and many features are not fully implemented. You are encouraged to open new PRs to improve the integration and functionality.
+## RLHF using Axolotl
 
-The various RL training methods are implemented in trl and wrapped via axolotl. Below are various examples with how you can use various preference datasets to train models that use ChatML
+::: {.callout-important}
+This is a BETA feature and many features are not fully implemented. You are encouraged to open new PRs to improve the integration and functionality.
+:::
+
+We rely on the [TRL](https://github.com/huggingface/trl) library for implementations of various RL training methods, which we wrap around to expose in axolotl. Each method has their own supported ways of loading datasets and prompt formats.
+
+::: {.callout-tip}
+You can find what each method supports by going into `src/axolotl/prompt_strategies/{method}` where `{method}` is one of our supported methods. The `type: ` can be retrieved from `{method}.{function_name}`.
+:::
+
+### DPO
+
+Example config:
 
-#### DPO
 ```yaml
 rl: dpo
 datasets:
@@ -29,15 +48,284 @@ datasets:
     type: chatml.intel
   - path: argilla/ultrafeedback-binarized-preferences
     split: train
-    type: chatml.argilla
+    type: chatml
+```
+
+DPO supports the following types with the following dataset format:
+
+#### chatml.argilla
+
+```json
+{
+    "system": "...", // optional
+    "instruction": "...",
+    "chosen_response": "...",
+    "rejected_response": "..."
+}
+```
+
+#### chatml.argilla_chat
+
+```json
+{
+    "chosen": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ],
+    "rejected": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ]
+}
+```
+
+#### chatml.icr
+
+```json
+{
+    "system": "...", // optional
+    "input": "...",
+    "chosen": "...",
+    "rejected": "..."
+}
+```
+
+#### chatml.intel
+
+```json
+{
+    "system": "...", // optional
+    "question": "...",
+    "chosen": "...",
+    "rejected": "..."
+}
+```
+
+#### chatml.prompt_pairs
+
+```json
+{
+    "system": "...", // optional
+    "prompt": "...",
+    "chosen": "...",
+    "rejected": "..."
+}
+```
+
+#### chatml.ultra
+
+```json
+{
+    "system": "...", // optional
+    "prompt": "...",
+    "chosen": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ],
+    "rejected": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ]
+}
+```
+
+#### llama3.argilla
+
+```json
+{
+    "system": "...", // optional
+    "instruction": "...",
+    "chosen_response": "...",
+    "rejected_response": "..."
+}
+```
+
+#### llama3.argilla_chat
+
+```json
+{
+    "chosen": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ],
+    "rejected": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ]
+}
+```
+
+#### llama3.icr
+
+```json
+{
+    "system": "...", // optional
+    "input": "...",
+    "chosen": "...",
+    "rejected": "..."
+}
+```
+
+#### llama3.intel
+
+```json
+{
+    "system": "...", // optional
+    "question": "...",
+    "chosen": "...",
+    "rejected": "..."
+}
+```
+
+#### llama3.prompt_pairs
+
+```json
+{
+    "system": "...", // optional
+    "prompt": "...",
+    "chosen": "...",
+    "rejected": "..."
+}
+```
+
+#### llama3.ultra
+
+```json
+{
+    "system": "...", // optional
+    "prompt": "...",
+    "chosen": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ],
+    "rejected": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ]
+}
+```
+
+#### zephyr.nectar
+
+```json
+{
+    "prompt": "...",
+    "answers": [
+        {
+            "answer": "...",
+            "rank": 1
+        },
+        {
+            "answer": "...",
+            "rank": 2
+        }
+        // ... more answers with ranks
+    ]
+}
+```
+
+#### chat_template.argilla_chat
+
+```json
+{
+    "chosen": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ],
+    "rejected": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ]
+}
+```
+
+#### chat_template.default
+
+```yaml
+rl: dpo
+datasets:
+  - path: ...
+    split: train
+    type: chat_template.default
+    field_messages: "messages"
+    field_chosen: "chosen"
+    field_rejected: "rejected"
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      user: ["user"]
+      assistant: ["assistant"]
+      system: ["system"]
 ```
 
-#### IPO
+Sample input format:
+
+```json
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": "..."
+        },
+        {
+            "role": "user",
+            "content": "..."
+        },
+        // ... more messages
+    ],
+    "chosen": {
+        "role": "assistant",
+        "content": "..."
+    },
+    "rejected": {
+        "role": "assistant",
+        "content": "..."
+    }
+}
+```
+
+#### user_defined.default
+
+For custom behaviors,
+
 ```yaml
-rl: ipo
+rl: dpo
+datasets:
+  - path: ...
+    split: train
+    type:
+      field_prompt: "prompt"
+      field_system: "system"
+      field_chosen: "chosen"
+      field_rejected: "rejected"
+      prompt_format: "{prompt}"
+      chosen_format: "{chosen}"
+      rejected_format: "{rejected}"
 ```
 
-#### ORPO
+The input format is a simple JSON input with customizable fields based on the above config.
+
+```json
+{
+    "system": "...",  // optional
+    "prompt": "...",
+    "chosen": "...",
+    "rejected": "..."
+}
+```
+
+### IPO
+
+As IPO is just DPO with a different loss function, all supported dataset formats for [DPO](#dpo) are also supported for IPO.
+
+```yaml
+rl: dpo
+dpo_loss_type: ["ipo"]
+```
+*Note:* Passing `rl: ipo` directly is still supported, but will soon be deprecated.
+
+### ORPO
 
 Paper: https://arxiv.org/abs/2403.07691
 
@@ -52,7 +340,1010 @@ datasets:
     type: chat_template.argilla
 ```
 
-#### Using local dataset files
+ORPO supports the following types with the following dataset format:
+
+#### chat_template.argilla
+
+```json
+{
+    "system": "...",  // optional
+    "prompt": "...",  // if available, will be taken as user message for single-turn instead of from list below
+
+    // chosen/rejected should be same till last content and only even-number of alternating user/assistant turns
+    "chosen": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ],
+    "rejected": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ]
+}
+```
+
+### KTO
+
+```yaml
+rl: kto
+rl_beta: 0.1  # default
+kto_desirable_weight: 1.0  # default
+kto_undesirable_weight: 1.0  # default
+
+remove_unused_columns: false
+
+datasets:
+  - path: argilla/ultrafeedback-binarized-preferences-cleaned-kto
+    type: llama3.ultra
+    split: train
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+```
+
+KTO supports the following types with the following dataset format:
+
+#### chatml.argilla
+
+```json
+{
+    "system": "...", // optional
+    "instruction": "...",
+    "completion": "..."
+}
+```
+
+#### chatml.argilla_chat
+
+```json
+{
+    "chosen": [
+        {"role": "user", "content": "..."}
+    ],
+    "completion": [
+        {"role": "assistant", "content": "..."}
+    ]
+}
+```
+
+#### chatml.intel
+
+```json
+{
+    "system": "...", // optional
+    "question": "...",
+    "completion": "..."
+}
+```
+
+#### chatml.prompt_pairs
+
+```json
+{
+    "system": "...", // optional
+    "prompt": "...",
+    "completion": "..."
+}
+```
+
+#### chatml.ultra
+
+```json
+{
+    "system": "...", // optional
+    "prompt": "...",
+    "completion": "..."
+}
+```
+
+#### llama3.argilla
+
+```json
+{
+    "system": "...", // optional
+    "instruction": "...",
+    "completion": "..."
+}
+```
+
+#### llama3.argilla_chat
+
+```json
+{
+    "completion": [
+        {"role": "user", "content": "..."},
+        {"role": "assistant", "content": "..."}
+    ]
+}
+```
+
+#### llama3.intel
+
+```json
+{
+    "system": "...", // optional
+    "question": "...",
+    "completion": "..."
+}
+```
+
+#### llama3.prompt_pairs
+
+```json
+{
+    "system": "...", // optional
+    "prompt": "...",
+    "completion": "..."
+}
+```
+
+#### llama3.ultra
+
+```json
+{
+    "system": "...", // optional
+    "prompt": "...",
+    "completion": "..."
+}
+```
+
+#### user_defined.default
+
+For custom behaviors,
+
+```yaml
+rl: kto
+datasets:
+  - path: ...
+    split: train
+    type:
+      field_prompt: "prompt"
+      field_system: "system"
+      field_completion: "completion"
+      field_label: "label"
+      prompt_format: "{prompt}"
+      completion_format: "{completion}"
+```
+
+The input format is a simple JSON input with customizable fields based on the above config.
+
+```json
+{
+    "system": "...",  // optional
+    "prompt": "...",
+    "completion": "...",
+    "label": "..."
+}
+```
+
+### GRPO
+
+::: {.callout-tip}
+Check out our [GRPO cookbook](https://github.com/axolotl-ai-cloud/grpo_code). For a comprehensive guide covering async training, custom rewards, importance sampling, and scaling, see the [GRPO deep dive](grpo.qmd).
+:::
+
+In the latest GRPO implementation, `vLLM` is used to significantly speedup trajectory generation during training. 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
+```
+
+::: {.callout-note}
+Due to TRL's implementation with vLLM, the vLLM instance must use the last N GPUs instead of the first N GPUs. This is why in the example above, we use `CUDA_VISIBLE_DEVICES=2,3` for the vLLM instance.
+:::
+
+#### 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:
+
+```python
+# rewards.py
+import random
+
+def rand_reward_func(completions, **kwargs) -> list[float]:
+    return [random.uniform(0, 1) for _ in completions]
+
+def oai_gsm8k_transform(cfg, *args, **kwargs):
+    def transform_fn(example, tokenizer=None):
+        label = example["answer"].split("####")[-1].strip().replace(",", "")
+        return {
+            "prompt": [{"role": "user", "content": example["question"]},],
+            "answer": label,
+        }
+    return transform_fn, {"remove_columns": ["question"]}
+```
+
+```yaml
+rl: grpo
+
+trl:
+    beta: 0.001
+    max_completion_length: 256
+    use_vllm: True
+    num_generations: 4
+    reward_funcs: ["rewards.rand_reward_func"]    # format: '{file_name}.{fn_name}'
+    reward_weights: [1.0]
+datasets:
+  - path: openai/gsm8k
+    name: main
+    type: rewards.oai_gsm8k_transform  # format: '{file_name}.{fn_name}'
+```
+
+To see other examples of custom reward functions, please see [TRL GRPO Docs](https://github.com/huggingface/trl/blob/main/docs/source/grpo_trainer.md#using-a-custom-reward-function).
+
+To see all configs, please see [TRLConfig](https://github.com/axolotl-ai-cloud/axolotl/blob/v0.9.2/src/axolotl/utils/schemas/trl.py).
+
+#### OpenEnv Rollout Functions
+
+GRPO supports custom rollout functions for OpenEnv-style environments, enabling interactive tasks like web browsing, code execution, or tool use. This allows you to implement custom generation logic that interacts with external environments.
+
+For example, to implement a simple math-solving environment with step-by-step verification:
+
+```python
+# math_env.py
+import re
+
+def math_solver_rollout(model, processing_class, prompts, generation_config=None):
+    """
+    Custom rollout function that generates step-by-step math solutions.
+
+    Args:
+        model: The language model
+        processing_class: The tokenizer/processing_class
+        prompts: List of prompt dicts (with 'messages' key for chat format)
+        generation_config: Optional generation configuration
+
+    Returns:
+        List of completion strings
+    """
+    completions = []
+
+    for prompt in prompts:
+        # Apply chat template to prompt
+        messages = prompt.get("messages", [])
+        formatted_prompt = processing_class.apply_chat_template(
+            messages, processing_class=False, add_generation_prompt=True
+        )
+
+        # Generate step-by-step solution
+        full_response = ""
+        for step in range(5):  # Max 5 reasoning steps
+            current_input = formatted_prompt + full_response + "\nNext step:"
+            inputs = processing_class(current_input, return_tensors="pt").to(model.device)
+
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=100,
+                generation_config=generation_config,
+            )
+            step_text = processing_class.decode(
+                outputs[0][inputs.input_ids.shape[1]:],
+                skip_special_tokens=True
+            )
+
+            # Check if solution is complete
+            if "FINAL ANSWER:" in step_text:
+                full_response += step_text
+                break
+            full_response += step_text + "\n"
+
+        completions.append(full_response)
+
+    return completions
+
+def math_reward(prompts, completions, answers, **kwargs):
+    """Reward function that checks mathematical correctness"""
+    rewards = []
+    for completion, correct_answer in zip(completions, answers):
+        # Extract predicted answer
+        match = re.search(r"FINAL ANSWER:\s*(.+)", completion)
+        predicted = match.group(1).strip() if match else ""
+
+        # Compare with correct answer
+        reward = 1.0 if predicted == str(correct_answer) else 0.0
+        rewards.append(reward)
+
+    return rewards
+
+def math_transform(cfg, *args, **kwargs):
+    """Transform dataset to GRPO format with answer field"""
+    def transform_fn(example, processing_class=None):
+        return {
+            "prompt": [{"role": "user", "content": example["question"]}],
+            "answer": str(example["answer"]),
+        }
+    return transform_fn, {"remove_columns": ["question"]}
+```
+
+```yaml
+rl: grpo
+
+trl:
+  beta: 0.001
+  max_completion_length: 512
+  num_generations: 4
+  rollout_func: "math_env.math_solver_rollout"  # Custom rollout function
+  reward_funcs: ["math_env.math_reward"]
+  reward_weights: [1.0]
+
+datasets:
+  - path: openai/gsm8k
+    name: main
+    type: math_env.math_transform
+```
+
+The `rollout_func` parameter accepts a fully qualified name (e.g., `module_name.function_name`) that points to a callable function in your local directory. The function receives:
+
+- `model`: The language model
+- `processing_class`: The tokenizer/processing class
+- `prompts`: List of prompt dictionaries
+- `generation_config` (optional): Generation configuration
+
+And should return a list of completion strings.
+
+For more OpenEnv examples, see [TRL OpenEnv Documentation](https://huggingface.co/docs/trl/main/en/openenv).
+
+#### GRPO with DAPO/Dr. GRPO loss
+
+The DAPO paper and subsequently Dr. GRPO paper proposed an alternative loss function for GRPO to remediate the penalty in longer responses.
+
+```yaml
+trl:
+  loss_type: dr_grpo
+  # Normalizes loss based on max completion length (default: 256)
+  max_completion_length:
+```
+
+For more information, see [GRPO docs](https://huggingface.co/docs/trl/v0.17.0/en/grpo_trainer#loss-types).
+
+#### Async GRPO
+
+Async GRPO overlaps vLLM generation with training by producing rollouts in a background thread. While the model trains on the current batch, the next batch is already being generated. This can significantly reduce wall-clock time per step.
+
+```yaml
+trl:
+  use_data_producer: true     # Enable data producer protocol
+  use_vllm: true
+  async_prefetch: true         # Generate rollouts in background thread
+  prefetch_depth: 1            # Number of rollouts to prefetch
+  vllm_sync_interval: 2        # Sync weights to vLLM every N steps
+```
+
+::: {.callout-note}
+Because the background thread generates completions with slightly stale model weights, async GRPO uses importance sampling correction to account for the distribution shift. This is controlled by `vllm_importance_sampling_correction: true` (default when async is enabled).
+:::
+
+##### vLLM LoRA Sync
+
+By default, weight sync to vLLM merges the LoRA adapter into the base model and broadcasts all parameters via NCCL. LoRA sync is a faster alternative that saves only the adapter weights to the filesystem and has vLLM load them natively using Punica kernels.
+
+```yaml
+adapter: lora
+lora_r: 32
+lora_alpha: 64
+lora_target_linear: true
+
+trl:
+  vllm_lora_sync: true         # Enable native LoRA sync
+```
+
+When `vllm_lora_sync: true` is set, axolotl automatically selects the LoRA-aware vLLM serve module. Start vLLM as usual:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
+```
+
+Then start training on a separate GPU:
+
+```bash
+CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml
+```
+
+::: {.callout-tip}
+LoRA sync is especially beneficial with multi-GPU training (FSDP/DeepSpeed), where NCCL merge-sync can cause GPU contention with vLLM generation.
+:::
+
+##### Streaming Partial Batch
+
+Instead of scoring the entire batch at once, streaming mode scores one prompt group at a time. This enables finer-grained zero-advantage skipping and reduces peak memory usage during scoring.
+
+```yaml
+trl:
+  streaming_partial_batch: true
+```
+
+##### Importance Sampling Correction
+
+When using async prefetch, completions are generated from a slightly older version of the model. Importance sampling (IS) correction adjusts the policy gradient to account for this distribution shift.
+
+```yaml
+trl:
+  vllm_importance_sampling_correction: true   # Enable IS correction
+  importance_sampling_level: token             # 'token' or 'sequence'
+  off_policy_mask_threshold: 0.5              # Mask sequences with IS ratio below this
+```
+
+- `importance_sampling_level: token` applies per-token IS ratios (recommended with Liger kernel)
+- `importance_sampling_level: sequence` applies per-sequence IS ratios
+- `off_policy_mask_threshold` masks out sequences where the IS ratio indicates they are too far off-policy
+
+##### Replay Buffer
+
+The replay buffer caches rollout groups that had learning signal (non-zero reward variance) and uses them to replace zero-signal groups in later batches.
+
+```yaml
+trl:
+  replay_buffer_size: 100       # Max cached groups (0 = disabled)
+  replay_recompute_logps: true  # Recompute log-probs for replayed data (recommended)
+```
+
+::: {.callout-note}
+When `replay_recompute_logps: true` (default), old log-probabilities are recomputed using the current model weights. This fixes the IS mismatch that would otherwise occur when replaying stale data.
+:::
+
+##### Deferred Re-rolling
+
+Failed prompts (where the model produces zero reward for all generations) are buffered and re-injected into later batches when the model may be better equipped to solve them.
+
+```yaml
+trl:
+  reroll_start_fraction: 0.5    # Start re-rolling after 50% of training
+  reroll_max_groups: 1          # Max groups to replace per batch
+```
+
+##### Zero-Advantage Batch Skipping
+
+When all advantages in a micro-batch are zero (no learning signal), the forward/backward pass is skipped entirely. This is enabled by default and logged as `skipped_zero_adv_batches=1`.
+
+```yaml
+trl:
+  skip_zero_advantage_batches: true   # default
+```
+
+##### Parallel Reward Workers
+
+Reward functions that use `signal.alarm()` (e.g., `math_verify`) must run in the main thread. Parallel reward workers use subprocesses to work around this limitation while enabling concurrent reward computation.
+
+```yaml
+trl:
+  reward_num_workers: 4         # Number of subprocess workers (1 = no parallelism)
+```
+
+##### Full Async GRPO Example
+
+```yaml
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+
+vllm:
+    host: 0.0.0.0
+    port: 8000
+    gpu_memory_utilization: 0.35
+    dtype: auto
+
+adapter: lora
+lora_r: 32
+lora_alpha: 64
+lora_target_linear: true
+
+rl: grpo
+trl:
+  use_data_producer: true
+  use_vllm: true
+  async_prefetch: true
+  prefetch_depth: 1
+  vllm_sync_interval: 2
+  vllm_lora_sync: true
+  streaming_partial_batch: true
+  vllm_importance_sampling_correction: true
+  off_policy_mask_threshold: 0.5
+  importance_sampling_level: token
+  num_generations: 8
+  max_completion_length: 512
+  reward_funcs:
+    - rewards.accuracy_reward
+  reroll_start_fraction: 0.5
+  replay_buffer_size: 100
+  reward_num_workers: 4
+  skip_zero_advantage_batches: true
+
+datasets:
+  - path: AI-MO/NuminaMath-TIR
+    type: rewards.prompt_transform
+    split: train
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+max_steps: 500
+learning_rate: 1e-5
+bf16: true
+gradient_checkpointing: true
+```
+
+```bash
+# Terminal 1: Start vLLM on GPU 0
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
+
+# Terminal 2: Train on GPU 1
+CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml
+```
+
+##### Multi-GPU Async GRPO
+
+Async GRPO supports FSDP and DeepSpeed ZeRO-3 for multi-GPU training. vLLM runs on one GPU while training is distributed across the remaining GPUs.
+
+**FSDP:**
+
+```yaml
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+```
+
+**DeepSpeed ZeRO-3:**
+
+```yaml
+deepspeed: deepspeed_configs/zero3_bf16.json
+gradient_checkpointing_kwargs:
+  use_reentrant: true   # Required for ZeRO-3
+```
+
+```bash
+# Terminal 1: Start vLLM on GPU 0
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
+
+# Terminal 2: Train on GPUs 0,1
+CUDA_VISIBLE_DEVICES=0,1 axolotl train config.yaml
+```
+
+::: {.callout-important}
+With multi-GPU async prefetch, only rank 0 generates completions in the background thread. Results are broadcast to all ranks on the main thread. This avoids FSDP/DeepSpeed collective deadlocks from unsynchronized background threads.
+:::
+
+### GDPO
+
+GDPO (Group Reward-Decoupled Policy Optimization) extends GRPO for multi-reward training. It addresses the **reward advantage collapse** problem by normalizing each reward function independently before combining them.
+
+::: {.callout-tip}
+Use GDPO when training with multiple reward functions. For single reward, GRPO and GDPO produce equivalent results.
+:::
+
+Paper: [https://arxiv.org/pdf/2501.05242](https://arxiv.org/pdf/2501.05242)
+
+GDPO uses TRL's native `multi_objective_aggregation` parameter under the hood. When you set `rl: gdpo`, axolotl automatically configures TRL to use `normalize_then_sum` aggregation.
+
+```yaml
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+
+vllm:
+    host: 0.0.0.0
+    port: 8000
+    tensor_parallel_size: 2
+    gpu_memory_utilization: 0.85
+
+rl: gdpo
+
+trl:
+    beta: 0.001
+    max_completion_length: 256
+    use_vllm: true
+    num_generations: 4
+    reward_funcs:
+        - rewards.format_reward
+        - rewards.correctness_reward
+    reward_weights: [1.0, 2.0]
+
+datasets:
+    - path: openai/gsm8k
+      name: main
+      type: rewards.oai_gsm8k_transform
+```
+
+You can also use GRPO with explicit aggregation control:
+
+```yaml
+rl: grpo
+trl:
+    multi_objective_aggregation: normalize_then_sum  # GDPO behavior
+    # or: sum_then_normalize  # Default GRPO behavior
+```
+
+#### GDPO vs GRPO
+
+| Aspect | GRPO | GDPO |
+|--------|------|------|
+| **Aggregation** | `sum_then_normalize` | `normalize_then_sum` |
+| **Multi-reward** | May collapse advantages | Preserves reward signals |
+| **Single reward** | Standard behavior | Equivalent to GRPO |
+
+#### Why GDPO?
+
+When using multiple rewards with GRPO, different reward combinations can produce identical advantages:
+
+```
+# Example: format + correctness rewards
+[format=0, correct=3] → sum=3
+[format=1, correct=2] → sum=3  ← GRPO sees these as equal!
+[format=2, correct=1] → sum=3
+[format=3, correct=0] → sum=3
+```
+
+GDPO normalizes each reward independently, preserving their relative differences.
+
+#### Reward Functions
+
+GDPO uses the same reward function format as GRPO:
+
+```python
+# rewards.py
+def format_reward(completions, **kwargs) -> list[float]:
+    return [1.0 if len(c) > 10 else 0.0 for c in completions]
+
+def correctness_reward(completions, answers, **kwargs) -> list[float]:
+    rewards = []
+    for completion, answer in zip(completions, answers):
+        # Your scoring logic here
+        rewards.append(score)
+    return rewards
+```
+
+#### Sequence Parallelism
+
+GDPO supports sequence parallelism for long-context training:
+
+```yaml
+rl: gdpo
+context_parallel_size: 2
+```
+
+### SimPO
+
+SimPO uses [CPOTrainer](https://huggingface.co/docs/trl/main/en/cpo_trainer) but with alternative loss function.
+
+```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).
+
+### EBFT {#ebft}
+
+::: {.callout-tip}
+For a detailed guide on EBFT modes, feature extraction, and configuration, see the [EBFT guide](ebft.qmd).
+:::
+
+EBFT (Energy-Based Fine-Tuning) fine-tunes language models by optimizing a **feature-matching loss** rather than relying on external reward functions. A frozen copy of the model extracts embeddings from both generated and ground-truth completions, and the generator is updated via REINFORCE to match the ground-truth feature moments.
+
+Paper: ["Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models"](https://arxiv.org/abs/2603.12248) (Jelassi et al., 2026)
+
+**Key advantages:**
+
+- No reward model or verifier required — works on any (prompt, completion) data
+- Applicable to non-verifiable tasks (code, translation, creative writing)
+- Operates on model rollouts (not teacher forcing), reducing distribution shift
+
+EBFT supports two modes:
+
+- **Structured mode**: For QA/instruction data with prompt + completion pairs. Uses vLLM for generation (like GRPO).
+- **Strided mode**: For unstructured text without prompt/completion splits. Uses strided block-parallel generation with flex_attention — no vLLM needed.
+
+#### Structured Mode
+
+```yaml
+base_model: Qwen/Qwen3-4B
+
+rl: ebft
+
+ebft:
+  feature_layers: [0.25, 0.5, 0.75]    # Extract features at 25%, 50%, 75% depth
+  embed_method: last_token
+  use_whitening: false
+  alignment_coef: 1.0                    # Cosine similarity reward weight
+  diversity_coef: 1.0                    # Pairwise dot product penalty
+  ce_coef: 0.0                          # Cross-entropy on GT tokens (0 = off)
+
+trl:
+  num_generations: 4
+  max_completion_length: 256
+  temperature: 0.7
+  use_vllm: true
+  vllm_server_host: 0.0.0.0
+  vllm_server_port: 8000
+  vllm_lora_sync: true                   # LoRA adapter sync (recommended)
+  vllm_sync_interval: 3
+  use_data_producer: true
+  async_prefetch: true                   # Set false for sync mode
+  scale_rewards: true
+  loss_type: grpo
+  epsilon: 0.2
+
+vllm:
+  gpu_memory_utilization: 0.5
+  max_model_len: 2048
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+    split: train[:500]
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_target_linear: true
+```
+
+```bash
+# Terminal 1: Start vLLM
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
+
+# Terminal 2: Train
+CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml
+```
+
+#### Strided Mode
+
+For unstructured text (raw code, prose). No vLLM needed — runs on a single GPU.
+
+```yaml
+base_model: meta-llama/Llama-3.2-1B
+
+rl: ebft
+
+ebft:
+  mode: strided
+  stride: 8
+  context_length: 8
+  generate_max_len: 8
+  n_samples_per_prompt: 4
+  temperature: 0.6
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: true
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  rl_coef: 1.0
+  ce_coef: 0.03
+  advantage_estimator: rloo
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_strided_structured.transform
+    split: train[:1%]
+
+attn_implementation: flex_attention   # Strided mode uses flex_attention
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true    # Required for flex_attention
+```
+
+```bash
+CUDA_VISIBLE_DEVICES=0 axolotl train config.yaml
+```
+
+::: {.callout-tip}
+See `examples/ebft/` for complete example configs covering Llama 1B/3B/8B and Qwen3 4B/8B models in both modes.
+:::
+
+#### EBFT Configuration Reference
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `ebft.feature_layers` | `[0.25, 0.5, 0.75]` | Layer depths for feature extraction (fractional) |
+| `ebft.embed_method` | `last_token` | Feature pooling: `last_token`, `mean_pooling`, `concat` |
+| `ebft.use_whitening` | `false` | SVD whitening of feature dimensions |
+| `ebft.alignment_coef` | `1.0` | Cosine similarity reward weight |
+| `ebft.diversity_coef` | `1.0` | Pairwise dot product penalty weight |
+| `ebft.ce_coef` | `0.0` | Cross-entropy loss on ground-truth tokens |
+| `ebft.mode` | `structured` | `structured` (vLLM) or `strided` (no vLLM) |
+| `ebft.stride` | — | Tokens between anchor points (strided mode) |
+| `ebft.context_length` | — | Context window per block (strided mode) |
+| `ebft.generate_max_len` | — | Tokens to generate per block (strided mode) |
+| `ebft.n_samples_per_prompt` | — | Rollouts per document (strided mode) |
+| `ebft.advantage_estimator` | `grpo` | `grpo` or `rloo` (strided mode) |
+
+### NeMo Gym Integration
+
+[NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) provides 50+ verified RL environments (math, coding, tool-use, reasoning) with deterministic reward signals. The axolotl integration supports both **single-turn** (call `/verify` after generation) and **multi-turn** (agent-based tool execution via `/run`).
+
+#### Single-Turn (Simplest)
+
+For environments that only need answer verification (math, coding challenges). No agent server needed — the reward function calls `/verify` directly on the resource server.
+
+```yaml
+base_model: Qwen/Qwen2.5-0.5B-Instruct
+
+rl: grpo
+chat_template: tokenizer_default
+
+trl:
+  use_vllm: false                          # Colocate mode (single GPU)
+  num_generations: 4
+  max_completion_length: 128
+  temperature: 0.9
+  reward_funcs:
+    - axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify
+
+plugins:
+  - axolotl.integrations.nemo_gym.NemoGymPlugin
+
+nemo_gym_enabled: true
+nemo_gym_dir: ~/Gym
+nemo_gym_auto_start: false
+nemo_gym_head_port: 11000
+nemo_gym_datasets:
+  - path: resources_servers/reasoning_gym/data/train_basic_arithmetic.jsonl
+    server_name: reasoning_gym
+
+datasets:
+  - path: ~/Gym/resources_servers/reasoning_gym/data/train_basic_arithmetic.jsonl
+    type: chat_template
+    field_messages: responses_create_params.input
+    message_field_content: content
+    message_field_role: role
+```
+
+```bash
+# Terminal 1: Start NeMo Gym resource server
+cd ~/Gym && .venv/bin/ng_run \
+    "+config_paths=[resources_servers/reasoning_gym/configs/resources_only.yaml]" \
+    "+skip_venv_if_present=true"
+
+# Terminal 2: Train
+CUDA_VISIBLE_DEVICES=0 axolotl train config.yaml
+```
+
+::: {.callout-note}
+`nemo_gym_datasets.path` is relative to `nemo_gym_dir`. Don't use absolute paths or they will be double-joined.
+:::
+
+#### Multi-Turn with Async GRPO (Recommended)
+
+For environments with tool-use (weather, search, databases). An agent server orchestrates multi-turn interactions: generate → parse tool calls → execute tools → feed results back → repeat until done.
+
+```yaml
+base_model: Qwen/Qwen3-0.6B
+
+rl: grpo
+chat_template: tokenizer_default
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj]
+
+trl:
+  use_vllm: true
+  vllm_mode: server
+  vllm_server_host: localhost
+  vllm_server_port: 8000
+  vllm_lora_sync: true
+  vllm_sync_interval: 5
+  use_data_producer: true
+  async_prefetch: true                     # 3x speedup
+  num_generations: 4
+  max_completion_length: 512
+  temperature: 0.8
+  reward_funcs:
+    - axolotl.integrations.nemo_gym.rewards.reward_env
+
+plugins:
+  - axolotl.integrations.nemo_gym.NemoGymPlugin
+
+nemo_gym_enabled: true
+nemo_gym_auto_start: false
+nemo_gym_head_port: 11000
+nemo_gym_multi_turn: true
+nemo_gym_verify_timeout: 120
+nemo_gym_datasets:
+  - path: resources_servers/example_single_tool_call/data/weather_tool_calling.jsonl
+    server_name: example_single_tool_call
+
+datasets:
+  - path: ~/Gym/resources_servers/example_single_tool_call/data/weather_tool_calling.jsonl
+    type: chat_template
+    field_messages: responses_create_params.input
+    message_field_content: content
+    message_field_role: role
+
+vllm:
+  gpu_memory_utilization: 0.85
+  max_model_len: 2048
+```
+
+Multi-turn requires three services running:
+
+```bash
+# Terminal 1: vLLM with LoRA + tool calling
+VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 CUDA_VISIBLE_DEVICES=0 \
+    python -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen3-0.6B --max-model-len 2048 \
+    --gpu-memory-utilization 0.85 \
+    --enable-lora --max-lora-rank 64 \
+    --enable-auto-tool-choice --tool-call-parser hermes
+
+# Terminal 2: NeMo Gym servers (resource + model proxy + agent)
+cd ~/Gym && .venv/bin/ng_run \
+    "+config_paths=[configs/axolotl_tool_calling.yaml]" \
+    "+skip_venv_if_present=true"
+
+# Terminal 3: Training
+CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml
+```
+
+::: {.callout-important}
+Multi-turn requires a NeMo Gym agent config YAML that defines three components: a resource server (tools + `/verify`), a model server proxy (forwards to your vLLM), and an agent server (orchestrates `/run`). See the [NeMo Gym README](https://github.com/NVIDIA-NeMo/Gym) for agent config format.
+:::
+
+#### NeMo Gym Prerequisites
+
+```bash
+# Clone and set up NeMo Gym
+git clone https://github.com/NVIDIA-NeMo/Gym.git ~/Gym
+cd ~/Gym
+uv venv --python 3.12 && source .venv/bin/activate && uv sync
+
+# Fix pycosat build (GCC 13+)
+CFLAGS="" uv pip install pycosat --python .venv/bin/python --no-build-isolation
+```
+
+#### NeMo Gym Configuration Reference
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `nemo_gym_enabled` | bool | — | Enable the NeMo Gym integration |
+| `nemo_gym_dir` | str | `~/Gym` | Path to NeMo Gym repo |
+| `nemo_gym_auto_start` | bool | `true` | Auto-start resource servers |
+| `nemo_gym_head_port` | int | `11000` | Head server port |
+| `nemo_gym_multi_turn` | bool | `false` | Enable multi-turn via agent `/run` |
+| `nemo_gym_verify_timeout` | int | `30` | Per-request timeout (seconds) |
+| `nemo_gym_datasets` | list | required | Dataset configs with `path` and `server_name` |
+
+#### Reward Functions
+
+| Function | Mode | Description |
+|----------|------|-------------|
+| `axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify` | Single-turn | Calls `/verify`, returns binary reward |
+| `axolotl.integrations.nemo_gym.rewards.reward_env` | Multi-turn | Passthrough reward from agent `/run` |
+
+### Using local dataset files
+
 ```yaml
 datasets:
   - ds_type: json
@@ -62,9 +1353,9 @@ datasets:
     type: chatml.intel
 ```
 
-#### Trl autounwrap for peft
+### TRL auto-unwrapping for PEFT
 
-Trl supports autounwrapping peft models, so that a ref model does not need to be additionally loaded, leading to less VRAM needed. This is on by default. To turn it off, pass the following config.
+TRL supports auto-unwrapping PEFT models for RL training paradigms which rely on a reference model. This significantly reduces memory pressure as an additional refreference model does not need to be loaded, and reference model log-probabilities can be obtained by disabling PEFT adapters. This is enabled by default. To turn it off, pass the following config:
 
 ```yaml
 # load ref model when adapter training.
diff --git a/docs/scripts/examples-allowlist.yml b/docs/scripts/examples-allowlist.yml
new file mode 100644
index 0000000000..f2fe88e08b
--- /dev/null
+++ b/docs/scripts/examples-allowlist.yml
@@ -0,0 +1,92 @@
+examples:
+  # December 2025
+  - name: kimi-linear
+    title: Kimi Linear
+  - name: plano
+    title: Plano Orchestrator
+  - name: mimo
+    title: MiMo
+  - name: internvl3_5
+    title: InternVL 3.5
+
+  # AllenAI
+  - name: olmo3
+    title: OLMo 3
+
+  # ArceeAI
+  - name: trinity
+    title: Trinity
+  - name: arcee
+    title: Arcee AFM
+
+  # MistralAI
+  - name: mistral-medium-3_5
+    title: Mistral Medium 3.5
+  - name: ministral3/think
+    title: Ministral 3 Thinking
+  - name: ministral3/vision
+    title: Ministral 3 Vision
+  - name: magistral/think
+    title: Magistral Thinking
+  - name: magistral/vision
+    title: Magistral Vision
+  - name: ministral
+    title: Ministral
+  - name: mistral-small
+    title: Mistral Small 3.1/3.2
+  - name: voxtral
+    title: Voxtral
+  - name: devstral
+    title: Devstral
+  - name: mistral
+    title: Mistral 7B
+
+  # Meta
+  - name: llama-4
+    title: Llama 4
+  - name: llama-2
+    title: Llama 2
+
+  # Alibaba
+  - name: qwen3-next
+    title: Qwen 3 Next
+  - name: qwen3
+    title: Qwen 3
+
+  # Google
+  - name: gemma3n
+    title: Gemma 3n
+
+  # Swiss AI
+  - name: apertus
+    title: Apertus
+
+  # GPT-OSS
+  - name: gpt-oss
+    title: GPT-OSS
+  - name: seed-oss
+    title: Seed-OSS
+
+  # Microsoft
+  - name: phi
+    title: Phi
+
+  # SmolVLM
+  - name: smolvlm2
+    title: SmolVLM 2
+
+  # IBM
+  - name: granite4
+    title: Granite 4
+
+  # LiquidAI
+  - name: LiquidAI
+    title: Liquid Foundation Models 2
+
+  # Other
+  - name: hunyuan
+    title: Hunyuan
+  - name: jamba
+    title: Jamba
+  - name: orpheus
+    title: Orpheus
diff --git a/docs/scripts/generate_config_docs.py b/docs/scripts/generate_config_docs.py
new file mode 100644
index 0000000000..6efa2038bf
--- /dev/null
+++ b/docs/scripts/generate_config_docs.py
@@ -0,0 +1,749 @@
+# type: ignore
+
+"""
+Quarto documentation generation from Pydantic models. Uses Pydantic model source code
+to automatically group fields, including inherited fields from parent classes.
+"""
+
+import ast
+import inspect
+import textwrap
+import types
+import typing
+from typing import Any, FrozenSet, Type, Union
+
+from pydantic import BaseModel
+
+from axolotl.utils.schemas.config import AxolotlInputConfig
+
+
+class QuartoGenerator:
+    """Generate Quarto documentation from Pydantic models."""
+
+    def __init__(self):
+        self._class_fields_cache = {}
+        self._inheritance_map_cache = {}
+        self._nested_models_cache = {}
+
+    def _get_direct_fields(self, cls: Type[BaseModel]) -> FrozenSet[str]:
+        """Get fields defined directly in a single class (not inherited)."""
+        if cls in self._class_fields_cache:
+            return self._class_fields_cache[cls]
+
+        fields = set()
+
+        # Get annotated fields
+        if hasattr(cls, "__annotations__"):
+            fields.update(cls.__annotations__.keys())
+
+        # Filter out private/special methods
+        fields = {f for f in fields if not f.startswith("_")}
+
+        result = frozenset(fields)
+        self._class_fields_cache[cls] = result
+        return result
+
+    def _is_pydantic_model(self, type_obj) -> bool:
+        """Check if a type is a Pydantic BaseModel."""
+        return inspect.isclass(type_obj) and issubclass(type_obj, BaseModel)
+
+    def _extract_nested_type(self, field_type) -> Any:
+        """Extract the actual type from complex type annotations."""
+        # Handle Annotated types (Python 3.9+)
+        if hasattr(typing, "get_origin") and hasattr(typing, "get_args"):
+            origin = typing.get_origin(field_type)
+            args = typing.get_args(field_type)
+
+            if origin is not None:
+                # Handle Annotated[SomeType, ...] - extract the first argument
+                if hasattr(typing, "Annotated") and origin is typing.Annotated:
+                    if args:
+                        return self._extract_nested_type(
+                            args[0]
+                        )  # Recursively process the actual type
+
+                # Handle list[SomeType], List[SomeType], etc.
+                elif origin in (list, typing.List):
+                    if args:
+                        return self._extract_nested_type(
+                            args[0]
+                        )  # Extract element type
+
+                # Handle Union types (including | syntax)
+                elif origin is typing.Union:
+                    # Get non-None types from the Union
+                    non_none_types = [arg for arg in args if arg is not type(None)]
+                    if len(non_none_types) >= 1:
+                        # Prioritize Pydantic models over primitive types
+                        pydantic_models = [
+                            arg
+                            for arg in non_none_types
+                            if self._is_pydantic_model(arg)
+                        ]
+                        if pydantic_models:
+                            # Return the first Pydantic model found
+                            return self._extract_nested_type(pydantic_models[0])
+
+                        # No Pydantic models, return the first non-None type
+                        return self._extract_nested_type(non_none_types[0])
+
+        # Handle new Python 3.10+ union syntax (PeftConfig | None)
+        if hasattr(field_type, "__class__") and field_type.__class__ is types.UnionType:
+            # Get non-None types from the Union
+            non_none_types = [
+                arg for arg in field_type.__args__ if arg is not type(None)
+            ]
+            if len(non_none_types) >= 1:
+                # Prioritize Pydantic models over primitive types
+                pydantic_models = [
+                    arg for arg in non_none_types if self._is_pydantic_model(arg)
+                ]
+                if pydantic_models:
+                    return self._extract_nested_type(pydantic_models[0])
+                return self._extract_nested_type(non_none_types[0])
+
+        # Handle old typing.Union syntax (fallback)
+        if hasattr(field_type, "__origin__"):
+            if field_type.__origin__ is Union:
+                # Get non-None types from the Union
+                non_none_types = [
+                    arg for arg in field_type.__args__ if arg is not type(None)
+                ]
+                if len(non_none_types) >= 1:
+                    # Prioritize Pydantic models over primitive types
+                    pydantic_models = [
+                        arg for arg in non_none_types if self._is_pydantic_model(arg)
+                    ]
+                    if pydantic_models:
+                        return self._extract_nested_type(pydantic_models[0])
+                    return self._extract_nested_type(non_none_types[0])
+            # Handle other generic types like dict[str, Any], etc.
+            elif hasattr(field_type, "__args__"):
+                return field_type
+
+        return field_type
+
+    def _extract_all_pydantic_models_from_type(
+        self, field_type
+    ) -> list[type[BaseModel]]:
+        """Extract all Pydantic models from a type annotation, including from Unions."""
+        models = []
+
+        if field_type is None:
+            return models
+
+        # Handle Annotated types
+        if hasattr(typing, "get_origin") and hasattr(typing, "get_args"):
+            origin = typing.get_origin(field_type)
+            args = typing.get_args(field_type)
+
+            if origin is not None:
+                # Handle Annotated[SomeType, ...] - extract from the first argument
+                if hasattr(typing, "Annotated") and origin is typing.Annotated:
+                    if args:
+                        models.extend(
+                            self._extract_all_pydantic_models_from_type(args[0])
+                        )
+                    return models
+
+                # Handle list[SomeType], List[SomeType], etc.
+                if origin in (list, typing.List):
+                    if args:
+                        models.extend(
+                            self._extract_all_pydantic_models_from_type(args[0])
+                        )
+                    return models
+
+                # Handle Union types
+                if origin is typing.Union:
+                    for arg in args:
+                        if arg is not type(None):  # Skip None type
+                            models.extend(
+                                self._extract_all_pydantic_models_from_type(arg)
+                            )
+                    return models
+
+        # Handle new Python 3.10+ union syntax
+        if hasattr(field_type, "__class__") and field_type.__class__ is types.UnionType:
+            for arg in field_type.__args__:
+                if arg is not type(None):  # Skip None type
+                    models.extend(self._extract_all_pydantic_models_from_type(arg))
+            return models
+
+        # Handle old typing.Union syntax (fallback)
+        if hasattr(field_type, "__origin__") and field_type.__origin__ is Union:
+            for arg in field_type.__args__:
+                if arg is not type(None):  # Skip None type
+                    models.extend(self._extract_all_pydantic_models_from_type(arg))
+            return models
+
+        # Check if this type itself is a Pydantic model
+        if self._is_pydantic_model(field_type):
+            models.append(field_type)
+
+        return models
+
+    def _get_nested_models(
+        self, model_class: type[BaseModel], visited=None
+    ) -> dict[str, type[BaseModel]]:
+        """Get all nested Pydantic models from a model class."""
+        if visited is None:
+            visited = set()
+
+        # Avoid infinite recursion
+        if model_class in visited:
+            return {}
+
+        if model_class in self._nested_models_cache:
+            return self._nested_models_cache[model_class]
+
+        visited.add(model_class)
+        nested_models = {}
+
+        # Check all fields in the model
+        for field_info in model_class.model_fields.values():
+            field_type = self._extract_nested_type(field_info.annotation)
+
+            if self._is_pydantic_model(field_type):
+                nested_models[field_type.__name__] = field_type
+                # Recursively get nested models from this nested model
+                deeper_nested = self._get_nested_models(field_type, visited.copy())
+                nested_models.update(deeper_nested)
+
+        self._nested_models_cache[model_class] = nested_models
+        return nested_models
+
+    def _build_inheritance_map(self, child_class: Type[BaseModel]):
+        """Build inheritance map for a class and all its parents."""
+        if child_class in self._inheritance_map_cache:
+            return self._inheritance_map_cache[child_class]
+
+        inheritance_map = {}
+
+        # Get MRO and filter out BaseModel and object
+        mro_classes = [
+            cls
+            for cls in child_class.__mro__
+            if cls not in (BaseModel, object) and hasattr(cls, "__annotations__")
+        ]
+
+        # Process each class in the MRO
+        for cls in mro_classes:
+            inheritance_map[cls] = self._get_direct_fields(cls)
+
+        self._inheritance_map_cache[child_class] = inheritance_map
+        return inheritance_map
+
+    def _wrap_comment(self, text: str, width: int = 88) -> list[str]:
+        """Wrap a comment to specified width, accounting for '# ' prefix."""
+        if not text.strip():
+            return ["#"]
+
+        # Account for "# " prefix (2 characters)
+        content_width = width - 2
+        wrapped_lines = textwrap.wrap(text, width=content_width)
+        return [f"# {line}" for line in wrapped_lines]
+
+    def _extract_type_from_source(
+        self, model_class: type[BaseModel], field_name: str
+    ) -> str:
+        """Extract the actual type annotation text from source code, checking inheritance chain."""
+        # Use inheritance map to check classes efficiently
+        inheritance_map = self._build_inheritance_map(model_class)
+
+        # Check classes in MRO order
+        for cls in model_class.__mro__:
+            if cls in inheritance_map and field_name in inheritance_map[cls]:
+                type_annotation = self._get_type_from_class_source(cls, field_name)
+                if type_annotation != "unknown":
+                    return type_annotation
+
+        return "unknown"
+
+    def _get_type_from_class_source(self, class_obj: type, field_name: str) -> str:
+        """Extract type annotation from a specific class's source code."""
+        try:
+            source = inspect.getsource(class_obj)
+            tree = ast.parse(source)
+        except (OSError, TypeError):
+            return "unknown"
+
+        # Find the class definition
+        for node in tree.body:
+            if isinstance(node, ast.ClassDef) and node.name == class_obj.__name__:
+                # Find the field assignment
+                for body_node in node.body:
+                    if isinstance(body_node, ast.AnnAssign) and isinstance(
+                        body_node.target, ast.Name
+                    ):
+                        if body_node.target.id == field_name and body_node.annotation:
+                            return ast.unparse(body_node.annotation)
+                break
+
+        return "unknown"
+
+    def _extract_field_groups_from_all_classes(
+        self, model_class: type[BaseModel]
+    ) -> list[dict]:
+        """Extract field groups from all classes in the inheritance hierarchy."""
+        all_groups = []
+        inheritance_map = self._build_inheritance_map(model_class)
+
+        # Get all Pydantic base classes in MRO order (most specific first)
+        # This puts AxolotlInputConfig fields first, then parent class fields
+        pydantic_classes = [
+            cls
+            for cls in model_class.__mro__
+            if cls in inheritance_map and inheritance_map[cls]
+        ]
+
+        # Extract groups from each class
+        for cls in pydantic_classes:
+            class_groups = self._extract_field_groups_from_source(cls)
+            for group in class_groups:
+                all_groups.append(group)
+
+        # If no groups found, create a default grouping by class
+        if not all_groups:
+            for cls in pydantic_classes:
+                fields_in_class = inheritance_map[cls]
+                if fields_in_class:
+                    all_groups.append(
+                        {
+                            "fields": list(fields_in_class),
+                        }
+                    )
+
+        return all_groups
+
+    def _extract_field_groups_from_source(
+        self, model_class: type[BaseModel]
+    ) -> list[dict]:
+        """Extract field groups from source code based on blank lines and comments."""
+        try:
+            source = inspect.getsource(model_class)
+            tree = ast.parse(source)
+        except (OSError, TypeError):
+            # Fallback if we can't get source code
+            fields_in_class = self._get_direct_fields(model_class)
+            if fields_in_class:
+                return [
+                    {
+                        "fields": list(fields_in_class),
+                    }
+                ]
+            return []
+
+        groups = []
+        current_group_fields = []
+        current_group_comment = None
+
+        # Find the class definition
+        class_node = None
+        for node in ast.walk(tree):
+            if isinstance(node, ast.ClassDef) and node.name == model_class.__name__:
+                class_node = node
+                break
+
+        if not class_node:
+            fields_in_class = self._get_direct_fields(model_class)
+            if fields_in_class:
+                return [
+                    {
+                        "fields": list(fields_in_class),
+                    }
+                ]
+            return []
+
+        # Parse the source lines to detect groupings
+        source_lines = source.split("\n")
+
+        # Get fields that are actually defined in this specific class
+        fields_in_class = self._get_direct_fields(model_class)
+
+        # Find assignments that correspond to model fields for THIS class only
+        field_assignments = []
+        for node in class_node.body:
+            if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
+                field_name = node.target.id
+                if field_name in fields_in_class:
+                    field_assignments.append(
+                        {
+                            "name": field_name,
+                            "lineno": node.lineno,
+                            "end_lineno": getattr(node, "end_lineno", node.lineno),
+                        }
+                    )
+
+        if not field_assignments:
+            if fields_in_class:
+                return [
+                    {
+                        "fields": list(fields_in_class),
+                    }
+                ]
+            return []
+
+        # Sort by line number
+        field_assignments.sort(key=lambda x: x["lineno"])
+
+        # Group fields based on blank lines and comments
+        for i, field_info in enumerate(field_assignments):
+            field_name = field_info["name"]
+            current_line = field_info["lineno"]
+
+            # Check if this starts a new group (blank line before or significant gap)
+            is_new_group = False
+
+            if i == 0:
+                is_new_group = True
+            else:
+                prev_end_line = field_assignments[i - 1]["end_lineno"]
+
+                # Check for blank lines or comments between fields
+                lines_between = source_lines[prev_end_line : current_line - 1]
+                has_blank_line = any(line.strip() == "" for line in lines_between)
+                has_comment = any(
+                    line.strip().startswith("#") for line in lines_between
+                )
+
+                # Start new group if there's a blank line or comment, or significant gap
+                if has_blank_line or has_comment or (current_line - prev_end_line > 3):
+                    is_new_group = True
+
+            if is_new_group and current_group_fields:
+                # Save the previous group
+                groups.append(
+                    {
+                        "fields": current_group_fields.copy(),
+                        "description": current_group_comment,
+                    }
+                )
+                current_group_fields = []
+                current_group_comment = None
+
+            current_group_fields.append(field_name)
+
+        # Add the final group
+        if current_group_fields:
+            groups.append(
+                {
+                    "fields": current_group_fields,
+                    "description": current_group_comment,
+                }
+            )
+
+        return groups
+
+    def _generate_field_documentation(
+        self,
+        model_class: type[BaseModel],
+        field_name: str,
+        field_info: dict,
+        field_type_str: str,
+        is_required: bool,
+        indent_level: int = 0,
+        visited_models: set = None,
+    ) -> list[str]:
+        """Generate documentation for a single field, expanding nested models inline."""
+        if visited_models is None:
+            visited_models = set()
+
+        lines = []
+        indent = "  " * indent_level
+
+        # Get the actual field type for nested model detection
+        if field_name in model_class.model_fields:
+            pydantic_field_info = model_class.model_fields[field_name]
+            actual_field_type = pydantic_field_info.annotation
+        else:
+            actual_field_type = None
+
+        # Add description comment if available
+        description = field_info.get("description", "")
+        if description:
+            wrapped_lines = self._wrap_comment(description, width=88 - len(indent))
+            for line in wrapped_lines:
+                lines.append(f"{indent}{line}")
+
+        # Extract nested Pydantic models from the type annotation
+        nested_models = self._extract_all_pydantic_models_from_type(actual_field_type)
+
+        # Filter out already visited models to prevent infinite recursion
+        expandable_models = [
+            model for model in nested_models if model not in visited_models
+        ]
+
+        if expandable_models:
+            # This field contains Pydantic models that can be expanded
+
+            # Show the field with its full type annotation
+            field_line = f"{indent}{field_name}: {field_type_str}"
+            if field_info.get("default") is not None:
+                field_line += f" = {field_info['default']}"
+            if is_required:
+                field_line += " (required)"
+            lines.append(field_line)
+
+            # Add to visited to prevent infinite recursion
+            new_visited = visited_models.copy()
+            new_visited.update(expandable_models)
+
+            # Expand each nested Pydantic model
+            for i, nested_model in enumerate(expandable_models):
+                if i > 0:
+                    lines.append("\n")
+                lines.append(f"{indent}  # For {nested_model.__name__}:")
+
+                # Get nested model schema
+                try:
+                    nested_schema = nested_model.model_json_schema()
+                    nested_properties = nested_schema.get("properties", {})
+                    nested_required = nested_schema.get("required", [])
+                except Exception:
+                    # Fallback: use model fields directly
+                    nested_properties = {}
+                    nested_required = []
+                    for (
+                        nested_field_name,
+                        nested_field_info,
+                    ) in nested_model.model_fields.items():
+                        nested_description = ""
+                        if (
+                            hasattr(nested_field_info, "json_schema_extra")
+                            and nested_field_info.json_schema_extra
+                        ):
+                            nested_description = (
+                                nested_field_info.json_schema_extra.get(
+                                    "description", ""
+                                )
+                            )
+                        elif (
+                            hasattr(nested_field_info, "description")
+                            and nested_field_info.description
+                        ):
+                            nested_description = nested_field_info.description
+
+                        nested_default_val = None
+                        if (
+                            hasattr(nested_field_info, "default")
+                            and nested_field_info.default is not None
+                        ):
+                            if str(nested_field_info.default) != "PydanticUndefined":
+                                nested_default_val = nested_field_info.default
+
+                        nested_properties[nested_field_name] = {
+                            "type": "unknown",
+                            "description": nested_description,
+                            "default": nested_default_val,
+                        }
+
+                        if nested_field_info.is_required():
+                            nested_required.append(nested_field_name)
+
+                # Get field groups for the nested model
+                nested_field_groups = self._extract_field_groups_from_all_classes(
+                    nested_model
+                )
+
+                # Generate nested fields with increased indentation
+                for i, group in enumerate(nested_field_groups):
+                    if not group["fields"]:
+                        continue
+
+                    # Add blank line between groups (except before first group)
+                    if i > 0:
+                        lines.append("")
+
+                    # Process nested fields
+                    for nested_field_name in group["fields"]:
+                        if nested_field_name not in nested_properties:
+                            continue
+
+                        nested_field_info = nested_properties[nested_field_name]
+                        nested_field_type = self._extract_type_from_source(
+                            nested_model, nested_field_name
+                        )
+                        nested_is_required = nested_field_name in nested_required
+
+                        # Recursively generate documentation for nested field
+                        nested_lines = self._generate_field_documentation(
+                            nested_model,
+                            nested_field_name,
+                            nested_field_info,
+                            nested_field_type,
+                            nested_is_required,
+                            indent_level + 1,
+                            new_visited,
+                        )
+                        lines.extend(nested_lines)
+        else:
+            # Regular field (no expandable nested models)
+            field_line = f"{indent}{field_name}: {field_type_str}"
+            if field_info.get("default") is not None:
+                field_line += f" = {field_info['default']}"
+            if is_required:
+                field_line += " (required)"
+            lines.append(field_line)
+
+        return lines
+
+    def generate_qmd(
+        self,
+        model_class: type[BaseModel],
+        title: str | None = None,
+        expand_nested: bool = True,
+    ) -> str:
+        """Auto-generate config reference documentation including inherited fields."""
+
+        if title is None:
+            title = f"{model_class.__name__} Reference"
+
+        # Try to get JSON schema, with fallback for serialization issues
+        try:
+            schema = model_class.model_json_schema()
+            properties = schema.get("properties", {})
+            required = schema.get("required", [])
+        except Exception as e:
+            print(
+                f"Warning: Could not generate JSON schema ({e}). Using model fields instead."
+            )
+            # Fallback: use model fields directly
+            properties = {}
+            required = []
+            for field_name, field_info in model_class.model_fields.items():
+                # Extract description from json_schema_extra or field info
+                description = ""
+                if (
+                    hasattr(field_info, "json_schema_extra")
+                    and field_info.json_schema_extra
+                ):
+                    description = field_info.json_schema_extra.get("description", "")
+                elif hasattr(field_info, "description") and field_info.description:
+                    description = field_info.description
+
+                # Get default value
+                default_val = None
+                if hasattr(field_info, "default") and field_info.default is not None:
+                    # Handle special Pydantic default markers
+                    if str(field_info.default) != "PydanticUndefined":
+                        default_val = field_info.default
+
+                properties[field_name] = {
+                    "type": "unknown",
+                    "description": description,
+                    "default": default_val,
+                }
+
+                if field_info.is_required():
+                    required.append(field_name)
+
+        # Extract field groups from all classes in inheritance hierarchy
+        field_groups = self._extract_field_groups_from_all_classes(model_class)
+
+        # Start building QMD content
+        qmd_lines = [
+            "---",
+            f"title: {title}",
+            "description: A complete list of all configuration options.",
+            "---",
+            "",
+        ]
+
+        # Generate one big code block with all fields (inline nested expansion)
+        qmd_lines.append("```yaml")
+
+        for i, group in enumerate(field_groups):
+            if not group["fields"]:
+                continue
+
+            # Add blank line between groups (except before first group)
+            if i > 0:
+                qmd_lines.append("")
+
+            # Process fields in the order they appear in source
+            for field_name in group["fields"]:
+                if field_name not in properties:
+                    continue
+
+                field_info = properties[field_name]
+                field_type = self._extract_type_from_source(model_class, field_name)
+                is_required = field_name in required
+
+                if expand_nested:
+                    # Check if this field has nested models
+                    if field_name in model_class.model_fields:
+                        pydantic_field_info = model_class.model_fields[field_name]
+                        nested_models = self._extract_all_pydantic_models_from_type(
+                            pydantic_field_info.annotation
+                        )
+                        has_nested = bool(nested_models)
+                    else:
+                        has_nested = False
+
+                    # Add blank line before nested config
+                    if has_nested:
+                        qmd_lines.append("")
+
+                    # Use the new inline generation method
+                    field_lines = self._generate_field_documentation(
+                        model_class,
+                        field_name,
+                        field_info,
+                        field_type,
+                        is_required,
+                        indent_level=0,
+                        visited_models=set(),
+                    )
+                    qmd_lines.extend(field_lines)
+
+                    # Add blank line after nested config
+                    if has_nested:
+                        qmd_lines.append("")
+                else:
+                    # Original simple approach
+                    description = field_info.get("description", "")
+                    default = field_info.get("default")
+
+                    # Add wrapped comment for description
+                    if description:
+                        wrapped_lines = self._wrap_comment(description)
+                        qmd_lines.extend(wrapped_lines)
+
+                    line = f"{field_name}: {field_type}"
+                    if default is not None:
+                        line += f" = {default}"
+                    if is_required:
+                        line += " (required)"
+                    qmd_lines.append(line)
+
+        qmd_lines.append("```")
+
+        # Join all lines and clean up any double newlines
+        content = "\n".join(qmd_lines)
+
+        # Replace multiple consecutive newlines with just two newlines (one blank line)
+        import re
+
+        content = re.sub(r"\n{3,}", "\n\n", content)
+
+        # Ensure single newline at the very end
+        content = content.rstrip("\n") + "\n"
+
+        return content
+
+
+def main():
+    generator = QuartoGenerator()
+
+    print("Generating config reference content...")
+    qmd_content = generator.generate_qmd(AxolotlInputConfig, "Config Reference", True)
+
+    print("Writing to file...")
+    with open("docs/config-reference.qmd", "w", encoding="utf-8") as f:
+        f.write(qmd_content)
+    print("Done!")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/docs/scripts/generate_examples_docs.py b/docs/scripts/generate_examples_docs.py
new file mode 100755
index 0000000000..ba01088f17
--- /dev/null
+++ b/docs/scripts/generate_examples_docs.py
@@ -0,0 +1,424 @@
+"""
+auto generate example docs from allowlist
+"""
+
+import re
+import shutil
+import sys
+from pathlib import Path
+
+import yaml
+
+# Paths
+THIS = Path(__file__).resolve()
+ROOT = THIS.parents[2]  # repo root (docs/scripts -> docs -> ROOT)
+EXAMPLES_DIR = ROOT / "examples"
+OUTPUT_DIR = ROOT / "docs" / "models"
+ALLOWLIST_YML = THIS.parent / "examples-allowlist.yml"
+
+
+def slugify(name: str) -> str:
+    """Convert a name to a slug (lowercase, hyphens for spaces)."""
+    s = re.sub(r"[^a-zA-Z0-9\s\-]+", "", name.strip())
+    s = re.sub(r"\s+", "-", s).strip("-").lower()
+    return s or "example"
+
+
+def read_allowlist():
+    with open(ALLOWLIST_YML, "r", encoding="utf-8") as f:
+        data = yaml.safe_load(f) or {}
+    items = data.get("examples", [])
+    if not isinstance(items, list):
+        raise ValueError("`examples` must be a list in examples-allowlist.yml")
+    return items
+
+
+def find_readme(folder: Path) -> Path | None:
+    for name in ("README.md", "Readme.md", "readme.md"):
+        p = folder / name
+        if p.exists():
+            return p
+    return None
+
+
+def remove_first_h1(md: str) -> tuple[str, str | None]:
+    """
+    Remove the first H1 from markdown and return (modified_md, h1_title).
+    The H1 is removed since we use the frontmatter title instead.
+    """
+    lines = md.splitlines()
+    result = []
+    h1_title = None
+    skipped_first = False
+
+    for line in lines:
+        if not skipped_first and line.startswith("# "):
+            h1_title = line[2:].strip()
+            skipped_first = True
+            continue
+        result.append(line)
+
+    return "\n".join(result), h1_title
+
+
+IMG_RE = re.compile(r"!\[[^\]]*\]\(([^)]+)\)")
+LINK_RE = re.compile(r"\[([^\]]+)\]\(([^)]+)\)")
+
+
+def rewrite_and_copy_assets(md: str, src_dir: Path, dest_assets_root: Path) -> str:
+    """
+    Copy local image assets referenced in markdown to
+    docs/examples/assets/... and rewrite the links.
+    """
+    dest_assets = dest_assets_root / "assets"
+
+    def repl(m):
+        url = m.group(1).strip()
+        if re.match(r"^(https?:)?//", url):
+            return m.group(0)  # leave remote URLs
+        src_path = (src_dir / url).resolve()
+        if not src_path.exists():
+            return m.group(0)  # leave as-is if not found
+        rel = src_path.relative_to(src_dir)
+        # Create a unique asset path based on source directory name
+        asset_name = src_dir.name.replace("/", "-")
+        dest_path = dest_assets / asset_name / rel
+        dest_path.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(src_path, dest_path)
+        new_rel = f"assets/{asset_name}/{rel.as_posix()}"
+        return m.group(0).replace(url, new_rel)
+
+    return IMG_RE.sub(repl, md)
+
+
+def rewrite_readme_links(
+    md: str,
+    src_dir: Path,
+    examples_dir: Path,
+    parent_index_only: set,
+    current_src_path: str,
+    allowlist_entries: set,
+    current_output_path: str,
+) -> str:
+    """
+    Rewrite links between README.md files to point to the correct .qmd files.
+    """
+
+    def repl(m):
+        text = m.group(1)
+        url = m.group(2).strip()
+
+        # Skip remote URLs and anchor links
+        if re.match(r"^(https?:)?//", url) or url.startswith("#"):
+            return m.group(0)
+
+        # Skip non-markdown files
+        if not url.lower().endswith(".md"):
+            return m.group(0)
+
+        # Resolve the target path
+        try:
+            target_path = (src_dir / url).resolve()
+
+            # Check if target is outside examples_dir
+            try:
+                rel_path = target_path.relative_to(examples_dir)
+            except ValueError:
+                # Target is outside examples_dir, leave as-is
+                return m.group(0)
+
+            parts = list(rel_path.parts)
+
+            # Determine the output path for the target
+            if len(parts) > 0 and parts[-1].lower() in ("readme.md", "readme"):
+                # This is a README link
+                if len(parts) == 1:
+                    # Link to root README -> index.qmd
+                    target_output = "index.qmd"
+                elif len(parts) == 2:
+                    if parts[0] == ".":
+                        # Current directory README
+                        target_output = "index.qmd"
+                    else:
+                        # subdir/README.md
+                        parent_dir = parts[0]
+                        if parent_dir in parent_index_only:
+                            target_output = f"{parent_dir}/index.qmd"
+                        else:
+                            target_output = f"{parent_dir}.qmd"
+                else:
+                    # Deeper nesting: parent/subdir/README.md
+                    # Build the full path like "parent/subdir"
+                    full_path = "/".join(parts[:-1])  # Remove README.md
+                    # Check if this exact path is in allowlist
+                    if full_path in allowlist_entries:
+                        # This is a sub-entry with its own entry -> use .qmd
+                        target_output = f"{full_path}.qmd"
+                    elif parts[0] == ".":
+                        # ./subdir/README.md -> check if subdir has own entry
+                        subdir = parts[1]
+                        if subdir in parent_index_only:
+                            target_output = f"{subdir}/index.qmd"
+                        else:
+                            target_output = f"{subdir}.qmd"
+                    else:
+                        # parent/subdir where parent doesn't have own entry
+                        target_output = f"{full_path}/index.qmd"
+            else:
+                # Regular .md file -> convert to .qmd, keep path structure
+                target_output = "/".join(parts)[:-2] + "qmd"
+
+            # Compute relative path from current output file to target
+            current_parts = current_output_path.split("/")
+            target_parts = target_output.split("/")
+
+            # Special case: if current is a subdir file and target is a single-component file at root
+            # Example: current="magistral/vision", target="magistral.qmd"
+            if len(current_parts) > 1 and len(target_parts) == 1:
+                # Current is in subdir, target is at root level
+                # Go up to root: ../ for each level
+                up_count = len(current_parts) - 1
+                rel_parts = [".."] * up_count + [target_parts[0]]
+                new_url = "/".join(rel_parts)
+            else:
+                # Find common prefix
+                i = 0
+                while (
+                    i < min(len(current_parts) - 1, len(target_parts))
+                    and current_parts[i] == target_parts[i]
+                ):
+                    i += 1
+
+                # Build relative path: go up (../) then down to target
+                up_count = len(current_parts) - 1 - i
+                rel_parts = [".."] * up_count + target_parts[i:]
+
+                if not rel_parts or rel_parts == [".."]:
+                    # Points to same directory or parent
+                    new_url = "/".join(rel_parts) if rel_parts else "."
+                else:
+                    new_url = "/".join(rel_parts)
+
+            return f"[{text}]({new_url})"
+        except (ValueError, IndexError):
+            return m.group(0)
+
+    return LINK_RE.sub(repl, md)
+
+
+def write_qmd(out_path: Path, title: str, body_md: str):
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    fm = f"---\ntitle: {title!r}\nexecute:\n  eval: false\nformat:\n  html:\n    toc: true\n---\n\n"
+    out_path.write_text(fm + body_md, encoding="utf-8")
+
+
+def update_quarto_yml(generated: list[tuple[str, str, str]]):
+    """
+    Update _quarto.yml with the generated example files in the correct order.
+    This keeps the sidebar in sync with the allowlist.
+
+    Model Guides is now nested under "Getting Started" section.
+    Creates nested sections for models with sub-entries (e.g., magistral, ministral3).
+    Parent pages are now flat files (e.g., ministral3.qmd) with sub-pages in subdirs.
+    """
+    quarto_yml = ROOT / "_quarto.yml"
+    if not quarto_yml.exists():
+        print(f"[WARN] {quarto_yml} not found, skipping update", file=sys.stderr)
+        return
+
+    content = quarto_yml.read_text(encoding="utf-8")
+
+    # First pass: find all parents that have sub-entries
+    parents_with_subs = set()
+    for path, _name, _title in generated:
+        if "/" in path:
+            parent = path.split("/")[0]
+            parents_with_subs.add(parent)
+
+    # Build the YAML contents while preserving allowlist order
+    lines = []
+    processed_sections = set()
+
+    for path, _name, title in generated:
+        # Check if this is a parent page that has sub-pages
+        if path in parents_with_subs:
+            # This is a parent page with sub-pages - create a nested section
+            if path not in processed_sections:
+                processed_sections.add(path)
+                section_title = (
+                    title or path.replace("-", " ").replace("_", " ").title()
+                )
+                lines.append(f'                - section: "{section_title}"')
+                lines.append("                  contents:")
+                # Add the parent page first
+                lines.append(f"                    - docs/models/{path}.qmd")
+                # Then add all sub-pages
+                for sub_path, _sub_name, _sub_title in generated:
+                    if "/" in sub_path and sub_path.split("/")[0] == path:
+                        lines.append(
+                            f"                    - docs/models/{sub_path}.qmd"
+                        )
+        elif "/" not in path:
+            # This is a flat item with no sub-pages
+            # Skip if it was already included as part of a parent section
+            if path not in processed_sections:
+                lines.append(f"                - docs/models/{path}.qmd")
+
+    yaml_content = "\n".join(lines) + "\n"
+
+    # Pattern to match only the Model Guides contents, stopping at the next item
+    # in Getting Started (lines starting with 12 spaces: same level as the section)
+    pattern = r'(            - section: "Model Guides"\n              contents:)([^\n]*|.*?)(?=\n            - |\n        - section:|\n\nformat:)'
+
+    def replacement(match):
+        prefix = match.group(1)
+        return prefix + "\n" + yaml_content
+
+    new_content = re.sub(pattern, replacement, content, flags=re.DOTALL)
+
+    if new_content != content:
+        quarto_yml.write_text(new_content, encoding="utf-8")
+        print(f"Updated {quarto_yml}")
+    else:
+        print(f"No changes needed for {quarto_yml}")
+
+
+def main():
+    allow = read_allowlist()
+    if not EXAMPLES_DIR.exists():
+        print(f"[WARN] {EXAMPLES_DIR} not found", file=sys.stderr)
+        return
+
+    (OUTPUT_DIR / "assets").mkdir(parents=True, exist_ok=True)
+
+    # First pass: identify which parents have their own entry vs only sub-entries
+    parent_entries = set()  # Parents that have their own entry
+    parent_with_subs = set()  # Parents that have sub-entries
+    allowlist_entries = set()  # All entries in allowlist
+
+    for item in allow:
+        if isinstance(item, str):
+            name = item
+        else:
+            name = item.get("name")
+
+        allowlist_entries.add(name)
+
+        if "/" in name:
+            parent = name.split("/")[0]
+            parent_with_subs.add(parent)
+        else:
+            parent_entries.add(name)
+
+    # Parents with subs that DON'T have their own entry -> use index.qmd
+    parent_index_only = parent_with_subs - parent_entries
+
+    generated = []
+    seen_dirs = set()  # Track which parent directories we've created index for
+
+    for item in allow:
+        if isinstance(item, str):
+            name = item
+            title = None
+        else:
+            name = item.get("name")
+            title = item.get("title")
+
+        if not name:
+            print(f"[WARN] Skipping item without name: {item}", file=sys.stderr)
+            continue
+
+        src_dir = EXAMPLES_DIR / name
+        if not src_dir.exists() or not src_dir.is_dir():
+            print(f"[WARN] Skipping {name} (not a directory)", file=sys.stderr)
+            continue
+
+        readme = find_readme(src_dir)
+        if not readme:
+            print(f"[WARN] Skipping {name} (no README.md)", file=sys.stderr)
+            continue
+
+        md = readme.read_text(encoding="utf-8")
+
+        # Determine output path first (needed for link rewriting)
+        parts = name.split("/")
+        if len(parts) == 1:
+            # Simple case: no subdirectory
+            out_path = OUTPUT_DIR / f"{parts[0]}.qmd"
+            sidebar_path = parts[0]
+        else:
+            # Has subdirectory: e.g., magistral/think
+            parent = parts[0]
+            child = "-".join(parts[1:])  # handle nested subdirs
+            out_path = OUTPUT_DIR / parent / f"{child}.qmd"
+            sidebar_path = f"{parent}/{child}"
+
+        # Remove the first H1 (we use frontmatter title instead)
+        md, _ = remove_first_h1(md)
+        # Rewrite links between README files
+        md = rewrite_readme_links(
+            md,
+            src_dir,
+            EXAMPLES_DIR,
+            parent_index_only,
+            name,
+            allowlist_entries,
+            sidebar_path,
+        )
+        md = rewrite_and_copy_assets(md, src_dir, OUTPUT_DIR)
+
+        # Handle parent page generation for sub-entries
+        if len(parts) > 1:
+            # Has subdirectory: e.g., magistral/think
+            parent = parts[0]
+
+            # Create parent.qmd if not already done and parent doesn't have own entry
+            if parent not in seen_dirs and parent in parent_index_only:
+                parent_readme = find_readme(EXAMPLES_DIR / parent)
+                if parent_readme:
+                    parent_md = parent_readme.read_text(encoding="utf-8")
+                    parent_md, _ = remove_first_h1(parent_md)
+                    parent_md = rewrite_readme_links(
+                        parent_md,
+                        EXAMPLES_DIR / parent,
+                        EXAMPLES_DIR,
+                        parent_index_only,
+                        parent,
+                        allowlist_entries,
+                        parent,
+                    )
+                    parent_md = rewrite_and_copy_assets(
+                        parent_md, EXAMPLES_DIR / parent, OUTPUT_DIR
+                    )
+                    parent_title = parent.replace("-", " ").replace("_", " ").title()
+                    write_qmd(OUTPUT_DIR / f"{parent}.qmd", parent_title, parent_md)
+                    generated.append((parent, parent, parent_title))
+                    seen_dirs.add(parent)
+
+        if not title:
+            title = name.replace("/", " ").replace("-", " ").title()
+
+        write_qmd(out_path, title, md)
+        generated.append((sidebar_path, name, title))
+
+    # Index page - preserve allowlist order
+    if generated:
+        listing = "\n".join(
+            [f"- [{title}]({path}.qmd)" for path, name, title in generated]
+        )
+        index_md = (
+            "# Model Guides\n\nBelow are the curated examples for training various model architectures:\n\n"
+            + listing
+            + "\n"
+        )
+        index_fm = (
+            "---\nexecute:\n  eval: false\nformat:\n  html:\n    toc: true\n---\n\n"
+        )
+        (OUTPUT_DIR / "index.qmd").write_text(index_fm + index_md, encoding="utf-8")
+
+        # Auto-update _quarto.yml to keep sidebar in sync
+        update_quarto_yml(generated)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/docs/sequence_parallelism.qmd b/docs/sequence_parallelism.qmd
new file mode 100644
index 0000000000..f03daf0d89
--- /dev/null
+++ b/docs/sequence_parallelism.qmd
@@ -0,0 +1,128 @@
+---
+title: Sequence Parallelism
+description: Train with long sequences split across multiple GPUs.
+---
+
+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
+context_parallel_size: 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
+# Optional; one of "varlen_llama3" or "batch_ring". Defaults to
+# "varlen_llama3" when `sample_packing: true`, and "batch_ring" otherwise.
+ring_attn_func:
+```
+
+The `context_parallel_size` 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
+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 (`attn_implementation: flash_attention_2` in config YAML)
+- May have a small performance overhead due to communication between GPUs
+
+### Mamba/SSM Hybrid Models
+
+Context parallelism is supported for hybrid models that combine attention and Mamba2 SSM
+layers. These models require special handling because:
+
+- **Attention layers** work correctly via ring flash attention (same as pure-attention models).
+- **SSM (Mamba2) layers** are recurrent and need cross-rank hidden-state propagation.
+
+Axolotl handles both aspects:
+
+1. **Sample packing boundaries** (`seq_idx`): When multiple sequences are packed into one
+   row, the SSM kernels need `seq_idx` to reset state at boundaries. Under CP, chunks may
+   start mid-sample, so `seq_idx` is derived from `position_ids` using a CP-safe cumsum
+   normalization (see `mamba_utils.get_seq_idx`).
+
+2. **Cross-rank SSM state passing**: After each SSM scan, the final hidden state is sent to
+   the next rank via P2P communication (`ring_shift_ssm_state`), and an additive correction
+   is applied to account for the missing initial state (`mamba2_cp_correction`). This uses
+   the linearity property of SSMs to avoid a second forward pass.
+
+#### Supported Architectures
+
+| Model family | `model_config_type` | Architecture notes |
+|---|---|---|
+| Nemotron-H | `nemotron_h` | Mamba2 / Attention / MoE hybrid; block type selected per layer |
+| Falcon-H1 | `falcon_h1` | Mamba2 and Attention run **in parallel** in every layer |
+| Granite MoE Hybrid | `granitemoehybrid` | Mamba2 / Attention / MoE hybrid |
+
+## Example
+
+```yaml
+base_model: meta-llama/Llama-3-8B-Instruct
+sequence_len: 8192
+
+...
+
+context_parallel_size: 4  # Split each sequence into 4 parts, one per GPU
+# Optional; strides across the key dimension. Larger values use more memory but should make training faster.
+heads_k_stride: 1
+# Optional; one of "varlen_llama3" or "batch_ring". Defaults to
+# "varlen_llama3" when `sample_packing: true`, and "batch_ring" otherwise.
+ring_attn_func:
+
+...
+```
+
+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 `context_parallel_size`. This happens because:
+
+- Each group of `context_parallel_size` 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 `context_parallel_size=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
diff --git a/docs/streaming.qmd b/docs/streaming.qmd
new file mode 100644
index 0000000000..2a233a4fc7
--- /dev/null
+++ b/docs/streaming.qmd
@@ -0,0 +1,120 @@
+---
+title: Streaming Datasets
+description: How to use streaming mode for large-scale datasets and memory-efficient training
+order: 10
+---
+
+Streaming enables memory-efficient training with large datasets by loading data
+incrementally rather than loading the entire dataset into memory at once.
+
+Use streaming when:
+
+- Your dataset is too large to fit in memory (e.g. when you're doing pretraining with massive text corpora)
+- You want to start training immediately without preprocessing the entire dataset
+
+Streaming works with both remote and locally stored datasets!
+
+::: {.callout-note}
+Streaming currently only supports a single dataset. Multi-dataset support will be added soon.
+:::
+
+
+## Configuration
+
+### Basic Streaming
+
+Enable streaming mode by setting the `streaming` flag:
+
+```yaml
+streaming: true
+```
+
+### Pretraining with Streaming
+
+For pretraining tasks, streaming is automatically enabled when using `pretraining_dataset`:
+
+```yaml
+pretraining_dataset:
+  - path: HuggingFaceFW/fineweb-edu
+    type: pretrain
+    text_column: text
+    split: train
+
+# Optionally, enable sample packing
+streaming_multipack_buffer_size: 10000
+sample_packing: true
+```
+
+### SFT with Streaming
+
+For supervised fine-tuning with streaming:
+
+```yaml
+streaming: true
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+    split: train
+
+# Optionally, enable sample packing
+streaming_multipack_buffer_size: 10000
+sample_packing: true
+```
+
+## Configuration Options
+
+### `streaming_multipack_buffer_size`
+
+Controls the buffer size for multipack streaming (default: 10,000). This determines how
+many samples are buffered before packing. Larger buffers can improve packing efficiency
+but use more memory.
+
+### `shuffle_merged_datasets`
+
+When enabled, shuffles the streaming dataset using the buffer. This requires additional
+memory for the shuffle buffer.
+
+## Sample Packing with Streaming
+
+Sample packing is supported for streaming datasets. When enabled, multiple samples are
+packed into a single sequence to maximize GPU utilization:
+
+```yaml
+sample_packing: true
+streaming_multipack_buffer_size: 10000
+
+# For SFT: attention is automatically isolated between packed samples
+# For pretraining: control with pretrain_multipack_attn
+pretrain_multipack_attn: true  # prevent cross-attention between packed samples
+```
+
+For more information, see our [documentation](multipack.qmd) on multipacking.
+
+## Important Considerations
+
+### Memory Usage
+
+While streaming reduces memory usage compared to loading entire datasets, you still need
+to consider:
+
+- You can control the memory usage by adjusting `streaming_multipack_buffer_size`
+- Sample packing requires buffering multiple samples
+- Shuffling requires additional memory for the shuffle buffer
+
+### Performance
+
+- Streaming may have slightly higher latency compared to preprocessed datasets, as samples are processed on-the-fly
+- Network speed and disk read speed are important when streaming from remote sources or a local dataset, respectively
+- Consider using `axolotl preprocess` for smaller or more frequently used datasets
+
+### Evaluation Datasets
+
+Evaluation datasets are not streamed to ensure consistent evaluation metrics. They're
+loaded normally even when training uses streaming.
+
+## Examples
+
+See the `examples/streaming/` directory for complete configuration examples:
+
+- `pretrain.yaml`: Pretraining with streaming dataset
+- `sft.yaml`: Supervised fine-tuning with streaming
diff --git a/docs/support-matrix.qmd b/docs/support-matrix.qmd
new file mode 100644
index 0000000000..f2db689bd8
--- /dev/null
+++ b/docs/support-matrix.qmd
@@ -0,0 +1,315 @@
+---
+title: Support Matrix
+description: What Axolotl supports, how features couple together, and what is not supported yet.
+---
+
+Axolotl is config-driven: every capability below is a YAML key, not a code change. **Any HuggingFace causal/seq2seq LM trains out of the box** through generic Transformers support. The tables list *first-class* features and the exact key that enables each.
+
+New here? You mostly need three tables: [Training methods](#training-methods) (what objective you are optimizing), [Fine-tuning strategies](#fine-tuning-strategies) (how much of the model updates), and [Model architectures](#model-architectures) (whether your model is covered). Everything else is reference for when you hit a specific wall.
+
+**I want to...**
+
+| Goal | Start here |
+|---|---|
+| Check if my model is supported | Any HuggingFace causal/seq2seq LM trains as-is; [Model architectures](#model-architectures) lists the ones with extra acceleration |
+| Fit a large model on one GPU | [QLoRA](#fine-tuning-strategies), plus [FSDP + QLoRA](#distributed-parallelism) for the biggest |
+| Choose a training objective (SFT, DPO, GRPO, ...) | [Training methods](#training-methods) |
+| Train faster or use less memory | [Performance integrations](#performance-integrations-custom-kernels); Liger and Cut Cross Entropy are the two most people enable first |
+| Know whether two features combine | [Compatibility rules](#compatibility-rules) |
+| See what is experimental, deprecated, or planned | [Maturity & roadmap](#maturity-roadmap) |
+
+The catalog below has one small table per axis, and most features are independent. The couplings that actually bite are collected in [Compatibility rules](#compatibility-rules), and maturity plus roadmap are rolled up in [Maturity & roadmap](#maturity-roadmap). Feature-specific limits (for example "not on Turing" or "LoRA unsupported") are marked inline on the row they affect.
+
+::: {.callout-tip}
+## Want faster training?
+Most people start with two, each a single plugin to turn on: [Liger](#fused-training-kernels) kernels and [Cut Cross Entropy](#fused-training-kernels). The full set of speed and memory work (ScatterMoE/SonicMoE, Expert Parallel, fused attention) lives in [Performance integrations & custom kernels](#performance-integrations-custom-kernels).
+:::
+
+::: {.callout-important}
+## Read the status markers
+"Supported" moves fast, especially for kernels and quantization. Rows are annotated:
+🟢 stable (default, unmarked) · 🟡 experimental/beta · 🚧 WIP / feature-branch only · 🔵 planned / not yet · 🔴 deprecated · ⚪ out-of-scope / won't-fix. Limits live on the row they affect; [Maturity & roadmap](#maturity-roadmap) is only the cross-cutting roll-up. This page reflects the **0.17.0 development line**; 🚧 items may not be in a tagged release yet.
+:::
+
+## Training methods
+
+*The "what objective are you optimizing" axis.*
+
+| Method | Config | Ref model | Data shape | Notes |
+|---|---|---|---|---|
+| Supervised FT (SFT) | *(default)* | n/a | prompt to response | `train_on_inputs`, chat/alpaca/completion |
+| Continued pretraining | `pretraining_dataset:` | n/a | raw text | streaming; needs `max_steps` |
+| Reward model (ORM) | `reward_model: true` | n/a | chosen/rejected | Bradley-Terry, sequence classification |
+| Process reward model (PRM) | `process_reward_model: true` | n/a | stepwise | token classification, `stepwise_supervised` |
+| DPO | `rl: dpo` | ✔ | paired | `dpo_loss_type`, `rl_beta`, liger/padding-free variants |
+| IPO | `rl: dpo` + `dpo_loss_type: [ipo]` | ✔ | paired | `rl: ipo` still works (🔴 will deprecate) |
+| ORPO | `rl: orpo` | No | paired | single-stage, ~half the VRAM of DPO |
+| SimPO / CPO | `rl: simpo` | No | paired | reference-free (TRL CPOTrainer) |
+| KTO | `rl: kto` | ✔ | unpaired + binary label | needs `remove_unused_columns: false` |
+| GDPO | `rl: gdpo` | ✔ | paired, multi-objective | multi-objective DPO (via GRPO strategy) |
+| GRPO | `rl: grpo` | ✔ | prompts (online) | vLLM generation, custom `reward_funcs`, async/replay |
+| EBFT | `rl: ebft` | ✔ | QA or raw text | energy/feature-matching rewards; needs `ebft:` block |
+
+Most runs start with **SFT** (the default). Reach for the preference methods (DPO and its variants) once you have chosen/rejected pairs, and GRPO when you have a verifiable reward. The rest are specialized.
+
+## Fine-tuning strategies
+
+*The "how much of the model updates" axis.*
+
+| Strategy | Config | Notes |
+|---|---|---|
+| Full fine-tune | `adapter:` *(omit)* | all params trainable; requires unquantized weights (see [Precision x trainable params](#precision-x-trainable-params)) |
+| LoRA | `adapter: lora` | `lora_r` / `lora_alpha` / `lora_target_modules` / `lora_target_linear` |
+| QLoRA | `adapter: qlora` | requires `load_in_4bit: true` |
+| DoRA | `+ peft_use_dora: true` | weight-decomposed |
+| rsLoRA | `+ peft_use_rslora: true` | rank-stabilized |
+| LoRA+ | `+ loraplus_lr_ratio:` | split A/B learning rates |
+| LoftQ | `+ peft.loftq_config:` | quant-aware init |
+| ReLoRA | `relora: true` + `jagged_restart_steps:` | periodic merge/restart (no FSDP/DeepSpeed) |
+| Train embeddings/head | `lora_modules_to_save: [...]` | needed when adding tokens |
+| Spectrum | plugin `spectrum` | SNR-selected frozen params |
+| LISA | `lisa_n_layers` + `lisa_step_interval` | rotating layer unfreeze |
+| MoRA / ReMoRA | plugin `mora` | high-rank LoRA alternative |
+
+## Quantization & precision
+
+Quantization shows up at three distinct points in the lifecycle; they are easy to confuse.
+
+| Phase | Purpose | Config | Trains? |
+|---|---|---|---|
+| **Load-time (frozen base)** | shrink the base so an adapter fits | `load_in_4bit`/`load_in_8bit`, `gptq`, AWQ, `model_quantization_config: FineGrainedFP8Config`/`Mxfp4Config`, NVFP4-modelopt (MoE) | adapter only |
+| **Train-time (compute)** | faster/leaner training, high-precision master weights kept | `fp8: true` (torchao float8, `fp8_enable_fsdp_float8_all_gather`), `qat:` (fake-quant during training) | full fine-tune ✅ |
+| **Post-training (export)** | quantize the finished model | `quantization:` PTQ via `axolotl quantize` (int4/int8/fp8/nvfp4/mxfp4) | n/a (after training) |
+
+::: {.callout-note}
+Two things named "FP8" are different: `fp8: true` is **mixed-precision compute** (master weights stay bf16/fp32, so full fine-tune works), whereas `model_quantization_config: FineGrainedFP8Config` loads a **frozen fp8 base** (adapter only). Likewise **QAT** (`qat:`) happens *during* training and is full-model; **PTQ** (`quantization:`) happens *after* training via `axolotl quantize`. QAT is mutually exclusive with any adapter.
+:::
+
+## Distributed & parallelism
+
+| Strategy | Config | Composes with |
+|---|---|---|
+| DDP | *(default multi-GPU)* | n/a |
+| DeepSpeed ZeRO 1/2/3 (+CPU offload) | `deepspeed: deepspeed_configs/zeroN*.json` | TP |
+| FSDP1 🔴 | `fsdp_version: 1` | *(deprecated, use FSDP2)* |
+| FSDP2 | `fsdp_version: 2` + `fsdp_config:` | TP, CP, EP |
+| FSDP + QLoRA | `adapter: qlora` + FSDP2 | 70B on consumer GPUs |
+| Tensor Parallel (TP) 🟡 | `tensor_parallel_size:` | FSDP2 |
+| Context/Sequence Parallel (CP) | `context_parallel_size:` | FSDP2, ring-flash-attn |
+| Expert Parallel (EP) 🟡 | plugin `expert_parallel` + `expert_parallel_size:` | FSDP2 (DeepEP backend) |
+| N-D (HSDP, FSDP+TP+CP, FSDP+EP) 🟡 | combine `dp_replicate`/`dp_shard`/`tp`/`cp`/`ep` | see [nd_parallelism](nd_parallelism.qmd) |
+
+Blocked combinations: **EP x TP/CP**, **DDP x TP/CP** (use FSDP2 instead), **DeepSpeed x FSDP**. See [Incompatible combinations](#incompatible).
+
+## Performance integrations & custom kernels
+
+The differentiators. All are opt-in; most need a specific GPU generation.
+
+### Fused training kernels
+
+| Kernel | Config | Covers |
+|---|---|---|
+| Liger | plugin `liger` + `liger_rope`, `liger_rms_norm(_gated)`, `liger_swiglu`/`liger_glu_activation`, `liger_cross_entropy`, `liger_fused_linear_cross_entropy` | RoPE, RMSNorm, SwiGLU, CE, fused-linear-CE |
+| Cut Cross Entropy (CCE) | plugin `cut_cross_entropy` + `cut_cross_entropy: true` | memory-lean logit-free CE (Apple fork) |
+| Chunked CE | `chunked_cross_entropy: true` | chunked loss for long sequences |
+| DenseMixer | plugin `densemixer` | fused MoE forward (OLMoE, Qwen2/3-MoE) |
+
+Only **one** cross-entropy optimization may be enabled at a time (CCE / Liger CE / chunked).
+
+### MoE expert kernels
+
+| Kernel | Config | GPU | Scope |
+|---|---|---|---|
+| ScatterMoE | `use_kernels: true` + `use_scattermoe: true` | any CUDA (Triton) | LoRA on experts (fused into grouped GEMM); **composes with EP** |
+| SonicMoE 🚧 | `use_kernels: true` + `use_sonicmoe: true` | Hopper/Blackwell, CUDA 12.9+ | LoRA; **NVFP4 path is WIP**; ⚪ **blocked under EP** (use ScatterMoE) |
+| DSv4 fused kernels | `use_dsv4_kernels: true` | SM90+ | DeepSeek-V4 attention/RoPE/MLP |
+| GLM-DSA kernels | `use_glm_dsa_kernels: true` | SM90+ | GLM-4.7/5.2 sparse-MLA (DSA) attention |
+| Grouped-GEMM backend | `moe_grouped_backend: auto\|marlin\|cutlass\|deepgemm` | capability-selected | fp4 grouped experts |
+
+**With Expert Parallel:** ScatterMoE composes with EP (registered as `deep_ep_scattermoe`). **SonicMoE + EP is blocked at config validation** (`use_sonicmoe: true` with `expert_parallel_size > 1` raises), use `use_scattermoe` under EP.
+
+::: {.callout-warning}
+**NVFP4 MoE training is narrow today.** It covers **DeepSeek-V4** and **GLM-4.7/5.2** (`glm_moe_dsa`) only, is **LoRA-only** (frozen fp4 experts + fused low-rank; no full fine-tune of the fp4 base), and the **SonicMoE NVFP4 path is 🚧 WIP** (ScatterMoE is the working path). See [Precision x trainable params](#precision-x-trainable-params).
+:::
+
+### Attention backends
+
+Canonical key `attn_implementation:` (legacy boolean flags in parentheses are 🔴 deprecated).
+
+| Backend | Value / flag | Notes |
+|---|---|---|
+| Flash Attention 2/3 | `flash_attention_2` / `flash_attention_3` (`flash_attention: true`) | FA2 auto-upgrades to FA3 on SM90+; ⚪ Ampere or newer only (not Turing) |
+| Flash Attention 4 🟡 | auto-upgrade when available | pre-release; 🔵 backward not in pip on Hopper (build from source) |
+| SDPA | `sdpa` (`sdp_attention: true`) | safe default; works everywhere |
+| SDPA varlen | `sdpa_varlen: true` | packing without a mask tensor, PyTorch >= 2.10 |
+| FlexAttention | `flex_attention` (`flex_attention: true`) | PyTorch >= 2.6; enables `scaling_softmax` |
+| xFormers | (`xformers_attention: true`) | varlen packing |
+| SageAttention | (`sage_attention: true`) | block-quantized |
+| FP8 attention 🟡 | `attn_implementation: fp8` | SM90+, PyTorch >= 2.11 |
+| Large head-dim (>256) | `large_head_attention: auto\|sdpa\|triton_flash` | Triton kernel up to 512 |
+
+### Throughput / memory
+
+| Feature | Config | Notes |
+|---|---|---|
+| Sample packing (multipack) | `sample_packing: true` | block-diagonal attention + position reset; needs a varlen backend; ⚪ not with RLHF or multimodal |
+| Fused LoRA-MLP/QKV/O kernels | `lora_mlp_kernel` / `lora_qkv_kernel` / `lora_o_kernel` | SFT + FSDP2 only; ⚪ not RLHF, not FSDP1, not `trust_remote_code` |
+| Tiled MLP | `tiled_mlp: true` | shards MLP to cut memory |
+| Ring attention | via `context_parallel_size` | long-context CP |
+| Activation offloading | `activation_offloading:` | requires gradient checkpointing; ⚪ CUDA only (no CPU-only training) |
+
+## Model architectures
+
+Generic HuggingFace support is universal. This table lists **added acceleration/patches only** (~69 example configs under `examples/`).
+
+| Family | Special support |
+|---|---|
+| Llama 2/3/4 | flash-attn hijack, SwiGLU, CCE, Liger; Llama-4 linearized experts |
+| Mistral / Mixtral / Ministral / Magistral | flash-attn hijack, CCE; Mixtral ZeRO-3 MoE patch |
+| Qwen 2/2.5/3/3.5 (+MoE, +VL) | fused attention kernels, multipack, gated RMSNorm, VL flash |
+| Gemma 2/3/4 (+unified VL) | hybrid sliding/global attention mask, fused attn, expert quant |
+| DeepSeek V2/V3/V4 | DSv4 fused attn/RoPE/MLP kernels, NVFP4 grouped experts (LoRA) |
+| GLM 4.x / 4.7 / 5.2 (+MoE-DSA) | DSA sparse-attention kernels, NVFP4 grouped experts (LoRA) |
+| Hybrid SSM (Mamba, Nemotron-H, Falcon-H1, GraniteMoE-Hybrid) | packing + CP for Mamba2 layers |
+| OLMo 2/3, Cohere, Phi, Hunyuan, Jamba, Kimi-Linear, Apertus, SEED-OSS | packing / tokenizer / activation patches |
+| Multimodal 🟡 (Qwen-VL, Pixtral, Llama-Vision, LLaVA, InternVL, SmolVLM2, Voxtral, LFM2-VL) | processor + VL attention support (no full feature parity) |
+| BitNet (1.58-bit) ⚪ | full fine-tune only, **LoRA not supported** (see [1_58bit_finetuning](1_58bit_finetuning.qmd)) |
+
+## Dataset formats
+
+| Format | `type:` | Notes |
+|---|---|---|
+| Chat template | `chat_template` | Jinja; per-turn/role/EOS loss masking, tools, reasoning traces |
+| Alpaca & variants | `alpaca`, `alpaca_chat`, ... | legacy instruction |
+| Input/output (template-free) | `input_output` | explicit `segments` masking |
+| Completion / raw | `completion` | pretraining |
+| Stepwise supervised | `stepwise_supervised` | PRM data |
+| Preference | `dpo_datasets:` / `kto_datasets:` | chosen/rejected or completion+label |
+| Multimodal 🟡 | `chat_template` + image/audio | resizing, role boundaries |
+
+## Optimizers & schedulers
+
+**Optimizers:** all HuggingFace / `bitsandbytes` optimizers, plus AdamW (torch-fused, optimi), TorchAO 4-bit/8-bit/FP8 AdamW, ADOPT, CAME, Muon, Dion, SinkGD, Flash AdamW/Adam/SGD/SGDW/Lion, Q-GaLore. (Muon / Flash / Q-GaLore require FSDP2, not DeepSpeed.)
+
+**Schedulers:** cosine (+min-lr, +constant-ratio, +quadratic warmup), REX, one-cycle, linear warmup, jagged-restart (ReLoRA).
+
+**Also:** gradient checkpointing (+CPU/disk offload), gradient accumulation, NEFTune, LR groups / embedding LR, loss watchdog, early stopping, `torch_compile`.
+
+## Method plug-ins & extensions
+
+| Plugin | Purpose |
+|---|---|
+| `kd` | knowledge distillation (offline + online vLLM/SGLang teacher) |
+| `diffusion` | diffusion-LM training |
+| `hatchery` | remote training (Tinker/Hatchery backends) |
+| `nemo_gym` | RL environments / verifiable rewards |
+| `llm_compressor` | sparse fine-tuning |
+| `grokfast` | grokking-accelerated optimizer |
+| `lm_eval` | post-train eval harness |
+
+## Experiment tracking
+
+Weights & Biases, MLflow, Comet, Trackio, SwanLab, OpenTelemetry/Prometheus, Ray (launcher), Gradio (inference UI).
+
+## Compatibility rules
+
+Most features compose. This section lists the exceptions that actually bite. The full, machine-checked set of rules lives in the config validators (`src/axolotl/utils/schemas/`); the tables below are the high-signal subset.
+
+### Precision x trainable params
+
+The dense corner. This is why, for example, DeepSeek-V4 (which ships in NVFP4) can be LoRA-trained but not full fine-tuned.
+
+| Weights loaded as | Full fine-tune | LoRA / QLoRA | Mechanism |
+|---|:---:|:---:|---|
+| bf16 / fp16 (unquantized) | ✅ | ✅ (LoRA) | standard; also supports `fp8:` compute and `qat:` |
+| bnb NF4 (`load_in_4bit`) | ❌ | ✅ → this *is* **QLoRA** | frozen 4-bit base + bf16 adapter |
+| bnb int8 (`load_in_8bit`) | ❌ | ✅ (LoRA) | frozen 8-bit base |
+| GPTQ / AWQ (prequantized) | ❌ | ✅ (LoRA) | frozen; adapter can't be merged back |
+| FP8 frozen base (`FineGrainedFP8Config`) | ❌ | ✅ (LoRA) | distinct from `fp8:` compute |
+| MXFP4 (`Mxfp4Config`) | ❌ | ✅ (LoRA) | frozen |
+| NVFP4 MoE: **DeepSeek-V4 & GLM-4.7/5.2 only** | ❌ *(no FFT-fp4 path)* | ✅ experts-only LoRA via `use_scattermoe` (`use_sonicmoe` 🚧 WIP) | frozen fp4 experts + fused grouped-GEMM LoRA |
+| **GGUF / llama.cpp K-quants** (Q4_K, Q6_K, ...) | ⚪ ❌ | ⚪ ❌ | **not trainable**, inference / merge-export format only |
+
+**Rule of thumb:** *quantized weights ⇒ frozen base ⇒ adapter-only. Full fine-tune needs unquantized weights.* The exceptions that *look* like quantized full fine-tune are actually train-time modes with high-precision master weights: `fp8: true` (mixed-precision compute) and `qat:` (fake-quant during training). GGUF and K-quant (QX) formats are deployment artifacts and cannot be trained at all.
+
+### Requires (feature -> hard dependency)
+
+| Feature | Requires |
+|---|---|
+| QLoRA | `load_in_4bit: true` |
+| Expert Parallel | FSDP2 + DeepEP + Ampere/Hopper (NVLink) |
+| Context/Sequence parallel | flash attention (FA2/FA3) |
+| SonicMoE | Hopper/Blackwell + CUDA 12.9 (else auto-falls back to ScatterMoE) |
+| FP8 attention | SM90+ and PyTorch >= 2.11 |
+| `quantize_moe_experts` | LoRA/QLoRA + 4/8-bit + CUDA (use `lora_target_parameters`, not `lora_target_linear`) |
+| ReLoRA | `jagged_restart_steps` (and *not* FSDP/DeepSpeed/one_cycle) |
+| Sample packing | varlen backend (FA2/3, flex, xformers, sage) |
+| Muon / Flash / Q-GaLore optimizers | FSDP2 (not DeepSpeed) |
+| EBFT (`rl: ebft`) | `ebft:` config block |
+
+### Incompatible {#incompatible}
+
+| ✕ | ✕ | Note |
+|---|---|---|
+| DeepSpeed | FSDP | pick one |
+| EP | TP / CP | raises `NotImplementedError` |
+| DDP | TP / CP | use FSDP2 instead |
+| EP | SonicMoE | blocked at config validation; use ScatterMoE under EP (ScatterMoE + EP *is* supported) |
+| QAT | any adapter, or 4/8-bit | full-model only |
+| Sample packing | any `rl:`, multimodal | RLHF & MM packing not supported |
+| Cut Cross Entropy | Liger CE / chunked CE | one CE optimization only |
+| `use_dsv4_kernels` | attention-level LoRA | experts-only (fused indexer is gradientless) |
+| FSDP2 | 4/8-bit + DPO/KTO/ORPO/IPO | use DeepSpeed or FSDP1 |
+| LoRA kernels | FSDP1 / `trust_remote_code` / RLHF | SFT + FSDP2 only |
+| `batch_flattening` | `sample_packing` | choose one |
+| FP8 + DDP | `torch.compile` | ⚪ known-broken; drop compile or use FSDP2 |
+
+## Maturity & roadmap
+
+At-a-glance roll-up. Per-feature limits are annotated *inline* on the row they affect (with the same 🟡/🚧/🔵/🔴/⚪ markers); this section is only the cross-cutting view and forward roadmap.
+
+### 🟡 Experimental / beta
+Works, but flagged unstable:
+
+- FP8 training
+- Flash Attention 4
+- Multimodal / VLM ("limited, no full feature parity")
+- N-D parallelism and Tensor Parallel
+- Expert Parallel / DeepEP
+
+### 🚧 WIP / feature-branch only
+Not in a tagged release yet:
+
+- **SonicMoE + NVFP4** grouped-experts training (ScatterMoE is the working path)
+
+### 🔴 Deprecated (still works)
+
+- `fsdp_version: 1` and bare `fsdp:`: use `fsdp_config`
+- Legacy attention booleans: use `attn_implementation`
+- `rl: ipo` direct: use `dpo_loss_type: [ipo]`
+- `noisy_embedding_alpha`, `dpo_beta`, `evaluation_strategy`
+- Removed entirely: `s2_attention`, `flash_attn_rms_norm`
+
+### 🔵 Not supported yet / planned
+
+- **Full fine-tune of NVFP4** (fp4 base FFT, "a separate, larger effort")
+- **SonicMoE + Expert Parallel** (ScatterMoE + EP works today; a partial sm_120 fallback exists but is gated by the config validator)
+- Multimodal + sample packing
+- LoRA kernels + RLHF
+- EP + TP/CP
+- Multiple DPO loss types (RPO)
+- Sample packing across multiple streaming datasets
+- FA4 backward on Hopper via pip
+
+### ⚪ Out-of-scope / won't-fix
+
+- **Ascend NPU** feature parity: attention, optimizers, and quantization are all unsupported there (cross-cutting, no single-row home)
+
+Everything else out-of-scope is marked on its own row above:
+
+- GGUF / K-quant training, see [Precision x trainable params](#precision-x-trainable-params)
+- FA2 / FA3 pre-Ampere (Turing), see [Attention backends](#attention-backends)
+- BitNet + LoRA, see [Model architectures](#model-architectures)
+- CPU-only + activation offloading, see [Throughput / memory](#throughput-memory)
+- FP8 + DDP + `torch.compile`, see [Incompatible combinations](#incompatible)
+- SonicMoE under Expert Parallel, see [MoE expert kernels](#moe-expert-kernels)
diff --git a/docs/telemetry.qmd b/docs/telemetry.qmd
new file mode 100644
index 0000000000..62d7c9bbc0
--- /dev/null
+++ b/docs/telemetry.qmd
@@ -0,0 +1,61 @@
+---
+title: Telemetry
+description: A description of the telemetry implementation in Axolotl.
+---
+
+# Telemetry in Axolotl
+
+Axolotl implements anonymous telemetry to help maintainers understand how the library
+is used and where users encounter issues. This data helps prioritize features, optimize
+performance, and fix bugs.
+
+## Data Collection
+
+We collect:
+
+- System info: OS, Python version, Axolotl version, PyTorch version, Transformers
+version, etc.
+- Hardware info: CPU count, memory, GPU count and models
+- Runtime metrics: Training progress, memory usage, timing information
+- Usage patterns: Models (from a whitelist) and configurations used
+- Error tracking: Stack traces and error messages (sanitized to remove personal
+information)
+
+Personally identifiable information (PII) is not collected.
+
+## Implementation
+
+Telemetry is implemented using PostHog and consists of:
+
+- `axolotl.telemetry.TelemetryManager`: A singleton class that initializes the
+telemetry system and provides methods for tracking events.
+- `axolotl.telemetry.errors.send_errors`: A decorator that captures exceptions and
+sends sanitized stack traces.
+- `axolotl.telemetry.runtime_metrics.RuntimeMetricsTracker`: A class that tracks
+runtime metrics during training.
+- `axolotl.telemetry.callbacks.TelemetryCallback`: A Trainer callback that sends
+runtime metrics telemetry.
+
+The telemetry system will block training startup for 10 seconds to ensure users are
+aware of data collection, unless telemetry is explicitly enabled or disabled.
+
+## Opt-Out Mechanism
+
+Telemetry is **enabled by default** on an opt-out basis. To disable it, set
+`AXOLOTL_DO_NOT_TRACK=1` or `DO_NOT_TRACK=1`.
+
+A warning message will be logged on start to clearly inform users about telemetry.
+We will remove this after some period.
+
+To hide the warning message about telemetry that is displayed on train, etc. startup,
+explicitly set: `AXOLOTL_DO_NOT_TRACK=0` (enable telemetry) or `AXOLOTL_DO_NOT_TRACK=1`
+(explicitly disable telemetry).
+
+## Privacy
+
+- All path-like config information is automatically redacted from telemetry data
+- Model information is only collected for whitelisted organizations
+    - See `axolotl/telemetry/whitelist.yaml` for the set of whitelisted organizations
+- Each run generates a unique anonymous ID
+    - This allows us to link different telemetry events in a single same training run
+- Telemetry is only sent from the main process to avoid duplicate events
diff --git a/docs/torchao.qmd b/docs/torchao.qmd
new file mode 100644
index 0000000000..f50fc9ee13
--- /dev/null
+++ b/docs/torchao.qmd
@@ -0,0 +1,25 @@
+---
+title: "PyTorch ao"
+description: "Custom data types and layouts for training and inference"
+---
+
+To use experimental optimizers (`AdamWFp8`, `AdamW4bit`, `AdamW8bit`) from Pytorch Ao, please install the package as shown below.
+
+::: {.callout-tip}
+Some experimental optimizers are already present in regular Pytorch, so please re-check if you actually need this package!
+:::
+
+### Installation
+
+Stable Release from the PyTorch index
+
+```bash
+pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124
+```
+
+
+Nightly release
+
+```bash
+pip install --pre torchao-nightly --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124
+```
diff --git a/docs/training_stability.qmd b/docs/training_stability.qmd
new file mode 100644
index 0000000000..9849a35d1b
--- /dev/null
+++ b/docs/training_stability.qmd
@@ -0,0 +1,399 @@
+---
+title: "Training Stability & Debugging"
+order: 15
+description: "Guide to monitoring, debugging, and stabilizing training runs in axolotl"
+---
+
+This guide covers practical techniques for monitoring training health, diagnosing instability, and resolving common failures in both supervised fine-tuning (SFT) and reinforcement learning (GRPO/EBFT) workflows.
+
+## Monitoring Training
+
+### Key Metrics for SFT
+
+Every SFT run should be monitored through at least these four metrics:
+
+| Metric | What It Tells You | Healthy Range |
+|--------|-------------------|---------------|
+| `train/loss` | How well the model fits training data | Decreasing; typically 0.5--2.0 for chat fine-tuning |
+| `eval/loss` | Generalization performance | Tracks train loss with small gap; divergence signals overfitting |
+| `grad_norm` | Gradient magnitude | 0.1--10.0; spikes above 100 indicate instability |
+| `learning_rate` | Current LR from scheduler | Should follow expected schedule (warmup then decay) |
+
+::: {.callout-tip}
+## Set Up Logging Early
+Enable W&B or TensorBoard from the start. Debugging a failed run without metrics is guesswork.
+
+```yaml
+wandb_project: my-project
+wandb_run_id:   # optional, for resuming
+logging_steps: 1
+```
+:::
+
+### Key Metrics for RL (GRPO)
+
+GRPO training logs a richer set of metrics. These are the critical ones:
+
+| Metric | Healthy Range | Red Flag |
+|--------|---------------|----------|
+| `rewards/<name>/mean` | > 0.15 within 20 steps | Stays at 0 -- reward function is broken or task is too hard |
+| `reward_std` | > 0 on most steps | Always 0 -- no learning signal (all completions get the same reward) |
+| `frac_reward_zero_std` | < 0.8 | 1.0 on every step -- zero-advantage skip fires constantly, no gradient updates |
+| `grad_norm` | 0.001--1.0 | 0.0 is acceptable occasionally (zero-adv skip); > 10.0 is unstable |
+| `entropy` | 0.05--0.5 | < 0.01 suggests mode collapse; > 1.0 suggests the model is not converging |
+| `kl` | 0.0--0.5 | > 2.0 suggests policy has diverged too far from reference |
+| `sampling/sampling_logp_difference/mean` | < 0.1 | > 1.0 means policy has diverged far from vLLM server weights |
+| `sampling/importance_sampling_ratio/min` | > 0.1 | Near 0 indicates stale off-policy data; increase `vllm_sync_interval` |
+| `clip_ratio/region_mean` | < 0.1 | > 0.3 means PPO clipping is too aggressive |
+| `completions/mean_length` | Task-dependent | Monotonically increasing to max length suggests reward hacking |
+| `completions/clipped_ratio` | < 0.3 | > 0.8 means most completions hit `max_completion_length` -- increase it |
+
+::: {.callout-note}
+## EBFT-Specific Metrics
+For EBFT training, also monitor `ebft/alignment` (should trend upward, healthy 0.3--0.9), `ebft/diversity` (healthy 0.01--0.1; > 1.0 indicates mode collapse), and `ebft/cfm_loss` (should trend downward, < 10).
+:::
+
+## SFT Stability
+
+### Loss Plateau
+
+**Symptom**: Loss stops decreasing early in training, well above expected values.
+
+**Causes and fixes**:
+
+- **Learning rate too low**: Increase by 2--5x. Typical ranges: full fine-tune 1e-5 to 5e-5, LoRA 1e-4 to 3e-4.
+- **Insufficient warmup**: Set `warmup_steps` to 5--10% of total steps. Too-aggressive learning at the start can push the model into a flat region.
+- **Data quality**: Check that labels are correctly masked. Use `axolotl preprocess` and inspect tokenized samples to confirm only the target tokens are trainable.
+- **Weight decay too high**: Default 0.01 is usually fine. Values above 0.1 can suppress learning in LoRA.
+
+### Loss Spikes
+
+**Symptom**: Loss suddenly jumps by 2--10x then (possibly) recovers.
+
+**Causes and fixes**:
+
+- **Bad data samples**: A single malformed or extremely long example can cause a spike. Enable `sample_packing: false` temporarily and check if spikes correlate with specific batches.
+- **Learning rate too high**: Reduce by 2--5x, or increase warmup.
+- **Gradient accumulation mismatch**: Effective batch size = `micro_batch_size * gradient_accumulation_steps * num_gpus`. Very large effective batch sizes amplify gradient noise.
+- **Mixed precision issues**: With `bf16: true`, some operations can lose precision. If spikes are severe, try `fp32` for diagnosis.
+
+### Overfitting
+
+**Symptom**: Train loss keeps decreasing but eval loss starts increasing.
+
+**Fixes**:
+
+- Increase `val_set_size` (e.g., 0.05) and monitor `eval/loss`.
+- Reduce `num_epochs` or `max_steps`.
+- Increase `weight_decay` (try 0.01--0.1).
+- Use a smaller LoRA rank (`lora_r`). Typical values: 8--32.
+- Increase dropout: `lora_dropout: 0.05`.
+
+## RL/GRPO Stability
+
+### Reward Never Increases
+
+If `rewards/*/mean` stays at 0 for more than 20 steps:
+
+1. **Test reward function standalone**: Run it outside training with known inputs to verify it returns nonzero values.
+   ```bash
+   cd experiments && python -c "import my_rewards; print(my_rewards.accuracy_reward(...))"
+   ```
+2. **Check dataset columns**: The reward function receives `**kwargs` containing dataset columns. Verify the columns it needs (e.g., `answer`) are not removed by the dataset transform.
+3. **Check completion content**: Enable `log_completions: true` in the `trl:` config and inspect logged completions in W&B. If completions are empty or incoherent, the model may be too weak for the task.
+4. **Verify vLLM is serving the right model**: Hit the vLLM health endpoint and confirm the model name matches your config.
+
+### Entropy Collapse (Mode Collapse)
+
+**Symptom**: `entropy` drops below 0.01; all completions become nearly identical.
+
+**Fixes**:
+
+- Increase `temperature` in generation kwargs (try 0.8--1.0).
+- Reduce learning rate.
+- Add a KL penalty term (`beta` parameter in GRPO config).
+- Check that `num_generations` is sufficient (16+ gives better advantage estimates).
+
+### IS Ratio Divergence
+
+**Symptom**: `sampling/importance_sampling_ratio/min` drops near 0, or `sampling/sampling_logp_difference/mean` exceeds 1.0.
+
+This means the policy has diverged significantly from the weights used by vLLM for generation. The importance sampling correction becomes unreliable.
+
+**Fixes**:
+
+- Decrease `vllm_sync_interval` (sync weights more often).
+- Enable `off_policy_mask_threshold` (e.g., 0.5) to mask stale off-policy samples.
+- Use `importance_sampling_level: token` for finer-grained correction.
+
+### Gradient Norm Instability
+
+**Symptom**: `grad_norm` oscillates wildly or exceeds 10.0 regularly.
+
+**Fixes**:
+
+- Enable gradient clipping: `max_grad_norm: 1.0` (default in most configs).
+- Reduce learning rate.
+- Increase `gradient_accumulation_steps` to smooth out noisy batches.
+- Check for NaN issues (see next section).
+
+## NaN and Inf Handling
+
+### Common Causes
+
+| Cause | Where It Manifests | Detection |
+|-------|-------------------|-----------|
+| FP8 zero-scale division | Forward pass logits | `grad_norm: nan`, loss becomes NaN immediately |
+| Gradient explosion | Backward pass | `grad_norm` spikes to inf, then loss goes NaN |
+| Bad data (empty sequences) | Logprob computation | NaN in specific batches only |
+| Numerical overflow in log-softmax | Loss computation | Large negative logprobs cause exp() overflow |
+
+### FP8-Specific NaN Issues
+
+FP8 quantization (`fp8: true`) can produce NaN when the activation quantization kernel divides by `max(abs(x)) / 448`. If the input tensor is all zeros (e.g., padding positions), the scale becomes 0, causing division by zero.
+
+**Fixes applied in axolotl**:
+
+- The `act_quant_kernel` has a zero-guard: `s = tl.where(s == 0, 1.0, s)`.
+- A safety net `nan_to_num(logits, nan=0.0)` is applied in `_get_per_token_logps_and_entropies`.
+- Embedding padding is zero-padded for FP8 compatibility.
+
+::: {.callout-important}
+## After Modifying Triton Kernels
+If you patch any Triton JIT kernel (e.g., the FP8 quantization kernels in transformers), you must clear the Triton cache for changes to take effect:
+
+```bash
+rm -rf ~/.triton/cache
+```
+:::
+
+### General NaN Debugging Steps
+
+1. **Enable anomaly detection** (slow, but pinpoints the source):
+   ```python
+   torch.autograd.set_detect_anomaly(True)
+   ```
+2. **Check grad_norm**: If it goes to NaN, the backward pass is the problem. If loss is NaN but grad_norm was fine on the previous step, the forward pass is the problem.
+3. **Reduce to single GPU, single batch**: Eliminate distributed training variables.
+4. **Inspect data**: Print the batch that triggers NaN. Look for empty sequences, extreme token IDs, or unexpected padding patterns.
+
+## OOM Debugging
+
+Out-of-memory errors are the most common training failure. Use this systematic approach, from least to most disruptive:
+
+### Step 1: Reduce Batch Size
+
+The single highest-impact change. VRAM scales roughly linearly with batch size.
+
+```yaml
+micro_batch_size: 1              # Start here
+gradient_accumulation_steps: 16  # Increase to maintain effective batch size
+```
+
+For GRPO specifically, the logits tensor for policy logprob computation can be very large. `batch_size * num_generations * seq_len * vocab_size` in bf16. For example, with `num_generations: 16` and `micro_batch_size: 8`, the logits tensor alone is:
+
+```
+8 * 16 * 2048 * 151936 * 2 bytes = ~75 GB  (way too large)
+```
+
+Reduce `micro_batch_size` to 2--4 for GRPO.
+
+### Step 2: Enable Gradient Checkpointing
+
+Trades compute for memory by recomputing activations during the backward pass instead of storing them.
+
+```yaml
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false     # Recommended default
+```
+
+::: {.callout-warning}
+## Reentrant Checkpointing Exceptions
+Some configurations require `use_reentrant: true`:
+
+- DeepSpeed ZeRO-3 (non-reentrant causes `CheckpointError`)
+- EBFT strided mode with flex_attention
+:::
+
+### Step 3: Use Quantization
+
+Load the base model in reduced precision:
+
+```yaml
+# 4-bit QLoRA
+adapter: qlora
+load_in_4bit: true
+
+# 8-bit
+load_in_8bit: true
+
+# FP8 (saves ~50% model VRAM, same compute speed as bf16)
+fp8: true
+```
+
+### Step 4: Reduce Sequence Length
+
+```yaml
+sequence_len: 1024     # Down from 2048 or 4096
+```
+
+For GRPO, also reduce `max_completion_length`. Memory scales quadratically with sequence length when using standard attention.
+
+### Step 5: Use Flash Attention
+
+Reduces attention memory from O(n^2) to O(n):
+
+```yaml
+attn_implementation: flash_attention_2
+```
+
+### Step 6: Offload with DeepSpeed
+
+For extreme cases, offload optimizer states or parameters to CPU:
+
+```yaml
+deepspeed: deepspeed_configs/zero3_bf16.json
+```
+
+### Diagnosing the Specific Culprit
+
+Use the `profiler_steps` config option to capture GPU memory snapshots:
+
+```yaml
+profiler_steps: [1, 2]
+```
+
+This generates PyTorch profiler traces you can inspect to see exactly which tensor allocation caused the OOM.
+
+## Common Errors
+
+| Error Message | Likely Cause | Fix |
+|---------------|-------------|-----|
+| `exitcode: -9` | System RAM exhaustion | Reduce dataset size, `dataset_num_proc`, or number of data workers |
+| `exitcode: -7` (DeepSpeed) | DeepSpeed version issue | `pip install -U deepspeed` |
+| `CUDA out of memory` | GPU VRAM exhaustion | Follow OOM debugging steps above |
+| `RuntimeError: NCCL communicator was aborted` | GPU communication failure | See [NCCL docs](nccl.qmd); check `NCCL_DEBUG=INFO` output |
+| `ValueError: Asking to pad but the tokenizer does not have a padding token` | Missing pad token | Add `special_tokens: { pad_token: "<\|endoftext\|>" }` to config |
+| `'DummyOptim' object has no attribute 'step'` | DeepSpeed on single GPU | Remove `deepspeed:` section from config |
+| `unable to load strategy X` then `None is not callable` | Reward module not importable | Run `cd experiments && python -c "import my_rewards"` to check |
+| `generation_batch_size not divisible by num_generations` | micro_batch_size too small | Set `micro_batch_size >= num_generations` and make it divisible |
+| `'weight' must be 2-D` | FSDP1 flattened parameters | Use `fsdp_version: 2` or skip `unwrap_model` when FSDP is enabled |
+| `CheckpointError` (tensor count mismatch) | Non-reentrant checkpointing + ZeRO-3 or flex_attention | Set `use_reentrant: true` in `gradient_checkpointing_kwargs` |
+| `BFloat16` TypeError during weight sync | NumPy does not support bf16 | Fixed in axolotl's `weight_serde.py` (auto bf16 to fp16 conversion) |
+| `Content end boundary is before start boundary` | Chat template parsing issue | Check `eos_token` matches template; file a GitHub issue if persistent |
+| `CAS service error` during data processing | HuggingFace XET issue | Set `export HF_HUB_DISABLE_XET=1` |
+| Training hangs (multi-GPU) | FSDP + async prefetch deadlock | Set `async_prefetch: false` with FSDP |
+
+## Profiling
+
+### PyTorch Profiler
+
+Axolotl supports PyTorch profiler integration via the config:
+
+```yaml
+profiler_steps: [1, 2, 3]
+```
+
+This captures profiler traces for the specified steps. View them in TensorBoard:
+
+```bash
+tensorboard --logdir output_dir/runs
+```
+
+Or open the `.json` trace file in `chrome://tracing`.
+
+### CUDA Memory Snapshots
+
+For detailed memory analysis, use PyTorch's memory snapshot API. Add this to your training script or use it interactively:
+
+```python
+import torch
+
+# Enable memory history tracking
+torch.cuda.memory._record_memory_history()
+
+# ... run your training step ...
+
+# Save snapshot
+torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")
+```
+
+Visualize with PyTorch's memory visualizer:
+
+```bash
+python -m torch.cuda.memory._viz memory_snapshot.pickle
+```
+
+### Quick GPU Memory Check
+
+During training, monitor GPU utilization in a separate terminal:
+
+```bash
+watch -n 1 nvidia-smi
+```
+
+For programmatic access within axolotl, the logged metrics `memory/max_alloc` and `memory/max_reserved` come from `torch.cuda.max_memory_allocated()` and `torch.cuda.max_memory_reserved()`. Note these report PyTorch's view of memory, which may differ from `nvidia-smi` (see [FAQ](faq.qmd)).
+
+## W&B and Logging
+
+### Enabling Logging
+
+```yaml
+wandb_project: my-project
+wandb_entity: my-team          # optional
+wandb_run_id: run-123          # optional, for resuming
+wandb_name: experiment-name    # optional
+logging_steps: 1               # log every step (recommended for RL)
+```
+
+### Debug Logging
+
+For detailed axolotl-internal debug output:
+
+```bash
+AXOLOTL_LOG_LEVEL=DEBUG axolotl train config.yaml 2>&1 | tee /tmp/training.log
+```
+
+::: {.callout-tip}
+## Always Log to a File
+Pipe training output to a log file so you can inspect it after the run:
+
+```bash
+axolotl train config.yaml 2>&1 | tee /tmp/my_run.log
+```
+:::
+
+### What Axolotl Logs
+
+**SFT metrics** (logged every `logging_steps`):
+
+- `train/loss`, `eval/loss` -- training and validation loss
+- `train/grad_norm` -- gradient L2 norm (before clipping)
+- `train/learning_rate` -- current learning rate
+- `memory/max_alloc`, `memory/max_reserved` -- peak GPU memory
+
+**GRPO/RL metrics** (logged every step):
+
+- `rewards/<name>/mean`, `rewards/<name>/std` -- per-reward-function statistics
+- `reward`, `reward_std` -- aggregated reward across all reward functions
+- `frac_reward_zero_std` -- fraction of prompt groups where all completions got the same reward
+- `completions/mean_length`, `completions/min_length`, `completions/max_length` -- completion token lengths
+- `completions/clipped_ratio` -- fraction of completions that hit the max length
+- `completions/mean_terminated_length`, `completions/min_terminated_length`, `completions/max_terminated_length` -- lengths of naturally terminated completions
+- `kl` -- KL divergence between policy and reference
+- `entropy` -- policy entropy (measure of output diversity)
+- `clip_ratio/region_mean`, `clip_ratio/low_mean`, `clip_ratio/high_mean` -- PPO clipping statistics
+- `sampling/sampling_logp_difference/mean`, `sampling/sampling_logp_difference/max` -- log-probability difference between policy and sampling distribution
+- `sampling/importance_sampling_ratio/min`, `sampling/importance_sampling_ratio/mean`, `sampling/importance_sampling_ratio/max` -- IS ratio statistics for off-policy correction
+- `num_tokens` -- total tokens processed
+
+### Reading W&B Charts
+
+For a healthy GRPO run, expect to see:
+
+1. **`reward/mean`**: Gradual upward trend. May start near 0 and reach 0.3--0.8 depending on task difficulty. Not monotonic -- fluctuations are normal.
+2. **`entropy`**: Gradual decrease from initial values (often 0.3--0.6) as the model becomes more confident. Should not collapse to near-zero.
+3. **`grad_norm`**: Mostly in the 0.001--1.0 range. Occasional 0.0 values are fine (zero-advantage skip). Persistent values above 10.0 need investigation.
+4. **`kl`**: Starts near 0 and grows slowly. If it shoots up rapidly, the policy is diverging from the reference.
+5. **`completions/mean_length`**: Should reflect the task's natural answer length. If it steadily increases to `max_completion_length`, the model may be reward-hacking by generating longer outputs.
diff --git a/docs/vllm_serving.qmd b/docs/vllm_serving.qmd
new file mode 100644
index 0000000000..642b723401
--- /dev/null
+++ b/docs/vllm_serving.qmd
@@ -0,0 +1,318 @@
+---
+title: "vLLM Serving for GRPO Training"
+description: "How to configure and run vLLM as a generation backend for GRPO reinforcement learning in Axolotl."
+format:
+  html:
+    toc: true
+    toc-depth: 3
+    number-sections: true
+execute:
+  enabled: false
+---
+
+## Overview {#sec-overview}
+
+GRPO (Group Relative Policy Optimization) trains a language model by generating completions, scoring them with reward functions, and updating the policy to favor higher-reward outputs. The generation step is the bottleneck: producing thousands of tokens per training step with the policy model is slow using standard HuggingFace generation.
+
+Axolotl uses [vLLM](https://github.com/vllm-project/vllm) as a high-throughput generation backend. vLLM runs as a separate process (either on a dedicated GPU or colocated on the training GPU) and serves completions via an HTTP API. The trainer sends prompts to vLLM, receives completions, scores them, and performs gradient updates.
+
+```
+┌──────────────────────┐       HTTP        ┌──────────────────────┐
+│   Trainer (GPU 1)    │ ───────────────── │   vLLM Server (GPU 0)│
+│                      │  prompts/compls   │                      │
+│  - Policy model      │ ◄──────────────── │  - Same base model   │
+│  - Reward scoring    │                   │  - Fast generation   │
+│  - Gradient updates  │  weight sync      │  - LoRA adapter      │
+│  - LoRA adapter      │ ─────────────────►│    (periodically     │
+│                      │  (every N steps)  │     updated)         │
+└──────────────────────┘                   └──────────────────────┘
+```
+
+::: {.callout-important}
+vLLM must serve the **same base model** specified in your training config. If the models do not match, weight synchronization will silently produce incorrect results.
+:::
+
+## Server Mode {#sec-server-mode}
+
+Server mode runs vLLM as an external process on dedicated GPU(s). This is the recommended configuration for most setups.
+
+### Starting the Server
+
+Use the `axolotl vllm-serve` command with your training config:
+
+```bash
+# Terminal 1: Start vLLM on GPU 0
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve grpo_config.yaml
+```
+
+```bash
+# Terminal 2: Start training on GPU 1
+CUDA_VISIBLE_DEVICES=1 axolotl train grpo_config.yaml
+```
+
+The server reads vLLM settings from the `vllm:` section of your config and starts an HTTP server (default: `http://0.0.0.0:8000`).
+
+::: {.callout-tip}
+Use `tmux` or `screen` to manage the vLLM server process. Typical startup time is 30-90 seconds depending on model size and whether CUDA graphs are captured.
+:::
+
+### Minimal Server Config
+
+```yaml
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+
+vllm:
+  host: 0.0.0.0
+  port: 8000
+  gpu_memory_utilization: 0.85
+  dtype: auto
+  max_model_len: 4096
+
+rl: grpo
+trl:
+  use_vllm: true
+  vllm_server_host: 0.0.0.0
+  vllm_server_port: 8000
+  vllm_server_timeout: 300
+```
+
+### Multi-GPU vLLM
+
+For larger models, use tensor parallelism across multiple GPUs:
+
+```yaml
+vllm:
+  tensor_parallel_size: 2
+  gpu_memory_utilization: 0.85
+```
+
+```bash
+# vLLM on GPUs 2,3; training on GPUs 0,1
+CUDA_VISIBLE_DEVICES=2,3 axolotl vllm-serve grpo_config.yaml
+CUDA_VISIBLE_DEVICES=0,1 axolotl train grpo_config.yaml --num-processes 2
+```
+
+::: {.callout-note}
+Due to how TRL maps vLLM device indices, the vLLM instance should use the **last** N GPUs (highest device indices), while training uses the first N.
+:::
+
+## Colocate Mode {#sec-colocate-mode}
+
+Colocate mode runs vLLM on the same GPU as the trainer. This is useful when you only have a single GPU.
+
+```yaml
+trl:
+  use_vllm: true
+  vllm_mode: colocate
+  vllm_enable_sleep_mode: true
+```
+
+With `vllm_enable_sleep_mode: true`, vLLM offloads its VRAM allocation when not actively generating, freeing memory for training. When the trainer needs new completions, vLLM wakes up and reclaims VRAM.
+
+::: {.callout-warning}
+Colocate mode is significantly slower than server mode because generation and training cannot overlap. The GPU alternates between the two workloads. This mode is practical only for smaller models (up to ~3B on a 24 GB GPU).
+:::
+
+**When to use colocate mode:**
+
+- You have exactly one GPU
+- The model fits in memory with both vLLM and training active (with sleep mode), or is small enough to time-share
+- You accept the performance tradeoff for simpler setup (no separate vLLM process to manage)
+
+**When to use server mode:**
+
+- You have two or more GPUs
+- You want maximum throughput (generation overlaps with training via async prefetch)
+- You are running larger models (7B+)
+
+## LoRA Sync {#sec-lora-sync}
+
+LoRA sync is the recommended weight synchronization method when training with LoRA adapters. Instead of merging adapter weights into the base model and broadcasting the full merged weights over NCCL, it saves only the LoRA adapter files to the filesystem and tells vLLM to load them natively.
+
+### How It Works
+
+1. The trainer calls `model.save_pretrained()` to write the LoRA adapter weights to a temporary directory
+2. The trainer sends an HTTP POST to `/set_lora_adapter/` on the vLLM server
+3. vLLM loads the adapter using its native LoRA support (Punica kernels)
+4. Generation uses the updated adapter on the next request
+
+### Benefits
+
+- **Smaller sync payload**: Transfers ~40 MB of LoRA weights instead of ~1.4 GB+ of merged model weights (for a typical 0.5-3B model)
+- **No NCCL communicator**: Eliminates the need for a cross-GPU NCCL communication channel, removing GPU contention between vLLM generation and weight sync
+- **Faster sync**: ~200 ms per sync vs. 350 ms to 5+ seconds for NCCL merge sync
+- **Simpler multi-GPU**: No need to set up NCCL groups between trainer and vLLM processes
+
+### Configuration
+
+```yaml
+adapter: lora
+lora_r: 32
+lora_alpha: 64
+lora_target_linear: true
+
+trl:
+  vllm_lora_sync: true    # Enables LoRA sync mode
+  vllm_sync_interval: 5   # Sync every 5 training steps
+```
+
+Setting `vllm_lora_sync: true` automatically selects the LoRA-aware vLLM serve script (`axolotl.scripts.vllm_serve_lora`). You do not need to set `vllm.serve_module` manually.
+
+::: {.callout-important}
+LoRA sync requires that you are training with a LoRA adapter (`adapter: lora` or `adapter: qlora`). It is not applicable to full fine-tuning.
+:::
+
+## Weight Synchronization {#sec-weight-sync}
+
+During GRPO training, the policy model on the trainer is continuously updated via gradient steps. The vLLM server, however, still holds the old weights. Periodically, the trainer must push updated weights to vLLM so that future generations reflect the improved policy.
+
+### Sync Interval
+
+The `vllm_sync_interval` parameter controls how often weights are synced:
+
+```yaml
+trl:
+  vllm_sync_interval: 5   # Sync every 5 optimizer steps
+```
+
+**Tradeoffs:**
+
+- **Lower interval** (e.g., 1-3): Fresher generations, better on-policy data, but more sync overhead per step
+- **Higher interval** (e.g., 5-10): Less overhead, but generations become increasingly off-policy between syncs
+- **Recommended**: 3-5 for most setups. Axolotl includes importance sampling correction (`vllm_importance_sampling_correction: true`) to handle mild distribution mismatch from stale vLLM weights.
+
+### Sync Methods
+
+| Method | Config | Payload | Mechanism | Typical Time |
+|--------|--------|---------|-----------|-------------|
+| **LoRA sync** | `vllm_lora_sync: true` | LoRA adapter only (~40 MB) | Filesystem + HTTP | ~200 ms |
+| **NCCL merge sync** | Default (no lora_sync) | Full merged weights (~1.4 GB+) | HTTP trigger + NCCL broadcast | 350 ms - 5 s |
+
+::: {.callout-tip}
+If you are training with LoRA (which is recommended for GRPO), always enable `vllm_lora_sync: true`. The performance difference is substantial, especially as training progresses and NCCL contention increases.
+:::
+
+### Importance Sampling Correction
+
+When vLLM weights are stale (between syncs), the generated data is slightly off-policy. Axolotl can correct for this:
+
+```yaml
+trl:
+  vllm_importance_sampling_correction: true
+  importance_sampling_level: token          # 'token' or 'sequence'
+  off_policy_mask_threshold: 0.5            # KL threshold for masking stale sequences
+```
+
+- **Token-level IS** is recommended when using Liger kernel (sequence-level has numerical issues with chunked computation)
+- **Off-policy sequence masking (OPSM)** drops sequences that have diverged too far from the current policy, providing a safety net against stale data
+
+## Restart Requirements {#sec-restart}
+
+::: {.callout-warning}
+**vLLM must be restarted between training runs.** Weight syncs from a previous run leave the server in a corrupted state. If you start a new training run against a stale vLLM server, the model may fail to learn.
+:::
+
+### When to Restart
+
+- Before every new training experiment
+- After a training run crashes or is interrupted
+- If you change the base model in your config
+
+### How to Restart
+
+Killing vLLM reliably requires terminating both the main process and its background EngineCore subprocess:
+
+```bash
+# Kill all vLLM-related processes
+pkill -9 -f "vllm|EngineCore"
+
+# Verify GPU memory is freed
+nvidia-smi
+
+# Restart the server
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve grpo_config.yaml
+```
+
+::: {.callout-tip}
+A single `kill` often does not fully stop vLLM. Always use `kill -9` and verify with `nvidia-smi` that GPU memory has been released before restarting.
+:::
+
+### Health Check
+
+The vLLM server exposes a health endpoint. Wait for it to return 200 before starting training:
+
+```bash
+# For the LoRA serve script (trailing slash required)
+curl http://localhost:8000/health/
+
+# For the default TRL serve script
+curl http://localhost:8000/health
+```
+
+## Configuration Reference {#sec-config-reference}
+
+### vLLM Server Options (`vllm:` section)
+
+These control the vLLM server process started by `axolotl vllm-serve`.
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `host` | str | `0.0.0.0` | Host address for the vLLM server |
+| `port` | int | `8000` | Port for the vLLM server |
+| `device` | str | `auto` | Device to use for vLLM |
+| `tensor_parallel_size` | int | `None` | Number of GPUs for tensor parallelism |
+| `data_parallel_size` | int | `None` | Number of data parallel replicas |
+| `gpu_memory_utilization` | float | `0.9` | Fraction of GPU memory for vLLM (0.0-1.0) |
+| `dtype` | str | `auto` | Data type (`auto`, `float16`, `bfloat16`) |
+| `max_model_len` | int | `None` | Maximum model context length. Set explicitly if the default is too large for your GPU |
+| `enable_prefix_caching` | bool | `None` | Enable prefix caching for repeated prompt prefixes |
+| `enable_reasoning` | bool | `None` | Enable reasoning mode for models with thinking tokens |
+| `reasoning_parser` | str | `None` | Parser for reasoning output |
+| `enforce_eager` | bool | `None` | Disable CUDA graph capture (required for some architectures like Qwen3.5 hybrid attention) |
+| `serve_module` | str | `None` | Python module for vLLM serve script. Auto-set when `vllm_lora_sync: true` |
+| `worker_extension_cls` | str | `None` | vLLM worker extension class for weight sync |
+
+### Trainer vLLM Options (`trl:` section)
+
+These control how the trainer interacts with vLLM.
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `use_vllm` | bool | `false` | Enable vLLM for generation |
+| `vllm_mode` | str | `None` | `server` (external process) or `colocate` (same GPU) |
+| `vllm_server_host` | str | `0.0.0.0` | Host of the vLLM server to connect to |
+| `vllm_server_port` | int | `8000` | Port of the vLLM server to connect to |
+| `vllm_server_timeout` | int | `None` | Timeout in seconds for vLLM requests |
+| `vllm_lora_sync` | bool | `false` | Sync LoRA adapters via filesystem instead of NCCL merge |
+| `vllm_sync_interval` | int | `None` | Sync weights every N optimizer steps |
+| `vllm_enable_sleep_mode` | bool | `None` | Offload vLLM VRAM when idle (colocate mode) |
+| `vllm_guided_decoding_regex` | str | `None` | Regex constraint for guided decoding |
+
+For async pipeline and off-policy correction options, see the [GRPO Configuration Reference](grpo.qmd#configuration-reference).
+
+## Complete Example {#sec-complete-example}
+
+For a full working GRPO config including vLLM, LoRA sync, async generation, rewards, and dataset setup, see the [GRPO Quick Start](grpo.qmd#quick-start). That config includes all the vLLM settings covered in this guide.
+
+```bash
+# Terminal 1: Start vLLM
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve grpo_config.yaml
+
+# Wait for health check to pass
+curl http://localhost:8000/health/
+
+# Terminal 2: Start training
+CUDA_VISIBLE_DEVICES=1 axolotl train grpo_config.yaml
+```
+
+## Troubleshooting {#sec-troubleshooting}
+
+| Problem | Likely Cause | Solution |
+|---------|-------------|----------|
+| Training hangs waiting for vLLM | Server not started or wrong port | Check `curl http://localhost:8000/health/` and verify `vllm_server_host`/`vllm_server_port` match |
+| OOM on vLLM GPU | `gpu_memory_utilization` too high or `max_model_len` too large | Reduce `gpu_memory_utilization` to 0.7 or set `max_model_len` explicitly |
+| OOM on training GPU | Batch too large for policy logprobs | Reduce `micro_batch_size` or `num_generations` |
+| Accuracy stays at zero | Stale vLLM from previous run | Restart vLLM: `pkill -9 -f "vllm\|EngineCore"`, verify with `nvidia-smi`, restart |
+| `ResponseValidationError` from vLLM | Missing logprobs in response | Ensure you are using the correct serve module (auto-selected with `vllm_lora_sync: true`) |
+| Weight sync takes 5+ seconds | NCCL contention with vLLM generation | Switch to `vllm_lora_sync: true` to eliminate NCCL |
+| `async_prefetch` deadlocks with FSDP | Background threads run unsynchronized FSDP collectives | Set `async_prefetch: false` when using FSDP or DeepSpeed multi-GPU |
diff --git a/examples/LiquidAI/README.md b/examples/LiquidAI/README.md
new file mode 100644
index 0000000000..9ac637e33e
--- /dev/null
+++ b/examples/LiquidAI/README.md
@@ -0,0 +1,66 @@
+# Finetune Liquid Foundation Models 2 (LFM2) with Axolotl
+
+[Liquid Foundation Models 2 (LFM2)](https://huggingface.co/collections/LiquidAI/lfm2-686d721927015b2ad73eaa38) are a family of small, open-weight models from [Liquid AI](https://www.liquid.ai/) focused on quality, speed, and memory efficiency. Liquid AI released text-only [LFM2](https://huggingface.co/collections/LiquidAI/lfm2-686d721927015b2ad73eaa38) and text+vision [LFM2-VL](https://huggingface.co/collections/LiquidAI/lfm2-vl-68963bbc84a610f7638d5ffa) models.
+
+LFM2 features a new hybrid Liquid architecture with multiplicative gates, short-range convolutions, and grouped query attention, enabling fast training and inference.
+
+This guide shows how to fine-tune both the LFM2 and LFM2-VL models with Axolotl.
+
+Thanks to the team at LiquidAI for giving us early access to prepare for these releases.
+
+## Getting Started
+
+1.  Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+    ```bash
+    # Ensure you have a compatible version of Pytorch installed
+    uv pip install --no-build-isolation 'axolotl>=0.16.1'
+    ```
+
+2.  Run one of the finetuning examples below.
+
+    **LFM2**
+    ```bash
+    # FFT SFT (1x48GB @ 25GiB)
+    axolotl train examples/LiquidAI/lfm2-350m-fft.yaml
+    ```
+
+    **LFM2-VL**
+    ```bash
+    # LoRA SFT (1x48GB @ 2.7GiB)
+    axolotl train examples/LiquidAI/lfm2-vl-lora.yaml
+    ```
+
+    **LFM2-MoE**
+    ```bash
+    uv pip install git+https://github.com/huggingface/transformers.git@0c9a72e4576fe4c84077f066e585129c97bfd4e6
+
+    # LoRA SFT (1x48GB @ 16.2GiB)
+    axolotl train examples/LiquidAI/lfm2-8b-a1b-lora.yaml
+    ```
+
+### TIPS
+
+- **Installation Error**: If you encounter `ImportError: ... undefined symbol ...` or `ModuleNotFoundError: No module named 'causal_conv1d_cuda'`, the `causal-conv1d` package may have been installed incorrectly. Try uninstalling it:
+  ```bash
+  uv pip uninstall causal-conv1d
+  ```
+
+- **Dataset Loading**: Read more on how to load your own dataset in our [documentation](https://docs.axolotl.ai/docs/dataset_loading.html).
+- **Dataset Formats**:
+  - For LFM2 models, the dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+  - For LFM2-VL models, Axolotl follows the multi-content Messages format. See our [Multimodal docs](https://docs.axolotl.ai/docs/multimodal.html#dataset-format) for details.
+
+## Optimization Guides
+
+- [Optimizations Guide](https://docs.axolotl.ai/docs/optimizations.html)
+
+## Related Resources
+
+- [LFM2 Blog](https://www.liquid.ai/blog/liquid-foundation-models-v2-our-second-series-of-generative-ai-models)
+- [LFM2-VL Blog](https://www.liquid.ai/blog/lfm2-vl-efficient-vision-language-models)
+- [LFM2-MoE Blog](https://www.liquid.ai/blog/lfm2-8b-a1b-an-efficient-on-device-mixture-of-experts)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/LiquidAI/lfm2-350m-fft.yaml b/examples/LiquidAI/lfm2-350m-fft.yaml
new file mode 100644
index 0000000000..cd5942206c
--- /dev/null
+++ b/examples/LiquidAI/lfm2-350m-fft.yaml
@@ -0,0 +1,50 @@
+base_model: LiquidAI/LFM2-350M
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+eot_tokens:
+  - "<|im_end|>"
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_field_role: from
+    message_field_content: value
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.05
+output_dir: ./outputs/out
+
+sequence_len: 4096
+sample_packing: true
+
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 4
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 5e-5
+
+bf16: true
+tf32: true
+
+gradient_checkpointing: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 2
+saves_per_epoch: 1
+
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/LiquidAI/lfm2-8b-a1b-lora.yaml b/examples/LiquidAI/lfm2-8b-a1b-lora.yaml
new file mode 100644
index 0000000000..4932ea06e9
--- /dev/null
+++ b/examples/LiquidAI/lfm2-8b-a1b-lora.yaml
@@ -0,0 +1,59 @@
+base_model: LiquidAI/LFM2-8B-A1B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: true
+
+eot_tokens:
+  - "<|im_end|>"
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_field_role: from
+    message_field_content: value
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.05
+output_dir: ./outputs/out
+
+sequence_len: 4096
+sample_packing: true
+
+adapter: lora
+lora_model_dir:
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: '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: 2
+micro_batch_size: 4
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 5e-5
+
+bf16: true
+tf32: true
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 2
+saves_per_epoch: 1
+
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/LiquidAI/lfm2-vl-lora.yaml b/examples/LiquidAI/lfm2-vl-lora.yaml
new file mode 100644
index 0000000000..9a125da5e0
--- /dev/null
+++ b/examples/LiquidAI/lfm2-vl-lora.yaml
@@ -0,0 +1,60 @@
+base_model: LiquidAI/LFM2-VL-450M
+trust_remote_code: true
+model_type: AutoModelForImageTextToText
+processor_type: AutoProcessor
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+# 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
+
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+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: 'model.language_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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/alst/README.md b/examples/alst/README.md
new file mode 100644
index 0000000000..6d201f8264
--- /dev/null
+++ b/examples/alst/README.md
@@ -0,0 +1,30 @@
+# Arctic Long Sequence Training (ALST)
+
+Artic Long Sequence Training (ALST) is a technique for training long context models using a variety of optimization
+techniques. It is a combination of:
+- TiledMLP: Leverage tiling over the sequence dimension on MLP layers to reduce memory usage
+- Tiled Loss: Using optimized loss functions like Liger-Kernel or Cut Cross Entropy to reduce memory usage
+- Activation Offloading: Offload activations to CPU RAM to reduce memory usage
+
+For more information, you can check out the ALST paper [here](https://www.arxiv.org/abs/2506.13996).
+
+## Usage
+
+```yaml
+tiled_mlp: true
+
+# See Sequence Parallelism docs
+# https://docs.axolotl.ai/docs/sequence_parallelism.html
+context_parallel_size: int
+
+plugins:
+# See Cut Cross Entropy docs
+# https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+# or Liger Kernel docs
+# https://docs.axolotl.ai/docs/custom_integrations.html#liger-kernels
+  - axolotl.integrations.liger.LigerPlugin
+# ...
+
+```
diff --git a/examples/alst/llama3-8b-deepspeed-alst.yaml b/examples/alst/llama3-8b-deepspeed-alst.yaml
new file mode 100644
index 0000000000..e844c68230
--- /dev/null
+++ b/examples/alst/llama3-8b-deepspeed-alst.yaml
@@ -0,0 +1,53 @@
+base_model: meta-llama/Llama-3.1-8B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+datasets:
+  - path: togethercomputer/Long-Data-Collections
+    type: completion
+    field: text
+    data_files:
+      - pretrain/rp_sub.jsonl.zst
+  - path: princeton-nlp/TextbookChapters
+    type: completion
+    field: chapter
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+sequence_len: 500_000
+min_sample_len: 200_000
+sample_packing: true
+
+tiled_mlp: true
+context_parallel_size: 8
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+activation_offloading: legacy
+
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_steps: 100
+saves_per_epoch: 1
+evals_per_epoch: 2
+weight_decay: 0.0
+special_tokens:
+  pad_token: <|end_of_text|>
+
+deepspeed: deepspeed_configs/zero3_bf16_cpuoffload_all.json
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/alst/llama3-8b-fsdp2-alst.yaml b/examples/alst/llama3-8b-fsdp2-alst.yaml
new file mode 100644
index 0000000000..a7da926375
--- /dev/null
+++ b/examples/alst/llama3-8b-fsdp2-alst.yaml
@@ -0,0 +1,59 @@
+base_model: meta-llama/Llama-3.1-8B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+datasets:
+  - path: togethercomputer/Long-Data-Collections
+    type: completion
+    field: text
+    data_files:
+      - pretrain/rp_sub.jsonl.zst
+  - path: princeton-nlp/TextbookChapters
+    type: completion
+    field: chapter
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+sequence_len: 500_000
+min_sample_len: 200_000
+sample_packing: true
+
+tiled_mlp: true
+context_parallel_size: 8
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+activation_offloading: legacy
+
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_steps: 100
+saves_per_epoch: 1
+evals_per_epoch: 2
+weight_decay: 0.0
+special_tokens:
+  pad_token: <|end_of_text|>
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false  # offloading is currently not compatible with SP + torchao optimizer
+  state_dict_type: SHARDED_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  reshard_after_forward: true
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/apertus/README.md b/examples/apertus/README.md
new file mode 100644
index 0000000000..9ff2d89924
--- /dev/null
+++ b/examples/apertus/README.md
@@ -0,0 +1,109 @@
+# Finetune Swiss-AI's Apertus with Axolotl
+
+[Apertus](https://huggingface.co/collections/swiss-ai/apertus-llm-68b699e65415c231ace3b059) is a family of opensource models trained by Swiss-ai.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as Apertus is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).
+
+    Here is an example of how to install from main for pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+git clone https://github.com/axolotl-ai-cloud/axolotl.git
+cd axolotl
+
+uv pip install --no-build-isolation -e '.'
+
+# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
+python scripts/cutcrossentropy_install.py | sh
+```
+
+2. (Optional, highly recommended) Install XIELU CUDA
+
+```bash
+## Recommended for reduced VRAM and faster speeds
+
+# Point to CUDA toolkit directory
+# For those using our Docker image, use the below path.
+export CUDA_HOME=/usr/local/cuda
+
+uv pip install git+https://github.com/nickjbrowning/XIELU@59d6031 --no-build-isolation --no-deps
+```
+
+For any installation errors, see [XIELU Installation Issues](#xielu-installation-issues)
+
+3. Run the finetuning example:
+
+```bash
+axolotl train examples/apertus/apertus-8b-qlora.yaml
+```
+
+This config uses about 8.7 GiB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### Tips
+
+- For inference, the official Apertus team recommends `top_p=0.9` and `temperature=0.8`.
+- You can instead use full paremter fine-tuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+### XIELU Installation Issues
+
+#### `ModuleNotFoundError: No module named 'torch'`
+
+Please check these one by one:
+- Running in correct environment
+- Env has PyTorch installed
+- CUDA toolkit is at `CUDA_HOME`
+
+If those didn't help, please try the below solutions:
+
+1. Pass env for CMAKE and try install again:
+
+    ```bash
+    Python_EXECUTABLE=$(which python) uv pip install git+https://github.com/nickjbrowning/XIELU@59d6031 --no-build-isolation --no-deps
+    ```
+
+2. Git clone the repo and manually hardcode python path:
+
+    ```bash
+    git clone https://github.com/nickjbrowning/XIELU
+    cd xielu
+    git checkout 59d6031
+
+    cd xielu
+    nano CMakeLists.txt  # or vi depending on your preference
+    ```
+
+    ```diff
+    execute_process(
+    -    COMMAND ${Python_EXECUTABLE} -c "import torch.utils; print(torch.utils.cmake_prefix_path)"
+    +    COMMAND /root/miniconda3/envs/py3.11/bin/python -c "import torch.utils; print(torch.utils.cmake_prefix_path)"
+        RESULT_VARIABLE TORCH_CMAKE_PATH_RESULT
+        OUTPUT_VARIABLE TORCH_CMAKE_PATH_OUTPUT
+        ERROR_VARIABLE TORCH_CMAKE_PATH_ERROR
+    )
+    ```
+
+    ```bash
+    uv pip install . --no-build-isolation --no-deps
+    ```
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Related Resources
+
+- [Apertus Tech Report](https://github.com/swiss-ai/apertus-tech-report/blob/main/Apertus_Tech_Report.pdf)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/apertus/apertus-8b-qlora.yaml b/examples/apertus/apertus-8b-qlora.yaml
new file mode 100644
index 0000000000..f43901363c
--- /dev/null
+++ b/examples/apertus/apertus-8b-qlora.yaml
@@ -0,0 +1,64 @@
+base_model: swiss-ai/Apertus-8B-Instruct-2509
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/arcee/README.md b/examples/arcee/README.md
new file mode 100644
index 0000000000..5296e022e3
--- /dev/null
+++ b/examples/arcee/README.md
@@ -0,0 +1,55 @@
+# Finetune ArceeAI's AFM with Axolotl
+
+[Arcee Foundation Models (AFM)](https://huggingface.co/collections/arcee-ai/afm-45b-68823397c351603014963473) are a family of 4.5B parameter open weight models trained by Arcee.ai.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+Thanks to the team at Arcee.ai for using Axolotl in supervised fine-tuning the AFM model.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as AFM is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).
+
+    Here is an example of how to install from main for pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+git clone https://github.com/axolotl-ai-cloud/axolotl.git
+cd axolotl
+
+uv pip install --no-build-isolation -e '.'
+
+# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
+python scripts/cutcrossentropy_install.py | sh
+```
+
+2. Run the finetuning example:
+
+```bash
+axolotl train examples/arcee/afm-4.5b-qlora.yaml
+```
+
+This config uses about 7.8GiB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- For inference, the official Arcee.ai team recommends `top_p: 0.95`, `temperature: 0.5`, `top_k: 50`, and `repeat_penalty: 1.1`.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Related Resources
+
+- [AFM Blog](https://docs.arcee.ai/arcee-foundation-models/introduction-to-arcee-foundation-models)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/arcee/afm-4.5b-qlora.yaml b/examples/arcee/afm-4.5b-qlora.yaml
new file mode 100644
index 0000000000..8e70847adf
--- /dev/null
+++ b/examples/arcee/afm-4.5b-qlora.yaml
@@ -0,0 +1,64 @@
+base_model: arcee-ai/AFM-4.5B
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/archived/README.md b/examples/archived/README.md
new file mode 100644
index 0000000000..da797c5525
--- /dev/null
+++ b/examples/archived/README.md
@@ -0,0 +1,5 @@
+# Archived Examples
+
+This directory contains examples that are no longer maintained and may no longer be functional.
+
+We keep them around for archival purposes in case they are useful to others.
diff --git a/examples/cerebras/btlm-ft.yml b/examples/archived/cerebras/btlm-ft.yml
similarity index 83%
rename from examples/cerebras/btlm-ft.yml
rename to examples/archived/cerebras/btlm-ft.yml
index 18dd86e6b4..5a5f8dc12c 100644
--- a/examples/cerebras/btlm-ft.yml
+++ b/examples/archived/cerebras/btlm-ft.yml
@@ -1,13 +1,13 @@
 base_model: cerebras/btlm-3b-8k-base
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: GPT2Tokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
 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
 datasets:
@@ -30,7 +30,6 @@ lora_alpha:
 lora_dropout:
 lora_target_modules:
 lora_target_linear:
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -38,11 +37,11 @@ wandb_watch:
 wandb_name:
 wandb_log_model:
 
-output_dir: btlm-out
+output_dir: ./outputs/btlm-out
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 1
-optimizer: adamw_torch
+optimizer: adamw_torch_fused
 adam_beta2: 0.95
 adam_eps: 0.000000001
 max_grad_norm: 1.0
@@ -54,30 +53,23 @@ 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:
+attn_implementation: flash_attention_2
 flash_optimum:
 
 gptq_groupsize:
 gptq_model_v1:
 
-warmup_steps: 32
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
 save_total_limit:
 
-debug:
-deepspeed:
 weight_decay: 0.1
 special_tokens:
   pad_token: "<|endoftext|>"
diff --git a/examples/cerebras/qlora.yml b/examples/archived/cerebras/qlora.yml
similarity index 76%
rename from examples/cerebras/qlora.yml
rename to examples/archived/cerebras/qlora.yml
index c4f44326c2..22f52e6821 100644
--- a/examples/cerebras/qlora.yml
+++ b/examples/archived/cerebras/qlora.yml
@@ -1,7 +1,9 @@
 base_model: cerebras/Cerebras-GPT-1.3B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 push_dataset_to_hub:
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -19,13 +21,12 @@ lora_target_modules:
   - c_attn
   - c_proj
 lora_target_linear:
-lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 batch_size: 4
 micro_batch_size: 4
 num_epochs: 2
@@ -33,27 +34,17 @@ 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:
+attn_implementation: xformers
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 special_tokens:
   pad_token: "<|endoftext|>"
diff --git a/examples/code-llama/13b/lora.yml b/examples/archived/code-llama/13b/lora.yml
similarity index 74%
rename from examples/code-llama/13b/lora.yml
rename to examples/archived/code-llama/13b/lora.yml
index ce5a892d08..43f6233579 100644
--- a/examples/code-llama/13b/lora.yml
+++ b/examples/archived/code-llama/13b/lora.yml
@@ -1,21 +1,23 @@
 base_model: codellama/CodeLlama-13b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
+# 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
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter: lora
 lora_model_dir:
@@ -23,7 +25,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -38,29 +39,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
-s2_attention:
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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>"
diff --git a/examples/code-llama/13b/qlora.yml b/examples/archived/code-llama/13b/qlora.yml
similarity index 74%
rename from examples/code-llama/13b/qlora.yml
rename to examples/archived/code-llama/13b/qlora.yml
index d822e68470..086f5e3d85 100644
--- a/examples/code-llama/13b/qlora.yml
+++ b/examples/archived/code-llama/13b/qlora.yml
@@ -1,31 +1,31 @@
 base_model: codellama/CodeLlama-13b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-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:
@@ -40,28 +40,18 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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>"
diff --git a/examples/code-llama/34b/lora.yml b/examples/archived/code-llama/34b/lora.yml
similarity index 74%
rename from examples/code-llama/34b/lora.yml
rename to examples/archived/code-llama/34b/lora.yml
index dfef2538b0..19aa898be7 100644
--- a/examples/code-llama/34b/lora.yml
+++ b/examples/archived/code-llama/34b/lora.yml
@@ -1,21 +1,23 @@
 base_model: codellama/CodeLlama-34b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
+# 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
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter: lora
 lora_model_dir:
@@ -23,7 +25,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -38,29 +39,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
-s2_attention:
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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>"
diff --git a/examples/code-llama/34b/qlora.yml b/examples/archived/code-llama/34b/qlora.yml
similarity index 74%
rename from examples/code-llama/34b/qlora.yml
rename to examples/archived/code-llama/34b/qlora.yml
index 77f821e1c8..2ec78f0d8a 100644
--- a/examples/code-llama/34b/qlora.yml
+++ b/examples/archived/code-llama/34b/qlora.yml
@@ -1,31 +1,31 @@
 base_model: codellama/CodeLlama-34b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-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:
@@ -40,28 +40,18 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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>"
diff --git a/examples/code-llama/7b/lora.yml b/examples/archived/code-llama/7b/lora.yml
similarity index 74%
rename from examples/code-llama/7b/lora.yml
rename to examples/archived/code-llama/7b/lora.yml
index 3e6c7fe620..30bc633553 100644
--- a/examples/code-llama/7b/lora.yml
+++ b/examples/archived/code-llama/7b/lora.yml
@@ -1,21 +1,23 @@
 base_model: codellama/CodeLlama-7b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
+# 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
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter: lora
 lora_model_dir:
@@ -23,7 +25,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -38,29 +39,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
-s2_attention:
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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>"
diff --git a/examples/code-llama/7b/qlora.yml b/examples/archived/code-llama/7b/qlora.yml
similarity index 74%
rename from examples/code-llama/7b/qlora.yml
rename to examples/archived/code-llama/7b/qlora.yml
index e817b113cc..0c3b385193 100644
--- a/examples/code-llama/7b/qlora.yml
+++ b/examples/archived/code-llama/7b/qlora.yml
@@ -1,31 +1,31 @@
 base_model: codellama/CodeLlama-7b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: CodeLlamaTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-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:
@@ -40,28 +40,18 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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>"
diff --git a/examples/code-llama/README.md b/examples/archived/code-llama/README.md
similarity index 100%
rename from examples/code-llama/README.md
rename to examples/archived/code-llama/README.md
diff --git a/examples/dbrx/16bit-lora.yaml b/examples/archived/dbrx/16bit-lora.yaml
similarity index 85%
rename from examples/dbrx/16bit-lora.yaml
rename to examples/archived/dbrx/16bit-lora.yaml
index e5e3ea9216..eca58f94cb 100644
--- a/examples/dbrx/16bit-lora.yaml
+++ b/examples/archived/dbrx/16bit-lora.yaml
@@ -1,16 +1,15 @@
 base_model: LnL-AI/dbrx-base-converted-v2
-trust_remote_code: true
+# 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
+trust_remote_code: true
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 512
 sample_packing: false
@@ -45,26 +44,20 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch:
 saves_per_epoch: 1
-debug:
+
 weight_decay: 0.0
 fsdp:
   - full_shard
diff --git a/examples/dbrx/8bit-lora.yaml b/examples/archived/dbrx/8bit-lora.yaml
similarity index 88%
rename from examples/dbrx/8bit-lora.yaml
rename to examples/archived/dbrx/8bit-lora.yaml
index 89e24db058..59f5241b45 100644
--- a/examples/dbrx/8bit-lora.yaml
+++ b/examples/archived/dbrx/8bit-lora.yaml
@@ -1,16 +1,18 @@
 base_model: LnL-AI/dbrx-base-converted-v2
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
 trust_remote_code: true
 
 load_in_8bit: true
 load_in_4bit: false
-strict: false
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 512
 sample_packing: false
@@ -45,26 +47,20 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch:
 saves_per_epoch: 1
-debug:
+
 weight_decay: 0.0
 fsdp:
   - full_shard
diff --git a/examples/dbrx/README.md b/examples/archived/dbrx/README.md
similarity index 100%
rename from examples/dbrx/README.md
rename to examples/archived/dbrx/README.md
diff --git a/examples/dbrx/fft-ds-zero3.yaml b/examples/archived/dbrx/fft-ds-zero3.yaml
similarity index 76%
rename from examples/dbrx/fft-ds-zero3.yaml
rename to examples/archived/dbrx/fft-ds-zero3.yaml
index 68292707a4..2cb3e6da1a 100644
--- a/examples/dbrx/fft-ds-zero3.yaml
+++ b/examples/archived/dbrx/fft-ds-zero3.yaml
@@ -1,16 +1,15 @@
 base_model: LnL-AI/dbrx-base-converted-v2
-trust_remote_code: true
+# 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
+trust_remote_code: true
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 512
 sample_packing: false
@@ -32,25 +31,19 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch:
 saves_per_epoch: 1
-debug:
+
 weight_decay: 0.0
 deepspeed: deepspeed_configs/zero3_bf16.json
diff --git a/examples/archived/deepcoder/deepcoder-14B-preview-lora.yml b/examples/archived/deepcoder/deepcoder-14B-preview-lora.yml
new file mode 100644
index 0000000000..b125e9e3f3
--- /dev/null
+++ b/examples/archived/deepcoder/deepcoder-14B-preview-lora.yml
@@ -0,0 +1,54 @@
+base_model: agentica-org/DeepCoder-14B-Preview
+# 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: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/lora-out
+
+sequence_len: 4096
+sample_packing: true
+eval_sample_packing: false
+
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/falcon/config-7b-lora.yml b/examples/archived/falcon/config-7b-lora.yml
similarity index 74%
rename from examples/falcon/config-7b-lora.yml
rename to examples/archived/falcon/config-7b-lora.yml
index 5be9c64253..71dd572b3c 100644
--- a/examples/falcon/config-7b-lora.yml
+++ b/examples/archived/falcon/config-7b-lora.yml
@@ -1,12 +1,16 @@
 base_model: tiiuae/falcon-7b
-trust_remote_code: true
+# 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
+
+# required by falcon custom model code: https://huggingface.co/tiiuae/falcon-7b/tree/main
+trust_remote_code: true
 
 load_in_8bit: true
 load_in_4bit: false
 gptq: false
-strict: false
 push_dataset_to_hub:
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -20,15 +24,13 @@ 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:
 wandb_name:
 wandb_log_model:
-output_dir: ./falcon-7b
+output_dir: ./outputs/falcon-7b
 batch_size: 2
 micro_batch_size: 1
 num_epochs: 4
@@ -36,28 +38,18 @@ 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:
+attn_implementation: xformers
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 40
+warmup_ratio: 0.1
 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|>"
diff --git a/examples/falcon/config-7b-qlora.yml b/examples/archived/falcon/config-7b-qlora.yml
similarity index 89%
rename from examples/falcon/config-7b-qlora.yml
rename to examples/archived/falcon/config-7b-qlora.yml
index eb1cdfcdba..edd6550a71 100644
--- a/examples/falcon/config-7b-qlora.yml
+++ b/examples/archived/falcon/config-7b-qlora.yml
@@ -1,16 +1,20 @@
 # 1b: tiiuae/falcon-rw-1b
 # 40b: tiiuae/falcon-40b
 base_model: tiiuae/falcon-7b
-# required by falcon custom model code: https://huggingface.co/tiiuae/falcon-7b/tree/main
-trust_remote_code: true
+# 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
+
+# required by falcon custom model code: https://huggingface.co/tiiuae/falcon-7b/tree/main
+trust_remote_code: true
+
 
 load_in_8bit: false
 # enable 4bit for QLoRA
 load_in_4bit: true
 gptq: false
-strict: false
 push_dataset_to_hub:
 datasets:
   - path: QingyiSi/Alpaca-CoT
@@ -33,16 +37,14 @@ 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:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 # QLoRA paper Table 9
 # - 16 for 7b & 13b
@@ -62,10 +64,7 @@ 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,20 +72,14 @@ 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:
+attn_implementation: xformers
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 10
+warmup_ratio: 0.1
 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|>"
diff --git a/examples/falcon/config-7b.yml b/examples/archived/falcon/config-7b.yml
similarity index 73%
rename from examples/falcon/config-7b.yml
rename to examples/archived/falcon/config-7b.yml
index 1dd46a93ff..6da39d7ab1 100644
--- a/examples/falcon/config-7b.yml
+++ b/examples/archived/falcon/config-7b.yml
@@ -1,12 +1,13 @@
 base_model: tiiuae/falcon-7b
-trust_remote_code: true
+# 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
 
-load_in_8bit: false
-load_in_4bit: false
+# required by falcon custom model code: https://huggingface.co/tiiuae/falcon-7b/tree/main
+trust_remote_code: true
 gptq: false
-strict: false
 push_dataset_to_hub:
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -20,15 +21,13 @@ 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:
 wandb_name:
 wandb_log_model:
-output_dir: ./falcon-7b
+output_dir: ./outputs/falcon-7b
 batch_size: 2
 micro_batch_size: 1
 num_epochs: 4
@@ -36,28 +35,18 @@ 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:
+attn_implementation: xformers
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 40
+warmup_ratio: 0.1
 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|>"
diff --git a/examples/gemma/qlora.yml b/examples/archived/gemma/qlora.yml
similarity index 75%
rename from examples/gemma/qlora.yml
rename to examples/archived/gemma/qlora.yml
index 619a401291..5b5ec4a9f9 100644
--- a/examples/gemma/qlora.yml
+++ b/examples/archived/gemma/qlora.yml
@@ -1,18 +1,20 @@
 # use google/gemma-7b if you have access
 base_model: mhenrichsen/gemma-7b
+# 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
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 # huggingface repo
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 val_set_size: 0.1
-output_dir: ./out
+output_dir: ./outputs/out
 
 adapter: qlora
 lora_r: 32
@@ -23,7 +25,7 @@ lora_target_linear: true
 sequence_len: 4096
 sample_packing: true
 eval_sample_packing: false
-pad_to_sequence_len: true
+
 
 wandb_project:
 wandb_entity:
@@ -39,28 +41,16 @@ 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
+attn_implementation: flash_attention_2
 
 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:
diff --git a/examples/gptj/qlora.yml b/examples/archived/gptj/qlora.yml
similarity index 74%
rename from examples/gptj/qlora.yml
rename to examples/archived/gptj/qlora.yml
index cd3f2e2ad7..7e10adeaa0 100644
--- a/examples/gptj/qlora.yml
+++ b/examples/archived/gptj/qlora.yml
@@ -1,7 +1,9 @@
 base_model: EleutherAI/gpt-j-6b
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 push_dataset_to_hub:
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -15,15 +17,13 @@ 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:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 gradient_accumulation_steps: 2
 micro_batch_size: 2
 num_epochs: 2
@@ -31,27 +31,17 @@ 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:
+attn_implementation: xformers
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 special_tokens:
   pad_token: "<|endoftext|>"
diff --git a/examples/jeopardy-bot/config.yml b/examples/archived/jeopardy-bot/config.yml
similarity index 77%
rename from examples/jeopardy-bot/config.yml
rename to examples/archived/jeopardy-bot/config.yml
index a672c7b94f..90ca3b4bca 100644
--- a/examples/jeopardy-bot/config.yml
+++ b/examples/archived/jeopardy-bot/config.yml
@@ -1,6 +1,10 @@
 base_model: huggyllama/llama-7b
+# optionally might have model_type or tokenizer_type
 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
 datasets:
   - path: openaccess-ai-collective/jeopardy
@@ -21,7 +25,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./jeopardy-bot-7b
+output_dir: ./outputs/jeopardy-bot-7b
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 4
@@ -29,26 +33,17 @@ 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:
+attn_implementation: xformers
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 20
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 tokens:
   bos_token: "<s>"
   eos_token: "</s>"
diff --git a/examples/mpt-7b/README.md b/examples/archived/mpt-7b/README.md
similarity index 100%
rename from examples/mpt-7b/README.md
rename to examples/archived/mpt-7b/README.md
diff --git a/examples/mpt-7b/config.yml b/examples/archived/mpt-7b/config.yml
similarity index 82%
rename from examples/mpt-7b/config.yml
rename to examples/archived/mpt-7b/config.yml
index 45e31266f1..588981bf71 100644
--- a/examples/mpt-7b/config.yml
+++ b/examples/archived/mpt-7b/config.yml
@@ -1,5 +1,9 @@
 base_model: mosaicml/mpt-7b
+# optionally might have model_type or tokenizer_type
 tokenizer_type: AutoTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
 trust_remote_code: true  # required for mpt as their model class is not merged into transformers yet
 load_in_8bit: false
 datasets:
@@ -23,7 +27,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./mpt-alpaca-7b
+output_dir: ./outputs/mpt-alpaca-7b
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 4
@@ -31,26 +35,16 @@ optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.0000002
-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:
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 20
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0001
-fsdp:
-fsdp_config:
 tokens:
   pad_token: "<|padding|>"
   bos_token: "<|endoftext|>"
diff --git a/examples/openllama-3b/README.md b/examples/archived/openllama-3b/README.md
similarity index 100%
rename from examples/openllama-3b/README.md
rename to examples/archived/openllama-3b/README.md
diff --git a/examples/openllama-3b/config.yml b/examples/archived/openllama-3b/config.yml
similarity index 75%
rename from examples/openllama-3b/config.yml
rename to examples/archived/openllama-3b/config.yml
index 0a404c79d8..14104ff4ba 100644
--- a/examples/openllama-3b/config.yml
+++ b/examples/archived/openllama-3b/config.yml
@@ -1,9 +1,9 @@
 base_model: openlm-research/open_llama_3b_v2
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 push_dataset_to_hub:
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -19,13 +19,12 @@ lora_alpha:
 lora_dropout:
 lora_target_modules:
 lora_target_linear:
-lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./openllama-out
+output_dir: ./outputs/openllama-out
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 4
@@ -33,29 +32,20 @@ optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.000003
-train_on_inputs: false
-group_by_length: false
 float16: true
 bf16: false
 fp16: false
 tf32: false
 gradient_checkpointing: true
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention: true
+attn_implementation: flash_attention_2
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 20
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 special_tokens:
   bos_token: "<s>"
   eos_token: "</s>"
diff --git a/examples/openllama-3b/lora.yml b/examples/archived/openllama-3b/lora.yml
similarity index 79%
rename from examples/openllama-3b/lora.yml
rename to examples/archived/openllama-3b/lora.yml
index b83b2db4e4..30d3888f11 100644
--- a/examples/openllama-3b/lora.yml
+++ b/examples/archived/openllama-3b/lora.yml
@@ -1,9 +1,12 @@
 base_model: openlm-research/open_llama_3b_v2
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
+# 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
 push_dataset_to_hub:
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -25,13 +28,12 @@ lora_target_modules:
   - v_proj
   - k_proj
   - o_proj
-lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 gradient_accumulation_steps: 1
 micro_batch_size: 2
 num_epochs: 4
@@ -39,29 +41,19 @@ optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.0002
-train_on_inputs: false
-group_by_length: false
 bf16: false
 fp16: true
 tf32: false
 gradient_checkpointing: true
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention: true
+attn_implementation: flash_attention_2
 gptq_groupsize:
-s2_attention:
 gptq_model_v1:
-warmup_steps: 20
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 special_tokens:
   bos_token: "<s>"
   eos_token: "</s>"
diff --git a/examples/openllama-3b/qlora.yml b/examples/archived/openllama-3b/qlora.yml
similarity index 77%
rename from examples/openllama-3b/qlora.yml
rename to examples/archived/openllama-3b/qlora.yml
index 3d6218b308..fc9d1d7035 100644
--- a/examples/openllama-3b/qlora.yml
+++ b/examples/archived/openllama-3b/qlora.yml
@@ -1,9 +1,12 @@
 base_model: openlm-research/open_llama_3b_v2
+# optionally might have model_type or tokenizer_type
 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: true
-strict: false
 push_dataset_to_hub:
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -17,15 +20,13 @@ sample_packing: true
 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:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 gradient_accumulation_steps: 1
 micro_batch_size: 2
 num_epochs: 4
@@ -33,28 +34,19 @@ optimizer: paged_adamw_32bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.0002
-train_on_inputs: false
-group_by_length: false
 bf16: false
 fp16: true
 tf32: false
 gradient_checkpointing: true
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention: true
+attn_implementation: flash_attention_2
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 20
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 special_tokens:
   bos_token: "<s>"
   eos_token: "</s>"
diff --git a/examples/pythia-12b/README.md b/examples/archived/pythia-12b/README.md
similarity index 100%
rename from examples/pythia-12b/README.md
rename to examples/archived/pythia-12b/README.md
diff --git a/examples/pythia-12b/config.yml b/examples/archived/pythia-12b/config.yml
similarity index 81%
rename from examples/pythia-12b/config.yml
rename to examples/archived/pythia-12b/config.yml
index e44bba7451..170d6ac595 100644
--- a/examples/pythia-12b/config.yml
+++ b/examples/archived/pythia-12b/config.yml
@@ -1,9 +1,10 @@
 base_model: EleutherAI/pythia-12b-deduped
 base_model_ignore_patterns: pytorch*  # prefer safetensors
+# optionally might have model_type or tokenizer_type
 model_type: GPTNeoXForCausalLM
 tokenizer_type: AutoTokenizer
-load_in_8bit: false
-load_in_4bit: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 gptq: false
 device_map: auto
 datasets:
@@ -18,7 +19,6 @@ max_packed_sequence_len: 2048
 lora_r: 64
 lora_alpha: 32
 lora_dropout: 0.0
-lora_target_modules:
 lora_target_linear: true
 lora_fan_in_fan_out: true  # pythia/GPTNeoX lora specific
 wandb_project:
@@ -26,23 +26,17 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./pythia-12b
+output_dir: ./outputs/pythia-12b
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 5
 learning_rate: 0.00003
 optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
-train_on_inputs: false
-group_by_length: false
 bf16: false
 fp16: false
 float16: true
 tf32: true
 flash_optimum: true
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 gradient_checkpointing: true
-fsdp:
-fsdp_config:
diff --git a/examples/pythia/lora.yml b/examples/archived/pythia/lora.yml
similarity index 81%
rename from examples/pythia/lora.yml
rename to examples/archived/pythia/lora.yml
index 7cb07fe258..0549967eca 100644
--- a/examples/pythia/lora.yml
+++ b/examples/archived/pythia/lora.yml
@@ -1,4 +1,7 @@
 base_model: EleutherAI/pythia-1.4b-deduped
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
 load_in_8bit: true
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
@@ -20,18 +23,14 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./lora-alpaca-pythia
+output_dir: ./outputs/lora-alpaca-pythia
 gradient_accumulation_steps: 1
 micro_batch_size: 4
 num_epochs: 4
 learning_rate: 0.00001
-train_on_inputs: false
-group_by_length: false
 bf16: auto
 tf32: true
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 weight_decay: 0.1
 evals_per_epoch: 4
 logging_steps: 1
diff --git a/examples/qwen/README.md b/examples/archived/qwen/README.md
similarity index 100%
rename from examples/qwen/README.md
rename to examples/archived/qwen/README.md
diff --git a/examples/qwen/lora.yml b/examples/archived/qwen/lora.yml
similarity index 74%
rename from examples/qwen/lora.yml
rename to examples/archived/qwen/lora.yml
index da4d784e0a..362a848a8d 100644
--- a/examples/qwen/lora.yml
+++ b/examples/archived/qwen/lora.yml
@@ -1,19 +1,21 @@
 base_model: Qwen/Qwen-7B
+# 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
 
 trust_remote_code: true
 
 load_in_8bit: true
 load_in_4bit: false
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 2048  # supports up to 8192
 sample_packing: false
@@ -25,7 +27,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -40,28 +41,15 @@ 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: false
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention:
 
-warmup_steps: 10
+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:
diff --git a/examples/qwen/qlora.yml b/examples/archived/qwen/qlora.yml
similarity index 74%
rename from examples/qwen/qlora.yml
rename to examples/archived/qwen/qlora.yml
index 501a866b2d..bce3012e72 100644
--- a/examples/qwen/qlora.yml
+++ b/examples/archived/qwen/qlora.yml
@@ -1,19 +1,21 @@
 base_model: Qwen/Qwen-7B
+# 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
 
 trust_remote_code: true
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 2048  # supports up to 8192
 sample_packing: false
@@ -25,7 +27,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -40,28 +41,15 @@ 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: false
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention:
 
-warmup_steps: 10
+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:
diff --git a/examples/qwen/qwen2-moe-lora.yaml b/examples/archived/qwen/qwen2-moe-lora.yaml
similarity index 73%
rename from examples/qwen/qwen2-moe-lora.yaml
rename to examples/archived/qwen/qwen2-moe-lora.yaml
index c59b282d0a..97c0d51a6b 100644
--- a/examples/qwen/qwen2-moe-lora.yaml
+++ b/examples/archived/qwen/qwen2-moe-lora.yaml
@@ -1,16 +1,15 @@
 base_model: Qwen/Qwen1.5-MoE-A2.7B
-trust_remote_code: true
+# 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
+trust_remote_code: true
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 1024  # supports up to 32k
 sample_packing: false
@@ -22,7 +21,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -37,28 +35,18 @@ optimizer: paged_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
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention: true
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
diff --git a/examples/qwen/qwen2-moe-qlora.yaml b/examples/archived/qwen/qwen2-moe-qlora.yaml
similarity index 77%
rename from examples/qwen/qwen2-moe-qlora.yaml
rename to examples/archived/qwen/qwen2-moe-qlora.yaml
index d6a835a0a3..a16089eed4 100644
--- a/examples/qwen/qwen2-moe-qlora.yaml
+++ b/examples/archived/qwen/qwen2-moe-qlora.yaml
@@ -1,16 +1,18 @@
 base_model: Qwen/Qwen1.5-MoE-A2.7B
+# 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: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 1024  # supports up to 32k
 sample_packing: false
@@ -22,7 +24,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -37,28 +38,18 @@ optimizer: paged_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
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention: true
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
diff --git a/examples/redpajama/README.md b/examples/archived/redpajama/README.md
similarity index 100%
rename from examples/redpajama/README.md
rename to examples/archived/redpajama/README.md
diff --git a/examples/redpajama/config-3b.yml b/examples/archived/redpajama/config-3b.yml
similarity index 81%
rename from examples/redpajama/config-3b.yml
rename to examples/archived/redpajama/config-3b.yml
index 5a42e2a952..676f314769 100644
--- a/examples/redpajama/config-3b.yml
+++ b/examples/archived/redpajama/config-3b.yml
@@ -1,6 +1,10 @@
 base_model: togethercomputer/RedPajama-INCITE-Chat-3B-v1
+# optionally might have model_type or tokenizer_type
 model_type: GPTNeoXForCausalLM
 tokenizer_type: AutoTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
 trust_remote_code:
 load_in_8bit: false
 datasets:
@@ -24,7 +28,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./redpajama-alpaca-3b
+output_dir: ./outputs/redpajama-alpaca-3b
 batch_size: 4
 micro_batch_size: 1
 num_epochs: 4
@@ -32,26 +36,16 @@ optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.0000002
-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:
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 20
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0001
-fsdp:
-fsdp_config:
 tokens:
   pad_token: "<|padding|>"
   bos_token: "<|endoftext|>"
diff --git a/examples/replit-3b/config-lora.yml b/examples/archived/replit-3b/config-lora.yml
similarity index 76%
rename from examples/replit-3b/config-lora.yml
rename to examples/archived/replit-3b/config-lora.yml
index bdfe1bd854..b0a0c9089d 100644
--- a/examples/replit-3b/config-lora.yml
+++ b/examples/archived/replit-3b/config-lora.yml
@@ -1,4 +1,7 @@
 base_model: replit/replit-code-v1-3b
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
 trust_remote_code: true
 load_in_8bit: false
 datasets:
@@ -17,13 +20,12 @@ lora_target_modules:
   - Wqkv
   - mlp_up
   - mlp_down
-lora_fan_in_fan_out:
 wandb_project: lora-replit
 wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./lora-replit
+output_dir: ./outputs/lora-replit
 batch_size: 8
 micro_batch_size: 1
 num_epochs: 4
@@ -31,25 +33,15 @@ optimizer:
 torchdistx_path:
 lr_scheduler:
 learning_rate: 0.00001
-train_on_inputs: false
-group_by_length: false
 bf16: auto
 tf32: true
 gradient_checkpointing:
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention:
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 20
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0
-fsdp:
-fsdp_config:
 #special_tokens:
diff --git a/examples/stablelm-2/1.6b/fft.yml b/examples/archived/stablelm-2/1.6b/fft.yml
similarity index 72%
rename from examples/stablelm-2/1.6b/fft.yml
rename to examples/archived/stablelm-2/1.6b/fft.yml
index f3fc16f867..05f59544c6 100644
--- a/examples/stablelm-2/1.6b/fft.yml
+++ b/examples/archived/stablelm-2/1.6b/fft.yml
@@ -1,22 +1,22 @@
 base_model: stabilityai/stablelm-2-1_6b
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: AutoTokenizer
-trust_remote_code: true
+# 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
+trust_remote_code: true
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter:
 lora_model_dir:
@@ -24,7 +24,6 @@ lora_r:
 lora_alpha:
 lora_dropout:
 lora_target_linear:
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -39,31 +38,21 @@ 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
+attn_implementation: flash_attention_2
 flash_attn_cross_entropy: false
 flash_attn_rms_norm: true
-flash_attn_fuse_qkv: false
 flash_attn_fuse_mlp: true
 
-warmup_steps: 100
+warmup_ratio: 0.1
 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:
diff --git a/examples/stablelm-2/1.6b/lora.yml b/examples/archived/stablelm-2/1.6b/lora.yml
similarity index 75%
rename from examples/stablelm-2/1.6b/lora.yml
rename to examples/archived/stablelm-2/1.6b/lora.yml
index c5051fab6e..1edb56e0c0 100644
--- a/examples/stablelm-2/1.6b/lora.yml
+++ b/examples/archived/stablelm-2/1.6b/lora.yml
@@ -1,22 +1,25 @@
 base_model: stabilityai/stablelm-2-1_6b
+# 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
+
 trust_remote_code: true
 
 load_in_8bit: true
 load_in_4bit: false
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter: lora
 lora_model_dir:
@@ -24,7 +27,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -39,28 +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
+attn_implementation: flash_attention_2
 flash_attn_cross_entropy: false
 flash_attn_rms_norm: true
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
diff --git a/examples/stablelm-2/README.md b/examples/archived/stablelm-2/README.md
similarity index 100%
rename from examples/stablelm-2/README.md
rename to examples/archived/stablelm-2/README.md
diff --git a/examples/starcoder2/qlora.yml b/examples/archived/starcoder2/qlora.yml
similarity index 71%
rename from examples/starcoder2/qlora.yml
rename to examples/archived/starcoder2/qlora.yml
index 1efdfbc8e0..0fd0f453cb 100644
--- a/examples/starcoder2/qlora.yml
+++ b/examples/archived/starcoder2/qlora.yml
@@ -1,8 +1,9 @@
 base_model: bigcode/starcoder2-3b
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
@@ -11,21 +12,19 @@ datasets:
 
 dataset_prepared_path:
 val_set_size: 0.2
-output_dir: ./qlora
+output_dir: ./outputs/qlora
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 8192
 sample_packing: true
-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:
@@ -40,30 +39,20 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 2e-5
 
-train_on_inputs: false
-group_by_length: false
 bf16: auto
 fp16: false
 tf32: false
 
 gradient_checkpointing: true
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention: true
+attn_implementation: flash_attention_2
 
-warmup_steps: 20
+warmup_ratio: 0.1
 evals_per_epoch: 4
 eval_steps:
-eval_table_size:
 saves_per_epoch: 4
 save_steps:
 save_total_limit: 2
-debug:
-deepspeed:
 weight_decay:
-fsdp:
-fsdp_config:
 special_tokens:
diff --git a/examples/tiny-llama/README.md b/examples/archived/tiny-llama/README.md
similarity index 100%
rename from examples/tiny-llama/README.md
rename to examples/archived/tiny-llama/README.md
diff --git a/examples/tiny-llama/lora-mps.yml b/examples/archived/tiny-llama/lora-mps.yml
similarity index 67%
rename from examples/tiny-llama/lora-mps.yml
rename to examples/archived/tiny-llama/lora-mps.yml
index fd7b02caca..bf3292c356 100644
--- a/examples/tiny-llama/lora-mps.yml
+++ b/examples/archived/tiny-llama/lora-mps.yml
@@ -1,21 +1,23 @@
-base_model: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
+base_model: TinyLlama/TinyLlama_v1.1
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
+# 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
 dataset_prepared_path:
 val_set_size: 0
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 eval_sample_packing: false
 
 adapter: lora
@@ -24,7 +26,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -35,30 +36,20 @@ wandb_log_model:
 gradient_accumulation_steps: 4
 micro_batch_size: 2
 num_epochs: 4
-optimizer: adamw_torch
+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
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 0
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
diff --git a/examples/tiny-llama/lora.yml b/examples/archived/tiny-llama/lora.yml
similarity index 63%
rename from examples/tiny-llama/lora.yml
rename to examples/archived/tiny-llama/lora.yml
index 4a16f14b92..a12d637466 100644
--- a/examples/tiny-llama/lora.yml
+++ b/examples/archived/tiny-llama/lora.yml
@@ -1,22 +1,23 @@
-base_model: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
-model_type: LlamaForCausalLM
-tokenizer_type: LlamaTokenizer
+base_model: TinyLlama/TinyLlama_v1.1
+# optionally might have model_type or tokenizer_type
+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
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
 eval_sample_packing: false
-pad_to_sequence_len: true
+
 
 adapter: lora
 lora_model_dir:
@@ -24,7 +25,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -39,26 +39,16 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
diff --git a/examples/tiny-llama/pretrain.yml b/examples/archived/tiny-llama/pretrain.yml
similarity index 65%
rename from examples/tiny-llama/pretrain.yml
rename to examples/archived/tiny-llama/pretrain.yml
index 3b68a7f547..4d16861381 100644
--- a/examples/tiny-llama/pretrain.yml
+++ b/examples/archived/tiny-llama/pretrain.yml
@@ -1,20 +1,18 @@
 base_model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
-
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 max_steps: 200
 pretraining_dataset:
-  path: c4
-  name: en
-  type: pretrain
+  - path: allenai/c4
+    name: en
+    type: pretrain
 dataset_prepared_path:
 val_set_size: 0.0
-output_dir: ./model-out
+output_dir: ./outputs/model-out
 
 sequence_len: 2048
 sample_packing: true
@@ -32,27 +30,16 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch:
-eval_table_size:
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
diff --git a/examples/tiny-llama/qlora.yml b/examples/archived/tiny-llama/qlora.yml
similarity index 67%
rename from examples/tiny-llama/qlora.yml
rename to examples/archived/tiny-llama/qlora.yml
index 3ea313c838..b1adcb2e61 100644
--- a/examples/tiny-llama/qlora.yml
+++ b/examples/archived/tiny-llama/qlora.yml
@@ -1,31 +1,32 @@
-base_model: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
+base_model: TinyLlama/TinyLlama_v1.1
+# optionally might have model_type or tokenizer_type
 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: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+eval_sample_packing: false
+
 
 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:
@@ -40,26 +41,16 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
diff --git a/examples/xgen-7b/xgen-7b-8k-qlora.yml b/examples/archived/xgen-7b/xgen-7b-8k-qlora.yml
similarity index 89%
rename from examples/xgen-7b/xgen-7b-8k-qlora.yml
rename to examples/archived/xgen-7b/xgen-7b-8k-qlora.yml
index e3faa01bdb..d548032b9a 100644
--- a/examples/xgen-7b/xgen-7b-8k-qlora.yml
+++ b/examples/archived/xgen-7b/xgen-7b-8k-qlora.yml
@@ -1,14 +1,18 @@
 # An example finetuning Saleforce's XGen-7b model with 8k context using qlora
 # on Tim Dettmer's Guanaco dataset.
 base_model: Salesforce/xgen-7b-8k-base
-trust_remote_code: true
+# 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
+
+trust_remote_code: true
+
 load_in_8bit: false
 # enable 4bit for QLoRA
 load_in_4bit: true
 gptq: false
-strict: false
 push_dataset_to_hub:
 datasets:
   - path: timdettmers/openassistant-guanaco
@@ -31,16 +35,14 @@ 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:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 # QLoRA paper Table 9
 # - 16 for 7b & 13b
@@ -60,10 +62,7 @@ lr_scheduler: cosine
 # - 2e-4 for 7b & 13b
 # - 1e-4 for 33b & 64b
 learning_rate: 0.00002
-train_on_inputs: false
-group_by_length: false
 bf16: auto
-fp16:
 tf32: false
 gradient_checkpointing: true
 # stop training after this many evaluation losses have increased in a row
@@ -71,17 +70,13 @@ 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:
+attn_implementation: xformers
 gptq_groupsize:
 gptq_model_v1:
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
 special_tokens:
   eos_token: "<|endoftext|>"
diff --git a/examples/yi-34B-chat/README.md b/examples/archived/yi-34B-chat/README.md
similarity index 100%
rename from examples/yi-34B-chat/README.md
rename to examples/archived/yi-34B-chat/README.md
diff --git a/examples/yi-34B-chat/qlora.yml b/examples/archived/yi-34B-chat/qlora.yml
similarity index 76%
rename from examples/yi-34B-chat/qlora.yml
rename to examples/archived/yi-34B-chat/qlora.yml
index dc8c37d187..5d3d54dc6e 100644
--- a/examples/yi-34B-chat/qlora.yml
+++ b/examples/archived/yi-34B-chat/qlora.yml
@@ -1,15 +1,16 @@
 base_model: 01-ai/Yi-34B-Chat
+# optionally might have model_type or tokenizer_type
 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: true
-strict: false
 sequence_len: 1024
 bf16: auto
-fp16:
 tf32: false
-flash_attention: true
+attn_implementation: flash_attention_2
 special_tokens:
   bos_token: "<|startoftext|>"
   eos_token: "<|endoftext|>"
@@ -19,7 +20,7 @@ special_tokens:
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
-warmup_steps: 10
+warmup_ratio: 0.1
 
 # Iterations
 num_epochs: 1
@@ -27,20 +28,17 @@ num_epochs: 1
 # Evaluation
 val_set_size: 0.1
 evals_per_epoch: 5
-eval_table_size:
-eval_max_new_tokens: 128
 eval_sample_packing: false
 eval_batch_size: 1
 
 # LoRA
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 adapter: qlora
 lora_model_dir:
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 lora_target_modules:
 
 # Sampling
@@ -61,15 +59,6 @@ lr_scheduler: cosine
 learning_rate: 0.0002
 
 # Misc
-train_on_inputs: false
-group_by_length: false
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-debug:
-deepspeed:
 weight_decay: 0
-fsdp:
-fsdp_config:
diff --git a/examples/cloud/baseten.yaml b/examples/cloud/baseten.yaml
new file mode 100644
index 0000000000..23c4b52d69
--- /dev/null
+++ b/examples/cloud/baseten.yaml
@@ -0,0 +1,10 @@
+provider: baseten
+project_name:
+
+secrets:
+  - HF_TOKEN
+  - WANDB_API_KEY
+
+gpu: h100
+gpu_count: 8
+node_count: 1
diff --git a/examples/cloud/modal.yaml b/examples/cloud/modal.yaml
new file mode 100644
index 0000000000..1950314948
--- /dev/null
+++ b/examples/cloud/modal.yaml
@@ -0,0 +1,28 @@
+project_name:
+volumes:
+  - name: axolotl-data
+    mount: /workspace/data
+  - name: axolotl-artifacts
+    mount: /workspace/artifacts
+
+# environment variables from local to set as secrets
+secrets:
+  - HF_TOKEN
+  - WANDB_API_KEY
+
+# Which branch of axolotl to use remotely
+branch:
+
+# additional custom commands when building the image
+dockerfile_commands:
+
+gpu: h100
+gpu_count: 1
+
+# Train specific configurations
+memory: 128
+timeout: 86400
+
+# Preprocess specific configurations
+memory_preprocess: 32
+timeout_preprocess: 14400
diff --git a/examples/cohere/command-r-7b-qlora.yml b/examples/cohere/command-r-7b-qlora.yml
new file mode 100644
index 0000000000..c4d03b0ec6
--- /dev/null
+++ b/examples/cohere/command-r-7b-qlora.yml
@@ -0,0 +1,59 @@
+base_model: CohereForAI/c4ai-command-r7b-12-2024
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+
+load_in_8bit: false
+load_in_4bit: true
+
+# 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
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/colab-notebooks/colab-axolotl-example.ipynb b/examples/colab-notebooks/colab-axolotl-example.ipynb
index 9adbe00047..68e341e4ac 100644
--- a/examples/colab-notebooks/colab-axolotl-example.ipynb
+++ b/examples/colab-notebooks/colab-axolotl-example.ipynb
@@ -1,216 +1,9955 @@
 {
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "AKjdG7tbTb-n"
-      },
-      "source": [
-        "# Example notebook for running Axolotl on google colab"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "RcbNpOgWRcii"
-      },
-      "outputs": [],
-      "source": [
-        "import torch\n",
-        "# Check so there is a gpu available, a T4(free tier) is enough to run this notebook\n",
-        "assert (torch.cuda.is_available()==True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "h3nLav8oTRA5"
-      },
-      "source": [
-        "## Install Axolotl and dependencies"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "3c3yGAwnOIdi",
-        "outputId": "e3777b5a-40ef-424f-e181-62dfecd1dd01"
-      },
-      "outputs": [],
-      "source": [
-        "!pip install torch==\"2.1.2\"\n",
-        "!pip install -e git+https://github.com/OpenAccess-AI-Collective/axolotl#egg=axolotl\n",
-        "!pip install flash-attn==\"2.5.0\"\n",
-        "!pip install deepspeed==\"0.13.1\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "BW2MFr7HTjub"
-      },
-      "source": [
-        "## Create an yaml config file"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "9pkF2dSoQEUN"
-      },
-      "outputs": [],
-      "source": [
-        "import yaml\n",
-        "\n",
-        "# Your YAML string\n",
-        "yaml_string = \"\"\"\n",
-        "base_model: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T\n",
-        "model_type: LlamaForCausalLM\n",
-        "tokenizer_type: LlamaTokenizer\n",
-        "is_llama_derived_model: true\n",
-        "\n",
-        "load_in_8bit: false\n",
-        "load_in_4bit: true\n",
-        "strict: false\n",
-        "\n",
-        "datasets:\n",
-        "  - path: mhenrichsen/alpaca_2k_test\n",
-        "    type: alpaca\n",
-        "dataset_prepared_path:\n",
-        "val_set_size: 0.05\n",
-        "output_dir: ./qlora-out\n",
-        "\n",
-        "adapter: qlora\n",
-        "lora_model_dir:\n",
-        "\n",
-        "sequence_len: 1096\n",
-        "sample_packing: true\n",
-        "pad_to_sequence_len: true\n",
-        "\n",
-        "lora_r: 32\n",
-        "lora_alpha: 16\n",
-        "lora_dropout: 0.05\n",
-        "lora_target_modules:\n",
-        "lora_target_linear: true\n",
-        "lora_fan_in_fan_out:\n",
-        "\n",
-        "wandb_project:\n",
-        "wandb_entity:\n",
-        "wandb_watch:\n",
-        "wandb_name:\n",
-        "wandb_log_model:\n",
-        "\n",
-        "mlflow_experiment_name: colab-example\n",
-        "\n",
-        "gradient_accumulation_steps: 1\n",
-        "micro_batch_size: 1\n",
-        "num_epochs: 4\n",
-        "max_steps: 20\n",
-        "optimizer: paged_adamw_32bit\n",
-        "lr_scheduler: cosine\n",
-        "learning_rate: 0.0002\n",
-        "\n",
-        "train_on_inputs: false\n",
-        "group_by_length: false\n",
-        "bf16: false\n",
-        "fp16: true\n",
-        "tf32: false\n",
-        "\n",
-        "gradient_checkpointing: true\n",
-        "early_stopping_patience:\n",
-        "resume_from_checkpoint:\n",
-        "local_rank:\n",
-        "logging_steps: 1\n",
-        "xformers_attention:\n",
-        "flash_attention: false\n",
-        "\n",
-        "warmup_steps: 10\n",
-        "evals_per_epoch:\n",
-        "saves_per_epoch:\n",
-        "debug:\n",
-        "deepspeed:\n",
-        "weight_decay: 0.0\n",
-        "fsdp:\n",
-        "fsdp_config:\n",
-        "special_tokens:\n",
-        "\n",
-        "\"\"\"\n",
-        "\n",
-        "# Convert the YAML string to a Python dictionary\n",
-        "yaml_dict = yaml.safe_load(yaml_string)\n",
-        "\n",
-        "# Specify your file path\n",
-        "file_path = 'test_axolotl.yaml'\n",
-        "\n",
-        "# Write the YAML file\n",
-        "with open(file_path, 'w') as file:\n",
-        "    yaml.dump(yaml_dict, file)\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bidoj8YLTusD"
-      },
-      "source": [
-        "## Launch the training"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "ydTI2Jk2RStU",
-        "outputId": "d6d0df17-4b53-439c-c802-22c0456d301b"
-      },
-      "outputs": [],
-      "source": [
-        "# Buy using the ! the comand will be executed as a bash command\n",
-        "!accelerate launch -m axolotl.cli.train /content/test_axolotl.yaml"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## Play with inference"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Buy using the ! the comand will be executed as a bash command\n",
-        "!accelerate launch -m axolotl.cli.inference /content/test_axolotl.yaml \\\n",
-        "    --qlora_model_dir=\"./qlora-out\" --gradio"
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "OPLSwmgdrB7g"
+   },
+   "source": [
+    "# Fine-Tune Qwen3 14B with Axolotl\n",
+    "\n",
+    "[<img src=\"https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/image/axolotl-badge-web.png\" alt=\"Built with Axolotl\" width=\"200\" height=\"32\"/>](https://github.com/axolotl-ai-cloud/axolotl)\n",
+    "\n",
+    "Axolotl is the most performant LLM post-training framework available, delivering faster training with efficient, consistent and stable performance. Train your workload and ship your product 30% faster; saving you both time and money.\n",
+    "\n",
+    "- ⭐ us on [GitHub](https://github.com/axolotl-ai-cloud/axolotl)\n",
+    "- 📜 Read the [Docs](http://docs.axolotl.ai/)\n",
+    "- 💬 Chat with us on [Discord](https://discord.gg/mnpEYgRUmD)\n",
+    "- 📰 Get updates on [X/Twitter](https://x.com/axolotl_ai)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "rVjKD7CbxIP3"
+   },
+   "source": [
+    "# Installation\n",
+    "\n",
+    "Axolotl is easy to install from [pip](https://pypi.org/project/axolotl/), or use our [pre-built Docker images](http://docs.axolotl.ai/docs/docker.html) for a hassle free dependency experience. See our [docs](http://docs.axolotl.ai/docs/installation.html) for more information."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "msOCO4NRmRLa"
+   },
+   "outputs": [],
+   "source": [
+    "%%capture\n",
+    "# This step can take ~5-10 minutes to install dependencies\n",
+    "!pip install --no-build-isolation \"axolotl>=0.16.1\"\n",
+    "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@5f0c7a7\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Colab ships an older numpy that gets half-upgraded during install; force a\n",
+    "# clean, consistent numpy so it doesn't crash with a missing '_blas_supports_fpe'.\n",
+    "# Restart the runtime (Runtime -> Restart session) after this before importing axolotl.\n",
+    "!pip install --force-reinstall --no-cache-dir \"numpy>=2.3\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "N0OW0YeksDLr"
+   },
+   "source": [
+    "## Demo: Talk Like a Pirate\n",
+    "\n",
+    "In this demo, we are training the model ***to respond like a pirate***. This was chosen as a way to easily show how to train a model to respond in a certain style of your choosing (without being prompted) and is quite easy to validate within the scope of a Colab."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "8Du2fANTsNCK"
+   },
+   "source": [
+    "### Upload your own dataset or use a Huggingface dataset\n",
+    "\n",
+    "You can choose to use your own JSONL file from your own [Google Drive](https://drive.google.com/drive/home); for example downloading the [Pirate-Ultrachat JSONL](https://huggingface.co/datasets/winglian/pirate-ultrachat-10k/blob/main/train.jsonl) to your Google Drive. JSONL datasets should be formatted similar to the [OpenAI dataset format](https://cookbook.openai.com/examples/chat_finetuning_data_prep).\n",
+    "\n",
+    "You can also simply use the [`winglian/pirate-ultrachat-10k`](https://huggingface.co/datasets/winglian/pirate-ultrachat-10k) dataset directly.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "fGEEjyQ-r_IV"
+   },
+   "outputs": [],
+   "source": [
+    "# Default to HF dataset location\n",
+    "dataset_id = \"winglian/pirate-ultrachat-10k\"\n",
+    "uploaded = {}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "c5MyYqk7vIsG"
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# Optionally, upload your own JSONL to your Google Drive\n",
+    "GOOGLE_DRIVE_PATH = \"\"  # ex: \"MyDrive/Colab\\ Notebooks/train.jsonl\"\n",
+    "\n",
+    "# \"Select All\" permissions, or you may get the error:\n",
+    "# \"MessageError: Error: credential propagation was unsuccessful\"\n",
+    "if GOOGLE_DRIVE_PATH:\n",
+    "    from google.colab import drive\n",
+    "\n",
+    "    # Mount your Google Drive\n",
+    "    GOOGLE_DRIVE_MNT = \"/content/drive/\"\n",
+    "    drive.mount(GOOGLE_DRIVE_MNT, force_remount=True)\n",
+    "    tmp_path = os.path.join(GOOGLE_DRIVE_MNT, GOOGLE_DRIVE_PATH.lstrip(\"/\"))\n",
+    "    # make sure file exists\n",
+    "    if not os.path.isfile(tmp_path):\n",
+    "        raise ValueError(f\"File {tmp_path} does not exist\")\n",
+    "    dataset_id = tmp_path"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "U6pTk3A9xj1W"
+   },
+   "source": [
+    "# Configure for Supervised Fine-Tuning (SFT)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 151,
+     "referenced_widgets": [
+      "388f618924274d21a066f098f4f1e744",
+      "7c95f85a2b1f47a1bd846d110c47bb3c",
+      "083f9cda8d754c168beee10d2f8955a2",
+      "62e1a65582f446a78612eaa804e08a7d",
+      "487a177d020f4605834878b2fdc7afa3",
+      "7fd44cf9ca6e4726bfd7ac21846d6a14",
+      "366a343b62fa47d8985a3bd464d99f9e",
+      "a0a11e929edd4189b79723d618522c33",
+      "e87ea87fcff247b5bbcc331ba79a8dc2",
+      "5e18768f7ad6434ba8b8b8a2e853e204",
+      "bb33aec33a6447078c31bfd728942994"
+     ]
+    },
+    "id": "fdRioqytmTtX",
+    "outputId": "f0acdcec-4b41-4a3f-ffed-c2d2d929158e"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-08 13:40:27,488] [INFO] [root.register:348] [PID:174] Attempting to load plugin: axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin\n",
+      "[2025-05-08 13:40:27,493] [INFO] [root.register:351] [PID:174] Plugin loaded successfully: axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin\n",
+      "[2025-05-08 13:40:27,959] [INFO] [axolotl.utils.schemas.config.check_eval_packing:721] [PID:174] [RANK:0] explicitly setting `eval_sample_packing` to match `sample_packing`\u001b[39m\n",
+      "[2025-05-08 13:40:27,960] [INFO] [axolotl.utils.schemas.config.hint_sample_packing_padding:514] [PID:174] [RANK:0] Setting `pad_to_sequence_len: true` to prevent memory leaks when sample_packing\u001b[39m\n",
+      "[2025-05-08 13:40:27,961] [INFO] [axolotl.utils.schemas.config.check_bf16:1251] [PID:174] [RANK:0] bf16 support detected, but not enabled for this configuration.\u001b[39m\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "388f618924274d21a066f098f4f1e744",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "config.json:   0%|          | 0.00/728 [00:00<?, ?B/s]"
       ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-08 13:40:28,590] [INFO] [axolotl.normalize_config:237] [PID:174] [RANK:0] cuda memory usage baseline: 0.000GB (+0.002GB cache, +0.359GB misc)\u001b[39m\n"
+     ]
     }
-  ],
-  "metadata": {
-    "accelerator": "GPU",
+   ],
+   "source": [
+    "from axolotl.cli.config import load_cfg\n",
+    "from axolotl.utils.dict import DictDefault\n",
+    "\n",
+    "# Axolotl provides full control and transparency over model and training configuration\n",
+    "config = DictDefault(\n",
+    "    base_model=\"Qwen/Qwen3-14B\",  # Use the instruct tuned model, but we're aligning it to be a pirate\n",
+    "    load_in_4bit=True,  # set to True for qLoRA\n",
+    "    adapter=\"qlora\",\n",
+    "    lora_r=32,\n",
+    "    lora_alpha=64,\n",
+    "    lora_target_modules=[\n",
+    "        \"q_proj\",\n",
+    "        \"k_proj\",\n",
+    "        \"v_proj\",\n",
+    "        \"o_proj\",  # train self_attn linear modules\n",
+    "        \"gate_proj\",\n",
+    "        \"down_proj\",\n",
+    "        \"up_proj\",  # train MLP linear modules\n",
+    "    ],\n",
+    "    lora_qkv_kernel=True,  # optimized triton kernels for LoRA\n",
+    "    lora_o_kernel=True,\n",
+    "    lora_mlp_kernel=True,\n",
+    "    embeddings_skip_upcast=True,  # keep embeddings in fp16 so the model fits in 15GB VRAM\n",
+    "    xformers_attention=True,  # use xformers on Colab w/ T4 for memory efficient attention, flash_attention only on Ampere or above\n",
+    "    plugins=[\n",
+    "        # more efficient training using Apple's Cut Cross Entropy; https://github.com/apple/ml-cross-entropy\n",
+    "        \"axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin\",\n",
+    "    ],\n",
+    "    sample_packing=True,  # 2-6x increase in tokens per micro-batch\n",
+    "    # when using packing, use a slightly higher learning rate to account for fewer steps\n",
+    "    # alternatively, reduce the micro_batch_size + gradient_accumulation_steps to achieve closer to the same number of steps/epoch\n",
+    "    learning_rate=0.00019,\n",
+    "    sequence_len=4096,  # larger sequence length improves packing efficiency for more tokens/sec\n",
+    "    micro_batch_size=1,\n",
+    "    gradient_accumulation_steps=1,\n",
+    "    gradient_checkpointing=True,  # tradeoff reduced VRAM for increased time\n",
+    "    gradient_checkpointing_kwargs={\n",
+    "        \"use_reentrant\": False,\n",
+    "    },\n",
+    "    optimizer=\"paged_adamw_8bit\",\n",
+    "    lr_scheduler=\"cosine\",\n",
+    "    warmup_steps=5,\n",
+    "    fp16=True,  # use float16 + automatic mixed precision, bfloat16 not supported on Colab w/ T4\n",
+    "    bf16=False,\n",
+    "    max_grad_norm=0.1,  # gradient clipping\n",
+    "    num_epochs=1,\n",
+    "    saves_per_epoch=2,  # how many checkpoints to save over one epoch\n",
+    "    logging_steps=1,\n",
+    "    output_dir=\"./outputs/qwen-sft-pirate-rrr\",\n",
+    "    chat_template=\"qwen3\",\n",
+    "    datasets=[\n",
+    "        {\n",
+    "            \"path\": dataset_id,  # Huggingface Dataset id or path to train.jsonl\n",
+    "            \"type\": \"chat_template\",\n",
+    "            \"split\": \"train\",\n",
+    "            \"eot_tokens\": [\"<|im_end|>\"],\n",
+    "        }\n",
+    "    ],\n",
+    "    dataloader_prefetch_factor=8,  # dataloader optimizations\n",
+    "    dataloader_num_workers=2,\n",
+    "    dataloader_pin_memory=True,\n",
+    ")\n",
+    "\n",
+    "# validates the configuration\n",
+    "cfg = load_cfg(config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "715UpvnSoBIS"
+   },
+   "outputs": [],
+   "source": [
+    "from axolotl.utils import set_pytorch_cuda_alloc_conf\n",
+    "\n",
+    "set_pytorch_cuda_alloc_conf()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Vc6MC-hwyH-n"
+   },
+   "source": [
+    "# Datasets\n",
+    "\n",
+    "Axolotl has a robust suite of loaders and transforms to parse most open datasets of any format into the appropriate chat template for your model. Axolotl will mask input tokens from the user's prompt so that the train loss is only calculated against the model's response. For more information, [see our documentation](http://docs.axolotl.ai/docs/dataset-formats/conversation.html) on dataset preparation.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
     "colab": {
-      "gpuType": "T4",
-      "provenance": []
+     "base_uri": "https://localhost:8080/",
+     "height": 1000,
+     "referenced_widgets": [
+      "b82aa8c57f7c422a9a9c90f333ed2a99",
+      "c0991cf63ee6458b96e9a75e7a88b61a",
+      "71c8af139cd248b1b51101fd46a93f35",
+      "1d5117195d4b49eb8f1a73b18419f7ce",
+      "3c21e4a511b4441192c03b7f1d0976e9",
+      "ed28e2e0410d4e0b855467e798e53d66",
+      "d93f134f802b4b69b575bdaf07dbd27c",
+      "d0e9dce55cec4c1ca619a0ccf209d924",
+      "4c727d40ef0443449afc31724ee79f0c",
+      "0dea5caa27384f5689e3cab51f558727",
+      "a6f48410b9964fefba0c3009a77dc838",
+      "95caff42f08a4c2aa14c867b8f37f231",
+      "de7c37ee83e24f0c889e84d07279c2ec",
+      "9d4897eefb5f48259ffb2d23e332f752",
+      "253017b0d0534e54ab44e181f6d7c82d",
+      "27beaf06e41b472abdb544a43c720c5a",
+      "34cf3df51fbc41cabfdbba153c007f0e",
+      "ac764024cf1c4e08ba7749afd2cd20ac",
+      "30a81da86f8043eca301e86a8651201a",
+      "e8b7a81040904c1e89e58978223b1737",
+      "1c6f1f10667545aaab958016ba7e2c94",
+      "e6e969610738449887259063967f82b0",
+      "a138859f19b74fc0928dc236ab5359db",
+      "9b42e08b3c9548818488268768a118b1",
+      "12b56912736849fea2ad8124456fdc5c",
+      "879c8ab5873847a8833bd74123be90a4",
+      "20352e5f58d24bb8b1f3940efd14fe4a",
+      "d955dcaa0e944e719f3a06139dd54a03",
+      "d3de2662c7964f1ba96e58da382af720",
+      "97e36007e1304e1583fd81bfb13f0edd",
+      "c65dc74c7d6f4bab8f7dd28455161dd8",
+      "ef223e8504b64e3592589880326aaf41",
+      "598da69727bd4fb8b1caf465ac736d7a",
+      "5f86cd894de94c3280fadc1e2fd0ee13",
+      "a20927bf5f2c41f58c1e31ac858ab36c",
+      "0a46ad75c198463d843fb35e813642cb",
+      "09007681cf8d42aeb8c1d2f6a74e470a",
+      "ebc80d1a55fa47f4a5ea2756588569ec",
+      "1811cda0644e4190a9469d1774435d82",
+      "35c811d2ae8e43f3b5cecbdd3cfa857f",
+      "b8e39e4dddc3497fbc29ae45c66da759",
+      "63b4e563e85c4f03b1b72beda9577bcc",
+      "b195f160ca20442fadd8b5aed0ee41af",
+      "ca65e32eb52f48c09a84b33cb18f22cd",
+      "7cd0b85ebd204b7aba908417811ce4e0",
+      "7baeab52d6694c32b1efd1ea1a0a7782",
+      "519a7b154022443db6703f04a9142bae",
+      "d4183e9715f34d249942b8271cca3bdf",
+      "da2347ac94764a3fa2743343cf0d3cd2",
+      "93a44a11aa4846fa8efc6c1413ef1627",
+      "a55060adc3564407ac81ad7297d34aaa",
+      "d02274afd47b462291c745f261209d42",
+      "0f417447a7bd4a33acca96fa37aec877",
+      "63580b6fb30642479fe3000915bf551a",
+      "8f726dbfb45d4528afa33e36a6313267",
+      "03b093d592ba4386aa61f7b8483da660",
+      "b8766a88716948cf968f4563531a76d9",
+      "6f3a28b912714c6e931003549664bfa3",
+      "16d1283741404b7bb319094c992fce01",
+      "2a5bb0e818ab47be8cf6465988328503",
+      "2b3a2659b12244bd8548320320016dbf",
+      "0cd7efffbb3c4c4b972e63749f61ab97",
+      "5ca240f31e6b44e3882c5eb37cd5a309",
+      "5eb06edeb58e4930b1affef2a59eae81",
+      "a4e5789584564049b83df7c6c54a3e08",
+      "ff3a94b146a948b6907f5d80c7157f99",
+      "258b7c635c1045329d4669e48c46ccd5",
+      "6f68ed9889f54ad2ae8a3b95ac263a83",
+      "80366349d81e4dcc892db6cd56e384f3",
+      "c73055099c084dca996159e23e162d0b",
+      "977f799afaac4a55b2dc1cffa7d5b63b",
+      "41f3b32c2f6b4034ae7a3b9124e28bc7",
+      "a10d0a76010f4e508c65a9b69ebc5156",
+      "f8ef805b776145c3bfa9ba8d90972058",
+      "cc587493c33c4f118d1b1170f85be24c",
+      "e40d1c1ac9494b3bade9858324e7ffdf",
+      "d65b6b060d9845779299491ac5599c31",
+      "0f6907ebbc6242c8bde059cef1e1bd29",
+      "5bdfd87fc6cd4f9dabef7cfee29c8060",
+      "64f54d4a744a4627a07c3c0120276f3b",
+      "65b75b9b8bc143cf997796af68ff6668",
+      "d6fe74e4255444368f8f90a62157d869",
+      "4d468f96ec924681ad65eb671674b93e",
+      "ad7599de524549c48bf2d3124ad4b299",
+      "0546d04aae644dde846c58a4afb598a6",
+      "897b77a56c09479bb11d7f2a30997e55",
+      "81c3db71ac704280ad030072655f1537",
+      "042e091f75694c47aee761e760e76773",
+      "ef0a3c7a6f14460fb4da096928ae249e",
+      "07fb3a2c8315494e97b447e672dfae06",
+      "ec030fc3c346426f9abc3a89892258d3",
+      "e3fb3fc6afe04b3c9b7ac61809ce78fa",
+      "c3be9109d63c485d9c0ef4f9bc0f9218",
+      "12815f401eba44658caa7b2e490137a8",
+      "30e02aa2d0d241979369e598287f2639",
+      "dfd2a2649b8341ef913207526708aff1",
+      "4f1977d7e4824ef1a14b65f0f42bba10",
+      "c6164e05a1914ae48083db9ad7f4ef7c",
+      "813621384dc748b0ad06775e22761c0b",
+      "dc892a596f6942d7973c616c38f0eebb",
+      "c84cc07789be48aebb322c23d355289e",
+      "bed8726b8069434687c75452e21f19e5",
+      "16a188a0b06d45f980dcf3933509fe0a",
+      "60c1a0d765c14a1d888317e6a507e4ea",
+      "0077aedc3d174560bce924ee89e9c006",
+      "00321cce58884f6f9b3855a21fcd9187",
+      "fa864b41586f4a7aa56aeafd1d84eb75",
+      "3225603166b54e7aab766b9964a2f660",
+      "349eee9f56d64f0cba6fc24ff2c50c9b",
+      "7e5d3774060e4589aa65982da5ea4ef4",
+      "7c2485c6cdfe463da6fdb35982a1070d",
+      "ad1236893754446881e153adc9d5c962",
+      "daee63fd167e4441a32324b51b00ad2b",
+      "fe41858c6bd04c58840112b67c19a336",
+      "d262c82138024169b9f3aa034ca756fa",
+      "62e302ebdad64aada0ffe64ae1c873f3",
+      "bd1b0dfed6d34d16af33a4a58330f5ec",
+      "d07c8b97d3314f1c852e44bdd40f61ed",
+      "ebb69a2c3d0a4299a484698287b3087c",
+      "e5a82df528bb4e408797a3b6c2758f4a",
+      "f113ebd8c1c34806bea4dd7ed3035173"
+     ]
+    },
+    "id": "KQQhgK8FoDfF",
+    "outputId": "f69441d8-95f9-4885-c306-6c8709090ff6"
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b82aa8c57f7c422a9a9c90f333ed2a99",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer_config.json:   0%|          | 0.00/9.68k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
     },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "95caff42f08a4c2aa14c867b8f37f231",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "vocab.json:   0%|          | 0.00/2.78M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
     },
-    "language_info": {
-      "name": "python"
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a138859f19b74fc0928dc236ab5359db",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "merges.txt:   0%|          | 0.00/1.67M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5f86cd894de94c3280fadc1e2fd0ee13",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer.json:   0%|          | 0.00/11.4M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-08 13:41:00,844] [DEBUG] [axolotl.utils.models.load_tokenizer:441] [PID:174] [RANK:0] EOS: 151645 / <|im_end|>\u001b[39m\n",
+      "[2025-05-08 13:41:00,845] [DEBUG] [axolotl.utils.models.load_tokenizer:442] [PID:174] [RANK:0] BOS: None / None\u001b[39m\n",
+      "[2025-05-08 13:41:00,846] [DEBUG] [axolotl.utils.models.load_tokenizer:443] [PID:174] [RANK:0] PAD: 151643 / <|endoftext|>\u001b[39m\n",
+      "[2025-05-08 13:41:00,847] [DEBUG] [axolotl.utils.models.load_tokenizer:444] [PID:174] [RANK:0] UNK: None / None\u001b[39m\n",
+      "[2025-05-08 13:41:00,869] [INFO] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:271] [PID:174] [RANK:0] Unable to find prepared dataset in last_run_prepared/97037817611d38b3a9c681753c3c4c95\u001b[39m\n",
+      "[2025-05-08 13:41:00,870] [INFO] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:272] [PID:174] [RANK:0] Loading raw datasets...\u001b[39m\n",
+      "\u001b[33m[2025-05-08 13:41:00,870] [WARNING] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:274] [PID:174] [RANK:0] Processing datasets during training can lead to VRAM instability. Please pre-process your dataset.\u001b[39m\n",
+      "[2025-05-08 13:41:00,871] [INFO] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:281] [PID:174] [RANK:0] No seed provided, using default seed of 42\u001b[39m\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7cd0b85ebd204b7aba908417811ce4e0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "train.jsonl:   0%|          | 0.00/27.3M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "03b093d592ba4386aa61f7b8483da660",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Generating train split: 0 examples [00:00, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-08 13:41:04,196] [INFO] [axolotl.utils.data.sft.get_dataset_wrapper:484] [PID:174] [RANK:0] Loading dataset with base_type: chat_template and prompt_style: None\u001b[39m\n",
+      "[2025-05-08 13:41:04,233] [INFO] [axolotl.__call__:761] [PID:174] [RANK:0] Using chat template:\n",
+      "---\n",
+      "{%- if tools %}\n",
+      "    {{- '<|im_start|>system\\n' }}\n",
+      "    {%- if messages[0].role == 'system' %}\n",
+      "        {{- messages[0].content + '\\n\\n' }}\n",
+      "    {%- endif %}\n",
+      "    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n",
+      "    {%- for tool in tools %}\n",
+      "        {{- \"\\n\" }}\n",
+      "        {{- tool | tojson }}\n",
+      "    {%- endfor %}\n",
+      "    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n",
+      "{%- else %}\n",
+      "    {%- if messages[0].role == 'system' %}\n",
+      "        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n",
+      "    {%- endif %}\n",
+      "{%- endif %}\n",
+      "{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n",
+      "{%- for message in messages[::-1] %}\n",
+      "    {%- set index = (messages|length - 1) - loop.index0 %}\n",
+      "    {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n",
+      "        {%- set ns.multi_step_tool = false %}\n",
+      "        {%- set ns.last_query_index = index %}\n",
+      "    {%- endif %}\n",
+      "{%- endfor %}\n",
+      "{%- for message in messages %}\n",
+      "    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n",
+      "        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n",
+      "    {%- elif message.role == \"assistant\" %}\n",
+      "        {%- set content = message.content %}\n",
+      "        {%- set reasoning_content = '' %}\n",
+      "        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n",
+      "            {%- set reasoning_content = message.reasoning_content %}\n",
+      "        {%- else %}\n",
+      "            {%- if '</think>' in message.content %}\n",
+      "                {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n",
+      "                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n",
+      "            {%- endif %}\n",
+      "        {%- endif %}\n",
+      "        {%- if loop.index0 > ns.last_query_index %}\n",
+      "            {%- if loop.last or (not loop.last and reasoning_content) %}\n",
+      "                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n",
+      "            {%- else %}\n",
+      "                {{- '<|im_start|>' + message.role + '\\n' + content }}\n",
+      "            {%- endif %}\n",
+      "        {%- else %}\n",
+      "            {{- '<|im_start|>' + message.role + '\\n' + content }}\n",
+      "        {%- endif %}\n",
+      "        {%- if message.tool_calls %}\n",
+      "            {%- for tool_call in message.tool_calls %}\n",
+      "                {%- if (loop.first and content) or (not loop.first) %}\n",
+      "                    {{- '\\n' }}\n",
+      "                {%- endif %}\n",
+      "                {%- if tool_call.function %}\n",
+      "                    {%- set tool_call = tool_call.function %}\n",
+      "                {%- endif %}\n",
+      "                {{- '<tool_call>\\n{\"name\": \"' }}\n",
+      "                {{- tool_call.name }}\n",
+      "                {{- '\", \"arguments\": ' }}\n",
+      "                {%- if tool_call.arguments is string %}\n",
+      "                    {{- tool_call.arguments }}\n",
+      "                {%- else %}\n",
+      "                    {{- tool_call.arguments | tojson }}\n",
+      "                {%- endif %}\n",
+      "                {{- '}\\n</tool_call>' }}\n",
+      "            {%- endfor %}\n",
+      "        {%- endif %}\n",
+      "        {{- '<|im_end|>\\n' }}\n",
+      "    {%- elif message.role == \"tool\" %}\n",
+      "        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n",
+      "            {{- '<|im_start|>user' }}\n",
+      "        {%- endif %}\n",
+      "        {{- '\\n<tool_response>\\n' }}\n",
+      "        {{- message.content }}\n",
+      "        {{- '\\n</tool_response>' }}\n",
+      "        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n",
+      "            {{- '<|im_end|>\\n' }}\n",
+      "        {%- endif %}\n",
+      "    {%- endif %}\n",
+      "{%- endfor %}\n",
+      "{%- if add_generation_prompt %}\n",
+      "    {{- '<|im_start|>assistant\\n' }}\n",
+      "    {%- if enable_thinking is defined and enable_thinking is false %}\n",
+      "        {{- '<think>\\n\\n</think>\\n\\n' }}\n",
+      "    {%- endif %}\n",
+      "{%- endif %}\n",
+      "---\u001b[39m\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "258b7c635c1045329d4669e48c46ccd5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Tokenizing Prompts (num_proc=2):   0%|          | 0/9985 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-08 13:42:09,195] [INFO] [axolotl.utils.data.utils.drop_long_seq_in_dataset:177] [PID:174] [RANK:0] min_input_len: 23\u001b[39m\n",
+      "[2025-05-08 13:42:09,196] [INFO] [axolotl.utils.data.utils.drop_long_seq_in_dataset:179] [PID:174] [RANK:0] max_input_len: 3380\u001b[39m\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0f6907ebbc6242c8bde059cef1e1bd29",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Dropping Long Sequences (num_proc=2):   0%|          | 0/9985 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ef0a3c7a6f14460fb4da096928ae249e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Drop Samples with Zero Trainable Tokens (num_proc=2):   0%|          | 0/9985 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "dc892a596f6942d7973c616c38f0eebb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Add position_id column (Sample Packing) (num_proc=2):   0%|          | 0/9985 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-08 13:42:21,651] [INFO] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:351] [PID:174] [RANK:0] Saving merged prepared dataset to disk... last_run_prepared/97037817611d38b3a9c681753c3c4c95\u001b[39m\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7c2485c6cdfe463da6fdb35982a1070d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Saving the dataset (0/1 shards):   0%|          | 0/9985 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-08 13:42:25,711] [INFO] [axolotl.utils.samplers.multipack.calc_min_len:411] [PID:174] [RANK:0] gather_len_batches: [1540]\u001b[39m\n",
+      "[2025-05-08 13:42:25,714] [INFO] [axolotl.calc_sample_packing_eff_est:491] [PID:174] [RANK:0] sample_packing_eff_est across ranks: [0.9987832601968344]\u001b[39m\n"
+     ]
     }
+   ],
+   "source": [
+    "from axolotl.common.datasets import load_datasets\n",
+    "\n",
+    "# Load, parse and tokenize the datasets to be formatted with qwen3 chat template\n",
+    "# Drop long samples from the dataset that overflow the max sequence length\n",
+    "dataset_meta = load_datasets(cfg=cfg)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "mrSNfHpk0EAe"
+   },
+   "source": [
+    "# Training\n",
+    "\n"
+   ]
   },
-  "nbformat": 4,
-  "nbformat_minor": 0
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 1000,
+     "referenced_widgets": [
+      "004d9177a6a14118a5930dc3cc13147b",
+      "a80410b919e442c49aea15acc1ce1a72",
+      "c6e00f5224364822bc4239b176686919",
+      "ec11d1e5ae7b42c883d9b1f38a65356e",
+      "734185351eb543fa9a00a881dcbb9fe7",
+      "fa1282ccc7544e4f818e2f03ccffe4a5",
+      "bbbf575d2a4b4c6ea8389be79b2a6039",
+      "2a51b36be41745468e4c2d7a21b1c0d2",
+      "4fd114abe9f5494ab59858949f5055f1",
+      "936d04b5fe1b4c63bf0b080e423d051b",
+      "f1cef8e8dc2646fb9fd09f3b09081074",
+      "cdebbc55a1164c018546c2ac6f8c620c",
+      "a44f630e099e43899f20a77084ae60cd",
+      "c3725c7f79fe415fbd1ea336f0cc9cf1",
+      "0e50870ed0c643e0b6c18cc5d7ddae7f",
+      "c33ced495f70464aa4a3a91922090853",
+      "ed5ca967ad5342929e578ac6aa4dc4c0",
+      "af401d117d5047629d3a6e2361757b62",
+      "b191ac001a2e4962bc9a245fcdf26e6b",
+      "054c8dffadba48c6b895a6cc62448ecc",
+      "bfcdbba993b74972a9e3e575f86908ff",
+      "6ebb2ec171414e47a14765505f64bb3c",
+      "500e272208a246089613bf788a165271",
+      "200df5e79b9244849e589ecb0250a520",
+      "cc94432d08464affa3e58b560bdad194",
+      "3036608c71904ce9ae4bb2a9fa8802d9",
+      "adacfdcc1b0140efac56918e9ccf064e",
+      "f4a1795dc7514a718f478245f521f0ba",
+      "5e746eb25bbe416fb585fa24e79f5177",
+      "b5b65414154544aa8a71b1a39164aad7",
+      "f0a58fbd0fca4340890041f99fa2f8c8",
+      "5ca6be24acb548cea130bd58e9954c7c",
+      "5cfb02ee044b4011a378efa8b54a370f",
+      "4d05314858354e729d76094b3b0ce761",
+      "c42acf646f344a88b8c11f81e67f7206",
+      "7be6f04c284e4326bb4ff3d301e7b3c6",
+      "ffdbb12a2f2c4d14911685e7683e0ef0",
+      "bee3501b2a17427784a717e50a85e7fa",
+      "8bc9d8ba866c442b9118d9630009939c",
+      "9f56a2d9979c4bd8928c644c22c3ecdf",
+      "9503a45960984adc97b58e16c50662e0",
+      "da6e93f3e4984780b930fe7a706983ea",
+      "ab93eabd7cea4b94b4b7a387f101e8a1",
+      "704f2f5a9b1c49d5a75a0025a5dda11b",
+      "dd0e646fad3f4a89ba23b39d162bd8d9",
+      "d43c6df07ddb466587807d6dbe1ff614",
+      "e0e8b840b8ea4d0d9db09afe99fa287d",
+      "9327977822be4b1294f80e876552e305",
+      "77304d1a46b3468a98483e02ec0ac4a4",
+      "8c4d4fc5a30f4e7cb3be53fe2adda33d",
+      "e90658f4bcb642baa78426012f863152",
+      "f7434f3e03124a1c938a39af79d7fa59",
+      "c1314f241a434c41b45d84dc4d3b30f8",
+      "37de928300e34184881039378bd75e7f",
+      "0e936d9dbf9c4fdd86bbfe9730dedc47",
+      "e21e180307e5485cbbe908672fd6639a",
+      "2e2b0c1599c341a198f632f46a40c90e",
+      "bff139df987d4a62abec6456cb27f3d4",
+      "ebe1cc366d324ad59b264c8b3c431441",
+      "114dece49dba437c8572ef94b23c3b1e",
+      "be724f04b03942b2a033a7e8898bb4fd",
+      "fcbab4d8dced41a18dfccce81e3a45a0",
+      "c1f9c267ba3f40039cdb5eb3267e8043",
+      "33b3b1d0295646edaac7b4822761aeb0",
+      "fba7aa824b38467ab3061b226114cdec",
+      "f3075dccbd2747b4a7913b66f44f2596",
+      "fe18bba7f3fb4c31bf840541f36b3425",
+      "fd4f333f7ece4450b04e1a9af1f9d2f6",
+      "f60a2bdb6b6b4e0e8c3508580e247132",
+      "c0892a1881de4eb4bfabc6a68f87ae99",
+      "1bec6297c90242a88672d195bc09d429",
+      "d1f9b10c130542f094c8fd3d1e23b5e9",
+      "e575d87a7efe4ec7b1efde489839d4a6",
+      "edc99591b9c747b689b94d0052fec14c",
+      "35cc989ca3374e7dba0cb166febc4bde",
+      "158c8b85dbf34de6a94b4e35e2fc7d5a",
+      "0b4c9753a7cb4354b8e5f187e6e1ad7c",
+      "4471ff62258549fba9514bb67050f965",
+      "9cd5211b5d8b457aa0002f1d17b80028",
+      "19127c7bb1554ccbac877059f9a82db0",
+      "f4667818b9d34a09891cd727a429a610",
+      "9ed02dc43412471a9ab47f3620ccf3a5",
+      "6932489232ec4ab18a160b1e7fbcdfe1",
+      "4540927d98f54466b434ba4c0edf045d",
+      "e400cbf14bcc446a9d33b210cd93550b",
+      "71002199df6b40c9a1ac40df5fb27a1b",
+      "4b27c267393640f28f6eae0875bd2ed9",
+      "9858cb74a09748a39e8149baac96702c",
+      "eb1c9535e6a546098b760528b2ea387c",
+      "18357b321ce44d7b8bd9d1c886f69275",
+      "279937fe03bc4e4eb25b472d7e9df163",
+      "bca2c7185b6749fd899c06a2ba4c5e46",
+      "1f7d30f71bbd4547a9150d21da071055",
+      "e366ae3fceec4566b9ed303d6c5f90af",
+      "5dd7d150dbe04f08b165ce7f2c27cd11",
+      "b634bb73cfa743d09a5999101b840976",
+      "742b1030acfd414bbd9d5327b7e3826d",
+      "0f480e3a0b0a45d2a2d2dec3cad923f3",
+      "fcb30372e7404c5d8a1ad4df91e6c7b2",
+      "2860e3bb3baf4f7da058465850e800c5",
+      "3efd18ea8eaa41918894883da9541bfa",
+      "e09f1bcbb9d94c09be53e5e1303642c2",
+      "82177df57a494de8900c14c2f5185175",
+      "ccfcdc95baf646f8aeb3d516742383f2",
+      "8f5bd719974e41c3a8dd9a5b0d3d71e6",
+      "b87c84de30e84b3abf4871461fb9cbd3",
+      "e7d8e4fe58384e93a106de546068c65e",
+      "0aa8ab56b85f4171a79c3bc210594025",
+      "67da6c4260574869aa24c3cbc1bc1654",
+      "94b9088614464f60a203de39dbcae853",
+      "fea1b70fb46745feb5111b3929175b5d",
+      "f365820a3d3c42b2948abfe32065de14",
+      "823f1c78f15043e38bbd4dca3932a86a",
+      "a1959759c5424da9961fb2a308d4dee4",
+      "34c9c0137b504cd799c6bd6de69507c2",
+      "735d4f225b24414294fc1b213c61223c",
+      "5e5e15b0569b474c9620083b3ec6af55",
+      "03a3c744d716431488163b4358b80f92",
+      "a5434ee714f9498d83870544b67c0cb7",
+      "3aaecbf540f54a2db9ab0931e3b1fe57",
+      "9e333ed3b5014069ac1dd969255dd591"
+     ]
+    },
+    "id": "IwrpurmloGOy",
+    "outputId": "84fa167f-ba27-4255-d508-dc9df56ad39b"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "     #@@ #@@      @@# @@#\n",
+      "    @@  @@          @@  @@           =@@#                               @@                 #@    =@@#.\n",
+      "    @@    #@@@@@@@@@    @@           #@#@=                              @@                 #@     .=@@\n",
+      "      #@@@@@@@@@@@@@@@@@            =@# @#     ##=     ##    =####=+    @@      =#####+  =#@@###.   @@\n",
+      "    @@@@@@@@@@/  +@@/  +@@          #@  =@=     #@=   @@   =@#+  +#@#   @@    =@#+  +#@#   #@.      @@\n",
+      "    @@@@@@@@@@  ##@@  ##@@         =@#   @#      =@# @#    @@      @@   @@    @@      #@   #@       @@\n",
+      "     @@@@@@@@@@@@@@@@@@@@          #@=+++#@=      =@@#     @@      @@   @@    @@      #@   #@       @@\n",
+      "                                  =@#=====@@     =@# @#    @@      @@   @@    @@      #@   #@       @@\n",
+      "    @@@@@@@@@@@@@@@@  @@@@        #@      #@=   #@=  +@@   #@#    =@#   @@.   =@#    =@#   #@.      @@\n",
+      "                                 =@#       @#  #@=     #@   =#@@@@#=    +#@@=  +#@@@@#=    .##@@+   @@\n",
+      "    @@@@  @@@@@@@@@@@@@@@@\n",
+      "\n",
+      "[2025-05-07 22:08:14,344] [INFO] [axolotl.monkeypatch.peft.utils.patch_peft_prep_code:76] [PID:1336] [RANK:0] patching prepare_model_for_kbit_training to allow for overrides\u001b[39m\n",
+      "[2025-05-07 22:08:14,549] [INFO] [axolotl.integrations.cut_cross_entropy.pre_model_load:80] [PID:1336] [RANK:0] Applying Cut Cross Entropy to model type: qwen3\u001b[39m\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "004d9177a6a14118a5930dc3cc13147b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model.safetensors.index.json:   0%|          | 0.00/36.5k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "cdebbc55a1164c018546c2ac6f8c620c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00001-of-00008.safetensors:   0%|          | 0.00/3.84G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "500e272208a246089613bf788a165271",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00002-of-00008.safetensors:   0%|          | 0.00/3.96G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4d05314858354e729d76094b3b0ce761",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00003-of-00008.safetensors:   0%|          | 0.00/3.96G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "dd0e646fad3f4a89ba23b39d162bd8d9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00004-of-00008.safetensors:   0%|          | 0.00/3.96G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e21e180307e5485cbbe908672fd6639a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00005-of-00008.safetensors:   0%|          | 0.00/3.96G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fe18bba7f3fb4c31bf840541f36b3425",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00006-of-00008.safetensors:   0%|          | 0.00/3.96G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4471ff62258549fba9514bb67050f965",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00007-of-00008.safetensors:   0%|          | 0.00/3.96G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "eb1c9535e6a546098b760528b2ea387c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00008-of-00008.safetensors:   0%|          | 0.00/1.91G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-07 22:09:49,798] [INFO] [accelerate.utils.modeling.get_balanced_memory:990] [PID:1336] We will use 90% of the memory on device 0 for storing the model, and 10% for the buffer to avoid OOM. You can set `max_memory` in to a higher value to use more memory (at your own risk).\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2860e3bb3baf4f7da058465850e800c5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/8 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fea1b70fb46745feb5111b3929175b5d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "generation_config.json:   0%|          | 0.00/239 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-07 22:11:37,521] [INFO] [axolotl.utils.models.load_model:1302] [PID:1336] [RANK:0] cuda memory usage after model load: 9.264GB (+1.721GB cache, +0.375GB misc)\u001b[39m\n",
+      "[2025-05-07 22:11:37,532] [INFO] [axolotl.utils.models.prepare_model:1205] [PID:1336] [RANK:0] converting PEFT model w/ prepare_model_for_kbit_training\u001b[39m\n",
+      "[2025-05-07 22:11:37,537] [INFO] [axolotl.utils.models.load_model:1341] [PID:1336] [RANK:0] Converting modules to torch.float16\u001b[39m\n",
+      "trainable params: 128,450,560 || all params: 14,896,757,760 || trainable%: 0.8623\n",
+      "[2025-05-07 22:11:40,170] [INFO] [axolotl.utils.models.load_model:1402] [PID:1336] [RANK:0] cuda memory usage after adapters: 9.743GB (+1.476GB cache, +0.375GB misc)\u001b[39m\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/lib/python3.11/dist-packages/axolotl/core/trainers/base.py:64: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `AxolotlTrainer.__init__`. Use `processing_class` instead.\n",
+      "  super().__init__(*_args, **kwargs)\n",
+      "No label_names provided for model class `PeftModelForCausalLM`. Since `PeftModel` hides base models input arguments, if label_names is not given, label_names can't be set automatically within `Trainer`. Note that empty label_names list will be used instead.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-07 22:11:41,755] [INFO] [axolotl.train.save_initial_configs:359] [PID:1336] [RANK:0] Pre-saving adapter config to ./outputs/qwen-sft-pirate-rrr...\u001b[39m\n",
+      "[2025-05-07 22:11:41,756] [INFO] [axolotl.train.save_initial_configs:363] [PID:1336] [RANK:0] Pre-saving tokenizer to ./outputs/qwen-sft-pirate-rrr...\u001b[39m\n",
+      "[2025-05-07 22:11:41,974] [INFO] [axolotl.train.save_initial_configs:366] [PID:1336] [RANK:0] Pre-saving model config to ./outputs/qwen-sft-pirate-rrr...\u001b[39m\n",
+      "[2025-05-07 22:11:41,982] [INFO] [axolotl.train.execute_training:211] [PID:1336] [RANK:0] Starting trainer...\u001b[39m\n",
+      "[2025-05-07 22:11:45,047] [INFO] [axolotl.utils.samplers.multipack.calc_min_len:411] [PID:1336] [RANK:0] gather_len_batches: [1540]\u001b[39m\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "You're using a Qwen2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
+      "You're using a Qwen2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "    <div>\n",
+       "      \n",
+       "      <progress value='25' max='25' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
+       "      [25/25 09:25, Epoch 0/1]\n",
+       "    </div>\n",
+       "    <table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       " <tr style=\"text-align: left;\">\n",
+       "      <th>Step</th>\n",
+       "      <th>Training Loss</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <td>1</td>\n",
+       "      <td>1.092300</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2</td>\n",
+       "      <td>1.554200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3</td>\n",
+       "      <td>1.041400</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>4</td>\n",
+       "      <td>1.733800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>5</td>\n",
+       "      <td>1.430000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>6</td>\n",
+       "      <td>1.258500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>7</td>\n",
+       "      <td>1.343600</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>8</td>\n",
+       "      <td>1.101700</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>9</td>\n",
+       "      <td>1.086500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>10</td>\n",
+       "      <td>0.813200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>11</td>\n",
+       "      <td>0.689600</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>12</td>\n",
+       "      <td>0.826700</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>13</td>\n",
+       "      <td>1.541800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>14</td>\n",
+       "      <td>0.948000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>15</td>\n",
+       "      <td>1.357000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>16</td>\n",
+       "      <td>1.085800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>17</td>\n",
+       "      <td>1.516800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>18</td>\n",
+       "      <td>1.146800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>19</td>\n",
+       "      <td>0.834800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>20</td>\n",
+       "      <td>0.968000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>21</td>\n",
+       "      <td>1.388800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>22</td>\n",
+       "      <td>1.511500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>23</td>\n",
+       "      <td>1.338500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>24</td>\n",
+       "      <td>1.206600</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>25</td>\n",
+       "      <td>1.504600</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table><p>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2025-05-07 22:12:42,746] [INFO] [axolotl.callbacks.on_step_end:128] [PID:1336] [RANK:0] cuda memory usage while training: 9.768GB (+3.287GB cache, +0.646GB misc)\u001b[39m\n",
+      "[2025-05-07 22:21:46,859] [INFO] [axolotl.train.save_trained_model:231] [PID:1336] [RANK:0] Training completed! Saving pre-trained model to ./outputs/qwen-sft-pirate-rrr.\u001b[39m\n"
+     ]
+    }
+   ],
+   "source": [
+    "from axolotl.train import train\n",
+    "\n",
+    "# just train the first 25 steps for demo.\n",
+    "# This is sufficient to align the model as we've used packing to maximize the trainable samples per step.\n",
+    "cfg.max_steps = 25\n",
+    "model, tokenizer, trainer = train(cfg=cfg, dataset_meta=dataset_meta)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "j1b9ypF78eCb"
+   },
+   "source": [
+    "# Inferencing the trained model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "r3_vHhif8YEs",
+    "outputId": "e5050605-f6c9-421c-98f9-bde56a281eae"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Ahoy there, matey! Shiver me timbers, ye be lookin' for the Pythagorean theorem, eh? Well, hold yer horses and listen up, for I'll be tellin' ye all about it in me own special way.\n",
+      "\n",
+      "The Pythagorean theorem be a real gem of a mathematical trick that helps ye find the length of a side of a right triangle. Now, a right triangle be a triangle with a right angle, which be that little corner that looks like a square. \n",
+      "\n",
+      "The theorem be named after a clever fellow named Pythagoras, who be a mathematician from ancient Greece. He discovered that if ye have a right triangle, the square of the length of the hypotenuse (that be the side opposite the right angle) be equal to the sum of the squares of the other two sides. \n",
+      "\n",
+      "In other words, if ye have a triangle with sides of length a, b, and c (\n"
+     ]
+    }
+   ],
+   "source": [
+    "from transformers import TextStreamer\n",
+    "\n",
+    "messages = [\n",
+    "    {\n",
+    "        \"role\": \"user\",\n",
+    "        \"content\": \"Explain the Pythagorean theorem to me.\",\n",
+    "    },\n",
+    "]\n",
+    "\n",
+    "prompt = tokenizer.apply_chat_template(\n",
+    "    messages,\n",
+    "    add_generation_prompt=True,\n",
+    "    tokenize=False,\n",
+    "    enable_thinking=False,\n",
+    ")\n",
+    "\n",
+    "outputs = model.generate(\n",
+    "    **tokenizer(prompt, return_tensors=\"pt\").to(\"cuda\"),\n",
+    "    max_new_tokens=192,\n",
+    "    temperature=1.0,\n",
+    "    top_p=0.8,\n",
+    "    top_k=32,\n",
+    "    streamer=TextStreamer(tokenizer, skip_prompt=True),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "HoGwT2JRSIjA"
+   },
+   "source": [
+    "# Saving your trained model\n",
+    "\n",
+    "Axolotl automatically saves checkpoints to the `output_dir` path.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "5BmSbiy6NaaS",
+    "outputId": "f5e1d913-7d55-42d2-8340-f9f1b0bc2b38"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "total 506M\n",
+      "-rw-r--r-- 1 root root  845 May  7 22:21 adapter_config.json\n",
+      "-rw-r--r-- 1 root root 491M May  7 22:21 adapter_model.safetensors\n",
+      "-rw-r--r-- 1 root root  707 May  7 22:11 added_tokens.json\n",
+      "drwxr-xr-x 2 root root 4.0K May  7 22:17 checkpoint-13\n",
+      "drwxr-xr-x 2 root root 4.0K May  7 22:21 checkpoint-25\n",
+      "-rw-r--r-- 1 root root 1.2K May  7 22:11 config.json\n",
+      "-rw-r--r-- 1 root root 1.6M May  7 22:11 merges.txt\n",
+      "-rw-r--r-- 1 root root 2.6K May  7 22:21 README.md\n",
+      "-rw-r--r-- 1 root root  613 May  7 22:11 special_tokens_map.json\n",
+      "-rw-r--r-- 1 root root 9.5K May  7 22:11 tokenizer_config.json\n",
+      "-rw-r--r-- 1 root root  11M May  7 22:11 tokenizer.json\n",
+      "-rw-r--r-- 1 root root 2.7M May  7 22:11 vocab.json\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Show the saved checkpoints in the output_dir\n",
+    "!ls -lh \"./outputs/qwen-sft-pirate-rrr\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "_PCIFWxuOZd6"
+   },
+   "source": [
+    "Setting `hub_model_id: ` in the original config would have automatically uploaded the model to HuggingFace Hub (e.g. `hub_model_id: username/model_id`)\n",
+    "\n",
+    "If you prefer to manually upload the training artifacts, we can still upload the entire final checkpoint to HuggingFace from the CLI."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 955,
+     "referenced_widgets": [
+      "c12ea43372ac4d57bb9605f1a429b397",
+      "86816687746246b4a6105e8010384e25",
+      "6f05e9bebf7b40c9835808e77de6c236",
+      "c7433acd3c4841e6958ae8f7e87b1808",
+      "19c1e38389fa46c7b7e2152a56e1df34",
+      "0e067d8db8ed48308a718d5f57683fd1",
+      "131065f118274a1586ac38e39ed84ef0",
+      "8640ac440fbc4644b9a3af7ba3ae7183",
+      "5cea7996f02040b187ece0bb2d6a8d1f",
+      "2e257c8be2da40b4bb67a9e4ab6811f3",
+      "56e3768bef5a4b9db4168c5c17f509c2",
+      "62c028fdef904dedb9cdeca2b3bda725",
+      "a7cf477e80fc43e0ad82c7997b076dce",
+      "835bcc28a5564fb9b3d651bc8e32dc46",
+      "9f1c9a0695384bdaa6f8b847ef89bee8",
+      "b1bea589efa14258a9982071b87938bf",
+      "590eef89881545aa8bbef9a8bbe7fb00",
+      "4b1f04ff63d14a118fdd15814dff50e4",
+      "39789237703c4a418134243055c9cbf5",
+      "a3a945817f684328b34651fe052393ec"
+     ]
+    },
+    "id": "2yw8pLvlSMl8",
+    "outputId": "6e489ab2-4abe-4e28-84ca-959f912433a4"
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c12ea43372ac4d57bb9605f1a429b397",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "It seems you are trying to upload a large folder at once. This might take some time and then fail if the folder is too large. For such cases, it is recommended to upload in smaller batches or to use `HfApi().upload_large_folder(...)`/`huggingface-cli upload-large-folder` instead. For more details, check out https://huggingface.co/docs/huggingface_hub/main/en/guides/upload#upload-a-large-folder.\n",
+      "Start hashing 40 files.\n",
+      "Finished hashing 40 files.\n",
+      "Uploading files using Xet Storage..\n",
+      "Uploading...:  87% 1.82G/2.10G [00:23<00:04, 67.3MB/s]Cancellation requested; stopping current tasks.\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/_commit_api.py\", line 598, in _upload_xet_files\n",
+      "    upload_files(\n",
+      "RuntimeError: Xet Runtime Error: Task cancelled; possible runtime shutdown in progress (task 9 was cancelled).\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/usr/local/bin/huggingface-cli\", line 8, in <module>\n",
+      "    sys.exit(main())\n",
+      "             ^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/commands/huggingface_cli.py\", line 57, in main\n",
+      "    service.run()\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/commands/upload.py\", line 207, in run\n",
+      "    print(self._upload())\n",
+      "          ^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/commands/upload.py\", line 302, in _upload\n",
+      "    return self.api.upload_folder(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py\", line 114, in _inner_fn\n",
+      "    return fn(*args, **kwargs)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/hf_api.py\", line 1633, in _inner\n",
+      "    return fn(self, *args, **kwargs)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/hf_api.py\", line 4942, in upload_folder\n",
+      "    commit_info = self.create_commit(\n",
+      "                  ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py\", line 114, in _inner_fn\n",
+      "    return fn(*args, **kwargs)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/hf_api.py\", line 1633, in _inner\n",
+      "    return fn(self, *args, **kwargs)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/hf_api.py\", line 4202, in create_commit\n",
+      "    self.preupload_lfs_files(\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/hf_api.py\", line 4483, in preupload_lfs_files\n",
+      "    _upload_xet_files(**upload_kwargs, create_pr=create_pr)  # type: ignore [arg-type]\n",
+      "    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py\", line 114, in _inner_fn\n",
+      "    return fn(*args, **kwargs)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/huggingface_hub/_commit_api.py\", line 592, in _upload_xet_files\n",
+      "    with progress_cm as progress:\n",
+      "  File \"/usr/local/lib/python3.11/dist-packages/tqdm/std.py\", line 1138, in __exit__\n",
+      "    def __exit__(self, exc_type, exc_value, traceback):\n",
+      "\n",
+      "KeyboardInterrupt\n",
+      "^C\n"
+     ]
+    }
+   ],
+   "source": [
+    "from huggingface_hub import notebook_login\n",
+    "\n",
+    "# remove the partial epoch checkpoints\n",
+    "!rm -rf \"./outputs/qwen-sft-pirate-rrr/checkpoint-*\"\n",
+    "\n",
+    "# HF Notebook login widget\n",
+    "notebook_login()\n",
+    "\n",
+    "# upload the LoRA adapter for your model to HF, remember to update the username/model-name below\n",
+    "!huggingface-cli upload --repo-type=model winglian/pirate-qwen-14B \"./outputs/qwen-sft-pirate-rrr\""
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "widgets": {
+   "application/vnd.jupyter.widget-state+json": {
+    "00321cce58884f6f9b3855a21fcd9187": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "004d9177a6a14118a5930dc3cc13147b": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_a80410b919e442c49aea15acc1ce1a72",
+       "IPY_MODEL_c6e00f5224364822bc4239b176686919",
+       "IPY_MODEL_ec11d1e5ae7b42c883d9b1f38a65356e"
+      ],
+      "layout": "IPY_MODEL_734185351eb543fa9a00a881dcbb9fe7"
+     }
+    },
+    "0077aedc3d174560bce924ee89e9c006": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "03a3c744d716431488163b4358b80f92": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "03b093d592ba4386aa61f7b8483da660": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_b8766a88716948cf968f4563531a76d9",
+       "IPY_MODEL_6f3a28b912714c6e931003549664bfa3",
+       "IPY_MODEL_16d1283741404b7bb319094c992fce01"
+      ],
+      "layout": "IPY_MODEL_2a5bb0e818ab47be8cf6465988328503"
+     }
+    },
+    "042e091f75694c47aee761e760e76773": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "0546d04aae644dde846c58a4afb598a6": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "054c8dffadba48c6b895a6cc62448ecc": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "07fb3a2c8315494e97b447e672dfae06": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_12815f401eba44658caa7b2e490137a8",
+      "placeholder": "​",
+      "style": "IPY_MODEL_30e02aa2d0d241979369e598287f2639",
+      "value": "Drop Samples with Zero Trainable Tokens (num_proc=2): 100%"
+     }
+    },
+    "083f9cda8d754c168beee10d2f8955a2": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_a0a11e929edd4189b79723d618522c33",
+      "max": 728,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_e87ea87fcff247b5bbcc331ba79a8dc2",
+      "value": 728
+     }
+    },
+    "09007681cf8d42aeb8c1d2f6a74e470a": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_b195f160ca20442fadd8b5aed0ee41af",
+      "placeholder": "​",
+      "style": "IPY_MODEL_ca65e32eb52f48c09a84b33cb18f22cd",
+      "value": " 11.4M/11.4M [00:00&lt;00:00, 21.8MB/s]"
+     }
+    },
+    "0a46ad75c198463d843fb35e813642cb": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_b8e39e4dddc3497fbc29ae45c66da759",
+      "max": 11422654,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_63b4e563e85c4f03b1b72beda9577bcc",
+      "value": 11422654
+     }
+    },
+    "0aa8ab56b85f4171a79c3bc210594025": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "0b4c9753a7cb4354b8e5f187e6e1ad7c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "0cd7efffbb3c4c4b972e63749f61ab97": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "0dea5caa27384f5689e3cab51f558727": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "0e067d8db8ed48308a718d5f57683fd1": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_b1bea589efa14258a9982071b87938bf",
+      "placeholder": "​",
+      "style": "IPY_MODEL_590eef89881545aa8bbef9a8bbe7fb00",
+      "value": "\n<b>Pro Tip:</b> If you don't already have one, you can create a dedicated\n'notebooks' token with 'write' access, that you can then easily reuse for all\nnotebooks. </center>"
+     }
+    },
+    "0e50870ed0c643e0b6c18cc5d7ddae7f": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_bfcdbba993b74972a9e3e575f86908ff",
+      "placeholder": "​",
+      "style": "IPY_MODEL_6ebb2ec171414e47a14765505f64bb3c",
+      "value": " 3.84G/3.84G [00:09&lt;00:00, 664MB/s]"
+     }
+    },
+    "0e936d9dbf9c4fdd86bbfe9730dedc47": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "0f417447a7bd4a33acca96fa37aec877": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "0f480e3a0b0a45d2a2d2dec3cad923f3": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "0f6907ebbc6242c8bde059cef1e1bd29": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_5bdfd87fc6cd4f9dabef7cfee29c8060",
+       "IPY_MODEL_64f54d4a744a4627a07c3c0120276f3b",
+       "IPY_MODEL_65b75b9b8bc143cf997796af68ff6668"
+      ],
+      "layout": "IPY_MODEL_d6fe74e4255444368f8f90a62157d869"
+     }
+    },
+    "114dece49dba437c8572ef94b23c3b1e": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "12815f401eba44658caa7b2e490137a8": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "12b56912736849fea2ad8124456fdc5c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_97e36007e1304e1583fd81bfb13f0edd",
+      "max": 1671853,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_c65dc74c7d6f4bab8f7dd28455161dd8",
+      "value": 1671853
+     }
+    },
+    "131065f118274a1586ac38e39ed84ef0": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": "center",
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": "flex",
+      "flex": null,
+      "flex_flow": "column",
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": "50%"
+     }
+    },
+    "158c8b85dbf34de6a94b4e35e2fc7d5a": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "16a188a0b06d45f980dcf3933509fe0a": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_349eee9f56d64f0cba6fc24ff2c50c9b",
+      "placeholder": "​",
+      "style": "IPY_MODEL_7e5d3774060e4589aa65982da5ea4ef4",
+      "value": " 9985/9985 [00:04&lt;00:00, 2604.11 examples/s]"
+     }
+    },
+    "16d1283741404b7bb319094c992fce01": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_a4e5789584564049b83df7c6c54a3e08",
+      "placeholder": "​",
+      "style": "IPY_MODEL_ff3a94b146a948b6907f5d80c7157f99",
+      "value": " 9985/0 [00:00&lt;00:00, 50763.46 examples/s]"
+     }
+    },
+    "1811cda0644e4190a9469d1774435d82": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "18357b321ce44d7b8bd9d1c886f69275": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_e366ae3fceec4566b9ed303d6c5f90af",
+      "placeholder": "​",
+      "style": "IPY_MODEL_5dd7d150dbe04f08b165ce7f2c27cd11",
+      "value": "model-00008-of-00008.safetensors: 100%"
+     }
+    },
+    "19127c7bb1554ccbac877059f9a82db0": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "danger",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_e400cbf14bcc446a9d33b210cd93550b",
+      "max": 3963750880,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_71002199df6b40c9a1ac40df5fb27a1b",
+      "value": 3963750502
+     }
+    },
+    "19c1e38389fa46c7b7e2152a56e1df34": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ButtonModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ButtonModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ButtonView",
+      "button_style": "",
+      "description": "Login",
+      "disabled": false,
+      "icon": "",
+      "layout": "IPY_MODEL_835bcc28a5564fb9b3d651bc8e32dc46",
+      "style": "IPY_MODEL_9f1c9a0695384bdaa6f8b847ef89bee8",
+      "tooltip": ""
+     }
+    },
+    "1bec6297c90242a88672d195bc09d429": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "1c6f1f10667545aaab958016ba7e2c94": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "1d5117195d4b49eb8f1a73b18419f7ce": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_0dea5caa27384f5689e3cab51f558727",
+      "placeholder": "​",
+      "style": "IPY_MODEL_a6f48410b9964fefba0c3009a77dc838",
+      "value": " 9.68k/9.68k [00:00&lt;00:00, 812kB/s]"
+     }
+    },
+    "1f7d30f71bbd4547a9150d21da071055": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "200df5e79b9244849e589ecb0250a520": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_f4a1795dc7514a718f478245f521f0ba",
+      "placeholder": "​",
+      "style": "IPY_MODEL_5e746eb25bbe416fb585fa24e79f5177",
+      "value": "model-00002-of-00008.safetensors: 100%"
+     }
+    },
+    "20352e5f58d24bb8b1f3940efd14fe4a": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "253017b0d0534e54ab44e181f6d7c82d": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_1c6f1f10667545aaab958016ba7e2c94",
+      "placeholder": "​",
+      "style": "IPY_MODEL_e6e969610738449887259063967f82b0",
+      "value": " 2.78M/2.78M [00:00&lt;00:00, 17.8MB/s]"
+     }
+    },
+    "258b7c635c1045329d4669e48c46ccd5": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_6f68ed9889f54ad2ae8a3b95ac263a83",
+       "IPY_MODEL_80366349d81e4dcc892db6cd56e384f3",
+       "IPY_MODEL_c73055099c084dca996159e23e162d0b"
+      ],
+      "layout": "IPY_MODEL_977f799afaac4a55b2dc1cffa7d5b63b"
+     }
+    },
+    "279937fe03bc4e4eb25b472d7e9df163": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "danger",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_b634bb73cfa743d09a5999101b840976",
+      "max": 1912371880,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_742b1030acfd414bbd9d5327b7e3826d",
+      "value": 1912371698
+     }
+    },
+    "27beaf06e41b472abdb544a43c720c5a": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "2860e3bb3baf4f7da058465850e800c5": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_3efd18ea8eaa41918894883da9541bfa",
+       "IPY_MODEL_e09f1bcbb9d94c09be53e5e1303642c2",
+       "IPY_MODEL_82177df57a494de8900c14c2f5185175"
+      ],
+      "layout": "IPY_MODEL_ccfcdc95baf646f8aeb3d516742383f2"
+     }
+    },
+    "2a51b36be41745468e4c2d7a21b1c0d2": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "2a5bb0e818ab47be8cf6465988328503": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "2b3a2659b12244bd8548320320016dbf": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "2e257c8be2da40b4bb67a9e4ab6811f3": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "2e2b0c1599c341a198f632f46a40c90e": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_be724f04b03942b2a033a7e8898bb4fd",
+      "placeholder": "​",
+      "style": "IPY_MODEL_fcbab4d8dced41a18dfccce81e3a45a0",
+      "value": "model-00005-of-00008.safetensors: 100%"
+     }
+    },
+    "3036608c71904ce9ae4bb2a9fa8802d9": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_5ca6be24acb548cea130bd58e9954c7c",
+      "placeholder": "​",
+      "style": "IPY_MODEL_5cfb02ee044b4011a378efa8b54a370f",
+      "value": " 3.96G/3.96G [00:10&lt;00:00, 531MB/s]"
+     }
+    },
+    "30a81da86f8043eca301e86a8651201a": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "30e02aa2d0d241979369e598287f2639": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "3225603166b54e7aab766b9964a2f660": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "33b3b1d0295646edaac7b4822761aeb0": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "349eee9f56d64f0cba6fc24ff2c50c9b": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "34c9c0137b504cd799c6bd6de69507c2": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "34cf3df51fbc41cabfdbba153c007f0e": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "35c811d2ae8e43f3b5cecbdd3cfa857f": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "35cc989ca3374e7dba0cb166febc4bde": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "366a343b62fa47d8985a3bd464d99f9e": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "37de928300e34184881039378bd75e7f": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "388f618924274d21a066f098f4f1e744": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_7c95f85a2b1f47a1bd846d110c47bb3c",
+       "IPY_MODEL_083f9cda8d754c168beee10d2f8955a2",
+       "IPY_MODEL_62e1a65582f446a78612eaa804e08a7d"
+      ],
+      "layout": "IPY_MODEL_487a177d020f4605834878b2fdc7afa3"
+     }
+    },
+    "39789237703c4a418134243055c9cbf5": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "3aaecbf540f54a2db9ab0931e3b1fe57": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "3c21e4a511b4441192c03b7f1d0976e9": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "3efd18ea8eaa41918894883da9541bfa": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_8f5bd719974e41c3a8dd9a5b0d3d71e6",
+      "placeholder": "​",
+      "style": "IPY_MODEL_b87c84de30e84b3abf4871461fb9cbd3",
+      "value": "Loading checkpoint shards: 100%"
+     }
+    },
+    "41f3b32c2f6b4034ae7a3b9124e28bc7": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "4471ff62258549fba9514bb67050f965": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_9cd5211b5d8b457aa0002f1d17b80028",
+       "IPY_MODEL_19127c7bb1554ccbac877059f9a82db0",
+       "IPY_MODEL_f4667818b9d34a09891cd727a429a610"
+      ],
+      "layout": "IPY_MODEL_9ed02dc43412471a9ab47f3620ccf3a5"
+     }
+    },
+    "4540927d98f54466b434ba4c0edf045d": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "487a177d020f4605834878b2fdc7afa3": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "4b1f04ff63d14a118fdd15814dff50e4": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "LabelModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "LabelModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "LabelView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_39789237703c4a418134243055c9cbf5",
+      "placeholder": "​",
+      "style": "IPY_MODEL_a3a945817f684328b34651fe052393ec",
+      "value": "Connecting..."
+     }
+    },
+    "4b27c267393640f28f6eae0875bd2ed9": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "4c727d40ef0443449afc31724ee79f0c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "4d05314858354e729d76094b3b0ce761": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_c42acf646f344a88b8c11f81e67f7206",
+       "IPY_MODEL_7be6f04c284e4326bb4ff3d301e7b3c6",
+       "IPY_MODEL_ffdbb12a2f2c4d14911685e7683e0ef0"
+      ],
+      "layout": "IPY_MODEL_bee3501b2a17427784a717e50a85e7fa"
+     }
+    },
+    "4d468f96ec924681ad65eb671674b93e": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "4f1977d7e4824ef1a14b65f0f42bba10": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "4fd114abe9f5494ab59858949f5055f1": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "500e272208a246089613bf788a165271": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_200df5e79b9244849e589ecb0250a520",
+       "IPY_MODEL_cc94432d08464affa3e58b560bdad194",
+       "IPY_MODEL_3036608c71904ce9ae4bb2a9fa8802d9"
+      ],
+      "layout": "IPY_MODEL_adacfdcc1b0140efac56918e9ccf064e"
+     }
+    },
+    "519a7b154022443db6703f04a9142bae": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_d02274afd47b462291c745f261209d42",
+      "max": 27341251,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_0f417447a7bd4a33acca96fa37aec877",
+      "value": 27341251
+     }
+    },
+    "56e3768bef5a4b9db4168c5c17f509c2": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "590eef89881545aa8bbef9a8bbe7fb00": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "598da69727bd4fb8b1caf465ac736d7a": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "5bdfd87fc6cd4f9dabef7cfee29c8060": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_4d468f96ec924681ad65eb671674b93e",
+      "placeholder": "​",
+      "style": "IPY_MODEL_ad7599de524549c48bf2d3124ad4b299",
+      "value": "Dropping Long Sequences (num_proc=2): 100%"
+     }
+    },
+    "5ca240f31e6b44e3882c5eb37cd5a309": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": "20px"
+     }
+    },
+    "5ca6be24acb548cea130bd58e9954c7c": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "5cea7996f02040b187ece0bb2d6a8d1f": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "5cfb02ee044b4011a378efa8b54a370f": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "5dd7d150dbe04f08b165ce7f2c27cd11": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "5e18768f7ad6434ba8b8b8a2e853e204": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "5e5e15b0569b474c9620083b3ec6af55": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "5e746eb25bbe416fb585fa24e79f5177": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "5eb06edeb58e4930b1affef2a59eae81": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "5f86cd894de94c3280fadc1e2fd0ee13": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_a20927bf5f2c41f58c1e31ac858ab36c",
+       "IPY_MODEL_0a46ad75c198463d843fb35e813642cb",
+       "IPY_MODEL_09007681cf8d42aeb8c1d2f6a74e470a"
+      ],
+      "layout": "IPY_MODEL_ebc80d1a55fa47f4a5ea2756588569ec"
+     }
+    },
+    "60c1a0d765c14a1d888317e6a507e4ea": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "62c028fdef904dedb9cdeca2b3bda725": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "62e1a65582f446a78612eaa804e08a7d": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_5e18768f7ad6434ba8b8b8a2e853e204",
+      "placeholder": "​",
+      "style": "IPY_MODEL_bb33aec33a6447078c31bfd728942994",
+      "value": " 728/728 [00:00&lt;00:00, 20.3kB/s]"
+     }
+    },
+    "62e302ebdad64aada0ffe64ae1c873f3": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "63580b6fb30642479fe3000915bf551a": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "63b4e563e85c4f03b1b72beda9577bcc": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "64f54d4a744a4627a07c3c0120276f3b": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_0546d04aae644dde846c58a4afb598a6",
+      "max": 9985,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_897b77a56c09479bb11d7f2a30997e55",
+      "value": 9985
+     }
+    },
+    "65b75b9b8bc143cf997796af68ff6668": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_81c3db71ac704280ad030072655f1537",
+      "placeholder": "​",
+      "style": "IPY_MODEL_042e091f75694c47aee761e760e76773",
+      "value": " 9985/9985 [00:02&lt;00:00, 3977.47 examples/s]"
+     }
+    },
+    "67da6c4260574869aa24c3cbc1bc1654": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "6932489232ec4ab18a160b1e7fbcdfe1": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "6ebb2ec171414e47a14765505f64bb3c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "6f05e9bebf7b40c9835808e77de6c236": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "PasswordModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "PasswordModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "PasswordView",
+      "continuous_update": true,
+      "description": "Token:",
+      "description_tooltip": null,
+      "disabled": false,
+      "layout": "IPY_MODEL_2e257c8be2da40b4bb67a9e4ab6811f3",
+      "placeholder": "​",
+      "style": "IPY_MODEL_56e3768bef5a4b9db4168c5c17f509c2",
+      "value": ""
+     }
+    },
+    "6f3a28b912714c6e931003549664bfa3": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_5ca240f31e6b44e3882c5eb37cd5a309",
+      "max": 1,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_5eb06edeb58e4930b1affef2a59eae81",
+      "value": 1
+     }
+    },
+    "6f68ed9889f54ad2ae8a3b95ac263a83": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_41f3b32c2f6b4034ae7a3b9124e28bc7",
+      "placeholder": "​",
+      "style": "IPY_MODEL_a10d0a76010f4e508c65a9b69ebc5156",
+      "value": "Tokenizing Prompts (num_proc=2): 100%"
+     }
+    },
+    "704f2f5a9b1c49d5a75a0025a5dda11b": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "71002199df6b40c9a1ac40df5fb27a1b": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "71c8af139cd248b1b51101fd46a93f35": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_d0e9dce55cec4c1ca619a0ccf209d924",
+      "max": 9675,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_4c727d40ef0443449afc31724ee79f0c",
+      "value": 9675
+     }
+    },
+    "734185351eb543fa9a00a881dcbb9fe7": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "735d4f225b24414294fc1b213c61223c": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "742b1030acfd414bbd9d5327b7e3826d": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "77304d1a46b3468a98483e02ec0ac4a4": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "7baeab52d6694c32b1efd1ea1a0a7782": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_93a44a11aa4846fa8efc6c1413ef1627",
+      "placeholder": "​",
+      "style": "IPY_MODEL_a55060adc3564407ac81ad7297d34aaa",
+      "value": "train.jsonl: 100%"
+     }
+    },
+    "7be6f04c284e4326bb4ff3d301e7b3c6": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "danger",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_9503a45960984adc97b58e16c50662e0",
+      "max": 3963750880,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_da6e93f3e4984780b930fe7a706983ea",
+      "value": 3963750502
+     }
+    },
+    "7c2485c6cdfe463da6fdb35982a1070d": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_ad1236893754446881e153adc9d5c962",
+       "IPY_MODEL_daee63fd167e4441a32324b51b00ad2b",
+       "IPY_MODEL_fe41858c6bd04c58840112b67c19a336"
+      ],
+      "layout": "IPY_MODEL_d262c82138024169b9f3aa034ca756fa"
+     }
+    },
+    "7c95f85a2b1f47a1bd846d110c47bb3c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_7fd44cf9ca6e4726bfd7ac21846d6a14",
+      "placeholder": "​",
+      "style": "IPY_MODEL_366a343b62fa47d8985a3bd464d99f9e",
+      "value": "config.json: 100%"
+     }
+    },
+    "7cd0b85ebd204b7aba908417811ce4e0": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_7baeab52d6694c32b1efd1ea1a0a7782",
+       "IPY_MODEL_519a7b154022443db6703f04a9142bae",
+       "IPY_MODEL_d4183e9715f34d249942b8271cca3bdf"
+      ],
+      "layout": "IPY_MODEL_da2347ac94764a3fa2743343cf0d3cd2"
+     }
+    },
+    "7e5d3774060e4589aa65982da5ea4ef4": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "7fd44cf9ca6e4726bfd7ac21846d6a14": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "80366349d81e4dcc892db6cd56e384f3": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_f8ef805b776145c3bfa9ba8d90972058",
+      "max": 9985,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_cc587493c33c4f118d1b1170f85be24c",
+      "value": 9985
+     }
+    },
+    "813621384dc748b0ad06775e22761c0b": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "81c3db71ac704280ad030072655f1537": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "82177df57a494de8900c14c2f5185175": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_67da6c4260574869aa24c3cbc1bc1654",
+      "placeholder": "​",
+      "style": "IPY_MODEL_94b9088614464f60a203de39dbcae853",
+      "value": " 8/8 [01:47&lt;00:00, 11.64s/it]"
+     }
+    },
+    "823f1c78f15043e38bbd4dca3932a86a": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_03a3c744d716431488163b4358b80f92",
+      "max": 239,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_a5434ee714f9498d83870544b67c0cb7",
+      "value": 239
+     }
+    },
+    "835bcc28a5564fb9b3d651bc8e32dc46": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "8640ac440fbc4644b9a3af7ba3ae7183": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "86816687746246b4a6105e8010384e25": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_8640ac440fbc4644b9a3af7ba3ae7183",
+      "placeholder": "​",
+      "style": "IPY_MODEL_5cea7996f02040b187ece0bb2d6a8d1f",
+      "value": "<center> <img\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.svg\nalt='Hugging Face'> <br> Copy a token from <a\nhref=\"https://huggingface.co/settings/tokens\" target=\"_blank\">your Hugging Face\ntokens page</a> and paste it below. <br> Immediately click login after copying\nyour token or it might be stored in plain text in this notebook file. </center>"
+     }
+    },
+    "879c8ab5873847a8833bd74123be90a4": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_ef223e8504b64e3592589880326aaf41",
+      "placeholder": "​",
+      "style": "IPY_MODEL_598da69727bd4fb8b1caf465ac736d7a",
+      "value": " 1.67M/1.67M [00:00&lt;00:00, 19.0MB/s]"
+     }
+    },
+    "897b77a56c09479bb11d7f2a30997e55": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "8bc9d8ba866c442b9118d9630009939c": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "8c4d4fc5a30f4e7cb3be53fe2adda33d": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "8f5bd719974e41c3a8dd9a5b0d3d71e6": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "8f726dbfb45d4528afa33e36a6313267": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "9327977822be4b1294f80e876552e305": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_37de928300e34184881039378bd75e7f",
+      "placeholder": "​",
+      "style": "IPY_MODEL_0e936d9dbf9c4fdd86bbfe9730dedc47",
+      "value": " 3.96G/3.96G [00:13&lt;00:00, 273MB/s]"
+     }
+    },
+    "936d04b5fe1b4c63bf0b080e423d051b": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "93a44a11aa4846fa8efc6c1413ef1627": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "94b9088614464f60a203de39dbcae853": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "9503a45960984adc97b58e16c50662e0": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "95caff42f08a4c2aa14c867b8f37f231": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_de7c37ee83e24f0c889e84d07279c2ec",
+       "IPY_MODEL_9d4897eefb5f48259ffb2d23e332f752",
+       "IPY_MODEL_253017b0d0534e54ab44e181f6d7c82d"
+      ],
+      "layout": "IPY_MODEL_27beaf06e41b472abdb544a43c720c5a"
+     }
+    },
+    "977f799afaac4a55b2dc1cffa7d5b63b": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "97e36007e1304e1583fd81bfb13f0edd": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "9858cb74a09748a39e8149baac96702c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "9b42e08b3c9548818488268768a118b1": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_d955dcaa0e944e719f3a06139dd54a03",
+      "placeholder": "​",
+      "style": "IPY_MODEL_d3de2662c7964f1ba96e58da382af720",
+      "value": "merges.txt: 100%"
+     }
+    },
+    "9cd5211b5d8b457aa0002f1d17b80028": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_6932489232ec4ab18a160b1e7fbcdfe1",
+      "placeholder": "​",
+      "style": "IPY_MODEL_4540927d98f54466b434ba4c0edf045d",
+      "value": "model-00007-of-00008.safetensors: 100%"
+     }
+    },
+    "9d4897eefb5f48259ffb2d23e332f752": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_30a81da86f8043eca301e86a8651201a",
+      "max": 2776833,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_e8b7a81040904c1e89e58978223b1737",
+      "value": 2776833
+     }
+    },
+    "9e333ed3b5014069ac1dd969255dd591": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "9ed02dc43412471a9ab47f3620ccf3a5": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "9f1c9a0695384bdaa6f8b847ef89bee8": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ButtonStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ButtonStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "button_color": null,
+      "font_weight": ""
+     }
+    },
+    "9f56a2d9979c4bd8928c644c22c3ecdf": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "a0a11e929edd4189b79723d618522c33": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "a10d0a76010f4e508c65a9b69ebc5156": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "a138859f19b74fc0928dc236ab5359db": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_9b42e08b3c9548818488268768a118b1",
+       "IPY_MODEL_12b56912736849fea2ad8124456fdc5c",
+       "IPY_MODEL_879c8ab5873847a8833bd74123be90a4"
+      ],
+      "layout": "IPY_MODEL_20352e5f58d24bb8b1f3940efd14fe4a"
+     }
+    },
+    "a1959759c5424da9961fb2a308d4dee4": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_3aaecbf540f54a2db9ab0931e3b1fe57",
+      "placeholder": "​",
+      "style": "IPY_MODEL_9e333ed3b5014069ac1dd969255dd591",
+      "value": " 239/239 [00:00&lt;00:00, 30.9kB/s]"
+     }
+    },
+    "a20927bf5f2c41f58c1e31ac858ab36c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_1811cda0644e4190a9469d1774435d82",
+      "placeholder": "​",
+      "style": "IPY_MODEL_35c811d2ae8e43f3b5cecbdd3cfa857f",
+      "value": "tokenizer.json: 100%"
+     }
+    },
+    "a3a945817f684328b34651fe052393ec": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "a44f630e099e43899f20a77084ae60cd": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_ed5ca967ad5342929e578ac6aa4dc4c0",
+      "placeholder": "​",
+      "style": "IPY_MODEL_af401d117d5047629d3a6e2361757b62",
+      "value": "model-00001-of-00008.safetensors: 100%"
+     }
+    },
+    "a4e5789584564049b83df7c6c54a3e08": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "a5434ee714f9498d83870544b67c0cb7": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "a55060adc3564407ac81ad7297d34aaa": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "a6f48410b9964fefba0c3009a77dc838": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "a7cf477e80fc43e0ad82c7997b076dce": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "a80410b919e442c49aea15acc1ce1a72": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_fa1282ccc7544e4f818e2f03ccffe4a5",
+      "placeholder": "​",
+      "style": "IPY_MODEL_bbbf575d2a4b4c6ea8389be79b2a6039",
+      "value": "model.safetensors.index.json: 100%"
+     }
+    },
+    "ab93eabd7cea4b94b4b7a387f101e8a1": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "ac764024cf1c4e08ba7749afd2cd20ac": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "ad1236893754446881e153adc9d5c962": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_62e302ebdad64aada0ffe64ae1c873f3",
+      "placeholder": "​",
+      "style": "IPY_MODEL_bd1b0dfed6d34d16af33a4a58330f5ec",
+      "value": "Saving the dataset (1/1 shards): 100%"
+     }
+    },
+    "ad7599de524549c48bf2d3124ad4b299": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "adacfdcc1b0140efac56918e9ccf064e": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "af401d117d5047629d3a6e2361757b62": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "b191ac001a2e4962bc9a245fcdf26e6b": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "b195f160ca20442fadd8b5aed0ee41af": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "b1bea589efa14258a9982071b87938bf": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "b5b65414154544aa8a71b1a39164aad7": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "b634bb73cfa743d09a5999101b840976": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "b82aa8c57f7c422a9a9c90f333ed2a99": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_c0991cf63ee6458b96e9a75e7a88b61a",
+       "IPY_MODEL_71c8af139cd248b1b51101fd46a93f35",
+       "IPY_MODEL_1d5117195d4b49eb8f1a73b18419f7ce"
+      ],
+      "layout": "IPY_MODEL_3c21e4a511b4441192c03b7f1d0976e9"
+     }
+    },
+    "b8766a88716948cf968f4563531a76d9": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_2b3a2659b12244bd8548320320016dbf",
+      "placeholder": "​",
+      "style": "IPY_MODEL_0cd7efffbb3c4c4b972e63749f61ab97",
+      "value": "Generating train split: "
+     }
+    },
+    "b87c84de30e84b3abf4871461fb9cbd3": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "b8e39e4dddc3497fbc29ae45c66da759": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "bb33aec33a6447078c31bfd728942994": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "bbbf575d2a4b4c6ea8389be79b2a6039": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "bca2c7185b6749fd899c06a2ba4c5e46": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_0f480e3a0b0a45d2a2d2dec3cad923f3",
+      "placeholder": "​",
+      "style": "IPY_MODEL_fcb30372e7404c5d8a1ad4df91e6c7b2",
+      "value": " 1.91G/1.91G [00:05&lt;00:00, 444MB/s]"
+     }
+    },
+    "bd1b0dfed6d34d16af33a4a58330f5ec": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "be724f04b03942b2a033a7e8898bb4fd": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "bed8726b8069434687c75452e21f19e5": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_fa864b41586f4a7aa56aeafd1d84eb75",
+      "max": 9985,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_3225603166b54e7aab766b9964a2f660",
+      "value": 9985
+     }
+    },
+    "bee3501b2a17427784a717e50a85e7fa": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "bfcdbba993b74972a9e3e575f86908ff": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "bff139df987d4a62abec6456cb27f3d4": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "danger",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_c1f9c267ba3f40039cdb5eb3267e8043",
+      "max": 3963750880,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_33b3b1d0295646edaac7b4822761aeb0",
+      "value": 3963750502
+     }
+    },
+    "c0892a1881de4eb4bfabc6a68f87ae99": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_158c8b85dbf34de6a94b4e35e2fc7d5a",
+      "placeholder": "​",
+      "style": "IPY_MODEL_0b4c9753a7cb4354b8e5f187e6e1ad7c",
+      "value": " 3.96G/3.96G [00:15&lt;00:00, 564MB/s]"
+     }
+    },
+    "c0991cf63ee6458b96e9a75e7a88b61a": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_ed28e2e0410d4e0b855467e798e53d66",
+      "placeholder": "​",
+      "style": "IPY_MODEL_d93f134f802b4b69b575bdaf07dbd27c",
+      "value": "tokenizer_config.json: 100%"
+     }
+    },
+    "c12ea43372ac4d57bb9605f1a429b397": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "VBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "VBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "VBoxView",
+      "box_style": "",
+      "children": [],
+      "layout": "IPY_MODEL_131065f118274a1586ac38e39ed84ef0"
+     }
+    },
+    "c1314f241a434c41b45d84dc4d3b30f8": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "c1f9c267ba3f40039cdb5eb3267e8043": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "c33ced495f70464aa4a3a91922090853": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "c3725c7f79fe415fbd1ea336f0cc9cf1": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "danger",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_b191ac001a2e4962bc9a245fcdf26e6b",
+      "max": 3841788544,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_054c8dffadba48c6b895a6cc62448ecc",
+      "value": 3841788178
+     }
+    },
+    "c3be9109d63c485d9c0ef4f9bc0f9218": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "c42acf646f344a88b8c11f81e67f7206": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_8bc9d8ba866c442b9118d9630009939c",
+      "placeholder": "​",
+      "style": "IPY_MODEL_9f56a2d9979c4bd8928c644c22c3ecdf",
+      "value": "model-00003-of-00008.safetensors: 100%"
+     }
+    },
+    "c6164e05a1914ae48083db9ad7f4ef7c": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "c65dc74c7d6f4bab8f7dd28455161dd8": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "c6e00f5224364822bc4239b176686919": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_2a51b36be41745468e4c2d7a21b1c0d2",
+      "max": 36514,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_4fd114abe9f5494ab59858949f5055f1",
+      "value": 36514
+     }
+    },
+    "c73055099c084dca996159e23e162d0b": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_e40d1c1ac9494b3bade9858324e7ffdf",
+      "placeholder": "​",
+      "style": "IPY_MODEL_d65b6b060d9845779299491ac5599c31",
+      "value": " 9985/9985 [01:04&lt;00:00, 189.08 examples/s]"
+     }
+    },
+    "c7433acd3c4841e6958ae8f7e87b1808": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "CheckboxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "CheckboxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "CheckboxView",
+      "description": "Add token as git credential?",
+      "description_tooltip": null,
+      "disabled": false,
+      "indent": true,
+      "layout": "IPY_MODEL_62c028fdef904dedb9cdeca2b3bda725",
+      "style": "IPY_MODEL_a7cf477e80fc43e0ad82c7997b076dce",
+      "value": false
+     }
+    },
+    "c84cc07789be48aebb322c23d355289e": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_0077aedc3d174560bce924ee89e9c006",
+      "placeholder": "​",
+      "style": "IPY_MODEL_00321cce58884f6f9b3855a21fcd9187",
+      "value": "Add position_id column (Sample Packing) (num_proc=2): 100%"
+     }
+    },
+    "ca65e32eb52f48c09a84b33cb18f22cd": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "cc587493c33c4f118d1b1170f85be24c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "cc94432d08464affa3e58b560bdad194": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "danger",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_b5b65414154544aa8a71b1a39164aad7",
+      "max": 3963750816,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_f0a58fbd0fca4340890041f99fa2f8c8",
+      "value": 3963750438
+     }
+    },
+    "ccfcdc95baf646f8aeb3d516742383f2": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "cdebbc55a1164c018546c2ac6f8c620c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_a44f630e099e43899f20a77084ae60cd",
+       "IPY_MODEL_c3725c7f79fe415fbd1ea336f0cc9cf1",
+       "IPY_MODEL_0e50870ed0c643e0b6c18cc5d7ddae7f"
+      ],
+      "layout": "IPY_MODEL_c33ced495f70464aa4a3a91922090853"
+     }
+    },
+    "d02274afd47b462291c745f261209d42": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "d07c8b97d3314f1c852e44bdd40f61ed": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "d0e9dce55cec4c1ca619a0ccf209d924": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "d1f9b10c130542f094c8fd3d1e23b5e9": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "d262c82138024169b9f3aa034ca756fa": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "d3de2662c7964f1ba96e58da382af720": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "d4183e9715f34d249942b8271cca3bdf": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_63580b6fb30642479fe3000915bf551a",
+      "placeholder": "​",
+      "style": "IPY_MODEL_8f726dbfb45d4528afa33e36a6313267",
+      "value": " 27.3M/27.3M [00:00&lt;00:00, 31.0MB/s]"
+     }
+    },
+    "d43c6df07ddb466587807d6dbe1ff614": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_8c4d4fc5a30f4e7cb3be53fe2adda33d",
+      "placeholder": "​",
+      "style": "IPY_MODEL_e90658f4bcb642baa78426012f863152",
+      "value": "model-00004-of-00008.safetensors: 100%"
+     }
+    },
+    "d65b6b060d9845779299491ac5599c31": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "d6fe74e4255444368f8f90a62157d869": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "d93f134f802b4b69b575bdaf07dbd27c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "d955dcaa0e944e719f3a06139dd54a03": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "da2347ac94764a3fa2743343cf0d3cd2": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "da6e93f3e4984780b930fe7a706983ea": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "daee63fd167e4441a32324b51b00ad2b": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_d07c8b97d3314f1c852e44bdd40f61ed",
+      "max": 9985,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_ebb69a2c3d0a4299a484698287b3087c",
+      "value": 9985
+     }
+    },
+    "dc892a596f6942d7973c616c38f0eebb": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_c84cc07789be48aebb322c23d355289e",
+       "IPY_MODEL_bed8726b8069434687c75452e21f19e5",
+       "IPY_MODEL_16a188a0b06d45f980dcf3933509fe0a"
+      ],
+      "layout": "IPY_MODEL_60c1a0d765c14a1d888317e6a507e4ea"
+     }
+    },
+    "dd0e646fad3f4a89ba23b39d162bd8d9": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_d43c6df07ddb466587807d6dbe1ff614",
+       "IPY_MODEL_e0e8b840b8ea4d0d9db09afe99fa287d",
+       "IPY_MODEL_9327977822be4b1294f80e876552e305"
+      ],
+      "layout": "IPY_MODEL_77304d1a46b3468a98483e02ec0ac4a4"
+     }
+    },
+    "de7c37ee83e24f0c889e84d07279c2ec": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_34cf3df51fbc41cabfdbba153c007f0e",
+      "placeholder": "​",
+      "style": "IPY_MODEL_ac764024cf1c4e08ba7749afd2cd20ac",
+      "value": "vocab.json: 100%"
+     }
+    },
+    "dfd2a2649b8341ef913207526708aff1": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "e09f1bcbb9d94c09be53e5e1303642c2": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_e7d8e4fe58384e93a106de546068c65e",
+      "max": 8,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_0aa8ab56b85f4171a79c3bc210594025",
+      "value": 8
+     }
+    },
+    "e0e8b840b8ea4d0d9db09afe99fa287d": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "danger",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_f7434f3e03124a1c938a39af79d7fa59",
+      "max": 3963750880,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_c1314f241a434c41b45d84dc4d3b30f8",
+      "value": 3963750502
+     }
+    },
+    "e21e180307e5485cbbe908672fd6639a": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_2e2b0c1599c341a198f632f46a40c90e",
+       "IPY_MODEL_bff139df987d4a62abec6456cb27f3d4",
+       "IPY_MODEL_ebe1cc366d324ad59b264c8b3c431441"
+      ],
+      "layout": "IPY_MODEL_114dece49dba437c8572ef94b23c3b1e"
+     }
+    },
+    "e366ae3fceec4566b9ed303d6c5f90af": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "e3fb3fc6afe04b3c9b7ac61809ce78fa": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_c6164e05a1914ae48083db9ad7f4ef7c",
+      "placeholder": "​",
+      "style": "IPY_MODEL_813621384dc748b0ad06775e22761c0b",
+      "value": " 9985/9985 [00:03&lt;00:00, 3622.89 examples/s]"
+     }
+    },
+    "e400cbf14bcc446a9d33b210cd93550b": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "e40d1c1ac9494b3bade9858324e7ffdf": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "e575d87a7efe4ec7b1efde489839d4a6": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "e5a82df528bb4e408797a3b6c2758f4a": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "e6e969610738449887259063967f82b0": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "e7d8e4fe58384e93a106de546068c65e": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "e87ea87fcff247b5bbcc331ba79a8dc2": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "e8b7a81040904c1e89e58978223b1737": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "e90658f4bcb642baa78426012f863152": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "eb1c9535e6a546098b760528b2ea387c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_18357b321ce44d7b8bd9d1c886f69275",
+       "IPY_MODEL_279937fe03bc4e4eb25b472d7e9df163",
+       "IPY_MODEL_bca2c7185b6749fd899c06a2ba4c5e46"
+      ],
+      "layout": "IPY_MODEL_1f7d30f71bbd4547a9150d21da071055"
+     }
+    },
+    "ebb69a2c3d0a4299a484698287b3087c": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "ebc80d1a55fa47f4a5ea2756588569ec": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "ebe1cc366d324ad59b264c8b3c431441": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_fba7aa824b38467ab3061b226114cdec",
+      "placeholder": "​",
+      "style": "IPY_MODEL_f3075dccbd2747b4a7913b66f44f2596",
+      "value": " 3.96G/3.96G [00:13&lt;00:00, 398MB/s]"
+     }
+    },
+    "ec030fc3c346426f9abc3a89892258d3": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "success",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_dfd2a2649b8341ef913207526708aff1",
+      "max": 9985,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_4f1977d7e4824ef1a14b65f0f42bba10",
+      "value": 9985
+     }
+    },
+    "ec11d1e5ae7b42c883d9b1f38a65356e": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_936d04b5fe1b4c63bf0b080e423d051b",
+      "placeholder": "​",
+      "style": "IPY_MODEL_f1cef8e8dc2646fb9fd09f3b09081074",
+      "value": " 36.5k/36.5k [00:00&lt;00:00, 4.32MB/s]"
+     }
+    },
+    "ed28e2e0410d4e0b855467e798e53d66": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "ed5ca967ad5342929e578ac6aa4dc4c0": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "edc99591b9c747b689b94d0052fec14c": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "ef0a3c7a6f14460fb4da096928ae249e": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_07fb3a2c8315494e97b447e672dfae06",
+       "IPY_MODEL_ec030fc3c346426f9abc3a89892258d3",
+       "IPY_MODEL_e3fb3fc6afe04b3c9b7ac61809ce78fa"
+      ],
+      "layout": "IPY_MODEL_c3be9109d63c485d9c0ef4f9bc0f9218"
+     }
+    },
+    "ef223e8504b64e3592589880326aaf41": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "f0a58fbd0fca4340890041f99fa2f8c8": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "ProgressStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "ProgressStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "bar_color": null,
+      "description_width": ""
+     }
+    },
+    "f113ebd8c1c34806bea4dd7ed3035173": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "f1cef8e8dc2646fb9fd09f3b09081074": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "f3075dccbd2747b4a7913b66f44f2596": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "f365820a3d3c42b2948abfe32065de14": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_735d4f225b24414294fc1b213c61223c",
+      "placeholder": "​",
+      "style": "IPY_MODEL_5e5e15b0569b474c9620083b3ec6af55",
+      "value": "generation_config.json: 100%"
+     }
+    },
+    "f4667818b9d34a09891cd727a429a610": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_4b27c267393640f28f6eae0875bd2ed9",
+      "placeholder": "​",
+      "style": "IPY_MODEL_9858cb74a09748a39e8149baac96702c",
+      "value": " 3.96G/3.96G [00:11&lt;00:00, 457MB/s]"
+     }
+    },
+    "f4a1795dc7514a718f478245f521f0ba": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "f60a2bdb6b6b4e0e8c3508580e247132": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "FloatProgressModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "FloatProgressModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "ProgressView",
+      "bar_style": "danger",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_edc99591b9c747b689b94d0052fec14c",
+      "max": 3963750880,
+      "min": 0,
+      "orientation": "horizontal",
+      "style": "IPY_MODEL_35cc989ca3374e7dba0cb166febc4bde",
+      "value": 3963750502
+     }
+    },
+    "f7434f3e03124a1c938a39af79d7fa59": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "f8ef805b776145c3bfa9ba8d90972058": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "fa1282ccc7544e4f818e2f03ccffe4a5": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "fa864b41586f4a7aa56aeafd1d84eb75": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "fba7aa824b38467ab3061b226114cdec": {
+     "model_module": "@jupyter-widgets/base",
+     "model_module_version": "1.2.0",
+     "model_name": "LayoutModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/base",
+      "_model_module_version": "1.2.0",
+      "_model_name": "LayoutModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "LayoutView",
+      "align_content": null,
+      "align_items": null,
+      "align_self": null,
+      "border": null,
+      "bottom": null,
+      "display": null,
+      "flex": null,
+      "flex_flow": null,
+      "grid_area": null,
+      "grid_auto_columns": null,
+      "grid_auto_flow": null,
+      "grid_auto_rows": null,
+      "grid_column": null,
+      "grid_gap": null,
+      "grid_row": null,
+      "grid_template_areas": null,
+      "grid_template_columns": null,
+      "grid_template_rows": null,
+      "height": null,
+      "justify_content": null,
+      "justify_items": null,
+      "left": null,
+      "margin": null,
+      "max_height": null,
+      "max_width": null,
+      "min_height": null,
+      "min_width": null,
+      "object_fit": null,
+      "object_position": null,
+      "order": null,
+      "overflow": null,
+      "overflow_x": null,
+      "overflow_y": null,
+      "padding": null,
+      "right": null,
+      "top": null,
+      "visibility": null,
+      "width": null
+     }
+    },
+    "fcb30372e7404c5d8a1ad4df91e6c7b2": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "fcbab4d8dced41a18dfccce81e3a45a0": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "fd4f333f7ece4450b04e1a9af1f9d2f6": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_d1f9b10c130542f094c8fd3d1e23b5e9",
+      "placeholder": "​",
+      "style": "IPY_MODEL_e575d87a7efe4ec7b1efde489839d4a6",
+      "value": "model-00006-of-00008.safetensors: 100%"
+     }
+    },
+    "fe18bba7f3fb4c31bf840541f36b3425": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_fd4f333f7ece4450b04e1a9af1f9d2f6",
+       "IPY_MODEL_f60a2bdb6b6b4e0e8c3508580e247132",
+       "IPY_MODEL_c0892a1881de4eb4bfabc6a68f87ae99"
+      ],
+      "layout": "IPY_MODEL_1bec6297c90242a88672d195bc09d429"
+     }
+    },
+    "fe41858c6bd04c58840112b67c19a336": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_e5a82df528bb4e408797a3b6c2758f4a",
+      "placeholder": "​",
+      "style": "IPY_MODEL_f113ebd8c1c34806bea4dd7ed3035173",
+      "value": " 9985/9985 [00:00&lt;00:00, 44264.88 examples/s]"
+     }
+    },
+    "fea1b70fb46745feb5111b3929175b5d": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HBoxModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HBoxModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HBoxView",
+      "box_style": "",
+      "children": [
+       "IPY_MODEL_f365820a3d3c42b2948abfe32065de14",
+       "IPY_MODEL_823f1c78f15043e38bbd4dca3932a86a",
+       "IPY_MODEL_a1959759c5424da9961fb2a308d4dee4"
+      ],
+      "layout": "IPY_MODEL_34c9c0137b504cd799c6bd6de69507c2"
+     }
+    },
+    "ff3a94b146a948b6907f5d80c7157f99": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "DescriptionStyleModel",
+     "state": {
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "DescriptionStyleModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/base",
+      "_view_module_version": "1.2.0",
+      "_view_name": "StyleView",
+      "description_width": ""
+     }
+    },
+    "ffdbb12a2f2c4d14911685e7683e0ef0": {
+     "model_module": "@jupyter-widgets/controls",
+     "model_module_version": "1.5.0",
+     "model_name": "HTMLModel",
+     "state": {
+      "_dom_classes": [],
+      "_model_module": "@jupyter-widgets/controls",
+      "_model_module_version": "1.5.0",
+      "_model_name": "HTMLModel",
+      "_view_count": null,
+      "_view_module": "@jupyter-widgets/controls",
+      "_view_module_version": "1.5.0",
+      "_view_name": "HTMLView",
+      "description": "",
+      "description_tooltip": null,
+      "layout": "IPY_MODEL_ab93eabd7cea4b94b4b7a387f101e8a1",
+      "placeholder": "​",
+      "style": "IPY_MODEL_704f2f5a9b1c49d5a75a0025a5dda11b",
+      "value": " 3.96G/3.96G [00:12&lt;00:00, 656MB/s]"
+     }
+    }
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
 }
diff --git a/examples/deepcogito/cogito-v1-preview-llama-3B-lora.yml b/examples/deepcogito/cogito-v1-preview-llama-3B-lora.yml
new file mode 100644
index 0000000000..c36b0e74a1
--- /dev/null
+++ b/examples/deepcogito/cogito-v1-preview-llama-3B-lora.yml
@@ -0,0 +1,56 @@
+base_model: deepcogito/cogito-v1-preview-llama-3B
+# 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: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/lora-out
+
+sequence_len: 4096
+sample_packing: true
+eval_sample_packing: false
+
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/deepcogito/cogito-v1-preview-qwen-14B-lora.yml b/examples/deepcogito/cogito-v1-preview-qwen-14B-lora.yml
new file mode 100644
index 0000000000..2b2aafd754
--- /dev/null
+++ b/examples/deepcogito/cogito-v1-preview-qwen-14B-lora.yml
@@ -0,0 +1,56 @@
+base_model: deepcogito/cogito-v1-preview-qwen-14B
+# 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: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/lora-out
+
+sequence_len: 4096
+sample_packing: true
+eval_sample_packing: false
+
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/deepseek-v2/fft-fsdp-16b.yaml b/examples/deepseek-v2/fft-fsdp-16b.yaml
new file mode 100644
index 0000000000..2eac9aea3f
--- /dev/null
+++ b/examples/deepseek-v2/fft-fsdp-16b.yaml
@@ -0,0 +1,59 @@
+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
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+sequence_len: 2048
+sample_packing: true
+
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 8
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 2
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: DeepseekV2DecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/deepseek-v2/qlora-fsdp-2_5.yaml b/examples/deepseek-v2/qlora-fsdp-2_5.yaml
new file mode 100644
index 0000000000..0e23a0266f
--- /dev/null
+++ b/examples/deepseek-v2/qlora-fsdp-2_5.yaml
@@ -0,0 +1,83 @@
+base_model: axolotl-quants/DeepSeek-V2.5-bnb-nf4-bf16
+# 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: true
+
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+liger_rms_norm: true
+liger_glu_activation: true
+liger_fused_linear_cross_entropy: true
+
+chat_template: deepseek_v2
+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
+
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+adapter: qlora
+lora_r: 256
+lora_alpha: 256
+lora_target_linear: true
+peft_use_rslora: true
+
+gradient_accumulation_steps: 1
+micro_batch_size: 8
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 2
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: DeepseekV2DecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/deepseek-v4/README.md b/examples/deepseek-v4/README.md
new file mode 100644
index 0000000000..664013f3ac
--- /dev/null
+++ b/examples/deepseek-v4/README.md
@@ -0,0 +1,43 @@
+# Finetune DeepSeek-V4-Flash with Axolotl
+
+[DeepSeek-V4-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash) is a sparse MoE model with NVFP4 experts and head_dim 512 eager-only attention.
+
+This guide trains the MoE experts (LoRA on the 3D expert parameters) on the NVFP4 checkpoint [nvidia/DeepSeek-V4-Flash-NVFP4](https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4) (~168GB).
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy).
+
+3. Run the finetuning example:
+
+```bash
+# 2xB200 (sm100): ~140GB/GPU, ~130 tok/s, train_loss ~1.09
+axolotl train examples/deepseek-v4/v4-flash-nvfp4-lora.yaml
+```
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- For single GPU, remove FSDP block. It would require >=180GB GPU.
+- On cloud, if using volume mount, pointing the HF cache (via `HF_HOME`) at a local disk keeps weight load fast.
+- Train the experts only. Do not add attention or module LoRA on a `use_dsv4_kernels` run: it is unsupported and breaks the experts-only FSDP2 invariant (data-independent backward collectives across ranks).
+- Keep `attn_implementation: eager` so the dsv4 kernels plugin owns the attention path.
+- `sample_packing` attends within the sliding window across packed documents, exactly as plain eager attention would. Drop packing if your data needs strict cross-document isolation.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Related Resources
+
+- [DeepSeek-V4-Flash on HuggingFace](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/deepseek-v4/v4-flash-nvfp4-lora.yaml b/examples/deepseek-v4/v4-flash-nvfp4-lora.yaml
new file mode 100644
index 0000000000..30c40896a7
--- /dev/null
+++ b/examples/deepseek-v4/v4-flash-nvfp4-lora.yaml
@@ -0,0 +1,58 @@
+# DeepSeek-V4-Flash MoE LoRA on an NVFP4 (modelopt) checkpoint via ScatterMoE. FSDP2 multi-GPU only.
+base_model: nvidia/DeepSeek-V4-Flash-NVFP4
+attn_implementation: eager  # required: the dsv4 kernels own the attention path (V4-Flash is eager-only)
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+
+use_kernels: true
+use_scattermoe: true
+use_dsv4_kernels: true
+lora_mlp_kernel: true
+dsv4_fp4_grouped_mode: nvfp4
+strict: false
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+val_set_size: 0.0
+output_dir: ./outputs/dsv4-flash-nvfp4-lora
+sequence_len: 4096
+sample_packing: true  # required: variable-length steps re-autotune the grouped GEMM and throughput collapses
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.0
+lora_target_parameters:
+  - mlp.experts.gate_up_proj
+  - mlp.experts.down_proj
+
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+num_epochs: 1
+optimizer: adamw_torch
+lr_scheduler: constant
+learning_rate: 0.0001
+bf16: true
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+activation_offloading: false
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  state_dict_type: SHARDED_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: DeepseekV4DecoderLayer
+  reshard_after_forward: true
+
+logging_steps: 1
+warmup_ratio: 0.1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/devstral/README.md b/examples/devstral/README.md
new file mode 100644
index 0000000000..9fbe1d4fcb
--- /dev/null
+++ b/examples/devstral/README.md
@@ -0,0 +1,72 @@
+# Finetune Devstral with Axolotl
+
+Devstral Small is a 24B parameter opensource model from MistralAI found on HuggingFace [Devstral-Small-2505](https://huggingface.co/mistralai/Devstral-Small-2505) and [Devstral-Small-2507](https://huggingface.co/mistralai/Devstral-Small-2507). `Devstral-Small-2507` is the latest version of the model and has [function calling](https://mistralai.github.io/mistral-common/usage/tools/) support.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations with proper masking.
+
+The model was fine-tuned ontop of [Mistral-Small-3.1](https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Base-2503) without the vision layer and has a context of up to 128k tokens.
+
+Thanks to the team at MistralAI for giving us early access to prepare for this release.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+uv pip install --no-build-isolation 'axolotl>=0.16.1'
+```
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage
+
+```bash
+python scripts/cutcrossentropy_install.py | sh
+```
+
+3. Run the finetuning example:
+
+```bash
+axolotl train examples/devstral/devstral-small-qlora.yml
+```
+
+This config uses about 21GB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+- Learn how to use function calling with Axolotl at [docs](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#using-tool-use).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+- [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy)
+- [Liger Kernel](https://docs.axolotl.ai/docs/custom_integrations.html#liger-kernels)
+
+## Limitations
+
+We only support the `mistral-common` tokenizer for Supervised Fine-tuning at the moment and for `type: chat_template` only.
+
+In addition, we do not support overriding tokens yet.
+
+## Related Resources
+
+- [MistralAI Devstral Blog](https://mistral.ai/news/devstral)
+- [MistralAI Devstral 1.1 Blog](https://mistral.ai/news/devstral-2507)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
+
+
+## Future Work
+
+- Add parity to Preference Tuning, RL, Multi-modal, etc.
+- Add parity to other tokenizer configs like overriding tokens.
diff --git a/examples/devstral/devstral-small-qlora.yml b/examples/devstral/devstral-small-qlora.yml
new file mode 100644
index 0000000000..6ee0e014d2
--- /dev/null
+++ b/examples/devstral/devstral-small-qlora.yml
@@ -0,0 +1,66 @@
+base_model: mistralai/Devstral-Small-2507
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+load_in_8bit: false
+load_in_4bit: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/qlora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_linear: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_torch
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+# scaling_softmax: true  # needs flex_attention
+
+loss_watchdog_threshold: 5.0
+loss_watchdog_patience: 3
+
+warmup_ratio: 0.05
+evals_per_epoch: 4
+saves_per_epoch: 1
+
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/distributed-parallel/README.md b/examples/distributed-parallel/README.md
new file mode 100644
index 0000000000..ad7c48d5fb
--- /dev/null
+++ b/examples/distributed-parallel/README.md
@@ -0,0 +1,52 @@
+# ND Parallelism Examples
+
+This directory contains example configurations for training models using ND Parallelism in Axolotl. These examples demonstrate how to compose different parallelism strategies (FSDP, TP, CP, HSDP) for efficient multi-GPU training.
+
+## Quick Start
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Run the command below:
+
+```bash
+# Train Qwen3 8B with FSDP + TP + CP on a single 8-GPU node
+axolotl train examples/distributed-parallel/qwen3-8b-fsdp-tp-cp.yaml
+
+# Train Llama 3.1 8B with HSDP + TP on 2 nodes (16 GPUs total)
+axolotl train examples/distributed-parallel/llama-3_1-8b-hsdp-tp.yaml
+```
+
+## Example Configurations
+
+### Single Node (8 GPUs)
+
+**Qwen3 8B with FSDP + TP + CP** ([qwen3-8b-fsdp-tp-cp.yaml](./qwen3-8b-fsdp-tp-cp.yaml))
+- Uses all 3 parallelism dimensions on a single node
+- Ideal for: when model weights, activations, and/or context are too large to fit on single GPU
+
+```yaml
+dp_shard_size: 2         # FSDP across 2 GPUs
+tensor_parallel_size: 2  # TP across 2 GPUs
+context_parallel_size: 2 # CP across 2 GPUs
+# Total: 2 × 2 × 2 = 8 GPUs
+```
+
+### Multi-Node
+
+**Llama 3.1 8B with HSDP + TP** ([llama-3_1-8b-hsdp-tp.yaml](./llama-3_1-8b-hsdp-tp.yaml))
+- FSDP & TP within nodes, DDP across nodes to minimize inter-node communication
+- Ideal for: Scaling to multiple nodes while maintaining training efficiency
+
+```yaml
+dp_shard_size: 4        # FSDP within each 4-GPU group
+tensor_parallel_size: 2 # TP within each node
+dp_replicate_size: 2    # DDP across 2 groups
+# Total: (4 × 2) × 2 = 16 GPUs (2 nodes)
+```
+
+## Learn More
+
+- [ND Parallelism Documentation](https://docs.axolotl.ai/docs/nd_parallelism.html)
+- [Blog: Accelerate ND-Parallel Guide](https://huggingface.co/blog/accelerate-nd-parallel)
+- [Multi-GPU Training Guide](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/distributed-parallel/llama-3_1-8b-hsdp-tp.yaml b/examples/distributed-parallel/llama-3_1-8b-hsdp-tp.yaml
new file mode 100644
index 0000000000..a99a6bef8e
--- /dev/null
+++ b/examples/distributed-parallel/llama-3_1-8b-hsdp-tp.yaml
@@ -0,0 +1,47 @@
+base_model: meta-llama/Llama-3.1-8B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+dp_shard_size: 4
+dp_replicate_size: 2
+tensor_parallel_size: 2
+# context_parallel_size: 2
+
+dataset_prepared_path: last_run_prepared
+
+special_tokens:
+  pad_token: <|end_of_text|>
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  state_dict_type: FULL_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  reshard_after_forward: true
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+
+output_dir: ./outputs/ndp-out/
+
+sequence_len: 2048
+sample_packing: true
+attn_implementation: flash_attention_2
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 2
+optimizer: adamw_torch_fused
+lr_scheduler: constant_with_warmup
+learning_rate: 2e-6
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
diff --git a/examples/distributed-parallel/qwen3-8b-fsdp-tp-cp.yaml b/examples/distributed-parallel/qwen3-8b-fsdp-tp-cp.yaml
new file mode 100644
index 0000000000..a12b524edd
--- /dev/null
+++ b/examples/distributed-parallel/qwen3-8b-fsdp-tp-cp.yaml
@@ -0,0 +1,46 @@
+base_model: Qwen/Qwen3-8B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+dp_shard_size: 2
+# dp_replicate_size: 1
+context_parallel_size: 2
+tensor_parallel_size: 2
+
+dataset_prepared_path: last_run_prepared
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  state_dict_type: FULL_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen3DecoderLayer
+  reshard_after_forward: true
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+
+output_dir: ./outputs/ndp-out/
+
+sequence_len: 8192
+sample_packing: true
+attn_implementation: flash_attention_2
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1  # must be 1 when using context parallel
+num_epochs: 2
+optimizer: adamw_torch_fused
+lr_scheduler: constant_with_warmup
+learning_rate: 2e-6
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+
+special_tokens:
diff --git a/examples/eaft/eaft-example.yml b/examples/eaft/eaft-example.yml
new file mode 100644
index 0000000000..b4b13a14c5
--- /dev/null
+++ b/examples/eaft/eaft-example.yml
@@ -0,0 +1,76 @@
+base_model: google/gemma-3-1b-it
+
+model_type: Gemma3ForCausalLM
+cls_model_config: Gemma3TextConfig
+
+# gemma3 doesn't seem to play nice with ddp
+ddp_find_unused_parameters: true
+
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+datasets:
+  - path: cgato/SlimOrcaDedupCleaned
+    type: chat_template
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+dataset_prepared_path:
+val_set_size: 0
+output_dir: ./outputs/eaft-gemma-3-1b
+
+use_eaft: true
+eaft_alpha: 1.0
+eaft_k: 20
+
+sequence_len: 1024
+sample_packing: false
+
+adapter:
+lora_model_dir:
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+eval_batch_size: 1
+max_steps: 1000
+evaluation_strategy: "no"
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 5e-5
+
+train_on_inputs: false
+group_by_length: false
+bf16: auto
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+debug:
+deepspeed:
+fsdp:
+fsdp_config:
+special_tokens:
diff --git a/examples/ebft/README.md b/examples/ebft/README.md
new file mode 100644
index 0000000000..24f2c582d4
--- /dev/null
+++ b/examples/ebft/README.md
@@ -0,0 +1,211 @@
+# Energy-Based Fine-Tuning (EBFT)
+
+EBFT is an integration of ["Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models"](https://arxiv.org/abs/2603.12248) (Jelassi et al., 2026) into axolotl.
+
+## Overview
+
+EBFT fine-tunes language models by optimizing a **feature-matching loss** rather than relying on external reward functions or verifiers. A frozen copy of the model (the "feature network") extracts embeddings from both generated and ground-truth completions, and the generator is updated via REINFORCE to match the ground-truth feature moments.
+
+**Key advantages over SFT:**
+- Operates on model rollouts (not teacher forcing), reducing distribution shift
+- Provides dense sequence-level supervision without a task-specific verifier
+- Improves both downstream accuracy and validation cross-entropy simultaneously
+
+**Key advantages over RLVR:**
+- No reward model or verifier required — works on any (prompt, completion) data
+- Applicable to non-verifiable tasks (e.g., raw code, translation, creative writing)
+- Maintains distributional calibration (low feature-matching loss)
+
+## Two Modes
+
+EBFT supports two modes depending on your data format:
+
+### Structured Mode (`mode: structured`, default)
+For **QA/instruction data** with prompt + completion pairs (e.g., OpenCodeInstruct, ALMA translation).
+- Extends GRPOTrainer — uses vLLM for fast rollout generation
+- RLOO advantages and clipped policy gradient from GRPO
+- Feature-matching rewards replace external reward functions
+
+### Strided Mode (`mode: strided`)
+For **unstructured text** without prompt/completion splits (e.g., raw code, prose, SwallowCode).
+- Uses **strided block-parallel generation** — multiple short rollouts at different anchor points within a document
+- No vLLM needed — generation uses custom strided attention masks
+- Uses **torch flex_attention** with compiled block masks for efficient fused attention kernels (~2x faster than eager attention)
+- Compatible with gradient checkpointing via automatic dtype normalization
+- This is the core EBFT algorithm from the paper (Section F)
+
+### Common to both modes:
+- **Frozen feature network** — deep copy of the model at initialization (frozen, eval mode)
+- **Feature extraction** — hidden states at configurable layer depths (default: 25%, 50%, 75%), L2-normalized per layer before concatenation
+- **Feature-matching rewards** — cosine similarity (alignment) minus pairwise dot-product (diversity), scaled by 2 per paper equation (7)
+- **SVD whitening** — decorrelates feature dimensions; the paper shows removing it causes the largest degradation
+- **CFM loss tracking** — conditional feature-matching loss (paper eq 2) logged as `ebft/cfm_loss`
+- **FSDP2 compatible** — feature network stays outside FSDP wrapping (frozen, inference-only)
+
+## Quick Start
+
+### Structured Mode (QA data + vLLM)
+
+```bash
+# 1. Start vLLM server (LoRA serve module auto-selected when vllm_lora_sync: true)
+CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve examples/ebft/qwen3-4b-ebft-structured-async.yaml
+
+# 2. Train on a separate GPU
+CUDA_VISIBLE_DEVICES=1 axolotl train examples/ebft/qwen3-4b-ebft-structured-async.yaml
+```
+
+### Strided Mode (unstructured text)
+
+```bash
+# No vLLM needed — strided generation is built-in
+axolotl train examples/ebft/llama-3b-ebft-strided-fft.yaml
+```
+
+## Configuration
+
+### Common EBFT Settings
+
+```yaml
+rl: ebft
+
+ebft:
+  # Feature network: which layers to extract hidden states from
+  # Values are fractions of total depth (0.0 = embedding, 1.0 = final layer)
+  feature_layers: [0.25, 0.5, 0.75]
+
+  # How to pool per-token hidden states into sequence embeddings
+  # Options: "last_token" (recommended), "mean_pooling", "concat"
+  embed_method: last_token
+
+  # SVD whitening — strongly recommended (paper shows largest degradation without it)
+  use_whitening: true
+
+  # Reward = alignment_coef * alignment - diversity_coef * diversity
+  # Per paper Variant (i) (eq 49): alignment uses cosine similarity (normalized),
+  # diversity uses raw dot product — both are bounded after whitening.
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+
+  # Cross-entropy loss on ground-truth tokens (mixed objective, paper Section 2.1)
+  # 0.0 = pure feature matching; 0.03 = recommended balance; 0.1 = CE-dominated
+  ce_coef: 0.0
+```
+
+### Strided Mode Settings
+
+```yaml
+ebft:
+  mode: strided
+  stride: 8              # tokens between anchor points (paper default: 8)
+  context_length: 8      # context window per block (paper default: 8)
+  generate_max_len: 8    # tokens generated per block (paper default: 8)
+  n_samples_per_prompt: 4  # independent rollouts per document (>= 2 for RLOO)
+  temperature: 0.6
+  rl_coef: 1.0           # RL loss weight
+  advantage_estimator: rloo  # rloo (recommended), group_norm, or reinforce
+```
+
+### Structured Mode Settings (via TRL)
+
+```yaml
+trl:
+  num_generations: 4           # samples per prompt
+  max_completion_length: 256   # max tokens to generate
+  temperature: 1.0
+  use_vllm: true
+  scale_rewards: true
+  loss_type: grpo
+  epsilon: 0.2
+```
+
+### Dataset Format
+
+**Structured mode** — QA data with prompt + ground-truth completion:
+```yaml
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+```
+Transform returns: `{"prompt": ..., "ground_truth": ...}`
+
+**Strided mode** — raw text tokenized to fixed length:
+```yaml
+datasets:
+  - path: sjelassi/swallow_code_20m
+    type: ebft_pretrain.transform
+```
+Transform returns: `{"input_ids": ..., "attention_mask": ..., "labels": ...}`
+
+## How It Works
+
+### Structured Mode
+1. **Generate**: For each prompt, generate `num_generations` completions via vLLM
+2. **Extract features**: Forward both generated and ground-truth sequences through the frozen feature network
+3. **Compute rewards**: `2 * alignment - 2 * diversity` (paper eq 7)
+4. **RLOO advantages**: subtract leave-one-out group mean
+5. **Policy gradient**: clipped PPO-style loss
+
+### Strided Mode
+1. **Anchor selection**: Pick `num_blocks = (seq_len - gen_len - ctx_len) / stride + 1` anchor points across the document
+2. **Block-parallel generation**: At each anchor, generate `gen_len` tokens using a custom strided attention mask via `flex_attention` compiled block masks
+3. **Feature extraction**: Forward the full sequence (prompt + generated) through the frozen feature network **with the strided attention mask** — this is critical for correct feature representations
+4. **Per-block rewards**:
+   - **Alignment** = `2 * cosine_similarity(gen_block_emb, gt_block_emb)` — normalized, bounded in [-2, 2]
+   - **Diversity** = `2 * mean_pairwise_dot_product(gen_block_embs)` — raw dot product on whitened vectors
+   - **Reward** = `alignment_coef * alignment - diversity_coef * diversity`
+5. **RLOO advantages**: leave-one-out baseline across `n_samples_per_prompt` rollouts per block
+6. **Policy gradient**: REINFORCE loss on generated tokens, weighted by per-block advantages
+
+### Tracked Metrics
+
+| Metric | Description |
+|--------|-------------|
+| `ebft/alignment` | Mean cosine similarity between generated and GT features (higher = better) |
+| `ebft/diversity` | Mean pairwise similarity between samples (lower = more diverse) |
+| `ebft/mean_reward` | alignment - diversity (should trend upward) |
+| `ebft/cfm_loss` | Conditional feature-matching loss ‖E[φ(ŷ)] - φ(y)‖² (paper eq 2, lower = better) |
+| `ebft/rl_loss` | REINFORCE policy gradient loss |
+| `ebft/ce_loss` | Cross-entropy loss on ground-truth tokens (when `ce_coef > 0`) |
+| `ebft/advantages_std` | RLOO advantage standard deviation (should be non-zero) |
+
+## Tips and Recommendations
+
+### Reward coefficients
+- **`use_whitening: true`**: Strongly recommended. The paper's ablation (Figure 7) shows removing whitening causes the largest performance degradation. Safe to use with `diversity_coef > 0`.
+- **`diversity_coef`**: Default 1.0. Per the paper's Variant (i) (eq 49), alignment uses cosine similarity while diversity uses raw dot product. After whitening, both are bounded and on compatible scales.
+- **`n_samples_per_prompt`**: Must be >= 2 for diversity and RLOO. 4 is the paper's default.
+- **`ce_coef`**: The paper ablates `γ ∈ {0, 0.03, 0.1}`. `0.03` balances CE and RL signals; `0.1` causes CE to dominate the gradient. `0.0` gives pure feature matching.
+
+### Feature extraction
+- **`feature_layers: [0.25, 0.5, 0.75]`**: Extracts and concatenates hidden states from 25%, 50%, 75% depth. Each layer is L2-normalized independently before concatenation. The paper shows this works better than mean pooling or single-layer extraction.
+- **`embed_method: last_token`**: Uses the last token's hidden state per block. The paper shows this outperforms mean pooling (Figure 7).
+
+### Performance
+- **`torch_compile: true`**: Recommended for strided mode. Provides additional speedup via graph compilation.
+- **flex_attention**: Strided mode automatically uses `flex_attention` with compiled block masks when available (~2x faster than eager attention). Works with gradient checkpointing via automatic dtype normalization. Falls back to eager attention with dense 4D masks if flex_attention is unavailable.
+
+### Memory
+- EBFT requires a frozen copy of the model (the feature network), roughly doubling model memory.
+- **LoRA** is recommended to reduce trainable parameter memory. The feature network is always a frozen copy of the base model (without LoRA adapters).
+- With 2 GPUs visible, the trainer automatically places the feature network on the second GPU.
+- **FSDP2** is supported — the feature network stays outside FSDP wrapping since it's frozen and inference-only. With `cpu_ram_efficient_loading`, the feature network is loaded separately from pretrained weights.
+
+## Example Configs
+
+| Config | Mode | Model | Description |
+|--------|------|-------|-------------|
+| `llama-1b-ebft-opencode.yaml` | Structured | Llama-3.2-1B | QA coding with vLLM |
+| `llama-1b-ebft-opencode-novllm.yaml` | Structured | Llama-3.2-1B | QA coding without vLLM |
+| `llama-3b-ebft-strided-fft.yaml` | Strided | Llama-3.2-3B | Unstructured code with LoRA |
+| `llama-1b-ebft-strided.yaml` | Strided | Llama-3.2-1B | Quick validation |
+
+## Citation
+
+```bibtex
+@article{jelassi2026matching,
+  title={Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models},
+  author={Jelassi, Samy and Kwun, Mujin and Zhao, Rosie and Li, Yuanzhi and Fusi, Nicolo and Du, Yilun and Kakade, Sham M. and Domingo-Enrich, Carles},
+  journal={arXiv preprint arXiv:2603.12248},
+  year={2026}
+}
+```
diff --git a/examples/ebft/ebft_opencode.py b/examples/ebft/ebft_opencode.py
new file mode 100644
index 0000000000..677949ba51
--- /dev/null
+++ b/examples/ebft/ebft_opencode.py
@@ -0,0 +1,28 @@
+"""
+Dataset transform for nvidia/OpenCodeInstruct with EBFT.
+
+Maps the dataset's `input` (prompt) and `output` (code solution) fields
+to the format expected by the EBFT trainer.
+"""
+
+
+def transform(cfg, *args, **kwargs):
+    def transform_fn(example, tokenizer=None):
+        return {
+            "prompt": [
+                {"role": "user", "content": example["input"]},
+            ],
+            "ground_truth": example["output"],
+        }
+
+    return transform_fn, {
+        "remove_columns": [
+            "id",
+            "domain",
+            "generation_algorithm",
+            "llm_judgement",
+            "unit_tests",
+            "tests_execution_status",
+            "average_test_score",
+        ]
+    }
diff --git a/examples/ebft/ebft_pretrain.py b/examples/ebft/ebft_pretrain.py
new file mode 100644
index 0000000000..27a1e54b90
--- /dev/null
+++ b/examples/ebft/ebft_pretrain.py
@@ -0,0 +1,31 @@
+"""
+Dataset transform for unstructured text data with strided EBFT.
+
+Tokenizes raw text into fixed-length input_ids for the strided trainer.
+Sequences are padded to sequence_len for uniform batching.
+"""
+
+
+def transform(cfg, *args, **kwargs):
+    seq_len = cfg.sequence_len
+
+    def transform_fn(example, tokenizer=None):
+        text = example.get("question", example.get("text", ""))
+        if tokenizer is None:
+            return {"prompt": text}
+
+        encoded = tokenizer(
+            text,
+            truncation=True,
+            max_length=seq_len,
+            padding="max_length",
+            add_special_tokens=True,
+            return_tensors=None,
+        )
+        return {
+            "input_ids": encoded["input_ids"],
+            "attention_mask": encoded["attention_mask"],
+            "labels": list(encoded["input_ids"]),
+        }
+
+    return transform_fn, {"remove_columns": ["question", "answer"]}
diff --git a/examples/ebft/ebft_strided_structured.py b/examples/ebft/ebft_strided_structured.py
new file mode 100644
index 0000000000..48743931aa
--- /dev/null
+++ b/examples/ebft/ebft_strided_structured.py
@@ -0,0 +1,80 @@
+"""
+Dataset transform for structured (prompt, completion) data with strided EBFT.
+
+Tokenizes prompt and completion separately, concatenates into a single
+input_ids sequence, and marks prompt tokens with labels=-100 so the
+strided trainer knows where to place anchors (completion span only).
+
+Works with datasets that have chat-style fields (e.g., nvidia/OpenCodeInstruct).
+"""
+
+
+def transform(cfg, *args, **kwargs):
+    seq_len = cfg.sequence_len
+
+    def transform_fn(example, tokenizer=None):
+        # Extract prompt and completion from the example
+        prompt_text = example.get(
+            "input", example.get("prompt", example.get("question", ""))
+        )
+        completion_text = example.get(
+            "output", example.get("completion", example.get("answer", ""))
+        )
+
+        if tokenizer is None:
+            return {"prompt": prompt_text}
+
+        pad_id = tokenizer.pad_token_id or tokenizer.eos_token_id
+
+        # Tokenize prompt and completion separately
+        prompt_enc = tokenizer(
+            prompt_text,
+            truncation=False,
+            add_special_tokens=True,
+            return_tensors=None,
+        )
+        completion_enc = tokenizer(
+            completion_text,
+            truncation=False,
+            add_special_tokens=False,
+            return_tensors=None,
+        )
+
+        prompt_ids = prompt_enc["input_ids"]
+        completion_ids = completion_enc["input_ids"]
+
+        # Truncate to fit within seq_len (prioritize keeping prompt + some completion)
+        total_len = len(prompt_ids) + len(completion_ids)
+        if total_len > seq_len:
+            # Truncate completion first, then prompt if needed
+            max_completion = seq_len - len(prompt_ids)
+            if max_completion < 1:
+                # Prompt alone exceeds seq_len — truncate prompt, keep at least 1 completion token
+                prompt_ids = prompt_ids[: seq_len - 1]
+                completion_ids = completion_ids[:1]
+            else:
+                completion_ids = completion_ids[:max_completion]
+
+        input_ids = prompt_ids + completion_ids
+        prompt_length = len(prompt_ids)
+
+        # Labels: -100 for prompt tokens, input_ids for completion tokens
+        labels = [-100] * prompt_length + completion_ids
+
+        # Pad to seq_len
+        pad_len = seq_len - len(input_ids)
+        attention_mask = [1] * len(input_ids) + [0] * pad_len
+        labels = labels + [-100] * pad_len
+        input_ids = input_ids + [pad_id] * pad_len
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "prompt_length": prompt_length,
+        }
+
+    # Signal to remove all original columns (filtered to existing ones at map time)
+    return transform_fn, {
+        "remove_columns": "__all__",
+    }
diff --git a/examples/ebft/llama-1b-ebft-opencode-novllm.yaml b/examples/ebft/llama-1b-ebft-opencode-novllm.yaml
new file mode 100644
index 0000000000..7d7edad339
--- /dev/null
+++ b/examples/ebft/llama-1b-ebft-opencode-novllm.yaml
@@ -0,0 +1,64 @@
+# EBFT validation config — no vLLM, uses HF generate for simplicity
+# Run: CUDA_VISIBLE_DEVICES=0 axolotl train examples/ebft/llama-1b-ebft-opencode-novllm.yaml
+
+base_model: meta-llama/Llama-3.2-1B
+chat_template: llama3
+rl: ebft
+
+ebft:
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: false
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  ce_coef: 0.0
+
+trl:
+  num_generations: 4
+  max_completion_length: 128
+  temperature: 1.0
+  use_vllm: false
+  scale_rewards: true
+  loss_type: grpo
+  epsilon: 0.2
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+    split: train[:1%]
+
+sequence_len: 512
+micro_batch_size: 2
+gradient_accumulation_steps: 2
+num_epochs: 1
+max_steps: 10
+
+learning_rate: 1.0e-5
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 2
+weight_decay: 0.01
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_linear: true
+
+bf16: auto
+attn_implementation: flash_attention_2
+gradient_checkpointing: true
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-validation
+
+wandb_project: ebft
+wandb_run_id:
+wandb_watch:
+wandb_log_model:
+
+logging_steps: 1
+save_steps: 100
diff --git a/examples/ebft/llama-1b-ebft-opencode.yaml b/examples/ebft/llama-1b-ebft-opencode.yaml
new file mode 100644
index 0000000000..c77c366779
--- /dev/null
+++ b/examples/ebft/llama-1b-ebft-opencode.yaml
@@ -0,0 +1,81 @@
+# EBFT: Energy-Based Fine-Tuning with Llama-3.2-1B on OpenCodeInstruct
+#
+# Paper: "Matching Features, Not Tokens" (Jelassi et al., 2026)
+# https://arxiv.org/abs/2603.12248
+#
+# Prerequisites:
+#   1. Start vLLM server on a separate GPU:
+#      CUDA_VISIBLE_DEVICES=1 python -m trl.scripts.vllm_serve \
+#          --model meta-llama/Llama-3.2-1B \
+#          --host 0.0.0.0 --port 8000 \
+#          --gpu-memory-utilization 0.4 --dtype bfloat16
+#
+#   2. Run training:
+#      CUDA_VISIBLE_DEVICES=0 axolotl train examples/ebft/llama-1b-ebft-opencode.yaml
+
+base_model: meta-llama/Llama-3.2-1B
+chat_template: llama3
+
+# --- Training method ---
+rl: ebft
+
+# --- EBFT configuration ---
+ebft:
+  feature_layers: [0.25, 0.5, 0.75]  # extract hidden states at 25%, 50%, 75% depth
+  embed_method: last_token            # pool to sequence embedding via last token
+  use_whitening: false                # SVD whitening (disable for speed in small runs)
+  alignment_coef: 1.0                 # cosine similarity with ground-truth features
+  diversity_coef: 1.0                 # pairwise similarity penalty
+  ce_coef: 0.0                        # cross-entropy on ground-truth (0 = pure feature matching)
+
+# --- Generation settings (via TRL/GRPO infrastructure) ---
+trl:
+  num_generations: 4                  # samples per prompt for RLOO
+  max_completion_length: 256          # max generated tokens
+  temperature: 1.0
+  use_vllm: true
+  scale_rewards: true
+  loss_type: grpo
+  epsilon: 0.2
+
+# --- Dataset ---
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+    split: train[:1%]        # first 1% for validation runs
+
+# --- Training hyperparameters ---
+sequence_len: 1024
+micro_batch_size: 2
+gradient_accumulation_steps: 4
+num_epochs: 1
+max_steps: 50
+
+learning_rate: 1.0e-5
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 5
+weight_decay: 0.01
+
+# --- LoRA (recommended to reduce memory with frozen feature network) ---
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_linear: true
+
+# --- Hardware ---
+bf16: auto
+attn_implementation: flash_attention_2
+gradient_checkpointing: true
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-llama-1b-opencode
+
+# --- Logging ---
+use_tensorboard: true
+logging_steps: 1
+save_steps: 25
diff --git a/examples/ebft/llama-1b-ebft-strided-structured.yaml b/examples/ebft/llama-1b-ebft-strided-structured.yaml
new file mode 100644
index 0000000000..02e89dea01
--- /dev/null
+++ b/examples/ebft/llama-1b-ebft-strided-structured.yaml
@@ -0,0 +1,64 @@
+# EBFT Strided Structured Mode: For structured (prompt, completion) data
+# Uses strided block-parallel generation on completion spans — no vLLM needed.
+#
+# Run: CUDA_VISIBLE_DEVICES=0 axolotl train examples/ebft/llama-1b-ebft-strided-structured.yaml
+
+base_model: meta-llama/Llama-3.2-1B
+rl: ebft
+
+ebft:
+  mode: strided                       # strided block-parallel generation
+  stride: 8                           # tokens between anchor points
+  context_length: 8                   # context window per block
+  generate_max_len: 8                 # tokens to generate per block
+  n_samples_per_prompt: 4             # rollouts per document
+  temperature: 0.6
+  top_p: 1.0
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: true
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  rl_coef: 1.0
+  ce_coef: 0.03                       # small CE weight for structured data
+  advantage_estimator: rloo
+  min_completion_prefix: 8            # skip anchors too close to prompt boundary
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_strided_structured.transform
+    split: train[:1%]
+
+sequence_len: 2048
+micro_batch_size: 1
+gradient_accumulation_steps: 2
+num_epochs: 1
+# max_steps: 10
+
+learning_rate: 1.0e-6
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 5
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_linear: true
+
+bf16: auto
+attn_implementation: flex_attention
+                        # (full-model compile conflicts with gradient checkpointing + flex_attention)
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true   # required for flex_attention (non-reentrant causes CheckpointError)
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-strided-structured
+
+wandb_project: ebft
+logging_steps: 1
+save_steps: 100
diff --git a/examples/ebft/llama-1b-ebft-strided.yaml b/examples/ebft/llama-1b-ebft-strided.yaml
new file mode 100644
index 0000000000..e3cfe8040e
--- /dev/null
+++ b/examples/ebft/llama-1b-ebft-strided.yaml
@@ -0,0 +1,59 @@
+# EBFT Strided Mode: For unstructured text data (raw code, prose)
+# Uses strided block-parallel generation — no vLLM needed.
+#
+# Run: CUDA_VISIBLE_DEVICES=0 axolotl train examples/ebft/llama-1b-ebft-strided.yaml
+
+base_model: meta-llama/Llama-3.2-1B
+rl: ebft
+
+ebft:
+  mode: strided                       # strided block-parallel generation
+  stride: 8                           # tokens between anchor points
+  context_length: 8                   # context window per block
+  generate_max_len: 8                 # tokens to generate per block
+  n_samples_per_prompt: 4             # rollouts per document
+  temperature: 0.6
+  top_p: 1.0
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: true
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  rl_coef: 1.0
+  ce_coef: 0.0
+  advantage_estimator: rloo
+
+datasets:
+  - path: sjelassi/swallow_code_20m
+    type: ebft_pretrain.transform
+    split: train[:100]
+
+sequence_len: 256
+micro_batch_size: 1
+gradient_accumulation_steps: 2
+num_epochs: 1
+max_steps: 5
+
+learning_rate: 1.0e-6
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 2
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_linear: true
+
+bf16: auto
+gradient_checkpointing: true
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-strided-validation
+
+wandb_project: ebft
+logging_steps: 1
+save_steps: 100
diff --git a/examples/ebft/llama-3b-ebft-strided-fft.yaml b/examples/ebft/llama-3b-ebft-strided-fft.yaml
new file mode 100644
index 0000000000..e39d3bcfa2
--- /dev/null
+++ b/examples/ebft/llama-3b-ebft-strided-fft.yaml
@@ -0,0 +1,68 @@
+# EBFT Strided: LoRA Llama-3.2-3B on SwallowCode, 100 steps
+# Actor on GPU 0, frozen feature network on GPU 1
+#
+# Run: CUDA_VISIBLE_DEVICES=0,1 python -m axolotl.cli.train examples/ebft/llama-3b-ebft-strided-fft.yaml
+
+base_model: meta-llama/Llama-3.2-3B
+rl: ebft
+
+ebft:
+  mode: strided
+  stride: 8
+  context_length: 8
+  generate_max_len: 8
+  n_samples_per_prompt: 4
+  temperature: 0.6
+  top_p: 1.0
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: true
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  rl_coef: 1.0
+  ce_coef: 0.0   # paper recommends 0.03 for mixed objective; 0.1 causes CE to dominate
+  advantage_estimator: rloo
+
+datasets:
+  - path: sjelassi/swallow_code_20m
+    type: ebft_pretrain.transform
+    split: train[:5000]
+
+sequence_len: 1024
+micro_batch_size: 1
+gradient_accumulation_steps: 4
+num_epochs: 1
+max_steps: 100
+
+learning_rate: 1.0e-5
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 10
+weight_decay: 0.01
+
+adapter: lora
+lora_r: 32
+lora_alpha: 64
+lora_dropout: 0.05
+lora_target_linear: true
+
+bf16: auto
+torch_dtype: bfloat16
+gradient_checkpointing: true
+torch_compile: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+ddp: false
+device_map:
+  "": 0
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-llama3b-strided
+
+wandb_project: ebft
+wandb_name: llama3b-strided-lora-100steps
+logging_steps: 1
+save_steps: 50
diff --git a/examples/ebft/llama-8b-ebft-strided-fft.yaml b/examples/ebft/llama-8b-ebft-strided-fft.yaml
new file mode 100644
index 0000000000..caed980855
--- /dev/null
+++ b/examples/ebft/llama-8b-ebft-strided-fft.yaml
@@ -0,0 +1,57 @@
+# EBFT Strided: Full-parameter Llama-3.1-8B on SwallowCode, 100 steps
+# Feature network is CPU-offloaded to fit in single 32GB GPU
+#
+# Run: CUDA_VISIBLE_DEVICES=0 python -m axolotl.cli.train examples/ebft/llama-8b-ebft-strided-fft.yaml
+
+base_model: meta-llama/Llama-3.1-8B
+rl: ebft
+
+ebft:
+  mode: strided
+  stride: 8
+  context_length: 8
+  generate_max_len: 8
+  n_samples_per_prompt: 4
+  temperature: 0.6
+  top_p: 1.0
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: true
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  rl_coef: 1.0
+  ce_coef: 0.0
+  advantage_estimator: rloo
+
+datasets:
+  - path: sjelassi/swallow_code_20m
+    type: ebft_pretrain.transform
+    split: train[:5000]
+
+sequence_len: 1024
+micro_batch_size: 1
+gradient_accumulation_steps: 4
+num_epochs: 1
+max_steps: 100
+
+learning_rate: 1.0e-6
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 10
+weight_decay: 0.01
+
+bf16: auto
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-llama8b-strided
+
+wandb_project: ebft
+wandb_name: llama8b-strided-fft-100steps
+logging_steps: 1
+save_steps: 50
diff --git a/examples/ebft/qwen35-4b-ebft-structured-async.yaml b/examples/ebft/qwen35-4b-ebft-structured-async.yaml
new file mode 100644
index 0000000000..daa77d6f67
--- /dev/null
+++ b/examples/ebft/qwen35-4b-ebft-structured-async.yaml
@@ -0,0 +1,86 @@
+# EBFT Structured Mode: Qwen3.5-4B (hybrid linear attention)
+#
+# Qwen3.5 uses hybrid attention: linear attention (conv1d) on 3/4 of layers,
+# full attention every 4th layer. This tests EBFT compatibility.
+#
+# Prerequisites:
+#   1. Start vLLM on GPU 0:
+#      CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve examples/ebft/qwen35-4b-ebft-structured-async.yaml
+#
+#   2. Run training on GPU 1:
+#      CUDA_VISIBLE_DEVICES=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+#          axolotl train examples/ebft/qwen35-4b-ebft-structured-async.yaml
+
+base_model: Qwen/Qwen3.5-4B
+
+rl: ebft
+
+ebft:
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: false
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  ce_coef: 0.0
+
+trl:
+  num_generations: 4
+  max_completion_length: 256
+  temperature: 0.7
+  use_vllm: true
+  vllm_server_host: 0.0.0.0
+  vllm_server_port: 8000
+  scale_rewards: true
+  loss_type: grpo
+  epsilon: 0.2
+  generation_kwargs:
+    stop_token_ids: [248044, 248046]  # <|endoftext|>, <|im_end|>
+  chat_template_kwargs:
+    enable_thinking: false
+  async_prefetch: true
+  vllm_server_timeout: 300
+
+vllm:
+  gpu_memory_utilization: 0.5
+  max_model_len: 2048
+  serve_module: axolotl.scripts.vllm_serve_lora
+  enforce_eager: true
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+    split: train[:500]
+
+sequence_len: 1024
+micro_batch_size: 1
+gradient_accumulation_steps: 4
+num_epochs: 1
+max_steps: 10
+
+learning_rate: 5.0e-6
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 3
+weight_decay: 0.01
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.0
+# Target full-attention q/k/v/o on layers 3,7,11,15,19,23,27,31 + MLP on all layers
+# Avoids linear_attn modules (in_proj_qkv, in_proj_z, etc.) which break vLLM LoRA
+lora_target_modules: ".*\\.layers\\.(3|7|11|15|19|23|27|31)\\.self_attn\\.(q|k|v|o)_proj|.*\\.mlp\\.(gate|up|down)_proj"
+
+bf16: auto
+attn_implementation: flash_attention_2
+gradient_checkpointing: true
+
+special_tokens:
+  pad_token: "<|endoftext|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-qwen35-4b-structured-async
+
+wandb_project: ebft
+logging_steps: 1
+save_steps: 50
diff --git a/examples/ebft/qwen35-4b-ebft-structured.yaml b/examples/ebft/qwen35-4b-ebft-structured.yaml
new file mode 100644
index 0000000000..d1b2a72f28
--- /dev/null
+++ b/examples/ebft/qwen35-4b-ebft-structured.yaml
@@ -0,0 +1,77 @@
+# EBFT Structured Mode: Qwen3.5-4B (hybrid linear attention)
+#
+# Qwen3.5 uses hybrid attention: linear attention (conv1d) on 3/4 of layers,
+# full attention every 4th layer. This tests EBFT compatibility.
+#
+# Prerequisites:
+#   1. Start vLLM on GPU 0:
+#      CUDA_VISIBLE_DEVICES=0 trl vllm-serve --model Qwen/Qwen3.5-4B \
+#          --gpu-memory-utilization 0.5 --max-model-len 2048 --enforce-eager
+#
+#   2. Run training on GPU 1:
+#      CUDA_VISIBLE_DEVICES=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+#          axolotl train examples/ebft/qwen35-4b-ebft-structured.yaml
+
+base_model: Qwen/Qwen3.5-4B
+
+rl: ebft
+
+ebft:
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: false
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  ce_coef: 0.0
+
+trl:
+  num_generations: 4
+  max_completion_length: 256
+  temperature: 0.7
+  use_vllm: true
+  vllm_server_host: 0.0.0.0
+  vllm_server_port: 8000
+  scale_rewards: true
+  loss_type: grpo
+  epsilon: 0.2
+  generation_kwargs:
+    stop_token_ids: [248044, 248046]  # <|endoftext|>, <|im_end|>
+  chat_template_kwargs:
+    enable_thinking: false  # disable Qwen3.5 thinking mode for shorter completions
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+    split: train[:500]
+
+sequence_len: 1024
+micro_batch_size: 1
+gradient_accumulation_steps: 4
+num_epochs: 1
+max_steps: 10
+
+learning_rate: 5.0e-6
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 3
+weight_decay: 0.01
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.0
+lora_target_modules: ".*\\.layers\\.(3|7|11|15|19|23|27|31)\\.self_attn\\.(q|k|v|o)_proj|.*\\.mlp\\.(gate|up|down)_proj"
+
+bf16: auto
+attn_implementation: flash_attention_2
+gradient_checkpointing: true
+
+special_tokens:
+  pad_token: "<|endoftext|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-qwen35-4b-structured
+
+wandb_project: ebft
+logging_steps: 1
+save_steps: 50
diff --git a/examples/ebft/qwen35-9b-ebft-structured.yaml b/examples/ebft/qwen35-9b-ebft-structured.yaml
new file mode 100644
index 0000000000..ad3b8538e7
--- /dev/null
+++ b/examples/ebft/qwen35-9b-ebft-structured.yaml
@@ -0,0 +1,82 @@
+# EBFT Structured Mode: Qwen3.5-9B (hybrid linear attention)
+#
+# Prerequisites:
+#   1. Start vLLM on GPU 0:
+#      CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve examples/ebft/qwen35-9b-ebft-structured.yaml
+#
+#   2. Run training on GPU 1:
+#      CUDA_VISIBLE_DEVICES=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+#          axolotl train examples/ebft/qwen35-9b-ebft-structured.yaml
+
+base_model: Qwen/Qwen3.5-9B
+
+rl: ebft
+
+ebft:
+  feature_layers: [0.25, 0.5, 0.75]
+  embed_method: last_token
+  use_whitening: false
+  alignment_coef: 1.0
+  diversity_coef: 1.0
+  ce_coef: 0.0
+
+trl:
+  num_generations: 4
+  max_completion_length: 256
+  temperature: 0.7
+  use_vllm: true
+  vllm_server_host: 0.0.0.0
+  vllm_server_port: 8000
+  scale_rewards: true
+  loss_type: grpo
+  epsilon: 0.2
+  generation_kwargs:
+    stop_token_ids: [248044, 248046]  # <|endoftext|>, <|im_end|>
+  chat_template_kwargs:
+    enable_thinking: false
+  vllm_server_timeout: 300
+
+vllm:
+  gpu_memory_utilization: 0.7
+  max_model_len: 2048
+  serve_module: axolotl.scripts.vllm_serve_lora
+  enforce_eager: true
+
+datasets:
+  - path: nvidia/OpenCodeInstruct
+    type: ebft_opencode.transform
+    split: train[:500]
+
+sequence_len: 1024
+micro_batch_size: 1
+gradient_accumulation_steps: 4
+num_epochs: 1
+max_steps: 10
+
+learning_rate: 3.0e-6
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+warmup_steps: 3
+weight_decay: 0.01
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.0
+# Target full-attention q/k/v/o on layers 3,7,11,15,19,23,27,31 + MLP on all layers
+# Avoids linear_attn modules (in_proj_qkv, in_proj_z, etc.) which break vLLM LoRA
+lora_target_modules: ".*\\.layers\\.(3|7|11|15|19|23|27|31)\\.self_attn\\.(q|k|v|o)_proj|.*\\.mlp\\.(gate|up|down)_proj"
+
+bf16: auto
+attn_implementation: flash_attention_2
+gradient_checkpointing: true
+
+special_tokens:
+  pad_token: "<|endoftext|>"
+
+val_set_size: 0.0
+output_dir: ./outputs/ebft-qwen35-9b-structured
+
+wandb_project: ebft
+logging_steps: 1
+save_steps: 50
diff --git a/examples/expert_parallel/qwen3_30ba3b_ep_fft_4gpu.yaml b/examples/expert_parallel/qwen3_30ba3b_ep_fft_4gpu.yaml
new file mode 100644
index 0000000000..76b0ce3706
--- /dev/null
+++ b/examples/expert_parallel/qwen3_30ba3b_ep_fft_4gpu.yaml
@@ -0,0 +1,37 @@
+base_model: Qwen/Qwen3-30B-A3B-Instruct-2507
+
+plugins:
+  - axolotl.integrations.expert_parallel.ExpertParallelPlugin
+
+experts_implementation: grouped_mm
+expert_parallel_size: 4 # world size 4
+
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+
+dataset_prepared_path:
+output_dir: ./outputs/out
+
+sequence_len: 2048
+sample_packing: true
+
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+max_steps: 50
+
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 5e-6
+warmup_steps: 0.1
+
+bf16: true
+tf32: false
+
+gradient_checkpointing: true
+flash_attention: true
+logging_steps: 1
+evals_per_epoch:
+saves_per_epoch: 1
+
+seed: 42
diff --git a/examples/expert_parallel/qwen3_30ba3b_ep_fsdp_fft_4gpu.yaml b/examples/expert_parallel/qwen3_30ba3b_ep_fsdp_fft_4gpu.yaml
new file mode 100644
index 0000000000..f3761eae24
--- /dev/null
+++ b/examples/expert_parallel/qwen3_30ba3b_ep_fsdp_fft_4gpu.yaml
@@ -0,0 +1,49 @@
+base_model: Qwen/Qwen3-30B-A3B-Instruct-2507
+
+plugins:
+  - axolotl.integrations.expert_parallel.ExpertParallelPlugin
+
+experts_implementation: grouped_mm
+# world size 4 minimum: expert_parallel_size x dp_shard_size
+expert_parallel_size: 2
+dp_shard_size: 2
+
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+
+dataset_prepared_path:
+output_dir: ./outputs/out
+
+sequence_len: 2048
+sample_packing: true
+
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+max_steps: 50
+
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 5e-6
+warmup_steps: 0.1
+
+bf16: true
+tf32: false
+
+gradient_checkpointing: true
+flash_attention: true
+logging_steps: 1
+evals_per_epoch:
+saves_per_epoch: 1
+
+seed: 42
+
+fsdp_config:
+  fsdp_version: 2
+  offload_params: false
+  cpu_ram_efficient_loading: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen3MoeDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
diff --git a/examples/expert_parallel/qwen3_30ba3b_ep_lora_4gpu.yaml b/examples/expert_parallel/qwen3_30ba3b_ep_lora_4gpu.yaml
new file mode 100644
index 0000000000..1f08ed2dc0
--- /dev/null
+++ b/examples/expert_parallel/qwen3_30ba3b_ep_lora_4gpu.yaml
@@ -0,0 +1,59 @@
+base_model: Qwen/Qwen3-30B-A3B-Instruct-2507
+
+plugins:
+  - axolotl.integrations.expert_parallel.ExpertParallelPlugin
+
+experts_implementation: grouped_mm
+expert_parallel_size: 4 # world size 4
+
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+
+dataset_prepared_path:
+output_dir: ./outputs/out
+
+sequence_len: 2048
+sample_packing: true
+
+adapter: lora
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  # Add gate_up_proj and down_proj to also target shared experts (nn.Linear):
+  # - gate_up_proj
+  # - down_proj
+
+# Target routed experts (3D nn.Parameter tensors, not nn.Linear — use lora_target_parameters):
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+max_steps: 50
+
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.003
+warmup_steps: 0.1
+
+bf16: true
+tf32: false
+
+gradient_checkpointing: true
+flash_attention: true
+logging_steps: 1
+evals_per_epoch:
+saves_per_epoch: 1
+
+seed: 42
diff --git a/examples/falcon-e/falcon-e-3b-dpo.yaml b/examples/falcon-e/falcon-e-3b-dpo.yaml
new file mode 100644
index 0000000000..72d1cc41af
--- /dev/null
+++ b/examples/falcon-e/falcon-e-3b-dpo.yaml
@@ -0,0 +1,93 @@
+base_model: axolotl-ai-co/Falcon-E-1.2-3B-Exp-prequantized
+output_dir: ./output
+
+plugins:
+  - axolotl.integrations.kernels.KernelsPlugin
+
+use_kernels: false
+use_scattermoe: false
+use_sonicmoe: false
+use_onebitllms: true
+
+load_in_8bit: false
+load_in_4bit: false
+
+chat_template: tokenizer_default
+
+rl: dpo
+datasets:
+  - path: allenai/Dolci-Think-DPO-7B
+    split: train
+    type: chatml.ultra
+
+dataset_prepared_path: ./axolotl_dataset_cache
+
+sequence_len: 8192
+trust_remote_code: false
+
+gradient_accumulation_steps: 4 # This can run on 4 GPUs
+
+# Very important to enable gradient accumulation with FSDP
+# https://github.com/huggingface/transformers/issues/29425
+accelerator_config:
+  gradient_accumulation_kwargs:
+    sync_each_batch: True
+
+
+micro_batch_size: 1
+num_epochs: 3
+optimizer: adamw_torch
+lr_scheduler: cosine
+learning_rate: 1.0e-5
+# adamw hyperparams
+adam_beta1: 0.9
+adam_beta2: 0.95
+
+bf16: true
+tf32: false
+
+logging_steps: 1
+
+flash_attention: true
+
+loss_watchdog_threshold: 15.0
+loss_watchdog_patience: 3
+
+warmup_steps: 128
+evals_per_epoch: 0
+
+save_steps: 500
+save_strategy: steps
+
+weight_decay: 0.01
+
+shuffle_merged_datasets: true
+experimental_skip_move_to_device: true
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+# Comment to disable CP
+# The number of GPUs to shard the model parameters across (FSDP dimension).
+dp_shard_size: 1
+
+# The number of times to replicate the sharded model (DDP dimension).
+dp_replicate_size: 1
+
+# Number of GPUs for Tensor Parallelism.
+tensor_parallel_size: 1  # (default is 1, no TP)
+
+# Number of GPUs for Context/Sequence Parallelism.
+context_parallel_size: 1 # (default is 1, no CP)
+
+special_tokens:
+  eos_token: <|end_of_text|>
+
+eot_tokens:
+  - <|im_end|>
diff --git a/examples/falcon-e/falcon-e-3b-ft.yaml b/examples/falcon-e/falcon-e-3b-ft.yaml
new file mode 100644
index 0000000000..0898271cfe
--- /dev/null
+++ b/examples/falcon-e/falcon-e-3b-ft.yaml
@@ -0,0 +1,100 @@
+base_model: tiiuae/Falcon-E-3B-Base-prequantized
+output_dir: ./output
+
+plugins:
+  - axolotl.integrations.kernels.KernelsPlugin
+
+use_kernels: false
+use_scattermoe: false
+use_sonicmoe: false
+use_onebitllms: true
+
+load_in_8bit: false
+load_in_4bit: false
+
+chat_template: tokenizer_default
+
+datasets:
+  - path: cgato/SlimOrcaDedupCleaned
+    type: chat_template
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+dataset_prepared_path: ./axolotl_dataset_cache
+
+sequence_len: 32768
+trust_remote_code: false
+
+
+gradient_accumulation_steps: 4 # This can run on 4 GPUs
+
+# Very important to enable gradient accumulation with FSDP
+# https://github.com/huggingface/transformers/issues/29425
+accelerator_config:
+  gradient_accumulation_kwargs:
+    sync_each_batch: True
+
+
+micro_batch_size: 1
+num_epochs: 3
+optimizer: adamw_torch
+lr_scheduler: cosine
+learning_rate: 5.0e-4
+# adamw hyperparams
+adam_beta1: 0.9
+adam_beta2: 0.95
+
+bf16: true
+tf32: false
+
+logging_steps: 1
+
+flash_attention: true
+
+loss_watchdog_threshold: 15.0
+loss_watchdog_patience: 3
+
+warmup_steps: 128
+evals_per_epoch: 0
+
+save_steps: 500
+save_strategy: steps
+
+weight_decay: 0.01
+
+sample_packing: true
+pad_to_sequence_len: true
+
+shuffle_merged_datasets: true
+experimental_skip_move_to_device: true
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  reshard_after_forward: true
+  activation_checkpointing: true
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
+
+# Comment to disable CP
+# The number of GPUs to shard the model parameters across (FSDP dimension).
+dp_shard_size: 1
+
+# The number of times to replicate the sharded model (DDP dimension).
+dp_replicate_size: 1
+
+# Number of GPUs for Tensor Parallelism.
+tensor_parallel_size: 1  # (default is 1, no TP)
+
+# Number of GPUs for Context/Sequence Parallelism.
+context_parallel_size: 1 # (default is 1, no CP)
+
+special_tokens:
+  eos_token: <|end_of_text|>
+
+eot_tokens:
+  - <|im_end|>
diff --git a/examples/falcon-h1/falcon-h1-1b-deep-qlora.yaml b/examples/falcon-h1/falcon-h1-1b-deep-qlora.yaml
new file mode 100644
index 0000000000..f59f0df5ca
--- /dev/null
+++ b/examples/falcon-h1/falcon-h1-1b-deep-qlora.yaml
@@ -0,0 +1,73 @@
+base_model: tiiuae/Falcon-H1-1.5B-Deep-Base
+# 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
+
+load_in_8bit: false
+load_in_4bit: true
+
+# huggingface repo
+chat_template: falcon_h1
+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_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - in_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/falcon-h1/falcon-h1-1b-qlora-cp.yaml b/examples/falcon-h1/falcon-h1-1b-qlora-cp.yaml
new file mode 100644
index 0000000000..606832d316
--- /dev/null
+++ b/examples/falcon-h1/falcon-h1-1b-qlora-cp.yaml
@@ -0,0 +1,75 @@
+base_model: tiiuae/Falcon-H1-1.5B-Deep-Base
+# 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
+
+load_in_8bit: false
+load_in_4bit: true
+
+# huggingface repo
+chat_template: falcon_h1
+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_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - in_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+context_parapllel_size: 2
+
+
+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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/falcon-h1/falcon-h1-1b-qlora.yaml b/examples/falcon-h1/falcon-h1-1b-qlora.yaml
new file mode 100644
index 0000000000..8c3eb080d2
--- /dev/null
+++ b/examples/falcon-h1/falcon-h1-1b-qlora.yaml
@@ -0,0 +1,72 @@
+base_model: tiiuae/Falcon-H1-1.5B-Base
+# 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
+
+load_in_8bit: false
+load_in_4bit: true
+
+# huggingface repo
+chat_template: falcon_h1
+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_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - in_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/falcon-h1/falcon-h1-34b-qlora.yaml b/examples/falcon-h1/falcon-h1-34b-qlora.yaml
new file mode 100644
index 0000000000..28e7de9566
--- /dev/null
+++ b/examples/falcon-h1/falcon-h1-34b-qlora.yaml
@@ -0,0 +1,73 @@
+base_model: tiiuae/Falcon-H1-34B-Base
+# 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
+
+load_in_8bit: false
+load_in_4bit: true
+
+# huggingface repo
+chat_template: falcon_h1
+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_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - in_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/falcon-h1/falcon-h1-3b-qlora.yaml b/examples/falcon-h1/falcon-h1-3b-qlora.yaml
new file mode 100644
index 0000000000..71b38e2f7c
--- /dev/null
+++ b/examples/falcon-h1/falcon-h1-3b-qlora.yaml
@@ -0,0 +1,73 @@
+base_model: tiiuae/Falcon-H1-3B-Base
+# 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
+
+load_in_8bit: false
+load_in_4bit: true
+
+# huggingface repo
+chat_template: falcon_h1
+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_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - in_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/falcon-h1/falcon-h1-500m-qlora.yaml b/examples/falcon-h1/falcon-h1-500m-qlora.yaml
new file mode 100644
index 0000000000..91602ae714
--- /dev/null
+++ b/examples/falcon-h1/falcon-h1-500m-qlora.yaml
@@ -0,0 +1,73 @@
+base_model: tiiuae/Falcon-H1-0.5B-Instruct
+# 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
+
+load_in_8bit: false
+load_in_4bit: true
+
+# huggingface repo
+chat_template: falcon_h1
+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_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - in_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/falcon-h1/falcon-h1-7b-qlora.yaml b/examples/falcon-h1/falcon-h1-7b-qlora.yaml
new file mode 100644
index 0000000000..cc7e8f6cdc
--- /dev/null
+++ b/examples/falcon-h1/falcon-h1-7b-qlora.yaml
@@ -0,0 +1,73 @@
+base_model: tiiuae/Falcon-H1-7B-Base
+# 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
+
+load_in_8bit: false
+load_in_4bit: true
+
+# huggingface repo
+chat_template: falcon_h1
+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_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - in_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/gemma2/qlora.yml b/examples/gemma2/qlora.yml
new file mode 100644
index 0000000000..8e79316493
--- /dev/null
+++ b/examples/gemma2/qlora.yml
@@ -0,0 +1,66 @@
+base_model: google/gemma-2-9b
+# 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
+
+load_in_8bit: false
+load_in_4bit: true
+
+# huggingface repo
+chat_template: gemma
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: cgato/SlimOrcaDedupCleaned
+    type: chat_template
+    drop_system_message: true
+    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
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/gemma2/reward-model.yaml b/examples/gemma2/reward-model.yaml
new file mode 100644
index 0000000000..f48bff6266
--- /dev/null
+++ b/examples/gemma2/reward-model.yaml
@@ -0,0 +1,54 @@
+base_model: google/gemma-2-2b
+# optionally might have model_type or tokenizer_type
+model_type: AutoModelForSequenceClassification
+num_labels: 1
+tokenizer_type: AutoTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+reward_model: true
+chat_template: gemma
+datasets:
+  - path: argilla/distilabel-intel-orca-dpo-pairs
+    type: bradley_terry.chat_template
+val_set_size: 0.0
+output_dir: ./outputs/out
+remove_unused_columns: false
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+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: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/gemma3/gemma-3-1b-qlora.yml b/examples/gemma3/gemma-3-1b-qlora.yml
new file mode 100644
index 0000000000..95b99a0da9
--- /dev/null
+++ b/examples/gemma3/gemma-3-1b-qlora.yml
@@ -0,0 +1,73 @@
+base_model: google/gemma-3-1b-it
+
+# 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
+
+# huggingface repo
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+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
+
+# Freeze vision tower
+unfrozen_parameters:
+  - ^model.language_model.*
+  - ^lm_head.*
+
+adapter: qlora
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_linear: true
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/gemma3/gemma-3-270m-qlora.yml b/examples/gemma3/gemma-3-270m-qlora.yml
new file mode 100644
index 0000000000..800a88a1b9
--- /dev/null
+++ b/examples/gemma3/gemma-3-270m-qlora.yml
@@ -0,0 +1,71 @@
+base_model: google/gemma-3-270m-it
+
+# 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
+
+# huggingface repo
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+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
+
+# Freeze vision tower
+unfrozen_parameters:
+  - ^model.language_model.*
+  - ^lm_head.*
+
+adapter: qlora
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_linear: true
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: false
+
+
+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: auto
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/gemma3/gemma-3-4b-qlora.yml b/examples/gemma3/gemma-3-4b-qlora.yml
new file mode 100644
index 0000000000..e7c43ddef8
--- /dev/null
+++ b/examples/gemma3/gemma-3-4b-qlora.yml
@@ -0,0 +1,68 @@
+base_model: google/gemma-3-4b-it
+
+load_in_4bit: true
+
+# gemma3 doesn't seem to play nice with ddp
+ddp_find_unused_parameters: true
+
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+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
+
+# Freeze vision tower
+unfrozen_parameters:
+  - ^model.language_model.*
+  - ^lm_head.*
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_linear: true
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/gemma3/gemma-3-4b-vision-qlora.yml b/examples/gemma3/gemma-3-4b-vision-qlora.yml
new file mode 100644
index 0000000000..790d9543a3
--- /dev/null
+++ b/examples/gemma3/gemma-3-4b-vision-qlora.yml
@@ -0,0 +1,65 @@
+base_model: google/gemma-3-4b-it
+processor_type: AutoProcessor
+
+load_in_4bit: true
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+# gemma3 doesn't seem to play nice with ddp
+ddp_find_unused_parameters: true
+
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+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
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/gemma3n/README.md b/examples/gemma3n/README.md
new file mode 100644
index 0000000000..db265d7581
--- /dev/null
+++ b/examples/gemma3n/README.md
@@ -0,0 +1,69 @@
+# Finetune Gemma-3n with Axolotl
+
+Gemma-3n is a family of multimodal models from Google found on [HuggingFace](https://huggingface.co/collections/google/gemma-3n-685065323f5984ef315c93f4). This guide shows how to fine-tune it with Axolotl.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+uv pip install --no-build-isolation 'axolotl>=0.16.1'
+```
+
+2. In addition to Axolotl's requirements, Gemma-3n requires:
+
+```bash
+uv pip install timm==1.0.17
+
+# for loading audio data
+uv pip install librosa==0.11.0
+```
+
+3. Download sample dataset files
+
+```bash
+# for text + vision + audio only
+wget https://huggingface.co/datasets/Nanobit/text-vision-audio-2k-test/resolve/main/African_elephant.jpg
+wget https://huggingface.co/datasets/Nanobit/text-vision-audio-2k-test/resolve/main/En-us-African_elephant.oga
+```
+
+4. Run the finetuning example:
+
+```bash
+# text only
+axolotl train examples/gemma3n/gemma-3n-e2b-qlora.yml
+
+# text + vision
+axolotl train examples/gemma3n/gemma-3n-e2b-vision-qlora.yml
+
+# text + vision + audio
+axolotl train examples/gemma3n/gemma-3n-e2b-vision-audio-qlora.yml
+```
+
+Let us know how it goes. Happy finetuning! 🚀
+
+WARNING: The loss and grad norm will be much higher than normal. We suspect this to be inherent to the model as of the moment. If anyone would like to submit a fix for this, we are happy to take a look.
+
+### TIPS
+
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+- The multimodal dataset format follows the OpenAI multi-content Messages format as seen [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Related Resources
+
+- [Gemma 3n Blog](https://ai.google.dev/gemma/docs/gemma-3n)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/gemma3n/gemma-3n-e2b-qlora.yml b/examples/gemma3n/gemma-3n-e2b-qlora.yml
new file mode 100644
index 0000000000..ad7ab5726b
--- /dev/null
+++ b/examples/gemma3n/gemma-3n-e2b-qlora.yml
@@ -0,0 +1,74 @@
+base_model: google/gemma-3n-E2B-it
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+cut_cross_entropy: true
+
+load_in_8bit: false
+load_in_4bit: true
+
+# for use with fft to only train on language model layers
+# unfrozen_parameters:
+  # - model.language_model.*
+  # - lm_head
+  # - embed_tokens
+
+
+chat_template: gemma3n
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: cgato/SlimOrcaDedupCleaned
+    type: chat_template
+    split: train[:1%]
+    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: # Does not work with gemma3n currently
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+pad_to_sequence_len: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+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  # Any attention impl does not work with gemma3n now
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/gemma3n/gemma-3n-e2b-vision-audio-qlora.yml b/examples/gemma3n/gemma-3n-e2b-vision-audio-qlora.yml
new file mode 100644
index 0000000000..d72d7fbc08
--- /dev/null
+++ b/examples/gemma3n/gemma-3n-e2b-vision-audio-qlora.yml
@@ -0,0 +1,78 @@
+base_model: google/gemma-3n-E2B-it
+processor_type: AutoProcessor
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+cut_cross_entropy: true
+
+# for use with fft to only train on language model layers
+# unfrozen_parameters:
+  # - model.language_model.*
+  # - lm_head
+  # - embed_tokens
+
+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: gemma3n
+eot_tokens:
+  - <end_of_turn>
+
+# sample dataset below requires downloading audio/image in advance
+# wget https://huggingface.co/datasets/Nanobit/text-vision-audio-2k-test/resolve/main/African_elephant.jpg
+# wget https://huggingface.co/datasets/Nanobit/text-vision-audio-2k-test/resolve/main/En-us-African_elephant.oga
+datasets:
+  - path: Nanobit/text-vision-audio-2k-test
+    type: chat_template
+dataset_prepared_path:
+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: 'model.language_model.layers.[\d]+.(mlp|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  # Any attention impl does not work with gemma3n now
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/gemma3n/gemma-3n-e2b-vision-qlora.yml b/examples/gemma3n/gemma-3n-e2b-vision-qlora.yml
new file mode 100644
index 0000000000..c87eca663a
--- /dev/null
+++ b/examples/gemma3n/gemma-3n-e2b-vision-qlora.yml
@@ -0,0 +1,75 @@
+base_model: google/gemma-3n-E2B-it
+processor_type: AutoProcessor
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+cut_cross_entropy: true
+
+# for use with fft to only train on language model layers
+# unfrozen_parameters:
+  # - model.language_model.*
+  # - lm_head
+  # - embed_tokens
+
+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: gemma3n
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+dataset_prepared_path:
+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: 'model.language_model.layers.[\d]+.(mlp|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  # Any attention impl does not work with gemma3n now
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/gemma4-unified/12b-text-lora.yaml b/examples/gemma4-unified/12b-text-lora.yaml
new file mode 100644
index 0000000000..f4f1cd60a0
--- /dev/null
+++ b/examples/gemma4-unified/12b-text-lora.yaml
@@ -0,0 +1,60 @@
+base_model: google/gemma-4-12B-it
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.liger.LigerPlugin
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+strict: false
+
+chat_template: gemma4_unified
+eot_tokens:
+  - "<turn|>"
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:10%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.05
+output_dir: ./outputs/gemma4-unified-12b-lora
+
+sequence_len: 2048
+sample_packing: true
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|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_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+logging_steps: 1
+
+gemma4_hybrid_attn_impl: true
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/gemma4-unified/12b-vision-lora.yaml b/examples/gemma4-unified/12b-vision-lora.yaml
new file mode 100644
index 0000000000..6090cb7c03
--- /dev/null
+++ b/examples/gemma4-unified/12b-vision-lora.yaml
@@ -0,0 +1,62 @@
+base_model: google/gemma-4-12B-it
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.liger.LigerPlugin
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+
+processor_type: AutoProcessor
+freeze_mm_modules: true
+strict: false
+
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: gemma4_unified
+eot_tokens:
+  - "<turn|>"
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:100]
+
+val_set_size: 0
+output_dir: ./outputs/gemma4-unified-12b-vision-lora
+
+adapter: lora
+sequence_len: 2048
+pad_to_sequence_len: false
+
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+max_steps: 10
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+logging_steps: 1
+attn_implementation: sdpa
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
diff --git a/examples/gemma4-unified/README.md b/examples/gemma4-unified/README.md
new file mode 100644
index 0000000000..5964956801
--- /dev/null
+++ b/examples/gemma4-unified/README.md
@@ -0,0 +1,42 @@
+# Finetune Google's Gemma 4 Unified with Axolotl
+
+[Gemma 4 Unified](https://huggingface.co/google/gemma-4-12B-it) is the **encoder-free** multimodal member of the [Gemma 4](https://huggingface.co/collections/google/gemma-4) family; no vision tower and no audio tower. Raw image patches and 16 kHz waveform frames are projected directly into the language model through lightweight `LayerNorm/RMSNorm → Linear` pipelines. The text backbone is the standard Gemma 4 decoder (mixed sliding/global attention, `global_head_dim=512`, optional KV sharing).
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+```bash
+# Text LoRA (1x96GB @ ~27.76 GiB)
+axolotl train examples/gemma4-unified/12b-text-lora.yaml
+
+# Vision LoRA (1x96GB @ ~26.5 GiB)
+axolotl train examples/gemma4-unified/12b-vision-lora.yaml
+```
+
+## Limitations
+
+- **Attention**: FA2 (max head_dim=256) / FA4 (max head_dim=128) cannot serve `global_head_dim=512` on their own. For `sample_packing: true`, use `flex_attention` or `gemma4_hybrid_attn_impl: true` (with `flash_attention` varlen for sliding-window layers + block-diagonal `sdpa` for global layers, packing-safe).
+- **lora_target_linear**: incompatible for multimodal. Use `lora_target_modules` as seen in vision example.
+
+### TIPS
+
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+- The multimodal dataset format follows the OpenAI multi-content Messages format as seen [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [Gemma 4 Blog](https://blog.google/innovation-and-ai/technology/developers-tools/introducing-gemma-4-12B/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/gemma4/26b-a4b-moe-bnb-lora.yaml b/examples/gemma4/26b-a4b-moe-bnb-lora.yaml
new file mode 100644
index 0000000000..9867080ac7
--- /dev/null
+++ b/examples/gemma4/26b-a4b-moe-bnb-lora.yaml
@@ -0,0 +1,103 @@
+# Gemma-4-26B-A4B MoE QLoRA on a bf16 checkpoint quantized to bnb-4bit at load — the throughput-
+# optimized recipe for the no-NVFP4-checkpoint case. sm120 (RTX PRO 6000), 4k seq, experts-only +
+# attn LoRA r=16. ~1940 tok/s / ~21 GiB on one GPU (vs nvfp4-marlin's ~2350 / ~24 GiB).
+#
+# The three levers that matter most here:
+#   1. attn_implementation: flash_attention_2 + gemma4_hybrid_attn_impl: true
+#        FA2 (via the `kernels` lib — no flash-attn source build) on the sliding-window layers +
+#        SDPA on the global layers. ~2x over plain sdpa, which otherwise computes full attention
+#        over the packed 4k sequence and dominates runtime.
+#   2. quantize_moe_experts: true (+ load_in_4bit)
+#        Stores the MoE experts as bnb-4bit at load. scattermoe dequants only the active experts
+#        per step and recomputes that bf16 in backward (recipe) instead of saving it — so the
+#        expert weights stay 4-bit-resident and never pin per-layer (no ~60 GiB blowup).
+#   3. moe_bnb_fast (default on)
+#        Routes bnb experts through the 1-launch parallel_linear (scatter2scatter) path. Faster
+#        than the chunked torch._grouped_mm path at the same low memory. Set false to fall back to
+#        chunked for large-expert MoEs / tiny GPUs where the single-shot active-expert dequant is
+#        too big.
+base_model: google/gemma-4-26B-A4B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+  - axolotl.integrations.liger.LigerPlugin
+use_kernels: true
+use_scattermoe: true
+experts_implementation: scattermoe
+liger_layer_norm: true
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_rms_norm_gated: true
+strict: false
+
+ddp_find_unused_parameters: true
+
+chat_template: gemma4
+eot_tokens:
+  - "<turn|>"
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:10%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.05
+output_dir: ./outputs/gemma4-26b-a4b-bnb-lora
+
+sequence_len: 4096
+sample_packing: true
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+# text-backbone attention + MLP projections (skip vision/audio encoders)
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+# MoE expert LoRA (3D Parameter tensors, not nn.Linear)
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+# Fused LoRA kernels for the NON-expert path: the dense shared MLP (Gemma4TextMLP) and the
+# attention q/k/v/o projections. The routed experts are handled by scattermoe; these fuse the
+# remaining per-layer LoRA. fused_attn_kernel lets qkv/o attach on the gemma4_unified variant too.
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+fused_attn_kernel: true
+
+bnb_config_kwargs:
+  bnb_4bit_use_double_quant: true
+
+# moe_bnb_fast: true is the default (1-launch parallel_linear, recompute-in-backward). Set false
+# to fall back to the chunked path for large-expert MoEs / tiny GPUs (bound via moe_dequant_chunk_size).
+# moe_bnb_fast: false
+# moe_dequant_chunk_size: 32
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_8bit       # GPU 8-bit Adam: ~2.6 GiB less than fp32 Adam at r=16
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+
+# gemma4 head_dim is 256. flex_attention OOMs sm120 shared memory here; the hybrid impl runs FA2 on
+# the sliding-window layers + SDPA on the global layers. Requires attn_implementation=flash_attention_2.
+attn_implementation: flash_attention_2
+gemma4_hybrid_attn_impl: true
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/gemma4/26b-a4b-moe-nvfp4-lora.yaml b/examples/gemma4/26b-a4b-moe-nvfp4-lora.yaml
new file mode 100644
index 0000000000..ba9871abb9
--- /dev/null
+++ b/examples/gemma4/26b-a4b-moe-nvfp4-lora.yaml
@@ -0,0 +1,72 @@
+# Gemma-4-26B-A4B MoE LoRA on an NVFP4 (modelopt) checkpoint via ScatterMoE. Single GPU (sm120).
+base_model: nvidia/Gemma-4-26B-A4B-NVFP4
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+  - axolotl.integrations.liger.LigerPlugin
+use_kernels: true
+use_scattermoe: true
+experts_implementation: scattermoe
+dsv4_fp4_grouped_mode: nvfp4
+liger_layer_norm: true
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_rms_norm_gated: true
+strict: false
+
+ddp_find_unused_parameters: true
+
+chat_template: gemma4
+eot_tokens:
+  - "<turn|>"
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:10%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.05
+output_dir: ./outputs/gemma4-26b-a4b-nvfp4-lora
+
+sequence_len: 4096
+sample_packing: true
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+# text backbone only (skip vision/audio encoders)
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+fused_attn_kernel: true
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+
+# flex_attention OOMs on sm120 here; the hybrid runs FA2 (sliding) + SDPA (global). Needs attn_implementation=flash_attention_2.
+attn_implementation: flash_attention_2
+gemma4_hybrid_attn_impl: true
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/gemma4/26b-a4b-moe-qlora.yaml b/examples/gemma4/26b-a4b-moe-qlora.yaml
new file mode 100644
index 0000000000..091437ed09
--- /dev/null
+++ b/examples/gemma4/26b-a4b-moe-qlora.yaml
@@ -0,0 +1,105 @@
+base_model: google/gemma-4-26B-A4B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+  - axolotl.integrations.liger.LigerPlugin
+use_kernels: true
+use_scattermoe: true
+experts_implementation: scattermoe
+torch_compile: true
+liger_layer_norm: true
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_rms_norm_gated: true
+strict: false
+
+# Multi-GPU (DDP) only: LoRA targets the text backbone, so the frozen vision/
+# audio encoders and Gemma4's KV-sharing layers leave some adapter params
+# without gradients. Required to avoid "parameters that were not used in
+# producing loss". No effect on single-GPU runs.
+ddp_find_unused_parameters: true
+
+chat_template: gemma4
+eot_tokens:
+  - "<turn|>"
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:10%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.05
+output_dir: ./outputs/gemma4-26b-a4b-qlora
+
+sequence_len: 2048
+sample_packing: true
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+
+# Restrict LoRA to text backbone only (skip vision/audio encoders)
+# using regex to match only the text decoder attention projections.
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+# MoE expert LoRA (3D Parameter tensors, not nn.Linear)
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+
+# Fused LoRA kernels for the NON-expert path: the dense shared MLP (Gemma4TextMLP) and the
+# attention q/k/v/o projections. The routed experts are handled by scattermoe; these fuse the
+# remaining per-layer LoRA. fused_attn_kernel lets qkv/o attach on the gemma4_unified variant too.
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+fused_attn_kernel: true
+
+bnb_config_kwargs:
+  bnb_4bit_use_double_quant: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+activation_offloading: true
+logging_steps: 1
+
+# Hybrid attention is the biggest throughput lever (~2x vs plain sdpa, which computes full attention
+# over the packed 4k sequence): FA2 (via the `kernels` lib, no source build) on the sliding-window
+# layers + SDPA on the global layers. flex_attention OOMs sm120 shared memory here.
+attn_implementation: flash_attention_2
+gemma4_hybrid_attn_impl: true
+
+# bnb-4bit MoE experts have no fused 4-bit-read kernel, so they're dequantized to bf16 by the
+# default 1-launch parallel_linear path (the dequant'd bf16 is recomputed in backward, not saved,
+# so it stays low-memory). For large-expert MoEs / tiny GPUs, set moe_bnb_fast: false to fall back
+# to the chunked path and bound the per-pass dequant transient with moe_dequant_chunk_size (32).
+# moe_bnb_fast: false
+# moe_dequant_chunk_size: 32
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/gemma4/31b-qlora.yaml b/examples/gemma4/31b-qlora.yaml
new file mode 100644
index 0000000000..b8613dd49e
--- /dev/null
+++ b/examples/gemma4/31b-qlora.yaml
@@ -0,0 +1,70 @@
+base_model: google/gemma-4-31B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.liger.LigerPlugin
+liger_layer_norm: true
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_rms_norm_gated: true
+strict: false
+
+chat_template: gemma4
+eot_tokens:
+  - "<turn|>"
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:10%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.05
+output_dir: ./outputs/gemma4-31b-qlora
+
+sequence_len: 2048
+sample_packing: true
+
+load_in_4bit: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+
+# Restrict LoRA to text backbone only (skip vision/audio encoders)
+# using regex to match only the text decoder attention projections.
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+bnb_config_kwargs:
+  bnb_4bit_use_double_quant: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+activation_offloading: true
+logging_steps: 1
+
+gemma4_hybrid_attn_impl: true
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/gemma4/README.md b/examples/gemma4/README.md
new file mode 100644
index 0000000000..36db373e46
--- /dev/null
+++ b/examples/gemma4/README.md
@@ -0,0 +1,50 @@
+# Finetune Google's Gemma 4 with Axolotl
+
+[Gemma 4](https://huggingface.co/collections/google/gemma-4) is a family of multimodal models from Google. This guide covers how to train them with Axolotl.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+```bash
+# 26B MoE QLoRA (1x80GB)
+axolotl train examples/gemma4/26b-a4b-moe-qlora.yaml
+
+# 31B Dense QLoRA (1x80GB @ ~25.2 GiB)
+axolotl train examples/gemma4/31b-qlora.yaml
+
+# E2B vision LoRA (1x80GB @ ~10.4 GiB)
+axolotl train examples/gemma4/e2b-vision-lora.yaml
+```
+
+### MoE Expert Quantization & Expert LoRA (26B-A4B only)
+
+The 26B-A4B config uses ScatterMoE kernels via the transformers `ExpertsInterface` and quantizes expert weights on load. To learn about expert quantization, expert LoRA targeting, and related limitations, see the [MoE Expert Quantization](https://docs.axolotl.ai/docs/expert_quantization.html) docs.
+
+## Limitations
+
+- **Flash Attention**: FA2 (max head_dim=256) and FA4 (max head_dim=128) cannot serve Gemma 4's `global_head_dim=512` on their own. Use `flex_attention`, or `gemma4_hybrid_attn_impl: true` to run the sliding-window layers under FA2 and the global (head_dim=512) layers under `sdpa` (requires `attn_implementation: flash_attention_2` and a flash-attn build for your GPU arch).
+- **LoRA kernels**: Not supported for models with KV-sharing layers.
+- **lora_target_linear**: Incompatible for multimodal models; use `lora_target_modules` with a regex to restrict LoRA to the text backbone.
+
+### TIPS
+
+- `gemma4_hybrid_attn_impl: true` trains ~2× faster than `flex_attention` on 31B (~25.2 GiB reserved, packing on) and avoids the flex `head_dim=512` kernel, which can exhaust shared memory on Blackwell.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- You can run full finetuning by removing `adapter: qlora`, `load_in_4bit: true`, and `quantize_moe_experts: true` from the config. This is heavy and has not been tested.
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [Gemma 4 Blog](https://huggingface.co/blog/gemma4)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/gemma4/e2b-vision-lora.yaml b/examples/gemma4/e2b-vision-lora.yaml
new file mode 100644
index 0000000000..e48b192105
--- /dev/null
+++ b/examples/gemma4/e2b-vision-lora.yaml
@@ -0,0 +1,59 @@
+base_model: google/gemma-4-E2B-it
+processor_type: AutoProcessor
+freeze_mm_modules: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+strict: false
+
+# Required for vision/multimodal training
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: gemma4
+eot_tokens:
+  - "<turn|>"
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:100]
+
+val_set_size: 0
+output_dir: ./outputs/gemma4-e2b-vision-lora
+
+adapter: lora
+sequence_len: 2048
+pad_to_sequence_len: false
+
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+# Target language model only — vision encoder is frozen via freeze_mm_modules
+lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+max_steps: 10
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+logging_steps: 1
+attn_implementation: sdpa
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
diff --git a/examples/glm4/qlora-32b.yaml b/examples/glm4/qlora-32b.yaml
new file mode 100644
index 0000000000..151820924d
--- /dev/null
+++ b/examples/glm4/qlora-32b.yaml
@@ -0,0 +1,64 @@
+base_model: THUDM/GLM-4-32B-0414
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+load_in_4bit: true
+
+datasets:
+  - path: teknium/GPT4-LLM-Cleaned
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/qlora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+loss_watchdog_threshold: 5.0
+loss_watchdog_patience: 3
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/glm45/README.md b/examples/glm45/README.md
new file mode 100644
index 0000000000..1659638255
--- /dev/null
+++ b/examples/glm45/README.md
@@ -0,0 +1,80 @@
+# Finetune Z.ai's GLM-4.5-Air with Axolotl
+
+[GLM-4.5-Air](https://huggingface.co/zai-org/GLM-4.5-Air) is a MoE model by Z.ai.
+
+This guide shows how to fine-tune it with Axolotl.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+```bash
+# QLoRA (1x80GB @ ~63.4GiB/GPU)
+axolotl train examples/glm45/glm-45-air-qlora.yaml
+```
+
+### Dataset
+
+In addition to the standard OpenAI Messages format, GLM-4.5 supports an extra parameter for thinking in the assistant section.
+
+```json
+{
+    "role": "assistant",
+    "reasoning_content": "...",  // or have </think>...</think> in `content`
+    "content": "..."
+}
+```
+
+Make sure you set the below extra attributes if needed:
+
+```yaml
+datasets:
+  - path: ...
+    type: chat_template
+    message_property_mappings:
+      role: role
+      content: content
+
+    #   tool_calls: tool_calls  # uncomment if using tools
+    #   reasoning_content: reasoning_content  # uncomment if have reasoning
+
+# Uncomment if training on tool role (you would rarely if ever need this)
+# eot_tokens:
+#   - <|observation|>
+```
+
+### Tips
+
+- The role name for tools in this template is `tool`.
+- You will see this Axolotl WARNING — this is expected as the template does not use EOS:
+  ```
+  EOS token '<|endoftext|>' not found in chat_template. Please check if your template/EOS token is correct.
+  ```
+- You can run a full finetuning by removing `adapter: qlora`, `load_in_4bit: true`, and `quantize_moe_experts: true` from the config.
+- **LoRA kernels**: Incompatible with this model. Must be explicitly disabled (`lora_*_kernel: false`).
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+
+### GGUF / llama.cpp loading error (missing tensors)
+
+If you see `missing tensor 'blk.X.attn_norm.weight'` when loading a GLM-4 / GLM4-MoE model in llama.cpp, this is likely
+caused by `num_nextn_predict_layers` being set to `1` in `config.json` while the MTP weights were not exported (possible
+after PEFT/QLoRA training).
+
+**Fix:** Set `"num_nextn_predict_layers": 0` in your `config.json` before converting to GGUF.
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [GLM-4.5-Air on HuggingFace](https://huggingface.co/zai-org/GLM-4.5-Air)
+- [GLM-4.5 Blog](https://z.ai/blog/glm-4.5)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/glm45/glm-45-air-qlora.yaml b/examples/glm45/glm-45-air-qlora.yaml
new file mode 100644
index 0000000000..5723d3c45c
--- /dev/null
+++ b/examples/glm45/glm-45-air-qlora.yaml
@@ -0,0 +1,64 @@
+base_model: zai-org/GLM-4.5-Air
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+quantize_moe_experts: true # important
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 16
+lora_alpha: 8
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/glm46v/README.md b/examples/glm46v/README.md
new file mode 100644
index 0000000000..965e08e517
--- /dev/null
+++ b/examples/glm46v/README.md
@@ -0,0 +1,44 @@
+# Finetune GLM-4.6V with Axolotl
+
+GLM-4.6V is a family of vision-language models from ZhipuAI found on [HuggingFace](https://huggingface.co/zai-org/GLM-4.6V). This guide shows how to fine-tune it with Axolotl for vision-language tasks.
+
+
+
+## Getting started
+
+1. Install Axolotl from source following the [installation guide](https://docs.axolotl.ai/docs/installation.html#sec-edge-build).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+
+3. Run the fine-tuning:
+
+    glm-4-6v-flash(9B)
+    ```bash
+    axolotl train examples/glm46v/glm-4-6v-flash-qlora.yaml
+    ```
+
+Let us know how it goes. Happy finetuning! 🚀
+
+## Tips
+
+- Vision datasets should follow the format described in the [multimodal docs](https://docs.axolotl.ai/docs/multimodal.html#dataset-format)
+- You can run a **full finetuning** by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset in the [dataset loading docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+
+## Supported Models
+
+- **GLM-4.6V**: Full vision-language model (`zai-org/GLM-4.6V`)
+- **GLM-4.6V-Flash**: Faster variant (`zai-org/GLM-4.6V-Flash`)
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [ZhipuAI GLM-4.6V](https://huggingface.co/zai-org/GLM-4.6V)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/glm46v/glm-4-6v-flash-ddp.yaml b/examples/glm46v/glm-4-6v-flash-ddp.yaml
new file mode 100644
index 0000000000..274f041a32
--- /dev/null
+++ b/examples/glm46v/glm-4-6v-flash-ddp.yaml
@@ -0,0 +1,53 @@
+base_model: zai-org/GLM-4.6V-Flash
+trust_remote_code: true
+
+processor_type: AutoProcessor
+load_in_4bit: true
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+ddp_find_unused_parameters: true
+
+output_dir: ./outputs/glm-4-6v-flash-qlora
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+sequence_len: 2048
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+logging_steps: 1
+attn_implementation: sdpa
+
+warmup_ratio: 0.1
+evals_per_epoch: 0
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/glm46v/glm-4-6v-flash-qlora.yaml b/examples/glm46v/glm-4-6v-flash-qlora.yaml
new file mode 100644
index 0000000000..9fe8d6e436
--- /dev/null
+++ b/examples/glm46v/glm-4-6v-flash-qlora.yaml
@@ -0,0 +1,50 @@
+base_model: zai-org/GLM-4.6V-Flash
+trust_remote_code: true
+
+processor_type: AutoProcessor
+load_in_4bit: true
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+output_dir: ./outputs/glm-4-6v-flash-qlora
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+sequence_len: 2048
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+logging_steps: 1
+attn_implementation: sdpa
+
+warmup_ratio: 0.1
+evals_per_epoch: 0
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/glm47-flash/README.md b/examples/glm47-flash/README.md
new file mode 100644
index 0000000000..2e5e21010a
--- /dev/null
+++ b/examples/glm47-flash/README.md
@@ -0,0 +1,65 @@
+# Finetune Z.ai's GLM-4.7-Flash with Axolotl
+
+[GLM-4.7-Flash](https://huggingface.co/zai-org/GLM-4.7-Flash) is a 30B-A3B MoE model by Z.ai.
+
+This guide shows how to fine-tune it with Axolotl.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+```bash
+# QLoRA
+# - no target experts (1x48GB @ ~24GiB/GPU)
+# - target experts (1x48GB @ ~34GiB/GPU)
+axolotl train examples/glm47-flash/qlora.yaml
+
+# QLoRA FSDP2 no target experts (2x48GB @ ~29GiB/GPU)
+axolotl train examples/glm47-flash/qlora_fsdp.yaml
+```
+
+```bash
+# LoRA
+# - no target experts (1x48GB @ ~35GiB/GPU)
+# - target experts (1x48GB @ OOM. Projected ~45-50GiB/GPU)
+axolotl train examples/glm47-flash/lora.yaml
+
+# LoRA FSDP2 no target experts (2x48GB @ ~43GiB/GPU)
+axolotl train examples/glm47-flash/lora_fsdp.yaml
+```
+
+### MoE Expert Quantization & Expert LoRA
+
+This model quantize expert weights on load. To learn about expert quantization, expert LoRA targeting, and related limitations, see the [MoE Expert Quantization](https://docs.axolotl.ai/docs/expert_quantization.html) docs.
+
+## Limitations
+
+- **lora_target_linear**: Incompatible for this model.
+- **LoRA kernels**: Incompatible with this model due to non-standard attention projections (DSA). Must be explicitly disabled (`lora_*_kernel: false`).
+
+
+### TIPS
+
+- For inference, the official Z.ai team recommends these default settings (most tasks):
+  - `temperature: 1.0`
+  - `top_p: 0.95`
+  - `max_new_tokens: 131072`
+- You can run a full finetuning by removing `adapter: qlora`, `load_in_4bit: true`, and `quantize_moe_experts: true` from the config. This is heavy, so we have not tested this.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [GLM-4.7-Flash on HuggingFace](https://huggingface.co/zai-org/GLM-4.7-Flash)
+- [GLM-4.7 Blog](https://z.ai/blog/glm-4.7)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/glm47-flash/lora.yaml b/examples/glm47-flash/lora.yaml
new file mode 100644
index 0000000000..5f3de36e9a
--- /dev/null
+++ b/examples/glm47-flash/lora.yaml
@@ -0,0 +1,65 @@
+base_model: zai-org/GLM-4.7-Flash
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: true
+quantize_moe_experts: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/glm4.7-flash-lora-8bit-out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+# Uncomment to also target MoE expert weights:
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+# LoRA kernels incompatible with DSA attention
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
diff --git a/examples/glm47-flash/lora_fsdp.yaml b/examples/glm47-flash/lora_fsdp.yaml
new file mode 100644
index 0000000000..cf1d2de55a
--- /dev/null
+++ b/examples/glm47-flash/lora_fsdp.yaml
@@ -0,0 +1,75 @@
+base_model: zai-org/GLM-4.7-Flash
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: true
+quantize_moe_experts: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/glm4.7-flash-lora-8bit-fsdp-out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+# Uncomment to also target MoE expert weights:
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+# LoRA kernels incompatible with DSA attention
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+fsdp_config:
+  fsdp_version: 2
+  offload_params: false
+  cpu_ram_efficient_loading: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Glm4MoeLiteDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
diff --git a/examples/glm47-flash/qlora.yaml b/examples/glm47-flash/qlora.yaml
new file mode 100644
index 0000000000..a05bf54d2f
--- /dev/null
+++ b/examples/glm47-flash/qlora.yaml
@@ -0,0 +1,65 @@
+base_model: zai-org/GLM-4.7-Flash
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_4bit: true
+quantize_moe_experts: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/glm4.7-flash-qlora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+# Uncomment to also target MoE expert weights:
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+# LoRA kernels incompatible with DSA attention
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
diff --git a/examples/glm47-flash/qlora_fsdp.yaml b/examples/glm47-flash/qlora_fsdp.yaml
new file mode 100644
index 0000000000..9ad5a6212c
--- /dev/null
+++ b/examples/glm47-flash/qlora_fsdp.yaml
@@ -0,0 +1,75 @@
+base_model: zai-org/GLM-4.7-Flash
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_4bit: true
+quantize_moe_experts: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/glm4.7-flash-qlora-fsdp-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+# Uncomment to also target MoE expert weights:
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+# LoRA kernels incompatible with DSA attention
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+fsdp_config:
+  fsdp_version: 2
+  offload_params: false
+  cpu_ram_efficient_loading: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Glm4MoeLiteDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
diff --git a/examples/glm_moe_dsa/glm-5.2-nvfp4-lora.yaml b/examples/glm_moe_dsa/glm-5.2-nvfp4-lora.yaml
new file mode 100644
index 0000000000..6c480f9017
--- /dev/null
+++ b/examples/glm_moe_dsa/glm-5.2-nvfp4-lora.yaml
@@ -0,0 +1,79 @@
+# GLM-5.2 (glm_moe_dsa) MoE LoRA on an NVFP4 (modelopt) checkpoint via ScatterMoE. ~500B, FSDP2 multi-GPU (2xB200 / 8xH100).
+base_model: Mapika/GLM-5.2-NVFP4
+
+plugins:
+  # Expert parallelism (requires the deep_ep package). Drop this plugin + the expert_parallel_* keys below for plain FSDP2.
+  - axolotl.integrations.expert_parallel.ExpertParallelPlugin
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+use_kernels: true
+use_scattermoe: true
+experts_implementation: scattermoe
+dsv4_fp4_grouped_mode: nvfp4
+use_glm_dsa_kernels: true
+strict: false
+
+attn_implementation: sdpa
+
+# Pure EP: set expert_parallel_size to the total GPU count (256 must divide it evenly), dp_shard_size 1.
+expert_parallel_size: 8                      # = total GPU count (8 for 8xH100; 2 for 2xB200)
+dp_shard_size: 1
+expert_parallel_num_nvl_bytes: 4294967296    # 4 GiB DeepEP intranode (NVLink) dispatch/combine buffer
+# Caps tokens per expert to guard the DeepEP NVLink combine (it hangs if an expert exceeds the buffer).
+# The 1M default is effectively off at normal lengths; lower it toward the hang threshold for long context (>=64k).
+expert_parallel_token_capacity: 1000000
+
+# Long context: shard the sequence with context parallelism (set ep * cp == world_size).
+# context_parallel_size: 2
+
+chat_template: tokenizer_default
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:5%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.02
+output_dir: ./outputs/glm-5.2-nvfp4-lora
+
+sequence_len: 4096
+sample_packing: true
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+# MLA attention + shared/dense MLP (not the no-grad indexer).
+lora_target_modules: 'model.layers.[\d]+.(self_attn.(q_a_proj|q_b_proj|kv_a_proj_with_mqa|kv_b_proj|o_proj)|mlp.(shared_experts.)?(gate_proj|up_proj|down_proj))'
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+gradient_checkpointing: true
+# Keep off under EP (headroom); set true for plain FSDP2 / tight memory.
+activation_offloading: false
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true   # required: rank0-only NVFP4 load
+  state_dict_type: SHARDED_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: GlmMoeDsaDecoderLayer
+  reshard_after_forward: true
diff --git a/examples/gpt-oss/README.md b/examples/gpt-oss/README.md
new file mode 100644
index 0000000000..ae71eec1ef
--- /dev/null
+++ b/examples/gpt-oss/README.md
@@ -0,0 +1,146 @@
+# Finetune OpenAI's GPT-OSS with Axolotl
+
+[GPT-OSS](https://huggingface.co/collections/openai/gpt-oss-68911959590a1634ba11c7a4) are a family of open-weight MoE models trained by OpenAI, released in August 2025. There are two variants: 20B and 120B.
+
+In October 2025, OpenAI released safeguard models built upon GPT-OSS called [GPT-OSS-Safeguard](https://huggingface.co/collections/openai/gpt-oss-safeguard). They use the same architecture, so the same examples below can be re-used.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+uv pip install --no-build-isolation 'axolotl>=0.16.1'
+```
+
+2. Choose one of the following configs below for training the 20B model. (for 120B, see [below](#training-120b))
+
+```bash
+# LoRA SFT linear layers (1x48GB @ ~44GiB)
+axolotl train examples/gpt-oss/gpt-oss-20b-sft-lora-singlegpu.yaml
+
+# FFT SFT with offloading (2x24GB @ ~21GiB/GPU)
+axolotl train examples/gpt-oss/gpt-oss-20b-fft-fsdp2-offload.yaml
+
+# FFT SFT (8x48GB @ ~36GiB/GPU or 4x80GB @ ~46GiB/GPU)
+axolotl train examples/gpt-oss/gpt-oss-20b-fft-fsdp2.yaml
+```
+
+Note: Memory usage taken from `device_mem_reserved(gib)` from logs.
+
+### Training 120B
+
+On 8xH100s, make sure you have ~3TB of free disk space. With each checkpoint clocking in at ~720GB, along with the base
+model, and final model output, you may need at least 3TB of free disk space to keep at least 2 checkpoints.
+
+```bash
+# FFT SFT with offloading (8x80GB @ ~49GiB/GPU)
+axolotl train examples/gpt-oss/gpt-oss-120b-fft-fsdp2-offload.yaml
+```
+
+To simplify fine-tuning across 2 nodes × 8x H100 (80GB) GPUs, we've partnered with [Baseten](https://baseten.co) to showcase multi-node
+training of the 120B model using Baseten Truss. You can read more about this recipe on
+[Baseten's blog](https://www.baseten.co/blog/how-to-fine-tune-gpt-oss-120b-with-baseten-and-axolotl/). The recipe can
+be found on their
+[GitHub](https://github.com/basetenlabs/ml-cookbook/tree/main/examples/oss-gpt-120b-axolotl/training).
+
+ERRATA: Transformers saves the model Architecture prefixed with `FSDP` which needs to be manually renamed in `config.json`.
+See https://github.com/huggingface/transformers/pull/40207 for the status of this issue.
+
+```bash
+sed -i 's/FSDPGptOssForCausalLM/GptOssForCausalLM/g' ./outputs/gpt-oss-out/config.json
+```
+
+When using SHARDED_STATE_DICT with FSDP, the final checkpoint should automatically merge the sharded weights to your
+configured `output_dir`. However, if that step fails due to a disk space error, you can take an additional step to
+merge the sharded weights.  This step will automatically determine the last checkpoint directory and merge the sharded
+weights to `{output_dir}/merged`.
+
+```bash
+axolotl merge-sharded-fsdp-weights examples/gpt-oss/gpt-oss-120b-fft-fsdp2-offload.yaml
+mv ./outputs/gpt-oss-out/merged/* ./outputs/gpt-oss-out/
+```
+
+### How to set reasoning_effort in template?
+
+The harmony template has a feature to set the `reasoning_effort` during prompt building. The default is `medium`. If you would like to adjust this, you can add the following to your config:
+
+```yaml
+chat_template_kwargs:
+  reasoning_effort: "high"  # low | medium | high
+```
+
+Currently, this applies globally. There is no method to apply per sample yet. If you are interested in adding this, please feel free to create an Issue to discuss.
+
+### Inferencing your fine-tuned model
+
+#### vLLM
+
+GPT-OSS support in vLLM does not exist in a stable release yet. See https://x.com/MaziyarPanahi/status/1955741905515323425
+for more information about using a special vllm-openai docker image for inferencing with vLLM.
+
+Optionally, vLLM can be installed from nightly:
+
+```bash
+uv pip install --no-build-isolation --pre -U vllm --extra-index-url https://wheels.vllm.ai/nightly
+```
+and the vLLM server can be started with the following command (modify `--tensor-parallel-size 8` to match your environment):
+```bash
+vllm serve ./outputs/gpt-oss-out/ --served-model-name axolotl/gpt-oss-20b --host 0.0.0.0 --port 8888  --tensor-parallel-size 8
+```
+
+#### SGLang
+
+SGLang has 0-day support in main, see https://github.com/sgl-project/sglang/issues/8833 for infomation on installing
+SGLang from source. Once you've installed SGLang, run the following command to launch a SGLang server:
+
+```bash
+python3 -m sglang.launch_server --model ./outputs/gpt-oss-out/ --served-model-name axolotl/gpt-oss-120b --host 0.0.0.0 --port 8888 --tp 8
+```
+
+### Tool use
+
+GPT-OSS has a comprehensive tool understanding. Axolotl supports tool calling datasets for Supervised Fine-tuning.
+
+Here is an example dataset config:
+```yaml
+datasets:
+  - path: Nanobit/text-tools-2k-test
+    type: chat_template
+```
+
+See [Nanobit/text-tools-2k-test](https://huggingface.co/datasets/Nanobit/text-tools-2k-test) for the sample dataset.
+
+Refer to [our docs](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#using-tool-use) for more info.
+
+### Thinking and chat_template masking conflict
+
+OpenAI’s Harmony template hides `thinking` in all non-final turns, which conflicts with Axolotl’s `chat_template` masking.
+
+If your dataset has `thinking` content mid-turn, there are two paths we recommend:
+
+- Train only on the last turn. This can be accomplished via chat_template's [train on last doc](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#training-on-last-message).
+
+- Adjust your dataset to only have `thinking` content in the last turn.
+
+### TIPS
+
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+
+## Related Resources
+
+- [GPT-OSS Blog](https://openai.com/index/introducing-gpt-oss/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/gpt-oss/gpt-oss-120b-fft-fsdp2-offload.yaml b/examples/gpt-oss/gpt-oss-120b-fft-fsdp2-offload.yaml
new file mode 100644
index 0000000000..71692958f0
--- /dev/null
+++ b/examples/gpt-oss/gpt-oss-120b-fft-fsdp2-offload.yaml
@@ -0,0 +1,71 @@
+# the original mxfp4 quantized model is not supported with FSDP cpu_ram_efficient_loading
+# FSDP cpu_ram_efficient_loading is used to reduce the initial CPU memory usage when loading the model
+base_model: axolotl-ai-co/gpt-oss-120b-dequantized
+
+use_kernels: false
+
+dp_shard_size: 16  # requires 2x8xH100 nodes
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+experimental_skip_move_to_device: true  # prevent OOM by NOT putting model to GPU before sharding
+
+datasets:
+  - path: HuggingFaceH4/Multilingual-Thinking
+    type: chat_template
+    field_thinking: thinking
+    template_thinking_key: thinking
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/gpt-oss-out/
+save_total_limit: 2  # the 120B model can use up to 720GB of disk space per checkpoint, so let's only keep the last 2
+
+sequence_len: 4096
+sample_packing: true
+pad_to_sequence_len: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+trackio_project_name:
+trackio_run_name:
+trackio_space_id:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+
+optimizer: adamw_torch_fused  # 8bit optimizers do not work with FSDP2 offload
+lr_scheduler: constant_with_warmup
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+attn_implementation: kernels-community/vllm-flash-attn3  # this is not needed if using flash_attn >= 2.8.3
+
+gradient_checkpointing: true
+activation_offloading: true
+
+logging_steps: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.03
+
+special_tokens:
+eot_tokens:
+  - "<|end|>"
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: true
+  state_dict_type: SHARDED_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: GptOssDecoderLayer
+  reshard_after_forward: true
+  cpu_ram_efficient_loading: true
diff --git a/examples/gpt-oss/gpt-oss-20b-fft-deepspeed-zero3.yaml b/examples/gpt-oss/gpt-oss-20b-fft-deepspeed-zero3.yaml
new file mode 100644
index 0000000000..5912f876bb
--- /dev/null
+++ b/examples/gpt-oss/gpt-oss-20b-fft-deepspeed-zero3.yaml
@@ -0,0 +1,61 @@
+base_model: openai/gpt-oss-20b
+use_kernels: false
+model_quantization_config: Mxfp4Config
+model_quantization_config_kwargs:
+  dequantize: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+experimental_skip_move_to_device: true  # prevent OOM by NOT putting model to GPU before sharding
+
+datasets:
+  - path: HuggingFaceH4/Multilingual-Thinking
+    type: chat_template
+    field_thinking: thinking
+    template_thinking_key: thinking
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/gpt-oss-out/
+
+sequence_len: 4096
+sample_packing: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+trackio_project_name:
+trackio_run_name:
+trackio_space_id:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+
+optimizer: adamw_torch_8bit
+lr_scheduler: constant_with_warmup
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+attn_implementation: kernels-community/vllm-flash-attn3  # this is not needed if using flash_attn >= 2.8.3
+
+gradient_checkpointing: true
+activation_offloading: true
+
+logging_steps: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.03
+
+special_tokens:
+eot_tokens:
+  - "<|end|>"
+
+# choose the zero3 configuration that best fits your system capabilities
+deepspeed: deepspeed_configs/zero3_bf16.json
diff --git a/examples/gpt-oss/gpt-oss-20b-fft-fsdp2-offload.yaml b/examples/gpt-oss/gpt-oss-20b-fft-fsdp2-offload.yaml
new file mode 100644
index 0000000000..b1a0fef4ab
--- /dev/null
+++ b/examples/gpt-oss/gpt-oss-20b-fft-fsdp2-offload.yaml
@@ -0,0 +1,71 @@
+base_model: openai/gpt-oss-20b
+use_kernels: true
+model_quantization_config: Mxfp4Config
+model_quantization_config_kwargs:
+  dequantize: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+experimental_skip_move_to_device: true  # prevent OOM by NOT putting model to GPU before sharding
+
+datasets:
+  - path: HuggingFaceH4/Multilingual-Thinking
+    type: chat_template
+    field_thinking: thinking
+    template_thinking_key: thinking
+
+dataset_prepared_path: ./outputs/last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/gpt-oss-out/
+
+sequence_len: 4096
+sample_packing: true
+pad_to_sequence_len: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+trackio_project_name:
+trackio_run_name:
+trackio_space_id:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+
+optimizer: adamw_torch_fused  # 8bit optimizers do not work with FSDP2 offload
+lr_scheduler: constant_with_warmup
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+attn_implementation: kernels-community/vllm-flash-attn3  # this is not needed if using flash_attn >= 2.8.3
+
+gradient_checkpointing: true
+activation_offloading: true
+
+logging_steps: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.03
+
+special_tokens:
+eot_tokens:
+  - "<|end|>"
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: true
+  state_dict_type: SHARDED_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: GptOssDecoderLayer
+  reshard_after_forward: true
+  #  cpu_ram_efficient_loading: true
+
+# cpu_ram_efficient_loading cannot be used with MXFP4 model quantization.
+# It can only be used with a dequantized model like `axolotl-ai-co/gpt-oss-120b-dequantized`
diff --git a/examples/gpt-oss/gpt-oss-20b-fft-fsdp2.yaml b/examples/gpt-oss/gpt-oss-20b-fft-fsdp2.yaml
new file mode 100644
index 0000000000..f97174cd90
--- /dev/null
+++ b/examples/gpt-oss/gpt-oss-20b-fft-fsdp2.yaml
@@ -0,0 +1,67 @@
+base_model: openai/gpt-oss-20b
+use_kernels: false
+model_quantization_config: Mxfp4Config
+model_quantization_config_kwargs:
+  dequantize: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+experimental_skip_move_to_device: true  # prevent OOM by NOT putting model to GPU before sharding
+
+datasets:
+  - path: HuggingFaceH4/Multilingual-Thinking
+    type: chat_template
+    field_thinking: thinking
+    template_thinking_key: thinking
+
+dataset_prepared_path: ./outputs/last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/gpt-oss-out/
+
+sequence_len: 4096
+sample_packing: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+trackio_project_name:
+trackio_run_name:
+trackio_space_id:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+
+optimizer: adamw_torch_8bit
+lr_scheduler: constant_with_warmup
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+attn_implementation: kernels-community/vllm-flash-attn3  # this is not needed if using flash_attn >= 2.8.3
+
+gradient_checkpointing: true
+activation_offloading: true
+
+logging_steps: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.03
+
+special_tokens:
+eot_tokens:
+  - "<|end|>"
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  state_dict_type: SHARDED_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: GptOssDecoderLayer
+  reshard_after_forward: true
+#  cpu_ram_efficient_loading: true
diff --git a/examples/gpt-oss/gpt-oss-20b-sft-lora-singlegpu.yaml b/examples/gpt-oss/gpt-oss-20b-sft-lora-singlegpu.yaml
new file mode 100644
index 0000000000..122fb0b6c3
--- /dev/null
+++ b/examples/gpt-oss/gpt-oss-20b-sft-lora-singlegpu.yaml
@@ -0,0 +1,70 @@
+base_model: openai/gpt-oss-20b
+use_kernels: true
+model_quantization_config: Mxfp4Config
+model_quantization_config_kwargs:
+  dequantize: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+experimental_skip_move_to_device: true  # prevent OOM by not putting model to GPU before sharding
+
+datasets:
+  - path: HuggingFaceH4/Multilingual-Thinking
+    type: chat_template
+    field_thinking: thinking
+    template_thinking_key: thinking
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/gpt-oss-out/
+
+sequence_len: 4096
+sample_packing: true
+
+adapter: lora
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.0  # dropout not supported when using LoRA over expert parameters
+lora_target_linear: true
+
+# TODO: not supported for now, see peft#2710
+#lora_target_parameters:  # target the experts in the last two layers
+#  - "22._checkpoint_wrapped_module.mlp.experts.gate_up_proj"
+#  - "22._checkpoint_wrapped_module.mlp.experts.down_proj"
+#  - "23._checkpoint_wrapped_module.mlp.experts.gate_up_proj"
+#  - "23._checkpoint_wrapped_module.mlp.experts.down_proj"
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+trackio_project_name:
+trackio_run_name:
+trackio_space_id:
+
+gradient_accumulation_steps: 8
+micro_batch_size: 1
+num_epochs: 1
+
+optimizer: adamw_torch_8bit
+lr_scheduler: constant_with_warmup
+learning_rate: 2e-4
+
+bf16: true
+tf32: true
+
+attn_implementation: kernels-community/vllm-flash-attn3  # this is not needed if using flash_attn >= 2.8.3
+
+gradient_checkpointing: true
+activation_offloading: true
+
+logging_steps: 1
+saves_per_epoch: 1
+warmup_ratio: 0.1
+
+special_tokens:
+eot_tokens:
+  - "<|end|>"
diff --git a/examples/gpt-oss/gpt-oss-safeguard-20b-sft-lora-singlegpu.yaml b/examples/gpt-oss/gpt-oss-safeguard-20b-sft-lora-singlegpu.yaml
new file mode 100644
index 0000000000..7ba5f29b59
--- /dev/null
+++ b/examples/gpt-oss/gpt-oss-safeguard-20b-sft-lora-singlegpu.yaml
@@ -0,0 +1,70 @@
+base_model: openai/gpt-oss-safeguard-20b
+use_kernels: true
+model_quantization_config: Mxfp4Config
+model_quantization_config_kwargs:
+  dequantize: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+experimental_skip_move_to_device: true  # prevent OOM by not putting model to GPU before sharding
+
+datasets:
+  - path: HuggingFaceH4/Multilingual-Thinking
+    type: chat_template
+    field_thinking: thinking
+    template_thinking_key: thinking
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/gpt-oss-safeguard-out/
+
+sequence_len: 4096
+sample_packing: true
+
+adapter: lora
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.0  # dropout not supported when using LoRA over expert parameters
+lora_target_linear: true
+
+# TODO: not supported for now, see peft#2710
+#lora_target_parameters:  # target the experts in the last two layers
+#  - "22._checkpoint_wrapped_module.mlp.experts.gate_up_proj"
+#  - "22._checkpoint_wrapped_module.mlp.experts.down_proj"
+#  - "23._checkpoint_wrapped_module.mlp.experts.gate_up_proj"
+#  - "23._checkpoint_wrapped_module.mlp.experts.down_proj"
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+trackio_project_name:
+trackio_run_name:
+trackio_space_id:
+
+gradient_accumulation_steps: 8
+micro_batch_size: 1
+num_epochs: 1
+
+optimizer: adamw_torch_8bit
+lr_scheduler: constant_with_warmup
+learning_rate: 2e-4
+
+bf16: true
+tf32: true
+
+attn_implementation: kernels-community/vllm-flash-attn3  # this is not needed if using flash_attn >= 2.8.3
+
+gradient_checkpointing: true
+activation_offloading: true
+
+logging_steps: 1
+saves_per_epoch: 1
+warmup_ratio: 0.1
+
+special_tokens:
+eot_tokens:
+  - "<|end|>"
diff --git a/examples/granite4/README.md b/examples/granite4/README.md
new file mode 100644
index 0000000000..85b6621a0b
--- /dev/null
+++ b/examples/granite4/README.md
@@ -0,0 +1,64 @@
+# Finetune IBM's Granite 4.0 with Axolotl
+
+[Granite 4.0](https://huggingface.co/collections/ibm-granite/granite-40-language-models) are a family of open source models trained by IBM Research.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as Granite4 is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).
+
+    Here is an example of how to install from main for pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+git clone https://github.com/axolotl-ai-cloud/axolotl.git
+cd axolotl
+
+uv pip install --no-build-isolation -e '.'
+
+# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
+python scripts/cutcrossentropy_install.py | sh
+```
+
+2. Run the finetuning example:
+
+```bash
+axolotl train examples/granite4/granite-4.0-tiny-fft.yaml
+```
+
+This config uses about 40.8GiB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+### Limitation
+
+Adapter finetuning does not work at the moment. It would error with
+
+```bash
+RuntimeError: mat1 and mat2 shapes cannot be multiplied (4096x3072 and 1x1179648)
+```
+
+In addition, if adapter training works, `lora_target_linear: true` will not work due to:
+```bash
+ValueError: Target module GraniteMoeHybridParallelExperts() is not supported.
+```
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Related Resources
+
+- [Granite Docs](https://www.ibm.com/granite/docs/models/granite)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/granite4/granite-4.0-tiny-fft.yaml b/examples/granite4/granite-4.0-tiny-fft.yaml
new file mode 100644
index 0000000000..fd7d2a3127
--- /dev/null
+++ b/examples/granite4/granite-4.0-tiny-fft.yaml
@@ -0,0 +1,45 @@
+base_model: ibm-granite/granite-4.0-tiny-preview
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/model-out
+
+sequence_len: 2048
+sample_packing: true
+
+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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/hunyuan/README.md b/examples/hunyuan/README.md
new file mode 100644
index 0000000000..d17752cc48
--- /dev/null
+++ b/examples/hunyuan/README.md
@@ -0,0 +1,84 @@
+# Finetune HunYuan with Axolotl
+
+Tencent released a family of opensource models called HunYuan with varying parameter scales of 0.5B, 1.8B, 4B, and 7B scale for both Pre-trained and Instruct variants. The models can be found at [HuggingFace](https://huggingface.co/collections/tencent/hunyuan-dense-model-6890632cda26b19119c9c5e7). This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as HunYuan is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).
+
+    Here is an example of how to install from main for pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+git clone https://github.com/axolotl-ai-cloud/axolotl.git
+cd axolotl
+
+uv pip install --no-build-isolation -e '.'
+
+# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
+python scripts/cutcrossentropy_install.py | sh
+```
+
+2. Run the finetuning example:
+
+```bash
+axolotl train examples/hunyuan/hunyuan-v1-dense-qlora.yaml
+```
+
+This config uses about 4.7 GB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### Dataset
+
+HunYuan Instruct models can choose to enter a slow think or fast think pattern. For best performance on fine-tuning their Instruct models, your dataset should be adjusted to match their pattern.
+
+```python
+# fast think pattern
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "/no_think What color is the sun?" },
+    {"role": "assistant", "content": "<think>\n\n</think>\n<answer>\nThe sun is yellow.\n</answer>"}
+]
+
+# slow think pattern
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "/no_think What color is the sun?" },
+    {"role": "assistant", "content": "<think>\nThe user is asking about the color of the sun. I need to ...\n</think>\n<answer>\nThe sun is yellow.\n</answer>"}
+]
+```
+
+### TIPS
+
+- For inference, the official Tencent team recommends
+
+```json
+
+{
+  "do_sample": true,
+  "top_k": 20,
+  "top_p": 0.8,
+  "repetition_penalty": 1.05,
+  "temperature": 0.7
+}
+
+```
+
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Related Resources
+
+- [Tencent HunYuan Blog](https://hunyuan.tencent.com/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/hunyuan/hunyuan-v1-dense-qlora.yaml b/examples/hunyuan/hunyuan-v1-dense-qlora.yaml
new file mode 100644
index 0000000000..1ae6b000d5
--- /dev/null
+++ b/examples/hunyuan/hunyuan-v1-dense-qlora.yaml
@@ -0,0 +1,64 @@
+base_model: tencent/Hunyuan-0.5B-Instruct
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/internvl3_5/README.md b/examples/internvl3_5/README.md
new file mode 100644
index 0000000000..7424385bb3
--- /dev/null
+++ b/examples/internvl3_5/README.md
@@ -0,0 +1,43 @@
+# Finetune OpenGV's InternVL with Axolotl
+
+[InternVL 3.5](https://huggingface.co/OpenGVLab/InternVL3_5-8B-HF) is a family of powerful vision-language models supporting dynamic resolution and multi-image understanding by OpenGV. It features a ViT-style vision encoder and strong language model backbone for tasks like visual question answering, OCR, and scene text understanding.
+
+This guide shows how to fine-tune it with Axolotl.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install `timm` for vision model support:
+
+    ```bash
+    uv pip install timm==1.0.19
+    ```
+
+3. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+4. Run the finetuning example:
+
+    ```bash
+    axolotl train examples/internvl3_5/internvl3_5-8b-qlora.yml
+    ```
+
+This config uses about 8.21 GiB VRAM. Let us know how it goes. Happy finetuning! 🚀
+
+### Tips
+
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the multi-modal format as seen [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [InternVL Paper](https://huggingface.co/papers/2508.18265)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/internvl3_5/internvl3_5-8b-qlora.yml b/examples/internvl3_5/internvl3_5-8b-qlora.yml
new file mode 100644
index 0000000000..2d924c6f18
--- /dev/null
+++ b/examples/internvl3_5/internvl3_5-8b-qlora.yml
@@ -0,0 +1,60 @@
+base_model: OpenGVLab/InternVL3_5-8B-HF
+processor_type: AutoProcessor
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+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
+
+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
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+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
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/jamba/README.md b/examples/jamba/README.md
index 54f5d1da9c..4c9dc85a06 100644
--- a/examples/jamba/README.md
+++ b/examples/jamba/README.md
@@ -6,5 +6,5 @@
 - ✅ qlora w/ deepspeed Zero-3 needs at least 2x GPUs and 67GiB VRAM (wtf?)
 - ✅ qlora single-gpu, ~51GiB VRAM
 - ✅ multipack
-- ❓ FSDP
+- ✅ FSDP
 - ❓ 8-bit LoRA
diff --git a/examples/jamba/qlora.yaml b/examples/jamba/qlora.yaml
index 41a3854fe1..f625fb6f5a 100644
--- a/examples/jamba/qlora.yaml
+++ b/examples/jamba/qlora.yaml
@@ -1,16 +1,18 @@
 base_model: ai21labs/Jamba-v0.1
+# 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: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 4096
 sample_packing: false
@@ -37,26 +39,20 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch:
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
 special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/jamba/qlora_deepspeed.yaml b/examples/jamba/qlora_deepspeed.yaml
index ef04fb53fe..8ec74f9052 100644
--- a/examples/jamba/qlora_deepspeed.yaml
+++ b/examples/jamba/qlora_deepspeed.yaml
@@ -1,16 +1,17 @@
 base_model: ai21labs/Jamba-v0.1
+# 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: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 4096
 sample_packing: false
@@ -37,26 +38,22 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch:
 saves_per_epoch: 1
-debug:
+
 deepspeed: deepspeed_configs/zero2.json
 weight_decay: 0.0
 special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/jamba/qlora_fsdp_large.yaml b/examples/jamba/qlora_fsdp_large.yaml
new file mode 100644
index 0000000000..76cc0ef186
--- /dev/null
+++ b/examples/jamba/qlora_fsdp_large.yaml
@@ -0,0 +1,67 @@
+base_model: ai21labs/AI21-Jamba-1.5-Large
+# optionally might have model_type or tokenizer_type
+tokenizer_type: AutoTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+load_in_4bit: true
+use_tensorboard: true
+chat_template: jamba
+datasets:
+  - path: cgato/SlimOrcaDedupCleaned
+    type: chat_template
+    drop_system_message: true
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: jamba-large-fsdp-qlora-ft
+adapter: qlora
+sequence_len: 2048
+sample_packing: true
+
+
+lora_r: 16
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: [down_proj,gate_proj,in_proj,k_proj,o_proj,out_proj,q_proj,up_proj,v_proj,x_proj]
+lora_target_linear: false
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 2
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 0.00001
+
+bf16: true
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: false
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: JambaAttentionDecoderLayer,JambaMambaDecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/kimi-linear/README.md b/examples/kimi-linear/README.md
new file mode 100644
index 0000000000..250f8bea1c
--- /dev/null
+++ b/examples/kimi-linear/README.md
@@ -0,0 +1,47 @@
+# Finetune MoonshotAI's Kimi Linear with Axolotl
+
+[Kimi Linear](https://huggingface.co/collections/moonshotai/kimi-linear-a3b) is a MoE model (48B total, 3B active) by MoonshotAI using a hybrid linear attention architecture to achieve a 1M token context length. It uses Kimi Delta Attention (KDA), a refined version of Gated DeltaNet that reduces KV cache size by up to 75% and boosts decoding throughput by up to 6x for long contexts.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+**Note:** Axolotl uses experimental training code for Kimi Linear as their original modeling code is inference-only.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install CCE via [docs](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy)
+
+3. Run the finetuning example:
+
+    ```bash
+    axolotl train examples/kimi-linear/kimi-48b-lora.yaml
+    ```
+
+This config uses about 98.7GiB VRAM.
+
+Let us know how it goes. Happy finetuning!
+
+### TIPS
+
+- Kimi Linear requires `trust_remote_code: true`.
+- You can run a full finetuning by removing the `adapter: lora` and `load_in_8bit: true`.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html)
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template)
+
+## Optimization Guides
+
+See 👉 [docs](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Limitations
+
+This is not yet compatible with MoE kernels from transformers v5.
+
+## Related Resources
+
+- [Kimi Linear Paper](https://huggingface.co/papers/2510.26692)
+- [Kimi Linear GitHub](https://github.com/MoonshotAI/Kimi-Linear)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/kimi-linear/kimi-48b-lora.yaml b/examples/kimi-linear/kimi-48b-lora.yaml
new file mode 100644
index 0000000000..befa298916
--- /dev/null
+++ b/examples/kimi-linear/kimi-48b-lora.yaml
@@ -0,0 +1,81 @@
+base_model: moonshotai/Kimi-Linear-48B-A3B-Instruct
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+trust_remote_code: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: true
+load_in_4bit: false
+strict: false
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+    split: train
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.2
+output_dir: ./outputs/lora-out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+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
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+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
+attn_implementation: flash_attention_2
+
+loss_watchdog_threshold: 5.0
+loss_watchdog_patience: 3
+
+warmup_ratio: 0.1
+evals_per_epoch: 2
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:
+special_tokens:
diff --git a/examples/llama-2/fft_optimized.yml b/examples/llama-2/fft_optimized.yml
index 74edc95e6b..eb46007b98 100644
--- a/examples/llama-2/fft_optimized.yml
+++ b/examples/llama-2/fft_optimized.yml
@@ -1,21 +1,20 @@
 base_model: NousResearch/Llama-2-7b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter:
 lora_model_dir:
@@ -23,7 +22,6 @@ lora_r:
 lora_alpha:
 lora_dropout:
 lora_target_linear:
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -38,31 +36,22 @@ 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
+attn_implementation: flash_attention_2
 flash_attn_cross_entropy: false
-flash_attn_rms_norm: true
-flash_attn_fuse_qkv: false
 flash_attn_fuse_mlp: true
 
-warmup_steps: 100
+warmup_ratio: 0.1
 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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-2/gptq-lora.yml b/examples/llama-2/gptq-lora.yml
index 68ca9ed31c..c4073b80ae 100644
--- a/examples/llama-2/gptq-lora.yml
+++ b/examples/llama-2/gptq-lora.yml
@@ -1,13 +1,15 @@
 base_model: TheBloke/Llama-2-7B-GPTQ
-gptq: true
-gptq_disable_exllama: true
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: LlamaTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+gptq: true
+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
 datasets:
@@ -28,16 +30,15 @@ lora_target_modules:
   - q_proj
   - v_proj
 lora_target_linear:
-lora_fan_in_fan_out:
 wandb_project:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./model-out
+output_dir: ./outputs/model-out
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 4
-optimizer: adamw_torch
+optimizer: adamw_torch_fused
 adam_beta2: 0.95
 adam_eps: 0.00001
 max_grad_norm: 1.0
@@ -45,28 +46,21 @@ 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
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
 special_tokens:
   bos_token: "<s>"
   eos_token: "</s>"
   unk_token: "<unk>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-2/lisa.yml b/examples/llama-2/lisa.yml
index e692c7ac1e..b125f83ca2 100644
--- a/examples/llama-2/lisa.yml
+++ b/examples/llama-2/lisa.yml
@@ -1,21 +1,20 @@
 base_model: NousResearch/Llama-2-7b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./lisa-out
+output_dir: ./outputs/lisa-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter:
 lora_model_dir:
@@ -23,7 +22,6 @@ lora_r:
 lora_alpha:
 lora_dropout:
 lora_target_linear:
-lora_fan_in_fan_out:
 
 lisa_n_layers: 4
 lisa_step_interval: 20
@@ -42,34 +40,23 @@ 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
+attn_implementation: flash_attention_2
 flash_attn_cross_entropy: false
-flash_attn_rms_norm: true
-flash_attn_fuse_qkv: false
 flash_attn_fuse_mlp: true
 
-warmup_steps: 100
+warmup_ratio: 0.1
 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>"
   unk_token: "<unk>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-2/loftq.yml b/examples/llama-2/loftq.yml
index 4529a912dc..f1562ec295 100644
--- a/examples/llama-2/loftq.yml
+++ b/examples/llama-2/loftq.yml
@@ -1,21 +1,20 @@
 base_model: NousResearch/Llama-2-7b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter: lora
 lora_model_dir:
@@ -23,7 +22,6 @@ 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
@@ -41,29 +39,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
-s2_attention:
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-2/lora.yml b/examples/llama-2/lora.yml
index a7793dce4c..8c2242b714 100644
--- a/examples/llama-2/lora.yml
+++ b/examples/llama-2/lora.yml
@@ -1,21 +1,23 @@
 base_model: NousResearch/Llama-2-7b-hf
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
+# 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
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter: lora
 lora_model_dir:
@@ -23,7 +25,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -38,29 +39,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
-s2_attention:
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-2/qlora-fsdp.yml b/examples/llama-2/qlora-fsdp.yml
index 93b3b2a60a..102eb7af7d 100644
--- a/examples/llama-2/qlora-fsdp.yml
+++ b/examples/llama-2/qlora-fsdp.yml
@@ -1,31 +1,31 @@
 base_model: NousResearch/Llama-2-7b-hf
+# optionally might have model_type or tokenizer_type
 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: true
-strict: false
 
 datasets:
   - path: yahma/alpaca-cleaned
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 512
 sample_packing: false
-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:
@@ -36,32 +36,23 @@ wandb_log_model:
 gradient_accumulation_steps: 4
 micro_batch_size: 4
 num_epochs: 4
-optimizer: adamw_torch
+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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
-eval_table_size:
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
 fsdp:
   - full_shard
@@ -75,4 +66,7 @@ fsdp_config:
   fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
   fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
   fsdp_state_dict_type: FULL_STATE_DICT
+  # fsdp_cpu_offload_pin_memory: false  # uncomment to enable swap memory usage when RAM is insufficient
 special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-2/qlora.yml b/examples/llama-2/qlora.yml
index 834dbfb33a..87e710792a 100644
--- a/examples/llama-2/qlora.yml
+++ b/examples/llama-2/qlora.yml
@@ -1,31 +1,31 @@
 base_model: NousResearch/Llama-2-7b-hf
+# optionally might have model_type or tokenizer_type
 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: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-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:
@@ -40,27 +40,18 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
-eval_table_size:
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-2/relora.yml b/examples/llama-2/relora.yml
index 9fd19953c6..8e3df58bf4 100644
--- a/examples/llama-2/relora.yml
+++ b/examples/llama-2/relora.yml
@@ -5,32 +5,32 @@ tokenizer_type: LlamaTokenizer
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: teknium/GPT4-LLM-Cleaned
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./relora-out
+output_dir: ./outputs/relora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-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
+relora: true
+relora_prune_ratio: 0.9
 relora_cpu_offload: false
+jagged_restart_steps: 150
+jagged_restart_warmup_steps: 10
+jagged_restart_anneal_steps: false
 
 wandb_project:
 wandb_entity:
@@ -45,29 +45,21 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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>"
   unk_token: "<unk>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3-vision/lora-11b.yaml b/examples/llama-3-vision/lora-11b.yaml
new file mode 100644
index 0000000000..4e5eb4c4ee
--- /dev/null
+++ b/examples/llama-3-vision/lora-11b.yaml
@@ -0,0 +1,60 @@
+base_model: alpindale/Llama-3.2-11B-Vision-Instruct
+# optionally might have model_type or tokenizer_type or processor_type
+processor_type: AutoProcessor
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+
+# 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: llama3_2_vision
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+dataset_prepared_path:
+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: 'model.language_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  # use for text-only mode
+attn_implementation: sdpa
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/3b-fp8-fsdp2.yaml b/examples/llama-3/3b-fp8-fsdp2.yaml
new file mode 100644
index 0000000000..cfc15870fd
--- /dev/null
+++ b/examples/llama-3/3b-fp8-fsdp2.yaml
@@ -0,0 +1,75 @@
+base_model: meta-llama/Llama-3.2-3B
+# 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
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+
+datasets:
+  - path: yahma/alpaca-cleaned
+    type: alpaca
+
+output_dir: ./outputs/fp8_out/
+
+sample_packing: true
+pad_to_sequence_len: true
+sequence_len: 512
+
+attn_implementation: flex_attention
+flex_attn_compile_kwargs:
+  dynamic: false
+  mode: max-autotune-no-cudagraphs
+torch_compile: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+num_epochs: 1
+optimizer: adamw_torch_fused
+
+cosine_constant_lr_ratio: 0
+cosine_min_lr_ratio: 1.0
+learning_rate: 2e-5
+save_only_model: true
+
+fp8: true
+fp8_enable_fsdp_float8_all_gather: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_steps: 10
+weight_decay: 0.0
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: false
+
+special_tokens:
+  pad_token: <|end_of_text|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/3b-qat-fsdp2.yaml b/examples/llama-3/3b-qat-fsdp2.yaml
new file mode 100644
index 0000000000..99c975351e
--- /dev/null
+++ b/examples/llama-3/3b-qat-fsdp2.yaml
@@ -0,0 +1,79 @@
+base_model: meta-llama/Llama-3.2-3B
+# 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
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+
+
+datasets:
+  - path: yahma/alpaca-cleaned
+    type: alpaca
+    split: train[:95%]
+
+output_dir: ./outputs/qat_out/
+dataset_prepared_path: ./outputs/qat_out/dataset_prepared
+
+sample_packing: false
+sequence_len: 8192
+attn_implementation: flash_attention_2
+
+qat:
+  activation_dtype: int8
+  weight_dtype: int4
+  group_size: 32
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+num_epochs: 1
+optimizer: adamw_torch_fused
+
+cosine_constant_lr_ratio: 0
+cosine_min_lr_ratio: 1.0
+learning_rate: 2e-5
+save_only_model: true
+bf16: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp:
+  - full_shard
+  - auto_wrap
+
+fsdp_config:
+  fsdp_version: 2
+  fsdp_offload_params: false
+  fsdp_cpu_ram_efficient_loading: false
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_reshard_after_forward: true
+  fsdp_activation_checkpointing: true
+
+special_tokens:
+  pad_token: <|finetune_right_pad_id|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/3b-qat-mxfp4.yaml b/examples/llama-3/3b-qat-mxfp4.yaml
new file mode 100644
index 0000000000..4e9f646857
--- /dev/null
+++ b/examples/llama-3/3b-qat-mxfp4.yaml
@@ -0,0 +1,65 @@
+base_model: meta-llama/Llama-3.2-3B
+# 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
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+
+datasets:
+  - path: yahma/alpaca-cleaned
+    type: alpaca
+    split: train[:95%]
+
+output_dir: ./outputs/qat_out/
+dataset_prepared_path: ./outputs/dataset_prepared
+
+sequence_len: 2048
+attn_implementation: flash_attention_2
+
+qat:
+  activation_dtype: mxfp4
+  weight_dtype: mxfp4
+  group_size: 32
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_checkpointing: true
+activation_offloading: true
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_8bit
+
+cosine_constant_lr_ratio: 0
+cosine_min_lr_ratio: 1.0
+learning_rate: 2e-5
+save_only_model: true
+bf16: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+special_tokens:
+  pad_token: <|finetune_right_pad_id|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/3b-qat-nvfp4.yaml b/examples/llama-3/3b-qat-nvfp4.yaml
new file mode 100644
index 0000000000..77cf2b19bc
--- /dev/null
+++ b/examples/llama-3/3b-qat-nvfp4.yaml
@@ -0,0 +1,64 @@
+base_model: meta-llama/Llama-3.2-3B
+# 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
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+
+datasets:
+  - path: yahma/alpaca-cleaned
+    type: alpaca
+    split: train[:95%]
+
+output_dir: ./outputs/qat_out/
+dataset_prepared_path: ./outputs/dataset_prepared
+
+sequence_len: 8192
+attn_implementation: flash_attention_2
+
+qat:
+  activation_dtype: nvfp4
+  weight_dtype: nvfp4
+  group_size: 16 # only group_size of 16 is supported with nvfp4
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_checkpointing: true
+gradient_accumulation_steps: 1
+micro_batch_size: 64
+num_epochs: 1
+optimizer: adamw_torch_fused
+
+cosine_constant_lr_ratio: 0
+cosine_min_lr_ratio: 1.0
+learning_rate: 2e-5
+save_only_model: true
+bf16: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+special_tokens:
+  pad_token: <|finetune_right_pad_id|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/diffusion/pretrain-1b.yaml b/examples/llama-3/diffusion/pretrain-1b.yaml
new file mode 100644
index 0000000000..1b488db7a1
--- /dev/null
+++ b/examples/llama-3/diffusion/pretrain-1b.yaml
@@ -0,0 +1,56 @@
+base_model: meta-llama/Llama-3.2-1B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+pretraining_dataset:
+  - path: wikitext
+    name: wikitext-103-raw-v1
+    type: completion
+    field: text
+
+plugins:
+  - axolotl.integrations.diffusion.DiffusionPlugin
+
+diffusion:
+  noise_schedule: cosine
+  min_mask_ratio: 0.15
+  max_mask_ratio: 0.85
+  num_diffusion_steps: 128
+  eps: 5e-4
+  importance_weighting: true
+  mask_token_id: 128002
+  generate_samples: true
+  generation_interval: 250
+
+output_dir: ./outputs/model-out
+
+sequence_len: 512
+sample_packing: true
+
+gradient_accumulation_steps: 8
+micro_batch_size: 4
+max_steps: 10000
+warmup_ratio: 0.1
+
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 3e-4
+attn_implementation: sdpa
+
+bf16: auto
+tf32: true
+
+logging_steps: 1
+save_strategy: steps
+save_steps: 1000
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/diffusion/sft-1b.yaml b/examples/llama-3/diffusion/sft-1b.yaml
new file mode 100644
index 0000000000..b6de76af3a
--- /dev/null
+++ b/examples/llama-3/diffusion/sft-1b.yaml
@@ -0,0 +1,59 @@
+base_model: meta-llama/Llama-3.2-1B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+datasets:
+  - path: teknium/GPT4-LLM-Cleaned
+    type: alpaca
+val_set_size: 0.05
+
+plugins:
+  - axolotl.integrations.diffusion.DiffusionPlugin
+
+diffusion:
+  noise_schedule: cosine
+  min_mask_ratio: 0.1
+  max_mask_ratio: 0.9
+  num_diffusion_steps: 128
+  eps: 1e-3
+  importance_weighting: true
+  mask_token_id: 128002
+  generate_samples: true
+  generation_interval: 250
+
+output_dir: ./outputs/model-out
+
+sequence_len: 512
+sample_packing: true
+eval_sample_packing: true
+
+gradient_accumulation_steps: 4
+micro_batch_size: 4
+num_epochs: 1
+warmup_steps: 0.1
+
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 1e-5
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+attn_implementation: sdpa
+
+logging_steps: 1
+save_strategy: best
+eval_strategy: epoch
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/fft-8b-liger-fsdp.yaml b/examples/llama-3/fft-8b-liger-fsdp.yaml
new file mode 100644
index 0000000000..b96bc920ed
--- /dev/null
+++ b/examples/llama-3/fft-8b-liger-fsdp.yaml
@@ -0,0 +1,76 @@
+base_model: NousResearch/Meta-Llama-3.1-8B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_fused_linear_cross_entropy: true
+
+
+chat_template: llama3
+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.02
+output_dir: ./outputs/out
+
+sequence_len: 4096
+sample_packing: true
+
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 2
+saves_per_epoch: 1
+weight_decay: 0.0
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_backward_prefetch: BACKWARD_PRE
+special_tokens:
+  pad_token: <|finetune_right_pad_id|>
+  eos_token: <|eot_id|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/fft-8b.yaml b/examples/llama-3/fft-8b.yaml
index 8c9ba90bfe..3e2809196b 100644
--- a/examples/llama-3/fft-8b.yaml
+++ b/examples/llama-3/fft-8b.yaml
@@ -1,21 +1,17 @@
-base_model: meta-llama/Meta-Llama-3-8B
-model_type: LlamaForCausalLM
-tokenizer_type: AutoTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+base_model: NousResearch/Meta-Llama-3.1-8B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 8192
 sample_packing: true
-pad_to_sequence_len: true
+
 
 wandb_project:
 wandb_entity:
@@ -30,29 +26,21 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 100
+warmup_ratio: 0.1
 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|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/instruct-dpo-lora-8b.yml b/examples/llama-3/instruct-dpo-lora-8b.yml
new file mode 100644
index 0000000000..b49ace2edd
--- /dev/null
+++ b/examples/llama-3/instruct-dpo-lora-8b.yml
@@ -0,0 +1,75 @@
+base_model: meta-llama/Meta-Llama-3-8B-Instruct
+# 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
+
+special_tokens:
+  pad_token: <|finetune_right_pad_id|>
+  eos_token: <|eot_id|>
+
+load_in_8bit: true
+load_in_4bit: false
+
+chat_template: llama3
+rl: dpo
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_dpo_test
+    type: chat_template.default
+    field_messages: conversation
+    field_chosen: chosen
+    field_rejected: rejected
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      system:
+        - system
+      user:
+        - user
+      assistant:
+        - assistant
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/lora-out
+
+sequence_len: 4096
+sample_packing: false
+
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/instruct-lora-8b.yml b/examples/llama-3/instruct-lora-8b.yml
new file mode 100644
index 0000000000..1c61ce9e42
--- /dev/null
+++ b/examples/llama-3/instruct-lora-8b.yml
@@ -0,0 +1,59 @@
+base_model: NousResearch/Meta-Llama-3-8B-Instruct
+# 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
+
+chat_template: llama3
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/lora-out
+
+sequence_len: 4096
+sample_packing: false
+
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+   pad_token: <|end_of_text|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/lora-1b-deduplicate-dpo.yml b/examples/llama-3/lora-1b-deduplicate-dpo.yml
new file mode 100644
index 0000000000..2be72c4d01
--- /dev/null
+++ b/examples/llama-3/lora-1b-deduplicate-dpo.yml
@@ -0,0 +1,87 @@
+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
+
+chat_template: llama3
+rl: dpo
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_dpo_test
+    type: chat_template.default
+    field_messages: conversation
+    field_chosen: chosen
+    field_rejected: rejected
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      system:
+        - system
+      user:
+        - user
+      assistant:
+        - assistant
+  - path: fozziethebeat/alpaca_messages_2k_dpo_test
+    type: chat_template.default
+    field_messages: conversation
+    field_chosen: chosen
+    field_rejected: rejected
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      system:
+        - system
+      user:
+        - user
+      assistant:
+        - assistant
+
+dataset_exact_deduplication: true
+dataset_prepared_path:
+val_set_size: 0
+output_dir: ./outputs/lora-out
+
+sequence_len: 4096
+sample_packing: false
+
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/lora-1b-deduplicate-sft.yml b/examples/llama-3/lora-1b-deduplicate-sft.yml
new file mode 100644
index 0000000000..ad21cb266a
--- /dev/null
+++ b/examples/llama-3/lora-1b-deduplicate-sft.yml
@@ -0,0 +1,65 @@
+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
+
+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
+
+dataset_exact_deduplication: true
+
+sequence_len: 4096
+sample_packing: true
+eval_sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+   pad_token: <|end_of_text|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/lora-1b-kernels.yml b/examples/llama-3/lora-1b-kernels.yml
new file mode 100644
index 0000000000..b0914f87ae
--- /dev/null
+++ b/examples/llama-3/lora-1b-kernels.yml
@@ -0,0 +1,69 @@
+base_model: NousResearch/Llama-3.2-1B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+datasets:
+  - path: teknium/GPT4-LLM-Cleaned
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+
+lora_r: 16
+lora_alpha: 32
+# Currently, we don't support dropout with our custom Triton kernels
+# lora_dropout: 0.05
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+# These options enable our custom Triton kernels / autograd
+# functions for MLP and attention calculations
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+loss_watchdog_threshold: 5.0
+loss_watchdog_patience: 3
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/lora-1b-ray.yml b/examples/llama-3/lora-1b-ray.yml
new file mode 100644
index 0000000000..a3aa1cf5e0
--- /dev/null
+++ b/examples/llama-3/lora-1b-ray.yml
@@ -0,0 +1,68 @@
+base_model: NousResearch/Llama-3.2-1B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+datasets:
+  - path: teknium/GPT4-LLM-Cleaned
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+loss_watchdog_threshold: 5.0
+loss_watchdog_patience: 3
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+
+deepspeed: deepspeed_configs/zero3.json
+weight_decay: 0.0
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+use_ray: true
+ray_num_workers: 4
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/lora-1b-sample-packing-sequentially.yml b/examples/llama-3/lora-1b-sample-packing-sequentially.yml
new file mode 100644
index 0000000000..f6c24bc74f
--- /dev/null
+++ b/examples/llama-3/lora-1b-sample-packing-sequentially.yml
@@ -0,0 +1,67 @@
+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
+
+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
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+  pad_token: <|end_of_text|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/lora-1b.yml b/examples/llama-3/lora-1b.yml
new file mode 100644
index 0000000000..d01c618bc8
--- /dev/null
+++ b/examples/llama-3/lora-1b.yml
@@ -0,0 +1,64 @@
+base_model: NousResearch/Llama-3.2-1B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+datasets:
+  - path: teknium/GPT4-LLM-Cleaned
+    type: alpaca
+
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+loss_watchdog_threshold: 5.0
+loss_watchdog_patience: 3
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/lora-8b.yml b/examples/llama-3/lora-8b.yml
index d60f8a3035..90084ec95d 100644
--- a/examples/llama-3/lora-8b.yml
+++ b/examples/llama-3/lora-8b.yml
@@ -1,21 +1,24 @@
-base_model: meta-llama/Meta-Llama-3-8B
+base_model: NousResearch/Meta-Llama-3-8B
+# 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
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+eval_sample_packing: false
+
 
 adapter: lora
 lora_model_dir:
@@ -23,7 +26,9 @@ 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
 
 wandb_project:
 wandb_entity:
@@ -38,30 +43,19 @@ 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:
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+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:
    pad_token: <|end_of_text|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/opentelemetry-qlora.yml b/examples/llama-3/opentelemetry-qlora.yml
new file mode 100644
index 0000000000..0c9995dae4
--- /dev/null
+++ b/examples/llama-3/opentelemetry-qlora.yml
@@ -0,0 +1,49 @@
+base_model: NousResearch/Llama-3.2-1B
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+
+load_in_4bit: true
+
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+
+output_dir: ./outputs/opentelemetry-example
+
+adapter: qlora
+sequence_len: 512
+sample_packing: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+# OpenTelemetry Configuration
+use_otel_metrics: true
+otel_metrics_host: "localhost"
+otel_metrics_port: 8000
+
+# Disable WandB
+use_wandb: false
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+optimizer: paged_adamw_32bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+logging_steps: 1
+
+warmup_ratio: 0.1
+evals_per_epoch: 2
+saves_per_epoch: 1
+weight_decay: 0.0
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
diff --git a/examples/llama-3/qlora-1b-gdpo.yaml b/examples/llama-3/qlora-1b-gdpo.yaml
new file mode 100644
index 0000000000..f754a68875
--- /dev/null
+++ b/examples/llama-3/qlora-1b-gdpo.yaml
@@ -0,0 +1,68 @@
+base_model: meta-llama/Llama-3.2-1B-Instruct
+
+chat_template: llama3
+
+rl: gdpo
+
+trl:
+  beta: 0.001
+  max_completion_length: 128
+  num_generations: 2
+  temperature: 0.7
+  top_p: 0.95
+
+  use_vllm: false
+
+
+  multi_objective_aggregation: normalize_then_sum
+
+  reward_funcs:
+    - rwd.format_reward
+    - rwd.correctness_reward
+  reward_weights: [1.0, 2.0]
+
+  log_completions: true
+  num_completions_to_print: 3
+  scale_rewards: true
+
+datasets:
+  - path: openai/gsm8k
+    name: main
+    split: train[:1000]
+    type: rwd.gsm8k_transform
+
+val_set_size: 0.0
+output_dir: ./outputs/llama3-gdpo-out
+
+sequence_len: 512
+sample_packing: false
+pad_to_sequence_len: false
+
+gradient_accumulation_steps: 8
+micro_batch_size: 1
+num_epochs: 1
+max_steps: 100
+
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 5e-5
+weight_decay: 0.01
+warmup_steps: 10
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+
+attn_implementation: flash_attention_2
+logging_steps: 1
+save_steps: 50
+save_safetensors: true
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+
+seed: 42
diff --git a/examples/llama-3/qlora-1b-kto.yaml b/examples/llama-3/qlora-1b-kto.yaml
new file mode 100644
index 0000000000..18c240d972
--- /dev/null
+++ b/examples/llama-3/qlora-1b-kto.yaml
@@ -0,0 +1,65 @@
+base_model: meta-llama/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: true
+
+rl: kto
+rl_beta: 0.5
+kto_desirable_weight: 0.2
+
+datasets:
+  - path: argilla/ultrafeedback-binarized-preferences-cleaned-kto
+    type: llama3.ultra
+    split: train
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/qlora-out
+
+remove_unused_columns: false
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: false  # not supported with kto
+eval_sample_packing: false
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 64
+lora_dropout: 0.05
+lora_target_linear: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/qlora-1b.yml b/examples/llama-3/qlora-1b.yml
new file mode 100644
index 0000000000..d1e5e18ae2
--- /dev/null
+++ b/examples/llama-3/qlora-1b.yml
@@ -0,0 +1,66 @@
+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: true
+
+datasets:
+  - path: teknium/GPT4-LLM-Cleaned
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/qlora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+loss_watchdog_threshold: 5.0
+loss_watchdog_patience: 3
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/qlora-fsdp-405b.yaml b/examples/llama-3/qlora-fsdp-405b.yaml
new file mode 100644
index 0000000000..b801af8458
--- /dev/null
+++ b/examples/llama-3/qlora-fsdp-405b.yaml
@@ -0,0 +1,63 @@
+base_model: hugging-quants/Meta-Llama-3.1-405B-BNB-NF4-BF16
+# optionally might have model_type or tokenizer_type
+tokenizer_type: AutoTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+load_in_4bit: true
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out/qlora-llama3_1-405b
+
+adapter: qlora
+
+sequence_len: 2048
+sample_packing: true
+
+
+lora_r: 16
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 2
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 0.00001
+
+bf16: true
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+special_tokens:
+  pad_token: <|finetune_right_pad_id|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/qlora-fsdp-70b.yaml b/examples/llama-3/qlora-fsdp-70b.yaml
index 8d8785bfd5..5ce774e182 100644
--- a/examples/llama-3/qlora-fsdp-70b.yaml
+++ b/examples/llama-3/qlora-fsdp-70b.yaml
@@ -1,31 +1,31 @@
 base_model: casperhansen/llama-3-70b-fp16
+# optionally might have model_type or tokenizer_type
 model_type: LlamaForCausalLM
 tokenizer_type: AutoTokenizer  # PreTrainedTokenizerFast
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out/qlora-llama3-70b
+output_dir: ./outputs/out/qlora-llama3-70b
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 512
 sample_packing: false
-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:
@@ -36,32 +36,23 @@ wandb_log_model:
 gradient_accumulation_steps: 4
 micro_batch_size: 1
 num_epochs: 4
-optimizer: adamw_torch
+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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
-eval_table_size:
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.0
 fsdp:
   - full_shard
@@ -78,3 +69,5 @@ fsdp_config:
   fsdp_sharding_strategy: FULL_SHARD
 special_tokens:
   pad_token: <|end_of_text|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/qlora.yml b/examples/llama-3/qlora.yml
index 9cedee8eec..fad507cd98 100644
--- a/examples/llama-3/qlora.yml
+++ b/examples/llama-3/qlora.yml
@@ -1,31 +1,31 @@
-base_model: meta-llama/Meta-Llama-3-8B
+base_model: NousResearch/Meta-Llama-3-8B
+# 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
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: aaditya/alpaca_subset_1
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-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:
@@ -40,28 +40,19 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-3/sparse-finetuning.yaml b/examples/llama-3/sparse-finetuning.yaml
new file mode 100644
index 0000000000..0ce4aa03d1
--- /dev/null
+++ b/examples/llama-3/sparse-finetuning.yaml
@@ -0,0 +1,78 @@
+base_model: neuralmagic/Sparse-Llama-3.1-8B-2of4
+
+plugins:
+  - axolotl.integrations.llm_compressor.LLMCompressorPlugin
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.05
+output_dir: ./outputs/out
+
+sequence_len: 4096
+sample_packing: true
+
+eval_sample_packing: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 8
+micro_batch_size: 1
+num_epochs: 1
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+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|>
+
+llmcompressor:
+  recipe:
+    finetuning_stage:
+      finetuning_modifiers:
+        ConstantPruningModifier:
+          targets: [
+            're:.*q_proj.weight',
+            're:.*k_proj.weight',
+            're:.*v_proj.weight',
+            're:.*o_proj.weight',
+            're:.*gate_proj.weight',
+            're:.*up_proj.weight',
+            're:.*down_proj.weight',
+          ]
+          start: 0
+  save_compressed: true
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-4/README.md b/examples/llama-4/README.md
new file mode 100644
index 0000000000..653da155cd
--- /dev/null
+++ b/examples/llama-4/README.md
@@ -0,0 +1,38 @@
+# Llama 4 by Meta AI
+
+## Flash Attention vs Flex Attention
+
+While Flash Attention to support is "enabled" for Llama-4, the upstream implementation is not correct and usage of Flex Attention is recommended.
+
+## Available Examples
+
+### Llama 4 Scout 17Bx16Experts (109B)
+
+Flex Attention
+- [Text Single GPU (H100) QLoRA](./scout-qlora-single-h100-flex.yaml)
+- [Text Multi GPU QLoRA w/ FSDP2](./scout-qlora-flexattn-fsdp2.yaml)
+
+[//]: # (Flash Attention &#40;Do not use&#41;)
+
+[//]: # (- [Multi-Modal/Vision QLoRA w/ FSDP1]&#40;./scout-vision-qlora-fsdp.yaml&#41;)
+
+[//]: # (- [Text Single GPU &#40;H100&#41; QLoRA]&#40;./scout-qlora-single-h100.yaml&#41;)
+
+[//]: # (- [Text Multi GPU QLoRA w/ FSDP1]&#40;./scout-qlora-fsdp1.yaml&#41;)
+
+Our Single H100 implementation for Llama 4 Scout uses only 64.5GB VRAM for post-training with 4k context length @ 519 tokens/second. [WandB logs here](https://wandb.ai/axolotl-ai/llama4-flexattn-qlora/runs/wpie7dkj)
+Multi-GPU (4xH100) for Llama 4 Scout uses 62.8GB VRAM/GPU @ 4k contenxt length @ 280tps/gpu, [WandB logs here](https://wandb.ai/axolotl-ai/llama4-flexattn-qlora/runs/2lkezdj8)
+
+### Llama 4 Maverick 17Bx128Experts (400B)
+
+Coming Soon
+
+## Delinearized Llama 4 Models
+
+We provide a script to delinearize Llama 4 linearized models into regular HuggingFace Llama 4 models.
+
+```bash
+axolotl delinearize-llama4 --model path/to/model_dir --output path/to/output_dir
+```
+
+Note: This only works with the non-quantized linearized model. If you have an adapter, merge it with the *non-quantized linearized* model before delinearizing.
diff --git a/examples/llama-4/do-no-use-fa2/maverick-qlora-fsdp1.yaml b/examples/llama-4/do-no-use-fa2/maverick-qlora-fsdp1.yaml
new file mode 100644
index 0000000000..2c701a2aa3
--- /dev/null
+++ b/examples/llama-4/do-no-use-fa2/maverick-qlora-fsdp1.yaml
@@ -0,0 +1,90 @@
+base_model: axolotl-quants/Llama-4-Maverick-17B-128E-Linearized-bnb-nf4-bf16
+model_type: Llama4ForConditionalGeneration
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_glu_activation: true
+liger_rms_norm: true
+liger_layer_norm: true
+
+llama4_linearized_experts: true
+load_in_4bit: true
+adapter: qlora
+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
+  - shared_expert.gate_proj
+  - shared_expert.up_proj
+  - shared_expert.down_proj
+  # - experts.gate_projs.[0-9]+$
+  # - experts.up_projs.[0-9]+$
+  # - experts.down_projs.[0-9]+$
+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
+
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 1e-4
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+gradient_checkpointing: offload
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+fsdp:
+  - auto_wrap
+  - full_shard
+fsdp_config:
+  fsdp_transformer_layer_cls_to_wrap: Llama4TextDecoderLayer
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+special_tokens:
+  pad_token: <|finetune_right_pad|>
+  eos_token: <|eot|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-4/do-no-use-fa2/scout-qlora-fsdp1.yaml b/examples/llama-4/do-no-use-fa2/scout-qlora-fsdp1.yaml
new file mode 100644
index 0000000000..8197d16297
--- /dev/null
+++ b/examples/llama-4/do-no-use-fa2/scout-qlora-fsdp1.yaml
@@ -0,0 +1,94 @@
+base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
+model_type: Llama4ForConditionalGeneration
+  # Automatically upload checkpoint and final model to HF
+  # hub_model_id: username/custom_model_name
+
+
+# torch_compile: true
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_glu_activation: true
+liger_rms_norm: true
+liger_layer_norm: true
+
+llama4_linearized_experts: true
+load_in_4bit: true
+adapter: qlora
+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
+  - shared_expert.gate_proj
+  - shared_expert.up_proj
+  - shared_expert.down_proj
+    # - experts.gate_projs.[0-9]+$
+    # - experts.up_projs.[0-9]+$
+    # - experts.down_projs.[0-9]+$
+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
+
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+fsdp:
+  - auto_wrap
+  - full_shard
+fsdp_config:
+  fsdp_transformer_layer_cls_to_wrap: Llama4TextDecoderLayer
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_activation_checkpointing: true
+special_tokens:
+  pad_token: <|finetune_right_pad|>
+  eos_token: <|eot|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-4/do-no-use-fa2/scout-qlora-single-h100.yaml b/examples/llama-4/do-no-use-fa2/scout-qlora-single-h100.yaml
new file mode 100644
index 0000000000..2dcff36cd5
--- /dev/null
+++ b/examples/llama-4/do-no-use-fa2/scout-qlora-single-h100.yaml
@@ -0,0 +1,87 @@
+base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
+model_type: Llama4ForConditionalGeneration
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_glu_activation: true
+liger_rms_norm: true
+liger_layer_norm: true
+
+llama4_linearized_experts: true
+load_in_4bit: true
+adapter: qlora
+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
+  - shared_expert.gate_proj
+  - shared_expert.up_proj
+  - shared_expert.down_proj
+  # - experts.gate_projs.[0-9]+$
+  # - experts.up_projs.[0-9]+$
+  # - experts.down_projs.[0-9]+$
+lora_modules_to_save:
+  # - lm_head
+  # - embed_tokens
+
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+
+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  # up to 8k will work on a single H100
+sample_packing: true
+
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 1e-4
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+gradient_checkpointing: offload
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+  pad_token: <|finetune_right_pad|>
+  eos_token: <|eot|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-4/do-no-use-fa2/scout-vision-qlora-fsdp.yaml b/examples/llama-4/do-no-use-fa2/scout-vision-qlora-fsdp.yaml
new file mode 100644
index 0000000000..de7ae5f50d
--- /dev/null
+++ b/examples/llama-4/do-no-use-fa2/scout-vision-qlora-fsdp.yaml
@@ -0,0 +1,89 @@
+base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
+model_type: Llama4ForConditionalGeneration
+processor_type: Llama4Processor
+  # Automatically upload checkpoint and final model to HF
+  # hub_model_id: username/custom_model_name
+
+
+# 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
+
+sequence_len: 4096
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_glu_activation: true
+liger_rms_norm: true
+liger_layer_norm: true
+
+llama4_linearized_experts: true  # use Axolotl's customized model
+load_in_4bit: true
+adapter: qlora
+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
+  - shared_expert.gate_proj
+  - shared_expert.up_proj
+  - shared_expert.down_proj
+  - vision_adapter.mlp.fc1
+  - vision_adapter.mlp.fc2
+  # - experts.gate_projs.[0-9]+$
+  # - experts.up_projs.[0-9]+$
+  # - experts.down_projs.[0-9]+$
+lora_modules_to_save:
+  - lm_head
+  - embed_tokens
+
+chat_template: llama4
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+fsdp:
+  - auto_wrap
+  - full_shard
+fsdp_config:
+  fsdp_transformer_layer_cls_to_wrap: Llama4TextDecoderLayer
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_activation_checkpointing: true
+special_tokens:
+  pad_token: <|finetune_right_pad|>
+  eos_token: <|eot|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-4/scout-qlora-flexattn-fsdp2.yaml b/examples/llama-4/scout-qlora-flexattn-fsdp2.yaml
new file mode 100644
index 0000000000..c5343fa2e0
--- /dev/null
+++ b/examples/llama-4/scout-qlora-flexattn-fsdp2.yaml
@@ -0,0 +1,88 @@
+base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
+model_type: Llama4ForConditionalGeneration
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_glu_activation: true
+liger_rms_norm: true
+liger_layer_norm: true
+
+llama4_linearized_experts: true
+load_in_4bit: true
+adapter: qlora
+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
+  - shared_expert.gate_proj
+  - shared_expert.up_proj
+  - shared_expert.down_proj
+  # - experts.gate_projs.[0-9]+$
+  # - experts.up_projs.[0-9]+$
+  # - experts.down_projs.[0-9]+$
+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
+
+
+gradient_accumulation_steps: 1
+micro_batch_size: 2
+num_epochs: 3
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 1e-4
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+attn_implementation: flex_attention
+flex_attn_compile_kwargs:
+  dynamic: false
+  mode: max-autotune-no-cudagraphs
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+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 # does not work with load_in_8bit/4bit
+  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|>
+  eos_token: <|eot|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-4/scout-qlora-single-h100-flex.yaml b/examples/llama-4/scout-qlora-single-h100-flex.yaml
new file mode 100644
index 0000000000..00491c3b16
--- /dev/null
+++ b/examples/llama-4/scout-qlora-single-h100-flex.yaml
@@ -0,0 +1,86 @@
+base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
+model_type: Llama4ForConditionalGeneration
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+liger_glu_activation: true
+liger_rms_norm: true
+liger_layer_norm: true
+
+llama4_linearized_experts: true  # needed with custom linearized experts model
+load_in_4bit: true
+adapter: qlora
+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
+  - shared_expert.gate_proj
+  - shared_expert.up_proj
+  - shared_expert.down_proj
+  # - experts.gate_projs.[0-9]+$  # optionally train the moe experts
+  # - experts.up_projs.[0-9]+$
+  # - experts.down_projs.[0-9]+$
+lora_modules_to_save:
+  # - lm_head  # needed if modifying vocabulary
+  # - embed_tokens
+
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+
+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  # up to 8k will work on a single H100
+sample_packing: true
+
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 1e-4
+
+bf16: true
+tf32: true
+
+torch_compile: true
+attn_implementation: flex_attention
+flex_attn_compile_kwargs:
+  dynamic: false
+  mode: max-autotune-no-cudagraphs
+
+gradient_checkpointing: offload
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+
+logging_steps: 1
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+weight_decay: 0.0
+special_tokens:
+  pad_token: <|finetune_right_pad|>
+  eos_token: <|eot|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llama-4/scout-vision-qlora-fsdp2-flex.yaml b/examples/llama-4/scout-vision-qlora-fsdp2-flex.yaml
new file mode 100644
index 0000000000..9b3e089b5a
--- /dev/null
+++ b/examples/llama-4/scout-vision-qlora-fsdp2-flex.yaml
@@ -0,0 +1,90 @@
+base_model: axolotl-quants/Llama-4-Scout-17B-16E-Linearized-bnb-nf4-bf16
+model_type: Llama4ForConditionalGeneration
+processor_type: Llama4Processor
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+# 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
+
+sequence_len: 4096
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_glu_activation: true
+liger_rms_norm: true
+liger_layer_norm: true
+
+llama4_linearized_experts: true  # use Axolotl's customized model
+load_in_4bit: true
+adapter: qlora
+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
+  - shared_expert.gate_proj
+  - shared_expert.up_proj
+  - shared_expert.down_proj
+  - vision_adapter.mlp.fc1
+  - vision_adapter.mlp.fc2
+  # - experts.gate_projs.[0-9]+$
+  # - experts.up_projs.[0-9]+$
+  # - experts.down_projs.[0-9]+$
+lora_modules_to_save:
+  - lm_head
+  - embed_tokens
+
+chat_template: llama4
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 1e-4
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+attn_implementation: flex_attention
+flex_attn_compile_kwargs:
+  dynamic: false
+  mode: max-autotune-no-cudagraphs
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+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|>
+  eos_token: <|eot|>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/llava/lora-7b.yaml b/examples/llava/lora-7b.yaml
new file mode 100644
index 0000000000..56b48fda9f
--- /dev/null
+++ b/examples/llava/lora-7b.yaml
@@ -0,0 +1,55 @@
+base_model: llava-hf/llava-1.5-7b-hf
+processor_type: AutoProcessor
+
+# 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%]
+dataset_prepared_path:
+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: 'model.language_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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/magistral/README.md b/examples/magistral/README.md
new file mode 100644
index 0000000000..5cce33ea77
--- /dev/null
+++ b/examples/magistral/README.md
@@ -0,0 +1,80 @@
+# Finetune Magistral Small with Axolotl
+
+Magistral Small is a 24B parameter opensource model from MistralAI found on HuggingFace at [2506](https://huggingface.co/mistralai/Magistral-Small-2506), [2507](https://huggingface.co/mistralai/Magistral-Small-2507) (see [Thinking](#thinking)), and [2509](https://huggingface.co/mistralai/Magistral-Small-2509) (see [Vision](#vision)). This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+MistralAI has also released a proprietary medium-sized version called Magistral Medium.
+
+Thanks to the team at MistralAI for giving us early access to prepare for these releases.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+uv pip install --no-build-isolation 'axolotl>=0.16.1'
+```
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage
+
+```bash
+python scripts/cutcrossentropy_install.py | sh
+```
+
+3. Run the finetuning example:
+
+```bash
+axolotl train examples/magistral/magistral-small-qlora.yaml
+```
+
+This config uses about 24GB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### Thinking
+
+MistralAI has released their [2507](https://huggingface.co/mistralai/Magistral-Small-2507) model with thinking capabilities, enabling Chain-of-Thought reasoning with explicit thinking steps.
+
+📚 **[See the Thinking fine-tuning guide →](./think/README.md)**
+
+### Vision
+
+MistralAI has released their [2509](https://huggingface.co/mistralai/Magistral-Small-2509) model with vision capabilities.
+
+📚 **[See the Vision fine-tuning guide →](./vision/README.md)**
+
+### Tips
+
+- We recommend adding the same/similar SystemPrompt that the model is tuned for. You can find this within the repo's files titled `SYSTEM_PROMPT.txt`.
+- For inference, the official MistralAI team recommends `top_p: 0.95` and `temperature: 0.7` with `max_tokens: 40960`.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Limitations
+
+We only support the `mistral-common` tokenizer for Supervised Fine-tuning at the moment and for `type: chat_template` only.
+
+In addition, we do not support overriding tokens yet.
+
+## Related Resources
+
+- [MistralAI Magistral Blog](https://mistral.ai/news/magistral/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
+
+
+## Future Work
+
+- Add parity to Preference Tuning, RL, etc.
+- Add parity to other tokenizer configs like overriding tokens.
diff --git a/examples/magistral/magistral-small-fsdp-qlora.yaml b/examples/magistral/magistral-small-fsdp-qlora.yaml
new file mode 100644
index 0000000000..f31ca7326d
--- /dev/null
+++ b/examples/magistral/magistral-small-fsdp-qlora.yaml
@@ -0,0 +1,76 @@
+base_model: mistralai/Magistral-Small-2506
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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_torch_fused
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing:
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: MistralDecoderLayer
+  fsdp_activation_checkpointing: true
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/magistral/magistral-small-qlora.yaml b/examples/magistral/magistral-small-qlora.yaml
new file mode 100644
index 0000000000..90f6b6f912
--- /dev/null
+++ b/examples/magistral/magistral-small-qlora.yaml
@@ -0,0 +1,67 @@
+base_model: mistralai/Magistral-Small-2506
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/magistral/think/README.md b/examples/magistral/think/README.md
new file mode 100644
index 0000000000..adb22bbba2
--- /dev/null
+++ b/examples/magistral/think/README.md
@@ -0,0 +1,74 @@
+# Magistral Small Thinking Fine-tuning
+
+This guide covers fine-tuning [Magistral Small 2507](https://huggingface.co/mistralai/Magistral-Small-2507) with thinking capabilities using Axolotl. The thinking model enables explicit Chain-of-Thought reasoning with separate thinking and response sections.
+
+## Prerequisites
+
+Before starting, ensure you have:
+
+- Installed Axolotl (see [main README](../README.md))
+
+## Getting Started
+
+Run the thinking model fine-tuning:
+
+```bash
+axolotl train examples/magistral/think/magistral-small-think-qlora.yaml
+```
+
+This config uses about 19.1 GiB VRAM.
+
+### Tips
+
+- Dataset uses multi-content format with `type: thinking` support. See [Dataset Format](#dataset-format) below.
+- You cannot mix `content: str` and `content: list[dict]`, otherwise, dataset loading will fail. Keep it consistent.
+
+## Dataset Format
+
+The thinking model requires the multi-content dataset format with support for an extra `role: thinking` within system and assistant messages.
+
+Example format:
+
+```json
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": [
+                { "type": "text", "text": "{SYSTEM_PROMPT}"}
+            ]
+        },
+        {
+            "role": "user",
+            "content": [
+                { "type": "text", "text": "Solve this step by step: What is 15% of 240?"}
+            ]
+        },
+        {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "thinking",
+                    "thinking": "I need to calculate 15% of 240. First, I'll convert 15% to decimal: 0.15. Then multiply: 0.15 × 240 = 36."
+                },
+                {
+                    "type": "text",
+                    "text": "To find 15% of 240, I'll multiply 240 by 0.15:\n\n240 × 0.15 = 36\n\nTherefore, 15% of 240 is 36."
+                }
+            ]
+        }
+    ]
+}
+```
+
+### Advanced Options
+
+The `thinking` section supports an optional `closed` parameter:
+
+```json
+{
+    "type": "thinking",
+    "thinking": "Internal reasoning here...",
+    "closed": true  // Default: true, controls adding the closing [/THINK] tag
+}
+```
diff --git a/examples/magistral/think/magistral-small-think-qlora.yaml b/examples/magistral/think/magistral-small-think-qlora.yaml
new file mode 100644
index 0000000000..85abe18daf
--- /dev/null
+++ b/examples/magistral/think/magistral-small-think-qlora.yaml
@@ -0,0 +1,67 @@
+base_model: mistralai/Magistral-Small-2507
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: Nanobit/text-think-2k-test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/magistral/vision/README.md b/examples/magistral/vision/README.md
new file mode 100644
index 0000000000..8babc904ac
--- /dev/null
+++ b/examples/magistral/vision/README.md
@@ -0,0 +1,61 @@
+# Magistral Small Vision Fine-tuning
+
+This guide covers fine-tuning [Magistral Small 2509](https://huggingface.co/mistralai/Magistral-Small-2509) with vision capabilities using Axolotl.
+
+## Prerequisites
+
+Before starting, ensure you have:
+
+- Installed Axolotl from source (see [main README](../README.md))
+
+## Getting started
+
+1. Install the required vision lib:
+    ```bash
+    uv pip install 'mistral-common[opencv]==1.8.5'
+    ```
+
+2. Download the example dataset image:
+   ```bash
+   wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+   ```
+
+3. Run the fine-tuning:
+   ```bash
+   axolotl train examples/magistral/vision/magistral-small-vision-24B-qlora.yml
+   ```
+
+This config uses about 17GiB VRAM.
+
+WARNING: The loss and grad norm will be much higher than normal at first. We suspect this to be inherent to the model as of the moment. If anyone would like to submit a fix for this, we are happy to take a look.
+
+### Tips
+
+Key differences from text-only model:
+- `max_tokens: 131072` for inference
+- Multi-modal dataset format required
+- Sample packing not supported
+
+## Dataset Format
+
+The vision model requires multi-modal dataset format as documented [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+One exception is that, passing `"image": PIL.Image` is not supported. MistralTokenizer only supports `path`, `url`, and `base64` for now.
+
+Example:
+```json
+{
+    "messages": [
+        {"role": "system", "content": [{ "type": "text", "text": "{SYSTEM_PROMPT}"}]},
+        {"role": "user", "content": [
+            { "type": "text", "text": "What's in this image?"},
+            {"type": "image", "path": "path/to/image.jpg" }
+        ]},
+        {"role": "assistant", "content": [{ "type": "text", "text": "..." }]},
+    ],
+}
+```
+
+## Limitations
+
+- Sample Packing is not supported for multi-modality training currently.
diff --git a/examples/magistral/vision/magistral-small-vision-24B-qlora.yml b/examples/magistral/vision/magistral-small-vision-24B-qlora.yml
new file mode 100644
index 0000000000..abd2446472
--- /dev/null
+++ b/examples/magistral/vision/magistral-small-vision-24B-qlora.yml
@@ -0,0 +1,64 @@
+base_model: mistralai/Magistral-Small-2509
+processor_type: AutoProcessor
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+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
+
+# sample dataset below requires downloading image in advance
+# wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+datasets:
+  - path: Nanobit/text-vision-2k-test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mamba/config.yml b/examples/mamba/config.yml
index 0a5223bcac..0c39768d8f 100644
--- a/examples/mamba/config.yml
+++ b/examples/mamba/config.yml
@@ -1,18 +1,17 @@
 base_model: state-spaces/mamba-2.8b
+# optionally might have model_type or tokenizer_type or tokenizer_config
 model_type: MambaLMHeadModel
 tokenizer_type: AutoTokenizer
 tokenizer_config: EleutherAI/gpt-neox-20b
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 2048
 sample_packing: false
@@ -35,27 +34,17 @@ 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
+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:
 tokens:
-save_safetensors: False
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mimo/README.md b/examples/mimo/README.md
new file mode 100644
index 0000000000..5ae2143438
--- /dev/null
+++ b/examples/mimo/README.md
@@ -0,0 +1,39 @@
+# Finetune Xiaomi's MiMo with Axolotl
+
+[MiMo](https://huggingface.co/XiaomiMiMo/MiMo-7B-RL) is a family of models trained from scratch for reasoning tasks, incorporating **Multiple-Token Prediction (MTP)** as an additional training objective for enhanced performance and faster inference. Pre-trained on ~25T tokens with a three-stage data mixture strategy and optimized reasoning pattern density.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Run the finetuning example:
+
+    ```bash
+    axolotl train examples/mimo/mimo-7b-qlora.yaml
+    ```
+
+This config uses about 17.2 GiB VRAM. Let us know how it goes. Happy finetuning! 🚀
+
+### Tips
+
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Limitations
+
+**Cut Cross Entropy (CCE)**: Currently not supported. We plan to include CCE support for MiMo in the near future.
+
+## Related Resources
+
+- [MiMo Paper](https://arxiv.org/abs/2505.07608)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/mimo/mimo-7b-qlora.yaml b/examples/mimo/mimo-7b-qlora.yaml
new file mode 100644
index 0000000000..7ced584e11
--- /dev/null
+++ b/examples/mimo/mimo-7b-qlora.yaml
@@ -0,0 +1,67 @@
+base_model: XiaomiMiMo/MiMo-7B-RL
+trust_remote_code: true
+revision_of_model: 6299b5a
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+# CCE - N/A as of now
+# plugins:
+#   - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/ministral/README.md b/examples/ministral/README.md
new file mode 100644
index 0000000000..f8af7bf27f
--- /dev/null
+++ b/examples/ministral/README.md
@@ -0,0 +1,50 @@
+# Finetune Ministral with Axolotl
+
+Ministral is a family of openweight models from MistralAI found on [HuggingFace](mistralai/Ministral-8B-Instruct-2410). This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+    ```bash
+    axolotl train examples/ministral/ministral-small-qlora.yaml
+    ```
+
+This config uses about 8.76 GiB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### Tips
+
+- We recommend adding the same/similar SystemPrompt that the model is tuned for. You can find this within the repo's files titled `SYSTEM_PROMPT.txt`.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Limitations
+
+We only support the `mistral-common` tokenizer for Supervised Fine-tuning at the moment and for `type: chat_template` only.
+
+In addition, we do not support overriding tokens yet.
+
+## Related Resources
+
+- [MistralAI Ministral Blog](https://mistral.ai/news/ministraux)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
+
+
+## Future Work
+
+- Add parity to Preference Tuning, RL, etc.
+- Add parity to other tokenizer configs like overriding tokens.
diff --git a/examples/ministral/ministral-small-qlora.yaml b/examples/ministral/ministral-small-qlora.yaml
new file mode 100644
index 0000000000..4c3bdfe94c
--- /dev/null
+++ b/examples/ministral/ministral-small-qlora.yaml
@@ -0,0 +1,67 @@
+base_model: mistralai/Ministral-8B-Instruct-2410
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/ministral3/README.md b/examples/ministral3/README.md
new file mode 100644
index 0000000000..72f21b746a
--- /dev/null
+++ b/examples/ministral3/README.md
@@ -0,0 +1,79 @@
+# Finetune Ministral3 with Axolotl
+
+Ministral3 is a family of open-weight models from MistralAI found on [HuggingFace](https://huggingface.co/collections/mistralai/ministral-3). This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+Please see [Thinking](#thinking) and [Vision](#vision) for their respective fine-tuning.
+
+Thanks to the team at MistralAI for giving us early access to prepare for these releases.
+
+Note: This is still experimental given it is based on transformers v5 RC.
+
+## Getting started
+
+1. Install Axolotl from source following the [installation guide](https://docs.axolotl.ai/docs/installation.html#sec-edge-build).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Swap to the Axolotl transformers v5 branch
+
+    ```bash
+    cp examples/ministral3/ministral3-3b-qlora.yaml ministral3-3b-qlora.yaml
+
+    git fetch
+    git checkout transformers-v5
+
+    # Install packages for transformers v5
+    uv pip install -e .
+    ```
+
+4. Run the fine-tuning:
+
+    ```bash
+    axolotl train ministral3-3b-qlora.yaml
+    ```
+
+Let us know how it goes. Happy finetuning! 🚀
+
+
+### Tips
+
+- We recommend adding the same/similar SystemPrompt that the model is tuned for. You can find this within the repo's files titled `SYSTEM_PROMPT.txt`.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+### Thinking
+
+Ministral3 2512 model supports thinking capabilities, enabling Chain-of-Thought reasoning with explicit thinking steps.
+
+📚 **[See the Thinking fine-tuning guide →](./think/README.md)**
+
+### Vision
+
+Ministral3 2512 model also supports vision capabilities.
+
+📚 **[See the Vision fine-tuning guide →](./vision/README.md)**
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Limitations
+
+We only support the `mistral-common` tokenizer for Supervised Fine-tuning at the moment and for `type: chat_template` only.
+
+In addition, we do not support overriding tokens yet.
+
+## Related Resources
+
+- [MistralAI Mistral3 Blog](https://mistral.ai/news/mistral-3)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
+
+
+## Future Work
+
+- Add parity to Preference Tuning, RL, etc.
+- Add parity to other tokenizer configs like overriding tokens.
diff --git a/examples/ministral3/ministral3-3b-qlora.yaml b/examples/ministral3/ministral3-3b-qlora.yaml
new file mode 100644
index 0000000000..985f2fad9d
--- /dev/null
+++ b/examples/ministral3/ministral3-3b-qlora.yaml
@@ -0,0 +1,68 @@
+base_model: mistralai/Ministral-3-3B-Reasoning-2512
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+# scaling_softmax: true  # needs flex_attention
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/ministral3/think/README.md b/examples/ministral3/think/README.md
new file mode 100644
index 0000000000..a116dd5ddf
--- /dev/null
+++ b/examples/ministral3/think/README.md
@@ -0,0 +1,74 @@
+# Ministral3 2512 Thinking Fine-tuning
+
+This guide covers fine-tuning [Ministral3 2512](https://huggingface.co/collections/mistralai/ministral-3) with thinking capabilities using Axolotl. The thinking model enables explicit Chain-of-Thought reasoning with separate thinking and response sections.
+
+## Prerequisites
+
+Before starting, ensure you have:
+
+- Installed Axolotl (see [main README](../README.md))
+
+## Getting Started
+
+Run the thinking model fine-tuning:
+
+```bash
+axolotl train examples/ministral3/think/ministral3-3b-think-qlora.yaml
+```
+
+This config uses about 4.76 GiB VRAM.
+
+### Tips
+
+- Dataset uses multi-content format with `type: thinking` support. See [Dataset Format](#dataset-format) below.
+- You cannot mix `content: str` and `content: list[dict]`, otherwise, dataset loading will fail. Keep it consistent.
+
+## Dataset Format
+
+The thinking model requires the multi-content dataset format with support for an extra `role: thinking` within system and assistant messages.
+
+Example format:
+
+```json
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": [
+                { "type": "text", "text": "{SYSTEM_PROMPT}"}
+            ]
+        },
+        {
+            "role": "user",
+            "content": [
+                { "type": "text", "text": "Solve this step by step: What is 15% of 240?"}
+            ]
+        },
+        {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "thinking",
+                    "thinking": "I need to calculate 15% of 240. First, I'll convert 15% to decimal: 0.15. Then multiply: 0.15 × 240 = 36."
+                },
+                {
+                    "type": "text",
+                    "text": "To find 15% of 240, I'll multiply 240 by 0.15:\n\n240 × 0.15 = 36\n\nTherefore, 15% of 240 is 36."
+                }
+            ]
+        }
+    ]
+}
+```
+
+### Advanced Options
+
+The `thinking` section supports an optional `closed` parameter:
+
+```json
+{
+    "type": "thinking",
+    "thinking": "Internal reasoning here...",
+    "closed": true  // Default: true, controls adding the closing [/THINK] tag
+}
+```
diff --git a/examples/ministral3/think/ministral3-3b-think-qlora.yaml b/examples/ministral3/think/ministral3-3b-think-qlora.yaml
new file mode 100644
index 0000000000..508575cac9
--- /dev/null
+++ b/examples/ministral3/think/ministral3-3b-think-qlora.yaml
@@ -0,0 +1,67 @@
+base_model: mistralai/Ministral-3-3B-Reasoning-2512
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: Nanobit/text-think-2k-test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/ministral3/vision/README.md b/examples/ministral3/vision/README.md
new file mode 100644
index 0000000000..cc9a0f38b3
--- /dev/null
+++ b/examples/ministral3/vision/README.md
@@ -0,0 +1,58 @@
+# Ministral3 2512 Vision Fine-tuning
+
+This guide covers fine-tuning [Ministral3 2512](https://huggingface.co/collections/mistralai/ministral-3) with vision capabilities using Axolotl.
+
+## Prerequisites
+
+Before starting, ensure you have:
+
+- Installed Axolotl from source (see [main README](../README.md))
+
+## Getting started
+
+1. Install the required vision lib:
+    ```bash
+    uv pip install 'mistral-common[opencv]==1.8.6'
+    ```
+
+2. Download the example dataset image:
+   ```bash
+   wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+   ```
+
+3. Run the fine-tuning:
+   ```bash
+   axolotl train examples/ministral3/vision/ministral3-3b-vision-qlora.yml
+   ```
+
+WARNING: The loss and grad norm will be much higher than normal at first. We suspect this to be inherent to the model as of the moment. If anyone would like to submit a fix for this, we are happy to take a look.
+
+### Tips
+
+Key differences from text-only model:
+- Multi-modal dataset format required
+- Sample packing not supported
+
+## Dataset Format
+
+The vision model requires multi-modal dataset format as documented [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+One exception is that, passing `"image": PIL.Image` is not supported. MistralTokenizer only supports `path`, `url`, and `base64` for now.
+
+Example:
+```json
+{
+    "messages": [
+        {"role": "system", "content": [{ "type": "text", "text": "{SYSTEM_PROMPT}"}]},
+        {"role": "user", "content": [
+            { "type": "text", "text": "What's in this image?"},
+            {"type": "image", "path": "path/to/image.jpg" }
+        ]},
+        {"role": "assistant", "content": [{ "type": "text", "text": "..." }]},
+    ],
+}
+```
+
+## Limitations
+
+- Sample Packing is not supported for multi-modality training currently.
diff --git a/examples/ministral3/vision/ministral3-3b-vision-qlora.yml b/examples/ministral3/vision/ministral3-3b-vision-qlora.yml
new file mode 100644
index 0000000000..f1430ba539
--- /dev/null
+++ b/examples/ministral3/vision/ministral3-3b-vision-qlora.yml
@@ -0,0 +1,64 @@
+base_model: mistralai/Ministral-3-3B-Reasoning-2512
+processor_type: AutoProcessor
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+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
+
+# sample dataset below requires downloading image in advance
+# wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+datasets:
+  - path: Nanobit/text-vision-2k-test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral-medium-3_5/README.md b/examples/mistral-medium-3_5/README.md
new file mode 100644
index 0000000000..f90397a030
--- /dev/null
+++ b/examples/mistral-medium-3_5/README.md
@@ -0,0 +1,78 @@
+# Finetune Mistral Medium 3.5 with Axolotl
+
+[Mistral Medium 3.5](https://huggingface.co/mistralai/Mistral-Medium-3.5-128B) is a 128B parameter dense multimodal model from MistralAI that unifies instruct, reasoning, and agentic capabilities into a single model.
+It shares the `mistral3` architecture (dense, YaRN RoPE, 256k context) with Ministral 3 and supports the same `reasoning_effort` toggle as Mistral Small 4.
+
+Thanks to the team at MistralAI for giving us early access to prepare for this release.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. (Text config only) Install [Flash Attention 4](https://docs.axolotl.ai/docs/attention.html#flash-attention-4) on Hopper/Blackwell.
+
+4. Run one of the example configs:
+
+    ```bash
+    # text-only
+    axolotl train examples/mistral-medium-3_5/qlora-text.yml  # ~83.1 GiB
+
+    # text + vision
+    # wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+    axolotl train examples/mistral-medium-3_5/qlora-vision.yml  # ~80.3 GiB
+    ```
+
+Note: vision training does not currently work with Flash Attention 4.
+
+## Reasoning Effort
+
+The chat template supports a `reasoning_effort` variable to control the model's reasoning depth:
+
+- `"none"` — instruct mode (default)
+- `"high"` — reasoning mode with explicit thinking steps
+
+Pass it via `chat_template_kwargs` under your dataset config:
+
+```yaml
+datasets:
+  - path: your/dataset
+    type: chat_template
+    chat_template_kwargs:
+      reasoning_effort: high
+```
+
+## Thinking Support
+
+The chat template supports a `thinking` content type in assistant messages for training on reasoning traces (rendered as `[THINK]...[/THINK]` blocks).
+
+To use thinking datasets, add the `thinking` mapping via `message_property_mappings`:
+
+```yaml
+datasets:
+  - path: your/thinking-dataset
+    type: chat_template
+    message_property_mappings:
+      role: role
+      content: content
+      thinking: thinking
+    chat_template_kwargs:
+      reasoning_effort: high
+```
+
+See the [Magistral thinking guide](../magistral/think/README.md) for dataset format details.
+
+## Tips
+
+- For smaller experiments on the same architecture, see [`examples/ministral3`](../ministral3/README.md) (Ministral 3, 3B).
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+- The vision model requires multi-modal dataset format as documented [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+## Related Resources
+
+- [Mistral Medium 3.5 Blog](https://mistral.ai/news/vibe-remote-agents-mistral-medium-3-5)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/mistral-medium-3_5/qlora-text.yml b/examples/mistral-medium-3_5/qlora-text.yml
new file mode 100644
index 0000000000..2ff5b1af38
--- /dev/null
+++ b/examples/mistral-medium-3_5/qlora-text.yml
@@ -0,0 +1,56 @@
+base_model: axolotl-ai-co/Mistral-Medium-3.5-128B-BF16
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir:
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+# prevents targeting vision layers
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+max_steps: 10
+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_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
diff --git a/examples/mistral-medium-3_5/qlora-vision.yml b/examples/mistral-medium-3_5/qlora-vision.yml
new file mode 100644
index 0000000000..a8a51116b6
--- /dev/null
+++ b/examples/mistral-medium-3_5/qlora-vision.yml
@@ -0,0 +1,61 @@
+base_model: axolotl-ai-co/Mistral-Medium-3.5-128B-BF16
+processor_type: AutoProcessor
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_4bit: true
+
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+# sample dataset below requires downloading image in advance
+# wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+datasets:
+  - path: Nanobit/text-vision-2k-test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir:
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+max_steps: 10
+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_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
diff --git a/examples/mistral-small/README.md b/examples/mistral-small/README.md
new file mode 100644
index 0000000000..c5120aab79
--- /dev/null
+++ b/examples/mistral-small/README.md
@@ -0,0 +1,52 @@
+# Mistral Small 3.1/3.2 Fine-tuning
+
+This guide covers fine-tuning [Mistral Small 3.1](mistralai/Mistral-Small-3.1-24B-Instruct-2503) and [Mistral Small 3.2](mistralai/Mistral-Small-3.2-24B-Instruct-2506) with vision capabilities using Axolotl.
+
+## Prerequisites
+
+Before starting, ensure you have:
+
+- Installed Axolotl (see [Installation docs](https://docs.axolotl.ai/docs/installation.html))
+
+## Getting Started
+
+1. Install the required vision lib:
+    ```bash
+    uv pip install 'mistral-common[opencv]==1.8.5'
+    ```
+
+2. Download the example dataset image:
+   ```bash
+   wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+   ```
+
+3. Run the fine-tuning:
+   ```bash
+   axolotl train examples/mistral/mistral-small/mistral-small-3.1-24B-lora.yml
+   ```
+
+This config uses about 29.4 GiB VRAM.
+
+## Dataset Format
+
+The vision model requires multi-modal dataset format as documented [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+One exception is that, passing `"image": PIL.Image` is not supported. MistralTokenizer only supports `path`, `url`, and `base64` for now.
+
+Example:
+```json
+{
+    "messages": [
+        {"role": "system", "content": [{ "type": "text", "text": "{SYSTEM_PROMPT}"}]},
+        {"role": "user", "content": [
+            { "type": "text", "text": "What's in this image?"},
+            {"type": "image", "path": "path/to/image.jpg" }
+        ]},
+        {"role": "assistant", "content": [{ "type": "text", "text": "..." }]},
+    ],
+}
+```
+
+## Limitations
+
+- Sample Packing is not supported for multi-modality training currently.
diff --git a/examples/mistral-small/mistral-small-3.1-24B-lora.yml b/examples/mistral-small/mistral-small-3.1-24B-lora.yml
new file mode 100644
index 0000000000..4d3f78a132
--- /dev/null
+++ b/examples/mistral-small/mistral-small-3.1-24B-lora.yml
@@ -0,0 +1,62 @@
+base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
+processor_type: AutoProcessor
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+load_in_8bit: 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
+
+# sample dataset below requires downloading image in advance
+# wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+datasets:
+  - path: Nanobit/text-vision-2k-test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+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
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/bigstral-ds-zero3.yaml b/examples/mistral/bigstral/bigstral-ds-zero3.yaml
similarity index 76%
rename from examples/mistral/bigstral-ds-zero3.yaml
rename to examples/mistral/bigstral/bigstral-ds-zero3.yaml
index cc0a44b2a4..4648ae4b4e 100644
--- a/examples/mistral/bigstral-ds-zero3.yaml
+++ b/examples/mistral/bigstral/bigstral-ds-zero3.yaml
@@ -1,11 +1,11 @@
 base_model: mistral-community/Mixtral-8x22B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: LlamaTokenizer
-trust_remote_code: true
+# 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
+trust_remote_code: true
 
 unfrozen_parameters:
   - ^lm_head.weight$
@@ -23,11 +23,11 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 2048
 sample_packing: true
-pad_to_sequence_len: true
+
 
 gradient_accumulation_steps: 1
 micro_batch_size: 1
@@ -36,28 +36,22 @@ 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
+attn_implementation: flash_attention_2
 
 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:
   - "<|im_start|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/config.yml b/examples/mistral/config.yml
index c909c63e22..aa10667337 100644
--- a/examples/mistral/config.yml
+++ b/examples/mistral/config.yml
@@ -1,21 +1,20 @@
 base_model: mistralai/Mistral-7B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: MistralForCausalLM
 tokenizer_type: LlamaTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 8192
 sample_packing: true
-pad_to_sequence_len: true
+
 eval_sample_packing: false
 
 wandb_project:
@@ -31,28 +30,18 @@ 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 10
+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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/dpo/mistral-dpo-qlora.yml b/examples/mistral/dpo/mistral-dpo-qlora.yml
new file mode 100644
index 0000000000..604eada744
--- /dev/null
+++ b/examples/mistral/dpo/mistral-dpo-qlora.yml
@@ -0,0 +1,83 @@
+#Note that we are switching from the regular chat template to chatml.
+#If you experience problems with the special tokens, training for more epochs can help.
+#After training, merge the model before inference otherwise you might
+#face problems with the special tokens.
+
+base_model: mistralai/Mistral-7B-Instruct-v0.2
+# optionally might have model_type or tokenizer_type
+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: true
+
+chat_template: chatml
+rl: dpo
+datasets:
+  - path: olivermolenschot/alpaca_messages_dpo_test
+    type: chat_template.default
+    field_messages: conversation
+    field_chosen: chosen
+    field_rejected: rejected
+    message_property_mappings:
+      role: role
+      content: content
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/dpo-qlora
+
+sequence_len: 2048
+sample_packing: false
+
+
+adapter: qlora
+lora_model_dir:
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.2
+lora_target_linear: true
+
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+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: 16
+num_epochs: 6
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0001
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+  bos_token: "<|im_start|>"
+  eos_token: "<|im_end|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/lora.yml b/examples/mistral/lora.yml
index ac9ac0dd98..b157fcc219 100644
--- a/examples/mistral/lora.yml
+++ b/examples/mistral/lora.yml
@@ -1,30 +1,31 @@
 base_model: mistralai/Mistral-7B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: MistralForCausalLM
 tokenizer_type: LlamaTokenizer
+# 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
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.1
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 adapter: lora
 lora_model_dir:
 
 sequence_len: 8192
 sample_packing: true
-pad_to_sequence_len: true
+
 
 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,31 +48,21 @@ 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
+attn_implementation: flash_attention_2
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
 
-warmup_steps: 10
+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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/mistral-qlora-fsdp.yml b/examples/mistral/mistral-qlora-fsdp.yml
index 71ac1e701f..27d8be3cd0 100644
--- a/examples/mistral/mistral-qlora-fsdp.yml
+++ b/examples/mistral/mistral-qlora-fsdp.yml
@@ -1,18 +1,21 @@
 base_model: mistralai/Mixtral-8x7B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: LlamaTokenizer
+# 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: true
-strict: false
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.02
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 model_config:
   output_router_logits: true
@@ -28,7 +31,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -43,29 +45,21 @@ 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
-early_stopping_patience:
 resume_from_checkpoint:
-local_rank:
 logging_steps: 1
-xformers_attention:
-flash_attention: true
+attn_implementation: flash_attention_2
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
-eval_table_size:
-eval_max_new_tokens: 128
 saves_per_epoch: 1
-debug:
+
 weight_decay: 0.0
 fsdp:
   - full_shard
@@ -80,3 +74,5 @@ fsdp_config:
   fsdp_state_dict_type: FULL_STATE_DICT
   fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
 special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/mixtral-8x22b-qlora-fsdp.yml b/examples/mistral/mixtral/mixtral-8x22b-qlora-fsdp.yml
similarity index 77%
rename from examples/mistral/mixtral-8x22b-qlora-fsdp.yml
rename to examples/mistral/mixtral/mixtral-8x22b-qlora-fsdp.yml
index ac80a2a756..1b66de8f0f 100644
--- a/examples/mistral/mixtral-8x22b-qlora-fsdp.yml
+++ b/examples/mistral/mixtral/mixtral-8x22b-qlora-fsdp.yml
@@ -1,17 +1,19 @@
 base_model: mistral-community/Mixtral-8x22B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 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: true
-strict: false
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.02
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 model_config:
   output_router_logits: true
@@ -27,7 +29,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -38,33 +39,25 @@ wandb_log_model:
 gradient_accumulation_steps: 4
 micro_batch_size: 2
 num_epochs: 1
-optimizer: adamw_torch
+optimizer: adamw_torch_fused
 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
+attn_implementation: flash_attention_2
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
-eval_table_size:
-eval_max_new_tokens: 128
 saves_per_epoch: 1
-debug:
+
 weight_decay: 0.0
 fsdp:
   - full_shard
@@ -79,3 +72,5 @@ fsdp_config:
   fsdp_state_dict_type: FULL_STATE_DICT
   fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
 special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/mixtral-qlora-fsdp.yml b/examples/mistral/mixtral/mixtral-qlora-fsdp.yml
similarity index 78%
rename from examples/mistral/mixtral-qlora-fsdp.yml
rename to examples/mistral/mixtral/mixtral-qlora-fsdp.yml
index b6a07ae51c..bd7c8620ef 100644
--- a/examples/mistral/mixtral-qlora-fsdp.yml
+++ b/examples/mistral/mixtral/mixtral-qlora-fsdp.yml
@@ -1,18 +1,21 @@
 base_model: mistralai/Mixtral-8x7B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: LlamaTokenizer
+# 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: true
-strict: false
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.02
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 model_config:
   output_router_logits: true
@@ -28,7 +31,6 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -39,33 +41,25 @@ wandb_log_model:
 gradient_accumulation_steps: 4
 micro_batch_size: 2
 num_epochs: 1
-optimizer: adamw_torch
+optimizer: adamw_torch_fused
 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
+attn_implementation: flash_attention_2
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
-eval_table_size:
-eval_max_new_tokens: 128
 saves_per_epoch: 1
-debug:
+
 weight_decay: 0.0
 fsdp:
   - full_shard
@@ -83,3 +77,5 @@ fsdp_config:
   fsdp_forward_prefetch: false
   fsdp_backward_prefetch: BACKWARD_PRE
 special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/mixtral.yml b/examples/mistral/mixtral/mixtral.yml
similarity index 80%
rename from examples/mistral/mixtral.yml
rename to examples/mistral/mixtral/mixtral.yml
index 5ee3da9d65..b493ed3172 100644
--- a/examples/mistral/mixtral.yml
+++ b/examples/mistral/mixtral/mixtral.yml
@@ -1,18 +1,21 @@
 base_model: mistralai/Mixtral-8x7B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: LlamaTokenizer
+# 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: true
-strict: false
 
 datasets:
   - path: tatsu-lab/alpaca
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 ## You can optionally freeze the entire model and unfreeze a subset of parameters
 unfrozen_parameters:
@@ -31,13 +34,12 @@ lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 #lora_target_modules:
 #  - gate
 #  - q_proj
@@ -61,31 +63,23 @@ 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
+attn_implementation: flash_attention_2
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
 
-warmup_steps: 10
+warmup_ratio: 0.1
 evals_per_epoch: 4
-eval_table_size:
-eval_max_new_tokens: 128
 saves_per_epoch: 1
-debug:
+
 deepspeed: deepspeed_configs/zero2.json
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/mixtral_22.yml b/examples/mistral/mixtral/mixtral_22.yml
similarity index 75%
rename from examples/mistral/mixtral_22.yml
rename to examples/mistral/mixtral/mixtral_22.yml
index 9abb6f407a..3b87af04e6 100644
--- a/examples/mistral/mixtral_22.yml
+++ b/examples/mistral/mixtral/mixtral_22.yml
@@ -1,11 +1,11 @@
 base_model: mistral-community/Mixtral-8x22B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: LlamaTokenizer
-trust_remote_code: true
+# 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
+trust_remote_code: true
 
 unfrozen_parameters:
   - ^lm_head.weight$
@@ -21,11 +21,11 @@ model_config:
 datasets:
   - path: yahma/alpaca-cleaned
     type: alpaca
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 8000
 sample_packing: true
-pad_to_sequence_len: true
+
 
 gradient_accumulation_steps: 1
 micro_batch_size: 1
@@ -34,28 +34,22 @@ 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
+attn_implementation: flash_attention_2
 
 save_total_limit: 1
 save_steps:
-debug:
+
 deepspeed: deepspeed_configs/zero3_bf16_cpuoffload_all.json
 weight_decay: 0.0
-fsdp:
-fsdp_config:
 special_tokens:
   eos_token: "<|im_end|>"
 tokens:
   - "<|im_start|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/lora-mps.yml b/examples/mistral/mps/lora-mps.yml
similarity index 69%
rename from examples/mistral/lora-mps.yml
rename to examples/mistral/mps/lora-mps.yml
index 31b0d527e2..1b80210852 100644
--- a/examples/mistral/lora-mps.yml
+++ b/examples/mistral/mps/lora-mps.yml
@@ -1,17 +1,16 @@
 base_model: mistralai/Mistral-7B-v0.1
+# optionally might have model_type or tokenizer_type
 model_type: MistralForCausalLM
 tokenizer_type: LlamaTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 eval_sample_packing: false
 
 adapter: lora
@@ -19,13 +18,12 @@ lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
-pad_to_sequence_len: true
+
 
 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
@@ -44,36 +42,26 @@ wandb_log_model:
 gradient_accumulation_steps: 8
 micro_batch_size: 1
 num_epochs: 2
-optimizer: adamw_torch
+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
+attn_implementation: sdpa
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
 
-warmup_steps: 10
+warmup_ratio: 0.1
 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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/mistral-qlora-orpo.yml b/examples/mistral/orpo/mistral-qlora-orpo.yml
similarity index 75%
rename from examples/mistral/mistral-qlora-orpo.yml
rename to examples/mistral/orpo/mistral-qlora-orpo.yml
index 7727fd7485..d1c0065e50 100644
--- a/examples/mistral/mistral-qlora-orpo.yml
+++ b/examples/mistral/orpo/mistral-qlora-orpo.yml
@@ -1,10 +1,12 @@
 base_model: mistralai/Mistral-7B-v0.1
+# optionally might have model_type or tokenizer_type
 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: true
-strict: false
 
 rl: orpo
 orpo_alpha: 0.1
@@ -16,20 +18,19 @@ datasets:
     type: chat_template.argilla
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.1
-output_dir: ./mistral-qlora-orpo-out
+output_dir: ./outputs/mistral-qlora-orpo-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: false
-pad_to_sequence_len: true
+
 
 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
@@ -52,31 +53,21 @@ 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
+attn_implementation: flash_attention_2
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
 
-warmup_steps: 10
+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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral/qlora.yml b/examples/mistral/qlora.yml
index 6fbbb96183..4fa82d11e0 100644
--- a/examples/mistral/qlora.yml
+++ b/examples/mistral/qlora.yml
@@ -1,30 +1,31 @@
 base_model: mistralai/Mistral-7B-v0.1
+# optionally might have model_type or tokenizer_type
 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: true
-strict: false
 
 datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.1
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 8192
 sample_packing: true
-pad_to_sequence_len: true
+
 
 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,31 +48,21 @@ 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
+attn_implementation: flash_attention_2
 
 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
 
-warmup_steps: 10
+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:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/mistral4/README.md b/examples/mistral4/README.md
new file mode 100644
index 0000000000..ccbb03b216
--- /dev/null
+++ b/examples/mistral4/README.md
@@ -0,0 +1,82 @@
+# Finetune Mistral Small 4 with Axolotl
+
+Mistral Small 4 is a 119B parameter (6.5B active) multimodal MoE model from MistralAI that unifies instruct, reasoning, and coding capabilities into a single model. It is available on HuggingFace at [Mistral-Small-4-119B-2603](https://huggingface.co/mistralai/Mistral-Small-4-119B-2603).
+
+Thanks to the team at MistralAI for giving us early access to prepare for this release.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage
+
+3. Install transformers from main
+
+  ```bash
+  uv pip install git+https://github.com/huggingface/transformers.git
+  ```
+
+4. Run one of the example configs:
+
+  ```bash
+  # text-only
+  axolotl train examples/mistral4/qlora-text.yml  # no experts ~69 GiB, experts ~93 GiB
+  axolotl train examples/mistral4/fft-text.yml
+
+  # text + vision
+  # run: wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+  axolotl train examples/mistral4/qlora-vision.yml  # no experts ~68 GiB
+  axolotl train examples/mistral4/fft-vision.yml
+  ```
+
+Note: FFT configs provided as reference. Please adjust hyperparameters as needed.
+
+## Reasoning Effort
+
+The chat template supports a `reasoning_effort` variable to control the model's reasoning depth:
+
+- `"none"` — instruct mode (default)
+- `"high"` — reasoning mode with explicit thinking steps
+
+Pass it via `chat_template_kwargs` under your dataset config:
+
+```yaml
+datasets:
+  - path: your/dataset
+    type: chat_template
+    chat_template_kwargs:
+      reasoning_effort: high
+```
+
+## Thinking Support
+
+The chat template supports a `thinking` content type in assistant messages for training on reasoning traces (rendered as `[THINK]...[/THINK]` blocks).
+
+To use thinking datasets, add the `thinking` mapping via `message_property_mappings`:
+
+```yaml
+datasets:
+  - path: your/thinking-dataset
+    type: chat_template
+    message_property_mappings:
+      role: role
+      content: content
+      thinking: thinking
+    chat_template_kwargs:
+      reasoning_effort: high
+```
+
+See the [Magistral thinking guide](../magistral/think/README.md) for dataset format details.
+
+## Tips
+
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+- The vision model requires multi-modal dataset format as documented [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+## Related Resources
+
+- [MistralAI Mistral Small 4 Blog](https://mistral.ai/news/mistral-small-4)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/mistral4/fft-text.yml b/examples/mistral4/fft-text.yml
new file mode 100644
index 0000000000..2cdab6a423
--- /dev/null
+++ b/examples/mistral4/fft-text.yml
@@ -0,0 +1,58 @@
+base_model: axolotl-ai-co/Mistral-Small-4-119B-2603-BF16
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+use_kernels: true
+use_sonicmoe: true
+
+# only train language model layers, freeze vision tower
+unfrozen_parameters:
+  - model.language_model.*
+  - lm_head
+  - embed_tokens
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+sequence_len: 2048
+sample_packing: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: false
+  state_dict_type: FULL_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Mistral4DecoderLayer
+  reshard_after_forward: true
+  activation_checkpointing: true
diff --git a/examples/mistral4/fft-vision.yml b/examples/mistral4/fft-vision.yml
new file mode 100644
index 0000000000..22262c55af
--- /dev/null
+++ b/examples/mistral4/fft-vision.yml
@@ -0,0 +1,57 @@
+base_model: axolotl-ai-co/Mistral-Small-4-119B-2603-BF16
+processor_type: AutoProcessor
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+use_kernels: true
+use_sonicmoe: true
+
+# vision requirements
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+datasets:
+  - path: Nanobit/text-vision-2k-test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+sequence_len: 2048
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+fsdp_version: 2
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: false
+  state_dict_type: FULL_STATE_DICT
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Mistral4DecoderLayer
+  reshard_after_forward: true
+  activation_checkpointing: true
diff --git a/examples/mistral4/qlora-text.yml b/examples/mistral4/qlora-text.yml
new file mode 100644
index 0000000000..887ce6da09
--- /dev/null
+++ b/examples/mistral4/qlora-text.yml
@@ -0,0 +1,58 @@
+base_model: axolotl-ai-co/Mistral-Small-4-119B-2603-BF16
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_4bit: true
+quantize_moe_experts: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+adapter: qlora
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+# uncomment to train on expert layers
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+# lora_mlp_kernel: false
+# lora_qkv_kernel: false
+# lora_o_kernel: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: true
+tf32: true
+
+gradient_checkpointing: true
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/mistral4/qlora-vision.yml b/examples/mistral4/qlora-vision.yml
new file mode 100644
index 0000000000..d01f8e85b5
--- /dev/null
+++ b/examples/mistral4/qlora-vision.yml
@@ -0,0 +1,63 @@
+base_model: axolotl-ai-co/Mistral-Small-4-119B-2603-BF16
+processor_type: AutoProcessor
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_4bit: true
+quantize_moe_experts: true
+
+# vision chat template requirements
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+datasets:
+  - path: Nanobit/text-vision-2k-test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+adapter: qlora
+
+sequence_len: 2048
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+# uncomment to train on expert layers
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+# lora_mlp_kernel: false
+# lora_qkv_kernel: false
+# lora_o_kernel: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: true
+tf32: true
+
+gradient_checkpointing: true
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/nemotron-h/120b-a12b-qlora.yaml b/examples/nemotron-h/120b-a12b-qlora.yaml
new file mode 100644
index 0000000000..1174cec21f
--- /dev/null
+++ b/examples/nemotron-h/120b-a12b-qlora.yaml
@@ -0,0 +1,82 @@
+base_model: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_layer_norm: true
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_rms_norm_gated: true
+
+# LoRA kernel patches are incompatible with this architecture — see README.
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+chat_template: tokenizer_default
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 4096
+sample_packing: true
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+
+# To also train MoE expert weights, add them via lora_target_parameters
+# (they are 3D nn.Parameter tensors, not nn.Linear — no gate_proj):
+# lora_target_parameters:
+#   - up_proj
+#   - down_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_torch_4bit
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 2
+saves_per_epoch: 1
+weight_decay: 0.0
+
+special_tokens:
diff --git a/examples/nemotron-h/README.md b/examples/nemotron-h/README.md
new file mode 100644
index 0000000000..3f90714310
--- /dev/null
+++ b/examples/nemotron-h/README.md
@@ -0,0 +1,48 @@
+# Nemotron-H (nvidia/NVIDIA-Nemotron-3-*)
+
+Hybrid Mamba2 / Attention / MoE architecture (`model_type: nemotron_h`).
+
+| Model | Total params | Active params | Layers |
+|---|---|---|---|
+| NVIDIA-Nemotron-3-Super-120B-A12B-BF16 | 120B | ~12B | 88 |
+| NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 | 30B | ~3B | — |
+
+## Requirements
+
+```bash
+pip install mamba-ssm causal-conv1d   # fast Mamba2 CUDA kernels
+```
+
+## Architecture notes
+
+- Three block types per layer: **Mamba2** (selective SSM), **Attention** (sparse), **MoE** (mixture-of-experts).
+- Only ~12 out of 88 blocks are attention layers (120B variant).
+- MLP activation is `relu2` via `mlp_hidden_act` (not the usual `hidden_act`).
+
+## LoRA kernel patches
+
+All three LoRA Triton kernel patches must be disabled:
+
+```yaml
+lora_qkv_kernel: false   # attention lives in NemotronHBlock.mixer, not layer.self_attn
+lora_o_kernel: false     # same reason
+lora_mlp_kernel: false   # relu2 (mlp_hidden_act) is not supported by lora_mlp_kernel
+```
+
+## MoE expert weights
+
+NemotronH experts store `up_proj` and `down_proj` as 3D `nn.Parameter` tensors
+(shape `[num_experts, out_dim, in_dim]`), **not** `nn.Linear` modules — there is no
+`gate_proj`. To fine-tune them alongside attention, use `lora_target_parameters`
+instead of `lora_target_modules`:
+
+```yaml
+lora_target_parameters:
+  - up_proj
+  - down_proj
+```
+
+## Limitations
+
+- **MoE Triton kernels**: `lora_mlp_kernel` is not supported for NemotronH's MoE expert layers. The expert weights are 3D `nn.Parameter` tensors (not `nn.Linear`), which the Triton kernel does not support. Keep `lora_mlp_kernel: false`.
+- **Gradient checkpointing**: Only supported when `sample_packing: true`. Without sample packing the upstream model marks `supports_gradient_checkpointing = False`.
diff --git a/examples/nemotron-h/nano-30b-a3b-qlora-cp.yaml b/examples/nemotron-h/nano-30b-a3b-qlora-cp.yaml
new file mode 100644
index 0000000000..cab402dccf
--- /dev/null
+++ b/examples/nemotron-h/nano-30b-a3b-qlora-cp.yaml
@@ -0,0 +1,85 @@
+# See examples/nemotron-h/README.md for architecture notes and requirements.
+base_model: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_layer_norm: true
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_rms_norm_gated: true
+
+# LoRA kernel patches are incompatible with this architecture — see README.
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+chat_template: tokenizer_default
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 4096
+sample_packing: true
+
+context_parallel_size: 2
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+
+# To also train MoE expert weights, add them via lora_target_parameters
+# (they are 3D nn.Parameter tensors, not nn.Linear — no gate_proj):
+# lora_target_parameters:
+#   - up_proj
+#   - down_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+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: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+
+special_tokens:
diff --git a/examples/nemotron-h/nano-30b-a3b-qlora.yaml b/examples/nemotron-h/nano-30b-a3b-qlora.yaml
new file mode 100644
index 0000000000..206bd5df87
--- /dev/null
+++ b/examples/nemotron-h/nano-30b-a3b-qlora.yaml
@@ -0,0 +1,83 @@
+# See examples/nemotron-h/README.md for architecture notes and requirements.
+base_model: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_layer_norm: true
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_rms_norm_gated: true
+
+# LoRA kernel patches are incompatible with this architecture — see README.
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+chat_template: tokenizer_default
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 4096
+sample_packing: true
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+
+# To also train MoE expert weights, add them via lora_target_parameters
+# (they are 3D nn.Parameter tensors, not nn.Linear — no gate_proj):
+# lora_target_parameters:
+#   - up_proj
+#   - down_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+
+special_tokens:
diff --git a/examples/nemotron/nemotron-mini-4b-qlora.yaml b/examples/nemotron/nemotron-mini-4b-qlora.yaml
new file mode 100644
index 0000000000..3f3772071e
--- /dev/null
+++ b/examples/nemotron/nemotron-mini-4b-qlora.yaml
@@ -0,0 +1,57 @@
+base_model: nvidia/Nemotron-Mini-4B-Instruct
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/nemotron-mini-4b-qlora
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 4096
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - up_proj
+  - down_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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+special_tokens:
diff --git a/examples/olmo3/README.md b/examples/olmo3/README.md
new file mode 100644
index 0000000000..1602236282
--- /dev/null
+++ b/examples/olmo3/README.md
@@ -0,0 +1,38 @@
+# Finetune Allenai's Olmo 3 with Axolotl
+
+[Olmo 3](https://huggingface.co/collections/allenai/olmo-3) are a family of 7B and 32B models open source models trained by The Allen Institute for Artificial Intelligence.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+    ```bash
+    axolotl train examples/olmo3/olmo3-7b-qlora.yaml
+    ```
+
+This uses about 11.3 GiB VRAM. Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- The example config can be re-used for Olmo and Olmo 2.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [Olmo 3 Blog](https://allenai.org/blog/olmo3)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/olmo3/olmo3-7b-qlora.yaml b/examples/olmo3/olmo3-7b-qlora.yaml
new file mode 100644
index 0000000000..b494699e0a
--- /dev/null
+++ b/examples/olmo3/olmo3-7b-qlora.yaml
@@ -0,0 +1,64 @@
+base_model: allenai/Olmo-3-7B-Instruct-SFT
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/orpheus/README.md b/examples/orpheus/README.md
new file mode 100644
index 0000000000..5fea05ecf4
--- /dev/null
+++ b/examples/orpheus/README.md
@@ -0,0 +1,341 @@
+# Finetuning LLMs to output audio
+
+In this example, we finetune Orpcanopylabs/orpheus-tts-0.1-pretrained (a LLaMA 3.2 3b model) to output audio.
+
+The `finetune.yml` withe current settings will run on any Nvidia GPU with 45GB VRAM or more. If you adjust the batch size it can easily run on any GPU under 24GB.
+
+## Dataset pre-processing for pre-training
+If you are adding another voice in English, please jump ahead to finetuning pre-processing.
+
+For this to work, we need to preprocess our dataset. Since we are expecting to output audio, we will need to add tokens to the tokenizer.
+
+Using this code, it will download the SNAC model and add the correct tokens and upload the final dataset.
+
+```python
+import torch
+from snac import SNAC
+from datasets import load_dataset
+from huggingface_hub import snapshot_download
+from datasets import load_dataset
+import random
+import torchaudio.transforms as T
+from transformers import AutoTokenizer
+import os
+
+my_original_dataset_name = "<huggingface-id-of-dataset-that-we-want-to-preprocess>"
+name_to_push_dataset_to = "<huggingface-id-of-where-to-save-dataset>"
+
+dsn = my_original_dataset_name
+
+snapshot_download(
+    repo_id=dsn,
+    repo_type="dataset",
+    revision="main",
+    max_workers=64,
+)
+
+
+ds = load_dataset(dsn, split="train")
+ds_sample_rate = ds[0]["audio"]["sampling_rate"]
+
+model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
+model = model.to("mps")
+
+def tokenise_audio(waveform):
+  waveform = torch.from_numpy(waveform).unsqueeze(0)
+  waveform = waveform.to(dtype=torch.float32)
+  resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=24000)
+  waveform = resample_transform(waveform)
+
+  waveform = waveform.unsqueeze(0).to("cuda")
+
+  #generate the codes from snac
+  with torch.inference_mode():
+    codes = model.encode(waveform)
+
+  all_codes = []
+  for i in range(codes[0].shape[1]):
+    all_codes.append(codes[0][0][i].item()+128266)
+    all_codes.append(codes[1][0][2*i].item()+128266+4096)
+    all_codes.append(codes[2][0][4*i].item()+128266+(2*4096))
+    all_codes.append(codes[2][0][(4*i)+1].item()+128266+(3*4096))
+    all_codes.append(codes[1][0][(2*i)+1].item()+128266+(4*4096))
+    all_codes.append(codes[2][0][(4*i)+2].item()+128266+(5*4096))
+    all_codes.append(codes[2][0][(4*i)+3].item()+128266+(6*4096))
+
+
+  return all_codes
+
+def add_codes(example):
+    # Always initialize codes_list to None
+    codes_list = None
+
+    try:
+        answer_audio = example.get("audio")
+        # If there's a valid audio array, tokenise it
+        if answer_audio and "array" in answer_audio:
+            audio_array = answer_audio["array"]
+            codes_list = tokenise_audio(audio_array)
+    except Exception as e:
+        print(f"Skipping row due to error: {e}")
+        # Keep codes_list as None if we fail
+    example["codes_list"] = codes_list
+
+    return example
+
+ds = ds.map(add_codes, remove_columns=["audio"])
+
+#@title Load Tokenizer
+tokeniser_length = 128256
+start_of_text = 128000
+end_of_text = 128009
+
+start_of_speech = tokeniser_length + 1
+end_of_speech = tokeniser_length + 2
+
+start_of_human = tokeniser_length + 3
+end_of_human = tokeniser_length + 4
+
+start_of_ai = tokeniser_length + 5
+end_of_ai =  tokeniser_length + 6
+pad_token = tokeniser_length + 7
+
+audio_tokens_start = tokeniser_length + 10
+
+tokenizer_name = "canopylabs/orpheus-3b-0.1-pretrained"
+
+
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+num_proc = os.cpu_count() - 2
+
+ds = ds.filter(lambda x: x["codes_list"] is not None)
+ds = ds.filter(lambda x: len(x["codes_list"]) > 0)
+
+#@title Create Input Ids
+def remove_duplicate_frames(example):
+    vals = example["codes_list"]
+    if len(vals) % 7 != 0:
+        raise ValueError("Input list length must be divisible by 7")
+
+    result = vals[:7]
+
+    removed_frames = 0
+
+    for i in range(7, len(vals), 7):
+        current_first = vals[i]
+        previous_first = result[-7]
+
+        if current_first != previous_first:
+            result.extend(vals[i:i+7])
+        else:
+            removed_frames += 1
+
+    example["codes_list"] = result
+
+    return example
+
+ds = ds.map(remove_duplicate_frames, num_proc=num_proc)
+
+
+def create_input_ids(example):
+    text_ids = tokenizer.encode({example['text']},  add_special_tokens=True)
+    text_ids.append(end_of_text)
+    example["text_tokens"] = text_ids
+    input_ids = (
+        [start_of_human]
+        + example["text_tokens"]
+        + [end_of_human]
+        + [start_of_ai]
+        + [start_of_speech]
+        + example["codes_list"]
+        + [end_of_speech]
+        + [end_of_ai]
+    )
+    example["input_ids"] = input_ids
+    example["labels"] = input_ids
+    example["attention_mask"] = [1] * len(input_ids)
+
+    return example
+
+ds = ds.map(create_input_ids, num_proc=num_proc, remove_columns=["text", "codes_list"])
+
+#@title Remove unnecessary columns
+columns_to_keep = ["input_ids", "labels", "attention_mask"]
+columns_to_remove = [col for col in ds.column_names if col not in columns_to_keep]
+
+ds = ds.remove_columns(columns_to_remove)
+
+ds.push_to_hub(name_to_push_dataset_to)
+```
+
+
+## Finetune pre-processing
+Use this code to add a new voice.
+
+```python
+import torch
+from snac import SNAC
+from datasets import load_dataset
+from huggingface_hub import snapshot_download
+from datasets import load_dataset
+import random
+import torchaudio.transforms as T
+from transformers import AutoTokenizer
+import os
+
+my_original_dataset_name = "<huggingface-id-of-dataset-that-we-want-to-preprocess>"
+name_to_push_dataset_to = "<huggingface-id-of-where-to-save-dataset>"
+
+dsn = my_original_dataset_name
+
+snapshot_download(
+    repo_id=dsn,
+    repo_type="dataset",
+    revision="main",
+    max_workers=64,
+)
+
+
+ds = load_dataset(dsn, split="train")
+ds_sample_rate = ds[0]["audio"]["sampling_rate"]
+
+model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
+model = model.to("mps")
+
+def tokenise_audio(waveform):
+  waveform = torch.from_numpy(waveform).unsqueeze(0)
+  waveform = waveform.to(dtype=torch.float32)
+  resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=24000)
+  waveform = resample_transform(waveform)
+
+  waveform = waveform.unsqueeze(0).to("cuda")
+
+  #generate the codes from snac
+  with torch.inference_mode():
+    codes = model.encode(waveform)
+
+  all_codes = []
+  for i in range(codes[0].shape[1]):
+    all_codes.append(codes[0][0][i].item()+128266)
+    all_codes.append(codes[1][0][2*i].item()+128266+4096)
+    all_codes.append(codes[2][0][4*i].item()+128266+(2*4096))
+    all_codes.append(codes[2][0][(4*i)+1].item()+128266+(3*4096))
+    all_codes.append(codes[1][0][(2*i)+1].item()+128266+(4*4096))
+    all_codes.append(codes[2][0][(4*i)+2].item()+128266+(5*4096))
+    all_codes.append(codes[2][0][(4*i)+3].item()+128266+(6*4096))
+
+
+  return all_codes
+
+def add_codes(example):
+    # Always initialize codes_list to None
+    codes_list = None
+
+    try:
+        answer_audio = example.get("audio")
+        # If there's a valid audio array, tokenise it
+        if answer_audio and "array" in answer_audio:
+            audio_array = answer_audio["array"]
+            codes_list = tokenise_audio(audio_array)
+    except Exception as e:
+        print(f"Skipping row due to error: {e}")
+        # Keep codes_list as None if we fail
+    example["codes_list"] = codes_list
+
+    return example
+
+ds = ds.map(add_codes, remove_columns=["audio"])
+
+#@title Load Tokenizer
+tokeniser_length = 128256
+start_of_text = 128000
+end_of_text = 128009
+
+start_of_speech = tokeniser_length + 1
+end_of_speech = tokeniser_length + 2
+
+start_of_human = tokeniser_length + 3
+end_of_human = tokeniser_length + 4
+
+start_of_ai = tokeniser_length + 5
+end_of_ai =  tokeniser_length + 6
+pad_token = tokeniser_length + 7
+
+audio_tokens_start = tokeniser_length + 10
+
+tokenizer_name = "canopylabs/orpheus-3b-0.1-pretrained"
+
+
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+num_proc = os.cpu_count() - 2
+
+ds = ds.filter(lambda x: x["codes_list"] is not None)
+ds = ds.filter(lambda x: len(x["codes_list"]) > 0)
+
+#@title Create Input Ids
+def remove_duplicate_frames(example):
+    vals = example["codes_list"]
+    if len(vals) % 7 != 0:
+        raise ValueError("Input list length must be divisible by 7")
+
+    result = vals[:7]
+
+    removed_frames = 0
+
+    for i in range(7, len(vals), 7):
+        current_first = vals[i]
+        previous_first = result[-7]
+
+        if current_first != previous_first:
+            result.extend(vals[i:i+7])
+        else:
+            removed_frames += 1
+
+    example["codes_list"] = result
+
+    return example
+
+ds = ds.map(remove_duplicate_frames, num_proc=num_proc)
+
+tok_info = '''*** HERE you can modify the text prompt
+i.e. if you wanted a multispeaker model like canopylabs/orpheus-3b-0.1-ft, you can pass:
+f"{example["source"]}:  {example["text"]}", as is passed.
+'''
+print(tok_info)
+
+def create_input_ids(example):
+    text_ids = tokenizer.encode(f"{example['speaker_id']}: {example['text']}",  add_special_tokens=True)
+    text_ids.append(end_of_text)
+    example["text_tokens"] = text_ids
+    input_ids = (
+        [start_of_human]
+        + example["text_tokens"]
+        + [end_of_human]
+        + [start_of_ai]
+        + [start_of_speech]
+        + example["codes_list"]
+        + [end_of_speech]
+        + [end_of_ai]
+    )
+    example["input_ids"] = input_ids
+    example["labels"] = input_ids
+    example["attention_mask"] = [1] * len(input_ids)
+
+    return example
+
+ds = ds.map(create_input_ids, num_proc=num_proc, remove_columns=["text", "codes_list"])
+
+#@title Remove unnecessary columns
+columns_to_keep = ["input_ids", "labels", "attention_mask"]
+columns_to_remove = [col for col in ds.column_names if col not in columns_to_keep]
+
+ds = ds.remove_columns(columns_to_remove)
+
+ds.push_to_hub(name_to_push_dataset_to)
+```
+
+## Training
+After preprocessing is done, fill out the blanks in finetune.yml and simply run `axolotl train finetune.yml`
+
+## Inference
+For inference, please refer to the original [orpheus github](https://github.com/canopyai/Orpheus-TTS/tree/main).
diff --git a/examples/orpheus/finetune.yml b/examples/orpheus/finetune.yml
new file mode 100644
index 0000000000..86a488c846
--- /dev/null
+++ b/examples/orpheus/finetune.yml
@@ -0,0 +1,54 @@
+base_model: canopylabs/orpheus-3b-0.1-pretrained
+
+hub_model_id: <your-hub-model-id>
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_fused_linear_cross_entropy: true
+
+datasets:
+  - path: <your-hf-dataset-id>
+    type:  # leave empty to load pre-tokenized
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+sequence_len: 8192
+sample_packing: true
+
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 8
+micro_batch_size: 4
+num_epochs: 3
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 5
+saves_per_epoch: 5
+weight_decay: 0.05
+
+special_tokens:
+  pad_token: <custom_token_7>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/phi/lora-3.5.yaml b/examples/phi/lora-3.5.yaml
new file mode 100644
index 0000000000..c10014dab1
--- /dev/null
+++ b/examples/phi/lora-3.5.yaml
@@ -0,0 +1,58 @@
+base_model: microsoft/Phi-3.5-mini-instruct
+# 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
+
+load_in_8bit: true
+load_in_4bit: false
+
+chat_template: phi_3
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/lora-out
+
+sequence_len: 4096
+sample_packing: false
+
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 4
+num_epochs: 2
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bfloat16: true
+bf16: true
+fp16:
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 4
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/phi/phi-ft.yml b/examples/phi/phi-ft.yml
index b21386f707..c16b15d8a9 100644
--- a/examples/phi/phi-ft.yml
+++ b/examples/phi/phi-ft.yml
@@ -1,10 +1,9 @@
 base_model: microsoft/phi-1_5
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: AutoTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: garage-bAInd/Open-Platypus
@@ -12,11 +11,11 @@ datasets:
 
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./phi-sft-out
+output_dir: ./outputs/phi-sft-out
 
 sequence_len: 2048
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter:
 lora_model_dir:
@@ -24,7 +23,6 @@ lora_r:
 lora_alpha:
 lora_dropout:
 lora_target_linear:
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -35,37 +33,29 @@ wandb_log_model:
 gradient_accumulation_steps: 1
 micro_batch_size: 2
 num_epochs: 4
-optimizer: adamw_torch
+optimizer: adamw_torch_fused
 adam_beta2: 0.95
 adam_epsilon: 0.00001
 max_grad_norm: 1.0
 lr_scheduler: cosine
 learning_rate: 0.000003
 
-train_on_inputs: false
-group_by_length: false
 bf16: auto
-fp16:
 tf32: true
 
 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 100
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 resize_token_embeddings_to_32x: true
 special_tokens:
   pad_token: "<|endoftext|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/phi/phi-qlora.yml b/examples/phi/phi-qlora.yml
index d2b5d661c9..ac49703554 100644
--- a/examples/phi/phi-qlora.yml
+++ b/examples/phi/phi-qlora.yml
@@ -1,10 +1,12 @@
 base_model: microsoft/phi-1_5
+# 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
 
 load_in_8bit: false
 load_in_4bit: true
-strict: false
 
 datasets:
   - path: garage-bAInd/Open-Platypus
@@ -12,11 +14,11 @@ datasets:
 
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./phi-sft-out
+output_dir: ./outputs/phi-sft-out
 
 sequence_len: 2048
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter: qlora
 lora_model_dir:
@@ -24,7 +26,6 @@ lora_r: 64
 lora_alpha: 32
 lora_dropout: 0.05
 lora_target_linear: true
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -35,37 +36,29 @@ wandb_log_model:
 gradient_accumulation_steps: 1
 micro_batch_size: 2
 num_epochs: 4
-optimizer: adamw_torch
+optimizer: adamw_torch_fused
 adam_beta2: 0.95
 adam_epsilon: 0.00001
 max_grad_norm: 1.0
 lr_scheduler: cosine
 learning_rate: 0.000003
 
-train_on_inputs: false
-group_by_length: false
 bf16: auto
-fp16:
 tf32: true
 
 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 100
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 resize_token_embeddings_to_32x: true
 special_tokens:
   pad_token: "<|endoftext|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/phi/phi2-ft.yml b/examples/phi/phi2-ft.yml
index 7a2d05d018..5702cc9b8d 100644
--- a/examples/phi/phi2-ft.yml
+++ b/examples/phi/phi2-ft.yml
@@ -1,10 +1,9 @@
 base_model: microsoft/phi-2
+# optionally might have model_type or tokenizer_type
 model_type: AutoModelForCausalLM
 tokenizer_type: AutoTokenizer
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
 
 datasets:
   - path: garage-bAInd/Open-Platypus
@@ -12,11 +11,11 @@ datasets:
 
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./phi-sft-out
+output_dir: ./outputs/phi-sft-out
 
 sequence_len: 2048
 sample_packing: true
-pad_to_sequence_len: true
+
 
 adapter:
 lora_model_dir:
@@ -24,7 +23,6 @@ lora_r:
 lora_alpha:
 lora_dropout:
 lora_target_linear:
-lora_fan_in_fan_out:
 
 wandb_project:
 wandb_entity:
@@ -35,37 +33,29 @@ wandb_log_model:
 gradient_accumulation_steps: 1
 micro_batch_size: 2
 num_epochs: 4
-optimizer: adamw_torch
+optimizer: adamw_torch_fused
 adam_beta2: 0.95
 adam_epsilon: 0.00001
 max_grad_norm: 1.0
 lr_scheduler: cosine
 learning_rate: 0.000003
 
-train_on_inputs: false
-group_by_length: false
 bf16: auto
-fp16:
 tf32: true
 
 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
+attn_implementation: flash_attention_2
 
-warmup_steps: 100
+warmup_ratio: 0.1
 evals_per_epoch: 4
 saves_per_epoch: 1
-debug:
-deepspeed:
 weight_decay: 0.1
-fsdp:
-fsdp_config:
 resize_token_embeddings_to_32x: true
 special_tokens:
   pad_token: "<|endoftext|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/phi/phi3-ft-fsdp.yml b/examples/phi/phi3-ft-fsdp.yml
new file mode 100644
index 0000000000..49d3e44cbc
--- /dev/null
+++ b/examples/phi/phi3-ft-fsdp.yml
@@ -0,0 +1,75 @@
+base_model: microsoft/Phi-3-mini-4k-instruct
+# 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
+
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+
+dataset_prepared_path:
+val_set_size: 0
+output_dir: ./phi-sft-out
+
+sequence_len: 4096
+sample_packing: true
+
+trust_remote_code: true
+
+adapter:
+lora_model_dir:
+lora_r:
+lora_alpha:
+lora_dropout:
+lora_target_linear:
+
+wandb_project: phi3
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 12
+num_epochs: 2
+optimizer: adamw_torch_fused
+adam_beta2: 0.95
+adam_epsilon: 0.00001
+max_grad_norm: 1.0
+lr_scheduler: cosine
+learning_rate: 0.000003
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.1
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: Phi3DecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+resize_token_embeddings_to_32x: true
+special_tokens:
+  pad_token: "<|endoftext|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/phi/phi3-ft.yml b/examples/phi/phi3-ft.yml
new file mode 100644
index 0000000000..d36317f7bb
--- /dev/null
+++ b/examples/phi/phi3-ft.yml
@@ -0,0 +1,63 @@
+base_model: microsoft/Phi-3-mini-4k-instruct
+# optionally might have model_type or tokenizer_type
+trust_remote_code: true
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+chat_template: phi_3
+
+datasets:
+  - path: garage-bAInd/Open-Platypus
+    type: alpaca:phi
+
+dataset_prepared_path:
+val_set_size: 0.01
+output_dir: ./out
+
+sequence_len: 4096
+sample_packing: true
+
+
+adapter: lora
+lora_model_dir:
+lora_r: 64
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_linear: true
+
+gradient_accumulation_steps: 1
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_torch_fused
+adam_beta2: 0.95
+adam_epsilon: 0.00001
+max_grad_norm: 1.0
+lr_scheduler: cosine
+learning_rate: 5.0e-6
+
+bf16: auto
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: True
+early_stopping_patience: 3
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+eval_steps: 1000
+save_steps: 5000
+eval_batch_size: 2
+eval_sample_packing: false
+eval_table_size: 2
+eval_max_new_tokens: 32
+eval_causal_lm_metrics: ["perplexity"]
+do_causal_lm_eval: true
+
+warmup_ratio: 0.2
+debug: true
+weight_decay: 0.1
+resize_token_embeddings_to_32x: true
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/pixtral/lora-12b.yml b/examples/pixtral/lora-12b.yml
new file mode 100644
index 0000000000..2e36688a13
--- /dev/null
+++ b/examples/pixtral/lora-12b.yml
@@ -0,0 +1,57 @@
+base_model: mistral-community/pixtral-12b
+processor_type: AutoProcessor
+
+# 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: pixtral
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+dataset_prepared_path:
+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: 'model.language_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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+  pad_token: <pad>
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/plano/README.md b/examples/plano/README.md
new file mode 100644
index 0000000000..e3f4d14dcc
--- /dev/null
+++ b/examples/plano/README.md
@@ -0,0 +1,42 @@
+# Finetune Katanemo's Plano-Orchestrator with Axolotl
+
+[Plano-Orchestrator](https://huggingface.co/collections/katanemo/plano-orchestrator) is a family of 4B and 30B-A3B routing and orchestration models designed for multi-agent systems. It analyzes user intent and conversation context to make precise routing decisions, excelling at multi-turn context understanding, multi-intent detection, and context-dependent routing.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+    ```bash
+    axolotl train examples/plano/plano-4b-qlora.yaml
+    ```
+
+This config uses about 5.1 GiB VRAM. Let us know how it goes. Happy finetuning! 🚀
+
+### Orchestration Prompt
+
+Plano-Orchestrator uses a specific orchestration prompt format for routing/agent decisions. Please check the [official model card](https://huggingface.co/katanemo/Plano-Orchestrator-4B) for proper prompt formatting and the `ORCHESTRATION_PROMPT` template.
+
+### Tips
+
+- To use the larger [Plano-Orchestrator-30B-A3B](https://huggingface.co/katanemo/Plano-Orchestrator-30B-A3B) MoE model, simply change `base_model: katanemo/Plano-Orchestrator-30B-A3B` in the config and enable multi-GPU training if needed.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [Plano GitHub](https://github.com/katanemo/plano)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/plano/plano-4b-qlora.yaml b/examples/plano/plano-4b-qlora.yaml
new file mode 100644
index 0000000000..30e0c36ff3
--- /dev/null
+++ b/examples/plano/plano-4b-qlora.yaml
@@ -0,0 +1,65 @@
+base_model: katanemo/Plano-Orchestrator-4B
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+chat_template: qwen3
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Gemma3-12B_baseline.yml b/examples/qat_nvfp4/Gemma3-12B_baseline.yml
new file mode 100644
index 0000000000..e1c7e998ae
--- /dev/null
+++ b/examples/qat_nvfp4/Gemma3-12B_baseline.yml
@@ -0,0 +1,67 @@
+base_model: google/gemma-3-12b-it
+# 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
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: gemma3
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+
+output_dir: ./outputs/out_gemma/
+
+sequence_len: 8096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 4e-5
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Gemma3DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Gemma3-12B_qat.yml b/examples/qat_nvfp4/Gemma3-12B_qat.yml
new file mode 100644
index 0000000000..061fd6061b
--- /dev/null
+++ b/examples/qat_nvfp4/Gemma3-12B_qat.yml
@@ -0,0 +1,72 @@
+base_model: google/gemma-3-12b-it
+# 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
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: gemma3
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+
+output_dir: ./outputs/qat_out_gemma/
+
+sequence_len: 8096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+qat:
+  activation_dtype: nvfp4
+  weight_dtype: nvfp4
+  group_size: 16 # only group_size of 16 is supported with nvfp4
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 4e-5
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Gemma3DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Math-Gemma3-12B_baseline.yml b/examples/qat_nvfp4/Math-Gemma3-12B_baseline.yml
new file mode 100644
index 0000000000..94a5cf39de
--- /dev/null
+++ b/examples/qat_nvfp4/Math-Gemma3-12B_baseline.yml
@@ -0,0 +1,69 @@
+base_model: google/gemma-3-12b-it
+# Math finetuning configuration for Gemma3-12B
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: AI-MO/NuminaMath-CoT
+    type: chat_template
+
+output_dir: ./outputs/out_math_gemma/
+
+sequence_len: 4096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 8
+
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 3e-5
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Gemma3DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Math-Gemma3-12B_qat.yml b/examples/qat_nvfp4/Math-Gemma3-12B_qat.yml
new file mode 100644
index 0000000000..c29ed5f4dd
--- /dev/null
+++ b/examples/qat_nvfp4/Math-Gemma3-12B_qat.yml
@@ -0,0 +1,74 @@
+base_model: google/gemma-3-12b-it
+# Math finetuning configuration for Gemma3-12B
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: AI-MO/NuminaMath-CoT
+    type: chat_template
+
+output_dir: ./outputs/qat_out_math_gemma/
+
+sequence_len: 4096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+qat:
+  activation_dtype: nvfp4
+  weight_dtype: nvfp4
+  group_size: 16 # only group_size of 16 is supported with nvfp4
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 8
+
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 3e-5
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Gemma3DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Math-Gemma3-27B_baseline.yml b/examples/qat_nvfp4/Math-Gemma3-27B_baseline.yml
new file mode 100644
index 0000000000..93184eaa85
--- /dev/null
+++ b/examples/qat_nvfp4/Math-Gemma3-27B_baseline.yml
@@ -0,0 +1,70 @@
+base_model: google/gemma-3-27b-it
+# Math finetuning configuration for Gemma3-27B
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: AI-MO/NuminaMath-CoT
+    type: chat_template
+
+output_dir: ./outputs/out_math_gemma27/
+
+sequence_len: 4096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 5e-6
+eta_min: 7e-7
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Gemma3DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Math-Gemma3-27B_qat.yml b/examples/qat_nvfp4/Math-Gemma3-27B_qat.yml
new file mode 100644
index 0000000000..9d69e325e4
--- /dev/null
+++ b/examples/qat_nvfp4/Math-Gemma3-27B_qat.yml
@@ -0,0 +1,75 @@
+base_model: google/gemma-3-27b-it
+# Math finetuning configuration for Gemma3-27B
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: gemma3
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: AI-MO/NuminaMath-CoT
+    type: chat_template
+
+output_dir: ./outputs/qat_out_math_gemma27/
+
+sequence_len: 4096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+qat:
+  activation_dtype: nvfp4
+  weight_dtype: nvfp4
+  group_size: 16 # only group_size of 16 is supported with nvfp4
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 5e-6
+eta_min: 7e-7
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Gemma3DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Math-Qwen2.5-72B_baseline.yml b/examples/qat_nvfp4/Math-Qwen2.5-72B_baseline.yml
new file mode 100644
index 0000000000..487fc8e4e1
--- /dev/null
+++ b/examples/qat_nvfp4/Math-Qwen2.5-72B_baseline.yml
@@ -0,0 +1,67 @@
+base_model: Qwen/Qwen2.5-72B
+# Math finetuning configuration for Qwen2.5-72B (non-instruct)
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: qwen_25
+datasets:
+  - path: AI-MO/NuminaMath-CoT
+    type: chat_template
+
+output_dir: ./outputs/out_math_72b/
+
+sequence_len: 4096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 8
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 5e-6
+eta_min: 7e-7
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Math-Qwen2.5-72B_qat.yml b/examples/qat_nvfp4/Math-Qwen2.5-72B_qat.yml
new file mode 100644
index 0000000000..12812d8593
--- /dev/null
+++ b/examples/qat_nvfp4/Math-Qwen2.5-72B_qat.yml
@@ -0,0 +1,72 @@
+base_model: Qwen/Qwen2.5-72B
+# Math finetuning configuration for Qwen2.5-72B (non-instruct)
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: qwen_25
+datasets:
+  - path: AI-MO/NuminaMath-CoT
+    type: chat_template
+
+output_dir: ./outputs/qat_out_math_72b/
+
+sequence_len: 4096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+qat:
+  activation_dtype: nvfp4
+  weight_dtype: nvfp4
+  group_size: 16 # only group_size of 16 is supported with nvfp4
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 8
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 5e-6
+eta_min: 7e-7
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Qwen2.5-72B_baseline.yml b/examples/qat_nvfp4/Qwen2.5-72B_baseline.yml
new file mode 100644
index 0000000000..c52fd6b0a9
--- /dev/null
+++ b/examples/qat_nvfp4/Qwen2.5-72B_baseline.yml
@@ -0,0 +1,67 @@
+base_model: Qwen/Qwen2.5-72B
+# Alpaca finetuning configuration for Qwen2.5-72B
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: qwen_25
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+
+output_dir: ./outputs/out_qwen72b/
+
+sequence_len: 8096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qat_nvfp4/Qwen2.5-72B_qat.yml b/examples/qat_nvfp4/Qwen2.5-72B_qat.yml
new file mode 100644
index 0000000000..cc67107c0e
--- /dev/null
+++ b/examples/qat_nvfp4/Qwen2.5-72B_qat.yml
@@ -0,0 +1,72 @@
+base_model: Qwen/Qwen2.5-72B
+# Alpaca finetuning configuration for Qwen2.5-72B
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+seed: 42
+chat_template: qwen_25
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+
+output_dir: ./outputs/qat_out_qwen72b/
+
+sequence_len: 8096
+sample_packing: true
+attn_implementation: flash_attention_2
+
+qat:
+  activation_dtype: nvfp4
+  weight_dtype: nvfp4
+  group_size: 16 # only group_size of 16 is supported with nvfp4
+
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 16
+
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+# evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp_version: 2
+
+fsdp_config:
+  offload_params: false
+  cpu_ram_efficient_loading: true
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen2-vl/lora-7b.yaml b/examples/qwen2-vl/lora-7b.yaml
new file mode 100644
index 0000000000..d9bc4826b1
--- /dev/null
+++ b/examples/qwen2-vl/lora-7b.yaml
@@ -0,0 +1,56 @@
+base_model: Qwen/Qwen2-VL-7B-Instruct
+processor_type: AutoProcessor
+
+# 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: qwen2_vl
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+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: 'model.language_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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen2/adamw-pretrain-fsdp2.yaml b/examples/qwen2/adamw-pretrain-fsdp2.yaml
new file mode 100644
index 0000000000..4129338db9
--- /dev/null
+++ b/examples/qwen2/adamw-pretrain-fsdp2.yaml
@@ -0,0 +1,70 @@
+base_model: Qwen/Qwen2.5-0.5B
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+
+# Use random initialization for fair comparison
+reinit_weights: true
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+# Pretraining dataset
+pretraining_dataset:
+  - path: allenai/c4
+    name: en
+    type: pretrain
+    split: train
+
+dataset_prepared_path:
+val_set_size: 0.0
+output_dir: ./outputs/compare-adamw-pretrain
+
+sequence_len: 2048
+sample_packing: true
+pad_to_sequence_len: true
+
+wandb_project: dist_muon
+wandb_entity:
+wandb_watch:
+wandb_name: adamw
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 4
+num_epochs: 1
+max_steps: 305
+
+# AdamW optimizer settings (standard LR for AdamW)
+optimizer: adamw_torch_fused
+learning_rate: 0.0002
+weight_decay: 0.01
+lr_scheduler: cosine
+
+train_on_inputs: true
+group_by_length: false
+bf16: auto
+fp16: false
+tf32: false
+
+gradient_checkpointing: false
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_steps: 10
+evals_per_epoch: 0
+saves_per_epoch: 1
+
+# Reproducibility
+seed: 42
+
+fsdp_config:
+  fsdp_version: 2
+  fsdp_offload_params: false
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_cpu_ram_efficient_loading: false
+  fsdp_reshard_after_forward: true
+
+special_tokens:
diff --git a/examples/qwen2/dpo.yaml b/examples/qwen2/dpo.yaml
new file mode 100644
index 0000000000..6096053fdb
--- /dev/null
+++ b/examples/qwen2/dpo.yaml
@@ -0,0 +1,58 @@
+base_model: Qwen/Qwen2.5-0.5B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+chat_template: qwen_25
+rl: dpo
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_dpo_test
+    type: chat_template.default
+    field_messages: conversation
+    field_chosen: chosen
+    field_rejected: rejected
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      system:
+        - system
+      user:
+        - user
+      assistant:
+        - assistant
+
+dataset_prepared_path:
+val_set_size: 0.0
+output_dir: ./outputs/dpo-out
+
+sequence_len: 2048
+sample_packing: false
+
+
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen2/muon-pretrain-fsdp2.yaml b/examples/qwen2/muon-pretrain-fsdp2.yaml
new file mode 100644
index 0000000000..40dcff7be2
--- /dev/null
+++ b/examples/qwen2/muon-pretrain-fsdp2.yaml
@@ -0,0 +1,70 @@
+base_model: Qwen/Qwen2.5-0.5B
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+
+# Use random initialization for fair comparison
+reinit_weights: true
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+# Pretraining dataset
+pretraining_dataset:
+  - path: allenai/c4
+    name: en
+    type: pretrain
+    split: train
+
+dataset_prepared_path:
+val_set_size: 0.0
+output_dir: ./outputs/compare-muon-pretrain
+
+sequence_len: 2048
+sample_packing: true
+pad_to_sequence_len: true
+
+wandb_project: dist_muon
+wandb_entity:
+wandb_watch:
+wandb_name: muon
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 4
+num_epochs: 1
+max_steps: 305
+
+# Muon optimizer settings
+optimizer: muon
+learning_rate: 0.02
+weight_decay: 0.01
+lr_scheduler: cosine
+
+train_on_inputs: true
+group_by_length: false
+bf16: auto
+fp16: false
+tf32: false
+
+gradient_checkpointing: false
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_steps: 10
+evals_per_epoch: 0
+saves_per_epoch: 1
+
+# Reproducibility
+seed: 42
+
+fsdp_config:
+  fsdp_version: 2
+  fsdp_offload_params: false
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_cpu_ram_efficient_loading: false
+  fsdp_reshard_after_forward: true
+
+special_tokens:
diff --git a/examples/qwen2/prm.yaml b/examples/qwen2/prm.yaml
new file mode 100644
index 0000000000..1b3579fd4f
--- /dev/null
+++ b/examples/qwen2/prm.yaml
@@ -0,0 +1,59 @@
+base_model: Qwen/Qwen2.5-3B
+# optionally might have model_type or tokenizer_type
+model_type: AutoModelForTokenClassification
+num_labels: 2
+tokenizer_type: AutoTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+process_reward_model: true
+chat_template:
+datasets:
+  - path: trl-lib/math_shepherd
+    type: stepwise_supervised
+    step_separator: "\n"
+    max_completion_length:
+    train_on_last_step_only: false
+
+val_set_size: 0.2
+output_dir: ./outputs/out
+remove_unused_columns: false
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+
+gradient_accumulation_steps: 1
+micro_batch_size: 8
+eval_batch_size: 8
+num_epochs: 1
+optimizer: adamw_torch
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: true
+fp16:
+tf32:
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+eval_steps: 100
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen2/qlora-fsdp.yaml b/examples/qwen2/qlora-fsdp.yaml
new file mode 100644
index 0000000000..7bb035c3ad
--- /dev/null
+++ b/examples/qwen2/qlora-fsdp.yaml
@@ -0,0 +1,71 @@
+base_model: Qwen/Qwen2-7B
+# 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: true
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/out
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+
+adapter: qlora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 64
+lora_dropout: 0.05
+lora_target_linear: 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_torch_fused
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: Qwen2DecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen2/reward-model.yaml b/examples/qwen2/reward-model.yaml
new file mode 100644
index 0000000000..b7039cba04
--- /dev/null
+++ b/examples/qwen2/reward-model.yaml
@@ -0,0 +1,53 @@
+base_model:  Qwen/Qwen2.5-0.5B
+# optionally might have model_type or tokenizer_type
+model_type: AutoModelForSequenceClassification
+num_labels: 1
+tokenizer_type: AutoTokenizer
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+reward_model: true
+chat_template: qwen_25
+datasets:
+  - path: argilla/distilabel-intel-orca-dpo-pairs
+    type: bradley_terry.chat_template
+val_set_size: 0.0
+output_dir: ./outputs/out
+remove_unused_columns: false
+
+sequence_len: 2048
+sample_packing: false
+eval_sample_packing: false
+
+
+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: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen2_5-vl/lora-7b.yaml b/examples/qwen2_5-vl/lora-7b.yaml
new file mode 100644
index 0000000000..e78aac78b8
--- /dev/null
+++ b/examples/qwen2_5-vl/lora-7b.yaml
@@ -0,0 +1,56 @@
+base_model: Qwen/Qwen2.5-VL-7B-Instruct
+processor_type: AutoProcessor
+
+# 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: qwen2_vl
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+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: 'model.language_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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen3-next/README.md b/examples/qwen3-next/README.md
new file mode 100644
index 0000000000..f570973858
--- /dev/null
+++ b/examples/qwen3-next/README.md
@@ -0,0 +1,49 @@
+# Finetune Qwen3-Next with Axolotl
+
+[Qwen3-Next](https://huggingface.co/collections/Qwen/qwen3-next-68c25fd6838e585db8eeea9d) represents the next-generation foundation models optimized for extreme context length and large-scale parameter efficiency. The series introduces architectural innovations including Hybrid Attention (Gated DeltaNet + Gated Attention), High-Sparsity MoE with 1:50 activation ratio, and Multi-Token Prediction for enhanced performance and inference acceleration.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Install FLA for improved performance
+```bash
+uv pip uninstall causal-conv1d && uv pip install flash-linear-attention==0.4.1
+```
+
+4. Run the finetuning example:
+
+```bash
+# ~47 GiB (no target experts) and ~71GiB (target experts) VRAM.
+axolotl train examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml
+
+# NVFP4 MoE-LoRA (~85 GiB at sequence_len 2048, ~106 GiB at 16k tokens/step)
+axolotl train examples/qwen3-next/qwen3-next-80b-a3b-nvfp4-lora.yaml
+```
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- For inference, you can experiment with `temperature: 0.7`, `top_p: 0.8`, `top_k: 20`, and `min_p: 0`.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config. See [Multi-GPU](#optimization-guides) section below.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Related Resources
+
+- [Qwen3-Next Blog](https://qwenlm.github.io/blog/qwen3_next/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/qwen3-next/qwen3-next-80b-a3b-nvfp4-lora.yaml b/examples/qwen3-next/qwen3-next-80b-a3b-nvfp4-lora.yaml
new file mode 100644
index 0000000000..2d5080d4b2
--- /dev/null
+++ b/examples/qwen3-next/qwen3-next-80b-a3b-nvfp4-lora.yaml
@@ -0,0 +1,57 @@
+# Qwen3-Next-80B-A3B MoE LoRA on an NVFP4 (modelopt) checkpoint via SonicMoE.
+# The Gated DeltaNet layers need flash-linear-attention installed; with sample_packing, micro_batch_size must be 1.
+base_model: nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4
+
+plugins:
+  - axolotl.integrations.kernels.KernelsPlugin
+use_sonicmoe: true
+strict: false
+
+chat_template: tokenizer_default
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:2%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/qwen3-next-80b-nvfp4-lora
+
+sequence_len: 2048
+sample_packing: true
+attn_implementation: flash_attention_2
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 0.0002
+warmup_ratio: 0.1
+
+bf16: auto
+tf32: true
+gradient_checkpointing: true
+
+logging_steps: 1
+saves_per_epoch: 1
+
+wandb_project:
+wandb_entity:
+wandb_name:
diff --git a/examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml b/examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml
new file mode 100644
index 0000000000..e8e7e08c79
--- /dev/null
+++ b/examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml
@@ -0,0 +1,77 @@
+base_model: Qwen/Qwen3-Next-80B-A3B-Instruct
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+quantize_moe_experts: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 16
+lora_alpha: 8
+lora_dropout: 0
+lora_target_modules:
+  - linear_attn.in_proj_ba
+  - linear_attn.in_proj_qkvz
+  - linear_attn.out_proj
+  - shared_expert.up_proj
+  - shared_expert.down_proj
+  - shared_expert.gate_proj
+  - shared_expert_gate
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen3.5/122b-a10b-moe-qlora-fsdp.yaml b/examples/qwen3.5/122b-a10b-moe-qlora-fsdp.yaml
new file mode 100644
index 0000000000..47842c5616
--- /dev/null
+++ b/examples/qwen3.5/122b-a10b-moe-qlora-fsdp.yaml
@@ -0,0 +1,85 @@
+base_model: Qwen/Qwen3.5-122B-A10B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+strict: false
+
+chat_template: qwen3_5
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 2048
+sample_packing: true
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  # Add gate_up_proj and down_proj to also target shared experts (nn.Linear):
+  # - gate_up_proj
+  # - down_proj
+
+# Target routed experts (3D nn.Parameter tensors, not nn.Linear — use lora_target_parameters):
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+fsdp_config:
+  fsdp_version: 2
+  offload_params: true
+  cpu_ram_efficient_loading: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen3_5MoeDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
diff --git a/examples/qwen3.5/122b-a10b-moe-qlora.yaml b/examples/qwen3.5/122b-a10b-moe-qlora.yaml
new file mode 100644
index 0000000000..f2675c7d78
--- /dev/null
+++ b/examples/qwen3.5/122b-a10b-moe-qlora.yaml
@@ -0,0 +1,74 @@
+base_model: Qwen/Qwen3.5-122B-A10B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+strict: false
+
+chat_template: qwen3_5
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 2048
+sample_packing: true
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  # Add gate_up_proj and down_proj to also target shared experts (nn.Linear):
+  # - gate_up_proj
+  # - down_proj
+
+# Target routed experts (3D nn.Parameter tensors, not nn.Linear — use lora_target_parameters):
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/qwen3.5/27b-fft.yaml b/examples/qwen3.5/27b-fft.yaml
new file mode 100644
index 0000000000..e00c94926b
--- /dev/null
+++ b/examples/qwen3.5/27b-fft.yaml
@@ -0,0 +1,70 @@
+base_model: Qwen/Qwen3.5-27B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+# Full fine-tune (FFT) of the text-only path of Qwen3.5-27B.
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+strict: false
+
+chat_template: qwen3_5
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 2048
+sample_packing: true
+
+# Freeze vision encoder
+unfrozen_parameters:
+  - model.language_model.*
+  - lm_head.*
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_fused
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+fsdp_config:
+  fsdp_version: 2
+  offload_params: false
+  cpu_ram_efficient_loading: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen3_5DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
diff --git a/examples/qwen3.5/27b-qlora-fsdp.yaml b/examples/qwen3.5/27b-qlora-fsdp.yaml
new file mode 100644
index 0000000000..7a5423c773
--- /dev/null
+++ b/examples/qwen3.5/27b-qlora-fsdp.yaml
@@ -0,0 +1,81 @@
+base_model: Qwen/Qwen3.5-27B
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+strict: false
+
+chat_template: qwen3_5
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 2048
+sample_packing: true
+
+load_in_4bit: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - down_proj
+  - up_proj
+  # Uncomment below to also target the linear attention projections.
+  # These use separate in_proj_qkv / in_proj_z / out_proj (Qwen3.5-specific).
+  # - linear_attn.in_proj_qkv
+  # - linear_attn.in_proj_z
+  # - linear_attn.out_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+fsdp_config:
+  fsdp_version: 2
+  offload_params: false
+  cpu_ram_efficient_loading: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen3_5DecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
diff --git a/examples/qwen3.5/27b-qlora.yaml b/examples/qwen3.5/27b-qlora.yaml
new file mode 100644
index 0000000000..2401a4865f
--- /dev/null
+++ b/examples/qwen3.5/27b-qlora.yaml
@@ -0,0 +1,70 @@
+base_model: Qwen/Qwen3.5-27B
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+strict: false
+
+chat_template: qwen3_5
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 2048
+sample_packing: true
+
+load_in_4bit: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - down_proj
+  - up_proj
+  # Uncomment below to also target the linear attention projections.
+  # These use separate in_proj_qkv / in_proj_z / out_proj (Qwen3.5-specific).
+  # - linear_attn.in_proj_qkv
+  # - linear_attn.in_proj_z
+  # - linear_attn.out_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/qwen3.5/35b-a3b-moe-qlora-fsdp.yaml b/examples/qwen3.5/35b-a3b-moe-qlora-fsdp.yaml
new file mode 100644
index 0000000000..2fb7f15f89
--- /dev/null
+++ b/examples/qwen3.5/35b-a3b-moe-qlora-fsdp.yaml
@@ -0,0 +1,85 @@
+base_model: Qwen/Qwen3.5-35B-A3B
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+strict: false
+
+chat_template: qwen3_5
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 2048
+sample_packing: true
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  # Add gate_up_proj and down_proj to also target shared experts (nn.Linear):
+  # - gate_up_proj
+  # - down_proj
+
+# Target routed experts (3D nn.Parameter tensors, not nn.Linear — use lora_target_parameters):
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+lora_mlp_kernel: false
+lora_qkv_kernel: false
+lora_o_kernel: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+fsdp_config:
+  fsdp_version: 2
+  offload_params: true
+  cpu_ram_efficient_loading: false
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen3_5MoeDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+  reshard_after_forward: true
+  activation_checkpointing: true
diff --git a/examples/qwen3.5/35b-a3b-moe-qlora.yaml b/examples/qwen3.5/35b-a3b-moe-qlora.yaml
new file mode 100644
index 0000000000..a6afc1aa2c
--- /dev/null
+++ b/examples/qwen3.5/35b-a3b-moe-qlora.yaml
@@ -0,0 +1,84 @@
+base_model: Qwen/Qwen3.5-35B-A3B-Base
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+  - axolotl.integrations.kernels.KernelsPlugin
+  - axolotl.integrations.liger.LigerPlugin
+use_kernels: true
+use_scattermoe: true
+liger_layer_norm: true
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_rms_norm_gated: true
+
+torch_compile: false
+
+chat_template: qwen3_5
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 2048
+sample_packing: true
+
+load_in_4bit: true
+quantize_moe_experts: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  # Add gate_up_proj and down_proj to also target shared experts (nn.Linear):
+  # - gate_up_proj
+  # - down_proj
+
+# Target routed experts (3D nn.Parameter tensors, not nn.Linear — use lora_target_parameters):
+# lora_target_parameters:
+#   - mlp.experts.gate_up_proj
+#   - mlp.experts.down_proj
+
+lora_qkv_kernel: true
+lora_o_kernel: true
+lora_mlp_kernel: false
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 4
+num_epochs: 1
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+activation_offloading: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/qwen3.5/35b-a3b-moe-vision-lora.yaml b/examples/qwen3.5/35b-a3b-moe-vision-lora.yaml
new file mode 100644
index 0000000000..7cfad32902
--- /dev/null
+++ b/examples/qwen3.5/35b-a3b-moe-vision-lora.yaml
@@ -0,0 +1,62 @@
+# Qwen 3.5 35B-A3B MoE Vision LoRA
+#
+# Vision fine-tuning of the hybrid DeltaNet + Attention MoE model.
+# 256 experts, 8 active per token, with early-fusion vision support.
+
+base_model: Qwen/Qwen3.5-35B-A3B
+processor_type: AutoProcessor
+
+# Required for vision/multimodal training
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: qwen3_5
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:100]
+
+val_set_size: 0
+output_dir: ./outputs/qwen35-35b-a3b-vision-lora
+
+adapter: lora
+sequence_len: 4096
+pad_to_sequence_len: false
+
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - down_proj
+  - up_proj
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+max_steps: 10
+optimizer: adamw_torch_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
diff --git a/examples/qwen3.5/9b-fft-vision.yaml b/examples/qwen3.5/9b-fft-vision.yaml
new file mode 100644
index 0000000000..e8427b8844
--- /dev/null
+++ b/examples/qwen3.5/9b-fft-vision.yaml
@@ -0,0 +1,49 @@
+base_model: Qwen/Qwen3.5-9B
+processor_type: AutoProcessor
+
+# Required for multimodal training
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: qwen3_5
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+sequence_len: 4096
+pad_to_sequence_len: false
+
+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: auto
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/examples/qwen3.5/9b-lora-vision.yaml b/examples/qwen3.5/9b-lora-vision.yaml
new file mode 100644
index 0000000000..9c2b9397e5
--- /dev/null
+++ b/examples/qwen3.5/9b-lora-vision.yaml
@@ -0,0 +1,66 @@
+base_model: Qwen/Qwen3.5-9B
+processor_type: AutoProcessor
+
+# These 3 lines are required for vision/multimodal training
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: qwen3_5
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+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
+# Targets the language model attention and MLP layers.
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - down_proj
+  - up_proj
+  # Uncomment to also target the linear attention (GatedDeltaNet) projections:
+  # - linear_attn.in_proj_qkv
+  # - linear_attn.in_proj_z
+  # - linear_attn.out_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
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/qwen3.5/README.md b/examples/qwen3.5/README.md
new file mode 100644
index 0000000000..22e9e360dc
--- /dev/null
+++ b/examples/qwen3.5/README.md
@@ -0,0 +1,95 @@
+# Finetune Qwen3.5 with Axolotl
+
+[Qwen3.5](https://huggingface.co/collections/Qwen/qwen35) is a hybrid architecture model series combining Gated DeltaNet linear attention with standard Transformer attention. All Qwen3.5 models are early-fusion vision-language models: dense variants use `Qwen3_5ForConditionalGeneration` and MoE variants use `Qwen3_5MoeForConditionalGeneration`.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Install FLA for sample packing support with the Gated DeltaNet linear attention layers:
+  ```bash
+  uv pip uninstall causal-conv1d && uv pip install flash-linear-attention==0.4.1
+  ```
+  > FLA is required when `sample_packing: true`. Without it, training raises a `RuntimeError` on packed sequences. Vision configs use `sample_packing: false` so FLA is optional there.
+
+4. Pick any config from the table below and run:
+
+    ```bash
+    axolotl train examples/qwen3.5/<config>.yaml
+    ```
+
+Available configs:
+
+| Config | Model | Type | Peak VRAM |
+|---|---|---|---|
+| `9b-lora-vision.yaml` | Qwen3.5-9B | Vision+text LoRA, single GPU | — |
+| `9b-fft-vision.yaml` | Qwen3.5-9B | Vision+text FFT, single GPU | ~61 GiB |
+| `27b-qlora.yaml` | Qwen3.5-27B | Dense, text-only QLoRA | ~47 GiB |
+| `27b-fft.yaml` | Qwen3.5-27B | Dense, text-only FFT (vision frozen) | ~53 GiB |
+| `27b-qlora-fsdp.yaml` | Qwen3.5-27B | Dense, text-only QLoRA + FSDP2 | — |
+| `35b-a3b-moe-qlora.yaml` | Qwen3.5-35B-A3B | MoE, text-only QLoRA | — |
+| `35b-a3b-moe-qlora-fsdp.yaml` | Qwen3.5-35B-A3B | MoE, text-only QLoRA + FSDP2 | — |
+| `122b-a10b-moe-qlora.yaml` | Qwen3.5-122B-A10B | MoE, text-only QLoRA | — |
+| `122b-a10b-moe-qlora-fsdp.yaml` | Qwen3.5-122B-A10B | MoE, text-only QLoRA + FSDP2 | — |
+
+### Gated DeltaNet Linear Attention
+
+Qwen3.5 interleaves standard attention with Gated DeltaNet linear attention layers. To apply LoRA to them, add to `lora_target_modules`:
+
+```yaml
+lora_target_modules:
+  # ... standard projections ...
+  - linear_attn.in_proj_qkv
+  - linear_attn.in_proj_z
+  - linear_attn.out_proj
+```
+
+### Routed Experts (MoE)
+
+To apply LoRA to routed expert parameters, add `lora_target_parameters`:
+
+```yaml
+lora_target_parameters:
+  - mlp.experts.gate_up_proj
+  - mlp.experts.down_proj
+#  - mlp.gate.weight  # router
+```
+
+### Shared Experts (MoE)
+
+Shared experts use `nn.Linear` (unlike routed experts which are 3D `nn.Parameter` tensors), so they can be targeted via `lora_target_modules`. To also train shared expert projections alongside attention, uncomment `gate_up_proj` and `down_proj` in `lora_target_modules`:
+
+```yaml
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  # Add gate_up_proj and down_proj to also target shared experts (nn.Linear):
+  # - gate_up_proj
+  # - down_proj
+```
+
+Use `lora_target_parameters` (see [Routed Experts](#routed-experts-moe) above) to target routed experts separately.
+
+### TIPS
+
+- For inference hyp, please see the respective model card details.
+- You can run a full finetuning of smaller configs by removing `adapter: qlora` and `load_in_4bit: true`. See [Multi-GPU](#optimization-guides) below.
+- Read more on loading your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+- For **multimodal** finetuning, set `processor_type: AutoProcessor`, `skip_prepare_dataset: true`, and `remove_unused_columns: false` as shown in `9b-lora-vision.yaml`.
+
+## Optimization Guides
+
+- [Optimizations Guide](https://docs.axolotl.ai/docs/optimizations.html)
+
+## Related Resources
+
+- [Qwen3.5 Blog](https://qwenlm.github.io/blog/qwen3.5/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/qwen3/30b-a3b-nvfp4-lora.yaml b/examples/qwen3/30b-a3b-nvfp4-lora.yaml
new file mode 100644
index 0000000000..628f5d1059
--- /dev/null
+++ b/examples/qwen3/30b-a3b-nvfp4-lora.yaml
@@ -0,0 +1,56 @@
+# Qwen3-30B-A3B MoE LoRA on an NVFP4 (modelopt) checkpoint via SonicMoE.
+base_model: nvidia/Qwen3-30B-A3B-NVFP4
+
+plugins:
+  - axolotl.integrations.kernels.KernelsPlugin
+use_sonicmoe: true
+strict: false
+
+chat_template: tokenizer_default
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:2%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+dataset_prepared_path: last_run_prepared
+val_set_size: 0
+output_dir: ./outputs/qwen3-30b-a3b-nvfp4-lora
+
+sequence_len: 2048
+sample_packing: true
+attn_implementation: flash_attention_2
+
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+lora_target_parameters:
+  - experts.gate_up_proj
+  - experts.down_proj
+
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 0.0002
+warmup_ratio: 0.1
+
+bf16: auto
+tf32: true
+gradient_checkpointing: true
+
+logging_steps: 1
+saves_per_epoch: 1
+
+wandb_project:
+wandb_entity:
+wandb_name:
diff --git a/examples/qwen3/32b-qlora.yaml b/examples/qwen3/32b-qlora.yaml
new file mode 100644
index 0000000000..dd5dd696ef
--- /dev/null
+++ b/examples/qwen3/32b-qlora.yaml
@@ -0,0 +1,71 @@
+base_model: Qwen/Qwen3-32B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+strict: false
+
+chat_template: qwen3
+datasets:
+  - path: mlabonne/FineTome-100k
+    type: chat_template
+    split: train[:20%]
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+val_set_size: 0.0
+output_dir: ./outputs/out
+dataset_prepared_path: last_run_prepared
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+
+load_in_4bit: true
+adapter: qlora
+lora_r: 16
+lora_alpha: 32
+lora_target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - down_proj
+  - up_proj
+lora_mlp_kernel: true
+lora_qkv_kernel: true
+lora_o_kernel: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 2
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_4bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: offload
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen3/8b-lora-fused-attn.yaml b/examples/qwen3/8b-lora-fused-attn.yaml
new file mode 100644
index 0000000000..70d8c36fa1
--- /dev/null
+++ b/examples/qwen3/8b-lora-fused-attn.yaml
@@ -0,0 +1,44 @@
+base_model: Qwen/Qwen3-8B
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/qwen3-8b-lora-fused
+
+sequence_len: 4096
+sample_packing: true
+eval_sample_packing: true
+
+adapter: lora
+lora_r: 32
+lora_alpha: 64
+lora_dropout: 0.0
+lora_target_linear: true
+
+# Opt-in fused RMSNorm + RoPE Triton kernel for Qwen3 attention.
+# fp32-internal, single round: matches an fp32 reference within bf16 rounding
+# (more accurate than the eager bf16 path). Speeds up the q/k norm+rope path.
+fused_attn_kernel: true
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+flash_attention: true
+
+warmup_ratio: 0.1
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/qwen3/8b-qat-fsdp2.yml b/examples/qwen3/8b-qat-fsdp2.yml
new file mode 100644
index 0000000000..3c9607a9ab
--- /dev/null
+++ b/examples/qwen3/8b-qat-fsdp2.yml
@@ -0,0 +1,80 @@
+base_model: Qwen/Qwen3-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
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+
+output_dir: ./outputs/qat_out/
+
+sequence_len: 2048
+sample_packing: true
+attn_implementation: flex_attention
+
+
+flex_attn_compile_kwargs:
+  dynamic: false
+  mode: max-autotune-no-cudagraphs
+
+qat:
+  activation_dtype: int8
+  weight_dtype: int4
+  group_size: 256
+  fake_quant_after_n_steps: 1000
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 2
+max_steps: 2000
+optimizer: adamw_torch_fused
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+bf16: true
+tf32: true
+
+resume_from_checkpoint:
+logging_steps: 1
+
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+warmup_ratio: 0.1
+weight_decay: 0.0
+fsdp:
+  - full_shard
+  - auto_wrap
+
+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: Qwen3DecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_reshard_after_forward: true
+  fsdp_activation_checkpointing: true
+
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen3/README.md b/examples/qwen3/README.md
new file mode 100644
index 0000000000..42b9a66302
--- /dev/null
+++ b/examples/qwen3/README.md
@@ -0,0 +1,49 @@
+# Finetune Qwen3 with Axolotl
+
+[Qwen3](https://huggingface.co/collections/Qwen/qwen3) are a family of open source models trained by Alibaba.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+    ```bash
+    axolotl train examples/qwen3/32b-qlora.yaml
+
+    # NVFP4 MoE-LoRA (~30 GiB at sequence_len 2048, ~58 GiB at 16k tokens/step)
+    axolotl train examples/qwen3/30b-a3b-nvfp4-lora.yaml
+    ```
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### Chat template masking a few tokens off
+
+If you notice that the `chat_template` masking for assistant prompts are off by a few tokens, please ensure that you are adding the below to the yaml.
+
+```yaml
+chat_template: qwen3
+```
+
+### TIPS
+
+- For inference, please check the official model card as it depends on your reasoning mode.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [Qwen3 Blog](https://qwenlm.github.io/blog/qwen3/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/qwen3/qlora-fsdp.yaml b/examples/qwen3/qlora-fsdp.yaml
new file mode 100644
index 0000000000..a3852d457f
--- /dev/null
+++ b/examples/qwen3/qlora-fsdp.yaml
@@ -0,0 +1,70 @@
+base_model: Qwen/Qwen3-8B
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: true
+strict: false
+
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/out
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+
+adapter: qlora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 64
+lora_dropout: 0.05
+lora_target_linear: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_torch_fused
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 4
+saves_per_epoch: 1
+weight_decay: 0.0
+fsdp:
+  - full_shard
+  - auto_wrap
+fsdp_config:
+  fsdp_limit_all_gathers: true
+  fsdp_sync_module_states: true
+  fsdp_offload_params: true
+  fsdp_use_orig_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: Qwen3DecoderLayer
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sharding_strategy: FULL_SHARD
+special_tokens:
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/qwen3/reward-model.yaml b/examples/qwen3/reward-model.yaml
new file mode 100644
index 0000000000..43c62ecc4c
--- /dev/null
+++ b/examples/qwen3/reward-model.yaml
@@ -0,0 +1,44 @@
+base_model: Skywork/Skywork-Reward-V2-Qwen3-8B
+model_type: AutoModelForSequenceClassification
+num_labels: 1
+
+reward_model: true
+center_rewards_coefficient: 0.01  # Incentivize mean-zero rewards for improved stability
+chat_template: qwen3
+datasets:
+  - path: argilla/distilabel-intel-orca-dpo-pairs
+    type: bradley_terry.chat_template
+
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+sequence_len: 8192
+sample_packing: false
+eval_sample_packing: false
+pad_to_sequence_len: true
+
+deepspeed: deepspeed_configs/zero1.json
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+eval_batch_size: 1
+num_epochs: 3
+optimizer: adamw_bnb_8bit
+lr_scheduler: linear
+learning_rate: 0.00002
+
+bf16: true
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+warmup_ratio: 0.1
+logging_steps: 1
+weight_decay: 0.01
diff --git a/examples/seed-oss/README.md b/examples/seed-oss/README.md
new file mode 100644
index 0000000000..d5d4baa895
--- /dev/null
+++ b/examples/seed-oss/README.md
@@ -0,0 +1,47 @@
+# Finetune ByteDance's Seed-OSS with Axolotl
+
+[Seed-OSS](https://huggingface.co/collections/ByteDance-Seed/seed-oss-68a609f4201e788db05b5dcd) are a series of 36B parameter open source models trained by ByteDance's Seed Team.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1.  Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+    ```bash
+    # Ensure you have a compatible version of Pytorch installed
+    uv pip install --no-build-isolation 'axolotl>=0.16.1'
+
+    # Install Cut Cross Entropy
+    python scripts/cutcrossentropy_install.py | sh
+    ```
+
+2. Run the finetuning example:
+
+```bash
+axolotl train examples/seed-oss/seed-oss-36b-qlora.yaml
+```
+
+This config uses about 27.7 GiB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- For inference, the official Seed Team recommends `top_p=0.95` and `temperature=1.1`.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [ByteDance Seed Website](https://seed.bytedance.com/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/seed-oss/seed-oss-36b-qlora.yaml b/examples/seed-oss/seed-oss-36b-qlora.yaml
new file mode 100644
index 0000000000..a8423f8512
--- /dev/null
+++ b/examples/seed-oss/seed-oss-36b-qlora.yaml
@@ -0,0 +1,56 @@
+base_model: ByteDance-Seed/Seed-OSS-36B-Instruct
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/slurm/README.md b/examples/slurm/README.md
new file mode 100644
index 0000000000..4c116b7132
--- /dev/null
+++ b/examples/slurm/README.md
@@ -0,0 +1,66 @@
+# SLURM Multi-Node Training
+
+This directory contains an example SLURM script for running Axolotl training jobs across multiple nodes in a SLURM cluster.
+
+## Prerequisites
+
+- Access to a SLURM cluster with GPU nodes
+- Axolotl installed on all nodes (see [installation docs](https://docs.axolotl.ai/docs/installation.html))
+
+## Usage
+
+### Standard SLURM Clusters
+
+1. Copy [`axolotl.slurm`](./axolotl.slurm) to your working directory.
+2. Place your Axolotl config file (`train.yaml`) in the same directory.
+3. Set the appropriate environment variables for the job:
+    ```bash
+    export HF_TOKEN="your-huggingface-token"
+
+    # metric tracking
+    # export WANDB_API_KEY="your-wandb-api-key"
+    # ...
+    ```
+4. Submit the job:
+   ```bash
+   sbatch --export=ALL,NUM_NODES=2,NUM_TRAINERS=8,PRIMARY_ADDR=<master-node>,PRIMARY_PORT=29400 axolotl.slurm
+   ```
+
+   Where:
+   - `NUM_NODES`: Number of nodes to use
+   - `NUM_TRAINERS`: GPUs per node (typically 8)
+   - `PRIMARY_ADDR`: Hostname/IP of the master node
+   - `PRIMARY_PORT`: Port for distributed training (default: 29400)
+
+5. (Optional) Run other slurm commands:
+    ```bash
+    # check job info
+    scontrol show job axolotl-cli
+
+    # check job queue
+    squeue
+
+    # check cluster status
+    sinfo
+    ```
+
+### RunPod Instant Clusters
+
+Axolotl works with RunPod Instant Clusters. This feature provides managed SLURM clusters with zero configuration.
+
+1. **Deploy a SLURM Cluster**:
+   - Go to [RunPod Instant Clusters](https://console.runpod.io/cluster)
+   - Click "Create a Cluster"
+   - Choose your GPU type, node count, and region
+   - Choose an [Axolotl cloud docker image](https://docs.axolotl.ai/docs/docker.html#cloud)
+   - Deploy the cluster
+
+2. **Connect to the Controller Node**: Find the controller node in the RunPod console and connect via SSH
+
+3. **Follow the instructions in [Standard SLURM Clusters](#standard-slurm-clusters)**
+
+## Additional Resources
+
+- [Axolotl Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [SLURM Documentation](https://slurm.schedmd.com/documentation.html)
+- [RunPod SLURM Clusters Guide](https://docs.runpod.io/instant-clusters/slurm-clusters)
diff --git a/examples/slurm/axolotl.slurm b/examples/slurm/axolotl.slurm
new file mode 100644
index 0000000000..741d68ced2
--- /dev/null
+++ b/examples/slurm/axolotl.slurm
@@ -0,0 +1,20 @@
+#!/bin/bash
+# Prior to running this script, export your HF_TOKEN and WANDB_API_KEY to your environment; i.e.
+# export HF_TOKEN="..."
+# export WANDB_API_KEY="..."
+#
+
+# ---------- SBATCH commands ---------- #
+#SBATCH --job-name=axolotl-slurm-multinode
+#SBATCH --ntasks-per-node=1
+#SBATCH --nodes=$NUM_NODES
+#SBATCH --gpus-per-task=8
+#SBATCH --cpus-per-task=128
+
+export TORCH_DIST_INIT_BARRIER=0
+
+srun axolotl preprocess train.yaml
+
+srun axolotl train train.yaml --launcher torchrun -- \
+    --nproc_per_node=$NUM_TRAINERS --nnodes=$NUM_NODES \
+    --rdzv_id axolotl-cli --rdzv_backend c10d --rdzv_endpoint "${PRIMARY_ADDR}:${PRIMARY_PORT}" --rdzv-conf="join_timeout=1800"
diff --git a/examples/smolvlm2/README.md b/examples/smolvlm2/README.md
new file mode 100644
index 0000000000..01ee7fa623
--- /dev/null
+++ b/examples/smolvlm2/README.md
@@ -0,0 +1,46 @@
+# Finetune SmolVLM2 with Axolotl
+
+[SmolVLM2](https://huggingface.co/collections/HuggingFaceTB/smolvlm2-smallest-video-lm-ever-67ab6b5e84bf8aaa60cb17c7) are a family of lightweight, open-source multimodal models from HuggingFace designed to analyze and understand video, image, and text content.
+
+These models are built for efficiency, making them well-suited for on-device applications where computational resources are limited. Models are available in multiple sizes, including 2.2B, 500M, and 256M.
+
+This guide shows how to fine-tune SmolVLM2 models with Axolotl.
+
+## Getting Started
+
+1.  Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+    ```bash
+    # Ensure you have a compatible version of Pytorch installed
+    uv pip install --no-build-isolation 'axolotl>=0.16.1'
+    ```
+
+2. Install an extra dependency:
+
+    ```bash
+    uv pip install num2words==0.5.14
+    ```
+
+3.  Run the finetuning example:
+
+    ```bash
+    # LoRA SFT (1x48GB @ 6.8GiB)
+    axolotl train examples/smolvlm2/smolvlm2-2B-lora.yaml
+    ```
+
+## TIPS
+
+- **Dataset Format**: For video finetuning, your dataset must be compatible with the multi-content Messages format. For more details, see our documentation on [Multimodal Formats](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+- **Dataset Loading**: Read more on how to prepare and load your own datasets in our [documentation](https://docs.axolotl.ai/docs/dataset_loading.html).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [SmolVLM2 Blog](https://huggingface.co/blog/smolvlm2)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/smolvlm2/smolvlm2-2B-lora.yaml b/examples/smolvlm2/smolvlm2-2B-lora.yaml
new file mode 100644
index 0000000000..4cd8d5b0d2
--- /dev/null
+++ b/examples/smolvlm2/smolvlm2-2B-lora.yaml
@@ -0,0 +1,55 @@
+base_model: HuggingFaceTB/SmolVLM2-2.2B-Instruct
+trust_remote_code: true
+processor_type: AutoProcessor
+
+# 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
+
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+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: 'model.text_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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/streaming/README.md b/examples/streaming/README.md
new file mode 100644
index 0000000000..cdbb5baea1
--- /dev/null
+++ b/examples/streaming/README.md
@@ -0,0 +1,50 @@
+# Streaming Dataset Examples
+
+This directory contains example configurations for using Axolotl's streaming dataset
+functionality, which enables memory-efficient training with large datasets.
+
+## Examples
+
+Run the following examples with e.g. `axolotl train examples/streaming/sft.yaml`; no
+`axolotl preprocess` required!
+
+### Pretraining (`pretrain.yaml`)
+
+Demonstrates streaming configuration for pretraining tasks using the fineweb-edu dataset
+with SmolLM2-135M.
+
+- Uses `pretraining_dataset` configuration for automatic streaming
+- Multipack attention control to prevent cross-attention between packed sequences
+- Buffer size configuration for memory management
+
+### SFT (`sft.yaml`)
+
+Shows how to use streaming for supervised fine-tuning with the Alpaca dataset.
+
+- Explicit `streaming: true` flag for SFT datasets
+- Memory-efficient training on instruction datasets
+- Evaluation datasets are currently not streamed
+
+## Key Configuration Options
+
+### `streaming`
+- Enables streaming mode for standard datasets
+- Automatically enabled for `pretraining_dataset`
+
+### `streaming_multipack_buffer_size`
+- Controls buffer size for sample packing (default: 10,000)
+- Larger values improve packing efficiency but use more memory
+- Adjust based on available memory
+
+### `shuffle_merged_datasets`
+- Enables shuffling of streaming datasets
+- Requires additional memory for shuffle buffer
+
+### `sample_packing`
+- Packs multiple samples into single sequences
+- Minimize per-step padding tokens
+
+## Performance Tips
+
+- Download small / frequently-used datasets locally for better performance
+- Larger buffer sizes improve packing efficiency
diff --git a/examples/streaming/pretrain.yaml b/examples/streaming/pretrain.yaml
new file mode 100644
index 0000000000..a0d8b17c04
--- /dev/null
+++ b/examples/streaming/pretrain.yaml
@@ -0,0 +1,57 @@
+base_model: HuggingFaceTB/SmolLM2-135M
+
+# Streaming pretraining configuration
+pretraining_dataset:
+  - path: HuggingFaceFW/fineweb-edu
+    name: sample-10BT
+    type: pretrain
+    text_column: text
+    split: train
+
+# Streaming-specific settings
+streaming_multipack_buffer_size: 10000
+shuffle_merged_datasets: true
+
+# Training configuration
+max_steps: 1000
+output_dir: ./outputs/smollm2-135m-pretrain-streaming
+
+# Sequence and packing settings
+sequence_len: 1024
+sample_packing: true
+pretrain_multipack_attn: true  # Prevent cross-attention between packed sequences
+attn_implementation: flash_attention_2
+
+# Batch size settings
+gradient_accumulation_steps: 8
+micro_batch_size: 1
+
+# Optimizer and scheduler
+optimizer: adamw_torch
+lr_scheduler: cosine
+learning_rate: 5e-4
+warmup_ratio: 0.1
+weight_decay: 0.01
+
+# Precision and performance
+bf16: auto
+tf32: true
+
+# Logging and checkpointing
+logging_steps: 10
+save_strategy: steps
+save_steps: 250
+save_total_limit: 3
+
+# Weights & Biases (optional)
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+# Special tokens
+special_tokens:
+  pad_token: "<|endoftext|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/streaming/sft.yaml b/examples/streaming/sft.yaml
new file mode 100644
index 0000000000..4a43c34ebb
--- /dev/null
+++ b/examples/streaming/sft.yaml
@@ -0,0 +1,55 @@
+base_model: HuggingFaceTB/SmolLM2-135M
+
+# Dataset configuration
+datasets:
+  - path: tatsu-lab/alpaca
+    type: alpaca
+    split: train
+
+# Streaming-specific settings
+streaming: true
+streaming_multipack_buffer_size: 10000
+shuffle_merged_datasets: true
+
+# Training configuration
+max_steps: 1000
+output_dir: ./outputs/smollm2-135m-sft-streaming
+
+# Sequence and packing settings
+sequence_len: 1024
+sample_packing: true
+attn_implementation: flash_attention_2
+
+# Batch size settings
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+
+# Optimizer and scheduler
+optimizer: adamw_torch
+lr_scheduler: cosine
+learning_rate: 2e-4
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+# Precision and performance
+bf16: auto
+tf32: true
+
+# Logging and checkpointing
+logging_steps: 10
+save_strategy: steps
+save_steps: 100
+save_total_limit: 3
+
+# Weights & Biases (optional)
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+# Special tokens
+special_tokens:
+  pad_token: "<|endoftext|>"
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/swanlab/README.md b/examples/swanlab/README.md
new file mode 100644
index 0000000000..c793ff2c28
--- /dev/null
+++ b/examples/swanlab/README.md
@@ -0,0 +1,285 @@
+# SwanLab Integration Examples
+
+This directory contains example configurations demonstrating SwanLab integration with Axolotl.
+
+## Examples Overview
+
+### 1. DPO with Completion Logging
+**File**: `dpo-swanlab-completions.yml`
+
+Demonstrates DPO (Direct Preference Optimization) training with RLHF completion table logging.
+
+**Features**:
+- Basic SwanLab experiment tracking
+- Completion table logging (prompts, chosen/rejected responses, rewards)
+- Memory-bounded buffer for long training runs
+- Cloud sync configuration
+
+**Best for**: RLHF practitioners who want to analyze model outputs qualitatively
+
+**Quick start**:
+```bash
+export SWANLAB_API_KEY=your-api-key
+accelerate launch -m axolotl.cli.train examples/swanlab/dpo-swanlab-completions.yml
+```
+
+---
+
+### 2. LoRA with Performance Profiling
+**File**: `lora-swanlab-profiling.yml`
+
+Demonstrates standard LoRA fine-tuning with performance profiling enabled.
+
+**Features**:
+- SwanLab experiment tracking
+- Automatic profiling of trainer methods
+- Profiling metrics visualization
+- Performance optimization guidance
+
+**Best for**: Engineers optimizing training performance and comparing different configurations
+
+**Quick start**:
+```bash
+export SWANLAB_API_KEY=your-api-key
+accelerate launch -m axolotl.cli.train examples/swanlab/lora-swanlab-profiling.yml
+```
+
+---
+
+### 3. Full-Featured DPO Production Setup
+**File**: `dpo-swanlab-full-featured.yml`
+
+Comprehensive production-ready configuration with ALL SwanLab features enabled.
+
+**Features**:
+- Experiment tracking with team workspace
+- RLHF completion logging
+- Performance profiling
+- Lark (Feishu) team notifications
+- Private deployment support
+- Production checklist and troubleshooting
+
+**Best for**: Production RLHF training with team collaboration
+
+**Quick start**:
+```bash
+export SWANLAB_API_KEY=your-api-key
+export SWANLAB_LARK_WEBHOOK_URL=https://open.feishu.cn/...
+export SWANLAB_LARK_SECRET=your-webhook-secret
+accelerate launch -m axolotl.cli.train examples/swanlab/dpo-swanlab-full-featured.yml
+```
+
+---
+
+### 4. Custom Trainer Profiling (Python)
+**File**: `custom_trainer_profiling.py`
+
+Python code examples showing how to add SwanLab profiling to custom trainers.
+
+**Features**:
+- `@swanlab_profile` decorator examples
+- Context manager profiling for fine-grained timing
+- `ProfilingConfig` for advanced filtering and throttling
+- Multiple profiling patterns and best practices
+
+**Best for**: Advanced users creating custom trainers
+
+**Usage**:
+```python
+from custom_trainer_profiling import CustomTrainerWithProfiling
+# See file for detailed examples and patterns
+```
+
+---
+
+## Feature Matrix
+
+| Example | Tracking | Completion Logging | Profiling | Lark Notifications | Team Workspace |
+|---------|----------|-------------------|-----------|-------------------|----------------|
+| dpo-swanlab-completions.yml | ✅ | ✅ | ✅ (auto) | ➖ (commented) | ➖ (commented) |
+| lora-swanlab-profiling.yml | ✅ | ➖ (disabled) | ✅ (auto) | ➖ (commented) | ➖ (commented) |
+| dpo-swanlab-full-featured.yml | ✅ | ✅ | ✅ (auto) | ✅ | ✅ |
+| custom_trainer_profiling.py | N/A | N/A | ✅ (manual) | N/A | N/A |
+
+---
+
+## Configuration Quick Reference
+
+### Basic SwanLab Setup
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: my-project
+swanlab_experiment_name: my-experiment
+swanlab_mode: cloud  # cloud, local, offline, disabled
+```
+
+### RLHF Completion Logging
+```yaml
+swanlab_log_completions: true
+swanlab_completion_log_interval: 100  # Log every 100 steps
+swanlab_completion_max_buffer: 128    # Memory-bounded buffer
+```
+
+### Lark Team Notifications
+```yaml
+swanlab_lark_webhook_url: https://open.feishu.cn/...
+swanlab_lark_secret: your-webhook-secret  # Required for production
+```
+
+### Team Workspace
+```yaml
+swanlab_workspace: my-research-team
+```
+
+### Private Deployment
+```yaml
+swanlab_web_host: https://swanlab.yourcompany.com
+swanlab_api_host: https://api.swanlab.yourcompany.com
+```
+
+---
+
+## Authentication
+
+### Recommended: Environment Variable
+```bash
+export SWANLAB_API_KEY=your-api-key
+export SWANLAB_LARK_WEBHOOK_URL=https://open.feishu.cn/...
+export SWANLAB_LARK_SECRET=your-webhook-secret
+```
+
+### Alternative: Config File (less secure)
+```yaml
+swanlab_api_key: your-api-key
+swanlab_lark_webhook_url: https://open.feishu.cn/...
+swanlab_lark_secret: your-webhook-secret
+```
+
+---
+
+## Common Use Cases
+
+### Use Case 1: Migrate from WandB to SwanLab
+Start with `lora-swanlab-profiling.yml`, add your model/dataset config, disable WandB:
+```yaml
+use_swanlab: true
+use_wandb: false
+```
+
+### Use Case 2: Analyze DPO Model Outputs
+Use `dpo-swanlab-completions.yml`, adjust completion logging interval based on your training length:
+```yaml
+swanlab_completion_log_interval: 50   # More frequent for short training
+swanlab_completion_log_interval: 200  # Less frequent for long training
+```
+
+### Use Case 3: Optimize Training Performance
+Use `lora-swanlab-profiling.yml`, run multiple experiments with different optimizations:
+- Baseline: `flash_attention: false, gradient_checkpointing: false`
+- Flash Attention: `flash_attention: true`
+- Gradient Checkpointing: `gradient_checkpointing: true`
+- Both: `flash_attention: true, gradient_checkpointing: true`
+
+Compare profiling metrics in SwanLab dashboard.
+
+### Use Case 4: Production RLHF with Team Collaboration
+Use `dpo-swanlab-full-featured.yml`, set up team workspace and Lark notifications:
+```yaml
+swanlab_workspace: ml-team
+swanlab_lark_webhook_url: ...
+swanlab_lark_secret: ...
+```
+
+---
+
+## Viewing Your Experiments
+
+### Cloud Mode
+Visit [https://swanlab.cn](https://swanlab.cn) and navigate to your project.
+
+**Dashboard sections**:
+- **Metrics**: Training loss, learning rate, profiling metrics
+- **Tables**: RLHF completions (for DPO/KTO/ORPO/GRPO)
+- **Config**: Hyperparameters and configuration
+- **System**: Resource usage (GPU, memory, CPU)
+- **Files**: Logged artifacts
+
+### Local Mode
+```bash
+swanlab watch ./swanlog
+# Open browser to http://localhost:5092
+```
+
+---
+
+## Troubleshooting
+
+### SwanLab not initializing
+```bash
+# Check API key
+echo $SWANLAB_API_KEY
+
+# Verify SwanLab is installed
+pip show swanlab
+
+# Check config
+grep -A 5 "use_swanlab" your-config.yml
+```
+
+### Completions not appearing
+- Verify you're using an RLHF trainer (DPO/KTO/ORPO/GRPO)
+- Check `swanlab_log_completions: true`
+- Wait for `swanlab_completion_log_interval` steps
+- Look for "Registered SwanLab RLHF completion logging" in logs
+
+### Lark notifications not working
+- Test webhook manually: `curl -X POST "$SWANLAB_LARK_WEBHOOK_URL" ...`
+- Verify `SWANLAB_LARK_SECRET` is set correctly
+- Check bot is added to Lark group chat
+- Look for "Registered Lark notification callback" in logs
+
+### Profiling metrics not appearing
+- Verify `use_swanlab: true`
+- Check SwanLab is initialized (look for init log message)
+- Profiling metrics are under "profiling/" namespace
+- Profiling auto-enabled when SwanLab is enabled
+
+---
+
+## Performance Notes
+
+### Overhead Comparison
+
+| Feature | Overhead per Step | Memory Usage |
+|---------|------------------|--------------|
+| Basic tracking | < 0.1% | ~10 MB |
+| Completion logging | < 0.5% | ~64 KB (buffer=128) |
+| Profiling | < 0.1% | ~1 KB |
+| **Total** | **< 0.7%** | **~10 MB** |
+
+### Best Practices
+1. Use ONE logging tool in production (disable WandB/MLflow when using SwanLab)
+2. Adjust completion log interval based on training length (100-200 steps)
+3. Keep completion buffer size reasonable (128-512)
+4. Profile critical path methods first (training_step, compute_loss)
+5. Use ProfilingConfig to throttle high-frequency operations
+
+---
+
+## Further Reading
+
+- **Full Documentation**: [src/axolotl/integrations/swanlab/README.md](../../src/axolotl/integrations/swanlab/README.md)
+- **SwanLab Docs**: [https://docs.swanlab.cn](https://docs.swanlab.cn)
+- **Axolotl Docs**: [https://axolotl-ai-cloud.github.io/axolotl/](https://axolotl-ai-cloud.github.io/axolotl/)
+- **DPO Paper**: [Direct Preference Optimization](https://arxiv.org/abs/2305.18290)
+
+---
+
+## Contributing
+
+Found an issue or have an improvement? Please submit a PR or open an issue:
+- [Axolotl Issues](https://github.com/axolotl-ai-cloud/axolotl/issues)
+- [SwanLab Issues](https://github.com/SwanHubX/SwanLab/issues)
diff --git a/examples/swanlab/custom_trainer_profiling.py b/examples/swanlab/custom_trainer_profiling.py
new file mode 100644
index 0000000000..65461c4e5a
--- /dev/null
+++ b/examples/swanlab/custom_trainer_profiling.py
@@ -0,0 +1,299 @@
+"""Example: Custom Trainer with SwanLab Profiling
+
+This example demonstrates how to add SwanLab profiling to your custom trainer.
+
+Features:
+- @swanlab_profile decorator for automatic profiling
+- swanlab_profiling_context for fine-grained profiling
+- ProfilingConfig for advanced filtering and throttling
+
+Usage:
+    1. Create your custom trainer extending AxolotlTrainer
+    2. Add @swanlab_profile decorators to methods you want to profile
+    3. Use swanlab_profiling_context for fine-grained profiling within methods
+    4. Enable SwanLab in your config (use_swanlab: true)
+
+See also:
+    - examples/swanlab/lora-swanlab-profiling.yml for config
+    - src/axolotl/integrations/swanlab/profiling.py for implementation
+"""
+
+from axolotl.core.trainers.base import AxolotlTrainer
+from axolotl.integrations.swanlab.profiling import (
+    ProfilingConfig,
+    swanlab_profile,
+    swanlab_profiling_context,
+    swanlab_profiling_context_advanced,
+)
+
+
+class CustomTrainerWithProfiling(AxolotlTrainer):
+    """Custom trainer with SwanLab profiling enabled.
+
+    This trainer demonstrates three profiling patterns:
+    1. Decorator-based profiling (@swanlab_profile)
+    2. Context manager profiling (swanlab_profiling_context)
+    3. Advanced profiling with filtering (ProfilingConfig)
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Create custom profiling config for high-frequency operations
+        self.fast_op_config = ProfilingConfig(
+            enabled=True,
+            min_duration_ms=0.5,  # Only log if duration > 0.5ms
+            log_interval=50,  # Log every 50th call
+        )
+
+    # ========================================================================
+    # Pattern 1: Decorator-based Profiling
+    # ========================================================================
+    # Best for: Methods you always want to profile
+    # Overhead: ~2-5 microseconds per call (negligible)
+
+    @swanlab_profile
+    def training_step(self, model, inputs):
+        """Main training step - always profile.
+
+        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.training_step
+        """
+        return super().training_step(model, inputs)
+
+    @swanlab_profile
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """Loss computation - always profile.
+
+        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.compute_loss
+        """
+        return super().compute_loss(model, inputs, return_outputs)
+
+    @swanlab_profile
+    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+        """Prediction step - always profile.
+
+        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.prediction_step
+        """
+        return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys)
+
+    # ========================================================================
+    # Pattern 2: Fine-grained Context Manager Profiling
+    # ========================================================================
+    # Best for: Profiling specific code blocks within a method
+    # Use case: When you want to profile forward vs backward separately
+
+    def complex_training_step(self, model, inputs):
+        """Training step with fine-grained profiling.
+
+        Profiling metrics:
+        - profiling/Time taken: CustomTrainerWithProfiling.forward_pass
+        - profiling/Time taken: CustomTrainerWithProfiling.backward_pass
+        - profiling/Time taken: CustomTrainerWithProfiling.optimizer_step
+        """
+        # Profile just the forward pass
+        with swanlab_profiling_context(self, "forward_pass"):
+            outputs = model(**inputs)
+            loss = outputs.loss
+
+        # Profile just the backward pass
+        with swanlab_profiling_context(self, "backward_pass"):
+            loss.backward()
+
+        # Profile optimizer step
+        with swanlab_profiling_context(self, "optimizer_step"):
+            self.optimizer.step()
+            self.optimizer.zero_grad()
+
+        return outputs
+
+    # ========================================================================
+    # Pattern 3: Advanced Profiling with Filtering
+    # ========================================================================
+    # Best for: High-frequency operations where you want to throttle logging
+    # Use case: Methods called 100+ times per step
+
+    def _prepare_inputs(self, inputs):
+        """Prepare inputs - throttled profiling.
+
+        This method is called frequently (once per batch), so we throttle
+        profiling to reduce overhead:
+        - Only log if duration > 0.5ms (skip very fast operations)
+        - Only log every 50th call (reduce logging frequency)
+
+        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.prepare_inputs
+        """
+        with swanlab_profiling_context_advanced(
+            self, "prepare_inputs", config=self.fast_op_config
+        ):
+            return super()._prepare_inputs(inputs)
+
+    def _prepare_input_for_model(self, input_ids):
+        """Another high-frequency operation - throttled profiling.
+
+        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.prepare_input_for_model
+        """
+        with swanlab_profiling_context_advanced(
+            self, "prepare_input_for_model", config=self.fast_op_config
+        ):
+            # Your custom input preparation logic
+            return input_ids
+
+    # ========================================================================
+    # Pattern 4: Exception-safe Profiling
+    # ========================================================================
+    # Profiling is exception-safe: duration is logged even if method raises
+
+    @swanlab_profile
+    def potentially_failing_method(self):
+        """This method may raise an exception.
+
+        SwanLab profiling will still log the duration before re-raising.
+        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.potentially_failing_method
+        """
+        # Do some work
+        result = self._do_risky_computation()
+
+        # If this raises, profiling duration is still logged
+        if result < 0:
+            raise ValueError("Invalid result")
+
+        return result
+
+    def _do_risky_computation(self):
+        """Placeholder for risky computation."""
+        return 42
+
+
+# ============================================================================
+# Advanced Example: Custom ProfilingConfig Per Method
+# ============================================================================
+
+
+class AdvancedProfilingTrainer(AxolotlTrainer):
+    """Trainer with method-specific profiling configurations."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Different profiling configs for different method types
+        self.critical_path_config = ProfilingConfig(
+            enabled=True,
+            min_duration_ms=0.0,  # Log everything on critical path
+            log_interval=1,  # Log every call
+        )
+
+        self.fast_path_config = ProfilingConfig(
+            enabled=True,
+            min_duration_ms=1.0,  # Only log if > 1ms
+            log_interval=100,  # Log every 100th call
+        )
+
+        self.debug_config = ProfilingConfig(
+            enabled=True,
+            min_duration_ms=0.0,  # Log everything
+            log_interval=1,  # Log every call
+        )
+
+    def training_step(self, model, inputs):
+        """Critical path - log everything."""
+        with swanlab_profiling_context_advanced(
+            self, "training_step", config=self.critical_path_config
+        ):
+            return super().training_step(model, inputs)
+
+    def _prepare_inputs(self, inputs):
+        """Fast path - throttle logging."""
+        with swanlab_profiling_context_advanced(
+            self, "prepare_inputs", config=self.fast_path_config
+        ):
+            return super()._prepare_inputs(inputs)
+
+    def _debug_method(self, data):
+        """Debug-only method - verbose logging."""
+        with swanlab_profiling_context_advanced(
+            self, "debug_method", config=self.debug_config
+        ):
+            # Your debug logic
+            pass
+
+
+# ============================================================================
+# How to Use This Custom Trainer
+# ============================================================================
+
+"""
+To use this custom trainer:
+
+1. Save this file to your project (e.g., my_custom_trainer.py)
+
+2. Create a config file that uses your custom trainer:
+
+    # config.yml
+    base_model: NousResearch/Llama-3.2-1B
+
+    # ... other config ...
+
+    plugins:
+      - axolotl.integrations.swanlab.SwanLabPlugin
+
+    use_swanlab: true
+    swanlab_project: my-profiling-experiment
+
+    # Optional: Specify custom trainer
+    # (Or modify axolotl to use your custom trainer class)
+
+3. Run training:
+
+    export SWANLAB_API_KEY=your-api-key
+    accelerate launch -m axolotl.cli.train config.yml
+
+4. View profiling metrics in SwanLab dashboard:
+   - profiling/Time taken: CustomTrainerWithProfiling.training_step
+   - profiling/Time taken: CustomTrainerWithProfiling.forward_pass
+   - profiling/Time taken: CustomTrainerWithProfiling.backward_pass
+   - etc.
+
+5. Compare profiling metrics across runs:
+   - Run baseline without optimizations
+   - Run with flash_attention enabled
+   - Run with gradient_checkpointing enabled
+   - Compare profiling metrics to see performance impact
+"""
+
+# ============================================================================
+# Tips for Effective Profiling
+# ============================================================================
+
+"""
+1. Profile the critical path first:
+   - training_step, compute_loss, prediction_step
+   - These methods are called most frequently and have biggest impact
+
+2. Use throttling for high-frequency operations:
+   - Methods called 100+ times per step
+   - Use log_interval=50 or log_interval=100
+   - Reduces profiling overhead and dashboard clutter
+
+3. Filter noise with min_duration_ms:
+   - Set min_duration_ms=1.0 to skip very fast operations
+   - Focus on operations that actually take time
+
+4. Compare across runs:
+   - Run same config multiple times to check consistency
+   - Compare different optimization strategies
+   - Track profiling trends over time
+
+5. Monitor distributed training:
+   - Check for per-rank timing differences
+   - Look for stragglers (slower ranks)
+   - Identify synchronization bottlenecks
+
+6. Disable profiling in production:
+   - from axolotl.integrations.swanlab.profiling import DEFAULT_PROFILING_CONFIG
+   - DEFAULT_PROFILING_CONFIG.enabled = False
+
+7. Exception handling:
+   - Profiling is exception-safe
+   - Duration logged even if method raises
+   - Useful for debugging methods that fail intermittently
+"""
diff --git a/examples/swanlab/dpo-swanlab-completions.yml b/examples/swanlab/dpo-swanlab-completions.yml
new file mode 100644
index 0000000000..fb21dbbba1
--- /dev/null
+++ b/examples/swanlab/dpo-swanlab-completions.yml
@@ -0,0 +1,168 @@
+# SwanLab DPO Training Example with Completion Logging
+#
+# This example demonstrates DPO (Direct Preference Optimization) training
+# with SwanLab integration for experiment tracking and completion table logging.
+#
+# Features enabled:
+# - SwanLab experiment tracking
+# - RLHF completion table logging (prompts, chosen/rejected responses, rewards)
+# - Lark (Feishu) team notifications (optional)
+#
+# To run:
+#   export SWANLAB_API_KEY=your-api-key
+#   accelerate launch -m axolotl.cli.train examples/swanlab/dpo-swanlab-completions.yml
+
+# Model Configuration
+base_model: meta-llama/Meta-Llama-3-8B-Instruct
+model_type: LlamaForCausalLM
+tokenizer_type: AutoTokenizer
+
+special_tokens:
+  pad_token: <|finetune_right_pad_id|>
+  eos_token: <|eot_id|>
+
+# Quantization
+load_in_8bit: true
+load_in_4bit: false
+
+# LoRA Configuration
+adapter: lora
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+# DPO Configuration
+chat_template: llama3
+rl: dpo
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_dpo_test
+    type: chat_template.default
+    field_messages: conversation
+    field_chosen: chosen
+    field_rejected: rejected
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      system:
+        - system
+      user:
+        - user
+      assistant:
+        - assistant
+
+# Dataset and Output
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/dpo-swanlab-out
+
+# Training Configuration
+sequence_len: 4096
+sample_packing: false
+micro_batch_size: 2
+gradient_accumulation_steps: 4
+num_epochs: 4
+
+# Optimization
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+# Precision
+bf16: auto
+tf32: false
+
+# Performance
+gradient_checkpointing: true
+attn_implementation: flash_attention_2
+
+# Checkpointing and Logging
+logging_steps: 1
+evals_per_epoch: 4
+saves_per_epoch: 1
+
+# ============================================================================
+# SwanLab Integration
+# ============================================================================
+
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+# Basic SwanLab Configuration
+use_swanlab: true
+swanlab_project: dpo-training
+swanlab_experiment_name: llama-3-dpo-completions-demo
+swanlab_description: "DPO training with completion table logging"
+swanlab_mode: cloud  # Options: cloud, local, offline, disabled
+
+# SwanLab Authentication
+# Recommended: Set via environment variable
+#   export SWANLAB_API_KEY=your-api-key
+# Or set in config (less secure):
+# swanlab_api_key: your-api-key
+
+# Optional: Team workspace
+# swanlab_workspace: my-research-team
+
+# ============================================================================
+# RLHF Completion Table Logging
+# ============================================================================
+#
+# Automatically logs model completions to SwanLab for qualitative analysis:
+# - Prompts from your DPO dataset
+# - Chosen responses (preferred)
+# - Rejected responses (non-preferred)
+# - Reward differences
+#
+# View the table in SwanLab dashboard under "rlhf_completions"
+
+swanlab_log_completions: true
+swanlab_completion_log_interval: 100  # Log every 100 training steps
+swanlab_completion_max_buffer: 128    # Keep last 128 completions in memory
+
+# Memory Usage Notes:
+# - Buffer size 128: ~64 KB (default, recommended)
+# - Buffer size 512: ~256 KB (for more historical completions)
+# - Buffer size 1024: ~512 KB (maximum for very long training runs)
+
+# Performance Notes:
+# - Completion logging overhead: < 0.5% per training step
+# - Only logs every N steps to minimize impact
+# - Memory-bounded buffer prevents memory leaks
+
+# ============================================================================
+# Optional: Lark (Feishu) Team Notifications
+# ============================================================================
+#
+# Get real-time training notifications in your team chat
+# Uncomment to enable:
+
+# swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
+# swanlab_lark_secret: your-webhook-secret  # Recommended for production
+
+# Notifications sent for:
+# - Training start
+# - Training completion
+# - Training errors
+# - Metric milestones (if configured)
+
+# ============================================================================
+# Optional: Private SwanLab Deployment
+# ============================================================================
+#
+# For enterprise users with private SwanLab deployment:
+
+# swanlab_web_host: https://swanlab.yourcompany.com
+# swanlab_api_host: https://api.swanlab.yourcompany.com
+
+# ============================================================================
+# Disable WandB if you're migrating from it
+# ============================================================================
+
+# wandb_project:
+# wandb_entity:
+# use_wandb: false
diff --git a/examples/swanlab/dpo-swanlab-full-featured.yml b/examples/swanlab/dpo-swanlab-full-featured.yml
new file mode 100644
index 0000000000..ac52e6a85b
--- /dev/null
+++ b/examples/swanlab/dpo-swanlab-full-featured.yml
@@ -0,0 +1,329 @@
+# SwanLab Full-Featured DPO Training Example
+#
+# This example demonstrates ALL SwanLab integration features:
+# - Experiment tracking with cloud sync
+# - RLHF completion table logging
+# - Performance profiling
+# - Lark (Feishu) team notifications
+# - Team workspace collaboration
+#
+# Use this as a reference for production RLHF training setups.
+#
+# To run:
+#   export SWANLAB_API_KEY=your-api-key
+#   export SWANLAB_LARK_WEBHOOK_URL=https://open.feishu.cn/...
+#   export SWANLAB_LARK_SECRET=your-webhook-secret
+#   accelerate launch -m axolotl.cli.train examples/swanlab/dpo-swanlab-full-featured.yml
+
+# ============================================================================
+# Model Configuration
+# ============================================================================
+
+base_model: meta-llama/Meta-Llama-3-8B-Instruct
+model_type: LlamaForCausalLM
+tokenizer_type: AutoTokenizer
+
+special_tokens:
+  pad_token: <|finetune_right_pad_id|>
+  eos_token: <|eot_id|>
+
+# Quantization for efficient training
+load_in_8bit: true
+load_in_4bit: false
+
+# ============================================================================
+# LoRA Configuration
+# ============================================================================
+
+adapter: lora
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true  # Target all linear layers
+
+# ============================================================================
+# DPO (Direct Preference Optimization) Configuration
+# ============================================================================
+
+chat_template: llama3
+rl: dpo  # Enable DPO trainer
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_dpo_test
+    type: chat_template.default
+    field_messages: conversation
+    field_chosen: chosen
+    field_rejected: rejected
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      system:
+        - system
+      user:
+        - user
+      assistant:
+        - assistant
+
+# ============================================================================
+# Dataset and Output Configuration
+# ============================================================================
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/dpo-swanlab-full-featured-out
+
+# ============================================================================
+# Training Configuration
+# ============================================================================
+
+sequence_len: 4096
+sample_packing: false
+
+micro_batch_size: 2
+gradient_accumulation_steps: 4
+num_epochs: 4
+
+# ============================================================================
+# Optimization
+# ============================================================================
+
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+# ============================================================================
+# Precision and Performance
+# ============================================================================
+
+bf16: auto
+tf32: false
+
+gradient_checkpointing: true
+attn_implementation: flash_attention_2
+
+# ============================================================================
+# Checkpointing and Logging
+# ============================================================================
+
+logging_steps: 1
+evals_per_epoch: 4
+saves_per_epoch: 1
+
+# ============================================================================
+# SwanLab Integration - Full Configuration
+# ============================================================================
+
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+# ------------------------------------------------------------------------------
+# Basic SwanLab Configuration
+# ------------------------------------------------------------------------------
+
+use_swanlab: true
+swanlab_project: dpo-production
+swanlab_experiment_name: llama-3-dpo-full-featured-v1
+swanlab_description: |
+  Production DPO training with all SwanLab features enabled:
+  - Completion table logging for qualitative analysis
+  - Performance profiling for optimization
+  - Lark notifications for team collaboration
+
+swanlab_mode: cloud  # Options: cloud, local, offline, disabled
+
+# ------------------------------------------------------------------------------
+# Team Collaboration
+# ------------------------------------------------------------------------------
+
+# Workspace for team collaboration (shared experiments)
+swanlab_workspace: ml-research-team
+
+# Authentication (recommended: use environment variable)
+#   export SWANLAB_API_KEY=your-api-key
+# Or set in config (less secure):
+# swanlab_api_key: your-api-key
+
+# ------------------------------------------------------------------------------
+# RLHF Completion Table Logging
+# ------------------------------------------------------------------------------
+# Automatically logs model completions for qualitative analysis:
+# - Prompts from your DPO dataset
+# - Chosen responses (preferred)
+# - Rejected responses (non-preferred)
+# - Reward differences
+#
+# View in SwanLab dashboard under "rlhf_completions" table
+
+swanlab_log_completions: true
+swanlab_completion_log_interval: 100  # Log every 100 steps
+swanlab_completion_max_buffer: 256    # Larger buffer for long training runs
+
+# Buffer size recommendations:
+# - 128: Default, ~64 KB memory (recommended for most cases)
+# - 256: ~128 KB memory (this config, good for longer training)
+# - 512: ~256 KB memory (maximum for very long runs)
+
+# ------------------------------------------------------------------------------
+# Lark (Feishu) Team Notifications
+# ------------------------------------------------------------------------------
+# Get real-time training notifications in your team chat
+#
+# Notifications sent for:
+# - Training start
+# - Training completion
+# - Training errors
+# - Metric milestones (if configured)
+
+# Recommended: Set via environment variables
+#   export SWANLAB_LARK_WEBHOOK_URL=https://open.feishu.cn/...
+#   export SWANLAB_LARK_SECRET=your-webhook-secret
+
+# Or set in config (less secure):
+# swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
+# swanlab_lark_secret: your-webhook-secret  # REQUIRED for production
+
+# Security note: ALWAYS use swanlab_lark_secret in production to prevent
+# unauthorized parties from sending fake notifications to your team chat.
+
+# ------------------------------------------------------------------------------
+# Performance Profiling
+# ------------------------------------------------------------------------------
+# Profiling is automatically enabled when SwanLab is enabled.
+# Metrics logged to SwanLab under "profiling/" namespace:
+#   profiling/Time taken: AxolotlTrainer.training_step
+#   profiling/Time taken: AxolotlTrainer.compute_loss
+#   profiling/Time taken: AxolotlTrainer.prediction_step
+#
+# Use these metrics to:
+# - Identify bottlenecks in training loop
+# - Compare performance across different configurations
+# - Monitor performance regressions over time
+# - Debug unexpected slowdowns
+
+# For custom profiling in your own trainer, see:
+#   examples/swanlab/custom_trainer_profiling.py
+
+# ------------------------------------------------------------------------------
+# Optional: Private SwanLab Deployment
+# ------------------------------------------------------------------------------
+# For enterprise users with private SwanLab deployment:
+
+# swanlab_web_host: https://swanlab.yourcompany.com
+# swanlab_api_host: https://api.swanlab.yourcompany.com
+
+# ------------------------------------------------------------------------------
+# Optional: Model Checkpointing to SwanLab
+# ------------------------------------------------------------------------------
+# Log model checkpoints to SwanLab (coming soon)
+
+swanlab_log_model: false
+
+# ============================================================================
+# Disable Other Logging Tools (Recommended)
+# ============================================================================
+# Using multiple logging tools simultaneously can impact performance:
+# - Expected overhead: ~1-2% per logger
+# - Potential config/callback conflicts
+#
+# For production training, use ONLY SwanLab:
+
+# wandb_project:
+# use_wandb: false
+#
+# use_mlflow: false
+#
+# use_comet: false
+
+# ============================================================================
+# Expected Training Behavior
+# ============================================================================
+
+# With this configuration, you should see:
+#
+# 1. SwanLab Initialization (rank 0 only):
+#    INFO: SwanLab initialized for project: dpo-production
+#    INFO: SwanLab experiment: llama-3-dpo-full-featured-v1
+#    INFO: SwanLab mode: cloud
+#    INFO: SwanLab workspace: ml-research-team
+#
+# 2. Completion Logging (rank 0 only):
+#    INFO: Registered SwanLab RLHF completion logging callback for DPOTrainer
+#          (log_interval=100, max_buffer=256)
+#
+# 3. Lark Notifications (rank 0 only):
+#    INFO: Registered Lark notification callback with HMAC authentication
+#
+# 4. Distributed Training Detection (if multi-GPU):
+#    INFO: Distributed training detected (world_size=N)
+#    INFO: Only rank 0 will initialize SwanLab
+#    INFO: Other ranks will skip SwanLab to avoid conflicts
+#
+# 5. Training Start Notification (Lark):
+#    Your team chat receives: "Training started: llama-3-dpo-full-featured-v1"
+#
+# 6. Periodic Completion Logging:
+#    Every 100 steps, completion table is updated in SwanLab dashboard
+#
+# 7. Training Complete Notification (Lark):
+#    Your team chat receives: "Training completed: llama-3-dpo-full-featured-v1"
+#    With link to SwanLab dashboard and final metrics
+#
+# 8. SwanLab Dashboard Shows:
+#    - Training metrics (loss, learning rate, etc.)
+#    - Completion table (rlhf_completions)
+#    - Profiling metrics (profiling/Time taken: ...)
+#    - Hyperparameters and configuration
+#    - System resource usage
+
+# ============================================================================
+# Production Checklist
+# ============================================================================
+
+# Before deploying to production, verify:
+# ✅ SwanLab API key is set via environment variable (not in config)
+# ✅ Lark webhook secret is set (required for HMAC authentication)
+# ✅ Workspace is set to your team's workspace
+# ✅ Experiment name is descriptive and unique
+# ✅ Only SwanLab is enabled (other loggers disabled)
+# ✅ Completion logging buffer size is appropriate for your training duration
+# ✅ Private deployment hosts are set (if using enterprise SwanLab)
+# ✅ Test run completes successfully and shows up in SwanLab dashboard
+# ✅ Lark notifications are received in team chat
+# ✅ Profiling metrics are logged correctly
+
+# ============================================================================
+# Troubleshooting
+# ============================================================================
+
+# If SwanLab initialization fails:
+# 1. Check SWANLAB_API_KEY environment variable is set
+# 2. Verify swanlab_project is set in config
+# 3. Check swanlab_mode is valid (cloud/local/offline/disabled)
+# 4. Verify internet connectivity (for cloud mode)
+
+# If Lark notifications not received:
+# 1. Check SWANLAB_LARK_WEBHOOK_URL is set correctly
+# 2. Verify SWANLAB_LARK_SECRET matches your Lark bot settings
+# 3. Test webhook manually: curl -X POST "$SWANLAB_LARK_WEBHOOK_URL" ...
+# 4. Check training logs for "Registered Lark notification callback"
+# 5. Verify bot is added to the target Lark group chat
+
+# If completions not appearing in SwanLab:
+# 1. Verify you're using an RLHF trainer (DPO/KTO/ORPO/GRPO)
+# 2. Check swanlab_log_completions is true
+# 3. Wait for log_interval steps (default: 100)
+# 4. Check training logs for "Registered SwanLab RLHF completion logging"
+
+# If profiling metrics not appearing:
+# 1. Verify use_swanlab is true
+# 2. Check SwanLab is initialized (check logs)
+# 3. Look under "profiling/" namespace in dashboard
+# 4. Profiling may be disabled if DEFAULT_PROFILING_CONFIG.enabled = False
+
+# For more help:
+# - SwanLab docs: https://docs.swanlab.cn
+# - Axolotl SwanLab integration: src/axolotl/integrations/swanlab/README.md
+# - GitHub issues: https://github.com/axolotl-ai-cloud/axolotl/issues
diff --git a/examples/swanlab/lora-swanlab-profiling.yml b/examples/swanlab/lora-swanlab-profiling.yml
new file mode 100644
index 0000000000..3dff6e315f
--- /dev/null
+++ b/examples/swanlab/lora-swanlab-profiling.yml
@@ -0,0 +1,178 @@
+# SwanLab LoRA Training Example with Performance Profiling
+#
+# This example demonstrates standard LoRA fine-tuning with SwanLab integration
+# for performance profiling and optimization.
+#
+# Features enabled:
+# - SwanLab experiment tracking
+# - Performance profiling (training step, forward/backward pass timing)
+# - Real-time metrics visualization
+#
+# To run:
+#   export SWANLAB_API_KEY=your-api-key
+#   accelerate launch -m axolotl.cli.train examples/swanlab/lora-swanlab-profiling.yml
+
+# Model Configuration
+base_model: NousResearch/Llama-3.2-1B
+
+# Dataset Configuration
+datasets:
+  - path: teknium/GPT4-LLM-Cleaned
+    type: alpaca
+
+val_set_size: 0.1
+output_dir: ./outputs/lora-swanlab-profiling-out
+
+# LoRA Configuration
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_dropout: 0.05
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+
+# Training Configuration
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+
+micro_batch_size: 2
+gradient_accumulation_steps: 2
+num_epochs: 1
+
+# Optimization
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+warmup_ratio: 0.1
+weight_decay: 0.0
+
+# Precision
+bf16: auto
+tf32: false
+
+# Performance
+gradient_checkpointing: true
+attn_implementation: flash_attention_2
+
+# Checkpointing and Logging
+logging_steps: 1
+evals_per_epoch: 4
+saves_per_epoch: 1
+
+# Loss Monitoring
+loss_watchdog_threshold: 5.0
+loss_watchdog_patience: 3
+
+special_tokens:
+  pad_token: "<|end_of_text|>"
+
+# ============================================================================
+# SwanLab Integration
+# ============================================================================
+
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+# Basic SwanLab Configuration
+use_swanlab: true
+swanlab_project: lora-profiling
+swanlab_experiment_name: llama-3.2-1b-profiling-demo
+swanlab_description: "LoRA fine-tuning with performance profiling"
+swanlab_mode: cloud  # Options: cloud, local, offline, disabled
+
+# SwanLab Authentication
+# Recommended: Set via environment variable
+#   export SWANLAB_API_KEY=your-api-key
+# Or set in config (less secure):
+# swanlab_api_key: your-api-key
+
+# Optional: Team workspace
+# swanlab_workspace: my-ml-team
+
+# ============================================================================
+# Performance Profiling
+# ============================================================================
+#
+# SwanLab automatically profiles trainer methods when enabled.
+# Profiling metrics appear in SwanLab dashboard under "profiling/" namespace.
+#
+# Built-in profiling:
+# - Minimal overhead (< 0.1% per step)
+# - High-precision timing (microsecond accuracy)
+# - Exception-safe (logs duration even if method fails)
+#
+# View profiling metrics in SwanLab dashboard:
+#   profiling/Time taken: AxolotlTrainer.training_step
+#   profiling/Time taken: AxolotlTrainer.compute_loss
+#   profiling/Time taken: AxolotlTrainer.prediction_step
+#
+# For custom profiling in your own trainer, see:
+#   examples/swanlab/custom_trainer_profiling.py
+
+# Completion logging is disabled for non-RLHF trainers
+swanlab_log_completions: false  # Only works with DPO/KTO/ORPO/GRPO
+
+# ============================================================================
+# Optional: Compare with Multiple Runs
+# ============================================================================
+#
+# To compare profiling metrics across different configurations:
+#
+# 1. Run baseline without flash attention:
+#    swanlab_experiment_name: llama-3.2-1b-no-flash-attn
+#    flash_attention: false
+#
+# 2. Run with gradient checkpointing:
+#    swanlab_experiment_name: llama-3.2-1b-grad-checkpoint
+#    gradient_checkpointing: true
+#
+# 3. Run with both:
+#    swanlab_experiment_name: llama-3.2-1b-optimized
+#    flash_attention: true
+#    gradient_checkpointing: true
+#
+# Then compare profiling metrics in SwanLab dashboard to see performance impact
+
+# ============================================================================
+# Optional: Lark (Feishu) Team Notifications
+# ============================================================================
+#
+# Get notified when profiling experiments complete:
+
+# swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
+# swanlab_lark_secret: your-webhook-secret
+
+# ============================================================================
+# Profiling Best Practices
+# ============================================================================
+#
+# 1. Run multiple epochs to see profiling trends over time
+# 2. Ignore first ~10 steps (warmup period, slower)
+# 3. Look for outliers (steps that take significantly longer)
+# 4. Compare profiling metrics before/after optimization changes
+# 5. Monitor per-rank profiling in distributed training
+#
+# Common bottlenecks to profile:
+# - training_step: Overall step time (should be consistent)
+# - compute_loss: Loss computation (scales with sequence length)
+# - prediction_step: Evaluation time (can be slow for large val sets)
+#
+# If you see inconsistent timing:
+# - Check for data loading bottlenecks
+# - Monitor GPU utilization (may be CPU-bound)
+# - Check for gradient accumulation effects
+# - Verify CUDA kernel synchronization
+
+# ============================================================================
+# Disable WandB if you're migrating from it
+# ============================================================================
+
+# wandb_project:
+# use_wandb: false
diff --git a/examples/trinity/README.md b/examples/trinity/README.md
new file mode 100644
index 0000000000..e9710915ca
--- /dev/null
+++ b/examples/trinity/README.md
@@ -0,0 +1,40 @@
+# Finetune ArceeAI's Trinity with Axolotl
+
+[Trinity](https://huggingface.co/collections/arcee-ai/trinity) is a family of open weight MoE models trained by Arcee.ai.
+
+This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+
+## Getting started
+
+1. Install Axolotl following the main from the [installation guide](https://docs.axolotl.ai/docs/installation.html#sec-edge-build).
+
+2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+
+3. Run the finetuning example:
+
+    ```bash
+    axolotl train examples/trinity/trinity-nano-preview-qlora.yaml
+    ```
+
+This config uses about 24.9 GiB VRAM (w/o CCE).
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- For inference, the official Arcee.ai team recommends `top_p: 0.75`, `temperature: 0.15`, `top_k: 50`, and `min_p: 0.06`.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+
+## Optimization Guides
+
+Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
+
+## Related Resources
+
+- [Trinity Blog](https://www.arcee.ai/blog/the-trinity-manifesto)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
diff --git a/examples/trinity/trinity-nano-preview-qlora.yaml b/examples/trinity/trinity-nano-preview-qlora.yaml
new file mode 100644
index 0000000000..52c0c0c60c
--- /dev/null
+++ b/examples/trinity/trinity-nano-preview-qlora.yaml
@@ -0,0 +1,67 @@
+base_model: arcee-ai/Trinity-Nano-Preview
+revision_of_model: 2ee94b0
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+# CCE - N/A as of now
+# plugins:
+#   - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.1
+output_dir: ./outputs/lora-out
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 2048
+sample_packing: true
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_target_modules:
+  - gate_proj
+  - down_proj
+  - up_proj
+  - q_proj
+  - v_proj
+  - k_proj
+  - 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: auto
+tf32: false
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+# flash_attention: true  # Not supported
+attn_implementation: sdpa
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
diff --git a/examples/voxtral/README.md b/examples/voxtral/README.md
new file mode 100644
index 0000000000..f8e7b51be1
--- /dev/null
+++ b/examples/voxtral/README.md
@@ -0,0 +1,82 @@
+# Finetune Voxtral with Axolotl
+
+Voxtral is a [3B](https://huggingface.co/mistralai/Voxtral-Mini-3B-2507)/[24B](https://huggingface.co/mistralai/Voxtral-Small-24B-2507) parameter opensource model from MistralAI found on HuggingFace. This guide shows how to fine-tune it with Axolotl.
+
+Thanks to the team at MistralAI for giving us early access to prepare for this release.
+
+## Getting started
+
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.9.1 min)
+uv pip install --no-build-isolation 'axolotl>=0.16.1'
+```
+
+2. Please install the below.
+
+```bash
+# audio
+uv pip install librosa==0.11.0
+uv pip install 'mistral_common[audio]==1.8.3'
+
+# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
+python scripts/cutcrossentropy_install.py | sh
+```
+
+3. Download sample dataset files
+
+```bash
+# for text + audio only
+wget https://huggingface.co/datasets/Nanobit/text-audio-2k-test/resolve/main/En-us-African_elephant.oga
+```
+
+4. Run the finetuning example:
+
+```bash
+# text only
+axolotl train examples/voxtral/voxtral-mini-qlora.yml
+
+# text + audio
+axolotl train examples/voxtral/voxtral-mini-audio-qlora.yml
+```
+
+These configs use about 4.8 GB VRAM.
+
+Let us know how it goes. Happy finetuning! 🚀
+
+### TIPS
+
+- For inference, the official MistralAI team recommends `temperature: 0.2` and `top_p: 0.95` for audio understanding and `temperature: 0.0` for transcription.
+- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
+- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
+- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+- The multimodal dataset format follows the OpenAI multi-content Messages format as seen [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+
+## Limitations
+
+We only support the `mistral-common` tokenizer for Supervised Fine-tuning at the moment and for `type: chat_template` only.
+
+In addition, we do not support overriding tokens yet.
+
+## Related Resources
+
+- [MistralAI Magistral Blog](https://mistral.ai/news/magistral/)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl Website](https://axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
+
+## Future Work
+
+- Add parity to Preference Tuning, RL, etc.
+- Add parity to other tokenizer configs like overriding tokens.
diff --git a/examples/voxtral/voxtral-mini-audio-qlora.yml b/examples/voxtral/voxtral-mini-audio-qlora.yml
new file mode 100644
index 0000000000..cfa351ccde
--- /dev/null
+++ b/examples/voxtral/voxtral-mini-audio-qlora.yml
@@ -0,0 +1,78 @@
+base_model: mistralai/Voxtral-Mini-3B-2507
+processor_type: VoxtralProcessor
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+# for use with fft to only train on language model layers
+# unfrozen_parameters:
+  # - language_model.model.*
+  # - lm_head
+  # - embed_tokens
+
+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
+
+eot_tokens:
+  - <end_of_turn>
+
+# sample dataset below requires downloading audio/image in advance
+# wget https://huggingface.co/datasets/Nanobit/text-audio-2k-test/resolve/main/En-us-African_elephant.oga
+datasets:
+  - path: NanoBit/text-audio-2k-test
+    type: chat_template
+dataset_prepared_path:
+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|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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
diff --git a/examples/voxtral/voxtral-mini-qlora.yml b/examples/voxtral/voxtral-mini-qlora.yml
new file mode 100644
index 0000000000..61e8933d06
--- /dev/null
+++ b/examples/voxtral/voxtral-mini-qlora.yml
@@ -0,0 +1,73 @@
+base_model: mistralai/Voxtral-Mini-3B-2507
+
+# Automatically upload checkpoint and final model to HF
+# hub_model_id: username/custom_model_name
+
+# Enable to use mistral-common tokenizer
+tokenizer_use_mistral_common: true
+
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+load_in_8bit: false
+load_in_4bit: true
+
+# for use with fft to only train on language model layers
+# unfrozen_parameters:
+  # - language_model.model.*
+  # - lm_head
+  # - embed_tokens
+
+eot_tokens:
+  - <end_of_turn>
+datasets:
+  - path: cgato/SlimOrcaDedupCleaned
+    type: chat_template
+    split: train[:1%]
+    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_modules: 'language_model.model.layers.[\d]+.(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: true
+pad_to_sequence_len: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+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
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch:
+saves_per_epoch: 1
+weight_decay: 0.0
+special_tokens:
diff --git a/favicon.jpg b/favicon.jpg
index 43c6902443..4ec3587464 100644
Binary files a/favicon.jpg and b/favicon.jpg differ
diff --git a/image/axolotl-badge-web-legacy.png b/image/axolotl-badge-web-legacy.png
new file mode 100644
index 0000000000..42217dca31
Binary files /dev/null and b/image/axolotl-badge-web-legacy.png differ
diff --git a/image/axolotl-badge-web.png b/image/axolotl-badge-web.png
index 42217dca31..55a5679436 100644
Binary files a/image/axolotl-badge-web.png and b/image/axolotl-badge-web.png differ
diff --git a/image/axolotl_logo_digital_black.svg b/image/axolotl_logo_digital_black.svg
new file mode 100644
index 0000000000..2b0ae35e5c
--- /dev/null
+++ b/image/axolotl_logo_digital_black.svg
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" version="1.1" viewBox="0 0 1113 283.5">
+    <path fill="#141310" d="M435,234.3l-12.1-48.8h-54.4l-12.1,48.8h-24.7l48.2-185.1h31.6l47.9,185.1h-24.5ZM417.7,164.9l-13.8-55.6c-2.7-10.7-4.8-19.7-6.3-26.9-.9-4.2-1.5-7.5-2-9.9-.5,2.5-1.2,5.8-2,9.9-1.5,7.1-3.6,16.1-6.3,26.7l-13.8,55.9h44.3Z"/>
+    <path fill="#141310" d="M568.2,234.3l-29.9-45.6c-1.2-1.9-2.4-4.1-3.5-6.5-.8-1.7-1.5-3.3-2.1-4.5-.6,1.3-1.4,2.8-2.3,4.5-1.3,2.4-2.6,4.6-4,6.5l-29.9,45.6h-28.5l49.6-71.9-46.5-67.9h28.5l27.6,43.1c1.2,1.9,2.3,3.9,3.4,6.1.7,1.4,1.4,2.7,1.9,3.8.5-1.1,1.1-2.4,1.8-3.8,1.1-2.2,2.2-4.2,3.4-6.1l27.6-43.1h28.5l-46.5,68.2,49.3,71.7h-28.5Z"/>
+    <path fill="#141310" d="M658.6,236.3c-16.7,0-30.2-5-40.1-14.8-9.9-9.8-14.9-23.7-14.9-41.3v-31.7c0-17.7,5-31.7,14.8-41.4,9.8-9.7,23.4-14.7,40.3-14.7s30.4,4.9,40.3,14.7c9.8,9.7,14.8,23.7,14.8,41.4v31.7c0,17.6-5,31.5-14.9,41.3-9.9,9.8-23.4,14.8-40.1,14.8ZM658.6,114.1c-9.5,0-17.1,2.7-22.6,8.1-5.5,5.4-8.3,13.4-8.3,23.8v36.7c0,10.5,2.8,18.5,8.3,23.8,5.5,5.4,13.1,8.1,22.6,8.1s17.3-2.7,22.7-8.1c5.4-5.4,8.2-13.4,8.2-23.9v-36.7c0-10.5-2.8-18.5-8.2-23.9-5.4-5.4-13.1-8.1-22.7-8.1Z"/>
+    <path fill="#141310" d="M860.6,236.3c-16.7,0-30.2-5-40.1-14.8-9.9-9.8-14.9-23.7-14.9-41.3v-31.7c0-17.7,5-31.7,14.8-41.4,9.8-9.7,23.4-14.7,40.3-14.7s30.4,4.9,40.3,14.7c9.8,9.7,14.8,23.7,14.8,41.4v31.7c0,17.6-5,31.5-14.9,41.3-9.9,9.8-23.4,14.8-40.1,14.8ZM860.6,114.1c-9.5,0-17.1,2.7-22.6,8.1-5.5,5.4-8.3,13.4-8.3,23.8v36.7c0,10.5,2.8,18.5,8.3,23.8,5.5,5.4,13.1,8.1,22.6,8.1s17.3-2.7,22.7-8.1c5.4-5.4,8.2-13.4,8.2-23.9v-36.7c0-10.5-2.8-18.5-8.2-23.9-5.4-5.4-13.1-8.1-22.7-8.1Z"/>
+    <path fill="#141310" d="M773.9,234c-18,0-32.6-14.6-32.6-32.6V48.8h24.1v152.6c0,4.7,3.8,8.5,8.5,8.5h16.8v24.1h-16.8Z"/>
+    <path fill="#141310" d="M1036.2,234.3V81.4c0-4.7-3.8-8.5-8.5-8.5h-16.8v-24.1h16.8c18,0,32.6,14.6,32.6,32.6v152.9h-24.1Z"/>
+    <path fill="#141310" d="M978.6,234.3c-18,0-32.6-14.6-32.6-32.6v-85.1h-20.3v-22.1h20.3v-45.3h24.1v45.3h30.2v22.1h-30.2v85.1c0,4.7,3.8,8.5,8.5,8.5h21.7v24.1h-21.7Z"/>
+    <path fill="#141310" d="M51.5,49h12.2v-20.6h-12.2c-16,0-29,13-29,29v32.8h20.6v-32.8c0-4.7,3.8-8.4,8.4-8.4Z"/>
+    <path fill="#141310" d="M92.8,49h12.2v-20.6h-12.2c-16,0-29,13-29,29v12.2h20.6v-12.2c0-4.7,3.8-8.4,8.4-8.4Z"/>
+    <path fill="#141310" d="M249.3,57.4c0-16-13-29-29-29h-12.2v20.6h12.2c4.7,0,8.4,3.8,8.4,8.4v32.8h20.6v-32.8Z"/>
+    <path fill="#141310" d="M187.4,90.2v-20.6h-103.1v20.6h-41.2v20.6h-20.6v41.2c0,11.4,9.2,20.6,20.6,20.6h185.5c11.4,0,20.6-9.2,20.6-20.6v-41.2h-20.6v-20.6h-41.2ZM166.8,141.7c0-5.7-4.6-10.3-10.3-10.3s-10.3,4.6-10.3,10.3v10.3h-20.6v-20.6c0-11.4,9.2-20.6,20.6-20.6s20.6,9.2,20.6,20.6v10.3ZM228.7,141.7c0-5.7-4.6-10.3-10.3-10.3s-10.3,4.6-10.3,10.3v10.3h-20.6v-20.6c0-11.4,9.2-20.6,20.6-20.6s20.6,9.2,20.6,20.6v10.3Z"/>
+    <path fill="#141310" d="M208,57.4c0-16-13-29-29-29h-12.2v20.6h12.2c4.7,0,8.4,3.8,8.4,8.4v12.2h20.6v-12.2Z"/>
+    <rect fill="#141310" x="22.5" y="234.5" width="41.2" height="20.6"/>
+    <rect fill="#141310" x="84.3" y="234.5" width="164.9" height="20.6"/>
+    <rect fill="#141310" x="208" y="193.3" width="41.2" height="20.6"/>
+    <rect fill="#141310" x="22.5" y="193.3" width="164.9" height="20.6"/>
+</svg>
diff --git a/image/axolotl_logo_digital_white.svg b/image/axolotl_logo_digital_white.svg
new file mode 100644
index 0000000000..fe954a6fd1
--- /dev/null
+++ b/image/axolotl_logo_digital_white.svg
@@ -0,0 +1,11 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" version="1.1" viewBox="0 0 1113 283.5">
+    <path fill="#fff" d="M462.9,234.2l-12.1-48.8h-54.4l-12.1,48.8h-24.7l48.2-185h31.6l47.9,185h-24.4ZM445.7,164.8l-13.8-55.6c-2.7-10.7-4.8-19.7-6.3-26.9-.9-4.2-1.5-7.5-2-9.9-.5,2.5-1.2,5.8-2,9.9-1.5,7.1-3.6,16.1-6.3,26.7l-13.8,55.9h44.3Z"/>
+    <path fill="#fff" d="M596.1,234.2l-29.9-45.6c-1.2-1.9-2.4-4.1-3.5-6.5-.8-1.7-1.5-3.3-2.1-4.5-.6,1.3-1.4,2.8-2.3,4.5-1.3,2.4-2.6,4.6-4,6.5l-29.9,45.6h-28.5l49.5-71.9-46.5-67.9h28.5l27.6,43.1c1.2,1.9,2.3,3.9,3.4,6.1.7,1.4,1.3,2.7,1.9,3.8.5-1.1,1.1-2.4,1.8-3.8,1.1-2.2,2.2-4.2,3.4-6.1l27.6-43.1h28.5l-46.5,68.1,49.3,71.6h-28.5Z"/>
+    <path fill="#fff" d="M686.4,236.2c-16.7,0-30.2-5-40.1-14.8-9.9-9.8-14.9-23.7-14.9-41.3v-31.7c0-17.7,5-31.6,14.8-41.4,9.8-9.7,23.4-14.7,40.2-14.7s30.4,4.9,40.2,14.7c9.8,9.7,14.8,23.7,14.8,41.4v31.7c0,17.6-5,31.4-14.9,41.3-9.9,9.8-23.4,14.8-40.1,14.8ZM686.4,114.1c-9.5,0-17.1,2.7-22.6,8.1-5.5,5.4-8.3,13.4-8.3,23.8v36.7c0,10.5,2.8,18.5,8.3,23.8,5.5,5.4,13.1,8.1,22.6,8.1s17.3-2.7,22.7-8.1c5.4-5.4,8.2-13.4,8.2-23.8v-36.7c0-10.5-2.8-18.5-8.2-23.8-5.4-5.4-13.1-8.1-22.7-8.1Z"/>
+    <path fill="#fff" d="M888.3,236.2c-16.7,0-30.2-5-40.1-14.8-9.9-9.8-14.9-23.7-14.9-41.3v-31.7c0-17.7,5-31.6,14.8-41.4,9.8-9.7,23.4-14.7,40.2-14.7s30.4,4.9,40.2,14.7c9.8,9.7,14.8,23.7,14.8,41.4v31.7c0,17.6-5,31.4-14.9,41.3-9.9,9.8-23.4,14.8-40.1,14.8ZM888.3,114.1c-9.5,0-17.1,2.7-22.6,8.1-5.5,5.4-8.3,13.4-8.3,23.8v36.7c0,10.5,2.8,18.5,8.3,23.8,5.5,5.4,13.1,8.1,22.6,8.1s17.3-2.7,22.7-8.1c5.4-5.4,8.2-13.4,8.2-23.8v-36.7c0-10.5-2.8-18.5-8.2-23.8-5.4-5.4-13.1-8.1-22.7-8.1Z"/>
+    <path fill="#fff" d="M801.7,234c-18,0-32.6-14.6-32.6-32.6V48.8h24.1v152.5c0,4.7,3.8,8.5,8.5,8.5h16.7v24.1h-16.7Z"/>
+    <path fill="#fff" d="M1063.8,234.2V81.4c0-4.7-3.8-8.5-8.5-8.5h-16.7v-24.1h16.7c18,0,32.6,14.6,32.6,32.6v152.8h-24.1Z"/>
+    <path fill="#fff" d="M1006.2,234.2c-18,0-32.6-14.6-32.6-32.6v-85h-20.3v-22.1h20.3v-45.2h24.1v45.2h30.2v22.1h-30.2v85c0,4.7,3.8,8.5,8.5,8.5h21.7v24.1h-21.7Z"/>
+    <path fill="#fff" d="M160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM277.3,57.4c0-23.8-19.3-43.1-43.1-43.1h-12.2c-3.9,0-7.6,1.6-10.2,4.4-5.9-2.9-12.3-4.4-18.9-4.4h-12.2c-7.7,0-14.1,6.3-14.1,14.1v20.6c0,2.4.6,4.6,1.6,6.6h-37c1-2,1.6-4.2,1.6-6.6v-20.6c0-7.7-6.3-14.1-14.1-14.1h-12.2c-6.5,0-13,1.5-18.9,4.4-2.6-2.8-6.3-4.4-10.2-4.4h-12.2c-23.8,0-43.1,19.3-43.1,43.1v32.8c0,4.1,1.7,7.7,4.5,10.3-2.8,2.6-4.5,6.2-4.5,10.3v41.2c0,11,5.2,20.8,13.2,27.2-7.3.4-13.2,6.6-13.2,14v20.6c0,4.1,1.7,7.7,4.5,10.3-2.8,2.6-4.5,6.2-4.5,10.3v20.6c0,7.7,6.3,14.1,14.1,14.1h41.2c4.1,0,7.7-1.7,10.3-4.5,2.6,2.8,6.2,4.5,10.3,4.5h164.9c7.7,0,14.1-6.3,14.1-14.1v-20.6c0-4.1-1.7-7.7-4.5-10.3,2.8-2.6,4.5-6.2,4.5-10.3v-20.6c0-7.5-5.8-13.6-13.2-14,8-6.4,13.2-16.2,13.2-27.2v-41.2c0-4.1-1.7-7.7-4.5-10.3,2.8-2.6,4.5-6.2,4.5-10.3v-32.8ZM77.8,255.1h-41.2v-20.6h41.2v20.6ZM36.5,213.9v-20.6h164.9v20.6H36.5ZM263.3,255.1H98.4v-20.6h164.9v20.6ZM263.3,213.9h-41.2v-20.6h41.2v20.6ZM263.3,90.2h-20.6v20.6h20.6v41.2c0,11.4-9.2,20.6-20.6,20.6H57.2c-11.4,0-20.6-9.2-20.6-20.6v-41.2h20.6v-20.6h-20.6v-32.8c0-16,13-29,29-29h12.2v20.6h-12.2c-4.7,0-8.4,3.8-8.4,8.4v32.8h41.2v-20.6h-20.6v-12.2c0-16,13-29,29-29h12.2v20.6h-12.2c-4.7,0-8.4,3.8-8.4,8.4v12.2h103.1v-12.2c0-4.7-3.8-8.4-8.4-8.4h-12.2v-20.6h12.2c16,0,29,13,29,29v12.2h-20.6v20.6h41.2v-32.8c0-4.7-3.8-8.4-8.4-8.4h-12.2v-20.6h12.2c16,0,29,13,29,29v32.8ZM201.4,152h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6s-20.6,9.2-20.6,20.6v20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM222,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM160.2,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6Z"/>
+</svg>
diff --git a/image/axolotl_symbol_digital_black.svg b/image/axolotl_symbol_digital_black.svg
new file mode 100644
index 0000000000..6e1f2960d4
--- /dev/null
+++ b/image/axolotl_symbol_digital_black.svg
@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" version="1.1" viewBox="0 0 283.5 283.5">
+  <defs>
+    <style>
+      .cls-1 {
+        fill: #141310;
+      }
+    </style>
+  </defs>
+  <!-- Generator: Adobe Illustrator 28.7.1, SVG Export Plug-In . SVG Version: 1.2.0 Build 142)  -->
+  <g>
+    <g id="Layer_1">
+      <g>
+        <path class="cls-1" d="M46.9,37.4h13.7V14.2h-13.7c-18,0-32.7,14.6-32.7,32.7v36.9h23.2v-36.9c0-5.2,4.2-9.5,9.5-9.5Z"/>
+        <path class="cls-1" d="M93.2,37.4h13.7V14.2h-13.7c-18,0-32.7,14.6-32.7,32.7v13.7h23.2v-13.7c0-5.2,4.2-9.5,9.5-9.5Z"/>
+        <path class="cls-1" d="M269.3,46.9c0-18-14.6-32.7-32.7-32.7h-13.7v23.2h13.7c5.2,0,9.5,4.2,9.5,9.5v36.9h23.2v-36.9Z"/>
+        <path class="cls-1" d="M199.7,83.8v-23.2h-116v23.2h-46.4v23.2H14.2v46.4c0,12.8,10.4,23.2,23.2,23.2h208.7c12.8,0,23.2-10.4,23.2-23.2v-46.4h-23.2v-23.2h-46.4ZM176.5,141.7c0-6.4-5.2-11.6-11.6-11.6s-11.6,5.2-11.6,11.6v11.6h-23.2v-23.2c0-12.8,10.4-23.2,23.2-23.2s23.2,10.4,23.2,23.2v11.6ZM246.1,141.7c0-6.4-5.2-11.6-11.6-11.6s-11.6,5.2-11.6,11.6v11.6h-23.2v-23.2c0-12.8,10.4-23.2,23.2-23.2s23.2,10.4,23.2,23.2v11.6Z"/>
+        <path class="cls-1" d="M222.9,46.9c0-18-14.6-32.7-32.7-32.7h-13.7v23.2h13.7c5.2,0,9.5,4.2,9.5,9.5v13.7h23.2v-13.7Z"/>
+        <rect class="cls-1" x="14.2" y="246.1" width="46.4" height="23.2"/>
+        <rect class="cls-1" x="83.8" y="246.1" width="185.5" height="23.2"/>
+        <rect class="cls-1" x="222.9" y="199.7" width="46.4" height="23.2"/>
+        <rect class="cls-1" x="14.2" y="199.7" width="185.5" height="23.2"/>
+      </g>
+    </g>
+  </g>
+</svg>
diff --git a/image/axolotl_symbol_digital_white.svg b/image/axolotl_symbol_digital_white.svg
new file mode 100644
index 0000000000..2752ffb0f8
--- /dev/null
+++ b/image/axolotl_symbol_digital_white.svg
@@ -0,0 +1,16 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" version="1.1" viewBox="0 0 283.5 283.5">
+  <defs>
+    <style>
+      .cls-1 {
+        fill: #fff;
+      }
+    </style>
+  </defs>
+  <!-- Generator: Adobe Illustrator 28.7.1, SVG Export Plug-In . SVG Version: 1.2.0 Build 142)  -->
+  <g>
+    <g id="Layer_1">
+      <path class="cls-1" d="M152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM269.3,57.3c0-23.8-19.4-43.1-43.1-43.1h-12.2c-3.9,0-7.6,1.6-10.2,4.4-5.9-2.9-12.3-4.4-18.9-4.4h-12.2c-7.8,0-14.1,6.3-14.1,14.1v20.6c0,2.4.6,4.6,1.6,6.6h-37c1-2,1.6-4.2,1.6-6.6v-20.6c0-7.8-6.3-14.1-14.1-14.1h-12.2c-6.6,0-13,1.5-18.9,4.4-2.6-2.8-6.3-4.4-10.2-4.4h-12.2c-23.8,0-43.1,19.4-43.1,43.1v32.8c0,4.1,1.7,7.7,4.5,10.3-2.8,2.6-4.5,6.2-4.5,10.3v41.3c0,11,5.2,20.9,13.2,27.2-7.4.4-13.2,6.6-13.2,14v20.6c0,4.1,1.7,7.7,4.5,10.3-2.8,2.6-4.5,6.2-4.5,10.3v20.6c0,7.8,6.3,14.1,14.1,14.1h41.3c4.1,0,7.7-1.7,10.3-4.5,2.6,2.8,6.2,4.5,10.3,4.5h165.1c7.8,0,14.1-6.3,14.1-14.1v-20.6c0-4.1-1.7-7.7-4.5-10.3,2.8-2.6,4.5-6.2,4.5-10.3v-20.6c0-7.5-5.9-13.6-13.2-14,8-6.4,13.2-16.2,13.2-27.2v-41.3c0-4.1-1.7-7.7-4.5-10.3,2.8-2.6,4.5-6.2,4.5-10.3v-32.8ZM69.5,255.2H28.2v-20.6h41.3v20.6ZM28.2,214v-20.6h165.1v20.6H28.2ZM255.2,255.2H90.1v-20.6h165.1v20.6ZM255.2,214h-41.3v-20.6h41.3v20.6ZM255.2,90.1h-20.6v20.6h20.6v41.3c0,11.4-9.2,20.6-20.6,20.6H48.9c-11.4,0-20.6-9.2-20.6-20.6v-41.3h20.6v-20.6h-20.6v-32.8c0-16.1,13-29.1,29.1-29.1h12.2v20.6h-12.2c-4.7,0-8.4,3.8-8.4,8.4v32.8h41.3v-20.6h-20.6v-12.2c0-16.1,13-29.1,29.1-29.1h12.2v20.6h-12.2c-4.7,0-8.4,3.8-8.4,8.4v12.2h103.2v-12.2c0-4.7-3.8-8.4-8.4-8.4h-12.2v-20.6h12.2c16.1,0,29.1,13,29.1,29.1v12.2h-20.6v20.6h41.3v-32.8c0-4.7-3.8-8.4-8.4-8.4h-12.2v-20.6h12.2c16.1,0,29.1,13,29.1,29.1v32.8ZM193.3,152h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6s-20.6,9.2-20.6,20.6v20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM214,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6ZM152,110.8c-11.4,0-20.6,9.2-20.6,20.6v20.6h20.6v-10.3c0-5.7,4.6-10.3,10.3-10.3s10.3,4.6,10.3,10.3v-10.3c0-11.4-9.2-20.6-20.6-20.6Z"/>
+    </g>
+  </g>
+</svg>
diff --git a/image/axolotl_wordmark_digital_black.svg b/image/axolotl_wordmark_digital_black.svg
new file mode 100644
index 0000000000..d8041b1a85
--- /dev/null
+++ b/image/axolotl_wordmark_digital_black.svg
@@ -0,0 +1,17 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" version="1.1" viewBox="0 0 765.4 212.6">
+  <!-- Generator: Adobe Illustrator 28.7.1, SVG Export Plug-In . SVG Version: 1.2.0 Build 142)  -->
+  <g>
+    <g id="Layer_1">
+      <g>
+        <path d="M121.6,198.1l-12.1-48.8h-54.4l-12.1,48.8h-24.7L66.6,12.9h31.6l47.9,185.1h-24.5ZM104.4,128.6l-13.8-55.6c-2.7-10.7-4.8-19.7-6.3-26.9-.9-4.2-1.5-7.5-2-9.9-.5,2.5-1.2,5.8-2,9.9-1.5,7.1-3.6,16.1-6.3,26.7l-13.8,55.9h44.3Z"/>
+        <path d="M254.9,198.1l-29.9-45.6c-1.2-1.9-2.4-4.1-3.5-6.5-.8-1.7-1.5-3.3-2.1-4.5-.6,1.3-1.4,2.8-2.3,4.5-1.3,2.4-2.6,4.6-4,6.5l-29.9,45.6h-28.5l49.6-71.9-46.5-67.9h28.5l27.6,43.1c1.2,1.9,2.3,3.9,3.4,6.1.7,1.4,1.4,2.7,1.9,3.8.5-1.1,1.1-2.4,1.8-3.8,1.1-2.2,2.2-4.2,3.4-6.1l27.6-43.1h28.5l-46.5,68.2,49.3,71.7h-28.5Z"/>
+        <path d="M345.2,200.1c-16.7,0-30.2-5-40.1-14.8-9.9-9.8-14.9-23.7-14.9-41.3v-31.7c0-17.7,5-31.7,14.8-41.4,9.8-9.7,23.4-14.7,40.3-14.7s30.4,4.9,40.3,14.7c9.8,9.7,14.8,23.7,14.8,41.4v31.7c0,17.6-5,31.5-14.9,41.3-9.9,9.8-23.4,14.8-40.1,14.8ZM345.2,77.8c-9.5,0-17.1,2.7-22.6,8.1-5.5,5.4-8.3,13.4-8.3,23.8v36.7c0,10.5,2.8,18.5,8.3,23.8,5.5,5.4,13.1,8.1,22.6,8.1s17.3-2.7,22.7-8.1c5.4-5.4,8.2-13.4,8.2-23.9v-36.7c0-10.5-2.8-18.5-8.2-23.9-5.4-5.4-13.1-8.1-22.7-8.1Z"/>
+        <path d="M547.3,200.1c-16.7,0-30.2-5-40.1-14.8-9.9-9.8-14.9-23.7-14.9-41.3v-31.7c0-17.7,5-31.7,14.8-41.4,9.8-9.7,23.4-14.7,40.3-14.7s30.4,4.9,40.3,14.7c9.8,9.7,14.8,23.7,14.8,41.4v31.7c0,17.6-5,31.5-14.9,41.3-9.9,9.8-23.4,14.8-40.1,14.8ZM547.3,77.8c-9.5,0-17.1,2.7-22.6,8.1-5.5,5.4-8.3,13.4-8.3,23.8v36.7c0,10.5,2.8,18.5,8.3,23.8,5.5,5.4,13.1,8.1,22.6,8.1s17.3-2.7,22.7-8.1c5.4-5.4,8.2-13.4,8.2-23.9v-36.7c0-10.5-2.8-18.5-8.2-23.9-5.4-5.4-13.1-8.1-22.7-8.1Z"/>
+        <path d="M460.6,197.8c-18,0-32.6-14.6-32.6-32.6V12.5h24.1v152.6c0,4.7,3.8,8.5,8.5,8.5h16.8v24.1h-16.8Z"/>
+        <path d="M722.8,198.1V45.2c0-4.7-3.8-8.5-8.5-8.5h-16.8V12.5h16.8c18,0,32.6,14.6,32.6,32.6v152.9h-24.1Z"/>
+        <path d="M665.2,198.1c-18,0-32.6-14.6-32.6-32.6v-85.1h-20.3v-22.1h20.3V12.9h24.1v45.3h30.2v22.1h-30.2v85.1c0,4.7,3.8,8.5,8.5,8.5h21.7v24.1h-21.7Z"/>
+      </g>
+    </g>
+  </g>
+</svg>
diff --git a/image/axolotl_wordmark_digital_white.svg b/image/axolotl_wordmark_digital_white.svg
new file mode 100644
index 0000000000..7cada6ae00
--- /dev/null
+++ b/image/axolotl_wordmark_digital_white.svg
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" version="1.1" viewBox="0 0 765.4 212.6">
+  <defs>
+    <style>
+      .cls-1 {
+        fill: #fff;
+      }
+    </style>
+  </defs>
+  <!-- Generator: Adobe Illustrator 28.7.1, SVG Export Plug-In . SVG Version: 1.2.0 Build 142)  -->
+  <g>
+    <g id="Layer_1">
+      <g>
+        <path class="cls-1" d="M121.6,198.1l-12.1-48.8h-54.4l-12.1,48.8h-24.7L66.6,12.9h31.6l47.9,185.1h-24.5ZM104.4,128.6l-13.8-55.6c-2.7-10.7-4.8-19.7-6.3-26.9-.9-4.2-1.5-7.5-2-9.9-.5,2.5-1.2,5.8-2,9.9-1.5,7.1-3.6,16.1-6.3,26.7l-13.8,55.9h44.3Z"/>
+        <path class="cls-1" d="M254.9,198.1l-29.9-45.6c-1.2-1.9-2.4-4.1-3.5-6.5-.8-1.7-1.5-3.3-2.1-4.5-.6,1.3-1.4,2.8-2.3,4.5-1.3,2.4-2.6,4.6-4,6.5l-29.9,45.6h-28.5l49.6-71.9-46.5-67.9h28.5l27.6,43.1c1.2,1.9,2.3,3.9,3.4,6.1.7,1.4,1.4,2.7,1.9,3.8.5-1.1,1.1-2.4,1.8-3.8,1.1-2.2,2.2-4.2,3.4-6.1l27.6-43.1h28.5l-46.5,68.2,49.3,71.7h-28.5Z"/>
+        <path class="cls-1" d="M345.2,200.1c-16.7,0-30.2-5-40.1-14.8-9.9-9.8-14.9-23.7-14.9-41.3v-31.7c0-17.7,5-31.7,14.8-41.4,9.8-9.7,23.4-14.7,40.3-14.7s30.4,4.9,40.3,14.7c9.8,9.7,14.8,23.7,14.8,41.4v31.7c0,17.6-5,31.5-14.9,41.3-9.9,9.8-23.4,14.8-40.1,14.8ZM345.2,77.8c-9.5,0-17.1,2.7-22.6,8.1-5.5,5.4-8.3,13.4-8.3,23.8v36.7c0,10.5,2.8,18.5,8.3,23.8,5.5,5.4,13.1,8.1,22.6,8.1s17.3-2.7,22.7-8.1c5.4-5.4,8.2-13.4,8.2-23.9v-36.7c0-10.5-2.8-18.5-8.2-23.9-5.4-5.4-13.1-8.1-22.7-8.1Z"/>
+        <path class="cls-1" d="M547.3,200.1c-16.7,0-30.2-5-40.1-14.8-9.9-9.8-14.9-23.7-14.9-41.3v-31.7c0-17.7,5-31.7,14.8-41.4,9.8-9.7,23.4-14.7,40.3-14.7s30.4,4.9,40.3,14.7c9.8,9.7,14.8,23.7,14.8,41.4v31.7c0,17.6-5,31.5-14.9,41.3-9.9,9.8-23.4,14.8-40.1,14.8ZM547.3,77.8c-9.5,0-17.1,2.7-22.6,8.1-5.5,5.4-8.3,13.4-8.3,23.8v36.7c0,10.5,2.8,18.5,8.3,23.8,5.5,5.4,13.1,8.1,22.6,8.1s17.3-2.7,22.7-8.1c5.4-5.4,8.2-13.4,8.2-23.9v-36.7c0-10.5-2.8-18.5-8.2-23.9-5.4-5.4-13.1-8.1-22.7-8.1Z"/>
+        <path class="cls-1" d="M460.6,197.8c-18,0-32.6-14.6-32.6-32.6V12.5h24.1v152.6c0,4.7,3.8,8.5,8.5,8.5h16.8v24.1h-16.8Z"/>
+        <path class="cls-1" d="M722.8,198.1V45.2c0-4.7-3.8-8.5-8.5-8.5h-16.8V12.5h16.8c18,0,32.6,14.6,32.6,32.6v152.9h-24.1Z"/>
+        <path class="cls-1" d="M665.2,198.1c-18,0-32.6-14.6-32.6-32.6v-85.1h-20.3v-22.1h20.3V12.9h24.1v45.3h30.2v22.1h-30.2v85.1c0,4.7,3.8,8.5,8.5,8.5h21.7v24.1h-21.7Z"/>
+      </g>
+    </g>
+  </g>
+</svg>
diff --git a/index.qmd b/index.qmd
index ed23d235ad..01572a8bec 100644
--- a/index.qmd
+++ b/index.qmd
@@ -1,7 +1,7 @@
 ---
-toc-location: right-body
-toc-title: Table Of Contents
-toc-expand: 2
+# toc-location: right-body
+# toc-title: Table Of Contents
+# toc-expand: 2
 ---
 
 ```{python}
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000000..82709debb3
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,282 @@
+[build-system]
+requires = ["setuptools>=64", "wheel", "setuptools_scm>=8"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "axolotl"
+dynamic = ["version"]
+description = "LLM Trainer"
+readme = "README.md"
+requires-python = ">=3.10"
+# license = "Apache-2.0"
+
+dependencies = [
+    # Core ML stack
+    "torch>=2.11.0,<=2.12.1",
+    "packaging==26.0",
+    "huggingface_hub==1.17.0",
+    "peft==0.19.1",
+    "tokenizers==0.22.2",
+    "transformers==5.12.1",
+    "accelerate==1.13.0",
+    "datasets==4.8.4",
+    "trl==1.7.0",
+    "hf_xet==1.4.3",
+    "kernels==0.13.0",
+    "trackio==0.19.0",
+    "typing-extensions==4.15.0",
+    "optimum==1.16.2",
+    "hf_transfer",
+    "sentencepiece",
+    "gradio==6.10.0",
+    "modal==1.3.0.post1",
+    "pydantic==2.12.5",
+    "addict",
+    "fire",
+    "PyYAML==6.0.3",
+    "requests",
+    "wandb",
+    "einops",
+    "colorama",
+    "numba==0.63.1",
+    "numpy==2.3.5",
+    "typer==0.25.1",
+
+    # Evaluation & metrics
+    "evaluate==0.4.1",
+    "scipy",
+    "nvidia-ml-py==12.560.30",
+    "art",
+    "tensorboard",
+    "python-dotenv==1.0.1",
+
+    # Remote filesystems
+    "s3fs==2025.10.0",
+    "gcsfs==2025.10.0",
+    "adlfs==2026.2.0",
+    "ocifs==1.3.2",
+
+    "zstandard==0.22.0",
+    "fastcore",
+
+    # lm eval harness
+    "lm_eval==0.4.11",
+    "langdetect==1.0.9",
+    "immutabledict==4.2.0",
+    "antlr4-python3-runtime==4.13.2",
+
+    "schedulefree==1.4.1",
+    "openenv-core==0.1.0",
+
+    # Axolotl contribs
+    "axolotl-contribs-lgpl==0.0.7",
+    "axolotl-contribs-mit==0.0.6",
+
+    # Telemetry
+    "posthog==6.7.11",
+
+    "mistral-common==1.11.5",
+
+    # Platform-specific (Linux only)
+    "bitsandbytes==0.49.1 ; sys_platform != 'darwin'",
+    "triton>=3.4.0 ; sys_platform != 'darwin'",
+    "xformers>=0.0.33.post2 ; sys_platform != 'darwin' and platform_machine != 'aarch64'",
+    "liger-kernel==0.8.0 ; sys_platform != 'darwin'",
+    "torchao==0.17.0 ; sys_platform != 'darwin' and platform_machine != 'aarch64'",
+
+    # Architecture-specific
+    "fla-core==0.4.1 ; platform_machine != 'aarch64'",
+    "flash-linear-attention==0.4.1 ; platform_machine != 'aarch64'",
+]
+
+[project.optional-dependencies]
+flash-attn = ["flash-attn==2.8.3"]
+ring-flash-attn = [
+    "flash-attn==2.8.3",
+    "ring-flash-attn>=0.1.7",
+]
+deepspeed = [
+    "deepspeed>=0.18.6,<0.19.0",
+    "deepspeed-kernels",
+]
+mamba-ssm = [
+    "mamba-ssm==1.2.0.post1",
+    "causal_conv1d",
+]
+auto-gptq = [
+    "auto-gptq==0.5.1",
+]
+mlflow = [
+    "mlflow",
+]
+galore = [
+    "galore_torch",
+]
+apollo = [
+    "apollo-torch",
+]
+optimizers = [
+    "galore_torch",
+    "apollo-torch",
+    "lomo-optim==0.1.1",
+    "torch-optimi==0.2.1",
+    "came_pytorch==0.1.3",
+    "q-galore-torch==1.0",
+]
+ray = [
+    "ray[train]>=2.52.1",
+]
+vllm = [
+    "vllm>=0.15.0",
+]
+llmcompressor = [
+    "llmcompressor>=0.10.0",
+]
+fbgemm-gpu = ["fbgemm-gpu-genai>=1.3.0"]
+opentelemetry = [
+    "opentelemetry-api",
+    "opentelemetry-sdk",
+    "opentelemetry-exporter-prometheus",
+    "prometheus-client",
+]
+
+[dependency-groups]
+dev = [
+    "black",
+    "mypy",
+    "pre-commit",
+    "types-requests",
+    "quartodoc",
+    "jupyter",
+    "blobfile",
+    "tiktoken",
+]
+test = [
+    "codecov",
+    "codecov-cli",
+    "pytest",
+    "pytest-cov",
+    "pytest-retry",
+    "pytest-sugar",
+    "pytest-xdist",
+    "tbparse",
+]
+
+[project.scripts]
+axolotl = "axolotl.cli.main:main"
+
+[project.urls]
+Homepage = "https://axolotl.ai/"
+Documentation = "https://docs.axolotl.ai/"
+Repository = "https://github.com/axolotl-ai-cloud/axolotl.git"
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.package-data]
+axolotl = [
+    "**/*.yaml",
+    "**/*.jinja",
+    "integrations/kernels/libs/scattermoe_lora/metadata.json",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.dynamic]
+version = { file = "VERSION" }
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "C90", "B", "I"]
+ignore = [
+    "E203",  # Whitespace before ':'
+    "E501",  # Line too long
+    "C901",  # Too complex
+    "B019",  # Use of functools.cache on methods
+    "E722",  # Bare except
+    "F821",  # Undefined name (for dynamic exec)
+    "E731",  # Lambda assignment (idiomatic for the kernel recipe/heuristic closures)
+    "E741",  # Ambiguous variable name (single-letter math/kernel vars: I, O, l)
+]
+
+[tool.ruff.lint.isort]
+known-third-party = ["wandb", "comet_ml"]
+known-local-folder = ["src", "tests"]
+# Black-compatible isort settings
+force-single-line = false
+combine-as-imports = true
+split-on-trailing-comma = true
+
+[tool.ruff.format]
+# Use black's formatting style exactly
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
+docstring-code-format = false
+
+[tool.pytest.ini_options]
+addopts = "-m 'not slow'"
+norecursedirs = [
+    ".claude",
+    ".git",
+    ".mypy_cache",
+    ".nox",
+    ".pytest_cache",
+    ".ruff_cache",
+    ".scratch_fp4",
+    ".tox",
+    ".venv",
+    "build",
+    "dist",
+    "outputs",
+    "unsloth_compiled_cache",
+]
+markers = [
+    "slow: marks tests as slow",
+    "perf: marks tests as performance/memory regression guards (Blackwell only, large E)",
+]
+
+# UV specific configuration
+[tool.uv]
+conflicts = [
+    [
+        { package = "axolotl" },
+        { extra = "vllm" },
+    ],
+    [
+        { package = "axolotl" },
+        { extra = "flash-attn" },
+    ],
+    [
+        { package = "axolotl" },
+        { extra = "ring-flash-attn" },
+    ],
+    [
+        { package = "axolotl" },
+        { extra = "mamba-ssm" },
+    ],
+    [
+        { package = "axolotl" },
+        { extra = "auto-gptq" },
+    ],
+    [
+        { package = "axolotl" },
+        { extra = "fbgemm-gpu" },
+    ],
+    [
+        { package = "axolotl" },
+        { extra = "llmcompressor" },
+    ],
+]
+
+[tool.uv.extra-build-dependencies]
+mamba-ssm = [{ requirement = "torch", match-runtime = true }]
+causal-conv1d = [{ requirement = "torch", match-runtime = true }]
+flash-attn = [{ requirement = "torch", match-runtime = true }]
+deepspeed = [{ requirement = "torch", match-runtime = true }]
+auto-gptq = [{ requirement = "torch", match-runtime = true }]
diff --git a/requirements-dev.txt b/requirements-dev.txt
deleted file mode 100644
index 4b5df167b6..0000000000
--- a/requirements-dev.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-pre-commit
-black
-mypy
-types-requests
diff --git a/requirements-tests.txt b/requirements-tests.txt
deleted file mode 100644
index e079f8a603..0000000000
--- a/requirements-tests.txt
+++ /dev/null
@@ -1 +0,0 @@
-pytest
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index b15a28c901..0000000000
--- a/requirements.txt
+++ /dev/null
@@ -1,44 +0,0 @@
---extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
-packaging==23.2
-peft==0.10.0
-transformers @ git+https://github.com/huggingface/transformers.git@43d17c18360ac9c3d3491389328e2fe55fe8f9ce
-tokenizers==0.15.0
-bitsandbytes==0.43.0
-accelerate==0.28.0
-deepspeed==0.13.1
-pydantic==2.6.3
-addict
-fire
-PyYAML>=6.0
-requests
-datasets==2.15.0
-flash-attn==2.5.5
-sentencepiece
-wandb
-einops
-xformers==0.0.22
-optimum==1.16.2
-hf_transfer
-colorama
-numba
-numpy>=1.24.4
-# qlora things
-evaluate==0.4.1
-scipy
-scikit-learn==1.2.2
-pynvml
-art
-fschat @ git+https://github.com/lm-sys/FastChat.git@27a05b04a35510afb1d767ae7e5990cbd278f8fe
-gradio==3.50.2
-tensorboard
-
-mamba-ssm==1.2.0.post1
-
-# remote filesystems
-s3fs
-gcsfs
-# adlfs
-
-trl==0.8.5
-zstandard==0.22.0
-fastcore
diff --git a/scripts/analyze_profile.py b/scripts/analyze_profile.py
new file mode 100644
index 0000000000..df922502fe
--- /dev/null
+++ b/scripts/analyze_profile.py
@@ -0,0 +1,1518 @@
+#!/usr/bin/env python3
+"""
+Axolotl Training Profiler Analyzer
+===================================
+
+Analyzes PyTorch profiler output from axolotl training runs (profiler_steps config).
+Produces breakdowns by CUDA kernel category, identifies bottlenecks, and optionally
+compares two traces (e.g. before/after optimization).
+
+Supports both:
+  - profiler_trace.json (torch.profiler Chrome trace -- timing analysis)
+  - snapshot.pickle (torch.cuda.memory._snapshot -- memory analysis)
+
+Usage:
+  # Analyze a single trace
+  python analyze_profile.py outputs/qwen35_moe_profile/
+
+  # Compare before vs after
+  python analyze_profile.py outputs/before/ --compare outputs/after/
+
+  # Include step 0 (warmup/compilation) in analysis
+  python analyze_profile.py outputs/run/ --include-warmup
+
+  # Memory-only analysis
+  python analyze_profile.py outputs/run/ --memory-only
+
+  # Quick mode (first 2M events only, for large traces)
+  python analyze_profile.py outputs/run/ --quick
+"""
+
+import argparse
+import json
+import pickle  # nosec B403
+import time
+from collections import defaultdict
+from pathlib import Path
+
+# ---- Kernel categorization ------------------------------------------------
+
+KERNEL_CATEGORIES = [
+    # ScatterMoE -- ordered most specific first
+    (
+        "ScatterMoE bwd LoRA (split)",
+        ["_group_bwd_lora_split", "_group_bwd_da", "_group_bwd_db"],
+    ),
+    ("ScatterMoE bwd LoRA (fused)", ["_group_bwd_lora"]),
+    ("ScatterMoE bwd dX", ["_scatter2scatter_lora_dx"]),
+    ("ScatterMoE fwd", ["_scatter2scatter_lora"]),
+    # Quantization
+    ("BnB Dequantization", ["dequantize", "kDequantizeBlockwise"]),
+    # Attention
+    ("Flash Attention", ["flash", "fmha"]),
+    # Loss
+    ("CCE Loss", ["_cce_"]),
+    # LoRA fused kernels (autograd.Function based)
+    ("LoRA QKV Kernel", ["lora_qkv"]),
+    ("LoRA O Kernel", ["lora_o"]),
+    ("LoRA MLP Kernel", ["lora_mlp"]),
+    # LoRA activation kernels (SwiGLU/GEGLU Triton kernels)
+    ("LoRA Activation (SwiGLU/GEGLU)", ["swiglu", "geglu"]),
+    # DoRA weight norm
+    ("DoRA Weight Norm", ["linalg_norm", "dora_scale"]),
+    # Compute
+    ("GEMM/CUTLASS", ["cutlass", "gemm", "gemv", "cublas"]),
+    ("Triton (norms etc)", ["triton"]),
+    ("Conv1d", ["conv1d", "causal_conv"]),
+    # Optimizer
+    ("Optimizer", ["adam", "optim"]),
+    # Dtype conversion (fp32→bf16 LoRA matrix casts etc)
+    ("Dtype Conversion", ["_to_copy", "to_copy"]),
+    # Memory
+    ("Elementwise/Fill", ["fill", "elementwise", "cast", "copy_kernel"]),
+    ("Memory ops", ["memcpy", "memset"]),
+    # Routing
+    ("TopK/Sort", ["topk", "sort"]),
+    ("Index/Gather/Scatter", ["index", "gather", "scatter"]),
+]
+
+# Categories to keep when pre-filtering during streaming load
+_KEEP_CATS = {"kernel", "gpu_memcpy", "cpu_op", "python_function", "ac2g", "Runtime"}
+
+
+def categorize_kernel(name):
+    nl = name.lower()
+    for cat_name, patterns in KERNEL_CATEGORIES:
+        if any(p in nl for p in patterns):
+            return cat_name
+    return "Other"
+
+
+# ---- Trace loading ---------------------------------------------------------
+
+
+def _try_ijson_load(trace_file, quick=False, max_events=2_000_000):
+    """Stream-parse trace JSON with ijson. Returns filtered events list."""
+    try:
+        import ijson
+    except ImportError:
+        return None
+
+    events = []
+    count = 0
+    with open(trace_file, "rb") as f:
+        for ev in ijson.items(f, "traceEvents.item"):
+            count += 1
+            cat = ev.get("cat", "")
+            # Pre-filter: only keep categories we care about
+            if cat in _KEEP_CATS or ev.get("ph") == "M":
+                events.append(ev)
+            if quick and count >= max_events:
+                print(f"  --quick: stopped after {count:,} events")
+                break
+    return events, count
+
+
+def load_trace(path, quick=False):
+    trace_file = (
+        Path(path) / "profiler_trace.json" if Path(path).is_dir() else Path(path)
+    )
+    if not trace_file.exists():
+        return None
+    size_gb = trace_file.stat().st_size / 1e9
+    print(f"Loading {trace_file.name} ({size_gb:.1f} GB)...")
+    t0 = time.monotonic()
+
+    # Try streaming parser first for large files
+    if size_gb > 0.5:
+        result = _try_ijson_load(trace_file, quick=quick)
+        if result is not None:
+            events, total_count = result
+            elapsed = time.monotonic() - t0
+            print(
+                f"  {total_count:,} total events, {len(events):,} kept "
+                f"(streamed in {elapsed:.1f}s)"
+            )
+            return events
+
+    # Fallback: standard json.load
+    with open(trace_file) as f:
+        data = json.load(f)
+    all_events = data.get("traceEvents", [])
+    elapsed = time.monotonic() - t0
+
+    if quick:
+        all_events = all_events[:2_000_000]
+        print(f"  --quick: limited to first {len(all_events):,} events")
+
+    # Pre-filter
+    events = [
+        ev
+        for ev in all_events
+        if ev.get("cat", "") in _KEEP_CATS or ev.get("ph") == "M"
+    ]
+    print(
+        f"  {len(all_events):,} total events, {len(events):,} kept "
+        f"(loaded in {elapsed:.1f}s)"
+    )
+    return events
+
+
+# ---- Trace analysis -------------------------------------------------------
+
+
+def _estimate_n_steps(cuda_events):
+    """Estimate the number of training steps from CUDA event timestamps.
+
+    Detects step boundaries by looking for large gaps (>2x median gap) in the
+    sorted timestamp sequence of CUDA kernels.
+    """
+    if len(cuda_events) < 100:
+        return 1
+    timestamps = sorted(float(ev.get("ts", 0)) for ev in cuda_events)
+    # Compute gaps between consecutive events
+    gaps = [timestamps[i + 1] - timestamps[i] for i in range(len(timestamps) - 1)]
+    if not gaps:
+        return 1
+    median_gap = sorted(gaps)[len(gaps) // 2]
+    # A step boundary has a gap much larger than the median inter-kernel gap
+    threshold = max(median_gap * 50, 100_000)  # at least 100ms
+    n_boundaries = sum(1 for g in gaps if g > threshold)
+    return max(n_boundaries + 1, 1)
+
+
+def analyze_trace(events, skip_warmup=True):
+    cuda_events = [
+        ev
+        for ev in events
+        if ev.get("ph") == "X" and ev.get("cat") in ("kernel", "gpu_memcpy")
+    ]
+    if not cuda_events:
+        print("  No CUDA kernel events found!")
+        return None
+
+    cutoff_ts = None
+
+    if skip_warmup and len(cuda_events) > 1000:
+        timestamps = sorted(set(float(ev.get("ts", 0)) for ev in cuda_events))
+        min_ts, max_ts = timestamps[0], timestamps[-1]
+        total_span = max_ts - min_ts
+
+        # Step 0 is warmup (Triton compilation + autotune). It's typically
+        # the slowest step by far. Use 45% of wall-clock as cutoff -- step 0
+        # usually takes >50% of total time when it includes compilation.
+        cutoff_ts = min_ts + total_span * 0.45
+        before = len(cuda_events)
+        cuda_events = [ev for ev in cuda_events if float(ev.get("ts", 0)) > cutoff_ts]
+        print(f"  Excluding step 0 (warmup): {before:,} -> {len(cuda_events):,} events")
+
+    n_steps_profiled = _estimate_n_steps(cuda_events)
+
+    # Aggregate by kernel (cast to float to handle ijson Decimal values)
+    kernel_stats = defaultdict(lambda: {"total_us": 0.0, "count": 0, "max_us": 0.0})
+    for ev in cuda_events:
+        name = ev.get("name", "unknown")
+        dur = float(ev.get("dur", 0))
+        kernel_stats[name]["total_us"] += dur
+        kernel_stats[name]["count"] += 1
+        kernel_stats[name]["max_us"] = max(kernel_stats[name]["max_us"], dur)
+
+    total_cuda = sum(v["total_us"] for v in kernel_stats.values())
+
+    # Group by category
+    cat_stats = defaultdict(lambda: {"total_us": 0, "count": 0})
+    for name, info in kernel_stats.items():
+        cat = categorize_kernel(name)
+        cat_stats[cat]["total_us"] += info["total_us"]
+        cat_stats[cat]["count"] += info["count"]
+
+    # Fill/zero_ analysis: find FillFunctor kernels and group by tensor size
+    fill_by_size = defaultdict(lambda: {"total_us": 0, "count": 0})
+    for ev in cuda_events:
+        name = ev.get("name", "")
+        if "FillFunctor" not in name and "fill" not in name.lower():
+            continue
+        # Extract Input Dims from args if present
+        args = ev.get("args", {})
+        input_dims = args.get("Input Dims", args.get("input_dims", "unknown"))
+        if isinstance(input_dims, list):
+            input_dims = str(input_dims)
+        dur = float(ev.get("dur", 0))
+        fill_by_size[input_dims]["total_us"] += dur
+        fill_by_size[input_dims]["count"] += 1
+
+    # CPU op analysis for wall-clock estimation (apply same warmup cutoff)
+    cpu_ops = [
+        ev
+        for ev in events
+        if ev.get("ph") == "X" and ev.get("cat") in ("cpu_op", "python_function")
+    ]
+    if cutoff_ts is not None:
+        cpu_ops = [ev for ev in cpu_ops if float(ev.get("ts", 0)) > cutoff_ts]
+    wall_clock_us = 0
+    if cpu_ops:
+        ts_sorted = sorted(cpu_ops, key=lambda e: float(e.get("ts", 0)))
+        min_cpu_ts = float(ts_sorted[0].get("ts", 0))
+        max_cpu_end = max(
+            float(e.get("ts", 0)) + float(e.get("dur", 0)) for e in ts_sorted
+        )
+        wall_clock_us = max_cpu_end - min_cpu_ts
+
+    return {
+        "total_cuda_us": total_cuda,
+        "n_steps": n_steps_profiled,
+        "categories": dict(cat_stats),
+        "kernel_stats": dict(kernel_stats),
+        "n_events": len(cuda_events),
+        "fill_by_size": dict(fill_by_size),
+        "wall_clock_us": wall_clock_us,
+    }
+
+
+def print_trace_analysis(result, label=""):
+    total = result["total_cuda_us"]
+    n = result["n_steps"]
+
+    if label:
+        print(f"\n{'=' * 75}")
+        print(f"  {label}")
+        print(f"{'=' * 75}")
+
+    print(
+        f"\n  CUDA kernel time: {total / 1e6:.2f}s over {n} steps "
+        f"(~{total / n / 1e6:.2f}s/step)"
+    )
+
+    if result.get("wall_clock_us"):
+        wc = result["wall_clock_us"]
+        print(
+            f"  Wall clock span:  {wc / 1e6:.2f}s over {n} steps "
+            f"(~{wc / n / 1e6:.2f}s/step)"
+        )
+
+    print(f"\n  {'Category':<40} {'Total':>9} {'%':>6} {'Count':>7} {'Per step':>9}")
+    print(f"  {'-' * 75}")
+    for cat, info in sorted(
+        result["categories"].items(), key=lambda x: x[1]["total_us"], reverse=True
+    ):
+        pct = info["total_us"] / total * 100
+        ps = info["total_us"] / n / 1000
+        print(
+            f"  {cat:<40} {info['total_us'] / 1000:>8.1f}ms {pct:>5.1f}% "
+            f"{info['count']:>7} {ps:>7.1f}ms"
+        )
+
+    print("\n  Top 15 individual kernels:")
+    for name, info in sorted(
+        result["kernel_stats"].items(), key=lambda x: x[1]["total_us"], reverse=True
+    )[:15]:
+        pct = info["total_us"] / total * 100
+        avg = info["total_us"] / info["count"] / 1000
+        print(
+            f"    {name[:62]:<62} {info['total_us'] / 1000:>7.1f}ms "
+            f"({pct:>4.1f}%) x{info['count']:<5} avg={avg:.3f}ms"
+        )
+
+    # Fill/zero_ breakdown
+    fill_data = result.get("fill_by_size", {})
+    if fill_data:
+        total_fill_us = sum(v["total_us"] for v in fill_data.values())
+        if total_fill_us > 0:
+            print(
+                f"\n  Fill/zero_ kernel breakdown by tensor size "
+                f"(total: {total_fill_us / 1000:.1f}ms, "
+                f"{total_fill_us / total * 100:.1f}% of CUDA time):"
+            )
+            print(f"    {'Input Dims':<50} {'Time':>9} {'Count':>7}")
+            print(f"    {'-' * 70}")
+            for dims, info in sorted(
+                fill_data.items(), key=lambda x: x[1]["total_us"], reverse=True
+            )[:10]:
+                dims_str = str(dims)[:50]
+                print(
+                    f"    {dims_str:<50} {info['total_us'] / 1000:>7.1f}ms "
+                    f"{info['count']:>7}"
+                )
+
+
+def print_summary(result, mem_result=None):
+    """Print optimization recommendations and summary."""
+    total = result["total_cuda_us"]
+    n = result["n_steps"]
+
+    print(f"\n{'=' * 75}")
+    print("  SUMMARY & RECOMMENDATIONS")
+    print(f"{'=' * 75}")
+
+    # Estimated per-step wall clock
+    wc = result.get("wall_clock_us", 0)
+    if wc > 0:
+        print(f"\n  Estimated per-step wall clock: {wc / n / 1e6:.2f}s")
+    else:
+        print(f"\n  Estimated per-step CUDA time:  {total / n / 1e6:.2f}s")
+
+    # Memory utilization
+    if mem_result:
+        reserved = mem_result["total_reserved"]
+        allocated = mem_result["total_allocated"]
+        if reserved > 0:
+            util_pct = allocated / reserved * 100
+            print(
+                f"  Memory utilization: {util_pct:.1f}% "
+                f"({allocated / 1e9:.2f} / {reserved / 1e9:.2f} GB)"
+            )
+
+    # Build recommendations
+    recommendations = []
+
+    cats_sorted = sorted(
+        result["categories"].items(), key=lambda x: x[1]["total_us"], reverse=True
+    )
+
+    # Check top category
+    if cats_sorted:
+        top_cat, top_info = cats_sorted[0]
+        top_pct = top_info["total_us"] / total * 100
+        if top_pct > 30:
+            if "GEMM" in top_cat or "CUTLASS" in top_cat:
+                recommendations.append(
+                    f"GEMM/matmul dominates ({top_pct:.0f}%). "
+                    f"Consider FP8 training, LoRA (fewer params), "
+                    f"or smaller batch size to reduce compute."
+                )
+            elif "Attention" in top_cat:
+                recommendations.append(
+                    f"Attention dominates ({top_pct:.0f}%). "
+                    f"Ensure FlashAttention v2/v3 is active. "
+                    f"Consider reducing sequence length or using sliding window."
+                )
+            elif "Fill" in top_cat or "Elementwise" in top_cat:
+                recommendations.append(
+                    f"Elementwise/Fill ops dominate ({top_pct:.0f}%). "
+                    f"Enable kernel fusion (Liger, torch.compile) to reduce "
+                    f"memory-bound elementwise operations."
+                )
+            elif "ScatterMoE" in top_cat:
+                recommendations.append(
+                    f"MoE routing/scatter dominates ({top_pct:.0f}%). "
+                    f"Check expert count and capacity factor. "
+                    f"Verify ScatterMoE kernels are using optimal block sizes."
+                )
+            else:
+                recommendations.append(
+                    f"'{top_cat}' dominates ({top_pct:.0f}%). "
+                    f"Focus optimization efforts here first."
+                )
+
+    # Check memory ops
+    for cat, info in cats_sorted:
+        pct = info["total_us"] / total * 100
+        if "Memory" in cat and pct > 10:
+            recommendations.append(
+                f"Memory ops are {pct:.0f}% of CUDA time. "
+                f"Consider gradient checkpointing, reducing activation "
+                f"recomputation, or pinned memory for data loading."
+            )
+            break
+
+    # Check fill overhead
+    fill_data = result.get("fill_by_size", {})
+    total_fill_us = sum(v["total_us"] for v in fill_data.values())
+    fill_pct = total_fill_us / total * 100 if total > 0 else 0
+    if fill_pct > 5:
+        recommendations.append(
+            f"Fill/zero_ kernels consume {fill_pct:.1f}% of CUDA time. "
+            f"Large zero-fills suggest excessive tensor allocation. "
+            f"Consider reusing buffers or lazy initialization."
+        )
+
+    # Check fragmentation
+    if mem_result and mem_result.get("fragmentation_pct", 0) > 20:
+        frag = mem_result["fragmentation_pct"]
+        recommendations.append(
+            f"Memory fragmentation is {frag:.0f}%. "
+            f"Use PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True "
+            f"or max_split_size_mb to reduce fragmentation."
+        )
+
+    # GPU utilization hint from wall clock vs CUDA time
+    if wc > 0 and total > 0:
+        gpu_util = total / wc * 100
+        if gpu_util < 50:
+            recommendations.append(
+                f"GPU utilization is ~{gpu_util:.0f}% (CUDA time vs wall clock). "
+                f"CPU-side bottleneck likely. Profile data loading, "
+                f"reward computation, or weight sync overhead."
+            )
+
+    # Print top 3
+    print("\n  Top optimization recommendations:")
+    if not recommendations:
+        print("    No major issues detected.")
+    for i, rec in enumerate(recommendations[:3], 1):
+        print(f"    {i}. {rec}")
+
+    print()
+
+
+def compare_traces(before, after):
+    print(f"\n{'=' * 75}")
+    print("  COMPARISON")
+    print(f"{'=' * 75}")
+
+    tb = before["total_cuda_us"]
+    ta = after["total_cuda_us"]
+    nb = before["n_steps"]
+    na = after["n_steps"]
+
+    print(
+        f"\n  Per-step CUDA time: {tb / nb / 1e6:.2f}s -> {ta / na / 1e6:.2f}s "
+        f"({(ta / na - tb / nb) / (tb / nb) * 100:+.1f}%)"
+    )
+
+    all_cats = sorted(
+        set(list(before["categories"]) + list(after["categories"])),
+        key=lambda c: before["categories"].get(c, {"total_us": 0})["total_us"],
+        reverse=True,
+    )
+
+    print(
+        f"\n  {'Category':<35} {'Before/step':>11} {'After/step':>11} "
+        f"{'Delta':>9} {'Speedup':>8}"
+    )
+    print(f"  {'-' * 78}")
+    for cat in all_cats:
+        b = before["categories"].get(cat, {"total_us": 0})["total_us"] / nb
+        a = after["categories"].get(cat, {"total_us": 0})["total_us"] / na
+        delta = a - b
+        speedup = b / a if a > 0 else float("inf")
+        if b > tb / nb * 0.003 or a > ta / na * 0.003:
+            print(
+                f"  {cat:<35} {b / 1000:>9.1f}ms {a / 1000:>9.1f}ms "
+                f"{delta / 1000:>+8.1f}ms {speedup:>7.2f}x"
+            )
+
+
+# ---- Allocation churn attribution -----------------------------------------
+
+
+def _short_path(p):
+    """Shorten a file path for display."""
+    if "/site-packages/" in p:
+        return "..." + p.split("/site-packages/")[1]
+    if "/axolotl/src/" in p:
+        return "axolotl/" + p.split("/axolotl/src/")[1]
+    if "/axolotl/" in p:
+        return "axolotl/" + p.split("/axolotl/")[1]
+    if len(p) > 60:
+        parts = p.split("/")
+        return ".../" + "/".join(parts[-3:])
+    return p
+
+
+def _get_top_python_frame(frames):
+    """Get the first meaningful Python frame from a trace frame list.
+
+    Skips internal torch/cuda/autograd frames to find the user-level code
+    that triggered the allocation.  Falls back to first Python frame.
+    """
+    skip_patterns = [
+        "torch/autograd/",
+        "torch/utils/checkpoint",
+        "torch/nn/modules/module.py",
+        "torch/_dynamo/",
+        "torch/_compile",
+        "torch/_ops",
+        "torch/_functorch/",
+        "torch/_inductor/",
+        "torch/library",
+        "torch/cuda",
+        "torch/_C",
+        "<frozen",
+        "<unknown>",
+        "fire/core.py",
+        "runpy",
+    ]
+    py_frames = [
+        f
+        for f in frames
+        if isinstance(f, dict) and f.get("filename", "").endswith(".py")
+    ]
+    for fr in py_frames:
+        fname = fr.get("filename", "")
+        if not any(p in fname for p in skip_patterns):
+            return fr
+    return py_frames[0] if py_frames else None
+
+
+def _categorize_source(fname, funcname):
+    """Assign a human-readable category to an allocation source."""
+    if "checkpoint" in fname.lower() or "backward" in funcname:
+        return "Gradient checkpoint recompute"
+    if "bitsandbytes" in fname:
+        return "BnB dequantization"
+    if "scattermoe" in fname or "scatter" in funcname.lower():
+        return "ScatterMoE LoRA"
+    if "adam" in funcname.lower() or "optim" in fname.lower() or "inductor" in fname:
+        return "Optimizer"
+    if "norm" in funcname.lower():
+        return "LayerNorm"
+    if "fla/" in fname or "gated_delta" in fname:
+        return "FLA linear attention"
+    return "Other"
+
+
+def analyze_allocation_churn(snapshot):
+    """Group allocation churn by Python source attribution.
+
+    For each major churn size, identifies which Python code creates
+    the tensors (gradient checkpointing, BnB dequant, LoRA conversion, etc).
+    """
+    traces = snapshot.get("device_traces", [[]])[0]
+    if not traces:
+        return None
+
+    # Find the top churn sizes
+    size_counts = defaultdict(int)
+    for ev in traces:
+        if ev.get("action") == "alloc":
+            size_counts[ev["size"]] += 1
+
+    # Top 5 sizes by total churn (count × size)
+    top_sizes = sorted(size_counts.items(), key=lambda x: x[0] * x[1], reverse=True)[:5]
+
+    results = {}
+    for target_sz, total_count in top_sizes:
+        mb = target_sz / 1e6
+        if mb < 1.0:
+            continue
+
+        frame_groups = defaultdict(lambda: {"count": 0})
+        for ev in traces:
+            if ev.get("action") == "alloc" and ev.get("size") == target_sz:
+                top = _get_top_python_frame(ev.get("frames", []))
+                if top:
+                    key = (top["filename"], top["name"], top.get("line", 0))
+                else:
+                    key = ("<unknown>", "<no attribution>", 0)
+                frame_groups[key]["count"] += 1
+
+        # Group into categories
+        categories = defaultdict(lambda: {"count": 0, "sources": []})
+        for (fname, funcname, lineno), info in frame_groups.items():
+            cat = _categorize_source(fname, funcname)
+            categories[cat]["count"] += info["count"]
+            categories[cat]["sources"].append(
+                (funcname, _short_path(fname), lineno, info["count"])
+            )
+
+        results[target_sz] = {
+            "total_count": total_count,
+            "total_churn_gb": total_count * target_sz / 1e9,
+            "categories": dict(categories),
+        }
+
+    return results
+
+
+def print_allocation_churn(results, label=""):
+    if not results:
+        return
+
+    if label:
+        print(f"\n{'=' * 75}")
+        print(f"  {label}")
+        print(f"{'=' * 75}")
+    else:
+        print("\n  ALLOCATION CHURN BY SOURCE")
+
+    for target_sz in sorted(
+        results, key=lambda s: results[s]["total_churn_gb"], reverse=True
+    ):
+        info = results[target_sz]
+        mb = target_sz / 1e6
+        print(
+            f"\n  {mb:.1f}MB × {info['total_count']:,} = {info['total_churn_gb']:.0f}GB churn:"
+        )
+        for cat, cinfo in sorted(
+            info["categories"].items(), key=lambda x: x[1]["count"], reverse=True
+        ):
+            pct = cinfo["count"] / info["total_count"] * 100
+            print(f"    {pct:>4.0f}%  {cat} ({cinfo['count']} allocs)")
+            for funcname, short, lineno, cnt in sorted(
+                cinfo["sources"], key=lambda x: x[3], reverse=True
+            )[:3]:
+                print(f"           {funcname:35s} {short}:{lineno} (x{cnt})")
+
+
+# ---- CPU overhead analysis ------------------------------------------------
+
+
+def analyze_cpu_overhead(events):
+    """Analyze memcpy, checkpoint recomputation, and GPU utilization from trace events."""
+    cuda_total_us = 0
+    cuda_count = 0
+    memcpy_stats = defaultdict(lambda: {"dur_us": 0, "count": 0, "bytes": 0})
+    checkpoint_us = 0
+    checkpoint_count = 0
+    min_ts = float("inf")
+    max_ts_end = 0
+
+    for ev in events:
+        if ev.get("ph") != "X":
+            continue
+        cat = ev.get("cat", "")
+        name = ev.get("name", "")
+        dur = float(ev.get("dur", 0))
+        ts = float(ev.get("ts", 0))
+        end = ts + dur
+        if ts < min_ts:
+            min_ts = ts
+        if end > max_ts_end:
+            max_ts_end = end
+
+        if cat == "kernel":
+            cuda_total_us += dur
+            cuda_count += 1
+        elif cat == "gpu_memcpy":
+            if "DtoH" in name:
+                direction = "GPU→CPU (offload)"
+            elif "HtoD" in name:
+                direction = "CPU→GPU (reload)"
+            elif "DtoD" in name:
+                direction = "GPU→GPU"
+            else:
+                direction = name
+            memcpy_stats[direction]["dur_us"] += dur
+            memcpy_stats[direction]["count"] += 1
+            nbytes = ev.get("args", {}).get("Bytes", 0)
+            if nbytes:
+                memcpy_stats[direction]["bytes"] += int(nbytes)
+        elif cat in ("cpu_op", "python_function"):
+            nl = name.lower()
+            if "checkpoint" in nl or "recompute" in nl:
+                checkpoint_us += dur
+                checkpoint_count += 1
+
+    wall_us = max_ts_end - min_ts if max_ts_end > min_ts else 0
+
+    return {
+        "wall_clock_us": wall_us,
+        "cuda_total_us": cuda_total_us,
+        "cuda_kernel_count": cuda_count,
+        "memcpy_stats": dict(memcpy_stats),
+        "checkpoint_us": checkpoint_us,
+        "checkpoint_count": checkpoint_count,
+    }
+
+
+def print_cpu_overhead(result, n_steps=2, label=""):
+    if not result:
+        return
+    if label:
+        print(f"\n{'=' * 75}")
+        print(f"  {label}")
+        print(f"{'=' * 75}")
+    else:
+        print("\n  CPU OVERHEAD ANALYSIS")
+
+    wall = result["wall_clock_us"]
+    cuda = result["cuda_total_us"]
+    count = result["cuda_kernel_count"]
+    gpu_util = cuda / wall * 100 if wall > 0 else 0
+    cpu_gap = wall - cuda
+
+    print(
+        f"\n  Wall clock:       {wall / 1e6:.2f}s (~{wall / n_steps / 1e6:.2f}s/step)"
+    )
+    print(f"  CUDA kernel time: {cuda / 1e6:.2f}s ({count:,} kernels)")
+    print(f"  GPU utilization:  {gpu_util:.1f}%")
+    print(f"  CPU overhead:     {cpu_gap / 1e6:.2f}s ({100 - gpu_util:.1f}%)")
+
+    memcpy = result["memcpy_stats"]
+    if memcpy:
+        total_memcpy = sum(v["dur_us"] for v in memcpy.values())
+        total_bytes = sum(v["bytes"] for v in memcpy.values())
+        print(
+            f"\n  Memory transfers: {total_memcpy / 1e6:.3f}s "
+            f"({total_bytes / 1e9:.2f}GB, {total_memcpy / wall * 100:.1f}% of wall)"
+        )
+        for direction, info in sorted(
+            memcpy.items(), key=lambda x: x[1]["dur_us"], reverse=True
+        ):
+            gb = info["bytes"] / 1e9
+            print(
+                f"    {direction:30s} {info['dur_us'] / 1e6:.3f}s  "
+                f"x{info['count']:>5}  {gb:.2f}GB"
+            )
+
+    if result["checkpoint_count"] > 0:
+        print(
+            f"\n  Gradient checkpoint CPU ops: {result['checkpoint_us'] / 1e6:.3f}s "
+            f"(x{result['checkpoint_count']})"
+        )
+
+
+# ---- Memory snapshot analysis ---------------------------------------------
+
+
+def load_snapshot(path):
+    """Load a PyTorch CUDA memory snapshot from a pickle file.
+
+    WARNING: This uses pickle.load() which can execute arbitrary code.
+    Only load snapshot files that you generated yourself from trusted
+    training runs. Never load snapshots from untrusted sources.
+    """
+    snap_file = Path(path) / "snapshot.pickle" if Path(path).is_dir() else Path(path)
+    if not snap_file.exists():
+        return None
+    print(f"Loading {snap_file.name} ({snap_file.stat().st_size / 1e6:.0f} MB)...")
+    with open(snap_file, "rb") as f:
+        return pickle.load(f)  # nosec B301
+
+
+def _extract_python_frames(snapshot):
+    """Extract Python source attribution from snapshot blocks with stacks='all'.
+
+    The snapshot structure (when stacks='all') stores frames in:
+      segments[i].blocks[j].history[k].frames = [(filename, lineno, name), ...]
+
+    Returns a dict mapping (filename, function_name) -> {"bytes": int, "count": int}
+    """
+    source_allocs = defaultdict(lambda: {"bytes": 0, "count": 0})
+
+    for seg in snapshot.get("segments", []):
+        for block in seg.get("blocks", []):
+            if block.get("state") != "active_allocated":
+                continue
+            size = block.get("size", 0)
+            history = block.get("history", [])
+            if not history:
+                continue
+
+            # Use the most recent allocation history entry
+            last_hist = history[-1]
+            frames = last_hist.get("frames", [])
+
+            # Find the first Python frame (skip C++ frames)
+            # Frames are tuples: (filename, lineno, name)
+            attributed = False
+            for frame in frames:
+                if not isinstance(frame, (list, tuple)) or len(frame) < 3:
+                    continue
+                filename, lineno, funcname = frame[0], frame[1], frame[2]
+                # Skip internal torch/cuda frames to find user-level attribution
+                fname_str = str(filename)
+                if any(
+                    skip in fname_str
+                    for skip in [
+                        "torch/cuda",
+                        "torch/_C",
+                        "torch/utils",
+                        "cuda/memory.py",
+                        "<unknown>",
+                    ]
+                ):
+                    continue
+                key = (fname_str, funcname, lineno)
+                source_allocs[key]["bytes"] += size
+                source_allocs[key]["count"] += 1
+                attributed = True
+                break
+
+            # If no user frame found, use first available frame
+            if not attributed and frames:
+                frame = frames[0]
+                if isinstance(frame, (list, tuple)) and len(frame) >= 3:
+                    key = (str(frame[0]), str(frame[2]), frame[1])
+                    source_allocs[key]["bytes"] += size
+                    source_allocs[key]["count"] += 1
+
+    return dict(source_allocs)
+
+
+def _extract_source_file_summary(source_allocs):
+    """Aggregate per-frame allocations to per-file level."""
+    file_allocs = defaultdict(lambda: {"bytes": 0, "count": 0, "functions": set()})
+    for (filename, funcname, _lineno), info in source_allocs.items():
+        file_allocs[filename]["bytes"] += info["bytes"]
+        file_allocs[filename]["count"] += info["count"]
+        file_allocs[filename]["functions"].add(funcname)
+    return dict(file_allocs)
+
+
+def analyze_snapshot(snapshot):
+    segments = snapshot.get("segments", [])
+    total_reserved = sum(s.get("total_size", 0) for s in segments)
+    total_allocated = sum(s.get("allocated_size", 0) for s in segments)
+
+    # Active blocks
+    active_blocks = []
+    for seg in segments:
+        for block in seg.get("blocks", []):
+            if block.get("state") == "active_allocated":
+                active_blocks.append(block.get("size", 0))
+
+    # Allocation churn from trace
+    trace = snapshot.get("device_traces", [[]])[0]
+    size_counts = defaultdict(lambda: {"count": 0, "total": 0})
+    for ev in trace:
+        if ev.get("action") == "alloc":
+            sz = ev.get("size", 0)
+            size_counts[sz]["count"] += 1
+            size_counts[sz]["total"] += sz
+
+    # Python frame attribution
+    source_allocs = _extract_python_frames(snapshot)
+    file_summary = _extract_source_file_summary(source_allocs)
+
+    return {
+        "total_reserved": total_reserved,
+        "total_allocated": total_allocated,
+        "fragmentation_pct": (total_reserved - total_allocated) / total_reserved * 100
+        if total_reserved > 0
+        else 0,
+        "n_segments": len(segments),
+        "n_active_blocks": len(active_blocks),
+        "active_bytes": sum(active_blocks),
+        "largest_active": sorted(active_blocks, reverse=True)[:10],
+        "alloc_churn": dict(
+            sorted(size_counts.items(), key=lambda x: x[1]["total"], reverse=True)[:15]
+        ),
+        "n_trace_events": len(trace),
+        "source_allocs": source_allocs,
+        "file_summary": file_summary,
+    }
+
+
+def print_memory_analysis(result, label=""):
+    if label:
+        print(f"\n{'=' * 75}")
+        print(f"  {label}")
+        print(f"{'=' * 75}")
+
+    reserved = result["total_reserved"]
+    allocated = result["total_allocated"]
+
+    print(f"\n  Reserved:    {reserved / 1e9:.2f} GB")
+    print(f"  Allocated:   {allocated / 1e9:.2f} GB")
+    print(f"  Utilization: {allocated / reserved * 100:.1f}%" if reserved > 0 else "")
+    print(f"  Fragmentation: {result['fragmentation_pct']:.1f}%")
+    print(
+        f"  Segments: {result['n_segments']}, Active blocks: {result['n_active_blocks']}"
+    )
+
+    print("\n  Largest active allocations:")
+    for sz in result["largest_active"]:
+        print(f"    {sz / 1e6:>10.1f} MB")
+
+    # Python source file attribution
+    file_summary = result.get("file_summary", {})
+    if file_summary:
+        print("\n  Top allocations by source file:")
+        print(f"    {'Source file':<55} {'Alloc':>10} {'Count':>7}")
+        print(f"    {'-' * 75}")
+        for fname, info in sorted(
+            file_summary.items(), key=lambda x: x[1]["bytes"], reverse=True
+        )[:15]:
+            # Shorten path for display
+            short = fname
+            if len(short) > 55:
+                parts = short.split("/")
+                # Keep last 3 path components
+                short = ".../" + "/".join(parts[-3:])
+                if len(short) > 55:
+                    short = short[:52] + "..."
+            funcs = ", ".join(sorted(info["functions"])[:3])
+            sz = info["bytes"]
+            unit = "MB"
+            val = sz / 1e6
+            if val >= 1000:
+                unit = "GB"
+                val = sz / 1e9
+            print(f"    {short:<55} {val:>8.1f}{unit} {info['count']:>7}")
+            if funcs:
+                print(f"      functions: {funcs[:70]}")
+
+    # Top allocations by function (more granular)
+    source_allocs = result.get("source_allocs", {})
+    if source_allocs:
+        print("\n  Top allocations by function (with line numbers):")
+        print(f"    {'Function':<35} {'File:Line':<35} {'Size':>10}")
+        print(f"    {'-' * 82}")
+        for (fname, funcname, lineno), info in sorted(
+            source_allocs.items(), key=lambda x: x[1]["bytes"], reverse=True
+        )[:15]:
+            # Shorten filename
+            short_file = fname
+            parts = short_file.split("/")
+            if len(parts) > 2:
+                short_file = "/".join(parts[-2:])
+            loc = f"{short_file}:{lineno}"
+            if len(loc) > 35:
+                loc = "..." + loc[-32:]
+            sz = info["bytes"]
+            if sz >= 1e9:
+                sz_str = f"{sz / 1e9:.2f}GB"
+            else:
+                sz_str = f"{sz / 1e6:.1f}MB"
+            print(f"    {funcname:<35} {loc:<35} {sz_str:>10}")
+
+    print("\n  Top allocation churn (alloc count x size):")
+    print(f"    {'Size':>12} {'Count':>8} {'Total churned':>14}")
+    print(f"    {'-' * 38}")
+    for sz, info in result["alloc_churn"].items():
+        if sz >= 1e6:
+            print(
+                f"    {sz / 1e6:>10.1f}MB {info['count']:>8} "
+                f"{info['total'] / 1e9:>12.2f}GB"
+            )
+
+
+# ---- Peak memory timeline from trace events --------------------------------
+
+
+def analyze_peak_memory(snapshot):
+    """Walk through device_traces chronologically to find peak concurrent memory usage.
+
+    The snapshot's segment data only captures end-of-step state.  The device_traces
+    record every alloc/free, letting us reconstruct peak usage and identify which
+    allocation sources were live at that moment.
+    """
+    traces = snapshot.get("device_traces", [[]])[0]
+    if not traces:
+        return None
+
+    current = 0
+    peak = 0
+    peak_idx = 0
+    live_allocs = {}  # addr -> (size, frames)
+    peak_live = {}
+
+    for i, ev in enumerate(traces):
+        action = ev.get("action")
+        addr = ev.get("addr", 0)
+        size = ev.get("size", 0)
+
+        if action == "alloc":
+            current += size
+            live_allocs[addr] = (size, ev.get("frames", []))
+            if current > peak:
+                peak = current
+                peak_idx = i
+                peak_live = dict(live_allocs)
+        elif action == "free_requested":
+            if addr in live_allocs:
+                current -= live_allocs[addr][0]
+                del live_allocs[addr]
+
+    # Categorize allocations at peak
+    peak_categories = defaultdict(lambda: {"bytes": 0, "count": 0})
+    for _addr, (size, frames) in peak_live.items():
+        top = _get_top_python_frame(frames)
+        if top:
+            cat = _categorize_source(top["filename"], top["name"])
+        else:
+            cat = "Unknown"
+        peak_categories[cat]["bytes"] += size
+        peak_categories[cat]["count"] += 1
+
+    return {
+        "peak_bytes": peak,
+        "peak_event_idx": peak_idx,
+        "total_events": len(traces),
+        "end_bytes": current,
+        "peak_categories": dict(peak_categories),
+    }
+
+
+def print_peak_memory(result, mem_result=None, label=""):
+    if not result:
+        return
+
+    if label:
+        print(f"\n{'=' * 75}")
+        print(f"  {label}")
+        print(f"{'=' * 75}")
+
+    peak_gb = result["peak_bytes"] / 1e9
+    end_gb = result["end_bytes"] / 1e9
+    # The device_traces only record allocations AFTER profiling starts.
+    # Model weights and other persistent allocations are not tracked.
+    # We can estimate the persistent baseline from snapshot allocated - peak_traced.
+    persistent_gb = 0
+    if mem_result:
+        persistent_gb = mem_result["total_allocated"] / 1e9 - end_gb
+    total_peak_gb = persistent_gb + peak_gb
+
+    print(
+        f"\n  Profiled peak (transient):  {peak_gb:.2f} GB  "
+        f"(at event {result['peak_event_idx']:,} / {result['total_events']:,})"
+    )
+    if persistent_gb > 0:
+        print(
+            f"  Persistent baseline:       {persistent_gb:.2f} GB  "
+            f"(model + optimizer, allocated before profiling)"
+        )
+        print(f"  Estimated total peak:      {total_peak_gb:.2f} GB")
+    print(f"  Transient headroom:        {peak_gb - end_gb:.2f} GB above end-of-trace")
+
+    cats = result.get("peak_categories", {})
+    if cats:
+        print("\n  Allocations live at peak:")
+        print(f"    {'Category':<35} {'Size':>10} {'Count':>7}")
+        print(f"    {'-' * 55}")
+        for cat, info in sorted(
+            cats.items(), key=lambda x: x[1]["bytes"], reverse=True
+        ):
+            sz = info["bytes"]
+            if sz >= 1e9:
+                sz_str = f"{sz / 1e9:.2f} GB"
+            else:
+                sz_str = f"{sz / 1e6:.1f} MB"
+            print(f"    {cat:<35} {sz_str:>10} {info['count']:>7}")
+
+
+# ---- Fragmentation diagnosis -----------------------------------------------
+
+
+def analyze_fragmentation(snapshot):
+    """Analyze segment-level memory layout to explain fragmentation.
+
+    Examines each CUDA segment for inactive (freed but unreturned) blocks,
+    pinned small allocations that prevent segment merging, and the overall
+    segment size distribution.
+    """
+    segments = snapshot.get("segments", [])
+    if not segments:
+        return None
+
+    total_reserved = 0
+    total_allocated = 0
+    total_inactive = 0
+    segment_sizes = []
+    inactive_gaps = []  # (gap_size, segment_size, active_around)
+    pinned_fragments = []  # small active blocks surrounded by inactive
+
+    for seg in segments:
+        seg_size = seg.get("total_size", 0)
+        total_reserved += seg_size
+        segment_sizes.append(seg_size)
+        blocks = seg.get("blocks", [])
+
+        seg_active = 0
+        seg_inactive = 0
+        for bi, block in enumerate(blocks):
+            bsize = block.get("size", 0)
+            if block.get("state") == "active_allocated":
+                seg_active += bsize
+                total_allocated += bsize
+                # Check if this small block is surrounded by inactive
+                if bsize < 2 * 1024 * 1024:  # < 2MB
+                    prev_inactive = bi > 0 and blocks[bi - 1].get("state") == "inactive"
+                    next_inactive = (
+                        bi < len(blocks) - 1
+                        and blocks[bi + 1].get("state") == "inactive"
+                    )
+                    if prev_inactive and next_inactive:
+                        pinned_fragments.append((bsize, seg_size))
+            elif block.get("state") == "inactive":
+                seg_inactive += bsize
+                total_inactive += bsize
+                inactive_gaps.append((bsize, seg_size))
+
+    # Classify segment sizes
+    size_buckets = defaultdict(lambda: {"count": 0, "total": 0})
+    for sz in segment_sizes:
+        if sz >= 1024 * 1024 * 1024:
+            bucket = ">=1 GB"
+        elif sz >= 256 * 1024 * 1024:
+            bucket = "256MB-1GB"
+        elif sz >= 64 * 1024 * 1024:
+            bucket = "64-256MB"
+        elif sz >= 2 * 1024 * 1024:
+            bucket = "2-64MB"
+        else:
+            bucket = "<2MB"
+        size_buckets[bucket]["count"] += 1
+        size_buckets[bucket]["total"] += sz
+
+    # Large inactive gaps that could be reclaimed
+    inactive_gaps.sort(key=lambda x: x[0], reverse=True)
+
+    return {
+        "total_reserved": total_reserved,
+        "total_allocated": total_allocated,
+        "total_inactive": total_inactive,
+        "n_segments": len(segments),
+        "segment_size_buckets": dict(size_buckets),
+        "large_inactive_gaps": inactive_gaps[:20],
+        "pinned_fragments": len(pinned_fragments),
+        "expandable_segments_would_help": (
+            total_inactive > 0.1 * total_reserved and len(segments) > 10
+        ),
+    }
+
+
+def print_fragmentation(result, gpu_capacity_gb=None, label=""):
+    if not result:
+        return
+
+    if label:
+        print(f"\n{'=' * 75}")
+        print(f"  {label}")
+        print(f"{'=' * 75}")
+
+    reserved = result["total_reserved"]
+    allocated = result["total_allocated"]
+    inactive = result["total_inactive"]
+    frag_pct = inactive / reserved * 100 if reserved > 0 else 0
+
+    print(
+        f"\n  Reserved:     {reserved / 1e9:.2f} GB across {result['n_segments']} segments"
+    )
+    print(f"  Allocated:    {allocated / 1e9:.2f} GB")
+    print(f"  Inactive:     {inactive / 1e9:.2f} GB ({frag_pct:.1f}% fragmentation)")
+
+    if result["pinned_fragments"] > 0:
+        print(
+            f"  Pinned small blocks (<2MB between inactive): "
+            f"{result['pinned_fragments']} (prevent segment merging)"
+        )
+
+    # Segment size distribution
+    print("\n  Segment size distribution:")
+    bucket_order = [">=1 GB", "256MB-1GB", "64-256MB", "2-64MB", "<2MB"]
+    for bucket in bucket_order:
+        info = result["segment_size_buckets"].get(bucket)
+        if info:
+            print(
+                f"    {bucket:<12} {info['count']:>4} segments  "
+                f"{info['total'] / 1e9:>6.2f} GB"
+            )
+
+    # Largest inactive gaps
+    gaps = result.get("large_inactive_gaps", [])
+    if gaps:
+        print("\n  Largest inactive gaps (freed but unreclaimable):")
+        shown = 0
+        for gap_sz, seg_sz in gaps:
+            if gap_sz >= 32 * 1024 * 1024 and shown < 10:
+                print(
+                    f"    {gap_sz / 1e6:>8.0f} MB gap in {seg_sz / 1e6:.0f} MB segment"
+                )
+                shown += 1
+
+    # OOM risk assessment
+    if gpu_capacity_gb:
+        gpu_bytes = gpu_capacity_gb * 1e9
+        usable = gpu_bytes - (reserved - allocated)  # capacity minus fragmented waste
+        print(f"\n  OOM Risk Assessment (GPU: {gpu_capacity_gb:.1f} GB):")
+        print(
+            f"    Usable capacity: {usable / 1e9:.2f} GB "
+            f"(GPU capacity minus {inactive / 1e9:.2f} GB fragmentation)"
+        )
+        headroom = gpu_bytes - reserved
+        print(f"    Current headroom: {headroom / 1e9:.2f} GB")
+        if headroom < 1.0e9:
+            print("    ⚠ CRITICAL: <1 GB headroom — high OOM risk!")
+        elif headroom < 2.0e9:
+            print("    ⚠ WARNING: <2 GB headroom — moderate OOM risk")
+
+    # Recommendation
+    if result.get("expandable_segments_would_help"):
+        print("\n  → FIX: Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True")
+        print("    This eliminates segment fragmentation by growing segments in-place,")
+        print(
+            f"    which would reclaim up to {inactive / 1e9:.1f} GB of wasted memory."
+        )
+
+
+# ---- Sequence-length scaling analysis --------------------------------------
+
+
+def analyze_scaling(mem_a, mem_b, churn_a, churn_b):
+    """Compare per-tensor allocation sizes between two runs.
+
+    When two profiles differ only in sequence length, this shows which tensor
+    categories scale with sequence length by comparing the dominant tensor sizes.
+    Total churn may differ due to different profiling windows, so we focus on
+    per-tensor size ratios instead.
+    """
+    if not churn_a or not churn_b:
+        return None
+
+    def _cat_sizes(churn):
+        """Map category -> {largest_size, total_bytes, count}."""
+        cat_data = defaultdict(lambda: {"max_size": 0, "sizes": [], "count": 0})
+        for sz, info in churn.items():
+            for cat, cinfo in info.get("categories", {}).items():
+                cat_data[cat]["count"] += cinfo["count"]
+                cat_data[cat]["sizes"].append((sz, cinfo["count"]))
+                if sz > cat_data[cat]["max_size"]:
+                    cat_data[cat]["max_size"] = sz
+        return dict(cat_data)
+
+    cats_a = _cat_sizes(churn_a)
+    cats_b = _cat_sizes(churn_b)
+
+    all_cats = sorted(
+        set(list(cats_a) + list(cats_b)),
+        key=lambda c: max(
+            cats_a.get(c, {"max_size": 0})["max_size"],
+            cats_b.get(c, {"max_size": 0})["max_size"],
+        ),
+        reverse=True,
+    )
+
+    scaling = []
+    for cat in all_cats:
+        a = cats_a.get(cat)
+        b = cats_b.get(cat)
+        if not a or not b:
+            continue
+        a_max = a["max_size"]
+        b_max = b["max_size"]
+        if a_max > 1e6 and b_max > 1e6:  # Only compare >1MB tensors
+            tensor_ratio = b_max / a_max if a_max > 0 else None
+            scaling.append(
+                {
+                    "category": cat,
+                    "size_a_mb": a_max / 1e6,
+                    "size_b_mb": b_max / 1e6,
+                    "tensor_ratio": tensor_ratio,
+                    "count_a": a["count"],
+                    "count_b": b["count"],
+                    "scales_with_seqlen": tensor_ratio is not None
+                    and tensor_ratio > 1.05,
+                }
+            )
+
+    scaling.sort(key=lambda x: x["size_b_mb"], reverse=True)
+    return scaling
+
+
+def print_scaling(scaling, label_a="Before", label_b="After", label=""):
+    if not scaling:
+        return
+
+    if label:
+        print(f"\n{'=' * 75}")
+        print(f"  {label}")
+        print(f"{'=' * 75}")
+
+    print("\n  Per-tensor size comparison (largest tensor per category):")
+    print(
+        f"  {'Category':<35} {'A size':>10} {'B size':>10} {'Ratio':>7} {'Scales?':>8}"
+    )
+    print(f"  {'-' * 73}")
+    for entry in scaling:
+        ratio_str = f"{entry['tensor_ratio']:.2f}x" if entry["tensor_ratio"] else "N/A"
+        scales = "YES" if entry["scales_with_seqlen"] else "no"
+        print(
+            f"  {entry['category']:<35} {entry['size_a_mb']:>8.1f}MB "
+            f"{entry['size_b_mb']:>8.1f}MB {ratio_str:>7} {scales:>8}"
+        )
+
+    # Summary
+    seq_scaling = [e for e in scaling if e["scales_with_seqlen"]]
+    constant = [e for e in scaling if not e["scales_with_seqlen"]]
+    if seq_scaling:
+        ratios = [e["tensor_ratio"] for e in seq_scaling if e["tensor_ratio"]]
+        avg_ratio = sum(ratios) / len(ratios) if ratios else 0
+        print(f"\n  Sequence-length scaling detected ({avg_ratio:.2f}x avg):")
+        for e in seq_scaling:
+            print(
+                f"    - {e['category']}: {e['size_a_mb']:.1f}MB -> "
+                f"{e['size_b_mb']:.1f}MB ({e['tensor_ratio']:.2f}x)"
+            )
+    if constant:
+        print("\n  Constant-size categories (do not scale with seq len):")
+        for e in constant:
+            print(f"    - {e['category']}: {e['size_a_mb']:.1f}MB")
+
+
+# ---- Main -----------------------------------------------------------------
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Analyze axolotl training profiler output",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    parser.add_argument(
+        "path",
+        help="Path to output directory (containing profiler_trace.json and/or "
+        "snapshot.pickle) or directly to a trace file. "
+        "Security note: snapshot.pickle uses pickle deserialization — "
+        "only use files from your own trusted training runs.",
+    )
+    parser.add_argument(
+        "--compare",
+        help="Path to second run for A/B comparison. "
+        "Same security note as path: only use trusted snapshot files.",
+    )
+    parser.add_argument(
+        "--include-warmup",
+        action="store_true",
+        help="Include step 0 (warmup/compilation) in timing analysis",
+    )
+    parser.add_argument(
+        "--memory-only",
+        action="store_true",
+        help="Only analyze memory snapshot, skip trace",
+    )
+    parser.add_argument(
+        "--quick",
+        action="store_true",
+        help="Only load first 2M events for rapid analysis of large traces",
+    )
+    parser.add_argument(
+        "--gpu-gb",
+        type=float,
+        default=None,
+        help="GPU total memory in GB (for OOM risk assessment). "
+        "Auto-detected if not specified.",
+    )
+    args = parser.parse_args()
+
+    # Auto-detect GPU capacity if not specified
+    gpu_capacity_gb = args.gpu_gb
+    if gpu_capacity_gb is None:
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                gpu_capacity_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+        except Exception:
+            pass
+
+    skip = not args.include_warmup
+    trace_result = None
+    mem_result = None
+    events = None
+
+    # -- Trace analysis --
+    if not args.memory_only:
+        events = load_trace(args.path, quick=args.quick)
+        if events:
+            trace_result = analyze_trace(events, skip_warmup=skip)
+            if trace_result:
+                print_trace_analysis(trace_result, label=f"Trace: {args.path}")
+
+                if args.compare:
+                    events2 = load_trace(args.compare, quick=args.quick)
+                    if events2:
+                        result2 = analyze_trace(events2, skip_warmup=skip)
+                        if result2:
+                            print_trace_analysis(
+                                result2, label=f"Trace: {args.compare}"
+                            )
+                            compare_traces(trace_result, result2)
+        else:
+            print(f"  No profiler_trace.json found in {args.path}")
+
+    # -- CPU overhead analysis (from trace) --
+    if events and trace_result:
+        cpu_result = analyze_cpu_overhead(events)
+        print_cpu_overhead(
+            cpu_result,
+            n_steps=trace_result["n_steps"],
+            label=f"CPU Overhead: {args.path}",
+        )
+
+    # -- Memory analysis --
+    snapshot = load_snapshot(args.path)
+    churn_result = None
+    churn2 = None
+    if snapshot:
+        mem_result = analyze_snapshot(snapshot)
+        print_memory_analysis(mem_result, label=f"Memory: {args.path}")
+
+        # Peak memory timeline
+        peak_result = analyze_peak_memory(snapshot)
+        print_peak_memory(
+            peak_result, mem_result=mem_result, label=f"Peak Memory: {args.path}"
+        )
+
+        # Fragmentation diagnosis
+        frag_result = analyze_fragmentation(snapshot)
+        print_fragmentation(
+            frag_result,
+            gpu_capacity_gb=gpu_capacity_gb,
+            label=f"Fragmentation: {args.path}",
+        )
+
+        # Allocation churn attribution
+        churn_result = analyze_allocation_churn(snapshot)
+        if churn_result:
+            print_allocation_churn(churn_result, label=f"Allocation Churn: {args.path}")
+
+        if args.compare:
+            snapshot2 = load_snapshot(args.compare)
+            if snapshot2:
+                mem2 = analyze_snapshot(snapshot2)
+                print_memory_analysis(mem2, label=f"Memory: {args.compare}")
+
+                # Peak memory for comparison
+                peak2 = analyze_peak_memory(snapshot2)
+                print_peak_memory(
+                    peak2, mem_result=mem2, label=f"Peak Memory: {args.compare}"
+                )
+
+                # Fragmentation for comparison
+                frag2 = analyze_fragmentation(snapshot2)
+                print_fragmentation(
+                    frag2,
+                    gpu_capacity_gb=gpu_capacity_gb,
+                    label=f"Fragmentation: {args.compare}",
+                )
+
+                churn2 = analyze_allocation_churn(snapshot2)
+                if churn2:
+                    print_allocation_churn(
+                        churn2, label=f"Allocation Churn: {args.compare}"
+                    )
+
+                # Memory comparison summary
+                print("\n  Memory comparison:")
+                print(
+                    f"    Reserved:  {mem_result['total_reserved'] / 1e9:.2f} -> "
+                    f"{mem2['total_reserved'] / 1e9:.2f} GB"
+                )
+                print(
+                    f"    Allocated: {mem_result['total_allocated'] / 1e9:.2f} -> "
+                    f"{mem2['total_allocated'] / 1e9:.2f} GB"
+                )
+                print(
+                    f"    Frag:      {mem_result['fragmentation_pct']:.1f}% -> "
+                    f"{mem2['fragmentation_pct']:.1f}%"
+                )
+                if peak_result and peak2:
+                    print(
+                        f"    Peak:      {peak_result['peak_bytes'] / 1e9:.2f} -> "
+                        f"{peak2['peak_bytes'] / 1e9:.2f} GB"
+                    )
+
+                # Scaling analysis
+                if churn_result and churn2:
+                    scaling = analyze_scaling(mem_result, mem2, churn_result, churn2)
+                    print_scaling(
+                        scaling,
+                        label_a=str(args.path),
+                        label_b=str(args.compare),
+                        label="Allocation Scaling Analysis",
+                    )
+    elif not args.memory_only:
+        pass  # trace-only is fine
+    else:
+        print(f"  No snapshot.pickle found in {args.path}")
+
+    # -- Summary --
+    if trace_result:
+        print_summary(trace_result, mem_result=mem_result)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/chat_datasets.py b/scripts/chat_datasets.py
new file mode 100644
index 0000000000..0c1e0bd037
--- /dev/null
+++ b/scripts/chat_datasets.py
@@ -0,0 +1,60 @@
+"""
+helper script to parse chat datasets into a usable yaml
+"""
+
+import click
+import yaml
+from datasets import load_dataset
+
+
+@click.command()
+@click.argument("dataset", type=str)
+@click.option("--split", type=str, default="train")
+def parse_dataset(dataset=None, split="train"):
+    ds_cfg = {}
+    ds_cfg["path"] = dataset
+    ds_cfg["split"] = split
+    ds_cfg["type"] = "chat_template"
+    ds_cfg["chat_template"] = "<<<Replace based on your model>>>"
+
+    dataset = load_dataset(dataset, split=split)
+    features = dataset.features
+    feature_keys = features.keys()
+    field_messages = None
+    for key in ["conversation", "conversations", "messages"]:
+        if key in feature_keys:
+            field_messages = key
+            break
+    if not field_messages:
+        raise ValueError(
+            f"No conversation field found in dataset: {', '.join(feature_keys)}"
+        )
+    ds_cfg["field_messages"] = field_messages
+
+    message_fields = features[field_messages][0].keys()
+
+    message_property_mappings = {"role": None, "content": None}
+    for key in ["from", "role"]:
+        if key in message_fields:
+            message_property_mappings["role"] = key
+            break
+    if not message_property_mappings["role"]:
+        raise ValueError(
+            f"No role field found in messages: {', '.join(message_fields)}"
+        )
+
+    for key in ["content", "text", "value"]:
+        if key in message_fields:
+            message_property_mappings["content"] = key
+            break
+    if not message_property_mappings["content"]:
+        raise ValueError(
+            f"No content field found in messages: {', '.join(message_fields)}"
+        )
+    ds_cfg["message_property_mappings"] = message_property_mappings
+
+    print(yaml.dump({"datasets": [ds_cfg]}))
+
+
+if __name__ == "__main__":
+    parse_dataset()
diff --git a/scripts/cloud-entrypoint-term.sh b/scripts/cloud-entrypoint-term.sh
new file mode 100755
index 0000000000..94511ec7c6
--- /dev/null
+++ b/scripts/cloud-entrypoint-term.sh
@@ -0,0 +1,82 @@
+#!/bin/bash
+
+# Export specific ENV variables to /etc/rp_environment
+echo "Exporting environment variables..."
+printenv | grep -E '^RUNPOD_|^PATH=|^_=' | sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >> /etc/rp_environment
+conda init
+# this needs to come after conda init
+echo 'source /etc/rp_environment' >> ~/.bashrc
+
+add_keys_to_authorized() {
+    local key_value=$1
+
+    # Create the ~/.ssh directory and set permissions
+    mkdir -p ~/.ssh
+    chmod 700 ~/.ssh
+
+    # Create the authorized_keys file if it doesn't exist
+    touch ~/.ssh/authorized_keys
+
+    # Initialize an empty key variable
+    local key=""
+
+    # Read the key variable word by word
+    for word in $key_value; do
+        # Check if the word looks like the start of a key
+        if [[ $word == ssh-* ]]; then
+            # If there's a key being built, add it to the authorized_keys file
+            if [[ -n $key ]]; then
+                echo $key >> ~/.ssh/authorized_keys
+            fi
+            # Start a new key
+            key=$word
+        else
+            # Append the word to the current key
+            key="$key $word"
+        fi
+    done
+
+    # Add the last key to the authorized_keys file
+    if [[ -n $key ]]; then
+        echo $key >> ~/.ssh/authorized_keys
+    fi
+
+    # Set the correct permissions
+    chmod 600 ~/.ssh/authorized_keys
+    chmod 700 -R ~/.ssh
+}
+
+if [[ $PUBLIC_KEY ]]; then
+    # runpod
+    add_keys_to_authorized "$PUBLIC_KEY"
+    # Start the SSH service in the background
+    service ssh start
+elif [[ $SSH_KEY ]]; then
+    # latitude.sh
+    add_keys_to_authorized "$SSH_KEY"
+    # Start the SSH service in the background
+    service ssh start
+else
+    echo "No PUBLIC_KEY or SSH_KEY environment variable provided, not starting openSSH daemon"
+fi
+
+# Check if JUPYTER_PASSWORD is set and not empty
+if [ -n "$JUPYTER_PASSWORD" ]; then
+    # Set JUPYTER_TOKEN to the value of JUPYTER_PASSWORD
+    export JUPYTER_TOKEN="$JUPYTER_PASSWORD"
+fi
+
+if [ "$JUPYTER_DISABLE" != "1" ]; then
+    # Run Jupyter Lab in the background
+    jupyter lab --port=8888 --ip=* --allow-root --ServerApp.allow_origin=* &
+fi
+
+if [ ! -d "/workspace/data/axolotl-artifacts" ]; then
+    mkdir -p /workspace/data/axolotl-artifacts
+fi
+if [ ! -L "/workspace/axolotl/outputs" ]; then
+    ln -sf /workspace/data/axolotl-artifacts /workspace/axolotl/outputs
+fi
+
+# Execute the passed arguments (CMD)
+exec "$@"
diff --git a/scripts/cloud-entrypoint.sh b/scripts/cloud-entrypoint.sh
index 3a15a884aa..7a4fb3ffbd 100755
--- a/scripts/cloud-entrypoint.sh
+++ b/scripts/cloud-entrypoint.sh
@@ -2,23 +2,67 @@
 
 # Export specific ENV variables to /etc/rp_environment
 echo "Exporting environment variables..."
-printenv | grep -E '^RUNPOD_|^PATH=|^_=' | sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >> /etc/rp_environment
+if [ -f /workspace/axolotl/scripts/cuda13_env.sh ]; then
+    source /workspace/axolotl/scripts/cuda13_env.sh
+    if [ -n "${LD_LIBRARY_PATH:-}" ]; then
+        echo "export LD_LIBRARY_PATH=\"${LD_LIBRARY_PATH}\"" >> /etc/rp_environment
+    fi
+fi
+printenv | grep -E '^HF_|^BNB_|^CUDA_|^NCCL_|^NV|^RUNPOD_|^PATH=|^_=' | sed 's/^\([^=]*\)=\(.*\)$/export \1="\2"/' | grep -v 'printenv' >> /etc/rp_environment
 echo 'source /etc/rp_environment' >> ~/.bashrc
 
-if [[ $PUBLIC_KEY ]]; then
-    # runpod
+add_keys_to_authorized() {
+    local key_value=$1
+
+    # Create the ~/.ssh directory and set permissions
     mkdir -p ~/.ssh
     chmod 700 ~/.ssh
-    echo $PUBLIC_KEY >> ~/.ssh/authorized_keys
+
+    # Create the authorized_keys file if it doesn't exist
+    touch ~/.ssh/authorized_keys
+
+    # Initialize an empty key variable
+    local key=""
+
+    # Read the key variable word by word
+    for word in $key_value; do
+        # Check if the word looks like the start of a key
+        if [[ $word == ssh-* ]]; then
+            # If there's a key being built, add it to the authorized_keys file
+            if [[ -n $key ]]; then
+                echo $key >> ~/.ssh/authorized_keys
+            fi
+            # Start a new key
+            key=$word
+        else
+            # Append the word to the current key
+            key="$key $word"
+        fi
+    done
+
+    # Add the last key to the authorized_keys file
+    if [[ -n $key ]]; then
+        echo $key >> ~/.ssh/authorized_keys
+    fi
+
+    # Set the correct permissions
+    chmod 600 ~/.ssh/authorized_keys
     chmod 700 -R ~/.ssh
+}
+
+# Set SSH port
+if [ ! -z "$SSH_PORT" ]; then
+    sed -i "s/#Port 22/Port $SSH_PORT/" /etc/ssh/sshd_config
+fi
+
+if [[ $PUBLIC_KEY ]]; then
+    # runpod, prime intellect
+    add_keys_to_authorized "$PUBLIC_KEY"
     # Start the SSH service in the background
     service ssh start
-elif [ -n "$SSH_KEY" ]; then
+elif [[ $SSH_KEY ]]; then
     # latitude.sh
-    mkdir -p ~/.ssh
-    chmod 700 ~/.ssh
-    echo $SSH_KEY >> ~/.ssh/authorized_keys
-    chmod 700 -R ~/.ssh
+    add_keys_to_authorized "$SSH_KEY"
     # Start the SSH service in the background
     service ssh start
 else
@@ -36,5 +80,20 @@ if [ "$JUPYTER_DISABLE" != "1" ]; then
     jupyter lab --port=8888 --ip=* --allow-root --ServerApp.allow_origin=* &
 fi
 
+if [ ! -d "/workspace/data/axolotl-artifacts" ]; then
+    mkdir -p /workspace/data/axolotl-artifacts
+fi
+if [ ! -L "/workspace/axolotl/outputs" ]; then
+    ln -sf /workspace/data/axolotl-artifacts /workspace/axolotl/outputs
+fi
+
+# start the runpod slurm init
+SLURM_INIT="${SLURM_INIT:-/slurm-init.sh}"
+
+if [[ -f "$SLURM_INIT" ]]; then
+  echo "[entrypoint] running $SLURM_INIT..."
+  bash "$SLURM_INIT"
+fi
+
 # Execute the passed arguments (CMD)
 exec "$@"
diff --git a/scripts/cuda13_env.sh b/scripts/cuda13_env.sh
new file mode 100644
index 0000000000..a739063a4e
--- /dev/null
+++ b/scripts/cuda13_env.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# Keep CUDA 13/cu130 uv images working when NVIDIA ships the cu13 package.
+# This is a no-op on images without nvidia.cu13.
+axolotl_prepend_cu13_ld_library_path() {
+    local updated_ld_library_path
+    if updated_ld_library_path="$(python - <<'PY'
+from axolotl.utils.cuda13 import prepend_cu13_ld_library_path
+import os
+
+print(prepend_cu13_ld_library_path(os.environ.get("LD_LIBRARY_PATH")))
+PY
+    )"; then
+        export LD_LIBRARY_PATH="$updated_ld_library_path"
+    fi
+}
+
+axolotl_prepend_cu13_ld_library_path
+unset -f axolotl_prepend_cu13_ld_library_path
diff --git a/scripts/cutcrossentropy_install.py b/scripts/cutcrossentropy_install.py
new file mode 100644
index 0000000000..201356e11d
--- /dev/null
+++ b/scripts/cutcrossentropy_install.py
@@ -0,0 +1,33 @@
+"""Script to output the correct installation command for cut-cross-entropy."""
+
+import importlib.util
+import sys
+
+try:
+    import torch
+except ImportError as exc:
+    raise ImportError("Install torch via `pip install torch`") from exc
+from packaging.version import Version as V
+
+USE_UV = "--uv" in sys.argv[1:]
+
+v = V(torch.__version__)
+
+# no cut-cross-entropy support for torch < 2.4.0
+if v < V("2.4.0"):
+    print("")
+    sys.exit(0)
+
+cce_spec = importlib.util.find_spec("cut_cross_entropy")
+
+UNINSTALL_PREFIX = ""
+if cce_spec:
+    if not importlib.util.find_spec("cut_cross_entropy.transformers"):
+        UNINSTALL_PREFIX = "pip uninstall -y cut-cross-entropy && "
+
+UV_PREFIX = "uv " if USE_UV else ""
+
+print(
+    UNINSTALL_PREFIX
+    + f'{UV_PREFIX}pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@5f0c7a7"'
+)
diff --git a/scripts/finetune.py b/scripts/finetune.py
deleted file mode 100644
index d5bbcaf8f0..0000000000
--- a/scripts/finetune.py
+++ /dev/null
@@ -1,52 +0,0 @@
-"""Prepare and train a model on a dataset. Can also infer from a model or merge lora"""
-import logging
-from pathlib import Path
-
-import fire
-import transformers
-
-from axolotl.cli import (
-    check_accelerate_default_config,
-    check_user_token,
-    do_inference,
-    do_merge_lora,
-    load_cfg,
-    load_datasets,
-    print_axolotl_text_art,
-)
-from axolotl.cli.shard import shard
-from axolotl.common.cli import TrainerCliArgs
-from axolotl.train import train
-
-LOG = logging.getLogger("axolotl.scripts.finetune")
-
-
-def do_cli(config: Path = Path("examples/"), **kwargs):
-    print_axolotl_text_art()
-    LOG.warning(
-        str(
-            PendingDeprecationWarning(
-                "scripts/finetune.py will be replaced with calling axolotl.cli.train"
-            )
-        )
-    )
-    parsed_cfg = load_cfg(config, **kwargs)
-    check_accelerate_default_config()
-    check_user_token()
-    parser = transformers.HfArgumentParser((TrainerCliArgs))
-    parsed_cli_args, _ = parser.parse_args_into_dataclasses(
-        return_remaining_strings=True
-    )
-    if parsed_cli_args.inference:
-        do_inference(cfg=parsed_cfg, cli_args=parsed_cli_args)
-    elif parsed_cli_args.merge_lora:
-        do_merge_lora(cfg=parsed_cfg, cli_args=parsed_cli_args)
-    elif parsed_cli_args.shard:
-        shard(cfg=parsed_cfg, cli_args=parsed_cli_args)
-    else:
-        dataset_meta = load_datasets(cfg=parsed_cfg, cli_args=parsed_cli_args)
-        train(cfg=parsed_cfg, cli_args=parsed_cli_args, dataset_meta=dataset_meta)
-
-
-if __name__ == "__main__":
-    fire.Fire(do_cli)
diff --git a/scripts/generate_cli_config_options.py b/scripts/generate_cli_config_options.py
new file mode 100644
index 0000000000..a3fef9f94d
--- /dev/null
+++ b/scripts/generate_cli_config_options.py
@@ -0,0 +1,6 @@
+"""Generate lightweight Click option metadata for the Axolotl CLI."""
+
+from axolotl.cli.generate_config_options import main
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/motd b/scripts/motd
index 060a5c5c6c..275a4fcba8 100644
--- a/scripts/motd
+++ b/scripts/motd
@@ -1,17 +1,24 @@
 
-                                 dP            dP   dP
-                                 88            88   88
-      .d8888b. dP.  .dP .d8888b. 88 .d8888b. d8888P 88
-      88'  `88  `8bd8'  88'  `88 88 88'  `88   88   88
-      88.  .88  .d88b.  88.  .88 88 88.  .88   88   88
-      `88888P8 dP'  `dP `88888P' dP `88888P'   dP   dP
+     #@@ #@@      @@# @@#
+    @@  @@          @@  @@           =@@#                               @@                 #@    =@@#.
+    @@    #@@@@@@@@@    @@           #@#@=                              @@                 #@     .=@@
+      #@@@@@@@@@@@@@@@@@            =@# @#     ##=     ##    =####=+    @@      =#####+  =#@@###.   @@
+    @@@@@@@@@@/  +@@/  +@@          #@  =@=     #@=   @@   =@#+  +#@#   @@    =@#+  +#@#   #@.      @@
+    @@@@@@@@@@  ##@@  ##@@         =@#   @#      =@# @#    @@      @@   @@    @@      #@   #@       @@
+     @@@@@@@@@@@@@@@@@@@@          #@=+++#@=      =@@#     @@      @@   @@    @@      #@   #@       @@
+                                  =@#=====@@     =@# @#    @@      @@   @@    @@      #@   #@       @@
+    @@@@@@@@@@@@@@@@  @@@@        #@      #@=   #@=  +@@   #@#    =@#   @@.   =@#    =@#   #@.      @@
+                                 =@#       @#  #@=     #@   =#@@@@#=    +#@@=  +#@@@@#=    .##@@+   @@
+    @@@@  @@@@@@@@@@@@@@@@
 
-Welcome to the axolotl cloud image! If the you've mounted a disk to /workspace and the axolotl directory ie empty, run the following commands:
+Welcome to the axolotl cloud image! If the you've mounted a disk to /workspace and the axolotl directory is empty, run the following commands:
+
+Need help with your post-training workloads? Reach out us at contact@axolotl.ai for assistance.
 
 ```
 cd /workspace
 rm -rf /workspace/axolotl
-git clone https://github.com/OpenAccess-AI-Collective/axolotl.git
+git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
-pip install --no-deps -e .
+pip install --no-build-isolation --no-deps -e .
 ```
diff --git a/scripts/uv-entrypoint.sh b/scripts/uv-entrypoint.sh
new file mode 100644
index 0000000000..8e386352ab
--- /dev/null
+++ b/scripts/uv-entrypoint.sh
@@ -0,0 +1,7 @@
+#!/bin/bash
+
+if [ -f /workspace/axolotl/scripts/cuda13_env.sh ]; then
+    source /workspace/axolotl/scripts/cuda13_env.sh
+fi
+
+exec "$@"
diff --git a/setup.py b/setup.py
deleted file mode 100644
index fbca5a360e..0000000000
--- a/setup.py
+++ /dev/null
@@ -1,96 +0,0 @@
-"""setup.py for axolotl"""
-
-import platform
-import re
-from importlib.metadata import PackageNotFoundError, version
-
-from setuptools import find_packages, setup
-
-
-def parse_requirements():
-    _install_requires = []
-    _dependency_links = []
-    with open("./requirements.txt", encoding="utf-8") as requirements_file:
-        lines = [r.strip() for r in requirements_file.readlines()]
-        for line in lines:
-            is_extras = (
-                "flash-attn" in line
-                or "flash-attention" in line
-                or "deepspeed" in line
-                or "mamba-ssm" in line
-                or "lion-pytorch" in line
-            )
-            if line.startswith("--extra-index-url"):
-                # Handle custom index URLs
-                _, url = line.split()
-                _dependency_links.append(url)
-            elif not is_extras and line and line[0] != "#":
-                # Handle standard packages
-                _install_requires.append(line)
-
-    try:
-        if "Darwin" in platform.system():
-            _install_requires.pop(_install_requires.index("xformers==0.0.22"))
-        else:
-            torch_version = version("torch")
-            _install_requires.append(f"torch=={torch_version}")
-
-            version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version)
-            if version_match:
-                major, minor, patch = version_match.groups()
-                major, minor = int(major), int(minor)
-                patch = (
-                    int(patch) if patch is not None else 0
-                )  # Default patch to 0 if not present
-            else:
-                raise ValueError("Invalid version format")
-
-            if (major, minor) >= (2, 1):
-                _install_requires.pop(_install_requires.index("xformers==0.0.22"))
-                _install_requires.append("xformers>=0.0.23")
-    except PackageNotFoundError:
-        pass
-
-    return _install_requires, _dependency_links
-
-
-install_requires, dependency_links = parse_requirements()
-
-
-setup(
-    name="axolotl",
-    version="0.4.0",
-    description="LLM Trainer",
-    long_description="Axolotl is a tool designed to streamline the fine-tuning of various AI models, offering support for multiple configurations and architectures.",
-    package_dir={"": "src"},
-    packages=find_packages(),
-    install_requires=install_requires,
-    dependency_links=dependency_links,
-    extras_require={
-        "flash-attn": [
-            "flash-attn==2.5.5",
-        ],
-        "fused-dense-lib": [
-            "fused-dense-lib  @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",
-        ],
-        "deepspeed": [
-            "deepspeed==0.13.1",
-            "deepspeed-kernels",
-        ],
-        "mamba-ssm": [
-            "mamba-ssm==1.2.0.post1",
-        ],
-        "auto-gptq": [
-            "auto-gptq==0.5.1",
-        ],
-        "mlflow": [
-            "mlflow",
-        ],
-        "lion-pytorch": [
-            "lion-pytorch==0.1.2",
-        ],
-        "galore": [
-            "galore_torch",
-        ],
-    },
-)
diff --git a/src/axolotl/__init__.py b/src/axolotl/__init__.py
index e69de29bb2..4b9f11bbd0 100644
--- a/src/axolotl/__init__.py
+++ b/src/axolotl/__init__.py
@@ -0,0 +1,11 @@
+"""Axolotl - Train and fine-tune large language models"""
+
+import pkgutil
+from importlib.metadata import PackageNotFoundError, version
+
+__path__ = pkgutil.extend_path(__path__, __name__)  # Make this a namespace package
+
+try:
+    __version__ = version("axolotl")
+except PackageNotFoundError:
+    __version__ = "unknown"
diff --git a/src/axolotl/cli/__init__.py b/src/axolotl/cli/__init__.py
index 7ec3f524ab..799d5694ef 100644
--- a/src/axolotl/cli/__init__.py
+++ b/src/axolotl/cli/__init__.py
@@ -1,490 +1,11 @@
-"""Prepare and train a model on a dataset. Can also infer from a model or merge lora"""
+"""Axolotl CLI module initialization."""
 
-import importlib
-import json
-import logging
-import math
 import os
-import random
-import sys
-import tempfile
-from pathlib import Path
-from threading import Thread
-from typing import Any, Dict, List, Optional, Union
-from urllib.parse import urlparse
 
-import requests
-import torch
-import yaml
-
-# add src to the pythonpath so we don't need to pip install this
-from accelerate.commands.config import config_args
-from art import text2art
-from huggingface_hub import HfApi
-from huggingface_hub.utils import LocalTokenNotFoundError
-from transformers import GenerationConfig, TextIteratorStreamer, TextStreamer
-from transformers.utils import is_torch_bf16_gpu_available
-from transformers.utils.import_utils import _is_package_available
-
-from axolotl.common.cli import TrainerCliArgs, load_model_and_tokenizer
 from axolotl.logging_config import configure_logging
-from axolotl.train import TrainDatasetMeta
-from axolotl.utils.config import (
-    normalize_cfg_datasets,
-    normalize_config,
-    validate_config,
-)
-from axolotl.utils.data import load_prepare_dpo_datasets, prepare_dataset
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.distributed import is_main_process
-from axolotl.utils.mlflow_ import setup_mlflow_env_vars
-from axolotl.utils.models import load_tokenizer
-from axolotl.utils.tokenization import check_dataset_labels
-from axolotl.utils.trainer import prepare_optim_env
-from axolotl.utils.wandb_ import setup_wandb_env_vars
 
-project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
-src_dir = os.path.join(project_root, "src")
-sys.path.insert(0, src_dir)
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
+os.environ.setdefault("TRL_EXPERIMENTAL_SILENCE", "1")
 
 configure_logging()
-LOG = logging.getLogger("axolotl.scripts")
-
-os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-
-
-def print_axolotl_text_art(suffix=None):
-    font = "nancyj"
-    ascii_text = "  axolotl"
-    if suffix:
-        ascii_text += f"  x  {suffix}"
-    ascii_art = text2art(ascii_text, font=font)
-
-    if is_main_process():
-        print(ascii_art)
-
-    print_dep_versions()
-
-
-def print_dep_versions():
-    packages = ["accelerate", "peft", "transformers", "trl", "torch", "bitsandbytes"]
-    max_len = max(len(pkg) for pkg in packages)
-    if is_main_process():
-        print("*" * 40)
-        print("**** Axolotl Dependency Versions *****")
-        for pkg in packages:
-            version = _is_package_available(pkg, return_version=True)
-            print(f"{pkg: >{max_len}}: {version[1]: <15}")
-        print("*" * 40)
-
-
-def check_remote_config(config: Union[str, Path]):
-    # Check if the config is a valid HTTPS URL to a .yml or .yaml file
-    if not (isinstance(config, str) and config.startswith("https://")):
-        return config  # Return the original value if it's not a valid URL
-
-    filename = os.path.basename(urlparse(config).path)
-    temp_dir = tempfile.mkdtemp()
-
-    try:
-        response = requests.get(config, timeout=30)
-        response.raise_for_status()  # Check for HTTP errors
-
-        content = response.content
-        try:
-            # Try parsing as JSON first to catch cases where JSON content is mistakenly considered YAML
-            json.loads(content)
-            # Log a warning but do not raise an error; JSON is technically valid YAML - this can happen when you forget to point to a raw github link
-            LOG.warning(
-                f"Warning: The content of the file at {config} is JSON, which is technically valid YAML but might not be intended."
-            )
-        except json.JSONDecodeError:
-            # If it's not valid JSON, verify it's valid YAML
-            try:
-                yaml.safe_load(content)
-            except yaml.YAMLError as err:
-                raise ValueError(
-                    f"Failed to parse the content at {config} as YAML: {err}"
-                ) from err
-
-        # Write the content to a file if it's valid YAML (or JSON treated as YAML)
-        output_path = Path(temp_dir) / filename
-        with open(output_path, "wb") as file:
-            file.write(content)
-        LOG.info(
-            f"Using the following config obtained from {config}:\n\n{content.decode('utf-8')}\n"
-        )
-        return output_path
-
-    except requests.RequestException as err:
-        # This catches all requests-related exceptions including HTTPError
-        raise RuntimeError(f"Failed to download {config}: {err}") from err
-    except Exception as err:
-        # Catch-all for any other exceptions
-        raise err
-
-
-def get_multi_line_input() -> Optional[str]:
-    print("Give me an instruction (Ctrl + D to submit): ")
-    instruction = ""
-    for line in sys.stdin:
-        instruction += line  # pylint: disable=consider-using-join
-    # instruction = pathlib.Path("/proc/self/fd/0").read_text()
-    return instruction
-
-
-def do_merge_lora(
-    *,
-    cfg: DictDefault,
-    cli_args: TrainerCliArgs,
-):
-    model, tokenizer = load_model_and_tokenizer(cfg=cfg, cli_args=cli_args)
-    safe_serialization = cfg.save_safetensors is True
-
-    LOG.info("running merge of LoRA with base model")
-    model = model.merge_and_unload(progressbar=True)
-    try:
-        model.to(dtype=cfg.torch_dtype)
-    except RuntimeError:
-        pass
-    model.generation_config.do_sample = True
-
-    if cfg.local_rank == 0:
-        LOG.info(f"saving merged model to: {str(Path(cfg.output_dir) / 'merged')}")
-        model.save_pretrained(
-            str(Path(cfg.output_dir) / "merged"),
-            safe_serialization=safe_serialization,
-            progressbar=True,
-        )
-        tokenizer.save_pretrained(str(Path(cfg.output_dir) / "merged"))
-
-
-def do_inference(
-    *,
-    cfg: DictDefault,
-    cli_args: TrainerCliArgs,
-):
-    model, tokenizer = load_model_and_tokenizer(cfg=cfg, cli_args=cli_args)
-    prompter = cli_args.prompter
-    default_tokens = {"unk_token": "<unk>", "bos_token": "<s>", "eos_token": "</s>"}
-
-    for token, symbol in default_tokens.items():
-        # If the token isn't already specified in the config, add it
-        if not (cfg.special_tokens and token in cfg.special_tokens):
-            tokenizer.add_special_tokens({token: symbol})
-
-    prompter_module = None
-    if prompter:
-        prompter_module = getattr(
-            importlib.import_module("axolotl.prompters"), prompter
-        )
-
-    model = model.to(cfg.device, dtype=cfg.torch_dtype)
-
-    while True:
-        print("=" * 80)
-        # support for multiline inputs
-        instruction = get_multi_line_input()
-        if not instruction:
-            return
-        if prompter_module:
-            prompt: str = next(
-                prompter_module().build_prompt(instruction=instruction.strip("\n"))
-            )
-        else:
-            prompt = instruction.strip()
-        batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
-
-        print("=" * 40)
-        model.eval()
-        with torch.no_grad():
-            generation_config = GenerationConfig(
-                repetition_penalty=1.1,
-                max_new_tokens=1024,
-                temperature=0.9,
-                top_p=0.95,
-                top_k=40,
-                bos_token_id=tokenizer.bos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                pad_token_id=tokenizer.pad_token_id,
-                do_sample=True,
-                use_cache=True,
-                return_dict_in_generate=True,
-                output_attentions=False,
-                output_hidden_states=False,
-                output_scores=False,
-            )
-            streamer = TextStreamer(tokenizer)
-            generated = model.generate(
-                inputs=batch["input_ids"].to(cfg.device),
-                generation_config=generation_config,
-                streamer=streamer,
-            )
-        print("=" * 40)
-        print(tokenizer.decode(generated["sequences"].cpu().tolist()[0]))
-
-
-def do_inference_gradio(
-    *,
-    cfg: DictDefault,
-    cli_args: TrainerCliArgs,
-):
-    import gradio as gr
-
-    model, tokenizer = load_model_and_tokenizer(cfg=cfg, cli_args=cli_args)
-    prompter = cli_args.prompter
-    default_tokens = {"unk_token": "<unk>", "bos_token": "<s>", "eos_token": "</s>"}
-
-    for token, symbol in default_tokens.items():
-        # If the token isn't already specified in the config, add it
-        if not (cfg.special_tokens and token in cfg.special_tokens):
-            tokenizer.add_special_tokens({token: symbol})
-
-    prompter_module = None
-    if prompter:
-        prompter_module = getattr(
-            importlib.import_module("axolotl.prompters"), prompter
-        )
-
-    model = model.to(cfg.device, dtype=cfg.torch_dtype)
-
-    def generate(instruction):
-        if not instruction:
-            return
-        if prompter_module:
-            # pylint: disable=stop-iteration-return
-            prompt: str = next(
-                prompter_module().build_prompt(instruction=instruction.strip("\n"))
-            )
-        else:
-            prompt = instruction.strip()
-        batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
-
-        model.eval()
-        with torch.no_grad():
-            generation_config = GenerationConfig(
-                repetition_penalty=1.1,
-                max_new_tokens=cfg.get("gradio_max_new_tokens", 1024),
-                temperature=cfg.get("gradio_temperature", 0.9),
-                top_p=0.95,
-                top_k=40,
-                bos_token_id=tokenizer.bos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                pad_token_id=tokenizer.pad_token_id,
-                do_sample=True,
-                use_cache=True,
-                return_dict_in_generate=True,
-                output_attentions=False,
-                output_hidden_states=False,
-                output_scores=False,
-            )
-            streamer = TextIteratorStreamer(tokenizer)
-            generation_kwargs = {
-                "inputs": batch["input_ids"].to(cfg.device),
-                "generation_config": generation_config,
-                "streamer": streamer,
-            }
-
-            thread = Thread(target=model.generate, kwargs=generation_kwargs)
-            thread.start()
-
-            all_text = ""
-
-            for new_text in streamer:
-                all_text += new_text
-                yield all_text
-
-    demo = gr.Interface(
-        fn=generate,
-        inputs="textbox",
-        outputs="text",
-        title=cfg.get("gradio_title", "Axolotl Gradio Interface"),
-    )
-
-    demo.queue().launch(
-        show_api=False,
-        share=cfg.get("gradio_share", True),
-        server_name=cfg.get("gradio_server_name", "127.0.0.1"),
-        server_port=cfg.get("gradio_server_port", None),
-    )
-
-
-def choose_config(path: Path):
-    yaml_files = list(path.glob("*.yml"))
-
-    if not yaml_files:
-        raise ValueError(
-            "No YAML config files found in the specified directory. Are you using a .yml extension?"
-        )
-
-    if len(yaml_files) == 1:
-        print(f"Using default YAML file '{yaml_files[0]}'")
-        return yaml_files[0]
-
-    print("Choose a YAML file:")
-    for idx, file in enumerate(yaml_files):
-        print(f"{idx + 1}. {file}")
-
-    chosen_file = None
-    while chosen_file is None:
-        try:
-            choice = int(input("Enter the number of your choice: "))
-            if 1 <= choice <= len(yaml_files):
-                chosen_file = yaml_files[choice - 1]
-            else:
-                print("Invalid choice. Please choose a number from the list.")
-        except ValueError:
-            print("Invalid input. Please enter a number.")
-
-    return chosen_file
-
-
-def check_not_in(list1: List[str], list2: Union[Dict[str, Any], List[str]]) -> bool:
-    return not any(el in list2 for el in list1)
-
-
-def load_cfg(config: Union[str, Path] = Path("examples/"), **kwargs):
-    config = check_remote_config(config)
-    if Path(config).is_dir():
-        config = choose_config(Path(config))
-
-    # load the config from the yaml file
-    with open(config, encoding="utf-8") as file:
-        cfg: DictDefault = DictDefault(yaml.safe_load(file))
-    # if there are any options passed in the cli, if it is something that seems valid from the yaml,
-    # then overwrite the value
-    cfg_keys = cfg.keys()
-    for k, _ in kwargs.items():
-        # if not strict, allow writing to cfg even if it's not in the yml already
-        if k in cfg_keys or not cfg.strict:
-            # handle booleans
-            if isinstance(cfg[k], bool):
-                cfg[k] = bool(kwargs[k])
-            else:
-                cfg[k] = kwargs[k]
-
-    cfg.axolotl_config_path = config
-
-    try:
-        device_props = torch.cuda.get_device_properties("cuda")
-        gpu_version = "sm_" + str(device_props.major) + str(device_props.minor)
-    except:  # pylint: disable=bare-except # noqa: E722
-        gpu_version = None
-
-    cfg = validate_config(
-        cfg,
-        capabilities={
-            "bf16": is_torch_bf16_gpu_available(),
-            "n_gpu": os.environ.get("WORLD_SIZE", 1),
-            "compute_capability": gpu_version,
-        },
-    )
-
-    prepare_optim_env(cfg)
-
-    normalize_config(cfg)
-
-    normalize_cfg_datasets(cfg)
-
-    setup_wandb_env_vars(cfg)
-
-    setup_mlflow_env_vars(cfg)
-
-    return cfg
-
-
-def load_datasets(
-    *,
-    cfg: DictDefault,
-    cli_args: TrainerCliArgs,
-) -> TrainDatasetMeta:
-    tokenizer = load_tokenizer(cfg)
-
-    train_dataset, eval_dataset, total_num_steps, prompters = prepare_dataset(
-        cfg, tokenizer
-    )
-
-    if cli_args.debug or cfg.debug:
-        LOG.info("check_dataset_labels...")
-        check_dataset_labels(
-            train_dataset.select(
-                [
-                    random.randrange(0, len(train_dataset) - 1)  # nosec
-                    for _ in range(cli_args.debug_num_examples)
-                ]
-            ),
-            tokenizer,
-            num_examples=cli_args.debug_num_examples,
-            text_only=cli_args.debug_text_only,
-        )
-
-        LOG.info("printing prompters...")
-        for prompter in prompters:
-            LOG.info(prompter)
-
-    return TrainDatasetMeta(
-        train_dataset=train_dataset,
-        eval_dataset=eval_dataset,
-        total_num_steps=total_num_steps,
-    )
-
-
-def load_rl_datasets(
-    *,
-    cfg: DictDefault,
-    cli_args: TrainerCliArgs,  # pylint: disable=unused-argument
-) -> TrainDatasetMeta:
-    train_dataset, eval_dataset = load_prepare_dpo_datasets(cfg)
-    total_num_steps = int(
-        math.ceil(len(train_dataset) * cfg.num_epochs / cfg.batch_size)
-    )
-
-    if cli_args.debug or cfg.debug:
-        LOG.info("check_dataset_labels...")
-
-        tokenizer = load_tokenizer(cfg)
-        check_dataset_labels(
-            train_dataset.select(
-                [
-                    random.randrange(0, len(train_dataset) - 1)  # nosec
-                    for _ in range(cli_args.debug_num_examples)
-                ]
-            ),
-            tokenizer,
-            num_examples=cli_args.debug_num_examples,
-            text_only=cli_args.debug_text_only,
-            rl_mode=True,
-        )
-
-    return TrainDatasetMeta(
-        train_dataset=train_dataset,
-        eval_dataset=eval_dataset,
-        total_num_steps=total_num_steps,
-    )
-
-
-def check_accelerate_default_config():
-    if Path(config_args.default_yaml_config_file).exists():
-        LOG.warning(
-            f"accelerate config file found at {config_args.default_yaml_config_file}. This can lead to unexpected errors"
-        )
-
-
-def check_user_token():
-    # Skip check if HF_HUB_OFFLINE is set to True
-    if os.getenv("HF_HUB_OFFLINE") == "1":
-        LOG.info(
-            "Skipping HuggingFace token verification because HF_HUB_OFFLINE is set to True. Only local files will be used."
-        )
-        return True
-
-    # Verify if token is valid
-    api = HfApi()
-    try:
-        user_info = api.whoami()
-        return bool(user_info)
-    except LocalTokenNotFoundError:
-        LOG.warning(
-            "Error verifying HuggingFace token. Remember to log in using `huggingface-cli login` and get your access token from https://huggingface.co/settings/tokens if you want to use gated models or datasets."
-        )
-        return False
diff --git a/src/axolotl/cli/agent_docs/__init__.py b/src/axolotl/cli/agent_docs/__init__.py
new file mode 100644
index 0000000000..14dbff32d3
--- /dev/null
+++ b/src/axolotl/cli/agent_docs/__init__.py
@@ -0,0 +1,108 @@
+"""Bundled agent documentation for axolotl.
+
+These docs are optimized for consumption by AI coding agents.
+The source of truth is docs/agents/*.md and AGENTS.md in the repo root.
+This module resolves those paths at runtime — no files are duplicated
+into the package.
+
+For pip-only installs (no repo checkout), run `axolotl fetch docs` first
+to download the docs locally.
+"""
+
+from pathlib import Path
+
+# Topic name -> (filename in docs/agents/, fallback filename for AGENTS.md)
+TOPICS = {
+    "overview": "AGENTS.md",
+    "sft": "docs/agents/sft.md",
+    "grpo": "docs/agents/grpo.md",
+    "preference_tuning": "docs/agents/preference_tuning.md",
+    "reward_modelling": "docs/agents/reward_modelling.md",
+    "pretraining": "docs/agents/pretraining.md",
+    "model_architectures": "docs/agents/model_architectures.md",
+    "new_model_support": "docs/agents/new_model_support.md",
+}
+
+
+def _find_repo_root() -> Path | None:
+    """Walk up from this file to find the repo root (contains AGENTS.md)."""
+    # In an editable install or repo checkout, walk up from
+    # src/axolotl/cli/agent_docs/ to find the repo root
+    current = Path(__file__).resolve().parent
+    while current != current.parent:
+        if (current / "AGENTS.md").exists() and (current / "docs" / "agents").is_dir():
+            return current
+        current = current.parent
+    return None
+
+
+def _find_docs_dir() -> Path | None:
+    """Find a fetched docs directory (from `axolotl fetch docs`)."""
+    # axolotl fetch docs --dest defaults to ./docs/ in cwd
+    cwd_docs = Path.cwd() / "docs" / "agents"
+    if cwd_docs.is_dir():
+        return Path.cwd()
+    return None
+
+
+def _resolve_path(topic: str) -> Path:
+    """Resolve a topic name to the actual file path."""
+    if topic not in TOPICS:
+        available = ", ".join(sorted(TOPICS.keys()))
+        raise FileNotFoundError(f"Unknown topic: {topic!r}. Available: {available}")
+
+    relative_path = TOPICS[topic]
+
+    # Try repo root first (editable install / repo checkout)
+    repo_root = _find_repo_root()
+    if repo_root:
+        candidate = repo_root / relative_path
+        if candidate.exists():
+            return candidate
+
+    # Try cwd (fetched docs via `axolotl fetch docs`)
+    docs_root = _find_docs_dir()
+    if docs_root:
+        candidate = docs_root / relative_path
+        if candidate.exists():
+            return candidate
+
+    # Also check cwd directly for AGENTS.md
+    if topic == "overview":
+        cwd_agents = Path.cwd() / "AGENTS.md"
+        if cwd_agents.exists():
+            return cwd_agents
+
+    raise FileNotFoundError(
+        f"Could not find {relative_path!r}. "
+        f"If you installed axolotl via pip, run `axolotl fetch docs` first "
+        f"to download the documentation."
+    )
+
+
+def get_doc(topic: str = "overview") -> str:
+    """Return the content of an agent doc by topic name.
+
+    Args:
+        topic: One of the keys in TOPICS, or "overview" (default).
+
+    Returns:
+        The markdown content of the doc.
+
+    Raises:
+        FileNotFoundError: If the topic can't be found.
+    """
+    return _resolve_path(topic).read_text()
+
+
+def list_topics() -> dict[str, str]:
+    """Return a dict of topic name -> first line (title) of each doc."""
+    result = {}
+    for topic in sorted(TOPICS.keys()):
+        try:
+            path = _resolve_path(topic)
+            first_line = path.read_text().split("\n", 1)[0].lstrip("# ").strip()
+            result[topic] = first_line
+        except FileNotFoundError:
+            result[topic] = "(not found — run `axolotl fetch docs`)"
+    return result
diff --git a/src/axolotl/cli/args.py b/src/axolotl/cli/args.py
new file mode 100644
index 0000000000..14dafa43f1
--- /dev/null
+++ b/src/axolotl/cli/args.py
@@ -0,0 +1,134 @@
+"""Module for axolotl CLI command arguments."""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class PreprocessCliArgs:
+    """Dataclass with CLI arguments for `axolotl preprocess` command."""
+
+    debug: bool = field(default=False)
+    debug_text_only: bool = field(default=False)
+    debug_num_examples: int = field(default=1)
+    prompter: Optional[str] = field(default=None)
+    download: Optional[bool] = field(default=True)
+    iterable: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": (
+                "Deprecated in v0.13.0, will be removed in v0.14.0. For streaming "
+                "datasets, use 'axolotl train' and set 'streaming: true' in your YAML "
+                "config, or pass --streaming instead in the CLI."
+            )
+        },
+    )
+
+
+@dataclass
+class TrainerCliArgs:
+    """Dataclass with CLI arguments for `axolotl train` command."""
+
+    debug: bool = field(default=False)
+    debug_text_only: bool = field(default=False)
+    debug_num_examples: int = field(default=0)
+    prompter: Optional[str] = field(default=None)
+    shard: bool = field(default=False)
+
+
+@dataclass
+class VllmServeCliArgs:
+    """Dataclass with CLI arguments for `axolotl vllm-serve` command."""
+
+    tensor_parallel_size: Optional[int] = field(
+        default=None,
+        metadata={"help": "Number of tensor parallel workers to use."},
+    )
+    data_parallel_size: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of data parallel workers to use for vLLM serving. This controls how many model replicas are used for parallel inference."
+        },
+    )
+    host: Optional[str] = field(
+        default=None,  # nosec B104
+        metadata={"help": "Host address to run the server on."},
+    )
+    port: Optional[int] = field(
+        default=None,
+        metadata={"help": "Port to run the server on."},
+    )
+    gpu_memory_utilization: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV "
+            "cache on the device dedicated to generation powered by vLLM. Higher values will increase the KV cache "
+            "size and thus improve the model's throughput. However, if the value is too high, it may cause "
+            "out-of-memory (OOM) errors during initialization."
+        },
+    )
+    dtype: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Data type to use for vLLM generation. If set to 'auto', the data type will be automatically "
+            "determined based on the model configuration. Find the supported values in the vLLM documentation."
+        },
+    )
+    max_model_len: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "If set, the `max_model_len` to use for vLLM. This can be useful when running with reduced "
+            "`vllm_gpu_memory_utilization`, leading to a reduced KV cache size. If not set, vLLM will use the model "
+            "context size, which might be much larger than the KV cache, leading to inefficiencies."
+        },
+    )
+    enable_prefix_caching: Optional[bool] = field(
+        default=None,
+        metadata={
+            "help": "Whether to enable prefix caching in vLLM. If set to `True`, ensure that the model and the "
+            "hardware support this feature."
+        },
+    )
+    serve_module: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Module to serve. If not set, the default module will be used."
+        },
+    )
+
+    enable_reasoning: Optional[bool] = field(
+        default=None,
+    )
+
+    reasoning_parser: Optional[str] = field(
+        default=None,
+    )
+
+
+@dataclass
+class QuantizeCliArgs:
+    """Dataclass with CLI arguments for `axolotl quantize` command."""
+
+    base_model: Optional[str] = field(default=None)
+    weight_dtype: Optional[str] = field(default=None)
+    activation_dtype: Optional[str] = field(default=None)
+    quantize_embedding: Optional[bool] = field(default=None)
+    group_size: Optional[int] = field(default=None)
+    output_dir: Optional[str] = field(default=None)
+    hub_model_id: Optional[str] = field(default=None)
+
+
+@dataclass
+class EvaluateCliArgs:
+    """Dataclass with CLI arguments for `axolotl evaluate` command."""
+
+    debug: bool = field(default=False)
+    debug_text_only: bool = field(default=False)
+    debug_num_examples: int = field(default=0)
+
+
+@dataclass
+class InferenceCliArgs:
+    """Dataclass with CLI arguments for `axolotl inference` command."""
+
+    prompter: Optional[str] = field(default=None)
diff --git a/src/axolotl/cli/art.py b/src/axolotl/cli/art.py
new file mode 100644
index 0000000000..81dbb98313
--- /dev/null
+++ b/src/axolotl/cli/art.py
@@ -0,0 +1,30 @@
+"""Axolotl ASCII logo utils."""
+
+from axolotl.utils.distributed import is_main_process
+
+AXOLOTL_LOGO = """
+     #@@ #@@      @@# @@#
+    @@  @@          @@  @@           =@@#                               @@                 #@    =@@#.
+    @@    #@@@@@@@@@    @@           #@#@=                              @@                 #@     .=@@
+      #@@@@@@@@@@@@@@@@@            =@# @#     ##=     ##    =####=+    @@      =#####+  =#@@###.   @@
+    @@@@@@@@@@/  +@@/  +@@          #@  =@=     #@=   @@   =@#+  +#@#   @@    =@#+  +#@#   #@.      @@
+    @@@@@@@@@@  ##@@  ##@@         =@#   @#      =@# @#    @@      @@   @@    @@      #@   #@       @@
+     @@@@@@@@@@@@@@@@@@@@          #@=+++#@=      =@@#     @@      @@   @@    @@      #@   #@       @@
+                                  =@#=====@@     =@# @#    @@      @@   @@    @@      #@   #@       @@
+    @@@@@@@@@@@@@@@@  @@@@        #@      #@=   #@=  +@@   #@#    =@#   @@.   =@#    =@#   #@.      @@
+                                 =@#       @#  #@=     #@   =#@@@@#=    +#@@=  +#@@@@#=    .##@@+   @@
+    @@@@  @@@@@@@@@@@@@@@@
+"""
+
+HAS_PRINTED_LOGO = False
+
+
+def print_axolotl_text_art():
+    """Prints axolotl ASCII art."""
+
+    global HAS_PRINTED_LOGO
+    if HAS_PRINTED_LOGO:
+        return
+    if is_main_process():
+        HAS_PRINTED_LOGO = True
+        print(AXOLOTL_LOGO)
diff --git a/src/axolotl/cli/chat.py b/src/axolotl/cli/chat.py
new file mode 100644
index 0000000000..c54fe262d2
--- /dev/null
+++ b/src/axolotl/cli/chat.py
@@ -0,0 +1,1279 @@
+"""Interactive multi-turn chat CLI for a trained model."""
+
+import difflib
+import json
+import shlex
+import sys
+import threading
+import time
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable
+
+import torch
+from rich.console import Console
+from rich.live import Live
+from rich.markup import escape
+from rich.text import Text
+from transformers import (
+    DynamicCache,
+    GenerationConfig,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer,
+)
+
+from axolotl.cli.args import InferenceCliArgs
+from axolotl.cli.utils import load_model_and_tokenizer, resolve_chat_template_str
+from axolotl.integrations.base import PluginManager
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+USER_PROMPT = ">>> "
+CONTINUATION_PROMPT = "... "
+
+# marker pairs used by the bundled chat templates (qwen3/exaone4/phi_4 vs command_a)
+THINK_MARKER_PAIRS: tuple[tuple[str, str], ...] = (
+    ("<think>", "</think>"),
+    ("<|START_THINKING|>", "<|END_THINKING|>"),
+)
+DEFAULT_THINK_MARKERS = THINK_MARKER_PAIRS[0]
+
+
+def detect_think_markers(chat_template_str: str | None) -> tuple[str, str]:
+    """Picks the thinking marker pair the template works with. Called once at startup."""
+    if chat_template_str:
+        for pair in THINK_MARKER_PAIRS:
+            if pair[1] in chat_template_str:
+                return pair
+    return DEFAULT_THINK_MARKERS
+
+
+def detect_think_toggle_key(chat_template_str: str | None) -> str | None:
+    """
+    Finds the jinja variable the template uses to toggle thinking at render time
+    (`enable_thinking` in our bundled gemma4/qwen3_5 templates; `thinking` on some
+    hub templates). Called once at startup.
+    """
+    if not chat_template_str:
+        return None
+    if "enable_thinking" in chat_template_str:
+        return "enable_thinking"
+    if "thinking" in chat_template_str:
+        return "thinking"
+    return None
+
+
+@dataclass(frozen=True)
+class GenParamSpec:
+    """Specification of a runtime-adjustable generation parameter."""
+
+    key: str
+    cast: Callable
+    lo: float
+    hi: float
+    default: Any
+    aliases: tuple[str, ...] = ()
+    nullable: bool = False
+    help: str = ""
+
+
+GEN_PARAMS: tuple[GenParamSpec, ...] = (
+    GenParamSpec(
+        "temperature", float, 0.0, 5.0, 0.9, ("temp",), help="0 = greedy decoding"
+    ),
+    GenParamSpec("top_p", float, 0.0, 1.0, 0.95),
+    GenParamSpec("top_k", int, 0, 1000, 40),
+    GenParamSpec("min_p", float, 0.0, 1.0, None, nullable=True),
+    GenParamSpec("max_new_tokens", int, 1, 1_000_000, 1024, ("max_tokens", "max")),
+    GenParamSpec("repetition_penalty", float, 0.5, 3.0, 1.1, ("rep",)),
+    GenParamSpec(
+        "seed", int, 0, 2**32 - 1, None, nullable=True, help="`/set seed none` clears"
+    ),
+)
+
+
+DIFFUSION_GEN_PARAMS: tuple[GenParamSpec, ...] = (
+    GenParamSpec(
+        "temperature", float, 0.0, 5.0, 0.0, ("temp",), help="0 = greedy denoising"
+    ),
+    GenParamSpec(
+        "max_new_tokens",
+        int,
+        1,
+        100_000,
+        256,
+        ("tokens", "max_tokens", "max"),
+        help="size of the denoised completion block",
+    ),
+    GenParamSpec("steps", int, 1, 10_000, 128, help="number of denoising steps"),
+    GenParamSpec(
+        "seed", int, 0, 2**32 - 1, None, nullable=True, help="`/set seed none` clears"
+    ),
+)
+
+
+def default_gen_params(
+    specs: tuple[GenParamSpec, ...] = GEN_PARAMS,
+) -> dict[str, Any]:
+    return {spec.key: spec.default for spec in specs}
+
+
+def resolve_gen_param(
+    name: str, specs: tuple[GenParamSpec, ...] = GEN_PARAMS
+) -> GenParamSpec | None:
+    for spec in specs:
+        if name == spec.key or name in spec.aliases:
+            return spec
+    return None
+
+
+def parse_gen_param_value(spec: GenParamSpec, raw: str) -> Any:
+    if spec.nullable and raw.lower() in ("none", "null", "off"):
+        return None
+    try:
+        value = spec.cast(raw)
+    except ValueError as err:
+        raise ValueError(f"{spec.key} expects a {spec.cast.__name__}") from err
+    if not spec.lo <= value <= spec.hi:
+        raise ValueError(f"{spec.key} must be in [{spec.lo}, {spec.hi}]")
+    return value
+
+
+def longest_common_prefix_len(a: list[int], b: list[int]) -> int:
+    n = min(len(a), len(b))
+    for i in range(n):
+        if a[i] != b[i]:
+            return i
+    return n
+
+
+def find_subsequence(haystack: list[int], needle: list[int], start: int = 0) -> int:
+    if not needle:
+        return -1
+    n = len(needle)
+    for i in range(start, len(haystack) - n + 1):
+        if haystack[i : i + n] == needle:
+            return i
+    return -1
+
+
+def partial_suffix_len(text: str, marker: str) -> int:
+    """Length of the longest suffix of `text` that is a proper prefix of `marker`."""
+    max_len = min(len(text), len(marker) - 1)
+    for k in range(max_len, 0, -1):
+        if text.endswith(marker[:k]):
+            return k
+    return 0
+
+
+def content_as_text(content: Any) -> str:
+    # parse_response may return content as a list of parts rather than a string
+    if isinstance(content, list):
+        return "".join(
+            part.get("text", "") for part in content if isinstance(part, dict)
+        )
+    return content or ""
+
+
+@dataclass
+class ChatSession:
+    """Holds the conversation state for a chat session."""
+
+    messages: list[dict] = field(default_factory=list)
+    system: str | None = None
+
+    def conversation(self) -> list[dict]:
+        prefix = [{"role": "system", "content": self.system}] if self.system else []
+        return prefix + self.messages
+
+    def add_user(self, content: str):
+        # merge into a trailing unanswered user message (e.g. after a failed
+        # generation) so strict templates never see consecutive user turns
+        if self.messages and self.messages[-1]["role"] == "user":
+            self.messages[-1]["content"] += "\n" + content
+            return
+        self.messages.append({"role": "user", "content": content})
+
+    def add_assistant(self, content: str):
+        self.add_assistant_message({"role": "assistant", "content": content})
+
+    def add_assistant_message(self, message: dict):
+        self.messages.append(message)
+
+    def clear(self):
+        self.messages = []
+
+    def undo(self) -> bool:
+        """Removes the last user/assistant exchange. Returns False if empty."""
+        if not self.messages:
+            return False
+        if self.messages[-1]["role"] == "assistant":
+            self.messages.pop()
+        if self.messages and self.messages[-1]["role"] == "user":
+            self.messages.pop()
+        return True
+
+    def drop_last_assistant(self) -> bool:
+        """Removes the trailing assistant message so the turn can be retried."""
+        if self.messages and self.messages[-1]["role"] == "assistant":
+            self.messages.pop()
+        return bool(self.messages) and self.messages[-1]["role"] == "user"
+
+    def save_jsonl(self, path: str):
+        # content-parts format: text-only today, but matches the multimodal
+        # dataset format so saved sessions stay usable as training data
+        messages = []
+        for message in self.conversation():
+            content = message.get("content")
+            parts = (
+                content
+                if isinstance(content, list)
+                else [{"type": "text", "text": content or ""}]
+            )
+            out = {
+                "role": message["role"],
+                "content": parts,
+            }
+            for key in ("reasoning_content", "thinking", "tool_calls"):
+                if message.get(key):
+                    out[key] = message[key]
+            messages.append(out)
+        with open(path, "a", encoding="utf-8") as file:
+            file.write(json.dumps({"messages": messages}, ensure_ascii=False) + "\n")
+
+
+@dataclass
+class TurnResult:
+    """Result of generating a single assistant turn."""
+
+    content: str
+    message: dict | None = None
+    interrupted: bool = False
+    prompt_tokens: int = 0
+    reused_tokens: int = 0
+    new_tokens: int = 0
+    thinking_tokens: int = 0
+    response_tokens: int = 0
+    seconds: float = 0.0
+
+
+class _StopOnEvent(StoppingCriteria):
+    """Stops generation when the given event is set (e.g. on Ctrl+C)."""
+
+    def __init__(self, event: threading.Event):
+        self.event = event
+
+    def __call__(self, input_ids, scores, **kwargs) -> bool:
+        return self.event.is_set()
+
+
+class TurnGenerator:
+    """Base for assistant-turn generators: template rendering and EOS handling."""
+
+    def __init__(self, model, tokenizer, chat_template_str: str | None, device):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.chat_template_str = chat_template_str
+        self.device = device
+        self.think_markers = detect_think_markers(
+            chat_template_str or getattr(tokenizer, "chat_template", None)
+        )
+        self._think_marker_ids: tuple[list[int], list[int]] | None = None
+        self._eos_strings: tuple[str, ...] | None = None
+
+        self.eos_token_ids: set[int] = set()
+        if tokenizer.eos_token_id is not None:
+            self.eos_token_ids.add(tokenizer.eos_token_id)
+        config_eos = getattr(model.generation_config, "eos_token_id", None)
+        if isinstance(config_eos, int):
+            self.eos_token_ids.add(config_eos)
+        elif isinstance(config_eos, (list, tuple)):
+            self.eos_token_ids.update(config_eos)
+
+    def render(
+        self, conversation: list[dict], render_kwargs: dict | None = None
+    ) -> list[int]:
+        kwargs = dict(render_kwargs or {})
+        if self.chat_template_str:
+            kwargs["chat_template"] = self.chat_template_str
+        batch = self.tokenizer.apply_chat_template(
+            conversation,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            **kwargs,
+        )
+        return list(batch["input_ids"])
+
+    def _split_think_token_ids(
+        self, generated: list[int]
+    ) -> tuple[list[int], list[int]]:
+        """Splits a generated sequence into (thinking, response) token ids."""
+        if self._think_marker_ids is None:
+            try:
+                self._think_marker_ids = (
+                    self.tokenizer.encode(
+                        self.think_markers[0], add_special_tokens=False
+                    ),
+                    self.tokenizer.encode(
+                        self.think_markers[1], add_special_tokens=False
+                    ),
+                )
+            except Exception:  # pylint: disable=broad-exception-caught
+                self._think_marker_ids = ([], [])
+
+        open_ids, close_ids = self._think_marker_ids
+        i = find_subsequence(generated, open_ids)
+        if i < 0:
+            return [], generated
+        j = find_subsequence(generated, close_ids, i + len(open_ids))
+        if j < 0:
+            return generated[i + len(open_ids) :], generated[:i]
+        thinking = generated[i + len(open_ids) : j]
+        response = generated[:i] + generated[j + len(close_ids) :]
+        return thinking, response
+
+    def split_think_token_counts(self, generated: list[int]) -> tuple[int, int]:
+        """Returns (thinking, response) token counts for a generated sequence."""
+        thinking, response = self._split_think_token_ids(generated)
+        return len(thinking), len(response)
+
+    def build_assistant_message(
+        self,
+        generated: list[int],
+        split: tuple[list[int], list[int]] | None = None,
+    ) -> dict:
+        """
+        Parses generated token ids into an assistant message dict. Prefers the
+        tokenizer's own `parse_response` schema (transformers v5); otherwise splits
+        thinking out of the content by marker and stores it under
+        `reasoning_content`, the key the bundled chat templates read. Special
+        tokens are kept out of the stored text either way — the template re-adds
+        them on render.
+        """
+        if getattr(self.tokenizer, "response_schema", None):
+            try:
+                text = self.tokenizer.decode(generated, skip_special_tokens=False)
+                message = self.tokenizer.parse_response(text)
+                if isinstance(message, dict):
+                    message.setdefault("role", "assistant")
+                    message.setdefault("content", "")
+                    return message
+            except Exception:  # pylint: disable=broad-exception-caught
+                LOG.warning(
+                    "tokenizer.parse_response failed; falling back to marker split",
+                    exc_info=True,
+                )
+
+        thinking_ids, response_ids = (
+            split if split is not None else self._split_think_token_ids(generated)
+        )
+        parsed: dict[str, Any] = {
+            "role": "assistant",
+            "content": self.tokenizer.decode(
+                response_ids, skip_special_tokens=True
+            ).strip(),
+        }
+        if thinking_ids:
+            parsed["reasoning_content"] = self.tokenizer.decode(
+                thinking_ids, skip_special_tokens=True
+            ).strip()
+        return parsed
+
+    def eos_strings(self) -> tuple[str, ...]:
+        if self._eos_strings is None:
+            self._eos_strings = tuple(
+                text
+                for token_id in sorted(self.eos_token_ids)
+                if (text := self.tokenizer.decode([token_id]))
+            )
+        return self._eos_strings
+
+    def generate_turn(
+        self,
+        conversation: list[dict],
+        params: dict[str, Any],
+        on_text: Callable[[str], None],
+        render_kwargs: dict | None = None,
+    ) -> TurnResult:
+        raise NotImplementedError
+
+
+class EosTextTrimmer:
+    """
+    Filters streamed text so terminal EOS markers (e.g. `<|im_end|>`) never reach
+    the display. Text that could be the start of an EOS string is held back until
+    disambiguated by the next chunk.
+    """
+
+    def __init__(self, eos_strings: tuple[str, ...], emit: Callable[[str], None]):
+        self.eos_strings = tuple(s for s in eos_strings if s)
+        self.emit = emit
+        self.pending = ""
+        self.done = False
+
+    def feed(self, text: str):
+        if self.done or not text:
+            return
+        self.pending += text
+        positions = [
+            idx for s in self.eos_strings if (idx := self.pending.find(s)) >= 0
+        ]
+        if positions:
+            if min(positions) > 0:
+                self.emit(self.pending[: min(positions)])
+            self.pending = ""
+            self.done = True
+            return
+        hold = max(
+            (partial_suffix_len(self.pending, s) for s in self.eos_strings),
+            default=0,
+        )
+        if len(self.pending) > hold:
+            self.emit(self.pending[: len(self.pending) - hold])
+            self.pending = self.pending[len(self.pending) - hold :]
+
+    def finish(self):
+        if not self.done and self.pending:
+            self.emit(self.pending)
+        self.pending = ""
+
+
+class CausalTurnGenerator(TurnGenerator):
+    """
+    Generates assistant turns with `model.generate`, re-using the KV cache across
+    turns when the rendered conversation extends the previously cached tokens.
+    """
+
+    def __init__(self, model, tokenizer, chat_template_str: str | None, device):
+        super().__init__(model, tokenizer, chat_template_str, device)
+        self._cache: DynamicCache | None = None
+        self._cached_ids: list[int] = []
+
+    def reset_cache(self):
+        self._cache = None
+        self._cached_ids = []
+
+    def _new_cache(self) -> DynamicCache:
+        return DynamicCache(config=self.model.config)
+
+    def _prepare_cache(self, ids: list[int]) -> int:
+        """
+        Crops or resets the cross-turn cache so it holds a strict prefix of `ids`.
+        Chat templates may rewrite earlier turns when re-rendering (e.g. stripping
+        prior-turn thinking blocks), so reuse is gated on a token-level prefix
+        check rather than assumed.
+
+        Returns the number of re-used prefix tokens.
+        """
+        common = longest_common_prefix_len(self._cached_ids, ids)
+        # generate() needs at least one uncached input token
+        keep = min(common, len(ids) - 1)
+
+        if self._cache is None or keep <= 0:
+            self._cache = self._new_cache()
+            self._cached_ids = []
+            return 0
+
+        if keep < len(self._cached_ids):
+            try:
+                self._cache.crop(keep)
+                self._cached_ids = self._cached_ids[:keep]
+            except Exception:  # pylint: disable=broad-exception-caught
+                # some cache layer types (e.g. sliding window) cannot crop
+                self._cache = self._new_cache()
+                self._cached_ids = []
+                return 0
+
+        return len(self._cached_ids)
+
+    def _build_generation_config(self, params: dict[str, Any]) -> GenerationConfig:
+        do_sample = params["temperature"] > 0
+        kwargs: dict[str, Any] = {
+            "max_new_tokens": params["max_new_tokens"],
+            "repetition_penalty": params["repetition_penalty"],
+            "do_sample": do_sample,
+            "bos_token_id": self.tokenizer.bos_token_id,
+            "eos_token_id": sorted(self.eos_token_ids) or None,
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "use_cache": True,
+            "return_dict_in_generate": True,
+        }
+        if do_sample:
+            kwargs["temperature"] = params["temperature"]
+            kwargs["top_p"] = params["top_p"]
+            kwargs["top_k"] = params["top_k"]
+            if params["min_p"] is not None:
+                kwargs["min_p"] = params["min_p"]
+        return GenerationConfig(**kwargs)
+
+    def generate_turn(
+        self,
+        conversation: list[dict],
+        params: dict[str, Any],
+        on_text: Callable[[str], None],
+        render_kwargs: dict | None = None,
+    ) -> TurnResult:
+        ids = self.render(conversation, render_kwargs)
+        reused = self._prepare_cache(ids)
+        cache = self._cache
+        assert cache is not None
+
+        if params["seed"] is not None:
+            torch.manual_seed(params["seed"])
+
+        input_ids = torch.tensor([ids], dtype=torch.long, device=self.device)
+        attention_mask = torch.ones_like(input_ids)
+        generation_config = self._build_generation_config(params)
+
+        stop_event = threading.Event()
+        streamer = TextIteratorStreamer(
+            self.tokenizer, skip_prompt=True, skip_special_tokens=False
+        )
+        holder: dict[str, Any] = {}
+
+        def _worker():
+            try:
+                with torch.no_grad():
+                    holder["output"] = self.model.generate(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        generation_config=generation_config,
+                        streamer=streamer,
+                        stopping_criteria=StoppingCriteriaList(
+                            [_StopOnEvent(stop_event)]
+                        ),
+                        past_key_values=cache,
+                    )
+            except Exception as err:  # pylint: disable=broad-exception-caught
+                holder["error"] = err
+                streamer.end()
+
+        start = time.monotonic()
+        thread = threading.Thread(target=_worker, daemon=True)
+        thread.start()
+
+        trimmer = EosTextTrimmer(self.eos_strings(), on_text)
+        interrupted = False
+        try:
+            for text in streamer:
+                trimmer.feed(text)
+        except KeyboardInterrupt:
+            interrupted = True
+            stop_event.set()
+            for text in streamer:
+                trimmer.feed(text)
+        finally:
+            # a 2nd Ctrl+C can escape the drain; join so the worker stops writing the cache
+            stop_event.set()
+            thread.join()
+            trimmer.finish()
+        seconds = time.monotonic() - start
+
+        if "error" in holder:
+            self.reset_cache()
+            raise holder["error"]
+
+        sequence = holder["output"].sequences[0].tolist()
+        self._cached_ids = sequence[: cache.get_seq_length()]
+
+        generated = sequence[len(ids) :]
+        while generated and generated[-1] in self.eos_token_ids:
+            generated.pop()
+        thinking_ids, response_ids = self._split_think_token_ids(generated)
+        message = self.build_assistant_message(generated, (thinking_ids, response_ids))
+
+        return TurnResult(
+            content=message.get("content") or "",
+            message=message,
+            interrupted=interrupted,
+            prompt_tokens=len(ids),
+            reused_tokens=reused,
+            new_tokens=len(generated),
+            thinking_tokens=len(thinking_ids),
+            response_tokens=len(response_ids),
+            seconds=seconds,
+        )
+
+
+class DiffusionTurnGenerator(TurnGenerator):
+    """
+    Generates assistant turns for diffusion LMs by appending a masked completion
+    block to the rendered conversation and denoising it. The whole block resolves
+    at once, so the reply is emitted in one piece rather than streamed.
+    """
+
+    def __init__(
+        self, model, tokenizer, chat_template_str: str | None, device, mask_token_id
+    ):
+        super().__init__(model, tokenizer, chat_template_str, device)
+        self.mask_token_id = int(mask_token_id)
+
+    def generate_turn(
+        self,
+        conversation: list[dict],
+        params: dict[str, Any],
+        on_text: Callable[[str], None],
+        render_kwargs: dict | None = None,
+    ) -> TurnResult:
+        from axolotl.integrations.diffusion import generate as diffusion_generate
+
+        ids = self.render(conversation, render_kwargs)
+        if params["seed"] is not None:
+            torch.manual_seed(params["seed"])
+
+        sequence = torch.tensor([ids], dtype=torch.long, device=self.device)
+
+        start = time.monotonic()
+        with torch.no_grad():
+            result = diffusion_generate(
+                self.model,
+                self.tokenizer,
+                original_sequence=sequence,
+                num_diffusion_steps=params["steps"],
+                temperature=params["temperature"],
+                mask_token_id=self.mask_token_id,
+                mode="completion",
+                completion_tokens=params["max_new_tokens"],
+            )
+        seconds = time.monotonic() - start
+
+        generated = result["generated_ids"][len(ids) :]
+        for i, token_id in enumerate(generated):
+            if token_id in self.eos_token_ids:
+                generated = generated[:i]
+                break
+        content = self.tokenizer.decode(generated, skip_special_tokens=False)
+        on_text(content)
+        thinking_ids, response_ids = self._split_think_token_ids(generated)
+
+        return TurnResult(
+            content=content,
+            message=self.build_assistant_message(
+                generated, (thinking_ids, response_ids)
+            ),
+            prompt_tokens=len(ids),
+            new_tokens=len(generated),
+            thinking_tokens=len(thinking_ids),
+            response_tokens=len(response_ids),
+            seconds=seconds,
+        )
+
+
+class ThinkStreamRenderer:
+    """
+    Renders one streamed turn. When collapsing, thinking is shown as a rolling
+    dim tail in a live region (so it never enters scrollback) and replaced by a
+    one-line summary when the block closes; the reply streams normally. When
+    collapse is off, this is a plain passthrough print.
+    """
+
+    LIVE_FPS = 12
+
+    def __init__(
+        self,
+        console: Console,
+        collapse: bool,
+        markers: tuple[str, str] = DEFAULT_THINK_MARKERS,
+        tail_lines: int = 6,
+    ):
+        self.console = console
+        self.collapse = collapse
+        self.open_marker, self.close_marker = markers
+        self.tail_lines = tail_lines
+        self.think_text = ""
+        self._mode = "detect"
+        self._pending = ""
+        self._start = time.monotonic()
+        self._live: Live | None = None
+        self._last_live_update = 0.0
+
+    def feed(self, text: str):
+        if not self.collapse:
+            print(text, end="", flush=True)
+            return
+        self._pending += text
+        self._process()
+
+    def finish(self, interrupted: bool = False):
+        if not self.collapse:
+            return
+        if self._mode == "think":
+            self.think_text += self._pending
+            self._pending = ""
+            reason = "interrupted" if interrupted else f"no {self.close_marker}"
+            self._end_think(f" ({escape(reason)})")
+        elif self._pending:
+            print(self._pending, end="", flush=True)
+            self._pending = ""
+
+    def _process(self):
+        while True:
+            if self._mode == "detect":
+                stripped = self._pending.lstrip()
+                if stripped.startswith(self.open_marker):
+                    idx = self._pending.find(self.open_marker)
+                    self._pending = self._pending[idx + len(self.open_marker) :]
+                    self._mode = "think"
+                    self._live = Live(
+                        Text(""),
+                        console=self.console,
+                        refresh_per_second=self.LIVE_FPS,
+                        transient=True,
+                    )
+                    self._live.start()
+                    continue
+                if not stripped or self.open_marker.startswith(stripped):
+                    return  # could still be a marker prefix; wait for more text
+                self._mode = "reply"
+                continue
+
+            if self._mode == "think":
+                idx = self._pending.find(self.close_marker)
+                if idx >= 0:
+                    self.think_text += self._pending[:idx]
+                    self._pending = self._pending[
+                        idx + len(self.close_marker) :
+                    ].lstrip("\n")
+                    self._end_think()
+                    self._mode = "reply"
+                    continue
+                keep = partial_suffix_len(self._pending, self.close_marker)
+                emit_until = len(self._pending) - keep
+                self.think_text += self._pending[:emit_until]
+                self._pending = self._pending[emit_until:]
+                self._update_live()
+                return
+
+            # reply mode
+            if self._pending:
+                print(self._pending, end="", flush=True)
+                self._pending = ""
+            return
+
+    def _update_live(self):
+        if self._live is None:
+            return
+        # Live repaints at LIVE_FPS; building renderables faster than that is wasted
+        now = time.monotonic()
+        if now - self._last_live_update < 1 / self.LIVE_FPS:
+            return
+        self._last_live_update = now
+        lines = self.think_text.splitlines()[-self.tail_lines :]
+        self._live.update(Text("\n".join(lines), style="dim"))
+
+    def _end_think(self, note: str = ""):
+        if self._live is not None:
+            self._live.stop()
+            self._live = None
+        seconds = time.monotonic() - self._start
+        self.console.print(
+            f"[dim]▸ thought for {seconds:.1f}s{note} · /expand to view[/dim]"
+        )
+
+
+@dataclass(frozen=True)
+class Command:
+    """A slash command with its aliases and handler."""
+
+    name: str
+    handler: str
+    help: str
+    aliases: tuple[str, ...] = ()
+    usage: str = ""
+
+
+COMMANDS: tuple[Command, ...] = (
+    Command("help", "cmd_help", "show this help", ("?",)),
+    Command(
+        "new",
+        "cmd_new",
+        "clear the conversation (keeps system prompt and params)",
+        ("clear", "reset"),
+    ),
+    Command(
+        "system",
+        "cmd_system",
+        "show, set, or clear the system prompt",
+        usage="/system [text|clear]",
+    ),
+    Command(
+        "set",
+        "cmd_set",
+        "set a generation parameter",
+        usage="/set <param> <value>",
+    ),
+    Command("status", "cmd_status", "show model and generation settings", ("params",)),
+    Command("history", "cmd_history", "show the conversation so far"),
+    Command("retry", "cmd_retry", "regenerate the last assistant reply", ("regen",)),
+    Command("undo", "cmd_undo", "remove the last exchange"),
+    Command(
+        "save",
+        "cmd_save",
+        "append conversation as a chat_template-format JSONL sample",
+        usage="/save [path]",
+    ),
+    Command(
+        "think",
+        "cmd_think",
+        "toggle template-level thinking, if the template supports it",
+        usage="/think [on|off|default]",
+    ),
+    Command(
+        "collapse",
+        "cmd_collapse",
+        "collapse thinking blocks in the display",
+        usage="/collapse [on|off]",
+    ),
+    Command("expand", "cmd_expand", "show the hidden thinking from the last reply"),
+    Command("quit", "cmd_quit", "exit chat", ("exit", "q")),
+)
+
+
+def resolve_command(name: str) -> Command | None:
+    for command in COMMANDS:
+        if name == command.name or name in command.aliases:
+            return command
+    return None
+
+
+class ChatRepl:
+    """Interactive chat loop: slash commands plus streamed model turns."""
+
+    def __init__(
+        self,
+        *,
+        generator: TurnGenerator,
+        session: ChatSession | None = None,
+        params: dict[str, Any] | None = None,
+        param_specs: tuple[GenParamSpec, ...] = GEN_PARAMS,
+        console: Console | None = None,
+        banner: dict[str, str] | None = None,
+        input_fn: Callable[[str], str] | None = None,
+        think_toggle_key: str | None = None,
+        collapse_thinking: bool = True,
+    ):
+        self.generator = generator
+        self.session = session or ChatSession()
+        self.param_specs = param_specs
+        self.params = params or default_gen_params(param_specs)
+        self.console = console or Console()
+        self.banner = banner or {}
+        self.input_fn = input_fn or input
+        self.think_toggle_key = think_toggle_key
+        self.collapse_thinking = collapse_thinking
+        self.render_kwargs: dict[str, Any] = {}
+        self.last_think_text: str | None = None
+
+    def run(self):
+        self._print_banner()
+        while True:
+            try:
+                line = self._read_line()
+            except EOFError:
+                break
+            except KeyboardInterrupt:
+                self.console.print("\n[dim]Use /quit to exit.[/dim]")
+                continue
+
+            line = line.strip()
+            if not line:
+                continue
+
+            if line.startswith("/"):
+                try:
+                    action = self._dispatch(line)
+                except Exception as err:  # pylint: disable=broad-exception-caught
+                    self.console.print(f"[red]Command failed: {escape(str(err))}[/red]")
+                    continue
+                if action == "quit":
+                    break
+                if action == "regenerate":
+                    self._generate_turn()
+                continue
+
+            self.session.add_user(line)
+            self._generate_turn()
+
+    def _read_line(self) -> str:
+        parts = []
+        prompt = USER_PROMPT
+        while True:
+            line = self.input_fn(prompt)
+            if line.endswith("\\"):
+                parts.append(line[:-1])
+                prompt = CONTINUATION_PROMPT
+                continue
+            parts.append(line)
+            break
+        return "\n".join(parts)
+
+    def _dispatch(self, line: str) -> str | None:
+        name, _, args = line[1:].partition(" ")
+        name = name.lower()
+        args = args.strip()
+
+        command = resolve_command(name)
+        if command:
+            return getattr(self, command.handler)(args)
+
+        # bare parameter shortcuts: /temp 0.7, /top_p 0.9, ...
+        spec = resolve_gen_param(name, self.param_specs)
+        if spec:
+            return self.cmd_set(f"{spec.key} {args}" if args else spec.key)
+
+        candidates = [
+            alias for command in COMMANDS for alias in (command.name, *command.aliases)
+        ] + [alias for spec in self.param_specs for alias in (spec.key, *spec.aliases)]
+        close = difflib.get_close_matches(name, candidates, n=1)
+        hint = f" Did you mean /{close[0]}?" if close else ""
+        self.console.print(f"[red]Unknown command /{name}.[/red]{hint}")
+        return None
+
+    def _generate_turn(self):
+        renderer = ThinkStreamRenderer(
+            self.console,
+            collapse=self.collapse_thinking,
+            markers=getattr(self.generator, "think_markers", DEFAULT_THINK_MARKERS),
+        )
+        try:
+            result = self.generator.generate_turn(
+                self.session.conversation(),
+                self.params,
+                renderer.feed,
+                render_kwargs=self.render_kwargs or None,
+            )
+        except KeyboardInterrupt:
+            # an escaped interrupt leaves the cache out of sync with _cached_ids
+            reset_cache = getattr(self.generator, "reset_cache", None)
+            if callable(reset_cache):
+                reset_cache()
+            renderer.finish(interrupted=True)
+            print()
+            self.console.print(
+                "[dim]Interrupted; reply discarded. Your message is kept —"
+                " /retry regenerates, /undo removes it.[/dim]"
+            )
+            return
+        except Exception as err:  # pylint: disable=broad-exception-caught
+            renderer.finish(interrupted=True)
+            self.console.print(f"\n[red]Generation failed: {escape(str(err))}[/red]")
+            self.console.print(
+                "[dim]Your message is kept — /retry regenerates, /undo removes it.[/dim]"
+            )
+            return
+
+        renderer.finish(interrupted=result.interrupted)
+        message = result.message or {"role": "assistant", "content": result.content}
+        self.last_think_text = (
+            renderer.think_text.strip()
+            or message.get("reasoning_content")
+            or message.get("thinking")
+            or None
+        )
+
+        print()
+        self.session.add_assistant_message(message)
+        token_summary = f"{result.new_tokens} tokens"
+        if result.thinking_tokens:
+            token_summary += (
+                f" ({result.thinking_tokens} thinking · {result.response_tokens} reply)"
+            )
+        stats = (
+            f"{token_summary} · {result.seconds:.1f}s · "
+            f"{result.prompt_tokens} prompt ({result.reused_tokens} cached)"
+        )
+        if result.interrupted:
+            stats += " · interrupted (partial reply kept; /retry to regenerate)"
+        self.console.print(f"[dim]{stats}[/dim]")
+
+    def _print_banner(self):
+        self.console.print("[bold]axolotl chat[/bold]")
+        for key, value in self.banner.items():
+            self.console.print(f"[dim]{key}:[/dim] {escape(value)}")
+        self.console.print(
+            "[dim]Type a message to chat, /help for commands, \\ at line end to"
+            " continue on the next line.[/dim]"
+        )
+
+    # --- command handlers (return "quit", "regenerate", or None) ---
+
+    def cmd_help(self, _args: str) -> None:
+        for command in COMMANDS:
+            names = "/" + command.name
+            if command.aliases:
+                names += " (" + ", ".join("/" + a for a in command.aliases) + ")"
+            usage = f" — {command.usage}" if command.usage else ""
+            self.console.print(f"  [bold]{names}[/bold]: {command.help}{usage}")
+        params = ", ".join(
+            "/" + s.key + ("".join(f" /{a}" for a in s.aliases))
+            for s in self.param_specs
+        )
+        self.console.print(f"  parameter shortcuts: {params}")
+        return None
+
+    def cmd_new(self, _args: str) -> None:
+        self.session.clear()
+        self.last_think_text = None
+        reset_cache = getattr(self.generator, "reset_cache", None)
+        if callable(reset_cache):
+            reset_cache()
+        self.console.print("[dim]Conversation cleared.[/dim]")
+        return None
+
+    def cmd_system(self, args: str) -> None:
+        if not args:
+            if self.session.system:
+                self.console.print(escape(self.session.system))
+            else:
+                self.console.print("[dim]No system prompt set.[/dim]")
+            return None
+        if args.lower() == "clear":
+            self.session.system = None
+            self.console.print("[dim]System prompt cleared.[/dim]")
+            return None
+        self.session.system = args
+        self.console.print("[dim]System prompt set.[/dim]")
+        return None
+
+    def cmd_set(self, args: str) -> str | None:
+        tokens = args.replace("=", " ").split()
+        if len(tokens) != 2:
+            self.console.print("[red]Usage: /set <param> <value>[/red]")
+            return None
+        spec = resolve_gen_param(tokens[0].lower(), self.param_specs)
+        if not spec:
+            valid = ", ".join(s.key for s in self.param_specs)
+            self.console.print(f"[red]Unknown parameter. Valid: {valid}[/red]")
+            return None
+        try:
+            self.params[spec.key] = parse_gen_param_value(spec, tokens[1])
+        except ValueError as err:
+            self.console.print(f"[red]{err}[/red]")
+            return None
+        self.console.print(f"[dim]{spec.key} = {self.params[spec.key]}[/dim]")
+        return None
+
+    def cmd_status(self, _args: str) -> None:
+        for key, value in self.banner.items():
+            self.console.print(f"[dim]{key}:[/dim] {escape(value)}")
+        for spec in self.param_specs:
+            self.console.print(f"[dim]{spec.key}:[/dim] {self.params[spec.key]}")
+        n_messages = len(self.session.messages)
+        self.console.print(f"[dim]messages:[/dim] {n_messages}")
+        return None
+
+    def cmd_history(self, _args: str) -> None:
+        conversation = self.session.conversation()
+        if not conversation:
+            self.console.print("[dim]No messages yet.[/dim]")
+            return None
+        for message in conversation:
+            self.console.print(f"[bold]{message['role']}:[/bold]")
+            reasoning = message.get("reasoning_content") or message.get("thinking")
+            if reasoning:
+                self.console.print(escape(content_as_text(reasoning)), style="dim")
+            self.console.print(escape(content_as_text(message.get("content"))))
+        return None
+
+    def cmd_retry(self, _args: str) -> str | None:
+        if not self.session.drop_last_assistant():
+            self.console.print("[dim]Nothing to retry yet.[/dim]")
+            return None
+        return "regenerate"
+
+    def cmd_undo(self, _args: str) -> None:
+        if self.session.undo():
+            self.last_think_text = None
+            self.console.print("[dim]Removed last exchange.[/dim]")
+        else:
+            self.console.print("[dim]Nothing to undo.[/dim]")
+        return None
+
+    def cmd_save(self, args: str) -> None:
+        if not self.session.messages:
+            self.console.print("[dim]Nothing to save yet.[/dim]")
+            return None
+        path = (
+            shlex.split(args)[0]
+            if args
+            else f"chat-{datetime.now().strftime('%Y%m%d-%H%M%S')}.jsonl"
+        )
+        self.session.save_jsonl(path)
+        self.console.print(f"[dim]Saved conversation to {path}[/dim]")
+        return None
+
+    def cmd_think(self, args: str) -> None:
+        if not args:
+            current = self.render_kwargs.get(self.think_toggle_key or "", "default")
+            self.console.print(f"[dim]template thinking: {current}[/dim]")
+            return None
+        if not self.think_toggle_key:
+            self.console.print(
+                "[dim]This chat template has no thinking toggle; thinking is"
+                " controlled by the model/template itself.[/dim]"
+            )
+            return None
+        value = args.lower()
+        if value in ("on", "true"):
+            self.render_kwargs[self.think_toggle_key] = True
+        elif value in ("off", "false"):
+            self.render_kwargs[self.think_toggle_key] = False
+        elif value == "default":
+            self.render_kwargs.pop(self.think_toggle_key, None)
+        else:
+            self.console.print("[red]Usage: /think [on|off|default][/red]")
+            return None
+        current = self.render_kwargs.get(self.think_toggle_key, "default")
+        self.console.print(
+            f"[dim]{self.think_toggle_key} = {current} (applies from next turn)[/dim]"
+        )
+        return None
+
+    def cmd_collapse(self, args: str) -> None:
+        value = args.lower()
+        if value in ("on", "true", ""):
+            self.collapse_thinking = True
+        elif value in ("off", "false"):
+            self.collapse_thinking = False
+        else:
+            self.console.print("[red]Usage: /collapse [on|off][/red]")
+            return None
+        state = "collapsed" if self.collapse_thinking else "shown raw"
+        self.console.print(f"[dim]Thinking blocks will be {state}.[/dim]")
+        return None
+
+    def cmd_expand(self, _args: str) -> None:
+        if self.last_think_text:
+            self.console.print(escape(self.last_think_text), style="dim")
+        else:
+            self.console.print(
+                "[dim]No hidden thinking recorded for the last reply.[/dim]"
+            )
+        return None
+
+    def cmd_quit(self, _args: str) -> str:
+        return "quit"
+
+
+def _build_banner(cfg: DictDefault) -> dict[str, str]:
+    banner = {"model": str(cfg.base_model)}
+
+    if cfg.lora_model_dir:
+        banner["adapter"] = f"{cfg.adapter or 'lora'} from {cfg.lora_model_dir}"
+    elif cfg.adapter:
+        banner["adapter"] = str(cfg.adapter)
+
+    quant = []
+    if cfg.load_in_4bit:
+        quant.append("4-bit (bnb)")
+    if cfg.load_in_8bit:
+        quant.append("8-bit (bnb)")
+    if cfg.qat:
+        quant.append("QAT fake-quant active")
+    if quant:
+        banner["quantization"] = ", ".join(quant)
+
+    if cfg.chat_template:
+        template_name = getattr(cfg.chat_template, "value", cfg.chat_template)
+        banner["chat template"] = f"config ({template_name})"
+    elif cfg.datasets and cfg.datasets[0].type == "chat_template":
+        banner["chat template"] = "dataset config"
+    else:
+        banner["chat template"] = "tokenizer default"
+
+    return banner
+
+
+@send_errors
+def do_chat(
+    *,
+    cfg: DictDefault,
+    cli_args: InferenceCliArgs,
+):
+    """
+    Runs an interactive multi-turn chat session on the command line. The chat
+    template is applied to the full conversation each turn, and generation
+    parameters can be adjusted at runtime via slash commands.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        cli_args: Inference-specific CLI arguments.
+    """
+    if cli_args.prompter:
+        raise ValueError(
+            "--chat does not support --prompter; legacy prompters are single-turn."
+            " Use the default inference mode instead."
+        )
+
+    plugin_manager = PluginManager.get_instance()
+    is_diffusion = any(
+        plugin.__class__.__name__ == "DiffusionPlugin"
+        for plugin in plugin_manager.plugins.values()
+    )
+
+    if not sys.stdin.isatty():
+        raise ValueError(
+            "--chat requires an interactive terminal. For piped input, use the"
+            " default inference mode."
+        )
+
+    try:
+        import readline  # noqa: F401  pylint: disable=unused-import
+    except ImportError:
+        pass
+
+    model, tokenizer, _ = load_model_and_tokenizer(cfg=cfg, inference=True)
+    if cfg.is_multimodal:
+        LOG.warning(
+            "Multimodal attachments are not supported in chat mode yet;"
+            " proceeding with text-only chat."
+        )
+
+    chat_template_str = resolve_chat_template_str(cfg, tokenizer)
+    if not chat_template_str and not tokenizer.chat_template:
+        raise ValueError(
+            "Chat mode requires a chat template. Set `chat_template` in your config"
+            " or use a tokenizer that provides one."
+        )
+
+    model = model.to(cfg.device, dtype=cfg.torch_dtype)
+    model.eval()
+
+    banner = _build_banner(cfg)
+    generator: TurnGenerator
+    param_specs = GEN_PARAMS
+
+    if is_diffusion:
+        from axolotl.integrations.diffusion import resolve_mask_token_id
+
+        mask_token_id = resolve_mask_token_id(tokenizer, cfg, allow_add=False)
+        generator = DiffusionTurnGenerator(
+            model, tokenizer, chat_template_str, cfg.device, mask_token_id
+        )
+        param_specs = DIFFUSION_GEN_PARAMS
+        params = default_gen_params(param_specs)
+        if cfg.diffusion.num_diffusion_steps:
+            params["steps"] = cfg.diffusion.num_diffusion_steps
+        if cfg.diffusion.generation_temperature is not None:
+            params["temperature"] = cfg.diffusion.generation_temperature
+        banner["mode"] = "diffusion (completion-block denoising)"
+    else:
+        generator = CausalTurnGenerator(model, tokenizer, chat_template_str, cfg.device)
+        params = default_gen_params(param_specs)
+
+    think_toggle_key = detect_think_toggle_key(
+        chat_template_str or tokenizer.chat_template
+    )
+    repl = ChatRepl(
+        generator=generator,
+        params=params,
+        param_specs=param_specs,
+        banner=banner,
+        think_toggle_key=think_toggle_key,
+    )
+    repl.run()
diff --git a/src/axolotl/cli/checks.py b/src/axolotl/cli/checks.py
new file mode 100644
index 0000000000..4d1ab63e2d
--- /dev/null
+++ b/src/axolotl/cli/checks.py
@@ -0,0 +1,56 @@
+"""Various checks for Axolotl CLI."""
+
+import os
+from pathlib import Path
+
+import httpcore
+import httpx
+from accelerate.commands.config import config_args
+from huggingface_hub import HfApi
+from huggingface_hub.utils import LocalTokenNotFoundError
+from requests import HTTPError
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def check_accelerate_default_config() -> None:
+    """Logs at warning level if no accelerate config file is found."""
+    if Path(config_args.default_yaml_config_file).exists():
+        LOG.warning(
+            f"accelerate config file found at {config_args.default_yaml_config_file}. This can lead to unexpected errors"
+        )
+
+
+def check_user_token() -> bool:
+    """Checks for HF user info. Check is skipped if HF_HUB_OFFLINE=1.
+
+    Returns:
+        Boolean indicating successful check (i.e., HF_HUB_OFFLINE=1 or HF user info is retrieved).
+
+    Raises:
+        LocalTokenNotFoundError: If HF user info can't be retrieved.
+    """
+    # Skip check if HF_HUB_OFFLINE is set to True
+    if os.getenv("HF_HUB_OFFLINE") == "1":
+        LOG.info(
+            "Skipping HuggingFace token verification because HF_HUB_OFFLINE is set to True. Only local files will be used."
+        )
+        return True
+
+    # Verify if token is valid
+    api = HfApi()
+    try:
+        user_info = api.whoami()
+        return bool(user_info)
+    except LocalTokenNotFoundError:
+        LOG.warning(
+            "Error verifying HuggingFace token. Remember to log in using `hf auth login` and get your access token from https://huggingface.co/settings/tokens if you want to use gated models or datasets."
+        )
+        return False
+    except (HTTPError, httpcore.ConnectError, httpx.ConnectError):
+        LOG.warning(
+            "Error accessing HuggingFace. This may be due to a network issue or rate limiting."
+        )
+        return False
diff --git a/src/axolotl/cli/cloud/__init__.py b/src/axolotl/cli/cloud/__init__.py
new file mode 100644
index 0000000000..60f6a51ce4
--- /dev/null
+++ b/src/axolotl/cli/cloud/__init__.py
@@ -0,0 +1,74 @@
+"""
+launch axolotl in supported cloud platforms
+"""
+
+from pathlib import Path
+from typing import Literal
+
+import yaml
+
+from axolotl.cli.cloud.base import Cloud
+from axolotl.cli.cloud.baseten import BasetenCloud
+from axolotl.cli.cloud.modal_ import ModalCloud
+from axolotl.utils.dict import DictDefault
+
+
+def load_cloud_cfg(cloud_config: Path | str) -> DictDefault:
+    """Load and validate cloud configuration."""
+    # Load cloud configuration.
+    with open(cloud_config, encoding="utf-8") as file:
+        cloud_cfg: DictDefault = DictDefault(yaml.safe_load(file))
+    return cloud_cfg
+
+
+def do_cli_preprocess(
+    cloud_config: Path | str,
+    config: Path | str,
+) -> None:
+    cloud_cfg = load_cloud_cfg(cloud_config)
+    cloud = ModalCloud(cloud_cfg)
+    with open(config, "r", encoding="utf-8") as file:
+        config_yaml = file.read()
+    cloud.preprocess(config_yaml)
+
+
+def do_cli_train(
+    cloud_config: Path | str,
+    config: Path | str,
+    launcher: Literal["accelerate", "torchrun", "python"] = "accelerate",
+    launcher_args: list[str] | None = None,
+    cwd=None,
+    **kwargs,
+) -> None:
+    cloud_cfg: DictDefault = load_cloud_cfg(cloud_config)
+    provider = cloud_cfg.provider or "modal"
+    cloud: Cloud | None
+    if provider == "modal":
+        cloud = ModalCloud(cloud_cfg)
+    elif provider == "baseten":
+        cloud = BasetenCloud(cloud_cfg.to_dict())
+    else:
+        raise ValueError(f"Unsupported cloud provider: {provider}")
+    with open(config, "r", encoding="utf-8") as file:
+        config_yaml = file.read()
+    local_dirs = {}
+    if cwd and not Path(cwd).joinpath("src", "axolotl").exists():
+        local_dirs = {"/workspace/mounts": cwd}
+    cloud.train(
+        config_yaml,
+        launcher=launcher,
+        launcher_args=launcher_args,
+        local_dirs=local_dirs,
+        **kwargs,
+    )
+
+
+def do_cli_lm_eval(
+    cloud_config: Path | str,
+    config: Path | str,
+) -> None:
+    cloud_cfg = load_cloud_cfg(cloud_config)
+    cloud = ModalCloud(cloud_cfg)
+    with open(config, "r", encoding="utf-8") as file:
+        config_yaml = file.read()
+    cloud.lm_eval(config_yaml)
diff --git a/src/axolotl/cli/cloud/base.py b/src/axolotl/cli/cloud/base.py
new file mode 100644
index 0000000000..c498e8691e
--- /dev/null
+++ b/src/axolotl/cli/cloud/base.py
@@ -0,0 +1,27 @@
+"""
+base class for cloud platforms from cli
+"""
+
+from abc import ABC, abstractmethod
+from typing import Literal
+
+
+class Cloud(ABC):
+    """
+    Abstract base class for cloud platforms.
+    """
+
+    @abstractmethod
+    def preprocess(self, config_yaml: str, *args, **kwargs) -> None:
+        pass
+
+    @abstractmethod
+    def train(
+        self,
+        config_yaml: str,
+        launcher: Literal["accelerate", "torchrun", "python"] = "accelerate",
+        launcher_args: list[str] | None = None,
+        local_dirs: dict[str, str] | None = None,
+        **kwargs,
+    ):
+        pass
diff --git a/src/axolotl/cli/cloud/baseten/__init__.py b/src/axolotl/cli/cloud/baseten/__init__.py
new file mode 100644
index 0000000000..914504de3a
--- /dev/null
+++ b/src/axolotl/cli/cloud/baseten/__init__.py
@@ -0,0 +1,48 @@
+"""Baseten Cloud CLI"""
+
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from os.path import dirname
+from typing import Literal
+
+import yaml
+
+from axolotl.cli.cloud.base import Cloud
+
+
+class BasetenCloud(Cloud):
+    """Baseten Cloud Axolotl CLI"""
+
+    def __init__(self, config: dict):
+        self.config = config
+
+    def preprocess(self, config_yaml: str, *args, **kwargs) -> None:
+        raise NotImplementedError(
+            "Separate preprocess function for Baseten is not "
+            "implemented and will happen during hte train step."
+        )
+
+    def train(
+        self,
+        config_yaml: str,
+        launcher: Literal["accelerate", "torchrun", "python"] = "accelerate",
+        launcher_args: list[str] | None = None,
+        local_dirs: dict[str, str] | None = None,  # pylint: disable=unused-argument
+        **kwargs,
+    ):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            config = self.config.copy()
+            config["launcher"] = launcher
+            config["launcher_args"] = launcher_args
+            with open(tmp_dir + "/cloud.yaml", "w", encoding="utf-8") as cloud_fout:
+                yaml.dump(config, cloud_fout)
+            with open(tmp_dir + "/train.yaml", "w", encoding="utf-8") as config_fout:
+                config_fout.write(config_yaml)
+            shutil.copyfile(dirname(__file__) + "/template/run.sh", tmp_dir + "/run.sh")
+            shutil.copyfile(
+                dirname(__file__) + "/template/train_sft.py", tmp_dir + "/train_sft.py"
+            )
+            subprocess.run(  # nosec B603 B607
+                ["truss", "train", "push", "train_sft.py"], cwd=tmp_dir, check=False
+            )
diff --git a/src/axolotl/cli/cloud/baseten/template/run.sh b/src/axolotl/cli/cloud/baseten/template/run.sh
new file mode 100644
index 0000000000..37dc9688f2
--- /dev/null
+++ b/src/axolotl/cli/cloud/baseten/template/run.sh
@@ -0,0 +1,9 @@
+#!/bin/bash
+set -eux
+
+export NCCL_SOCKET_IFNAME="^docker0,lo"
+export NCCL_IB_DISABLE=0
+export NCCL_TIMEOUT=1800000
+
+axolotl preprocess train.yaml
+axolotl train train.yaml --launcher ${AXOLOTL_LAUNCHER} ${AXOLOTL_LAUNCHER_ARGS}
diff --git a/src/axolotl/cli/cloud/baseten/template/train_sft.py b/src/axolotl/cli/cloud/baseten/template/train_sft.py
new file mode 100644
index 0000000000..6dcf477c79
--- /dev/null
+++ b/src/axolotl/cli/cloud/baseten/template/train_sft.py
@@ -0,0 +1,70 @@
+"""
+Baseten Training Script for Axolotl
+"""
+
+# pylint: skip-file
+import yaml
+from truss.base import truss_config
+
+# Import necessary classes from the Baseten Training SDK
+from truss_train import definitions
+
+cloud_config = yaml.safe_load(open("cloud.yaml", "r"))
+gpu = cloud_config.get("gpu", "h100")
+gpu_count = int(cloud_config.get("gpu_count", 1))
+node_count = int(cloud_config.get("node_count", 1))
+project_name = cloud_config.get("project_name", "axolotl-project") or "axolotl-project"
+secrets = cloud_config.get("secrets", [])
+launcher = cloud_config.get("launcher", "accelerate")
+launcher_args = cloud_config.get("launcher_args", [])
+script_name = "run.sh"
+
+launcher_args_str = ""
+if launcher_args:
+    launcher_args_str = "-- " + " ".join(launcher_args)
+
+# 1. Define a base image for your training job
+BASE_IMAGE = "axolotlai/axolotl:main-py3.11-cu128-2.9.1"
+
+# 2. Define the Runtime Environment for the Training Job
+# This includes start commands and environment variables.a
+# Secrets from the baseten workspace like API keys are referenced using
+# `SecretReference`.
+
+env_vars = {
+    "AXOLOTL_LAUNCHER": launcher,
+    "AXOLOTL_LAUNCHER_ARGS": launcher_args_str,
+}
+for secret_name in secrets:
+    env_vars[secret_name] = definitions.SecretReference(name=secret_name)
+
+training_runtime = definitions.Runtime(
+    start_commands=[  # Example: list of commands to run your training script
+        f"/bin/sh -c 'chmod +x ./{script_name} && ./{script_name}'"
+    ],
+    environment_variables=env_vars,
+)
+
+# 3. Define the Compute Resources for the Training Job
+training_compute = definitions.Compute(
+    node_count=node_count,
+    accelerator=truss_config.AcceleratorSpec(
+        accelerator=truss_config.Accelerator.H100,
+        count=gpu_count,
+    ),
+)
+
+# 4. Define the Training Job
+# This brings together the image, compute, and runtime configurations.
+my_training_job = definitions.TrainingJob(
+    image=definitions.Image(base_image=BASE_IMAGE),
+    compute=training_compute,
+    runtime=training_runtime,
+)
+
+
+# This config will be pushed using the Truss CLI.
+# The association of the job to the project happens at the time of push.
+first_project_with_job = definitions.TrainingProject(
+    name=project_name, job=my_training_job
+)
diff --git a/src/axolotl/cli/cloud/modal_.py b/src/axolotl/cli/cloud/modal_.py
new file mode 100644
index 0000000000..df8f8f215c
--- /dev/null
+++ b/src/axolotl/cli/cloud/modal_.py
@@ -0,0 +1,316 @@
+"""
+Modal Cloud support from CLI
+"""
+
+import copy
+import json
+import os
+import subprocess  # nosec B404
+from pathlib import Path
+from random import randint
+from typing import Literal
+
+import modal
+
+from axolotl.cli.cloud.base import Cloud
+
+
+def run_cmd(cmd: str, run_folder: str, volumes=None):
+    """Run a command inside a folder, with Modal Volume reloading before and commit on success."""
+    # Ensure volumes contain latest files.
+    if volumes:
+        for _, vol in volumes.items():
+            vol.reload()
+
+    # modal workaround so it doesn't use the automounted axolotl
+    new_env = copy.deepcopy(os.environ)
+
+    if "PYTHONPATH" in new_env:
+        paths = ["/workspace/mounts"]
+        for sub_python_path_str in new_env["PYTHONPATH"].split(":"):
+            sub_python_path = Path(sub_python_path_str)
+            if not sub_python_path.joinpath("src", "axolotl").exists():
+                # we don't want to use the automounted axolotl or unexpected behavior happens
+                paths.append(str(sub_python_path))
+        if paths:
+            new_env["PYTHONPATH"] = ":".join(paths)
+        else:
+            del new_env["PYTHONPATH"]
+
+    # Propagate errors from subprocess.
+    if exit_code := subprocess.call(  # nosec B603
+        cmd.split(), cwd=run_folder, env=new_env
+    ):
+        exit(exit_code)
+
+    # Commit writes to volume.
+    if volumes:
+        for _, vol in volumes.items():
+            vol.commit()
+
+
+class ModalCloud(Cloud):
+    """
+    Modal Cloud implementation.
+    """
+
+    def __init__(self, config, app=None):
+        self.config = config
+        if not app:
+            app = modal.App()
+        self.app = app
+
+        self.volumes = {}
+        if config.volumes:
+            for volume_config in config.volumes:
+                _, mount, vol = self.create_volume(volume_config)
+                self.volumes[mount] = (vol, volume_config)
+
+    def get_env(self):
+        res = {
+            "HF_DATASETS_CACHE": "/workspace/data/huggingface-cache/datasets",
+            "HF_HUB_CACHE": "/workspace/data/huggingface-cache/hub",
+        }
+
+        for key in self.config.get("env", []):
+            if isinstance(key, str):
+                if val := os.environ.get(key, ""):
+                    res[key] = val
+            elif isinstance(key, dict):
+                (key_, val) = list(key.items())[0]
+                res[key_] = val
+        return res
+
+    def get_image(self):
+        docker_tag = "main-py3.11-cu128-2.9.1"
+        if self.config.docker_tag:
+            docker_tag = self.config.docker_tag
+        docker_image = f"axolotlai/axolotl:{docker_tag}"
+
+        # grab the sha256 hash from docker hub for this image+tag
+        # this ensures that we always get the latest image for this tag, even if it's already cached
+        try:
+            manifest = subprocess.check_output(  # nosec
+                ["docker", "manifest", "inspect", docker_image],
+            ).decode("utf-8")
+            sha256_hash = json.loads(manifest)["manifests"][0]["digest"]
+        except subprocess.CalledProcessError:
+            sha256_hash = None
+
+        # create the image
+        if sha256_hash:
+            image = modal.Image.from_registry(f"axolotlai/axolotl@{sha256_hash}")
+        else:
+            image = modal.Image.from_registry(docker_image)
+
+        dockerfile_commands = []
+        if self.config.dockerfile_commands:
+            dockerfile_commands.extend(self.config.dockerfile_commands)
+
+        # branch
+        if self.config.branch:
+            dockerfile_commands.extend(
+                [
+                    # Random id for cache busting of branch commits
+                    f"RUN echo '{str(randint(0, 1000000))}'",  # nosec B311
+                    f"RUN cd /workspace/axolotl && git fetch && git checkout {self.config.branch} && git pull",
+                ]
+            )
+
+        if dockerfile_commands:
+            image = image.dockerfile_commands(dockerfile_commands)
+
+        if env := self.get_env():
+            image = image.env(env)
+
+        return image
+
+    def get_secrets(self):
+        res = []
+        if self.config.secrets:
+            for key in self.config.get("secrets", []):
+                if isinstance(key, str):
+                    if val := os.environ.get(key, ""):
+                        res.append(modal.Secret.from_dict({key: val}))
+                elif isinstance(key, dict):
+                    (key_, val) = list(key.items())[0]
+                    res.append(modal.Secret.from_dict({key_: val}))
+        return res
+
+    def create_volume(self, volume_config):
+        name = volume_config.name
+        mount = volume_config.mount
+        return name, mount, modal.Volume.from_name(name, create_if_missing=True)
+
+    def get_ephemeral_disk_size(self):
+        return 1000 * 525  # 1 TiB
+
+    def get_preprocess_timeout(self):
+        if self.config.timeout_preprocess:
+            return int(self.config.timeout_preprocess)
+        return 60 * 60 * 3  # 3 hours
+
+    def get_preprocess_memory(self):
+        memory = 128  # default to 128GiB
+        if self.config.memory:
+            memory = int(self.config.memory)
+        if self.config.memory_preprocess:
+            memory = int(self.config.memory_preprocess)
+        return 1024 * memory
+
+    def get_preprocess_env(self):
+        return self.app.function(
+            image=self.get_image(),
+            volumes={k: v[0] for k, v in self.volumes.items()},
+            cpu=8.0,
+            ephemeral_disk=self.get_ephemeral_disk_size(),
+            memory=self.get_preprocess_memory(),
+            timeout=self.get_preprocess_timeout(),
+            secrets=self.get_secrets(),
+        )
+
+    def preprocess(self, config_yaml: str, *args, **kwargs):
+        modal_fn = self.get_preprocess_env()(_preprocess)
+        with modal.enable_output():
+            with self.app.run(detach=True):
+                modal_fn.remote(
+                    config_yaml,
+                    *args,
+                    volumes={k: v[0] for k, v in self.volumes.items()},
+                    **kwargs,
+                )
+
+    def get_train_timeout(self):
+        if self.config.timeout:
+            return int(self.config.timeout)
+        return 60 * 60 * 24  # 24 hours
+
+    def get_train_gpu(self):
+        count = self.config.gpu_count or 1
+        family = self.config.gpu.lower() or "l40s"
+
+        if family == "l40s":
+            return modal.gpu.L40S(count=count)
+        if family in ["a100", "a100-40gb"]:
+            return modal.gpu.A100(count=count, size="40GB")
+        if family == "a100-80gb":
+            return modal.gpu.A100(count=count, size="80GB")
+        if family in ["a10", "a10g"]:
+            return modal.gpu.A10G(count=count)
+        if family == "h100":
+            return f"H100:{count}"
+        if family == "t4":
+            return modal.gpu.T4(count=count)
+        if family == "l4":
+            return modal.gpu.L4(count=count)
+        raise ValueError(f"Unsupported GPU family: {family}")
+
+    def get_train_memory(self):
+        memory = 128  # default to 128GiB
+        if self.config.memory:
+            memory = int(self.config.memory)
+        return 1024 * memory
+
+    def get_train_env(self, local_dirs=None):
+        image = self.get_image()
+        for mount, local_dir in (local_dirs or {}).items():
+            image = image.add_local_dir(local_dir, mount)
+        return self.app.function(
+            image=image,
+            volumes={k: v[0] for k, v in self.volumes.items()},
+            cpu=16.0,
+            gpu=self.get_train_gpu(),
+            memory=self.get_train_memory(),
+            timeout=self.get_train_timeout(),
+            secrets=self.get_secrets(),
+        )
+
+    def train(
+        self,
+        config_yaml: str,
+        launcher: Literal["accelerate", "torchrun", "python"] = "accelerate",
+        launcher_args: list[str] | None = None,
+        local_dirs: dict[str, str] | None = None,
+        **kwargs,
+    ):
+        modal_fn = self.get_train_env(local_dirs)(_train)
+        with modal.enable_output():
+            with self.app.run(detach=True):
+                modal_fn.remote(
+                    config_yaml,
+                    launcher=launcher,
+                    launcher_args=launcher_args,
+                    volumes={k: v[0] for k, v in self.volumes.items()},
+                    **kwargs,
+                )
+
+    def lm_eval(self, config_yaml: str):
+        modal_fn = self.get_train_env()(_lm_eval)
+        with modal.enable_output():
+            with self.app.run(detach=True):
+                if self.config.get("spawn", False):
+                    modal_fn_exec = modal_fn.spawn
+                else:
+                    modal_fn_exec = modal_fn.remote
+                modal_fn_exec(
+                    config_yaml,
+                    volumes={k: v[0] for k, v in self.volumes.items()},
+                )
+
+
+def _preprocess(config_yaml: str, volumes=None):
+    Path("/workspace/mounts").mkdir(parents=True, exist_ok=True)
+    with open("/workspace/mounts/config.yaml", "w", encoding="utf-8") as f_out:
+        f_out.write(config_yaml)
+    run_folder = "/workspace/mounts"
+    run_cmd(
+        "axolotl preprocess /workspace/mounts/config.yaml --dataset-processes=8",
+        run_folder,
+        volumes,
+    )
+
+
+def _train(
+    config_yaml: str,
+    launcher: Literal["accelerate", "torchrun", "python"] = "accelerate",
+    launcher_args: list[str] | None = None,
+    volumes=None,
+    **kwargs,
+):
+    Path("/workspace/mounts").mkdir(parents=True, exist_ok=True)
+    with open("/workspace/mounts/config.yaml", "w", encoding="utf-8") as f_out:
+        f_out.write(config_yaml)
+    run_folder = "/workspace/mounts"
+
+    launcher_args = launcher_args or []
+
+    # Build the base command
+    if launcher == "accelerate":
+        launcher_arg = "--launcher accelerate"
+    elif launcher == "torchrun":
+        launcher_arg = "--launcher torchrun"
+    else:
+        launcher_arg = "--launcher python"
+
+    # Build launcher args string
+    launcher_args_str = ""
+    if launcher_args:
+        launcher_args_str = "-- " + " ".join(launcher_args)
+
+    run_cmd(
+        f"axolotl train {launcher_arg} /workspace/mounts/config.yaml {launcher_args_str}".strip(),
+        run_folder,
+        volumes,
+    )
+
+
+def _lm_eval(config_yaml: str, volumes=None):
+    Path("/workspace/mounts").mkdir(parents=True, exist_ok=True)
+    with open("/workspace/mounts/config.yaml", "w", encoding="utf-8") as f_out:
+        f_out.write(config_yaml)
+    run_folder = "/workspace/mounts"
+    run_cmd(
+        "axolotl lm-eval /workspace/mounts/config.yaml",
+        run_folder,
+        volumes,
+    )
diff --git a/src/axolotl/cli/config.py b/src/axolotl/cli/config.py
new file mode 100644
index 0000000000..666970b3e8
--- /dev/null
+++ b/src/axolotl/cli/config.py
@@ -0,0 +1,375 @@
+"""Configuration loading and processing."""
+
+import json
+import os
+import tempfile
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Any, Optional, Union
+from urllib.parse import urlparse
+
+import requests
+import torch
+import yaml
+from transformers.utils import is_torch_bf16_gpu_available, is_torch_tf32_available
+
+from axolotl.telemetry.errors import send_errors
+from axolotl.telemetry.manager import TelemetryManager
+from axolotl.utils.comet_ import setup_comet_env_vars
+from axolotl.utils.config import (
+    normalize_cfg_datasets,
+    normalize_config,
+    validate_config,
+)
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+from axolotl.utils.mlflow_ import setup_mlflow_env_vars
+from axolotl.utils.tee import prepare_debug_log
+from axolotl.utils.trackio_ import setup_trackio_env_vars
+from axolotl.utils.trainer import prepare_optim_env
+from axolotl.utils.wandb_ import setup_wandb_env_vars
+
+LOG = get_logger(__name__)
+
+
+def _coerce_value(value: Any, existing: Optional[Any] = None) -> Any:
+    """Coerce a string CLI value to its most likely Python type.
+
+    If an existing value is present in the config, its type is used to guide
+    casting.  Otherwise, YAML-style inference is applied: booleans, ints,
+    floats, and None literals are recognised automatically.
+
+    Args:
+        value: The raw value (typically a string from the CLI).
+        existing: An optional existing config value whose type guides coercion.
+
+    Returns:
+        The value cast to the inferred or expected type.
+    """
+    if not isinstance(value, str):
+        return value
+
+    # If the config already has a typed value, cast to match
+    if existing is not None:
+        if isinstance(existing, bool):
+            return value.lower() in ("true", "1", "yes")
+        if isinstance(existing, int):
+            try:
+                return int(value)
+            except (ValueError, TypeError):
+                return value
+        if isinstance(existing, float):
+            try:
+                return float(value)
+            except (ValueError, TypeError):
+                return value
+        # For other types (str, list, dict, etc.), return as-is
+        return value
+
+    # No existing value -- use YAML-style inference
+    lower = value.lower()
+    if lower in ("true", "yes"):
+        return True
+    if lower in ("false", "no"):
+        return False
+    if lower in ("null", "none", "~"):
+        return None
+
+    # Try int then float
+    try:
+        return int(value)
+    except ValueError:
+        pass
+    try:
+        return float(value)
+    except ValueError:
+        pass
+
+    return value
+
+
+API_KEY_FIELDS = {"comet_api_key"}
+
+TELEMETRY_MANAGER = TelemetryManager.get_instance()
+
+
+def check_remote_config(config: Union[str, Path]) -> Union[str, Path]:
+    """
+    First, determines if the passed config is a valid HTTPS URL. Then, attempts to query
+    for it and parse its content, first as JSON, then as YAML (YAML is preferred).
+    Finally, the parsed content is written to a local file and its path is returned.
+
+    Args:
+        config: HTTPS URL to a YAML or JSON file.
+
+    Returns:
+        Either the original `config` if it's not a valid HTTPS URL, or the path to the
+        downloaded remote config.
+
+    Raises:
+        ValueError: If the remote configuration is neither valid JSON or YAML.
+        RuntimeError: If some request-related exception occurs from the file download.
+        Exception: Catch-all for any other exception.
+    """
+    # Check if the config is a valid HTTPS URL to a .yml or .yaml file
+    if not (isinstance(config, str) and config.startswith("https://")):
+        return config  # Return the original value if it's not a valid URL
+
+    filename = os.path.basename(urlparse(config).path)
+    temp_dir = tempfile.mkdtemp()
+
+    try:
+        response = requests.get(config, timeout=30)
+        response.raise_for_status()  # Check for HTTP errors
+
+        content = response.content
+        try:
+            # Try parsing as JSON first to catch cases where JSON content is mistakenly
+            # considered YAML.
+            json.loads(content)
+
+            # Log a warning but do not raise an error; JSON is technically valid YAML.
+            # This can happen when you forget to point to a raw GitHub link.
+            LOG.warning(
+                f"Warning: The content of the file at {config} is JSON, which is technically valid YAML but might not be intended."
+            )
+        except json.JSONDecodeError:
+            # If it's not valid JSON, verify it's valid YAML
+            try:
+                yaml.safe_load(content)
+            except yaml.YAMLError as err:
+                raise ValueError(
+                    f"Failed to parse the content at {config} as YAML: {err}"
+                ) from err
+
+        # Write the content to a file if it's valid YAML (or JSON treated as YAML)
+        output_path = Path(temp_dir) / filename
+        with open(output_path, "wb") as file:
+            file.write(content)
+        LOG.info(
+            f"Using the following config obtained from {config}: \n\n{content.decode('utf-8')}\n"
+        )
+        return output_path
+
+    except requests.RequestException as err:
+        # This catches all requests-related exceptions including HTTPError
+        raise RuntimeError(f"Failed to download {config}: {err}") from err
+    except Exception as err:
+        # Catch-all for any other exceptions
+        raise err
+
+
+def choose_config(path: Path) -> str:
+    """
+    Helper method for choosing a `axolotl` config YAML file (considering only files
+    ending with `.yml` or `.yaml`). If more than one config file exists in the passed
+    `path`, the user is prompted to choose one.
+
+    Args:
+        path: Directory in which config file(s) are stored.
+
+    Returns:
+        Path to either (1) the sole YAML file, or (2) if more than one YAML files exist,
+        the user-selected YAML file.
+
+    Raises:
+        ValueError: If no YAML files are found in the given `path`.
+    """
+    yaml_files = list(path.glob("*.yml")) + list(path.glob("*.yaml"))
+
+    if not yaml_files:
+        raise ValueError(
+            "No YAML config files found in the specified directory. Are you using a .yml extension?"
+        )
+
+    if len(yaml_files) == 1:
+        LOG.info(f"Using default YAML file '{yaml_files[0]}'")
+        return str(yaml_files[0])
+
+    LOG.info("Choose a YAML file:")
+    for idx, file in enumerate(yaml_files):
+        LOG.info(f"{idx + 1}. {file}")
+
+    chosen_file = None
+    while chosen_file is None:
+        try:
+            choice = int(input("Enter the number of your choice: "))
+            if 1 <= choice <= len(yaml_files):
+                chosen_file = str(yaml_files[choice - 1])
+            else:
+                LOG.info("Invalid choice. Please choose a number from the list.")
+        except ValueError:
+            LOG.info("Invalid input. Please enter a number.")
+
+    return chosen_file
+
+
+def prepare_plugins(cfg: DictDefault):
+    """
+    Registers the plugins for the given configuration.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+    """
+    if cfg.get("plugins"):
+        from axolotl.integrations.base import PluginManager
+
+        plugin_manager = PluginManager.get_instance()
+        for plugin_name in cfg["plugins"]:
+            plugin_manager.register(plugin_name)
+        for plugin in plugin_manager.plugins.values():
+            plugin.register(cfg)
+
+
+def plugin_set_cfg(cfg: DictDefault):
+    if cfg.get("plugins"):
+        from axolotl.integrations.base import PluginManager
+
+        plugin_manager = PluginManager.get_instance()
+        plugin_manager.cfg = cfg
+
+
+@send_errors
+def load_cfg(
+    config: str | Path | DictDefault = Path("examples/"), **kwargs
+) -> DictDefault:
+    """
+    Loads the `axolotl` configuration stored at `config`, validates it, and performs
+    various setup.
+
+    Args:
+        config: Path (local or remote) to `axolotl` config YAML file.
+        kwargs: Additional keyword arguments to override config file values.
+
+    Returns:
+        `DictDefault` mapping configuration keys to values.
+    """
+    if isinstance(config, (str, Path)):
+        config = check_remote_config(config)
+        if Path(config).is_dir():
+            config = choose_config(Path(config))
+
+        # Load the config from the yaml file
+        with open(config, encoding="utf-8") as file:
+            cfg: DictDefault = DictDefault(yaml.safe_load(file))
+
+        cfg.axolotl_config_path = config
+    else:
+        cfg = config
+        with NamedTemporaryFile(
+            mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
+        ) as temp_file:
+            temp_file.write(yaml.dump(config.to_dict()))
+            temp_file.close()
+        cfg.axolotl_config_path = temp_file.name
+
+    TELEMETRY_MANAGER.send_event(event_type="config-loaded", properties=cfg)
+
+    # If there are any options passed in the cli, if it is something that seems valid
+    # from the yaml, then overwrite the value
+    cfg_keys = cfg.keys()
+
+    # Separate nested (dot-notation) kwargs from flat kwargs
+    nested_kwargs: dict[str, dict[str, Any]] = {}
+    flat_kwargs: dict[str, Any] = {}
+    for key, value in kwargs.items():
+        if "__" in key:
+            parent, child = key.split("__", 1)
+            nested_kwargs.setdefault(parent, {})[child] = value
+        else:
+            flat_kwargs[key] = value
+
+    # Apply flat kwargs
+    for key, value in flat_kwargs.items():
+        # If not strict, allow writing to cfg even if it's not in the yml already
+        if key in cfg_keys or not cfg.strict:
+            cfg[key] = _coerce_value(value, cfg.get(key))
+
+    # Apply nested kwargs (e.g., trl__beta -> cfg.trl.beta)
+    for parent, children in nested_kwargs.items():
+        if parent not in cfg_keys and cfg.strict:
+            continue
+        if cfg[parent] is None:
+            cfg[parent] = {}
+        if not isinstance(cfg[parent], dict):
+            LOG.warning(
+                "Overwriting non-dict value for '%s' with nested CLI overrides", parent
+            )
+            cfg[parent] = {}
+        for child_key, child_value in children.items():
+            existing_child = cfg[parent].get(child_key)
+            cfg[parent][child_key] = _coerce_value(child_value, existing_child)
+
+    prepare_plugins(cfg)
+
+    if cfg.use_ray:
+        # Ray drivers typically have no GPU; defer capability checks to the worker.
+        capabilities, env_capabilities = None, None
+    else:
+        capabilities, env_capabilities = gpu_capabilities()
+
+    cfg = validate_config(
+        cfg,
+        capabilities=capabilities,
+        env_capabilities=env_capabilities,
+    )
+
+    # NOTE(djsaunde): We start outputting to output_dir/debug.log at this point since we
+    # have to wait for cfg.output to be resolved. We could call this earlier if we write
+    # to a temporary file, and then move it later.
+    prepare_debug_log(cfg)
+    prepare_optim_env(cfg)
+    normalize_config(cfg)
+    normalize_cfg_datasets(cfg)
+    setup_wandb_env_vars(cfg)
+    setup_mlflow_env_vars(cfg)
+    setup_comet_env_vars(cfg)
+    setup_trackio_env_vars(cfg)
+    plugin_set_cfg(cfg)
+
+    TELEMETRY_MANAGER.send_event(event_type="config-processed", properties=cfg)
+    cfg_to_log = {
+        k: "[REDACTED]" if k in API_KEY_FIELDS else v
+        for k, v in cfg.items()
+        if v is not None
+    }
+    LOG.info(
+        "config:\n%s",
+        json.dumps(cfg_to_log, indent=2, default=str, sort_keys=True),
+    )
+
+    return cfg
+
+
+def compute_supports_fp8() -> bool:
+    try:
+        compute_capability = torch.cuda.get_device_capability()
+        return compute_capability >= (9, 0)
+    except (RuntimeError, AssertionError):
+        return False
+
+
+def gpu_capabilities() -> tuple[dict, dict]:
+    """Probe the local GPU and return ``(capabilities, env_capabilities)`` dicts
+    suitable for :func:`axolotl.utils.config.validate_config`.
+
+    Must be called on a GPU-enabled host (e.g. a Ray worker), otherwise the
+    detected values reflect the driver/CPU node and not the training device.
+    """
+    try:
+        device_props = torch.cuda.get_device_properties("cuda")
+        gpu_version = "sm_" + str(device_props.major) + str(device_props.minor)
+    except (RuntimeError, AssertionError):
+        gpu_version = None
+
+    capabilities = {
+        "bf16": is_torch_bf16_gpu_available(),
+        "fp8": compute_supports_fp8(),
+        "tf32": is_torch_tf32_available(),
+        "n_gpu": int(os.environ.get("WORLD_SIZE", 1)),
+        "compute_capability": gpu_version,
+    }
+    env_capabilities = {
+        "torch_version": str(torch.__version__).split("+", maxsplit=1)[0]
+    }
+    return capabilities, env_capabilities
diff --git a/src/axolotl/cli/config_options.py b/src/axolotl/cli/config_options.py
new file mode 100644
index 0000000000..d5fb17c4b2
--- /dev/null
+++ b/src/axolotl/cli/config_options.py
@@ -0,0 +1,3007 @@
+"""Generated Click option metadata for Axolotl config overrides."""
+
+# Generated by axolotl generate-cli-config-options.
+# Do not edit by hand.
+
+from __future__ import annotations
+
+AXOLOTL_CONFIG_CLI_OPTIONS = (
+    (
+        ("--max-packed-sequence-len",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--rope-scaling",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--noisy-embedding-alpha",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dpo-beta",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--evaluation-strategy",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--eval-table-size",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--eval-max-new-tokens",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dpo-use-logits-to-keep/--no-dpo-use-logits-to-keep",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dpo-generate-during-eval/--no-dpo-generate-during-eval",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dpo-norm-loss/--no-dpo-norm-loss",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--rpo-alpha",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--s2-attention/--no-s2-attention",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--flash-attn-rms-norm/--no-flash-attn-rms-norm",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--overrides-of-model-config",),
+        None,
+        None,
+        "optional overrides to the base model configuration",
+    ),
+    (
+        ("--overrides-of-model-kwargs",),
+        None,
+        None,
+        "optional overrides the base model loading from_pretrained",
+    ),
+    (
+        ("--type-of-model",),
+        None,
+        None,
+        "If you want to specify the type of model to load, AutoModelForCausalLM is a good choice too",
+    ),
+    (
+        ("--revision-of-model",),
+        None,
+        None,
+        "You can specify to choose a specific model revision from huggingface hub",
+    ),
+    (
+        ("--image-size",),
+        None,
+        None,
+        "The size of the image to resize to. It can be an integer (resized into padded-square image) or a tuple (width, height).If not provided, we will attempt to load from preprocessor.size, otherwise, images won't be resized.",
+    ),
+    (
+        ("--image-resize-algorithm",),
+        None,
+        None,
+        "The resampling algorithm to use for image resizing. Default is bilinear. Please refer to PIL.Image.Resampling for more details.",
+    ),
+    (
+        ("--role-boundaries",),
+        None,
+        None,
+        "Opt-in override for the MM mask scanner's per-role boundary markers. Non-empty list replaces built-ins wholesale; unset or empty falls back to built-ins. See docs/multimodal_assistant_mask.md.",
+    ),
+    (
+        ("--use-ray/--no-use-ray",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--ray-run-name",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--ray-num-workers",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--resources-per-worker",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--gradio-title",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--gradio-share/--no-gradio-share",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--gradio-server-name",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--gradio-server-port",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--gradio-max-new-tokens",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--gradio-temperature",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--lisa-n-layers",),
+        None,
+        None,
+        "the number of activate layers in LISA",
+    ),
+    (
+        ("--lisa-step-interval",),
+        None,
+        None,
+        "how often to switch layers in LISA",
+    ),
+    (
+        ("--lisa-layers-attribute",),
+        None,
+        None,
+        "path under the model to access the layers",
+    ),
+    (
+        ("--use-otel-metrics/--no-use-otel-metrics",),
+        None,
+        None,
+        "Enable OpenTelemetry metrics collection and Prometheus export",
+    ),
+    (
+        ("--otel-metrics-host",),
+        None,
+        None,
+        "Host to bind the OpenTelemetry metrics server to",
+    ),
+    (
+        ("--otel-metrics-port",),
+        None,
+        None,
+        "Port for the Prometheus metrics HTTP server",
+    ),
+    (
+        ("--use-trackio/--no-use-trackio",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--trackio-project-name",),
+        None,
+        None,
+        "Your trackio project name",
+    ),
+    (
+        ("--trackio-run-name",),
+        None,
+        None,
+        "Set the name of your trackio run",
+    ),
+    (
+        ("--trackio-space-id",),
+        None,
+        None,
+        "Hugging Face Space ID to sync dashboard to (optional, runs locally if not provided)",
+    ),
+    (
+        ("--use-comet/--no-use-comet",),
+        None,
+        None,
+        "Enable or disable Comet integration.",
+    ),
+    (
+        ("--comet-api-key",),
+        None,
+        None,
+        "API key for Comet. Recommended to set via `comet login`.",
+    ),
+    (
+        ("--comet-workspace",),
+        None,
+        None,
+        "Workspace name in Comet. Defaults to the user's default workspace.",
+    ),
+    (
+        ("--comet-project-name",),
+        None,
+        None,
+        "Project name in Comet. Defaults to Uncategorized.",
+    ),
+    (
+        ("--comet-experiment-key",),
+        None,
+        None,
+        "Identifier for the experiment. Used to append data to an existing experiment or control the key of new experiments. Default to a random key.",
+    ),
+    (
+        ("--comet-mode",),
+        None,
+        None,
+        'Create a new experiment ("create") or log to an existing one ("get"). Default ("get_or_create") auto-selects based on configuration.',
+    ),
+    (
+        ("--comet-online/--no-comet-online",),
+        None,
+        None,
+        "Set to True to log data to Comet server, or False for offline storage. Default is True.",
+    ),
+    (
+        ("--comet-experiment-config",),
+        None,
+        None,
+        "Dictionary for additional configuration settings, see the doc for more details.",
+    ),
+    (
+        ("--use-mlflow/--no-use-mlflow",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--mlflow-tracking-uri",),
+        None,
+        None,
+        "URI to mlflow",
+    ),
+    (
+        ("--mlflow-experiment-name",),
+        None,
+        None,
+        "Your experiment name",
+    ),
+    (
+        ("--mlflow-run-name",),
+        None,
+        None,
+        "Your run name",
+    ),
+    (
+        ("--hf-mlflow-log-artifacts/--no-hf-mlflow-log-artifacts",),
+        None,
+        None,
+        "set to true to copy each saved checkpoint on each save to mlflow artifact registry",
+    ),
+    (
+        ("--use-wandb/--no-use-wandb",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--wandb-name",),
+        None,
+        None,
+        "Set the name of your wandb run",
+    ),
+    (
+        ("--wandb-run-id",),
+        None,
+        None,
+        "Set the ID of your wandb run",
+    ),
+    (
+        ("--wandb-mode",),
+        None,
+        None,
+        '"offline" to save run metadata locally and not sync to the server, "disabled" to turn off wandb',
+    ),
+    (
+        ("--wandb-project",),
+        None,
+        None,
+        "Your wandb project name",
+    ),
+    (
+        ("--wandb-entity",),
+        None,
+        None,
+        "A wandb Team name if using a Team",
+    ),
+    (
+        ("--wandb-watch",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--wandb-log-model",),
+        None,
+        None,
+        '"checkpoint" to log model to wandb Artifacts every `save_steps` or "end" to log only at the end of training',
+    ),
+    (
+        ("--gradient-accumulation-steps",),
+        None,
+        None,
+        "If greater than 1, backpropagation will be skipped and the gradients will be accumulated for the given number of steps.",
+    ),
+    (
+        ("--micro-batch-size",),
+        None,
+        None,
+        "The number of samples to include in each batch. This is the number of samples sent to each GPU. Batch size per gpu = micro_batch_size * gradient_accumulation_steps",
+    ),
+    (
+        ("--batch-size",),
+        None,
+        None,
+        "Total batch size, we do not recommended setting this manually",
+    ),
+    (
+        ("--eval-batch-size",),
+        None,
+        None,
+        "per gpu micro batch size for evals, defaults to value of micro_batch_size",
+    ),
+    (
+        ("--auto-find-batch-size/--no-auto-find-batch-size",),
+        None,
+        None,
+        "whether to find batch size that fits in memory. Passed to underlying transformers Trainer",
+    ),
+    (
+        ("--train-on-inputs/--no-train-on-inputs",),
+        None,
+        None,
+        "Whether to mask out or include the human's prompt from the training labels",
+    ),
+    (
+        ("--group-by-length/--no-group-by-length",),
+        None,
+        None,
+        "Group similarly sized data to minimize padding. May be slower to start, as it must download and sort the entire dataset. Note that training loss may have an oscillating pattern with this enabled.",
+    ),
+    (
+        ("--learning-rate",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--embedding-lr",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--embedding-lr-scale",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--weight-decay",),
+        None,
+        None,
+        "Specify weight decay",
+    ),
+    (
+        ("--optimizer",),
+        None,
+        None,
+        "Specify optimizer",
+    ),
+    (
+        ("--optim-args",),
+        None,
+        None,
+        "Dictionary of arguments to pass to the optimizer",
+    ),
+    (
+        ("--optim-target-modules",),
+        None,
+        None,
+        "The target modules to optimize, i.e. the module names that you would like to train, right now this is used only for GaLore algorithm",
+    ),
+    (
+        ("--torchdistx-path",),
+        None,
+        None,
+        "Path to torch distx for optim 'adamw_anyprecision'",
+    ),
+    (
+        ("--lr-scheduler",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--lr-scheduler-kwargs",),
+        None,
+        None,
+        "Specify a scheduler and kwargs to use with the optimizer",
+    ),
+    (
+        ("--lr-quadratic-warmup/--no-lr-quadratic-warmup",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--cosine-min-lr-ratio",),
+        None,
+        None,
+        "decay lr to some percentage of the peak lr, e.g. cosine_min_lr_ratio=0.1 for 10% of peak lr",
+    ),
+    (
+        ("--cosine-constant-lr-ratio",),
+        None,
+        None,
+        "freeze lr at some percentage of the step, e.g. cosine_constant_lr_ratio=0.8 means start cosine_min_lr at 80% of training step",
+    ),
+    (
+        ("--lr-div-factor",),
+        None,
+        None,
+        "Learning rate div factor",
+    ),
+    (
+        ("--lr-groups",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--adam-epsilon",),
+        None,
+        None,
+        "adamw hyperparams",
+    ),
+    (
+        ("--adam-epsilon2",),
+        None,
+        None,
+        "only used for CAME Optimizer",
+    ),
+    (
+        ("--adam-beta1",),
+        None,
+        None,
+        "adamw hyperparams",
+    ),
+    (
+        ("--adam-beta2",),
+        None,
+        None,
+        "adamw hyperparams",
+    ),
+    (
+        ("--adam-beta3",),
+        None,
+        None,
+        "only used for CAME Optimizer",
+    ),
+    (
+        ("--dion-lr",),
+        None,
+        None,
+        "Dion Optimizer learning rate",
+    ),
+    (
+        ("--dion-momentum",),
+        None,
+        None,
+        "Dion Optimizer momentum",
+    ),
+    (
+        ("--dion-rank-fraction",),
+        None,
+        None,
+        "Dion Optimizer: r/d fraction for low-rank approximation. Used to compute the low-rank dimension.",
+    ),
+    (
+        ("--dion-rank-multiple-of",),
+        None,
+        None,
+        "Dion Optimizer: Round up the low-rank dimension to a multiple of this number. This may be useful to ensure even sharding.",
+    ),
+    (
+        ("--qgalore-rank",),
+        None,
+        None,
+        "Q-GaLore: rank r of the low-rank gradient projection. Smaller r reduces optimizer state but loses gradient information.",
+    ),
+    (
+        ("--qgalore-update-proj-gap",),
+        None,
+        None,
+        "Q-GaLore: maximum number of steps between SVD recomputations of the projection matrix. The adaptive scheduler may skip updates earlier based on cos_threshold.",
+    ),
+    (
+        ("--qgalore-scale",),
+        None,
+        None,
+        "Q-GaLore: scaling factor applied to the projected gradient after project_back. Equivalent to GaLore's `galore_scale`.",
+    ),
+    (
+        ("--qgalore-proj-type",),
+        None,
+        None,
+        "Q-GaLore: projection type for the GaLoreProjector. One of 'std', 'reverse_std', 'right', 'left', 'full'.",
+    ),
+    (
+        ("--qgalore-proj-quant/--no-qgalore-proj-quant",),
+        None,
+        None,
+        "Q-GaLore: enable INT-quantization of the projection matrix P (the resilient-to-quantization observation from the paper).",
+    ),
+    (
+        ("--qgalore-proj-bits",),
+        None,
+        None,
+        "Q-GaLore: bitwidth for the quantized projection matrix when qgalore_proj_quant is True (paper default: 4).",
+    ),
+    (
+        ("--qgalore-proj-group-size",),
+        None,
+        None,
+        "Q-GaLore: group size for projection-matrix quantization. Must evenly divide the projection's last dimension.",
+    ),
+    (
+        ("--qgalore-cos-threshold",),
+        None,
+        None,
+        "Q-GaLore: cosine-similarity threshold for the lazy subspace update. If the new P is within this similarity of the previous one, the SVD is skipped.",
+    ),
+    (
+        ("--qgalore-gamma-proj",),
+        None,
+        None,
+        "Q-GaLore: multiplicative factor by which update_proj_gap grows once a layer's subspace is judged stable.",
+    ),
+    (
+        ("--qgalore-queue-size",),
+        None,
+        None,
+        "Q-GaLore: length of the moving-average queue used by the adaptive-frequency scheduler.",
+    ),
+    (
+        ("--max-grad-norm",),
+        None,
+        None,
+        "Gradient clipping max norm",
+    ),
+    (
+        ("--num-epochs",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--jagged-restart-steps",),
+        None,
+        None,
+        "how often to reset for jagged restarts",
+    ),
+    (
+        ("--jagged-restart-warmup-steps",),
+        None,
+        None,
+        "how many warmup steps to take after reset for jagged restarts",
+    ),
+    (
+        ("--jagged-restart-anneal-steps",),
+        None,
+        None,
+        "how many anneal steps to take before reset for jagged restarts",
+    ),
+    (
+        ("--relora/--no-relora",),
+        None,
+        None,
+        "Whether to use ReLoRA. Use with jagged_restart_*steps options.",
+    ),
+    (
+        ("--relora-prune-ratio",),
+        None,
+        None,
+        "Fraction of optimizer state values to zero on each ReLoRA restart. When relora_prune_method='reset' and this is omitted, defaults to 0.999 (paper-style near-full reset). For other methods, defaults to 0.9.",
+    ),
+    (
+        ("--relora-prune-method",),
+        None,
+        None,
+        "Optimizer state pruning method on each ReLoRA restart. 'magnitude' (default) keeps top-k by absolute value; 'random' keeps a random subset at relora_prune_ratio; 'reset' uses near-full random pruning (default ratio 0.999, honoring relora_prune_ratio when explicitly set). Paper-style recipe: relora_prune_method='reset' with no relora_prune_ratio, equivalent to 'random' with ratio=0.999.",
+    ),
+    (
+        ("--relora-cpu-offload/--no-relora-cpu-offload",),
+        None,
+        None,
+        "True to perform lora weight merges on cpu during restarts, for modest gpu memory savings",
+    ),
+    (
+        ("--load-in-8bit/--no-load-in-8bit",),
+        None,
+        None,
+        "This will attempt to quantize the model down to 8 bits and use adam 8 bit optimizer",
+    ),
+    (
+        ("--load-in-4bit/--no-load-in-4bit",),
+        None,
+        None,
+        "Use bitsandbytes 4 bit",
+    ),
+    (
+        ("--adapter",),
+        None,
+        None,
+        "If you want to use a built-in or plugin adapter, or leave blank to train all parameters in original model",
+    ),
+    (
+        ("--lora-model-dir",),
+        None,
+        None,
+        "If you already have a lora model trained that you want to load, put that here. This means after training, if you want to test the model, you should set this to the value of `output_dir`. Note that if you merge an adapter to the base model, a new subdirectory `merged` will be created under the `output_dir`.",
+    ),
+    (
+        ("--lora-r",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--lora-alpha",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--lora-fan-in-fan-out/--no-lora-fan-in-fan-out",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--lora-target-modules",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--lora-target-parameters",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--lora-target-linear/--no-lora-target-linear",),
+        None,
+        None,
+        "If true, will target all linear modules",
+    ),
+    (
+        ("--lora-modules-to-save",),
+        None,
+        None,
+        "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. `embed_tokens` converts tokens to embeddings, and `lm_head` converts embeddings to token probabilities.",
+    ),
+    (
+        ("--lora-rank-pattern",),
+        None,
+        None,
+        "Per-module LoRA rank overrides (regex \u2192 int > 0). Forwarded to PEFT LoraConfig.rank_pattern.",
+    ),
+    (
+        ("--lora-alpha-pattern",),
+        None,
+        None,
+        "Per-module LoRA alpha overrides (regex \u2192 int > 0). Forwarded to PEFT LoraConfig.alpha_pattern.",
+    ),
+    (
+        ("--lora-dropout",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--peft-layers-to-transform",),
+        None,
+        None,
+        "The layer indices to transform, otherwise, apply to all layers",
+    ),
+    (
+        ("--peft-layers-pattern",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--peft.loftq-config",),
+        "peft__loftq_config",
+        None,
+        "Configuration options for loftq initialization for LoRA",
+    ),
+    (
+        ("--peft-use-dora/--no-peft-use-dora",),
+        None,
+        None,
+        "Whether to use DoRA.",
+    ),
+    (
+        ("--peft-use-rslora/--no-peft-use-rslora",),
+        None,
+        None,
+        "Whether to use RSLoRA.",
+    ),
+    (
+        ("--peft-layer-replication",),
+        None,
+        None,
+        "List of layer indices to replicate.",
+    ),
+    (
+        ("--peft-init-lora-weights/--no-peft-init-lora-weights",),
+        None,
+        None,
+        "How to initialize LoRA weights. Default to True which is MS original implementation.",
+    ),
+    (
+        ("--peft-trainable-token-indices",),
+        None,
+        None,
+        "A list of token indices to fine-tune on the `embed_tokens` layer.\nOtherwise, a dict mapping an embedding layer name to its trainable token indices.\nSee https://huggingface.co/docs/peft/v0.17.0/en/developer_guides/lora#efficiently-train-tokens-alongside-lora",
+    ),
+    (
+        ("--peft-ensure-weight-tying/--no-peft-ensure-weight-tying",),
+        None,
+        None,
+        "Whether to tie adapter weights for tied model weights. See https://github.com/huggingface/peft/issues/2864",
+    ),
+    (
+        ("--peft-autocast-adapter-dtype/--no-peft-autocast-adapter-dtype",),
+        None,
+        None,
+        "Whether to upcast the LoRA adapter to fp32. This is enabled by default in PEFT.",
+    ),
+    (
+        ("--qlora-sharded-model-loading/--no-qlora-sharded-model-loading",),
+        None,
+        None,
+        "load qlora model in sharded format for FSDP using answer.ai technique.",
+    ),
+    (
+        ("--lora-on-cpu/--no-lora-on-cpu",),
+        None,
+        None,
+        "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",
+    ),
+    (
+        ("--gptq/--no-gptq",),
+        None,
+        None,
+        "Whether you are training a 4-bit GPTQ quantized model",
+    ),
+    (
+        ("--bnb-config-kwargs",),
+        None,
+        None,
+        "optional overrides to the bnb 4bit quantization configuration",
+    ),
+    (
+        ("--loraplus-lr-ratio",),
+        None,
+        None,
+        "loraplus learning rate ratio lr_B / lr_A. Recommended value is 2^4.",
+    ),
+    (
+        ("--loraplus-lr-embedding",),
+        None,
+        None,
+        "loraplus learning rate for lora embedding layers. Default value is 1e-6.",
+    ),
+    (
+        ("--merge-lora/--no-merge-lora",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--merge-method",),
+        None,
+        None,
+        "Method to use for LoRA merging. 'memory_efficient' (default) processes shards individually to reduce memory usage, 'legacy' loads the full model into memory.",
+    ),
+    (
+        ("--output-dir",),
+        None,
+        None,
+        "Where to save the full-finetuned model to",
+    ),
+    (
+        ("--hub-model-id",),
+        None,
+        None,
+        "push checkpoints to hub",
+    ),
+    (
+        ("--hub-strategy",),
+        None,
+        None,
+        "how to push checkpoints to hub",
+    ),
+    (
+        ("--hub-revision",),
+        None,
+        None,
+        "branch/revision to push to on hub (default: main)",
+    ),
+    (
+        ("--save-safetensors/--no-save-safetensors",),
+        None,
+        None,
+        "Whether to save the model using safetensors format. Defaults to True.",
+    ),
+    (
+        ("--base-model",),
+        None,
+        None,
+        "This is the huggingface model that contains *.pt, *.safetensors, or *.bin files. This can also be a relative path to a model on disk",
+    ),
+    (
+        ("--base-model-config",),
+        None,
+        None,
+        "If the base_model repo on hf hub doesn't include configuration .json files, You can set that here, or leave this empty to default to base_model",
+    ),
+    (
+        ("--cls-model-config",),
+        None,
+        None,
+        "transformers config class (e.g., 'LlamaConfig', 'MistralConfig'). Defaults to AutoConfig.",
+    ),
+    (
+        ("--tokenizer-config",),
+        None,
+        None,
+        "Optional tokenizer configuration path in case you want to use a different tokenizer than the one defined in the base model",
+    ),
+    (
+        ("--tokenizer-use-fast/--no-tokenizer-use-fast",),
+        None,
+        None,
+        "use_fast option for tokenizer loading from_pretrained, default to True",
+    ),
+    (
+        ("--tokenizer-legacy/--no-tokenizer-legacy",),
+        None,
+        None,
+        "Whether to use the legacy tokenizer setting, defaults to True",
+    ),
+    (
+        ("--tokenizer-use-mistral-common/--no-tokenizer-use-mistral-common",),
+        None,
+        None,
+        "Whether to use mistral-common tokenizer. If set to True, it will use the mistral-common tokenizer.",
+    ),
+    (
+        ("--tokenizer-type",),
+        None,
+        None,
+        "Corresponding tokenizer for the model AutoTokenizer is a good choice",
+    ),
+    (
+        ("--processor-type",),
+        None,
+        None,
+        "transformers processor class",
+    ),
+    (
+        ("--processor-kwargs",),
+        None,
+        None,
+        "kwargs forwarded to the processor's from_pretrained(), overriding processor config (e.g. image_seq_length, min_pixels, etc.).",
+    ),
+    (
+        ("--tokenizer-save-jinja-files/--no-tokenizer-save-jinja-files",),
+        None,
+        None,
+        "Whether to save jinja files for tokenizer, transformers default is True",
+    ),
+    (
+        ("--trust-remote-code/--no-trust-remote-code",),
+        None,
+        None,
+        "Trust remote code for untrusted source",
+    ),
+    (
+        ("--experimental-skip-move-to-device/--no-experimental-skip-move-to-device",),
+        None,
+        None,
+        "Don't move the model to the device before sharding. Set to `false` to revert to legacy behavior.",
+    ),
+    (
+        ("--use-kernels/--no-use-kernels",),
+        None,
+        None,
+        "Use custom kernels, e.g. MegaBlocks.",
+    ),
+    (
+        ("--model-quantization-config",),
+        None,
+        None,
+        "Model loading quantization config",
+    ),
+    (
+        ("--model-quantization-config-kwargs",),
+        None,
+        None,
+        "kwargs for model quantization config",
+    ),
+    (
+        ("--use-onebitllms/--no-use-onebitllms",),
+        None,
+        None,
+        "Whether to use `onebitllms` for 1.58bit training (only for bitnet models).",
+    ),
+    (
+        ("--strict/--no-strict",),
+        None,
+        None,
+        "Allow overwrite yml config using from cli",
+    ),
+    (
+        ("--resume-from-checkpoint",),
+        None,
+        None,
+        "Resume from a specific checkpoint dir",
+    ),
+    (
+        ("--auto-resume-from-checkpoints/--no-auto-resume-from-checkpoints",),
+        None,
+        None,
+        "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.",
+    ),
+    (
+        ("--resize-token-embeddings-to-32x/--no-resize-token-embeddings-to-32x",),
+        None,
+        None,
+        "Resize the model embeddings when new tokens are added to multiples of 32. This is reported to improve training speed on some models",
+    ),
+    (
+        ("--mean-resizing-embeddings/--no-mean-resizing-embeddings",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--shrink-embeddings/--no-shrink-embeddings",),
+        None,
+        None,
+        "Whether to shrink the embeddings to len(tokenizer). By default, we won't shrink.",
+    ),
+    (
+        ("--embeddings-skip-upcast/--no-embeddings-skip-upcast",),
+        None,
+        None,
+        "Don't upcast the embeddings to float32 when using PEFT. Useful for low-VRAM GPUs",
+    ),
+    (
+        ("--reinit-weights/--no-reinit-weights",),
+        None,
+        None,
+        "Reinitialize model weights randomly instead of loading pretrained weights",
+    ),
+    (
+        ("--trainer-cls",),
+        None,
+        None,
+        "module to custom trainer class to use for training",
+    ),
+    (
+        ("--rl",),
+        None,
+        None,
+        "Use RL training: 'dpo', 'ipo', 'kto', 'simpo', 'orpo', 'grpo', 'ebft'",
+    ),
+    (
+        ("--trl.beta",),
+        "trl__beta",
+        "float",
+        "Beta parameter for the RL training. Same as `rl_beta`. Use",
+    ),
+    (
+        ("--trl.max-completion-length",),
+        "trl__max_completion_length",
+        "int",
+        "Maximum length of the completion for RL training.",
+    ),
+    (
+        ("--trl.use-vllm/--no-trl.use-vllm",),
+        "trl__use_vllm",
+        None,
+        "Whether to use VLLM for RL training.",
+    ),
+    (
+        ("--trl.vllm-mode",),
+        "trl__vllm_mode",
+        None,
+        "VLLM mode to use, one of 'server' or 'colocate'",
+    ),
+    (
+        ("--trl.vllm-server-host",),
+        "trl__vllm_server_host",
+        "str",
+        "Host of the vLLM server to connect to.",
+    ),
+    (
+        ("--trl.vllm-server-port",),
+        "trl__vllm_server_port",
+        "int",
+        "Port of the vLLM server to connect to.",
+    ),
+    (
+        ("--trl.vllm-server-timeout",),
+        "trl__vllm_server_timeout",
+        "int",
+        "Total timeout (in seconds) to wait for the vLLM server to respond.",
+    ),
+    (
+        ("--trl.vllm-guided-decoding-regex",),
+        "trl__vllm_guided_decoding_regex",
+        "str",
+        "Regex for vLLM guided decoding.",
+    ),
+    (
+        ("--trl.reward-funcs",),
+        "trl__reward_funcs",
+        None,
+        "List of reward functions to load. Paths must be importable from current dir.",
+    ),
+    (
+        ("--trl.reward-weights",),
+        "trl__reward_weights",
+        None,
+        "List of reward weights for the reward functions.",
+    ),
+    (
+        ("--trl.generation-batch-size",),
+        "trl__generation_batch_size",
+        "int",
+        "Batch size for generation. Controls how many unique prompts are generated per step. Should be num_generations * data_parallel_size for full DP utilization.",
+    ),
+    (
+        ("--trl.num-generations",),
+        "trl__num_generations",
+        "int",
+        "Number of generations to sample.",
+    ),
+    (
+        ("--trl.log-completions/--no-trl.log-completions",),
+        "trl__log_completions",
+        None,
+        "Whether to log completions.",
+    ),
+    (
+        ("--trl.num-completions-to-print",),
+        "trl__num_completions_to_print",
+        "int",
+        "Number of completions to print when log_completions is True.",
+    ),
+    (
+        ("--trl.importance-sampling-level",),
+        "trl__importance_sampling_level",
+        None,
+        "Controls whether importance sampling ratios are computed at the `'token'` or `'sequence'` level. For GSPO, use `sequence`, default is None which corresponds to the original GRPO paper.",
+    ),
+    (
+        ("--trl.sync-ref-model/--no-trl.sync-ref-model",),
+        "trl__sync_ref_model",
+        None,
+        "Whether to sync the reference model.",
+    ),
+    (
+        ("--trl.ref-model-mixup-alpha",),
+        "trl__ref_model_mixup_alpha",
+        "float",
+        "Mixup alpha for the reference model.",
+    ),
+    (
+        ("--trl.ref-model-sync-steps",),
+        "trl__ref_model_sync_steps",
+        "int",
+        "Sync steps for the reference model.",
+    ),
+    (
+        ("--trl.scale-rewards/--no-trl.scale-rewards",),
+        "trl__scale_rewards",
+        None,
+        "Whether to scale rewards by their standard deviation.",
+    ),
+    (
+        ("--trl.temperature",),
+        "trl__temperature",
+        "float",
+        "Sampling temperature for the GRPO policy.",
+    ),
+    (
+        ("--trl.top-p",),
+        "trl__top_p",
+        "float",
+        "Top-p sampling probability for the generation policy.",
+    ),
+    (
+        ("--trl.top-k",),
+        "trl__top_k",
+        "int",
+        "Top-k sampling for the generation policy.",
+    ),
+    (
+        ("--trl.min-p",),
+        "trl__min_p",
+        "float",
+        "Minimum probability for the generation policy.",
+    ),
+    (
+        ("--trl.repetition-penalty",),
+        "trl__repetition_penalty",
+        "float",
+        "Penalty for tokens that appear in prompt and generated text.",
+    ),
+    (
+        ("--trl.generation-kwargs",),
+        "trl__generation_kwargs",
+        None,
+        "Additional generation parameters passed to vLLM SamplingParams. Useful for stop_token_ids, seed, frequency_penalty, etc.",
+    ),
+    (
+        ("--trl.chat-template-kwargs",),
+        "trl__chat_template_kwargs",
+        None,
+        "Additional kwargs for the chat template. E.g., {enable_thinking: false} for Qwen3.5 models.",
+    ),
+    (
+        ("--trl.num-iterations",),
+        "trl__num_iterations",
+        "int",
+        "Number of iterations per batch (\u03bc) for GRPO.",
+    ),
+    (
+        ("--trl.epsilon",),
+        "trl__epsilon",
+        "float",
+        "Epsilon value for clipping in the GRPO algorithm.",
+    ),
+    (
+        ("--trl.epsilon-high",),
+        "trl__epsilon_high",
+        "float",
+        "Upper-bound epsilon value for clipping in the GRPO algorithm.",
+    ),
+    (
+        ("--trl.use-liger-loss/--no-trl.use-liger-loss",),
+        "trl__use_liger_loss",
+        None,
+        "Whether to use Liger loss for GRPO.",
+    ),
+    (
+        ("--trl.loss-type",),
+        "trl__loss_type",
+        "str",
+        "Loss formulation to use. Supported values: grpo, bnpo, dr_grpo.",
+    ),
+    (
+        ("--trl.mask-truncated-completions/--no-trl.mask-truncated-completions",),
+        "trl__mask_truncated_completions",
+        None,
+        "Whether to exclude truncated completions from loss calculation.",
+    ),
+    (
+        ("--trl.vllm-enable-sleep-mode/--no-trl.vllm-enable-sleep-mode",),
+        "trl__vllm_enable_sleep_mode",
+        None,
+        "Enable sleep mode for vLLM to offload VRAM when idle",
+    ),
+    (
+        ("--trl.rollout-func",),
+        "trl__rollout_func",
+        "str",
+        "Path to custom rollout function. Must be importable from current dir.",
+    ),
+    (
+        ("--trl.multi-objective-aggregation",),
+        "trl__multi_objective_aggregation",
+        None,
+        "Multi-objective reward aggregation strategy. 'sum_then_normalize' (GRPO default): weights and sums rewards first, then normalizes. 'normalize_then_sum' (GDPO): normalizes each reward independently, then sums.",
+    ),
+    (
+        ("--trl.use-data-producer/--no-trl.use-data-producer",),
+        "trl__use_data_producer",
+        None,
+        "Use the GRPODataProducer protocol for online data generation.",
+    ),
+    (
+        ("--trl.async-prefetch/--no-trl.async-prefetch",),
+        "trl__async_prefetch",
+        None,
+        "Generate rollouts in a background thread while training on the previous rollout.",
+    ),
+    (
+        ("--trl.prefetch-depth",),
+        "trl__prefetch_depth",
+        "int",
+        "Number of rollouts to prefetch ahead of training.",
+    ),
+    (
+        ("--trl.vllm-sync-interval",),
+        "trl__vllm_sync_interval",
+        "int",
+        "Sync model weights to vLLM every N optimizer steps (async mode only).",
+    ),
+    (
+        ("--trl.streaming-partial-batch/--no-trl.streaming-partial-batch",),
+        "trl__streaming_partial_batch",
+        None,
+        "Score prompt groups incrementally instead of the full batch at once.",
+    ),
+    (
+        ("--trl.streaming-min-groups",),
+        "trl__streaming_min_groups",
+        "int",
+        "Minimum prompt groups to score per streaming chunk.",
+    ),
+    (
+        (
+            "--trl.vllm-importance-sampling-correction/--no-trl.vllm-importance-sampling-correction",
+        ),
+        "trl__vllm_importance_sampling_correction",
+        None,
+        "Apply IS correction for distribution mismatch between vLLM and training model.",
+    ),
+    (
+        ("--trl.vllm-importance-sampling-mode",),
+        "trl__vllm_importance_sampling_mode",
+        None,
+        "IS mode: token_truncate, token_mask, sequence_truncate, or sequence_mask.",
+    ),
+    (
+        ("--trl.vllm-importance-sampling-cap",),
+        "trl__vllm_importance_sampling_cap",
+        "float",
+        "Cap C for IS ratio clipping/masking.",
+    ),
+    (
+        ("--trl.off-policy-mask-threshold",),
+        "trl__off_policy_mask_threshold",
+        "float",
+        "KL threshold for off-policy sequence masking (OPSM). None = disabled.",
+    ),
+    (
+        ("--trl.use-bias-correction-kl/--no-trl.use-bias-correction-kl",),
+        "trl__use_bias_correction_kl",
+        None,
+        "Apply IS correction to KL divergence term.",
+    ),
+    (
+        ("--trl.reward-num-workers",),
+        "trl__reward_num_workers",
+        "int",
+        "Number of persistent subprocess workers for parallel reward computation. Each worker has its own main thread so signal.alarm() (used by math_verify) works correctly. Work is sharded across workers by prompt groups. Only used with use_data_producer=True and non-nn.Module reward functions.",
+    ),
+    (
+        ("--trl.replay-buffer-size",),
+        "trl__replay_buffer_size",
+        "int",
+        "[Experimental, disabled by default] Size of the replay buffer for storing high-signal rollout groups. When > 0, groups with reward variance are cached and used to replace zero-signal groups (where all rewards are identical). Set to 0 to disable. Only used with use_data_producer=True.",
+    ),
+    (
+        ("--trl.replay-recompute-logps/--no-trl.replay-recompute-logps",),
+        "trl__replay_recompute_logps",
+        None,
+        "When True (default), recompute old_per_token_logps for replayed groups using the current training model. This fixes the importance sampling mismatch that occurs when replaying stale data. Only relevant when replay_buffer_size > 0.",
+    ),
+    (
+        ("--trl.reroll-start-fraction",),
+        "trl__reroll_start_fraction",
+        "float",
+        "Fraction of total training steps after which deferred re-rolling begins. Zero-signal prompts (where all rewards in a group are identical) are buffered and re-injected into later batches when the model is more likely to solve them. Set to 1.0 to disable. Only used with use_data_producer=True.",
+    ),
+    (
+        ("--trl.reroll-max-groups",),
+        "trl__reroll_max_groups",
+        "int",
+        "Maximum number of prompt groups to replace with re-roll candidates per batch. Higher values increase data utilization but reduce prompt diversity. Only used with use_data_producer=True.",
+    ),
+    (
+        ("--trl.skip-zero-advantage-batches/--no-trl.skip-zero-advantage-batches",),
+        "trl__skip_zero_advantage_batches",
+        None,
+        "When True, skip gradient computation for micro-batches where all advantages are zero (no learning signal). This avoids the forward/backward pass entirely when no learning signal is present. The step is logged with skipped_zero_adv_batches=1 for monitoring.",
+    ),
+    (
+        ("--trl.vllm-lora-sync/--no-trl.vllm-lora-sync",),
+        "trl__vllm_lora_sync",
+        None,
+        "Sync LoRA adapter to vLLM via filesystem instead of merging + NCCL broadcast. Auto-selects vllm_serve_lora serve module. Syncs only LoRA adapter weights vs full merged model.",
+    ),
+    (
+        ("--vllm.device",),
+        "vllm__device",
+        "str",
+        "Device to use for VLLM",
+    ),
+    (
+        ("--vllm.tensor-parallel-size",),
+        "vllm__tensor_parallel_size",
+        "int",
+        "Tensor parallel size for VLLM",
+    ),
+    (
+        ("--vllm.data-parallel-size",),
+        "vllm__data_parallel_size",
+        "int",
+        "Data parallel size for VLLM",
+    ),
+    (
+        ("--vllm.gpu-memory-utilization",),
+        "vllm__gpu_memory_utilization",
+        "float",
+        "GPU memory utilization for VLLM",
+    ),
+    (
+        ("--vllm.dtype",),
+        "vllm__dtype",
+        "str",
+        "Data type for VLLM",
+    ),
+    (
+        ("--vllm.max-model-len",),
+        "vllm__max_model_len",
+        "int",
+        "Maximum length of the model context for VLLM",
+    ),
+    (
+        ("--vllm.enable-prefix-caching/--no-vllm.enable-prefix-caching",),
+        "vllm__enable_prefix_caching",
+        None,
+        "Enable prefix caching for VLLM",
+    ),
+    (
+        ("--vllm.host",),
+        "vllm__host",
+        "str",
+        "Host for the vLLM server to start on",
+    ),
+    (
+        ("--vllm.port",),
+        "vllm__port",
+        "int",
+        "Port of the vLLM server to start on",
+    ),
+    (
+        ("--vllm.enable-reasoning/--no-vllm.enable-reasoning",),
+        "vllm__enable_reasoning",
+        None,
+        "Enable reasoning for VLLM",
+    ),
+    (
+        ("--vllm.reasoning-parser",),
+        "vllm__reasoning_parser",
+        "str",
+        "Reasoning parser for VLLM",
+    ),
+    (
+        ("--vllm.enforce-eager/--no-vllm.enforce-eager",),
+        "vllm__enforce_eager",
+        None,
+        "Disable CUDA graph capture in vLLM. Required for models with causal_conv1d (e.g., Qwen3.5 hybrid linear attention).",
+    ),
+    (
+        ("--vllm.serve-module",),
+        "vllm__serve_module",
+        "str",
+        "Python module for vLLM serve script. Set to 'axolotl.scripts.vllm_serve_lora' for native LoRA support, or leave None for default TRL serve.",
+    ),
+    (
+        ("--vllm.worker-extension-cls",),
+        "vllm__worker_extension_cls",
+        "str",
+        "vLLM worker extension class for weight synchronization. Defaults to 'trl.scripts.vllm_serve.WeightSyncWorkerExtension'.",
+    ),
+    (
+        ("--ebft.feature-layers",),
+        "ebft__feature_layers",
+        None,
+        "Fractional layer depths for feature extraction (e.g., [0.25, 0.5, 0.75])",
+    ),
+    (
+        ("--ebft.embed-method",),
+        "ebft__embed_method",
+        None,
+        "Embedding method: 'last_token', 'mean_pooling', 'completion_mean', or 'concat'",
+    ),
+    (
+        ("--ebft.use-whitening/--no-ebft.use-whitening",),
+        "ebft__use_whitening",
+        None,
+        "Apply SVD whitening to feature embeddings",
+    ),
+    (
+        ("--ebft.alignment-coef",),
+        "ebft__alignment_coef",
+        "float",
+        "Coefficient for alignment reward (cosine similarity with ground truth)",
+    ),
+    (
+        ("--ebft.diversity-coef",),
+        "ebft__diversity_coef",
+        "float",
+        "Coefficient for diversity penalty (pairwise similarity between samples)",
+    ),
+    (
+        ("--ebft.ce-coef",),
+        "ebft__ce_coef",
+        "float",
+        "Cross-entropy loss coefficient on ground-truth tokens",
+    ),
+    (
+        ("--ebft.adaptive-max-tokens/--no-ebft.adaptive-max-tokens",),
+        "ebft__adaptive_max_tokens",
+        None,
+        "Set per-batch max_tokens based on ground-truth length",
+    ),
+    (
+        ("--ebft.gt-length-multiplier",),
+        "ebft__gt_length_multiplier",
+        "float",
+        "Multiplier for ground-truth token count when computing adaptive max_tokens",
+    ),
+    (
+        ("--ebft.mode",),
+        "ebft__mode",
+        None,
+        "EBFT mode: 'structured' (QA with vLLM) or 'strided' (unstructured text)",
+    ),
+    (
+        ("--ebft.stride",),
+        "ebft__stride",
+        "int",
+        "Stride between anchor points (tokens)",
+    ),
+    (
+        ("--ebft.context-length",),
+        "ebft__context_length",
+        "int",
+        "Context window size per block",
+    ),
+    (
+        ("--ebft.generate-max-len",),
+        "ebft__generate_max_len",
+        "int",
+        "Tokens to generate per block",
+    ),
+    (
+        ("--ebft.n-samples-per-prompt",),
+        "ebft__n_samples_per_prompt",
+        "int",
+        "Independent rollouts per document",
+    ),
+    (
+        ("--ebft.temperature",),
+        "ebft__temperature",
+        "float",
+        "Sampling temperature for strided generation",
+    ),
+    (
+        ("--ebft.top-p",),
+        "ebft__top_p",
+        "float",
+        "Top-p nucleus sampling threshold",
+    ),
+    (
+        ("--ebft.rl-coef",),
+        "ebft__rl_coef",
+        "float",
+        "RL policy gradient loss coefficient",
+    ),
+    (
+        ("--ebft.advantage-estimator",),
+        "ebft__advantage_estimator",
+        None,
+        "Advantage estimator: 'rloo', 'group_norm', 'reinforce'",
+    ),
+    (
+        ("--ebft.min-completion-prefix",),
+        "ebft__min_completion_prefix",
+        "int",
+        "Minimum tokens into completion before placing anchors. Skips anchors too close to the prompt boundary where features are dominated by prompt context.",
+    ),
+    (
+        ("--qat.activation-dtype",),
+        "qat__activation_dtype",
+        None,
+        "Fake quantization layout to use for activation quantization.",
+    ),
+    (
+        ("--qat.weight-dtype",),
+        "qat__weight_dtype",
+        None,
+        "Fake quantization layout to use for weight quantization.",
+    ),
+    (
+        ("--qat.quantize-embedding/--no-qat.quantize-embedding",),
+        "qat__quantize_embedding",
+        None,
+        "Quantize embedding",
+    ),
+    (
+        ("--qat.group-size",),
+        "qat__group_size",
+        "int",
+        "The number of elements in each group for per-group fake quantization",
+    ),
+    (
+        ("--qat.fake-quant-after-n-steps",),
+        "qat__fake_quant_after_n_steps",
+        "int",
+        "The number of steps to apply fake quantization after",
+    ),
+    (
+        ("--quantization.weight-dtype",),
+        "quantization__weight_dtype",
+        None,
+        "Fake quantization layout to use for weight quantization.",
+    ),
+    (
+        ("--quantization.activation-dtype",),
+        "quantization__activation_dtype",
+        None,
+        "Fake quantization layout to use for activation quantization.",
+    ),
+    (
+        ("--quantization.quantize-embedding/--no-quantization.quantize-embedding",),
+        "quantization__quantize_embedding",
+        None,
+        "Whether to quantize the embedding layer.",
+    ),
+    (
+        ("--quantization.group-size",),
+        "quantization__group_size",
+        "int",
+        "The number of elements in each group for per-group fake quantization",
+    ),
+    (
+        ("--reward-model/--no-reward-model",),
+        None,
+        None,
+        "Reward modelling: `True` or `False`",
+    ),
+    (
+        ("--dynamic-checkpoint.enabled/--no-dynamic-checkpoint.enabled",),
+        "dynamic_checkpoint__enabled",
+        None,
+        "Enable dynamic checkpoint triggering during training. Create a file 'axolotl_checkpoint.save' in the configured `output_dir` to trigger. ",
+    ),
+    (
+        ("--dynamic-checkpoint.check-interval",),
+        "dynamic_checkpoint__check_interval",
+        "int",
+        "Check for trigger file every N steps (reduces I/O overhead). Default: 100",
+    ),
+    (
+        ("--dynamic-checkpoint.trigger-file-path",),
+        "dynamic_checkpoint__trigger_file_path",
+        "str",
+        "Custom trigger filename (optional). If not specified, defaults to 'axolotl_checkpoint.save'. Specify a filename (not a full path) to override the default.",
+    ),
+    (
+        ("--process-reward-model/--no-process-reward-model",),
+        None,
+        None,
+        "Process reward modelling: `True` or `False`",
+    ),
+    (
+        ("--center-rewards-coefficient",),
+        None,
+        None,
+        "Coefficient to incentivize the reward model to output mean-zero rewards (proposed by https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`.",
+    ),
+    (
+        ("--num-labels",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dpo-use-weighting/--no-dpo-use-weighting",),
+        None,
+        None,
+        "Whether to perform weighting in DPO trainer",
+    ),
+    (
+        ("--dpo-label-smoothing",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--precompute-ref-log-probs/--no-precompute-ref-log-probs",),
+        None,
+        None,
+        "Precompute reference model log probabilities for DPO",
+    ),
+    (
+        ("--dpo-use-liger-kernel/--no-dpo-use-liger-kernel",),
+        None,
+        None,
+        "Whether to use Liger kernel for DPO loss.",
+    ),
+    (
+        ("--dpo-padding-free/--no-dpo-padding-free",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dpo-loss-type",),
+        None,
+        None,
+        "List of DPO losses to use.",
+    ),
+    (
+        ("--dpo-loss-weights",),
+        None,
+        None,
+        "Weights for each DPO loss.",
+    ),
+    (
+        ("--datasets",),
+        None,
+        None,
+        "A list of one or more datasets to finetune the model with",
+    ),
+    (
+        ("--test-datasets",),
+        None,
+        None,
+        "A list of one or more datasets to eval the model with. You can use either test_datasets, or val_set_size, but not both.",
+    ),
+    (
+        ("--shuffle-merged-datasets/--no-shuffle-merged-datasets",),
+        None,
+        None,
+        "If false, the datasets will not be shuffled and will keep their original order in `datasets`. The same applies to the `test_datasets` option and the `pretraining_dataset` option. Default is true.",
+    ),
+    (
+        ("--shuffle-before-merging-datasets/--no-shuffle-before-merging-datasets",),
+        None,
+        None,
+        "If true, each dataset in `datasets` will be shuffled before merging. This allows curriculum learning strategies to be applied at the dataset level. Default is false.",
+    ),
+    (
+        ("--dataset-prepared-path",),
+        None,
+        None,
+        "Axolotl attempts to save the dataset as an arrow after packing the data together so subsequent training attempts load faster, relative path",
+    ),
+    (
+        ("--dataset-shard-num",),
+        None,
+        None,
+        "Num shards for whole dataset",
+    ),
+    (
+        ("--dataset-shard-idx",),
+        None,
+        None,
+        "Index of shard to use for whole dataset",
+    ),
+    (
+        ("--skip-prepare-dataset/--no-skip-prepare-dataset",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--num-dataset-shards-to-save",),
+        None,
+        None,
+        "Number of shards to save the prepared dataset",
+    ),
+    (
+        ("--pretraining-dataset",),
+        None,
+        None,
+        "Set to HF dataset for type: 'completion' for streaming instead of pre-tokenize",
+    ),
+    (
+        ("--dataset-processes",),
+        None,
+        None,
+        "The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()` if not set.\nFor Runpod VMs, it will default to number of vCPUs via RUNPOD_CPU_COUNT.",
+    ),
+    (
+        ("--dataset-num-proc",),
+        None,
+        None,
+        "The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()` if not set.\nFor Runpod VMs, it will default to number of vCPUs via RUNPOD_CPU_COUNT.",
+    ),
+    (
+        ("--dataset-exact-deduplication/--no-dataset-exact-deduplication",),
+        None,
+        None,
+        "Deduplicates datasets and test_datasets with identical entries",
+    ),
+    (
+        ("--dataset-keep-in-memory/--no-dataset-keep-in-memory",),
+        None,
+        None,
+        "Keep dataset in memory while preprocessing. Only needed if cached dataset is taking too much storage",
+    ),
+    (
+        ("--dataloader-pin-memory/--no-dataloader-pin-memory",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dataloader-num-workers",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dataloader-prefetch-factor",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--dataloader-drop-last/--no-dataloader-drop-last",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--accelerator-config",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--remove-unused-columns/--no-remove-unused-columns",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--push-dataset-to-hub",),
+        None,
+        None,
+        "Push prepared dataset to hub - repo_org/repo_name",
+    ),
+    (
+        ("--hf-use-auth-token/--no-hf-use-auth-token",),
+        None,
+        None,
+        "Whether to use hf `use_auth_token` for loading datasets. Useful for fetching private datasets. Required to be true when used in combination with `push_dataset_to_hub`",
+    ),
+    (
+        ("--device",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--device-map",),
+        None,
+        None,
+        "Passed through to transformers when loading the model when launched without accelerate. Use `sequential` when training w/ model parallelism to limit memory",
+    ),
+    (
+        ("--world-size",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--local-rank",),
+        None,
+        None,
+        "Don't mess with this, it's here for accelerate and torchrun",
+    ),
+    (
+        ("--ddp/--no-ddp",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--seed",),
+        None,
+        None,
+        "Seed for reproducibility",
+    ),
+    (
+        ("--ddp-timeout",),
+        None,
+        None,
+        "Advanced DDP Arguments - timeout",
+    ),
+    (
+        ("--ddp-bucket-cap-mb",),
+        None,
+        None,
+        "Advanced DDP Arguments - bucket cap in MB",
+    ),
+    (
+        ("--ddp-broadcast-buffers/--no-ddp-broadcast-buffers",),
+        None,
+        None,
+        "Advanced DDP Arguments - broadcast buffers",
+    ),
+    (
+        ("--ddp-find-unused-parameters/--no-ddp-find-unused-parameters",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--do-causal-lm-eval/--no-do-causal-lm-eval",),
+        None,
+        None,
+        "Whether to run causal language model evaluation for metrics in `eval_causal_lm_metrics`",
+    ),
+    (
+        ("--eval-causal-lm-metrics",),
+        None,
+        None,
+        "HF evaluate metrics used during evaluation. Default is ['sacrebleu', 'comet', 'ter', 'chrf', 'perplexity']",
+    ),
+    (
+        ("--do-bench-eval/--no-do-bench-eval",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--bench-dataset",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--bench-split",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--metric-for-best-model",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--greater-is-better/--no-greater-is-better",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--loss-watchdog-threshold",),
+        None,
+        None,
+        "High loss value, indicating the learning has broken down (a good estimate is ~2 times the loss at the start of training)",
+    ),
+    (
+        ("--loss-watchdog-patience",),
+        None,
+        None,
+        "Number of high-loss steps in a row before the trainer aborts (default: 3)",
+    ),
+    (
+        ("--gc-steps",),
+        None,
+        None,
+        "Deprecated. Run garbage collection every `gc_steps` steps. Use `torch_empty_cache_steps` and `gc_collect_steps` instead.",
+    ),
+    (
+        ("--torch-empty-cache-steps",),
+        None,
+        None,
+        "Steps between native HF Trainer `torch.cuda.empty_cache()` calls.",
+    ),
+    (
+        ("--gc-collect-steps",),
+        None,
+        None,
+        "Steps between Python `gc.collect()` calls. -1 runs on epoch end and before evals only; None disables.",
+    ),
+    (
+        ("--bf16",),
+        None,
+        None,
+        "Use CUDA bf16. bool or 'full' for `bf16_full_eval`, or 'auto' for automatic detection. require >=ampere",
+    ),
+    (
+        ("--fp16/--no-fp16",),
+        None,
+        None,
+        "Use CUDA fp16",
+    ),
+    (
+        ("--fp8/--no-fp8",),
+        None,
+        None,
+        "Enable FP8 mixed precision training using TorchAO. Best used in combination with torch.compile.",
+    ),
+    (
+        ("--fp8-enable-fsdp-float8-all-gather/--no-fp8-enable-fsdp-float8-all-gather",),
+        None,
+        None,
+        "Enable FSDP float8 all-gather optimization for FP8 training. Can improve training speed by 10-15% when FSDP is enabled.",
+    ),
+    (
+        ("--bfloat16/--no-bfloat16",),
+        None,
+        None,
+        "No AMP (automatic mixed precision) - require >=ampere",
+    ),
+    (
+        ("--float16/--no-float16",),
+        None,
+        None,
+        "No AMP (automatic mixed precision)",
+    ),
+    (
+        ("--tf32",),
+        None,
+        None,
+        "bool to use CUDA tf32 or 'auto' for automatic detection - require >=ampere",
+    ),
+    (
+        ("--float32/--no-float32",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--gradient-checkpointing",),
+        None,
+        None,
+        "Whether to use gradient checkpointing. Available options are: true, false, 'offload', 'offload_disk'. https://huggingface.co/docs/transformers/v4.18.0/en/performance#gradient-checkpointing",
+    ),
+    (
+        ("--gradient-checkpointing-kwargs",),
+        None,
+        None,
+        "Additional kwargs to pass to the trainer for gradient checkpointing",
+    ),
+    (
+        ("--selective-checkpointing.save",),
+        "selective_checkpointing__save",
+        None,
+        "Ops to save during forward instead of recomputing in backward. 'attention' matches SDPA and flash-attention forward ops; other entries are substring-matched against qualified torch op names (e.g. 'aten::mm').",
+    ),
+    (
+        (
+            "--selective-checkpointing.save-sliding-window/--no-selective-checkpointing.save-sliding-window",
+        ),
+        "selective_checkpointing__save_sliding_window",
+        None,
+        "In hybrid full/sliding-window attention models, also save sliding-window attention calls. Default false: SWA is cheap to recompute, so only full-attention calls are saved.",
+    ),
+    (
+        ("--selective-checkpointing.recompute-layer-types",),
+        "selective_checkpointing__recompute_layer_types",
+        None,
+        "Layer types (config.layer_types values) whose attention is recomputed instead of saved. Defaults to ['sliding_attention', 'chunked_attention']. Ignored when save_sliding_window is true. Linear-attention layers never dispatch a matchable attention op, so they need no entry.",
+    ),
+    (
+        ("--selective-checkpointing.offload/--no-selective-checkpointing.offload",),
+        "selective_checkpointing__offload",
+        None,
+        "Offload saved tensors to pinned CPU memory (side-stream copies with backward prefetch) instead of keeping them on GPU.",
+    ),
+    (
+        ("--activation-offloading",),
+        None,
+        None,
+        "Whether to offload activations. Options: true/false, 'legacy', 'disk' (TRL offloader), or 'hidden_states' (checkpoint input offload; non-reentrant by default, ALST-style when gradient_checkpointing_kwargs.use_reentrant is true).",
+    ),
+    (
+        ("--layer-offloading/--no-layer-offloading",),
+        None,
+        None,
+        "Offload model layer parameters to CPU during forward, prefetch back during backward.",
+    ),
+    (
+        ("--freeze-mm-modules/--no-freeze-mm-modules",),
+        None,
+        None,
+        "Freeze multimodal encoder parameters (vision, audio, etc.) for text-only training of multimodal models. When True, parameters belonging to vision towers, audio towers, multimodal projectors, and similar non-language modules are frozen (requires_grad=False). This allows DDP training without ddp_find_unused_parameters=True.",
+    ),
+    (
+        ("--unfrozen-parameters",),
+        None,
+        None,
+        "List of regex patterns for parameter names to keep unfrozen. All other parameters will be frozen via requires_grad=False. Note: range-based patterns (e.g. embed_tokens.weight$[:32000]) use gradient zeroing rather than a true freeze, so weight decay will still apply to the frozen portion and optimizer states are allocated for the full parameter.",
+    ),
+    (
+        ("--sequence-len",),
+        None,
+        None,
+        "The maximum length of an input to train with, this should typically be less than 2048 as most models have a token/context limit of 2048",
+    ),
+    (
+        ("--excess-length-strategy",),
+        None,
+        None,
+        "What to do when a tokenized row exceeds sequence_len. 'drop' removes the row; 'truncate' slices tensors to sequence_len; 'raise' raises a ValueError. Defaults to 'drop' for backward compatibility.",
+    ),
+    (
+        ("--eval-sequence-len",),
+        None,
+        None,
+        "The maximum length of an input for evaluation. If not specified, defaults to sequence_len",
+    ),
+    (
+        ("--min-sample-len",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--max-prompt-len",),
+        None,
+        None,
+        "maximum prompt length for RL training",
+    ),
+    (
+        ("--sample-packing/--no-sample-packing",),
+        None,
+        None,
+        "Use efficient multi-packing with block diagonal attention and per sequence position_ids. Recommend set to 'true'",
+    ),
+    (
+        ("--sample-packing-group-size",),
+        None,
+        None,
+        "The number of samples packed at a time. Increasing the following values helps with packing, but usually only slightly (<%1.)",
+    ),
+    (
+        ("--sample-packing-bin-size",),
+        None,
+        None,
+        "The number of samples which can be packed into one sequence. Increase if using a large sequence_len with many short samples.",
+    ),
+    (
+        ("--sample-packing-sequentially/--no-sample-packing-sequentially",),
+        None,
+        None,
+        "Whether to pack samples sequentially",
+    ),
+    (
+        ("--sample-packing-mp-start-method",),
+        None,
+        None,
+        "The multiprocessing start method to use for packing. Should be 'fork', 'spawn' or 'forkserver'",
+    ),
+    (
+        ("--eval-sample-packing/--no-eval-sample-packing",),
+        None,
+        None,
+        "Set to 'false' if getting errors during eval with sample_packing on",
+    ),
+    (
+        ("--pad-to-sequence-len/--no-pad-to-sequence-len",),
+        None,
+        None,
+        "Pad inputs so each step uses constant sized buffers. This will reduce memory fragmentation and may prevent OOMs, by re-using memory more efficiently. Defaults to True if `sample_packing` enabled",
+    ),
+    (
+        ("--pad-to-multiple-of",),
+        None,
+        None,
+        "Pad each batch to a multiple of this value.",
+    ),
+    (
+        ("--curriculum-sampling/--no-curriculum-sampling",),
+        None,
+        None,
+        "Whether to use sequential sampling for curriculum learning",
+    ),
+    (
+        ("--multipack-real-batches/--no-multipack-real-batches",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--batch-flattening",),
+        None,
+        None,
+        "Use batch flattening for speedups when not using sample_packing",
+    ),
+    (
+        ("--use-pose/--no-use-pose",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--pose-split-on-token-ids",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--pose-max-context-len",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--pose-num-chunks",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--pretrain-multipack-buffer-size",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--pretrain-multipack-attn/--no-pretrain-multipack-attn",),
+        None,
+        None,
+        "whether to prevent cross attention for packed sequences during pretraining",
+    ),
+    (
+        ("--pretraining-sample-concatenation/--no-pretraining-sample-concatenation",),
+        None,
+        None,
+        "whether to concatenate samples during pretraining",
+    ),
+    (
+        ("--streaming/--no-streaming",),
+        None,
+        None,
+        "Use streaming mode for loading datasets",
+    ),
+    (
+        ("--streaming-multipack-buffer-size",),
+        None,
+        None,
+        "Buffer size for multipack streaming datasets",
+    ),
+    (
+        ("--xformers-attention/--no-xformers-attention",),
+        None,
+        None,
+        "[DEPRECATED] Use `attn_implementation: xformers`. https://github.com/facebookresearch/xformers",
+    ),
+    (
+        ("--sdp-attention/--no-sdp-attention",),
+        None,
+        None,
+        "[DEPRECATED] Use `attn_implementation: sdpa`.",
+    ),
+    (
+        ("--flex-attention/--no-flex-attention",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--flex-attn-compile-kwargs",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--flash-attention/--no-flash-attention",),
+        None,
+        None,
+        "[DEPRECATED] Use `attn_implementation: flash_attention_2`. https://github.com/Dao-AILab/flash-attention",
+    ),
+    (
+        ("--flash-attn-cross-entropy/--no-flash-attn-cross-entropy",),
+        None,
+        None,
+        "Whether to use flash-attention cross entropy implementation - advanced use only",
+    ),
+    (
+        ("--flash-attn-fuse-mlp/--no-flash-attn-fuse-mlp",),
+        None,
+        None,
+        "Whether to fuse part of the MLP into a single operation",
+    ),
+    (
+        ("--flash-optimum/--no-flash-optimum",),
+        None,
+        None,
+        "Whether to use bettertransformers",
+    ),
+    (
+        ("--sage-attention/--no-sage-attention",),
+        None,
+        None,
+        "[DEPRECATED] Use `attn_implementation: sage`. https://github.com/thu-ml/SageAttention",
+    ),
+    (
+        ("--eager-attention/--no-eager-attention",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--attn-implementation",),
+        None,
+        None,
+        "Attention backend. Canonical values: eager, sdpa, flash_attention_2, flash_attention_3, flex_attention, xformers, sage, fp8. Hub-kernel paths (e.g. kernels-community/flash-attn3) are also accepted and passed through to transformers.",
+    ),
+    (
+        ("--gemma4-hybrid-attn-impl/--no-gemma4-hybrid-attn-impl",),
+        None,
+        None,
+        "Use hybrid attention for Gemma 4: flash_attention_2 for sliding window layers and sdpa for global (full_attention) layers. Global layers have head_dim=512 which exceeds flash attention's supported size.",
+    ),
+    (
+        ("--large-head-attention",),
+        None,
+        None,
+        "Generic capability for attention layers with head_dim > 256 (which flash/cuDNN SDPA can't serve): 'sdpa' (stock SDPA, default) | 'auto' (Triton flash kernel on packed rows, SDPA otherwise) | 'triton_flash' (prefer the Triton kernel). Applies to any SDPA model; Gemma-4's head_dim=512 global layers use it automatically.",
+    ),
+    (
+        ("--flash-attn-d512/--no-flash-attn-d512",),
+        None,
+        None,
+        "Deprecated: use `large_head_attention: auto`. Routes head_dim>256 attention through the Triton flash kernel.",
+    ),
+    (
+        ("--sdpa-varlen/--no-sdpa-varlen",),
+        None,
+        None,
+        "With sample packing + attn_implementation=sdpa, route packed rows through torch.nn.attention.varlen.varlen_attn (cu_seqlens) instead of an explicit 4D block-diagonal mask. Skips cross-document blocks (faster + lower memory) with no flash_attn dependency. Left unset (null) it auto-enables when supported (torch >= 2.10, head_dim <= 256, no sliding window); set true/false to force. When it can't be used, packed rows still isolate documents via the block-diagonal mask. Sliding-window attention needs torch >= 2.11 (varlen_attn window_size).",
+    ),
+    (
+        ("--fused-attn-kernel/--no-fused-attn-kernel",),
+        None,
+        None,
+        "Replace ``q_norm + apply_rotary_pos_emb`` (and the matching k path) with a single fused RMSNorm+RoPE Triton kernel launch. Currently implemented for Qwen3, Qwen3-MoE, Qwen3.5, and Qwen3.5-MoE full-attention layers; Gemma 4 always uses the fused path. Disabled (None/False) falls back to the eager transformers implementation. Compile-safe via torch.library.triton_op \u2014 traces under torch.compile(fullgraph=True). Per-step wins are arch-dependent: ~+7-12% across sm_86 and sm_120. Combining with torch_compile=true is a clear win on sm_120 (+9% extra) but currently regresses on sm_86 due to Inductor autotune biases \u2014 flip them on independently and benchmark.",
+    ),
+    (
+        ("--experts-implementation",),
+        None,
+        None,
+        "Which experts implementation to use for MoE models,",
+    ),
+    (
+        ("--quantize-moe-experts/--no-quantize-moe-experts",),
+        None,
+        None,
+        "Quantize MoE expert weights on load to reduce VRAM. Requires adapter (lora/qlora) with load_in_4bit or load_in_8bit. Requires CUDA (not compatible with ROCm or other backends). Note: total parameter count may be reported incorrectly when enabled (trainable param count is correct).",
+    ),
+    (
+        ("--scaling-softmax/--no-scaling-softmax",),
+        None,
+        None,
+        "Whether to use Scaled Softmax (SSMax) attention. Ref: https://arxiv.org/abs/2501.19399",
+    ),
+    (
+        ("--scaling-softmax-factor",),
+        None,
+        None,
+        "Scaling factor for SSMax attention. Default is 0.43",
+    ),
+    (
+        ("--scaling-softmax-bias",),
+        None,
+        None,
+        "Bias for SSMax attention. Default is 0.0. Note: The paper recommends bias=0 for better length generalization.",
+    ),
+    (
+        ("--lora-mlp-kernel/--no-lora-mlp-kernel",),
+        None,
+        None,
+        "Apply custom LoRA autograd functions and activation function Triton kernels for speed and memory savings. See: https://docs.axolotl.ai/docs/lora_optims.html",
+    ),
+    (
+        ("--lora-qkv-kernel/--no-lora-qkv-kernel",),
+        None,
+        None,
+        "Apply custom LoRA autograd functions and activation function Triton kernels for speed and memory savings. See: https://docs.axolotl.ai/docs/lora_optims.html",
+    ),
+    (
+        ("--lora-o-kernel/--no-lora-o-kernel",),
+        None,
+        None,
+        "Apply custom LoRA autograd functions and activation function Triton kernels for speed and memory savings. See: https://docs.axolotl.ai/docs/lora_optims.html",
+    ),
+    (
+        ("--lora-embedding-kernel/--no-lora-embedding-kernel",),
+        None,
+        None,
+        "Apply custom LoRA autograd function for embedding layers. See: https://docs.axolotl.ai/docs/lora_optims.html",
+    ),
+    (
+        ("--chunked-cross-entropy/--no-chunked-cross-entropy",),
+        None,
+        None,
+        "Whether to use chunked cross entropy loss for memory efficiency",
+    ),
+    (
+        ("--chunked-cross-entropy-num-chunks",),
+        None,
+        None,
+        "Number of chunks to use for chunked cross entropy loss",
+    ),
+    (
+        ("--use-eaft/--no-use-eaft",),
+        None,
+        None,
+        "Enable Entropy-Aware Focal Training loss (EAFT)",
+    ),
+    (
+        ("--eaft-alpha",),
+        None,
+        None,
+        "Exponent for entropy weighting in EAFT (default: 1.0)",
+    ),
+    (
+        ("--eaft-k",),
+        None,
+        None,
+        "Number of top logits for entropy approximation (default: 20)",
+    ),
+    (
+        ("--tiled-mlp/--no-tiled-mlp",),
+        None,
+        None,
+        "Whether to use ALST tiled mlp for memory efficient long context",
+    ),
+    (
+        ("--tiled-mlp-num-shards",),
+        None,
+        None,
+        "Number of shards to use for ALST tiled mlp. If unset, it will be set based on seqlen/hidden_size",
+    ),
+    (
+        ("--tiled-mlp-use-original-mlp/--no-tiled-mlp-use-original-mlp",),
+        None,
+        None,
+        "Whether to use original mlp for ALST tiled mlp. Otherwise uses a generic MLP based on llama.",
+    ),
+    (
+        ("--llama4-linearized-experts/--no-llama4-linearized-experts",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--deepspeed",),
+        None,
+        None,
+        "Deepspeed config path. e.g., deepspeed_configs/zero3.json",
+    ),
+    (
+        ("--deepcompile/--no-deepcompile",),
+        None,
+        None,
+        "Whether to use deepcompile for faster training with deepspeed",
+    ),
+    (
+        ("--fsdp",),
+        None,
+        None,
+        "FSDP configuration",
+    ),
+    (
+        ("--fsdp-config.fsdp-version",),
+        "fsdp_config__fsdp_version",
+        "int",
+        "FSDP version",
+    ),
+    (
+        (
+            "--fsdp-config.activation-checkpointing/--no-fsdp-config.activation-checkpointing",
+        ),
+        "fsdp_config__activation_checkpointing",
+        None,
+        "Enable activation checkpointing to reduce memory usage during forward passes",
+    ),
+    (
+        ("--fsdp-config.offload-params/--no-fsdp-config.offload-params",),
+        "fsdp_config__offload_params",
+        None,
+        "Offload parameters to CPU to reduce GPU memory usage",
+    ),
+    (
+        ("--fsdp-config.sync-module-states/--no-fsdp-config.sync-module-states",),
+        "fsdp_config__sync_module_states",
+        None,
+        "Synchronize module states across all processes",
+    ),
+    (
+        (
+            "--fsdp-config.cpu-ram-efficient-loading/--no-fsdp-config.cpu-ram-efficient-loading",
+        ),
+        "fsdp_config__cpu_ram_efficient_loading",
+        None,
+        "Enable CPU RAM efficient loading to reduce memory usage during model loading",
+    ),
+    (
+        (
+            "--fsdp-config.cpu-offload-pin-memory/--no-fsdp-config.cpu-offload-pin-memory",
+        ),
+        "fsdp_config__cpu_offload_pin_memory",
+        None,
+        "Disabling this enables swap memory usage for resource-constrained setups when offload_params is enabled.",
+    ),
+    (
+        ("--fsdp-config.use-orig-params/--no-fsdp-config.use-orig-params",),
+        "fsdp_config__use_orig_params",
+        None,
+        "Use original parameters instead of flattened parameters",
+    ),
+    (
+        ("--fsdp-config.state-dict-type",),
+        "fsdp_config__state_dict_type",
+        None,
+        "Type of state dict to use for saving/loading checkpoints",
+    ),
+    (
+        ("--fsdp-config.final-state-dict-type",),
+        "fsdp_config__final_state_dict_type",
+        None,
+        "Final state dict type to use after training completion",
+    ),
+    (
+        ("--fsdp-config.auto-wrap-policy",),
+        "fsdp_config__auto_wrap_policy",
+        None,
+        "Policy for automatically wrapping modules with FSDP",
+    ),
+    (
+        ("--fsdp-config.transformer-layer-cls-to-wrap",),
+        "fsdp_config__transformer_layer_cls_to_wrap",
+        "str",
+        "Class name of transformer layers to wrap (e.g., 'LlamaDecoderLayer')",
+    ),
+    (
+        ("--fsdp-config.min-num-params",),
+        "fsdp_config__min_num_params",
+        "int",
+        "Minimum parameter count a module must have to be wrapped under the SIZE_BASED_WRAP policy",
+    ),
+    (
+        ("--fsdp-config.reshard-after-forward/--no-fsdp-config.reshard-after-forward",),
+        "fsdp_config__reshard_after_forward",
+        None,
+        "Reshard parameters after forward pass to save memory",
+    ),
+    (
+        ("--fsdp-config.mixed-precision-policy",),
+        "fsdp_config__mixed_precision_policy",
+        "str",
+        "Mixed precision policy for FSDP (e.g., 'fp16', 'bf16')",
+    ),
+    (
+        ("--fsdp-version",),
+        None,
+        None,
+        "FSDP version",
+    ),
+    (
+        ("--fp32-norms/--no-fp32-norms",),
+        None,
+        None,
+        "Keep norm modules (RMSNorm/LayerNorm) in fp32 by sharding them under their own FSDP2 MixedPrecisionPolicy. Requires fsdp_version: 2.",
+    ),
+    (
+        ("--fp32-norm-classes",),
+        None,
+        None,
+        "Class-name patterns to match for fp32 norm sharding. Patterns without a '.' match against type(module).__name__ as a suffix. Patterns containing a '.' match the fully qualified class path exactly. Defaults to ['RMSNorm', 'LayerNorm'] when fp32_norms is true and this is unset.",
+    ),
+    (
+        ("--fsdp-final-state-dict-type",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--val-set-size",),
+        None,
+        None,
+        "How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc. 0 for no eval.",
+    ),
+    (
+        ("--dp-shard-size",),
+        None,
+        None,
+        "Number of devices to shard across. If not set, will use all available devices.",
+    ),
+    (
+        ("--dp-replicate-size",),
+        None,
+        None,
+        "Number of devices to replicate across.",
+    ),
+    (
+        ("--sequence-parallel-degree",),
+        None,
+        None,
+        "Deprecated: use `context_parallel_size` instead",
+    ),
+    (
+        ("--context-parallel-size",),
+        None,
+        None,
+        "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://docs.axolotl.ai/docs/sequence_parallelism.html for more details.",
+    ),
+    (
+        ("--heads-k-stride",),
+        None,
+        None,
+        "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.",
+    ),
+    (
+        ("--ring-attn-func",),
+        None,
+        None,
+        "One of 'varlen_llama3', 'batch_ring', 'batch_zigzag', 'batch_stripe'. Defaults to 'varlen_llama3' in the sample packing case, and 'batch_ring' in the non-sample packing case.",
+    ),
+    (
+        ("--tensor-parallel-size",),
+        None,
+        None,
+        "Number of tensor parallel processes in TP group. Only supported with DeepSpeed AutoTP.",
+    ),
+    (
+        ("--special-tokens.bos-token",),
+        "special_tokens__bos_token",
+        "str",
+        None,
+    ),
+    (
+        ("--special-tokens.eos-token",),
+        "special_tokens__eos_token",
+        "str",
+        None,
+    ),
+    (
+        ("--special-tokens.pad-token",),
+        "special_tokens__pad_token",
+        "str",
+        None,
+    ),
+    (
+        ("--special-tokens.unk-token",),
+        "special_tokens__unk_token",
+        "str",
+        None,
+    ),
+    (
+        ("--special-tokens.additional-special-tokens",),
+        "special_tokens__additional_special_tokens",
+        None,
+        None,
+    ),
+    (
+        ("--tokens",),
+        None,
+        None,
+        "Add extra tokens to the tokenizer",
+    ),
+    (
+        ("--added-tokens-overrides",),
+        None,
+        None,
+        "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.",
+    ),
+    (
+        ("--torch-compile",),
+        None,
+        None,
+        "Whether to use torch.compile and which backend to use.",
+    ),
+    (
+        ("--torch-compile-backend",),
+        None,
+        None,
+        "Backend to use for torch.compile",
+    ),
+    (
+        ("--torch-compile-mode",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--max-steps",),
+        None,
+        None,
+        "Maximum number of iterations to train for. It precedes num_epochs which means that if both are set, num_epochs will not be guaranteed. e.g., when 1 epoch is 1000 steps => `num_epochs: 2` and `max_steps: 100` will train for 100 steps",
+    ),
+    (
+        ("--warmup-steps",),
+        None,
+        None,
+        "Number of warmup steps. Cannot use with warmup_ratio",
+    ),
+    (
+        ("--warmup-ratio",),
+        None,
+        None,
+        "Warmup ratio. Cannot use with warmup_steps",
+    ),
+    (
+        ("--eval-steps",),
+        None,
+        None,
+        "Leave empty to eval at each epoch, integer for every N steps. float for fraction of total steps",
+    ),
+    (
+        ("--evals-per-epoch",),
+        None,
+        None,
+        "Number of times per epoch to run evals, mutually exclusive with eval_steps",
+    ),
+    (
+        ("--eval-strategy",),
+        None,
+        None,
+        "Set to `no` to skip evaluation, `epoch` at end of each epoch, leave empty to infer from `eval_steps`",
+    ),
+    (
+        ("--save-steps",),
+        None,
+        None,
+        "Leave empty to save at each epoch, integer for every N steps. float for fraction of total steps",
+    ),
+    (
+        ("--saves-per-epoch",),
+        None,
+        None,
+        "Number of times per epoch to save a checkpoint, mutually exclusive with save_steps",
+    ),
+    (
+        ("--save-strategy",),
+        None,
+        None,
+        "Set to `no` to skip checkpoint saves, `epoch` at end of each epoch, `best` when better result is achieved, leave empty to infer from `save_steps`",
+    ),
+    (
+        ("--save-total-limit",),
+        None,
+        None,
+        "Checkpoints saved at a time",
+    ),
+    (
+        ("--save-first-step/--no-save-first-step",),
+        None,
+        None,
+        "Whether to checkpoint a model after the first step of training. Defaults to False.",
+    ),
+    (
+        ("--logging-steps",),
+        None,
+        None,
+        "Logging frequency",
+    ),
+    (
+        ("--early-stopping-patience",),
+        None,
+        None,
+        "Stop training after this many evaluation losses have increased in a row. https://huggingface.co/transformers/v4.2.2/_modules/transformers/trainer_callback.html#EarlyStoppingCallback",
+    ),
+    (
+        ("--load-best-model-at-end/--no-load-best-model-at-end",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--save-only-model/--no-save-only-model",),
+        None,
+        None,
+        "Save only the model weights, skipping the optimizer. Using this means you can't resume from checkpoints.",
+    ),
+    (
+        ("--use-tensorboard/--no-use-tensorboard",),
+        None,
+        None,
+        "Use tensorboard for logging",
+    ),
+    (
+        ("--profiler-steps",),
+        None,
+        None,
+        "Enable the pytorch profiler to capture the first N steps of training to the output_dir. see https://pytorch.org/blog/understanding-gpu-memory-1/ for more information. Snapshots can be visualized @ https://pytorch.org/memory_viz",
+    ),
+    (
+        ("--profiler-steps-start",),
+        None,
+        None,
+        "Which step to start the profiler at. Useful for only capturing a few steps mid-run.",
+    ),
+    (
+        ("--include-tokens-per-second/--no-include-tokens-per-second",),
+        None,
+        None,
+        "bool of whether to report tokens per second at the end of training. This is not supported with pre-training datasets.",
+    ),
+    (
+        ("--include-tkps/--no-include-tkps",),
+        None,
+        None,
+        "bool of whether to report tokens per second per-gpu during training by measuring throughput of non-padding tokens.",
+    ),
+    (
+        ("--neftune-noise-alpha",),
+        None,
+        None,
+        "NEFT https://arxiv.org/abs/2310.05914, set this to a number (paper default is 5) to add noise to embeddings. Currently only supported on Llama and Mistral",
+    ),
+    (
+        ("--orpo-alpha",),
+        None,
+        None,
+        "Parameter controlling the relative ratio loss weight in the ORPO loss. Passed to `beta` in `ORPOConfig` due to trl mapping.",
+    ),
+    (
+        ("--simpo-gamma",),
+        None,
+        None,
+        "Target reward margin for the SimPO loss",
+    ),
+    (
+        ("--cpo-alpha",),
+        None,
+        None,
+        "Weight of the BC regularizer",
+    ),
+    (
+        ("--kto-desirable-weight",),
+        None,
+        None,
+        "Factor for desirable loss term in KTO loss",
+    ),
+    (
+        ("--kto-undesirable-weight",),
+        None,
+        None,
+        "Factor for undesirable loss term in KTO loss",
+    ),
+    (
+        ("--rl-beta",),
+        None,
+        None,
+        "The beta parameter for the RL training",
+    ),
+    (
+        ("--max-memory",),
+        None,
+        None,
+        "Defines the max memory usage per gpu on the system. Passed through to transformers when loading the model.",
+    ),
+    (
+        ("--gpu-memory-limit",),
+        None,
+        None,
+        "Limit the memory for all available GPUs to this amount (if an integer, expressed in gigabytes); default: unset",
+    ),
+    (
+        ("--low-cpu-mem-usage/--no-low-cpu-mem-usage",),
+        None,
+        None,
+        "Whether to use low_cpu_mem_usage",
+    ),
+    (
+        ("--chat-template",),
+        None,
+        None,
+        "The name of the chat template to use for training, following values are supported: tokenizer_default: Uses the chat template that is available in the tokenizer_config.json. If the chat template is not available in the tokenizer, it will raise an error. This is the default value. alpaca/inst/chatml/gemma/cohere/llama3/phi_3/deepseek_v2/jamba: These chat templates are available in the axolotl codebase at src/axolotl/utils/chat_templates.py. tokenizer_default_fallback_*: where * is the name of the chat template to fallback to. E.g. tokenizer_default_fallback_chatml. This is useful when the chat template is not available in the tokenizer. jinja: Uses a custom jinja template for the chat template. The custom jinja template should be provided in the chat_template_jinja field. The selected chat template will be saved to the tokenizer_config.json for easier inferencing",
+    ),
+    (
+        ("--chat-template-jinja",),
+        None,
+        None,
+        "Custom jinja template or path to jinja file for chat template. This will be only used if chat_template is set to `jinja` or `null` (in which case chat_template is automatically set to `jinja`). Default is null.",
+    ),
+    (
+        ("--chat-template-kwargs",),
+        None,
+        None,
+        "Additional kwargs to pass to the chat template. This is useful for customizing the chat template. For example, you can pass `thinking=False` to add a generation prompt to the chat template.",
+    ),
+    (
+        ("--eot-tokens",),
+        None,
+        None,
+        "Custom EOT (End-of-Turn) tokens to mask/unmask during training. These tokens mark the boundaries between conversation turns. For example: ['/INST', '</s>', '[/SYSTEM_PROMPT]']. If not specified, defaults to just the model's eos_token. This is useful for templates that use multiple delimiter tokens.",
+    ),
+    (
+        ("--default-system-message",),
+        None,
+        None,
+        "Changes the default system message. Currently only supports chatml.",
+    ),
+    (
+        ("--fix-untrained-tokens",),
+        None,
+        None,
+        "Token index or indices to adjust embedding weights to the mean of the other tokens. This is useful when the model has untrained embeddings.",
+    ),
+    (
+        ("--is-preprocess/--no-is-preprocess",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--preprocess-iterable/--no-preprocess-iterable",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--total-num-tokens",),
+        None,
+        None,
+        "Total number of tokens - internal use",
+    ),
+    (
+        ("--total-supervised-tokens",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--sample-packing-eff-est",),
+        None,
+        None,
+        "You can set these packing optimizations AFTER starting a training at least once. The trainer will provide recommended values for these values.",
+    ),
+    (
+        ("--axolotl-config-path",),
+        None,
+        None,
+        None,
+    ),
+    (
+        ("--is-falcon-derived-model/--no-is-falcon-derived-model",),
+        None,
+        None,
+        "Internal use only - Used to identify which the model is based on",
+    ),
+    (
+        ("--is-llama-derived-model/--no-is-llama-derived-model",),
+        None,
+        None,
+        "Internal use only - Used to identify which the model is based on",
+    ),
+    (
+        ("--is-mistral-derived-model/--no-is-mistral-derived-model",),
+        None,
+        None,
+        "Internal use only - Used to identify which the model is based on. Please note that if you set this to true, `padding_side` will be set to 'left' by default",
+    ),
+    (
+        ("--is-qwen-derived-model/--no-is-qwen-derived-model",),
+        None,
+        None,
+        "Internal use only - Used to identify which the model is based on",
+    ),
+    (
+        ("--plugins",),
+        None,
+        None,
+        "Add plugins to extend the pipeline. See `src/axolotl/integrations` for the available plugins or doc below for more details. https://docs.axolotl.ai/docs/custom_integrations.html",
+    ),
+    (
+        ("--generate-samples/--no-generate-samples",),
+        None,
+        None,
+        "Enable sample generation during training for monitoring",
+    ),
+    (
+        ("--num-generation-samples",),
+        None,
+        None,
+        "Number of samples to generate at each interval",
+    ),
+    (
+        ("--generation-max-new-tokens",),
+        None,
+        None,
+        "Maximum new tokens to generate per sample",
+    ),
+    (
+        ("--generation-temperature",),
+        None,
+        None,
+        "Temperature for sample generation (0.0 = greedy)",
+    ),
+    (
+        ("--generation-top-p",),
+        None,
+        None,
+        "Nucleus sampling parameter for generation",
+    ),
+    (
+        ("--generation-top-k",),
+        None,
+        None,
+        "Top-k sampling parameter for generation",
+    ),
+    (
+        ("--generation-prompt-ratio",),
+        None,
+        None,
+        "Ratio of input to use as prompt (0.0-1.0)",
+    ),
+    (
+        ("--generation-do-sample/--no-generation-do-sample",),
+        None,
+        None,
+        "Whether to use sampling (vs greedy decoding)",
+    ),
+)
diff --git a/src/axolotl/cli/delinearize_llama4.py b/src/axolotl/cli/delinearize_llama4.py
new file mode 100644
index 0000000000..4f5448a14b
--- /dev/null
+++ b/src/axolotl/cli/delinearize_llama4.py
@@ -0,0 +1,152 @@
+"""
+CLI tool to delinearize quantized/Linearized Llama-4 models.
+"""
+
+import os
+from pathlib import Path
+from typing import Generator, Union
+
+import fire
+import torch
+from accelerate import init_empty_weights
+from transformers import AutoProcessor
+
+
+def iter_convert_patched_to_hf(model_state_dict, num_experts) -> Generator:
+    keys = list(model_state_dict.keys())
+    for key in keys:
+        if ".feed_forward.experts." not in key:
+            yield key, model_state_dict[key]
+        if ".feed_forward.experts.gate_projs" in key:
+            # gate gets fused with up so skip the yield on this and we'll fuse it when asking for the up
+            continue
+        if ".feed_forward.experts.up_projs" in key:
+            if ".feed_forward.experts.up_projs.0." in key:
+                # handle the re-shape and fusing of gate and up, and conversion from linear to parameter
+                prefix = key.split(".up_projs.0.")[0]
+                key = f"{prefix}.gate_up_proj"
+                # grab all the up_projs and gate_projs across all experts
+                gate_stacked = torch.stack(
+                    [
+                        model_state_dict[
+                            f"{prefix}.gate_projs.{expert_idx}.weight"
+                        ].transpose(0, 1)
+                        for expert_idx in range(num_experts)
+                    ]
+                )
+                up_stacked = torch.stack(
+                    [
+                        model_state_dict[
+                            f"{prefix}.up_projs.{expert_idx}.weight"
+                        ].transpose(0, 1)
+                        for expert_idx in range(num_experts)
+                    ]
+                )
+                gate_up_proj = torch.cat((gate_stacked, up_stacked), dim=-1)
+                del gate_stacked, up_stacked
+                yield key, gate_up_proj
+            else:
+                del model_state_dict[key]
+                continue
+        if ".feed_forward.experts.down_projs" in key:
+            if ".feed_forward.experts.down_projs.0." in key:
+                # handle the re-shape and fusing of gate and up, and conversion from linear to parameter
+                prefix = key.split(".down_projs.0.")[0]
+                key = f"{prefix}.down_proj"
+                # grab all the down_projs across all experts
+                down_stacked = torch.stack(
+                    [
+                        model_state_dict[
+                            f"{prefix}.down_projs.{expert_idx}.weight"
+                        ].transpose(0, 1)
+                        for expert_idx in range(num_experts)
+                    ]
+                )
+                yield key, down_stacked
+            else:
+                del model_state_dict[key]
+                continue
+
+
+def do_cli(model: Union[Path, str], output: Union[Path, str]) -> None:
+    """
+    Convert a patched HF format Llama4 model (with separated projections)
+    back to the original HF format (with fused projections).
+
+    Args:
+        model: Path to the patched HF model
+        output: Path to save the converted model
+    """
+    print(f"Loading model from {model}")
+    from axolotl.monkeypatch.models.llama4.modeling import (
+        patch_llama4_linearized_modeling,
+    )
+
+    unpatch_llama4 = patch_llama4_linearized_modeling()
+    from transformers import Llama4ForConditionalGeneration
+
+    model_ = Llama4ForConditionalGeneration.from_pretrained(model, dtype=torch.bfloat16)
+    processor = AutoProcessor.from_pretrained(model)
+    processor.save_pretrained(output)
+
+    device = model_.device.type
+    if device == "cuda":
+        print(
+            f"peak memory allocated: {torch.cuda.max_memory_allocated() / 1024**2} MB"
+        )
+        print(f"peak memory reserved: {torch.cuda.max_memory_reserved() / 1024**2} MB")
+    model_config = model_.config
+    config = model_.config.get_text_config()
+
+    # Get key dimensions from the config
+    hidden_size = config.hidden_size
+    intermediate_size = config.intermediate_size
+    num_experts = config.num_local_experts
+
+    print(
+        f"Model dimensions: hidden_size={hidden_size}, intermediate_size={intermediate_size}, num_experts={num_experts}"
+    )
+
+    # Create output directory if it doesn't exist
+    os.makedirs(output, exist_ok=True)
+
+    # Get state dict
+    state_dict = model_.state_dict()
+    del model_
+
+    # Create a new state dict for the converted model
+    converted_state_dict = {}
+
+    # First, copy all keys that don't need modification
+    for key, value in iter_convert_patched_to_hf(state_dict, num_experts):
+        converted_state_dict[key] = value
+
+    del state_dict
+    if device == "cuda":
+        torch.cuda.empty_cache()
+        print("State dict converted.")
+        print(
+            f"peak memory allocated: {torch.cuda.max_memory_allocated() / 1024**2} MB"
+        )
+        print(f"peak memory reserved: {torch.cuda.max_memory_reserved() / 1024**2} MB")
+    # Ideally re-load the model import to load the converted state dict
+    # Save the converted model
+    with init_empty_weights():
+        unpatch_llama4()
+        model_ = Llama4ForConditionalGeneration(model_config)
+
+    if device == "cuda":
+        print("State dict loaded into model.")
+        print(
+            f"peak memory allocated: {torch.cuda.max_memory_allocated() / 1024**2} MB"
+        )
+        print(f"peak memory reserved: {torch.cuda.max_memory_reserved() / 1024**2} MB")
+    model_.load_state_dict(converted_state_dict, strict=False, assign=True)
+    print(f"Saving converted model to {output}...")
+    model_.save_pretrained(output)
+
+    print(f"Model successfully converted and saved to {output}")
+
+
+if __name__ == "__main__":
+    fire.Fire(do_cli)
diff --git a/src/axolotl/cli/evaluate.py b/src/axolotl/cli/evaluate.py
new file mode 100644
index 0000000000..1a73937a2c
--- /dev/null
+++ b/src/axolotl/cli/evaluate.py
@@ -0,0 +1,62 @@
+"""CLI to run evaluation on a model."""
+
+import os
+from pathlib import Path
+from typing import Union
+
+import fire
+from transformers.hf_argparser import HfArgumentParser
+
+from axolotl.cli.args import TrainerCliArgs
+from axolotl.cli.checks import check_accelerate_default_config, check_user_token
+from axolotl.cli.config import load_cfg
+from axolotl.common.datasets import load_datasets, load_preference_datasets
+from axolotl.evaluate import evaluate
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def do_evaluate(cfg: DictDefault, cli_args: TrainerCliArgs) -> None:
+    """
+    Evaluates a `transformers` model by first loading the dataset(s) specified in the
+    `axolotl` config, and then calling `axolotl.evaluate.evaluate`, which computes
+    evaluation metrics on the given dataset(s) and writes them to disk.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        cli_args: CLI arguments.
+    """
+
+    check_accelerate_default_config()
+    if int(os.getenv("LOCAL_RANK", "0")) == 0:
+        check_user_token()
+
+    if cfg.rl:
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
+    else:
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+    evaluate(cfg=cfg, dataset_meta=dataset_meta)
+
+
+def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs) -> None:
+    """
+    Parses `axolotl` config, CLI args, and calls `do_evaluate`.
+
+    Args:
+        config: Path to `axolotl` config YAML file.
+        kwargs: Additional keyword arguments to override config file values.
+    """
+
+    parsed_cfg = load_cfg(config, **kwargs)
+    parser = HfArgumentParser(TrainerCliArgs)
+    parsed_cli_args, _ = parser.parse_args_into_dataclasses(
+        return_remaining_strings=True
+    )
+    do_evaluate(parsed_cfg, parsed_cli_args)
+
+
+if __name__ == "__main__":
+    fire.Fire(do_cli)
diff --git a/src/axolotl/cli/generate_config_options.py b/src/axolotl/cli/generate_config_options.py
new file mode 100644
index 0000000000..7e3e97832a
--- /dev/null
+++ b/src/axolotl/cli/generate_config_options.py
@@ -0,0 +1,145 @@
+"""Generate lightweight Click option metadata for the Axolotl CLI."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from axolotl.cli.utils.args import ConfigOption, _strip_optional_type
+from axolotl.utils.schemas.config import AxolotlInputConfig
+
+DEFAULT_OUTPUT = Path("src/axolotl/cli/config_options.py")
+
+
+def _is_pydantic_model(field_type: type) -> bool:
+    try:
+        return isinstance(field_type, type) and issubclass(field_type, BaseModel)
+    except TypeError:
+        return False
+
+
+def _get_field_description(field) -> str | None:
+    if field.description:
+        return field.description
+    if field.json_schema_extra and isinstance(field.json_schema_extra, dict):
+        return field.json_schema_extra.get("description")
+    return None
+
+
+def _click_type_name(field_type: type) -> str | None:
+    return {str: "str", int: "int", float: "float"}.get(field_type)
+
+
+def _option_for_scalar(name: str, field) -> ConfigOption:
+    field_type = _strip_optional_type(field.annotation)
+    description = _get_field_description(field)
+    option_name = f"--{name.replace('_', '-')}"
+    if field_type is bool:
+        return (
+            (f"{option_name}/--no-{name.replace('_', '-')}",),
+            None,
+            None,
+            description,
+        )
+    return ((option_name,), None, None, description)
+
+
+def _option_for_nested(parent_name: str, sub_name: str, sub_field) -> ConfigOption:
+    sub_type = _strip_optional_type(sub_field.annotation)
+    cli_name = f"{parent_name}.{sub_name}".replace("_", "-")
+    param_name = f"{parent_name}__{sub_name}"
+    description = _get_field_description(sub_field)
+
+    if sub_type is bool:
+        return ((f"--{cli_name}/--no-{cli_name}",), param_name, None, description)
+    return ((f"--{cli_name}",), param_name, _click_type_name(sub_type), description)
+
+
+def _format_string(value: str | None) -> str:
+    if value is None:
+        return "None"
+    if '"' in value and "'" not in value:
+        return ascii(value)
+    return json.dumps(value)
+
+
+def _format_param_decls(param_decls: tuple[str, ...]) -> list[str]:
+    decls = ", ".join(_format_string(decl) for decl in param_decls)
+    if len(param_decls) == 1:
+        decls += ","
+    single_line = f"        ({decls}),"
+    if len(single_line) <= 88:
+        return [single_line]
+    return [
+        "        (",
+        *(f"            {_format_string(decl)}," for decl in param_decls),
+        "        ),",
+    ]
+
+
+def _format_option(option: ConfigOption) -> str:
+    param_decls, param_name, click_type_name, description = option
+    param_name_literal = _format_string(param_name)
+    click_type_literal = _format_string(click_type_name)
+    description_literal = _format_string(description)
+    return "\n".join(
+        [
+            "    (",
+            *_format_param_decls(param_decls),
+            f"        {param_name_literal},",
+            f"        {click_type_literal},",
+            f"        {description_literal},",
+            "    ),",
+        ]
+    )
+
+
+def generate_options() -> list[ConfigOption]:
+    options = []
+    for name, field in AxolotlInputConfig.model_fields.items():
+        field_type = _strip_optional_type(field.annotation)
+        if _is_pydantic_model(field_type):
+            for sub_name, sub_field in field_type.model_fields.items():
+                options.append(_option_for_nested(name, sub_name, sub_field))
+        else:
+            options.append(_option_for_scalar(name, field))
+    return options
+
+
+def generate_content() -> str:
+    options = generate_options()
+    return "\n".join(
+        [
+            '"""Generated Click option metadata for Axolotl config overrides."""',
+            "",
+            "# Generated by axolotl generate-cli-config-options.",
+            "# Do not edit by hand.",
+            "",
+            "from __future__ import annotations",
+            "",
+            "AXOLOTL_CONFIG_CLI_OPTIONS = (",
+            *(_format_option(option) for option in options),
+            ")",
+            "",
+        ]
+    )
+
+
+def write_options(output: Path = DEFAULT_OUTPUT, check: bool = False) -> None:
+    content = generate_content()
+    if check:
+        current = output.read_text(encoding="utf-8") if output.exists() else ""
+        if current != content:
+            raise SystemExit(f"{output} is out of date")
+        print(f"{output} is up to date")
+        return
+
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(content, encoding="utf-8")
+    print(f"Wrote {output}")
+
+
+def main() -> None:
+    write_options()
diff --git a/src/axolotl/cli/inference.py b/src/axolotl/cli/inference.py
index 86ad8409ff..c6fd8d1dd8 100644
--- a/src/axolotl/cli/inference.py
+++ b/src/axolotl/cli/inference.py
@@ -1,32 +1,318 @@
-"""
-CLI to run inference on a trained model
-"""
+"""CLI to run inference on a trained model."""
+
+import importlib
+import sys
 from pathlib import Path
+from threading import Thread
+from typing import Union
 
 import fire
+import torch
 import transformers
+from transformers import GenerationConfig, TextIteratorStreamer, TextStreamer
 
-from axolotl.cli import (
-    do_inference,
-    do_inference_gradio,
-    load_cfg,
-    print_axolotl_text_art,
+from axolotl.cli.args import InferenceCliArgs
+from axolotl.cli.config import load_cfg
+from axolotl.cli.utils import load_model_and_tokenizer, resolve_chat_template_str
+from axolotl.cli.utils.diffusion import (
+    diffusion_inference,
+    launch_diffusion_gradio_ui,
 )
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.integrations.base import PluginManager
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def get_multi_line_input() -> str:
+    """
+    Gets multi-line input from terminal.
+
+    Returns:
+        Possibly multi-line, possibly empty stdin input as a string.
+    """
+    print("Give me an instruction (Ctrl + D to submit): ")
+    print("=" * 80)
+
+    instruction = ""
+    for line in sys.stdin:
+        instruction += line
+
+    return instruction
+
+
+@send_errors
+def do_inference(
+    *,
+    cfg: DictDefault,
+    cli_args: InferenceCliArgs,
+):
+    """
+    Runs inference on the command line in a loop. User input is accepted, a chat
+    template is (optionally) applied, and the model specified in the `axolotl` config is
+    used to generate completions according to a default generation config.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        cli_args: Inference-specific CLI arguments.
+    """
+    model, tokenizer, _ = load_model_and_tokenizer(cfg=cfg, inference=True)
+    prompter = cli_args.prompter
+
+    prompter_module = None
+    chat_template_str = None
+    if prompter:
+        prompter_module = getattr(
+            importlib.import_module("axolotl.prompters"), prompter
+        )
+    else:
+        chat_template_str = resolve_chat_template_str(cfg, tokenizer)
+
+    model = model.to(cfg.device, dtype=cfg.torch_dtype)
+
+    # Detect diffusion mode
+    plugin_manager = PluginManager.get_instance()
+    is_diffusion = any(
+        plugin.__class__.__name__ == "DiffusionPlugin"
+        for plugin in plugin_manager.plugins.values()
+    )
+
+    if is_diffusion:
+        print("=" * 80)
+        print("Commands:")
+        print(":complete N -> completion mode with N tokens (default 64)")
+        print(":mask R     -> random masking with ratio R (0.0–1.0)")
+
+    while True:
+        print("=" * 80)
+        instruction = get_multi_line_input()
+        if not instruction:
+            return
+
+        if prompter_module:
+            prompt: str = next(
+                prompter_module().build_prompt(instruction=instruction.strip("\n"))
+            )
+        else:
+            prompt = instruction.strip()
+
+        if chat_template_str:
+            batch = tokenizer.apply_chat_template(
+                [
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    }
+                ],
+                return_tensors="pt",
+                add_special_tokens=True,
+                add_generation_prompt=True,
+                chat_template=chat_template_str,
+                tokenize=True,
+                return_dict=True,
+            )
+        else:
+            batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
+
+        print("=" * 80)
+        model.eval()
+        with torch.no_grad():
+            if is_diffusion:
+                diffusion_inference(
+                    model=model,
+                    tokenizer=tokenizer,
+                    cfg=cfg,
+                    prompt=prompt,
+                    chat_template_str=chat_template_str,
+                )
+                continue
+
+            generation_config = GenerationConfig(
+                repetition_penalty=1.1,
+                max_new_tokens=1024,
+                temperature=0.9,
+                top_p=0.95,
+                top_k=40,
+                bos_token_id=tokenizer.bos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
+                do_sample=True,
+                use_cache=True,
+                return_dict_in_generate=True,
+                output_attentions=False,
+                output_hidden_states=False,
+                output_scores=False,
+            )
+            streamer = TextStreamer(tokenizer)
+            generated = model.generate(
+                inputs=batch["input_ids"].to(cfg.device),
+                generation_config=generation_config,
+                streamer=streamer,
+            )
+        print("=" * 80)
+        print(tokenizer.decode(generated["sequences"].cpu().tolist()[0]))
+
+
+@send_errors
+def do_inference_gradio(
+    *,
+    cfg: DictDefault,
+    cli_args: InferenceCliArgs,
+):
+    """
+    Runs inference in a Gradio interface. User input is accepted, a chat template is
+    (optionally) applied, and the model specified in the `axolotl` config is used to
+    generate completions according to a default generation config.
 
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        cli_args: Inference-specific CLI arguments.
+    """
+    import gradio as gr
 
-def do_cli(config: Path = Path("examples/"), gradio=False, **kwargs):
-    # pylint: disable=duplicate-code
-    print_axolotl_text_art()
-    parsed_cfg = load_cfg(config, **kwargs)
+    model, tokenizer, _ = load_model_and_tokenizer(cfg=cfg, inference=True)
+    prompter = cli_args.prompter
+
+    prompter_module = None
+    chat_template_str = None
+    if prompter:
+        prompter_module = getattr(
+            importlib.import_module("axolotl.prompters"), prompter
+        )
+    else:
+        chat_template_str = resolve_chat_template_str(cfg, tokenizer)
+
+    model = model.to(cfg.device, dtype=cfg.torch_dtype)
+
+    # Detect diffusion mode
+    plugin_manager = PluginManager.get_instance()
+    is_diffusion = any(
+        plugin.__class__.__name__ == "DiffusionPlugin"
+        for plugin in plugin_manager.plugins.values()
+    )
+
+    if is_diffusion:
+        launch_diffusion_gradio_ui(
+            model=model,
+            tokenizer=tokenizer,
+            cfg=cfg,
+            prompter_module=prompter_module,
+            chat_template_str=chat_template_str,
+        )
+        return
+
+    def generate(instruction):
+        if not instruction:
+            return
+        if prompter_module:
+            prompt: str = next(
+                prompter_module().build_prompt(instruction=instruction.strip("\n"))
+            )
+        else:
+            prompt = instruction.strip()
+
+        if chat_template_str:
+            batch = tokenizer.apply_chat_template(
+                [
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    }
+                ],
+                return_tensors="pt",
+                add_special_tokens=True,
+                add_generation_prompt=True,
+                chat_template=chat_template_str,
+                tokenize=True,
+                return_dict=True,
+            )
+        else:
+            batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
+
+        model.eval()
+        with torch.no_grad():
+            generation_config = GenerationConfig(
+                repetition_penalty=1.1,
+                max_new_tokens=cfg.get("gradio_max_new_tokens", 1024),
+                temperature=cfg.get("gradio_temperature", 0.9),
+                top_p=0.95,
+                top_k=40,
+                bos_token_id=tokenizer.bos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
+                do_sample=True,
+                use_cache=True,
+                return_dict_in_generate=True,
+                output_attentions=False,
+                output_hidden_states=False,
+                output_scores=False,
+            )
+            streamer = TextIteratorStreamer(tokenizer)
+            generation_kwargs = {
+                "inputs": batch["input_ids"].to(cfg.device),
+                "attention_mask": batch["attention_mask"].to(cfg.device),
+                "generation_config": generation_config,
+                "streamer": streamer,
+            }
+
+            thread = Thread(target=model.generate, kwargs=generation_kwargs)
+            thread.start()
+
+            all_text = ""
+
+            for new_text in streamer:
+                all_text += new_text
+                yield all_text
+
+    demo = gr.Interface(
+        fn=generate,
+        inputs="textbox",
+        outputs="text",
+        title=cfg.get("gradio_title", "Axolotl Gradio Interface"),
+    )
+
+    demo.launch(
+        footer_links=["gradio", "settings"],
+        share=cfg.get("gradio_share", True),
+        server_name=cfg.get("gradio_server_name", "127.0.0.1"),
+        server_port=cfg.get("gradio_server_port", None),
+    )
+
+
+def do_cli(
+    config: Union[Path, str] = Path("examples/"),
+    gradio: bool = False,
+    chat: bool = False,
+    **kwargs,
+) -> None:
+    """
+    Parses axolotl config, CLI args, and calls `do_inference`, `do_inference_gradio`,
+    or `do_chat`.
+
+    Args:
+        config: Path to `axolotl` config YAML file.
+        gradio: Whether to launch the Gradio browser interface.
+        chat: Whether to launch the interactive multi-turn chat interface.
+        kwargs: Additional keyword arguments to override config file values.
+    """
+
+    parsed_cfg = load_cfg(config, inference=True, rl=None, **kwargs)
     parsed_cfg.sample_packing = False
-    parser = transformers.HfArgumentParser((TrainerCliArgs))
+    parser = transformers.HfArgumentParser(InferenceCliArgs)
     parsed_cli_args, _ = parser.parse_args_into_dataclasses(
         return_remaining_strings=True
     )
-    parsed_cli_args.inference = True
 
-    if gradio:
+    if gradio and chat:
+        raise ValueError("--gradio and --chat are mutually exclusive.")
+
+    if chat:
+        from axolotl.cli.chat import do_chat
+
+        do_chat(cfg=parsed_cfg, cli_args=parsed_cli_args)
+    elif gradio:
         do_inference_gradio(cfg=parsed_cfg, cli_args=parsed_cli_args)
     else:
         do_inference(cfg=parsed_cfg, cli_args=parsed_cli_args)
diff --git a/src/axolotl/cli/main.py b/src/axolotl/cli/main.py
new file mode 100644
index 0000000000..9b5322b2a5
--- /dev/null
+++ b/src/axolotl/cli/main.py
@@ -0,0 +1,499 @@
+"""Click CLI definitions for various axolotl commands."""
+
+import os
+import subprocess  # nosec B404
+from pathlib import Path
+from typing import Literal, Optional
+
+import click
+from dotenv import load_dotenv
+
+import axolotl
+from axolotl.cli.args import (
+    EvaluateCliArgs,
+    PreprocessCliArgs,
+    QuantizeCliArgs,
+    TrainerCliArgs,
+    VllmServeCliArgs,
+)
+from axolotl.cli.art import print_axolotl_text_art
+from axolotl.cli.config_options import AXOLOTL_CONFIG_CLI_OPTIONS
+from axolotl.cli.utils import (
+    add_options_from_config_options,
+    add_options_from_dataclass,
+    build_command,
+    fetch_from_github,
+    filter_none_kwargs,
+    generate_config_files,
+    launch_training,
+)
+from axolotl.utils import set_misc_env, set_pytorch_cuda_alloc_conf
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+LAUNCHER_COMMAND_MAPPING = {
+    "accelerate": ["accelerate", "launch"],
+    "torchrun": ["torchrun"],
+}
+
+
+@click.group()
+@click.version_option(version=axolotl.__version__, prog_name="axolotl")
+def cli():
+    """Axolotl CLI - Train and fine-tune large language models"""
+    print_axolotl_text_art()
+    load_dotenv()
+    set_pytorch_cuda_alloc_conf()
+    set_misc_env()
+
+
+@cli.command()
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@click.option("--cloud", default=None, type=click.Path(exists=True, path_type=str))
+@add_options_from_dataclass(PreprocessCliArgs)
+@add_options_from_config_options(AXOLOTL_CONFIG_CLI_OPTIONS)
+@filter_none_kwargs
+def preprocess(config: str, cloud: Optional[str] = None, **kwargs):
+    """
+    Preprocess datasets before training.
+
+    Args:
+        config: Path to `axolotl` config YAML file.
+        cloud: Path to a cloud accelerator configuration file.
+        kwargs: Additional keyword arguments which correspond to CLI args or `axolotl`
+            config options.
+    """
+
+    if cloud:
+        from axolotl.cli.cloud import do_cli_preprocess
+
+        do_cli_preprocess(cloud_config=cloud, config=config)
+    else:
+        from axolotl.cli.preprocess import do_cli
+
+        do_cli(config=config, **kwargs)
+
+
+@cli.command(
+    context_settings={"ignore_unknown_options": True, "allow_extra_args": True}
+)
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@click.option(
+    "--launcher",
+    type=click.Choice(["accelerate", "torchrun", "python"]),
+    default="accelerate",
+    help="Launcher to use for multi-GPU training",
+)
+@click.option("--cloud", default=None, type=click.Path(exists=True, path_type=str))
+@click.option(
+    "--sweep",
+    type=click.Path(exists=True, path_type=str),
+    help="YAML config for sweeping hyperparameters",
+)
+@add_options_from_dataclass(TrainerCliArgs)
+@add_options_from_config_options(AXOLOTL_CONFIG_CLI_OPTIONS)
+@filter_none_kwargs
+@click.pass_context
+def train(
+    ctx: click.Context,
+    config: str,
+    launcher: Literal["accelerate", "torchrun", "python"] = "accelerate",
+    cloud: str | None = None,
+    sweep: str | None = None,
+    **kwargs,
+):
+    """
+    Train or fine-tune a model.
+
+    Args:
+        ctx: Click context for extra args.
+        config: Path to `axolotl` config YAML file.
+        launcher: Launcher to use for multi-GPU training ("accelerate", "torchrun", or "python").
+        cloud: Path to a cloud accelerator configuration file
+        sweep: Path to YAML config for sweeping hyperparameters.
+        kwargs: Additional keyword arguments which correspond to CLI args or `axolotl`
+            config options.
+    """
+    # Extract launcher args from extra args (after --)
+    launcher_args = ctx.args if ctx.args else []
+
+    # Handle Ray launcher override
+    _launcher = None if kwargs.get("use_ray") else launcher
+
+    # Process each configuration
+    for cfg_file, is_group in generate_config_files(config, sweep):
+        try:
+            use_exec = is_group is not True
+            launch_training(cfg_file, _launcher, cloud, kwargs, launcher_args, use_exec)
+        except subprocess.CalledProcessError as exc:
+            LOG.error(f"Failed to train/fine-tune config '{cfg_file}': {exc}")
+            if not sweep:
+                raise exc
+        finally:
+            # Only delete temp files, not the original config
+            if cfg_file != config:
+                os.unlink(cfg_file)
+
+
+@cli.command(
+    context_settings={"ignore_unknown_options": True, "allow_extra_args": True}
+)
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@click.option(
+    "--launcher",
+    type=click.Choice(["accelerate", "torchrun", "python"]),
+    default="accelerate",
+    help="Launcher to use for multi-GPU evaluation",
+)
+@add_options_from_dataclass(EvaluateCliArgs)
+@add_options_from_config_options(AXOLOTL_CONFIG_CLI_OPTIONS)
+@filter_none_kwargs
+@click.pass_context
+def evaluate(ctx: click.Context, config: str, launcher: str, **kwargs):
+    """
+    Evaluate a model.
+
+    Args:
+        ctx: Click context for extra args.
+        config: Path to `axolotl` config YAML file.
+        launcher: Launcher to use for multi-GPU evaluation ("accelerate", "torchrun", or "python").
+        kwargs: Additional keyword arguments which correspond to CLI args or `axolotl`
+            config options.
+    """
+    # Extract launcher args from extra args (after --)
+    launcher_args = ctx.args if ctx.args else []
+
+    if launcher in LAUNCHER_COMMAND_MAPPING:
+        base_cmd = (
+            LAUNCHER_COMMAND_MAPPING[launcher]
+            + launcher_args
+            + ["-m", "axolotl.cli.evaluate"]
+        )
+        if config:
+            base_cmd.append(config)
+        cmd = build_command(base_cmd, kwargs)
+        subprocess.run(cmd, check=True)  # nosec B603
+    else:
+        from axolotl.cli.evaluate import do_cli
+
+        do_cli(config=config, **kwargs)
+
+
+@cli.command(
+    context_settings={"ignore_unknown_options": True, "allow_extra_args": True}
+)
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@click.option(
+    "--launcher",
+    type=click.Choice(["accelerate", "torchrun", "python"]),
+    default="accelerate",
+    help="Launcher to use for multi-GPU inference",
+)
+@click.option("--gradio", is_flag=True, help="Launch Gradio interface")
+@click.option(
+    "--chat", is_flag=True, help="Launch interactive multi-turn chat interface"
+)
+@add_options_from_dataclass(TrainerCliArgs)
+@add_options_from_config_options(AXOLOTL_CONFIG_CLI_OPTIONS)
+@filter_none_kwargs
+@click.pass_context
+def inference(
+    ctx: click.Context, config: str, launcher: str, gradio: bool, chat: bool, **kwargs
+):
+    """
+    Run inference with a trained model.
+
+    Args:
+        ctx: Click context for extra args.
+        config: Path to `axolotl` config YAML file.
+        launcher: Launcher to use for multi-GPU inference ("accelerate", "torchrun", or "python").
+        gradio: Whether to use Gradio browser interface or command line for inference.
+        chat: Whether to use the interactive multi-turn chat interface.
+        kwargs: Additional keyword arguments which correspond to CLI args or `axolotl`
+            config options.
+    """
+    if gradio and chat:
+        raise click.UsageError("--gradio and --chat are mutually exclusive.")
+
+    # Extract launcher args from extra args (after --)
+    launcher_args = ctx.args if ctx.args else []
+
+    if launcher in LAUNCHER_COMMAND_MAPPING:
+        base_cmd = (
+            LAUNCHER_COMMAND_MAPPING[launcher]
+            + launcher_args
+            + ["-m", "axolotl.cli.inference"]
+        )
+        if config:
+            base_cmd.append(config)
+        if gradio:
+            base_cmd.append("--gradio")
+        if chat:
+            base_cmd.append("--chat")
+        cmd = build_command(base_cmd, kwargs)
+        subprocess.run(cmd, check=True)  # nosec B603
+    else:
+        from axolotl.cli.inference import do_cli
+
+        do_cli(config=config, gradio=gradio, chat=chat, **kwargs)
+
+
+@cli.command(
+    context_settings={"ignore_unknown_options": True, "allow_extra_args": True}
+)
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@click.option(
+    "--launcher",
+    type=click.Choice(["accelerate", "torchrun", "python"]),
+    default="accelerate",
+    help="Launcher to use for weight merging",
+)
+@add_options_from_dataclass(TrainerCliArgs)
+@add_options_from_config_options(AXOLOTL_CONFIG_CLI_OPTIONS)
+@filter_none_kwargs
+@click.pass_context
+def merge_sharded_fsdp_weights(
+    ctx: click.Context, config: str, launcher: str, **kwargs
+):
+    """
+    Merge sharded FSDP model weights.
+
+    Args:
+        ctx: Click context for extra args.
+        config: Path to `axolotl` config YAML file.
+        launcher: Launcher to use for weight merging ("accelerate", "torchrun", or "python").
+        kwargs: Additional keyword arguments which correspond to CLI args or `axolotl`
+            config options.
+    """
+    # Extract launcher args from extra args (after --)
+    launcher_args = ctx.args if ctx.args else []
+
+    if launcher in LAUNCHER_COMMAND_MAPPING:
+        base_cmd = (
+            LAUNCHER_COMMAND_MAPPING[launcher]
+            + launcher_args
+            + ["-m", "axolotl.cli.merge_sharded_fsdp_weights"]
+        )
+        if config:
+            base_cmd.append(config)
+        cmd = build_command(base_cmd, kwargs)
+        subprocess.run(cmd, check=True)  # nosec B603
+    else:
+        from axolotl.cli.merge_sharded_fsdp_weights import do_cli
+
+        do_cli(config=config, **kwargs)
+
+
+@cli.command()
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@add_options_from_dataclass(TrainerCliArgs)
+@add_options_from_config_options(AXOLOTL_CONFIG_CLI_OPTIONS)
+@filter_none_kwargs
+def merge_lora(config: str, **kwargs):
+    """
+    Merge trained LoRA adapters into a base model.
+
+    Args:
+        config: Path to `axolotl` config YAML file.
+        kwargs: Additional keyword arguments which correspond to CLI args or `axolotl`
+            config options.
+    """
+    from axolotl.cli.merge_lora import do_cli
+
+    do_cli(config=config, **kwargs)
+
+
+@cli.command()
+@click.argument(
+    "directory", type=click.Choice(["examples", "deepspeed_configs", "docs"])
+)
+@click.option("--dest", help="Destination directory")
+def fetch(directory: str, dest: Optional[str]):
+    """
+    Fetch example configs or other resources.
+
+    Available directories:
+    - examples: Example configuration files
+    - deepspeed_configs: DeepSpeed configuration files
+    - docs: Full documentation (Quarto markdown files)
+
+    Args:
+        directory: One of `examples`, `deepspeed_configs`, `docs`.
+        dest: Optional destination directory.
+    """
+    fetch_from_github(f"{directory}/", dest)
+
+
+@cli.command()
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@add_options_from_dataclass(VllmServeCliArgs)
+@filter_none_kwargs
+def vllm_serve(config: str, **cli_args: VllmServeCliArgs):
+    from axolotl.cli.vllm_serve import do_vllm_serve
+
+    do_vllm_serve(config, cli_args)
+
+
+@cli.command()
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@add_options_from_dataclass(QuantizeCliArgs)
+@filter_none_kwargs
+def quantize(config: str, **cli_args: QuantizeCliArgs):
+    from axolotl.cli.quantize import do_quantize
+
+    do_quantize(config, cli_args)
+
+
+@cli.command()
+@click.argument("model", type=click.Path(exists=True, path_type=str))
+@click.argument("output", type=click.Path(exists=False, path_type=str))
+def delinearize_llama4(model: str, output: str):
+    from axolotl.cli.delinearize_llama4 import do_cli as do_delinearize_llama4
+
+    do_delinearize_llama4(model, output)
+
+
+@cli.command("generate-cli-config-options")
+@click.option(
+    "--output",
+    type=click.Path(dir_okay=False, path_type=Path),
+    default=Path("src/axolotl/cli/config_options.py"),
+    show_default=True,
+    help="Path to write generated option metadata.",
+)
+@click.option("--check", is_flag=True, help="Fail if generated metadata is stale.")
+def generate_cli_config_options(output: Path, check: bool):
+    """Regenerate CLI config override option metadata."""
+    from axolotl.cli.generate_config_options import write_options
+
+    write_options(output, check=check)
+
+
+@cli.command("lm-eval")
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@click.option("--cloud", default=None, type=click.Path(exists=True, path_type=str))
+@click.pass_context
+def lm_eval(ctx: click.Context, config: str, cloud: Optional[str] = None):
+    """
+    use lm eval to evaluate a trained language model
+    """
+    from axolotl.integrations.lm_eval.cli import lm_eval as _lm_eval
+
+    ctx.invoke(_lm_eval, config=config, cloud=cloud)
+
+
+@cli.command("agent-docs")
+@click.argument("topic", required=False, default=None)
+@click.option("--list", "list_topics", is_flag=True, help="List available topics")
+def agent_docs(topic: Optional[str], list_topics: bool):
+    """Show agent-optimized documentation.
+
+    Prints reference docs designed for AI coding agents.
+    These docs are bundled with the package — no network access needed.
+
+    \b
+    Examples:
+        axolotl agent-docs              # overview (start here)
+        axolotl agent-docs grpo         # GRPO reference
+        axolotl agent-docs sft          # SFT reference
+        axolotl agent-docs --list       # list all topics
+    """
+    from axolotl.cli.agent_docs import get_doc, list_topics as _list_topics
+
+    if list_topics:
+        for name, title in _list_topics().items():
+            click.echo(f"  {name:25s} {title}")
+        return
+
+    if topic is None:
+        topic = "overview"
+
+    try:
+        click.echo(get_doc(topic))
+    except FileNotFoundError as exc:
+        raise click.BadParameter(str(exc)) from exc
+
+
+@cli.command("config-schema")
+@click.option(
+    "--format",
+    "output_format",
+    type=click.Choice(["json", "yaml"]),
+    default="json",
+    help="Output format (default: json)",
+)
+@click.option("--field", help="Show schema for a specific field only")
+def config_schema(output_format: str, field: Optional[str]):
+    """Dump the full config JSON schema.
+
+    Useful for AI agents and tooling to discover all available config options,
+    their types, defaults, and descriptions.
+
+    \b
+    Examples:
+        axolotl config-schema                    # full JSON schema
+        axolotl config-schema --format yaml      # YAML format
+        axolotl config-schema --field adapter     # single field
+    """
+    import json
+
+    from axolotl.utils.schemas.config import AxolotlInputConfig
+
+    try:
+        schema = AxolotlInputConfig.model_json_schema()
+    except (TypeError, ValueError, AttributeError) as exc:
+        # Fallback: dump field names, types, and defaults when full schema
+        # generation fails (e.g. torch.dtype not JSON-serializable)
+        LOG.warning(
+            "Full JSON schema generation failed, using simplified fallback: %s", exc
+        )
+        fields = {}
+        for name, field_info in AxolotlInputConfig.model_fields.items():
+            entry = {}
+            if field_info.description:
+                entry["description"] = field_info.description
+            if field_info.default is not None:
+                try:
+                    json.dumps(field_info.default)
+                    entry["default"] = field_info.default
+                except (TypeError, ValueError):
+                    entry["default"] = str(field_info.default)
+            annotation = field_info.annotation
+            if annotation is not None:
+                entry["type"] = str(annotation)
+            fields[name] = entry
+        schema = {
+            "properties": fields,
+            "_note": "simplified schema (full generation failed)",
+        }
+
+    if field:
+        props = schema.get("properties", {})
+        if field not in props:
+            # Try case-insensitive match
+            matches = [k for k in props if k.lower() == field.lower()]
+            if matches:
+                field = matches[0]
+            else:
+                raise click.BadParameter(
+                    f"Unknown field: {field!r}. "
+                    f"Omit --field to dump the full schema, "
+                    f"or pipe to jq: axolotl config-schema | jq '.properties | keys'"
+                )
+        schema = {field: props[field]}
+
+    if output_format == "yaml":
+        import yaml  # pylint: disable=import-outside-toplevel
+
+        click.echo(yaml.dump(schema, default_flow_style=False, sort_keys=False))
+    else:
+        click.echo(json.dumps(schema, indent=2))
+
+
+def main():
+    cli()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/axolotl/cli/merge_lora.py b/src/axolotl/cli/merge_lora.py
index 8db3fa9897..8253584b5c 100644
--- a/src/axolotl/cli/merge_lora.py
+++ b/src/axolotl/cli/merge_lora.py
@@ -1,50 +1,179 @@
-"""
-CLI to run merge a trained LoRA into a base model
-"""
+"""CLI to merge a trained LoRA into a base model."""
+
 from pathlib import Path
+from typing import Union
 
 import fire
-import transformers
+import torch
+
+from axolotl.cli.config import load_cfg
+from axolotl.cli.utils import load_model_and_tokenizer
+from axolotl.cli.utils.lora_merge import merge_lora_sharded_efficient
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+@send_errors
+def do_merge_lora(*, cfg: DictDefault) -> None:
+    """
+    Merges LoRA adapters with base model using either memory-efficient or legacy approach.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+    """
+    merge_method = str(getattr(cfg, "merge_method", "memory_efficient"))
+    if merge_method == "legacy":
+        LOG.debug("Using legacy LoRA merging method...")
+        _do_merge_lora_legacy(cfg=cfg)
+    else:
+        LOG.debug("Using memory-efficient LoRA merging method...")
+        _do_merge_lora_efficient(cfg=cfg)
+
+
+def _do_merge_lora_legacy(*, cfg: DictDefault) -> None:
+    """
+    Legacy LoRA merging using merge_and_unload.
+    Loads the full model into memory before merging.
+    """
+    LOG.debug("Using legacy LoRA merging method...")
+    model, tokenizer, processor = load_model_and_tokenizer(cfg=cfg)
+
+    LOG.info("Running merge of LoRA with base model...")
+    model = model.merge_and_unload(progressbar=True)
+    try:
+        model.to(dtype=cfg.torch_dtype)
+    except ValueError as e:
+        LOG.warning("Failed to convert model to dtype %s", cfg.torch_dtype)
+        LOG.warning("Ignore this if the base_model is pre-quantized.")
+        LOG.warning("Error raised: %s", e)
+
+    model.generation_config.do_sample = True
+    model.config.use_cache = True
+
+    if cfg.local_rank == 0:
+        LOG.info(f"Saving merged model to: {str(Path(cfg.output_dir) / 'merged')}...")
+        model.save_pretrained(
+            str(Path(cfg.output_dir) / "merged"),
+            progressbar=True,
+        )
+        tokenizer.save_pretrained(
+            str(Path(cfg.output_dir) / "merged"),
+            save_jinja_files=cfg.tokenizer_save_jinja_files,
+        )
+
+        if processor:
+            processor.save_pretrained(str(Path(cfg.output_dir) / "merged"))
+
 
-from axolotl.cli import do_merge_lora, load_cfg, print_axolotl_text_art
-from axolotl.common.cli import TrainerCliArgs
+def _do_merge_lora_efficient(*, cfg: DictDefault) -> None:
+    """
+    Memory-efficient LoRA merging using shard-by-shard processing.
+    Does not load the full model into memory.
 
+    Supports standard LoRA, RSLoRA, and DoRA. Unsupported methods (AdaLoRA, VeRA)
+    will raise NotImplementedError — use legacy method for those.
+    """
+    LOG.debug("Using memory-efficient LoRA merging method...")
 
-def do_cli(config: Path = Path("examples/"), **kwargs):
-    # pylint: disable=duplicate-code
-    print_axolotl_text_art()
-    parser = transformers.HfArgumentParser((TrainerCliArgs))
-    parsed_cli_args, _ = parser.parse_args_into_dataclasses(
-        return_remaining_strings=True
+    output_path = Path(cfg.output_dir) / "merged"
+    safe_tensors = getattr(cfg, "save_safetensors", True)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Detect NF4 quantization from config to simulate QLoRA training dynamics.
+    # Check both current and original (pre-override) config values since do_cli
+    # forces load_in_4bit=False for the legacy path.
+    simulate_nf4 = bool(
+        getattr(cfg, "load_in_4bit", False)
+        or getattr(cfg, "_original_load_in_4bit", False)
+        or getattr(cfg, "adapter", None) == "qlora"
+        or getattr(cfg, "_original_adapter", None) == "qlora"
+    )
+
+    bnb_config_kwargs = getattr(cfg, "bnb_config_kwargs", None) or {}
+    nf4_blocksize = bnb_config_kwargs.get("blocksize", None)
+    nf4_double_quant = bnb_config_kwargs.get(
+        "bnb_4bit_use_double_quant",
+        getattr(cfg, "bnb_4bit_use_double_quant", True),
+    )
+
+    # Detect MoE expert quantization
+    simulate_nf4_experts = bool(
+        getattr(cfg, "quantize_moe_experts", False)
+        or getattr(cfg, "_original_quantize_moe_experts", False)
+    )
+
+    merge_lora_sharded_efficient(
+        base_model_path=cfg.base_model,
+        lora_adapter_path=cfg.lora_model_dir,
+        output_path=output_path,
+        safe_tensors=safe_tensors,
+        device=device,
+        simulate_nf4=simulate_nf4,
+        simulate_nf4_experts=simulate_nf4_experts,
+        nf4_blocksize=nf4_blocksize,
+        nf4_double_quant=nf4_double_quant,
+        trust_remote_code=bool(getattr(cfg, "trust_remote_code", False)),
+        dequant=bool(getattr(cfg, "merge_dequant", False)),
     )
-    parsed_cli_args.merge_lora = True
+
+    LOG.debug("Memory-efficient LoRA merge completed successfully!")
+
+
+def do_cli(
+    config: Union[Path, str] = Path("examples/"), dequant: bool = False, **kwargs
+) -> None:
+    """
+    Parses `axolotl` config, CLI args, and calls `do_merge_lora`. Note that various
+    config values will be overwritten to allow the LoRA merge logic to work as expected
+    (`load_in_8bit=False`, `load_in4bit=False`, `flash_attention=False`, etc.).
+
+    Args:
+        config: Path to `axolotl` config YAML file.
+        dequant: If set (``--dequant``), dequantize a quantized base to bf16 in the merged
+            checkpoint. Default keeps the base's quantized dtypes (fp8/nvfp4) intact.
+        kwargs: Additional keyword arguments to override config file values.
+
+    Raises:
+        ValueError: If target directory for LoRA merged model does not exist.
+    """
+
+    # Pre-load config to detect original quantization settings before overrides
+    raw_cfg = load_cfg(config, **kwargs)
+    original_load_in_4bit = getattr(raw_cfg, "load_in_4bit", False)
+    original_adapter = getattr(raw_cfg, "adapter", None)
+    original_quantize_moe_experts = getattr(raw_cfg, "quantize_moe_experts", False)
 
     parsed_cfg = load_cfg(
         config,
         merge_lora=True,
         load_in_8bit=False,
         load_in_4bit=False,
-        flash_attention=False,
+        quantize_moe_experts=False,
+        attn_implementation=None,
+        context_parallel_size=None,
         deepspeed=None,
         fsdp=None,
+        fsdp_config=None,
         **kwargs,
     )
 
+    # Stash original quantization settings for NF4 simulation in efficient merge
+    parsed_cfg._original_load_in_4bit = original_load_in_4bit
+    parsed_cfg._original_adapter = original_adapter
+    parsed_cfg._original_quantize_moe_experts = original_quantize_moe_experts
+    parsed_cfg.merge_dequant = bool(dequant)
+
     if not parsed_cfg.lora_model_dir and parsed_cfg.output_dir:
         parsed_cfg.lora_model_dir = parsed_cfg.output_dir
     if not Path(parsed_cfg.lora_model_dir).exists():
         raise ValueError(
-            f"Target directory for merge: `{parsed_cfg.lora_model_dir}` does not exist."
+            f"Target directory for LoRA adapter weights does not exist: `{parsed_cfg.lora_model_dir}`"
         )
 
-    parsed_cfg.load_in_4bit = False
-    parsed_cfg.load_in_8bit = False
-    parsed_cfg.flash_attention = False
-    parsed_cfg.deepspeed = None
-    parsed_cfg.fsdp = None
-    parsed_cfg.fsdp_config = None
-
-    do_merge_lora(cfg=parsed_cfg, cli_args=parsed_cli_args)
+    do_merge_lora(cfg=parsed_cfg)
 
 
 if __name__ == "__main__":
diff --git a/src/axolotl/cli/merge_sharded_fsdp_weights.py b/src/axolotl/cli/merge_sharded_fsdp_weights.py
new file mode 100644
index 0000000000..f4779f5d9b
--- /dev/null
+++ b/src/axolotl/cli/merge_sharded_fsdp_weights.py
@@ -0,0 +1,207 @@
+"""CLI to merge sharded FSDP model checkpoints into a single combined checkpoint."""
+
+import json
+import os
+import shutil
+from pathlib import Path
+from typing import Dict, Union
+
+import fire
+import torch
+import torch.distributed.checkpoint as dist_cp
+import torch.distributed.checkpoint.format_utils as dist_cp_format_utils
+from accelerate import PartialState
+from accelerate.utils import (
+    SAFE_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_NAME,
+    is_torch_version,
+)
+from huggingface_hub import split_torch_state_dict_into_shards
+from safetensors.torch import save_file as safe_save_file
+from torch.distributed.checkpoint.format_utils import _EmptyStateDictLoadPlanner
+
+from axolotl.cli.config import load_cfg
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.logging import get_logger
+from axolotl.utils.train import determine_last_checkpoint
+
+LOG = get_logger(__name__)
+
+
+class BFloat16CastPlanner(_EmptyStateDictLoadPlanner):
+    """A custom planner to cast tensors to bfloat16 on the fly during loading."""
+
+    def commit_tensor(self, read_item, tensor):
+        tensor.copy_(tensor.to(torch.bfloat16))
+
+
+def _distributed_checkpoint_to_merged_weights(
+    checkpoint_dir: Union[str, Path],
+    save_path: str,
+    max_shard_size: str = "5GB",
+) -> Path:
+    """
+    Passthrough to `torch.distributed.checkpoint.format_utils.dcp_to_torch_save`. Will
+    save under `save_path` as `model.safetensors`.
+
+    Args:
+        checkpoint_dir: Directory where distributed checkpoint is saved.
+        save_path: Path to save model to.
+        max_shard_size: Max size of model shards to save.
+
+    Returns:
+        Path where model is saved.
+    """
+
+    state_dict: Dict = {}
+    save_path_ = Path(save_path)
+    save_path_.mkdir(exist_ok=True)
+    dist_cp_format_utils._load_state_dict(
+        state_dict,
+        storage_reader=dist_cp.FileSystemReader(checkpoint_dir),
+        planner=BFloat16CastPlanner(),
+        no_dist=True,
+    )
+
+    # To handle if state is a dict like {model: {...}}
+    if len(state_dict.keys()) == 1:
+        state_dict = state_dict[list(state_dict)[0]]
+
+    # Ensure all tensors are in bfloat16
+    for key, value in state_dict.items():
+        if isinstance(value, torch.Tensor) and value.dtype != torch.bfloat16:
+            state_dict[key] = value.to(torch.bfloat16)
+
+    filename_pattern = SAFE_WEIGHTS_NAME.replace(".safetensors", "{suffix}.safetensors")
+    state_dict_split = split_torch_state_dict_into_shards(
+        state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size
+    )
+
+    # Save index if sharded
+    index = None
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+
+    # Save the model
+    filename_to_tensors = state_dict_split.filename_to_tensors.items()
+
+    for shard_file, tensors in filename_to_tensors:
+        shard = {tensor: state_dict[tensor] for tensor in tensors}
+        safe_save_file(
+            shard, os.path.join(save_path_, shard_file), metadata={"format": "pt"}
+        )
+
+    if index is not None:
+        save_index_file = os.path.join(save_path_, SAFE_WEIGHTS_INDEX_NAME)
+        # Save the index as well
+        with open(save_index_file, "w", encoding="utf-8") as fout:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            fout.write(content)
+
+    return save_path_
+
+
+@send_errors
+def merge_fsdp_weights(
+    checkpoint_dir: str,
+    output_path: str,
+    remove_checkpoint_dir: bool = False,
+):
+    """
+    Merge the weights from sharded FSDP model checkpoints into a single combined checkpoint. Should be used if
+    `SHARDED_STATE_DICT` was used for the model. Weights will be saved to `{output_path}/model.safetensors`.
+
+    Note: this is a CPU-bound process.
+
+    Args:
+        checkpoint_dir (`str`):
+            The directory containing the FSDP checkpoints (can be either the model or optimizer).
+        output_path (`str`):
+            The path to save the merged checkpoint.
+        remove_checkpoint_dir (`bool`, *optional*, defaults to `False`):
+            Whether to remove the checkpoint directory after merging.
+
+    Raises:
+        ValueError: If torch version < 2.3.0, or if `checkpoint_dir` does not exist.
+    """
+    checkpoint_dir_ = Path(checkpoint_dir)
+
+    if not is_torch_version(">=", "2.3.0"):
+        raise ValueError("`merge_fsdp_weights` requires PyTorch >= 2.3.0`")
+
+    # Verify that the checkpoint directory exists
+    if not checkpoint_dir_.exists():
+        model_path_exists = (checkpoint_dir_ / "pytorch_model_fsdp_0").exists()
+        optimizer_path_exists = (checkpoint_dir_ / "optimizer_0").exists()
+        err = f"Tried to load from {checkpoint_dir_} but couldn't find a valid metadata file."
+        if model_path_exists and optimizer_path_exists:
+            err += (
+                " However, potential model and optimizer checkpoint directories exist."
+            )
+            err += f"Please pass in either {checkpoint_dir_}/pytorch_model_fsdp_0 or {checkpoint_dir_}/optimizer_0"
+            err += "instead."
+        elif model_path_exists:
+            err += " However, a potential model checkpoint directory exists."
+            err += (
+                f"Please try passing in {checkpoint_dir_}/pytorch_model_fsdp_0 instead."
+            )
+        elif optimizer_path_exists:
+            err += " However, a potential optimizer checkpoint directory exists."
+            err += f"Please try passing in {checkpoint_dir_}/optimizer_0 instead."
+        raise ValueError(err)
+
+    # To setup `save` to work
+    state = PartialState()
+    if state.is_main_process:
+        LOG.info(f"Merging FSDP weights from {checkpoint_dir_}")
+        save_path = _distributed_checkpoint_to_merged_weights(
+            checkpoint_dir_, output_path
+        )
+        LOG.info(f"Successfully merged FSDP weights and saved to {save_path}")
+        if remove_checkpoint_dir:
+            LOG.info(f"Removing old checkpoint directory {checkpoint_dir_}")
+            shutil.rmtree(checkpoint_dir_)
+
+
+def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
+    """
+    Parses `axolotl` config, CLI args, and calls `merge_fsdp_weights`.
+
+    Args:
+        config: Path to `axolotl` config YAML file.
+        kwargs: Additional keyword arguments to override config file values.
+    """
+
+    parsed_cfg = load_cfg(config, **kwargs)
+
+    fsdp_dir = Path(parsed_cfg.output_dir) / "pytorch_model_fsdp_0"
+    if not fsdp_dir.exists():
+        checkpoint_dir = determine_last_checkpoint(parsed_cfg, update=False)
+        if checkpoint_dir:
+            fsdp_dir = Path(checkpoint_dir) / "pytorch_model_fsdp_0"
+        if not fsdp_dir.exists():
+            raise ValueError(
+                f"Could not find FSDP checkpoint `pytorch_model_fsdp_0` in {checkpoint_dir}"
+            )
+
+    output_path = str(Path(parsed_cfg.output_dir) / "merged")
+    merge_fsdp_weights(
+        checkpoint_dir=str(fsdp_dir),
+        output_path=output_path,
+    )
+    state = PartialState()
+    state.wait_for_everyone()
+    LOG.info(
+        f"FSDP SHARDED_STATE_DICT weights successfully merged to: {output_path}",
+    )
+    LOG.info(
+        "Merged weights are only the safetensors and doesn't include the model configuration "
+        f"or tokenizer which may be found in {parsed_cfg.output_dir}.",
+    )
+
+
+if __name__ == "__main__":
+    fire.Fire(do_cli)
diff --git a/src/axolotl/cli/preprocess.py b/src/axolotl/cli/preprocess.py
index e7b3596a4f..af35dd8010 100644
--- a/src/axolotl/cli/preprocess.py
+++ b/src/axolotl/cli/preprocess.py
@@ -1,62 +1,60 @@
-"""
-CLI to run training on a model
-"""
-import logging
+"""CLI to run preprocessing of a dataset."""
+
+import os
+import warnings
 from pathlib import Path
 from typing import Union
 
 import fire
 import transformers
+from accelerate import init_empty_weights
 from colorama import Fore
+from transformers import AutoModelForCausalLM
 
-from axolotl.cli import (
-    check_accelerate_default_config,
-    check_user_token,
-    load_cfg,
-    load_datasets,
-    load_rl_datasets,
-    print_axolotl_text_art,
-)
-from axolotl.common.cli import PreprocessCliArgs
+from axolotl.cli.args import PreprocessCliArgs
+from axolotl.cli.checks import check_accelerate_default_config, check_user_token
+from axolotl.cli.config import load_cfg
 from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
-from axolotl.prompt_strategies.sharegpt import (
-    register_chatml_template,
-    register_llama3_template,
-)
+from axolotl.common.datasets import load_datasets, load_preference_datasets
+from axolotl.integrations.base import PluginManager
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+from axolotl.utils.trainer import disable_datasets_caching
 
-LOG = logging.getLogger("axolotl.cli.preprocess")
+LOG = get_logger(__name__)
 
 
-def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
-    # pylint: disable=duplicate-code
-    print_axolotl_text_art()
-    parsed_cfg = load_cfg(config, **kwargs)
-    parsed_cfg.is_preprocess = True
+@send_errors
+def do_preprocess(cfg: DictDefault, cli_args: PreprocessCliArgs) -> None:
+    """
+    Preprocesses dataset specified in axolotl config.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        cli_args: Preprocessing-specific CLI arguments.
+    """
     check_accelerate_default_config()
     check_user_token()
-    parser = transformers.HfArgumentParser((PreprocessCliArgs))
-    parsed_cli_args, _ = parser.parse_args_into_dataclasses(
-        return_remaining_strings=True
-    )
 
-    if parsed_cfg.chat_template == "chatml":
-        if parsed_cfg.default_system_message:
-            LOG.info(
-                f"ChatML set. Adding default system message: {parsed_cfg.default_system_message}"
-            )
-            register_chatml_template(parsed_cfg.default_system_message)
-        else:
-            register_chatml_template()
-    elif parsed_cfg.chat_template == "llama3":
-        if parsed_cfg.default_system_message:
-            LOG.info(
-                f"LLaMA-3 set. Adding default system message: {parsed_cfg.default_system_message}"
+    if cli_args.iterable:
+        LOG.error(
+            "The --iterable CLI argument for 'axolotl preprocess' is no longer "
+            "supported. For training, set 'streaming: true' in your YAML config or "
+            "pass '--streaming' in your 'axolotl train' command for on-the-fly "
+            "preprocessing."
+        )
+        return
+
+    for key in ["skip_prepare_dataset", "pretraining_dataset"]:
+        if cfg.get(key):
+            LOG.error(
+                f"You have set `{key}:`. `preprocess` is not needed. Run the 'axolotl "
+                "train' CLI directly instead."
             )
-            register_llama3_template(parsed_cfg.default_system_message)
-        else:
-            register_llama3_template()
+            return
 
-    if not parsed_cfg.dataset_prepared_path:
+    if not cfg.dataset_prepared_path:
         msg = (
             Fore.RED
             + "preprocess CLI called without dataset_prepared_path set, "
@@ -64,19 +62,63 @@ def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
             + Fore.RESET
         )
         LOG.warning(msg)
-        parsed_cfg.dataset_prepared_path = DEFAULT_DATASET_PREPARED_PATH
+        cfg.dataset_prepared_path = DEFAULT_DATASET_PREPARED_PATH
 
-    if parsed_cfg.rl:  # and parsed_cfg.rl != "orpo":
-        load_rl_datasets(cfg=parsed_cfg, cli_args=parsed_cli_args)
-    else:
-        load_datasets(cfg=parsed_cfg, cli_args=parsed_cli_args)
+    with disable_datasets_caching():
+        plugin_manager = PluginManager.get_instance()
+        if plugin_manager.load_datasets(cfg, preprocess=True):
+            pass
+        elif cfg.rl:
+            load_preference_datasets(cfg=cfg, cli_args=cli_args)
+        else:
+            load_datasets(cfg=cfg, cli_args=cli_args)
+
+    if cli_args.download:
+        model_name = cfg.base_model
+        with warnings.catch_warnings():
+            # there are a bunch of useless UserWarnings about
+            # "copying from a non-meta parameter in the checkpoint to a meta parameter in the current model"
+            warnings.simplefilter("ignore")
+            with init_empty_weights(include_buffers=True):
+                # fmt: off
+                try:
+                    AutoModelForCausalLM.from_pretrained(
+                        model_name, trust_remote_code=True
+                    )
+                except Exception:  # nosec B110
+                    pass
+                # fmt: on
 
     LOG.info(
         Fore.GREEN
-        + f"Success! Preprocessed data path: `dataset_prepared_path: {parsed_cfg.dataset_prepared_path}`"
+        + f"Success! Preprocessed data path: `dataset_prepared_path: {cfg.dataset_prepared_path}`"
         + Fore.RESET
     )
 
 
+def do_cli(
+    config: Union[Path, str] = Path("examples/"),
+    **kwargs,
+) -> None:
+    """
+    Parses `axolotl` config, CLI args, and calls `do_preprocess`.
+
+    Args:
+        config: Path to `axolotl` config YAML file.
+        kwargs: Additional keyword arguments to override config file values.
+    """
+
+    os.environ["AXOLOTL_IS_PREPROCESS"] = "1"
+    is_preprocess = kwargs.pop("is_preprocess", True)
+    parsed_cfg = load_cfg(config, is_preprocess=is_preprocess, **kwargs)
+    parsed_cfg.is_preprocess = True
+    parser = transformers.HfArgumentParser(PreprocessCliArgs)
+    parsed_cli_args, _ = parser.parse_args_into_dataclasses(
+        return_remaining_strings=True
+    )
+
+    do_preprocess(parsed_cfg, parsed_cli_args)
+
+
 if __name__ == "__main__":
     fire.Fire(do_cli)
diff --git a/src/axolotl/cli/quantize.py b/src/axolotl/cli/quantize.py
new file mode 100644
index 0000000000..177d04de6d
--- /dev/null
+++ b/src/axolotl/cli/quantize.py
@@ -0,0 +1,128 @@
+"""
+CLI to post-training quantize a model using torchao
+"""
+
+from pathlib import Path
+from typing import Union
+
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from axolotl.cli.config import load_cfg
+from axolotl.loaders import load_processor, load_tokenizer
+from axolotl.utils.logging import get_logger
+from axolotl.utils.quantization import (
+    TorchAOQuantDType,
+    get_quantization_config,
+    quantization_config_to_str,
+    quantize_model,
+    save_quantized_model,
+)
+
+LOG = get_logger(__name__)
+
+
+def do_quantize(
+    config: Union[Path, str],
+    cli_args: dict,
+):
+    """
+    Quantizes a model's model's weights
+
+    Args:
+        config (Union[Path, str]): The path to the config file
+        cli_args (dict): Additional command-line arguments
+    """
+
+    cfg = load_cfg(config)
+
+    if cfg.qat and cfg.quantization:
+        raise ValueError(
+            "QAT and quantization cannot be used together. Please specify only one of qat or quantization in your config file."
+        )
+
+    if cfg.qat:
+        quantize_cfg = cfg.qat
+    elif cfg.quantization:
+        quantize_cfg = cfg.quantization
+    else:
+        raise ValueError(
+            "No quantization configuration found. Please specify either qat or quantization in your config file."
+        )
+
+    model_path = cli_args.get("base_model") or cfg.output_dir
+    if weight_dtype := cli_args.get("weight_dtype"):
+        weight_dtype = TorchAOQuantDType.from_string(weight_dtype)
+    else:
+        weight_dtype = quantize_cfg.weight_dtype
+    if activation_dtype := cli_args.get("activation_dtype"):
+        activation_dtype = TorchAOQuantDType.from_string(activation_dtype)
+    else:
+        activation_dtype = quantize_cfg.activation_dtype
+    group_size = cli_args.get("group_size") or quantize_cfg.group_size
+    quantize_embedding = (
+        cli_args.get("quantize_embedding") or quantize_cfg.quantize_embedding
+    )
+    output_dir = cli_args.get("output_dir") or cfg.output_dir
+    hub_model_id = cli_args.get("hub_model_id") or cfg.hub_model_id
+
+    LOG.info(f"Loading model from {model_path}.")
+    tokenizer = load_tokenizer(cfg)
+
+    processor = None
+    if cfg.is_multimodal:
+        processor = load_processor(cfg, tokenizer)
+
+    config = AutoConfig.from_pretrained(model_path)
+    torch_dtype = config.torch_dtype if hasattr(config, "torch_dtype") else None
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path, device_map="auto", dtype=torch_dtype
+    )
+
+    LOG.info(
+        f"Quantizing model with configuration: \n"
+        f"\tweight_dtype: {weight_dtype}\n"
+        f"\tactivation_dtype: {activation_dtype}\n"
+        f"\tgroup_size: {group_size}\n"
+        f"\tquantize_embedding: {quantize_embedding}"
+    )
+
+    quantize_model(
+        model, weight_dtype, group_size, activation_dtype, quantize_embedding
+    )
+
+    quantization_config = get_quantization_config(
+        weight_dtype, activation_dtype, group_size
+    )
+
+    save_path = str(Path(output_dir) / "quantized")
+    is_mx = getattr(model, "_is_mx_quantized", False)
+    if is_mx:
+        LOG.info(
+            f"Saving MX-quantized model to: {save_path} "
+            "(using torch.save — MXTensor does not support safetensors yet)."
+        )
+    else:
+        LOG.info(f"Saving quantized model to: {save_path}.")
+    save_quantized_model(model, save_path, progressbar=True)
+    tokenizer.save_pretrained(
+        str(Path(output_dir) / "quantized"),
+        progressbar=True,
+        save_jinja_files=cfg.tokenizer_save_jinja_files,
+    )
+
+    if processor:
+        LOG.info(f"Saving processor to: {str(Path(output_dir) / 'quantized')}.")
+        processor.save_pretrained(str(Path(output_dir) / "quantized"))
+
+    if hub_model_id:
+        hub_model_id = (
+            hub_model_id.rstrip("-")
+            + f"-{quantization_config_to_str[type(quantization_config)]}"
+        )
+        model.push_to_hub(hub_model_id)
+        tokenizer.push_to_hub(hub_model_id)
+        if processor:
+            processor.push_to_hub(hub_model_id)
+        LOG.info(f"Quantized model pushed to: {hub_model_id}.")
+
+    LOG.info(f"Quantized model saved to: {str(Path(output_dir) / 'quantized')}.")
diff --git a/src/axolotl/cli/shard.py b/src/axolotl/cli/shard.py
deleted file mode 100644
index 48f22790ac..0000000000
--- a/src/axolotl/cli/shard.py
+++ /dev/null
@@ -1,43 +0,0 @@
-"""
-CLI to shard a trained model into 10GiB chunks
-"""
-import logging
-from pathlib import Path
-from typing import Union
-
-import fire
-import transformers
-
-from axolotl.cli import load_cfg, print_axolotl_text_art
-from axolotl.common.cli import TrainerCliArgs, load_model_and_tokenizer
-from axolotl.utils.dict import DictDefault
-
-LOG = logging.getLogger("axolotl.scripts")
-
-
-def shard(
-    *,
-    cfg: DictDefault,
-    cli_args: TrainerCliArgs,
-):
-    model, _ = load_model_and_tokenizer(cfg=cfg, cli_args=cli_args)
-    safe_serialization = cfg.save_safetensors is True
-    LOG.debug("Re-saving model w/ sharding")
-    model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
-
-
-def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
-    # pylint: disable=duplicate-code
-    print_axolotl_text_art()
-    parsed_cfg = load_cfg(config, **kwargs)
-    parser = transformers.HfArgumentParser((TrainerCliArgs))
-    parsed_cli_args, _ = parser.parse_args_into_dataclasses(
-        return_remaining_strings=True
-    )
-    parsed_cli_args.shard = True
-
-    shard(cfg=parsed_cfg, cli_args=parsed_cli_args)
-
-
-if __name__ == "__main__":
-    fire.Fire(do_cli)
diff --git a/src/axolotl/cli/train.py b/src/axolotl/cli/train.py
index 7bb4a51844..72b88151c9 100644
--- a/src/axolotl/cli/train.py
+++ b/src/axolotl/cli/train.py
@@ -1,69 +1,205 @@
-"""
-CLI to run training on a model
-"""
-import logging
+"""CLI to run training on a model."""
+
+import gc
+import os
 from pathlib import Path
-from typing import Tuple, Union
+from typing import Any, Union
 
 import fire
-from transformers.hf_argparser import HfArgumentParser
-from transformers.modeling_utils import PreTrainedModel
-from transformers.tokenization_utils import PreTrainedTokenizer
-
-from axolotl.cli import (
-    check_accelerate_default_config,
-    check_user_token,
-    load_cfg,
-    load_datasets,
-    load_rl_datasets,
-    print_axolotl_text_art,
-)
-from axolotl.common.cli import TrainerCliArgs
-from axolotl.prompt_strategies.sharegpt import (
-    register_chatml_template,
-    register_llama3_template,
-)
-from axolotl.train import train
-
-LOG = logging.getLogger("axolotl.cli.train")
+
+from axolotl.cli.args import TrainerCliArgs
+from axolotl.utils import make_lazy_getattr
+from axolotl.utils.dict import DictDefault
+
+_LAZY_IMPORTS = {
+    "Accelerator": "accelerate",
+    "HfArgumentParser": "transformers.hf_argparser",
+    "gpu_capabilities": "axolotl.cli.config",
+    "load_cfg": "axolotl.cli.config",
+    "load_datasets": "axolotl.common.datasets",
+    "load_preference_datasets": "axolotl.common.datasets",
+    "normalize_config": "axolotl.utils.config",
+    "plugin_set_cfg": "axolotl.cli.config",
+    "prepare_optim_env": "axolotl.utils.trainer",
+    "prepare_plugins": "axolotl.cli.config",
+    "resolve_dtype": "axolotl.utils.config",
+    "train": "axolotl.train",
+    "validate_config": "axolotl.utils.config",
+}
+
+__getattr__ = make_lazy_getattr(_LAZY_IMPORTS, __name__, globals())
+
+
+def _lazy_attr(name: str) -> Any:
+    return globals().get(name) or __getattr__(name)
+
+
+def do_train(cfg: DictDefault, cli_args: TrainerCliArgs):
+    """
+    Trains a `transformers` model by first loading the dataset(s) specified in the
+    `axolotl` config, and then calling `axolotl.train.train`. Also runs the plugin
+    manager's `post_train_unload` once training completes.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        cli_args: Training-specific CLI arguments.
+    """
+    from axolotl.cli.checks import check_accelerate_default_config, check_user_token
+
+    check_accelerate_default_config()
+    if int(os.getenv("LOCAL_RANK", "0")) == 0:
+        check_user_token()
+
+    load_datasets_fn: Any = globals().get("load_datasets")
+    load_preference_datasets_fn: Any = globals().get("load_preference_datasets")
+    train_fn: Any = globals().get("train")
+    if load_datasets_fn is None or load_preference_datasets_fn is None:
+        from axolotl.common import datasets as datasets_module
+
+        load_datasets_fn = load_datasets_fn or datasets_module.load_datasets
+        load_preference_datasets_fn = (
+            load_preference_datasets_fn or datasets_module.load_preference_datasets
+        )
+    if train_fn is None:
+        from axolotl.train import train as train_fn
+
+    dataset_meta = None
+    if cfg.get("plugins"):
+        from axolotl.integrations.base import PluginManager
+
+        plugin_manager = PluginManager.get_instance()
+        dataset_meta = plugin_manager.load_datasets(cfg, preprocess=False)
+
+    if dataset_meta is None:
+        if cfg.rl:
+            dataset_meta = load_preference_datasets_fn(cfg=cfg, cli_args=cli_args)
+        else:
+            dataset_meta = load_datasets_fn(cfg=cfg, cli_args=cli_args)
+
+    model, tokenizer, trainer = train_fn(cfg=cfg, dataset_meta=dataset_meta)
+
+    del model, tokenizer, trainer
+
+    gc.collect()
+
+    if cfg.get("plugins"):
+        from axolotl.integrations.base import PluginManager
+
+        plugin_manager = PluginManager.get_instance()
+        plugin_manager.post_train_unload(cfg)
 
 
 def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
-    # pylint: disable=duplicate-code
-    parsed_cfg = load_cfg(config, **kwargs)
-    parser = HfArgumentParser((TrainerCliArgs))
+    """
+    Parses `axolotl` config, CLI args, and calls `do_train`.
+
+    Args:
+        config: Path to `axolotl` config YAML file.
+        kwargs: Additional keyword arguments to override config file values.
+    """
+    parser_cls: Any = globals().get("HfArgumentParser")
+    load_cfg_fn: Any = globals().get("load_cfg")
+    if parser_cls is None:
+        from transformers.hf_argparser import HfArgumentParser
+
+        parser_cls = HfArgumentParser
+    if load_cfg_fn is None:
+        from axolotl.cli.config import load_cfg
+
+        load_cfg_fn = load_cfg
+
+    parsed_cfg = load_cfg_fn(config, **kwargs)
+    parser = parser_cls(TrainerCliArgs)
     parsed_cli_args, _ = parser.parse_args_into_dataclasses(
         return_remaining_strings=True
     )
-    return do_train(parsed_cfg, parsed_cli_args)
 
+    if parsed_cfg.use_ray:
+        from ray.train import RunConfig, ScalingConfig
+        from ray.train.torch import TorchTrainer
 
-def do_train(cfg, cli_args) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
-    print_axolotl_text_art()
-    check_accelerate_default_config()
-    check_user_token()
-    if cfg.chat_template == "chatml" and cfg.default_system_message:
-        LOG.info(
-            f"ChatML set. Adding default system message: {cfg.default_system_message}"
+        train_loop_config = {"cfg": parsed_cfg.to_dict(), "cli_args": parsed_cli_args}
+        trainer = TorchTrainer(
+            ray_train_func,
+            train_loop_config=train_loop_config,
+            scaling_config=ScalingConfig(
+                num_workers=parsed_cfg.ray_num_workers,
+                resources_per_worker=parsed_cfg.resources_per_worker.to_dict(),
+                use_gpu=True,
+            ),
+            run_config=RunConfig(
+                name=parsed_cfg.ray_run_name,
+                storage_path=Path(parsed_cfg.output_dir).absolute().as_posix(),
+            ),
         )
-        register_chatml_template(cfg.default_system_message)
-    else:
-        register_chatml_template()
 
-    if cfg.chat_template == "llama3" and cfg.default_system_message:
-        LOG.info(
-            f"LLaMA-3 set. Adding default system message: {cfg.default_system_message}"
-        )
-        register_llama3_template(cfg.default_system_message)
-    else:
-        register_llama3_template()
+        trainer.fit()
+        return
+
+    do_train(parsed_cfg, parsed_cli_args)
+
+
+def ray_train_func(kwargs: dict):
+    """Ray Train entrypoint executed on each GPU worker.
+
+    Re-validates the config against worker-local GPU capabilities (deferred from
+    the driver), then runs the standard training pipeline.
+    """
+    # cast `cfg` back to DictDefault (ray tune deepcopy has issues with DictDefault so needed it to be dict)
+    # also renormalize the config now that TorchTrainer has spawned distributed workers
+    cfg = DictDefault(kwargs["cfg"])
+
+    # Plugins must be registered before `validate_config` so the plugin-extended
+    # pydantic schema is in scope on this worker; otherwise plugin-specific cfg
+    # fields are silently dropped by `model_dump(exclude_none=True)`.
+    if cfg.get("plugins"):
+        prepare_plugins_fn: Any = _lazy_attr("prepare_plugins")
+        prepare_plugins_fn(cfg)
+
+    # GPU capability detection was deferred from the driver; run the checks now
+    # that we are on a worker that actually has the training device attached.
+    accelerator_cls: Any = _lazy_attr("Accelerator")
+    gpu_capabilities_fn: Any = _lazy_attr("gpu_capabilities")
+    normalize_config_fn: Any = _lazy_attr("normalize_config")
+    prepare_optim_env_fn: Any = _lazy_attr("prepare_optim_env")
+    resolve_dtype_fn: Any = _lazy_attr("resolve_dtype")
+    validate_config_fn: Any = _lazy_attr("validate_config")
+
+    capabilities, env_capabilities = gpu_capabilities_fn()
+    cfg = validate_config_fn(
+        cfg,
+        capabilities=capabilities,
+        env_capabilities=env_capabilities,
+    )
+
+    # Derive here (not in controller normalize_config) so the worker's
+    # validate_config above doesn't see both set and trip check_gas_bsz.
+    cfg.gradient_accumulation_steps = cfg.gradient_accumulation_steps or (
+        cfg.batch_size // cfg.micro_batch_size
+    )
+    cfg.batch_size = (
+        cfg.batch_size or cfg.micro_batch_size * cfg.gradient_accumulation_steps
+    )
+
+    prepare_optim_env_fn(cfg)
+    normalize_config_fn(cfg)
+    resolve_dtype_fn(cfg)
+
+    # ray serializing objects gets rid of frozen attribute - HF expects dict not DefaultDict
+    if cfg.deepspeed and hasattr(cfg.deepspeed, "to_dict"):
+        cfg.deepspeed = cfg.deepspeed.to_dict()
+
+    # initialize accelerator before model instantiation
+    accelerator_cls(gradient_accumulation_steps=cfg.gradient_accumulation_steps)
+
+    # Bind the post-validation cfg to the plugin manager.
+    if cfg.get("plugins"):
+        plugin_set_cfg_fn: Any = _lazy_attr("plugin_set_cfg")
+        plugin_set_cfg_fn(cfg)
 
-    if cfg.rl:  # and cfg.rl != "orpo":
-        dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
-    else:
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+    kwargs["cfg"] = cfg
 
-    return train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+    do_train(**kwargs)
 
 
 if __name__ == "__main__":
diff --git a/src/axolotl/cli/utils/__init__.py b/src/axolotl/cli/utils/__init__.py
new file mode 100644
index 0000000000..cb1b4f0911
--- /dev/null
+++ b/src/axolotl/cli/utils/__init__.py
@@ -0,0 +1,34 @@
+"""Init for axolotl.cli.utils module."""
+
+from axolotl.utils import make_lazy_getattr
+
+from .args import (
+    add_options_from_config,
+    add_options_from_config_options,
+    add_options_from_dataclass,
+    filter_none_kwargs,
+)
+from .fetch import fetch_from_github
+from .sweeps import generate_sweep_configs
+from .train import build_command, generate_config_files, launch_training
+
+__all__ = [
+    "filter_none_kwargs",
+    "add_options_from_dataclass",
+    "add_options_from_config",
+    "add_options_from_config_options",
+    "build_command",
+    "generate_config_files",
+    "generate_sweep_configs",
+    "load_model_and_tokenizer",
+    "resolve_chat_template_str",
+    "launch_training",
+    "fetch_from_github",
+]
+
+_LAZY_IMPORTS = {
+    "load_model_and_tokenizer": ".load",
+    "resolve_chat_template_str": ".load",
+}
+
+__getattr__ = make_lazy_getattr(_LAZY_IMPORTS, __name__, globals())
diff --git a/src/axolotl/cli/utils/args.py b/src/axolotl/cli/utils/args.py
new file mode 100644
index 0000000000..8b9fa8b023
--- /dev/null
+++ b/src/axolotl/cli/utils/args.py
@@ -0,0 +1,219 @@
+"""Utilities for axolotl CLI args."""
+
+import dataclasses
+from functools import wraps
+from types import NoneType, UnionType
+from typing import Any, Callable, Type, Union, get_args, get_origin
+
+import click
+from pydantic import BaseModel
+
+ConfigOption = tuple[tuple[str, ...], str | None, str | None, str | None]
+
+
+def _strip_optional_type(field_type: type | str | None):
+    """
+    Extracts the non-`None` type from an `Optional` / `Union` type.
+
+    Args:
+        field_type: Type of field for Axolotl CLI command.
+
+    Returns:
+        If the input type is `Union[T, None]` or `Optional[T]`, returns `T`. Otherwise
+            returns the input type unchanged.
+    """
+    is_union = get_origin(field_type) is Union or isinstance(field_type, UnionType)
+    if is_union and type(None) in get_args(field_type):
+        field_type = next(
+            t for t in get_args(field_type) if not isinstance(t, NoneType)
+        )
+
+    return field_type
+
+
+def filter_none_kwargs(func: Callable) -> Callable:
+    """
+    Wraps function to remove `None`-valued `kwargs`.
+
+    Args:
+        func: Function to wrap.
+
+    Returns:
+        Wrapped function.
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs) -> Callable:
+        """Filters out `None`-valued `kwargs`."""
+        filtered_kwargs = {k: v for k, v in kwargs.items() if v is not None}
+
+        return func(*args, **filtered_kwargs)
+
+    return wrapper
+
+
+def add_options_from_dataclass(config_class: Type[Any]) -> Callable:
+    """
+    Create Click options from the fields of a dataclass.
+
+    Args:
+        config_class: Dataclass with fields to parse from the CLI.
+
+    Returns:
+        Function decorator for Axolotl CLI command.
+    """
+
+    def decorator(function: Callable) -> Callable:
+        # Process dataclass fields in reverse order for correct option ordering
+        for field in reversed(dataclasses.fields(config_class)):
+            field_type = _strip_optional_type(field.type)
+
+            if field_type is bool:
+                field_name = field.name.replace("_", "-")
+                option_name = f"--{field_name}/--no-{field_name}"
+                function = click.option(
+                    option_name,
+                    default=field.default,
+                    help=field.metadata.get("description"),
+                )(function)
+            else:
+                option_name = f"--{field.name.replace('_', '-')}"
+                function = click.option(
+                    option_name,
+                    type=field_type,
+                    default=field.default,
+                    help=field.metadata.get("description"),
+                )(function)
+
+        return function
+
+    return decorator
+
+
+def _is_pydantic_model(field_type: type) -> bool:
+    """Check if a type is a Pydantic BaseModel subclass."""
+    try:
+        return isinstance(field_type, type) and issubclass(field_type, BaseModel)
+    except TypeError:
+        return False
+
+
+def _get_field_description(field) -> str | None:
+    """Get description from a Pydantic field, checking both .description and json_schema_extra."""
+    if field.description:
+        return field.description
+    if field.json_schema_extra and isinstance(field.json_schema_extra, dict):
+        return field.json_schema_extra.get("description")
+    return None
+
+
+def _add_nested_model_options(
+    function: Callable, parent_name: str, model_class: Type[BaseModel]
+) -> Callable:
+    """
+    Add Click options for all fields of a nested Pydantic model using dot-notation.
+
+    Note: Only single-level nesting is supported (e.g., ``--trl.beta``).
+    Deeper nesting (e.g., ``--trl.scheduler.warmup``) is not handled.
+
+    Args:
+        function: Click command function to add options to.
+        parent_name: Parent field name (e.g., "trl").
+        model_class: Nested Pydantic model class.
+
+    Returns:
+        Function with added Click options.
+    """
+    for sub_name, sub_field in reversed(model_class.model_fields.items()):
+        sub_type = _strip_optional_type(sub_field.annotation)
+        # Use dot notation: --parent.sub_field
+        cli_name = f"{parent_name}.{sub_name}".replace("_", "-")
+        # The kwarg name uses double-underscore as separator
+        param_name = f"{parent_name}__{sub_name}"
+        description = _get_field_description(sub_field)
+
+        if sub_type is bool:
+            option_name = f"--{cli_name}/--no-{cli_name}"
+            function = click.option(
+                option_name, param_name, default=None, help=description
+            )(function)
+        else:
+            option_name = f"--{cli_name}"
+            click_type = {str: str, int: int, float: float}.get(sub_type)
+            function = click.option(
+                option_name, param_name, default=None, type=click_type, help=description
+            )(function)
+
+    return function
+
+
+def add_options_from_config(config_class: Type[BaseModel]) -> Callable:
+    """
+    Create Click options from the fields of a Pydantic model.
+
+    For fields whose type is itself a Pydantic BaseModel, dot-notation CLI options are
+    generated for each sub-field (e.g., ``--trl.beta=0.1``).
+
+    Args:
+        config_class: PyDantic model with fields to parse from the CLI
+
+    Returns:
+        Function decorator for Axolotl CLI command.
+    """
+
+    def decorator(function: Callable) -> Callable:
+        # Process model fields in reverse order for correct option ordering
+        for name, field in reversed(config_class.model_fields.items()):
+            field_type = _strip_optional_type(field.annotation)
+
+            # Handle nested Pydantic models with dot-notation options
+            if _is_pydantic_model(field_type):
+                function = _add_nested_model_options(function, name, field_type)
+                continue
+
+            if field_type is bool:
+                field_name = name.replace("_", "-")
+                option_name = f"--{field_name}/--no-{field_name}"
+                function = click.option(
+                    option_name, default=None, help=field.description
+                )(function)
+            else:
+                option_name = f"--{name.replace('_', '-')}"
+                function = click.option(
+                    option_name, default=None, help=field.description
+                )(function)
+
+        return function
+
+    return decorator
+
+
+def add_options_from_config_options(options: tuple[ConfigOption, ...]) -> Callable:
+    """
+    Create Click options from precomputed config option metadata.
+
+    This mirrors :func:`add_options_from_config` without importing the full Pydantic
+    config schema at CLI startup.
+    """
+
+    click_types = {"str": str, "int": int, "float": float}
+
+    def decorator(function: Callable) -> Callable:
+        for param_decls, param_name, click_type_name, description in reversed(options):
+            option_args: tuple[str, ...]
+            if param_name:
+                option_args = (*param_decls, param_name)
+            else:
+                option_args = param_decls
+            click_type = click_types[click_type_name] if click_type_name else None
+
+            function = click.option(
+                *option_args,
+                default=None,
+                type=click_type,
+                help=description,
+            )(function)
+
+        return function
+
+    return decorator
diff --git a/src/axolotl/cli/utils/diffusion.py b/src/axolotl/cli/utils/diffusion.py
new file mode 100644
index 0000000000..7bf68048eb
--- /dev/null
+++ b/src/axolotl/cli/utils/diffusion.py
@@ -0,0 +1,374 @@
+"""Helpers for diffusion-mode inference in CLI and Gradio."""
+
+from __future__ import annotations
+
+import gradio as gr
+from colorama import Fore, Style
+
+from axolotl.integrations.diffusion import generate, resolve_mask_token_id
+from axolotl.utils.dict import DictDefault
+
+
+def diffusion_inference(
+    model,
+    tokenizer,
+    cfg,
+    prompt: str,
+    chat_template_str: str | None = None,
+):
+    """Diffusion inference helper method."""
+    mode = "random"
+    completion_tokens = 0
+    target_mask_ratio = None
+    mode, completion_tokens, target_mask_ratio, cleaned = _parse_commands(prompt)
+
+    if cleaned:
+        prompt = cleaned
+
+    info = run_diffusion(
+        model=model,
+        tokenizer=tokenizer,
+        cfg=cfg,
+        prompt=prompt,
+        chat_template_str=chat_template_str,
+        mode=mode,
+        target_mask_ratio=target_mask_ratio,
+        completion_tokens=completion_tokens,
+    )
+    masked_text = info["masked_text"]
+    mask_ratio = info["mask_ratio"]
+    generated_ids = info["generated_ids"]
+    masked_positions = info["masked_positions"]
+    orig_ids = info["orig_ids"]
+
+    # Display with masked preview and colored diff
+    if masked_text is not None and mask_ratio is not None:
+        print(f"Masked ({mask_ratio:.1%}):\n{masked_text}\n")
+    if generated_ids is not None:
+        # Compute per-token style
+        styles: list[str] = []
+        for i, tid in enumerate(generated_ids):
+            if i in masked_positions:
+                if i < len(orig_ids) and tid == orig_ids[i]:
+                    styles.append("green")  # correct fill
+                elif i < len(orig_ids):
+                    styles.append("red")  # incorrect fill
+                else:
+                    styles.append("normal")  # appended
+            else:
+                same = i < len(orig_ids) and tid == orig_ids[i]
+                styles.append("dim" if same else "normal")
+
+        # Group contiguous spans by style
+        styled_spans: list[tuple[str, int, int]] = []
+        if generated_ids:
+            current_style = styles[0]
+            start = 0
+            for i in range(1, len(generated_ids)):
+                s = styles[i]
+                if s != current_style:
+                    styled_spans.append((current_style, start, i))
+                    current_style, start = s, i
+            styled_spans.append((current_style, start, len(generated_ids)))
+
+        out_parts = []
+        for style_name, a, b in styled_spans:
+            chunk_text = tokenizer.decode(generated_ids[a:b], skip_special_tokens=False)
+            if style_name == "green":
+                out_parts.append(Fore.GREEN + chunk_text + Style.RESET_ALL)
+            elif style_name == "red":
+                out_parts.append(Fore.RED + chunk_text + Style.RESET_ALL)
+            else:
+                if style_name == "dim":
+                    out_parts.append(Style.DIM + chunk_text + Style.RESET_ALL)
+                else:
+                    out_parts.append(chunk_text)
+        print("Generated:\n" + "".join(out_parts))
+    else:
+        print("Generated:\n(no output)")
+
+
+def _parse_commands(text: str):
+    """
+    Parse leading diffusion commands.
+
+    Supported at start of input (can be chained):
+      :complete N  -> completion mode with N tokens (default 64)
+      :mask R      -> random masking with ratio R in [0, 1]
+    """
+    tokens = text.strip().split()
+    i = 0
+    mode = "random"
+    completion_tokens = 0
+    target_mask_ratio = None
+    consumed = 0
+    while i < len(tokens) and tokens[i].startswith(":"):
+        cmd = tokens[i]
+        i += 1
+        consumed = i
+        if cmd == ":complete":
+            mode = "completion"
+            if i < len(tokens):
+                try:
+                    completion_tokens = int(tokens[i])
+                    i += 1
+                    consumed = i
+                except Exception:
+                    completion_tokens = 64
+            else:
+                completion_tokens = 64
+        elif cmd == ":mask":
+            mode = "random"
+            if i < len(tokens):
+                try:
+                    target_mask_ratio = float(tokens[i])
+                    i += 1
+                    consumed = i
+                except Exception:
+                    target_mask_ratio = None
+        else:
+            i -= 1
+            consumed = i
+            break
+
+    cleaned = " ".join(tokens[consumed:])
+
+    return mode, completion_tokens, target_mask_ratio, cleaned
+
+
+def run_diffusion(
+    *,
+    model,
+    tokenizer,
+    cfg: DictDefault,
+    prompt: str,
+    chat_template_str: str | None,
+    mode: str = "random",
+    target_mask_ratio: float | None = None,
+    completion_tokens: int = 0,
+):
+    """Run a single diffusion generation and return a structured result dict."""
+    if chat_template_str:
+        batch = tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            return_tensors="pt",
+            add_special_tokens=True,
+            add_generation_prompt=True,
+            chat_template=chat_template_str,
+            tokenize=True,
+            return_dict=True,
+        )
+    else:
+        batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
+
+    mask_token_id = resolve_mask_token_id(tokenizer, cfg, allow_add=False)
+
+    seq = batch["input_ids"].to(cfg.device)
+    gen_mode = "completion" if mode == "completion" else "random"
+    comp_tokens = int(completion_tokens) if gen_mode == "completion" else 0
+
+    result = generate(
+        model,
+        tokenizer,
+        original_sequence=seq[:1],
+        num_diffusion_steps=cfg.diffusion.num_diffusion_steps,
+        temperature=cfg.diffusion.generation_temperature,
+        mask_token_id=int(mask_token_id),
+        mode=gen_mode,  # type: ignore[arg-type]
+        completion_tokens=comp_tokens,
+        target_mask_ratio=target_mask_ratio,
+    )
+
+    masked_text = result.get("masked") if isinstance(result, dict) else None
+    mask_ratio = result.get("mask_ratio") if isinstance(result, dict) else None
+    generated_ids = result.get("generated_ids") if isinstance(result, dict) else None
+    masked_positions = (
+        set(result.get("masked_positions") or []) if isinstance(result, dict) else set()
+    )
+    orig_ids = seq[0].detach().cpu().tolist()
+
+    return {
+        "masked_text": masked_text,
+        "mask_ratio": mask_ratio,
+        "generated_ids": generated_ids,
+        "masked_positions": masked_positions,
+        "orig_ids": orig_ids,
+    }
+
+
+def render_html(
+    *,
+    generated_ids: list[int] | None,
+    orig_ids: list[int],
+    masked_positions: set[int],
+    tokenizer,
+) -> str:
+    """Render HTML visualizing diffusion outputs."""
+    if not generated_ids:
+        return "<pre>Generated:\n(no output)</pre>"
+
+    def _style_for(i: int, tid: int) -> str:
+        if i in masked_positions:
+            if i < len(orig_ids) and tid == orig_ids[i]:
+                return "green"
+            if i < len(orig_ids):
+                return "red"
+            return "normal"
+        same = i < len(orig_ids) and tid == orig_ids[i]
+        return "dim" if same else "normal"
+
+    # Group contiguous spans by style to reduce HTML size
+    spans: list[tuple[str, int, int]] = []
+    if generated_ids:
+        cur = _style_for(0, generated_ids[0])
+        start = 0
+        for i in range(1, len(generated_ids)):
+            s = _style_for(i, generated_ids[i])
+            if s != cur:
+                spans.append((cur, start, i))
+                cur, start = s, i
+        spans.append((cur, start, len(generated_ids)))
+
+    html_parts = []
+    for style_name, a, b in spans:
+        txt = tokenizer.decode(generated_ids[a:b], skip_special_tokens=False)
+        if style_name == "green":
+            html_parts.append(f'<span style="color:#2e7d32">{txt}</span>')
+        elif style_name == "red":
+            html_parts.append(f'<span style="color:#c62828">{txt}</span>')
+        elif style_name == "dim":
+            html_parts.append(f'<span style="opacity:0.6">{txt}</span>')
+        else:
+            html_parts.append(txt)
+
+    legend = (
+        '<div style="font-size:0.9em;margin-bottom:4px">'
+        '<span style="color:#2e7d32">correct</span>, '
+        '<span style="color:#c62828">incorrect</span>, '
+        '<span style="opacity:0.6">unchanged</span>'
+        "</div>"
+    )
+
+    return (
+        legend
+        + '<pre style="white-space:pre-wrap">Generated:\n'
+        + "".join(html_parts)
+        + "</pre>"
+    )
+
+
+def launch_diffusion_gradio_ui(
+    *,
+    model,
+    tokenizer,
+    cfg: DictDefault,
+    prompter_module=None,
+    chat_template_str: str | None = None,
+):
+    """Build and launch a simple Gradio UI for diffusion inference."""
+    with gr.Blocks(
+        title=cfg.get("gradio_title", "Axolotl Diffusion Interface")
+    ) as demo:
+        gr.Markdown(
+            """
+            ## Axolotl Diffusion Inference
+            - Mode "Random" masks tokens at a target ratio and fills them.
+            - Mode "Completion" appends N masked tokens at the end and fills them.
+            """
+        )
+
+        with gr.Row():
+            mode = gr.Radio(
+                choices=["random", "completion"],
+                value="random",
+                label="Mode",
+            )
+            mask_ratio = gr.Slider(
+                minimum=0.0,
+                maximum=1.0,
+                step=0.05,
+                value=0.4,
+                label="Mask ratio (random mode)",
+                interactive=True,
+            )
+            completion_tokens = gr.Number(
+                value=64,
+                precision=0,
+                label="Completion tokens (completion mode)",
+                interactive=True,
+                visible=False,
+            )
+
+        instruction = gr.Textbox(label="Instruction", lines=6)
+        run_btn = gr.Button("Generate")
+
+        masked_preview = gr.Textbox(label="Masked preview", lines=6)
+        html_out = gr.HTML(label="Generated")
+
+        def _toggle_controls(selected_mode: str):
+            return (
+                gr.update(visible=(selected_mode == "random")),
+                gr.update(visible=(selected_mode == "completion")),
+            )
+
+        mode.change(
+            _toggle_controls,
+            inputs=[mode],
+            outputs=[mask_ratio, completion_tokens],
+        )
+
+        def _gen(instruction_text: str, selected_mode: str, mratio: float, ctoks: int):
+            if not instruction_text:
+                return "", "<pre>Generated:\n(no output)</pre>"
+
+            if prompter_module:
+                prompt: str = next(
+                    prompter_module().build_prompt(
+                        instruction=instruction_text.strip("\n")
+                    )
+                )
+            else:
+                prompt = instruction_text.strip()
+
+            info = run_diffusion(
+                model=model,
+                tokenizer=tokenizer,
+                cfg=cfg,
+                prompt=prompt,
+                chat_template_str=chat_template_str,
+                mode=selected_mode,
+                target_mask_ratio=mratio if selected_mode == "random" else None,
+                completion_tokens=int(ctoks) if selected_mode == "completion" else 0,
+            )
+
+            masked_text = info.get("masked_text")
+            mask_ratio_val = info.get("mask_ratio")
+            generated_ids = info.get("generated_ids")
+            masked_positions = info.get("masked_positions") or set()
+            orig_ids = info.get("orig_ids") or []
+
+            preview = (
+                f"Masked ({mask_ratio_val:.1%}):\n{masked_text}"
+                if masked_text is not None and mask_ratio_val is not None
+                else ""
+            )
+            html = render_html(
+                generated_ids=generated_ids,
+                orig_ids=orig_ids,
+                masked_positions=masked_positions,
+                tokenizer=tokenizer,
+            )
+            return preview, html
+
+        run_btn.click(
+            _gen,
+            inputs=[instruction, mode, mask_ratio, completion_tokens],
+            outputs=[masked_preview, html_out],
+        )
+
+        demo.launch(
+            footer_links=["gradio", "settings"],
+            share=cfg.get("gradio_share", True),
+            server_name=cfg.get("gradio_server_name", "127.0.0.1"),
+            server_port=cfg.get("gradio_server_port", None),
+        )
diff --git a/src/axolotl/cli/utils/fetch.py b/src/axolotl/cli/utils/fetch.py
new file mode 100644
index 0000000000..441b7f6f74
--- /dev/null
+++ b/src/axolotl/cli/utils/fetch.py
@@ -0,0 +1,142 @@
+"""Utilities for axolotl fetch CLI command."""
+
+import concurrent.futures
+import hashlib
+import json
+from pathlib import Path
+
+import click
+import requests
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _download_file(
+    file_info: tuple, raw_base_url: str, dest_path: Path, dir_prefix: str
+) -> tuple[str, str]:
+    """
+    Download a single file and return its processing status.
+
+    Args:
+        file_info: Tuple of (file_path, remote_sha).
+        raw_base_url: Base URL for raw GitHub content.
+        dest_path: Local destination directory.
+        dir_prefix: Directory prefix to filter files.
+
+    Returns:
+        Tuple of (file_path, status) where status is 'new', 'updated', or 'unchanged'.
+    """
+    file_path, remote_sha = file_info
+    raw_url = f"{raw_base_url}/{file_path}"
+    dest_file = dest_path / file_path.split(dir_prefix)[-1]
+
+    # Check if file exists and needs updating
+    if dest_file.exists():
+        with open(dest_file, "rb") as file:
+            content = file.read()
+            # Calculate git blob SHA
+            blob = b"blob " + str(len(content)).encode() + b"\0" + content
+            local_sha = hashlib.sha1(blob, usedforsecurity=False).hexdigest()
+
+        if local_sha == remote_sha:
+            print(f"Skipping {file_path} (unchanged)")
+            return file_path, "unchanged"
+
+        print(f"Updating {file_path}")
+        status = "updated"
+    else:
+        print(f"Downloading {file_path}")
+        status = "new"
+
+    # Create directories if needed
+    dest_file.parent.mkdir(parents=True, exist_ok=True)
+
+    # Download and save file
+    try:
+        response = requests.get(raw_url, timeout=30)
+        response.raise_for_status()
+
+        with open(dest_file, "wb") as file:
+            file.write(response.content)
+
+        return file_path, status
+    except (requests.RequestException, IOError) as request_error:
+        print(f"Error downloading {file_path}: {str(request_error)}")
+        return file_path, "error"
+
+
+def fetch_from_github(
+    dir_prefix: str, dest_dir: str | None = None, max_workers: int = 5
+) -> None:
+    """
+    Sync files from a specific directory in the GitHub repository.
+    Only downloads files that don't exist locally or have changed.
+
+    Args:
+        dir_prefix: Directory prefix to filter files (e.g., 'examples/',
+            'deepspeed_configs/').
+        dest_dir: Local destination directory.
+        max_workers: Maximum number of concurrent downloads.
+    """
+    api_url = "https://api.github.com/repos/axolotl-ai-cloud/axolotl/git/trees/main?recursive=1"
+    raw_base_url = "https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main"
+
+    # Get repository tree with timeout
+    response = requests.get(api_url, timeout=30)
+    response.raise_for_status()
+    tree = json.loads(response.text)
+
+    # Filter for files and get their SHA
+    files = {
+        item["path"]: item["sha"]
+        for item in tree["tree"]
+        if item["type"] == "blob" and item["path"].startswith(dir_prefix)
+    }
+
+    if not files:
+        raise click.ClickException(f"No files found in {dir_prefix}")
+
+    # Default destination directory is the last part of dir_prefix
+    default_dest = Path(dir_prefix.rstrip("/"))
+    dest_path = Path(dest_dir) if dest_dir else default_dest
+
+    # Keep track of processed files for summary
+    files_processed: dict[str, list[str]] = {
+        "new": [],
+        "updated": [],
+        "unchanged": [],
+        "error": [],
+    }
+
+    # Process files in parallel using ThreadPoolExecutor
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        future_to_file = {
+            executor.submit(
+                _download_file,
+                (file_path, remote_sha),
+                raw_base_url,
+                dest_path,
+                dir_prefix,
+            ): file_path
+            for file_path, remote_sha in files.items()
+        }
+
+        # Process completed tasks as they finish
+        for future in concurrent.futures.as_completed(future_to_file):
+            file_path = future_to_file[future]
+            try:
+                file_path, status = future.result()
+                files_processed[status].append(file_path)
+            except (requests.RequestException, IOError) as request_error:
+                print(f"Error processing {file_path}: {str(request_error)}")
+                files_processed["error"].append(file_path)
+
+    # Log summary
+    LOG.info("\nSync Summary:")
+    LOG.info(f"New files: {len(files_processed['new'])}")
+    LOG.info(f"Updated files: {len(files_processed['updated'])}")
+    LOG.info(f"Unchanged files: {len(files_processed['unchanged'])}")
+    if files_processed["error"]:
+        LOG.info(f"Failed files: {len(files_processed['error'])}")
diff --git a/src/axolotl/cli/utils/load.py b/src/axolotl/cli/utils/load.py
new file mode 100644
index 0000000000..4fb6f06f32
--- /dev/null
+++ b/src/axolotl/cli/utils/load.py
@@ -0,0 +1,79 @@
+"""Utilities for model, tokenizer, etc. loading."""
+
+from typing import Any
+
+from transformers import (
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+    ProcessorMixin,
+)
+
+from axolotl.loaders import load_processor, load_tokenizer
+from axolotl.loaders.model import ModelLoader
+from axolotl.utils.chat_templates import get_chat_template_from_config
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def resolve_chat_template_str(
+    cfg: DictDefault,
+    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast | Any,
+) -> str | None:
+    """
+    Resolves the chat template string for inference from the `axolotl` config,
+    mirroring how it would be resolved at training time: an explicit
+    `chat_template` config takes precedence, then the first dataset's
+    `chat_template` if that dataset is of type `chat_template`.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        tokenizer: Tokenizer to fall back to for tokenizer-default templates.
+
+    Returns:
+        Chat template string, or None if the config does not specify one.
+    """
+    if cfg.chat_template:
+        return get_chat_template_from_config(cfg, ds_cfg=None, tokenizer=tokenizer)
+    if cfg.datasets and cfg.datasets[0].type == "chat_template":
+        return get_chat_template_from_config(
+            cfg=cfg, ds_cfg=cfg.datasets[0], tokenizer=tokenizer
+        )
+    return None
+
+
+def load_model_and_tokenizer(
+    *,
+    cfg: DictDefault,
+    inference: bool = False,
+) -> tuple[
+    PreTrainedModel,
+    PreTrainedTokenizer | PreTrainedTokenizerFast | Any,
+    ProcessorMixin | None,
+]:
+    """
+    Helper function for loading a model, tokenizer, and processor specified in the
+    given `axolotl` config.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        inference: Boolean denoting inference mode.
+
+    Returns:
+        Tuple of (PreTrainedModel, PreTrainedTokenizer, ProcessorMixin).
+    """
+    LOG.info(f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}")
+    tokenizer = load_tokenizer(cfg)
+
+    LOG.info("loading model...")
+    model_loader = ModelLoader(cfg, tokenizer, inference=inference)
+    model, _ = model_loader.load()
+
+    processor = None
+    if cfg.is_multimodal:
+        LOG.info("loading processor...")
+        processor = load_processor(cfg, tokenizer)
+
+    return model, tokenizer, processor
diff --git a/src/axolotl/cli/utils/lora_merge.py b/src/axolotl/cli/utils/lora_merge.py
new file mode 100644
index 0000000000..92d8548945
--- /dev/null
+++ b/src/axolotl/cli/utils/lora_merge.py
@@ -0,0 +1,2129 @@
+import gc
+import math
+import os
+import re
+import shutil
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+import safetensors
+import safetensors.torch
+import torch
+from huggingface_hub import snapshot_download
+from peft import LoraConfig
+from peft.utils.other import get_pattern_key
+from tqdm import tqdm
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _resolve_lora_alpha_for_key(
+    weight_key: str,
+    lora_config_dict: Dict,
+    weight_renamings: Optional[Dict[str, str]] = None,
+) -> Optional[int]:
+    # Mirror PEFT's get_pattern_key matching so merge uses the same per-module alpha as training.
+    alpha_pattern = lora_config_dict.get("alpha_pattern") or {}
+    if not alpha_pattern:
+        return None
+    module_path = (
+        weight_key.rsplit(".weight", 1)[0]
+        if weight_key.endswith(".weight")
+        else weight_key
+    )
+    pattern_keys = list(alpha_pattern.keys())
+    matched_key = get_pattern_key(pattern_keys, module_path)
+    if matched_key in alpha_pattern:
+        return alpha_pattern[matched_key]
+    # Fall back to renamed path so alpha lookup follows the same key resolution as find_lora_weights.
+    if weight_renamings:
+        import re
+
+        for src_pattern, tgt_pattern in weight_renamings.items():
+            renamed = re.sub(src_pattern, tgt_pattern, module_path)
+            if renamed != module_path:
+                matched_key = get_pattern_key(pattern_keys, renamed)
+                if matched_key in alpha_pattern:
+                    return alpha_pattern[matched_key]
+    return None
+
+
+def _build_layer_type_map(
+    base_model_path: Path, trust_remote_code: bool = False
+) -> dict[str, str]:
+    """Build a map of module_name -> layer_type using a meta-device model.
+
+    Instantiates the model architecture on the meta device (zero memory)
+    to inspect which modules are Linear vs Conv1d/Conv2d/Conv3d.
+    This avoids relying on weight tensor ndim heuristics.
+    """
+    import json as _json
+
+    import torch.nn as nn
+    from transformers import AutoConfig
+
+    config_path = base_model_path / "config.json"
+    if not config_path.exists():
+        return {}
+
+    try:
+        with open(config_path) as f:
+            model_config = _json.load(f)
+    except (OSError, _json.JSONDecodeError):
+        return {}
+
+    architectures = model_config.get("architectures", [])
+    if not architectures:
+        return {}
+
+    try:
+        config = AutoConfig.from_pretrained(
+            str(base_model_path), trust_remote_code=trust_remote_code
+        )
+    except Exception:
+        LOG.debug("Could not load config for layer type introspection")
+        return {}
+
+    # Determine the right Auto class from architectures
+    from transformers import (
+        AutoModel,
+        AutoModelForCausalLM,
+    )
+
+    auto_classes = [AutoModelForCausalLM, AutoModel]
+    try:
+        from transformers import AutoModelForImageTextToText
+
+        auto_classes.insert(0, AutoModelForImageTextToText)
+    except ImportError:
+        pass
+
+    model = None
+    for auto_cls in auto_classes:
+        try:
+            with torch.device("meta"):
+                model = auto_cls.from_config(
+                    config, trust_remote_code=trust_remote_code
+                )
+            break
+        except Exception:  # noqa: BLE001
+            LOG.debug(
+                "Could not instantiate meta model with %s, trying next",
+                auto_cls.__name__,
+            )
+
+    if model is None:
+        LOG.debug("Could not instantiate meta model for layer type introspection")
+        return {}
+
+    layer_types = {}
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Conv3d):
+            layer_types[name] = "Conv3d"
+        elif isinstance(module, nn.Conv2d):
+            layer_types[name] = "Conv2d"
+        elif isinstance(module, nn.Conv1d):
+            layer_types[name] = "Conv1d"
+        elif isinstance(module, nn.Linear):
+            layer_types[name] = "Linear"
+
+    del model
+    LOG.debug(
+        f"Layer type map: {len(layer_types)} modules "
+        f"({sum(1 for v in layer_types.values() if 'Conv' in v)} conv layers)"
+    )
+    return layer_types
+
+
+def _simulate_nf4_roundtrip(
+    tensor: torch.Tensor,
+    blocksize: Optional[int] = None,
+    compress_statistics: bool = True,
+    device: Optional[Union[str, torch.device]] = None,
+) -> torch.Tensor:
+    """
+    Simulate NF4 quantization roundtrip to match QLoRA training dynamics.
+
+    During QLoRA training, base weights are quantized to NF4 and dequantized on-the-fly
+    for each forward pass. The LoRA adapters learn to compensate for the quantization
+    noise in the dequantized weights. To match this at merge time, we apply the same
+    quantize → dequantize roundtrip so the merged result reflects what the model saw
+    during training.
+
+    Args:
+        tensor: Base model weight tensor (fp16/bf16/fp32)
+        blocksize: NF4 quantization block size (default: bitsandbytes default)
+        compress_statistics: Whether to use double quantization
+        device: Device for quantization computation.  bitsandbytes requires a
+            CUDA device; defaults to "cuda" when available.
+
+    Returns:
+        Tensor after NF4 quantize → dequantize roundtrip, in original dtype
+    """
+    import bitsandbytes.functional as bnb_F
+
+    quant_device: torch.device
+    if device is None:
+        quant_device = torch.device("cuda")
+    elif isinstance(device, str):
+        quant_device = torch.device(device)
+    else:
+        quant_device = device
+
+    if quant_device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(
+            "NF4 simulation requires CUDA but no GPU is available. "
+            "Either run on a machine with a GPU or disable NF4 simulation."
+        )
+
+    original_dtype = tensor.dtype
+    original_shape = tensor.shape
+
+    # bitsandbytes requires float32 input for quantization and contiguous+CUDA tensor
+    flat = tensor.reshape(-1).to(torch.float32).contiguous().to(quant_device)
+
+    quant_kwargs = {
+        "quant_type": "nf4",
+        "compress_statistics": compress_statistics,
+    }
+    if blocksize is not None:
+        quant_kwargs["blocksize"] = blocksize
+
+    quantized, quant_state = bnb_F.quantize_4bit(flat, **quant_kwargs)
+    dequantized = bnb_F.dequantize_4bit(quantized, quant_state, quant_type="nf4")
+
+    return dequantized.reshape(original_shape).to(original_dtype).cpu()
+
+
+def find_lora_weights(
+    lora_state: Dict[str, torch.Tensor],
+    key: str,
+    weight_renamings: Optional[Dict[str, str]] = None,
+) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+    """
+    Find corresponding LoRA A and B weights for a given key.
+
+    Also tries keys after applying weight renamings (from transformers v5
+    conversion mappings) in case the checkpoint key names differ from the
+    runtime model key names used by the LoRA adapter.
+    """
+    import re
+
+    clean_key = key[:-7] if key.endswith(".weight") else key
+
+    # Try the direct key first
+    a_key = f"base_model.model.{clean_key}.lora_A.weight"
+    b_key = f"base_model.model.{clean_key}.lora_B.weight"
+
+    lora_a = lora_state.get(a_key)
+    lora_b = lora_state.get(b_key)
+
+    if lora_a is not None and lora_b is not None:
+        return lora_a, lora_b
+
+    # Try renamed keys (checkpoint format → runtime format)
+    if weight_renamings:
+        for src_pattern, tgt_pattern in weight_renamings.items():
+            renamed_key = re.sub(src_pattern, tgt_pattern, clean_key)
+            if renamed_key != clean_key:
+                a_key = f"base_model.model.{renamed_key}.lora_A.weight"
+                b_key = f"base_model.model.{renamed_key}.lora_B.weight"
+                lora_a = lora_state.get(a_key)
+                lora_b = lora_state.get(b_key)
+                if lora_a is not None and lora_b is not None:
+                    return lora_a, lora_b
+
+    return None, None
+
+
+def _find_param_wrapper_lora(
+    lora_state: Dict[str, torch.Tensor],
+    key: str,
+    tensor_shape: Optional[tuple] = None,
+) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[str]]:
+    """
+    Find LoRA weights from a ParamWrapper (lora_target_parameters) that targets
+    a parent module containing this weight as a sub-parameter.
+
+    For example, base weight key 'model.layers.0.mlp.experts.down_proj' may have
+    LoRA at 'base_model.model.model.layers.0.mlp.experts.lora_A.weight' (targeting
+    the 'experts' module with 'down_proj' as the parameter_name).
+
+    When tensor_shape is provided, validates that the LoRA dimensions match the
+    target tensor (important when multiple ParamWrappers are nested and each
+    nesting level has different LoRA dimensions).
+
+    Returns (lora_A, lora_B, parameter_name) or (None, None, None).
+    """
+    clean_key = key[:-7] if key.endswith(".weight") else key
+    # Strip trailing parameter name to get the parent module path
+    # e.g., "model.layers.0.mlp.experts.down_proj" → parent="model.layers.0.mlp.experts", param="down_proj"
+    parts = clean_key.rsplit(".", 1)
+    if len(parts) != 2:
+        return None, None, None
+
+    parent_key, param_name = parts
+
+    # PEFT's ParamWrapper nesting: when multiple parameters are targeted on
+    # the same module, it nests wrappers. The outer wrapper's LoRA is at
+    # parent.lora_A/B and inner wrappers use parent.base_layer.lora_A/B,
+    # parent.base_layer.base_layer.lora_A/B, etc.
+    prefixes_to_try = [
+        f"base_model.model.{parent_key}",
+    ]
+    # Walk up .base_layer nesting levels (typically 1-2 deep)
+    for depth in range(1, 4):
+        bl = ".base_layer" * depth
+        prefixes_to_try.append(f"base_model.model.{parent_key}{bl}")
+
+    # Both 3D orientations exist: gpt-oss-style [E, in, out] pairs with
+    # (A_in, B_out) = (shape[1], shape[2]); Qwen3-style [E, out, in] with
+    # (A_in, B_out) = (shape[2], shape[1]). Exhaust every nesting level in the
+    # exact orientation before falling back to the transposed one, so a
+    # transposed outer LoRA cannot shadow an exact inner match.
+    orientations: tuple = (None,)
+    if tensor_shape is not None and len(tensor_shape) >= 3:
+        orientations = (
+            (tensor_shape[1], tensor_shape[2]),
+            (tensor_shape[2], tensor_shape[1]),
+        )
+
+    for orientation in orientations:
+        for prefix in prefixes_to_try:
+            a_key = f"{prefix}.lora_A.weight"
+            b_key = f"{prefix}.lora_B.weight"
+            lora_a = lora_state.get(a_key)
+            lora_b = lora_state.get(b_key)
+            if lora_a is None or lora_b is None:
+                continue
+
+            # When tensor_shape is given, verify dimensions match before returning.
+            # This prevents returning a mismatched LoRA from a different nesting level.
+            if orientation is not None and tensor_shape is not None:
+                num_experts = tensor_shape[0]
+                if not (
+                    lora_a.shape[0] == lora_b.shape[1]
+                    and lora_a.shape[0] % num_experts == 0
+                    and (lora_a.shape[1], lora_b.shape[0]) == orientation
+                ):
+                    continue  # Dimensions don't match, try next nesting level
+
+            return lora_a, lora_b, param_name
+
+    return None, None, None
+
+
+def _build_peft_layer_and_get_delta(
+    lora_a: torch.Tensor,
+    lora_b: torch.Tensor,
+    lora_config_dict: Dict,
+    base_tensor: torch.Tensor,
+    adapter_name: str = "default",
+    is_param_wrapper: bool = False,
+    magnitude: Optional[torch.Tensor] = None,
+    layer_type: Optional[str] = None,
+    lora_alpha_override: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Use PEFT's own layer classes to compute the LoRA delta weight.
+
+    Instead of re-implementing the merge math for every LoRA variant, this
+    constructs a lightweight PEFT layer, loads the A/B weights, and calls
+    ``get_delta_weight`` (or ``merge`` for DoRA) which handles standard LoRA,
+    RSLoRA, DoRA, and ParamWrapper (expert-blocked) LoRA.
+
+    Returns the delta tensor (same shape as base_tensor).
+    """
+    import warnings
+
+    import torch.nn as nn
+
+    r_total = lora_a.shape[0]
+    in_features = lora_a.shape[1]
+    out_features = lora_b.shape[0]
+    # Per-module override from alpha_pattern wins over the global alpha so merge matches training scale.
+    if lora_alpha_override is not None:
+        lora_alpha = lora_alpha_override
+    else:
+        lora_alpha = lora_config_dict.get("lora_alpha", lora_config_dict.get("r", 1))
+    use_rslora = bool(lora_config_dict.get("use_rslora", False))
+    use_dora = bool(lora_config_dict.get("use_dora", False))
+
+    if is_param_wrapper:
+        from peft.tuners.lora.layer import ParamWrapper
+
+        num_experts = base_tensor.shape[0]
+        r = r_total // num_experts
+
+        class _FakeModule(nn.Module):
+            pass
+
+        fake = _FakeModule()
+        fake.register_parameter(
+            "weight", nn.Parameter(base_tensor.clone(), requires_grad=False)
+        )
+
+        # ParamWrapper rejects dropout/fan_in_fan_out/lora_bias/use_dora, so
+        # build a minimal config with only the fields it accepts.
+        pw_config = LoraConfig(
+            r=r,
+            lora_alpha=lora_alpha,
+            lora_dropout=0.0,
+            fan_in_fan_out=False,
+            use_rslora=use_rslora,
+            use_dora=False,
+            lora_bias=False,
+        )
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", UserWarning)
+            layer = ParamWrapper(
+                fake,
+                adapter_name=adapter_name,
+                parameter_name="weight",
+                config=pw_config,
+                r=r,
+                lora_alpha=lora_alpha,
+            )
+        layer.lora_A[adapter_name].weight.data = lora_a
+        layer.lora_B[adapter_name].weight.data = lora_b
+        delta = layer.get_delta_weight(adapter_name)
+        # peft >=0.19.1 may return delta with transposed dims for 3D params
+        if delta.shape != base_tensor.shape and delta.ndim == 3:
+            delta = delta.transpose(1, 2).contiguous()
+        return delta
+    elif (
+        layer_type and "Conv" in layer_type or (layer_type is None and lora_a.ndim > 2)
+    ):
+        # Conv layer detected via model introspection (or ndim fallback)
+
+        from peft.tuners.lora import layer as peft_lora_layer
+
+        # Determine conv type from layer_type map or fall back to ndim
+        if layer_type and "Conv" in layer_type:
+            conv_type: str = layer_type
+        else:
+            ndim = lora_a.ndim
+            _conv_map = {3: "Conv1d", 4: "Conv2d", 5: "Conv3d"}
+            if ndim not in _conv_map:
+                raise ValueError(
+                    f"Unsupported LoRA weight dimensionality {ndim} for conv layer"
+                )
+            conv_type = _conv_map[ndim]
+            LOG.warning(
+                f"Using ndim-based fallback for conv detection (ndim={ndim}). "
+                f"Consider providing layer_type from meta-device introspection."
+            )
+
+        conv_cls_map = {"Conv1d": nn.Conv1d, "Conv2d": nn.Conv2d, "Conv3d": nn.Conv3d}
+        ConvCls = conv_cls_map[conv_type]
+        PeftConvCls = getattr(peft_lora_layer, conv_type)
+
+        # Reconstruct conv parameters from base tensor and lora_a shapes
+        # base_tensor: [out_channels, in_channels/groups, *kernel_size]
+        # lora_a:      [r, in_channels/groups, *kernel_size]
+        # lora_b:      [out_channels, r, *ones]
+        out_channels = base_tensor.shape[0]
+        in_channels = base_tensor.shape[1]
+        kernel_size = tuple(base_tensor.shape[2:])
+        stride = (1,) * (base_tensor.ndim - 2)
+        padding = (0,) * (base_tensor.ndim - 2)
+
+        base_layer = ConvCls(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=False,
+        )
+        base_layer.weight.data = base_tensor.clone()
+
+        conv_config = LoraConfig(
+            r=r_total,
+            lora_alpha=lora_alpha,
+            use_rslora=use_rslora,
+            use_dora=use_dora,
+        )
+        layer = PeftConvCls(
+            base_layer,
+            adapter_name=adapter_name,
+            config=conv_config,
+            r=r_total,
+            lora_alpha=lora_alpha,
+        )
+        layer.lora_A[adapter_name].weight.data = lora_a
+        layer.lora_B[adapter_name].weight.data = lora_b
+
+        if use_dora:
+            if magnitude is None:
+                raise ValueError(
+                    f"DoRA merge requires a magnitude vector but none was found "
+                    f"for conv layer (adapter={adapter_name}). Check that the "
+                    f"adapter checkpoint contains lora_magnitude_vector weights."
+                )
+            mag_layer = layer.lora_magnitude_vector[adapter_name]
+            mag_layer.weight = nn.Parameter(magnitude)
+            layer.merge(adapter_names=[adapter_name])
+            return base_layer.weight.data - base_tensor
+
+        return layer.get_delta_weight(adapter_name)
+    else:
+        from peft.tuners.lora.layer import Linear as LoraLinear
+
+        base_layer = nn.Linear(in_features, out_features, bias=False)
+        base_layer.weight.data = base_tensor.clone()
+
+        fan_in_fan_out = bool(
+            lora_config_dict.get("fan_in_fan_out", False)
+            or lora_config_dict.get("lora_fan_in_fan_out", False)
+        )
+
+        linear_config = LoraConfig(
+            r=r_total,
+            lora_alpha=lora_alpha,
+            fan_in_fan_out=fan_in_fan_out,
+            use_rslora=use_rslora,
+            use_dora=use_dora,
+        )
+        layer = LoraLinear(
+            base_layer,
+            adapter_name=adapter_name,
+            config=linear_config,
+            r=r_total,
+            lora_alpha=lora_alpha,
+        )
+        layer.lora_A[adapter_name].weight.data = lora_a
+        layer.lora_B[adapter_name].weight.data = lora_b
+
+        if use_dora:
+            # DoRA merges magnitude normalization into the weight directly.
+            # Use PEFT's merge() which handles DoRA internally, then
+            # compute the delta as merged_weight - original_weight.
+            if magnitude is None:
+                raise ValueError(
+                    f"DoRA merge requires a magnitude vector but none was found "
+                    f"for linear layer (adapter={adapter_name}). Check that the "
+                    f"adapter checkpoint contains lora_magnitude_vector weights."
+                )
+            mag_layer = layer.lora_magnitude_vector[adapter_name]
+            mag_layer.weight = nn.Parameter(magnitude)
+            layer.merge(adapter_names=[adapter_name])
+            return base_layer.weight.data - base_tensor
+
+        return layer.get_delta_weight(adapter_name)
+
+
+def get_model_shards(model_path: Path) -> list[Path]:
+    """Find all model shards in the given path."""
+    shards: list[Path] = []
+
+    patterns = ["model*.safetensors", "pytorch_model*.bin"]
+
+    for pattern in patterns:
+        shards.extend(model_path.glob(pattern))
+        if shards:
+            break
+
+    return sorted(shards)
+
+
+def copy_non_model_files(
+    input_path: Path, output_path: Path, model_shards: list[Path]
+) -> None:
+    """
+    Copy all non-model files to the output directory.
+
+    Args:
+        input_path: Source directory
+        output_path: Destination directory
+        model_shards: List of model shard files to skip
+    """
+    LOG.info("Copying non-model files to output directory...")
+
+    shard_names = {shard.name for shard in model_shards}
+
+    for filepath in input_path.glob("*"):
+        if filepath.is_dir():
+            continue
+        if filepath.name in shard_names:
+            continue
+        if (
+            filepath.name.startswith("model") and filepath.suffix == ".safetensors"
+        ) or (filepath.name.startswith("pytorch_model") and filepath.suffix == ".bin"):
+            continue
+        if filepath.suffix == ".gguf":
+            continue
+        # Skip weight-map index files — they reference shard filenames that may
+        # change during the merge (e.g. .bin → .safetensors).  A correct index
+        # is regenerated after all shards have been written.
+        if filepath.name.endswith(".index.json"):
+            continue
+
+        LOG.debug(f"Copying {filepath.name} to output")
+        shutil.copy2(filepath, output_path)
+
+
+def _find_dora_magnitude(
+    lora_state: Dict[str, torch.Tensor],
+    key: str,
+    weight_renamings: Optional[Dict[str, str]] = None,
+) -> Optional[torch.Tensor]:
+    """
+    Find DoRA magnitude vector for a given key.
+    """
+    import re
+
+    clean_key = key[:-7] if key.endswith(".weight") else key
+    mag_key = f"base_model.model.{clean_key}.lora_magnitude_vector"
+    result = lora_state.get(mag_key)
+    if result is not None:
+        return result
+
+    if weight_renamings:
+        for src_pattern, tgt_pattern in weight_renamings.items():
+            renamed_key = re.sub(src_pattern, tgt_pattern, clean_key)
+            if renamed_key != clean_key:
+                mag_key = f"base_model.model.{renamed_key}.lora_magnitude_vector"
+                result = lora_state.get(mag_key)
+                if result is not None:
+                    return result
+
+    return None
+
+
+def _dequant_block_fp8(w: torch.Tensor, si: torch.Tensor, dev: str) -> torch.Tensor:
+    """FineGrainedFP8 (block-fp8): ``W`` e4m3 * ``scale_inv`` fp32 (128x128 blocks). 2D or fused-3D
+    (block axes = last two dims)."""
+    wf = w.to(dev).float()
+    s = si.to(dev).float()
+    *lead, N, K = w.shape
+    sr, sc = s.shape[-2], s.shape[-1]
+    # FineGrainedFP8 dims are always block-aligned; a ragged grid would need a fixed 128-block
+    # pad/slice (not a floor-spread), so refuse rather than silently mis-scale.
+    if N % sr or K % sc:
+        raise ValueError(
+            f"block-fp8 weight {tuple(w.shape)} is not aligned to its scale grid "
+            f"({sr}x{sc}); FineGrainedFP8 requires block-aligned dims"
+        )
+    bn, bk = N // sr, K // sc
+    return (wf.reshape(*lead, sr, bn, sc, bk) * s.reshape(*lead, sr, 1, sc, 1)).reshape(
+        *lead, N, K
+    )
+
+
+# Standard OCP-MX / NVFP4 FP4 E2M1 codebook, indexed by raw 4-bit nibble (sign|2-exp|1-mantissa).
+_FP4_E2M1_LUT = (
+    0.0,
+    0.5,
+    1.0,
+    1.5,
+    2.0,
+    3.0,
+    4.0,
+    6.0,
+    -0.0,
+    -0.5,
+    -1.0,
+    -1.5,
+    -2.0,
+    -3.0,
+    -4.0,
+    -6.0,
+)
+
+
+_MX_BLOCK = 32  # OCP-MX block width (one e8m0 scale per 32 elements)
+
+
+def _mx_block_scale(sb_bytes: torch.Tensor, N: int, K: int, dev: str) -> torch.Tensor:
+    """e8m0 biased-exponent bytes [.., N, nb] -> per-element multiplier [.., N, K]. Each scale covers
+    a fixed 32-wide MX block; a ragged final block (K not a multiple of 32) is trimmed."""
+    s = torch.exp2(sb_bytes.to(dev).float() - 127.0)
+    if s.shape[-1] * _MX_BLOCK == K:  # perfectly aligned -> memory-efficient expand
+        return s[..., None].expand(*s.shape, _MX_BLOCK).reshape(*s.shape[:-1], K)
+    return s.repeat_interleave(_MX_BLOCK, dim=-1)[..., :K]
+
+
+def _dequant_mxfp8(w: torch.Tensor, s: torch.Tensor, dev: str) -> torch.Tensor:
+    """OCP-MX fp8: ``W`` e4m3 * ``2^(e8m0_byte-127)`` (32-wide blocks along the last dim). ``s`` is the
+    e8m0 scale (``float8_e8m0fnu`` or raw ``uint8``). Ragged-K safe."""
+    wf = w.to(dev).float()
+    *lead, N, K = w.shape
+    scale = _mx_block_scale(s.view(torch.uint8), N, K, dev).reshape(*lead, N, K)
+    return wf * scale
+
+
+def _unpack_fp4(packed: torch.Tensor, dev: str) -> torch.Tensor:
+    """Unpack a packed-e2m1 uint8 tensor [.., K/2] -> fp32 values [.., K] via the OCP LUT. Low nibble =
+    element 2i, high nibble = element 2i+1 (torchao / OCP convention)."""
+    lut = torch.tensor(_FP4_E2M1_LUT, dtype=torch.float32, device=dev)
+    b = packed.to(dev)
+    lo = (b & 0xF).long()
+    hi = ((b >> 4) & 0xF).long()
+    nib = torch.stack([lo, hi], dim=-1).reshape(*b.shape[:-1], b.shape[-1] * 2)
+    return lut[nib]
+
+
+def _dequant_mxfp4(w: torch.Tensor, s: torch.Tensor, dev: str) -> torch.Tensor:
+    """OCP-MX fp4: packed e2m1 ``W`` [.., K/2] uint8 * ``2^(e8m0-127)`` (32-wide blocks). Ragged-K safe.
+    Matches the codebook + low/high nibble order the ScatterMoE MX forward uses, so the merged weight
+    equals what the model computes."""
+    vals = _unpack_fp4(w, dev)
+    *lead, N, K = vals.shape
+    scale = _mx_block_scale(s.view(torch.uint8), N, K, dev).reshape(*lead, N, K)
+    return vals * scale
+
+
+def _dequant_nvfp4(w, scale, scale2, dev: str) -> torch.Tensor:
+    """NVFP4: packed e2m1 ``W`` [..,K/2] uint8 + e4m3 block-16 ``scale`` [..,K/16] + optional per-tensor
+    ``scale_2`` (two-level; single-level when absent -> per_tensor_scale=1). Uses torchao's own
+    dequantize (matches the loader). Auto-detects swizzled scales by shape and passes the flag. Raises
+    if torchao is unavailable (caller degrades)."""
+    from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+    wt = w.to(dev)
+    st = scale.to(dev)
+    p = scale2.to(dev).float().reshape(()) if scale2 is not None else None
+    *lead, N, K = wt.shape
+    # a padded swizzled scale has more elements than the plain block-16 grid; torchao must unswizzle it
+    expect = 1
+    for d in (*lead, N, (K * 2) // 16):
+        expect *= d
+    swizzled = st.numel() != expect
+    nv = NVFP4Tensor(
+        wt, st, 16, torch.bfloat16, per_tensor_scale=p, is_swizzled_scales=swizzled
+    )
+    return nv.dequantize(torch.bfloat16)
+
+
+def _detect_quant_format(key, w, shard_tensors, e8m0):
+    """Return (fmt, {suffix: scale_tensor}) for a quantized weight ``key``, else (None, {}).
+    fmt in {block_fp8, mxfp8, nvfp4, mxfp4}."""
+    si = shard_tensors.get(key + "_scale_inv")
+    sc = shard_tensors.get(key + "_scale")
+    sc2 = shard_tensors.get(key + "_scale_2")
+    is_e8m0 = sc is not None and (
+        sc.dtype == torch.uint8 or (e8m0 is not None and sc.dtype == e8m0)
+    )
+    if w.dtype == torch.float8_e4m3fn and si is not None:
+        return "block_fp8", {"_scale_inv": si}
+    if w.dtype == torch.float8_e4m3fn and is_e8m0:
+        return "mxfp8", {"_scale": sc}
+    if w.dtype == torch.uint8 and sc is not None and sc.dtype == torch.float8_e4m3fn:
+        scales = {"_scale": sc}
+        if sc2 is not None:
+            scales["_scale_2"] = sc2
+        return "nvfp4", scales
+    if w.dtype == torch.uint8 and is_e8m0:
+        return "mxfp4", {"_scale": sc}
+    return None, {}
+
+
+def _dequant_by_format(fmt, w, scales, dev):
+    if fmt == "block_fp8":
+        return _dequant_block_fp8(w, scales["_scale_inv"], dev)
+    if fmt == "mxfp8":
+        return _dequant_mxfp8(w, scales["_scale"], dev)
+    if fmt == "nvfp4":
+        return _dequant_nvfp4(w, scales["_scale"], scales.get("_scale_2"), dev)
+    return _dequant_mxfp4(w, scales["_scale"], dev)  # mxfp4
+
+
+def _requant_by_format(fmt, w_bf16, scales, dev):
+    """Re-quantize a merged bf16 weight back to its original format. Returns
+    ``{"": qweight, "_scale*": scale_tensors}`` matching the original scale dtypes/shapes so the
+    merged checkpoint loads exactly like the base did.
+
+    nvfp4 reuses the base scales verbatim and only re-rounds the codes: recomputing scales shifts
+    the whole dequant grid, re-rounding EVERY element and burying a small LoRA delta under
+    uncorrelated noise, while on the original grid only elements the delta pushes across a code
+    boundary change. It also keeps gate/up outer scales equal, which the loader's fuse relies on."""
+    w = w_bf16.to(dev).float()
+    if fmt == "block_fp8":
+        si = scales["_scale_inv"]
+        *lead, N, K = w.shape
+        sr, sc = si.shape[-2], si.shape[-1]
+        bn, bk = N // sr, K // sc
+        wb = w.reshape(*lead, sr, bn, sc, bk)
+        amax = wb.abs().amax(dim=(-3, -1), keepdim=True).clamp_min(1e-12)
+        scale_inv = amax / 448.0
+        q = (wb / scale_inv).clamp_(-448, 448).to(torch.float8_e4m3fn)
+        return {
+            "": q.reshape(*lead, N, K).cpu(),
+            "_scale_inv": scale_inv.reshape(*lead, sr, sc).to(si.dtype).cpu(),
+        }
+    if fmt == "mxfp8":
+        q, ebyte = _quant_mx(w, 32)
+        s = scales["_scale"]
+        return {"": q.cpu(), "_scale": ebyte.view(s.dtype).cpu()}
+    if fmt == "mxfp4":
+        packed, ebyte = _quant_mxfp4(w)
+        s = scales["_scale"]
+        return {"": packed.cpu(), "_scale": ebyte.view(s.dtype).cpu()}
+    # nvfp4: original grid, fresh codes; bump only the block scales the delta outgrew
+    sc = scales["_scale"]
+    sc2 = scales.get("_scale_2")
+    pts = sc2.to(dev).float().reshape(()) if sc2 is not None else 1.0
+    sc_f = sc.to(dev).float()
+    amax = w.unflatten(-1, (sc_f.shape[-1], 16)).abs().amax(-1)
+    need = amax > 6.0 * sc_f * pts
+    if need.any():
+        sc_f = torch.where(need, (amax / (6.0 * pts)).clamp(max=448.0), sc_f)
+    sc_out = sc_f.to(sc.dtype)
+    denom = (sc_out.float() * pts).repeat_interleave(16, dim=-1).clamp_min(1e-30)
+    lut = torch.tensor(_FP4_E2M1_LUT, dtype=torch.float32, device=w.device)
+    idx = ((w / denom).unsqueeze(-1) - lut).abs().argmin(-1).to(torch.uint8)
+    packed = idx[..., 0::2] | (idx[..., 1::2] << 4)
+    out = {"": packed.cpu(), "_scale": sc_out.cpu()}
+    if sc2 is not None:
+        out["_scale_2"] = sc2.cpu()
+    return out
+
+
+def _quant_mx(w_f32, block):
+    """bf16/f32 -> (e4m3 qdata, uint8 e8m0 exponent byte), FLOOR e8m0, 32-wide blocks along last dim."""
+    *lead, N, K = w_f32.shape
+    nb = K // block
+    wb = w_f32.reshape(*lead, N, nb, block)
+    amax = wb.abs().amax(-1).clamp_min(1e-12)
+    exp = torch.floor(torch.log2(amax)) - 8.0
+    q = (wb / torch.exp2(exp)[..., None]).clamp_(-448, 448).to(torch.float8_e4m3fn)
+    ebyte = (exp + 127.0).clamp_(0, 254).to(torch.uint8)
+    return q.reshape(*lead, N, K), ebyte
+
+
+def _quant_mxfp4(w_f32):
+    """f32 -> (packed e2m1 uint8 [.., K/2], uint8 e8m0 [.., K/32]); nearest-codebook, low/high nibble."""
+    lut = torch.tensor(_FP4_E2M1_LUT, dtype=torch.float32, device=w_f32.device)
+    *lead, N, K = w_f32.shape
+    nb = K // _MX_BLOCK
+    wb = w_f32.reshape(*lead, N, nb, _MX_BLOCK)
+    amax = wb.abs().amax(-1).clamp_min(1e-6)
+    exp = torch.floor(torch.log2(amax / 6.0))
+    wn = (wb / torch.exp2(exp)[..., None]).reshape(*lead, N, K)
+    idx = (wn.unsqueeze(-1) - lut).abs().argmin(-1).to(torch.uint8)  # [.., N, K]
+    packed = idx[..., 0::2] | (idx[..., 1::2] << 4)
+    ebyte = (exp + 127.0).clamp_(0, 254).to(torch.uint8)
+    return packed, ebyte
+
+
+def _key_has_lora(key, shape, lora_state, weight_renamings):
+    a, b = find_lora_weights(lora_state, key, weight_renamings)
+    if a is not None and b is not None:
+        return True
+    if len(shape) >= 3:
+        pa, pb, _ = _find_param_wrapper_lora(lora_state, key, tuple(shape))
+        return pa is not None and pb is not None
+    return False
+
+
+def _dequantize_quantized_shard(
+    shard_tensors: Dict[str, torch.Tensor],
+    device: str,
+    lora_state: Optional[Dict[str, torch.Tensor]] = None,
+    weight_renamings: Optional[Dict[str, str]] = None,
+    dequant_all: bool = True,
+) -> tuple[Dict[str, torch.Tensor], bool, bool, Dict]:
+    """Dequantize quantized weights to bf16 so the LoRA delta folds into the true value.
+
+    Returns ``(new_shard, any_dequantized, any_left_quantized, requant_plan)``.
+
+    ``dequant_all=True`` dequantizes EVERY quantized weight to bf16 (the ``--dequant`` merge: output is
+    bf16). Default (``dequant_all=False``) is FORMAT-PRESERVING: only LoRA-targeted quantized weights
+    are dequantized (so the delta can fold correctly); ``requant_plan[key] = (fmt, scales)`` records how
+    to re-quantize each back to its original format after the merge, and quantized weights with no LoRA
+    are left untouched (same dtype + scales). ``any_left_quantized`` = a quantized weight was detected
+    but could not be dequantized (unsupported format / torchao missing) -> keep the config.
+
+    Formats (detected by scale sibling): block-fp8 (``_scale_inv`` fp32, 128x128), mxfp8 (``_scale``
+    e8m0/32), nvfp4 (``_scale`` e4m3/16 [+ ``_scale_2``]), mxfp4 (``_scale`` e8m0/32). Covers 2D linears
+    and fused-3D experts; native per-expert-unfused nvfp4/mxfp4 is not handled here."""
+    dev = device if (device != "cpu" and torch.cuda.is_available()) else "cpu"
+    e8m0 = getattr(torch, "float8_e8m0fnu", None)
+    out: Dict[str, torch.Tensor] = dict(shard_tensors)
+    drop: set = set()
+    plan: Dict = {}
+    did = False
+    left = False
+    lora_state = lora_state or {}
+    for key, w in shard_tensors.items():
+        if key in drop or key.endswith(("_scale", "_scale_inv", "_scale_2")):
+            continue
+        if w.ndim not in (2, 3):
+            continue
+        fmt, scales = _detect_quant_format(key, w, shard_tensors, e8m0)
+        if fmt is None:
+            # a low-bit weight with a scale sibling but an unrecognized format -> keep the config
+            if w.dtype in (torch.float8_e4m3fn, torch.uint8) and (
+                shard_tensors.get(key + "_scale_inv") is not None
+                or shard_tensors.get(key + "_scale") is not None
+            ):
+                left = True
+            continue
+        # format-preserving: only touch quantized weights a LoRA actually targets
+        if not dequant_all and not _key_has_lora(
+            key, w.shape, lora_state, weight_renamings
+        ):
+            continue
+        try:
+            deq = _dequant_by_format(fmt, w, scales, dev)
+        except Exception as ex:  # torchao missing / shape mismatch -> leave quantized
+            LOG.warning("%s dequant skipped for %s: %s", fmt, key, ex)
+            left = True
+            continue
+        out[key] = deq.to(torch.bfloat16).cpu()
+        did = True
+        drop.update(key + suf for suf in scales)
+        if not dequant_all:
+            plan[key] = (fmt, scales)  # re-quantize after the LoRA fold
+    for sk in drop:
+        out.pop(sk, None)
+    return out, did, left, plan
+
+
+_FUSED_EXPERT_LORA_RE = re.compile(r"\.experts\.(?:base_layer\.)?lora_[AB]\.weight$")
+_PER_EXPERT_WEIGHT_RE = re.compile(
+    r"\.experts\.\d+\.(?:gate_proj|up_proj|down_proj|w1|w2|w3|gate_up_proj)\.weight$"
+)
+
+
+def _detect_per_expert_unfused_mismatch(model_shards, lora_state) -> bool:
+    """True if the adapter carries a FUSED expert LoRA (``experts.lora_A``, targeting a fused
+    ``experts.gate_up_proj`` param) but the base stores experts PER-EXPERT and unfused
+    (``experts.0.gate_proj.weight`` ...). In that case the fused adapter keys match no base tensor,
+    so the expert LoRA would be *silently dropped* by the shard-by-shard merge. Peeks shard keys
+    only (no tensor loads)."""
+    if not any(_FUSED_EXPERT_LORA_RE.search(k) for k in lora_state):
+        return False
+    for shard in model_shards:
+        try:
+            if str(shard).endswith(".safetensors"):
+                with safetensors.safe_open(shard, framework="pt") as f:
+                    keys = list(f.keys())
+            else:
+                keys = list(
+                    torch.load(shard, map_location="meta", weights_only=True).keys()  # nosec B614
+                )
+        except Exception:  # noqa: BLE001  # nosec B112 - unreadable shard: skip the peek, not fatal
+            continue
+        if any(_PER_EXPERT_WEIGHT_RE.search(k) for k in keys):
+            return True
+    return False
+
+
+def _per_expert_weights_are_packed(model_shards) -> bool:
+    """True if any per-expert unfused expert weight is packed uint8 (NVFP4/MXFP4 qdata) — the layout
+    the expert-merge writer understands. Reads safetensors headers only."""
+    for shard in model_shards:
+        try:
+            if str(shard).endswith(".safetensors"):
+                with safetensors.safe_open(shard, framework="pt") as f:
+                    for k in f.keys():
+                        if _PER_EXPERT_WEIGHT_RE.search(k) and ".experts." in k:
+                            if f.get_slice(k).get_dtype() == "U8":
+                                return True
+            else:
+                tensors = torch.load(shard, map_location="meta", weights_only=True)  # nosec B614
+                for k, t in tensors.items():
+                    if _PER_EXPERT_WEIGHT_RE.search(k) and t.dtype == torch.uint8:
+                        return True
+        except Exception:  # noqa: BLE001  # nosec B112 - unreadable shard: skip the peek, not fatal
+            continue
+    return False
+
+
+_EXPERT_TRIPLE_RE = re.compile(
+    r"^(?P<prefix>.*\.experts)\.(?P<e>\d+)\."
+    r"(?P<proj>gate_proj|up_proj|down_proj|w1|w2|w3)\."
+    r"(?P<leaf>weight|weight_scale|weight_scale_2)$"
+)
+_EXPERT_LEAVES = ("weight", "weight_scale", "weight_scale_2")
+# (fused runtime param name, per-expert checkpoint proj names, concatenated on the row axis in order)
+_EXPERT_FUSED_GROUPS = (
+    ("gate_up_proj", ("gate_proj", "up_proj")),
+    ("gate_up_proj", ("w1", "w3")),
+    ("down_proj", ("down_proj",)),
+    ("down_proj", ("w2",)),
+)
+
+
+class _Nvfp4ExpertMergeWriter:
+    """Folds a FUSED expert LoRA (PEFT ParamWrapper over ``experts.gate_up_proj``/``down_proj``)
+    into a base that stores experts PER-EXPERT unfused as modelopt NVFP4
+    (``experts.<i>.<proj>.{weight,weight_scale,weight_scale_2}``), where no base key matches the
+    adapter and the shard merge would otherwise drop the expert LoRA.
+
+    ``consume`` claims the per-expert quantized tensors of LoRA-targeted layers out of each shard
+    (buffering across shard boundaries, since a layer's expert list can be split). When a fused
+    group is complete it dequantizes each expert (torchao, the same path training saw), fuses
+    to the runtime 3D layout (stack experts, concat gate-then-up rows), folds the ParamWrapper
+    delta, then unfuses and re-quantizes each expert back to NVFP4 with fresh scales — so the
+    merged checkpoint keeps the base's exact per-expert layout. Under ``dequant=True`` it emits
+    the merged FUSED bf16 param instead (matching the bf16 fuse pass convention).
+    """
+
+    def __init__(
+        self,
+        lora_state: Dict[str, torch.Tensor],
+        lora_config_dict: Dict,
+        expected_num_experts: int,
+        device: str,
+        dequant: bool = False,
+    ):
+        self.lora_state = lora_state
+        self.lora_config_dict = lora_config_dict
+        self.num_experts = expected_num_experts
+        self.dequant = dequant
+        self._dev = device if (device != "cpu" and torch.cuda.is_available()) else "cpu"
+        # prefix -> proj -> expert_idx -> {leaf: tensor}
+        self.pending: Dict[str, Dict[str, Dict[int, Dict[str, torch.Tensor]]]] = {}
+        self._prefix_lora: Dict[str, bool] = {}
+        self.merged_groups = 0
+
+    def _prefix_has_fused_lora(self, prefix: str) -> bool:
+        has = self._prefix_lora.get(prefix)
+        if has is None:
+            has = any(
+                f"base_model.model.{prefix}{'.base_layer' * d}.lora_A.weight"
+                in self.lora_state
+                for d in range(4)
+            )
+            self._prefix_lora[prefix] = has
+        return has
+
+    @torch.no_grad()
+    def consume(
+        self, shard_tensors: Dict[str, torch.Tensor]
+    ) -> tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], int]:
+        """Claim this shard's per-expert NVFP4 tensors for LoRA-targeted expert modules and emit
+        the merged tensors of every layer group that is now complete. Returns
+        ``(remaining_shard_tensors, emitted_tensors, merged_lora_count)``."""
+        remaining = dict(shard_tensors)
+        for key in list(remaining):
+            m = _EXPERT_TRIPLE_RE.match(key)
+            if m is None or not self._prefix_has_fused_lora(m["prefix"]):
+                continue
+            # only the packed-uint8 layout is understood; a float weight (bf16 base) flows through
+            if m["leaf"] == "weight" and remaining[key].dtype != torch.uint8:
+                continue
+            self.pending.setdefault(m["prefix"], {}).setdefault(
+                m["proj"], {}
+            ).setdefault(int(m["e"]), {})[m["leaf"]] = remaining.pop(key)
+
+        emitted: Dict[str, torch.Tensor] = {}
+        merged = 0
+        for prefix in list(self.pending):
+            for fused_name, members in _EXPERT_FUSED_GROUPS:
+                if not self._group_complete(prefix, members):
+                    continue
+                out, n = self._process_group(prefix, fused_name, members)
+                emitted.update(out)
+                merged += n
+                for mp in members:
+                    del self.pending[prefix][mp]
+            if not self.pending[prefix]:
+                del self.pending[prefix]
+        self.merged_groups += merged
+        return remaining, emitted, merged
+
+    def _group_complete(self, prefix: str, members: tuple) -> bool:
+        projs = self.pending[prefix]
+        if not any(mp in projs for mp in members):
+            return False
+        expected = set(range(self.num_experts))
+        return all(
+            mp in projs
+            and set(projs[mp]) == expected
+            and all(
+                all(leaf in projs[mp][e] for leaf in _EXPERT_LEAVES) for e in expected
+            )
+            for mp in members
+        )
+
+    def _process_group(
+        self, prefix: str, fused_name: str, members: tuple
+    ) -> tuple[Dict[str, torch.Tensor], int]:
+        projs = self.pending[prefix]
+        E = self.num_experts
+        # fused shape from the packed qdata headers (rows N, packed K/2 -> K), no dequant needed
+        row_counts = [projs[mp][0]["weight"].shape[0] for mp in members]
+        k_dim = projs[members[0]][0]["weight"].shape[1] * 2
+        fused_key = f"{prefix}.{fused_name}"
+        lora_a, lora_b, _ = _find_param_wrapper_lora(
+            self.lora_state, fused_key, tensor_shape=(E, sum(row_counts), k_dim)
+        )
+        emitted: Dict[str, torch.Tensor] = {}
+        if lora_a is None or lora_b is None:
+            LOG.warning(
+                "expert-merge writer: no shape-matching fused LoRA for %s; "
+                "passing its per-expert tensors through unchanged",
+                fused_key,
+            )
+            for mp in members:
+                for e in range(E):
+                    for leaf, t in projs[mp][e].items():
+                        emitted[f"{prefix}.{e}.{mp}.{leaf}"] = t
+            return emitted, 0
+
+        dev = self._dev
+        per_proj = [
+            torch.stack(
+                [
+                    _dequant_nvfp4(
+                        projs[mp][e]["weight"],
+                        projs[mp][e]["weight_scale"],
+                        projs[mp][e]["weight_scale_2"],
+                        dev,
+                    )
+                    for e in range(E)
+                ],
+                dim=0,
+            )
+            for mp in members
+        ]
+        fused = per_proj[0] if len(per_proj) == 1 else torch.cat(per_proj, dim=1)
+        del per_proj
+        delta = _build_peft_layer_and_get_delta(
+            lora_a.to(dev),
+            lora_b.to(dev),
+            self.lora_config_dict,
+            fused,
+            is_param_wrapper=True,
+        )
+        merged_t = (fused.to(torch.float32) + delta.to(torch.float32)).to(
+            torch.bfloat16
+        )
+        del fused, delta
+
+        if self.dequant:
+            emitted[fused_key] = merged_t.detach().cpu()
+            return emitted, 1
+
+        start = 0
+        for mp, rows in zip(members, row_counts, strict=True):
+            part = merged_t[:, start : start + rows, :]
+            start += rows
+            for e in range(E):
+                leaves = projs[mp][e]
+                wkey = f"{prefix}.{e}.{mp}.weight"
+                requant = _requant_by_format(
+                    "nvfp4",
+                    part[e],
+                    {
+                        "_scale": leaves["weight_scale"],
+                        "_scale_2": leaves["weight_scale_2"],
+                    },
+                    dev,
+                )
+                emitted[wkey] = requant.pop("")
+                for suf, t in requant.items():
+                    emitted[wkey + suf] = t
+        return emitted, 1
+
+    def assert_drained(self) -> None:
+        if not self.pending:
+            return
+        detail = {
+            prefix: {mp: len(ed) for mp, ed in projs.items()}
+            for prefix, projs in self.pending.items()
+        }
+        raise RuntimeError(
+            f"expert-merge writer: expert groups never completed (experts seen per projection: "
+            f"{detail}; expected {self.num_experts} per projection with weight/weight_scale/"
+            f"weight_scale_2 each). The base checkpoint is missing per-expert tensors."
+        )
+
+
+def _update_config_vocab_size(output_path: Path, vocab_size: int) -> None:
+    """After carrying a resized ``embed_tokens`` into the merge, set ``vocab_size`` in the merged
+    ``config.json`` (and any nested text config) so the checkpoint loads with the enlarged vocab."""
+    import json as _json
+
+    cfg_path = output_path / "config.json"
+    if not cfg_path.exists():
+        return
+    cfg = _json.loads(cfg_path.read_text())
+    changed = cfg.get("vocab_size") != vocab_size
+    cfg["vocab_size"] = vocab_size
+    for sub in ("text_config", "llm_config"):
+        if isinstance(cfg.get(sub), dict):
+            cfg[sub]["vocab_size"] = vocab_size
+            changed = True
+    if changed:
+        cfg_path.write_text(_json.dumps(cfg, indent=2))
+        LOG.info(
+            "Set merged config.json vocab_size=%d (resized embeddings carried)",
+            vocab_size,
+        )
+
+
+def _find_full_override(
+    lora_state: Dict[str, torch.Tensor], key: str
+) -> Optional[torch.Tensor]:
+    """Return the adapter's FULL-weight override for a base ``key`` if present, else None.
+
+    PEFT saves trainable non-LoRA modules — ``modules_to_save`` and resized ``embed_tokens`` /
+    ``lm_head`` — as plain ``base_model.model.<key>`` tensors (no ``lora_A``/``lora_B``). These must
+    REPLACE the base weight at merge, not be ignored (otherwise a resized vocab / trained head is
+    silently dropped). Matches ``.weight`` overrides only; ignores any ``.lora_`` tensor."""
+    if not key.endswith(".weight") or ".lora_" in key:
+        return None
+    for cand in (
+        "base_model.model." + key,
+        "base_model.model.model." + key,
+        "base_model.model." + key.replace("modules_to_save.", ""),
+    ):
+        t = lora_state.get(cand)
+        if t is not None and ".lora_" not in cand:
+            return t
+    # PEFT modules_to_save layout: ...<module>.modules_to_save.default.weight
+    ms = (
+        "base_model.model." + key[: -len(".weight")] + ".modules_to_save.default.weight"
+    )
+    return lora_state.get(ms)
+
+
+def _strip_quantization_config(output_path: Path) -> None:
+    """After a block-fp8 -> bf16 merge, remove ``quantization_config`` from the merged ``config.json``
+    so the merged checkpoint loads as bf16 (not FineGrainedFP8) and set ``torch_dtype`` to bfloat16."""
+    import json as _json
+
+    cfg_path = output_path / "config.json"
+    if not cfg_path.exists():
+        return
+    cfg = _json.loads(cfg_path.read_text())
+    changed = cfg.pop("quantization_config", None) is not None
+    if cfg.get("torch_dtype") not in (None, "bfloat16"):
+        cfg["torch_dtype"] = "bfloat16"
+        changed = True
+    if changed:
+        cfg_path.write_text(_json.dumps(cfg, indent=2))
+        LOG.info(
+            "Stripped quantization_config from merged config.json (block-fp8 -> bf16 merge)"
+        )
+
+
+_QUANT_DTYPES = {
+    torch.float8_e4m3fn,
+    getattr(torch, "float8_e5m2", torch.float8_e4m3fn),
+    torch.uint8,
+    torch.int8,
+}
+_WARNED_UNDEQUANT: set = set()
+
+
+def _warn_if_quant_undequantized(key: str, tensor: torch.Tensor, do_nf4: bool) -> None:
+    """Loudly warn (once/key) if a LoRA delta is about to be folded into a still-quantized weight
+    that the shard dequant did NOT convert to a real dtype (an unhandled format — e.g. mxfp4, a
+    per-tensor-fp8 with a scalar scale, or a per-expert-unfused nvfp4/mxfp4 layout). Folding into raw
+    packed bytes silently corrupts the weight; NF4 is handled by the simulate_nf4 roundtrip and is
+    exempt."""
+    if do_nf4 or tensor.dtype not in _QUANT_DTYPES or key in _WARNED_UNDEQUANT:
+        return
+    _WARNED_UNDEQUANT.add(key)
+    LOG.warning(
+        "LoRA merge: '%s' is still %s (a quantized format the merge did not dequantize) yet a LoRA "
+        "delta targets it — folding into raw quantized data is WRONG. This format is unsupported by "
+        "the efficient merge (handled: bf16, nf4-sim, block-fp8, mxfp8, fused-nvfp4). For per-expert "
+        "nvfp4/mxfp4 experts use the nvfp4 expert-merge writer; otherwise use merge_method: legacy.",
+        key,
+        tensor.dtype,
+    )
+
+
+def _should_nf4_roundtrip(
+    key: str,
+    tensor: torch.Tensor,
+    simulate_nf4: bool,
+    simulate_nf4_experts: bool,
+) -> bool:
+    """Determine if a tensor should undergo NF4 quantization roundtrip."""
+    if tensor.ndim < 2:
+        return False
+    if simulate_nf4:
+        return True
+    if simulate_nf4_experts and tensor.ndim >= 3 and "expert" in key.lower():
+        return True
+    return False
+
+
+def _merge_tensor_with_lora(
+    tensor: torch.Tensor,
+    key: str,
+    lora_state: Dict[str, torch.Tensor],
+    scale: float,
+    lora_config_dict: Dict,
+    device: str,
+    simulate_nf4: bool = False,
+    simulate_nf4_experts: bool = False,
+    nf4_blocksize: Optional[int] = None,
+    nf4_double_quant: bool = True,
+    use_dora: bool = False,
+    weight_renamings: Optional[Dict[str, str]] = None,
+    layer_type_map: Optional[Dict[str, str]] = None,
+) -> tuple[torch.Tensor, bool]:
+    """
+    Helper function to merge a single tensor with its corresponding LoRA weights.
+
+    Args:
+        tensor: Base model tensor
+        key: Tensor key/name
+        lora_state: Dictionary containing LoRA weights
+        scale: LoRA scaling factor (alpha/r)
+        lora_config_dict: LoRA configuration dictionary
+        device: Device to perform computations on
+        simulate_nf4: Whether to simulate NF4 quantization roundtrip for all weights
+        simulate_nf4_experts: Whether to simulate NF4 roundtrip for MoE expert tensors only
+        nf4_blocksize: Block size for NF4 quantization
+        nf4_double_quant: Whether to use double quantization
+        use_dora: Whether to apply DoRA (Weight-Decomposed LoRA) merging
+        weight_renamings: Optional key renamings from transformers conversion mapping
+
+    Returns:
+        Tuple of (merged tensor, whether LoRA was applied)
+    """
+    lora_a, lora_b = find_lora_weights(lora_state, key, weight_renamings)
+
+    do_nf4 = _should_nf4_roundtrip(key, tensor, simulate_nf4, simulate_nf4_experts)
+
+    if lora_a is not None and lora_b is not None:
+        LOG.debug(f"Merging LoRA for {key}: {lora_a.shape}, {lora_b.shape}")
+        _warn_if_quant_undequantized(key, tensor, do_nf4)
+
+        original_dtype = tensor.dtype
+
+        # Simulate NF4 quantization roundtrip to match QLoRA training dynamics
+        if do_nf4:
+            tensor = _simulate_nf4_roundtrip(
+                tensor,
+                blocksize=nf4_blocksize,
+                compress_statistics=nf4_double_quant,
+                device=device,
+            )
+
+        magnitude = (
+            _find_dora_magnitude(lora_state, key, weight_renamings)
+            if use_dora
+            else None
+        )
+
+        # Look up layer type from meta-device model introspection
+        _layer_type = None
+        if layer_type_map:
+            mod_path = key.rsplit(".weight", 1)[0] if key.endswith(".weight") else key
+            _layer_type = layer_type_map.get(mod_path)
+            # Try common prefix variations (e.g. with/without "model." prefix)
+            if _layer_type is None:
+                for prefix in [
+                    "model.",
+                    "model.language_model.",
+                    "model.language_model.model.",
+                ]:
+                    _layer_type = layer_type_map.get(prefix + mod_path)
+                    if _layer_type:
+                        break
+
+        delta = _build_peft_layer_and_get_delta(
+            lora_a.to(device),
+            lora_b.to(device),
+            lora_config_dict,
+            tensor.to(device),
+            magnitude=magnitude.to(device) if magnitude is not None else None,
+            layer_type=_layer_type,
+            lora_alpha_override=_resolve_lora_alpha_for_key(
+                key, lora_config_dict, weight_renamings
+            ),
+        )
+        merged_tensor = (
+            (tensor.to(device).to(torch.float32) + delta.to(torch.float32))
+            .to(original_dtype)
+            .detach()
+            .cpu()
+        )
+        return merged_tensor, True
+    else:
+        # Try ParamWrapper LoRA (lora_target_parameters) — the LoRA targets a
+        # parent module and this weight is a sub-parameter of that module.
+        if tensor.ndim >= 3:
+            pw_a, pw_b, param_name = _find_param_wrapper_lora(
+                lora_state, key, tensor_shape=tuple(tensor.shape)
+            )
+            if pw_a is not None and pw_b is not None:
+                LOG.debug(
+                    f"Merging ParamWrapper LoRA for {key} "
+                    f"(param={param_name}): {pw_a.shape}, {pw_b.shape}"
+                )
+                _warn_if_quant_undequantized(key, tensor, do_nf4)
+                if do_nf4:
+                    tensor = _simulate_nf4_roundtrip(
+                        tensor,
+                        blocksize=nf4_blocksize,
+                        compress_statistics=nf4_double_quant,
+                        device=device,
+                    )
+                original_dtype = tensor.dtype
+                delta = _build_peft_layer_and_get_delta(
+                    pw_a.to(device),
+                    pw_b.to(device),
+                    lora_config_dict,
+                    tensor.to(device),
+                    is_param_wrapper=True,
+                    lora_alpha_override=_resolve_lora_alpha_for_key(
+                        key, lora_config_dict, weight_renamings
+                    ),
+                )
+                merged = (
+                    (tensor.to(device).to(torch.float32) + delta.to(torch.float32))
+                    .to(original_dtype)
+                    .detach()
+                    .cpu()
+                )
+                return merged, True
+
+        if do_nf4:
+            tensor = _simulate_nf4_roundtrip(
+                tensor,
+                blocksize=nf4_blocksize,
+                compress_statistics=nf4_double_quant,
+                device=device,
+            )
+        return tensor.detach().cpu(), False
+
+
+def _get_conversion_info(base_model_path: Path) -> tuple[Dict[str, str], list]:
+    """
+    Load the model's config.json and check if transformers has WeightRenaming
+    or WeightConverter mappings for this model type.
+
+    Returns:
+        - dict of {source_pattern: target_pattern} for simple renamings
+        - list of WeightConverter objects for fuse/unfuse operations
+    """
+    import json as _json
+
+    config_path = base_model_path / "config.json"
+    if not config_path.exists():
+        return {}, []
+
+    try:
+        with open(config_path) as f:
+            model_config = _json.load(f)
+    except (OSError, _json.JSONDecodeError):
+        return {}, []
+
+    model_type = model_config.get("model_type")
+    if not model_type:
+        return {}, []
+
+    try:
+        from transformers.conversion_mapping import get_checkpoint_conversion_mapping
+        from transformers.core_model_loading import WeightConverter, WeightRenaming
+    except ImportError:
+        return {}, []
+
+    conversions = get_checkpoint_conversion_mapping(model_type)
+    if not conversions:
+        return {}, []
+
+    renamings = {}
+    weight_converters = []
+    for conv in conversions:
+        if isinstance(conv, WeightRenaming):
+            # WeightRenaming stores patterns as lists internally
+            src_list = (
+                conv.source_patterns
+                if isinstance(conv.source_patterns, list)
+                else [conv.source_patterns]
+            )
+            tgt_list = (
+                conv.target_patterns
+                if isinstance(conv.target_patterns, list)
+                else [conv.target_patterns]
+            )
+            if len(src_list) == 1 and len(tgt_list) == 1:
+                renamings[src_list[0]] = tgt_list[0]
+        elif isinstance(conv, WeightConverter):
+            weight_converters.append(conv)
+
+    return renamings, weight_converters
+
+
+def _get_expected_num_experts(base_model_path: Path) -> Optional[int]:
+    """Expert count from config.json, used to detect expert lists split across shards."""
+    import json as _json
+
+    config_path = base_model_path / "config.json"
+    if not config_path.exists():
+        return None
+    try:
+        cfg = _json.loads(config_path.read_text())
+    except (OSError, _json.JSONDecodeError):
+        return None
+    for sub in (cfg, cfg.get("text_config"), cfg.get("llm_config")):
+        if not isinstance(sub, dict):
+            continue
+        for key in (
+            "num_experts",
+            "num_local_experts",
+            "n_routed_experts",
+            "num_routed_experts",
+        ):
+            val = sub.get(key)
+            if isinstance(val, int) and val > 0:
+                return val
+    return None
+
+
+def _fuse_and_unfuse_with_merge(
+    shard_tensors: Dict[str, torch.Tensor],
+    weight_converters: list,
+    lora_state: Dict[str, torch.Tensor],
+    scale: float,
+    lora_config_dict: Dict,
+    device: str,
+    simulate_nf4: bool = False,
+    simulate_nf4_experts: bool = False,
+    nf4_blocksize: Optional[int] = None,
+    nf4_double_quant: bool = True,
+    use_dora: bool = False,
+    weight_renamings: Optional[Dict[str, str]] = None,
+    layer_type_map: Optional[Dict[str, str]] = None,
+    expected_num_experts: Optional[int] = None,
+) -> tuple[Dict[str, torch.Tensor], int, set]:
+    """
+    For tensors matching WeightConverter patterns (MoE expert weights):
+    1. Fuse checkpoint-format tensors into runtime-format (e.g., per-expert → fused 3D)
+    2. Apply NF4 roundtrip + LoRA merge on the fused tensor
+    3. Unfuse back to checkpoint format for saving
+
+    Returns:
+        - Updated tensor dict
+        - Count of merged LoRA targets
+        - Set of keys that were processed (fused/merged/unfused) and should be
+          skipped by the per-tensor merge pass to avoid double NF4 roundtrip
+    """
+    import re
+
+    from transformers.core_model_loading import Concatenate, MergeModulelist
+
+    result = dict(shard_tensors)  # Start with all tensors
+    merged_count = 0
+    processed_keys: set = set()  # Keys that were fuse/unfuse processed
+
+    for converter in weight_converters:
+        src_patterns = (
+            converter.source_patterns
+            if isinstance(converter.source_patterns, list)
+            else [converter.source_patterns]
+        )
+        tgt_patterns = (
+            converter.target_patterns
+            if isinstance(converter.target_patterns, list)
+            else [converter.target_patterns]
+        )
+
+        # Build regex for each source pattern
+        pattern_regexes = []
+        for pat in src_patterns:
+            regex_str = re.escape(pat).replace(r"\.\*\.", r"\.(\d+)\.")
+            regex_str = (
+                regex_str.rstrip(r"\$") if regex_str.endswith(r"\$") else regex_str
+            )
+            pattern_regexes.append(re.compile(r"(.*\.)?" + regex_str + "$"))
+
+        # Group matching keys by layer prefix and source pattern
+        # {layer_prefix: {pat_idx: {expert_idx: (key, tensor)}}}
+        layer_groups: Dict[str, Dict[int, Dict[int, tuple[str, torch.Tensor]]]] = {}
+
+        for key in list(result.keys()):
+            for pat_idx, pat_regex in enumerate(pattern_regexes):
+                match = pat_regex.match(key)
+                if match:
+                    prefix = match.group(1) or ""
+                    # Extract expert index from the matched portion
+                    remaining = key[len(prefix) :]
+                    expert_match = re.search(r"\.(\d+)\.", remaining)
+                    expert_idx = int(expert_match.group(1)) if expert_match else 0
+
+                    layer_groups.setdefault(prefix, {}).setdefault(pat_idx, {})[
+                        expert_idx
+                    ] = (key, result[key])
+                    break
+
+        is_expert_list = any(
+            isinstance(op, MergeModulelist) for op in converter.operations
+        )
+
+        # Process each layer group
+        for prefix, pat_groups in layer_groups.items():
+            # Check we have all source patterns for this layer
+            if not pat_groups:
+                continue
+
+            # Shards are processed one at a time, so a layer's expert list can be
+            # split across shard boundaries. Fusing a partial list either crashes
+            # (gate/up count mismatch in torch.cat) or silently fuses a subset of
+            # experts under the fused key. Fusing still-quantized tensors is also
+            # wrong: it stacks raw qdata and orphans the per-expert scale siblings.
+            # Skip fusion in both cases; the per-tensor pass carries the tensors
+            # through unchanged.
+            index_sets = [set(g.keys()) for g in pat_groups.values()]
+            complete = (
+                len(pat_groups) == len(pattern_regexes)
+                and all(s == index_sets[0] for s in index_sets[1:])
+                and index_sets[0] == set(range(len(index_sets[0])))
+                and (
+                    not is_expert_list
+                    or expected_num_experts is None
+                    or len(index_sets[0]) == expected_num_experts
+                )
+            )
+            skip_reason = None
+            if not complete:
+                expected_str = (
+                    f", expected {expected_num_experts}"
+                    if is_expert_list and expected_num_experts is not None
+                    else ""
+                )
+                skip_reason = (
+                    "expert list incomplete in this shard (found "
+                    f"{[len(pat_groups[p]) for p in sorted(pat_groups)]} experts "
+                    f"per pattern{expected_str})"
+                )
+            elif any(
+                t.dtype in _QUANT_DTYPES
+                or k + "_scale" in result
+                or k + "_scale_inv" in result
+                for g in pat_groups.values()
+                for (k, t) in g.values()
+            ):
+                skip_reason = "tensors are still quantized (raw qdata cannot be fused)"
+            if skip_reason:
+                LOG.info(
+                    "Skipping fuse for '%s%s': %s; leaving per-expert tensors "
+                    "unchanged",
+                    prefix,
+                    tgt_patterns[0],
+                    skip_reason,
+                )
+                continue
+
+            # Step 1: Fuse — MergeModulelist (stack experts) per source pattern
+            fused_per_pattern = {}
+            original_keys_per_pattern: Dict[int, list[str]] = {}
+            num_experts = None
+
+            for pat_idx in sorted(pat_groups.keys()):
+                expert_data = pat_groups[pat_idx]
+                sorted_indices = sorted(expert_data.keys())
+                if num_experts is None:
+                    num_experts = len(sorted_indices)
+
+                sorted_tensors = [expert_data[idx][1] for idx in sorted_indices]
+                original_keys_per_pattern[pat_idx] = [
+                    expert_data[idx][0] for idx in sorted_indices
+                ]
+                fused_per_pattern[src_patterns[pat_idx]] = torch.stack(
+                    sorted_tensors, dim=0
+                )
+
+            # Apply remaining operations (Concatenate)
+            fused_tensor = None
+            has_concat = False
+            concat_dim = 1  # default
+
+            for op in converter.operations:
+                if isinstance(op, MergeModulelist):
+                    pass  # Already handled
+                elif isinstance(op, Concatenate):
+                    has_concat = True
+                    concat_dim = op.dim
+                    tensors_to_cat = [
+                        fused_per_pattern[sp]
+                        for sp in src_patterns
+                        if sp in fused_per_pattern
+                    ]
+                    if len(tensors_to_cat) > 1:
+                        fused_tensor = torch.cat(tensors_to_cat, dim=concat_dim)
+                    elif tensors_to_cat:
+                        fused_tensor = tensors_to_cat[0]
+
+            if not has_concat and len(fused_per_pattern) == 1:
+                fused_tensor = next(iter(fused_per_pattern.values()))
+
+            if fused_tensor is None:
+                continue
+
+            # Step 2: Build the fused key name and merge LoRA
+            fused_key = prefix + tgt_patterns[0]
+
+            # Apply NF4 roundtrip on the fused tensor (matching training dynamics)
+            do_nf4 = _should_nf4_roundtrip(
+                fused_key, fused_tensor, simulate_nf4, simulate_nf4_experts
+            )
+            if do_nf4:
+                fused_tensor = _simulate_nf4_roundtrip(
+                    fused_tensor,
+                    blocksize=nf4_blocksize,
+                    compress_statistics=nf4_double_quant,
+                    device=device,
+                )
+
+            # Try to find and merge LoRA weights for the fused key
+            lora_a, lora_b = find_lora_weights(lora_state, fused_key, weight_renamings)
+            if lora_a is not None and lora_b is not None:
+                LOG.debug(
+                    f"Merging LoRA for fused key {fused_key}: {lora_a.shape}, {lora_b.shape}"
+                )
+                original_dtype = fused_tensor.dtype
+                magnitude = (
+                    _find_dora_magnitude(lora_state, fused_key, weight_renamings)
+                    if use_dora
+                    else None
+                )
+                # Look up layer type for the fused key
+                _layer_type = None
+                if layer_type_map:
+                    mod_path = (
+                        fused_key.rsplit(".weight", 1)[0]
+                        if fused_key.endswith(".weight")
+                        else fused_key
+                    )
+                    _layer_type = layer_type_map.get(mod_path)
+                    if _layer_type is None:
+                        for prefix in [
+                            "model.",
+                            "model.language_model.",
+                            "model.language_model.model.",
+                        ]:
+                            _layer_type = layer_type_map.get(prefix + mod_path)
+                            if _layer_type:
+                                break
+
+                delta = _build_peft_layer_and_get_delta(
+                    lora_a.to(device),
+                    lora_b.to(device),
+                    lora_config_dict,
+                    fused_tensor.to(device),
+                    magnitude=magnitude.to(device) if magnitude is not None else None,
+                    layer_type=_layer_type,
+                    lora_alpha_override=_resolve_lora_alpha_for_key(
+                        fused_key, lora_config_dict, weight_renamings
+                    ),
+                )
+                fused_tensor = (
+                    (
+                        fused_tensor.to(device).to(torch.float32)
+                        + delta.to(torch.float32)
+                    )
+                    .to(original_dtype)
+                    .detach()
+                    .cpu()
+                )
+                merged_count += 1
+
+            # Step 3: Save in fused format (runtime format) so that the merged
+            # model can be loaded directly without needing WeightConverter
+            # fusion during from_pretrained (which can OOM for large MoE models).
+            # Remove the original per-expert keys and save the fused tensor
+            # under the runtime key name.
+            for pat_idx in sorted(original_keys_per_pattern.keys()):
+                for ok in original_keys_per_pattern[pat_idx]:
+                    result.pop(ok, None)
+                    processed_keys.add(ok)
+
+            result[fused_key] = fused_tensor.detach().cpu()
+            processed_keys.add(fused_key)
+
+    return result, merged_count, processed_keys
+
+
+def merge_lora_sharded_efficient(
+    base_model_path: Union[str, Path],
+    lora_adapter_path: Union[str, Path],
+    output_path: Union[str, Path],
+    device: str = "cpu",
+    safe_tensors: bool = True,
+    simulate_nf4: bool = False,
+    simulate_nf4_experts: bool = False,
+    nf4_blocksize: Optional[int] = None,
+    nf4_double_quant: bool = True,
+    trust_remote_code: bool = False,
+    dequant: bool = False,
+) -> None:
+    """
+    Memory-efficient LoRA merging that processes shards individually
+    without loading the full model into memory.
+
+    dequant: if True, dequantize every quantized weight and write a bf16 checkpoint (strips
+        quantization_config). Default False = FORMAT-PRESERVING: LoRA-targeted quantized weights are
+        dequantized, the delta folded, then re-quantized back to the SAME format (fp8 stays fp8,
+        nvfp4 stays nvfp4), so a large quantized base does not double in size.
+
+    Args:
+        simulate_nf4: Apply NF4 roundtrip to ALL weight tensors (for QLoRA)
+        simulate_nf4_experts: Apply NF4 roundtrip only to MoE expert tensors
+            (for quantize_moe_experts). Expert tensors are identified by having
+            "expert" in the key name and ndim >= 3.
+        trust_remote_code: Whether to trust remote code when loading model
+            config for layer-type introspection. Defaults to False for safety.
+    """
+    base_model_path = Path(base_model_path)
+    lora_adapter_path = Path(lora_adapter_path)
+    output_path = Path(output_path)
+
+    if "/" in str(base_model_path) and not base_model_path.exists():
+        base_model_path = Path(snapshot_download(str(base_model_path)))
+
+    # Check for weight conversion requirements (transformers v5)
+    weight_renamings, weight_converters = _get_conversion_info(base_model_path)
+    if weight_renamings:
+        LOG.debug(f"Found {len(weight_renamings)} weight renamings for this model type")
+    expected_num_experts = None
+    if weight_converters:
+        LOG.debug(
+            f"Found {len(weight_converters)} weight converters (fuse/unfuse) for this model type. "
+            f"Will fuse→merge→unfuse within each shard."
+        )
+        expected_num_experts = _get_expected_num_experts(base_model_path)
+
+    os.makedirs(output_path, exist_ok=True)
+
+    config_file = lora_adapter_path / "adapter_config.json"
+    if not config_file.exists():
+        raise FileNotFoundError(f"LoRA config not found: {config_file}")
+
+    lora_config_dict = LoraConfig.from_json_file(str(config_file))
+    if not lora_config_dict.get("r") or lora_config_dict["r"] <= 0:
+        raise ValueError("LoRA config 'r' must be > 0")
+
+    use_dora = bool(lora_config_dict.get("use_dora", False))
+
+    # Build layer type map via meta-device model introspection
+    layer_type_map = _build_layer_type_map(
+        base_model_path, trust_remote_code=trust_remote_code
+    )
+    unsupported_methods = []
+
+    # Check for AdaLoRA (Adaptive LoRA)
+    if lora_config_dict.get("use_adalora", False):
+        unsupported_methods.append("AdaLoRA (Adaptive LoRA)")
+
+    # Check for VeRA (Vector-based Random Matrix Adaptation)
+    if lora_config_dict.get("use_vera", False):
+        unsupported_methods.append("VeRA (Vector-based Random Matrix Adaptation)")
+
+    # Check for other advanced LoRA variants by task_type
+    task_type = lora_config_dict.get("task_type", "")
+    if task_type and task_type not in [
+        "CAUSAL_LM",
+        "SEQ_2_SEQ_LM",
+        "TOKEN_CLS",
+        "SEQ_CLS",
+        "QUESTION_ANS",
+    ]:
+        unsupported_methods.append(f"Task type: {task_type}")
+
+    # PEFT writes peft_type=ADALORA and target_r for AdaLoRA; rank_pattern/alpha_pattern alone are plain LoRA.
+    if (
+        lora_config_dict.get("peft_type") == "ADALORA"
+        or lora_config_dict.get("target_r") is not None
+    ):
+        unsupported_methods.append("AdaLoRA (rank adaptation detected)")
+
+    # Check for advanced initialization methods
+    init_lora_weights = lora_config_dict.get("init_lora_weights", "")
+    if init_lora_weights and init_lora_weights not in [
+        "gaussian",
+        "loftq",
+        True,
+        False,
+    ]:
+        unsupported_methods.append(f"Advanced initialization: {init_lora_weights}")
+
+    if unsupported_methods:
+        methods_str = ", ".join(unsupported_methods)
+        raise NotImplementedError(
+            f"Memory-efficient LoRA merge only supports standard LoRA. "
+            f"Detected unsupported methods: {methods_str}. "
+            f"Please use the legacy merge method for advanced LoRA variants."
+        )
+
+    use_rslora = bool(lora_config_dict.get("use_rslora", False))
+    if use_rslora:
+        scale = float(lora_config_dict["lora_alpha"]) / math.sqrt(
+            float(lora_config_dict["r"])
+        )
+    else:
+        scale = float(lora_config_dict["lora_alpha"]) / float(lora_config_dict["r"])
+
+    LOG.debug(f"LoRA scale factor: {scale} (rslora={use_rslora})")
+
+    if simulate_nf4:
+        LOG.info(
+            "NF4 simulation enabled: base weights will undergo quantize→dequantize "
+            "roundtrip before LoRA merge to match QLoRA training dynamics"
+        )
+
+    lora_file = lora_adapter_path / "adapter_model.safetensors"
+    if not lora_file.exists():
+        lora_file = lora_adapter_path / "adapter_model.bin"
+        if not lora_file.exists():
+            raise FileNotFoundError(
+                f"LoRA adapter weights not found in {lora_adapter_path}"
+            )
+
+    LOG.debug(f"Loading LoRA weights from {lora_file}")
+
+    if lora_file.suffix == ".safetensors":
+        lora_state = safetensors.torch.load_file(lora_file)
+    else:
+        lora_state = torch.load(lora_file, map_location="cpu", weights_only=True)  # nosec B614
+    LOG.debug("Keeping LoRA weights on CPU; will move per-tensor during merge")
+
+    model_shards = get_model_shards(base_model_path)
+    if not model_shards:
+        raise FileNotFoundError(f"No model shards found in {base_model_path}")
+
+    LOG.debug(f"Found {len(model_shards)} model shards in {base_model_path}")
+    copy_non_model_files(base_model_path, output_path, model_shards)
+
+    expert_writer = None
+    if _detect_per_expert_unfused_mismatch(model_shards, lora_state):
+        try:
+            import torchao  # noqa: F401
+
+            has_torchao = True
+        except ImportError:
+            has_torchao = False
+        if not has_torchao or not _per_expert_weights_are_packed(model_shards):
+            LOG.warning(
+                "MERGE INCOMPLETE: the adapter has a FUSED expert LoRA (experts.gate_up_proj/down_proj) "
+                "but this base stores experts PER-EXPERT and unfused (experts.<i>.gate_proj.weight ...) "
+                "in a layout the nvfp4 expert-merge writer cannot handle (%s), so the EXPERT LoRA is "
+                "being DROPPED (non-expert LoRA still merges). Use merge_method: legacy (loads the "
+                "full model so PEFT fuses the experts itself).",
+                "torchao is not installed"
+                if not has_torchao
+                else "expert weights are not packed uint8 NVFP4",
+            )
+        else:
+            n_experts = expected_num_experts or _get_expected_num_experts(
+                base_model_path
+            )
+            if n_experts is None:
+                raise RuntimeError(
+                    "The adapter has a fused expert LoRA over a per-expert unfused base, but "
+                    "config.json exposes no expert count (num_experts/num_local_experts/"
+                    "n_routed_experts), so the expert-merge writer cannot validate layer "
+                    "completeness across shards. Add the expert count to config.json or use "
+                    "merge_method: legacy."
+                )
+            expert_writer = _Nvfp4ExpertMergeWriter(
+                lora_state, lora_config_dict, n_experts, device, dequant=dequant
+            )
+            LOG.info(
+                "Adapter has a FUSED expert LoRA over a PER-EXPERT unfused NVFP4 base: using "
+                "the expert-merge writer (dequant -> fuse -> fold delta -> unfuse -> requant)."
+            )
+
+    merged_count = 0
+    total_tensors = 0
+    block_fp8_dequantized = False
+    left_quantized = False
+    resized_vocab = None
+    # Track weight_map for index regeneration: {tensor_key: shard_filename}
+    weight_map: Dict[str, str] = {}
+
+    for shard_path in tqdm(model_shards, desc="Merging shards"):
+        merged_tensors = {}
+        metadata = {}
+
+        # Load all tensors from the shard
+        if shard_path.suffix == ".safetensors":
+            with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f:
+                if hasattr(f, "metadata") and f.metadata():
+                    metadata = f.metadata()
+                shard_tensors = {key: f.get_tensor(key) for key in f.keys()}
+        else:
+            shard_tensors = torch.load(  # nosec B614: loading trusted model weights
+                shard_path, map_location="cpu", weights_only=True
+            )
+
+        total_tensors += len(shard_tensors)
+
+        # Per-expert unfused NVFP4 experts with a fused adapter: the writer claims those tensors
+        # (buffered across shard boundaries) and emits the merged per-expert keys itself, so the
+        # dequant/fuse/per-tensor passes below never see them.
+        if expert_writer is not None:
+            shard_tensors, expert_emitted, expert_merged = expert_writer.consume(
+                shard_tensors
+            )
+            merged_tensors.update(expert_emitted)
+            merged_count += expert_merged
+
+        # Step 0: dequantize quantized weights so the LoRA delta folds into the TRUE weight (a raw read
+        # misses the block scale). Default (dequant=False) touches only LoRA-targeted weights and plans
+        # to re-quantize them back to their original format after the fold; dequant=True dequants all.
+        shard_tensors, _shard_deq, _shard_left, requant_plan = (
+            _dequantize_quantized_shard(
+                shard_tensors,
+                device,
+                lora_state=lora_state,
+                weight_renamings=weight_renamings,
+                dequant_all=dequant,
+            )
+        )
+        block_fp8_dequantized = block_fp8_dequantized or _shard_deq
+        left_quantized = left_quantized or _shard_left
+
+        # Step 1: Handle fused weight conversions (MoE experts) if applicable
+        fused_keys: set = set()
+        if weight_converters:
+            shard_tensors, fused_merged, fused_keys = _fuse_and_unfuse_with_merge(
+                shard_tensors,
+                weight_converters,
+                lora_state,
+                scale,
+                lora_config_dict,
+                device,
+                simulate_nf4=simulate_nf4,
+                simulate_nf4_experts=simulate_nf4_experts,
+                nf4_blocksize=nf4_blocksize,
+                nf4_double_quant=nf4_double_quant,
+                use_dora=use_dora,
+                weight_renamings=weight_renamings,
+                layer_type_map=layer_type_map,
+                expected_num_experts=expected_num_experts,
+            )
+            merged_count += fused_merged
+
+        # Step 2: Merge remaining (non-fused) tensors with LoRA
+        # Skip keys already processed by fuse/unfuse to avoid double NF4 roundtrip
+        for key, tensor in shard_tensors.items():
+            if key in fused_keys:
+                merged_tensors[key] = tensor.detach().cpu()
+                continue
+            # Full-weight override (modules_to_save / resized embed_tokens+lm_head): replace, don't
+            # fold LoRA. Track a vocab resize so the merged config.json stays consistent.
+            override = _find_full_override(lora_state, key)
+            if override is not None:
+                merged_tensors[key] = override.to(tensor.dtype).detach().cpu()
+                merged_count += 1
+                # a resized vocab shows up as a row-count change on embed_tokens AND/OR lm_head
+                # (they may be untied); catch either so the merged config.json stays consistent.
+                if (
+                    key.endswith(("embed_tokens.weight", "lm_head.weight"))
+                    and override.shape[0] != tensor.shape[0]
+                ):
+                    resized_vocab = int(override.shape[0])
+                continue
+            merged_tensor, was_merged = _merge_tensor_with_lora(
+                tensor,
+                key,
+                lora_state,
+                scale,
+                lora_config_dict,
+                device,
+                simulate_nf4=simulate_nf4,
+                simulate_nf4_experts=simulate_nf4_experts,
+                nf4_blocksize=nf4_blocksize,
+                nf4_double_quant=nf4_double_quant,
+                use_dora=use_dora,
+                weight_renamings=weight_renamings,
+                layer_type_map=layer_type_map,
+            )
+            merged_tensors[key] = merged_tensor
+            if was_merged:
+                merged_count += 1
+
+        # Step 3 (format-preserving): re-quantize each merged weight back to its original format
+        # (fresh block scales) so the output dtype matches the input (fp8 -> fp8, nvfp4 -> nvfp4).
+        for key, (fmt, scales) in requant_plan.items():
+            if key not in merged_tensors:
+                continue
+            try:
+                requant = _requant_by_format(fmt, merged_tensors[key], scales, device)
+            except (
+                Exception
+            ) as ex:  # torchao missing etc. -> keep the bf16 merge for this weight
+                LOG.warning("%s re-quant skipped for %s: %s", fmt, key, ex)
+                left_quantized = True
+                continue
+            merged_tensors[key] = requant.pop("")
+            for suf, t in requant.items():
+                merged_tensors[key + suf] = t
+
+        output_shard_path = output_path / shard_path.name
+        merged_tensors = {k: v.detach().cpu() for k, v in merged_tensors.items()}
+
+        if safe_tensors:
+            if not str(output_shard_path).endswith(".safetensors"):
+                output_shard_path = output_path / (shard_path.stem + ".safetensors")
+            safetensors.torch.save_file(
+                merged_tensors, output_shard_path, metadata=metadata
+            )
+        else:
+            if shard_path.suffix == ".safetensors":
+                safetensors.torch.save_file(
+                    merged_tensors, output_shard_path, metadata=metadata
+                )
+            else:
+                torch.save(merged_tensors, output_shard_path)
+
+        for tensor_key in merged_tensors:
+            weight_map[tensor_key] = output_shard_path.name
+
+        del merged_tensors, shard_tensors
+        if device != "cpu" and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+    if expert_writer is not None:
+        expert_writer.assert_drained()
+        if expert_writer.merged_groups == 0:
+            LOG.warning(
+                "MERGE INCOMPLETE: the expert-merge writer matched no fused expert LoRA to any "
+                "per-expert layer group (unrecognized checkpoint naming or shape mismatch); the "
+                "expert LoRA was NOT merged."
+            )
+
+    # Regenerate weight-map index if the model was sharded
+    if len(model_shards) > 1 and weight_map:
+        import json as _json
+
+        index_name = (
+            "model.safetensors.index.json"
+            if safe_tensors
+            else "pytorch_model.bin.index.json"
+        )
+        index = {
+            "metadata": {"total_size": total_tensors},
+            "weight_map": weight_map,
+        }
+        with open(output_path / index_name, "w") as f:
+            _json.dump(index, f, indent=2)
+        LOG.debug(f"Wrote weight-map index: {index_name}")
+
+    # Only the --dequant merge emits a bf16 checkpoint; the format-preserving default keeps the
+    # quantized dtypes (and their quantization_config) intact.
+    if dequant and block_fp8_dequantized and not left_quantized:
+        _strip_quantization_config(output_path)
+    elif block_fp8_dequantized and left_quantized:
+        LOG.warning(
+            "Merged some quantized weights to bf16 but others were left quantized (unsupported "
+            "format); keeping quantization_config so the checkpoint still loads its quantized tensors."
+        )
+    if resized_vocab is not None:
+        _update_config_vocab_size(output_path, resized_vocab)
+
+    if merged_count == 0:
+        LOG.warning(
+            "No LoRA weights were matched to base model tensors. "
+            "This may indicate a key name mismatch between the checkpoint format "
+            "and the LoRA adapter. Consider using merge_method: legacy."
+        )
+    LOG.info(f"Applied LoRA to {merged_count}/{total_tensors} tensors")
diff --git a/src/axolotl/cli/utils/sweeps.py b/src/axolotl/cli/utils/sweeps.py
new file mode 100644
index 0000000000..2a0aa1367a
--- /dev/null
+++ b/src/axolotl/cli/utils/sweeps.py
@@ -0,0 +1,81 @@
+"""Utilities for handling sweeps over configs for axolotl train CLI command"""
+
+import random
+from copy import deepcopy
+from itertools import product
+from typing import Any
+
+
+def generate_sweep_configs(
+    base_config: dict[str, list], sweeps_config: dict[str, list]
+) -> list[dict[str, Any]]:
+    """
+    Recursively generates all possible configurations by applying sweeps to the base config.
+
+    Args:
+        base_config (dict): The original configuration dictionary
+        sweeps_config (dict): Dictionary where keys are parameters and values are either:
+            - lists of values to sweep independently
+            - or for paired values, a list of dicts under the '_' key
+
+    Returns:
+        list: List of all possible configuration dictionaries
+
+    Example:
+        sweeps_config = {
+            'learning_rate': [0.1, 0.01],
+            '_': [
+                {'load_in_8bit': True, 'adapter': 'lora'},
+                {'load_in_4bit': True, 'adapter': 'qlora'}
+            ]
+        }
+    """
+    # Separate paired values from regular sweeps
+    paired_values = sweeps_config.get("_", [])
+    regular_sweeps = {k: v for k, v in sweeps_config.items() if k != "_"}
+
+    # Process regular sweeps
+    param_names = list(regular_sweeps.keys())
+    param_values = list(regular_sweeps.values())
+
+    # Generate combinations for regular sweeps
+    regular_combinations = list(product(*param_values)) if param_values else [()]
+
+    # Combine regular sweeps with paired values
+    all_combinations = []
+    for reg_combo in regular_combinations:
+        if paired_values:
+            for paired_set in paired_values:
+                new_config = {}
+                # new_config = deepcopy(base_config)
+                # Combine regular parameters with paired parameters
+                full_combo = {
+                    **dict(zip(param_names, reg_combo, strict=False)),
+                    **paired_set,
+                }
+                for param_name, param_value in full_combo.items():
+                    new_config[param_name] = param_value
+                print(new_config)
+                all_combinations.append(new_config)
+        else:
+            # If no paired values, just use regular combinations
+            # new_config = deepcopy(base_config)
+            new_config = {}
+            for param_name, param_value in zip(param_names, reg_combo, strict=False):
+                new_config[param_name] = param_value
+            print(new_config)
+            all_combinations.append(new_config)
+
+    # randomize the order of trials
+    random.seed(42)
+    random.shuffle(all_combinations)
+
+    # Generate a new config for each combination
+    result_configs = []
+    for combination in all_combinations:
+        new_config = deepcopy(base_config)
+        for param_name, param_value in combination.items():
+            new_config[param_name] = param_value
+        result_configs.append(new_config)
+
+    return result_configs
diff --git a/src/axolotl/cli/utils/train.py b/src/axolotl/cli/utils/train.py
new file mode 100644
index 0000000000..6ce7d8df30
--- /dev/null
+++ b/src/axolotl/cli/utils/train.py
@@ -0,0 +1,225 @@
+"""Utilities for axolotl train CLI command."""
+
+import os
+import subprocess  # nosec
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Iterator, Literal
+
+import yaml
+
+from axolotl.cli.utils.sweeps import generate_sweep_configs
+
+
+def _add_default_rdzv_args(launcher_args: list[str]) -> list[str]:
+    """
+    Add default RDZV arguments if rdzv_endpoint is set but rdzv_backend/rdzv_id are missing.
+
+    Args:
+        launcher_args: List of launcher arguments
+
+    Returns:
+        Updated launcher args with defaults added if needed
+    """
+    args = launcher_args.copy()
+
+    # Check if rdzv_endpoint is present
+    has_rdzv_endpoint = any("--rdzv_endpoint" in arg for arg in args)
+
+    if has_rdzv_endpoint:
+        # Check if rdzv_backend is already provided
+        has_rdzv_backend = any("--rdzv_backend" in arg for arg in args)
+        if not has_rdzv_backend:
+            args.extend(["--rdzv_backend", "c10d"])
+
+        # Check if rdzv_id is already provided
+        has_rdzv_id = any("--rdzv_id" in arg for arg in args)
+        if not has_rdzv_id:
+            import uuid
+
+            args.extend(["--rdzv_id", str(uuid.uuid4())[:8]])
+
+    return args
+
+
+def build_command(base_cmd: list[str], options: dict[str, Any]) -> list[str]:
+    """
+    Build command list from base command and options.
+
+    Args:
+        base_cmd: Command without options.
+        options: Options to parse and append to base command.
+
+    Returns:
+        List of strings giving shell command.
+    """
+    cmd = base_cmd.copy()
+
+    for key, value in options.items():
+        if value is None:
+            continue
+
+        key = key.replace("_", "-")
+        cmd.append(f"--{key}={value}")
+
+    return cmd
+
+
+def generate_config_files(config: str, sweep: str | None) -> Iterator[tuple[str, bool]]:
+    """
+    Generate list of configuration files to process. Yields a tuple of the configuration file name and a boolean indicating
+    whether this is a group of configurations (i.e., a sweep).
+
+    Args:
+        config: Base configuration file
+        sweep: Sweep configuration file
+    """
+
+    if not sweep:
+        yield config, False
+        return
+
+    # Load sweep and base configurations
+    with open(sweep, "r", encoding="utf-8") as fin:
+        sweep_config: dict[str, list] = yaml.safe_load(fin)
+    with open(config, "r", encoding="utf-8") as fin:
+        base_config: dict[str, list] = yaml.safe_load(fin)
+
+    # Generate all possible configurations
+    permutations = generate_sweep_configs(base_config, sweep_config)
+    is_group = len(permutations) > 1
+    base_output_dir = base_config.get("output_dir", "./model-out")
+    for idx, permutation in enumerate(permutations, start=1):
+        permutation_dir = Path(permutation.get("output_dir", base_output_dir))
+        permutation_id = f"sweep{idx:04d}"
+        permutation["output_dir"] = str(permutation_dir / permutation_id)
+
+        temp_file = tempfile.NamedTemporaryFile(
+            mode="w",
+            suffix=".yaml",
+            delete=False,
+            encoding="utf-8",
+        )
+        yaml.dump(permutation, temp_file)
+        temp_file.close()
+        yield temp_file.name, is_group
+
+
+def launch_training(
+    cfg_file: str,
+    launcher: Literal["accelerate", "torchrun", "python"] | None,
+    cloud: str | None,
+    kwargs: dict,
+    launcher_args: list[str] | None = None,
+    use_exec: bool = False,
+) -> None:
+    """Execute training with the given configuration."""
+    launcher_args = launcher_args or []
+
+    if cloud:
+        _launch_cloud_training(cloud, cfg_file, launcher, kwargs, launcher_args)
+    elif launcher:
+        if launcher == "accelerate":
+            _launch_accelerate_training(cfg_file, kwargs, launcher_args, use_exec)
+        elif launcher == "torchrun":
+            _launch_torchrun_training(cfg_file, kwargs, launcher_args, use_exec)
+        elif launcher == "python":
+            _launch_python_training(cfg_file, kwargs)
+    elif launcher is None:
+        # handle ray train launch
+        _launch_python_training(cfg_file, kwargs)
+
+
+def _launch_cloud_training(
+    cloud: str,
+    cfg_file: str,
+    launcher: Literal["accelerate", "torchrun", "python"] | None,
+    kwargs: dict,
+    launcher_args: list[str] | None = None,
+) -> None:
+    """Execute training via cloud launcher."""
+    from axolotl.cli.cloud import do_cli_train
+
+    launcher_args = launcher_args or []
+    cwd = os.getcwd() if launcher else None
+
+    do_cli_train(
+        cloud_config=cloud,
+        config=cfg_file,
+        launcher=launcher or "accelerate",
+        launcher_args=launcher_args,
+        cwd=cwd,
+        **kwargs,
+    )
+
+
+def _launch_accelerate_training(
+    cfg_file: str,
+    kwargs: dict,
+    launcher_args: list[str] | None = None,
+    use_exec: bool = False,
+) -> None:
+    """Execute training via accelerate launcher."""
+    launcher_args = launcher_args or []
+    internal_launcher_args = []
+
+    # Extract launcher-specific arguments from kwargs (legacy support)
+    if "main_process_port" in kwargs:
+        main_process_port = kwargs.pop("main_process_port")
+        internal_launcher_args.extend(["--main_process_port", str(main_process_port)])
+
+    if "num_processes" in kwargs:
+        num_processes = kwargs.pop("num_processes")
+        internal_launcher_args.extend(["--num_processes", str(num_processes)])
+
+    # Combine internal args with user-provided launcher args
+    all_launcher_args = internal_launcher_args + launcher_args
+
+    base_cmd = (
+        ["accelerate", "launch"] + all_launcher_args + ["-m", "axolotl.cli.train"]
+    )
+    if cfg_file:
+        base_cmd.append(cfg_file)
+
+    cmd = build_command(base_cmd, kwargs)
+    if use_exec:
+        # make sure to flush stdout and stderr before replacing the process
+        sys.stdout.flush()
+        sys.stderr.flush()
+        os.execvpe(cmd[0], cmd, os.environ)  # nosec B606
+    else:
+        subprocess.run(cmd, check=True)  # nosec B603
+
+
+def _launch_torchrun_training(
+    cfg_file: str,
+    kwargs: dict,
+    launcher_args: list[str] | None = None,
+    use_exec: bool = False,
+) -> None:
+    """Execute training via torchrun launcher."""
+    launcher_args = launcher_args or []
+
+    # Add default RDZV arguments if rdzv_endpoint is set
+    launcher_args = _add_default_rdzv_args(launcher_args)
+
+    base_cmd = ["torchrun"] + launcher_args + ["-m", "axolotl.cli.train"]
+    if cfg_file:
+        base_cmd.append(cfg_file)
+
+    cmd = build_command(base_cmd, kwargs)
+    if use_exec:
+        # make sure to flush stdout and stderr before replacing the process
+        sys.stdout.flush()
+        sys.stderr.flush()
+        os.execvpe(cmd[0], cmd, os.environ)  # nosec B606
+    else:
+        subprocess.run(cmd, check=True)  # nosec B603
+
+
+def _launch_python_training(cfg_file: str, kwargs: dict) -> None:
+    """Execute training via python launcher."""
+    from axolotl.cli.train import do_cli
+
+    do_cli(config=cfg_file, **kwargs)
diff --git a/src/axolotl/cli/vllm_serve.py b/src/axolotl/cli/vllm_serve.py
new file mode 100644
index 0000000000..a5281b2dfa
--- /dev/null
+++ b/src/axolotl/cli/vllm_serve.py
@@ -0,0 +1,122 @@
+"""
+CLI to start the vllm server for online RL
+"""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Union
+
+from trl.scripts.vllm_serve import ScriptArguments
+
+from axolotl.cli.config import load_cfg
+
+
+@dataclass
+class AxolotlScriptArguments(ScriptArguments):
+    """
+    Additional arguments for the VLLM server
+    """
+
+    reasoning_parser: str = field(default="", kw_only=True)
+    enable_reasoning: bool | None = field(default=None, kw_only=True)
+
+
+def do_vllm_serve(
+    config: Union[Path, str],
+    cli_args: dict,
+):
+    """
+    Starts the VLLM server for serving LLM models used for online RL
+
+    Args
+        :param cfg: Parsed doct of the YAML config
+        :param cli_args: dict of additional command-line arguments of type VllmServeCliArgs
+
+    Returns:
+        process_id: the process id of the started VLLM server
+    """
+    cfg = load_cfg(config)
+    model = cfg.base_model
+
+    # Determine serve module: explicit CLI/config > default (axolotl's LoRA-aware serve).
+    # We default to axolotl's serve module instead of TRL's because TRL's sends
+    # truncate_prompt_tokens which is unsupported in vLLM 0.17+.
+    serve_module = cli_args.get("serve_module") or getattr(
+        cfg.vllm, "serve_module", None
+    )
+    if serve_module is None:
+        serve_module = "axolotl.scripts.vllm_serve_lora"
+    vllm_serve_main = __import__(serve_module, fromlist=["main"]).main
+    tensor_parallel_size = 1
+    data_parallel_size = 1
+
+    if cli_args.get("tensor_parallel_size") or cfg.vllm.tensor_parallel_size:
+        tensor_parallel_size = (
+            cli_args.get("tensor_parallel_size") or cfg.vllm.tensor_parallel_size
+        )
+    if cli_args.get("data_parallel_size") or cfg.vllm.data_parallel_size:
+        data_parallel_size = (
+            cli_args.get("data_parallel_size") or cfg.vllm.data_parallel_size
+        )
+    host = cli_args.get("host") or cfg.vllm.host
+    port = cli_args.get("port") or cfg.vllm.port
+    gpu_memory_utilization = (
+        cli_args.get("gpu_memory_utilization") or cfg.vllm.gpu_memory_utilization
+    )
+    dtype = cli_args.get("dtype") or cfg.vllm.dtype
+    max_model_len = cli_args.get("max_model_len") or cfg.vllm.max_model_len
+    # Booleans check for None so an explicit CLI False can disable a config-
+    # enabled option (`cli or cfg` would let a falsy CLI value fall through).
+    cli_prefix = cli_args.get("enable_prefix_caching")
+    enable_prefix_caching = (
+        cfg.vllm.enable_prefix_caching if cli_prefix is None else cli_prefix
+    )
+    reasoning_parser = (
+        cli_args.get("reasoning_parser") or cfg.vllm.reasoning_parser or ""
+    )
+    cli_reasoning = cli_args.get("enable_reasoning")
+    enable_reasoning = (
+        cfg.vllm.enable_reasoning if cli_reasoning is None else cli_reasoning
+    ) or False
+
+    cli_enforce_eager = cli_args.get("enforce_eager")
+    cfg_enforce_eager = getattr(cfg.vllm, "enforce_eager", None)
+    raw_enforce_eager = (
+        cfg_enforce_eager if cli_enforce_eager is None else cli_enforce_eager
+    )
+    enforce_eager = bool(raw_enforce_eager) if raw_enforce_eager is not None else False
+    base_kwargs = dict(
+        model=model,
+        tensor_parallel_size=tensor_parallel_size,
+        data_parallel_size=data_parallel_size,
+        host=host,
+        port=port,
+        gpu_memory_utilization=gpu_memory_utilization,
+        dtype=dtype,
+        max_model_len=max_model_len,
+        enable_prefix_caching=enable_prefix_caching,
+        enforce_eager=enforce_eager,
+    )
+
+    # Use LoRAScriptArguments when serving with native LoRA support
+    if serve_module == "axolotl.scripts.vllm_serve_lora":
+        from axolotl.scripts.vllm_serve_lora import LoRAScriptArguments
+
+        lora_kwargs = {}
+        if hasattr(cfg, "lora_r") and cfg.lora_r:
+            lora_kwargs["max_lora_rank"] = cfg.lora_r
+        # Disable native LoRA in vLLM if not using vllm_lora_sync
+        # (merged weight sync via batch_update doesn't need vLLM LoRA mode)
+        if not getattr(cfg.trl, "vllm_lora_sync", False):
+            lora_kwargs["enable_lora"] = False
+        if getattr(cfg.vllm, "worker_extension_cls", None):
+            lora_kwargs["worker_extension_cls"] = cfg.vllm.worker_extension_cls
+        vllm_script_args = LoRAScriptArguments(**base_kwargs, **lora_kwargs)
+    else:
+        vllm_script_args = AxolotlScriptArguments(
+            **base_kwargs,
+            reasoning_parser=reasoning_parser,
+            enable_reasoning=enable_reasoning,
+        )
+
+    vllm_serve_main(vllm_script_args)
diff --git a/src/axolotl/common/architectures.py b/src/axolotl/common/architectures.py
new file mode 100644
index 0000000000..1f943852ae
--- /dev/null
+++ b/src/axolotl/common/architectures.py
@@ -0,0 +1,27 @@
+"""
+Common architecture specific constants
+"""
+
+MOE_ARCH_BLOCK = {
+    "dbrx": "DbrxFFN",
+    "jamba": "JambaSparseMoeBlock",
+    "jetmoe": [
+        "JetMoeMoA",
+        "JetMoeMoE",
+    ],
+    "mixtral": "MixtralSparseMoeBlock",
+    "qwen2_moe": "Qwen2MoeSparseMoeBlock",
+    "qwen3_moe": "Qwen3MoeSparseMoeBlock",
+    "qwen3_5_moe": "Qwen3_5MoeSparseMoeBlock",
+    "qwen3_vl_moe": "Qwen3VLMoeTextSparseMoeBlock",
+    "deepseek_v2": "DeepseekV2MoE",
+    "deepseek_v3": "DeepseekV3MoE",
+    "mistral4": "Mistral4MoE",
+    "gpt_oss": "GptOssDecoderLayer",
+    "lfm2_moe": "Lfm2MoeSparseMoeBlock",
+    "afmoe": "AfmoeMoE",
+    "glm4_moe": "Glm4MoeDecoderLayer",
+    "glm4_moe_lite": "Glm4MoeLiteDecoderLayer",
+    "glm_moe_dsa": "GlmMoeDsaDecoderLayer",
+    "nemotron_h": "NemotronHMoE",
+}
diff --git a/src/axolotl/common/cli.py b/src/axolotl/common/cli.py
deleted file mode 100644
index 636a23ba52..0000000000
--- a/src/axolotl/common/cli.py
+++ /dev/null
@@ -1,55 +0,0 @@
-"""
-shared module for cli specific things
-"""
-
-import logging
-from dataclasses import dataclass, field
-from typing import Optional
-
-import axolotl.monkeypatch.data.batch_dataset_fetcher  # pylint: disable=unused-import  # noqa: F401
-from axolotl.logging_config import configure_logging
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model, load_tokenizer
-
-configure_logging()
-LOG = logging.getLogger("axolotl.common.cli")
-
-
-@dataclass
-class TrainerCliArgs:
-    """
-    dataclass representing the various non-training arguments
-    """
-
-    debug: bool = field(default=False)
-    debug_text_only: bool = field(default=False)
-    debug_num_examples: int = field(default=5)
-    inference: bool = field(default=False)
-    merge_lora: bool = field(default=False)
-    prompter: Optional[str] = field(default=None)
-    shard: bool = field(default=False)
-
-
-@dataclass
-class PreprocessCliArgs:
-    """
-    dataclass representing arguments for preprocessing only
-    """
-
-    debug: bool = field(default=False)
-    debug_text_only: bool = field(default=False)
-    debug_num_examples: int = field(default=1)
-    prompter: Optional[str] = field(default=None)
-
-
-def load_model_and_tokenizer(
-    *,
-    cfg: DictDefault,
-    cli_args: TrainerCliArgs,
-):
-    LOG.info(f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}")
-    tokenizer = load_tokenizer(cfg)
-    LOG.info("loading model and (optionally) peft_config...")
-    model, _ = load_model(cfg, tokenizer, inference=cli_args.inference)
-
-    return model, tokenizer
diff --git a/src/axolotl/common/const.py b/src/axolotl/common/const.py
index fd34ad4694..8aae06e997 100644
--- a/src/axolotl/common/const.py
+++ b/src/axolotl/common/const.py
@@ -1,5 +1,3 @@
-"""
-Various shared constants
-"""
+"""Various shared constants"""
 
 DEFAULT_DATASET_PREPARED_PATH = "last_run_prepared"
diff --git a/src/axolotl/common/datasets.py b/src/axolotl/common/datasets.py
new file mode 100644
index 0000000000..b661e74c95
--- /dev/null
+++ b/src/axolotl/common/datasets.py
@@ -0,0 +1,146 @@
+"""Dataset loading utilities."""
+
+import math
+import random
+from dataclasses import dataclass
+
+from datasets import Dataset
+
+import axolotl.monkeypatch.data.batch_dataset_fetcher  # noqa: F401
+from axolotl.cli.args import PreprocessCliArgs, TrainerCliArgs
+from axolotl.loaders import load_processor, load_tokenizer
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.data import prepare_datasets, prepare_preference_datasets
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.enums import RLType
+from axolotl.utils.tokenization import check_dataset_labels
+
+LOG = get_logger(__name__)
+
+
+@dataclass
+class TrainDatasetMeta:
+    """Dataclass with fields for training and validation datasets and metadata."""
+
+    train_dataset: Dataset
+    eval_dataset: Dataset | None = None
+    total_num_steps: int | None = None
+
+
+def sample_dataset(dataset: Dataset, num_samples: int) -> Dataset:
+    """Randomly sample `num_samples` samples with replacement from `dataset`."""
+    return dataset.select(
+        [random.randrange(0, len(dataset) - 1) for _ in range(num_samples)]  # nosec
+    )
+
+
+@send_errors
+def load_datasets(
+    *,
+    cfg: DictDefault,
+    cli_args: PreprocessCliArgs | TrainerCliArgs | None = None,
+    debug: bool = False,
+) -> TrainDatasetMeta:
+    """Loads one or more training or evaluation datasets, calling
+    `axolotl.utils.data.prepare_datasets`. Optionally, logs out debug information.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        cli_args: Command-specific CLI arguments.
+        debug: Whether to print out tokenization of sample. This is duplicated in
+            `cfg` and `cli_args`, but is kept due to use in our Colab notebooks.
+
+    Returns:
+        Dataclass with fields for training and evaluation datasets and the computed
+            `total_num_steps`.
+    """
+    tokenizer = load_tokenizer(cfg)
+    processor = load_processor(cfg, tokenizer=tokenizer) if cfg.processor_type else None
+
+    train_dataset, eval_dataset, total_num_steps, prompters = prepare_datasets(
+        cfg,
+        tokenizer,
+        processor=processor,
+    )
+
+    if (
+        cfg.debug
+        or getattr(cli_args, "debug", False)
+        or getattr(cli_args, "debug_text_only", False)
+        or getattr(cli_args, "debug_num_examples", 0) > 0
+        or debug
+    ):
+        LOG.info("check_dataset_labels...")
+
+        num_examples = cli_args.debug_num_examples if cli_args else 1
+        text_only = cli_args.debug_text_only if cli_args else False
+        try:
+            train_samples = sample_dataset(train_dataset, num_examples)
+            check_dataset_labels(
+                train_samples,
+                tokenizer,
+                num_examples=num_examples,
+                text_only=text_only,
+            )
+        except AttributeError:
+            # can't sample iterable datasets
+            pass
+
+        LOG.info("printing prompters...")
+        for prompter in prompters:
+            LOG.info(prompter)
+
+    return TrainDatasetMeta(
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        total_num_steps=total_num_steps,
+    )
+
+
+@send_errors
+def load_preference_datasets(
+    *, cfg: DictDefault, cli_args: PreprocessCliArgs | TrainerCliArgs | None = None
+) -> TrainDatasetMeta:
+    """Loads one or more training or evaluation datasets for RL training using paired
+    preference data, calling `axolotl.utils.data.rl.prepare_preference_datasets`.
+    Optionally, logs out debug information.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        cli_args: Command-specific CLI arguments.
+
+    Returns:
+        Dataclass with fields for training and evaluation datasets and the computed
+        `total_num_steps`.
+    """
+    tokenizer = load_tokenizer(cfg)
+    train_dataset, eval_dataset = prepare_preference_datasets(cfg, tokenizer)
+
+    total_num_steps: int | None = None
+    if cfg.rl not in {RLType.GRPO, RLType.EBFT}:
+        total_num_steps = int(
+            math.ceil(len(train_dataset) * cfg.num_epochs / cfg.batch_size)
+        )
+
+    if ((cli_args and cli_args.debug) or cfg.debug) and cfg.rl != RLType.ORPO:
+        LOG.info("check_dataset_labels...")
+
+        num_examples = cli_args.debug_num_examples if cli_args else 1
+        text_only = cli_args.debug_text_only if cli_args else False
+
+        tokenizer = load_tokenizer(cfg)
+        train_samples = sample_dataset(train_dataset, num_examples)
+        check_dataset_labels(
+            dataset=train_samples,
+            tokenizer=tokenizer,
+            num_examples=num_examples,
+            text_only=text_only,
+            rl_mode=True,
+        )
+
+    return TrainDatasetMeta(
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        total_num_steps=total_num_steps,
+    )
diff --git a/src/axolotl/convert.py b/src/axolotl/convert.py
index 357e0ec50e..f8d8b25f46 100644
--- a/src/axolotl/convert.py
+++ b/src/axolotl/convert.py
@@ -1,6 +1,5 @@
 """Module containing File Reader, File Writer, Json Parser, and Jsonl Serializer classes"""
 
-
 import json
 import sys
 
@@ -68,9 +67,7 @@ def __init__(self, file_reader, file_writer, json_parser, jsonl_serializer):
         self.json_parser = json_parser
         self.jsonl_serializer = jsonl_serializer
 
-    def convert(
-        self, input_file_path, output_file_path
-    ):  # pylint: disable=unused-argument
+    def convert(self, input_file_path):
         content = self.file_reader.read(input_file_path)
         data = self.json_parser.parse(content)
         # data = [r for r in data if r["conversations"]]  # vicuna cleaned has rows with empty conversations
diff --git a/src/axolotl/utils/config/models/input/next/__init__.py b/src/axolotl/core/attention/__init__.py
similarity index 100%
rename from src/axolotl/utils/config/models/input/next/__init__.py
rename to src/axolotl/core/attention/__init__.py
diff --git a/src/axolotl/core/builders/__init__.py b/src/axolotl/core/builders/__init__.py
new file mode 100644
index 0000000000..5bd2444341
--- /dev/null
+++ b/src/axolotl/core/builders/__init__.py
@@ -0,0 +1,6 @@
+"""Trainer builder classes"""
+
+from .causal import HFCausalTrainerBuilder
+from .rl import HFRLTrainerBuilder
+
+__all__ = ["HFCausalTrainerBuilder", "HFRLTrainerBuilder"]
diff --git a/src/axolotl/core/builders/base.py b/src/axolotl/core/builders/base.py
new file mode 100644
index 0000000000..2c91f893d8
--- /dev/null
+++ b/src/axolotl/core/builders/base.py
@@ -0,0 +1,684 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base class for trainer builder"""
+
+import abc
+import importlib
+import sys
+from abc import abstractmethod
+from contextlib import suppress
+from pathlib import Path
+from typing import Any
+
+import torch
+from transformers import TrainerCallback
+from transformers.trainer_pt_utils import AcceleratorConfig
+
+from axolotl.integrations.base import PluginManager
+from axolotl.monkeypatch.trainer.lr import patch_trainer_get_lr
+from axolotl.telemetry.callbacks import TelemetryCallback
+from axolotl.telemetry.manager import TelemetryManager
+from axolotl.utils import (
+    is_comet_available,
+    is_mlflow_available,
+    is_opentelemetry_available,
+    is_trackio_available,
+)
+from axolotl.utils.callbacks import (
+    GCCallback,
+    SaveAxolotlConfigtoWandBCallback,
+    SaveModelOnFirstStepCallback,
+    SkipEvalOnResumeCallback,
+)
+from axolotl.utils.callbacks.profiler import PytorchProfilerCallback
+from axolotl.utils.distributed import build_parallelism_config
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.enums import CustomSupportedOptimizers
+
+LOG = get_logger(__name__)
+
+with suppress(ImportError):
+    import torch._dynamo
+
+
+class TrainerBuilderBase(abc.ABC):
+    """Base class for trainer builder."""
+
+    def __init__(self, cfg, model, tokenizer, processor=None):
+        self.cfg = cfg
+        self.model = model
+        self.tokenizer = tokenizer
+        self.processor = processor
+
+        self._train_dataset = None
+        self._eval_dataset = None
+        self._model_ref = None
+        self._peft_config = None
+
+        # If the model supports tagging, add the axolotl tag.
+        # This makes sure the tag is correctly pushed even if a user calls
+        # model.push_to_hub instead of trainer.push_to_hub.
+        if hasattr(model, "add_model_tags"):
+            model.add_model_tags(["axolotl"])
+
+        patch_trainer_get_lr()
+
+    @property
+    def model_ref(self):
+        return self._model_ref
+
+    @model_ref.setter
+    def model_ref(self, model):
+        self._model_ref = model
+
+    @property
+    def train_dataset(self):
+        return self._train_dataset
+
+    @train_dataset.setter
+    def train_dataset(self, dataset):
+        self._train_dataset = dataset
+
+    @property
+    def eval_dataset(self):
+        return self._eval_dataset
+
+    @eval_dataset.setter
+    def eval_dataset(self, dataset):
+        self._eval_dataset = dataset
+
+    @property
+    def peft_config(self):
+        return self._peft_config
+
+    @peft_config.setter
+    def peft_config(self, peft_config):
+        self._peft_config = peft_config
+
+    @abstractmethod
+    def build(self, total_num_steps):
+        pass
+
+    def get_callbacks(self) -> list[TrainerCallback]:
+        callbacks = []
+
+        plugin_manager = PluginManager.get_instance()
+        callbacks.extend(
+            plugin_manager.add_callbacks_pre_trainer(cfg=self.cfg, model=self.model)
+        )
+
+        if self.cfg.resume_from_checkpoint:
+            callbacks.append(SkipEvalOnResumeCallback())
+
+        gc_collect_steps = self.cfg.gc_collect_steps or self.cfg.gc_steps
+        if gc_collect_steps:
+            callbacks.append(GCCallback(gc_collect_steps=gc_collect_steps))
+
+        if self.cfg.dynamic_checkpoint and self.cfg.dynamic_checkpoint.enabled:
+            from axolotl.utils.callbacks.dynamic_checkpoint import (
+                DynamicCheckpointCallback,
+            )
+
+            callbacks.append(DynamicCheckpointCallback(self.cfg))
+
+        if self.cfg.use_wandb:
+            callbacks.append(
+                SaveAxolotlConfigtoWandBCallback(self.cfg.axolotl_config_path)
+            )
+        if self.cfg.use_mlflow and is_mlflow_available():
+            from axolotl.utils.callbacks.mlflow_ import (
+                SaveAxolotlConfigtoMlflowCallback,
+            )
+
+            callbacks.extend(
+                [
+                    SaveAxolotlConfigtoMlflowCallback(self.cfg.axolotl_config_path),
+                ]
+            )
+        if self.cfg.use_comet and is_comet_available():
+            from axolotl.utils.callbacks.comet_ import SaveAxolotlConfigtoCometCallback
+
+            callbacks.append(
+                SaveAxolotlConfigtoCometCallback(self.cfg.axolotl_config_path)
+            )
+        if self.cfg.use_trackio and is_trackio_available():
+            from axolotl.utils.callbacks.trackio_ import (
+                SaveAxolotlConfigtoTrackioCallback,
+            )
+
+            callbacks.append(
+                SaveAxolotlConfigtoTrackioCallback(self.cfg.axolotl_config_path)
+            )
+        if self.cfg.use_otel_metrics and is_opentelemetry_available():
+            from axolotl.utils.callbacks.opentelemetry import (
+                OpenTelemetryMetricsCallback,
+            )
+
+            callbacks.append(OpenTelemetryMetricsCallback(self.cfg))
+        if self.cfg.save_first_step:
+            callbacks.append(SaveModelOnFirstStepCallback())
+
+        if self.cfg.profiler_steps:
+            callbacks.append(
+                PytorchProfilerCallback(
+                    steps_to_profile=self.cfg.profiler_steps,
+                    profiler_steps_start=self.cfg.profiler_steps_start,
+                )
+            )
+
+        telemetry_manager = TelemetryManager.get_instance()
+        if telemetry_manager.enabled:
+            callbacks.append(TelemetryCallback())
+
+            # Report the fused RMSNorm+RoPE autotune selection + GPU identity so
+            # per-hardware tuning can be aggregated (mirrors scattermoe-lora).
+            if self.cfg.fused_attn_kernel or self.cfg.model_config_type in (
+                "gemma4",
+                "gemma4_text",
+                "gemma4_unified",
+                "gemma4_unified_text",
+            ):
+                from axolotl.kernels.autotune_telemetry import (
+                    FusedRopeAutotuneReportCallback,
+                )
+
+                callbacks.append(FusedRopeAutotuneReportCallback())
+
+        return callbacks
+
+    def get_post_trainer_create_callbacks(self, trainer):
+        """
+        Callbacks added after the trainer is created, usually b/c these need access to the trainer
+        """
+        callbacks = []
+        if self.cfg.plugins:
+            plugin_manager = PluginManager.get_instance()
+            callbacks.extend(
+                [
+                    cb
+                    for cb in plugin_manager.add_callbacks_post_trainer(
+                        self.cfg, trainer
+                    )
+                    if cb
+                ]
+            )
+        return callbacks
+
+    def hook_pre_create_training_args(self, training_arguments_kwargs):
+        # TODO
+        return training_arguments_kwargs
+
+    def hook_post_create_training_args(self, training_arguments):
+        # TODO
+        return training_arguments
+
+    def hook_pre_create_trainer(self, trainer_kwargs, trainer_cls):
+        # TODO
+        return trainer_kwargs, trainer_cls
+
+    def hook_post_create_trainer(self, trainer):
+        # TODO
+        return trainer
+
+    def _configure_warmup_and_logging(
+        self, total_num_steps: int, training_args_kwargs: dict
+    ):
+        warmup_steps: int | float = 0
+        warmup_ratio = 0.0
+        if self.cfg.warmup_steps is not None:
+            warmup_steps = self.cfg.warmup_steps
+        elif self.cfg.warmup_ratio is not None:
+            if total_num_steps:
+                warmup_steps = max(int(self.cfg.warmup_ratio * total_num_steps), 0)
+            else:
+                warmup_ratio = self.cfg.warmup_ratio
+        elif total_num_steps:
+            warmup_steps = min(int(0.03 * total_num_steps), 100)
+        else:
+            warmup_ratio = 0.03
+
+        # transformers v5
+        if warmup_ratio > 0.0 and warmup_steps == 0:
+            warmup_steps = warmup_ratio
+
+        if warmup_steps == 1:
+            warmup_steps = 2
+
+        if self.cfg.logging_steps is not None:
+            training_args_kwargs["logging_steps"] = self.cfg.logging_steps
+        else:
+            training_args_kwargs["logging_steps"] = (
+                500  # transformers defaults to 500
+                if not total_num_steps
+                else max(min(int(0.005 * total_num_steps), 10), 1)
+            )
+
+        training_args_kwargs["warmup_steps"] = warmup_steps
+
+    def _configure_precision_settings(self, training_args_kwargs: dict):
+        training_args_kwargs["fp16"] = (self.cfg.fp16 and not self.cfg.bf16) or False
+        training_args_kwargs["tf32"] = True if self.cfg.tf32 is True else False
+        if self.cfg.bf16 == "full":
+            training_args_kwargs["bf16_full_eval"] = True
+        else:
+            bf16 = self.cfg.bf16 or self.cfg.bfloat16
+            bf16 = bf16 if bf16 is not None else False
+            training_args_kwargs["bf16"] = bf16
+
+    def _configure_scheduler(self, training_args_kwargs: dict):
+        if self.cfg.lr_scheduler in ["one_cycle", "rex"]:
+            training_args_kwargs["lr_scheduler_type"] = "cosine"
+            training_args_kwargs["alternate_lr_scheduler_type"] = self.cfg.lr_scheduler
+        else:
+            training_args_kwargs["lr_scheduler_type"] = (
+                self.cfg.lr_scheduler if self.cfg.lr_scheduler else "cosine"
+            )
+        training_args_kwargs["lr_scheduler_kwargs"] = (
+            self.cfg.lr_scheduler_kwargs if self.cfg.lr_scheduler_kwargs else {}
+        )
+
+    def _configure_optimizer(self, training_args_kwargs: dict, trainer_kwargs: dict):
+        def _configure_custom_optimizer(
+            training_args_kwargs: dict, trainer_kwargs: dict
+        ):
+            # Common optimizer kwargs
+            optimizer_kwargs = {
+                "lr": training_args_kwargs["learning_rate"],
+                "weight_decay": training_args_kwargs["weight_decay"],
+            }
+
+            # Adam-specific kwargs
+            adam_kwargs: dict = {}
+            if training_args_kwargs.get("adam_beta1") and training_args_kwargs.get(
+                "adam_beta2"
+            ):
+                adam_kwargs["betas"] = (
+                    training_args_kwargs.get("adam_beta1"),
+                    training_args_kwargs.get("adam_beta2"),
+                )
+            if training_args_kwargs.get("adam_epsilon"):
+                adam_kwargs["eps"] = training_args_kwargs.get("adam_epsilon")
+
+            if self.cfg.optimizer == "muon":
+                _, device_mesh = build_parallelism_config(self.cfg)
+
+                if device_mesh is not None:
+                    from axolotl.contribs.mit.muon.dist_muon import (
+                        DistMuonOptimizerFactory,
+                    )
+
+                    optimizer_cls = DistMuonOptimizerFactory
+                    optimizer_kwargs["device_mesh"] = device_mesh
+                else:
+                    from axolotl.contribs.mit.muon import (
+                        MuonOptimizerFactory,
+                    )
+
+                    optimizer_cls = MuonOptimizerFactory
+
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "dion":
+                from axolotl.contribs.mit.dion import (
+                    DionOptimizerFactory,
+                )
+
+                optimizer_cls = DionOptimizerFactory
+                optimizer_kwargs["dion_lr"] = training_args_kwargs["dion_learning_rate"]
+                optimizer_kwargs["dion_mu"] = training_args_kwargs["dion_momentum"]
+                optimizer_kwargs.update(adam_kwargs)
+                _, device_mesh = build_parallelism_config(self.cfg)
+                if device_mesh is not None:
+                    optimizer_kwargs["device_mesh"] = device_mesh
+            elif self.cfg.optimizer == "sinkgd":
+                _, device_mesh = build_parallelism_config(self.cfg)
+
+                if device_mesh is not None:
+                    from axolotl.utils.optimizers.sinkgd import (
+                        DistSinkGDOptimizerFactory,
+                    )
+
+                    optimizer_cls = DistSinkGDOptimizerFactory
+                    optimizer_kwargs["device_mesh"] = device_mesh
+                else:
+                    from axolotl.utils.optimizers.sinkgd import SinkGDOptimizerFactory
+
+                    optimizer_cls = SinkGDOptimizerFactory
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "optimi_adamw":
+                from optimi import AdamW
+
+                optimizer_kwargs["foreach"] = False
+                optimizer_cls = AdamW
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "ao_adamw_fp8":
+                from torchao.optim.adam import AdamWFp8
+
+                optimizer_cls = AdamWFp8
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "adopt_adamw":
+                from axolotl.utils.optimizers.adopt import ADOPT
+
+                optimizer_cls = ADOPT
+                adam_kwargs["decouple"] = True
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "came_pytorch":
+                from came_pytorch import CAME
+
+                optimizer_cls = CAME
+
+                beta1 = training_args_kwargs.get("adam_beta1", 0.9)
+                beta2 = training_args_kwargs.get("adam_beta2", 0.999)
+                beta3 = training_args_kwargs.get("adam_beta3", 0.9999)
+                eps1 = training_args_kwargs.get("adam_epsilon", 1e-30)
+                eps2 = training_args_kwargs.get("adam_epsilon2", 1e-16)
+                adam_kwargs["betas"] = (beta1, beta2, beta3)
+                adam_kwargs["eps"] = (eps1, eps2)
+
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "q_galore_adamw8bit":
+                from axolotl.utils.optimizers.qgalore import (
+                    build_qgalore_param_groups,
+                    patch_q_galore_for_modern_bnb,
+                )
+
+                patch_q_galore_for_modern_bnb()
+                from q_galore_torch import QGaLoreAdamW8bit
+
+                optimizer_cls = QGaLoreAdamW8bit
+                optimizer_kwargs["params"] = build_qgalore_param_groups(
+                    self.model,
+                    self.cfg.optim_target_modules,
+                    rank=self.cfg.qgalore_rank,
+                    update_proj_gap=self.cfg.qgalore_update_proj_gap,
+                    scale=self.cfg.qgalore_scale,
+                    proj_type=self.cfg.qgalore_proj_type,
+                    proj_quant=self.cfg.qgalore_proj_quant,
+                    proj_bits=self.cfg.qgalore_proj_bits,
+                    proj_group_size=self.cfg.qgalore_proj_group_size,
+                    cos_threshold=self.cfg.qgalore_cos_threshold,
+                    gamma_proj=self.cfg.qgalore_gamma_proj,
+                    queue_size=self.cfg.qgalore_queue_size,
+                )
+
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "flash_adamw":
+                from flashoptim import FlashAdamW
+
+                optimizer_cls = FlashAdamW
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "flash_adam":
+                from flashoptim import FlashAdam
+
+                optimizer_cls = FlashAdam
+                optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "flash_sgd":
+                from flashoptim import FlashSGD
+
+                optimizer_cls = FlashSGD
+            elif self.cfg.optimizer == "flash_sgdw":
+                from flashoptim import FlashSGDW
+
+                optimizer_cls = FlashSGDW
+            elif self.cfg.optimizer == "flash_lion":
+                from flashoptim import FlashLion
+
+                optimizer_cls = FlashLion
+                if "betas" in adam_kwargs:
+                    optimizer_kwargs["betas"] = adam_kwargs["betas"]
+            else:
+                raise ValueError(
+                    f"Unhandled optimizer: {self.cfg.optimizer}. Please raise an Issue."
+                )
+
+            # Parse any additional optimizer args from config
+            if self.cfg.optim_args:
+                if isinstance(self.cfg.optim_args, dict):
+                    optimizer_kwargs.update(self.cfg.optim_args)
+                else:
+                    # Parse string format "key1=value1,key2=value2"
+                    for mapping in self.cfg.optim_args.replace(" ", "").split(","):
+                        key, value = mapping.split("=")
+                        optimizer_kwargs[key] = value
+
+            # Note: This is not used in training_args_kwargs, but in trainer_kwargs
+            trainer_kwargs["optimizer_cls_and_kwargs"] = (
+                optimizer_cls,
+                optimizer_kwargs,
+            )
+
+        # Handle custom optimizer
+        custom_supported_optimizers = [opt.value for opt in CustomSupportedOptimizers]
+        if self.cfg.optimizer in custom_supported_optimizers:
+            _configure_custom_optimizer(training_args_kwargs, trainer_kwargs)
+        else:
+            # Use transformers' optimizer
+            training_args_kwargs["optim"] = self.cfg.optimizer
+
+            # Parse any additional optimizer args from config
+            if self.cfg.optim_args:
+                if isinstance(self.cfg.optim_args, dict):
+                    optim_args = ",".join(
+                        [f"{key}={value}" for key, value in self.cfg.optim_args.items()]
+                    )
+                else:
+                    optim_args = self.cfg.optim_args
+                training_args_kwargs["optim_args"] = optim_args
+
+            if (
+                self.cfg.optimizer == "adamw_anyprecision"
+                and Path(self.cfg.torchdistx_path).exists()
+            ):
+                sys.path.append(self.cfg.torchdistx_path)
+                importlib.import_module("torchdistx")
+
+    def _configure_hub_parameters(self, training_args_kwargs: dict):
+        if self.cfg.hub_model_id:
+            training_args_kwargs["hub_model_id"] = self.cfg.hub_model_id
+            training_args_kwargs["push_to_hub"] = True
+            training_args_kwargs["hub_private_repo"] = True
+            training_args_kwargs["hub_always_push"] = True
+
+            if self.cfg.hub_strategy:
+                training_args_kwargs["hub_strategy"] = self.cfg.hub_strategy
+
+            if self.cfg.hub_revision:
+                training_args_kwargs["hub_revision"] = self.cfg.hub_revision
+
+    def _configure_save_and_eval_strategy(self, training_args_kwargs: dict):
+        # save_strategy and save_steps
+        if self.cfg.save_steps:
+            training_args_kwargs["save_strategy"] = "steps"
+            training_args_kwargs["save_steps"] = self.cfg.save_steps
+        elif self.cfg.save_strategy:
+            training_args_kwargs["save_strategy"] = self.cfg.save_strategy
+        else:
+            # default to saving each epoch if not defined
+            training_args_kwargs["save_strategy"] = "epoch"
+
+        training_args_kwargs["save_total_limit"] = (
+            self.cfg.save_total_limit if self.cfg.save_total_limit else 4
+        )
+
+        # eval_strategy and eval_steps
+        if not self.eval_dataset and self.cfg.val_set_size == 0:
+            # do not eval if no eval_dataset and val_set_size=0
+            training_args_kwargs["eval_strategy"] = "no"
+        elif self.cfg.eval_steps:
+            training_args_kwargs["eval_strategy"] = "steps"
+            training_args_kwargs["eval_steps"] = self.cfg.eval_steps
+            training_args_kwargs["eval_on_start"] = True
+        elif self.cfg.eval_strategy:
+            training_args_kwargs["eval_strategy"] = self.cfg.eval_strategy
+            training_args_kwargs["eval_on_start"] = True
+
+    def _configure_reporting(self, training_args_kwargs: dict):
+        report_to = []
+        if self.cfg.use_wandb:
+            report_to.append("wandb")
+        if self.cfg.use_mlflow:
+            report_to.append("mlflow")
+        if self.cfg.use_tensorboard:
+            report_to.append("tensorboard")
+        if self.cfg.use_comet:
+            report_to.append("comet_ml")
+        if self.cfg.use_trackio:
+            report_to.append("trackio")
+
+        training_args_kwargs["report_to"] = report_to
+
+        if self.cfg.use_wandb:
+            training_args_kwargs["run_name"] = self.cfg.wandb_name
+        elif self.cfg.use_mlflow:
+            training_args_kwargs["run_name"] = self.cfg.mlflow_run_name
+        elif self.cfg.use_trackio:
+            training_args_kwargs["run_name"] = self.cfg.trackio_run_name
+        else:
+            training_args_kwargs["run_name"] = None
+
+    def _configure_torch_compile(self, training_args_kwargs: dict):
+        if self.cfg.torch_compile and getattr(torch, "_dynamo", None):
+            torch._dynamo.config.suppress_errors = True
+            torch._dynamo.config.accumulated_cache_size_limit = 256
+            training_args_kwargs["torch_compile"] = self.cfg.torch_compile
+            if self.cfg.torch_compile_backend:
+                training_args_kwargs["torch_compile_backend"] = (
+                    self.cfg.torch_compile_backend
+                )
+            if self.cfg.torch_compile_mode:
+                training_args_kwargs["torch_compile_mode"] = self.cfg.torch_compile_mode
+
+    def _configure_accelerator_config(self, training_args_kwargs: dict):
+        if self.cfg.accelerator_config:
+            training_args_kwargs["accelerator_config"] = AcceleratorConfig(
+                **self.cfg.accelerator_config
+            )
+        else:
+            training_args_kwargs["accelerator_config"] = AcceleratorConfig()
+
+    def _configure_gradient_checkpointing(self, training_args_kwargs: dict):
+        if self.cfg.layer_offloading:
+            training_args_kwargs["layer_offloading"] = True
+        if self.cfg.activation_offloading == "hidden_states":
+            training_args_kwargs["gradient_checkpointing"] = True
+            gc_kwargs = dict(self.cfg.gradient_checkpointing_kwargs or {})
+            training_args_kwargs["gradient_checkpointing_kwargs"] = gc_kwargs
+            if gc_kwargs["use_reentrant"] is False:
+                training_args_kwargs["activation_offloading"] = (
+                    self.cfg.activation_offloading
+                )
+        elif self.cfg.activation_offloading:
+            # TRL offloader replaces HF recompute (re-added for full finetune in the
+            # model loader), so disable HF checkpointing and pass the mode through.
+            training_args_kwargs["gradient_checkpointing"] = False
+            training_args_kwargs["activation_offloading"] = (
+                self.cfg.activation_offloading
+            )
+        elif self.cfg.gradient_checkpointing is not None:
+            training_args_kwargs["gradient_checkpointing"] = (
+                self.cfg.gradient_checkpointing
+            )
+            if self.cfg.gradient_checkpointing_kwargs is not None:
+                training_args_kwargs["gradient_checkpointing_kwargs"] = (
+                    self.cfg.gradient_checkpointing_kwargs
+                )
+            else:
+                training_args_kwargs["gradient_checkpointing_kwargs"] = {
+                    "use_reentrant": False
+                }
+
+    def _set_base_training_args(
+        self, total_num_steps
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        training_args_kwargs: dict[str, Any] = {}
+        trainer_kwargs: dict[str, Any] = {}
+
+        self._configure_warmup_and_logging(total_num_steps, training_args_kwargs)
+        self._configure_precision_settings(training_args_kwargs)
+        self._configure_save_and_eval_strategy(training_args_kwargs)
+        self._configure_gradient_checkpointing(training_args_kwargs)
+
+        # set arg into trainer_args_kwargs with same name if value not None
+        for arg in [
+            # optim/scheduler
+            "adam_beta1",
+            "adam_beta2",
+            "adam_beta3",
+            "adam_epsilon",
+            "adam_epsilon2",
+            "cosine_min_lr_ratio",
+            "cosine_constant_lr_ratio",
+            "optim_target_modules",
+            # trainer
+            "max_grad_norm",
+            "dataloader_num_workers",
+            "dataloader_pin_memory",
+            "dataloader_prefetch_factor",
+            "gradient_accumulation_steps",
+            "learning_rate",
+            "embedding_lr",
+            "embedding_lr_scale",
+            "lr_groups",
+            "loraplus_lr_ratio",
+            "loraplus_lr_embedding",
+            "output_dir",
+            "save_only_model",
+            "weight_decay",
+            "seed",
+            "dion_momentum",
+            "dion_rank_fraction",
+            "dion_rank_multiple_of",
+            "dataset_num_proc",
+            # memory management
+            "torch_empty_cache_steps",
+        ]:
+            if hasattr(self.cfg, arg) and getattr(self.cfg, arg) is not None:
+                training_args_kwargs[arg] = getattr(self.cfg, arg)
+
+        arg_map = {
+            "dion_learning_rate": "dion_lr",
+            "include_num_input_tokens_seen": "include_tokens_per_second",
+        }
+        for kwarg, cfg_arg in arg_map.items():
+            if hasattr(self.cfg, cfg_arg) and getattr(self.cfg, cfg_arg) is not None:
+                training_args_kwargs[kwarg] = getattr(self.cfg, cfg_arg)
+
+        training_args_kwargs["per_device_train_batch_size"] = self.cfg.micro_batch_size
+        training_args_kwargs["average_tokens_across_devices"] = False
+
+        if self.cfg.eval_batch_size:
+            training_args_kwargs["per_device_eval_batch_size"] = (
+                self.cfg.eval_batch_size
+            )
+
+        training_args_kwargs["include_tkps"] = self.cfg.include_tkps
+        training_args_kwargs["max_steps"] = self.cfg.max_steps or total_num_steps or -1
+        training_args_kwargs["num_train_epochs"] = self.cfg.num_epochs
+
+        # max_length is not used in CausalTrainer
+        if self.cfg.reward_model or self.cfg.rl:
+            training_args_kwargs["max_length"] = self.cfg.sequence_len
+
+        if self.cfg.fsdp_config or self.cfg.fsdp:
+            training_args_kwargs["fsdp_config"] = self.cfg.fsdp_config
+            training_args_kwargs["fsdp"] = self.cfg.fsdp if self.cfg.fsdp else True
+
+        self._configure_reporting(training_args_kwargs)
+        self._configure_hub_parameters(training_args_kwargs)
+        self._configure_scheduler(training_args_kwargs)
+        self._configure_optimizer(training_args_kwargs, trainer_kwargs)
+        self._configure_torch_compile(training_args_kwargs)
+        self._configure_accelerator_config(training_args_kwargs)
+
+        return training_args_kwargs, trainer_kwargs
diff --git a/src/axolotl/core/builders/causal.py b/src/axolotl/core/builders/causal.py
new file mode 100644
index 0000000000..6f7866720f
--- /dev/null
+++ b/src/axolotl/core/builders/causal.py
@@ -0,0 +1,607 @@
+"""Builder for causal trainers"""
+
+import inspect
+import math
+import os
+from pathlib import Path
+from typing import Type, Union
+
+import transformers
+from transformers import (
+    DataCollatorWithFlattening,
+    EarlyStoppingCallback,
+    Trainer,
+)
+from trl.trainer.reward_trainer import DataCollatorForPreference
+
+from axolotl.core.builders.base import TrainerBuilderBase
+from axolotl.core.trainers import (
+    AxolotlMambaTrainer,
+    AxolotlPRMTrainer,
+    AxolotlRewardTrainer,
+    AxolotlTrainer,
+)
+from axolotl.integrations.base import PluginManager
+from axolotl.monkeypatch.multipack import SUPPORTED_MULTIPACK_MODEL_TYPES
+from axolotl.monkeypatch.relora import ReLoRACallback
+from axolotl.processing_strategies import get_processing_strategy
+from axolotl.utils import is_comet_available, is_mlflow_available
+from axolotl.utils.callbacks import (
+    LossWatchDogCallback,
+    bench_eval_callback_factory,
+    causal_lm_bench_eval_callback_factory,
+    colab_inference_post_train_callback,
+    log_prediction_callback_factory,
+)
+from axolotl.utils.callbacks.lisa import lisa_callback_factory
+from axolotl.utils.callbacks.qat import QATCallback
+from axolotl.utils.callbacks.tokens_per_second import TokensPerSecondCallback
+from axolotl.utils.chat_templates import get_chat_template_from_config
+from axolotl.utils.collators import (
+    BatchSamplerDataCollatorForSeq2Seq,
+    DataCollatorForSeq2Seq,
+    MambaDataCollator,
+    V2BatchSamplerDataCollatorForSeq2Seq,
+)
+from axolotl.utils.collators.mm_chat import MultiModalChatDataCollator
+from axolotl.utils.import_helper import get_cls_from_module_str
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_MM_NUM_WORKERS_WARNED: set = set()
+
+
+def _warn_if_num_workers_zero_for_mm(cfg, log) -> None:
+    if not getattr(cfg, "processor_type", None):
+        return
+    if getattr(cfg, "dataloader_num_workers", None) not in (None, 0):
+        return
+    if getattr(cfg, "train_on_inputs", False):
+        return
+    if "mm_num_workers_zero" in _MM_NUM_WORKERS_WARNED:
+        return
+    _MM_NUM_WORKERS_WARNED.add("mm_num_workers_zero")
+    log.warning(
+        "Increase dataloader_num_workers to speed up multimodal training with assistant-only loss masking (currently dataloader_num_workers=0)."
+    )
+
+
+class HFCausalTrainerBuilder(TrainerBuilderBase):
+    """
+    Build the HuggingFace training args/trainer for causal models and reward modeling
+    using TRL.
+    """
+
+    def get_callbacks(self):
+        callbacks = super().get_callbacks()
+
+        if self.cfg.relora:
+            callbacks.append(ReLoRACallback(self.cfg))
+
+        # TODO: check if can move to base class
+        if self.cfg.loss_watchdog_threshold is not None:
+            callbacks.append(LossWatchDogCallback(self.cfg))
+
+        if self.cfg.qat:
+            callbacks.append(QATCallback(self.cfg.qat))
+
+        if self.cfg.include_tkps:
+            callbacks.append(
+                TokensPerSecondCallback(
+                    resume_from_checkpoint=self.cfg.resume_from_checkpoint,
+                )
+            )
+        return callbacks
+
+    def get_post_trainer_create_callbacks(self, trainer):
+        callbacks = []
+        if self.cfg.use_wandb and self.cfg.eval_table_size > 0:
+            LogPredictionCallback = log_prediction_callback_factory(
+                trainer, self.tokenizer, "wandb"
+            )
+            callbacks.append(LogPredictionCallback(self.cfg))
+        if (
+            self.cfg.use_mlflow
+            and is_mlflow_available()
+            and self.cfg.eval_table_size > 0
+        ):
+            LogPredictionCallback = log_prediction_callback_factory(
+                trainer, self.tokenizer, "mlflow"
+            )
+            callbacks.append(LogPredictionCallback(self.cfg))
+        if self.cfg.use_comet and is_comet_available() and self.cfg.eval_table_size > 0:
+            LogPredictionCallback = log_prediction_callback_factory(
+                trainer, self.tokenizer, "comet_ml"
+            )
+            callbacks.append(LogPredictionCallback(self.cfg))
+
+        if self.cfg.do_bench_eval:
+            callbacks.append(bench_eval_callback_factory(trainer, self.tokenizer))
+        if self.cfg.do_causal_lm_eval:
+            CausalLMBenchEvalCallback = causal_lm_bench_eval_callback_factory(
+                trainer, self.tokenizer
+            )
+            callbacks.append(CausalLMBenchEvalCallback(self.cfg))
+
+        if self.cfg.early_stopping_patience:
+            early_stop_cb = EarlyStoppingCallback(
+                self.cfg.early_stopping_patience,
+            )
+            callbacks.append(early_stop_cb)
+
+        if self.cfg.lisa_step_interval and self.cfg.lisa_n_layers:
+            callbacks.append(lisa_callback_factory(trainer))
+
+        if any("COLAB_" in key for key in os.environ):
+            ColabCallback = colab_inference_post_train_callback(trainer)
+            callbacks.append(ColabCallback(self.cfg))
+
+        if getattr(self.cfg, "generate_samples", False):
+            from axolotl.utils.callbacks.generation import SFTGenerationCallback
+
+            callbacks.append(SFTGenerationCallback(trainer))
+            LOG.info("SFT sample generation enabled")
+
+        callbacks.extend(super().get_post_trainer_create_callbacks(trainer=trainer))
+        return callbacks
+
+    def _get_trainer_cls(self):
+        """
+        Gets the trainer class for the given configuration.
+        """
+        if self.cfg.plugins:
+            plugin_manager = PluginManager.get_instance()
+            trainer_cls = plugin_manager.get_trainer_cls(self.cfg)
+            if trainer_cls:
+                return trainer_cls
+        if self.cfg.model_config_type == "mamba":
+            return AxolotlMambaTrainer
+        if self.cfg.reward_model:
+            return AxolotlRewardTrainer
+        if self.cfg.process_reward_model:
+            return AxolotlPRMTrainer
+
+        if self.cfg.trainer_cls:
+            # override the trainer cls
+            try:
+                trainer_cls = get_cls_from_module_str(self.cfg.trainer_cls)
+                LOG.debug(f"Using custom trainer class: {self.cfg.trainer_cls}")
+                return trainer_cls
+            except (ImportError, AttributeError, ValueError) as e:
+                raise ValueError(
+                    f"Failed to load custom trainer class '{self.cfg.trainer_cls}': {e}"
+                ) from e
+
+        return AxolotlTrainer
+
+    def build(self, total_num_steps):
+        from axolotl.core.training_args import (
+            AxolotlPRMConfig,
+            AxolotlRewardConfig,
+            AxolotlTrainingArguments,
+        )
+
+        training_arguments_kwargs, trainer_kwargs = self._set_base_training_args(
+            total_num_steps
+        )
+        if self.cfg.adapter == "qlora":
+            training_arguments_kwargs["qlora"] = True
+
+        # deepspeed
+        if self.cfg.deepspeed:
+            training_arguments_kwargs["deepspeed"] = self.cfg.deepspeed
+
+        if self.cfg.lr_quadratic_warmup is not None:
+            training_arguments_kwargs["lr_quadratic_warmup"] = (
+                self.cfg.lr_quadratic_warmup
+            )
+
+        if self.cfg.dataloader_drop_last is not None:
+            training_arguments_kwargs["dataloader_drop_last"] = (
+                self.cfg.dataloader_drop_last
+            )
+        elif self.cfg.sample_packing and self.cfg.eval_sample_packing is False:
+            training_arguments_kwargs["dataloader_drop_last"] = True
+
+        if self.cfg.remove_unused_columns is not None:
+            training_arguments_kwargs["remove_unused_columns"] = (
+                self.cfg.remove_unused_columns
+            )
+
+        if self.cfg.do_bench_eval:
+            training_arguments_kwargs["do_bench_eval"] = self.cfg.do_bench_eval
+            if self.cfg.bench_dataset:
+                training_arguments_kwargs["bench_dataset"] = self.cfg.bench_dataset
+        if self.cfg.do_causal_lm_eval:
+            training_arguments_kwargs["do_causal_lm_eval"] = self.cfg.do_causal_lm_eval
+        if self.cfg.metric_for_best_model:
+            training_arguments_kwargs["metric_for_best_model"] = (
+                self.cfg.metric_for_best_model
+            )
+        if self.cfg.greater_is_better:
+            training_arguments_kwargs["greater_is_better"] = self.cfg.greater_is_better
+
+        # DDP Config
+        if self.cfg.ddp_timeout:
+            training_arguments_kwargs["ddp_timeout"] = self.cfg.ddp_timeout
+        # see https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html
+        if self.cfg.ddp_bucket_cap_mb:
+            training_arguments_kwargs["ddp_bucket_cap_mb"] = self.cfg.ddp_bucket_cap_mb
+        if self.cfg.ddp_broadcast_buffers is not None:
+            training_arguments_kwargs["ddp_broadcast_buffers"] = (
+                self.cfg.ddp_broadcast_buffers
+            )
+
+        # these are all the "standard" kwargs that are def used
+        training_arguments_kwargs["max_seq_length"] = self.cfg.sequence_len
+
+        if self.cfg.auto_find_batch_size is not None:
+            training_arguments_kwargs["auto_find_batch_size"] = (
+                self.cfg.auto_find_batch_size
+            )
+
+        training_arguments_kwargs["eval_accumulation_steps"] = (
+            self.cfg.gradient_accumulation_steps
+        )
+
+        training_arguments_kwargs["load_best_model_at_end"] = (
+            (
+                self.cfg.load_best_model_at_end is not False
+                or self.cfg.early_stopping_patience
+            )
+            and (
+                (not self.cfg.test_datasets and self.cfg.val_set_size > 0)
+                or (self.cfg.test_datasets and self.cfg.val_set_size == 0)
+            )
+            and self.cfg.save_steps
+            and self.cfg.eval_steps
+            and self.cfg.save_steps % self.cfg.eval_steps == 0
+        ) or False
+
+        # handle ddp
+        ddp_find_unused_parameters = None
+        if self.cfg.ddp:
+            ddp_find_unused_parameters = bool(self.cfg.ddp_find_unused_parameters)
+        training_arguments_kwargs["ddp_find_unused_parameters"] = (
+            ddp_find_unused_parameters
+        )
+
+        if self.cfg.group_by_length:
+            training_arguments_kwargs["train_sampling_strategy"] = "group_by_length"
+        training_arguments_kwargs["curriculum_sampling"] = self.cfg.curriculum_sampling
+
+        training_arguments_kwargs["sample_packing"] = bool(self.cfg.sample_packing)
+        training_arguments_kwargs["sample_packing_drop_attention_mask"] = (
+            self.cfg.attn_decontaminates_packing
+        )
+        training_arguments_kwargs["multipack_real_batches"] = (
+            self.cfg.multipack_real_batches
+            if self.cfg.multipack_real_batches is not None
+            else not self.cfg.attn_supports_packing
+        )
+        training_arguments_kwargs["eval_sample_packing"] = bool(
+            self.cfg.eval_sample_packing
+        )
+        if self.cfg.sample_packing_sequentially is not None:
+            training_arguments_kwargs["sample_packing_sequentially"] = (
+                self.cfg.sample_packing_sequentially
+            )
+        if self.cfg.sample_packing_bin_size is not None:
+            training_arguments_kwargs["sample_packing_bin_size"] = (
+                self.cfg.sample_packing_bin_size
+            )
+        if self.cfg.sample_packing_group_size is not None:
+            training_arguments_kwargs["sample_packing_group_size"] = (
+                self.cfg.sample_packing_group_size
+            )
+        if self.cfg.sample_packing_eff_est:
+            training_arguments_kwargs["sample_packing_efficiency"] = (
+                self.cfg.sample_packing_eff_est
+            )
+
+        if self.cfg.relora and self.cfg.jagged_restart_steps:
+            if self.cfg.relora_prune_ratio is not None:
+                training_arguments_kwargs["relora_prune_ratio"] = (
+                    self.cfg.relora_prune_ratio
+                )
+            if self.cfg.relora_prune_method:
+                training_arguments_kwargs["relora_prune_method"] = (
+                    self.cfg.relora_prune_method
+                )
+
+        if self.cfg.jagged_restart_steps:
+            training_arguments_kwargs["jagged_restart_steps"] = (
+                self.cfg.jagged_restart_steps
+            )
+            if self.cfg.jagged_restart_warmup_steps:
+                training_arguments_kwargs["jagged_restart_warmup_steps"] = (
+                    self.cfg.jagged_restart_warmup_steps
+                )
+            if self.cfg.jagged_restart_anneal_steps:
+                training_arguments_kwargs["jagged_restart_anneal_steps"] = (
+                    self.cfg.jagged_restart_anneal_steps
+                )
+
+        if self.cfg.lisa_step_interval and self.cfg.lisa_n_layers:
+            training_arguments_kwargs["lisa_n_layers"] = self.cfg.lisa_n_layers
+            training_arguments_kwargs["lisa_step_interval"] = (
+                self.cfg.lisa_step_interval
+            )
+            training_arguments_kwargs["lisa_layers_attribute"] = (
+                self.cfg.lisa_layers_attribute
+            )
+
+        training_arguments_kwargs = self.hook_pre_create_training_args(
+            training_arguments_kwargs
+        )
+        training_arguments_kwargs["model_type"] = self.cfg.model_config_type
+        training_arguments_kwargs["pretraining"] = bool(self.cfg.pretraining_dataset)
+        if self.cfg.chat_template:
+            training_arguments_kwargs["chat_template"] = get_chat_template_from_config(
+                cfg=self.cfg,
+                tokenizer=self.tokenizer,
+            )
+
+        if self.cfg.neftune_noise_alpha is not None:
+            training_arguments_kwargs["neftune_noise_alpha"] = (
+                self.cfg.neftune_noise_alpha
+            )
+
+        if self.cfg.image_size:
+            training_arguments_kwargs["image_size"] = self.cfg.image_size
+        if self.cfg.image_resize_algorithm:
+            training_arguments_kwargs["image_resize_algorithm"] = (
+                self.cfg.image_resize_algorithm
+            )
+
+        if self.cfg.plugins:
+            plugin_manager = PluginManager.get_instance()
+            plugin_training_args = plugin_manager.get_training_args(self.cfg)
+            if plugin_training_args:
+                training_arguments_kwargs.update(plugin_training_args)
+
+        if self.cfg.reward_model:
+            training_args_cls = AxolotlRewardConfig
+            if self.cfg.center_rewards_coefficient is not None:
+                training_arguments_kwargs["center_rewards_coefficient"] = (
+                    self.cfg.center_rewards_coefficient
+                )
+        elif self.cfg.process_reward_model:
+            training_args_cls = AxolotlPRMConfig
+        else:
+            training_args_cls = AxolotlTrainingArguments
+        training_args = training_args_cls(
+            **training_arguments_kwargs,
+        )
+        training_args = self.hook_post_create_training_args(training_args)
+
+        # unset run_name so wandb sets up experiment names
+        if self.cfg.use_wandb and training_args.run_name == training_args.output_dir:
+            training_args.run_name = None
+
+        data_collator_kwargs = {
+            "padding": True,  # True/"longest" is the default
+        }
+        multiple = getattr(self.cfg, "pad_to_multiple_of", None) or 64
+        if self.cfg.pad_to_sequence_len:
+            data_collator_kwargs["pad_to_multiple_of"] = multiple * math.ceil(
+                self.cfg.sequence_len / multiple
+            )
+        elif self.cfg.pad_to_sequence_len is None:
+            # A100 is best at 64, while others at 8. Let's use the larger so we don't have to check
+            # https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
+            data_collator_kwargs["pad_to_multiple_of"] = multiple
+
+        if self.cfg.use_eaft:
+            from functools import partial
+
+            from axolotl.monkeypatch.loss.eaft import eaft_loss
+
+            configured_eaft_loss = partial(
+                eaft_loss,
+                alpha=self.cfg.eaft_alpha if self.cfg.eaft_alpha is not None else 1.0,
+                k=self.cfg.eaft_k if self.cfg.eaft_k is not None else 20,
+            )
+            trainer_kwargs["compute_loss_func"] = configured_eaft_loss
+
+        trainer_cls = self._get_trainer_cls()
+
+        trainer_kwargs, trainer_cls = self.hook_pre_create_trainer(
+            trainer_kwargs, trainer_cls
+        )
+        if eval_data_collator := self.build_collator(
+            training_args, is_eval=True, **data_collator_kwargs
+        ):
+            if not (self.cfg.reward_model or self.cfg.process_reward_model):
+                trainer_kwargs["eval_data_collator"] = eval_data_collator
+        if not (self.cfg.reward_model or self.cfg.process_reward_model):
+            trainer_kwargs["bench_data_collator"] = transformers.DataCollatorForSeq2Seq(
+                self.tokenizer,
+                return_tensors="pt",
+                **data_collator_kwargs,
+            )
+        sig = inspect.signature(trainer_cls)
+        if "processing_class" in sig.parameters or issubclass(trainer_cls, Trainer):
+            trainer_kwargs["processing_class"] = self.tokenizer
+        elif "tokenizer" in sig.parameters:
+            trainer_kwargs["tokenizer"] = self.tokenizer
+
+        if (
+            trainer_cls not in [AxolotlRewardTrainer, AxolotlPRMTrainer]
+            and self.cfg.datasets is not None
+        ):
+            trainer_kwargs["dataset_tags"] = [
+                d["path"] for d in self.cfg.datasets if not Path(d["path"]).is_dir()
+            ]
+        # TRL's RewardTrainer validates num_labels=1 on pre-loaded models; ensure the
+        # config reflects this regardless of how the model was instantiated.
+        if (
+            self.cfg.reward_model
+            and getattr(self.model.config, "num_labels", None) != 1
+        ):
+            self.model.config.num_labels = 1
+        trainer = trainer_cls(
+            model=self.model,
+            train_dataset=self.train_dataset,
+            eval_dataset=self.eval_dataset,
+            args=training_args,
+            data_collator=self.build_collator(training_args, **data_collator_kwargs),
+            callbacks=self.get_callbacks(),
+            **trainer_kwargs,
+        )
+        trainer = self.hook_post_create_trainer(trainer)
+        # if the trainer has the `axolotl_cfg` property, set it
+        if hasattr(trainer, "axolotl_cfg"):
+            trainer.axolotl_cfg = self.cfg
+        for callback in self.get_post_trainer_create_callbacks(trainer):
+            trainer.add_callback(callback)
+
+        if self.cfg.deepspeed and self.cfg.sample_packing:
+            trainer.accelerator.state.deepspeed_plugin.deepspeed_config[
+                "train_micro_batch_size_per_gpu"
+            ] = self.cfg.micro_batch_size
+
+        return trainer
+
+    def build_collator(
+        self,
+        training_args,  # type: "AxolotlTrainingArguments"  # type: ignore
+        is_eval=False,
+        **kwargs,
+    ):
+        if training_args.pretraining:
+            if (
+                self.cfg.pretraining_sample_concatenation is False
+                or self.cfg.micro_batch_size > 1
+            ):
+                return DataCollatorForSeq2Seq(self.tokenizer, **kwargs)
+            if not (self.cfg.sample_packing and self.cfg.pretrain_multipack_attn) or (
+                self.cfg.micro_batch_size == 1 and is_eval is False
+            ):
+                return None
+
+        if self.cfg.model_config_type == "mamba":
+            return MambaDataCollator(tokenizer=self.tokenizer)
+
+        use_batch_sampler_collator = False
+        if is_eval is False and training_args.sample_packing:
+            use_batch_sampler_collator = True
+        if is_eval and training_args.eval_sample_packing:
+            use_batch_sampler_collator = True
+
+        collator: Type[
+            Union[
+                V2BatchSamplerDataCollatorForSeq2Seq,
+                BatchSamplerDataCollatorForSeq2Seq,
+                DataCollatorForSeq2Seq,
+                DataCollatorWithFlattening,
+                DataCollatorForPreference,
+            ]
+        ]
+        collator_args = [self.tokenizer]
+
+        collator_cls_and_kwargs = None
+        if self.cfg.plugins:
+            plugin_manager = PluginManager.get_instance()
+            collator_cls_and_kwargs = plugin_manager.get_collator_cls_and_kwargs(
+                self.cfg, is_eval=is_eval
+            )
+
+        if collator_cls_and_kwargs:
+            collator = collator_cls_and_kwargs[0]
+            if kwargs and isinstance(kwargs, dict):
+                kwargs.update(collator_cls_and_kwargs[1])
+        elif self.cfg.reward_model:
+            collator = DataCollatorForPreference
+            tokenizer = collator_args.pop(0)
+            kwargs["pad_token_id"] = tokenizer.pad_token_id
+            kwargs.pop("padding")
+        elif use_batch_sampler_collator:
+            # Use V2BatchSamplerDataCollatorForSeq2Seq for flex attention,
+            # supported multipack models, or non-flash-attention llama
+            if (
+                self.cfg.attn_implementation == "flex_attention"
+                or self.cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES
+                or (
+                    self.cfg.model_config_type in ["llama"]
+                    and self.cfg.attn_implementation != "flash_attention_2"
+                )
+            ):
+                collator = V2BatchSamplerDataCollatorForSeq2Seq
+            else:
+                collator = BatchSamplerDataCollatorForSeq2Seq
+        else:
+            if self.cfg.processor_type and self.processor:
+                collator = MultiModalChatDataCollator
+                # Mirror ChatTemplateStrategy: per-dataset masking knobs from first MM dataset, else global cfg.
+                # NOTE: Multi-dataset configs use the first dataset's masking knobs for all datasets;
+                # heterogeneous per-dataset overrides are not supported in the MM path today.
+                ds_entries = self.cfg.datasets or []
+                ds_cfg = ds_entries[0] if ds_entries else None
+
+                def _ds_get(cfg_obj, key):
+                    # Handle DictDefault / dict / pydantic uniformly:
+                    # dict-style .get first, then attribute access.
+                    if cfg_obj is None:
+                        return None
+                    if hasattr(cfg_obj, "get"):
+                        try:
+                            return cfg_obj.get(key)
+                        except (AttributeError, KeyError, TypeError):
+                            pass
+                    return getattr(cfg_obj, key, None)
+
+                roles_to_train = _ds_get(ds_cfg, "roles_to_train")
+                train_on_eos = _ds_get(ds_cfg, "train_on_eos")
+
+                # cfg.role_boundaries replaces the strategy's built-in markers.
+                role_boundaries_override = None
+                if self.cfg.role_boundaries:
+                    role_boundaries_override = list(self.cfg.role_boundaries)
+
+                # Deduped union of per-dataset `field_messages` for the MM collator.
+                field_messages = []
+                for dataset_cfg in ds_entries:
+                    field_message = _ds_get(dataset_cfg, "field_messages")
+                    if field_message and field_message not in field_messages:
+                        field_messages.append(field_message)
+
+                # build() calls build_collator twice (eval + train); log once.
+                if not is_eval:
+                    LOG.info(
+                        "MM collator: train_on_inputs=%s roles_to_train=%s "
+                        "train_on_eos=%s role_boundaries_override=%s",
+                        bool(self.cfg.train_on_inputs),
+                        roles_to_train,
+                        train_on_eos,
+                        "set" if role_boundaries_override else "none",
+                    )
+                    _warn_if_num_workers_zero_for_mm(self.cfg, LOG)
+
+                kwargs["processing_strategy"] = get_processing_strategy(
+                    self.processor,
+                    training_args.chat_template,
+                    self.cfg.chat_template,
+                    image_size=training_args.image_size,
+                    image_resize_algorithm=training_args.image_resize_algorithm,
+                    train_on_inputs=bool(self.cfg.train_on_inputs),
+                    roles_to_train=roles_to_train,
+                    train_on_eos=train_on_eos,
+                    role_boundaries_override=role_boundaries_override,
+                    field_messages=field_messages or None,
+                )
+            elif self.cfg.batch_flattening:
+                collator = DataCollatorWithFlattening
+                collator_args.pop(0)
+                kwargs.pop("pad_to_multiple_of", None)
+                kwargs.pop("padding", None)
+            else:
+                collator = DataCollatorForSeq2Seq
+
+        kwargs["return_tensors"] = "pt"
+
+        return collator(
+            *collator_args,
+            **kwargs,
+        )
diff --git a/src/axolotl/core/builders/rl.py b/src/axolotl/core/builders/rl.py
new file mode 100644
index 0000000000..25340809a3
--- /dev/null
+++ b/src/axolotl/core/builders/rl.py
@@ -0,0 +1,339 @@
+"""Builder for RLHF trainers"""
+
+import inspect
+from pathlib import Path
+
+from axolotl.core.builders.base import TrainerBuilderBase
+from axolotl.core.trainers import (
+    AxolotlCPOTrainer,
+    AxolotlKTOTrainer,
+    AxolotlORPOTrainer,
+)
+from axolotl.core.trainers.dpo import DPOStrategy
+from axolotl.core.trainers.dpo.args import AxolotlDPOConfig
+from axolotl.integrations.base import PluginManager
+from axolotl.loaders.utils import ensure_dtype
+from axolotl.utils.callbacks.qat import QATCallback
+from axolotl.utils.import_helper import get_cls_from_module_str
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.enums import RLType
+
+LOG = get_logger(__name__)
+
+
+class HFRLTrainerBuilder(TrainerBuilderBase):
+    """Trainer factory class for TRL-based RLHF trainers (e.g. DPO)"""
+
+    def get_callbacks(self):
+        callbacks = super().get_callbacks()
+
+        if self.cfg.qat:
+            callbacks.append(QATCallback(self.cfg.qat))
+
+        return callbacks
+
+    def get_post_trainer_create_callbacks(self, trainer):
+        callbacks = super().get_post_trainer_create_callbacks(trainer=trainer)
+        return callbacks
+
+    def _get_trainer_cls(self, trainer_kwargs: dict):
+        """
+        Returns trainer_cls and trainer_cls_args
+        """
+        if self.cfg.plugins:
+            plugin_manager = PluginManager.get_instance()
+            trainer_cls = plugin_manager.get_trainer_cls(self.cfg)
+            trainer_cls_args = []  # type: ignore
+
+            if trainer_cls is not None:
+                return trainer_cls, trainer_cls_args
+
+        trainer_cls = None
+        trainer_cls_args = [self.model]
+
+        if self.cfg.rl in {RLType.GRPO, RLType.GDPO}:
+            from axolotl.core.trainers.grpo import GRPOStrategy
+
+            async_grpo = bool(
+                self.cfg.trl
+                and (
+                    getattr(self.cfg.trl, "async_prefetch", False)
+                    or getattr(self.cfg.trl, "use_data_producer", False)
+                )
+            )
+            trainer_cls = GRPOStrategy.get_trainer_class(
+                sequence_parallel=self.cfg.context_parallel_size > 1,
+                async_grpo=async_grpo,
+            )
+            trainer_cls_args.extend(GRPOStrategy.set_trainer_args(self.cfg))
+            trainer_kwargs.update(GRPOStrategy.set_trainer_kwargs(self.cfg))
+
+        elif self.cfg.rl in [RLType.DPO, RLType.IPO]:
+            trainer_cls = DPOStrategy.get_trainer_class()
+            trainer_cls_args.append(self.model_ref)
+
+        elif self.cfg.rl is RLType.ORPO:
+            trainer_cls = AxolotlORPOTrainer
+        elif self.cfg.rl is RLType.KTO:
+            trainer_cls = AxolotlKTOTrainer
+        elif self.cfg.rl is RLType.SIMPO:
+            trainer_cls = AxolotlCPOTrainer
+        elif self.cfg.rl is RLType.EBFT:
+            from axolotl.core.trainers.ebft import EBFTStrategy
+
+            trainer_cls = EBFTStrategy.get_trainer_class(self.cfg)
+            trainer_kwargs.update(EBFTStrategy.set_trainer_kwargs(self.cfg))
+        else:
+            raise ValueError(f"Unsupported RL: {self.cfg.rl}")
+
+        if self.cfg.trainer_cls:
+            # override the trainer cls
+            try:
+                trainer_cls = get_cls_from_module_str(self.cfg.trainer_cls)
+                LOG.debug(f"Using custom trainer class: {self.cfg.trainer_cls}")
+            except (ImportError, AttributeError, ValueError) as e:
+                raise ValueError(
+                    f"Failed to load custom trainer class '{self.cfg.trainer_cls}': {e}"
+                ) from e
+
+        return trainer_cls, trainer_cls_args
+
+    def _build_training_arguments(self, total_num_steps):
+        """
+        Returns training_args and trainer_kwargs
+        """
+        from axolotl.core.training_args import (
+            AxolotlCPOConfig,
+            AxolotlKTOConfig,
+            AxolotlORPOConfig,
+        )
+
+        training_args_kwargs, trainer_kwargs = self._set_base_training_args(
+            total_num_steps=total_num_steps
+        )
+
+        if self.cfg.remove_unused_columns is not None:
+            training_args_kwargs["remove_unused_columns"] = (
+                self.cfg.remove_unused_columns
+            )
+        else:
+            training_args_kwargs["remove_unused_columns"] = False
+
+        if self.cfg.trl and self.cfg.trl.beta is not None:
+            training_args_kwargs["beta"] = self.cfg.trl.beta
+        elif self.cfg.rl_beta is not None:
+            training_args_kwargs["beta"] = self.cfg.rl_beta
+        elif self.cfg.orpo_alpha is not None:
+            # trl does some odd mapping of alpha to beta to reuse the beta parameter ???
+            training_args_kwargs["beta"] = self.cfg.orpo_alpha
+
+        if self.cfg.use_wandb:
+            training_args_kwargs["run_name"] = self.cfg.wandb_name
+
+        training_args_cls = None
+        blocklist_args_kwargs = []
+        if self.cfg.rl is RLType.SIMPO:
+            training_args_cls = AxolotlCPOConfig
+            training_args_kwargs["loss_type"] = "simpo"
+            training_args_kwargs["simpo_gamma"] = self.cfg.simpo_gamma
+            if self.cfg.cpo_alpha is not None:
+                training_args_kwargs["cpo_alpha"] = self.cfg.cpo_alpha
+
+            blocklist_args_kwargs.append("max_prompt_length")
+
+        elif self.cfg.rl is RLType.ORPO:
+            training_args_cls = AxolotlORPOConfig
+
+            blocklist_args_kwargs.append("max_prompt_length")
+
+        elif self.cfg.rl is RLType.KTO:
+            training_args_cls = AxolotlKTOConfig
+            # KTOConfig in TRL >= 0.27.0 no longer accepts max_prompt_length
+            blocklist_args_kwargs.append("max_prompt_length")
+
+            training_args_kwargs["desirable_weight"] = (
+                self.cfg.kto_desirable_weight or 1.0
+            )
+            training_args_kwargs["undesirable_weight"] = (
+                self.cfg.kto_undesirable_weight or 1.0
+            )
+
+        elif self.cfg.rl in {RLType.GRPO, RLType.GDPO}:
+            from axolotl.core.trainers.grpo import GRPOStrategy
+
+            async_grpo = bool(
+                self.cfg.trl
+                and (
+                    getattr(self.cfg.trl, "async_prefetch", False)
+                    or getattr(self.cfg.trl, "use_data_producer", False)
+                )
+            )
+            training_args_cls = GRPOStrategy.get_training_args_class(
+                async_grpo=async_grpo
+            )
+            training_args_kwargs.update(GRPOStrategy.set_training_args_kwargs(self.cfg))
+            blocklist_args_kwargs = GRPOStrategy.get_blocklist_args_kwargs()
+            if not async_grpo:
+                # Filter out async/fast-async-only fields not in standard GRPOConfig.
+                # These are defined in FastAsyncGRPOConfig and only used by
+                # AxolotlAsyncGRPOConfig. Standard GRPOConfig rejects them.
+                import dataclasses
+
+                from trl import GRPOConfig as _BaseGRPOConfig
+
+                from axolotl.core.trainers.grpo.fast_async_trainer import (
+                    FastAsyncGRPOConfig,
+                )
+
+                async_only_fields = {
+                    f.name for f in dataclasses.fields(FastAsyncGRPOConfig)
+                } - {f.name for f in dataclasses.fields(_BaseGRPOConfig)}
+                blocklist_args_kwargs.extend(list(async_only_fields))
+            if self.cfg.rl is RLType.GDPO:
+                training_args_kwargs.setdefault(
+                    "multi_objective_aggregation", "normalize_then_sum"
+                )
+
+        elif self.cfg.rl in [RLType.DPO, RLType.IPO]:
+            training_args_cls = AxolotlDPOConfig
+            training_args_kwargs.update(DPOStrategy.set_training_args_kwargs(self.cfg))
+
+        elif self.cfg.rl is RLType.EBFT:
+            from axolotl.core.trainers.ebft import EBFTStrategy
+
+            training_args_cls = EBFTStrategy.get_training_args_class(self.cfg)
+            training_args_kwargs.update(EBFTStrategy.set_training_args_kwargs(self.cfg))
+            blocklist_args_kwargs = EBFTStrategy.get_blocklist_args_kwargs(self.cfg)
+        else:
+            raise ValueError(f"Unsupported RL: {self.cfg.rl}")
+
+        for blocklist_key in blocklist_args_kwargs:
+            if blocklist_key in training_args_kwargs:
+                del training_args_kwargs[blocklist_key]
+
+        if self.cfg.plugins:
+            plugin_manager = PluginManager.get_instance()
+            plugin_training_args = plugin_manager.get_training_args(self.cfg)
+            if plugin_training_args:
+                training_args_kwargs.update(plugin_training_args)
+
+        training_args = training_args_cls(
+            logging_first_step=True,
+            **training_args_kwargs,
+        )
+
+        # unset run_name so wandb sets up experiment names
+        if self.cfg.use_wandb and training_args.run_name == training_args.output_dir:
+            training_args.run_name = None
+
+        return training_args, trainer_kwargs
+
+    def build_collator(self, **kwargs):
+        """Build a data collator for preference-tuning trainers.
+
+        Returns None for RL types that provide their own collator (e.g. GRPO,
+        KTO), letting the trainer construct its default. For DPO/IPO/ORPO/SIMPO
+        returns an ``AxolotlDPODataCollatorWithPadding`` when
+        ``pad_to_multiple_of`` is set, otherwise None (so the trainer
+        falls back to the TRL default).
+        """
+        if self.cfg.rl not in (
+            RLType.DPO,
+            RLType.IPO,
+            RLType.ORPO,
+            RLType.SIMPO,
+        ):
+            return None
+
+        pad_to_multiple_of = getattr(self.cfg, "pad_to_multiple_of", None)
+        if not pad_to_multiple_of:
+            return None
+
+        from axolotl.utils.collators.dpo import AxolotlDPODataCollatorWithPadding
+
+        LOG.info(
+            f"Using AxolotlDPODataCollatorWithPadding with pad_to_multiple_of="
+            f"{pad_to_multiple_of}"
+        )
+        is_enc_dec = getattr(self.model.config, "is_encoder_decoder", False)
+        return AxolotlDPODataCollatorWithPadding(
+            pad_token_id=self.tokenizer.pad_token_id,
+            is_encoder_decoder=is_enc_dec,
+            pad_to_multiple_of=pad_to_multiple_of,
+            **kwargs,
+        )
+
+    def build(self, total_num_steps):
+        training_args, trainer_kwargs = self._build_training_arguments(total_num_steps)
+
+        if (data_collator := self.build_collator()) is not None:
+            trainer_kwargs["data_collator"] = data_collator
+
+        if self.eval_dataset:
+            trainer_kwargs["eval_dataset"] = self.eval_dataset
+        if (
+            self.cfg.adapter
+            and self.peft_config
+            and self.cfg.rl not in (RLType.GRPO, RLType.ORPO, RLType.EBFT, RLType.SIMPO)
+        ):
+            trainer_kwargs["peft_config"] = self.peft_config
+
+        trainer_cls, trainer_cls_args = self._get_trainer_cls(trainer_kwargs)
+
+        sig = inspect.signature(trainer_cls)
+        if "tokenizer" in sig.parameters:
+            trainer_kwargs["tokenizer"] = self.tokenizer
+        else:
+            trainer_kwargs["processing_class"] = self.tokenizer
+
+        if self.cfg.datasets is not None and (
+            trainer_cls is DPOStrategy.get_trainer_class()
+        ):
+            trainer_kwargs["dataset_tags"] = [
+                d["path"] for d in self.cfg.datasets if not Path(d["path"]).is_dir()
+            ]
+
+        trainer_kwargs, trainer_cls = self.hook_pre_create_trainer(
+            trainer_kwargs, trainer_cls
+        )
+
+        # Allow FP8-quantized models to be fine-tuned with LoRA adapters.
+        # transformers' validate_quantization_for_training blocks FP8 because
+        # hf_quantizer.is_trainable is False, but LoRA only trains the adapters
+        # (base weights stay frozen in FP8).
+        _orig_validate_quant = None
+        if (
+            self.cfg.adapter
+            and hasattr(self.model, "is_quantized")
+            and self.model.is_quantized
+        ):
+            import transformers.trainer as _trainer_module
+
+            _orig_validate_quant = _trainer_module.validate_quantization_for_training
+            _trainer_module.validate_quantization_for_training = lambda model: None
+
+        try:
+            trainer = trainer_cls(
+                *trainer_cls_args,
+                args=training_args,
+                train_dataset=self.train_dataset,
+                callbacks=self.get_callbacks(),
+                **trainer_kwargs,
+            )
+        finally:
+            if _orig_validate_quant is not None:
+                import transformers.trainer as _trainer_module
+
+                _trainer_module.validate_quantization_for_training = (
+                    _orig_validate_quant
+                )
+        if self.cfg.fsdp_config or self.cfg.fsdp:
+            ensure_dtype(trainer.model, dtype=self.cfg.torch_dtype)
+            if self.cfg.rl in [RLType.DPO, RLType.IPO] and trainer.ref_model:
+                ensure_dtype(trainer.ref_model, dtype=self.cfg.torch_dtype)
+
+        trainer = self.hook_post_create_trainer(trainer)
+        for callback in self.get_post_trainer_create_callbacks(trainer):
+            trainer.add_callback(callback)
+
+        return trainer
diff --git a/src/axolotl/core/chat/__init__.py b/src/axolotl/core/chat/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/core/chat/format/__init__.py b/src/axolotl/core/chat/format/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/core/chat/format/chatml.py b/src/axolotl/core/chat/format/chatml.py
new file mode 100644
index 0000000000..deb8a9997b
--- /dev/null
+++ b/src/axolotl/core/chat/format/chatml.py
@@ -0,0 +1,35 @@
+"""
+ChatML transformation functions for MessageContents
+"""
+
+from typing import Optional
+
+from ..messages import MessageContents, Messages
+from .shared import wrap_tools
+
+
+def format_message(
+    message: Messages,
+    message_index: Optional[int] = None,
+) -> Messages:
+    if message.is_chat_formatted:
+        return message
+
+    # prepend the role prefix within a MessageContents to message.content
+    message.content.insert(
+        0,
+        MessageContents(
+            type="text",
+            value=f"<|im_start|>{message.role}\n",
+            weight=0,
+        ),
+    )
+    message.content.append(
+        MessageContents(type="text", value="<|im_end|>", weight=message.weight)
+    )
+    message.content.append(MessageContents(type="text", value="\n", weight=0))
+
+    message = wrap_tools(message)
+
+    message.is_chat_formatted = True
+    return message
diff --git a/src/axolotl/core/chat/format/llama3x.py b/src/axolotl/core/chat/format/llama3x.py
new file mode 100644
index 0000000000..a0ce053e51
--- /dev/null
+++ b/src/axolotl/core/chat/format/llama3x.py
@@ -0,0 +1,46 @@
+"""
+Llama 3.x chat formatting functions for MessageContents
+"""
+
+from typing import Optional
+
+from ..messages import MessageContents, Messages
+from .shared import wrap_tools
+
+
+def format_message(message: Messages, message_index: Optional[int] = None) -> Messages:
+    if message.is_chat_formatted:
+        return message
+
+    message_role = message.role
+    if message.role == "tool":
+        message_role = "ipython"
+
+    # prepend the role prefix within a MessageContents to message.content
+    message.content.insert(
+        0,
+        MessageContents(
+            type="text",
+            value=f"<|start_header_id|>{message_role}<|end_header_id|>\n\n",
+            weight=0,
+        ),
+    )
+
+    message.content.append(
+        MessageContents(type="text", value="<|eot_id|>", weight=message.weight)
+    )
+
+    message = wrap_tools(message)
+
+    if message_index == 0:
+        message.content.insert(
+            0,
+            MessageContents(
+                type="text",
+                value="<|begin_of_text|>",
+                weight=0,
+            ),
+        )
+
+    message.is_chat_formatted = True
+    return message
diff --git a/src/axolotl/core/chat/format/shared.py b/src/axolotl/core/chat/format/shared.py
new file mode 100644
index 0000000000..0a0f56f3ad
--- /dev/null
+++ b/src/axolotl/core/chat/format/shared.py
@@ -0,0 +1,48 @@
+"""
+shared functions for format transforms
+"""
+
+from axolotl.core.chat.messages import MessageContents, Messages
+
+
+def wrap_tools(message: Messages):
+    # loop over message.content by index to find tool calls, we need to wrap each with tags,
+    # so be wary of indexing issues when changing the list while iterating.
+    # iterate over the range in reverse order to avoid index shifting
+    for i in range(len(message.content) - 1, -1, -1):
+        if message.content[i].type == "tool_call":
+            # append a </tool_call> MessageContents text tag after
+            message.content.insert(
+                i + 1,
+                MessageContents(
+                    type="text", value="</tool_call>\n", weight=message.weight
+                ),
+            )
+            # make sure the actual tool call content ends with a newline
+            message.content[i].has_newline = True
+            # prepend a <tool_call> MessageContents text tag before
+            message.content.insert(
+                i,
+                MessageContents(
+                    type="text", value="<tool_call>\n", weight=message.weight
+                ),
+            )
+        elif message.content[i].type == "tool_response":
+            # append a </tool_call> MessageContents text tag after
+            message.content.insert(
+                i + 1,
+                MessageContents(
+                    type="text", value="</tool_response>\n", weight=message.weight
+                ),
+            )
+            # make sure the actual tool response content ends with a newline
+            message.content[i].has_newline = True
+            # prepend a <tool_call> MessageContents text tag before
+            message.content.insert(
+                i,
+                MessageContents(
+                    type="text", value="<tool_response>\n", weight=message.weight
+                ),
+            )
+
+    return message
diff --git a/src/axolotl/core/chat/messages.py b/src/axolotl/core/chat/messages.py
new file mode 100644
index 0000000000..912a12ca1e
--- /dev/null
+++ b/src/axolotl/core/chat/messages.py
@@ -0,0 +1,230 @@
+"""
+internal message representations of chat messages
+"""
+
+import json
+from enum import Enum
+from typing import Any, Callable, List, Optional, Union
+
+from pydantic import BaseModel
+from transformers import PreTrainedTokenizer
+
+
+class MessageRoles(str, Enum):
+    """
+    Message roles for the system, user, assistant, and tools
+    """
+
+    system = "system"
+    user = "user"
+    assistant = "assistant"
+    tool = "tool"
+    ipython = (
+        # for responses from builtin tools
+        "ipython"
+    )
+
+
+class MessageContentTypes(str, Enum):
+    """
+    Message content types for text, image, audio, tool calls, and tool responses
+    """
+
+    special_token = "special_token"  # nosec B105
+    text = "text"
+    image = "image"
+    audio = "audio"
+    tool_call = "tool_call"
+    tool_response = "tool_response"
+
+
+class SpecialToken(str, Enum):
+    """
+    Special tokens for beginning of string and end of string
+    """
+
+    bos_token = "bos_token"  # nosec B105
+    eos_token = "eos_token"  # nosec B105
+
+
+class ToolCallFunction(BaseModel):
+    """
+    Tool call function with name and arguments
+    """
+
+    name: str
+    arguments: dict[str, str]
+
+
+class Tool(BaseModel):
+    """
+    Tool with description, function, and parameters
+    """
+
+    description: str
+    function: ToolCallFunction
+    parameters: dict[str, str]  # .properties
+
+
+class ToolCallContents(BaseModel):
+    """
+    Tool call contents with name, arguments, and optional id
+    """
+
+    name: str
+    arguments: dict[str, Union[str, int]]
+    id: Optional[str] = None
+
+    def __str__(self) -> str:
+        data = {"name": self.name, "arguments": self.arguments}
+        if self.id is not None:
+            data["id"] = self.id
+        return json.dumps(data)
+
+
+class ToolResponseContents(BaseModel):
+    """
+    Tool response contents with name, content, and optional id
+    """
+
+    name: str
+    content: Union[str, dict[str, Union[str, int, float]]]
+    id: Optional[str] = None
+
+    def __str__(self) -> str:
+        data = {"name": self.name, "content": self.content}
+        if self.id is not None:
+            data["id"] = self.id
+        return json.dumps(data)
+
+
+class MessageContents(BaseModel):
+    """
+    Message contents with type, value, metadata, weight, newline, and end of contents
+    """
+
+    type: Union[str, MessageContentTypes]
+    value: Union[str, ToolCallContents, ToolResponseContents, SpecialToken]
+    meta: Optional[dict[str, Any]] = None  # support additional arbitrary metadata
+    weight: Optional[Union[int, float]] = None
+    has_newline: bool = False
+    eoc: bool = False  # end of contents
+
+    def __str__(self) -> str:
+        str_val = str(self.value)
+        if self.has_newline and not str_val.endswith("\n"):
+            str_val += "\n"
+        return str_val
+
+
+class Messages(BaseModel):
+    """
+    Messages with role, content, metadata, weight, and chat formatting
+    """
+
+    role: Union[MessageRoles, str]  # allows for arbitrary roles
+    content: List["MessageContents"]
+    meta: Optional[dict[str, Any]] = None  # support additional arbitrary metadata
+    weight: Optional[Union[int, float]] = None
+    is_chat_formatted: bool = False
+
+    def __str__(self) -> str:
+        return "".join(str(c) for c in self.content)
+
+    def tokenized(
+        self, tokenizer: PreTrainedTokenizer, ignore_index=-100
+    ) -> dict[str, List[int]]:
+        # iterate over the contents, tokenizing the concatenated string values up to the current MessageContents
+        # returns a dictionary mapping w input_ids, attention_mask, and labels
+        input_ids: List[int] = []
+        labels: List[int] = []
+        pending_input_ids: List[int] = []
+        pending_weight = self.weight
+        running_content = ""
+        for _, msg_content in enumerate(self.content):
+            # TODO also handle non-text content types
+            if msg_content.type in [
+                MessageContentTypes.text.value,
+                MessageContentTypes.tool_call.value,
+                MessageContentTypes.tool_response.value,
+            ]:
+                running_content += str(msg_content)
+                tok_results = tokenizer(running_content, add_special_tokens=False)
+                tok_input_ids = tok_results["input_ids"]
+                if pending_input_ids:
+                    new_pending_inputs = tok_input_ids[
+                        len(input_ids) : len(input_ids) + len(pending_input_ids)
+                    ]
+                    if new_pending_inputs != pending_input_ids:
+                        pending_input_ids = new_pending_inputs
+                    input_ids.extend(pending_input_ids)
+                    if pending_weight:
+                        labels.extend(pending_input_ids)
+                    else:
+                        labels.extend([ignore_index] * len(pending_input_ids))
+                pending_input_ids = tok_results["input_ids"][len(input_ids) :]
+                pending_weight = self.weight and msg_content.weight not in [0, 0.0]
+        input_ids.extend(pending_input_ids)
+        if pending_weight:
+            labels.extend(pending_input_ids)
+        else:
+            labels.extend([ignore_index] * len(pending_input_ids))
+        attention_mask = [1] * len(input_ids)
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+        }
+
+
+class Chats(BaseModel):
+    """
+    top level data structure for chat conversations
+    """
+
+    conversation: List[Messages]
+
+    def __str__(self) -> str:
+        return "".join(str(c) for c in self.conversation)
+
+    def tokenized(
+        self, tokenizer: Callable[[str], dict[str, List[int]]], ignore_index=-100
+    ) -> dict[str, List[int]]:
+        input_ids = []
+        attention_mask = []
+        labels = []
+        for msg in self.conversation:
+            msg_results = msg.tokenized(tokenizer, ignore_index)
+            input_ids.extend(msg_results["input_ids"])
+            attention_mask.extend(msg_results["attention_mask"])
+            labels.extend(msg_results["labels"])
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+        }
+
+
+class ChatFormattedChats(Chats):
+    """
+    Chat formatted chats with formatter and optional train on inputs
+    """
+
+    formatter: Callable  # [[Union[dict, Chats]], Chats]
+    train_on_inputs: bool = False
+
+    def model_post_init(self, __context):
+        for i, msg in enumerate(self.conversation):
+            self.conversation[i] = self.formatter(msg, message_index=i)
+            if self.train_on_inputs:
+                self.conversation[i].weight = 1
+
+
+class PreferenceChats(BaseModel):
+    """
+    representation for preference data for chat
+    """
+
+    prompt: List[Messages]
+    chosen: Messages
+    rejected: Messages
diff --git a/src/axolotl/core/datasets/__init__.py b/src/axolotl/core/datasets/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/core/datasets/chat.py b/src/axolotl/core/datasets/chat.py
new file mode 100644
index 0000000000..a4dc300d9c
--- /dev/null
+++ b/src/axolotl/core/datasets/chat.py
@@ -0,0 +1,51 @@
+"""
+chat dataset module
+"""
+
+from typing import Callable, Optional, Union
+
+from datasets import Dataset
+from transformers import PreTrainedTokenizer
+
+from axolotl.core.chat.messages import ChatFormattedChats
+
+
+class TokenizedChatDataset(Dataset):
+    """
+    Tokenized chat dataset
+    """
+
+    def __init__(
+        self,
+        data: Dataset,
+        model_transform: Union[PreTrainedTokenizer, Callable],
+        *args,
+        message_transform: Optional[Callable] = None,
+        formatter=None,
+        process_count: Optional[int] = None,
+        keep_in_memory: Optional[bool] = False,
+        **kwargs,
+    ):
+        def map_fn(ex):
+            if message_transform is not None:
+                ex = message_transform(ex)
+            if formatter is not None:
+                ex = ChatFormattedChats(
+                    formatter=formatter,
+                    **ex,
+                )
+            else:
+                ex = ChatFormattedChats(
+                    **ex,
+                )
+            return ex.tokenized(model_transform)
+
+        features = data.features.keys()
+        tokenized_data = data.map(
+            map_fn,
+            num_proc=process_count,
+            keep_in_memory=keep_in_memory,
+            remove_columns=features,
+            desc="Tokenizing Chats",
+        )
+        super().__init__(tokenized_data.data, *args, **kwargs)
diff --git a/src/axolotl/core/datasets/transforms/__init__.py b/src/axolotl/core/datasets/transforms/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/core/datasets/transforms/chat_builder.py b/src/axolotl/core/datasets/transforms/chat_builder.py
new file mode 100644
index 0000000000..0de0ecb403
--- /dev/null
+++ b/src/axolotl/core/datasets/transforms/chat_builder.py
@@ -0,0 +1,151 @@
+"""
+This module contains a function that builds a transform that takes a row from the
+dataset and converts it to a Chat.
+"""
+
+from typing import Any, Mapping
+
+
+def chat_message_transform_builder(
+    train_on_inputs=False,
+    conversations_field: str = "messages",
+    message_field_role: str | list[str] | None = None,  # commonly "role"
+    message_field_content: str | list[str] | None = None,  # commonly "content"
+    message_field_training: str | list[str] | None = None,  # commonly "weight"
+):
+    """Builds a transform that takes a row from the dataset and converts it to a Chat
+
+    Args:
+        train_on_inputs (bool, optional):
+            If True, the transform will train on the inputs. If False, the transform will train on the targets.
+            Defaults to False.
+        conversations_field (str, optional):
+            The field name of the conversations. Defaults to "messages".
+        message_field_role (str | list[str], optional):
+            The field name of the role.
+        message_field_content (str | list[str], optional):
+            The field name of the message content.
+        message_field_training (str | list[str], optional):
+            The field name of the train/weight.
+
+    Returns:
+        Callable:
+            A function that takes a list of conversations and returns a list of messages.
+    """
+
+    if message_field_training is None:
+        message_field_training = ["train", "weight"]
+    if message_field_content is None:
+        message_field_content = ["value", "text", "content"]
+    if message_field_role is None:
+        message_field_role = ["role", "from"]
+    message_field_role = (
+        [message_field_role]
+        if isinstance(message_field_role, str)
+        else message_field_role
+    )
+    message_field_content = (
+        [message_field_content]
+        if isinstance(message_field_content, str)
+        else message_field_content
+    )
+    message_weight_fields = (
+        [message_field_training]
+        if isinstance(message_field_training, str)
+        else message_field_training
+    )
+
+    role_value_mappings = {
+        "system": "system",
+        "user": "user",
+        "human": "user",
+        "assistant": "assistant",
+        "gpt": "assistant",
+        "tool": "tool",
+        "ipython": "ipython",
+    }
+    if train_on_inputs:
+        role_default_weights_mappings = {
+            "system": 1,
+            "user": 1,
+            "assistant": 1,
+            "tool": 1,
+            "ipython": 1,
+        }
+    else:
+        role_default_weights_mappings = {
+            "system": 0,
+            "user": 0,
+            "assistant": 1,
+            "tool": 0,
+            "ipython": 0,
+        }
+
+    def transform_builder(sample: Mapping[str, Any]):
+        if conversations_field not in sample:
+            raise ValueError(f"Field '{conversations_field}' not found in sample.")
+        # if none of the role fields are in the message, raise an error
+        if not any(
+            role in sample[conversations_field][0] for role in message_field_role
+        ):
+            raise ValueError("No role field found in message.")
+        role_field = next(
+            role
+            for role in message_field_role
+            if role in sample[conversations_field][0]
+        )
+        if not any(
+            field in sample[conversations_field][0] for field in message_field_content
+        ):
+            raise ValueError("No message_content field found in message.")
+        message_content_field = next(
+            field
+            for field in message_field_content
+            if field in sample[conversations_field][0]
+        )
+        if not any(
+            field in sample[conversations_field][0] for field in message_field_training
+        ):
+            message_weight_field = None
+        else:
+            message_weight_field = next(
+                field
+                for field in message_weight_fields
+                if field in sample[conversations_field][0]
+            )
+
+        messages = []
+        for message in sample[conversations_field]:
+            role = role_value_mappings[message[role_field]]
+            weight = (
+                int(message[message_weight_field])
+                if message_weight_field
+                else role_default_weights_mappings[role]
+            )
+
+            # TODO if "tool_calls" in message[message_content_field]: then convert tool call to ToolCallContents
+            if isinstance(message[message_content_field], str):
+                messages.append(
+                    {
+                        "role": role,
+                        "content": [
+                            {
+                                "type": "text",
+                                "value": message[message_content_field],
+                            }
+                        ],
+                        "weight": weight,
+                    }
+                )
+            else:
+                messages.append(
+                    {
+                        "role": role,
+                        "content": message[message_content_field],
+                        "weight": weight,
+                    }
+                )
+
+        return {"conversation": messages}
+
+    return transform_builder
diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
deleted file mode 100644
index 2f38b12dc1..0000000000
--- a/src/axolotl/core/trainer_builder.py
+++ /dev/null
@@ -1,1616 +0,0 @@
-# pylint: disable=too-many-lines
-"""
-Builder for the training args and trainer
-"""
-
-import abc
-import importlib
-import importlib.util
-import logging
-import math
-import sys
-from abc import abstractmethod
-from collections import defaultdict
-from dataclasses import dataclass, field
-from functools import wraps
-from pathlib import Path
-from typing import Dict, List, Literal, Optional, Type, Union
-
-import torch
-import transformers
-from datasets import Dataset
-from torch.optim.lr_scheduler import OneCycleLR
-from torch.utils.data import BatchSampler, DataLoader, RandomSampler, SequentialSampler
-from transformers import (
-    EarlyStoppingCallback,
-    PreTrainedModel,
-    Trainer,
-    TrainerCallback,
-    TrainingArguments,
-)
-from transformers.trainer_utils import seed_worker
-from transformers.utils import is_sagemaker_mp_enabled
-from trl import DPOTrainer, ORPOConfig, ORPOTrainer
-from trl.trainer.utils import pad_to_length
-
-from axolotl.loraplus import create_loraplus_optimizer
-from axolotl.monkeypatch.multipack import SUPPORTED_MULTIPACK_MODEL_TYPES
-from axolotl.monkeypatch.relora import ReLoRACallback, ReLoRAScheduler
-from axolotl.utils import is_mlflow_available
-from axolotl.utils.callbacks import (
-    EvalFirstStepCallback,
-    GPUStatsCallback,
-    LossWatchDogCallback,
-    SaveAxolotlConfigtoWandBCallback,
-    SaveBetterTransformerModelCallback,
-    SaveModelOnTrainEndCallback,
-    bench_eval_callback_factory,
-    causal_lm_bench_eval_callback_factory,
-    log_prediction_callback_factory,
-)
-from axolotl.utils.callbacks.lisa import lisa_callback_factory
-from axolotl.utils.collators import (
-    BatchSamplerDataCollatorForSeq2Seq,
-    DataCollatorForSeq2Seq,
-    MambaDataCollator,
-    V2BatchSamplerDataCollatorForSeq2Seq,
-)
-from axolotl.utils.models import ensure_dtype
-from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
-from axolotl.utils.schedulers import (
-    get_cosine_schedule_with_min_lr,
-    get_cosine_schedule_with_quadratic_warmup,
-    get_cosine_schedule_with_warmup_decay_constant,
-)
-
-if is_sagemaker_mp_enabled():
-    import smdistributed.modelparallel.torch as smp
-
-try:
-    import torch._dynamo  # pylint: disable=ungrouped-imports
-except ImportError:
-    pass
-
-LOG = logging.getLogger("axolotl.core.trainer_builder")
-
-
-def _sanitize_kwargs_for_tagging(tag_names, kwargs=None):
-    if isinstance(tag_names, str):
-        tag_names = [tag_names]
-
-    if kwargs is not None:
-        if "tags" not in kwargs:
-            kwargs["tags"] = tag_names
-        elif "tags" in kwargs and isinstance(kwargs["tags"], list):
-            kwargs["tags"].extend(tag_names)
-        elif "tags" in kwargs and isinstance(kwargs["tags"], str):
-            tag_names.append(kwargs["tags"])
-            kwargs["tags"] = tag_names
-
-    return kwargs
-
-
-@dataclass
-class AxolotlTrainingArguments(TrainingArguments):
-    """
-    Extend the base TrainingArguments for axolotl helpers
-    """
-
-    model_type: Optional[str] = field(
-        default=None, metadata={"help": "HF model configuration model_type."}
-    )
-    lr_quadratic_warmup: bool = field(
-        default=False,
-        metadata={"help": "Use quadratic warmup for cosine scheduling."},
-    )
-    pretraining: bool = field(
-        default=False,
-        metadata={
-            "help": "Indicates to trainer whether we are doing continued pretraining."
-        },
-    )
-    sample_packing: bool = field(
-        default=False,
-        metadata={"help": "Use sample packing for efficient training."},
-    )
-    multipack_real_batches: bool = field(
-        default=False,
-        metadata={"help": "Use real batches for efficient training."},
-    )
-    eval_sample_packing: Optional[bool] = field(
-        default=None,
-        metadata={"help": "Use sample packing for efficient evals."},
-    )
-    sample_packing_efficiency: float = field(
-        default=1.0,
-        metadata={"help": "Sample packing efficiency for calculating batch length."},
-    )
-    max_seq_length: int = field(
-        default=2048,
-        metadata={"help": "The maximum sequence length the model can handle"},
-    )
-    sample_packing_seq_len_multiplier: int = field(
-        default=1,
-        metadata={"help": "the multiplier for the max len for packed sequences"},
-    )
-    relora_steps: Optional[int] = field(
-        default=None,
-        metadata={"help": "how often to reset for ReLoRA"},
-    )
-    relora_warmup_steps: Optional[int] = field(
-        default=None,
-        metadata={"help": "how many warmup steps to take after reset for ReLoRA"},
-    )
-    relora_anneal_steps: Optional[int] = field(
-        default=None,
-        metadata={"help": "how many warmup steps to take after reset for ReLoRA"},
-    )
-    relora_prune_ratio: Optional[float] = field(
-        default=0.9,
-        metadata={"help": "prune ratio for magnitude pruning of the optimizer"},
-    )
-    bench_split: Optional[str] = field(
-        default="eval", metadata={"help": "The benchmark split to run on"}
-    )
-    bench_dataset: Optional[str] = field(
-        default="pharaouk/dharma-1/dharma_1_mini.json",
-        metadata={
-            "help": "Benchmark dataset to use: options are `mmlu-zs`, `mmlu-fs`, or the full path to the dataset file"
-        },
-    )
-    do_bench_eval: Optional[bool] = field(
-        default=False, metadata={"help": "Whether to run the Benchmark evaluation."}
-    )
-    do_causal_lm_eval: Optional[bool] = field(
-        default=False, metadata={"help": "Whether to run the Causal LM evaluation."}
-    )
-    max_bench_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "If set, only evaluates on `max_bench_samples` of the benchmark dataset."
-        },
-    )
-    bench_source_max_len: int = field(
-        default=2048, metadata={"help": "Maximum source sequence length for bench."}
-    )
-    dataloader_prefetch_factor: Optional[int] = field(
-        default=None,
-        metadata={"help": "prefetch_factor argument to the dataloader"},
-    )
-    cosine_min_lr_ratio: Optional[float] = field(
-        default=None,
-        metadata={"help": "Minimum learning rate is min_lr_ratio * learning_rate"},
-    )
-    cosine_constant_lr_ratio: Optional[float] = field(
-        default=None,
-        metadata={
-            "help": "Starting constant learning rate step is cosine_constant_lr_ratio * max_steps"
-        },
-    )
-    loraplus_lr_ratio: Optional[float] = field(
-        default=None, metadata={"help": "loraplus learning rate ratio lr_B / lr_A."}
-    )
-    loraplus_lr_embedding: Optional[float] = field(
-        default=1e-6,
-        metadata={"help": "loraplus learning rate for lora embedding layers."},
-    )
-    qlora: bool = field(
-        default=False,
-        metadata={"help": "whether this is a qlora training"},
-    )
-    orpo_alpha: Optional[float] = field(
-        default=None,
-    )
-    lisa_n_layers: Optional[int] = field(
-        default=None,
-        metadata={"help": "the number of activate layers in LISA"},
-    )
-    lisa_step_interval: Optional[int] = field(
-        default=None,
-        metadata={"help": "how often to switch layers in LISA"},
-    )
-    lisa_layers_attribute: Optional[str] = field(
-        default=None,
-        metadata={"help": "path under the model to access the layers"},
-    )
-    curriculum_sampling: Optional[bool] = field(
-        default=None,
-        metadata={"help": "whether to use sequential sampling for curriculum learning"},
-    )
-
-
-class AxolotlTrainer(Trainer):
-    """
-    Extend the base Trainer for axolotl helpers
-    """
-
-    args = None  # type: AxolotlTrainingArguments
-    tag_names = ["axolotl"]
-
-    def __init__(
-        self,
-        *_args,
-        num_epochs=1,
-        bench_data_collator=None,
-        eval_data_collator=None,
-        **kwargs,
-    ):
-        self.num_epochs = num_epochs
-        self.bench_data_collator = bench_data_collator
-        self.eval_data_collator = eval_data_collator
-        super().__init__(*_args, **kwargs)
-        self.train_data_collator = self.data_collator
-        self._stored_metrics = defaultdict(lambda: defaultdict(list))
-        if self.args.orpo_alpha:
-            self.loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
-
-    def create_optimizer(self):
-        if self.args.loraplus_lr_ratio is None:
-            return super().create_optimizer()
-
-        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
-        if self.optimizer is None:  # pylint: disable=access-member-before-definition
-            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
-                self.args,
-                opt_model,
-            )
-
-            loraplus_lr_ratio = getattr(self.args, "loraplus_lr_ratio", None)
-            loraplus_lr_embedding = getattr(self.args, "loraplus_lr_embedding", None)
-            self.optimizer = create_loraplus_optimizer(  # pylint: disable=attribute-defined-outside-init
-                opt_model,
-                optimizer_cls,
-                optimizer_kwargs,
-                loraplus_lr_ratio,
-                loraplus_lr_embedding,
-            )
-
-        if is_sagemaker_mp_enabled():
-            self.optimizer = smp.DistributedOptimizer(  # pylint: disable=attribute-defined-outside-init
-                self.optimizer
-            )
-
-        return self.optimizer
-
-    def create_scheduler(
-        self, num_training_steps: int, optimizer: torch.optim.Optimizer = None
-    ):
-        """
-        Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or
-        passed as an argument.
-
-        Args:
-            num_training_steps (int): The number of training steps to do.
-            optimizer (torch.optim.Optimizer): The training optimizer
-        """
-        use_cosine_quadratic = (
-            self.args.lr_scheduler_type == "cosine"
-            and self.args.lr_quadratic_warmup is True
-        )
-
-        use_cosine_min_lr = (
-            self.args.lr_scheduler_type == "cosine"
-            and self.args.cosine_min_lr_ratio is not None
-        )
-
-        # fmt: off
-        if self.lr_scheduler is None:  # type: ignore  # pylint: disable=access-member-before-definition
-            # fmt: on
-            if use_cosine_quadratic:
-                if use_cosine_min_lr:
-                    LOG.warning("Both cosine quadratic warmup and min lr detected. Using quadratic warmup.")
-
-                self.lr_scheduler = get_cosine_schedule_with_quadratic_warmup(  # pylint: disable=attribute-defined-outside-init
-                    optimizer,
-                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
-                    num_training_steps=num_training_steps,
-                )
-            elif self.args.cosine_min_lr_ratio and self.args.cosine_constant_lr_ratio and use_cosine_min_lr:
-                assert 0 <= self.args.cosine_min_lr_ratio <= 1.0, "cosine_min_lr_ratio must be between 0.0 and 1.0"
-                assert 0 <= self.args.cosine_constant_lr_ratio <= 1.0, "cosine_constant_lr_ratio must be between 0.0 and 1.0"
-                self.lr_scheduler = get_cosine_schedule_with_warmup_decay_constant(  # pylint: disable=attribute-defined-outside-init
-                    optimizer,
-                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
-                    num_training_steps=num_training_steps,
-                    min_lr_ratio=self.args.cosine_min_lr_ratio,
-                    constant_lr_ratio=self.args.cosine_constant_lr_ratio,
-                )
-            elif self.args.cosine_min_lr_ratio and use_cosine_min_lr:
-                assert 0 <= self.args.cosine_min_lr_ratio <= 1.0, "cosine_min_lr_ratio must be between 0.0 and 1.0"
-                self.lr_scheduler = get_cosine_schedule_with_min_lr(  # pylint: disable=attribute-defined-outside-init
-                    optimizer,
-                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
-                    num_training_steps=num_training_steps,
-                    min_lr_ratio=self.args.cosine_min_lr_ratio,
-                )
-            else:
-                return super().create_scheduler(num_training_steps, optimizer)
-        else:
-            if use_cosine_quadratic:
-                LOG.warning("axolotl's cosine scheduler with quadratic warmup not used (e.g., because of deepspeed).")
-
-            if use_cosine_min_lr:
-                LOG.warning("axolotl's cosine scheduler with min lr not used (e.g., because of deepspeed).")
-
-        return self.lr_scheduler
-
-    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
-        if self.args.sample_packing and not self.args.pretraining:
-            if self.args.multipack_real_batches:
-                batch_size = self.args.per_device_train_batch_size
-                batch_max_len = self.args.max_seq_length
-            else:
-                batch_size = 1
-                batch_max_len = (
-                    self.args.per_device_train_batch_size * self.args.max_seq_length
-                )
-            return MultipackBatchSampler(
-                RandomSampler(self.train_dataset),
-                batch_size=batch_size,
-                drop_last=True,
-                batch_max_len=batch_max_len,
-                lengths=get_dataset_lengths(self.train_dataset),
-                packing_efficiency_estimate=self.args.sample_packing_efficiency,
-            )
-        if self.args.curriculum_sampling:
-            return SequentialSampler(self.train_dataset)
-        return super()._get_train_sampler()
-
-    def _get_eval_sampler(
-        self, eval_dataset: Dataset
-    ) -> Optional[torch.utils.data.Sampler]:
-        if self.args.sample_packing and self.args.eval_sample_packing is not False:
-            if self.args.multipack_real_batches:
-                batch_size = self.args.per_device_eval_batch_size
-                batch_max_len = self.args.max_seq_length
-            else:
-                batch_size = 1
-                batch_max_len = (
-                    self.args.per_device_eval_batch_size * self.args.max_seq_length
-                )
-            return MultipackBatchSampler(
-                SequentialSampler(eval_dataset),
-                batch_size=batch_size,
-                drop_last=True,
-                batch_max_len=batch_max_len,
-                lengths=get_dataset_lengths(eval_dataset),
-                packing_efficiency_estimate=self.args.sample_packing_efficiency,
-            )
-        return super()._get_eval_sampler(eval_dataset)
-
-    def get_train_dataloader(self) -> DataLoader:
-        if self.args.sample_packing and not self.args.pretraining:
-            train_dataset = self.train_dataset
-            if "length" in train_dataset.features.keys():
-                train_dataset = train_dataset.remove_columns(["length"])
-            data_collator = self.data_collator
-            dataloader_params = {
-                "batch_size": self._train_batch_size,
-                "collate_fn": data_collator,
-                "num_workers": self.args.dataloader_num_workers,
-                "pin_memory": self.args.dataloader_pin_memory,
-            }
-            if self.args.dataloader_prefetch_factor:
-                dataloader_params[
-                    "prefetch_factor"
-                ] = self.args.dataloader_prefetch_factor
-
-            sampler = self._get_train_sampler()
-            if isinstance(sampler, BatchSampler):
-                dataloader_params["batch_sampler"] = sampler
-                del dataloader_params["batch_size"]
-            else:
-                dataloader_params["sampler"] = sampler
-                dataloader_params["drop_last"] = self.args.dataloader_drop_last
-            dataloader_params["worker_init_fn"] = seed_worker
-
-            self.accelerator.even_batches = False
-            return self.accelerator.prepare_data_loader(
-                DataLoader(train_dataset, **dataloader_params)
-            )
-        return super().get_train_dataloader()
-
-    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
-        if self.args.sample_packing and self.args.eval_sample_packing is False:
-            self.data_collator = (  # pylint: disable=attribute-defined-outside-init
-                self.eval_data_collator
-            )
-            dataloader = super().get_eval_dataloader(eval_dataset)
-            self.data_collator = (  # pylint: disable=attribute-defined-outside-init
-                self.train_data_collator
-            )
-            return dataloader
-
-        if self.args.sample_packing and self.args.eval_sample_packing is not False:
-            eval_dataset = (
-                eval_dataset if eval_dataset is not None else self.eval_dataset
-            )
-
-            eval_sampler = self._get_eval_sampler(eval_dataset)
-            eval_dataset = eval_dataset.remove_columns(["length"])
-            data_collator = self.data_collator
-            dataloader_params = {
-                "batch_size": self.args.eval_batch_size,
-                "collate_fn": data_collator,
-                "num_workers": self.args.dataloader_num_workers,
-                "pin_memory": self.args.dataloader_pin_memory,
-            }
-            if self.args.dataloader_prefetch_factor:
-                dataloader_params[
-                    "prefetch_factor"
-                ] = self.args.dataloader_prefetch_factor
-
-            if isinstance(eval_sampler, BatchSampler):
-                dataloader_params["batch_sampler"] = eval_sampler
-                del dataloader_params["batch_size"]
-            else:
-                dataloader_params["sampler"] = eval_sampler
-                dataloader_params["drop_last"] = self.args.dataloader_drop_last
-
-            self.accelerator.even_batches = False
-            return self.accelerator.prepare_data_loader(
-                DataLoader(eval_dataset, **dataloader_params)
-            )
-
-        return super().get_eval_dataloader(eval_dataset)
-
-    def _get_bench_sampler(
-        self, bench_dataset: Dataset
-    ) -> Optional[torch.utils.data.Sampler]:
-        if self.args.world_size <= 1:
-            return SequentialSampler(bench_dataset)
-        return None
-
-    def get_bench_dataloader(
-        self,
-        bench_dataset: Dataset,
-    ) -> DataLoader:
-        dataloader_params = {
-            "batch_size": self.args.eval_batch_size,
-            "collate_fn": self.bench_data_collator,
-            "num_workers": self.args.dataloader_num_workers,
-            "pin_memory": self.args.dataloader_pin_memory,
-        }
-        if self.args.dataloader_prefetch_factor:
-            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
-
-        if not isinstance(bench_dataset, torch.utils.data.IterableDataset):
-            dataloader_params["sampler"] = self._get_bench_sampler(bench_dataset)
-            dataloader_params["drop_last"] = self.args.dataloader_drop_last
-
-        return DataLoader(bench_dataset, **dataloader_params)
-        # return self.accelerator.prepare(DataLoader(bench_dataset, **dataloader_params))
-
-    def compute_loss(self, model, inputs, return_outputs=False):
-        # use one's weighted cross entropy loss calc
-        # if self.args.sample_packing:
-        #     labels = inputs.pop("labels")
-        #     outputs = model(**inputs)
-        #     loss = trainer_weighted_loss(outputs, labels, shift_labels=True)
-        #     return (loss, outputs) if return_outputs else loss
-        if self.args.orpo_alpha:
-            return self.orpo_compute_loss(model, inputs, return_outputs=return_outputs)
-        return super().compute_loss(model, inputs, return_outputs=return_outputs)
-
-    @staticmethod
-    def orpo_concatenate_inputs(inputs, label_pad_token=-100, pad_token=0, device=None):
-        concatenated_batch = {}
-
-        max_length = max(
-            inputs["input_ids"].shape[1], inputs["rejected_input_ids"].shape[1]
-        )
-        # Concatenate positive and negative inputs
-        concatenated_batch["input_ids"] = pad_to_length(
-            inputs["input_ids"], max_length, pad_token
-        )
-        concatenated_batch["rejected_input_ids"] = pad_to_length(
-            inputs["rejected_input_ids"], max_length, pad_token
-        )
-        concatenated_batch["labels"] = pad_to_length(
-            inputs["labels"], max_length, label_pad_token
-        )
-        concatenated_batch["rejected_labels"] = pad_to_length(
-            inputs["rejected_labels"], max_length, label_pad_token
-        )
-        concatenated_batch["attention_mask"] = pad_to_length(
-            inputs["attention_mask"], max_length, 0
-        )
-        concatenated_batch["rejected_attention_mask"] = pad_to_length(
-            inputs["rejected_attention_mask"], max_length, 0
-        )
-        concatenated_batch["prompt_attention_mask"] = pad_to_length(
-            inputs["prompt_attention_mask"], max_length, 0
-        ).to(device=device)
-
-        input_ids = torch.cat(
-            [concatenated_batch["input_ids"], concatenated_batch["rejected_input_ids"]],
-            dim=0,
-        ).to(device=device)
-        attention_mask = torch.cat(
-            [
-                concatenated_batch["attention_mask"],
-                concatenated_batch["rejected_attention_mask"],
-            ],
-            dim=0,
-        ).to(device=device)
-        labels = torch.cat(
-            [concatenated_batch["labels"], concatenated_batch["rejected_labels"]], dim=0
-        ).to(device=device)
-
-        return {
-            "input_ids": input_ids,
-            "labels": labels,
-            "attention_mask": attention_mask,
-            "prompt_attention_mask": concatenated_batch["prompt_attention_mask"],
-        }
-
-    def orpo_compute_custom_loss(self, logits, labels):
-        logits = logits.contiguous()
-        loss = 0.0
-
-        if labels is not None:
-            # move labels to correct device to enable model parallelism
-            labels = labels.to(logits.device)
-            # Shift so that tokens < n predict n
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-
-            # Flatten the tokens
-            loss = self.loss_fct(shift_logits.transpose(2, 1), shift_labels).mean(
-                dim=-1
-            )
-
-        return loss
-
-    def orpo_compute_logps(
-        self, prompt_attention_mask, chosen_inputs, chosen_attention_mask, logits
-    ):
-        # Get the shape of chosen_attention_mask[:, :-1]
-        chosen_shape = chosen_attention_mask[:, :-1].shape
-
-        # Calculate the padding size
-        pad_length = chosen_shape[1] - (prompt_attention_mask.shape[1] - 1)
-
-        # Pad prompt_attention_mask with zeros to match the desired shape
-        prompt_attention_mask_padded = torch.nn.functional.pad(
-            prompt_attention_mask[:, 1:], (0, pad_length), mode="constant", value=0
-        )
-
-        # Perform the subtraction operation
-        mask = chosen_attention_mask[:, :-1] > prompt_attention_mask_padded
-
-        per_token_logps = torch.gather(
-            logits[:, :-1, :].log_softmax(-1),
-            dim=2,
-            index=(mask * chosen_inputs[:, 1:]).unsqueeze(2),
-        ).squeeze(2)
-        return torch.mul(per_token_logps, mask).sum(dim=1) / mask.sum(dim=1)
-
-    def orpo_compute_loss(self, model, inputs, return_outputs=False):
-        concat_inputs = AxolotlTrainer.orpo_concatenate_inputs(
-            inputs,
-            label_pad_token=-100,
-            pad_token=self.tokenizer.pad_token_id,
-            device=self.accelerator.device,
-        )
-
-        # Perform a single forward pass
-        outputs = model(
-            **{
-                "input_ids": concat_inputs["input_ids"],
-                "attention_mask": concat_inputs["attention_mask"],
-                "labels": concat_inputs["labels"],
-            },
-            output_hidden_states=True,
-        )
-
-        # Split the outputs for positive and negative examples
-        outputs_pos, outputs_neg = outputs.logits.chunk(2)
-
-        # Calculate NLL loss
-        pos_loss = self.orpo_compute_custom_loss(
-            logits=outputs_pos, labels=concat_inputs["input_ids"].chunk(2)[0]
-        )
-
-        # Calculate Log Probability
-        pos_prob = self.orpo_compute_logps(
-            prompt_attention_mask=concat_inputs["prompt_attention_mask"],
-            chosen_inputs=concat_inputs["input_ids"].chunk(2)[0],
-            chosen_attention_mask=concat_inputs["attention_mask"].chunk(2)[0],
-            logits=outputs_pos,
-        )
-        neg_prob = self.orpo_compute_logps(
-            prompt_attention_mask=concat_inputs["prompt_attention_mask"],
-            chosen_inputs=concat_inputs["input_ids"].chunk(2)[1],
-            chosen_attention_mask=concat_inputs["attention_mask"].chunk(2)[1],
-            logits=outputs_neg,
-        )
-
-        # Calculate log odds
-        log_odds = (pos_prob - neg_prob) - (
-            torch.log(1 - torch.exp(pos_prob)) - torch.log(1 - torch.exp(neg_prob))
-        )
-        sig_ratio = torch.nn.functional.sigmoid(log_odds)
-        ratio = torch.log(sig_ratio)
-
-        # Calculate the Final Loss
-        loss = torch.mean(pos_loss - self.args.orpo_alpha * ratio).to(
-            dtype=torch.bfloat16
-        )
-
-        metrics = {}
-        metrics["chosen_geometric_mean"] = torch.mean(pos_prob).cpu().item()
-        metrics["rejected_geometric_mean"] = torch.mean(neg_prob).cpu().item()
-        metrics["log_odds_ratio"] = torch.mean(ratio).cpu().item()
-        metrics["log_odds"] = torch.mean(log_odds).cpu().item()
-        self.store_metrics(metrics, train_eval="train")
-
-        return (loss, outputs_pos) if return_outputs else loss
-
-    @wraps(Trainer.push_to_hub)
-    def push_to_hub(self, *args, **kwargs) -> str:
-        """
-        Overwrite the `push_to_hub` method in order to force-add the tags when pushing the
-        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.
-        """
-        kwargs = _sanitize_kwargs_for_tagging(tag_names=self.tag_names, kwargs=kwargs)
-
-        return super().push_to_hub(*args, **kwargs)
-
-    @wraps(Trainer.create_accelerator_and_postprocess)
-    def create_accelerator_and_postprocess(self):
-        res = super().create_accelerator_and_postprocess()
-
-        if self.is_fsdp_enabled:
-            if (
-                "limit_all_gathers" in self.args.fsdp_config
-                and self.args.fsdp_config["limit_all_gathers"]
-            ):
-                self.accelerator.state.fsdp_plugin.limit_all_gathers = True
-
-        return res
-
-    def log(self, logs: Dict[str, float]) -> None:
-        """
-        Log `logs` on the various objects watching training, including stored metrics.
-
-        Args:
-            logs (`Dict[str, float]`):
-                The values to log.
-        """
-        # logs either has 'loss' or 'eval_loss'
-        train_eval = "train" if "loss" in logs else "eval"
-        # Add averaged stored metrics to logs
-        for key, metrics in self._stored_metrics[train_eval].items():
-            logs[key] = torch.tensor(metrics).mean().item()
-        del self._stored_metrics[train_eval]
-        return super().log(logs)
-
-    def store_metrics(
-        self, metrics: Dict[str, float], train_eval: Literal["train", "eval"] = "train"
-    ) -> None:
-        for key, value in metrics.items():
-            self._stored_metrics[train_eval][key].append(value)
-
-
-class AxolotlMambaTrainer(AxolotlTrainer):
-    """
-    Mamba specific trainer to handle loss calculation
-    """
-
-    tag_names = ["axolotl", "mamba"]
-
-    def compute_loss(
-        self,
-        model,
-        inputs,
-        return_outputs=False,  # pylint: disable=unused-argument
-    ):
-        input_ids = inputs.pop("input_ids")
-        lm_logits = model(input_ids).logits
-
-        labels = input_ids.to(lm_logits.device)
-        shift_logits = lm_logits[:, :-1, :].contiguous()
-        labels = labels[:, 1:].contiguous()
-
-        loss_fct = torch.nn.CrossEntropyLoss()
-        lm_loss = loss_fct(
-            shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1)
-        )
-
-        return lm_loss
-
-
-class OneCycleLRSchedulerTrainer(AxolotlTrainer):
-    """
-    Trainer subclass that uses the OneCycleLR scheduler
-    """
-
-    tag_names = ["axolotl", "onecycle"]
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.lr_scheduler = None
-
-    def create_scheduler(
-        self,
-        num_training_steps: int,
-        optimizer: Optional[torch.optim.Optimizer] = None,
-    ):
-        optimizer = self.optimizer if optimizer is None else optimizer
-        num_warmup_steps = self.args.get_warmup_steps(num_training_steps)
-        pct_start = num_warmup_steps / num_training_steps
-
-        self.lr_scheduler = OneCycleLR(
-            optimizer,
-            max_lr=self.args.learning_rate,
-            total_steps=num_training_steps,
-            pct_start=pct_start,
-            div_factor=6,
-        )
-
-        return self.lr_scheduler
-
-
-class ReLoRATrainer(AxolotlTrainer):
-    """
-    Trainer subclass that uses the OneCycleLR scheduler
-    """
-
-    tag_names = ["axolotl", "relora"]
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.lr_scheduler = None
-
-    def create_scheduler(
-        self,
-        num_training_steps: int,
-        optimizer: Optional[torch.optim.Optimizer] = None,
-    ):
-        optimizer = self.optimizer if optimizer is None else optimizer
-        lr_scheduler = super().create_scheduler(num_training_steps, optimizer)
-
-        if self.args.relora_steps:
-            warmup_steps = (
-                self.args.relora_warmup_steps if self.args.relora_warmup_steps else 10
-            )
-            anneal_steps = (
-                self.args.relora_anneal_steps if self.args.relora_anneal_steps else 1
-            )
-            self.lr_scheduler = ReLoRAScheduler(
-                optimizer,
-                lr_scheduler,
-                self.args.relora_steps,
-                anneal_steps,
-                warmup_steps,
-            )
-        else:
-            self.lr_scheduler = lr_scheduler
-
-        return self.lr_scheduler
-
-
-class AxolotlDPOTrainer(DPOTrainer):
-    """
-    Extend the base DPOTrainer for axolotl helpers
-    """
-
-    tag_names = ["axolotl", "dpo"]
-
-    @wraps(DPOTrainer.push_to_hub)
-    def push_to_hub(self, *args, **kwargs) -> str:
-        """
-        Overwrite the `push_to_hub` method in order to force-add the tags when pushing the
-        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.
-        """
-        kwargs = _sanitize_kwargs_for_tagging(tag_names=self.tag_names, kwargs=kwargs)
-
-        return super().push_to_hub(*args, **kwargs)
-
-    def tokenize_row(
-        self, feature, model: Optional[Union[PreTrainedModel, torch.nn.Module]] = None
-    ) -> Dict:
-        res = super().tokenize_row(feature, model=model)
-        if self.tokenizer.bos_token_id is None and res["prompt_input_ids"][0] is None:
-            for key in res.keys():
-                res[key] = res[key][1:]
-        return res
-
-
-class AxolotlORPOTrainer(ORPOTrainer):
-    """
-    Extend the base ORPOTrainer for axolotl helpers
-    """
-
-    tag_names = ["axolotl", "orpo"]
-
-
-class TrainerBuilderBase(abc.ABC):
-    """
-    Base class for trainer builder
-    """
-
-    _train_dataset = None
-    _eval_dataset = None
-    _model_ref = None
-    _peft_config = None
-
-    def __init__(self, cfg, model, tokenizer):
-        self.cfg = cfg
-        self.model = model
-        self.tokenizer = tokenizer
-
-        # in case the model supports tagging, add the axolotl tag.
-        # This makes sure the tag is correctly pushed even if a user calls
-        # model.push_to_hub instad of  trainer.push_to_hub.
-        if hasattr(model, "add_model_tags"):
-            model.add_model_tags(["axolotl"])
-
-    @property
-    def model_ref(self):
-        return self._model_ref
-
-    @model_ref.setter
-    def model_ref(self, model):
-        self._model_ref = model
-
-    @property
-    def train_dataset(self):
-        return self._train_dataset
-
-    @train_dataset.setter
-    def train_dataset(self, dataset):
-        self._train_dataset = dataset
-
-    @property
-    def eval_dataset(self):
-        return self._eval_dataset
-
-    @eval_dataset.setter
-    def eval_dataset(self, dataset):
-        self._eval_dataset = dataset
-
-    @property
-    def peft_config(self):
-        return self._peft_config
-
-    @peft_config.setter
-    def peft_config(self, peft_config):
-        self._peft_config = peft_config
-
-    @abstractmethod
-    def build(self, total_num_steps):
-        pass
-
-    def get_callbacks(self) -> List[TrainerCallback]:
-        callbacks = []
-        if self.cfg.use_wandb:
-            callbacks.append(
-                SaveAxolotlConfigtoWandBCallback(self.cfg.axolotl_config_path)
-            )
-        if self.cfg.use_mlflow and is_mlflow_available():
-            from axolotl.utils.callbacks.mlflow_ import (
-                SaveAxolotlConfigtoMlflowCallback,
-            )
-
-            callbacks.append(
-                SaveAxolotlConfigtoMlflowCallback(self.cfg.axolotl_config_path)
-            )
-
-        return callbacks
-
-    @abstractmethod
-    def get_post_trainer_create_callbacks(self, trainer):
-        """
-        Callbacks added after the trainer is created, usually b/c these need access to the trainer
-        """
-
-    def hook_pre_create_training_args(self, training_arguments_kwargs):
-        # TODO
-        return training_arguments_kwargs
-
-    def hook_post_create_training_args(self, training_arguments):
-        # TODO
-        return training_arguments
-
-    def hook_pre_create_trainer(self, trainer_kwargs, trainer_cls):
-        # TODO
-        return trainer_kwargs, trainer_cls
-
-    def hook_post_create_trainer(self, trainer):
-        # TODO
-        return trainer
-
-
-class HFCausalTrainerBuilder(TrainerBuilderBase):
-    """
-    Build the HuggingFace training args/trainer for Causal models
-    """
-
-    def get_callbacks(self):
-        callbacks = super().get_callbacks()
-        callbacks.append(GPUStatsCallback(self.cfg))
-        callbacks.append(EvalFirstStepCallback())
-
-        if self.cfg.relora_steps:
-            callbacks.append(ReLoRACallback(self.cfg))
-
-        if (
-            hasattr(self.model, "use_bettertransformer")
-            and self.model.use_bettertransformer is True
-        ):
-            callbacks.append(SaveBetterTransformerModelCallback())
-
-        if self.cfg.loss_watchdog_threshold is not None:
-            callbacks.append(LossWatchDogCallback(self.cfg))
-
-        callbacks.append(SaveModelOnTrainEndCallback())
-
-        return callbacks
-
-    def get_post_trainer_create_callbacks(self, trainer):
-        callbacks = []
-        if self.cfg.use_wandb and self.cfg.eval_table_size > 0:
-            LogPredictionCallback = log_prediction_callback_factory(
-                trainer, self.tokenizer, "wandb"
-            )
-            callbacks.append(LogPredictionCallback(self.cfg))
-        if (
-            self.cfg.use_mlflow
-            and is_mlflow_available()
-            and self.cfg.eval_table_size > 0
-        ):
-            LogPredictionCallback = log_prediction_callback_factory(
-                trainer, self.tokenizer, "mlflow"
-            )
-            callbacks.append(LogPredictionCallback(self.cfg))
-
-        if self.cfg.do_bench_eval:
-            callbacks.append(bench_eval_callback_factory(trainer, self.tokenizer))
-        if self.cfg.do_causal_lm_eval:
-            CausalLMBenchEvalCallback = causal_lm_bench_eval_callback_factory(
-                trainer, self.tokenizer
-            )
-            callbacks.append(CausalLMBenchEvalCallback(self.cfg))
-
-        if self.cfg.early_stopping_patience:
-            early_stop_cb = EarlyStoppingCallback(
-                self.cfg.early_stopping_patience,
-            )
-            callbacks.append(early_stop_cb)
-
-        if self.cfg.lisa_step_interval and self.cfg.lisa_n_layers:
-            callbacks.append(lisa_callback_factory(trainer))
-        return callbacks
-
-    def _get_trainer_cls(self):
-        if self.cfg.lr_scheduler == "one_cycle" and (
-            self.cfg.fsdp or self.cfg.adapter == "qlora"
-        ):
-            return OneCycleLRSchedulerTrainer
-        if self.cfg.relora_steps:
-            return ReLoRATrainer
-        if self.cfg.model_config_type == "mamba":
-            return AxolotlMambaTrainer
-        return AxolotlTrainer
-
-    def build(self, total_num_steps):
-        warmup_steps = None
-        if self.cfg.warmup_steps is not None:
-            warmup_steps = self.cfg.warmup_steps
-        elif self.cfg.warmup_ratio is not None:
-            warmup_steps = max(int(self.cfg.warmup_ratio * total_num_steps), 0)
-        else:
-            warmup_steps = min(int(0.03 * total_num_steps), 100)
-
-        logging_steps = (
-            self.cfg.logging_steps
-            if self.cfg.logging_steps is not None
-            else max(min(int(0.005 * total_num_steps), 10), 1)
-        )
-
-        training_arguments_kwargs = {}
-        if self.cfg.bf16 == "full":
-            training_arguments_kwargs["bf16_full_eval"] = True
-        else:
-            training_arguments_kwargs["bf16"] = self.cfg.bf16
-        training_arguments_kwargs["fp16"] = (
-            self.cfg.fp16 and not self.cfg.bf16
-        ) or False
-        training_arguments_kwargs["tf32"] = self.cfg.tf32
-        training_arguments_kwargs["warmup_steps"] = warmup_steps
-        training_arguments_kwargs["logging_steps"] = logging_steps
-
-        if self.cfg.seed:
-            training_arguments_kwargs["seed"] = self.cfg.seed
-
-        if self.cfg.gradient_checkpointing:
-            training_arguments_kwargs[
-                "gradient_checkpointing"
-            ] = self.cfg.gradient_checkpointing
-            if self.cfg.gradient_checkpointing_kwargs is not None:
-                training_arguments_kwargs[
-                    "gradient_checkpointing_kwargs"
-                ] = self.cfg.gradient_checkpointing_kwargs
-        if self.cfg.fsdp:
-            training_arguments_kwargs["fsdp"] = self.cfg.fsdp
-            if self.cfg.fsdp_config:
-                training_arguments_kwargs["fsdp_config"] = dict(self.cfg.fsdp_config)
-
-        if self.cfg.adapter == "qlora":
-            training_arguments_kwargs["qlora"] = True
-
-        # deepspeed
-        if self.cfg.deepspeed:
-            training_arguments_kwargs["deepspeed"] = self.cfg.deepspeed
-
-        if self.cfg.lr_quadratic_warmup is not None:
-            training_arguments_kwargs[
-                "lr_quadratic_warmup"
-            ] = self.cfg.lr_quadratic_warmup
-
-        if self.cfg.adam_beta1:
-            training_arguments_kwargs["adam_beta1"] = self.cfg.adam_beta1
-        if self.cfg.adam_beta2:
-            training_arguments_kwargs["adam_beta2"] = self.cfg.adam_beta2
-        if self.cfg.adam_epsilon:
-            training_arguments_kwargs["adam_epsilon"] = self.cfg.adam_epsilon
-        if self.cfg.max_grad_norm:
-            training_arguments_kwargs["max_grad_norm"] = self.cfg.max_grad_norm
-
-        if self.cfg.hub_model_id:
-            training_arguments_kwargs["hub_model_id"] = self.cfg.hub_model_id
-            training_arguments_kwargs["push_to_hub"] = True
-            training_arguments_kwargs["hub_private_repo"] = True
-            training_arguments_kwargs["hub_always_push"] = True
-
-            if self.cfg.hub_strategy:
-                training_arguments_kwargs["hub_strategy"] = self.cfg.hub_strategy
-
-        if self.cfg.save_safetensors is not None:
-            training_arguments_kwargs["save_safetensors"] = self.cfg.save_safetensors
-
-        if self.cfg.sample_packing_eff_est:
-            training_arguments_kwargs[
-                "sample_packing_efficiency"
-            ] = self.cfg.sample_packing_eff_est
-
-        if self.cfg.dataloader_pin_memory is not None:
-            training_arguments_kwargs[
-                "dataloader_pin_memory"
-            ] = self.cfg.dataloader_pin_memory
-        if self.cfg.dataloader_num_workers is not None:
-            training_arguments_kwargs[
-                "dataloader_num_workers"
-            ] = self.cfg.dataloader_num_workers
-        if self.cfg.dataloader_prefetch_factor is not None:
-            training_arguments_kwargs[
-                "dataloader_prefetch_factor"
-            ] = self.cfg.dataloader_prefetch_factor
-        if self.cfg.dataloader_drop_last is not None:
-            training_arguments_kwargs[
-                "dataloader_drop_last"
-            ] = self.cfg.dataloader_drop_last
-        elif self.cfg.sample_packing and self.cfg.eval_sample_packing is False:
-            training_arguments_kwargs["dataloader_drop_last"] = True
-
-        if self.cfg.remove_unused_columns is not None:
-            training_arguments_kwargs[
-                "remove_unused_columns"
-            ] = self.cfg.remove_unused_columns
-
-        if not self.cfg.test_datasets and self.cfg.val_set_size == 0:
-            # no eval set, so don't eval
-            training_arguments_kwargs["evaluation_strategy"] = "no"
-        elif self.cfg.eval_steps:
-            training_arguments_kwargs["evaluation_strategy"] = "steps"
-            training_arguments_kwargs["eval_steps"] = self.cfg.eval_steps
-        elif self.cfg.evaluation_strategy:
-            training_arguments_kwargs[
-                "evaluation_strategy"
-            ] = self.cfg.evaluation_strategy
-        else:
-            # we have an eval set, but no steps defined, default to use epoch
-            training_arguments_kwargs["evaluation_strategy"] = "epoch"
-
-        if self.cfg.save_steps:
-            training_arguments_kwargs["save_strategy"] = "steps"
-            training_arguments_kwargs["save_steps"] = self.cfg.save_steps
-        elif self.cfg.save_strategy:
-            training_arguments_kwargs["save_strategy"] = self.cfg.save_strategy
-        else:
-            # default to saving each epoch if not defined
-            training_arguments_kwargs["save_strategy"] = "epoch"
-
-        if self.cfg.do_bench_eval:
-            training_arguments_kwargs["do_bench_eval"] = self.cfg.do_bench_eval
-            if self.cfg.bench_dataset:
-                training_arguments_kwargs["bench_dataset"] = self.cfg.bench_dataset
-        if self.cfg.do_causal_lm_eval:
-            training_arguments_kwargs["do_causal_lm_eval"] = self.cfg.do_causal_lm_eval
-        if self.cfg.metric_for_best_model:
-            training_arguments_kwargs[
-                "metric_for_best_model"
-            ] = self.cfg.metric_for_best_model
-        if self.cfg.greater_is_better:
-            training_arguments_kwargs["greater_is_better"] = self.cfg.greater_is_better
-
-        if self.cfg.torch_compile:
-            if torch.__version__ < "2.1.0":  # pylint: disable=protected-access
-                LOG.warning("torch>=2.1.0 required for torch_compile to work properly")
-            elif torch._dynamo:  # pylint: disable=protected-access
-                torch._dynamo.config.suppress_errors = (  # pylint: disable=protected-access
-                    True
-                )
-                training_arguments_kwargs["torch_compile"] = self.cfg.torch_compile
-                if self.cfg.torch_compile_backend:
-                    training_arguments_kwargs[
-                        "torch_compile_backend"
-                    ] = self.cfg.torch_compile_backend
-
-        # DDP Config
-        if self.cfg.ddp_timeout:
-            training_arguments_kwargs["ddp_timeout"] = self.cfg.ddp_timeout
-        # see https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html
-        if self.cfg.ddp_bucket_cap_mb:
-            training_arguments_kwargs["ddp_bucket_cap_mb"] = self.cfg.ddp_bucket_cap_mb
-        if self.cfg.ddp_broadcast_buffers is not None:
-            training_arguments_kwargs[
-                "ddp_broadcast_buffers"
-            ] = self.cfg.ddp_broadcast_buffers
-
-        # these are all the "standard" kwargs that are def used
-        training_arguments_kwargs["max_steps"] = (
-            total_num_steps if self.cfg.max_steps else -1
-        )
-        training_arguments_kwargs["max_seq_length"] = self.cfg.sequence_len
-        training_arguments_kwargs[
-            "per_device_train_batch_size"
-        ] = self.cfg.micro_batch_size
-        if self.cfg.eval_batch_size:
-            training_arguments_kwargs[
-                "per_device_eval_batch_size"
-            ] = self.cfg.eval_batch_size
-        training_arguments_kwargs[
-            "gradient_accumulation_steps"
-        ] = self.cfg.gradient_accumulation_steps
-        training_arguments_kwargs[
-            "eval_accumulation_steps"
-        ] = self.cfg.gradient_accumulation_steps
-        training_arguments_kwargs["num_train_epochs"] = self.cfg.num_epochs
-        training_arguments_kwargs["learning_rate"] = self.cfg.learning_rate
-        training_arguments_kwargs["output_dir"] = self.cfg.output_dir
-        training_arguments_kwargs["save_total_limit"] = (
-            self.cfg.save_total_limit if self.cfg.save_total_limit else 4
-        )
-        training_arguments_kwargs["load_best_model_at_end"] = (
-            (
-                self.cfg.load_best_model_at_end is not False
-                or self.cfg.early_stopping_patience
-            )
-            and (
-                (not self.cfg.test_datasets and self.cfg.val_set_size > 0)
-                or (self.cfg.test_datasets and self.cfg.val_set_size == 0)
-            )
-            and self.cfg.save_steps
-            and self.cfg.eval_steps
-            and self.cfg.save_steps % self.cfg.eval_steps == 0
-        ) or False
-        training_arguments_kwargs["ddp_find_unused_parameters"] = (
-            False if self.cfg.ddp else None
-        )
-        training_arguments_kwargs["group_by_length"] = self.cfg.group_by_length
-        training_arguments_kwargs["curriculum_sampling"] = self.cfg.curriculum_sampling
-        report_to = None
-        if self.cfg.use_wandb:
-            report_to = "wandb"
-        if self.cfg.use_mlflow:
-            report_to = "mlflow"
-        training_arguments_kwargs["report_to"] = report_to
-        training_arguments_kwargs["run_name"] = (
-            self.cfg.wandb_name if self.cfg.use_wandb else None
-        )
-        training_arguments_kwargs["optim"] = (
-            self.cfg.optimizer if self.cfg.optimizer else "adamw_hf"
-        )
-        if self.cfg.optim_args:
-            if isinstance(self.cfg.optim_args, dict):
-                optim_args = ",".join(
-                    [f"{key}={value}" for key, value in self.cfg.optim_args.items()]
-                )
-            else:
-                optim_args = self.cfg.optim_args
-            training_arguments_kwargs["optim_args"] = optim_args
-        if self.cfg.optim_target_modules:
-            training_arguments_kwargs[
-                "optim_target_modules"
-            ] = self.cfg.optim_target_modules
-        training_arguments_kwargs["loraplus_lr_ratio"] = self.cfg.loraplus_lr_ratio
-        training_arguments_kwargs[
-            "loraplus_lr_embedding"
-        ] = self.cfg.loraplus_lr_embedding
-        training_arguments_kwargs["lr_scheduler_type"] = (
-            self.cfg.lr_scheduler
-            if self.cfg.lr_scheduler
-            and self.cfg.lr_scheduler not in ("one_cycle", "log_sweep")
-            else "cosine"
-        )
-        training_arguments_kwargs["lr_scheduler_kwargs"] = (
-            self.cfg.lr_scheduler_kwargs if self.cfg.lr_scheduler_kwargs else {}
-        )
-        training_arguments_kwargs["cosine_min_lr_ratio"] = self.cfg.cosine_min_lr_ratio
-        training_arguments_kwargs[
-            "cosine_constant_lr_ratio"
-        ] = self.cfg.cosine_constant_lr_ratio
-        training_arguments_kwargs["weight_decay"] = (
-            self.cfg.weight_decay if self.cfg.weight_decay is not None else 0.0
-        )
-        training_arguments_kwargs["sample_packing"] = (
-            self.cfg.sample_packing if self.cfg.sample_packing else False
-        )
-        training_arguments_kwargs["multipack_real_batches"] = (
-            self.cfg.flash_attention is not True
-        )
-        training_arguments_kwargs["eval_sample_packing"] = (
-            self.cfg.sample_packing
-            if self.cfg.eval_sample_packing is not False
-            else False
-        )
-        training_arguments_kwargs[
-            "sample_packing_seq_len_multiplier"
-        ] = self.cfg.micro_batch_size
-        if self.cfg.relora_steps:
-            training_arguments_kwargs["relora_steps"] = self.cfg.relora_steps
-            training_arguments_kwargs[
-                "relora_warmup_steps"
-            ] = self.cfg.relora_warmup_steps
-            if self.cfg.relora_anneal_steps:
-                training_arguments_kwargs[
-                    "relora_anneal_steps"
-                ] = self.cfg.relora_anneal_steps
-            if self.cfg.relora_prune_ratio:
-                training_arguments_kwargs[
-                    "relora_prune_ratio"
-                ] = self.cfg.relora_prune_ratio
-
-        if self.cfg.lisa_step_interval and self.cfg.lisa_n_layers:
-            training_arguments_kwargs["lisa_n_layers"] = self.cfg.lisa_n_layers
-            training_arguments_kwargs[
-                "lisa_step_interval"
-            ] = self.cfg.lisa_step_interval
-            training_arguments_kwargs[
-                "lisa_layers_attribute"
-            ] = self.cfg.lisa_layers_attribute
-
-        training_arguments_kwargs = self.hook_pre_create_training_args(
-            training_arguments_kwargs
-        )
-        training_arguments_kwargs["model_type"] = self.cfg.model_config_type
-        training_arguments_kwargs["pretraining"] = bool(self.cfg.pretraining_dataset)
-
-        if self.cfg.rl == "orpo":
-            training_arguments_kwargs["orpo_alpha"] = self.cfg.orpo_alpha
-
-        if self.cfg.neftune_noise_alpha is not None:
-            training_arguments_kwargs[
-                "neftune_noise_alpha"
-            ] = self.cfg.neftune_noise_alpha
-
-        trainer_kwargs = {}
-
-        if self.cfg.optimizer == "lion_pytorch":
-            from lion_pytorch import Lion
-
-            lion_kwargs = {"lr": training_arguments_kwargs["learning_rate"]}
-            if "weight_decay" in training_arguments_kwargs:
-                lion_kwargs["weight_decay"] = training_arguments_kwargs["weight_decay"]
-
-            if (
-                "adam_beta1" in training_arguments_kwargs
-                and "adam_beta2" in training_arguments_kwargs
-            ):
-                lion_kwargs["betas"] = (
-                    training_arguments_kwargs["adam_beta1"],
-                    training_arguments_kwargs["adam_beta2"],
-                )
-
-            trainer_kwargs["optimizers"] = (
-                Lion(params=self.model.parameters(), **lion_kwargs),
-                None,
-            )
-            # Set default so transformers doesn't throw
-            training_arguments_kwargs["optim"] = "adamw_hf"
-
-        if self.cfg.optimizer == "adamw_anyprecision":
-            if Path(self.cfg.torchdistx_path).exists():
-                sys.path.append(self.cfg.torchdistx_path)
-                importlib.import_module("torchdistx")
-
-        training_args = (
-            AxolotlTrainingArguments(  # pylint: disable=unexpected-keyword-arg
-                **training_arguments_kwargs,
-            )
-        )
-        training_args = self.hook_post_create_training_args(training_args)
-
-        data_collator_kwargs = {
-            "padding": True,  # True/"longest" is the default
-        }
-        if self.cfg.pad_to_sequence_len:
-            data_collator_kwargs["pad_to_multiple_of"] = 64 * math.ceil(
-                self.cfg.sequence_len / 64
-            )
-        else:
-            # A100 is best at 64, while others at 8. Let's use the larger so we don't have to check
-            # https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
-            data_collator_kwargs["pad_to_multiple_of"] = 64
-
-        trainer_cls = self._get_trainer_cls()
-        trainer_kwargs, trainer_cls = self.hook_pre_create_trainer(
-            trainer_kwargs, trainer_cls
-        )
-        trainer = trainer_cls(
-            model=self.model,
-            train_dataset=self.train_dataset,
-            eval_dataset=self.eval_dataset,
-            args=training_args,
-            tokenizer=self.tokenizer,
-            data_collator=self.build_collator(training_args, **data_collator_kwargs),
-            eval_data_collator=self.build_collator(
-                training_args, is_eval=True, **data_collator_kwargs
-            ),
-            bench_data_collator=transformers.DataCollatorForSeq2Seq(
-                self.tokenizer,
-                return_tensors="pt",
-                **data_collator_kwargs,
-            ),
-            callbacks=self.get_callbacks(),
-            num_epochs=self.cfg.num_epochs,
-            **trainer_kwargs,
-        )
-        trainer = self.hook_post_create_trainer(trainer)
-        for callback in self.get_post_trainer_create_callbacks(trainer):
-            trainer.add_callback(callback)
-
-        if self.cfg.deepspeed and self.cfg.sample_packing:
-            trainer.accelerator.state.deepspeed_plugin.deepspeed_config[
-                "train_micro_batch_size_per_gpu"
-            ] = self.cfg.micro_batch_size
-
-        return trainer
-
-    def build_collator(
-        self, training_args: AxolotlTrainingArguments, is_eval=False, **kwargs
-    ):
-        if training_args.pretraining:
-            return None
-
-        if self.cfg.model_config_type == "mamba":
-            return MambaDataCollator(tokenizer=self.tokenizer)
-
-        use_batch_sampler_collator = False
-        if is_eval is False and training_args.sample_packing:
-            use_batch_sampler_collator = True
-        if is_eval and training_args.eval_sample_packing:
-            use_batch_sampler_collator = True
-
-        collator: Type[
-            Union[
-                V2BatchSamplerDataCollatorForSeq2Seq,
-                BatchSamplerDataCollatorForSeq2Seq,
-                DataCollatorForSeq2Seq,
-            ]
-        ]
-        if use_batch_sampler_collator:
-            if self.cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES:
-                collator = V2BatchSamplerDataCollatorForSeq2Seq
-            elif (
-                self.cfg.model_config_type in ["llama"]
-                and self.cfg.flash_attention is not True
-            ):
-                collator = V2BatchSamplerDataCollatorForSeq2Seq
-            else:
-                collator = BatchSamplerDataCollatorForSeq2Seq
-        else:
-            collator = DataCollatorForSeq2Seq
-
-        return collator(
-            self.tokenizer,
-            return_tensors="pt",
-            **kwargs,
-        )
-
-
-class HFRLTrainerBuilder(TrainerBuilderBase):
-    """
-    Trainer factory class for DPO Trainer
-    """
-
-    def get_callbacks(self):
-        callbacks = super().get_callbacks()
-        callbacks.append(SaveModelOnTrainEndCallback())
-
-        return callbacks
-
-    def get_post_trainer_create_callbacks(self, trainer):
-        callbacks = []
-        return callbacks
-
-    def build_training_arguments(self, total_num_steps):
-        training_args_kwargs = {}
-        for arg in [
-            "adam_beta1",
-            "adam_beta2",
-            "adam_epsilon",
-            "dataloader_num_workers",
-            "dataloader_pin_memory",
-        ]:
-            if hasattr(self.cfg, arg) and getattr(self.cfg, arg) is not None:
-                training_args_kwargs[arg] = getattr(self.cfg, arg)
-
-        if self.cfg.hub_model_id:
-            training_args_kwargs["hub_model_id"] = self.cfg.hub_model_id
-            training_args_kwargs["push_to_hub"] = True
-            training_args_kwargs["hub_private_repo"] = True
-            training_args_kwargs["hub_always_push"] = True
-
-            if self.cfg.hub_strategy:
-                training_args_kwargs["hub_strategy"] = self.cfg.hub_strategy
-
-        if self.cfg.save_safetensors is not None:
-            training_args_kwargs["save_safetensors"] = self.cfg.save_safetensors
-
-        if self.eval_dataset:
-            training_args_kwargs["evaluation_strategy"] = "steps"
-            training_args_kwargs["eval_steps"] = self.cfg.eval_steps
-        else:
-            training_args_kwargs["evaluation_strategy"] = "no"
-
-        if self.cfg.bf16 or self.cfg.bfloat16:
-            training_args_kwargs["bf16"] = True
-
-        training_args_kwargs["lr_scheduler_type"] = (
-            self.cfg.lr_scheduler if self.cfg.lr_scheduler else "cosine"
-        )
-        training_args_kwargs["lr_scheduler_kwargs"] = (
-            self.cfg.lr_scheduler_kwargs if self.cfg.lr_scheduler_kwargs else {}
-        )
-        if self.cfg.remove_unused_columns is not None:
-            training_args_kwargs[
-                "remove_unused_columns"
-            ] = self.cfg.remove_unused_columns
-        else:
-            training_args_kwargs["remove_unused_columns"] = False
-
-        if self.cfg.dataloader_pin_memory is not None:
-            training_args_kwargs[
-                "dataloader_pin_memory"
-            ] = self.cfg.dataloader_pin_memory
-        if self.cfg.dataloader_num_workers is not None:
-            training_args_kwargs[
-                "dataloader_num_workers"
-            ] = self.cfg.dataloader_num_workers
-        if self.cfg.dataloader_prefetch_factor is not None:
-            training_args_kwargs[
-                "dataloader_prefetch_factor"
-            ] = self.cfg.dataloader_prefetch_factor
-        if self.cfg.gradient_checkpointing:
-            training_args_kwargs[
-                "gradient_checkpointing"
-            ] = self.cfg.gradient_checkpointing
-            if self.cfg.gradient_checkpointing_kwargs is not None:
-                training_args_kwargs[
-                    "gradient_checkpointing_kwargs"
-                ] = self.cfg.gradient_checkpointing_kwargs
-            else:
-                training_args_kwargs["gradient_checkpointing_kwargs"] = {
-                    "use_reentrant": False
-                }
-
-        # set save_strategy and save_steps
-        if self.cfg.save_steps:
-            training_args_kwargs["save_strategy"] = "steps"
-            training_args_kwargs["save_steps"] = self.cfg.save_steps
-        elif self.cfg.save_strategy:
-            training_args_kwargs["save_strategy"] = self.cfg.save_strategy
-        else:
-            # default to saving each epoch if not defined
-            training_args_kwargs["save_strategy"] = "epoch"
-
-        if self.cfg.orpo_alpha:
-            # trl does some odd mapping of alpha to beta to reuse the beta parameter ???
-            training_args_kwargs["beta"] = self.cfg.orpo_alpha
-
-        training_args_cls = TrainingArguments
-        if self.cfg.rl == "orpo":
-            training_args_cls = ORPOConfig
-            training_args_kwargs["dataset_num_proc"] = self.cfg.dataset_processes
-            training_args_kwargs["max_length"] = self.cfg.sequence_len
-            if self.cfg.max_prompt_len:
-                training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
-
-        training_args = training_args_cls(
-            per_device_train_batch_size=self.cfg.micro_batch_size,
-            max_steps=self.cfg.max_steps or total_num_steps,
-            gradient_accumulation_steps=self.cfg.gradient_accumulation_steps,
-            learning_rate=self.cfg.learning_rate,
-            output_dir=self.cfg.output_dir,
-            warmup_steps=self.cfg.warmup_steps,
-            logging_first_step=True,
-            logging_steps=1,
-            optim=self.cfg.optimizer,
-            save_total_limit=self.cfg.save_total_limit or 5,
-            **training_args_kwargs,
-        )
-
-        return training_args
-
-    def build(self, total_num_steps):
-        training_args = self.build_training_arguments(total_num_steps)
-        dpo_trainer_kwargs = {}
-        if self.cfg.rl == "ipo":
-            dpo_trainer_kwargs["loss_type"] = "ipo"
-            if self.cfg.dpo_label_smoothing:
-                dpo_trainer_kwargs["label_smoothing"] = self.cfg.dpo_label_smoothing
-        elif self.cfg.rl == "kto_pair":
-            dpo_trainer_kwargs["loss_type"] = "kto_pair"
-        if self.eval_dataset:
-            dpo_trainer_kwargs["eval_dataset"] = self.eval_dataset
-        if self.cfg.adapter and self.peft_config:
-            dpo_trainer_kwargs["peft_config"] = self.peft_config
-        if self.cfg.precompute_ref_log_probs is not None:
-            dpo_trainer_kwargs[
-                "precompute_ref_log_probs"
-            ] = self.cfg.precompute_ref_log_probs
-        if self.cfg.rl in ["dpo", "ipo", "kto_pair"]:
-            trainer_cls = AxolotlDPOTrainer
-            dpo_trainer_kwargs["beta"] = self.cfg.dpo_beta or 0.1
-            trainer_cls_args = [self.model, self.model_ref]
-
-            # these aren't used for the ORPO trainer
-            dpo_trainer_kwargs["max_length"] = self.cfg.sequence_len
-            dpo_trainer_kwargs["max_target_length"] = None
-            dpo_trainer_kwargs["max_prompt_length"] = self.cfg.sequence_len
-            dpo_trainer_kwargs["generate_during_eval"] = True
-            if self.cfg.rl == "dpo":
-                dpo_trainer_kwargs["dataset_num_proc"] = self.cfg.dataset_processes
-        elif self.cfg.rl == "orpo":
-            trainer_cls = AxolotlORPOTrainer
-            trainer_cls_args = [self.model]
-        else:
-            raise ValueError(f"Unsupported RL: {self.cfg.rl}")
-        dpo_trainer = trainer_cls(
-            *trainer_cls_args,
-            args=training_args,
-            train_dataset=self.train_dataset,
-            tokenizer=self.tokenizer,
-            callbacks=self.get_callbacks(),
-            **dpo_trainer_kwargs,
-        )
-        if self.cfg.fsdp:
-            ensure_dtype(dpo_trainer.model, dtype=self.cfg.torch_dtype)
-
-        dpo_trainer = self.hook_post_create_trainer(dpo_trainer)
-        for callback in self.get_post_trainer_create_callbacks(dpo_trainer):
-            dpo_trainer.add_callback(callback)
-
-        return dpo_trainer
-
-
-class HFPPOTrainerBuilder(TrainerBuilderBase):
-    """
-    HF Factory class for PPO Trainer
-    """
-
-    def get_callbacks(self):
-        callbacks = []
-        return callbacks
-
-    def get_post_trainer_create_callbacks(self, trainer):
-        callbacks = []
-        return callbacks
-
-    def build(self, total_num_steps):
-        # build PPOConfig
-        pass
diff --git a/src/axolotl/core/trainers/__init__.py b/src/axolotl/core/trainers/__init__.py
index e69de29bb2..3efe45a51a 100644
--- a/src/axolotl/core/trainers/__init__.py
+++ b/src/axolotl/core/trainers/__init__.py
@@ -0,0 +1,35 @@
+"""Init for axolotl.core.trainers"""
+
+# flake8: noqa
+
+from axolotl.utils import make_lazy_getattr
+
+from .base import AxolotlTrainer
+
+# noinspection PyUnresolvedReferences
+__all__ = [
+    "AxolotlTrainer",
+    "AxolotlCPOTrainer",
+    "AxolotlDPOTrainer",
+    "AxolotlEBFTTrainer",
+    "AxolotlKTOTrainer",
+    "AxolotlMambaTrainer",
+    "AxolotlORPOTrainer",
+    "AxolotlPRMTrainer",
+    "AxolotlRewardTrainer",
+    "AxolotlStridedEBFTTrainer",
+]
+
+_LAZY_IMPORTS = {
+    "AxolotlDPOTrainer": ".dpo.trainer",
+    "AxolotlStridedEBFTTrainer": ".ebft.strided",
+    "AxolotlEBFTTrainer": ".ebft.trainer",
+    "AxolotlMambaTrainer": ".mamba",
+    "AxolotlCPOTrainer": ".trl",
+    "AxolotlKTOTrainer": ".trl",
+    "AxolotlORPOTrainer": ".trl",
+    "AxolotlPRMTrainer": ".trl",
+    "AxolotlRewardTrainer": ".trl",
+}
+
+__getattr__ = make_lazy_getattr(_LAZY_IMPORTS, __name__, globals())
diff --git a/src/axolotl/core/trainers/base.py b/src/axolotl/core/trainers/base.py
new file mode 100644
index 0000000000..1c27b97adb
--- /dev/null
+++ b/src/axolotl/core/trainers/base.py
@@ -0,0 +1,994 @@
+"""Module for customized trainers"""
+
+from __future__ import annotations
+
+import contextlib
+import gc
+import json
+import math
+import os
+import time
+from collections import defaultdict
+from functools import partial, wraps
+from typing import Any, Callable, Literal, Optional
+
+import datasets
+import safetensors
+import torch
+from accelerate.state import AcceleratorState
+from datasets import Dataset
+from peft import PeftModel
+from torch.utils.data import (
+    BatchSampler,
+    DataLoader,
+    RandomSampler,
+    Sampler,
+    SequentialSampler,
+)
+from transformers import PreTrainedModel, Trainer
+from transformers.trainer import TRAINING_ARGS_NAME
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, has_length, seed_worker
+from transformers.utils import SAFE_WEIGHTS_NAME, is_peft_available
+from trl.experimental.utils import pad_to_length
+from typing_extensions import override
+
+from axolotl.core.trainers.constants import TOKENS_STATE_FILE
+from axolotl.core.trainers.mixins import (
+    ActivationOffloadingMixin,
+    CheckpointSaveMixin,
+    DistributedParallelMixin,
+    LayerOffloadingMixin,
+    OptimizerMixin,
+    PackingMixin,
+    RngLoaderMixin,
+    SchedulerMixin,
+)
+from axolotl.core.trainers.utils import (
+    sanitize_kwargs_for_ds_tagging,
+    sanitize_kwargs_for_tagging,
+    trainable_tokens_per_sec_per_gpu,
+)
+from axolotl.utils import get_not_null
+from axolotl.utils.bench import get_gpu_memory_usage
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.distributed import (
+    get_world_size,
+    is_distributed,
+    is_main_process,
+)
+from axolotl.utils.logging import get_logger
+from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
+
+LOG = get_logger(__name__)
+
+REDUCTION_FNS = {
+    "mean": torch.mean,
+    "min": torch.min,
+    "max": torch.max,
+    "sum": torch.sum,
+}
+
+
+class AxolotlTrainer(
+    PackingMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    RngLoaderMixin,
+    CheckpointSaveMixin,
+    LayerOffloadingMixin,
+    ActivationOffloadingMixin,
+    DistributedParallelMixin,
+    Trainer,
+):
+    """Extend the base Trainer for axolotl helpers"""
+
+    args = None  # type: "AxolotlTrainingArguments"  # type: ignore[name-defined]
+    tag_names = ["axolotl"]
+    _axolotl_cfg: DictDefault | None = None
+
+    @property
+    def axolotl_cfg(self):
+        return self._axolotl_cfg
+
+    @axolotl_cfg.setter
+    def axolotl_cfg(self, cfg):
+        self._axolotl_cfg = cfg
+
+    def __init__(
+        self,
+        *_args,
+        bench_data_collator=None,
+        eval_data_collator=None,
+        dataset_tags=None,
+        **kwargs,
+    ):
+        self.bench_data_collator = bench_data_collator
+        self.eval_data_collator = eval_data_collator
+        self.dataset_tags = dataset_tags
+        self._signature_columns = None  # workaround for pylint
+
+        super().__init__(*_args, **kwargs)
+
+        # Gemma4 (and similar multimodal models) declare **kwargs in forward() for
+        # extra inputs like mm_token_type_ids.  HF Trainer interprets VAR_KEYWORD as
+        # "the model handles num_items_in_batch internally" and skips the loss ÷
+        # gradient_accumulation_steps normalisation, which inflates the *logged* loss
+        # (the gradient itself is still correct). Override to False when the model
+        # doesn't actually consume num_items_in_batch.
+        if self.model_accepts_loss_kwargs:
+            model_to_check = self.accelerator.unwrap_model(self.model)
+            if hasattr(model_to_check, "base_model"):  # PEFT wrapper
+                model_to_check = model_to_check.base_model
+            if hasattr(model_to_check, "model"):
+                model_to_check = model_to_check.model
+            fwd = getattr(model_to_check, "forward", None)
+            if fwd is not None:
+                import inspect
+
+                params = inspect.signature(fwd).parameters
+                if "num_items_in_batch" not in params:
+                    self.model_accepts_loss_kwargs = False
+
+        self.train_data_collator = self.data_collator
+        self._tkps_prev_trainable: float | None = None
+        self._tkps_prev_time: float | None = None
+        self._stored_metrics = defaultdict(
+            lambda: defaultdict(lambda: {"values": [], "reduction": "mean"})
+        )
+        if self.args.orpo_alpha:
+            self.loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+
+    def _create_multipack_sampler(
+        self, base_sampler: Sampler, dataset: Dataset
+    ) -> MultipackBatchSampler:
+        """
+        Helper method to create a `MultipackBatchSampler` for multipacking sequences
+        for training.
+
+        Args:
+            base_sampler: Sampler to wrap with `MultipackBatchSampler`.
+            dataset: Dataset to sample from.
+
+        Returns:
+            Multipack (sample packing) batch sampler.
+        """
+        if self.args.multipack_real_batches:
+            batch_size = self.args.per_device_train_batch_size
+            batch_max_len = self.args.max_seq_length
+        else:
+            batch_size = 1
+            train_batch_size = (
+                self.state.train_batch_size or self.args.per_device_train_batch_size
+            )
+            batch_max_len = train_batch_size * self.args.max_seq_length
+
+        sampler = MultipackBatchSampler(
+            base_sampler,
+            lengths=get_dataset_lengths(dataset),
+            packing_efficiency_estimate=self.args.sample_packing_efficiency,
+            batch_max_len=batch_max_len,
+            batch_size=batch_size,
+            group_size=self.args.sample_packing_group_size,
+            bin_size=self.args.sample_packing_bin_size,
+            sequential=self.args.sample_packing_sequentially,
+            drop_last=True,
+            num_processes=self.args.dataset_num_proc,
+            mp_start_method=self.args.sample_packing_mp_start_method or "fork",
+        )
+
+        len(sampler)
+        return sampler
+
+    def _get_train_sampler(
+        self, train_dataset: Dataset | None = None
+    ) -> Sampler | None:
+        """
+        Helper method to get the sampler for training. Handles cases for sample packing
+        and curriculum sampling (sequential).
+
+        Returns:
+            If the dataset is non-empty, a sampler is returned, the type of which
+                depends on the passed training args.
+        """
+        # from https://github.com/huggingface/transformers/blob/2166b6b4ff09f6dd3867ab982f262f66482aa968/src/transformers/trainer.py#L969C1-L972C24
+        if train_dataset is None:
+            train_dataset = self.train_dataset
+        if train_dataset is None or not has_length(train_dataset):
+            return None
+
+        use_sample_packing = self.args.sample_packing and not self.args.pretraining
+
+        # Determine the base sampler first
+        if self.args.curriculum_sampling:
+            base_sampler = SequentialSampler(train_dataset)
+        elif use_sample_packing:
+            base_sampler = RandomSampler(train_dataset)
+        else:
+            # Default to parent class implementation for standard random sampling
+            return super()._get_train_sampler(train_dataset)
+
+        # Apply multipack wrapper if needed
+        if use_sample_packing:
+            return self._create_multipack_sampler(
+                base_sampler=base_sampler,
+                dataset=train_dataset,
+            )
+
+        return base_sampler
+
+    def _get_eval_sampler(self, eval_dataset: Dataset | None = None) -> Sampler | None:
+        """
+        Helper method to get the sampler for evaluation. Handles sample packing case.
+
+        Returns:
+            If the dataset is non-empty, a sampler is returned, the type of which
+                depends on the passed training args.
+        """
+        # from https://github.com/huggingface/transformers/blob/2166b6b4ff09f6dd3867ab982f262f66482aa968/src/transformers/trainer.py#L1065C9-L1066C24
+        if eval_dataset is None or not has_length(eval_dataset):
+            return None
+
+        # Multipacking enabled if training is enabled and eval is not explicitly disabled
+        use_multipack = (
+            self.args.sample_packing and self.args.eval_sample_packing is not False
+        )
+
+        # Determine the base sampler
+        if use_multipack:
+            base_sampler = SequentialSampler(eval_dataset)
+        else:
+            return super()._get_eval_sampler(eval_dataset)
+
+        # Apply multipack wrapper if needed
+        if use_multipack:
+            return self._create_multipack_sampler(
+                base_sampler=base_sampler,
+                dataset=eval_dataset,
+            )
+
+        return base_sampler
+
+    def _get_dataloader(
+        self,
+        dataset: Dataset,
+        description: str,
+        batch_size: int,
+        sampler_fn: Optional[Callable[[Dataset], torch.utils.data.Sampler]] = None,
+        is_training: bool = False,
+        dataloader_key: Optional[str] = None,
+    ) -> DataLoader:
+        """Create a [`~torch.utils.data.DataLoader`] from the given dataset."""
+
+        data_collator = self.data_collator if is_training else self.eval_data_collator
+
+        if isinstance(dataset, datasets.Dataset):
+            if is_training:
+                if not self.args.sample_packing or self.args.pretraining:
+                    dataset = self._remove_unused_columns(
+                        dataset, description="training"
+                    )
+            elif (
+                not is_training
+                and self.args.sample_packing
+                and self.args.eval_sample_packing is not False
+            ):
+                batch_size = (
+                    batch_size
+                    if self.args.sample_packing
+                    else self.args.per_device_eval_batch_size
+                )
+            else:
+                dataset = self._remove_unused_columns(dataset, description=description)
+        else:
+            data_collator = self._get_collator_with_removed_columns(
+                self.data_collator, description=description
+            )
+
+        dataloader_params = {
+            "batch_size": batch_size,
+            "collate_fn": data_collator,
+            "num_workers": self.args.dataloader_num_workers,
+            "pin_memory": self.args.dataloader_pin_memory,
+            "persistent_workers": self.args.dataloader_persistent_workers,
+        }
+
+        if not isinstance(dataset, torch.utils.data.IterableDataset):
+            dataloader_params["drop_last"] = get_not_null(
+                self.args.dataloader_drop_last, True
+            )
+            if sampler_fn is not None:
+                sampler = sampler_fn(dataset)
+                if isinstance(sampler, BatchSampler):
+                    # batch_size and batch_sampler are mutually exclusive
+                    dataloader_params["batch_sampler"] = sampler
+                    del dataloader_params["batch_size"]
+                    del dataloader_params["drop_last"]
+                else:
+                    dataloader_params["sampler"] = sampler
+
+            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
+            if is_training:
+                dataloader_params["worker_init_fn"] = partial(
+                    seed_worker,
+                    num_workers=self.args.dataloader_num_workers,
+                    rank=self.args.process_index,
+                )
+        if self.args.sample_packing and (
+            (is_training and not self.args.pretraining)
+            or (not is_training and self.args.eval_sample_packing is not False)
+        ):
+            self.accelerator.even_batches = False
+
+        if dataset.column_names and "length" in dataset.column_names:
+            dataset = dataset.remove_columns(["length"])
+
+        if (
+            dataset.column_names
+            and "position_ids" in dataset.column_names
+            and "attention_mask" in dataset.column_names
+            and self.args.sample_packing
+            and self.args.sample_packing_drop_attention_mask
+        ):
+            dataset = dataset.remove_columns(["attention_mask"])
+
+        dataloader = DataLoader(dataset, **dataloader_params)
+
+        # Accelerator.free_memory() will destroy the references, so
+        # we need to store the non-prepared version for eval dataloaders.
+        # fmt: off
+        if dataloader_key is not None and self.args.dataloader_persistent_workers:
+            if hasattr(self, "_eval_dataloaders"):
+                self._eval_dataloaders[dataloader_key] = dataloader  # type: ignore
+            else:
+                self._eval_dataloaders = {dataloader_key: dataloader}
+        # fmt: on
+
+        return self.accelerator.prepare(dataloader)
+
+    def _get_bench_sampler(
+        self, bench_dataset: Dataset
+    ) -> torch.utils.data.Sampler | None:
+        if self.args.world_size <= 1:
+            return SequentialSampler(bench_dataset)
+        return None
+
+    def get_bench_dataloader(
+        self,
+        bench_dataset: Dataset,
+    ) -> DataLoader:
+        dataloader_params = {
+            "batch_size": self.args.eval_batch_size,
+            "collate_fn": self.bench_data_collator,
+            "num_workers": self.args.dataloader_num_workers,
+            "pin_memory": self.args.dataloader_pin_memory,
+        }
+        if self.args.dataloader_prefetch_factor:
+            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
+
+        if not isinstance(bench_dataset, torch.utils.data.IterableDataset):
+            dataloader_params["sampler"] = self._get_bench_sampler(bench_dataset)
+            dataloader_params["drop_last"] = self.args.dataloader_drop_last
+
+        return DataLoader(bench_dataset, **dataloader_params)
+        # return self.accelerator.prepare(DataLoader(bench_dataset, **dataloader_params))
+
+    @override
+    def compute_loss(
+        self, model, inputs, return_outputs=False, num_items_in_batch=None
+    ):
+        # use one's weighted cross entropy loss calc
+        # if self.args.sample_packing:
+        #     labels = inputs.pop("labels")
+        #     outputs = model(**inputs)
+        #     loss = trainer_weighted_loss(outputs, labels, shift_labels=True)
+        #     return (loss, outputs) if return_outputs else loss
+
+        # track number of tokens for tokens per second calculation
+        if self.args.include_tkps and model.training:
+            inputs_key = "labels" if "labels" in inputs else "input_ids"
+            trainable_tokens = (inputs[inputs_key] != -100).sum()
+            total_tokens = inputs[inputs_key].numel()
+            total_tokens = torch.tensor(total_tokens, device=inputs[inputs_key].device)
+
+            if is_distributed():
+                torch.distributed.all_reduce(
+                    trainable_tokens, op=torch.distributed.ReduceOp.SUM
+                )
+                torch.distributed.all_reduce(
+                    total_tokens, op=torch.distributed.ReduceOp.SUM
+                )
+
+            if not hasattr(self.state, "tokens"):
+                self.state.tokens = {
+                    "trainable": torch.zeros(1),
+                    "total": torch.zeros(1),
+                }
+
+            # trainable tokens for throughput and total token slots for summaries
+            self.state.tokens["trainable"] = (
+                self.state.tokens["trainable"] + trainable_tokens.detach().cpu()
+            )
+            self.state.tokens["total"] = self.state.tokens["total"] + total_tokens.cpu()
+
+        # Gemma4 requires mm_token_type_ids during training (even for text-only).
+        # Inject zeros (= text token type) when not provided by the data collator.
+        # Use unwrap_model to handle DDP/FSDP wrappers that don't proxy .config.
+        _unwrapped = self.accelerator.unwrap_model(model)
+        _model_type = getattr(getattr(_unwrapped, "config", None), "model_type", None)
+        if (
+            "mm_token_type_ids" not in inputs
+            and "input_ids" in inputs
+            and _model_type in ("gemma4", "gemma4_unified")
+        ):
+            inputs["mm_token_type_ids"] = torch.zeros_like(inputs["input_ids"])
+
+        # Gemma4 (and Gemma3): transformers' masking_utils detects packed sequences
+        # from position_ids, but only when attention_mask is None.  When sample
+        # packing is active the collator provides an all-ones attention_mask that
+        # prevents this detection — remove it so the model builds the correct
+        # per-sequence causal masks.
+        if (
+            self.args.sample_packing
+            and _model_type in ("gemma4", "gemma3", "gemma4_unified")
+            and "attention_mask" in inputs
+            and "position_ids" in inputs
+        ):
+            del inputs["attention_mask"]
+
+        if self.args.orpo_alpha:
+            return self.orpo_compute_loss(
+                model,
+                inputs,
+                return_outputs=return_outputs,
+                num_items_in_batch=num_items_in_batch,
+            )
+
+        return super().compute_loss(
+            model,
+            inputs,
+            return_outputs=return_outputs,
+            num_items_in_batch=num_items_in_batch,
+        )
+
+    @override
+    def _prepare_context_parallel_inputs(self, model, inputs):
+        """Disable HF Trainer's CP splitting when Axolotl's ring_attn handles it."""
+        from axolotl.monkeypatch.models.mamba_utils import is_cp_active
+
+        if is_cp_active():
+            return contextlib.nullcontext, inputs
+        return super()._prepare_context_parallel_inputs(model, inputs)
+
+    @override
+    def evaluate(self, *args, **kwargs):
+        LOG.info("Running evaluation step...")
+        return super().evaluate(*args, **kwargs)
+
+    @override
+    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+        # Gemma4 requires mm_token_type_ids even during evaluation.
+        _unwrapped = self.accelerator.unwrap_model(model)
+        _model_type = getattr(getattr(_unwrapped, "config", None), "model_type", None)
+        if (
+            "mm_token_type_ids" not in inputs
+            and "input_ids" in inputs
+            and _model_type in ("gemma4", "gemma4_unified")
+        ):
+            inputs["mm_token_type_ids"] = torch.zeros_like(inputs["input_ids"])
+        return super().prediction_step(
+            model, inputs, prediction_loss_only, ignore_keys=ignore_keys
+        )
+
+    @staticmethod
+    def orpo_concatenate_inputs(inputs, label_pad_token=-100, pad_token=0, device=None):
+        concatenated_batch = {}
+
+        max_length = max(inputs["input_ids"].shape[1], inputs["rejected_ids"].shape[1])
+        # Concatenate positive and negative inputs
+        concatenated_batch["input_ids"] = pad_to_length(
+            inputs["input_ids"], max_length, pad_token
+        )
+        concatenated_batch["rejected_ids"] = pad_to_length(
+            inputs["rejected_ids"], max_length, pad_token
+        )
+        concatenated_batch["labels"] = pad_to_length(
+            inputs["labels"], max_length, label_pad_token
+        )
+        concatenated_batch["rejected_labels"] = pad_to_length(
+            inputs["rejected_labels"], max_length, label_pad_token
+        )
+        concatenated_batch["attention_mask"] = pad_to_length(
+            inputs["attention_mask"], max_length, 0
+        )
+        concatenated_batch["rejected_attention_mask"] = pad_to_length(
+            inputs["rejected_attention_mask"], max_length, 0
+        )
+        concatenated_batch["prompt_attention_mask"] = pad_to_length(
+            inputs["prompt_attention_mask"], max_length, 0
+        ).to(device=device)
+
+        input_ids = torch.cat(
+            [concatenated_batch["input_ids"], concatenated_batch["rejected_ids"]],
+            dim=0,
+        ).to(device=device)
+        attention_mask = torch.cat(
+            [
+                concatenated_batch["attention_mask"],
+                concatenated_batch["rejected_attention_mask"],
+            ],
+            dim=0,
+        ).to(device=device)
+        labels = torch.cat(
+            [concatenated_batch["labels"], concatenated_batch["rejected_labels"]], dim=0
+        ).to(device=device)
+
+        return {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": attention_mask,
+            "prompt_attention_mask": concatenated_batch["prompt_attention_mask"],
+        }
+
+    def orpo_compute_custom_loss(self, logits, labels):
+        logits = logits.contiguous()
+        loss = 0.0
+
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            # Flatten the tokens
+            loss = self.loss_fct(shift_logits.transpose(2, 1), shift_labels).mean(
+                dim=-1
+            )
+
+        return loss
+
+    def orpo_compute_logps(
+        self, prompt_attention_mask, chosen_inputs, chosen_attention_mask, logits
+    ):
+        # Get the shape of chosen_attention_mask[:, :-1]
+        chosen_shape = chosen_attention_mask[:, :-1].shape
+
+        # Calculate the padding size
+        pad_length = chosen_shape[1] - (prompt_attention_mask.shape[1] - 1)
+
+        # Pad prompt_attention_mask with zeros to match the desired shape
+        prompt_attention_mask_padded = torch.nn.functional.pad(
+            prompt_attention_mask[:, 1:], (0, pad_length), mode="constant", value=0
+        )
+
+        # Perform the subtraction operation
+        mask = chosen_attention_mask[:, :-1] > prompt_attention_mask_padded
+
+        per_token_logps = torch.gather(
+            logits[:, :-1, :].log_softmax(-1),
+            dim=2,
+            index=(mask * chosen_inputs[:, 1:]).unsqueeze(2),
+        ).squeeze(2)
+        return torch.mul(per_token_logps, mask).sum(dim=1) / mask.sum(dim=1)
+
+    def orpo_compute_loss(
+        self,
+        model,
+        inputs,
+        return_outputs=False,
+        num_items_in_batch=None,
+    ):
+        concat_inputs = AxolotlTrainer.orpo_concatenate_inputs(
+            inputs,
+            label_pad_token=-100,
+            pad_token=self.tokenizer.pad_token_id,
+            device=self.accelerator.device,
+        )
+
+        # Perform a single forward pass
+        forward_kwargs = {
+            "input_ids": concat_inputs["input_ids"],
+            "attention_mask": concat_inputs["attention_mask"],
+            "labels": concat_inputs["labels"],
+        }
+        # Gemma4 requires mm_token_type_ids during training (even for text-only).
+        # Unwrap to read .config (DDP/DeepSpeed wrappers don't proxy it).
+        _orpo_model_type = getattr(
+            getattr(self.accelerator.unwrap_model(model), "config", None),
+            "model_type",
+            None,
+        )
+        if (
+            _orpo_model_type in ("gemma4", "gemma4_unified")
+            and "mm_token_type_ids" not in concat_inputs
+        ):
+            forward_kwargs["mm_token_type_ids"] = torch.zeros_like(
+                concat_inputs["input_ids"]
+            )
+        elif "mm_token_type_ids" in concat_inputs:
+            forward_kwargs["mm_token_type_ids"] = concat_inputs["mm_token_type_ids"]
+
+        outputs = model(
+            **forward_kwargs,
+            output_hidden_states=True,
+        )
+
+        # Split the outputs for positive and negative examples
+        outputs_pos, outputs_neg = outputs.logits.chunk(2)
+
+        # Calculate NLL loss
+        pos_loss = self.orpo_compute_custom_loss(
+            logits=outputs_pos, labels=concat_inputs["input_ids"].chunk(2)[0]
+        )
+
+        # Calculate Log Probability
+        pos_prob = self.orpo_compute_logps(
+            prompt_attention_mask=concat_inputs["prompt_attention_mask"],
+            chosen_inputs=concat_inputs["input_ids"].chunk(2)[0],
+            chosen_attention_mask=concat_inputs["attention_mask"].chunk(2)[0],
+            logits=outputs_pos,
+        )
+        neg_prob = self.orpo_compute_logps(
+            prompt_attention_mask=concat_inputs["prompt_attention_mask"],
+            chosen_inputs=concat_inputs["input_ids"].chunk(2)[1],
+            chosen_attention_mask=concat_inputs["attention_mask"].chunk(2)[1],
+            logits=outputs_neg,
+        )
+
+        # Calculate log odds
+        log_odds = (pos_prob - neg_prob) - (
+            torch.log(1 - torch.exp(pos_prob)) - torch.log(1 - torch.exp(neg_prob))
+        )
+        sig_ratio = torch.nn.functional.sigmoid(log_odds)
+        ratio = torch.log(sig_ratio)
+
+        # Calculate the Final Loss
+        loss = torch.mean(pos_loss - self.args.orpo_alpha * ratio).to(
+            dtype=torch.bfloat16
+        )
+
+        metrics = {}
+        metrics["chosen_geometric_mean"] = torch.mean(pos_prob).cpu().item()
+        metrics["rejected_geometric_mean"] = torch.mean(neg_prob).cpu().item()
+        metrics["log_odds_ratio"] = torch.mean(ratio).cpu().item()
+        metrics["log_odds"] = torch.mean(log_odds).cpu().item()
+        self.store_metrics(metrics, train_eval="train")
+
+        return (loss, outputs_pos) if return_outputs else loss
+
+    @wraps(Trainer.push_to_hub)
+    def push_to_hub(self, *args, **kwargs) -> str:
+        """
+        Overwrite the `push_to_hub` method in order to force-add the tags when pushing the
+        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.
+        """
+        kwargs = sanitize_kwargs_for_ds_tagging(
+            dataset_tags=self.dataset_tags, kwargs=kwargs
+        )
+        kwargs = sanitize_kwargs_for_tagging(tag_names=self.tag_names, kwargs=kwargs)
+
+        return super().push_to_hub(*args, **kwargs)
+
+    @wraps(Trainer.create_accelerator_and_postprocess)
+    def create_accelerator_and_postprocess(self):
+        # cleanup the PartialState states so Accelerate automatically configures everything from the env vars
+        accelerator_config = self.args.accelerator_config.to_dict()
+        use_configured_state = accelerator_config.get("use_configured_state", False)
+        if not use_configured_state:
+            AcceleratorState._reset_state(reset_partial_state=True)
+
+        super().create_accelerator_and_postprocess()
+
+    def additional_accelerator_args(
+        self, fp8: bool = False, enable_fsdp_float8_all_gather: bool = False, **kwargs
+    ) -> dict[str, Any]:
+        ret_kwargs = {}
+        if fp8:
+            from accelerate.utils import AORecipeKwargs
+            from torchao.float8 import Float8LinearConfig
+
+            # By default, Float8LinearConfig is instantiated using the "tensorwise"
+            # scaling strategy. See more details here:
+            # https://github.com/pytorch/ao/tree/main/torchao/float8.
+            config = Float8LinearConfig(
+                enable_fsdp_float8_all_gather=enable_fsdp_float8_all_gather,
+                force_recompute_fp8_weight_in_bwd=enable_fsdp_float8_all_gather is True,
+            )
+
+            ret_kwargs["mixed_precision"] = "fp8"
+            ret_kwargs["kwargs_handlers"] = [AORecipeKwargs(config=config)]  # type: ignore
+            os.environ["ACCELERATE_MIXED_PRECISION"] = "fp8"
+
+        return ret_kwargs
+
+    def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
+        """
+        Log `logs` on the various objects watching training, including stored metrics.
+
+        Args:
+            logs: The values to log.
+            start_time: The start of training.
+        """
+        # logs either has 'loss' or 'eval_loss'
+        train_eval = "train" if "loss" in logs else "eval"
+        metric_ndigits = int(os.getenv("AXOLOTL_METRIC_NDIGITS", "5"))
+
+        for key, metric_data in self._stored_metrics[train_eval].items():
+            values = torch.tensor(metric_data["values"])  # type: ignore[arg-type]
+            reduction_type = metric_data["reduction"]
+
+            fn = REDUCTION_FNS.get(reduction_type)
+            if fn is None:
+                raise NotImplementedError(
+                    "Metric reduction must be one of [mean, min, max, sum]"
+                )
+            logs[key] = round(fn(values).item(), metric_ndigits)
+
+        if "loss" in logs:
+            try:
+                logs["ppl"] = round(math.exp(logs["loss"]), metric_ndigits)
+            except OverflowError:
+                logs["ppl"] = float("inf")
+        if "eval_loss" in logs:
+            try:
+                logs["eval_ppl"] = round(math.exp(logs["eval_loss"]), metric_ndigits)
+            except OverflowError:
+                logs["eval_ppl"] = float("inf")
+
+        if is_main_process():
+            # Add memory usage
+            try:
+                active, allocated, reserved = get_gpu_memory_usage()
+                logs["memory/max_active (GiB)"] = round(active, 2)
+                logs["memory/max_allocated (GiB)"] = round(allocated, 2)
+                logs["memory/device_reserved (GiB)"] = round(reserved, 2)
+            except (ValueError, TypeError, FileNotFoundError):
+                pass
+
+        if (
+            self.args.include_tkps
+            and train_eval == "train"
+            and hasattr(self.state, "tokens")
+        ):
+            if "trainable" in self.state.tokens:
+                now = time.perf_counter()
+                curr_trainable = float(self.state.tokens["trainable"].item())
+                elapsed = (
+                    now - self._tkps_prev_time
+                    if self._tkps_prev_time is not None
+                    else 0.0
+                )
+                rate = trainable_tokens_per_sec_per_gpu(
+                    self._tkps_prev_trainable,
+                    curr_trainable,
+                    get_world_size(),
+                    elapsed,
+                )
+                if rate is not None:
+                    logs["tokens/train_per_sec_per_gpu"] = round(rate, 2)
+                self._tkps_prev_trainable = curr_trainable
+                self._tkps_prev_time = now
+                logs["tokens/trainable"] = int(curr_trainable)
+            if "total" in self.state.tokens:
+                logs["tokens/total"] = int(self.state.tokens["total"].item())
+
+        del self._stored_metrics[train_eval]
+
+        return super().log(logs, start_time)
+
+    def store_metrics(
+        self,
+        metrics: dict[str, float] | dict[str, tuple[int | float, str]],
+        train_eval: Literal["train", "eval"] = "train",
+        reduction: Literal["mean", "min", "max", "sum"] = "mean",
+    ) -> None:
+        """
+        Store metrics with specified reduction type.
+
+        Args:
+            metrics: Dictionary of metric names to values, or metric names to (value,
+                reduction_type) tuples.
+            train_eval: Whether this is for training or evaluation.
+        """
+        for key, value in metrics.items():
+            if isinstance(value, tuple):
+                value, _reduction = value  # type: ignore[assignment]
+            else:
+                value, _reduction = value, reduction
+
+            self._stored_metrics[train_eval][key]["values"].append(value)
+            self._stored_metrics[train_eval][key]["reduction"] = _reduction
+
+    def _is_fsdp2_checkpoint_save_enabled(self) -> bool:
+        cfg = getattr(self, "axolotl_cfg", None)
+        cfg_fsdp2 = bool(
+            cfg
+            and str(getattr(cfg, "fsdp_version", "")) == "2"
+            and (getattr(cfg, "fsdp_config", None) or getattr(cfg, "fsdp", None))
+        )
+        return bool(getattr(self, "is_fsdp_enabled", False) or cfg_fsdp2)
+
+    @staticmethod
+    def _is_fsdp2_quantized_param(param) -> bool:
+        quant_names = {"NVFP4Tensor", "Float8Tensor", "MXTensor"}
+        tensor = getattr(param, "_local_tensor", param)
+        return (
+            type(tensor).__name__ in quant_names
+            or type(getattr(tensor, "data", None)).__name__ in quant_names
+        )
+
+    def _save_fsdp2_quantized_lora_adapter(self, model, output_dir) -> bool:
+        """Save just the LoRA adapter (gathered via DTensor.full_tensor) when the run is FSDP2 + a
+        quantized (NVFP4/Float8) frozen base — the case where the DCP sharded save raises
+        "Failed to validate global plan". Returns True if it handled the save, else False (caller
+        falls back to the normal checkpoint path). No-op for non-PEFT / non-FSDP2 / non-quantized runs.
+        """
+        if not self._is_fsdp2_checkpoint_save_enabled():
+            return False
+        try:
+            from peft import PeftModel
+
+            unwrapped = self.accelerator.unwrap_model(model)
+            if not isinstance(unwrapped, PeftModel):
+                return False
+            # quantized base? (torchao tensor-subclass DTensors — what breaks DCP). Handle DTensor
+            # by inspecting the local tensor.
+            has_quant = any(
+                self._is_fsdp2_quantized_param(p) for p in unwrapped.parameters()
+            )
+            if not has_quant:
+                return False
+            from axolotl.integrations.expert_parallel.shard import (
+                save_fsdp2_lora_adapter,
+            )
+
+            cfg = getattr(self, "axolotl_cfg", None)
+            if cfg and (getattr(cfg, "expert_parallel_size", 1) or 1) > 1:
+                from axolotl.integrations.expert_parallel.plugin import (
+                    ExpertParallelPlugin,
+                )
+                from axolotl.integrations.expert_parallel.shard import (
+                    save_ep_lora_adapter,
+                )
+
+                ep_group = ExpertParallelPlugin._resolve_ep_group(cfg)
+                if save_ep_lora_adapter(unwrapped, output_dir, ep_group):
+                    return True
+
+            return bool(save_fsdp2_lora_adapter(unwrapped, output_dir))
+        except Exception as exc:  # pylint: disable=broad-except
+            LOG.warning(
+                "FSDP2 quantized-LoRA adapter save failed (%s); falling back to default save.",
+                exc,
+            )
+            return False
+
+    def _save_checkpoint(self, model, trial, **kwargs):
+        # make sure the checkpoint dir exists, since trainer is flakey
+        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
+        run_dir = self._get_output_dir(trial=trial)
+        output_dir = os.path.join(run_dir, checkpoint_folder)
+        os.makedirs(output_dir, exist_ok=True)
+
+        # FSDP2 + a quantized (NVFP4/Float8) frozen base breaks the DCP sharded save
+        # ("Failed to validate global plan" — the torchao tensor-subclass DTensors are unvalidatable).
+        # For a PEFT (LoRA) run we only need the adapter, so gather it directly and skip the DCP save.
+        if self._save_fsdp2_quantized_lora_adapter(model, output_dir):
+            # The adapter is written above in place of the DCP model state. Still persist the
+            # optimizer/scheduler/scaler/RNG and trainer state so the checkpoint stays RESUMABLE —
+            # only the trainable adapter carries optimizer state (the quantized base is frozen), so
+            # these saves don't touch the unvalidatable NVFP4 DTensors. Defensive: if any of them
+            # fails under FSDP2, keep the (already-written) adapter rather than aborting the save.
+            try:
+                if not self.args.save_only_model:
+                    self._save_optimizer_and_scheduler(output_dir)
+                    self._save_scaler(output_dir)
+                    self._save_rng_state(output_dir)
+                if self.args.should_save:
+                    # "trainer_state.json" is HF's canonical resume file (global_step, epoch, ...).
+                    self.state.save_to_json(
+                        os.path.join(output_dir, "trainer_state.json")
+                    )
+            except Exception as exc:  # pylint: disable=broad-except
+                LOG.warning(
+                    "Could not persist optimizer/RNG/trainer state for %s (%s); the checkpoint is "
+                    "adapter-only and not resumable.",
+                    output_dir,
+                    exc,
+                )
+            gc.collect()
+            return None
+
+        # Save total_tokens state if tracking is enabled
+        if self.args.include_tkps and hasattr(self.state, "tokens"):
+            tokens_state = {
+                "total": int(torch.as_tensor(self.state.tokens.get("total", 0)).item()),
+                "trainable": int(
+                    torch.as_tensor(self.state.tokens.get("trainable", 0)).item()
+                ),
+            }
+            tokens_state_path = os.path.join(output_dir, TOKENS_STATE_FILE)
+            with open(tokens_state_path, "w", encoding="utf-8") as f:
+                json.dump(tokens_state, f)
+
+        result = super()._save_checkpoint(model, trial, **kwargs)
+
+        # Reclaim VRAM held by the FSDP full-state-dict gather.
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        return result
+
+    # TODO(wing): remove once https://github.com/huggingface/transformers/pull/39866/files is merged
+    def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        # If we are executing this function, we are the process zero, so we don't check for that.
+        output_dir = output_dir if output_dir is not None else self.args.output_dir
+        os.makedirs(output_dir, exist_ok=True)
+        LOG.info(f"Saving model checkpoint to {output_dir}")
+
+        # fix for Context Parallel save: CP eval invalidates tensor storage
+        # pointers, so clone to CPU to get fresh valid storage for safetensors
+        if (
+            state_dict is not None
+            and self.axolotl_cfg
+            and self.axolotl_cfg.context_parallel_size
+            and self.axolotl_cfg.context_parallel_size > 1
+        ):
+            state_dict = {
+                k: v.detach().cpu() if isinstance(v, torch.Tensor) else v
+                for k, v in state_dict.items()
+            }
+
+        supported_classes = (
+            (PreTrainedModel,)
+            if not is_peft_available()
+            else (PreTrainedModel, PeftModel)
+        )
+        # Save a trained model and configuration using `save_pretrained()`.
+        # They can then be reloaded using `from_pretrained()`
+        if not isinstance(self.model, supported_classes):
+            if state_dict is None:
+                state_dict = self.model.state_dict()
+
+            if isinstance(
+                self.accelerator.unwrap_model(self.model, keep_torch_compile=False),
+                supported_classes,
+            ):
+                self.accelerator.unwrap_model(
+                    self.model, keep_torch_compile=False
+                ).save_pretrained(
+                    output_dir,
+                    state_dict=state_dict,
+                    is_main_process=self.accelerator.is_main_process,
+                )
+            else:
+                LOG.info(
+                    "Trainer.model is not a `PreTrainedModel`, only saving its state dict."
+                )
+                safetensors.torch.save_file(
+                    state_dict,
+                    os.path.join(output_dir, SAFE_WEIGHTS_NAME),
+                    metadata={"format": "pt"},
+                )
+        else:
+            self.model.save_pretrained(
+                output_dir,
+                state_dict=state_dict,
+                is_main_process=self.accelerator.is_main_process,
+            )
+
+        if self.processing_class is not None:
+            self.processing_class.save_pretrained(output_dir)
+        elif (
+            self.data_collator is not None
+            and hasattr(self.data_collator, "tokenizer")
+            and self.data_collator.tokenizer is not None
+        ):
+            LOG.info(
+                "Saving Trainer.data_collator.tokenizer by default as Trainer.processing_class is `None`"
+            )
+            self.data_collator.tokenizer.save_pretrained(output_dir)
+
+        # Good practice: save your training arguments together with the trained model
+        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
diff --git a/src/axolotl/core/trainers/constants.py b/src/axolotl/core/trainers/constants.py
new file mode 100644
index 0000000000..ccd7d39b91
--- /dev/null
+++ b/src/axolotl/core/trainers/constants.py
@@ -0,0 +1 @@
+TOKENS_STATE_FILE = "tokens_state.json"
diff --git a/src/axolotl/core/trainers/dpo/__init__.py b/src/axolotl/core/trainers/dpo/__init__.py
new file mode 100644
index 0000000000..a73ec3149e
--- /dev/null
+++ b/src/axolotl/core/trainers/dpo/__init__.py
@@ -0,0 +1,47 @@
+"""DPO Specific Strategy for training"""
+
+from axolotl.core.trainers.dpo.trainer import AxolotlDPOTrainer
+from axolotl.utils.schemas.enums import RLType
+
+
+class DPOStrategy:
+    """Strategy for DPO training"""
+
+    @classmethod
+    def get_trainer_class(cls):
+        return AxolotlDPOTrainer
+
+    @classmethod
+    def get_training_args_class(cls):
+        from axolotl.core.trainers.dpo.args import AxolotlDPOConfig
+
+        return AxolotlDPOConfig
+
+    @classmethod
+    def set_training_args_kwargs(cls, cfg):
+        training_args_kwargs = {}
+        if cfg.rl is RLType.DPO:
+            if cfg.dpo_loss_type is not None:
+                training_args_kwargs["loss_type"] = cfg.dpo_loss_type
+
+            if cfg.dpo_loss_weights is not None:
+                training_args_kwargs["loss_weights"] = cfg.dpo_loss_weights
+
+        if cfg.rl is RLType.IPO:
+            training_args_kwargs["loss_type"] = ["ipo"]
+
+        # Label smoothing is not compatible with IPO
+        if cfg.rl is RLType.DPO and cfg.dpo_label_smoothing:
+            training_args_kwargs["label_smoothing"] = cfg.dpo_label_smoothing
+        training_args_kwargs["max_length"] = cfg.sequence_len
+        if cfg.dpo_use_weighting is not None:
+            training_args_kwargs["use_weighting"] = cfg.dpo_use_weighting
+        if cfg.dpo_padding_free is not None:
+            training_args_kwargs["padding_free"] = cfg.dpo_padding_free
+        if cfg.dpo_use_liger_kernel is not None:
+            training_args_kwargs["use_liger_kernel"] = cfg.dpo_use_liger_kernel
+        if cfg.precompute_ref_log_probs is not None:
+            training_args_kwargs["precompute_ref_log_probs"] = (
+                cfg.precompute_ref_log_probs
+            )
+        return training_args_kwargs
diff --git a/src/axolotl/core/trainers/dpo/args.py b/src/axolotl/core/trainers/dpo/args.py
new file mode 100644
index 0000000000..de1758ed09
--- /dev/null
+++ b/src/axolotl/core/trainers/dpo/args.py
@@ -0,0 +1,16 @@
+"""
+Axolotl specific DPO args
+"""
+
+from dataclasses import dataclass
+
+from trl import DPOConfig
+
+from axolotl.core.training_args import AxolotlTrainingMixins
+
+
+@dataclass
+class AxolotlDPOConfig(AxolotlTrainingMixins, DPOConfig):
+    """
+    DPO config for DPO training
+    """
diff --git a/src/axolotl/core/trainers/dpo/trainer.py b/src/axolotl/core/trainers/dpo/trainer.py
new file mode 100644
index 0000000000..69e18080fd
--- /dev/null
+++ b/src/axolotl/core/trainers/dpo/trainer.py
@@ -0,0 +1,102 @@
+"""DPO trainer for axolotl"""
+
+import gc
+import inspect
+from functools import wraps
+from typing import Any, Dict, Union
+
+import torch
+from torch import nn
+from transformers import PreTrainedTokenizerBase, ProcessorMixin
+from trl import DPOTrainer
+
+from axolotl.core.trainers.mixins import (
+    DistributedParallelMixin,
+    RngLoaderMixin,
+    SchedulerMixin,
+)
+from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin
+from axolotl.core.trainers.utils import (
+    sanitize_kwargs_for_ds_tagging,
+    sanitize_kwargs_for_tagging,
+)
+from axolotl.utils.data.utils import remove_double_bos_token
+
+
+class AxolotlDPOTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DPOTrainer,
+    DistributedParallelMixin,
+):
+    """Extend the base DPOTrainer for axolotl helpers."""
+
+    tag_names = ["axolotl", "dpo"]
+
+    def __init__(self, *args, dataset_tags=None, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.dataset_tags = dataset_tags
+        self.optimizer = None
+        self.model_accepts_loss_kwargs = False
+
+    @wraps(DPOTrainer.push_to_hub)
+    def push_to_hub(self, *args, **kwargs) -> str:
+        """
+        Overwrite the `push_to_hub` method in order to force-add the tags when pushing
+        the model on the Hub. Please refer to `~transformers.Trainer.push_to_hub`
+        for more details.
+        """
+        kwargs = sanitize_kwargs_for_ds_tagging(
+            dataset_tags=self.dataset_tags, kwargs=kwargs
+        )
+        kwargs = sanitize_kwargs_for_tagging(tag_names=self.tag_names, kwargs=kwargs)
+
+        return super().push_to_hub(*args, **kwargs)
+
+    def _tokenize(
+        self,
+        processing_class: PreTrainedTokenizerBase | ProcessorMixin,
+        input: str | list,
+        is_vlm: bool | None = None,
+        **kwargs,
+    ) -> dict[str, list]:
+        """
+        Override TRL's tokenization in DPO trainer to fix double bos_token bug (eg. llama).
+        """
+        parent_tokenize = super()._tokenize
+        try:
+            accepts_is_vlm = "is_vlm" in inspect.signature(parent_tokenize).parameters
+        except (TypeError, ValueError):
+            accepts_is_vlm = False
+        if accepts_is_vlm:
+            kwargs["is_vlm"] = is_vlm
+        result = parent_tokenize(
+            processing_class=processing_class, input=input, **kwargs
+        )
+
+        # Handle multimodal models
+        tokenizer = (
+            getattr(processing_class, "tokenizer", None)
+            if isinstance(processing_class, ProcessorMixin)
+            else processing_class
+        )
+
+        bos_token_id = getattr(tokenizer, "bos_token_id", None) if tokenizer else None
+        if bos_token_id is not None:
+            result = remove_double_bos_token(result, bos_token_id)
+
+        return result
+
+    def training_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        num_items_in_batch=None,
+    ) -> torch.Tensor:
+        loss: torch.Tensor = super().training_step(model, inputs, num_items_in_batch)
+        gc.collect()
+        torch.cuda.empty_cache()
+        return loss
diff --git a/src/axolotl/core/trainers/ebft/__init__.py b/src/axolotl/core/trainers/ebft/__init__.py
new file mode 100644
index 0000000000..92abe9f269
--- /dev/null
+++ b/src/axolotl/core/trainers/ebft/__init__.py
@@ -0,0 +1,229 @@
+"""EBFT (Energy-Based Fine-Tuning) Strategy for training
+
+Two modes:
+- structured: For QA data with prompt/completion splits. Uses GRPOTrainer + vLLM.
+- strided: For unstructured text (raw code, prose). Uses strided block-parallel generation.
+"""
+
+from typing import Any
+
+from axolotl.core.trainers.ebft.args import (
+    AxolotlAsyncEBFTConfig,
+    AxolotlEBFTConfig,
+    AxolotlStridedEBFTConfig,
+)
+from axolotl.utils.dict import DictDefault
+
+
+def _get_ebft_mode(cfg: DictDefault) -> str:
+    """Determine EBFT mode from config."""
+    if cfg.ebft and hasattr(cfg.ebft, "mode") and cfg.ebft.mode:
+        return cfg.ebft.mode
+    return "structured"
+
+
+class EBFTStrategy:
+    """Strategy for EBFT training — dispatches between structured and strided modes."""
+
+    @classmethod
+    def get_trainer_class(cls, cfg: DictDefault | None = None):
+        mode = _get_ebft_mode(cfg) if cfg else "structured"
+        if mode == "strided":
+            from axolotl.core.trainers.ebft.strided import AxolotlStridedEBFTTrainer
+
+            return AxolotlStridedEBFTTrainer
+
+        # Structured mode: async or sync
+        # use_data_producer also triggers async trainer (needed for LoRA sync
+        # without async_prefetch, since sync trainer lacks LoRA sync support)
+        use_async = (
+            cfg
+            and cfg.trl
+            and (
+                getattr(cfg.trl, "async_prefetch", False)
+                or getattr(cfg.trl, "use_data_producer", False)
+            )
+        )
+        if use_async:
+            from axolotl.core.trainers.ebft.trainer import AxolotlAsyncEBFTTrainer
+
+            return AxolotlAsyncEBFTTrainer
+        from axolotl.core.trainers.ebft.trainer import AxolotlEBFTTrainer
+
+        return AxolotlEBFTTrainer
+
+    @classmethod
+    def get_training_args_class(cls, cfg: DictDefault | None = None):
+        mode = _get_ebft_mode(cfg) if cfg else "structured"
+        if mode == "strided":
+            return AxolotlStridedEBFTConfig
+
+        # Structured mode: async or sync config
+        use_async = (
+            cfg
+            and cfg.trl
+            and (
+                getattr(cfg.trl, "async_prefetch", False)
+                or getattr(cfg.trl, "use_data_producer", False)
+            )
+        )
+        if use_async:
+            return AxolotlAsyncEBFTConfig
+        return AxolotlEBFTConfig
+
+    @classmethod
+    def is_strided(cls, cfg: DictDefault) -> bool:
+        return _get_ebft_mode(cfg) == "strided"
+
+    @classmethod
+    def set_training_args_kwargs(cls, cfg: DictDefault) -> dict[str, Any]:
+        """Map axolotl YAML config fields to training args kwargs."""
+        kwargs: dict[str, Any] = {}
+        mode = _get_ebft_mode(cfg)
+
+        # Common EBFT fields
+        ebft = cfg.ebft
+        if ebft:
+            if ebft.feature_layers is not None:
+                kwargs["ebft_feature_layers"] = ebft.feature_layers
+            if ebft.embed_method is not None:
+                kwargs["ebft_embed_method"] = ebft.embed_method
+            if ebft.use_whitening is not None:
+                kwargs["ebft_use_whitening"] = ebft.use_whitening
+            if ebft.alignment_coef is not None:
+                kwargs["ebft_alignment_coef"] = ebft.alignment_coef
+            if ebft.diversity_coef is not None:
+                kwargs["ebft_diversity_coef"] = ebft.diversity_coef
+            if ebft.ce_coef is not None:
+                kwargs["ebft_ce_coef"] = ebft.ce_coef
+            if getattr(ebft, "adaptive_max_tokens", None) is not None:
+                kwargs["ebft_adaptive_max_tokens"] = ebft.adaptive_max_tokens
+            if getattr(ebft, "gt_length_multiplier", None) is not None:
+                kwargs["ebft_gt_length_multiplier"] = ebft.gt_length_multiplier
+
+        if mode == "strided":
+            # Strided-specific fields
+            if ebft:
+                if ebft.stride is not None:
+                    kwargs["ebft_stride"] = ebft.stride
+                if ebft.context_length is not None:
+                    kwargs["ebft_context_length"] = ebft.context_length
+                if ebft.generate_max_len is not None:
+                    kwargs["ebft_generate_max_len"] = ebft.generate_max_len
+                if ebft.n_samples_per_prompt is not None:
+                    kwargs["ebft_n_samples_per_prompt"] = ebft.n_samples_per_prompt
+                if ebft.temperature is not None:
+                    kwargs["ebft_temperature"] = ebft.temperature
+                if ebft.top_p is not None:
+                    kwargs["ebft_top_p"] = ebft.top_p
+                if ebft.rl_coef is not None:
+                    kwargs["ebft_rl_coef"] = ebft.rl_coef
+                if ebft.advantage_estimator is not None:
+                    kwargs["ebft_advantage_estimator"] = ebft.advantage_estimator
+                if ebft.min_completion_prefix is not None:
+                    kwargs["ebft_min_completion_prefix"] = ebft.min_completion_prefix
+        else:
+            # Structured mode: map TRL config fields
+            trl = cfg.trl
+            if trl:
+                if trl.use_vllm:
+                    kwargs["use_vllm"] = trl.use_vllm
+                    if trl.vllm_mode:
+                        kwargs["vllm_mode"] = trl.vllm_mode
+                    if trl.vllm_mode == "colocate":
+                        kwargs["vllm_enable_sleep_mode"] = trl.vllm_enable_sleep_mode
+                        vllm_cfg = cfg.vllm
+                        if vllm_cfg:
+                            kwargs["vllm_gpu_memory_utilization"] = (
+                                vllm_cfg.gpu_memory_utilization
+                            )
+                            kwargs["vllm_tensor_parallel_size"] = (
+                                vllm_cfg.tensor_parallel_size
+                            )
+                    kwargs["vllm_server_host"] = trl.vllm_server_host or (
+                        trl.vllm.host if trl.vllm else None
+                    )
+                    kwargs["vllm_server_port"] = trl.vllm_server_port or (
+                        trl.vllm.port if trl.vllm else None
+                    )
+                    if trl.vllm_server_timeout:
+                        kwargs["vllm_server_timeout"] = trl.vllm_server_timeout
+
+                if trl.num_generations:
+                    kwargs["num_generations"] = trl.num_generations
+                if trl.max_completion_length is not None:
+                    kwargs["max_completion_length"] = trl.max_completion_length
+                if trl.temperature is not None:
+                    kwargs["temperature"] = trl.temperature
+                if trl.top_p is not None:
+                    kwargs["top_p"] = trl.top_p
+                if trl.top_k is not None:
+                    kwargs["top_k"] = trl.top_k
+                if trl.min_p is not None:
+                    kwargs["min_p"] = trl.min_p
+                if trl.num_iterations is not None:
+                    kwargs["num_iterations"] = trl.num_iterations
+                if trl.epsilon is not None:
+                    kwargs["epsilon"] = trl.epsilon
+                if trl.epsilon_high is not None:
+                    kwargs["epsilon_high"] = trl.epsilon_high
+                if trl.scale_rewards is not None:
+                    kwargs["scale_rewards"] = trl.scale_rewards
+                if trl.loss_type is not None:
+                    kwargs["loss_type"] = trl.loss_type
+                if trl.mask_truncated_completions is not None:
+                    kwargs["mask_truncated_completions"] = (
+                        trl.mask_truncated_completions
+                    )
+                if trl.log_completions is not None:
+                    kwargs["log_completions"] = trl.log_completions
+                if trl.num_completions_to_print is not None:
+                    kwargs["num_completions_to_print"] = trl.num_completions_to_print
+                if trl.sync_ref_model:
+                    kwargs["sync_ref_model"] = trl.sync_ref_model
+                if trl.repetition_penalty is not None:
+                    kwargs["repetition_penalty"] = trl.repetition_penalty
+                if trl.generation_kwargs is not None:
+                    kwargs["generation_kwargs"] = trl.generation_kwargs
+                if trl.chat_template_kwargs is not None:
+                    kwargs["chat_template_kwargs"] = trl.chat_template_kwargs
+
+                # Async prefetch fields (only pass when enabled — sync config doesn't have these)
+                if getattr(trl, "async_prefetch", False):
+                    kwargs["async_prefetch"] = trl.async_prefetch
+                    if getattr(trl, "vllm_sync_interval", None) is not None:
+                        kwargs["vllm_sync_interval"] = trl.vllm_sync_interval
+                if getattr(trl, "vllm_lora_sync", False):
+                    kwargs["vllm_lora_sync"] = trl.vllm_lora_sync
+
+        return kwargs
+
+    @classmethod
+    def set_trainer_args(cls, cfg: DictDefault) -> list[Any]:
+        return []
+
+    @classmethod
+    def set_trainer_kwargs(cls, cfg: DictDefault) -> dict[str, Any]:
+        return {}
+
+    @classmethod
+    def get_blocklist_args_kwargs(cls, cfg: DictDefault | None = None) -> list[str]:
+        mode = _get_ebft_mode(cfg) if cfg else "structured"
+        if mode == "strided":
+            return [
+                "dataset_num_proc",
+                "max_length",
+                "max_prompt_length",
+                "include_tokens_per_second",
+                "beta",
+            ]
+        return [
+            "dataset_num_proc",
+            "max_length",
+            "include_tokens_per_second",
+            "max_prompt_length",
+        ]
+
+    @classmethod
+    def get_collator(cls, *args, **kwargs):
+        return None
diff --git a/src/axolotl/core/trainers/ebft/args.py b/src/axolotl/core/trainers/ebft/args.py
new file mode 100644
index 0000000000..4a31b6b6b4
--- /dev/null
+++ b/src/axolotl/core/trainers/ebft/args.py
@@ -0,0 +1,133 @@
+"""
+EBFT-specific training arguments.
+
+Two config classes:
+- AxolotlEBFTConfig: extends GRPOConfig for structured QA data (uses vLLM generation)
+- AxolotlStridedEBFTConfig: extends TrainingArguments for unstructured text (strided generation)
+"""
+
+from dataclasses import dataclass, field
+
+from transformers import TrainingArguments
+from trl import GRPOConfig
+
+from axolotl.core.trainers.grpo.fast_async_trainer import FastAsyncGRPOConfig
+from axolotl.core.training_args import AxolotlTrainingMixins
+
+
+# -- Shared EBFT fields as a mixin --
+@dataclass
+class EBFTFieldsMixin:
+    """Common fields shared between structured and strided EBFT configs."""
+
+    ebft_feature_layers: list[float] = field(
+        default_factory=lambda: [0.25, 0.5, 0.75],
+        metadata={"help": "Fractional layer depths for feature extraction"},
+    )
+    ebft_embed_method: str = field(
+        default="last_token",
+        metadata={"help": "Pooling method: 'last_token', 'mean_pooling', or 'concat'"},
+    )
+    ebft_use_whitening: bool = field(
+        default=False,
+        metadata={"help": "Apply SVD whitening to feature embeddings"},
+    )
+    ebft_alignment_coef: float = field(
+        default=1.0,
+        metadata={"help": "Coefficient for alignment reward (cosine similarity)"},
+    )
+    ebft_diversity_coef: float = field(
+        default=1.0,
+        metadata={"help": "Coefficient for diversity penalty"},
+    )
+    ebft_ce_coef: float = field(
+        default=0.0,
+        metadata={"help": "Cross-entropy loss coefficient on ground-truth tokens"},
+    )
+    ebft_adaptive_max_tokens: bool = field(
+        default=True,
+        metadata={"help": "Set per-batch max_tokens based on ground-truth length"},
+    )
+    ebft_gt_length_multiplier: float = field(
+        default=1.5,
+        metadata={
+            "help": "Multiplier for ground-truth token count when computing adaptive max_tokens"
+        },
+    )
+
+
+# -- Structured mode: extends GRPOTrainer for QA data with vLLM --
+@dataclass
+class AxolotlEBFTConfig(EBFTFieldsMixin, AxolotlTrainingMixins, GRPOConfig):
+    """EBFT config for structured QA data — extends GRPOConfig."""
+
+    vllm_lora_sync: bool = field(
+        default=False,
+        metadata={
+            "help": "Sync LoRA adapters to vLLM via filesystem instead of NCCL weight merge."
+        },
+    )
+
+
+# -- Async structured mode: extends FastAsyncGRPOConfig --
+@dataclass
+class AxolotlAsyncEBFTConfig(
+    EBFTFieldsMixin, AxolotlTrainingMixins, FastAsyncGRPOConfig
+):
+    """EBFT config for async structured QA data — extends FastAsyncGRPOConfig.
+
+    Includes all async fields: async_prefetch, vllm_lora_sync,
+    skip_zero_advantage_batches, streaming_partial_batch, replay_buffer_size, etc.
+    """
+
+    vllm_lora_sync: bool = field(
+        default=False,
+        metadata={
+            "help": "Sync LoRA adapters to vLLM via filesystem instead of NCCL weight merge."
+        },
+    )
+
+
+# -- Strided mode: extends TrainingArguments for unstructured text --
+@dataclass
+class AxolotlStridedEBFTConfig(
+    EBFTFieldsMixin, AxolotlTrainingMixins, TrainingArguments
+):
+    """EBFT config for unstructured text with strided block-parallel generation."""
+
+    ebft_stride: int = field(
+        default=8,
+        metadata={"help": "Stride between anchor points (in tokens)"},
+    )
+    ebft_context_length: int = field(
+        default=8,
+        metadata={"help": "Context window size for each block"},
+    )
+    ebft_generate_max_len: int = field(
+        default=8,
+        metadata={"help": "Number of tokens to generate per block"},
+    )
+    ebft_n_samples_per_prompt: int = field(
+        default=4,
+        metadata={"help": "Number of independent rollouts per document"},
+    )
+    ebft_temperature: float = field(
+        default=0.6,
+        metadata={"help": "Sampling temperature for strided generation"},
+    )
+    ebft_top_p: float = field(
+        default=1.0,
+        metadata={"help": "Top-p nucleus sampling threshold"},
+    )
+    ebft_rl_coef: float = field(
+        default=1.0,
+        metadata={"help": "RL policy gradient loss coefficient"},
+    )
+    ebft_advantage_estimator: str = field(
+        default="rloo",
+        metadata={"help": "Advantage estimator: 'rloo', 'group_norm', or 'reinforce'"},
+    )
+    ebft_min_completion_prefix: int = field(
+        default=0,
+        metadata={"help": "Minimum tokens into completion before placing anchors"},
+    )
diff --git a/src/axolotl/core/trainers/ebft/kernels.py b/src/axolotl/core/trainers/ebft/kernels.py
new file mode 100644
index 0000000000..03d3940727
--- /dev/null
+++ b/src/axolotl/core/trainers/ebft/kernels.py
@@ -0,0 +1,336 @@
+"""
+Fused Triton kernels for strided EBFT.
+
+These kernels eliminate intermediate tensor materializations that dominate
+the elementwise/fill category (~40% of CUDA time in profiling).
+
+Kernels:
+  1. fused_log_softmax_gather: log_softmax + gather in one pass (no full vocab materialization)
+  2. fused_masked_reinforce_loss: -logp * advantage * mask, reduced to scalar
+  3. fused_cosine_similarity: batched cosine similarity without intermediate tensors
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+# ---------------------------------------------------------------------------
+# 1. Fused log_softmax + gather (selective log softmax)
+# ---------------------------------------------------------------------------
+# Instead of: log_softmax(logits, dim=-1)  → (B, S, V)  → gather(index=labels)
+# We compute: for each (b, s), the log_softmax value at logits[b, s, labels[b, s]]
+# This avoids materializing the full (B, S, V) log_softmax output.
+
+
+@triton.jit
+def _fused_log_softmax_gather_kernel(
+    logits_ptr,  # (B*S, V) row-major
+    labels_ptr,  # (B*S,) int64
+    output_ptr,  # (B*S,) float32
+    V: tl.constexpr,  # vocab size
+    BLOCK_V: tl.constexpr,  # tile width over vocab
+):
+    """Compute log_softmax(logits)[label] for each row without materializing full output."""
+    row = tl.program_id(0)
+
+    logits_row_ptr = logits_ptr + row * V
+    label = tl.load(labels_ptr + row)
+
+    # Pass 1: find max for numerical stability
+    max_val = -float("inf")
+    for v_start in range(0, V, BLOCK_V):
+        v_offsets = v_start + tl.arange(0, BLOCK_V)
+        mask = v_offsets < V
+        vals = tl.load(logits_row_ptr + v_offsets, mask=mask, other=-float("inf"))
+        max_val = tl.maximum(max_val, tl.max(vals, axis=0))
+
+    # Pass 2: compute sum(exp(x - max))
+    sum_exp = 0.0
+    for v_start in range(0, V, BLOCK_V):
+        v_offsets = v_start + tl.arange(0, BLOCK_V)
+        mask = v_offsets < V
+        vals = tl.load(logits_row_ptr + v_offsets, mask=mask, other=-float("inf"))
+        sum_exp += tl.sum(tl.exp(vals - max_val), axis=0)
+
+    log_sum_exp = tl.log(sum_exp)
+
+    # Gather: log_softmax[label] = logits[label] - max - log_sum_exp
+    target_logit = tl.load(logits_row_ptr + label)
+    result = target_logit - max_val - log_sum_exp
+
+    tl.store(output_ptr + row, result)
+
+
+@register_kernel_op("ebft_fused_log_softmax_gather")
+def fused_log_softmax_gather(
+    logits: torch.Tensor, labels: torch.Tensor
+) -> torch.Tensor:
+    """Compute log_softmax(logits, dim=-1).gather(-1, labels) without materializing full output.
+
+    Args:
+        logits: (B, S, V) or (B*S, V) float tensor (bf16 or fp32)
+        labels: (B, S) or (B*S,) int64 tensor of target indices
+
+    Returns:
+        (B, S) or (B*S,) float32 tensor of selected log probabilities
+    """
+    orig_shape = logits.shape[:-1]
+    V = logits.shape[-1]
+    logits_2d = logits.reshape(-1, V).contiguous()
+    labels_1d = labels.reshape(-1).contiguous()
+    n_rows = logits_2d.shape[0]
+
+    output = torch.empty(n_rows, device=logits.device, dtype=torch.float32)
+
+    # Choose BLOCK_V: must be power of 2, large enough for good occupancy
+    BLOCK_V = min(triton.next_power_of_2(V), 65536)
+
+    _fused_log_softmax_gather_kernel[(n_rows,)](
+        logits_2d,
+        labels_1d,
+        output,
+        V=V,
+        BLOCK_V=BLOCK_V,
+    )
+
+    return output.view(orig_shape)
+
+
+@fused_log_softmax_gather.register_fake
+def _(logits, labels):
+    return torch.empty(logits.shape[:-1], device=logits.device, dtype=torch.float32)
+
+
+# ---------------------------------------------------------------------------
+# 2. Fused masked REINFORCE loss reduction
+# ---------------------------------------------------------------------------
+# Instead of: (-logp * adv * mask).sum() / mask.sum()
+# We do the masked multiply-accumulate in one kernel, returning (sum, count).
+
+
+@triton.jit
+def _fused_reinforce_loss_kernel(
+    logps_ptr,  # (N,) float32 per-token log probs
+    advs_ptr,  # (N,) float32 advantages
+    mask_ptr,  # (N,) bool action mask
+    partial_sum_ptr,  # (n_blocks,) partial sums
+    partial_cnt_ptr,  # (n_blocks,) partial counts
+    N: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    block_id = tl.program_id(0)
+    offsets = block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+    valid = offsets < N
+
+    logps = tl.load(logps_ptr + offsets, mask=valid, other=0.0)
+    advs = tl.load(advs_ptr + offsets, mask=valid, other=0.0)
+    m = tl.load(mask_ptr + offsets, mask=valid, other=0).to(tl.float32)
+
+    # -logp * advantage * mask
+    loss = -logps * advs * m
+    block_sum = tl.sum(loss, axis=0)
+    block_cnt = tl.sum(m, axis=0)
+
+    tl.store(partial_sum_ptr + block_id, block_sum)
+    tl.store(partial_cnt_ptr + block_id, block_cnt)
+
+
+@register_kernel_op("ebft_fused_reinforce_loss")
+def fused_reinforce_loss(
+    per_token_logps: torch.Tensor,
+    advantages: torch.Tensor,
+    action_mask: torch.Tensor,
+) -> torch.Tensor:
+    """Compute masked REINFORCE loss: (-logp * adv * mask).sum() / mask.sum().
+
+    All inputs should be flat or will be flattened. Returns scalar loss tensor.
+    """
+    logps_flat = per_token_logps.reshape(-1).contiguous()
+    advs_flat = advantages.reshape(-1).contiguous()
+    mask_flat = action_mask.reshape(-1).contiguous()
+    N = logps_flat.shape[0]
+
+    BLOCK_N = 1024
+    n_blocks = triton.cdiv(N, BLOCK_N)
+
+    partial_sum = torch.empty(n_blocks, device=logps_flat.device, dtype=torch.float32)
+    partial_cnt = torch.empty(n_blocks, device=logps_flat.device, dtype=torch.float32)
+
+    _fused_reinforce_loss_kernel[(n_blocks,)](
+        logps_flat,
+        advs_flat,
+        mask_flat,
+        partial_sum,
+        partial_cnt,
+        N=N,
+        BLOCK_N=BLOCK_N,
+    )
+
+    total_sum = partial_sum.sum()
+    total_cnt = partial_cnt.sum().clamp(min=1)
+    return total_sum / total_cnt
+
+
+@fused_reinforce_loss.register_fake
+def _(per_token_logps, advantages, action_mask):
+    return torch.empty((), device=per_token_logps.device, dtype=torch.float32)
+
+
+# ---------------------------------------------------------------------------
+# 3. Fused cosine similarity (batched, for EBFT rewards)
+# ---------------------------------------------------------------------------
+# Instead of: F.cosine_similarity(gen, gt, dim=-1) which normalizes then dots,
+# we fuse the dot product, norm computation, and division into one kernel.
+
+
+@triton.jit
+def _fused_cosine_sim_kernel(
+    a_ptr,  # (N, D) first set of vectors
+    b_ptr,  # (N, D) second set of vectors
+    out_ptr,  # (N,) cosine similarities
+    D: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    row = tl.program_id(0)
+    a_row_ptr = a_ptr + row * D
+    b_row_ptr = b_ptr + row * D
+
+    dot = 0.0
+    norm_a = 0.0
+    norm_b = 0.0
+
+    for d_start in range(0, D, BLOCK_D):
+        d_offsets = d_start + tl.arange(0, BLOCK_D)
+        mask = d_offsets < D
+        a_vals = tl.load(a_row_ptr + d_offsets, mask=mask, other=0.0).to(tl.float32)
+        b_vals = tl.load(b_row_ptr + d_offsets, mask=mask, other=0.0).to(tl.float32)
+
+        dot += tl.sum(a_vals * b_vals, axis=0)
+        norm_a += tl.sum(a_vals * a_vals, axis=0)
+        norm_b += tl.sum(b_vals * b_vals, axis=0)
+
+    denom = tl.sqrt(norm_a) * tl.sqrt(norm_b)
+    denom = tl.where(denom > 1e-8, denom, 1e-8)
+    result = dot / denom
+
+    tl.store(out_ptr + row, result)
+
+
+@register_kernel_op("ebft_fused_cosine_similarity")
+def fused_cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Compute cosine similarity along the last dimension.
+
+    Args:
+        a, b: (..., D) tensors of the same shape
+
+    Returns:
+        (...,) tensor of cosine similarities
+    """
+    orig_shape = a.shape[:-1]
+    D = a.shape[-1]
+    a_2d = a.reshape(-1, D).contiguous()
+    b_2d = b.reshape(-1, D).contiguous()
+    N = a_2d.shape[0]
+
+    output = torch.empty(N, device=a.device, dtype=torch.float32)
+
+    BLOCK_D = min(triton.next_power_of_2(D), 4096)
+
+    _fused_cosine_sim_kernel[(N,)](
+        a_2d,
+        b_2d,
+        output,
+        D=D,
+        BLOCK_D=BLOCK_D,
+    )
+
+    return output.view(orig_shape)
+
+
+@fused_cosine_similarity.register_fake
+def _(a, b):
+    return torch.empty(a.shape[:-1], device=a.device, dtype=torch.float32)
+
+
+# ---------------------------------------------------------------------------
+# 4. Fused pairwise diversity penalty
+# ---------------------------------------------------------------------------
+# Instead of: bmm(gen, gen.T) → mask diagonal → sum / (n-1)
+# We compute the pairwise dot products and exclusion in one kernel.
+
+
+@triton.jit
+def _fused_diversity_kernel(
+    emb_ptr,  # (B, N, D) embeddings, row-major
+    out_ptr,  # (B, N) diversity penalties
+    N: tl.constexpr,  # n_samples
+    D: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    """For each (b, i), compute mean dot product to all j != i."""
+    b = tl.program_id(0)
+    i = tl.program_id(1)
+
+    # Pointer to emb[b, i, :]
+    emb_bi_ptr = emb_ptr + (b * N + i) * D
+
+    total_sim = 0.0
+    for j in range(N):
+        emb_bj_ptr = emb_ptr + (b * N + j) * D
+
+        dot = 0.0
+        for d_start in range(0, D, BLOCK_D):
+            d_offsets = d_start + tl.arange(0, BLOCK_D)
+            d_mask = d_offsets < D
+            a_vals = tl.load(emb_bi_ptr + d_offsets, mask=d_mask, other=0.0).to(
+                tl.float32
+            )
+            b_vals = tl.load(emb_bj_ptr + d_offsets, mask=d_mask, other=0.0).to(
+                tl.float32
+            )
+            dot += tl.sum(a_vals * b_vals, axis=0)
+
+        # Exclude self-similarity (j == i)
+        is_other = j != i
+        total_sim += dot * is_other
+
+    result = total_sim / (N - 1)
+    tl.store(out_ptr + b * N + i, result)
+
+
+@register_kernel_op("ebft_fused_diversity_penalty")
+def fused_diversity_penalty(embeddings: torch.Tensor) -> torch.Tensor:
+    """Compute mean pairwise dot product (excluding self) per sample.
+
+    Args:
+        embeddings: (B, N, D) tensor where N is n_samples
+
+    Returns:
+        (B, N) tensor of diversity penalties
+    """
+    B, N, D = embeddings.shape
+    embeddings = embeddings.contiguous()
+    output = torch.zeros(B, N, device=embeddings.device, dtype=torch.float32)
+    if N <= 1:
+        return output  # diversity is 0 when there's only one sample
+
+    BLOCK_D = min(triton.next_power_of_2(D), 4096)
+
+    _fused_diversity_kernel[(B, N)](
+        embeddings,
+        output,
+        N=N,
+        D=D,
+        BLOCK_D=BLOCK_D,
+    )
+
+    return output
+
+
+@fused_diversity_penalty.register_fake
+def _(embeddings):
+    return torch.empty(
+        embeddings.shape[:-1], device=embeddings.device, dtype=torch.float32
+    )
diff --git a/src/axolotl/core/trainers/ebft/rewards.py b/src/axolotl/core/trainers/ebft/rewards.py
new file mode 100644
index 0000000000..993650bbe9
--- /dev/null
+++ b/src/axolotl/core/trainers/ebft/rewards.py
@@ -0,0 +1,264 @@
+"""
+Feature-matching reward utilities for Energy-Based Fine-Tuning (EBFT).
+
+Ported from: feature-002/ebft_openrlhf/openrlhf/utils/embedding_utils.py
+Paper: "Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models"
+       (Jelassi et al., 2026) https://arxiv.org/abs/2603.12248
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@torch.no_grad()
+def extract_hidden_states(
+    model: nn.Module,
+    input_ids: torch.Tensor,
+    attention_mask: torch.Tensor,
+    layer_indices: list[int],
+    batch_size: int | None = None,
+) -> torch.Tensor:
+    """
+    Forward pass through model, extracting and concatenating hidden states
+    at specified layer indices.
+
+    Args:
+        model: The frozen feature network
+        input_ids: (B, S) token ids
+        attention_mask: (B, S) attention mask
+        layer_indices: List of layer indices to extract (e.g., [8, 16, 24] for 32-layer model)
+        batch_size: If set, process in chunks to reduce peak memory
+
+    Returns:
+        Concatenated hidden states: (B, S, num_layers * H)
+    """
+    if batch_size is None:
+        batch_size = input_ids.shape[0]
+
+    # Use the inner transformer body (skips lm_head) when available.
+    # This avoids the expensive hidden_dim × vocab_size matmul whose
+    # output (logits) is never used — only hidden_states are needed.
+    body = getattr(model, "model", None)
+    if body is not None and hasattr(body, "forward"):
+        forward_model = body
+    else:
+        forward_model = model
+
+    all_features = []
+    for i in range(0, input_ids.shape[0], batch_size):
+        chunk_ids = input_ids[i : i + batch_size]
+        chunk_mask = attention_mask[i : i + batch_size]
+
+        outputs = forward_model(
+            chunk_ids,
+            attention_mask=chunk_mask,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+
+        # hidden_states is a tuple of (num_layers + 1) tensors, each (B, S, H)
+        # index 0 is the embedding layer output
+        hidden_states = outputs.hidden_states
+        chunk_features = []
+        for idx in layer_indices:
+            chunk_features.append(hidden_states[idx])
+
+        # Concatenate across feature dimension: (B, S, num_layers * H)
+        all_features.append(torch.cat(chunk_features, dim=-1))
+
+    return torch.cat(all_features, dim=0)
+
+
+def apply_embed_method(
+    hidden_states: torch.Tensor,
+    method: str,
+    attention_mask: torch.Tensor | None = None,
+    prompt_lengths: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """
+    Pool per-token hidden states into per-sequence embeddings.
+
+    Args:
+        hidden_states: (B, S, D) concatenated hidden states
+        method: One of "last_token", "mean_pooling", "completion_mean", "concat"
+        attention_mask: (B, S) mask for mean pooling
+        prompt_lengths: (B,) number of prompt tokens per sample (for completion_mean)
+
+    Returns:
+        Sequence embeddings: (B, D) for last_token/mean_pooling/completion_mean,
+                             (B, 3*D) for concat
+    """
+    if method == "last_token":
+        if attention_mask is not None:
+            # Find last non-padding position per sample
+            last_idx = attention_mask.sum(dim=1).long() - 1  # (B,)
+            return hidden_states[torch.arange(hidden_states.shape[0]), last_idx]
+        return hidden_states[:, -1, :]
+
+    if method == "mean_pooling":
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).float()  # (B, S, 1)
+            return (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+        return hidden_states.mean(dim=1)
+
+    if method == "completion_mean":
+        # Mean pool over completion tokens only (exclude prompt)
+        if prompt_lengths is None:
+            raise ValueError("completion_mean requires prompt_lengths")
+        B, S, _ = hidden_states.shape
+        positions = torch.arange(S, device=hidden_states.device).unsqueeze(0)  # (1, S)
+        comp_mask = positions >= prompt_lengths.unsqueeze(1)  # (B, S)
+        if attention_mask is not None:
+            comp_mask = comp_mask & attention_mask.bool()
+        mask = comp_mask.unsqueeze(-1).float()  # (B, S, 1)
+        return (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+
+    if method == "concat":
+        B, S, D = hidden_states.shape
+        if attention_mask is not None:
+            valid_lens = attention_mask.sum(dim=1).long()  # (B,)
+        else:
+            valid_lens = torch.full(
+                (B,), S, device=hidden_states.device, dtype=torch.long
+            )
+        # Compute quartile positions relative to valid length per sample
+        # First valid position index for each sample (handles right-padding)
+        q1 = (valid_lens // 4).clamp(min=0, max=S - 1)
+        q2 = (valid_lens // 2).clamp(min=0, max=S - 1)
+        q3 = (3 * valid_lens // 4).clamp(min=0, max=S - 1)
+        batch_idx = torch.arange(B, device=hidden_states.device)
+        return torch.cat(
+            [
+                hidden_states[batch_idx, q1],
+                hidden_states[batch_idx, q2],
+                hidden_states[batch_idx, q3],
+            ],
+            dim=-1,
+        )
+
+    raise ValueError(f"Unknown embed_method: {method}")
+
+
+@torch.no_grad()
+def get_alignment_rewards(
+    gen_embedding: torch.Tensor,
+    gt_embedding: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Compute alignment reward as cosine similarity between generated
+    and ground-truth feature embeddings.
+
+    Args:
+        gen_embedding: (B, D) generated sequence embeddings
+        gt_embedding: (B, D) ground-truth sequence embeddings
+            If num_generations > 1, gt_embedding should be repeated
+            to match gen_embedding's batch dim.
+
+    Returns:
+        Alignment rewards: (B,) cosine similarities in [-1, 1]
+    """
+    return F.cosine_similarity(gen_embedding, gt_embedding, dim=-1)
+
+
+@torch.no_grad()
+def get_diversity_rewards(
+    gen_embedding: torch.Tensor,
+    num_generations: int,
+) -> torch.Tensor:
+    """
+    Compute diversity penalty as mean pairwise dot-product similarity
+    between samples from the same prompt (excluding self-similarity).
+
+    Args:
+        gen_embedding: (B, D) generated embeddings where B = num_prompts * num_generations
+        num_generations: Number of generations per prompt
+
+    Returns:
+        Diversity penalties: (B,) mean similarity to other samples from same prompt
+    """
+    if num_generations <= 1:
+        return torch.zeros(gen_embedding.shape[0], device=gen_embedding.device)
+
+    num_prompts = gen_embedding.shape[0] // num_generations
+
+    # Reshape to (num_prompts, num_generations, D)
+    reshaped = gen_embedding.view(num_prompts, num_generations, -1)
+
+    # Pairwise dot products within each group: (num_prompts, num_generations, num_generations)
+    sims = torch.bmm(reshaped, reshaped.transpose(1, 2))
+
+    # Zero out self-similarity (diagonal)
+    eye = torch.eye(num_generations, device=sims.device, dtype=torch.bool)
+    sims = sims.masked_fill(eye.unsqueeze(0), 0.0)
+
+    # Mean similarity to other samples: (num_prompts, num_generations)
+    diversity = sims.sum(dim=-1) / (num_generations - 1)
+
+    # Flatten back to (B,)
+    return diversity.view(-1)
+
+
+def whiten_embeddings_batched(
+    phi: torch.Tensor,
+    phi_gt: torch.Tensor,
+    whiten_tol: float = 1e-5,
+    normalize: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Whiten generated embeddings using SVD, then apply same transform to ground-truth.
+
+    Whitening decorrelates feature dimensions so no single direction dominates
+    the feature-matching loss. Uses pseudo-inverse for rank-deficient cases.
+
+    Note: Singular values scale with sqrt(B), so reward magnitudes are
+    batch-size dependent. This is acceptable because B = n_samples_per_prompt
+    which is fixed during training (typically 2-4).
+
+    Args:
+        phi: (B, D) generated embeddings (used to estimate covariance)
+        phi_gt: (B, D) ground-truth embeddings
+        whiten_tol: Tolerance for singular value cutoff
+        normalize: If True, L2-normalize after whitening
+
+    Returns:
+        Whitened (phi, phi_gt) tuple, each (B, D)
+    """
+    phi_f = phi.float()
+    phi_gt_f = phi_gt.float()
+
+    # Feature-space SVD: operate on phi_f.T (D, B) so U is (D, D)
+    try:
+        U, S, _ = torch.linalg.svd(phi_f.T.unsqueeze(0), full_matrices=False)
+    except torch._C._LinAlgError:
+        # Fallback: add small noise
+        noise = 1e-6 * phi_f.abs().mean()
+        try:
+            U, S, _ = torch.linalg.svd(
+                (phi_f.T + noise * torch.randn_like(phi_f.T)).unsqueeze(0),
+                full_matrices=False,
+            )
+        except torch._C._LinAlgError:
+            if normalize:
+                return (
+                    F.normalize(phi, p=2, dim=-1),
+                    F.normalize(phi_gt, p=2, dim=-1),
+                )
+            return phi, phi_gt
+
+    U, S = U.squeeze(0), S.squeeze(0)  # U: (D, min(D,B)), S: (min(D,B),)
+
+    # Safe inverse of singular values
+    s_max = S.max()
+    inv_s = torch.where(S > whiten_tol * s_max, 1.0 / (S + 1e-12), torch.zeros_like(S))
+
+    # W = U @ diag(inv_s) @ U^T  — feature-space whitening matrix (D, D)
+    W = (U * inv_s.unsqueeze(0)) @ U.T  # (D, D)
+    phi_w = (phi_f @ W).to(phi.dtype)  # (B, D)
+    phi_gt_w = (phi_gt_f @ W).to(phi_gt.dtype)  # (B, D)
+
+    if normalize:
+        phi_w = F.normalize(phi_w, p=2, dim=-1)
+        phi_gt_w = F.normalize(phi_gt_w, p=2, dim=-1)
+
+    return phi_w, phi_gt_w
diff --git a/src/axolotl/core/trainers/ebft/strided.py b/src/axolotl/core/trainers/ebft/strided.py
new file mode 100644
index 0000000000..5cfc5b99bd
--- /dev/null
+++ b/src/axolotl/core/trainers/ebft/strided.py
@@ -0,0 +1,1152 @@
+"""
+Strided block-parallel EBFT trainer for unstructured text data.
+
+This trainer implements the full EBFT algorithm from the paper, including
+strided block-parallel generation where multiple short rollouts are generated
+at different anchor points within a single document. This is essential for
+training on raw text data (code, prose, etc.) without prompt/completion splits.
+
+Uses torch flex_attention with a compiled block mask for efficient strided
+attention patterns. Falls back to eager attention with dense 4D masks when
+flex_attention is not available.
+
+Paper: "Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models"
+       (Jelassi et al., 2026) https://arxiv.org/abs/2603.12248
+"""
+
+import contextlib
+import copy
+
+import torch
+import torch.nn.functional as F
+from transformers import Trainer
+
+from axolotl.core.trainers.ebft.kernels import (  # noqa: F401 — available for future use
+    fused_cosine_similarity,
+    fused_diversity_penalty,
+    fused_log_softmax_gather,
+    fused_reinforce_loss,
+)
+from axolotl.core.trainers.ebft.rewards import (
+    whiten_embeddings_batched,
+)
+from axolotl.core.trainers.mixins import (
+    DistributedParallelMixin,
+    RngLoaderMixin,
+    SchedulerMixin,
+)
+from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# Check flex_attention availability
+_FLEX_ATTENTION_AVAILABLE = False
+try:
+    from torch.nn.attention.flex_attention import (
+        create_block_mask,
+    )
+
+    _FLEX_ATTENTION_AVAILABLE = True
+except ImportError:
+    pass
+
+
+def _patch_flex_attention_dtype():
+    """
+    Patch HF's flex_attention_forward to cast q/k/v to a uniform dtype.
+
+    This fixes the incompatibility between flex_attention and gradient
+    checkpointing, where recomputation can produce q/k in float32 while
+    v stays in bfloat16. flex_attention requires all three to match.
+    """
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    original_fn = ALL_ATTENTION_FUNCTIONS.get("flex_attention")
+    if original_fn is None:
+        return
+
+    def patched_flex_attention_forward(
+        module, query, key, value, attention_mask, **kwargs
+    ):
+        # Cast q/k/v to the same dtype (use value's dtype as reference,
+        # since model weights stay in the original dtype)
+        target_dtype = value.dtype
+        if query.dtype != target_dtype:
+            query = query.to(target_dtype)
+        if key.dtype != target_dtype:
+            key = key.to(target_dtype)
+        return original_fn(module, query, key, value, attention_mask, **kwargs)
+
+    ALL_ATTENTION_FUNCTIONS["flex_attention"] = patched_flex_attention_forward
+
+
+@contextlib.contextmanager
+def override_attn_implementation(model, implementation: str):
+    """Temporarily override a model's attention implementation.
+
+    Useful for forcing eager attention during generation (where sequence
+    lengths change each step, causing dynamo recompiles) while keeping
+    flex_attention for the fixed-size training forward pass.
+
+    Usage::
+
+        with override_attn_implementation(model, "eager"):
+            output = model(input_ids, attention_mask=dense_4d_mask, ...)
+    """
+    config = getattr(model, "config", None)
+    if config is None or not hasattr(config, "_attn_implementation"):
+        yield
+        return
+
+    saved = config._attn_implementation
+    config._attn_implementation = implementation
+    try:
+        yield
+    finally:
+        config._attn_implementation = saved
+
+
+# ---------------------------------------------------------------------------
+# Strided attention mask builders
+# ---------------------------------------------------------------------------
+
+
+def _strided_mask_mod(
+    b,
+    h,
+    q_idx,
+    kv_idx,
+    prompt_length,
+    context_length,
+    max_generation_length,
+    stride,
+    num_blocks,
+):
+    """
+    Mask mod function for flex_attention's create_block_mask.
+
+    Defines the strided block-parallel attention pattern:
+    - Prompt region: standard causal
+    - Generated region: each block attends to its context window + same-block predecessors
+    """
+    # --- Prompt region: standard causal ---
+    is_prompt_q = q_idx < prompt_length
+    is_prompt_kv = kv_idx < prompt_length
+    prompt_causal = is_prompt_q & is_prompt_kv & (q_idx >= kv_idx)
+
+    # --- Generated region ---
+    is_gen_q = ~is_prompt_q
+    # Which generation step and block does this query token belong to?
+    gen_offset = q_idx - prompt_length
+    gen_step = gen_offset // num_blocks
+    block_idx = gen_offset % num_blocks
+
+    # Context window end for this block.
+    # Note: if prompt_length < max_generation_length, context_end clamps to 0 for all
+    # blocks. This is safe because compute_loss guards with num_blocks <= 0 → zero loss.
+    context_end = torch.clamp(
+        block_idx * stride + context_length,
+        max=prompt_length - max_generation_length,
+    )
+
+    # Rule 1: Generated token attends to its context window in the prompt
+    in_context = is_gen_q & is_prompt_kv & (kv_idx < context_end)
+
+    # Rule 2: Self-attention
+    is_self = q_idx == kv_idx
+
+    # Rule 3: Attend to earlier tokens in the SAME block (same block_idx, earlier gen_step)
+    is_gen_kv = ~is_prompt_kv
+    kv_gen_offset = kv_idx - prompt_length
+    kv_gen_step = kv_gen_offset // num_blocks
+    kv_block_idx = kv_gen_offset % num_blocks
+    same_block_prev = (
+        is_gen_q & is_gen_kv & (kv_block_idx == block_idx) & (kv_gen_step < gen_step)
+    )
+
+    return prompt_causal | in_context | is_self | same_block_prev
+
+
+def create_strided_block_mask(
+    prompt_length: int,
+    context_length: int,
+    max_generation_length: int,
+    stride: int,
+    num_blocks: int,
+    full_sequence_length: int,
+    batch_size: int,
+    num_heads: int,
+    device: torch.device,
+):
+    """
+    Create a BlockMask for flex_attention using the strided EBFT pattern.
+
+    Returns a BlockMask that can be passed directly to model.forward()
+    when using attn_implementation="flex_attention".
+
+    Parameters that vary across training steps (context_length, num_blocks)
+    are captured as tensors so torch.compile/dynamo treats them as dynamic
+    values rather than guarding on literal int values (which causes recompiles).
+    """
+    # Wrap ALL mask params as 0-d tensors to prevent dynamo from guarding
+    # on their int values. Without this, each new anchor_offset or num_blocks
+    # triggers a recompile until the limit is hit → unfused fallback → OOM.
+    _prompt_length = torch.tensor(prompt_length, device=device)
+    _context_length = torch.tensor(context_length, device=device)
+    _max_gen_len = torch.tensor(max_generation_length, device=device)
+    _stride = torch.tensor(stride, device=device)
+    _num_blocks = torch.tensor(num_blocks, device=device)
+
+    def mask_mod(b, h, q_idx, kv_idx):
+        return _strided_mask_mod(
+            b,
+            h,
+            q_idx,
+            kv_idx,
+            prompt_length=_prompt_length,
+            context_length=_context_length,
+            max_generation_length=_max_gen_len,
+            stride=_stride,
+            num_blocks=_num_blocks,
+        )
+
+    block_mask = create_block_mask(
+        mask_mod,
+        B=batch_size,
+        H=None,  # broadcast across heads
+        Q_LEN=full_sequence_length,
+        KV_LEN=full_sequence_length,
+        device=device,
+    )
+    return block_mask
+
+
+def build_strided_position_ids(
+    full_sequence_length: int,
+    prompt_length: int,
+    context_length: int,
+    generation_step: int,
+    stride: int,
+    num_blocks: int,
+    device: torch.device,
+    batch_size: int = 1,
+):
+    """Build position IDs for strided generation (shared between flex and eager modes)."""
+    position_ids = torch.empty(
+        (batch_size, full_sequence_length), dtype=torch.long, device=device
+    )
+    position_ids[:, :prompt_length] = torch.arange(prompt_length, device=device)
+
+    block_starting_positions = (
+        torch.arange(num_blocks, device=device) * stride + context_length
+    )
+    for gen_step in range(generation_step):
+        start = prompt_length + gen_step * num_blocks
+        end = start + num_blocks
+        position_ids[:, start:end] = block_starting_positions + gen_step
+
+    return position_ids
+
+
+def build_strided_dense_mask_and_positions(
+    full_sequence_length: int,
+    prompt_length: int,
+    context_length: int,
+    generation_step: int,
+    max_generation_length: int,
+    stride: int,
+    num_blocks: int,
+    device: torch.device,
+    batch_size: int = 1,
+    dtype: torch.dtype = torch.bfloat16,
+):
+    """Build dense 4D attention mask (eager fallback) + position IDs."""
+    min_value = torch.finfo(dtype).min
+    attention_mask = torch.full(
+        (batch_size, 1, full_sequence_length, full_sequence_length),
+        min_value,
+        dtype=dtype,
+        device=device,
+    )
+
+    causal_mask = torch.tril(
+        torch.ones((prompt_length, prompt_length), dtype=torch.bool, device=device)
+    )
+    attention_mask[:, :, :prompt_length, :prompt_length].masked_fill_(
+        causal_mask.view(1, 1, prompt_length, prompt_length), 0.0
+    )
+
+    for gen_step in range(generation_step):
+        for block_idx in range(num_blocks):
+            gen_pos = prompt_length + gen_step * num_blocks + block_idx
+            context_end = min(
+                block_idx * stride + context_length,
+                prompt_length - max_generation_length,
+            )
+            attention_mask[:, 0, gen_pos, :context_end] = 0.0
+            attention_mask[:, 0, gen_pos, gen_pos] = 0.0
+            if gen_step > 0:
+                for prev_s in range(gen_step):
+                    prev_pos = prompt_length + prev_s * num_blocks + block_idx
+                    attention_mask[:, 0, gen_pos, prev_pos] = 0.0
+
+    position_ids = build_strided_position_ids(
+        full_sequence_length,
+        prompt_length,
+        context_length,
+        generation_step,
+        stride,
+        num_blocks,
+        device,
+        batch_size,
+    )
+    return attention_mask, position_ids
+
+
+# ---------------------------------------------------------------------------
+# Trainer
+# ---------------------------------------------------------------------------
+
+
+class AxolotlStridedEBFTTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DistributedParallelMixin,
+    Trainer,
+):
+    """
+    Strided block-parallel EBFT trainer for unstructured text data.
+
+    Takes full text documents (no prompt/completion split needed), generates
+    short rollouts at multiple anchor points via strided attention, and trains
+    with feature-matching rewards.
+
+    When flex_attention is available (torch >= 2.5), uses compiled block masks
+    for efficient fused attention kernels. Otherwise falls back to eager
+    attention with dense 4D masks.
+    """
+
+    _tag_names = ["ebft", "strided", "axolotl"]
+
+    def __init__(self, model, args, train_dataset, **kwargs):
+        super().__init__(model=model, args=args, train_dataset=train_dataset, **kwargs)
+
+        # EBFT config
+        self.ebft_stride = getattr(args, "ebft_stride", 8)
+        self.ebft_context_length = getattr(args, "ebft_context_length", 8)
+        self.ebft_generate_max_len = getattr(args, "ebft_generate_max_len", 8)
+        self.ebft_n_samples = getattr(args, "ebft_n_samples_per_prompt", 4)
+        self.ebft_temperature = getattr(args, "ebft_temperature", 0.6)
+        self.ebft_top_p = getattr(args, "ebft_top_p", 1.0)
+        self.ebft_alignment_coef = getattr(args, "ebft_alignment_coef", 1.0)
+        self.ebft_diversity_coef = getattr(args, "ebft_diversity_coef", 1.0)
+        self.ebft_rl_coef = getattr(args, "ebft_rl_coef", 1.0)
+        self.ebft_ce_coef = getattr(args, "ebft_ce_coef", 0.0)
+        self.ebft_use_whitening = getattr(args, "ebft_use_whitening", False)
+        self.ebft_advantage_estimator = getattr(
+            args, "ebft_advantage_estimator", "rloo"
+        )
+        self.ebft_min_completion_prefix = getattr(args, "ebft_min_completion_prefix", 0)
+
+        # Validate config combinations
+        if self.ebft_use_whitening and self.ebft_diversity_coef > 0:
+            LOG.info(
+                "ebft: whitening + diversity enabled. Per paper Variant (i) (eq 49): "
+                "alignment uses cosine similarity (normalized), diversity uses raw dot product. "
+                "Both are bounded after whitening."
+            )
+        if self.ebft_n_samples == 1 and self.ebft_diversity_coef > 0:
+            LOG.warning(
+                "ebft.n_samples_per_prompt=1 with diversity_coef > 0: diversity penalty requires "
+                "multiple samples. Setting diversity_coef to 0."
+            )
+            self.ebft_diversity_coef = 0.0
+        if self.ebft_n_samples == 1 and self.ebft_advantage_estimator == "rloo":
+            LOG.warning(
+                "ebft.n_samples_per_prompt=1 with advantage_estimator='rloo': RLOO requires "
+                "multiple samples for baseline. Falling back to 'reinforce'."
+            )
+            self.ebft_advantage_estimator = "reinforce"
+
+        # Feature network config
+        feature_layers_frac = getattr(args, "ebft_feature_layers", [0.25, 0.5, 0.75])
+        embed_method = getattr(args, "ebft_embed_method", "last_token")
+        self.ebft_embed_method = embed_method
+
+        # Attention implementation selection
+        unwrapped = self.accelerator.unwrap_model(self.model)
+        self.use_flex_attention = (
+            _FLEX_ATTENTION_AVAILABLE and torch.cuda.is_available()
+        )
+
+        if self.use_flex_attention:
+            _patch_flex_attention_dtype()
+            LOG.info("Using flex_attention for strided EBFT (compiled block masks)")
+            if hasattr(unwrapped.config, "_attn_implementation"):
+                unwrapped.config._attn_implementation = "flex_attention"
+            self._num_heads = unwrapped.config.num_attention_heads
+        else:
+            LOG.info("Using eager attention for strided EBFT (dense 4D masks)")
+            if hasattr(unwrapped.config, "_attn_implementation"):
+                unwrapped.config._attn_implementation = "eager"
+            self._num_heads = None
+
+        # Feature network setup: either share weights with actor (PEFT models)
+        # or deepcopy (full-parameter models / multi-GPU).
+        first_param = next(unwrapped.parameters())
+        original_device = first_param.device
+        actor_gpu = (
+            original_device.index
+            if (original_device.type == "cuda" and original_device.index is not None)
+            else 0
+        )
+        visible_gpus = torch.cuda.device_count()
+
+        # Check if we can share weights (PEFT model on single GPU)
+        from peft import PeftModel
+
+        self._share_feature_weights = (
+            isinstance(unwrapped, PeftModel)
+            and visible_gpus == 1
+            and original_device.type != "meta"
+        )
+
+        if self._share_feature_weights:
+            # Share weights: use actor's base model with adapters disabled for
+            # feature extraction. Saves ~2.5 GB (no deepcopy of base weights).
+            self.feature_network = None  # no separate network
+            self._feature_device = torch.device(f"cuda:{actor_gpu}")
+            self._feature_use_flex = self.use_flex_attention
+            LOG.info(
+                "Feature network shares actor weights (PEFT disable_adapter). "
+                f"Saving {sum(p.numel() for p in unwrapped.parameters()) * 2 / 1e9:.1f} GB"
+            )
+        elif visible_gpus > 1 and original_device.type != "meta":
+            # Multi-GPU: deepcopy to a separate device
+            self.feature_network = copy.deepcopy(unwrapped)
+            self.feature_network.to(dtype=torch.bfloat16)
+            self._feature_device = torch.device(
+                f"cuda:{(actor_gpu + 1) % visible_gpus}"
+            )
+            LOG.info(f"Creating frozen feature network on {self._feature_device}...")
+            self.feature_network.to(device=self._feature_device)
+            if _FLEX_ATTENTION_AVAILABLE and self._feature_device.type == "cuda":
+                if hasattr(self.feature_network.config, "_attn_implementation"):
+                    self.feature_network.config._attn_implementation = "flex_attention"
+                self._feature_use_flex = True
+                LOG.info("Feature network using flex_attention")
+            else:
+                if hasattr(self.feature_network.config, "_attn_implementation"):
+                    self.feature_network.config._attn_implementation = "eager"
+                self._feature_use_flex = False
+            for param in self.feature_network.parameters():
+                param.requires_grad = False
+            self.feature_network.eval()
+        elif original_device.type == "meta":
+            # FSDP2 with cpu_ram_efficient_loading
+            from transformers import AutoModelForCausalLM
+
+            feature_model_name = (
+                getattr(args, "model_name_or_path", None)
+                or unwrapped.config._name_or_path
+            )
+            self.feature_network = AutoModelForCausalLM.from_pretrained(
+                feature_model_name,
+                torch_dtype=torch.bfloat16,
+                attn_implementation="eager",
+            )
+            self._feature_device = torch.device(f"cuda:{actor_gpu}")
+            self.feature_network.to(device=self._feature_device)
+            self._feature_use_flex = False
+            for param in self.feature_network.parameters():
+                param.requires_grad = False
+            self.feature_network.eval()
+            LOG.warning("Feature network loaded from pretrained (meta device)")
+        else:
+            # Single-GPU, non-PEFT: deepcopy on same device
+            self.feature_network = copy.deepcopy(unwrapped)
+            self.feature_network.to(dtype=torch.bfloat16)
+            self._feature_device = torch.device(f"cuda:{actor_gpu}")
+            self.feature_network.to(device=self._feature_device)
+            if _FLEX_ATTENTION_AVAILABLE:
+                if hasattr(self.feature_network.config, "_attn_implementation"):
+                    self.feature_network.config._attn_implementation = "flex_attention"
+                self._feature_use_flex = True
+            else:
+                if hasattr(self.feature_network.config, "_attn_implementation"):
+                    self.feature_network.config._attn_implementation = "eager"
+                self._feature_use_flex = False
+            for param in self.feature_network.parameters():
+                param.requires_grad = False
+            self.feature_network.eval()
+            LOG.info(
+                f"Created frozen feature network (deepcopy) on {self._feature_device}"
+            )
+
+        num_layers = unwrapped.config.num_hidden_layers
+        self.feature_layer_indices = [
+            int(frac * num_layers) for frac in feature_layers_frac
+        ]
+        LOG.info(
+            f"Strided EBFT: layers={self.feature_layer_indices}, "
+            f"stride={self.ebft_stride}, ctx={self.ebft_context_length}, "
+            f"gen_len={self.ebft_generate_max_len}, n_samples={self.ebft_n_samples}, "
+            f"embed={embed_method}, flex_attn={self.use_flex_attention}, "
+            f"min_completion_prefix={self.ebft_min_completion_prefix}"
+        )
+
+    def _build_strided_mask(
+        self,
+        full_seq_len,
+        seq_len,
+        generation_step,
+        num_blocks,
+        batch_size,
+        device,
+        dtype,
+        anchor_offset=None,
+    ):
+        """Build strided attention mask + position IDs using flex or eager.
+
+        Args:
+            anchor_offset: Position where anchors start. For unstructured data this
+                equals context_length; for structured data it equals
+                max(prompt_length + min_completion_prefix, context_length).
+                Defaults to self.ebft_context_length if not provided.
+        """
+        if anchor_offset is None:
+            anchor_offset = self.ebft_context_length
+
+        pos_ids = build_strided_position_ids(
+            full_seq_len,
+            seq_len,
+            anchor_offset,
+            generation_step,
+            self.ebft_stride,
+            num_blocks,
+            device,
+            batch_size,
+        )
+
+        if self.use_flex_attention:
+            block_mask = create_strided_block_mask(
+                prompt_length=seq_len,
+                context_length=anchor_offset,
+                max_generation_length=self.ebft_generate_max_len,
+                stride=self.ebft_stride,
+                num_blocks=num_blocks,
+                full_sequence_length=full_seq_len,
+                batch_size=batch_size,
+                num_heads=self._num_heads,
+                device=device,
+            )
+            return block_mask, pos_ids
+
+        dense_mask, pos_ids = build_strided_dense_mask_and_positions(
+            full_sequence_length=full_seq_len,
+            prompt_length=seq_len,
+            context_length=anchor_offset,
+            generation_step=generation_step,
+            max_generation_length=self.ebft_generate_max_len,
+            stride=self.ebft_stride,
+            num_blocks=num_blocks,
+            device=device,
+            batch_size=batch_size,
+            dtype=dtype,
+        )
+        return dense_mask, pos_ids
+
+    def compute_loss(
+        self, model, inputs, return_outputs=False, num_items_in_batch=None
+    ):
+        """
+        Full strided EBFT training step.
+
+        1. Take tokenized documents from inputs
+        2. Generate n_samples short rollouts at strided anchor points
+        3. Extract features from frozen network for both generated and GT blocks
+        4. Compute alignment/diversity rewards per block
+        5. Compute RLOO advantages
+        6. Policy gradient loss on the strided forward pass
+
+        Supports both unstructured text (no prompt/completion split) and
+        structured data (prompt + completion with labels masking). For structured
+        data, anchors are placed only within the completion span.
+        """
+        outputs = None
+        device = next(model.parameters()).device
+        input_ids = inputs["input_ids"].to(device)  # (B, seq_len)
+        B, seq_len = input_ids.shape
+
+        stride = self.ebft_stride
+        ctx_len = self.ebft_context_length
+        gen_len = self.ebft_generate_max_len
+        n_samples = self.ebft_n_samples
+
+        # --- Detect structured data and compute anchor_offset ---
+        # For structured data, anchors must start within the completion span.
+        # anchor_offset replaces ctx_len as the starting position for anchors.
+        is_structured = False
+        if "prompt_length" in inputs:
+            # Explicit prompt_length from dataset transform
+            prompt_lengths = inputs["prompt_length"].to(device)  # (B,)
+            is_structured = True
+        elif "labels" in inputs:
+            # Derive prompt_length from labels: first position where labels != -100
+            labels = inputs["labels"].to(device)
+            non_masked = labels != -100
+            # prompt_length = index of first non-masked token (or seq_len if all masked)
+            has_completion = non_masked.any(dim=1)
+            prompt_lengths = torch.where(
+                has_completion,
+                non_masked.float().argmax(dim=1),
+                torch.tensor(seq_len, device=device, dtype=torch.float),
+            ).long()
+            is_structured = prompt_lengths.min().item() > 0
+
+        if is_structured:
+            # Use max prompt_length across batch for uniform anchor_offset
+            max_prompt_len = prompt_lengths.max().item()
+            anchor_offset = max(
+                max_prompt_len + self.ebft_min_completion_prefix, ctx_len
+            )
+        else:
+            anchor_offset = ctx_len
+
+        num_blocks = (seq_len - gen_len - anchor_offset) // stride + 1
+        if num_blocks <= 0:
+            LOG.warning(
+                f"Sequence too short for strided EBFT: seq_len={seq_len}, "
+                f"anchor_offset={anchor_offset}, "
+                f"need >= {gen_len + anchor_offset + stride}. Returning zero loss."
+            )
+            dummy_loss = input_ids.float().mean() * 0.0
+            return (dummy_loss, None) if return_outputs else dummy_loss
+
+        # --- Step 1: Generate strided blocks for n_samples ---
+        repeated_ids = input_ids.repeat_interleave(n_samples, dim=0)
+
+        with torch.no_grad():
+            full_sequences = self._generate_strided_blocks(
+                model,
+                repeated_ids,
+                num_blocks,
+                anchor_offset=anchor_offset,
+            )
+
+        # --- Step 2: Build strided mask for full generation ---
+        full_seq_len = full_sequences.shape[1]
+        model_dtype = next(model.parameters()).dtype
+
+        # Free generation-phase memory before training forward pass
+        torch.cuda.empty_cache()
+
+        attn_mask, pos_ids = self._build_strided_mask(
+            full_seq_len,
+            seq_len,
+            gen_len,
+            num_blocks,
+            B * n_samples,
+            device,
+            model_dtype,
+            anchor_offset=anchor_offset,
+        )
+
+        # --- Step 3: Forward pass through actor for log probs ---
+        # Memory optimization: process one sample at a time through the backbone
+        # to avoid B*N × S × H activation memory. For Llama-1B at S=3900, each
+        # sample's backbone forward takes ~8.7 GB with grad checkpointing.
+        # Processing B*N=4 at once would need ~35 GB → OOM.
+        # Instead, we accumulate per-token logprobs sample-by-sample.
+        gen_start = seq_len - 1  # shifted index where generated tokens start
+        compute_start = 0 if self.ebft_ce_coef > 0 else gen_start
+        BN = B * n_samples
+
+        unwrapped_model = self.accelerator.unwrap_model(model)
+        # Navigate through PEFT wrapper to get backbone + lm_head
+        base_model = getattr(unwrapped_model, "model", unwrapped_model)
+        if hasattr(base_model, "model") and hasattr(base_model, "lm_head"):
+            backbone = base_model.model
+            lm_head = base_model.lm_head
+        else:
+            backbone = None
+
+        per_token_logps_list = []
+        shift_labels = full_sequences[:, 1:]
+
+        if backbone is not None:
+            # Process one sample at a time: backbone → chunked lm_head → logprobs
+            # This keeps peak memory at ~1 sample's activations instead of B*N.
+            for s_idx in range(BN):
+                seq_s = full_sequences[s_idx : s_idx + 1]  # (1, full_seq_len)
+                # Handle attention mask format (BlockMask vs dense 4D)
+                if isinstance(attn_mask, torch.Tensor) and attn_mask.dim() == 4:
+                    mask_s = attn_mask[s_idx : s_idx + 1]
+                else:
+                    mask_s = attn_mask  # BlockMask broadcasts over batch
+                pos_s = pos_ids[s_idx : s_idx + 1]
+
+                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                    backbone_out = backbone(
+                        seq_s,
+                        attention_mask=mask_s,
+                        position_ids=pos_s,
+                        return_dict=True,
+                    )
+                hidden_s = backbone_out.last_hidden_state  # (1, full_seq_len, H)
+                labels_s = shift_labels[s_idx : s_idx + 1]
+
+                logps_s = torch.zeros(
+                    1,
+                    hidden_s.shape[1] - 1,
+                    device=device,
+                    dtype=torch.float32,
+                )
+
+                region_h = hidden_s[:, compute_start:-1, :]
+                region_l = labels_s[:, compute_start:]
+                chunk_size = 256
+                for i in range(0, region_h.shape[1], chunk_size):
+                    h_chunk = region_h[:, i : i + chunk_size, :]
+                    l_chunk = region_l[:, i : i + chunk_size]
+                    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                        logits_chunk = lm_head(h_chunk)
+                    chunk_lp = F.log_softmax(logits_chunk.float(), dim=-1)
+                    logps_s[
+                        :, compute_start + i : compute_start + i + h_chunk.shape[1]
+                    ] = chunk_lp.gather(-1, l_chunk.unsqueeze(-1)).squeeze(-1)
+                    del logits_chunk, chunk_lp
+                per_token_logps_list.append(logps_s)
+                del hidden_s, backbone_out, region_h
+
+            per_token_logps = torch.cat(per_token_logps_list, dim=0)
+        else:
+            # Fallback: full forward (non-standard model architecture)
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                outputs = model(
+                    full_sequences,
+                    attention_mask=attn_mask,
+                    position_ids=pos_ids,
+                    return_dict=True,
+                )
+            logits = outputs.logits
+            per_token_logps = torch.zeros(
+                logits.shape[0],
+                logits.shape[1] - 1,
+                device=device,
+                dtype=torch.float32,
+            )
+            region_logits = logits[:, compute_start:-1, :]
+            region_labels = shift_labels[:, compute_start:]
+            chunk_size = 256
+            for i in range(0, region_logits.shape[1], chunk_size):
+                chunk_logits = region_logits[:, i : i + chunk_size, :]
+                chunk_labels = region_labels[:, i : i + chunk_size]
+                chunk_lp = F.log_softmax(chunk_logits.float(), dim=-1)
+                per_token_logps[
+                    :, compute_start + i : compute_start + i + chunk_logits.shape[1]
+                ] = chunk_lp.gather(-1, chunk_labels.unsqueeze(-1)).squeeze(-1)
+            del logits, region_logits
+
+        action_mask = torch.zeros(
+            per_token_logps.shape, dtype=torch.bool, device=device
+        )
+        # Only mark actual generated tokens (not padding beyond num_blocks * gen_len)
+        gen_end = gen_start + num_blocks * gen_len
+        action_mask[:, gen_start:gen_end] = True
+
+        # --- Step 4: Extract features and compute rewards ---
+        with torch.no_grad():
+            block_rewards = self._compute_block_rewards(
+                full_sequences,
+                attn_mask,
+                pos_ids,
+                input_ids,
+                num_blocks,
+                B,
+                n_samples,
+                anchor_offset=anchor_offset,
+            )
+
+        del attn_mask, pos_ids
+        torch.cuda.empty_cache()
+
+        # --- Step 5: Compute advantages ---
+        advantages_per_block = self._compute_advantages(
+            block_rewards, B, n_samples, num_blocks
+        )
+
+        token_advantages = advantages_per_block.repeat_interleave(gen_len, dim=1)
+        full_advantages = torch.zeros_like(per_token_logps)
+        # Only fill actual generated region (not padding beyond num_blocks * gen_len)
+        adv_len = token_advantages.shape[1]  # = num_blocks * gen_len
+        full_advantages[:, gen_start : gen_start + adv_len] = token_advantages
+
+        # --- Step 6: Compute loss ---
+        # RL loss: REINFORCE on generated tokens (needs grad through per_token_logps)
+        rl_loss_per_token = -per_token_logps * full_advantages.detach()
+        rl_loss = (
+            rl_loss_per_token * action_mask.float()
+        ).sum() / action_mask.float().sum().clamp(min=1)
+
+        # CE loss: For structured data, only compute on completion ground-truth tokens
+        # (labels != -100 in the original input). For unstructured data, compute on
+        # all non-action (prompt) tokens as before.
+        ce_loss = torch.tensor(0.0, device=device)
+        if self.ebft_ce_coef > 0:
+            if is_structured and "labels" in inputs:
+                labels = inputs["labels"].to(device)  # (B, seq_len)
+                shifted_labels = labels[:, 1:]  # (B, seq_len - 1)
+                ce_mask_base = shifted_labels != -100  # (B, seq_len - 1)
+                ce_mask_repeated = ce_mask_base.repeat_interleave(n_samples, dim=0)
+                ce_mask = torch.zeros(
+                    per_token_logps.shape, dtype=torch.bool, device=device
+                )
+                ce_mask[:, : ce_mask_repeated.shape[1]] = ce_mask_repeated
+                ce_mask[:, gen_start:] = False
+            else:
+                ce_mask = ~action_mask
+            ce_loss = (
+                -per_token_logps * ce_mask.float()
+            ).sum() / ce_mask.float().sum().clamp(min=1)
+
+        loss = self.ebft_rl_coef * rl_loss + self.ebft_ce_coef * ce_loss
+
+        # --- Log metrics ---
+        if self.state.global_step % self.args.logging_steps == 0:
+            _alignment = getattr(self, "_last_alignment", 0.0)
+            _diversity = getattr(self, "_last_diversity", 0.0)
+            _cfm = getattr(self, "_last_cfm", 0.0)
+            _mean_reward = block_rewards.mean().item()
+            _adv_std = advantages_per_block.std().item()
+
+            log_dict = {
+                "ebft/rl_loss": rl_loss.item(),
+                "ebft/ce_loss": ce_loss.item(),
+                "ebft/cfm_loss": _cfm,
+                "ebft/mean_reward": _mean_reward,
+                "ebft/alignment": _alignment,
+                "ebft/diversity": _diversity,
+                "ebft/num_blocks": num_blocks,
+                "ebft/advantages_std": _adv_std,
+            }
+            if is_structured:
+                log_dict["ebft/anchor_offset"] = anchor_offset
+            self.log(log_dict)
+
+            # Human-readable summary with direction arrows:
+            #   alignment (^ better)  — cosine sim to GT features, range [-2, 2]
+            #   diversity (v better)  — pairwise sim penalty, lower = more diverse
+            #   cfm_loss  (v better)  — ||E[phi(y_hat)] - phi(y)||^2
+            #   reward    (^ better)  — alignment - diversity
+            LOG.info(
+                f"step {self.state.global_step} | "
+                f"align {_alignment:+.3f} ^ | "
+                f"divers {_diversity:+.3f} v | "
+                f"cfm {_cfm:.3f} v | "
+                f"reward {_mean_reward:+.3f} ^ | "
+                f"adv_std {_adv_std:.3f} | "
+                f"blocks {num_blocks}"
+            )
+
+        return (loss, outputs) if return_outputs else loss
+
+    @torch._dynamo.disable
+    @torch.no_grad()
+    def _generate_strided_blocks(
+        self, model, prompt_ids, num_blocks, anchor_offset=None
+    ):
+        """Generate tokens using strided block-parallel attention.
+
+        Uses eager attention (dense 4D masks) during generation to avoid dynamo
+        recompilation — each generation step has a different sequence length.
+        The training forward pass (fixed size) uses flex_attention when available.
+
+        Args:
+            anchor_offset: Position where anchors start. Defaults to context_length.
+        """
+        B, seq_len = prompt_ids.shape
+        gen_len = self.ebft_generate_max_len
+        stride = self.ebft_stride
+        if anchor_offset is None:
+            anchor_offset = self.ebft_context_length
+        temperature = self.ebft_temperature
+        top_p = self.ebft_top_p
+        device = prompt_ids.device
+        model_dtype = next(model.parameters()).dtype
+
+        full_sequence = prompt_ids.clone()
+
+        # Force eager attention during generation to avoid dynamo recompiles from:
+        # 1. Variable sequence lengths per gen step → size-mismatch recompiles
+        # 2. no_grad vs grad toggling → grad_mode recompiles
+        # Both cause dynamo to hit the recompile limit → unfused fallback → OOM
+        unwrapped = self.accelerator.unwrap_model(model)
+        with override_attn_implementation(unwrapped, "eager"):
+            for generation_step in range(gen_len):
+                cur_len = full_sequence.shape[1]
+
+                dense_mask, pos_ids = build_strided_dense_mask_and_positions(
+                    full_sequence_length=cur_len,
+                    prompt_length=seq_len,
+                    context_length=anchor_offset,
+                    generation_step=generation_step,
+                    max_generation_length=gen_len,
+                    stride=stride,
+                    num_blocks=num_blocks,
+                    device=device,
+                    batch_size=B,
+                    dtype=model_dtype,
+                )
+
+                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                    output = model(
+                        full_sequence,
+                        attention_mask=dense_mask,
+                        position_ids=pos_ids,
+                        return_dict=True,
+                    )
+                all_logits = output.logits
+
+                logit_positions = []
+                for block_idx in range(num_blocks):
+                    if generation_step == 0:
+                        # Last token of the context window predicts the first rollout token
+                        pos = anchor_offset + block_idx * stride - 1
+                    else:
+                        pos = seq_len + (generation_step - 1) * num_blocks + block_idx
+                    logit_positions.append(pos)
+
+                position_indices = torch.tensor(logit_positions, device=device)
+                block_logits = all_logits.index_select(1, position_indices)
+
+                if temperature > 0:
+                    block_logits = block_logits / temperature
+                    probs = torch.softmax(block_logits, dim=-1)
+
+                    if top_p < 1.0:
+                        sorted_probs, sorted_idx = torch.sort(
+                            probs, descending=True, dim=-1
+                        )
+                        cumulative = torch.cumsum(sorted_probs, dim=-1)
+                        remove = cumulative > top_p
+                        remove[..., 1:] = remove[..., :-1].clone()
+                        remove[..., 0] = False
+                        mask = torch.zeros_like(probs, dtype=torch.bool)
+                        mask.scatter_(-1, sorted_idx, remove)
+                        probs[mask] = 0
+                        probs = probs / probs.sum(dim=-1, keepdim=True)
+
+                    flat_probs = probs.view(-1, probs.shape[-1])
+                    sampled = torch.multinomial(flat_probs, 1).squeeze(-1)
+                    sampled = sampled.view(B, num_blocks)
+                else:
+                    sampled = torch.argmax(block_logits, dim=-1)
+
+                full_sequence = torch.cat([full_sequence, sampled], dim=1)
+
+        return full_sequence
+
+    @torch._dynamo.disable
+    @torch.no_grad()
+    def _compute_block_rewards(
+        self,
+        full_sequences,
+        attn_mask,
+        pos_ids,
+        original_ids,
+        num_blocks,
+        batch_size,
+        n_samples,
+        anchor_offset=None,
+    ):
+        """Extract features and compute per-block rewards. Returns (B, N, NB).
+
+        Args:
+            anchor_offset: Position where anchors start. For structured data this
+                is after the prompt; for unstructured it equals context_length.
+        """
+        device = full_sequences.device
+        seq_len = original_ids.shape[1]
+        gen_len = self.ebft_generate_max_len
+        stride = self.ebft_stride
+        if anchor_offset is None:
+            anchor_offset = self.ebft_context_length
+
+        # Run feature network on its device WITH the strided attention mask.
+        # Without the strided mask, generated tokens see tokens from other blocks
+        # via default causal attention, corrupting the feature representations.
+        fd = self._feature_device
+        fn_seqs = full_sequences.to(fd)
+        fn_pos = pos_ids.to(fd)
+
+        # Determine which model to use for feature extraction
+        if self._share_feature_weights:
+            # Use actor's base weights with adapters disabled.
+            # Force eager attention to avoid grad_mode recompiles on the shared
+            # compiled flex_attention kernel (feature extraction is no_grad,
+            # training forward is with grad — each switch recompiles).
+            unwrapped_actor = self.accelerator.unwrap_model(self.model)
+            feat_model = unwrapped_actor
+            feature_ctx = unwrapped_actor.disable_adapter()
+            # Use SDPA (flash attention) instead of flex to avoid grad_mode recompiles
+            # on the shared compiled flex kernel. SDPA is fused (no score matrix
+            # materialization) and needs no compilation — ideal for no_grad feature extraction.
+            attn_ctx = override_attn_implementation(unwrapped_actor, "sdpa")
+            use_flex_for_features = False
+        else:
+            feat_model = self.feature_network
+            feature_ctx = contextlib.nullcontext()
+            attn_ctx = contextlib.nullcontext()
+            use_flex_for_features = self._feature_use_flex
+
+        # Build strided mask — flex block mask if available, else dense 4D
+        if use_flex_for_features:
+            fn_attn_mask = create_strided_block_mask(
+                prompt_length=seq_len,
+                context_length=anchor_offset,
+                max_generation_length=gen_len,
+                stride=stride,
+                num_blocks=num_blocks,
+                full_sequence_length=full_sequences.shape[1],
+                batch_size=full_sequences.shape[0],
+                num_heads=feat_model.config.num_attention_heads,
+                device=fd,
+            )
+        else:
+            fn_attn_mask, _ = build_strided_dense_mask_and_positions(
+                full_sequence_length=full_sequences.shape[1],
+                prompt_length=seq_len,
+                context_length=anchor_offset,
+                generation_step=gen_len,
+                max_generation_length=gen_len,
+                stride=stride,
+                num_blocks=num_blocks,
+                device=fd,
+                batch_size=full_sequences.shape[0],
+                dtype=torch.bfloat16,
+            )
+
+        with (
+            feature_ctx,
+            attn_ctx,
+            torch.autocast(device_type="cuda", dtype=torch.bfloat16),
+        ):
+            was_training = feat_model.training
+            feat_model.eval()
+            fn_outputs = feat_model(
+                fn_seqs,
+                attention_mask=fn_attn_mask,
+                position_ids=fn_pos,
+                output_hidden_states=True,
+                return_dict=True,
+            )
+            if was_training:
+                feat_model.train()
+        hidden_states_cpu = [
+            fn_outputs.hidden_states[idx].to(device)
+            for idx in self.feature_layer_indices
+        ]
+        del fn_outputs, fn_seqs, fn_pos, fn_attn_mask
+
+        # Normalize each layer's hidden states separately (like the reference critic),
+        # then concatenate. This prevents one dominant layer from suppressing others.
+        normalized_layers = [F.normalize(h, p=2, dim=-1) for h in hidden_states_cpu]
+        features = torch.cat(normalized_layers, dim=-1).to(device)
+        del hidden_states_cpu, normalized_layers
+
+        # Ground-truth features start from anchor_offset (not ctx_len) so they
+        # align with where anchors are actually placed.
+        gt_features = features[:, anchor_offset:seq_len, :]
+        # Only take actual generated tokens (exclude padding beyond num_blocks * gen_len)
+        gen_features = features[:, seq_len : seq_len + num_blocks * gen_len, :]
+
+        gt_block_features = gt_features.unfold(1, gen_len, stride).permute(0, 1, 3, 2)
+        gen_block_features = gen_features.reshape(
+            batch_size * n_samples, gen_len, num_blocks, -1
+        ).transpose(1, 2)
+
+        if self.ebft_embed_method == "mean_pooling":
+            gt_emb = gt_block_features.mean(dim=2)
+            gen_emb = gen_block_features.mean(dim=2)
+        else:  # last_token
+            gt_emb = gt_block_features[:, :, -1, :]
+            gen_emb = gen_block_features[:, :, -1, :]
+
+        gt_emb = gt_emb.view(batch_size, n_samples, num_blocks, -1)
+        gen_emb = gen_emb.view(batch_size, n_samples, num_blocks, -1)
+
+        if self.ebft_use_whitening:
+            whitened_gen, whitened_gt = [], []
+            for b in range(batch_size):
+                for nb in range(num_blocks):
+                    w_gen, w_gt = whiten_embeddings_batched(
+                        gen_emb[b, :, nb, :],
+                        gt_emb[b, :, nb, :],
+                    )
+                    whitened_gen.append(w_gen)
+                    whitened_gt.append(w_gt)
+            gen_emb = (
+                torch.stack(whitened_gen)
+                .view(batch_size, num_blocks, n_samples, -1)
+                .transpose(1, 2)
+            )
+            gt_emb = (
+                torch.stack(whitened_gt)
+                .view(batch_size, num_blocks, n_samples, -1)
+                .transpose(1, 2)
+            )
+
+        alignment = F.cosine_similarity(gen_emb, gt_emb, dim=-1)
+
+        # Batched diversity: reshape to avoid per-block Python loop
+        diversity = torch.zeros_like(alignment)
+        if n_samples > 1:
+            # (B, N, NB, D) → (B*NB, N, D) for a single batched bmm
+            gen_for_div = gen_emb.permute(0, 2, 1, 3).reshape(
+                batch_size * num_blocks, n_samples, -1
+            )
+            sims = torch.bmm(gen_for_div, gen_for_div.transpose(1, 2))  # (B*NB, N, N)
+            eye = torch.eye(n_samples, device=device, dtype=torch.bool)
+            sims = sims.masked_fill(eye.unsqueeze(0), 0.0)
+            div_flat = sims.sum(dim=-1) / (n_samples - 1)  # (B*NB, N)
+            diversity = div_flat.view(batch_size, num_blocks, n_samples).permute(
+                0, 2, 1
+            )  # (B, N, NB)
+
+        # Scale by 2 per paper equation (7):
+        #   r_j = 2*φ(ŷ_j)^T*φ(y) - 2/(n-1) * Σ_{j'≠j} φ(ŷ_j)^T*φ(ŷ_{j'})
+        alignment = alignment * 2
+        diversity = diversity * 2
+
+        # Compute CFM loss: ||E[φ(ŷ)] - φ(y)||^2 (paper eq 2)
+        # Mean generated embedding per prompt, squared distance to GT
+        mean_gen_emb = gen_emb.mean(dim=1, keepdim=True)  # (B, 1, NB, D)
+        gt_for_cfm = gt_emb[:, 0:1, :, :]  # (B, 1, NB, D) — one GT per prompt
+        cfm_loss = ((mean_gen_emb - gt_for_cfm) ** 2).sum(dim=-1).mean()
+
+        # Store for logging
+        self._last_alignment = alignment.mean().item()
+        self._last_diversity = diversity.mean().item()
+        self._last_cfm = cfm_loss.item()
+
+        return (
+            self.ebft_alignment_coef * alignment - self.ebft_diversity_coef * diversity
+        )
+
+    def _compute_advantages(self, rewards, batch_size, n_samples, num_blocks):
+        """Compute RLOO advantages. rewards: (B, N, NB) → (B*N, NB)."""
+        if self.ebft_advantage_estimator == "rloo" and n_samples > 1:
+            total = rewards.sum(dim=1, keepdim=True)
+            baseline = (total - rewards) / (n_samples - 1)
+            advantages = rewards - baseline
+        elif self.ebft_advantage_estimator == "group_norm" and n_samples > 1:
+            mean = rewards.mean(dim=1, keepdim=True)
+            std = rewards.std(dim=1, keepdim=True) + 1e-8
+            advantages = (rewards - mean) / std
+        else:
+            advantages = rewards
+        return advantages.view(batch_size * n_samples, num_blocks)
diff --git a/src/axolotl/core/trainers/ebft/trainer.py b/src/axolotl/core/trainers/ebft/trainer.py
new file mode 100644
index 0000000000..1c27fa91f0
--- /dev/null
+++ b/src/axolotl/core/trainers/ebft/trainer.py
@@ -0,0 +1,531 @@
+"""
+EBFT Trainer — Energy-Based Fine-Tuning integrated via GRPOTrainer.
+
+Extends AxolotlGRPOTrainer by plugging feature-matching rewards into
+the standard GRPO reward function interface.
+
+Paper: "Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models"
+       (Jelassi et al., 2026) https://arxiv.org/abs/2603.12248
+"""
+
+import contextlib
+import copy
+from typing import TYPE_CHECKING, Any
+
+import torch
+from datasets import Dataset, IterableDataset
+from peft import PeftModel
+from transformers import PreTrainedModel, PreTrainedTokenizerBase, TrainerCallback
+
+from axolotl.core.trainers.ebft.args import AxolotlEBFTConfig
+from axolotl.core.trainers.ebft.rewards import (
+    apply_embed_method,
+    extract_hidden_states,
+    get_alignment_rewards,
+    get_diversity_rewards,
+    whiten_embeddings_batched,
+)
+from axolotl.core.trainers.grpo.trainer import (
+    AxolotlAsyncGRPOTrainer,
+    AxolotlGRPOTrainer,
+)
+from axolotl.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections import defaultdict
+
+    from accelerate import Accelerator
+    from trl.generation.vllm_generation import VLLMGeneration
+
+LOG = get_logger(__name__)
+
+
+class EBFTMixin:
+    """
+    Mixin that adds EBFT feature-matching reward logic to any GRPO-based trainer.
+
+    Provides:
+    - Frozen feature network setup (shared weights for PEFT, deepcopy otherwise)
+    - _feature_matching_reward() callable for GRPO reward function interface
+    - _sequential_rollout() for multi-turn conversations
+    """
+
+    # Type stubs for attributes provided by the composed GRPOTrainer base class.
+    # These are not defined here but accessed via cooperative multiple inheritance.
+    if TYPE_CHECKING:
+        accelerator: Accelerator
+        model: PreTrainedModel
+        args: AxolotlEBFTConfig
+        processing_class: PreTrainedTokenizerBase
+        num_generations: int
+        vllm_generation: VLLMGeneration
+        _metrics: defaultdict
+
+    _tag_names = ["trl", "ebft", "axolotl"]
+
+    def __init__(
+        self,
+        model: str | PreTrainedModel,
+        args: AxolotlEBFTConfig | None = None,
+        train_dataset: Dataset | IterableDataset | None = None,
+        eval_dataset: Dataset
+        | IterableDataset
+        | dict[str, Dataset | IterableDataset]
+        | None = None,
+        processing_class: PreTrainedTokenizerBase | None = None,
+        callbacks: list[TrainerCallback] | None = None,
+        optimizers: tuple[
+            torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None
+        ] = (None, None),
+        peft_config: Any | None = None,
+    ):
+        # Pass our feature-matching reward function to GRPOTrainer
+        # It will be called with (prompts, completions, **kwargs) where
+        # kwargs includes all extra dataset fields like "ground_truth"
+        super().__init__(  # type: ignore[call-arg]
+            model=model,
+            reward_funcs=[self._feature_matching_reward],
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=processing_class,
+            callbacks=callbacks,
+            optimizers=optimizers,
+            peft_config=peft_config,
+        )
+        assert args is not None
+
+        # --- Feature network setup ---
+        unwrapped = self.accelerator.unwrap_model(self.model)
+        # Check for PEFT model — use hasattr for robustness across DDP/FSDP wrapping
+        self._share_feature_weights = isinstance(unwrapped, PeftModel) or hasattr(
+            unwrapped, "disable_adapter"
+        )
+
+        if self._share_feature_weights:
+            # Share weights: use actor's base model with adapters disabled.
+            # Saves a full model copy (~8 GB for 4B model).
+            self.feature_network = None
+            param_gb = sum(p.numel() for p in unwrapped.parameters()) * 2 / 1e9
+            LOG.info(
+                f"EBFT feature network shares actor weights (PEFT disable_adapter). "
+                f"Saving ~{param_gb:.1f} GB"
+            )
+        else:
+            LOG.info("Creating frozen feature network for EBFT (deepcopy)...")
+            self.feature_network = copy.deepcopy(unwrapped)
+            for param in self.feature_network.parameters():
+                param.requires_grad = False
+            self.feature_network.eval()
+
+        # Compute layer indices from fractional depths
+        # Handle VLM models where num_hidden_layers is on text_config
+        config = unwrapped.config
+        if hasattr(config, "text_config") and hasattr(
+            config.text_config, "num_hidden_layers"
+        ):
+            config = config.text_config
+        num_layers = config.num_hidden_layers
+        self.feature_layer_indices = [
+            int(frac * num_layers) for frac in args.ebft_feature_layers
+        ]
+        LOG.info(
+            f"EBFT feature extraction from layers {self.feature_layer_indices} "
+            f"(of {num_layers} total), embed_method={args.ebft_embed_method}"
+        )
+        if args.ebft_adaptive_max_tokens:
+            LOG.info(
+                f"EBFT adaptive max_tokens enabled "
+                f"(gt_length_multiplier={args.ebft_gt_length_multiplier})"
+            )
+
+    _adaptive_max_lock = None  # initialized lazily
+
+    def _generate_only(self, inputs, rank0_only=False):
+        """Override to set per-batch max_tokens based on ground-truth length.
+
+        Uses a lock to prevent race conditions in async mode where concurrent
+        BG threads could interleave mutations of max_completion_length.
+        """
+        import threading
+
+        args = self.args
+        if (
+            args.ebft_adaptive_max_tokens
+            and hasattr(self, "vllm_generation")
+            and inputs
+        ):
+            gt_texts = [
+                x.get("ground_truth", "") for x in inputs if x.get("ground_truth")
+            ]
+            if gt_texts:
+                gt_token_counts = [
+                    len(self.processing_class.encode(gt, add_special_tokens=False))
+                    for gt in gt_texts
+                ]
+                multiplier = args.ebft_gt_length_multiplier
+                max_completion = self.vllm_generation.max_completion_length
+                adaptive_max = max(
+                    min(int(c * multiplier), max_completion) for c in gt_token_counts
+                )
+                adaptive_max = max(adaptive_max, 64)
+
+                if self._adaptive_max_lock is None:
+                    self._adaptive_max_lock = threading.Lock()
+                with self._adaptive_max_lock:
+                    original = self.vllm_generation.max_completion_length
+                    self.vllm_generation.max_completion_length = adaptive_max
+                    try:
+                        return super()._generate_only(inputs, rank0_only)
+                    finally:
+                        self.vllm_generation.max_completion_length = original
+
+        return super()._generate_only(inputs, rank0_only)
+
+    @torch.no_grad()
+    def _feature_matching_reward(
+        self,
+        prompts: list,
+        completions: list,
+        ground_truth: list[str] | None = None,
+        remaining_turns: list | None = None,
+        **kwargs,
+    ) -> list[float]:
+        """
+        Compute feature-matching rewards for generated completions.
+
+        This is called by GRPOTrainer's _generate_and_score_completions()
+        as a standard reward function. The `ground_truth` field comes from
+        the dataset via reward_kwargs.
+
+        For multi-turn conversations, `remaining_turns` contains the subsequent
+        user/assistant turn pairs. When present, we do sequential rollouts:
+        generate each assistant turn conditioned on history + previous generations,
+        then compute feature-matching rewards on the full generated conversation.
+
+        Args:
+            prompts: List of prompt strings/messages
+            completions: List of generated completion strings
+            ground_truth: List of reference completion strings (from dataset)
+            remaining_turns: List of remaining conversation turns after the
+                first assistant turn (for multi-turn rollouts)
+
+        Returns:
+            List of scalar rewards, one per completion
+        """
+        if ground_truth is None:
+            LOG.warning("No ground_truth field in dataset — using zero rewards")
+            return [0.0] * len(prompts)
+
+        device = self.accelerator.device
+        args = self.args
+        num_gens = self.num_generations
+
+        # --- Multi-turn sequential rollout ---
+        # If remaining_turns is provided, generate subsequent assistant turns
+        # by calling vLLM for each turn, building up the full conversation.
+        if remaining_turns is not None and hasattr(self, "vllm_generation"):
+            completions = self._sequential_rollout(
+                prompts, completions, remaining_turns, num_gens
+            )
+
+        # --- Tokenize generated sequences: prompt + completion ---
+        gen_texts = []
+        gen_prompt_texts = []
+        for p, c in zip(prompts, completions, strict=True):
+            if isinstance(p, list):
+                prompt_text = self.processing_class.apply_chat_template(
+                    p, tokenize=False, add_generation_prompt=True
+                )
+            else:
+                prompt_text = p
+            if isinstance(c, list):
+                comp_text = c[0].get("content", "") if c else ""
+            else:
+                comp_text = c
+            gen_texts.append(prompt_text + comp_text)
+            gen_prompt_texts.append(prompt_text)
+
+        gen_encoded = self.processing_class(
+            text=gen_texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=getattr(self.args, "max_length", None)
+            or getattr(self.args, "max_seq_length", None)
+            or 2048,
+            add_special_tokens=False,
+        )
+        gen_ids = gen_encoded["input_ids"].to(device)
+        gen_mask = gen_encoded["attention_mask"].to(device)
+
+        # Compute prompt lengths for completion_mean pooling
+        gen_prompt_lengths = torch.tensor(
+            [
+                len(self.processing_class.encode(pt, add_special_tokens=False))
+                for pt in gen_prompt_texts
+            ],
+            device=device,
+        )
+
+        # --- Tokenize ground-truth sequences: prompt + ground_truth ---
+        # For multi-turn (remaining_turns present), render the full GT conversation
+        # through the chat template to preserve role markers between turns.
+        gt_texts = []
+        gt_prompt_texts = []
+        for i, (p, gt) in enumerate(zip(prompts, ground_truth, strict=True)):
+            if i % num_gens != 0:
+                continue  # Only need one GT per prompt group
+            if isinstance(p, list):
+                prompt_text = self.processing_class.apply_chat_template(
+                    p, tokenize=False, add_generation_prompt=True
+                )
+                # Multi-turn: build full GT conversation with remaining turns
+                if remaining_turns is not None:
+                    prompt_idx = i // num_gens
+                    turns = (
+                        remaining_turns[prompt_idx]
+                        if prompt_idx < len(remaining_turns)
+                        else []
+                    )
+                    if turns:
+                        gt_conv = list(p) + [{"role": "assistant", "content": gt}]
+                        gt_conv.extend(turns)
+                        full_gt_text = self.processing_class.apply_chat_template(
+                            gt_conv, tokenize=False, add_generation_prompt=False
+                        )
+                        gt_texts.append(full_gt_text)
+                        gt_prompt_texts.append(prompt_text)
+                        continue
+            else:
+                prompt_text = p
+            gt_texts.append(prompt_text + gt)
+            gt_prompt_texts.append(prompt_text)
+
+        gt_encoded = self.processing_class(
+            text=gt_texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=getattr(self.args, "max_length", None)
+            or getattr(self.args, "max_seq_length", None)
+            or 2048,
+            add_special_tokens=False,
+        )
+        gt_ids = gt_encoded["input_ids"].to(device)
+        gt_mask = gt_encoded["attention_mask"].to(device)
+
+        gt_prompt_lengths = torch.tensor(
+            [
+                len(self.processing_class.encode(pt, add_special_tokens=False))
+                for pt in gt_prompt_texts
+            ],
+            device=device,
+        )
+
+        # --- Extract features from frozen feature network ---
+        # INVARIANT: disable_adapter() yields the unmodified base weights because
+        # _sync_peft_weights_no_merge and _sync_lora_adapter never call
+        # merge_adapter() — they compute merged weights as new tensors or save
+        # the adapter to filesystem. Base weights are never modified in-place.
+        if self._share_feature_weights:
+            unwrapped = self.accelerator.unwrap_model(self.model)
+            feature_ctx = unwrapped.disable_adapter()
+        else:
+            unwrapped = self.feature_network
+            feature_ctx = contextlib.nullcontext()
+
+        with feature_ctx:
+            was_training = unwrapped.training
+            unwrapped.eval()
+            gen_hidden = extract_hidden_states(
+                unwrapped, gen_ids, gen_mask, self.feature_layer_indices
+            )
+            gt_hidden = extract_hidden_states(
+                unwrapped, gt_ids, gt_mask, self.feature_layer_indices
+            )
+            if was_training:
+                unwrapped.train()
+
+        # --- Pool to sequence-level embeddings ---
+        gen_emb = apply_embed_method(
+            gen_hidden,
+            args.ebft_embed_method,
+            gen_mask,
+            prompt_lengths=gen_prompt_lengths,
+        )
+        gt_emb = apply_embed_method(
+            gt_hidden,
+            args.ebft_embed_method,
+            gt_mask,
+            prompt_lengths=gt_prompt_lengths,
+        )
+
+        # --- Optional whitening ---
+        batch_size = gen_emb.shape[0]
+        if args.ebft_use_whitening and batch_size > 1:
+            num_prompts = batch_size // num_gens
+            gen_reshaped = gen_emb.view(num_prompts, num_gens, -1)
+            whitened_gen_list = []
+            whitened_gt_list = []
+            for i in range(num_prompts):
+                w_gen, w_gt = whiten_embeddings_batched(
+                    gen_reshaped[i], gt_emb[i : i + 1]
+                )
+                whitened_gen_list.append(w_gen)
+                whitened_gt_list.append(w_gt)
+            gen_emb = torch.cat(whitened_gen_list, dim=0)
+            gt_emb = torch.cat(whitened_gt_list, dim=0)
+        else:
+            gen_emb = torch.nn.functional.normalize(gen_emb, p=2, dim=-1)
+            gt_emb = torch.nn.functional.normalize(gt_emb, p=2, dim=-1)
+
+        # Repeat gt_emb: each GT repeated num_generations times
+        gt_emb_expanded = gt_emb.repeat_interleave(num_gens, dim=0)
+
+        # --- Compute rewards ---
+        alignment = get_alignment_rewards(gen_emb, gt_emb_expanded)
+        diversity = get_diversity_rewards(gen_emb, num_gens)
+
+        # Scale by 2 per paper equation (7):
+        #   r_j = 2*φ(ŷ_j)^T*φ(y) - 2/(n-1) * Σ_{j'≠j} φ(ŷ_j)^T*φ(ŷ_{j'})
+        alignment = alignment * 2
+        diversity = diversity * 2
+
+        rewards = (
+            args.ebft_alignment_coef * alignment - args.ebft_diversity_coef * diversity
+        )
+
+        # Compute CFM loss: ||E[φ(ŷ)] - φ(y)||^2 (paper eq 2)
+        gen_reshaped = gen_emb.view(-1, num_gens, gen_emb.shape[-1])
+        mean_gen = gen_reshaped.mean(dim=1)  # (num_prompts, D)
+        cfm_loss = ((mean_gen - gt_emb) ** 2).sum(dim=-1).mean()
+
+        # Log feature-matching metrics to console and wandb
+        _align = alignment.mean().item()
+        _divers = diversity.mean().item()
+        _reward = rewards.mean().item()
+        _cfm = cfm_loss.item()
+
+        LOG.info(
+            f"ebft reward | "
+            f"align {_align:+.3f} ^ | "
+            f"divers {_divers:+.3f} v | "
+            f"cfm {_cfm:.3f} v | "
+            f"reward {_reward:+.3f} ^"
+        )
+
+        # Log to wandb via trainer's _metrics (picked up by GRPO's logging)
+        mode = "train" if self.model.training else "eval"
+        if hasattr(self, "_metrics"):
+            self._metrics[mode]["ebft/alignment"].append(_align)
+            self._metrics[mode]["ebft/diversity"].append(_divers)
+            self._metrics[mode]["ebft/cfm_loss"].append(_cfm)
+            self._metrics[mode]["ebft/reward"].append(_reward)
+
+        return rewards.cpu().tolist()
+
+    @torch.no_grad()
+    def _sequential_rollout(
+        self,
+        prompts: list,
+        first_completions: list,
+        remaining_turns: list,
+        num_gens: int,
+    ) -> list:
+        """
+        Extend single-turn completions into multi-turn conversations.
+
+        For each prompt group, takes the first generated assistant turn and
+        sequentially generates subsequent assistant turns by calling vLLM,
+        building up a full multi-turn conversation.
+
+        Args:
+            prompts: List of prompt message lists (repeated num_gens times)
+            first_completions: List of generated first-turn completions
+            remaining_turns: List of remaining turn pairs after first assistant turn.
+                Each element is a list of dicts: [{"role": "user", "content": "..."},
+                {"role": "assistant", "content": "...GT..."}]
+            num_gens: Number of generations per prompt
+
+        Returns:
+            Extended completions incorporating all generated turns
+        """
+        vllm_client = self.vllm_generation.vllm_client
+        max_tokens = getattr(self.args, "max_completion_length", 256)
+        temperature = getattr(self.args, "temperature", 0.7)
+        gen_kwargs = getattr(self.args, "generation_kwargs", None) or {}
+
+        extended_completions = []
+
+        for idx in range(len(prompts)):
+            prompt_msgs = prompts[idx] if isinstance(prompts[idx], list) else []
+            first_comp = first_completions[idx]
+
+            # Extract first completion text
+            if isinstance(first_comp, list):
+                first_text = first_comp[0].get("content", "") if first_comp else ""
+            else:
+                first_text = first_comp
+
+            # Get remaining turns for this prompt (same for all num_gens copies)
+            prompt_idx = idx // num_gens
+            turns = (
+                remaining_turns[prompt_idx] if prompt_idx < len(remaining_turns) else []
+            )
+
+            if not turns:
+                extended_completions.append(first_text)
+                continue
+
+            # Build conversation with generated first turn
+            conv = list(prompt_msgs) + [{"role": "assistant", "content": first_text}]
+
+            # Generate subsequent turns
+            for turn in turns:
+                if turn["role"] == "user":
+                    conv.append(turn)
+                elif turn["role"] == "assistant":
+                    try:
+                        result = vllm_client.chat(
+                            messages=[conv],
+                            n=1,
+                            max_tokens=max_tokens,
+                            temperature=temperature,
+                            generation_kwargs=gen_kwargs,
+                        )
+                        gen_ids = result.get("completion_ids", [[]])[0]
+                        gen_text = self.processing_class.decode(
+                            gen_ids, skip_special_tokens=True
+                        )
+                    except Exception as e:
+                        LOG.warning(f"Multi-turn rollout generation failed: {e}")
+                        gen_text = ""
+
+                    conv.append({"role": "assistant", "content": gen_text})
+
+            # Render full conversation through chat template, then extract
+            # everything after the original prompt as the "completion" text.
+            # This preserves role markers and formatting between turns.
+            full_rendered = self.processing_class.apply_chat_template(
+                conv, tokenize=False, add_generation_prompt=False
+            )
+            prompt_rendered = self.processing_class.apply_chat_template(
+                prompt_msgs, tokenize=False, add_generation_prompt=True
+            )
+            completion_text = full_rendered[len(prompt_rendered) :]
+            extended_completions.append(completion_text)
+
+        return extended_completions
+
+
+class AxolotlEBFTTrainer(EBFTMixin, AxolotlGRPOTrainer):
+    """EBFT trainer using synchronous GRPO (standard vLLM generation)."""
+
+    pass
+
+
+class AxolotlAsyncEBFTTrainer(EBFTMixin, AxolotlAsyncGRPOTrainer):
+    """EBFT trainer using async GRPO (prefetches next batch during training)."""
+
+    pass
diff --git a/src/axolotl/core/trainers/grpo/__init__.py b/src/axolotl/core/trainers/grpo/__init__.py
new file mode 100644
index 0000000000..bb0046e57b
--- /dev/null
+++ b/src/axolotl/core/trainers/grpo/__init__.py
@@ -0,0 +1,318 @@
+"""GRPO Specific Strategy for training"""
+
+import importlib
+import inspect
+import os
+from typing import Any
+
+from huggingface_hub import snapshot_download
+from requests import HTTPError
+from trl.trainer.grpo_trainer import RewardFunc
+
+from axolotl.core.trainers.grpo.args import AxolotlAsyncGRPOConfig, AxolotlGRPOConfig
+from axolotl.core.trainers.grpo.trainer import (
+    AxolotlAsyncGRPOTrainer,
+    AxolotlGRPOSequenceParallelTrainer,
+    AxolotlGRPOTrainer,
+)
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.trl import TRLConfig
+from axolotl.utils.schemas.vllm import VllmConfig
+
+LOG = get_logger(__name__)
+
+
+class GRPOStrategy:
+    """Strategy for GRPO training"""
+
+    @classmethod
+    def get_trainer_class(
+        cls,
+        sequence_parallel: bool = False,
+        async_grpo: bool = False,
+    ) -> (
+        type[AxolotlGRPOTrainer]
+        | type[AxolotlGRPOSequenceParallelTrainer]
+        | type[AxolotlAsyncGRPOTrainer]
+    ):
+        if sequence_parallel and async_grpo:
+            raise ValueError(
+                "sequence_parallel and async_grpo cannot both be enabled. "
+                "Disable one of context_parallel_size > 1 or async_prefetch/use_data_producer."
+            )
+        if sequence_parallel:
+            return AxolotlGRPOSequenceParallelTrainer
+        if async_grpo:
+            return AxolotlAsyncGRPOTrainer
+        return AxolotlGRPOTrainer
+
+    @classmethod
+    def get_training_args_class(
+        cls, async_grpo: bool = False
+    ) -> type[AxolotlGRPOConfig] | type[AxolotlAsyncGRPOConfig]:
+        if async_grpo:
+            return AxolotlAsyncGRPOConfig
+        return AxolotlGRPOConfig
+
+    @classmethod
+    def set_training_args_kwargs(cls, cfg: DictDefault) -> dict[str, Any]:
+        grpo_args_kwargs: dict[str, Any] = {}
+
+        if not hasattr(cfg, "trl") or not cfg.trl:
+            return grpo_args_kwargs
+
+        trl: TRLConfig = cfg.trl  # type: ignore
+        vllm_cfg: VllmConfig = cfg.vllm  # type: ignore
+
+        if trl.use_vllm:
+            grpo_args_kwargs["use_vllm"] = trl.use_vllm
+            if trl.vllm_mode:
+                grpo_args_kwargs["vllm_mode"] = trl.vllm_mode
+            if trl.vllm_mode == "colocate":
+                grpo_args_kwargs["vllm_enable_sleep_mode"] = trl.vllm_enable_sleep_mode  # type: ignore[attr-defined]
+                grpo_args_kwargs["vllm_gpu_memory_utilization"] = (
+                    vllm_cfg.gpu_memory_utilization
+                )
+                grpo_args_kwargs["vllm_tensor_parallel_size"] = (
+                    vllm_cfg.tensor_parallel_size
+                )
+            grpo_args_kwargs["vllm_server_host"] = trl.vllm_server_host or trl.vllm.host  # type: ignore[attr-defined]
+            grpo_args_kwargs["vllm_server_port"] = trl.vllm_server_port or trl.vllm.port  # type: ignore[attr-defined]
+            if trl.vllm_server_timeout:
+                grpo_args_kwargs["vllm_server_timeout"] = trl.vllm_server_timeout
+            if trl.vllm_guided_decoding_regex:
+                grpo_args_kwargs["vllm_guided_decoding_regex"] = (
+                    trl.vllm_guided_decoding_regex
+                )
+
+        if trl.num_generations:
+            grpo_args_kwargs["num_generations"] = trl.num_generations
+        if trl.generation_batch_size is not None:
+            grpo_args_kwargs["generation_batch_size"] = trl.generation_batch_size
+
+        if trl.sync_ref_model:
+            grpo_args_kwargs["sync_ref_model"] = trl.sync_ref_model
+
+            if trl.ref_model_mixup_alpha:
+                grpo_args_kwargs["ref_model_mixup_alpha"] = trl.ref_model_mixup_alpha
+
+            if trl.ref_model_sync_steps:
+                grpo_args_kwargs["ref_model_sync_steps"] = trl.ref_model_sync_steps
+
+        grpo_args_kwargs["max_completion_length"] = trl.max_completion_length
+        grpo_args_kwargs["log_completions"] = trl.log_completions
+        grpo_args_kwargs["num_completions_to_print"] = trl.num_completions_to_print
+
+        if cfg.context_parallel_size > 1:
+            grpo_args_kwargs["context_parallel_size"] = cfg.context_parallel_size
+
+        if trl.importance_sampling_level is not None:
+            grpo_args_kwargs["importance_sampling_level"] = (
+                trl.importance_sampling_level
+            )
+
+        if trl.reward_weights:
+            grpo_args_kwargs["reward_weights"] = trl.reward_weights
+
+        if trl.scale_rewards is not None:
+            grpo_args_kwargs["scale_rewards"] = trl.scale_rewards
+
+        if trl.loss_type is not None:
+            grpo_args_kwargs["loss_type"] = trl.loss_type
+        if trl.mask_truncated_completions is not None:
+            grpo_args_kwargs["mask_truncated_completions"] = (
+                trl.mask_truncated_completions
+            )
+
+        if trl.temperature is not None:
+            grpo_args_kwargs["temperature"] = trl.temperature
+        if trl.top_p is not None:
+            grpo_args_kwargs["top_p"] = trl.top_p
+        if trl.top_k is not None:
+            grpo_args_kwargs["top_k"] = trl.top_k
+        if trl.min_p is not None:
+            grpo_args_kwargs["min_p"] = trl.min_p
+        if trl.repetition_penalty is not None:
+            grpo_args_kwargs["repetition_penalty"] = trl.repetition_penalty
+
+        if trl.num_iterations is not None:
+            grpo_args_kwargs["num_iterations"] = trl.num_iterations
+        if trl.epsilon is not None:
+            grpo_args_kwargs["epsilon"] = trl.epsilon
+        if trl.epsilon_high is not None:
+            grpo_args_kwargs["epsilon_high"] = trl.epsilon_high
+
+        if trl.use_liger_loss is not None:
+            grpo_args_kwargs["use_liger_kernel"] = trl.use_liger_loss
+
+        if trl.multi_objective_aggregation is not None:
+            grpo_args_kwargs["multi_objective_aggregation"] = (
+                trl.multi_objective_aggregation
+            )
+
+        # Async GRPO fields
+        if getattr(trl, "use_data_producer", None) is not None:
+            grpo_args_kwargs["use_data_producer"] = trl.use_data_producer
+        if getattr(trl, "async_prefetch", None) is not None:
+            grpo_args_kwargs["async_prefetch"] = trl.async_prefetch
+        if getattr(trl, "prefetch_depth", None) is not None:
+            grpo_args_kwargs["prefetch_depth"] = trl.prefetch_depth
+        if getattr(trl, "vllm_sync_interval", None) is not None:
+            grpo_args_kwargs["vllm_sync_interval"] = trl.vllm_sync_interval
+        if getattr(trl, "streaming_partial_batch", None) is not None:
+            grpo_args_kwargs["streaming_partial_batch"] = trl.streaming_partial_batch
+        if getattr(trl, "streaming_min_groups", None) is not None:
+            grpo_args_kwargs["streaming_min_groups"] = trl.streaming_min_groups
+        if getattr(trl, "vllm_importance_sampling_correction", None) is not None:
+            grpo_args_kwargs["vllm_importance_sampling_correction"] = (
+                trl.vllm_importance_sampling_correction
+            )
+        if getattr(trl, "vllm_importance_sampling_mode", None) is not None:
+            grpo_args_kwargs["vllm_importance_sampling_mode"] = (
+                trl.vllm_importance_sampling_mode
+            )
+        if getattr(trl, "vllm_importance_sampling_cap", None) is not None:
+            grpo_args_kwargs["vllm_importance_sampling_cap"] = (
+                trl.vllm_importance_sampling_cap
+            )
+        if getattr(trl, "off_policy_mask_threshold", None) is not None:
+            grpo_args_kwargs["off_policy_mask_threshold"] = (
+                trl.off_policy_mask_threshold
+            )
+        if getattr(trl, "use_bias_correction_kl", None) is not None:
+            grpo_args_kwargs["use_bias_correction_kl"] = trl.use_bias_correction_kl
+
+        # Fast Async GRPO fields
+        if getattr(trl, "reward_num_workers", None) is not None:
+            grpo_args_kwargs["reward_num_workers"] = trl.reward_num_workers
+        if getattr(trl, "replay_buffer_size", None) is not None:
+            grpo_args_kwargs["replay_buffer_size"] = trl.replay_buffer_size
+        if getattr(trl, "replay_recompute_logps", None) is not None:
+            grpo_args_kwargs["replay_recompute_logps"] = trl.replay_recompute_logps
+        if getattr(trl, "reroll_start_fraction", None) is not None:
+            grpo_args_kwargs["reroll_start_fraction"] = trl.reroll_start_fraction
+        if getattr(trl, "reroll_max_groups", None) is not None:
+            grpo_args_kwargs["reroll_max_groups"] = trl.reroll_max_groups
+        if getattr(trl, "skip_zero_advantage_batches", None) is not None:
+            grpo_args_kwargs["skip_zero_advantage_batches"] = (
+                trl.skip_zero_advantage_batches
+            )
+        if getattr(trl, "vllm_lora_sync", None) is not None:
+            grpo_args_kwargs["vllm_lora_sync"] = trl.vllm_lora_sync
+
+        # Batch flattening (top-level config, not under trl)
+        if getattr(cfg, "batch_flattening", None):
+            grpo_args_kwargs["batch_flattening"] = cfg.batch_flattening
+
+        return grpo_args_kwargs
+
+    @classmethod
+    def set_trainer_args(cls, cfg: DictDefault) -> list[Any]:
+        trainer_args = []
+        if cfg.trl and cfg.trl.reward_funcs:
+            reward_funcs = []
+            for reward_func_fqn in cfg.trl.reward_funcs:
+                reward_funcs.append(cls.get_reward_func(reward_func_fqn))
+            trainer_args.append(reward_funcs)
+
+        return trainer_args
+
+    @classmethod
+    def set_trainer_kwargs(cls, cfg: DictDefault) -> dict[str, Any]:
+        trainer_kwargs = {}
+        if cfg.trl and cfg.trl.reward_processing_classes:
+            trainer_kwargs["reward_processing_classes"] = (
+                cfg.trl.reward_processing_classes
+            )
+        if cfg.trl and cfg.trl.rollout_func:
+            trainer_kwargs["rollout_func"] = cls.get_rollout_func(cfg.trl.rollout_func)
+
+        return trainer_kwargs
+
+    @classmethod
+    def get_collator(cls, *args, **kwargs):
+        # No data collation is needed in GRPO, handled by trl's trainer __init__
+        return None
+
+    @classmethod
+    def get_blocklist_args_kwargs(cls) -> list[str]:
+        return [
+            "dataset_num_proc",
+            "max_length",
+            "include_tokens_per_second",
+            "max_prompt_length",
+        ]
+
+    @classmethod
+    def get_reward_func(cls, reward_func_fqn: str) -> RewardFunc:
+        """
+        Returns the reward function from the given fully qualified name, or the path to the reward function model.
+
+        Args:
+            reward_func_fqn (str): Fully qualified name of the reward function (e.g. r1_grpo.gsm8k_transform),
+                or a HF hub path to the reward model.
+
+        Returns:
+            RewardFunc: A callable that accepts prompts and completions and returns rewards,
+                or a path to a reward model.
+
+        Raises:
+            ValueError: If the reward function does not accept at least two arguments.
+        """
+        try:
+            # use importlib to dynamically load the reward function from the module
+            reward_func_module_name = reward_func_fqn.split(".")[-1]
+            reward_func_module = importlib.import_module(
+                ".".join(reward_func_fqn.split(".")[:-1])
+            )
+            reward_func = getattr(reward_func_module, reward_func_module_name)
+            if not len(inspect.signature(reward_func).parameters) >= 2:
+                raise ValueError(
+                    "Reward function must accept at least two arguments: prompts: list and completions: list"
+                )
+            return reward_func
+        except ModuleNotFoundError as exc:
+            # the user has passed a string (ideally indicating the path of a reward model)
+            # check if it's a local dir path and not empty dir to a reward model
+            pretrained_log_msg = f"Reward function {reward_func_fqn} is a pre-trained model path - if this is unexpected, please check the reward function path."
+            if os.path.isdir(reward_func_fqn) and os.listdir(reward_func_fqn):
+                LOG.info(pretrained_log_msg)
+                return reward_func_fqn
+            try:
+                snapshot_download(reward_func_fqn, repo_type="model")
+                LOG.info(pretrained_log_msg)
+                return reward_func_fqn
+            except HTTPError:
+                raise ValueError(
+                    f"Reward function {reward_func_fqn} not found."
+                ) from exc
+
+    @classmethod
+    def get_rollout_func(cls, rollout_func_fqn: str):
+        """
+        Returns the rollout function from the given fully qualified name.
+
+        Args:
+            rollout_func_fqn (str): Fully qualified name of the rollout function
+                                    (e.g. my_module.my_rollout_func)
+
+        Returns:
+            Callable rollout function
+        """
+        try:
+            rollout_func_module_name = rollout_func_fqn.split(".")[-1]
+            rollout_func_module = importlib.import_module(
+                ".".join(rollout_func_fqn.split(".")[:-1])
+            )
+            rollout_func = getattr(rollout_func_module, rollout_func_module_name)
+
+            if not callable(rollout_func):
+                raise ValueError(
+                    f"Rollout function {rollout_func_fqn} must be callable"
+                )
+
+            return rollout_func
+
+        except ModuleNotFoundError as exc:
+            raise ValueError(f"Rollout function {rollout_func_fqn} not found.") from exc
diff --git a/src/axolotl/core/trainers/grpo/args.py b/src/axolotl/core/trainers/grpo/args.py
new file mode 100644
index 0000000000..f1dd5a6e7a
--- /dev/null
+++ b/src/axolotl/core/trainers/grpo/args.py
@@ -0,0 +1,24 @@
+"""
+Axolotl Specific Training Args
+"""
+
+from dataclasses import dataclass
+
+from trl import GRPOConfig
+
+from axolotl.core.trainers.grpo.fast_async_trainer import FastAsyncGRPOConfig
+from axolotl.core.training_args import AxolotlTrainingMixins
+
+
+@dataclass
+class AxolotlGRPOConfig(AxolotlTrainingMixins, GRPOConfig):
+    """Axolotl GRPO Config for GRPO training"""
+
+    context_parallel_size: int | None = None
+
+
+@dataclass
+class AxolotlAsyncGRPOConfig(AxolotlTrainingMixins, FastAsyncGRPOConfig):
+    """Axolotl Async GRPO Config — adds async prefetch, streaming scoring, and IS correction."""
+
+    context_parallel_size: int | None = None
diff --git a/src/axolotl/core/trainers/grpo/async_trainer.py b/src/axolotl/core/trainers/grpo/async_trainer.py
new file mode 100644
index 0000000000..79b7512537
--- /dev/null
+++ b/src/axolotl/core/trainers/grpo/async_trainer.py
@@ -0,0 +1,3257 @@
+"""
+Async GRPO training with streaming scoring and IS correction.
+
+Works on stock TRL v0.29.0 and transformers v5.3.0 — no custom branches needed.
+
+Features:
+  - Async prefetch: background thread generates completions via vLLM while the main
+    thread trains on the previous rollout.
+  - Deferred scoring: rewards, advantages, and policy logprobs computed on the main
+    thread (thread-safe with GPU forward passes).
+  - Streaming group scoring: scores prompt groups incrementally so that reward
+    computation overlaps with the next group's logprob computation.
+  - Importance sampling (IS) correction: corrects for stale vLLM weights.
+  - Off-Policy Sequence Mask (OPSM): drops sequences with high KL + negative advantage.
+  - Configurable vLLM weight sync interval.
+
+Classes exported:
+  - AsyncGRPOConfig: GRPOConfig extended with async/streaming/IS fields
+  - AsyncGRPOTrainer: GRPOTrainer with async prefetch and IS correction
+  - ProducerConfig, DataProducer, BaseDataProducer, AsyncDataProducer: data producer protocol
+"""
+
+import atexit
+import concurrent.futures
+import queue
+import threading
+from abc import ABC, abstractmethod
+from collections import deque
+from contextlib import nullcontext
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from trl.extras.profiling import profiling_decorator
+from trl.trainer import GRPOConfig, GRPOTrainer
+from trl.trainer.utils import (
+    RepeatSampler,
+    entropy_from_logits,
+    nanmax,
+    nanmin,
+    nanstd,
+    pad,
+    selective_log_softmax,
+    shuffle_sequence_dict,
+    split_pixel_values_by_grid,
+    split_tensor_dict,
+    unsplit_pixel_values_by_grid,
+)
+
+try:
+    from trl.data_utils import (
+        apply_chat_template,
+        is_conversational,
+        prepare_multimodal_messages,
+    )
+except ImportError:
+    from trl.chat_template_utils import apply_chat_template
+    from trl.data_utils import is_conversational, prepare_multimodal_messages
+
+try:
+    from trl.models.utils import disable_gradient_checkpointing
+except ImportError:
+    from contextlib import contextmanager
+
+    @contextmanager
+    def disable_gradient_checkpointing(model, kwargs):
+        yield
+
+
+try:
+    from accelerate.utils import gather_object
+except ImportError:
+    gather_object = None
+
+try:
+    from peft import PeftModel
+    from trl.trainer.utils import use_adapter
+except ImportError:
+    PeftModel = None
+    use_adapter = nullcontext
+
+try:
+    from liger_kernel.ops.grpo_loss import (
+        fused_selective_log_softmax as _fused_selective_log_softmax,
+    )
+except ImportError:
+    _fused_selective_log_softmax = None
+
+from axolotl.utils.logging import get_logger
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class AsyncGRPOConfig(GRPOConfig):
+    """GRPOConfig extended with async prefetch, streaming scoring, and IS correction fields.
+
+    Fields already present in stock GRPOConfig (e.g. ``importance_sampling_level``,
+    ``multi_objective_aggregation``) are listed here for safety: if the stock version
+    does not define them, the defaults below ensure everything works.
+    """
+
+    # --- Data producer ---
+    use_data_producer: bool = field(
+        default=False,
+        metadata={
+            "help": "Use the GRPODataProducer protocol for online data generation."
+        },
+    )
+
+    # --- Async data production ---
+    async_prefetch: bool = field(
+        default=False,
+        metadata={
+            "help": "Generate rollouts in a background thread while training on the previous rollout."
+        },
+    )
+    prefetch_depth: int = field(
+        default=1,
+        metadata={"help": "Number of rollouts to prefetch ahead of training."},
+    )
+    vllm_sync_interval: int = field(
+        default=1,
+        metadata={
+            "help": "Sync model weights to vLLM every N optimizer steps (async mode only)."
+        },
+    )
+
+    # --- Batch flattening ---
+    batch_flattening: bool = field(
+        default=False,
+        metadata={
+            "help": "Use batch flattening for the scoring forward pass. Removes padding tokens "
+            "before the forward pass, reducing attention FLOPs proportional to the padding ratio. "
+            "Requires flash_attention_2 attention implementation. Incompatible with FSDP and "
+            "multimodal models. The per-token logprob results differ by bf16 precision (~0.03 mean) "
+            "but produce equivalent loss and gradients."
+        },
+    )
+
+    # --- Streaming scoring ---
+    streaming_partial_batch: bool = field(
+        default=False,
+        metadata={
+            "help": "Score prompt groups incrementally instead of the full batch at once."
+        },
+    )
+    streaming_min_groups: int = field(
+        default=1,
+        metadata={"help": "Minimum prompt groups to score per streaming chunk."},
+    )
+
+    # --- vLLM importance sampling correction ---
+    vllm_importance_sampling_correction: bool = field(
+        default=True,
+        metadata={
+            "help": "Apply IS correction for distribution mismatch between vLLM and training model."
+        },
+    )
+    vllm_importance_sampling_mode: str = field(
+        default="token_truncate",
+        metadata={
+            "help": "IS mode: token_truncate, token_mask, sequence_truncate, or sequence_mask."
+        },
+    )
+    vllm_importance_sampling_cap: float = field(
+        default=3.0,
+        metadata={"help": "Cap C for IS ratio clipping/masking."},
+    )
+
+    # --- Off-policy sequence mask (OPSM) ---
+    off_policy_mask_threshold: float | None = field(
+        default=None,
+        metadata={"help": "KL threshold for OPSM (DeepSeek-V3.2). None = disabled."},
+    )
+
+    # --- Bias-corrected KL ---
+    use_bias_correction_kl: bool = field(
+        default=False,
+        metadata={"help": "Apply IS correction to KL divergence term."},
+    )
+
+
+# ---------------------------------------------------------------------------
+# Data Producer Protocol (standalone — no transformers branch needed)
+# ---------------------------------------------------------------------------
+
+logger = get_logger(__name__)
+_dp_logger = get_logger(__name__ + ".data_producer")
+
+
+@dataclass
+class ProducerConfig:
+    """Configuration for a :class:`DataProducer`.
+
+    Args:
+        mini_epochs: Number of training passes over each produced dataset.
+        max_rollouts: Maximum number of produce-then-train rounds (None = unlimited).
+        steps_per_generation: Optimisation steps per produced dataset before regenerating.
+        num_iterations: Number of times to reuse each generation across optimisation steps.
+        async_prefetch: Produce the next dataset in a background thread.
+        prefetch_depth: How many rollouts to queue ahead when async.
+        sync_warmup_rollouts: Initial on-policy rollouts before switching to async.
+        eval_during_produce: Switch model to eval() during produce().
+        empty_cache_before_produce: torch.cuda.empty_cache() before produce().
+        empty_cache_after_produce: torch.cuda.empty_cache() after produce().
+    """
+
+    mini_epochs: int = 1
+    max_rollouts: int | None = None
+    steps_per_generation: int | None = None
+    num_iterations: int = 1
+    async_prefetch: bool = False
+    prefetch_depth: int = 1
+    sync_warmup_rollouts: int = 0
+    eval_during_produce: bool = True
+    empty_cache_before_produce: bool = False
+    empty_cache_after_produce: bool = False
+
+    def __post_init__(self):
+        if self.mini_epochs < 1:
+            raise ValueError(f"mini_epochs must be >= 1, got {self.mini_epochs}")
+        if self.max_rollouts is not None and self.max_rollouts < 1:
+            raise ValueError(
+                f"max_rollouts must be >= 1 or None, got {self.max_rollouts}"
+            )
+        if self.num_iterations < 1:
+            raise ValueError(f"num_iterations must be >= 1, got {self.num_iterations}")
+        if self.steps_per_generation is not None and self.steps_per_generation < 1:
+            raise ValueError(
+                f"steps_per_generation must be >= 1 or None, got {self.steps_per_generation}"
+            )
+        if self.prefetch_depth < 1:
+            raise ValueError(f"prefetch_depth must be >= 1, got {self.prefetch_depth}")
+        if self.sync_warmup_rollouts < 0:
+            raise ValueError(
+                f"sync_warmup_rollouts must be >= 0, got {self.sync_warmup_rollouts}"
+            )
+
+
+class _GroupShardedSampler:
+    """Rank-aware shard of a ``RepeatSampler`` that preserves GRPO groups.
+
+    ``RepeatSampler`` yields ``num_generations`` consecutive copies of
+    each prompt, forming a GRPO group. For distributed training each
+    rank must see a disjoint slice of prompts (otherwise every rank
+    dogpiles on the first 1/world_size of the batch) while keeping each
+    group intact on a single rank so advantage normalization sees all
+    peer generations.
+
+    ``accelerator.prepare(DataLoader)`` does not handle this correctly
+    for custom samplers with ``split_batches=False`` (the default): it
+    leaves the sampler alone and every rank replays identical indices.
+    This wrapper fixes that by consuming the inner sampler's full
+    output, chunking it into ``num_generations``-sized groups, and
+    round-robining whole groups across ranks.
+
+    Intended to be used ONLY when distributed training is active
+    (``num_replicas > 1``); for single-rank it is a no-op but still
+    correct.
+    """
+
+    def __init__(
+        self,
+        inner: Any,
+        num_generations: int,
+        rank: int,
+        num_replicas: int,
+    ):
+        if num_generations < 1:
+            raise ValueError(f"num_generations must be >= 1, got {num_generations}")
+        if num_replicas < 1:
+            raise ValueError(f"num_replicas must be >= 1, got {num_replicas}")
+        if not (0 <= rank < num_replicas):
+            raise ValueError(f"rank must be in [0, {num_replicas}), got {rank}")
+        self.inner = inner
+        self.num_generations = num_generations
+        self.rank = rank
+        self.num_replicas = num_replicas
+
+    def __iter__(self):
+        all_indices = list(self.inner)
+        if len(all_indices) % self.num_generations != 0:
+            raise ValueError(
+                f"inner sampler yielded {len(all_indices)} indices, "
+                f"not a multiple of num_generations={self.num_generations}"
+            )
+        # Chunk the flat index sequence into groups of num_generations
+        # consecutive indices. ``RepeatSampler`` guarantees that each
+        # group contains num_generations copies of the same prompt id.
+        groups = [
+            all_indices[i : i + self.num_generations]
+            for i in range(0, len(all_indices), self.num_generations)
+        ]
+        # Round-robin whole groups across ranks. Round-robin (vs.
+        # contiguous chunking) preserves approximate shuffled order on
+        # each rank even when the group count is small relative to the
+        # world size.
+        for group in groups[self.rank :: self.num_replicas]:
+            yield from group
+
+    def __len__(self):
+        try:
+            inner_len = len(self.inner)
+        except TypeError:
+            # Non-sized inner sampler — we can't know the per-rank
+            # length without materializing. Return 0 as a hint that the
+            # DataLoader should fall back to iteration.
+            return 0
+        total_groups = inner_len // self.num_generations
+        # Ceiling division for the trailing groups that don't divide
+        # evenly — extra groups go to the first ``total_groups %
+        # num_replicas`` ranks, matching the round-robin above.
+        my_groups = (
+            total_groups + self.num_replicas - self.rank - 1
+        ) // self.num_replicas
+        return my_groups * self.num_generations
+
+
+class DataProducer(ABC):
+    """Abstract base class for online data producers.
+
+    Subclass this and implement :meth:`produce` to supply fresh training data
+    each rollout round.
+    """
+
+    config: ProducerConfig
+
+    @abstractmethod
+    def produce(
+        self,
+        model: Any,
+        global_step: int,
+        *,
+        processing_class: Any = None,
+        accelerator: Any = None,
+        args: Any = None,
+        **kwargs,
+    ) -> Dataset:
+        """Generate a fresh training dataset."""
+        ...
+
+
+class BaseDataProducer(DataProducer):
+    """Convenience base class with a default :class:`ProducerConfig` and lifecycle hooks."""
+
+    def __init__(self, config: ProducerConfig | None = None):
+        self.config = config or ProducerConfig()
+
+    def on_rollout_begin(self, global_step: int) -> None:
+        """Called before each produce() invocation."""
+
+    def on_rollout_end(self, dataset: Dataset, global_step: int) -> None:
+        """Called after each produce() invocation with the produced dataset."""
+
+
+class AsyncDataProducer:
+    """Wraps a synchronous :class:`DataProducer` for background-thread data generation.
+
+    While the Trainer trains on the current rollout, this wrapper produces upcoming
+    datasets in a background thread.
+
+    FSDP compatibility: Background threads must NOT call cross-rank collectives
+    (gather_object, broadcast_object_list, FSDP all-gather) because the main thread
+    may be doing FSDP forward/backward concurrently, causing deadlocks. When
+    ``num_processes > 1``, only rank 0 runs BG generation; results are broadcast
+    to other ranks on the main thread when ``produce()`` is next called.
+    """
+
+    def __init__(
+        self, inner: DataProducer, background_produce_kwargs: dict | None = None
+    ):
+        self._inner = inner
+        self._depth = inner.config.prefetch_depth
+        self._warmup_remaining = inner.config.sync_warmup_rollouts
+        self._background_kwargs = background_produce_kwargs or {}
+        self._executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="async-producer"
+        )
+        self._queue: deque[concurrent.futures.Future] = deque()
+        self._initialized = False
+        # Lock held by the background thread during vLLM generation.
+        # The main thread acquires this lock for weight sync to ensure
+        # merge_adapter/unmerge_adapter don't overlap with generation.
+        self._generate_lock = threading.Lock()
+        # Detected at first produce() call
+        self._num_processes: int | None = None
+        self._is_main: bool | None = None
+
+    @property
+    def config(self) -> ProducerConfig:
+        return self._inner.config
+
+    def produce(self, model: Any, global_step: int, **kwargs) -> Dataset:
+        """Return the next dataset, blocking if the prefetch hasn't finished."""
+        # Detect multi-process on first call
+        if self._num_processes is None:
+            accelerator = kwargs.get("accelerator")
+            if accelerator is not None:
+                self._num_processes = accelerator.num_processes
+                self._is_main = accelerator.is_main_process
+            else:
+                self._num_processes = 1
+                self._is_main = True
+
+        # During warmup, produce synchronously (on-policy)
+        if self._warmup_remaining > 0:
+            self._warmup_remaining -= 1
+            _dp_logger.info(
+                f"AsyncDataProducer: sync warmup rollout (remaining={self._warmup_remaining})"
+            )
+            return self._inner.produce(model, global_step, **kwargs)
+
+        if not self._initialized:
+            dataset = self._inner.produce(model, global_step, **kwargs)
+            bg_kwargs = {**kwargs, **self._background_kwargs}
+            # With FSDP (multi-process), only submit BG tasks on rank 0.
+            # Non-rank-0 processes will receive data via broadcast.
+            if self._num_processes > 1:
+                bg_kwargs["_rank0_only"] = True
+            for i in range(1, self._depth + 1):
+                self._queue.append(
+                    self._executor.submit(
+                        self._locked_produce, model, global_step + i, **bg_kwargs
+                    )
+                )
+            self._initialized = True
+            return dataset
+
+        # Get the pre-generated dataset from the BG thread
+        dataset = self._queue.popleft().result()
+
+        # With FSDP: BG thread only ran on rank 0. Broadcast to all ranks.
+        if self._num_processes > 1:
+            dataset = self._broadcast_dataset(dataset)
+
+        bg_kwargs = {**kwargs, **self._background_kwargs}
+        if self._num_processes > 1:
+            bg_kwargs["_rank0_only"] = True
+        next_step = global_step + self._depth
+        self._queue.append(
+            self._executor.submit(self._locked_produce, model, next_step, **bg_kwargs)
+        )
+        return dataset
+
+    def _broadcast_dataset(self, dataset) -> Dataset:
+        """Broadcast a prefetched dataset from rank 0 to all ranks (main thread).
+
+        Rank 0 has a full RolloutDataset from BG generation; other ranks have None.
+        After broadcast, tensors are moved to each rank's local device.
+        """
+        import torch.distributed as dist
+
+        if not dist.is_initialized():
+            return dataset
+
+        # Rank 0 sends _data dict; others receive it
+        obj_list = [dataset._data if self._is_main else None]
+        dist.broadcast_object_list(obj_list, src=0)
+
+        data: dict[str, Any] = obj_list[0]  # type: ignore[assignment]
+
+        # Move tensors to local device (broadcast_object_list deserializes to CPU)
+        accelerator = self._inner._trainer.accelerator  # type: ignore[attr-defined]
+        device = accelerator.device
+        for key, val in data.items():
+            if isinstance(val, torch.Tensor) and val.device != device:
+                data[key] = val.to(device)
+
+        if not self._is_main:
+            from axolotl.core.trainers.grpo.async_trainer import RolloutDataset
+
+            dataset = RolloutDataset(data)
+        else:
+            # Rank 0 already has the dataset, but update _data with device-moved tensors
+            dataset._data = data
+        return dataset
+
+    def _locked_produce(self, model: Any, global_step: int, **kwargs) -> Dataset:
+        """Run produce while holding the generate lock."""
+        with self._generate_lock:
+            return self._inner.produce(model, global_step, **kwargs)
+
+    def on_rollout_begin(self, global_step: int) -> None:
+        if hasattr(self._inner, "on_rollout_begin"):
+            self._inner.on_rollout_begin(global_step)
+
+    def on_rollout_end(self, dataset: Dataset, global_step: int) -> None:
+        if hasattr(self._inner, "on_rollout_end"):
+            self._inner.on_rollout_end(dataset, global_step)
+
+    def shutdown(self) -> None:
+        """Shut down the background thread pool and cancel pending futures."""
+        for future in self._queue:
+            future.cancel()
+        self._queue.clear()
+        self._executor.shutdown(wait=False)
+
+
+class DataProducerCallback:
+    """Marker class: if a DataProducer also inherits from this, the Trainer will
+    automatically register it as a callback."""
+
+    pass
+
+
+# ---------------------------------------------------------------------------
+# RolloutDataset + GRPODataProducer
+# ---------------------------------------------------------------------------
+
+
+class RolloutDataset(Dataset):
+    """A Dataset wrapping the output dict from _generate_and_score_completions.
+
+    Per-sample tensors are sliced by index; shared metadata is passed through.
+    """
+
+    _ALWAYS_SHARED = frozenset(
+        {"num_items_in_batch", "_pending_policy_logps", "_rank0_only"}
+    )
+
+    def __init__(self, data: dict[str, Any]):
+        self._data = data
+        self._shared_keys: set[str] = set()
+        self._sample_keys: set[str] = set()
+
+        for key, val in data.items():
+            if key in self._ALWAYS_SHARED:
+                self._shared_keys.add(key)
+            elif not isinstance(val, torch.Tensor):
+                self._shared_keys.add(key)
+            elif val.dim() == 0:
+                self._shared_keys.add(key)
+            else:
+                self._sample_keys.add(key)
+
+        self._num_samples = 0
+        for key in self._sample_keys:
+            n = data[key].size(0)
+            if self._num_samples == 0:
+                self._num_samples = n
+            elif n != self._num_samples:
+                raise ValueError(
+                    f"Inconsistent sample count: key '{key}' has {n}, expected {self._num_samples}"
+                )
+        if self._num_samples == 0:
+            raise ValueError("No per-sample tensors found in rollout data")
+
+    def __len__(self) -> int:
+        return self._num_samples
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        item: dict[str, Any] = {}
+        for key in self._sample_keys:
+            item[key] = self._data[key][idx]
+        for key in self._shared_keys:
+            item[key] = self._data[key]
+        return item
+
+
+def make_rollout_collator(shared_keys: set[str]):
+    """Return a collator that stacks per-sample tensors and passes shared keys through."""
+
+    def _collate(batch: list[dict[str, Any]]) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        for key in batch[0]:
+            if key in shared_keys:
+                result[key] = batch[0][key]
+            else:
+                values = [item[key] for item in batch]
+                if isinstance(values[0], torch.Tensor):
+                    result[key] = torch.stack(values)
+                else:
+                    result[key] = values
+        return result
+
+    return _collate
+
+
+class GRPODataProducer(BaseDataProducer):
+    """Produces GRPO training rollouts using the trainer's generation pipeline.
+
+    Created before Trainer.__init__ completes; the trainer reference is injected
+    later via set_trainer().
+    """
+
+    def __init__(
+        self,
+        config: ProducerConfig,
+        prompt_dataset,
+        *,
+        num_generations: int,
+        generation_batch_size: int,
+        train_batch_size: int,
+        steps_per_generation: int,
+        shuffle_dataset: bool,
+        seed: int,
+    ):
+        super().__init__(config)
+        self._dataset = prompt_dataset
+        self._num_generations = num_generations
+        self._generation_batch_size = generation_batch_size
+        self._train_batch_size = train_batch_size
+        self._steps_per_generation = steps_per_generation
+        self._shuffle_dataset = shuffle_dataset
+        self._seed = seed
+        self._trainer: Any = None
+        self._prompt_dl: Any = None
+        self._prompt_iter: Any = None
+
+    def set_trainer(self, trainer) -> None:
+        """Inject the live trainer reference and create the prompt DataLoader."""
+        self._trainer = trainer
+        # Defer _init_prompt_dataloader if trainer.args is not yet set
+        # (happens when set_trainer is called from _create_data_producer during __init__)
+        if getattr(trainer, "args", None) is not None:
+            self._init_prompt_dataloader()
+
+    def _init_prompt_dataloader(self) -> None:
+        from functools import partial
+
+        from transformers.trainer_utils import seed_worker
+
+        trainer = self._trainer
+        sampler = RepeatSampler(
+            data_source=self._dataset,
+            mini_repeat_count=self._num_generations,
+            batch_size=self._generation_batch_size // self._num_generations,
+            repeat_count=1,
+            shuffle=self._shuffle_dataset,
+            seed=self._seed,
+        )
+
+        # Shard the sampler across distributed ranks so each rank sees
+        # a disjoint slice of prompts. ``RepeatSampler`` groups each
+        # prompt with ``num_generations`` consecutive copies — our
+        # wrapper round-robins WHOLE groups across ranks so all
+        # generations of a given prompt stay on the same rank (needed
+        # for GRPO advantage normalization within a group).
+        #
+        # Without this, ``accelerator.prepare(dl)`` with the default
+        # ``split_batches=False`` leaves the custom sampler alone, so
+        # every rank iterates the identical index sequence and the
+        # cluster dogpiles on the first 1/world_size of the prompts.
+        num_replicas = max(1, trainer.accelerator.num_processes)
+        if num_replicas > 1:
+            sampler = _GroupShardedSampler(
+                inner=sampler,
+                num_generations=self._num_generations,
+                rank=trainer.accelerator.process_index,
+                num_replicas=num_replicas,
+            )
+            logger.info(
+                "[RANK:%d] _GroupShardedSampler active "
+                "(num_replicas=%d, num_generations=%d, gen_batch=%d)",
+                trainer.accelerator.process_index,
+                num_replicas,
+                self._num_generations,
+                self._generation_batch_size,
+            )
+
+        # Use identity collator (same as stock GRPOTrainer)
+        def _identity(x):
+            return x
+
+        dl = DataLoader(
+            self._dataset,
+            batch_size=self._train_batch_size * self._steps_per_generation,
+            sampler=sampler,
+            collate_fn=_identity,
+            num_workers=trainer.args.dataloader_num_workers,
+            pin_memory=trainer.args.dataloader_pin_memory,
+            persistent_workers=trainer.args.dataloader_persistent_workers,
+            worker_init_fn=partial(
+                seed_worker,
+                num_workers=trainer.args.dataloader_num_workers,
+                rank=trainer.args.process_index,
+            ),
+        )
+        # Skip accelerator.prepare — we're handling per-rank sharding
+        # ourselves via ``_GroupShardedSampler``. ``prepare()`` would
+        # otherwise try to wrap the DataLoader with its own sharding
+        # logic which does not understand our group structure.
+        self._prompt_dl = dl
+
+        self._prompt_iter = iter(self._prompt_dl)
+
+    def produce(
+        self,
+        model: Any,
+        global_step: int,
+        *,
+        skip_policy_logps: bool = False,
+        processing_class: Any = None,
+        accelerator: Any = None,
+        args: Any = None,
+        _rank0_only: bool = False,
+        **kwargs,
+    ) -> RolloutDataset | None:
+        """Generate a fresh GRPO training rollout."""
+        # Lazy init: create prompt DataLoader if deferred from set_trainer
+        if self._prompt_dl is None and self._trainer is not None:
+            self._init_prompt_dataloader()
+
+        is_main = self._trainer.accelerator.is_main_process
+
+        # FSDP rank0-only mode: non-rank-0 returns None (broadcast fills it later)
+        if _rank0_only and not is_main:
+            return None
+
+        try:
+            inputs = next(self._prompt_iter)
+        except StopIteration:
+            self._prompt_iter = iter(self._prompt_dl)
+            inputs = next(self._prompt_iter)
+
+        if skip_policy_logps:
+            # Async path: use _generate_only (generation without scoring) which
+            # works on stock TRL (no skip_policy_logps parameter needed).
+            output = self._trainer._generate_only(inputs, rank0_only=_rank0_only)
+        else:
+            # Sync path: full generation + scoring
+            output = self._trainer._generate_and_score_completions(inputs)
+
+            # Strip non-sequence metadata before shuffling
+            metadata = {}
+            for key in list(output.keys()):
+                val = output[key]
+                if not isinstance(val, (torch.Tensor, list)):
+                    metadata[key] = output.pop(key)
+                elif isinstance(val, torch.Tensor) and val.dim() == 0:
+                    metadata[key] = output.pop(key)
+
+            output = shuffle_sequence_dict(output)
+            output.update(metadata)
+
+        return RolloutDataset(output)
+
+
+# ---------------------------------------------------------------------------
+# Trainer
+# ---------------------------------------------------------------------------
+
+
+class AsyncGRPOTrainer(GRPOTrainer):
+    """GRPOTrainer with async prefetch, streaming scoring, and IS correction.
+
+    Drop-in replacement: pass ``AsyncGRPOConfig`` as ``args`` and use this trainer
+    instead of ``GRPOTrainer``.
+    """
+
+    def __init__(self, *args, **kwargs):
+        # Skip NCCL communicator init when using LoRA sync (filesystem) or HTTP-only
+        # merged weight sync. NCCL is only needed for the standard update_named_param
+        # path which broadcasts tensors through the communicator.
+        training_args = kwargs.get("args") or (args[1] if len(args) > 1 else None)
+        _skip_nccl = False
+        if training_args is not None:
+            if getattr(training_args, "vllm_lora_sync", False):
+                _skip_nccl = True  # LoRA sync uses filesystem + HTTP
+            elif getattr(training_args, "async_prefetch", False):
+                # Skip NCCL at init to avoid DDP param count mismatch in multi-GPU.
+                # init_communicator allocates device tensors on rank 0 only, which
+                # causes DDP to see different param counts across ranks.
+                # The communicator is initialized lazily on first weight sync instead.
+                _skip_nccl = True
+        if _skip_nccl:
+            from trl.generation.vllm_generation import VLLMGeneration
+
+            _orig_init_vllm = VLLMGeneration._init_vllm
+
+            def _init_vllm_no_communicator(self_vllm):
+                """Init vLLM client without NCCL communicator (LoRA sync uses filesystem)."""
+                if self_vllm.mode == "server" and self_vllm.accelerator.is_main_process:
+                    from trl.generation.vllm_client import VLLMClient
+
+                    if self_vllm.server_base_url is not None:
+                        base_url = self_vllm.server_base_url
+                    else:
+                        base_url = (
+                            f"http://{self_vllm.server_host}:{self_vllm.server_port}"
+                        )
+                    self_vllm.vllm_client = VLLMClient(
+                        base_url=base_url,
+                        group_port=self_vllm.group_port,
+                        connection_timeout=self_vllm.server_timeout,
+                    )
+                    # Deliberately skip init_communicator — no NCCL needed
+                elif self_vllm.mode != "server":
+                    _orig_init_vllm(self_vllm)
+
+            VLLMGeneration._init_vllm = _init_vllm_no_communicator
+
+        try:
+            super().__init__(*args, **kwargs)
+        finally:
+            # Restore original _init_vllm so other trainers aren't affected
+            if _skip_nccl:
+                VLLMGeneration._init_vllm = _orig_init_vllm  # type: ignore[possibly-undefined]
+
+        # FP8 models: zero out the pad token embedding so that padding
+        # positions have zero hidden states throughout the network.
+        # FP8 linear layers produce NaN on non-zero inputs at masked
+        # positions (the Triton fp8 matmul kernel can't handle the
+        # extreme values that accumulate at unattended positions).
+        self._zero_pad_embedding_for_fp8()
+
+        # Ensure custom attributes exist (stock GRPOTrainer.__init__ may not set them).
+        for attr, cfg_key, default in [
+            (
+                "vllm_importance_sampling_correction",
+                "vllm_importance_sampling_correction",
+                True,
+            ),
+            (
+                "vllm_importance_sampling_mode",
+                "vllm_importance_sampling_mode",
+                "token_truncate",
+            ),
+            ("vllm_importance_sampling_cap", "vllm_importance_sampling_cap", 3.0),
+            ("off_policy_mask_threshold", "off_policy_mask_threshold", None),
+        ]:
+            if not hasattr(self, attr):
+                setattr(self, attr, getattr(self.args, cfg_key, default))
+
+        # Async state
+        self._async_queue: queue.Queue | None = None
+        self._executor: concurrent.futures.ThreadPoolExecutor | None = None
+        self._prompt_iter = None
+        self._last_synced_step = -1
+        self._buffered_inputs: list | None = None  # override stock attr
+
+        # Data producer (the proper architecture for async generation)
+        self.data_producer = None
+        if getattr(self.args, "use_data_producer", False):
+            self.data_producer = self._create_data_producer(
+                kwargs["args"], kwargs["train_dataset"]
+            )
+
+        if self.args.async_prefetch and self.data_producer is None:
+            # Legacy path: direct _prepare_inputs override without data producer
+            self._setup_async()
+
+    def _create_data_producer(self, args, train_dataset):
+        """Create and return the GRPODataProducer (possibly wrapped in AsyncDataProducer)."""
+        producer_config = ProducerConfig(
+            mini_epochs=args.num_iterations,
+            max_rollouts=None,
+            eval_during_produce=False,
+            empty_cache_before_produce=True,
+            empty_cache_after_produce=True,
+            async_prefetch=args.async_prefetch,
+            prefetch_depth=args.prefetch_depth,
+        )
+        data_producer = GRPODataProducer(
+            config=producer_config,
+            prompt_dataset=train_dataset,
+            num_generations=self.num_generations,
+            generation_batch_size=args.generation_batch_size,
+            train_batch_size=args.per_device_train_batch_size,
+            steps_per_generation=args.steps_per_generation,
+            shuffle_dataset=getattr(self, "shuffle_dataset", True),
+            seed=args.seed,
+        )
+        data_producer.set_trainer(self)
+
+        if args.async_prefetch:
+            data_producer = AsyncDataProducer(
+                data_producer,
+                background_produce_kwargs={"skip_policy_logps": True},
+            )
+        return data_producer
+
+    # ------------------------------------------------------------------
+    # Async setup / teardown
+    # ------------------------------------------------------------------
+
+    def _setup_async(self):
+        """Create background thread pool, prompt iterator, and pre-fill the async queue."""
+        gen_batch_size = getattr(
+            self.args,
+            "generation_batch_size",
+            self._train_batch_size * self.args.gradient_accumulation_steps,
+        )
+        # RepeatSampler groups prompts with num_generations repetitions each.
+        # DataLoader batches the yielded indices into generation-sized batches.
+        sampler = RepeatSampler(
+            data_source=self.train_dataset,
+            mini_repeat_count=self.num_generations,
+            batch_size=gen_batch_size // self.num_generations,
+            repeat_count=10_000,  # effectively infinite
+            shuffle=True,
+            seed=self.args.seed,
+        )
+        self._prompt_dataloader = DataLoader(
+            self.train_dataset,
+            batch_size=gen_batch_size,
+            sampler=sampler,
+            collate_fn=self.data_collator,
+            num_workers=0,
+        )
+        self._prompt_iter = iter(self._prompt_dataloader)
+        self._async_queue = queue.Queue(maxsize=self.args.prefetch_depth)
+        self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+
+        # Pre-submit generations to fill the queue
+        for _ in range(self.args.prefetch_depth):
+            self._submit_generation()
+
+        atexit.register(self._shutdown_async)
+
+    def _shutdown_async(self):
+        if self._executor is not None:
+            self._executor.shutdown(wait=False, cancel_futures=True)
+            self._executor = None
+
+    def _submit_generation(self):
+        """Submit the next background generation job.
+
+        With multi-process (DDP/FSDP), only rank 0 generates to avoid
+        cross-rank NCCL collectives from background threads.  Non-rank-0
+        processes enqueue a sentinel ``None`` that is replaced by a
+        broadcast in ``_prepare_inputs_legacy_async``.
+        """
+        rank0_only = self.accelerator.num_processes > 1
+        if rank0_only and not self.accelerator.is_main_process:
+            # Non-rank-0: nothing to generate; enqueue a resolved None future
+            f: concurrent.futures.Future = concurrent.futures.Future()
+            f.set_result(None)
+            self._async_queue.put(f)
+            return
+        batch = next(self._prompt_iter)
+        future = self._executor.submit(self._generate_only, batch, rank0_only)
+        self._async_queue.put(future)
+
+    # ------------------------------------------------------------------
+    # Broadcast rollout (legacy async, multi-process)
+    # ------------------------------------------------------------------
+
+    def _broadcast_rollout(self, rollout: dict | None) -> dict:
+        """Broadcast a rank0-only rollout dict to all ranks (main thread).
+
+        Rank 0 has the full rollout dict from ``_generate_only``; other ranks
+        have ``None``.  After broadcast, tensors are moved to each rank's
+        local device.
+        """
+        import torch.distributed as dist
+
+        obj_list = [rollout if self.accelerator.is_main_process else None]
+        dist.broadcast_object_list(obj_list, src=0)
+        rollout = obj_list[0]
+        assert rollout is not None, "broadcast_object_list failed to deliver rollout"
+
+        # Move tensors to local device (broadcast deserializes to CPU)
+        device = self.accelerator.device
+        for key, val in rollout.items():
+            if isinstance(val, torch.Tensor) and val.device != device:
+                rollout[key] = val.to(device)
+
+        return rollout
+
+    # ------------------------------------------------------------------
+    # Weight sync
+    # ------------------------------------------------------------------
+
+    def _sync_peft_weights_no_merge(self):
+        """Thread-safe weight sync: compute merged LoRA weights without in-place modification.
+
+        Required for FP8 models where merge_adapter() fails (addmm not implemented
+        for Float8), and also safe for concurrent use since it never modifies base
+        weights in-place.
+        """
+        accelerator = self.vllm_generation.accelerator
+        if not (self.vllm_generation.mode == "server" and accelerator.is_main_process):
+            return
+
+        # In multi-GPU async mode, we skip NCCL communicator init to avoid
+        # DDP param count mismatch and NCCL device conflicts. Weight sync
+        # uses the HTTP-only fallback in batch_update_named_params instead.
+
+        model = self.vllm_generation.model
+        vllm_client = self.vllm_generation.vllm_client
+        fix_name = self.vllm_generation._fix_param_name_to_vllm
+
+        # Build lookup: module_path -> (A, B, scaling) for all active LoRA layers
+        lora_info = {}
+        for mod_name, module in model.base_model.model.named_modules():
+            if not hasattr(module, "lora_A") or not hasattr(module, "active_adapters"):
+                continue
+            active = module.active_adapters[0]
+            if active not in module.lora_A:
+                continue
+            lora_info[mod_name] = (
+                module.lora_A[active].weight.data,
+                module.lora_B[active].weight.data,
+                module.scaling[active],
+            )
+
+        # Build lookup for FP8 scale_inv parameters (needed for dequantization)
+        scale_inv_lookup = {}
+        for pname, pparam in model.named_parameters():
+            if "weight_scale_inv" in pname:
+                # Map weight name -> scale_inv tensor
+                weight_name = pname.replace(".weight_scale_inv", ".weight")
+                scale_inv_lookup[weight_name] = pparam.data
+
+        # Only sync parameters that have LoRA modifications — skip unchanged
+        # base weights to avoid OOM on the vLLM GPU from allocating the entire
+        # model's worth of NCCL receive buffers.
+        params_to_sync = []
+        compute_dtype = torch.bfloat16
+        for name, param in model.named_parameters():
+            vllm_name = name.removeprefix("base_model.model.").replace(
+                ".base_layer", ""
+            )
+            if model.prefix in vllm_name:
+                continue
+            if "original_module" in vllm_name:
+                continue
+            if "weight_scale_inv" in vllm_name or "input_scale" in vllm_name:
+                continue
+            if not vllm_name.endswith(".weight"):
+                continue
+            # fix_name strips modules_to_save.default. prefix
+            raw_mod_path = vllm_name[: -len(".weight")]
+            vllm_name = fix_name(vllm_name, extra_prefixes=["modules_to_save.default."])
+            mod_path = vllm_name[: -len(".weight")]
+
+            # Sync weights that have LoRA adapters OR are modules_to_save
+            is_lora = mod_path in lora_info
+            is_modules_to_save = raw_mod_path != mod_path  # fix_name stripped a prefix
+            if not is_lora and not is_modules_to_save:
+                continue
+
+            data = param.data
+
+            # Dequantize FP8 weights before merging
+            if data.dtype == torch.float8_e4m3fn and name in scale_inv_lookup:
+                scale_inv = scale_inv_lookup[name]
+                fp8_bf16 = data.to(compute_dtype)
+                if scale_inv.dim() == 2 and fp8_bf16.dim() == 2:
+                    sr, sc = scale_inv.shape
+                    br = fp8_bf16.shape[0] // sr
+                    bc = fp8_bf16.shape[1] // sc
+                    data = (
+                        fp8_bf16.reshape(sr, br, sc, bc)
+                        * scale_inv[:, None, :, None].to(compute_dtype)
+                    ).reshape(fp8_bf16.shape)
+                elif scale_inv.dim() <= 1:
+                    data = fp8_bf16 * scale_inv.to(compute_dtype)
+                else:
+                    data = fp8_bf16
+            elif data.dtype == torch.float8_e4m3fn:
+                data = data.to(compute_dtype)
+
+            if is_lora:
+                A, B, s = lora_info[mod_path]
+                merged = data.to(compute_dtype) + s * (
+                    B.to(compute_dtype) @ A.to(compute_dtype)
+                )
+                params_to_sync.append((vllm_name, merged))
+            else:
+                # modules_to_save: send raw weight (no LoRA merge needed)
+                params_to_sync.append((vllm_name, data.to(compute_dtype)))
+
+        # Batch sync only LoRA-modified params via HTTP+NCCL
+        if params_to_sync:
+            sync_mb = sum(t.numel() * t.element_size() for _, t in params_to_sync) / 1e6
+            logger.info(
+                f"Syncing {len(params_to_sync)} LoRA-modified params ({sync_mb:.0f} MB)"
+            )
+            vllm_client.batch_update_named_params(params_to_sync)
+
+        # Reset prefix cache after weight update
+        vllm_client.reset_prefix_cache()
+
+    def _sync_lora_adapter(self):
+        """Sync LoRA adapter to vLLM via filesystem (native LoRA mode).
+
+        Saves the PEFT adapter to a temp directory and POSTs the path to vLLM's
+        /set_lora_adapter/ endpoint. vLLM loads the adapter natively using Punica
+        kernels, avoiding the need to merge weights and NCCL-broadcast the full model.
+
+        Syncs only the LoRA adapter weights via filesystem instead of the full merged model via NCCL.
+
+        FSDP/DeepSpeed: All ranks must participate in the state_dict gather.
+        accelerator.get_state_dict() handles this (FSDP uses FullStateDictConfig
+        with rank0_only=True). Only rank 0 gets the full dict, writes files, and
+        does the HTTP POST.
+        """
+        import os
+        import tempfile
+
+        accelerator = self.vllm_generation.accelerator
+        model = self.vllm_generation.model
+
+        if self.vllm_generation.mode != "server":
+            return
+
+        is_main = accelerator.is_main_process
+
+        # Increment adapter version (all ranks, kept in sync)
+        if not hasattr(self, "_lora_sync_version"):
+            self._lora_sync_version = 0
+            if is_main:
+                self._lora_sync_dir = tempfile.mkdtemp(prefix="lora_sync_")
+            else:
+                self._lora_sync_dir = None
+            # Broadcast sync dir from rank 0 to all ranks
+            if accelerator.num_processes > 1:
+                import torch.distributed as dist
+
+                if dist.is_initialized():
+                    obj_list = [self._lora_sync_dir]
+                    dist.broadcast_object_list(obj_list, src=0)
+                    self._lora_sync_dir = obj_list[0]
+        self._lora_sync_version += 1
+
+        adapter_path = os.path.join(self._lora_sync_dir, f"v{self._lora_sync_version}")
+
+        # Gather state dict from all ranks (FSDP/DeepSpeed gather, rank0_only)
+        # All ranks must participate even though only rank 0 gets the result.
+        # Use self.model_wrapped (the DeepSpeed/FSDP engine) for get_state_dict,
+        # since it has the necessary hooks (e.g. zero_gather_16bit_weights_on_model_save).
+        # self.vllm_generation.model is the unwrapped PEFT model which lacks these.
+        wrapped_model = getattr(self, "model_wrapped", model)
+        state_dict = accelerator.get_state_dict(wrapped_model)
+
+        if is_main:
+            # Unwrap to access PEFT's save_pretrained
+            unwrapped = accelerator.unwrap_model(model)
+            unwrapped.save_pretrained(adapter_path, state_dict=state_dict)
+
+            import requests
+
+            vllm_client = self.vllm_generation.vllm_client
+            base_url = vllm_client.base_url
+            base_model = getattr(self.args, "model_name_or_path", "axolotl-lora")
+            sync_timeout = getattr(self.args, "vllm_server_timeout", 300) or 300
+
+            # Try standard vLLM /v1/load_lora_adapter first, fall back to custom endpoint
+            response = requests.post(
+                f"{base_url}/v1/load_lora_adapter",
+                json={
+                    "lora_name": base_model,
+                    "lora_path": adapter_path,
+                    "load_inplace": True,
+                },
+                timeout=sync_timeout,
+            )
+            if response.status_code != 200:
+                # Fallback: try custom /set_lora_adapter/ endpoint
+                response = requests.post(
+                    f"{base_url}/set_lora_adapter/",
+                    json={
+                        "lora_name": "active_lora",
+                        "lora_int_id": self._lora_sync_version,
+                        "lora_path": adapter_path,
+                    },
+                    timeout=30,
+                )
+                if response.status_code != 200:
+                    logger.warning(
+                        "Failed to set LoRA adapter: %s %s",
+                        response.status_code,
+                        response.text,
+                    )
+                    return
+
+            # Reset prefix cache after adapter update
+            try:
+                vllm_client.reset_prefix_cache()
+            except Exception as exc:
+                logger.warning("Failed to reset prefix cache: %s", exc)
+
+            # Clean up old adapter versions (keep only current)
+            if self._lora_sync_version > 1:
+                old_path = os.path.join(
+                    self._lora_sync_dir, f"v{self._lora_sync_version - 1}"
+                )
+                if os.path.exists(old_path):
+                    import shutil
+
+                    shutil.rmtree(old_path, ignore_errors=True)
+
+            logger.info(
+                "Synced LoRA adapter v%d to vLLM (%s)",
+                self._lora_sync_version,
+                adapter_path,
+            )
+
+        # Barrier to ensure all ranks complete before resuming forward passes.
+        # Without this, rank 1 may start a forward pass (triggering FSDP unshard)
+        # while rank 0 is still doing save_pretrained, causing FSDP all-gather deadlock.
+        if accelerator.num_processes > 1:
+            import torch.distributed as dist
+
+            if dist.is_initialized():
+                dist.barrier()
+
+    def _maybe_sync_vllm_weights(self):
+        """Sync model weights to vLLM if the interval has elapsed.
+
+        Dispatches to one of three strategies:
+        - vllm_lora_sync: saves adapter to filesystem, vLLM loads natively
+        - PEFT no-merge: computes merged weights as new tensors, NCCL broadcast
+        - Non-PEFT: stock sync_weights via merge_adapter + NCCL
+
+        This is the canonical sync trigger and runs in BOTH async and
+        synchronous modes from ``_prepare_inputs_with_data_producer`` /
+        ``_prepare_inputs_legacy_async``. The ``_generate_single_turn``
+        patch is a parallel backup for non-data-producer paths (vanilla
+        GRPO without NeMo Gym), where the data producer is bypassed
+        entirely and TRL's stock generate-then-sync flow is used instead.
+        """
+        if not self.use_vllm:
+            return
+        step = self.state.global_step
+        # Default to syncing every step when no interval is configured —
+        # otherwise ``step % None`` would TypeError, and the previous
+        # behavior of crashing on the first sync was strictly worse than
+        # the standard "sync every optimizer step".
+        interval = self.args.vllm_sync_interval or 1
+        if step != self._last_synced_step and step % interval == 0:
+            if step == 0:
+                logger.info("Skipping vLLM weight sync at step 0 (no training yet)")
+                self._last_synced_step = step
+                return
+            if getattr(self.args, "vllm_lora_sync", False):
+                # Native LoRA sync: save adapter to filesystem, vLLM loads it directly
+                self._sync_lora_adapter()
+            else:
+                from accelerate.utils import is_peft_model
+
+                use_no_merge = is_peft_model(self.vllm_generation.model)
+
+                if use_no_merge:
+                    # No-merge sync: computes merged weights as new tensors
+                    # (doesn't modify base weights in-place), so it's safe to
+                    # run concurrently with BG generation — no lock needed.
+                    self._sync_peft_weights_no_merge()
+                else:
+                    # Non-PEFT: use stock sync (acquires lock to avoid overlap)
+                    if self.data_producer is not None and hasattr(
+                        self.data_producer, "_generate_lock"
+                    ):
+                        with self.data_producer._generate_lock:
+                            self.vllm_generation.sync_weights()
+                    elif self._async_queue is not None:
+                        pending = list(self._async_queue.queue)
+                        for f in pending:
+                            if isinstance(f, concurrent.futures.Future):
+                                f.result()
+                        self.vllm_generation.sync_weights()
+                    else:
+                        self.vllm_generation.sync_weights()
+            self._last_synced_step = step
+
+    def _zero_pad_embedding_for_fp8(self):
+        """Zero out the pad token embedding for FP8 models.
+
+        FP8 linear layers produce NaN when processing positions with
+        attention_mask=0 (the hidden states at those positions have
+        unconstrained values that overflow FP8 range during
+        quantization). By setting the pad token embedding to zeros,
+        padding positions start with zero hidden states and stay zero
+        through masked attention, preventing NaN from FP8 matmul.
+        """
+        model = self.accelerator.unwrap_model(self.model)
+        # Check if model has FP8 weights
+        has_fp8 = any(
+            p.dtype == torch.float8_e4m3fn
+            for p in model.parameters()
+            if not p.requires_grad
+        )
+        if not has_fp8:
+            return
+
+        # Find the embedding layer
+        if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
+            embed = model.model.embed_tokens
+        elif hasattr(model, "base_model") and hasattr(model.base_model, "model"):
+            m = model.base_model.model
+            if hasattr(m, "model") and hasattr(m.model, "embed_tokens"):
+                embed = m.model.embed_tokens
+            else:
+                return
+        else:
+            return
+
+        pad_id = self.processing_class.pad_token_id
+        if pad_id is not None and pad_id < embed.weight.shape[0]:
+            with torch.no_grad():
+                embed.weight.data[pad_id].zero_()
+            get_logger("async_grpo").info(
+                f"Zeroed pad token embedding (id={pad_id}) for FP8 NaN prevention"
+            )
+
+    # ------------------------------------------------------------------
+    # Background-thread generation (no scoring)
+    # ------------------------------------------------------------------
+
+    def _generate_single_turn(self, prompts, *args, **kwargs):
+        """Override to prevent weight sync from background thread and to use
+        no-merge sync for PEFT models (FP8 models can't merge_adapter)."""
+        is_bg = threading.current_thread() is not threading.main_thread()
+        saved_step = None
+
+        if is_bg and self.use_vllm:
+            # Trick: match _last_loaded_step so the stock sync check is a no-op
+            saved_step = getattr(self, "_last_loaded_step", None)
+            self._last_loaded_step = self.state.global_step
+
+        # Permanently replace vllm_generation.sync_weights with our custom
+        # sync to avoid merge_adapter (fails on FP8 / races with training).
+        #
+        # The design has two modes that have to be threaded carefully:
+        #
+        #   - Async prefetch ON: BG generation thread can't safely call
+        #     sync_weights mid-rollout (it races with the trainer's optimizer
+        #     step and can corrupt weights). We no-op the stock sync hook and
+        #     drive sync ourselves from ``_maybe_sync_vllm_weights`` after the
+        #     optimizer step on the main thread.
+        #
+        #   - Async prefetch OFF (synchronous mode): TRL's stock
+        #     ``_generate_single_turn`` calls ``sync_weights`` once per step
+        #     boundary. There's no BG thread to race with, and
+        #     ``_maybe_sync_vllm_weights`` short-circuits with
+        #     ``if not async_prefetch: return``, so we MUST wire the stock
+        #     hook directly to our LoRA sync helper — otherwise nothing ever
+        #     pushes weights to vLLM and the trainer becomes a no-op (vLLM
+        #     keeps serving the base model, every rollout in every group
+        #     produces identical outputs, advantages are zero, optimizer
+        #     step gets skipped, repeat).
+        if not getattr(self, "_patched_sync_weights", False):
+            if self.use_vllm and hasattr(self, "vllm_generation"):
+                if getattr(self.args, "vllm_lora_sync", False):
+                    if getattr(self.args, "async_prefetch", False):
+                        # Async: drive sync from main thread via
+                        # _maybe_sync_vllm_weights instead.
+                        self.vllm_generation.sync_weights = lambda: None
+                    else:
+                        # Sync mode: TRL's _generate_single_turn already
+                        # calls sync_weights once per step boundary. Wire
+                        # it directly to our LoRA filesystem sync helper.
+                        sync_helper = self._sync_lora_adapter
+
+                        def _lora_filesystem_sync():
+                            sync_helper()
+
+                        self.vllm_generation.sync_weights = _lora_filesystem_sync
+                    self._patched_sync_weights = True
+                else:
+                    from accelerate.utils import is_peft_model
+
+                    if is_peft_model(self.vllm_generation.model):
+
+                        def _no_merge_sync():
+                            self._sync_peft_weights_no_merge()
+
+                        self.vllm_generation.sync_weights = _no_merge_sync
+                        self._patched_sync_weights = True
+
+        try:
+            return super()._generate_single_turn(prompts, *args, **kwargs)
+        finally:
+            if saved_step is not None:
+                self._last_loaded_step = saved_step
+
+    def _generate_rank0_only(self, prompts):
+        """Generate using vLLM directly on rank 0 without cross-rank collectives.
+
+        Called from BG thread in FSDP mode. Bypasses ``gather_object`` /
+        ``broadcast_object_list`` since the main thread may be running FSDP
+        collectives concurrently.
+
+        Returns the same tuple as ``_generate``.
+        """
+        import copy
+
+        prompts = copy.deepcopy(prompts)
+
+        # Duplicate prompts for num_generations (same as TRL's gather+unique pattern)
+        num_generations = self.num_generations
+        unique_prompts = prompts[::num_generations]
+
+        # Build sampling params
+        vg = self.vllm_generation
+        sampling_params = {
+            "n": num_generations,
+            "repetition_penalty": vg.repetition_penalty,
+            "temperature": vg.temperature,
+            "top_p": vg.top_p,
+            "top_k": vg.top_k,
+            "min_p": 0.0 if vg.min_p is None else vg.min_p,
+            "max_tokens": vg.max_completion_length,
+            "logprobs": vg.logprobs,
+            "structured_outputs_regex": vg.structured_outputs_regex,
+            "generation_kwargs": vg.generation_kwargs,
+        }
+
+        # Call vLLM directly (no collectives)
+        from trl.data_utils import is_conversational
+
+        if is_conversational({"prompt": unique_prompts[0]}):
+            output = vg.vllm_client.chat(
+                messages=unique_prompts,
+                **sampling_params,
+                chat_template_kwargs=self.chat_template_kwargs,
+                tools=self.tools,
+                chat_template=getattr(self, "chat_template", None),
+            )
+        else:
+            output = vg.vllm_client.generate(prompts=unique_prompts, **sampling_params)
+
+        # vLLM returns 1 prompt_ids per unique prompt, but num_generations completion_ids.
+        # Duplicate prompt_ids to match completions (one per generation).
+        raw_prompt_ids = output["prompt_ids"]
+        prompt_ids = [pid for pid in raw_prompt_ids for _ in range(num_generations)]
+        completion_ids = output["completion_ids"]
+        logprobs_raw = output["logprobs"]
+        extra_fields = {
+            k: v
+            for k, v in output.items()
+            if k
+            not in {"prompt_ids", "completion_ids", "logprobs", "logprob_token_ids"}
+        }
+
+        # Extract top-1 logprob per token
+        logprobs = [[lp[0] for lp in seq] for seq in logprobs_raw]
+
+        # Decode completions
+        if is_conversational({"prompt": prompts[0]}):
+            contents = self.processing_class.batch_decode(
+                completion_ids, skip_special_tokens=True
+            )
+            completions = [[{"role": "assistant", "content": c}] for c in contents]
+        else:
+            completions = self.processing_class.batch_decode(
+                completion_ids, skip_special_tokens=True
+            )
+
+        tool_mask = extra_fields.pop("env_mask", None)
+
+        # Compute total completion tokens locally (no gather)
+        total_completion_tokens = sum(len(ids) for ids in completion_ids)
+
+        return (
+            prompt_ids,
+            completion_ids,
+            tool_mask,
+            completions,
+            total_completion_tokens,
+            logprobs,
+            extra_fields,
+        )
+
+    def _generate_only(self, inputs, rank0_only=False):
+        """Generate completions without scoring.  Runs on background thread.
+
+        Mirrors the first half of ``_generate_and_score_completions`` (prompt
+        extraction → vLLM generation → tensor padding) and returns a deferred
+        output dict for main-thread scoring.
+
+        When ``rank0_only=True`` (FSDP mode), bypasses ``gather_object`` /
+        ``broadcast_object_list`` collectives and calls vLLM directly on rank 0.
+        Results are broadcast to other ranks on the main thread later.
+
+        Args:
+            inputs: list of dicts (one per sample), as yielded by the DataLoader
+                    with ``identity`` collate_fn.
+        """
+        device = self.accelerator.device
+
+        prompts = [x["prompt"] for x in inputs]
+
+        # --- Handle images (multimodal) ---
+        if "images" in inputs[0]:
+            images = [ex.get("images") for ex in inputs]
+        elif "image" in inputs[0]:
+            images = [
+                [ex.get("image")] if ex.get("image") is not None else None
+                for ex in inputs
+            ]
+        else:
+            images = None
+        if images is not None and all(img == [] for img in images):
+            images = None
+
+        if images is not None:
+            if not is_conversational(inputs[0]):
+                raise ValueError("Multimodal training requires conversational prompts.")
+            prompts = [
+                prepare_multimodal_messages(p, il)
+                for p, il in zip(prompts, images, strict=True)
+            ]
+
+        # --- Generate completions ---
+        if rank0_only:
+            # FSDP mode: call vLLM directly without cross-rank collectives
+            (
+                prompt_ids_list,
+                completion_ids_list,
+                tool_mask_list,
+                completions,
+                num_items_in_batch,
+                sampling_per_token_logps_list,
+                extra_fields,
+            ) = self._generate_rank0_only(prompts)
+        else:
+            (
+                prompt_ids_list,
+                completion_ids_list,
+                tool_mask_list,
+                completions,
+                num_items_in_batch,
+                sampling_per_token_logps_list,
+                extra_fields,
+            ) = self._generate(prompts)
+            # _generate gathers prompts from all ranks internally. Gather inputs
+            # to match the full-batch output size.
+            if self.accelerator.num_processes > 1:
+                from accelerate.utils import gather_object
+
+                inputs = gather_object(inputs)
+                prompts = [x["prompt"] for x in inputs]
+
+        # --- Pad to tensors ---
+        prompt_ids = [torch.tensor(ids, device=device) for ids in prompt_ids_list]
+        prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids]
+        prompt_ids = pad(
+            prompt_ids, padding_value=self.pad_token_id, padding_side="left"
+        )
+        prompt_mask = pad(prompt_mask, padding_value=0, padding_side="left")
+
+        completion_ids = [
+            torch.tensor(ids, device=device) for ids in completion_ids_list
+        ]
+        completion_mask = [
+            torch.ones_like(ids, dtype=torch.long) for ids in completion_ids
+        ]
+        completion_ids = pad(
+            completion_ids, padding_value=self.pad_token_id, padding_side="right"
+        )
+        completion_mask = pad(completion_mask, padding_value=0, padding_side="right")
+
+        if sampling_per_token_logps_list is not None:
+            sampling_logps = [
+                torch.tensor(lp, device=device) for lp in sampling_per_token_logps_list
+            ]
+            sampling_per_token_logps = pad(
+                sampling_logps, padding_value=0.0, padding_side="right"
+            )
+        else:
+            sampling_per_token_logps = None
+
+        if tool_mask_list is not None:
+            tool_mask = [torch.tensor(m, device=device) for m in tool_mask_list]
+            tool_mask = pad(tool_mask, padding_value=1, padding_side="right")
+        else:
+            tool_mask = None
+
+        # --- Mask truncated completions ---
+        if self.mask_truncated_completions:
+            eos_and_pad = [self.eos_token_id, self.pad_token_id]
+            is_trunc = torch.tensor(
+                [ids[-1] not in eos_and_pad for ids in completion_ids_list],
+                device=device,
+            )
+            completion_mask = completion_mask * (~is_trunc).unsqueeze(1).int()
+            if tool_mask is not None:
+                tool_mask = tool_mask * (~is_trunc).unsqueeze(1).int()
+
+        # --- Multimodal forward kwargs ---
+        num_images = [len(il) for il in images] if images is not None else None
+        if images is not None:
+            prompts_text = [
+                apply_chat_template(
+                    {"prompt": p},
+                    self.processing_class,
+                    tools=self.tools,
+                    **self.chat_template_kwargs,
+                )["prompt"]
+                for p in prompts
+            ]
+            prompt_inputs = self.processing_class(
+                images=images, text=prompts_text, padding=True, return_tensors="pt"
+            )
+            forward_kwargs = {
+                k: v.to(device) if isinstance(v, torch.Tensor) else v
+                for k, v in prompt_inputs.items()
+                if k not in ("input_ids", "attention_mask")
+            }
+        else:
+            forward_kwargs = {}
+
+        # Extend token_type_ids / mm_token_type_ids for completion tokens
+        for ttid_key in ("token_type_ids", "mm_token_type_ids"):
+            if ttid_key in forward_kwargs:
+                tt = forward_kwargs[ttid_key]
+                forward_kwargs[ttid_key] = torch.cat(
+                    [tt, tt.new_zeros(completion_ids.shape)], dim=1
+                )
+
+        # Merge extra_fields from rollout_func into inputs
+        if extra_fields:
+            for i, inp in enumerate(inputs):
+                for key, values in extra_fields.items():
+                    if isinstance(values, list) and i < len(values):
+                        inp[key] = values[i]
+                    elif not isinstance(values, list):
+                        inp[key] = values
+
+        # No explicit CUDA sync needed here — both threads share the
+        # default stream, so operations are naturally ordered.
+
+        # --- Construct deferred output ---
+        output = {
+            "prompt_ids": prompt_ids,
+            "prompt_mask": prompt_mask,
+            "completion_ids": completion_ids,
+            "completion_mask": completion_mask,
+            "num_items_in_batch": num_items_in_batch,
+            "advantages": torch.zeros(completion_ids.size(0), device=device),
+            # Sentinels for deferred scoring
+            "_pending_policy_logps": True,
+            "_deferred_inputs": inputs,
+            "_deferred_prompts": prompts,
+            "_deferred_completions": completions,
+            "_deferred_completion_ids_list": completion_ids_list,
+            "_rank0_only": rank0_only,
+        }
+        if sampling_per_token_logps is not None:
+            output["sampling_per_token_logps"] = sampling_per_token_logps
+        if tool_mask is not None:
+            output["tool_mask"] = tool_mask
+        if images is not None:
+            output["num_images"] = num_images
+        for k in (
+            "pixel_values",
+            "image_grid_thw",
+            "pixel_attention_mask",
+            "image_sizes",
+            "token_type_ids",
+            "mm_token_type_ids",
+        ):
+            if k in forward_kwargs:
+                output[k] = forward_kwargs[k]
+        return output
+
+    # ------------------------------------------------------------------
+    # Hooks (overridden by subclasses like FastAsyncGRPOTrainer)
+    # ------------------------------------------------------------------
+
+    def _compute_rewards_for_batch(
+        self, inputs, prompts, completions, completion_ids_list
+    ):
+        """Compute rewards for a batch. Override for parallel workers, caching, etc."""
+        return self._calculate_rewards(
+            inputs, prompts, completions, completion_ids_list
+        )
+
+    def _launch_reward_workers(self, inputs, prompts, completions, completion_ids_list):
+        """Launch reward computation in background. Override for parallel dispatch.
+
+        Default: no-op (rewards computed synchronously in _collect_reward_workers).
+        """
+        self._pending_reward_args = (inputs, prompts, completions, completion_ids_list)
+
+    def _collect_reward_workers(
+        self, inputs, prompts, completions, completion_ids_list
+    ):
+        """Collect reward results. Override to collect from parallel workers.
+
+        Default: compute rewards synchronously now.
+        """
+        args = getattr(self, "_pending_reward_args", None)
+        if args is not None:
+            self._pending_reward_args = None
+            return self._compute_rewards_for_batch(*args)
+        return self._compute_rewards_for_batch(
+            inputs, prompts, completions, completion_ids_list
+        )
+
+    def _post_advantage_hook(
+        self,
+        data: dict,
+        rewards_per_func,
+        advantages,
+        inputs: list,
+        num_generations: int,
+        mode: str,
+        s_start: int | None = None,
+        s_end: int | None = None,
+        is_last_chunk: bool = True,
+    ) -> None:
+        """Called after advantages are computed. Override for replay buffer, re-roll, etc."""
+
+    def _notify_rollouts_scored(
+        self,
+        prompts: list[str],
+        completions: list[str],
+        rewards: dict[str, list[float]],
+        advantages: list[float],
+    ):
+        """Dispatch on_rollouts_scored to all registered plugins (rank 0 only)."""
+        if not self.accelerator.is_main_process:
+            return
+
+        from axolotl.integrations.base import PluginManager
+
+        pm = PluginManager.get_instance()
+        if pm and pm.plugins:
+            # Try _axolotl_cfg first (set by causal builder), fall back to
+            # PluginManager's stored cfg (set during register phase).
+            cfg = getattr(self, "_axolotl_cfg", None) or getattr(pm, "_cfg", None)
+            if cfg is not None:
+                pm.on_rollouts_scored(
+                    cfg, self, prompts, completions, rewards, advantages
+                )
+
+    # ------------------------------------------------------------------
+    # Main-thread scoring
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def _compute_deferred_scores(self, rollout: dict) -> dict:
+        """Compute rewards, advantages, policy logprobs, and IS ratio on the main thread.
+
+        Takes the deferred output from ``_generate_only`` and produces a fully
+        scored dict ready for ``split_tensor_dict`` → micro-batches.
+        """
+        device = self.accelerator.device
+        batch_size = self.args.per_device_train_batch_size
+        num_generations = self.num_generations
+        mode = "train"
+
+        # --- Extract deferred data ---
+        data = rollout
+        inputs = data.pop("_deferred_inputs")
+        prompts = data.pop("_deferred_prompts")
+        completions = data.pop("_deferred_completions")
+        completion_ids_list = data.pop("_deferred_completion_ids_list")
+        rank0_only = data.pop("_rank0_only", False)
+        del data["_pending_policy_logps"]
+
+        prompt_ids = data["prompt_ids"]
+        completion_ids = data["completion_ids"]
+        prompt_mask = data["prompt_mask"]
+        completion_mask = data["completion_mask"]
+        prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
+        logits_to_keep = completion_ids.size(1)
+
+        # Multimodal forward kwargs
+        forward_kwargs = {}
+        for key in (
+            "pixel_values",
+            "image_grid_thw",
+            "pixel_attention_mask",
+            "image_sizes",
+            "token_type_ids",
+            "mm_token_type_ids",
+        ):
+            if key in data:
+                forward_kwargs[key] = data[key]
+        num_images = data.get("num_images")
+
+        # --- Launch rewards in parallel with logprobs ---
+        self._launch_reward_workers(inputs, prompts, completions, completion_ids_list)
+
+        # --- Policy logprobs ---
+        # When batch_flattening is enabled, use the flattened (padding-free) forward
+        # pass for the scoring path. This removes padding tokens before the forward
+        # pass, reducing attention FLOPs proportional to the padding ratio (20-34%
+        # faster in benchmarks). Requires flash_attention_2 and no multimodal inputs.
+        can_flatten = (
+            getattr(self.args, "batch_flattening", False)
+            and not forward_kwargs  # no multimodal inputs
+            and not self.is_fsdp_enabled  # FSDP needs wrapped model
+        )
+
+        logprob_batch_size = min(batch_size * 4, len(prompt_ids))
+        with disable_gradient_checkpointing(
+            self.model, self.args.gradient_checkpointing_kwargs
+        ):
+            generate_every = self.args.steps_per_generation * self.num_iterations
+            if self.args.gradient_accumulation_steps % generate_every != 0 or (
+                self.use_vllm
+                and getattr(self, "vllm_importance_sampling_correction", False)
+            ):
+                if can_flatten:
+                    old_per_token_logps = self._get_per_token_logps_flattened(
+                        self.model,
+                        prompt_completion_ids,
+                        attention_mask,
+                        logits_to_keep,
+                        batch_size=logprob_batch_size,
+                        prompt_mask=prompt_mask,
+                    )
+                else:
+                    old_per_token_logps, _ = self._get_per_token_logps_and_entropies(
+                        self.model,
+                        prompt_completion_ids,
+                        attention_mask,
+                        logits_to_keep,
+                        logprob_batch_size,
+                        num_images=num_images,
+                        **forward_kwargs,
+                    )
+                data["old_per_token_logps"] = old_per_token_logps
+            else:
+                old_per_token_logps = None
+
+            # Reference model logprobs
+            if self.beta != 0.0:
+                if self.ref_model is not None:
+                    ref_logps, _ = self._get_per_token_logps_and_entropies(
+                        self.ref_model,
+                        prompt_completion_ids,
+                        attention_mask,
+                        logits_to_keep,
+                        batch_size,
+                        num_images=num_images,
+                        **forward_kwargs,
+                    )
+                else:
+                    unwrapped = self.accelerator.unwrap_model(self.model)
+                    adapter_name = (
+                        "ref"
+                        if hasattr(unwrapped, "peft_config")
+                        and "ref" in unwrapped.peft_config
+                        else None
+                    )
+                    with use_adapter(unwrapped, adapter_name=adapter_name):
+                        ref_logps, _ = self._get_per_token_logps_and_entropies(
+                            self.model,
+                            prompt_completion_ids,
+                            attention_mask,
+                            logits_to_keep,
+                            batch_size,
+                            num_images=num_images,
+                            **forward_kwargs,
+                        )
+                data["ref_per_token_logps"] = ref_logps
+
+        # --- IS ratio ---
+        if (
+            self.use_vllm
+            and getattr(self, "vllm_importance_sampling_correction", False)
+            and old_per_token_logps is not None
+            and "sampling_per_token_logps" in data
+        ):
+            sampling_logps = data["sampling_per_token_logps"]
+            is_mask = (
+                completion_mask
+                if "tool_mask" not in data
+                else completion_mask * data["tool_mask"]
+            )
+            per_token_logps_diff = (old_per_token_logps - sampling_logps) * is_mask
+
+            is_mode = getattr(self, "vllm_importance_sampling_mode", "token_truncate")
+            is_cap = getattr(self, "vllm_importance_sampling_cap", 3.0)
+            sequence_level_is = is_mode in ("sequence_mask", "sequence_truncate")
+            if sequence_level_is:
+                logps_diff = per_token_logps_diff.sum(dim=-1, keepdim=True)
+            else:
+                logps_diff = per_token_logps_diff
+
+            is_ratio = torch.exp(logps_diff)
+            is_floor = 1.0 / is_cap  # symmetric floor (e.g., cap=3.0 -> floor=0.333)
+            if is_mode in ("sequence_truncate", "token_truncate"):
+                is_ratio = torch.clamp(is_ratio, min=is_floor, max=is_cap)
+            elif is_mode in ("sequence_mask", "token_mask"):
+                is_ratio = is_ratio.masked_fill(is_ratio > is_cap, value=0.0)
+                is_ratio = is_ratio.clamp(min=is_floor)
+            data["importance_sampling_ratio"] = is_ratio
+
+        # --- Collect rewards (launched before logprobs, should be done) ---
+        rewards_per_func = self._collect_reward_workers(
+            inputs, prompts, completions, completion_ids_list
+        )
+        # In rank0_only mode, all ranks compute the same rewards on identical data.
+        # _calculate_rewards / _collect_reward_workers always `gather()` across ranks,
+        # which duplicates the rows (N_local * num_processes).  De-duplicate so that
+        # rewards_per_func matches the data dict (which has N_local rows).
+        if rank0_only and rewards_per_func.size(0) > len(prompts):
+            rewards_per_func = rewards_per_func[: len(prompts)]
+
+        # --- Advantages ---
+        if self.multi_objective_aggregation == "sum_then_normalize":
+            rewards = (
+                rewards_per_func * self.reward_weights.to(device).unsqueeze(0)
+            ).nansum(dim=1)
+            mean_grouped = (
+                rewards.view(-1, num_generations)
+                .mean(dim=1)
+                .repeat_interleave(num_generations)
+            )
+            if self.scale_rewards in ("group", "none"):
+                if num_generations > 1:
+                    std_rewards = (
+                        rewards.view(-1, num_generations)
+                        .std(dim=1)
+                        .repeat_interleave(num_generations)
+                    )
+                else:
+                    std_rewards = torch.zeros_like(rewards)
+            elif self.scale_rewards == "batch":
+                std_rewards = (
+                    rewards.std().expand_as(rewards)
+                    if rewards.numel() > 1
+                    else torch.zeros_like(rewards)
+                )
+            else:
+                raise ValueError(f"Invalid scale_rewards: {self.scale_rewards}")
+            advantages = rewards - mean_grouped
+            if self.scale_rewards != "none":
+                advantages = advantages / (std_rewards + 1e-4)
+            is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards))
+
+        elif self.multi_objective_aggregation == "normalize_then_sum":
+            grouped = rewards_per_func.view(-1, num_generations, len(self.reward_funcs))
+            mean_k = torch.nanmean(grouped, dim=1, keepdim=True)
+            std_k = (
+                nanstd(grouped, dim=1, keepdim=True)
+                if num_generations > 1
+                else torch.zeros_like(mean_k)
+            )
+            reward_k = (grouped - mean_k) / (std_k + 1e-4)
+            reward_k = reward_k.view(-1, len(self.reward_funcs))
+            rewards = (reward_k * self.reward_weights.to(device).unsqueeze(0)).nansum(
+                dim=1
+            )
+            std_rewards = (
+                rewards.std().expand_as(rewards)
+                if rewards.numel() > 1
+                else torch.zeros_like(rewards)
+            )
+            advantages = (rewards - rewards.mean()) / (std_rewards + 1e-4)
+            is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards))
+        else:
+            raise ValueError(
+                f"Invalid multi_objective_aggregation: {self.multi_objective_aggregation}"
+            )
+
+        # Slice for local process
+        # In rank0_only mode, all ranks already have identical data from broadcast,
+        # so no slicing needed. Otherwise, each rank takes its portion.
+        if rank0_only:
+            process_slice = slice(0, len(prompts))
+        else:
+            process_slice = slice(
+                self.accelerator.process_index * len(prompts),
+                (self.accelerator.process_index + 1) * len(prompts),
+            )
+        all_advantages = advantages.clone()
+        advantages = advantages[process_slice]
+        data["advantages"] = advantages
+
+        # --- Post-advantage hook (for replay buffer, re-roll, etc.) ---
+        self._post_advantage_hook(
+            data,
+            rewards_per_func,
+            advantages,
+            inputs,
+            num_generations,
+            mode,
+        )
+
+        # --- Metrics ---
+        for i, name in enumerate(self.reward_func_names):
+            self._metrics[mode][f"rewards/{name}/mean"].append(
+                torch.nanmean(rewards_per_func[:, i]).item()
+            )
+            self._metrics[mode][f"rewards/{name}/std"].append(
+                nanstd(rewards_per_func[:, i]).item()
+            )
+        agg_rewards = rewards_per_func.nansum(dim=1)
+        self._metrics[mode]["reward"].append(agg_rewards.mean().item())
+        self._metrics[mode]["reward_std"].append(agg_rewards.std().item())
+        self._metrics[mode]["frac_reward_zero_std"].append(
+            is_std_zero.float().mean().item()
+        )
+
+        # Token counting
+        total_prompt = self.accelerator.gather(prompt_mask.sum()).sum()
+        total_completion = self.accelerator.gather(completion_mask.sum()).sum()
+        self.state.num_input_tokens_seen += (total_prompt + total_completion).item()
+        self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen]
+
+        # Completion length metrics
+        comp_lengths = completion_mask.sum(dim=1)
+        agg_lengths = self.accelerator.gather(comp_lengths)
+        self._metrics[mode]["completions/mean_length"].append(
+            agg_lengths.float().mean().item()
+        )
+        self._metrics[mode]["completions/min_length"].append(
+            agg_lengths.float().min().item()
+        )
+        self._metrics[mode]["completions/max_length"].append(
+            agg_lengths.float().max().item()
+        )
+
+        eos_and_pad = [self.eos_token_id, self.pad_token_id]
+        is_trunc = torch.tensor(
+            [ids[-1].item() not in eos_and_pad for ids in completion_ids], device=device
+        )
+        agg_trunc = self.accelerator.gather(is_trunc)
+        self._metrics[mode]["completions/clipped_ratio"].append(
+            agg_trunc.float().mean().item()
+        )
+        term_lengths = agg_lengths[~agg_trunc]
+        if len(term_lengths) == 0:
+            term_lengths = torch.zeros(1, device=device)
+        self._metrics[mode]["completions/mean_terminated_length"].append(
+            term_lengths.float().mean().item()
+        )
+        self._metrics[mode]["completions/min_terminated_length"].append(
+            term_lengths.float().min().item()
+        )
+        self._metrics[mode]["completions/max_terminated_length"].append(
+            term_lengths.float().max().item()
+        )
+
+        # IS metrics
+        if "importance_sampling_ratio" in data and "sampling_per_token_logps" in data:
+            old_lp = data["old_per_token_logps"]
+            samp_lp = data["sampling_per_token_logps"]
+            mask = completion_mask.bool()
+            delta = torch.abs(old_lp - samp_lp)
+            delta_m = delta[mask]
+            md = (
+                torch.mean(delta_m)
+                if delta_m.numel() > 0
+                else torch.tensor(0.0, device=device)
+            )
+            xd = (
+                torch.max(delta_m)
+                if delta_m.numel() > 0
+                else torch.tensor(0.0, device=device)
+            )
+            self._metrics[mode]["sampling/sampling_logp_difference/mean"].append(
+                self.accelerator.gather(md).mean().item()
+            )
+            self._metrics[mode]["sampling/sampling_logp_difference/max"].append(
+                self.accelerator.gather(xd).max().item()
+            )
+            isr = data["importance_sampling_ratio"]
+            is_mode = getattr(self, "vllm_importance_sampling_mode", "token_truncate")
+            if is_mode in ("sequence_mask", "sequence_truncate"):
+                flat_isr = isr.flatten()
+            else:
+                flat_isr = isr[mask]
+            if flat_isr.numel() > 0:
+                self._metrics[mode]["sampling/importance_sampling_ratio/min"].append(
+                    nanmin(self.accelerator.gather(torch.min(flat_isr))).item()
+                )
+                self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append(
+                    self.accelerator.gather(torch.mean(flat_isr)).nanmean().item()
+                )
+                self._metrics[mode]["sampling/importance_sampling_ratio/max"].append(
+                    nanmax(self.accelerator.gather(torch.max(flat_isr))).item()
+                )
+
+        # Log prompt/completion texts.
+        # NB: gather_object merges per-rank local texts into a full-batch list
+        # matching rewards_per_func and all_advantages which are already full-batch
+        # tensors (gathered/computed earlier in this method). Lengths stay aligned.
+        prompts_text = self.processing_class.batch_decode(
+            prompt_ids, skip_special_tokens=True
+        )
+        completions_text = self.processing_class.batch_decode(
+            completion_ids, skip_special_tokens=True
+        )
+        if gather_object is not None:
+            gathered_prompts = gather_object(prompts_text)
+            gathered_completions = gather_object(completions_text)
+            self._logs["prompt"].extend(gathered_prompts)
+            self._logs["completion"].extend(gathered_completions)
+        else:
+            gathered_prompts = prompts_text
+            gathered_completions = completions_text
+        rewards_dict = {}
+        for i, name in enumerate(self.reward_func_names):
+            reward_list = rewards_per_func[:, i].tolist()  # already full-batch
+            self._logs["rewards"][name].extend(reward_list)
+            rewards_dict[name] = reward_list
+        adv_list = all_advantages.tolist()  # already full-batch
+        self._logs["advantages"].extend(adv_list)
+
+        # Notify plugins of scored rollouts
+        self._notify_rollouts_scored(
+            gathered_prompts, gathered_completions, rewards_dict, adv_list
+        )
+
+        # Remove deferred keys
+        for k in list(data.keys()):
+            if k.startswith("_deferred") or k == "_pending_policy_logps":
+                data.pop(k, None)
+
+        return data
+
+    @torch.no_grad()
+    def _compute_streaming_group_scores(
+        self,
+        data,
+        s_start,
+        s_end,
+        inputs,
+        prompts,
+        completions,
+        completion_ids_list,
+        is_last_chunk,
+        rank0_only=False,
+    ):
+        """Score a chunk of prompt groups: rewards, policy logprobs, advantages.
+
+        Called during streaming scoring to incrementally score groups.
+        Writes results directly into ``data`` at positions ``s_start:s_end``.
+        """
+        device = self.accelerator.device
+        batch_size = self.args.per_device_train_batch_size
+        num_generations = self.num_generations
+        mode = "train"
+        chunk_size = s_end - s_start
+
+        # --- Policy logprobs for this chunk ---
+        chunk_prompt_ids = data["prompt_ids"][s_start:s_end]
+        chunk_completion_ids = data["completion_ids"][s_start:s_end]
+        chunk_prompt_mask = data["prompt_mask"][s_start:s_end]
+        chunk_completion_mask = data["completion_mask"][s_start:s_end]
+        prompt_completion_ids = torch.cat(
+            [chunk_prompt_ids, chunk_completion_ids], dim=1
+        )
+        attention_mask = torch.cat([chunk_prompt_mask, chunk_completion_mask], dim=1)
+        logits_to_keep = chunk_completion_ids.size(1)
+
+        # Slice multimodal forward kwargs for this chunk
+        forward_kwargs = {}
+        for key in (
+            "pixel_values",
+            "image_grid_thw",
+            "pixel_attention_mask",
+            "image_sizes",
+            "token_type_ids",
+            "mm_token_type_ids",
+        ):
+            if key in data:
+                val = data[key]
+                if (
+                    isinstance(val, torch.Tensor)
+                    and val.dim() > 0
+                    and val.size(0) == len(data["prompt_ids"])
+                ):
+                    forward_kwargs[key] = val[s_start:s_end]
+                else:
+                    forward_kwargs[key] = val
+        num_images = data.get("num_images")
+        if (
+            num_images is not None
+            and hasattr(num_images, "__getitem__")
+            and len(num_images) == len(data["prompt_ids"])
+        ):
+            num_images = num_images[s_start:s_end]
+
+        # --- Launch rewards in parallel with logprobs ---
+        self._launch_reward_workers(inputs, prompts, completions, completion_ids_list)
+
+        # --- Policy logprobs for this chunk (GPU, overlaps with BG rewards) ---
+        can_flatten = (
+            getattr(self.args, "batch_flattening", False)
+            and not forward_kwargs
+            and not self.is_fsdp_enabled
+        )
+        logprob_batch_size = min(batch_size * 2, chunk_size)
+        with disable_gradient_checkpointing(
+            self.model, self.args.gradient_checkpointing_kwargs
+        ):
+            generate_every = self.args.steps_per_generation * self.num_iterations
+            if self.args.gradient_accumulation_steps % generate_every != 0 or (
+                self.use_vllm
+                and getattr(self, "vllm_importance_sampling_correction", False)
+            ):
+                if can_flatten:
+                    old_logps = self._get_per_token_logps_flattened(
+                        self.model,
+                        prompt_completion_ids,
+                        attention_mask,
+                        logits_to_keep,
+                        batch_size=logprob_batch_size,
+                        prompt_mask=chunk_prompt_mask,
+                    )
+                else:
+                    old_logps, _ = self._get_per_token_logps_and_entropies(
+                        self.model,
+                        prompt_completion_ids,
+                        attention_mask,
+                        logits_to_keep,
+                        logprob_batch_size,
+                        num_images=num_images,
+                        **forward_kwargs,
+                    )
+                if "old_per_token_logps" not in data:
+                    total = len(data["prompt_ids"])
+                    data["old_per_token_logps"] = torch.zeros(
+                        total, old_logps.size(1), device=device, dtype=old_logps.dtype
+                    )
+                data["old_per_token_logps"][s_start:s_end] = old_logps
+
+                # Compute IS ratio for this chunk
+                if "sampling_per_token_logps" in data:
+                    samp_chunk = data["sampling_per_token_logps"][s_start:s_end]
+                    is_mask = (
+                        chunk_completion_mask
+                        if "tool_mask" not in data
+                        else (chunk_completion_mask * data["tool_mask"][s_start:s_end])
+                    )
+                    diff = (old_logps - samp_chunk) * is_mask
+                    is_mode = getattr(
+                        self, "vllm_importance_sampling_mode", "token_truncate"
+                    )
+                    is_cap = getattr(self, "vllm_importance_sampling_cap", 3.0)
+                    seq_is = is_mode in ("sequence_mask", "sequence_truncate")
+                    logps_diff = diff.sum(dim=-1, keepdim=True) if seq_is else diff
+                    is_ratio = torch.exp(logps_diff)
+                    # Symmetric floor clamp (matches non-streaming path at line ~1651)
+                    is_floor = 1.0 / is_cap
+                    if is_mode in ("sequence_truncate", "token_truncate"):
+                        is_ratio = torch.clamp(is_ratio, min=is_floor, max=is_cap)
+                    elif is_mode in ("sequence_mask", "token_mask"):
+                        is_ratio = is_ratio.masked_fill(is_ratio > is_cap, value=0.0)
+                        is_ratio = is_ratio.clamp(min=is_floor)
+                    if "importance_sampling_ratio" not in data:
+                        total = len(data["prompt_ids"])
+                        shape = (total, 1) if seq_is else (total, is_ratio.size(1))
+                        data["importance_sampling_ratio"] = torch.ones(
+                            *shape, device=device, dtype=is_ratio.dtype
+                        )
+                    data["importance_sampling_ratio"][s_start:s_end] = is_ratio
+
+            # Reference logprobs
+            if self.beta != 0.0:
+                if self.ref_model is not None:
+                    ref_logps, _ = self._get_per_token_logps_and_entropies(
+                        self.ref_model,
+                        prompt_completion_ids,
+                        attention_mask,
+                        logits_to_keep,
+                        batch_size,
+                        num_images=num_images,
+                        **forward_kwargs,
+                    )
+                else:
+                    unwrapped = self.accelerator.unwrap_model(self.model)
+                    adapter_name = (
+                        "ref"
+                        if hasattr(unwrapped, "peft_config")
+                        and "ref" in unwrapped.peft_config
+                        else None
+                    )
+                    with use_adapter(unwrapped, adapter_name=adapter_name):
+                        ref_logps, _ = self._get_per_token_logps_and_entropies(
+                            self.model,
+                            prompt_completion_ids,
+                            attention_mask,
+                            logits_to_keep,
+                            batch_size,
+                            num_images=num_images,
+                            **forward_kwargs,
+                        )
+                if "ref_per_token_logps" not in data:
+                    total = len(data["prompt_ids"])
+                    data["ref_per_token_logps"] = torch.zeros(
+                        total, ref_logps.size(1), device=device, dtype=ref_logps.dtype
+                    )
+                data["ref_per_token_logps"][s_start:s_end] = ref_logps
+
+        # --- Collect rewards (should already be done, ran in parallel with logprobs) ---
+        rewards_per_func = self._collect_reward_workers(
+            inputs, prompts, completions, completion_ids_list
+        )
+        # De-duplicate gathered rewards when all ranks computed the same data.
+        # _calculate_rewards always gather()s, which duplicates rows in rank0_only mode.
+        if rewards_per_func.size(0) > chunk_size:
+            rewards_per_func = rewards_per_func[:chunk_size]
+
+        # --- Advantages (group-level normalization) ---
+        if self.multi_objective_aggregation == "sum_then_normalize":
+            rewards = (
+                rewards_per_func * self.reward_weights.to(device).unsqueeze(0)
+            ).nansum(dim=1)
+            mean_g = (
+                rewards.view(-1, num_generations)
+                .mean(dim=1)
+                .repeat_interleave(num_generations)
+            )
+            if num_generations > 1:
+                std_r = (
+                    rewards.view(-1, num_generations)
+                    .std(dim=1)
+                    .repeat_interleave(num_generations)
+                )
+            else:
+                std_r = torch.zeros_like(rewards)
+            advantages = rewards - mean_g
+            if self.scale_rewards != "none":
+                advantages = advantages / (std_r + 1e-4)
+            is_std_zero = torch.isclose(std_r, torch.zeros_like(std_r))
+
+        elif self.multi_objective_aggregation == "normalize_then_sum":
+            grouped = rewards_per_func.view(-1, num_generations, len(self.reward_funcs))
+            mean_k = torch.nanmean(grouped, dim=1, keepdim=True)
+            std_k = (
+                nanstd(grouped, dim=1, keepdim=True)
+                if num_generations > 1
+                else torch.zeros_like(mean_k)
+            )
+            reward_k = ((grouped - mean_k) / (std_k + 1e-4)).view(
+                -1, len(self.reward_funcs)
+            )
+            rewards = (reward_k * self.reward_weights.to(device).unsqueeze(0)).nansum(
+                dim=1
+            )
+            std_r = (
+                rewards.view(-1, num_generations)
+                .std(dim=1)
+                .repeat_interleave(num_generations)
+            )
+            mean_r = (
+                rewards.view(-1, num_generations)
+                .mean(dim=1)
+                .repeat_interleave(num_generations)
+            )
+            advantages = (rewards - mean_r) / (std_r + 1e-4)
+            is_std_zero = torch.isclose(std_r, torch.zeros_like(std_r))
+        else:
+            raise ValueError(
+                f"Invalid multi_objective_aggregation: {self.multi_objective_aggregation}"
+            )
+
+        if rank0_only:
+            process_slice = slice(0, len(prompts))
+        else:
+            process_slice = slice(
+                self.accelerator.process_index * len(prompts),
+                (self.accelerator.process_index + 1) * len(prompts),
+            )
+        advantages = advantages[process_slice]
+
+        if "advantages" not in data or not isinstance(data["advantages"], torch.Tensor):
+            data["advantages"] = torch.zeros(len(data["prompt_ids"]), device=device)
+        data["advantages"][s_start:s_end] = advantages
+
+        # --- Post-advantage hook (for replay buffer, re-roll, etc.) ---
+        self._post_advantage_hook(
+            data,
+            rewards_per_func,
+            advantages,
+            inputs,
+            num_generations,
+            mode,
+            s_start=s_start,
+            s_end=s_end,
+            is_last_chunk=is_last_chunk,
+        )
+
+        # --- Chunk metrics ---
+        for i, name in enumerate(self.reward_func_names):
+            self._metrics[mode][f"rewards/{name}/mean"].append(
+                torch.nanmean(rewards_per_func[:, i]).item()
+            )
+            self._metrics[mode][f"rewards/{name}/std"].append(
+                nanstd(rewards_per_func[:, i]).item()
+            )
+        agg_rewards = rewards_per_func.nansum(dim=1)
+        self._metrics[mode]["reward"].append(agg_rewards.mean().item())
+        self._metrics[mode]["reward_std"].append(agg_rewards.std().item())
+        self._metrics[mode]["frac_reward_zero_std"].append(
+            is_std_zero.float().mean().item()
+        )
+
+        # --- Full-batch metrics on last chunk ---
+        if is_last_chunk:
+            all_prompt_mask = data["prompt_mask"]
+            all_completion_mask = data["completion_mask"]
+            all_completion_ids = data["completion_ids"]
+            total_p = self.accelerator.gather(all_prompt_mask.sum()).sum()
+            total_c = self.accelerator.gather(all_completion_mask.sum()).sum()
+            self.state.num_input_tokens_seen += (total_p + total_c).item()
+            self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen]
+
+            comp_lengths = all_completion_mask.sum(dim=1)
+            agg_lengths = self.accelerator.gather(comp_lengths)
+            self._metrics[mode]["completions/mean_length"].append(
+                agg_lengths.float().mean().item()
+            )
+            self._metrics[mode]["completions/min_length"].append(
+                agg_lengths.float().min().item()
+            )
+            self._metrics[mode]["completions/max_length"].append(
+                agg_lengths.float().max().item()
+            )
+
+            eos_and_pad = [self.eos_token_id, self.pad_token_id]
+            is_trunc = torch.tensor(
+                [ids[-1].item() not in eos_and_pad for ids in all_completion_ids],
+                device=device,
+            )
+            agg_trunc = self.accelerator.gather(is_trunc)
+            self._metrics[mode]["completions/clipped_ratio"].append(
+                agg_trunc.float().mean().item()
+            )
+            term = agg_lengths[~agg_trunc]
+            if len(term) == 0:
+                term = torch.zeros(1, device=device)
+            self._metrics[mode]["completions/mean_terminated_length"].append(
+                term.float().mean().item()
+            )
+            self._metrics[mode]["completions/min_terminated_length"].append(
+                term.float().min().item()
+            )
+            self._metrics[mode]["completions/max_terminated_length"].append(
+                term.float().max().item()
+            )
+
+            # IS metrics
+            if (
+                self.use_vllm
+                and getattr(self, "vllm_importance_sampling_correction", False)
+                and "sampling_per_token_logps" in data
+                and "old_per_token_logps" in data
+            ):
+                old_lp = data["old_per_token_logps"]
+                samp_lp = data["sampling_per_token_logps"]
+                mask = all_completion_mask.bool()
+                delta = torch.abs(old_lp - samp_lp)[mask]
+                md = (
+                    torch.mean(delta)
+                    if delta.numel() > 0
+                    else torch.tensor(0.0, device=device)
+                )
+                xd = (
+                    torch.max(delta)
+                    if delta.numel() > 0
+                    else torch.tensor(0.0, device=device)
+                )
+                self._metrics[mode]["sampling/sampling_logp_difference/mean"].append(
+                    self.accelerator.gather(md).mean().item()
+                )
+                self._metrics[mode]["sampling/sampling_logp_difference/max"].append(
+                    self.accelerator.gather(xd).max().item()
+                )
+                is_mode = getattr(
+                    self, "vllm_importance_sampling_mode", "token_truncate"
+                )
+                isr = data["importance_sampling_ratio"]
+                flat = (
+                    isr.flatten()
+                    if is_mode in ("sequence_mask", "sequence_truncate")
+                    else isr[mask]
+                )
+                if flat.numel() > 0:
+                    self._metrics[mode][
+                        "sampling/importance_sampling_ratio/min"
+                    ].append(nanmin(self.accelerator.gather(torch.min(flat))).item())
+                    self._metrics[mode][
+                        "sampling/importance_sampling_ratio/mean"
+                    ].append(self.accelerator.gather(torch.mean(flat)).nanmean().item())
+                    self._metrics[mode][
+                        "sampling/importance_sampling_ratio/max"
+                    ].append(nanmax(self.accelerator.gather(torch.max(flat))).item())
+
+    def _score_streaming(self, rollout: dict) -> list[dict]:
+        """Score a rollout using streaming group scoring.  Returns list of micro-batches."""
+        data = rollout
+        num_gen = self.num_generations
+        n_groups = len(data["prompt_ids"]) // num_gen
+        batch_size = self.args.per_device_train_batch_size
+        min_groups = max(1, self.args.streaming_min_groups)
+
+        # Extract deferred data
+        inputs = data.pop("_deferred_inputs")
+        prompts = data.pop("_deferred_prompts")
+        completions = data.pop("_deferred_completions")
+        completion_ids_list = data.pop("_deferred_completion_ids_list")
+        rank0_only = data.pop("_rank0_only", False)
+        del data["_pending_policy_logps"]
+
+        all_micro_batches = []
+        shared_keys = {"num_items_in_batch"}
+
+        for chunk_start_g in range(0, n_groups, min_groups):
+            chunk_end_g = min(chunk_start_g + min_groups, n_groups)
+            s_start = chunk_start_g * num_gen
+            s_end = chunk_end_g * num_gen
+
+            self._compute_streaming_group_scores(
+                data=data,
+                s_start=s_start,
+                s_end=s_end,
+                inputs=inputs[s_start:s_end],
+                prompts=prompts[s_start:s_end],
+                completions=completions[s_start:s_end],
+                completion_ids_list=completion_ids_list[s_start:s_end],
+                is_last_chunk=(chunk_end_g == n_groups),
+                rank0_only=rank0_only,
+            )
+
+            # Yield micro-batches from this scored chunk
+            chunk_size = s_end - s_start
+            perm = torch.randperm(chunk_size)
+            for mb_off in range(0, chunk_size, batch_size):
+                mb_idx = perm[mb_off : mb_off + batch_size]
+                abs_idx = mb_idx + s_start
+                mb = {}
+                for key in data:
+                    if key.startswith("_"):
+                        continue
+                    val = data[key]
+                    if key in shared_keys:
+                        mb[key] = val
+                    elif isinstance(val, torch.Tensor) and val.dim() > 0:
+                        mb[key] = val[abs_idx]
+                    else:
+                        mb[key] = val
+                all_micro_batches.append(mb)
+
+        # Repeat for num_iterations
+        return all_micro_batches * self.num_iterations
+
+    # ------------------------------------------------------------------
+    # _prepare_inputs override
+    # ------------------------------------------------------------------
+
+    def _prepare_inputs(self, generation_batch):
+        """Override to support data producer and async prefetch paths."""
+        mode = "train" if self.model.training else "eval"
+
+        # --- Data producer path ---
+        if mode == "train" and self.data_producer is not None:
+            return self._prepare_inputs_data_producer(generation_batch)
+
+        # --- Legacy async prefetch path (no data producer) ---
+        if mode == "train" and self.args.async_prefetch:
+            return self._prepare_inputs_legacy_async(generation_batch)
+
+        # --- Stock path ---
+        return super()._prepare_inputs(generation_batch)
+
+    def _prepare_inputs_data_producer(self, generation_batch):
+        """Data producer path: produce rollout, score deferred logps, split into micro-batches.
+
+        Architecture (with async_prefetch=True):
+          BG thread:  produce(skip_policy_logps=True) → vLLM generation + reward computation
+          Main thread: deferred scoring (policy logprobs via GPU forward pass) → training
+
+        Why deferred scoring is necessary for stable training:
+          The policy logprobs (old_per_token_logps) must come from the CURRENT
+          training model, not the vLLM model (which is N steps behind). Using
+          stale vLLM logprobs as old_logps causes the importance sampling ratio
+          to start far from 1.0, leading to:
+            - Immediate PPO clipping → wasted samples
+            - High-variance gradients from IS correction
+            - Compounding per-token ratio errors on long sequences
+            - In extreme cases, complete training failure (exp-003: accuracy=0)
+
+          Deferred scoring computes old_logps with the latest model weights, so
+          the IS ratio starts at exactly 1.0 and drifts gradually — giving
+          maximum useful gradient signal before clipping activates.
+
+          Cost: one additional forward pass per scoring round (GPU-bound, cannot
+          overlap with training on the same GPU). Use ``batch_flattening: true``
+          to reduce this cost by eliminating padding tokens from the forward pass.
+
+        Pipeline:
+          [produce(BG)] → [deferred_scores(GPU)] → [train×GA(GPU)] → [weight_sync]
+                           ↑ can't overlap with train (same GPU)
+
+        Bottleneck: the produce() wait (generation-limited) dominates when
+        generation is slower than training + scoring. Async prefetch hides
+        part of this by generating in the BG thread while training runs.
+        """
+        # Return from buffer if available
+        if self._buffered_inputs:
+            return self._buffered_inputs.pop(0)
+
+        # Produce a new rollout
+        self._maybe_sync_vllm_weights()
+
+        rollout_dataset = self.data_producer.produce(
+            self.model,
+            self.state.global_step,
+            processing_class=self.processing_class,
+            accelerator=self.accelerator,
+            args=self.args,
+        )
+
+        rollout = rollout_dataset._data
+
+        if rollout.get("_pending_policy_logps"):
+            if self.args.streaming_partial_batch:
+                micro_batches = self._score_streaming(rollout)
+            else:
+                scored = self._compute_deferred_scores(rollout)
+                scored = split_pixel_values_by_grid(scored)
+                scored = shuffle_sequence_dict(scored)
+                batches = split_tensor_dict(scored, self.args.steps_per_generation)
+                micro_batches = [unsplit_pixel_values_by_grid(b) for b in batches]
+                micro_batches = micro_batches * self.num_iterations
+        else:
+            rollout = split_pixel_values_by_grid(rollout)
+            batches = split_tensor_dict(rollout, self.args.steps_per_generation)
+            micro_batches = [unsplit_pixel_values_by_grid(b) for b in batches]
+            micro_batches = micro_batches * self.num_iterations
+
+        self._buffered_inputs = micro_batches[1:]
+        return micro_batches[0]
+
+    def _prepare_inputs_legacy_async(self, generation_batch):
+        """Legacy async path: direct queue-based prefetch without data producer."""
+        # Return from buffer if available
+        if self._buffered_inputs:
+            return self._buffered_inputs.pop(0)
+
+        # Need a new rollout
+        self._maybe_sync_vllm_weights()
+        future = self._async_queue.get()
+        rollout = future.result()
+        self._submit_generation()
+
+        # With multi-process, only rank 0 generated. Broadcast to all ranks.
+        if self.accelerator.num_processes > 1:
+            rollout = self._broadcast_rollout(rollout)
+
+        if self.args.streaming_partial_batch:
+            micro_batches = self._score_streaming(rollout)
+        else:
+            scored = self._compute_deferred_scores(rollout)
+            scored = split_pixel_values_by_grid(scored)
+            scored = shuffle_sequence_dict(scored)
+            batches = split_tensor_dict(scored, self.args.steps_per_generation)
+            micro_batches = [unsplit_pixel_values_by_grid(b) for b in batches]
+            micro_batches = micro_batches * self.num_iterations
+
+        self._buffered_inputs = micro_batches[1:]
+
+        # Release cached CUDA memory from scoring
+        # before training allocations begin, reducing peak reserved memory.
+        torch.cuda.empty_cache()
+
+        return micro_batches[0]
+
+    def _get_per_token_logps_flattened(
+        self,
+        model,
+        input_ids,
+        attention_mask,
+        logits_to_keep,
+        batch_size=None,
+        prompt_mask=None,
+    ) -> torch.Tensor:
+        """Compute per-token log-probs using batch flattening (padding-free).
+
+        Instead of processing padded batches where attention wastes compute on
+        padding tokens, this method:
+        1. Chunks the batch into sub-batches of ``batch_size`` sequences
+        2. For each chunk, flattens non-padding tokens into [1, chunk_tokens]
+        3. Uses FlashAttentionKwargs (cu_seq_lens) for varlen attention
+        4. Computes selective_log_softmax on the flat logits
+        5. Gathers completion logprobs back to (B, logits_to_keep) padded format
+
+        Args:
+            prompt_mask: (B, L) mask where 1 = prompt token, 0 = completion/padding.
+                Used to determine the exact prompt length per sequence for correct
+                logprob gathering. If None, inferred as seq_len - logits_to_keep.
+
+        Chunking prevents OOM when the total flattened sequence is too long
+        (e.g., 32 sequences × 2048 tokens = 65K tokens → 20GB logits tensor).
+
+        Requires flash_attention_2 attention implementation.
+        """
+        if not self.is_fsdp_enabled:
+            model = self.accelerator.unwrap_model(model, keep_fp32_wrapper=False)
+
+        device = input_ids.device
+        B, L = input_ids.shape
+        if batch_size is None:
+            batch_size = max(1, B)
+
+        autocast_ctx = torch.autocast(device_type=device.type, dtype=torch.bfloat16)
+        all_logps = torch.zeros(B, logits_to_keep, device=device)
+
+        for chunk_start in range(0, B, batch_size):
+            chunk_end = min(chunk_start + batch_size, B)
+            chunk_ids = input_ids[chunk_start:chunk_end]
+            chunk_mask = attention_mask[chunk_start:chunk_end]
+            n = chunk_end - chunk_start
+
+            seq_lens = chunk_mask.sum(dim=1).to(torch.int32)
+            total_tokens = seq_lens.sum().item()
+            cu_seqlens = torch.zeros(n + 1, dtype=torch.int32, device=device)
+            cu_seqlens[1:] = seq_lens.cumsum(0)
+
+            valid = chunk_mask.bool()
+            flat_ids = chunk_ids[valid].unsqueeze(0)
+            positions = torch.arange(L, device=device).unsqueeze(0).expand(n, L)
+            flat_pos = positions[valid].unsqueeze(0)
+
+            with autocast_ctx:
+                logits = model(
+                    input_ids=flat_ids,
+                    position_ids=flat_pos,
+                    use_cache=False,
+                    cu_seq_lens_q=cu_seqlens,
+                    cu_seq_lens_k=cu_seqlens,
+                    max_length_q=seq_lens.max().item(),
+                    max_length_k=seq_lens.max().item(),
+                ).logits
+                logits = torch.nan_to_num(logits, nan=0.0)
+
+                # Compute logprobs on the flat shifted tensor
+                flat_logits = logits[0, :-1, :] / self.temperature
+                flat_targets = flat_ids[0, 1:]
+                flat_logps = selective_log_softmax(
+                    flat_logits.unsqueeze(0), flat_targets.unsqueeze(0)
+                )[0]
+
+                # Mask out cross-sequence boundary positions. In the shifted
+                # tensor, position cu_seqlens[i]-1 (for i>0) is where sequence
+                # i-1's last token "predicts" sequence i's first token — garbage.
+                for boundary in cu_seqlens[1:-1]:
+                    idx = boundary.item() - 1
+                    if 0 <= idx < flat_logps.size(0):
+                        flat_logps[idx] = 0.0
+
+            # Gather completion logprobs per sequence.
+            # Use prompt_mask to determine exact prompt length (not logits_to_keep,
+            # which is the padded completion dimension and may exceed the actual
+            # completion length for shorter sequences).
+            for i in range(n):
+                slen = seq_lens[i].item()
+                abs_i = chunk_start + i  # absolute index in the full batch
+                if prompt_mask is not None:
+                    plen = int(prompt_mask[abs_i].sum().item())
+                else:
+                    plen = max(1, slen - logits_to_keep)
+                n_compl = slen - plen
+                start = cu_seqlens[i].item() + plen - 1
+                start = max(0, start)
+                actual = min(n_compl, total_tokens - 1 - start)
+                if actual > 0:
+                    all_logps[chunk_start + i, :actual] = flat_logps[
+                        start : start + actual
+                    ]
+
+            del logits, flat_logits, flat_logps, flat_ids
+            torch.cuda.empty_cache()
+
+        return all_logps
+
+    def _get_per_token_logps_and_entropies_flattened(
+        self,
+        model,
+        input_ids,
+        attention_mask,
+        logits_to_keep,
+        batch_size=None,
+        prompt_mask=None,
+        compute_entropy=True,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Flattened forward pass for training (with gradients).
+
+        Same padding removal as the scoring path, but:
+        - Gradients flow through for backward pass
+        - Computes entropy alongside logprobs
+        - Per-sequence logprob/entropy extraction preserves grad graph
+        """
+        device = input_ids.device
+        B, L = input_ids.shape
+        if batch_size is None:
+            batch_size = max(1, B)
+
+        autocast_ctx = torch.autocast(device_type=device.type, dtype=torch.bfloat16)
+
+        # Pre-allocate output containers (will be filled with grad-carrying slices)
+        all_logps_list: list[torch.Tensor] = []
+        all_entropy_list: list[torch.Tensor] = []
+
+        for chunk_start in range(0, B, batch_size):
+            chunk_end = min(chunk_start + batch_size, B)
+            chunk_ids = input_ids[chunk_start:chunk_end]
+            chunk_mask = attention_mask[chunk_start:chunk_end]
+            n = chunk_end - chunk_start
+
+            seq_lens = chunk_mask.sum(dim=1).to(torch.int32)
+            cu_seqlens = torch.zeros(n + 1, dtype=torch.int32, device=device)
+            cu_seqlens[1:] = seq_lens.cumsum(0)
+
+            valid = chunk_mask.bool()
+            flat_ids = chunk_ids[valid].unsqueeze(0)
+            positions = torch.arange(L, device=device).unsqueeze(0).expand(n, L)
+            flat_pos = positions[valid].unsqueeze(0)
+
+            with autocast_ctx:
+                logits = model(
+                    input_ids=flat_ids,
+                    position_ids=flat_pos,
+                    use_cache=False,
+                    cu_seq_lens_q=cu_seqlens,
+                    cu_seq_lens_k=cu_seqlens,
+                    max_length_q=seq_lens.max().item(),
+                    max_length_k=seq_lens.max().item(),
+                ).logits
+                logits = torch.nan_to_num(logits, nan=0.0)
+
+            # Extract logprobs and entropy per-sequence (avoids cross-sequence targets,
+            # preserves gradient graph through selective_log_softmax → logits → model)
+            for i in range(n):
+                slen = seq_lens[i].item()
+                abs_i = chunk_start + i
+                if prompt_mask is not None:
+                    plen = int(prompt_mask[abs_i].sum().item())
+                else:
+                    plen = max(1, slen - logits_to_keep)
+                n_compl = slen - plen
+                s = cu_seqlens[i].item()
+
+                if n_compl <= 0:
+                    # No completion tokens — append zeros
+                    all_logps_list.append(torch.zeros(logits_to_keep, device=device))
+                    if compute_entropy:
+                        all_entropy_list.append(
+                            torch.zeros(logits_to_keep, device=device)
+                        )
+                    continue
+
+                with autocast_ctx:
+                    # Shifted logits and targets for this sequence only
+                    seq_logits = logits[0, s + plen - 1 : s + slen - 1, :]
+                    seq_logits = seq_logits / self.temperature
+                    seq_targets = flat_ids[0, s + plen : s + slen]
+
+                    # Log probs (differentiable)
+                    lps = selective_log_softmax(
+                        seq_logits.unsqueeze(0), seq_targets.unsqueeze(0)
+                    )[0]  # (n_compl,)
+
+                    # Pad to logits_to_keep
+                    if n_compl < logits_to_keep:
+                        lps = F.pad(lps, (0, logits_to_keep - n_compl))
+                    all_logps_list.append(lps[:logits_to_keep])
+
+                    if compute_entropy:
+                        ent = entropy_from_logits(seq_logits)  # (n_compl,)
+                        if n_compl < logits_to_keep:
+                            ent = F.pad(ent, (0, logits_to_keep - n_compl))
+                        all_entropy_list.append(ent[:logits_to_keep])
+
+        # Stack per-sequence results into (B, logits_to_keep) tensors
+        all_logps = torch.stack(all_logps_list, dim=0)
+        all_entropies = (
+            torch.stack(all_entropy_list, dim=0) if compute_entropy else None
+        )
+        return all_logps, all_entropies
+
+    @profiling_decorator
+    def _get_per_token_logps_and_entropies(
+        self,
+        model,
+        input_ids,
+        attention_mask,
+        logits_to_keep,
+        batch_size=None,
+        compute_entropy=False,
+        pixel_values=None,
+        image_grid_thw=None,
+        num_images=None,
+        pixel_attention_mask=None,
+        image_sizes=None,
+        token_type_ids=None,
+        mm_token_type_ids=None,
+    ) -> tuple[Any, torch.Tensor | None]:
+        """Compute log-probs and (optionally) entropies for each token.
+
+        When running under no_grad (scoring path), bypasses accelerate's
+        ConvertOutputsToFp32 wrapper to avoid a fp32 copy of the
+        logits tensor.
+        """
+        # Bypass accelerate's ConvertOutputsToFp32 wrapper which converts the
+        # entire (B, L, V) logits tensor from bf16 to fp32 — unnecessary and
+        # extremely wasteful for large vocabularies.
+        # Skip unwrapping for FSDP — parameters are only valid inside FSDP's
+        # forward context; unwrapping exposes flattened/sharded tensors.
+        if not self.is_fsdp_enabled:
+            model = self.accelerator.unwrap_model(model, keep_fp32_wrapper=False)
+        autocast_ctx = torch.autocast(
+            device_type=input_ids.device.type, dtype=torch.bfloat16
+        )
+
+        # Use Liger's Triton kernel in scoring path (no grad): fuses
+        # temperature + log_softmax + gather into a single kernel pass.
+        use_fused = (
+            self.use_liger_kernel
+            and _fused_selective_log_softmax is not None
+            and not torch.is_grad_enabled()
+        )
+
+        batch_size = batch_size or input_ids.size(0)
+        all_logps = []
+        all_entropies = []
+        with autocast_ctx:
+            for start in range(0, input_ids.size(0), batch_size):
+                input_ids_batch = input_ids[start : start + batch_size]
+                attention_mask_batch = attention_mask[start : start + batch_size]
+
+                # Build model inputs
+                model_inputs = {
+                    "input_ids": input_ids_batch,
+                    "attention_mask": attention_mask_batch,
+                }
+                if image_grid_thw is not None and pixel_values is not None:
+                    rows_per_image = image_grid_thw.prod(dim=-1)
+                    rows_per_sample = torch.split(rows_per_image, num_images)
+                    rows_per_sample = torch.stack([s.sum() for s in rows_per_sample])
+                    cum_rows = torch.cat(
+                        [
+                            torch.tensor([0], device=rows_per_sample.device),
+                            rows_per_sample.cumsum(0),
+                        ]
+                    )
+                    row_start, row_end = (
+                        cum_rows[start].item(),
+                        cum_rows[start + batch_size].item(),
+                    )
+                    model_inputs["pixel_values"] = pixel_values[row_start:row_end]
+                    cum_imgs = torch.tensor([0] + num_images).cumsum(0)
+                    img_start, img_end = cum_imgs[start], cum_imgs[start + batch_size]
+                    model_inputs["image_grid_thw"] = image_grid_thw[img_start:img_end]
+                elif pixel_values is not None:
+                    model_inputs["pixel_values"] = pixel_values[
+                        start : start + batch_size
+                    ]
+                if pixel_attention_mask is not None:
+                    model_inputs["pixel_attention_mask"] = pixel_attention_mask[
+                        start : start + batch_size
+                    ]
+                if image_sizes is not None:
+                    model_inputs["image_sizes"] = image_sizes[
+                        start : start + batch_size
+                    ]
+                if token_type_ids is not None:
+                    model_inputs["token_type_ids"] = token_type_ids[
+                        start : start + batch_size
+                    ]
+                if mm_token_type_ids is not None:
+                    model_inputs["mm_token_type_ids"] = mm_token_type_ids[
+                        start : start + batch_size
+                    ]
+
+                if "logits_to_keep" in self.model_kwarg_keys:
+                    model_inputs["logits_to_keep"] = logits_to_keep + 1
+
+                model_inputs["use_cache"] = False
+
+                logits = model(**model_inputs).logits
+                completion_ids = input_ids_batch[:, -logits_to_keep:]
+                # FP8 models produce NaN logits at positions where
+                # attention_mask=0 (padding). Replace NaN with 0 so
+                # log_softmax yields uniform distribution for those positions.
+                # The completion_mask ensures these don't affect the loss.
+                logits = torch.nan_to_num(logits, nan=0.0)
+
+                if use_fused:
+                    logits = logits[:, -(logits_to_keep + 1) :, :]
+                    if not logits.is_contiguous():
+                        logits = logits.contiguous()
+                    logps = _fused_selective_log_softmax(
+                        logits, completion_ids, self.temperature
+                    )
+                    all_logps.append(logps)
+                    # Liger fused path doesn't compute entropy — append zeros
+                    if compute_entropy:
+                        all_entropies.append(torch.zeros_like(logps))
+                else:
+                    logits = logits[:, :-1, :]
+                    logits = logits[:, -logits_to_keep:, :]
+                    logits.div_(self.temperature)
+                    logps = selective_log_softmax(logits, completion_ids)
+                    all_logps.append(logps)
+
+                    if compute_entropy:
+                        with torch.no_grad():
+                            entropies = entropy_from_logits(logits)
+                        all_entropies.append(entropies)
+
+        logps = torch.cat(all_logps, dim=0)
+        entropies = torch.cat(all_entropies, dim=0) if compute_entropy else None
+        return logps, entropies
+
+    # ------------------------------------------------------------------
+    # Loss override (adds IS ratio + OPSM)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def get_off_policy_mask(
+        advantages,
+        per_token_logps,
+        sampling_per_token_logps,
+        mask,
+        off_policy_threshold,
+    ):
+        """OPSM from DeepSeek-V3.2: drop sequences with negative advantage + high KL."""
+        kl_div = sampling_per_token_logps - per_token_logps.detach()
+        seq_kl = (kl_div * mask).sum(dim=1, keepdim=True) / mask.sum(
+            dim=1, keepdim=True
+        ).clamp(min=1.0)
+        is_pos_adv = advantages >= 0
+        is_low_kl = seq_kl <= off_policy_threshold
+        return (is_pos_adv | is_low_kl).to(dtype=mask.dtype)
+
+    def _compute_loss(self, model, inputs):
+        """Override to add IS ratio correction and off-policy sequence masking."""
+        prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"]
+        completion_ids, completion_mask = (
+            inputs["completion_ids"],
+            inputs["completion_mask"],
+        )
+        input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
+        logits_to_keep = completion_ids.size(1)
+        mask = (
+            completion_mask
+            if "tool_mask" not in inputs
+            else completion_mask * inputs["tool_mask"]
+        )
+
+        # Check for multimodal inputs
+        forward_kwargs = {
+            k: inputs[k]
+            for k in (
+                "pixel_values",
+                "image_grid_thw",
+                "num_images",
+                "pixel_attention_mask",
+                "image_sizes",
+                "token_type_ids",
+                "mm_token_type_ids",
+            )
+            if k in inputs and inputs[k] is not None
+        }
+
+        can_flatten = (
+            getattr(self.args, "batch_flattening", False)
+            and not forward_kwargs
+            and not self.is_fsdp_enabled
+        )
+
+        if can_flatten:
+            per_token_logps, entropies = (
+                self._get_per_token_logps_and_entropies_flattened(
+                    model,
+                    input_ids,
+                    attention_mask,
+                    logits_to_keep,
+                    prompt_mask=prompt_mask,
+                    compute_entropy=True,
+                )
+            )
+        else:
+            per_token_logps, entropies = self._get_per_token_logps_and_entropies(
+                model,
+                input_ids,
+                attention_mask,
+                logits_to_keep,
+                compute_entropy=True,
+                **forward_kwargs,
+            )
+        if self.top_entropy_quantile < 1.0:
+            entropy_mask = self.get_high_entropy_mask(
+                entropies, mask, 1 - self.top_entropy_quantile
+            )
+        else:
+            entropy_mask = None
+
+        advantages = inputs["advantages"]
+        if advantages.dim() == 1:
+            advantages = advantages.unsqueeze(1)
+
+        old_per_token_logps = inputs.get("old_per_token_logps")
+        old_per_token_logps = (
+            per_token_logps.detach()
+            if old_per_token_logps is None
+            else old_per_token_logps
+        )
+
+        # --- OPSM (off-policy sequence mask) ---
+        off_policy_mask = None
+        if getattr(self, "off_policy_mask_threshold", None) is not None:
+            sampling_per_token_logps = inputs.get(
+                "sampling_per_token_logps", old_per_token_logps
+            )
+            off_policy_mask = self.get_off_policy_mask(
+                advantages=advantages,
+                per_token_logps=per_token_logps,
+                sampling_per_token_logps=sampling_per_token_logps,
+                mask=mask,
+                off_policy_threshold=self.off_policy_mask_threshold,
+            )
+
+        # --- Importance weights ---
+        log_ratio = per_token_logps - old_per_token_logps
+        is_level = getattr(
+            self,
+            "importance_sampling_level",
+            getattr(self.args, "importance_sampling_level", "token"),
+        )
+        if is_level == "token":
+            log_importance_weights = log_ratio
+        elif is_level == "sequence":
+            log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(
+                min=1.0
+            )
+            log_importance_weights = log_importance_weights.unsqueeze(-1)
+        else:
+            raise ValueError(f"Unknown importance sampling level: {is_level}")
+
+        coef_1 = torch.exp(log_importance_weights)
+
+        # --- KL divergence ---
+        if self.beta != 0.0:
+            ref_per_token_logps = inputs["ref_per_token_logps"]
+            per_token_kl = (
+                torch.exp(ref_per_token_logps - per_token_logps)
+                - (ref_per_token_logps - per_token_logps)
+                - 1
+            )
+            if getattr(self.args, "use_bias_correction_kl", False):
+                per_token_kl = per_token_kl * coef_1
+
+        # --- Per-token loss ---
+        if self.loss_type == "cispo":
+            clamped = torch.clamp(coef_1, max=self.epsilon_high).detach()
+            per_token_loss = -clamped * advantages * per_token_logps
+        elif self.loss_type in ("grpo", "bnpo", "dr_grpo", "dapo", "luspo"):
+            coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high)
+            if self.args.delta is not None:
+                coef_1_c = torch.clamp(coef_1, max=self.args.delta)
+            else:
+                coef_1_c = coef_1
+            per_token_loss = -torch.min(coef_1_c * advantages, coef_2 * advantages)
+        elif self.loss_type == "sapo":
+            temps = torch.where(
+                advantages > 0,
+                self.args.sapo_temperature_pos,
+                self.args.sapo_temperature_neg,
+            )
+            soft = torch.sigmoid(temps * (coef_1 - 1)) * 4 / temps
+            per_token_loss = -soft * advantages
+        else:
+            raise ValueError(f"Unknown loss type: {self.loss_type}")
+
+        # --- Apply masks ---
+        if off_policy_mask is not None:
+            per_token_loss = per_token_loss * off_policy_mask
+        if entropy_mask is not None:
+            per_token_loss = per_token_loss * entropy_mask
+
+        # --- IS ratio correction (vLLM distribution mismatch) ---
+        if (
+            self.use_vllm
+            and getattr(self, "vllm_importance_sampling_correction", False)
+            and "importance_sampling_ratio" in inputs
+        ):
+            per_token_loss = per_token_loss * inputs["importance_sampling_ratio"]
+
+        if self.beta != 0.0:
+            per_token_loss = per_token_loss + self.beta * per_token_kl
+
+        # --- Aggregate loss ---
+        mode = "train" if self.model.training else "eval"
+        normalizer = (
+            self.current_gradient_accumulation_steps if mode == "train" else 1.0
+        )
+
+        if self.loss_type in ("grpo", "sapo"):
+            loss = (
+                (per_token_loss * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)
+            ).mean() / normalizer
+        elif self.loss_type == "bnpo":
+            loss = (
+                (per_token_loss * mask).sum() / mask.sum().clamp(min=1.0) / normalizer
+            )
+        elif self.loss_type == "dr_grpo":
+            loss = (
+                (per_token_loss * mask).sum()
+                / (per_token_loss.size(0) * self.max_completion_length)
+                / normalizer
+            )
+        elif self.loss_type in ("cispo", "dapo"):
+            norm = inputs["num_items_in_batch"] / self.accelerator.num_processes
+            loss = (per_token_loss * mask).sum() / norm
+        elif self.loss_type == "luspo":
+            loss = (per_token_loss * mask.sum(1, keepdim=True)).mean() / normalizer
+        else:
+            raise ValueError(f"Unknown loss type: {self.loss_type}")
+
+        # --- Metrics ---
+        completion_token_count = mask.sum().clamp(min=1.0)
+
+        def masked_batch_mean(x):
+            return (
+                x.mean()
+                if x.shape[1] == 1
+                else (x * mask).sum() / completion_token_count
+            )
+
+        if self.beta != 0.0:
+            mean_kl = masked_batch_mean(per_token_kl)
+            self._metrics[mode]["kl"].append(
+                self.accelerator.gather(mean_kl).nanmean().item()
+            )
+
+        mean_entropy = masked_batch_mean(entropies)
+        self._metrics[mode]["entropy"].append(
+            self.accelerator.gather(mean_entropy).nanmean().item()
+        )
+
+        if self.loss_type in ("grpo", "bnpo", "dr_grpo", "dapo", "luspo"):
+            is_low = (coef_1 < 1 - self.epsilon_low) & (advantages < 0)
+            is_high = (coef_1 > 1 + self.epsilon_high) & (advantages > 0)
+            is_region = is_low | is_high
+            low_clip = masked_batch_mean(is_low.float())
+            high_clip = masked_batch_mean(is_high.float())
+            clip_ratio = masked_batch_mean(is_region.float())
+            g_low = self.accelerator.gather(low_clip)
+            self._metrics[mode]["clip_ratio/low_mean"].append(g_low.nanmean().item())
+            self._metrics[mode]["clip_ratio/low_min"].append(nanmin(g_low).item())
+            g_high = self.accelerator.gather(high_clip)
+            self._metrics[mode]["clip_ratio/high_mean"].append(g_high.nanmean().item())
+            self._metrics[mode]["clip_ratio/high_max"].append(nanmax(g_high).item())
+            g_clip = self.accelerator.gather(clip_ratio)
+            self._metrics[mode]["clip_ratio/region_mean"].append(
+                g_clip.nanmean().item()
+            )
+        elif self.loss_type == "cispo":
+            is_cispo = (coef_1 > self.epsilon_high) & (advantages > 0)
+            cr = masked_batch_mean(is_cispo.float())
+            self._metrics[mode]["cispo_clip_ratio"].append(
+                self.accelerator.gather(cr).nanmean().item()
+            )
+
+        return loss
diff --git a/src/axolotl/core/trainers/grpo/fast_async_trainer.py b/src/axolotl/core/trainers/grpo/fast_async_trainer.py
new file mode 100644
index 0000000000..26b1789fc7
--- /dev/null
+++ b/src/axolotl/core/trainers/grpo/fast_async_trainer.py
@@ -0,0 +1,768 @@
+# Copyright 2020-2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Experimental GRPO extensions: parallel reward workers, replay buffer,
+deferred re-roll, and zero-advantage skipping.
+
+These features are built as subclasses of GRPOTrainer and GRPODataProducer,
+using the hook system (_compute_rewards_for_batch, _post_advantage_hook,
+_pre_produce_hook) defined in the base classes.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from dataclasses import dataclass, field
+
+import torch
+from torch import nn
+from trl import GRPOTrainer
+
+from axolotl.core.trainers.grpo.async_trainer import (
+    AsyncGRPOConfig,
+    AsyncGRPOTrainer,
+    GRPODataProducer,
+)
+from axolotl.core.trainers.grpo.replay_buffer import ReplayBuffer
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Extended config
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class FastAsyncGRPOConfig(AsyncGRPOConfig):
+    """GRPOConfig with additional experimental parameters."""
+
+    reward_num_workers: int = field(
+        default=1,
+        metadata={
+            "help": "Number of persistent subprocess workers for parallel reward computation. Each worker has its "
+            "own main thread so signal.alarm() (used by math_verify) works correctly. Work is sharded across "
+            "workers by prompt groups. Only used with use_data_producer=True and non-nn.Module reward functions."
+        },
+    )
+    replay_buffer_size: int = field(
+        default=0,
+        metadata={
+            "help": "[Experimental, disabled by default] Size of the replay buffer for storing high-signal rollout "
+            "groups. When > 0, groups with reward variance are cached and used to replace zero-signal groups "
+            "(where all rewards are identical). Set to 0 to disable. Only used with use_data_producer=True."
+        },
+    )
+    replay_recompute_logps: bool = field(
+        default=True,
+        metadata={
+            "help": "When True (default), recompute old_per_token_logps for replayed groups using the current "
+            "training model. This fixes the importance sampling mismatch that occurs when replaying stale data. "
+            "Only relevant when replay_buffer_size > 0."
+        },
+    )
+    reroll_start_fraction: float = field(
+        default=0.5,
+        metadata={
+            "help": "Fraction of total training steps after which deferred re-rolling begins. Zero-signal prompts "
+            "(where all rewards in a group are identical) are buffered and re-injected into later batches when the "
+            "model is more likely to solve them. Set to 1.0 to disable. Only used with use_data_producer=True."
+        },
+    )
+    reroll_max_groups: int = field(
+        default=1,
+        metadata={
+            "help": "Maximum number of prompt groups to replace with re-roll candidates per batch. Higher values "
+            "increase data utilization but reduce prompt diversity. Only used with use_data_producer=True."
+        },
+    )
+    skip_zero_advantage_batches: bool = field(
+        default=True,
+        metadata={
+            "help": "When True, skip gradient computation for micro-batches where all advantages are zero (no learning "
+            "signal). This avoids the forward/backward pass entirely when no learning signal is present. The step is "
+            "logged with skipped_zero_adv_batches=1 for monitoring."
+        },
+    )
+    vllm_lora_sync: bool = field(
+        default=False,
+        metadata={
+            "help": "When True, sync LoRA adapter weights to vLLM via filesystem instead of merging into base model "
+            "and NCCL-broadcasting all parameters. vLLM loads the adapter natively using Punica kernels. "
+            "Requires vllm_serve_lora serve module (auto-selected when this is True). "
+            "Syncs only LoRA adapter weights (much smaller) vs full merged model. Legacy merge behavior is used when False."
+        },
+    )
+
+
+# ---------------------------------------------------------------------------
+# Extended data producer with re-roll injection
+# ---------------------------------------------------------------------------
+
+
+class RerollDataProducer(GRPODataProducer):
+    """GRPODataProducer that injects re-roll candidates into prompt batches.
+
+    Reads from the trainer's ``_reroll_buffer`` (populated by
+    ``GRPOExperimentalTrainer._post_advantage_hook``) and replaces the
+    last N prompt groups with previously-failed prompts.
+    """
+
+    def _pre_produce_hook(self, inputs: list, global_step: int) -> list:
+        trainer = self._trainer
+        reroll_buf = getattr(trainer, "_reroll_buffer", None)
+        reroll_lock = getattr(trainer, "_reroll_lock", None)
+        if reroll_buf is None or reroll_lock is None:
+            return inputs
+
+        max_steps = getattr(trainer.args, "max_steps", -1)
+        start_frac = getattr(trainer.args, "reroll_start_fraction", 1.0)
+        max_groups = getattr(trainer.args, "reroll_max_groups", 1)
+        reroll_start_step = (
+            max(1, int(max_steps * start_frac)) if max_steps > 0 else float("inf")
+        )
+
+        if global_step < reroll_start_step:
+            return inputs
+
+        with reroll_lock:
+            n_to_take = min(max_groups, len(reroll_buf))
+            reroll_prompts = [reroll_buf.pop(0) for _ in range(n_to_take)]
+
+        if reroll_prompts:
+            num_gen = self._num_generations
+            n_groups = len(inputs) // num_gen
+            for i, reroll_prompt in enumerate(reroll_prompts):
+                group_idx = n_groups - 1 - i
+                if group_idx < 0:
+                    break
+                start = group_idx * num_gen
+                for j in range(num_gen):
+                    inputs[start + j] = reroll_prompt
+            logger.info(
+                f"[REROLL] Step {global_step}: replaced {len(reroll_prompts)}/{n_groups} prompt groups "
+                f"with deferred re-roll candidates ({len(reroll_buf)} remaining)"
+            )
+
+        return inputs
+
+
+# ---------------------------------------------------------------------------
+# Persistent reward subprocess pool
+# ---------------------------------------------------------------------------
+
+
+def _persistent_reward_worker(conn):
+    """Long-lived reward worker. Receives work items, returns results."""
+    while True:
+        try:
+            msg = conn.recv()
+        except EOFError:
+            break
+        if msg is None:  # Shutdown signal
+            break
+        (
+            reward_funcs,
+            prompts,
+            completions,
+            completion_ids_list,
+            inputs,
+            reward_func_names,
+        ) = msg
+        try:
+            keys = [
+                key
+                for key in inputs[0]
+                if key not in ["prompt", "completion", "completion_ids"]
+            ]
+            reward_kwargs = {key: [example[key] for example in inputs] for key in keys}
+            results = []
+            for reward_func, _reward_func_name in zip(
+                reward_funcs, reward_func_names, strict=True
+            ):
+                output = reward_func(
+                    prompts=prompts,
+                    completions=completions,
+                    completion_ids=completion_ids_list,
+                    **reward_kwargs,
+                )
+                results.append(
+                    [float(r) if r is not None else float("nan") for r in output]
+                )
+            conn.send(results)
+        except Exception:
+            conn.send(None)
+
+
+# ---------------------------------------------------------------------------
+# Extended trainer
+# ---------------------------------------------------------------------------
+
+
+class FastAsyncGRPOTrainer(AsyncGRPOTrainer):
+    """GRPOTrainer with experimental extensions.
+
+    Adds:
+    - Parallel reward subprocess workers (``reward_num_workers``)
+    - Replay buffer for high-signal group reuse (``replay_buffer_size``)
+    - Deferred re-roll of failed prompts (``reroll_start_fraction``)
+    - Zero-advantage micro-batch skipping
+    """
+
+    def __init__(self, *args, **kwargs):
+        # These must be initialized before super().__init__() because
+        # _create_data_producer (called during super().__init__) needs them.
+        self._reroll_buffer: list = []
+        self._reroll_lock = threading.Lock()
+
+        # Temporarily suppress the base class's Liger + OPSM validation check,
+        # since this subclass supports it via a custom compute_liger_loss override.
+        grpo_args = kwargs.get("args")
+        if grpo_args is None:
+            for a in args:
+                if hasattr(a, "off_policy_mask_threshold"):
+                    grpo_args = a
+                    break
+        saved_threshold = None
+        if grpo_args is not None and getattr(grpo_args, "use_liger_kernel", False):
+            saved_threshold = grpo_args.off_policy_mask_threshold
+            grpo_args.off_policy_mask_threshold = None
+
+        super().__init__(*args, **kwargs)
+
+        if saved_threshold is not None:
+            grpo_args.off_policy_mask_threshold = saved_threshold
+            self.off_policy_mask_threshold = saved_threshold
+
+        # Replay buffer
+        if getattr(self.args, "replay_buffer_size", 0) > 0:
+            self._replay_buffer = ReplayBuffer(max_size=self.args.replay_buffer_size)
+        else:
+            self._replay_buffer = None
+        self._replay_recompute_logps = getattr(
+            self.args, "replay_recompute_logps", True
+        )
+
+        # Reward worker pool (lazy-initialized)
+        self._reward_workers = None
+
+    # -- Factory override: use RerollDataProducer ----------------------------
+
+    def _create_data_producer(self, args, train_dataset):
+        """Override to use RerollDataProducer for re-roll prompt injection."""
+        from axolotl.core.trainers.grpo.async_trainer import (
+            AsyncDataProducer,
+            ProducerConfig,
+        )
+
+        producer_config = ProducerConfig(
+            mini_epochs=args.num_iterations,
+            max_rollouts=None,
+            eval_during_produce=False,
+            empty_cache_before_produce=True,
+            empty_cache_after_produce=True,
+            async_prefetch=args.async_prefetch,
+            prefetch_depth=args.prefetch_depth,
+        )
+        data_producer = RerollDataProducer(
+            config=producer_config,
+            prompt_dataset=train_dataset,
+            num_generations=self.num_generations,
+            generation_batch_size=args.generation_batch_size,
+            train_batch_size=args.per_device_train_batch_size,
+            steps_per_generation=args.steps_per_generation,
+            shuffle_dataset=self.shuffle_dataset,
+            seed=args.seed,
+        )
+        data_producer.set_trainer(self)
+        if args.async_prefetch:
+            data_producer = AsyncDataProducer(
+                data_producer,
+                background_produce_kwargs={"skip_policy_logps": True},
+            )
+        return data_producer
+
+    # -- Reward worker pool --------------------------------------------------
+
+    def _get_reward_workers(self):
+        """Return a list of persistent reward worker subprocesses (lazy-initialized)."""
+        import multiprocessing as _mp
+
+        num_workers = getattr(self.args, "reward_num_workers", 1)
+        if num_workers < 1:
+            num_workers = 1
+
+        if self._reward_workers is not None:
+            alive = all(proc.is_alive() for conn, proc in self._reward_workers)
+            if alive and len(self._reward_workers) == num_workers:
+                return self._reward_workers
+            self._shutdown_reward_workers()
+
+        workers = []
+        for _ in range(num_workers):
+            parent_conn, child_conn = _mp.Pipe()
+            proc = _mp.Process(
+                target=_persistent_reward_worker, args=(child_conn,), daemon=True
+            )
+            proc.start()
+            child_conn.close()
+            workers.append((parent_conn, proc))
+
+        self._reward_workers = workers
+        return workers
+
+    def _shutdown_reward_workers(self):
+        """Shut down all persistent reward workers."""
+        if self._reward_workers is None:
+            return
+        for conn, proc in self._reward_workers:
+            try:
+                conn.send(None)
+                proc.join(timeout=5)
+            except Exception:
+                pass
+            try:
+                conn.close()
+            except Exception:
+                pass
+        self._reward_workers = None
+
+    # -- Hook overrides ------------------------------------------------------
+
+    def _compute_rewards_for_batch(
+        self, inputs, prompts, completions, completion_ids_list
+    ):
+        """Dispatch rewards to parallel subprocess workers (synchronous wrapper)."""
+        self._launch_reward_workers(inputs, prompts, completions, completion_ids_list)
+        return self._collect_reward_workers(
+            inputs, prompts, completions, completion_ids_list
+        )
+
+    def _launch_reward_workers(self, inputs, prompts, completions, completion_ids_list):
+        """Send reward work to subprocess workers (non-blocking).
+
+        Results are collected later by _collect_reward_workers, allowing GPU
+        logprob computation to overlap with CPU reward computation.
+        """
+        reward_can_bg = all(
+            callable(rf)
+            and not isinstance(rf, nn.Module)
+            and not asyncio.iscoroutinefunction(rf)
+            for rf in self.reward_funcs
+        )
+        num_workers = getattr(self.args, "reward_num_workers", 1)
+
+        if not reward_can_bg or num_workers <= 1:
+            # Can't parallelize — store args for sync fallback in collect
+            self._reward_workers_used = None
+            self._pending_reward_args = (
+                inputs,
+                prompts,
+                completions,
+                completion_ids_list,
+            )
+            return
+
+        workers = self._get_reward_workers()
+        num_generations = self.num_generations
+        num_prompts = len(prompts)
+        num_groups = num_prompts // num_generations
+
+        # Shard by prompt groups across workers
+        groups_per_worker = max(1, (num_groups + len(workers) - 1) // len(workers))
+        workers_used = []
+        for w_idx, (conn, _proc) in enumerate(workers):
+            g_start = w_idx * groups_per_worker
+            g_end = min((w_idx + 1) * groups_per_worker, num_groups)
+            if g_start >= num_groups:
+                break
+            s_start = g_start * num_generations
+            s_end = g_end * num_generations
+            conn.send(
+                (
+                    self.reward_funcs,
+                    prompts[s_start:s_end],
+                    completions[s_start:s_end],
+                    completion_ids_list[s_start:s_end],
+                    inputs[s_start:s_end],
+                    self.reward_func_names,
+                )
+            )
+            workers_used.append(conn)
+
+        self._reward_workers_used = workers_used
+        self._pending_reward_args = (inputs, prompts, completions, completion_ids_list)
+
+    def _collect_reward_workers(
+        self, inputs, prompts, completions, completion_ids_list
+    ):
+        """Collect reward results from subprocess workers (blocks until done)."""
+        from accelerate.utils import gather
+
+        workers_used = getattr(self, "_reward_workers_used", None)
+        args = getattr(self, "_pending_reward_args", None)
+        self._reward_workers_used = None
+        self._pending_reward_args = None
+
+        if workers_used is None:
+            # Sync fallback — compute on main thread
+            if args is not None:
+                return self._calculate_rewards(*args)
+            return self._calculate_rewards(
+                inputs, prompts, completions, completion_ids_list
+            )
+
+        device = self.accelerator.device
+        num_prompts = len(args[1]) if args else len(prompts)
+
+        # Collect results from workers
+        all_worker_results = []
+        any_failed = False
+        for conn in workers_used:
+            result = conn.recv()
+            if result is None:
+                any_failed = True
+                # Drain remaining workers to prevent stale results in pipes
+                for remaining_conn in workers_used:
+                    if remaining_conn is not conn:
+                        try:
+                            remaining_conn.recv()
+                        except Exception:
+                            pass
+                break
+            all_worker_results.append(result)
+
+        if not any_failed:
+            rewards_per_func = torch.zeros(
+                num_prompts, len(self.reward_funcs), device=device
+            )
+            offset = 0
+            for worker_result in all_worker_results:
+                chunk_size = len(worker_result[0])
+                for i, result in enumerate(worker_result):
+                    rewards_per_func[offset : offset + chunk_size, i] = torch.tensor(
+                        result, dtype=torch.float32, device=device
+                    )
+                offset += chunk_size
+            return gather(rewards_per_func)
+
+        # Fallback to main thread on failure
+        if args is not None:
+            return self._calculate_rewards(*args)
+        return self._calculate_rewards(
+            inputs, prompts, completions, completion_ids_list
+        )
+
+    def _post_advantage_hook(
+        self,
+        data: dict,
+        rewards_per_func,
+        advantages,
+        inputs: list,
+        num_generations: int,
+        mode: str,
+        s_start: int | None = None,
+        s_end: int | None = None,
+        is_last_chunk: bool = True,
+    ) -> None:
+        """Replay buffer store/replace + re-roll buffering."""
+        from trl.models.utils import disable_gradient_checkpointing
+
+        # -- Replay buffer: store high-signal groups --
+        if self._replay_buffer is not None:
+            local_grouped = rewards_per_func.view(
+                -1, num_generations, len(self.reward_funcs)
+            )
+            per_group_std = local_grouped.std(dim=1)
+            has_signal = (per_group_std > 0).any(dim=1)
+            offset = s_start or 0
+
+            if has_signal.any():
+                grouped_adv = advantages.view(-1, num_generations)
+                replay_scores = grouped_adv.abs().sum(dim=1) * per_group_std.sum(dim=1)
+                for group_idx in has_signal.nonzero(as_tuple=True)[0]:
+                    gi = group_idx.item()
+                    start = offset + gi * num_generations
+                    end = start + num_generations
+                    group_data = {}
+                    for key in data:
+                        val = data[key]
+                        if (
+                            isinstance(val, torch.Tensor)
+                            and val.dim() > 0
+                            and val.size(0) >= end
+                        ):
+                            group_data[key] = val[start:end].clone()
+                    self._replay_buffer.add(replay_scores[gi].item(), group_data)
+
+            # Replace zero-signal groups with high-signal replay buffer entries
+            # Only in non-streaming path (s_start is None) — streaming scores
+            # groups incrementally, so replacement + logprob recompute would be
+            # too expensive per chunk.
+            n_replaced = 0
+            if s_start is None:
+                no_signal = ~has_signal
+                replaced_ranges = []
+                if no_signal.any() and len(self._replay_buffer) > 0:
+                    for group_idx in no_signal.nonzero(as_tuple=True)[0]:
+                        sampled = self._replay_buffer.sample(1)
+                        if sampled is None:
+                            break
+                        sampled_group = sampled[0]
+                        gi = group_idx.item()
+                        start = offset + gi * num_generations
+                        end = start + num_generations
+                        for key, val in sampled_group.items():
+                            if key in data and isinstance(data[key], torch.Tensor):
+                                src = val.to(data[key].device)
+                                tgt_seq_len = (
+                                    data[key].size(1) if data[key].dim() > 1 else None
+                                )
+                                if start >= data[key].size(0) or end > data[key].size(
+                                    0
+                                ):
+                                    continue
+                                if tgt_seq_len is not None:
+                                    if src.size(1) <= tgt_seq_len:
+                                        data[key][start:end] = 0
+                                        data[key][start:end, : src.size(1)] = src
+                                    else:
+                                        data[key][start:end] = src[:, :tgt_seq_len]
+                                else:
+                                    data[key][start:end] = src
+                        replaced_ranges.append((start, end))
+                        n_replaced += 1
+
+                # Recompute old_per_token_logps for replayed groups
+                if (
+                    n_replaced > 0
+                    and self._replay_recompute_logps
+                    and "old_per_token_logps" in data
+                ):
+                    with (
+                        torch.no_grad(),
+                        disable_gradient_checkpointing(
+                            self.model, self.args.gradient_checkpointing_kwargs
+                        ),
+                    ):
+                        for r_start, r_end in replaced_ranges:
+                            r_ids = torch.cat(
+                                [
+                                    data["prompt_ids"][r_start:r_end],
+                                    data["completion_ids"][r_start:r_end],
+                                ],
+                                dim=1,
+                            )
+                            r_mask = torch.cat(
+                                [
+                                    data["prompt_mask"][r_start:r_end],
+                                    data["completion_mask"][r_start:r_end],
+                                ],
+                                dim=1,
+                            )
+                            r_logits_to_keep = data["completion_ids"].size(1)
+                            r_fwd_kwargs = {}
+                            for fk in (
+                                "pixel_values",
+                                "image_grid_thw",
+                                "pixel_attention_mask",
+                                "image_sizes",
+                                "token_type_ids",
+                                "mm_token_type_ids",
+                            ):
+                                if fk in data:
+                                    r_fwd_kwargs[fk] = data[fk]
+                            r_logps, _ = self._get_per_token_logps_and_entropies(
+                                self.model,
+                                r_ids,
+                                r_mask,
+                                r_logits_to_keep,
+                                r_end - r_start,
+                                **r_fwd_kwargs,
+                            )
+                            data["old_per_token_logps"][r_start:r_end] = r_logps
+
+                if n_replaced > 0:
+                    self._metrics[mode]["replay_buffer_replacements"].append(
+                        float(n_replaced)
+                    )
+
+            if is_last_chunk:
+                self._metrics[mode]["replay_buffer_size"].append(
+                    float(len(self._replay_buffer))
+                )
+
+        # -- Re-roll buffer: store failed prompts --
+        if getattr(self.args, "reroll_start_fraction", 1.0) < 1.0:
+            grouped_rewards = rewards_per_func.view(
+                -1, num_generations, len(self.reward_funcs)
+            )
+            per_group_std = grouped_rewards.std(dim=1)
+            per_group_mean = grouped_rewards.mean(dim=1)
+            zero_signal = (per_group_std == 0).all(dim=1)
+            all_failed = (per_group_mean.abs() < 1e-6).all(dim=1)
+            should_reroll = zero_signal & all_failed
+            _n_buffered = 0
+            with self._reroll_lock:
+                for group_idx in should_reroll.nonzero(as_tuple=True)[0]:
+                    idx = group_idx.item() * num_generations
+                    if idx >= len(inputs):
+                        continue
+                    prompt_input = inputs[idx]
+                    self._reroll_buffer.append(prompt_input)
+                    _n_buffered += 1
+            if _n_buffered > 0:
+                self._metrics[mode]["reroll_buffered"].append(float(_n_buffered))
+            if is_last_chunk:
+                self._metrics[mode]["reroll_buffer_size"].append(
+                    float(len(self._reroll_buffer))
+                )
+
+    # -- Zero-advantage skipping + Liger OPSM ---------------------------------
+
+    def compute_liger_loss(self, unwrapped_model, inputs):
+        """Liger loss with zero-adv skipping and off-policy sequence masking (OPSM).
+
+        The base class Liger path doesn't support OPSM because the fused kernel
+        doesn't expose per-token logprobs needed for the KL computation. This
+        override computes them via chunked lm_head matmul (no grad, low memory)
+        and applies the OPSM to the loss mask before calling the kernel.
+        """
+        if self.args.skip_zero_advantage_batches and torch.all(
+            inputs["advantages"] == 0
+        ):
+            mode = "train" if self.model.training else "eval"
+            self._metrics[mode]["skipped_zero_adv_batches"].append(1.0)
+            return torch.tensor(
+                0.0, device=inputs["advantages"].device, requires_grad=True
+            )
+
+        if self.off_policy_mask_threshold is None:
+            return super().compute_liger_loss(unwrapped_model, inputs)
+
+        # OPSM path: need per_token_logps for KL, which Liger kernel doesn't provide
+        prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"]
+        completion_ids, completion_mask = (
+            inputs["completion_ids"],
+            inputs["completion_mask"],
+        )
+        input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
+        logits_to_keep = completion_ids.size(1)
+
+        last_hidden_state = self._get_last_hidden_state(
+            unwrapped_model,
+            input_ids,
+            attention_mask,
+            logits_to_keep,
+            inputs.get("pixel_values"),
+            inputs.get("image_grid_thw"),
+            inputs.get("pixel_attention_mask"),
+            inputs.get("image_sizes"),
+        )
+
+        loss_mask = (
+            completion_mask
+            if "tool_mask" not in inputs
+            else completion_mask * inputs["tool_mask"]
+        )
+
+        # Compute per_token_logps via chunked lm_head matmul (no grad, low memory)
+        lm_weight = unwrapped_model.lm_head.weight
+        lm_bias = unwrapped_model.lm_head.bias
+        with torch.no_grad():
+            per_token_logps_chunks = []
+            for i in range(last_hidden_state.size(0)):
+                chunk_logits = torch.matmul(last_hidden_state[i : i + 1], lm_weight.t())
+                if lm_bias is not None:
+                    chunk_logits = chunk_logits + lm_bias
+                chunk_lps = (
+                    chunk_logits.float()
+                    .log_softmax(-1)
+                    .gather(-1, completion_ids[i : i + 1].unsqueeze(-1))
+                    .squeeze(-1)
+                )
+                per_token_logps_chunks.append(chunk_lps)
+                del chunk_logits
+            per_token_logps = torch.cat(per_token_logps_chunks, dim=0)
+
+        advantages = inputs["advantages"]
+        if advantages.dim() == 1:
+            advantages_2d = advantages.unsqueeze(1)
+        else:
+            advantages_2d = advantages
+
+        sampling_per_token_logps = inputs.get("sampling_per_token_logps")
+        if sampling_per_token_logps is None:
+            sampling_per_token_logps = inputs.get("old_per_token_logps")
+        if sampling_per_token_logps is None:
+            sampling_per_token_logps = per_token_logps
+
+        off_policy_mask = GRPOTrainer.get_off_policy_mask(
+            advantages=advantages_2d,
+            per_token_logps=per_token_logps,
+            sampling_per_token_logps=sampling_per_token_logps,
+            mask=loss_mask,
+            off_policy_threshold=self.off_policy_mask_threshold,
+        )
+        loss_mask = loss_mask * off_policy_mask
+
+        # Call the Liger fused kernel with OPSM-modified mask
+        loss, metrics = self.liger_grpo_loss(
+            _input=last_hidden_state,
+            lin_weight=unwrapped_model.lm_head.weight,
+            selected_token_ids=completion_ids,
+            attention_mask=loss_mask,
+            advantages=inputs["advantages"],
+            bias=unwrapped_model.lm_head.bias,
+            old_per_token_logps=inputs.get("old_per_token_logps"),
+            ref_per_token_logps=inputs.get("ref_per_token_logps"),
+            vllm_is_ratio=inputs.get("importance_sampling_ratio"),
+        )
+
+        mean_kl = metrics[0] if self.beta != 0.0 else None
+        clip_ratio = metrics[-1]
+
+        mode = "train" if self.model.training else "eval"
+        if self.beta != 0.0:
+            self._metrics[mode]["kl"].append(
+                self.accelerator.gather(mean_kl).mean().item()
+            )
+        self._metrics[mode]["clip_ratio"].append(
+            self.accelerator.gather(clip_ratio).mean().item()
+        )
+        normalizer = (
+            self.current_gradient_accumulation_steps if mode == "train" else 1.0
+        )
+        return loss / normalizer
+
+    def _compute_loss(self, model, inputs):
+        if self.args.skip_zero_advantage_batches and torch.all(
+            inputs["advantages"] == 0
+        ):
+            mode = "train" if self.model.training else "eval"
+            self._metrics[mode]["skipped_zero_adv_batches"].append(1.0)
+            # Create zero loss with grad_fn. DeepSpeed requires grad_fn != None.
+            # With ZeRO-3, parameters are partitioned (shape=[0], requires_grad=False)
+            # so we can't just do `(p * 0).sum()`. Instead, do a tiny forward pass
+            # with a single token to create a proper computation graph.
+            prompt_ids = inputs["prompt_ids"][:1, :1]  # (1, 1)
+            attn = torch.ones_like(prompt_ids)
+            with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+                out = model(input_ids=prompt_ids, attention_mask=attn)
+            return out.logits.sum() * 0
+        return super()._compute_loss(model, inputs)
diff --git a/src/axolotl/core/trainers/grpo/replay_buffer.py b/src/axolotl/core/trainers/grpo/replay_buffer.py
new file mode 100644
index 0000000000..1e35cb56a7
--- /dev/null
+++ b/src/axolotl/core/trainers/grpo/replay_buffer.py
@@ -0,0 +1,44 @@
+"""Simple replay buffer for storing and sampling high-signal rollout groups."""
+
+import heapq
+
+import torch
+
+
+class ReplayBuffer:
+    """Min-heap replay buffer that keeps the highest-scoring rollout groups.
+    Groups are scored by signal quality (advantage magnitude * reward variance).
+    When sampling, groups are drawn proportional to their scores.
+    """
+
+    def __init__(self, max_size: int):
+        self.max_size = max_size
+        self._heap: list[tuple[float, int, dict]] = []  # min-heap of (score, id, data)
+        self._counter = 0  # unique tiebreaker for heap
+
+    def __len__(self):
+        return len(self._heap)
+
+    def add(self, score: float, data: dict):
+        """Add a group to the buffer. If full, replaces lowest-scoring entry."""
+        if self.max_size <= 0:
+            return
+        self._counter += 1
+        if len(self._heap) < self.max_size:
+            heapq.heappush(self._heap, (score, self._counter, data))
+        elif score > self._heap[0][0]:
+            heapq.heapreplace(self._heap, (score, self._counter, data))
+
+    def sample(self, num_samples: int) -> list[dict] | None:
+        """Sample groups weighted by their scores. Returns None if buffer is empty."""
+        if self.max_size <= 0 or not self._heap:
+            return None
+
+        scores = torch.tensor([item[0] for item in self._heap], dtype=torch.float32)
+        scores = scores.clamp(min=1e-8)  # avoid zero probabilities
+        probs = scores / scores.sum()
+        replacement = num_samples > len(self._heap)
+        indices = torch.multinomial(
+            probs, num_samples, replacement=replacement
+        ).tolist()
+        return [self._heap[i][2] for i in indices]
diff --git a/src/axolotl/core/trainers/grpo/sampler.py b/src/axolotl/core/trainers/grpo/sampler.py
new file mode 100644
index 0000000000..df679a6d2f
--- /dev/null
+++ b/src/axolotl/core/trainers/grpo/sampler.py
@@ -0,0 +1,172 @@
+"""Repeat random sampler (similar to the one implemented in
+https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py) that adds
+sequence parallelism functionality; i.e., duplicating data across ranks in the same
+sequence parallel group.
+"""
+
+from typing import Iterator, Sized
+
+import torch
+from torch.utils.data import Sampler
+
+
+class SequenceParallelRepeatRandomSampler(Sampler):
+    """Sampler for GRPO training with sequence parallelism.
+
+    This sampler ensures:
+    - Ranks in the same sequence parallel (SP) group receive identical data.
+    - Each index is repeated multiple times for sampling different completions.
+    - Entire batches are repeated for reuse in multiple updates.
+    - Data is properly distributed across SP groups.
+
+    In the table below, the values represent dataset indices. Each SP group has
+    `context_parallel_size = 2` GPUs working together on the same data. There are 2
+    SP groups (SP0 and SP1), with `world_size = 4` total GPUs.
+
+                                               Sequence Parallel Groups
+                                        |       SP0        |       SP1        |
+                                        |  GPU 0  |  GPU 1 |  GPU 2  |  GPU 3 |
+                    global_step  step    <---> mini_repeat_count=3
+                                            <----------> batch_size=2 per SP group
+    grad_accum=2   ▲  ▲  0       0         [0 0 0  1 1 1]     [2 2 2  3 3 3]   <- SP groups get different data
+                   ▼  |  0       1         [0 0 0  1 1 1]     [2 2 2  3 3 3]   <- Same data for each SP group GPU
+                      |
+                      |  1       2         [0 0 0  1 1 1]     [2 2 2  3 3 3]   <- Repeat same indices for iterations
+    num_iterations=2  ▼  1       3         [0 0 0  1 1 1]     [2 2 2  3 3 3]   <- When using gradient accumulation
+
+                         2       4         [4 4 4  5 5 5]     [6 6 6  7 7 7]   <- New batch of data indices
+                         2       5         [4 4 4  5 5 5]     [6 6 6  7 7 7]
+                                            ...
+
+    Args:
+        dataset: Dataset to sample from.
+        mini_repeat_count: How many times to repeat each sample immediately.
+        world_size: Total number of processes.
+        rank: Rank of current process.
+        batch_size: Number of samples per batch.
+        repeat_count: How many times to repeat the full sampling process.
+        context_parallel_size: Number of ranks in a sequence parallel group.
+        shuffle: Whether to shuffle the dataset.
+        seed: Random seed for shuffling.
+        drop_last: Whether to drop the last incomplete batch.
+    """
+
+    def __init__(
+        self,
+        dataset: Sized,
+        mini_repeat_count: int,
+        world_size: int,
+        rank: int,
+        batch_size: int = 1,
+        repeat_count: int = 1,
+        context_parallel_size: int = 1,
+        shuffle: bool = True,
+        seed: int = 0,
+        drop_last: bool = False,
+    ):
+        self.dataset = dataset
+        self.mini_repeat_count = mini_repeat_count
+        self.batch_size = batch_size
+        self.repeat_count = repeat_count
+        self.shuffle = shuffle
+        self.seed = seed
+        self.drop_last = drop_last
+        self.epoch = 0
+
+        self.world_size = world_size
+        self.rank = rank
+
+        # Sequence parallelism parameters
+        self.context_parallel_size = context_parallel_size
+        self.num_sp_groups = world_size // context_parallel_size
+        self.sp_group_id = rank // context_parallel_size
+
+        # Adjust dataset size for distributed sampling
+        self.num_samples = len(self.dataset)
+        self.total_size = self.num_samples
+
+        # Calculate effective number of samples per SP group
+        if (
+            self.drop_last
+            and self.total_size % (self.num_sp_groups * self.batch_size) != 0
+        ):
+            # Drop last incomplete batch if drop_last is True
+            self.num_samples_per_sp_group = (
+                self.total_size // self.batch_size // self.num_sp_groups
+            ) * self.batch_size
+        else:
+            # Round up to include last batch if drop_last is False
+            self.num_samples_per_sp_group = (
+                (self.total_size + self.batch_size * self.num_sp_groups - 1)
+                // (self.batch_size * self.num_sp_groups)
+                * self.batch_size
+            )
+
+        if shuffle:
+            self.generator = torch.Generator()
+            self.generator.manual_seed(seed)
+
+    def __iter__(self) -> Iterator[int]:
+        """Creates iterator over dataset indices.
+
+        Returns:
+            Iterator that yields indices into the dataset.
+        """
+        # Deterministically shuffle based on epoch and seed
+        if self.shuffle:
+            indices = torch.randperm(
+                self.num_samples, generator=self.generator
+            ).tolist()
+        else:
+            indices = list(range(self.num_samples))
+
+        # Add extra samples to make it evenly divisible by batch_size
+        if len(indices) % self.batch_size != 0:
+            padding = indices[: self.batch_size - len(indices) % self.batch_size]
+            indices += padding
+
+        # Subsample based on SP group ID
+        # Each SP group gets distinct batches of data
+        batch_indices = []
+        for i in range(0, len(indices), self.batch_size * self.num_sp_groups):
+            start_idx = i + self.sp_group_id * self.batch_size
+            end_idx = min(start_idx + self.batch_size, len(indices))
+            if start_idx < len(indices):
+                for j in range(self.batch_size):
+                    if start_idx + j < end_idx:
+                        batch_indices.append(indices[start_idx + j])
+
+        # Make sure batch_indices is exactly batch_size * num_batches_per_sp_group
+        if self.drop_last:
+            num_batches_per_sp_group = self.num_samples_per_sp_group // self.batch_size
+            target_len = self.batch_size * num_batches_per_sp_group
+            if len(batch_indices) > target_len:
+                batch_indices = batch_indices[:target_len]
+
+        # Apply the GRPO repeat pattern
+        final_indices = []
+        for _ in range(self.repeat_count):
+            for idx in batch_indices:
+                for _ in range(self.mini_repeat_count):
+                    final_indices.append(idx)
+
+        return iter(final_indices)
+
+    def __len__(self) -> int:
+        """Returns the total length of the iterable including repetitions.
+
+        Returns:
+            Total number of samples.
+        """
+        # Total length including all repetitions
+        return (
+            self.num_samples_per_sp_group * self.mini_repeat_count * self.repeat_count
+        )
+
+    def set_epoch(self, epoch: int) -> None:
+        """Sets the epoch for this sampler.
+
+        Args:
+            epoch: Epoch number to use for shuffling.
+        """
+        self.epoch = epoch
diff --git a/src/axolotl/core/trainers/grpo/trainer.py b/src/axolotl/core/trainers/grpo/trainer.py
new file mode 100644
index 0000000000..118aa3e9ad
--- /dev/null
+++ b/src/axolotl/core/trainers/grpo/trainer.py
@@ -0,0 +1,698 @@
+"""Axolotl GRPO trainers (with and without sequence parallelism handling)"""
+
+import warnings
+from functools import partial
+from typing import Any
+
+import datasets
+import torch
+import torch.distributed as dist
+import torch.utils.data
+from accelerate.utils import (
+    broadcast_object_list,
+    gather,
+    gather_object,
+    is_peft_available,
+)
+from datasets import Dataset, IterableDataset
+from torch import nn
+from torch.utils.data import (
+    BatchSampler,
+    DataLoader,
+    Sampler,
+)
+from transformers import (
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+)
+from transformers.trainer_utils import seed_worker
+from trl import GRPOTrainer
+from trl.data_utils import (
+    apply_chat_template,
+    is_conversational,
+    maybe_apply_chat_template,
+)
+from trl.extras.profiling import profiling_context
+from trl.models import unwrap_model_for_generation
+from trl.trainer.grpo_config import GRPOConfig
+from trl.trainer.grpo_trainer import RewardFunc, nanstd
+from trl.trainer.utils import pad
+
+from axolotl.core.trainers.grpo.fast_async_trainer import FastAsyncGRPOTrainer
+from axolotl.core.trainers.grpo.sampler import SequenceParallelRepeatRandomSampler
+from axolotl.core.trainers.mixins import (
+    DistributedParallelMixin,
+    RngLoaderMixin,
+    SchedulerMixin,
+)
+from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin
+from axolotl.monkeypatch.ring_attn import get_ring_attn_group
+
+if is_peft_available():
+    from peft import PeftConfig
+
+
+class AxolotlGRPOTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DistributedParallelMixin,
+    GRPOTrainer,
+):
+    """Extend the base GRPOTrainer for axolotl helpers"""
+
+    _tag_names = ["trl", "grpo", "axolotl"]
+
+
+class AxolotlAsyncGRPOTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DistributedParallelMixin,
+    FastAsyncGRPOTrainer,
+):
+    """Extend AsyncGRPOTrainer with axolotl helpers"""
+
+    _tag_names = ["trl", "grpo", "async", "axolotl"]
+
+
+class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
+    """Extend the base GRPOTrainer for sequence parallelism handling"""
+
+    def __init__(
+        self,
+        model: str | PreTrainedModel,
+        reward_funcs: RewardFunc | list[RewardFunc],
+        args: GRPOConfig | None = None,
+        train_dataset: Dataset | IterableDataset | None = None,
+        eval_dataset: (
+            Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None
+        ) = None,
+        processing_class: PreTrainedTokenizerBase | None = None,
+        reward_processing_classes: (
+            PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None
+        ) = None,
+        callbacks: list[TrainerCallback] | None = None,
+        optimizers: tuple[
+            torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None
+        ] = (None, None),
+        peft_config: "PeftConfig | None" = None,
+        optimizer_cls_and_kwargs: tuple[type, dict] | None = None,
+    ):
+        # First call the superclass constructor with all arguments
+        super().__init__(
+            model=model,
+            reward_funcs=reward_funcs,
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=processing_class,
+            reward_processing_classes=reward_processing_classes,
+            callbacks=callbacks,
+            optimizers=optimizers,
+            peft_config=peft_config,
+            optimizer_cls_and_kwargs=optimizer_cls_and_kwargs,
+        )
+
+        # Get number of SP groups (number of processes divided by SP degree)
+        num_processes = self.accelerator.num_processes
+        num_sp_groups = num_processes // self.args.context_parallel_size
+
+        # Calculate batch size per SP group (not per process)
+        sp_group_batch_size = self.args.per_device_train_batch_size * num_sp_groups
+        possible_values = [
+            n_gen
+            for n_gen in range(2, sp_group_batch_size + 1)
+            if (sp_group_batch_size) % n_gen == 0
+        ]
+
+        if self.num_generations not in possible_values:
+            raise ValueError(
+                f"The batch size per SP group ({num_sp_groups} x "
+                f"{self.args.per_device_train_batch_size}) must be evenly divisible by "
+                f"the number of generations per prompt ({self.num_generations}). Given "
+                "the current configuration, the valid values for the number of "
+                f"generations are: {possible_values}."
+            )
+
+        if self.args.eval_strategy != "no":
+            # If sequence parallelism is enabled, calculate batch size per SP group
+            sp_group_eval_batch_size = args.per_device_eval_batch_size * num_sp_groups  # type: ignore[union-attr]
+            possible_values = [
+                n_gen
+                for n_gen in range(2, sp_group_eval_batch_size + 1)
+                if (sp_group_eval_batch_size) % n_gen == 0
+            ]
+
+            if self.num_generations not in possible_values:
+                raise ValueError(
+                    f"With sequence parallelism (degree {self.args.context_parallel_size}), "
+                    f"the eval batch size per SP group ({num_sp_groups} x {self.args.per_device_eval_batch_size}) "
+                    f"must be evenly divisible by the number of generations per prompt "
+                    f"({self.num_generations}). Given the current eval batch size, "
+                    f"the valid values for the number of generations are: {possible_values}."
+                )
+
+        self.sp_group = None
+        self.rank = dist.get_rank()
+        self.world_size = dist.get_world_size()
+        self.local_rank = 0
+        self.local_world_size = 1
+
+    def train(self, *args, **kwargs):
+        # Initialize the SP group
+        self.sp_group = get_ring_attn_group()
+        self.rank = dist.get_rank()
+        self.world_size = dist.get_world_size()
+        self.local_rank = dist.get_rank(group=self.sp_group)
+        self.local_world_size = dist.get_world_size(group=self.sp_group)
+
+        return super().train(*args, **kwargs)
+
+    def _get_train_sampler(self) -> Sampler:
+        effective_batch_size = (
+            self.args.per_device_train_batch_size
+            * self.world_size
+            * self.args.gradient_accumulation_steps
+        )
+
+        return SequenceParallelRepeatRandomSampler(
+            dataset=self.train_dataset,
+            mini_repeat_count=self.num_generations,
+            world_size=self.world_size,
+            rank=self.rank,
+            batch_size=effective_batch_size
+            // self.num_generations
+            // self.args.context_parallel_size,
+            repeat_count=self.num_iterations * self.args.gradient_accumulation_steps,
+            context_parallel_size=self.args.context_parallel_size,
+            shuffle=not self.args.curriculum_sampling,
+            seed=self.args.seed,
+            drop_last=True,
+        )
+
+    def _create_dataloader_params(self, is_eval=False, custom_batch_size=None):
+        """Create common dataloader parameters for train or eval."""
+        batch_size = custom_batch_size or (
+            self.args.eval_batch_size if is_eval else self._train_batch_size
+        )
+
+        params = {
+            "batch_size": batch_size,
+            "collate_fn": self.data_collator,
+            "num_workers": self.args.dataloader_num_workers,
+            "pin_memory": self.args.dataloader_pin_memory,
+        }
+
+        # Add persistent workers only for training
+        if not is_eval and hasattr(self.args, "dataloader_persistent_workers"):
+            params["persistent_workers"] = self.args.dataloader_persistent_workers
+
+        # Add prefetch factor if specified
+        if self.args.dataloader_prefetch_factor:
+            params["prefetch_factor"] = self.args.dataloader_prefetch_factor
+
+        return params
+
+    def _prepare_dataloader(
+        self, dataset, sampler, is_eval=False, custom_batch_size=None
+    ):
+        """Prepare a dataloader with the given dataset and sampler."""
+        # Get base parameters
+        dataloader_params = self._create_dataloader_params(is_eval, custom_batch_size)
+
+        # Add sampler configuration
+        if not isinstance(dataset, torch.utils.data.IterableDataset):
+            if isinstance(sampler, BatchSampler):
+                # batch_size and batch_sampler are mutually exclusive
+                dataloader_params["batch_sampler"] = sampler
+                del dataloader_params["batch_size"]
+            else:
+                dataloader_params["sampler"] = sampler
+                dataloader_params["drop_last"] = self.args.dataloader_drop_last
+
+            if not is_eval:
+                dataloader_params["worker_init_fn"] = partial(
+                    seed_worker,
+                    num_workers=self.args.dataloader_num_workers,
+                    rank=self.args.process_index,
+                )
+
+        # Create the dataloader
+        dataloader = DataLoader(dataset, **dataloader_params)
+
+        if self.args.sample_packing and (
+            (not is_eval and not self.args.pretraining)
+            or (is_eval and self.args.eval_sample_packing is not False)
+        ):
+            self.accelerator.even_batches = False
+
+        # Return unprepared dataloader if using sequence parallelism
+        # TODO(djsaunde): We might be able to use `accelerate`'s dataloader preparation
+        # if we use `dispatch_batches` and `slice_fn_for_dispatch` properly (i.e.,
+        # slice each batch along the sequence dimension).
+        if self.args.context_parallel_size > 1:
+            return dataloader
+
+        # Otherwise prepare with accelerator
+        return self.accelerator.prepare_data_loader(dataloader)
+
+    def get_train_dataloader(self) -> DataLoader:
+        """Get dataloader for training"""
+        train_dataset = self.train_dataset
+
+        data_collator = self.data_collator  # type: ignore
+
+        # Handle dataset preprocessing
+        if isinstance(train_dataset, datasets.Dataset):
+            # Add debug print before any modifications
+            if self.args.sample_packing and not self.args.pretraining:
+                train_dataset = train_dataset.remove_columns(["length"])
+            if not self.args.sample_packing or self.args.pretraining:
+                train_dataset = self._remove_unused_columns(
+                    train_dataset, description="training"
+                )
+        else:
+            self.data_collator = self._get_collator_with_removed_columns(
+                data_collator,
+                description="training",
+            )
+
+        # Get sampler and create dataloader
+        sampler = self._get_train_sampler()
+        dataloader = self._prepare_dataloader(train_dataset, sampler, is_eval=False)
+
+        return dataloader
+
+    def _generate_and_score_completions(
+        self, inputs: list[dict[str, torch.Tensor | Any]]
+    ) -> dict[str, torch.Tensor | Any]:
+        device = self.accelerator.device
+        mode = "eval" if self.control.should_evaluate else "train"
+
+        prompts = [x["prompt"] for x in inputs]
+        prompts_text = [
+            maybe_apply_chat_template(example, self.processing_class)["prompt"]
+            for example in inputs
+        ]
+        prompt_inputs = self.processing_class(
+            text=prompts_text,
+            return_tensors="pt",
+            padding=True,
+            padding_side="left",
+            add_special_tokens=False,
+        )
+        prompt_inputs = Trainer._prepare_inputs(self, prompt_inputs)
+        prompt_ids, prompt_mask = (
+            prompt_inputs["input_ids"],
+            prompt_inputs["attention_mask"],
+        )
+
+        if self.max_prompt_length is not None:
+            prompt_ids = prompt_ids[:, -self.max_prompt_length :]
+            prompt_mask = prompt_mask[:, -self.max_prompt_length :]
+
+        # Generate completions using either vLLM or regular generation
+        if self.args.use_vllm:
+            # First, have main process load weights if needed
+
+            if self.state.global_step != self._last_loaded_step:  # type: ignore[has-type]
+                self._move_model_to_vllm()
+
+                self._last_loaded_step = self.state.global_step
+
+            # Generate completions using vLLM: gather all prompts and use them in a single call in the main process
+            all_prompts_text = gather_object(prompts_text)
+            if self.accelerator.is_main_process:
+                if self.args.context_parallel_size > 1:
+                    # Calculate sequence parallel group information
+                    world_size = self.accelerator.num_processes
+                    context_parallel_size = self.args.context_parallel_size
+                    num_sp_groups = world_size // context_parallel_size
+
+                    # Since processes in the same SP group have the same prompts, we need to ensure
+                    # we only take one copy of each prompt from each SP group
+                    ordered_set_of_prompts = []
+                    for sp_group_id in range(num_sp_groups):
+                        # Get the first process from each SP group (typically the group leader)
+                        group_leader_rank = sp_group_id * context_parallel_size
+
+                        # Extract prompts from this SP group, accounting for num_generations duplicates
+                        # We only need prompts from one rank in each SP group
+                        group_prompts = all_prompts_text[
+                            group_leader_rank * len(prompts_text) : (
+                                group_leader_rank + 1
+                            )
+                            * len(prompts_text) : self.num_generations
+                        ]
+
+                        ordered_set_of_prompts.extend(group_prompts)
+                else:
+                    # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate
+                    # num_generations outputs for each one. This is faster than generating outputs for each duplicate
+                    # prompt individually.
+                    ordered_set_of_prompts = all_prompts_text[
+                        :: self.num_generations * self.args.context_parallel_size
+                    ]
+
+                with profiling_context(self, "vLLM.generate"):
+                    completion_ids = self.vllm_client.generate(
+                        prompts=ordered_set_of_prompts,
+                        n=self.num_generations,
+                        repetition_penalty=self.repetition_penalty,
+                        temperature=self.temperature,
+                        top_p=self.top_p,
+                        top_k=-1 if self.top_k is None else self.top_k,
+                        min_p=0.0 if self.min_p is None else self.min_p,
+                        max_tokens=self.max_completion_length,
+                        guided_decoding_regex=self.guided_decoding_regex,
+                    )
+            else:
+                completion_ids = [None] * (
+                    len(all_prompts_text) // self.args.context_parallel_size
+                )
+
+            # Broadcast the completions from the main process to all processes
+            completion_ids = broadcast_object_list(completion_ids, from_process=0)
+
+            # Determine the appropriate slice based on sequence parallelism
+            if self.args.context_parallel_size > 1:
+                # Calculate SP group ID (which group of ranks this rank belongs to)
+                sp_group_id = self.accelerator.process_index // self.local_world_size
+
+                # Calculate the start index for this SP group
+                sp_group_start = sp_group_id * len(prompts) * self.local_world_size
+
+                # All ranks in the same SP group get the same data slice
+                process_slice = slice(
+                    sp_group_start,
+                    sp_group_start + len(prompts),
+                )
+                completion_ids = completion_ids[process_slice]
+            else:
+                # Original behavior for non-sequence parallel case
+                process_slice = slice(
+                    self.accelerator.process_index * len(prompts),
+                    (self.accelerator.process_index + 1) * len(prompts),
+                )
+                completion_ids = completion_ids[process_slice]
+
+            # Pad the completions, and concatenate them with the prompts
+            completion_ids = [
+                torch.tensor(ids, device=device) for ids in completion_ids
+            ]
+            completion_ids = pad(
+                completion_ids, padding_value=self.processing_class.pad_token_id
+            )
+            prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+        else:
+            # Regular generation path
+            with unwrap_model_for_generation(
+                self.model_wrapped,
+                self.accelerator,
+                gather_deepspeed3_params=self.args.ds3_gather_for_generation,
+            ) as unwrapped_model:
+                prompt_completion_ids = unwrapped_model.generate(
+                    prompt_ids,
+                    attention_mask=prompt_mask,
+                    generation_config=self.generation_config,
+                )
+
+            # Compute prompt length and extract completion ids
+            prompt_length = prompt_ids.size(1)
+            prompt_ids = prompt_completion_ids[:, :prompt_length]
+            completion_ids = prompt_completion_ids[:, prompt_length:]
+
+        # Mask everything after the first EOS token
+        is_eos = completion_ids == self.processing_class.eos_token_id
+        eos_idx = torch.full(
+            (is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device
+        )
+        eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
+        sequence_indices = torch.arange(is_eos.size(1), device=device).expand(
+            is_eos.size(0), -1
+        )
+        completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int()
+
+        # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask
+        if self.args.mask_truncated_completions:
+            truncated_completions = ~is_eos.any(dim=1)
+            completion_mask = (
+                completion_mask * (~truncated_completions).unsqueeze(1).int()
+            )
+
+        # Concatenate prompt_mask with completion_mask for logit computation
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)  # (B, P+C)
+
+        logits_to_keep = completion_ids.size(
+            1
+        )  # we only need to compute the logits for the completion tokens
+        batch_size = (
+            self.args.per_device_train_batch_size
+            if mode == "train"
+            else self.args.per_device_eval_batch_size
+        )
+
+        with torch.no_grad():
+            # When using num_iterations == 1, old_per_token_logps == per_token_logps, so we can skip it's
+            # computation here, and use per_token_logps.detach() instead.
+            if self.num_iterations > 1:
+                old_per_token_logps = self._get_per_token_logps(
+                    self.model,
+                    prompt_completion_ids,
+                    attention_mask,
+                    logits_to_keep,
+                    batch_size,
+                )
+            else:
+                old_per_token_logps = None
+
+            if self.beta == 0.0:
+                ref_per_token_logps = None
+            elif self.ref_model is not None:
+                ref_per_token_logps = self._get_per_token_logps(
+                    self.ref_model,
+                    prompt_completion_ids,
+                    attention_mask,
+                    logits_to_keep,
+                    batch_size,
+                )
+            else:
+                with self.accelerator.unwrap_model(self.model).disable_adapter():
+                    ref_per_token_logps = self._get_per_token_logps(
+                        self.model,
+                        prompt_completion_ids,
+                        attention_mask,
+                        logits_to_keep,
+                        batch_size,
+                    )
+
+        # Decode the generated completions
+        completions_text = self.processing_class.batch_decode(
+            completion_ids, skip_special_tokens=True
+        )
+        if is_conversational(inputs[0]):
+            completions = []
+            for prompt, completion in zip(prompts, completions_text, strict=False):
+                bootstrap = (
+                    prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else ""
+                )
+                completions.append(
+                    [{"role": "assistant", "content": bootstrap + completion}]
+                )
+        else:
+            completions = completions_text
+
+        rewards_per_func = torch.zeros(
+            len(prompts), len(self.reward_funcs), device=device
+        )
+        for i, (reward_func, reward_processing_class, reward_func_name) in enumerate(
+            zip(
+                self.reward_funcs,
+                self.reward_processing_classes,
+                self.reward_func_names,
+                strict=False,
+            )
+        ):
+            with profiling_context(self, reward_func_name):
+                if isinstance(
+                    reward_func, nn.Module
+                ):  # Module instead of PretrainedModel for compat with compiled models
+                    if is_conversational(inputs[0]):
+                        messages = [
+                            {"messages": p + c}
+                            for p, c in zip(prompts, completions, strict=False)
+                        ]
+                        texts = [
+                            apply_chat_template(x, reward_processing_class)["text"]
+                            for x in messages
+                        ]
+                    else:
+                        texts = [
+                            p + c for p, c in zip(prompts, completions, strict=False)
+                        ]
+                    reward_inputs = reward_processing_class(
+                        text=texts,
+                        return_tensors="pt",
+                        padding=True,
+                        padding_side="right",
+                        add_special_tokens=False,
+                    )
+                    reward_inputs = Trainer._prepare_inputs(self, reward_inputs)
+                    with torch.inference_mode():
+                        rewards_per_func[:, i] = reward_func(**reward_inputs).logits[
+                            :, 0
+                        ]  # Shape (B*G,)
+                else:
+                    # Repeat all input columns (but "prompt" and "completion") to match the number of generations
+                    keys = [
+                        key for key in inputs[0] if key not in ["prompt", "completion"]
+                    ]
+                    reward_kwargs = {
+                        key: [example[key] for example in inputs] for key in keys
+                    }
+                    output_reward_func = reward_func(
+                        prompts=prompts, completions=completions, **reward_kwargs
+                    )
+                    # Convert None values to NaN
+                    output_reward_func = [
+                        reward if reward is not None else torch.nan
+                        for reward in output_reward_func
+                    ]
+
+                    rewards_per_func[:, i] = torch.tensor(
+                        output_reward_func, dtype=torch.float32, device=device
+                    )
+
+        # If all reward functions return None for a given row, issue a detailed warning
+        if torch.isnan(rewards_per_func).all(dim=1).any():
+            nan_row_idx = (
+                torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0]
+            )
+            row_reward_kwargs = {
+                key: value[nan_row_idx] for key, value in reward_kwargs.items()
+            }
+            row_reward_kwargs["prompt"] = prompts[nan_row_idx]
+            row_reward_kwargs["completion"] = completions[nan_row_idx]
+            warnings.warn(
+                f"All reward functions returned None for the following kwargs: {row_reward_kwargs}. "
+                "Please ensure that at least one reward function returns a valid reward.",
+                stacklevel=2,
+            )
+
+        # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the
+        # completions may be distributed across processes
+        rewards_per_func = gather(rewards_per_func)
+
+        # Apply weights to each reward function's output and sum
+        rewards = (
+            rewards_per_func * self.reward_weights.to(device).unsqueeze(0)
+        ).nansum(dim=1)
+
+        # Compute grouped-wise rewards
+        mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1)
+        std_grouped_rewards = rewards.view(-1, self.num_generations).std(dim=1)
+
+        # Normalize the rewards to compute the advantages
+        mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(
+            self.num_generations, dim=0
+        )
+        std_grouped_rewards = std_grouped_rewards.repeat_interleave(
+            self.num_generations, dim=0
+        )
+        advantages = rewards - mean_grouped_rewards
+        if self.args.scale_rewards:
+            advantages = advantages / (std_grouped_rewards + 1e-4)
+
+        # Slice to keep only the local part of the data
+        if self.args.context_parallel_size > 1:
+            # Calculate SP group ID (which group of ranks this rank belongs to)
+            sp_group_id = self.accelerator.process_index // self.local_world_size
+
+            # Calculate the start index for this SP group
+            sp_group_start = sp_group_id * len(prompts) * self.local_world_size
+
+            # All ranks in the same SP group get the same data slice
+            process_slice = slice(
+                sp_group_start,
+                sp_group_start + len(prompts),
+            )
+        else:
+            # Original behavior for non-sequence parallel case
+            process_slice = slice(
+                self.accelerator.process_index * len(prompts),
+                (self.accelerator.process_index + 1) * len(prompts),
+            )
+        advantages = advantages[process_slice]
+
+        # Log the metrics
+        if mode == "train":
+            self._total_train_tokens += (
+                self.accelerator.gather_for_metrics(attention_mask.sum()).sum().item()
+            )
+        self._metrics[mode]["num_tokens"] = [self._total_train_tokens]
+
+        # log completion lengths, mean, min, max
+        agg_completion_mask = self.accelerator.gather_for_metrics(
+            completion_mask.sum(1)
+        )
+        self._metrics[mode]["completions/mean_length"].append(
+            agg_completion_mask.float().mean().item()
+        )
+        self._metrics[mode]["completions/min_length"].append(
+            agg_completion_mask.float().min().item()
+        )
+        self._metrics[mode]["completions/max_length"].append(
+            agg_completion_mask.float().max().item()
+        )
+
+        # identify sequences that terminated with EOS and log their lengths
+        agg_terminated_with_eos = self.accelerator.gather_for_metrics(is_eos.any(dim=1))
+        term_completion_mask = agg_completion_mask[agg_terminated_with_eos]
+        clipped_completions_ratio = 1 - len(term_completion_mask) / len(
+            agg_completion_mask
+        )
+        self._metrics[mode]["completions/clipped_ratio"].append(
+            clipped_completions_ratio
+        )
+        if len(term_completion_mask) == 0:
+            # edge case where no completed sequences are found
+            term_completion_mask = torch.zeros(1, device=device)
+        self._metrics[mode]["completions/mean_terminated_length"].append(
+            term_completion_mask.float().mean().item()
+        )
+        self._metrics[mode]["completions/min_terminated_length"].append(
+            term_completion_mask.float().min().item()
+        )
+        self._metrics[mode]["completions/max_terminated_length"].append(
+            term_completion_mask.float().max().item()
+        )
+
+        # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values)
+        for i, reward_func_name in enumerate(self.reward_func_names):
+            mean_rewards = torch.nanmean(rewards_per_func[:, i]).item()
+            self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards)
+            std_rewards = nanstd(rewards_per_func[:, i]).item()
+            self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_rewards)
+        self._metrics[mode]["reward"].append(mean_grouped_rewards.mean().item())
+        self._metrics[mode]["reward_std"].append(std_grouped_rewards.mean().item())
+
+        # Log prompt and completion texts
+        self._textual_logs["prompt"].extend(gather_object(prompts_text))
+        self._textual_logs["completion"].extend(gather_object(completions_text))
+        for i, name in enumerate(self.reward_func_names):
+            self._textual_logs["rewards"][name].extend(rewards_per_func[:, i].tolist())
+
+        return {
+            "prompt_ids": prompt_ids,
+            "prompt_mask": prompt_mask,
+            "completion_ids": completion_ids,
+            "completion_mask": completion_mask,
+            "advantages": advantages,
+            "old_per_token_logps": old_per_token_logps,
+            "ref_per_token_logps": ref_per_token_logps,
+        }
diff --git a/src/axolotl/core/trainers/mamba.py b/src/axolotl/core/trainers/mamba.py
new file mode 100644
index 0000000000..dedda1b290
--- /dev/null
+++ b/src/axolotl/core/trainers/mamba.py
@@ -0,0 +1,32 @@
+"""Module for mamba trainer"""
+
+import torch
+
+from axolotl.core.trainers.base import AxolotlTrainer
+
+
+class AxolotlMambaTrainer(AxolotlTrainer):
+    """Mamba specific trainer to handle loss calculation"""
+
+    tag_names = ["axolotl", "mamba"]
+
+    def compute_loss(
+        self,
+        model,
+        inputs,
+        return_outputs=False,
+        num_items_in_batch=None,
+    ):
+        input_ids = inputs.pop("input_ids")
+        lm_logits = model(input_ids).logits
+
+        labels = input_ids.to(lm_logits.device)
+        shift_logits = lm_logits[:, :-1, :].contiguous()
+        labels = labels[:, 1:].contiguous()
+
+        loss_fct = torch.nn.CrossEntropyLoss()
+        lm_loss = loss_fct(
+            shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1)
+        )
+
+        return lm_loss
diff --git a/src/axolotl/core/trainers/mixins/__init__.py b/src/axolotl/core/trainers/mixins/__init__.py
new file mode 100644
index 0000000000..241694e443
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/__init__.py
@@ -0,0 +1,12 @@
+"""Init for axolotl.core.trainers.mixins"""
+
+# flake8: noqa
+
+from .activation_checkpointing import ActivationOffloadingMixin
+from .checkpoints import CheckpointSaveMixin
+from .layer_offloading import LayerOffloadingMixin
+from .distributed_parallel import DistributedParallelMixin
+from .optimizer import OptimizerMixin
+from .packing import PackingMixin
+from .rng_state_loader import RngLoaderMixin
+from .scheduler import SchedulerMixin
diff --git a/src/axolotl/core/trainers/mixins/activation_checkpointing.py b/src/axolotl/core/trainers/mixins/activation_checkpointing.py
new file mode 100644
index 0000000000..e635eb39d7
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/activation_checkpointing.py
@@ -0,0 +1,271 @@
+"""
+Trainer mixin for activation checkpointing w offloading
+"""
+
+import contextlib
+
+import torch
+from peft import PeftModel
+from torch import nn
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    apply_activation_checkpointing,
+)
+from torch.distributed.fsdp.wrap import ModuleWrapPolicy
+from transformers import GradientCheckpointingLayer, Trainer
+from trl.models.activation_offloading import (
+    NoOpManager,
+    OffloadActivations,
+    get_act_offloading_ctx_manager,
+)
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _patch_trl_offload_current_stream() -> None:
+    """Make TRL ``OffloadActivations`` sync its CPU<->GPU copies against the LIVE compute stream
+    instead of the once-captured ``torch.cuda.default_stream()``.
+
+    ``OffloadActivations`` caches ``self.s0 = torch.cuda.default_stream()`` in ``__init__`` and uses
+    it for every ``s0.wait_stream(s1)`` / ``s0.wait_event`` / ``record_stream(s0)``. When the actual
+    forward/backward compute runs on a NON-default stream — FSDP2, the checkpoint recompute, or any
+    custom autograd Function whose backward launches kernels (e.g. the scattermoe NVFP4 MoE) — that
+    sync targets the wrong stream, so a streamed-in activation is read before its H2D copy finishes:
+    garbage saved tensors -> NaN/illegal-memory in the backward (silent, finite forward). Querying
+    ``current_stream()`` at each use site fixes the ordering (xpu/npu already use ``current_stream``).
+    """
+    from trl.models.activation_offloading import OffloadActivations
+
+    if getattr(OffloadActivations, "_axolotl_live_stream", False):
+        return
+
+    def _live_compute_stream(self):
+        t = getattr(self, "accelerator_type", "cuda")
+        if t == "xpu":
+            return torch.xpu.current_stream()
+        if t == "npu":
+            import torch_npu  # noqa: F401
+
+            return torch.npu.current_stream()
+        return torch.cuda.current_stream()
+
+    # property read returns the live stream; the no-op setter swallows __init__'s `self.s0 = ...`.
+    OffloadActivations.s0 = property(_live_compute_stream, lambda self, _v: None)
+    OffloadActivations._axolotl_live_stream = True
+    LOG.info(
+        "Patched TRL OffloadActivations to sync against the live compute stream "
+        "(fixes streamed activation offloading under FSDP2 / custom-Function backward)"
+    )
+
+
+class ActivationOffloadingMixin(Trainer):
+    """
+    Trainer mixin class for activation checkpointing w offloading
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.args.activation_offloading:
+            # "legacy" uses the previous synchronous implementation (no CUDA
+            # streams), which keeps far fewer activations resident on-GPU than
+            # the stream-overlapped path (True/"disk"); the streams path stashes
+            # several activations to overlap copies, inflating peak reserved.
+            use_streams = self.args.activation_offloading != "legacy"
+            if use_streams:
+                _patch_trl_offload_current_stream()
+            if self.args.activation_offloading == "hidden_states":
+                from axolotl.monkeypatch.checkpoint_activation_offload import (
+                    get_checkpoint_hidden_states_offloading_ctx_manager,
+                )
+
+                self.activation_offload_context = (
+                    get_checkpoint_hidden_states_offloading_ctx_manager(
+                        use_streams=use_streams
+                    )
+                )
+            elif isinstance(self.model, PeftModel):
+                self.activation_offload_context = get_lora_act_offloading_ctx_manager(
+                    self.model, use_streams=use_streams
+                )
+            else:
+                self.activation_offload_context = get_act_offloading_ctx_manager(
+                    self.model, use_streams=use_streams
+                )
+        else:
+            self.activation_offload_context = contextlib.nullcontext()
+
+    def training_step(self, *args, **kwargs):
+        with self.activation_offload_context:
+            return super().training_step(*args, **kwargs)
+
+
+def ac_wrap_hf_model(model: nn.Module, **kwargs):
+    auto_wrap_policy = ModuleWrapPolicy(set((GradientCheckpointingLayer,)))
+    apply_activation_checkpointing(model, auto_wrap_policy=auto_wrap_policy, **kwargs)
+
+
+def get_lora_act_offloading_ctx_manager(
+    model: nn.Module,
+    use_pin_memory: bool = True,
+    use_streams: bool = True,
+    min_offload_size: int = 1024,
+    max_fwd_stash_size: int = 5,
+    warn_if_no_head: bool = True,
+) -> OffloadActivations:
+    """
+    Returns the activation offloading context manager for the model. All but the last output Linear in every step will
+    be offloaded.
+
+    If activation offloading is enabled, we return the OffloadActivations context manager. If activation offloading is
+    disabled, we return a NoOpManager context manager.
+
+    Args:
+        model (`nn.Module`):
+            Model to wrap with the activation offloading context manager.
+        use_pin_memory (`bool`, *optional*, defaults to `True`):
+            Whether to offloaded Tensor will be placed in pinned memory on the CPU. Pinned memory allows the Tensor to
+            be moved back onto GPU more quickly but is a limited resource.
+        use_streams (`bool`, *optional*, defaults to `True`):
+            Whether to use streams for performance optimization where the communications get overlapped with the
+            computation. Requires a torch build after torch-2.5.0.
+        min_offload_size (`int`, *optional*, defaults to `1024`):
+            Minimum number of bytes a Tensor must be in order to qualify for offloading. If the tensor is too small, we
+            do not want to waste bandwidth and resources moving it to CPU and back.
+        max_fwd_stash_size (`int`, *optional*, defaults to `5`):
+            Maximum size of the forward stash, or the maximum number of consecutive activations to keep alive during
+            the forward pass. This number must be at least 1. Keeping alive more activations will potentially allow
+            more overlap between the communication and compute streams at the cost of increasing memory usage. Keeping
+            alive fewer activations will conserve memory, but may cause poor overlap between the streams, increasing
+            runtime.
+        warn_if_no_head (`bool`, *optional*, defaults to `True`):
+            Whether to warn if no output head is detected. If set to `False`, no warning will be raised if no output
+            head is detected.
+
+    Returns:
+        `contextlib.ContextDecorator`:
+            Activation offloading context manager for the model.
+    """
+
+    activations_handling_ctx = OffloadActivations(
+        use_pin_memory=use_pin_memory,
+        use_streams=use_streams,
+        min_offload_size=min_offload_size,
+        max_fwd_stash_size=max_fwd_stash_size,
+    )
+
+    # Below is our hack to disable offloading the last output Linear in every
+    # step, as the cost for offloading the activation and then soon after bringing
+    # it back is expensive.
+    output_head_detected = False
+    noop_ctx = NoOpManager()
+
+    # Try to get the actual model if it's wrapped
+    unwrapped_model = model
+    if hasattr(unwrapped_model, "module"):
+        unwrapped_model = unwrapped_model.module
+    # check for PEFT models
+    if hasattr(unwrapped_model, "base_model") and hasattr(
+        unwrapped_model, "peft_config"
+    ):
+        unwrapped_model = unwrapped_model.base_model
+
+    # Check for different types of output heads
+    if hasattr(unwrapped_model, "output"):
+        if isinstance(unwrapped_model.output, nn.Module):
+            unwrapped_model.output.register_forward_pre_hook(
+                lambda *args: noop_ctx.__enter__()
+            )
+            unwrapped_model.output.register_forward_hook(
+                lambda *args: noop_ctx.__exit__(), always_call=True
+            )
+            output_head_detected = True
+        elif hasattr(unwrapped_model.output, "linear") and isinstance(
+            unwrapped_model.output.linear, nn.Module
+        ):
+            unwrapped_model.output.linear.register_forward_pre_hook(
+                lambda *args: noop_ctx.__enter__()
+            )
+            unwrapped_model.output.linear.register_forward_hook(
+                lambda *args: noop_ctx.__exit__(), always_call=True
+            )
+            output_head_detected = True
+
+    # Check for HuggingFace model output heads
+    elif hasattr(unwrapped_model, "lm_head"):
+        unwrapped_model.lm_head.register_forward_pre_hook(
+            lambda *args: noop_ctx.__enter__()
+        )
+        unwrapped_model.lm_head.register_forward_hook(
+            lambda *args: noop_ctx.__exit__(), always_call=True
+        )
+        output_head_detected = True
+
+    # Check for decoder-based models
+    elif hasattr(unwrapped_model, "decoder"):
+        decoder = unwrapped_model.decoder
+        if hasattr(decoder, "output"):
+            decoder.output.register_forward_pre_hook(lambda *args: noop_ctx.__enter__())
+            decoder.output.register_forward_hook(
+                lambda *args: noop_ctx.__exit__(), always_call=True
+            )
+            output_head_detected = True
+        # Some models have lm_head in the decoder
+        elif hasattr(decoder, "lm_head"):
+            decoder.lm_head.register_forward_pre_hook(
+                lambda *args: noop_ctx.__enter__()
+            )
+            decoder.lm_head.register_forward_hook(
+                lambda *args: noop_ctx.__exit__(), always_call=True
+            )
+            output_head_detected = True
+
+    # Check for transformer models with final layer norm
+    elif hasattr(unwrapped_model, "final_layer_norm") or hasattr(
+        unwrapped_model, "ln_f"
+    ):
+        final_norm = (
+            getattr(unwrapped_model, "final_layer_norm", None) or unwrapped_model.ln_f
+        )
+        final_norm.register_forward_pre_hook(lambda *args: noop_ctx.__enter__())
+        final_norm.register_forward_hook(
+            lambda *args: noop_ctx.__exit__(), always_call=True
+        )
+        output_head_detected = True
+
+    # Check for models with head module
+    elif hasattr(unwrapped_model, "head") and isinstance(
+        unwrapped_model.head, nn.Module
+    ):
+        unwrapped_model.head.register_forward_pre_hook(
+            lambda *args: noop_ctx.__enter__()
+        )
+        unwrapped_model.head.register_forward_hook(
+            lambda *args: noop_ctx.__exit__(), always_call=True
+        )
+        output_head_detected = True
+
+    if not output_head_detected and warn_if_no_head:
+        LOG.warning(
+            "During activation offloading, no output head was detected. If your model has an output head, it will be "
+            "offloaded. This usually greatly slows training, given the large vocabulary size. To change this "
+            "behavior, set your output head as model.output and make it an nn.Module. You can disable this warning by "
+            "passing `warn_if_no_head=False`."
+        )
+
+    for name, module in unwrapped_model.named_modules():
+        # Disable offloading for any Liger modules
+        if "liger" in name.lower():
+            module.register_forward_pre_hook(lambda *args: noop_ctx.__enter__())
+            module.register_forward_hook(
+                lambda *args: noop_ctx.__exit__(), always_call=True
+            )
+        # disable offloading for any submodules to fix LoRA training
+        if name.endswith("._checkpoint_wrapped_module"):
+            for _, sub_module in module.named_modules():
+                sub_module.register_forward_pre_hook(lambda *args: noop_ctx.__enter__())
+                sub_module.register_forward_hook(
+                    lambda *args: noop_ctx.__exit__(), always_call=True
+                )
+
+    return activations_handling_ctx
diff --git a/src/axolotl/core/trainers/mixins/checkpoints.py b/src/axolotl/core/trainers/mixins/checkpoints.py
new file mode 100644
index 0000000000..fe13cd8ac2
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/checkpoints.py
@@ -0,0 +1,22 @@
+"""Custom handling to not fail training if fsdp optimizer is not savable"""
+
+from transformers import Trainer
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class CheckpointSaveMixin(Trainer):
+    """Mixin to handle saving the optimizer and scheduler if they are not savable."""
+
+    def _save_optimizer_and_scheduler(self, output_dir):
+        try:
+            super()._save_optimizer_and_scheduler(output_dir)
+        except (NotImplementedError, KeyError) as exc:
+            # TODO: fix fsdp2 optimizer saving
+            LOG.warning_once(
+                f"Trainer does not support saving optimizer and scheduler:  {exc}\n"
+                "Optimizer and scheduler states were not saved - resuming from checkpoints "
+                "for this training run will not be possible.",
+            )
diff --git a/src/axolotl/core/trainers/mixins/distributed_parallel.py b/src/axolotl/core/trainers/mixins/distributed_parallel.py
new file mode 100644
index 0000000000..77aee5236b
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/distributed_parallel.py
@@ -0,0 +1,32 @@
+"""
+Mixin for correctly saving fsdp
+"""
+
+from accelerate import PartialState
+from transformers import Trainer
+
+
+class DistributedParallelMixin(Trainer):
+    """
+    Mixin for correctly saving fsdp
+    """
+
+    def _save(self, output_dir: str | None = None, state_dict=None):
+        if (
+            state_dict is None
+            and self.accelerator.parallelism_config
+            and self.accelerator.parallelism_config.dp_shard_enabled
+        ):
+            state_dict = self.accelerator.get_state_dict(self.model)
+        super()._save(output_dir, state_dict=state_dict)
+
+    def create_accelerator_and_postprocess(self):
+        super().create_accelerator_and_postprocess()
+        if (
+            self.accelerator.distributed_type == "FSDP"
+            and self.accelerator.state.fsdp_plugin is None
+        ):
+            # handle Context Parallelism without FSDP
+            self.accelerator.state.distributed_type = "MULTI_GPU"
+            self.accelerator.state._shared_state["distributed_type"] = "MULTI_GPU"
+            PartialState().distributed_type = "MULTI_GPU"
diff --git a/src/axolotl/core/trainers/mixins/layer_offloading.py b/src/axolotl/core/trainers/mixins/layer_offloading.py
new file mode 100644
index 0000000000..83a9feff58
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/layer_offloading.py
@@ -0,0 +1,304 @@
+"""
+Trainer mixin for layer-wise parameter offloading to CPU.
+
+Offloads frozen (non-trainable) parameters in decoder layers to CPU, then uses
+forward/backward hooks to stream them on/off GPU one layer at a time with CUDA
+stream prefetching. Trainable parameters (e.g. LoRA weights) stay on GPU always.
+
+Forward:  pre-hook loads layer N's frozen params to GPU (prefetches N+1 on
+          transfer stream), post-hook offloads layer N-1's frozen params.
+Backward: same in reverse order.
+"""
+
+import contextlib
+
+import torch
+import torch.nn as nn
+from transformers import Trainer
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _find_decoder_layers(model: nn.Module) -> tuple[nn.ModuleList | None, list[str]]:
+    """Recursively search the model for the decoder layer ModuleList.
+
+    Finds any ModuleList whose children have 'DecoderLayer' in their class name.
+    Handles all common HF architectures including VLM wrappers (e.g. Qwen3.5-MoE
+    where layers are at model.language_model.layers).
+    """
+    # BFS to find the first ModuleList containing decoder layers
+    queue = [model]
+    while queue:
+        m = queue.pop(0)
+        for _name, child in m.named_children():
+            if isinstance(child, nn.ModuleList) and len(child) > 0:
+                first_type = type(child[0]).__name__
+                if "DecoderLayer" in first_type or "TransformerBlock" in first_type:
+                    layer_types = list({type(layer).__name__ for layer in child})
+                    return child, layer_types
+            else:
+                queue.append(child)
+
+    return None, []
+
+
+def _get_frozen_params(layer: nn.Module) -> list[tuple[str, nn.Parameter]]:
+    """Get all non-trainable parameters in a layer."""
+    return [(n, p) for n, p in layer.named_parameters() if not p.requires_grad]
+
+
+class LayerOffloadManager:
+    """Manages offloading frozen decoder layer params to CPU and streaming
+    them back during forward/backward with CUDA stream overlap.
+
+    Only frozen (requires_grad=False) parameters are offloaded.
+    Trainable parameters (LoRA weights, etc.) remain on GPU at all times.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        num_prefetch: int = 1,
+    ):
+        self.model = model
+        self.num_prefetch = num_prefetch
+        self._hooks: list = []
+        self._device = None
+
+        # Find decoder layers
+        self.layers, layer_types = _find_decoder_layers(model)
+        if self.layers is None:
+            LOG.warning(
+                "LayerOffloadManager: no decoder layers found, offloading disabled"
+            )
+            self.enabled = False
+            return
+
+        self.enabled = True
+        self.n_layers = len(self.layers)
+        LOG.info(
+            f"Layer offloading: found {self.n_layers} layers ({', '.join(layer_types)})"
+        )
+
+        # Determine GPU device
+        for p in model.parameters():
+            if p.device.type == "cuda":
+                self._device = p.device
+                break
+        if self._device is None:
+            LOG.warning("LayerOffloadManager: no CUDA parameters found")
+            self.enabled = False
+            return
+
+        # Transfer stream for async prefetch
+        self._transfer_stream = torch.cuda.Stream(device=self._device)
+
+        # Track which layers have their frozen params on GPU
+        self._on_gpu: set[int] = set(range(self.n_layers))
+
+        # Cache: frozen param references per layer (list of (name, param) tuples)
+        self._frozen_params: list[list[tuple[str, nn.Parameter]]] = [
+            _get_frozen_params(self.layers[i]) for i in range(self.n_layers)
+        ]
+
+        # CPU storage: pinned tensors for each layer's frozen params
+        # Populated on first offload
+        self._cpu_data: list[dict[str, torch.Tensor]] = [
+            {} for _ in range(self.n_layers)
+        ]
+
+        # Offload all layers upfront
+        self._offload_all()
+
+        # Release cached memory blocks back to the driver
+        torch.cuda.empty_cache()
+
+    def _offload_all(self):
+        """Move all frozen params in all decoder layers to CPU."""
+        mem_before = torch.cuda.memory_allocated(self._device)
+        for i in range(self.n_layers):
+            self._offload_layer(i)
+        mem_after = torch.cuda.memory_allocated(self._device)
+        freed = (mem_before - mem_after) / 1e6
+        LOG.info(
+            f"Layer offloading: offloaded frozen params from {self.n_layers} layers, "
+            f"freed {freed:.0f} MB GPU memory"
+        )
+
+    def _offload_layer(self, idx: int):
+        """Move frozen params of layer idx to CPU pinned memory."""
+        if idx not in self._on_gpu:
+            return
+        for name, param in self._frozen_params[idx]:
+            if param.device.type != "cuda":
+                continue
+            # Allocate pinned CPU tensor on first offload
+            if name not in self._cpu_data[idx]:
+                self._cpu_data[idx][name] = torch.empty_like(
+                    param.data, device="cpu", pin_memory=True
+                )
+            cpu_buf = self._cpu_data[idx][name]
+            # Async copy GPU -> CPU (on transfer stream for overlap)
+            cpu_buf.copy_(param.data, non_blocking=True)
+            # Point parameter at a dummy CPU tensor to free GPU memory
+            param.data = cpu_buf
+        self._on_gpu.discard(idx)
+
+    def _load_layer(self, idx: int, stream=None):
+        """Move frozen params of layer idx back to GPU."""
+        if idx in self._on_gpu or idx < 0 or idx >= self.n_layers:
+            return
+        ctx = (
+            torch.cuda.stream(stream)
+            if stream is not None
+            else contextlib.nullcontext()
+        )
+        with ctx:
+            for _name, param in self._frozen_params[idx]:
+                if param.device.type == "cuda":
+                    continue
+                gpu_data = param.data.to(self._device, non_blocking=True)
+                param.data = gpu_data
+        self._on_gpu.add(idx)
+
+    def _prefetch_layer(self, idx: int):
+        """Async prefetch layer idx on the transfer stream."""
+        if idx in self._on_gpu or idx < 0 or idx >= self.n_layers:
+            return
+        self._transfer_stream.wait_stream(torch.cuda.default_stream(self._device))
+        self._load_layer(idx, stream=self._transfer_stream)
+
+    def _wait_transfer(self):
+        """Make default stream wait for any in-flight transfers."""
+        torch.cuda.default_stream(self._device).wait_stream(self._transfer_stream)
+
+    def setup_hooks(self):
+        """Register forward and backward hooks on each decoder layer."""
+        if not self.enabled:
+            return
+
+        for idx in range(self.n_layers):
+            layer = self.layers[idx]
+
+            def make_pre_fwd(i):
+                def hook(module, args):
+                    # Ensure this layer is on GPU
+                    if i not in self._on_gpu:
+                        self._load_layer(i)
+                    self._wait_transfer()
+                    # Prefetch next layer(s)
+                    for offset in range(1, self.num_prefetch + 1):
+                        self._prefetch_layer(i + offset)
+
+                return hook
+
+            def make_post_fwd(i):
+                def hook(module, args, output):
+                    # Offload previous layer (no longer needed in forward)
+                    if i > 0:
+                        self._offload_layer(i - 1)
+                    # Offload last layer after forward
+                    if i == self.n_layers - 1:
+                        self._offload_layer(i)
+
+                return hook
+
+            def make_pre_bwd(i):
+                def hook(module, grad_output):
+                    # Load this layer for backward
+                    if i not in self._on_gpu:
+                        self._load_layer(i)
+                    self._wait_transfer()
+                    # Prefetch previous layer(s)
+                    for offset in range(1, self.num_prefetch + 1):
+                        self._prefetch_layer(i - offset)
+
+                return hook
+
+            def make_post_bwd(i):
+                def hook(module, grad_input, grad_output):
+                    # Offload the layer above
+                    if i < self.n_layers - 1:
+                        self._offload_layer(i + 1)
+                    # Offload first layer after backward
+                    if i == 0:
+                        self._offload_layer(i)
+
+                return hook
+
+            h1 = layer.register_forward_pre_hook(make_pre_fwd(idx))
+            h2 = layer.register_forward_hook(make_post_fwd(idx))
+            h3 = layer.register_full_backward_pre_hook(make_pre_bwd(idx))
+            h4 = layer.register_full_backward_hook(make_post_bwd(idx))
+            self._hooks.extend([h1, h2, h3, h4])
+
+    def remove_hooks(self):
+        """Remove all hooks and restore layers to GPU."""
+        for h in self._hooks:
+            h.remove()
+        self._hooks.clear()
+        if self.enabled:
+            for i in range(self.n_layers):
+                if i not in self._on_gpu:
+                    self._load_layer(i)
+
+    def pre_step(self):
+        """Called before each training step — ensure layers start offloaded."""
+        if not self.enabled:
+            return
+        for i in list(self._on_gpu):
+            self._offload_layer(i)
+        # Prefetch layer 0 for forward
+        self._prefetch_layer(0)
+
+    def post_step(self):
+        """Called after each training step — ensure layers are offloaded."""
+        if not self.enabled:
+            return
+        for i in list(self._on_gpu):
+            self._offload_layer(i)
+        # Prefetch layer 0 for next step
+        self._prefetch_layer(0)
+
+
+class _LayerOffloadContext:
+    """Context manager wrapping pre_step / post_step around a training step."""
+
+    def __init__(self, manager: LayerOffloadManager):
+        self.manager = manager
+
+    def __enter__(self):
+        self.manager.pre_step()
+        return self
+
+    def __exit__(self, *args):
+        self.manager.post_step()
+
+
+class LayerOffloadingMixin(Trainer):
+    """
+    Trainer mixin class for layer-wise parameter offloading to CPU.
+
+    Offloads frozen decoder layer params to CPU at init, then streams them
+    on/off GPU one layer at a time during each training step.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if getattr(self.args, "layer_offloading", False):
+            LOG.info("Layer parameter offloading enabled")
+            self._layer_offload_manager = LayerOffloadManager(
+                model=self.model,
+                num_prefetch=1,
+            )
+            self._layer_offload_manager.setup_hooks()
+            self._layer_offload_ctx = _LayerOffloadContext(self._layer_offload_manager)
+        else:
+            self._layer_offload_manager = None
+            self._layer_offload_ctx = contextlib.nullcontext()
+
+    def training_step(self, *args, **kwargs):
+        with self._layer_offload_ctx:
+            return super().training_step(*args, **kwargs)
diff --git a/src/axolotl/core/trainers/mixins/optimizer.py b/src/axolotl/core/trainers/mixins/optimizer.py
new file mode 100644
index 0000000000..985b756219
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/optimizer.py
@@ -0,0 +1,216 @@
+"""Module for Axolotl trainer optimizer mixin"""
+
+from peft.optimizers import create_loraplus_optimizer
+from torch import nn
+from transformers.trainer import Trainer
+from transformers.trainer_optimizer import is_optimizer_factory
+from transformers.utils import is_sagemaker_mp_enabled
+
+from axolotl.utils.logging import get_logger
+
+if is_sagemaker_mp_enabled():
+    import smdistributed.modelparallel.torch as smp
+
+LOG = get_logger(__name__)
+
+
+class OptimizerMixin(Trainer):
+    """Mixin class for shared handling of building custom optimizers"""
+
+    args = None  # type: "AxolotlTrainingArguments"  # type: ignore[name-defined]
+
+    def create_optimizer_grouped_parameters(
+        self, opt_model, optimizer_kwargs
+    ) -> list[dict]:
+        decay_parameters = self.get_decay_parameter_names(opt_model)
+        params: dict = {
+            "to_weight_decay": {},  # LayerNorm and bias
+            "embeddings": {},  # lm_head, embed_tokens,
+            "no_weight_decay": {},
+        }
+        lr_groups_lookup = {}
+        lr_groups_learning_rates = {}
+        if self.args.lr_groups:
+            for lr_group in self.args.lr_groups:
+                group_name = lr_group["name"]
+                group_modules = lr_group["modules"]
+                for module in group_modules:
+                    lr_groups_lookup[module] = group_name
+                lr_groups_learning_rates[group_name] = lr_group["lr"]
+                params[f"to_weight_decay_{group_name}"] = {}
+
+        for name, param in opt_model.named_parameters():
+            if not param.requires_grad:
+                continue
+            if name.endswith("modules_to_save.default.weight") or any(
+                embed_name in name for embed_name in ["embed_tokens", "lm_head"]
+            ):
+                params["embeddings"][name] = param
+            elif name in decay_parameters:
+                lr_group_modules = [
+                    group_modules
+                    for group_modules in lr_groups_lookup
+                    if group_modules in name
+                ]
+                if lr_groups_lookup and any(lr_group_modules):
+                    lr_group_module = lr_group_modules[0]
+                    group_name = lr_groups_lookup[lr_group_module]
+                    params[f"to_weight_decay_{group_name}"][name] = param
+                else:
+                    params["to_weight_decay"][name] = param
+            else:
+                params["no_weight_decay"][name] = param
+        optimizer_grouped_parameters = []
+        if params["to_weight_decay"]:
+            optimizer_grouped_parameters.append(
+                {
+                    "params": list(params["to_weight_decay"].values()),
+                    "weight_decay": self.args.weight_decay,
+                    "lr": optimizer_kwargs["lr"],
+                }
+            )
+        if params["embeddings"]:
+            lr = optimizer_kwargs["lr"]
+            if self.args.embedding_lr_scale:
+                lr *= self.args.embedding_lr_scale
+            elif self.args.embedding_lr:
+                lr = self.args.embedding_lr
+            optimizer_grouped_parameters.append(
+                {
+                    "params": list(params["embeddings"].values()),
+                    "weight_decay": 0.0,
+                    "lr": lr,
+                }
+            )
+        if params["no_weight_decay"]:
+            optimizer_grouped_parameters.append(
+                {
+                    "params": list(params["no_weight_decay"].values()),
+                    "weight_decay": 0.0,
+                    "lr": optimizer_kwargs["lr"],
+                }
+            )
+        for group_name, group_lr in lr_groups_learning_rates.items():
+            if params[f"to_weight_decay_{group_name}"]:
+                optimizer_grouped_parameters.append(
+                    {
+                        "params": list(
+                            params[f"to_weight_decay_{group_name}"].values()
+                        ),
+                        "weight_decay": self.args.weight_decay,
+                        "lr": group_lr,
+                    }
+                )
+
+        return optimizer_grouped_parameters
+
+    def create_optimizer(self, model=None):
+        if (
+            self.args.loraplus_lr_ratio is None
+            and self.args.embedding_lr_scale is None
+            and self.args.embedding_lr is None
+            and self.args.lr_groups is None
+            and self.optimizer_cls_and_kwargs is None
+        ):
+            return super().create_optimizer(model=model)
+
+        opt_model = self.model if model is None else model
+
+        if (
+            not self.optimizer
+            and self.optimizer_cls_and_kwargs is not None
+            and is_optimizer_factory(self.optimizer_cls_and_kwargs[0])
+        ):
+            optimizer_factory_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
+            # forward training_args for back-compat with our factory signatures
+            self.optimizer = optimizer_factory_cls()(
+                opt_model, self.args, **optimizer_kwargs
+            )
+
+        if not self.optimizer:
+            if self.optimizer_cls_and_kwargs is not None:
+                optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
+            else:
+                optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(
+                    self.args, opt_model
+                )
+
+            optimizer_grouped_parameters = self.create_optimizer_grouped_parameters(
+                opt_model, optimizer_kwargs
+            )
+
+            if self.args.loraplus_lr_ratio is not None:
+                loraplus_lr_ratio = getattr(self.args, "loraplus_lr_ratio", None)
+                loraplus_lr_embedding = getattr(
+                    self.args, "loraplus_lr_embedding", 1e-6
+                )
+                self.optimizer = create_loraplus_optimizer(
+                    opt_model,
+                    optimizer_cls,
+                    loraplus_lr_ratio=loraplus_lr_ratio,
+                    loraplus_lr_embedding=loraplus_lr_embedding,
+                    **optimizer_kwargs,
+                )
+            else:
+                # Overwrite `params` in case it's created by `get_optimizer_cls_and_kwargs`
+                # e.g. for GaLore optimizer.
+                if "params" in optimizer_kwargs:
+                    optimizer_grouped_parameters = optimizer_kwargs.pop("params")
+
+                # Overwrite `model` in case it's created by `get_optimizer_cls_and_kwargs`
+                # e.g. for LOMO optimizer.
+                if "model" in optimizer_kwargs:
+                    optimizer_grouped_parameters = optimizer_kwargs.pop("model")
+
+                # For layer-wise dummy optimizers we overwrite optimizer_grouped_parameters with `optimizer_dict`
+                # to avoid arguments conflicts.
+                if "optimizer_dict" in optimizer_kwargs:
+                    optimizer_grouped_parameters = optimizer_kwargs.pop(
+                        "optimizer_dict"
+                    )
+
+                self.optimizer = optimizer_cls(
+                    optimizer_grouped_parameters, **optimizer_kwargs
+                )
+
+            if optimizer_cls.__name__ == "Adam8bit":
+                import bitsandbytes
+
+                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
+
+                skipped = 0
+                for module in opt_model.modules():
+                    if isinstance(module, nn.Embedding):
+                        skipped += sum(
+                            {
+                                p.data_ptr(): p.numel() for p in module.parameters()
+                            }.values()
+                        )
+                        LOG.info(f"skipped {module}: {skipped / 2**20}M params")
+                        manager.register_module_override(
+                            module, "weight", {"optim_bits": 32}
+                        )
+                        LOG.debug(f"bitsandbytes: will optimize {module} in fp32")
+                LOG.info(f"skipped: {skipped / 2**20}M params")
+
+        if is_sagemaker_mp_enabled():
+            self.optimizer = smp.DistributedOptimizer(self.optimizer)
+
+        return self.optimizer
+
+
+class OptimizerInitMixin:
+    """
+    Mixin to handle common optimizer initialization logic for Trainers (mostly TRL) that do not
+    accept optimizer_cls_and_kwargs as kwarg in constructor.
+    """
+
+    def __init__(self, *args, **kwargs):
+        optimizer_cls_and_kwargs = kwargs.pop("optimizer_cls_and_kwargs", None)
+        super().__init__(*args, **kwargs)
+        if (
+            optimizer_cls_and_kwargs
+            and self.optimizer_cls_and_kwargs is None
+            and self.optimizer is None
+        ):
+            self.optimizer_cls_and_kwargs = optimizer_cls_and_kwargs
diff --git a/src/axolotl/core/trainers/mixins/packing.py b/src/axolotl/core/trainers/mixins/packing.py
new file mode 100644
index 0000000000..249ceeb4f8
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/packing.py
@@ -0,0 +1,20 @@
+"""Trainer mixin to support packing"""
+
+from transformers import Trainer
+
+
+class PackingMixin(Trainer):
+    """
+    Trainer mixin to support packing
+    """
+
+    def _set_signature_columns_if_needed(self):
+        super()._set_signature_columns_if_needed()
+        if (
+            self._signature_columns
+            and self.args.sample_packing
+            and self.args.sample_packing_drop_attention_mask
+        ):
+            set_sig_columns = set(self._signature_columns)
+            set_sig_columns.remove("attention_mask")
+            self._signature_columns = list(set_sig_columns)
diff --git a/src/axolotl/core/trainers/mixins/rng_state_loader.py b/src/axolotl/core/trainers/mixins/rng_state_loader.py
new file mode 100644
index 0000000000..f248394b2e
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/rng_state_loader.py
@@ -0,0 +1,68 @@
+"""
+Temporary fix/override for bug in resume from checkpoint
+
+See https://github.com/huggingface/transformers/pull/37162
+
+TODO: Remove when upstream added PR to release
+"""
+
+import os
+import random
+
+import numpy as np
+import torch
+from transformers import Trainer, is_torch_npu_available
+from transformers.trainer import safe_globals
+from transformers.trainer_pt_utils import set_rng_state_for_device
+from transformers.training_args import ParallelMode
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class RngLoaderMixin(Trainer):
+    """
+    mixin for method override to load RNG states from a checkpoint
+    """
+
+    def _load_rng_state(self, checkpoint):
+        # Load RNG states from `checkpoint`
+        if checkpoint is None:
+            return
+
+        if self.args.world_size > 1:
+            process_index = self.args.process_index
+            rng_file = os.path.join(checkpoint, f"rng_state_{process_index}.pth")
+            if not os.path.isfile(rng_file):
+                LOG.info(
+                    f"Didn't find an RNG file for process {process_index}, if you are resuming a training that "
+                    "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
+                )
+                return
+        else:
+            rng_file = os.path.join(checkpoint, "rng_state.pth")
+            if not os.path.isfile(rng_file):
+                LOG.info(
+                    "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
+                    "fashion, reproducibility is not guaranteed."
+                )
+                return
+
+        # Use safe_globals to ensure numpy RNG states can be deserialized safely under PyTorch 2.6+,
+        # which requires allowlisted classes when loading with weights_only=True.
+        with safe_globals():
+            checkpoint_rng_state = torch.load(rng_file)  # nosec B614
+        random.setstate(checkpoint_rng_state["python"])
+        np.random.set_state(checkpoint_rng_state["numpy"])
+        torch.random.set_rng_state(checkpoint_rng_state["cpu"])
+
+        is_distributed = self.args.parallel_mode == ParallelMode.DISTRIBUTED
+        if torch.cuda.is_available():
+            set_rng_state_for_device(
+                "CUDA", torch.cuda, checkpoint_rng_state, is_distributed
+            )
+        if is_torch_npu_available():
+            set_rng_state_for_device(
+                "NPU", torch.npu, checkpoint_rng_state, is_distributed
+            )
diff --git a/src/axolotl/core/trainers/mixins/scheduler.py b/src/axolotl/core/trainers/mixins/scheduler.py
new file mode 100644
index 0000000000..23e4f6c255
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/scheduler.py
@@ -0,0 +1,152 @@
+"""Module for Axolotl trainer scheduler mixin"""
+
+import torch
+from torch.optim.lr_scheduler import LRScheduler, OneCycleLR
+from transformers.trainer import Trainer
+
+from axolotl.integrations.base import PluginManager
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schedulers import (
+    JaggedLRRestartScheduler,
+    RexLR,
+    get_cosine_schedule_with_min_lr,
+    get_cosine_schedule_with_quadratic_warmup,
+    get_cosine_schedule_with_warmup_decay_constant,
+)
+
+LOG = get_logger(__name__)
+
+
+class SchedulerMixin(Trainer):
+    """
+    Mixin class for scheduler setup in CausalTrainer.
+    """
+
+    args = None  # type: "AxolotlTrainingArguments"  # type: ignore[name-defined]
+
+    def create_scheduler(
+        self, num_training_steps: int, optimizer: None | torch.optim.Optimizer = None
+    ) -> LRScheduler:
+        """
+        Set up the scheduler. The optimizer of the trainer must have been set up either before this method is called or
+        passed as an argument.
+
+        Args:
+            num_training_steps (int): The number of training steps to do.
+            optimizer (torch.optim.Optimizer): The training optimizer
+        """
+        use_cosine_quadratic = (
+            self.args.lr_scheduler_type == "cosine"
+            and self.args.lr_quadratic_warmup is True
+        )
+
+        use_cosine_min_lr = (
+            self.args.lr_scheduler_type == "cosine"
+            and self.args.cosine_min_lr_ratio is not None
+        )
+
+        if optimizer is None:
+            if self.optimizer is None:
+                raise ValueError(
+                    "Optimizer must be set before calling create_scheduler or passed as an argument."
+                )
+            optimizer = self.optimizer
+
+        # fmt: off
+        if self.lr_scheduler is None:  # type: ignore
+            # fmt: on
+            plugin_manager = PluginManager.get_instance()
+            lr_scheduler: LRScheduler | None = plugin_manager.create_lr_scheduler(
+                trainer=self,
+                optimizer=optimizer,
+                num_training_steps=num_training_steps
+            )
+            if lr_scheduler is not None:
+                LOG.info(f"Using plugin-created lr_scheduler: {lr_scheduler}")
+                self.lr_scheduler = lr_scheduler
+            elif self.args.alternate_lr_scheduler_type == "one_cycle":
+                num_warmup_steps = self.args.get_warmup_steps(num_training_steps)
+                pct_start = num_warmup_steps / num_training_steps
+                extra_lr_kwargs = {}
+                if "pct_start" not in self.args.lr_scheduler_kwargs:
+                    extra_lr_kwargs["pct_start"] = pct_start
+                if "anneal_strategy" not in self.args.lr_scheduler_kwargs:
+                    extra_lr_kwargs["anneal_strategy"] = "cos"
+
+                self.lr_scheduler = OneCycleLR(
+                    optimizer,
+                    max_lr=self.args.learning_rate,
+                    total_steps=num_training_steps,
+                    **extra_lr_kwargs,
+                    **self.args.lr_scheduler_kwargs,
+                )
+            elif self.args.alternate_lr_scheduler_type == "rex":
+                if use_cosine_min_lr:
+                    assert 0 <= self.args.cosine_min_lr_ratio <= 1.0, "cosine_min_lr_ratio must be between 0.0 and 1.0"
+
+                self.lr_scheduler = RexLR(
+                    optimizer=optimizer,
+                    max_lr=self.args.learning_rate,
+                    min_lr=0 if not use_cosine_min_lr else (
+                        self.args.learning_rate * self.args.cosine_min_lr_ratio),
+                    total_steps=num_training_steps,
+                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
+                )
+            elif use_cosine_quadratic:
+                if use_cosine_min_lr:
+                    LOG.warning(
+                        "Both cosine quadratic warmup and min lr detected. Using quadratic warmup.")
+
+                self.lr_scheduler = get_cosine_schedule_with_quadratic_warmup(
+                    optimizer,
+                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
+                    num_training_steps=num_training_steps,
+                )
+            elif self.args.cosine_min_lr_ratio and self.args.cosine_constant_lr_ratio and use_cosine_min_lr:
+                assert 0 <= self.args.cosine_min_lr_ratio <= 1.0, "cosine_min_lr_ratio must be between 0.0 and 1.0"
+                assert 0 <= self.args.cosine_constant_lr_ratio <= 1.0, "cosine_constant_lr_ratio must be between 0.0 and 1.0"
+                self.lr_scheduler = get_cosine_schedule_with_warmup_decay_constant(
+                    optimizer,
+                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
+                    num_training_steps=num_training_steps,
+                    min_lr_ratio=self.args.cosine_min_lr_ratio,
+                    constant_lr_ratio=self.args.cosine_constant_lr_ratio,
+                )
+            elif self.args.cosine_min_lr_ratio and use_cosine_min_lr:
+                assert 0 <= self.args.cosine_min_lr_ratio <= 1.0, "cosine_min_lr_ratio must be between 0.0 and 1.0"
+                self.lr_scheduler = get_cosine_schedule_with_min_lr(
+                    optimizer,
+                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
+                    num_training_steps=num_training_steps,
+                    min_lr_ratio=self.args.cosine_min_lr_ratio,
+                )
+            else:
+                super().create_scheduler(num_training_steps, optimizer=optimizer)
+        else:
+            if use_cosine_quadratic:
+                LOG.warning(
+                    "axolotl's cosine scheduler with quadratic warmup not used (e.g., because of deepspeed).")
+
+            if use_cosine_min_lr:
+                LOG.warning(
+                    "axolotl's cosine scheduler with min lr not used (e.g., because of deepspeed).")
+
+        if self.args.jagged_restart_steps:
+            warmup_steps = (
+                self.args.jagged_restart_warmup_steps or 10
+            )
+            anneal_steps = (
+                self.args.jagged_restart_anneal_steps or 1
+            )
+            if not self.lr_scheduler:
+                super().create_scheduler(num_training_steps, optimizer)
+            self.lr_scheduler = JaggedLRRestartScheduler(
+                optimizer,
+                self.lr_scheduler,
+                self.args.jagged_restart_steps,
+                warmup_steps,
+                anneal_steps,
+                min_lr_scale=self.args.cosine_min_lr_ratio or 0.001,
+            )
+
+        return self.lr_scheduler  # type: ignore
diff --git a/src/axolotl/core/trainers/trl.py b/src/axolotl/core/trainers/trl.py
index 24c0b04123..bc49754bea 100644
--- a/src/axolotl/core/trainers/trl.py
+++ b/src/axolotl/core/trainers/trl.py
@@ -1,66 +1,86 @@
-"""
-module for TRL PPO training
-"""
-import torch
-from tqdm import tqdm
-from trl import PPOTrainer
+"""Module for TRL RL trainers"""
 
+from trl import RewardTrainer
+from trl.experimental.cpo import CPOTrainer
+from trl.experimental.kto import KTOTrainer
+from trl.experimental.orpo import ORPOTrainer
+from trl.experimental.prm import PRMTrainer
 
-class TRLPPOTrainer(PPOTrainer):
+from axolotl.core.trainers.mixins import DistributedParallelMixin, RngLoaderMixin
+from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin
+from axolotl.core.trainers.mixins.scheduler import SchedulerMixin
+
+
+class AxolotlORPOTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DistributedParallelMixin,
+    ORPOTrainer,
+):
+    """
+    Extend the base ORPOTrainer for axolotl helpers
+    """
+
+    tag_names = ["axolotl", "orpo"]
+
+
+class AxolotlKTOTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DistributedParallelMixin,
+    KTOTrainer,
+):
+    """
+    Extend the base KTOTrainer for axolotl helpers
+    """
+
+    tag_names = ["axolotl", "kto"]
+
+
+class AxolotlCPOTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DistributedParallelMixin,
+    CPOTrainer,
+):
+    """
+    Extend the base CPOTrainer for axolotl helpers
+    """
+
+    tag_names = ["axolotl", "cpo"]
+
+
+class AxolotlRewardTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DistributedParallelMixin,
+    RewardTrainer,
+):
+    """
+    Extend the base RewardTrainer for axolotl helpers
+    """
+
+    tag_names = ["axolotl", "reward"]
+
+
+class AxolotlPRMTrainer(
+    RngLoaderMixin,
+    SchedulerMixin,
+    OptimizerMixin,
+    OptimizerInitMixin,
+    DistributedParallelMixin,
+    PRMTrainer,
+):
     """
-    wrapper for ppo trainer to handle customizations
+    Extend the base trl.PRMTrainer for axolotl helpers
     """
 
-    def train(
-        self,
-        reward_pipe,
-        resume_from_checkpoint=None,  # pylint: disable=unused-argument
-    ):
-        generation_kwargs = {
-            "min_length": -1,
-            "top_k": 0.0,
-            "top_p": 1.0,
-            "do_sample": True,
-            "pad_token_id": self.tokenizer.eos_token_id,
-            "max_new_tokens": 32,
-        }
-        sent_kwargs = {
-            "return_all_scores": True,
-            "function_to_apply": "none",
-            "batch_size": 16,
-        }
-
-        for epoch, batch in tqdm(  # pylint: disable=unused-variable
-            enumerate(self.dataloader)
-        ):
-            query_tensors = batch["input_ids"]
-
-            # generate model response
-            response_tensors, ref_response_tensors = self.generate(
-                query_tensors,
-                return_prompt=False,
-                generate_ref_response=True,
-                **generation_kwargs
-            )
-            batch["response"] = self.tokenizer.batch_decode(response_tensors)
-            batch["ref_response"] = self.tokenizer.batch_decode(ref_response_tensors)
-
-            # Compute sentiment score
-            texts = [q + r for q, r in zip(batch["query"], batch["response"])]
-            pipe_outputs = reward_pipe(texts, **sent_kwargs)
-            rewards = [torch.tensor(output[1]["score"]) for output in pipe_outputs]
-            ref_texts = [q + r for q, r in zip(batch["query"], batch["ref_response"])]
-            ref_pipe_outputs = reward_pipe(ref_texts, **sent_kwargs)
-            ref_rewards = [
-                torch.tensor(output[1]["score"]) for output in ref_pipe_outputs
-            ]
-            batch["ref_rewards"] = ref_rewards
-
-            # Run PPO step
-            stats = self.step(query_tensors, response_tensors, rewards)
-            self.log_stats(
-                stats,
-                batch,
-                rewards,
-                columns_to_log=["query", "response", "ref_response", "ref_rewards"],
-            )
+    tag_names = ["axolotl", "prm"]
diff --git a/src/axolotl/core/trainers/utils.py b/src/axolotl/core/trainers/utils.py
new file mode 100644
index 0000000000..55959c6a41
--- /dev/null
+++ b/src/axolotl/core/trainers/utils.py
@@ -0,0 +1,51 @@
+"""Utils for Axolotl trainers"""
+
+
+def trainable_tokens_per_sec_per_gpu(
+    prev_trainable: float | None,
+    curr_trainable: float,
+    world_size: int,
+    elapsed: float,
+) -> float | None:
+    """Effective per-GPU trainable-token throughput over a logging window.
+
+    ``curr_trainable``/``prev_trainable`` are the cumulative trainable-token
+    counter (SUM-reduced across all ranks) at this log and the previous one, so
+    the delta covers every gradient-accumulation microbatch and the elapsed wall
+    time captures in-window overhead. Returns None when there is no prior window.
+    """
+    if prev_trainable is None or elapsed <= 0:
+        return None
+    return (curr_trainable - prev_trainable) / max(1, world_size) / elapsed
+
+
+def sanitize_kwargs_for_tagging(tag_names, kwargs=None):
+    if isinstance(tag_names, str):
+        tag_names = [tag_names]
+
+    if kwargs is not None:
+        if "tags" not in kwargs:
+            kwargs["tags"] = tag_names
+        elif "tags" in kwargs and isinstance(kwargs["tags"], list):
+            kwargs["tags"].extend(tag_names)
+        elif "tags" in kwargs and isinstance(kwargs["tags"], str):
+            tag_names.append(kwargs["tags"])
+            kwargs["tags"] = tag_names
+
+    return kwargs
+
+
+def sanitize_kwargs_for_ds_tagging(dataset_tags, kwargs=None):
+    if isinstance(dataset_tags, str):
+        dataset_tags = [dataset_tags]
+
+    if (dataset_tags is not None) and (kwargs is not None):
+        if "dataset_tags" not in kwargs:
+            kwargs["dataset_tags"] = dataset_tags
+        elif "dataset_tags" in kwargs and isinstance(kwargs["dataset_tags"], list):
+            kwargs["dataset_tags"].extend(dataset_tags)
+        elif "dataset_tags" in kwargs and isinstance(kwargs["dataset_tags"], str):
+            dataset_tags.append(kwargs["dataset_tags"])
+            kwargs["dataset_tags"] = dataset_tags
+
+    return kwargs
diff --git a/src/axolotl/core/training_args.py b/src/axolotl/core/training_args.py
new file mode 100644
index 0000000000..2a155e5ef8
--- /dev/null
+++ b/src/axolotl/core/training_args.py
@@ -0,0 +1,69 @@
+"""
+extra axolotl specific training args
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional, Type
+
+from transformers import TrainingArguments
+from trl import RewardConfig
+from trl.experimental.cpo import CPOConfig
+from trl.experimental.kto import KTOConfig
+from trl.experimental.orpo import ORPOConfig
+from trl.experimental.prm import PRMConfig
+
+from axolotl.integrations.config import merge_training_args
+
+AxolotlTrainingMixins: Type = merge_training_args()
+
+
+@dataclass
+class AxolotlTrainingArguments(AxolotlTrainingMixins, TrainingArguments):
+    """
+    Training arguments for Causal trainer
+
+    This code is duplicated due to HF TrainingArguments not setting output_dir with a
+    default value so it can't be used as a mixin.
+    """
+
+
+@dataclass
+class AxolotlORPOConfig(AxolotlTrainingMixins, ORPOConfig):
+    """
+    ORPO config for ORPO training
+    """
+
+
+@dataclass
+class AxolotlKTOConfig(AxolotlTrainingMixins, KTOConfig):
+    """
+    KTO config for KTO training
+    """
+
+
+@dataclass
+class AxolotlCPOConfig(AxolotlTrainingMixins, CPOConfig):
+    """
+    CPO config for CPO training
+    """
+
+    simpo_gamma: Optional[float] = field(
+        default=None,
+        metadata={"help": "simpo gamma parameter"},
+    )
+
+
+@dataclass
+class AxolotlRewardConfig(AxolotlTrainingMixins, RewardConfig):
+    """
+    Reward config for Reward training
+    """
+
+
+@dataclass
+class AxolotlPRMConfig(AxolotlTrainingMixins, PRMConfig):
+    """
+    PRM config for PRM training
+    """
diff --git a/src/axolotl/core/training_args_base.py b/src/axolotl/core/training_args_base.py
new file mode 100644
index 0000000000..5611e73ec4
--- /dev/null
+++ b/src/axolotl/core/training_args_base.py
@@ -0,0 +1,282 @@
+"""
+Base Axolotl Training Mixins shared across various trainer configs
+"""
+
+from dataclasses import dataclass, field
+from typing import Literal, Optional
+
+from PIL.Image import Resampling
+
+
+@dataclass
+class AxolotlTrainingMixins:
+    """
+    Mixin class for the Axolotl training args.
+    """
+
+    model_type: Optional[str] = field(
+        default=None, metadata={"help": "HF model configuration model_type."}
+    )
+    lr_quadratic_warmup: bool = field(
+        default=False,
+        metadata={"help": "Use quadratic warmup for cosine scheduling."},
+    )
+    pretraining: bool = field(
+        default=False,
+        metadata={
+            "help": "Indicates to trainer whether we are doing continued pretraining."
+        },
+    )
+    sample_packing: bool = field(
+        default=False,
+        metadata={"help": "Use sample packing for efficient training."},
+    )
+    sample_packing_sequentially: bool = field(
+        default=False,
+        metadata={
+            "help": "Use next-fit sample packing that preserves the order of samples coming from the sampler. Use in combination with curriculum_sampling for fully sequential packing."
+        },
+    )
+    sample_packing_mp_start_method: str | None = field(
+        default=None,
+        metadata={"help": "The multiprocessing start method to use."},
+    )
+    sample_packing_drop_attention_mask: bool = field(
+        default=False,
+        metadata={"help": "Drop attention mask from inputs when using packing."},
+    )
+    multipack_real_batches: bool = field(
+        default=False,
+        metadata={"help": "Use real batches for efficient training."},
+    )
+    include_tkps: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to include tokens per second in the training metrics."
+        },
+    )
+    eval_sample_packing: Optional[bool] = field(
+        default=None,
+        metadata={"help": "Use sample packing for efficient evals."},
+    )
+    sample_packing_efficiency: float = field(
+        default=1.0,
+        metadata={"help": "Sample packing efficiency for calculating batch length."},
+    )
+    sample_packing_bin_size: int = field(
+        default=200,
+        metadata={
+            "help": "The max number of samples that packed sample can contain after packing. Increase for better packing."
+        },
+    )
+    sample_packing_group_size: int = field(
+        default=100000,
+        metadata={
+            "help": "The number of samples to group together for packing. Increase for better packing."
+        },
+    )
+    max_seq_length: int = field(
+        default=2048,
+        metadata={"help": "The maximum sequence length the model can handle"},
+    )
+    dataset_num_proc: int | None = field(
+        default=None,
+        metadata={"help": "The number of processes to use for data processing"},
+    )
+    relora_prune_ratio: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": (
+                "prune ratio for optimizer state pruning; "
+                "defaults to 0.999 for reset method, 0.9 for others"
+            )
+        },
+    )
+    relora_prune_method: Optional[str] = field(
+        default=None,
+        metadata={"help": "optimizer state pruning method: magnitude | random | reset"},
+    )
+    jagged_restart_steps: Optional[int] = field(
+        default=None,
+        metadata={"help": "how often to reset for jagged restarts"},
+    )
+    jagged_restart_warmup_steps: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "how many warmup steps to take after reset for jagged restarts"
+        },
+    )
+    jagged_restart_anneal_steps: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "how many anneal steps to take before reset for jagged restarts"
+        },
+    )
+    bench_split: Optional[str] = field(
+        default="eval", metadata={"help": "The benchmark split to run on"}
+    )
+    bench_dataset: Optional[str] = field(
+        default="pharaouk/dharma-1/dharma_1_mini.json",
+        metadata={
+            "help": "Benchmark dataset to use: options are `mmlu-zs`, `mmlu-fs`, or the full path to the dataset file"
+        },
+    )
+    do_bench_eval: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to run the Benchmark evaluation."}
+    )
+    do_causal_lm_eval: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to run the Causal LM evaluation."}
+    )
+    max_bench_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "If set, only evaluates on `max_bench_samples` of the benchmark dataset."
+        },
+    )
+    bench_source_max_len: int = field(
+        default=2048, metadata={"help": "Maximum source sequence length for bench."}
+    )
+    dataloader_prefetch_factor: Optional[int] = field(
+        default=None,
+        metadata={"help": "prefetch_factor argument to the dataloader"},
+    )
+    cosine_min_lr_ratio: Optional[float] = field(
+        default=None,
+        metadata={"help": "Minimum learning rate is min_lr_ratio * learning_rate"},
+    )
+    cosine_constant_lr_ratio: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "Starting constant learning rate step is cosine_constant_lr_ratio * max_steps"
+        },
+    )
+    loraplus_lr_ratio: Optional[float] = field(
+        default=None, metadata={"help": "loraplus learning rate ratio lr_B / lr_A."}
+    )
+    loraplus_lr_embedding: Optional[float] = field(
+        default=1e-6,
+        metadata={"help": "loraplus learning rate for lora embedding layers."},
+    )
+    embedding_lr_scale: Optional[float] = field(
+        default=None,
+        metadata={"help": "Scale the learning rate for the embedding layers."},
+    )
+    lr_groups: Optional[list[dict]] = field(
+        default=None,
+        metadata={"help": "Specify learning rate groups for with different LRs."},
+    )
+    embedding_lr: Optional[float] = field(
+        default=None,
+        metadata={"help": "absolute learning rate for the embedding layers."},
+    )
+    qlora: bool = field(
+        default=False,
+        metadata={"help": "whether this is a qlora training"},
+    )
+    orpo_alpha: Optional[float] = field(
+        default=None,
+    )
+    lisa_n_layers: Optional[int] = field(
+        default=None,
+        metadata={"help": "the number of activate layers in LISA"},
+    )
+    lisa_step_interval: Optional[int] = field(
+        default=None,
+        metadata={"help": "how often to switch layers in LISA"},
+    )
+    lisa_layers_attribute: Optional[str] = field(
+        default=None,
+        metadata={"help": "path under the model to access the layers"},
+    )
+    curriculum_sampling: Optional[bool] = field(
+        default=None,
+        metadata={"help": "whether to use sequential sampling for curriculum learning"},
+    )
+    alternate_lr_scheduler_type: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "workaround to pass an alternate lr scheduler to the HF trainer"
+        },
+    )
+    chat_template: Optional[str] = field(
+        default=None,
+        metadata={"help": "Chat template converting chat messages to text"},
+    )
+
+    # kd_ce_alpha: Optional[float] = field(
+    #     default=None,
+    #     metadata={
+    #         "help": "The alpha scaling parameter for SFT cross entropy loss when using KD"
+    #     },
+    # )
+    #
+    # kd_alpha: Optional[float] = field(
+    #     default=1.0,
+    #     metadata={"help": "The alpha scaling parameter for KD loss"},
+    # )
+    #
+    # kd_temperature: Optional[float] = field(
+    #     default=1.0,
+    #     metadata={
+    #         "help": "the temperature parameter for KL divergence loss when using KD"
+    #     },
+    # )
+
+    adam_beta3: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "The beta3 hyperparameter used in some optimizers such as CAME"
+        },
+    )
+    adam_epsilon2: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "The epsilon2 hyperparameter used in some optimizers such as CAME"
+        },
+    )
+
+    activation_offloading: Literal["legacy", "disk", "hidden_states"] | bool | None = (
+        field(
+            default=None,
+            metadata={
+                "help": "Activation offloading mode: True (stream-overlapped), "
+                "'legacy' (synchronous), 'disk', 'hidden_states', or False."
+            },
+        )
+    )
+
+    layer_offloading: bool | None = field(
+        default=None,
+        metadata={
+            "help": "Offload model layer parameters to CPU during forward, prefetch back during backward."
+        },
+    )
+
+    # multi-modal section
+
+    image_size: int | tuple[int, int] | None = field(
+        default=None,
+        metadata={"help": "The size of the image to resize to"},
+    )
+
+    image_resize_algorithm: Resampling | None = field(
+        default=None,
+        metadata={"help": "The algorithm to use for image resizing"},
+    )
+
+    # end of multi-modal section
+
+    dion_learning_rate: float | None = field(
+        default=None,
+        metadata={"help": "The learning rate for Dion"},
+    )
+    dion_momentum: float | None = field(
+        default=None,
+        metadata={"help": "The momentum for Dion"},
+    )
+    dion_rank_fraction: float | None = field(
+        default=None,
+    )
+    dion_rank_multiple_of: int | None = field(
+        default=None,
+    )
diff --git a/src/axolotl/datasets.py b/src/axolotl/datasets.py
index b5638a614d..20acb85219 100644
--- a/src/axolotl/datasets.py
+++ b/src/axolotl/datasets.py
@@ -1,39 +1,36 @@
-"""Module containing Dataset functionality"""
+"""
+Module containing dataset functionality.
 
-import logging
-import os
-from typing import List, Optional
+We want this to be a wrapper for an existing dataset that we have loaded. Lets use the
+concept of middlewares to wrap each dataset. We'll use the collators later on to pad the
+datasets.
+"""
 
-import torch
 from datasets import Dataset, IterableDataset
 
-from .prompt_tokenizers import PromptTokenizingStrategy
+from axolotl.utils.logging import get_logger
 
-# We want this to be a wrapper for an existing dataset that we have loaded
-# lets use the concept of middlewares to wrap each dataset, for example
-# ConstantLengthDataset(ShuffledDataset([TokenizedPromptDataset(alpaca_dataset)]))
-# let's check to ensure we don't truncate an item in the middle, we'll use
-# the collators later on to pad the datasets
+from .prompt_tokenizers import PromptTokenizingStrategy
 
-LOG = logging.getLogger("axolotl")
+LOG = get_logger(__name__)
 
 
 class TokenizedPromptDataset(Dataset):
-    """
-    Dataset that returns tokenized prompts from a stream of text files.
-        Args:
-            prompt_tokenizer (PromptTokenizingStrategy): The prompt tokenizing method for processing the data.
-            dataset (dataset.Dataset): Dataset with text files.
-            process_count (int): Number of processes to use for tokenizing.
-            keep_in_memory (bool): Whether to keep the tokenized dataset in memory.
+    """Dataset that returns tokenized prompts from a stream of text files.
+
+    Args:
+        prompt_tokenizer: The prompt tokenizing method for processing the data.
+        dataset: Dataset with text files.
+        process_count: Number of processes to use for tokenizing.
+        keep_in_memory: Whether to keep the tokenized dataset in memory.
     """
 
-    def __init__(  # pylint: disable=super-init-not-called
+    def __init__(
         self,
         prompt_tokenizer: PromptTokenizingStrategy,
         dataset: Dataset,
-        process_count: Optional[int] = None,
-        keep_in_memory: Optional[bool] = False,
+        process_count: int | None = None,
+        keep_in_memory: bool | None = False,
         **kwargs,
     ):
         self.prompt_tokenizer = prompt_tokenizer
@@ -46,15 +43,25 @@ def __init__(  # pylint: disable=super-init-not-called
 
     def process(self, dataset):
         features = dataset.features.keys()
-        num_proc = min(64, self.process_count if self.process_count else os.cpu_count())
 
         map_kwargs = {}
         if self.prompt_tokenizer.supports_batched:
             map_kwargs["batched"] = True
-            map_kwargs["batch_size"] = 100
+            map_kwargs["batch_size"] = 1_000
+
+        if (
+            hasattr(self.prompt_tokenizer, "filter_rows")
+            and self.prompt_tokenizer.filter_rows
+        ):
+            dataset = dataset.filter(
+                self.prompt_tokenizer.filter_rows,
+                num_proc=self.process_count,
+                desc="Strategy Filtering Rows",
+            )
+
         return dataset.map(
             self.prompt_tokenizer.tokenize_prompt,
-            num_proc=num_proc,
+            num_proc=self.process_count,
             remove_columns=features,
             keep_in_memory=self.keep_in_memory,
             desc="Tokenizing Prompts",
@@ -62,127 +69,19 @@ def process(self, dataset):
         )
 
 
-# TODO this isn't the best since it can't interleave datasets
-class ConstantLengthDataset(IterableDataset):
-    """
-    Iterable dataset that returns constant length chunks of tokens from stream of text files.
-        Args:
-            tokenizer (Tokenizer): The processor used for processing the data.
-            dataset (dataset.Dataset): Dataset with text files.
-            seq_length (int): Length of token sequences to return.
-    """
-
-    def __init__(  # pylint: disable=super-init-not-called
-        self,
-        tokenizer,
-        datasets,
-        seq_length=2048,
-    ):
-        self.tokenizer = tokenizer
-        self.concat_token_id = tokenizer.eos_token_id
-        self.datasets: List[IterableDataset] = datasets
-        self.seq_length = seq_length
-
-        vocab_size = len(tokenizer.get_vocab())
-
-        if vocab_size <= torch.iinfo(torch.int16).max:
-            self.tokens_dtype = torch.int16
-        elif vocab_size <= torch.iinfo(torch.int32).max:
-            self.tokens_dtype = torch.int32
-        else:
-            self.tokens_dtype = torch.int64
-
-    def __iter__(self):
-        buffer = {
-            "input_ids": [],
-            "attention_mask": [],
-            "labels": [],
-            "position_ids": [],
-        }
-        buffer_len = 0
-        for dataset in self.datasets:
-            idx = 0
-            iterator = iter(dataset)
-            more_examples = True
-            while more_examples:
-                try:
-                    example = next(iterator)
-                    idx += 1
-                except StopIteration:
-                    more_examples = False
-                    example = None
-
-                add_concat_token = False
-                if example:
-                    example_len = len(example["input_ids"])
-                    add_concat_token = example["input_ids"][-1] != self.concat_token_id
-                else:
-                    example_len = 0
-
-                if not example_len or (
-                    buffer_len + int(add_concat_token) + example_len > self.seq_length
-                ):
-                    if buffer["input_ids"]:
-                        input_ids = torch.cat(buffer["input_ids"], dim=-1)[
-                            : self.seq_length
-                        ]
-                        attention_mask = torch.cat(buffer["attention_mask"], dim=-1)[
-                            : self.seq_length
-                        ]
-                        position_ids = torch.cat(buffer["position_ids"], dim=-1)[
-                            : self.seq_length
-                        ]
-                        labels = torch.cat(buffer["labels"], dim=-1)[: self.seq_length]
-                        if labels.size() == input_ids.size() and (
-                            attention_mask.size() == input_ids.size()
-                        ):
-                            yield {
-                                "input_ids": input_ids,
-                                "labels": labels,
-                                "attention_mask": attention_mask,
-                                "position_ids": position_ids,
-                            }
-                        else:
-                            LOG.warning(
-                                f"dropping batch due to tensor size mismatch input_ids: {input_ids.size()}, labels: {labels.size()}, attention_mask: {attention_mask.size()}"
-                            )
-                    buffer = {
-                        "input_ids": [],
-                        "attention_mask": [],
-                        "labels": [],
-                        "position_ids": [],
-                    }
-                    buffer_len = 0
-                    idx = 1
-
-                if example:
-                    # FIXME
-                    # just going to drop data points that are too long
-                    if len(example["input_ids"]) <= self.seq_length:
-                        input_ids = example["input_ids"]
-                        attention_mask = example["attention_mask"]
-                        labels = example["labels"]
-
-                        if add_concat_token:
-                            input_ids.append(self.concat_token_id)
-                            attention_mask.append(1)
-                            labels.append(self.concat_token_id)
-
-                        input_ids_with_concat = torch.tensor(
-                            input_ids, dtype=self.tokens_dtype
-                        )
-                        attention_mask_with_concat = torch.tensor(
-                            [idx * m for m in attention_mask], dtype=torch.int16
-                        )
-                        labels_with_concat = torch.tensor(
-                            labels, dtype=self.tokens_dtype
-                        )
-                        position_ids = torch.arange(
-                            len(input_ids), dtype=self.tokens_dtype
-                        )
-
-                        buffer["input_ids"].append(input_ids_with_concat)
-                        buffer["attention_mask"].append(attention_mask_with_concat)
-                        buffer["labels"].append(labels_with_concat)
-                        buffer["position_ids"].append(position_ids)
-                        buffer_len += len(input_ids)
+def wrap_dataset_for_tokenized_prompt(
+    prompt_tokenizer: PromptTokenizingStrategy,
+    dataset: Dataset | IterableDataset,
+    **kwargs,
+):
+    if isinstance(dataset, IterableDataset):
+        map_kwargs = {}
+        if prompt_tokenizer.supports_batched:
+            map_kwargs["batched"] = True
+        features = list(dataset.features.keys())
+        return dataset.map(
+            prompt_tokenizer.tokenize_prompt,
+            remove_columns=features,
+            **map_kwargs,
+        )
+    return TokenizedPromptDataset(prompt_tokenizer, dataset, **kwargs)
diff --git a/src/axolotl/evaluate.py b/src/axolotl/evaluate.py
new file mode 100644
index 0000000000..db6fb3f161
--- /dev/null
+++ b/src/axolotl/evaluate.py
@@ -0,0 +1,151 @@
+"""Module for evaluating models."""
+
+import csv
+import os
+import sys
+from pathlib import Path
+from typing import Dict, Optional
+
+import torch
+from datasets import Dataset
+from transformers.trainer import Trainer
+
+from axolotl.telemetry.errors import send_errors
+from axolotl.train import (
+    TrainDatasetMeta,
+    setup_model_and_tokenizer,
+)
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.distributed import cleanup_distributed
+from axolotl.utils.logging import get_logger
+from axolotl.utils.trainer import setup_trainer
+
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+src_dir = os.path.join(project_root, "src")
+sys.path.insert(0, src_dir)
+
+LOG = get_logger(__name__)
+
+
+def evaluate_dataset(
+    trainer: Trainer, dataset: Dataset, dataset_type: str, flash_optimum: bool = False
+) -> Optional[Dict[str, float]]:
+    """Helper function to evaluate a single dataset.
+
+    Args:
+        trainer: The trainer instance.
+        dataset: Dataset to evaluate.
+        dataset_type: Type of dataset ('train' or 'eval').
+        flash_optimum: Whether to use flash optimum.
+
+    Returns:
+        Dictionary of metrics or None if dataset is None.
+    """
+    if dataset is None:
+        return None
+
+    LOG.info(f"Starting {dataset_type} set evaluation...")
+
+    if flash_optimum:
+        with torch.backends.cuda.sdp_kernel(
+            enable_flash=True,
+            enable_math=True,
+            enable_mem_efficient=True,
+        ):
+            metrics = trainer.evaluate(dataset, metric_key_prefix=dataset_type)
+    else:
+        metrics = trainer.evaluate(dataset, metric_key_prefix=dataset_type)
+
+    LOG.info(f"{dataset_type.capitalize()} set evaluation completed!")
+    LOG.info(f"{dataset_type.capitalize()} Metrics:")
+    for key, value in metrics.items():
+        LOG.info(f"{key}: {value}")
+
+    return metrics
+
+
+@send_errors
+def evaluate(*, cfg: DictDefault, dataset_meta: TrainDatasetMeta) -> Dict[str, float]:
+    """
+    Evaluate a model on training and validation datasets.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        dataset_meta: Dataset metadata containing training and evaluation datasets.
+
+    Returns:
+        Dictionary mapping metric names to their values.
+    """
+    # Load tokenizer, processor and model
+    LOG.debug("loading model for evaluation...")
+    model, tokenizer, _, processor = setup_model_and_tokenizer(cfg)
+
+    # Get datasets
+
+    train_dataset = dataset_meta.train_dataset
+    eval_dataset = dataset_meta.eval_dataset
+    total_num_steps = dataset_meta.total_num_steps
+
+    # Set up trainer
+    trainer = setup_trainer(
+        cfg=cfg,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        model=model,
+        tokenizer=tokenizer,
+        processor=processor,
+        total_num_steps=total_num_steps,
+    )
+
+    # Evaluate datasets
+    all_metrics = {}
+    train_metrics = evaluate_dataset(trainer, train_dataset, "train", cfg.flash_optimum)
+    eval_metrics = evaluate_dataset(trainer, eval_dataset, "eval", cfg.flash_optimum)
+
+    if train_metrics:
+        all_metrics.update(train_metrics)
+    if eval_metrics:
+        all_metrics.update(eval_metrics)
+
+    # Save metrics to CSV if output directory is specified and we have metrics
+    if cfg.output_dir and (train_metrics or eval_metrics):
+        output_dir = Path(cfg.output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        metrics_file = output_dir / "eval_summary.csv"
+        with metrics_file.open("w", newline="", encoding="utf-8") as file:
+            writer = csv.writer(file)
+            writer.writerow(["metric", "training", "validation"])
+
+            # Get unique metric names (removing prefixes) from available metrics
+            train_metric_names = {
+                k.replace("train_", ""): k for k in (train_metrics or {})
+            }
+            eval_metric_names = {
+                k.replace("eval_", ""): k for k in (eval_metrics or {})
+            }
+            all_metric_names = sorted(
+                set(train_metric_names.keys()) | set(eval_metric_names.keys())
+            )
+
+            for metric_name in all_metric_names:
+                train_value = (
+                    train_metrics.get(train_metric_names.get(metric_name, ""), "")
+                    if train_metrics
+                    else ""
+                )
+                eval_value = (
+                    eval_metrics.get(eval_metric_names.get(metric_name, ""), "")
+                    if eval_metrics
+                    else ""
+                )
+                writer.writerow([metric_name, train_value, eval_value])
+
+        LOG.info(f"Evaluation results saved to {metrics_file}")
+
+    del model
+    del tokenizer
+
+    cleanup_distributed()
+
+    return all_metrics
diff --git a/src/axolotl/integrations/LICENSE.md b/src/axolotl/integrations/LICENSE.md
new file mode 100644
index 0000000000..435d36d75b
--- /dev/null
+++ b/src/axolotl/integrations/LICENSE.md
@@ -0,0 +1,58 @@
+### AXOLOTL COMMUNITY LICENSE AGREEMENT
+
+This Axolotl Community License Agreement (“Agreement”) is entered into by and between Axolotl AI Corp. (“Axolotl”) and
+any individual or entity (“Licensee”) who wishes to use the Software (as defined below) in accordance with the terms
+and conditions set forth in this Agreement.
+
+1.  Definitions
+    1.1 “Licensee” refers to any individual or entity who has obtained a copy of the Software under this Agreement.
+    1.2 “Plugin Integration” means independent integration software modules which may or may not be offered by Axolotl,
+        which may be licensed separately by their respective  authors and/or licensors.
+    1.3 “Software” refers to the specific sub-directory of the Axolotl, Inc. software located at
+        https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations and its subdirectories which
+        permits Plugin Integrations to integrate with the Axolotl service.
+2.  Grant of License
+    2.1	Axolotl hereby grants Licensee a worldwide, non-exclusive, royalty-free, license to use, copy, modify, merge,
+        publish, distribute, sublicense, and/or otherwise exploit the Software, subject to the following conditions:
+        - Licensee must comply with all the terms and conditions of this Agreement.
+        - Licensee must include the original copyright notice and disclaimer of warranty in all copies or substantial
+          portions of the Software.
+    2.2 Licensee may use the Software for any lawful purpose, except as restricted in Section 3.
+3.  Restrictions
+    3.1 Licensee shall not use the Software for any activity that constitutes a commercial activity of offering for
+        free or for sale any services, platform, or equivalent  to third parties for the purposes of allowing such
+        third parties to fine-tune artificial intelligence models.
+    3.2 Licensee shall not:
+        - Use the Software for any illegal or unauthorized purpose.
+        - Reverse engineer, decompile, or disassemble the Software.
+        - Remove or modify any copyright, trademark, or other proprietary notices contained in the Software.
+        - Use the Software in a way that could damage, disable, overburden, or impair the functionality of the
+          Software or interfere with any third-party use of the Software.
+    3.3 Axolotl reserves the right to restrict certain Plugin Integrations for use with the Software. To the extent Licensee integrates a permitted, applicable Plugin Integration with the Software, Licensee shall comply with any additional terms and conditions imposed by the licensors of such Plugin Integration for use of such Plugin Integrations. Licensee shall contact Axolotl if it has questions about whether its use of the Software falls beyond the scope of this Agreement.
+4.  Intellectual Property Rights
+    4.1 Axolotl and its contributors retain all intellectual property rights in and to the Software. Licensee
+        acknowledges that this Agreement does not transfer any ownership rights or intellectual property rights to
+        Licensee.
+5.  Disclaimer of Warranty
+    5.1 THE SOFTWARE IS PROVIDED “AS IS,” WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+        TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT. IN NO EVENT SHALL
+        THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES, OR OTHER LIABILITY, WHETHER IN AN ACTION OF
+        CONTRACT, TORT, OR OTHERWISE, ARISING FROM, OUT OF, OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+        DEALINGS IN THE SOFTWARE.
+6.  Termination
+    6.1 Axolotl may terminate this Agreement at any time if Licensee fails to comply with any of the terms and
+        conditions set forth herein. Upon termination, Licensee shall cease all use of the Software and destroy any
+        copies in its possession.
+7.  Governing Law
+    7.1 This Agreement shall be governed by and construed in accordance with the laws of the State of California,
+        without regards to conflicts of laws provisions thereof.
+8.  Entire Agreement
+    8.1 This Agreement constitutes the entire agreement between Axolotl and Licensee with respect to the subject matter
+        hereof and supersedes all prior or contemporaneous understandings or agreements between the parties concerning
+        the Software, whether written or oral. Axolotl may update the terms of this Agreement from time to time, and
+        Licensee’s continued use of the Software after any such updates shall constitute acceptance of updated terms
+        on a go-forward basis.  Axolotl will use commercially reasonable efforts to provide Licensee notice of any
+        material updates. By using the Software, Licensee acknowledges that it has read, understood, and agrees to be
+        bound by the terms and conditions of this Agreement.
+
+This Agreement was last updated on August 23, 2024.
diff --git a/src/axolotl/integrations/__init__.py b/src/axolotl/integrations/__init__.py
new file mode 100644
index 0000000000..f77af49c2f
--- /dev/null
+++ b/src/axolotl/integrations/__init__.py
@@ -0,0 +1,3 @@
+import pkgutil
+
+__path__ = pkgutil.extend_path(__path__, __name__)
diff --git a/src/axolotl/integrations/base.py b/src/axolotl/integrations/base.py
new file mode 100644
index 0000000000..5d6517d8dc
--- /dev/null
+++ b/src/axolotl/integrations/base.py
@@ -0,0 +1,789 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+
+"""Base class for all plugins.
+
+A plugin is a reusable, modular, and self-contained piece of code that extends the functionality of Axolotl.
+Plugins can be used to integrate third-party models, modify the training process, or add new features.
+
+To create a new plugin, you need to inherit from the BasePlugin class and implement the required methods.
+"""
+
+from __future__ import annotations
+
+import collections
+import importlib
+import traceback
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, OrderedDict, Union
+
+from peft import PeftConfig, PeftMixedModel, PeftModel
+from torch import nn
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+from transformers import PreTrainedModel, Trainer
+from transformers.trainer_pt_utils import get_parameter_names
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from axolotl.common.datasets import TrainDatasetMeta
+
+
+@dataclass(frozen=True)
+class AdapterCapabilities:
+    """Capabilities for an adapter contributed by a plugin."""
+
+    name: str
+    lora_like: bool = False
+    relora: bool = False
+
+
+class BasePlugin:
+    """Base class for all plugins. Defines the interface for plugin methods.
+
+    A plugin is a reusable, modular, and self-contained piece of code that extends
+    the functionality of Axolotl. Plugins can be used to integrate third-party models,
+    modify the training process, or add new features.
+
+    To create a new plugin, you need to inherit from the BasePlugin class and
+    implement the required methods.
+
+    Note:
+        Plugin methods include:
+        - register(cfg): Registers the plugin with the given configuration.
+        - load_datasets(cfg): Loads and preprocesses the dataset for training.
+        - pre_model_load(cfg): Performs actions before the model is loaded.
+        - post_model_build(cfg, model): Performs actions after the model is loaded, but
+            before LoRA adapters are applied.
+        - pre_lora_load(cfg, model): Performs actions before LoRA weights are loaded.
+        - post_lora_load(cfg, model): Performs actions after LoRA weights are loaded.
+        - post_model_load(cfg, model): Performs actions after the model is loaded,
+            inclusive of any adapters.
+        - post_trainer_create(cfg, trainer): Performs actions after the trainer is
+            created.
+        - create_optimizer(cfg, trainer): Creates and returns an optimizer for training.
+        - create_lr_scheduler(cfg, trainer, optimizer, num_training_steps): Creates and
+            returns a learning rate scheduler.
+        - add_callbacks_pre_trainer(cfg, model): Adds callbacks to the trainer before
+            training.
+        - add_callbacks_post_trainer(cfg, trainer): Adds callbacks to the trainer after
+            training.
+    """
+
+    def __init__(self):
+        """Initializes the BasePlugin."""
+
+    def register(self, cfg: dict):
+        """Registers the plugin with the given configuration as an unparsed dict.
+
+        Args:
+            cfg: The configuration for the plugin.
+        """
+
+    def get_input_args(self) -> str | None:
+        """Returns a pydantic model for the plugin's input arguments."""
+
+    def get_training_args_mixin(self) -> str | None:
+        """
+        Returns a dataclass model for the plugin's training arguments.
+        """
+
+    def get_adapter_capabilities(self) -> list[AdapterCapabilities]:
+        """Returns adapter capabilities contributed by the plugin."""
+        return []
+
+    def get_lora_config_kwargs(self, cfg: DictDefault) -> dict:
+        """Returns extra PEFT LoraConfig kwargs for plugin LoRA-like adapters."""
+        return {}
+
+    def load_adapter(
+        self,
+        model: PreTrainedModel,
+        cfg: DictDefault,
+        inference: bool = False,
+        config_only: bool = False,
+    ) -> (
+        tuple[PreTrainedModel | PeftModel | PeftMixedModel | None, PeftConfig | None]
+        | None
+    ):
+        """Optionally load a plugin adapter instead of the generic loader."""
+
+    def load_datasets(
+        self, cfg: DictDefault, preprocess: bool = False
+    ) -> Union["TrainDatasetMeta", None]:
+        """Loads and preprocesses the dataset for training.
+
+        Args:
+            cfg: The configuration for the plugin.
+            preprocess: Whether this is the preprocess step of the datasets.
+
+        Returns:
+            dataset_meta: The metadata for the training dataset.
+        """
+
+    def pre_model_load(self, cfg: DictDefault):
+        """Performs actions before the model is loaded.
+
+        Args:
+            cfg: The configuration for the plugin.
+        """
+
+    def post_model_build(self, cfg: DictDefault, model: PreTrainedModel):
+        """Performs actions after the model is built/loaded, but before any adapters are applied.
+
+        Args:
+            cfg: The configuration for the plugin.
+        """
+
+    def pre_lora_load(self, cfg: DictDefault, model: PreTrainedModel):
+        """Performs actions before LoRA weights are loaded.
+
+        Args:
+            cfg: The configuration for the plugin.
+            model: The loaded model.
+        """
+
+    def post_lora_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        """Performs actions after LoRA weights are loaded.
+
+        Args:
+            cfg: The configuration for the plugin.
+            model: The loaded model.
+        """
+
+    def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        """Performs actions after the model is loaded.
+
+        Args:
+            cfg: The configuration for the plugin.
+            model: The loaded model.
+        """
+
+    def get_trainer_cls(self, cfg: DictDefault) -> type[Trainer] | None:
+        """Returns a custom class for the trainer.
+
+        Args:
+            cfg: The global axolotl configuration.
+
+        Returns:
+            The first non-`None` trainer class returned by a plugin.
+        """
+
+    def post_trainer_create(self, cfg: DictDefault, trainer: Trainer):
+        """Performs actions after the trainer is created.
+
+        Args:
+            cfg: The configuration for the plugin.
+            trainer: The trainer object for training.
+        """
+
+    def get_training_args(self, cfg: DictDefault):
+        """
+        Returns custom training arguments to set on TrainingArgs.
+
+        Args:
+            cfg: The global axolotl configuration.
+
+        Returns:
+            object: dict containing the training arguments.
+        """
+
+    def get_collator_cls_and_kwargs(self, cfg: DictDefault, is_eval: bool = False):
+        """
+        Returns a custom class for the collator.
+
+        Args:
+            cfg: The global axolotl configuration.
+            is_eval: Whether this is an eval split.
+
+        Returns:
+            class: The class for the collator.
+        """
+
+    def create_optimizer(self, cfg: DictDefault, trainer: Trainer) -> Optimizer | None:
+        """Creates and returns an optimizer for training.
+
+        Args:
+            cfg: The configuration for the plugin.
+            trainer: The trainer object for training.
+
+        Returns:
+            The created optimizer.
+        """
+
+    def create_lr_scheduler(
+        self,
+        cfg: DictDefault,
+        trainer: Trainer,
+        optimizer: Optimizer,
+        num_training_steps: int,
+    ) -> LRScheduler | None:
+        """Creates and returns a learning rate scheduler.
+
+        Args:
+            cfg: The configuration for the plugin.
+            trainer: The trainer object for training.
+            optimizer: The optimizer for training.
+            num_training_steps: Total number of training steps
+
+        Returns:
+            The created learning rate scheduler.
+        """
+
+    def add_callbacks_pre_trainer(
+        self, cfg: DictDefault, model: PreTrainedModel
+    ) -> list[Callable]:
+        """Set up callbacks before creating the trainer.
+
+        Args:
+            cfg: The configuration for the plugin.
+            model: The loaded model.
+
+        Returns:
+            A list of callback functions to be added to the `TrainingArgs`.
+        """
+        return []
+
+    def add_callbacks_post_trainer(
+        self, cfg: DictDefault, trainer: Trainer
+    ) -> list[Callable]:
+        """Adds callbacks to the trainer after creating the trainer. This is useful for
+        callbacks that require access to the model or trainer.
+
+        Args:
+            cfg: The configuration for the plugin.
+            trainer: The trainer object for training.
+
+        Returns:
+            A list of callback functions to be added
+        """
+        return []
+
+    def on_rollouts_scored(
+        self,
+        cfg: DictDefault,
+        trainer,
+        prompts: list[str],
+        completions: list[str],
+        rewards: dict[str, list[float]],
+        advantages: list[float],
+    ):
+        """Called after rollouts are scored during online RL (GRPO/PPO).
+
+        Provides access to the full scored rollout data for logging, trace
+        storage, or analysis. Called once per scoring step with all samples
+        from that step.
+
+        Args:
+            cfg: The axolotl configuration.
+            trainer: The trainer instance.
+            prompts: List of prompt texts (one per sample).
+            completions: List of completion texts (one per sample).
+            rewards: Dict mapping reward function name to list of reward values.
+            advantages: List of advantage values (one per sample).
+        """
+
+    def post_train(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        """Performs actions after training is complete.
+
+        Args:
+            cfg: The axolotl configuration.
+            model: The loaded model.
+        """
+
+    def post_train_unload(self, cfg: DictDefault):
+        """Performs actions after training is complete and the model is unloaded.
+
+        Args:
+            cfg: The configuration for the plugin.
+        """
+
+
+def load_plugin(plugin_name: str) -> BasePlugin:
+    """Loads a plugin based on the given plugin name.
+
+    The plugin name should be in the format "module_name.class_name". This function
+    splits the plugin name into module and class, imports the module, retrieves the
+    class from the module, and creates an instance of the class.
+
+    Args:
+        plugin_name: The name of the plugin to be loaded. The name should be in the
+            format "module_name.class_name".
+
+    Returns:
+        An instance of the loaded plugin.
+
+    Raises:
+        ImportError: If the plugin module cannot be imported.
+    """
+    # split the plugin name into module and class
+    module_name, class_name = plugin_name.rsplit(".", 1)
+
+    # import the module
+    try:
+        module = importlib.import_module(module_name)
+    except ModuleNotFoundError as orig_exc:
+        try:
+            if not module_name.startswith("axolotl.integrations."):
+                module = importlib.import_module("axolotl.integrations." + module_name)
+            else:
+                raise orig_exc
+        except ModuleNotFoundError as exc:
+            raise orig_exc from exc
+
+    # instantiate the class
+    plugin_class = getattr(module, class_name)
+    # create an instance of the class
+    plugin = plugin_class()
+
+    return plugin
+
+
+class PluginManager:
+    """The `PluginManager` class is responsible for loading and managing plugins. It
+    should be a singleton so it can be accessed from anywhere in the codebase.
+
+    Attributes:
+        plugins: A list of loaded plugins.
+
+    Note:
+        Key methods include:
+        - get_instance(): Static method to get the singleton instance of `PluginManager`.
+        - register(plugin_name: str): Registers a new plugin by its name.
+        - pre_model_load(cfg): Calls the pre_model_load method of all registered plugins.
+    """
+
+    plugins: OrderedDict[str, BasePlugin] = collections.OrderedDict()
+
+    _instance: PluginManager | None = None
+    _cfg: DictDefault | None = None
+
+    def __new__(cls):
+        """Creates a new instance of PluginManager if it doesn't exist yet."""
+        if cls._instance is None:
+            cls._instance = super(PluginManager, cls).__new__(cls)
+            cls._instance.plugins: OrderedDict[str, BasePlugin] = (
+                collections.OrderedDict()
+            )
+        return cls._instance
+
+    @staticmethod
+    def get_instance() -> "PluginManager":
+        """Returns the singleton instance of PluginManager. If the instance doesn't
+        exist, it creates a new one.
+        """
+        if PluginManager._instance is None:
+            PluginManager()
+        return PluginManager._instance  # type: ignore
+
+    @property
+    def cfg(self):
+        return self._cfg
+
+    @cfg.setter
+    def cfg(self, cfg):
+        self._cfg = cfg
+
+    def register(self, plugin_name: str):
+        """Registers a new plugin by its name.
+
+        Args:
+            plugin_name: The name of the plugin to be registered.
+
+        Raises:
+            ImportError: If the plugin module cannot be imported.
+        """
+        try:
+            LOG.info(f"Attempting to load plugin: {plugin_name}")
+            plugin = load_plugin(plugin_name)
+            self.plugins[plugin_name] = plugin
+            LOG.info(f"Plugin loaded successfully: {plugin_name}")
+        except ImportError as exc:
+            LOG.error(f"Failed to load plugin: {plugin_name}")
+            # print stacktrace
+            traceback.print_exc()
+            print(f"Error: {exc}")
+
+    def get_input_args(self) -> list[str]:
+        """Returns a list of Pydantic classes for all registered plugins' input arguments.'
+
+        Returns:
+            A list of Pydantic classes for all registered plugins' input arguments.'
+        """
+        input_args = []
+        for plugin in self.plugins.values():
+            input_args_from_plugin = plugin.get_input_args()
+            if input_args_from_plugin is not None:
+                input_args.append(input_args_from_plugin)
+        return input_args
+
+    def get_training_args_mixin(self):
+        """
+        Returns a list of dataclasses for all registered plugins' training args mixins'
+
+        Returns:
+        list[str]: A list of dataclsses
+        """
+        training_args = []
+        for plugin in self.plugins.values():
+            training_args_from_plugin = plugin.get_training_args_mixin()
+            if training_args_from_plugin is not None:
+                training_args.append(training_args_from_plugin)
+        return training_args
+
+    def adapter_capabilities(self) -> dict[str, AdapterCapabilities]:
+        """Returns adapter capabilities by adapter name."""
+        capabilities = {}
+        for plugin in self.plugins.values():
+            for adapter_capability in plugin.get_adapter_capabilities():
+                capabilities[adapter_capability.name] = adapter_capability
+        return capabilities
+
+    def get_adapter_capability(self, adapter: str) -> AdapterCapabilities | None:
+        """Returns capabilities for a registered plugin adapter."""
+        return self.adapter_capabilities().get(adapter)
+
+    def supports_adapter(self, adapter: str) -> bool:
+        """Returns whether a plugin has registered the adapter name."""
+        return adapter in self.adapter_capabilities()
+
+    def adapter_supports_relora(self, adapter: str) -> bool:
+        """Returns whether a plugin adapter supports ReLoRA restart semantics."""
+        capability = self.get_adapter_capability(adapter)
+        return bool(capability and capability.relora)
+
+    def get_lora_config_kwargs(self, cfg: DictDefault) -> dict:
+        """Returns extra LoraConfig kwargs from plugins for the configured adapter."""
+        lora_config_kwargs = {}
+        for plugin in self.plugins.values():
+            plugin_kwargs = plugin.get_lora_config_kwargs(cfg)
+            if plugin_kwargs:
+                lora_config_kwargs.update(plugin_kwargs)
+        return lora_config_kwargs
+
+    def load_adapter(
+        self,
+        model: PreTrainedModel,
+        cfg: DictDefault,
+        inference: bool = False,
+        config_only: bool = False,
+    ) -> (
+        tuple[PreTrainedModel | PeftModel | PeftMixedModel | None, PeftConfig | None]
+        | None
+    ):
+        """Returns the first plugin adapter loader result, if any."""
+        for plugin in self.plugins.values():
+            loaded = plugin.load_adapter(
+                model,
+                cfg,
+                inference=inference,
+                config_only=config_only,
+            )
+            if loaded is not None:
+                return loaded
+        return None
+
+    def load_datasets(
+        self, cfg: DictDefault, preprocess: bool = False
+    ) -> Union["TrainDatasetMeta", None]:
+        """Calls the load_datasets method of each registered plugin.
+
+        Args:
+            cfg: The configuration for the plugins.
+            preprocess: Whether this is preprocess step of the datasets.
+
+        Returns:
+            The dataset metadata loaded from all registered plugins.
+        """
+        return_ds_meta = None
+        for plugin in self.plugins.values():
+            dataset_meta = plugin.load_datasets(cfg, preprocess)
+            if dataset_meta is not None:
+                if return_ds_meta is None:
+                    return_ds_meta = dataset_meta
+                else:
+                    raise RuntimeError("Multiple plugins loaded datasets")
+        return return_ds_meta
+
+    def pre_model_load(self, cfg: DictDefault):
+        """Calls the pre_model_load method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+        """
+        for plugin in self.plugins.values():
+            plugin.pre_model_load(cfg)
+
+    def post_model_build(self, cfg: DictDefault, model: PreTrainedModel):
+        """Calls the `post_model_build` method of all registered plugins after the
+        model has been built / loaded, but before any adapters have been applied.
+
+        Args:
+            cfg: The configuration for the plugins.
+            model: The loaded model.
+        """
+        for plugin in self.plugins.values():
+            plugin.post_model_build(cfg, model)
+
+    def pre_lora_load(self, cfg: DictDefault, model: PreTrainedModel):
+        """Calls the `pre_lora_load` method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+            model: The loaded model.
+        """
+        for plugin in self.plugins.values():
+            plugin.pre_lora_load(cfg, model)
+
+    def post_lora_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        """Calls the `post_lora_load` method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+            model: The loaded model.
+        """
+        for plugin in self.plugins.values():
+            plugin.post_lora_load(cfg, model)
+
+    def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        """Calls the `post_model_load` method of all registered plugins after the model
+        has been loaded inclusive of any adapters.
+
+        Args:
+            cfg: The configuration for the plugins.
+            model: The loaded model.
+        """
+        for plugin in self.plugins.values():
+            plugin.post_model_load(cfg, model)
+
+    def get_trainer_cls(self, cfg: DictDefault) -> Trainer | None:
+        """Calls the `get_trainer_cls` method of all registered plugins and returns the
+        first non-`None` trainer class.
+
+        Args:
+            cfg: The configuration for the plugins.
+
+        Returns:
+            The first non-`None` trainer class returned by a plugin.
+        """
+        for plugin in self.plugins.values():
+            trainer_cls = plugin.get_trainer_cls(cfg)
+            if trainer_cls is not None:
+                return trainer_cls
+        return None
+
+    def get_training_args(self, cfg):
+        """
+        Calls the get_training_args method of all registered plugins and returns the combined training arguments.
+
+        Parameters:
+        cfg (dict): The configuration for the plugins.
+
+        Returns:
+        object: The training arguments
+        """
+        training_args_kwargs = {}
+        for plugin in self.plugins.values():
+            training_args = plugin.get_training_args(cfg)
+            if training_args is not None:
+                training_args_kwargs.update(training_args)
+
+        return training_args_kwargs
+
+    def get_collator_cls_and_kwargs(self, cfg, is_eval=False):
+        """
+        Calls the get_collator_cls_and_kwargs method of all registered plugins and returns the first non-None collator class.
+
+        Parameters:
+        cfg (dict): The configuration for the plugins.
+        is_eval (bool): Whether this is an eval split.
+
+        Returns:
+        object: The collator class, or None if none was found.
+        """
+        for plugin in self.plugins.values():
+            collator = plugin.get_collator_cls_and_kwargs(cfg, is_eval=is_eval)
+            if collator is not None:
+                collator_cls, collator_kwargs = collator
+                return collator_cls, collator_kwargs
+        return None
+
+    def post_trainer_create(self, cfg: DictDefault, trainer: Trainer):
+        """Calls the `post_trainer_create` method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+            trainer: The trainer object for training.
+        """
+        for plugin in self.plugins.values():
+            plugin.post_trainer_create(cfg, trainer)
+
+    def create_optimizer(self, trainer: Trainer) -> Optimizer | None:
+        """Calls the `create_optimizer` method of all registered plugins and returns
+        the first non-`None` optimizer.
+
+        Args:
+            trainer: The trainer object for training.
+
+        Returns:
+            The created optimizer, or `None` if none was found.
+        """
+        for plugin in self.plugins.values():
+            optimizer = plugin.create_optimizer(self.cfg, trainer)
+            if optimizer is not None:
+                return optimizer
+        return None
+
+    def create_lr_scheduler(
+        self, trainer: Trainer, optimizer: Optimizer, num_training_steps: int
+    ) -> LRScheduler | None:
+        """Calls the `create_lr_scheduler` method of all registered plugins and returns
+        the first non-`None` scheduler.
+
+        Args:
+            trainer: The trainer object for training.
+            optimizer: The optimizer for training.
+
+        Returns:
+            The created learning rate scheduler, or `None` if not found.
+        """
+        for plugin in self.plugins.values():
+            scheduler: LRScheduler | None = plugin.create_lr_scheduler(
+                self.cfg,
+                trainer=trainer,
+                optimizer=optimizer,
+                num_training_steps=num_training_steps,
+            )
+            if scheduler is not None:
+                return scheduler
+        return None
+
+    def add_callbacks_pre_trainer(
+        self, cfg: DictDefault, model: PreTrainedModel
+    ) -> list[Callable]:
+        """Calls the add_callbacks_pre_trainer method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+            model: The loaded model.
+
+        Returns:
+            A list of callback functions to be added to the `TrainingArgs`.
+        """
+        callbacks = []
+        for plugin in self.plugins.values():
+            plugin_callbacks = plugin.add_callbacks_pre_trainer(cfg, model)
+            if plugin_callbacks:  # if the plugin returned a list of callbacks
+                callbacks.extend(plugin_callbacks)
+        return callbacks
+
+    def add_callbacks_post_trainer(
+        self, cfg: DictDefault, trainer: Trainer
+    ) -> list[Callable]:
+        """Calls the `add_callbacks_post_trainer` method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+            trainer: The trainer object for training.
+
+        Returns:
+            A list of callback functions to be added to the `TrainingArgs`.
+        """
+        callbacks = []
+        for plugin in self.plugins.values():
+            plugin_callbacks = plugin.add_callbacks_post_trainer(cfg, trainer)
+            if plugin_callbacks:
+                callbacks.extend(plugin_callbacks)
+        return callbacks
+
+    def post_train(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        """Calls the post_train method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+            model: The loaded model.
+        """
+        for plugin in self.plugins.values():
+            plugin.post_train(cfg, model)
+
+    def on_rollouts_scored(
+        self,
+        cfg: DictDefault,
+        trainer,
+        prompts: list[str],
+        completions: list[str],
+        rewards: dict[str, list[float]],
+        advantages: list[float],
+    ):
+        """Calls the on_rollouts_scored method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+            trainer: The trainer instance.
+            prompts: List of prompt texts.
+            completions: List of completion texts.
+            rewards: Dict mapping reward function name to list of rewards.
+            advantages: List of advantage values.
+        """
+        for plugin in self.plugins.values():
+            try:
+                plugin.on_rollouts_scored(
+                    cfg, trainer, prompts, completions, rewards, advantages
+                )
+            except Exception:
+                LOG.warning(
+                    f"Plugin {plugin.__class__.__name__}.on_rollouts_scored failed",
+                    exc_info=True,
+                )
+
+    def post_train_unload(self, cfg: DictDefault):
+        """Calls the post_train_unload method of all registered plugins.
+
+        Args:
+            cfg: The configuration for the plugins.
+        """
+        for plugin in self.plugins.values():
+            plugin.post_train_unload(cfg)
+
+
+class BaseOptimizerFactory:
+    """Base class for factories to create custom optimizers"""
+
+    def __call__(
+        self, opt_model, training_args, **optimizer_kwargs
+    ) -> Optimizer | None:
+        pass
+
+    # duplicated from transformers
+    def get_decay_parameter_names(self, model) -> list[str]:
+        """
+        Get all parameter names that weight decay will be applied to.
+
+        This function filters out parameters in two ways:
+        1. By layer type (instances of layers specified in ALL_LAYERNORM_LAYERS)
+        2. By parameter name patterns (containing 'bias', or variation of 'norm')
+        """
+        forbidden_name_patterns = [
+            r"bias",
+            r"layernorm",
+            r"rmsnorm",
+            r"(?:^|\.)norm(?:$|\.)",
+            r"_norm(?:$|\.)",
+        ]
+        decay_parameters = get_parameter_names(
+            model, [nn.LayerNorm], forbidden_name_patterns
+        )
+        return decay_parameters
diff --git a/src/axolotl/integrations/config.py b/src/axolotl/integrations/config.py
new file mode 100644
index 0000000000..8ae8aab39b
--- /dev/null
+++ b/src/axolotl/integrations/config.py
@@ -0,0 +1,93 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+
+"""
+Module to handle merging the plugins' input arguments with the base configurations.
+
+This was moved here to prevent circular imports.
+"""
+
+from typing import Any, Dict, List, Type
+
+from axolotl.utils.schemas.config import (
+    AxolotlConfigWCapabilities as AxolotlConfigWCapabilitiesBase,
+    AxolotlInputConfig as AxolotlInputConfigBase,
+)
+
+
+def merge_input_args():
+    """
+    Merges input arguments from registered plugins with the base configurations.
+
+    This function retrieves the input arguments from registered plugins using the PluginManager.
+    It then dynamically creates new classes, AxolotlConfigWCapabilities and AxolotlInputConfig,
+    that inherit from the base configurations and include the input arguments from the plugins.
+
+    Returns:
+    tuple: A tuple containing the newly created classes, AxolotlConfigWCapabilities and AxolotlInputConfig.
+    """
+    from axolotl.integrations.base import PluginManager
+
+    plugin_manager = PluginManager.get_instance()
+    input_args: List[str] = plugin_manager.get_input_args()
+    plugin_classes = []
+    dynamic_input = ""
+    for plugin_args in input_args:
+        plugin_module, plugin_cls = plugin_args.rsplit(".", 1)
+        dynamic_input += f"from {plugin_module} import {plugin_cls}\n"
+        plugin_classes.append(plugin_cls)
+    if dynamic_input:
+        dynamic_input += f"class AxolotlConfigWCapabilities(AxolotlConfigWCapabilitiesBase, {', '.join(plugin_classes)}):\n    pass\n"
+        dynamic_input += f"class AxolotlInputConfig(AxolotlInputConfigBase, {', '.join(plugin_classes)}):\n    pass\n"
+
+        namespace: Dict[Any, Any] = {}
+        exec(dynamic_input, globals(), namespace)  # nosec B102
+        AxolotlInputConfig = namespace["AxolotlInputConfig"]
+        AxolotlConfigWCapabilities = namespace["AxolotlConfigWCapabilities"]
+        return AxolotlConfigWCapabilities, AxolotlInputConfig
+    return AxolotlConfigWCapabilitiesBase, AxolotlInputConfigBase
+
+
+def merge_training_args() -> Type:
+    """
+    Merges training arguments from registered plugins with the base TrainingArguments.
+
+    This function retrieves the training arguments from registered plugins using the PluginManager.
+    It then dynamically creates new classes, AxolotlTrainingMixins,
+    that inherit from the base configurations and include the training arguments from the plugins.
+
+    Returns:
+    tuple: A tuple containing the newly created classes, AxolotlTrainingMixins.
+    """
+
+    from axolotl.core.training_args_base import (
+        AxolotlTrainingMixins as AxolotlTrainingMixinsBase,
+    )
+    from axolotl.integrations.base import PluginManager
+
+    plugin_manager = PluginManager.get_instance()
+    training_args_mixins: List[str] = plugin_manager.get_training_args_mixin()
+    mixin_classes = []
+    dynamic_input = ""
+    for plugin_args in training_args_mixins:
+        plugin_module, plugin_cls = plugin_args.rsplit(".", 1)
+        dynamic_input += f"from {plugin_module} import {plugin_cls}\n"
+        mixin_classes.append(plugin_cls)
+    if dynamic_input:
+        dynamic_input += f"class AxolotlTrainingMixins(AxolotlTrainingMixinsBase, {', '.join(mixin_classes)}):\n    pass\n"
+
+        namespace: Dict[Any, Any] = {}
+        local_vars = {"AxolotlTrainingMixinsBase": AxolotlTrainingMixinsBase}
+        exec(dynamic_input, {**globals(), **local_vars}, namespace)  # nosec B102
+        AxolotlTrainingMixins = namespace["AxolotlTrainingMixins"]
+        return AxolotlTrainingMixins
+    return AxolotlTrainingMixinsBase
diff --git a/src/axolotl/integrations/cut_cross_entropy/ACKNOWLEDGEMENTS.md b/src/axolotl/integrations/cut_cross_entropy/ACKNOWLEDGEMENTS.md
new file mode 100644
index 0000000000..03d1cbfb00
--- /dev/null
+++ b/src/axolotl/integrations/cut_cross_entropy/ACKNOWLEDGEMENTS.md
@@ -0,0 +1,325 @@
+Acknowledgements
+
+Portions of this Cut Cross Entropy Software may utilize the following copyrighted
+material, the use of which is hereby acknowledged.
+
+
+------
+
+
+PyTorch
+
+    From PyTorch:
+
+    Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
+    Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
+    Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
+    Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
+    Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
+    Copyright (c) 2011-2013 NYU                      (Clement Farabet)
+    Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
+    Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
+    Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
+
+    From Caffe2:
+
+    Copyright (c) 2016-present, Facebook Inc. All rights reserved.
+
+    All contributions by Facebook:
+    Copyright (c) 2016 Facebook Inc.
+
+    All contributions by Google:
+    Copyright (c) 2015 Google Inc.
+    All rights reserved.
+
+    All contributions by Yangqing Jia:
+    Copyright (c) 2015 Yangqing Jia
+    All rights reserved.
+
+    All contributions by Kakao Brain:
+    Copyright 2019-2020 Kakao Brain
+
+    All contributions by Cruise LLC:
+    Copyright (c) 2022 Cruise LLC.
+    All rights reserved.
+
+    All contributions by Arm:
+    Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates
+
+    All contributions from Caffe:
+    Copyright(c) 2013, 2014, 2015, the respective contributors
+    All rights reserved.
+
+    All other contributions:
+    Copyright(c) 2015, 2016 the respective contributors
+    All rights reserved.
+
+    Caffe2 uses a copyright model similar to Caffe: each contributor holds
+    copyright over their contributions to Caffe2. The project versioning records
+    all such contribution and copyright details. If a contributor wants to further
+    mark their specific copyright on a particular contribution, they should
+    indicate their copyright solely in the commit message of the change when it is
+    committed.
+
+    All rights reserved.
+
+    Redistribution and use in source and binary forms, with or without
+    modification, are permitted provided that the following conditions are met:
+
+    1. Redistributions of source code must retain the above copyright
+    notice, this list of conditions and the following disclaimer.
+
+    2. Redistributions in binary form must reproduce the above copyright
+    notice, this list of conditions and the following disclaimer in the
+    documentation and/or other materials provided with the distribution.
+
+    3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
+    and IDIAP Research Institute nor the names of its contributors may be
+    used to endorse or promote products derived from this software without
+    specific prior written permission.
+
+    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+    AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+    IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+    ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+    LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+    CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+    SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+    INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+    CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+    ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+    POSSIBILITY OF SUCH DAMAGE.
+
+
+Triton
+
+    /*
+    * Copyright 2018-2020 Philippe Tillet
+    * Copyright 2020-2022 OpenAI
+    *
+    * Permission is hereby granted, free of charge, to any person obtaining
+    * a copy of this software and associated documentation files
+    * (the "Software"), to deal in the Software without restriction,
+    * including without limitation the rights to use, copy, modify, merge,
+    * publish, distribute, sublicense, and/or sell copies of the Software,
+    * and to permit persons to whom the Software is furnished to do so,
+    * subject to the following conditions:
+    *
+    * The above copyright notice and this permission notice shall be
+    * included in all copies or substantial portions of the Software.
+    *
+    * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+    * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+    * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+    * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+    * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+    * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+    * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+    */
+
+
+Transformers
+
+    Copyright 2018- The Hugging Face team. All rights reserved.
+
+                                    Apache License
+                            Version 2.0, January 2004
+                            http://www.apache.org/licenses/
+
+    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+    1. Definitions.
+
+        "License" shall mean the terms and conditions for use, reproduction,
+        and distribution as defined by Sections 1 through 9 of this document.
+
+        "Licensor" shall mean the copyright owner or entity authorized by
+        the copyright owner that is granting the License.
+
+        "Legal Entity" shall mean the union of the acting entity and all
+        other entities that control, are controlled by, or are under common
+        control with that entity. For the purposes of this definition,
+        "control" means (i) the power, direct or indirect, to cause the
+        direction or management of such entity, whether by contract or
+        otherwise, or (ii) ownership of fifty percent (50%) or more of the
+        outstanding shares, or (iii) beneficial ownership of such entity.
+
+        "You" (or "Your") shall mean an individual or Legal Entity
+        exercising permissions granted by this License.
+
+        "Source" form shall mean the preferred form for making modifications,
+        including but not limited to software source code, documentation
+        source, and configuration files.
+
+        "Object" form shall mean any form resulting from mechanical
+        transformation or translation of a Source form, including but
+        not limited to compiled object code, generated documentation,
+        and conversions to other media types.
+
+        "Work" shall mean the work of authorship, whether in Source or
+        Object form, made available under the License, as indicated by a
+        copyright notice that is included in or attached to the work
+        (an example is provided in the Appendix below).
+
+        "Derivative Works" shall mean any work, whether in Source or Object
+        form, that is based on (or derived from) the Work and for which the
+        editorial revisions, annotations, elaborations, or other modifications
+        represent, as a whole, an original work of authorship. For the purposes
+        of this License, Derivative Works shall not include works that remain
+        separable from, or merely link (or bind by name) to the interfaces of,
+        the Work and Derivative Works thereof.
+
+        "Contribution" shall mean any work of authorship, including
+        the original version of the Work and any modifications or additions
+        to that Work or Derivative Works thereof, that is intentionally
+        submitted to Licensor for inclusion in the Work by the copyright owner
+        or by an individual or Legal Entity authorized to submit on behalf of
+        the copyright owner. For the purposes of this definition, "submitted"
+        means any form of electronic, verbal, or written communication sent
+        to the Licensor or its representatives, including but not limited to
+        communication on electronic mailing lists, source code control systems,
+        and issue tracking systems that are managed by, or on behalf of, the
+        Licensor for the purpose of discussing and improving the Work, but
+        excluding communication that is conspicuously marked or otherwise
+        designated in writing by the copyright owner as "Not a Contribution."
+
+        "Contributor" shall mean Licensor and any individual or Legal Entity
+        on behalf of whom a Contribution has been received by Licensor and
+        subsequently incorporated within the Work.
+
+    2. Grant of Copyright License. Subject to the terms and conditions of
+        this License, each Contributor hereby grants to You a perpetual,
+        worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+        copyright license to reproduce, prepare Derivative Works of,
+        publicly display, publicly perform, sublicense, and distribute the
+        Work and such Derivative Works in Source or Object form.
+
+    3. Grant of Patent License. Subject to the terms and conditions of
+        this License, each Contributor hereby grants to You a perpetual,
+        worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+        (except as stated in this section) patent license to make, have made,
+        use, offer to sell, sell, import, and otherwise transfer the Work,
+        where such license applies only to those patent claims licensable
+        by such Contributor that are necessarily infringed by their
+        Contribution(s) alone or by combination of their Contribution(s)
+        with the Work to which such Contribution(s) was submitted. If You
+        institute patent litigation against any entity (including a
+        cross-claim or counterclaim in a lawsuit) alleging that the Work
+        or a Contribution incorporated within the Work constitutes direct
+        or contributory patent infringement, then any patent licenses
+        granted to You under this License for that Work shall terminate
+        as of the date such litigation is filed.
+
+    4. Redistribution. You may reproduce and distribute copies of the
+        Work or Derivative Works thereof in any medium, with or without
+        modifications, and in Source or Object form, provided that You
+        meet the following conditions:
+
+        (a) You must give any other recipients of the Work or
+            Derivative Works a copy of this License; and
+
+        (b) You must cause any modified files to carry prominent notices
+            stating that You changed the files; and
+
+        (c) You must retain, in the Source form of any Derivative Works
+            that You distribute, all copyright, patent, trademark, and
+            attribution notices from the Source form of the Work,
+            excluding those notices that do not pertain to any part of
+            the Derivative Works; and
+
+        (d) If the Work includes a "NOTICE" text file as part of its
+            distribution, then any Derivative Works that You distribute must
+            include a readable copy of the attribution notices contained
+            within such NOTICE file, excluding those notices that do not
+            pertain to any part of the Derivative Works, in at least one
+            of the following places: within a NOTICE text file distributed
+            as part of the Derivative Works; within the Source form or
+            documentation, if provided along with the Derivative Works; or,
+            within a display generated by the Derivative Works, if and
+            wherever such third-party notices normally appear. The contents
+            of the NOTICE file are for informational purposes only and
+            do not modify the License. You may add Your own attribution
+            notices within Derivative Works that You distribute, alongside
+            or as an addendum to the NOTICE text from the Work, provided
+            that such additional attribution notices cannot be construed
+            as modifying the License.
+
+        You may add Your own copyright statement to Your modifications and
+        may provide additional or different license terms and conditions
+        for use, reproduction, or distribution of Your modifications, or
+        for any such Derivative Works as a whole, provided Your use,
+        reproduction, and distribution of the Work otherwise complies with
+        the conditions stated in this License.
+
+    5. Submission of Contributions. Unless You explicitly state otherwise,
+        any Contribution intentionally submitted for inclusion in the Work
+        by You to the Licensor shall be under the terms and conditions of
+        this License, without any additional terms or conditions.
+        Notwithstanding the above, nothing herein shall supersede or modify
+        the terms of any separate license agreement you may have executed
+        with Licensor regarding such Contributions.
+
+    6. Trademarks. This License does not grant permission to use the trade
+        names, trademarks, service marks, or product names of the Licensor,
+        except as required for reasonable and customary use in describing the
+        origin of the Work and reproducing the content of the NOTICE file.
+
+    7. Disclaimer of Warranty. Unless required by applicable law or
+        agreed to in writing, Licensor provides the Work (and each
+        Contributor provides its Contributions) on an "AS IS" BASIS,
+        WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+        implied, including, without limitation, any warranties or conditions
+        of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+        PARTICULAR PURPOSE. You are solely responsible for determining the
+        appropriateness of using or redistributing the Work and assume any
+        risks associated with Your exercise of permissions under this License.
+
+    8. Limitation of Liability. In no event and under no legal theory,
+        whether in tort (including negligence), contract, or otherwise,
+        unless required by applicable law (such as deliberate and grossly
+        negligent acts) or agreed to in writing, shall any Contributor be
+        liable to You for damages, including any direct, indirect, special,
+        incidental, or consequential damages of any character arising as a
+        result of this License or out of the use or inability to use the
+        Work (including but not limited to damages for loss of goodwill,
+        work stoppage, computer failure or malfunction, or any and all
+        other commercial damages or losses), even if such Contributor
+        has been advised of the possibility of such damages.
+
+    9. Accepting Warranty or Additional Liability. While redistributing
+        the Work or Derivative Works thereof, You may choose to offer,
+        and charge a fee for, acceptance of support, warranty, indemnity,
+        or other liability obligations and/or rights consistent with this
+        License. However, in accepting such obligations, You may act only
+        on Your own behalf and on Your sole responsibility, not on behalf
+        of any other Contributor, and only if You agree to indemnify,
+        defend, and hold each Contributor harmless for any liability
+        incurred by, or claims asserted against, such Contributor by reason
+        of your accepting any such warranty or additional liability.
+
+    END OF TERMS AND CONDITIONS
+
+    APPENDIX: How to apply the Apache License to your work.
+
+        To apply the Apache License to your work, attach the following
+        boilerplate notice, with the fields enclosed by brackets "[]"
+        replaced with your own identifying information. (Don't include
+        the brackets!)  The text should be enclosed in the appropriate
+        comment syntax for the file format. We also recommend that a
+        file or class name and description of purpose be included on the
+        same "printed page" as the copyright notice for easier
+        identification within third-party archives.
+
+    Copyright [yyyy] [name of copyright owner]
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+        http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
diff --git a/src/axolotl/integrations/cut_cross_entropy/LICENSE b/src/axolotl/integrations/cut_cross_entropy/LICENSE
new file mode 100644
index 0000000000..7ab90b7d9c
--- /dev/null
+++ b/src/axolotl/integrations/cut_cross_entropy/LICENSE
@@ -0,0 +1,47 @@
+Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+IMPORTANT:  This Apple software is supplied to you by Apple
+Inc. ("Apple") in consideration of your agreement to the following
+terms, and your use, installation, modification or redistribution of
+this Apple software constitutes acceptance of these terms.  If you do
+not agree with these terms, please do not use, install, modify or
+redistribute this Apple software.
+
+In consideration of your agreement to abide by the following terms, and
+subject to these terms, Apple grants you a personal, non-exclusive
+license, under Apple's copyrights in this original Apple software (the
+"Apple Software"), to use, reproduce, modify and redistribute the Apple
+Software, with or without modifications, in source and/or binary forms;
+provided that if you redistribute the Apple Software in its entirety and
+without modifications, you must retain this notice and the following
+text and disclaimers in all such redistributions of the Apple Software.
+Neither the name, trademarks, service marks or logos of Apple Inc. may
+be used to endorse or promote products derived from the Apple Software
+without specific prior written permission from Apple.  Except as
+expressly stated in this notice, no other rights or licenses, express or
+implied, are granted by Apple herein, including but not limited to any
+patent rights that may be infringed by your derivative works or by other
+works in which the Apple Software may be incorporated.
+
+The Apple Software is provided by Apple on an "AS IS" basis.  APPLE
+MAKES NO WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION
+THE IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS
+FOR A PARTICULAR PURPOSE, REGARDING THE APPLE SOFTWARE OR ITS USE AND
+OPERATION ALONE OR IN COMBINATION WITH YOUR PRODUCTS.
+
+IN NO EVENT SHALL APPLE BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL
+OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) ARISING IN ANY WAY OUT OF THE USE, REPRODUCTION,
+MODIFICATION AND/OR DISTRIBUTION OF THE APPLE SOFTWARE, HOWEVER CAUSED
+AND WHETHER UNDER THEORY OF CONTRACT, TORT (INCLUDING NEGLIGENCE),
+STRICT LIABILITY OR OTHERWISE, EVEN IF APPLE HAS BEEN ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.
+
+
+-------------------------------------------------------------------------------
+SOFTWARE DISTRIBUTED WITH CUT CROSS ENTROPY:
+
+The Cut Cross Entropy software includes a number of subcomponents with separate
+copyright notices and license terms - please see the file ACKNOWLEDGEMENTS.md.
+-------------------------------------------------------------------------------
diff --git a/src/axolotl/integrations/cut_cross_entropy/README.md b/src/axolotl/integrations/cut_cross_entropy/README.md
new file mode 100644
index 0000000000..6ec8e7cc5a
--- /dev/null
+++ b/src/axolotl/integrations/cut_cross_entropy/README.md
@@ -0,0 +1,131 @@
+# Cut Cross Entropy
+
+Cut Cross Entropy (CCE) reduces VRAM usage through optimization on the cross-entropy operation during loss calculation.
+
+See https://github.com/apple/ml-cross-entropy
+
+## Requirements
+
+- PyTorch 2.4.0 or higher
+
+## Installation
+
+Run the following command to install `cut_cross_entropy[transformers]` if you don't have it already.
+
+- If you are in dev environment
+```bash
+python scripts/cutcrossentropy_install.py | sh
+```
+
+- If you are installing from pip
+```bash
+pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@5f0c7a7"
+```
+
+## Usage
+
+```yaml
+plugins:
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+```
+
+## Supported Models
+
+- afmoe
+- apertus
+- arcee
+- cohere
+- cohere2
+- cohere2_moe
+- deepseek_v2
+- deepseek_v3
+- deepseek_v4
+- exaone4
+- exaone4_5
+- exaone_moe
+- gemma
+- gemma2
+- gemma3
+- gemma3_text
+- gemma3n
+- gemma3n_text
+- gemma4
+- gemma4_text
+- gemma4_unified
+- gemma4_unified_text
+- glm
+- glm4
+- glm4_moe
+- glm4_moe_lite
+- glm46v
+- glm4v
+- glm4v_moe
+- glm_image
+- glm_moe_dsa
+- gpt_oss
+- granite
+- granitemoe
+- granitemoehybrid
+- granitemoeshared
+- hunyuan_v1_dense
+- hunyuan_v1_moe
+- internvl
+- kimi_linear
+- lfm2
+- lfm2_moe
+- lfm2_vl
+- llama
+- llama4
+- llama4_text
+- llava
+- minimax
+- minimax_m2
+- ministral
+- ministral3
+- mistral
+- mistral3
+- mistral4
+- mixtral
+- mllama
+- nemotron_h
+- olmo
+- olmo2
+- olmo3
+- olmoe
+- phi
+- phi3
+- phi4_multimodal
+- qwen2
+- qwen2_5_vl
+- qwen2_moe
+- qwen2_vl
+- qwen3
+- qwen3_5
+- qwen3_5_text
+- qwen3_5_moe
+- qwen3_5_moe_text
+- qwen3_moe
+- qwen3_next
+- qwen3_vl
+- qwen3_vl_moe
+- seed_oss
+- smollm3
+- step3p5
+- step3p7
+- voxtral
+
+## Citation
+
+```bib
+@article{wijmans2024cut,
+  author       = {Erik Wijmans and
+                  Brody Huval and
+                  Alexander Hertzberg and
+                  Vladlen Koltun and
+                  Philipp Kr\"ahenb\"uhl},
+  title        = {Cut Your Losses in Large-Vocabulary Language Models},
+  journal      = {arXiv},
+  year         = {2024},
+  url          = {https://arxiv.org/abs/2411.09009},
+}
+```
diff --git a/src/axolotl/integrations/cut_cross_entropy/__init__.py b/src/axolotl/integrations/cut_cross_entropy/__init__.py
new file mode 100644
index 0000000000..94a6938c7a
--- /dev/null
+++ b/src/axolotl/integrations/cut_cross_entropy/__init__.py
@@ -0,0 +1,151 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Module for the Plugin for Cut Cross Entropy integration with Axolotl.
+
+Cut Cross Entropy is an optimized implementation of cross entropy loss
+from Apple's ML team.
+"""
+
+import importlib
+from functools import partial
+
+import torch
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils import get_pytorch_version
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
+from axolotl.utils.logging import get_logger
+
+from .args import CutCrossEntropyArgs as CutCrossEntropyArgs
+
+LOG = get_logger(__name__)
+
+_CCE_INSTALL_MESSAGE = (
+    "Please install Axolotl's fork of cut_cross_entropy with transformers support using "
+    '`pip uninstall -y cut-cross-entropy && pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@5f0c7a7"`'
+)
+
+
+class CutCrossEntropyPlugin(BasePlugin):
+    """
+    Plugin for Cut Cross Entropy integration with Axolotl.
+    """
+
+    def get_input_args(self):
+        return "axolotl.integrations.cut_cross_entropy.CutCrossEntropyArgs"
+
+    def _check_requirements(self):
+        """Check if all requirements are met."""
+        # Check PyTorch version
+
+        major, minor, _ = get_pytorch_version()
+        if (major, minor) < (2, 4):
+            raise ImportError(
+                "Cut Cross Entropy requires PyTorch >= 2.4.0. "
+                f"Current version: {torch.__version__}"
+            )
+
+        # Check if cut_cross_entropy is installed
+        cce_spec = importlib.util.find_spec("cut_cross_entropy")
+        if cce_spec is None:
+            raise ImportError(_CCE_INSTALL_MESSAGE)
+
+        cce_spec_transformers = importlib.util.find_spec(
+            "cut_cross_entropy.transformers"
+        )
+        if cce_spec_transformers is None:
+            raise ImportError(
+                "Transformers support is not installed. " + _CCE_INSTALL_MESSAGE
+            )
+
+        # Check if Axolotl's cce fork is installed
+        try:
+            from cut_cross_entropy.transformers.patch import AXOLOTL_CCE_FORK
+
+            if not AXOLOTL_CCE_FORK:
+                raise ImportError
+        except ImportError as e:
+            raise ImportError(
+                "Axolotl's fork of cut_cross_entropy is not installed. "
+                + _CCE_INSTALL_MESSAGE
+            ) from e
+
+    def pre_model_load(self, cfg):
+        """Apply cut cross entropy before model loading if enabled."""
+        if cfg.cut_cross_entropy:
+            self._check_requirements()
+            self.patch_llama_like(cfg.model_config_type)
+
+            from cut_cross_entropy.transformers.patch import cce_patch
+
+            LOG.info(
+                f"Applying Cut Cross Entropy to model type: {cfg.model_config_type}"
+            )
+
+            # The patch checks model_type internally
+
+            cce_patch(
+                cfg.model_config_type,
+                remote_model_id=cfg.base_model if cfg.trust_remote_code else None,
+            )
+
+    def patch_llama_like(
+        self,
+        model_type_to_patch: str,
+    ) -> None:
+        """
+        Generic patch for model architectures with causal lm similar to llama
+        """
+        from cut_cross_entropy.transformers.patch import PATCH_FNS
+
+        def patch_generic(
+            maybe_model,
+            patch_options,
+            remote_model_id: str | None,
+            model_type: str,
+        ):
+            import cut_cross_entropy.transformers.llama
+            from cut_cross_entropy.transformers.llama import cce_forward
+
+            try:
+                # Dynamically import the module and CausalLM class
+                module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+                model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+                module = __import__(
+                    module_path, fromlist=[f"{model_cls_prefix}ForCausalLM"]
+                )
+                model_cls = getattr(module, f"{model_cls_prefix}ForCausalLM")
+
+                cut_cross_entropy.transformers.llama._PATCH_OPTS = patch_options
+
+                model_cls.forward = cce_forward
+
+            except (ImportError, AttributeError) as e:
+                raise RuntimeError(
+                    f"Could not import ForCausalLM class for model_type: {model_type}. "
+                    f"Error: {str(e)}"
+                ) from e
+
+        if model_type_to_patch not in PATCH_FNS:
+            LOG.warning_once(
+                "Setting up generic cce patch for model type: %s", model_type_to_patch
+            )
+            LOG.warning_once(
+                f"Generic Cut Cross Entropy + {model_type_to_patch} support is experimental and may not work as expected."
+            )
+            PATCH_FNS[model_type_to_patch] = partial(
+                patch_generic, model_type=model_type_to_patch
+            )
diff --git a/src/axolotl/integrations/cut_cross_entropy/args.py b/src/axolotl/integrations/cut_cross_entropy/args.py
new file mode 100644
index 0000000000..3eeb9fac73
--- /dev/null
+++ b/src/axolotl/integrations/cut_cross_entropy/args.py
@@ -0,0 +1,54 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Module for handling Cut Cross Entropy input arguments.
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel, model_validator
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class CutCrossEntropyArgs(BaseModel):
+    """
+    Input args for Cut Cross Entropy.
+    """
+
+    cut_cross_entropy: Optional[bool] = True
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_dtype_is_half(cls, data):
+        if data.get("cut_cross_entropy") and not (data.get("bf16") or data.get("fp16")):
+            raise ValueError(
+                "Cut Cross Entropy requires fp16/bf16 training for backward pass. "
+                "Please set `bf16` or `fp16` to `True`."
+            )
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_chunked_cross_entropy_not_set(cls, data):
+        if data.get("chunked_cross_entropy"):
+            raise ValueError(
+                "Cut Cross Entropy does not support chunked cross entropy. "
+                "Please set `chunked_cross_entropy` to `False` or disable Cut Cross Entropy."
+            )
+        return data
diff --git a/src/axolotl/integrations/densemixer/README.md b/src/axolotl/integrations/densemixer/README.md
new file mode 100644
index 0000000000..62da1bb072
--- /dev/null
+++ b/src/axolotl/integrations/densemixer/README.md
@@ -0,0 +1,12 @@
+# DenseMixer
+
+See [DenseMixer](https://github.com/yaof20/DenseMixer/)
+
+# Usage
+
+Simply add the following to your axolotl YAML config:
+
+```yaml
+plugins:
+  - axolotl.integrations.densemixer.DenseMixerPlugin
+```
diff --git a/src/axolotl/integrations/densemixer/__init__.py b/src/axolotl/integrations/densemixer/__init__.py
new file mode 100644
index 0000000000..901bdc1c11
--- /dev/null
+++ b/src/axolotl/integrations/densemixer/__init__.py
@@ -0,0 +1,5 @@
+"""Integration entry point for the DenseMixer plugin."""
+
+from .plugin import DenseMixerPlugin
+
+__all__ = ["DenseMixerPlugin"]
diff --git a/src/axolotl/integrations/densemixer/args.py b/src/axolotl/integrations/densemixer/args.py
new file mode 100644
index 0000000000..c8bf549319
--- /dev/null
+++ b/src/axolotl/integrations/densemixer/args.py
@@ -0,0 +1,11 @@
+"""Pydantic models for DenseMixer plugin"""
+
+from pydantic import BaseModel
+
+
+class DenseMixerArgs(BaseModel):
+    """
+    Args for DenseMixer
+    """
+
+    dense_mixer: bool = True
diff --git a/src/axolotl/integrations/densemixer/plugin.py b/src/axolotl/integrations/densemixer/plugin.py
new file mode 100644
index 0000000000..9548fd19ac
--- /dev/null
+++ b/src/axolotl/integrations/densemixer/plugin.py
@@ -0,0 +1,42 @@
+"""DenseMixer plugin for Axolotl"""
+
+import importlib
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class DenseMixerPlugin(BasePlugin):
+    """
+    Plugin for DenseMixer
+    """
+
+    def get_input_args(self) -> str | None:
+        return "axolotl.integrations.densemixer.args.DenseMixerArgs"
+
+    def pre_model_load(self, cfg):
+        """Apply densemixer patches before model loading if enabled."""
+        if cfg.dense_mixer:
+            if not importlib.util.find_spec("densemixer"):
+                raise RuntimeError(
+                    "DenseMixer is not installed. Install it with `pip install densemixer`"
+                )
+
+            from densemixer.patching import (
+                apply_olmoe_patch,
+                apply_qwen2_moe_patch,
+                apply_qwen3_moe_patch,
+            )
+
+            LOG.info(
+                f"Applying DenseMixer patches for model type: {cfg.model_config_type}"
+            )
+
+            if cfg.model_config_type == "olmoe":
+                apply_olmoe_patch()
+            if cfg.model_config_type == "qwen2_moe":
+                apply_qwen2_moe_patch()
+            if cfg.model_config_type == "qwen3_moe":
+                apply_qwen3_moe_patch()
diff --git a/src/axolotl/integrations/diffusion/README.md b/src/axolotl/integrations/diffusion/README.md
new file mode 100644
index 0000000000..c27f33de15
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/README.md
@@ -0,0 +1,154 @@
+# Diffusion LM Training Plugin for Axolotl
+
+This plugin enables diffusion language model training using an approach inspired by
+LLaDA (Large Language Diffusion Models) within Axolotl.
+
+## Overview
+
+LLaDA is a diffusion-based approach to language model training that uses:
+- **Random token masking** during training instead of next-token prediction
+- **Bidirectional attention** to allow the model to attend to the full context
+- **Importance weighting** based on masking probabilities for stable training
+
+This approach can lead to more robust language models with better understanding of
+bidirectional context.
+
+## Installation
+
+The plugin is included with Axolotl. See our
+[installation docs](https://docs.axolotl.ai/docs/installation.html).
+
+## Quickstart
+
+Train with an example config (Llama‑3.2 1B):
+   - Pretrain: `axolotl train examples/llama-3/diffusion-3.2-1b-pretrain.yaml`
+   - SFT: `axolotl train examples/llama-3/diffusion-3.2-1b-sft.yaml`
+
+### Basic Configuration
+
+You can also modify your existing configs to enable / customize diffusion training.
+
+Add the following to your Axolotl config:
+
+```yaml
+# Enable diffusion LM training plugin
+plugins:
+  - axolotl.integrations.diffusion.DiffusionPlugin
+```
+
+And, configure the nested `diffusion` block (defaults shown):
+
+```yaml
+diffusion:
+  noise_schedule: linear  # or "cosine"
+  min_mask_ratio: 0.1
+  max_mask_ratio: 0.9
+  num_diffusion_steps: 128
+  eps: 1e-3
+  importance_weighting: true
+
+  # Mask token (training auto-adds if missing, avoid pad/eos)
+  mask_token_str: "<|diffusion_mask|>"
+  # Or use an existing special token id (e.g., 128002 for Llama-3.x)
+  # mask_token_id: 128002
+
+  # Sample generation during training (optional)
+  generate_samples: true
+  generation_interval: 100
+  num_generation_samples: 3
+  generation_steps: 128
+  generation_temperature: 0.0
+  generation_max_length: 100
+```
+
+## Supported Models
+
+Any models that support 4D attention masks should work out of the box. If not, please
+create an [issue](https://github.com/axolotl-ai-cloud/axolotl/issues) or open a
+[PR](https://github.com/axolotl-ai-cloud/axolotl/compare)!
+
+## How It Works
+
+### Random Masking
+During training, tokens are randomly masked:
+- Sample timestep `t` uniformly from [0, 1]
+- Calculate masking probability: `p = (1 - eps) * t + eps`
+- Randomly mask tokens with probability `p`
+
+### Diffusion Loss
+
+Loss is computed only on masked tokens with (optional) importance weighting:
+
+```python
+loss = sum(cross_entropy(pred, target) / p_mask) / total_tokens
+```
+
+## Sample Generation
+
+When `diffusion.generate_samples: true`, the plugin generates samples during training:
+
+```
+Sample 1:
+   Original (45 tokens): The quick brown fox jumps over the lazy dog...
+   Masked (18/45 tokens, 40.0%): The [MASK] [MASK] fox [MASK] over [MASK] lazy [MASK]...
+   Generated: The quick brown fox jumps over the lazy dog...
+```
+
+Samples are logged to console and wandb (if enabled).
+
+## Inference
+
+Diffusion inference is integrated into the standard Axolotl CLI. Use the same config
+you trained with and run:
+
+```
+axolotl inference path/to/your-config.yaml
+```
+
+Optionally, pass `--gradio` to use a simple web interface.
+
+Interactive controls (prefix the prompt with commands):
+- `:complete N` → completion mode with N new masked tokens appended (default 64)
+- `:mask R` → random masking mode with target mask ratio R in [0.0, 1.0]
+
+Example session:
+
+```
+================================================================================
+Commands:
+:complete N -> completion mode with N tokens (default 64)
+:mask R     -> random masking with ratio R (0.0–1.0)
+================================================================================
+Give me an instruction (Ctrl + D to submit):
+
+:mask 0.4 The quick brown fox jumps over the lazy dog
+
+Masked (40.0%):
+The [MASK] brown [MASK] jumps over the [MASK] dog
+
+Generated:
+The quick brown fox jumps over the loud dog
+```
+
+## Metrics and Monitoring
+
+The plugin adds (or modifies) several metrics to track diffusion training:
+
+- `train/loss`: Weighted diffusion loss
+- `train/accuracy`: Accuracy on masked tokens
+- `train/mask_ratio`: Average fraction of tokens masked
+- `train/num_masked_tokens`: Number of tokens masked
+- `train/avg_p_mask`: Average masking probability
+- `train/ce_loss`: Unweighted cross-entropy loss
+- `train/importance_weight_avg`: Average importance weight
+
+## Limitations
+
+- No flash attention support
+- No RL training support
+
+## References
+
+- [LLaDA Paper](https://arxiv.org/abs/2404.10406)
+- [Axolotl Documentation](https://docs.axolotl.ai/)
+- [API reference for plugin](https://docs.axolotl.ai/docs/api/integrations.diffusion.args.html#axolotl.integrations.diffusion.args)
diff --git a/src/axolotl/integrations/diffusion/__init__.py b/src/axolotl/integrations/diffusion/__init__.py
new file mode 100644
index 0000000000..9e38cc5c1d
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/__init__.py
@@ -0,0 +1,19 @@
+"""Diffusion LM training plugin init."""
+
+from .args import DiffusionArgs, DiffusionConfig
+from .callbacks import DiffusionGenerationCallback
+from .generation import generate
+from .plugin import DiffusionPlugin
+from .trainer import DiffusionTrainer
+from .utils import create_bidirectional_attention_mask, resolve_mask_token_id
+
+__all__ = [
+    "DiffusionArgs",
+    "DiffusionPlugin",
+    "DiffusionTrainer",
+    "generate",
+    "resolve_mask_token_id",
+    "create_bidirectional_attention_mask",
+    "DiffusionGenerationCallback",
+    "DiffusionConfig",
+]
diff --git a/src/axolotl/integrations/diffusion/args.py b/src/axolotl/integrations/diffusion/args.py
new file mode 100644
index 0000000000..4f5bfe499f
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/args.py
@@ -0,0 +1,95 @@
+"""Config args for diffusion LM training (nested under `diffusion:`)."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class DiffusionConfig(BaseModel):
+    """Nested diffusion configuration available under the `diffusion` key."""
+
+    # Noise schedule config
+    noise_schedule: Literal["linear", "cosine"] = Field(
+        default="linear", description="Type of noise schedule for diffusion training"
+    )
+    min_mask_ratio: float = Field(
+        default=0.1,
+        ge=0.0,
+        le=1.0,
+        description="Minimum masking ratio for diffusion noise schedule",
+    )
+    max_mask_ratio: float = Field(
+        default=0.9,
+        ge=0.0,
+        le=1.0,
+        description="Maximum masking ratio for diffusion noise schedule",
+    )
+    num_diffusion_steps: int = Field(
+        default=128, ge=1, description="Number of diffusion timesteps"
+    )
+    eps: float = Field(
+        default=1e-3,
+        ge=0.0,
+        le=1.0,
+        description="Epsilon value for minimum masking probability in forward process",
+    )
+
+    # Training config
+    importance_weighting: bool = Field(
+        default=True,
+        description="Apply importance weighting to loss based on masking probability",
+    )
+    mask_token_id: int | None = Field(
+        default=None,
+        description=(
+            "Token ID to use for masking. Unset by default; can use one of the "
+            "tokenizer's special tokens here."
+        ),
+    )
+    mask_token_str: str | None = Field(
+        default=None,
+        description=(
+            "Token string to use as a mask. If `mask_token_id` is invalid or unset, "
+            "this token will be ensured to exist as an additional special token and "
+            "used. If absent, a default '<|diffusion_mask|>' will be added."
+        ),
+    )
+
+    # Sample generation config
+    generate_samples: bool = Field(
+        default=True, description="Enable sample generation during training"
+    )
+    generation_interval: int = Field(
+        default=100, ge=1, description="Generate samples every N steps"
+    )
+    num_generation_samples: int = Field(
+        default=3, ge=1, description="Number of samples to generate each time"
+    )
+    generation_steps: int = Field(
+        default=128, ge=1, description="Number of diffusion steps for generation"
+    )
+    generation_temperature: float = Field(
+        default=0.0,
+        ge=0.0,
+        description="Temperature for generation sampling (0.0 = deterministic)",
+    )
+    generation_max_length: int = Field(
+        default=100, ge=1, description="Maximum sequence length for generation"
+    )
+
+    @model_validator(mode="after")
+    def _validate_mask_ratios(self) -> "DiffusionConfig":
+        if self.min_mask_ratio > self.max_mask_ratio:
+            raise ValueError("min_mask_ratio must be ≤ max_mask_ratio")
+        return self
+
+
+class DiffusionArgs(BaseModel):
+    """Plugin entry that exposes the nested `diffusion` block to the core config."""
+
+    diffusion: DiffusionConfig = Field(
+        default_factory=DiffusionConfig,
+        description="Diffusion training configuration. Only nested block is supported.",
+    )
diff --git a/src/axolotl/integrations/diffusion/callbacks.py b/src/axolotl/integrations/diffusion/callbacks.py
new file mode 100644
index 0000000000..61464ee7df
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/callbacks.py
@@ -0,0 +1,176 @@
+"""Callbacks for diffusion training."""
+
+import logging
+import sys
+
+import wandb
+from colorama import Fore, Style
+from transformers.trainer_callback import TrainerCallback, TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+
+from .generation import generate_samples
+
+# Simpler logger for more readable sample generation
+logger = logging.getLogger(__name__)
+if not logger.handlers:
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setFormatter(logging.Formatter("%(message)s"))
+    logger.addHandler(handler)
+    logger.propagate = False
+logger.setLevel(logging.INFO)
+
+
+class DiffusionGenerationCallback(TrainerCallback):
+    """Callback for generating samples during diffusion training."""
+
+    def __init__(self, trainer):
+        self.trainer = trainer
+
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Generate samples at specified intervals."""
+        if (
+            state.global_step > 0
+            and state.global_step
+            % self.trainer.axolotl_cfg.diffusion.generation_interval
+            == 0
+        ):
+            if not self.trainer.state.is_world_process_zero:
+                return
+
+            # Use eval dataloader if available, otherwise use train dataloader
+            dataloader = None
+            try:
+                if getattr(self.trainer, "eval_dataset", None) is not None:
+                    dataloader = self.trainer.get_eval_dataloader()
+            except Exception:
+                dataloader = None
+            if dataloader is None:
+                dataloader = self.trainer.get_train_dataloader()
+
+            # Generate samples
+            diffusion_cfg = self.trainer.axolotl_cfg.diffusion
+            samples = generate_samples(
+                model=self.trainer.model,
+                tokenizer=self.trainer.processing_class,
+                dataloader=dataloader,
+                num_generation_samples=diffusion_cfg.num_generation_samples,
+                max_length=diffusion_cfg.generation_max_length,
+                num_diffusion_steps=diffusion_cfg.generation_steps,
+                temperature=diffusion_cfg.generation_temperature,
+                mask_token_id=diffusion_cfg.mask_token_id,
+            )
+
+            # Log samples
+            self._log_samples(samples, state.global_step)
+
+    def _log_samples(self, samples: list, step: int):
+        """Log generated samples."""
+        if not samples:
+            return
+
+        logger.info("=" * 60)
+        logger.info("GENERATED SAMPLES")
+        logger.info("=" * 60)
+
+        for i, sample_data in enumerate(samples, 1):
+            original = sample_data["original"]
+            masked = sample_data["masked"]
+            generated = sample_data["generated"]
+            mask_ratio = sample_data["mask_ratio"]
+            masked_tokens = sample_data["masked_tokens"]
+            total_tokens = sample_data["total_tokens"]
+
+            logger.info(f"\nSample {i}:")
+            logger.info(f"\tOriginal ({total_tokens} tokens): {original}")
+            logger.info(
+                f"\tMasked ({masked_tokens}/{total_tokens} tokens, "
+                f"{mask_ratio:.1%}): {masked}"
+            )
+
+            try:
+                gen_ids = sample_data.get("generated_ids")
+                orig_ids = sample_data.get("orig_ids")
+                masked_positions = set(sample_data.get("masked_positions") or [])
+                if isinstance(gen_ids, list) and isinstance(orig_ids, list):
+                    styles: list[str] = []
+                    for i, tid in enumerate(gen_ids):
+                        if i in masked_positions:
+                            if i < len(orig_ids) and tid == orig_ids[i]:
+                                styles.append("green")
+                            elif i < len(orig_ids):
+                                styles.append("red")
+                            else:
+                                styles.append("normal")
+                        else:
+                            same = i < len(orig_ids) and tid == orig_ids[i]
+                            styles.append("dim" if same else "normal")
+
+                    spans: list[tuple[str, int, int]] = []
+                    if gen_ids:
+                        cur = styles[0]
+                        start = 0
+                        for i in range(1, len(gen_ids)):
+                            s = styles[i]
+                            if s != cur:
+                                spans.append((cur, start, i))
+                                cur, start = s, i
+                        spans.append((cur, start, len(gen_ids)))
+
+                    parts = []
+                    for style_name, a, b in spans:
+                        chunk_text = self.trainer.processing_class.decode(
+                            gen_ids[a:b], skip_special_tokens=False
+                        )
+                        if style_name == "green":
+                            parts.append(Fore.GREEN + chunk_text + Style.RESET_ALL)
+                        elif style_name == "red":
+                            parts.append(Fore.RED + chunk_text + Style.RESET_ALL)
+                        else:
+                            if style_name == "dim":
+                                parts.append(Style.DIM + chunk_text + Style.RESET_ALL)
+                            else:
+                                parts.append(chunk_text)
+                    logger.info("\tGenerated:\n%s", "".join(parts))
+                else:
+                    logger.info(f"\tGenerated: {generated}")
+            except Exception:
+                logger.info(f"\tGenerated: {generated}")
+
+        logger.info("=" * 60)
+
+        if self.trainer.axolotl_cfg.use_wandb:
+            if wandb.run is not None:  # type: ignore[attr-defined]
+                wandb.log(  # type: ignore[attr-defined]
+                    {
+                        "generated_samples": wandb.Table(  # type: ignore[attr-defined]
+                            columns=[
+                                "step",
+                                "original",
+                                "masked",
+                                "generated",
+                                "mask_ratio",
+                                "masked_tokens",
+                                "total_tokens",
+                            ],
+                            data=[
+                                [
+                                    step,
+                                    sample["original"],
+                                    sample["masked"],
+                                    sample["generated"],
+                                    f"{sample['mask_ratio']:.1%}",
+                                    sample["masked_tokens"],
+                                    sample["total_tokens"],
+                                ]
+                                for sample in samples
+                            ],
+                        )
+                    },
+                    step=step,
+                )
diff --git a/src/axolotl/integrations/diffusion/generation.py b/src/axolotl/integrations/diffusion/generation.py
new file mode 100644
index 0000000000..ec517fd238
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/generation.py
@@ -0,0 +1,409 @@
+"""Sample generation utilities for diffusion training."""
+
+import re
+from typing import Any, List, Literal, Optional
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+from .utils import create_bidirectional_attention_mask, shift_logits_to_input_positions
+
+LOG = get_logger(__name__)
+
+
+def generate_samples(
+    model: torch.nn.Module,
+    tokenizer: Any,
+    dataloader: Optional[Any] = None,
+    num_generation_samples: int = 3,
+    max_length: int = 100,
+    num_diffusion_steps: int = 128,
+    temperature: float = 0.0,
+    mask_token_id: int = 32000,
+    mode: Literal["random", "completion"] = "random",
+    completion_tokens: int = 0,
+    target_mask_ratio: Optional[float] = None,
+) -> List[dict]:
+    """
+    Generate text samples using the diffusion model by randomly masking sequences from
+    the given dataset and running the reverse diffusion process.
+
+    Args:
+        model: The wrapped or unwrapped model
+        tokenizer: Tokenizer for encoding/decoding
+        dataloader: Validation dataloader (for sampling sequences)
+        num_generation_samples: Number of samples to generate
+        max_length: Maximum length of sequences to use
+        num_diffusion_steps: Number of diffusion steps for generation
+        temperature: Temperature for sampling (0.0 = deterministic)
+        mask_token_id: Token ID used for masking
+
+    Returns:
+        List of dictionaries with original text, masked text, and generated text
+    """
+    if dataloader is None:
+        LOG.warning("No validation dataloader provided, cannot generate samples")
+        return []
+
+    unwrapped_model = model.module if hasattr(model, "module") else model
+    training = unwrapped_model.training
+    unwrapped_model.eval()
+
+    # Resolve device robustly (some modules don't expose `.device`)
+    device = getattr(unwrapped_model, "device", None)
+    if device is None:
+        try:
+            device = next(unwrapped_model.parameters()).device
+        except StopIteration:
+            device = torch.device("cpu")
+    generations = []
+
+    # Sample sequences from validation dataset
+    sampled_sequences = _sample_sequences_from_dataloader(
+        dataloader, num_generation_samples, max_length, device
+    )
+    LOG.info(f"Sampled {len(sampled_sequences)} sequences from validation dataset")
+
+    # Generate samples using reverse diffusion process
+    with torch.no_grad():
+        for sample in sampled_sequences:
+            if isinstance(sample, dict):
+                original_sequence = sample.get("input_ids")
+                labels_seq = sample.get("labels")
+                attn_seq = sample.get("attention_mask")
+            else:
+                original_sequence = sample
+                labels_seq = None
+                attn_seq = None
+            generation_result = generate(
+                unwrapped_model,
+                tokenizer,
+                original_sequence,
+                num_diffusion_steps,
+                temperature,
+                mask_token_id,
+                mode=mode,
+                completion_tokens=completion_tokens,
+                target_mask_ratio=target_mask_ratio,
+                labels=labels_seq,
+                attention_mask=attn_seq,
+            )
+            generations.append(generation_result)
+
+    # Restore prior training state
+    if training:
+        unwrapped_model.train()
+    else:
+        unwrapped_model.eval()
+
+    return generations
+
+
+def _sample_sequences_from_dataloader(
+    dataloader: Any, num_samples: int, max_length: int, device: torch.device
+) -> List[Any]:
+    """Sample sequences from validation dataloader."""
+    sampled_sequences: list[dict[str, torch.Tensor] | torch.Tensor] = []
+    sample_count = 0
+
+    # Skip a random number of batches (we could be more clever about this)
+    skip_batches = torch.randint(0, 10, (1,)).item()
+    batch_count = 0
+
+    for batch in dataloader:
+        # Skip some batches for variety
+        if batch_count < skip_batches:
+            batch_count += 1
+            continue
+
+        if sample_count >= num_samples:
+            break
+
+        batch_count += 1
+        input_ids = batch["input_ids"]
+        attention_mask = batch.get("attention_mask")
+        labels = batch.get("labels")
+
+        # Randomly sample from sequences in this batch
+        batch_indices = torch.randperm(input_ids.size(0)).tolist()
+
+        for i in batch_indices:
+            if sample_count >= num_samples:
+                break
+
+            # Get actual sequence length (non-padded)
+            if attention_mask is not None:
+                seq_len = attention_mask[i].sum().item()
+            else:
+                seq_len = input_ids.size(1)
+
+            if seq_len < 10:
+                continue
+
+            # Determine truncation length
+            max_total = min(seq_len, max_length)
+            if labels is not None:
+                labels_i = labels[i][:seq_len]
+                answer_mask = labels_i != -100
+                if not answer_mask.any():
+                    # No answer tokens; skip for SFT masking
+                    continue
+                first_ans_idx = int(
+                    torch.nonzero(answer_mask, as_tuple=False)[0].item()
+                )
+                prompt_len = first_ans_idx
+                if prompt_len >= max_total:
+                    # Prompt alone reaches cap; cannot include any answer
+                    continue
+                remaining_answer = int(answer_mask[prompt_len:].sum().item())
+                allowed_answer = max_total - prompt_len
+                take_answer = min(remaining_answer, allowed_answer)
+                if take_answer <= 0:
+                    continue
+                actual_length = prompt_len + take_answer
+            else:
+                actual_length = max_total
+
+            # Extract the (possibly truncated) sequence
+            sequence = input_ids[i][:actual_length].unsqueeze(0).to(device)
+            attn_seq = (
+                attention_mask[i][:actual_length].unsqueeze(0).to(device)
+                if attention_mask is not None
+                else None
+            )
+            if labels is not None:
+                labels_seq = labels[i][:actual_length].unsqueeze(0).to(device)
+                sampled_sequences.append(
+                    {
+                        "input_ids": sequence,
+                        "labels": labels_seq,
+                        "attention_mask": attn_seq,
+                    }
+                )
+            else:
+                if attn_seq is not None:
+                    sampled_sequences.append(
+                        {"input_ids": sequence, "attention_mask": attn_seq}
+                    )
+                else:
+                    sampled_sequences.append(sequence)
+            sample_count += 1
+
+    return sampled_sequences
+
+
+def generate(
+    model: torch.nn.Module,
+    tokenizer: Any,
+    original_sequence: torch.Tensor,
+    num_diffusion_steps: int,
+    temperature: float,
+    mask_token_id: int,
+    *,
+    mode: Literal["random", "completion"] = "random",
+    completion_tokens: int = 0,
+    target_mask_ratio: Optional[float] = None,
+    labels: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> dict:
+    """Generate a single sample using reverse diffusion."""
+    # Get original text for comparison
+    original_text = tokenizer.decode(
+        original_sequence[0].cpu(), skip_special_tokens=True
+    )
+
+    # Build masked sequence
+    if (
+        labels is not None
+        and labels.numel() > 0
+        and (labels == -100).any()
+        and (labels != -100).any()
+    ):
+        # SFT case: completely mask all answer tokens (labels != -100)
+        total_tokens = original_sequence.size(1)
+        masked_indices = (labels != -100).to(dtype=torch.bool)
+        masked_sequence = original_sequence.clone()
+        masked_sequence[masked_indices] = mask_token_id
+        masked_tokens = int(masked_indices.sum().item())
+        mask_ratio = masked_tokens / max(int(total_tokens), 1)
+    elif mode == "completion" and completion_tokens > 0:
+        # Append mask tokens to the right for completion
+        total_tokens = original_sequence.size(1) + int(completion_tokens)
+        masked_indices = torch.zeros(
+            1, total_tokens, dtype=torch.bool, device=original_sequence.device
+        )
+        masked_indices[0, -int(completion_tokens) :] = True
+
+        append = torch.full(
+            (1, int(completion_tokens)), mask_token_id, device=original_sequence.device
+        )
+        masked_sequence = torch.cat([original_sequence, append], dim=1)
+        masked_tokens = int(completion_tokens)
+        mask_ratio = masked_tokens / total_tokens
+    else:
+        # Apply random masking with optional fixed ratio
+        total_tokens = original_sequence.size(1)
+        if target_mask_ratio is None:
+            min_ratio, max_ratio = 0.1, 0.7
+            target_mask_ratio = (
+                torch.rand(1).item() * (max_ratio - min_ratio) + min_ratio
+            )
+        target_masked_tokens = max(1, int(total_tokens * float(target_mask_ratio)))
+
+        # Create random mask indices
+        mask_positions = torch.randperm(total_tokens)[:target_masked_tokens]
+        masked_indices = torch.zeros(
+            1, total_tokens, dtype=torch.bool, device=original_sequence.device
+        )
+        masked_indices[0, mask_positions] = True
+
+        # Create masked sequence
+        masked_sequence = original_sequence.clone()
+        masked_sequence[masked_indices] = mask_token_id
+
+        # Calculate actual mask ratio
+        masked_tokens = masked_indices.sum().item()
+        mask_ratio = masked_tokens / total_tokens
+
+    # Get masked text for comparison
+    masked_text = tokenizer.decode(masked_sequence[0].cpu(), skip_special_tokens=False)
+    masked_text = _clean_masked_text(masked_text, tokenizer, mask_token_id)
+
+    # Run reverse diffusion process
+    sequence = masked_sequence.clone()
+    attention_mask = create_bidirectional_attention_mask(
+        sequence, attention_mask, sample_packing=attention_mask is not None
+    )
+    for step in range(num_diffusion_steps):
+        sequence = _diffusion_step(
+            model,
+            sequence,
+            step,
+            num_diffusion_steps,
+            temperature,
+            mask_token_id,
+            attention_mask,
+        )
+    generated_text = tokenizer.decode(sequence[0].cpu(), skip_special_tokens=True)
+
+    # Collect diagnostic info
+    final_ids = sequence[0].detach().cpu().tolist()
+    orig_ids_for_render = original_sequence[0].detach().cpu().tolist()
+    if masked_indices is not None:
+        masked_positions = (
+            torch.where(masked_indices[0])[0].detach().cpu().tolist()
+            if masked_indices.ndim == 2
+            else []
+        )
+    else:
+        masked_positions = []
+
+    result = {
+        "original": original_text,
+        "masked": masked_text,
+        "generated": generated_text,
+        "mask_ratio": mask_ratio,
+        "masked_tokens": masked_tokens,
+        "total_tokens": total_tokens,
+        "generated_ids": final_ids,
+        "masked_positions": masked_positions,
+        "orig_ids": orig_ids_for_render,
+        "formatted": (
+            f"Original: '{original_text}' → Masked: '{masked_text}' "
+            f"({mask_ratio:.1%}) → Generated: '{generated_text}'"
+        ),
+    }
+
+    return result
+
+
+def _clean_masked_text(masked_text: str, tokenizer: Any, mask_token_id: int) -> str:
+    """Clean up masked text for display."""
+    mask_token_repr = tokenizer.decode([mask_token_id], skip_special_tokens=False)
+    cleaned = masked_text.replace(mask_token_repr, "[MASK]")
+
+    # Remove literal special token strings
+    if hasattr(tokenizer, "special_tokens_map"):
+        for token_value in tokenizer.special_tokens_map.values():
+            if token_value and isinstance(token_value, str):
+                cleaned = cleaned.replace(token_value, "")
+
+    # Normalize whitespace but preserve newlines
+    cleaned = cleaned.replace("\r\n", "\n").replace("\r", "\n")
+    cleaned = re.sub(r"[ \t]+", " ", cleaned)
+    cleaned = "\n".join(line.rstrip() for line in cleaned.split("\n")).strip()
+    return cleaned
+
+
+def _diffusion_step(
+    model: torch.nn.Module,
+    sequence: torch.Tensor,
+    step: int,
+    num_diffusion_steps: int,
+    temperature: float,
+    mask_token_id: int,
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Perform a single diffusion step with remasking."""
+    # Only process if there are masked tokens remaining
+    current_mask = sequence == mask_token_id
+    if not current_mask.any():
+        return sequence
+
+    # Create or use provided attention mask
+    if attention_mask is None:
+        batch_size, seq_len = sequence.shape
+        attention_mask = torch.ones(
+            batch_size, 1, seq_len, seq_len, dtype=torch.bool, device=sequence.device
+        )
+
+    # Forward pass
+    outputs = model(input_ids=sequence, attention_mask=attention_mask)
+    logits = shift_logits_to_input_positions(outputs.logits)
+
+    # Only sample at currently masked positions
+    if current_mask.any():
+        masked_logits = logits[current_mask]
+
+        # Apply temperature scaling
+        if temperature > 0:
+            scaled_logits = masked_logits / temperature
+        else:
+            scaled_logits = masked_logits
+
+        # Suppress mask token in outputs
+        scaled_logits[:, mask_token_id] = -float("inf")
+
+        if temperature > 0:
+            # Add Gumbel noise for sampling
+            gumbel_noise = -torch.log(
+                -torch.log(torch.rand_like(scaled_logits, dtype=torch.float32))
+            )
+            gumbel_logits = scaled_logits + gumbel_noise
+            predicted_tokens = torch.argmax(gumbel_logits, dim=-1)
+        else:
+            predicted_tokens = torch.argmax(scaled_logits, dim=-1)
+
+        # Calculate probabilities for confidence scoring
+        probs = torch.softmax(scaled_logits, dim=-1)
+        predicted_token_probs = probs[range(len(predicted_tokens)), predicted_tokens]
+
+        # Determine how many tokens to unmask this step
+        remaining_masked = current_mask.sum().item()
+        if step == num_diffusion_steps - 1:
+            num_to_unmask = remaining_masked
+        else:
+            unmask_ratio = 1.0 / (num_diffusion_steps - step)
+            num_to_unmask = max(1, int(remaining_masked * unmask_ratio))
+
+        # Select highest confidence predictions to unmask
+        if num_to_unmask >= remaining_masked:
+            sequence[current_mask] = predicted_tokens
+        else:
+            _, top_indices = predicted_token_probs.topk(num_to_unmask)
+            mask_positions = torch.where(current_mask)[1]
+            positions_to_unmask = mask_positions[top_indices]
+            sequence[0, positions_to_unmask] = predicted_tokens[top_indices]
+
+    return sequence
diff --git a/src/axolotl/integrations/diffusion/plugin.py b/src/axolotl/integrations/diffusion/plugin.py
new file mode 100644
index 0000000000..cbfc2bde36
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/plugin.py
@@ -0,0 +1,43 @@
+"""Diffusion LM training plugin for Axolotl."""
+
+from peft import PeftModel
+from transformers import PreTrainedModel
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+from .trainer import DiffusionTrainer
+
+LOG = get_logger(__name__)
+
+
+class DiffusionPlugin(BasePlugin):
+    """
+    Plugin for diffusion language model training.
+
+    This plugin enables diffusion-based training using the LLaDA approach, which uses
+    random masking and bidirectional attention to train language models.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.cfg = None
+
+    def get_input_args(self) -> str:
+        """Returns the pydantic model for LLaDA plugin arguments."""
+        return "axolotl.integrations.diffusion.DiffusionArgs"
+
+    def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        """Perform actions after model is loaded."""
+        self.cfg = cfg
+
+    def get_trainer_cls(self, cfg: DictDefault) -> type[DiffusionTrainer] | None:
+        """Return custom trainer class for diffusion training."""
+        return DiffusionTrainer
+
+    def post_trainer_create(self, cfg: DictDefault, trainer: DiffusionTrainer):
+        """Configure trainer after creation."""
+        if hasattr(trainer, "axolotl_cfg"):
+            trainer.axolotl_cfg = cfg
+        trainer.post_set_axolotl_cfg()
diff --git a/src/axolotl/integrations/diffusion/trainer.py b/src/axolotl/integrations/diffusion/trainer.py
new file mode 100644
index 0000000000..7e35616f3d
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/trainer.py
@@ -0,0 +1,300 @@
+"""Custom trainer for diffusion LM training."""
+
+from typing import Any, Literal
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from axolotl.core.trainers.base import AxolotlTrainer
+from axolotl.utils.logging import get_logger
+
+from .callbacks import DiffusionGenerationCallback
+from .utils import create_bidirectional_attention_mask, shift_logits_to_input_positions
+
+LOG = get_logger(__name__)
+
+
+class DiffusionTrainer(AxolotlTrainer):
+    """Custom trainer for diffusion LM training that overrides loss computation."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._special_token_ids = None
+
+    def post_set_axolotl_cfg(self):
+        """Set config for diffusion training."""
+        self._cache_special_token_ids()
+        self._resolve_mask_token_id()
+
+        token_id = int(getattr(self.axolotl_cfg.diffusion, "mask_token_id", 0))
+        LOG.info(f"Diffusion: using mask_token_id={token_id}")
+
+        if getattr(self.axolotl_cfg.diffusion, "generate_samples", True):
+            generation_callback = DiffusionGenerationCallback(self)
+            self.add_callback(generation_callback)
+
+    def _resolve_mask_token_id(self) -> None:
+        """Ensure mask_token_id is valid for the current tokenizer."""
+        from .utils import resolve_mask_token_id
+
+        assert self.axolotl_cfg is not None, "axolotl_cfg is not set yet"
+
+        tokenizer = getattr(self, "processing_class", None)
+        if tokenizer is None:
+            return
+
+        mid = resolve_mask_token_id(
+            tokenizer,
+            self.axolotl_cfg,
+            allow_add=True,
+            model=getattr(self, "model", None),
+        )
+        try:
+            self.axolotl_cfg.diffusion.mask_token_id = int(mid)
+        except Exception:
+            pass
+
+    def compute_loss(
+        self,
+        model: nn.Module,
+        inputs: dict[str, torch.Tensor],
+        return_outputs: bool = False,
+        num_items_in_batch: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """Override compute_loss to use diffusion loss."""
+        input_ids = inputs.get("input_ids")
+        attention_mask = inputs.get("attention_mask")
+        labels = inputs.get("labels")
+
+        if input_ids is None:
+            raise ValueError("input_ids is required for diffusion training")
+
+        loss, outputs = self._compute_diffusion_loss(
+            model, input_ids, attention_mask, labels
+        )
+
+        if return_outputs:
+            return loss, outputs
+        return loss
+
+    def _cache_special_token_ids(self):
+        """Cache special token IDs to avoid repeated tokenizer access."""
+        if self.processing_class is None:
+            self._special_token_ids = set()
+            return
+
+        tokenizer = self.processing_class
+        special_tokens = set()
+
+        if hasattr(tokenizer, "bos_token_id") and tokenizer.bos_token_id is not None:
+            special_tokens.add(tokenizer.bos_token_id)
+        if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
+            special_tokens.add(tokenizer.eos_token_id)
+        if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is not None:
+            special_tokens.add(tokenizer.pad_token_id)
+
+        self._special_token_ids = special_tokens
+
+    def _forward_process(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        eps: float = 1e-3,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Forward noising process. A timestep is sampled along the process, and tokens are
+        masked with probability determined by the configured noise schedule.
+
+        Args:
+            input_ids: Input token ids [batch_size, seq_len].
+            attention_mask: Attention mask [batch_size, seq_len].
+            labels: Labels for SFT training [batch_size, seq_len].
+            eps: Small epsilon value for minimum masking probability.
+
+        Returns:
+            noisy_batch: Input with some tokens masked.
+            masked_indices: Boolean mask indicating which tokens were masked.
+            p_mask: Masking probabilities for each token [batch_size, seq_len].
+        """
+        batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+
+        # Sample random timesteps for each sample in batch
+        t = torch.rand(batch_size, device=device)
+        p_mask = (1 - eps) * t + eps  # [batch_size]
+        p_mask = p_mask[:, None].repeat(1, seq_len)  # [batch_size, seq_len]
+
+        # Don't mask padding tokens if attention_mask is provided
+        if attention_mask is not None:
+            valid_mask = attention_mask.bool()
+            p_mask = p_mask * valid_mask.float()
+
+        # Create mask to exclude special tokens
+        special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+        if self._special_token_ids:
+            for token_id in self._special_token_ids:
+                special_token_mask |= input_ids == token_id
+
+        # Create random mask based on p_mask
+        masked_indices = torch.rand((batch_size, seq_len), device=device) < p_mask
+        masked_indices = masked_indices & ~special_token_mask
+        if attention_mask is not None:
+            masked_indices = masked_indices & attention_mask.bool()
+
+        # For SFT data, only mask answer tokens
+        if labels is not None:
+            answer_mask = labels != -100
+            masked_indices = masked_indices & answer_mask
+
+        # Create masked input
+        mask_token_id = int(self.axolotl_cfg.diffusion.mask_token_id)
+        mask_value = torch.full_like(input_ids, mask_token_id)
+        noisy_batch = torch.where(masked_indices, mask_value, input_ids)
+
+        return noisy_batch, masked_indices, p_mask
+
+    def _compute_diffusion_loss(
+        self,
+        model: nn.Module,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | Any]:
+        """
+        Compute diffusion loss.
+
+        Args:
+            model: The model to compute loss for.
+            input_ids: Ground truth token ids [batch_size, seq_len].
+            attention_mask: Attention mask [batch_size, seq_len].
+            labels: Labels for SFT training [batch_size, seq_len].
+
+        Returns:
+            loss: Cross-entropy loss.
+            metrics: Dictionary of metrics.
+        """
+        # Short-circuit empty sequences
+        if input_ids is None or input_ids.numel() == 0 or input_ids.shape[1] == 0:
+            zero = torch.tensor(
+                0.0,
+                device=(input_ids.device if input_ids is not None else None),
+                requires_grad=True,
+            )
+            return zero, {}
+
+        # If an attention_mask is provided and all positions are padding for every
+        # sample in this batch, skip the step.
+        if attention_mask is not None:
+            if attention_mask.dim() == 2 and (attention_mask.sum(dim=1) == 0).all():
+                zero = torch.tensor(0.0, device=input_ids.device, requires_grad=True)
+                return zero, {}
+
+        # Apply forward process
+        noisy_batch, masked_indices, p_mask = self._forward_process(
+            input_ids, attention_mask, labels, self.axolotl_cfg.diffusion.eps
+        )
+
+        # Create bidirectional attention mask
+        bidirectional_mask = create_bidirectional_attention_mask(
+            input_ids, attention_mask, sample_packing=self.axolotl_cfg.sample_packing
+        )
+
+        # Forward pass
+        outputs = model(
+            input_ids=noisy_batch.long(),
+            attention_mask=bidirectional_mask,
+        )
+        logits = shift_logits_to_input_positions(outputs.logits)
+
+        if masked_indices.sum() > 0:
+            valid_indices = torch.where(masked_indices)
+            batch_indices, seq_indices = valid_indices
+
+            masked_logits = logits[batch_indices, seq_indices]
+            masked_targets = input_ids[batch_indices, seq_indices]
+            masked_p_mask = p_mask[batch_indices, seq_indices]
+
+            # Compute cross-entropy loss without reduction
+            token_loss = F.cross_entropy(
+                masked_logits.float(), masked_targets, reduction="none"
+            )
+
+            if self.axolotl_cfg.diffusion.importance_weighting:
+                masked_p_mask = masked_p_mask.float()
+                weighted_loss = token_loss / masked_p_mask
+            else:
+                weighted_loss = token_loss
+
+            if labels is not None:
+                # For SFT data: normalize by answer token count per sample
+                answer_mask = labels != -100
+                answer_lengths = answer_mask.sum(dim=1).float()  # [batch_size]
+
+                # Get batch indices for masked tokens
+                masked_batch_indices = batch_indices
+
+                # Sum losses per sample and divide by answer length
+                batch_size = input_ids.shape[0]
+                loss_per_sample = torch.zeros(batch_size, device=input_ids.device)
+                for i in range(batch_size):
+                    sample_mask = masked_batch_indices == i
+                    if sample_mask.sum() > 0:
+                        sample_loss = weighted_loss[sample_mask].sum()
+                        denom = answer_lengths[i].clamp(min=1.0)
+                        loss_per_sample[i] = sample_loss / denom
+
+                loss = loss_per_sample.mean()
+            else:
+                # Non-SFT: when importance weighting is enabled, use unbiased estimator
+                # (sum(loss/p) / total_tokens). Otherwise, average over masked tokens
+                # for stable scaling across varying mask ratios.
+                if self.axolotl_cfg.diffusion.importance_weighting:
+                    loss = weighted_loss.sum() / (
+                        input_ids.shape[0] * input_ids.shape[1]
+                    )
+                else:
+                    loss = weighted_loss.mean()
+
+            ce_loss = token_loss.mean()
+
+            # Compute accuracy on masked tokens
+            with torch.no_grad():
+                pred_tokens = masked_logits.argmax(dim=-1)
+                accuracy = (pred_tokens == masked_targets).float().mean()
+        else:
+            loss = torch.tensor(0.0, device=input_ids.device, requires_grad=True)
+            accuracy = torch.tensor(0.0, device=input_ids.device)
+            ce_loss = torch.tensor(0.0, device=input_ids.device)
+            masked_p_mask = torch.tensor(1.0, device=input_ids.device)
+
+        avg_p_mask = (
+            p_mask[masked_indices].mean().item() if masked_indices.any() else 0.0
+        )
+        metrics = {
+            "loss": loss.item(),
+            "accuracy": accuracy.item(),
+            "mask_ratio": masked_indices.float().mean().item(),
+            "num_masked_tokens": (masked_indices.sum().item(), "sum"),
+            "avg_p_mask": avg_p_mask,
+            "ce_loss": ce_loss.item(),
+        }
+
+        # If doing SFT training, log answer-specific metrics
+        if self.axolotl_cfg.datasets is not None:
+            with torch.no_grad():
+                answer_mask = labels != -100
+                answer_lengths = answer_mask.sum(dim=1).float()  # type: ignore
+                total_answer_tokens = answer_mask.sum().item()  # type: ignore
+                total_tokens = labels.numel()  # type: ignore
+                metrics["answer_ratio"] = total_answer_tokens / max(total_tokens, 1)
+                metrics["avg_answer_length"] = answer_lengths.mean().item()
+
+        if self.axolotl_cfg.diffusion.importance_weighting:
+            metrics["importance_weight_avg"] = (1.0 / masked_p_mask).mean().item()
+
+        train_eval: Literal["train", "eval"] = "train" if model.training else "eval"
+        self.store_metrics(metrics, train_eval=train_eval)
+
+        return loss, outputs
diff --git a/src/axolotl/integrations/diffusion/utils.py b/src/axolotl/integrations/diffusion/utils.py
new file mode 100644
index 0000000000..b6f71c07b9
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/utils.py
@@ -0,0 +1,166 @@
+"""Shared utilities for diffusion integration."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+
+from axolotl.utils.dict import DictDefault
+
+
+def resolve_mask_token_id(
+    tokenizer: Any,
+    cfg: DictDefault,
+    *,
+    allow_add: bool,
+    model: Any | None = None,
+    default_token: str = "<|diffusion_mask|>",
+) -> int:
+    """Resolve mask token id. Training may add a new special token; inference won't."""
+    # Determine vocab size if available
+    vocab_size = None
+    if tokenizer is not None:
+        if hasattr(tokenizer, "vocab_size") and tokenizer.vocab_size is not None:
+            try:
+                vocab_size = int(tokenizer.vocab_size)  # type: ignore[arg-type]
+            except Exception:
+                vocab_size = None
+        elif hasattr(tokenizer, "__len__"):
+            try:
+                vocab_size = int(len(tokenizer))
+            except Exception:
+                vocab_size = None
+
+    # Use explicit id from config if provided
+    diffusion_cfg = getattr(cfg, "diffusion", None)
+    # Fallback to top-level attr names only if nested missing (shouldn't happen)
+    cfg_id = (
+        getattr(diffusion_cfg, "mask_token_id", None)
+        if diffusion_cfg is not None
+        else getattr(cfg, "diffusion_mask_token_id", None)
+    )
+    if isinstance(cfg_id, int) and cfg_id >= 0:
+        if vocab_size is None or cfg_id < vocab_size:
+            return int(cfg_id)
+
+    def _existing_special_token_id(token_str: str | None) -> int | None:
+        """Attempt to resolve an existing special token string to a real ID."""
+        if not token_str or not hasattr(tokenizer, "convert_tokens_to_ids"):
+            return None
+        try:
+            token_id = tokenizer.convert_tokens_to_ids(token_str)
+        except Exception:
+            return None
+
+        if not isinstance(token_id, int) or token_id < 0:
+            return None
+
+        # Ensure it's registered as special and not UNK, and within vocab
+        unk_id = getattr(tokenizer, "unk_token_id", None)
+        specials = set(getattr(tokenizer, "all_special_tokens", []) or [])
+        addl = set(getattr(tokenizer, "additional_special_tokens", []) or [])
+        is_special = token_str in specials or token_str in addl
+        in_vocab = vocab_size is None or token_id < vocab_size
+        if (
+            (unk_id is not None and token_id == unk_id)
+            or not is_special
+            or not in_vocab
+        ):
+            return None
+        return token_id
+
+    # Try mask token string if provided
+    token_str = (
+        getattr(diffusion_cfg, "mask_token_str", None)
+        if diffusion_cfg is not None
+        else getattr(cfg, "diffusion_mask_token_str", None)
+    )
+    for candidate in (token_str, default_token):
+        token_id = _existing_special_token_id(candidate)
+        if isinstance(token_id, int):
+            try:
+                if diffusion_cfg is None:
+                    cfg.diffusion_mask_token_id = int(token_id)  # legacy fallback
+                else:
+                    diffusion_cfg.mask_token_id = int(token_id)
+            except Exception:
+                pass
+            return int(token_id)
+
+    # Optionally add and return a dedicated special token during training
+    if allow_add and hasattr(tokenizer, "add_special_tokens"):
+        token_to_add = token_str or default_token
+        try:
+            tokenizer.add_special_tokens({"additional_special_tokens": [token_to_add]})
+
+            # Resize embeddings if possible
+            if (
+                model is not None
+                and hasattr(tokenizer, "__len__")
+                and hasattr(model, "resize_token_embeddings")
+            ):
+                try:
+                    model.resize_token_embeddings(len(tokenizer))
+                except Exception:
+                    pass
+            new_id = tokenizer.convert_tokens_to_ids(token_to_add)
+            if isinstance(new_id, int) and new_id >= 0:
+                try:
+                    if diffusion_cfg is None:
+                        cfg.diffusion_mask_token_id = int(new_id)  # legacy fallback
+                    else:
+                        diffusion_cfg.mask_token_id = int(new_id)
+                except Exception:
+                    pass
+                return int(new_id)
+        except Exception:
+            pass
+
+    # Fallback to unk or 0 (do not update cfg)
+    fallback = getattr(tokenizer, "unk_token_id", 0) or 0
+    return int(fallback)
+
+
+def create_bidirectional_attention_mask(
+    input_ids: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    sample_packing: bool = False,
+) -> torch.Tensor:
+    """
+    Create bidirectional attention mask to override default causal masking.
+    Handles sample-packed sequences where different samples are identified
+    by different attention mask values.
+
+    Args:
+        input_ids: Input token ids [batch_size, seq_len]
+        attention_mask: Attention mask [batch_size, seq_len]
+        sample_packing: Whether sample packing is enabled
+
+    Returns:
+        bidirectional_mask: 4D attention mask [batch_size, 1, seq_len, seq_len]
+    """
+    batch_size, seq_len = input_ids.shape
+    device = input_ids.device
+
+    if attention_mask is None or not sample_packing:
+        return torch.ones(
+            batch_size, 1, seq_len, seq_len, dtype=torch.bool, device=device
+        )
+
+    # Handle sample packing: tokens can only attend within their sample
+    mask_i = attention_mask.unsqueeze(2)  # [batch_size, seq_len, 1]
+    mask_j = attention_mask.unsqueeze(1)  # [batch_size, 1, seq_len]
+
+    # Tokens can attend to each other if they have the same non-zero sample ID
+    bidirectional_mask = (mask_i == mask_j) & (mask_i > 0)
+
+    # Add head dimension: [batch_size, 1, seq_len, seq_len]
+    return bidirectional_mask.unsqueeze(1)
+
+
+def shift_logits_to_input_positions(logits: torch.Tensor) -> torch.Tensor:
+    """Align next-token logits with their input token positions for diffusion."""
+    if logits.size(1) <= 1:
+        return logits
+    return torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
diff --git a/src/axolotl/integrations/expert_parallel/README.md b/src/axolotl/integrations/expert_parallel/README.md
new file mode 100644
index 0000000000..bf726f7c49
--- /dev/null
+++ b/src/axolotl/integrations/expert_parallel/README.md
@@ -0,0 +1,208 @@
+# Expert Parallelism Integration
+
+Replaces the MoE dispatch/combine path with DeepEP's fused kernels.
+
+## Requirements
+
+Ampere (sm_80, A100) or Hopper (sm_90, H100), all-pairs NVLink.
+
+## Installation
+
+**Hopper (sm_90, H100), multi-node with NCCL 2.29+ (torch 2.11+) and OFED:**
+
+```bash
+git clone --depth 1 https://github.com/deepseek-ai/DeepEP.git
+cd DeepEP
+TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=16 uv pip install --no-build-isolation .
+python -c "import deep_ep; print(deep_ep.Buffer)"
+```
+
+**Hopper (sm_90, H100), single-node intranode-only (no OFED):**
+
+```bash
+git clone https://github.com/deepseek-ai/DeepEP.git
+cd DeepEP
+git checkout v1.2.1
+# Patch 1: setup.py — honor DISABLE_NVSHMEM=1
+git apply <<'EOF'
+--- a/setup.py
++++ b/setup.py
+@@ -19,7 +19,10 @@ if __name__ == '__main__':
+     disable_nvshmem = False
+     nvshmem_dir = os.getenv('NVSHMEM_DIR', None)
+     nvshmem_host_lib = 'libnvshmem_host.so'
+-    if nvshmem_dir is None:
++    if int(os.getenv('DISABLE_NVSHMEM', '0')):
++        disable_nvshmem = True
++        nvshmem_dir = None
++    elif nvshmem_dir is None:
+         try:
+             nvshmem_dir = importlib.util.find_spec("nvidia.nvshmem").submodule_search_locations[0]
+             nvshmem_host_lib = get_nvshmem_host_lib_name(nvshmem_dir)
+EOF
+
+# Patch 2: setup.py — drop -rdc=true when NVSHMEM is disabled
+git apply <<'EOF'
+--- a/setup.py
++++ b/setup.py
+@@ -71,7 +71,9 @@ if __name__ == '__main__':
+         os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '9.0')
+
+         # CUDA 12 flags
+-        nvcc_flags.extend(['-rdc=true', '--ptxas-options=--register-usage-level=10'])
++        nvcc_flags.append('--ptxas-options=--register-usage-level=10')
++        if not disable_nvshmem:
++            nvcc_flags.append('-rdc=true')
+
+     # Disable LD/ST tricks, as some CUDA version does not support `.L1::no_allocate`
+     if os.environ['TORCH_CUDA_ARCH_LIST'].strip() != '9.0':
+EOF
+
+DISABLE_NVSHMEM=1 TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=16 \
+  uv pip install --no-build-isolation .
+python -c "import deep_ep; print(deep_ep.Buffer)"
+```
+
+Notes:
+
+- Hopper kernels (FP8, TMA, etc.) are preserved; only intranode dispatch/combine is built — appropriate for single-node H100×{4,8}.
+- Patch 1 lets `DISABLE_NVSHMEM=1` skip the NVSHMEM build path, which would otherwise need Mellanox OFED dev headers (`infiniband/mlx5dv.h`).
+- Patch 2 drops `-rdc=true` when NVSHMEM is off; otherwise the device-link step has nothing to link against and import fails with `__cudaRegisterLinkedBinary_*` undefined symbol.
+- The `v1.2.1` pin is required: it is the last release whose `setup.py` still carries the `disable_nvshmem` path the two patches edit. DeepEP `main` restructured `setup.py` and removed that path ([DeepEP #664](https://github.com/deepseek-ai/DeepEP/pull/664)), so against HEAD the patches fail with `patch failed: setup.py:19`. The pin also sidesteps DeepEP HEAD's `csrc/elastic/` (Engram/EPv2, commit `b306af0`), which needs `ncclGinRequest_t` from NCCL 2.29+.
+
+**Ampere (sm_80, A100, intranode-only)** — needs two small source patches gated on `DISABLE_NVSHMEM=1`:
+
+```bash
+git clone https://github.com/deepseek-ai/DeepEP.git
+cd DeepEP
+git checkout v1.2.1
+
+# Patch 1: setup.py — honor DISABLE_NVSHMEM=1
+git apply <<'EOF'
+--- a/setup.py
++++ b/setup.py
+@@ -19,7 +19,10 @@ if __name__ == '__main__':
+     disable_nvshmem = False
+     nvshmem_dir = os.getenv('NVSHMEM_DIR', None)
+     nvshmem_host_lib = 'libnvshmem_host.so'
+-    if nvshmem_dir is None:
++    if int(os.getenv('DISABLE_NVSHMEM', '0')):
++        disable_nvshmem = True
++        nvshmem_dir = None
++    elif nvshmem_dir is None:
+         try:
+             nvshmem_dir = importlib.util.find_spec("nvidia.nvshmem").submodule_search_locations[0]
+             nvshmem_host_lib = get_nvshmem_host_lib_name(nvshmem_dir)
+EOF
+
+# Patch 2: csrc/deep_ep.cpp — gate the three mask_buffer methods
+git apply <<'EOF'
+--- a/csrc/deep_ep.cpp
++++ b/csrc/deep_ep.cpp
+@@ -1823,22 +1823,34 @@ bool is_sm90_compiled() {
+ }
+
+ void Buffer::low_latency_update_mask_buffer(int rank_to_mask, bool mask) {
++#ifndef DISABLE_NVSHMEM
+     EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
+     EP_HOST_ASSERT(rank_to_mask >= 0 and rank_to_mask < num_ranks);
+     internode_ll::update_mask_buffer(mask_buffer_ptr, rank_to_mask, mask, at::cuda::getCurrentCUDAStream());
++#else
++    EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation");
++#endif
+ }
+
+ void Buffer::low_latency_query_mask_buffer(const torch::Tensor& mask_status) {
++#ifndef DISABLE_NVSHMEM
+     EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
+     EP_HOST_ASSERT(mask_status.numel() == num_ranks && mask_status.scalar_type() == torch::kInt32);
+
+     internode_ll::query_mask_buffer(
+         mask_buffer_ptr, num_ranks, reinterpret_cast<int*>(mask_status.data_ptr()), at::cuda::getCurrentCUDAStream());
++#else
++    EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation");
++#endif
+ }
+
+ void Buffer::low_latency_clean_mask_buffer() {
++#ifndef DISABLE_NVSHMEM
+     EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
+     internode_ll::clean_mask_buffer(mask_buffer_ptr, num_ranks, at::cuda::getCurrentCUDAStream());
++#else
++    EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation");
++#endif
+ }
+
+ }  // namespace deep_ep
+EOF
+
+DISABLE_NVSHMEM=1 DISABLE_SM90_FEATURES=1 TORCH_CUDA_ARCH_LIST=8.0 MAX_JOBS=16 \
+  uv pip install --no-build-isolation .
+python -c "import deep_ep; print(deep_ep.Buffer)"
+```
+
+## Usage
+
+```yaml
+plugins:
+  - axolotl.integrations.expert_parallel.ExpertParallelPlugin
+
+expert_parallel_size: 2  # 1 = disabled (default); > 1 = enabled
+```
+
+For composition with FSDP at 4+ GPUs, set both `expert_parallel_size` and `dp_shard_size`. The product must equal `world_size`:
+
+```yaml
+# 4-GPU example: ep × dp_shard = 2 × 2 = 4
+expert_parallel_size: 2
+dp_shard_size: 2
+fsdp_config:
+  fsdp_version: 2
+  auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  transformer_layer_cls_to_wrap: Qwen3MoeDecoderLayer
+  state_dict_type: FULL_STATE_DICT
+  sharding_strategy: FULL_SHARD
+```
+
+See full example configs at [`examples/expert_parallel/`](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/expert_parallel).
+
+#### Implementation notes
+
+EP composes with the local-experts kernel you've already configured: ScatterMoE, ~SonicMoE~ (WIP), grouped_mm, or eager.
+
+EP composes with FSDP on orthogonal mesh axes: experts are sharded across the `ep` axis, non-expert params across `dp_shard`. The two collectives run on disjoint process groups, so they don't conflict. Layout follows [*Expert Parallelism with FSDP* (tinkerings.dev)](https://tinkerings.dev/posts/expert_parallel.html) — "rows share weights, columns move tokens."
+
+| Your existing config                                | Local kernel under DeepEP |
+|-----------------------------------------------------|---------------------------|
+| `use_scattermoe: true`                              | ScatterMoE (Triton)       |
+| `use_sonicmoe: true` (WIP)                          | SonicMoE (Gemma4)         |
+| `experts_implementation: grouped_mm` / `batched_mm` | grouped_mm (transformers) |
+| `experts_implementation: eager`                     | eager Python loop         |
+| (unset)                                             | grouped_mm (default)      |
+
+## Limitations
+
+- Models' modeling code must use `@use_experts_implementation` (canonical 3D `gate_up_proj` / `down_proj`). `ModuleList` as used in Mixtral is not supported.
+- `num_experts` must be divisible by `expert_parallel_size`.
+- Supported mesh axes: EP, EP × dp_shard, **EP × cp**, EP × cp × dp_shard (experts shard on `ep`,
+  the sequence on `cp`, non-expert weights on `dp_shard`). EP × **TP** is not yet supported and
+  raises `NotImplementedError`. EP × CP requires the model's attention to be context-parallel-aware
+  on the `cp` axis (e.g. GLM-5.2 DSA via the kernels plugin); stock attention uses accelerate CP.
+- DeepEP limitation: Low-latency (LL) kernels are inter-node only by design (pure RDMA via IBGDA). Single-node + intranode setups always use the standard kernels and don't benefit from LL.
+- FP8 dispatch needs Hopper + DISABLE_SM90_FEATURES=0.
+
+## Troubleshooting
+
+### `CUBLAS_STATUS_INVALID_VALUE` on a basic bf16 GEMM after `import deep_ep`
+
+The system's `libcublas.so.13` is older than what cu130 torch expects. Put the cu13 lib that ships with the torch wheel on `LD_LIBRARY_PATH`:
+
+```bash
+export LD_LIBRARY_PATH="$(python -c 'import nvidia.cu13 as m; print(list(m.__path__)[0] + "/lib")'):$LD_LIBRARY_PATH"
+```
+
+Unrelated to DeepEP itself, but anyone on cu130 torch hits it on boxes with a system CUDA toolkit older than 13.0.
+
+### `CUDA error 803` (system not yet initialized)
+
+On driver `< 580`, also prepend `/usr/local/cuda-13.0/compat` to `LD_LIBRARY_PATH`. **Do not** add the compat dir on driver `≥ 580` (its `libcuda` is older than the running driver and triggers `CUDA error 803`).
diff --git a/src/axolotl/integrations/expert_parallel/__init__.py b/src/axolotl/integrations/expert_parallel/__init__.py
new file mode 100644
index 0000000000..188c42d794
--- /dev/null
+++ b/src/axolotl/integrations/expert_parallel/__init__.py
@@ -0,0 +1,26 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+"""Expert-Parallel (DeepEP) integration for axolotl.
+
+Replaces the dispatch/combine path in transformers MoE blocks with DeepEP's
+fused kernels. Registers four names in `transformers.integrations.moe.ALL_EXPERTS_FUNCTIONS`:
+
+- `deep_ep`               — eager local expert MLP (reference)
+- `deep_ep_grouped_mm`    — transformers' grouped_mm kernel (default)
+- `deep_ep_scattermoe`    — axolotl's ScatterMoE kernel
+- `deep_ep_sonicmoe`      — axolotl's SonicMoE kernel
+"""
+
+from .args import ExpertParallelArgs
+from .plugin import ExpertParallelPlugin
+
+__all__ = [
+    "ExpertParallelArgs",
+    "ExpertParallelPlugin",
+]
diff --git a/src/axolotl/integrations/expert_parallel/args.py b/src/axolotl/integrations/expert_parallel/args.py
new file mode 100644
index 0000000000..79fa58c3c2
--- /dev/null
+++ b/src/axolotl/integrations/expert_parallel/args.py
@@ -0,0 +1,62 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+"""Pydantic args for the Expert-Parallel (DeepEP) plugin."""
+
+from typing import Literal
+
+from pydantic import BaseModel, model_validator
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class ExpertParallelArgs(BaseModel):
+    """Input args for the expert_parallel plugin. See the integration README."""
+
+    expert_parallel_size: int = 1
+    """Number of EP ranks. 1 = disabled (default), > 1 = enabled."""
+
+    expert_parallel_backend: Literal["deep_ep"] = "deep_ep"
+
+    expert_parallel_num_nvl_bytes: int = 256 << 20
+
+    expert_parallel_num_rdma_bytes: int = 0
+
+    expert_parallel_token_capacity: int | None = None
+    """Max tokens routed to any single expert per forward; the lowest-weight excess (token,expert)
+    assignments are dropped. DeepEP's intranode combine deadlocks once one expert is overloaded, and
+    GLM-style routers concentrate more with depth — set this (e.g. 1024) for them. ``None`` = no cap."""
+
+    expert_parallel_fallback_on_unsupported: bool = True
+
+    @model_validator(mode="after")
+    def _validate(self):
+        if self.expert_parallel_size < 1:
+            raise ValueError(
+                f"expert_parallel_size must be >= 1 (got {self.expert_parallel_size!r}). "
+                f"Use 1 to disable EP."
+            )
+
+        if (
+            self.expert_parallel_token_capacity is not None
+            and self.expert_parallel_token_capacity < 1
+        ):
+            raise ValueError(
+                f"expert_parallel_token_capacity must be >= 1 when set "
+                f"(got {self.expert_parallel_token_capacity!r}). Use None to disable the cap."
+            )
+
+        if self.expert_parallel_size > 1 and self.expert_parallel_num_rdma_bytes != 0:
+            LOG.warning(
+                "expert_parallel_num_rdma_bytes != 0 — RDMA path requires "
+                "Hopper + IBGDA-capable InfiniBand. Will fail on Ampere/intranode."
+            )
+
+        return self
diff --git a/src/axolotl/integrations/expert_parallel/buffer.py b/src/axolotl/integrations/expert_parallel/buffer.py
new file mode 100644
index 0000000000..1ef4508177
--- /dev/null
+++ b/src/axolotl/integrations/expert_parallel/buffer.py
@@ -0,0 +1,110 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+"""DeepEP `Buffer` singleton, lazily constructed on first call.
+
+A single Buffer is reused across all MoE layers in a model, since DeepEP's
+intranode kernels are sized by `num_nvl_bytes` which we set conservatively at
+plugin init. Per-layer Buffer construction would burn memory.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch.distributed as dist
+
+_BUFFER = None
+_EP_GROUP: Optional[dist.ProcessGroup] = None
+_NUM_NVL_BYTES = 256 << 20
+_DISPATCH_CONFIG = None
+_COMBINE_CONFIG = None
+_NUM_RDMA_BYTES = 0
+
+
+def configure_buffer(
+    ep_group: Optional[dist.ProcessGroup],
+    num_nvl_bytes: int = 256 << 20,
+    num_rdma_bytes: int = 0,
+) -> None:
+    """Stash params for lazy Buffer construction. Call from `post_model_build`."""
+    global _EP_GROUP, _NUM_NVL_BYTES, _NUM_RDMA_BYTES, _BUFFER
+    _EP_GROUP = ep_group
+    _NUM_NVL_BYTES = num_nvl_bytes
+    _NUM_RDMA_BYTES = num_rdma_bytes
+    _BUFFER = None  # invalidate any prior buffer
+
+
+def get_buffer():
+    """Return the (lazily constructed) DeepEP Buffer."""
+    global _BUFFER
+    if _BUFFER is not None:
+        return _BUFFER
+
+    import deep_ep
+
+    group = _EP_GROUP if _EP_GROUP is not None else dist.group.WORLD
+    _BUFFER = deep_ep.Buffer(
+        group=group,
+        num_nvl_bytes=_NUM_NVL_BYTES,
+        num_rdma_bytes=_NUM_RDMA_BYTES,
+        low_latency_mode=False,
+    )
+    return _BUFFER
+
+
+def _ep_world() -> int:
+    grp = _EP_GROUP if _EP_GROUP is not None else dist.group.WORLD
+    return dist.get_world_size(grp)
+
+
+def get_dispatch_config():
+    """Recommended DeepEP dispatch Config for the EP group. DeepEP's own tests ALWAYS pass an
+    explicit config to dispatch/combine; the no-config default deadlocks / launch-fails the combine
+    on the singleton buffer's reuse across MoE layers (cudaErrorLaunchFailure)."""
+    global _DISPATCH_CONFIG
+    if _DISPATCH_CONFIG is None:
+        import deep_ep
+
+        _DISPATCH_CONFIG = deep_ep.Buffer.get_dispatch_config(_ep_world())
+    return _DISPATCH_CONFIG
+
+
+def get_combine_config():
+    """Recommended DeepEP combine Config for the EP group (see get_dispatch_config)."""
+    global _COMBINE_CONFIG
+    if _COMBINE_CONFIG is None:
+        import deep_ep
+
+        _COMBINE_CONFIG = deep_ep.Buffer.get_combine_config(_ep_world())
+    return _COMBINE_CONFIG
+
+
+def barrier_ep() -> None:
+    """Barrier on the EP group. Placed before each DeepEP ``combine`` so the fast ranks wait (on
+    NCCL, no short timeout) for any rank still AUTOTUNING the local expert kernel — otherwise the
+    combine collective hits DeepEP's short internal timeout (``value=0``) and aborts. Negligible cost
+    once kernels are cached (all ranks arrive together). Disable with AXOLOTL_EP_NO_BARRIER=1."""
+    import os
+
+    if os.environ.get("AXOLOTL_EP_NO_BARRIER"):
+        return
+    if not dist.is_initialized():
+        return
+    import torch
+
+    torch.cuda.synchronize()
+    dist.barrier(_EP_GROUP if _EP_GROUP is not None else dist.group.WORLD)
+
+
+def reset_buffer() -> None:
+    """Drop the cached Buffer + configs. Used in tests."""
+    global _BUFFER, _DISPATCH_CONFIG, _COMBINE_CONFIG
+    _BUFFER = None
+    _DISPATCH_CONFIG = None
+    _COMBINE_CONFIG = None
diff --git a/src/axolotl/integrations/expert_parallel/experts_fn.py b/src/axolotl/integrations/expert_parallel/experts_fn.py
new file mode 100644
index 0000000000..5d3046d1f7
--- /dev/null
+++ b/src/axolotl/integrations/expert_parallel/experts_fn.py
@@ -0,0 +1,400 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+"""DeepEP-backed registered functions for `ALL_EXPERTS_FUNCTIONS`.
+
+Four names registered (eager / grouped_mm / scattermoe / sonicmoe) sharing one
+`_deep_ep_forward` body. Templates ported from `bench_deep_ep.py` Stage 1
+modes 3 and 4.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from .buffer import barrier_ep, get_buffer, get_combine_config, get_dispatch_config
+
+# Per-forward [B*S] bool mask of real (non-padding) tokens, set by the EP plugin's model
+# pre-hook from the batch `attention_mask` and consumed by `_deep_ep_forward` to exclude
+# padding from the dispatch. ``None`` when there is no padding (packed / no attention_mask).
+_VALID_TOKEN_MASK: torch.Tensor | None = None
+
+# Max tokens routed to any single expert per forward (``expert_parallel_token_capacity``), set by the
+# EP plugin. ``None`` disables the cap. Overloaded experts deadlock DeepEP's intranode combine, so
+# GLM-style routers (concentration grows with depth) need this.
+_TOKEN_CAPACITY: int | None = None
+
+
+def set_token_capacity(cap: int | None) -> None:
+    global _TOKEN_CAPACITY
+    _TOKEN_CAPACITY = cap
+
+
+def _apply_expert_capacity(topk_idx, topk_w, cap):
+    """Cap tokens-per-expert to ``cap`` by sentinelling (-1) the lowest-weight excess (token,expert)
+    assignments, then rescale each token's surviving weights back to its pre-drop gate sum. DeepEP's
+    intranode combine hangs once one expert receives too many tokens (the GLM router concentrates more
+    with depth, crossing the hang threshold ~layer 31); standard MoE capacity-dropping keeps every
+    expert under the limit. The GLM router normalizes the top-k weights to sum to 1, so dropping an
+    expert without rescaling would silently attenuate that token's combined expert output. Already-(-1)
+    slots are left untouched. Returns ``(capped_idx, rescaled_w)``."""
+    import torch
+
+    ntok, K = topk_idx.shape
+    flat_e = topk_idx.reshape(-1)
+    flat_w = topk_w.reshape(-1)
+    # sort assignments by (expert asc, weight desc) via two stable passes
+    o1 = torch.argsort(flat_w, descending=True, stable=True)
+    o2 = torch.argsort(flat_e[o1], stable=True)
+    order = o1[o2]
+    se = flat_e[order]
+    pos = torch.arange(se.numel(), device=se.device)
+    is_new = torch.ones_like(se, dtype=torch.bool)
+    is_new[1:] = se[1:] != se[:-1]
+    grp_start = torch.cummax(torch.where(is_new, pos, torch.zeros_like(pos)), 0).values
+    within = pos - grp_start
+    drop_sorted = (within >= cap) & (se >= 0)
+    drop = torch.zeros_like(flat_e, dtype=torch.bool)
+    drop[order] = drop_sorted
+    capped_idx = flat_e.masked_fill(drop, -1).reshape(ntok, K)
+
+    orig_sum = (topk_w * (topk_idx >= 0)).sum(dim=-1, keepdim=True)
+    kept = (capped_idx >= 0).to(topk_w.dtype)
+    kept_sum = (topk_w * kept).sum(dim=-1, keepdim=True)
+    # double-where: guard the divisor so a fully-dropped token (kept_sum==0) can't backprop 0*inf=NaN
+    safe_kept_sum = torch.where(kept_sum > 0, kept_sum, torch.ones_like(kept_sum))
+    rescale = torch.where(
+        kept_sum > 0, orig_sum / safe_kept_sum, torch.ones_like(kept_sum)
+    )
+    rescaled_w = topk_w * kept * rescale
+    return capped_idx, rescaled_w
+
+
+def set_valid_token_mask(mask: torch.Tensor | None) -> None:
+    global _VALID_TOKEN_MASK
+    _VALID_TOKEN_MASK = mask
+
+
+def _get_valid_token_mask() -> torch.Tensor | None:
+    return _VALID_TOKEN_MASK
+
+
+def _eager_local(experts, recv_x, recv_topk_idx, recv_topk_weights):
+    """Eager Python loop over local experts. Reference for numerics."""
+    out = torch.zeros_like(recv_x)
+    num_local = getattr(experts, "num_local_experts", experts.num_experts)
+    for e in range(num_local):
+        rows, ks = (recv_topk_idx == e).nonzero(as_tuple=True)
+        if rows.numel() == 0:
+            continue
+        x_e = recv_x[rows]
+        gate_up = F.linear(x_e, experts.gate_up_proj[e])
+        gate, up = gate_up.chunk(2, dim=-1)
+        h = experts.act_fn(gate) * up
+        y = F.linear(h, experts.down_proj[e])
+        weighted = (y * recv_topk_weights[rows, ks].unsqueeze(-1)).to(out.dtype)
+        out.index_add_(0, rows, weighted)
+    return out
+
+
+def _maybe_install_decorator_attrs(experts):
+    """`@use_experts_implementation` injects has_gate/has_bias/is_transposed.
+    For models loaded outside the decorator path (or in tests), patch them in.
+    """
+    if not hasattr(experts, "has_gate"):
+        experts.has_gate = True
+    if not hasattr(experts, "has_bias"):
+        experts.has_bias = hasattr(experts, "gate_up_proj_bias")
+    if not hasattr(experts, "is_transposed"):
+        experts.is_transposed = False
+
+
+def _mask_sentinels(recv_topk_idx, recv_topk_weights):
+    """Map ``-1`` remote sentinels to expert 0 / weight 0 for kernels that index by
+    expert id and don't filter (grouped_mm). The zero weight nulls their contribution."""
+    safe_idx = torch.where(
+        recv_topk_idx >= 0, recv_topk_idx, torch.zeros_like(recv_topk_idx)
+    )
+    valid = (recv_topk_idx >= 0).to(recv_topk_weights.dtype)
+    return safe_idx, recv_topk_weights * valid
+
+
+def _grouped_mm_local(experts, recv_x, recv_topk_idx, recv_topk_weights):
+    from transformers.integrations.moe import grouped_mm_experts_forward
+
+    _maybe_install_decorator_attrs(experts)
+    safe_idx, safe_w = _mask_sentinels(recv_topk_idx, recv_topk_weights)
+    return grouped_mm_experts_forward(experts, recv_x, safe_idx, safe_w)
+
+
+def _scattermoe_local(experts, recv_x, recv_topk_idx, recv_topk_weights):
+    # scattermoe skips sentinel rows natively (only valid rows hit the grouped GEMM
+    # + per-row LoRA) -- pass the raw -1-tagged routing, not the masked version.
+    from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+        scattermoe_experts_forward_ep,
+    )
+
+    return scattermoe_experts_forward_ep(
+        experts, recv_x, recv_topk_idx, recv_topk_weights
+    )
+
+
+def _sonicmoe_local(experts, recv_x, recv_topk_idx, recv_topk_weights):
+    # The sonic-moe CUTLASS kernel can't run on sm_120 (Blackwell); for standard-layout
+    # experts there, use the vendored scattermoe EP path (sentinel-skip + fused MXFP4),
+    # which runs on sm_120 and gives sonicmoe + LoRA + EP. Elsewhere sonicmoe+EP is not
+    # yet wired (needs the upstream EP-sentinel kernel).
+    from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+        scattermoe_experts_forward_ep,
+        scattermoe_supports_layout,
+    )
+    from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+        _sonicmoe_kernel_supported,
+    )
+
+    if not _sonicmoe_kernel_supported() and scattermoe_supports_layout(experts):
+        return scattermoe_experts_forward_ep(
+            experts, recv_x, recv_topk_idx, recv_topk_weights
+        )
+    raise NotImplementedError(
+        "Sonicmoe + EP is not yet implemented on this device/layout. On sm_120 with a "
+        "standard expert layout it falls back to the scattermoe EP path automatically; "
+        "otherwise use use_scattermoe."
+    )
+
+
+_LOCAL_KERNELS = {
+    "eager": _eager_local,
+    "grouped_mm": _grouped_mm_local,
+    "scattermoe": _scattermoe_local,
+    "sonicmoe": _sonicmoe_local,
+}
+
+
+class _DeepEPDispatch(torch.autograd.Function):
+    """Autograd wrapper for `Buffer.dispatch`.
+
+    Forward: routes `x` to ranks owning its top-k experts.
+    Backward of dispatch is combine — gradients on `recv_x` are reduced back
+    across ranks using the same handle. Only `x`'s grad is returned (idx/weights
+    are non-differentiable bookkeeping).
+    """
+
+    @staticmethod
+    def forward(
+        ctx, x, topk_idx, topk_weights, num_per_rank, num_per_expert, is_in_rank
+    ):
+        buffer = get_buffer()
+        recv_x, recv_topk_idx, recv_topk_weights, _, handle, _ = buffer.dispatch(
+            x,
+            num_tokens_per_rank=num_per_rank,
+            is_token_in_rank=is_in_rank,
+            num_tokens_per_expert=num_per_expert,
+            topk_idx=topk_idx,
+            topk_weights=topk_weights,
+            config=get_dispatch_config(),
+        )
+        ctx.handle = handle
+        return recv_x, recv_topk_idx, recv_topk_weights, _DeepEPHandleHolder(handle)
+
+    @staticmethod
+    def backward(ctx, grad_recv_x, _grad_idx, grad_recv_w, _grad_handle):
+        buffer = get_buffer()
+        # Backward-of-dispatch is combine: reduce grad on recv_x (and recv_topk_weights)
+        # back to the source rank via the same handle.
+        topk_w_grad = (
+            grad_recv_w.contiguous()
+            if (grad_recv_w is not None and grad_recv_w.numel() > 0)
+            else None
+        )
+        barrier_ep()
+        grad_x, grad_topk_w, _ = buffer.combine(
+            grad_recv_x.contiguous(),
+            ctx.handle,
+            topk_weights=topk_w_grad,
+            config=get_combine_config(),
+        )
+        return grad_x, None, grad_topk_w, None, None, None
+
+
+class _DeepEPCombine(torch.autograd.Function):
+    """Autograd wrapper for `Buffer.combine`.
+
+    Forward: sums per-rank partial outputs back into the source token.
+    Backward of combine is dispatch — reuses the handle from the dispatch call.
+    """
+
+    @staticmethod
+    def forward(ctx, x, handle_holder):
+        buffer = get_buffer()
+        ctx.handle = handle_holder.handle
+        combined, _, _ = buffer.combine(
+            x.contiguous(), handle_holder.handle, config=get_combine_config()
+        )
+        return combined
+
+    @staticmethod
+    def backward(ctx, grad_combined):
+        buffer = get_buffer()
+        # backward-of-combine is dispatch: reuse the cached handle to send
+        # gradients to the ranks that produced partial outputs.
+        recv_grad, _, _, _, _, _ = buffer.dispatch(
+            grad_combined.contiguous(), handle=ctx.handle, config=get_dispatch_config()
+        )
+        return recv_grad, None
+
+
+class _DeepEPHandleHolder:
+    """Carries the dispatch handle through to combine without breaking autograd
+    (autograd.Function output must be Tensors or non-Tensor objects; tuples
+    of opaque handles work as a passthrough since combine doesn't differentiate
+    against them)."""
+
+    def __init__(self, handle):
+        self.handle = handle
+
+
+def _deep_ep_forward(self, hidden_states, top_k_index, top_k_weights, *, kernel_name):
+    """Shared dispatch -> local-experts -> combine pipeline.
+
+    Inputs come in with **global** routing indices (we do not run
+    `transformers.RouterParallel`; see DEEP_EP.md §2.4 for why). Dispatch returns
+    local expert ids in `[0, E_local)` with `-1` for slots routed to remote experts;
+    the `-1` sentinels are passed through to the local kernel, which decides whether
+    to skip them (eager/scattermoe) or mask them (grouped_mm).
+    """
+    if hidden_states.dtype != torch.bfloat16:
+        original_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.bfloat16)
+    else:
+        original_dtype = None
+
+    buffer = get_buffer()
+    E_global = getattr(self, "num_experts_global", self.num_experts)
+
+    topk_idx_i64 = top_k_index.to(torch.int64)
+    topk_w_f32 = top_k_weights.to(torch.float32)
+
+    # Cap tokens-per-expert before building the dispatch layout — an overloaded expert deadlocks
+    # DeepEP's intranode combine (the GLM router concentrates with depth).
+    if _TOKEN_CAPACITY is not None:
+        topk_idx_i64, topk_w_f32 = _apply_expert_capacity(
+            topk_idx_i64, topk_w_f32, _TOKEN_CAPACITY
+        )
+
+    # Drop padding tokens from the dispatch. Under non-packed training with
+    # `pad_to_sequence_len`, padding rows carry identical embeddings, so the router
+    # sends them all to the same one or two experts — a routing imbalance DeepEP's
+    # intranode dispatch can't survive ('unspecified launch failure'). Sentinelling
+    # their routing to -1 makes `get_dispatch_layout` skip them entirely (they get a
+    # zero expert output, which is correct: their loss is masked anyway). Real long
+    # sequences (packed, or a single long sample) have no padding and are unaffected.
+    valid = _get_valid_token_mask()
+    if valid is not None and valid.numel() == topk_idx_i64.shape[0]:
+        import os as _os
+
+        if _os.environ.get("AXOLOTL_EP_DEBUG"):
+            from axolotl.utils.logging import get_logger
+
+            get_logger(__name__).info(
+                f"EP padding sentinel: {int(valid.sum())}/{valid.numel()} real tokens dispatched"
+            )
+        topk_idx_i64 = topk_idx_i64.masked_fill(~valid.view(-1, 1), -1)
+
+    # Layout is non-differentiable (bookkeeping only).
+    with torch.no_grad():
+        num_per_rank, _, num_per_expert, is_in_rank, _ = buffer.get_dispatch_layout(
+            topk_idx_i64, E_global
+        )
+
+    recv_x, recv_topk_idx, recv_topk_weights, handle_holder = _DeepEPDispatch.apply(
+        hidden_states,
+        topk_idx_i64,
+        topk_w_f32,
+        num_per_rank,
+        num_per_expert,
+        is_in_rank,
+    )
+
+    # Pass the raw -1-tagged routing through; each local kernel handles sentinels
+    # its own way (eager/scattermoe skip them, grouped_mm masks internally).
+    local_out = _LOCAL_KERNELS[kernel_name](
+        self, recv_x, recv_topk_idx, recv_topk_weights
+    )
+
+    barrier_ep()  # wait out the local-kernel autotune skew before the combine collective
+    combined = _DeepEPCombine.apply(local_out, handle_holder)
+
+    if original_dtype is not None:
+        combined = combined.to(original_dtype)
+    return combined
+
+
+def deep_ep_experts_forward(self, hidden_states, top_k_index, top_k_weights):
+    return _deep_ep_forward(
+        self, hidden_states, top_k_index, top_k_weights, kernel_name="eager"
+    )
+
+
+def deep_ep_grouped_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights):
+    return _deep_ep_forward(
+        self, hidden_states, top_k_index, top_k_weights, kernel_name="grouped_mm"
+    )
+
+
+def deep_ep_scattermoe_experts_forward(self, hidden_states, top_k_index, top_k_weights):
+    return _deep_ep_forward(
+        self, hidden_states, top_k_index, top_k_weights, kernel_name="scattermoe"
+    )
+
+
+def deep_ep_sonicmoe_experts_forward(self, hidden_states, top_k_index, top_k_weights):
+    return _deep_ep_forward(
+        self, hidden_states, top_k_index, top_k_weights, kernel_name="sonicmoe"
+    )
+
+
+REGISTRY = {
+    "deep_ep": deep_ep_experts_forward,
+    "deep_ep_grouped_mm": deep_ep_grouped_mm_experts_forward,
+    "deep_ep_scattermoe": deep_ep_scattermoe_experts_forward,
+    "deep_ep_sonicmoe": deep_ep_sonicmoe_experts_forward,
+}
+
+
+def register_all() -> None:
+    """Register the four names in `ALL_EXPERTS_FUNCTIONS` and whitelist them."""
+    from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+    from transformers.modeling_utils import PreTrainedModel
+
+    for name, fn in REGISTRY.items():
+        ALL_EXPERTS_FUNCTIONS.register(name, fn)
+
+    if not getattr(
+        PreTrainedModel.get_correct_experts_implementation, "_deep_ep_patched", False
+    ):
+        original = PreTrainedModel.get_correct_experts_implementation
+
+        def patched(self, requested_experts):
+            if requested_experts in REGISTRY:
+                return requested_experts
+            return original(self, requested_experts)
+
+        patched._deep_ep_patched = True  # type: ignore[attr-defined]
+        PreTrainedModel.get_correct_experts_implementation = patched  # type: ignore[assignment]
+
+
+def kernel_to_registered_name(kernel: str) -> str:
+    """Map `expert_parallel_local_kernel` -> registered name."""
+    return {
+        "eager": "deep_ep",
+        "grouped_mm": "deep_ep_grouped_mm",
+        "scattermoe": "deep_ep_scattermoe",
+        "sonicmoe": "deep_ep_sonicmoe",
+    }[kernel]
diff --git a/src/axolotl/integrations/expert_parallel/plugin.py b/src/axolotl/integrations/expert_parallel/plugin.py
new file mode 100644
index 0000000000..8e9e5359d9
--- /dev/null
+++ b/src/axolotl/integrations/expert_parallel/plugin.py
@@ -0,0 +1,491 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+"""Expert-Parallel (DeepEP) plugin for axolotl."""
+
+from __future__ import annotations
+
+from importlib.util import find_spec
+
+import torch
+import torch.distributed as dist
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def expert_shard_axis(mesh_dim_names) -> str | None:
+    """The non-``ep`` mesh axis the routed experts FSDP-shard on under EP composition, or ``None``.
+
+    Prefers ``dp_shard`` (EP×dp_shard: experts shard on the data axis); falls back to ``cp`` (EP×cp,
+    where the cp ranks of an ep-group hold the SAME experts since cp shards the sequence, not the
+    experts, so FSDP-sharding them on cp keeps each rank from holding the full ep-group slice). Returns
+    ``None`` for pure EP (no secondary axis) or when there is no ``ep`` axis to compose with — those
+    paths don't pre-wrap the experts here.
+    """
+    names = tuple(mesh_dim_names or ())
+    if "ep" not in names:
+        return None
+    if "dp_shard" in names:
+        return "dp_shard"
+    if "cp" in names:
+        return "cp"
+    return None
+
+
+class ExpertParallelPlugin(BasePlugin):
+    """Plugin that swaps MoE dispatch/combine for DeepEP-fused kernels."""
+
+    def get_input_args(self):
+        return "axolotl.integrations.expert_parallel.ExpertParallelArgs"
+
+    def pre_model_load(self, cfg):
+        if not self._is_ep_enabled(cfg):
+            return
+
+        if not self._deep_ep_available(cfg):
+            return  # already-warned fallback path
+
+        # Cross-cfg validation that args.py can't do (it only sees its own fields).
+        self._validate_mesh_axes(cfg)
+
+        from .experts_fn import kernel_to_registered_name, register_all
+
+        register_all()
+
+        # Upgrade the user's chosen local kernel to its DeepEP-wrapped variant.
+        local_kernel = self._infer_local_kernel(cfg)
+        composite = kernel_to_registered_name(local_kernel)
+        previous = getattr(cfg, "experts_implementation", None)
+        cfg.experts_implementation = composite
+        LOG.debug(
+            f"expert_parallel: experts_implementation {previous!r} -> {composite!r} "
+            f"(local kernel: {local_kernel!r})"
+        )
+
+    def post_model_build(self, cfg, model):
+        if not self._is_ep_enabled(cfg):
+            return
+        if not self._deep_ep_available(cfg):
+            return
+
+        from .buffer import configure_buffer
+        from .shard import shard_expert_weights
+
+        ep_group = self._resolve_ep_group(cfg)
+        sharded = shard_expert_weights(model, ep_group)
+
+        if sharded == 0:
+            LOG.warning(
+                "expert_parallel_enabled=true but no Experts modules were detected "
+                "for sharding (model uses a non-canonical layout, or single-rank). "
+                "DeepEP dispatch/combine will run as a no-op."
+            )
+
+        configure_buffer(
+            ep_group=ep_group,
+            num_nvl_bytes=cfg.expert_parallel_num_nvl_bytes,
+            num_rdma_bytes=cfg.expert_parallel_num_rdma_bytes,
+        )
+        from .experts_fn import set_token_capacity
+
+        set_token_capacity(getattr(cfg, "expert_parallel_token_capacity", None))
+        # Pure-EP path: register the grad-scale hook now. FSDP+EP defers
+        # registration to `fully_shard_experts` (after experts become DTensors).
+        if (cfg.dp_shard_size or 1) <= 1:
+            ep_size = cfg.expert_parallel_size or 1
+            self._register_expert_grad_scale(model, ep_size)
+
+    def post_model_load(self, cfg, model):
+        """Propagate DDP-ignored params to the outermost model wrapper.
+
+        `post_model_build` set `_ddp_params_and_buffers_to_ignore` on the inner
+        model. After PEFT wraps it (in `PeftModel`), DDP wraps `PeftModel`, but
+        DDP looks for the attribute on the top-level module — which is now
+        `PeftModel`, not our inner model. Mirror the list up.
+        """
+        if not self._is_ep_enabled(cfg):
+            return
+
+        self._register_padding_dispatch_hook(model)
+
+        # Find the inner module that has the attribute (shard set it on whatever
+        # was the top-level model at post_model_build time).
+        inner = getattr(model, "base_model", model)
+        # base_model may itself be wrapped (e.g., LoraModel.model). Recurse.
+        while not hasattr(inner, "_ddp_params_and_buffers_to_ignore"):
+            next_inner = getattr(inner, "model", None) or getattr(
+                inner, "base_model", None
+            )
+            if next_inner is None or next_inner is inner:
+                break
+            inner = next_inner
+
+        ignore_list = getattr(inner, "_ddp_params_and_buffers_to_ignore", None)
+        if not ignore_list:
+            return
+
+        # PEFT prefixes parameter names. Re-resolve the list against the wrapper
+        # so DDP can match by name.
+        resolved: list[str] = []
+        wrapper_param_names = {n for n, _ in model.named_parameters()}
+        wrapper_buffer_names = {n for n, _ in model.named_buffers()}
+        all_names = wrapper_param_names | wrapper_buffer_names
+
+        for short_name in ignore_list:
+            # Match either an exact suffix or with PEFT's `base_model.model.` prefix.
+            for full in all_names:
+                if (
+                    full == short_name
+                    or full.endswith("." + short_name)
+                    or full.endswith(short_name)
+                ):
+                    resolved.append(full)
+
+        # De-dup while preserving order.
+        seen = set()
+        resolved = [n for n in resolved if not (n in seen or seen.add(n))]
+
+        existing = list(getattr(model, "_ddp_params_and_buffers_to_ignore", []))
+        model._ddp_params_and_buffers_to_ignore = existing + resolved
+        LOG.debug(
+            f"expert_parallel: propagated {len(resolved)} DDP-ignored param "
+            f"name(s) onto outer wrapper {type(model).__name__}."
+        )
+
+    @staticmethod
+    def _register_padding_dispatch_hook(model) -> None:
+        """Feed the batch's real-token mask to the DeepEP dispatch so padding tokens are
+        not routed (they'd otherwise pile onto one expert and break intranode dispatch).
+
+        A model-level forward pre-hook reads the 2D ``attention_mask`` (1=real, 0=pad) and
+        stashes a flattened ``[B*S]`` bool mask; ``_deep_ep_forward`` sentinels those rows.
+        Under sample packing there is no 2D mask, but the multipack collator pads partial
+        packs to ``seq_len`` — those identical pad embeddings still pile onto one expert and
+        break DeepEP intranode dispatch — so fall back to ``input_ids != pad_token_id``."""
+        from .experts_fn import set_valid_token_mask
+
+        if getattr(model, "_ep_padding_hook", False):
+            return
+
+        pad_id = getattr(getattr(model, "config", None), "pad_token_id", None)
+
+        def _pre_hook(_module, args, kwargs):
+            am = kwargs.get("attention_mask")
+            if am is None and args:
+                am = next(
+                    (a for a in args if torch.is_tensor(a) and a.dim() == 2), None
+                )
+            if am is not None and am.dim() == 2:
+                set_valid_token_mask((am != 0).reshape(-1))
+                return args, kwargs
+            # Packing (no 2D mask): exclude pad rows so they don't overload one expert.
+            ids = kwargs.get("input_ids")
+            if ids is None and args:
+                ids = next(
+                    (
+                        a
+                        for a in args
+                        if torch.is_tensor(a)
+                        and a.dim() == 2
+                        and a.dtype in (torch.long, torch.int, torch.int32)
+                    ),
+                    None,
+                )
+            set_valid_token_mask(
+                (ids != pad_id).reshape(-1)
+                if (ids is not None and pad_id is not None)
+                else None
+            )
+            return args, kwargs
+
+        model.register_forward_pre_hook(_pre_hook, with_kwargs=True)
+        model._ep_padding_hook = True
+
+    @staticmethod
+    def _infer_local_kernel(cfg) -> str:
+        """Decide which local-experts kernel runs under DeepEP dispatch.
+
+        `use_scattermoe` / `use_sonicmoe` are the master flags from
+        `kernels/args.py` and take precedence; otherwise fall back to
+        `experts_implementation` (`eager` / `grouped_mm` / `batched_mm`).
+        """
+        if getattr(cfg, "use_scattermoe", False):
+            return "scattermoe"
+
+        if getattr(cfg, "use_sonicmoe", False):
+            return "sonicmoe"
+
+        ei = getattr(cfg, "experts_implementation", None)
+        if ei in ("grouped_mm", "batched_mm"):
+            return "grouped_mm"
+        if ei == "eager":
+            return "eager"
+        # default: upstream-shipped fast kernel
+        return "grouped_mm"
+
+    # Cached 2D DeviceMesh when EP composes with FSDP. Set by `_resolve_ep_group`.
+    _device_mesh = None
+
+    @staticmethod
+    def _accelerate_mesh():
+        """Return accelerate's device_mesh, force-creating the Accelerator if
+        needed. The AcceleratorState singleton makes this idempotent w.r.t.
+        the trainer's later `Accelerator()` call.
+        """
+        from accelerate import Accelerator
+        from accelerate.state import AcceleratorState
+
+        try:
+            state = AcceleratorState()
+        except (RuntimeError, AttributeError, ValueError) as e:
+            # ValueError occurs from unittest due to no trainer constructing Accelerator()
+            LOG.debug(f"expert_parallel: AcceleratorState() not ready: {e}")
+            return None
+        mesh = getattr(state, "device_mesh", None)
+        if mesh is not None:
+            return mesh
+        try:
+            Accelerator()
+        except (RuntimeError, ValueError) as e:
+            LOG.debug(f"expert_parallel: Accelerator() force-init failed: {e}")
+            return None
+        return getattr(AcceleratorState(), "device_mesh", None)
+
+    @staticmethod
+    def _resolve_ep_group(cfg):
+        """Return the EP ProcessGroup.
+
+        For FSDP+EP, returns `accelerate_mesh["ep"].get_group()` — the same
+        process group that accelerate's parallelism_config built. For pure EP
+        (ep_size == world_size, no FSDP), returns `dist.group.WORLD`.
+        """
+        if not dist.is_available() or not dist.is_initialized():
+            return None
+
+        world_size = dist.get_world_size()
+        ep_size = getattr(cfg, "expert_parallel_size", 1) or 1
+        dp_shard_size = getattr(cfg, "dp_shard_size", None) or 1
+        tp_size = getattr(cfg, "tensor_parallel_size", None) or 1
+        cp_size = getattr(cfg, "context_parallel_size", None) or 1
+
+        if ep_size <= 1:
+            return dist.group.WORLD
+
+        # Validate the world_size = product check.
+        product = ep_size * dp_shard_size * tp_size * cp_size
+        if product != world_size:
+            raise ValueError(
+                f"expert_parallel_size ({ep_size}) * dp_shard_size ({dp_shard_size}) "
+                f"* tensor_parallel_size ({tp_size}) * context_parallel_size ({cp_size}) "
+                f"= {product}, but world_size = {world_size}. The product must equal "
+                f"the world size for orthogonal mesh axes to be valid."
+            )
+
+        if ep_size == world_size:
+            return dist.group.WORLD
+
+        # EP composed with FSDP (`dp_shard`) and/or context parallel (`cp`) on orthogonal mesh
+        # axes — read the ep group from accelerate's mesh, or build one ourselves if accelerate
+        # hasn't (e.g., topology unit tests that drive `_resolve_ep_group` directly). Experts shard
+        # on `ep` (tokens move via all-to-all); the sequence shards on `cp` (DSA attention gathers
+        # the compressed KV on that axis); non-expert weights shard on `dp_shard`. TP is still
+        # unsupported in composition.
+        if dp_shard_size > 1 or cp_size > 1:
+            if tp_size > 1:
+                raise NotImplementedError(
+                    "EP × TP composition not yet supported. Got "
+                    f"ep={ep_size}, dp_shard={dp_shard_size}, tp={tp_size}, cp={cp_size}. "
+                    "Supported: EP, EP × dp_shard, EP × cp, EP × cp × dp_shard."
+                )
+            mesh = ExpertParallelPlugin._accelerate_mesh()
+            if mesh is None or "ep" not in (mesh.mesh_dim_names or ()):
+                from torch.distributed.device_mesh import init_device_mesh
+
+                # Fallback mesh from the >1 axes (ep outermost). Orthogonality of the ep/cp/dp
+                # groups is what matters; accelerate's mesh is preferred when present so the ep
+                # group matches the one used for the experts' FSDP exclusion.
+                axes = [("ep", ep_size)]
+                if cp_size > 1:
+                    axes.append(("cp", cp_size))
+                if dp_shard_size > 1:
+                    axes.append(("dp_shard", dp_shard_size))
+                mesh = init_device_mesh(
+                    "cuda" if torch.cuda.is_available() else "cpu",
+                    tuple(s for _, s in axes),
+                    mesh_dim_names=tuple(n for n, _ in axes),
+                )
+            ExpertParallelPlugin._device_mesh = mesh
+            LOG.debug(
+                f"expert_parallel: ep mesh shape={tuple(mesh.shape)} "
+                f"axes={mesh.mesh_dim_names}; ep group "
+                f"members={dist.get_process_group_ranks(mesh['ep'].get_group())}"
+            )
+            return mesh["ep"].get_group()
+
+        # ep_size > 1, ep_size < world_size, no dp_shard/cp to fill the rest — invalid.
+        raise ValueError(
+            f"expert_parallel_size ({ep_size}) < world_size ({world_size}) "
+            "without dp_shard_size/context_parallel_size > 1 to fill the remaining axes is not "
+            "supported. Set dp_shard_size and/or context_parallel_size such that "
+            "ep × cp × dp_shard == world_size, or set expert_parallel_size = world_size for pure EP."
+        )
+
+    @staticmethod
+    def _resolve_cp_group(cfg):
+        """Return the context-parallel ProcessGroup (the `cp` axis of the EP mesh), or None when
+        ``context_parallel_size <= 1``. The DSA attention shards the sequence on this axis (gathering
+        the compressed KV across it); experts shard on the orthogonal ``ep`` axis. Reads the mesh
+        built by ``_resolve_ep_group`` / accelerate."""
+        cp_size = getattr(cfg, "context_parallel_size", None) or 1
+        if cp_size <= 1:
+            return None
+        mesh = (
+            ExpertParallelPlugin._device_mesh or ExpertParallelPlugin._accelerate_mesh()
+        )
+        if mesh is not None and "cp" in (mesh.mesh_dim_names or ()):
+            return mesh["cp"].get_group()
+        return None
+
+    @staticmethod
+    def fully_shard_experts(model, dp_shard_mesh, fsdp2_kwargs):
+        """Pre-wrap each Experts module with FSDP on the `dp_shard` axis.
+
+        Called from the patched `fsdp2_prepare_model` BEFORE the outer auto-wrap
+        so experts become FSDPModules and the auto-wrap walker skips them.
+        Inherits the outer wrap's policy (mp, offload, reshard) so inner/outer
+        collective dtypes line up; only `mesh` is overridden.
+        """
+        from torch.distributed.fsdp import fully_shard
+
+        from .shard import (
+            _detect_experts_modules,
+            _is_param_wrapper,
+            _real_experts_base,
+        )
+
+        kwargs = dict(fsdp2_kwargs)
+        kwargs["mesh"] = dp_shard_mesh
+        kwargs.pop("ignored_params", None)
+
+        for _name, module in _detect_experts_modules(model):
+            fully_shard(module, **kwargs)
+
+        # `target_parameters` expert LoRA lives on the ParamWrapper chain wrapping the experts
+        # module (which `_detect_experts_modules` skips). Left to the outer decoder-layer auto-wrap
+        # it shards on the FULL ep×dp mesh — i.e. ACROSS the ep axis — corrupting the per-ep-rank
+        # expert slice (grads averaged over ranks owning different experts; save reconstructs the
+        # wrong shape). Wrap the OUTERMOST expert ParamWrapper as its own FSDP unit on dp_shard:
+        # its forward IS the fused-LoRA fastpath, so FSDP unshards the adapter (incl. the nested
+        # inner wrapper's, which is not a separate unit) to plain tensors right before the kernel
+        # reads them — sharded on the same axis as the weights, but gathered during use.
+        all_pws = [m for _n, m in model.named_modules() if _is_param_wrapper(m)]
+        inner = {getattr(pw, "base_layer", None) for pw in all_pws}
+        outer_expert_pws = [
+            pw
+            for pw in all_pws
+            if pw not in inner
+            and _real_experts_base(pw) is not None
+            and getattr(_real_experts_base(pw), "num_local_experts", None) is not None
+        ]
+        for pw in outer_expert_pws:
+            fully_shard(pw, **kwargs)
+
+        LOG.debug(
+            f"expert_parallel: pre-wrapped Experts modules + {len(outer_expert_pws)} expert "
+            f"ParamWrapper(s) on dp_shard mesh (size={dp_shard_mesh.size()})."
+        )
+
+        root = dp_shard_mesh._get_root_mesh()
+        ep_size = (
+            root["ep"].size()
+            if root is not None and "ep" in (root.mesh_dim_names or ())
+            else 1
+        )
+        ExpertParallelPlugin._register_expert_grad_scale(model, ep_size)
+
+    @staticmethod
+    def _register_expert_grad_scale(model, ep_size: int) -> int:
+        """Scale expert weight grads by `1/ep_size` so EP / FSDP / FSDP+EP
+        produce the same effective gradient.
+        """
+        from .shard import _detect_experts_modules
+
+        if ep_size <= 1:
+            return 0
+        scale = 1.0 / ep_size
+
+        def _scale(p):
+            if p.grad is not None:
+                p.grad.mul_(scale)
+
+        n_hooks = 0
+        for _name, module in _detect_experts_modules(model):
+            for p in module.parameters(recurse=True):
+                # Only trainable params have a grad to scale — and a hook can only be
+                # registered on a tensor that requires grad. Under LoRA the base expert
+                # weights are frozen (only the adapters train), so skip them.
+                if not p.requires_grad:
+                    continue
+                p.register_post_accumulate_grad_hook(_scale)
+                n_hooks += 1
+        LOG.debug(
+            f"expert_parallel: registered {n_hooks} expert grad-scale hooks "
+            f"(scale = 1/{ep_size})"
+        )
+        return n_hooks
+
+    @staticmethod
+    def _is_ep_enabled(cfg) -> bool:
+        """EP is enabled when expert_parallel_size > 1 (mirrors TP / DP UX)."""
+        ep_size = getattr(cfg, "expert_parallel_size", 1) or 1
+        return ep_size > 1
+
+    @staticmethod
+    def _validate_mesh_axes(cfg) -> None:
+        """Sanity-check the mesh-axis sizes early, with a clear error.
+
+        `_resolve_ep_group` re-validates at process-group construction time;
+        this catches misconfigured YAMLs before model loading wastes minutes.
+        """
+        ep_size = getattr(cfg, "expert_parallel_size", 1) or 1
+        if ep_size <= 1:
+            return
+
+        if not (dist.is_available() and dist.is_initialized()):
+            return  # validated at process-group time
+        world_size = dist.get_world_size()
+        if world_size <= 1:
+            return  # single-rank context; mesh shapes are meaningless
+        dp_shard_size = getattr(cfg, "dp_shard_size", None) or 1
+        tp_size = getattr(cfg, "tensor_parallel_size", None) or 1
+        cp_size = getattr(cfg, "context_parallel_size", None) or 1
+
+        product = ep_size * dp_shard_size * tp_size * cp_size
+        if product != world_size:
+            raise ValueError(
+                f"expert_parallel: world_size ({world_size}) must equal "
+                f"expert_parallel_size ({ep_size}) * dp_shard_size ({dp_shard_size}) "
+                f"* tensor_parallel_size ({tp_size}) * context_parallel_size ({cp_size}) "
+                f"= {product}."
+            )
+
+    @staticmethod
+    def _deep_ep_available(cfg) -> bool:
+        if find_spec("deep_ep") is not None:
+            return True
+        msg = (
+            "expert_parallel_enabled=true but `deep_ep` is not importable. "
+            "See the integration README for install instructions."
+        )
+        if cfg.expert_parallel_fallback_on_unsupported:
+            LOG.warning(msg + " Falling back to standard experts implementation.")
+            return False
+        raise ImportError(msg)
diff --git a/src/axolotl/integrations/expert_parallel/shard.py b/src/axolotl/integrations/expert_parallel/shard.py
new file mode 100644
index 0000000000..06837924f9
--- /dev/null
+++ b/src/axolotl/integrations/expert_parallel/shard.py
@@ -0,0 +1,605 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+"""Generic expert-weight sharding for `@use_experts_implementation` modules.
+
+After this runs (in `post_model_build`, before FSDP wraps), each rank's Experts
+modules hold only their local slice of the experts dim. The registered
+`deep_ep_*` forward function then handles dispatch -> local compute -> combine.
+"""
+
+from __future__ import annotations
+
+import gc
+
+import torch
+import torch.distributed as dist
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _replace_with_slice(module, attr_name: str, start: int, end: int) -> None:
+    """Replace `module.{attr_name}` with a fresh-allocation slice [start:end].
+
+    `tensor[start:end]` on a contiguous storage returns a VIEW that shares the
+    same underlying allocation. `.contiguous()` is a no-op on an already-
+    contiguous view, so wrapping the view in a new Parameter does NOT release
+    the original full-size storage — refcount on the storage stays >= 1 from
+    the new view.
+
+    `.clone()` forces a fresh allocation. After we drop the old Parameter,
+    the original storage's refcount drops to zero and PyTorch's caching
+    allocator reclaims the memory on the next `empty_cache()`.
+    """
+    old = getattr(module, attr_name)
+    new_data = old.data[start:end].detach().clone().contiguous()
+    requires_grad = old.requires_grad
+    # Drop the old Parameter from the module's _parameters registry first so
+    # nothing is holding a reference when we assign the new one.
+    if attr_name in module._parameters:
+        del module._parameters[attr_name]
+    setattr(
+        module,
+        attr_name,
+        torch.nn.Parameter(new_data, requires_grad=requires_grad),
+    )
+    # `old` goes out of scope on return.
+
+
+def _scatter_expert_from_rank0(module, attr_name, e_local, dp_size):
+    """Populate ``module.{attr_name}`` with THIS rank's real ``[e_local]`` expert slice by scattering
+    GLOBAL rank-0's full expert tensor over the WORLD group. Under cpu_ram_efficient_loading only global
+    rank 0 materializes real weights, so it is the single source (sourcing from an ep-subgroup's rank-0
+    would crash on the meta ranks). Each rank ``r`` receives its ep-group's slice
+    ``full[(r//dp_size)*e_local : ((r//dp_size)+1)*e_local]`` — the ``dp_size`` ranks within an ep-group
+    get the SAME ep slice (the dp axis FSDP-shards it across them later). ``dp_size == 1`` is pure EP
+    (each rank its own [e_local]); ``dp_size > 1`` is EP×dp_shard / EP×cp composition. Handles torchao
+    NVFP4Tensor (qdata/scale/per_tensor_scale) and plain tensors; runs on GPU (NCCL), result moved to
+    the param's original device."""
+    import torch.distributed as dist
+
+    old = getattr(module, attr_name)
+    nv = old.data
+    dev = torch.device("cuda", torch.cuda.current_device())
+    world = dist.get_world_size()
+    is_src = dist.get_rank() == 0
+    dst_device = (
+        old.data.device if old.data.device.type != "meta" else torch.device("cpu")
+    )
+    requires_grad = old.requires_grad
+
+    def _scatter_dim0(full_comp, like):
+        out = torch.empty((e_local, *like.shape[1:]), dtype=like.dtype, device=dev)
+        chunks = None
+        if is_src:
+            chunks = [
+                full_comp[(r // dp_size) * e_local : (r // dp_size + 1) * e_local]
+                .contiguous()
+                .to(dev)
+                for r in range(world)
+            ]
+        dist.scatter(out, scatter_list=chunks, src=0)
+        return out
+
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+    except ImportError:
+        NVFP4Tensor = None
+
+    if NVFP4Tensor is not None and isinstance(nv, NVFP4Tensor):
+        full_q = nv.qdata if is_src else None
+        local_q = _scatter_dim0(full_q, nv.qdata)
+        # scale (e4m3) isn't a collective dtype on all backends -> scatter the uint8 view.
+        full_s = nv.scale.view(torch.uint8) if is_src else None
+        local_s = _scatter_dim0(full_s, nv.scale.view(torch.uint8)).view(nv.scale.dtype)
+        pts = getattr(nv, "per_tensor_scale", None)
+        local_pts = None
+        if pts is not None:
+            if pts.dim() >= 1 and pts.shape[0] == nv.qdata.shape[0]:
+                local_pts = _scatter_dim0(pts if is_src else None, pts)
+            else:  # replicated scalar
+                local_pts = pts.to(dev).clone()
+                dist.broadcast(local_pts, src=0)
+        local_nv = NVFP4Tensor(
+            local_q.to(dst_device),
+            local_s.to(dst_device),
+            nv.block_size,
+            nv.dtype,
+            per_tensor_scale=(
+                local_pts.to(dst_device) if local_pts is not None else None
+            ),
+        )
+        new_param = torch.nn.Parameter(local_nv, requires_grad=requires_grad)
+    else:
+        local = _scatter_dim0(nv if is_src else None, nv)
+        new_param = torch.nn.Parameter(
+            local.to(dst_device), requires_grad=requires_grad
+        )
+
+    if attr_name in module._parameters:
+        del module._parameters[attr_name]
+    setattr(module, attr_name, new_param)
+
+
+def _detect_experts_modules(model):
+    """Yield (name, module) pairs for every module that looks like an Experts class.
+
+    Detection: 3D `gate_up_proj` and `down_proj` parameters with experts on dim 0.
+    This is the canonical layout enforced by `@use_experts_implementation`.
+    Mixtral's `ModuleList[MixtralBlockSparseTop2MLP]` does NOT match — out of scope
+    for v1.
+    """
+    for name, module in model.named_modules():
+        # A PEFT ParamWrapper delegates `gate_up_proj` to its `base_layer`, so it
+        # would match too and double-count the experts (double grad-scale, redundant
+        # fully_shard). Yield only the real experts module (the wrapped base_layer).
+        if _is_param_wrapper(module):
+            continue
+        gp = getattr(module, "gate_up_proj", None)
+        dp = getattr(module, "down_proj", None)
+        if gp is None or dp is None:
+            continue
+        if not (
+            isinstance(gp, torch.nn.Parameter) and isinstance(dp, torch.nn.Parameter)
+        ):
+            continue
+        if gp.dim() != 3 or dp.dim() != 3:
+            continue
+        yield name, module
+
+
+def shard_expert_weights(model, ep_group) -> int:
+    """Slice expert weights along dim 0 per the EP rank.
+
+    Args:
+        model: A built (but not yet FSDP-wrapped) HuggingFace model.
+        ep_group: `torch.distributed.ProcessGroup` for EP, or `None` for
+            single-rank (no-op).
+
+    Returns:
+        Number of Experts modules sharded (0 if EP disabled or none found).
+
+    Raises:
+        ValueError: if any Experts module's `num_experts` is not divisible by
+            the EP world size.
+
+    DDP composition: the sharded params hold DIFFERENT content per rank, so we
+    add their fully-qualified names to `model._ddp_params_and_buffers_to_ignore`
+    to prevent the startup broadcast from copying rank 0's slice everywhere.
+    FSDP composition is handled in `ExpertParallelPlugin.fully_shard_experts`.
+    """
+    if ep_group is None:
+        return 0
+
+    ep_size = dist.get_world_size(ep_group)
+    if ep_size <= 1:
+        return 0
+
+    ep_rank = dist.get_rank(ep_group)
+    sharded = 0
+    ignore_names: list[str] = []
+
+    for name, module in _detect_experts_modules(model):
+        gp = module.gate_up_proj
+        E = gp.shape[0]
+        if E % ep_size != 0:
+            raise ValueError(
+                f"Expert module {name!r}: num_experts={E} not divisible by "
+                f"ep_size={ep_size}. Adjust the model config or ep_size."
+            )
+        E_local = E // ep_size
+        start = ep_rank * E_local
+
+        # Scatter global rank-0's REAL experts to every rank's ep-group slice. Under
+        # cpu_ram_efficient_loading only global rank 0 has data; plain-slicing (the old path) kept
+        # only rank 0's own ep-group's slice and zeroed the rest, so ep-groups 1..N had dead experts.
+        # dp_size>1 (EP×dp_shard / EP×cp) gives the dp ranks within an ep-group the same ep slice; the
+        # FSDP dp-axis shards it across them afterwards. See _scatter_expert_from_rank0.
+        dp_size = dist.get_world_size() // ep_size
+        with torch.no_grad():
+            _scatter_expert_from_rank0(module, "gate_up_proj", E_local, dp_size)
+            _scatter_expert_from_rank0(module, "down_proj", E_local, dp_size)
+            for bias_name in ("gate_up_proj_bias", "down_proj_bias"):
+                bias = getattr(module, bias_name, None)
+                if (
+                    isinstance(bias, torch.nn.Parameter)
+                    and bias.dim() >= 1
+                    and bias.shape[0] == E
+                ):
+                    _scatter_expert_from_rank0(module, bias_name, E_local, dp_size)
+
+        # Stash metadata the registered fn needs.
+        module.local_expert_offset = start
+        module.num_local_experts = E_local
+        module.num_experts_global = E
+        # Single global expert count for the cpu_ram_efficient load path (all routed-expert modules
+        # share it); used to reshape the global expert-LoRA adapter when slicing per-rank shards.
+        model._ep_num_experts_global = E
+        # Override the kernel's view of num_experts for grouped_mm/scattermoe bucketing.
+        module.num_experts = E_local
+
+        # Mark sharded params as DDP-ignored — they hold rank-specific content
+        # and must NOT be broadcast from rank 0 at DDP construction.
+        ignore_names.append(f"{name}.gate_up_proj")
+        ignore_names.append(f"{name}.down_proj")
+        for bias_name in ("gate_up_proj_bias", "down_proj_bias"):
+            if isinstance(getattr(module, bias_name, None), torch.nn.Parameter):
+                ignore_names.append(f"{name}.{bias_name}")
+
+        sharded += 1
+
+    if sharded:
+        # Drop refs to any leftover full-size storage and return cached memory
+        # to the allocator. Without this, the original `from_pretrained`
+        # allocations stay reserved and `memory_allocated()` doesn't drop.
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        # Append (don't overwrite) — other systems may have set this too.
+        existing = list(getattr(model, "_ddp_params_and_buffers_to_ignore", []))
+        model._ddp_params_and_buffers_to_ignore = existing + ignore_names
+        LOG.info(
+            f"Sharded {sharded} Experts module(s) along the experts dim "
+            f"(ep_rank={ep_rank}, ep_size={ep_size}). "
+            f"Marked {len(ignore_names)} param(s) as DDP-ignored."
+        )
+    return sharded
+
+
+def _is_param_wrapper(module) -> bool:
+    """True for a PEFT ParamWrapper, including after ``fully_shard`` renames its class
+    (e.g. ``FSDPParamWrapper``) — FSDP2 sets ``__class__`` to a subclass, so ``isinstance``
+    still holds while a ``type(...).__name__ == 'ParamWrapper'`` check would miss it."""
+    try:
+        from peft.tuners.lora.layer import ParamWrapper
+
+        if isinstance(module, ParamWrapper):
+            return True
+    except (ImportError, AttributeError):
+        pass
+    return "ParamWrapper" in type(module).__name__
+
+
+def _real_experts_base(wrapper):
+    """Walk a (possibly chained) PEFT ParamWrapper down to the real experts module."""
+    base = getattr(wrapper, "base_layer", None)
+    while base is not None and _is_param_wrapper(base):
+        base = getattr(base, "base_layer", None)
+    return base
+
+
+def _slice_expert_lora_param(linear, dim: int, e_global: int, start: int, end: int):
+    """Replace ``linear.weight`` with its local-experts slice.
+
+    ``dim`` is the axis carrying the ``E*r`` packing:
+      * lora_A: ``[E*r, in]`` expert-major  -> slice rows ``[start*r:end*r]`` (dim 0).
+      * lora_B: ``[out, r*E]`` rank-major    -> reshape ``[out, r, E]``, take experts, flatten.
+    Returns the rank ``r`` (so the caller can keep ``in_features``/``out_features`` consistent).
+    """
+    w = linear.weight
+    if dim == 0:  # lora_A: expert-major rows
+        r = w.shape[0] // e_global
+        new = w.data[start * r : end * r, :].detach().clone()
+        linear.out_features = new.shape[0]
+    else:  # lora_B: rank-major columns [out, r, E]
+        out_dim = w.shape[0]
+        r = w.shape[1] // e_global
+        new = (
+            w.data.reshape(out_dim, r, e_global)[:, :, start:end]
+            .reshape(out_dim, r * (end - start))
+            .detach()
+            .clone()
+        )
+        linear.in_features = new.shape[1]
+    linear.weight = torch.nn.Parameter(new, requires_grad=w.requires_grad)
+    return r
+
+
+def ep_adapter_load_local_shard(
+    global_adapter, ep_dim, e_global, ep_coord, ep_size, placements, dp_size, dp_rank
+):
+    """Slice an EP-composition expert-LoRA adapter from rank-0's GLOBAL (all-experts) tensor down to
+    THIS rank's local FSDP shard — the inverse of ``shard_expert_lora`` + the FSDP dp/cp sharding, used
+    by the cpu_ram_efficient load path. First take this ep-group's experts, then this rank's dp/cp shard
+    along each Shard placement.
+
+    ``ep_dim`` is the adapter's expert axis: 0 for lora_A's expert-major ``[E*r, in]`` rows, 1 for
+    lora_B's ``[out, r*E]`` columns. lora_B's experts are the LAST axis of the ``[out, r, E]`` view (NOT
+    contiguous in the flat ``r*E`` dim), so a plain ``chunk`` on dim 1 would pick a rank-component, not
+    this ep-group's experts — hence the reshape-slice that mirrors :func:`_slice_expert_lora_param`.
+    """
+    from torch.distributed.tensor import Shard
+
+    e_local = e_global // ep_size
+    start, end = ep_coord * e_local, (ep_coord + 1) * e_local
+    if ep_dim == 0:  # lora_A: [E*r, in] expert-major rows
+        r = global_adapter.shape[0] // e_global
+        ep_slice = global_adapter[start * r : end * r, :]
+    else:  # lora_B: [out, r*E] -> [out, r, E], experts on the last axis
+        out_dim = global_adapter.shape[0]
+        r = global_adapter.shape[1] // e_global
+        ep_slice = global_adapter.reshape(out_dim, r, e_global)[
+            :, :, start:end
+        ].reshape(out_dim, r * e_local)
+    local = ep_slice
+    for placement in placements:
+        if isinstance(placement, Shard):
+            local = local.chunk(dp_size, dim=placement.dim)[dp_rank]
+    return local.contiguous()
+
+
+def shard_expert_lora(model, ep_size: int) -> int:
+    """Slice PEFT ``target_parameters`` expert LoRA to each rank's local experts.
+
+    PEFT sizes the LoRA for a 3D ``experts.{gate_up,down}_proj`` from the parameter's
+    own dim-0 (the *global* expert count) at adapter-application time, before EP's
+    weight slice takes effect on the parameter PEFT wrapped. Left alone, the fused
+    EP kernel (``num_experts = E_local``) and the FSDP2 parametrize merge both see a
+    full-expert LoRA against a local-expert weight -> shape mismatch. This realigns
+    the LoRA with the EP-sharded weights (same ``[offset:offset+E_local]`` slice) and
+    registers the ``1/ep_size`` expert grad-scale on the new params. Idempotent.
+
+    Run AFTER PEFT applies the adapter and BEFORE FSDP wraps. Returns the count of
+    LoRA params sliced.
+    """
+    if ep_size <= 1:
+        return 0
+
+    scale = 1.0 / ep_size
+
+    def _scale_hook(p):
+        if p.grad is not None:
+            p.grad.mul_(scale)
+
+    n = 0
+    for _name, wrapper in model.named_modules():
+        if not _is_param_wrapper(wrapper):
+            continue
+        if getattr(wrapper, "_ep_lora_sharded", False):
+            continue
+        base = _real_experts_base(wrapper)
+        e_local = getattr(base, "num_local_experts", None) if base is not None else None
+        e_global = (
+            getattr(base, "num_experts_global", None) if base is not None else None
+        )
+        if e_local is None or e_global is None or e_local >= e_global:
+            continue
+        start = base.local_expert_offset
+        end = start + e_local
+
+        for adapters, dim in (
+            (getattr(wrapper, "lora_A", {}), 0),
+            (getattr(wrapper, "lora_B", {}), 1),
+        ):
+            for ad in list(adapters.keys()):
+                _slice_expert_lora_param(adapters[ad], dim, e_global, start, end)
+                adapters[ad].weight.register_post_accumulate_grad_hook(_scale_hook)
+                n += 1
+        wrapper._ep_lora_sharded = True
+
+    if n:
+        LOG.info(
+            f"Sharded {n} expert-LoRA param(s) to local experts "
+            f"(ep_size={ep_size}, grad-scale=1/{ep_size})."
+        )
+    return n
+
+
+def save_ep_lora_adapter(model, output_dir: str, ep_group) -> bool:
+    """Write a complete LoRA adapter when experts are EP-sharded.
+
+    The attention/router LoRA is replicated across EP, but ``target_parameters`` expert
+    LoRA is EP-sharded (each rank holds ``[offset:offset+E_local]``). A plain save would
+    persist only the local rank's experts. This gathers each adapter param to a full
+    tensor (FSDP all-gather via ``full_tensor`` + EP all-gather for expert LoRA), renames
+    to PEFT adapter keys, and writes ``adapter_model.safetensors`` on rank 0. Returns
+    ``True`` if it handled the save.
+    """
+    from pathlib import Path
+
+    from peft.utils.save_and_load import get_peft_model_state_dict
+    from safetensors.torch import save_file
+
+    # The EP-sharded 3D expert params and the global expert count. Gather straight from the
+    # ParamWrappers by parameter_name (chaining `base_layer` through FSDP-wrapped units to reach
+    # the experts module via `_real_experts_base` is brittle once the outer wrapper is its own
+    # FSDP unit).
+    expert_param_names: set = set()
+    e_global = None
+    for _n, m in model.named_modules():
+        if (
+            getattr(m, "num_local_experts", None) is not None
+            and getattr(m, "num_experts_global", None) is not None
+            and m.num_local_experts < m.num_experts_global
+        ):
+            e_global = m.num_experts_global
+            for pn in ("gate_up_proj", "down_proj"):
+                if hasattr(m, pn):
+                    expert_param_names.add(pn)
+    if e_global is None or not expert_param_names:
+        return False
+
+    expert_wrappers = [
+        (wname, wrapper)
+        for wname, wrapper in model.named_modules()
+        if _is_param_wrapper(wrapper)
+        and getattr(wrapper, "parameter_name", None) in expert_param_names
+    ]
+    if not expert_wrappers:
+        return False
+
+    active = model.active_adapter if hasattr(model, "active_adapter") else "default"
+    if isinstance(active, (list, tuple)):
+        active = active[0]
+
+    # Prefer the ep group captured at FSDP-setup time; re-resolving from cfg at save can
+    # return a stale/degenerate (size-1) mesh.
+    stashed = getattr(model, "_ep_lora_group", None)
+    if stashed is not None:
+        ep_group = stashed
+
+    # Replicated (attention/router) LoRA: full tensors via FSDP all-gather, canonical PEFT keys.
+    sd = {
+        name: (p.full_tensor() if type(p).__name__ == "DTensor" else p.data).detach()
+        for name, p in model.named_parameters()
+        if "lora_" in name
+    }
+    adapter_sd = get_peft_model_state_dict(model, state_dict=sd)
+
+    # Expert LoRA: gather each wrapper's adapter across FSDP (dp_shard) + EP, key by module name.
+    gathered = 0
+    for wname, wrapper in expert_wrappers:
+        # Only EP×dp_shard/cp composition physically slices the adapter to E_local (shard_expert_lora
+        # sets _ep_lora_sharded), so it must be EP-all-gathered back to E_global. Pure EP keeps the
+        # adapter GLOBAL (E_global, forward-sliced at runtime) — gathering it would replicate every
+        # expert ep_size times into an oversized adapter, so write the FSDP-gathered tensor as-is.
+        ep_sharded = getattr(wrapper, "_ep_lora_sharded", False)
+        for sub, kind in (("lora_A", "A"), ("lora_B", "B")):
+            for w in (mod.weight for mod in getattr(wrapper, sub, {}).values()):
+                full_local = (
+                    (w.full_tensor() if type(w).__name__ == "DTensor" else w.data)
+                    .detach()
+                    .contiguous()
+                )
+                full = (
+                    gather_expert_lora_full(full_local, kind, e_global, ep_group)
+                    if ep_sharded
+                    else full_local
+                )
+                key = f"{wname}.{sub}.weight"
+                target = (
+                    key
+                    if key in adapter_sd
+                    else next(
+                        (
+                            k
+                            for k in adapter_sd
+                            if k.endswith(key.split("base_model.model.")[-1])
+                        ),
+                        None,
+                    )
+                )
+                if target is not None:
+                    adapter_sd[target] = full
+                    gathered += 1
+                else:
+                    LOG.warning(f"EP-LoRA save: no adapter key for {key!r}.")
+
+    if dist.get_rank() == 0:
+        out = Path(output_dir)
+        out.mkdir(parents=True, exist_ok=True)
+        save_file(
+            {k: v.cpu().contiguous() for k, v in adapter_sd.items()},
+            str(out / "adapter_model.safetensors"),
+        )
+        model.peft_config[active].save_pretrained(str(out))
+        LOG.info(
+            f"Saved EP-gathered LoRA adapter ({len(adapter_sd)} tensors, "
+            f"{gathered} expert params gathered across EP) to {out}."
+        )
+    dist.barrier()
+    return True
+
+
+def save_fsdp2_lora_adapter(model, output_dir: str) -> bool:
+    """Write a complete LoRA adapter under FSDP2 WITHOUT expert parallelism.
+
+    The DCP ``SHARDED_STATE_DICT`` save fails ("Failed to validate global plan") on the frozen NVFP4
+    base params (torchao tensor-subclass DTensors the planner can't validate). For a LoRA run we only
+    need the (tiny) adapter, so gather each ``lora_`` param to a full tensor via FSDP all-gather
+    (``DTensor.full_tensor``) and write ``adapter_model.safetensors`` on rank 0. ``target_parameters``
+    expert LoRA lives on PEFT ParamWrappers (``lora_A``/``lora_B`` submodules) — gather those too and
+    key by module name (no EP axis to gather here, unlike :func:`save_ep_lora_adapter`).
+
+    Returns ``True`` if it handled the save (model has LoRA params), else ``False``.
+    """
+    from pathlib import Path
+
+    from peft.utils.save_and_load import get_peft_model_state_dict
+    from safetensors.torch import save_file
+
+    if not any("lora_" in n for n, _ in model.named_parameters()):
+        return False
+
+    active = model.active_adapter if hasattr(model, "active_adapter") else "default"
+    if isinstance(active, (list, tuple)):
+        active = active[0]
+
+    # Replicated + dp-sharded LoRA: full tensors via FSDP all-gather (collective — same iteration
+    # order on every rank). Canonical PEFT keys via get_peft_model_state_dict.
+    sd = {
+        name: (p.full_tensor() if type(p).__name__ == "DTensor" else p.data).detach()
+        for name, p in model.named_parameters()
+        if "lora_" in name
+    }
+    adapter_sd = get_peft_model_state_dict(model, state_dict=sd)
+
+    gathered = 0
+    for wname, wrapper in model.named_modules():
+        if not _is_param_wrapper(wrapper):
+            continue
+        for sub in ("lora_A", "lora_B"):
+            for w in (mod.weight for mod in getattr(wrapper, sub, {}).values()):
+                full = (
+                    w.full_tensor() if type(w).__name__ == "DTensor" else w.data
+                ).detach()
+                key = f"{wname}.{sub}.weight"
+                target = (
+                    key
+                    if key in adapter_sd
+                    else next(
+                        (
+                            k
+                            for k in adapter_sd
+                            if k.endswith(key.split("base_model.model.")[-1])
+                        ),
+                        None,
+                    )
+                )
+                if target is not None:
+                    adapter_sd[target] = full
+                    gathered += 1
+
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        out = Path(output_dir)
+        out.mkdir(parents=True, exist_ok=True)
+        save_file(
+            {k: v.cpu().contiguous() for k, v in adapter_sd.items()},
+            str(out / "adapter_model.safetensors"),
+        )
+        model.peft_config[active].save_pretrained(str(out))
+        LOG.info(
+            f"Saved FSDP2-gathered LoRA adapter ({len(adapter_sd)} tensors, "
+            f"{gathered} expert params) to {out}."
+        )
+    if dist.is_initialized():
+        dist.barrier()
+    return True
+
+
+def gather_expert_lora_full(local: torch.Tensor, kind: str, e_global: int, ep_group):
+    """Inverse of the EP LoRA slice: all-gather a local-experts LoRA tensor across the
+    EP group and reassemble the full ``e_global``-expert tensor in the PEFT layout.
+
+      * ``kind="A"`` (expert-major ``[E*r, in]``): concat gathered slices along rows.
+      * ``kind="B"`` (rank-major ``[out, r*E]``): place each rank's experts into the
+        ``E`` axis of ``[out, r, E]`` and flatten.
+    """
+    ep_size = dist.get_world_size(ep_group)
+    gathered = [torch.empty_like(local) for _ in range(ep_size)]
+    dist.all_gather(gathered, local.contiguous(), group=ep_group)
+    if kind == "A":
+        return torch.cat(gathered, dim=0)
+    out_dim = local.shape[0]
+    e_local = e_global // ep_size
+    r = local.shape[1] // e_local
+    parts = [g.reshape(out_dim, r, e_local) for g in gathered]
+    return torch.cat(parts, dim=2).reshape(out_dim, r * e_global)
diff --git a/src/axolotl/integrations/grokfast/LICENSE b/src/axolotl/integrations/grokfast/LICENSE
new file mode 100644
index 0000000000..21e35ee968
--- /dev/null
+++ b/src/axolotl/integrations/grokfast/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Jaerin Lee, Bong Gyun Kang, Kihoon Kim, Kyoung Mu Lee
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/src/axolotl/integrations/grokfast/README.md b/src/axolotl/integrations/grokfast/README.md
new file mode 100644
index 0000000000..7c678b07e7
--- /dev/null
+++ b/src/axolotl/integrations/grokfast/README.md
@@ -0,0 +1,24 @@
+# Grokfast Optimizer
+
+See https://github.com/ironjr/grokfast
+
+## Usage
+
+```yaml
+plugins:
+  - axolotl.integrations.grokfast.GrokfastPlugin
+
+grokfast_alpha: 2.0
+grokfast_lamb: 0.98
+```
+
+## Citation
+
+```bib
+@article{lee2024grokfast,
+    title={{Grokfast}: Accelerated Grokking by Amplifying Slow Gradients},
+    author={Lee, Jaerin and Kang, Bong Gyun and Kim, Kihoon and Lee, Kyoung Mu},
+    journal={arXiv preprint arXiv:2405.20233},
+    year={2024}
+}
+```
diff --git a/src/axolotl/integrations/grokfast/__init__.py b/src/axolotl/integrations/grokfast/__init__.py
new file mode 100644
index 0000000000..df8cf2cf39
--- /dev/null
+++ b/src/axolotl/integrations/grokfast/__init__.py
@@ -0,0 +1,49 @@
+"""
+Grokfast plugin for Axolotl
+"""
+
+from transformers.trainer_callback import TrainerCallback
+
+from axolotl.utils.logging import get_logger
+
+from ..base import BasePlugin
+from .args import GrokfastArgs as GrokfastArgs
+from .optimizer import gradfilter_ema
+
+LOG = get_logger(__name__)
+
+
+class GrokfastCallbackHandler(TrainerCallback):
+    """
+    Transformer trainer callbacks for Grokfast
+    """
+
+    def __init__(self, *args_, alpha=0.98, lamb=2.0, **kwargs):
+        super().__init__(*args_, **kwargs)
+        self.grads = None
+        self.alpha = alpha
+        self.lamb = lamb
+
+    def on_train_begin(self, *args_, **kwargs):
+        self.grads = None
+
+    def on_pre_optimizer_step(self, args_, state, control, **kwargs):
+        model = kwargs.pop("model")
+        self.grads = gradfilter_ema(model, self.grads, alpha=self.alpha, lamb=self.lamb)
+        return control
+
+
+class GrokfastPlugin(BasePlugin):
+    """
+    Plugin for Grokfast optimizer integraton with Axolotl.
+    """
+
+    def get_input_args(self):
+        return "axolotl.integrations.grokfast.GrokfastArgs"
+
+    def add_callbacks_post_trainer(self, cfg, trainer):
+        LOG.info("Adding Grokfast callback to the trainer")
+        callback = GrokfastCallbackHandler(
+            alpha=cfg.grokfast_alpha, lamb=cfg.grokfast_lamb
+        )
+        return [callback]
diff --git a/src/axolotl/integrations/grokfast/args.py b/src/axolotl/integrations/grokfast/args.py
new file mode 100644
index 0000000000..ac91c7395c
--- /dev/null
+++ b/src/axolotl/integrations/grokfast/args.py
@@ -0,0 +1,16 @@
+"""
+config args for grokfast plugin
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class GrokfastArgs(BaseModel):
+    """
+    Input args for Grokfast optimizer.
+    """
+
+    grokfast_alpha: Optional[float] = 0.98
+    grokfast_lamb: Optional[float] = 2.0
diff --git a/src/axolotl/integrations/grokfast/optimizer.py b/src/axolotl/integrations/grokfast/optimizer.py
new file mode 100644
index 0000000000..c83ef43bc1
--- /dev/null
+++ b/src/axolotl/integrations/grokfast/optimizer.py
@@ -0,0 +1,62 @@
+# Copyright: MIT License (c) 2024 Jaerin Lee, Bong Gyun Kang, Kihoon Kim, Kyoung Mu Lee
+# Reference: https://github.com/ironjr/grokfast
+
+from collections import deque
+from typing import Dict, Literal, Optional
+
+import torch
+import torch.nn as nn
+
+
+def gradfilter_ma(
+    m: nn.Module,
+    grads: Optional[Dict[str, deque]] = None,
+    window_size: int = 100,
+    lamb: float = 5.0,
+    filter_type: Literal["mean", "sum"] = "mean",
+    warmup: bool = True,
+    trigger: bool = False,  # For ablation study.
+) -> Dict[str, deque]:
+    if grads is None:
+        grads = {
+            n: deque(maxlen=window_size)
+            for n, p in m.named_parameters()
+            if p.requires_grad and p.grad is not None
+        }
+
+    for n, p in m.named_parameters():
+        if p.requires_grad and p.grad is not None:
+            grads[n].append(p.grad.data.detach())  # .cpu())
+
+            # Modify the gradients.
+            if not warmup or len(grads[n]) == window_size and not trigger:
+                if filter_type == "mean":
+                    avg = sum(grads[n]) / len(grads[n])
+                elif filter_type == "sum":
+                    avg = sum(grads[n])
+                else:
+                    raise ValueError(f"Unrecognized filter_type {filter_type}")
+                p.grad.data = p.grad.data + avg * lamb
+
+    return grads
+
+
+def gradfilter_ema(
+    m: nn.Module,
+    grads: Optional[Dict[str, torch.Tensor]] = None,
+    alpha: float = 0.98,
+    lamb: float = 2.0,
+) -> Dict[str, torch.Tensor]:
+    if grads is None:
+        grads = {
+            n: p.grad.data.detach()
+            for n, p in m.named_parameters()
+            if p.requires_grad and p.grad is not None
+        }
+
+    for n, p in m.named_parameters():
+        if p.requires_grad and p.grad is not None:
+            grads[n] = grads[n] * alpha + p.grad.data.detach() * (1 - alpha)
+            p.grad.data = p.grad.data + grads[n] * lamb
+
+    return grads
diff --git a/src/axolotl/integrations/hatchery/__init__.py b/src/axolotl/integrations/hatchery/__init__.py
new file mode 100644
index 0000000000..c0d8510dbc
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/__init__.py
@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Hatchery/Tinker remote training integration for Axolotl.
+
+Routes axolotl's preprocessed data to a remote training API (Tinker or
+Hatchery) instead of running forward/backward locally. The remote
+service handles model weights, LoRA adapters, and gradient updates.
+"""
+
+from .args import HatcheryArgs, HatcheryConfig
+from .plugin import HatcheryPlugin
+
+__all__ = ["HatcheryArgs", "HatcheryConfig", "HatcheryPlugin"]
+
+# Usage:
+#   plugins:
+#     - axolotl.integrations.hatchery.HatcheryPlugin
+#
+#   hatchery:
+#     backend: tinker  # or "hatchery"
+#     lora_rank: 32
+#     loss_fn: cross_entropy  # SFT
+#     # loss_fn: ppo         # RL (auto-selects HatcheryRLTrainer)
+#
+#   learning_rate: 1e-4  # top-level, not under hatchery:
diff --git a/src/axolotl/integrations/hatchery/args.py b/src/axolotl/integrations/hatchery/args.py
new file mode 100644
index 0000000000..e3fdb95a24
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/args.py
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Pydantic config schema for the Hatchery integration."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, Optional
+
+from pydantic import BaseModel, Field
+
+
+class HatcheryConfig(BaseModel):
+    """Nested config under `hatchery:` in the axolotl YAML.
+
+    Only contains hatchery-specific settings. Standard training params
+    (learning_rate, weight_decay, adam_beta1/2, max_grad_norm,
+    gradient_accumulation_steps) are read from axolotl's top-level config.
+    """
+
+    # Backend & connection
+    backend: Literal["tinker", "hatchery"] = "tinker"
+    base_url: Optional[str] = None
+    api_key: Optional[str] = None
+    project_id: Optional[str] = None
+
+    # LoRA config sent to remote
+    lora_rank: int = Field(32, ge=1, le=256)
+    train_attn: bool = True
+    train_mlp: bool = True
+    train_unembed: bool = True
+
+    # Loss function
+    loss_fn: Literal["cross_entropy", "importance_sampling", "ppo", "cispo", "dro"] = (
+        "cross_entropy"
+    )
+    loss_fn_config: Optional[dict[str, Any]] = None
+
+    # Pipelining: submit next batch before awaiting previous result
+    pipeline: bool = True
+
+    # Sampling params (for RL flows)
+    max_sample_tokens: int = 256
+    sample_temperature: float = 1.0
+    num_samples: int = 4
+
+    # Reward functions (for RL) — list of fully qualified names
+    reward_funcs: Optional[list[str]] = None
+
+    # Checkpointing
+    save_steps: Optional[int] = None
+    save_name_prefix: str = "checkpoint"
+
+    # Timeout per future (seconds)
+    future_timeout: float = 600.0
+
+
+class HatcheryArgs(BaseModel):
+    """Top-level mixin that adds the nested `hatchery:` field."""
+
+    hatchery: Optional[HatcheryConfig] = None
diff --git a/src/axolotl/integrations/hatchery/data.py b/src/axolotl/integrations/hatchery/data.py
new file mode 100644
index 0000000000..f7baa2cca4
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/data.py
@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Convert axolotl batch tensors to Tinker/Hatchery Datum format.
+
+Both Tinker and Hatchery expect the client to apply the causal LM shift:
+
+  Original tokens:  [t0, t1, t2, ..., t_{L-1}]
+  model_input:      [t0, t1, ..., t_{L-2}]       (last token dropped)
+  target_tokens:    [t1, t2, ..., t_{L-1}]        (first token dropped)
+  weights:          [w1, w2, ..., w_{L-1}]        (aligned to targets)
+
+At position i, the model sees t_i and predicts target_tokens[i] = t_{i+1}.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def _tensor_to_wire(t: torch.Tensor) -> dict[str, Any]:
+    """Serialize a tensor to the TensorData wire dict."""
+    flat = t.detach().cpu().flatten()
+    dtype_map = {
+        torch.float32: "float32",
+        torch.float16: "float16",
+        torch.bfloat16: "bfloat16",
+        torch.int64: "int64",
+        torch.int32: "int32",
+    }
+    return {
+        "dtype": dtype_map.get(flat.dtype, "float32"),
+        "shape": list(t.shape),
+        "data": flat.tolist(),
+    }
+
+
+def _make_datum(
+    tokens: list[int],
+    loss_fn_inputs: dict[str, torch.Tensor],
+) -> dict[str, Any]:
+    """Build a Datum as a plain dict (wire-compatible with both Tinker and Hatchery)."""
+    return {
+        "model_input": {
+            "chunks": [{"type": "encoded_text", "tokens": tokens}],
+        },
+        "loss_fn_inputs": {
+            key: _tensor_to_wire(tensor) for key, tensor in loss_fn_inputs.items()
+        },
+    }
+
+
+def datums_to_tinker(datums: list[dict[str, Any]]):
+    """Wrap plain-dict datums into tinker.types.Datum objects.
+
+    Both the Tinker SDK and updated Hatchery client accept these.
+    """
+    import tinker.types as tt
+
+    result = []
+    for d in datums:
+        tokens = d["model_input"]["chunks"][0]["tokens"]
+        tinker_inputs = {}
+        for key, wire in d["loss_fn_inputs"].items():
+            tinker_inputs[key] = tt.TensorData(
+                data=wire["data"],
+                dtype=wire["dtype"],
+                shape=wire["shape"],
+            )
+        result.append(
+            tt.Datum(
+                model_input=tt.ModelInput.from_ints(tokens),
+                loss_fn_inputs=tinker_inputs,
+            )
+        )
+    return result
+
+
+def batch_to_datums_sft(
+    input_ids: torch.Tensor,
+    labels: torch.Tensor,
+    attention_mask: torch.Tensor | None = None,
+) -> list[dict[str, Any]]:
+    """Convert an axolotl SFT batch to Datum dicts with causal shift."""
+    batch_size = input_ids.size(0)
+    datums = []
+
+    for i in range(batch_size):
+        ids = input_ids[i]
+        lbl = labels[i]
+
+        if attention_mask is not None:
+            seq_len = int(attention_mask[i].sum().item())
+            ids = ids[:seq_len]
+            lbl = lbl[:seq_len]
+
+        model_tokens = ids[:-1].tolist()
+        shifted_labels = lbl[1:]
+
+        target_tokens = shifted_labels.clone()
+        weights = (shifted_labels != -100).float()
+        target_tokens[target_tokens == -100] = 0
+
+        datums.append(
+            _make_datum(
+                model_tokens,
+                {
+                    "target_tokens": target_tokens,
+                    "weights": weights,
+                },
+            )
+        )
+
+    return datums
+
+
+def batch_to_datums_rl(
+    input_ids: torch.Tensor,
+    labels: torch.Tensor,
+    logprobs: torch.Tensor,
+    advantages: torch.Tensor,
+    attention_mask: torch.Tensor | None = None,
+) -> list[dict[str, Any]]:
+    """Convert an RL batch to importance_sampling/ppo Datum dicts with causal shift."""
+    batch_size = input_ids.size(0)
+    datums = []
+
+    for i in range(batch_size):
+        ids = input_ids[i]
+        lbl = labels[i]
+
+        if attention_mask is not None:
+            seq_len = int(attention_mask[i].sum().item())
+        else:
+            seq_len = ids.size(0)
+        ids = ids[:seq_len]
+        lbl = lbl[:seq_len]
+        lp = logprobs[i, :seq_len]
+        adv = advantages[i, :seq_len]
+
+        model_tokens = ids[:-1].tolist()
+
+        target_tokens = lbl[1:].clone()
+        target_tokens[target_tokens == -100] = 0
+
+        datums.append(
+            _make_datum(
+                model_tokens,
+                {
+                    "target_tokens": target_tokens,
+                    "logprobs": lp[1:],
+                    "advantages": adv[1:],
+                },
+            )
+        )
+
+    return datums
diff --git a/src/axolotl/integrations/hatchery/examples/prep_math_rl.py b/src/axolotl/integrations/hatchery/examples/prep_math_rl.py
new file mode 100644
index 0000000000..1838159075
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/examples/prep_math_rl.py
@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Prepare hendrycks_math for RL training with Hatchery/Tinker.
+
+Creates a dataset with chat-formatted prompts that include
+a hidden gold answer tag for the reward function.
+
+Run:
+  python src/axolotl/integrations/hatchery/examples/prep_math_rl.py
+"""
+
+import os
+import re
+
+from datasets import Dataset, load_dataset
+from transformers import AutoTokenizer
+
+
+def extract_boxed(text: str) -> str:
+    match = re.search(r"\\boxed\{", text)
+    if not match:
+        return ""
+    start = match.end()
+    depth = 1
+    i = start
+    while i < len(text) and depth > 0:
+        if text[i] == "{":
+            depth += 1
+        elif text[i] == "}":
+            depth -= 1
+        i += 1
+    return text[start : i - 1] if depth == 0 else ""
+
+
+def main():
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B", trust_remote_code=True)
+
+    ds = load_dataset("EleutherAI/hendrycks_math", "algebra", split="test")
+    level = os.environ.get("MATH_LEVEL", "Level 1")
+    filtered_rows = [x for x in ds if x["level"] == level]
+    print(f"{level} algebra: {len(filtered_rows)} problems")
+
+    rows = []
+    for prob in filtered_rows:
+        gold = extract_boxed(prob["solution"])
+        if not gold:
+            continue
+
+        # Format as chat prompt with hidden gold tag
+        prompt = (
+            f"Solve the following math problem. "
+            f"Show your work and put your final answer in \\boxed{{}}.\n\n"
+            f"{prob['problem']}"
+            f"<|gold|>{gold}<|/gold|>"
+        )
+
+        # Tokenize the prompt
+        text = tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        prompt_ids = tokenizer.encode(text, add_special_tokens=False)
+
+        rows.append(
+            {
+                "input_ids": prompt_ids,
+                "labels": [-100] * len(prompt_ids),
+                "attention_mask": [1] * len(prompt_ids),
+            }
+        )
+
+    out = Dataset.from_list(rows)
+    out_dir = f"./data/math_rl_{level.lower().replace(' ', '')}"
+    out.save_to_disk(out_dir)
+    print(f"Saved {len(out)} examples to {out_dir}")
+    if rows:
+        print(
+            f"Prompt length range: {min(len(r['input_ids']) for r in rows)}"
+            f"-{max(len(r['input_ids']) for r in rows)}"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/axolotl/integrations/hatchery/examples/tinker_rl.yaml b/src/axolotl/integrations/hatchery/examples/tinker_rl.yaml
new file mode 100644
index 0000000000..caab3fe0d8
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/examples/tinker_rl.yaml
@@ -0,0 +1,47 @@
+# RL (GRPO): hendrycks_math Level 1 via Tinker with Qwen3-8B
+#
+# Prep:
+#   python src/axolotl/integrations/hatchery/examples/prep_math_rl.py
+#
+# Run:
+#   export TINKER_API_KEY="your-key"
+#   axolotl train src/axolotl/integrations/hatchery/examples/tinker_rl.yaml
+
+base_model: Qwen/Qwen3-8B
+
+plugins:
+  - axolotl.integrations.hatchery.HatcheryPlugin
+
+hatchery:
+  backend: tinker
+  lora_rank: 16
+  loss_fn: importance_sampling
+  max_sample_tokens: 2048
+  sample_temperature: 0.7
+  num_samples: 4
+  pipeline: true
+  save_steps: 5
+  reward_funcs:
+    - axolotl.integrations.hatchery.rewards.math_reward.math_reward
+
+datasets:
+  - path: ./data/math_rl_level1
+    ds_type: arrow
+    type: completion
+
+sequence_len: 2048
+
+learning_rate: 5.0e-5
+optimizer: adamw_torch
+adam_beta1: 0.9
+adam_beta2: 0.95
+weight_decay: 0.01
+max_grad_norm: 1.0
+
+max_steps: 10
+num_epochs: 1
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+logging_steps: 1
+
+output_dir: ./outputs/tinker-rl-math
diff --git a/src/axolotl/integrations/hatchery/examples/tinker_sft.yaml b/src/axolotl/integrations/hatchery/examples/tinker_sft.yaml
new file mode 100644
index 0000000000..d99f043ae1
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/examples/tinker_sft.yaml
@@ -0,0 +1,42 @@
+# SFT: KIMI-K2 thinking data via Tinker remote API with Qwen3-8B
+#
+# Usage:
+#   export TINKER_API_KEY="your-key"
+#   axolotl train src/axolotl/integrations/hatchery/examples/tinker_sft.yaml
+
+base_model: Qwen/Qwen3-8B
+
+plugins:
+  - axolotl.integrations.hatchery.HatcheryPlugin
+
+hatchery:
+  backend: tinker
+  lora_rank: 16
+  loss_fn: cross_entropy
+  pipeline: true
+  save_steps: 10
+
+datasets:
+  - path: TeichAI/kimi-k2-thinking-1000x
+    split: train[:50]
+    type: chat_template
+    chat_template: qwen3
+    split_thinking: true
+
+chat_template: qwen3
+sequence_len: 2048
+
+learning_rate: 3.0e-4
+optimizer: adamw_torch
+adam_beta1: 0.9
+adam_beta2: 0.95
+weight_decay: 0.01
+max_grad_norm: 1.0
+
+num_epochs: 1
+max_steps: 20
+micro_batch_size: 2
+gradient_accumulation_steps: 1
+logging_steps: 1
+
+output_dir: ./outputs/tinker-sft
diff --git a/src/axolotl/integrations/hatchery/plugin.py b/src/axolotl/integrations/hatchery/plugin.py
new file mode 100644
index 0000000000..1546958e80
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/plugin.py
@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Axolotl plugin that routes training to a remote Hatchery/Tinker API."""
+
+from __future__ import annotations
+
+import torch
+from peft import PeftModel
+from transformers import AutoConfig, PreTrainedModel, Trainer
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class HatcheryPlugin(BasePlugin):
+    """Plugin that replaces local training with remote API calls.
+
+    Activated by adding to the axolotl YAML:
+
+        plugins:
+          - axolotl.integrations.hatchery.HatcheryPlugin
+
+        hatchery:
+          backend: tinker  # or "hatchery"
+          lora_rank: 32
+          loss_fn: cross_entropy
+          # ... see HatcheryConfig for full options
+    """
+
+    def get_input_args(self) -> str:
+        return "axolotl.integrations.hatchery.args.HatcheryArgs"
+
+    def register(self, cfg: dict):
+        """Auto-set config values needed for remote training."""
+        if cfg.get("remove_unused_columns") is None:
+            cfg["remove_unused_columns"] = False
+
+    def pre_model_load(self, cfg: DictDefault):
+        """Replace model loading with a tiny stub."""
+        hcfg = cfg.hatchery or {}
+        backend = (
+            hcfg.get("backend", "tinker")
+            if isinstance(hcfg, dict)
+            else getattr(hcfg, "backend", "tinker")
+        )
+        LOG.info(
+            f"Hatchery plugin active: training dispatched to remote "
+            f"{backend} API. Skipping local model weight loading."
+        )
+
+        from axolotl.loaders import ModelLoader
+
+        def _stub_build_model(loader_self) -> bool:
+            base_model = loader_self.cfg.base_model
+            LOG.info(f"Skipping model weight loading for: {base_model}")
+
+            config = AutoConfig.from_pretrained(
+                base_model,
+                trust_remote_code=loader_self.cfg.get("trust_remote_code", False),
+            )
+
+            class _Stub(PreTrainedModel):
+                config_class = type(config)
+                _no_split_modules: list[str] = []
+                supports_gradient_checkpointing = False
+
+                def __init__(self, cfg):
+                    super().__init__(cfg)
+                    vocab_size = getattr(cfg, "vocab_size", 32000)
+                    self.embed_tokens = torch.nn.Embedding(vocab_size, 1)
+
+                def get_input_embeddings(self):
+                    return self.embed_tokens
+
+                def set_input_embeddings(self, value):
+                    pass
+
+                def get_output_embeddings(self):
+                    return None
+
+            loader_self.model = _Stub(config)
+            return True
+
+        ModelLoader._build_model = _stub_build_model  # type: ignore[method-assign,assignment]
+
+    def get_trainer_cls(self, cfg: DictDefault) -> type[Trainer] | None:
+        """Return the appropriate remote trainer class."""
+        hcfg = cfg.hatchery
+        loss_fn = getattr(hcfg, "loss_fn", "cross_entropy") if hcfg else "cross_entropy"
+
+        if loss_fn in ("importance_sampling", "ppo", "cispo", "dro"):
+            from .rl_trainer import HatcheryRLTrainer
+
+            return HatcheryRLTrainer
+
+        from .trainer import HatcheryTrainer
+
+        return HatcheryTrainer
+
+    def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        model._hatchery_remote = True
+
+    def post_train(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        LOG.info(
+            "Hatchery: skipping local model save (weights are on remote API). "
+            "Use `tinker checkpoint download` or hatchery CLI to retrieve."
+        )
+
+    def post_trainer_create(self, cfg: DictDefault, trainer: Trainer):
+        """Inject hatchery config + axolotl training params into the trainer."""
+        from .args import HatcheryConfig
+        from .rl_trainer import HatcheryRLTrainer
+        from .trainer import HatcheryTrainer
+
+        if not isinstance(trainer, (HatcheryTrainer, HatcheryRLTrainer)):
+            return
+
+        hcfg = cfg.hatchery
+        if isinstance(hcfg, dict):
+            hatchery_config = HatcheryConfig(**hcfg)
+        elif hcfg is None:
+            hatchery_config = HatcheryConfig()
+        else:
+            hatchery_config = hcfg
+
+        trainer.hatchery_args = hatchery_config
+        trainer._base_model_name = cfg.base_model
+
+        # Pull standard training params from axolotl config so they
+        # don't need to be duplicated under hatchery:
+        trainer._optim_params = {
+            "learning_rate": cfg.learning_rate
+            if cfg.learning_rate is not None
+            else 1e-4,
+            "beta1": cfg.adam_beta1 if cfg.adam_beta1 is not None else 0.9,
+            "beta2": cfg.adam_beta2 if cfg.adam_beta2 is not None else 0.95,
+            "eps": cfg.adam_epsilon if cfg.adam_epsilon is not None else 1e-12,
+            "weight_decay": cfg.weight_decay if cfg.weight_decay is not None else 0.0,
+            "grad_clip_norm": cfg.max_grad_norm
+            if cfg.max_grad_norm is not None
+            else 0.0,
+        }
diff --git a/src/axolotl/integrations/hatchery/rewards/__init__.py b/src/axolotl/integrations/hatchery/rewards/__init__.py
new file mode 100644
index 0000000000..1cfe76e777
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/rewards/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
diff --git a/src/axolotl/integrations/hatchery/rewards/math_reward.py b/src/axolotl/integrations/hatchery/rewards/math_reward.py
new file mode 100644
index 0000000000..a0e71b7c90
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/rewards/math_reward.py
@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Math reward function for hendrycks_math GRPO training.
+
+Uses math_verify for robust answer comparison. Falls back to
+exact string match of \\boxed{} content only when math_verify
+is unavailable.
+"""
+
+from __future__ import annotations
+
+import re
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def extract_boxed(text: str) -> str | None:
+    """Extract \\boxed{...} answer handling nested braces."""
+    match = re.search(r"\\boxed\{", text)
+    if not match:
+        return None
+    start = match.end()
+    depth = 1
+    i = start
+    while i < len(text) and depth > 0:
+        if text[i] == "{":
+            depth += 1
+        elif text[i] == "}":
+            depth -= 1
+        i += 1
+    return text[start : i - 1] if depth == 0 else None
+
+
+def math_reward(prompts: list[str], completions: list[str], **kwargs) -> list[float]:
+    """Score completions by checking if \\boxed{} answer matches the gold answer.
+
+    The gold answer is extracted from the prompt (appended as a hidden
+    tag by the dataset preprocessing). Format:
+      ... <|gold|>ANSWER<|/gold|>
+    """
+    rewards = []
+    for prompt, completion in zip(prompts, completions, strict=True):
+        gold_match = re.search(r"<\|gold\|>(.*?)<\|/gold\|>", prompt)
+        if not gold_match:
+            rewards.append(0.0)
+            continue
+
+        gold_answer = gold_match.group(1).strip()
+        pred_answer = extract_boxed(completion)
+
+        if pred_answer is None:
+            rewards.append(0.0)
+            continue
+
+        verified = None
+        try:
+            from math_verify import parse, verify
+
+            gold_parsed = parse(gold_answer)
+            pred_parsed = parse(pred_answer)
+            verified = verify(gold_parsed, pred_parsed)
+        except Exception:
+            LOG.debug(
+                "math_verify unavailable or failed, using string fallback",
+                exc_info=True,
+            )
+
+        if verified is not None:
+            rewards.append(1.0 if verified else 0.0)
+        elif pred_answer.strip() == gold_answer.strip():
+            rewards.append(1.0)
+        else:
+            rewards.append(0.0)
+
+    return rewards
diff --git a/src/axolotl/integrations/hatchery/rl_trainer.py b/src/axolotl/integrations/hatchery/rl_trainer.py
new file mode 100644
index 0000000000..fc6d32d6a6
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/rl_trainer.py
@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Remote RL trainer (GRPO/PPO) using Tinker or Hatchery API.
+
+Full RL loop per step:
+  1. Extract prompts from dataset batch
+  2. Sample N completions per prompt via remote SamplingClient
+  3. Score completions with local reward functions
+  4. Compute GRPO-style advantages (per-group normalization)
+  5. Send (prompt+completion, logprobs, advantages) as forward_backward
+  6. Optimizer step
+"""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+import re
+import time
+from typing import Any, Callable, Optional
+
+import torch
+from transformers.trainer_utils import TrainOutput
+
+from axolotl.core.trainers.base import AxolotlTrainer
+from axolotl.utils.logging import get_logger
+
+from .args import HatcheryConfig
+from .data import batch_to_datums_rl, datums_to_tinker
+from .trainer import _create_training_client
+
+LOG = get_logger(__name__)
+
+
+def _load_reward_func(fqn: str) -> Callable:
+    """Load a reward function from a fully qualified name like 'module.func'."""
+    module_path = ".".join(fqn.split(".")[:-1])
+    func_name = fqn.split(".")[-1]
+    mod = importlib.import_module(module_path)
+    func = getattr(mod, func_name)
+    if len(inspect.signature(func).parameters) < 2:
+        raise ValueError(f"Reward function {fqn} must accept (prompts, completions)")
+    return func
+
+
+class HatcheryRLTrainer(AxolotlTrainer):
+    """Remote RL trainer using Tinker/Hatchery for sampling and training."""
+
+    hatchery_args: Optional[HatcheryConfig]
+    _base_model_name: Optional[str]
+    _training_client: Any
+    _reward_functions: list[Callable]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.hatchery_args = None
+        self._base_model_name = None
+        self._training_client = None
+        self._reward_functions = []
+
+    def _ensure_reward_functions(self):
+        if self._reward_functions:
+            return
+        args = self.hatchery_args
+        if not args or not args.reward_funcs:
+            raise ValueError(
+                "No reward functions configured. Set hatchery.reward_funcs "
+                "in YAML, e.g. reward_funcs: ['my_module.my_reward']"
+            )
+        for fqn in args.reward_funcs:
+            self._reward_functions.append(_load_reward_func(fqn))
+        LOG.info(f"Loaded {len(self._reward_functions)} reward function(s)")
+
+    def _get_training_client(self):
+        if self._training_client is not None:
+            return self._training_client
+
+        self._training_client = _create_training_client(
+            self.hatchery_args, self._base_model_name
+        )
+        LOG.info(
+            f"Remote RL session created: backend={self.hatchery_args.backend}, "
+            f"model={self._base_model_name}, rank={self.hatchery_args.lora_rank}"
+        )
+        return self._training_client
+
+    def _sample_completions(self, prompt_ids_list: list[list[int]]):
+        """Sample completions for prompts via remote API."""
+        import tinker.types as tt
+
+        tc = self._get_training_client()
+        args = self.hatchery_args
+        assert args is not None  # validated by _get_training_client
+        results = []
+
+        sc = tc.save_weights_and_get_sampling_client()
+
+        for prompt_ids in prompt_ids_list:
+            if hasattr(sc, "sampling_session_id"):
+                sample_result = sc.sample(
+                    prompt_ids,
+                    max_tokens=args.max_sample_tokens,
+                    temperature=args.sample_temperature,
+                    n=args.num_samples,
+                ).result(timeout=args.future_timeout)
+            else:
+                mi = tt.ModelInput.from_ints(prompt_ids)
+                sp = tt.SamplingParams(
+                    max_tokens=args.max_sample_tokens,
+                    temperature=args.sample_temperature,
+                    top_p=0.95,
+                    top_k=-1,
+                )
+                sample_result = sc.sample(
+                    prompt=mi,
+                    num_samples=args.num_samples,
+                    sampling_params=sp,
+                ).result(timeout=args.future_timeout)
+
+            sequences = (
+                sample_result.sequences
+                if hasattr(sample_result, "sequences")
+                else sample_result.get("sequences", [])
+            )
+            for seq in sequences:
+                tokens = (
+                    list(seq.tokens)
+                    if hasattr(seq, "tokens")
+                    else seq.get("tokens", [])
+                )
+                logprobs = (
+                    list(seq.logprobs)
+                    if hasattr(seq, "logprobs") and seq.logprobs
+                    else seq.get("logprobs", [])
+                )
+                results.append(
+                    {
+                        "tokens": list(prompt_ids) + tokens,
+                        "completion_tokens": tokens,
+                        "logprobs": logprobs,
+                        "prompt_len": len(prompt_ids),
+                    }
+                )
+
+        return results
+
+    def _compute_rewards(
+        self, prompts: list[str], completions: list[str]
+    ) -> list[float]:
+        total_rewards = [0.0] * len(completions)
+        for reward_fn in self._reward_functions:
+            rewards = reward_fn(prompts, completions)
+            for i, r in enumerate(rewards):
+                total_rewards[i] += r
+        return total_rewards
+
+    @staticmethod
+    def _compute_advantages(rewards: list[float], group_size: int) -> list[float]:
+        advantages = []
+        for i in range(0, len(rewards), group_size):
+            group = rewards[i : i + group_size]
+            mean = sum(group) / len(group)
+            var = sum((r - mean) ** 2 for r in group) / max(len(group), 1)
+            std = var**0.5 if var > 1e-8 else 1.0
+            advantages.extend([(r - mean) / std for r in group])
+        return advantages
+
+    def _do_optim_step(self):
+        import tinker.types as tt
+
+        tc = self._get_training_client()
+        return tc.optim_step(tt.AdamParams(**self._optim_params))
+
+    def train(
+        self,
+        resume_from_checkpoint: Optional[str] = None,
+        trial: Any = None,
+        ignore_keys_for_eval: Optional[list[str]] = None,
+        **kwargs,
+    ) -> TrainOutput:
+        args = self.hatchery_args
+        if args is None:
+            raise RuntimeError("hatchery_args not configured")
+
+        self._ensure_reward_functions()
+
+        train_dataloader = self.get_train_dataloader()
+        num_train_epochs = int(self.args.num_train_epochs)
+        max_steps = self.args.max_steps if self.args.max_steps > 0 else 1000
+
+        LOG.info(
+            f"Remote RL training: max_steps={max_steps}, "
+            f"loss_fn={args.loss_fn}, samples/prompt={args.num_samples}"
+        )
+
+        self.state.max_steps = max_steps
+        self.state.num_train_epochs = num_train_epochs
+        self.state.is_local_process_zero = True
+        self.state.is_world_process_zero = True
+
+        self.control = self.callback_handler.on_train_begin(
+            self.args,
+            self.state,
+            self.control,  # type: ignore[has-type]
+        )
+
+        tokenizer = self.processing_class
+        global_step = 0
+        total_loss = 0.0
+        total_reward = 0.0
+        start_time = time.time()
+
+        for _epoch in range(num_train_epochs):
+            if global_step >= max_steps:
+                break
+
+            for batch in train_dataloader:
+                if global_step >= max_steps:
+                    break
+
+                self.control = self.callback_handler.on_step_begin(
+                    self.args, self.state, self.control
+                )
+
+                prompt_ids_batch = batch["input_ids"]
+                # Full prompt text (with gold tag) for reward scoring
+                prompt_texts = tokenizer.batch_decode(
+                    prompt_ids_batch, skip_special_tokens=False
+                )
+
+                # Strip <|gold|>...<|/gold|> from token ids before
+                # sending to the model for sampling — the gold answer
+                # must only be visible to the local reward function.
+                sampling_prompts = []
+                for prompt_text in prompt_texts:
+                    clean = re.sub(r"<\|gold\|>.*?<\|/gold\|>", "", prompt_text)
+                    clean_ids = tokenizer.encode(clean, add_special_tokens=False)
+                    sampling_prompts.append(clean_ids)
+
+                # 1. Sample completions (without gold answer)
+                t0 = time.time()
+                samples = self._sample_completions(sampling_prompts)
+                t_sample = time.time() - t0
+
+                if not samples:
+                    LOG.warning("No samples generated, skipping step")
+                    continue
+                LOG.info(
+                    f"Sampled {len(samples)} completions, "
+                    f"avg_len={sum(len(s['completion_tokens']) for s in samples) / len(samples):.0f}tok"
+                )
+
+                # 2. Decode and score
+                completion_texts = [
+                    tokenizer.decode(s["completion_tokens"], skip_special_tokens=False)
+                    for s in samples
+                ]
+                sample_prompts = []
+                for prompt_text in prompt_texts:
+                    sample_prompts.extend([prompt_text] * args.num_samples)
+
+                rewards = self._compute_rewards(sample_prompts, completion_texts)
+
+                # 3. GRPO advantages
+                advantages_list = self._compute_advantages(
+                    rewards, group_size=args.num_samples
+                )
+
+                # 4. Build training data
+                all_datums = []
+                for i, sample in enumerate(samples):
+                    full_tokens = sample["tokens"]
+                    prompt_len = sample["prompt_len"]
+                    seq_len = len(full_tokens)
+
+                    input_ids = torch.tensor([full_tokens], dtype=torch.long)
+                    labels = torch.full((1, seq_len), -100, dtype=torch.long)
+                    labels[0, prompt_len:] = torch.tensor(full_tokens[prompt_len:])
+
+                    logprobs_t = torch.zeros(1, seq_len)
+                    if sample["logprobs"]:
+                        lp = sample["logprobs"][: seq_len - prompt_len]
+                        logprobs_t[0, prompt_len : prompt_len + len(lp)] = torch.tensor(
+                            lp
+                        )
+
+                    adv_t = torch.zeros(1, seq_len)
+                    adv_t[0, prompt_len:] = advantages_list[i]
+
+                    all_datums.extend(
+                        batch_to_datums_rl(input_ids, labels, logprobs_t, adv_t)
+                    )
+
+                # 5. Forward backward (one datum at a time for memory) + optim
+                t0 = time.time()
+                tc = self._get_training_client()
+                step_loss = 0.0
+                for datum in all_datums:
+                    fb_future = tc.forward_backward(
+                        datums_to_tinker([datum]),
+                        loss_fn=args.loss_fn,
+                        loss_fn_config=args.loss_fn_config,
+                    )
+                    fb_result = fb_future.result(timeout=args.future_timeout)
+                    if hasattr(fb_result, "metrics"):
+                        step_loss += float(
+                            (fb_result.metrics or {}).get("loss:sum", 0.0)
+                        )
+                    elif isinstance(fb_result, dict):
+                        step_loss += float(
+                            fb_result.get("metrics", {}).get("loss:sum", 0.0)
+                        )
+                optim_future = self._do_optim_step()
+                if not args.pipeline:
+                    optim_future.result(timeout=args.future_timeout)
+                t_train = time.time() - t0
+
+                mean_reward = sum(rewards) / len(rewards)
+                accuracy = sum(1 for r in rewards if r > 0) / len(rewards)
+                mean_adv = sum(abs(a) for a in advantages_list) / len(advantages_list)
+                global_step += 1
+                total_loss += step_loss
+                total_reward += mean_reward
+                self.state.global_step = global_step
+
+                log_interval = self.args.logging_steps or 1
+                if global_step % log_interval == 0:
+                    elapsed = time.time() - start_time
+                    LOG.info(
+                        f"[step {global_step}/{max_steps}] "
+                        f"acc={accuracy:.2f} reward={mean_reward:.3f} "
+                        f"|adv|={mean_adv:.3f} loss:sum={step_loss:.1f} "
+                        f"sample={t_sample:.1f}s train={t_train:.1f}s "
+                        f"{elapsed / global_step:.1f}s/step"
+                    )
+                    self.log(
+                        {
+                            "loss": step_loss,
+                            "reward": mean_reward,
+                            "accuracy": accuracy,
+                            "mean_abs_advantage": mean_adv,
+                            "learning_rate": self._optim_params["learning_rate"],
+                        }
+                    )
+
+                if args.save_steps and global_step % args.save_steps == 0:
+                    self._save_remote_checkpoint(global_step)
+
+                self.control = self.callback_handler.on_step_end(
+                    self.args, self.state, self.control
+                )
+                if self.control.should_training_stop:
+                    break
+
+            if self.control.should_training_stop:
+                break
+
+        if global_step > 0:
+            self._save_remote_checkpoint(global_step, name="final")
+
+        elapsed = time.time() - start_time
+        avg_loss = total_loss / max(global_step, 1)
+        avg_reward = total_reward / max(global_step, 1)
+
+        LOG.info(
+            f"RL training complete: {global_step} steps, {elapsed:.1f}s, "
+            f"avg_reward={avg_reward:.4f}"
+        )
+
+        self.control = self.callback_handler.on_train_end(
+            self.args, self.state, self.control
+        )
+
+        return TrainOutput(
+            global_step=global_step,
+            training_loss=avg_loss,
+            metrics={
+                "train_loss": avg_loss,
+                "train_reward": avg_reward,
+                "train_runtime": elapsed,
+            },
+        )
+
+    def _save_remote_checkpoint(self, step: int, name: Optional[str] = None):
+        tc = self._get_training_client()
+        args = self.hatchery_args
+        assert args is not None  # validated by _get_training_client
+        ckpt_name = name or f"{args.save_name_prefix}-{step:06d}"
+        try:
+            future = tc.save_state(ckpt_name)
+            future.result(timeout=args.future_timeout)
+            LOG.info(f"Remote checkpoint saved: {ckpt_name}")
+        except Exception:
+            LOG.exception(f"Failed to save checkpoint {ckpt_name}")
+            if name == "final":
+                raise
+
+    def save_model(self, output_dir=None, _internal_call=False):
+        self._save_remote_checkpoint(
+            step=self.state.global_step,
+            name=output_dir or "hf-save",
+        )
+
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        raise NotImplementedError(
+            "HatcheryRLTrainer uses remote API; compute_loss not called locally."
+        )
diff --git a/src/axolotl/integrations/hatchery/trainer.py b/src/axolotl/integrations/hatchery/trainer.py
new file mode 100644
index 0000000000..4eb632db37
--- /dev/null
+++ b/src/axolotl/integrations/hatchery/trainer.py
@@ -0,0 +1,327 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Remote trainer that dispatches to Tinker or Hatchery API."""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Any, Optional
+
+import torch
+from transformers.trainer_utils import TrainOutput
+
+from axolotl.core.trainers.base import AxolotlTrainer
+from axolotl.utils.logging import get_logger
+
+from .args import HatcheryConfig
+from .data import batch_to_datums_sft, datums_to_tinker
+
+LOG = get_logger(__name__)
+
+
+def _extract_loss(result) -> float:
+    """Extract loss:sum from a forward_backward result.
+
+    Tinker's cross_entropy (and other losses) return the SUM of per-token
+    losses, not the mean. This is by design — it lets users control
+    normalization via the weights tensor. The trainer logs this raw sum;
+    users who want per-token loss should divide by number of active tokens.
+    """
+    if hasattr(result, "metrics"):
+        metrics = result.metrics or {}
+        return float(metrics.get("loss:sum", metrics.get("loss", 0.0)))
+    if isinstance(result, dict):
+        metrics = result.get("metrics", {})
+        return float(metrics.get("loss:sum", metrics.get("loss", 0.0)))
+    return 0.0
+
+
+def _create_training_client(args: HatcheryConfig, base_model: str):
+    """Create a training client for either Tinker or Hatchery backend."""
+    if args.backend == "tinker":
+        import tinker
+
+        api_key = args.api_key or os.environ.get("TINKER_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "Tinker API key required. Set `hatchery.api_key` in config "
+                "or TINKER_API_KEY env var."
+            )
+        os.environ["TINKER_API_KEY"] = api_key
+
+        service = tinker.ServiceClient(project_id=args.project_id)
+        return service.create_lora_training_client(
+            base_model=base_model,
+            rank=args.lora_rank,
+            train_mlp=args.train_mlp,
+            train_attn=args.train_attn,
+            train_unembed=args.train_unembed,
+        )
+
+    from hatchery.core.client import HatcheryClient
+
+    base_url = args.base_url or os.environ.get("HATCHERY_URL", "http://127.0.0.1:8420")
+    token = args.api_key or os.environ.get("HATCHERY_API_KEY", "dev")
+
+    client = HatcheryClient(base_url=base_url, token=token, timeout=args.future_timeout)
+    return client.create_lora_training_client(
+        base_model=base_model,
+        rank=args.lora_rank,
+        train_attn=args.train_attn,
+        train_mlp=args.train_mlp,
+        train_unembed=args.train_unembed,
+    )
+
+
+class HatcheryTrainer(AxolotlTrainer):
+    """Trainer that sends preprocessed batches to a remote training API.
+
+    Replaces local forward/backward with remote API calls to Tinker or
+    Hatchery. Uses axolotl's full data preprocessing pipeline (tokenization,
+    chat templates, packing, etc.) but offloads compute to remote GPUs.
+    """
+
+    hatchery_args: Optional[HatcheryConfig]
+    _base_model_name: Optional[str]
+    _training_client: Any
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.hatchery_args = None
+        self._base_model_name = None
+        self._training_client = None
+
+    def _get_training_client(self):
+        """Lazily create the remote training session."""
+        if self._training_client is not None:
+            return self._training_client
+
+        args = self.hatchery_args
+        if args is None:
+            raise RuntimeError(
+                "HatcheryTrainer.hatchery_args not set. "
+                "Ensure the HatcheryPlugin is registered."
+            )
+
+        base_model = self._base_model_name
+        if not base_model:
+            raise RuntimeError("HatcheryTrainer._base_model_name not set.")
+
+        self._training_client = _create_training_client(args, base_model)
+
+        LOG.info(
+            f"Remote training session created: backend={args.backend}, "
+            f"model={base_model}, rank={args.lora_rank}"
+        )
+        return self._training_client
+
+    def _send_batch(self, batch: dict[str, torch.Tensor]):
+        """Convert batch to datums and send forward_backward to remote.
+
+        Returns (future, n_active_tokens) where n_active_tokens counts
+        the completion tokens in this batch (for loss normalization).
+        """
+        input_ids = batch["input_ids"]
+        labels = batch["labels"]
+        attention_mask = batch.get("attention_mask")
+
+        n_active = int((labels[:, 1:] != -100).sum().item())
+        datums = batch_to_datums_sft(input_ids, labels, attention_mask)
+
+        tc = self._get_training_client()
+        args = self.hatchery_args
+        assert args is not None  # validated by _get_training_client
+        send_datums = datums_to_tinker(datums)
+
+        future = tc.forward_backward(
+            send_datums,
+            loss_fn=args.loss_fn,
+            loss_fn_config=args.loss_fn_config,
+        )
+        return future, n_active
+
+    def _do_optim_step(self):
+        """Send optimizer step to remote using axolotl's training params."""
+        import tinker.types as tt
+
+        tc = self._get_training_client()
+        return tc.optim_step(tt.AdamParams(**self._optim_params))
+
+    def train(
+        self,
+        resume_from_checkpoint: Optional[str] = None,
+        trial: Any = None,
+        ignore_keys_for_eval: Optional[list[str]] = None,
+        **kwargs,
+    ) -> TrainOutput:
+        """Main training loop — sends batches to remote API."""
+        args = self.hatchery_args
+        if args is None:
+            raise RuntimeError("hatchery_args not configured")
+
+        train_dataloader = self.get_train_dataloader()
+        num_batches = len(train_dataloader)
+
+        grad_accum = self.args.gradient_accumulation_steps
+        num_train_epochs = int(self.args.num_train_epochs)
+        steps_per_epoch = max(num_batches // grad_accum, 1)
+        max_steps = (
+            self.args.max_steps
+            if self.args.max_steps > 0
+            else steps_per_epoch * num_train_epochs
+        )
+
+        LOG.info(
+            f"Remote training: {num_batches} batches/epoch, "
+            f"{grad_accum} grad_accum, {max_steps} max steps, "
+            f"{num_train_epochs} epochs"
+        )
+
+        self.state.max_steps = max_steps
+        self.state.num_train_epochs = num_train_epochs
+        self.state.is_local_process_zero = True
+        self.state.is_world_process_zero = True
+
+        self.control = self.callback_handler.on_train_begin(
+            self.args,
+            self.state,
+            self.control,  # type: ignore[has-type]
+        )
+
+        global_step = 0
+        total_loss = 0.0
+        start_time = time.time()
+
+        for _epoch in range(num_train_epochs):
+            if global_step >= max_steps:
+                break
+
+            self.control = self.callback_handler.on_epoch_begin(
+                self.args, self.state, self.control
+            )
+
+            pending_fb_futures = []
+            accum_count = 0
+
+            for batch_idx, batch in enumerate(train_dataloader):
+                if global_step >= max_steps:
+                    break
+
+                self.control = self.callback_handler.on_step_begin(
+                    self.args, self.state, self.control
+                )
+
+                fb_future, n_active = self._send_batch(batch)
+                pending_fb_futures.append((fb_future, n_active))
+                accum_count += 1
+
+                if accum_count >= grad_accum:
+                    step_loss_sum = 0.0
+                    step_active = 0
+                    for fut, n_act in pending_fb_futures:
+                        result = fut.result(timeout=args.future_timeout)
+                        step_loss_sum += _extract_loss(result)
+                        step_active += n_act
+
+                    optim_future = self._do_optim_step()
+                    if not args.pipeline:
+                        optim_future.result(timeout=args.future_timeout)
+
+                    step_loss = (
+                        step_loss_sum / step_active
+                        if step_active > 0
+                        else step_loss_sum
+                    )
+
+                    global_step += 1
+                    total_loss += step_loss
+                    self.state.global_step = global_step
+                    self.state.epoch = _epoch + (batch_idx + 1) / num_batches
+
+                    log_interval = self.args.logging_steps or 1
+                    if global_step % log_interval == 0:
+                        elapsed = time.time() - start_time
+                        avg_loss = total_loss / global_step
+                        LOG.info(
+                            f"[step {global_step}/{max_steps}] "
+                            f"loss/tok={step_loss:.4f} avg={avg_loss:.4f} "
+                            f"active={step_active} "
+                            f"{elapsed / global_step:.2f}s/step"
+                        )
+                        self.log(
+                            {
+                                "loss": step_loss,
+                                "learning_rate": self._optim_params["learning_rate"],
+                                "epoch": self.state.epoch,
+                            }
+                        )
+
+                    if args.save_steps and global_step % args.save_steps == 0:
+                        self._save_remote_checkpoint(global_step)
+
+                    self.control = self.callback_handler.on_step_end(
+                        self.args, self.state, self.control
+                    )
+
+                    pending_fb_futures = []
+                    accum_count = 0
+
+                    if self.control.should_training_stop:
+                        break
+
+            self.control = self.callback_handler.on_epoch_end(
+                self.args, self.state, self.control
+            )
+            if self.control.should_training_stop:
+                break
+
+        if global_step > 0:
+            self._save_remote_checkpoint(global_step, name="final")
+
+        elapsed = time.time() - start_time
+        avg_loss = total_loss / max(global_step, 1)
+
+        LOG.info(
+            f"Training complete: {global_step} steps, {elapsed:.1f}s total, "
+            f"{elapsed / max(global_step, 1):.2f}s/step, avg_loss={avg_loss:.4f}"
+        )
+
+        self.control = self.callback_handler.on_train_end(
+            self.args, self.state, self.control
+        )
+
+        return TrainOutput(
+            global_step=global_step,
+            training_loss=avg_loss,
+            metrics={"train_loss": avg_loss, "train_runtime": elapsed},
+        )
+
+    def _save_remote_checkpoint(self, step: int, name: Optional[str] = None):
+        """Save a checkpoint on the remote service."""
+        tc = self._get_training_client()
+        args = self.hatchery_args
+        assert args is not None  # validated by _get_training_client
+        ckpt_name = name or f"{args.save_name_prefix}-{step:06d}"
+        try:
+            future = tc.save_state(ckpt_name)
+            future.result(timeout=args.future_timeout)
+            LOG.info(f"Remote checkpoint saved: {ckpt_name}")
+        except Exception:
+            LOG.exception(f"Failed to save checkpoint {ckpt_name}")
+            if name == "final":
+                raise
+
+    def save_model(self, output_dir=None, _internal_call=False):
+        """Delegate to remote checkpoint save so HF callbacks create checkpoints."""
+        self._save_remote_checkpoint(
+            step=self.state.global_step,
+            name=output_dir or "hf-save",
+        )
+
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        raise NotImplementedError(
+            "HatcheryTrainer uses remote API; compute_loss should not be called."
+        )
diff --git a/src/axolotl/integrations/kd/README.md b/src/axolotl/integrations/kd/README.md
new file mode 100644
index 0000000000..1f24fc8a6d
--- /dev/null
+++ b/src/axolotl/integrations/kd/README.md
@@ -0,0 +1,23 @@
+# Knowledge Distillation
+
+## Usage
+
+```yaml
+plugins:
+  - "axolotl.integrations.kd.KDPlugin"
+
+kd_trainer: True
+kd_ce_alpha: 0.1
+kd_alpha: 0.9
+kd_temperature: 1.0
+
+torch_compile: True  # recommended to reduce vram
+
+datasets:
+  - path: ...
+    type: "axolotl.integrations.kd.chat_template"
+    field_messages: "messages_combined"
+    logprobs_field: "llm_text_generation_vllm_logprobs"  # for kd only, field of logprobs
+```
+
+An example dataset can be found at [`axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample`](https://huggingface.co/datasets/axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample)
diff --git a/src/axolotl/integrations/kd/__init__.py b/src/axolotl/integrations/kd/__init__.py
new file mode 100644
index 0000000000..25a8537253
--- /dev/null
+++ b/src/axolotl/integrations/kd/__init__.py
@@ -0,0 +1,103 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Plugin init to add KD support to Axolotl.
+"""
+
+from typing import Any
+
+from transformers import Trainer
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.integrations.kd.callbacks import KDTemperatureSchedulerCallback
+
+from .args import KDArgs as KDArgs
+
+
+class KDPlugin(BasePlugin):
+    """
+    Plugin for KD support in Axolotl.
+    """
+
+    def get_input_args(self):
+        return "axolotl.integrations.kd.KDArgs"
+
+    def get_training_args_mixin(self):
+        return "axolotl.integrations.kd.args.KDTrainingArgsMixin"
+
+    def get_trainer_cls(self, cfg):
+        if cfg.kd_trainer:
+            from .trainer import AxolotlKDTrainer
+
+            return AxolotlKDTrainer
+        return None
+
+    def get_training_args(self, cfg):
+        return {
+            "kd_ce_alpha": cfg.kd_ce_alpha,
+            "kd_alpha": cfg.kd_alpha,
+            "kd_temperature": cfg.kd_temperature,
+            "kd_beta": cfg.kd_beta,
+            "kd_normalize_topk": cfg.kd_normalize_topk,
+        }
+
+    def get_collator_cls_and_kwargs(self, cfg, is_eval=False):
+        if not cfg.kd_trainer:
+            return None, None
+
+        from .collator import DataCollatorForKD, KDBatchSamplerDataCollatorForSeq2Seq
+
+        use_batch_sampler_collator = False
+        if is_eval is False and cfg.sample_packing:
+            use_batch_sampler_collator = True
+        if cfg.eval_sample_packing and is_eval:
+            use_batch_sampler_collator = True
+
+        if cfg.kd_online_server_base_url:
+            from .collator_online_teacher import OnlineTeacherCollator
+
+            return OnlineTeacherCollator, {
+                "kd_online_server_base_url": cfg.kd_online_server_base_url,
+                "kd_online_topk": cfg.kd_online_topk,
+                "kd_temperature": cfg.kd_temperature,
+                "kd_online_server": cfg.kd_online_server,
+                "kd_online_timeout": cfg.kd_online_timeout,
+                "kd_normalize_topk": cfg.kd_normalize_topk,
+            }
+
+        if use_batch_sampler_collator:
+            return KDBatchSamplerDataCollatorForSeq2Seq, {}
+        return DataCollatorForKD, {}
+
+    def add_callbacks_post_trainer(self, cfg: Any, trainer: Trainer) -> list:
+        """
+        Adds temp scheduler callback to the Trainer instance.
+
+        Args:
+            cfg (Any): Configuration object containing the sparse recipe.
+            trainer (Trainer): Huggingface Trainer instance.
+
+        Returns:
+            list: List containing the configured callback instances.
+        """
+        if cfg.kd_temperature_min is not None and cfg.kd_online_server_base_url:
+            callback = KDTemperatureSchedulerCallback(
+                cfg.kd_temperature,
+                cfg.kd_temperature_min,
+                trainer,
+            )
+            return [callback]
+
+        return []
diff --git a/src/axolotl/integrations/kd/args.py b/src/axolotl/integrations/kd/args.py
new file mode 100644
index 0000000000..425d8ddf60
--- /dev/null
+++ b/src/axolotl/integrations/kd/args.py
@@ -0,0 +1,76 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Plugin args for KD support.
+"""
+
+from dataclasses import dataclass
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class InferenceServerType(str, Enum):
+    """
+    Online inferences server types to handle different request args
+    """
+
+    vllm = "vllm"
+    sglang = "sglang"
+
+
+class KDArgs(BaseModel):
+    """
+    Input args for knowledge distillation.
+    """
+
+    kd_trainer: float | None = None  # whether to use KD trainer
+    kd_ce_alpha: float | None = (
+        None  # loss coefficient for cross-entropy loss during KD
+    )
+    kd_alpha: float | None = None  # loss coefficient for KD loss
+    kd_temperature: float | None = None  # temperature for sampling during KD
+    kd_beta: float | None = 0.0  # beta coefficient for ratio of fwd and reverse KL
+    kd_normalize_topk: bool | None = (
+        None  # whether to normalize student logits during KD
+    )
+
+    # TODO online kd
+    kd_online_server_base_url: str | None = None
+    kd_online_topk: int | None = None
+    kd_online_server: InferenceServerType | None = Field(
+        default_factory=lambda: InferenceServerType.vllm
+    )
+    kd_online_timeout: int | None = 120
+    kd_temperature_min: float | None = (
+        None  # kd temperature scheduling during online kd
+    )
+
+
+@dataclass
+class KDTrainingArgsMixin:
+    """
+    Additional args for KD training.
+    """
+
+    kd_ce_alpha: float | None = (
+        None  # loss coefficient for cross-entropy loss during KD
+    )
+    kd_alpha: float | None = None  # loss coefficient for KD loss
+    kd_temperature: float | None = None  # temperature for sampling during KD
+    kd_beta: float | None = None  # beta coefficient for ratio of fwd and reverse KL
+    kd_normalize_topk: float | None = (
+        None  # whether to normalize student logits during KD
+    )
diff --git a/src/axolotl/integrations/kd/callbacks.py b/src/axolotl/integrations/kd/callbacks.py
new file mode 100644
index 0000000000..c73d8a8bbe
--- /dev/null
+++ b/src/axolotl/integrations/kd/callbacks.py
@@ -0,0 +1,34 @@
+"""
+Transformers trainer callbacks to schedule the KD temperature during training
+"""
+
+import math
+
+from transformers.trainer_callback import TrainerCallback
+
+
+class KDTemperatureSchedulerCallback(TrainerCallback):
+    """
+    KD temperature scheduler callback for the trainer.
+    """
+
+    def __init__(self, temperature_start, temperature_min, trainer):
+        self.temperature_start = temperature_start
+        self.temperature_min = temperature_min
+        self.temperature = temperature_start
+
+        self.trainer = trainer
+
+    def on_step_end(self, args, state, control, **kwargs):
+        # cosine decay temperature over the max steps
+
+        progress = state.global_step / state.max_steps
+        # Cosine decay factor: 0.5 * (1 + cos(pi * progress))
+        # This factor goes from 1 (at progress=0) to 0 (at progress=1)
+        decay_factor = 0.5 * (1.0 + math.cos(math.pi * progress))
+        self.temperature = self.temperature_start - (
+            (self.temperature_start - self.temperature_min) * (1.0 - decay_factor)
+        )
+
+        if hasattr(self.trainer.data_collator, "kd_temperature"):
+            self.trainer.data_collator.kd_temperature = self.temperature
diff --git a/src/axolotl/integrations/kd/chat_template.py b/src/axolotl/integrations/kd/chat_template.py
new file mode 100644
index 0000000000..5330a80f1e
--- /dev/null
+++ b/src/axolotl/integrations/kd/chat_template.py
@@ -0,0 +1,335 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Chat template prompt strategy loader with KD support
+"""
+
+from typing import Any, Dict
+
+import torch
+
+from axolotl.prompt_strategies.chat_template import ChatTemplateStrategy, StrategyLoader
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class ChatTemplateStrategyWithKD(ChatTemplateStrategy):
+    """
+    Handle fields for logprob KD
+    """
+
+    def __init__(
+        self,
+        prompter,
+        tokenizer,
+        train_on_inputs,
+        sequence_len,
+        roles_to_train=None,
+        train_on_eos=None,
+        train_on_eot=None,
+        eot_tokens=None,
+        split_thinking: bool | None = False,
+        logprobs_field="logprobs",
+        gen_temperature=1.0,
+        kd_temperature=1.0,
+    ):
+        self.logprobs_field = logprobs_field
+        self.gen_temperature = gen_temperature
+        self.kd_temperature = kd_temperature
+
+        super().__init__(
+            prompter,
+            tokenizer,
+            train_on_inputs,
+            sequence_len,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            train_on_eot=train_on_eot,
+            eot_tokens=eot_tokens,
+            split_thinking=split_thinking,
+        )
+
+    @property
+    def supports_batched(self) -> bool:
+        # batching doesn't work well for logprob data
+        return False
+
+    def transform_logprobs(self, sample):
+        """
+        Transform logprobs to target format for KD training
+        """
+
+        logprobs = sample.pop(self.logprobs_field)
+        target_seq_len = len(logprobs)
+        input_seq_len = len(sample["input_ids"])
+        input_padding_len = input_seq_len - target_seq_len
+        # get non-zero top-k (prune None logprobs from vllm data step)
+        top_k_vals = [
+            len(logprobs[i])
+            for i in range(len(logprobs))
+            if logprobs[i] is not None and len(logprobs[i])
+        ]
+        max_top_k = max(set(top_k_vals), key=top_k_vals.count)
+        min_top_k = min(set(top_k_vals), key=top_k_vals.count)
+        top_k = min(max_top_k, min_top_k)
+        if top_k == 0:
+            raise ValueError("No non-zero top-k logprobs found.")
+
+        target_logprobs = []
+        target_token_ids = []
+        target_mask = []
+
+        if input_padding_len < 0:
+            # logprobs is longer than target_seq_len,
+            # so we need to slice from the left/beginning of logprobs
+            logprobs = logprobs[:-input_seq_len]
+            input_padding_len = 0
+            # target_seq_len = input_seq_len
+
+        # truncate the second dimension of the logprobs to top_k
+        logprobs = [row[:top_k] for row in logprobs]
+
+        # fill with -inf for padding_len tokens for top_k tokens
+        # extend target_logprobs with a padding_len x top_k 2D list filled with -inf
+
+        # we shift for causal models in the trainer, so start the range from 0
+        for _ in range(0, input_padding_len):
+            target_logprobs.append([-float("inf")] * top_k)
+            target_token_ids.append(list(range(top_k)))
+            target_mask.append([0] * top_k)
+
+        for position in range(input_padding_len, input_seq_len):
+            if sample["labels"][position] == -100:
+                target_mask.append([0] * top_k)
+            else:
+                target_mask.append([1] * top_k)
+
+        for _, token_pos_logprobs in enumerate(logprobs):
+            # Initialize collections for logprobs and token_ids
+            position_logprobs = []
+            position_token_ids = []
+
+            # Process each token probability entry
+            for entry in token_pos_logprobs:
+                # Extract logprob value
+                logprob = entry["logprob"]
+
+                # Parse token_id from the "token_id:###" format
+                token_id = int(entry["token"].split(":")[1])
+
+                # Append to our collections
+                position_logprobs.append(logprob)
+                position_token_ids.append(token_id)
+
+            # Convert to a tensor for easier manipulation
+            position_logprobs_tensor = torch.tensor(
+                position_logprobs, dtype=torch.float
+            )
+
+            # Now we have distribution at T1 in log form, i.e. log p_{T1}(k).
+            # Next, re-scale to T2 = self.kd_temperature via exponent-based trick
+            # p_{T2}(k) = [p_{T1}(k)]^(T1 / T2) / Z
+            #
+            # Convert from log to probability
+            teacher_probs_t1 = position_logprobs_tensor.exp()
+            # normalize probabilities to sum to 1 in case they aren't already
+            teacher_probs_t1_sum = teacher_probs_t1.sum(dim=0, keepdim=True)
+            if teacher_probs_t1_sum > 1e-9:
+                teacher_probs_t1 = teacher_probs_t1 / teacher_probs_t1_sum
+            if self.kd_temperature != self.gen_temperature:
+                # Exponentiate by factor (T1 / T2)
+                exponent = self.gen_temperature / self.kd_temperature
+                teacher_probs_t2 = teacher_probs_t1**exponent
+            else:
+                teacher_probs_t2 = teacher_probs_t1
+            # Re-normalize
+            teacher_probs_t2 = teacher_probs_t2 / teacher_probs_t2.sum(
+                dim=0, keepdim=True
+            )
+            # Convert back to log
+            position_logprobs_tensor = torch.log(teacher_probs_t2)
+
+            # Now we have log p_{teacher, T2}(k) stored in position_logprobs_tensor
+            position_logprobs_scaled = position_logprobs_tensor.tolist()
+
+            target_logprobs.append(position_logprobs_scaled)
+            target_token_ids.append(position_token_ids)
+
+        # Update sample with transformed logprobs
+        sample["target_logprobs"] = target_logprobs
+        sample["target_token_ids"] = target_token_ids
+        sample["target_mask"] = target_mask
+
+        return sample
+
+    def _tokenize_single_prompt(self, prompt):
+        logprobs = prompt.pop(self.logprobs_field)
+        tokenized_prompt = super()._tokenize_single_prompt(prompt)
+        tokenized_prompt[self.logprobs_field] = logprobs
+
+        # let subclasses add fields before transform
+        tokenized_prompt = self._prepare_kd_fields(tokenized_prompt, prompt)
+
+        tokenized_prompt = self.transform_logprobs(tokenized_prompt)
+        return tokenized_prompt
+
+    def _prepare_kd_fields(self, tokenized_prompt, original_prompt):
+        """
+        Hook for subclasses to prepare additional KD fields before transform
+        """
+        return tokenized_prompt
+
+
+class ChatTemplateStrategyWithKDv2(ChatTemplateStrategyWithKD):
+    """
+    Strat for datasets with complete structured KD logprob data
+    """
+
+    def transform_logprobs(self, sample):
+        """
+        Transform logprobs to target format for KD training
+        """
+
+        logprobs = sample.pop(self.logprobs_field)
+        target_seq_len = len(logprobs)
+        input_seq_len = len(sample["input_ids"])
+        input_padding_len = input_seq_len - target_seq_len
+        # get non-zero top-k (prune None logprobs from vllm data step)
+        top_k_vals = [
+            len(logprobs[i])
+            for i in range(len(logprobs))
+            if logprobs[i] is not None and len(logprobs[i])
+        ]
+        max_top_k = max(set(top_k_vals), key=top_k_vals.count)
+        min_top_k = min(set(top_k_vals), key=top_k_vals.count)
+        top_k = min(max_top_k, min_top_k)
+        if top_k == 0:
+            raise ValueError("No non-zero top-k logprobs found.")
+
+        target_logprobs = []
+        target_token_ids = []
+        target_mask = []
+
+        if input_padding_len < 0:
+            # logprobs is longer than target_seq_len,
+            # so we need to slice from the left/beginning of logprobs
+            logprobs = logprobs[:-input_seq_len]
+            input_padding_len = 0
+            # target_seq_len = input_seq_len
+
+        # truncate the second dimension of the logprobs to top_k
+        logprobs = [row[:top_k] for row in logprobs]
+
+        # fill with -inf for padding_len tokens for top_k tokens
+        # extend target_logprobs with a padding_len x top_k 2D list filled with -inf
+
+        # we shift for causal models in the trainer, so start the range from 0
+        for _ in range(0, input_padding_len):
+            target_logprobs.append([-float("inf")] * top_k)
+            target_token_ids.append(list(range(top_k)))
+            target_mask.append([0] * top_k)
+
+        for position in range(input_padding_len, input_seq_len):
+            if sample["labels"][position] == -100:
+                target_mask.append([0] * top_k)
+            else:
+                target_mask.append([1] * top_k)
+
+        for token_pos_logprobs, pos_target_token_ids in zip(
+            logprobs, sample["target_token_ids"], strict=False
+        ):
+            # Convert to a tensor for easier manipulation
+            position_logprobs_tensor = torch.tensor(
+                token_pos_logprobs, dtype=torch.float
+            )
+
+            # Now we have distribution at T1 in log form, i.e. log p_{T1}(k).
+            # Next, re-scale to T2 = self.kd_temperature via exponent-based trick
+            # p_{T2}(k) = [p_{T1}(k)]^(T1 / T2) / Z
+            #
+            # Convert from log to probability
+            teacher_probs_t1 = position_logprobs_tensor.exp()
+            # normalize probabilities to sum to 1 in case they aren't already
+            teacher_probs_t1_sum = teacher_probs_t1.sum(dim=0, keepdim=True)
+            if teacher_probs_t1_sum > 1e-9:
+                teacher_probs_t1 = teacher_probs_t1 / teacher_probs_t1_sum
+            if self.kd_temperature != self.gen_temperature:
+                # Exponentiate by factor (T1 / T2)
+                exponent = self.gen_temperature / self.kd_temperature
+                teacher_probs_t2 = teacher_probs_t1**exponent
+            else:
+                teacher_probs_t2 = teacher_probs_t1
+            # Re-normalize
+            teacher_probs_t2 = teacher_probs_t2 / teacher_probs_t2.sum(
+                dim=0, keepdim=True
+            )
+            # Convert back to log
+            position_logprobs_tensor = torch.log(teacher_probs_t2)
+
+            # Now we have log p_{teacher, T2}(k) stored in position_logprobs_tensor
+            position_logprobs_scaled = position_logprobs_tensor.tolist()
+
+            target_logprobs.append(position_logprobs_scaled)
+            target_token_ids.append(pos_target_token_ids)
+
+        # Update sample with transformed logprobs
+        sample["target_logprobs"] = target_logprobs
+        sample["target_token_ids"] = target_token_ids
+        sample["target_mask"] = target_mask
+
+        return sample
+
+    def _prepare_kd_fields(self, tokenized_prompt, original_prompt):
+        """
+        Add pre-tokenized target_token_ids for v2 format
+        """
+        target_token_ids = original_prompt.pop("target_token_ids", None)
+        if target_token_ids is not None:
+            tokenized_prompt["target_token_ids"] = target_token_ids
+        return tokenized_prompt
+
+
+class KDStrategyLoader(StrategyLoader):
+    """
+    Load ChatTemplateStrategy with KD support using StrategyLoader.
+    """
+
+    def _get_strategy_cls(self, cfg):
+        return ChatTemplateStrategyWithKD
+
+    def _get_strategy_params(self, cfg, ds_cfg: Dict[str, Any]):
+        strategy_params = super()._get_strategy_params(cfg, ds_cfg)
+        if logprobs_field := ds_cfg.get("logprobs_field"):
+            strategy_params["logprobs_field"] = logprobs_field
+        if gen_temperature := ds_cfg.get("temperature"):
+            strategy_params["gen_temperature"] = gen_temperature
+        if kd_temperature := cfg.get("kd_temperature"):
+            strategy_params["kd_temperature"] = kd_temperature
+
+        return strategy_params
+
+
+class KDStrategyLoaderV2(KDStrategyLoader):
+    """
+    Load KD chat template datasets with pre-tokenized logprob data
+    """
+
+    def _get_strategy_cls(self, cfg):
+        return ChatTemplateStrategyWithKDv2
+
+
+load_legacy = KDStrategyLoader()
+load = KDStrategyLoaderV2()
diff --git a/src/axolotl/integrations/kd/collator.py b/src/axolotl/integrations/kd/collator.py
new file mode 100644
index 0000000000..675485d9df
--- /dev/null
+++ b/src/axolotl/integrations/kd/collator.py
@@ -0,0 +1,267 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+DataCollator for axolotl to handle KD fields without using -inf for padding,
+and with a teacher_mask to identify padded positions.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+from transformers import PreTrainedTokenizerBase
+from transformers.utils import PaddingStrategy
+
+from axolotl.utils.collators.batching import DataCollatorForSeq2Seq
+
+
+@dataclass
+class DataCollatorForKD(DataCollatorForSeq2Seq):
+    """
+    Data collator for KD, including handling KD-specific fields.
+
+    This version avoids using -inf and instead uses a large negative value for padding
+    target_logprobs. It also creates a teacher_mask to indicate which entries are valid.
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    model: Optional[Any] = None
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    label_pad_token_id: int = -100
+    position_pad_token_id: int = 0
+    return_tensors: str = "pt"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+
+    def __call__(self, features, return_tensors=None):
+        if return_tensors is None:
+            return_tensors = self.return_tensors
+
+        padding_side = self.tokenizer.padding_side
+        max_len = 0
+
+        # Pad labels and position_ids first
+        for feature_name, pad_token_id in [
+            ("labels", self.label_pad_token_id),
+            ("position_ids", self.position_pad_token_id),
+        ]:
+            if feature_name in features[0]:
+                feat = [f[feature_name] for f in features]
+                max_len = max(len(x) for x in feat)
+                if self.pad_to_multiple_of is not None:
+                    max_len = (
+                        (max_len + self.pad_to_multiple_of - 1)
+                        // self.pad_to_multiple_of
+                    ) * self.pad_to_multiple_of
+
+                for f in features:
+                    remainder = [pad_token_id] * (max_len - len(f[feature_name]))
+                    if isinstance(f[feature_name], list):
+                        f[feature_name] = (
+                            f[feature_name] + remainder
+                            if padding_side == "right"
+                            else remainder + f[feature_name]
+                        )
+                    else:
+                        # If they are numpy arrays
+                        if padding_side == "right":
+                            f[feature_name] = np.concatenate(
+                                [f[feature_name], remainder]
+                            ).astype(np.int64)
+                        else:
+                            f[feature_name] = np.concatenate(
+                                [remainder, f[feature_name]]
+                            ).astype(np.int64)
+
+        # Handle target_logprobs and target_token_ids manually
+        target_logprobs_list = []
+        target_token_ids_list = []
+        target_mask_list = []
+        has_teacher_data = ("target_logprobs" in features[0]) and (
+            "target_token_ids" in features[0]
+        )
+
+        if has_teacher_data:
+            # Extract and remove from features
+            for f in features:
+                target_logprobs_list.append(f.pop("target_logprobs"))
+                target_token_ids_list.append(f.pop("target_token_ids"))
+                target_mask_list.append(f.pop("target_mask"))
+
+            # Determine max lengths
+            max_teacher_seq_len = max_len or max(
+                len(seq) for seq in target_logprobs_list
+            )
+            max_k = max(len(seq_k) for seq in target_logprobs_list for seq_k in seq)
+
+            padded_target_logprobs = []
+            padded_target_token_ids = []
+            padded_teacher_mask_list = []
+
+            for t_logprobs, t_ids, t_mask in zip(
+                target_logprobs_list,
+                target_token_ids_list,
+                target_mask_list,
+                strict=False,
+            ):
+                t_logprobs_padded = []
+                t_ids_padded = []
+                t_mask_padded = []
+
+                for lp, ids, mask in zip(t_logprobs, t_ids, t_mask, strict=False):
+                    lp_len = len(lp)
+                    if lp_len < max_k:
+                        # Use -1e9 for padding logprobs and 0 for token_ids
+                        pad_len = max_k - lp_len
+                        lp = lp + [-1e9] * pad_len
+                        ids = ids + [0] * pad_len
+                        mask = mask + [0] * pad_len
+                    else:
+                        lp = lp[:max_k]
+                        ids = ids[:max_k]
+                        mask = mask[:max_k]
+
+                    t_logprobs_padded.append(lp)
+                    t_ids_padded.append(ids)
+                    t_mask_padded.append(mask)
+
+                seq_len_diff = max_teacher_seq_len - len(t_logprobs_padded)
+                if seq_len_diff > 0:
+                    # Pad sequences fully if needed
+                    t_logprobs_padded.extend(
+                        [[-1e9] * max_k for _ in range(seq_len_diff)]
+                    )
+                    t_ids_padded.extend([[0] * max_k for _ in range(seq_len_diff)])
+                    t_mask_padded.extend([[0] * max_k for _ in range(seq_len_diff)])
+
+                padded_target_logprobs.append(t_logprobs_padded)
+                padded_target_token_ids.append(t_ids_padded)
+                padded_teacher_mask_list.append(t_mask_padded)
+
+            # Convert to tensors
+            padded_target_logprobs = torch.tensor(
+                padded_target_logprobs, dtype=torch.float
+            )
+            padded_target_token_ids = torch.tensor(
+                padded_target_token_ids, dtype=torch.long
+            )
+            padded_teacher_mask_list = torch.tensor(
+                padded_teacher_mask_list, dtype=torch.int
+            )
+
+        # Pad using tokenizer for regular fields
+        features = self.tokenizer.pad(
+            features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=return_tensors,
+        )
+
+        # Add back teacher data if present
+        if has_teacher_data:
+            features["target_logprobs"] = padded_target_logprobs
+            features["target_token_ids"] = padded_target_token_ids
+            features["target_mask"] = padded_teacher_mask_list
+
+        # Prepare decoder_input_ids if the model supports it
+        if (
+            "labels" in features
+            and self.model is not None
+            and hasattr(self.model, "prepare_decoder_input_ids_from_labels")
+        ):
+            decoder_input_ids = self.model.prepare_decoder_input_ids_from_labels(
+                labels=features["labels"]
+            )
+            features["decoder_input_ids"] = decoder_input_ids
+
+        return features
+
+
+class KDBatchSamplerDataCollatorForSeq2Seq(DataCollatorForKD):
+    """
+    Collator for multipack (batch of sub-batches) specifically for KD.
+    Adapts DataCollatorForKD so it can pack multiple sequences in a single batch item.
+    """
+
+    def __call__(self, features, return_tensors=None):
+        """
+        Expects that `features` could be either:
+          - a single list of dicts, OR
+          - a list of lists of dicts (the "sub-batches" to be packed).
+        """
+        # 1) If we are *not* dealing with multiple sequences per batch element,
+        #    just pass straight to parent.
+        if not isinstance(features[0], list):
+            return super().__call__(features, return_tensors=return_tensors)
+
+        # 2) Otherwise, we *are* dealing with multiple sequences in each batch item.
+        #    We want to produce a single "merged" feature dict for each sub-batch.
+        out_features = [{} for _ in features]
+
+        for i, sub_features in enumerate(features):
+            # sub_features is a list of dicts, each dict = one sequence’s features
+            # We'll merge them into out_features[i].
+            #
+            # NOTE: You can customize how you combine fields as needed (e.g. summation
+            # or offset for attention_mask). Below is a straightforward concatenation/extension.
+
+            for field_name in sub_features[0].keys():
+                # Some fields you might want to skip or treat specially:
+                if field_name == "length":
+                    continue
+
+                # If it’s a KD field that’s a list-of-lists (e.g. target_logprobs),
+                # you typically just want to flatten them by extending.
+                if field_name in ["target_logprobs", "target_token_ids", "target_mask"]:
+                    combined = []
+                    for feat in sub_features:
+                        combined.extend(feat[field_name])
+                    out_features[i][field_name] = combined
+
+                elif field_name == "attention_mask":
+                    # Here we apply the (j+1) factor to differentiate each sub-sample
+                    # within this merged batch item.
+                    arrays = []
+                    for j, feat in enumerate(sub_features):
+                        if field_name in feat:
+                            arrays.append((j + 1) * np.array(feat[field_name]))
+                    out_features[i][field_name] = np.concatenate(arrays)
+                else:
+                    # By default, just concatenate them if they are arrays
+                    # or extend them if they are lists.
+                    # For example, input_ids or labels are often arrays.
+                    arrays = []
+                    for feat in sub_features:
+                        if field_name in feat and isinstance(
+                            feat[field_name], (list, torch.Tensor)
+                        ):
+                            if isinstance(feat[field_name][0], (dict, str)):
+                                continue
+                            arr = np.array(feat[field_name])
+                            arrays.append(arr)
+                    if arrays:
+                        out_features[i][field_name] = np.concatenate(arrays)
+
+        # 3) Now call the parent collator, which will do:
+        #    - padding of labels/position_ids
+        #    - KD-specific padding for target_logprobs, target_token_ids, etc.
+        #    - final conversion to return_tensors
+        return super().__call__(out_features, return_tensors=return_tensors)
diff --git a/src/axolotl/integrations/kd/collator_online_teacher.py b/src/axolotl/integrations/kd/collator_online_teacher.py
new file mode 100644
index 0000000000..2cc48cf86a
--- /dev/null
+++ b/src/axolotl/integrations/kd/collator_online_teacher.py
@@ -0,0 +1,564 @@
+"""
+Packed data loader for online teacher training supporting vllm and sglang.
+"""
+
+import hashlib
+import hmac
+from typing import Any, Dict, List, Optional
+
+import requests
+import torch
+from orjson import orjson
+
+from axolotl.integrations.kd.collator import KDBatchSamplerDataCollatorForSeq2Seq
+from axolotl.integrations.kd.utils import normalize_logprobs
+from axolotl.utils.data.utils import retry_on_request_exceptions
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def hmac_sha_from_int_list(int_list, key, hash_func=hashlib.sha256):
+    """
+    Create HMAC-SHA hash from a list of integers
+
+    Args:
+        int_list: List of integers
+        key: Secret key (string or bytes)
+        hash_func: Hash function (default: sha256)
+
+    Returns:
+        HMAC digest as hex string
+    """
+    # Convert key to bytes if it's a string
+    if isinstance(key, str):
+        key = key.encode("utf-8")
+
+    # Convert list of ints to bytes
+    # Method 1: Convert each int to bytes and concatenate
+    data = b"".join(i.to_bytes(4, byteorder="big") for i in int_list)
+
+    # Create HMAC
+    h = hmac.new(key, data, hash_func)
+    return h.hexdigest()
+
+
+class OnlineTeacherCollator(KDBatchSamplerDataCollatorForSeq2Seq):
+    """
+    Collator for online teacher training.
+    """
+
+    DEFAULT_LABEL_PAD_TOKEN_ID: int = -100
+
+    def __init__(
+        self,
+        *args: Any,
+        kd_online_server_base_url: Optional[str] = None,
+        kd_online_topk: Optional[int] = None,
+        kd_temperature: Optional[float] = 1.0,
+        kd_online_server: Optional[str] = "vllm",
+        kd_online_timeout: Optional[int] = 120,
+        kd_cache_dir: Optional[str] = None,
+        kd_normalize_topk: Optional[bool] = True,
+        **kwargs: Any,
+    ):
+        super().__init__(*args, **kwargs)
+
+        if kd_online_server_base_url is None:
+            raise ValueError(
+                "kd_online_server_base_url must be provided for OnlineTeacherDataloader"
+            )
+        if kd_online_topk is None or kd_online_topk <= 0:
+            raise ValueError(
+                "kd_online_topk must be a positive integer for OnlineTeacherDataloader"
+            )
+
+        self.kd_online_server_base_url = kd_online_server_base_url.rstrip("/")
+        self.kd_online_topk = kd_online_topk
+        self.kd_temperature = kd_temperature
+        self.kd_online_server = kd_online_server
+        self.http_session = requests.Session()
+        self.kd_online_timeout = kd_online_timeout
+        self.kd_cache_dir = kd_cache_dir
+        self.kd_normalize_topk = kd_normalize_topk
+
+    def _normalize_logprobs(self, raw_logprobs: List[float]) -> List[float]:
+        """
+        Re-normalizes top-k raw logprobs as probabilities, and converts back to logprobs.
+        """
+        if not raw_logprobs or self.kd_online_topk == 0:
+            return (
+                [-float("inf")] * self.kd_online_topk if self.kd_online_topk > 0 else []
+            )
+
+        raw_logprobs_tensor = torch.tensor(raw_logprobs, dtype=torch.float32)
+        return normalize_logprobs(raw_logprobs_tensor, self.kd_online_topk).tolist()
+
+    @retry_on_request_exceptions(max_retries=10, delay=5)
+    def fetch_online_logprobs_sglang(
+        self, batch_input_ids: List[List[int]], labels: List[List[int]]
+    ):
+        """
+        Fetches logprobs from an online teacher served by sglang for a batch of input_ids.
+        Assumes API returns token IDs as strings in logprob dictionary keys.
+        """
+        api_endpoint = f"{self.kd_online_server_base_url}/generate"
+
+        payload = {
+            "input_ids": batch_input_ids,
+            "return_logprob": True,
+            "top_logprobs_num": self.kd_online_topk,
+            "logprob_start_len": 0,
+            "return_text_in_logprobs": True,
+            "echo": True,
+            "sampling_params": {
+                "max_new_tokens": 0,
+                "temperature": self.kd_temperature,
+                "skip_special_tokens": False,
+            },
+        }
+
+        # Initialize with empty lists, so if API call fails, these are returned.
+        ret_data_target_token_ids: List[List[List[int]]] = []
+        ret_data_target_logprobs: List[List[List[float]]] = []
+        ret_data_target_mask: List[List[List[int]]] = []
+
+        try:
+            response = self.http_session.post(
+                api_endpoint, json=payload, timeout=self.kd_online_timeout
+            )
+            response.raise_for_status()
+            api_data: list[dict] = response.json()
+
+            # Ensure api_data is a list, and its length matches batch_input_ids
+            if not isinstance(api_data, list) or len(api_data) != len(batch_input_ids):
+                LOG.error(
+                    f"API response format error. Expected a list of {len(batch_input_ids)} "
+                    f"items, got {type(api_data)} with length {len(api_data) if isinstance(api_data, list) else 'N/A'}."
+                )
+                # Return empty data; items processed later will get default empty KD fields
+                return {
+                    "target_token_ids": ret_data_target_token_ids,
+                    "target_logprobs": ret_data_target_logprobs,
+                    "target_mask": ret_data_target_mask,
+                }
+
+            for sequence_data, seq_input_ids, seq_labels in zip(
+                api_data, batch_input_ids, labels, strict=False
+            ):
+                current_target_logprobs = []
+                current_target_token_ids = []
+                current_target_mask = []
+
+                meta_info = sequence_data.pop("meta_info", {})
+                # Ensure input_top_logprobs is a list
+                input_top_logprobs: Optional[list[None | list[tuple]]] = meta_info.pop(
+                    "input_top_logprobs", []
+                )
+                if not isinstance(input_top_logprobs, list):
+                    LOG.warning(
+                        f"Received non-list input_top_logprobs: {input_top_logprobs}. Skipping sequence."
+                    )
+                    input_top_logprobs = []  # Treat as empty
+
+                # basic check that the logprob data len matches the input len, so no need to handle padding
+                assert len(seq_input_ids) == len(input_top_logprobs)
+
+                for i, _, label in zip(
+                    range(len(seq_input_ids)), seq_input_ids, seq_labels, strict=False
+                ):
+                    if i < len(input_top_logprobs) and input_top_logprobs[i] is None:
+                        # this is always the case for the first token.
+                        # there is never logprob data for the first token since that's a true input
+                        # so we replace the None value with padding data
+                        current_target_logprobs.append(
+                            [-float("inf")] * self.kd_online_topk
+                        )
+                        current_target_token_ids.append([0] * self.kd_online_topk)
+                        current_target_mask.append([0] * self.kd_online_topk)
+                    elif (
+                        i < len(input_top_logprobs)
+                        and input_top_logprobs[i] is not None
+                    ):
+                        pos_top_logprobs_data = input_top_logprobs[i]
+                        # Ensure pos_top_logprobs_data is a list of lists as expected
+                        if not (
+                            isinstance(pos_top_logprobs_data, list)
+                            and all(
+                                isinstance(item, list) for item in pos_top_logprobs_data
+                            )
+                            and len(pos_top_logprobs_data) > 0
+                            and len(pos_top_logprobs_data[0]) == 3
+                        ):  # [logprob, token_id, token_str]
+                            LOG.warning(
+                                f"Malformed pos_top_logprobs_data: {pos_top_logprobs_data}. Padding this position."
+                            )
+                            current_target_logprobs.append(
+                                [-float("inf")] * self.kd_online_topk
+                            )
+                            current_target_token_ids.append([0] * self.kd_online_topk)
+                            current_target_mask.append([0] * self.kd_online_topk)
+                            continue
+
+                        # pos_top_logprobs: list of logprobs, pos_token_ids: list of token_ids
+                        pos_logprobs_raw, pos_token_ids, _ = [
+                            list(row)
+                            for row in zip(*pos_top_logprobs_data, strict=False)
+                        ]
+
+                        # Ensure correct length (top_k)
+                        if len(pos_logprobs_raw) < self.kd_online_topk:
+                            pad_len = self.kd_online_topk - len(pos_logprobs_raw)
+                            pos_logprobs_raw.extend([-float("inf")] * pad_len)
+                            pos_token_ids.extend([0] * pad_len)  # Pad with 0 token_id
+
+                        # truncate to top_k in case the response was longer
+                        current_target_token_ids.append(
+                            pos_token_ids[: self.kd_online_topk]
+                        )
+
+                        if self.kd_normalize_topk:
+                            normalized_logprobs_for_position = self._normalize_logprobs(
+                                pos_logprobs_raw[: self.kd_online_topk]
+                            )
+                            current_target_logprobs.append(
+                                normalized_logprobs_for_position
+                            )
+                        else:
+                            current_target_logprobs.append(
+                                pos_logprobs_raw[: self.kd_online_topk]
+                            )
+
+                        # Mask depends on the corresponding label for the student
+                        if label == self.DEFAULT_LABEL_PAD_TOKEN_ID:
+                            current_target_mask.append([0] * self.kd_online_topk)
+                        else:
+                            current_target_mask.append([1] * self.kd_online_topk)
+                    else:
+                        # Pad if no logprobs for this position (either due to length mismatch or None entry)
+                        current_target_logprobs.append(
+                            [-float("inf")] * self.kd_online_topk
+                        )
+                        current_target_token_ids.append([0] * self.kd_online_topk)
+                        current_target_mask.append([0] * self.kd_online_topk)
+
+                ret_data_target_token_ids.append(current_target_token_ids)
+                ret_data_target_logprobs.append(current_target_logprobs)
+                ret_data_target_mask.append(current_target_mask)
+
+        except requests.exceptions.RequestException as e:
+            LOG.error(f"Error fetching logprobs from online teacher: {e}")
+            raise e
+            # ret_logprobs_data will be returned with empty lists, handled by the caller.
+        except Exception as e:  # Catch other potential errors during processing
+            LOG.error(
+                f"Unexpected error processing API response in fetch_online_logprobs: {e}",
+                exc_info=True,
+            )
+            raise e
+
+        return {
+            "target_token_ids": ret_data_target_token_ids,
+            "target_logprobs": ret_data_target_logprobs,
+            "target_mask": ret_data_target_mask,
+        }
+
+    @retry_on_request_exceptions(max_retries=10, delay=5)
+    def fetch_online_logprobs_vllm(
+        self, batch_input_ids: List[List[int]], labels: List[List[int]]
+    ):
+        """
+        Fetches logprobs from an online teacher served by vllm for a batch of input_ids.
+        Assumes API returns token IDs as strings in logprob dictionary keys.
+        """
+        api_endpoint = f"{self.kd_online_server_base_url}/v1/completions"
+
+        payload = {
+            "prompt": batch_input_ids,
+            "echo": True,
+            "logprobs": True,
+            "prompt_logprobs": self.kd_online_topk,
+            "top_logprobs": self.kd_online_topk,
+            "max_new_tokens": 0,
+            "skip_special_tokens": False,
+            "temperature": self.kd_temperature,
+            "sampling_params": {
+                "max_tokens": 0,
+            },
+        }
+
+        # Initialize with empty lists, so if API call fails, these are returned.
+        ret_data_target_token_ids: List[List[List[int]]] = []
+        ret_data_target_logprobs: List[List[List[float]]] = []
+        ret_data_target_mask: List[List[List[int]]] = []
+
+        try:
+            headers = {"Accept-Encoding": "deflate, gzip, br, zstd"}
+            response = self.http_session.post(
+                api_endpoint,
+                json=payload,
+                headers=headers,
+                timeout=self.kd_online_timeout,
+            )
+            response.raise_for_status()
+            api_data: dict = orjson.loads(response.content)
+            choices: list[dict] = api_data["choices"]
+
+            # Ensure api_data is a list, and its length matches batch_input_ids
+            if not isinstance(choices, list) or len(choices) != len(batch_input_ids):
+                LOG.error(
+                    f"API response format error. Expected a list of {len(batch_input_ids)} "
+                    f"items, got {type(api_data)} with length {len(api_data) if isinstance(api_data, list) else 'N/A'}."
+                )
+                # Return empty data; items processed later will get default empty KD fields
+                return {
+                    "target_token_ids": ret_data_target_token_ids,
+                    "target_logprobs": ret_data_target_logprobs,
+                    "target_mask": ret_data_target_mask,
+                }
+
+            for sequence_data, seq_input_ids, seq_labels in zip(
+                choices, batch_input_ids, labels, strict=False
+            ):
+                # seq_input_ids: List[int]
+                # seq_labels: List[int]
+
+                current_target_logprobs = []
+                current_target_token_ids = []
+                current_target_mask = []
+
+                # Ensure input_top_logprobs is a list
+                input_top_logprobs: Optional[list[None | dict[str, dict]]] = (
+                    sequence_data.pop("prompt_logprobs", [])
+                )
+
+                if not isinstance(input_top_logprobs, list):
+                    LOG.warning(
+                        f"Received non-list input_top_logprobs: {input_top_logprobs}. Skipping sequence."
+                    )
+                    input_top_logprobs = []  # Treat as empty
+
+                # basic check that the logprob data len matches the input len, so no need to handle padding
+                assert len(seq_input_ids) == len(input_top_logprobs)
+
+                seq_len = len(seq_input_ids)
+
+                for i, _, label in zip(
+                    range(seq_len), seq_input_ids, seq_labels, strict=False
+                ):
+                    if i < len(input_top_logprobs) and input_top_logprobs[i] is None:
+                        # this is always the case for the first token.
+                        # there is never logprob data for the first token since that's a true input
+                        continue
+                    if (
+                        i < len(input_top_logprobs)
+                        and input_top_logprobs[i] is not None
+                    ):
+                        pos_top_logprobs_data: dict[str, dict] = input_top_logprobs[i]  # type: ignore[assignment]
+                        # Ensure pos_top_logprobs_data is a list of lists as expected
+                        if not (
+                            isinstance(pos_top_logprobs_data, dict)
+                            and all(
+                                isinstance(item, dict)
+                                for item in pos_top_logprobs_data.values()
+                            )
+                            and len(pos_top_logprobs_data.keys()) > 0
+                        ):  # [logprob, token_id, token_str]
+                            LOG.warning(
+                                f"Malformed pos_top_logprobs_data: {pos_top_logprobs_data}. Padding this position."
+                            )
+                            current_target_logprobs.append(
+                                [-float("inf")] * self.kd_online_topk
+                            )
+                            current_target_token_ids.append(
+                                list(range(self.kd_online_topk))
+                            )
+                            current_target_mask.append([0] * self.kd_online_topk)
+                            continue
+
+                        # pos_top_logprobs: list of logprobs, pos_token_ids: list of token_ids
+                        pos_token_ids_str = list(pos_top_logprobs_data.keys())
+                        pos_logprobs_dict = pos_top_logprobs_data.values()
+                        pos_token_ids = [
+                            int(token_id) for token_id in pos_token_ids_str
+                        ]
+                        pos_logprobs_raw = [
+                            float(logprob.get("logprob", -float("inf")))
+                            for logprob in pos_logprobs_dict
+                        ]
+
+                        # Ensure correct length (top_k)
+                        if len(pos_logprobs_raw) < self.kd_online_topk:
+                            pad_len = self.kd_online_topk - len(pos_logprobs_raw)
+                            LOG.warning(
+                                f"Padding position {i} with {pad_len} top-k tokens and logprobs."
+                            )
+                            pos_logprobs_raw.extend([-float("inf")] * pad_len)
+                            pos_token_ids.extend([0] * pad_len)  # Pad with 0 token_id
+
+                        # truncate to top_k in case the response was longer
+                        current_target_token_ids.append(
+                            pos_token_ids[: self.kd_online_topk]
+                        )
+
+                        if self.kd_normalize_topk:
+                            normalized_logprobs_for_position = self._normalize_logprobs(
+                                pos_logprobs_raw[: self.kd_online_topk]
+                            )
+                            current_target_logprobs.append(
+                                normalized_logprobs_for_position
+                            )
+                        else:
+                            current_target_logprobs.append(
+                                pos_logprobs_raw[: self.kd_online_topk]
+                            )
+
+                        # Mask depends on the corresponding label for the student
+                        if label == self.DEFAULT_LABEL_PAD_TOKEN_ID:
+                            current_target_mask.append([0] * self.kd_online_topk)
+                        else:
+                            current_target_mask.append([1] * self.kd_online_topk)
+                    else:
+                        # Pad if no logprobs for this position (either due to length mismatch or None entry)
+                        current_target_logprobs.append(
+                            [-float("inf")] * self.kd_online_topk
+                        )
+                        current_target_token_ids.append(
+                            list(range(self.kd_online_topk))
+                        )
+                        current_target_mask.append([0] * self.kd_online_topk)
+                for _ in range(max(0, seq_len - len(current_target_logprobs))):
+                    current_target_logprobs.append(
+                        [-float("inf")] * self.kd_online_topk
+                    )
+                    current_target_token_ids.append(list(range(self.kd_online_topk)))
+                    current_target_mask.append([0] * self.kd_online_topk)
+
+                ret_data_target_token_ids.append(current_target_token_ids)
+                ret_data_target_logprobs.append(current_target_logprobs)
+                ret_data_target_mask.append(current_target_mask)
+
+                # TODO save and load targets to disk for caching for next epoch
+                # generate a hmac SHA256 hash over the list seq_input_ids and convert it to an int
+                # if self.kd_cache_dir:
+                #     hash_input_ids = hmac_sha_from_int_list(
+                #         seq_input_ids, f"{self.kd_online_server_base_url}:{self.kd_online_topk}"
+                #     )
+                #     with open(f"{self.kd_cache_dir}/{hash_input_ids}.parquet", "wb") as f:
+                #         pd.DataFrame(ret_logprobs_data).to_parquet(f, index=False)
+
+        except requests.exceptions.RequestException as e:
+            LOG.error(f"Error fetching logprobs from online teacher: {e}")
+            raise e
+            # ret_logprobs_data will be returned with empty lists, handled by the caller.
+        except Exception as e:  # Catch other potential errors during processing
+            LOG.error(
+                f"Unexpected error processing API response in fetch_online_logprobs: {e}",
+                exc_info=True,
+            )
+            raise e
+
+        return {
+            "target_token_ids": ret_data_target_token_ids,
+            "target_logprobs": ret_data_target_logprobs,
+            "target_mask": ret_data_target_mask,
+        }
+
+    def __call__(
+        self, features: List[List[Dict[str, Any]]], return_tensors: Optional[str] = None
+    ) -> Dict[str, Any]:
+        if not features:
+            return super().__call__(features, return_tensors=return_tensors)
+
+        for (
+            sub_batch_features
+        ) in features:  # sub_batch_features is List[Dict[str, Any]]
+            if not sub_batch_features:
+                continue
+
+            input_ids_for_api_call: List[List[int]] = []
+            labels_for_api_call: List[List[int]] = []
+            # Store references to the original item dictionaries to update them in-place
+            items_for_api_call: List[Dict[str, Any]] = []
+
+            for item_dict in sub_batch_features:
+                if not isinstance(item_dict, dict):
+                    LOG.warning(
+                        f"Skipping non-dict item in sub_batch_features: {item_dict}"
+                    )
+                    continue
+
+                current_input_ids = item_dict.get("input_ids")
+                current_labels = item_dict.get("labels")
+
+                if current_input_ids is not None and current_labels is not None:
+                    # Ensure input_ids and labels are lists of ints for JSON serialization
+                    input_ids_list = (
+                        current_input_ids.tolist()
+                        if hasattr(current_input_ids, "tolist")
+                        else list(current_input_ids)
+                    )
+                    labels_list = (
+                        current_labels.tolist()
+                        if hasattr(current_labels, "tolist")
+                        else list(current_labels)
+                    )
+
+                    input_ids_for_api_call.append(input_ids_list)
+                    labels_for_api_call.append(labels_list)
+                    items_for_api_call.append(item_dict)
+                else:
+                    # This item will not get teacher logprobs from the API.
+                    # Initialize KD fields to empty lists so downstream collators handle them uniformly.
+                    item_dict.setdefault("target_token_ids", [])
+                    item_dict.setdefault("target_logprobs", [])
+                    item_dict.setdefault("target_mask", [])
+
+            # print(items_for_api_call)
+            if items_for_api_call:  # Only call API if there's something to process
+                if self.kd_online_server == "sglang":
+                    api_responses_for_sub_batch = self.fetch_online_logprobs_sglang(
+                        input_ids_for_api_call, labels_for_api_call
+                    )
+                else:
+                    api_responses_for_sub_batch = self.fetch_online_logprobs_vllm(
+                        input_ids_for_api_call, labels_for_api_call
+                    )
+
+                # api_responses_for_sub_batch has keys: "target_token_ids", "target_logprobs", "target_mask"
+                # Each value is a list, corresponding to items_for_api_call
+                for i, item_to_update in enumerate(items_for_api_call):
+                    # TODO make sure to figure out which input in sub_batch_features to update the batch in the original `features` object so the super class can handle it properly.
+                    if api_responses_for_sub_batch and i < len(
+                        api_responses_for_sub_batch["target_token_ids"]
+                    ):  # Check bounds
+                        assert len(
+                            api_responses_for_sub_batch["target_token_ids"][i]
+                        ) == len(item_to_update["input_ids"])
+                        assert len(
+                            api_responses_for_sub_batch["target_logprobs"][i]
+                        ) == len(item_to_update["input_ids"])
+                        assert len(
+                            api_responses_for_sub_batch["target_mask"][i]
+                        ) == len(item_to_update["labels"])
+                        item_to_update["target_token_ids"] = (
+                            api_responses_for_sub_batch["target_token_ids"][i]
+                        )
+                        item_to_update["target_logprobs"] = api_responses_for_sub_batch[
+                            "target_logprobs"
+                        ][i]
+                        item_to_update["target_mask"] = api_responses_for_sub_batch[
+                            "target_mask"
+                        ][i]
+                    else:
+                        # API call failed for this item, or response was shorter than expected.
+                        # Ensure KD fields are initialized as empty lists.
+                        LOG.warning(
+                            f" (index {i}), or API response was too short. "
+                            f"API response keys: {list(api_responses_for_sub_batch.keys()) if api_responses_for_sub_batch else 'None'}"
+                        )
+                        item_to_update.setdefault("target_token_ids", [])
+                        item_to_update.setdefault("target_logprobs", [])
+                        item_to_update.setdefault("target_mask", [])
+
+        return super().__call__(features, return_tensors=return_tensors)
diff --git a/src/axolotl/integrations/kd/kernels/__init__.py b/src/axolotl/integrations/kd/kernels/__init__.py
new file mode 100644
index 0000000000..0a967144a0
--- /dev/null
+++ b/src/axolotl/integrations/kd/kernels/__init__.py
@@ -0,0 +1,7 @@
+"""
+Liger Chunked loss optimizations module
+"""
+
+from .liger import LigerFusedLinearKLTopKLogprobLoss
+
+__all__ = ["LigerFusedLinearKLTopKLogprobLoss"]
diff --git a/src/axolotl/integrations/kd/kernels/liger.py b/src/axolotl/integrations/kd/kernels/liger.py
new file mode 100644
index 0000000000..a150420855
--- /dev/null
+++ b/src/axolotl/integrations/kd/kernels/liger.py
@@ -0,0 +1,494 @@
+"""
+Liger Kernels for Chunked Top-K Log-Prob Distillation
+"""
+
+import torch
+import torch.nn.functional as F
+from liger_kernel.chunked_loss.fused_linear_distillation import (
+    LigerFusedLinearDistillationBase,
+)
+
+from axolotl.integrations.kd.utils import normalize_logprobs
+
+
+class LigerFusedLinearKLTopKLogprobFunction(LigerFusedLinearDistillationBase):
+    """
+    Chunked kl-div loss for top-k logprobs
+    """
+
+    @staticmethod
+    def distillation_loss_fn(
+        student_logits_temp_scaled: torch.Tensor,  # [chunk_size, vocab_size], already temp-scaled
+        target_token_ids_chunk: torch.Tensor,  # [chunk_size, top_k]
+        target_logprobs_chunk: torch.Tensor,  # [chunk_size, top_k], already temp-scaled and normalized logprobs
+        target_mask_chunk: torch.Tensor,  # [chunk_size, top_k]
+        beta: float = 0.0,
+        normalize_topk: bool = True,
+    ) -> torch.Tensor:
+        """
+        Compute Top-K KL divergence loss for a chunk.
+        Args:
+            student_logits_temp_scaled: Student logits, scaled by temperature. Shape: (N, V).
+            target_token_ids_chunk: Top-k teacher token IDs. Shape: (N, K).
+            target_logprobs_chunk: Top-k teacher log probabilities (temp-scaled, normalized). Shape: (N, K).
+            target_mask_chunk: Mask for valid top-k tokens. Shape: (N, K).
+            beta: Controls the type of KL divergence.
+                  0.0 for Forward KL (P_teacher || P_student).
+                  1.0 for Reverse KL (P_student || P_teacher).
+                  0.5 for Symmetric KL (average of Forward and Reverse).
+            normalize_topk: Whether to normalize the log probabilities
+        Returns:
+            Sum of KL divergence losses for the chunk.
+        """
+        topk = target_token_ids_chunk.shape[-1]
+        student_logits_temp_scaled = (  # [chunk_size, vocab_size]
+            student_logits_temp_scaled.float()
+        )
+        target_logprobs_chunk = target_logprobs_chunk.float()
+
+        # Gather student logits for the top-k teacher token IDs
+        # target_token_ids_chunk: [chunk_size, top_k]
+        # student_logits_topk_temp_scaled: [chunk_size, top_k]
+        student_logits_topk_temp_scaled = torch.gather(
+            student_logits_temp_scaled, dim=-1, index=target_token_ids_chunk
+        )
+
+        # Student log-probabilities for the gathered top-k tokens
+        student_lse = torch.logsumexp(
+            student_logits_temp_scaled, dim=-1, keepdim=True
+        )  # [chunk_size, 1]
+        student_logprobs_topk_temp_scaled = (
+            student_logits_topk_temp_scaled - student_lse
+        )
+
+        # we have the top-k student logprobs, normalize them
+        if normalize_topk:
+            student_logprobs_topk_temp_scaled = normalize_logprobs(
+                student_logprobs_topk_temp_scaled, topk
+            )
+
+        valid_mask = target_mask_chunk.to(torch.bool)  # [chunk_size, top_k]
+
+        student_logprobs_topk_valid = student_logprobs_topk_temp_scaled[valid_mask]
+        teacher_logprobs_valid = target_logprobs_chunk[valid_mask]
+
+        # Teacher probabilities P(y|x_teacher) from logprobs
+        # target_logprobs_valid are already normalized (log(softmax(teacher_logits/T)))
+        teacher_probs_valid = teacher_logprobs_valid.exp()
+        # Student probabilities P_student from log P_student
+        student_probs_topk_valid = student_logprobs_topk_valid.exp()
+
+        # kd_loss_per_token = torch.zeros_like(target_logprobs_valid)
+
+        # KL divergence: sum(P_teacher * (log P_teacher - log P_student))
+        # = sum(P_teacher * log P_teacher) - sum(P_teacher * log P_student)
+        # The distillation loss is often formulated as -sum(P_teacher * log P_student)
+        # or as sum(P_teacher * (log_softmax_teacher - log_softmax_student))
+        # Here, target_logprobs_valid are log_softmax_teacher.
+        # student_logprobs_topk_valid are log_softmax_student (for the selected K indices).
+        if beta == 0.0:  # Contribution from Forward KL
+            fwd_kl_per_token = teacher_probs_valid * (
+                teacher_logprobs_valid - student_logprobs_topk_valid
+            )
+            kd_loss = fwd_kl_per_token.sum()
+        elif beta == 1.0:  # Contribution from Reverse KL
+            rev_kl_per_token = student_probs_topk_valid * (
+                student_logprobs_topk_valid - teacher_logprobs_valid
+            )
+            kd_loss = rev_kl_per_token.sum()
+        else:
+            # JSD - Jensen-Shannon Divergence / Symmetric
+            mean_probs = (
+                1 - beta
+            ) * student_probs_topk_valid + beta * teacher_probs_valid
+            log_mean_probs = mean_probs.log()
+            student_kl = F.kl_div(
+                log_mean_probs,
+                student_logprobs_topk_valid,
+                reduction="sum",
+                log_target=True,
+            )
+            teacher_kl = F.kl_div(
+                log_mean_probs, teacher_logprobs_valid, reduction="sum", log_target=True
+            )
+            jsd_loss = beta * teacher_kl + (1 - beta) * student_kl
+            kd_loss = jsd_loss
+
+        return kd_loss
+
+    @staticmethod
+    def _compute_loss_kl_topk(
+        student_input_chunk: torch.Tensor,
+        student_weight: torch.Tensor,
+        # Args for student_bias, target_token_ids_chunk etc. are passed to the lambda wrapped by grad_and_value
+        # or through `partial`. Let's make them explicit here for clarity.
+        target_token_ids_chunk: torch.Tensor,
+        target_logprobs_chunk: torch.Tensor,
+        target_mask_chunk: torch.Tensor,
+        target_chunk: torch.Tensor,  # For hard loss (true labels)
+        student_bias: torch.Tensor = None,  # This will be one of the grad targets
+        # Other params passed via `partial` from `forward`
+        distillation_loss_fn=None,
+        ignore_index: int = -100,
+        weight_hard_loss: float = 0.5,
+        weight_soft_loss: float = 0.5,
+        compute_ce_loss: bool = True,
+        temperature: float = 1.0,
+        beta: float = 0.0,
+        normalize_topk: bool = True,
+    ):
+        # Compute student logits for the chunk from hidden states and LM head
+        # student_input_chunk: [chunk_size, hidden_dim]
+        # student_lm_head_weight: [vocab_size, hidden_dim]
+        # student_logits_chunk: [chunk_size, vocab_size]
+        student_logits_chunk = F.linear(
+            student_input_chunk, student_weight, student_bias
+        )
+
+        ce_loss = torch.tensor(
+            0.0, device=student_logits_chunk.device, dtype=student_logits_chunk.dtype
+        )
+        if compute_ce_loss and weight_hard_loss > 0.0:
+            ce_loss = F.cross_entropy(
+                student_logits_chunk.view(-1, student_logits_chunk.shape[-1]),
+                target_chunk.view(-1),
+                reduction="sum",
+                ignore_index=ignore_index,
+            )
+
+        soft_loss = torch.tensor(
+            0.0, device=student_logits_chunk.device, dtype=student_logits_chunk.dtype
+        )
+        if weight_soft_loss > 0.0:
+            student_logits_chunk_temp_scaled = student_logits_chunk / temperature
+
+            # Assuming student_weight.shape[0] (vocab_size) is adequate for target_token_ids_chunk.max()
+            # No explicit padding here; user must ensure vocab alignment or pre-pad student_weight.
+
+            soft_loss = distillation_loss_fn(
+                student_logits_chunk_temp_scaled,
+                target_token_ids_chunk,
+                target_logprobs_chunk,
+                target_mask_chunk,
+                beta=beta,
+                normalize_topk=normalize_topk,
+            )
+
+        return soft_loss, ce_loss
+
+    @classmethod
+    def forward(
+        cls,
+        ctx,
+        student_input: torch.Tensor,  # [batch_size, seq_len, dim]
+        student_lm_head_weight: torch.Tensor,  # [dim, vocab_size]
+        target_token_ids: torch.Tensor,  # [batch_size, seq_len, top_k]
+        target_logprobs: torch.Tensor,  # [batch_size, seq_len, top_k]
+        target_mask: torch.Tensor,  # [batch_size, seq_len, top_k]
+        true_labels: torch.Tensor,  # [batch_size, seq_len]
+        student_lm_head_bias: torch.Tensor = None,
+        weight_hard_loss: float = 0.5,
+        weight_soft_loss: float = 0.5,
+        ignore_index: int = -100,
+        temperature: float = 1.0,
+        beta: float = 0.0,
+        compiled: bool = False,
+        chunk_size: int = 1024,
+        compute_ce_loss: bool = True,
+        normalize_topk: bool = True,
+    ):
+        CHUNK_SIZE = chunk_size
+        grad_weight_acc = torch.zeros_like(student_lm_head_weight)
+        grad_inputs_list = []
+        grad_bias_acc = (
+            torch.zeros_like(student_lm_head_bias)
+            if student_lm_head_bias is not None
+            else None
+        )
+        kd_loss_acc = torch.zeros(
+            (), device=student_input.device, dtype=student_input.dtype
+        )
+        ce_loss_acc = torch.zeros(
+            (), device=student_input.device, dtype=student_input.dtype
+        )
+
+        # This function will be what torch.func.grad_and_value differentiates.
+        # It takes student_input_chunk, student_weight (full), student_bias (full) as primals.
+        # Other necessary data (target_*, etc.) are passed as non-differentiable arguments.
+        def loss_fn_for_grad(
+            _student_input_chunk,
+            _student_lm_head_weight,  # full weight
+            _student_lm_head_bias,  # full bias
+            # Fixed arguments for a given chunk, not differentiated:
+            _target_token_ids_chunk,
+            _target_logprobs_chunk,
+            _target_mask_chunk,
+            _true_labels_chunk,
+        ):
+            soft_loss, ce_loss = cls._compute_loss_kl_topk(
+                student_input_chunk=_student_input_chunk,
+                student_weight=_student_lm_head_weight,
+                target_token_ids_chunk=_target_token_ids_chunk,
+                target_logprobs_chunk=_target_logprobs_chunk,
+                target_mask_chunk=_target_mask_chunk,
+                target_chunk=_true_labels_chunk,
+                student_bias=_student_lm_head_bias,
+                distillation_loss_fn=cls.distillation_loss_fn,
+                ignore_index=ignore_index,
+                weight_hard_loss=weight_hard_loss,
+                weight_soft_loss=weight_soft_loss,
+                compute_ce_loss=compute_ce_loss,
+                temperature=temperature,
+                beta=beta,
+                normalize_topk=normalize_topk,
+            )
+            # has_aux=True: first return is the differentiated scalar, rest is aux.
+            # Combine here so both terms contribute to backward; aux carries the
+            # unweighted soft/ce values for reporting.
+            combined = weight_soft_loss * soft_loss + weight_hard_loss * ce_loss
+            return combined, (soft_loss, ce_loss)
+
+        def accumulate_chunk_grads(
+            student_input_chunk_ac,
+            target_token_ids_chunk_ac,
+            target_logprobs_chunk_ac,
+            target_mask_chunk_ac,
+            true_labels_chunk_ac,
+        ):
+            # student_weight and student_bias are closed over from the outer scope (full tensors)
+            if student_lm_head_bias is not None:
+                (
+                    (chunk_grad_input, chunk_grad_weight, chunk_grad_bias),
+                    # grad_and_value(has_aux=True) returns (grads, (value, aux));
+                    # value = combined scalar, aux = (soft_loss, ce_loss) from loss_fn_for_grad.
+                    (_, (chunk_kd_loss, chunk_ce_loss)),
+                ) = torch.func.grad_and_value(
+                    loss_fn_for_grad, argnums=(0, 1, 2), has_aux=True
+                )(
+                    student_input_chunk_ac,
+                    student_lm_head_weight,
+                    student_lm_head_bias,  # primals
+                    target_token_ids_chunk_ac,
+                    target_logprobs_chunk_ac,
+                    target_mask_chunk_ac,
+                    true_labels_chunk_ac,
+                )  # non-primals
+                grad_bias_acc.add_(chunk_grad_bias)
+            else:
+                argnums_for_grad = (0, 1)  # Differentiate wrt input_chunk, weight
+                (
+                    (chunk_grad_input, chunk_grad_weight),  # No grad for bias
+                    # grad_and_value(has_aux=True) returns (grads, (value, aux));
+                    # value = combined scalar, aux = (soft_loss, ce_loss) from loss_fn_for_grad.
+                    (_, (chunk_kd_loss, chunk_ce_loss)),
+                ) = torch.func.grad_and_value(
+                    loss_fn_for_grad, argnums=argnums_for_grad, has_aux=True
+                )(
+                    student_input_chunk_ac,
+                    student_lm_head_weight,
+                    None,  # Pass None for student_bias primal
+                    target_token_ids_chunk_ac,
+                    target_logprobs_chunk_ac,
+                    target_mask_chunk_ac,
+                    true_labels_chunk_ac,
+                )
+
+            grad_weight_acc.add_(chunk_grad_weight)
+            kd_loss_acc.add_(chunk_kd_loss)
+            ce_loss_acc.add_(chunk_ce_loss)
+
+            return chunk_grad_input
+
+        if compiled:
+            accumulate_chunk_grads_compiled = torch.compile(
+                accumulate_chunk_grads, dynamic=True, backend="inductor"
+            )  # dynamic=True often helpful
+        else:
+            accumulate_chunk_grads_compiled = accumulate_chunk_grads
+
+        # Use the same chunking logic as LigerFusedLinearDistillationBase.forward
+        B, N, D = student_input.shape
+        K = target_token_ids.shape[-1]
+
+        student_input_flat = student_input.reshape(-1, student_input.shape[-1])
+        target_token_ids_flat = target_token_ids.reshape(-1, target_token_ids.shape[-1])
+        target_logprobs_flat = target_logprobs.reshape(-1, target_logprobs.shape[-1])
+        target_mask_flat = target_mask.reshape(-1, target_mask.shape[-1])
+        # pad and shift for cross entropy loss
+        true_labels = torch.nn.functional.pad(true_labels, (0, 1), value=ignore_index)
+        true_labels_flat = true_labels[:, 1:].contiguous().view(-1)
+
+        num_chunks = max(1, student_input_flat.shape[0] // CHUNK_SIZE)
+
+        _student_input_chunks = torch.chunk(
+            student_input_flat, chunks=num_chunks, dim=0
+        )
+        _target_token_ids_chunks = torch.chunk(
+            target_token_ids_flat, chunks=num_chunks, dim=0
+        )
+        _target_logprobs_chunks = torch.chunk(
+            target_logprobs_flat, chunks=num_chunks, dim=0
+        )
+        _target_mask_chunks = torch.chunk(target_mask_flat, chunks=num_chunks, dim=0)
+        _true_labels_chunks = torch.chunk(true_labels_flat, chunks=num_chunks, dim=0)
+
+        for i in range(num_chunks):
+            grad_input_chunk = accumulate_chunk_grads_compiled(
+                _student_input_chunks[i],
+                _target_token_ids_chunks[i],
+                _target_logprobs_chunks[i],
+                _target_mask_chunks[i],
+                _true_labels_chunks[i],
+            )
+            grad_inputs_list.append(grad_input_chunk)
+
+        grad_inputs_combined = torch.cat(grad_inputs_list, dim=0)
+        ctx.save_for_backward(grad_inputs_combined, grad_weight_acc, grad_bias_acc)
+
+        # For matching None returns in backward for non-tensor/non-grad_requiring inputs
+        ctx.hyperparams_count = 9  # Corresponds to number of hyperparams after main tensors in fwd signature
+        ctx.bias_was_none = student_lm_head_bias is None
+        ctx.orig_dims = (B, N, D, K)
+
+        # since this is packed, there is simply a single batch, so batchmean reduction of kl-div is simply the accumulated sum
+        # we still need to scale the kd_loss by the temp^2
+        kd_loss_acc = kd_loss_acc * (temperature**2)
+        final_loss = weight_soft_loss * kd_loss_acc + weight_hard_loss * ce_loss_acc
+
+        return final_loss
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input_flat, grad_weight, grad_bias_maybe = (
+            ctx.saved_tensors
+        )  # grad_input_flat is (B*N, D)
+
+        # Scale gradients by grad_output if it's not 1.0
+        if not torch.equal(
+            grad_output,
+            torch.tensor(1.0, device=grad_output.device, dtype=grad_output.dtype),
+        ):
+            grad_input_flat = grad_input_flat * grad_output
+            grad_weight = grad_weight * grad_output
+            if grad_bias_maybe is not None:
+                grad_bias_maybe = grad_bias_maybe * grad_output
+
+        # Reshape grad_input_flat to match original student_input shape (B, N, D)
+        # ctx.orig_dims stores (B, N, D, K)
+        # We need the first three dimensions for student_input's shape.
+        # Ensure that orig_dims are not (0,0,0,K) for empty inputs leading to view errors
+        if (
+            ctx.orig_dims[0] * ctx.orig_dims[1] * ctx.orig_dims[2] == 0
+            and grad_input_flat.numel() == 0
+        ):
+            # If original input was empty, gradient should also be empty with correct shape
+            grad_input_reshaped = torch.zeros(
+                ctx.orig_dims[0],
+                ctx.orig_dims[1],
+                ctx.orig_dims[2],
+                dtype=grad_input_flat.dtype,
+                device=grad_input_flat.device,
+            )
+        elif grad_input_flat.numel() == 0 and not (
+            ctx.orig_dims[0] * ctx.orig_dims[1] * ctx.orig_dims[2] == 0
+        ):
+            # This case should ideally not happen if forward path is correct (non-empty input -> non-empty flat grad)
+            # but as a safeguard:
+            grad_input_reshaped = torch.zeros(
+                ctx.orig_dims[0],
+                ctx.orig_dims[1],
+                ctx.orig_dims[2],
+                dtype=grad_input_flat.dtype,
+                device=grad_input_flat.device,
+            )
+        else:
+            grad_input_reshaped = grad_input_flat.view(
+                ctx.orig_dims[0], ctx.orig_dims[1], ctx.orig_dims[2]
+            )
+
+        nones_for_hyperparams = [None] * ctx.hyperparams_count
+        grad_bias_return = grad_bias_maybe if not ctx.bias_was_none else None
+
+        return (
+            grad_input_reshaped,  # Gradient for student_input (reshaped)
+            grad_weight,  # Gradient for student_lm_head_weight
+            None,  # Gradient for target_token_ids
+            None,  # Gradient for target_logprobs
+            None,  # Gradient for target_mask
+            None,  # Gradient for true_labels
+            grad_bias_return,  # Gradient for student_lm_head_bias
+            *nones_for_hyperparams,  # Grads for weight_hard_loss, ..., compute_ce_loss
+        )
+
+
+class LigerFusedLinearKLTopKLogprobLoss(torch.nn.Module):
+    """
+    wrapper for chunked top-k logprob kl-d
+    """
+
+    def __init__(
+        self,
+        weight_hard_loss: float = 0.5,
+        weight_soft_loss: float = 0.5,
+        temperature: float = 1.0,  # This is the kd_temperature
+        beta: float = 1.0,
+        ignore_index: int = -100,
+        compiled: bool = True,
+        chunk_size: int = 1024,
+        compute_ce_loss: bool = True,
+        normalize_topk: bool = True,
+    ):
+        super().__init__()
+        if not (0.0 <= weight_hard_loss <= 1.0 and 0.0 <= weight_soft_loss <= 1.0):
+            raise ValueError("Loss weights must be between 0.0 and 1.0.")
+        if temperature <= 0:
+            raise ValueError("Temperature must be positive.")
+
+        self.weight_hard_loss = weight_hard_loss
+        self.weight_soft_loss = weight_soft_loss
+        self.temperature = temperature
+        self.beta = beta
+        self.ignore_index = ignore_index
+        self.compiled = compiled
+        self.chunk_size = chunk_size
+        self.compute_ce_loss = compute_ce_loss
+        self.normalize_topk = normalize_topk
+
+        if not self.compute_ce_loss and self.weight_hard_loss > 0.0:
+            print(
+                f"Warning: compute_ce_loss is False, but weight_hard_loss ({self.weight_hard_loss}) > 0. Hard loss will effectively be zero."
+            )
+            # self.weight_hard_loss = 0.0 # Or let user manage this
+        if self.weight_soft_loss == 0.0:
+            print(
+                "Warning: weight_soft_loss is 0.0. Soft (KD) loss will not be computed."
+            )
+
+    def forward(
+        self,
+        lm_head_weight: torch.Tensor,  # Weights of the linear layer in the LM head
+        student_hidden_states: torch.Tensor,  # student_hidden_states before the lm_head
+        target_token_ids: torch.Tensor,
+        target_logprobs: torch.Tensor,
+        target_mask: torch.Tensor,
+        true_labels: torch.Tensor,
+        student_bias: torch.Tensor = None,
+    ) -> torch.Tensor:
+        return LigerFusedLinearKLTopKLogprobFunction.apply(
+            student_hidden_states,
+            lm_head_weight,
+            target_token_ids,
+            target_logprobs,
+            target_mask,
+            true_labels,
+            student_bias,
+            self.weight_hard_loss,
+            self.weight_soft_loss,
+            self.ignore_index,
+            self.temperature,
+            self.beta,
+            self.compiled,
+            self.chunk_size,
+            self.compute_ce_loss,
+            self.normalize_topk,
+        )
diff --git a/src/axolotl/integrations/kd/topk_logprob/__init__.py b/src/axolotl/integrations/kd/topk_logprob/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/integrations/kd/topk_logprob/forward_kl.py b/src/axolotl/integrations/kd/topk_logprob/forward_kl.py
new file mode 100644
index 0000000000..b79ba26f37
--- /dev/null
+++ b/src/axolotl/integrations/kd/topk_logprob/forward_kl.py
@@ -0,0 +1,173 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+loss for top_k KL divergence
+"""
+
+import torch
+from torch import nn
+
+
+@torch.jit.script
+def loss(
+    student_logits: torch.Tensor,
+    target_token_ids: torch.Tensor,
+    target_logprobs: torch.Tensor,
+    target_mask: torch.Tensor,
+    num_items_in_batch: int = -1,  # Use -1 to indicate "None"
+    kd_temperature: float = 1.0,
+) -> torch.Tensor:
+    """
+    A KD loss function that is TorchScript-friendly.
+
+    Arguments:
+        student_logits (torch.Tensor): The logits of the student model.
+            Shape: [B, student_seq_len, vocab_size]
+        target_token_ids (torch.Tensor): The top-k teacher/target token IDs
+            Shape: [B, teacher_seq_len, top_k]
+        target_logprobs (torch.Tensor): The top-k teacher/target logprobs, these should already be re-normalized.
+            Shape: [B, teacher_seq_len, top_k]
+        target_mask (torch.Tensor): The mask for valid tokens.
+            Shape: [B, teacher_seq_len, top_k]
+        num_items_in_batch (int, optional): The number of items in the batch.
+        kd_temperature (float, optional): The temperature for KD.
+            Default: 1.0
+    """
+
+    target_logprobs = target_logprobs.float()
+
+    # Determine the teacher sequence length
+    # target_token_ids shape: [B, teacher_seq_len, K]
+    # student_logits shape:   [B, student_seq_len, vocab_size]
+    teacher_seq_len = target_token_ids.shape[1]
+
+    # Slice student logits to match teacher-provided sequence length
+    student_logits_for_kd = (
+        student_logits[:, :teacher_seq_len, :] / kd_temperature
+    )  # [B, teacher_seq_len, vocab_size]
+
+    # keep in full precision for numerical stability of loss
+    student_logits_for_kd = student_logits_for_kd.float()
+
+    # Gather student logits for teacher's top-K tokens
+    student_logits_topk = torch.gather(
+        student_logits_for_kd, dim=-1, index=target_token_ids
+    )  # [B, teacher_seq_len, K]
+
+    # Compute logsumexp across full vocabulary
+    student_lse = torch.logsumexp(student_logits_for_kd, dim=-1, keepdim=True)
+
+    #  Convert just the top-k logits to logprobs
+    student_logprobs_topk = student_logits_topk - student_lse
+
+    # Convert teacher_mask to boolean for indexing
+    # In TorchScript, .bool() is sometimes unsupported, so we do:
+    valid_mask = target_mask.to(torch.bool)
+
+    # Prune tensors to only keep valid tokens
+    student_logprobs_topk = student_logprobs_topk[valid_mask]
+    target_logprobs = target_logprobs[valid_mask]
+
+    # Convert teacher logprobs to probabilities
+    teacher_probs = target_logprobs.exp()
+
+    # Compute forward KL
+    kd_loss_per_token = teacher_probs * (target_logprobs - student_logprobs_topk)
+    kd_loss = kd_loss_per_token.sum()
+
+    # Normalize by number of items (if provided) or by valid tokens
+    if num_items_in_batch > 0:
+        kd_loss = kd_loss / float(num_items_in_batch)
+    else:
+        # Fall back to average over valid tokens
+        kd_loss = kd_loss / float(kd_loss_per_token.size(0))
+
+    return kd_loss
+
+
+class ChunkedTopKKDLoss(nn.Module):
+    """
+    A wrapper that chunks (splits) the student and teacher outputs along the time dimension
+    to reduce peak memory usage when upcasting from bf16 to fp32, especially for large vocabularies.
+
+    Usage is analogous to ForwardKLWithChunkedOutputLoss but adapted to top-K teacher logprobs.
+    """
+
+    def __init__(self, num_output_chunks: int = 8, kd_temperature: float = 1.0):
+        super().__init__()
+        self.num_output_chunks = num_output_chunks
+        self.kd_temperature = kd_temperature
+
+    def forward(
+        self,
+        student_logits: torch.Tensor,  # [B, seq_len, vocab_size]
+        target_token_ids: torch.Tensor,  # [B, seq_len, K]
+        target_logprobs: torch.Tensor,  # [B, seq_len, K]
+        target_mask: torch.Tensor,  # [B, seq_len, K]
+        num_items_in_batch: int = -1,  # optional batch size for normalization
+    ) -> torch.Tensor:
+        # 1. Split along the "token" dimension (dim=1).
+        student_logits_chunks = student_logits.chunk(self.num_output_chunks, dim=1)
+        token_ids_chunks = target_token_ids.chunk(self.num_output_chunks, dim=1)
+        logprobs_chunks = target_logprobs.chunk(self.num_output_chunks, dim=1)
+        mask_chunks = target_mask.chunk(self.num_output_chunks, dim=1)
+
+        # We'll accumulate a global "sum of losses" and "sum of valid tokens"
+        # so that our final average is consistent with the entire sequence/batch.
+        total_loss = 0.0
+        total_valid_tokens = 0
+
+        # 2. Loop over each chunk and compute a chunk-specific loss.
+        for st_chunk, tid_chunk, lp_chunk, msk_chunk in zip(
+            student_logits_chunks,
+            token_ids_chunks,
+            logprobs_chunks,
+            mask_chunks,
+            strict=False,
+        ):
+            # We pass num_items_in_batch=-1 so that the kd_loss
+            # will average over *this chunk's* valid tokens only.
+            chunk_loss = loss(
+                student_logits=st_chunk,
+                target_token_ids=tid_chunk,
+                target_logprobs=lp_chunk,
+                target_mask=msk_chunk,
+                num_items_in_batch=-1,  # ensure per-chunk averaging by valid tokens
+                kd_temperature=self.kd_temperature,
+            )
+
+            # kd_loss returns an average over the chunk's valid tokens.
+            # We want a global average in the end, so we need to re‐weight
+            # by the number of valid tokens in this chunk and keep track of the total.
+            chunk_valid_mask = msk_chunk.to(torch.bool)
+            chunk_valid_count = chunk_valid_mask.sum()  # scalar tensor
+
+            # Re-scale "chunk average" back to "chunk sum"
+            chunk_loss_sum = chunk_loss * chunk_valid_count
+
+            total_loss += chunk_loss_sum
+            total_valid_tokens += chunk_valid_count
+
+        # 3. Normalize *once* at the end.
+        if num_items_in_batch > 0:
+            # If the user gave us a manual denominator (e.g. total items in batch),
+            # we divide by it. Typically used if each item is of different length.
+            final_loss = total_loss / float(num_items_in_batch)
+        else:
+            # Otherwise, divide by total valid tokens across all chunks.
+            # to get the same result as a non-chunked approach.
+            final_loss = total_loss / float(total_valid_tokens)
+
+        return final_loss
diff --git a/src/axolotl/integrations/kd/trainer.py b/src/axolotl/integrations/kd/trainer.py
new file mode 100644
index 0000000000..e5d1504eb1
--- /dev/null
+++ b/src/axolotl/integrations/kd/trainer.py
@@ -0,0 +1,121 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+KD trainer
+"""
+
+import torch.nn as nn
+from typing_extensions import override
+
+from axolotl.core.trainers.base import AxolotlTrainer
+
+from .kernels.liger import LigerFusedLinearKLTopKLogprobLoss
+
+
+def _resolve_lm_head(model: nn.Module) -> nn.Module:
+    base = model
+    if hasattr(base, "get_base_model"):
+        base = base.get_base_model()
+    if hasattr(base, "language_model") and hasattr(base.language_model, "lm_head"):
+        return base.language_model.lm_head
+    if hasattr(base, "lm_head"):
+        return base.lm_head
+    raise AttributeError(f"could not find lm_head on {type(model).__name__}")
+
+
+class AxolotlKDTrainer(AxolotlTrainer):
+    """
+    Custom trainer subclass for Knowledge Distillation (KD)
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model_accepts_loss_kwargs = True
+
+        self._kd_loss_fn = LigerFusedLinearKLTopKLogprobLoss(
+            self.args.kd_ce_alpha,  # hard label loss
+            self.args.kd_alpha,  # kd loss
+            self.args.kd_temperature,
+            self.args.kd_beta or 0.0,
+            compute_ce_loss=bool(self.args.kd_ce_alpha),
+            normalize_topk=self.args.kd_normalize_topk,
+        )
+
+    def _set_signature_columns_if_needed(self):
+        super()._set_signature_columns_if_needed()
+        columns_to_add = []
+        if self._signature_columns:
+            if "target_logprobs" not in self._signature_columns:
+                columns_to_add.append("target_logprobs")
+            if "target_token_ids" not in self._signature_columns:
+                columns_to_add.append("target_token_ids")
+            if "target_mask" not in self._signature_columns:
+                columns_to_add.append("target_mask")
+            if columns_to_add:
+                self._signature_columns += columns_to_add
+
+    @override
+    def compute_loss(
+        self,
+        model,
+        inputs,
+        return_outputs=False,
+        num_items_in_batch=None,
+    ):
+        inputs = dict(inputs)
+
+        required_keys = ("labels", "target_token_ids", "target_logprobs", "target_mask")
+        missing = [k for k in required_keys if k not in inputs]
+        if missing:
+            raise KeyError(f"KD batch missing required keys: {missing}")
+
+        if num_items_in_batch is None and "labels" in inputs:
+            num_items_in_batch = (inputs["labels"] != -100).sum().item()
+
+        labels = inputs.pop("labels")
+        target_token_ids = inputs.pop("target_token_ids")
+        target_logprobs = inputs.pop("target_logprobs")
+        target_mask = inputs.pop("target_mask")
+
+        # num_items_in_batch is a loss kwarg, not a forward kwarg.
+        inputs.pop("num_items_in_batch", None)
+
+        inputs["output_hidden_states"] = True
+        inputs["return_dict"] = True
+        inputs["logits_to_keep"] = 1
+        outputs = model(**inputs)
+        hidden_states = getattr(outputs, "hidden_states", None)
+        if hidden_states is None:
+            raise RuntimeError(
+                f"{type(model).__name__}.forward did not return hidden_states"
+            )
+        hidden_states = hidden_states[-1]
+
+        lm_head = _resolve_lm_head(model)
+        hidden_states = hidden_states.to(lm_head.weight.dtype)
+
+        loss = self._kd_loss_fn(
+            lm_head.weight,
+            hidden_states,
+            target_token_ids,
+            target_logprobs,
+            target_mask,
+            true_labels=labels,
+        )
+
+        if num_items_in_batch is not None and num_items_in_batch > 0:
+            loss = loss / num_items_in_batch
+
+        return (loss, outputs) if return_outputs else loss
diff --git a/src/axolotl/integrations/kd/utils.py b/src/axolotl/integrations/kd/utils.py
new file mode 100644
index 0000000000..ba60694a56
--- /dev/null
+++ b/src/axolotl/integrations/kd/utils.py
@@ -0,0 +1,100 @@
+"""Helper KD utils"""
+
+import math
+from typing import List, Union
+
+import numpy as np
+import torch
+from torch import FloatTensor, Tensor
+
+
+def normalize_logprobs(logprobs: FloatTensor, topk: int) -> FloatTensor:
+    """
+    Re-normalizes top-k raw logprobs as probabilities, and converts back to logprobs.
+    """
+    # Ensure raw_logprobs matches kd_online_topk length for tensor operations
+    # This should ideally be handled by the caller ensuring correct padding/truncation first
+    if logprobs.shape[-1] != topk:
+        # pad last dimension of logprobs to match topk length with -inf
+        padding_len = topk - logprobs.shape[-1]
+        padding_tensor = torch.full(
+            (
+                *logprobs.shape[:-1],
+                padding_len,
+            ),  # Takes all dimensions of logprobs except the last, then appends padding_needed
+            float("-inf"),
+            dtype=logprobs.dtype,
+            device=logprobs.device,
+        )
+        logprobs = torch.cat((logprobs, padding_tensor), dim=-1)
+
+    # Convert logprobs at T_online to probabilities
+    # use log sum exp trick to avoid underflow
+    position_logprobs_lse = torch.logsumexp(logprobs, dim=-1, keepdim=True)
+    teacher_probs_t_online = torch.exp(logprobs - position_logprobs_lse)
+
+    # Normalize probabilities (sum to 1)
+    # This is important if the top-k from server aren't a full distribution
+    teacher_probs_t_online_sum = teacher_probs_t_online.sum(dim=-1, keepdim=True)
+    teacher_probs_t_online = teacher_probs_t_online / teacher_probs_t_online_sum
+
+    final_logprobs_tensor = torch.log(teacher_probs_t_online)
+
+    return final_logprobs_tensor
+
+
+def strided_chunk_views(
+    tensor: Union[np.ndarray, torch.Tensor],
+    chunks: int,
+    dim: int = 0,
+    stride: int = 1,
+    chunk_size: int | None = None,
+) -> List[Union[np.ndarray, torch.Tensor]]:
+    """
+    Split a tensor into chunks along a dimension with striding, prioritizing views over copies.
+
+    Args:
+        tensor: Input tensor (numpy array or torch tensor)
+        chunks: Number of chunks to create
+        dim: Dimension along which to chunk (default: 0)
+        stride: Stride between chunk starting positions (default: 1)
+        chunk_size: Size of each chunk. If None, calculated automatically (default: None)
+
+    Returns:
+        List of tensor chunks (views when possible, copies when necessary)
+    """
+
+    # Get the size of the specified dimension
+    dim_size = tensor.shape[dim]
+
+    # Calculate chunk size if not provided
+    if chunk_size is None:
+        chunk_size = (dim_size + chunks - 1) // chunks  # Ceiling division
+
+    chunks_list = []
+
+    for i in range(chunks):
+        start_idx = i * stride
+        end_idx = min(start_idx + chunk_size, dim_size)
+
+        # Break if we've gone beyond the tensor
+        if start_idx >= dim_size:
+            break
+
+        # Create slice objects for all dimensions
+        slices = [slice(None)] * tensor.ndim
+        slices[dim] = slice(start_idx, end_idx)
+
+        chunk = tensor[tuple(slices)]
+        chunks_list.append(chunk)
+
+    return chunks_list
+
+
+def chunk_overlap(input_tensor: Tensor, chunks: int, dim: int = 0, overlap: int = 1):
+    dim_size = input_tensor.shape[dim]
+    stride = math.ceil(dim_size / chunks)
+
+    return strided_chunk_views(
+        input_tensor, chunks, dim, stride=stride, chunk_size=stride + overlap
+    )
diff --git a/src/axolotl/integrations/kernels/README.md b/src/axolotl/integrations/kernels/README.md
new file mode 100644
index 0000000000..ff8e4aee5e
--- /dev/null
+++ b/src/axolotl/integrations/kernels/README.md
@@ -0,0 +1,147 @@
+# Kernels Integration
+
+MoE (Mixture of Experts) kernels speed up training for MoE layers and reduce VRAM costs. Transformers v5 introduced a uniform dispatch point for the per-expert grouped GEMMs via the `experts_implementation` config kwarg:
+
+```python
+class ExpertsInterface(GeneralInterface):
+    _global_mapping = {
+        "batched_mm": batched_mm_experts_forward,
+        "grouped_mm": grouped_mm_experts_forward,
+        "sonicmoe":   sonicmoe_experts_forward,   # upstream HF integration
+    }
+```
+
+Axolotl registers two additional implementations into this same global registry: **ScatterMoE** (Triton, runs on any CUDA GPU) and a LoRA-aware **SonicMoE** variant (CUTLASS / cute-DSL, Hopper or newer). Routing — softmax/sigmoid top-k, group selection, shared experts, bias correction, etc. — stays in each model's `SparseMoEBlock`, where transformers handles all per-architecture variation. Axolotl only swaps the experts forward.
+
+## Usage
+
+Add the following to your axolotl YAML config:
+
+```yaml
+plugins:
+  - axolotl.integrations.kernels.KernelsPlugin
+
+use_kernels: true
+
+# Choose one (mutually exclusive):
+use_scattermoe: true
+# OR
+use_sonicmoe: true
+```
+
+`experts_implementation` is auto-set to `scattermoe` / `sonicmoe` from the kernel flag, but you can override to `eager` / `batched_mm` / `grouped_mm` to compare against the transformers reference implementations.
+
+### SonicMoE installation
+
+**Prerequisites:**
+- NVIDIA Hopper (H100/H200) or Blackwell (B200/GB200/B300) GPU
+- CUDA 12.9+ (13.0+ for B300)
+- PyTorch 2.7+
+- For B300: Triton 3.6.x
+
+The sonic-moe kernel ships through the HF [`kernels`](https://github.com/huggingface/kernels) package. Transformers v5.8+ auto-fetches a prebuilt kernel from [`kernels-community/sonic-moe`](https://huggingface.co/kernels-community/sonic-moe) on first use:
+
+```bash
+pip install kernels "nvidia-cutlass-dsl==4.5.2"
+```
+
+**Note:** Blackwell support is in upstream beta. On Blackwell GPUs Axolotl automatically sets `USE_QUACK_GEMM=1` to enable the Blackwell kernels.
+
+
+## How It Works
+
+The `KernelsPlugin` runs once before model loading and:
+
+1. Calls `register_scattermoe_experts()` or `register_sonicmoe_experts()`, which inserts the kernel forward into `transformers.integrations.moe.ALL_EXPERTS_FUNCTIONS`.
+2. Sets `cfg.experts_implementation` to the matching name.
+3. When the model loads, transformers' `@use_experts_implementation` decorator on each model's `Experts` class reads `config._experts_implementation` and dispatches to our registered forward.
+
+That's the entire integration — there is no per-architecture SparseMoEBlock monkey-patch, no per-model routing code, and no weight-layout conversion. As new MoE models adopt the decorator upstream they immediately benefit from both kernels.
+
+## BF16 LoRA Support
+
+Both kernels train PEFT adapters on `gate_up_proj` / `down_proj` (and `gate` for the router) end-to-end:
+
+- **ScatterMoE** fuses the LoRA `B @ A` product into the per-expert grouped GEMM via custom Triton kernels (`parallel_linear_lora`). No extra materialization pass.
+- **SonicMoE** materializes `W_eff = W + scaling * (B @ A)` per expert inside a custom `MoELoRAMaterialize` `autograd.Function` and passes the effective weight into the CUTLASS kernel. Backward decomposes `dW_eff` into `dA` and `dB` via the chain rule, so LoRA parameters train without modifying the kernel.
+
+Both paths detect PEFT `ParamWrapper` on individual expert parameters (`target_parameters` API) and unwrap them before dispatch.
+
+### ScatterMoE NVFP4 (W4A16) LoRA
+
+Train LoRA on ModelOpt NVFP4 checkpoint via ScatterMoE. Routed experts are dequantized to bf16 (W4A16).
+
+Requires:
+- CUDA GPU with Triton.
+- `qwen3_moe`, `qwen3_next`, `deepseek_v4`, `glm_moe_dsa`, and `gemma4_text`
+
+**Tip:** in our tests, the Triton `dequant` path below is currently faster end to end than the ScatterMoE NVFP4 (if the arch is supported).
+
+### SonicMoE NVFP4 (W4A4) LoRA
+
+Train LoRA on ModelOpt NVFP4 checkpoint via Quack (e.g. `nvidia/Qwen3-30B-A3B-NVFP4`).
+
+Requires:
+- Blackwell SM100 for W4A4, others for W4A16
+- `qwen3_moe` / `qwen3_next`
+
+Install the pinned quack kernels (other versions untested):
+
+```bash
+pip install "quack-kernels==0.5.0" "nvidia-cutlass-dsl==4.5.2"
+```
+
+`AXOLOTL_SONICMOE_NVFP4_BACKEND` picks the expert GEMM (unset = auto):
+
+| Backend | Compute | Runs on |
+|---------|---------|---------|
+| `fp4_cute` | native W4A4 tensor cores (quack) | Blackwell B200/GB200 (SM100/110) |
+| `dequant` | dequant to bf16 (W4A16) | any CUDA GPU with Triton |
+
+Advanced tuning knobs (fused up-proj, per-tensor-scale fold, fp8 DeepGEMM backward) are exposed as other `AXOLOTL_SONICMOE_NVFP4_*` env vars; defaults are correct for normal training.
+
+**Tip:** in our tests, the Triton `dequant` path is currently faster (and lower mem) end to end than the `fp4_cute` path for < 100B param MoE model. However, `fp4_cute` should overtake when expert matmul become more dominant cost.
+
+## Model Support
+
+Any model whose `Experts` class is decorated with `@use_experts_implementation` upstream works automatically. As of transformers 5.8 this includes (verified). The `bf16` columns are base (unquantized) support; the `NVFP4` columns mark which kernel trains LoRA on a ModelOpt NVFP4 checkpoint of that arch:
+
+| Model Type        | ScatterMoE (bf16) | SonicMoE (bf16) | NVFP4 (ScatterMoE) | NVFP4 (SonicMoE) |
+|-------------------|:---------:|:--------:|:---:|:---:|
+| `mixtral`         |    Yes    |   Yes    |  -  |  -  |
+| `qwen2_moe`       |    Yes    |   Yes    |  -  |  -  |
+| `qwen3_moe`       |    Yes    |   Yes    | Yes | Yes |
+| `qwen3_next`      |    Yes    |   Yes    | Yes | Yes |
+| `qwen3_5_moe`     |    Yes    |   Yes    |  -  |  -  |
+| `olmoe`           |    Yes    |   Yes    |  -  |  -  |
+| `mistral4`        |    Yes    |   Yes    |  -  |  -  |
+| `glm_moe_dsa`     |    Yes    |   Yes    | Yes |  -  |
+| `deepseek_v3`     |    Yes    |   Yes    |  -  |  -  |
+| `minimax_m2`      |    Yes    |   Yes    |  -  |  -  |
+| `ernie4_5_moe`    |    Yes    |   Yes    |  -  |  -  |
+| `hunyuan_v1_moe`  |    Yes    |   Yes    |  -  |  -  |
+| `gemma4_text`     |    Yes    |   Yes    | Yes |  -  |
+| `gpt_oss`         |    Yes    |   Yes    |  -  |  -  |
+
+NVFP4 for `deepseek_v4` is supported via ScatterMoE with `use_dsv4_kernels` (its own fused-kernel path), so it is not a row above.
+
+`gpt_oss` carries the decorator with `is_concatenated=False, is_transposed=True, has_bias=True` and uses a sigmoid-GLU activation with clamping. Both forwards read these flags off `self` and dispatch accordingly: the ScatterMoE forward handles the transposed/interleaved/biased layout and clamped sigmoid-GLU via its Triton path (no weight transpose, interleaved gate/up, per-expert bias folded into the grouped GEMM); the SonicMoE forward uses the upstream CUTLASS kernel.
+
+### Blackwell (sm_120) note
+
+`use_sonicmoe` errors at load on consumer Blackwell (sm_120): the published `kernels-community/sonic-moe` prebuilt bundles quack 0.3.11, which lacks sm_120 support (we are waiting for a rebuild against quack >= 0.4.0). Use `use_scattermoe` instead, including for `gpt_oss`.
+
+## Feature comparison
+
+| Feature                          | ScatterMoE | SonicMoE |
+|----------------------------------|:----------:|:--------:|
+| Kernel backend                   | Triton     | CUTLASS / cute-DSL |
+| GPU requirement                  | Any CUDA   | Hopper+ |
+| LoRA path                        | Fused in Triton kernel | `MoELoRAMaterialize` + custom autograd |
+| LoRA overhead                    | Lower (fused) | Higher (materialization pass) |
+| Selective expert dequantization  | Yes (~97% memory savings) | No |
+| Weight format                    | Standard `[E, 2*I, H]` | Standard `[E, 2*I, H]` (concat layout, no interleave) |
+
+## Note on MegaBlocks
+
+We tested [MegaBlocks](https://huggingface.co/kernels-community/megablocks) but were unable to ensure numerical accuracy, so we did not integrate it. It was also incompatible with many newer model architectures in transformers.
diff --git a/src/axolotl/integrations/kernels/__init__.py b/src/axolotl/integrations/kernels/__init__.py
new file mode 100644
index 0000000000..d879424359
--- /dev/null
+++ b/src/axolotl/integrations/kernels/__init__.py
@@ -0,0 +1,7 @@
+from .args import KernelsArgs
+from .plugin import KernelsPlugin
+
+__all__ = [
+    "KernelsArgs",
+    "KernelsPlugin",
+]
diff --git a/src/axolotl/integrations/kernels/adapters/__init__.py b/src/axolotl/integrations/kernels/adapters/__init__.py
new file mode 100644
index 0000000000..a2e045ae75
--- /dev/null
+++ b/src/axolotl/integrations/kernels/adapters/__init__.py
@@ -0,0 +1,68 @@
+"""Model adapters for the kernels plugin.
+
+``KernelsPlugin`` orchestrates plugin hooks and generic capabilities (expert-kernel
+registration, the quantized-training guard, grouped NVFP4 MoE dispatch). Model-family
+specifics (DeepSeek-V4 fused kernels / quantizer / dtype policy, Gemma-4 NVFP4 converters /
+non-expert quantization) live in ``ModelAdapter`` subclasses here, so the plugin stays a thin
+orchestrator and new models opt in by adding an adapter rather than another inline branch.
+"""
+
+from __future__ import annotations
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class ModelAdapter:
+    """Base class for model-family kernel adapters. All hooks are no-ops by default.
+
+    An adapter is *active* for a run when :meth:`matches` returns True. The plugin calls the
+    hooks of every active adapter at the corresponding lifecycle point.
+    """
+
+    #: short name for logging
+    name: str = "base"
+
+    def matches(self, cfg) -> bool:  # pragma: no cover - trivial
+        return False
+
+    def consumes_nonexpert_quantization(self, cfg) -> bool:
+        """Whether this adapter acts on the ``nonexpert_quantization`` intent for ``cfg``.
+
+        The plugin warns when a non-expert quantization policy is configured but no active
+        adapter claims it (so a future model can't silently no-op the setting). Default False.
+        """
+        return False
+
+    def pre_model_load(self, cfg) -> None:
+        """Before the model is constructed (register converters/quantizers, patch modules)."""
+
+    def pre_lora_load(self, cfg, model) -> None:
+        """After the model loads, before PEFT wraps it (fix expert loading, quantize non-experts)."""
+
+    def post_model_load(self, cfg, model) -> None:
+        """After PEFT wraps the model (dtype policy, fused-LoRA kernel swaps)."""
+
+
+def _all_adapters() -> list[ModelAdapter]:
+    from axolotl.integrations.kernels.adapters.dsv4 import DSV4Adapter
+    from axolotl.integrations.kernels.adapters.gemma4 import Gemma4Adapter
+    from axolotl.integrations.kernels.adapters.glm_moe_dsa import GlmMoeDsaAdapter
+    from axolotl.integrations.kernels.adapters.qwen3_moe import Qwen3MoeAdapter
+
+    return [DSV4Adapter(), Gemma4Adapter(), GlmMoeDsaAdapter(), Qwen3MoeAdapter()]
+
+
+def get_active_adapters(cfg) -> list[ModelAdapter]:
+    """Return the adapters whose ``matches(cfg)`` is True (order = registration order)."""
+    active = []
+    for adapter in _all_adapters():
+        try:
+            if adapter.matches(cfg):
+                active.append(adapter)
+        except Exception:  # pragma: no cover - matching must never break loading
+            LOG.debug("adapter %s match check failed; skipping", adapter.name)
+    if active:
+        LOG.info("kernels: active model adapters: %s", [a.name for a in active])
+    return active
diff --git a/src/axolotl/integrations/kernels/adapters/dsv4.py b/src/axolotl/integrations/kernels/adapters/dsv4.py
new file mode 100644
index 0000000000..bc064f678f
--- /dev/null
+++ b/src/axolotl/integrations/kernels/adapters/dsv4.py
@@ -0,0 +1,181 @@
+"""DeepSeek-V4 kernel adapter.
+
+Owns everything DSV4-specific: the NVFP4/FP8 quantizer install, fused attention/RoPE/mHC
+kernel patching, the post-PEFT fp32->compute dtype policy (keeping the mHC/keep_in_fp32 modules
+in fp32), the fused clamped-SwiGLU shared-expert MLP LoRA patch, and the fp8 attention LoRA
+patch. Also registers the mHC module class names that FSDP2's quantized-mixed-dtype path shards
+in their own fp32 group.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from axolotl.integrations.kernels.adapters import ModelAdapter
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# mHC modules kept in fp32 and sharded separately by the FSDP2 quantized path.
+DSV4_FP32_SHARD_CLASSES = ("DeepseekV4HyperConnection", "DeepseekV4HyperHead")
+
+
+def _maybe_truncate_layers(model) -> None:
+    """Debug-only (``DSV4_TRUNCATE_LAYERS=N``): truncate the decoder stack to the first N layers
+    before PEFT so a small slice of the model fits a couple of GPUs for fast local bring-up /
+    memory iteration. The first layers span all attention types (sliding/CSA/HCA). No-op unless the
+    env var is set. Keeps HF's ``len(layer_types) == num_hidden_layers`` validators happy by
+    truncating per-layer config lists (length ``orig`` -> n; ``orig+1`` e.g. compress_ratios -> n+1)."""
+    import os
+
+    trunc = os.environ.get("DSV4_TRUNCATE_LAYERS")
+    if not trunc:
+        return
+    import gc
+
+    import torch.nn as nn
+
+    n = int(trunc)
+    for mod in model.modules():
+        layers = getattr(mod, "layers", None)
+        if (
+            isinstance(layers, nn.ModuleList)
+            and len(layers) > n
+            and "DecoderLayer" in type(layers[0]).__name__
+        ):
+            orig = len(layers)
+            mod.layers = layers[:n]
+            for cfg_obj in (model.config, getattr(mod, "config", None)):
+                if cfg_obj is None:
+                    continue
+                if hasattr(cfg_obj, "num_hidden_layers"):
+                    cfg_obj.num_hidden_layers = n
+                for k, v in list(vars(cfg_obj).items()):
+                    if isinstance(v, list) and len(v) in (orig, orig + 1):
+                        setattr(cfg_obj, k, v[: n + (len(v) - orig)])
+            gc.collect()
+            torch.cuda.empty_cache()
+            LOG.warning("DSV4_TRUNCATE_LAYERS=%d: truncated decoder stack (debug)", n)
+            break
+
+
+class DSV4Adapter(ModelAdapter):
+    name = "deepseek_v4"
+
+    def matches(self, cfg) -> bool:
+        return bool(cfg.use_dsv4_kernels)
+
+    def pre_model_load(self, cfg) -> None:
+        # NVFP4 MoE checkpoints declare `quant_method: fp8` but transformers' finegrained-FP8
+        # quantizer has no NVFP4 expert path. Install an NVFP4-aware subclass so experts load as
+        # NVFP4Tensor for the scattermoe fused path. Only relevant with the scattermoe expert path.
+        if cfg.use_scattermoe:
+            from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_fp8_quantizer import (
+                configure_nonexpert_mode,
+                install_nvfp4_fp8_quantizer,
+            )
+
+            configure_nonexpert_mode(cfg.get("dsv4_fp8_nonexpert_mode"))
+            install_nvfp4_fp8_quantizer()
+
+        # Fused Triton training kernels for DSV4 (attention / RoPE / mHC).
+        from axolotl.integrations.kernels.libs.dsv4 import patch_deepseek_v4_kernels
+
+        patch_deepseek_v4_kernels()
+        # Make the FSDP2 quantized path keep these mHC modules in their own fp32 shard group.
+        from axolotl.monkeypatch.accelerate.fsdp2_quantized import (
+            register_fp32_shard_classes,
+        )
+
+        register_fp32_shard_classes(DSV4_FP32_SHARD_CLASSES)
+
+    def pre_lora_load(self, cfg, model) -> None:
+        # MIXED_PRECISION checkpoints load NVFP4 experts as plain packed uint8 with dropped scales;
+        # rebuild them as NVFP4Tensor so scattermoe dequantizes correctly.
+        if not cfg.use_scattermoe:
+            return
+        _maybe_truncate_layers(model)
+        has_packed_experts = any(
+            isinstance(getattr(m, "gate_up_proj", None), torch.Tensor)
+            and m.gate_up_proj.dtype == torch.uint8
+            and m.gate_up_proj.ndim == 3
+            for m in model.modules()
+        )
+        if has_packed_experts:
+            from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_moe_loading import (
+                attach_nvfp4_expert_scales,
+            )
+
+            attach_nvfp4_expert_scales(model, cfg.base_model)
+
+    def post_model_load(self, cfg, model) -> None:
+        self._apply_fp32_dtype_policy(cfg, model)
+        self._patch_shared_mlp_lora(cfg, model)
+        self._patch_attn_fp8_lora(cfg, model)
+
+    # --- dtype policy -------------------------------------------------------
+    @staticmethod
+    def _apply_fp32_dtype_policy(cfg, model) -> None:
+        """Cast residual fp32 params (PEFT-upcast LoRA) to the compute dtype while keeping the
+        model's ``_keep_in_fp32_modules[_strict]`` params (mHC/norms/sinks) in fp32. The fused
+        DSV4 kernels demote fp32 activations at their input boundary, so this gives one consistent
+        compute dtype. ``dsv4_bf16_all: true`` reverts to a blanket cast including keep_in_fp32."""
+        dt = cfg.torch_dtype or torch.bfloat16
+        keep_all = bool(cfg.get("dsv4_bf16_all"))
+        keep_patterns: list[str] = []
+        if not keep_all:
+            seen: set[str] = set()
+            for m in model.modules():
+                for attr in ("_keep_in_fp32_modules_strict", "_keep_in_fp32_modules"):
+                    for pat in getattr(m, attr, None) or ():
+                        if pat not in seen:
+                            seen.add(pat)
+                            keep_patterns.append(pat)
+        n = kept = 0
+        for name, p in model.named_parameters():
+            if p.dtype != torch.float32:
+                continue
+            if keep_patterns and any(pat in name for pat in keep_patterns):
+                kept += 1
+                continue
+            p.data = p.data.to(dt)
+            n += 1
+        if n:
+            LOG.info(
+                "dsv4: cast %d residual fp32 params to %s for fused kernels (kept %d keep_in_fp32 fp32)",
+                n,
+                dt,
+                kept,
+            )
+
+    # --- fused LoRA kernel swaps -------------------------------------------
+    @staticmethod
+    def _patch_shared_mlp_lora(cfg, model) -> None:
+        """Swap DSV4 shared-expert MLPs for the fused clamped-SwiGLU LoRA kernel.
+
+        Gated by ``dsv4_shared_mlp_lora_kernel`` (NOT the generic ``lora_mlp_kernel``, which the
+        MoE-kernel validator force-disables). The validator translates a legacy
+        ``lora_mlp_kernel: true`` on a DSV4 run into this flag (see KernelsArgs.disable_mlp_kernel).
+        """
+        if cfg.get("adapter") != "lora" or not cfg.get("dsv4_shared_mlp_lora_kernel"):
+            return
+        from axolotl.integrations.kernels.libs.dsv4.lora_mlp import (
+            patch_dsv4_shared_mlp_lora,
+        )
+
+        n = patch_dsv4_shared_mlp_lora(model)
+        LOG.info(
+            "dsv4: patched %d shared-expert MLPs with fused clamped-SwiGLU LoRA", n
+        )
+
+    @staticmethod
+    def _patch_attn_fp8_lora(cfg, model) -> None:
+        """Native blockwise-fp8 fused LoRA for the large attention projections (q_b/o_b).
+        No-op unless ``dsv4_fp8_lora_kernel``."""
+        if cfg.get("adapter") != "lora":
+            return
+        from axolotl.integrations.kernels.libs.dsv4.lora_fp8 import (
+            patch_dsv4_attn_fp8_lora,
+        )
+
+        patch_dsv4_attn_fp8_lora(model, enabled=bool(cfg.get("dsv4_fp8_lora_kernel")))
diff --git a/src/axolotl/integrations/kernels/adapters/gemma4.py b/src/axolotl/integrations/kernels/adapters/gemma4.py
new file mode 100644
index 0000000000..76e629166a
--- /dev/null
+++ b/src/axolotl/integrations/kernels/adapters/gemma4.py
@@ -0,0 +1,116 @@
+"""Gemma-4 NVFP4 kernel adapter.
+
+Owns Gemma-4 specifics: detecting the NVFP4-modelopt checkpoint, registering the native NVFP4
+expert WeightConverters, and applying the non-expert quantization policy (fp8 / nf4). The hybrid
+global-attention mask and the large-head flash kernel are generic capabilities wired in the
+loader/patch_manager and are intentionally NOT owned here.
+"""
+
+from __future__ import annotations
+
+from axolotl.integrations.kernels.adapters import ModelAdapter
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def is_gemma4_nvfp4_modelopt(cfg) -> bool:
+    """True iff the base model is a Gemma-4 NVFP4-modelopt checkpoint (quant_method=modelopt,
+    quant_algo=NVFP4, model_type startswith gemma4). Any failure returns False."""
+    try:
+        from transformers import AutoConfig
+
+        hf_cfg = AutoConfig.from_pretrained(cfg.base_model, trust_remote_code=True)
+    except Exception:
+        return False
+    if not str(getattr(hf_cfg, "model_type", "")).startswith("gemma4"):
+        return False
+    qcfg = getattr(hf_cfg, "quantization_config", None)
+    if qcfg is None:
+        return False
+    if isinstance(qcfg, dict):
+        return (
+            qcfg.get("quant_method") == "modelopt" and qcfg.get("quant_algo") == "NVFP4"
+        )
+    return (
+        getattr(qcfg, "quant_method", None) == "modelopt"
+        and getattr(qcfg, "quant_algo", None) == "NVFP4"
+    )
+
+
+def resolve_nonexpert_quantization(cfg) -> str | None:
+    """Resolve the non-expert quantization policy to one of {None, 'fp8', 'nf4', 'nvfp4'}.
+
+    Prefers the intent field ``nonexpert_quantization`` (none/bf16/fp8_blockwise/nf4/nvfp4); falls
+    back to the legacy gemma4-specific flags. Returns None for no quantization (bf16 non-experts).
+    'nvfp4' routes non-experts through the Marlin W4A16 kernel (same path as the experts); 'nf4'
+    uses bitsandbytes.
+    """
+    policy = cfg.get("nonexpert_quantization")
+    if policy:
+        p = str(policy).lower()
+        if p in ("none", "bf16"):
+            return None
+        if p in ("fp8", "fp8_blockwise"):
+            return "fp8"
+        if p == "nf4":
+            return "nf4"
+        if p == "nvfp4":
+            return "nvfp4"
+    if cfg.get("gemma4_fp8_nonexpert"):
+        return "fp8"
+    if cfg.get("gemma4_nf4_nonexpert"):
+        return "nf4"
+    return None
+
+
+class Gemma4Adapter(ModelAdapter):
+    name = "gemma4_nvfp4"
+
+    def matches(self, cfg) -> bool:
+        return bool(cfg.use_scattermoe) and is_gemma4_nvfp4_modelopt(cfg)
+
+    def consumes_nonexpert_quantization(self, cfg) -> bool:
+        return True
+
+    def pre_model_load(self, cfg) -> None:
+        from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_weight_converter import (
+            register_gemma4_nvfp4_converters,
+        )
+
+        register_gemma4_nvfp4_converters()
+        LOG.info(
+            "gemma4: registered NVFP4 expert WeightConverters (modelopt checkpoint)"
+        )
+
+    def pre_lora_load(self, cfg, model) -> None:
+        policy = resolve_nonexpert_quantization(cfg)
+        if policy == "fp8":
+            from axolotl.integrations.kernels.libs.scattermoe_lora.gemma4_fp8_nonexpert import (
+                quantize_gemma4_nonexpert_linears,
+            )
+
+            n = quantize_gemma4_nonexpert_linears(model)
+            LOG.info(
+                "gemma4: fp8-quantized %d non-expert linears (experts stay NVFP4)", n
+            )
+        elif policy == "nf4":
+            from axolotl.integrations.kernels.libs.scattermoe_lora.gemma4_nf4_nonexpert import (
+                quantize_gemma4_nonexpert_nf4,
+            )
+
+            n = quantize_gemma4_nonexpert_nf4(model)
+            LOG.info(
+                "gemma4: nf4-quantized %d non-expert linears (experts stay NVFP4)", n
+            )
+        elif policy == "nvfp4":
+            from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_nonexpert import (
+                quantize_gemma4_nonexpert_nvfp4,
+            )
+
+            n = quantize_gemma4_nonexpert_nvfp4(model)
+            LOG.info(
+                "gemma4: nvfp4-quantized %d non-expert linears via Marlin W4A16 "
+                "(experts also NVFP4)",
+                n,
+            )
diff --git a/src/axolotl/integrations/kernels/adapters/glm_moe_dsa.py b/src/axolotl/integrations/kernels/adapters/glm_moe_dsa.py
new file mode 100644
index 0000000000..5d1d99dd1f
--- /dev/null
+++ b/src/axolotl/integrations/kernels/adapters/glm_moe_dsa.py
@@ -0,0 +1,225 @@
+"""GLM-5.2 (``glm_moe_dsa``) NVFP4 kernel adapter.
+
+GLM-5.2 is the DeepSeek-V3.2 sparse-MLA (DSA / Lightning-Indexer) lineage. A ``quant_method:
+modelopt`` / ``quant_algo: NVFP4`` checkpoint is not a recognized transformers quantizer, so the
+model loads as a bf16 skeleton; we then register ``WeightConverter``s so the NVFP4 tensors load
+correctly. Rather than assume a fixed layout, the adapter inspects the checkpoint's safetensors
+index (:func:`inspect_nvfp4_layout`) to discover what is actually quantized — the per-expert
+projections (fused into 3D ``NVFP4Tensor`` for the scattermoe path) and any non-routed NVFP4
+linears (dequantized to bf16) — and registers exactly those converters for THIS checkpoint.
+
+When ``use_glm_dsa_kernels`` is set, ``post_model_load`` also patches the DSA attention with the
+fused absorbed-MLA sparse-gather kernels + fused Lightning-Indexer (``libs/glm_dsa``), keeps the MoE
+router fp32, and wires the context-parallel group so the attention shards the sequence on the ``cp``
+axis (composing with EP on the orthogonal ``ep`` axis).
+"""
+
+from __future__ import annotations
+
+from axolotl.integrations.kernels.adapters import ModelAdapter
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _disable_cudnn_sdp() -> None:
+    """Disable the cuDNN SDPA backend for the stock-sdpa GLM path.
+
+    GLM-5.2 develops "massive activations" (residual-stream outliers ~600+ in the top layers). The
+    cuDNN flash-attention BACKWARD produces NaN on the resulting extreme q/k (the math/mem-efficient
+    backends are robust), which propagates to every gradient → grad_norm=nan. Forward is finite, so
+    it only shows in training. Disabling the cuDNN backend makes SDPA fall back to flash/mem-efficient.
+    No-op cost on the ``use_glm_dsa_kernels`` path (that replaces SDPA with the fused DSA flash kernel).
+    """
+    import torch
+
+    try:
+        torch.backends.cuda.enable_cudnn_sdp(False)
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+
+def is_glm_moe_dsa_nvfp4_modelopt(cfg) -> bool:
+    """True iff the base model is a GLM-5.2 (``glm_moe_dsa``) NVFP4-modelopt checkpoint
+    (``quant_method=modelopt``, ``quant_algo=NVFP4``). Any failure returns False."""
+    try:
+        from transformers import AutoConfig
+
+        # Honor the user's opt-in rather than forcing remote code: glm_moe_dsa is a native
+        # transformers model, so config loading needs no remote code by default.
+        hf_cfg = AutoConfig.from_pretrained(
+            cfg.base_model,
+            trust_remote_code=bool(getattr(cfg, "trust_remote_code", False)),
+        )
+    except Exception:
+        return False
+    if str(getattr(hf_cfg, "model_type", "")) != "glm_moe_dsa":
+        return False
+    qcfg = getattr(hf_cfg, "quantization_config", None)
+    if qcfg is None:
+        return False
+    if isinstance(qcfg, dict):
+        return (
+            qcfg.get("quant_method") == "modelopt" and qcfg.get("quant_algo") == "NVFP4"
+        )
+    return (
+        getattr(qcfg, "quant_method", None) == "modelopt"
+        and getattr(qcfg, "quant_algo", None) == "NVFP4"
+    )
+
+
+def _is_glm_moe_dsa(cfg) -> bool:
+    """True iff the base model is a glm_moe_dsa checkpoint (any quant)."""
+    try:
+        from transformers import AutoConfig
+
+        # Honor the user's opt-in rather than forcing remote code: glm_moe_dsa is a native
+        # transformers model, so config loading needs no remote code by default.
+        hf_cfg = AutoConfig.from_pretrained(
+            cfg.base_model,
+            trust_remote_code=bool(getattr(cfg, "trust_remote_code", False)),
+        )
+    except Exception:
+        return False
+    return str(getattr(hf_cfg, "model_type", "")) == "glm_moe_dsa"
+
+
+def _resolve_cp_group(cfg):
+    """The context-parallel ProcessGroup (the `cp` axis of accelerate's mesh / the EP mesh), or None
+    when ``context_parallel_size <= 1``. The DSA attention shards the sequence on this axis."""
+    if (getattr(cfg, "context_parallel_size", None) or 1) <= 1:
+        return None
+    try:  # the EP plugin owns the mesh when EP is on; else read accelerate's mesh
+        from axolotl.integrations.expert_parallel.plugin import ExpertParallelPlugin
+
+        grp = ExpertParallelPlugin._resolve_cp_group(cfg)
+        if grp is not None:
+            return grp
+    except Exception:
+        pass
+    try:
+        from accelerate.state import AcceleratorState
+
+        mesh = getattr(AcceleratorState(), "device_mesh", None)
+        if mesh is not None and "cp" in (mesh.mesh_dim_names or ()):
+            return mesh["cp"].get_group()
+    except Exception:
+        pass
+    return None
+
+
+class GlmMoeDsaAdapter(ModelAdapter):
+    name = "glm_moe_dsa"
+
+    def matches(self, cfg) -> bool:
+        # NVFP4 loading needs scattermoe + a modelopt-NVFP4 checkpoint; the DSA kernels need only
+        # use_glm_dsa_kernels on a glm_moe_dsa model. Activate for either.
+        if bool(cfg.use_scattermoe) and is_glm_moe_dsa_nvfp4_modelopt(cfg):
+            return True
+        return bool(cfg.get("use_glm_dsa_kernels")) and _is_glm_moe_dsa(cfg)
+
+    def pre_model_load(self, cfg) -> None:
+        # The cuDNN SDPA backward NaNs on GLM's massive activations (see _disable_cudnn_sdp).
+        _disable_cudnn_sdp()
+        if not (cfg.use_scattermoe and is_glm_moe_dsa_nvfp4_modelopt(cfg)):
+            return  # NVFP4 converter registration is only for modelopt-NVFP4 checkpoints
+        from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_moe_loading import (
+            inspect_nvfp4_layout,
+            patch_nvfp4_tensor_meta_ops,
+            patch_skip_missing_expert_init,
+        )
+        from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_weight_converter import (
+            patch_conversion_loader_rank0_only,
+            register_nvfp4_converters_for_layout,
+        )
+
+        # FSDP2 cpu_ram_efficient_loading materializes meta receive-buffers via zeros_like/empty_like.
+        patch_nvfp4_tensor_meta_ops()
+        # transformers' conversion loader ignores cpu_ram_efficient_loading (loads on every rank);
+        # gate it to rank0-only so the full model doesn't blow up CPU RAM by the world size. Only
+        # safe when the FSDP broadcast will later fill the non-rank-0 meta params — i.e. when
+        # cpu_ram_efficient_loading is set — so DDP ranks (which each need real weights) are spared.
+        if (cfg.get("fsdp_config") or {}).get("cpu_ram_efficient_loading"):
+            patch_conversion_loader_rank0_only()
+            # transformers gates the META SKELETON (and its own load path) on is_fsdp_enabled(),
+            # which requires the process group to be initialized. axolotl doesn't init it until
+            # AFTER model load, so without this non-rank-0 builds a full real-storage skeleton
+            # (world-size× CPU blowup) before any weights even load. Init it now.
+            from transformers.integrations.fsdp import is_fsdp_enabled
+
+            from axolotl.utils.distributed import init_distributed_state
+
+            init_distributed_state()
+            LOG.info(
+                "glm_moe_dsa: initialized distributed state for rank0-only loading "
+                "(is_fsdp_enabled=%s)",
+                is_fsdp_enabled(),
+            )
+
+        import os
+
+        layout = inspect_nvfp4_layout(cfg.base_model)
+        LOG.info(
+            "glm_moe_dsa: detected NVFP4 layout — routed experts: %s (projs=%s); "
+            "non-routed NVFP4 linears: %s",
+            layout["routed_present"],
+            layout["routed_projs"],
+            layout["nonrouted_suffixes"],
+        )
+        # FAST routed-expert load (opt-in): skip the routed converters here and read+fuse the experts
+        # DIRECTLY in post_model_load — bypasses transformers' ~7-min per-tensor conversion loop over
+        # ~240k expert source tensors (direct path ~25s). Non-routed converters still register.
+        self._direct_expert_load = (
+            bool(os.environ.get("AXOLOTL_DIRECT_EXPERT_LOAD"))
+            and layout["routed_present"]
+        )
+        self._routed_projs = layout.get("routed_projs", [])
+        if self._direct_expert_load:
+            reg_layout = dict(layout)
+            reg_layout["routed_present"] = (
+                False  # don't register the slow routed converters
+            )
+            register_nvfp4_converters_for_layout("glm_moe_dsa", reg_layout)
+            patch_skip_missing_expert_init()
+            LOG.info("glm_moe_dsa: routed experts will be DIRECT-loaded (fast path)")
+        else:
+            register_nvfp4_converters_for_layout("glm_moe_dsa", layout)
+
+    def post_model_load(self, cfg, model) -> None:
+        """Patch the DSA attention with the fused absorbed-MLA kernels + keep the MoE router fp32.
+        Gated on ``use_glm_dsa_kernels``. Wires the context-parallel group so the attention shards
+        the sequence on the `cp` axis (composes with EP on the orthogonal `ep` axis)."""
+        if getattr(self, "_direct_expert_load", False):
+            import time
+
+            from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_moe_loading import (
+                direct_load_nvfp4_experts,
+            )
+
+            t0 = time.time()
+            n = direct_load_nvfp4_experts(model, cfg.base_model, self._routed_projs)
+            LOG.info(
+                "glm_moe_dsa: direct-loaded %d fused expert params in %.1fs (fast path)",
+                n,
+                time.time() - t0,
+            )
+
+        if not cfg.get("use_glm_dsa_kernels"):
+            return
+        from axolotl.integrations.kernels.libs.glm_dsa import (
+            keep_router_fp32,
+            patch_glm_moe_dsa_attention,
+        )
+
+        cp_group = _resolve_cp_group(cfg)
+        n = patch_glm_moe_dsa_attention(
+            model, use_fused_indexer=True, cp_group=cp_group
+        )
+        r = keep_router_fp32(model)
+        LOG.info(
+            "glm_moe_dsa: patched %d attention modules with fused DSA kernels (fused indexer, "
+            "context_parallel=%s); kept %d routers fp32",
+            n,
+            cp_group is not None,
+            r,
+        )
diff --git a/src/axolotl/integrations/kernels/adapters/qwen3_moe.py b/src/axolotl/integrations/kernels/adapters/qwen3_moe.py
new file mode 100644
index 0000000000..71e30db95c
--- /dev/null
+++ b/src/axolotl/integrations/kernels/adapters/qwen3_moe.py
@@ -0,0 +1,148 @@
+"""Qwen3-MoE / Qwen3-Next NVFP4 kernel adapter.
+
+A ``quant_method: modelopt`` / ``quant_algo: NVFP4`` Qwen3-MoE-family checkpoint (e.g.
+nvidia/Qwen3-30B-A3B-NVFP4, nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4) is not a recognized
+transformers quantizer, so the model loads as a bf16 skeleton; we then register
+``WeightConverter``s so the NVFP4 tensors load correctly. Like the glm_moe_dsa adapter (this is
+that adapter minus the DSA attention patching), the checkpoint's safetensors index is inspected
+(:func:`inspect_nvfp4_layout`) to discover what is actually quantized: the per-expert
+projections (fused into 3D ``NVFP4Tensor`` for the grouped MoE paths) and any non-routed NVFP4
+linears (dequantized to bf16), and exactly those converters are registered for THIS checkpoint.
+qwen3_next differs only outside the MoE block (Gated DeltaNet linear attention on 3/4 layers,
+a shared expert, per-projection ``input_scale`` tensors, an ``mtp.*`` head the causal-LM class
+never loads); the layout inspection handles all of that without model-specific code.
+"""
+
+from __future__ import annotations
+
+from axolotl.integrations.kernels.adapters import ModelAdapter
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def is_qwen3_moe_nvfp4_modelopt(cfg) -> bool:
+    """True iff the base model is a ``qwen3_moe``/``qwen3_next`` NVFP4-modelopt checkpoint
+    (``quant_method=modelopt``, ``quant_algo=NVFP4``). Any failure returns False."""
+    try:
+        from transformers import AutoConfig
+
+        # Honor the user's opt-in rather than forcing remote code: qwen3_moe is a native
+        # transformers model, so config loading needs no remote code by default.
+        hf_cfg = AutoConfig.from_pretrained(
+            cfg.base_model,
+            trust_remote_code=bool(getattr(cfg, "trust_remote_code", False)),
+        )
+    except Exception:
+        return False
+    if str(getattr(hf_cfg, "model_type", "")) not in ("qwen3_moe", "qwen3_next"):
+        return False
+    qcfg = getattr(hf_cfg, "quantization_config", None)
+    if qcfg is None:
+        return False
+    if isinstance(qcfg, dict):
+        return (
+            qcfg.get("quant_method") == "modelopt" and qcfg.get("quant_algo") == "NVFP4"
+        )
+    return (
+        getattr(qcfg, "quant_method", None) == "modelopt"
+        and getattr(qcfg, "quant_algo", None) == "NVFP4"
+    )
+
+
+class Qwen3MoeAdapter(ModelAdapter):
+    name = "qwen3_moe"
+
+    def matches(self, cfg) -> bool:
+        return bool(
+            cfg.use_scattermoe or cfg.use_sonicmoe
+        ) and is_qwen3_moe_nvfp4_modelopt(cfg)
+
+    def pre_model_load(self, cfg) -> None:
+        from transformers import AutoConfig
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_moe_loading import (
+            inspect_nvfp4_layout,
+            patch_nvfp4_tensor_meta_ops,
+            patch_skip_missing_expert_init,
+        )
+
+        # the conversion mapping is looked up under the model's actual type
+        model_type = str(
+            AutoConfig.from_pretrained(
+                cfg.base_model,
+                trust_remote_code=bool(getattr(cfg, "trust_remote_code", False)),
+            ).model_type
+        )
+        from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_weight_converter import (
+            patch_conversion_loader_rank0_only,
+            register_nvfp4_converters_for_layout,
+        )
+
+        # FSDP2 cpu_ram_efficient_loading materializes meta receive-buffers via zeros_like/empty_like.
+        patch_nvfp4_tensor_meta_ops()
+        # transformers' conversion loader ignores cpu_ram_efficient_loading (loads on every rank);
+        # gate it to rank0-only so the full model doesn't blow up CPU RAM by the world size. Only
+        # safe when the FSDP broadcast will later fill the non-rank-0 meta params — i.e. when
+        # cpu_ram_efficient_loading is set — so DDP ranks (which each need real weights) are spared.
+        if (cfg.get("fsdp_config") or {}).get("cpu_ram_efficient_loading"):
+            patch_conversion_loader_rank0_only()
+            # transformers gates the META SKELETON (and its own load path) on is_fsdp_enabled(),
+            # which requires the process group to be initialized. axolotl doesn't init it until
+            # AFTER model load, so without this non-rank-0 builds a full real-storage skeleton
+            # (world-size× CPU blowup) before any weights even load. Init it now.
+            from transformers.integrations.fsdp import is_fsdp_enabled
+
+            from axolotl.utils.distributed import init_distributed_state
+
+            init_distributed_state()
+            LOG.info(
+                "qwen3_moe: initialized distributed state for rank0-only loading "
+                "(is_fsdp_enabled=%s)",
+                is_fsdp_enabled(),
+            )
+
+        import os
+
+        layout = inspect_nvfp4_layout(cfg.base_model)
+        LOG.info(
+            "qwen3_moe: detected NVFP4 layout — routed experts: %s (projs=%s); "
+            "non-routed NVFP4 linears: %s",
+            layout["routed_present"],
+            layout["routed_projs"],
+            layout["nonrouted_suffixes"],
+        )
+        # FAST routed-expert load (opt-in): skip the routed converters here and read+fuse the experts
+        # DIRECTLY in post_model_load — bypasses transformers' per-tensor conversion loop over the
+        # per-expert source tensors. Non-routed converters still register.
+        self._direct_expert_load = (
+            bool(os.environ.get("AXOLOTL_DIRECT_EXPERT_LOAD"))
+            and layout["routed_present"]
+        )
+        self._routed_projs = layout.get("routed_projs", [])
+        if self._direct_expert_load:
+            reg_layout = dict(layout)
+            reg_layout["routed_present"] = (
+                False  # don't register the slow routed converters
+            )
+            register_nvfp4_converters_for_layout(model_type, reg_layout)
+            patch_skip_missing_expert_init()
+            LOG.info("qwen3_moe: routed experts will be DIRECT-loaded (fast path)")
+        else:
+            register_nvfp4_converters_for_layout(model_type, layout)
+
+    def post_model_load(self, cfg, model) -> None:
+        if getattr(self, "_direct_expert_load", False):
+            import time
+
+            from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_moe_loading import (
+                direct_load_nvfp4_experts,
+            )
+
+            t0 = time.time()
+            n = direct_load_nvfp4_experts(model, cfg.base_model, self._routed_projs)
+            LOG.info(
+                "qwen3_moe: direct-loaded %d fused expert params in %.1fs (fast path)",
+                n,
+                time.time() - t0,
+            )
diff --git a/src/axolotl/integrations/kernels/args.py b/src/axolotl/integrations/kernels/args.py
new file mode 100644
index 0000000000..e4df68e130
--- /dev/null
+++ b/src/axolotl/integrations/kernels/args.py
@@ -0,0 +1,351 @@
+from pydantic import BaseModel, model_validator
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+# deep_ep[_*] are EP-plugin composites, passed through when expert_parallel_size > 1.
+_BUILTIN_EXPERTS_IMPLS = {"eager", "batched_mm", "grouped_mm"}
+_KERNEL_EXPERTS_IMPLS = {"scattermoe", "sonicmoe"}
+_EP_EXPERTS_IMPLS = {
+    "deep_ep",
+    "deep_ep_grouped_mm",
+    "deep_ep_scattermoe",
+    "deep_ep_sonicmoe",
+}
+_VALID_EXPERTS_IMPLS = (
+    _BUILTIN_EXPERTS_IMPLS | _KERNEL_EXPERTS_IMPLS | _EP_EXPERTS_IMPLS
+)
+
+
+class KernelsArgs(BaseModel):
+    # --- intent-based capability surface (preferred) ----------------------------------------
+    # MoE expert backend, an alias for use_scattermoe/use_sonicmoe so new models opt into a
+    # capability by name rather than a vendor flag: scattermoe | sonicmoe | eager | builtin.
+    expert_backend: str | None = None
+    # Non-expert linear quantization policy (replaces the per-model gemma4_*_nonexpert flags):
+    # none | bf16 | fp8_blockwise | nf4. Resolved per-model by the adapters.
+    nonexpert_quantization: str | None = None
+    # Grouped NVFP4 MoE base-GEMM backend selection: auto (capability-select; default) | marlin |
+    # cutlass | deepgemm. An unavailable choice warns + falls back to auto.
+    moe_grouped_backend: str | None = None
+
+    # bnb-4bit MoE experts (quantize_moe_experts + load_in_4bit): number of experts dequantized to
+    # bf16 per chunk in the chunked-dequant grouped path. None = fixed default (memory-safe for
+    # smaller GPUs). Raise it on large GPUs to trade VRAM for throughput (bigger grouped GEMMs).
+    moe_dequant_chunk_size: int | None = None
+
+    # bnb-4bit MoE experts: route through the 1-launch parallel_linear (scatter2scatter) path instead
+    # of the chunked torch._grouped_mm path. Faster (fewer kernel launches) at the same low memory --
+    # the dequant'd bf16 is recomputed in backward via a recipe, not saved. None defaults to True; set
+    # False to force the chunked path (bounds the per-pass transient for large-expert MoEs / tiny GPUs).
+    moe_bnb_fast: bool | None = None
+
+    # --- legacy / low-level flags (kept for backwards compatibility) -------------------------
+    use_scattermoe: bool | None = None
+    use_sonicmoe: bool | None = None
+    # Fused Triton training kernels for DeepSeek-V4 (attention / RoPE / mHC).
+    use_dsv4_kernels: bool | None = None
+    # GLM-5.2 (glm_moe_dsa) DSA fused attention: MLA-absorption head-batched sparse-gather attn
+    # (fwd+bwd) + fused Lightning-Indexer + length-aware dense/sparse dispatch, replacing the dense
+    # [B,S,T]-mask eager/sdpa path. The router is kept fp32. Composes with use_scattermoe (experts)
+    # and context_parallel_size (the attention shards the sequence on the cp axis).
+    use_glm_dsa_kernels: bool | None = None
+    # DeepSeek-V4 FP8 non-expert weight storage: "float8tensor" (default, 1-byte torchao
+    # Float8Tensor base) or "bf16" (dequantize to bf16 at load).
+    dsv4_fp8_nonexpert_mode: str | None = None
+    # Native blockwise-fp8 fused LoRA kernel for the large attention projections (q_b/o_b).
+    # Off by default; the e2e gain is small for this expert-dominated model.
+    dsv4_fp8_lora_kernel: bool | None = None
+    # Fused clamped-SwiGLU LoRA kernel for the DSV4 shared-expert MLPs. Distinct from the generic
+    # `lora_mlp_kernel` (dense-MLP fusion, force-disabled under MoE kernels). A legacy
+    # `lora_mlp_kernel: true` on a DSV4 run is translated into this flag (see disable_mlp_kernel).
+    dsv4_shared_mlp_lora_kernel: bool | None = None
+    # Fallback: cast ALL residual fp32 params (incl. keep_in_fp32 mHC/norms) to the compute
+    # dtype for the fused kernels, instead of preserving keep_in_fp32 in fp32.
+    dsv4_bf16_all: bool | None = None
+    # Grouped fp4 MoE experts path (variable-M contiguous-grouped, base-fp4 GEMM + LoRA): off by
+    # default (existing fused/chunked paths unchanged). "nvfp4" = fp4 act x fp4 weight (fastest,
+    # lossy acts); "fp8" = fp8 act x mxfp4 weight (accurate). Auto-dispatches the base GEMM:
+    # DeepGEMM (sm90/sm100) -> CUTLASS grouped (sm120) -> chunked-dequant fallback.
+    dsv4_fp4_grouped_mode: str | None = None
+    # Gemma-4 frankenstein: fp8-quantize non-expert linears in-place after loading (per-channel
+    # e4m3, dequant-in-forward).  Experts remain NVFP4Tensor.  ~2 GB resident savings.
+    gemma4_fp8_nonexpert: bool | None = None
+    # Gemma-4: NF4-quantize non-expert linears (bnb 4-bit, double-quant) after loading, the same
+    # non-expert compute path unsloth uses, for apples-to-apples experts-only LoRA comparison.
+    gemma4_nf4_nonexpert: bool | None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_dsv4_fp4_grouped_mode(cls, data):
+        mode = data.get("dsv4_fp4_grouped_mode")
+        if mode is None:
+            return data
+        m = str(mode).lower()
+        if m == "fp8":
+            # Documented historically but never implemented in the training path
+            # (grouped_fp4_available/_train_backend only support 'nvfp4'). Reject loudly rather
+            # than silently no-op.
+            raise ValueError(
+                "dsv4_fp4_grouped_mode='fp8' is not implemented for training (only 'nvfp4' is "
+                "supported). Use 'nvfp4' or omit dsv4_fp4_grouped_mode."
+            )
+        if m != "nvfp4":
+            raise ValueError(f"dsv4_fp4_grouped_mode must be 'nvfp4', got {mode!r}")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_dsv4_fp8_nonexpert_mode(cls, data):
+        mode = data.get("dsv4_fp8_nonexpert_mode")
+        if mode is not None and str(mode).lower() not in ("float8tensor", "bf16"):
+            raise ValueError(
+                f"dsv4_fp8_nonexpert_mode must be 'float8tensor' or 'bf16', got {mode!r}"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_moe_dequant_chunk_size(cls, data):
+        chunk = data.get("moe_dequant_chunk_size")
+        if chunk is None:
+            return data
+        err = ValueError(
+            f"moe_dequant_chunk_size must be a positive integer, got {chunk!r}"
+        )
+        if isinstance(chunk, bool):
+            raise err
+        try:
+            chunk_int = int(chunk)
+        except (TypeError, ValueError):
+            raise err from None
+        if isinstance(chunk, float) and not chunk.is_integer():
+            raise err
+        if chunk_int <= 0:
+            raise err
+        return data
+
+    @staticmethod
+    def _canonicalize_expert_backend(data):
+        """Canonicalize the intent ``expert_backend`` onto use_scattermoe/use_sonicmoe (the rest of
+        the stack reads those). ``expert_backend`` is authoritative: it sets the chosen backend True
+        and REJECTS an explicit legacy flag that contradicts it (so the result is never ambiguous
+        regardless of validator ordering). Idempotent: every consumer below calls this first because
+        pydantic runs same-mode validators in REVERSE definition order, so the before-validator alone
+        would run AFTER its consumers and they'd never see its writes."""
+        eb = data.get("expert_backend")
+        if eb is None:
+            return data
+        eb = str(eb).lower()
+        valid = {"scattermoe", "sonicmoe", "eager", "builtin"}
+        if eb not in valid:
+            raise ValueError(
+                f"expert_backend must be one of {sorted(valid)}, got {eb!r}"
+            )
+        want_scatter = eb == "scattermoe"
+        want_sonic = eb == "sonicmoe"
+        if data.get("use_scattermoe") is True and not want_scatter:
+            raise ValueError(
+                f"expert_backend={eb!r} conflicts with use_scattermoe=true; set only one."
+            )
+        if data.get("use_sonicmoe") is True and not want_sonic:
+            raise ValueError(
+                f"expert_backend={eb!r} conflicts with use_sonicmoe=true; set only one."
+            )
+        # Only the chosen backend is written, leaving the other untouched (None), so the end state is
+        # byte-identical to setting the equivalent legacy flag directly. eager/builtin write nothing.
+        if want_scatter:
+            data["use_scattermoe"] = True
+        elif want_sonic:
+            data["use_sonicmoe"] = True
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def normalize_expert_backend(cls, data):
+        return cls._canonicalize_expert_backend(data)
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_moe_grouped_backend(cls, data):
+        backend = data.get("moe_grouped_backend")
+        if backend is None:
+            return data
+        b = str(backend).lower()
+        if b == "dequant":
+            # The chunked-dequant fallback has no training/autograd path (the training dispatch only
+            # wires marlin/deepgemm/cutlass); accepting it would silently run cutlass. Reject loudly.
+            raise ValueError(
+                "moe_grouped_backend='dequant' is not implemented for training (the chunked-dequant "
+                "fallback has no autograd path). Use 'auto' (default), 'marlin', 'deepgemm', or "
+                "'cutlass', or omit moe_grouped_backend."
+            )
+        valid = {"auto", "marlin", "cutlass", "deepgemm"}
+        if b not in valid:
+            raise ValueError(
+                f"moe_grouped_backend must be one of {sorted(valid)}, got {backend!r}"
+            )
+        # The override only takes effect once the grouped NVFP4 MoE path is enabled.
+        if not data.get("dsv4_fp4_grouped_mode"):
+            LOG.warning(
+                "moe_grouped_backend=%r has no effect unless the grouped NVFP4 MoE path is enabled "
+                "(set dsv4_fp4_grouped_mode: nvfp4).",
+                backend,
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_nonexpert_quantization(cls, data):
+        """Validate the non-expert quantization intent and warn on the deprecated per-model flags."""
+        nq = data.get("nonexpert_quantization")
+        if nq is not None:
+            valid = {"none", "bf16", "fp8", "fp8_blockwise", "nf4", "nvfp4"}
+            if str(nq).lower() not in valid:
+                raise ValueError(
+                    f"nonexpert_quantization must be one of {sorted(valid)}, got {nq!r}"
+                )
+        if data.get("gemma4_fp8_nonexpert") or data.get("gemma4_nf4_nonexpert"):
+            LOG.warning(
+                "gemma4_fp8_nonexpert / gemma4_nf4_nonexpert are deprecated; prefer "
+                "`nonexpert_quantization: fp8_blockwise` or `nonexpert_quantization: nf4`."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_mutually_exclusive(cls, data):
+        data = cls._canonicalize_expert_backend(data)
+        if data.get("use_scattermoe") and data.get("use_sonicmoe"):
+            raise ValueError(
+                "Cannot use both ScatterMoE and SonicMoE simultaneously. "
+                "Please set only one of `use_scattermoe` or `use_sonicmoe` to true."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_use_kernels(cls, data):
+        if data.get("use_kernels") is not True:
+            LOG.warning(
+                "`use_kernels` must be set to True to use this. Automatically setting it to True."
+            )
+            data["use_kernels"] = True
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_dsv4_attention_lora_unsupported(cls, data):
+        """Reject module-level (attention) LoRA on a DSV4 fused-kernel run.
+
+        The fused indexer feeds only gradientless topk indices, so LoRA on the indexer/scorer
+        projections trains against no gradient; and attention-level LoRA reintroduces data-dependent
+        FSDP2 backward collectives that break the experts-only invariant the DSV4 recipe relies on.
+        Experts-only LoRA (lora_target_parameters) is the supported surface. This also covers
+        lora_target_linear: true, which expands (find_all_linear_names) to every Linear including
+        attention q/k/v/o. lora_exclude_modules is the explicit opt-out: setting it signals the user
+        has excluded the indexer scorer projections, so we downgrade to a warning."""
+        if not data.get("use_dsv4_kernels"):
+            return data
+        if not (data.get("lora_target_modules") or data.get("lora_target_linear")):
+            return data
+        if data.get("lora_exclude_modules"):
+            LOG.warning(
+                "attention/module-level LoRA (lora_target_modules / lora_target_linear) with "
+                "use_dsv4_kernels: ensure lora_exclude_modules excludes the indexer scorer "
+                "projections (the fused indexer is gradientless); keep LoRA experts-only "
+                "(lora_target_parameters) where possible."
+            )
+            return data
+        raise ValueError(
+            "attention/module-level LoRA is not supported with use_dsv4_kernels (this includes "
+            "lora_target_modules and lora_target_linear: true, which expands to all linear layers "
+            "incl. attention q/k/v/o): the fused indexer is gradientless (topk indices only) and "
+            "attention LoRA reintroduces data-dependent FSDP2 backward collectives that break the "
+            "experts-only invariant. Either (1) keep LoRA experts-only via lora_target_parameters, "
+            "or (2) explicitly exclude the indexer scorer projections via lora_exclude_modules."
+        )
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_sonicmoe_ep_unsupported(cls, data):
+        """SonicMoE + EP is not yet implemented (EP `_sonicmoe_local` raises)."""
+        data = cls._canonicalize_expert_backend(data)
+        if data.get("use_sonicmoe") and (data.get("expert_parallel_size") or 1) > 1:
+            raise ValueError(
+                "use_sonicmoe=true is not supported with expert_parallel_size > 1. "
+                "Use use_scattermoe=true under EP, or set expert_parallel_size=1."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_experts_implementation(cls, data):
+        """Auto-select impl from kernel flags; reject mismatched/unknown values."""
+        data = cls._canonicalize_expert_backend(data)
+        experts_implementation = data.get("experts_implementation")
+        use_scattermoe = bool(data.get("use_scattermoe"))
+        use_sonicmoe = bool(data.get("use_sonicmoe"))
+
+        if experts_implementation is None:
+            if use_scattermoe:
+                data["experts_implementation"] = "scattermoe"
+            elif use_sonicmoe:
+                data["experts_implementation"] = "sonicmoe"
+            else:
+                # Transformers defaults to a non-eager backend when unset; pin to
+                # eager unless the user explicitly opts in.
+                data["experts_implementation"] = "eager"
+            return data
+
+        if experts_implementation == "scattermoe" and not use_scattermoe:
+            LOG.warning(
+                "`experts_implementation='scattermoe'` requires `use_scattermoe: true`. "
+                "Automatically setting to 'eager'."
+            )
+            data["experts_implementation"] = "eager"
+        elif experts_implementation == "sonicmoe" and not use_sonicmoe:
+            LOG.warning(
+                "`experts_implementation='sonicmoe'` requires `use_sonicmoe: true`. "
+                "Automatically setting to 'eager'."
+            )
+            data["experts_implementation"] = "eager"
+        elif experts_implementation not in _VALID_EXPERTS_IMPLS:
+            LOG.warning(
+                f"`experts_implementation={experts_implementation!r}` is not recognized. "
+                f"Valid options: {sorted(_VALID_EXPERTS_IMPLS)}. "
+                f"Automatically setting to 'eager'."
+            )
+            data["experts_implementation"] = "eager"
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def disable_mlp_kernel(cls, data):
+        data = cls._canonicalize_expert_backend(data)
+        if data.get("use_scattermoe") is True or data.get("use_sonicmoe") is True:
+            # DSV4's shared/routed expert MLP needs the dedicated clamped-SwiGLU kernel, not the
+            # generic dense-MLP one; translate the intent and disable the generic path for DSV4.
+            if (
+                data.get("lora_mlp_kernel") is True
+                and data.get("use_dsv4_kernels") is True
+            ):
+                if data.get("dsv4_shared_mlp_lora_kernel") is None:
+                    data["dsv4_shared_mlp_lora_kernel"] = True
+                    LOG.warning(
+                        "Translated lora_mlp_kernel -> dsv4_shared_mlp_lora_kernel for the DSV4 MoE "
+                        "run (the generic lora_mlp_kernel is disabled under DSV4 kernels)."
+                    )
+                data["lora_mlp_kernel"] = False
+            # Otherwise keep lora_mlp_kernel: under the custom MoE expert kernels it only fuses the
+            # DENSE shared MLP (layer.mlp via find_mlp_in_layer), which is a plain gated Linear MLP
+            # separate from the routed experts (those are handled by the MoE kernel and aren't PEFT
+            # nn.Linear, so the can_patch_mlp lora_A guard skips them). The fused kernel dequantizes
+            # bnb-4bit / fp8 bases, so it composes with quantized non-experts.
+            data["mlp_kernel"] = False  # the non-LoRA mlp kernel stays off under MoE
+
+        return data
diff --git a/src/axolotl/integrations/kernels/autotune_callback.py b/src/axolotl/integrations/kernels/autotune_callback.py
new file mode 100644
index 0000000000..c0444777fa
--- /dev/null
+++ b/src/axolotl/integrations/kernels/autotune_callback.py
@@ -0,0 +1,149 @@
+"""Trainer callback for reporting Triton autotune results from scattermoe-lora kernels."""
+
+import torch
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# Give up looking for autotune data after this many training steps.
+_MAX_POLL_STEP = 5
+
+
+def _get_gpu_info() -> dict:
+    """Return basic GPU identification for the current device."""
+    if not torch.cuda.is_available():
+        return {}
+    try:
+        idx = torch.cuda.current_device()
+        props = torch.cuda.get_device_properties(idx)
+        return {
+            "gpu_name": props.name,
+            "gpu_compute_capability": f"{props.major}.{props.minor}",
+            "gpu_memory_bytes": props.total_memory,
+        }
+    except Exception:  # pylint: disable=broad-exception-caught
+        return {}
+
+
+def _get_smem_capacity() -> dict:
+    """Return shared memory capacity (bytes). Prefers the runtime lora_ops helper, falls
+    back to the Triton active-driver device properties, then to torch device properties —
+    so a dsv4-only run (no scattermoe) still reports smem."""
+    # 1) lora_ops helper (matches what the scattermoe autotuner actually saw)
+    try:
+        from axolotl.integrations.kernels.autotune_collector import (
+            _find_lora_ops_module,
+        )
+
+        lora_ops = _find_lora_ops_module()
+        fn = getattr(lora_ops, "_get_smem_capacity", None) if lora_ops else None
+        if fn is not None:
+            return {"smem_capacity_bytes": fn()}
+    except Exception:  # pylint: disable=broad-exception-caught
+        pass
+    # 2) Triton active driver
+    try:
+        import triton
+
+        idx = torch.cuda.current_device()
+        props = triton.runtime.driver.active.utils.get_device_properties(idx)
+        smem = props.get("max_shared_mem")
+        if smem:
+            return {"smem_capacity_bytes": int(smem)}
+    except Exception:  # pylint: disable=broad-exception-caught
+        pass
+    # 3) torch device properties
+    try:
+        props = torch.cuda.get_device_properties(torch.cuda.current_device())
+        for name in (
+            "shared_memory_per_block_optin",
+            "shared_memory_per_multiprocessor",
+        ):
+            val = getattr(props, name, None)
+            if val:
+                return {"smem_capacity_bytes": int(val), "smem_source": name}
+    except Exception:  # pylint: disable=broad-exception-caught
+        pass
+    return {}
+
+
+class AutotuneReportCallback(TrainerCallback):
+    """Reports Triton kernel autotune selections via telemetry.
+
+    Fires **once** after the first training step completes (step 1), at
+    which point the forward and backward passes have both run and the
+    autotuned kernels have populated their caches.  If for some reason
+    the caches are still empty (e.g. the kernel was never invoked), the
+    callback retries on subsequent steps up to ``_MAX_POLL_STEP`` and
+    then stops polling.
+
+    After reporting (or giving up) every subsequent ``on_step_end``
+    call short-circuits on the ``_reported`` flag — zero hot-path cost.
+    """
+
+    def __init__(self):
+        self._reported = False
+
+    # pylint: disable=unused-argument
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        if self._reported:
+            return
+
+        # Lazy import: Triton / scattermoe kernels may not be installed.
+        from axolotl.integrations.kernels.autotune_collector import (
+            collect_autotune_configs,
+        )
+
+        configs = collect_autotune_configs()
+
+        if not configs:
+            if state.global_step >= _MAX_POLL_STEP:
+                LOG.debug(
+                    "No autotune data found after %d steps; giving up.",
+                    state.global_step,
+                )
+                self._reported = True
+            return
+
+        self._reported = True
+
+        from axolotl.telemetry.manager import TelemetryManager
+
+        telemetry_manager = TelemetryManager.get_instance()
+        if not telemetry_manager.enabled:
+            return
+
+        properties = {
+            "kernel_count": len(configs),
+            "kernels": configs,
+        }
+        properties.update(_get_gpu_info())
+        properties.update(_get_smem_capacity())
+
+        telemetry_manager.send_event(
+            event_type="triton-autotune",
+            properties=properties,
+        )
+
+        names = sorted({c.get("kernel", "?") for c in configs})
+        LOG.info(
+            "Reported %d Triton autotune config(s) to telemetry on %s (sm %s, smem %s B): %s",
+            len(configs),
+            properties.get("gpu_name", "?"),
+            properties.get("gpu_compute_capability", "?"),
+            properties.get("smem_capacity_bytes", "?"),
+            ", ".join(names),
+        )
diff --git a/src/axolotl/integrations/kernels/autotune_collector.py b/src/axolotl/integrations/kernels/autotune_collector.py
new file mode 100644
index 0000000000..34433f2ddc
--- /dev/null
+++ b/src/axolotl/integrations/kernels/autotune_collector.py
@@ -0,0 +1,170 @@
+"""Collect Triton autotune results from scattermoe-lora kernels.
+
+This module reads the ``.cache`` attribute from Triton ``@triton.autotune``
+decorated kernel objects and returns structured dicts describing the selected
+configurations.  It has **no** telemetry dependency — callers decide what to
+do with the data.
+"""
+
+import sys
+from types import ModuleType
+from typing import Any
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# (human-readable name, attribute on the lora_ops module)
+_KERNEL_REGISTRY: list[tuple[str, str]] = [
+    ("scatter2scatter_lora_fwd", "_scatter2scatter_lora"),
+    ("scatter2scatter_lora_dX", "_scatter2scatter_lora_dX"),
+    ("group_bwd_lora", "_group_bwd_lora"),
+    ("group_bwd_lora_fused", "_group_bwd_lora_fused"),
+]
+
+# The autotune key declared on every kernel: key=["M_BUCKET", "N", "K"].
+# M_BUCKET is the seqlen-bucketed M (see _bucket_m in lora_ops.py) so cache
+# entries don't churn with every distinct M.
+_KEY_NAMES: list[str] = ["M_BUCKET", "N", "K"]
+
+
+def _parse_key_tuple(key_tuple: tuple) -> dict[str, Any]:
+    """Turn the autotune cache key tuple into a labelled dict.
+
+    Triton builds the cache key from the values of the declared ``key``
+    args (``M_BUCKET``, ``N``, ``K``) followed by dtype signature elements.
+    We label the first three and store the rest under ``_extra``.
+    """
+    result: dict[str, Any] = {}
+    for i, name in enumerate(_KEY_NAMES):
+        if i < len(key_tuple):
+            result[name] = key_tuple[i]
+    if len(key_tuple) > len(_KEY_NAMES):
+        result["_extra"] = [str(v) for v in key_tuple[len(_KEY_NAMES) :]]
+    return result
+
+
+def _find_lora_ops_module() -> ModuleType | None:
+    """Locate the *runtime* ``lora_ops`` module in ``sys.modules``.
+
+    Normally there is a single canonical instance (the axolotl import path); a duplicate can
+    appear if the same file is imported under a second module name (its kernel objects would
+    carry separate ``.cache`` dicts). ``register_scattermoe_experts`` calls
+    ``_ensure_single_lora_ops`` to collapse such duplicates, but this lookup stays defensive:
+    it returns the first ``sys.modules`` entry whose name contains ``lora_ops`` and that
+    exposes the ``_scatter2scatter_lora`` kernel.
+    """
+    for name, module in list(sys.modules.items()):
+        if (
+            module is not None
+            and "lora_ops" in name
+            and hasattr(module, "_scatter2scatter_lora")
+        ):
+            return module
+    return None
+
+
+# Substrings identifying the kernel-bearing modules to scan in sys.modules. Covers the HF-`kernels`
+# hash-suffixed scattermoe copies, the normally-imported dsv4 kernels, and the glm_dsa (GLM-5.2
+# sparse-MLA) kernels.
+_KERNEL_MODULE_HINTS: tuple[str, ...] = (
+    "lora_ops",
+    "scattermoe_lora.kernels",
+    "libs.dsv4.attention",
+    "libs.dsv4.attention_csa",
+    "libs.dsv4.attention_gather",
+    "libs.dsv4.rope",
+    "libs.dsv4.mhc",
+    "libs.dsv4.gated_pool",
+    "libs.dsv4.indexer",
+    ".dsv4.",
+    "libs.glm_dsa.attention_topk",
+    "libs.glm_dsa.attention_mla_absorb",
+    "libs.glm_dsa.indexer",
+    ".glm_dsa.",
+)
+
+
+def _is_autotuner(obj: Any) -> bool:
+    """A Triton ``Autotuner`` exposes a ``.cache`` dict and the wrapped ``.base_fn``/``.fn``."""
+    return isinstance(getattr(obj, "cache", None), dict) and (
+        getattr(obj, "base_fn", None) is not None
+        or getattr(obj, "fn", None) is not None
+    )
+
+
+def _config_to_dict(config: Any) -> dict[str, Any]:
+    out = dict(getattr(config, "kwargs", {}) or {})
+    out["num_warps"] = getattr(config, "num_warps", None)
+    out["num_stages"] = getattr(config, "num_stages", None)
+    if getattr(config, "num_ctas", None) is not None:
+        out["num_ctas"] = config.num_ctas
+    return out
+
+
+def _label_key(autotuner: Any, key_tuple: tuple) -> dict[str, Any]:
+    """Label the cache key tuple by the kernel's own declared autotune ``key`` arg names
+    (e.g. ``["M_BUCKET","N","K"]`` or ``["S","T","H"]``); extras (dtype specialization) go
+    under ``_extra``."""
+    names = list(getattr(autotuner, "keys", None) or _KEY_NAMES)
+    result: dict[str, Any] = {}
+    for i, name in enumerate(names):
+        if i < len(key_tuple):
+            result[name] = key_tuple[i]
+    if len(key_tuple) > len(names):
+        result["_extra"] = [str(v) for v in key_tuple[len(names) :]]
+    return result
+
+
+def collect_autotune_configs() -> list[dict[str, Any]]:
+    """Read autotune caches from ALL dsv4 + scattermoe ``@triton.autotune`` kernels.
+
+    Returns a (possibly empty) list of dicts, each containing:
+
+    * ``kernel`` – kernel function name
+    * ``module`` – short module name it lives in
+    * ``key``    – dict of the autotune-key args (problem shapes), labeled by the kernel's
+      own declared ``key`` names
+    * ``config`` – selected tile sizes, ``num_warps``, ``num_stages`` (and ``num_ctas``)
+
+    Scans ``sys.modules`` for both the HF-``kernels`` hash-suffixed scattermoe copies (whose
+    caches are the populated runtime ones) and the normally-imported dsv4 kernel modules.
+    """
+    results: list[dict[str, Any]] = []
+    # Dedup by the FULL module path so two distinct module instances of the same kernel
+    # (e.g. a duplicate import) stay visible as separate entries; that duplication is exactly
+    # what telemetry should surface (a single Autotuner.cache is a dict and can't hold a key
+    # twice, so any same-(kernel,key) duplicate means >1 module instance).
+    seen: set[tuple[str, str, tuple]] = set()
+
+    for modname, module in list(sys.modules.items()):
+        if module is None or not any(h in modname for h in _KERNEL_MODULE_HINTS):
+            continue
+        for attr in dir(module):
+            obj = getattr(module, attr, None)
+            if not _is_autotuner(obj):
+                continue
+            cache = obj.cache  # type: ignore[union-attr]
+            if not cache:
+                continue
+            base = getattr(obj, "base_fn", None) or getattr(obj, "fn", None)
+            kname = getattr(base, "__name__", attr)
+            for key_tuple, config in cache.items():
+                dedup = (modname, kname, tuple(key_tuple))
+                if dedup in seen:
+                    continue
+                seen.add(dedup)
+                results.append(
+                    {
+                        "kernel": kname,
+                        "module": modname.rsplit(".", 1)[-1],
+                        # Fully-qualified dotted module name (NOT a filesystem path). Named
+                        # `module_fqn` rather than `module_path` so the telemetry PII redactor
+                        # (which scrubs any key containing "path") doesn't blank this safe value.
+                        "module_fqn": modname,
+                        "key": _label_key(obj, key_tuple),
+                        "config": _config_to_dict(config),
+                    }
+                )
+
+    return results
diff --git a/src/axolotl/integrations/kernels/constants.py b/src/axolotl/integrations/kernels/constants.py
new file mode 100644
index 0000000000..7373fa5ef2
--- /dev/null
+++ b/src/axolotl/integrations/kernels/constants.py
@@ -0,0 +1,36 @@
+"""Diagnostic helpers for MoE kernel integrations (kernel dispatch itself
+is architecture-agnostic via the ExpertsInterface)."""
+
+import importlib
+
+# Models where MoE is embedded in the decoder layer (no separate SparseMoeBlock).
+EXPERTS_ONLY_BLOCK = {
+    "gemma4_text": "Gemma4TextExperts",
+}
+
+
+def resolve_experts_class(model_type: str):
+    """Resolve the Experts class for a known model type, or ``None``."""
+    entry = EXPERTS_ONLY_BLOCK.get(model_type)
+    if entry is None:
+        return None
+
+    module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+    try:
+        module = importlib.import_module(module_path)
+    except ModuleNotFoundError:
+        if model_type.endswith("_text"):
+            parent_type = model_type.removesuffix("_text")
+            module_path = f"transformers.models.{parent_type}.modeling_{parent_type}"
+            module = importlib.import_module(module_path)
+        else:
+            raise
+
+    cls = getattr(module, entry, None)
+    if cls is None:
+        raise ValueError(f"Could not find class '{entry}' in '{module_path}'")
+    return cls
+
+
+def is_experts_only_model(model_type: str) -> bool:
+    return model_type in EXPERTS_ONLY_BLOCK
diff --git a/src/axolotl/integrations/kernels/libs/__init__.py b/src/axolotl/integrations/kernels/libs/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/__init__.py b/src/axolotl/integrations/kernels/libs/dsv4/__init__.py
new file mode 100644
index 0000000000..51a0452d47
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/__init__.py
@@ -0,0 +1,22 @@
+from .attention import sliding_attn
+from .attention_csa import csa_attn
+from .attention_gather import csa_attn_topk
+from .gated_pool import gated_softmax_pool
+from .indexer import indexer_scores
+from .lora_mlp import apply_lora_mlp_clamped_swiglu, patch_dsv4_shared_mlp_lora
+from .mhc import hyperconnection_forward
+from .patch import patch_deepseek_v4_kernels
+from .rope import apply_rotary_pos_emb_triton
+
+__all__ = [
+    "sliding_attn",
+    "csa_attn",
+    "csa_attn_topk",
+    "gated_softmax_pool",
+    "indexer_scores",
+    "hyperconnection_forward",
+    "apply_rotary_pos_emb_triton",
+    "patch_deepseek_v4_kernels",
+    "apply_lora_mlp_clamped_swiglu",
+    "patch_dsv4_shared_mlp_lora",
+]
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/attention.py b/src/axolotl/integrations/kernels/libs/dsv4/attention.py
new file mode 100644
index 0000000000..7a1cb1fe85
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/attention.py
@@ -0,0 +1,588 @@
+"""Fused flash attention for DeepSeek-V4 sliding-window layers (fwd + bwd).
+
+Replaces the forced-eager path in
+``transformers.models.deepseek_v4.modeling_deepseek_v4.eager_attention_forward``
+for ``sliding_attention`` layers. V4 cannot use FA2/3/4 (head_dim 512 > 256 cap),
+SDPA (no per-head sink), or flex (compressor KV-concat), so eager — which
+materializes ``[B, H, S, KV]`` scores and broadcasts the single MQA KV head ×64 via
+``repeat_kv`` — is the only baseline. This kernel:
+
+  * keeps the single MQA KV head in SRAM (no ``repeat_kv`` blow-up),
+  * never materializes the score matrix,
+  * masks the sliding window *implicitly* from indices, so work is O(S * window)
+    instead of O(S * KV),
+  * folds the gpt-oss-style per-head learnable **sink** into the online softmax by
+    seeding ``m = sink`` / ``l = 1`` (the sink is a logit-only column: it enters the
+    denominator, contributes nothing to the output),
+  * exposes an ``acc_dtype`` knob (fp32 vs bf16 online-softmax/PV accumulation) so the
+    speed/memory vs error trade-off can be measured per path.
+
+Sliding-window semantics (verified against ``create_sliding_window_causal_mask``):
+query ``i`` attends to key ``j`` iff ``j <= i`` and ``i - j < sliding_window``.
+"""
+
+import functools
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+# Autotuner prunes configs that overflow SMEM per dtype/head_dim. ``EL`` (element size) is
+# in the autotune key so fp32 and bf16 are tuned/cached separately.
+_CONFIGS = [
+    triton.Config({"BLOCK_M": bm, "BLOCK_N": bn}, num_warps=w, num_stages=s)
+    for bm in (16, 32, 64, 128)
+    for bn in (16, 32, 64)
+    for w in (4, 8)
+    for s in (1, 2)
+    # N=16 wgmma at large M illegal-accesses on Hopper/Blackwell; keep BN=16 only for small
+    # BM (the fp32 D=512 path needs the small tile).
+    if not (bn == 16 and bm >= 64)
+]
+
+
+@functools.lru_cache(maxsize=None)
+def _smem_limit(dev):
+    # Real per-block SMEM of this GPU (portable: sm_120 ~99KB, H100/B200 ~228KB).
+    try:
+        return triton.runtime.driver.active.utils.get_device_properties(dev)[
+            "max_shared_mem"
+        ]
+    except Exception:
+        return 101376
+
+
+@functools.lru_cache(maxsize=None)
+def _max_m(dev):
+    # Blackwell (cc >= 10, tcgen05) keeps the MMA accumulator in tensor memory (512 cols).
+    # With D=512 the [M,512] acc nearly fills tmem, so M>32 misaligns (sm_100) / OOMs tmem;
+    # cap M<=32 there. Hopper (cc 9, no tmem) handles up to 128.
+    try:
+        return 32 if torch.cuda.get_device_capability(dev)[0] >= 10 else 128
+    except Exception:
+        return 128
+
+
+def _prune_smem(n_qtiles, n_kvtiles):
+    """Drop configs that would overflow SMEM (D-wide tiles ×num_stages) or, on Blackwell,
+    exceed the tensor-memory accumulator limit (M>32). Portable; the autotuner's own
+    OutOfResources catch is the backstop."""
+
+    def prune(configs, nargs, **kwargs):
+        D, EL = (
+            kwargs["D"],
+            nargs["EL"],
+        )  # D is constexpr (kwargs); EL is runtime (nargs)
+        dev = torch.cuda.current_device()
+        budget = int(_smem_limit(dev) * 0.9)
+        max_m = _max_m(dev)
+        kept = []
+        for c in configs:
+            bm, bn, st = c.kwargs["BLOCK_M"], c.kwargs["BLOCK_N"], c.num_stages
+            est = (n_qtiles * bm + n_kvtiles * st * bn) * D * EL
+            if est <= budget and bm <= max_m:
+                kept.append(c)
+        return kept or [
+            min(
+                configs,
+                key=lambda c: c.kwargs["BLOCK_M"] * c.kwargs["BLOCK_N"] * c.num_stages,
+            )
+        ]
+
+    return prune
+
+
+@triton.autotune(
+    configs=_CONFIGS,
+    key=["H", "EL"],
+    prune_configs_by={"early_config_prune": _prune_smem(1, 2)},
+)  # q + (k,v)
+@triton.jit
+def _fwd_kernel(
+    Q,
+    K,
+    V,
+    sinks,
+    Out,
+    L,
+    scale,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vn,
+    stride_vd,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    H,
+    S,
+    KV,
+    WINDOW,
+    EL,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    ACC: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, D)
+
+    q_ptr = (
+        Q
+        + b * stride_qb
+        + h * stride_qh
+        + offs_m[:, None] * stride_qm
+        + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptr, mask=offs_m[:, None] < S, other=0.0)
+
+    sink = tl.load(sinks + h).to(tl.float32)
+    m_i = tl.zeros([BLOCK_M], tl.float32) + sink
+    l_i = tl.zeros([BLOCK_M], tl.float32) + 1.0  # exp(sink - sink) seed
+    acc = tl.zeros([BLOCK_M, D], tl.float32)
+
+    # sliding window: keys in (m - window, m]
+    lo = tl.maximum(0, (pid_m * BLOCK_M - WINDOW + 1))
+    lo = (lo // BLOCK_N) * BLOCK_N
+    hi = tl.minimum(KV, pid_m * BLOCK_M + BLOCK_M)
+
+    kbase = K + b * stride_kb
+    vbase = V + b * stride_vb
+    for start_n in range(lo, hi, BLOCK_N):
+        offs_n = start_n + tl.arange(0, BLOCK_N)
+        nmask = offs_n < KV
+        k = tl.load(
+            kbase + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(k), input_precision=PREC) * scale
+
+        valid = (
+            (offs_n[None, :] <= offs_m[:, None])
+            & (offs_m[:, None] - offs_n[None, :] < WINDOW)
+            & nmask[None, :]
+        )
+        qk = tl.where(valid, qk, float("-inf"))
+
+        m_new = tl.maximum(m_i, tl.max(qk, 1))
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new[:, None])
+        l_i = l_i * alpha + tl.sum(p, 1)
+        acc = acc * alpha[:, None]
+        v = tl.load(
+            vbase + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        acc += tl.dot(p.to(ACC), v.to(ACC), input_precision=PREC).to(tl.float32)
+        m_i = m_new
+
+    acc = acc / l_i[:, None]
+    o_ptr = (
+        Out
+        + b * stride_ob
+        + h * stride_oh
+        + offs_m[:, None] * stride_om
+        + offs_d[None, :] * stride_od
+    )
+    tl.store(o_ptr, acc.to(Out.dtype.element_ty), mask=offs_m[:, None] < S)
+    tl.store(L + pid_bh * S + offs_m, m_i + tl.log(l_i), mask=offs_m < S)
+
+
+@register_kernel_op("dsv4_sliding_attn_fwd")
+def _sliding_attn_fwd_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    sinks: torch.Tensor,
+    scale: float,
+    window: int,
+    acc_fp32: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, H, S, D = q.shape
+    KV = k.shape[1]
+    out = torch.empty(B, H, S, D, device=q.device, dtype=q.dtype)
+    L = torch.empty(B * H, S, device=q.device, dtype=torch.float32)
+    ACC = tl.float32 if acc_fp32 else tl.bfloat16
+    grid = lambda meta: (triton.cdiv(S, meta["BLOCK_M"]), B * H)
+    _fwd_kernel[grid](
+        q,
+        k,
+        v,
+        sinks,
+        out,
+        L,
+        scale,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        v.stride(0),
+        v.stride(1),
+        v.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        H,
+        S,
+        KV,
+        window,
+        q.element_size(),
+        D=D,
+        ACC=ACC,
+        PREC=("ieee" if q.element_size() == 4 else "tf32"),
+    )
+    return out, L
+
+
+@_sliding_attn_fwd_op.register_fake
+def _(q, k, v, sinks, scale, window, acc_fp32):
+    B, H, S, _ = q.shape
+    return torch.empty(q.shape, dtype=q.dtype, device=q.device), torch.empty(
+        B * H, S, device=q.device, dtype=torch.float32
+    )
+
+
+# Two-kernel backward (dq pass, dk/dv pass): splitting halves the resident D-wide dot
+# operands per kernel so 16x16 tiles fit ~99KB SMEM.
+@triton.autotune(
+    configs=_CONFIGS,
+    key=["H", "EL"],
+    reset_to_zero=["DSINK"],
+    prune_configs_by={"early_config_prune": _prune_smem(2, 2)},
+)
+@triton.jit
+def _bwd_dq_kernel(
+    Q,
+    K,
+    V,
+    sinks,
+    DO,
+    L,
+    Delta,
+    DQ,
+    DSINK,
+    scale,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kn,
+    stride_kd,
+    H,
+    S,
+    KV,
+    WINDOW,
+    EL,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    ACC: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, D)
+    mmask = offs_m < S
+
+    q = tl.load(
+        Q
+        + b * stride_qb
+        + h * stride_qh
+        + offs_m[:, None] * stride_qm
+        + offs_d[None, :] * stride_qd,
+        mask=mmask[:, None],
+        other=0.0,
+    )
+    do = tl.load(
+        DO
+        + b * stride_qb
+        + h * stride_qh
+        + offs_m[:, None] * stride_qm
+        + offs_d[None, :] * stride_qd,
+        mask=mmask[:, None],
+        other=0.0,
+    )
+    l_i = tl.load(L + pid_bh * S + offs_m, mask=mmask, other=0.0)
+    delta = tl.load(Delta + pid_bh * S + offs_m, mask=mmask, other=0.0)
+    sink = tl.load(sinks + h).to(tl.float32)
+
+    dq = tl.zeros([BLOCK_M, D], tl.float32)
+    lo = (tl.maximum(0, pid_m * BLOCK_M - WINDOW + 1) // BLOCK_N) * BLOCK_N
+    hi = tl.minimum(KV, pid_m * BLOCK_M + BLOCK_M)
+    for start_n in range(lo, hi, BLOCK_N):
+        offs_n = start_n + tl.arange(0, BLOCK_N)
+        nmask = offs_n < KV
+        k = tl.load(
+            K
+            + b * stride_kb
+            + offs_n[:, None] * stride_kn
+            + offs_d[None, :] * stride_kd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        v = tl.load(
+            V
+            + b * stride_kb
+            + offs_n[:, None] * stride_kn
+            + offs_d[None, :] * stride_kd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(k), input_precision=PREC) * scale
+        valid = (
+            (offs_n[None, :] <= offs_m[:, None])
+            & (offs_m[:, None] - offs_n[None, :] < WINDOW)
+            & nmask[None, :]
+        )
+        p = tl.where(valid, tl.exp(qk - l_i[:, None]), 0.0)
+        dp = tl.dot(do, tl.trans(v), input_precision=PREC)
+        ds = (p * (dp - delta[:, None]) * scale).to(ACC)
+        dq += tl.dot(ds, k.to(ACC), input_precision=PREC).to(tl.float32)
+
+    p_sink = tl.exp(sink - l_i)
+    tl.atomic_add(DSINK + h, tl.sum(tl.where(mmask, -p_sink * delta, 0.0)))
+    tl.store(
+        DQ
+        + b * stride_qb
+        + h * stride_qh
+        + offs_m[:, None] * stride_qm
+        + offs_d[None, :] * stride_qd,
+        dq.to(DQ.dtype.element_ty),
+        mask=mmask[:, None],
+    )
+
+
+@triton.autotune(
+    configs=_CONFIGS,
+    key=["H", "EL"],
+    reset_to_zero=["DK", "DV"],
+    prune_configs_by={"early_config_prune": _prune_smem(2, 2)},
+)
+@triton.jit
+def _bwd_dkdv_kernel(
+    Q,
+    K,
+    V,
+    DO,
+    L,
+    Delta,
+    DK,
+    DV,
+    scale,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kn,
+    stride_kd,
+    H,
+    S,
+    KV,
+    WINDOW,
+    EL,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    ACC: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, D)
+    nmask = offs_n < KV
+
+    k = tl.load(
+        K + b * stride_kb + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd,
+        mask=nmask[:, None],
+        other=0.0,
+    )
+    v = tl.load(
+        V + b * stride_kb + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd,
+        mask=nmask[:, None],
+        other=0.0,
+    )
+    dk = tl.zeros([BLOCK_N, D], tl.float32)
+    dv = tl.zeros([BLOCK_N, D], tl.float32)
+
+    # queries that attend to this kv block: i in [n0, n0 + BLOCK_N - 1 + WINDOW - 1]
+    lo = (pid_n * BLOCK_N // BLOCK_M) * BLOCK_M
+    hi = tl.minimum(S, pid_n * BLOCK_N + BLOCK_N + WINDOW - 1)
+    for start_m in range(lo, hi, BLOCK_M):
+        offs_m = start_m + tl.arange(0, BLOCK_M)
+        mmask = offs_m < S
+        q = tl.load(
+            Q
+            + b * stride_qb
+            + h * stride_qh
+            + offs_m[:, None] * stride_qm
+            + offs_d[None, :] * stride_qd,
+            mask=mmask[:, None],
+            other=0.0,
+        )
+        do = tl.load(
+            DO
+            + b * stride_qb
+            + h * stride_qh
+            + offs_m[:, None] * stride_qm
+            + offs_d[None, :] * stride_qd,
+            mask=mmask[:, None],
+            other=0.0,
+        )
+        l_i = tl.load(L + pid_bh * S + offs_m, mask=mmask, other=0.0)
+        delta = tl.load(Delta + pid_bh * S + offs_m, mask=mmask, other=0.0)
+        qk = tl.dot(q, tl.trans(k), input_precision=PREC) * scale
+        valid = (
+            (offs_n[None, :] <= offs_m[:, None])
+            & (offs_m[:, None] - offs_n[None, :] < WINDOW)
+            & nmask[None, :]
+        )
+        p = tl.where(valid, tl.exp(qk - l_i[:, None]), 0.0)
+        dp = tl.dot(do, tl.trans(v), input_precision=PREC)
+        ds = (p * (dp - delta[:, None]) * scale).to(ACC)
+        dk += tl.dot(tl.trans(ds), q.to(ACC), input_precision=PREC).to(tl.float32)
+        dv += tl.dot(tl.trans(p.to(ACC)), do.to(ACC), input_precision=PREC).to(
+            tl.float32
+        )
+
+    # shared MQA head: sum across q heads via atomics
+    dk_ptr = (
+        DK + b * stride_kb + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+    )
+    dv_ptr = (
+        DV + b * stride_kb + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+    )
+    tl.atomic_add(dk_ptr, dk, mask=nmask[:, None])
+    tl.atomic_add(dv_ptr, dv, mask=nmask[:, None])
+
+
+@register_kernel_op("dsv4_sliding_attn_bwd")
+def _sliding_attn_bwd_op(
+    do: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    sinks: torch.Tensor,
+    out: torch.Tensor,
+    L: torch.Tensor,
+    scale: float,
+    window: int,
+    acc_fp32: bool,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, H, S, D = q.shape
+    KV = k.shape[1]
+    do = do.contiguous()
+    delta = (do.to(torch.float32) * out.to(torch.float32)).sum(-1).reshape(B * H, S)
+    dq = torch.empty_like(q)
+    dk = torch.zeros_like(k, dtype=torch.float32)
+    dv = torch.zeros_like(v, dtype=torch.float32)
+    dsink = torch.zeros_like(sinks, dtype=torch.float32)
+    ACC = tl.float32 if acc_fp32 else tl.bfloat16
+    PREC = "ieee" if q.element_size() == 4 else "tf32"
+    EL = q.element_size()
+    args = (
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        H,
+        S,
+        KV,
+        window,
+        EL,
+    )
+    _bwd_dq_kernel[lambda m: (triton.cdiv(S, m["BLOCK_M"]), B * H)](
+        q, k, v, sinks, do, L, delta, dq, dsink, scale, *args, D=D, ACC=ACC, PREC=PREC
+    )
+    _bwd_dkdv_kernel[lambda m: (triton.cdiv(KV, m["BLOCK_N"]), B * H)](
+        q, k, v, do, L, delta, dk, dv, scale, *args, D=D, ACC=ACC, PREC=PREC
+    )
+    return dq, dk.to(k.dtype), dv.to(v.dtype), dsink.to(sinks.dtype)
+
+
+@_sliding_attn_bwd_op.register_fake
+def _(do, q, k, v, sinks, out, L, scale, window, acc_fp32):
+    return (
+        torch.empty_like(q),
+        torch.empty_like(k),
+        torch.empty_like(v),
+        torch.empty_like(sinks),
+    )
+
+
+class _SlidingAttn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, sinks, scale, window, acc_dtype):
+        out, L = _sliding_attn_fwd_op(
+            q, k, v, sinks, scale, window, acc_dtype == torch.float32
+        )
+        ctx.save_for_backward(q, k, v, sinks, out, L)
+        ctx.scale, ctx.window, ctx.acc_dtype = scale, window, acc_dtype
+        return out
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, sinks, out, L = ctx.saved_tensors
+        dq, dk, dv, dsink = _sliding_attn_bwd_op(
+            do,
+            q,
+            k,
+            v,
+            sinks,
+            out,
+            L,
+            ctx.scale,
+            ctx.window,
+            ctx.acc_dtype == torch.float32,
+        )
+        return dq, dk, dv, dsink, None, None, None
+
+
+def sliding_attn(
+    q, k, v, sinks, scale=None, sliding_window=128, acc_dtype=torch.float32
+):
+    """q: [B, H, S, D]; k, v: [B, 1, KV, D] (single MQA head); sinks: [H].
+    Returns attn output [B, H, S, D]. ``acc_dtype`` controls PV/dq/dk/dv accumulation."""
+    B, H, S, D = q.shape
+    if scale is None:
+        scale = D**-0.5
+    # match k/v/sinks to q so the kernel never sees mixed dtypes
+    dt = q.dtype
+    k = k.to(dt) if k.dtype != dt else k
+    v = v.to(dt) if v.dtype != dt else v
+    sinks = sinks.to(dt) if sinks.dtype != dt else sinks
+    k2 = k[:, 0].contiguous()
+    v2 = v[:, 0].contiguous()
+    return _SlidingAttn.apply(
+        q.contiguous(), k2, v2, sinks.contiguous(), scale, sliding_window, acc_dtype
+    )
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/attention_csa.py b/src/axolotl/integrations/kernels/libs/dsv4/attention_csa.py
new file mode 100644
index 0000000000..0594c3a5bb
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/attention_csa.py
@@ -0,0 +1,601 @@
+"""Core attention for DeepSeek-V4 CSA/HCA layers (sparse-MLA), fwd + bwd.
+
+These layers call ``eager_attention_forward`` with
+``kv = cat([sliding_kv (S), compressed_kv (T)], dim=2)`` and
+``mask = [sliding_window_causal (S) | block_bias (S, T)]``. ``block_bias`` carries the
+per-query causality + Lightning-Indexer top-k selection over the T compressed entries
+and is a constant w.r.t. attention (top-k / argmax are non-differentiable).
+
+This kernel does one fused online softmax over: the windowed sliding keys (implicit
+mask, O(S*window)), the T compressed keys (dense additive ``block_bias``), and the
+per-head sink — never materializing the [B,H,S,S+T] scores eager builds, and keeping
+the single MQA KV head in SRAM. fp32 online-softmax accumulation.
+
+Shared-KV MQA: ``kv_slide`` is used as both K and V, ``kv_comp`` likewise; their grads
+are dk+dv summed. ``block_bias`` is dense [B,S,T] (smaller than eager's full mask but
+still O(S*T) for CSA where T~S/4 — exploiting the top-k to gather only index_topk keys
+per query is a further optimization, not done here).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+from .attention import _CONFIGS, _prune_smem
+
+
+@triton.autotune(
+    configs=_CONFIGS,
+    key=["H", "EL"],
+    prune_configs_by={"early_config_prune": _prune_smem(1, 2)},
+)
+@triton.jit
+def _csa_fwd_kernel(
+    Q,
+    KS,
+    KC,
+    BB,
+    sinks,
+    Out,
+    L,
+    scale,
+    sqb,
+    sqh,
+    sqm,
+    sqd,
+    skb,
+    skn,
+    skd,
+    scb,
+    scn,
+    scd,
+    sbb,
+    sbm,
+    sbt,
+    sob,
+    soh,
+    som,
+    sod,
+    H,
+    S,
+    T,
+    WINDOW,
+    EL,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    ACC: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, D)
+    q = tl.load(
+        Q + b * sqb + h * sqh + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+        mask=offs_m[:, None] < S,
+        other=0.0,
+    )
+    sink = tl.load(sinks + h).to(tl.float32)
+    m_i = tl.zeros([BLOCK_M], tl.float32) + sink
+    l_i = tl.zeros([BLOCK_M], tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, D], tl.float32)
+
+    lo = (tl.maximum(0, pid_m * BLOCK_M - WINDOW + 1) // BLOCK_N) * BLOCK_N
+    hi = tl.minimum(S, pid_m * BLOCK_M + BLOCK_M)
+    for start_n in range(lo, hi, BLOCK_N):
+        offs_n = start_n + tl.arange(0, BLOCK_N)
+        nmask = offs_n < S
+        k = tl.load(
+            KS + b * skb + offs_n[:, None] * skn + offs_d[None, :] * skd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(k), input_precision=PREC) * scale
+        valid = (
+            (offs_n[None, :] <= offs_m[:, None])
+            & (offs_m[:, None] - offs_n[None, :] < WINDOW)
+            & nmask[None, :]
+        )
+        qk = tl.where(valid, qk, float("-inf"))
+        m_new = tl.maximum(m_i, tl.max(qk, 1))
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new[:, None])
+        l_i = l_i * alpha + tl.sum(p, 1)
+        acc = acc * alpha[:, None]
+        v = tl.load(
+            KS + b * skb + offs_n[:, None] * skn + offs_d[None, :] * skd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        acc += tl.dot(p.to(ACC), v.to(ACC), input_precision=PREC).to(tl.float32)
+        m_i = m_new
+
+    for start_t in range(0, T, BLOCK_N):
+        offs_t = start_t + tl.arange(0, BLOCK_N)
+        tmask = offs_t < T
+        kc = tl.load(
+            KC + b * scb + offs_t[:, None] * scn + offs_d[None, :] * scd,
+            mask=tmask[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(kc), input_precision=PREC) * scale
+        bb = tl.load(
+            BB + b * sbb + offs_m[:, None] * sbm + offs_t[None, :] * sbt,
+            mask=(offs_m[:, None] < S) & tmask[None, :],
+            other=float("-inf"),
+        )
+        qk = tl.where(tmask[None, :], qk + bb, float("-inf"))
+        m_new = tl.maximum(m_i, tl.max(qk, 1))
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new[:, None])
+        l_i = l_i * alpha + tl.sum(p, 1)
+        acc = acc * alpha[:, None]
+        vc = tl.load(
+            KC + b * scb + offs_t[:, None] * scn + offs_d[None, :] * scd,
+            mask=tmask[:, None],
+            other=0.0,
+        )
+        acc += tl.dot(p.to(ACC), vc.to(ACC), input_precision=PREC).to(tl.float32)
+        m_i = m_new
+
+    acc = acc / l_i[:, None]
+    tl.store(
+        Out + b * sob + h * soh + offs_m[:, None] * som + offs_d[None, :] * sod,
+        acc.to(Out.dtype.element_ty),
+        mask=offs_m[:, None] < S,
+    )
+    tl.store(L + pid_bh * S + offs_m, m_i + tl.log(l_i), mask=offs_m < S)
+
+
+@register_kernel_op("dsv4_csa_attn_fwd")
+def _csa_attn_fwd_op(
+    q: torch.Tensor,
+    ks: torch.Tensor,
+    kc: torch.Tensor,
+    bb: torch.Tensor,
+    sinks: torch.Tensor,
+    scale: float,
+    window: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, H, S, D = q.shape
+    T = kc.shape[1]
+    out = torch.empty(B, H, S, D, device=q.device, dtype=q.dtype)
+    L = torch.empty(B * H, S, device=q.device, dtype=torch.float32)
+    ACC = tl.bfloat16 if q.dtype == torch.bfloat16 else tl.float32
+    PREC = "ieee" if q.element_size() == 4 else "tf32"
+    grid = lambda m: (triton.cdiv(S, m["BLOCK_M"]), B * H)
+    _csa_fwd_kernel[grid](
+        q,
+        ks,
+        kc,
+        bb,
+        sinks,
+        out,
+        L,
+        scale,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        ks.stride(0),
+        ks.stride(1),
+        ks.stride(2),
+        kc.stride(0),
+        kc.stride(1),
+        kc.stride(2),
+        bb.stride(0),
+        bb.stride(1),
+        bb.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        H,
+        S,
+        T,
+        window,
+        q.element_size(),
+        D=D,
+        ACC=ACC,
+        PREC=PREC,
+    )
+    return out, L
+
+
+@_csa_attn_fwd_op.register_fake
+def _(q, ks, kc, bb, sinks, scale, window):
+    B, H, S, _ = q.shape
+    return torch.empty(q.shape, dtype=q.dtype, device=q.device), torch.empty(
+        B * H, S, device=q.device, dtype=torch.float32
+    )
+
+
+@triton.autotune(
+    configs=_CONFIGS,
+    key=["H", "EL"],
+    reset_to_zero=["DSINK"],
+    prune_configs_by={"early_config_prune": _prune_smem(2, 2)},
+)
+@triton.jit
+def _csa_bwd_dq_kernel(
+    Q,
+    KS,
+    KC,
+    BB,
+    sinks,
+    DO,
+    L,
+    Delta,
+    DQ,
+    DSINK,
+    scale,
+    sqb,
+    sqh,
+    sqm,
+    sqd,
+    skb,
+    skn,
+    skd,
+    scb,
+    scn,
+    scd,
+    sbb,
+    sbm,
+    sbt,
+    H,
+    S,
+    T,
+    WINDOW,
+    EL,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    ACC: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, D)
+    mmask = offs_m < S
+    q = tl.load(
+        Q + b * sqb + h * sqh + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+        mask=mmask[:, None],
+        other=0.0,
+    )
+    do = tl.load(
+        DO + b * sqb + h * sqh + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+        mask=mmask[:, None],
+        other=0.0,
+    )
+    l_i = tl.load(L + pid_bh * S + offs_m, mask=mmask, other=0.0)
+    delta = tl.load(Delta + pid_bh * S + offs_m, mask=mmask, other=0.0)
+    sink = tl.load(sinks + h).to(tl.float32)
+    dq = tl.zeros([BLOCK_M, D], tl.float32)
+
+    lo = (tl.maximum(0, pid_m * BLOCK_M - WINDOW + 1) // BLOCK_N) * BLOCK_N
+    hi = tl.minimum(S, pid_m * BLOCK_M + BLOCK_M)
+    for start_n in range(lo, hi, BLOCK_N):
+        offs_n = start_n + tl.arange(0, BLOCK_N)
+        nmask = offs_n < S
+        k = tl.load(
+            KS + b * skb + offs_n[:, None] * skn + offs_d[None, :] * skd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        v = k
+        qk = tl.dot(q, tl.trans(k), input_precision=PREC) * scale
+        valid = (
+            (offs_n[None, :] <= offs_m[:, None])
+            & (offs_m[:, None] - offs_n[None, :] < WINDOW)
+            & nmask[None, :]
+        )
+        p = tl.where(valid, tl.exp(qk - l_i[:, None]), 0.0)
+        dp = tl.dot(do, tl.trans(v), input_precision=PREC)
+        ds = (p * (dp - delta[:, None]) * scale).to(ACC)
+        dq += tl.dot(ds, k.to(ACC), input_precision=PREC).to(tl.float32)
+
+    for start_t in range(0, T, BLOCK_N):
+        offs_t = start_t + tl.arange(0, BLOCK_N)
+        tmask = offs_t < T
+        kc = tl.load(
+            KC + b * scb + offs_t[:, None] * scn + offs_d[None, :] * scd,
+            mask=tmask[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(kc), input_precision=PREC) * scale
+        bb = tl.load(
+            BB + b * sbb + offs_m[:, None] * sbm + offs_t[None, :] * sbt,
+            mask=mmask[:, None] & tmask[None, :],
+            other=float("-inf"),
+        )
+        p = tl.where(tmask[None, :], tl.exp(qk + bb - l_i[:, None]), 0.0)
+        dp = tl.dot(do, tl.trans(kc), input_precision=PREC)
+        ds = (p * (dp - delta[:, None]) * scale).to(ACC)
+        dq += tl.dot(ds, kc.to(ACC), input_precision=PREC).to(tl.float32)
+
+    p_sink = tl.exp(sink - l_i)
+    tl.atomic_add(DSINK + h, tl.sum(tl.where(mmask, -p_sink * delta, 0.0)))
+    tl.store(
+        DQ + b * sqb + h * sqh + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+        dq.to(DQ.dtype.element_ty),
+        mask=mmask[:, None],
+    )
+
+
+@triton.autotune(
+    configs=_CONFIGS,
+    key=["H", "EL"],
+    reset_to_zero=["DKS", "DKC"],
+    prune_configs_by={"early_config_prune": _prune_smem(2, 2)},
+)
+@triton.jit
+def _csa_bwd_dkv_kernel(
+    Q,
+    KS,
+    KC,
+    BB,
+    DO,
+    L,
+    Delta,
+    DKS,
+    DKC,
+    scale,
+    sqb,
+    sqh,
+    sqm,
+    sqd,
+    skb,
+    skn,
+    skd,
+    scb,
+    scn,
+    scd,
+    sbb,
+    sbm,
+    sbt,
+    H,
+    S,
+    T,
+    WINDOW,
+    EL,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    ACC: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    # pid_n indexes a key block over the combined [sliding (S) ; compressed (T)] axis.
+    pid_n = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+    offs_d = tl.arange(0, D)
+    NSLIDE_BLK = tl.cdiv(S, BLOCK_N)
+    is_comp = pid_n >= NSLIDE_BLK
+
+    if is_comp:
+        start_t = (pid_n - NSLIDE_BLK) * BLOCK_N
+        offs_n = start_t + tl.arange(0, BLOCK_N)
+        nmask = offs_n < T
+        kc = tl.load(
+            KC + b * scb + offs_n[:, None] * scn + offs_d[None, :] * scd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        dk = tl.zeros([BLOCK_N, D], tl.float32)
+        dv = tl.zeros([BLOCK_N, D], tl.float32)
+        for start_m in range(0, S, BLOCK_M):  # any query may attend any compressed key
+            offs_m = start_m + tl.arange(0, BLOCK_M)
+            mmask = offs_m < S
+            q = tl.load(
+                Q + b * sqb + h * sqh + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+                mask=mmask[:, None],
+                other=0.0,
+            )
+            do = tl.load(
+                DO + b * sqb + h * sqh + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+                mask=mmask[:, None],
+                other=0.0,
+            )
+            l_i = tl.load(L + pid_bh * S + offs_m, mask=mmask, other=0.0)
+            delta = tl.load(Delta + pid_bh * S + offs_m, mask=mmask, other=0.0)
+            qk = tl.dot(q, tl.trans(kc), input_precision=PREC) * scale
+            bb = tl.load(
+                BB + b * sbb + offs_m[:, None] * sbm + offs_n[None, :] * sbt,
+                mask=mmask[:, None] & nmask[None, :],
+                other=float("-inf"),
+            )
+            p = tl.where(
+                mmask[:, None] & nmask[None, :], tl.exp(qk + bb - l_i[:, None]), 0.0
+            )
+            dp = tl.dot(do, tl.trans(kc), input_precision=PREC)
+            ds = (p * (dp - delta[:, None]) * scale).to(ACC)
+            dk += tl.dot(tl.trans(ds), q.to(ACC), input_precision=PREC).to(tl.float32)
+            dv += tl.dot(tl.trans(p.to(ACC)), do.to(ACC), input_precision=PREC).to(
+                tl.float32
+            )
+        ptr = DKC + b * scb + offs_n[:, None] * scn + offs_d[None, :] * scd
+        tl.atomic_add(ptr, (dk + dv), mask=nmask[:, None])
+    else:
+        start_n = pid_n * BLOCK_N
+        offs_n = start_n + tl.arange(0, BLOCK_N)
+        nmask = offs_n < S
+        ks = tl.load(
+            KS + b * skb + offs_n[:, None] * skn + offs_d[None, :] * skd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        dk = tl.zeros([BLOCK_N, D], tl.float32)
+        dv = tl.zeros([BLOCK_N, D], tl.float32)
+        lo = (start_n // BLOCK_M) * BLOCK_M
+        hi = tl.minimum(S, start_n + BLOCK_N + WINDOW - 1)
+        for start_m in range(lo, hi, BLOCK_M):
+            offs_m = start_m + tl.arange(0, BLOCK_M)
+            mmask = offs_m < S
+            q = tl.load(
+                Q + b * sqb + h * sqh + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+                mask=mmask[:, None],
+                other=0.0,
+            )
+            do = tl.load(
+                DO + b * sqb + h * sqh + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+                mask=mmask[:, None],
+                other=0.0,
+            )
+            l_i = tl.load(L + pid_bh * S + offs_m, mask=mmask, other=0.0)
+            delta = tl.load(Delta + pid_bh * S + offs_m, mask=mmask, other=0.0)
+            qk = tl.dot(q, tl.trans(ks), input_precision=PREC) * scale
+            valid = (
+                (offs_n[None, :] <= offs_m[:, None])
+                & (offs_m[:, None] - offs_n[None, :] < WINDOW)
+                & nmask[None, :]
+            )
+            p = tl.where(valid, tl.exp(qk - l_i[:, None]), 0.0)
+            dp = tl.dot(do, tl.trans(ks), input_precision=PREC)
+            ds = (p * (dp - delta[:, None]) * scale).to(ACC)
+            dk += tl.dot(tl.trans(ds), q.to(ACC), input_precision=PREC).to(tl.float32)
+            dv += tl.dot(tl.trans(p.to(ACC)), do.to(ACC), input_precision=PREC).to(
+                tl.float32
+            )
+        ptr = DKS + b * skb + offs_n[:, None] * skn + offs_d[None, :] * skd
+        tl.atomic_add(ptr, (dk + dv), mask=nmask[:, None])
+
+
+@register_kernel_op("dsv4_csa_attn_bwd")
+def _csa_attn_bwd_op(
+    do: torch.Tensor,
+    q: torch.Tensor,
+    ks: torch.Tensor,
+    kc: torch.Tensor,
+    bb: torch.Tensor,
+    sinks: torch.Tensor,
+    out: torch.Tensor,
+    L: torch.Tensor,
+    scale: float,
+    window: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, H, S, D = q.shape
+    T = kc.shape[1]
+    do = do.contiguous()
+    delta = (do.to(torch.float32) * out.to(torch.float32)).sum(-1).reshape(B * H, S)
+    dq = torch.empty_like(q)
+    dks = torch.zeros_like(ks, dtype=torch.float32)
+    dkc = torch.zeros_like(kc, dtype=torch.float32)
+    dsink = torch.zeros_like(sinks, dtype=torch.float32)
+    ACC = tl.bfloat16 if q.dtype == torch.bfloat16 else tl.float32
+    PREC = "ieee" if q.element_size() == 4 else "tf32"
+    EL = q.element_size()
+    qs = (q.stride(0), q.stride(1), q.stride(2), q.stride(3))
+    ss = (ks.stride(0), ks.stride(1), ks.stride(2))
+    cs = (kc.stride(0), kc.stride(1), kc.stride(2))
+    bs = (bb.stride(0), bb.stride(1), bb.stride(2))
+    _csa_bwd_dq_kernel[lambda m: (triton.cdiv(S, m["BLOCK_M"]), B * H)](
+        q,
+        ks,
+        kc,
+        bb,
+        sinks,
+        do,
+        L,
+        delta,
+        dq,
+        dsink,
+        scale,
+        *qs,
+        *ss,
+        *cs,
+        *bs,
+        H,
+        S,
+        T,
+        window,
+        EL,
+        D=D,
+        ACC=ACC,
+        PREC=PREC,
+    )
+    _csa_bwd_dkv_kernel[
+        lambda m: (triton.cdiv(S, m["BLOCK_N"]) + triton.cdiv(T, m["BLOCK_N"]), B * H)
+    ](
+        q,
+        ks,
+        kc,
+        bb,
+        do,
+        L,
+        delta,
+        dks,
+        dkc,
+        scale,
+        *qs,
+        *ss,
+        *cs,
+        *bs,
+        H,
+        S,
+        T,
+        window,
+        EL,
+        D=D,
+        ACC=ACC,
+        PREC=PREC,
+    )
+    return dq, dks.to(ks.dtype), dkc.to(kc.dtype), dsink.to(sinks.dtype)
+
+
+@_csa_attn_bwd_op.register_fake
+def _(do, q, ks, kc, bb, sinks, out, L, scale, window):
+    return (
+        torch.empty_like(q),
+        torch.empty_like(ks),
+        torch.empty_like(kc),
+        torch.empty_like(sinks),
+    )
+
+
+class _CSAAttn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, ks, kc, bb, sinks, scale, window):
+        out, L = _csa_attn_fwd_op(q, ks, kc, bb, sinks, scale, window)
+        ctx.save_for_backward(q, ks, kc, bb, sinks, out, L)
+        ctx.scale, ctx.window = scale, window
+        return out
+
+    @staticmethod
+    def backward(ctx, do):
+        q, ks, kc, bb, sinks, out, L = ctx.saved_tensors
+        dq, dks, dkc, dsink = _csa_attn_bwd_op(
+            do, q, ks, kc, bb, sinks, out, L, ctx.scale, ctx.window
+        )
+        return dq, dks, dkc, None, dsink, None, None
+
+
+def csa_attn(q, kv_slide, kv_comp, block_bias, sinks, scale=None, sliding_window=128):
+    """Core attention for CSA/HCA layers.
+    q: [B,H,S,D]; kv_slide: [B,1,S,D] (K==V); kv_comp: [B,1,T,D] (K==V);
+    block_bias: [B,1,S,T] additive; sinks: [H]. Returns [B,H,S,D].
+
+    NOTE: the dk/dv backward over the compressed block uses ``BLOCK_N=16`` for the grid's
+    block count (matches the smallest autotune tile); the kernel masks any tail."""
+    B, H, S, D = q.shape
+    if scale is None:
+        scale = D**-0.5
+    # compressed KV can arrive fp32 (keep_in_fp32 compressor); match all to q (no mixed dtypes).
+    dt = q.dtype
+    kv_slide = kv_slide.to(dt) if kv_slide.dtype != dt else kv_slide
+    kv_comp = kv_comp.to(dt) if kv_comp.dtype != dt else kv_comp
+    sinks = sinks.to(dt) if sinks.dtype != dt else sinks
+    ks = kv_slide[:, 0].contiguous()
+    kc = kv_comp[:, 0].contiguous()
+    bb = block_bias[:, 0].contiguous()
+    return _CSAAttn.apply(
+        q.contiguous(), ks, kc, bb, sinks.contiguous(), scale, sliding_window
+    )
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/attention_gather.py b/src/axolotl/integrations/kernels/libs/dsv4/attention_gather.py
new file mode 100644
index 0000000000..34674a5d88
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/attention_gather.py
@@ -0,0 +1,469 @@
+"""Top-k GATHER sparse-MLA for DeepSeek-V4 CSA/HCA core attention (fwd + bwd).
+
+Head-batched, MMA-efficient. An earlier per-query-position layout lost tensor cores
+(each position has a distinct top-k set, so gathered keys can't be shared across a
+position-block -> scalar dots). DeepSeek's TileLang sparse-MLA avoids this by putting the
+**H query heads on the MMA M-axis**: in MLA/MQA all H heads at one position share the same
+KV and the same per-position top-k, so a gathered tile of BN keys feeds a full
+``q[BH,D] @ Kᵀ[D,BN]`` MMA. We do the same.
+
+One program = (query position s, head-tile of BH heads, batch b). It attends to the
+sliding window (≤ ``window`` causal keys, MMA) and the position's ``K`` gathered top-k
+compressed keys (gathered in BN-chunks, MMA), plus the per-head sink. Turns the compressed
+term from O(S*T) (dense) to O(S*K) while keeping tensor cores busy.
+
+``topk_idx`` is [B, S, K] int32 (-1 = invalid, as ``DeepseekV4Indexer`` returns); kv_comp
+is shared K==V (grad scatter-added). int64 offsets so it's correct past 64k context.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+from .attention import _max_m, _smem_limit
+
+# BH (head tile, MMA M-axis) must be >= 16 (tensor-core min M). bwd also holds do +
+# transposes, so it needs smaller tiles; SMEM-pruned per GPU (sm_120 ~99KB vs H100/B200 ~228KB).
+_GCFGS = [
+    triton.Config({"BH": bh, "BN": bn}, num_warps=w, num_stages=s)
+    for bh in (16, 32, 64)
+    for bn in (16, 32, 64)
+    for w in (4, 8)
+    for s in (1, 2)
+    if not (bn == 16 and bh >= 64)
+]  # N=16 wgmma at large M crashes on Hopper/Blackwell
+
+
+def _gprune(n_tiles):
+    def prune(configs, nargs, **kwargs):
+        D = kwargs["D"]
+        EL = nargs["EL"]
+        dev = torch.cuda.current_device()
+        budget = int(_smem_limit(dev) * 0.9)
+        max_m = _max_m(dev)  # Blackwell tmem caps the [BH,512] acc at BH<=32
+        kept = [
+            c
+            for c in configs
+            if (n_tiles * c.kwargs["BH"] + n_tiles * c.num_stages * c.kwargs["BN"])
+            * D
+            * EL
+            <= budget
+            and c.kwargs["BH"] <= max_m
+        ]
+        return kept or [
+            min(configs, key=lambda c: c.kwargs["BH"] * c.kwargs["BN"] * c.num_stages)
+        ]
+
+    return prune
+
+
+@triton.autotune(
+    configs=_GCFGS,
+    key=["H", "K", "EL"],
+    prune_configs_by={"early_config_prune": _gprune(1)},
+)
+@triton.jit
+def _gather_fwd_kernel(
+    Q,
+    KS,
+    KC,
+    IDX,
+    sinks,
+    Out,
+    L,
+    scale,
+    sqb,
+    sqh,
+    sqm,
+    sqd,
+    skb,
+    skn,
+    skd,
+    scb,
+    scn,
+    scd,
+    sib,
+    sim,
+    sik,
+    H,
+    S,
+    K,
+    WINDOW,
+    EL,
+    D: tl.constexpr,
+    BH: tl.constexpr,
+    BN: tl.constexpr,
+    ACC: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    NHB = tl.cdiv(H, BH)
+    s = tl.program_id(0).to(tl.int64)
+    pid = tl.program_id(1)
+    b = (pid // NHB).to(tl.int64)
+    hb = pid % NHB
+    offs_h = (hb * BH + tl.arange(0, BH)).to(tl.int64)
+    offs_d = tl.arange(0, D).to(tl.int64)
+    hmask = offs_h < H
+
+    q = tl.load(
+        Q + b * sqb + offs_h[:, None] * sqh + s * sqm + offs_d[None, :] * sqd,
+        mask=hmask[:, None],
+        other=0.0,
+    )
+    sink = tl.load(sinks + offs_h, mask=hmask, other=0.0).to(tl.float32)
+    m_i = sink
+    l_i = tl.zeros([BH], tl.float32) + 1.0
+    acc = tl.zeros([BH, D], tl.float32)
+
+    # causal sliding window for position s: j in (s-window, s]
+    lo = (tl.maximum(0, s - WINDOW + 1) // BN) * BN
+    hi = s + 1
+    for start_n in range(lo, hi, BN):
+        offs_n = start_n + tl.arange(0, BN)
+        nmask = (offs_n <= s) & (s - offs_n < WINDOW) & (offs_n < S)
+        k = tl.load(
+            KS + b * skb + offs_n[:, None] * skn + offs_d[None, :] * skd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(k), input_precision=PREC) * scale
+        qk = tl.where(nmask[None, :], qk, float("-inf"))
+        m_new = tl.maximum(m_i, tl.max(qk, 1))
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new[:, None])
+        l_i = l_i * alpha + tl.sum(p, 1)
+        acc = acc * alpha[:, None] + tl.dot(
+            p.to(ACC), k.to(ACC), input_precision=PREC
+        ).to(tl.float32)
+        m_i = m_new
+
+    for start_k in range(0, K, BN):
+        kidx = start_k + tl.arange(0, BN)
+        idx = tl.load(IDX + b * sib + s * sim + kidx * sik, mask=kidx < K, other=-1)
+        gv = idx >= 0
+        idxs = tl.where(gv, idx, 0).to(tl.int64)
+        kc = tl.load(
+            KC + b * scb + idxs[:, None] * scn + offs_d[None, :] * scd,
+            mask=gv[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(kc), input_precision=PREC) * scale
+        qk = tl.where(gv[None, :], qk, float("-inf"))
+        m_new = tl.maximum(m_i, tl.max(qk, 1))
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new[:, None])
+        l_i = l_i * alpha + tl.sum(p, 1)
+        acc = acc * alpha[:, None] + tl.dot(
+            p.to(ACC), kc.to(ACC), input_precision=PREC
+        ).to(tl.float32)
+        m_i = m_new
+
+    acc = acc / l_i[:, None]
+    tl.store(
+        Out + b * sqb + offs_h[:, None] * sqh + s * sqm + offs_d[None, :] * sqd,
+        acc.to(Out.dtype.element_ty),
+        mask=hmask[:, None],
+    )
+    tl.store(L + b * (H * S) + offs_h * S + s, m_i + tl.log(l_i), mask=hmask)
+
+
+@triton.autotune(
+    configs=_GCFGS,
+    key=["H", "K", "EL"],
+    reset_to_zero=["DKS", "DKC", "DSINK"],
+    prune_configs_by={"early_config_prune": _gprune(2)},
+)
+@triton.jit
+def _gather_bwd_kernel(
+    Q,
+    KS,
+    KC,
+    IDX,
+    sinks,
+    DO,
+    L,
+    Delta,
+    DQ,
+    DKS,
+    DKC,
+    DSINK,
+    scale,
+    sqb,
+    sqh,
+    sqm,
+    sqd,
+    skb,
+    skn,
+    skd,
+    scb,
+    scn,
+    scd,
+    sib,
+    sim,
+    sik,
+    H,
+    S,
+    K,
+    WINDOW,
+    EL,
+    D: tl.constexpr,
+    BH: tl.constexpr,
+    BN: tl.constexpr,
+    ACC: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    NHB = tl.cdiv(H, BH)
+    s = tl.program_id(0).to(tl.int64)
+    pid = tl.program_id(1)
+    b = (pid // NHB).to(tl.int64)
+    hb = pid % NHB
+    offs_h = (hb * BH + tl.arange(0, BH)).to(tl.int64)
+    offs_d = tl.arange(0, D).to(tl.int64)
+    hmask = offs_h < H
+
+    q = tl.load(
+        Q + b * sqb + offs_h[:, None] * sqh + s * sqm + offs_d[None, :] * sqd,
+        mask=hmask[:, None],
+        other=0.0,
+    )
+    do = tl.load(
+        DO + b * sqb + offs_h[:, None] * sqh + s * sqm + offs_d[None, :] * sqd,
+        mask=hmask[:, None],
+        other=0.0,
+    )
+    l_i = tl.load(L + b * (H * S) + offs_h * S + s, mask=hmask, other=0.0)
+    delta = tl.load(Delta + b * (H * S) + offs_h * S + s, mask=hmask, other=0.0)
+    sink = tl.load(sinks + offs_h, mask=hmask, other=0.0).to(tl.float32)
+    dq = tl.zeros([BH, D], tl.float32)
+
+    lo = (tl.maximum(0, s - WINDOW + 1) // BN) * BN
+    hi = s + 1
+    for start_n in range(lo, hi, BN):
+        offs_n = start_n + tl.arange(0, BN)
+        nmask = (offs_n <= s) & (s - offs_n < WINDOW) & (offs_n < S)
+        k = tl.load(
+            KS + b * skb + offs_n[:, None] * skn + offs_d[None, :] * skd,
+            mask=nmask[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(k), input_precision=PREC) * scale
+        p = tl.where(nmask[None, :], tl.exp(qk - l_i[:, None]), 0.0)
+        dp = tl.dot(do, tl.trans(k), input_precision=PREC)
+        ds = (p * (dp - delta[:, None]) * scale).to(ACC)
+        dq += tl.dot(ds, k.to(ACC), input_precision=PREC).to(tl.float32)
+        dkv = tl.dot(tl.trans(ds), q.to(ACC), input_precision=PREC).to(
+            tl.float32
+        ) + tl.dot(tl.trans(p.to(ACC)), do.to(ACC), input_precision=PREC).to(tl.float32)
+        tl.atomic_add(
+            DKS + b * skb + offs_n[:, None] * skn + offs_d[None, :] * skd,
+            dkv,
+            mask=nmask[:, None],
+        )
+
+    for start_k in range(0, K, BN):
+        kidx = start_k + tl.arange(0, BN)
+        idx = tl.load(IDX + b * sib + s * sim + kidx * sik, mask=kidx < K, other=-1)
+        gv = idx >= 0
+        idxs = tl.where(gv, idx, 0).to(tl.int64)
+        kc = tl.load(
+            KC + b * scb + idxs[:, None] * scn + offs_d[None, :] * scd,
+            mask=gv[:, None],
+            other=0.0,
+        )
+        qk = tl.dot(q, tl.trans(kc), input_precision=PREC) * scale
+        p = tl.where(gv[None, :], tl.exp(qk - l_i[:, None]), 0.0)
+        dp = tl.dot(do, tl.trans(kc), input_precision=PREC)
+        ds = (p * (dp - delta[:, None]) * scale).to(ACC)
+        dq += tl.dot(ds, kc.to(ACC), input_precision=PREC).to(tl.float32)
+        dkv = tl.dot(tl.trans(ds), q.to(ACC), input_precision=PREC).to(
+            tl.float32
+        ) + tl.dot(tl.trans(p.to(ACC)), do.to(ACC), input_precision=PREC).to(tl.float32)
+        tl.atomic_add(
+            DKC + b * scb + idxs[:, None] * scn + offs_d[None, :] * scd,
+            dkv,
+            mask=gv[:, None],
+        )
+
+    p_sink = tl.exp(sink - l_i)
+    # mask the atomic itself: zeroing the value isn't enough, offs_h is out-of-bounds for
+    # padded heads (H % BH != 0).
+    tl.atomic_add(DSINK + offs_h, tl.where(hmask, -p_sink * delta, 0.0), mask=hmask)
+    tl.store(
+        DQ + b * sqb + offs_h[:, None] * sqh + s * sqm + offs_d[None, :] * sqd,
+        dq.to(DQ.dtype.element_ty),
+        mask=hmask[:, None],
+    )
+
+
+@register_kernel_op("dsv4_csa_topk_attn_fwd")
+def _csa_topk_attn_fwd_op(
+    q: torch.Tensor,
+    ks: torch.Tensor,
+    kc: torch.Tensor,
+    idx: torch.Tensor,
+    sinks: torch.Tensor,
+    scale: float,
+    window: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, H, S, D = q.shape
+    K = idx.shape[2]
+    out = torch.empty(B, H, S, D, device=q.device, dtype=q.dtype)
+    L = torch.empty(B, H, S, device=q.device, dtype=torch.float32)
+    ACC = tl.bfloat16 if q.dtype == torch.bfloat16 else tl.float32
+    PREC = "ieee" if q.element_size() == 4 else "tf32"
+    args = (
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        ks.stride(0),
+        ks.stride(1),
+        ks.stride(2),
+        kc.stride(0),
+        kc.stride(1),
+        kc.stride(2),
+        idx.stride(0),
+        idx.stride(1),
+        idx.stride(2),
+    )
+    grid = lambda m: (S, B * triton.cdiv(H, m["BH"]))
+    _gather_fwd_kernel[grid](
+        q,
+        ks,
+        kc,
+        idx,
+        sinks,
+        out,
+        L,
+        scale,
+        *args,
+        H,
+        S,
+        K,
+        window,
+        q.element_size(),
+        D=D,
+        ACC=ACC,
+        PREC=PREC,
+    )
+    return out, L
+
+
+@_csa_topk_attn_fwd_op.register_fake
+def _(q, ks, kc, idx, sinks, scale, window):
+    B, H, S, _ = q.shape
+    return torch.empty(q.shape, dtype=q.dtype, device=q.device), torch.empty(
+        B, H, S, device=q.device, dtype=torch.float32
+    )
+
+
+@register_kernel_op("dsv4_csa_topk_attn_bwd")
+def _csa_topk_attn_bwd_op(
+    do: torch.Tensor,
+    q: torch.Tensor,
+    ks: torch.Tensor,
+    kc: torch.Tensor,
+    idx: torch.Tensor,
+    sinks: torch.Tensor,
+    out: torch.Tensor,
+    L: torch.Tensor,
+    scale: float,
+    window: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, H, S, D = q.shape
+    K = idx.shape[2]
+    do = do.contiguous()
+    delta = (do.to(torch.float32) * out.to(torch.float32)).sum(-1)  # [B,H,S]
+    dq = torch.empty_like(q)
+    dks = torch.zeros_like(ks, dtype=torch.float32)
+    dkc = torch.zeros_like(kc, dtype=torch.float32)
+    dsink = torch.zeros_like(sinks, dtype=torch.float32)
+    ACC = tl.bfloat16 if q.dtype == torch.bfloat16 else tl.float32
+    PREC = "ieee" if q.element_size() == 4 else "tf32"
+    args = (
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        ks.stride(0),
+        ks.stride(1),
+        ks.stride(2),
+        kc.stride(0),
+        kc.stride(1),
+        kc.stride(2),
+        idx.stride(0),
+        idx.stride(1),
+        idx.stride(2),
+    )
+    grid = lambda m: (S, B * triton.cdiv(H, m["BH"]))
+    _gather_bwd_kernel[grid](
+        q,
+        ks,
+        kc,
+        idx,
+        sinks,
+        do,
+        L,
+        delta,
+        dq,
+        dks,
+        dkc,
+        dsink,
+        scale,
+        *args,
+        H,
+        S,
+        K,
+        window,
+        q.element_size(),
+        D=D,
+        ACC=ACC,
+        PREC=PREC,
+    )
+    return dq, dks.to(ks.dtype), dkc.to(kc.dtype), dsink.to(sinks.dtype)
+
+
+@_csa_topk_attn_bwd_op.register_fake
+def _(do, q, ks, kc, idx, sinks, out, L, scale, window):
+    return (
+        torch.empty_like(q),
+        torch.empty_like(ks),
+        torch.empty_like(kc),
+        torch.empty_like(sinks),
+    )
+
+
+class _CSATopK(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, ks, kc, idx, sinks, scale, window):
+        out, L = _csa_topk_attn_fwd_op(q, ks, kc, idx, sinks, scale, window)
+        ctx.save_for_backward(q, ks, kc, idx, sinks, out, L)
+        ctx.scale, ctx.window = scale, window
+        return out
+
+    @staticmethod
+    def backward(ctx, do):
+        q, ks, kc, idx, sinks, out, L = ctx.saved_tensors
+        dq, dks, dkc, dsink = _csa_topk_attn_bwd_op(
+            do, q, ks, kc, idx, sinks, out, L, ctx.scale, ctx.window
+        )
+        return dq, dks, dkc, None, dsink, None, None
+
+
+def csa_attn_topk(
+    q, kv_slide, kv_comp, topk_idx, sinks, scale=None, sliding_window=128
+):
+    """Top-k gather Compressed Sparse Attention (head-batched, MMA, linear in compressed dim).
+    q: [B,H,S,D]; kv_slide/kv_comp: [B,1,*,D] (K==V); topk_idx: [B,S,K] int32 (-1=invalid);
+    sinks: [H]. Returns [B,H,S,D]."""
+    B, H, S, D = q.shape
+    if scale is None:
+        scale = D**-0.5
+    ks = kv_slide[:, 0].contiguous()
+    kc = kv_comp[:, 0].contiguous()
+    idx = topk_idx.contiguous().to(torch.int32)
+    return _CSATopK.apply(
+        q.contiguous(), ks, kc, idx, sinks.contiguous(), scale, sliding_window
+    )
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/gated_pool.py b/src/axolotl/integrations/kernels/libs/dsv4/gated_pool.py
new file mode 100644
index 0000000000..82c6a5c36d
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/gated_pool.py
@@ -0,0 +1,104 @@
+"""Fused gated-softmax pooling for DeepSeek-V4 compressors (fwd + bwd).
+
+Reference (CSA/HCA compressor §2.3.1/§2.3.2 and the indexer):
+    compressed = (kv * gate.softmax(dim=-2, dtype=fp32)).sum(dim=-2)
+i.e. per (window, channel) a softmax over the W window tokens, then a weighted sum.
+Eager runs softmax (fp32) + multiply + reduce as separate passes over [.., W, Dh];
+this fuses them into one kernel, fp32 throughout.
+
+Per output element:  out[m,d] = sum_w p[m,w,d] * kv[m,w,d],  p = softmax_w(gate)
+Backward:            d_kv = p * d_out;   d_gate = p * d_out * (kv - out)
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BD": bd}, num_warps=w)
+        for bd in (32, 64, 128)
+        for w in (2, 4, 8)
+    ],
+    key=["M", "D"],
+)
+@triton.jit
+def _pool_fwd_kernel(KV, GATE, OUT, M, D, W: tl.constexpr, BD: tl.constexpr):
+    m = tl.program_id(0)
+    d = tl.program_id(1) * BD + tl.arange(0, BD)
+    dmask = d < D
+    wkv = tl.arange(0, W)
+    off = m * (W * D) + wkv[:, None] * D + d[None, :]
+    cmask = dmask[None, :]
+    g = tl.load(GATE + off, mask=cmask, other=float("-inf")).to(tl.float32)
+    kv = tl.load(KV + off, mask=cmask, other=0.0).to(tl.float32)
+    p = tl.exp(g - tl.max(g, axis=0, keep_dims=True))
+    p = p / tl.sum(p, axis=0, keep_dims=True)
+    out = tl.sum(p * kv, axis=0)
+    tl.store(OUT + m * D + d, out, mask=dmask)
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BD": bd}, num_warps=w)
+        for bd in (32, 64, 128)
+        for w in (2, 4, 8)
+    ],
+    key=["M", "D"],
+)
+@triton.jit
+def _pool_bwd_kernel(
+    KV, GATE, DOUT, DKV, DGATE, M, D, W: tl.constexpr, BD: tl.constexpr
+):
+    m = tl.program_id(0)
+    d = tl.program_id(1) * BD + tl.arange(0, BD)
+    dmask = d < D
+    wkv = tl.arange(0, W)
+    off = m * (W * D) + wkv[:, None] * D + d[None, :]
+    cmask = dmask[None, :]
+    g = tl.load(GATE + off, mask=cmask, other=float("-inf")).to(tl.float32)
+    kv = tl.load(KV + off, mask=cmask, other=0.0).to(tl.float32)
+    p = tl.exp(g - tl.max(g, axis=0, keep_dims=True))
+    p = p / tl.sum(p, axis=0, keep_dims=True)
+    out = tl.sum(p * kv, axis=0, keep_dims=True)
+    dout = tl.load(DOUT + m * D + d, mask=dmask, other=0.0).to(tl.float32)[None, :]
+    dkv = p * dout
+    dgate = p * dout * (kv - out)
+    tl.store(DKV + off, dkv, mask=cmask)
+    tl.store(DGATE + off, dgate, mask=cmask)
+
+
+class _GatedPool(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, kv, gate):
+        M, W, D = kv.shape
+        kv = kv.contiguous()
+        gate = gate.contiguous()
+        out = torch.empty(M, D, device=kv.device, dtype=torch.float32)
+        grid = lambda m: (M, triton.cdiv(D, m["BD"]))
+        _pool_fwd_kernel[grid](kv, gate, out, M, D, W=W)
+        ctx.save_for_backward(kv, gate)
+        ctx.dims = (M, W, D)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        kv, gate = ctx.saved_tensors
+        M, W, D = ctx.dims
+        dkv = torch.empty_like(kv, dtype=torch.float32)
+        dgate = torch.empty_like(gate, dtype=torch.float32)
+        grid = lambda m: (M, triton.cdiv(D, m["BD"]))
+        _pool_bwd_kernel[grid](kv, gate, dout.contiguous(), dkv, dgate, M, D, W=W)
+        return dkv.to(kv.dtype), dgate.to(gate.dtype)
+
+
+def gated_softmax_pool(kv: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+    """kv, gate: [..., W, D]. Returns [..., D] = sum_w softmax_w(gate) * kv (fp32).
+    Drop-in for ``(kv * gate.softmax(dim=-2, dtype=fp32)).sum(dim=-2)``. W must be a power
+    of 2 (compress_rate / 2*compress_rate are: HCA 128, CSA-overlap 8)."""
+    if gate.dtype != kv.dtype:
+        gate = gate.to(kv.dtype)
+    *lead, W, D = kv.shape
+    out = _GatedPool.apply(kv.reshape(-1, W, D), gate.reshape(-1, W, D))
+    return out.reshape(*lead, D)
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/indexer.py b/src/axolotl/integrations/kernels/libs/dsv4/indexer.py
new file mode 100644
index 0000000000..73d6719cbd
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/indexer.py
@@ -0,0 +1,138 @@
+"""Fused Lightning-Indexer scorer for DeepSeek-V4 CSA (forward-only).
+
+Reference (DeepseekV4IndexerScorer.forward):
+    scores  = matmul(q.float(), ck.float().transpose).unsqueeze(1)   # [B, S, H, T]
+    scores  = relu(scores) * softmax_scale
+    weights = weights_proj(hidden).float() * weights_scaling          # [B, S, H]
+    index_scores = (scores * weights.unsqueeze(-1)).sum(dim=2)        # [B, S, T]
+
+i.e.  out[b,s,t] = softmax_scale * Σ_h w[b,s,h] · relu(Σ_d q[b,s,h,d]·ck[b,t,d]).
+
+Eager materializes the full [B,S,H,T] fp32 score tensor twice (matmul output + the
+``scores*weights`` product) — with H=64 index heads this is the #1 transient memory
+hotspot. Fusing the H-reduction collapses it straight to [B,S,T], never holding the
+H axis.
+
+FORWARD ONLY: ``index_scores`` is consumed solely by ``topk(...).indices`` (a non-
+differentiable LongTensor) to build the gather mask, so no gradient ever flows back
+through the scorer. We return a plain (no-grad) tensor — matching eager semantics — and
+skip the backward entirely.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BS": bs, "BT": bt}, num_warps=w, num_stages=s)
+        for bs in (32, 64, 128)
+        for bt in (32, 64, 128)
+        for w in (4, 8)
+        for s in (2, 3)
+    ],
+    key=["S", "T", "H"],
+)
+@triton.jit
+def _indexer_score_kernel(
+    Q,
+    CK,
+    W,
+    OUT,
+    softmax_scale,
+    S,
+    T,
+    H,
+    sq_b,
+    sq_s,
+    sq_h,
+    sq_d,
+    sck_b,
+    sck_t,
+    sck_d,
+    sw_b,
+    sw_s,
+    sw_h,
+    so_b,
+    so_s,
+    so_t,
+    DH: tl.constexpr,
+    BS: tl.constexpr,
+    BT: tl.constexpr,
+):
+    b = tl.program_id(0)
+    s0 = tl.program_id(1) * BS
+    t0 = tl.program_id(2) * BT
+    offs_s = s0 + tl.arange(0, BS)
+    offs_t = t0 + tl.arange(0, BT)
+    offs_d = tl.arange(0, DH)
+    smask = offs_s < S
+    tmask = offs_t < T
+
+    # ck is independent of h, so load it once outside the head loop
+    ck_ptr = CK + b * sck_b + offs_d[:, None] * sck_d + offs_t[None, :] * sck_t
+    ck = tl.load(ck_ptr, mask=tmask[None, :], other=0.0).to(tl.float32)
+
+    acc = tl.zeros((BS, BT), dtype=tl.float32)
+    for h in range(H):
+        q_ptr = (
+            Q + b * sq_b + offs_s[:, None] * sq_s + h * sq_h + offs_d[None, :] * sq_d
+        )
+        q = tl.load(q_ptr, mask=smask[:, None], other=0.0).to(tl.float32)
+        qk = tl.dot(q, ck, input_precision="ieee")
+        qk = tl.maximum(qk, 0.0) * softmax_scale
+        w = tl.load(W + b * sw_b + offs_s * sw_s + h * sw_h, mask=smask, other=0.0).to(
+            tl.float32
+        )
+        acc += qk * w[:, None]
+
+    out_ptr = OUT + b * so_b + offs_s[:, None] * so_s + offs_t[None, :] * so_t
+    tl.store(out_ptr, acc, mask=smask[:, None] & tmask[None, :])
+
+
+def indexer_scores(
+    q: torch.Tensor,
+    compressed_kv: torch.Tensor,
+    weights: torch.Tensor,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """q [B,S,H,Dh], compressed_kv [B,T,Dh], weights [B,S,H] (already ·weights_scaling).
+    Returns index_scores [B,S,T] (fp32, no grad). Drop-in for DeepseekV4IndexerScorer."""
+    B, S, H, DH = q.shape
+    T = compressed_kv.shape[1]
+    out = torch.empty(B, S, T, device=q.device, dtype=torch.float32)
+    if T == 0 or S == 0:
+        return out
+    # compressed_kv may arrive fp32 (keep_in_fp32 compressor) while q is bf16.
+    if compressed_kv.dtype != q.dtype:
+        compressed_kv = compressed_kv.to(q.dtype)
+    q = q.contiguous()
+    ck = compressed_kv.contiguous()
+    w = weights.contiguous()
+    grid = lambda m: (B, triton.cdiv(S, m["BS"]), triton.cdiv(T, m["BT"]))
+    _indexer_score_kernel[grid](
+        q,
+        ck,
+        w,
+        out,
+        float(softmax_scale),
+        S,
+        T,
+        H,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        ck.stride(0),
+        ck.stride(1),
+        ck.stride(2),
+        w.stride(0),
+        w.stride(1),
+        w.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        DH=DH,
+    )
+    return out
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/lora_fp8.py b/src/axolotl/integrations/kernels/libs/dsv4/lora_fp8.py
new file mode 100644
index 0000000000..896afeb818
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/lora_fp8.py
@@ -0,0 +1,147 @@
+"""Native blockwise-FP8 fused LoRA for DeepSeek-V4 non-expert projections.
+
+PEFT's default path runs the frozen FP8 base as a dequant→bf16 GEMM (no fp8 throughput) and
+the LoRA as a separate pass. For the *large-output* attention projections (``q_b_proj`` 1024→
+32768, ``o_b_proj`` 8192→4096) a native blockwise-FP8 GEMM is 2.4–3.2× faster than the
+dequant path (benched); the small ones (``q_a``/``kv``) are a wash, so we patch only the big
+two.
+
+``torch._scaled_mm`` (and DeepGEMM) have no autograd formula, so the native fp8 forward has to
+live inside a custom ``autograd.Function``:
+  * forward  — ``fp8_linear`` (DeepGEMM 128×128 on B200, Triton fallback elsewhere) for the
+    frozen base, fused with the LoRA update via ``addmm_``;
+  * backward — ``dX`` through a bf16 dequant of the weight (Option A: the block scale sits on
+    the contraction axis, so an fp8 ``dX`` would need a transposed fp8 copy = 2× weight mem),
+    plus the LoRA ``dA``/``dB``. The base weight is frozen (no weight grad).
+
+Weight stays 128×128 blockwise (lossless vs the checkpoint); only the activation is quantized
+(1×128), matching the DeepSeek/NVIDIA recipe. Gated by config ``dsv4_fp8_lora_kernel``.
+"""
+
+from __future__ import annotations
+
+import types
+
+import torch
+
+from axolotl.kernels.lora import get_lora_parameters
+from axolotl.kernels.quantize import dequantize_fp8
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# Only the large-output projections where native fp8 beats the dequant GEMM.
+_TARGET_SUFFIXES = ("self_attn.q_b_proj", "self_attn.o_b_proj")
+
+
+def _fp8_linear():
+    from transformers.integrations.finegrained_fp8 import fp8_linear
+
+    return fp8_linear
+
+
+class _NativeFp8Lora(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, qdata, scale_inv, block_size, A, B, scaling, bias):
+        fp8_linear = _fp8_linear()
+        shape = x.shape
+        x2 = x.reshape(-1, shape[-1])
+        # frozen base: native blockwise fp8 GEMM (activation quantized 1x128 inside)
+        out = fp8_linear(
+            x2,
+            qdata,
+            scale_inv,
+            block_size=list(block_size),
+            bias=bias,
+            output_dtype=x.dtype,
+        )
+        # fused LoRA update via addmm_ (avoids a [M,out] temp)
+        xA = x2 @ A.t()
+        out = out.addmm_(xA, B.t(), alpha=float(scaling))
+        ctx.save_for_backward(x2, qdata, scale_inv, A, B, xA)
+        ctx.scaling = float(scaling)
+        ctx.block_size = tuple(block_size)
+        ctx.shape = shape
+        return out.view(*shape[:-1], -1)
+
+    @staticmethod
+    def backward(ctx, grad):
+        x2, qdata, scale_inv, A, B, xA = ctx.saved_tensors
+        s = ctx.scaling
+        g = grad.reshape(-1, grad.shape[-1])
+        # dX through a bf16 dequant of W (Option A: block scale on the contraction axis, so an
+        # fp8 dX would need a transposed fp8 weight copy = 2x weight mem). W transient.
+        w_deq = dequantize_fp8(qdata, scale_inv, g.dtype)
+        dxA = g @ B
+        dX = (g @ w_deq).addmm_(dxA, A, alpha=s)
+        dA = (dxA.t() * s) @ x2  # [r, in]
+        dB = (g.t() * s) @ xA  # [out, r]
+        return dX.view(*ctx.shape), None, None, None, dA, dB, None, None
+
+
+def _fp8_from_weight(w, base_layer):
+    """Return (qdata, scale_inv, block_size) for a Float8Tensor- or raw-FP8 weight, else None."""
+    qdata = getattr(w, "qdata", None)
+    if qdata is not None:  # torchao Float8Tensor
+        return qdata, w.scale, list(getattr(w, "block_size", None) or [128, 128])
+    if (
+        isinstance(w, torch.Tensor) and w.dtype == torch.float8_e4m3fn
+    ):  # raw transformers FP8
+        si = getattr(base_layer, "weight_scale_inv", None)
+        return (
+            (w, si, list(getattr(base_layer, "block_size", None) or [128, 128]))
+            if si is not None
+            else None
+        )
+    return None
+
+
+def _base_is_fp8(base_layer):
+    return _fp8_from_weight(getattr(base_layer, "weight", None), base_layer) is not None
+
+
+def _make_forward(lora_layer):
+    def forward(self, x, *args, **kwargs):
+        # get_lora_parameters unshards the LoRA A/B DTensors (FSDP2), grad-tracked. The base
+        # Float8Tensor is all-gathered by FSDP before forward, so this reads the live weight.
+        W, bias, _qs, A, B, scaling, _lora_bias, dropout, _mag = get_lora_parameters(
+            self
+        )
+        fp8 = _fp8_from_weight(W, self.base_layer)
+        if (
+            fp8 is None or A is None
+        ):  # disabled/merged adapter, or base not fp8: fall back
+            return self._dsv4_orig_forward(x, *args, **kwargs)
+        qdata, scale_inv, block_size = fp8
+        xin = dropout(x) if dropout is not None else x
+        return _NativeFp8Lora.apply(
+            xin, qdata, scale_inv, block_size, A, B, scaling, bias
+        )
+
+    return forward
+
+
+def patch_dsv4_attn_fp8_lora(model, enabled: bool = False) -> int:
+    """Swap the forward of LoRA-wrapped large attention projections (q_b/o_b) on an FP8 base
+    for the native-blockwise-fp8 fused-LoRA Function. No-op unless ``enabled``
+    (config ``dsv4_fp8_lora_kernel``)."""
+    if not enabled:
+        return 0
+    from peft.tuners.lora.layer import LoraLayer
+
+    n = 0
+    for name, mod in model.named_modules():
+        if not isinstance(mod, LoraLayer):
+            continue
+        if not any(name.endswith(suf) for suf in _TARGET_SUFFIXES):
+            continue
+        if not _base_is_fp8(mod.base_layer) or hasattr(mod, "_dsv4_orig_forward"):
+            continue
+        mod._dsv4_orig_forward = mod.forward  # fallback for disabled/merged adapter
+        mod.forward = types.MethodType(_make_forward(mod), mod)
+        n += 1
+    if n:
+        LOG.info(
+            "Patched %d DeepSeek-V4 attention projections with native-fp8 fused LoRA", n
+        )
+    return n
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/lora_mlp.py b/src/axolotl/integrations/kernels/libs/dsv4/lora_mlp.py
new file mode 100644
index 0000000000..7823fc9c30
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/lora_mlp.py
@@ -0,0 +1,144 @@
+"""Fused LoRA MLP for DeepSeek-V4's shared-expert (clamped SwiGLU).
+
+V4's ``DeepseekV4MLP`` is ``down(silu(gate(x).clamp(max=L)) * up(x).clamp(-L, L))`` — the
+``swiglu_limit`` clamp is what stops axolotl's stock ``apply_lora_mlp_swiglu`` (no clamp)
+from applying. ``LoRA_MLP`` is parameterized by the activation fns, so we reuse all of its
+LoRA-matmul fusion and just swap in a clamped SwiGLU activation (forward + backward).
+
+Only the shared expert is a plain SwiGLU MLP; the 256 routed experts go through
+``scattermoe_lora`` (use_scattermoe). The MLA attention projections have no fused-LoRA
+form and stay on PEFT default.
+"""
+
+import functools
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.lora import LoRA_MLP, _apply_dropout, get_lora_parameters
+
+
+@triton.jit
+def _clamped_swiglu_fwd_kernel(
+    gate_ptr, up_ptr, out_ptr, n_elements, LIMIT, block_size: tl.constexpr
+):
+    off = tl.program_id(0) * block_size + tl.arange(0, block_size)
+    mask = off < n_elements
+    gate = tl.load(gate_ptr + off, mask=mask, other=0).to(tl.float32)
+    up = tl.load(up_ptr + off, mask=mask, other=0).to(tl.float32)
+    gate = tl.minimum(gate, LIMIT)  # clamp(max=L)
+    up = tl.minimum(tl.maximum(up, -LIMIT), LIMIT)  # clamp(-L, L)
+    f = gate * tl.sigmoid(gate)
+    tl.store(out_ptr + off, (f * up).to(out_ptr.dtype.element_ty), mask=mask)
+
+
+@triton.jit
+def _clamped_swiglu_bwd_kernel(
+    grad_out_ptr, gate_ptr, up_ptr, n_elements, LIMIT, block_size: tl.constexpr
+):
+    off = tl.program_id(0) * block_size + tl.arange(0, block_size)
+    mask = off < n_elements
+    grad_out = tl.load(grad_out_ptr + off, mask=mask, other=0).to(tl.float32)
+    g_raw = tl.load(gate_ptr + off, mask=mask, other=0).to(tl.float32)
+    u_raw = tl.load(up_ptr + off, mask=mask, other=0).to(tl.float32)
+    gate = tl.minimum(g_raw, LIMIT)
+    up = tl.minimum(tl.maximum(u_raw, -LIMIT), LIMIT)
+    sig = tl.sigmoid(gate)
+    silu = sig * gate
+    h = silu * up
+    # clamp gradients are 0 where the input was clamped
+    g_mask = (g_raw < LIMIT).to(tl.float32)
+    u_mask = ((u_raw > -LIMIT) & (u_raw < LIMIT)).to(tl.float32)
+    grad_up = grad_out * silu * u_mask
+    grad_gate = grad_out * up * sig * (1.0 + gate * (1.0 - sig)) * g_mask
+    ety = grad_out_ptr.dtype.element_ty
+    tl.store(grad_out_ptr + off, h.to(ety), mask=mask)
+    tl.store(gate_ptr + off, grad_gate.to(ety), mask=mask)
+    tl.store(up_ptr + off, grad_up.to(ety), mask=mask)
+
+
+def _clamped_swiglu_forward(gate, up, limit):
+    n = gate.numel()
+    out = torch.empty_like(gate)
+    grid = lambda m: (triton.cdiv(n, m["block_size"]),)
+    _clamped_swiglu_fwd_kernel[grid](gate, up, out, n, float(limit), block_size=1024)
+    return out
+
+
+def _clamped_swiglu_backward(grad_output, gate, up, limit):
+    n = grad_output.numel()
+    grid = lambda m: (triton.cdiv(n, m["block_size"]),)
+    _clamped_swiglu_bwd_kernel[grid](
+        grad_output, gate, up, n, float(limit), block_size=1024
+    )
+    return grad_output, gate, up
+
+
+def apply_lora_mlp_clamped_swiglu(self, X, inplace: bool = False):
+    """Drop-in ``DeepseekV4MLP.forward`` using fused LoRA + clamped SwiGLU (``self.limit``).
+
+    ``inplace=False``: the shared expert's input aliases the MoE residual/routed-experts
+    input, so in-place intermediates would corrupt autograd."""
+    gateW, gateb, gateQ, gateA, gateB, gateS, gateLB, gateDrop, gateMag = (
+        get_lora_parameters(self.gate_proj)
+    )
+    upW, upb, upQ, upA, upB, upS, upLB, upDrop, upMag = get_lora_parameters(
+        self.up_proj
+    )
+    downW, downb, downQ, downA, downB, downS, downLB, downDrop, downMag = (
+        get_lora_parameters(self.down_proj)
+    )
+    X_drop = _apply_dropout(gateDrop, X, self.training)
+    fwd = functools.partial(_clamped_swiglu_forward, limit=self.limit)
+    bwd = functools.partial(_clamped_swiglu_backward, limit=self.limit)
+    return LoRA_MLP.apply(
+        X,
+        X_drop,
+        gateW,
+        gateb,
+        gateQ,
+        gateA,
+        gateB,
+        gateS,
+        gateLB,
+        gateMag,
+        upW,
+        upb,
+        upQ,
+        upA,
+        upB,
+        upS,
+        upLB,
+        upMag,
+        downW,
+        downb,
+        downQ,
+        downA,
+        downB,
+        downS,
+        downLB,
+        downMag,
+        fwd,
+        bwd,
+        inplace,
+    )
+
+
+def patch_dsv4_shared_mlp_lora(model):
+    """Swap each LoRA'd shared-expert ``DeepseekV4MLP.forward`` for the fused clamped-SwiGLU
+    LoRA kernel. Per-instance (only where gate/up/down carry LoRA), like axolotl's core
+    lora_kernels. Returns the count patched."""
+    import types
+
+    n = 0
+    for module in model.modules():
+        if type(module).__name__ != "DeepseekV4MLP":
+            continue
+        if all(
+            hasattr(getattr(module, p, None), "lora_A")
+            for p in ("gate_proj", "up_proj", "down_proj")
+        ):
+            module.forward = types.MethodType(apply_lora_mlp_clamped_swiglu, module)
+            n += 1
+    return n
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/mhc.py b/src/axolotl/integrations/kernels/libs/dsv4/mhc.py
new file mode 100644
index 0000000000..d2e17e30e9
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/mhc.py
@@ -0,0 +1,609 @@
+"""Fused Triton mHC mixer for DeepSeek-V4 HyperConnection (fwd + bwd).
+
+``DeepseekV4HyperConnection`` runs twice per layer. Its FLOPs are tiny but the
+forward/backward are dominated by ~40 tiny kernel launches from the 20-iteration
+Sinkhorn-Knopp loop over [B,S,4,4] matrices. torch.compile fuses the launches but
+mis-compiles ``collapsed`` to NaN on torch 2.11, so we fuse the
+sigmoid/softmax/Sinkhorn "mixer" (mixes -> pre, post, comb) into a single Triton
+kernel pair instead. The rmsnorm + fn-linear and the ``collapsed`` reduction stay in
+eager torch (correct, memory-bound, not the launch bottleneck).
+
+Sinkhorn (matches reference): comb = softmax(cl, -1) + eps; colnorm; then
+(iters-1)x [rownorm, colnorm]. Backward recomputes the forward, stores each
+normalize's input, and reverse-modes through them (normalize backward:
+dx = (dy - sum(dy*y, axis)) / (sum(x,axis)+eps)).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BM": bm}, num_warps=w) for bm in (8, 16, 32, 64) for w in (2, 4)
+    ],
+    key=["M"],
+)
+@triton.jit
+def _mhc_fwd_kernel(
+    MIX,
+    SCALE,
+    BASE,
+    PRE,
+    POST,
+    COMB,
+    STATES,
+    M,
+    HC: tl.constexpr,
+    MIX_N: tl.constexpr,
+    ITERS: tl.constexpr,
+    NSTATES: tl.constexpr,
+    EPS: tl.constexpr,
+    POSTMULT: tl.constexpr,
+    SAVE: tl.constexpr,
+    BM: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BM + tl.arange(0, BM)
+    mask = offs < M
+    j = tl.arange(0, HC)
+    jj = tl.arange(0, HC * HC)
+    s0 = tl.load(SCALE + 0)
+    s1 = tl.load(SCALE + 1)
+    s2 = tl.load(SCALE + 2)
+
+    pre_w = tl.load(
+        MIX + offs[:, None] * MIX_N + j[None, :], mask=mask[:, None], other=0.0
+    )
+    post_w = tl.load(
+        MIX + offs[:, None] * MIX_N + (HC + j)[None, :], mask=mask[:, None], other=0.0
+    )
+    comb_w = tl.load(
+        MIX + offs[:, None] * MIX_N + (2 * HC + jj)[None, :],
+        mask=mask[:, None],
+        other=0.0,
+    )
+    pre_b = tl.load(BASE + j)
+    post_b = tl.load(BASE + HC + j)
+    comb_b = tl.load(BASE + 2 * HC + jj)
+
+    pre = tl.sigmoid(pre_w * s0 + pre_b[None, :]) + EPS
+    post = POSTMULT * tl.sigmoid(post_w * s1 + post_b[None, :])
+    cl = tl.reshape(comb_w * s2 + comb_b[None, :], (BM, HC, HC))
+
+    ex = tl.exp(cl - tl.max(cl, axis=2, keep_dims=True))
+    comb = ex / tl.sum(ex, axis=2, keep_dims=True) + EPS
+
+    # Store each normalize's INPUT to scratch so backward can reverse-mode the Sinkhorn.
+    t = 0
+    if SAVE:
+        tl.store(
+            STATES + offs[:, None] * (NSTATES * HC * HC) + t * (HC * HC) + jj[None, :],
+            tl.reshape(comb, (BM, HC * HC)),
+            mask=mask[:, None],
+        )
+    t += 1
+    comb = comb / (tl.sum(comb, axis=1, keep_dims=True) + EPS)  # colnorm
+    for _ in range(ITERS - 1):
+        if SAVE:
+            tl.store(
+                STATES
+                + offs[:, None] * (NSTATES * HC * HC)
+                + t * (HC * HC)
+                + jj[None, :],
+                tl.reshape(comb, (BM, HC * HC)),
+                mask=mask[:, None],
+            )
+        t += 1
+        comb = comb / (tl.sum(comb, axis=2, keep_dims=True) + EPS)  # rownorm
+        if SAVE:
+            tl.store(
+                STATES
+                + offs[:, None] * (NSTATES * HC * HC)
+                + t * (HC * HC)
+                + jj[None, :],
+                tl.reshape(comb, (BM, HC * HC)),
+                mask=mask[:, None],
+            )
+        t += 1
+        comb = comb / (tl.sum(comb, axis=1, keep_dims=True) + EPS)  # colnorm
+
+    tl.store(PRE + offs[:, None] * HC + j[None, :], pre, mask=mask[:, None])
+    tl.store(POST + offs[:, None] * HC + j[None, :], post, mask=mask[:, None])
+    tl.store(
+        COMB + offs[:, None] * (HC * HC) + jj[None, :],
+        tl.reshape(comb, (BM, HC * HC)),
+        mask=mask[:, None],
+    )
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BM": bm}, num_warps=w) for bm in (8, 16, 32, 64) for w in (2, 4)
+    ],
+    key=["M"],
+    reset_to_zero=["DSCALE", "DBASE"],
+)
+@triton.jit
+def _mhc_bwd_kernel(
+    MIX,
+    SCALE,
+    BASE,
+    STATES,
+    DPRE,
+    DPOST,
+    DCOMB,
+    DMIX,
+    DSCALE,
+    DBASE,
+    M,
+    HC: tl.constexpr,
+    MIX_N: tl.constexpr,
+    ITERS: tl.constexpr,
+    NSTATES: tl.constexpr,
+    EPS: tl.constexpr,
+    POSTMULT: tl.constexpr,
+    BM: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BM + tl.arange(0, BM)
+    mask = offs < M
+    j = tl.arange(0, HC)
+    jj = tl.arange(0, HC * HC)
+    s0 = tl.load(SCALE + 0)
+    s1 = tl.load(SCALE + 1)
+    s2 = tl.load(SCALE + 2)
+
+    pre_w = tl.load(
+        MIX + offs[:, None] * MIX_N + j[None, :], mask=mask[:, None], other=0.0
+    )
+    post_w = tl.load(
+        MIX + offs[:, None] * MIX_N + (HC + j)[None, :], mask=mask[:, None], other=0.0
+    )
+    comb_w = tl.load(
+        MIX + offs[:, None] * MIX_N + (2 * HC + jj)[None, :],
+        mask=mask[:, None],
+        other=0.0,
+    )
+    pre_b = tl.load(BASE + j)
+    post_b = tl.load(BASE + HC + j)
+    _comb_b = tl.load(BASE + 2 * HC + jj)
+
+    pre_sig = tl.sigmoid(pre_w * s0 + pre_b[None, :])
+    post_sig = tl.sigmoid(post_w * s1 + post_b[None, :])
+    # sm = softmax(cl) = STATES[0] - eps (state 0 saved before colnorm)
+    sm = (
+        tl.reshape(
+            tl.load(
+                STATES
+                + offs[:, None] * (NSTATES * HC * HC)
+                + 0 * (HC * HC)
+                + jj[None, :],
+                mask=mask[:, None],
+                other=0.0,
+            ),
+            (BM, HC, HC),
+        )
+        - EPS
+    )
+
+    dy = tl.reshape(
+        tl.load(
+            DCOMB + offs[:, None] * (HC * HC) + jj[None, :],
+            mask=mask[:, None],
+            other=0.0,
+        ),
+        (BM, HC, HC),
+    )
+    # reverse the Sinkhorn normalizes. Forward: colnorm(idx0), then (iters-1)x [rownorm,
+    # colnorm]; undo last-first (colnorm=axis1, rownorm=axis2).
+    base_st = offs[:, None] * (NSTATES * HC * HC) + jj[None, :]
+    t = 2 * (ITERS - 1)  # last (col) state index
+    xc = tl.reshape(
+        tl.load(STATES + base_st + t * (HC * HC), mask=mask[:, None], other=0.0),
+        (BM, HC, HC),
+    )
+    Sc = tl.sum(xc, axis=1, keep_dims=True) + EPS
+    dy = (dy - tl.sum(dy * (xc / Sc), axis=1, keep_dims=True)) / Sc
+    t -= 1
+    for _ in range(ITERS - 1):
+        xr = tl.reshape(
+            tl.load(STATES + base_st + t * (HC * HC), mask=mask[:, None], other=0.0),
+            (BM, HC, HC),
+        )
+        Sr = tl.sum(xr, axis=2, keep_dims=True) + EPS
+        dy = (dy - tl.sum(dy * (xr / Sr), axis=2, keep_dims=True)) / Sr
+        t -= 1
+        xc2 = tl.reshape(
+            tl.load(STATES + base_st + t * (HC * HC), mask=mask[:, None], other=0.0),
+            (BM, HC, HC),
+        )
+        Sc2 = tl.sum(xc2, axis=1, keep_dims=True) + EPS
+        dy = (dy - tl.sum(dy * (xc2 / Sc2), axis=1, keep_dims=True)) / Sc2
+        t -= 1
+    # dy = dA = d(sm); softmax backward over axis=2
+    dcl = sm * (dy - tl.sum(dy * sm, axis=2, keep_dims=True))
+    dcl_flat = tl.reshape(dcl, (BM, HC * HC))
+
+    dpre = tl.load(
+        DPRE + offs[:, None] * HC + j[None, :], mask=mask[:, None], other=0.0
+    )
+    dpost = tl.load(
+        DPOST + offs[:, None] * HC + j[None, :], mask=mask[:, None], other=0.0
+    )
+    dpre_z = dpre * pre_sig * (1.0 - pre_sig)
+    dpost_z = dpost * POSTMULT * post_sig * (1.0 - post_sig)
+
+    d_pre_w = dpre_z * s0
+    d_post_w = dpost_z * s1
+    d_comb_w = dcl_flat * s2
+    tl.store(DMIX + offs[:, None] * MIX_N + j[None, :], d_pre_w, mask=mask[:, None])
+    tl.store(
+        DMIX + offs[:, None] * MIX_N + (HC + j)[None, :], d_post_w, mask=mask[:, None]
+    )
+    tl.store(
+        DMIX + offs[:, None] * MIX_N + (2 * HC + jj)[None, :],
+        d_comb_w,
+        mask=mask[:, None],
+    )
+
+    m2 = mask[:, None]
+    tl.atomic_add(DSCALE + 0, tl.sum(tl.where(m2, dpre_z * pre_w, 0.0)))
+    tl.atomic_add(DSCALE + 1, tl.sum(tl.where(m2, dpost_z * post_w, 0.0)))
+    tl.atomic_add(DSCALE + 2, tl.sum(tl.where(m2, dcl_flat * comb_w, 0.0)))
+    tl.atomic_add(DBASE + j, tl.sum(tl.where(m2, dpre_z, 0.0), axis=0))
+    tl.atomic_add(DBASE + HC + j, tl.sum(tl.where(m2, dpost_z, 0.0), axis=0))
+    tl.atomic_add(DBASE + 2 * HC + jj, tl.sum(tl.where(m2, dcl_flat, 0.0), axis=0))
+
+
+class _SinkhornMix(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, mixes, scale, base, hc, iters, eps, post_mult):
+        M, MIX_N = mixes.shape
+        mixes = mixes.contiguous().float()
+        scale = scale.contiguous().float()
+        base = base.contiguous().float()
+        nstates = 1 + 2 * (iters - 1)
+        pre = torch.empty(M, hc, device=mixes.device, dtype=torch.float32)
+        post = torch.empty(M, hc, device=mixes.device, dtype=torch.float32)
+        comb = torch.empty(M, hc * hc, device=mixes.device, dtype=torch.float32)
+        # Allocate state history if ANY trainable input needs grad (backward reads it for the
+        # scale/base grads too, not just mixes). is_grad_enabled() is False in Function.forward.
+        save = any(ctx.needs_input_grad)
+        states = (
+            torch.empty(M, nstates, hc * hc, device=mixes.device, dtype=torch.float32)
+            if save
+            else torch.empty(1, device=mixes.device, dtype=torch.float32)
+        )
+        grid = lambda m: (triton.cdiv(M, m["BM"]),)
+        _mhc_fwd_kernel[grid](
+            mixes,
+            scale,
+            base,
+            pre,
+            post,
+            comb,
+            states,
+            M,
+            HC=hc,
+            MIX_N=MIX_N,
+            ITERS=iters,
+            NSTATES=nstates,
+            EPS=eps,
+            POSTMULT=post_mult,
+            SAVE=save,
+        )
+        ctx.save_for_backward(mixes, scale, base, states)
+        ctx.cfg = (hc, MIX_N, iters, nstates, eps, post_mult)
+        return pre, post, comb
+
+    @staticmethod
+    def backward(ctx, dpre, dpost, dcomb):
+        mixes, scale, base, states = ctx.saved_tensors
+        hc, MIX_N, iters, nstates, eps, post_mult = ctx.cfg
+        M = mixes.shape[0]
+        dmix = torch.empty_like(mixes)
+        dscale = torch.zeros_like(scale)
+        dbase = torch.zeros_like(base)
+        grid = lambda m: (triton.cdiv(M, m["BM"]),)
+        _mhc_bwd_kernel[grid](
+            mixes,
+            scale,
+            base,
+            states,
+            dpre.contiguous().float(),
+            dpost.contiguous().float(),
+            dcomb.contiguous().float(),
+            dmix,
+            dscale,
+            dbase,
+            M,
+            HC=hc,
+            MIX_N=MIX_N,
+            ITERS=iters,
+            NSTATES=nstates,
+            EPS=eps,
+            POSTMULT=post_mult,
+        )
+        return dmix, dscale, dbase, None, None, None, None
+
+
+def sinkhorn_mix(mixes, scale, base, hc, iters, eps, post_mult=2.0):
+    """mixes: [M, hc*(2+hc)]; scale: [3]; base: [hc*(2+hc)]. Returns (pre, post, comb)
+    with pre/post [M, hc], comb [M, hc*hc]. Fused sigmoid/softmax/Sinkhorn (fp32)."""
+    return _SinkhornMix.apply(mixes, scale, base, hc, iters, eps, post_mult)
+
+
+# Fused RMSNorm + fn-linear. RMS scale r is per-row, so mixes = r * (streams @ fn^T): one
+# pass over the 16384-wide streams covers both streams@fn^T and sum-of-squares (eager does
+# ~3 passes + a 536MB fp32 flat).
+
+_RL_CFGS = [
+    triton.Config({"BK": bk}, num_warps=w, num_stages=s)
+    for bk in (64, 128, 256)
+    for w in (4, 8)
+    for s in (2, 3)
+]
+
+
+@triton.autotune(configs=_RL_CFGS, key=["M", "K", "N"])
+@triton.jit
+def _rmsln_fwd_kernel(
+    STREAMS,
+    FN,
+    MIXES,
+    R,
+    G,
+    M,
+    K,
+    EPS,
+    N: tl.constexpr,
+    NP: tl.constexpr,
+    BM: tl.constexpr,
+    BK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs_m = pid * BM + tl.arange(0, BM)
+    mmask = offs_m < M
+    n = tl.arange(0, NP)
+    nmask = n < N
+    acc = tl.zeros([BM, NP], tl.float32)
+    ssq = tl.zeros([BM], tl.float32)
+    for k0 in range(0, K, BK):
+        kk = k0 + tl.arange(0, BK)
+        kmask = kk < K
+        x = tl.load(
+            STREAMS + offs_m[:, None] * K + kk[None, :],
+            mask=mmask[:, None] & kmask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        ssq += tl.sum(x * x, axis=1)
+        w = tl.load(
+            FN + n[:, None] * K + kk[None, :],
+            mask=nmask[:, None] & kmask[None, :],
+            other=0.0,
+        )
+        acc += tl.dot(x, tl.trans(w), input_precision="ieee")
+    r = tl.rsqrt(ssq / K + EPS)
+    mixes = acc * r[:, None]
+    tl.store(
+        MIXES + offs_m[:, None] * N + n[None, :],
+        mixes,
+        mask=mmask[:, None] & nmask[None, :],
+    )
+    tl.store(R + offs_m, r, mask=mmask)
+    tl.store(
+        G + offs_m[:, None] * N + n[None, :], acc, mask=mmask[:, None] & nmask[None, :]
+    )
+
+
+@triton.autotune(configs=_RL_CFGS, key=["M", "K", "N"], reset_to_zero=["DFN"])
+@triton.jit
+def _rmsln_bwd_kernel(
+    STREAMS,
+    FN,
+    R,
+    G,
+    DMIX,
+    DSTREAMS,
+    DFN,
+    M,
+    K,
+    N: tl.constexpr,
+    NP: tl.constexpr,
+    BM: tl.constexpr,
+    BK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs_m = pid * BM + tl.arange(0, BM)
+    mmask = offs_m < M
+    n = tl.arange(0, NP)
+    nmask = n < N
+    dmix = tl.load(
+        DMIX + offs_m[:, None] * N + n[None, :],
+        mask=mmask[:, None] & nmask[None, :],
+        other=0.0,
+    )
+    g = tl.load(
+        G + offs_m[:, None] * N + n[None, :],
+        mask=mmask[:, None] & nmask[None, :],
+        other=0.0,
+    )
+    r = tl.load(R + offs_m, mask=mmask, other=0.0)
+    c = tl.sum(dmix * g, axis=1)  # dL/dr per row
+    dG = dmix * r[:, None]  # [BM, NP]
+    coef = (r * r * r) / K * c  # [BM]
+    for k0 in range(0, K, BK):
+        kk = k0 + tl.arange(0, BK)
+        kmask = kk < K
+        w = tl.load(
+            FN + n[:, None] * K + kk[None, :],
+            mask=nmask[:, None] & kmask[None, :],
+            other=0.0,
+        )
+        x = tl.load(
+            STREAMS + offs_m[:, None] * K + kk[None, :],
+            mask=mmask[:, None] & kmask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        dmf = tl.dot(dmix, w, input_precision="ieee")  # (d_mixes @ fn) [BM, BK]
+        dx = r[:, None] * dmf - x * coef[:, None]
+        tl.store(
+            DSTREAMS + offs_m[:, None] * K + kk[None, :],
+            dx,
+            mask=mmask[:, None] & kmask[None, :],
+        )
+        dfn = tl.dot(tl.trans(dG), x, input_precision="ieee")  # [NP, BK]
+        tl.atomic_add(
+            DFN + n[:, None] * K + kk[None, :],
+            dfn,
+            mask=nmask[:, None] & kmask[None, :],
+        )
+
+
+class _RMSNormLinear(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, streams_flat, fn, eps, BM):
+        M, K = streams_flat.shape
+        N = fn.shape[0]
+        NP = triton.next_power_of_2(N)
+        sf = streams_flat.contiguous()
+        fn = fn.contiguous().float()
+        mixes = torch.empty(M, N, device=sf.device, dtype=torch.float32)
+        R = torch.empty(M, device=sf.device, dtype=torch.float32)
+        G = torch.empty(M, N, device=sf.device, dtype=torch.float32)
+        grid = lambda m: (triton.cdiv(M, m["BM"]),)
+        _rmsln_fwd_kernel[grid](sf, fn, mixes, R, G, M, K, eps, N=N, NP=NP, BM=BM)
+        ctx.save_for_backward(sf, fn, R, G)
+        ctx.dims = (M, K, N, NP, BM)
+        return mixes
+
+    @staticmethod
+    def backward(ctx, dmix):
+        sf, fn, R, G = ctx.saved_tensors
+        M, K, N, NP, BM = ctx.dims
+        dstreams = torch.empty_like(sf, dtype=torch.float32)
+        dfn = torch.zeros_like(fn)
+        grid = lambda m: (triton.cdiv(M, m["BM"]),)
+        _rmsln_bwd_kernel[grid](
+            sf,
+            fn,
+            R,
+            G,
+            dmix.contiguous().float(),
+            dstreams,
+            dfn,
+            M,
+            K,
+            N=N,
+            NP=NP,
+            BM=BM,
+        )
+        return dstreams.to(sf.dtype), dfn, None, None
+
+
+def _rmsnorm_linear(streams_flat, fn, eps, BM=16):
+    return _RMSNormLinear.apply(streams_flat, fn, eps, BM)
+
+
+# fused collapse: collapsed[m,d] = sum_h pre[m,h] * streams[m,h,d]
+@triton.autotune(
+    configs=[
+        triton.Config({"BD": bd}, num_warps=w)
+        for bd in (256, 512, 1024)
+        for w in (4, 8)
+    ],
+    key=["M", "D"],
+)
+@triton.jit
+def _collapse_fwd_kernel(PRE, STREAMS, OUT, M, D, HC: tl.constexpr, BD: tl.constexpr):
+    m = tl.program_id(0)
+    dblk = tl.program_id(1) * BD + tl.arange(0, BD)
+    dmask = dblk < D
+    acc = tl.zeros([BD], tl.float32)
+    for h in range(HC):
+        p = tl.load(PRE + m * HC + h)
+        s = tl.load(STREAMS + m * (HC * D) + h * D + dblk, mask=dmask, other=0.0).to(
+            tl.float32
+        )
+        acc += p * s
+    tl.store(OUT + m * D + dblk, acc, mask=dmask)
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BD": bd}, num_warps=w)
+        for bd in (256, 512, 1024)
+        for w in (4, 8)
+    ],
+    key=["M", "D"],
+    reset_to_zero=["DPRE"],
+)
+@triton.jit
+def _collapse_bwd_kernel(
+    PRE, STREAMS, GOUT, DPRE, DSTREAMS, M, D, HC: tl.constexpr, BD: tl.constexpr
+):
+    m = tl.program_id(0)
+    dblk = tl.program_id(1) * BD + tl.arange(0, BD)
+    dmask = dblk < D
+    g = tl.load(GOUT + m * D + dblk, mask=dmask, other=0.0).to(tl.float32)
+    for h in range(HC):
+        p = tl.load(PRE + m * HC + h)
+        s = tl.load(STREAMS + m * (HC * D) + h * D + dblk, mask=dmask, other=0.0).to(
+            tl.float32
+        )
+        tl.store(DSTREAMS + m * (HC * D) + h * D + dblk, p * g, mask=dmask)
+        tl.atomic_add(DPRE + m * HC + h, tl.sum(tl.where(dmask, g * s, 0.0)))
+
+
+class _Collapse(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, pre, streams):
+        M, HC, D = streams.shape
+        pre = pre.contiguous()
+        streams = streams.contiguous()
+        out = torch.empty(M, D, device=streams.device, dtype=torch.float32)
+        grid = lambda m: (M, triton.cdiv(D, m["BD"]))
+        _collapse_fwd_kernel[grid](pre, streams, out, M, D, HC=HC)
+        ctx.save_for_backward(pre, streams)
+        ctx.dims = (M, HC, D)
+        return out
+
+    @staticmethod
+    def backward(ctx, gout):
+        pre, streams = ctx.saved_tensors
+        M, HC, D = ctx.dims
+        dpre = torch.zeros_like(pre)
+        dstreams = torch.empty_like(streams, dtype=torch.float32)
+        grid = lambda m: (M, triton.cdiv(D, m["BD"]))
+        _collapse_bwd_kernel[grid](
+            pre, streams, gout.contiguous().float(), dpre, dstreams, M, D, HC=HC
+        )
+        return dpre, dstreams.to(streams.dtype)
+
+
+def hyperconnection_forward(
+    hidden_streams, input_norm, fn, base, scale, hc, iters, eps, post_mult=2.0
+):
+    """Drop-in for DeepseekV4HyperConnection.forward. Returns (post, comb, collapsed).
+    Fully fused Triton: rmsnorm+fn-linear (one pass over streams), Sinkhorn mixer, and the
+    weighted collapse — eager does these in ~3 passes over the 16384-wide streams + a 536MB
+    fp32 intermediate."""
+    dtype = hidden_streams.dtype
+    # match the fused-linear / Sinkhorn params to the stream compute dtype (they may be fp32).
+    fn = fn.to(dtype) if fn.dtype != dtype else fn
+    scale = scale.to(dtype) if scale.dtype != dtype else scale
+    base = base.to(dtype) if base.dtype != dtype else base
+    *lead, _, hidden = hidden_streams.shape
+    streams_flat = hidden_streams.reshape(-1, hc * hidden)
+    mixes = _rmsnorm_linear(streams_flat, fn, input_norm.eps)
+    pre, post, comb = sinkhorn_mix(mixes, scale, base, hc, iters, eps, post_mult)
+    collapsed = (
+        _Collapse.apply(pre, hidden_streams.reshape(-1, hc, hidden))
+        .reshape(*lead, hidden)
+        .to(dtype)
+    )
+    post = post.reshape(*lead, hc)
+    comb = comb.reshape(*lead, hc, hc)
+    return post, comb, collapsed
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/patch.py b/src/axolotl/integrations/kernels/libs/dsv4/patch.py
new file mode 100644
index 0000000000..c310c6a044
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/patch.py
@@ -0,0 +1,348 @@
+"""Monkeypatch DeepSeek-V4 to use the fused Triton training kernels.
+
+V4 ships eager-only attention (head_dim 512 > FlashAttention's cap, plus a per-head sink
+and an in-block compressor KV-concat). This registers a custom attention backend that
+routes sliding-attention layers to ``sliding_attn`` and CSA/HCA layers to
+``csa_attn`` (splitting the eager-combined ``[sliding | compressed]`` KV and
+``[sliding_mask | block_bias]`` mask back apart), and swaps the interleaved partial-RoPE
+and the mHC HyperConnection for their fused kernels. o_a_proj stays on cuBLAS (a fused
+grouped GEMM doesn't beat it).
+
+Kernels assume ``attention_dropout == 0`` (the V4 default).
+"""
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+from .attention import sliding_attn
+from .attention_csa import csa_attn
+from .gated_pool import gated_softmax_pool
+from .indexer import indexer_scores
+from .mhc import hyperconnection_forward
+from .rope import apply_rotary_pos_emb_triton
+
+LOG = get_logger(__name__)
+
+_MOD = None  # the transformers deepseek_v4 module; set by patch_deepseek_v4_kernels
+
+
+def _dsv4_attention(
+    module,
+    query,
+    key,
+    value,
+    attention_mask,
+    scaling=None,
+    dropout=0.0,
+    sliding_window=None,
+    s_aux=None,
+    **kwargs,
+):
+    """Drop-in for ``eager_attention_forward``. ``query`` [B,H,S,D]; ``key``/``value``
+    [B,1,KV,D] (shared MQA); ``s_aux`` = per-head sinks; ``attention_mask`` is the
+    eager-built [sliding | block_bias] mask. Returns (attn_output [B,S,H,D], None).
+
+    Patched onto the module's ``eager_attention_forward`` (V4's default backend) so the
+    mask — and thus the compressor's ``block_bias`` — is still built and handed to us."""
+    B, H, S, D = query.shape
+    KV = key.shape[2]
+    if scaling is None:
+        scaling = D**-0.5
+    window = sliding_window if sliding_window is not None else module.sliding_window
+
+    if KV == S:  # sliding-attention layer (no compressor entries)
+        out = sliding_attn(query, key, value, s_aux, scaling, window)
+    else:  # CSA / HCA: keys are [sliding (S) | compressed (KV-S)], mask carries block_bias
+        kv_slide = key[:, :, :S]
+        kv_comp = key[:, :, S:]
+        block_bias = attention_mask[..., S:]
+        out = csa_attn(query, kv_slide, kv_comp, block_bias, s_aux, scaling, window)
+    return out.transpose(1, 2).contiguous(), None
+
+
+def _gated_pool_norm(new_kv, new_gate, kv_norm):
+    """Fused replacement for ``kv_norm((new_kv * new_gate.softmax(2,fp32)).sum(2))`` —
+    the compressor's per-window gated-softmax pool. Returns the compressor dtype (the
+    pool is fp32 internally; we cast back so the downstream RoPE/concat stay bf16)."""
+    return kv_norm(gated_softmax_pool(new_kv, new_gate).to(new_kv.dtype))
+
+
+def _indexer_scorer_forward(self, q, compressed_kv, hidden_states):
+    """Fused DeepseekV4IndexerScorer: never materializes the [B,S,H,T] fp32 score
+    tensor (H=64 index heads). Output feeds only topk().indices, so no grad needed."""
+    weights = self.weights_proj(hidden_states).float() * self.weights_scaling  # [B,S,H]
+    return indexer_scores(q, compressed_kv, weights, self.softmax_scale)
+
+
+def _hca_compressor_forward(
+    self, hidden_states, q_residual, position_ids, past_key_values, layer_idx
+):
+    mod = _MOD
+    batch, _, _ = hidden_states.shape
+    cache_layer = (
+        past_key_values.layers[layer_idx] if past_key_values is not None else None
+    )
+    kv = self.kv_proj(hidden_states)
+    gate = self.gate_proj(hidden_states)
+    if cache_layer is None:
+        usable = (kv.shape[1] // self.compress_rate) * self.compress_rate
+        chunk_kv, chunk_gate, first_window_position = (
+            kv[:, :usable],
+            gate[:, :usable],
+            0,
+        )
+    else:
+        chunk_kv, chunk_gate, first_window_position = (
+            cache_layer.store_compression_weights("compressor", kv, gate)
+        )
+
+    if chunk_kv.shape[1] > 0:
+        n_windows = chunk_kv.shape[1] // self.compress_rate
+        chunk_kv = chunk_kv.view(batch, n_windows, self.compress_rate, -1)
+        chunk_gate = (
+            chunk_gate.view(batch, n_windows, self.compress_rate, -1)
+            + self.position_bias
+        )
+        compressed = _gated_pool_norm(chunk_kv, chunk_gate, self.kv_norm)
+        positions = torch.arange(n_windows, device=compressed.device)
+        positions = (
+            (positions * self.compress_rate + first_window_position)
+            .unsqueeze(0)
+            .expand(batch, -1)
+        )
+        cos, sin = self.rotary_emb(
+            compressed, position_ids=positions, layer_type=self.rope_layer_type
+        )
+        compressed = mod.apply_rotary_pos_emb(
+            compressed.unsqueeze(1), cos, sin
+        ).squeeze(1)
+    else:
+        compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))
+
+    if cache_layer is not None:
+        compressed = cache_layer.update_compressor_states("compressor", compressed)
+    compressed_kv = compressed.unsqueeze(1)
+
+    compressed_len = compressed_kv.shape[2]
+    seq_len = position_ids.shape[1]
+    if seq_len == 1 or compressed_len == 0:
+        return compressed_kv, None
+
+    entry_indices = torch.arange(compressed_len, device=compressed_kv.device)
+    causal_threshold = (position_ids + 1) // self.compress_rate
+    block_bias = compressed_kv.new_zeros((batch, 1, seq_len, compressed_len))
+    block_bias = block_bias.masked_fill(
+        entry_indices.view(1, 1, 1, -1) >= causal_threshold.unsqueeze(1).unsqueeze(-1),
+        float("-inf"),
+    )
+    return compressed_kv, block_bias
+
+
+def _indexer_forward(
+    self, hidden_states, q_residual, position_ids, past_key_values, layer_idx
+):
+    mod = _MOD
+    batch, seq_len, _ = hidden_states.shape
+    cache_layer = (
+        past_key_values.layers[layer_idx] if past_key_values is not None else None
+    )
+    kv = self.kv_proj(hidden_states)
+    gate = self.gate_proj(hidden_states)
+
+    if cache_layer is None:
+        usable = (kv.shape[1] // self.compress_rate) * self.compress_rate
+        chunk_kv, chunk_gate, first_window_position = (
+            kv[:, :usable],
+            gate[:, :usable],
+            0,
+        )
+    else:
+        chunk_kv, chunk_gate, first_window_position = (
+            cache_layer.store_compression_weights("indexer", kv, gate)
+        )
+
+    if chunk_kv.shape[1] > 0:
+        n_windows = chunk_kv.shape[1] // self.compress_rate
+        ratio = self.compress_rate
+        chunk_kv = chunk_kv.view(batch, n_windows, ratio, -1)
+        chunk_gate = chunk_gate.view(batch, n_windows, ratio, -1) + self.position_bias
+
+        new_kv = chunk_kv.new_zeros((batch, n_windows, 2 * ratio, self.head_dim))
+        new_gate = chunk_gate.new_full(
+            (batch, n_windows, 2 * ratio, self.head_dim), float("-inf")
+        )
+        new_kv[:, :, ratio:] = chunk_kv[..., self.head_dim :]
+        new_gate[:, :, ratio:] = chunk_gate[..., self.head_dim :]
+        if n_windows > 1:
+            new_kv[:, 1:, :ratio] = chunk_kv[:, :-1, :, : self.head_dim]
+            new_gate[:, 1:, :ratio] = chunk_gate[:, :-1, :, : self.head_dim]
+        if cache_layer is not None:
+            prior_kv, prior_gate = cache_layer.update_overlap_state(
+                "indexer", chunk_kv, chunk_gate, self.head_dim
+            )
+            if prior_kv is not None:
+                new_kv[:, 0, :ratio] = prior_kv.to(new_kv.dtype)
+                new_gate[:, 0, :ratio] = prior_gate.to(new_gate.dtype)
+
+        compressed = _gated_pool_norm(new_kv, new_gate, self.kv_norm)
+        positions = torch.arange(n_windows, device=compressed.device)
+        positions = positions * self.compress_rate + first_window_position
+        positions = positions.unsqueeze(0).expand(batch, -1)
+        cos, sin = self.rotary_emb(
+            compressed, position_ids=positions, layer_type=self.rope_layer_type
+        )
+        compressed = mod.apply_rotary_pos_emb(
+            compressed.unsqueeze(1), cos, sin
+        ).squeeze(1)
+    else:
+        compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))
+
+    compressed_kv = (
+        compressed
+        if cache_layer is None
+        else cache_layer.update_compressor_states("indexer", compressed)
+    )
+
+    cos_q, sin_q = self.rotary_emb(
+        hidden_states, position_ids=position_ids, layer_type=self.rope_layer_type
+    )
+    q = (
+        self.q_b_proj(q_residual)
+        .view(batch, seq_len, -1, self.head_dim)
+        .transpose(1, 2)
+    )
+    q = mod.apply_rotary_pos_emb(q, cos_q, sin_q).transpose(1, 2)
+
+    index_scores = self.scorer(q, compressed_kv, hidden_states)
+    compressed_len = compressed_kv.shape[1]
+    top_k = min(self.index_topk, compressed_len)
+
+    if compressed_len > 0:
+        causal_threshold = (position_ids + 1) // self.compress_rate
+        entry_indices = torch.arange(compressed_len, device=index_scores.device)
+        future_mask = entry_indices.view(1, 1, -1) >= causal_threshold.unsqueeze(-1)
+        index_scores = index_scores.masked_fill(future_mask, float("-inf"))
+        top_k_indices = index_scores.topk(top_k, dim=-1).indices
+        invalid = top_k_indices >= causal_threshold.unsqueeze(-1)
+        return torch.where(invalid, torch.full_like(top_k_indices, -1), top_k_indices)
+
+    return index_scores.topk(top_k, dim=-1).indices
+
+
+def _csa_compressor_forward(
+    self, hidden_states, q_residual, position_ids, past_key_values, layer_idx
+):
+    mod = _MOD
+    batch, seq_len, _ = hidden_states.shape
+    cache_layer = (
+        past_key_values.layers[layer_idx] if past_key_values is not None else None
+    )
+    kv = self.kv_proj(hidden_states)
+    gate = self.gate_proj(hidden_states)
+
+    if cache_layer is None:
+        usable = (kv.shape[1] // self.compress_rate) * self.compress_rate
+        chunk_kv, chunk_gate, first_window_position = (
+            kv[:, :usable],
+            gate[:, :usable],
+            0,
+        )
+    else:
+        chunk_kv, chunk_gate, first_window_position = (
+            cache_layer.store_compression_weights("compressor", kv, gate)
+        )
+
+    if chunk_kv.shape[1] > 0:
+        n_windows = chunk_kv.shape[1] // self.compress_rate
+        ratio = self.compress_rate
+        chunk_kv = chunk_kv.view(batch, n_windows, ratio, -1)
+        chunk_gate = chunk_gate.view(batch, n_windows, ratio, -1) + self.position_bias
+
+        new_kv = chunk_kv.new_zeros((batch, n_windows, 2 * ratio, self.head_dim))
+        new_gate = chunk_gate.new_full(
+            (batch, n_windows, 2 * ratio, self.head_dim), float("-inf")
+        )
+        new_kv[:, :, ratio:] = chunk_kv[..., self.head_dim :]
+        new_gate[:, :, ratio:] = chunk_gate[..., self.head_dim :]
+        if n_windows > 1:
+            new_kv[:, 1:, :ratio] = chunk_kv[:, :-1, :, : self.head_dim]
+            new_gate[:, 1:, :ratio] = chunk_gate[:, :-1, :, : self.head_dim]
+        if cache_layer is not None:
+            prior_kv, prior_gate = cache_layer.update_overlap_state(
+                "compressor", chunk_kv, chunk_gate, self.head_dim
+            )
+            if prior_kv is not None:
+                new_kv[:, 0, :ratio] = prior_kv.to(new_kv.dtype)
+                new_gate[:, 0, :ratio] = prior_gate.to(new_gate.dtype)
+
+        compressed = _gated_pool_norm(new_kv, new_gate, self.kv_norm)
+        positions = torch.arange(n_windows, device=compressed.device)
+        positions = positions * self.compress_rate + first_window_position
+        positions = positions.unsqueeze(0).expand(batch, -1)
+        cos, sin = self.rotary_emb(
+            compressed, position_ids=positions, layer_type=self.rope_layer_type
+        )
+        compressed = mod.apply_rotary_pos_emb(
+            compressed.unsqueeze(1), cos, sin
+        ).squeeze(1)
+    else:
+        compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))
+
+    if cache_layer is not None:
+        compressed = cache_layer.update_compressor_states("compressor", compressed)
+    compressed_kv = compressed.unsqueeze(1)
+
+    top_k_indices = self.indexer(
+        hidden_states, q_residual, position_ids, past_key_values, layer_idx
+    )
+    compressed_len = compressed_kv.shape[2]
+    valid = top_k_indices >= 0
+    safe_indices = torch.where(
+        valid, top_k_indices, torch.full_like(top_k_indices, compressed_len)
+    )
+    block_bias = compressed_kv.new_full(
+        (batch, 1, seq_len, compressed_len + 1), float("-inf")
+    )
+    block_bias.scatter_(-1, safe_indices.unsqueeze(1), 0.0)
+    return compressed_kv, block_bias[..., :compressed_len]
+
+
+def _hyperconnection_forward(self, hidden_streams):
+    return hyperconnection_forward(
+        hidden_streams,
+        self.input_norm,
+        self.fn,
+        self.base,
+        self.scale,
+        self.hc_mult,
+        self.hc_sinkhorn_iters,
+        self.hc_eps,
+        post_mult=2.0,
+    )
+
+
+def patch_deepseek_v4_kernels():
+    """Patch the eager attention backend + RoPE + HyperConnection. Idempotent.
+
+    V4 stays on ``_attn_implementation='eager'`` (head_dim 512 bans the real flash
+    backends), so transformers still builds the mask and cats the compressor's
+    ``block_bias`` — we just replace the eager kernel itself."""
+    from transformers.models.deepseek_v4 import modeling_deepseek_v4 as mod
+
+    global _MOD
+    _MOD = mod
+
+    def _rope(x, cos, sin, unsqueeze_dim=1):
+        return apply_rotary_pos_emb_triton(x, cos, sin, unsqueeze_dim)
+
+    mod.eager_attention_forward = _dsv4_attention
+    mod.apply_rotary_pos_emb = _rope
+    mod.DeepseekV4HyperConnection.forward = _hyperconnection_forward
+    mod.DeepseekV4IndexerScorer.forward = _indexer_scorer_forward
+    mod.DeepseekV4HCACompressor.forward = _hca_compressor_forward
+    mod.DeepseekV4Indexer.forward = _indexer_forward
+    mod.DeepseekV4CSACompressor.forward = _csa_compressor_forward
+    LOG.info(
+        "Patched DeepSeek-V4 with fused Triton kernels (attention/rope/mHC/compressor/indexer)"
+    )
diff --git a/src/axolotl/integrations/kernels/libs/dsv4/rope.py b/src/axolotl/integrations/kernels/libs/dsv4/rope.py
new file mode 100644
index 0000000000..08a24eb1b7
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/dsv4/rope.py
@@ -0,0 +1,121 @@
+"""Fused interleaved partial-RoPE kernel for DeepSeek-V4 (fwd + bwd).
+
+Reference: ``transformers.models.deepseek_v4.modeling_deepseek_v4.apply_rotary_pos_emb``.
+
+V4 uses *interleaved* RoPE (GPT-J style: consecutive channel pairs) on the trailing
+``rope_dim`` channels of each head, leaving the leading ``nope`` channels untouched.
+``cos``/``sin`` arrive half-width (one entry per pair, shape ``[B, S, rope_dim//2]``)
+and are expanded by ``repeat_interleave(2)`` in the reference. Per pair ``i``:
+
+    out[2i]   = x[2i]*cos[i] - x[2i+1]*sin[i]
+    out[2i+1] = x[2i+1]*cos[i] + x[2i]*sin[i]
+
+The rotation is orthogonal, so the backward is the same map with ``sin -> -sin``.
+The model also calls the reference with a pre-negated ``sin`` (output inverse-rope);
+that composes naturally — this Function just rotates by whatever ``(cos, sin)`` it gets
+and negates ``sin`` for its own backward.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _rope_kernel(
+    x_ptr,
+    cos_ptr,
+    sin_ptr,
+    out_ptr,
+    S,
+    H,
+    cs_b,  # cos/sin strides: may be non-contiguous (the rotary's transpose view)
+    cs_s,
+    cs_i,
+    D: tl.constexpr,
+    ROPE: tl.constexpr,
+    Rh: tl.constexpr,
+    BLOCK_NOPE: tl.constexpr,
+):
+    row = tl.program_id(0)  # one program per (b, h, s) row of length D
+    s = row % S
+    b = (row // S) // H
+    nope = D - ROPE
+    cs_off = b * cs_b + s * cs_s
+
+    x_row = x_ptr + row * D
+    out_row = out_ptr + row * D
+
+    # passthrough leading nope channels
+    off = tl.arange(0, BLOCK_NOPE)
+    m = off < nope
+    tl.store(out_row + off, tl.load(x_row + off, mask=m, other=0.0), mask=m)
+
+    i = tl.arange(0, Rh)
+    even = nope + 2 * i
+    odd = even + 1
+    e = tl.load(x_row + even).to(tl.float32)
+    o = tl.load(x_row + odd).to(tl.float32)
+    c = tl.load(cos_ptr + cs_off + i * cs_i).to(tl.float32)
+    sn = tl.load(sin_ptr + cs_off + i * cs_i).to(tl.float32)
+    tl.store(out_row + even, e * c - o * sn)
+    tl.store(out_row + odd, o * c + e * sn)
+
+
+def _launch(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """x: [B, H, S, D] contiguous; cos/sin: [B, S, rope_dim//2] (any strides). Returns
+    [B, H, S, D]. cos/sin come from the rotary as a ``transpose(1,2)`` view, so they are
+    typically *not* contiguous — we pass their strides instead of assuming a layout."""
+    B, H, S, D = x.shape
+    Rh = cos.shape[-1]
+    ROPE = 2 * Rh
+    if sin.stride() != cos.stride():
+        sin = sin.contiguous()
+        cos = cos.contiguous()
+    out = torch.empty_like(x)
+    M = B * H * S
+    _rope_kernel[(M,)](
+        x,
+        cos,
+        sin,
+        out,
+        S,
+        H,
+        cos.stride(0),
+        cos.stride(1),
+        cos.stride(2),
+        D=D,
+        ROPE=ROPE,
+        Rh=Rh,
+        BLOCK_NOPE=triton.next_power_of_2(D - ROPE),
+        num_warps=4,
+    )
+    return out
+
+
+class _RoPE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, cos, sin):
+        x = x.contiguous()
+        ctx.save_for_backward(cos, sin)
+        return _launch(x, cos, sin)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        cos, sin = ctx.saved_tensors
+        grad_x = _launch(grad_out.contiguous(), cos, -sin)
+        return grad_x, None, None
+
+
+def apply_rotary_pos_emb_triton(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1
+) -> torch.Tensor:
+    """Drop-in for ``apply_rotary_pos_emb``. ``x``: [B, H, S, D]; ``cos``/``sin``:
+    [B, S, rope_dim//2]. ``unsqueeze_dim`` is accepted for signature parity (the
+    reference broadcasts cos/sin over the head axis; this kernel indexes it directly)."""
+    # rotary cos/sin are typically fp32; match them to x's compute dtype.
+    if cos.dtype != x.dtype:
+        cos = cos.to(x.dtype)
+    if sin.dtype != x.dtype:
+        sin = sin.to(x.dtype)
+    return _RoPE.apply(x, cos, sin)
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/__init__.py b/src/axolotl/integrations/kernels/libs/glm_dsa/__init__.py
new file mode 100644
index 0000000000..e3e0267fe1
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/__init__.py
@@ -0,0 +1,19 @@
+"""GLM-5.2 (glm_moe_dsa) DSA fused training kernels.
+
+DeepSeek-V3.2 sparse-MLA (DSA / Lightning-Indexer) attention for GLM-5.2: MLA weight absorption +
+head-batched sparse-gather flash attention (fwd+bwd), a fused Lightning-Indexer scorer + top-k,
+length-aware dense/sparse dispatch, and a context-parallel path that all-gathers only the 576-wide
+compressed KV. See ``patch.patch_glm_moe_dsa_attention``.
+"""
+
+from .patch import (
+    fused_indexer_topk,
+    keep_router_fp32,
+    patch_glm_moe_dsa_attention,
+)
+
+__all__ = [
+    "patch_glm_moe_dsa_attention",
+    "keep_router_fp32",
+    "fused_indexer_topk",
+]
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/_autotune.py b/src/axolotl/integrations/kernels/libs/glm_dsa/_autotune.py
new file mode 100644
index 0000000000..e7ac411506
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/_autotune.py
@@ -0,0 +1,42 @@
+"""Autotune helpers: portable shared-memory pruning.
+
+Triton's autotune raises ``OutOfResources`` on a config whose tiles exceed the device's shared
+memory (it does NOT silently skip it), so a config grid with big tiles crashes on a small-SMEM
+GPU. ``smem_prune`` drops the over-budget configs up front by estimating each config's SMEM and
+comparing to the *device's* real ``max_shared_mem`` (queried + memoized, not hardcoded) — so the
+same grid runs on sm120 (~99 KB) and uses bigger tiles on Hopper/Blackwell-datacenter (~228 KB).
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+
+_SMEM: dict[int, int] = {}
+
+
+def max_smem(dev: int | None = None) -> int:
+    dev = torch.cuda.current_device() if dev is None else dev
+    if dev not in _SMEM:
+        props = triton.runtime.driver.active.utils.get_device_properties(dev)
+        _SMEM[dev] = int(props["max_shared_mem"])
+    return _SMEM[dev]
+
+
+def smem_prune(est_bytes, headroom: float = 0.95):
+    """Build an ``early_config_prune`` callback. ``est_bytes(config_kwargs, num_stages, **constexprs)``
+    returns a config's estimated SMEM in bytes; configs over ``headroom * device_limit`` are dropped
+    (keeping the single smallest if all exceed, so autotune never gets an empty set)."""
+
+    def prune(configs, named_args, **kwargs):
+        limit = int(max_smem() * headroom)
+        keep = [
+            c for c in configs if est_bytes(c.kwargs, c.num_stages, **kwargs) <= limit
+        ]
+        if keep:
+            return keep
+        return sorted(
+            configs, key=lambda c: est_bytes(c.kwargs, c.num_stages, **kwargs)
+        )[:1]
+
+    return prune
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/attention_mla_absorb.py b/src/axolotl/integrations/kernels/libs/glm_dsa/attention_mla_absorb.py
new file mode 100644
index 0000000000..f4d76f4533
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/attention_mla_absorb.py
@@ -0,0 +1,565 @@
+"""MLA weight-absorption + head-batched sparse gather attention for GLM-5.2 DSA.
+
+In MLA the per-head key/value are produced from a SHARED compressed latent ``c_kv`` [S, kv_lora=512]:
+    k_nope[t,h] = c_kv[t] @ W_kb_k[h]        value[t,h] = c_kv[t] @ W_kb_v[h]
+so the nope score absorbs kv_b into the query:
+    q_nope[h]·k_nope[t,h] = (q_nope[h] @ W_kb_kᵀ[h]) · c_kv[t]
+With the rope part (k_rot shared across heads) appended, every head attends the SAME key tensor
+    K_shared[t] = cat(c_kv[t] [512], k_rot[t] [64])   (576-wide, shared across heads)
+    Q_abs[s,h]  = cat(q_nope[s,h] @ W_kb_kᵀ[h] [512], q_rot[s,h] [64])
+    score[s,h,t] = Q_abs[s,h] · K_shared[t] * scale
+    out_latent[s,h] = Σ_t softmax(score)[h,t] · c_kv[t]        (512-wide)
+    out[s,h] = out_latent[s,h] @ W_kb_v[h]                     (-> v_head_dim 256, then o_proj)
+
+Now all heads share K_shared -> a head-batched MMA (M=H), and the sparse gather loads the 576-wide
+shared KV ONCE per query position instead of per head (the win that the per-(b,h,s) GEMV lacked).
+The absorption/value GEMMs (q_nope@W_kb_kᵀ, out_latent@W_kb_v) stay in torch and are differentiable,
+so gradients flow to kv_b_proj in full-parameter training and to its LoRA adapter in LoRA training;
+the kernel differentiates Q_abs / K_shared.
+
+This module currently provides the absorption helpers + the head-batched gather FORWARD kernel,
+validated against the eager dense-mask MLA reference. Backward is added next.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+from ._autotune import smem_prune
+from .config import KV_LORA_RANK, QK_NOPE_HEAD_DIM, QK_ROPE_HEAD_DIM, V_HEAD_DIM
+
+# Measured: actual SMEM ≈ the single-count tile estimate + ~2KB (tl.trans is handled by the MMA,
+# no extra buffer). Use the raw estimate + a small fixed overhead; headroom 0.97 leaves margin
+# without pruning borderline-fitting configs (e.g. BH=16,BN=32 which is the fast one on sm120).
+_SMEM_OVERHEAD = 3072
+
+
+def _fwd_smem(cfg, num_stages, **kw):
+    # bf16: qn+qr [BH,w] loaded once; c_kv+k_rot [BN,w] pipelined x num_stages.
+    w = kw["DL"] + kw["DR"]
+    return (cfg["BH"] * w + num_stages * cfg["BN"] * w) * 2 + _SMEM_OVERHEAD
+
+
+def _bwd_smem(cfg, num_stages, **kw):
+    # bwd also holds out_l/dout_l [BH,DL] (loaded once).
+    DL, DR = kw["DL"], kw["DR"]
+    w = DL + DR
+    return (cfg["BH"] * (w + 2 * DL) + num_stages * cfg["BN"] * w) * 2 + _SMEM_OVERHEAD
+
+
+_ABSORB_CONFIGS = [
+    triton.Config({"BH": bh, "BN": bn}, num_warps=warps, num_stages=st)
+    for bh in (16, 32)
+    for bn in (16, 32, 64, 128)
+    for warps in (4, 8)
+    for st in (1, 2, 3)
+]
+
+DQK = KV_LORA_RANK + QK_ROPE_HEAD_DIM  # 576 (compressed score dim)
+
+
+def split_kv_b(kv_b_weight: torch.Tensor, n_heads: int):
+    """Split kv_b_proj.weight [H*(qk_nope+v_head), kv_lora] into per-head
+    W_kb_k [H, qk_nope, kv_lora] and W_kb_v [H, v_head, kv_lora]."""
+    w = kv_b_weight.view(n_heads, QK_NOPE_HEAD_DIM + V_HEAD_DIM, KV_LORA_RANK)
+    return w[:, :QK_NOPE_HEAD_DIM, :], w[:, QK_NOPE_HEAD_DIM:, :]
+
+
+def absorb_query(q_nope, q_rot, w_kb_k):
+    """q_nope [B,H,S,qk_nope], q_rot [B,H,S,qk_rope], w_kb_k [H,qk_nope,kv_lora] ->
+    Q_abs [B,H,S,DQK] = cat(q_nope @ W_kb_kᵀ, q_rot). Differentiable. ``w_kb_k`` is cast to the
+    activation dtype (under LoRA the effective kv_b weight promotes to fp32 via the adapters)."""
+    q_abs_nope = torch.einsum("bhsn,hnl->bhsl", q_nope, w_kb_k.to(q_nope.dtype))
+    return torch.cat([q_abs_nope, q_rot], dim=-1)
+
+
+def project_value(out_latent, w_kb_v):
+    """out_latent [B,H,S,kv_lora] -> out [B,H,S,v_head] = out_latent @ W_kb_vᵀ. Differentiable."""
+    return torch.einsum("bhsl,hvl->bhsv", out_latent, w_kb_v.to(out_latent.dtype))
+
+
+@triton.autotune(
+    configs=_ABSORB_CONFIGS,
+    key=["TOPK"],
+    prune_configs_by={"early_config_prune": smem_prune(_fwd_smem)},
+)
+@triton.jit
+def _absorb_fwd_kernel(
+    QABS,
+    KSH,
+    IDX,
+    OUT,
+    LSE,
+    scale,
+    S,
+    TOPK,
+    H,
+    q_offset,
+    sqa_b,
+    sqa_h,
+    sqa_s,
+    sqa_d,
+    sks_b,
+    sks_s,
+    sks_d,
+    si_b,
+    si_s,
+    si_t,
+    so_b,
+    so_h,
+    so_s,
+    so_d,
+    sl_b,
+    sl_h,
+    sl_s,
+    SEQQ,
+    SEQK,
+    ssq_b,
+    ssq_s,
+    ssk_b,
+    ssk_s,
+    DL: tl.constexpr,  # kv_lora_rank (=value latent dim), power of 2
+    DR: tl.constexpr,  # qk_rope_head_dim, power of 2
+    BH: tl.constexpr,
+    BN: tl.constexpr,
+    HAS_DOC: tl.constexpr,  # sample-packing: forbid cross-document keys
+):
+    # S on axis 0 (X, up to 2^31-1); CUDA caps grid Y/Z at 65535 so S must not be on Y/Z.
+    s = tl.program_id(0)
+    b = tl.program_id(1)
+    h0 = tl.program_id(2) * BH
+    offs_h = h0 + tl.arange(0, BH)
+    hmask = offs_h < H
+    offs_l = tl.arange(0, DL)  # compressed-latent dims [0, kv_lora)
+    offs_r = tl.arange(0, DR)  # rope dims, stored at [kv_lora, kv_lora+rope)
+
+    # split q_abs into its nope-latent (DL) and rope (DR) parts (avoids padding 576 to a pow2)
+    # keep tiles in their native (bf16) dtype for the tensor-core dot (fp32 accumulation)
+    base_q = QABS + b * sqa_b + offs_h[:, None] * sqa_h + s * sqa_s
+    qn = tl.load(base_q + offs_l[None, :] * sqa_d, mask=hmask[:, None], other=0.0)
+    qr = tl.load(
+        base_q + (DL + offs_r)[None, :] * sqa_d, mask=hmask[:, None], other=0.0
+    )
+
+    m_i = tl.full((BH,), -float("inf"), tl.float32)
+    l_i = tl.zeros((BH,), tl.float32)
+    acc = tl.zeros((BH, DL), tl.float32)
+
+    s_global = (
+        s + q_offset
+    )  # CP: query's global position for the causal check (q_offset=0 single-GPU)
+    seqq_s = tl.load(SEQQ + b * ssq_b + s * ssq_s) if HAS_DOC else 0
+    for t0 in range(0, TOPK, BN):
+        offs_t = t0 + tl.arange(0, BN)
+        tmask = offs_t < TOPK
+        idx = tl.load(
+            IDX + b * si_b + s * si_s + offs_t * si_t, mask=tmask, other=0
+        ).to(tl.int64)
+        valid = tmask & (idx <= s_global)
+        if HAS_DOC:
+            seqk = tl.load(SEQK + b * ssk_b + idx * ssk_s, mask=tmask, other=-1)
+            valid = valid & (seqk == seqq_s)
+
+        # gather the shared KV once: c_kv [BN, DL] (also the value latent) + k_rot [BN, DR]
+        base_k = KSH + b * sks_b + idx[:, None] * sks_s
+        c_kv = tl.load(base_k + offs_l[None, :] * sks_d, mask=valid[:, None], other=0.0)
+        k_rot = tl.load(
+            base_k + (DL + offs_r)[None, :] * sks_d, mask=valid[:, None], other=0.0
+        )
+
+        scores = tl.dot(qn, tl.trans(c_kv))  # bf16 tensor core, fp32 accum
+        scores += tl.dot(qr, tl.trans(k_rot))
+        scores = tl.where(valid[None, :], scores * scale, -float("inf"))  # [BH, BN]
+
+        m_new = tl.maximum(m_i, tl.max(scores, axis=1))
+        # A query whose leading topk block is entirely causally/doc-masked has m_new == -inf, so
+        # exp(m_i - m_new) = exp(-inf + inf) = NaN. Subtract a finite max (0 there; all scores are
+        # -inf so every p is exp(-inf)=0 regardless) — matches the dense path on masked rows.
+        m_safe = tl.where(m_new == -float("inf"), 0.0, m_new)
+        p = tl.exp(scores - m_safe[:, None])  # [BH, BN]
+        alpha = tl.exp(m_i - m_safe)
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None] + tl.dot(p.to(c_kv.dtype), c_kv)
+        m_i = m_new
+
+    acc = (
+        acc / tl.where(l_i == 0.0, 1.0, l_i)[:, None]
+    )  # fully-masked query -> 0 output (loss masked)
+    o_ptr = OUT + b * so_b + offs_h[:, None] * so_h + s * so_s + offs_l[None, :] * so_d
+    tl.store(o_ptr, acc.to(OUT.dtype.element_ty), mask=hmask[:, None])
+    tl.store(LSE + b * sl_b + offs_h * sl_h + s * sl_s, m_i + tl.log(l_i), mask=hmask)
+
+
+def _bwd_prune_sm90_safe(configs, named_args, **kwargs):
+    """sm90 (Hopper) trips a nondeterministic invalid-PC (CUDA 718) while autotune trials the full
+    bwd grid (each config is individually clean — it's the trialing that corrupts the context). Force
+    a single sanitizer-clean config there; elsewhere fall back to SMEM pruning."""
+    import torch
+
+    base = smem_prune(_bwd_smem)
+    kept = base(configs, named_args, **kwargs)
+    if torch.cuda.get_device_capability() == (9, 0):
+        safe = [
+            c
+            for c in kept
+            if c.kwargs.get("BH") == 32
+            and c.kwargs.get("BN") == 64
+            and c.num_warps == 8
+            and c.num_stages == 3
+        ]
+        if safe:
+            return safe[:1]
+        return sorted(kept, key=lambda c: (c.kwargs.get("BN", 0), c.num_stages))[:1]
+    return kept
+
+
+@triton.autotune(
+    configs=_ABSORB_CONFIGS,
+    key=["TOPK"],
+    reset_to_zero=[
+        "DKSH"
+    ],  # dk_shared is atomic-accumulated across positions + head-tiles
+    prune_configs_by={"early_config_prune": _bwd_prune_sm90_safe},
+)
+@triton.jit
+def _absorb_bwd_kernel(
+    QABS,
+    KSH,
+    IDX,
+    OUTL,
+    DOUTL,
+    LSE,
+    DQABS,
+    DKSH,
+    scale,
+    S,
+    TOPK,
+    H,
+    q_offset,
+    sqa_b,
+    sqa_h,
+    sqa_s,
+    sqa_d,
+    sks_b,
+    sks_s,
+    sks_d,
+    si_b,
+    si_s,
+    si_t,
+    so_b,
+    so_h,
+    so_s,
+    so_d,
+    sl_b,
+    sl_h,
+    sl_s,
+    SEQQ,
+    SEQK,
+    ssq_b,
+    ssq_s,
+    ssk_b,
+    ssk_s,
+    DL: tl.constexpr,
+    DR: tl.constexpr,
+    BH: tl.constexpr,
+    BN: tl.constexpr,
+    HAS_DOC: tl.constexpr,
+):
+    # S on axis 0 (X, up to 2^31-1); CUDA caps grid Y/Z at 65535 so S must not be on Y/Z.
+    s = tl.program_id(0)
+    b = tl.program_id(1)
+    h0 = tl.program_id(2) * BH
+    offs_h = h0 + tl.arange(0, BH)
+    hmask = offs_h < H
+    offs_l = tl.arange(0, DL)
+    offs_r = tl.arange(0, DR)
+
+    base_q = QABS + b * sqa_b + offs_h[:, None] * sqa_h + s * sqa_s
+    qn = tl.load(base_q + offs_l[None, :] * sqa_d, mask=hmask[:, None], other=0.0)
+    qr = tl.load(
+        base_q + (DL + offs_r)[None, :] * sqa_d, mask=hmask[:, None], other=0.0
+    )
+    base_o = (
+        OUTL + b * so_b + offs_h[:, None] * so_h + s * so_s + offs_l[None, :] * so_d
+    )
+    out_l = tl.load(base_o, mask=hmask[:, None], other=0.0).to(tl.float32)  # [BH, DL]
+    dout_l = tl.load(
+        DOUTL + b * so_b + offs_h[:, None] * so_h + s * so_s + offs_l[None, :] * so_d,
+        mask=hmask[:, None],
+        other=0.0,
+    )  # [BH, DL]
+    lse = tl.load(LSE + b * sl_b + offs_h * sl_h + s * sl_s, mask=hmask, other=0.0)
+    delta = tl.sum(dout_l.to(tl.float32) * out_l, axis=1)  # [BH]
+
+    s_global = s + q_offset
+    seqq_s = tl.load(SEQQ + b * ssq_b + s * ssq_s) if HAS_DOC else 0
+    dqn = tl.zeros((BH, DL), tl.float32)
+    dqr = tl.zeros((BH, DR), tl.float32)
+    for t0 in range(0, TOPK, BN):
+        offs_t = t0 + tl.arange(0, BN)
+        tmask = offs_t < TOPK
+        idx = tl.load(
+            IDX + b * si_b + s * si_s + offs_t * si_t, mask=tmask, other=0
+        ).to(tl.int64)
+        valid = tmask & (idx <= s_global)
+        if HAS_DOC:
+            seqk = tl.load(SEQK + b * ssk_b + idx * ssk_s, mask=tmask, other=-1)
+            valid = valid & (seqk == seqq_s)
+        base_k = KSH + b * sks_b + idx[:, None] * sks_s
+        c_kv = tl.load(base_k + offs_l[None, :] * sks_d, mask=valid[:, None], other=0.0)
+        k_rot = tl.load(
+            base_k + (DL + offs_r)[None, :] * sks_d, mask=valid[:, None], other=0.0
+        )
+
+        score = tl.dot(qn, tl.trans(c_kv)) + tl.dot(qr, tl.trans(k_rot))
+        score = tl.where(valid[None, :], score * scale, -float("inf"))
+        lse_safe = tl.where(
+            lse == -float("inf"), 0.0, lse
+        )  # fully-masked query: avoid exp(-inf+inf)=NaN
+        p = tl.exp(score - lse_safe[:, None])  # [BH, BN]
+        dp = tl.dot(dout_l, tl.trans(c_kv))  # [BH, BN] = dout_latent·c_kv
+        ds = (p * (dp - delta[:, None])) * scale  # [BH, BN]  grad wrt scaled score
+        ds_b = ds.to(c_kv.dtype)
+        p_b = p.to(c_kv.dtype)
+
+        dqn += tl.dot(ds_b, c_kv)  # [BH, DL]
+        dqr += tl.dot(ds_b, k_rot)  # [BH, DR]
+        # dc_kv = value-path (Σ_h p·dout_latent) + score-path (Σ_h ds·qn); dk_rot = Σ_h ds·qr
+        dc = tl.dot(tl.trans(p_b), dout_l.to(c_kv.dtype)) + tl.dot(
+            tl.trans(ds_b), qn
+        )  # [BN, DL]
+        dkr = tl.dot(tl.trans(ds_b), qr)  # [BN, DR]
+        dk_ptr = DKSH + b * sks_b + idx[:, None] * sks_s
+        tl.atomic_add(
+            dk_ptr + offs_l[None, :] * sks_d, tl.where(valid[:, None], dc, 0.0)
+        )
+        tl.atomic_add(
+            dk_ptr + (DL + offs_r)[None, :] * sks_d, tl.where(valid[:, None], dkr, 0.0)
+        )
+
+    dq_ptr = DQABS + b * sqa_b + offs_h[:, None] * sqa_h + s * sqa_s
+    tl.store(
+        dq_ptr + offs_l[None, :] * sqa_d,
+        dqn.to(DQABS.dtype.element_ty),
+        mask=hmask[:, None],
+    )
+    tl.store(
+        dq_ptr + (DL + offs_r)[None, :] * sqa_d,
+        dqr.to(DQABS.dtype.element_ty),
+        mask=hmask[:, None],
+    )
+
+
+def _grid_fn(B, S, H):
+    def grid(m):
+        return (S, B, triton.cdiv(H, m["BH"]))
+
+    return grid
+
+
+# zero strides signal dummy seq_q/seq_k when has_doc is False (keeps op signature tensor-only)
+def _doc_strides(seq_q, seq_k, has_doc):
+    if has_doc:
+        return (seq_q.stride(0), seq_q.stride(1), seq_k.stride(0), seq_k.stride(1))
+    return (0, 0, 0, 0)
+
+
+@register_kernel_op("glm_dsa_mla_absorb_attn_fwd")
+def _mla_absorb_fwd_op(
+    q_abs: torch.Tensor,
+    k_shared: torch.Tensor,
+    idx: torch.Tensor,
+    seq_q: torch.Tensor,
+    seq_k: torch.Tensor,
+    scale: float,
+    q_offset: int,
+    has_doc: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, H, S, _ = (
+        q_abs.shape
+    )  # S = local queries; k_shared may be longer (global, under CP)
+    DV = KV_LORA_RANK
+    TOPK = idx.shape[-1]
+    out = torch.empty(B, H, S, DV, device=q_abs.device, dtype=q_abs.dtype)
+    lse = torch.empty(B, H, S, device=q_abs.device, dtype=torch.float32)
+    _absorb_fwd_kernel[_grid_fn(B, S, H)](
+        q_abs,
+        k_shared,
+        idx,
+        out,
+        lse,
+        scale,
+        S,
+        TOPK,
+        H,
+        q_offset,
+        q_abs.stride(0),
+        q_abs.stride(1),
+        q_abs.stride(2),
+        q_abs.stride(3),
+        k_shared.stride(0),
+        k_shared.stride(1),
+        k_shared.stride(2),
+        idx.stride(0),
+        idx.stride(1),
+        idx.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        seq_q,
+        seq_k,
+        *_doc_strides(seq_q, seq_k, has_doc),
+        DL=KV_LORA_RANK,
+        DR=QK_ROPE_HEAD_DIM,
+        HAS_DOC=has_doc,
+    )
+    return out, lse
+
+
+@_mla_absorb_fwd_op.register_fake
+def _(q_abs, k_shared, idx, seq_q, seq_k, scale, q_offset, has_doc):
+    B, H, S, _ = q_abs.shape
+    return (
+        torch.empty(B, H, S, KV_LORA_RANK, device=q_abs.device, dtype=q_abs.dtype),
+        torch.empty(B, H, S, device=q_abs.device, dtype=torch.float32),
+    )
+
+
+@register_kernel_op("glm_dsa_mla_absorb_attn_bwd")
+def _mla_absorb_bwd_op(
+    dout: torch.Tensor,
+    q_abs: torch.Tensor,
+    k_shared: torch.Tensor,
+    idx: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    seq_q: torch.Tensor,
+    seq_k: torch.Tensor,
+    scale: float,
+    q_offset: int,
+    has_doc: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, H, S, DQk = q_abs.shape
+    Skv = k_shared.shape[1]  # global key count (>= S under CP)
+    TOPK = idx.shape[-1]
+    dout = dout.contiguous()
+    dq_abs = torch.zeros(B, H, S, DQk, device=q_abs.device, dtype=q_abs.dtype)
+    # dk_shared scatters at GLOBAL key positions -> size it like k_shared, not the local queries.
+    dk_shared = torch.zeros(B, Skv, DQk, device=q_abs.device, dtype=torch.float32)
+    _absorb_bwd_kernel[_grid_fn(B, S, H)](
+        q_abs,
+        k_shared,
+        idx,
+        out,
+        dout,
+        lse,
+        dq_abs,
+        dk_shared,
+        scale,
+        S,
+        TOPK,
+        H,
+        q_offset,
+        q_abs.stride(0),
+        q_abs.stride(1),
+        q_abs.stride(2),
+        q_abs.stride(3),
+        dk_shared.stride(0),
+        dk_shared.stride(1),
+        dk_shared.stride(2),
+        idx.stride(0),
+        idx.stride(1),
+        idx.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        seq_q,
+        seq_k,
+        *_doc_strides(seq_q, seq_k, has_doc),
+        DL=KV_LORA_RANK,
+        DR=QK_ROPE_HEAD_DIM,
+        HAS_DOC=has_doc,
+    )
+    return dq_abs, dk_shared
+
+
+@_mla_absorb_bwd_op.register_fake
+def _(dout, q_abs, k_shared, idx, out, lse, seq_q, seq_k, scale, q_offset, has_doc):
+    B, H, S, DQk = q_abs.shape
+    return (
+        torch.empty(B, H, S, DQk, device=q_abs.device, dtype=q_abs.dtype),
+        torch.empty(
+            B, k_shared.shape[1], DQk, device=q_abs.device, dtype=torch.float32
+        ),
+    )
+
+
+class _MlaAbsorbAttn(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx, q_abs, k_shared, topk_idx, scale, q_offset, seq_q=None, seq_k=None
+    ):
+        q_abs, k_shared = q_abs.contiguous(), k_shared.contiguous()
+        idx = topk_idx.contiguous()
+        has_doc = seq_q is not None and seq_k is not None
+        if has_doc:
+            seq_q = seq_q.contiguous()
+            seq_k = seq_k.contiguous()
+        else:
+            seq_q, seq_k = idx, idx
+        out, lse = _mla_absorb_fwd_op(
+            q_abs, k_shared, idx, seq_q, seq_k, float(scale), int(q_offset), has_doc
+        )
+        ctx.save_for_backward(q_abs, k_shared, idx, out, lse, seq_q, seq_k)
+        ctx.scale = float(scale)
+        ctx.q_offset = int(q_offset)
+        ctx.has_doc = has_doc
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        q_abs, k_shared, idx, out, lse, seq_q, seq_k = ctx.saved_tensors
+        dq_abs, dk_shared = _mla_absorb_bwd_op(
+            dout,
+            q_abs,
+            k_shared,
+            idx,
+            out,
+            lse,
+            seq_q,
+            seq_k,
+            ctx.scale,
+            ctx.q_offset,
+            ctx.has_doc,
+        )
+        return dq_abs, dk_shared.to(k_shared.dtype), None, None, None, None, None
+
+
+def mla_absorb_attn(
+    q_abs, k_shared, topk_idx, scale, q_offset=0, seq_q=None, seq_k=None
+):
+    """Differentiable head-batched MLA-absorption sparse attention. q_abs [B,H,S,576] (local
+    queries), k_shared [B,Skv,576] (first kv_lora cols = c_kv; Skv>=S under context parallel),
+    topk_idx [B,S,T] int32 referencing GLOBAL key positions. ``q_offset`` is the global position of
+    local query 0 (for causal masking under CP). Returns out_latent [B,H,S,kv_lora]."""
+    return _MlaAbsorbAttn.apply(
+        q_abs, k_shared, topk_idx, scale, q_offset, seq_q, seq_k
+    )
+
+
+def mla_absorb_attn_fwd(q_abs, k_shared, topk_idx, scale, q_offset=0):
+    """Forward-only (no autograd) — for benchmarking. Returns (out_latent, None)."""
+    with torch.no_grad():
+        return _MlaAbsorbAttn.apply(q_abs, k_shared, topk_idx, scale, q_offset), None
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/attention_topk.py b/src/axolotl/integrations/kernels/libs/glm_dsa/attention_topk.py
new file mode 100644
index 0000000000..3898cf58ce
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/attention_topk.py
@@ -0,0 +1,394 @@
+"""Sparse top-k MLA attention for GLM-5.2 DSA (forward + backward, training-correct).
+
+Not on the live training path: the runtime dispatcher (``dispatch.py``/``patch.py``) uses the
+weight-absorption kernel in ``attention_mla_absorb`` instead. This head-per-program baseline is
+retained only as an autotune/benchmarking reference (see ``autotune_collector``).
+
+The DSA training path (eager/sdpa) scatters the per-query top-k indices into a dense ``[B,S,T]``
+additive mask and runs full attention over all T keys — O(S^2) memory + compute even though only
+``index_topk`` (2048) keys per query matter. This kernel **gathers** each query's selected keys and
+runs flash online-softmax over just those: O(S·topk).
+
+``sparse_attn`` is a ``torch.autograd.Function`` — fwd + bwd — so it composes with training. It
+differentiates wrt q, k, v (the MLA projection activations), so gradients flow to the projection
+weights in **full-parameter** training and to the LoRA adapters in **LoRA** training identically;
+the kernel is agnostic to how q/k/v were produced. The top-k indices are treated as constants
+(``topk`` is non-differentiable, exactly as in the eager model), and the DSA indexer itself is
+``@torch.no_grad`` upstream, so no gradient flows through key selection.
+
+Layout: q,k ``[B,H,S,D]`` (D=qk_head_dim=256), v ``[B,H,S,Dv]`` (Dv=v_head_dim=256), ``topk_idx``
+``[B,S,T]`` int32 (selected key positions, shared across heads). Gathered indices beyond the query
+position are causal-masked so over-selected top-k padding contributes nothing (fwd) and gets no
+gradient (bwd). One program = (b, h, query position s); MMA-efficient batching is a later
+optimization — this is the correct, gradient-validated baseline.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BT": bt}, num_warps=w, num_stages=s)
+        for bt in (64, 128, 256)
+        for w in (4, 8)
+        for s in (2, 3)
+    ],
+    key=["TOPK", "D", "DV"],
+)
+@triton.jit
+def _fwd_kernel(
+    Q,
+    K,
+    V,
+    IDX,
+    OUT,
+    LSE,
+    scale,
+    S,
+    TOPK,
+    sq_b,
+    sq_h,
+    sq_s,
+    sq_d,
+    sk_b,
+    sk_h,
+    sk_s,
+    sk_d,
+    sv_b,
+    sv_h,
+    sv_s,
+    sv_d,
+    si_b,
+    si_s,
+    si_t,
+    so_b,
+    so_h,
+    so_s,
+    so_d,
+    sl_b,
+    sl_h,
+    sl_s,
+    H,
+    D: tl.constexpr,
+    DV: tl.constexpr,
+    BT: tl.constexpr,
+):
+    pid_bh = tl.program_id(0)
+    s = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+
+    offs_d = tl.arange(0, D)
+    offs_dv = tl.arange(0, DV)
+    q = tl.load(Q + b * sq_b + h * sq_h + s * sq_s + offs_d * sq_d).to(tl.float32)
+
+    m_i = -float("inf")
+    l_i = 0.0
+    acc = tl.zeros((DV,), dtype=tl.float32)
+
+    for t0 in range(0, TOPK, BT):
+        offs_t = t0 + tl.arange(0, BT)
+        tmask = offs_t < TOPK
+        idx = tl.load(
+            IDX + b * si_b + s * si_s + offs_t * si_t, mask=tmask, other=0
+        ).to(tl.int64)
+        valid = tmask & (idx <= s)
+
+        k_ptr = K + b * sk_b + h * sk_h + idx[:, None] * sk_s + offs_d[None, :] * sk_d
+        k = tl.load(k_ptr, mask=valid[:, None], other=0.0).to(tl.float32)  # [BT, D]
+        qk = tl.sum(q[None, :] * k, axis=1) * scale  # [BT]
+        qk = tl.where(valid, qk, -float("inf"))
+
+        m_new = tl.maximum(m_i, tl.max(qk, axis=0))
+        p = tl.exp(qk - m_new)
+        alpha = tl.exp(m_i - m_new)
+        l_i = l_i * alpha + tl.sum(p, axis=0)
+        v_ptr = V + b * sv_b + h * sv_h + idx[:, None] * sv_s + offs_dv[None, :] * sv_d
+        v = tl.load(v_ptr, mask=valid[:, None], other=0.0).to(tl.float32)  # [BT, DV]
+        acc = acc * alpha + tl.sum(p[:, None] * v, axis=0)
+        m_i = m_new
+
+    acc = acc / l_i
+    tl.store(
+        OUT + b * so_b + h * so_h + s * so_s + offs_dv * so_d,
+        acc.to(OUT.dtype.element_ty),
+    )
+    tl.store(LSE + b * sl_b + h * sl_h + s * sl_s, m_i + tl.log(l_i))
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BT": bt}, num_warps=w, num_stages=s)
+        for bt in (64, 128, 256)
+        for w in (4, 8)
+        for s in (2, 3)
+    ],
+    key=["TOPK", "D", "DV"],
+    # dK/dV are atomic-accumulated; zero them between autotune trials or benchmarking corrupts them.
+    reset_to_zero=["DKO", "DVO"],
+)
+@triton.jit
+def _bwd_kernel(
+    Q,
+    K,
+    V,
+    IDX,
+    OUT,
+    DOUT,
+    LSE,
+    DQO,
+    DKO,
+    DVO,
+    scale,
+    S,
+    TOPK,
+    sq_b,
+    sq_h,
+    sq_s,
+    sq_d,
+    sk_b,
+    sk_h,
+    sk_s,
+    sk_d,
+    sv_b,
+    sv_h,
+    sv_s,
+    sv_d,
+    si_b,
+    si_s,
+    si_t,
+    so_b,
+    so_h,
+    so_s,
+    so_d,
+    sl_b,
+    sl_h,
+    sl_s,
+    H,
+    D: tl.constexpr,
+    DV: tl.constexpr,
+    BT: tl.constexpr,
+):
+    pid_bh = tl.program_id(0)
+    s = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+
+    offs_d = tl.arange(0, D)
+    offs_dv = tl.arange(0, DV)
+    q = tl.load(Q + b * sq_b + h * sq_h + s * sq_s + offs_d * sq_d).to(tl.float32)
+    dout = tl.load(DOUT + b * so_b + h * so_h + s * so_s + offs_dv * so_d).to(
+        tl.float32
+    )
+    out = tl.load(OUT + b * so_b + h * so_h + s * so_s + offs_dv * so_d).to(tl.float32)
+    lse = tl.load(LSE + b * sl_b + h * sl_h + s * sl_s)
+    delta = tl.sum(dout * out, axis=0)  # = Σ_j p_j (dout·V_j)
+
+    dq = tl.zeros((D,), dtype=tl.float32)
+    for t0 in range(0, TOPK, BT):
+        offs_t = t0 + tl.arange(0, BT)
+        tmask = offs_t < TOPK
+        idx = tl.load(
+            IDX + b * si_b + s * si_s + offs_t * si_t, mask=tmask, other=0
+        ).to(tl.int64)
+        valid = tmask & (idx <= s)
+
+        k_ptr = K + b * sk_b + h * sk_h + idx[:, None] * sk_s + offs_d[None, :] * sk_d
+        k = tl.load(k_ptr, mask=valid[:, None], other=0.0).to(tl.float32)  # [BT, D]
+        v_ptr = V + b * sv_b + h * sv_h + idx[:, None] * sv_s + offs_dv[None, :] * sv_d
+        v = tl.load(v_ptr, mask=valid[:, None], other=0.0).to(tl.float32)  # [BT, DV]
+
+        s_j = tl.sum(q[None, :] * k, axis=1) * scale  # [BT]
+        p = tl.where(valid, tl.exp(s_j - lse), 0.0)  # [BT]
+        dp = tl.sum(dout[None, :] * v, axis=1)  # [BT] = dout·V_j
+        ds = p * (dp - delta)  # [BT] grad wrt pre-softmax score
+
+        dq += tl.sum((ds * scale)[:, None] * k, axis=0)  # [D]
+        # scatter-add dK_j = scale·ds_j·q, dV_j = p_j·dout  (race across query positions -> atomic)
+        dk_j = (ds * scale)[:, None] * q[None, :]  # [BT, D]
+        dv_j = p[:, None] * dout[None, :]  # [BT, DV]
+        dk_ptr = (
+            DKO + b * sk_b + h * sk_h + idx[:, None] * sk_s + offs_d[None, :] * sk_d
+        )
+        dv_ptr = (
+            DVO + b * sv_b + h * sv_h + idx[:, None] * sv_s + offs_dv[None, :] * sv_d
+        )
+        tl.atomic_add(dk_ptr, tl.where(valid[:, None], dk_j, 0.0))
+        tl.atomic_add(dv_ptr, tl.where(valid[:, None], dv_j, 0.0))
+
+    tl.store(DQO + b * sq_b + h * sq_h + s * sq_s + offs_d * sq_d, dq)
+
+
+@register_kernel_op("glm_dsa_sparse_topk_attn_fwd")
+def _sparse_topk_fwd_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    idx: torch.Tensor,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, H, S, D = q.shape
+    DV = v.shape[-1]
+    TOPK = idx.shape[-1]
+    out = torch.empty(B, H, S, DV, device=q.device, dtype=q.dtype)
+    lse = torch.empty(B, H, S, device=q.device, dtype=torch.float32)
+    grid = (B * H, S)
+    _fwd_kernel[grid](
+        q,
+        k,
+        v,
+        idx,
+        out,
+        lse,
+        scale,
+        S,
+        TOPK,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        k.stride(3),
+        v.stride(0),
+        v.stride(1),
+        v.stride(2),
+        v.stride(3),
+        idx.stride(0),
+        idx.stride(1),
+        idx.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        H,
+        D=D,
+        DV=DV,
+    )
+    return out, lse
+
+
+@_sparse_topk_fwd_op.register_fake
+def _(q, k, v, idx, scale):
+    B, H, S, _ = q.shape
+    return (
+        torch.empty(B, H, S, v.shape[-1], device=q.device, dtype=q.dtype),
+        torch.empty(B, H, S, device=q.device, dtype=torch.float32),
+    )
+
+
+@register_kernel_op("glm_dsa_sparse_topk_attn_bwd")
+def _sparse_topk_bwd_op(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    idx: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, H, S, D = q.shape
+    DV = v.shape[-1]
+    TOPK = idx.shape[-1]
+    dout = dout.contiguous()
+    dq = torch.zeros_like(q, dtype=torch.float32)
+    # dK/dV accumulate across query positions selecting the same key -> fp32 atomics.
+    dk = torch.zeros(B, H, S, D, device=q.device, dtype=torch.float32)
+    dv = torch.zeros(B, H, S, DV, device=q.device, dtype=torch.float32)
+    grid = (B * H, S)
+    _bwd_kernel[grid](
+        q,
+        k,
+        v,
+        idx,
+        out,
+        dout,
+        lse,
+        dq,
+        dk,
+        dv,
+        scale,
+        S,
+        TOPK,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        dk.stride(0),
+        dk.stride(1),
+        dk.stride(2),
+        dk.stride(3),
+        dv.stride(0),
+        dv.stride(1),
+        dv.stride(2),
+        dv.stride(3),
+        idx.stride(0),
+        idx.stride(1),
+        idx.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        H,
+        D=D,
+        DV=DV,
+    )
+    return dq, dk, dv
+
+
+@_sparse_topk_bwd_op.register_fake
+def _(dout, q, k, v, idx, out, lse, scale):
+    B, H, S, D = q.shape
+    DV = v.shape[-1]
+    return (
+        torch.empty_like(q, dtype=torch.float32),
+        torch.empty(B, H, S, D, device=q.device, dtype=torch.float32),
+        torch.empty(B, H, S, DV, device=q.device, dtype=torch.float32),
+    )
+
+
+class _SparseAttnTopK(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, topk_idx, scale):
+        H = q.shape[1]
+        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+        idx = topk_idx.contiguous()
+        out, lse = _sparse_topk_fwd_op(q, k, v, idx, float(scale))
+        ctx.save_for_backward(q, k, v, idx, out, lse)
+        ctx.scale = float(scale)
+        ctx.H = H
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        q, k, v, idx, out, lse = ctx.saved_tensors
+        dq, dk, dv = _sparse_topk_bwd_op(dout, q, k, v, idx, out, lse, ctx.scale)
+        return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), None, None
+
+
+def sparse_attn(q, k, v, topk_idx, scale):
+    """Differentiable sparse top-k MLA attention. q,k [B,H,S,D], v [B,H,S,Dv], topk_idx [B,S,T]
+    int32. Returns out [B,H,S,Dv]; backprops dq/dk/dv (topk_idx is a constant)."""
+    return _SparseAttnTopK.apply(q, k, v, topk_idx, scale)
+
+
+def sparse_attn_fwd(q, k, v, topk_idx, scale):
+    """Forward-only (no autograd) — for benchmarking the kernel in isolation."""
+    with torch.no_grad():
+        return _SparseAttnTopK.apply(
+            q.detach(), k.detach(), v.detach(), topk_idx, scale
+        )
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/config.py b/src/axolotl/integrations/kernels/libs/glm_dsa/config.py
new file mode 100644
index 0000000000..df23457b55
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/config.py
@@ -0,0 +1,54 @@
+"""Real GLM-5.2 (glm_moe_dsa) shapes for kernel development + validation.
+
+Kernels are validated at the REAL model dims (head_dim 256, index_head_dim 128, top-k 2048, ...)
+with a truncated layer count / per-op isolation — shrinking the dims would hide shape-specific
+behavior (SMEM pressure at head_dim 256, the 128-wide indexer dot, the 2048 top-k gather). These
+mirror ``GlmMoeDsaConfig`` defaults == the published GLM-5.2 checkpoint.
+"""
+
+from __future__ import annotations
+
+# --- attention (DeepSeek-V3.2 MLA) ---
+HIDDEN = 6144
+N_HEADS = 64
+N_KV_HEADS = 64
+Q_LORA_RANK = 2048
+KV_LORA_RANK = 512
+QK_NOPE_HEAD_DIM = 192
+QK_ROPE_HEAD_DIM = 64
+QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM  # 256
+V_HEAD_DIM = 256
+ATTN_SCALE = QK_HEAD_DIM**-0.5
+ROPE_THETA = 8_000_000.0
+
+# --- DSA Lightning indexer ---
+INDEX_N_HEADS = 32
+INDEX_HEAD_DIM = 128
+INDEX_TOPK = 2048
+INDEX_SOFTMAX_SCALE = INDEX_HEAD_DIM**-0.5
+
+# --- MoE (not exercised by the attention/indexer/rope kernels) ---
+N_ROUTED_EXPERTS = 256
+N_SHARED_EXPERTS = 1
+NUM_EXPERTS_PER_TOK = 8
+MOE_INTERMEDIATE = 2048
+
+NUM_HIDDEN_LAYERS = 78
+FIRST_K_DENSE_REPLACE = 3
+INDEX_TOPK_FREQ = 4
+
+
+def probe_shapes(seq: int = 4096, batch: int = 1) -> dict:
+    """A single (batch, seq) probe at real dims for kernel tests/benches."""
+    return {
+        "B": batch,
+        "S": seq,
+        "H": N_HEADS,
+        "D": QK_HEAD_DIM,
+        "Dv": V_HEAD_DIM,
+        "IDX_H": INDEX_N_HEADS,
+        "IDX_D": INDEX_HEAD_DIM,
+        "TOPK": min(INDEX_TOPK, seq),
+        "scale": ATTN_SCALE,
+        "idx_scale": INDEX_SOFTMAX_SCALE,
+    }
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/context_parallel.py b/src/axolotl/integrations/kernels/libs/glm_dsa/context_parallel.py
new file mode 100644
index 0000000000..26e2710d61
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/context_parallel.py
@@ -0,0 +1,107 @@
+"""Context-parallel DSA attention for GLM-5.2.
+
+GLM-5.2 is always multi-GPU; at long context the sequence is sharded across a context-parallel (CP)
+group. The MLA compression makes CP cheap: each rank needs every key, but a key is only the 576-wide
+*compressed* ``k_shared`` (kv_lora 512 + rope 64) — ~1.1 KB/token — NOT the per-head 64x256 expanded
+keys (~32 KB/token). So we all-gather the compressed KV (a ~28x smaller collective than gathering
+expanded K/V) and each rank attends its local queries against the global KV with the absorbed kernel
+(already global-KV-aware via ``q_offset`` + global top-k indices).
+
+The all-gather is differentiable (its backward reduce-scatters the per-rank ``dk_shared`` grads), so
+training is correct. ``topk_idx`` must reference GLOBAL key positions (the indexer scores against the
+same gathered KV). Sequence sharding here is contiguous (rank r owns ``[r·L, (r+1)·L)``).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.distributed as dist
+
+from .attention_mla_absorb import absorb_query
+from .dispatch import mla_attn
+
+_COMM_STREAM = {}
+
+
+def _comm_stream(dev):
+    if dev not in _COMM_STREAM:
+        _COMM_STREAM[dev] = torch.cuda.Stream(device=dev)
+    return _COMM_STREAM[dev]
+
+
+class _OverlapGather(torch.autograd.Function):
+    """Differentiable seq all-gather whose data movement was already started (async) before this op,
+    so it overlaps with intervening compute. Backward all-reduces the global grad and slices the
+    rank-local part (== the reduce-scatter the standard differentiable all-gather does)."""
+
+    @staticmethod
+    def forward(ctx, k_local, gathered, work, group, world, rank):
+        if work is not None:
+            work.wait()
+        ctx.group, ctx.world, ctx.rank, ctx.sl = group, world, rank, k_local.shape[1]
+        return torch.cat(gathered, dim=1)
+
+    @staticmethod
+    def backward(ctx, dg):
+        dg = dg.contiguous()
+        dist.all_reduce(dg, op=dist.ReduceOp.SUM, group=ctx.group)
+        s = slice(ctx.rank * ctx.sl, (ctx.rank + 1) * ctx.sl)
+        return dg[:, s].contiguous(), None, None, None, None, None
+
+
+def cp_mla_attn_overlapped(
+    q_pass, q_rot, w_kb_k, k_shared, topk_idx, scale, group=None, crossover=None
+):
+    """Comm-overlapped CP attention: issue the compressed-KV all-gather on a side stream, compute the
+    absorption GEMM (local, ~hidden behind the gather), then attend. Same result as cp_mla_attn but
+    the absorption is hidden behind the collective. Takes the RAW query projections so the GEMM can
+    run in the overlap window. Returns the local out_latent."""
+    world = dist.get_world_size(group) if dist.is_initialized() else 1
+    if world == 1:
+        q_abs = absorb_query(q_pass, q_rot, w_kb_k)
+        return mla_attn(
+            q_abs, k_shared, topk_idx, scale, q_offset=0, crossover=crossover
+        )
+    rank = dist.get_rank(group)
+    s_local = q_pass.shape[2]
+    gathered = [torch.empty_like(k_shared) for _ in range(world)]
+    stream = _comm_stream(k_shared.device)
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        work = dist.all_gather(
+            gathered, k_shared.contiguous(), group=group, async_op=True
+        )
+    q_abs = absorb_query(
+        q_pass, q_rot, w_kb_k
+    )  # overlaps with the gather on the side stream
+    torch.cuda.current_stream().wait_stream(stream)
+    k_global = _OverlapGather.apply(k_shared, gathered, work, group, world, rank)
+    return mla_attn(
+        q_abs, k_global, topk_idx, scale, q_offset=rank * s_local, crossover=crossover
+    )
+
+
+def all_gather_seq(x: torch.Tensor, group=None) -> torch.Tensor:
+    """Differentiable all-gather along the sequence dim (1): [B,S_local,D] -> [B,S_global,D]."""
+    world = dist.get_world_size(group) if dist.is_initialized() else 1
+    if world == 1:
+        return x
+    import torch.distributed.nn as dist_nn
+
+    parts = dist_nn.all_gather(x.contiguous(), group=group)
+    return torch.cat(list(parts), dim=1)
+
+
+def cp_mla_attn(q_abs, k_shared, topk_idx, scale, group=None, crossover=None):
+    """Context-parallel absorbed DSA attention. ``q_abs`` [B,H,S_local,576], ``k_shared``
+    [B,S_local,576], ``topk_idx`` [B,S_local,T] (GLOBAL key positions) are the rank-local shards.
+    All-gathers the compressed KV (cheap), then runs the local queries against the global KV with the
+    correct ``q_offset``. Returns the local out_latent [B,H,S_local,kv_lora]."""
+    rank = dist.get_rank(group) if dist.is_initialized() else 0
+    s_local = q_abs.shape[2]
+    k_global = all_gather_seq(
+        k_shared, group
+    )  # differentiable; backward reduce-scatters dk
+    return mla_attn(
+        q_abs, k_global, topk_idx, scale, q_offset=rank * s_local, crossover=crossover
+    )
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/dispatch.py b/src/axolotl/integrations/kernels/libs/glm_dsa/dispatch.py
new file mode 100644
index 0000000000..2968843055
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/dispatch.py
@@ -0,0 +1,160 @@
+"""Length-aware dispatch for GLM-5.2 DSA attention — robust across short and long sequences.
+
+The sparse gather kernel only wins when ``S >> index_topk`` (it trades MMA locality for skipping
+non-selected keys). Below a hardware-dependent crossover, a DENSE flash over the (absorbed) shared
+KV with the top-k applied as an additive mask is faster; and at ``S <= index_topk`` the top-k
+selects every causal key, so it degenerates to plain causal attention anyway.
+
+``mla_attn`` dispatches on ``S``: dense flash below the crossover, gather above. Both produce the
+same ``out_latent`` and are differentiable, so training is correct in either regime. The crossover
+is **auto-calibrated once** per ``(topk, H, dtype, device)`` by microbenchmarking the two paths —
+so it adapts to sm120 (~22k) vs a datacenter GPU (much lower) instead of hardcoding a threshold.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from .attention_mla_absorb import DQK, mla_absorb_attn
+from .config import KV_LORA_RANK
+
+_CROSSOVER: dict[tuple, int] = {}
+_GATHER_OK: dict[int, bool] = {}
+
+
+def _gather_supported(device) -> bool:
+    """Whether to use the sparse absorbed-MLA gather kernel (vs the dense fallback).
+
+    Validated on sm120 (Ada / consumer-Blackwell) and sm90 (Hopper). The earlier sm90 blockers are
+    fixed: the backward's nondeterministic invalid-PC (the full autotune grid trialing — each config
+    is individually clean) is avoided by forcing a single config on sm90 (``_bwd_prune_sm90_safe`` in
+    ``attention_mla_absorb``), and the leading-all-masked-block softmax NaN is guarded in the kernel.
+    Set ``GLM_DSA_DISABLE_GATHER=1`` to force the dense path, or ``GLM_DSA_FORCE_GATHER=1`` to use the
+    gather on an unlisted arch."""
+    import os
+
+    if os.environ.get("GLM_DSA_DISABLE_GATHER"):
+        return False
+    if os.environ.get("GLM_DSA_FORCE_GATHER"):
+        return True
+    dev = getattr(device, "index", None)
+    if dev is None:
+        dev = torch.cuda.current_device()
+    if dev not in _GATHER_OK:
+        _GATHER_OK[dev] = torch.cuda.get_device_capability(dev) in ((9, 0), (12, 0))
+    return _GATHER_OK[dev]
+
+
+def _topk_causal_mask(topk_idx, Skv, q_offset, dtype, device, seq_q=None, seq_k=None):
+    """Additive [B,1,S,Skv] mask: 0 at selected-and-causal global keys, -inf elsewhere. ``topk_idx``
+    references GLOBAL key positions; query i (local) is global position ``q_offset + i``. Under
+    sample packing, ``seq_q`` [B,S] / ``seq_k`` [B,Skv] further forbid cross-document keys (the
+    selected set may include earlier-document keys when ``topk == Skv``)."""
+    B, S = topk_idx.shape[:2]
+    sel = torch.zeros(B, S, Skv, dtype=torch.bool, device=device)
+    sel.scatter_(-1, topk_idx.long(), True)
+    kpos = torch.arange(Skv, device=device)
+    qpos = q_offset + torch.arange(S, device=device)
+    keep = sel & (kpos[None, None, :] <= qpos[None, :, None])
+    if seq_q is not None and seq_k is not None:
+        keep = keep & (seq_k[:, None, :] == seq_q[:, :, None])
+    mask = torch.zeros(B, 1, S, Skv, dtype=dtype, device=device)
+    return mask.masked_fill_(~keep.unsqueeze(1), float("-inf"))
+
+
+def dense_masked_out_latent(
+    q_abs, k_shared, topk_idx, scale, q_offset=0, seq_q=None, seq_k=None
+):
+    """Dense flash over the absorbed shared KV with the top-k as an additive mask (O(S) memory via
+    SDPA; the mask itself is O(S·Skv)). Differentiable. Faster than the gather below the crossover.
+    ``k_shared`` may be the global (gathered) KV under context parallel; ``q_offset`` is local q0's
+    global position. ``seq_q``/``seq_k`` enable per-document masking under sample packing."""
+    B, H, S, _ = q_abs.shape
+    Skv = k_shared.shape[1]
+    c_kv = k_shared[..., :KV_LORA_RANK]
+    k_e = k_shared.unsqueeze(1).expand(B, H, Skv, k_shared.shape[-1])
+    v_e = c_kv.unsqueeze(1).expand(B, H, Skv, KV_LORA_RANK)
+    mask = _topk_causal_mask(
+        topk_idx, Skv, q_offset, q_abs.dtype, q_abs.device, seq_q, seq_k
+    )
+    return F.scaled_dot_product_attention(q_abs, k_e, v_e, attn_mask=mask, scale=scale)
+
+
+def _bench(fn, it=8):
+    import time
+
+    for _ in range(3):
+        fn()
+    torch.cuda.synchronize()
+    t = time.time()
+    for _ in range(it):
+        fn()
+    torch.cuda.synchronize()
+    return time.time() - t
+
+
+def calibrate_crossover(topk, H, dtype, device, B=1) -> int:
+    """Microbench gather vs dense at increasing S; return the smallest S where the gather wins
+    (memoized). Falls back to a sparsity heuristic if a probe OOMs."""
+    key = (topk, H, str(dtype), torch.cuda.get_device_name(device))
+    if key in _CROSSOVER:
+        return _CROSSOVER[key]
+    crossover = None
+    for mult in (2, 3, 4, 6, 8, 12, 16):
+        S = topk * mult
+        try:
+            q = torch.randn(B, H, S, DQK, device=device, dtype=dtype)
+            k = torch.randn(B, S, DQK, device=device, dtype=dtype)
+            idx = torch.stack(
+                [torch.arange(topk, device=device).clamp(max=s).int() for s in range(S)]
+            ).unsqueeze(0)
+            tg = _bench(lambda q=q, k=k, idx=idx: mla_absorb_attn(q, k, idx, 1.0).sum())
+            td = _bench(
+                lambda q=q, k=k, idx=idx: dense_masked_out_latent(q, k, idx, 1.0).sum()
+            )
+        except torch.cuda.OutOfMemoryError:
+            crossover = crossover or S  # dense OOMed -> must use gather from here
+            break
+        if tg < td:
+            crossover = S
+            break
+    crossover = crossover or topk * 16
+    _CROSSOVER[key] = crossover
+    return crossover
+
+
+def mla_attn(
+    q_abs,
+    k_shared,
+    topk_idx,
+    scale,
+    q_offset=0,
+    crossover: int | None = None,
+    seq_q=None,
+    seq_k=None,
+):
+    """Differentiable DSA attention, dispatched by total key length (the sparsity that matters):
+    dense flash below the (auto-calibrated) crossover, sparse gather above. ``k_shared`` may be the
+    global gathered KV under context parallel; ``q_offset`` is local query 0's global position.
+    Returns out_latent [B,H,S,kv_lora].
+
+    Under sample packing (``seq_q``/``seq_k`` document ids set) the gather is doc-aware: it masks
+    cross-document keys in-kernel (``seq_k[idx] == seq_q[s]``), so packing keeps the sparse path
+    instead of falling back to the dense per-document mask."""
+    Skv = k_shared.shape[1]
+    topk = topk_idx.shape[-1]
+    # Check arch support BEFORE calibrating — calibrate_crossover benchmarks the gather, so on an
+    # unsupported GPU calibration would itself compile/run the gather Triton path we're avoiding.
+    if _gather_supported(q_abs.device):
+        if crossover is None:
+            crossover = calibrate_crossover(
+                topk, q_abs.shape[1], q_abs.dtype, q_abs.device
+            )
+        if Skv >= crossover:
+            return mla_absorb_attn(
+                q_abs, k_shared, topk_idx, scale, q_offset, seq_q, seq_k
+            )
+    return dense_masked_out_latent(
+        q_abs, k_shared, topk_idx, scale, q_offset, seq_q, seq_k
+    )
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/indexer.py b/src/axolotl/integrations/kernels/libs/glm_dsa/indexer.py
new file mode 100644
index 0000000000..c791e27205
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/indexer.py
@@ -0,0 +1,176 @@
+"""Fused GLM-5.2 DSA Lightning-Indexer scorer + top-k (forward-only).
+
+Reference (GlmMoeDsaIndexer.forward, post-RoPE):
+    scores       = relu(softmax_scale * q @ k^T)                         # [B, S, H, T]
+    index_scores = Σ_h weights[b,s,h] * scores[b,s,h,t]                  # [B, S, T]
+    (+ causal mask) -> topk(index_topk).indices                          # [B, S, topk] int32
+
+where q [B,S,H,D] (H=index_n_heads=32, D=index_head_dim=128), k [B,S,D] (shared across heads),
+weights[b,s,h] = weights_proj(h)[b,s,h] * H**-0.5.
+
+Eager materializes the full [B,S,H,T] fp32 score tensor (H=32) — the transient hotspot. Fusing
+the H-reduction collapses straight to [B,S,T], never holding the H axis. The indexer runs under
+``@torch.no_grad`` (its output only feeds ``topk().indices``), so this is FORWARD-ONLY.
+
+The scoring kernel is the DeepSeek-V4 ``indexer_scores`` kernel (identical math: relu(sm·qk) summed
+over weighted heads); GLM adds the causal mask + top-k here. ``relu(sm·x) == sm·relu(x)`` for sm>0,
+so applying the scale inside or outside relu is equivalent.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BS": bs, "BT": bt}, num_warps=w, num_stages=s)
+        for bs in (32, 64, 128)
+        for bt in (32, 64, 128)
+        for w in (4, 8)
+        for s in (2, 3)
+    ],
+    key=["S", "T", "H"],
+)
+@triton.jit
+def _indexer_score_kernel(
+    Q,
+    K,
+    W,
+    OUT,
+    softmax_scale,
+    S,
+    T,
+    H,
+    sq_b,
+    sq_s,
+    sq_h,
+    sq_d,
+    sk_b,
+    sk_t,
+    sk_d,
+    sw_b,
+    sw_s,
+    sw_h,
+    so_b,
+    so_s,
+    so_t,
+    DH: tl.constexpr,
+    BS: tl.constexpr,
+    BT: tl.constexpr,
+):
+    b = tl.program_id(0)
+    s0 = tl.program_id(1) * BS
+    t0 = tl.program_id(2) * BT
+    offs_s = s0 + tl.arange(0, BS)
+    offs_t = t0 + tl.arange(0, BT)
+    offs_d = tl.arange(0, DH)
+    smask = offs_s < S
+    tmask = offs_t < T
+
+    # k block [DH, BT] (transposed for q @ kᵀ), shared across heads
+    k_ptr = K + b * sk_b + offs_d[:, None] * sk_d + offs_t[None, :] * sk_t
+    k = tl.load(k_ptr, mask=tmask[None, :], other=0.0).to(tl.float32)
+
+    acc = tl.zeros((BS, BT), dtype=tl.float32)
+    for h in range(H):
+        q_ptr = (
+            Q + b * sq_b + offs_s[:, None] * sq_s + h * sq_h + offs_d[None, :] * sq_d
+        )
+        q = tl.load(q_ptr, mask=smask[:, None], other=0.0).to(tl.float32)
+        qk = tl.dot(q, k, input_precision="ieee")  # [BS, BT]
+        qk = tl.maximum(qk, 0.0) * softmax_scale  # relu · scale
+        w = tl.load(W + b * sw_b + offs_s * sw_s + h * sw_h, mask=smask, other=0.0).to(
+            tl.float32
+        )
+        acc += qk * w[:, None]
+
+    out_ptr = OUT + b * so_b + offs_s[:, None] * so_s + offs_t[None, :] * so_t
+    tl.store(out_ptr, acc, mask=smask[:, None] & tmask[None, :])
+
+
+def indexer_scores(
+    q: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, softmax_scale: float
+) -> torch.Tensor:
+    """q [B,S,H,D], k [B,S,D] (shared across heads), weights [B,S,H] (already ·H**-0.5).
+    Returns index_scores [B,S,T=S] (fp32, no grad). Fuses the H-reduction."""
+    B, S, H, DH = q.shape
+    T = k.shape[1]
+    out = torch.empty(B, S, T, device=q.device, dtype=torch.float32)
+    if T == 0 or S == 0:
+        return out
+    if k.dtype != q.dtype:
+        k = k.to(q.dtype)
+    q, k, w = q.contiguous(), k.contiguous(), weights.contiguous()
+
+    def grid(m):
+        return (B, triton.cdiv(S, m["BS"]), triton.cdiv(T, m["BT"]))
+
+    _indexer_score_kernel[grid](
+        q,
+        k,
+        w,
+        out,
+        float(softmax_scale),
+        S,
+        T,
+        H,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        w.stride(0),
+        w.stride(1),
+        w.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        DH=DH,
+    )
+    return out
+
+
+@torch.no_grad()
+def indexer_topk(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    weights: torch.Tensor,
+    softmax_scale: float,
+    index_topk: int,
+    attention_mask: torch.Tensor | None = None,
+    position_ids: torch.Tensor | None = None,
+    seq_q: torch.Tensor | None = None,
+    seq_k: torch.Tensor | None = None,
+    q_offset: int = 0,
+) -> torch.Tensor:
+    """Fused indexer scoring + causal mask + top-k. Returns int32 indices [B,S,topk].
+
+    ``attention_mask`` (additive, broadcastable to [B,S,T]) is added if given; otherwise a causal
+    mask from ``position_ids`` (or arange) is applied. Mirrors GlmMoeDsaIndexer's masking + topk.
+
+    Under sample packing, ``seq_q`` [B,S] / ``seq_k`` [B,T] give each query/key its document id;
+    the mask then forbids cross-document keys (a key in an earlier packed document is causally
+    "before" the query by global index but must not be attended). ``q_offset`` shifts local query
+    indices to global positions (context parallel). The selected indices still rank same-document
+    keys first; the attention kernel re-applies the same document mask (``mla_attn`` forces the
+    dense path under packing) so cross-document keys returned when ``topk == T`` are dropped.
+    """
+    scores = indexer_scores(q, k, weights, softmax_scale)  # [B,S,T]
+    B, S, T = scores.shape
+    packed = seq_q is not None and seq_k is not None
+    if attention_mask is not None and not packed:
+        scores = scores + attention_mask.to(scores.dtype)
+    else:
+        kpos = torch.arange(T, device=scores.device)
+        qpos = q_offset + torch.arange(S, device=scores.device)
+        causal = kpos[None, None, :] > qpos[None, :, None]
+        if seq_q is not None and seq_k is not None:
+            causal = causal | (seq_k[:, None, :] != seq_q[:, :, None])
+        scores = scores.masked_fill(causal, float("-inf"))
+    topk = min(index_topk, T)
+    return scores.topk(topk, dim=-1).indices.to(torch.int32)
diff --git a/src/axolotl/integrations/kernels/libs/glm_dsa/patch.py b/src/axolotl/integrations/kernels/libs/glm_dsa/patch.py
new file mode 100644
index 0000000000..c570cb75c1
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/glm_dsa/patch.py
@@ -0,0 +1,279 @@
+"""Patch GlmMoeDsaAttention.forward to use the absorbed-MLA + dispatched sparse-gather kernels.
+
+Replaces the model's dense-mask eager/sdpa attention with: build q_abs / k_shared from the MLA
+projections (compressed latent + interleaved-RoPE rope part), call ``mla_attn`` (dense flash below
+the auto-calibrated crossover, sparse gather above), then ``project_value`` + ``o_proj``. The DSA
+top-k selection is taken from the model's own indexer (``"full"`` layers) or the previous full
+layer's indices (``"shared"`` layers), unchanged. The kv_b_proj weight is read through
+``effective_weight`` so the absorption is correct (and differentiable) under both full-parameter
+and LoRA training.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.distributed as dist
+
+from .attention_mla_absorb import absorb_query, project_value, split_kv_b
+from .context_parallel import all_gather_seq
+from .dispatch import mla_attn
+from .indexer import indexer_topk
+
+
+@torch.no_grad()
+def fused_indexer_topk(
+    indexer,
+    hidden_states,
+    q_resid,
+    position_embeddings,
+    attention_mask,
+    position_ids,
+    group=None,
+    seq_q=None,
+    seq_k=None,
+    q_offset=0,
+):
+    """Replicates GlmMoeDsaIndexer.forward with the fused scorer (collapses the [B,S,32,T] reduction
+    to [B,S,T], the long-context memory win). Forward-only. Under context parallel (``group`` set),
+    all-gathers the indexer's keys so local queries score against the GLOBAL keys, and uses the
+    global ``position_ids`` for the causal mask. Returns top-k of GLOBAL key positions."""
+    from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import (
+        apply_rotary_pos_emb,
+    )
+
+    B, S, _ = hidden_states.shape
+    q = indexer.wq_b(q_resid).view(B, S, indexer.n_heads, indexer.head_dim)
+    q_rot, q_pass = torch.split(
+        q,
+        [indexer.qk_rope_head_dim, indexer.head_dim - indexer.qk_rope_head_dim],
+        dim=-1,
+    )
+    k = indexer.k_norm(indexer.wk(hidden_states)).unsqueeze(2)  # [B,S,1,D]
+    k_rot, k_pass = torch.split(
+        k,
+        [indexer.qk_rope_head_dim, indexer.head_dim - indexer.qk_rope_head_dim],
+        dim=-1,
+    )
+    cos, sin = position_embeddings
+    q_rot, k_rot = apply_rotary_pos_emb(
+        q_rot, k_rot, cos, sin, unsqueeze_dim=2
+    )  # non-interleaved
+    q = torch.cat([q_rot, q_pass], dim=-1)  # [B,S_local,32,128]
+    k = torch.cat([k_rot, k_pass], dim=-1).squeeze(2)  # [B,S_local,128]
+    weights = indexer.weights_proj(
+        hidden_states.to(indexer.weights_proj.weight.dtype)
+    ) * (indexer.n_heads**-0.5)  # [B,S_local,32]
+    cp = group is not None and dist.is_initialized() and dist.get_world_size(group) > 1
+    if cp:
+        k = all_gather_seq(
+            k, group
+        )  # [B,S_global,128]; indexer is no_grad so no bwd comm needed
+        attention_mask = None  # use global position_ids for causality instead
+    return indexer_topk(
+        q,
+        k,
+        weights,
+        indexer.softmax_scale,
+        indexer.index_topk,
+        attention_mask,
+        position_ids,
+        seq_q=seq_q,
+        seq_k=seq_k,
+        q_offset=q_offset,
+    )
+
+
+def _full(t):
+    """All-gather a FSDP2 DTensor param to a plain tensor (sharded weights can't mix with plain in
+    the absorption arithmetic). No-op for plain tensors."""
+    return t.full_tensor() if type(t).__name__ == "DTensor" else t
+
+
+def _seq_idx_from_position_ids(position_ids):
+    """Per-token document id [B,S] under sample packing, or ``None`` if not packed.
+
+    Multipack resets ``position_ids`` to 0 at each packed document's start, so a non-increasing
+    step marks a boundary; ``cumsum(position_ids == 0)`` numbers the documents. Returns ``None``
+    when ``position_ids`` is absent or strictly increasing (a single unpacked sequence) so the
+    fast sparse path is preserved for ordinary long-context training."""
+    if position_ids is None:
+        return None
+    if position_ids.dim() == 1:
+        position_ids = position_ids.unsqueeze(0)
+    if (
+        position_ids.shape[1] < 2
+        or not (position_ids[:, 1:] <= position_ids[:, :-1]).any()
+    ):
+        return None
+    return (position_ids == 0).to(torch.int32).cumsum(-1)
+
+
+def effective_weight(linear) -> torch.Tensor:
+    """The weight the projection actually applies, differentiable wrt trainable params:
+    full-parameter -> ``linear.weight``; PEFT-LoRA -> ``base.weight + scaling·(Bᵀ? )`` delta so the
+    absorption flows gradients to the adapter. Falls back to ``.weight``."""
+    if hasattr(linear, "base_layer"):  # PEFT LoRA layer
+        w = _full(linear.base_layer.weight)
+        active = getattr(linear, "active_adapters", None) or list(
+            getattr(linear, "lora_A", {}).keys()
+        )
+        for name in active:
+            A = _full(linear.lora_A[name].weight)  # [r, in]
+            Bm = _full(linear.lora_B[name].weight)  # [out, r]
+            w = w + linear.scaling[name] * (Bm @ A)
+        return w
+    return _full(linear.weight)
+
+
+def _glm_dsa_attention_forward(
+    self,
+    hidden_states,
+    position_embeddings,
+    attention_mask=None,
+    past_key_values=None,
+    position_ids=None,
+    prev_topk_indices=None,
+    **kwargs,
+):
+    from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import (
+        apply_rotary_pos_emb_interleave,
+    )
+
+    group = getattr(self, "_cp_group", None)
+    cp = group is not None and dist.is_initialized() and dist.get_world_size(group) > 1
+    rank = dist.get_rank(group) if cp else 0
+
+    B, S = hidden_states.shape[:-1]  # S = local query count under CP
+    q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states))
+    q_states = self.q_b_proj(q_resid).view(B, S, -1, self.qk_head_dim).transpose(1, 2)
+    q_pass, q_rot = torch.split(
+        q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+    )
+
+    compressed = self.kv_a_proj_with_mqa(hidden_states)
+    latent, k_rot = torch.split(
+        compressed, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+    )
+    c_kv = self.kv_a_layernorm(latent)  # [B, S, kv_lora]
+    k_rot = k_rot.view(B, 1, S, self.qk_rope_head_dim)
+    cos, sin = position_embeddings
+    q_rot, k_rot = apply_rotary_pos_emb_interleave(
+        q_rot, k_rot, cos, sin
+    )  # [B,H,S,64],[B,1,S,64]
+
+    # Per-document ids under sample packing (None for ordinary sequences). seq_q is the local
+    # queries' doc ids; seq_k is the keys' doc ids (aligned with k_shared).
+    # Under CP the doc ids MUST be derived globally: cumsum(position_ids == 0) on a local chunk
+    # restarts at 0 on every rank, so a document crossing a CP boundary would get different (and
+    # colliding) ids per rank, masking valid remote keys or letting attention cross documents. Gather
+    # the global position_ids, number the documents ONCE, keep the full ids for the keys and take this
+    # rank's slice for the queries — so a boundary-spanning doc shares one id across ranks. (Also
+    # handles a chunk that is entirely mid-document, where the local cumsum would see no reset at all.)
+    if cp and position_ids is not None:
+        global_position_ids = all_gather_seq(position_ids.unsqueeze(-1), group).squeeze(
+            -1
+        )
+        seq_k = _seq_idx_from_position_ids(global_position_ids)
+        seq_q = None if seq_k is None else seq_k[:, rank * S : (rank + 1) * S]
+    else:
+        seq_q = _seq_idx_from_position_ids(position_ids)
+        seq_k = seq_q
+
+    # DSA top-k: this layer's own indexer (full) or the previous full layer's selection (shared).
+    # Under CP the indexer gathers global keys; topk indices reference GLOBAL key positions.
+    packed = seq_q is not None
+    if self.indexer is not None:
+        idx_mask = (
+            attention_mask[:, 0, :, :]
+            if (attention_mask is not None and not cp and not packed)
+            else None
+        )
+        if cp or packed or getattr(self, "_use_fused_indexer", False):
+            topk_indices = fused_indexer_topk(
+                self.indexer,
+                hidden_states,
+                q_resid,
+                position_embeddings,
+                idx_mask,
+                position_ids,
+                group=group if cp else None,
+                seq_q=seq_q,
+                seq_k=seq_k,
+                q_offset=(rank * S if cp else 0),
+            )
+        else:
+            topk_indices = self.indexer(
+                hidden_states,
+                q_resid,
+                position_embeddings,
+                idx_mask,
+                position_ids,
+                past_key_values=past_key_values,
+            )
+    else:
+        if prev_topk_indices is None:
+            raise ValueError("shared DSA layer needs prev_topk_indices")
+        topk_indices = prev_topk_indices
+
+    w_kb_k, w_kb_v = split_kv_b(effective_weight(self.kv_b_proj), self.num_heads)
+    q_abs = absorb_query(q_pass, q_rot, w_kb_k)  # [B,H,S_local,576]
+    k_shared = torch.cat([c_kv, k_rot.squeeze(1)], dim=-1)  # [B,S_local,576]
+    q_offset = 0
+    if cp:
+        k_shared = all_gather_seq(k_shared, group)  # [B,S_global,576] (differentiable)
+        q_offset = rank * S
+    out_latent = mla_attn(
+        q_abs,
+        k_shared,
+        topk_indices,
+        self.scaling,
+        q_offset=q_offset,
+        seq_q=seq_q,
+        seq_k=seq_k,
+    )
+    attn = project_value(out_latent, w_kb_v)  # [B,H,S,v_head]
+    attn = attn.transpose(1, 2).reshape(B, S, -1).contiguous()
+    return self.o_proj(attn), None, topk_indices
+
+
+def keep_router_fp32(model) -> int:
+    """Force the MoE router (gate weight + e_score_correction_bias) to fp32 STORAGE, guaranteeing the
+    router never operates below fp32 regardless of the model's compute dtype or any kernelized path.
+
+    Expert selection is discrete (sigmoid -> group top-k), so router precision matters for
+    correctness. GlmMoeDsaTopkRouter.forward already upcasts to fp32 for the logit matmul and the
+    model keeps e_score_correction_bias fp32, so on a bf16 checkpoint (whose gate weight is *already*
+    bf16-rounded at source) this changes no values — it is a guard that the integration's bf16 cast
+    can't downcast the router, and it skips the per-forward upcast. The dominant source of routing
+    variance is the bf16 RESIDUAL-STREAM router INPUT (inherent to bf16 training), not the router;
+    that is addressed model-wide, not here. Returns the router count."""
+    import torch as _t
+
+    n = 0
+    for mod in model.modules():
+        if type(mod).__name__ == "GlmMoeDsaTopkRouter":
+            mod.weight.data = mod.weight.data.to(_t.float32)
+            bias = getattr(mod, "e_score_correction_bias", None)
+            if bias is not None:
+                mod.e_score_correction_bias.data = bias.data.to(_t.float32)
+            n += 1
+    return n
+
+
+def patch_glm_moe_dsa_attention(
+    model, use_fused_indexer: bool = False, cp_group=None
+) -> int:
+    """Swap every GlmMoeDsaAttention.forward for the absorbed-kernel version. ``use_fused_indexer``
+    replaces the eager indexer scoring with the fused [B,S,T] scorer. ``cp_group`` (a torch.distributed
+    process group) enables context parallelism: the attention all-gathers the compressed KV +
+    indexer-k so each rank's local queries attend the global keys. Returns the count."""
+    import types
+
+    n = 0
+    for mod in model.modules():
+        if type(mod).__name__ == "GlmMoeDsaAttention":
+            mod._use_fused_indexer = use_fused_indexer
+            mod._cp_group = cp_group
+            mod.forward = types.MethodType(_glm_dsa_attention_forward, mod)
+            n += 1
+    return n
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/__init__.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/__init__.py
new file mode 100644
index 0000000000..25f7855f7d
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/__init__.py
@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+from . import layers
+from .lora_ops import ParallelExperts
+from .parallel_experts import flatten_sort_count, parallel_linear
+from .parallel_linear_lora import ScatterMoELoRA, parallel_linear_lora
+
+__all__ = [
+    "layers",
+    "ParallelExperts",
+    "flatten_sort_count",
+    "parallel_linear",
+    "ScatterMoELoRA",
+    "parallel_linear_lora",
+    "lora_ops",
+]
+from .multi_lora import (  # noqa: E402,F401
+    ScatterMoEMultiLoRA,
+    build_multilora_routing,
+    scatter2scatter_multilora,
+)
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/chunked_bnb.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/chunked_bnb.py
new file mode 100644
index 0000000000..e5d2d88df5
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/chunked_bnb.py
@@ -0,0 +1,158 @@
+"""Chunked-dequant grouped MoE for bnb-4bit experts (portable, any GPU).
+
+scattermoe's fused grouped path reads NVFP4/MXFP4 experts at 4-bit, but bnb-nf4 experts have no
+4-bit-read kernel, so the naive path dequantizes ALL E experts to bf16 every forward (~48 GB for
+gemma4-26B → ~2.7x the VRAM of the nvfp4 path; at 4k seq ~all experts are active so "selective"
+dequant doesn't help). This processes experts in CHUNKS: per chunk, dequant only that chunk's
+experts to bf16, run the grouped gate_up (+LoRA) → gated activation → grouped down (+LoRA) for that
+chunk's tokens under activation checkpointing (re-dequant in backward), then free. Peak bf16
+transient is bounded to ``chunk_size / E`` of the full expert set, bringing bnb-MoE memory in line
+with the nvfp4 path. Portable: ``torch._grouped_mm`` + bnb dequant, no custom CUDA.
+
+Only the bnb path routes here (gated on ``module.parametrizations`` in the experts forward); the
+NVFP4/MXFP4/bf16 paths are untouched.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+from .runtime import DEFAULT_CHUNK, RUNTIME  # noqa: E402
+from .selective_dequant import selective_expert_weights
+
+
+def set_chunk_size_override(n) -> None:
+    RUNTIME.dequant_chunk_size = int(n) if n else None
+
+
+def set_layer_gc_active(flag) -> None:
+    RUNTIME.layer_gc_active = bool(flag)
+
+
+def set_bnb_fast(flag) -> None:
+    # None (config unset) keeps the fast default; only an explicit False forces chunked.
+    RUNTIME.bnb_fast = True if flag is None else bool(flag)
+
+
+def bnb_fast_enabled() -> bool:
+    return RUNTIME.bnb_fast
+
+
+def _plan_chunking(num_experts):
+    """(chunk_size, use_checkpoint) from the balanced default or the cfg.moe_dequant_chunk_size
+    override. Lower it for large-expert MoEs on small GPUs; raise it for max throughput. Per-chunk
+    checkpointing is used ONLY when layer GC is off (with layer GC on, the chunk loop already bounds
+    the transient and an extra checkpoint would just nest a redundant recompute)."""
+    override = RUNTIME.dequant_chunk_size
+    chunk = override if override is not None else DEFAULT_CHUNK
+    chunk = max(1, min(num_experts, chunk))
+    use_ckpt = (chunk < num_experts) and not RUNTIME.layer_gc_active
+    return chunk, use_ckpt
+
+
+def _gated_act(gate, up, act_type, limit):
+    if limit is not None:  # DSV4-style clamped SwiGLU
+        gate = gate.clamp(max=limit)
+        up = up.clamp(min=-limit, max=limit)
+    if act_type == "gelu_tanh":  # Gemma4 GeGLU
+        return F.gelu(gate, approximate="tanh") * up
+    return F.silu(gate) * up
+
+
+def _chunk_forward(
+    x, experts, chunk_idx, coff, gA, gB, gs, dA, dB, ds, act_type, limit
+):
+    """One expert-chunk end-to-end: dequant chunk → grouped gate_up (+LoRA) → act → grouped down
+    (+LoRA). Run under ``checkpoint`` so the chunk's bf16 weights live only during fwd/recompute."""
+    gub = selective_expert_weights(experts, "gate_up_proj", chunk_idx).transpose(1, 2)
+    dnb = selective_expert_weights(experts, "down_proj", chunk_idx).transpose(1, 2)
+
+    gu = torch._grouped_mm(x, gub, offs=coff)
+    if gA is not None:
+        xa = torch._grouped_mm(x, gA.transpose(1, 2), offs=coff)
+        gu = gu + gs * torch._grouped_mm(xa, gB.transpose(1, 2), offs=coff)
+
+    gate, up = gu.chunk(2, dim=-1)
+    h = _gated_act(gate, up, act_type, limit).contiguous()
+
+    dn = torch._grouped_mm(h, dnb, offs=coff)
+    if dA is not None:
+        ha = torch._grouped_mm(h, dA.transpose(1, 2), offs=coff)
+        dn = dn + ds * torch._grouped_mm(ha, dB.transpose(1, 2), offs=coff)
+    return dn
+
+
+def chunked_bnb_moe(
+    hidden,
+    idx,
+    wts,
+    experts,
+    gup_lora,
+    down_lora,
+    num_experts,
+    act_type="silu",
+    limit=None,
+):
+    """Grouped MoE over bnb-4bit experts via chunked dequant. ``hidden`` [N,H]; ``idx``/``wts``
+    [N,topk]; ``experts`` is the bnb-quantized experts module; ``*_lora`` are scattermoe-layout
+    (A[r*E,K], B[out,r*E], scaling) or None. Returns [N,H], differentiable to hidden + LoRA A/B."""
+    from .grouped_train import _lora_stack
+
+    N, H = hidden.shape
+    dev = hidden.device
+    chunk_size, use_ckpt = _plan_chunking(num_experts)
+
+    Agu = Bgu = sgu = Adn = Bdn = sdn = None
+    if gup_lora is not None and down_lora is not None:
+        two_i = gup_lora[1].shape[0]  # gate_up out = 2I (B is [2I, r*E])
+        inter = down_lora[0].shape[1]  # down in = I (A is [r*E, I])
+        Agu, Bgu, sgu = _lora_stack(gup_lora, num_experts, H, two_i)
+        Adn, Bdn, sdn = _lora_stack(down_lora, num_experts, inter, H)
+
+    # Sort tokens by expert so each expert's tokens are contiguous for grouped_mm.
+    flat_exp = idx.reshape(-1)
+    order = flat_exp.argsort()
+    rep = torch.arange(N, device=dev).repeat_interleave(idx.size(1))[order]
+    wflat = wts.reshape(-1)[order]
+    counts = torch.bincount(flat_exp, minlength=num_experts)
+    tok_off = torch.cat([counts.new_zeros(1), counts.cumsum(0)]).tolist()
+    flat = hidden[rep]
+
+    out = hidden.new_zeros(N, H)
+    pieces = []
+    for c0 in range(0, num_experts, chunk_size):
+        c1 = min(c0 + chunk_size, num_experts)
+        t0, t1 = int(tok_off[c0]), int(tok_off[c1])
+        if t1 == t0:  # no tokens routed here
+            continue
+        chunk_idx = torch.arange(c0, c1, device=dev)
+        coff = counts[c0:c1].cumsum(0).to(torch.int32)
+        gA = Agu[c0:c1] if Agu is not None else None
+        gB = Bgu[c0:c1] if Bgu is not None else None
+        dA = Adn[c0:c1] if Adn is not None else None
+        dB = Bdn[c0:c1] if Bdn is not None else None
+        args = (
+            flat[t0:t1],
+            experts,
+            chunk_idx,
+            coff,
+            gA,
+            gB,
+            sgu,
+            dA,
+            dB,
+            sdn,
+            act_type,
+            limit,
+        )
+        o = (
+            checkpoint(_chunk_forward, *args, use_reentrant=False)
+            if use_ckpt
+            else _chunk_forward(*args)
+        )
+        pieces.append(o * wflat[t0:t1, None].to(o.dtype))
+    if not pieces:
+        return out
+    return out.index_add(0, rep, torch.cat(pieces, 0))
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/__init__.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/__init__.py
new file mode 100644
index 0000000000..9636a605a2
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/__init__.py
@@ -0,0 +1,31 @@
+"""sm120 CUTLASS NVFP4 grouped-GEMM MoE path (CuTe DSL).
+
+Optional: requires an sm120 (consumer Blackwell) GPU and nvidia-cutlass-dsl>=4.6 (the sm120
+block-scaled SF helpers, partition_fragment_SFA etc., landed in the 4.6 dev line). On any other
+environment `cutlass_fp4_available()` returns False and the dsv4 MoE forward falls back to the
+DeepGEMM (sm90/sm100) or chunked-dequant path. Lazy imports keep non-Blackwell / no-cutlass-dsl
+environments (incl. CI) clean.
+"""
+
+from __future__ import annotations
+
+import functools
+
+
+@functools.lru_cache(maxsize=1)
+def cutlass_fp4_available() -> bool:
+    """True iff the sm120 CUTLASS NVFP4 grouped path can run here."""
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return False
+        if torch.cuda.get_device_capability()[0] != 12:  # sm120 consumer Blackwell
+            return False
+        import cutlass.cute.nvgpu.warp.mma as _wm  # noqa: F401
+        import cutlass.utils.blackwell_helpers as _bh
+
+        # sm120 block-scaled SF helpers landed in 4.6; absent in 4.5.x
+        return hasattr(_bh, "partition_fragment_SFA") and hasattr(_wm, "MmaMXF4NVF4Op")
+    except Exception:
+        return False
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/blockscaled_gemm_dispatch.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/blockscaled_gemm_dispatch.py
new file mode 100644
index 0000000000..0ec5bea524
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/blockscaled_gemm_dispatch.py
@@ -0,0 +1,271 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""SM120 block-scaled GEMM dispatch helpers shared by the cooperative and
+pingpong examples in this directory. Both examples are otherwise standalone;
+they only depend on this file (and the rest of the cutlass DSL library) and
+not on each other.
+
+Dispatch table:
+    Same-dtype paths:
+        (FP4, FP4, *, Float8E4M3FN, 16) -> MmaMXF4NVF4Op,  use_mxf8f6f4=False, mma_K=64
+        (FP4, FP4, *, Float8E8M0FNU, 32) -> MmaMXF4Op,    use_mxf8f6f4=False, mma_K=64
+        (FP8, FP8, *, Float8E8M0FNU, 32) -> MmaMXF8Op,    use_mxf8f6f4=True,  mma_K=32
+            (FP8 same-dtype is restricted to a_dtype == b_dtype.)
+    Mixed-precision:
+        (FP4, FP8, *, Float8E8M0FNU, 32) -> MmaMXF8F6F4Op, use_mxf8f6f4=True, mma_K=32
+        (FP8, FP4, *, Float8E8M0FNU, 32) -> MmaMXF8F6F4Op, use_mxf8f6f4=True, mma_K=32
+            (Same-width mixed-FP8 (E4M3 + E5M2) and FP6 mixed pairs are not supported.)
+
+`use_mxf8f6f4` is the third argument to
+`cutlass.utils.blackwell_helpers.get_permutation_mnk` and selects perm_k = 32
+(FP8 / mixed warp-MMA shape (16,8,32)) instead of perm_k = 64 (FP4 shape (16,8,64)).
+"""
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.nvgpu.warp.mma import MXF8F6F4_SUPPORTED_PAIRS
+
+# ldsm.b4x16_p64 puts the FP4 nibble in the low half of the byte; mma.sync.kind::mxf8f6f4
+# reads it from the middle, so shift left by 2 (cute::fp4_shift_A/B in mma_traits_sm120.hpp).
+FP4_SHIFT_BITS = 2
+
+_FP8_DTYPES = (cutlass.Float8E4M3FN, cutlass.Float8E5M2)
+
+
+def make_ldmatrix_atom(operand_dtype, transpose, num_matrices=4, mixed_mode=False):
+    """Build a warp-level ldmatrix copy atom for the SM120 block-scaled GEMM
+    A/B SMEM->RMEM path.
+    """
+    if mixed_mode and operand_dtype.width == 4:
+        assert not transpose, "LdMatrix8x16x8bOp does not support transpose"
+        return cute.make_copy_atom(
+            cute.nvgpu.warp.LdMatrix8x16x8bOp(
+                transpose=False,
+                num_matrices=num_matrices,
+                unpack_bits=4,
+            ),
+            cutlass.Int8,
+        )
+    atom_element_type = operand_dtype
+    if mixed_mode and operand_dtype.width < 8:
+        atom_element_type = cutlass.Int8
+    return cute.make_copy_atom(
+        cute.nvgpu.warp.LdMatrix8x8x16bOp(
+            transpose=transpose,
+            num_matrices=num_matrices,
+        ),
+        atom_element_type,
+    )
+
+
+def make_sm120_blockscaled_mma_op(a_dtype, b_dtype, acc_dtype, sf_dtype, sf_vec_size):
+    """Dispatch the SM120 block-scaled MMA op for the given operand / scale
+    factor combination.
+
+    Returns
+    -------
+    Tuple[MmaOp, bool]
+        (mma_op, use_mxf8f6f4) where use_mxf8f6f4 is True for FP8 / mixed
+        paths (atom shape (16,8,32)) and False for FP4 same-dtype paths
+        (atom shape (16,8,64)). `use_mxf8f6f4` is consumed by
+        `cutlass.utils.blackwell_helpers.get_permutation_mnk`.
+
+    Raises
+    ------
+    ValueError
+        On any unsupported combination, with a diagnostic that names the
+        offending combination and the supported alternatives.
+    """
+    if a_dtype == b_dtype:
+        if a_dtype == cutlass.Float4E2M1FN:
+            if sf_vec_size == 16:
+                if sf_dtype != cutlass.Float8E4M3FN:
+                    raise ValueError(
+                        f"Float4E2M1FN + sf_vec_size=16 requires sf_dtype=Float8E4M3FN, "
+                        f"got sf_dtype={sf_dtype}"
+                    )
+                return (
+                    cute.nvgpu.warp.MmaMXF4NVF4Op(a_dtype, acc_dtype, sf_dtype),
+                    False,
+                )
+            if sf_vec_size == 32:
+                if sf_dtype != cutlass.Float8E8M0FNU:
+                    raise ValueError(
+                        f"Float4E2M1FN + sf_vec_size=32 requires sf_dtype=Float8E8M0FNU, "
+                        f"got sf_dtype={sf_dtype}"
+                    )
+                return (
+                    cute.nvgpu.warp.MmaMXF4Op(a_dtype, acc_dtype, sf_dtype),
+                    False,
+                )
+            raise ValueError(
+                f"Float4E2M1FN requires sf_vec_size in (16, 32), got {sf_vec_size}"
+            )
+        if a_dtype in _FP8_DTYPES:
+            if sf_vec_size != 32:
+                raise ValueError(
+                    f"FP8 ab_dtype ({a_dtype}) requires sf_vec_size=32, "
+                    f"got {sf_vec_size}"
+                )
+            if sf_dtype != cutlass.Float8E8M0FNU:
+                raise ValueError(
+                    f"FP8 ab_dtype + sf_vec_size=32 requires sf_dtype=Float8E8M0FNU, "
+                    f"got sf_dtype={sf_dtype}"
+                )
+            return (
+                cute.nvgpu.warp.MmaMXF8Op(a_dtype, acc_dtype, sf_dtype),
+                True,
+            )
+        raise ValueError(
+            f"Unsupported same-dtype ab_dtype={a_dtype} for SM120 block-scaled GEMM. "
+            f"Supported same-dtype: Float4E2M1FN, Float8E4M3FN, Float8E5M2."
+        )
+    is_a_fp4_b_fp8 = a_dtype == cutlass.Float4E2M1FN and b_dtype in _FP8_DTYPES
+    is_a_fp8_b_fp4 = a_dtype in _FP8_DTYPES and b_dtype == cutlass.Float4E2M1FN
+    if is_a_fp4_b_fp8 or is_a_fp8_b_fp4:
+        if sf_vec_size != 32:
+            raise ValueError(
+                f"FP4 x FP8 mixed-precision requires sf_vec_size=32, got {sf_vec_size}"
+            )
+        if sf_dtype != cutlass.Float8E8M0FNU:
+            raise ValueError(
+                f"FP4 x FP8 mixed-precision requires sf_dtype=Float8E8M0FNU, "
+                f"got sf_dtype={sf_dtype}"
+            )
+        return (
+            cute.nvgpu.warp.MmaMXF8F6F4Op(a_dtype, b_dtype, acc_dtype, sf_dtype),
+            True,
+        )
+    if a_dtype in _FP8_DTYPES and b_dtype in _FP8_DTYPES:
+        raise ValueError(
+            f"same-width mixed-FP8 (a_dtype={a_dtype}, b_dtype={b_dtype}) is not supported. "
+            f"Supported FP4 x FP8 mixed pairs: (Float4E2M1FN, Float8E4M3FN), "
+            f"(Float4E2M1FN, Float8E5M2), (Float8E4M3FN, Float4E2M1FN), "
+            f"(Float8E5M2, Float4E2M1FN). Same-dtype FP8 uses MmaMXF8Op."
+        )
+    raise ValueError(
+        f"Unsupported (a_dtype={a_dtype}, b_dtype={b_dtype}, sf_dtype={sf_dtype}, "
+        f"sf_vec_size={sf_vec_size}) for SM120 block-scaled GEMM. FP6 mixed pairs "
+        f"are not supported; supported mixed pairs are FP4 x FP8 only."
+    )
+
+
+def validate_blockscaled_args(args, fp4_allowed_tiles, fp8_allowed_tiles):
+    """Post-argparse validation of (a_dtype, b_dtype, sf_dtype, sf_vec_size, tile_shape_mnk).
+
+    Same-dtype paths use the FP4 / FP8 same-dtype atoms. Mixed-precision FP4 x FP8
+    is permitted only for the four supported pairs. Same-width mixed-FP8 and FP6
+    are explicitly rejected with named diagnostics.
+
+    Tile-K constraints come from the BlockScaled SF SMEM layout
+    (`sm120_make_smem_layout_sfa`), which requires
+    ``tile_K >= sf_vec_size * blk_sf == sf_vec_size * 4``:
+      * sf_vec_size=16 (NVFP4): tile_K must be a multiple of 64
+      * sf_vec_size=32 (MXFP4 / MXFP8 / mixed): tile_K must be a multiple of 128
+    A K=64 SF block cannot be filled at sf_vec_size=32 (only 2 SFs along K
+    fit in the K=128-required basic chunk), so tile_K=64 is rejected for
+    sf_vec_size=32 even though the FP4 same-dtype path otherwise allows it.
+    """
+    tile = tuple(args.tile_shape_mnk)
+    a_dtype = args.a_dtype
+    b_dtype = args.b_dtype
+    if args.sf_vec_size not in (16, 32):
+        raise ValueError(f"--sf_vec_size must be 16 or 32, got {args.sf_vec_size}")
+    if a_dtype != b_dtype:
+        if (a_dtype, b_dtype) not in MXF8F6F4_SUPPORTED_PAIRS:
+            if a_dtype in _FP8_DTYPES and b_dtype in _FP8_DTYPES:
+                raise ValueError(
+                    f"same-width mixed-FP8 (--a_dtype {a_dtype} --b_dtype {b_dtype}) "
+                    f"is not supported. Supported mixed pairs: FP4 x FP8 only."
+                )
+            raise ValueError(
+                f"unsupported mixed (--a_dtype {a_dtype} --b_dtype {b_dtype}). "
+                f"Supported mixed pairs are FP4 x FP8 only: "
+                f"{sorted(repr(p) for p in MXF8F6F4_SUPPORTED_PAIRS)}. "
+                f"FP6 mixed pairs are not supported."
+            )
+        if args.sf_vec_size != 32:
+            raise ValueError(
+                f"FP4 x FP8 mixed-precision requires --sf_vec_size 32, "
+                f"got {args.sf_vec_size}"
+            )
+        if args.sf_dtype != cutlass.Float8E8M0FNU:
+            raise ValueError(
+                f"FP4 x FP8 mixed-precision requires --sf_dtype Float8E8M0FNU, "
+                f"got --sf_dtype {args.sf_dtype}"
+            )
+        if tile not in fp8_allowed_tiles:
+            raise ValueError(
+                f"tile_shape {tile} is not supported for FP4 x FP8 mixed-precision. "
+                f"Allowed mixed tile shapes: {sorted(fp8_allowed_tiles)}."
+            )
+        return
+    if a_dtype in _FP8_DTYPES:
+        if args.sf_vec_size != 32:
+            raise ValueError(
+                f"FP8 a_dtype ({a_dtype}) requires --sf_vec_size 32, "
+                f"got {args.sf_vec_size}"
+            )
+        if args.sf_dtype != cutlass.Float8E8M0FNU:
+            raise ValueError(
+                f"FP8 a_dtype + sf_vec_size=32 requires --sf_dtype Float8E8M0FNU, "
+                f"got --sf_dtype {args.sf_dtype}"
+            )
+        if tile not in fp8_allowed_tiles:
+            raise ValueError(
+                f"tile_shape {tile} is not supported for MXFP8. "
+                f"Allowed FP8 tile shapes: {sorted(fp8_allowed_tiles)}."
+            )
+    elif a_dtype == cutlass.Float4E2M1FN:
+        if args.sf_vec_size == 16 and args.sf_dtype != cutlass.Float8E4M3FN:
+            raise ValueError(
+                f"FP4 + --sf_vec_size 16 requires --sf_dtype Float8E4M3FN, "
+                f"got {args.sf_dtype}"
+            )
+        if args.sf_vec_size == 32 and args.sf_dtype != cutlass.Float8E8M0FNU:
+            raise ValueError(
+                f"FP4 + --sf_vec_size 32 requires --sf_dtype Float8E8M0FNU, "
+                f"got {args.sf_dtype}"
+            )
+        if tile not in fp4_allowed_tiles:
+            raise ValueError(
+                f"tile_shape {tile} is not supported for FP4 path. "
+                f"Allowed FP4 tile shapes: {sorted(fp4_allowed_tiles)}."
+            )
+        if args.sf_vec_size == 32 and tile[2] % 128 != 0:
+            raise ValueError(
+                f"FP4 + sf_vec_size=32 (MXFP4) requires tile_K to be a "
+                f"multiple of 128, "
+                f"got tile_K={tile[2]}."
+            )
+    else:
+        raise ValueError(
+            f"--a_dtype must be Float4E2M1FN, Float8E4M3FN, or Float8E5M2; "
+            f"got {a_dtype}"
+        )
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/fused_dx.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/fused_dx.py
new file mode 100644
index 0000000000..ff99d57939
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/fused_dx.py
@@ -0,0 +1,121 @@
+"""Fused in-kernel-dequant grouped dX Triton kernel (#23): dx[Mt,K] = A[Mt,N] @ Wdeq[expert(m),N,K]
+with bf16 grad A + NVFP4 weight decoded IN-KERNEL (no bf16 weight materialization, vectorized, no
+Python loop). Contiguous-grouped: each 128-row M-tile belongs to one expert via m_indices.
+This is the memory-optimal backward dX (vs dequant+cuBLAS which materializes bf16 W)."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+def _codebook(dev):
+    return torch.tensor(
+        [
+            0.0,
+            0.5,
+            1.0,
+            1.5,
+            2.0,
+            3.0,
+            4.0,
+            6.0,
+            0.0,
+            -0.5,
+            -1.0,
+            -1.5,
+            -2.0,
+            -3.0,
+            -4.0,
+            -6.0,
+        ],
+        device=dev,
+        dtype=torch.float32,
+    )
+
+
+@triton.jit
+def _fused_dx(
+    A,
+    Wq,
+    Ws,
+    MI,
+    PT,
+    OUT,
+    CB,
+    N,
+    K,
+    sa0,
+    sa1,
+    sq0,
+    sq1,
+    sq2,
+    ss0,
+    ss1,
+    ss2,
+    so0,
+    so1,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+):
+    pid_m = tl.program_id(0)  # contiguous-grouped: each 128-row M-tile is one expert
+    pid_k = tl.program_id(1)
+    e = tl.load(MI + pid_m)
+    pt_e = tl.load(PT + e).to(tl.float32)
+    m = pid_m * 128 + tl.arange(0, 128)
+    k = pid_k * BK + tl.arange(0, BK)
+    km = k < K
+    acc = tl.zeros((128, BK), tl.float32)
+    for n0 in range(0, N, BN):
+        n = n0 + tl.arange(0, BN)
+        a = tl.load(A + m[:, None] * sa0 + n[None, :] * sa1)
+        wq = tl.load(
+            Wq + e * sq0 + n[:, None] * sq1 + (k[None, :] // 2) * sq2,
+            mask=km[None, :],
+            other=0,
+        ).to(tl.int32)
+        nib = tl.where(k[None, :] % 2 == 1, (wq >> 4) & 0xF, wq & 0xF)
+        ws = tl.load(
+            Ws + e * ss0 + n[:, None] * ss1 + (k[None, :] // 16) * ss2,
+            mask=km[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        w = tl.load(CB + nib) * ws * pt_e
+        acc += tl.dot(a, w.to(tl.bfloat16))
+    tl.store(
+        OUT + m[:, None] * so0 + k[None, :] * so1, acc.to(tl.bfloat16), mask=km[None, :]
+    )
+
+
+def fused_dx(A, Wq, Ws, m_indices, pt, K, BN=64, BK=64):
+    """A[Mt,N] bf16, Wq[E,N,K/2] u8, Ws[E,N,K/16] e4m3, pt[E] -> dx[Mt,K] bf16."""
+    Mt, N = A.shape
+    # Mt//128 grid covers every row only if 128-padded; otherwise trailing rows stay uninitialized.
+    assert Mt % 128 == 0, (
+        f"fused_dx expects 128-padded Mt (contiguous-grouped), got Mt={Mt}"
+    )
+    out = torch.empty(Mt, K, device=A.device, dtype=torch.bfloat16)
+    grid = (Mt // 128, triton.cdiv(K, BK))
+    _fused_dx[grid](
+        A,
+        Wq,
+        Ws,
+        m_indices,
+        pt.contiguous(),
+        out,
+        _codebook(A.device),
+        N,
+        K,
+        A.stride(0),
+        A.stride(1),
+        Wq.stride(0),
+        Wq.stride(1),
+        Wq.stride(2),
+        Ws.stride(0),
+        Ws.stride(1),
+        Ws.stride(2),
+        out.stride(0),
+        out.stride(1),
+        BN=BN,
+        BK=BK,
+    )
+    return out
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/grouped.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/grouped.py
new file mode 100644
index 0000000000..55e7542c72
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/grouped.py
@@ -0,0 +1,195 @@
+"""sm120 CUTLASS NVFP4 grouped-GEMM host wrapper.
+
+Contiguous-grouped MoE: tokens sorted by expert + padded per-expert to TILE(128); ONE grouped
+GEMM where each m-tile reads its expert's weight via ``m_indices`` (grouped_kernel.py). Base-only
+fp4xfp4 (nvfp4) or fp8xfp4 (mxfp). LoRA + swiglu + routing live in the unified dispatcher; this
+file is the base GEMM engine. All cutlass/cute imports are lazy (module loads clean off sm120).
+"""
+
+from __future__ import annotations
+
+import torch
+
+TILE = 128
+
+# mode -> (a-dtype-name, sf_vec, sf-dtype-name)
+_MODES = {
+    "nvfp4": ("Float4E2M1FN", 16, "Float8E4M3FN"),
+    "fp8": ("Float8E4M3FN", 32, "Float8E8M0FNU"),
+}
+
+
+def _cd(a, b):
+    return (a + b - 1) // b
+
+
+def _helpers():
+    """Lazy cute helpers (operand/SF build + inject). Imported only when the sm120 path runs."""
+    import cutlass
+    import cutlass.cute as cute
+    import cutlass.torch as cutlass_torch
+    import cutlass.utils.blackwell_helpers as sm120_utils
+    from cutlass.cute.runtime import from_dlpack
+
+    from .grouped_kernel import cvt_sf_MKL_to_M32x4xrm_K4xrk_L  # noqa: F401
+
+    return (
+        cutlass,
+        cute,
+        cutlass_torch,
+        from_dlpack,
+        sm120_utils,
+        cvt_sf_MKL_to_M32x4xrm_K4xrk_L,
+    )
+
+
+class GroupedFp4Gemm:
+    """C[Mt,N] = sorted A[Mt,K] @ B[m_indices[m_tile]]^T. One launch; expert per m-tile.
+
+    Build once per (Mt, N, K, E, mode); inject weights once (set_weights), activation per-forward.
+    """
+
+    def __init__(self, Mt, N, K, E, mode="nvfp4"):
+        cutlass, cute, cutlass_torch, from_dlpack, sm120_utils, _ = _helpers()
+        from cutlass import BFloat16, Float4E2M1FN, Float32
+
+        from .grouped_kernel import Sm120BlockScaledGemmKernel
+
+        a_name, sf_vec, sf_name = _MODES[mode]
+        a_dtype = getattr(cutlass, a_name)
+        sf_dtype = getattr(cutlass, sf_name)
+        self.Mt, self.N, self.K, self.E, self.mode = Mt, N, K, E, mode
+        self.sf_vec, self._fp4_a = sf_vec, a_dtype is Float4E2M1FN
+
+        self.a_ct, self.a_st = self._operand(Mt, K, 1, a_dtype)
+        self.b_ct, self.b_st = self._operand(N, K, E, Float4E2M1FN)
+        self.sfa_ct, self.sfa_st, self.sfa_g, self.sfa_shape = self._sf(
+            Mt, K, 1, sf_vec, sf_dtype
+        )
+        self.sfb_ct, self.sfb_st, self.sfb_g, self.sfb_shape = self._sf(
+            N, K, E, sf_vec, sf_dtype
+        )
+        c_ref = cutlass_torch.matrix(1, Mt, N, False, BFloat16)
+        self.c_ct, self.c_st = cutlass_torch.cute_tensor_like(
+            c_ref, BFloat16, is_dynamic_layout=True, assumed_align=16
+        )
+        self.c_ct.mark_compact_shape_dynamic(
+            mode=1, stride_order=(2, 0, 1), divisibility=1
+        )
+        self.c_ct = cutlass_torch.convert_cute_tensor(
+            c_ref.cuda(), self.c_ct, BFloat16, is_dynamic_layout=True
+        )
+        self.mi = torch.zeros(Mt // TILE, dtype=torch.int32, device="cuda")
+        self.mi_ct = from_dlpack(self.mi, assumed_align=4)
+        gemm = Sm120BlockScaledGemmKernel(Float32, sf_vec, (128, 128, 128), (128, 128))
+        mac = cutlass.utils.HardwareInfo().get_max_active_clusters(1)
+        self.stream = cutlass_torch.default_stream()
+        self.compiled = cute.compile(
+            gemm,
+            self.a_ct,
+            self.b_ct,
+            self.sfa_ct,
+            self.sfb_ct,
+            self.c_ct,
+            self.mi_ct,
+            mac,
+            self.stream,
+        )
+
+    def _operand(self, mn, k, E, op_dtype):
+        cutlass, cute, cutlass_torch, _, _, _ = _helpers()
+        from cutlass import Float4E2M1FN, Float32
+
+        ref = cutlass_torch.matrix(E, mn, k, False, Float32)
+        ct, st = cutlass_torch.cute_tensor_like(
+            ref, op_dtype, is_dynamic_layout=True, assumed_align=16
+        )
+        ct.mark_compact_shape_dynamic(
+            mode=1,
+            stride_order=(2, 0, 1),
+            divisibility=2 if op_dtype is Float4E2M1FN else 1,
+        )
+        ct = cutlass_torch.convert_cute_tensor(
+            ref.cuda(), ct, op_dtype, is_dynamic_layout=True
+        )
+        return ct, st
+
+    def _sf(self, mn, k, E, sf_vec, sf_dtype):
+        cutlass, cute, cutlass_torch, from_dlpack, _, cvt = _helpers()
+        sfk = _cd(k, sf_vec)
+        mma_shape = (E, _cd(mn, 128), _cd(sfk, 4), 32, 4, 4)
+        nat = cutlass_torch.create_and_permute_torch_tensor(
+            (E, mn, sfk),
+            torch.float64,
+            permute_order=(1, 2, 0),
+            init_type=cutlass_torch.TensorInitType.SKIP,
+        )
+        nat.copy_(
+            torch.arange(E * mn * sfk, dtype=torch.float64)
+            .reshape(E, mn, sfk)
+            .permute(1, 2, 0)
+        )
+        cf = cutlass_torch.create_and_permute_torch_tensor(
+            mma_shape,
+            torch.float64,
+            permute_order=(3, 4, 1, 5, 2, 0),
+            init_type=cutlass_torch.TensorInitType.SKIP,
+        )
+        cvt(from_dlpack(nat), from_dlpack(cf))
+        gather = cf.reshape(-1).round().long().cuda()
+        cf_e = cutlass_torch.create_and_permute_torch_tensor(
+            mma_shape,
+            torch.float32,
+            permute_order=(3, 4, 1, 5, 2, 0),
+            init_type=cutlass_torch.TensorInitType.SKIP,
+        )
+        ct, st = cutlass_torch.cute_tensor_like(
+            cf_e, sf_dtype, is_dynamic_layout=True, assumed_align=16
+        )
+        ct = cutlass_torch.convert_cute_tensor(
+            cf_e.cuda(), ct, sf_dtype, is_dynamic_layout=True
+        )
+        return ct, st, gather, tuple(st.shape)
+
+    @staticmethod
+    def _inject_operand(st, qdata):
+        n = qdata.numel()
+        st.view(torch.uint8).permute(2, 0, 1).reshape(-1)[:n].copy_(
+            qdata.reshape(-1).view(torch.uint8)
+        )
+
+    @staticmethod
+    def _inject_sf(st, gather, shape, scale_flat):
+        st.view(torch.uint8).copy_(scale_flat.view(torch.uint8)[gather].reshape(shape))
+
+    def set_weights(self, q, s):  # q:[E,N,K/2], s:[E,N,sfk]; one-time per weight gather
+        self._inject_operand(self.b_st, q)
+        self._inject_sf(self.sfb_st, self.sfb_g, self.sfb_shape, s.reshape(-1))
+
+    def forward(
+        self, a_q, a_s, m_indices
+    ):  # a_q:[1,Mt,*], a_s:[1,Mt,sfk], m_indices:[Mt/128]
+        self._inject_operand(self.a_st, a_q)
+        self._inject_sf(self.sfa_st, self.sfa_g, self.sfa_shape, a_s.reshape(-1))
+        self.mi.copy_(m_indices)
+        self.compiled(
+            self.a_ct,
+            self.b_ct,
+            self.sfa_ct,
+            self.sfb_ct,
+            self.c_ct,
+            self.mi_ct,
+            self.stream,
+        )
+        return self.c_st[:, :, 0]
+
+
+def quant_act(x, mode):
+    """Fast Triton activation quant for the grouped path -> (qdata[Mt,*], scale[Mt,sfk])."""
+    if mode == "nvfp4":
+        from .quant_nvfp4 import nvfp4_quant
+
+        return nvfp4_quant(x)
+    from .quant_mxfp8 import mxfp8_quant
+
+    return mxfp8_quant(x)
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/grouped_kernel.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/grouped_kernel.py
new file mode 100644
index 0000000000..48365588e8
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/grouped_kernel.py
@@ -0,0 +1,1917 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import argparse
+from typing import Optional, Tuple, Type
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm120_utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+import cutlass.utils.hopper_helpers as sm90_utils
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cute.runtime import from_dlpack
+
+from .blockscaled_gemm_dispatch import (  # noqa: F401
+    FP4_SHIFT_BITS,
+    make_ldmatrix_atom,
+    make_sm120_blockscaled_mma_op,
+    validate_blockscaled_args,
+)
+
+"""
+A high-performance batched dense blockscaled GEMM (C = A*SF_A * B*SF_B) example for the NVIDIA Blackwell Geforce architecture using CUTE DSL.
+- Matrix A is MxKxL, L is batch dimension, A can only be row-major("K")
+- Matrix B is NxKxL, L is batch dimension, B can only be column-major("K")
+- Matrix C is MxNxL, L is batch dimension, C can only be row-major("N")
+- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has M×ceil_div(K, sf_vec_size)×L elements respectively
+- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has N×ceil_div(K, sf_vec_size)×L elements respectively
+- Source formats for matrices A and B: The only supported source format in this example is E2M1.
+- Source formats for matrices SF_A and SF_B are controlled separately. With sf_vec_size=32, the only supported source format is E8. With sf_vec_size=16, the supported source formats are E8 and E4M3.
+
+This GEMM kernel supports the following features:
+    - Utilizes Tensor Memory Access (TMA) for efficient memory operations
+    - Utilizes warp-level block-scaled MMA for matrix multiply-accumulate (MMA) operations
+    - Supports persistent tile scheduling to better overlap memory load/store with MMA between tiles
+    - Supports warp specialization to avoid explicit pipelining between mainloop load and MMA
+    - Uses a cooperative schedule: both warp groups work together to process a single output tile.
+      The two warp groups collaboratively execute the MMA mainloop, each computing a portion of
+      the tile, and then jointly perform the epilogue. This increases computational throughput
+      per tile by leveraging both warp groups simultaneously.
+
+This GEMM works as follows:
+1. DMA warp group:
+    - Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
+    - Load scale factor A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
+2. MMA warp groups (both warp groups cooperate on the same tile):
+    - Load A/B from shared memory (SMEM) to registers (RMEM) using ldmatrix instruction.
+    - Load scale factor A/B from shared memory (SMEM) to registers (RMEM) using universal copy.
+    - Perform matrix multiply-accumulate (MMA) operations using warp-level block-scaled MMA instruction.
+    - Store C matrix from registers (RMEM) to shared memory (SMEM), then to global memory (GMEM) with TMA operations.
+    Note: Both MMA warp groups jointly handle MMA and epilogue for the same tile.
+
+Warp-level block-scaled MMA instructions operate as follows:
+- Set matrix scale factor A/B from registers
+- Read matrix A/B from registers
+- Perform MMA operation and store the result in Accumulator(register)
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/cute/blackwell_geforce/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent_cooperative.py      \
+      --mnkl 1024,1024,1024,1 --tile_shape_mnk 128,128,128                      \
+      --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN                             \
+      --c_dtype Float16 --acc_dtype Float32                                     \
+      --sf_dtype Float8E4M3FN --sf_vec_size 16
+
+The above example command compute batched gemm with M=1024, N=1024, K=1024,
+batch_count=1. The tile shape is 128x128x128 and the cluster shape is (1,1).
+The input, mma accumulator and output data type are set as fp4, fp32
+and fp16, respectively.
+
+To collect performance with NCU profiler:
+
+.. code-block:: bash
+
+    ncu python examples/cute/blackwell_geforce/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent_cooperative.py  \
+      --mnkl 1024,1024,1024,1 --tile_shape_mnk 128,128,128                      \
+      --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN                             \
+      --c_dtype Float16 --acc_dtype Float32                                     \
+      --sf_dtype Float8E4M3FN --sf_vec_size 16
+
+Constraints:
+* Supported input data types: Float4E2M1FN
+* Only Float32 accumulation is supported in FP4 mma
+* CTA tile shape M/N/K:
+ - tile_shape_m should be divisible by 128
+ - tile_shape_n should be divisible by 128
+ - tile_shape_k should be divisible by 64 (sf_vec_size=16) or 128 (sf_vec_size=32)
+* Cluster shape M/N must be [1, 1] for Blackwell Gefore
+"""
+
+
+def parse_comma_separated_ints(s: str):
+    try:
+        return tuple([int(x.strip()) for x in s.split(",")])
+    except ValueError:
+        raise argparse.ArgumentTypeError(
+            "Invalid format. Expected comma-separated integers."
+        ) from None
+
+
+class Sm120BlockScaledGemmKernel:
+    def __init__(
+        self,
+        acc_dtype,
+        sf_vec_size,
+        tile_shape_mnk,
+        epi_tile,
+    ):
+        self.acc_dtype = acc_dtype
+        self.sf_vec_size = sf_vec_size
+        self.cluster_shape_mnk = (1, 1, 1)
+        self.tile_shape_mnk = tuple(tile_shape_mnk)
+        self.epi_tile = tuple(epi_tile)
+        self.tiled_mma = None
+
+        self.occupancy = 1
+        self.num_mma_warps = 8
+        self.tma_load_warp_id = self.num_mma_warps
+        self.num_threads_per_warp = 32
+        self.threads_per_cta = (
+            self.num_mma_warps + 1  # 1 warp for DMA
+        ) * self.num_threads_per_warp
+        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_120")
+
+        self.ab_stage = None
+        self.epi_stage = None
+
+        self.a_smem_layout_staged = None
+        self.b_smem_layout_staged = None
+        self.epi_smem_layout_staged = None
+
+        self.buffer_align_bytes = 1024
+
+        self.mma_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=self.num_mma_warps * self.num_threads_per_warp,
+        )
+        self.epilog_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=2,
+            num_threads=self.num_mma_warps * self.num_threads_per_warp,
+        )
+        self.load_register_requirement = 40
+        self.mma_register_requirement = 232
+
+    def _setup_attributes(self):
+        mma_op, use_mxf8f6f4 = make_sm120_blockscaled_mma_op(
+            self.a_dtype,
+            self.b_dtype,
+            self.acc_dtype,
+            self.sf_dtype,
+            self.sf_vec_size,
+        )
+        self.mixed_mode = self.a_dtype != self.b_dtype
+        # a_fp4_in_mixed / b_fp4_in_mixed: this side carries an FP4 operand in the
+        # mixed FP4 x FP8 mode, so SMEM/TMA see Int8 storage and the mma.sync
+        # consumer needs the LDSM b4x16_p64 unpack + register `<< FP4_SHIFT_BITS`.
+        a_fp4_in_mixed = self.mixed_mode and self.a_dtype.width < 8
+        b_fp4_in_mixed = self.mixed_mode and self.b_dtype.width < 8
+        self.smem_alloc_a_dtype = cutlass.Int8 if a_fp4_in_mixed else self.a_dtype
+        self.smem_alloc_b_dtype = cutlass.Int8 if b_fp4_in_mixed else self.b_dtype
+        # `internal_type` for `_make_tma_atoms_and_tensors`: None when the dtype
+        # already matches (TMA sees the native dtype), Int8 when we recast for FP4.
+        self.tma_internal_a_dtype = cutlass.Int8 if a_fp4_in_mixed else None
+        self.tma_internal_b_dtype = cutlass.Int8 if b_fp4_in_mixed else None
+        atom_shape = (4, 2, 1)
+        atom_layout = cute.make_layout(atom_shape)
+        permutation_mnk = sm120_utils.get_permutation_mnk(
+            self.tile_shape_mnk, self.sf_vec_size, use_mxf8f6f4
+        )
+        self.tiled_mma = cute.make_tiled_mma(
+            mma_op,
+            atom_layout,
+            permutation_mnk=permutation_mnk,
+        )
+
+        self.cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk)
+
+        sfa_smem_layout_per_stage = blockscaled_utils.sm120_make_smem_layout_sfa(
+            self.tiled_mma,
+            self.tile_shape_mnk,
+            self.sf_vec_size,
+            1,
+        )
+
+        sfb_smem_layout_per_stage = blockscaled_utils.sm120_make_smem_layout_sfb(
+            self.tiled_mma,
+            self.tile_shape_mnk,
+            self.sf_vec_size,
+            1,
+        )
+
+        self.ab_stage, self.epi_stage = self._compute_stages(
+            self.tile_shape_mnk,
+            self.smem_alloc_a_dtype,
+            self.smem_alloc_b_dtype,
+            self.sf_dtype,
+            sfa_smem_layout_per_stage,
+            sfb_smem_layout_per_stage,
+            self.epi_tile,
+            self.c_dtype,
+            self.smem_capacity,
+            self.occupancy,
+        )
+
+        assert self.epi_stage > 0, (
+            "epi_stage <= 0, no enough shared memory. This case will be skipped."
+        )
+
+        (
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.sfa_smem_layout_staged,
+            self.sfb_smem_layout_staged,
+            self.epi_smem_layout_staged,
+        ) = self._make_smem_layouts(
+            self.tile_shape_mnk,
+            self.epi_tile,
+            self.smem_alloc_a_dtype,
+            self.a_layout,
+            self.smem_alloc_b_dtype,
+            self.b_layout,
+            self.ab_stage,
+            self.c_dtype,
+            self.c_layout,
+            self.epi_stage,
+            self.sf_vec_size,
+            self.tiled_mma,
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        a: cute.Tensor,
+        b: cute.Tensor,
+        sfa: cute.Tensor,
+        sfb: cute.Tensor,
+        c: cute.Tensor,
+        m_indices: cute.Tensor,
+        max_active_clusters: cutlass.Constexpr,
+        stream: cuda.CUstream,
+    ):
+        """Execute the GEMM operation in steps:
+        - Setup static attributes
+        - Setup TMA load/store atoms and tensors
+        - Compute grid size
+        - Define shared storage for kernel
+        - Launch the kernel synchronously
+
+        :param a: Input tensor A
+        :type a: cute.Tensor
+        :param b: Input tensor B
+        :type b: cute.Tensor
+        :param c: Output tensor C
+        :type c: cute.Tensor
+        :param stream: CUDA stream for asynchronous execution
+        :type stream: cuda.CUstream
+        """
+
+        # setup static attributes before smem/grid/tma computation
+        self.a_dtype = a.element_type
+        self.b_dtype = b.element_type
+        self.c_dtype = c.element_type
+        self.sf_dtype = sfa.element_type
+
+        self.a_layout = utils.LayoutEnum.from_tensor(a)
+        self.b_layout = utils.LayoutEnum.from_tensor(b)
+        self.c_layout = utils.LayoutEnum.from_tensor(c)
+
+        self._setup_attributes()
+
+        # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
+        self.sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
+            a.shape, self.sf_vec_size
+        )
+        sfa_tensor = cute.make_tensor(sfa.iterator, self.sfa_layout)
+        # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
+        self.sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
+            b.shape, self.sf_vec_size
+        )
+        sfb_tensor = cute.make_tensor(sfb.iterator, self.sfb_layout)
+
+        tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors(
+            a,
+            self.a_smem_layout_staged,
+            (self.tile_shape_mnk[0], self.tile_shape_mnk[2]),
+            1,
+            internal_type=self.tma_internal_a_dtype,
+        )
+
+        tma_atom_b, tma_tensor_b = self._make_tma_atoms_and_tensors(
+            b,
+            self.b_smem_layout_staged,
+            (self.tile_shape_mnk[1], self.tile_shape_mnk[2]),
+            1,
+            internal_type=self.tma_internal_b_dtype,
+        )
+
+        tma_atom_sfa, tma_tensor_sfa = self._make_tma_atoms_and_tensors(
+            sfa_tensor,
+            self.sfa_smem_layout_staged,
+            (self.tile_shape_mnk[0], self.tile_shape_mnk[2]),
+            1,
+            internal_type=cutlass.Int16,
+        )
+
+        tma_atom_sfb, tma_tensor_sfb = self._make_tma_atoms_and_tensors(
+            sfb_tensor,
+            self.sfb_smem_layout_staged,
+            (self.tile_shape_mnk[1], self.tile_shape_mnk[2]),
+            1,
+            internal_type=cutlass.Int16,
+        )
+
+        tma_atom_c, tma_tensor_c = self._make_tma_store_atoms_and_tensors(
+            c,
+            self.epi_smem_layout_staged,
+            self.epi_tile,
+        )
+
+        tile_sched_params, grid = self._compute_grid(
+            c,
+            self.tile_shape_mnk,
+            max_active_clusters,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            mainloop_pipeline_array_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.ab_stage * 2
+            ]
+            sA: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.smem_alloc_a_dtype, cute.cosize(self.a_smem_layout_staged)
+                ],
+                self.buffer_align_bytes,
+            ]
+            sB: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.smem_alloc_b_dtype, cute.cosize(self.b_smem_layout_staged)
+                ],
+                self.buffer_align_bytes,
+            ]
+            sSFA: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
+                ],
+                self.buffer_align_bytes,
+            ]
+            sSFB: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
+                ],
+                self.buffer_align_bytes,
+            ]
+            sC: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.c_dtype, cute.cosize(self.epi_smem_layout_staged)
+                ],
+                self.buffer_align_bytes,
+            ]
+
+        self.shared_storage = SharedStorage
+
+        self.kernel(
+            tma_atom_a,
+            tma_tensor_a,
+            tma_atom_b,
+            tma_tensor_b,
+            tma_atom_sfa,
+            tma_tensor_sfa,
+            tma_atom_sfb,
+            tma_tensor_sfb,
+            tma_atom_c,
+            tma_tensor_c,
+            self.tiled_mma,
+            self.cta_layout_mnk,
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.sfa_smem_layout_staged,
+            self.sfb_smem_layout_staged,
+            self.epi_smem_layout_staged,
+            m_indices,
+            tile_sched_params,
+        ).launch(
+            grid=grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=[1, 1, 1],
+            stream=stream,
+        )
+        return
+
+    @cute.kernel
+    def kernel(
+        self,
+        tma_atom_a: cute.CopyAtom,
+        mA_mkl: cute.Tensor,
+        tma_atom_b: cute.CopyAtom,
+        mB_nkl: cute.Tensor,
+        tma_atom_sfa: cute.CopyAtom,
+        mSFA_mkl: cute.Tensor,
+        tma_atom_sfb: cute.CopyAtom,
+        mSFB_nkl: cute.Tensor,
+        tma_atom_c: cute.CopyAtom,
+        mC_mnl: cute.Tensor,
+        tiled_mma: cute.TiledMma,
+        cta_layout_mnk: cute.Layout,
+        a_smem_layout_staged: cute.ComposedLayout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        sfa_smem_layout_staged: cute.Layout,
+        sfb_smem_layout_staged: cute.Layout,
+        epi_smem_layout_staged: cute.ComposedLayout,
+        m_indices: cute.Tensor,
+        tile_sched_params: utils.PersistentTileSchedulerParams,
+    ):
+        """
+        GPU device kernel performing the batched GEMM computation.
+
+        :param tma_atom_a: TMA copy atom for A tensor
+        :type tma_atom_a: cute.CopyAtom
+        :param mA_mkl: Input tensor A
+        :type mA_mkl: cute.Tensor
+        :param tma_atom_b: TMA copy atom for B tensor
+        :type tma_atom_b: cute.CopyAtom
+        :param mB_nkl: Input tensor B
+        :type mB_nkl: cute.Tensor
+        :param tma_atom_c: TMA copy atom for C tensor
+        :type tma_atom_c: cute.CopyAtom
+        :param mC_mnl: Output tensor C
+        :type mC_mnl: cute.Tensor
+        :param tiled_mma: Tiled MMA object
+        :type tiled_mma: cute.TiledMma
+        :param cta_layout_mnk: CTA layout
+        :type cta_layout_mnk: cute.Layout
+        :param a_smem_layout_staged: Shared memory layout for A
+        :type a_smem_layout_staged: cute.ComposedLayout
+        :param b_smem_layout_staged: Shared memory layout for B
+        :type b_smem_layout_staged: cute.ComposedLayout
+        :param epi_smem_layout_staged: Shared memory layout for epilogue
+        :type epi_smem_layout_staged: cute.ComposedLayout
+        """
+
+        tidx, _, _ = cute.arch.thread_idx()
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+        if warp_idx == 0:
+            cpasync.prefetch_descriptor(tma_atom_a)
+            cpasync.prefetch_descriptor(tma_atom_b)
+            cpasync.prefetch_descriptor(tma_atom_sfa)
+            cpasync.prefetch_descriptor(tma_atom_sfb)
+            cpasync.prefetch_descriptor(tma_atom_c)
+
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(
+            cute.arch.block_idx_in_cluster()
+        )
+        cluster_coord_mnk = cta_layout_mnk.get_flat_coord(cta_rank_in_cluster)
+
+        a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, 0))
+        b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, 0))
+        sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, 0))
+        sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, 0))
+        tma_copy_bytes = (
+            cute.size_in_bytes(self.a_dtype, a_smem_layout)
+            + cute.size_in_bytes(self.b_dtype, b_smem_layout)
+            + cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
+            + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
+        )
+
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+
+        mainloop_pipeline_array_ptr = storage.mainloop_pipeline_array_ptr.data_ptr()
+
+        mainloop_pipeline_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread
+        )
+        mainloop_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self.num_mma_warps
+        )
+
+        cta_layout_vmnk = cute.make_layout((1, *cta_layout_mnk.shape))
+        mainloop_pipeline = pipeline.PipelineTmaAsync.create(
+            num_stages=self.ab_stage,
+            producer_group=mainloop_pipeline_producer_group,
+            consumer_group=mainloop_pipeline_consumer_group,
+            tx_count=tma_copy_bytes,
+            barrier_storage=mainloop_pipeline_array_ptr,
+            cta_layout_vmnk=cta_layout_vmnk,
+        )
+
+        if cute.size(self.cluster_shape_mnk) > 1:
+            cute.arch.cluster_arrive_relaxed()
+
+        sA = storage.sA.get_tensor(
+            a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
+        )
+        sB = storage.sB.get_tensor(
+            b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
+        )
+        sC = storage.sC.get_tensor(
+            epi_smem_layout_staged.outer, swizzle=epi_smem_layout_staged.inner
+        )
+        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+        sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
+
+        # (bM, bK, loopM, loopK, loopL)
+        gA_mkl = cute.local_tile(
+            mA_mkl,
+            cute.slice_(self.tile_shape_mnk, (None, 0, None)),
+            (None, None, None),
+        )
+        # (bN, bK, loopN, loopK, loopL)
+        gB_nkl = cute.local_tile(
+            mB_nkl,
+            cute.slice_(self.tile_shape_mnk, (0, None, None)),
+            (None, None, None),
+        )
+        # (tM, tK, loopM, loopK, loopL)
+        gSFA_mkl = cute.local_tile(
+            mSFA_mkl,
+            cute.slice_(self.tile_shape_mnk, (None, 0, None)),
+            (None, None, None),
+        )
+        # (tN, tK, loopN, loopK, loopL)
+        gSFB_nkl = cute.local_tile(
+            mSFB_nkl,
+            cute.slice_(self.tile_shape_mnk, (0, None, None)),
+            (None, None, None),
+        )
+        # (bM, bN, loopM, loopN, loopL)
+        gC_mnl = cute.local_tile(
+            mC_mnl,
+            cute.slice_(self.tile_shape_mnk, (None, None, 0)),
+            (None, None, None),
+        )
+
+        thr_mma = tiled_mma.get_slice(tidx)
+
+        a_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (0, None, 0)).shape)
+        a_cta_crd = cluster_coord_mnk[1]
+        tAsA, tAgA = cpasync.tma_partition(
+            tma_atom_a,
+            a_cta_crd,
+            a_cta_layout,
+            cute.group_modes(sA, 0, 2),
+            cute.group_modes(gA_mkl, 0, 2),
+        )
+
+        b_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (None, 0, 0)).shape)
+        b_cta_crd = cluster_coord_mnk[0]
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_b,
+            b_cta_crd,
+            b_cta_layout,
+            cute.group_modes(sB, 0, 2),
+            cute.group_modes(gB_nkl, 0, 2),
+        )
+
+        tAsSFA, tAgSFA = cpasync.tma_partition(
+            tma_atom_sfa,
+            a_cta_crd,
+            a_cta_layout,
+            cute.group_modes(sSFA, 0, 2),
+            cute.group_modes(gSFA_mkl, 0, 2),
+        )
+        tAsSFA = cute.filter_zeros(tAsSFA)
+        tAgSFA = cute.filter_zeros(tAgSFA)
+
+        tBsSFB, tBgSFB = cpasync.tma_partition(
+            tma_atom_sfb,
+            b_cta_crd,
+            b_cta_layout,
+            cute.group_modes(sSFB, 0, 2),
+            cute.group_modes(gSFB_nkl, 0, 2),
+        )
+        tBsSFB = cute.filter_zeros(tBsSFB)
+        tBgSFB = cute.filter_zeros(tBgSFB)
+
+        tCsA = thr_mma.partition_A(sA)
+        tCsB = thr_mma.partition_B(sB)
+
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])
+        tCrSFA = sm120_utils.partition_fragment_SFA(sSFA[None, None, 0], thr_mma, tidx)
+        tCrSFB = sm120_utils.partition_fragment_SFB(sSFB[None, None, 0], thr_mma, tidx)
+        # Keep residual K modes nested to match the C++ SM120 block-scaled mainloop.
+        tCrSFA = cute.group_modes(tCrSFA, 2, cute.rank(tCrSFA))
+        tCrSFB = cute.group_modes(tCrSFB, 2, cute.rank(tCrSFB))
+
+        tCgC = thr_mma.partition_C(gC_mnl)
+        acc_shape = tCgC.shape[:3]
+        accumulators = cute.make_rmem_tensor(acc_shape, self.acc_dtype)
+
+        if cute.size(self.cluster_shape_mnk) > 1:
+            cute.arch.cluster_wait()
+        else:
+            cute.arch.sync_threads()
+
+        k_tile_cnt = cute.size(gA_mkl, mode=[3])
+
+        tile_sched = utils.StaticPersistentTileScheduler.create(
+            tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+        )
+        work_tile = tile_sched.initial_work_tile_info()
+
+        mainloop_producer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.ab_stage
+        )
+        mainloop_consumer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.ab_stage
+        )
+
+        # MMA warp group
+        if warp_idx < self.num_mma_warps:
+            cute.arch.setmaxregister_increase(self.mma_register_requirement)
+
+            num_k_blocks = cute.size(tCrA, mode=[2])
+
+            atom_copy_ldmatrix_A = make_ldmatrix_atom(
+                self.a_dtype,
+                transpose=self.a_layout.is_m_major_a(),
+                num_matrices=4,
+                mixed_mode=self.mixed_mode,
+            )
+            atom_copy_ldmatrix_B = make_ldmatrix_atom(
+                self.b_dtype,
+                transpose=self.b_layout.is_n_major_b(),
+                num_matrices=4,
+                mixed_mode=self.mixed_mode,
+            )
+            smem_tiled_copy_A = cute.make_tiled_copy_A(atom_copy_ldmatrix_A, tiled_mma)
+            smem_tiled_copy_B = cute.make_tiled_copy_B(atom_copy_ldmatrix_B, tiled_mma)
+
+            atom_copy_ldmatrix_SF = cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(),
+                self.sf_dtype,
+            )
+            smem_tiled_copy_SFA = cute.make_tiled_copy(
+                atom_copy_ldmatrix_SF,
+                sm120_utils.get_layoutSFA_TV(tiled_mma),
+                (
+                    cute.size(tiled_mma.permutation_mnk[0]),
+                    cute.size(tiled_mma.permutation_mnk[2]),
+                ),
+            )
+            smem_tiled_copy_SFB = cute.make_tiled_copy(
+                atom_copy_ldmatrix_SF,
+                sm120_utils.get_layoutSFB_TV(tiled_mma),
+                (
+                    cute.size(tiled_mma.permutation_mnk[1]),
+                    cute.size(tiled_mma.permutation_mnk[2]),
+                ),
+            )
+
+            thr_copy_ldmatrix_A = smem_tiled_copy_A.get_slice(tidx)
+            thr_copy_ldmatrix_B = smem_tiled_copy_B.get_slice(tidx)
+            tCsA_copy_view = thr_copy_ldmatrix_A.partition_S(sA)
+            tCrA_copy_view = thr_copy_ldmatrix_A.retile(tCrA)
+            tCsB_copy_view = thr_copy_ldmatrix_B.partition_S(sB)
+            tCrB_copy_view = thr_copy_ldmatrix_B.retile(tCrB)
+
+            thr_copy_ldmatrix_SFA = smem_tiled_copy_SFA.get_slice(tidx)
+            thr_copy_ldmatrix_SFB = smem_tiled_copy_SFB.get_slice(tidx)
+            tCsSFA_copy_view = thr_copy_ldmatrix_SFA.partition_S(sSFA)
+            tCrSFA_copy_view = thr_copy_ldmatrix_SFA.retile(tCrSFA)
+            tCsSFB_copy_view = thr_copy_ldmatrix_SFB.partition_S(sSFB)
+            tCrSFB_copy_view = thr_copy_ldmatrix_SFB.retile(tCrSFB)
+
+            epi_buffer = cutlass.Int32(0)
+
+            while work_tile.is_valid_tile:
+                tile_coord_mnl = work_tile.tile_idx
+                gC_mnl_slice = gC_mnl[(None, None, *tile_coord_mnl)]
+                accumulators.fill(0.0)
+
+                mainloop_consumer_state.reset_count()
+
+                peek_ab_full_status = cutlass.Boolean(1)
+                if mainloop_consumer_state.count < k_tile_cnt:
+                    peek_ab_full_status = mainloop_pipeline.consumer_try_wait(
+                        mainloop_consumer_state
+                    )
+
+                mainloop_pipeline.consumer_wait(
+                    mainloop_consumer_state, peek_ab_full_status
+                )
+                # tCsA_p: (MMA, (4, MMA_M / 4), MMA_K), tCsA_p: (MMA, (4, MMA_N / 4), MMA_K)
+                tCsA_p = tCsA_copy_view[None, None, None, mainloop_consumer_state.index]
+                tCsB_p = tCsB_copy_view[None, None, None, mainloop_consumer_state.index]
+                tCsSFA_p = tCsSFA_copy_view[
+                    None, None, None, mainloop_consumer_state.index
+                ]
+                tCsSFB_p = tCsSFB_copy_view[
+                    None, None, None, mainloop_consumer_state.index
+                ]
+                cute.copy(
+                    smem_tiled_copy_A,
+                    tCsA_p[None, None, 0],
+                    tCrA_copy_view[None, None, 0],
+                )
+                cute.copy(
+                    smem_tiled_copy_B,
+                    tCsB_p[None, None, 0],
+                    tCrB_copy_view[None, None, 0],
+                )
+
+                tCsSFA_p_filtered = cute.filter_zeros(tCsSFA_p)
+                tCsSFB_p_filtered = cute.filter_zeros(tCsSFB_p)
+                tCrSFA_copy_view_filtered = cute.filter_zeros(tCrSFA_copy_view)
+                tCrSFB_copy_view_filtered = cute.filter_zeros(tCrSFB_copy_view)
+
+                cute.copy(
+                    smem_tiled_copy_SFA,
+                    tCsSFA_p_filtered[None, None, 0],
+                    tCrSFA_copy_view_filtered[None, None, 0],
+                )
+                cute.copy(
+                    smem_tiled_copy_SFB,
+                    tCsSFB_p_filtered[None, None, 0],
+                    tCrSFB_copy_view_filtered[None, None, 0],
+                )
+
+                for _k_tile in range(0, k_tile_cnt - 1, 1, unroll=1):
+                    for k_block_idx in cutlass.range_constexpr(num_k_blocks):
+                        k_block_next = (
+                            0 if k_block_idx + 1 == num_k_blocks else k_block_idx + 1
+                        )
+
+                        if k_block_idx == num_k_blocks - 1:
+                            mainloop_pipeline.consumer_release(mainloop_consumer_state)
+                            mainloop_consumer_state.advance()
+
+                            peek_ab_full_status = cutlass.Boolean(1)
+                            peek_ab_full_status = mainloop_pipeline.consumer_try_wait(
+                                mainloop_consumer_state
+                            )
+
+                            tCsA_p = tCsA_copy_view[
+                                None, None, None, mainloop_consumer_state.index
+                            ]
+                            tCsB_p = tCsB_copy_view[
+                                None, None, None, mainloop_consumer_state.index
+                            ]
+                            tCsSFA_p = tCsSFA_copy_view[
+                                None, None, None, mainloop_consumer_state.index
+                            ]
+                            tCsSFB_p = tCsSFB_copy_view[
+                                None, None, None, mainloop_consumer_state.index
+                            ]
+                            mainloop_pipeline.consumer_wait(
+                                mainloop_consumer_state, peek_ab_full_status
+                            )
+
+                        # Mixed FP4 x FP8 register-side bit shift before mma.sync
+                        # to move the FP4 nibble (in low half of each byte after
+                        # b4x16_p64 ldmatrix) into the middle as mxf8f6f4 expects.
+                        if cutlass.const_expr(
+                            self.mixed_mode and self.a_dtype.width < 8
+                        ):
+                            a_view = cute.recast_tensor(
+                                tCrA[None, None, k_block_idx], cutlass.Int8
+                            )
+                            for _i in cutlass.range_constexpr(cute.size(a_view)):
+                                a_view[_i] = cutlass.Int8(a_view[_i] << FP4_SHIFT_BITS)
+                        if cutlass.const_expr(
+                            self.mixed_mode and self.b_dtype.width < 8
+                        ):
+                            b_view = cute.recast_tensor(
+                                tCrB[None, None, k_block_idx], cutlass.Int8
+                            )
+                            for _i in cutlass.range_constexpr(cute.size(b_view)):
+                                b_view[_i] = cutlass.Int8(b_view[_i] << FP4_SHIFT_BITS)
+                        cute.gemm(
+                            tiled_mma,
+                            accumulators,
+                            [
+                                tCrA[None, None, k_block_idx],
+                                tCrSFA[None, None, k_block_idx],
+                            ],
+                            [
+                                tCrB[None, None, k_block_idx],
+                                tCrSFB[None, None, k_block_idx],
+                            ],
+                            accumulators,
+                        )
+                        cute.copy(
+                            smem_tiled_copy_A,
+                            tCsA_p[None, None, k_block_next],
+                            tCrA_copy_view[None, None, k_block_next],
+                        )
+                        cute.copy(
+                            smem_tiled_copy_B,
+                            tCsB_p[None, None, k_block_next],
+                            tCrB_copy_view[None, None, k_block_next],
+                        )
+
+                        tCsSFA_p_filtered = cute.filter_zeros(tCsSFA_p)
+                        tCsSFB_p_filtered = cute.filter_zeros(tCsSFB_p)
+                        tCrSFA_copy_view_filtered = cute.filter_zeros(tCrSFA_copy_view)
+                        tCrSFB_copy_view_filtered = cute.filter_zeros(tCrSFB_copy_view)
+                        cute.copy(
+                            smem_tiled_copy_SFA,
+                            tCsSFA_p_filtered[None, None, k_block_next],
+                            tCrSFA_copy_view_filtered[None, None, k_block_next],
+                        )
+                        cute.copy(
+                            smem_tiled_copy_SFB,
+                            tCsSFB_p_filtered[None, None, k_block_next],
+                            tCrSFB_copy_view_filtered[None, None, k_block_next],
+                        )
+
+                # Hoist out last k_tile
+                for k_block_idx in cutlass.range_constexpr(num_k_blocks):
+                    k_block_next = (
+                        0 if k_block_idx + 1 == num_k_blocks else k_block_idx + 1
+                    )
+
+                    if k_block_idx == num_k_blocks - 1:
+                        cute.arch.fence_proxy("async.shared", space="cta")
+                        mainloop_pipeline.consumer_release(mainloop_consumer_state)
+                        mainloop_consumer_state.advance()
+
+                    if k_block_next > 0:
+                        cute.copy(
+                            smem_tiled_copy_A,
+                            tCsA_p[None, None, k_block_next],
+                            tCrA_copy_view[None, None, k_block_next],
+                        )
+                        cute.copy(
+                            smem_tiled_copy_B,
+                            tCsB_p[None, None, k_block_next],
+                            tCrB_copy_view[None, None, k_block_next],
+                        )
+                        tCsSFA_p_filtered = cute.filter_zeros(tCsSFA_p)
+                        tCsSFB_p_filtered = cute.filter_zeros(tCsSFB_p)
+                        tCrSFA_copy_view_filtered = cute.filter_zeros(tCrSFA_copy_view)
+                        tCrSFB_copy_view_filtered = cute.filter_zeros(tCrSFB_copy_view)
+                        cute.copy(
+                            smem_tiled_copy_SFA,
+                            tCsSFA_p_filtered[None, None, k_block_next],
+                            tCrSFA_copy_view_filtered[None, None, k_block_next],
+                        )
+                        cute.copy(
+                            smem_tiled_copy_SFB,
+                            tCsSFB_p_filtered[None, None, k_block_next],
+                            tCrSFB_copy_view_filtered[None, None, k_block_next],
+                        )
+                    # Mixed FP4 x FP8 register-side bit shift before mma.sync (hoisted tail).
+                    if cutlass.const_expr(self.mixed_mode and self.a_dtype.width < 8):
+                        a_view_h = cute.recast_tensor(
+                            tCrA[None, None, k_block_idx], cutlass.Int8
+                        )
+                        for _i in cutlass.range_constexpr(cute.size(a_view_h)):
+                            a_view_h[_i] = cutlass.Int8(a_view_h[_i] << FP4_SHIFT_BITS)
+                    if cutlass.const_expr(self.mixed_mode and self.b_dtype.width < 8):
+                        b_view_h = cute.recast_tensor(
+                            tCrB[None, None, k_block_idx], cutlass.Int8
+                        )
+                        for _i in cutlass.range_constexpr(cute.size(b_view_h)):
+                            b_view_h[_i] = cutlass.Int8(b_view_h[_i] << FP4_SHIFT_BITS)
+                    cute.gemm(
+                        tiled_mma,
+                        accumulators,
+                        [
+                            tCrA[None, None, k_block_idx],
+                            tCrSFA[None, None, k_block_idx],
+                        ],
+                        [
+                            tCrB[None, None, k_block_idx],
+                            tCrSFB[None, None, k_block_idx],
+                        ],
+                        accumulators,
+                    )
+
+                copy_atom_r2s = sm120_utils.sm120_get_smem_store_op(
+                    self.c_layout,
+                    elem_ty_d=self.c_dtype,
+                    elem_ty_acc=self.acc_dtype,
+                )
+
+                copy_atom_C = cute.make_copy_atom(
+                    cute.nvgpu.warp.StMatrix8x8x16bOp(
+                        self.c_layout.is_m_major_c(),
+                        2,
+                    ),
+                    self.c_dtype,
+                )
+
+                tiled_copy_C_Atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma)
+
+                tiled_copy_r2s = cute.make_tiled_copy_S(
+                    copy_atom_r2s,
+                    tiled_copy_C_Atom,
+                )
+
+                thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+                # (R2S, R2S_M, R2S_N, PIPE_D)
+                tRS_sD = thr_copy_r2s.partition_D(sC)
+                # (R2S, R2S_M, R2S_N)
+                tRS_rAcc = tiled_copy_r2s.retile(accumulators)
+
+                rD_shape = cute.shape(thr_copy_r2s.partition_S(sC))
+                tRS_rD_layout = cute.make_layout(rD_shape[:3])
+                tRS_rD = cute.make_rmem_tensor(tRS_rD_layout.shape, self.acc_dtype)
+                _size_tRS_rD = cute.size(tRS_rD)
+
+                sepi_for_tma_partition = cute.group_modes(sC, 0, 2)
+                tcgc_for_tma_partition = cute.zipped_divide(gC_mnl_slice, self.epi_tile)
+
+                bSG_sD, bSG_gD = cpasync.tma_partition(
+                    tma_atom_c,
+                    0,
+                    cute.make_layout(1),
+                    sepi_for_tma_partition,
+                    tcgc_for_tma_partition,
+                )
+
+                tma_store_producer_group = pipeline.CooperativeGroup(
+                    pipeline.Agent.Thread,
+                    self.num_mma_warps * self.num_threads_per_warp,
+                )
+                tma_store_pipeline = pipeline.PipelineTmaStore.create(
+                    num_stages=self.epi_stage,
+                    producer_group=tma_store_producer_group,
+                )
+
+                epi_rest_m = bSG_gD.shape[1][0]
+                epi_rest_n = bSG_gD.shape[1][1]
+                epi_tile_m = self.epi_tile[0]
+                epi_tile_n = self.epi_tile[1]
+                mma_tile_m = self.tile_shape_mnk[0] // cute.size(tRS_rAcc, mode=[1])
+                mma_tile_n = self.tile_shape_mnk[1] // cute.size(tRS_rAcc, mode=[2])
+
+                for epi_m in cutlass.range_constexpr(epi_rest_m):
+                    for epi_n in cutlass.range_constexpr(epi_rest_n):
+                        MmaMPerEpiM = epi_tile_m // mma_tile_m
+                        MmaNPerEpiN = epi_tile_n // mma_tile_n
+                        for mma_n_in_epi in cutlass.range_constexpr(MmaNPerEpiN):
+                            for mma_m_in_epi in cutlass.range_constexpr(MmaMPerEpiM):
+                                mma_n = (epi_n * MmaNPerEpiN) + mma_n_in_epi
+                                mma_m = (epi_m * MmaMPerEpiM) + mma_m_in_epi
+                                tRS_rD_slice = tRS_rD[
+                                    (None, mma_m_in_epi, mma_n_in_epi)
+                                ]
+                                tRS_rAcc_slice = tRS_rAcc[(None, mma_m, mma_n)]
+                                for elem_idx in cutlass.range_constexpr(
+                                    cute.size(tRS_rD_slice)
+                                ):
+                                    tRS_rD_slice[elem_idx] = tRS_rAcc_slice[elem_idx]
+
+                        tRS_rD_out = cute.make_rmem_tensor(
+                            tRS_rD_layout.shape, self.c_dtype
+                        )
+                        acc_vec = tRS_rD.load()
+                        tRS_rD_out.store(acc_vec.to(self.c_dtype))
+
+                        epi_buffer = epi_buffer + 1
+                        epi_buffer = epi_buffer % cute.size(tRS_sD, mode=[3])
+                        self.epilog_sync_barrier.arrive_and_wait()
+                        cute.copy(
+                            tiled_copy_r2s,
+                            tRS_rD_out,
+                            tRS_sD[(None, None, None, epi_buffer)],
+                        )
+                        cute.arch.fence_proxy(
+                            "async.shared",
+                            space="cta",
+                        )
+                        self.epilog_sync_barrier.arrive_and_wait()
+
+                        gmem_coord = (epi_m, epi_n)
+                        if warp_idx == 0:
+                            cute.copy(
+                                tma_atom_c,
+                                bSG_sD[(None, epi_buffer)],
+                                bSG_gD[(None, gmem_coord)],
+                            )
+                            tma_store_pipeline.producer_commit()
+                            tma_store_pipeline.producer_acquire()
+
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+        # DMA warp group
+        elif warp_idx == self.tma_load_warp_id:
+            cute.arch.setmaxregister_decrease(self.load_register_requirement)
+
+            while work_tile.is_valid_tile:
+                tile_coord_mnl = work_tile.tile_idx
+                # grouped: A/C span all sorted tokens (l=0); B/SFB pick this m-tile's expert
+                expert = m_indices[tile_coord_mnl[0]]
+                tAgA_mkl = tAgA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])]
+                tBgB_nkl = tBgB[(None, tile_coord_mnl[1], None, expert)]
+                tAgSFA_mkl = tAgSFA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])]
+                tBgSFB_nkl = tBgSFB[(None, tile_coord_mnl[1], None, expert)]
+
+                mainloop_producer_state.reset_count()
+
+                for _k_tile in range(0, k_tile_cnt, 1, unroll=1):
+                    # acquire also sets the transaction barrier for the A/B buffers
+                    mainloop_pipeline.producer_acquire(mainloop_producer_state)
+
+                    tAgA_k = tAgA_mkl[(None, mainloop_producer_state.count)]
+                    tAsA_pipe = tAsA[(None, mainloop_producer_state.index)]
+
+                    tBgB_k = tBgB_nkl[(None, mainloop_producer_state.count)]
+                    tBsB_pipe = tBsB[(None, mainloop_producer_state.index)]
+
+                    tAgSFA_k = tAgSFA_mkl[(None, mainloop_producer_state.count)]
+                    tAsSFA_pipe = tAsSFA[(None, mainloop_producer_state.index)]
+
+                    tBgSFB_k = tBgSFB_nkl[(None, mainloop_producer_state.count)]
+                    tBsSFB_pipe = tBsSFB[(None, mainloop_producer_state.index)]
+
+                    cute.copy(
+                        tma_atom_a,
+                        tAgA_k,
+                        tAsA_pipe,
+                        tma_bar_ptr=mainloop_pipeline.producer_get_barrier(
+                            mainloop_producer_state
+                        ),
+                    )
+                    cute.copy(
+                        tma_atom_b,
+                        tBgB_k,
+                        tBsB_pipe,
+                        tma_bar_ptr=mainloop_pipeline.producer_get_barrier(
+                            mainloop_producer_state
+                        ),
+                    )
+                    cute.copy(
+                        tma_atom_sfa,
+                        tAgSFA_k,
+                        tAsSFA_pipe,
+                        tma_bar_ptr=mainloop_pipeline.producer_get_barrier(
+                            mainloop_producer_state
+                        ),
+                    )
+                    cute.copy(
+                        tma_atom_sfb,
+                        tBgSFB_k,
+                        tBsSFB_pipe,
+                        tma_bar_ptr=mainloop_pipeline.producer_get_barrier(
+                            mainloop_producer_state
+                        ),
+                    )
+                    # Mainloop pipeline's producer commit is a NOP
+                    mainloop_pipeline.producer_commit(mainloop_producer_state)
+                    mainloop_producer_state.advance()
+
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+            mainloop_pipeline.producer_tail(mainloop_producer_state)
+        return
+
+    @staticmethod
+    def _compute_stages(
+        tile_shape_mnk: tuple[int, int, int],
+        a_dtype: type[cutlass.Numeric],
+        b_dtype: type[cutlass.Numeric],
+        sf_dtype: type[cutlass.Numeric],
+        sfa_smem_layout: cute.Layout,
+        sfb_smem_layout: cute.Layout,
+        epi_tile: tuple[int, int],
+        c_dtype: type[cutlass.Numeric],
+        smem_capacity: int,
+        occupancy: int,
+    ) -> tuple[int, int]:
+        """Computes the number of stages for A/B/C operands based on heuristics.
+
+        :param tile_shape_mnk: The shape (M, N, K) of the CTA tile.
+        :type tile_shape_mnk: tuple[int, int, int]
+        :param a_dtype: Data type of operand A.
+        :type a_dtype: type[cutlass.Numeric]
+        :param b_dtype: Data type of operand B.
+        :type b_dtype: type[cutlass.Numeric]
+        :param smem_capacity: Total available shared memory capacity in bytes.
+        :type smem_capacity: int
+        :param occupancy: Target number of CTAs per SM (occupancy).
+        :type occupancy: int
+
+        :return: A tuple containing the computed number of stages for:
+                 (A/B operand stages, epilogue stages)
+        :rtype: tuple[int, int]
+        """
+
+        epi_stage_max = (tile_shape_mnk[1] // epi_tile[1]) * (
+            tile_shape_mnk[0] // epi_tile[0]
+        )
+        epi_stage = min(epi_stage_max, 4)
+        c_bytes_per_stage = cute.size(epi_tile) * c_dtype.width // 8
+        epi_bytes = c_bytes_per_stage * epi_stage
+
+        a_shape = cute.slice_(tile_shape_mnk, (None, 0, None))
+        b_shape = cute.slice_(tile_shape_mnk, (0, None, None))
+        ab_bytes_per_stage = (
+            cute.size(a_shape) * a_dtype.width // 8
+            + cute.size(b_shape) * b_dtype.width // 8
+        )
+        sf_bytes_per_stage = (
+            cute.size(cute.filter_zeros(sfa_smem_layout).shape) * sf_dtype.width // 8
+            + cute.size(cute.filter_zeros(sfb_smem_layout).shape) * sf_dtype.width // 8
+        )
+        mbar_helpers_bytes = 1024
+
+        ab_stage = (
+            (smem_capacity - occupancy * 1024) // occupancy
+            - mbar_helpers_bytes
+            - epi_bytes
+        ) // (ab_bytes_per_stage + sf_bytes_per_stage)
+        return ab_stage, epi_stage
+
+    @staticmethod
+    def _make_smem_layouts(
+        tile_shape_mnk: tuple[int, int, int],
+        epi_tile: tuple[int, int],
+        a_dtype: type[cutlass.Numeric],
+        a_layout: cute.Layout,
+        b_dtype: type[cutlass.Numeric],
+        b_layout: cute.Layout,
+        ab_stage: int,
+        c_dtype: type[cutlass.Numeric],
+        c_layout: cute.Layout,
+        epi_stage: int,
+        sf_vec_size: int,
+        tiled_mma: cute.TiledMma,
+    ) -> tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout]:
+        """Create shared memory layouts for A, B, and C tensors.
+
+        :param tile_shape_mnk: CTA tile shape (M,N,K)
+        :type tile_shape_mnk: Tuple[int, int, int]
+        :param epi_tile: Epilogue tile shape
+        :type epi_tile: Tuple[int, int]
+        :param a_dtype: Data type for matrix A
+        :type a_dtype: type[cutlass.Numeric]
+        :param a_layout: Layout for matrix A
+        :type a_layout: Layout
+        :param b_dtype: Data type for matrix B
+        :type b_dtype: type[cutlass.Numeric]
+        :param b_layout: Layout for matrix B
+        :type b_layout: Layout
+        :param ab_stage: Number of stages for A/B tensors
+        :type ab_stage: int
+        :param c_dtype: Data type for output matrix C
+        :type c_dtype: type[cutlass.Numeric]
+        :param c_layout: leading dimension of the output matrix C
+        :type c_layout: Layout
+        :param epi_stage: Number of epilogue stages
+        :type epi_stage: int
+
+        :return: Tuple of shared memory layouts for A, B, and C
+        :rtype: Tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout]
+        """
+        a_smem_shape = cute.slice_(tile_shape_mnk, (None, 0, None))
+
+        a_is_k_major = a_layout.is_k_major_a()
+        b_is_k_major = b_layout.is_k_major_b()
+        a_major_mode_size = tile_shape_mnk[2 if a_is_k_major else 0]
+
+        a_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom(
+            sm90_utils.get_smem_layout_atom(
+                a_layout,
+                a_dtype,
+                a_major_mode_size,
+            ),
+            a_dtype,
+        )
+        a_smem_layout_staged = cute.tile_to_shape(
+            a_smem_layout_atom,
+            cute.append(a_smem_shape, ab_stage),
+            order=(0, 1, 2) if a_is_k_major else (1, 0, 2),
+        )
+
+        b_smem_shape = cute.slice_(tile_shape_mnk, (0, None, None))
+
+        b_major_mode_size = tile_shape_mnk[2 if b_is_k_major else 1]
+        b_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom(
+            sm90_utils.get_smem_layout_atom(
+                b_layout,
+                b_dtype,
+                b_major_mode_size,
+            ),
+            b_dtype,
+        )
+        b_smem_layout_staged = cute.tile_to_shape(
+            b_smem_layout_atom,
+            cute.append(b_smem_shape, ab_stage),
+            order=(0, 1, 2) if b_is_k_major else (1, 0, 2),
+        )
+
+        sfa_smem_layout_staged = blockscaled_utils.sm120_make_smem_layout_sfa(
+            tiled_mma,
+            tile_shape_mnk,
+            sf_vec_size,
+            ab_stage,
+        )
+
+        sfb_smem_layout_staged = blockscaled_utils.sm120_make_smem_layout_sfb(
+            tiled_mma,
+            tile_shape_mnk,
+            sf_vec_size,
+            ab_stage,
+        )
+
+        c_smem_shape = epi_tile
+        c_major_mode_size = epi_tile[1] if c_layout.is_n_major_c() else epi_tile[0]
+        c_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom(
+            sm90_utils.get_smem_layout_atom(
+                c_layout,
+                c_dtype,
+                c_major_mode_size,
+            ),
+            c_dtype,
+        )
+        epi_smem_layout_staged = cute.tile_to_shape(
+            c_smem_layout_atom,
+            cute.append(c_smem_shape, epi_stage),
+            order=(1, 0, 2) if c_layout.is_m_major_c() else (0, 1, 2),
+        )
+
+        return (
+            a_smem_layout_staged,
+            b_smem_layout_staged,
+            sfa_smem_layout_staged,
+            sfb_smem_layout_staged,
+            epi_smem_layout_staged,
+        )
+
+    @staticmethod
+    def _compute_grid(
+        c: cute.Tensor,
+        tile_shape_mnk: tuple[int, int, int],
+        max_active_clusters: cutlass.Constexpr,
+    ) -> tuple[int, int, int]:
+        """Compute grid shape for the output tensor C.
+
+        :param c: The output tensor C
+        :type c: cute.Tensor
+        :param tile_shape_mnk: The shape (M, N, K) of the CTA tile.
+        :type tile_shape_mnk: tuple[int, int, int]
+
+        :return: Grid shape for kernel launch.
+        :rtype: tuple[int, int, int]
+        """
+
+        c_shape = cute.slice_(tile_shape_mnk, (None, None, 0))
+        gc = cute.zipped_divide(c, tiler=c_shape)
+        num_ctas_mnl = gc[(0, (None, None, None))].shape
+        cluster_shape_mnl = (1, 1, 1)
+        tile_sched_params = utils.PersistentTileSchedulerParams(
+            num_ctas_mnl, cluster_shape_mnl
+        )
+        grid = utils.StaticPersistentTileScheduler.get_grid_shape(
+            tile_sched_params, max_active_clusters
+        )
+        return tile_sched_params, grid
+
+    @staticmethod
+    def _make_tma_store_atoms_and_tensors(
+        tensor_c: cute.Tensor,
+        epi_smem_layout_staged: cute.ComposedLayout,
+        epi_tile: tuple[int, int],
+    ) -> tuple[cute.CopyAtom, cute.Tensor]:
+        """Create TMA atoms and tensors for C tensor storage.
+
+        :param tensor_c: Output tensor C
+        :type tensor_c: cute.Tensor
+        :param epi_smem_layout_staged: Shared memory layout for epilogue
+        :type epi_smem_layout_staged: cute.ComposedLayout
+        :param epi_tile: Epilogue tile shape
+        :type epi_tile: Tuple[int, int]
+
+        :return: TMA atom and tensor for C
+        :rtype: Tuple[cute.CopyAtom, cute.Tensor]
+        """
+        epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0))
+        tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileS2GOp(),
+            tensor_c,
+            epi_smem_layout,
+            epi_tile,
+        )
+
+        return tma_atom_c, tma_tensor_c
+
+    @staticmethod
+    def _make_tma_atoms_and_tensors(
+        tensor: cute.Tensor,
+        smem_layout_staged: cute.ComposedLayout,
+        smem_tile: tuple[int, int],
+        mcast_dim: int,
+        internal_type: Optional[Type[cutlass.Numeric]] = None,
+    ) -> tuple[cute.CopyAtom, cute.Tensor]:
+        """Create TMA atoms and tensors for input tensors.
+
+        :param tensor: Input tensor (A or B)
+        :type tensor: cute.Tensor
+        :param smem_layout_staged: Shared memory layout for the tensor
+        :type smem_layout_staged: cute.ComposedLayout
+        :param smem_tile: Shared memory tile shape
+        :type smem_tile: Tuple[int, int]
+        :param mcast_dim: Multicast dimension
+        :type mcast_dim: int
+
+        :return: TMA atom and tensor
+        :rtype: Tuple[cute.CopyAtom, cute.Tensor]
+        """
+        op = (
+            cpasync.CopyBulkTensorTileG2SOp()
+            if mcast_dim == 1
+            else cpasync.CopyBulkTensorTileG2SMulticastOp()
+        )
+
+        smem_layout = cute.slice_(smem_layout_staged, (None, None, 0))
+        tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(
+            op,
+            tensor,
+            smem_layout,
+            smem_tile,
+            num_multicast=mcast_dim,
+            internal_type=internal_type,
+        )
+
+        return tma_atom, tma_tensor
+
+    @staticmethod
+    def is_valid_tensor_alignment(
+        m: int,
+        n: int,
+        k: int,
+        l: int,
+        ab_dtype: Type[cutlass.Numeric],
+        c_dtype: Type[cutlass.Numeric],
+        a_major: str,
+        b_major: str,
+        c_major: str,
+    ) -> bool:
+        """
+        Check if the tensor alignment is valid
+
+        :param m: The number of rows in the A tensor
+        :type m: int
+        :param n: The number of columns in the B tensor
+        :type n: int
+        :param k: The number of columns in the A tensor
+        :type k: int
+        :param l: The number of columns in the C tensor
+        :type l: int
+        :param ab_dtype: The data type of the A and B operands
+        :type ab_dtype: Type[cutlass.Numeric]
+        :param c_dtype: The data type of the output tensor
+        :type c_dtype: Type[cutlass.Numeric]
+        :param a_major: The major axis of the A tensor
+        :type a_major: str
+        :param b_major: The major axis of the B tensor
+        :type b_major: str
+        :param c_major: The major axis of the C tensor
+        :type c_major: str
+
+        :return: True if the problem shape is valid, False otherwise
+        :rtype: bool
+        """
+        is_valid = True
+
+        def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
+            major_mode_idx = 0 if is_mode0_major else 1
+            num_major_elements = tensor_shape[major_mode_idx]
+            num_contiguous_elements = 16 * 8 // dtype.width
+            return num_major_elements % num_contiguous_elements == 0
+
+        if (
+            not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l))
+            or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l))
+            or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l))
+        ):
+            is_valid = False
+        return is_valid
+
+
+@cute.jit
+def cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
+    sf_ref_tensor: cute.Tensor,
+    sf_mma_tensor: cute.Tensor,
+):
+    """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout"""
+    # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l)
+    # group to ((32, 4, rest_m), (4, rest_k), l)
+    sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3)
+    sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3)
+    for i in cutlass.range(cute.size(sf_ref_tensor)):
+        mkl_coord = sf_ref_tensor.layout.get_hier_coord(i)
+        sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord]
+
+
+def run(
+    mnkl: Tuple[int, int, int, int],
+    a_dtype: Type[cutlass.Numeric],
+    b_dtype: Type[cutlass.Numeric],
+    sf_dtype: Type[cutlass.Numeric],
+    sf_vec_size: int,
+    c_dtype: Type[cutlass.Numeric],
+    a_major: str = "k",
+    b_major: str = "k",
+    c_major: str = "n",
+    tile_shape_mnk: Tuple[int, int, int] = (128, 128, 128),
+    epi_tile: Tuple[int, int] = (128, 128),
+    tolerance: float = 1e-01,
+    warmup_iterations: int = 0,
+    iterations: int = 1,
+    skip_ref_check: bool = False,
+    use_cold_l2: bool = False,
+    **kwargs,
+):
+    """Perf-framework-compatible entry point.
+
+    FP4 MMA only supports Float32 accumulation, so acc_dtype is always
+    Float32 regardless of what the caller passes.
+    """
+    return run_bs(
+        mnkl=mnkl,
+        a_dtype=a_dtype,
+        b_dtype=b_dtype,
+        sf_dtype=sf_dtype,
+        sf_vec_size=sf_vec_size,
+        c_dtype=c_dtype,
+        acc_dtype=cutlass.Float32,
+        a_major=a_major,
+        b_major=b_major,
+        c_major=c_major,
+        tile_shape_mnk=tile_shape_mnk,
+        epi_tile=epi_tile,
+        tolerance=tolerance,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+        skip_ref_check=skip_ref_check,
+        use_cold_l2=use_cold_l2,
+    )
+
+
+def run_bs(
+    mnkl: Tuple[int, int, int, int],
+    a_dtype: Type[cutlass.Numeric],
+    b_dtype: Type[cutlass.Numeric],
+    sf_dtype: Type[cutlass.Numeric],
+    sf_vec_size: int,
+    c_dtype: Type[cutlass.Numeric],
+    acc_dtype: Type[cutlass.Numeric],
+    a_major: str,
+    b_major: str,
+    c_major: str,
+    tile_shape_mnk: Tuple[int, int, int],
+    epi_tile: Tuple[int, int],
+    tolerance: float,
+    warmup_iterations: int,
+    iterations: int,
+    skip_ref_check: bool,
+    use_cold_l2: bool = False,
+    **kwargs,
+):
+    import cutlass.torch as cutlass_torch
+    import torch
+
+    print("Running Blackwell Geforce Blockscaled Dense GEMM with:")
+    print(f"mnkl: {mnkl}")
+    print(
+        f"A dtype: {a_dtype}, B dtype: {b_dtype}, A/B scale factor dtype: {sf_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}"
+    )
+    print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}")
+    print(f"Tile Shape: {tile_shape_mnk}")
+    print(f"Epilogue tile: {epi_tile}")
+    print(f"Tolerance: {tolerance}")
+    print(f"Warmup iterations: {warmup_iterations}")
+    print(f"Iterations: {iterations}")
+    print(f"Skip reference checking: {skip_ref_check}")
+    print(f"Use cold L2: {use_cold_l2}")
+
+    m, n, k, l = mnkl
+
+    if not Sm120BlockScaledGemmKernel.is_valid_tensor_alignment(
+        m, n, k, l, a_dtype, c_dtype, a_major, b_major, c_major
+    ):
+        raise ValueError("Invalid tensor alignment")
+
+    a_dtype = getattr(cutlass, a_dtype) if isinstance(a_dtype, str) else a_dtype
+    b_dtype = getattr(cutlass, b_dtype) if isinstance(b_dtype, str) else b_dtype
+    c_dtype = getattr(cutlass, c_dtype) if isinstance(c_dtype, str) else c_dtype
+    acc_dtype = getattr(cutlass, acc_dtype) if isinstance(acc_dtype, str) else acc_dtype
+
+    m, n, k, l = mnkl
+    cluster_shape_mnk = (1, 1, 1)
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("GPU is required to run this example!")
+
+    a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32)
+    b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32)
+    c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32)
+
+    a_tensor, a_torch = cutlass_torch.cute_tensor_like(
+        a_ref, a_dtype, is_dynamic_layout=True, assumed_align=16
+    )
+    b_tensor, b_torch = cutlass_torch.cute_tensor_like(
+        b_ref, b_dtype, is_dynamic_layout=True, assumed_align=16
+    )
+    c_tensor, c_torch = cutlass_torch.cute_tensor_like(
+        c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16
+    )
+
+    a_tensor.mark_compact_shape_dynamic(
+        mode=1 if a_major == "k" else 0,
+        stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0),
+        divisibility=2 if a_dtype == cutlass.Float4E2M1FN else 1,
+    )
+    b_tensor.mark_compact_shape_dynamic(
+        mode=1 if b_major == "k" else 0,
+        stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0),
+        divisibility=2 if b_dtype == cutlass.Float4E2M1FN else 1,
+    )
+    c_tensor.mark_compact_shape_dynamic(
+        mode=1 if c_major == "n" else 0,
+        stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0),
+        divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1,
+    )
+
+    def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype):
+        def ceil_div(a, b):
+            return (a + b - 1) // b
+
+        sf_k = ceil_div(k, sf_vec_size)
+        ref_shape = (l, mn, sf_k)
+
+        # The A, C, and D matrices are row-major whereas the B matrix is column-major.
+        # So only k-major (A/B) is supported.
+
+        atom_m = (32, 4)
+        atom_k = 4
+        mma_shape = (
+            l,
+            ceil_div(mn, atom_m[0] * atom_m[1]),
+            ceil_div(sf_k, atom_k),
+            atom_m[0],
+            atom_m[1],
+            atom_k,
+        )
+
+        ref_permute_order = (1, 2, 0)
+        mma_permute_order = (3, 4, 1, 5, 2, 0)
+
+        ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor(
+            ref_shape,
+            torch.float32,
+            permute_order=ref_permute_order,
+            init_type=cutlass_torch.TensorInitType.RANDOM,
+            init_config=cutlass_torch.RandomInitConfig(
+                min_val=1,
+                max_val=3,
+            ),
+        )
+
+        cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor(
+            mma_shape,
+            torch.float32,
+            permute_order=mma_permute_order,
+            init_type=cutlass_torch.TensorInitType.RANDOM,
+            init_config=cutlass_torch.RandomInitConfig(
+                min_val=0,
+                max_val=1,
+            ),
+        )
+
+        cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
+            from_dlpack(ref_f32_torch_tensor_cpu),
+            from_dlpack(cute_f32_torch_tensor_cpu),
+        )
+        cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda()
+
+        # reshape makes memory contiguous
+        ref_f32_torch_tensor_cpu = (
+            ref_f32_torch_tensor_cpu.permute(2, 0, 1)
+            .unsqueeze(-1)
+            .expand(l, mn, sf_k, sf_vec_size)
+            .reshape(l, mn, sf_k * sf_vec_size)
+            .permute(*ref_permute_order)
+        )
+        # prune to mkl for reference check.
+        ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :]
+
+        cute_tensor, torch_tensor = cutlass_torch.cute_tensor_like(
+            cute_f32_torch_tensor_cpu,
+            dtype,
+            is_dynamic_layout=True,
+            assumed_align=16,
+        )
+
+        cute_tensor = cutlass_torch.convert_cute_tensor(
+            cute_f32_torch_tensor,
+            cute_tensor,
+            dtype,
+            is_dynamic_layout=True,
+        )
+
+        return ref_f32_torch_tensor_cpu, cute_tensor, torch_tensor
+
+    sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor(
+        l, m, k, sf_vec_size, sf_dtype
+    )
+
+    sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor(
+        l, n, k, sf_vec_size, sf_dtype
+    )
+
+    gemm = Sm120BlockScaledGemmKernel(
+        acc_dtype,
+        sf_vec_size,
+        tile_shape_mnk,
+        epi_tile,
+    )
+
+    hardware_info = cutlass.utils.HardwareInfo()
+    max_active_clusters = hardware_info.get_max_active_clusters(
+        cluster_shape_mnk[0] * cluster_shape_mnk[1]
+    )
+
+    stream = cutlass_torch.default_stream()
+
+    compiled_gemm = cute.compile(
+        gemm,
+        a_tensor,
+        b_tensor,
+        sfa_tensor,
+        sfb_tensor,
+        c_tensor,
+        max_active_clusters,
+        stream,
+    )
+
+    if not skip_ref_check:
+        print("Reference checking ...")
+        compiled_gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, stream)
+        torch.cuda.synchronize()
+
+        res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref)
+        res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref)
+        ref = torch.einsum("mkl,nkl->mnl", res_a, res_b)
+
+        c_ref_device = c_ref.cuda()
+        cute.testing.convert(
+            c_tensor,
+            from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic(
+                leading_dim=(1 if c_major == "n" else 0)
+            ),
+        )
+        c_ref = c_ref_device.cpu()
+
+        if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16):
+            torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02)
+        elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN):
+            # Convert ref : f32 -> f8 -> f32
+            ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute(
+                1, 2, 0
+            )
+            ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic(
+                leading_dim=1
+            )
+            ref_f8.element_type = c_dtype
+            ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda()
+            ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic(
+                leading_dim=1
+            )
+            cute.testing.convert(ref_tensor, ref_f8)
+            cute.testing.convert(ref_f8, ref_tensor)
+            ref = ref_device.cpu()
+            torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02)
+        elif c_dtype is cutlass.Float4E2M1FN:
+            # Convert ref : f32 -> f4 -> f32
+            ref_f4_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute(
+                1, 2, 0
+            )
+            ref_f4 = from_dlpack(ref_f4_, assumed_align=16).mark_layout_dynamic(
+                leading_dim=1
+            )
+            ref_f4.element_type = c_dtype
+            ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda()
+            ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic(
+                leading_dim=1
+            )
+            cute.testing.convert(ref_tensor, ref_f4)
+            cute.testing.convert(ref_f4, ref_tensor)
+            ref = ref_device.cpu()
+            torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02)
+
+    def generate_tensors():
+        a_tensor, _ = cutlass_torch.cute_tensor_like(
+            a_ref, a_dtype, is_dynamic_layout=True, assumed_align=16
+        )
+        b_tensor, _ = cutlass_torch.cute_tensor_like(
+            b_ref, b_dtype, is_dynamic_layout=True, assumed_align=16
+        )
+        c_tensor, _ = cutlass_torch.cute_tensor_like(
+            c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16
+        )
+        a_tensor.mark_compact_shape_dynamic(
+            mode=1 if a_major == "k" else 0,
+            stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0),
+            divisibility=2 if a_dtype == cutlass.Float4E2M1FN else 1,
+        )
+        b_tensor.mark_compact_shape_dynamic(
+            mode=1 if b_major == "k" else 0,
+            stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0),
+            divisibility=2 if b_dtype == cutlass.Float4E2M1FN else 1,
+        )
+        c_tensor.mark_compact_shape_dynamic(
+            mode=1 if c_major == "n" else 0,
+            stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0),
+            divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1,
+        )
+
+        _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype)
+        _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype)
+        return cute.testing.JitArguments(
+            a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, stream
+        )
+
+    workspace_count = 1
+    if use_cold_l2:
+        one_workspace_bytes = (
+            a_ref.numel() * a_ref.element_size()
+            + b_ref.numel() * b_ref.element_size()
+            + sfa_ref.numel() * sfa_ref.element_size()
+            + sfb_ref.numel() * sfb_ref.element_size()
+            + c_ref.numel() * c_ref.element_size()
+        )
+        workspace_count = testing.get_workspace_count(
+            one_workspace_bytes, warmup_iterations, iterations
+        )
+
+    exec_time = testing.benchmark(
+        compiled_gemm,
+        workspace_generator=generate_tensors,
+        workspace_count=workspace_count,
+        stream=stream,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+    )
+
+    gflop = 2 * m * n * k / 1e9
+    gflops = gflop / exec_time * 1e6
+
+    print(f"Execution time: {exec_time} microseconds per iteration")
+    print(f"GFLOPS: {gflops}")
+
+    return exec_time  # Return execution time in microseconds
+
+
+def parse_arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Example of MxNxKxL GEMM on Blackwell Geforce."
+    )
+
+    parser.add_argument(
+        "--mnkl",
+        type=parse_comma_separated_ints,
+        default=(
+            1024,
+            1024,
+            1024,
+            1,
+        ),
+        help="mnkl dimensions (comma-separated)",
+    )
+    parser.add_argument(
+        "--tile_shape_mnk",
+        type=parse_comma_separated_ints,
+        choices=[
+            (128, 128, 128),
+            (128, 128, 256),
+        ],
+        default=(128, 128, 128),
+        help="CTA tile shape (comma-separated)",
+    )
+    parser.add_argument(
+        "--epi_tile",
+        type=parse_comma_separated_ints,
+        choices=[
+            (128, 128),
+            (64, 32),
+        ],
+        default=(128, 128),
+        help="Epilogue tile shape (comma-separated)",
+    )
+    parser.add_argument(
+        "--a_dtype",
+        type=cutlass.dtype,
+        default=cutlass.Float4E2M1FN,
+    )
+    parser.add_argument(
+        "--b_dtype",
+        type=cutlass.dtype,
+        default=cutlass.Float4E2M1FN,
+    )
+    parser.add_argument(
+        "--sf_dtype",
+        type=cutlass.dtype,
+        default=cutlass.Float8E4M3FN,
+    )
+    parser.add_argument(
+        "--sf_vec_size",
+        type=int,
+        choices=[16, 32],  # 16 for NVFP4, 32 for MXFP4 / MXFP8.
+        default=16,
+    )
+    parser.add_argument(
+        "--c_dtype",
+        type=cutlass.dtype,
+        default=cutlass.Float16,
+    )
+    parser.add_argument(
+        "--acc_dtype",
+        type=cutlass.dtype,
+        default=cutlass.Float32,
+    )
+    parser.add_argument("--a_major", choices=["k"], type=str, default="k")
+    parser.add_argument("--b_major", choices=["k"], type=str, default="k")
+    parser.add_argument("--c_major", choices=["n"], type=str, default="n")
+    parser.add_argument(
+        "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
+    )
+    parser.add_argument(
+        "--warmup_iterations", type=int, default=0, help="Warmup iterations"
+    )
+    parser.add_argument(
+        "--iterations",
+        type=int,
+        default=1,
+        help="Number of iterations to run the kernel",
+    )
+    parser.add_argument(
+        "--skip_ref_check",
+        action="store_true",
+        default=False,
+        help="Skip reference checking",
+    )
+    parser.add_argument(
+        "--use_cold_l2",
+        action="store_true",
+        default=False,
+        help="Use circular buffer tensor sets to ensure L2 cold cache",
+    )
+
+    args = parser.parse_args()
+
+    if len(args.mnkl) != 4:
+        parser.error("--mnkl must contain exactly 4 values")
+
+    fp4_allowed_tiles = {(128, 128, 128), (128, 128, 256)}
+    fp8_allowed_tiles = {(128, 128, 128)}
+    try:
+        validate_blockscaled_args(args, fp4_allowed_tiles, fp8_allowed_tiles)
+    except ValueError as e:
+        parser.error(str(e))
+
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    run_bs(
+        args.mnkl,
+        args.a_dtype,
+        args.b_dtype,
+        args.sf_dtype,
+        args.sf_vec_size,
+        args.c_dtype,
+        args.acc_dtype,
+        args.a_major,
+        args.b_major,
+        args.c_major,
+        args.tile_shape_mnk,
+        args.epi_tile,
+        args.tolerance,
+        args.warmup_iterations,
+        args.iterations,
+        args.skip_ref_check,
+        args.use_cold_l2,
+    )
+    print("PASS")
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/quant_mxfp8.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/quant_mxfp8.py
new file mode 100644
index 0000000000..122991c4a6
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/quant_mxfp8.py
@@ -0,0 +1,52 @@
+"""Fast Triton mxfp8 quantizer for activations (fp8 mode): bf16 [M,K] -> qdata e4m3 [M,K] +
+e8m0 block scale [M,K/32]. Matches torchao MXTensor.to_mx(x, float8_e4m3fn, 32, FLOOR)."""
+
+import torch
+import triton
+import triton.language as tl
+
+E4M3_MAX = 448.0
+
+
+@triton.jit
+def _qk(X, Q, S, sx0, sx1, sq0, sq1, ss0, ss1, BPM: tl.constexpr):
+    m = tl.program_id(0)
+    b0 = tl.program_id(1) * BPM
+    blk = tl.arange(0, BPM)
+    e = tl.arange(0, 32)
+    k = (b0 + blk)[:, None] * 32 + e[None, :]
+    x = tl.load(X + m * sx0 + k * sx1).to(tl.float32)
+    amax = tl.max(tl.abs(x), axis=1)
+    # FLOOR e8m0 scale: 2^(floor(log2(amax)) - floor(log2(448))), floor(log2(448))=8
+    exp = tl.floor(tl.log2(tl.where(amax > 0, amax, 1.0))) - 8.0
+    sc = tl.exp2(exp)
+    scb = sc[:, None]
+    q = tl.clamp(x / scb, -448.0, 448.0).to(tl.float8e4nv)
+    tl.store(Q + m * sq0 + k * sq1, q)
+    # e8m0 stores the biased exponent byte (bias 127); scale = 2^(E-127)
+    Ebiased = (exp + 127.0).to(tl.int32)
+    tl.store(S + m * ss0 + (b0 + blk) * ss1, Ebiased.to(tl.uint8))
+
+
+def mxfp8_quant(x, bpm=4):
+    M, K = x.shape
+    # K must be a whole number of 32-wide MX blocks, else nblk drops the tail and the kernel OOBs.
+    assert K % 32 == 0, f"mxfp8_quant expects K divisible by 32, got K={K}"
+    nblk = K // 32
+    Q = torch.empty(M, K, device=x.device, dtype=torch.float8_e4m3fn)
+    S = torch.empty(
+        M, nblk, device=x.device, dtype=torch.uint8
+    )  # e8m0 biased-exponent bytes
+    _qk[(M, triton.cdiv(nblk, bpm))](
+        x,
+        Q,
+        S,
+        x.stride(0),
+        x.stride(1),
+        Q.stride(0),
+        Q.stride(1),
+        S.stride(0),
+        S.stride(1),
+        BPM=bpm,
+    )
+    return Q, S.view(torch.float8_e8m0fnu)
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/quant_nvfp4.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/quant_nvfp4.py
new file mode 100644
index 0000000000..d2a5f5d3f7
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/quant_nvfp4.py
@@ -0,0 +1,69 @@
+"""Fast Triton nvfp4 quantizer: bf16 [M,K] -> qdata u8 [M,K/2] (2 fp4/byte, low-first) +
+e4m3 block scale [M,K/16]. Same qdata/scale layout as torchao (drop-in for the grouped FP4 GEMM),
+but an approximate fast encode: NOT byte-identical to torchao's to_nvfp4 (arithmetic-encode cascade
++ 1e-30 scale floor here vs torchao's RNE + 2**-6 e4m3 floor)."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _qk(X, Q, S, sx0, sx1, sq0, sq1, ss0, ss1, BPM: tl.constexpr):
+    m = tl.program_id(0)
+    b0 = tl.program_id(1) * BPM
+    blk = tl.arange(0, BPM)
+    e = tl.arange(0, 8)  # 8 bytes per 16-wide block
+    # even (lo) and odd (hi) fp4 positions interleaved within each byte
+    k_lo = (b0 + blk)[:, None] * 16 + 2 * e[None, :]
+    k_hi = k_lo + 1
+    lo = tl.load(X + m * sx0 + k_lo * sx1).to(tl.float32)
+    hi = tl.load(X + m * sx0 + k_hi * sx1).to(tl.float32)
+    amax = tl.maximum(tl.max(tl.abs(lo), axis=1), tl.max(tl.abs(hi), axis=1))
+    sc1d = tl.clamp(amax / 6.0, 1e-30, 1e30).to(tl.float8e4nv).to(tl.float32)
+    sc1d = tl.where(sc1d > 0, sc1d, 1.0)
+    sc = sc1d[:, None]
+    alo = tl.abs(lo) / sc
+    clo = (
+        tl.where(alo > 0.25, 1, 0)
+        + tl.where(alo > 0.75, 1, 0)
+        + tl.where(alo > 1.25, 1, 0)
+        + tl.where(alo > 1.75, 1, 0)
+        + tl.where(alo > 2.5, 1, 0)
+        + tl.where(alo > 3.5, 1, 0)
+        + tl.where(alo > 5.0, 1, 0)
+    ).to(tl.int32) + tl.where(lo < 0, 8, 0).to(tl.int32)
+    ahi = tl.abs(hi) / sc
+    chi = (
+        tl.where(ahi > 0.25, 1, 0)
+        + tl.where(ahi > 0.75, 1, 0)
+        + tl.where(ahi > 1.25, 1, 0)
+        + tl.where(ahi > 1.75, 1, 0)
+        + tl.where(ahi > 2.5, 1, 0)
+        + tl.where(ahi > 3.5, 1, 0)
+        + tl.where(ahi > 5.0, 1, 0)
+    ).to(tl.int32) + tl.where(hi < 0, 8, 0).to(tl.int32)
+    byte = (clo + chi * 16).to(tl.uint8)
+    qpos = (b0 + blk)[:, None] * 8 + e[None, :]
+    tl.store(Q + m * sq0 + qpos * sq1, byte)
+    tl.store(S + m * ss0 + (b0 + blk) * ss1, sc1d.to(tl.float8e4nv))
+
+
+def nvfp4_quant(x, bpm=8):
+    M, K = x.shape
+    nblk = K // 16
+    Q = torch.empty(M, K // 2, device=x.device, dtype=torch.uint8)
+    S = torch.empty(M, nblk, device=x.device, dtype=torch.float8_e4m3fn)
+    _qk[(M, triton.cdiv(nblk, bpm))](
+        x,
+        Q,
+        S,
+        x.stride(0),
+        x.stride(1),
+        Q.stride(0),
+        Q.stride(1),
+        S.stride(0),
+        S.stride(1),
+        BPM=bpm,
+    )
+    return Q, S
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/swiglu.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/swiglu.py
new file mode 100644
index 0000000000..daabfc6e3a
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/cutlass_fp4/swiglu.py
@@ -0,0 +1,148 @@
+"""Fused gated-activation forward + backward Triton kernels (memory-lean, no [Mt,I] intermediates).
+
+Two activation variants:
+  silu  (default, DSV4): fwd = silu(clamp(g,max=L)) * clamp(u,-L,L)
+  gelu_tanh (Gemma4):    fwd = gelu_pytorch_tanh(g) * u  (no clamp)
+
+bwd: (gu[Mt,2I], dh[Mt,I]) -> dgu[Mt,2I] in ONE pass (~8 intermediates become register-local).
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+
+@triton.jit
+def _fwd(GU, H, I, L, sg0, sg1, sh0, sh1, BLK: tl.constexpr):
+    r = tl.program_id(0)
+    c = tl.program_id(1) * BLK + tl.arange(0, BLK)
+    m = c < I
+    g = tl.load(GU + r * sg0 + c * sg1, mask=m, other=0.0).to(tl.float32)
+    u = tl.load(GU + r * sg0 + (I + c) * sg1, mask=m, other=0.0).to(tl.float32)
+    gc = tl.minimum(g, L)
+    uc = tl.minimum(tl.maximum(u, -L), L)
+    h = (gc * tl.sigmoid(gc)) * uc
+    tl.store(H + r * sh0 + c * sh1, h.to(H.dtype.element_ty), mask=m)
+
+
+@triton.jit
+def _bwd(GU, DH, DGU, I, L, sg0, sg1, sd0, sd1, so0, so1, BLK: tl.constexpr):
+    r = tl.program_id(0)
+    c = tl.program_id(1) * BLK + tl.arange(0, BLK)
+    m = c < I
+    g = tl.load(GU + r * sg0 + c * sg1, mask=m, other=0.0).to(tl.float32)
+    u = tl.load(GU + r * sg0 + (I + c) * sg1, mask=m, other=0.0).to(tl.float32)
+    dh = tl.load(DH + r * sd0 + c * sd1, mask=m, other=0.0).to(tl.float32)
+    gc = tl.minimum(g, L)
+    sg = tl.sigmoid(gc)
+    silu = gc * sg
+    dsilu = sg * (1.0 + gc * (1.0 - sg))
+    uc = tl.minimum(tl.maximum(u, -L), L)
+    dg = dh * uc * dsilu * (g <= L)
+    du = dh * silu * ((u >= -L) & (u <= L))
+    tl.store(DGU + r * so0 + c * so1, dg.to(DGU.dtype.element_ty), mask=m)
+    tl.store(DGU + r * so0 + (I + c) * so1, du.to(DGU.dtype.element_ty), mask=m)
+
+
+# gelu_tanh(x) = 0.5 * x * (1 + tanh(k * (x + 0.044715 * x^3))), k = sqrt(2/pi)
+# gelu_tanh'(x) = 0.5*(1+t) + 0.5*x*(1-t^2)*k*(1 + 3*0.044715*x^2)
+
+
+@triton.jit
+def _fwd_gelu(GU, H, I, sg0, sg1, sh0, sh1, BLK: tl.constexpr):
+    r = tl.program_id(0)
+    c = tl.program_id(1) * BLK + tl.arange(0, BLK)
+    m = c < I
+    g = tl.load(GU + r * sg0 + c * sg1, mask=m, other=0.0).to(tl.float32)
+    u = tl.load(GU + r * sg0 + (I + c) * sg1, mask=m, other=0.0).to(tl.float32)
+    k = 0.7978845608028654  # sqrt(2/pi)
+    inner = k * (g + 0.044715 * g * g * g)
+    t = libdevice.tanh(inner)
+    gelu_g = 0.5 * g * (1.0 + t)
+    h = gelu_g * u
+    tl.store(H + r * sh0 + c * sh1, h.to(H.dtype.element_ty), mask=m)
+
+
+@triton.jit
+def _bwd_gelu(GU, DH, DGU, I, sg0, sg1, sd0, sd1, so0, so1, BLK: tl.constexpr):
+    r = tl.program_id(0)
+    c = tl.program_id(1) * BLK + tl.arange(0, BLK)
+    m = c < I
+    g = tl.load(GU + r * sg0 + c * sg1, mask=m, other=0.0).to(tl.float32)
+    u = tl.load(GU + r * sg0 + (I + c) * sg1, mask=m, other=0.0).to(tl.float32)
+    dh = tl.load(DH + r * sd0 + c * sd1, mask=m, other=0.0).to(tl.float32)
+    k = 0.7978845608028654
+    g2 = g * g
+    inner = k * (g + 0.044715 * g2 * g)
+    t = libdevice.tanh(inner)
+    gelu_g = 0.5 * g * (1.0 + t)
+    # gelu'(g) = 0.5*(1+t) + 0.5*g*(1-t^2)*k*(1 + 3*0.044715*g^2)
+    dgelu = 0.5 * (1.0 + t) + 0.5 * g * (1.0 - t * t) * k * (1.0 + 3.0 * 0.044715 * g2)
+    dg = dh * u * dgelu
+    du = dh * gelu_g
+    tl.store(DGU + r * so0 + c * so1, dg.to(DGU.dtype.element_ty), mask=m)
+    tl.store(DGU + r * so0 + (I + c) * so1, du.to(DGU.dtype.element_ty), mask=m)
+
+
+def swiglu_fwd(gu, limit, act_type="silu", bpm=512):
+    """Gated-activation forward: gu[Mt,2I] -> h[Mt,I].
+    act_type='silu': clamped SwiGLU (DSV4). act_type='gelu_tanh': GeGLU (Gemma4, limit ignored)."""
+    Mt, twoI = gu.shape
+    I = twoI // 2
+    h = torch.empty(Mt, I, device=gu.device, dtype=gu.dtype)
+    if act_type == "gelu_tanh":
+        _fwd_gelu[(Mt, triton.cdiv(I, bpm))](
+            gu, h, I, gu.stride(0), gu.stride(1), h.stride(0), h.stride(1), BLK=bpm
+        )
+    else:
+        _fwd[(Mt, triton.cdiv(I, bpm))](
+            gu,
+            h,
+            I,
+            float(limit),
+            gu.stride(0),
+            gu.stride(1),
+            h.stride(0),
+            h.stride(1),
+            BLK=bpm,
+        )
+    return h
+
+
+def swiglu_bwd(gu, dh, limit, act_type="silu", bpm=512):
+    """Gated-activation backward: (gu[Mt,2I], dh[Mt,I]) -> dgu[Mt,2I].
+    act_type='silu': clamped SwiGLU (DSV4). act_type='gelu_tanh': GeGLU (Gemma4, limit ignored)."""
+    Mt, twoI = gu.shape
+    I = twoI // 2
+    dgu = torch.empty(Mt, twoI, device=gu.device, dtype=gu.dtype)
+    if act_type == "gelu_tanh":
+        _bwd_gelu[(Mt, triton.cdiv(I, bpm))](
+            gu,
+            dh,
+            dgu,
+            I,
+            gu.stride(0),
+            gu.stride(1),
+            dh.stride(0),
+            dh.stride(1),
+            dgu.stride(0),
+            dgu.stride(1),
+            BLK=bpm,
+        )
+    else:
+        _bwd[(Mt, triton.cdiv(I, bpm))](
+            gu,
+            dh,
+            dgu,
+            I,
+            float(limit),
+            gu.stride(0),
+            gu.stride(1),
+            dh.stride(0),
+            dh.stride(1),
+            dgu.stride(0),
+            dgu.stride(1),
+            BLK=bpm,
+        )
+    return dgu
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/dequant_grouped.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/dequant_grouped.py
new file mode 100644
index 0000000000..75afa0d5fa
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/dequant_grouped.py
@@ -0,0 +1,462 @@
+"""Dequant + grouped-GEMM MoE experts path — the fast alternative to the fully-fused
+``scatter2scatter_lora_mx`` Triton kernel.
+
+Benchmarks (DSV4-Flash expert shapes) showed the fused in-kernel-decode Triton kernel is
+~4x slower than dequantizing NVFP4->bf16 and running a cuBLAS-grade grouped GEMM, because a
+hand-written Triton MoE GEMM can't match cuBLAS/tensor-core throughput. This module provides:
+
+  * primary (SM90/SM100 + DeepGEMM): native fused-decode fp4 grouped GEMM (no bf16 weight
+    materialization, tensor-core speed) — ``deepgemm_grouped_available()`` gates it;
+  * fallback (any other GPU, e.g. SM120): a fast custom-Triton NVFP4->bf16 dequant (~35x
+    faster than torchao's eager ``.dequantize()``) + ``torch._grouped_mm``, tiled per
+    expert-chunk + checkpointed so the bf16 weight transient stays bounded.
+
+``torch._grouped_mm`` has no working autograd (stride bug), so ``_GroupedMM`` supplies the
+backward manually (``dX = grouped_mm(grad, Wᵀ)``; experts are frozen → no weight grad).
+
+STATUS: the chunked dequant+grouped BASE path is implemented + parity-validated here. Wiring
+into ``parallel_linear_lora`` (LoRA-on-experts + the DeepGEMM primary call) is the remaining
+integration step.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from torch.utils.checkpoint import checkpoint
+
+from .mx_weights import fp4_codebook
+
+_DG = None
+
+
+def _dg():
+    """Cached DeepGEMM kernel module. ``get_kernel`` is not idempotent — calling it twice
+    re-runs the build's ``register_fake`` and raises, so resolve it once and reuse."""
+    global _DG
+    if _DG is None:
+        from kernels import get_kernel
+
+        _DG = get_kernel("kernels-community/deep-gemm")
+    return _DG
+
+
+def deepgemm_grouped_available() -> bool:
+    """True iff DeepGEMM's grouped fp8xfp4 MoE kernel can run here.
+
+    SM100 ONLY (Blackwell datacenter, e.g. B200). The kernel symbol resolves on SM90 (Hopper) too,
+    but ``m_grouped_fp8_fp4_gemm_nt_contiguous`` asserts ``arch_major == 10`` at call time, so SM90
+    must report unavailable and fall back to Marlin (which supports SM80-90,120). Gating on the
+    capability up front avoids a hard CUDA assert mid-forward."""
+    if not torch.cuda.is_available():
+        return False
+    major, _minor = torch.cuda.get_device_capability()
+    if major != 10:  # fp8xfp4 grouped kernel is sm100-only
+        return False
+    try:
+        return getattr(_dg(), "m_grouped_fp8_fp4_gemm_nt_contiguous", None) is not None
+    except Exception:
+        return False
+
+
+@triton.jit
+def _nvfp4_deq_kernel(
+    Qp, Sc, Pt, Out, Kd, ROWS_PER_E, CB, sq0, sq1, ss0, ss1, so0, so1, BK: tl.constexpr
+):
+    rid = tl.program_id(0).to(tl.int64)  # c*nrows*stride overflows int32 at E=256
+    kk = tl.program_id(1) * BK + tl.arange(0, BK)
+    km = kk < Kd
+    pt = tl.load(Pt + rid // ROWS_PER_E).to(tl.float32)
+    packed = tl.load(Qp + rid * sq0 + (kk // 2) * sq1, mask=km, other=0).to(tl.int32)
+    nib = tl.where((kk % 2) == 1, (packed >> 4) & 0xF, packed & 0xF)
+    sc = tl.load(Sc + rid * ss0 + (kk // 16) * ss1, mask=km, other=0.0).to(tl.float32)
+    tl.store(
+        Out + rid * so0 + kk * so1,
+        (tl.load(CB + nib) * sc * pt).to(tl.bfloat16),
+        mask=km,
+    )
+
+
+def nvfp4_dequant_bf16(
+    qdata: torch.Tensor, scale: torch.Tensor, per_tensor: torch.Tensor
+) -> torch.Tensor:
+    """NVFP4 -> bf16 via a memory-bound Triton kernel (~35x faster than torchao eager).
+    ``qdata`` [c,N,K/2] uint8, ``scale`` [c,N,K/16] e4m3, ``per_tensor`` [c] fp32 → [c,N,K] bf16."""
+    c, nrows, kh = qdata.shape
+    kd = kh * 2
+    out = torch.empty(c, nrows, kd, device=qdata.device, dtype=torch.bfloat16)
+    q2, s2, o2 = (
+        qdata.reshape(c * nrows, kh),
+        scale.reshape(c * nrows, kd // 16),
+        out.reshape(c * nrows, kd),
+    )
+    BK = 1024  # larger blocks lift HBM utilization (memory-bound kernel)
+    _nvfp4_deq_kernel[(c * nrows, triton.cdiv(kd, BK))](
+        q2,
+        s2,
+        per_tensor.contiguous(),
+        o2,
+        kd,
+        nrows,
+        fp4_codebook(qdata.device),
+        q2.stride(0),
+        q2.stride(1),
+        s2.stride(0),
+        s2.stride(1),
+        o2.stride(0),
+        o2.stride(1),
+        BK=BK,
+    )
+    return out
+
+
+@triton.jit
+def _fp8_e8m0_q_kernel(X, O, SF, sx0, sx1, so0, so1, ss0, ss1, G: tl.constexpr):
+    row = tl.program_id(0)
+    kb = tl.program_id(1)
+    offs = kb * G + tl.arange(0, G)
+    x = tl.load(X + row * sx0 + offs * sx1).to(tl.float32)
+    amax = tl.maximum(tl.max(tl.abs(x)), 1e-4)
+    sf = tl.exp2(
+        tl.ceil(tl.log2(amax / 448.0))
+    )  # amax/448 rounded up to E8M0 (power of 2)
+    q = (x / sf).to(tl.float8e4nv)
+    tl.store(O + row * so0 + offs * so1, q)
+    tl.store(SF + row * ss0 + kb * ss1, sf)
+
+
+def fp8_e8m0_cast_128(x: torch.Tensor):
+    """Fast Triton replica of DeepGEMM's ``per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=128)``
+    (exact bit-match, ~4x faster). x[M,K] bf16 -> (fp8_e4m3 [M,K], scale fp32 [M,K/128])."""
+    M, K = x.shape
+    assert K % 128 == 0
+    o = torch.empty(M, K, device=x.device, dtype=torch.float8_e4m3fn)
+    sf = torch.empty(M, K // 128, device=x.device, dtype=torch.float32)
+    _fp8_e8m0_q_kernel[(M, K // 128)](
+        x,
+        o,
+        sf,
+        x.stride(0),
+        x.stride(1),
+        o.stride(0),
+        o.stride(1),
+        sf.stride(0),
+        sf.stride(1),
+        G=128,
+    )
+    return o, sf
+
+
+def nvfp4_to_mxfp4_weight(
+    qdata: torch.Tensor, scale: torch.Tensor, per_tensor: torch.Tensor, chunk: int = 16
+):
+    """One-time per-expert weight requant NVFP4 (E4M3/16) -> MXFP4 (E8M0/128) for DeepGEMM's
+    grouped fp8xfp4 kernel (which only accepts E8M0/128). qdata [E,N,K/2], scale [E,N,K/16],
+    per_tensor [E] -> (wq [E,N,K/2] uint8, ws [E,N,K/128] fp32).
+
+    Dequant + cast in expert CHUNKS into preallocated outputs so the bf16 intermediate is bounded to
+    ``chunk`` experts (the full [E,N,K] bf16 was ~E/chunk x larger and OOMed at E=256). The output is
+    full-size by construction (the grouped GEMM needs all experts present)."""
+    cast = _dg().utils.per_token_cast_to_fp4
+    E = int(qdata.size(0))
+    q0, s0 = cast(
+        nvfp4_dequant_bf16(qdata[:1], scale[:1], per_tensor[:1])[0].contiguous(),
+        True,
+        128,
+    )
+    wq = q0.new_empty((E, *q0.shape))
+    ws = s0.new_empty((E, *s0.shape))
+    for c0 in range(0, E, chunk):
+        c1 = min(c0 + chunk, E)
+        Wb = nvfp4_dequant_bf16(qdata[c0:c1], scale[c0:c1], per_tensor[c0:c1])
+        for i in range(c1 - c0):
+            q, s = cast(Wb[i].contiguous(), True, 128)
+            wq[c0 + i] = q
+            ws[c0 + i] = s
+        del Wb
+    return wq, ws
+
+
+def deepgemm_grouped_fp8_fp4(
+    a_bf16: torch.Tensor, wq: torch.Tensor, ws: torch.Tensor, m_indices: torch.Tensor
+) -> torch.Tensor:
+    """Contiguous-grouped fp8(act) x mxfp4(weight) base GEMM via DeepGEMM (SM90/SM100).
+    a_bf16 [Mt,K], wq [E,N,K/2] uint8, ws [E,N,K/128] fp32, m_indices [Mt] int32 -> [Mt,N] bf16.
+    Acts quantized to fp8 (E8M0/128) here; the kernel transforms the raw scales internally."""
+    dg = _dg()
+    a = fp8_e8m0_cast_128(a_bf16)
+    d = torch.empty(
+        a_bf16.size(0), wq.size(1), device=a_bf16.device, dtype=torch.bfloat16
+    )
+    dg.m_grouped_fp8_fp4_gemm_nt_contiguous(
+        a,
+        (wq, ws),
+        d,
+        m_indices,
+        recipe=None,
+        recipe_a=(1, 128),
+        recipe_b=(1, 128),
+        disable_ue8m0_cast=False,
+    )
+    return d
+
+
+def _cached_mxfp4(w_nv, per_tensor, cache=None, key=None):
+    """Return (wq, ws) MXFP4 weight for a (frozen) NVFP4Tensor, requantized once and cached.
+    Prefer a persistent ``cache`` dict (e.g. on the owning module) keyed by ``key`` — under
+    FSDP2 the gathered param is a fresh tensor each step, so a module-level cache (not a
+    per-tensor attribute) is what avoids recomputing the requant every forward. Falls back to a
+    per-tensor attribute, then to a one-shot recompute."""
+    from .runtime import RUNTIME
+
+    if not RUNTIME.mxfp4_cache_persist:
+        # FSDP: any attached cache accumulates a full-model mxfp4 copy across layers -> OOM.
+        # Recompute uncached so resident mxfp4 stays bounded to one layer.
+        return nvfp4_to_mxfp4_weight(w_nv.qdata, w_nv.scale, per_tensor)
+    if cache is not None and key is not None and key in cache:
+        return cache[key]
+    cached = getattr(w_nv, "_dg_mxfp4", None)
+    if cached is None:
+        cached = nvfp4_to_mxfp4_weight(w_nv.qdata, w_nv.scale, per_tensor)
+        try:
+            w_nv._dg_mxfp4 = cached
+        except (AttributeError, RuntimeError):
+            pass
+    if cache is not None and key is not None:
+        cache[key] = cached
+    return cached
+
+
+@triton.jit
+def _nvfp4_deq_fp8_kernel(
+    Qp, Sc, Pt, Out, Kd, ROWS_PER_E, CB, sq0, sq1, ss0, ss1, so0, so1, BK: tl.constexpr
+):
+    rid = tl.program_id(0).to(tl.int64)
+    kk = tl.program_id(1) * BK + tl.arange(0, BK)
+    km = kk < Kd
+    pt = tl.load(Pt + rid // ROWS_PER_E).to(tl.float32)
+    packed = tl.load(Qp + rid * sq0 + (kk // 2) * sq1, mask=km, other=0).to(tl.int32)
+    nib = tl.where((kk % 2) == 1, (packed >> 4) & 0xF, packed & 0xF)
+    sc = tl.load(Sc + rid * ss0 + (kk // 16) * ss1, mask=km, other=0.0).to(tl.float32)
+    tl.store(
+        Out + rid * so0 + kk * so1,
+        (tl.load(CB + nib) * sc * pt).to(tl.float8e4nv),
+        mask=km,
+    )
+
+
+def nvfp4_dequant_fp8(
+    qdata: torch.Tensor, scale: torch.Tensor, per_tensor: torch.Tensor
+) -> torch.Tensor:
+    """NVFP4 -> fp8 (e4m3) — half the bytes of the bf16 dequant. For the fp8-read backward dX
+    (#3744): the grouped GEMM reads fp8 and upcasts to bf16 in-register, halving the weight's
+    write+read bandwidth (a win on bandwidth-bound sm120; ~neutral speed but still half-memory
+    on sm100)."""
+    c, nrows, kh = qdata.shape
+    kd = kh * 2
+    out = torch.empty(c, nrows, kd, device=qdata.device, dtype=torch.float8_e4m3fn)
+    q2, s2, o2 = (
+        qdata.reshape(c * nrows, kh),
+        scale.reshape(c * nrows, kd // 16),
+        out.reshape(c * nrows, kd),
+    )
+    BK = 1024
+    _nvfp4_deq_fp8_kernel[(c * nrows, triton.cdiv(kd, BK))](
+        q2,
+        s2,
+        per_tensor.contiguous(),
+        o2,
+        kd,
+        nrows,
+        fp4_codebook(qdata.device),
+        q2.stride(0),
+        q2.stride(1),
+        s2.stride(0),
+        s2.stride(1),
+        o2.stride(0),
+        o2.stride(1),
+        BK=BK,
+    )
+    return out
+
+
+# BM is the routing pad TILE (128 cutlass, 64 marlin); a constexpr autotune key so BN/BK/warps
+# still tune per (N,K,BM).
+@triton.autotune(
+    configs=[
+        triton.Config({"BN": bn, "BK": bk}, num_warps=w, num_stages=st)
+        for bn in (64, 128)
+        for bk in (128, 256)
+        for w in (4, 8)
+        for st in (3, 4)
+    ],
+    key=["N", "K", "BM"],
+)
+@triton.jit
+def _grouped_dx_fp8_kernel(
+    GRAD,
+    W,
+    MIDX,
+    OUT,
+    M,
+    N,
+    K,
+    sg0,
+    sg1,
+    sw0,
+    sw1,
+    sw2,
+    so0,
+    so1,
+    BM: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    e = tl.load(MIDX + pid_m).to(tl.int64)
+    rm = pid_m * BM + tl.arange(0, BM)
+    rk = pid_k * BK + tl.arange(0, BK)
+    mk = rk < K
+    acc = tl.zeros((BM, BK), tl.float32)
+    for n0 in range(0, N, BN):
+        rn = n0 + tl.arange(0, BN)
+        a = tl.load(
+            GRAD + rm[:, None] * sg0 + rn[None, :] * sg1,
+            mask=rm[:, None] < M,
+            other=0.0,
+        )
+        # mask K dim to handle non-BK-aligned K (e.g. I=704 which is not a multiple of 128)
+        w = tl.load(
+            W + e * sw0 + rn[:, None] * sw1 + rk[None, :] * sw2,
+            mask=mk[None, :],
+            other=0.0,
+        ).to(tl.bfloat16)
+        acc += tl.dot(a, w)
+    tl.store(
+        OUT + rm[:, None] * so0 + rk[None, :] * so1,
+        acc.to(tl.bfloat16),
+        mask=(rm[:, None] < M) & mk[None, :],
+    )
+
+
+def grouped_dx_fp8(
+    grad: torch.Tensor, w_fp8: torch.Tensor, m_indices: torch.Tensor, block_m: int = 128
+) -> torch.Tensor:
+    """dX = grad @ W (contract N) for contiguous-grouped experts, reading the fp8 weight and
+    upcasting in-register. grad[Mt,N] bf16, w_fp8[E,N,K] e4m3, m_indices[Mt/block_m] (one expert id
+    per block_m-row tile; block_m = the routing pad TILE: 128 cutlass, 64 marlin) -> [Mt,K] bf16."""
+    Mt, N = grad.shape
+    K = w_fp8.size(2)
+    out = torch.empty(Mt, K, device=grad.device, dtype=torch.bfloat16)
+    grid = lambda meta: (Mt // block_m, triton.cdiv(K, meta["BK"]))  # noqa: E731
+    _grouped_dx_fp8_kernel[grid](
+        grad,
+        w_fp8,
+        m_indices,
+        out,
+        Mt,
+        N,
+        K,
+        grad.stride(0),
+        grad.stride(1),
+        w_fp8.stride(0),
+        w_fp8.stride(1),
+        w_fp8.stride(2),
+        out.stride(0),
+        out.stride(1),
+        BM=block_m,
+    )
+    return out
+
+
+class _GroupedMM(torch.autograd.Function):
+    """``torch._grouped_mm`` with a manual backward (its autograd is stride-broken). Frozen
+    weight ``bT`` (the dequantized experts) → only the input grad ``dX`` is returned."""
+
+    @staticmethod
+    def forward(ctx, a, bT, offs):
+        ctx.save_for_backward(bT, offs)
+        return torch._grouped_mm(a, bT, offs=offs)
+
+    @staticmethod
+    def backward(ctx, g):
+        bT, offs = ctx.saved_tensors
+        dA = torch._grouped_mm(
+            g.contiguous(), bT.transpose(1, 2).contiguous(), offs=offs
+        )
+        return dA, None, None
+
+
+def _chunk_forward(sorted_tok, gqd, gsc, dqd, dsc, pt, coff, limit):
+    """One expert-chunk: dequant -> grouped gate_up -> clamped-SwiGLU -> grouped down."""
+    gub = nvfp4_dequant_bf16(gqd, gsc, pt).transpose(1, 2)
+    dnb = nvfp4_dequant_bf16(dqd, dsc, pt).transpose(1, 2)
+    x = _GroupedMM.apply(sorted_tok, gub, coff)
+    g, u = x.chunk(2, dim=-1)
+    h = (
+        torch.nn.functional.silu(g.clamp(max=limit)) * u.clamp(min=-limit, max=limit)
+    ).contiguous()
+    return _GroupedMM.apply(h, dnb, coff)
+
+
+def chunked_dequant_grouped_base(
+    hidden, idx, wts, gate_up_nv, down_nv, per_tensor, limit, chunk_size=8
+):
+    """Base clamped-SwiGLU MoE on NVFP4 experts via chunked dequant + grouped_mm.
+
+    hidden [N,H]; idx [N,topk]; wts [N,topk]; gate_up_nv/down_nv: NVFP4Tensor
+    ([E,2I,H]/[E,H,I]); per_tensor: scalar or [E] fp32. Returns [N,H]. Differentiable to
+    ``hidden``. Per-chunk ``checkpoint`` keeps the bf16 weight transient bounded (re-dequant
+    in backward).
+    """
+    E = gate_up_nv.qdata.size(0)
+    Hdim = hidden.size(1)
+    # raw-row-major scale indexing only; deswizzling is not implemented here
+    assert not getattr(gate_up_nv, "is_swizzled_scales", False), (
+        "swizzled NVFP4 scales unsupported"
+    )
+    assert not getattr(down_nv, "is_swizzled_scales", False), (
+        "swizzled NVFP4 scales unsupported"
+    )
+    # per_tensor_scale may be a global scalar or per-expert; normalize to [E]
+    per_tensor = per_tensor.reshape(-1).to(torch.float32)
+    if per_tensor.numel() == 1:
+        per_tensor = per_tensor.expand(E)
+    flat_exp = idx.reshape(-1)
+    order = flat_exp.argsort()
+    rep = torch.arange(hidden.size(0), device=hidden.device).repeat_interleave(
+        idx.size(1)
+    )[order]
+    wflat = wts.reshape(-1)[order]
+    counts = torch.bincount(flat_exp, minlength=E)
+    tok_off = torch.cat([counts.new_zeros(1), counts.cumsum(0)]).tolist()
+    flat = hidden[rep]
+
+    out = hidden.new_zeros(hidden.size(0), Hdim)
+    pieces = []
+    gq, gs = gate_up_nv.qdata, gate_up_nv.scale
+    dq, ds = down_nv.qdata, down_nv.scale
+    for c0 in range(0, E, chunk_size):
+        c1 = min(c0 + chunk_size, E)
+        t0, t1 = int(tok_off[c0]), int(tok_off[c1])
+        if t1 == t0:
+            continue
+        coff = counts[c0:c1].cumsum(0).to(torch.int32)
+        o = checkpoint(
+            _chunk_forward,
+            flat[t0:t1],
+            gq[c0:c1],
+            gs[c0:c1],
+            dq[c0:c1],
+            ds[c0:c1],
+            per_tensor[c0:c1],
+            coff,
+            limit,
+            use_reentrant=False,
+        )
+        pieces.append(o * wflat[t0:t1, None].to(o.dtype))
+    if not pieces:  # no routed tokens (e.g. empty batch / expert-parallel shard)
+        return out
+    return out.index_add(0, rep, torch.cat(pieces, 0))
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/experts.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/experts.py
new file mode 100644
index 0000000000..9276dcd123
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/experts.py
@@ -0,0 +1,924 @@
+"""ScatterMoE experts forward for the transformers ExpertsInterface.
+
+PEFT LoRA on ``gate_up_proj`` / ``down_proj`` is fused into the
+ScatterMoE Triton call via ``parallel_linear_lora``.
+"""
+
+import functools
+
+import torch
+import torch.nn.functional as F
+
+from .mx_weights import selective_mx_weights_fwd, selective_nvfp4_weights_fwd
+from .parallel_experts import flatten_sort_count, parallel_linear
+from .parallel_linear_lora import get_lora_params_from_wrapper, parallel_linear_lora
+from .selective_dequant import (
+    get_active_experts,
+    is_mxfp4_param,
+    is_nvfp4_param,
+    remap_expert_indices,
+    selective_lora_weights,
+)
+
+
+def _has_peft_wrapper(module):
+    """Check if a module's parameter has been wrapped by PEFT ParamWrapper."""
+    try:
+        from peft.tuners.param_wrapper import ParamWrapper
+
+        for attr in ("gate_up_proj", "down_proj"):
+            param = getattr(module, attr, None)
+            if isinstance(param, ParamWrapper):
+                return True
+    except ImportError:
+        pass
+    return False
+
+
+def _unwrap_experts_lora(experts):
+    """Extract base weights and LoRA params from a PEFT-wrapped Experts module.
+
+    Returns:
+        (base_experts, gup_lora, down_lora) where each lora is
+        (lora_A, lora_B, scaling) or None.
+    """
+    try:
+        from peft.tuners.param_wrapper import ParamWrapper
+    except ImportError:
+        return experts, None, None
+
+    if not isinstance(getattr(experts, "gate_up_proj", None), ParamWrapper):
+        return experts, None, None
+
+    base_experts = experts
+    gup_lora = None
+    down_lora = None
+
+    for param, which in ((experts.gate_up_proj, "gup"), (experts.down_proj, "down")):
+        if not isinstance(param, ParamWrapper):
+            continue
+        lora_A, lora_B, scaling = get_lora_params_from_wrapper(param)
+        if lora_A is None:
+            continue
+        from .layers import peft_lora_to_scattermoe
+
+        a, b, n_local, rank = _ep_local_expert_lora(lora_A, lora_B, experts)
+        sm_A, sm_B = peft_lora_to_scattermoe(a, b, n_local, rank)
+        if which == "gup":
+            gup_lora = (sm_A, sm_B, scaling)
+        else:
+            down_lora = (sm_A, sm_B, scaling)
+
+    return base_experts, gup_lora, down_lora
+
+
+def _ep_local_expert_lora(lora_A, lora_B, experts):
+    """Slice the routed-expert LoRA down to THIS EP rank's local experts.
+
+    Under expert parallelism the base weights are EP-sharded to ``num_experts`` (= E_local) but the
+    PEFT adapter param stays a single GLOBAL tensor over all ``num_experts_global`` experts (it is
+    FSDP-managed/gathered to full here — keeping it global avoids a plain-vs-DTensor mismatch in the
+    optimizer/grad-clip). So compute the rank from the GLOBAL expert count and take this rank's
+    ``[offset : offset+E_local]`` expert block at forward time:
+      * lora_A ``[r*E_global, in]`` expert-major  -> rows ``[offset*r : (offset+E_local)*r]``
+      * lora_B ``[out, r*E_global]`` rank-major    -> reshape ``[out, r, E_global]``, take experts.
+    Non-EP modules (E_local == E_global) are a no-op. Returns ``(A, B, E_local, rank)``."""
+    e_local = experts.num_experts
+    e_global = getattr(experts, "num_experts_global", e_local)
+    rank = lora_A.shape[0] // e_global
+    if e_global == e_local:
+        return lora_A, lora_B, e_local, rank
+    offset = getattr(experts, "local_expert_offset", 0)
+    a = lora_A[offset * rank : (offset + e_local) * rank, :]
+    out = lora_B.shape[0]
+    b = lora_B.reshape(out, rank, e_global)[:, :, offset : offset + e_local].reshape(
+        out, rank * e_local
+    )
+    return a, b, e_local, rank
+
+
+def _get_base_param(param):
+    """Get the base tensor from a PEFT ParamWrapper or regular Parameter."""
+    try:
+        from peft.tuners.param_wrapper import ParamWrapper
+
+        while isinstance(param, ParamWrapper):
+            param = param.original_parameter
+    except ImportError:
+        pass
+    return param
+
+
+def _parallel_linear_maybe_lora(
+    x,
+    weight,
+    top_k,
+    sorted_expert_idxs,
+    sorted_scattered_idxs,
+    expert_offsets,
+    lora_tuple,
+    grouped_in,
+    grouped_out,
+    gates=None,
+    expert_biases=None,
+):
+    """Call parallel_linear or parallel_linear_lora depending on whether LoRA is active."""
+    if lora_tuple is not None:
+        lora_A, lora_B, scaling = lora_tuple
+        return parallel_linear_lora(
+            x,
+            weight,
+            top_k,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            lora_A,
+            lora_B,
+            scaling,
+            expert_biases=expert_biases,
+            grouped_in=grouped_in,
+            grouped_out=grouped_out,
+            gates=gates,
+        )
+    return parallel_linear(
+        x,
+        weight,
+        top_k,
+        sorted_expert_idxs,
+        sorted_scattered_idxs,
+        expert_offsets,
+        expert_biases=expert_biases,
+        grouped_in=grouped_in,
+        grouped_out=grouped_out,
+        gates=gates,
+    )
+
+
+def _prepare_weights_and_lora(
+    gu_param,
+    dn_param,
+    sorted_expert_idxs,
+    expert_offsets,
+    num_experts,
+    gup_lora,
+    down_lora,
+    dtype,
+    module=None,
+):
+    """Resolve the gate_up / down weights (+ routing + LoRA) for the grouped GEMM.
+
+    For an MXFP4 or NVFP4 base with LoRA on both projections, keep the weights packed
+    (4-bit) and route through the fused kernel: select the active experts and remap the
+    routing / LoRA to that compact set (dequant happens inside the kernel K-loop; NVFP4
+    reuses the MXWeights container with linear E4M3 block scales). Otherwise return bf16
+    ``[E, in, out]`` weights (dequantizing FP4 explicitly when LoRA is absent, since the
+    fused kernel is LoRA-only and the FP4 tensors have no transpose).
+
+    Returns ``(gate_up_weight, down_weight, sorted_expert_idxs, expert_offsets,
+    gup_lora, down_lora)``.
+    """
+    is_mx, is_nv = is_mxfp4_param(gu_param), is_nvfp4_param(gu_param)
+    if (is_mx or is_nv) and gup_lora is not None and down_lora is not None:
+        select = selective_mx_weights_fwd if is_mx else selective_nvfp4_weights_fwd
+        active = get_active_experts(sorted_expert_idxs, num_experts)
+        sorted_expert_idxs, expert_offsets = remap_expert_indices(
+            sorted_expert_idxs, expert_offsets, active, num_experts
+        )
+        gate_up_weight = select(gu_param, active)
+        down_weight = select(dn_param, active)
+        # Recompute-in-backward recipe: the selective gather is a per-layer copy (~2 GB/layer at
+        # dense routing) that ScatterMoELoRA would otherwise pin on ctx across all layers (peak
+        # grows with depth). It is reconstructible from the frozen resident param, so hand the
+        # Function a recipe to rebuild it in backward (cheap to recompute, expensive to store).
+        # FSDP2-safe: re-read the LIVE module param at backward (via `module`), not the forward-time
+        # object: FSDP reshards after forward and re-gathers for backward, so a captured reference
+        # would be the half-size shard (wrong dX shape). On single-GPU it's the same param.
+        if module is not None:
+            gate_up_weight.recipe = lambda m=module, a=active, f=select: f(
+                _get_base_param(m.gate_up_proj), a
+            )
+            down_weight.recipe = lambda m=module, a=active, f=select: f(
+                _get_base_param(m.down_proj), a
+            )
+        else:
+            gate_up_weight.recipe = lambda p=gu_param, a=active, f=select: f(p, a)
+            down_weight.recipe = lambda p=dn_param, a=active, f=select: f(p, a)
+        gup_lora = (
+            *selective_lora_weights(gup_lora[0], gup_lora[1], active, num_experts),
+            gup_lora[2],
+        )
+        down_lora = (
+            *selective_lora_weights(down_lora[0], down_lora[1], active, num_experts),
+            down_lora[2],
+        )
+    elif is_mx or is_nv:
+        # quantized base without LoRA: the fused kernel is LoRA-only and the FP4 tensors have no
+        # transpose, so dequantize to bf16 here.
+        gate_up_weight = gu_param.dequantize(dtype).transpose(2, 1)
+        down_weight = dn_param.dequantize(dtype).transpose(2, 1)
+    else:
+        gate_up_weight = gu_param.transpose(2, 1)
+        down_weight = dn_param.transpose(2, 1)
+    return (
+        gate_up_weight,
+        down_weight,
+        sorted_expert_idxs,
+        expert_offsets,
+        gup_lora,
+        down_lora,
+    )
+
+
+def _detect_act_type(module) -> str:
+    """Detect gated-activation type from an experts module's act_fn.
+
+    Returns 'gelu_tanh' for Gemma4-style GeGLU (gelu_pytorch_tanh * up),
+    'silu' for DSV4-style clamped SwiGLU (silu(clamp(gate)) * clamp(up)).
+    Falls back to 'silu' for any unrecognized activation.
+    """
+    act_fn = getattr(module, "act_fn", None)
+    if act_fn is None:
+        return "silu"
+    fn_name = (
+        getattr(act_fn, "__name__", "") or getattr(type(act_fn), "__name__", "") or ""
+    )
+    if "gelu" in fn_name.lower():
+        return "gelu_tanh"
+    try:
+        if isinstance(act_fn, functools.partial) and act_fn.func is F.gelu:
+            return "gelu_tanh"
+    except Exception:
+        pass
+    # last resort: compare act_fn output to gelu_pytorch_tanh numerically
+    try:
+        x = torch.tensor([0.5], dtype=torch.float32, device="cpu")
+        ref = F.gelu(x, approximate="tanh")
+        got = act_fn(x.clone())
+        if torch.allclose(ref, got, atol=1e-5):
+            return "gelu_tanh"
+    except Exception:
+        pass
+    return "silu"
+
+
+def scattermoe_supports_layout(self) -> bool:
+    """True iff this experts module uses the standard layout scattermoe handles:
+    gate_up concatenated as [E, 2I, H], gated SwiGLU, no expert bias. gpt_oss-style
+    experts (interleaved gate/up, transposed [E, H, 2I], expert bias) return False."""
+    return not (
+        getattr(self, "is_transposed", False)
+        or not getattr(self, "is_concatenated", True)
+        or getattr(self, "has_bias", False)
+        or not getattr(self, "has_gate", True)
+    )
+
+
+def _check_supported_layout(self):
+    """Reject expert layouts the fixed transpose/chunk below would miscompute."""
+    if not scattermoe_supports_layout(self):
+        raise NotImplementedError(
+            "scattermoe supports only concatenated, non-transposed, gated, biasless "
+            "experts (qwen/mixtral/deepseek/glm/...). This model's experts use an "
+            "unsupported layout; use use_sonicmoe or a built-in experts_implementation."
+        )
+
+
+def _is_gptoss_layout(self) -> bool:
+    """gpt_oss expert layout: transposed [E, H, 2I] weights, interleaved gate/up,
+    per-expert bias, clamped sigmoid-GLU. Handled by the Triton path below."""
+    return (
+        getattr(self, "is_transposed", False)
+        and not getattr(self, "is_concatenated", True)
+        and getattr(self, "has_bias", False)
+        and getattr(self, "has_gate", True)
+        and hasattr(self, "gate_up_proj_bias")
+        and hasattr(self, "down_proj_bias")
+    )
+
+
+def _gptoss_glu(gate_up: torch.Tensor, alpha: float, limit: float) -> torch.Tensor:
+    """gpt_oss clamped sigmoid-GLU over interleaved gate/up columns (matches
+    ``GptOssExperts._apply_gate``)."""
+    gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+    gate = gate.clamp(max=limit)
+    up = up.clamp(min=-limit, max=limit)
+    glu = gate * torch.sigmoid(gate * alpha)
+    return (up + 1) * glu
+
+
+def _scattermoe_gptoss_forward(
+    self,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    """ScatterMoE forward for gpt_oss-style experts (transposed/interleaved/biased).
+
+    gpt_oss stores weights already in ``[E, in, out]`` form, so no transpose; gate/up
+    are interleaved (``[..., ::2]`` / ``[..., 1::2]``); per-expert bias is folded into
+    the grouped GEMM; the activation is the clamped sigmoid-GLU. LoRA fuses exactly as
+    in the standard path (in/out dims match), and the down bias is added per row before
+    the routing-weight combine.
+    """
+    K = top_k_index.shape[1]
+    routing_weights = top_k_weights.to(hidden_states.dtype)
+    sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+        top_k_index, num_experts=self.num_experts
+    )
+
+    gate_up_weight = _get_base_param(
+        self.gate_up_proj
+    )  # already [E, in, out], no transpose
+    down_weight = _get_base_param(self.down_proj)
+    gate_up_bias = _get_base_param(self.gate_up_proj_bias)
+    down_bias = _get_base_param(self.down_proj_bias)
+
+    gup_lora, down_lora = None, None
+    sm_lora = getattr(self, "_scattermoe_lora", None)
+    if sm_lora:
+        gup_lora = sm_lora.get("gate_up_proj")
+        down_lora = sm_lora.get("down_proj")
+    elif _has_peft_wrapper(self):
+        _, gup_lora, down_lora = _unwrap_experts_lora(self)
+
+    gate_up = _parallel_linear_maybe_lora(
+        hidden_states,
+        gate_up_weight,
+        K,
+        sorted_expert_idxs,
+        sorted_scattered_idxs,
+        expert_offsets,
+        gup_lora,
+        grouped_in=False,
+        grouped_out=True,
+        expert_biases=gate_up_bias,
+    )
+    h = _gptoss_glu(gate_up, getattr(self, "alpha", 1.702), getattr(self, "limit", 7.0))
+
+    output = _parallel_linear_maybe_lora(
+        h,
+        down_weight,
+        1,
+        sorted_expert_idxs,
+        sorted_scattered_idxs,
+        expert_offsets,
+        down_lora,
+        grouped_in=True,
+        grouped_out=False,
+        gates=routing_weights,
+        expert_biases=down_bias,
+    )
+    return output
+
+
+# thin compat wrappers over the centralized runtime module (runtime.py)
+from .runtime import RUNTIME  # noqa: E402
+
+
+def set_fp4_grouped_mode(mode):
+    RUNTIME.fp4_grouped_mode = mode
+
+
+def set_fp4_dx_prefer_fp8(flag):
+    RUNTIME.fp4_dx_prefer_fp8 = bool(flag)
+
+
+def scattermoe_experts_forward(
+    self,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    """ScatterMoE experts forward with fused-LoRA support."""
+    if not scattermoe_supports_layout(self):
+        if _is_gptoss_layout(self):
+            return _scattermoe_gptoss_forward(
+                self, hidden_states, top_k_index, top_k_weights
+            )
+        _check_supported_layout(self)  # raises for any other unsupported layout
+
+    K = top_k_index.shape[1]
+
+    routing_weights = top_k_weights.to(hidden_states.dtype)
+    sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+        top_k_index, num_experts=self.num_experts
+    )
+
+    gup_lora, down_lora = None, None
+    sm_lora = getattr(self, "_scattermoe_lora", None)
+    if sm_lora:
+        gup_lora = sm_lora.get("gate_up_proj")
+        down_lora = sm_lora.get("down_proj")
+    elif _has_peft_wrapper(self):
+        _, gup_lora, down_lora = _unwrap_experts_lora(self)
+
+    # Config-gated grouped fp4 path (NVFP4 experts + LoRA + available fp4 backend); faster +
+    # lower-memory than the fused-Triton path at scale.
+    _fp4_grouped_mode = RUNTIME.fp4_grouped_mode
+    if _fp4_grouped_mode is not None and gup_lora is not None and down_lora is not None:
+        gu_base = _get_base_param(self.gate_up_proj)
+        dn_base = _get_base_param(self.down_proj)
+        if is_nvfp4_param(gu_base):
+            from .grouped_train import grouped_fp4_available, grouped_fp4_moe_train
+
+            # No LoRA-capable base GEMM here; the legacy fused-MX kernel SIGSEGVs on Blackwell, so
+            # hard-error rather than crash.
+            if not grouped_fp4_available(_fp4_grouped_mode):
+                _cap = (
+                    "sm%d%d" % torch.cuda.get_device_capability()
+                    if torch.cuda.is_available()
+                    else "cpu"
+                )
+                raise RuntimeError(
+                    f"dsv4_fp4_grouped_mode={_fp4_grouped_mode!r} with LoRA selected the grouped "
+                    f"NVFP4 MoE path, but no fused grouped backend resolved on this arch ({_cap}): "
+                    "marlin, cutlass, and deepgemm are all unavailable. The legacy fused-MX kernel "
+                    "(scatter2scatter_lora_mx) is unsafe on this arch (SIGSEGV on Blackwell), so it "
+                    "is not used as a fallback. Run on a supported GPU (sm80+/sm90/sm100/sm120) or "
+                    "disable dsv4_fp4_grouped_mode."
+                )
+            # FSDP-safe: backward re-reads the (re-gathered) params via this recipe
+            _recipe = lambda: (  # noqa: E731
+                _get_base_param(self.gate_up_proj),
+                _get_base_param(self.down_proj),
+            )
+            # persistent per-module cache so the backend requantizes the frozen weight to mxfp4
+            # once, surviving FSDP re-gathers (the gathered param is fresh each step).
+            cache = self.__dict__.setdefault("_dg_mxfp4_cache", {})
+            return grouped_fp4_moe_train(
+                hidden_states,
+                top_k_index,
+                routing_weights,
+                gu_base,
+                dn_base,
+                gup_lora,
+                down_lora,
+                getattr(self, "limit", None),
+                _fp4_grouped_mode,
+                act_type=_detect_act_type(self),
+                weight_recipe=_recipe,
+                mxfp4_cache=cache,
+                prefer_fp8_dx=RUNTIME.fp4_dx_prefer_fp8,
+            )
+
+    # NVFP4 experts + LoRA without dsv4_fp4_grouped_mode set would fall through to
+    # _prepare_weights_and_lora -> selective_nvfp4_weights_fwd -> scatter2scatter_lora_mx, which
+    # SIGSEGVs on Blackwell. Turn that silent crash into an actionable error (the grouped path is the
+    # supported route for NVFP4+LoRA); non-Blackwell archs keep the legacy path.
+    if (
+        _fp4_grouped_mode is None
+        and gup_lora is not None
+        and down_lora is not None
+        and torch.cuda.is_available()
+        and torch.cuda.get_device_capability()[0] >= 10
+        and is_nvfp4_param(_get_base_param(self.gate_up_proj))
+    ):
+        raise RuntimeError(
+            "NVFP4 experts with LoRA require the grouped fp4 MoE path on Blackwell (sm100/sm120): "
+            "the legacy fused-MX kernel (scatter2scatter_lora_mx) SIGSEGVs on this arch. Set "
+            "dsv4_fp4_grouped_mode: nvfp4 in your config to enable the supported grouped GEMM path."
+        )
+
+    # BnB-4bit experts have no 4-bit-read kernel; the naive path full-dequants all E experts every
+    # forward (~2.7x VRAM). Route to the chunked-dequant grouped MoE (bounded transient). Detect
+    # WITHOUT touching self.gate_up_proj (would full-dequant).
+    _bnb_experts = (
+        hasattr(self, "parametrizations")
+        and "gate_up_proj" in self.parametrizations
+        and "down_proj" in self.parametrizations
+    )
+    _bnb_fast = False
+    if _bnb_experts:
+        from .chunked_bnb import bnb_fast_enabled
+
+        _bnb_fast = bnb_fast_enabled()
+    if _bnb_experts and not _bnb_fast:
+        from .chunked_bnb import chunked_bnb_moe
+
+        return chunked_bnb_moe(
+            hidden_states,
+            top_k_index,
+            routing_weights,
+            self,
+            gup_lora,
+            down_lora,
+            self.num_experts,
+            act_type=_detect_act_type(self),
+            limit=getattr(self, "limit", None),
+        )
+
+    if _bnb_experts:
+        # bnb_fast: selective dequant of active experts -> 1-launch parallel_linear instead of the
+        # chunked torch._grouped_mm storm. Holds bf16 for backward.
+        from .selective_dequant import selective_expert_weights
+
+        active = get_active_experts(sorted_expert_idxs, self.num_experts)
+        sorted_expert_idxs, expert_offsets = remap_expert_indices(
+            sorted_expert_idxs, expert_offsets, active, self.num_experts
+        )
+        gate_up_weight = selective_expert_weights(
+            self, "gate_up_proj", active
+        ).transpose(2, 1)
+        down_weight = selective_expert_weights(self, "down_proj", active).transpose(
+            2, 1
+        )
+        # Recompute-in-backward recipe: the selective dequant is a per-layer bf16 copy of the active
+        # experts that ScatterMoELoRA would otherwise pin across all layers (~40 GB). The frozen
+        # 4-bit param is resident, so re-run the dequant in backward via the closure.
+        gate_up_weight.recipe = lambda m=self, a=active: selective_expert_weights(
+            m, "gate_up_proj", a
+        ).transpose(2, 1)
+        down_weight.recipe = lambda m=self, a=active: selective_expert_weights(
+            m, "down_proj", a
+        ).transpose(2, 1)
+        if gup_lora is not None and down_lora is not None:
+            gup_lora = (
+                *selective_lora_weights(
+                    gup_lora[0], gup_lora[1], active, self.num_experts
+                ),
+                gup_lora[2],
+            )
+            down_lora = (
+                *selective_lora_weights(
+                    down_lora[0], down_lora[1], active, self.num_experts
+                ),
+                down_lora[2],
+            )
+    else:
+        (
+            gate_up_weight,
+            down_weight,
+            sorted_expert_idxs,
+            expert_offsets,
+            gup_lora,
+            down_lora,
+        ) = _prepare_weights_and_lora(
+            _get_base_param(self.gate_up_proj),
+            _get_base_param(self.down_proj),
+            sorted_expert_idxs,
+            expert_offsets,
+            self.num_experts,
+            gup_lora,
+            down_lora,
+            hidden_states.dtype,
+            module=self,
+        )
+
+    gates_h = _parallel_linear_maybe_lora(
+        hidden_states,
+        gate_up_weight,
+        K,
+        sorted_expert_idxs,
+        sorted_scattered_idxs,
+        expert_offsets,
+        gup_lora,
+        grouped_in=False,
+        grouped_out=True,
+    )
+    gates, h = gates_h.chunk(2, dim=-1)
+    # Clamped SwiGLU when the model defines a swiglu_limit (e.g. DeepSeek-V4 limit=10) must match the
+    # eager experts' `_apply_gate` (gate.clamp(max=L); up.clamp(-L, L)), else outliers blow up.
+    _limit = getattr(self, "limit", None)
+    if _limit is not None:
+        gates = gates.clamp(max=_limit)
+        h = h.clamp(min=-_limit, max=_limit)
+    h = self.act_fn(gates) * h
+
+    output = _parallel_linear_maybe_lora(
+        h,
+        down_weight,
+        1,
+        sorted_expert_idxs,
+        sorted_scattered_idxs,
+        expert_offsets,
+        down_lora,
+        grouped_in=True,
+        grouped_out=False,
+        gates=routing_weights,
+    )
+
+    return output
+
+
+def _grouped_fp4_ep_maybe_tiled(
+    hidden,
+    idx,
+    wts,
+    gu_base,
+    dn_base,
+    gup_lora,
+    down_lora,
+    limit,
+    mode,
+    act_type,
+    recipe,
+    cache,
+    prefer_fp8_dx,
+):
+    """Run the grouped-fp4 expert GEMM, optionally TILED over the dispatched-token (sequence) dim.
+
+    The grouped path stores ``gu[Mt,2I]`` + ``h[Mt,I]`` for backward (the activation-memory cost).
+    With ``AXOLOTL_EP_EXPERT_SHARD_TOKENS=T`` set, slice the dispatched tokens into T-row chunks and
+    checkpoint each chunk (recompute its forward in backward) so peak saved activation is one chunk's
+    intermediates, not all of Mt — trading ~1 extra forward/chunk for memory. The frozen-weight mxfp4
+    requant ``cache`` is shared across chunks (and across the recompute), so weights are quantized once.
+    T<=0 or T>=N keeps the single un-tiled call. Chunks own disjoint token rows so the per-chunk [N_c,H]
+    outputs concatenate back in order."""
+    import os as _os_t
+
+    from .grouped_train import grouped_fp4_moe_train
+
+    def _call(h_c, i_c, w_c):
+        return grouped_fp4_moe_train(
+            h_c,
+            i_c,
+            w_c,
+            gu_base,
+            dn_base,
+            gup_lora,
+            down_lora,
+            limit,
+            mode,
+            act_type=act_type,
+            weight_recipe=recipe,
+            mxfp4_cache=cache,
+            prefer_fp8_dx=prefer_fp8_dx,
+        )
+
+    shard = int(_os_t.environ.get("AXOLOTL_EP_EXPERT_SHARD_TOKENS", "0") or 0)
+    n = hidden.size(0)
+    if shard <= 0 or n <= shard:
+        return _call(hidden, idx, wts)
+
+    import torch.utils.checkpoint as _ckpt
+
+    outs = [
+        _ckpt.checkpoint(
+            _call,
+            hidden[s : s + shard],
+            idx[s : s + shard],
+            wts[s : s + shard],
+            use_reentrant=False,
+        )
+        for s in range(0, n, shard)
+    ]
+    return torch.cat(outs, dim=0)
+
+
+def scattermoe_experts_forward_ep(
+    self,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    """ScatterMoE experts forward for the DeepEP local path, skipping EP sentinels.
+
+    After DeepEP dispatch ``top_k_index`` holds local expert ids in ``[0, E_local)``
+    for slots this rank owns and ``-1`` for slots routed to remote ranks. Rather than
+    map sentinels to expert 0 / weight 0 and run the full grouped GEMM over all
+    ``N*K`` rows (compute-and-mask), drop the sentinel rows so only the valid routed
+    rows hit the GEMM + per-row LoRA. Output matches the masked path since sentinel
+    slots carry weight 0.
+
+    Runs both projections fully grouped (the sentinel-compacted routing breaks the
+    ``L_scattered == X.rows * k`` fan-out contract of the scattered path), with the
+    weighted token-combine done via ``index_add_``.
+    """
+    _check_supported_layout(self)
+
+    # Prefer the grouped NVFP4 GEMM (the same backend the non-EP forward uses) over the triton
+    # scatter2scatter LoRA kernel: it's faster AND has no Triton autotune, so there's no per-rank
+    # autotune-timing skew to trip DeepEP's combine timeout (which is why the scatter2scatter path
+    # needed AXOLOTL_EP_SINGLE_CONFIG). grouped_fp4_moe_train tolerates the DeepEP -1 sentinels
+    # (remote-routed slots are dropped). Falls through to scatter2scatter for non-NVFP4 / mode-off /
+    # no-LoRA layouts.
+    _gup_lora_g, _down_lora_g = None, None
+    _sm_lora_g = getattr(self, "_scattermoe_lora", None)
+    if _sm_lora_g:
+        _gup_lora_g = _sm_lora_g.get("gate_up_proj")
+        _down_lora_g = _sm_lora_g.get("down_proj")
+    elif _has_peft_wrapper(self):
+        _, _gup_lora_g, _down_lora_g = _unwrap_experts_lora(self)
+
+    _fp4_grouped_mode_g = RUNTIME.fp4_grouped_mode
+    if (
+        _fp4_grouped_mode_g is not None
+        and _gup_lora_g is not None
+        and _down_lora_g is not None
+    ):
+        _gu_base_g = _get_base_param(self.gate_up_proj)
+        _dn_base_g = _get_base_param(self.down_proj)
+        if is_nvfp4_param(_gu_base_g):
+            from .grouped_train import grouped_fp4_available
+
+            if grouped_fp4_available(_fp4_grouped_mode_g):
+                _recipe_g = lambda: (  # noqa: E731
+                    _get_base_param(self.gate_up_proj),
+                    _get_base_param(self.down_proj),
+                )
+                _cache_g = self.__dict__.setdefault("_dg_mxfp4_cache", {})
+                return _grouped_fp4_ep_maybe_tiled(
+                    hidden_states,
+                    top_k_index,
+                    top_k_weights.to(hidden_states.dtype),
+                    _gu_base_g,
+                    _dn_base_g,
+                    _gup_lora_g,
+                    _down_lora_g,
+                    getattr(self, "limit", None),
+                    _fp4_grouped_mode_g,
+                    _detect_act_type(self),
+                    _recipe_g,
+                    _cache_g,
+                    RUNTIME.fp4_dx_prefer_fp8,
+                )
+
+    N = hidden_states.size(0)
+    K = top_k_index.shape[1]
+    E_local = self.num_experts
+
+    idx_flat = top_k_index.reshape(-1)
+    valid = idx_flat >= 0
+    e_v = idx_flat[valid].to(torch.long)
+    if e_v.numel() == 0:
+        return torch.zeros_like(hidden_states)
+
+    tok = torch.arange(N, device=hidden_states.device).repeat_interleave(K)
+    tok_v = tok[valid]
+    w_v = top_k_weights.reshape(-1)[valid].to(hidden_states.dtype)
+
+    # Sort valid rows by expert so each expert's rows are contiguous (grouped layout).
+    se, order = torch.sort(e_v)
+    tok_sorted = tok_v[order].to(torch.int32)
+    w_sorted = w_v[order]
+    expert_offsets = torch.bincount(e_v, minlength=E_local).cumsum(-1)
+    M = se.size(0)
+    ss = torch.arange(M, device=hidden_states.device, dtype=torch.int32)
+
+    gup_lora, down_lora = None, None
+    sm_lora = getattr(self, "_scattermoe_lora", None)
+    if sm_lora:
+        gup_lora = sm_lora.get("gate_up_proj")
+        down_lora = sm_lora.get("down_proj")
+    elif _has_peft_wrapper(self):
+        _, gup_lora, down_lora = _unwrap_experts_lora(self)
+
+    gate_up_weight, down_weight, se, expert_offsets, gup_lora, down_lora = (
+        _prepare_weights_and_lora(
+            _get_base_param(self.gate_up_proj),
+            _get_base_param(self.down_proj),
+            se,
+            expert_offsets,
+            E_local,
+            gup_lora,
+            down_lora,
+            hidden_states.dtype,
+        )
+    )
+
+    # Pre-gather token rows into expert-grouped order; both projections run grouped.
+    grouped_x = hidden_states.index_select(0, tok_sorted.to(torch.long))
+
+    gates_h = _parallel_linear_maybe_lora(
+        grouped_x,
+        gate_up_weight,
+        1,
+        se,
+        ss,
+        expert_offsets,
+        gup_lora,
+        grouped_in=True,
+        grouped_out=True,
+    )
+    gates, h = gates_h.chunk(2, dim=-1)
+    _limit = getattr(self, "limit", None)
+    if _limit is not None:
+        gates = gates.clamp(max=_limit)
+        h = h.clamp(min=-_limit, max=_limit)
+    h = self.act_fn(gates) * h
+
+    down_out = _parallel_linear_maybe_lora(
+        h,
+        down_weight,
+        1,
+        se,
+        ss,
+        expert_offsets,
+        down_lora,
+        grouped_in=True,
+        grouped_out=True,
+    )
+    down_out = down_out * w_sorted.unsqueeze(-1)
+
+    output = hidden_states.new_zeros((N, down_out.size(-1)))
+    output.index_add_(0, tok_sorted.to(torch.long), down_out)
+    return output
+
+
+_SCATTERMOE_PATCHED = False
+
+
+def _ensure_single_lora_ops():
+    """Guarantee a single module instance of the autotuned ``kernels.lora_ops``.
+
+    A second instance (same file imported under a different module name) would carry its own
+    Triton ``Autotuner`` caches — duplicate autotuning (wasted) and confusing telemetry where
+    the same kernel+shape+dtype appears twice with different configs. We canonicalize on the
+    axolotl import path and alias any duplicate in ``sys.modules`` to it, so later imports
+    resolve to one module. Idempotent; warns only if a real duplicate is found."""
+    import sys
+
+    from axolotl.utils.logging import get_logger
+
+    log = get_logger(__name__)
+    canon_name = "axolotl.integrations.kernels.libs.scattermoe_lora.kernels.lora_ops"
+    canon = sys.modules.get(canon_name)
+    if canon is None:
+        return
+    canon_file = getattr(canon, "__file__", None)
+    aliased = 0
+    for name, mod in list(sys.modules.items()):
+        if mod is None or mod is canon or name == canon_name or "lora_ops" not in name:
+            continue
+        if getattr(mod, "__file__", None) == canon_file:
+            sys.modules[name] = canon
+            aliased += 1
+    if aliased:
+        log.warning(
+            "Aliased %d duplicate lora_ops module instance(s) to %s (single autotune cache)",
+            aliased,
+            canon_name,
+        )
+
+
+def register_scattermoe_experts():
+    """Register ``"scattermoe"`` in the ExpertsInterface and the validator allowlist.
+
+    Idempotent.
+    """
+    global _SCATTERMOE_PATCHED
+
+    _ensure_single_lora_ops()
+
+    from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+    from transformers.modeling_utils import PreTrainedModel
+
+    ALL_EXPERTS_FUNCTIONS.register("scattermoe", scattermoe_experts_forward)
+
+    # transformers' lazy-import can leave TWO `integrations.moe` registries.
+    # `@use_experts_implementation` binds its registry as a default argument, which can be the OTHER
+    # object than the one imported above, so register into the decorator's default interface too.
+    try:
+        import inspect
+
+        from transformers.integrations import use_experts_implementation
+
+        canon = (
+            inspect.signature(use_experts_implementation)
+            .parameters["experts_interface"]
+            .default
+        )
+        if canon is not None and "scattermoe" not in canon:
+            canon.register("scattermoe", scattermoe_experts_forward)
+    except (ImportError, KeyError, AttributeError):
+        pass
+
+    # Route PEFT target_parameters expert LoRA to the fused kernel (bypassing the parametrization
+    # merge) for experts-interface MoEs.
+    try:
+        from .experts_lora_fastpath import patch_paramwrapper_fastpath
+
+        patch_paramwrapper_fastpath()
+    except (ImportError, AttributeError):
+        pass
+
+    # Safety net for any path that still merges `base + delta` on a frozen FP4 base
+    # (e.g. merge_and_unload): make aten.add dequantize the FP4 tensor.
+    try:
+        from .torchao_fp4_add import patch_torchao_fp4_add
+
+        patch_torchao_fp4_add()
+    except (ImportError, AttributeError):  # torchao optional / API drift
+        pass
+
+    # FSDP2 support for NVFP4Tensor (split/view/as_strided + all-gather hooks) so the quantized
+    # expert weights can be sharded across GPUs (torchao ships none).
+    try:
+        from .nvfp4_fsdp import patch_nvfp4_fsdp
+
+        patch_nvfp4_fsdp()
+    except (ImportError, AttributeError):
+        pass
+
+    if _SCATTERMOE_PATCHED:
+        return
+
+    _original_get_correct = PreTrainedModel.get_correct_experts_implementation
+
+    def _patched_get_correct(self_model, requested_experts: str | None) -> str:
+        if requested_experts == "scattermoe":
+            return "scattermoe"
+        return _original_get_correct(self_model, requested_experts)
+
+    PreTrainedModel.get_correct_experts_implementation = _patched_get_correct
+    _SCATTERMOE_PATCHED = True
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/experts_lora_fastpath.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/experts_lora_fastpath.py
new file mode 100644
index 0000000000..a84d7782d0
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/experts_lora_fastpath.py
@@ -0,0 +1,123 @@
+"""Route PEFT ``target_parameters`` LoRA on experts-interface MoEs to the fused kernel.
+
+For models whose MoE is dispatched through transformers' ExpertsInterface
+(``@use_experts_implementation`` — e.g. Gemma 4, DiffusionGemma), PEFT wraps the
+*experts module* in a ``ParamWrapper`` chain. The decoder layer then calls
+``experts(...)`` → ``ParamWrapper.forward`` → ``_activate_lora`` which materializes the
+full merged weight ``base + delta`` before ScatterMoE runs. That defeats ScatterMoE's
+fused LoRA kernel (which takes A/B separately, dequantizing only the active experts) and
+fails outright on quantized bases (``aten.add`` on NVFP4/MXFP4).
+
+This patches ``ParamWrapper.forward`` so that, when the wrapped base is a ScatterMoE
+experts module, it walks the wrapper chain, hands the LoRA A/B to the base via
+``_scattermoe_lora`` (in ScatterMoE layout), and calls the base forward directly —
+bypassing the parametrization merge entirely. The wrappers stay in the module tree, so
+PEFT save/load and optimizer tracking are unaffected.
+
+This mirrors what ``HFScatterMoEGatedMLP`` already does for SparseMoeBlock models (walk
+the wrapper, fuse the LoRA), bringing the experts-interface path to parity.
+"""
+
+from __future__ import annotations
+
+# Implementations whose experts forward consumes ``module._scattermoe_lora`` (the fused
+# per-row LoRA kernel). ``deep_ep_scattermoe`` is the EP composite: its ``_scattermoe_local``
+# stage calls ``scattermoe_experts_forward_ep``, which reads the same attribute — so the
+# fastpath must engage under EP too, else the LoRA falls back to PEFT's parametrize merge
+# (which can't add a full-expert delta onto the EP-sharded weight).
+_SCATTERMOE_IMPLS = frozenset({"scattermoe", "deep_ep_scattermoe"})
+
+
+def _is_scattermoe_experts(module) -> bool:
+    cfg = getattr(module, "config", None)
+    impl = getattr(cfg, "_experts_implementation", None)
+    return impl in _SCATTERMOE_IMPLS and hasattr(module, "gate_up_proj")
+
+
+def patch_paramwrapper_fastpath() -> None:
+    """Idempotently patch ``peft...ParamWrapper.forward`` for ScatterMoE experts."""
+    try:
+        from peft.tuners.lora.layer import ParamWrapper
+    except (ImportError, AttributeError):
+        return
+
+    if getattr(ParamWrapper.forward, "_scattermoe_fastpath", False):
+        return
+
+    from .layers import _convert_smoe_lora
+    from .parallel_linear_lora import get_lora_params_from_wrapper
+
+    _orig_forward = ParamWrapper.forward
+
+    def _fusable(wrapper) -> bool:
+        # Fused kernel needs exactly one adapter on a raw (un-merged) base; else defer to PEFT.
+        if getattr(wrapper, "disable_adapters", False) or getattr(
+            wrapper, "merged", False
+        ):
+            return False
+        return len(getattr(wrapper, "active_adapters", [])) == 1
+
+    def forward(self, x, *args, **kwargs):
+        # Mixed-batch inference (adapter_names) is PEFT's domain; defer to it.
+        if kwargs.get("adapter_names") is not None:
+            return _orig_forward(self, x, *args, **kwargs)
+
+        # one wrapper per targeted parameter
+        wrappers = {}
+        base = self
+        while hasattr(base, "base_layer") and hasattr(base, "lora_A"):
+            name = getattr(base, "parameter_name", None)
+            if name is not None:
+                wrappers[name] = base
+            base = base.base_layer
+
+        if not (
+            _is_scattermoe_experts(base) and all(_fusable(w) for w in wrappers.values())
+        ):
+            return _orig_forward(self, x, *args, **kwargs)
+
+        num_experts = getattr(base, "num_experts", None)
+        sm_lora = {}
+        for name, wrapper in wrappers.items():
+            lora_A, lora_B, scaling = get_lora_params_from_wrapper(wrapper)
+            if lora_A is None or num_experts is None:
+                continue
+            # Under EP the base is sliced to E_local experts but the adapter stays a single global
+            # tensor over E_global experts; take THIS rank's local-expert block and the true LoRA
+            # rank so the fused kernel reads the right experts. No-op when not EP-sharded.
+            from .experts import _ep_local_expert_lora
+
+            lora_A, lora_B, num_experts_local, rank = _ep_local_expert_lora(
+                lora_A, lora_B, base
+            )
+            # PEFT keeps LoRA fp32; cast to activation dtype (grads still route to the fp32 params).
+            lora_A = lora_A.to(x.dtype)
+            lora_B = lora_B.to(x.dtype)
+            sm_lora[name] = _convert_smoe_lora(
+                lora_A, lora_B, num_experts_local, rank, scaling
+            )
+
+        base._scattermoe_lora = sm_lora
+        # Loading can create multiple ExpertsInterface registries; the experts forward binds one as
+        # a closure default that may differ from the one register_scattermoe_experts() populated, so
+        # register into the one THIS module dispatches through.
+        if not getattr(base, "_scattermoe_iface_ok", False):
+            _fn = type(base).forward
+            _cells = dict(
+                zip(_fn.__code__.co_freevars, _fn.__closure__ or (), strict=False)
+            )
+            _ei = _cells.get("experts_interface")
+            if _ei is not None:
+                _ei = _ei.cell_contents
+                if "scattermoe" not in _ei:
+                    from .experts import scattermoe_experts_forward
+
+                    _ei.register("scattermoe", scattermoe_experts_forward)
+            base._scattermoe_iface_ok = True
+        try:
+            return base(x, *args, **kwargs)
+        finally:
+            base._scattermoe_lora = None
+
+    forward._scattermoe_fastpath = True  # type: ignore[attr-defined]
+    ParamWrapper.forward = forward
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/gemma4_fp8_nonexpert.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/gemma4_fp8_nonexpert.py
new file mode 100644
index 0000000000..b7e66a7c09
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/gemma4_fp8_nonexpert.py
@@ -0,0 +1,206 @@
+"""In-place fp8 per-channel quantization of gemma4 non-expert linears.
+
+Targets every ``nn.Linear`` that is NOT inside a ``Gemma4TextExperts`` block
+(experts stay NVFP4Tensor, untouched).  Norms, embeddings, router projections,
+vision tower weights, and lm_head are also skipped by class or name.
+
+After this runs each targeted linear has:
+  - ``weight``           : fp8_e4m3fn  [out, in]   (1 byte/param)
+  - ``weight_scale_inv`` : bfloat16    [out, 1]    (per-row scale = max_abs/fp8_max)
+  - ``forward``          : overridden to dequant on-the-fly and call F.linear
+
+The name ``weight_scale_inv`` matches the axolotl LoRA kernel's lookup so the
+backward dequantize path works without further changes.
+
+The dequant-in-forward path is the same as the DSV4 bf16-mode fallback —
+frozen fp8 base, bf16 activations, bf16 LoRA adapters layered on top by PEFT.
+The fused kernel rewrite (``torch._scaled_mm``) is a follow-up.
+
+Usage (in plugin.py ``pre_lora_load``):
+    from .gemma4_fp8_nonexpert import quantize_gemma4_nonexpert_linears
+    quantize_gemma4_nonexpert_linears(model)
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# Names/substrings that should NEVER be fp8-quantized (precision-sensitive or tiny).
+_SKIP_NAME_SUBSTRINGS = (
+    "norm",
+    "embed",
+    "lm_head",
+    "router",
+    "vision_tower",
+    "audio",
+)
+
+_FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+
+def _should_skip_by_name(name: str) -> bool:
+    return any(sub in name for sub in _SKIP_NAME_SUBSTRINGS)
+
+
+def _patch_linear_forward(mod: nn.Linear) -> None:
+    """Monkey-patch a single nn.Linear whose weight was already replaced with
+    an fp8 tensor + weight_scale_inv bfloat16 parameter [out, 1].
+
+    The attribute name ``weight_scale_inv`` is chosen to match what the axolotl
+    LoRA kernel (``lora.py:get_lora_parameters``) probes for on an fp8 weight::
+
+        quant_state = getattr(base_layer, "weight_scale_inv", None)
+
+    so the LoRA backward's ``dequantize(q_weight, q_quant)`` call receives the
+    per-channel scale and correctly dequantizes to bf16.
+    """
+
+    def _forward(self, x: torch.Tensor) -> torch.Tensor:
+        w_fp8: torch.Tensor = self.weight  # float8_e4m3fn [out, in]
+        scale: torch.Tensor = self.weight_scale_inv  # bfloat16 [out, 1]
+        w_bf16 = w_fp8.to(torch.bfloat16) * scale
+        x_compute = x.to(torch.bfloat16) if x.dtype != torch.bfloat16 else x
+        out = F.linear(x_compute, w_bf16, self.bias)
+        return out.to(x.dtype) if x.dtype != torch.bfloat16 else out
+
+    import types
+
+    mod.forward = types.MethodType(_forward, mod)
+
+
+def quantize_gemma4_nonexpert_linears(model: nn.Module) -> int:
+    """Quantize all non-expert nn.Linear weights in a loaded gemma4 model to
+    fp8_e4m3fn per-channel in-place.  Returns the number of modules quantized.
+
+    Safe to call multiple times: already-quantized modules (weight.dtype ==
+    float8_e4m3fn) are skipped.
+    """
+    try:
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts
+
+        _expert_cls = Gemma4TextExperts
+    except ImportError:
+        _expert_cls = None
+
+    expert_paths: set[str] = set()
+    if _expert_cls is not None:
+        for path, mod in model.named_modules():
+            if isinstance(mod, _expert_cls):
+                expert_paths.add(path)
+
+    def _is_under_expert(path: str) -> bool:
+        return any(path == ep or path.startswith(ep + ".") for ep in expert_paths)
+
+    quantized = 0
+    by_type: dict[str, int] = {}
+
+    for name, mod in model.named_modules():
+        if not isinstance(mod, nn.Linear):
+            continue
+        if _is_under_expert(name):
+            continue
+        if _should_skip_by_name(name):
+            continue
+        w = mod.weight
+        if not isinstance(w, nn.Parameter):
+            continue
+        if w.dtype == torch.float8_e4m3fn:
+            continue  # already quantized
+        if w.ndim != 2:
+            continue  # embeddings are sometimes Linear subclasses
+
+        w_data = w.data.to(torch.bfloat16) if w.dtype != torch.bfloat16 else w.data
+        scale_1d = w_data.abs().amax(dim=1).clamp(min=1e-12)
+        # "weight_scale_inv" name + [out, 1] shape match the axolotl LoRA kernel's fp8 quant_state lookup.
+        scale = (scale_1d / _FP8_MAX).to(torch.bfloat16).unsqueeze(1)
+        w_fp8 = (
+            (w_data / scale.to(w_data.dtype))
+            .clamp(-_FP8_MAX, _FP8_MAX)
+            .to(torch.float8_e4m3fn)
+        )
+
+        mod.weight = nn.Parameter(w_fp8, requires_grad=False)
+        mod.register_parameter(
+            "weight_scale_inv", nn.Parameter(scale, requires_grad=False)
+        )
+        _patch_linear_forward(mod)
+
+        key = type(mod).__name__
+        by_type[key] = by_type.get(key, 0) + 1
+        quantized += 1
+
+    if quantized:
+        LOG.info(
+            "Gemma4 frankenstein: fp8-quantized %d non-expert linears in-place "
+            "(per-channel e4m3, dequant-in-forward): %s",
+            quantized,
+            by_type,
+        )
+    return quantized
+
+
+def verify_gemma4_frankenstein(model: nn.Module) -> dict:
+    """Sanity-check the frankenstein state.  Returns a dict of counts."""
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+        _has_nvfp4 = True
+    except ImportError:
+        _has_nvfp4 = False
+        NVFP4Tensor = None
+
+    n_nvfp4 = n_fp8 = n_bf16_lin = n_skip = 0
+    fp8_bytes = bf16_lin_bytes = expert_bytes = 0
+
+    try:
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts
+
+        _expert_cls = Gemma4TextExperts
+    except ImportError:
+        _expert_cls = None
+
+    expert_paths: set[str] = set()
+    if _expert_cls is not None:
+        for path, mod in model.named_modules():
+            if isinstance(mod, _expert_cls):
+                expert_paths.add(path)
+
+    def _is_under_expert(path):
+        return any(path == ep or path.startswith(ep + ".") for ep in expert_paths)
+
+    for name, param in model.named_parameters():
+        if _has_nvfp4 and isinstance(param, NVFP4Tensor):
+            n_nvfp4 += 1
+            expert_bytes += param.data.numel() * 1  # uint8
+            continue
+        if param.dtype == torch.float8_e4m3fn:
+            n_fp8 += 1
+            fp8_bytes += param.numel()
+            continue
+        if param.dtype in (torch.bfloat16, torch.float16, torch.float32):
+            if name.endswith(".weight") and not _is_under_expert(
+                name.rsplit(".", 1)[0]
+            ):
+                if not _should_skip_by_name(name):
+                    n_bf16_lin += 1
+                    bf16_lin_bytes += param.numel() * 2
+                else:
+                    n_skip += 1
+            else:
+                n_skip += 1
+
+    return {
+        "nvfp4_expert_params": n_nvfp4,
+        "fp8_nonexpert_params": n_fp8,
+        "bf16_linear_unexpected": n_bf16_lin,
+        "skipped_params": n_skip,
+        "expert_bytes_GB": expert_bytes / 1e9,
+        "fp8_bytes_GB": fp8_bytes / 1e9,
+        "bf16_linear_unexpected_bytes_GB": bf16_lin_bytes / 1e9,
+    }
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/gemma4_nf4_nonexpert.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/gemma4_nf4_nonexpert.py
new file mode 100644
index 0000000000..abf337d3f5
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/gemma4_nf4_nonexpert.py
@@ -0,0 +1,108 @@
+"""In-place NF4 (bitsandbytes) quantization of gemma4 non-expert linears.
+
+This mirrors :mod:`gemma4_fp8_nonexpert` but uses bitsandbytes 4-bit NF4 — the
+SAME non-expert compute path unsloth's QLoRA uses — so an experts-only LoRA run
+can be compared apples-to-apples against unsloth: identical frozen-NF4 non-expert
+matmuls, the only difference being the routed-expert path (our NVFP4
+grouped/marlin + scatter LoRA vs unsloth's NF4 experts).
+
+Targets every ``nn.Linear`` that is NOT inside a ``Gemma4TextExperts`` block
+(experts stay NVFP4Tensor, untouched).  Norms, embeddings, router projections,
+vision/audio towers, and lm_head are skipped by name.  Each target becomes a
+frozen ``bnb.nn.Linear4bit`` (nf4, double-quant, bf16 compute) — no LoRA is
+attached there (experts-only), exactly like unsloth's non-expert state.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from axolotl.utils.logging import get_logger
+
+from .gemma4_fp8_nonexpert import _should_skip_by_name
+
+LOG = get_logger(__name__)
+
+
+def _expert_paths(model: nn.Module) -> set[str]:
+    try:
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts
+
+        cls = Gemma4TextExperts
+    except ImportError:
+        cls = None
+    paths: set[str] = set()
+    if cls is not None:
+        for path, mod in model.named_modules():
+            if isinstance(mod, cls):
+                paths.add(path)
+    return paths
+
+
+def quantize_gemma4_nonexpert_nf4(
+    model: nn.Module, compute_dtype: torch.dtype = torch.bfloat16
+) -> int:
+    """Swap all non-expert ``nn.Linear`` modules for frozen bnb NF4 ``Linear4bit``
+    (double-quant, bf16 compute), in place.  Returns the count swapped.  Experts
+    (NVFP4Tensor) are untouched.  Idempotent: existing ``Linear4bit`` are skipped.
+    """
+    import bitsandbytes as bnb
+
+    expert_paths = _expert_paths(model)
+
+    def _under_expert(path: str) -> bool:
+        return any(path == ep or path.startswith(ep + ".") for ep in expert_paths)
+
+    device = next(
+        (p.device for p in model.parameters() if p.is_cuda),
+        torch.device("cuda", torch.cuda.current_device()),
+    )
+
+    targets: list[tuple[str, nn.Linear]] = []
+    for name, mod in model.named_modules():
+        if isinstance(mod, bnb.nn.Linear4bit):
+            continue
+        if not isinstance(mod, nn.Linear):
+            continue
+        if _under_expert(name) or _should_skip_by_name(name):
+            continue
+        if not isinstance(mod.weight, nn.Parameter) or mod.weight.ndim != 2:
+            continue
+        targets.append((name, mod))
+
+    by_type: dict[str, int] = {}
+    for name, mod in targets:
+        has_bias = mod.bias is not None
+        new = bnb.nn.Linear4bit(
+            mod.in_features,
+            mod.out_features,
+            bias=has_bias,
+            compute_dtype=compute_dtype,
+            quant_type="nf4",
+        )
+        new.weight = bnb.nn.Params4bit(
+            mod.weight.data.to(torch.bfloat16),
+            requires_grad=False,
+            quant_type="nf4",
+            compress_statistics=True,
+        )
+        if has_bias:
+            new.bias = nn.Parameter(
+                mod.bias.data.to(compute_dtype), requires_grad=False
+            )
+        new = new.to(device)  # triggers NF4 quantization of Params4bit
+
+        parent_path, _, attr = name.rpartition(".")
+        parent = model.get_submodule(parent_path) if parent_path else model
+        setattr(parent, attr, new)
+        by_type[type(mod).__name__] = by_type.get(type(mod).__name__, 0) + 1
+
+    if targets:
+        LOG.info(
+            "Gemma4 NF4 frankenstein: swapped %d non-expert linears to bnb NF4 "
+            "(double-quant, bf16 compute): %s",
+            len(targets),
+            by_type,
+        )
+    return len(targets)
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_lora.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_lora.py
new file mode 100644
index 0000000000..508ad3822b
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_lora.py
@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Encapsulated scatter2scatter grouped LoRA ops for the TILE-padded marlin/grouped path.
+
+Replaces per-expert torch._grouped_mm loops in grouped_train.py with single-launch
+scatter2scatter calls.  All implementation choices (which kernel variant, dispatch)
+live inside this module — the training code calls one function per phase.
+
+Layout bridge: grouped_train uses a TILE-padded expert-sorted buffer A[Mt, K] with
+per-TILE expert ids (m_indices) and cumulative padded offsets (offs).  The
+scatter2scatter kernels in kernels/ops.py work on the same sorted layout when called
+with x_grouped=True / y_grouped=True.
+
+LoRA weight layout (stacked, from _lora_stack in grouped_train):
+  A : [E, r, in_K]   (A-adapter)
+  B : [E, out_N, r]  (B-adapter)
+
+scatter2scatter weight layout: W[E, in, out].  Mapping:
+  W_A = A.permute(0, 2, 1)  -> [E, in_K, r]    for X @ A^T step
+  W_B = B.permute(0, 2, 1)  -> [E, r, out_N]   for XA @ B^T step
+  W_Bt = B                  -> [E, out_N, r]    for dY @ B step (B already [E, out_N, r])
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .kernels.grouped_gram import grouped_lora_weight_grads
+from .kernels.ops import scatter2scatter
+
+
+def _mk_routing(offs: torch.Tensor, E: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build scatter2scatter routing tensors from the TILE-padded grouped layout.
+
+    Args:
+        offs: cumulative TILE-padded expert offsets [E] (int32), as produced by
+              grouped_fp4_moe_train — offs[e] = sum of tile-padded row counts
+              for experts 0..e.
+        E: number of experts.
+
+    Returns:
+        (sorted_expert_idxs [Mt], sorted_scattered_idxs [Mt])
+        sorted_expert_idxs[i] = expert id for row i of the grouped buffer.
+        sorted_scattered_idxs = identity [0..Mt-1]; with x_grouped=y_grouped=True
+        the scatter2scatter kernel reads/writes M_block directly and ignores it
+        for X/Y addressing, but it must have the right length.
+    """
+    Mt = int(offs[-1].item())
+    dev = offs.device
+    starts = torch.cat([offs.new_zeros(1), offs[:-1]])
+    sizes = offs - starts  # per-expert padded row count
+    e_ids = torch.arange(E, device=dev, dtype=torch.int32)
+    sorted_expert_idxs = e_ids.repeat_interleave(sizes)
+    sorted_scattered_idxs = torch.arange(Mt, device=dev, dtype=torch.int32)
+    return sorted_expert_idxs, sorted_scattered_idxs
+
+
+def grouped_lora_fwd(
+    x: torch.Tensor,  # [Mt, in_K], expert-grouped (TILE-padded)
+    A: torch.Tensor,  # [E, r, in_K]
+    B: torch.Tensor,  # [E, out_N, r]
+    scaling: float,
+    offs: torch.Tensor,  # [E] int32 cumulative TILE-padded offsets
+    E: int,
+    residual: torch.Tensor
+    | None = None,  # [Mt, out_N] base output to fold into the epilogue
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """LoRA forward on the TILE-padded grouped layout.
+
+    Computes Y_lora = scaling * (X @ A^T) @ B^T, row-for-row with the marlin
+    base output.  Two scatter2scatter launches with x_grouped=y_grouped=True keep
+    the result in the padded grouped layout for direct in-place addition to base.
+
+    If ``residual`` is given (the base expert GEMM output), it is added in the LoRA-B GEMM epilogue,
+    so the returned tensor is ``base + scaling*lora`` with NO separate add pass or temp tensor.
+
+    Returns:
+        (y [Mt, out_N], xa [Mt, r])  — y = (residual + scaling*lora) if residual else scaling*lora;
+        xa is saved (unscaled) for the backward dB/dA.
+    """
+    sei, ssi = _mk_routing(offs, E)
+
+    W_A = A.permute(0, 2, 1).contiguous()  # [E, in_K, r]
+    W_B = B.permute(0, 2, 1).contiguous()  # [E, r, out_N]
+
+    xa = scatter2scatter(
+        X=x,
+        W=W_A,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=1,
+        x_grouped=True,
+        y_grouped=True,
+    )  # [Mt, r]
+
+    # Scale the tiny [Mt, r] inner activation, not the large [Mt, out_N] output (identical result,
+    # far fewer elements). xa is returned UNSCALED; the backward derives scaling separately.
+    y_lora = scatter2scatter(
+        X=xa * scaling,
+        W=W_B,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=1,
+        x_grouped=True,
+        y_grouped=True,
+        residual=residual,
+    )  # [Mt, out_N], = (residual + scaling*lora) if residual else scaling*lora
+    return y_lora, xa
+
+
+def grouped_lora_bwd(
+    dy: torch.Tensor,  # [Mt, out_N], grad w.r.t. LoRA output, grouped
+    x: torch.Tensor,  # [Mt, in_K], forward input, grouped
+    A: torch.Tensor,  # [E, r, in_K]
+    B: torch.Tensor,  # [E, out_N, r]
+    xa: torch.Tensor,  # [Mt, r], X@A^T from forward (saved)
+    scaling: float,
+    offs: torch.Tensor,  # [E] int32 cumulative TILE-padded offsets
+    E: int,
+    residual: torch.Tensor
+    | None = None,  # [Mt, in_K] base dX to fold into the dX_lora epilogue
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """LoRA backward on the TILE-padded grouped layout.
+
+    Computes:
+        dX_lora = scaling * (dY @ B) @ A   [Mt, in_K]
+        dA                                  [E, r, in_K]
+        dB                                  [E, out_N, r]
+
+    If ``residual`` is given (the base dX), it is added in the dX_lora GEMM epilogue, so the returned
+    dx is ``base_dX + scaling*lora_dX`` with no separate add pass.
+
+    Returns:
+        (dx, dA, dB)  — dx = (residual + scaling*lora_dX) if residual else scaling*lora_dX
+    """
+    sei, ssi = _mk_routing(offs, E)
+
+    r = A.size(1)
+    in_K = A.size(2)
+    out_N = B.size(1)
+
+    # yb = dY @ B; B is already [E, out_N, r] = [E, K=out_N, N=r]
+    yb = scatter2scatter(
+        X=dy,
+        W=B.contiguous(),
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=1,
+        x_grouped=True,
+        y_grouped=True,
+    )  # [Mt, r]
+
+    # dx_lora = scaling * yb @ A; A is [E, r, in_K] = [E, K=r, N=in_K]. Fold `scaling` into the tiny
+    # [Mt, r] yb, not the large [Mt, in_K] output (identical result, far fewer elements).
+    dx_lora = scatter2scatter(
+        X=yb * scaling,
+        W=A.contiguous(),
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=1,
+        x_grouped=True,
+        y_grouped=True,
+        residual=residual,
+    )  # [Mt, in_K], = (residual + scaling*lora) if residual else scaling*lora
+
+    # grouped-Gram kernel expects flat scattermoe layout: lora_A [r*E, in_K], lora_B [out_N, r*E]
+    lora_A_flat = A.reshape(E * r, in_K)
+    lora_B_flat = B.permute(1, 0, 2).reshape(out_N, E * r)
+
+    dA_flat, dB_flat = grouped_lora_weight_grads(
+        grouped_grad_out=dy,
+        grouped_x=x,
+        yb=yb,
+        xa=xa,
+        lora_A=lora_A_flat,
+        lora_B=lora_B_flat,
+        combined_offsets=offs,
+        e_total=E,
+        scaling=scaling,
+    )
+    dA = dA_flat.reshape(E, r, in_K)
+    dB = dB_flat.reshape(out_N, E, r).permute(1, 0, 2).contiguous()
+
+    return dx_lora, dA, dB
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_moe.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_moe.py
new file mode 100644
index 0000000000..33aee70a9b
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_moe.py
@@ -0,0 +1,201 @@
+"""Unified grouped NVFP4 MoE forward for DeepSeek-V4 experts (config-gated, dsv4_fp4_grouped_mode).
+
+Contiguous-grouped: tokens sorted by expert, padded per-expert to TILE; ONE grouped gate_up GEMM
+-> clamped-SwiGLU -> ONE grouped down GEMM, with GPU-vectorized routing/pack/scatter (no per-expert
+Python loop). The base GEMM auto-dispatches to the best fp4 path:
+    Marlin W4A16 (sm120, pad-64)  ->  DeepGEMM (sm90/sm100)  ->  CUTLASS grouped (sm120, pad-128)
+    ->  chunked-dequant (any GPU, fallback).
+Marlin W4A16 (bf16 activations, bit-correct) is preferred on sm120: ~1.79x faster than CUTLASS and
+no activation-quant error. See marlin_w4a16/.
+
+OFF unless cfg.dsv4_fp4_grouped_mode is set; existing fused-Triton/eager paths untouched.
+Bench (RTX PRO 6000, H=4096 I=2048 top6, FORWARD): vs chunked-dequant E=32 1.6-1.9x, E=256
+3.5-3.7x (single-launch grouped vs chunked's per-chunk loop).
+
+STATUS: forward path (base + clamped-swiglu + routing) implemented + validated. TRAINING is the
+remaining piece — the cutlass/deepgemm engines are forward-only, so dX (hidden grad) and
+LoRA-on-experts backward need autograd.Function wrappers (tasks #14, #18, #19). Until then this
+path is correct for inference/eval; training still uses the fused-Triton path.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+TILE = 128
+_ENGINE_CACHE: dict = {}
+
+
+def grouped_fp4_backend(mode: str) -> str | None:
+    """Best available base-GEMM backend: 'marlin' (sm120 W4A16, preferred) | 'cutlass' (sm120) |
+    'deepgemm' (sm90/100) | 'chunked'. On sm120 Marlin W4A16 is ~1.79x faster than CUTLASS and
+    bit-correct (bf16 activations, no activation quantization)."""
+    try:
+        from .marlin_w4a16 import marlin_w4a16_available
+
+        if mode == "nvfp4" and marlin_w4a16_available():
+            return "marlin"
+    except Exception:
+        pass
+    try:
+        from .cutlass_fp4 import cutlass_fp4_available
+
+        if mode == "nvfp4" and cutlass_fp4_available():
+            return "cutlass"
+    except Exception:
+        pass
+    try:
+        from .dequant_grouped import deepgemm_grouped_available
+
+        if deepgemm_grouped_available():
+            return "deepgemm"
+    except Exception:
+        pass
+    return "chunked" if torch.cuda.is_available() else None
+
+
+def _route(idx, E, dev, tile=TILE):
+    flat = idx.reshape(-1)
+    order = flat.argsort()
+    rep = torch.arange(idx.size(0), device=dev).repeat_interleave(idx.size(1))[order]
+    exp_sorted = flat[order]
+    counts = torch.bincount(flat, minlength=E)
+    ptiles = (counts + tile - 1) // tile
+    roff = torch.cat([ptiles.new_zeros(1), ptiles.cumsum(0)]) * tile
+    coff = torch.cat([counts.new_zeros(1), counts.cumsum(0)])
+    local = torch.arange(exp_sorted.numel(), device=dev) - coff[exp_sorted]
+    padded_row = roff[exp_sorted] + local
+    m_indices = torch.repeat_interleave(
+        torch.arange(E, dtype=torch.int32, device=dev), ptiles
+    )
+    Mt = int(ptiles.sum()) * tile
+    return rep, padded_row, m_indices, counts, Mt
+
+
+def _engine(Mt, N, K, E, mode):
+    key = (Mt, N, K, E, mode)
+    eng = _ENGINE_CACHE.get(key)
+    if eng is None:
+        from .cutlass_fp4.grouped import GroupedFp4Gemm
+
+        _ENGINE_CACHE[key] = eng = GroupedFp4Gemm(Mt, N, K, E, mode)
+    return eng
+
+
+def _quant_weight(W_nv, mode):
+    """Per-expert weight quant for the grouped engines. W_nv: NVFP4Tensor [E,N,K]."""
+    E = W_nv.qdata.size(0)
+    if mode == "nvfp4":
+        from .cutlass_fp4.quant_nvfp4 import nvfp4_quant
+
+        deq = W_nv.dequantize(torch.bfloat16)
+        qs = [nvfp4_quant(deq[e].contiguous()) for e in range(E)]
+        return torch.stack([q for q, _ in qs]), torch.stack([s for _, s in qs])
+    from torchao.prototype.mx_formats.mx_tensor import MXTensor
+
+    deq = W_nv.dequantize(torch.bfloat16)
+    ts = [
+        MXTensor.to_mx(deq[e].contiguous(), torch.float4_e2m1fn_x2, 32)
+        for e in range(E)
+    ]
+    return torch.stack([t.qdata for t in ts]), torch.stack([t.scale for t in ts])
+
+
+@torch.no_grad()
+def grouped_fp4_moe_forward(
+    hidden, idx, wts, gate_up_nv, down_nv, limit, mode, backend=None, act_type="silu"
+):
+    """Forward-only grouped NVFP4 MoE. hidden[N,H], idx/wts[N,topk], experts NVFP4Tensor.
+
+    Returns [N,H]. backend auto-selected if None. TRAINING NOT YET (forward-only engines).
+    act_type: 'silu' (DSV4 clamped SwiGLU, default) or 'gelu_tanh' (Gemma4 GeGLU, no clamp).
+    """
+    N, H = hidden.shape
+    E = gate_up_nv.qdata.size(0)
+    _Idim = down_nv.qdata.size(2) * 2  # down K = I (packed K/2)
+    dev = hidden.device
+    backend = backend or grouped_fp4_backend(mode)
+    if backend == "marlin":
+        from .marlin_w4a16.backend import MARLIN_TILE
+
+        tile = MARLIN_TILE
+    else:
+        tile = TILE
+    rep, padded_row, m_indices, counts, Mt = _route(idx, E, dev, tile)
+    wflat = wts.reshape(-1)[idx.reshape(-1).argsort()]
+
+    A = hidden.new_zeros(Mt, H)
+    A[padded_row] = hidden[rep]
+
+    marlin_base = None
+    if backend == "marlin":
+        from .marlin_w4a16.backend import build_marlin_forward_base, marlin_base_forward
+
+        marlin_base = build_marlin_forward_base(gate_up_nv, down_nv)
+        gu = marlin_base_forward(marlin_base, 0, A, m_indices).float()
+    elif backend == "cutlass":
+        from .cutlass_fp4.grouped import quant_act
+
+        gu_eng = _engine(Mt, 2 * (down_nv.qdata.size(2) * 2), H, E, mode)  # N = 2I
+        gu_eng.set_weights(*_quant_weight(gate_up_nv, mode))
+        aq, as_ = quant_act(A, mode)
+        gu = gu_eng.forward(aq.unsqueeze(0), as_.unsqueeze(0), m_indices).float()
+    elif (
+        backend == "deepgemm"
+    ):  # native fp8(act) x mxfp4(weight) grouped GEMM (SM90/SM100)
+        from .dequant_grouped import _cached_mxfp4, deepgemm_grouped_fp8_fp4
+
+        wq, ws = _cached_mxfp4(gate_up_nv, _per_tensor(gate_up_nv, E))
+        gu = deepgemm_grouped_fp8_fp4(
+            A, wq, ws, m_indices.repeat_interleave(TILE)
+        ).float()
+    else:  # chunked fallback: dequant weight -> grouped bf16 matmul
+        from .dequant_grouped import nvfp4_dequant_bf16
+
+        pt = _per_tensor(gate_up_nv, E)
+        Wb = nvfp4_dequant_bf16(gate_up_nv.qdata, gate_up_nv.scale, pt)
+        offs = (torch.bincount(m_indices, minlength=E) * TILE).cumsum(0).to(torch.int32)
+        gu = torch._grouped_mm(A, Wb.transpose(1, 2), offs=offs).float()
+
+    g, u = gu.chunk(2, dim=-1)
+    if act_type == "gelu_tanh":
+        h = (F.gelu(g, approximate="tanh") * u).to(hidden.dtype)
+    else:
+        h = (F.silu(g.clamp(max=limit)) * u.clamp(min=-limit, max=limit)).to(
+            hidden.dtype
+        )
+
+    if backend == "marlin":
+        dn = marlin_base_forward(marlin_base, 1, h.contiguous(), m_indices)
+    elif backend == "cutlass":
+        from .cutlass_fp4.grouped import quant_act
+
+        dn_eng = _engine(Mt, H, down_nv.qdata.size(2) * 2, E, mode)  # K = I
+        dn_eng.set_weights(*_quant_weight(down_nv, mode))
+        hq, hs = quant_act(h.contiguous(), mode)
+        dn = dn_eng.forward(hq.unsqueeze(0), hs.unsqueeze(0), m_indices)
+    elif backend == "deepgemm":
+        from .dequant_grouped import _cached_mxfp4, deepgemm_grouped_fp8_fp4
+
+        wq, ws = _cached_mxfp4(down_nv, _per_tensor(down_nv, E))
+        dn = deepgemm_grouped_fp8_fp4(
+            h.contiguous(), wq, ws, m_indices.repeat_interleave(TILE)
+        )
+    else:
+        from .dequant_grouped import nvfp4_dequant_bf16
+
+        pt = _per_tensor(down_nv, E)
+        Wb = nvfp4_dequant_bf16(down_nv.qdata, down_nv.scale, pt)
+        offs = (torch.bincount(m_indices, minlength=E) * TILE).cumsum(0).to(torch.int32)
+        dn = torch._grouped_mm(h.contiguous(), Wb.transpose(1, 2), offs=offs)
+
+    out = hidden.new_zeros(N, H)
+    return out.index_add(
+        0, rep, (dn[padded_row] * wflat[:, None].to(dn.dtype)).to(out.dtype)
+    )
+
+
+def _per_tensor(w_nv, E):
+    pt = getattr(w_nv, "per_tensor_scale", None)
+    return torch.ones(E, device="cuda") if pt is None else pt.reshape(-1).float()
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_train.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_train.py
new file mode 100644
index 0000000000..710ea92cd2
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/grouped_train.py
@@ -0,0 +1,596 @@
+"""Training-capable grouped NVFP4 MoE (fwd+bwd) for DeepSeek-V4 experts — the config-gated
+cutlass-fp4 path. Forward: cutlass fp4 grouped GEMM (fast). Backward: chunked bf16-dequant +
+cuBLAS grouped_mm (accurate, bounded memory) — this beat the existing fused Triton kernel at the
+real E=256 scale on BOTH speed (2.24x) and memory (0.87x). LoRA-on-experts fused via
+scatter2scatter single-launch grouped GEMMs. All cute imports lazy (module loads clean off sm120).
+
+Backward design rationale (explored exhaustively): bf16 grad is REQUIRED (fp8 grad = 33% error);
+dequant+cuBLAS beats the fused in-kernel-decode Triton dX 5.7x (Triton GEMM << cuBLAS); chunking
+over expert-groups bounds the bf16-weight transient.
+
+LoRA GEMM design: scatter2scatter kernels (kernels/ops.py) operate directly on the TILE-padded
+expert-sorted layout with x_grouped=y_grouped=True — single launch, ragged-native, no padding
+transient.  grouped_lora_fwd/bwd in grouped_lora.py encapsulate the entire implementation;
+no selection or threshold logic lives here.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+from .grouped_lora import grouped_lora_bwd, grouped_lora_fwd
+
+LOG = get_logger(__name__)
+
+_BACKEND_LOGGED = False  # one-time log of the resolved base-GEMM backend
+
+
+# thin compat wrapper over the centralized runtime module. grouped_backend (cfg.moe_grouped_backend):
+# None/"auto" = capability auto-select; marlin|cutlass|deepgemm = force if available (else warn +
+# auto); "dequant" = force the chunked-dequant fallback.
+from .runtime import RUNTIME  # noqa: E402
+
+
+def set_grouped_backend_override(backend) -> None:
+    RUNTIME.grouped_backend = str(backend).lower() if backend else None
+
+
+def _backend_available(name: str) -> bool:
+    try:
+        if name == "marlin":
+            from .marlin_w4a16 import marlin_w4a16_available
+
+            return marlin_w4a16_available()
+        if name == "cutlass":
+            from .cutlass_fp4 import cutlass_fp4_available
+
+            return cutlass_fp4_available()
+        if name == "deepgemm":
+            from .dequant_grouped import deepgemm_grouped_available
+
+            return deepgemm_grouped_available()
+    except Exception:
+        return False
+    return False
+
+
+def _auto_backend() -> str | None:
+    """Capability + arch auto-select. Order is arch-aware so each GPU class gets its tuned default
+    while Marlin (the only fused W4A16 path that runs on Ampere/Ada) backs up everything:
+      - sm120 (consumer Blackwell): Marlin (~1.79x CUTLASS, bit-correct) > CUTLASS > DeepGEMM
+      - sm100 (datacenter Blackwell, e.g. B200): DeepGEMM (tuned fp8-act x mxfp4) > Marlin
+      - sm90 (Hopper) / sm80 / sm89 (Ampere / Ada): Marlin only (the DeepGEMM fp8xfp4 grouped
+        kernel is sm100-only; deepgemm_grouped_available() returns False below sm100 -> Marlin)
+    Marlin stays force-selectable everywhere via cfg.moe_grouped_backend."""
+    import torch
+
+    major = torch.cuda.get_device_capability()[0] if torch.cuda.is_available() else 0
+    order: tuple[str, ...]
+    if major >= 11:
+        order = ("marlin", "cutlass", "deepgemm")
+    elif major in (9, 10):
+        order = ("deepgemm", "marlin")
+    else:
+        order = ("marlin",)
+    for name in order:
+        if _backend_available(name):
+            return name
+    return None
+
+
+TILE = 128
+# backward base-dX dequant chunk; bounds the bf16-weight transient.
+# E=256 knee (B200, BK=1024 dequant): CHUNK_E=16 = 11.2ms bwd / 2.6GB peak vs 10.1ms floor.
+CHUNK_E = 16
+_ENGINES: dict = {}
+
+
+def grouped_fp4_available(mode: str) -> bool:
+    """True iff the grouped fp4 training path can run here for `mode`: nvfp4 on sm120 (CUTLASS
+    fused-decode) or sm90/sm100 (DeepGEMM fp8-act x mxfp4-weight). Backward is GPU-agnostic
+    (chunked bf16-dequant). chunked-only fallback tracked separately."""
+    if mode != "nvfp4":
+        return False
+    try:
+        from .marlin_w4a16 import marlin_w4a16_available
+
+        if marlin_w4a16_available():
+            return True
+    except Exception:
+        pass
+    try:
+        from .cutlass_fp4 import cutlass_fp4_available
+
+        if cutlass_fp4_available():
+            return True
+    except Exception:
+        pass
+    try:
+        from .dequant_grouped import deepgemm_grouped_available
+
+        return deepgemm_grouped_available()
+    except Exception:
+        return False
+
+
+def _train_backend(mode: str) -> str | None:
+    """Base-GEMM backend for the training forward: 'marlin' (sm120, W4A16 bf16-act — preferred) |
+    'cutlass' (sm120, W4A4 fp4-act) | 'deepgemm' (sm90/100), or None for the chunked-dequant
+    fallback. Auto-selects by capability unless cfg.moe_grouped_backend forced one (see
+    set_grouped_backend_override): an unavailable forced backend warns and falls back to auto;
+    'dequant' forces the fallback (None)."""
+    if mode != "nvfp4":
+        return None
+    override = RUNTIME.grouped_backend
+    if override and override != "auto":
+        if override == "dequant":
+            return None  # chunked-dequant fallback
+        if _backend_available(override):
+            return override
+        LOG.warning(
+            "moe_grouped_backend=%r is not available on this GPU; falling back to auto-select.",
+            override,
+        )
+    return _auto_backend()
+
+
+def _gmm(a, b, offs):
+    return torch._grouped_mm(a, b, offs=offs)
+
+
+def _swiglu(gu, limit, act_type="silu"):
+    from .cutlass_fp4.swiglu import swiglu_fwd
+
+    return swiglu_fwd(gu, limit, act_type)
+
+
+def _swiglu_bwd(dh, gu, limit, act_type="silu"):
+    from .cutlass_fp4.swiglu import swiglu_bwd
+
+    return swiglu_bwd(gu, dh, limit, act_type)
+
+
+def _base_forward(base, which, x, m_indices, mode):
+    """Frozen-expert base GEMM for gate_up (which=0) or down (which=1). `base` is
+    ('marlin', (gu_w, dn_w), dims, ws) | ('cutlass', gu_eng, dn_eng) | ('deepgemm', (guq,gus), (dnq,dns))."""
+    if base[0] == "marlin":
+        from .marlin_w4a16.backend import marlin_base_forward
+
+        return marlin_base_forward(base, which, x, m_indices)
+    backend, gw, dw = base
+    if backend == "cutlass":
+        from .cutlass_fp4.grouped import quant_act
+
+        eng = gw if which == 0 else dw
+        aq, as_ = quant_act(x, mode)
+        return eng.forward(aq.unsqueeze(0), as_.unsqueeze(0), m_indices)
+    from .dequant_grouped import deepgemm_grouped_fp8_fp4
+
+    wq, ws = gw if which == 0 else dw
+    # cutlass uses per-tile m_indices; DeepGEMM's grouped_layout is per-row (length Mt)
+    return deepgemm_grouped_fp8_fp4(x, wq, ws, m_indices.repeat_interleave(TILE))
+
+
+def _engine(Mt, N, K, E, mode):
+    key = (Mt, N, K, E, mode)
+    eng = _ENGINES.get(key)
+    if eng is None:
+        from .cutlass_fp4.grouped import GroupedFp4Gemm
+
+        _ENGINES[key] = eng = GroupedFp4Gemm(Mt, N, K, E, mode)
+    return eng
+
+
+def _fp8_read_dx_ok():
+    """fp8-read backward dX (#3744) wins on sm120 (bandwidth-bound: half weight bytes -> ~1.5x +
+    half memory). On sm100 it's speed-neutral (cuBLAS bf16 is fast) so keep the bf16 path there."""
+    if not torch.cuda.is_available():
+        return False
+    return torch.cuda.get_device_capability()[0] == 12
+
+
+def _base_dx(
+    w_nv, pt, g, out_k, offs, m_indices, prefer_fp8=True, tile=TILE, _marlin_raw=None
+):
+    """Chunked base-weight contraction g @ W (frozen NVFP4 experts). bf16 path: dequant CHUNK_E
+    experts to bf16 + cuBLAS grouped_mm. fp8-read path (sm120): dequant to fp8 (half bytes) + a
+    Triton grouped GEMM that reads fp8 and upcasts in-register (~1.5x faster, half the transient).
+    g[Mt,N], W[E,N,out_k] -> [Mt,out_k]; m_indices is the per-`tile` expert id (for the fp8 GEMM).
+    `tile` is the routing pad granularity (128 cutlass, 64 marlin) and is passed to the fp8 dX
+    kernel as its row-block (BM) so one expert maps to each block. Both fp8-read and bf16-dequant
+    are gradient-consistent (they dequant the FORWARD weight); fp8-read is the faster default.
+    _marlin_raw: optional (qw_flat_E, orig_scale, pt_E) for the marlin memory-free path — reads from
+    the marlin qweight cache via fused Triton dequant instead of the freed nv.qdata."""
+    from .dequant_grouped import nvfp4_dequant_bf16
+
+    fp8 = _fp8_read_dx_ok() and prefer_fp8
+    if fp8:
+        from .dequant_grouped import grouped_dx_fp8, nvfp4_dequant_fp8
+
+    if _marlin_raw is not None:
+        # Marlin memory-free path: dequant from marlin int32 layout via fused Triton kernel.
+        # _marlin_raw = (qw [E, words_per_expert] int32, orig_scale [E, N, K//16] fp8, pt [E] f32)
+        from .marlin_w4a16.backend import _build_base_scatter
+        from .marlin_w4a16.fused_dequant import marlin_dequant_bf16, marlin_dequant_fp8
+        from .mx_weights import fp4_codebook
+
+        qw, orig_scale, pt_e = _marlin_raw
+        # pt_e may have stride(0) from .expand(); Triton needs element-stride=1 for a correct load.
+        if pt_e.stride(0) == 0:
+            pt_e = pt_e.contiguous()
+        E, N, Kg = orig_scale.shape
+        K = Kg * 16
+        scatter_lut = _build_base_scatter(qw.device)
+        cb = fp4_codebook(qw.device).float()
+        starts = torch.cat([offs.new_zeros(1), offs]).tolist()
+        out = g.new_empty(g.size(0), out_k)
+        for c0 in range(0, E, CHUNK_E):
+            c1 = min(c0 + CHUNK_E, E)
+            t0, t1 = starts[c0], starts[c1]
+            if t1 == t0:
+                continue
+            C = c1 - c0
+            if fp8:
+                Wc = marlin_dequant_fp8(
+                    qw[c0:c1].reshape(C, -1),
+                    orig_scale[c0:c1],
+                    pt_e[c0:c1],
+                    scatter_lut,
+                    cb,
+                    N,
+                    K,
+                    C,
+                )
+                mi = (m_indices[t0 // tile : t1 // tile] - c0).to(torch.int32)
+                out[t0:t1] = grouped_dx_fp8(g[t0:t1], Wc, mi, tile)
+            else:
+                Wc = marlin_dequant_bf16(
+                    qw[c0:c1].reshape(C, -1),
+                    orig_scale[c0:c1],
+                    pt_e[c0:c1],
+                    scatter_lut,
+                    cb,
+                    N,
+                    K,
+                    C,
+                )
+                loc = (offs[c0:c1] - t0).to(torch.int32)
+                out[t0:t1] = _gmm(g[t0:t1], Wc, loc)
+        return out
+
+    qdata, scale = w_nv.qdata, w_nv.scale
+    E = qdata.size(0)
+    starts = torch.cat([offs.new_zeros(1), offs]).tolist()  # one sync, not per-iter
+    out = g.new_empty(g.size(0), out_k)
+    for c0 in range(0, E, CHUNK_E):
+        c1 = min(c0 + CHUNK_E, E)
+        t0, t1 = starts[c0], starts[c1]
+        if t1 == t0:
+            continue
+        if fp8:
+            Wc = nvfp4_dequant_fp8(qdata[c0:c1], scale[c0:c1], pt[c0:c1])
+            mi = (m_indices[t0 // tile : t1 // tile] - c0).to(torch.int32)
+            out[t0:t1] = grouped_dx_fp8(g[t0:t1], Wc, mi, tile)
+        else:
+            loc = (offs[c0:c1] - t0).to(torch.int32)
+            Wc = nvfp4_dequant_bf16(qdata[c0:c1], scale[c0:c1], pt[c0:c1])
+            out[t0:t1] = _gmm(g[t0:t1], Wc, loc)
+    return out
+
+
+def _pt(nv, E, dev):
+    p = getattr(nv, "per_tensor_scale", None)
+    if p is None:
+        return torch.ones(E, device=dev)
+    p = p.reshape(-1).float()
+    return p.expand(E) if p.numel() == 1 else p
+
+
+def _has_nonunit_pt(*nvs) -> bool:
+    """True iff any NVFP4 weight carries a non-unit per_tensor_scale (weight_scale_2). The cutlass
+    forward (set_weights) folds only the E4M3 block scale, dropping weight_scale_2, while the
+    backward applies it, giving a wrong forward + grad mismatch when it != 1 (the real DSV4 ckpt).
+    Folding weight_scale_2 into the cutlass weight + saved swiglu input is unimplemented, so the
+    backend selection skips cutlass in that case."""
+    for nv in nvs:
+        p = getattr(nv, "per_tensor_scale", None)
+        if p is None:
+            continue
+        if not torch.allclose(
+            p.reshape(-1).float(), torch.ones((), device=p.device, dtype=torch.float32)
+        ):
+            return True
+    return False
+
+
+class _GroupedExperts(torch.autograd.Function):
+    """x[Mt,H] -> gate_up(cutlass fp4 base + LoRA) -> gated-activation -> down(...) -> [Mt,H].
+    Frozen NVFP4 experts; trainable LoRA A/B (stacked [E,r,K]/[E,N,r]).
+    act_type: 'silu' (DSV4 clamped SwiGLU) or 'gelu_tanh' (Gemma4 GeGLU)."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        x,
+        base,
+        weight_recipe,
+        Agu,
+        Bgu,
+        Adn,
+        Bdn,
+        m_indices,
+        offs,
+        scaling,
+        limit,
+        mode,
+        act_type,
+        prefer_fp8_dx=True,
+    ):
+        E = Agu.size(0)
+        gu = _base_forward(base, 0, x, m_indices, mode)
+        # LoRA-B GEMM folds the base output in via residual=gu -> gu = base + lora, no separate add.
+        gu, xAg = grouped_lora_fwd(x, Agu, Bgu, scaling, offs, E, residual=gu)
+        h = _swiglu(gu, limit, act_type).to(x.dtype)
+        dn = _base_forward(base, 1, h, m_indices, mode)
+        dn, hAd = grouped_lora_fwd(h, Adn, Bdn, scaling, offs, E, residual=dn)
+        ctx.save_for_backward(x, Agu, Bgu, Adn, Bdn, offs, gu, h, xAg, hAd, m_indices)
+        # FSDP-safe: don't pin the gathered NVFP4 weight; re-read the (re-gathered) param in backward
+        ctx.weight_recipe, ctx.scaling, ctx.limit = weight_recipe, scaling, limit
+        ctx.act_type = act_type
+        # base dX: fp8-read (fast, ~2% grad) vs bf16-dequant (~0.5%)
+        ctx.prefer_fp8_dx = prefer_fp8_dx
+        # Marlin memory-free path: build_marlin_forward_base freed nv.qdata and saved (qdata, scale,
+        # pt) in the cache (single-GPU only); stash them so backward skips weight_recipe() (which
+        # would return the emptied NVFP4 tensor).
+        ctx.bwd_marlin = None
+        if base[0] == "marlin":
+            from .marlin_w4a16.backend import marlin_bwd_data
+
+            ctx.bwd_marlin = marlin_bwd_data(base)
+        return dn
+
+    @staticmethod
+    def backward(ctx, d_dn):
+        x, Agu, Bgu, Adn, Bdn, offs, gu, h, xAg, hAd, m_indices = ctx.saved_tensors
+        s, lim, act_type = ctx.scaling, ctx.limit, ctx.act_type
+        pf8 = _fp8_read_dx_ok() and ctx.prefer_fp8_dx
+        tile = (
+            x.size(0) // m_indices.numel()
+        )  # routing pad granularity (128 cutlass, 64 marlin)
+
+        d_dn = d_dn.contiguous().to(x.dtype)
+
+        if ctx.bwd_marlin is not None:
+            # Marlin memory-free path: nv.qdata was freed after repack; backward reads from the
+            # marlin qweight cache via fused Triton dequant (marlin int32 + original scales).
+            # bwd_marlin = ((gu_qw, (gu_scale, gu_pt, E, N, K)), (dn_qw, (dn_scale, ...)))
+            (gu_qw, gu_bwd), (dn_qw, dn_bwd) = ctx.bwd_marlin
+            dn_scale, dn_pt_e, _, _, _ = dn_bwd
+            dn_marlin_raw = (dn_qw, dn_scale, dn_pt_e)
+            dh = _base_dx(
+                None,
+                None,
+                d_dn,
+                h.size(1),
+                offs,
+                m_indices,
+                pf8,
+                tile,
+                _marlin_raw=dn_marlin_raw,
+            )
+        else:
+            gu_nv, dn_nv = ctx.weight_recipe()
+            E, dev = gu_nv.qdata.size(0), x.device
+            ptg, ptd = _pt(gu_nv, E, dev), _pt(dn_nv, E, dev)
+            dh = _base_dx(dn_nv, ptd, d_dn, h.size(1), offs, m_indices, pf8, tile)
+
+        E = Agu.size(0)
+        # dX_lora GEMM folds base dh in via residual=dh -> dh = base_dX + lora_dX.
+        dh, dAdn, dBdn = grouped_lora_bwd(
+            d_dn, h, Adn, Bdn, hAd, s, offs, E, residual=dh
+        )
+        dgu = _swiglu_bwd(dh, gu, lim, act_type).to(x.dtype)
+        del dh
+
+        if ctx.bwd_marlin is not None:
+            gu_scale, gu_pt_e, _, _, _ = ctx.bwd_marlin[0][1]
+            gu_marlin_raw = (gu_qw, gu_scale, gu_pt_e)
+            dx = _base_dx(
+                None,
+                None,
+                dgu,
+                x.size(1),
+                offs,
+                m_indices,
+                pf8,
+                tile,
+                _marlin_raw=gu_marlin_raw,
+            )
+        else:
+            dx = _base_dx(gu_nv, ptg, dgu, x.size(1), offs, m_indices, pf8, tile)
+
+        dx, dAgu, dBgu = grouped_lora_bwd(
+            dgu, x, Agu, Bgu, xAg, s, offs, E, residual=dx
+        )
+        # grads align to forward args: x, base, weight_recipe, Agu, Bgu, Adn, Bdn,
+        # m_indices, offs, scaling, limit, mode, act_type, prefer_fp8_dx
+        return (
+            dx,
+            None,
+            None,
+            dAgu,
+            dBgu,
+            dAdn,
+            dBdn,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+def _lora_stack(lora, E, K, out):
+    """scattermoe LoRA (A[r*E,K], B[out,r*E], scaling) -> stacked (Agu[E,r,K], Bgu[E,out,r], s)."""
+    A, B, scaling = lora
+    r = A.shape[0] // E
+    As = A.reshape(E, r, K).contiguous()
+    Bs = B.reshape(out, E, r).permute(1, 0, 2).contiguous()
+    return As, Bs, float(scaling)
+
+
+def grouped_fp4_moe_train(
+    hidden,
+    idx,
+    wts,
+    gate_up_nv,
+    down_nv,
+    gup_lora,
+    down_lora,
+    limit,
+    mode,
+    act_type="silu",
+    weight_recipe=None,
+    mxfp4_cache=None,
+    prefer_fp8_dx=True,
+):
+    """Training-capable grouped NVFP4 MoE forward. hidden[N,H], idx/wts[N,topk]; experts NVFP4Tensor;
+    *_lora = (A,B,scaling) scattermoe layout. Returns [N,H]; differentiable to hidden + LoRA A/B.
+    act_type: 'silu' (DSV4 clamped SwiGLU, default) or 'gelu_tanh' (Gemma4 GeGLU, no clamp).
+    weight_recipe: optional callable -> (gate_up_nv, down_nv) re-read for the FSDP-safe backward
+    (defaults to the forward tensors). mxfp4_cache: optional persistent dict (e.g. on the owning
+    module) so the DeepGEMM backend requantizes the frozen weight once across FSDP re-gathers.
+    prefer_fp8_dx: base dX backward — True = fp8-read (fast, ~2% grad error, sm120 default);
+    False = bf16-dequant (slower, ~0.5%, max gradient fidelity)."""
+    if weight_recipe is None:
+        weight_recipe = lambda: (gate_up_nv, down_nv)  # noqa: E731
+    N, H = hidden.shape
+    # NVFP4Tensor.shape reports full-precision dims and survives the marlin qdata-free path (qdata
+    # replaced with empty tensor). SimpleNamespace mocks (unit tests) lack .shape, so fall through
+    # to the qdata.size() branch.
+    _qd = gate_up_nv.qdata
+    if _qd.numel() > 0:
+        E = _qd.size(0)
+        twoI = _qd.size(1)
+        I = down_nv.qdata.size(2) * 2  # noqa: E741  (down K dim, packed K/2 * 2)
+    else:
+        # Marlin qdata-free path: qdata was freed; use the NVFP4Tensor wrapper shape.
+        _gu_shape = gate_up_nv.shape
+        E, twoI = int(_gu_shape[0]), int(_gu_shape[1])
+        I = int(down_nv.shape[2])  # noqa: E741
+    dev = hidden.device
+    backend = _train_backend(mode)
+    # cutlass weight_scale_2 (per_tensor_scale) folding is unimplemented: set_weights drops it on the
+    # forward while the backward applies it, giving a wrong forward + grad mismatch when it != 1.
+    # Marlin and DeepGEMM both fold _pt() into the weight, so prefer one of those (else hard-error)
+    # for non-unit scales.
+    if backend == "cutlass" and _has_nonunit_pt(gate_up_nv, down_nv):
+        for _alt in ("marlin", "deepgemm"):
+            if _backend_available(_alt):
+                backend = _alt
+                break
+        else:
+            major = (
+                torch.cuda.get_device_capability()[0]
+                if torch.cuda.is_available()
+                else 0
+            )
+            raise RuntimeError(
+                "grouped NVFP4 MoE: cutlass was selected but the weight has a non-unit "
+                "per_tensor_scale (weight_scale_2) that the cutlass forward cannot fold "
+                f"(unimplemented). No weight_scale_2-correct backend resolved on sm{major}x "
+                "(marlin/deepgemm unavailable). Run on a GPU with marlin or deepgemm, or "
+                "requantize the experts with a unit per_tensor_scale."
+            )
+    if not _BACKEND_LOGGED:
+        globals()["_BACKEND_LOGGED"] = True
+        LOG.info(
+            "grouped fp4 MoE base-GEMM backend: %s",
+            backend or "cutlass",
+        )
+    # Marlin (sm120 W4A16) pads to 64, half CUTLASS's 128 at thin-M (the padding is the cost since
+    # each expert weight is read once either way), and its bf16-act kernel is bit-correct + faster.
+    if backend == "marlin":
+        from .marlin_w4a16.backend import MARLIN_TILE
+
+        tile = MARLIN_TILE
+    else:
+        tile = TILE
+    flat = idx.reshape(-1)
+    # The DeepEP local path tags remote-routed slots with -1; drop them (bincount/grouping require
+    # expert ids >= 0, and a dropped slot must contribute nothing to its token — correct, since that
+    # (token, slot) is handled on the rank that owns the expert). No-op for the non-EP path (no -1).
+    _tok = torch.arange(N, device=dev).repeat_interleave(idx.size(1))
+    _wf = wts.reshape(-1)
+    if (flat < 0).any():
+        keep = flat >= 0
+        flat = flat[keep]
+        _tok = _tok[keep]
+        _wf = _wf[keep]
+    order = flat.argsort()
+    rep = _tok[order]
+    wflat = _wf[order]
+    exp_sorted = flat[order]
+    counts = torch.bincount(flat, minlength=E)
+    ptiles = (counts + tile - 1) // tile
+    roff = torch.cat([ptiles.new_zeros(1), ptiles.cumsum(0)]) * tile
+    coff = torch.cat([counts.new_zeros(1), counts.cumsum(0)])
+    padded_row = roff[exp_sorted] + (
+        torch.arange(exp_sorted.numel(), device=dev) - coff[exp_sorted]
+    )
+    m_indices = torch.repeat_interleave(
+        torch.arange(E, dtype=torch.int32, device=dev), ptiles
+    )
+    offs = (ptiles * tile).cumsum(0).to(torch.int32)
+    Mt = int(ptiles.sum()) * tile
+
+    if backend == "marlin":
+        from .marlin_w4a16.backend import build_marlin_forward_base
+
+        base = build_marlin_forward_base(gate_up_nv, down_nv, mxfp4_cache)
+    elif backend == "deepgemm":
+        from .dequant_grouped import _cached_mxfp4
+
+        base = (
+            "deepgemm",
+            _cached_mxfp4(gate_up_nv, _pt(gate_up_nv, E, dev), mxfp4_cache, "gate_up"),
+            _cached_mxfp4(down_nv, _pt(down_nv, E, dev), mxfp4_cache, "down"),
+        )
+    else:
+        gu_eng = _engine(Mt, twoI, H, E, mode)
+        gu_eng.set_weights(gate_up_nv.qdata, gate_up_nv.scale)
+        dn_eng = _engine(Mt, H, I, E, mode)
+        dn_eng.set_weights(down_nv.qdata, down_nv.scale)
+        base = ("cutlass", gu_eng, dn_eng)
+    Agu, Bgu, sgu = _lora_stack(gup_lora, E, H, twoI)
+    Adn, Bdn, sdn = _lora_stack(down_lora, E, I, H)
+    assert sgu == sdn, "gate_up/down LoRA scaling must match"
+
+    lim = (
+        float(limit) if limit is not None else 1e30
+    )  # no clamp when the model has no swiglu_limit
+    A = hidden.new_zeros(Mt, H).index_copy(0, padded_row, hidden[rep])
+    dn = _GroupedExperts.apply(
+        A,
+        base,
+        weight_recipe,
+        Agu,
+        Bgu,
+        Adn,
+        Bdn,
+        m_indices,
+        offs,
+        sgu,
+        lim,
+        mode,
+        act_type,
+        prefer_fp8_dx,
+    )
+    out = hidden.new_zeros(N, H)
+    return out.index_add(
+        0, rep, (dn[padded_row] * wflat[:, None].to(dn.dtype)).to(out.dtype)
+    )
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/__init__.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/__init__.py
new file mode 100644
index 0000000000..eb502db712
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/__init__.py
@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Original work Copyright (c) Shawn Tan and ScatterMoE Contributors
+# Adapted from https://github.com/shawntan/scattermoe
+# See https://github.com/shawntan/scattermoe/blob/main/LICENSE
+#
+# Modifications and LoRA adaptation Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+from . import lora_ops, ops
+
+__all__ = ["ops", "lora_ops"]
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/grouped_gram.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/grouped_gram.py
new file mode 100644
index 0000000000..db65e7e644
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/grouped_gram.py
@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Grouped-Gram LoRA weight gradients (dA/dB) without per-output-block recompute.
+
+The split kernel in ``lora_ops`` recomputes XA = X@A^T (resp. YB = dY@B) inside the
+kernel, once per output dim-block, to avoid materializing it. For LoRA's small rank
+those intermediates are tiny (``[M*k, R]``), so materializing them once and reducing
+the per-output-block recompute is a large net win at high expert counts. dA/dB then
+reduce to plain grouped Gram products: dA[g] = scaling * YB_g^T @ X_g,
+dB[g] = scaling * DY_g^T @ XA_g (reduction over the group's tokens).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from . import lora_ops
+from .lora_ops import _block_r_for_rank, _bucket_m
+
+
+def _device_shared_mem_optin() -> int:
+    """Opt-in max dynamic shared memory per block for the current CUDA device (bytes). Returns a
+    large sentinel when CUDA is unavailable so the big-SMEM config is chosen by default."""
+    try:
+        if torch.cuda.is_available():
+            return torch.cuda.get_device_properties(
+                torch.cuda.current_device()
+            ).shared_memory_per_block_optin
+    except Exception:  # pylint: disable=broad-except
+        pass
+    return 1 << 30
+
+
+# The large tile (BLOCK_M=128, BLOCK_WIDE=128, num_stages=3) needs ~144 KiB of shared memory per
+# block: it fits sm_80 (A100, ~163 KiB) and sm_90/sm_100 (H100/B200, ~228 KiB) but NOT sm_89 (L40S)
+# or consumer sm_120 (RTX 6000 / 5090, ~99 KiB), where Triton raises OutOfResources.
+_GRAM_LARGE_CONFIG_SMEM = 147456  # bytes
+
+
+def _grouped_gram_configs():
+    # ONE config sized to the device's shared memory. A single config keeps the autotune "sweep"
+    # instant — sweeping across configs re-times per rank at scattered layers and trips DeepEP's
+    # combine barrier (why this isn't a full per-shape search). Big-SMEM GPUs (H100/B200/A100) run
+    # the larger tile; small-SMEM GPUs (L40S / RTX 6000 / 5090) run a halved BLOCK_WIDE + double
+    # buffering that fits ~99 KiB (<= ~64 KiB even at BLOCK_R=128) instead of OOM-ing.
+    if _device_shared_mem_optin() >= _GRAM_LARGE_CONFIG_SMEM:
+        cfg = triton.Config(
+            {"BLOCK_WIDE": 128, "BLOCK_M": 128},
+            num_warps=4,
+            num_stages=3,
+            num_ctas=1,
+        )
+    else:
+        cfg = triton.Config(
+            {"BLOCK_WIDE": 64, "BLOCK_M": 128},
+            num_warps=4,
+            num_stages=2,
+            num_ctas=1,
+        )
+    return [cfg]
+
+
+@triton.autotune(configs=_grouped_gram_configs(), key=["M_BUCKET", "WIDE", "RANK_IS_I"])
+@triton.jit
+def _grouped_gram_kernel(
+    P_ptr,
+    stride_pm,
+    stride_pd,
+    Q_ptr,
+    stride_qm,
+    stride_qd,
+    OUT_ptr,
+    stride_og,
+    stride_oi,
+    stride_oj,
+    offsets_ptr,
+    M_BUCKET,
+    WIDE: tl.constexpr,
+    RANK: tl.constexpr,
+    scaling,
+    RANK_IS_I: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    BLOCK_WIDE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    allow_tf32: tl.constexpr,
+):
+    """Grouped Gram product: out[g] = scaling * (P_g^T @ Q_g), summed over the
+    tokens of combined-group g (a contiguous [start, end) slice).
+
+    Both LoRA weight gradients are this op once XA/YB are precomputed:
+      dA[g] = scaling * YB_g^T @ X_g   (P=YB[M,R], Q=X[M,K]; rank is the I/row dim)
+      dB[g] = scaling * DY_g^T @ XA_g  (P=DY[M,N], Q=XA[M,R]; rank is the J/col dim)
+
+    The wide dim (K for dA, N for dB) is tiled; the rank dim fits one BLOCK_R.
+    Every (group, wide-block) writes its full tile, so empty groups self-zero --
+    no output pre-init needed. Output strides place the [I, J] tile into either
+    lora_A-grad ([E*T*R, K]) or lora_B-grad ([N, E*T*R]) layout directly.
+    """
+    g = tl.program_id(0)
+    w_blk = tl.program_id(1)
+
+    start = tl.where(g == 0, 0, tl.load(offsets_ptr + g - 1, mask=g > 0, other=0))
+    end = tl.load(offsets_ptr + g)
+    num = end - start
+
+    w_idx = w_blk * BLOCK_WIDE + tl.arange(0, BLOCK_WIDE)
+    w_mask = w_idx < WIDE
+    r_idx = tl.arange(0, BLOCK_R)
+    r_mask = r_idx < RANK
+    input_dtype = P_ptr.dtype.element_ty
+
+    if RANK_IS_I:
+        acc = tl.zeros((BLOCK_R, BLOCK_WIDE), dtype=tl.float32)
+    else:
+        acc = tl.zeros((BLOCK_WIDE, BLOCK_R), dtype=tl.float32)
+
+    for i in range(tl.cdiv(num, BLOCK_M)):
+        m_idx = start + i * BLOCK_M + tl.arange(0, BLOCK_M)
+        m_mask = m_idx < end
+        if RANK_IS_I:
+            p = tl.load(
+                P_ptr + m_idx[:, None] * stride_pm + r_idx[None, :] * stride_pd,
+                mask=m_mask[:, None] & r_mask[None, :],
+                other=0.0,
+            ).to(input_dtype)
+            q = tl.load(
+                Q_ptr + m_idx[:, None] * stride_qm + w_idx[None, :] * stride_qd,
+                mask=m_mask[:, None] & w_mask[None, :],
+                other=0.0,
+            ).to(input_dtype)
+            acc += tl.dot(tl.trans(p), q, allow_tf32=allow_tf32)
+        else:
+            p = tl.load(
+                P_ptr + m_idx[:, None] * stride_pm + w_idx[None, :] * stride_pd,
+                mask=m_mask[:, None] & w_mask[None, :],
+                other=0.0,
+            ).to(input_dtype)
+            q = tl.load(
+                Q_ptr + m_idx[:, None] * stride_qm + r_idx[None, :] * stride_qd,
+                mask=m_mask[:, None] & r_mask[None, :],
+                other=0.0,
+            ).to(input_dtype)
+            acc += tl.dot(tl.trans(p), q, allow_tf32=allow_tf32)
+
+    acc = acc * scaling
+    if RANK_IS_I:
+        out_ptrs = (
+            OUT_ptr
+            + g * stride_og
+            + r_idx[:, None] * stride_oi
+            + w_idx[None, :] * stride_oj
+        )
+        out_mask = r_mask[:, None] & w_mask[None, :]
+    else:
+        out_ptrs = (
+            OUT_ptr
+            + g * stride_og
+            + w_idx[:, None] * stride_oi
+            + r_idx[None, :] * stride_oj
+        )
+        out_mask = w_mask[:, None] & r_mask[None, :]
+    tl.store(out_ptrs, acc.to(OUT_ptr.dtype.element_ty), mask=out_mask)
+
+
+def grouped_lora_weight_grads(
+    grouped_grad_out: torch.Tensor,  # DY  [M*k, N], grouped
+    grouped_x: torch.Tensor,  # X   [M*k, K], grouped
+    yb: torch.Tensor,  # DY@B [M*k, R], grouped (reused from dX path)
+    xa: torch.Tensor,  # X@A^T [M*k, R], grouped (saved from forward)
+    lora_A: torch.Tensor,  # [E*T*R, K]
+    lora_B: torch.Tensor,  # [N, E*T*R]
+    combined_offsets: torch.Tensor,
+    e_total: int,
+    scaling: float,
+):
+    """dA/dB via grouped Gram GEMMs over precomputed XA/YB.
+
+    Avoids the shared split kernel's per-output-block recompute of XA/YB (an
+    inner reduction over K or N, repeated cdiv(wide_dim, BLOCK) times per group);
+    that recompute is what makes the split path scale poorly as tenants grow.
+    """
+    rank = lora_A.size(0) // e_total
+    k_dim = lora_A.size(1)
+    n_dim = lora_B.size(0)
+    block_r = _block_r_for_rank(rank)
+    m_bucket = _bucket_m(max(1, grouped_x.size(0) // e_total))
+
+    dA = torch.empty_like(lora_A)
+    dB = torch.empty_like(lora_B)
+
+    grid_a = lambda meta: (e_total, triton.cdiv(k_dim, meta["BLOCK_WIDE"]))  # noqa: E731
+    _grouped_gram_kernel[grid_a](
+        yb,
+        yb.stride(0),
+        yb.stride(1),
+        grouped_x,
+        grouped_x.stride(0),
+        grouped_x.stride(1),
+        dA,
+        rank * dA.stride(0),
+        dA.stride(0),
+        dA.stride(1),
+        combined_offsets,
+        m_bucket,
+        WIDE=k_dim,
+        RANK=rank,
+        scaling=scaling,
+        RANK_IS_I=True,
+        BLOCK_R=block_r,
+        allow_tf32=lora_ops.ALLOW_TF32,
+    )
+
+    grid_b = lambda meta: (e_total, triton.cdiv(n_dim, meta["BLOCK_WIDE"]))  # noqa: E731
+    _grouped_gram_kernel[grid_b](
+        grouped_grad_out,
+        grouped_grad_out.stride(0),
+        grouped_grad_out.stride(1),
+        xa,
+        xa.stride(0),
+        xa.stride(1),
+        dB,
+        rank * dB.stride(1),
+        dB.stride(0),
+        dB.stride(1),
+        combined_offsets,
+        m_bucket,
+        WIDE=n_dim,
+        RANK=rank,
+        scaling=scaling,
+        RANK_IS_I=False,
+        BLOCK_R=block_r,
+        allow_tf32=lora_ops.ALLOW_TF32,
+    )
+    return dA, dB
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/lora_ops.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/lora_ops.py
new file mode 100644
index 0000000000..672917a6fb
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/lora_ops.py
@@ -0,0 +1,3462 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Fused ScatterMoE + LoRA Triton Kernels
+=======================================
+
+Provides fused forward and backward kernels for ScatterMoE with LoRA adapters.
+
+Forward: Y = X @ W + scaling * (X @ A^T) @ B^T
+Backward (LoRA training, W frozen):
+  - dX = dY @ W^T + scaling * (dY @ B) @ A    (input gradient)
+  - dA = scaling * (dY @ B)^T @ X              (LoRA A gradient)
+  - dB = scaling * dY^T @ (X @ A^T)            (LoRA B gradient)
+
+LoRA weight layout (from PEFT ParamWrapper):
+  - A: [r*E, K]  -- for expert e, rows [e*r : (e+1)*r] give A_e of shape [r, K]
+  - B: [N, r*E]  -- for expert e, cols [e*r : (e+1)*r] give B_e of shape [N, r]
+
+Key design decisions:
+  - The forward kernel fuses X@W and X@A^T in the same K-loop for data reuse on X,
+    then computes (X@A^T) @ B^T in the epilogue.
+  - The backward dA/dB kernel operates on grouped (expert-contiguous) data and
+    iterates over tokens per expert, accumulating gradients in registers.
+  - R (LoRA rank) is a tl.constexpr, allowing tl.arange(0, R). We pad R to a
+    power-of-2 for Triton tile compatibility; typical ranks (4, 8, 16, 32, 64)
+    already satisfy this.
+"""
+
+from itertools import product
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+BLOCK_M = 128
+ALLOW_TF32 = True
+
+
+def _next_power_of_2(n: int) -> int:
+    """Round up to next power of 2."""
+    n -= 1
+    n |= n >> 1
+    n |= n >> 2
+    n |= n >> 4
+    n |= n >> 8
+    n |= n >> 16
+    return n + 1
+
+
+# Granularity for the autotune cache key on M. The kernel still runs on the
+# real M (loop bounds + masks); only the @triton.autotune key is bucketed so
+# that varying seqlens/routing don't keep invalidating the cache.
+_M_BUCKET_GRANULARITY = 1024
+# Large-M / expert-parallel regime: bucket coarsely so per-rank token-count variation (DeepEP
+# dispatch is imbalanced) collapses to ONE autotune key. Otherwise ranks land on different
+# M_BUCKETs, re-autotune at scattered layers, and the resulting per-rank timing skew trips DeepEP's
+# combine barrier. The kernel still runs on the real M; only the @triton.autotune key is bucketed.
+_M_BUCKET_COARSE = 131072
+_M_BUCKET_COARSE_THRESHOLD = 16384
+
+
+def _bucket_m(m: int) -> int:
+    g = _M_BUCKET_COARSE if m >= _M_BUCKET_COARSE_THRESHOLD else _M_BUCKET_GRANULARITY
+    return ((m + g - 1) // g) * g
+
+
+# Triton tl.dot requires minimum tile dimensions of 16 on modern GPUs.
+MIN_TRITON_DOT_SIZE = 16
+
+
+def _block_r_for_rank(r: int) -> int:
+    """Compute BLOCK_R: next power-of-2 >= max(r, MIN_TRITON_DOT_SIZE)."""
+    return _next_power_of_2(max(r, MIN_TRITON_DOT_SIZE))
+
+
+# =============================================================================
+# Token Rounding: pad expert counts to BLOCK_M multiples
+# =============================================================================
+
+
+def round_expert_counts(
+    sorted_expert_idxs: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    E: int,
+    block_m: int = BLOCK_M,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Pad each expert's token count to a multiple of block_m to eliminate
+    partial-tile waste in the backward kernel.
+
+    Padding is done by duplicating the last valid token index for each expert.
+    The kernel's M_mask = M_idx < real_end_idx masks these padding entries, so
+    correctness is preserved (they contribute 0 to the accumulation via other=0.0).
+
+    This only helps the backward dA/dB kernel where per-expert iteration is
+    explicit. The forward scatter2scatter kernel handles partial tiles via masking.
+
+    Args:
+        sorted_expert_idxs: Expert assignments sorted [M*k]
+        sorted_scattered_idxs: Original indices sorted [M*k]
+        expert_offsets: Cumulative token counts per expert [E]
+        E: Number of experts
+        block_m: Block size for token dimension (default: BLOCK_M)
+
+    Returns:
+        padded_expert_idxs: [M_padded] expert assignments with padding
+        padded_scattered_idxs: [M_padded] original indices with padding
+        padded_offsets: [E] cumulative padded counts (for kernel iteration range)
+        real_offsets: [E] original cumulative counts (for M_mask in kernel)
+    """
+    device = sorted_expert_idxs.device
+
+    # Compute per-expert counts
+    counts = torch.zeros(E, dtype=torch.int64, device=device)
+    prev = 0
+    for e in range(E):
+        curr = expert_offsets[e].item()
+        counts[e] = curr - prev
+        prev = curr
+
+    # Round up each count to multiple of block_m
+    padded_counts = ((counts + block_m - 1) // block_m) * block_m
+    # Experts with 0 tokens stay at 0
+    padded_counts = torch.where(
+        counts > 0, padded_counts, torch.zeros_like(padded_counts)
+    )
+    total_padded = padded_counts.sum().item()
+
+    padded_expert_idxs = torch.empty(
+        total_padded, dtype=sorted_expert_idxs.dtype, device=device
+    )
+    padded_scattered_idxs = torch.empty(
+        total_padded, dtype=sorted_scattered_idxs.dtype, device=device
+    )
+
+    src_offset = 0
+    dst_offset = 0
+    for e in range(E):
+        count = counts[e].item()
+        padded_count = padded_counts[e].item()
+
+        if count > 0:
+            # Copy original tokens
+            padded_expert_idxs[dst_offset : dst_offset + count] = sorted_expert_idxs[
+                src_offset : src_offset + count
+            ]
+            padded_scattered_idxs[dst_offset : dst_offset + count] = (
+                sorted_scattered_idxs[src_offset : src_offset + count]
+            )
+
+            # Pad with last valid token (masked out by kernel via M_mask)
+            if padded_count > count:
+                padded_expert_idxs[dst_offset + count : dst_offset + padded_count] = (
+                    sorted_expert_idxs[src_offset + count - 1]
+                )
+                padded_scattered_idxs[
+                    dst_offset + count : dst_offset + padded_count
+                ] = sorted_scattered_idxs[src_offset + count - 1]
+
+        src_offset += count
+        dst_offset += padded_count
+
+    # Padded offsets: cumulative padded counts (for iteration range in kernel)
+    padded_offsets = padded_counts.cumsum(-1).to(expert_offsets.dtype)
+    # Real offsets: original cumulative counts (for M_mask in kernel)
+    real_offsets = expert_offsets.clone()
+
+    return padded_expert_idxs, padded_scattered_idxs, padded_offsets, real_offsets
+
+
+# =============================================================================
+# Autotuning: SMEM estimation and config pruning
+# =============================================================================
+
+_SMEM_CAPACITY: int | None = None
+
+
+def _get_smem_capacity() -> int:
+    """Get device shared memory capacity (bytes). Cached after first call."""
+    global _SMEM_CAPACITY
+    if _SMEM_CAPACITY is None:
+        props = triton.runtime.driver.active.utils.get_device_properties(
+            torch.cuda.current_device()
+        )
+        _SMEM_CAPACITY = props["max_shared_mem"]
+    return _SMEM_CAPACITY
+
+
+def _estimate_smem_usage(
+    num_stages: int, BLOCK_M: int, BLOCK_N: int, BLOCK_K: int, dtype_bytes: int = 2
+) -> int:
+    """Estimate shared memory in bytes for a GEMM-style tile.
+
+    Formula: stages * BLOCK_K * (BLOCK_M + BLOCK_N) + BLOCK_M * BLOCK_N
+    Multiply by dtype_bytes (2 for fp16/bf16).
+    """
+    return (
+        num_stages * BLOCK_K * (BLOCK_M + BLOCK_N) + BLOCK_M * BLOCK_N
+    ) * dtype_bytes
+
+
+# Conservative margin (bytes) subtracted from SMEM capacity to account for
+# estimation inaccuracies and kernel overhead (registers spilled to SMEM, etc.)
+_SMEM_SLACK = 10_000
+
+
+def _estimate_register_pressure(
+    num_warps: int,
+    *tile_sizes: tuple[int, int],
+) -> float:
+    """Rough estimate of per-thread register footprint from live tile sizes.
+
+    This is a heuristic, NOT an accurate register count.  Triton uses tensor
+    core MMA fragments that pack multiple elements per register, and can spill
+    to local memory when the hardware limit (255 regs/thread) is exceeded.
+
+    The estimate is used to prune only truly extreme configs that would cause
+    excessive spilling or compilation failures.  The threshold is set high
+    (``_MAX_REGS_SOFT_LIMIT``) because the heuristic overestimates — it
+    doesn't account for MMA fragment packing.  Configs like M=64,N=64,K=64
+    (est ~520) work fine in practice via spilling.
+
+    Returns estimated registers per thread.
+    """
+    # Each thread in a warp holds ~1/32 of the tile elements
+    tile_regs = sum(r * c for r, c in tile_sizes) / 32
+    scalar_overhead = 40
+    return tile_regs + scalar_overhead
+
+
+# Soft limit for register pressure pruning.  Only prune configs with extreme
+# tile products (e.g. M=128,K=256,N=256) that reliably crash on Blackwell.
+# Moderate configs (M=64,N=64,K=64, est ~520) work via register spilling.
+_MAX_REGS_SOFT_LIMIT = 1024
+
+
+# =============================================================================
+# Forward Kernel: scatter2scatter with fused LoRA
+# =============================================================================
+
+
+@triton.jit
+def _compute_expert_block_lora(
+    E_idx,
+    E_mask,
+    M_in_idx,
+    N_block,
+    N_mask,
+    # Base weight
+    X_ptr,
+    stride_xm,
+    stride_xk,
+    W_ptr,
+    stride_we,
+    stride_wk,
+    stride_wn,
+    # LoRA weights
+    A_ptr,
+    stride_ar,
+    stride_ak,  # A: [r*E, K], stride_ar = stride for r*E dim, stride_ak = stride for K dim
+    B_ptr,
+    stride_bn,
+    stride_br,  # B: [N, r*E], stride_bn = stride for N dim, stride_br = stride for r*E dim
+    # Dimensions
+    K,
+    ACTUAL_R: tl.constexpr,  # True LoRA rank (for indexing into weight arrays)
+    acc,
+    no_k_mask,
+    BLOCK_M: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_R: tl.constexpr,  # Padded tile size >= max(ACTUAL_R, 16)
+    scaling,
+    allow_tf32: tl.constexpr,
+):
+    """
+    Compute Y_block = X_block @ W_e + scaling * (X_block @ A_e^T) @ B_e^T
+
+    for tokens in this M-block assigned to expert E_idx.
+
+    ACTUAL_R is the true LoRA rank used for indexing into A[e*r:(e+1)*r, :].
+    BLOCK_R >= ACTUAL_R is the padded tile dimension (must be >= 16 for tl.dot).
+    When BLOCK_R > ACTUAL_R, loads are masked on the R dimension.
+    """
+    K_block = tl.arange(0, BLOCK_K)
+    R_block = tl.arange(0, BLOCK_R)
+    R_mask = R_block < ACTUAL_R  # Mask for padding when BLOCK_R > ACTUAL_R
+
+    # Base weight pointers: W[E_idx, :, :] is [K, N], load [BLOCK_K, BLOCK_N]
+    X_blk_ptrs = X_ptr + M_in_idx[:, None] * stride_xm + K_block[None, :] * stride_xk
+    W_blk_ptrs = (
+        W_ptr
+        + E_idx * stride_we
+        + K_block[:, None] * stride_wk
+        + N_block[None, :] * stride_wn
+    )
+
+    # LoRA A pointers: A[e*ACTUAL_R:(e+1)*ACTUAL_R, :] for expert e, shape [r, K]
+    A_expert_offset = E_idx * ACTUAL_R
+    A_blk_ptrs = (
+        A_ptr
+        + (A_expert_offset + R_block)[:, None] * stride_ar
+        + K_block[None, :] * stride_ak
+    )
+
+    iters = tl.cdiv(K, BLOCK_K)
+
+    # Accumulator for X @ A^T: [BLOCK_M, BLOCK_R]
+    xa_acc = tl.zeros((BLOCK_M, BLOCK_R), dtype=tl.float32)
+
+    # Determine the input element type for consistent casting.
+    # Masked tl.load with other=0.0 can upcast bf16->fp32 in some Triton versions,
+    # causing dtype mismatches in tl.dot.  We cast all tiles to the same type.
+    INPUT_DTYPE = X_ptr.dtype.element_ty
+
+    for i in range(iters):
+        if no_k_mask:
+            x = tl.load(X_blk_ptrs, mask=E_mask[:, None], other=0.0).to(INPUT_DTYPE)
+            w = tl.load(W_blk_ptrs, mask=N_mask[None, :], other=0.0).to(INPUT_DTYPE)
+            a = tl.load(A_blk_ptrs, mask=R_mask[:, None], other=0.0).to(INPUT_DTYPE)
+        else:
+            K_mask = (i * BLOCK_K + K_block) < K
+            x = tl.load(
+                X_blk_ptrs, mask=E_mask[:, None] & K_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+            w = tl.load(
+                W_blk_ptrs, mask=K_mask[:, None] & N_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+            a = tl.load(
+                A_blk_ptrs, mask=R_mask[:, None] & K_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+
+        # Base: acc += X @ W  ([M, K] @ [K, N] -> [M, N])
+        acc += tl.dot(x, w, allow_tf32=allow_tf32).to(tl.float32)
+
+        # LoRA: xa_acc += X @ A^T  ([M, K] @ [K, R] -> [M, R])
+        xa_acc += tl.dot(x, tl.trans(a), allow_tf32=allow_tf32).to(tl.float32)
+
+        X_blk_ptrs += BLOCK_K * stride_xk
+        W_blk_ptrs += BLOCK_K * stride_wk
+        A_blk_ptrs += BLOCK_K * stride_ak
+
+    # Epilogue: load B[e] and compute (X @ A^T) @ B^T
+    # B[e] is B[:, e*ACTUAL_R:(e+1)*ACTUAL_R], shape [N, r]. Load [BLOCK_N, BLOCK_R].
+    B_expert_offset = E_idx * ACTUAL_R
+    B_blk_ptrs = (
+        B_ptr
+        + N_block[:, None] * stride_bn
+        + (B_expert_offset + R_block)[None, :] * stride_br
+    )
+    b = tl.load(
+        B_blk_ptrs, mask=N_mask[:, None] & R_mask[None, :], other=0.0
+    )  # [BLOCK_N, BLOCK_R]
+
+    # tl.dot requires non-float32 inputs (tensor cores); cast back to input dtype
+    b_inp = b.to(INPUT_DTYPE)
+
+    # (X @ A^T) @ B^T: [M, R] @ [R, N] -> [M, N]
+    lora_out = tl.dot(xa_acc.to(INPUT_DTYPE), tl.trans(b_inp), allow_tf32=allow_tf32)
+
+    acc += scaling * lora_out
+    return acc
+
+
+def _scatter2scatter_lora_configs():
+    """Generate forward kernel autotune configs.
+
+    Search space includes BLOCK_M to allow trading token-tile size for
+    larger BLOCK_K/BLOCK_N tiles.  On GPUs with ~99KB SMEM, BLOCK_M=128
+    forces BLOCK_K=32 and BLOCK_N=32; BLOCK_M=64 allows BLOCK_K=128
+    (4× fewer inner-loop iterations).
+
+    Search space:
+      BLOCK_M:    {32, 64, 128}
+      BLOCK_N:    {32, 64}
+      BLOCK_K:    {32, 64, 128}
+      num_warps:  {4, 8}
+      num_stages: {3, 4, 5}
+    """
+    configs = []
+    for block_m, block_n, block_k, warps, stages in product(
+        [32, 64, 128],  # BLOCK_M
+        [32, 64],  # BLOCK_N
+        [32, 64, 128],  # BLOCK_K
+        [4, 8],  # num_warps
+        [3, 4, 5],  # num_stages
+    ):
+        configs.append(
+            triton.Config(
+                {"BLOCK_M": block_m, "BLOCK_N": block_n, "BLOCK_K": block_k},
+                num_stages=stages,
+                num_warps=warps,
+            )
+        )
+    return configs
+
+
+def _prune_fwd_configs(configs, named_args, **kwargs):
+    """Prune forward configs based on SMEM capacity and register pressure.
+
+    The forward kernel inner loop loads three tiles per pipeline stage:
+      X[BLOCK_M, BLOCK_K], W[BLOCK_K, BLOCK_N], A[BLOCK_R, BLOCK_K].
+    The base estimate only accounts for X and W. We add:
+      - A tile [BLOCK_R, BLOCK_K] per pipeline stage (loaded in the inner loop)
+      - B tile [BLOCK_N, BLOCK_R] loaded once in the epilogue
+      - Extra headroom for compiler overhead (register spills, metadata)
+    """
+    smem_cap = _get_smem_capacity()
+
+    # Get BLOCK_R from named_args if available, else assume worst case
+    block_r = named_args.get("BLOCK_R", 64)
+
+    scored = []
+    for config in configs:
+        block_m = config.kwargs["BLOCK_M"]
+        block_n = config.kwargs["BLOCK_N"]
+        block_k = config.kwargs["BLOCK_K"]
+        # Base: stages * BLOCK_K * (BLOCK_M + BLOCK_N) + BLOCK_M * BLOCK_N
+        smem_base = _estimate_smem_usage(config.num_stages, block_m, block_n, block_k)
+        # A tile [BLOCK_R, BLOCK_K] loaded per stage in the inner loop
+        smem_lora_loop = config.num_stages * block_r * block_k * 2
+        # B tile [BLOCK_N, BLOCK_R] loaded once in epilogue
+        smem_lora_epilogue = block_n * block_r * 2
+        smem = smem_base + smem_lora_loop + smem_lora_epilogue
+
+        # Register pressure: live tiles are acc[M,N], xa_acc[M,R],
+        # x[M,K], w[K,N], a[R,K], plus epilogue b[N,R]
+        est_regs = _estimate_register_pressure(
+            config.num_warps,
+            (block_m, block_n),  # acc
+            (block_m, block_r),  # xa_acc
+            (block_m, block_k),  # x tile
+            (block_k, block_n),  # w tile
+            (block_r, block_k),  # a tile
+            (block_n, block_r),  # b tile (epilogue)
+        )
+        if est_regs > _MAX_REGS_SOFT_LIMIT:
+            continue
+
+        scored.append((smem, config))
+
+    pruned = [c for s, c in scored if s <= smem_cap - _SMEM_SLACK]
+    if pruned:
+        return pruned
+    if scored:
+        # All surviving configs exceed SMEM — return the one with smallest usage
+        scored.sort(key=lambda x: x[0])
+        return [scored[0][1]]
+    # All configs pruned by register pressure — fall back to smallest tiles
+    return [
+        min(
+            configs,
+            key=lambda c: (
+                c.kwargs["BLOCK_M"] * c.kwargs["BLOCK_N"] * c.kwargs["BLOCK_K"]
+            ),
+        )
+    ]
+
+
+@triton.autotune(
+    configs=_scatter2scatter_lora_configs(),
+    key=["M_BUCKET", "N", "K"],
+    prune_configs_by={"early_config_prune": _prune_fwd_configs},
+)
+@triton.heuristics(
+    {
+        "NO_K_MASK": lambda args: (args["K"] % args["BLOCK_K"]) == 0,
+        "NO_N_MASK": lambda args: (args["N"] % args["BLOCK_N"]) == 0,
+    }
+)
+@triton.jit
+def _scatter2scatter_lora(
+    # Input/Output
+    X_ptr,
+    stride_xm: tl.constexpr,
+    stride_xk: tl.constexpr,
+    W_ptr,
+    stride_we,
+    stride_wk: tl.constexpr,
+    stride_wn: tl.constexpr,
+    Y_ptr,
+    stride_ym: tl.constexpr,
+    stride_yn: tl.constexpr,
+    # Bias
+    Bias_ptr,
+    stride_bias_e: tl.constexpr,
+    stride_bias_n: tl.constexpr,
+    # LoRA weights
+    LA_ptr,
+    stride_la_r,
+    stride_la_k,  # A: [r*E, K]
+    LB_ptr,
+    stride_lb_n,
+    stride_lb_r,  # B: [N, r*E]
+    # Routing
+    grouped_idx_ptr,
+    expert_idxs_ptr,
+    # Dimensions
+    FAN_OUT: tl.constexpr,
+    M,
+    M_BUCKET,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    E: tl.constexpr,
+    ACTUAL_R: tl.constexpr,  # True LoRA rank (for weight indexing)
+    # Block sizes
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_R: tl.constexpr,  # Padded tile size >= max(ACTUAL_R, 16)
+    # Config
+    ACC_TYPE: tl.constexpr,
+    scaling,
+    allow_tf32: tl.constexpr,
+    x_grouped: tl.constexpr,
+    y_grouped: tl.constexpr,
+    NO_K_MASK: tl.constexpr,
+    NO_N_MASK: tl.constexpr,
+    INT64_INDICES: tl.constexpr = False,
+):
+    """
+    Fused scatter2scatter with LoRA: Y = X @ W + scaling * (X @ A^T) @ B^T + bias
+    """
+    pid = tl.program_id(axis=0)
+
+    N_BLOCK_COUNT = tl.cdiv(N, BLOCK_N)
+    M_block_id = pid // N_BLOCK_COUNT
+    N_block_id = pid % N_BLOCK_COUNT
+
+    M_block = M_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    if INT64_INDICES:
+        M_block = M_block.to(tl.int64)
+    N_block = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+    N_mask = N_block < N
+    M_boundary_mask = M_block < (FAN_OUT * M)
+
+    E_idxs = tl.load(expert_idxs_ptr + M_block, mask=M_boundary_mask, other=E)
+
+    no_k_mask = NO_K_MASK
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
+
+    E_first_idx = tl.min(E_idxs)
+    E_last_idx = tl.minimum(tl.max(E_idxs), E - 1)
+    if INT64_INDICES:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int64)
+    else:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int32)
+
+    for E_idx in range(E_first_idx, E_last_idx + 1):
+        E_mask = E_idxs == E_idx
+        if x_grouped:
+            M_in_idx = M_block
+        else:
+            M_in_idx = M_idx // FAN_OUT
+
+        acc = _compute_expert_block_lora(
+            E_idx,
+            E_mask,
+            M_in_idx,
+            N_block,
+            N_mask,
+            X_ptr,
+            stride_xm,
+            stride_xk,
+            W_ptr,
+            stride_we,
+            stride_wk,
+            stride_wn,
+            LA_ptr,
+            stride_la_r,
+            stride_la_k,
+            LB_ptr,
+            stride_lb_n,
+            stride_lb_r,
+            K,
+            ACTUAL_R,
+            acc,
+            no_k_mask,
+            BLOCK_M,
+            BLOCK_K,
+            BLOCK_N,
+            BLOCK_R,
+            scaling,
+            allow_tf32=allow_tf32,
+        )
+
+    # Add bias if present
+    if Bias_ptr is not None:
+        B_blk_ptrs = (
+            Bias_ptr
+            + E_idxs[:, None] * stride_bias_e
+            + N_block[None, :] * stride_bias_n
+        )
+        acc += tl.load(B_blk_ptrs, mask=M_boundary_mask[:, None] & N_mask[None, :])
+
+    # Store output
+    if y_grouped:
+        M_out_idx = M_block
+    else:
+        M_out_idx = M_idx
+    Y_blk_ptrs = Y_ptr + (M_out_idx[:, None] * stride_ym + N_block[None, :] * stride_yn)
+    tl.store(Y_blk_ptrs, acc, mask=M_boundary_mask[:, None] & N_mask[None, :])
+
+
+def _scatter2scatter_lora_split(
+    X: torch.Tensor,
+    W: torch.Tensor,
+    sorted_expert_idxs: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    k: int,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    scaling: float,
+    b: Optional[torch.Tensor] = None,
+    x_grouped: bool = False,
+    y_grouped: bool = False,
+    out: Optional[torch.Tensor] = None,
+    int64_indices: bool = False,
+) -> torch.Tensor:
+    """Split base+LoRA forward: 3 scatter2scatter calls, no fused LoRA kernel.
+
+    Faster for models with few large experts (e.g. Mixtral E=8, I=14336)
+    because the base kernel runs at full speed without LoRA SMEM overhead,
+    and the LoRA matmuls (R=16) are tiny separate passes.
+
+    Y = scatter(X, W) + scaling * scatter(scatter(X, A^T), B^T)
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops import (
+        scatter2scatter,
+    )
+
+    E = W.size(0)
+    R = lora_A.size(0) // E
+    K = W.size(1)
+    N = W.size(2)
+
+    # 1. Base: Y_base = X @ W  (uses base kernel with optimal tile sizes)
+    output = scatter2scatter(
+        X=X,
+        W=W,
+        b=b,
+        sorted_expert_idxs=sorted_expert_idxs,
+        sorted_scattered_idxs=sorted_scattered_idxs,
+        k=k,
+        x_grouped=x_grouped,
+        y_grouped=y_grouped,
+        out=out,
+        int64_indices=int64_indices,
+    )
+
+    # 2. XA = X @ A^T  (tiny: output is [M*k, R])
+    # Reshape A: [R*E, K] → [E, K, R] (expert weights for scatter2scatter)
+    W_A = lora_A.reshape(E, R, K).permute(0, 2, 1).contiguous()
+    XA = scatter2scatter(
+        X=X,
+        W=W_A,
+        sorted_expert_idxs=sorted_expert_idxs,
+        sorted_scattered_idxs=sorted_scattered_idxs,
+        k=k,
+        x_grouped=x_grouped,
+        y_grouped=True,
+        int64_indices=int64_indices,
+    )
+
+    # 3. Y_lora = XA @ B^T  (R is tiny, so this is very fast)
+    # Reshape B: [N, R*E] → [E, R, N]
+    W_B = lora_B.T.reshape(E, R, N).contiguous()
+    Y_lora = scatter2scatter(
+        X=XA,
+        W=W_B,
+        sorted_expert_idxs=sorted_expert_idxs,
+        sorted_scattered_idxs=sorted_scattered_idxs,
+        k=1,
+        x_grouped=True,
+        y_grouped=y_grouped,
+        int64_indices=int64_indices,
+    )
+
+    # 4. Y = Y_base + scaling * Y_lora
+    output.add_(Y_lora, alpha=scaling)
+    return output
+
+
+# Threshold for switching from fused to split LoRA forward.
+# Split wins when per-expert matmul is large (bandwidth-bound LoRA tile
+# loads dominate in the fused kernel's inner loop).
+# Empirically: split wins for E<=32 with K*N > 20M (e.g. Mixtral, Phi-MoE).
+_SPLIT_LORA_FWD_THRESHOLD = 20_000_000  # per-expert K*N
+_SPLIT_LORA_FWD_MAX_EXPERTS = 32
+
+
+def scatter2scatter_lora(
+    X: torch.Tensor,
+    W: torch.Tensor,
+    sorted_expert_idxs: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    k: int,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    scaling: float,
+    b: Optional[torch.Tensor] = None,
+    x_grouped: bool = False,
+    y_grouped: bool = False,
+    out: Optional[torch.Tensor] = None,
+    int64_indices: bool = False,
+) -> torch.Tensor:
+    """
+    Scatter2scatter with LoRA: Y[i] = X[i] @ W[e] + scaling * (X[i] @ A[e]^T) @ B[e]^T + b[e]
+
+    Automatically selects between:
+    - Fused kernel: single Triton kernel with LoRA in the inner loop.
+      Best for many small experts (E>=64, small K*N).
+    - Split dispatch: 3 separate scatter2scatter calls (base + XA + lora).
+      Best for few large experts (E<=32, large K*N like Mixtral).
+
+    Args:
+        X: Input [M, K] or [M*k, K] if x_grouped
+        W: Expert weights [E, K, N]
+        sorted_expert_idxs: Expert assignments sorted [M*k]
+        sorted_scattered_idxs: Original indices sorted [M*k]
+        k: Fan-out (top-k)
+        lora_A: LoRA A weights [r*E, K]
+        lora_B: LoRA B weights [N, r*E]
+        scaling: LoRA scaling factor (alpha/r)
+        b: Optional bias [E, N]
+        x_grouped: Input pre-grouped by expert
+        y_grouped: Keep output grouped
+        out: Optional pre-allocated output buffer
+
+    Returns:
+        Y: Output [M*k, N]
+    """
+    E = W.size(0)
+    K = W.size(1)
+    N = W.size(2)
+
+    # Dispatch: split for few large experts, fused for many small experts
+    if E <= _SPLIT_LORA_FWD_MAX_EXPERTS and K * N >= _SPLIT_LORA_FWD_THRESHOLD:
+        return _scatter2scatter_lora_split(
+            X,
+            W,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            k,
+            lora_A,
+            lora_B,
+            scaling,
+            b,
+            x_grouped,
+            y_grouped,
+            out,
+            int64_indices=int64_indices,
+        )
+
+    assert sorted_scattered_idxs.size(0) == sorted_expert_idxs.size(0)
+    assert sorted_scattered_idxs.size(0) == X.size(0) * k
+
+    R = lora_A.size(0) // E
+
+    # Pad R to power of 2 for Triton tile size
+    BLOCK_R = _block_r_for_rank(R)
+
+    L_scattered = sorted_expert_idxs.size(0)
+
+    if out is None:
+        output = torch.empty((L_scattered, N), device=X.device, dtype=X.dtype)
+    else:
+        assert out.size(0) == L_scattered and out.size(1) == N
+        output = out
+
+    def grid(META):
+        return (
+            triton.cdiv(L_scattered, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),
+        )
+
+    if b is None:
+        stride_be = stride_bn = 0
+        b_ptr = None
+    else:
+        stride_be, stride_bn = b.stride()
+        b_ptr = b
+
+    _scatter2scatter_lora[grid](
+        X,
+        X.stride(0),
+        X.stride(1),
+        W,
+        W.stride(0),
+        W.stride(1),
+        W.stride(2),
+        output,
+        output.stride(0),
+        output.stride(1),
+        b_ptr,
+        stride_be,
+        stride_bn,
+        lora_A,
+        lora_A.stride(0),
+        lora_A.stride(1),
+        lora_B,
+        lora_B.stride(0),
+        lora_B.stride(1),
+        sorted_scattered_idxs,
+        sorted_expert_idxs,
+        FAN_OUT=k,
+        M=X.size(0),
+        M_BUCKET=_bucket_m(X.size(0)),
+        K=K,
+        N=N,
+        E=E,
+        ACTUAL_R=R,
+        BLOCK_R=BLOCK_R,
+        ACC_TYPE=tl.float32,
+        scaling=scaling,
+        allow_tf32=ALLOW_TF32,
+        x_grouped=x_grouped,
+        y_grouped=y_grouped,
+        INT64_INDICES=int64_indices,
+    )
+
+    return output
+
+
+# =============================================================================
+# Backward Kernel: Fused dX = dY @ W^T + scaling * (dY @ B) @ A
+# =============================================================================
+
+
+@triton.jit
+def _compute_expert_block_lora_dX(
+    E_idx,
+    E_mask,
+    M_in_idx,
+    K_block,
+    K_mask,
+    # Input: DY (gradient w.r.t. output)
+    DY_ptr,
+    stride_dym,
+    stride_dyn,
+    # Base weight W^T: we load W[e] as [K, N] and index as W^T[e] = [N, K]
+    W_ptr,
+    stride_we,
+    stride_wk,
+    stride_wn,
+    # LoRA weights
+    A_ptr,
+    stride_ar,
+    stride_ak,  # A: [r*E, K]
+    B_ptr,
+    stride_bn,
+    stride_br,  # B: [N, r*E]
+    # Dimensions
+    N,
+    ACTUAL_R: tl.constexpr,
+    acc,
+    no_n_mask,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    scaling,
+    allow_tf32: tl.constexpr,
+):
+    """
+    Compute dX_block = DY_block @ W_e^T + scaling * (DY_block @ B_e) @ A_e
+
+    for tokens in this M-block assigned to expert E_idx.
+
+    Inner loop over N dimension (reduction dim for dY @ W^T and dY @ B).
+    Output dimension is K.
+    Epilogue computes (dY @ B) @ A.
+
+    Transpose mapping from forward:
+      Forward: X@W (K-loop), X@A^T (K-loop), (X@A^T)@B^T (epilogue)
+      Backward: DY@W^T (N-loop), DY@B (N-loop), (DY@B)@A (epilogue)
+    """
+    N_block = tl.arange(0, BLOCK_N)
+    R_block = tl.arange(0, BLOCK_R)
+    R_mask = R_block < ACTUAL_R
+
+    # DY pointers: DY is [M_total, N], load [BLOCK_M, BLOCK_N]
+    DY_blk_ptrs = (
+        DY_ptr + M_in_idx[:, None] * stride_dym + N_block[None, :] * stride_dyn
+    )
+
+    # W^T pointers: W[e] is [K, N], W^T[e] is [N, K]. We load W^T as [BLOCK_N, BLOCK_K].
+    # W stored as [E, K, N], so W^T[e][n, k] = W[e][k, n] = W_ptr + e*stride_we + k*stride_wk + n*stride_wn
+    # As [BLOCK_N, BLOCK_K] tile: row=n, col=k
+    WT_blk_ptrs = (
+        W_ptr
+        + E_idx * stride_we
+        + N_block[:, None] * stride_wn  # row = n dimension
+        + K_block[None, :] * stride_wk
+    )  # col = k dimension
+
+    # B pointers: B[e] is B[:, e*R:(e+1)*R], shape [N, R]. Load [BLOCK_N, BLOCK_R].
+    B_expert_offset = E_idx * ACTUAL_R
+    B_blk_ptrs = (
+        B_ptr
+        + N_block[:, None] * stride_bn
+        + (B_expert_offset + R_block)[None, :] * stride_br
+    )
+
+    iters = tl.cdiv(N, BLOCK_N)
+
+    # Accumulator for DY @ B: [BLOCK_M, BLOCK_R]
+    dy_b_acc = tl.zeros((BLOCK_M, BLOCK_R), dtype=tl.float32)
+
+    # Determine the input element type for consistent casting.
+    INPUT_DTYPE = DY_ptr.dtype.element_ty
+
+    for i in range(iters):
+        if no_n_mask:
+            dy = tl.load(DY_blk_ptrs, mask=E_mask[:, None], other=0.0).to(INPUT_DTYPE)
+            wt = tl.load(WT_blk_ptrs, mask=K_mask[None, :], other=0.0).to(INPUT_DTYPE)
+            b = tl.load(B_blk_ptrs, mask=R_mask[None, :], other=0.0).to(INPUT_DTYPE)
+        else:
+            N_mask_iter = (i * BLOCK_N + N_block) < N
+            dy = tl.load(
+                DY_blk_ptrs, mask=E_mask[:, None] & N_mask_iter[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+            wt = tl.load(
+                WT_blk_ptrs, mask=N_mask_iter[:, None] & K_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+            b = tl.load(
+                B_blk_ptrs, mask=N_mask_iter[:, None] & R_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+
+        # Base: acc += DY @ W^T  ([M, N] @ [N, K] -> [M, K])
+        acc += tl.dot(dy, wt, allow_tf32=allow_tf32).to(tl.float32)
+
+        # LoRA: dy_b_acc += DY @ B  ([M, N] @ [N, R] -> [M, R])
+        dy_b_acc += tl.dot(dy, b, allow_tf32=allow_tf32).to(tl.float32)
+
+        DY_blk_ptrs += BLOCK_N * stride_dyn
+        WT_blk_ptrs += BLOCK_N * stride_wn
+        B_blk_ptrs += BLOCK_N * stride_bn
+
+    # Epilogue: load A[e] and compute (DY @ B) @ A
+    # A[e] is A[e*R:(e+1)*R, :], shape [R, K]. Load [BLOCK_R, BLOCK_K].
+    A_expert_offset = E_idx * ACTUAL_R
+    A_blk_ptrs = (
+        A_ptr
+        + (A_expert_offset + R_block)[:, None] * stride_ar
+        + K_block[None, :] * stride_ak
+    )
+    a_e = tl.load(A_blk_ptrs, mask=R_mask[:, None] & K_mask[None, :], other=0.0).to(
+        INPUT_DTYPE
+    )
+
+    # (DY @ B) @ A: [M, R] @ [R, K] -> [M, K]
+    # tl.dot requires non-float32 inputs (tensor cores); cast accumulator back to input dtype
+    lora_dx = tl.dot(dy_b_acc.to(INPUT_DTYPE), a_e, allow_tf32=allow_tf32)
+
+    acc += scaling * lora_dx
+    return acc
+
+
+def _scatter2scatter_lora_dX_configs():
+    """Generate backward dX kernel autotune configs.
+
+    The inner loop is over N (not K as in forward). The output dimension is K.
+    So BLOCK_K tiles the output and BLOCK_N tiles the reduction.
+
+    BLOCK_M is now autotunable (was fixed at 128).
+
+    Search space:
+      BLOCK_M:    {32, 64, 128}        (token tile)
+      BLOCK_K:    {32, 64, 128}   (output tile)
+      BLOCK_N:    {32, 64}   (reduction tile)
+      num_warps:  {4, 8}
+      num_stages: {3, 4, 5}
+    """
+    configs = []
+    for block_m, block_k, block_n, warps, stages in product(
+        [32, 64, 128],  # BLOCK_M
+        [32, 64, 128],  # BLOCK_K (output dimension)
+        [32, 64],  # BLOCK_N (reduction dimension)
+        [4, 8],  # num_warps
+        [3, 4, 5],  # num_stages
+    ):
+        configs.append(
+            triton.Config(
+                {"BLOCK_M": block_m, "BLOCK_K": block_k, "BLOCK_N": block_n},
+                num_stages=stages,
+                num_warps=warps,
+            )
+        )
+    return configs
+
+
+def _prune_dX_configs(configs, named_args, **kwargs):
+    """Prune backward dX configs based on SMEM capacity and register pressure.
+
+    The dX kernel inner loop loads three tiles per pipeline stage:
+      DY[BLOCK_M, BLOCK_N], W^T[BLOCK_N, BLOCK_K], B[BLOCK_N, BLOCK_R].
+    The base estimate only accounts for DY and W^T. We add:
+      - B tile [BLOCK_N, BLOCK_R] per pipeline stage (loaded in the inner loop)
+      - A tile [BLOCK_R, BLOCK_K] loaded once in the epilogue
+      - Extra headroom for compiler overhead (register spills, metadata)
+    """
+    smem_cap = _get_smem_capacity()
+
+    # Get BLOCK_R from named_args if available, else assume worst case
+    block_r = named_args.get("BLOCK_R", 64)
+
+    scored = []
+    for config in configs:
+        block_m = config.kwargs["BLOCK_M"]
+        block_k = config.kwargs["BLOCK_K"]
+        block_n = config.kwargs["BLOCK_N"]
+        # Base: stages * BLOCK_N * (BLOCK_M + BLOCK_K) + BLOCK_M * BLOCK_K
+        smem_base = _estimate_smem_usage(config.num_stages, block_m, block_k, block_n)
+        # B tile [BLOCK_N, BLOCK_R] loaded per stage in the inner loop
+        smem_lora_loop = config.num_stages * block_n * block_r * 2
+        # A tile [BLOCK_R, BLOCK_K] loaded once in epilogue
+        smem_lora_epilogue = block_r * block_k * 2
+        smem = smem_base + smem_lora_loop + smem_lora_epilogue
+
+        # Register pressure: live tiles are acc[M,K], dy_b_acc[M,R],
+        # dy[M,N], wt[N,K], b[N,R], plus epilogue a[R,K]
+        est_regs = _estimate_register_pressure(
+            config.num_warps,
+            (block_m, block_k),  # acc
+            (block_m, block_r),  # dy_b_acc
+            (block_m, block_n),  # dy tile
+            (block_n, block_k),  # wt tile
+            (block_n, block_r),  # b tile
+            (block_r, block_k),  # a tile (epilogue)
+        )
+        if est_regs > _MAX_REGS_SOFT_LIMIT:
+            continue
+
+        scored.append((smem, config))
+
+    pruned = [c for s, c in scored if s <= smem_cap - _SMEM_SLACK]
+    if pruned:
+        return pruned
+    if scored:
+        # All surviving configs exceed SMEM — return the one with smallest usage
+        scored.sort(key=lambda x: x[0])
+        return [scored[0][1]]
+    # All configs pruned by register pressure — fall back to smallest tiles
+    return [
+        min(
+            configs,
+            key=lambda c: (
+                c.kwargs["BLOCK_M"] * c.kwargs["BLOCK_K"] * c.kwargs["BLOCK_N"]
+            ),
+        )
+    ]
+
+
+@triton.autotune(
+    configs=_scatter2scatter_lora_dX_configs(),
+    key=["M_BUCKET", "N", "K"],
+    prune_configs_by={"early_config_prune": _prune_dX_configs},
+)
+@triton.heuristics(
+    {
+        "NO_K_MASK": lambda args: (args["K"] % args["BLOCK_K"]) == 0,
+        "NO_N_MASK": lambda args: (args["N"] % args["BLOCK_N"]) == 0,
+    }
+)
+@triton.jit
+def _scatter2scatter_lora_dX(
+    # Input: DY (gradient w.r.t. output, grouped)
+    DY_ptr,
+    stride_dym: tl.constexpr,
+    stride_dyn: tl.constexpr,
+    # Base weight: W [E, K, N] (we compute DY @ W^T)
+    W_ptr,
+    stride_we,
+    stride_wk: tl.constexpr,
+    stride_wn: tl.constexpr,
+    # Output: dX
+    DX_ptr,
+    stride_dxm: tl.constexpr,
+    stride_dxk: tl.constexpr,
+    # LoRA weights
+    LA_ptr,
+    stride_la_r,
+    stride_la_k,  # A: [r*E, K]
+    LB_ptr,
+    stride_lb_n,
+    stride_lb_r,  # B: [N, r*E]
+    # Routing
+    grouped_idx_ptr,
+    expert_idxs_ptr,
+    # Dimensions
+    FAN_OUT: tl.constexpr,
+    M,
+    M_BUCKET,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    E: tl.constexpr,
+    ACTUAL_R: tl.constexpr,
+    # Block sizes
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    # Config
+    ACC_TYPE: tl.constexpr,
+    scaling,
+    allow_tf32: tl.constexpr,
+    dy_grouped: tl.constexpr,
+    dx_grouped: tl.constexpr,
+    NO_K_MASK: tl.constexpr,
+    NO_N_MASK: tl.constexpr,
+    INT64_INDICES: tl.constexpr = False,
+):
+    """
+    Fused backward dX = DY @ W^T + scaling * (DY @ B) @ A
+
+    DY is in expert-grouped order (x_grouped=True).
+    dX is output in ungrouped or grouped order based on dx_grouped.
+
+    Grid: (cdiv(M_total, BLOCK_M) * cdiv(K, BLOCK_K),)
+    """
+    pid = tl.program_id(axis=0)
+
+    K_BLOCK_COUNT = tl.cdiv(K, BLOCK_K)
+    M_block_id = pid // K_BLOCK_COUNT
+    K_block_id = pid % K_BLOCK_COUNT
+
+    M_block = M_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    if INT64_INDICES:
+        M_block = M_block.to(tl.int64)
+    K_block = K_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
+    K_mask = K_block < K
+    M_boundary_mask = M_block < (FAN_OUT * M)
+
+    E_idxs = tl.load(expert_idxs_ptr + M_block, mask=M_boundary_mask, other=E)
+
+    no_n_mask = NO_N_MASK
+
+    acc = tl.zeros((BLOCK_M, BLOCK_K), dtype=ACC_TYPE)
+
+    E_first_idx = tl.min(E_idxs)
+    E_last_idx = tl.minimum(tl.max(E_idxs), E - 1)
+    if INT64_INDICES:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int64)
+    else:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int32)
+
+    for E_idx in range(E_first_idx, E_last_idx + 1):
+        E_mask = E_idxs == E_idx
+        if dy_grouped:
+            M_in_idx = M_block
+        else:
+            M_in_idx = M_idx // FAN_OUT
+
+        acc = _compute_expert_block_lora_dX(
+            E_idx,
+            E_mask,
+            M_in_idx,
+            K_block,
+            K_mask,
+            DY_ptr,
+            stride_dym,
+            stride_dyn,
+            W_ptr,
+            stride_we,
+            stride_wk,
+            stride_wn,
+            LA_ptr,
+            stride_la_r,
+            stride_la_k,
+            LB_ptr,
+            stride_lb_n,
+            stride_lb_r,
+            N,
+            ACTUAL_R,
+            acc,
+            no_n_mask,
+            BLOCK_M,
+            BLOCK_N,
+            BLOCK_K,
+            BLOCK_R,
+            scaling,
+            allow_tf32=allow_tf32,
+        )
+
+    # Store output
+    if dx_grouped:
+        M_out_idx = M_block
+    else:
+        M_out_idx = M_idx
+    DX_blk_ptrs = DX_ptr + (
+        M_out_idx[:, None] * stride_dxm + K_block[None, :] * stride_dxk
+    )
+    tl.store(DX_blk_ptrs, acc, mask=M_boundary_mask[:, None] & K_mask[None, :])
+
+
+def scatter2scatter_lora_dX(
+    DY: torch.Tensor,
+    W: torch.Tensor,
+    sorted_expert_idxs: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    k: int,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    scaling: float,
+    dy_grouped: bool = True,
+    dx_grouped: bool = False,
+    out: Optional[torch.Tensor] = None,
+    int64_indices: bool = False,
+) -> torch.Tensor:
+    """
+    Fused backward dX = DY @ W^T + scaling * (DY @ B) @ A
+
+    Replaces the separate:
+      1. base_ops.scatter2scatter(DY, W^T, x_grouped=True, ...)
+      2. _compute_lora_input_grad(DY, A, B, ...)
+
+    Args:
+        DY: Gradient w.r.t. output [M*k, N] (grouped by expert)
+        W: Expert weights [E, K, N] (NOT transposed — kernel handles W^T internally)
+        sorted_expert_idxs: Expert assignments sorted [M*k]
+        sorted_scattered_idxs: Original indices sorted [M*k]
+        k: Fan-out (top-k)
+        lora_A: LoRA A weights [r*E, K]
+        lora_B: LoRA B weights [N, r*E]
+        scaling: LoRA scaling factor
+        dy_grouped: Whether DY is in grouped (expert-sorted) order (default True)
+        dx_grouped: Whether to output dX in grouped order (default False)
+        out: Optional pre-allocated output buffer
+
+    Returns:
+        dX: Input gradient [M*k, K]
+    """
+    assert sorted_scattered_idxs.size(0) == sorted_expert_idxs.size(0)
+
+    E = W.size(0)
+    K = W.size(1)
+    N = W.size(2)
+    R = lora_A.size(0) // E
+
+    BLOCK_R = _block_r_for_rank(R)
+
+    L_scattered = sorted_expert_idxs.size(0)
+
+    # M for the kernel is DY.size(0) when dy_grouped, else the original M
+    if dy_grouped:
+        M = DY.size(0)
+        fan_out = 1  # DY is already expanded
+    else:
+        M = DY.size(0)
+        fan_out = k
+
+    if out is None:
+        output = torch.empty((L_scattered, K), device=DY.device, dtype=DY.dtype)
+    else:
+        assert out.size(0) == L_scattered and out.size(1) == K
+        output = out
+
+    def grid(META):
+        return (
+            triton.cdiv(L_scattered, META["BLOCK_M"]) * triton.cdiv(K, META["BLOCK_K"]),
+        )
+
+    _scatter2scatter_lora_dX[grid](
+        DY,
+        DY.stride(0),
+        DY.stride(1),
+        W,
+        W.stride(0),
+        W.stride(1),
+        W.stride(2),
+        output,
+        output.stride(0),
+        output.stride(1),
+        lora_A,
+        lora_A.stride(0),
+        lora_A.stride(1),
+        lora_B,
+        lora_B.stride(0),
+        lora_B.stride(1),
+        sorted_scattered_idxs,
+        sorted_expert_idxs,
+        FAN_OUT=fan_out,
+        M=M,
+        M_BUCKET=_bucket_m(M),
+        K=K,
+        N=N,
+        E=E,
+        ACTUAL_R=R,
+        # BLOCK_M is autotuned (injected by triton.autotune from Config kwargs)
+        BLOCK_R=BLOCK_R,
+        ACC_TYPE=tl.float32,
+        scaling=scaling,
+        allow_tf32=ALLOW_TF32,
+        dy_grouped=dy_grouped,
+        dx_grouped=dx_grouped,
+        INT64_INDICES=int64_indices,
+    )
+
+    return output
+
+
+# =============================================================================
+# Backward Kernel: LoRA gradient computation (dA, dB)
+# =============================================================================
+
+
+def _group_bwd_lora_configs():
+    """Generate backward (dA/dB) kernel autotune configs.
+
+    Search space includes smaller tile sizes and fewer pipeline stages to
+    support GPUs with limited shared memory (e.g. ~99KB on some GPUs).
+
+    Search space:
+      BLOCK_M:    {32, 64, 128}   (token-loop tile)
+      BLOCK_K:    {32, 64, 128}
+      BLOCK_N:    {32, 64}
+      num_warps:  {4, 8}
+      num_stages: {3, 4, 5}
+
+    The backward kernel also uses BLOCK_R (from LoRA rank), but that is
+    determined by the rank and not autotunable.
+    """
+    configs = []
+    for block_m, block_k, block_n, warps, stages in product(
+        [32, 64, 128],  # BLOCK_M
+        [32, 64, 128],  # BLOCK_K
+        [32, 64],  # BLOCK_N
+        [4, 8],  # num_warps
+        [3, 4, 5],  # num_stages
+    ):
+        configs.append(
+            triton.Config(
+                {"BLOCK_M": block_m, "BLOCK_K": block_k, "BLOCK_N": block_n},
+                num_stages=stages,
+                num_warps=warps,
+            )
+        )
+    return configs
+
+
+def _prune_bwd_lora_configs(configs, named_args, **kwargs):
+    """Prune backward configs based on SMEM capacity and register pressure.
+
+    The backward kernel loads X[BLOCK_M, BLOCK_K] and DY[BLOCK_M, BLOCK_N]
+    in the inner loop, plus holds A[BLOCK_R, BLOCK_K] and B[BLOCK_N, BLOCK_R]
+    for the full expert. We estimate SMEM based on the dominant terms.
+    """
+    smem_cap = _get_smem_capacity()
+    block_r = named_args.get("BLOCK_R", 64)
+
+    scored = []
+    for config in configs:
+        block_m = config.kwargs["BLOCK_M"]
+        block_k = config.kwargs["BLOCK_K"]
+        block_n = config.kwargs["BLOCK_N"]
+        # Inner loop loads X[M,K] and DY[M,N], pipeline over M iterations
+        smem_base = _estimate_smem_usage(config.num_stages, block_m, block_n, block_k)
+        # A[BLOCK_R, BLOCK_K] and B[BLOCK_N, BLOCK_R] held for the full expert
+        smem_lora = (block_r * block_k + block_n * block_r) * 2
+        smem = smem_base + smem_lora
+
+        # Register pressure: dA_acc[R,K], dB_acc[N,R], x[M,K], dy[M,N],
+        # a[R,K], b[N,R], xa[M,R], dy_b[M,R]
+        est_regs = _estimate_register_pressure(
+            config.num_warps,
+            (block_r, block_k),  # dA_acc
+            (block_n, block_r),  # dB_acc
+            (block_m, block_k),  # x tile
+            (block_m, block_n),  # dy tile
+            (block_r, block_k),  # a tile
+            (block_n, block_r),  # b tile
+            (block_m, block_r),  # xa intermediate
+        )
+        if est_regs > _MAX_REGS_SOFT_LIMIT:
+            continue
+
+        scored.append((smem, config))
+
+    pruned = [c for s, c in scored if s <= smem_cap - _SMEM_SLACK]
+    if pruned:
+        return pruned
+    if scored:
+        # All surviving configs exceed SMEM — return the one with smallest usage
+        scored.sort(key=lambda x: x[0])
+        return [scored[0][1]]
+    # All configs pruned by register pressure — fall back to smallest tiles
+    return [
+        min(
+            configs,
+            key=lambda c: (
+                c.kwargs["BLOCK_M"] * c.kwargs["BLOCK_K"] * c.kwargs["BLOCK_N"]
+            ),
+        )
+    ]
+
+
+@triton.autotune(
+    configs=_group_bwd_lora_configs(),
+    key=["M_BUCKET", "N", "K"],
+    prune_configs_by={"early_config_prune": _prune_bwd_lora_configs},
+    reset_to_zero=["DLA_ptr", "DLB_ptr"],
+)
+@triton.heuristics(
+    {
+        "NO_K_MASK": lambda args: (args["K"] % args["BLOCK_K"]) == 0,
+        "NO_N_MASK": lambda args: (args["N"] % args["BLOCK_N"]) == 0,
+    }
+)
+@triton.jit
+def _group_bwd_lora(
+    # Inputs
+    DY_ptr,
+    stride_dym,
+    stride_dyn,
+    X_ptr,
+    stride_xm,
+    stride_xk,
+    # LoRA weights (needed for cross-terms)
+    LA_ptr,
+    stride_la_r,
+    stride_la_k,  # A: [r*E, K]
+    LB_ptr,
+    stride_lb_n,
+    stride_lb_r,  # B: [N, r*E]
+    # Gradient outputs
+    DLA_ptr,
+    stride_dla_r,
+    stride_dla_k,
+    DLB_ptr,
+    stride_dlb_n,
+    stride_dlb_r,
+    # Expert offsets
+    expert_offsets_ptr,
+    # Dimensions
+    M,
+    M_BUCKET,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    ACTUAL_R: tl.constexpr,  # True LoRA rank (for weight indexing)
+    BLOCK_R: tl.constexpr,  # Padded tile size >= max(ACTUAL_R, 16)
+    scaling,
+    # Block sizes
+    BLOCK_M: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    allow_tf32: tl.constexpr,
+    NO_K_MASK: tl.constexpr,
+    NO_N_MASK: tl.constexpr,
+    INT64_INDICES: tl.constexpr = False,
+):
+    """
+    Compute LoRA gradients for each expert on grouped data.
+
+    Grid: (E * cdiv(K, BLOCK_K), cdiv(N, BLOCK_N))
+
+    For expert e:
+      dA[e] = scaling * (dY @ B[e])^T @ X   -> [r, K], accumulate over M tokens
+      dB[e] = scaling * dY^T @ (X @ A[e]^T)  -> [N, r], accumulate over M tokens
+
+    ACTUAL_R is the true LoRA rank. BLOCK_R >= ACTUAL_R is padded for tl.dot min size.
+    """
+    pid0 = tl.program_id(axis=0)
+    pid1 = tl.program_id(axis=1)
+
+    K_BLOCK_COUNT = tl.cdiv(K, BLOCK_K)
+    E_idx = pid0 // K_BLOCK_COUNT
+    K_block_id = pid0 % K_BLOCK_COUNT
+    N_block_id = pid1
+
+    # Get expert's token range from cumulative offsets
+    if INT64_INDICES:
+        if E_idx == 0:
+            start_idx = tl.zeros([], dtype=tl.int64)
+        else:
+            start_idx = tl.load(expert_offsets_ptr + E_idx - 1).to(tl.int64)
+        end_idx = tl.load(expert_offsets_ptr + E_idx).to(tl.int64)
+    else:
+        if E_idx == 0:
+            start_idx = tl.zeros([], dtype=tl.int32)
+        else:
+            start_idx = tl.load(expert_offsets_ptr + E_idx - 1).to(tl.int32)
+        end_idx = tl.load(expert_offsets_ptr + E_idx).to(tl.int32)
+    num_tokens = end_idx - start_idx
+
+    if num_tokens > 0:
+        M_block = tl.arange(0, BLOCK_M)
+        if INT64_INDICES:
+            M_block = M_block.to(tl.int64)
+        K_block = K_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
+        K_mask = K_block < K
+        N_block = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+        N_mask = N_block < N
+        R_block = tl.arange(0, BLOCK_R)
+        R_mask = R_block < ACTUAL_R  # Mask for padding
+
+        lora_offset = E_idx * ACTUAL_R
+
+        # Determine input element type for consistent casting.
+        INPUT_DTYPE = X_ptr.dtype.element_ty
+
+        # Load B[e]: [BLOCK_N, BLOCK_R] (masked on R and N, other=0 for padding)
+        B_blk_ptrs = (
+            LB_ptr
+            + N_block[:, None] * stride_lb_n
+            + (lora_offset + R_block)[None, :] * stride_lb_r
+        )
+        b_e = tl.load(B_blk_ptrs, mask=N_mask[:, None] & R_mask[None, :], other=0.0).to(
+            INPUT_DTYPE
+        )
+
+        # Load A[e]: [BLOCK_R, BLOCK_K] (masked on R and K, other=0 for padding)
+        A_blk_ptrs = (
+            LA_ptr
+            + (lora_offset + R_block)[:, None] * stride_la_r
+            + K_block[None, :] * stride_la_k
+        )
+        a_e = tl.load(A_blk_ptrs, mask=R_mask[:, None] & K_mask[None, :], other=0.0).to(
+            INPUT_DTYPE
+        )
+
+        # Accumulators
+        dA_acc = tl.zeros((BLOCK_R, BLOCK_K), dtype=ACC_TYPE)
+        dB_acc = tl.zeros((BLOCK_N, BLOCK_R), dtype=ACC_TYPE)
+
+        iters = tl.cdiv(num_tokens, BLOCK_M)
+        for i in range(iters):
+            M_idx = start_idx + i * BLOCK_M + M_block
+            M_mask = M_idx < end_idx
+
+            # Load X: [BLOCK_M, BLOCK_K]
+            X_blk_ptrs = (
+                X_ptr + M_idx[:, None] * stride_xm + K_block[None, :] * stride_xk
+            )
+            x = tl.load(
+                X_blk_ptrs, mask=M_mask[:, None] & K_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+
+            # Load dY: [BLOCK_M, BLOCK_N]
+            DY_blk_ptrs = (
+                DY_ptr + M_idx[:, None] * stride_dym + N_block[None, :] * stride_dyn
+            )
+            dy = tl.load(
+                DY_blk_ptrs, mask=M_mask[:, None] & N_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+
+            # X @ A[e]^T: [M, K] @ [K, R] -> [M, R]
+            xa = tl.dot(x, tl.trans(a_e), allow_tf32=allow_tf32)
+
+            # dY @ B[e]: [M, N] @ [N, R] -> [M, R]
+            dy_b = tl.dot(dy, b_e, allow_tf32=allow_tf32)
+
+            # Cast intermediates to input dtype for subsequent tl.dot calls
+            # (tl.dot requires both operands to have the same dtype)
+            dy_b_cast = dy_b.to(INPUT_DTYPE)
+            xa_cast = xa.to(INPUT_DTYPE)
+
+            # dA += (dY @ B)^T @ X: [R, M] @ [M, K] -> [R, K]
+            dA_acc += tl.dot(tl.trans(dy_b_cast), x, allow_tf32=allow_tf32)
+
+            # dB += dY^T @ (X @ A^T): [N, M] @ [M, R] -> [N, R]
+            dB_acc += tl.dot(tl.trans(dy), xa_cast, allow_tf32=allow_tf32)
+
+        # Store dA with scaling (atomic add since multiple N_blocks contribute)
+        # Only store the actual R rows, not the padded ones
+        DLA_blk_ptrs = (
+            DLA_ptr
+            + (lora_offset + R_block)[:, None] * stride_dla_r
+            + K_block[None, :] * stride_dla_k
+        )
+        tl.atomic_add(
+            DLA_blk_ptrs,
+            (dA_acc * scaling).to(DLA_ptr.dtype.element_ty),
+            mask=R_mask[:, None] & K_mask[None, :],
+        )
+
+        # Store dB with scaling (atomic add since multiple K_blocks contribute)
+        DLB_blk_ptrs = (
+            DLB_ptr
+            + N_block[:, None] * stride_dlb_n
+            + (lora_offset + R_block)[None, :] * stride_dlb_r
+        )
+        tl.atomic_add(
+            DLB_blk_ptrs,
+            (dB_acc * scaling).to(DLB_ptr.dtype.element_ty),
+            mask=N_mask[:, None] & R_mask[None, :],
+        )
+
+
+def _group_bwd_split_configs():
+    """Autotune configs for split dA/dB kernels."""
+    configs = []
+    for block_m, block_dim, warps, stages in product(
+        [32, 64, 128],  # BLOCK_M (token tile)
+        [32, 64, 128, 256],  # BLOCK_DIM (K for dA, N for dB — output tile)
+        [4, 8],  # num_warps
+        [3, 4, 5],  # num_stages
+    ):
+        configs.append(
+            triton.Config(
+                {"BLOCK_M": block_m, "BLOCK_DIM": block_dim},
+                num_stages=stages,
+                num_warps=warps,
+            )
+        )
+    return configs
+
+
+def _prune_split_configs(configs, named_args, **kwargs):
+    """Prune split kernel configs based on SMEM capacity and register pressure."""
+    smem_cap = _get_smem_capacity()
+    block_r = named_args.get("BLOCK_R", 64)
+
+    # Fixed inner tile for reduction dimension
+    BLOCK_INNER = 64
+
+    pruned = []
+    for config in configs:
+        block_m = config.kwargs["BLOCK_M"]
+        block_dim = config.kwargs["BLOCK_DIM"]
+        # Inner loop loads: input[M, INNER] and other[M, INNER_or_DIM]
+        smem = config.num_stages * BLOCK_INNER * (block_m + block_dim) * 2
+        # LoRA weights held in registers: [INNER, R] or [R, DIM]
+        smem += (block_r * max(block_dim, BLOCK_INNER)) * 2
+
+        # Register pressure check
+        est_regs = _estimate_register_pressure(
+            config.num_warps,
+            (block_r, block_dim),  # acc
+            (block_m, BLOCK_INNER),  # input tile
+            (block_m, block_dim),  # other tile
+            (block_r, BLOCK_INNER),  # lora weight
+        )
+        if est_regs > _MAX_REGS_SOFT_LIMIT:
+            continue
+
+        if smem <= smem_cap - _SMEM_SLACK:
+            pruned.append(config)
+
+    if pruned:
+        return pruned
+    configs.sort(key=lambda c: c.kwargs["BLOCK_M"] * c.kwargs["BLOCK_DIM"])
+    return [configs[0]]
+
+
+@triton.autotune(
+    configs=_group_bwd_split_configs(),
+    key=["M_BUCKET", "K", "N"],
+    prune_configs_by={"early_config_prune": _prune_split_configs},
+)
+@triton.heuristics(
+    {
+        "NO_DIM_MASK": lambda args: (
+            (args["K"] % args["BLOCK_DIM"]) == 0
+            if args["COMPUTE_DA"]
+            else (args["N"] % args["BLOCK_DIM"]) == 0
+        ),
+    }
+)
+@triton.jit
+def _group_bwd_lora_split(
+    # Data tensors (DY and X are always present)
+    DY_ptr,
+    stride_dym,
+    stride_dyn,
+    X_ptr,
+    stride_xm,
+    stride_xk,
+    # LoRA weight for the inner reduction (B for dA, A for dB)
+    LW_ptr,
+    stride_lw0,
+    stride_lw1,
+    # Output gradient tensor (dA or dB)
+    OUT_ptr,
+    stride_out0,
+    stride_out1,
+    # Expert offsets
+    expert_offsets_ptr,
+    # Dimensions
+    M,
+    M_BUCKET,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    ACTUAL_R: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    INNER_DIM: tl.constexpr,  # reduction dimension (N for dA, K for dB)
+    scaling,
+    # Mode flag
+    COMPUTE_DA: tl.constexpr,  # True = compute dA, False = compute dB
+    # Tile sizes
+    BLOCK_M: tl.constexpr,
+    BLOCK_DIM: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    allow_tf32: tl.constexpr,
+    NO_DIM_MASK: tl.constexpr,
+    INT64_INDICES: tl.constexpr = False,
+):
+    """
+    Unified split kernel for LoRA gradient computation.
+
+    When COMPUTE_DA=True:
+      dA[e] = scaling * (dY @ B[e])^T @ X  →  [R, K]
+      Grid: (E, cdiv(K, BLOCK_DIM))
+      - outer_ptr/stride = X (read [M, K_block])
+      - inner reduction over N using DY and B
+      - output shape [BLOCK_R, BLOCK_DIM]
+
+    When COMPUTE_DA=False:
+      dB[e] = scaling * dY^T @ (X @ A[e]^T)  →  [N, R]
+      Grid: (E, cdiv(N, BLOCK_DIM))
+      - outer_ptr/stride = DY (read [M, N_block])
+      - inner reduction over K using X and A
+      - output shape [BLOCK_DIM, BLOCK_R]
+
+    No atomic adds — each (E, dim_block) pair is written by exactly one block.
+    """
+    E_idx = tl.program_id(0)
+    dim_block_id = tl.program_id(1)
+
+    if INT64_INDICES:
+        if E_idx == 0:
+            start_idx = tl.zeros([], dtype=tl.int64)
+        else:
+            start_idx = tl.load(expert_offsets_ptr + E_idx - 1).to(tl.int64)
+        end_idx = tl.load(expert_offsets_ptr + E_idx).to(tl.int64)
+    else:
+        if E_idx == 0:
+            start_idx = tl.zeros([], dtype=tl.int32)
+        else:
+            start_idx = tl.load(expert_offsets_ptr + E_idx - 1).to(tl.int32)
+        end_idx = tl.load(expert_offsets_ptr + E_idx).to(tl.int32)
+    num_tokens = end_idx - start_idx
+
+    # Output dimension tile (K for dA, N for dB)
+    if COMPUTE_DA:
+        OUT_DIM: tl.constexpr = K  # type: ignore[no-redef]
+    else:
+        OUT_DIM: tl.constexpr = N  # type: ignore[no-redef]
+    dim_block = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM)
+    dim_mask = dim_block < OUT_DIM
+    R_block = tl.arange(0, BLOCK_R)
+    R_mask = R_block < ACTUAL_R
+    lora_offset = E_idx * ACTUAL_R
+
+    # Output pointers — layout differs: dA is [R, K], dB is [N, R]
+    if COMPUTE_DA:
+        out_blk_ptrs = (
+            OUT_ptr
+            + (lora_offset + R_block)[:, None] * stride_out0
+            + dim_block[None, :] * stride_out1
+        )
+        out_mask = R_mask[:, None] & dim_mask[None, :]
+    else:
+        out_blk_ptrs = (
+            OUT_ptr
+            + dim_block[:, None] * stride_out0
+            + (lora_offset + R_block)[None, :] * stride_out1
+        )
+        out_mask = dim_mask[:, None] & R_mask[None, :]
+
+    if num_tokens > 0:
+        M_block = tl.arange(0, BLOCK_M)
+        if INT64_INDICES:
+            M_block = M_block.to(tl.int64)
+        INPUT_DTYPE = X_ptr.dtype.element_ty
+        BLOCK_INNER: tl.constexpr = 64
+        inner_iters = tl.cdiv(INNER_DIM, BLOCK_INNER)
+
+        if COMPUTE_DA:
+            acc = tl.zeros((BLOCK_R, BLOCK_DIM), dtype=ACC_TYPE)
+        else:
+            acc = tl.zeros((BLOCK_DIM, BLOCK_R), dtype=ACC_TYPE)
+
+        M_iters = tl.cdiv(num_tokens, BLOCK_M)
+        for i in range(M_iters):
+            M_idx = start_idx + i * BLOCK_M + M_block
+            M_mask = M_idx < end_idx
+
+            if COMPUTE_DA:
+                # Load X[M, K_block] (the "outer" tensor for dA)
+                outer = tl.load(
+                    X_ptr + M_idx[:, None] * stride_xm + dim_block[None, :] * stride_xk,
+                    mask=M_mask[:, None] & dim_mask[None, :],
+                    other=0.0,
+                ).to(INPUT_DTYPE)
+
+                # Reduce DY[M, :] @ B[e][:, R] over N → [M, R]
+                reduced = tl.zeros((BLOCK_M, BLOCK_R), dtype=ACC_TYPE)
+                inner_range = tl.arange(0, BLOCK_INNER)
+                for j in range(inner_iters):
+                    inn_off = j * BLOCK_INNER + inner_range
+                    inn_mask = inn_off < N
+
+                    dy_tile = tl.load(
+                        DY_ptr
+                        + M_idx[:, None] * stride_dym
+                        + inn_off[None, :] * stride_dyn,
+                        mask=M_mask[:, None] & inn_mask[None, :],
+                        other=0.0,
+                    ).to(INPUT_DTYPE)
+                    # B layout: [N, r*E] → stride_lw0=N stride, stride_lw1=r*E stride
+                    lw_tile = tl.load(
+                        LW_ptr
+                        + inn_off[:, None] * stride_lw0
+                        + (lora_offset + R_block)[None, :] * stride_lw1,
+                        mask=inn_mask[:, None] & R_mask[None, :],
+                        other=0.0,
+                    ).to(INPUT_DTYPE)
+                    reduced += tl.dot(dy_tile, lw_tile, allow_tf32=allow_tf32)
+
+                # dA += (DY@B)^T @ X: [R, M] @ [M, K_block] → [R, K_block]
+                acc += tl.dot(
+                    tl.trans(reduced.to(INPUT_DTYPE)), outer, allow_tf32=allow_tf32
+                )
+            else:
+                # Load DY[M, N_block] (the "outer" tensor for dB)
+                outer = tl.load(
+                    DY_ptr
+                    + M_idx[:, None] * stride_dym
+                    + dim_block[None, :] * stride_dyn,
+                    mask=M_mask[:, None] & dim_mask[None, :],
+                    other=0.0,
+                ).to(INPUT_DTYPE)
+
+                # Reduce X[M, :] @ A[e][:, :].T over K → [M, R]
+                reduced = tl.zeros((BLOCK_M, BLOCK_R), dtype=ACC_TYPE)
+                inner_range = tl.arange(0, BLOCK_INNER)
+                for j in range(inner_iters):
+                    inn_off = j * BLOCK_INNER + inner_range
+                    inn_mask = inn_off < K
+
+                    x_tile = tl.load(
+                        X_ptr
+                        + M_idx[:, None] * stride_xm
+                        + inn_off[None, :] * stride_xk,
+                        mask=M_mask[:, None] & inn_mask[None, :],
+                        other=0.0,
+                    ).to(INPUT_DTYPE)
+                    # A layout: [r*E, K] → stride_lw0=r*E stride, stride_lw1=K stride
+                    # We want A[e]^T: [K, R], so load as [K_inner, R]
+                    lw_tile = tl.load(
+                        LW_ptr
+                        + (lora_offset + R_block)[None, :] * stride_lw0
+                        + inn_off[:, None] * stride_lw1,
+                        mask=inn_mask[:, None] & R_mask[None, :],
+                        other=0.0,
+                    ).to(INPUT_DTYPE)
+                    reduced += tl.dot(x_tile, lw_tile, allow_tf32=allow_tf32)
+
+                # dB += DY^T @ (X@A^T): [N_block, M] @ [M, R] → [N_block, R]
+                acc += tl.dot(
+                    tl.trans(outer), reduced.to(INPUT_DTYPE), allow_tf32=allow_tf32
+                )
+
+        tl.store(
+            out_blk_ptrs, (acc * scaling).to(OUT_ptr.dtype.element_ty), mask=out_mask
+        )
+    else:
+        # Zero out this expert's slice — needed because output uses empty_like
+        if COMPUTE_DA:
+            tl.store(
+                out_blk_ptrs,
+                tl.zeros((BLOCK_R, BLOCK_DIM), dtype=OUT_ptr.dtype.element_ty),
+                mask=out_mask,
+            )
+        else:
+            tl.store(
+                out_blk_ptrs,
+                tl.zeros((BLOCK_DIM, BLOCK_R), dtype=OUT_ptr.dtype.element_ty),
+                mask=out_mask,
+            )
+
+
+def group_bwd_lora(
+    DY: torch.Tensor,
+    X: torch.Tensor,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    E: int,
+    scaling: float,
+    sorted_scattered_idxs: Optional[torch.Tensor] = None,
+    k: int = 1,
+    int64_indices: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute LoRA gradients for A and B on expert-grouped data.
+
+    Uses split dA/dB kernels that eliminate atomic adds by giving each
+    (expert, output_block) pair its own thread block.
+
+    Args:
+        DY: Gradient w.r.t. output [M_total, N] (grouped by expert)
+        X: Input [M_total, K] (grouped by expert)
+        lora_A: LoRA A weights [r*E, K]
+        lora_B: LoRA B weights [N, r*E]
+        expert_offsets: Cumulative token counts per expert [E]
+        E: Number of experts
+        scaling: LoRA scaling factor
+
+    Returns:
+        dA: Gradient for A [r*E, K]
+        dB: Gradient for B [N, r*E]
+    """
+    R = lora_A.size(0) // E
+    K = X.size(1)
+    N = DY.size(1)
+
+    # No zero-init needed: the split kernels write zeros for experts with
+    # zero routed tokens directly in the kernel (else branch).
+    dA = torch.empty_like(lora_A)
+    dB = torch.empty_like(lora_B)
+
+    BLOCK_R = _block_r_for_rank(R)
+
+    def grid_dA(META):
+        return (E, triton.cdiv(K, META["BLOCK_DIM"]))
+
+    _group_bwd_lora_split[grid_dA](
+        DY,
+        DY.stride(0),
+        DY.stride(1),
+        X,
+        X.stride(0),
+        X.stride(1),
+        lora_B,
+        lora_B.stride(0),
+        lora_B.stride(1),
+        dA,
+        dA.stride(0),
+        dA.stride(1),
+        expert_offsets,
+        M=DY.size(0),
+        M_BUCKET=_bucket_m(DY.size(0)),
+        K=K,
+        N=N,
+        ACTUAL_R=R,
+        BLOCK_R=BLOCK_R,
+        INNER_DIM=N,
+        scaling=scaling,
+        COMPUTE_DA=True,
+        ACC_TYPE=tl.float32,
+        allow_tf32=ALLOW_TF32,
+        INT64_INDICES=int64_indices,
+    )
+
+    def grid_dB(META):
+        return (E, triton.cdiv(N, META["BLOCK_DIM"]))
+
+    _group_bwd_lora_split[grid_dB](
+        DY,
+        DY.stride(0),
+        DY.stride(1),
+        X,
+        X.stride(0),
+        X.stride(1),
+        lora_A,
+        lora_A.stride(0),
+        lora_A.stride(1),
+        dB,
+        dB.stride(0),
+        dB.stride(1),
+        expert_offsets,
+        M=DY.size(0),
+        M_BUCKET=_bucket_m(DY.size(0)),
+        K=K,
+        N=N,
+        ACTUAL_R=R,
+        BLOCK_R=BLOCK_R,
+        INNER_DIM=K,
+        scaling=scaling,
+        COMPUTE_DA=False,
+        ACC_TYPE=tl.float32,
+        allow_tf32=ALLOW_TF32,
+        INT64_INDICES=int64_indices,
+    )
+
+    return dA, dB
+
+
+# =============================================================================
+# Backward Kernel: Fused gather + LoRA gradient (dA, dB) — eliminates group()
+# =============================================================================
+
+
+@triton.autotune(
+    configs=_group_bwd_lora_configs(),
+    key=["M_BUCKET", "N", "K"],
+    prune_configs_by={"early_config_prune": _prune_bwd_lora_configs},
+    reset_to_zero=["DLA_ptr", "DLB_ptr"],
+)
+@triton.heuristics(
+    {
+        "NO_K_MASK": lambda args: (args["K"] % args["BLOCK_K"]) == 0,
+        "NO_N_MASK": lambda args: (args["N"] % args["BLOCK_N"]) == 0,
+    }
+)
+@triton.jit
+def _group_bwd_lora_fused(
+    # Inputs (ungrouped or grouped)
+    DY_ptr,
+    stride_dym,
+    stride_dyn,
+    X_ptr,
+    stride_xm,
+    stride_xk,
+    # Scatter indices for gather-on-load
+    sorted_scattered_idxs_ptr,
+    FAN_OUT: tl.constexpr,
+    # LoRA weights (needed for cross-terms)
+    LA_ptr,
+    stride_la_r,
+    stride_la_k,  # A: [r*E, K]
+    LB_ptr,
+    stride_lb_n,
+    stride_lb_r,  # B: [N, r*E]
+    # Gradient outputs
+    DLA_ptr,
+    stride_dla_r,
+    stride_dla_k,
+    DLB_ptr,
+    stride_dlb_n,
+    stride_dlb_r,
+    # Expert offsets
+    expert_offsets_ptr,
+    # Real expert offsets (for M_mask when using token rounding, else same as expert_offsets_ptr)
+    real_expert_offsets_ptr,
+    # Dimensions
+    M,
+    M_BUCKET,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    ACTUAL_R: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    scaling,
+    # Block sizes
+    BLOCK_M: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    allow_tf32: tl.constexpr,
+    NO_K_MASK: tl.constexpr,
+    NO_N_MASK: tl.constexpr,
+    # Whether DY is already in grouped (expert-sorted) order
+    dy_grouped: tl.constexpr = False,
+    INT64_INDICES: tl.constexpr = False,
+):
+    """
+    Fused gather + LoRA gradient computation. Same as _group_bwd_lora but
+    reads X from ungrouped buffers using sorted_scattered_idxs for indirect
+    indexing, eliminating the need for a separate group(X) call.
+
+    When dy_grouped=False (default): both X and DY are read via indirect
+    indexing through sorted_scattered_idxs.  This eliminates both group()
+    calls entirely.
+
+    When dy_grouped=True: DY is already in grouped order (e.g. gate_up_proj
+    backward where grouped_out=True) and is read directly.  Only X uses
+    indirect indexing.  This avoids the group(X) allocation while
+    still supporting the grouped DY case.
+
+    Grid: (E * cdiv(K, BLOCK_K), cdiv(N, BLOCK_N))
+
+    For expert e:
+      dA[e] = scaling * (dY @ B[e])^T @ X   -> [r, K]
+      dB[e] = scaling * dY^T @ (X @ A[e]^T)  -> [N, r]
+
+    Supports token rounding: expert_offsets_ptr gives the iteration range
+    (padded to BLOCK_M multiples), real_expert_offsets_ptr gives the real
+    token count for M_mask (to exclude padding tokens).
+    """
+    pid0 = tl.program_id(axis=0)
+    pid1 = tl.program_id(axis=1)
+
+    K_BLOCK_COUNT = tl.cdiv(K, BLOCK_K)
+    E_idx = pid0 // K_BLOCK_COUNT
+    K_block_id = pid0 % K_BLOCK_COUNT
+    N_block_id = pid1
+
+    # Get expert's token range from cumulative offsets
+    # start_idx/end_idx from expert_offsets_ptr: iteration range (possibly padded)
+    # real_end_idx from real_expert_offsets_ptr: for M_mask (real token count)
+    if INT64_INDICES:
+        if E_idx == 0:
+            start_idx = tl.zeros([], dtype=tl.int64)
+            real_start_idx = tl.zeros([], dtype=tl.int64)
+        else:
+            start_idx = tl.load(expert_offsets_ptr + E_idx - 1).to(tl.int64)
+            real_start_idx = tl.load(real_expert_offsets_ptr + E_idx - 1).to(tl.int64)
+        end_idx = tl.load(expert_offsets_ptr + E_idx).to(tl.int64)
+        real_end_idx = tl.load(real_expert_offsets_ptr + E_idx).to(tl.int64)
+    else:
+        if E_idx == 0:
+            start_idx = tl.zeros([], dtype=tl.int32)
+            real_start_idx = tl.zeros([], dtype=tl.int32)
+        else:
+            start_idx = tl.load(expert_offsets_ptr + E_idx - 1).to(tl.int32)
+            real_start_idx = tl.load(real_expert_offsets_ptr + E_idx - 1).to(tl.int32)
+        end_idx = tl.load(expert_offsets_ptr + E_idx).to(tl.int32)
+        real_end_idx = tl.load(real_expert_offsets_ptr + E_idx).to(tl.int32)
+    num_tokens = end_idx - start_idx
+
+    if num_tokens > 0:
+        M_block = tl.arange(0, BLOCK_M)
+        if INT64_INDICES:
+            M_block = M_block.to(tl.int64)
+        K_block = K_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
+        K_mask = K_block < K
+        N_block = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+        N_mask = N_block < N
+        R_block = tl.arange(0, BLOCK_R)
+        R_mask = R_block < ACTUAL_R
+
+        lora_offset = E_idx * ACTUAL_R
+
+        # Determine input element type for consistent casting.
+        INPUT_DTYPE = X_ptr.dtype.element_ty
+
+        # Load B[e] and A[e] — same as non-fused kernel
+        B_blk_ptrs = (
+            LB_ptr
+            + N_block[:, None] * stride_lb_n
+            + (lora_offset + R_block)[None, :] * stride_lb_r
+        )
+        b_e = tl.load(B_blk_ptrs, mask=N_mask[:, None] & R_mask[None, :], other=0.0).to(
+            INPUT_DTYPE
+        )
+
+        A_blk_ptrs = (
+            LA_ptr
+            + (lora_offset + R_block)[:, None] * stride_la_r
+            + K_block[None, :] * stride_la_k
+        )
+        a_e = tl.load(A_blk_ptrs, mask=R_mask[:, None] & K_mask[None, :], other=0.0).to(
+            INPUT_DTYPE
+        )
+
+        # Accumulators
+        dA_acc = tl.zeros((BLOCK_R, BLOCK_K), dtype=ACC_TYPE)
+        dB_acc = tl.zeros((BLOCK_N, BLOCK_R), dtype=ACC_TYPE)
+
+        real_num_tokens = real_end_idx - real_start_idx
+        iters = tl.cdiv(num_tokens, BLOCK_M)
+        for i in range(iters):
+            M_idx = start_idx + i * BLOCK_M + M_block
+            # Use real token count for masking (excludes padding tokens)
+            M_local = i * BLOCK_M + M_block
+            M_mask = M_local < real_num_tokens
+
+            # Fused gather: load scatter indices for indirect X access
+            if INT64_INDICES:
+                scatter_idx = tl.load(
+                    sorted_scattered_idxs_ptr + M_idx, mask=M_mask, other=0
+                ).to(tl.int64)
+            else:
+                scatter_idx = tl.load(
+                    sorted_scattered_idxs_ptr + M_idx, mask=M_mask, other=0
+                ).to(tl.int32)
+            X_token_idx = scatter_idx // FAN_OUT  # X is [M, K], not expanded by k
+
+            # Load X via indirect index: [BLOCK_M, BLOCK_K]
+            X_blk_ptrs = (
+                X_ptr + X_token_idx[:, None] * stride_xm + K_block[None, :] * stride_xk
+            )
+            x = tl.load(
+                X_blk_ptrs, mask=M_mask[:, None] & K_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+
+            # Load DY: indirect via scatter_idx when ungrouped, direct via M_idx when grouped
+            if dy_grouped:
+                DY_blk_ptrs = (
+                    DY_ptr + M_idx[:, None] * stride_dym + N_block[None, :] * stride_dyn
+                )
+            else:
+                DY_blk_ptrs = (
+                    DY_ptr
+                    + scatter_idx[:, None] * stride_dym
+                    + N_block[None, :] * stride_dyn
+                )
+            dy = tl.load(
+                DY_blk_ptrs, mask=M_mask[:, None] & N_mask[None, :], other=0.0
+            ).to(INPUT_DTYPE)
+
+            # X @ A[e]^T: [M, K] @ [K, R] -> [M, R]
+            xa = tl.dot(x, tl.trans(a_e), allow_tf32=allow_tf32)
+
+            # dY @ B[e]: [M, N] @ [N, R] -> [M, R]
+            dy_b = tl.dot(dy, b_e, allow_tf32=allow_tf32)
+
+            dy_b_cast = dy_b.to(INPUT_DTYPE)
+            xa_cast = xa.to(INPUT_DTYPE)
+
+            # dA += (dY @ B)^T @ X: [R, M] @ [M, K] -> [R, K]
+            dA_acc += tl.dot(tl.trans(dy_b_cast), x, allow_tf32=allow_tf32)
+
+            # dB += dY^T @ (X @ A^T): [N, M] @ [M, R] -> [N, R]
+            dB_acc += tl.dot(tl.trans(dy), xa_cast, allow_tf32=allow_tf32)
+
+        # Store dA with scaling (atomic add since multiple N_blocks contribute)
+        DLA_blk_ptrs = (
+            DLA_ptr
+            + (lora_offset + R_block)[:, None] * stride_dla_r
+            + K_block[None, :] * stride_dla_k
+        )
+        tl.atomic_add(
+            DLA_blk_ptrs,
+            (dA_acc * scaling).to(DLA_ptr.dtype.element_ty),
+            mask=R_mask[:, None] & K_mask[None, :],
+        )
+
+        # Store dB with scaling (atomic add since multiple K_blocks contribute)
+        DLB_blk_ptrs = (
+            DLB_ptr
+            + N_block[:, None] * stride_dlb_n
+            + (lora_offset + R_block)[None, :] * stride_dlb_r
+        )
+        tl.atomic_add(
+            DLB_blk_ptrs,
+            (dB_acc * scaling).to(DLB_ptr.dtype.element_ty),
+            mask=N_mask[:, None] & R_mask[None, :],
+        )
+
+
+def group_bwd_lora_fused(
+    DY: torch.Tensor,
+    X: torch.Tensor,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    E: int,
+    k: int,
+    scaling: float,
+    real_expert_offsets: Optional[torch.Tensor] = None,
+    dy_grouped: bool = False,
+    int64_indices: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Fused gather + LoRA gradient computation. Same result as
+    group(X) + group(DY) + group_bwd_lora(DY, X, ...) but without
+    the intermediate grouped buffers.
+
+    Args:
+        DY: Gradient w.r.t. output [M*k, N].
+            If dy_grouped=False: ungrouped (original token order), read via
+            indirect indexing through sorted_scattered_idxs.
+            If dy_grouped=True: already in grouped (expert-sorted) order,
+            read directly.
+        X: Input [M, K] (ungrouped, original token order).  Always read via
+            indirect indexing through sorted_scattered_idxs.
+        lora_A: LoRA A weights [r*E, K]
+        lora_B: LoRA B weights [N, r*E]
+        expert_offsets: Cumulative token counts per expert [E]
+            (or padded offsets if using token rounding)
+        sorted_scattered_idxs: Maps grouped position -> original position [M*k]
+            (or padded version if using token rounding)
+        E: Number of experts
+        k: Fan-out (top-k)
+        scaling: LoRA scaling factor
+        real_expert_offsets: Original cumulative counts for M_mask when using
+            token rounding. If None, expert_offsets is used for both.
+        dy_grouped: Whether DY is already in grouped order (default False).
+            When True, avoids indirect indexing for DY, used for gate_up_proj
+            backward where grouped_out=True.
+
+    Returns:
+        dA: Gradient for A [r*E, K]
+        dB: Gradient for B [N, r*E]
+    """
+    R = lora_A.size(0) // E
+    K = X.size(1)
+    N = DY.size(1)
+
+    # Zero-init for atomic accumulation
+    dA = torch.zeros_like(lora_A)
+    dB = torch.zeros_like(lora_B)
+
+    BLOCK_R = _block_r_for_rank(R)
+
+    if real_expert_offsets is None:
+        real_expert_offsets = expert_offsets
+
+    def grid(META):
+        return (
+            E * triton.cdiv(K, META["BLOCK_K"]),
+            triton.cdiv(N, META["BLOCK_N"]),
+        )
+
+    _group_bwd_lora_fused[grid](
+        DY,
+        DY.stride(0),
+        DY.stride(1),
+        X,
+        X.stride(0),
+        X.stride(1),
+        sorted_scattered_idxs,
+        FAN_OUT=k,
+        LA_ptr=lora_A,
+        stride_la_r=lora_A.stride(0),
+        stride_la_k=lora_A.stride(1),
+        LB_ptr=lora_B,
+        stride_lb_n=lora_B.stride(0),
+        stride_lb_r=lora_B.stride(1),
+        DLA_ptr=dA,
+        stride_dla_r=dA.stride(0),
+        stride_dla_k=dA.stride(1),
+        DLB_ptr=dB,
+        stride_dlb_n=dB.stride(0),
+        stride_dlb_r=dB.stride(1),
+        expert_offsets_ptr=expert_offsets,
+        real_expert_offsets_ptr=real_expert_offsets,
+        M=sorted_scattered_idxs.size(0),
+        M_BUCKET=_bucket_m(sorted_scattered_idxs.size(0)),
+        K=K,
+        N=N,
+        ACTUAL_R=R,
+        BLOCK_R=BLOCK_R,
+        scaling=scaling,
+        ACC_TYPE=tl.float32,
+        allow_tf32=ALLOW_TF32,
+        dy_grouped=dy_grouped,
+        INT64_INDICES=int64_indices,
+    )
+
+    return dA, dB
+
+
+# =============================================================================
+# Fused MXFP4 Forward / dX Kernels
+# =============================================================================
+#
+# These mirror ``_scatter2scatter_lora`` and ``_scatter2scatter_lora_dX`` but
+# load the base weight tile from a packed MXFP4 buffer + E8M0 scale buffer
+# instead of a dense bf16 tile. The K-loop unpacks two fp4 values per uint8
+# byte, looks them up in a 16-entry fp32 codebook, multiplies by
+# ``2^(scale_byte - 127)``, and casts back to bf16 for ``tl.dot``.
+#
+# Layout conventions (kernel coordinates):
+#   * Forward kernel: logical W is ``[E, K, N]`` (block axis = K).
+#     - packed: ``[E, N, K/2]`` uint8 — stored with the contraction axis K
+#       contiguous (matches torchao MXTensor's natural last-dim block layout
+#       once you treat the W storage as ``[E, N, K]``).
+#     - scale:  ``[E, N, K/32]`` uint8 (E8M0).
+#   * dX kernel: reuses the *forward* MX layout ``[E, N, K/2]`` (no
+#     pre-transpose). The kernel iterates the N reduction in outer tiles and,
+#     for each (K_tile, N_tile), decodes nibbles along the K rows of the
+#     packed tile and broadcasts scales within each ``MX_BLOCK_SIZE`` K-block,
+#     yielding the same dequantized ``W[e, k, n]`` values the forward path
+#     consumes. This deliberately avoids a "pre-transpose for dX" step that
+#     would dequantize + transpose + re-quantize the active-experts slice:
+#     that round-trip introduces a second MX rounding error on top of the
+#     forward quantization, perturbing dX in ways that are hard to bound.
+#     Reusing the forward buffer keeps numerics bitwise-comparable to a
+#     dequant-then-MMA dX reference.
+#
+# ``BLOCK_K`` must be a multiple of the OCP block size (32) so that each
+# K-tile aligns with whole scale blocks for both the forward and dX
+# kernels. The autotune config search space is pruned accordingly in
+# ``_prune_fwd_mx_configs`` / ``_prune_dX_mx_configs``.
+
+_MX_BLOCK_SIZE = 32
+
+
+@triton.jit
+def _compute_expert_block_lora_mxfp4(
+    E_idx,
+    E_mask,
+    M_in_idx,
+    N_block,
+    N_mask,
+    # X
+    X_ptr,
+    stride_xm,
+    stride_xk,
+    # Packed MXFP4 weight: [E, N, K/2] uint8
+    Wp_ptr,
+    stride_wpe,
+    stride_wpn,
+    stride_wpk,
+    # E8M0 scale: [E, N, K/32] uint8
+    Ws_ptr,
+    stride_wse,
+    stride_wsn,
+    stride_wsk,
+    # FP4 -> fp32 codebook (16 values)
+    Codebook_ptr,
+    # LoRA
+    A_ptr,
+    stride_ar,
+    stride_ak,
+    B_ptr,
+    stride_bn,
+    stride_br,
+    K,
+    ACTUAL_R: tl.constexpr,
+    acc,
+    no_k_mask,
+    BLOCK_M: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    MX_BLOCK_SIZE: tl.constexpr,
+    SCALE_LINEAR: tl.constexpr,
+    scaling,
+    allow_tf32: tl.constexpr,
+):
+    """Forward inner loop for MXFP4 expert weights.
+
+    Computes ``acc += X @ dequant(W_e) + scaling * (X @ A_e^T) @ B_e^T`` for
+    the active token rows in this M-tile assigned to expert ``E_idx``.
+
+    Each K-loop iteration loads a ``[BLOCK_N, BLOCK_K/2]`` packed tile and a
+    ``[BLOCK_N, BLOCK_K/MX_BLOCK_SIZE]`` scale tile, unpacks to bf16, and
+    transposes for the matmul.
+    """
+    K_block = tl.arange(0, BLOCK_K)
+    K_byte_block = K_block // 2
+    K_is_high = (K_block % 2) == 1
+    K_scale_block = K_block // MX_BLOCK_SIZE
+    R_block = tl.arange(0, BLOCK_R)
+    R_mask = R_block < ACTUAL_R
+
+    INPUT_DTYPE = X_ptr.dtype.element_ty
+
+    # X pointers
+    X_blk_ptrs = X_ptr + M_in_idx[:, None] * stride_xm + K_block[None, :] * stride_xk
+
+    # LoRA A pointers
+    A_expert_offset = E_idx * ACTUAL_R
+    A_blk_ptrs = (
+        A_ptr
+        + (A_expert_offset + R_block)[:, None] * stride_ar
+        + K_block[None, :] * stride_ak
+    )
+
+    # Packed W pointers: tile shape [BLOCK_N, BLOCK_K] (each byte loaded twice)
+    Wp_blk_ptrs = (
+        Wp_ptr
+        + E_idx * stride_wpe
+        + N_block[:, None] * stride_wpn
+        + K_byte_block[None, :] * stride_wpk
+    )
+    # Scale pointers: tile shape [BLOCK_N, BLOCK_K] (broadcast within block)
+    Ws_blk_ptrs = (
+        Ws_ptr
+        + E_idx * stride_wse
+        + N_block[:, None] * stride_wsn
+        + K_scale_block[None, :] * stride_wsk
+    )
+
+    iters = tl.cdiv(K, BLOCK_K)
+    xa_acc = tl.zeros((BLOCK_M, BLOCK_R), dtype=tl.float32)
+
+    for i in range(iters):
+        if no_k_mask:
+            K_mask_iter = K_block >= 0  # all-true [BLOCK_K]
+        else:
+            K_mask_iter = (i * BLOCK_K + K_block) < K
+        x_mask = E_mask[:, None] & K_mask_iter[None, :]
+        a_mask = R_mask[:, None] & K_mask_iter[None, :]
+        w_mask = N_mask[:, None] & K_mask_iter[None, :]
+
+        x = tl.load(X_blk_ptrs, mask=x_mask, other=0.0).to(INPUT_DTYPE)
+        a = tl.load(A_blk_ptrs, mask=a_mask, other=0.0).to(INPUT_DTYPE)
+
+        # MXFP4 dequant
+        packed = tl.load(Wp_blk_ptrs, mask=w_mask, other=0).to(tl.int32)
+        nibble = tl.where(K_is_high[None, :], (packed >> 4) & 0xF, packed & 0xF)
+        codebook_val = tl.load(Codebook_ptr + nibble)  # [BLOCK_N, BLOCK_K] fp32
+        if SCALE_LINEAR:
+            # NVFP4: scale is a linear fp multiplier (E4M3 block scale * per-tensor)
+            scale_fp = tl.load(Ws_blk_ptrs, mask=w_mask, other=0.0).to(tl.float32)
+        else:
+            # MXFP4: E8M0 scale, decode as 2^(byte-127)
+            scale_byte = tl.load(Ws_blk_ptrs, mask=w_mask, other=0).to(tl.int32)
+            scale_fp = tl.exp2((scale_byte - 127).to(tl.float32))
+        w_dq_nk = (codebook_val * scale_fp).to(INPUT_DTYPE)  # [BLOCK_N, BLOCK_K]
+        w_tile = tl.trans(w_dq_nk)  # [BLOCK_K, BLOCK_N]
+
+        # Base: acc += X @ W
+        acc += tl.dot(x, w_tile, allow_tf32=allow_tf32).to(tl.float32)
+        # LoRA: xa_acc += X @ A^T
+        xa_acc += tl.dot(x, tl.trans(a), allow_tf32=allow_tf32).to(tl.float32)
+
+        X_blk_ptrs += BLOCK_K * stride_xk
+        A_blk_ptrs += BLOCK_K * stride_ak
+        Wp_blk_ptrs += (BLOCK_K // 2) * stride_wpk
+        Ws_blk_ptrs += (BLOCK_K // MX_BLOCK_SIZE) * stride_wsk
+
+    # Epilogue (B @ xa_acc^T) — identical to dense path
+    B_expert_offset = E_idx * ACTUAL_R
+    B_blk_ptrs = (
+        B_ptr
+        + N_block[:, None] * stride_bn
+        + (B_expert_offset + R_block)[None, :] * stride_br
+    )
+    b = tl.load(B_blk_ptrs, mask=N_mask[:, None] & R_mask[None, :], other=0.0)
+    b_inp = b.to(INPUT_DTYPE)
+    lora_out = tl.dot(xa_acc.to(INPUT_DTYPE), tl.trans(b_inp), allow_tf32=allow_tf32)
+    acc += scaling * lora_out
+    return acc
+
+
+def _scatter2scatter_lora_mx_configs():
+    """Forward MX kernel configs. BLOCK_K must be a multiple of MX_BLOCK_SIZE.
+
+    ``num_stages=2`` is included: a single-GPU microbench found it optimal for the GLM expert
+    shapes on sm90, and it was absent from the original [3,4,5] search. For recognized
+    (arch, shape) profiles, :func:`_prune_fwd_mx_configs` narrows this full matrix to a small
+    pre-tuned subset (see ``_FWD_MX_PROFILES``); all other hardware/shapes keep the full
+    SMEM/register-pruned search.
+    """
+    configs = []
+    for block_m, block_n, block_k, warps, stages in product(
+        [32, 64, 128],
+        [32, 64],
+        [32, 64, 128],  # all multiples of MX_BLOCK_SIZE=32
+        [4, 8],
+        [2, 3, 4, 5],
+    ):
+        configs.append(
+            triton.Config(
+                {"BLOCK_M": block_m, "BLOCK_N": block_n, "BLOCK_K": block_k},
+                num_stages=stages,
+                num_warps=warps,
+            )
+        )
+    return configs
+
+
+# Pre-tuned forward-MX config subsets for known (GPU arch, expert-GEMM shape) profiles. Emitting a
+# small *measured* subset instead of running the full autotune avoids per-rank sweep-time skew under
+# expert parallelism (DeepEP) — where each rank's recv-token count lands on a different M_BUCKET, so
+# a large per-rank sweep desyncs ranks and trips DeepEP's combine barrier. Any arch/shape not in the
+# table keeps the full SMEM/register-pruned search. spec = (BLOCK_M, BLOCK_N, BLOCK_K, warps, stages).
+_GLM_MX_SUBSET = [
+    (
+        128,
+        64,
+        32,
+        4,
+        2,
+    ),  # H200/sm90 microbench winner (invariant over M_BUCKET 32k-237k)
+    (128, 32, 64, 4, 2),  # near-tie alternate
+    (128, 32, 64, 8, 3),  # robust fallback within the stock search space
+    (128, 64, 64, 4, 2),  # neighbor
+]
+_FWD_MX_PROFILES = {
+    # sm90 (H200), GLM-5.2 expert GEMMs: gate_up (N=4096, K=6144), down (N=6144, K=2048).
+    (9, 0): {(4096, 6144): _GLM_MX_SUBSET, (6144, 2048): _GLM_MX_SUBSET},
+}
+
+
+def _profile_fwd_mx_configs(configs, meta):
+    """If the running GPU + this kernel's (N, K) match a known profile and M is in the tuned
+    (large) regime, return the pre-measured config subset (filtered from ``configs``). Else None.
+
+    ``meta`` is the kernel's compile-time kwargs (N, K, M_BUCKET, ...) that Triton passes to the
+    ``early_config_prune`` hook — the constexpr values are NOT in ``named_args`` (pointers/strides)."""
+    try:
+        cap = torch.cuda.get_device_capability()
+    except Exception:  # pragma: no cover
+        return None
+    table = _FWD_MX_PROFILES.get(cap)
+    if not table:
+        return None
+    specs = table.get((meta.get("N"), meta.get("K")))
+    mb = meta.get("M_BUCKET")
+    if (
+        not specs or mb is None or mb < 16384
+    ):  # only the large-M EP regime this was tuned for
+        return None
+    want = set(specs)
+    sel = [
+        c
+        for c in configs
+        if (
+            c.kwargs["BLOCK_M"],
+            c.kwargs["BLOCK_N"],
+            c.kwargs["BLOCK_K"],
+            c.num_warps,
+            c.num_stages,
+        )
+        in want
+    ]
+    return sel or None
+
+
+def _prune_fwd_mx_configs(configs, named_args, **kwargs):
+    """Prune MX forward configs by SMEM and register pressure.
+
+    MX-aware accounting adds the packed W tile and the scale tile per
+    pipeline stage to the base GEMM SMEM estimate. Both tiles are sized
+    [BLOCK_N, BLOCK_K] uint8 in the kernel: the packed buffer reads each
+    byte twice (because K_byte = K // 2 indexes a [BLOCK_K]-wide vector
+    into a K/2-stride buffer), and the scale buffer reads each byte
+    MX_BLOCK_SIZE times (broadcast within each K-block). This matches
+    the conservative full-tile accounting in ``_prune_dX_mx_configs``.
+    Also require BLOCK_K % MX_BLOCK_SIZE == 0.
+    """
+    # Known (arch, shape) profile -> emit the small pre-tuned subset, skipping the full sweep.
+    # N/K/M_BUCKET are constexpr -> they're in **kwargs, not named_args (pointers/strides).
+    profiled = _profile_fwd_mx_configs(configs, kwargs)
+    if profiled is not None:
+        return profiled
+
+    smem_cap = _get_smem_capacity()
+    block_r = named_args.get("BLOCK_R", 64)
+
+    scored = []
+    for config in configs:
+        block_m = config.kwargs["BLOCK_M"]
+        block_n = config.kwargs["BLOCK_N"]
+        block_k = config.kwargs["BLOCK_K"]
+        if block_k % _MX_BLOCK_SIZE != 0:
+            continue
+        # Base GEMM tiles (X, dequantized W, acc)
+        smem_base = _estimate_smem_usage(config.num_stages, block_m, block_n, block_k)
+        # MX-specific loads per pipeline stage: packed and scale tiles are
+        # both [BLOCK_N, BLOCK_K] bytes (see docstring for why each is full
+        # tile size, not BLOCK_K/2 or BLOCK_K/MX_BLOCK_SIZE).
+        smem_packed = config.num_stages * block_n * block_k * 1
+        smem_scale = config.num_stages * block_n * block_k * 1
+        # LoRA tiles
+        smem_lora_loop = config.num_stages * block_r * block_k * 2
+        smem_lora_epilogue = block_n * block_r * 2
+        smem = (
+            smem_base + smem_packed + smem_scale + smem_lora_loop + smem_lora_epilogue
+        )
+
+        est_regs = _estimate_register_pressure(
+            config.num_warps,
+            (block_m, block_n),  # acc
+            (block_m, block_r),  # xa_acc
+            (block_m, block_k),  # x tile
+            (block_n, block_k),  # dequantized w (before transpose)
+            (block_r, block_k),  # a tile
+            (block_n, block_r),  # b tile (epilogue)
+        )
+        if est_regs > _MAX_REGS_SOFT_LIMIT:
+            continue
+
+        scored.append((smem, config))
+
+    pruned = [c for s, c in scored if s <= smem_cap - _SMEM_SLACK]
+    if pruned:
+        return pruned
+    if scored:
+        scored.sort(key=lambda x: x[0])
+        return [scored[0][1]]
+    return [
+        min(
+            configs,
+            key=lambda c: (
+                c.kwargs["BLOCK_M"] * c.kwargs["BLOCK_N"] * c.kwargs["BLOCK_K"]
+            ),
+        )
+    ]
+
+
+@triton.autotune(
+    configs=_scatter2scatter_lora_mx_configs(),
+    key=["M_BUCKET", "N", "K"],
+    prune_configs_by={"early_config_prune": _prune_fwd_mx_configs},
+)
+@triton.heuristics(
+    {
+        "NO_K_MASK": lambda args: (args["K"] % args["BLOCK_K"]) == 0,
+        "NO_N_MASK": lambda args: (args["N"] % args["BLOCK_N"]) == 0,
+    }
+)
+@triton.jit
+def _scatter2scatter_lora_mx(
+    # X
+    X_ptr,
+    stride_xm: tl.constexpr,
+    stride_xk: tl.constexpr,
+    # Packed MXFP4 W [E, N, K/2]
+    Wp_ptr,
+    stride_wpe,
+    stride_wpn: tl.constexpr,
+    stride_wpk: tl.constexpr,
+    # E8M0 scale [E, N, K/32]
+    Ws_ptr,
+    stride_wse,
+    stride_wsn: tl.constexpr,
+    stride_wsk: tl.constexpr,
+    # FP4 codebook (16 fp32 values)
+    Codebook_ptr,
+    # Output
+    Y_ptr,
+    stride_ym: tl.constexpr,
+    stride_yn: tl.constexpr,
+    # Bias
+    Bias_ptr,
+    stride_bias_e: tl.constexpr,
+    stride_bias_n: tl.constexpr,
+    # LoRA
+    LA_ptr,
+    stride_la_r,
+    stride_la_k,
+    LB_ptr,
+    stride_lb_n,
+    stride_lb_r,
+    # Routing
+    grouped_idx_ptr,
+    expert_idxs_ptr,
+    # Dimensions
+    FAN_OUT: tl.constexpr,
+    M,
+    M_BUCKET,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    E: tl.constexpr,
+    ACTUAL_R: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    MX_BLOCK_SIZE: tl.constexpr,
+    SCALE_LINEAR: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    scaling,
+    allow_tf32: tl.constexpr,
+    x_grouped: tl.constexpr,
+    y_grouped: tl.constexpr,
+    NO_K_MASK: tl.constexpr,
+    NO_N_MASK: tl.constexpr,
+    INT64_INDICES: tl.constexpr = False,
+):
+    """Fused scatter2scatter forward with MXFP4/NVFP4 base weights + LoRA."""
+    pid = tl.program_id(axis=0)
+    N_BLOCK_COUNT = tl.cdiv(N, BLOCK_N)
+    M_block_id = pid // N_BLOCK_COUNT
+    N_block_id = pid % N_BLOCK_COUNT
+
+    M_block = M_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    if INT64_INDICES:
+        M_block = M_block.to(tl.int64)
+    N_block = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+    N_mask = N_block < N
+    M_boundary_mask = M_block < (FAN_OUT * M)
+
+    E_idxs = tl.load(expert_idxs_ptr + M_block, mask=M_boundary_mask, other=E)
+    no_k_mask = NO_K_MASK
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
+
+    E_first_idx = tl.min(E_idxs)
+    E_last_idx = tl.minimum(tl.max(E_idxs), E - 1)
+    if INT64_INDICES:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int64)
+    else:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int32)
+
+    for E_idx in range(E_first_idx, E_last_idx + 1):
+        E_mask = E_idxs == E_idx
+        if x_grouped:
+            M_in_idx = M_block
+        else:
+            M_in_idx = M_idx // FAN_OUT
+
+        acc = _compute_expert_block_lora_mxfp4(
+            E_idx,
+            E_mask,
+            M_in_idx,
+            N_block,
+            N_mask,
+            X_ptr,
+            stride_xm,
+            stride_xk,
+            Wp_ptr,
+            stride_wpe,
+            stride_wpn,
+            stride_wpk,
+            Ws_ptr,
+            stride_wse,
+            stride_wsn,
+            stride_wsk,
+            Codebook_ptr,
+            LA_ptr,
+            stride_la_r,
+            stride_la_k,
+            LB_ptr,
+            stride_lb_n,
+            stride_lb_r,
+            K,
+            ACTUAL_R,
+            acc,
+            no_k_mask,
+            BLOCK_M,
+            BLOCK_K,
+            BLOCK_N,
+            BLOCK_R,
+            MX_BLOCK_SIZE,
+            SCALE_LINEAR,
+            scaling,
+            allow_tf32=allow_tf32,
+        )
+
+    if Bias_ptr is not None:
+        B_blk_ptrs = (
+            Bias_ptr
+            + E_idxs[:, None] * stride_bias_e
+            + N_block[None, :] * stride_bias_n
+        )
+        acc += tl.load(B_blk_ptrs, mask=M_boundary_mask[:, None] & N_mask[None, :])
+
+    if y_grouped:
+        M_out_idx = M_block
+    else:
+        M_out_idx = M_idx
+    Y_blk_ptrs = Y_ptr + (M_out_idx[:, None] * stride_ym + N_block[None, :] * stride_yn)
+    tl.store(Y_blk_ptrs, acc, mask=M_boundary_mask[:, None] & N_mask[None, :])
+
+
+def scatter2scatter_lora_mx(
+    X: torch.Tensor,
+    W_mx,  # MXWeights with layout=FWD
+    sorted_expert_idxs: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    k: int,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    scaling: float,
+    b: Optional[torch.Tensor] = None,
+    x_grouped: bool = False,
+    y_grouped: bool = False,
+    out: Optional[torch.Tensor] = None,
+    int64_indices: bool = False,
+) -> torch.Tensor:
+    """Forward dispatcher for the fused MXFP4 + LoRA kernel.
+
+    ``W_mx`` is an ``MXWeights`` instance in ``FWD`` layout (block axis = K).
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+        MXLayout,
+        fp4_codebook,
+    )
+
+    assert W_mx.layout == MXLayout.FWD, (
+        f"scatter2scatter_lora_mx requires FWD layout, got {W_mx.layout}"
+    )
+    assert sorted_scattered_idxs.size(0) == sorted_expert_idxs.size(0)
+    assert sorted_scattered_idxs.size(0) == X.size(0) * k
+
+    K = W_mx.K
+    N = W_mx.N
+    E = W_mx.packed.size(0)
+    R = lora_A.size(0) // E
+    BLOCK_R = _block_r_for_rank(R)
+    L_scattered = sorted_expert_idxs.size(0)
+
+    if out is None:
+        output = torch.empty((L_scattered, N), device=X.device, dtype=X.dtype)
+    else:
+        assert out.size(0) == L_scattered and out.size(1) == N
+        output = out
+
+    def grid(META):
+        return (
+            triton.cdiv(L_scattered, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),
+        )
+
+    if b is None:
+        stride_be = stride_bn = 0
+        b_ptr = None
+    else:
+        stride_be, stride_bn = b.stride()
+        b_ptr = b
+
+    codebook = fp4_codebook(X.device)
+
+    _scatter2scatter_lora_mx[grid](
+        X,
+        X.stride(0),
+        X.stride(1),
+        W_mx.packed,
+        W_mx.packed.stride(0),
+        W_mx.packed.stride(1),
+        W_mx.packed.stride(2),
+        W_mx.scales,
+        W_mx.scales.stride(0),
+        W_mx.scales.stride(1),
+        W_mx.scales.stride(2),
+        codebook,
+        output,
+        output.stride(0),
+        output.stride(1),
+        b_ptr,
+        stride_be,
+        stride_bn,
+        lora_A,
+        lora_A.stride(0),
+        lora_A.stride(1),
+        lora_B,
+        lora_B.stride(0),
+        lora_B.stride(1),
+        sorted_scattered_idxs,
+        sorted_expert_idxs,
+        FAN_OUT=k,
+        M=X.size(0),
+        M_BUCKET=_bucket_m(X.size(0)),
+        K=K,
+        N=N,
+        E=E,
+        ACTUAL_R=R,
+        BLOCK_R=BLOCK_R,
+        MX_BLOCK_SIZE=W_mx.block_size,
+        SCALE_LINEAR=W_mx.scale_is_linear,
+        ACC_TYPE=tl.float32,
+        scaling=scaling,
+        allow_tf32=ALLOW_TF32,
+        x_grouped=x_grouped,
+        y_grouped=y_grouped,
+        INT64_INDICES=int64_indices,
+    )
+    return output
+
+
+@triton.jit
+def _compute_expert_block_lora_dX_mxfp4(
+    E_idx,
+    E_mask,
+    M_in_idx,
+    K_block,
+    K_mask,
+    # dY [M, N]
+    DY_ptr,
+    stride_dym,
+    stride_dyn,
+    # Packed MXFP4 W in FWD layout: [E, N, K/2] uint8 (block axis = K)
+    Wp_ptr,
+    stride_wpe,
+    stride_wpn,
+    stride_wpk,
+    # Scale [E, N, K/32]
+    Ws_ptr,
+    stride_wse,
+    stride_wsn,
+    stride_wsk,
+    # FP4 codebook
+    Codebook_ptr,
+    # LoRA
+    A_ptr,
+    stride_ar,
+    stride_ak,
+    B_ptr,
+    stride_bn,
+    stride_br,
+    N,
+    ACTUAL_R: tl.constexpr,
+    acc,
+    no_n_mask,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    MX_BLOCK_SIZE: tl.constexpr,
+    SCALE_LINEAR: tl.constexpr,
+    scaling,
+    allow_tf32: tl.constexpr,
+):
+    """dX inner loop for MXFP4 base weights, FWD layout (block axis = K).
+
+    Computes ``acc += DY @ dequant(W_e)^T + scaling * (DY @ B_e) @ A_e`` for
+    the active token rows in this M-tile assigned to expert ``E_idx``.
+
+    W storage is the *forward* MX layout: packed ``[E, N, K/2]`` and scale
+    ``[E, N, K/32]`` — the same buffer used by the forward kernel, no
+    pre-transpose / re-quantize required. Per N-loop iter we load packed
+    ``[BLOCK_K, BLOCK_N]`` and scale ``[BLOCK_K, BLOCK_N]`` tiles indexed
+    K-as-row × N-as-col; nibbles are extracted along the K axis (the byte
+    is shared by adjacent K rows) and scales broadcast within their
+    ``MX_BLOCK_SIZE``-element K block.
+    """
+    N_block = tl.arange(0, BLOCK_N)
+    # K-axis decode tables (K-along-rows of the W tile)
+    K_byte_block = K_block // 2
+    K_is_high = (K_block % 2) == 1
+    K_scale_block = K_block // MX_BLOCK_SIZE
+    R_block = tl.arange(0, BLOCK_R)
+    R_mask = R_block < ACTUAL_R
+
+    INPUT_DTYPE = DY_ptr.dtype.element_ty
+
+    DY_blk_ptrs = (
+        DY_ptr + M_in_idx[:, None] * stride_dym + N_block[None, :] * stride_dyn
+    )
+    # Packed W in FWD layout [E, N, K/2]
+    # Tile shape [BLOCK_K, BLOCK_N]: row=K_byte (each byte loaded twice across
+    # adjacent K rows), col=N — note N is the *fast* axis here because
+    # stride_wpn (= K/2) is large, but the K row stride is 1.
+    Wp_blk_ptrs = (
+        Wp_ptr
+        + E_idx * stride_wpe
+        + N_block[None, :] * stride_wpn
+        + K_byte_block[:, None] * stride_wpk
+    )
+    # Scale [E, N, K/32]; row=K_scale_idx (broadcast within MX_BLOCK_SIZE), col=N
+    Ws_blk_ptrs = (
+        Ws_ptr
+        + E_idx * stride_wse
+        + N_block[None, :] * stride_wsn
+        + K_scale_block[:, None] * stride_wsk
+    )
+    B_expert_offset = E_idx * ACTUAL_R
+    B_blk_ptrs = (
+        B_ptr
+        + N_block[:, None] * stride_bn
+        + (B_expert_offset + R_block)[None, :] * stride_br
+    )
+
+    iters = tl.cdiv(N, BLOCK_N)
+    dy_b_acc = tl.zeros((BLOCK_M, BLOCK_R), dtype=tl.float32)
+
+    for i in range(iters):
+        if no_n_mask:
+            N_mask_iter = N_block >= 0  # all-true [BLOCK_N]
+        else:
+            N_mask_iter = (i * BLOCK_N + N_block) < N
+        dy_mask = E_mask[:, None] & N_mask_iter[None, :]
+        w_mask = K_mask[:, None] & N_mask_iter[None, :]
+        b_mask = N_mask_iter[:, None] & R_mask[None, :]
+
+        dy = tl.load(DY_blk_ptrs, mask=dy_mask, other=0.0).to(INPUT_DTYPE)
+
+        # MXFP4 dequant of W tile [BLOCK_K, BLOCK_N]
+        packed = tl.load(Wp_blk_ptrs, mask=w_mask, other=0).to(tl.int32)
+        nibble = tl.where(K_is_high[:, None], (packed >> 4) & 0xF, packed & 0xF)
+        codebook_val = tl.load(Codebook_ptr + nibble)
+        if SCALE_LINEAR:
+            scale_fp = tl.load(Ws_blk_ptrs, mask=w_mask, other=0.0).to(tl.float32)
+        else:
+            scale_byte = tl.load(Ws_blk_ptrs, mask=w_mask, other=0).to(tl.int32)
+            scale_fp = tl.exp2((scale_byte - 127).to(tl.float32))
+        w_dq = (codebook_val * scale_fp).to(INPUT_DTYPE)  # [BLOCK_K, BLOCK_N]
+
+        # Base: acc += DY @ W^T  ([M, N] @ [N, K] -> [M, K])
+        # W tile is [BLOCK_K, BLOCK_N]; W^T = tl.trans(w_dq) -> [BLOCK_N, BLOCK_K]
+        acc += tl.dot(dy, tl.trans(w_dq), allow_tf32=allow_tf32).to(tl.float32)
+
+        # LoRA: dy_b_acc += DY @ B
+        b = tl.load(B_blk_ptrs, mask=b_mask, other=0.0).to(INPUT_DTYPE)
+        dy_b_acc += tl.dot(dy, b, allow_tf32=allow_tf32).to(tl.float32)
+
+        DY_blk_ptrs += BLOCK_N * stride_dyn
+        Wp_blk_ptrs += BLOCK_N * stride_wpn
+        Ws_blk_ptrs += BLOCK_N * stride_wsn
+        B_blk_ptrs += BLOCK_N * stride_bn
+
+    # Epilogue: (DY @ B) @ A  ([M, R] @ [R, K] -> [M, K])
+    A_expert_offset = E_idx * ACTUAL_R
+    A_blk_ptrs = (
+        A_ptr
+        + (A_expert_offset + R_block)[:, None] * stride_ar
+        + K_block[None, :] * stride_ak
+    )
+    a_e = tl.load(A_blk_ptrs, mask=R_mask[:, None] & K_mask[None, :], other=0.0).to(
+        INPUT_DTYPE
+    )
+    lora_dx = tl.dot(dy_b_acc.to(INPUT_DTYPE), a_e, allow_tf32=allow_tf32)
+    acc += scaling * lora_dx
+    return acc
+
+
+def _scatter2scatter_lora_dX_mx_configs():
+    """dX MX kernel configs. BLOCK_K must be a multiple of MX_BLOCK_SIZE
+    because scales broadcast within MX_BLOCK_SIZE-element K-blocks.
+
+    The grid is deliberately conservative: the largest configs (BLOCK_M=128 /
+    BLOCK_K=128 / BLOCK_N=64 / num_stages=5) only fit on big-SMEM GPUs (sm90's
+    228 KB) and at least one of them issues an out-of-bounds access in the
+    grouped (expert-parallel) backward there — while the SMEM *estimate* used to
+    prune is too loose to exclude it reliably. Restricting the search space to
+    the validated low-resource configs keeps the kernel correct on every arch;
+    the dX (backward) path is far less autotune-sensitive than the forward.
+    ``AXOLOTL_MX_DX_SAFE_CONFIG`` pins a single config as an escape hatch."""
+    import os
+
+    if os.environ.get("AXOLOTL_MX_DX_SAFE_CONFIG"):
+        return [
+            triton.Config(
+                {"BLOCK_M": 64, "BLOCK_K": 32, "BLOCK_N": 32},
+                num_stages=3,
+                num_warps=4,
+            )
+        ]
+    configs = []
+    for block_m, block_k, block_n, warps, stages in product(
+        [32, 64],
+        [32, 64],  # all multiples of MX_BLOCK_SIZE=32
+        [32],
+        [4, 8],
+        [2, 3],
+    ):
+        configs.append(
+            triton.Config(
+                {"BLOCK_M": block_m, "BLOCK_K": block_k, "BLOCK_N": block_n},
+                num_stages=stages,
+                num_warps=warps,
+            )
+        )
+    return configs
+
+
+# dX MX configs needing more shared memory than this are excluded from autotune even on
+# GPUs that physically have more (e.g. sm90's 228 KB). The larger configs are only reachable
+# on those GPUs, were never exercised on the 100 KB-class GPUs this kernel was validated on,
+# and at least one of them issues an out-of-bounds access on sm90 in the grouped (EP) layout.
+# Capping to the validated budget keeps autotune among known-good configs across all arches.
+_DX_MX_VALIDATED_SMEM = 101_376  # 99 KB — sm120 / Ada / sm89 opt-in budget
+
+
+# Pre-tuned dX-MX subsets for known (arch, shape) profiles — see _FWD_MX_PROFILES for the rationale.
+# spec = (BLOCK_M, BLOCK_K, BLOCK_N, warps, stages).
+_GLM_DX_MX_SUBSET = [
+    (32, 64, 32, 4, 2),  # H200/sm90 microbench winner
+    (64, 64, 32, 4, 2),  # larger BLOCK_M
+    (32, 32, 32, 4, 2),  # smaller BLOCK_K
+    (64, 32, 32, 8, 3),  # fallback
+]
+_DX_MX_PROFILES = {
+    (9, 0): {(4096, 6144): _GLM_DX_MX_SUBSET, (6144, 2048): _GLM_DX_MX_SUBSET},
+}
+
+
+def _profile_dX_mx_configs(configs, meta):
+    """Known GPU arch + (N, K) + large M -> pre-measured dX subset (filtered from ``configs``)."""
+    try:
+        cap = torch.cuda.get_device_capability()
+    except Exception:  # pragma: no cover
+        return None
+    table = _DX_MX_PROFILES.get(cap)
+    if not table:
+        return None
+    specs = table.get((meta.get("N"), meta.get("K")))
+    mb = meta.get("M_BUCKET")
+    if not specs or mb is None or mb < 16384:
+        return None
+    want = set(specs)
+    sel = [
+        c
+        for c in configs
+        if (
+            c.kwargs["BLOCK_M"],
+            c.kwargs["BLOCK_K"],
+            c.kwargs["BLOCK_N"],
+            c.num_warps,
+            c.num_stages,
+        )
+        in want
+    ]
+    return sel or None
+
+
+def _prune_dX_mx_configs(configs, named_args, **kwargs):
+    """Prune dX MX configs by SMEM and register pressure (MX-aware)."""
+    profiled = _profile_dX_mx_configs(configs, kwargs)
+    if profiled is not None:
+        return profiled
+
+    smem_cap = min(_get_smem_capacity(), _DX_MX_VALIDATED_SMEM)
+    block_r = named_args.get("BLOCK_R", 64)
+
+    scored = []
+    for config in configs:
+        block_m = config.kwargs["BLOCK_M"]
+        block_k = config.kwargs["BLOCK_K"]
+        block_n = config.kwargs["BLOCK_N"]
+        if block_k % _MX_BLOCK_SIZE != 0:
+            continue
+        smem_base = _estimate_smem_usage(config.num_stages, block_m, block_k, block_n)
+        # Per stage: packed W tile [BLOCK_K, BLOCK_N] (bytes — each byte
+        # serves two K rows) and scale tile [BLOCK_K, BLOCK_N] (bytes, broadcast
+        # within each K-block of size MX_BLOCK_SIZE). Approximate to BLOCK_K *
+        # BLOCK_N bytes for each (overestimates SMEM slightly, conservative).
+        smem_packed = config.num_stages * block_k * block_n * 1
+        smem_scale = config.num_stages * block_k * block_n * 1
+        smem_lora_loop = config.num_stages * block_n * block_r * 2
+        smem_lora_epilogue = block_r * block_k * 2
+        smem = (
+            smem_base + smem_packed + smem_scale + smem_lora_loop + smem_lora_epilogue
+        )
+
+        est_regs = _estimate_register_pressure(
+            config.num_warps,
+            (block_m, block_k),
+            (block_m, block_r),
+            (block_m, block_n),
+            (block_k, block_n),
+            (block_n, block_r),
+            (block_r, block_k),
+        )
+        if est_regs > _MAX_REGS_SOFT_LIMIT:
+            continue
+
+        scored.append((smem, config))
+
+    pruned = [c for s, c in scored if s <= smem_cap - _SMEM_SLACK]
+    if pruned:
+        return pruned
+    if scored:
+        scored.sort(key=lambda x: x[0])
+        return [scored[0][1]]
+    return [
+        min(
+            configs,
+            key=lambda c: (
+                c.kwargs["BLOCK_M"] * c.kwargs["BLOCK_K"] * c.kwargs["BLOCK_N"]
+            ),
+        )
+    ]
+
+
+@triton.autotune(
+    configs=_scatter2scatter_lora_dX_mx_configs(),
+    key=["M_BUCKET", "N", "K"],
+    prune_configs_by={"early_config_prune": _prune_dX_mx_configs},
+)
+@triton.heuristics(
+    {
+        "NO_N_MASK": lambda args: (args["N"] % args["BLOCK_N"]) == 0,
+    }
+)
+@triton.jit
+def _scatter2scatter_lora_dX_mx(
+    DY_ptr,
+    stride_dym: tl.constexpr,
+    stride_dyn: tl.constexpr,
+    # Packed W in FWD layout [E, N, K/2]
+    Wp_ptr,
+    stride_wpe,
+    stride_wpn: tl.constexpr,
+    stride_wpk: tl.constexpr,
+    # Scale [E, N, K/32]
+    Ws_ptr,
+    stride_wse,
+    stride_wsn: tl.constexpr,
+    stride_wsk: tl.constexpr,
+    Codebook_ptr,
+    DX_ptr,
+    stride_dxm: tl.constexpr,
+    stride_dxk: tl.constexpr,
+    LA_ptr,
+    stride_la_r,
+    stride_la_k,
+    LB_ptr,
+    stride_lb_n,
+    stride_lb_r,
+    grouped_idx_ptr,
+    expert_idxs_ptr,
+    FAN_OUT: tl.constexpr,
+    M,
+    M_BUCKET,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    E: tl.constexpr,
+    ACTUAL_R: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    MX_BLOCK_SIZE: tl.constexpr,
+    SCALE_LINEAR: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    scaling,
+    allow_tf32: tl.constexpr,
+    dy_grouped: tl.constexpr,
+    dx_grouped: tl.constexpr,
+    NO_N_MASK: tl.constexpr,
+    INT64_INDICES: tl.constexpr = False,
+):
+    """Fused MXFP4/NVFP4 dX kernel."""
+    pid = tl.program_id(axis=0)
+    K_BLOCK_COUNT = tl.cdiv(K, BLOCK_K)
+    M_block_id = pid // K_BLOCK_COUNT
+    K_block_id = pid % K_BLOCK_COUNT
+
+    M_block = M_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    if INT64_INDICES:
+        M_block = M_block.to(tl.int64)
+    K_block = K_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
+    K_mask = K_block < K
+    M_boundary_mask = M_block < (FAN_OUT * M)
+
+    E_idxs = tl.load(expert_idxs_ptr + M_block, mask=M_boundary_mask, other=E)
+    no_n_mask = NO_N_MASK
+    acc = tl.zeros((BLOCK_M, BLOCK_K), dtype=ACC_TYPE)
+
+    E_first_idx = tl.min(E_idxs)
+    E_last_idx = tl.minimum(tl.max(E_idxs), E - 1)
+    if INT64_INDICES:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int64)
+    else:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int32)
+
+    for E_idx in range(E_first_idx, E_last_idx + 1):
+        E_mask = E_idxs == E_idx
+        if dy_grouped:
+            M_in_idx = M_block
+        else:
+            M_in_idx = M_idx // FAN_OUT
+
+        acc = _compute_expert_block_lora_dX_mxfp4(
+            E_idx,
+            E_mask,
+            M_in_idx,
+            K_block,
+            K_mask,
+            DY_ptr,
+            stride_dym,
+            stride_dyn,
+            Wp_ptr,
+            stride_wpe,
+            stride_wpn,
+            stride_wpk,
+            Ws_ptr,
+            stride_wse,
+            stride_wsn,
+            stride_wsk,
+            Codebook_ptr,
+            LA_ptr,
+            stride_la_r,
+            stride_la_k,
+            LB_ptr,
+            stride_lb_n,
+            stride_lb_r,
+            N,
+            ACTUAL_R,
+            acc,
+            no_n_mask,
+            BLOCK_M,
+            BLOCK_N,
+            BLOCK_K,
+            BLOCK_R,
+            MX_BLOCK_SIZE,
+            SCALE_LINEAR,
+            scaling,
+            allow_tf32=allow_tf32,
+        )
+
+    if dx_grouped:
+        M_out_idx = M_block
+    else:
+        M_out_idx = M_idx
+    DX_blk_ptrs = DX_ptr + (
+        M_out_idx[:, None] * stride_dxm + K_block[None, :] * stride_dxk
+    )
+    tl.store(DX_blk_ptrs, acc, mask=M_boundary_mask[:, None] & K_mask[None, :])
+
+
+def scatter2scatter_lora_dX_mx(
+    DY: torch.Tensor,
+    W_mx,  # MXWeights in FWD layout (same buffer as forward)
+    sorted_expert_idxs: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    k: int,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    scaling: float,
+    dy_grouped: bool = True,
+    dx_grouped: bool = False,
+    out: Optional[torch.Tensor] = None,
+    int64_indices: bool = False,
+) -> torch.Tensor:
+    """Backward-dX dispatcher for the fused MXFP4 kernel.
+
+    Reuses the *forward* MX layout (block axis = K). The kernel iterates the
+    N reduction in tiles, and for each (K_tile, N_tile) sub-tile, decodes
+    nibbles along the K rows of the tile (K_byte = K // 2) and broadcasts
+    scales within each ``MX_BLOCK_SIZE`` K-block. This avoids the
+    dequant + re-quantize "pre-transpose" round-trip and the extra MX
+    rounding error that would have introduced.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+        MXLayout,
+        fp4_codebook,
+    )
+
+    assert W_mx.layout == MXLayout.FWD, (
+        f"scatter2scatter_lora_dX_mx requires FWD layout, got {W_mx.layout}"
+    )
+    assert sorted_scattered_idxs.size(0) == sorted_expert_idxs.size(0)
+
+    K = W_mx.K
+    N = W_mx.N
+    E = W_mx.packed.size(0)
+    R = lora_A.size(0) // E
+    BLOCK_R = _block_r_for_rank(R)
+    L_scattered = sorted_expert_idxs.size(0)
+
+    if dy_grouped:
+        M = DY.size(0)
+        fan_out = 1
+    else:
+        M = DY.size(0)
+        fan_out = k
+
+    if out is None:
+        output = torch.empty((L_scattered, K), device=DY.device, dtype=DY.dtype)
+    else:
+        assert out.size(0) == L_scattered and out.size(1) == K
+        output = out
+
+    def grid(META):
+        return (
+            triton.cdiv(L_scattered, META["BLOCK_M"]) * triton.cdiv(K, META["BLOCK_K"]),
+        )
+
+    codebook = fp4_codebook(DY.device)
+
+    _scatter2scatter_lora_dX_mx[grid](
+        DY,
+        DY.stride(0),
+        DY.stride(1),
+        W_mx.packed,
+        W_mx.packed.stride(0),
+        W_mx.packed.stride(1),
+        W_mx.packed.stride(2),
+        W_mx.scales,
+        W_mx.scales.stride(0),
+        W_mx.scales.stride(1),
+        W_mx.scales.stride(2),
+        codebook,
+        output,
+        output.stride(0),
+        output.stride(1),
+        lora_A,
+        lora_A.stride(0),
+        lora_A.stride(1),
+        lora_B,
+        lora_B.stride(0),
+        lora_B.stride(1),
+        sorted_scattered_idxs,
+        sorted_expert_idxs,
+        FAN_OUT=fan_out,
+        M=M,
+        M_BUCKET=_bucket_m(M),
+        K=K,
+        N=N,
+        E=E,
+        ACTUAL_R=R,
+        BLOCK_R=BLOCK_R,
+        MX_BLOCK_SIZE=W_mx.block_size,
+        SCALE_LINEAR=W_mx.scale_is_linear,
+        ACC_TYPE=tl.float32,
+        scaling=scaling,
+        allow_tf32=ALLOW_TF32,
+        dy_grouped=dy_grouped,
+        dx_grouped=dx_grouped,
+        INT64_INDICES=int64_indices,
+    )
+    return output
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/ops.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/ops.py
new file mode 100644
index 0000000000..db919d45e2
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/ops.py
@@ -0,0 +1,680 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/shawntan/scattermoe
+# Copyright (c) Shawn Tan and ScatterMoE Contributors
+# Licensed under the Apache License, Version 2.0
+# See https://github.com/shawntan/scattermoe/blob/main/LICENSE
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+BLOCK_M = 128
+ALLOW_TF32 = True
+
+
+@triton.jit
+def _compute_expert_block(
+    E_idx,
+    E_mask,
+    M_in_idx,
+    N_block,
+    N_mask,
+    X_ptr,
+    stride_xm,
+    stride_xk,
+    W_ptr,
+    stride_we,
+    stride_wk,
+    stride_wn,
+    K,
+    acc,
+    no_k_mask,
+    BLOCK_K,
+    allow_tf32=True,
+):
+    K_block = tl.arange(0, BLOCK_K)
+    X_blk_ptrs = X_ptr + M_in_idx[:, None] * stride_xm + K_block[None, :] * stride_xk
+    W_blk_ptrs = (
+        W_ptr
+        + K_block[:, None] * stride_wk
+        + N_block[None, :] * stride_wn
+        + E_idx * stride_we
+    )
+    iters = tl.cdiv(K, BLOCK_K)
+
+    for K_block_id in range(iters):
+        if no_k_mask:
+            x = tl.load(X_blk_ptrs, mask=E_mask[:, None])
+            w = tl.load(W_blk_ptrs, mask=N_mask[None, :])
+        else:
+            K_mask = (K_block_id * BLOCK_K + K_block) < K
+            x = tl.load(X_blk_ptrs, mask=E_mask[:, None] & K_mask[None, :])
+            w = tl.load(W_blk_ptrs, mask=K_mask[:, None] & N_mask[None, :])
+
+        X_blk_ptrs += BLOCK_K * stride_xk
+        W_blk_ptrs += BLOCK_K * stride_wk
+        acc = tl.dot(x, w, acc, allow_tf32=allow_tf32)
+    return acc
+
+
+def _scatter2scatter_configs():
+    return [
+        triton.Config({"BLOCK_N": 128, "BLOCK_K": 32}, num_stages=4, num_warps=4),
+    ]
+
+
+@triton.autotune(
+    configs=_scatter2scatter_configs(),
+    key=["M", "N", "K"],
+)
+@triton.heuristics(
+    {
+        "NO_K_MASK": lambda args: (args["K"] % args["BLOCK_K"]) == 0,
+        "NO_N_MASK": lambda args: (args["N"] % args["BLOCK_N"]) == 0,
+    }
+)
+@triton.jit
+def _scatter2scatter(
+    X_ptr,
+    stride_xm: tl.constexpr,
+    stride_xk: tl.constexpr,
+    W_ptr,
+    stride_we,
+    stride_wk: tl.constexpr,
+    stride_wn: tl.constexpr,
+    Y_ptr,
+    stride_ym: tl.constexpr,
+    stride_yn: tl.constexpr,
+    B_ptr,
+    stride_be: tl.constexpr,
+    stride_bn: tl.constexpr,
+    R_ptr,
+    stride_rm,
+    stride_rn,
+    grouped_idx_ptr,
+    expert_idxs_ptr,
+    # block_start_idx_ptr,
+    FAN_OUT: tl.constexpr,
+    M,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    E: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    # OUT_M,
+    allow_tf32: tl.constexpr,
+    x_grouped: tl.constexpr,
+    y_grouped: tl.constexpr,
+    NO_K_MASK: tl.constexpr,
+    NO_N_MASK: tl.constexpr,
+    INT64_INDICES: tl.constexpr = False,
+):
+    pid = tl.program_id(axis=0)
+
+    N_BLOCK_COUNT = tl.cdiv(N, BLOCK_N)
+    M_block_id = pid // N_BLOCK_COUNT
+    N_block_id = pid % N_BLOCK_COUNT
+
+    M_block = M_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    if INT64_INDICES:
+        M_block = M_block.to(tl.int64)
+    N_block = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+    N_mask = N_block < N
+    M_boundary_mask = M_block < (FAN_OUT * M)
+    E_idxs = tl.load(expert_idxs_ptr + M_block, mask=M_boundary_mask, other=E)
+
+    no_k_mask = K % BLOCK_K == 0
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
+    E_first_idx = tl.min(E_idxs)
+    E_last_idx = tl.minimum(tl.max(E_idxs), E - 1)
+    if INT64_INDICES:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int64)
+    else:
+        M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int32)
+    for E_idx in range(E_first_idx, E_last_idx + 1):
+        E_mask = E_idxs == E_idx
+        E_M_idx = M_idx
+        if x_grouped:
+            M_in_idx = M_block
+        else:
+            M_in_idx = E_M_idx // FAN_OUT
+        acc = _compute_expert_block(
+            E_idx,
+            E_mask,
+            M_in_idx,
+            N_block,
+            N_mask,
+            X_ptr,
+            stride_xm,
+            stride_xk,
+            W_ptr,
+            stride_we,
+            stride_wk,
+            stride_wn,
+            K,
+            acc,
+            no_k_mask,
+            BLOCK_K,
+            allow_tf32=allow_tf32,
+        )
+
+    if B_ptr is not None:
+        B_blk_ptrs = B_ptr + E_idxs[:, None] * stride_be + N_block[None, :] * stride_bn
+        acc += tl.load(B_blk_ptrs, mask=M_boundary_mask[:, None] & N_mask[None, :])
+
+    if y_grouped:
+        M_out_idx = M_block
+    else:
+        M_out_idx = M_idx
+    out_mask = M_boundary_mask[:, None] & N_mask[None, :]
+    if R_ptr is not None:
+        # Fused per-row residual add (e.g. base expert GEMM output) so the LoRA-B GEMM writes
+        # (base + lora) directly, no separate add pass.
+        R_blk_ptrs = R_ptr + (
+            M_out_idx[:, None] * stride_rm + N_block[None, :] * stride_rn
+        )
+        acc += tl.load(R_blk_ptrs, mask=out_mask).to(ACC_TYPE)
+    Y_blk_ptrs = Y_ptr + (M_out_idx[:, None] * stride_ym + N_block[None, :] * stride_yn)
+    tl.store(Y_blk_ptrs, acc, mask=out_mask)
+
+
+def scatter2scatter(
+    X,
+    W,
+    sorted_expert_idxs,
+    sorted_scattered_idxs,
+    k,
+    b=None,
+    x_grouped=False,
+    y_grouped=False,
+    out=None,
+    int64_indices=False,
+    residual=None,
+):
+    assert sorted_scattered_idxs.size(0) == sorted_expert_idxs.size(0)
+    assert sorted_scattered_idxs.size(0) == X.size(0) * k
+    # Pre-kernel setup
+    y_dim = W.size(-1)
+    L_scattered = sorted_expert_idxs.size(0)
+    if out is None:
+        output = torch.empty((L_scattered, y_dim), device=X.device, dtype=X.dtype)
+    else:
+        assert out.size(0) == L_scattered and out.size(1) == y_dim
+        output = out
+    if residual is not None:
+        assert residual.size(0) == L_scattered and residual.size(1) == y_dim
+
+    scatter2scatter_compileable(
+        output,
+        W,
+        X,
+        k,
+        sorted_expert_idxs,
+        sorted_scattered_idxs,
+        b,
+        x_grouped,
+        y_grouped,
+        int64_indices,
+        residual,
+    )
+    return output
+
+
+@torch.library.custom_op("scattermoe::scatter2scatter", mutates_args={"output"})
+def scatter2scatter_compileable(
+    output: torch.Tensor,
+    W: torch.Tensor,
+    X: torch.Tensor,
+    k: int,
+    sorted_expert_idxs: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    b: Optional[torch.Tensor],
+    x_grouped: bool,
+    y_grouped: bool,
+    int64_indices: bool = False,
+    residual: Optional[torch.Tensor] = None,
+) -> None:
+    def grid(META):
+        grid_num = (
+            triton.cdiv(sorted_expert_idxs.size(0), META["BLOCK_M"])
+            * triton.cdiv(META["N"], META["BLOCK_N"]),
+        )
+        return grid_num
+
+    if b is None:
+        b = None
+        stride_be = stride_bn = 0
+    else:
+        stride_be, stride_bn = b.stride()
+
+    if residual is None:
+        stride_rm = stride_rn = 0
+    else:
+        stride_rm, stride_rn = residual.stride()
+
+    _scatter2scatter[grid](
+        # X_ptr, stride_xm, stride_xk,
+        X,
+        X.stride(0),
+        X.stride(1),
+        # W_ptr, stride_we, stride_wk, stride_wn,
+        W,
+        W.stride(0),
+        W.stride(1),
+        W.stride(2),
+        # Y_ptr, stride_ym, stride_yn,
+        output,
+        output.stride(0),
+        output.stride(1),
+        # B_ptr, stride_be, stride_bn
+        b,
+        stride_be,
+        stride_bn,
+        # R_ptr, stride_rm, stride_rn (per-row residual added in epilogue)
+        residual,
+        stride_rm,
+        stride_rn,
+        grouped_idx_ptr=sorted_scattered_idxs,
+        expert_idxs_ptr=sorted_expert_idxs,
+        # block_start_idx_ptr=padded_block_idxs,
+        FAN_OUT=k,
+        M=X.size(0),
+        K=X.size(1),
+        N=output.size(1),
+        E=W.size(0),
+        BLOCK_M=BLOCK_M,
+        ACC_TYPE=tl.float32,
+        allow_tf32=ALLOW_TF32,
+        x_grouped=x_grouped,
+        y_grouped=y_grouped,
+        INT64_INDICES=int64_indices,
+    )
+
+
+def _config_XtY():
+    return [
+        triton.Config(
+            {"BLOCK_N": 128, "BLOCK_K": 128, "BLOCK_M": 32}, num_stages=4, num_warps=4
+        ),
+    ]
+
+
+def group_bwd_W(DY, X, expert_offsets, E, has_bias=False):
+    DWt = torch.zeros((E, DY.size(-1), X.size(-1)), device=DY.device, dtype=DY.dtype)
+    DW = DWt.permute(0, 2, 1)
+    if has_bias:
+        Db = torch.zeros((E, DY.size(-1)), device=DY.device, dtype=DY.dtype)
+    else:
+        Db = None
+    groupXtY_compileable(E, DW, Db, DY, X, expert_offsets)
+    return DW, Db
+
+
+@torch.library.custom_op("scattermoe::groupXtY", mutates_args={"DW", "Db"})
+def groupXtY_compileable(
+    E: int,
+    DW: torch.Tensor,
+    Db: Optional[torch.Tensor],
+    DY: torch.Tensor,
+    X: torch.Tensor,
+    expert_offsets: torch.Tensor,
+) -> None:
+    def grid(META):
+        grid = (
+            E * triton.cdiv(META["K"], META["BLOCK_K"]),
+            triton.cdiv(META["N"], META["BLOCK_N"]),
+        )
+        return grid
+
+    if Db is None:
+        stride_dbe = 0
+        stride_dbn = 0
+    else:
+        stride_dbe, stride_dbn = Db.stride()
+
+    _groupXtY[grid](
+        # DY_ptr, stride_dym, stride_dyk,
+        DY,
+        DY.stride(0),
+        DY.stride(1),
+        # X_ptr, stride_xm, stride_xn,
+        X,
+        X.stride(0),
+        X.stride(1),
+        # DW_ptr, stride_dwe, stride_dwk, stride_dwn,
+        DW,
+        DW.stride(0),
+        DW.stride(1),
+        DW.stride(2),
+        # Db_ptr, stride_dwe, stride_dbn,
+        Db,
+        stride_dbe,
+        stride_dbn,
+        # expert_offsets_ptr,
+        expert_offsets,
+        # K: tl.constexpr, N: tl.constexpr,
+        M=DY.size(0),
+        N=DY.size(-1),
+        K=X.size(-1),
+        # ACC_TYPE: tl.constexpr,
+        ACC_TYPE=tl.float32,
+        allow_tf32=ALLOW_TF32,
+    )
+
+
+@triton.autotune(
+    configs=_config_XtY(),
+    key=["M", "N", "K"],
+)
+@triton.heuristics(
+    {
+        "NO_K_MASK": lambda args: (args["K"] % args["BLOCK_K"]) == 0,
+        "NO_N_MASK": lambda args: (args["N"] % args["BLOCK_N"]) == 0,
+    }
+)
+@triton.jit
+def _groupXtY(
+    DY_ptr,
+    stride_dym,
+    stride_dyk,
+    X_ptr,
+    stride_xm,
+    stride_xn,
+    DW_ptr,
+    stride_dwe,
+    stride_dwk,
+    stride_dwn,
+    Db_ptr,
+    stride_dbe,
+    stride_dbn,
+    expert_offsets_ptr,
+    M,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    allow_tf32: tl.constexpr,
+    NO_K_MASK: tl.constexpr,
+    NO_N_MASK: tl.constexpr,
+):
+    pid0 = tl.program_id(axis=0)
+    pid1 = tl.program_id(axis=1)
+    num0 = tl.num_programs(0)
+    num1 = tl.num_programs(1)
+    # pid1, pid0 = tl.swizzle2d(pid1, pid0, num1, num0, 128)
+    pid0, pid1 = tl.swizzle2d(pid0, pid1, num0, num1, 4)
+
+    K_BLOCK_COUNT = tl.cdiv(K, BLOCK_K)
+    E_idx = pid0 // K_BLOCK_COUNT
+    K_block_id = pid0 % K_BLOCK_COUNT
+    N_block_id = pid1
+
+    if E_idx == 0:
+        start_idx = 0
+    else:
+        start_idx = tl.load(expert_offsets_ptr + E_idx - 1).to(tl.int32)
+    end_idx = tl.load(expert_offsets_ptr + E_idx).to(tl.int32)
+
+    if end_idx > start_idx:
+        M_block = tl.max_contiguous(start_idx + tl.arange(0, BLOCK_M), BLOCK_M)
+
+        K_block = K_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
+        K_mask = K_block < K
+        K_block = tl.max_contiguous(tl.multiple_of(K_block % K, BLOCK_K), BLOCK_K)
+
+        N_block = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+        N_mask = N_block < N
+        N_block = tl.max_contiguous(tl.multiple_of(N_block % N, BLOCK_N), BLOCK_N)
+
+        M_idxs = M_block
+        xt_blk_ptrs = X_ptr + K_block[:, None] * stride_xn + M_idxs[None, :] * stride_xm
+        dy_blk_ptrs = (
+            DY_ptr + M_idxs[:, None] * stride_dym + N_block[None, :] * stride_dyk
+        )
+        if (Db_ptr is not None) and (K_block_id == 0):
+            _xty_and_bias(
+                E_idx,
+                start_idx,
+                end_idx,
+                M_block,
+                K_block,
+                K_mask,
+                N_block,
+                N_mask,
+                dy_blk_ptrs,
+                stride_dym,
+                xt_blk_ptrs,
+                stride_xm,
+                DW_ptr,
+                stride_dwe,
+                stride_dwk,
+                stride_dwn,
+                Db_ptr,
+                stride_dbe,
+                stride_dbn,
+                BLOCK_M,
+                BLOCK_N,
+                BLOCK_K,
+                ACC_TYPE,
+                allow_tf32,
+                NO_K_MASK,
+                NO_N_MASK,
+                compute_bias=True,
+            )
+        else:
+            _xty_and_bias(
+                E_idx,
+                start_idx,
+                end_idx,
+                M_block,
+                K_block,
+                K_mask,
+                N_block,
+                N_mask,
+                dy_blk_ptrs,
+                stride_dym,
+                xt_blk_ptrs,
+                stride_xm,
+                DW_ptr,
+                stride_dwe,
+                stride_dwk,
+                stride_dwn,
+                Db_ptr,
+                stride_dbe,
+                stride_dbn,
+                BLOCK_M,
+                BLOCK_N,
+                BLOCK_K,
+                ACC_TYPE,
+                allow_tf32,
+                NO_K_MASK,
+                NO_N_MASK,
+                compute_bias=False,
+            )
+
+
+@triton.jit
+def _xty_and_bias(
+    E_idx,
+    start_idx,
+    end_idx,
+    M_block,
+    K_block,
+    K_mask,
+    N_block,
+    N_mask,
+    dy_blk_ptrs,
+    stride_dym,
+    xt_blk_ptrs,
+    stride_xm,
+    DW_ptr,
+    stride_dwe,
+    stride_dwk,
+    stride_dwn,
+    Db_ptr,
+    stride_dbe,
+    stride_dbn,
+    BLOCK_M,
+    BLOCK_N,
+    BLOCK_K,
+    ACC_TYPE,
+    allow_tf32,
+    NO_K_MASK,
+    NO_N_MASK,
+    compute_bias: tl.constexpr,
+):
+    if compute_bias:
+        db_acc = tl.zeros((BLOCK_N,), dtype=ACC_TYPE)
+    else:
+        db_acc = None
+
+    acc = tl.zeros((BLOCK_K, BLOCK_N), dtype=ACC_TYPE)
+    iters = tl.cdiv(end_idx - start_idx, BLOCK_M)
+    for i in range(0, iters):
+        M_mask = (i * BLOCK_M + M_block) < end_idx
+        if NO_K_MASK:
+            xt = tl.load(xt_blk_ptrs, mask=M_mask[None, :])
+        else:
+            xt = tl.load(xt_blk_ptrs, mask=K_mask[:, None] & M_mask[None, :])
+        if NO_N_MASK:
+            dy = tl.load(dy_blk_ptrs, mask=M_mask[:, None])
+        else:
+            dy = tl.load(dy_blk_ptrs, mask=M_mask[:, None] & N_mask[None, :])
+
+        acc += tl.dot(xt, dy, out_dtype=ACC_TYPE, allow_tf32=allow_tf32)
+
+        xt_blk_ptrs += BLOCK_M * stride_xm
+        dy_blk_ptrs += BLOCK_M * stride_dym
+
+        if compute_bias:
+            db_acc += tl.sum(dy, axis=0)
+
+    DW_blk_ptrs = (
+        DW_ptr
+        + E_idx * stride_dwe
+        + K_block[:, None] * stride_dwk
+        + N_block[None, :] * stride_dwn
+    )
+    acc = acc.to(DW_blk_ptrs.dtype.element_ty)
+    tl.store(DW_blk_ptrs, acc, mask=K_mask[:, None] & N_mask[None, :])
+    if compute_bias:
+        Db_blk_ptrs = Db_ptr + E_idx * stride_dbe + N_block * stride_dbn
+        tl.store(Db_blk_ptrs, db_acc, mask=N_mask)
+
+
+def _config_grouping():
+    return [
+        triton.Config({"BLOCK_N": 256, "BLOCK_K": 128}, num_stages=4, num_warps=4),
+        # triton.Config({'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=4, num_warps=4),
+        # triton.Config({'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
+    ]
+
+
+def group(A, sorted_expert_idxs, coeff=None, fan_out=1, out=None):
+    N = sorted_expert_idxs.size(0)
+    K = A.size(1)
+    assert A.size(0) * fan_out == N
+    if out is not None:
+        Y = out
+    else:
+        Y = torch.empty((N, K), dtype=A.dtype, device=A.device)
+    group_compileable(A, K, N, Y, coeff, coeff is not None, fan_out, sorted_expert_idxs)
+    return Y
+
+
+@torch.library.custom_op("scattermoe::group", mutates_args={"Y"})
+def group_compileable(
+    A: torch.Tensor,
+    K: int,
+    N: int,
+    Y: torch.Tensor,
+    coeff: Optional[torch.Tensor],
+    has_coeff: bool,
+    fan_out: int,
+    sorted_expert_idxs: torch.Tensor,
+) -> None:
+    def grid(META):
+        grid_num = (triton.cdiv(META["N"], META["BLOCK_N"]),)
+        return grid_num
+
+    _group[grid](
+        # A_ptr, stride_an, stride_ai,
+        A,
+        A.stride(0),
+        A.stride(1),
+        has_coeff,
+        coeff,
+        fan_out,
+        # Y_ptr, stride_yn, stride_yk,
+        Y,
+        Y.stride(0),
+        Y.stride(1),
+        # grouped_idx_ptr,
+        sorted_expert_idxs,
+        # N: tl.constexpr, K: tl.constexpr,
+        N,
+        K,
+    )
+
+
+@triton.autotune(configs=_config_grouping(), key=["K"])
+@triton.heuristics({"NO_K_MASK": lambda args: (args["K"] % args["BLOCK_K"]) == 0})
+@triton.jit
+def _group(
+    src_ptr,
+    stride_sn,
+    stride_sk,
+    has_coeff: tl.constexpr,
+    coeff_ptr,
+    FAN_OUT: tl.constexpr,
+    tgt_ptr,
+    stride_tn,
+    stride_ti,
+    grouped_idx_ptr,
+    N,
+    K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    NO_K_MASK: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+
+    N_block_id = pid
+    N_blk = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+    N_mask = N_blk < N
+    N_blk = tl.max_contiguous(tl.multiple_of(N_blk % N, BLOCK_N), BLOCK_N)
+    N_idx = tl.load(grouped_idx_ptr + N_blk, mask=N_mask, other=0)
+
+    K_blk = tl.arange(0, BLOCK_K)
+    src_blk_ptrs = (
+        src_ptr + (N_idx // FAN_OUT)[:, None] * stride_sn + K_blk[None, :] * stride_sk
+    )
+    tgt_blk_ptrs = tgt_ptr + N_blk[:, None] * stride_tn + K_blk[None, :] * stride_ti
+
+    if has_coeff:
+        c = tl.load(coeff_ptr + N_idx, mask=N_mask)[:, None]
+
+    iters = tl.cdiv(K, BLOCK_K)
+    for i in range(0, iters):
+        if NO_K_MASK or i < iters - 1:
+            block = tl.load(src_blk_ptrs, mask=N_mask[:, None])
+            if has_coeff:
+                block *= c
+            tl.store(tgt_blk_ptrs, block, mask=N_mask[:, None])
+
+        else:
+            K_mask = (i * BLOCK_K + K_blk) < K
+            mask = N_mask[:, None] & K_mask[None, :]
+            block = tl.load(src_blk_ptrs, mask=mask)
+            if has_coeff:
+                block *= c
+            tl.store(tgt_blk_ptrs, block, mask=mask)
+        src_blk_ptrs += BLOCK_K * stride_sk
+        tgt_blk_ptrs += BLOCK_K * stride_ti
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/single.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/single.py
new file mode 100644
index 0000000000..9f0270aa67
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/kernels/single.py
@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/shawntan/scattermoe
+# Copyright (c) Shawn Tan and ScatterMoE Contributors
+# Licensed under the Apache License, Version 2.0
+# See https://github.com/shawntan/scattermoe/blob/main/LICENSE
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _single2scatter(
+    X_ptr,
+    stride_xm,
+    stride_xk,
+    W_ptr,
+    stride_we,
+    stride_wk,
+    stride_wn,
+    Y_ptr,
+    stride_ym,
+    stride_yn,
+    expert_idxs_ptr,
+    FAN_OUT: tl.constexpr,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    E: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+):
+    pid0 = tl.program_id(axis=0)
+    pid1 = tl.program_id(axis=1)
+
+    N_block_id = pid0
+    if FAN_OUT == 1:
+        in_idx = pid1
+    else:
+        in_idx = 0
+    out_idx = pid1
+
+    K_block = tl.arange(0, BLOCK_K)
+    N_block = tl.max_contiguous(
+        tl.multiple_of((N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)) % N, BLOCK_N),
+        BLOCK_N,
+    )
+    E_idx = tl.load(expert_idxs_ptr + pid1)
+    X_blk_ptrs = X_ptr + in_idx * stride_xm + K_block[:, None] * stride_xk
+    W_blk_ptrs = (
+        W_ptr
+        + E_idx * stride_we
+        + K_block[:, None] * stride_wk
+        + N_block[None, :] * stride_wn
+    )
+    N_mask = N_block < N
+    acc = tl.zeros((1, BLOCK_N), dtype=ACC_TYPE)
+    for _K_block_id in range(0, tl.cdiv(K, BLOCK_K)):
+        K_mask = K_block < K
+        x = tl.load(X_blk_ptrs, mask=K_mask[:, None], other=0.0)
+        w = tl.load(W_blk_ptrs, mask=K_mask[:, None] & N_mask[None, :], other=0.0)
+        acc += tl.sum(x * w, axis=0)[None, :]
+        X_blk_ptrs += BLOCK_K * stride_xk
+        W_blk_ptrs += BLOCK_K * stride_wk
+        K_block += BLOCK_K
+    Y_blk_ptrs = Y_ptr + out_idx * stride_ym + N_block[None, :] * stride_yn
+    tl.store(Y_blk_ptrs, acc, mask=N_mask[None, :])
+
+
+def single2scatter(X, W, expert_idxs):
+    E, xdim, ydim = W.size()
+    k = expert_idxs.size(1)
+    assert X.size(0) == k or X.size(0) == 1
+    Y = torch.empty((k, ydim), device=X.device, dtype=X.dtype)
+    BLOCK_N = 128
+    BLOCK_K = 128
+    grid = triton.cdiv(ydim, BLOCK_N), k
+    _single2scatter[grid](
+        X,
+        X.stride(0),
+        X.stride(1),
+        W,
+        W.stride(0),
+        W.stride(1),
+        W.stride(2),
+        Y,
+        Y.stride(0),
+        Y.stride(1),
+        expert_idxs,
+        FAN_OUT=Y.size(0) // X.size(0),
+        K=xdim,
+        N=ydim,
+        E=E,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        ACC_TYPE=tl.float32,
+    )
+    return Y
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/layers.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/layers.py
new file mode 100644
index 0000000000..de3f57af62
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/layers.py
@@ -0,0 +1,617 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Original work Copyright (c) Shawn Tan and ScatterMoE Contributors
+# Adapted from https://github.com/shawntan/scattermoe
+# See https://github.com/shawntan/scattermoe/blob/main/LICENSE
+#
+# Modifications and LoRA adaptation Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+ScatterMoE layer replacements for HuggingFace MoE architectures.
+
+Provides drop-in forward replacements that use ScatterMoE kernels for
+acceleration. When used via the HF ``kernels`` library
+(``replace_kernel_forward_from_hub``), these classes replace the forward
+method of the original MoE block.
+
+LoRA support
+------------
+When peft wraps parameters via ``target_parameters``, the ``self.experts``
+submodule becomes a chain of ``ParamWrapper`` objects and the ``self.gate``
+router may also become a ``ParamWrapper``.  The ``HFScatterMoEGatedMLP``
+forward detects this and automatically:
+
+1. Unwraps ``self.gate`` to the base router, applying gate LoRA delta
+2. Unwraps ``self.experts`` to the base ``OlmoeExperts`` module
+3. Extracts LoRA A/B weights and scaling from each wrapper
+4. Converts B layout from peft rank-major to scattermoe expert-major
+5. Routes to ``parallel_linear_lora`` for fused LoRA computation
+6. Passes through ``self.shared_expert`` / ``self.shared_expert_gate``
+   (peft wraps their linear layers with standard LoRA, no special handling)
+"""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .parallel_experts import flatten_sort_count, parallel_linear
+from .parallel_linear_lora import get_lora_params_from_wrapper, parallel_linear_lora
+from .selective_dequant import is_mxfp4_param
+
+# =============================================================================
+# LoRA layout conversion utilities (peft <-> scattermoe)
+# =============================================================================
+
+
+def peft_lora_B_to_scattermoe(peft_B, num_experts, rank):
+    """Convert peft rank-major lora_B ``[out, E*r]`` to scattermoe
+    expert-major ``[N, r*E]``.
+
+    peft reshapes B to ``[out, r, E]`` (rank-major).
+    scattermoe slices B as ``[:, e*r:(e+1)*r]`` (expert-major).
+    """
+    N = peft_B.shape[0]
+    return (
+        peft_B.reshape(N, rank, num_experts)
+        .permute(0, 2, 1)
+        .contiguous()
+        .reshape(N, num_experts * rank)
+    )
+
+
+def peft_lora_to_scattermoe(peft_A, peft_B, num_experts, rank):
+    """Convert peft LoRA weights to scattermoe layout.
+
+    peft >=0.19.1 assigns in/out features for 3D params such that
+    A and B already align with scattermoe's convention (no A<->B swap).
+    Only B needs rank-major → expert-major layout conversion.
+    """
+    smoe_A = peft_A
+    smoe_B = peft_lora_B_to_scattermoe(peft_B, num_experts, rank)
+    return smoe_A, smoe_B
+
+
+def peft_down_proj_lora_to_scattermoe(peft_A, peft_B, num_experts, rank):
+    """Deprecated alias for :func:`peft_lora_to_scattermoe`."""
+    return peft_lora_to_scattermoe(peft_A, peft_B, num_experts, rank)
+
+
+# =============================================================================
+# ParamWrapper unwrapping
+# =============================================================================
+
+
+def _unwrap_gate_lora(gate_module):
+    """Unwrap peft ``ParamWrapper`` on the router gate.
+
+    When peft targets ``gate.weight``, ``self.gate`` becomes::
+
+        ParamWrapper(weight)
+          -> base_layer: OlmoeTopKRouter (the real module)
+
+    This function detects the wrapping and returns the base router, its
+    weight tensor, and an optional LoRA delta tensor.
+
+    Returns:
+        (base_gate, gate_weight, gate_lora_delta_or_None)
+
+        ``base_gate`` is the original router module (with ``.top_k``,
+        ``.num_experts``, ``.norm_topk_prob``).
+        ``gate_weight`` is the base router weight (may be a DTensor under FSDP).
+        ``gate_lora_delta_or_None`` is the LoRA delta tensor if LoRA is active,
+        else ``None``.  Kept separate to avoid mixing DTensor + Tensor in an add.
+    """
+    if hasattr(gate_module, "base_layer") and hasattr(gate_module, "lora_A"):
+        base_gate = gate_module.base_layer
+        lora_A, lora_B, scaling = get_lora_params_from_wrapper(gate_module)
+        if lora_A is not None:
+            # gate weight: [num_experts, hidden_size]
+            # lora_A: [r, hidden_size], lora_B: [num_experts, r]
+            # delta = scaling * B @ A = [num_experts, hidden_size]
+            delta = scaling * (lora_B @ lora_A)
+            return base_gate, base_gate.weight, delta
+        else:
+            return base_gate, base_gate.weight, None
+    else:
+        # No wrapping — gate is the original module
+        return gate_module, gate_module.weight, None
+
+
+def _convert_smoe_lora(lora_A, lora_B, num_experts, rank, scaling):
+    """Convert peft LoRA weights to scattermoe layout."""
+    smoe_A, smoe_B = peft_lora_to_scattermoe(lora_A, lora_B, num_experts, rank)
+    return (smoe_A, smoe_B, scaling)
+
+
+def _unwrap_experts_lora(experts_module):
+    """Walk a peft ``ParamWrapper`` chain on ``self.experts``.
+
+    When peft targets ``experts.gate_up_proj`` and ``experts.down_proj`` via
+    ``target_parameters``, ``self.experts`` becomes a nested chain::
+
+        ParamWrapper(down_proj)
+          -> base_layer: ParamWrapper(gate_up_proj)
+              -> base_layer: OlmoeExperts (the real module)
+
+    This function walks the chain, collects LoRA params keyed by
+    ``parameter_name``, and returns the base experts module.
+
+    Returns:
+        (base_experts, gup_lora, down_lora)
+
+        Each ``*_lora`` is either ``(smoe_A, smoe_B, scaling)`` or ``None``.
+        A/B are already in scattermoe layout.
+    """
+    # Collect ParamWrapper layers by their parameter_name
+    wrappers = {}
+    module = experts_module
+    while hasattr(module, "base_layer") and hasattr(module, "lora_A"):
+        param_name = getattr(module, "parameter_name", None)
+        if param_name is not None:
+            wrappers[param_name] = module
+        module = module.base_layer
+
+    base_experts = module
+
+    if not wrappers:
+        return base_experts, None, None
+
+    # Determine num_experts from base module
+    num_experts = getattr(base_experts, "num_experts", None)
+    if num_experts is None:
+        # Fallback: infer from parameter shape
+        gup = getattr(base_experts, "gate_up_proj", None)
+        if gup is not None:
+            num_experts = gup.shape[0]
+
+    # Extract gate_up_proj LoRA (needs A<->B swap due to transposition)
+    gup_lora = None
+    gup_wrapper = wrappers.get("gate_up_proj")
+    if gup_wrapper is not None:
+        lora_A, lora_B, scaling = get_lora_params_from_wrapper(gup_wrapper)
+        if lora_A is not None:
+            rank = lora_A.shape[0] // num_experts
+            gup_lora = _convert_smoe_lora(lora_A, lora_B, num_experts, rank, scaling)
+
+    # Extract down_proj LoRA (needs A<->B swap due to transposition)
+    down_lora = None
+    down_wrapper = wrappers.get("down_proj")
+    if down_wrapper is not None:
+        lora_A, lora_B, scaling = get_lora_params_from_wrapper(down_wrapper)
+        if lora_A is not None:
+            rank = lora_A.shape[0] // num_experts
+            down_lora = _convert_smoe_lora(lora_A, lora_B, num_experts, rank, scaling)
+
+    return base_experts, gup_lora, down_lora
+
+
+# =============================================================================
+# Routing helpers
+# =============================================================================
+
+
+def _softmax_topk_route(
+    moe_block, base_gate, hidden_states, gate_weight, gate_lora_delta
+):
+    """Softmax→topk routing (Qwen, OLMoE, Mixtral, MiniMax).
+
+    Returns:
+        (routing_weights [T, K], selected_experts [T, K], top_k, num_experts)
+    """
+    router_logits = F.linear(hidden_states, gate_weight)
+    if gate_lora_delta is not None:
+        router_logits = router_logits + F.linear(hidden_states, gate_lora_delta)
+    routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+
+    top_k = base_gate.top_k
+    num_experts = base_gate.num_experts
+    routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+    if getattr(base_gate, "norm_topk_prob", True):
+        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+
+    return routing_weights, selected_experts, top_k, num_experts
+
+
+def _sigmoid_topk_route(
+    moe_block, base_gate, hidden_states, gate_weight, gate_lora_delta
+):
+    """Sigmoid→topk routing (GLM, DeepSeek V3, MiniMax M2).
+
+    Supports:
+    - ``e_score_correction_bias`` on gate or moe_block
+    - Group-based expert selection when ``n_group > 1``
+    - ``routed_scaling_factor`` applied to final weights
+    - Final weights gathered from original sigmoid probs (not bias-corrected)
+
+    Returns:
+        (routing_weights [T, K], selected_experts [T, K], top_k, num_experts)
+    """
+    router_logits = F.linear(hidden_states.float(), gate_weight.float())
+    if gate_lora_delta is not None:
+        router_logits = router_logits + F.linear(
+            hidden_states.float(), gate_lora_delta.float()
+        )
+    router_probs = router_logits.sigmoid()  # [T, E]
+
+    top_k = getattr(moe_block, "top_k", getattr(base_gate, "top_k", None))
+    num_experts = getattr(moe_block, "n_routed_experts", gate_weight.shape[0])
+
+    # Bias-corrected scores for expert selection (not used for final weights).
+    # glm_moe_dsa/deepseek_v3 store the bias on gate; minimax_m2 on the block.
+    e_score_correction_bias = getattr(base_gate, "e_score_correction_bias", None)
+    if e_score_correction_bias is None:
+        e_score_correction_bias = getattr(moe_block, "e_score_correction_bias", None)
+    if e_score_correction_bias is not None:
+        scores_for_choice = router_probs + e_score_correction_bias
+    else:
+        scores_for_choice = router_probs
+
+    # Group-based selection: pick top groups, mask the rest
+    n_group = getattr(moe_block, "n_group", 1)
+    if n_group > 1:
+        group_scores = (
+            scores_for_choice.view(-1, n_group, num_experts // n_group)
+            .topk(2, dim=-1)[0]
+            .sum(dim=-1)
+        )  # [T, n_group]
+        topk_group = getattr(moe_block, "topk_group", n_group)
+        group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1)
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(-1, n_group, num_experts // n_group)
+            .reshape(-1, num_experts)
+        )
+        scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0)
+
+    # Final topk from (possibly masked) scores
+    topk_indices = torch.topk(scores_for_choice, k=top_k, dim=-1, sorted=False)[1]
+
+    # Gather weights from original sigmoid scores (not bias-corrected)
+    topk_weights = router_probs.gather(1, topk_indices)
+
+    # Optional renormalization + scaling
+    if getattr(moe_block, "norm_topk_prob", True):
+        topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20)
+    routed_scaling_factor = getattr(moe_block, "routed_scaling_factor", 1.0)
+    topk_weights = topk_weights * routed_scaling_factor
+
+    return topk_weights, topk_indices, top_k, num_experts
+
+
+def _route(moe_block, base_gate, hidden_states, gate_weight, gate_lora_delta):
+    """Dispatch to the correct routing strategy based on block attributes.
+
+    Detects sigmoid routing by the presence of ``e_score_correction_bias``
+    on either the gate or the moe_block.
+    """
+    has_sigmoid = (
+        getattr(base_gate, "e_score_correction_bias", None) is not None
+        or getattr(moe_block, "e_score_correction_bias", None) is not None
+    )
+    if has_sigmoid:
+        return _sigmoid_topk_route(
+            moe_block, base_gate, hidden_states, gate_weight, gate_lora_delta
+        )
+    return _softmax_topk_route(
+        moe_block, base_gate, hidden_states, gate_weight, gate_lora_delta
+    )
+
+
+# =============================================================================
+# Shared expert helpers
+# =============================================================================
+
+
+def _compute_shared_expert(moe_block, hidden_states_flat):
+    """Compute shared expert output if the block has one.
+
+    Handles singular (qwen2_moe: ``shared_expert``), plural
+    (glm_moe_dsa/deepseek_v3: ``shared_experts``), and MLP
+    (hunyuan_v1_moe: ``shared_mlp``) attribute names.
+
+    peft wraps individual linear layers inside the shared expert with
+    standard LoRA — calling forward() handles this transparently.
+    """
+    shared_expert = (
+        getattr(moe_block, "shared_expert", None)
+        or getattr(moe_block, "shared_experts", None)
+        or getattr(moe_block, "shared_mlp", None)
+    )
+    if shared_expert is None:
+        return None
+
+    shared_expert_output = shared_expert(hidden_states_flat)
+
+    # Optional sigmoid gate (Qwen2MoE pattern).
+    # shared_expert_gate may also be peft-wrapped (standard LoRA
+    # on nn.Linear), its forward() applies LoRA automatically.
+    shared_expert_gate = getattr(moe_block, "shared_expert_gate", None)
+    if shared_expert_gate is not None:
+        shared_expert_output = (
+            F.sigmoid(shared_expert_gate(hidden_states_flat)) * shared_expert_output
+        )
+
+    return shared_expert_output
+
+
+# =============================================================================
+# Layer classes
+# =============================================================================
+
+
+class ScatterMoEGatedMLP(nn.Module):
+    def forward(self, layer_input):
+        """
+        Forward pass of the mixture of experts layer.
+
+        Args:
+            layer_input (Tensor):
+                Input tensor.
+
+        Returns:
+            Tensor:
+                Output tensor.
+        """
+        bsz, length, emb_size = layer_input.size()
+        layer_input = layer_input.reshape(-1, emb_size)
+        # compute the top_k routing decision
+        router_logits = self.router.layer(layer_input)
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(
+            routing_weights, self.router.top_k, dim=-1
+        )
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(layer_input.dtype)
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+            selected_experts, num_experts=self.router.num_experts
+        )
+
+        # compute experts
+        gates, h = parallel_linear(
+            layer_input,
+            self.input_linear.weight.transpose(2, 1),
+            self.router.top_k,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=False,
+            grouped_out=True,
+        ).chunk(2, dim=-1)
+        h = self.activation(gates) * h
+        layer_output = parallel_linear(
+            h,
+            self.output_linear.weight.transpose(2, 1),
+            1,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True,
+            grouped_out=False,
+            gates=routing_weights,
+        )
+        layer_output = layer_output.view(bsz, length, emb_size)
+        return layer_output
+
+
+class HFScatterMoEGatedMLP(nn.Module):
+    """
+    ScatterMoE-accelerated forward pass for HF MoEs.
+
+    Used as a kernel layer via the HF ``kernels`` library.  The ``forward``
+    method replaces the original SparseMoeBlock.forward.
+
+    Supports:
+
+    * **Softmax→topk routing**: OLMoE, Qwen2/3MoE, Mixtral, MiniMax
+    * **Sigmoid→topk routing**: GLM, DeepSeek V3, MiniMax M2
+    * **Full-parameter training**: uses ``parallel_linear`` (base ScatterMoE)
+    * **LoRA fine-tuning**: detects peft ``ParamWrapper`` on ``self.experts``,
+      extracts adapter weights, and uses ``parallel_linear_lora`` (fused kernel)
+    """
+
+    @staticmethod
+    def forward(self: nn.Module, layer_input: torch.Tensor):
+        """
+        Forward pass using ScatterMoE kernels.
+
+        Args:
+            self: The MoeSparseMoeBlock module containing:
+                - self.gate: Router (or peft ParamWrapper wrapping it)
+                - self.experts: Experts module (or peft ParamWrapper chain)
+                - self.shared_expert(s): Optional shared expert
+                - self.shared_expert_gate: Optional shared expert gate
+            layer_input: Input tensor [batch_size, seq_len, hidden_size]
+
+        Returns:
+            Tensor: [batch_size, seq_len, hidden_size]
+        """
+        batch_size, sequence_length, hidden_dim = layer_input.shape
+        hidden_states_flat = layer_input.view(-1, hidden_dim)
+
+        # ====================================================================
+        # Shared Expert (if present, e.g. Qwen2MoE, DeepSeek V3)
+        # ====================================================================
+        shared_expert_output = _compute_shared_expert(self, hidden_states_flat)
+
+        # ====================================================================
+        # Router Computation (with optional gate LoRA)
+        # ====================================================================
+        base_gate, gate_weight, gate_lora_delta = _unwrap_gate_lora(self.gate)
+        routing_weights, selected_experts, top_k, num_experts = _route(
+            self, base_gate, hidden_states_flat, gate_weight, gate_lora_delta
+        )
+        routing_weights = routing_weights.to(hidden_states_flat.dtype)
+
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+            selected_experts, num_experts=num_experts
+        )
+
+        # ====================================================================
+        # Detect LoRA (peft ParamWrapper) and extract adapter weights
+        # ====================================================================
+        experts, gup_lora, down_lora = _unwrap_experts_lora(self.experts)
+
+        # ====================================================================
+        # Selective expert weight dequantization
+        # ====================================================================
+        # When experts are BnB-quantized (quantize_moe_experts) or MXFP4
+        # (torchao MXTensor), dequantize only the active experts instead of
+        # all E. This saves ~97% memory for the transient dequant buffer when
+        # few experts are active. MXFP4 always routes through selective
+        # dequant because the kernel needs bf16 weights and full-tensor
+        # dequant of 256-expert MX params is prohibitive.
+        has_bnb_param = (
+            hasattr(experts, "parametrizations")
+            and "gate_up_proj" in experts.parametrizations
+        )
+        has_mxfp4_param = is_mxfp4_param(getattr(experts, "gate_up_proj", None))
+        use_selective = (
+            getattr(self, "_use_selective_dequant", False) and has_bnb_param
+        ) or has_mxfp4_param
+
+        if use_selective:
+            from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (
+                get_active_experts,
+                remap_expert_indices,
+                selective_expert_weights,
+                selective_lora_weights,
+            )
+
+            active_experts = get_active_experts(sorted_expert_idxs, num_experts)
+            remapped_expert_idxs, compact_offsets = remap_expert_indices(
+                sorted_expert_idxs,
+                expert_offsets,
+                active_experts,
+                num_experts,
+            )
+            # Dequantize only active experts' weights
+            gate_up_W = selective_expert_weights(
+                experts,
+                "gate_up_proj",
+                active_experts,
+            ).transpose(2, 1)  # [num_active, hidden, 2*inter]
+
+            # Remap LoRA weights to match compact expert indices
+            if gup_lora is not None:
+                gup_A, gup_B, gup_scaling = gup_lora
+                gup_A, gup_B = selective_lora_weights(
+                    gup_A,
+                    gup_B,
+                    active_experts,
+                    num_experts,
+                )
+                gup_lora = (gup_A, gup_B, gup_scaling)
+
+            # Use remapped indices for ScatterMoE kernels
+            sei_gup = remapped_expert_idxs
+            eo_gup = compact_offsets
+        else:
+            gate_up_W = experts.gate_up_proj.transpose(2, 1)  # [E, hidden, 2*inter]
+            sei_gup = sorted_expert_idxs
+            eo_gup = expert_offsets
+
+        # ====================================================================
+        # Gate + Up projection
+        # ====================================================================
+        if gup_lora is not None:
+            gup_A, gup_B, gup_scaling = gup_lora
+            gup = parallel_linear_lora(
+                hidden_states_flat,
+                gate_up_W,
+                top_k,
+                sei_gup,
+                sorted_scattered_idxs,
+                eo_gup,
+                lora_A=gup_A,
+                lora_B=gup_B,
+                scaling=gup_scaling,
+                grouped_in=False,
+                grouped_out=True,
+                use_fused_dX=True,
+                use_fused_gather=True,
+            )
+        else:
+            gup = parallel_linear(
+                hidden_states_flat,
+                gate_up_W,
+                top_k,
+                sei_gup,
+                sorted_scattered_idxs,
+                eo_gup,
+                grouped_in=False,
+                grouped_out=True,
+            )
+
+        gates, h = gup.chunk(2, dim=-1)
+        h = experts.act_fn(gates) * h
+
+        # ====================================================================
+        # Down projection
+        # ====================================================================
+        if use_selective:
+            down_W = selective_expert_weights(
+                experts,
+                "down_proj",
+                active_experts,
+            ).transpose(2, 1)  # [num_active, inter, hidden]
+
+            if down_lora is not None:
+                down_A, down_B, down_scaling = down_lora
+                down_A, down_B = selective_lora_weights(
+                    down_A,
+                    down_B,
+                    active_experts,
+                    num_experts,
+                )
+                down_lora = (down_A, down_B, down_scaling)
+
+            sei_down = remapped_expert_idxs
+            eo_down = compact_offsets
+        else:
+            down_W = experts.down_proj.transpose(2, 1)  # [E, inter, hidden]
+            sei_down = sorted_expert_idxs
+            eo_down = expert_offsets
+
+        if down_lora is not None:
+            down_A, down_B, down_scaling = down_lora
+            expert_output = parallel_linear_lora(
+                h,
+                down_W,
+                1,
+                sei_down,
+                sorted_scattered_idxs,
+                eo_down,
+                lora_A=down_A,
+                lora_B=down_B,
+                scaling=down_scaling,
+                gates=routing_weights,
+                grouped_in=True,
+                grouped_out=False,
+                use_fused_dX=True,
+                use_fused_gather=True,
+            )
+        else:
+            expert_output = parallel_linear(
+                h,
+                down_W,
+                1,
+                sei_down,
+                sorted_scattered_idxs,
+                eo_down,
+                grouped_in=True,
+                grouped_out=False,
+                gates=routing_weights,
+            )
+
+        # ====================================================================
+        # Combine with shared expert and reshape
+        # ====================================================================
+        if shared_expert_output is not None:
+            expert_output = expert_output + shared_expert_output
+
+        expert_output = expert_output.view(batch_size, sequence_length, hidden_dim)
+        return expert_output
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/lora_ops.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/lora_ops.py
new file mode 100644
index 0000000000..aec68311b6
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/lora_ops.py
@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+ParallelExperts module with LoRA support.
+
+Provides a drop-in replacement for ScatterMoE's ParallelExperts that
+uses the fused LoRA kernel when adapter weights are attached.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from .parallel_linear_lora import parallel_linear_lora
+
+
+class ParallelExperts(nn.Module):
+    """
+    Parallel Experts with fused LoRA support.
+
+    Drop-in replacement for the original ParallelExperts. When LoRA parameters
+    are attached via set_lora(), the forward pass uses a fused kernel:
+        Y = X @ W + scaling * (X @ A^T) @ B^T
+    """
+
+    def __init__(
+        self,
+        num_experts: int,
+        input_size: int,
+        output_size: int,
+        bias: bool = False,
+    ) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(num_experts, output_size, input_size))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(num_experts, output_size))
+        else:
+            self.bias = None
+        self.num_experts = num_experts
+        self.input_size = input_size
+        self.output_size = output_size
+        self._lora_A: torch.Tensor | None = None
+        self._lora_B: torch.Tensor | None = None
+        self._lora_scaling: float | None = None
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.weight, std=0.02)
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def extra_repr(self) -> str:
+        return (
+            f"num_experts={self.num_experts}, "
+            f"input_size={self.input_size}, "
+            f"output_size={self.output_size}"
+        )
+
+    def set_lora(self, lora_A: torch.Tensor, lora_B: torch.Tensor, scaling: float):
+        """Attach LoRA parameters for fused computation."""
+        self._lora_A = lora_A
+        self._lora_B = lora_B
+        self._lora_scaling = scaling
+
+    def clear_lora(self):
+        """Remove LoRA parameters."""
+        self._lora_A = None
+        self._lora_B = None
+        self._lora_scaling = None
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        k: int,
+        sorted_expert_idxs: torch.Tensor,
+        sorted_scattered_idxs: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        gates: Optional[torch.Tensor] = None,
+        grouped_in: bool = False,
+        grouped_out: bool = False,
+    ) -> torch.Tensor:
+        return parallel_linear_lora(
+            inputs,
+            self.weight.permute(0, 2, 1),  # [E, input, output]
+            k,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            lora_A=self._lora_A,
+            lora_B=self._lora_B,
+            scaling=self._lora_scaling if self._lora_scaling is not None else 1.0,
+            expert_biases=self.bias,
+            gates=gates,
+            grouped_in=grouped_in,
+            grouped_out=grouped_out,
+        )
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/__init__.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/__init__.py
new file mode 100644
index 0000000000..a35a482a34
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/__init__.py
@@ -0,0 +1,103 @@
+"""Marlin NVFP4 **W4A16** forward backend for the grouped DeepSeek-V4 MoE (Ampere and newer).
+
+Replaces the CUTLASS FP4xFP4 (lossy fp4 *activations*) base FORWARD with a Marlin weight-only NVFP4
+GEMM (bf16 activations, bit-correct). At the thin-M LoRA-FT shape (E=256, ~48 tok/expert) this is
+~1.79x faster than CUTLASS on sm120 AND removes the ~9.3% activation-quant error of the W4A4 path.
+The Marlin MoE GEMM + ``gptq_marlin_repack`` are vendored under ``_csrc/`` (extracted from vLLM,
+ported to a regular ``torch::Tensor`` ABI, bit-exact to vLLM) so there is NO vLLM runtime dependency.
+JIT-built once (for the current device's arch) via ``torch.utils.cpp_extension.load`` and cached.
+
+Portability: the kernel is the standard Ampere-class Marlin (``mma.m16n8k16`` bf16 + bit-twiddle FP4
+decode) — NO Hopper/Blackwell-only intrinsics. The W4A16 (bf16-act x NVFP4-weight) config runs on
+any sm80+ GPU; the ``__CUDA_ARCH__ < 890`` guard in the template only disables *fp8 activations*,
+which this path never uses. This makes Marlin the fused base GEMM that also covers Ampere (A100,
+sm80) and Ada (RTX 4090 / L40S, sm89), where DeepGEMM (sm90+) and CUTLASS-fp4 (sm120) don't run.
+
+This backend provides the FORWARD only; the backward keeps the existing gradient-consistent
+``_base_dx`` (fp8-read on sm120, now enabled for the pad-64 layout via the BM=tile dX kernel;
+bf16-dequant elsewhere). A Marlin NVFP4 transpose-repack backward was evaluated and rejected: the
+backward contracts the other axis, so reusing the weight needs an independent 4-bit re-quantization
+of W^T that disagrees with the forward weight by ~18% — measured to land ~13% in the LoRA gradients
+(it is a structured weight mismatch, not noise the low-rank projection averages out). fp8-read keeps
+the backward gradient-consistent at 1 byte/wt, which is the floor for block-scaled 4-bit.
+
+Lazy build/import keeps no-nvcc / pre-Ampere environments (incl. CI) clean; ``marlin_w4a16_available``
+returns False there and the dsv4 MoE forward falls back to CUTLASS (sm120) / DeepGEMM (sm90/100) /
+chunked dequant.
+"""
+
+from __future__ import annotations
+
+import functools
+import os
+
+import torch
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_CSRC = os.path.join(_HERE, "_csrc")
+_MOE = os.path.join(_CSRC, "libtorch_stable", "moe", "marlin_moe_wna16")
+_MARLIN_INC = os.path.join(_CSRC, "libtorch_stable", "quantization", "marlin")
+_EXT = None
+
+
+@functools.lru_cache(maxsize=1)
+def marlin_w4a16_available() -> bool:
+    """True iff the standalone Marlin W4A16 ext can build + run here: Ampere or newer (sm80+, the
+    bf16 ``mma`` the W4A16 path needs), CUDA available, and an nvcc present to JIT-build it."""
+    try:
+        if not torch.cuda.is_available():
+            return False
+        if torch.cuda.get_device_capability()[0] < 8:  # bf16 mma needs sm80+
+            return False
+        from torch.utils.cpp_extension import CUDA_HOME
+
+        return CUDA_HOME is not None
+    except Exception:
+        return False
+
+
+def _gencode_for_current_device() -> str:
+    """``-gencode`` for the running GPU's arch (build only what we run). Hopper+ uses the arch-
+    specific ('a') target to match the vendored vLLM build; Ampere/Ada use the plain target."""
+    major, minor = torch.cuda.get_device_capability()
+    arch = f"{major}{minor}"
+    suffix = "a" if major >= 9 else ""
+    return f"-gencode=arch=compute_{arch}{suffix},code=sm_{arch}{suffix}"
+
+
+def load_ext():
+    """JIT-build (cached) and return the vendored standalone Marlin ext module for this GPU's arch.
+
+    Exposes ``moe_wna16_marlin_gemm`` (NVFP4 W4A16 MoE GEMM, bf16 act/out) and ``gptq_marlin_repack``
+    (NVFP4 weight -> Marlin tile layout). Compiled for the current device's compute capability."""
+    global _EXT
+    if _EXT is None:
+        from torch.utils.cpp_extension import load
+
+        major, minor = torch.cuda.get_device_capability()
+        # Arch in the ext name so two arches in one torch-extensions cache don't collide.
+        _EXT = load(
+            name=f"axolotl_marlin_w4a16_sm{major}{minor}",
+            sources=[
+                os.path.join(_MOE, "ops_standalone.cu"),
+                os.path.join(_MOE, "sm80_kernel_bfloat16_fe2m1f_bfloat16.cu"),
+                os.path.join(_MOE, "repack_standalone.cu"),
+            ],
+            extra_cuda_cflags=[
+                "-O3",
+                _gencode_for_current_device(),
+                "-DMARLIN_NAMESPACE_NAME=marlin_moe_wna16",
+                "--expt-relaxed-constexpr",
+                "--expt-extended-lambda",
+                # Keep external linkage on the __global__ template instantiations so
+                # ops_standalone.cu's selector resolves them at link, not static stubs.
+                # Requires CUDA >= 12.4 nvcc.
+                "-static-global-template-stub=false",
+                "-Xcompiler",
+                "-fPIC",
+            ],
+            extra_cflags=["-O3", "-DMARLIN_NAMESPACE_NAME=marlin_moe_wna16"],
+            extra_include_paths=[_CSRC, _MARLIN_INC, _MOE],
+            verbose=False,
+        )
+    return _EXT
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/PROVENANCE.md b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/PROVENANCE.md
new file mode 100644
index 0000000000..766c57a68a
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/PROVENANCE.md
@@ -0,0 +1,34 @@
+# Provenance of the vendored Marlin W4A16 CUDA sources
+
+These files implement the NVFP4 weight-only (W4A16) Marlin MoE GEMM and the
+`gptq_marlin_repack` weight-prep op. They are **extracted from
+[vLLM](https://github.com/vllm-project/vllm)** (`csrc/`), which in turn builds on the
+**Marlin** kernel by Elias Frantar (`Copyright (C) Marlin.2024 Elias Frantar`, modified by
+Neural Magic). Licensing follows the upstream sources: **Apache-2.0** (vLLM) and the Marlin
+license headers retained verbatim in the individual files.
+
+Vendored here (rather than importing vLLM) so this optional sm120 backend adds **no vLLM
+runtime dependency** — consistent with the cutlass-dsl / DeepGEMM optional-backend pattern.
+
+## What was changed vs upstream
+
+Only the two host-wrapper translation units were edited; the kernels and headers are otherwise
+verbatim:
+
+- `libtorch_stable/moe/marlin_moe_wna16/ops_standalone.cu` — the `moe_wna16_marlin_gemm` host
+  wrapper, ported from vLLM's `torch::stable::Tensor` ABI to a regular `torch::Tensor` ABI, and
+  the `STABLE_TORCH_LIBRARY_IMPL` registration replaced by a `PYBIND11_MODULE`.
+- `libtorch_stable/moe/marlin_moe_wna16/repack_standalone.cu` — the `gptq_marlin_repack` host
+  wrapper, same `torch::stable` -> `torch::Tensor` port; compiled in the default `marlin`
+  namespace (the GEMM TU uses `marlin_moe_wna16`).
+
+Everything else (`marlin_template.h`, `marlin.cuh`, `marlin_mma.h`, `marlin_dtypes.cuh`,
+`dequant.h`, `kernel.h`, `kernel_selector.h`, `launcher_body.inc`, `core/scalar_type.hpp`,
+`sm80_kernel_bfloat16_fe2m1f_bfloat16.cu`) is copied verbatim from vLLM.
+
+## Equivalence
+
+The ported `moe_wna16_marlin_gemm` and `gptq_marlin_repack` were validated **bit-exact** against
+the same vLLM ops (identical inputs -> identical outputs). The Python prep (`../prep.py`) copies
+vLLM's pure-torch scale helpers verbatim, so the full NVFP4 -> Marlin weight prep is bit-identical
+to vLLM's `prepare_nvfp4_moe_layer_for_marlin`.
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/core/scalar_type.hpp b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/core/scalar_type.hpp
new file mode 100644
index 0000000000..b6f39ed795
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/core/scalar_type.hpp
@@ -0,0 +1,360 @@
+#pragma once
+
+#include <cstdint>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <variant>
+
+// For STD_TORCH_CHECK
+#include <torch/headeronly/util/Exception.h>
+
+namespace vllm {
+
+//
+//  ScalarType can represent a wide range of floating point and integer types,
+//  in particular it can be used to represent sub-byte data types (something
+//  that torch.dtype currently does not support).
+//
+//  The type definitions on the Python side can be found in: vllm/scalar_type.py
+//  these type definitions should be kept up to date with any Python API changes
+//  here.
+//
+class ScalarType {
+ public:
+  enum NanRepr : uint8_t {
+    NAN_NONE = 0,                // nans are not supported
+    NAN_IEEE_754 = 1,            // nans are: exp all 1s, mantissa not all 0s
+    NAN_EXTD_RANGE_MAX_MIN = 2,  // nans are: exp all 1s, mantissa all 1s
+
+    NAN_REPR_ID_MAX
+  };
+
+  constexpr ScalarType(uint8_t exponent, uint8_t mantissa, bool signed_,
+                       int32_t bias, bool finite_values_only = false,
+                       NanRepr nan_repr = NAN_IEEE_754)
+      : exponent(exponent),
+        mantissa(mantissa),
+        signed_(signed_),
+        bias(bias),
+        finite_values_only(finite_values_only),
+        nan_repr(nan_repr) {};
+
+  static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) {
+    return ScalarType(0, size_bits - 1, true, bias);
+  }
+
+  static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) {
+    return ScalarType(0, size_bits, false, bias);
+  }
+
+  // IEEE 754 compliant floating point type
+  static constexpr ScalarType float_IEEE754(uint8_t exponent,
+                                            uint8_t mantissa) {
+    STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
+    return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754);
+  }
+
+  // IEEE 754 non-compliant floating point type
+  static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
+                                     bool finite_values_only,
+                                     NanRepr nan_repr) {
+    STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
+    STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
+    STD_TORCH_CHECK(
+        nan_repr != NAN_IEEE_754,
+        "use `float_IEEE754` constructor for floating point types that "
+        "follow IEEE 754 conventions");
+    return ScalarType(exponent, mantissa, true, 0, finite_values_only,
+                      nan_repr);
+  }
+
+  uint8_t const exponent;  // size of the exponent field (0 for integer types)
+  uint8_t const mantissa;  // size of the mantissa field (size of the integer
+                           // excluding the sign bit for integer types)
+  bool const signed_;  // flag if the type supports negative numbers (i.e. has a
+                       // sign bit)
+  int32_t const bias;  // stored values equal value + bias,
+                       // used for quantized type
+
+  // Extra Floating point info
+  bool const finite_values_only;  // i.e. no +/-inf if true
+  NanRepr const nan_repr;         // how NaNs are represented
+                                  // (not applicable for integer types)
+
+  using Id = int64_t;
+
+ private:
+  // Field size in id
+  template <typename T_>
+  static constexpr size_t member_id_field_width() {
+    using T = std::decay_t<T_>;
+    return std::is_same_v<T, bool> ? 1 : sizeof(T) * 8;
+  }
+
+  template <typename Fn, typename Init, typename Member, typename... Rest>
+  static constexpr auto reduce_members_helper(Fn f, Init val, Member member,
+                                              Rest... rest) {
+    auto new_val = f(val, member);
+    if constexpr (sizeof...(rest) > 0) {
+      return reduce_members_helper(f, new_val, rest...);
+    } else {
+      return new_val;
+    };
+  }
+
+  template <typename Fn, typename Init>
+  constexpr auto reduce_members(Fn f, Init init) const {
+    // Should be in constructor order for `from_id`
+    return reduce_members_helper(f, init, exponent, mantissa, signed_, bias,
+                                 finite_values_only, nan_repr);
+  };
+
+  template <typename Fn, typename Init>
+  static constexpr auto reduce_member_types(Fn f, Init init) {
+    constexpr auto dummy_type = ScalarType(0, 0, false, 0, false, NAN_NONE);
+    return dummy_type.reduce_members(f, init);
+  };
+
+  static constexpr auto id_size_bits() {
+    return reduce_member_types(
+        [](int acc, auto member) -> int {
+          return acc + member_id_field_width<decltype(member)>();
+        },
+        0);
+  }
+
+ public:
+  // unique id for this scalar type that can be computed at compile time for
+  //  c++17 template specialization this is not needed once we migrate to
+  //  c++20 and can pass literal classes as template parameters
+  constexpr Id id() const {
+    static_assert(id_size_bits() <= sizeof(Id) * 8,
+                  "ScalarType id is too large to be stored");
+
+    auto or_and_advance = [](std::pair<Id, uint32_t> result,
+                             auto member) -> std::pair<Id, uint32_t> {
+      auto [id, bit_offset] = result;
+      auto constexpr bits = member_id_field_width<decltype(member)>();
+      return {id | (int64_t(member) & ((uint64_t(1) << bits) - 1))
+                       << bit_offset,
+              bit_offset + bits};
+    };
+    return reduce_members(or_and_advance, std::pair<Id, uint32_t>{}).first;
+  }
+
+  // create a ScalarType from an id, for c++17 template specialization,
+  //  this is not needed once we migrate to c++20 and can pass literal
+  //  classes as template parameters
+  static constexpr ScalarType from_id(Id id) {
+    auto extract_and_advance = [id](auto result, auto member) {
+      using T = decltype(member);
+      auto [tuple, bit_offset] = result;
+      auto constexpr bits = member_id_field_width<T>();
+      auto extracted_val = static_cast<T>((int64_t(id) >> bit_offset) &
+                                          ((uint64_t(1) << bits) - 1));
+      auto new_tuple = std::tuple_cat(tuple, std::make_tuple(extracted_val));
+      return std::pair<decltype(new_tuple), int>{new_tuple, bit_offset + bits};
+    };
+
+    auto [tuple_args, _] = reduce_member_types(extract_and_advance,
+                                               std::pair<std::tuple<>, int>{});
+    return std::apply([](auto... args) { return ScalarType(args...); },
+                      tuple_args);
+  }
+
+  constexpr int64_t size_bits() const {
+    return mantissa + exponent + is_signed();
+  }
+  constexpr bool is_signed() const { return signed_; }
+  constexpr bool is_integer() const { return exponent == 0; }
+  constexpr bool is_floating_point() const { return exponent > 0; }
+  constexpr bool is_ieee_754() const {
+    return is_floating_point() && finite_values_only == false &&
+           nan_repr == NAN_IEEE_754;
+  }
+  constexpr bool has_nans() const {
+    return is_floating_point() && nan_repr != NAN_NONE;
+  }
+  constexpr bool has_infs() const {
+    return is_floating_point() && finite_values_only == false;
+  }
+  constexpr bool has_bias() const { return bias != 0; }
+
+ private:
+  double _floating_point_max() const {
+    STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11,
+                    "Cannot represent max/min as a double for type ", str());
+
+    uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
+    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
+      max_mantissa -= 1;
+    }
+
+    uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
+    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
+      STD_TORCH_CHECK(exponent < 11,
+                      "Cannot represent max/min as a double for type ", str());
+      max_exponent += 1;
+    }
+
+    // adjust the exponent to match that of a double
+    //  for now we assume the exponent bias is the standard 2^(e-1) -1, (where e
+    //  is the exponent bits), there is some precedent for non-standard biases,
+    //  example `float8_e4m3b11fnuz` here: https://github.com/jax-ml/ml_dtypes
+    //  but to avoid premature over complication we are just assuming the
+    //  standard exponent bias until there is a need to support non-standard
+    //  biases
+    uint64_t exponent_bias = (uint64_t(1) << (exponent - 1)) - 1;
+    uint64_t exponent_bias_double = (uint64_t(1) << 10) - 1;  // double e = 11
+
+    uint64_t max_exponent_double =
+        max_exponent - exponent_bias + exponent_bias_double;
+
+    // shift the mantissa into the position for a double and
+    // the exponent
+    uint64_t double_raw =
+        (max_mantissa << (52 - mantissa)) | (max_exponent_double << 52);
+
+    return *reinterpret_cast<double*>(&double_raw);
+  }
+
+  constexpr std::variant<int64_t, double> _raw_max() const {
+    if (is_floating_point()) {
+      return {_floating_point_max()};
+    } else {
+      STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
+                      "Cannot represent max as a int64_t");
+      return {(int64_t(1) << mantissa) - 1};
+    }
+  }
+
+  constexpr std::variant<int64_t, double> _raw_min() const {
+    if (is_floating_point()) {
+      STD_TORCH_CHECK(
+          is_signed(),
+          "We currently assume all floating point types are signed");
+      constexpr uint64_t sign_bit_double = (uint64_t(1) << 63);
+
+      double max = _floating_point_max();
+      uint64_t max_raw = *reinterpret_cast<uint64_t*>(&max);
+      uint64_t min_raw = max_raw | sign_bit_double;
+      return {*reinterpret_cast<double*>(&min_raw)};
+    } else {
+      STD_TORCH_CHECK(!is_signed() || size_bits() <= 64,
+                      "Cannot represent min as a int64_t");
+      if (is_signed()) {
+        // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0
+        // then perform an arithmetic shift right to set all the bits above
+        // (size_bits() - 1) to 1
+        return {INT64_MIN >> (64 - size_bits())};
+      } else {
+        return {int64_t(0)};
+      }
+    }
+  }
+
+ public:
+  // Max representable value for this scalar type.
+  // (accounting for bias if there is one)
+  constexpr std::variant<int64_t, double> max() const {
+    return std::visit(
+        [this](auto x) -> std::variant<int64_t, double> { return {x - bias}; },
+        _raw_max());
+  }
+
+  // Min representable value for this scalar type.
+  // (accounting for bias if there is one)
+  constexpr std::variant<int64_t, double> min() const {
+    return std::visit(
+        [this](auto x) -> std::variant<int64_t, double> { return {x - bias}; },
+        _raw_min());
+  }
+
+  std::string str() const {
+    /* naming generally follows: https://github.com/jax-ml/ml_dtypes
+     * for floating point types (leading f) the scheme is:
+     *  `float<size_bits>_e<exponent_bits>m<mantissa_bits>[flags]`
+     *  flags:
+     *  - no-flags: means it follows IEEE 754 conventions
+     *  - f: means finite values only (no infinities)
+     *  - n: means nans are supported (non-standard encoding)
+     * for integer types the scheme is:
+     *  `[u]int<size_bits>[b<bias>]`
+     *  - if bias is not present it means its zero
+     */
+    if (is_floating_point()) {
+      auto ret = "float" + std::to_string(size_bits()) + "_e" +
+                 std::to_string(exponent) + "m" + std::to_string(mantissa);
+      if (!is_ieee_754()) {
+        if (finite_values_only) {
+          ret += "f";
+        }
+        if (nan_repr != NAN_NONE) {
+          ret += "n";
+        }
+      }
+      return ret;
+    } else {
+      auto ret = ((is_signed()) ? "int" : "uint") + std::to_string(size_bits());
+      if (has_bias()) {
+        ret += "b" + std::to_string(bias);
+      }
+      return ret;
+    }
+  }
+
+  constexpr bool operator==(ScalarType const& other) const {
+    return mantissa == other.mantissa && exponent == other.exponent &&
+           bias == other.bias && signed_ == other.signed_ &&
+           finite_values_only == other.finite_values_only &&
+           nan_repr == other.nan_repr;
+  }
+};
+
+using ScalarTypeId = ScalarType::Id;
+
+// "rust style" names generally following:
+//   https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L60-L70
+static inline constexpr auto kS4 = ScalarType::int_(4);
+static inline constexpr auto kU4 = ScalarType::uint(4);
+static inline constexpr auto kU4B8 = ScalarType::uint(4, 8);
+static inline constexpr auto kS8 = ScalarType::int_(8);
+static inline constexpr auto kU8 = ScalarType::uint(8);
+static inline constexpr auto kU8B128 = ScalarType::uint(8, 128);
+
+static inline constexpr auto kFE2M1f =
+    ScalarType::float_(2, 1, true, ScalarType::NAN_NONE);
+static inline constexpr auto kFE3M2f =
+    ScalarType::float_(3, 2, true, ScalarType::NAN_NONE);
+static inline constexpr auto kFE4M3fn =
+    ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE8M0fnu =
+    ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2);
+static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7);
+static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10);
+
+// Fixed width style names, generally following:
+//  https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57
+static inline constexpr auto kInt4 = kS4;
+static inline constexpr auto kUint4 = kU4;
+static inline constexpr auto kUint4b8 = kU4B8;
+static inline constexpr auto kInt8 = kS8;
+static inline constexpr auto kUint8 = kU8;
+static inline constexpr auto kUint8b128 = kU8B128;
+
+static inline constexpr auto kFloat4_e2m1f = kFE2M1f;
+static inline constexpr auto kFloat6_e3m2f = kFE3M2f;
+static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn;
+static inline constexpr auto kFloat8_e5m2 = kFE5M2;
+static inline constexpr auto kFloat16_e8m7 = kFE8M7;
+static inline constexpr auto kFloat16_e5m10 = kFE5M10;
+
+// colloquial names
+static inline constexpr auto kHalf = kFE5M10;
+static inline constexpr auto kFloat16 = kHalf;
+static inline constexpr auto kBFloat16 = kFE8M7;
+
+static inline constexpr auto kFloat16Id = kFloat16.id();
+};  // namespace vllm
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/launcher_body.inc b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/launcher_body.inc
new file mode 100644
index 0000000000..e47705270b
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/launcher_body.inc
@@ -0,0 +1,500 @@
+namespace MARLIN_NAMESPACE_NAME {
+
+__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};
+
+using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);
+
+// For a given "a" of size [M,K] performs a permutation of the K columns based
+// on the given "perm" indices.
+template <int moe_block_size>
+__global__ void permute_cols_kernel(
+    int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr,
+    int4* __restrict__ out_int4_ptr,
+    const int32_t* __restrict__ sorted_token_ids_ptr,
+    const int32_t* __restrict__ expert_ids_ptr,
+    const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m,
+    int size_k, int top_k) {
+  int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
+  int num_moe_blocks = div_ceil(num_tokens_past_padded, moe_block_size);
+  int32_t block_sorted_ids[moe_block_size];
+  int block_num_valid_tokens = 0;
+  int64_t old_expert_id = 0;
+  int64_t expert_id = 0;
+  int row_stride = size_k * sizeof(half) / 16;
+
+  auto read_moe_block_data = [&](int block_id) {
+    block_num_valid_tokens = moe_block_size;
+    int4* tmp_block_sorted_ids = reinterpret_cast<int4*>(block_sorted_ids);
+    for (int i = 0; i < moe_block_size / 4; i++) {
+      tmp_block_sorted_ids[i] =
+          ((int4*)sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i];
+    }
+    for (int i = 0; i < moe_block_size; i++) {
+      if (block_sorted_ids[i] >= size_m * top_k) {
+        block_num_valid_tokens = i;
+        break;
+      };
+    }
+  };
+
+  auto permute_row = [&](int row) {
+    int iters = size_k / default_threads;
+    int rest = size_k % default_threads;
+
+    int in_offset = (row / top_k) * row_stride;
+    int out_offset = row * row_stride;
+
+    half const* a_row_half =
+        reinterpret_cast<half const*>(a_int4_ptr + in_offset);
+    half* out_half = reinterpret_cast<half*>(out_int4_ptr + out_offset);
+
+    int base_k = 0;
+
+    for (int i = 0; i < iters; i++) {
+      auto cur_k = base_k + threadIdx.x;
+      int src_pos = perm_int_ptr[cur_k];
+
+      out_half[cur_k] = a_row_half[src_pos];
+
+      base_k += default_threads;
+    }
+
+    if (rest) {
+      if (threadIdx.x < rest) {
+        auto cur_k = base_k + threadIdx.x;
+        int src_pos = perm_int_ptr[cur_k];
+
+        out_half[cur_k] = a_row_half[src_pos];
+      }
+    }
+  };
+
+  for (int index = blockIdx.x; index < num_moe_blocks; index += gridDim.x) {
+    old_expert_id = expert_id;
+    int tmp_expert_id = expert_ids_ptr[index];
+    if (tmp_expert_id == -1) continue;
+    expert_id = tmp_expert_id;
+    perm_int_ptr += (expert_id - old_expert_id) * size_k;
+    read_moe_block_data(index);
+
+    for (int i = 0; i < block_num_valid_tokens; i++)
+      permute_row(block_sorted_ids[i]);
+  }
+}
+
+typedef struct {
+  int thread_k;
+  int thread_n;
+  int num_threads;
+} thread_config_t;
+
+thread_config_t small_batch_thread_configs[] = {
+    // Ordered by priority
+
+    // thread_k, thread_n, num_threads
+    {128, 128, 256},
+    {64, 128, 128},
+    {128, 64, 128}};
+
+thread_config_t large_batch_thread_configs[] = {
+    // Ordered by priority
+
+    // thread_k, thread_n, num_threads
+    {64, 256, 256},
+    {64, 128, 128},
+    {128, 64, 128}};
+
+typedef struct {
+  int blocks_per_sm;
+  thread_config_t tb_cfg;
+} exec_config_t;
+
+int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
+                          int prob_n, int prob_k, int num_bits, int group_size,
+                          bool has_act_order, bool is_k_full, int stages) {
+  bool cache_scales_chunk = has_act_order && !is_k_full;
+
+  int tb_n = th_config.thread_n;
+  int tb_k = th_config.thread_k;
+
+  // Get max scale groups per thread-block
+  int tb_groups;
+  if (group_size == -1) {
+    tb_groups = 1;
+  } else if (group_size == 0) {
+    tb_groups = div_ceil(tb_k, 32);  // Worst case is 32 group size
+  } else {
+    tb_groups = div_ceil(tb_k, group_size);
+  }
+
+  if (cache_scales_chunk) {
+    int load_groups =
+        tb_groups * stages * 2;          // Chunk size is 2x pipeline over dim K
+    load_groups = max(load_groups, 32);  // We load at least 32 scale groups
+    return load_groups * tb_n * 2;
+  } else {
+    int tb_scales = tb_groups * tb_n * 2;
+
+    return tb_scales * stages;
+  }
+}
+
+int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
+                          int thread_m_blocks, int prob_m, int prob_n,
+                          int prob_k, int num_bits, int group_size,
+                          bool has_act_order, bool is_k_full, int has_zp,
+                          int is_zp_float, bool is_a_8bit, int stages) {
+  int pack_factor = 32 / num_bits;
+
+  // Get B size
+  int tb_k = th_config.thread_k;
+  int tb_n = th_config.thread_n;
+  int tb_m = thread_m_blocks * 16;
+
+  // shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights
+  // both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32)
+  int sh_block_meta_size = tb_m * 16;
+  int sh_a_size = stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2);
+  int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4;
+  int sh_red_size = tb_m * (tb_n + 8) * 2;
+  int sh_bias_size = tb_n * 2;
+  int tmp_size =
+      (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size;
+  tmp_size = max(max(sh_b_size, sh_red_size), tmp_size);
+
+  int sh_s_size =
+      get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
+                            group_size, has_act_order, is_k_full, stages);
+  int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0;
+  int sh_zp_size = 0;
+  if (has_zp) {
+    if (is_zp_float)
+      sh_zp_size = sh_s_size;
+    else if (num_bits == 4)
+      sh_zp_size = sh_s_size / 4;
+    else if (num_bits == 8)
+      sh_zp_size = sh_s_size / 2;
+  }
+
+  int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size +
+                   sh_g_idx_size + sh_block_meta_size;
+
+  return total_size;
+}
+
+bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
+                     int thread_m_blocks, int prob_m, int prob_n, int prob_k,
+                     int num_bits, int group_size, bool has_act_order,
+                     bool is_k_full, int has_zp, int is_zp_float,
+                     bool is_a_8bit, int stages, int max_shared_mem) {
+  // Sanity
+  if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
+      th_config.num_threads == -1) {
+    return false;
+  }
+
+  // Verify K/N are divisible by thread K/N
+  if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) {
+    return false;
+  }
+
+  // Verify min for thread K/N
+  if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) {
+    return false;
+  }
+
+  // num_threads must be at least 128 (= 4 warps)
+  if (th_config.num_threads < 128) {
+    return false;
+  }
+
+  // Check that pipeline fits into cache
+  int cache_size =
+      get_kernel_cache_size(th_config, m_block_size_8, thread_m_blocks, prob_m,
+                            prob_n, prob_k, num_bits, group_size, has_act_order,
+                            is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
+  return cache_size <= max_shared_mem;
+}
+
+MarlinFuncPtr get_marlin_kernel(
+    const vllm::ScalarType a_type, const vllm::ScalarType b_type,
+    const vllm::ScalarType c_type, const vllm::ScalarType s_type,
+    int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
+    bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
+    int threads, bool is_zp_float, int stages) {
+  int num_bits = b_type.size_bits();
+  auto kernel = MarlinDefault;
+
+#include "kernel_selector.h"
+
+  return kernel;
+}
+
+exec_config_t determine_exec_config(
+    const vllm::ScalarType& a_type, const vllm::ScalarType& b_type,
+    const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
+    int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks,
+    bool m_block_size_8, int num_bits, int group_size, bool has_act_order,
+    bool is_k_full, bool has_zp, bool is_zp_float, bool is_a_8bit, int stages,
+    int max_shared_mem, int sms) {
+  exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
+  thread_config_t* thread_configs = thread_m_blocks > 1
+                                        ? large_batch_thread_configs
+                                        : small_batch_thread_configs;
+  int thread_configs_size =
+      thread_m_blocks > 1
+          ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t)
+          : sizeof(small_batch_thread_configs) / sizeof(thread_config_t);
+
+  int count = 0;
+  constexpr int device_max_reg_size = 255 * 1024;
+  for (int i = 0; i < thread_configs_size; i++) {
+    thread_config_t th_config = thread_configs[i];
+
+    if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m,
+                         prob_n, prob_k, num_bits, group_size, has_act_order,
+                         is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
+                         max_shared_mem - 512)) {
+      continue;
+    }
+
+    int cache_size = get_kernel_cache_size(
+        th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
+        num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
+        is_a_8bit, stages);
+
+    int group_blocks = 0;
+    if (!has_act_order) {
+      group_blocks = group_size == -1 ? -1 : (group_size / 16);
+    }
+
+    auto kernel =
+        get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
+                          th_config.thread_n / 16, th_config.thread_k / 16,
+                          m_block_size_8, has_act_order, has_zp, group_blocks,
+                          th_config.num_threads, is_zp_float, stages);
+
+    if (kernel == MarlinDefault) continue;
+
+    cudaFuncAttributes attr;
+    cudaFuncGetAttributes(&attr, kernel);
+    int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
+    int allow_count = min(device_max_reg_size / reg_size,
+                          max_shared_mem / (cache_size + 1536));
+    if (thread_m_blocks == 1)
+      allow_count = max(min(allow_count, 4), 1);
+    else
+      allow_count = max(min(allow_count, 2), 1);
+
+    if (prob_n / th_config.thread_n * prob_m * top_k * 4 < sms * allow_count) {
+      allow_count =
+          max(prob_n / th_config.thread_n * prob_m * top_k * 4 / sms, 1);
+    }
+
+    if (allow_count > count) {
+      count = allow_count;
+      exec_cfg = {count, th_config};
+    };
+  }
+
+  return exec_cfg;
+}
+
+void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
+               void* a_s, void* b_s, void* g_s, void* zp, void* g_idx,
+               void* perm, void* a_tmp, void* sorted_token_ids,
+               void* expert_ids, void* num_tokens_past_padded,
+               void* topk_weights, int moe_block_size, int num_experts,
+               int top_k, bool mul_topk_weights, int prob_m, int prob_n,
+               int prob_k, void* workspace, vllm::ScalarType const& a_type,
+               vllm::ScalarType const& b_type, vllm::ScalarType const& c_type,
+               vllm::ScalarType const& s_type, bool has_bias,
+               bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
+               int group_size, int dev, cudaStream_t stream, int thread_k,
+               int thread_n, int sms, int blocks_per_sm, bool use_atomic_add,
+               bool use_fp32_reduce, bool is_zp_float) {
+  int thread_m_blocks = div_ceil(moe_block_size, 16);
+  bool m_block_size_8 = moe_block_size == 8;
+  bool is_a_8bit = a_type.size_bits() == 8;
+
+  STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [",
+                  prob_m, ", ", prob_n, ", ", prob_k, "]");
+
+  int group_blocks = 0;
+  if (has_act_order) {
+    if (is_k_full) {
+      STD_TORCH_CHECK(group_size != -1);
+      group_blocks = group_size / 16;
+      STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
+                      " is not divisible by group_blocks = ", group_blocks);
+    } else {
+      STD_TORCH_CHECK(group_size == 0);
+      group_blocks = 0;
+    }
+  } else {
+    if (group_size == -1) {
+      group_blocks = -1;
+    } else {
+      group_blocks = group_size / 16;
+      STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
+                      " is not divisible by group_blocks = ", group_blocks);
+    }
+  }
+
+  int num_bits = b_type.size_bits();
+  const int4* A_ptr = (const int4*)A;
+  const int4* B_ptr = (const int4*)B;
+  int4* C_ptr = (int4*)C;
+  int4* C_tmp_ptr = (int4*)C_tmp;
+  const int4* bias_ptr = (const int4*)b_bias;
+  const float* a_s_ptr = (const float*)a_s;
+  const int4* b_s_ptr = (const int4*)b_s;
+  const float* g_s_ptr = (const float*)g_s;
+  const int4* zp_ptr = (const int4*)zp;
+  const int* g_idx_ptr = (const int*)g_idx;
+  const int* perm_ptr = (const int*)perm;
+  int4* a_tmp_ptr = (int4*)a_tmp;
+  const int32_t* sorted_token_ids_ptr = (const int32_t*)sorted_token_ids;
+  const int32_t* expert_ids_ptr = (const int32_t*)expert_ids;
+  const int32_t* num_tokens_past_padded_ptr =
+      (const int32_t*)num_tokens_past_padded;
+  const float* topk_weights_ptr = (const float*)topk_weights;
+  int* locks = (int*)workspace;
+
+  if (has_act_order) {
+    // Permute A columns
+    auto kernel = permute_cols_kernel<8>;
+    if (moe_block_size == 8) {
+    } else if (moe_block_size == 16)
+      kernel = permute_cols_kernel<16>;
+    else if (moe_block_size == 32)
+      kernel = permute_cols_kernel<32>;
+    else if (moe_block_size == 48)
+      kernel = permute_cols_kernel<48>;
+    else if (moe_block_size == 64)
+      kernel = permute_cols_kernel<64>;
+    else
+      STD_TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size);
+
+    // avoid ">>>" being formatted to "> > >"
+    // clang-format off
+    kernel<<<sms, default_threads, 0, stream>>>(
+        A_ptr, perm_ptr, a_tmp_ptr, sorted_token_ids_ptr, expert_ids_ptr,
+        num_tokens_past_padded_ptr, prob_m, prob_k, top_k);
+    // clang-format on
+    A_ptr = a_tmp_ptr;
+    prob_m = prob_m * top_k;
+    top_k = 1;
+
+    // If we have a full K, then we can run the non-act-order version of Marlin
+    // (since the weight rows are reordered by increasing group ids, and by
+    // having a full K, we have full original groups)
+    if (is_k_full) has_act_order = false;
+  }
+
+  int max_shared_mem = 0;
+  cudaDeviceGetAttribute(&max_shared_mem,
+                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
+  STD_TORCH_CHECK(max_shared_mem > 0);
+
+  int major_capability, minor_capability;
+  cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
+                         dev);
+  cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
+                         dev);
+  STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
+                  "marlin kernel only support Turing or newer GPUs.");
+  int stages = 4;
+  if (major_capability == 7 && minor_capability == 5) {
+    stages = 2;
+    STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
+                    "Turing only support FP16 or INT8 activation.");
+  }
+  if (a_type == vllm::kFE4M3fn) {
+    STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
+                    "FP8 only support Ada Lovelace or newer GPUs.");
+    STD_TORCH_CHECK(
+        major_capability * 10 + minor_capability == 89 ||
+            major_capability == 12,
+        "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
+        "Marlin W4A16 on other devices).");
+  }
+
+  // Set thread config
+  exec_config_t exec_cfg;
+  thread_config_t thread_tfg;
+  if (thread_k != -1 && thread_n != -1) {
+    thread_tfg = thread_config_t{thread_k, thread_n, thread_k * thread_n / 64};
+    if (blocks_per_sm == -1) blocks_per_sm = 1;
+    exec_cfg = exec_config_t{blocks_per_sm, thread_tfg};
+    STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
+                    " is not divisible by thread_n = ", thread_n);
+    STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
+                    " is not divisible by thread_k = ", thread_k);
+  } else {
+    // Auto config
+    exec_cfg = determine_exec_config(
+        a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts,
+        top_k, thread_m_blocks, m_block_size_8, num_bits, group_size,
+        has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
+        max_shared_mem, sms);
+    thread_tfg = exec_cfg.tb_cfg;
+  }
+
+  int num_threads = thread_tfg.num_threads;
+  thread_k = thread_tfg.thread_k;
+  thread_n = thread_tfg.thread_n;
+  int blocks = sms * exec_cfg.blocks_per_sm;
+  if (exec_cfg.blocks_per_sm > 1)
+    max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - 1024;
+
+  int thread_k_blocks = thread_k / 16;
+  int thread_n_blocks = thread_n / 16;
+
+  STD_TORCH_CHECK(
+      is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
+                      prob_n, prob_k, num_bits, group_size, has_act_order,
+                      is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
+                      max_shared_mem),
+      "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
+      ", thread_k = ", thread_tfg.thread_k,
+      ", thread_n = ", thread_tfg.thread_n,
+      ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
+      prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
+      ", group_size = ", group_size, ", has_act_order = ", has_act_order,
+      ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
+      ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);
+
+  int sh_cache_size =
+      get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
+                            prob_n, prob_k, num_bits, group_size, has_act_order,
+                            is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
+
+  auto kernel = get_marlin_kernel(
+      a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
+      thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
+      num_threads, is_zp_float, stages);
+
+  if (kernel == MarlinDefault) {
+    STD_TORCH_CHECK(
+        false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ",
+        prob_k, "]", ", has_act_order = ", has_act_order,
+        ", num_groups = ", num_groups, ", group_size = ", group_size,
+        ", thread_m_blocks = ", thread_m_blocks,
+        ", thread_n_blocks = ", thread_n_blocks,
+        ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits);
+  }
+
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       max_shared_mem);
+  // avoid ">>>" being formatted to "> > >"
+  // clang-format off
+  kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
+      A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, g_idx_ptr,
+      sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
+      topk_weights_ptr, top_k, mul_topk_weights, num_groups, prob_m,
+      prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce);
+  // clang-format on
+}
+
+}  // namespace MARLIN_NAMESPACE_NAME
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h
new file mode 100644
index 0000000000..783736ab50
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h
@@ -0,0 +1,47 @@
+
+#ifndef MARLIN_NAMESPACE_NAME
+  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#endif
+
+#include "libtorch_stable/quantization/marlin/marlin.cuh"
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
+#include "core/scalar_type.hpp"
+
+#define MARLIN_KERNEL_PARAMS                                          \
+  const int4 *__restrict__ A, const int4 *__restrict__ B,             \
+      int4 *__restrict__ C, int4 *__restrict__ C_tmp,                 \
+      const int4 *__restrict__ b_bias_ptr,                            \
+      const float *__restrict__ a_scales_ptr,                         \
+      const int4 *__restrict__ scales_ptr,                            \
+      const float *__restrict__ global_scale_ptr,                     \
+      const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
+      const int32_t *__restrict__ sorted_token_ids_ptr,               \
+      const int32_t *__restrict__ expert_ids_ptr,                     \
+      const int32_t *__restrict__ num_tokens_past_padded_ptr,         \
+      const float *__restrict__ topk_weights_ptr, int top_k,          \
+      bool mul_topk_weights, int num_groups, int prob_m, int prob_n,  \
+      int prob_k, int *locks, bool has_bias, bool use_atomic_add,     \
+      bool use_fp32_reduce
+
+namespace MARLIN_NAMESPACE_NAME {
+template <const vllm::ScalarTypeId a_type_id,  // A ScalarType id
+          const vllm::ScalarTypeId b_type_id,  // B ScalarType id
+          const vllm::ScalarTypeId c_type_id,  // C ScalarType id
+          const vllm::ScalarTypeId s_type_id,  // B_SCALE ScalarType id
+          const int threads,          // number of threads in a threadblock
+          const int thread_m_blocks,  // number of 16x16 blocks in the m
+                                      // dimension (batchsize) of the
+                                      // threadblock
+          const int thread_n_blocks,  // same for n dimension (output)
+          const int thread_k_blocks,  // same for k dimension (reduction)
+          const bool m_block_size_8,  // whether m_block_size == 8
+                                      // only works when thread_m_blocks == 1
+          const int stages,  // number of stages for the async global->shared
+                             // fetch pipeline
+          const int group_blocks,  // number of consecutive 16x16 blocks
+                                   // with a separate quantization scale
+          const bool is_zp_float   // is zero point of float16 type?
+          >
+__global__ void Marlin(MARLIN_KERNEL_PARAMS);
+
+}
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h
new file mode 100644
index 0000000000..fd6905de1d
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h
@@ -0,0 +1,32 @@
+// auto generated by generate_kernels.py
+// clang-format off
+if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, true, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, true, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, true, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 2, 16, 4, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 8, 4, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 4, 8, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 3, 16, 4, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 8, 4, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 4, 8, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 4, 16, 4, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 8, 4, false, 4, 1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kFE2M1f && c_type == vllm::kBFloat16 && s_type == vllm::kFE4M3fn && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 4, 8, false, 4, 1, false>;
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h
new file mode 100644
index 0000000000..04f90101be
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h
@@ -0,0 +1,2241 @@
+/*
+ * Modified by Neural Magic
+ * Copyright (C) Marlin.2024 Elias Frantar
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *         http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * Adapted from https://github.com/IST-DASLab/marlin
+ */
+
+#ifndef MARLIN_NAMESPACE_NAME
+  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#endif
+
+#include "libtorch_stable/quantization/marlin/marlin.cuh"
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
+#include "libtorch_stable/quantization/marlin/dequant.h"
+#include "libtorch_stable/quantization/marlin/marlin_mma.h"
+#include "core/scalar_type.hpp"
+
+#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t)               \
+  static_assert(std::is_same<scalar_t, half>::value ||          \
+                    std::is_same<scalar_t, nv_bfloat16>::value, \
+                "only float16 and bfloat16 is supported");
+
+namespace MARLIN_NAMESPACE_NAME {
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
+
+template <typename scalar_t,  // compute dtype, half or nv_float16
+          const vllm::ScalarTypeId b_type_id,  // weight MarlinScalarType id
+          const int threads,          // number of threads in a threadblock
+          const int thread_m_blocks,  // number of 16x16 blocks in the m
+                                      // dimension (batchsize) of the
+                                      // threadblock
+          const int thread_n_blocks,  // same for n dimension (output)
+          const int thread_k_blocks,  // same for k dimension (reduction)
+          const bool m_block_size_8,  // whether m_block_size == 8
+                                      // only works when thread_m_blocks == 1
+          const int stages,  // number of stages for the async global->shared
+                             // fetch pipeline
+          const bool has_act_order,  // whether act_order is enabled
+          const int group_blocks,    // number of consecutive 16x16 blocks
+                                     // with a separate quantization scale
+          const bool is_zp_float     // is zero point of float16 type?
+          >
+__global__ void Marlin(
+    const int4* __restrict__ A,  // fp16 input matrix of shape mxk
+    const int4* __restrict__ B,  // 4bit quantized weight matrix of shape kxn
+    int4* __restrict__ C,        // fp16 output buffer of shape mxn
+    int4* __restrict__ C_tmp,    // fp32 tmp output buffer (for reduce)
+    const int4* __restrict__ scales_ptr,  // fp16 quantization scales of shape
+                                          // (k/groupsize)xn
+    const int4* __restrict__ zp_ptr,      // 4bit packed zero-points of shape
+                                          // (k/groupsize)x(n/pack_factor)
+    const int* __restrict__ g_idx,        // int32 group indices of shape k
+    const int32_t* __restrict__ sorted_token_ids_ptr,        // moe sorted_ids
+    const int32_t* __restrict__ expert_ids_ptr,              // moe expert ids
+    const int32_t* __restrict__ num_tokens_past_padded_ptr,  // moe num tokens
+    const float* __restrict__ topk_weights_ptr,              // moe top weights
+    int top_k,              // num of experts per token
+    bool mul_topk_weights,  // mul topk weights or not
+    int num_groups,         // number of scale groups per output channel
+    int prob_m,             // batch dimension m
+    int prob_n,             // output dimension n
+    int prob_k,             // reduction dimension k
+    int* locks,             // extra global storage for barrier synchronization
+    bool use_atomic_add,    // whether to use atomic add to reduce
+    bool use_fp32_reduce    // whether to use fp32 global reduce
+) {}
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+#else
+
+// Instruction for loading a full 16x16 matrix fragment of operand A from shared
+// memory, directly in tensor core layout.
+template <int count, vllm::ScalarTypeId type_id>
+__device__ inline void ldsm(typename MarlinScalarType<type_id>::FragA& frag_a,
+                            const void* smem_ptr) {
+  uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  if constexpr (count == 4) {
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
+        : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])
+        : "r"(smem));
+  } else if constexpr (count == 2) {
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n"
+                 : "=r"(a[0]), "=r"(a[1])
+                 : "r"(smem));
+  } else if constexpr (count == 1) {
+    asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n"
+                 : "=r"(a[0])
+                 : "r"(smem));
+  } else {
+    static_assert(count == 1 || count == 2 || count == 4, "invalid count");
+  }
+}
+
+// Multiply dequantized values by the corresponding quantization scale; used
+// only for grouped quantization.
+template <vllm::ScalarTypeId type_id>
+__device__ inline void scale(typename MarlinScalarType<type_id>::FragB& frag_b,
+                             typename MarlinScalarType<type_id>::FragS& frag_s,
+                             int i) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
+  scalar_t2 s = MarlinScalarType<type_id>::num2num2(
+      reinterpret_cast<scalar_t*>(&frag_s)[i]);
+  frag_b[0] = __hmul2(frag_b[0], s);
+  frag_b[1] = __hmul2(frag_b[1], s);
+}
+
+template <vllm::ScalarTypeId type_id>
+__device__ inline void scale_and_sub(
+    typename MarlinScalarType<type_id>::FragB& frag_b,
+    typename MarlinScalarType<type_id>::scalar_t s,
+    typename MarlinScalarType<type_id>::scalar_t zp) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
+  scalar_t2 s2 = MarlinScalarType<type_id>::num2num2(s);
+  scalar_t2 zp2 = MarlinScalarType<type_id>::num2num2(zp);
+  frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2));
+  frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2));
+}
+
+template <vllm::ScalarTypeId type_id>
+__device__ inline void sub_zp(
+    typename MarlinScalarType<type_id>::FragB& frag_b,
+    typename MarlinScalarType<type_id>::scalar_t2& frag_zp, int i) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
+  scalar_t2 zp = MarlinScalarType<type_id>::num2num2(
+      reinterpret_cast<scalar_t*>(&frag_zp)[i]);
+  frag_b[0] = __hsub2(frag_b[0], zp);
+  frag_b[1] = __hsub2(frag_b[1], zp);
+}
+
+// Same as above, but for act_order (each K is multiplied individually)
+template <vllm::ScalarTypeId type_id>
+__device__ inline void scale4(
+    typename MarlinScalarType<type_id>::FragB& frag_b,
+    typename MarlinScalarType<type_id>::FragS& frag_s_1,
+    typename MarlinScalarType<type_id>::FragS& frag_s_2,
+    typename MarlinScalarType<type_id>::FragS& frag_s_3,
+    typename MarlinScalarType<type_id>::FragS& frag_s_4, int i) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
+
+  scalar_t2 s_val_1_2;
+  s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
+  s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
+
+  scalar_t2 s_val_3_4;
+  s_val_3_4.x = reinterpret_cast<scalar_t*>(&frag_s_3)[i];
+  s_val_3_4.y = reinterpret_cast<scalar_t*>(&frag_s_4)[i];
+
+  frag_b[0] = __hmul2(frag_b[0], s_val_1_2);
+  frag_b[1] = __hmul2(frag_b[1], s_val_3_4);
+}
+
+// Given 2 floats multiply by 2 scales (halves)
+template <vllm::ScalarTypeId type_id>
+__device__ inline void scale_float(
+    float* c, typename MarlinScalarType<type_id>::FragS& s) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s);
+  c[0] = __fmul_rn(c[0], MarlinScalarType<type_id>::num2float(s_ptr[0]));
+  c[1] = __fmul_rn(c[1], MarlinScalarType<type_id>::num2float(s_ptr[1]));
+}
+
+// Wait until barrier reaches `count`, then lock for current threadblock.
+__device__ inline void barrier_acquire(int* lock, int count) {
+  if (threadIdx.x == 0) {
+    int state = -1;
+    do
+      // Guarantee that subsequent writes by this threadblock will be visible
+      // globally.
+      asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
+                   : "=r"(state)
+                   : "l"(lock));
+    while (state != count);
+  }
+  __syncthreads();
+}
+
+// Release barrier and increment visitation count.
+__device__ inline void barrier_release(int* lock, bool reset = false) {
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    if (reset) {
+      lock[0] = 0;
+      return;
+    }
+    int val = 1;
+    // Make sure that all writes since acquiring this barrier are visible
+    // globally, while releasing the barrier.
+    asm volatile("fence.acq_rel.gpu;\n");
+    asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
+                 :
+                 : "l"(lock), "r"(val));
+  }
+}
+
+// Wait until value of lock to be negative, and then add 1
+__device__ inline void wait_negative_and_add(int* lock) {
+  if (threadIdx.x == 0) {
+    int state = 0;
+    do
+      // Guarantee that subsequent writes by this threadblock will be visible
+      // globally.
+      asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
+                   : "=r"(state)
+                   : "l"(lock));
+    while (state >= 0);
+    atomicAdd(lock, 1);
+  }
+  __syncthreads();
+}
+
+template <const vllm::ScalarTypeId a_type_id,  // A ScalarType id
+          const vllm::ScalarTypeId b_type_id,  // B ScalarType id
+          const vllm::ScalarTypeId c_type_id,  // C ScalarType id
+          const vllm::ScalarTypeId s_type_id,  // B_SCALE ScalarType id
+          const int threads,          // number of threads in a threadblock
+          const int thread_m_blocks,  // number of 16x16 blocks in the m
+                                      // dimension (batchsize) of the
+                                      // threadblock
+          const int thread_n_blocks,  // same for n dimension (output)
+          const int thread_k_blocks,  // same for k dimension (reduction)
+          const bool m_block_size_8,  // whether m_block_size == 8
+                                      // only works when thread_m_blocks == 1
+          const int stages,  // number of stages for the async global->shared
+                             // fetch pipeline
+          const int group_blocks,  // number of consecutive 16x16 blocks
+                                   // with a separate quantization scale
+          const bool is_zp_float   // is zero point of float16 type?
+          >
+__global__ void Marlin(
+    const int4* __restrict__ A,  // fp16 input matrix of shape mxk
+    const int4* __restrict__ B,  // 4bit quantized weight matrix of shape kxn
+    int4* __restrict__ C,        // fp16 output buffer of shape mxn
+    int4* __restrict__ C_tmp,    // fp32 tmp output buffer (for reduce)
+    const int4* __restrict__ b_bias_ptr,
+    // float scales of input matrix, only used when is_a_8bit == true.
+    // shape (m,)
+    const float* __restrict__ a_scales_ptr,
+    // fp16 quantization scales. shape (k/groupsize, n)
+    const int4* __restrict__ scales_ptr,
+    // fp16 global scale (for nvfp4// only)
+    const float* __restrict__ global_scale_ptr,
+    // 4bit packed zero-points of shape
+    // (k/groupsize, n/pack_factor)
+    const int4* __restrict__ zp_ptr,
+    // int32 group indices of shape k
+    const int* __restrict__ g_idx,
+    const int32_t* __restrict__ sorted_token_ids_ptr,        // moe sorted_ids
+    const int32_t* __restrict__ expert_ids_ptr,              // moe expert ids
+    const int32_t* __restrict__ num_tokens_past_padded_ptr,  // moe num tokens
+    const float* __restrict__ topk_weights_ptr,              // moe top weights
+    int top_k,              // num of experts per token
+    bool mul_topk_weights,  // mul topk weights or not
+    int num_groups,         // number of scale groups per output channel
+    int prob_m,             // batch dimension m
+    int prob_n,             // output dimension n
+    int prob_k,             // reduction dimension k
+    int* locks,             // extra global storage for barrier synchronization
+    bool has_bias,
+    bool use_atomic_add,  // whether to use atomic add to reduce
+    bool use_fp32_reduce  // whether to use fp32 global reduce
+) {
+  // Each threadblock processes one "stripe" of the B matrix with (roughly) the
+  // same size, which might involve multiple column "slices" (of width 16 *
+  // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
+  // example:
+  //   0 1 3
+  //   0 2 3
+  //   1 2 4
+  // While this kind of partitioning makes things somewhat more complicated, it
+  // ensures good utilization of all SMs for many kinds of shape and GPU
+  // configurations, while requiring as few slow global cross-threadblock
+  // reductions as possible.
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890
+  // FP8 computation is only supported for Ada Lovelace or newer architectures.
+  if constexpr (a_type_id == vllm::kFE4M3fn.id()) return;
+  #endif
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+  // Turing TensorCore only supports fp16 and int8
+  if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id())
+    return;
+  #endif
+
+  int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
+  constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+  static constexpr auto num_bits =
+      vllm::ScalarType::from_id(b_type_id).size_bits();
+  // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 &&
+  // num_bits == 4
+  constexpr bool use_fp16_accum =
+      a_type_id == vllm::kFloat16.id() &&
+      (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) &&
+       !(group_blocks == -1 && num_bits == 4));
+  #else
+  constexpr bool use_fp16_accum = false;
+  #endif
+  using Adtype = MarlinScalarType<a_type_id>;
+  using Cdtype = MarlinScalarType<c_type_id>;
+
+  using scalar_t = typename MarlinScalarType<a_type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<a_type_id>::scalar_t2;
+  using scalar_32bit_t = typename MarlinScalarType<a_type_id>::scalar_32bit_t;
+
+  using c_scalar_t = typename MarlinScalarType<c_type_id>::scalar_t;
+  using c_scalar_t2 = typename MarlinScalarType<c_type_id>::scalar_t2;
+
+  using FragA = typename MarlinScalarType<a_type_id>::FragA;
+  using FragB = typename MarlinScalarType<a_type_id>::FragB;
+  using FragC = typename MarlinScalarType<a_type_id>::FragC;
+  using FragS = typename MarlinScalarType<c_type_id>::FragS;
+  using FragZP = typename MarlinScalarType<c_type_id>::FragZP;
+
+  extern __shared__ int4 sh[];
+  static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id);
+  static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id);
+  static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id);
+  static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id);
+  if constexpr (b_type == vllm::kFE2M1f) {
+    static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 ||
+                  s_type == vllm::kFE8M0fnu && group_blocks == 2);
+  } else if constexpr (b_type == vllm::kFE4M3fn && s_type == vllm::kFE8M0fnu) {
+    static_assert(group_blocks == 2);
+  } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
+    static_assert(s_type == vllm::kBFloat16);
+  } else if constexpr (std::is_same<scalar_t, half>::value) {
+    static_assert(s_type == vllm::kFloat16);
+  }
+
+  constexpr bool is_a_8bit = a_type.size_bits() == 8;
+  if constexpr (!is_a_8bit) {
+    static_assert(std::is_same<scalar_t, c_scalar_t>::value);
+  }
+  constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8;
+  constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 ||
+                               b_type == vllm::kS4 || b_type == vllm::kS8 ||
+                               b_type == vllm::kU4B8 || b_type == vllm::kU8B128;
+  constexpr bool is_8bit_scale = s_type.size_bits() == 8;
+  // see comments of dequant.h for more details
+  constexpr bool dequant_skip_flop =
+      is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) ||
+      b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
+      has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
+      has_zp && !is_zp_float && !(b_type == vllm::kU8);
+
+  float global_scale_f32 = 1.0f;
+
+  constexpr bool has_act_order = group_blocks == 0;
+
+  constexpr int pack_factor = 32 / b_type.size_bits();
+  static_assert(thread_m_blocks == 1 || !m_block_size_8);
+  const int group_size =
+      (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups;
+  const int scales_expert_stride =
+      prob_n * prob_k / group_size / (is_8bit_scale ? 16 : 8);
+  const int zp_expert_stride =
+      is_zp_float ? prob_n * prob_k / group_size / 8
+                  : prob_n * prob_k / group_size / (pack_factor * 4);
+  const int b_bias_expert_stride = prob_n / 8;
+
+  // parallel: num valid moe blocks
+  int parallel = num_tokens_past_padded / moe_block_size;
+
+  int k_tiles = prob_k / 16 / thread_k_blocks;
+  int n_tiles = prob_n / 16 / thread_n_blocks;
+
+  int global_mn_tiles = parallel * n_tiles;
+  int part2_mn_tiles = global_mn_tiles;
+  int part1_mn_iters = 0;
+  bool in_part2 = false;
+
+  // we use DP + two-tile SK here
+  // part1: DP
+  // part2: two-tile SK
+  // see https://github.com/vllm-project/vllm/pull/24722 for more details
+  if (global_mn_tiles > gridDim.x) {
+    part2_mn_tiles = global_mn_tiles % gridDim.x;
+    if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x;
+    part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x;
+  }
+
+  int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x);
+
+  if constexpr (!has_act_order && group_blocks != -1) {
+    if (group_blocks >= thread_k_blocks) {
+      // Ensure that the number of tiles in each stripe is a multiple of the
+      // groupsize; this avoids an annoying special case where a stripe starts
+      // in the middle of group.
+      iters = (group_blocks / thread_k_blocks) *
+              div_ceil(iters, (group_blocks / thread_k_blocks));
+    }
+  }
+
+  int slice_row = 0;
+  int slice_col_par = blockIdx.x;
+  int slice_col;
+  int slice_iters =
+      k_tiles;  // number of threadblock tiles in the current slice
+  // total number of active threadblocks in the current slice
+  int slice_count = 1;
+  // index of threadblock in current slice; numbered bottom to top
+  int slice_idx = 0;
+
+  int par_id = 0;
+  int block_id = -1;
+  int64_t expert_id = 0;  // use int64 to avoid computation result overflow
+  int old_expert_id = 0;
+  int64_t B_expert_off = 0;
+
+  float* sh_a_s = reinterpret_cast<float*>(sh);
+  int4* sh_block_sorted_ids_int4 = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0);
+  int4* sh_rd_block_sorted_ids_int4 =
+      sh_block_sorted_ids_int4 + moe_block_size / 4;
+  int4* sh_block_topk_weights_int4 =
+      sh_rd_block_sorted_ids_int4 + moe_block_size / 4;
+  // sh_block_topk_weights_int4 only need (moe_block_size / 4);
+  // but we pad to align to 256 bytes
+  int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 2;
+  int32_t* sh_block_sorted_ids =
+      reinterpret_cast<int*>(sh_block_sorted_ids_int4);
+  int32_t* sh_rd_block_sorted_ids =
+      reinterpret_cast<int*>(sh_rd_block_sorted_ids_int4);
+  c_scalar_t2* sh_block_topk_weights =
+      reinterpret_cast<c_scalar_t2*>(sh_block_topk_weights_int4);
+
+  int32_t block_num_valid_tokens = 0;
+  int32_t locks_off = 0;
+
+  // We can easily implement parallel problem execution by just remapping
+  // indices and advancing global pointers
+  if (part2_mn_tiles >= gridDim.x) {
+    // when part2_mn_tiles >= sms
+    // then there are at most $sms$ conflict tile blocks
+    locks_off = blockIdx.x;
+  } else {
+    locks_off = (iters * blockIdx.x) / k_tiles - 1;
+  }
+
+  int prob_m_top_k = prob_m * top_k;
+  // read moe block data given block_id
+  // block_sorted_ids / block_num_valid_tokens / block_topk_weights
+  auto read_moe_block_data = [&](int block_id) {
+    block_num_valid_tokens = moe_block_size;
+
+    cp_async4_pred(sh_block_sorted_ids_int4 + threadIdx.x,
+                   reinterpret_cast<const int4*>(sorted_token_ids_ptr) +
+                       (block_id * moe_block_size / 4 + threadIdx.x),
+                   threadIdx.x < moe_block_size / 4);
+
+    cp_async_fence();
+    cp_async_wait<0>();
+
+    __syncthreads();
+
+    if (threadIdx.x >= threads - 32) {
+      constexpr int size_per_thread = div_ceil(moe_block_size, 32);
+      int lane_id = threadIdx.x - (threads - 32);
+
+      int local_count = 0;
+  #pragma unroll
+      for (int i = 0; i < size_per_thread; i++) {
+        int j = lane_id * size_per_thread + i;
+        if (j < moe_block_size) {
+          int idx = sh_block_sorted_ids[j];
+          if (idx < prob_m_top_k) local_count++;
+        }
+      }
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+
+      if constexpr (moe_block_size >= 16)
+        local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 16);
+      if constexpr (moe_block_size >= 8)
+        local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 8);
+      if constexpr (moe_block_size >= 4)
+        local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 4);
+      if constexpr (moe_block_size >= 2)
+        local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 2);
+
+      local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 1);
+      block_num_valid_tokens = local_count;
+  #else
+      block_num_valid_tokens = __reduce_add_sync(0xffffffff, local_count);
+  #endif
+
+      if (lane_id == 0)
+        reinterpret_cast<int*>(sh_new)[0] = block_num_valid_tokens;
+    }
+
+    if (threadIdx.x < moe_block_size) {
+      int idx = sh_block_sorted_ids[threadIdx.x];
+      sh_rd_block_sorted_ids[threadIdx.x] = idx / top_k;
+
+      if (mul_topk_weights) {
+        idx = idx < prob_m_top_k ? idx : 0;
+        float topk_weight_tmp = topk_weights_ptr[idx];
+        if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
+          topk_weight_tmp *= global_scale_f32;
+        }
+        c_scalar_t2 topk_weight_val =
+            Cdtype::num2num2(Cdtype::float2num(topk_weight_tmp));
+        sh_block_topk_weights[threadIdx.x] = topk_weight_val;
+      }
+    }
+
+    __syncthreads();
+
+    block_num_valid_tokens = reinterpret_cast<int*>(sh_new)[0];
+    __syncthreads();
+  };
+
+  // when move to next moe block, find the next block_id and expert_id
+  // and then read moe block data
+  auto update_next_moe_block_data = [&]() {
+    if (par_id >= parallel) return;
+
+    old_expert_id = expert_id;
+    block_id = par_id;
+    expert_id = expert_ids_ptr[block_id];
+
+    if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
+      global_scale_f32 = global_scale_ptr[expert_id];
+    }
+
+    B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4);
+    scales_ptr += (expert_id - old_expert_id) * scales_expert_stride;
+    if constexpr (has_zp) {
+      zp_ptr += (expert_id - old_expert_id) * zp_expert_stride;
+    }
+    if constexpr (has_act_order) {
+      g_idx += (expert_id - old_expert_id) * prob_k;
+    }
+    if (has_bias) {
+      b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride;
+    }
+
+    read_moe_block_data(block_id);
+  };
+
+  // Compute all information about the current slice which is required for
+  // synchronization.
+  bool first_init = true;
+  auto init_part2_slice = [&]() {
+    slice_iters =
+        iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
+    if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) slice_iters = 0;
+    if (slice_iters == 0) return;
+    if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
+    slice_count = 1;
+    slice_idx = 0;
+    int col_first = iters * div_ceil(k_tiles * slice_col_par, iters);
+    if (col_first <= k_tiles * (slice_col_par + 1)) {
+      int col_off = col_first - k_tiles * slice_col_par;
+      slice_count = div_ceil(k_tiles - col_off, iters);
+      if (col_off > 0) slice_count++;
+      int delta_first = iters * blockIdx.x - col_first;
+      if (delta_first < 0 || (col_off == 0 && delta_first == 0))
+        slice_idx = slice_count - 1;
+      else {
+        slice_idx = slice_count - 1 - delta_first / iters;
+        if (col_off > 0) slice_idx--;
+      }
+    }
+    if (part2_mn_tiles >= gridDim.x) {
+      if (slice_count > 1 && slice_idx == slice_count - 1) {
+        locks_off++;
+      }
+    } else {
+      locks_off++;
+    }
+
+    if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) {
+      constexpr int threads_per_m = 16 * thread_n_blocks / 8;
+      int m_per_thread =
+          div_ceil(block_num_valid_tokens, threads / threads_per_m);
+      for (int i = 0; i < m_per_thread; i++) {
+        int row = threads / threads_per_m * i + threadIdx.x / threads_per_m;
+        if (row < block_num_valid_tokens) {
+          int64_t sorted_row = sh_block_sorted_ids[row];
+          int col = slice_col * 16 * thread_n_blocks / 8 +
+                    threadIdx.x % threads_per_m;
+          C[sorted_row * prob_n / 8 + col] = {0, 0, 0, 0};
+        }
+      }
+      // After write zero to output, write a negative value to lock.
+      // Every SM that processes the same slice would wait for
+      // the negative value, and then atomicAdd 1 to it.
+      // After all SMs are processed, the lock value would back to 0 again.
+      __syncthreads();
+      if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count;
+    }
+
+    if (slice_col == n_tiles) {
+      slice_col = 0;
+      par_id++;
+      update_next_moe_block_data();
+    }
+    if (is_a_8bit && (first_init || slice_col == 0)) {
+      __syncthreads();
+      cp_async1_ca_pred(&sh_a_s[threadIdx.x],
+                        &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]],
+                        threadIdx.x < block_num_valid_tokens);
+    }
+  };
+
+  auto init_part1_slice = [&]() {
+    if (part1_mn_iters) {
+      part1_mn_iters--;
+      par_id = slice_col_par / n_tiles;
+      slice_col = slice_col_par % n_tiles;
+      slice_iters = k_tiles;
+      update_next_moe_block_data();
+      if (is_a_8bit) {
+        __syncthreads();
+        cp_async1_ca_pred(&sh_a_s[threadIdx.x],
+                          &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]],
+                          threadIdx.x < block_num_valid_tokens);
+      }
+    }
+  };
+
+  auto init_slice = [&]() {
+    if (!in_part2 && !part1_mn_iters) {
+      in_part2 = true;
+      slice_col_par = (iters * blockIdx.x) / k_tiles;
+      slice_row = (iters * blockIdx.x) % k_tiles;
+      slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles;
+      par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles;
+      update_next_moe_block_data();
+    }
+    if (!in_part2) {
+      init_part1_slice();
+    } else {
+      init_part2_slice();
+      first_init = false;
+    }
+  };
+
+  init_slice();
+
+  // A sizes/strides
+
+  // stride of the A matrix in global memory
+  int a_gl_stride = prob_k / (is_a_8bit ? 16 : 8);
+  // stride of an A matrix tile in shared memory
+  constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8);
+  // delta between subsequent A tiles in global memory
+  constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8);
+  // between subsequent accesses within a tile
+  int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);
+  // between shared memory writes
+  constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);
+  // within a shared memory tile
+  constexpr int a_sh_rd_delta_i = a_sh_stride * 16;
+  // overall size of a tile
+  constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks);
+  // number of shared write iterations for a tile
+  constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta);
+
+  // B sizes/strides
+  int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4));
+  constexpr int b_sh_stride =
+      ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4);
+  constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2;
+  constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;
+
+  int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
+  constexpr int b_sh_wr_delta = threads * b_thread_vecs;
+  constexpr int b_sh_stage =
+      b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
+  constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
+
+  // Scale sizes/strides without act_order
+  int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8);
+  constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8);
+  constexpr int s_tb_groups =
+      !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
+          ? thread_k_blocks / group_blocks
+          : 1;
+  constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
+  int s_gl_rd_delta = s_gl_stride;
+
+  // Scale size/strides with act_order
+  constexpr int tb_k = 16 * thread_k_blocks;
+  constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;
+  // constexpr int act_s_row_stride      = 1;
+  // int           act_s_col_stride      = act_s_row_stride * num_groups;
+  constexpr int act_s_max_num_groups = 32;
+  int act_s_col_stride = 1;
+  int act_s_col_warp_stride = act_s_col_stride * 8;
+
+  constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4);
+  int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
+
+  // Zero-points sizes/strides
+  int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4;
+  constexpr int zp_sh_stride = is_zp_float
+                                   ? 16 * thread_n_blocks / 8
+                                   : ((16 * thread_n_blocks) / pack_factor) / 4;
+  constexpr int zp_tb_groups = s_tb_groups;
+  constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0;
+  int zp_gl_rd_delta = zp_gl_stride;
+
+  // Global A read index of current thread.
+  int a_gl_rd_row = threadIdx.x / a_gl_rd_delta_o;
+  int a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o;
+  // Shared write index of current thread.
+  int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
+                (threadIdx.x % a_gl_rd_delta_o);
+  // Shared read index.
+  int a_sh_rd =
+      a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) +
+      (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1));
+  a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters;
+
+  int b_gl_rd;
+  if (threads <= b_sh_stride) {
+    b_gl_rd = threadIdx.x;
+  } else {
+    b_gl_rd =
+        b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride);
+  }
+
+  b_gl_rd += B_expert_off + b_sh_stride * slice_col;
+  b_gl_rd += b_gl_rd_delta_o * slice_row;
+  auto b_sh_rd = threadIdx.x * b_thread_vecs;
+  b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1));
+
+  // For act_order
+  int slice_k_start = tb_k * slice_row;
+  int slice_k_finish = slice_k_start + tb_k * slice_iters;
+  int slice_k_start_shared_fetch = slice_k_start;
+  int slice_n_offset = act_s_col_tb_stride * slice_col;
+
+  // No act_order
+  int s_gl_rd;
+  if constexpr (!has_act_order) {
+    if constexpr (group_blocks == -1) {
+      s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
+    } else if constexpr (group_blocks >= thread_k_blocks) {
+      s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
+                s_sh_stride * slice_col + threadIdx.x;
+    } else {
+      s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
+                               threadIdx.x / s_sh_stride) +
+                s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
+    }
+  }
+  auto s_sh_wr = threadIdx.x;
+  bool s_sh_wr_pred = threadIdx.x < s_sh_stage;
+
+  // Zero-points
+  int zp_gl_rd;
+  if constexpr (has_zp) {
+    if constexpr (group_blocks == -1) {
+      zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
+    } else if constexpr (group_blocks >= thread_k_blocks) {
+      zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
+                 zp_sh_stride * slice_col + threadIdx.x;
+    } else {
+      zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
+                                 threadIdx.x / zp_sh_stride) +
+                 zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
+    }
+  }
+  auto zp_sh_wr = threadIdx.x;
+  bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage;
+
+  // We use a different scale layout for grouped and column-wise quantization as
+  // we scale a `half2` tile in column-major layout in the former and in
+  // row-major in the latter case.
+  int s_sh_rd;
+  if constexpr (is_a_8bit) {
+    s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4);
+  } else if constexpr (group_blocks != -1)
+    s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4;
+  else if constexpr (group_blocks == -1 &&
+                     (m_block_size_8 || (has_zp && !dequant_skip_flop)))
+    s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8;
+  else
+    s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4;
+
+  int bias_sh_rd;
+  if constexpr (m_block_size_8) {
+    bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8;
+  } else {
+    bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) +
+                 (threadIdx.x % 32) % 4;
+  }
+
+  int bias_sh_wr = threadIdx.x;
+  int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
+
+  // Zero-points have the same read layout as the scales
+  // (without column-wise case)
+  constexpr int num_col_threads = 8;
+  constexpr int num_row_threads = 4;
+  constexpr int num_ints_per_thread = 8 / pack_factor;
+  int zp_sh_rd;
+  if constexpr (has_zp) {
+    if constexpr (is_zp_float) {
+      if constexpr (group_blocks != -1) {
+        zp_sh_rd =
+            8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4;
+      }
+    } else if (is_a_8bit) {
+      zp_sh_rd = num_ints_per_thread * num_col_threads *
+                     ((threadIdx.x / 32) % tb_n_warps / 2) +
+                 num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
+    } else {
+      zp_sh_rd = num_ints_per_thread * num_col_threads *
+                     ((threadIdx.x / 32) % tb_n_warps) +
+                 num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
+    }
+  }
+
+  // To ensure that writing and reading A tiles to/from shared memory, the
+  // latter in fragment format, is fully bank conflict free, we need to use a
+  // rather fancy XOR-based layout. The key here is that neither reads nor
+  // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the
+  // same shared memory banks. Further, it seems (based on NSight-Compute) that
+  // each warp must also write a consecutive memory segment?
+  auto transform_a = [&](int i) {
+    int row = i / a_gl_rd_delta_o;
+    return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8);
+  };
+  // Since the computation of this remapping is non-trivial and, due to our main
+  // loop unrolls, all shared memory accesses are static, we simply precompute
+  // both transformed reads and writes.
+  int a_sh_wr_trans[a_sh_wr_iters];
+  #pragma unroll
+  for (int i = 0; i < a_sh_wr_iters; i++)
+    a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
+  int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];
+  #pragma unroll
+  for (int i = 0; i < b_sh_wr_iters; i++) {
+  #pragma unroll
+    for (int j = 0; j < thread_m_blocks; j++)
+      a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd);
+  }
+
+  // Since B-accesses have non-constant stride they have to be computed at
+  // runtime; we break dependencies between subsequent accesses with a tile by
+  // maintining multiple pointers (we have enough registers), a tiny
+  // optimization.
+
+  // Shared memory storage for global fetch pipelines.
+  constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks;
+  constexpr int sh_b_size = stages * b_sh_stage;
+  int4* sh_b = sh_new;
+  int4* sh_red = sh_new;
+
+  constexpr int sh_size_b_red_min =
+      (sh_red_size < sh_b_size ? sh_red_size : sh_b_size);
+  constexpr int sh_size_b_red_max =
+      (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
+  constexpr int sh_bias_size = (thread_n_blocks * 16 / 8);
+  constexpr int sh_b_red_bias_size =
+      sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size)
+          ? sh_size_b_red_max
+          : (sh_size_b_red_min + sh_bias_size);
+
+  int4* sh_bias = sh_new + sh_size_b_red_min;
+  int4* sh_g_idx = sh_new + sh_b_red_bias_size;
+  int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
+  constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
+                                          : (stages * s_sh_stage);
+  int4* sh_s = sh_zp + (stages * zp_sh_stage);
+  int4* sh_a = sh_s + sh_s_size;
+
+  // Register storage for double buffer of shared memory reads.
+  FragA frag_a[2][thread_m_blocks];
+  I4 frag_b_quant[2][b_thread_vecs];
+  FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2];
+  FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2];
+  FragS frag_s[2][4];  // No act-order
+  FragS frag_bias[2][4];
+  FragS act_frag_s[2][4][4];             // For act-order
+  int frag_qzp[2][num_ints_per_thread];  // Zero-points
+  FragZP frag_zp;                        // Zero-points in fp16
+  FragZP frag_zpf[2];                    // Zero-points in fp16 in HQQ
+
+  if constexpr (is_a_8bit && group_blocks != -1) {
+  #pragma unroll
+    for (int j = 0; j < 2; j++) {
+  #pragma unroll
+      for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+        for (int g = 0; g < 4; g++) {
+          frag_c_tmp[i][j][0][g] = 0.0f;
+        }
+
+  #pragma unroll
+        for (int g = 0; g < 4; g++) {
+          frag_c_tmp[i][j][1][g] = 0.0f;
+        }
+      }
+    }
+  }
+
+  // Zero accumulators.
+  auto zero_accums = [&]() {
+  #pragma unroll
+    for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
+      reinterpret_cast<float*>(frag_c)[i] = 0;
+  };
+
+  int sh_first_group_id = -1;
+  int sh_num_groups = -1;
+
+  auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id,
+                                              int last_group_id) {
+    sh_first_group_id = first_group_id;
+    sh_num_groups = last_group_id - first_group_id + 1;
+
+    if (sh_num_groups > act_s_max_num_groups) {
+      sh_num_groups = act_s_max_num_groups;
+    }
+
+    if (sh_first_group_id + sh_num_groups > num_groups) {
+      sh_num_groups = num_groups - sh_first_group_id;
+    }
+
+    int row_offset = first_group_id * s_gl_stride;
+
+    if (is_async) {
+      for (int i = 0; i < sh_num_groups; i++) {
+        if (threadIdx.x < s_sh_stride) {
+          cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
+                         &scales_ptr[row_offset + (i * s_gl_stride) +
+                                     slice_n_offset + threadIdx.x]);
+        }
+      }
+    } else {
+      for (int i = 0; i < sh_num_groups; i++) {
+        if (threadIdx.x < s_sh_stride) {
+          sh_s[(i * s_sh_stride) + threadIdx.x] =
+              scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
+                         threadIdx.x];
+        }
+      }
+    }
+  };
+  // Asynchronously fetch the next A, B and s tile from global to the next
+  // shared memory pipeline location.
+  auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
+    if (pred) {
+      int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe;
+  #pragma unroll
+      for (int i = 0; i < a_sh_wr_iters; i++) {
+        int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row;
+        int64_t sorted_row = 0;
+        if (!m_block_size_8 || row < 8)
+          sorted_row = sh_rd_block_sorted_ids[row];
+        int64_t true_idx =
+            sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off;
+        cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx],
+                       row < block_num_valid_tokens);
+      }
+
+      int4* sh_b_stage = sh_b + b_sh_stage * pipe;
+  #pragma unroll
+      for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) {
+        constexpr int count = div_ceil(b_sh_stride, threads);
+        int b_gl_idx =
+            b_gl_rd + (i % count) * threads +
+            b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride);
+
+        cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]);
+      }
+
+      b_gl_rd += b_gl_rd_delta_o;
+
+      if constexpr (has_act_order) {
+        // Fetch g_idx thread-block portion
+        int full_pipe = a_off;
+        int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
+        if (cur_k < prob_k && cur_k < slice_k_finish) {
+          int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
+
+          int4 const* cur_g_idx_stage_ptr =
+              reinterpret_cast<int4 const*>(&g_idx[cur_k]);
+
+          if (threadIdx.x < g_idx_stage) {
+            cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
+                           &cur_g_idx_stage_ptr[threadIdx.x]);
+          }
+        }
+      } else {
+        if constexpr (group_blocks != -1) {
+          int4* sh_s_stage = sh_s + s_sh_stage * pipe;
+
+          // Only fetch scales if this tile starts a new group
+          if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
+            if (s_sh_wr_pred) {
+              cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
+            }
+            s_gl_rd += s_gl_rd_delta * s_tb_groups;
+          }
+        }
+
+        if constexpr (has_zp && group_blocks != -1) {
+          int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
+
+          // Only fetch zero points if this tile starts a new group
+          if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
+            if (zp_sh_wr_pred) {
+              cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
+            }
+            zp_gl_rd += zp_gl_rd_delta * zp_tb_groups;
+          }
+        }
+      }
+    }
+    // Insert a fence even when we are winding down the pipeline to ensure that
+    // waiting is also correct at this point.
+    cp_async_fence();
+  };
+
+  auto fetch_col_zp_to_shared = [&]() {
+    if (zp_sh_wr_pred) {
+      cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]);
+    }
+  };
+
+  auto fetch_col_scale_to_shared = [&]() {
+    if (s_sh_wr_pred) {
+      cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
+    }
+  };
+
+  // Wait until the next thread tile has been loaded to shared memory.
+  auto wait_for_stage = [&]() {
+    // We only have `stages - 2` active fetches since we are double buffering
+    // and can only issue the next fetch when it is guaranteed that the previous
+    // shared memory load is fully complete (as it may otherwise be
+    // overwritten).
+    cp_async_wait<stages - 2>();
+    __syncthreads();
+  };
+
+  // Load the next sub-tile from the current location in the shared memory pipe
+  // into the current register buffer.
+  auto fetch_to_registers = [&](int k, int pipe) {
+    int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe;
+  #pragma unroll
+    for (int i = 0; i < thread_m_blocks; i++)
+      ldsm<m_block_size_8 ? 2 : 4, a_type_id>(
+          frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
+    int4* sh_b_stage = sh_b + b_sh_stage * pipe;
+
+  #pragma unroll
+    for (int i = 0; i < b_thread_vecs; i++) {
+      frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
+          &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]);
+    }
+  };
+
+  bool is_same_group[stages];
+  int same_group_id[stages];
+
+  auto init_same_group = [&](int pipe) {
+    if constexpr (!has_act_order) {
+      return;
+    }
+
+    int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
+    int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
+
+    int group_id_1 = sh_g_idx_int_ptr[0];
+    int group_id_2 = sh_g_idx_int_ptr[tb_k - 1];
+
+    is_same_group[pipe] = group_id_1 == group_id_2;
+    same_group_id[pipe] = group_id_1;
+  };
+
+  auto fetch_scales_to_registers = [&](int k, int full_pipe) {
+    int pipe = full_pipe % stages;
+    using IT1 = typename std::conditional_t<is_a_8bit, int2, int4>;
+    using IT0 = typename std::conditional_t<is_a_8bit, int, int2>;
+    constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1);
+
+    if constexpr (!has_act_order) {
+      // No act-order case
+      if constexpr (group_blocks == -1) {
+        // load only when starting a new slice
+        if (k == 0 && full_pipe == 0 && dequant_skip_flop) {
+          reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
+          reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
+        }
+      } else if constexpr (group_blocks != -1) {
+        if constexpr (group_blocks >= thread_k_blocks) {
+          constexpr int g = group_blocks / thread_k_blocks;
+          if (pipe % g == 0) {
+            if (k % b_sh_wr_iters == 0) {
+              int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g));
+              reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
+            } else {
+              reinterpret_cast<int4*>(&frag_s[1])[0] =
+                  reinterpret_cast<int4*>(&frag_s[0])[0];
+            }
+          }
+        } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
+          auto warp_id = threadIdx.x / 32;
+          int warp_row = warp_id / tb_n_warps;
+
+          int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
+          int cur_group_id = k_blocks / group_blocks2;
+
+          int4* sh_s_stage = sh_s + s_sh_stage * pipe;
+
+          if constexpr (!is_8bit_scale) {
+            reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
+                sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
+          } else {
+            reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
+                reinterpret_cast<int2*>(
+                    sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
+          }
+        } else if (group_blocks >= b_sh_wr_iters) {
+          if constexpr (!is_8bit_scale) {
+            reinterpret_cast<int4*>(&frag_s[1])[0] =
+                reinterpret_cast<int4*>(&frag_s[0])[0];
+          } else {
+            reinterpret_cast<int2*>(&frag_s[1])[0] =
+                reinterpret_cast<int2*>(&frag_s[0])[0];
+          }
+        }
+      }
+
+      return;
+    }
+
+    // Act-order case
+
+    // Determine K of the "current" thread-block
+    int cur_k = slice_k_start + tb_k * full_pipe;
+    if (cur_k >= prob_k || cur_k >= slice_k_finish) {
+      return;
+    }
+
+    // Reset (to current thread-block) since we read g_idx portion from the
+    // shared memory
+    cur_k = 0;
+
+    // Progress to current iteration
+    cur_k += k % b_sh_wr_iters;
+
+    // Determine "position" inside the thread-block (based on warp and
+    // thread-id)
+    auto warp_id = threadIdx.x / 32;
+    int warp_row = warp_id / tb_n_warps;
+    int warp_col = warp_id % tb_n_warps;
+
+    cur_k += warp_row * 16 * b_sh_wr_iters;
+
+    auto th_id = threadIdx.x % 32;
+    cur_k += (th_id % 4) * 2;  // Due to tensor-core layout for fp16 B matrix
+
+    int s_col_shift =
+        /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) +
+        (th_id / 4) * act_s_col_stride;
+
+    if (is_same_group[pipe]) {
+      if (k % 2 == 0) {
+        *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
+            sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride +
+                 s_col_shift];
+      } else {
+        *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
+            *(reinterpret_cast<int4*>(&(act_frag_s[(k - 1) % 2][0][0])));
+      }
+
+      for (int i = 1; i < 4; i++) {
+        *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
+            *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0])));
+      }
+      return;
+    }
+
+    int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
+    int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
+
+    constexpr int k_frag_offsets[4] = {0, 1, 8,
+                                       9};  // Tensor core offsets per thread
+
+  #pragma unroll
+    for (int i = 0; i < 4; i++) {
+      int actual_k = cur_k + k_frag_offsets[i];
+
+      int group_id = sh_g_idx_int_ptr[actual_k];
+      int rel_group_id = group_id - sh_first_group_id;
+
+      *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
+          sh_s[rel_group_id * s_sh_stride + s_col_shift];
+    }
+  };
+
+  auto fetch_zp_to_registers = [&](int k, int full_pipe) {
+    // This code does not handle group_blocks == 0,
+    // which signifies act_order.
+    // has_zp implies AWQ, which doesn't have act_order,
+    static_assert(!has_zp || group_blocks != 0);
+
+    if constexpr (has_zp && !is_zp_float) {
+      int pipe = full_pipe % stages;
+
+      if constexpr (group_blocks == -1) {
+        // load only when starting a new slice
+        if (k == 0 && full_pipe == 0 || is_a_8bit) {
+  #pragma unroll
+          for (int i = 0; i < num_ints_per_thread; i++) {
+            frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];
+          }
+        }
+      } else if constexpr (group_blocks >= thread_k_blocks) {
+        constexpr int g = group_blocks / thread_k_blocks;
+        if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) {
+          int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g));
+  #pragma unroll
+          for (int i = 0; i < num_ints_per_thread; i++) {
+            frag_qzp[k % 2][i] =
+                (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
+          }
+        }
+      } else {
+        auto warp_id = threadIdx.x / 32;
+
+        int warp_row = warp_id / tb_n_warps;
+
+        int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
+        int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1);
+
+        int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
+
+        sh_zp_stage += cur_group_id * zp_sh_stride;
+
+  #pragma unroll
+        for (int i = 0; i < num_ints_per_thread; i++) {
+          frag_qzp[k % 2][i] =
+              (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
+        }
+      }
+    }
+
+    else if constexpr (has_zp && is_zp_float) {
+      int pipe = full_pipe % stages;
+
+      if constexpr (group_blocks != -1) {
+        if constexpr (group_blocks >= thread_k_blocks) {
+          constexpr int g = group_blocks / thread_k_blocks;
+          if (pipe % g == 0 && k % b_sh_wr_iters == 0) {
+            int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g));
+            reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
+                sh_zp_stage[zp_sh_rd];
+          }
+        } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
+          auto warp_id = threadIdx.x / 32;
+
+          int warp_row = warp_id / tb_n_warps;
+          int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
+          int cur_group_id = k_blocks / group_blocks;
+
+          int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
+
+          reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
+              sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride];
+        }
+      }
+    }
+  };
+
+  auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) {
+    if constexpr (a_type.size_bits() != b_type.size_bits()) {
+      if constexpr (is_a_8bit && has_zp) {
+        sub_zp_and_dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(
+            q, frag_b_ptr, zp);
+      } else {
+        dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(q, frag_b_ptr);
+      }
+    }
+  };
+
+  // Execute the actual tensor core matmul of a sub-tile.
+  bool is_first_matmul_in_slice = true;
+  auto matmul = [&](int k, int pipe) {
+    if (is_a_8bit) return;
+    int k2 = k % 2;
+    constexpr int g =
+        group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1;
+    const bool is_new_zp =
+        (group_blocks == 0) ||
+        ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) &&
+            (pipe % g == 0) ||
+        (group_blocks == -1 && is_first_matmul_in_slice);
+    if constexpr (has_zp && !is_zp_float) {
+      if (is_new_zp) {
+        if constexpr (group_blocks == -1) is_first_matmul_in_slice = false;
+        int zp_quant_0, zp_quant_1;
+
+        if constexpr (b_type.size_bits() == 4) {
+          zp_quant_0 = frag_qzp[k2][0];
+          zp_quant_1 = zp_quant_0 >> 8;
+        } else {
+          static_assert(b_type.size_bits() == 8);
+          zp_quant_0 = frag_qzp[k2][0];
+          zp_quant_1 = frag_qzp[k2][1];
+        }
+
+        dequant_data(zp_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_zp));
+        dequant_data(zp_quant_1,
+                     reinterpret_cast<scalar_32bit_t*>(&frag_zp) + 2);
+      }
+    }
+    if constexpr (!dequant_skip_flop && has_zp && is_zp_float) {
+      if (is_new_zp) {
+        reinterpret_cast<int4*>(&frag_zp)[0] =
+            reinterpret_cast<int4*>(&frag_zpf[k2])[0];
+      }
+    }
+
+    if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) {
+      int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
+      int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
+
+      dequant_fp8_scales<c_scalar_t2, s_type_id>(
+          s_quant_0, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]));
+      dequant_fp8_scales<c_scalar_t2, s_type_id>(
+          s_quant_1, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]) + 2);
+    }
+
+  // We have the m dimension as the inner loop in order to encourage overlapping
+  // dequantization and matmul operations.
+  #pragma unroll
+    for (int j = 0; j < 4; j++) {
+      FragB frag_b0;
+      FragB frag_b1;
+      int b_quant_0, b_quant_1;
+
+      if constexpr (b_type_id == vllm::kFE2M1f.id()) {
+        b_quant_1 = frag_b_quant[k2][0][j];
+        b_quant_0 = b_quant_1 << 8;
+      } else if constexpr (b_type.size_bits() == 4) {
+        b_quant_0 = frag_b_quant[k2][0][j];
+        b_quant_1 = b_quant_0 >> 8;
+      } else {
+        static_assert(b_type.size_bits() == 8);
+        int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k2]);
+        b_quant_0 = frag_b_quant_ptr[j * 2 + 0];
+        b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
+      }
+
+      dequant_data(b_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_b0));
+      dequant_data(b_quant_1, reinterpret_cast<scalar_32bit_t*>(&frag_b1));
+
+      if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) {
+        sub_zp<a_type_id>(frag_b0, frag_zp[j], 0);
+        sub_zp<a_type_id>(frag_b1, frag_zp[j], 1);
+      }
+
+      // Apply scale to frag_b0
+      if constexpr (has_act_order && !is_a_8bit) {
+        static_assert(group_blocks != -1);
+        scale4<a_type_id>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
+                          act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
+        scale4<a_type_id>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
+                          act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
+      } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
+                           group_blocks == -1 && !is_a_8bit) {
+        int idx = (threadIdx.x / 4) % 2;
+        scalar_t2 s2 = Adtype::nums2num2(
+            reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
+            reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 1])[idx]);
+        if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
+        scale_and_sub<a_type_id>(frag_b0, s2.x, frag_zp[j].x);
+        scale_and_sub<a_type_id>(frag_b1, s2.y, frag_zp[j].y);
+      } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 &&
+                           !is_a_8bit) {
+        if (is_new_zp)
+          frag_zp[j] = __hmul2(frag_zp[j],
+                               *reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
+        scale_and_sub<a_type_id>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x);
+        scale_and_sub<a_type_id>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y);
+      } else if constexpr (group_blocks != -1 && !is_a_8bit) {
+        scale<a_type_id>(frag_b0, frag_s[k2][j], 0);
+        scale<a_type_id>(frag_b1, frag_s[k2][j], 1);
+      }
+
+  #pragma unroll
+      for (int i = 0; i < thread_m_blocks; i++) {
+        if constexpr (m_block_size_8) {
+          mma_trans<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0, frag_b1,
+                                               frag_c[i][j][0]);
+        } else {
+          mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0,
+                                         frag_c[i][j][0]);
+          mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b1,
+                                         frag_c[i][j][1]);
+        }
+      }
+    }
+  };
+
+  auto matmul_a8 = [&](int k) {
+    int k2 = k % 2;
+  #pragma unroll
+    for (int j = 0; j < 2; j++) {
+      FragB frag_b[2];
+
+      if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) {
+        dequant_data(frag_b_quant[k2][0][j * 2],
+                     reinterpret_cast<scalar_32bit_t*>(&frag_b));
+        dequant_data(frag_b_quant[k2][0][j * 2 + 1],
+                     reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2);
+      } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) {
+        int off = (threadIdx.x / 32) % 2 * 2 + j;
+        int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF;
+        dequant_data(frag_b_quant[k2][0][j * 2],
+                     reinterpret_cast<scalar_32bit_t*>(&frag_b), zp);
+        zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF;
+        dequant_data(frag_b_quant[k2][0][j * 2 + 1],
+                     reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2, zp);
+      } else {
+        reinterpret_cast<int2*>(&frag_b)[0] =
+            reinterpret_cast<int2*>(&frag_b_quant[k2][j])[0];
+        reinterpret_cast<int2*>(&frag_b)[1] =
+            reinterpret_cast<int2*>(&frag_b_quant[k2][j])[1];
+      }
+
+  #pragma unroll
+      for (int i = 0; i < thread_m_blocks; i++) {
+        mma<a_type_id, false, 32>(
+            frag_a[k2][i], frag_b[0],
+            (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]);
+        mma<a_type_id, false, 32>(
+            frag_a[k2][i], frag_b[1],
+            (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]);
+      }
+
+      if constexpr (group_blocks != -1) {
+        if (group_blocks == 2 || k == 1) {
+          if constexpr (a_type == vllm::kS8) {
+            int2 s_vals[2];
+            s_vals[0] = {
+                (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[0],
+                (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[1]};
+            s_vals[1] = {
+                (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[0],
+                (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[1]};
+
+  #pragma unroll
+            for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+              for (int g = 0; g < 4; g++) {
+                int scale = reinterpret_cast<int*>(&s_vals[0])[g % 2];
+                *reinterpret_cast<int32_t*>(&frag_c[i][j][0][g]) +=
+                    *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][0][g]) *
+                    scale;
+                frag_c_tmp[i][j][0][g] = 0.0f;
+              }
+
+  #pragma unroll
+              for (int g = 0; g < 4; g++) {
+                int scale = reinterpret_cast<int*>(&s_vals[1])[g % 2];
+                *reinterpret_cast<int32_t*>(&frag_c[i][j][1][g]) +=
+                    *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][1][g]) *
+                    scale;
+                frag_c_tmp[i][j][1][g] = 0.0f;
+              }
+            }
+          } else {
+            float2 s_vals[2];
+            if constexpr (s_type_id != vllm::kFE8M0fnu.id()) {
+              static_assert(a_type.size_bits() == 16 ||
+                            s_type.size_bits() == 16);
+              s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]);
+              s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]);
+            } else {
+              int32_t* s_vals_int = reinterpret_cast<int32_t*>(&s_vals[0]);
+              int32_t s_vals_e8m0 =
+                  *reinterpret_cast<int32_t*>(&frag_s[k2][j][0]);
+
+              s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23;
+              s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15;
+              s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7;
+              s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1;
+            }
+
+  #pragma unroll
+            for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+              for (int g = 0; g < 4; g++) {
+                float scale = reinterpret_cast<float*>(&s_vals[0])[g % 2];
+                frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale;
+                frag_c_tmp[i][j][0][g] = 0.0f;
+              }
+
+  #pragma unroll
+              for (int g = 0; g < 4; g++) {
+                float scale = reinterpret_cast<float*>(&s_vals[1])[g % 2];
+                frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale;
+                frag_c_tmp[i][j][1][g] = 0.0f;
+              }
+            }
+          }
+        }
+      }
+    }
+  };
+
+  // Since we slice across the k dimension of a tile in order to increase the
+  // number of warps while keeping the n dimension of a tile reasonable, we have
+  // multiple warps that accumulate their partial sums of the same output
+  // location; which we have to reduce over in the end. We do in shared memory.
+  auto thread_block_reduce = [&]() {
+    constexpr int red_off = threads / b_sh_stride_threads / 2;
+    if (red_off >= 1) {
+      auto red_idx = threadIdx.x / b_sh_stride_threads;
+      constexpr int red_sh_stride =
+          b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2;
+      constexpr int red_sh_delta = b_sh_stride_threads;
+      int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
+                      (threadIdx.x % b_sh_stride_threads);
+
+      // Parallel logarithmic shared memory reduction. We make sure to avoid any
+      // unnecessary read or write iterations, e.g., for two warps we write only
+      // once by warp 1 and read only once by warp 0.
+
+  #pragma unroll
+      for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
+  #pragma unroll
+        for (int i = red_off; i > 0; i /= 2) {
+          if (i <= red_idx && red_idx < 2 * i) {
+  #pragma unroll
+            for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2;
+                 j += (m_block_size_8 ? 2 : 1)) {
+              int red_sh_wr =
+                  red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
+              if (i < red_off) {
+                float* c_rd = reinterpret_cast<float*>(
+                    &sh_red[red_sh_delta * j + red_sh_rd]);
+                float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]);
+  #pragma unroll
+                for (int k = 0; k < 4; k++)
+                  reinterpret_cast<FragC*>(
+                      frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] +=
+                      c_rd[k] + c_wr[k];
+              }
+              sh_red[red_sh_wr] = reinterpret_cast<int4*>(
+                  &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j];
+            }
+          }
+          __syncthreads();
+        }
+        if (red_idx == 0) {
+  #pragma unroll
+          for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2;
+               i += (m_block_size_8 ? 2 : 1)) {
+            float* c_rd =
+                reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]);
+  #pragma unroll
+            for (int j = 0; j < 4; j++)
+              reinterpret_cast<FragC*>(
+                  frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j];
+          }
+        }
+        __syncthreads();
+      }
+    }
+  };
+
+  // Since multiple threadblocks may process parts of the same column slice, we
+  // finally have to globally reduce over the results. As the striped
+  // partitioning minimizes the number of such reductions and our outputs are
+  // usually rather small, we perform this reduction serially in L2 cache.
+  auto global_reduce_fp16 = [&](bool first = false, bool last = false) {
+    // We are very careful here to reduce directly in the output buffer to
+    // maximize L2 cache utilization in this step. To do this, we write out
+    // results in FP16 (but still reduce with FP32 compute).
+    constexpr int active_threads = 32 * tb_n_warps;
+    bool is_th_active = threadIdx.x < active_threads;
+    if (!is_th_active) {
+      return;
+    }
+
+    int c_gl_stride = prob_n / 8 * (is_a_8bit ? 2 : 1);
+    int c_gl_wr_delta_o = 8 * c_gl_stride;
+    int c_gl_wr_delta_i = 4 * (active_threads / 32);
+    int c_gl_wr;
+    if constexpr (m_block_size_8) {
+      c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) +
+                (threadIdx.x % 32) / 8;
+      c_gl_wr += (2 * thread_n_blocks) * slice_col;
+    } else {
+      c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
+                4 * (threadIdx.x / 32) + threadIdx.x % 4;
+      c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1);
+    }
+    constexpr int c_sh_wr_delta = active_threads;
+    int c_sh_wr = threadIdx.x;
+
+    if (!first) {
+
+  #pragma unroll
+      for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
+        int c_idx;
+        if constexpr (m_block_size_8)
+          c_idx = c_gl_wr + i * c_gl_stride +
+                  (threadIdx.x % 8) / 4 * c_gl_wr_delta_i;
+        else
+          c_idx =
+              c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2);
+        if (c_idx / c_gl_stride < block_num_valid_tokens) {
+          int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride];
+          int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;
+          if constexpr (is_a_8bit) {
+            int2* sh_red_int2 = reinterpret_cast<int2*>(sh_red);
+            int2* c_int2 = reinterpret_cast<int2*>(C);
+            sh_red_int2[c_sh_wr + c_sh_wr_delta * i] = c_int2[true_idx];
+          } else {
+            sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx];
+          }
+        }
+      }
+    }
+
+  #pragma unroll
+    for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
+      if (!first) {
+        c_scalar_t* c_red_f16;
+        if constexpr (is_a_8bit) {
+          int2 tmp =
+              reinterpret_cast<int2*>(sh_red)[c_sh_wr + i * c_sh_wr_delta];
+          c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp);
+        } else {
+          int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta];
+          c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp);
+        }
+  #pragma unroll
+        for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) {
+          int delta = 0;
+          if constexpr (m_block_size_8) {
+            delta = j % 2 == 1 ? -2 : 0;
+          }
+          reinterpret_cast<float*>(
+              &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) +
+                       delta] += Cdtype::num2float(c_red_f16[j]);
+        }
+      }
+      if (!last) {
+        c_scalar_t c_f16[is_a_8bit ? 4 : 8];
+  #pragma unroll
+        for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) {
+          int delta = 0;
+          if constexpr (m_block_size_8) {
+            delta = j % 2 == 1 ? -2 : 0;
+          }
+          c_f16[j] = Cdtype::float2num(reinterpret_cast<float*>(
+              &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) +
+                       delta]);
+        }
+
+        int c_idx;
+        if constexpr (m_block_size_8)
+          c_idx = c_gl_wr + i * c_gl_stride +
+                  (threadIdx.x % 8) / 4 * c_gl_wr_delta_i;
+        else
+          c_idx =
+              c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2);
+        if (c_idx / c_gl_stride < block_num_valid_tokens) {
+          int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride];
+          int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;
+          if constexpr (is_a_8bit) {
+            int2* c_int2 = reinterpret_cast<int2*>(C);
+            c_int2[true_idx] = *reinterpret_cast<int2*>(c_f16);
+          } else {
+            C[true_idx] = *reinterpret_cast<int4*>(c_f16);
+          }
+        }
+      }
+    }
+  };
+
+  // Globally reduce over threadblocks that compute the same column block.
+  // We use a tmp C buffer to reduce in full fp32 precision.
+  auto global_reduce_fp32 = [&](bool first = false, bool last = false) {
+    constexpr int tb_m = thread_m_blocks * 16;
+    constexpr int tb_n = thread_n_blocks * 16;
+
+    constexpr int c_size = tb_m * tb_n * sizeof(float) / 16;
+
+    constexpr int active_threads = 32 * tb_n_warps;
+    bool is_th_active = threadIdx.x < active_threads;
+
+    constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4;
+    constexpr int th_size = num_floats * sizeof(float) / 16;
+
+    int c_cur_offset = locks_off * c_size;
+
+    if (!is_th_active) {
+      return;
+    }
+
+    if (!first) {
+      float* frag_c_ptr = reinterpret_cast<float*>(&frag_c);
+  #pragma unroll
+      for (int k = 0; k < th_size; k++) {
+        if constexpr (m_block_size_8) {
+          if (k % 2) continue;
+        } else {
+          if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens)
+            continue;
+        }
+
+        sh_red[threadIdx.x] =
+            C_tmp[c_cur_offset + active_threads * k + threadIdx.x];
+
+        float* sh_c_ptr = reinterpret_cast<float*>(&sh_red[threadIdx.x]);
+  #pragma unroll
+        for (int f = 0; f < 4; f++) {
+          frag_c_ptr[k * 4 + f] += sh_c_ptr[f];
+        }
+      }
+    }
+
+    if (!last) {
+      int4* frag_c_ptr = reinterpret_cast<int4*>(&frag_c);
+  #pragma unroll
+      for (int k = 0; k < th_size; k++) {
+        if constexpr (m_block_size_8) {
+          if (k % 2) continue;
+        } else {
+          if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens)
+            continue;
+        }
+
+        C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k];
+      }
+    }
+  };
+
+  // Write out the reduce final result in the correct layout. We only actually
+  // reshuffle matrix fragments in this step, the reduction above is performed
+  // in fragment layout.
+  auto write_result = [&](bool last) {
+    int c_gl_stride = prob_n / 8;
+    constexpr int c_sh_stride = 2 * thread_n_blocks + 1;
+    int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
+    constexpr int c_sh_rd_delta =
+        c_sh_stride * (threads / (2 * thread_n_blocks));
+
+    int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +
+                  (threadIdx.x % (2 * thread_n_blocks));
+    c_gl_wr += (2 * thread_n_blocks) * slice_col;
+    int c_sh_wr;
+    if constexpr (m_block_size_8) {
+      c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) +
+                (threadIdx.x % 32) / 4;
+      c_sh_wr += 64 * (threadIdx.x / 32);
+    } else {
+      c_sh_wr =
+          (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
+      c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32);
+    }
+
+    int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
+                  (threadIdx.x % (2 * thread_n_blocks));
+
+    // We first reorder in shared memory to guarantee the most efficient final
+    // global write patterns
+    auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
+      if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
+        if (!mul_topk_weights) {
+          c0 *= global_scale_f32;
+          c1 *= global_scale_f32;
+        }
+      }
+
+      c_scalar_t2 res =
+          Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
+
+      // For per-column quantization we finally apply the scale here (only for
+      // 4-bit)
+      if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit &&
+                    b_type.size_bits() == 4 &&
+                    (has_zp && dequant_skip_flop || !has_zp)) {
+        c_scalar_t2 tmp_scale = s[0];
+        if constexpr (m_block_size_8) {
+          tmp_scale = Cdtype::num2num2(
+              reinterpret_cast<scalar_t*>(&s[0])[(threadIdx.x % 8) / 4]);
+        }
+        res = __hmul2(res, tmp_scale);
+      }
+
+      if (has_bias && last) {
+        c_scalar_t2 tmp_bias = b_bias[0];
+        if constexpr (m_block_size_8) {
+          tmp_bias = Cdtype::num2num2(
+              reinterpret_cast<scalar_t*>(&b_bias[0])[(threadIdx.x % 8) / 4]);
+        }
+        res = __hadd2(res, tmp_bias);
+      }
+
+      if constexpr (m_block_size_8) {
+        ((c_scalar_t*)sh_red)[idx] = res.x;
+        ((c_scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
+      } else {
+        ((c_scalar_t2*)sh_red)[idx] = res;
+      }
+    };
+
+    if (threadIdx.x / 32 < tb_n_warps) {
+  #pragma unroll
+      for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+        for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) {
+          if constexpr (m_block_size_8) {
+            int wr = c_sh_wr + 16 * j;
+            write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
+                  frag_s[j / 2][2 * (j % 2) + 0],
+                  frag_bias[j / 2][2 * (j % 2) + 0]);
+            write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3],
+                  frag_s[j / 2][2 * (j % 2) + 1],
+                  frag_bias[j / 2][2 * (j % 2) + 1]);
+          } else {
+            int wr = c_sh_wr + 8 * j;
+            write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
+                  frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0],
+                  frag_bias[j / 2][2 * (j % 2) + 0]);
+            write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
+                  frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0],
+                  frag_bias[j / 2][2 * (j % 2) + 0]);
+            write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
+                  frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1],
+                  frag_bias[j / 2][2 * (j % 2) + 1]);
+            write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
+                  frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1],
+                  frag_bias[j / 2][2 * (j % 2) + 1]);
+          }
+        }
+        c_sh_wr += 16 * (4 * c_sh_stride);
+      }
+    }
+    __syncthreads();
+
+  #pragma unroll
+    for (int i = 0;
+         i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
+         i++) {
+      int row = c_gl_wr / c_gl_stride;
+      if (row < block_num_valid_tokens) {
+        int64_t sorted_row = sh_block_sorted_ids[row];
+        int64_t true_idx = sorted_row * c_gl_stride + c_gl_wr % c_gl_stride;
+        c_scalar_t2 topk_weight_score;
+        if (mul_topk_weights) topk_weight_score = sh_block_topk_weights[row];
+        if (use_atomic_add && slice_count > 1 || mul_topk_weights) {
+          c_scalar_t2* C_half2 = reinterpret_cast<c_scalar_t2*>(&C[true_idx]);
+          c_scalar_t2* sh_red_half2 =
+              reinterpret_cast<c_scalar_t2*>(&sh_red[c_sh_rd]);
+          if (mul_topk_weights) {
+  #pragma unroll
+            for (int a = 0; a < 4; a++) {
+              sh_red_half2[a] = __hmul2(sh_red_half2[a], topk_weight_score);
+            }
+          }
+
+          if (use_atomic_add && slice_count > 1) {
+  #pragma unroll
+            for (int a = 0; a < 4; a++) {
+              atomicAdd(&C_half2[a], sh_red_half2[a]);
+            }
+          } else {
+            C[true_idx] = *reinterpret_cast<int4*>(sh_red_half2);
+          }
+        } else {
+          C[true_idx] = sh_red[c_sh_rd];
+        }
+        c_gl_wr += c_gl_wr_delta;
+        c_sh_rd += c_sh_rd_delta;
+      }
+    }
+    __syncthreads();
+  };
+
+  // Start global fetch and register load pipelines.
+  auto start_pipes = [&]() {
+
+  #pragma unroll
+    for (int i = 0; i < stages - 1; i++) {
+      if (has_act_order && i == 0) {
+        int last_g_idx = slice_k_start + stages * tb_k * 2;
+        if (last_g_idx >= prob_k) {
+          last_g_idx = prob_k - 1;
+        }
+        fetch_act_order_scales_to_shared(true, g_idx[slice_k_start],
+                                         g_idx[last_g_idx]);
+      }
+
+      if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
+        if (i == 0) {
+          fetch_col_zp_to_shared();
+          if constexpr (!dequant_skip_flop) {
+            fetch_col_scale_to_shared();
+          }
+        }
+      }
+      fetch_to_shared(i, i, i < slice_iters);
+    }
+
+    zero_accums();
+    wait_for_stage();
+    init_same_group(0);
+    fetch_to_registers(0, 0);
+    fetch_scales_to_registers(0, 0);
+    fetch_zp_to_registers(0, 0);
+    a_gl_rd_col += a_gl_rd_delta_o * (stages - 1);
+    if constexpr (has_act_order) {
+      slice_k_start_shared_fetch += tb_k * (stages - 1);
+    }
+  };
+  if (slice_iters) {
+    start_pipes();
+  }
+
+  // Main loop.
+  while (slice_iters) {
+    // We unroll over both the global fetch and the register load pipeline to
+    // ensure all shared memory accesses are static. Note that both pipelines
+    // have even length meaning that the next iteration will always start at
+    // index 0.
+
+  #pragma unroll
+    for (int pipe = 0; pipe < stages;) {
+  #pragma unroll
+      for (int k = 0; k < b_sh_wr_iters; k++) {
+        fetch_to_registers(k + 1, pipe % stages);
+        fetch_scales_to_registers(k + 1, pipe);
+        fetch_zp_to_registers(k + 1, pipe);
+        if (k == b_sh_wr_iters - 2) {
+          fetch_to_shared((pipe + stages - 1) % stages, pipe,
+                          slice_iters >= stages);
+          pipe++;
+          wait_for_stage();
+          init_same_group(pipe % stages);
+        }
+
+        if constexpr (!is_a_8bit) {
+          matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0));
+        } else {
+          static_assert(group_blocks != 0 && group_blocks != 1);
+          matmul_a8(k);
+        }
+      }
+      slice_iters--;
+      if (slice_iters == 0) {
+        break;
+      }
+    }
+
+    a_gl_rd_col += a_gl_rd_delta_o * stages;
+
+    if constexpr (has_act_order) {
+      slice_k_start += tb_k * stages;
+
+      if (slice_k_start < prob_k) {
+        slice_k_start_shared_fetch += tb_k * stages;
+        int first_group_id = g_idx[slice_k_start];
+        int last_g_idx = slice_k_start + stages * tb_k * 2;
+        if (last_g_idx >= prob_k) {
+          last_g_idx = prob_k - 1;
+        }
+        int last_group_id = g_idx[last_g_idx];
+        if (last_group_id >= sh_first_group_id + sh_num_groups) {
+          fetch_act_order_scales_to_shared(false, first_group_id,
+                                           last_group_id);
+          __syncthreads();
+        }
+      }
+    }
+
+    // Process results and, if necessary, proceed to the next column slice.
+    // While this pattern may not be the most readable, other ways of writing
+    // the loop seemed to noticeably worse performance after compilation.
+    if (slice_iters == 0) {
+      // convert fp16 accum to fp32 for reduction
+      if constexpr (use_fp16_accum) {
+  #pragma unroll
+        for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) {
+          float* frag_c_part_float = reinterpret_cast<float*>(frag_c) + i * 4;
+          scalar_t* frag_c_part_half =
+              reinterpret_cast<scalar_t*>(frag_c_part_float);
+
+  #pragma unroll
+          for (int i = 3; i >= 0; i--) {
+            frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]);
+          }
+        }
+      }
+
+      if constexpr (is_a_8bit) {
+        float frag_a_s[2 * thread_m_blocks];
+
+        for (int i = 0; i < 2 * thread_m_blocks; i++)
+          frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4];
+
+  #pragma unroll
+        for (int j = 0; j < 2; j++) {
+  #pragma unroll
+          for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+            for (int g = 0; g < 4; g++) {
+              float c_val = frag_c[i][j][0][g];
+
+              if constexpr (a_type == vllm::kS8) {
+                c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val));
+              }
+              float s_val = frag_a_s[i * 2 + g / 2];
+              frag_c[i][j][0][g] = c_val * s_val;
+            }
+  #pragma unroll
+            for (int g = 0; g < 4; g++) {
+              float c_val = frag_c[i][j][1][g];
+
+              if constexpr (a_type == vllm::kS8) {
+                c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val));
+              }
+              float s_val = frag_a_s[i * 2 + g / 2];
+              frag_c[i][j][1][g] = c_val * s_val;
+            }
+          }
+        }
+      }
+
+      cp_async_wait<0>();
+      bool last = slice_idx == slice_count - 1;
+      // For per-column scales, we only fetch them here in the final step before
+      // write-out
+      if constexpr (!has_act_order && group_blocks == -1 &&
+                    (has_zp && dequant_skip_flop || !has_zp)) {
+        if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) {
+          if (s_sh_wr_pred) {
+            cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
+          }
+          cp_async_fence();
+        }
+      }
+
+      thread_block_reduce();
+
+      if (has_bias && last) {
+        __syncthreads();
+        cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd],
+                       threadIdx.x < 16 * thread_n_blocks / 8);
+        cp_async_fence();
+      }
+
+      if constexpr (!has_act_order && group_blocks == -1 &&
+                    (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) {
+        if constexpr (is_a_8bit) {
+          cp_async_wait<0>();
+          __syncthreads();
+          if (threadIdx.x / 32 < tb_n_warps) {
+            reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
+          }
+        } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) {
+          cp_async_wait<0>();
+          __syncthreads();
+          if (threadIdx.x / 32 < tb_n_warps) {
+            reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
+            reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
+            if constexpr (m_block_size_8) {
+              int idx = (threadIdx.x / 4) % 2;
+              c_scalar_t2* frag_s_half2 =
+                  reinterpret_cast<c_scalar_t2*>(frag_s);
+  #pragma unroll
+              for (int i = 0; i < 8; i++) {
+                frag_s_half2[i] = Cdtype::num2num2(
+                    reinterpret_cast<c_scalar_t*>(&frag_s_half2[i])[idx]);
+              }
+            }
+          }
+        }
+      }
+
+      // For 8-bit channelwise, we apply the scale before the global reduction
+      // that converts the fp32 results to fp16 (so that we avoid possible
+      // overflow in fp16)
+      if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) {
+  #pragma unroll
+        for (int j = 0; j < 2; j++) {
+          float2 aa[2];
+          aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]);
+          aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]);
+
+  #pragma unroll
+          for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+            for (int g = 0; g < 4; g++) {
+              float scale = reinterpret_cast<float*>(&aa[0])[g % 2];
+              frag_c[i][j][0][g] *= scale;
+            }
+
+  #pragma unroll
+            for (int g = 0; g < 4; g++) {
+              float scale = reinterpret_cast<float*>(&aa[1])[g % 2];
+              frag_c[i][j][1][g] *= scale;
+            }
+          }
+        }
+      } else if (!has_act_order && group_blocks == -1 &&
+                 b_type.size_bits() == 8 &&
+                 (has_zp && dequant_skip_flop || !has_zp)) {
+        if (threadIdx.x / 32 < tb_n_warps) {
+  #pragma unroll
+          for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+            for (int j = 0; j < 4; j++) {
+              scale_float<c_type_id>(
+                  reinterpret_cast<float*>(&frag_c[i][j][0][0]),
+                  frag_s[j / 2][2 * (j % 2) + 0]);
+              scale_float<c_type_id>(
+                  reinterpret_cast<float*>(&frag_c[i][j][0][2]),
+                  frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]);
+
+              if constexpr (!m_block_size_8) {
+                scale_float<c_type_id>(
+                    reinterpret_cast<float*>(&frag_c[i][j][1][0]),
+                    frag_s[j / 2][2 * (j % 2) + 1]);
+                scale_float<c_type_id>(
+                    reinterpret_cast<float*>(&frag_c[i][j][1][2]),
+                    frag_s[j / 2][2 * (j % 2) + 1]);
+              }
+            }
+          }
+        }
+      }
+
+      if (slice_count > 1 && !use_atomic_add) {
+        // only globally reduce if there is more than one block in a slice
+        barrier_acquire(&locks[locks_off], slice_idx);
+        if (use_fp32_reduce) {
+          global_reduce_fp32(slice_idx == 0, last);
+        } else {
+          global_reduce_fp16(slice_idx == 0, last);
+        }
+        barrier_release(&locks[locks_off], last);
+      }
+
+      if (has_bias && last) {
+        cp_async_wait<0>();
+        __syncthreads();
+        reinterpret_cast<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd];
+        if constexpr (!is_a_8bit)
+          reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
+        __syncthreads();
+      }
+
+      if (use_atomic_add && slice_count > 1 && slice_idx != 0)
+        wait_negative_and_add(&locks[locks_off]);
+      if (last || use_atomic_add)
+        // only the last block in a slice actually writes the result
+        write_result(last);
+      slice_row = 0;
+      if (!in_part2) {
+        slice_col_par += gridDim.x;
+      } else {
+        slice_col_par++;
+        slice_col++;
+      }
+      is_first_matmul_in_slice = true;
+      init_slice();
+
+      if (slice_iters) {
+        a_gl_rd_col =
+            a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o;
+        b_gl_rd = B_expert_off + b_gl_stride * (threadIdx.x / b_sh_stride) +
+                  (threadIdx.x % b_sh_stride);
+        b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row;
+
+        bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
+        // Update slice k/n for scales loading
+        if constexpr (has_act_order) {
+          slice_k_start = tb_k * slice_row;
+          slice_k_finish = slice_k_start + tb_k * slice_iters;
+          slice_k_start_shared_fetch = slice_k_start;
+          slice_n_offset = act_s_col_tb_stride * slice_col;
+        } else {
+          if constexpr (group_blocks == -1) {
+            s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
+            zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
+          } else if constexpr (group_blocks >= thread_k_blocks) {
+            s_gl_rd =
+                s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
+                s_sh_stride * slice_col + threadIdx.x;
+            zp_gl_rd =
+                zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
+                zp_sh_stride * slice_col + threadIdx.x;
+          } else {
+            s_gl_rd =
+                s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
+                               threadIdx.x / s_sh_stride) +
+                s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
+            zp_gl_rd =
+                zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
+                                threadIdx.x / zp_sh_stride) +
+                zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
+          }
+        }
+        start_pipes();
+      }
+    }
+  }
+}
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+#endif
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/ops_standalone.cu b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/ops_standalone.cu
new file mode 100644
index 0000000000..7cff87bffa
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/ops_standalone.cu
@@ -0,0 +1,405 @@
+/*
+ * Standalone extraction of vLLM's NVFP4 Marlin MoE GEMM (W4A16, bf16 act/out).
+ * Ported from vllm csrc .../moe/marlin_moe_wna16/ops.cu:
+ *   - kept the marlin_mm launcher + helpers verbatim (launcher_body.inc),
+ *   - ported the moe_wna16_marlin_gemm wrapper from torch::stable::Tensor to
+ *     regular torch::Tensor,
+ *   - replaced STABLE_TORCH_LIBRARY_IMPL registration with a PYBIND11_MODULE.
+ */
+
+#ifndef MARLIN_NAMESPACE_NAME
+  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#endif
+
+#include "kernel.h"
+
+#include <torch/extension.h>
+#include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <ATen/cuda/CUDAContext.h>
+
+// STD_TORCH_CHECK comes from the header-only torch ABI; the kernel headers and
+// the launcher use it. Keep it available for launcher_body.inc.
+#include <torch/headeronly/util/Exception.h>
+
+#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t)               \
+  static_assert(std::is_same<scalar_t, half>::value ||          \
+                    std::is_same<scalar_t, nv_bfloat16>::value, \
+                "only float16 and bfloat16 is supported");
+
+// ---- launcher + helpers (verbatim from upstream ops.cu lines 42..541) ----
+#include "launcher_body.inc"
+
+// -------------------------------------------------------------------------
+// Ported wrapper: regular torch::Tensor signature.
+// b_q_type is received as an int64 ScalarType id (vllm::ScalarTypeId) and
+// reconstructed via vllm::ScalarType::from_id().
+// -------------------------------------------------------------------------
+
+torch::Tensor moe_wna16_marlin_gemm(
+    torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
+    torch::Tensor& b_q_weight,
+    std::optional<torch::Tensor> const& b_bias_or_none,
+    torch::Tensor& b_scales,
+    std::optional<torch::Tensor> const& a_scales_or_none,
+    std::optional<torch::Tensor> const& global_scale_or_none,
+    std::optional<torch::Tensor> const& b_zeros_or_none,
+    std::optional<torch::Tensor> const& g_idx_or_none,
+    std::optional<torch::Tensor> const& perm_or_none,
+    torch::Tensor& workspace, torch::Tensor& sorted_token_ids,
+    torch::Tensor& expert_ids, torch::Tensor& num_tokens_past_padded,
+    torch::Tensor& topk_weights, int64_t moe_block_size, int64_t top_k,
+    bool mul_topk_weights, int64_t b_type_id, int64_t size_m, int64_t size_n,
+    int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
+    bool is_zp_float, int64_t thread_k, int64_t thread_n,
+    int64_t blocks_per_sm) {
+  vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
+
+  auto c_dtype = a.scalar_type();
+  if (a.scalar_type() == at::ScalarType::Half) {
+    a_type_id = vllm::kFloat16.id();
+    c_type_id = vllm::kFloat16.id();
+  } else if (a.scalar_type() == at::ScalarType::BFloat16) {
+    a_type_id = vllm::kBFloat16.id();
+    c_type_id = vllm::kBFloat16.id();
+  } else {
+    c_dtype = b_scales.scalar_type();
+    if (b_scales.scalar_type() == at::ScalarType::Half) {
+      c_type_id = vllm::kFloat16.id();
+    } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) {
+      c_type_id = vllm::kBFloat16.id();
+    } else {
+      c_type_id = vllm::kBFloat16.id();
+
+      TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
+      torch::Tensor c = c_or_none.value();
+      c_dtype = c.scalar_type();
+
+      if (c.scalar_type() == at::ScalarType::Half) {
+        c_type_id = vllm::kFloat16.id();
+      } else if (c.scalar_type() == at::ScalarType::BFloat16) {
+        c_type_id = vllm::kBFloat16.id();
+      } else {
+        TORCH_CHECK(false, "unsupported c dtype");
+      }
+    }
+
+    if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) {
+      a_type_id = vllm::kFE4M3fn.id();
+    } else if (a.scalar_type() == at::ScalarType::Char) {
+      a_type_id = vllm::kS8.id();
+    } else {
+      TORCH_CHECK(false, "unsupported `a` scalar_type");
+    }
+  }
+
+  s_type_id = c_type_id;
+  if (b_type_id == vllm::kFE2M1f.id()) {
+    if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) {
+      s_type_id = vllm::kFE4M3fn.id();
+    } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
+      s_type_id = vllm::kFE8M0fnu.id();
+    } else {
+      TORCH_CHECK(false,
+                  "When b_type = float4_e2m1f, b_scale scalar type must be",
+                  "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
+    }
+  } else if (b_type_id == vllm::kFE4M3fn.id() &&
+             b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
+    s_type_id = vllm::kFE8M0fnu.id();
+  }
+
+  vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id);
+  vllm::ScalarType b_type = vllm::ScalarType::from_id(b_type_id);
+  vllm::ScalarType c_type = vllm::ScalarType::from_id(c_type_id);
+  vllm::ScalarType s_type = vllm::ScalarType::from_id(s_type_id);
+
+  int pack_factor = 32 / b_type.size_bits();
+  int num_experts = b_q_weight.size(0);
+
+  if (moe_block_size != 8) {
+    TORCH_CHECK(moe_block_size % 16 == 0,
+                "unsupported moe_block_size=", moe_block_size);
+    TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64,
+                "unsupported moe_block_size=", moe_block_size);
+  }
+
+  // Verify A
+  TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0),
+              ", size_m = ", size_m);
+  TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1),
+              ", size_k = ", size_k);
+
+  // Verify B
+  TORCH_CHECK(
+      size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k,
+      " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
+  TORCH_CHECK(
+      (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1),
+      "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
+      ", size_k = ", size_k,
+      ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
+  TORCH_CHECK(
+      b_q_weight.size(2) % MARLIN_NAMESPACE_NAME::tile_size == 0,
+      "b_q_weight.size(2) = ", b_q_weight.size(2),
+      " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
+  int actual_size_n =
+      (b_q_weight.size(2) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor;
+  TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
+              ", actual_size_n = ", actual_size_n);
+
+  // Verify device and strides
+  TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
+  TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
+
+  TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
+  TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
+
+  TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
+  TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
+
+  torch::Tensor a_scales;
+  auto opts_f32 = a.options().dtype(at::kFloat);
+
+  if (a_scales_or_none.has_value()) {
+    a_scales = a_scales_or_none.value();
+    TORCH_CHECK(a_type.size_bits() == 8,
+                "a_scales can only be used for 8bit activation.");
+  } else {
+    a_scales = torch::empty({0}, opts_f32);
+    TORCH_CHECK(a_type.size_bits() != 8,
+                "the a_scales parameter must be passed for 8bit activation.");
+  }
+
+  // sms: number of SMs to use for the kernel
+  int sms = -1;
+  cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
+
+  // Alloc buffers
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+  torch::Tensor c;
+  if (c_or_none.has_value()) {
+    c = c_or_none.value();
+    TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
+    TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
+    TORCH_CHECK(c.size(0) == size_m * top_k,
+                "Shape mismatch: c.size(0) = ", c.size(0),
+                ", size_m * topk = ", size_m * top_k);
+    TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1),
+                ", size_n = ", size_n);
+  } else {
+    c = torch::empty({size_m * top_k, size_n}, a.options().dtype(c_dtype));
+  }
+
+  // Alloc C tmp buffer that is going to be used for the global reduce
+  torch::Tensor c_tmp;
+  auto options_fp32 = a.options().dtype(at::kFloat);
+  if (use_fp32_reduce && !use_atomic_add) {
+    // max num of threadblocks is sms * 4
+    long max_c_tmp_size = min(
+        (long)size_n * sorted_token_ids.size(0),
+        (long)sms * 4 * moe_block_size * MARLIN_NAMESPACE_NAME::max_thread_n);
+    if (moe_block_size == 8) max_c_tmp_size *= 2;
+    c_tmp = torch::empty({max_c_tmp_size}, options_fp32);
+  } else {
+    c_tmp = torch::empty({0}, options_fp32);
+  }
+
+  // Detect groupsize and act_order
+  int num_groups = -1;
+  int group_size = -1;
+
+  int rank = b_scales.dim();
+  TORCH_CHECK(rank == 3, "b_scales rank = ", rank, " is not 3");
+  TORCH_CHECK(b_scales.size(2) == size_n, "b_scales dim 2 = ", b_scales.size(2),
+              " is not size_n = ", size_n);
+  num_groups = b_scales.size(1);
+
+  torch::Tensor g_idx, perm, a_tmp;
+  auto opts_a = a.options().dtype(c_dtype);
+  if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
+    g_idx = g_idx_or_none.value();
+    perm = perm_or_none.value();
+
+    TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
+    TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
+    TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
+    TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
+
+    // Verify g_idx and perm
+    TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
+                    (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
+                "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
+                " and perm.size(-1) = ", perm.size(-1),
+                ", where size_k = ", size_k);
+  } else {
+    g_idx = torch::empty({0}, opts_a);
+    perm = torch::empty({0}, opts_a);
+    a_tmp = torch::empty({0}, opts_a);
+  }
+  bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
+
+  if (has_act_order) {
+    a_tmp = torch::empty({size_m * top_k, size_k}, opts_a);
+    if (is_k_full) {
+      TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
+      TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
+                  ", is not divisible by num_groups = ", num_groups);
+      group_size = size_k / num_groups;
+    } else {
+      group_size = 0;
+    }
+
+  } else {
+    a_tmp = torch::empty({0}, opts_a);
+    if (num_groups > 1) {
+      TORCH_CHECK(
+          size_k % num_groups == 0, "size_k = ", size_k,
+          ", is not divisible by b_scales.size(1) = ", b_scales.size(1));
+      group_size = size_k / num_groups;
+    } else {
+      group_size = -1;
+    }
+  }
+
+  torch::Tensor global_scale;
+  if (global_scale_or_none.has_value()) {
+    global_scale = global_scale_or_none.value();
+    TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
+                "global_scale can only be used for nvfp4 format.");
+  } else {
+    global_scale = torch::empty({0}, options_fp32);
+    TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
+                "the global_scale parameter must be passed for nvfp4 format.");
+  }
+
+  bool has_bias = b_bias_or_none.has_value();
+  torch::Tensor b_bias;
+  if (has_bias) {
+    b_bias = b_bias_or_none.value();
+    TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
+    TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
+    TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(1) != size_n");
+    TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1");
+  } else {
+    b_bias = torch::empty({0}, opts_a);
+  }
+
+  torch::Tensor b_zeros;
+  if (b_zeros_or_none.has_value()) {
+    b_zeros = b_zeros_or_none.value();
+    TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
+    TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
+  } else {
+    b_zeros = torch::empty({0}, opts_a);
+  }
+  bool has_zp = b_zeros.size(-1) > 0;
+  if (has_zp) {
+    TORCH_CHECK(b_type == vllm::kU4 || b_type == vllm::kU8,
+                "b_type must be u4 or u8 when has_zp = True. Got = ",
+                b_type.str());
+  } else {
+    TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
+                    b_type == vllm::kS4 || b_type == vllm::kS8 ||
+                    b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
+                "b_type must be uint4b8, uint8b128, int4, int8, "
+                "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
+                b_type.str());
+  }
+
+  if (has_zp && is_zp_float) {
+    TORCH_CHECK(a.scalar_type() == at::ScalarType::Half,
+                "Computation type must be float16 (half) when using float zero "
+                "points.");
+  }
+
+  // Verify b_zeros
+  if (has_zp) {
+    int rank2 = b_zeros.dim();
+    TORCH_CHECK(rank2 == 3, "b_zeros rank = ", rank2, " is not 3");
+    if (is_zp_float) {
+      TORCH_CHECK(b_zeros.size(2) == size_n,
+                  "b_zeros dim 2 = ", b_zeros.size(2),
+                  " is not size_n = ", size_n);
+      TORCH_CHECK(num_groups == b_zeros.size(1),
+                  "b_zeros dim 1 = ", b_zeros.size(1),
+                  " is not num_groups = ", num_groups);
+      TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
+    } else {
+      TORCH_CHECK(b_zeros.size(1) == num_groups,
+                  "b_zeros dim 1 = ", b_zeros.size(1),
+                  " is not num_groups = ", num_groups);
+      TORCH_CHECK(b_zeros.size(2) == size_n / pack_factor,
+                  "b_zeros dim 2 = ", b_zeros.size(2),
+                  " is not size_n / pack_factor = ", size_n / pack_factor);
+    }
+  }
+
+  // Verify workspace size
+  TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
+              "size_n = ", size_n, ", is not divisible by min_thread_n = ",
+              MARLIN_NAMESPACE_NAME::min_thread_n);
+
+  int max_n_tiles = size_n / MARLIN_NAMESPACE_NAME::min_thread_n;
+  int min_workspace_size = min(
+      max_n_tiles * (int)(sorted_token_ids.size(0) / moe_block_size), sms * 4);
+  TORCH_CHECK(workspace.numel() >= min_workspace_size,
+              "workspace.numel = ", workspace.numel(),
+              " is below min_workspace_size = ", min_workspace_size);
+
+  int dev = a.get_device();
+
+  TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
+              "scalar type of a_scales must be float");
+  TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
+              "scalar type of global_scale must be float");
+  if (a_type.size_bits() == 16) {
+    TORCH_CHECK(a.scalar_type() == c.scalar_type(),
+                "scalar type of a must be the same with c for 16 bit "
+                "activation");
+  }
+
+  MARLIN_NAMESPACE_NAME::marlin_mm(
+      a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(),
+      c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(),
+      a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(),
+      global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(),
+      g_idx.mutable_data_ptr(), perm.mutable_data_ptr(),
+      a_tmp.mutable_data_ptr(), sorted_token_ids.mutable_data_ptr(),
+      expert_ids.mutable_data_ptr(), num_tokens_past_padded.mutable_data_ptr(),
+      topk_weights.mutable_data_ptr(), moe_block_size, num_experts, top_k,
+      mul_topk_weights, size_m, size_n, size_k, workspace.mutable_data_ptr(),
+      a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full,
+      has_zp, num_groups, group_size, dev,
+      at::cuda::getCurrentCUDAStream(dev).stream(), thread_k, thread_n, sms,
+      blocks_per_sm, use_atomic_add, use_fp32_reduce, is_zp_float);
+
+  return c;
+}
+
+// Defined in repack_standalone.cu (ported gptq_marlin_repack, torch::Tensor ABI).
+torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
+                                 int64_t size_k, int64_t size_n,
+                                 int64_t num_bits, bool is_a_8bit);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  namespace py = pybind11;
+  m.def("gptq_marlin_repack", &gptq_marlin_repack,
+        "NVFP4 weight -> Marlin tile layout repack (standalone)",
+        py::arg("b_q_weight"), py::arg("perm"), py::arg("size_k"),
+        py::arg("size_n"), py::arg("num_bits"), py::arg("is_a_8bit") = false);
+  m.def("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm,
+        "NVFP4 W4A16 Marlin MoE GEMM (standalone)", py::arg("a"),
+        py::arg("c_or_none"), py::arg("b_q_weight"), py::arg("b_bias_or_none"),
+        py::arg("b_scales"), py::arg("a_scales_or_none"),
+        py::arg("global_scale_or_none"), py::arg("b_zeros_or_none"),
+        py::arg("g_idx_or_none"), py::arg("perm_or_none"), py::arg("workspace"),
+        py::arg("sorted_token_ids"), py::arg("expert_ids"),
+        py::arg("num_tokens_past_padded"), py::arg("topk_weights"),
+        py::arg("moe_block_size"), py::arg("top_k"), py::arg("mul_topk_weights"),
+        py::arg("b_type_id"), py::arg("size_m"), py::arg("size_n"),
+        py::arg("size_k"), py::arg("is_k_full"), py::arg("use_atomic_add"),
+        py::arg("use_fp32_reduce"), py::arg("is_zp_float"),
+        py::arg("thread_k") = -1, py::arg("thread_n") = -1,
+        py::arg("blocks_per_sm") = -1);
+}
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/repack_standalone.cu b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/repack_standalone.cu
new file mode 100644
index 0000000000..6651b60852
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/repack_standalone.cu
@@ -0,0 +1,371 @@
+/* Standalone port of vLLM's gptq_marlin_repack (NVFP4 weight -> Marlin tile layout).
+ * Source: vllm csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu.
+ * Kernel body kept VERBATIM; host wrapper ported torch::stable::Tensor -> torch::Tensor;
+ * STABLE_TORCH_LIBRARY_IMPL registration dropped (a pybind def is added in ops_standalone.cu).
+ * Compiled in the default `marlin` namespace: undef the build's MARLIN_NAMESPACE_NAME for this
+ * TU so marlin.cuh's constants land in `marlin` (the GEMM TU uses marlin_moe_wna16; constants are
+ * static/internal-linkage so the two TUs don't clash). */
+#ifdef MARLIN_NAMESPACE_NAME
+  #undef MARLIN_NAMESPACE_NAME
+#endif
+
+#include "marlin.cuh"
+
+#include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/headeronly/util/Exception.h>
+
+namespace marlin {
+
+template <int const num_threads, int const num_bits, bool const has_perm,
+          bool is_a_8bit>
+__global__ void gptq_marlin_repack_kernel(
+    uint32_t const* __restrict__ b_q_weight_ptr,
+    uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
+    int size_k, int size_n) {
+  constexpr int pack_factor = 32 / num_bits;
+
+  constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1);
+  constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1);
+  int k_tiles = size_k / target_tile_k_size;
+  int n_tiles = size_n / target_tile_n_size;
+  int block_k_tiles = div_ceil(k_tiles, gridDim.x);
+
+  auto start_k_tile = blockIdx.x * block_k_tiles;
+  if (start_k_tile >= k_tiles) {
+    return;
+  }
+
+  int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles);
+
+  // Wait until the next thread tile has been loaded to shared memory.
+  auto wait_for_stage = [&]() {
+    // We only have `stages - 2` active fetches since we are double buffering
+    // and can only issue the next fetch when it is guaranteed that the previous
+    // shared memory load is fully complete (as it may otherwise be
+    // overwritten).
+    cp_async_wait<repack_stages - 2>();
+    __syncthreads();
+  };
+
+  extern __shared__ int4 sh[];
+
+  constexpr int perm_size = target_tile_k_size / 4;
+
+  int4* sh_perm_ptr = sh;
+  int4* sh_pipe_ptr = sh_perm_ptr;
+  if constexpr (has_perm) {
+    sh_pipe_ptr += perm_size;
+  }
+
+  constexpr int tile_ints = target_tile_k_size / pack_factor;
+
+  constexpr int stage_n_threads = target_tile_n_size / 4;
+  constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints;
+  constexpr int stage_size = stage_k_threads * stage_n_threads;
+
+  auto load_perm_to_shared = [&](int k_tile_id) {
+    int first_k_int4 = (k_tile_id * target_tile_k_size) / 4;
+
+    int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);
+
+    if (threadIdx.x < perm_size) {
+      sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x];
+    }
+    __syncthreads();
+  };
+
+  auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      cp_async_fence();
+      return;
+    }
+
+    int first_n = n_tile_id * target_tile_n_size;
+
+    int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;
+
+    if constexpr (has_perm) {
+      if (threadIdx.x < stage_size) {
+        auto k_id = threadIdx.x / stage_n_threads;
+        auto n_id = threadIdx.x % stage_n_threads;
+
+        uint32_t const* sh_perm_int_ptr =
+            reinterpret_cast<uint32_t const*>(sh_perm_ptr);
+
+        int src_k = sh_perm_int_ptr[k_id];
+        int src_k_packed = src_k / pack_factor;
+
+        cp_async4(
+            &sh_ptr[k_id * stage_n_threads + n_id],
+            reinterpret_cast<int4 const*>(&(
+                b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)])));
+      }
+
+    } else {
+      if (threadIdx.x < stage_size) {
+        auto k_id = threadIdx.x / stage_n_threads;
+        auto n_id = threadIdx.x % stage_n_threads;
+
+        int first_k = k_tile_id * target_tile_k_size;
+        int first_k_packed = first_k / pack_factor;
+
+        cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
+                  reinterpret_cast<int4 const*>(
+                      &(b_q_weight_ptr[(first_k_packed + k_id) * size_n +
+                                       first_n + (n_id * 4)])));
+      }
+    }
+
+    cp_async_fence();
+  };
+
+  auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      return;
+    }
+
+    auto warp_id = threadIdx.x / 32;
+    auto th_id = threadIdx.x % 32;
+
+    if (warp_id >= 4) {
+      return;
+    }
+
+    int tc_col = th_id / 4;
+    int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2);
+
+    constexpr int tc_offsets[4] = {0, 1, 8, 9};
+
+    int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col;
+
+    constexpr int sh_stride = target_tile_n_size;
+    constexpr uint32_t mask = (1 << num_bits) - 1;
+
+    int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
+    uint32_t* sh_stage_int_ptr = reinterpret_cast<uint32_t*>(sh_stage_ptr);
+
+    uint32_t* sh_perm_int_ptr = reinterpret_cast<uint32_t*>(sh_perm_ptr);
+
+    uint32_t vals[8];
+
+    if constexpr (has_perm) {
+      static_assert(!is_a_8bit);
+      for (int i = 0; i < 4; i++) {
+        int k_idx = tc_row + tc_offsets[i];
+
+        uint32_t src_k = sh_perm_int_ptr[k_idx];
+        uint32_t src_k_pos = src_k % pack_factor;
+
+        uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n];
+        uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask;
+
+        uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8];
+        uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask;
+
+        vals[i] = b1_cur_val;
+        vals[4 + i] = b2_cur_val;
+      }
+
+    } else {
+      uint32_t b1_vals[tile_ints];
+      uint32_t b2_vals[tile_ints];
+
+#pragma unroll
+      for (int i = 0; i < tile_ints; i++) {
+        if constexpr (is_a_8bit) {
+          b1_vals[i] =
+              sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8];
+        } else {
+          b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
+          b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
+        }
+      }
+
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]);
+        int cur_int = cur_elem / pack_factor;
+        int cur_pos = cur_elem % pack_factor;
+
+        vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+        if constexpr (is_a_8bit)
+          vals[4 + i] =
+              (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask;
+        else
+          vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+      }
+    }
+
+    constexpr int tile_size =
+        target_tile_k_size * target_tile_n_size / pack_factor;
+    int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;
+
+    // Result of:
+    // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
+    if constexpr (!is_a_8bit && num_bits == 4) {
+      int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
+
+      uint32_t res = 0;
+#pragma unroll
+      for (int i = 0; i < 8; i++) {
+        res |= vals[pack_idx[i]] << (i * 4);
+      }
+
+      out_ptr[out_offset + th_id * 4 + warp_id] = res;
+
+    } else if constexpr (is_a_8bit && num_bits == 4) {
+      int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7};
+
+      uint32_t res = 0;
+#pragma unroll
+      for (int i = 0; i < 8; i++) {
+        res |= vals[pack_idx[i]] << (i * 4);
+      }
+
+      out_ptr[out_offset + th_id * 4 + warp_id] = res;
+
+    } else {
+      constexpr int pack_idx[4] = {0, 2, 1, 3};
+
+      uint32_t res1 = 0;
+      uint32_t res2 = 0;
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        const int ii = is_a_8bit ? i : pack_idx[i];
+        res1 |= vals[ii] << (i * 8);
+        res2 |= vals[4 + ii] << (i * 8);
+      }
+
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2;
+    }
+  };
+
+  auto start_pipes = [&](int k_tile_id, int n_tile_id) {
+#pragma unroll
+    for (int pipe = 0; pipe < repack_stages - 1; pipe++) {
+      fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe);
+    }
+
+    wait_for_stage();
+  };
+#pragma unroll
+  for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) {
+    int n_tile_id = 0;
+
+    if constexpr (has_perm) {
+      load_perm_to_shared(k_tile_id);
+    }
+
+    start_pipes(k_tile_id, n_tile_id);
+
+    while (n_tile_id < n_tiles) {
+#pragma unroll
+      for (int pipe = 0; pipe < repack_stages; pipe++) {
+        fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id,
+                        n_tile_id + pipe + repack_stages - 1);
+        repack_tile(pipe, k_tile_id, n_tile_id + pipe);
+        wait_for_stage();
+      }
+      n_tile_id += repack_stages;
+    }
+  }
+}
+
+}  // namespace marlin
+
+#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT)                              \
+  else if (num_bits == NUM_BITS && has_perm == HAS_PERM &&                  \
+           is_a_8bit == IS_A_8BIT) {                                        \
+    cudaFuncSetAttribute(                                                   \
+        marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
+                                          HAS_PERM, IS_A_8BIT>,             \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);       \
+    marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS,     \
+                                      HAS_PERM, IS_A_8BIT>                  \
+        <<<blocks, marlin::repack_threads, max_shared_mem, stream>>>(       \
+            b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n);             \
+  }
+
+// Ported wrapper: regular torch::Tensor signature (was torch::stable::Tensor).
+torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
+                                 int64_t size_k, int64_t size_n,
+                                 int64_t num_bits, bool is_a_8bit) {
+  // Verify compatibility with marlin tile of 16x64
+  TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
+              " is not divisible by tile_k_size = ", marlin::tile_k_size);
+  TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
+              " is not divisible by tile_n_size = ", marlin::tile_n_size);
+
+  TORCH_CHECK(num_bits == 4 || num_bits == 8,
+              "num_bits must be 4 or 8. Got = ", num_bits);
+  int const pack_factor = 32 / num_bits;
+
+  // Verify B
+  TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0),
+              "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
+              ", size_k = ", size_k, ", pack_factor = ", pack_factor);
+  TORCH_CHECK(b_q_weight.size(1) == size_n,
+              "b_q_weight.size(1) = ", b_q_weight.size(1),
+              " is not size_n = ", size_n);
+
+  // Verify device and strides
+  TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU");
+  TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
+  TORCH_CHECK(b_q_weight.scalar_type() == at::kInt,
+              "b_q_weight type is not kInt");
+
+  TORCH_CHECK(perm.is_cuda(), "perm is not on GPU");
+  TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
+  TORCH_CHECK(perm.scalar_type() == at::kInt, "perm type is not at::kInt");
+
+  int const device_index = b_q_weight.get_device();
+  const c10::cuda::CUDAGuard device_guard(device_index);
+  const cudaStream_t stream =
+      at::cuda::getCurrentCUDAStream(device_index).stream();
+
+  // Alloc buffers
+  auto options = torch::TensorOptions()
+                     .dtype(b_q_weight.scalar_type())
+                     .device(b_q_weight.device());
+  torch::Tensor out = torch::empty(
+      {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
+      options);
+
+  // Detect if there is act_order
+  bool has_perm = perm.size(0) != 0;
+
+  // Get ptrs
+  uint32_t const* b_q_weight_ptr =
+      reinterpret_cast<uint32_t const*>(b_q_weight.const_data_ptr());
+  uint32_t const* perm_ptr =
+      reinterpret_cast<uint32_t const*>(perm.const_data_ptr());
+  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.mutable_data_ptr());
+
+  int blocks;
+  cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index);
+
+  int max_shared_mem = 0;
+  cudaDeviceGetAttribute(&max_shared_mem,
+                         cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index);
+  TORCH_CHECK(max_shared_mem > 0);
+
+  if (false) {
+  }
+  CALL_IF(4, false, false)
+  CALL_IF(4, true, false)
+  CALL_IF(8, false, false)
+  CALL_IF(8, true, false)
+
+  CALL_IF(4, false, true)
+  CALL_IF(8, false, true)
+
+  else {
+    TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
+                ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit);
+  }
+
+  return out;
+}
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu
new file mode 100644
index 0000000000..63e98e79c6
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_fe2m1f_bfloat16.cu
@@ -0,0 +1,40 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, true, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 1, 8, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 1, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 2, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 2, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 3, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 3, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 256, 4, 16, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 8, 4, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kFE2M1f.id(), vllm::kBFloat16.id(), vllm::kFE4M3fn.id(), 128, 4, 4, 8, false, 4, 1, false>( MARLIN_KERNEL_PARAMS );
+
+}
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/dequant.h b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/dequant.h
new file mode 100644
index 0000000000..edd97dbfcd
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/dequant.h
@@ -0,0 +1,609 @@
+/*
+Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16)
+
+The process of fast dequantization can be summarized as a combination
+of bitwise operations and floating-point computations:
+
+weight =>(bit_op / bitwise operations)=>
+f16_value =>(flop / floating-point computation)=>
+dequantized_weight
+
+Since the dequantized weights typically require subtracting the zero point and
+applying a scale factor, the floating-point computation step can be fused with
+the zero-point subtraction and scaling operations.
+
+The following are the parts that need to be modified for the fused operation
+of zero-point subtraction and scaling.
+
+## INT4 => FP16/BF16 or INT8 => FP16
+
+The floating-point computation is `__hsub2`
+
+If has zero points:
+
+    flop(bit_op(weight)) - flop(bit_op(zp))
+  = sub(bit_op(weight), bias) - sub(bit_op(zp), bias)
+  = bit_op(weight) - bit_op(zp)
+
+so we don't need additional modification.
+
+If has float zero points:
+
+    flop(bit_op(weight)) - fzp
+  = sub(bit_op(weight), bias) - fzp
+  = bit_op(weight) - (fzp + bias)
+
+where the `fzp + bias` can be computed at weight loading. But this
+may have accuracy issue, so we should not use this in most cases.
+
+If has not zero points:
+
+    scale(flop(bit_op(weight)))
+  = scale(sub(bit_op(weight), bias))
+  = scale(bit_op(weight)) - scale(bias)
+  = fma(bit_op(weight), scale_factor, scale(bias))
+
+where the `scale(bias)` can be cached. But this may have accuracy issue,
+so we should not use this in most cases.
+
+
+## INT8 => BF16
+
+INT8 => BF16 is a special case, it use byte_perm instead of flop.
+We cannot fused byte_perm with scaling.
+
+
+## FP4/FP8 => FP16/BF16
+
+    scale(flop(bit_op(weight)))
+  = scale(mul(bit_op(weight), multiplier))
+  = mul(bit_op(weight), scale_factor * multiplier)
+
+where `scale_factor * multiplier` can be computed at weight loading.
+
+*/
+
+#include "marlin_dtypes.cuh"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 750
+// Lookup-table based 3-input logical operation; explicitly used for
+// dequantization as the compiler does not seem to automatically recognize it in
+// all cases.
+template <int lut>
+__device__ inline int lop3(int a, int b, int c) {
+  int res;
+  asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+               : "=r"(res)
+               : "r"(a), "r"(b), "r"(c), "n"(lut));
+  return res;
+}
+
+// Constructs destination register by taking bytes from 2 sources (based on
+// mask)
+template <int start_byte, int mask>
+__device__ inline uint32_t prmt(uint32_t a) {
+  uint32_t res;
+  asm volatile("prmt.b32 %0, %1, %2, %3;\n"
+               : "=r"(res)
+               : "r"(a), "n"(start_byte), "n"(mask));
+  return res;
+}
+
+template <typename scalar_t2, vllm::ScalarTypeId w_type_id,
+          bool skip_flop = false>
+__device__ inline void dequant(int q, scalar_t2* frag_b);
+
+//
+// Efficiently dequantize 4bit values packed in an int32 value into a full
+// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below,
+// with some small changes:
+// - FP16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287
+// - BF16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
+//
+template <>
+__device__ inline void dequant<half2, vllm::kU4B8.id(), true>(int q,
+                                                              half2* frag_b) {
+  const int MASK = 0x000f000f;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  q >>= 4;
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+
+  frag_b[0] = *reinterpret_cast<half2*>(&lo);
+  frag_b[1] = *reinterpret_cast<half2*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU4B8.id(), false>(int q,
+                                                               half2* frag_b) {
+  const int LO = 0x000f000f;
+  const int HI = 0x00f000f0;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
+  // clang-format on
+  // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
+  // directly into `SUB` and `ADD`.
+  const int SUB = 0x64086408;
+  const int MUL = 0x2c002c00;
+  const int ADD = 0xd480d480;
+  frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
+                      *reinterpret_cast<const half2*>(&SUB));
+  frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
+                      *reinterpret_cast<const half2*>(&MUL),
+                      *reinterpret_cast<const half2*>(&ADD));
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU4.id(), true>(int q,
+                                                            half2* frag_b) {
+  dequant<half2, vllm::kU4B8.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU4.id(), false>(int q,
+                                                             half2* frag_b) {
+  const int LO = 0x000f000f;
+  const int HI = 0x00f000f0;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
+  // clang-format on
+  // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
+  // directly into `SUB` and `ADD`.
+  const int SUB = 0x64006400;
+  const int MUL = 0x2c002c00;
+  const int ADD = 0xd400d400;
+  frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
+                      *reinterpret_cast<const half2*>(&SUB));
+  frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
+                      *reinterpret_cast<const half2*>(&MUL),
+                      *reinterpret_cast<const half2*>(&ADD));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id(), true>(
+    int q, nv_bfloat162* frag_b) {
+  static constexpr uint32_t MASK = 0x000f000f;
+  static constexpr uint32_t EX = 0x43004300;
+
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  q >>= 4;
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  // clang-format on
+
+  frag_b[0] = *reinterpret_cast<nv_bfloat162*>(&lo);
+  frag_b[1] = *reinterpret_cast<nv_bfloat162*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kU4B8.id(), true>(q, frag_b);
+
+  static constexpr uint32_t SUB = 0x43084308;
+
+  frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+  frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU4.id(), true>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kU4B8.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU4.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kU4.id(), true>(q, frag_b);
+
+  static constexpr uint32_t SUB = 0x43004300;
+
+  frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+  frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+}
+
+//
+// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or
+// bf16 Reference:
+// - FP16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85
+// - BF16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
+//
+template <>
+__device__ inline void dequant<half2, vllm::kU8B128.id(), true>(int q,
+                                                                half2* frag_b) {
+  static constexpr uint32_t mask_for_elt_01 = 0x5250;
+  static constexpr uint32_t mask_for_elt_23 = 0x5351;
+  static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
+
+  uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
+  uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
+
+  frag_b[0] = *reinterpret_cast<half2*>(&lo);
+  frag_b[1] = *reinterpret_cast<half2*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU8B128.id(), false>(
+    int q, half2* frag_b) {
+  dequant<half2, vllm::kU8B128.id(), true>(q, frag_b);
+
+  static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
+  frag_b[0] = __hsub2(frag_b[0],
+                      *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+  frag_b[1] = __hsub2(frag_b[1],
+                      *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU8.id(), true>(int q,
+                                                            half2* frag_b) {
+  dequant<half2, vllm::kU8B128.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU8.id(), false>(int q,
+                                                             half2* frag_b) {
+  dequant<half2, vllm::kU8.id(), true>(q, frag_b);
+
+  static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
+  frag_b[0] = __hsub2(frag_b[0],
+                      *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+  frag_b[1] = __hsub2(frag_b[1],
+                      *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU8B128.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  float fp32_intermediates[4];
+  uint32_t* fp32_intermediates_casted =
+      reinterpret_cast<uint32_t*>(fp32_intermediates);
+
+  static constexpr uint32_t fp32_base = 0x4B000000;
+  fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
+  fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
+  fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
+  fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
+
+  fp32_intermediates[0] -= 8388736.f;
+  fp32_intermediates[1] -= 8388736.f;
+  fp32_intermediates[2] -= 8388736.f;
+  fp32_intermediates[3] -= 8388736.f;
+
+  uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
+  bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
+                                   fp32_intermediates_casted[1], 0x7632);
+  bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
+                                   fp32_intermediates_casted[3], 0x7632);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU8.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  float fp32_intermediates[4];
+  uint32_t* fp32_intermediates_casted =
+      reinterpret_cast<uint32_t*>(fp32_intermediates);
+
+  static constexpr uint32_t fp32_base = 0x4B000000;
+  fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
+  fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
+  fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
+  fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
+
+  fp32_intermediates[0] -= 8388608.f;
+  fp32_intermediates[1] -= 8388608.f;
+  fp32_intermediates[2] -= 8388608.f;
+  fp32_intermediates[3] -= 8388608.f;
+
+  uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
+  bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
+                                   fp32_intermediates_casted[1], 0x7632);
+  bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
+                                   fp32_intermediates_casted[3], 0x7632);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kFE4M3fn.id(), true>(
+    int q, half2* frag_b) {
+  // Constants for FP8 (E4M3) and FP16 formats
+  constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
+  constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT;
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kFE4M3fn.id(), false>(
+    int q, half2* frag_b) {
+  dequant<half2, vllm::kFE4M3fn.id(), true>(q, frag_b);
+
+  // Constants for FP8 (E4M3) and FP16 formats
+  constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET =
+      (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
+  const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id(), true>(
+    int q, nv_bfloat162* frag_b) {
+  // Constants for FP8 (E4M3) and BF16 formats
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
+
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to BF16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kFE4M3fn.id(), true>(q, frag_b);
+
+  // Constants for FP8 (E4M3) and BF16 formats
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET =
+      (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
+  // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent
+  // position
+  constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
+  const nv_bfloat162 bias_reg =
+      __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
+
+  // Convert to bfloat162 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kFE2M1f.id(), true>(int q,
+                                                                half2* frag_b) {
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
+  constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT;
+  constexpr int MASK = 0x70007000;
+
+  // Extract and shift FP4 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 4;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kFE2M1f.id(), false>(
+    int q, half2* frag_b) {
+  dequant<half2, vllm::kFE2M1f.id(), true>(q, frag_b);
+
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET =
+      (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1));
+  const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), true>(
+    int q, nv_bfloat162* frag_b) {
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT;
+  constexpr int MASK = 0x70007000;
+
+  // Extract and shift FP4 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 4;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kFE2M1f.id(), true>(q, frag_b);
+
+  // Constants for FP4 (E2M1) and BF16 formats
+  constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET =
+      (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1));
+  // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent
+  // position
+  constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
+  const nv_bfloat162 bias_reg =
+      __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kFE2M1f.id(), true>(
+    int q, __nv_fp8x4_e4m3* frag_b) {
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, FP8_EXPONENT = 4;
+  constexpr int RIGHT_SHIFT = FP8_EXPONENT - FP4_EXPONENT;
+  constexpr int MASK = 0x70707070;
+
+  // Extract and shift FP4 values to FP16 format
+  int Out1 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 4;
+  int Out2 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note1: reverse indexing is intentional because weights are permuted
+  // Note2: when dequant to 8bit type, we write to `frag_b[2]` instead of
+  //        `frag_b[1]` to fit the layout of tensorcore
+  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<int32_t, vllm::kU4B8.id(), true>(
+    int q, int32_t* frag_b) {
+  constexpr int repeated_zp = 0x08080808;
+  constexpr int MASK = 0x80808080;
+
+  frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
+  q >>= 4;
+  frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
+}
+
+template <>
+__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kU4B8.id(), true>(
+    int q, __nv_fp8x4_e4m3* frag_b) {
+  int s = q & 0x08080808;
+  int Out1 = ((q & 0x07070707) | (s << 4)) + (s >> 3);
+  q >>= 4;
+  s = q & 0x08080808;
+  int Out2 = ((q & 0x07070707) | (s << 4)) + (s >> 3);
+
+  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
+  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
+}
+
+template <typename scalar_t2, vllm::ScalarTypeId s_type_id>
+__device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b);
+
+template <>
+__device__ inline void dequant_fp8_scales<half2, vllm::kFE4M3fn.id()>(
+    int q, half2* frag_b) {
+  int Out1 = (q & 0xFF00FF00) >> 1;
+  ;
+  q <<= 8;
+  int Out2 = (q & 0xFF00FF00) >> 1;
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+};
+
+template <>
+__device__ inline void dequant_fp8_scales<nv_bfloat162, vllm::kFE4M3fn.id()>(
+    int q, nv_bfloat162* frag_b) {
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to BF16 format
+  int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant_fp8_scales<nv_bfloat162, vllm::kFE8M0fnu.id()>(
+    int q, nv_bfloat162* frag_b) {
+  // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16,
+  // but we assume that such a extreme value would not occur in real models.
+  int Out1 = (q & 0xFF00FF00) >> 1;
+  q <<= 7;
+  int Out2 = q & 0x7F807F80;
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+};
+
+// subtract zero point in quanted format and then dequant
+template <typename scalar_t2, vllm::ScalarTypeId w_type_id,
+          bool skip_flop = false>
+__device__ inline void sub_zp_and_dequant(int q, scalar_t2* frag_b, int zp);
+
+template <>
+__device__ inline void sub_zp_and_dequant<int32_t, vllm::kU4.id(), true>(
+    int q, int32_t* frag_b, int zp) {
+  // INT4 with zp -> INT8
+  // see https://github.com/vllm-project/vllm/pull/24722
+  int repeated_zp = 0x01010101 * zp;
+  int MASK = 0x80808080;
+
+  frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
+  q >>= 4;
+  frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
+}
+
+template <>
+__device__ inline void sub_zp_and_dequant<__nv_fp8x4_e4m3, vllm::kU4.id(),
+                                          true>(int q, __nv_fp8x4_e4m3* frag_b,
+                                                int zp) {
+  // INT4 with zp -> FP8
+  // see https://github.com/vllm-project/vllm/pull/24722
+  uint32_t u_q = *reinterpret_cast<uint32_t*>(&q);
+  uint32_t u_zp = *reinterpret_cast<uint32_t*>(&zp);
+  uint32_t u_zp1 = u_zp + 1;
+  uint32_t repeated_zp = 0x01010101 * u_zp;
+
+  uint32_t q0, s;
+  q0 = (u_q & 0x0F0F0F0F) | 0x70707070;
+  s = (q0 + repeated_zp) & 0x80808080;
+  uint32_t Out1 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s;
+
+  u_q >>= 4;
+  q0 = (u_q & 0x0F0F0F0F) | 0x70707070;
+  s = (q0 + repeated_zp) & 0x80808080;
+  uint32_t Out2 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s;
+
+  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
+  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
+}
+
+#endif
+
+}  // namespace MARLIN_NAMESPACE_NAME
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin.cuh b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin.cuh
new file mode 100644
index 0000000000..88057dd2d1
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin.cuh
@@ -0,0 +1,173 @@
+#pragma once
+
+#ifndef _marlin_cuh
+  #define _marlin_cuh
+  #include <cuda.h>
+  #include <cuda_fp16.h>
+  #include <cuda_runtime.h>
+  #include <iostream>
+
+  #ifndef MARLIN_NAMESPACE_NAME
+    #define MARLIN_NAMESPACE_NAME marlin
+  #endif
+
+namespace MARLIN_NAMESPACE_NAME {
+
+// Marlin params
+
+// 8 warps are a good choice since every SM has 4 schedulers and having more
+// than 1 warp per schedule allows some more latency hiding. At the same time,
+// we want relatively few warps to have many registers per warp and small tiles.
+static constexpr int default_threads = 256;
+
+static constexpr int pipe_stages =
+    4;  // 4 pipeline stages fit into shared memory
+
+static constexpr int min_thread_n = 64;
+static constexpr int min_thread_k = 64;
+static constexpr int max_thread_n = 256;
+
+static constexpr int tile_size = 16;
+static constexpr int max_par = 16;
+
+// Repack params
+static constexpr int repack_stages = 8;
+
+static constexpr int repack_threads = 256;
+
+static constexpr int tile_k_size = tile_size;
+static constexpr int tile_n_size = tile_k_size * 4;
+
+// Helpers
+template <typename T, int n>
+struct Vec {
+  T elems[n];
+  __device__ T& operator[](int i) { return elems[i]; }
+};
+
+using I4 = Vec<int, 4>;
+
+constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+
+__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  if (pred) {
+    reinterpret_cast<int32_t*>(smem_ptr)[0] =
+        reinterpret_cast<const int32_t*>(glob_ptr)[0];
+  }
+}
+
+__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  if (pred) {
+    reinterpret_cast<int64_t*>(smem_ptr)[0] =
+        reinterpret_cast<const int64_t*>(glob_ptr)[0];
+  }
+}
+
+__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  if (pred) {
+    reinterpret_cast<int4*>(smem_ptr)[0] =
+        reinterpret_cast<const int4*>(glob_ptr)[0];
+  }
+}
+
+__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
+                                      bool pred = true) {
+  if (pred) {
+    reinterpret_cast<int4*>(smem_ptr)[0] =
+        reinterpret_cast<const int4*>(glob_ptr)[0];
+  }
+}
+
+__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
+  reinterpret_cast<int4*>(smem_ptr)[0] =
+      reinterpret_cast<const int4*>(glob_ptr)[0];
+}
+
+__device__ inline void cp_async_fence() {}
+
+template <int n>
+__device__ inline void cp_async_wait() {}
+
+  #else
+
+__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  const int BYTES = 4;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  const int BYTES = 8;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  const int BYTES = 16;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
+                                      bool pred = true) {
+  const int BYTES = 16;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.cg.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
+  const int BYTES = 16;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   cp.async.cg.shared.global [%0], [%1], %2;\n"
+      "}\n" ::"r"(smem),
+      "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async_fence() {
+  asm volatile("cp.async.commit_group;\n" ::);
+}
+
+template <int n>
+__device__ inline void cp_async_wait() {
+  asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
+}
+
+  #endif
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+#endif
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh
new file mode 100644
index 0000000000..a4807a6887
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh
@@ -0,0 +1,149 @@
+
+#ifndef _data_types_cuh
+#define _data_types_cuh
+#include "marlin.cuh"
+#include "core/scalar_type.hpp"
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#ifndef MARLIN_NAMESPACE_NAME
+  #define MARLIN_NAMESPACE_NAME marlin
+#endif
+
+namespace MARLIN_NAMESPACE_NAME {
+
+template <long scalar_type_id>
+class MarlinScalarType {};
+
+template <>
+class MarlinScalarType<vllm::kFloat16.id()> {
+ public:
+  using scalar_t = half;
+  using scalar_t2 = half2;
+  using scalar_t4 = half2;
+  using scalar_32bit_t = half2;
+
+  // Matrix fragments for tensor core instructions; their precise layout is
+  // documented here:
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type
+  using FragA = Vec<half2, 4>;
+  using FragB = Vec<half2, 2>;
+  using FragC = Vec<float, 4>;
+  using FragS = Vec<half2, 1>;
+  using FragS0 = Vec<__nv_fp8x2_e4m3, 1>;
+  using FragZP = Vec<half2, 4>;
+
+  static __device__ float inline num2float(const half x) {
+    return __half2float(x);
+  }
+
+  static __device__ half2 inline num2num2(const half x) {
+    return __half2half2(x);
+  }
+
+  static __device__ half2 inline nums2num2(const half x1, const half x2) {
+    return __halves2half2(x1, x2);
+  }
+
+  static __host__ __device__ half inline float2num(const float x) {
+    return __float2half(x);
+  }
+
+  static __host__ __device__ float2 inline num22float2(const half2 x) {
+    return __half22float2(x);
+  }
+};
+
+template <>
+class MarlinScalarType<vllm::kBFloat16.id()> {
+ public:
+  using scalar_t = nv_bfloat16;
+  using scalar_t2 = nv_bfloat162;
+  using scalar_t4 = nv_bfloat162;
+  using scalar_32bit_t = nv_bfloat162;
+
+  using FragA = Vec<nv_bfloat162, 4>;
+  using FragB = Vec<nv_bfloat162, 2>;
+  using FragC = Vec<float, 4>;
+  using FragS = Vec<nv_bfloat162, 1>;
+  using FragS0 = Vec<__nv_fp8x2_e4m3, 1>;
+  using FragZP = Vec<nv_bfloat162, 4>;
+
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+  static __device__ float inline num2float(const nv_bfloat16 x) {
+    return __bfloat162float(x);
+  }
+
+  static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
+    return __bfloat162bfloat162(x);
+  }
+
+  static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
+                                                  const nv_bfloat16 x2) {
+    return __halves2bfloat162(x1, x2);
+  }
+
+  static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
+    return __float2bfloat16(x);
+  }
+
+  static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) {
+    return __bfloat1622float2(x);
+  }
+#endif
+};
+
+template <>
+class MarlinScalarType<vllm::kFE4M3fn.id()> {
+ public:
+  using scalar_t = __nv_fp8_e4m3;
+  using scalar_t2 = __nv_fp8x2_e4m3;
+  using scalar_t4 = __nv_fp8x4_e4m3;
+  using scalar_32bit_t = __nv_fp8x4_e4m3;
+
+  using FragA = Vec<__nv_fp8x4_e4m3, 4>;
+  using FragB = Vec<__nv_fp8x4_e4m3, 2>;
+  using FragC = Vec<float, 4>;
+  using FragZP = Vec<__nv_fp8x2_e4m3, 4>;
+
+  static __host__ __device__
+      float2 inline num22float2(const __nv_fp8x2_e4m3 x) {
+    return (float2)x;
+  }
+};
+
+template <>
+class MarlinScalarType<vllm::kS8.id()> {
+ public:
+  using scalar_t = int8_t;
+  using scalar_t2 = int16_t;
+  using scalar_t4 = int32_t;
+  using scalar_32bit_t = int32_t;
+
+  using FragA = Vec<int32_t, 4>;
+  using FragB = Vec<int32_t, 2>;
+  using FragC = Vec<float, 4>;
+  using FragZP = Vec<int16_t, 4>;
+};
+
+template <typename scalar_t>
+class MarlinScalarType2 {};
+
+template <>
+class MarlinScalarType2<half> : public MarlinScalarType<vllm::kFloat16.id()> {};
+
+template <>
+class MarlinScalarType2<nv_bfloat16>
+    : public MarlinScalarType<vllm::kBFloat16.id()> {};
+
+template <>
+class MarlinScalarType2<__nv_fp8_e4m3>
+    : public MarlinScalarType<vllm::kFE4M3fn.id()> {};
+
+template <>
+class MarlinScalarType2<int8_t> : public MarlinScalarType<vllm::kS8.id()> {};
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+#endif
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin_mma.h b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin_mma.h
new file mode 100644
index 0000000000..41c9074077
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/_csrc/libtorch_stable/quantization/marlin/marlin_mma.h
@@ -0,0 +1,269 @@
+
+#include "marlin_dtypes.cuh"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
+// output/accumulation.
+template <vllm::ScalarTypeId type_id, bool use_fp16_accum, int k_size = 16>
+__device__ inline void mma(
+    const typename MarlinScalarType<type_id>::FragA& a_frag,
+    const typename MarlinScalarType<type_id>::FragB& frag_b,
+    typename MarlinScalarType<type_id>::FragC& frag_c, int idx = 0) {
+  const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
+  const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  if constexpr (!std::is_same<scalar_t, half>::value || k_size != 16) {
+    static_assert(!use_fp16_accum);
+  }
+
+  if constexpr (k_size == 16) {
+    if constexpr (std::is_same<scalar_t, half>::value && !use_fp16_accum) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(b[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[2]), "r"(a[3]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+#else
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+#endif
+    } else if constexpr (std::is_same<scalar_t, half>::value &&
+                         use_fp16_accum) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[0]), "r"(a[1]), "r"(b[0]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[2]), "r"(a[3]), "r"(b[1]), "r"(c[0]), "r"(c[1]));
+#else
+      uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "r"(c[0]), "r"(c[1]));
+#endif
+    } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "f"(c[0]),
+            "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
+      int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
+          : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "r"(c[0]),
+            "r"(c[1]), "r"(c[2]), "r"(c[3]));
+    }
+  } else if (k_size == 32) {
+    if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
+      int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[0]), "r"(b[0]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[2]), "=r"(c[3])
+          : "r"(a[1]), "r"(b[0]), "r"(c[2]), "r"(c[3]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[2]), "r"(b[1]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[2]), "=r"(c[3])
+          : "r"(a[3]), "r"(b[1]), "r"(c[2]), "r"(c[3]));
+#else
+      asm volatile(
+          "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
+#endif
+    }
+  }
+}
+
+template <vllm::ScalarTypeId type_id, bool use_fp16_accum, int k_size = 16>
+__device__ inline void mma_trans(
+    const typename MarlinScalarType<type_id>::FragA& a_frag,
+    const typename MarlinScalarType<type_id>::FragB& frag_b,
+    const typename MarlinScalarType<type_id>::FragB& frag_b2,
+    typename MarlinScalarType<type_id>::FragC& frag_c) {
+  const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
+  const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
+  const uint32_t* b2 = reinterpret_cast<const uint32_t*>(&frag_b2);
+  float* c = reinterpret_cast<float*>(&frag_c);
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  if constexpr (!std::is_same<scalar_t, half>::value || k_size != 16) {
+    static_assert(!use_fp16_accum);
+  }
+
+  if constexpr (k_size == 16) {
+    if constexpr (std::is_same<scalar_t, half>::value && !use_fp16_accum) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+#else
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+#endif
+    } else if constexpr (std::is_same<scalar_t, half>::value &&
+                         use_fp16_accum) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "r"(c[0]), "r"(c[1]));
+#else
+      uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "r"(c[0]), "r"(c[1]));
+#endif
+    } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
+      int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]),
+            "r"(c[3]));
+    }
+  } else {
+    if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
+      int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[2]), "=r"(c[3])
+          : "r"(b2[1]), "r"(a[0]), "r"(c[2]), "r"(c[3]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[0]), "r"(a[1]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[2]), "=r"(c[3])
+          : "r"(b2[1]), "r"(a[1]), "r"(c[2]), "r"(c[3]));
+#else
+      asm volatile(
+          "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
+#endif
+    }
+  }
+}
+
+}  // namespace MARLIN_NAMESPACE_NAME
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/backend.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/backend.py
new file mode 100644
index 0000000000..49f4826cb4
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/backend.py
@@ -0,0 +1,284 @@
+"""Marlin W4A16 forward base for the grouped NVFP4 MoE (Ampere+; sm80/89/90/100/120).
+
+Prep the frozen NVFP4 experts into Marlin layout ONCE (``build_marlin_forward_base``); run the
+bf16-activation forward GEMM (``marlin_base_forward``) on the pre-scattered, pad-to-TILE grouped
+layout that ``grouped_fp4_moe_train`` already builds. Marlin reads the activation directly in bf16
+and gathers per expert via identity ``sorted_token_ids`` + per-block ``expert_ids`` (derived from the
+per-tile ``m_indices``) + ``top_k=1`` — validated bit-correct on that layout.
+
+The marlin path pads to TILE=64 (vs CUTLASS's forced 128): at thin-M that halves the padded rows
+(2.7x -> 1.4x), which is where the speedup comes from (each expert weight is read once either way).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from . import load_ext
+from .prep import (
+    marlin_make_workspace_new,
+    marlin_moe_gemm,
+    prepare_nvfp4_weight_for_marlin,
+)
+
+# Must equal the marlin MoE block (BSM) so each padded tile maps to exactly one expert; 64
+# halves the thin-M padding vs CUTLASS's 128.
+MARLIN_TILE = 64
+_BSM = 64
+
+_BASE_SCATTER_LUT: dict[torch.device, torch.Tensor] = {}
+
+
+def _pt(nv, E, dev):
+    p = getattr(nv, "per_tensor_scale", None)
+    if p is None:
+        return torch.ones(E, device=dev)
+    p = p.reshape(-1).float()
+    return p.expand(E) if p.numel() == 1 else p
+
+
+def _build_base_scatter(dev: torch.device) -> torch.Tensor:
+    """Build and cache the 1024-entry scatter LUT: qdata flat nibble pos -> marlin flat nibble pos.
+    Probes the base tile (N_t=64, K_t=16) using gptq_marlin_repack; result is device-resident."""
+    if dev in _BASE_SCATTER_LUT:
+        return _BASE_SCATTER_LUT[dev]
+    ext = load_ext()
+    perm = torch.empty(0, dtype=torch.int, device=dev)
+    N_t, K_t = 64, 16
+    scatter = torch.full((N_t * K_t,), -1, dtype=torch.int32, device="cpu")
+    for n_b in range(N_t):
+        for k_b in range(K_t):
+            qd = torch.zeros(N_t, K_t // 2, dtype=torch.uint8, device=dev)
+            byte_idx = k_b // 2
+            qd[n_b, byte_idx] = 0x0F if k_b % 2 == 0 else 0xF0
+            qw_i = qd.view(torch.int32).T.contiguous()
+            qw_m = ext.gptq_marlin_repack(qw_i, perm, K_t, N_t, 4, False)
+            flat = qw_m.reshape(-1).cpu().tolist()
+            for wi, word in enumerate(flat):
+                if word == 0:
+                    continue
+                for bit in range(8):
+                    if (word >> (bit * 4)) & 0xF == 0xF:
+                        scatter[n_b * K_t + k_b] = wi * 8 + bit
+                        break
+                else:
+                    continue
+                break
+    lut = scatter.to(dev)
+    _BASE_SCATTER_LUT[dev] = lut
+    return lut
+
+
+def _marlin_cache_under_fsdp(device) -> bool:
+    """Whether to persist marlin-prepped experts in the per-layer module cache under FSDP.
+
+    DEFAULT: NO. The per-layer cache holds each layer's gathered (un-sharded) marlin experts, so it
+    accumulates the whole model across depth (~0.6 GiB/proj × 2 × num_layers — e.g. ~94 GiB for GLM-5.2)
+    on top of the resident sharded model → OOM. Profiled: a layer's marlin re-prep is only a ~17 GiB
+    transient that frees immediately, so re-prepping every forward is far cheaper than caching, and the
+    per-layer MoE fwd+bwd peak is ~17 GiB — fits easily beside the ~60 GiB resident shard.
+
+    OPT-IN for roomy GPUs (e.g. B200, or small models) where the full cache fits: set
+    ``AXOLOTL_MARLIN_CACHE_MIN_FREE_GB=<N>`` to cache while ≥ N GiB VRAM is free (``0`` = always cache).
+    Unset → never cache under FSDP."""
+    import os
+
+    val = os.environ.get("AXOLOTL_MARLIN_CACHE_MIN_FREE_GB")
+    if val is None:
+        return False
+    try:
+        free, _ = torch.cuda.mem_get_info(device)
+    except Exception:  # pylint: disable=broad-except
+        return True
+    return free > float(val) * (1024**3)
+
+
+def _cached_prep(nv, size_n, size_k, ext, cache, key):
+    """Prep one NVFP4 weight -> Marlin layout, cached (experts are frozen). Prefer a persistent
+    module-level ``cache`` dict keyed by ``key`` — under FSDP2 the gathered param is a fresh tensor
+    each step, so a module cache (not a per-tensor attr) avoids re-repacking 256 experts/forward.
+    Falls back to a per-tensor attribute, then a one-shot compute.
+
+    On the first build (single-GPU only), saves original scales + per-tensor scale in cache under
+    key+"_bwd" for the fused backward dequant, then frees nv.qdata to drop the duplicate 4-bit copy.
+    Under FSDP the gathered param is fresh each step so weight_recipe() re-gathers it anyway; skips
+    the free when distributed is active."""
+    import torch.distributed as dist
+
+    _is_distributed = (
+        dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1
+    )
+    bwd_key = key + "_bwd"
+    if cache is not None and key in cache:
+        return cache[key]
+    assert not getattr(nv, "is_swizzled_scales", False), (
+        "marlin_w4a16 needs raw (non-swizzled) NVFP4 scales; got is_swizzled_scales=True"
+    )
+    cached = getattr(nv, "_marlin_w4a16", None)
+    qdata_fresh = cached is None
+    if cached is None:
+        E, dev = nv.qdata.size(0), nv.qdata.device
+        cached = prepare_nvfp4_weight_for_marlin(
+            nv.qdata,
+            nv.scale,
+            _pt(nv, E, dev),
+            size_n,
+            size_k,
+            torch.bfloat16,
+            ext.gptq_marlin_repack,
+        )
+        # Stash the full-layer marlin weight on the gathered tensor ONLY single-GPU. Under FSDP2 the
+        # gathered param is a fresh tensor every forward AND every backward recompute, so this attr
+        # never produces a cache hit — it only pins ~9 GiB/proj alive on each layer's gathered tensor,
+        # and during the backward sweep FSDP keeps several layers' gathered tensors live at once →
+        # the marlin weights accumulate across depth and OOM (bypassing the module-cache guard below).
+        if not _is_distributed:
+            try:
+                nv._marlin_w4a16 = cached
+            except (AttributeError, RuntimeError):
+                pass
+    if cache is not None and not _is_distributed:
+        # Single-GPU: persist the marlin weight and drop the duplicate 4-bit qdata.
+        cache[key] = cached
+        if qdata_fresh and bwd_key not in cache:
+            E, dev = nv.qdata.size(0), nv.qdata.device
+            # Reference (don't clone) the original scales for the backward dequant: the
+            # marlin scales subsample [:, 1::2] so they're lossy and unusable for backward.
+            # Single-GPU: nv.scale stays live on the frozen param, so a clone just dups ~1.3 GiB.
+            # _pt may return .expand() with stride(0)==0; Triton needs a real stride.
+            pt_bwd = _pt(nv, E, dev)
+            if pt_bwd.stride(0) == 0:
+                pt_bwd = pt_bwd.contiguous()
+            cache[bwd_key] = (nv.scale, pt_bwd, E, size_n, size_k)
+            # 3-D [E, N, 0] placeholder (not flat empty(0)): NVFP4Tensor clone/save needs
+            # ndim==3, and numel()==0 still routes through the marlin qdata-free paths.
+            try:
+                _N = nv.qdata.size(1)
+                nv.qdata.data = torch.empty(
+                    (E, _N, 0), dtype=nv.qdata.dtype, device=dev
+                )
+            except (AttributeError, RuntimeError):
+                pass
+    elif cache is not None and _marlin_cache_under_fsdp(nv.qdata.device):
+        # FSDP2: the gathered weight is a fresh, transient tensor per forward (meant to reshard), so
+        # persisting its marlin copy for EVERY layer re-materializes the whole un-sharded model on
+        # each rank — OOMing large checkpoints. The cache is still a real perf win (no re-repacking
+        # 256 experts/forward) when VRAM is roomy, so cache only while there's headroom; once it
+        # tightens, fall through to per-forward re-prep (the transient ``_marlin_w4a16`` attr handles
+        # in-forward reuse and dies on reshard). Don't free qdata here — backward re-gathers it.
+        cache[key] = cached
+    return cached
+
+
+def _qdata_sizes(nv, cache_key, cache):
+    """Return (size_n, size_k_packed, device) for the NVFP4 weight.
+    Falls back to bwd cache when nv.qdata has been freed (single-GPU memory-free path)."""
+    qdata = nv.qdata
+    if qdata.numel() > 0:
+        return qdata.size(1), qdata.size(2), torch.device(qdata.device)
+    if cache is not None:
+        bwd = cache.get(cache_key + "_bwd")
+        if bwd is not None:
+            # bwd = (scale, pt, E, size_n, size_k); size_k_packed = size_k // 2
+            size_n, size_k = bwd[3], bwd[4]
+            dev = bwd[0].device
+            return size_n, size_k // 2, dev
+    raise RuntimeError(
+        f"marlin_w4a16: nv.qdata has been freed but no backward cache found for '{cache_key}'. "
+        "Pass a module-level mxfp4_cache to grouped_fp4_moe_train."
+    )
+
+
+def build_marlin_forward_base(gate_up_nv, down_nv, cache=None):
+    """Prep frozen NVFP4 experts -> Marlin forward weights (cached; experts are frozen).
+
+    gate_up_nv: NVFP4Tensor [E, 2I, H]; down_nv: NVFP4Tensor [E, H, I]. ``cache`` is the optional
+    persistent dict (the same ``mxfp4_cache`` threaded through ``grouped_fp4_moe_train``).
+    Returns the base tuple ``('marlin', (gu_w, dn_w), (twoI, H, I), workspace, cache)`` consumed by
+    ``marlin_base_forward`` and dispatched in ``grouped_train._base_forward``. The cache is threaded
+    through so the backward can retrieve the (scale, pt, E, N, K) bwd tuple from cache[key+"_bwd"]."""
+    ext = load_ext()
+    twoI, H_packed, dev = _qdata_sizes(gate_up_nv, "marlin_gate_up", cache)
+    H = H_packed * 2
+    _, I_packed, _ = _qdata_sizes(down_nv, "marlin_down", cache)
+    I = I_packed * 2  # noqa: E741
+    gu_w = _cached_prep(gate_up_nv, twoI, H, ext, cache, "marlin_gate_up")
+    dn_w = _cached_prep(down_nv, H, I, ext, cache, "marlin_down")
+    ws = marlin_make_workspace_new(dev, 4)
+    return ("marlin", (gu_w, dn_w), (twoI, H, I), ws, cache)
+
+
+def marlin_route(m_indices, Mt, dev):
+    """Build Marlin's routing for the pre-scattered pad-TILE layout: identity ``sorted_token_ids``
+    (the activation is already grouped+padded), per-BSM-block ``expert_ids`` from the per-tile
+    ``m_indices``, ``num_tokens_post_padded`` = Mt, and unit ``topk_weights`` (top_k=1)."""
+    tile = Mt // m_indices.numel()
+    si = torch.arange(Mt, dtype=torch.int32, device=dev)
+    ei = m_indices.repeat_interleave(tile // _BSM).to(torch.int32)
+    ntpp = torch.tensor([Mt], dtype=torch.int32, device=dev)
+    tw = torch.ones(Mt, 1, dtype=torch.float32, device=dev)
+    return si, ei, ntpp, tw
+
+
+def marlin_base_forward(base, which, x, m_indices):
+    """Frozen-expert base GEMM via Marlin (bf16 act). which=0 gate_up (x[Mt,H] -> [Mt,2I]);
+    which=1 down (x[Mt,I] -> [Mt,H]). ``m_indices`` is the per-TILE expert id from the routing.
+    Returns bf16."""
+    ext = load_ext()
+    _, (gu_w, dn_w), (twoI, H, I), ws = base[0], base[1], base[2], base[3]  # noqa: E741
+    Mt = x.size(0)
+    si, ei, ntpp, tw = marlin_route(m_indices, Mt, x.device)
+    if which == 0:
+        return marlin_moe_gemm(
+            ext,
+            x,
+            gu_w[0],
+            gu_w[1],
+            gu_w[2],
+            ws,
+            si,
+            ei,
+            ntpp,
+            tw,
+            _BSM,
+            1,
+            False,
+            Mt,
+            twoI,
+            H,
+        )
+    return marlin_moe_gemm(
+        ext,
+        x,
+        dn_w[0],
+        dn_w[1],
+        dn_w[2],
+        ws,
+        si,
+        ei,
+        ntpp,
+        tw,
+        _BSM,
+        1,
+        False,
+        Mt,
+        H,
+        I,
+    )
+
+
+def marlin_bwd_data(base):
+    """Return (gu_bwd, dn_bwd) tuples from the cache, or None if not available.
+    Each bwd tuple is (original_scale, pt, E, size_n, size_k); marlin qw is at cache[fwd_key][0]."""
+    cache = base[4] if len(base) > 4 else None
+    if cache is None:
+        return None
+    gu_bwd = cache.get("marlin_gate_up_bwd")
+    dn_bwd = cache.get("marlin_down_bwd")
+    if gu_bwd is None or dn_bwd is None:
+        return None
+    gu_qw = cache.get("marlin_gate_up")
+    dn_qw = cache.get("marlin_down")
+    if gu_qw is None or dn_qw is None:
+        return None
+    return (gu_qw[0], gu_bwd), (dn_qw[0], dn_bwd)
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/fused_dequant.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/fused_dequant.py
new file mode 100644
index 0000000000..aff1e715c4
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/fused_dequant.py
@@ -0,0 +1,179 @@
+"""Fused Triton kernels: marlin int32 layout -> bf16 / fp8_e4m3, without intermediate reconstruction.
+
+The backward dequant in the marlin training path needs to reconstruct the expert weights from the
+marlin qweight cache (int32, tile-scattered nibbles) + the original NVFP4 block scales. A two-step
+approach (reconstruct nibble array + call nvfp4_dequant_bf16) is ~12x slower due to the scattered
+gather; these kernels fold both steps into one Triton pass with only sequential reads on qweight and
+sequential writes on output, bringing cost to ~1.2-1.6x the direct nvfp4_dequant.
+
+The marlin tile layout (N_t=64, K_t=16) scatters nibbles according to gptq_marlin_repack. A 1024-entry
+LUT (base scatter, built once per process from _build_base_scatter in backend.py) maps each
+(local_n, local_k) position within a base tile to its flat nibble position in the marlin output.
+Larger (N, K) tile: offset = (k_tile * N_TILES + n_tile) * 1024 nibbles.
+"""
+
+from __future__ import annotations
+
+import triton
+import triton.language as tl
+
+# Grid (C, N // 64, K // 16): one block per (expert-chunk, n-tile, k-tile), each a 64x16 tile.
+
+
+@triton.jit
+def _marlin_dequant_bf16_kernel(
+    QW,  # [C, K*N//8] int32  (marlin flat)
+    SC,  # [C, N, K//16] fp8_e4m3
+    PT,  # [C] float32  per-tensor scale
+    SCATTER,  # [1024] int32  base scatter LUT
+    CB,  # [16] float32  fp4 codebook
+    OUT,  # [C, N, K] bfloat16
+    qw_stride_c: tl.constexpr,
+    sc_stride_c: tl.constexpr,
+    sc_stride_n: tl.constexpr,
+    out_stride_c: tl.constexpr,
+    out_stride_n: tl.constexpr,
+    N_TILES: tl.constexpr,
+    N_t: tl.constexpr,  # 64
+    K_t: tl.constexpr,  # 16
+    TILE_SIZE: tl.constexpr,  # N_t * K_t = 1024
+):
+    c_id = tl.program_id(0)
+    n_tile = tl.program_id(1)
+    k_tile = tl.program_id(2)
+
+    local_idx = tl.arange(0, TILE_SIZE)
+    local_n = local_idx // K_t
+    local_k = local_idx % K_t
+
+    local_marlin_pos = tl.load(SCATTER + local_n * K_t + local_k)
+    tile_off = (k_tile * N_TILES + n_tile) * TILE_SIZE
+    nib_pos = tile_off + local_marlin_pos
+
+    words = tl.load(QW + c_id * qw_stride_c + nib_pos // 8)
+    nibs = (words >> ((nib_pos % 8) * 4)) & 0xF
+
+    pt = tl.load(PT + c_id)
+    global_n = n_tile * N_t + local_n
+    sc = tl.load(SC + c_id * sc_stride_c + global_n * sc_stride_n + k_tile).to(
+        tl.float32
+    )
+    cb_val = tl.load(CB + nibs)
+
+    out_val = (cb_val * sc * pt).to(tl.bfloat16)
+    tl.store(
+        OUT + c_id * out_stride_c + global_n * out_stride_n + k_tile * K_t + local_k,
+        out_val,
+    )
+
+
+@triton.jit
+def _marlin_dequant_fp8_kernel(
+    QW,
+    SC,
+    PT,
+    SCATTER,
+    CB,
+    OUT,
+    qw_stride_c: tl.constexpr,
+    sc_stride_c: tl.constexpr,
+    sc_stride_n: tl.constexpr,
+    out_stride_c: tl.constexpr,
+    out_stride_n: tl.constexpr,
+    N_TILES: tl.constexpr,
+    N_t: tl.constexpr,
+    K_t: tl.constexpr,
+    TILE_SIZE: tl.constexpr,
+):
+    c_id = tl.program_id(0)
+    n_tile = tl.program_id(1)
+    k_tile = tl.program_id(2)
+
+    local_idx = tl.arange(0, TILE_SIZE)
+    local_n = local_idx // K_t
+    local_k = local_idx % K_t
+
+    local_marlin_pos = tl.load(SCATTER + local_n * K_t + local_k)
+    tile_off = (k_tile * N_TILES + n_tile) * TILE_SIZE
+    nib_pos = tile_off + local_marlin_pos
+
+    words = tl.load(QW + c_id * qw_stride_c + nib_pos // 8)
+    nibs = (words >> ((nib_pos % 8) * 4)) & 0xF
+
+    pt = tl.load(PT + c_id)
+    global_n = n_tile * N_t + local_n
+    sc = tl.load(SC + c_id * sc_stride_c + global_n * sc_stride_n + k_tile).to(
+        tl.float32
+    )
+    cb_val = tl.load(CB + nibs)
+
+    out_val = (cb_val * sc * pt).to(tl.float8e4nv)
+    tl.store(
+        OUT + c_id * out_stride_c + global_n * out_stride_n + k_tile * K_t + local_k,
+        out_val,
+    )
+
+
+def marlin_dequant_bf16(qw_flat, orig_scale, pt, scatter_lut, cb, N, K, C):
+    """Dequantize marlin int32 weight -> bf16 using the original NVFP4 block scales.
+
+    qw_flat: [C, K*N//8] int32  (marlin flat layout)
+    orig_scale: [C, N, K//16] fp8_e4m3  (original block scales, NOT marlin-processed)
+    pt: [C] float32  per-tensor scale
+    scatter_lut: [1024] int32  base scatter LUT (from _build_base_scatter)
+    cb: [16] float32  fp4 codebook
+    Returns: [C, N, K] bfloat16"""
+    import torch
+
+    out = torch.empty(C, N, K, device=qw_flat.device, dtype=torch.bfloat16)
+    N_t, K_t = 64, 16
+    N_TILES = N // N_t
+    grid = (C, N_TILES, K // K_t)
+    _marlin_dequant_bf16_kernel[grid](
+        qw_flat,
+        orig_scale,
+        pt,
+        scatter_lut,
+        cb,
+        out,
+        int(qw_flat.stride(0)),
+        int(orig_scale.stride(0)),
+        int(orig_scale.stride(1)),
+        int(out.stride(0)),
+        int(out.stride(1)),
+        N_TILES,
+        N_t,
+        K_t,
+        N_t * K_t,
+    )
+    return out
+
+
+def marlin_dequant_fp8(qw_flat, orig_scale, pt, scatter_lut, cb, N, K, C):
+    """Dequantize marlin int32 weight -> fp8_e4m3 using the original NVFP4 block scales.
+
+    Same signature as marlin_dequant_bf16; output dtype is float8_e4m3fn."""
+    import torch
+
+    out = torch.empty(C, N, K, device=qw_flat.device, dtype=torch.float8_e4m3fn)
+    N_t, K_t = 64, 16
+    N_TILES = N // N_t
+    grid = (C, N_TILES, K // K_t)
+    _marlin_dequant_fp8_kernel[grid](
+        qw_flat,
+        orig_scale,
+        pt,
+        scatter_lut,
+        cb,
+        out,
+        int(qw_flat.stride(0)),
+        int(orig_scale.stride(0)),
+        int(orig_scale.stride(1)),
+        int(out.stride(0)),
+        int(out.stride(1)),
+        N_TILES,
+        N_t,
+        K_t,
+        N_t * K_t,
+    )
+    return out
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/nonexpert_linear.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/nonexpert_linear.py
new file mode 100644
index 0000000000..d6163a41a6
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/nonexpert_linear.py
@@ -0,0 +1,120 @@
+"""Frozen NVFP4 W4A16 dense ``Linear`` via the Marlin grouped kernel (single expert).
+
+Non-expert quantization: swap a frozen ``nn.Linear`` for 4-bit NVFP4 weights that run through the
+SAME validated Marlin W4A16 kernel the grouped MoE experts use — 4-bit weight memory + bf16 tensor-
+core compute on any sm80+ GPU (no FP4 hardware needed). A dense Linear is just the grouped GEMM with
+one expert (E=1): all rows route to expert 0, M padded up to the Marlin block (64).
+
+Weights are quantized once (at swap time, before FSDP wrap) via torchao's two-level NVFP4 quant and
+prepped into Marlin layout eagerly; the bf16 weight is then dropped. The module is intentionally NOT
+an ``nn.Linear`` subclass so PEFT leaves it frozen (non-experts carry no LoRA).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from . import load_ext
+from .prep import (
+    marlin_make_workspace_new,
+    marlin_moe_gemm,
+    prepare_nvfp4_weight_for_marlin,
+)
+
+_BSM = 64  # Marlin MoE block size; M is padded to a multiple of this
+
+
+class MarlinW4A16Linear(nn.Module):
+    """Frozen NVFP4 W4A16 replacement for ``nn.Linear`` (bf16 act/out, 4-bit weight)."""
+
+    def __init__(self, weight: torch.Tensor, bias: torch.Tensor | None):
+        super().__init__()
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+        self.out_features, self.in_features = weight.shape
+        # The Marlin repack is CUDA-only; non-expert weights may still be on CPU at quant time.
+        if weight.device.type != "cuda":
+            weight = weight.to("cuda")
+        dev = weight.device
+        ext = load_ext()
+
+        # Two-level NVFP4 quant (per-tensor fp32 + per-16 e4m3 block scales), one [1,N,K] expert.
+        # Per-tensor scale needed for accuracy (rel err ~0.10 two-level vs ~0.25 without):
+        # pts = amax / (F4_E2M1_MAX * F8E4M3_MAX) = amax / (6 * 448).
+        w_bf16 = weight.to(torch.bfloat16)
+        pts = (w_bf16.abs().max() / (6.0 * 448.0)).reshape(1).float()
+        nv = NVFP4Tensor.to_nvfp4(
+            w_bf16[None], per_tensor_scale=pts, is_swizzled_scales=False
+        )
+        gscale = nv.per_tensor_scale
+        qw, sc, g = prepare_nvfp4_weight_for_marlin(
+            nv.qdata,
+            nv.scale,
+            gscale,
+            self.out_features,
+            self.in_features,
+            torch.bfloat16,
+            ext.gptq_marlin_repack,
+        )
+        # frozen -> buffers (move with .to()/.cuda(); FSDP replicates non-experts)
+        self.register_buffer("qweight", qw)  # [1, ...] int32 marlin layout
+        self.register_buffer("scales", sc)  # [1, ...] bf16
+        self.register_buffer("gscale", g)  # [1] fp32
+        self.register_buffer(
+            "bias", bias.to(torch.bfloat16) if bias is not None else None
+        )
+        # Device-specific scratch; persistent=False keeps it out of state_dict/checkpoints.
+        self.register_buffer(
+            "_workspace", marlin_make_workspace_new(dev, 4), persistent=False
+        )
+        self._route_cache: dict[tuple[int, torch.device], tuple] = {}
+
+    def _route(self, Mt: int, dev: torch.device):
+        key = (Mt, dev)
+        cached = self._route_cache.get(key)
+        if cached is None:
+            si = torch.arange(Mt, dtype=torch.int32, device=dev)
+            ei = torch.zeros(
+                Mt // _BSM, dtype=torch.int32, device=dev
+            )  # all -> expert 0
+            ntpp = torch.tensor([Mt], dtype=torch.int32, device=dev)
+            tw = torch.ones(Mt, 1, dtype=torch.float32, device=dev)
+            cached = (si, ei, ntpp, tw)
+            self._route_cache[key] = cached
+        return cached
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        ext = load_ext()
+        orig = x.shape
+        xf = x.reshape(-1, self.in_features).to(torch.bfloat16)
+        M = xf.shape[0]
+        Mt = (M + _BSM - 1) // _BSM * _BSM
+        if Mt != M:
+            xf = torch.nn.functional.pad(xf, (0, 0, 0, Mt - M))
+        si, ei, ntpp, tw = self._route(Mt, xf.device)
+        out = marlin_moe_gemm(
+            ext,
+            xf,
+            self.qweight,
+            self.scales,
+            self.gscale,
+            self._workspace,
+            si,
+            ei,
+            ntpp,
+            tw,
+            _BSM,
+            1,
+            False,
+            Mt,
+            self.out_features,
+            self.in_features,
+        )
+        out = out[:M]
+        if self.bias is not None:
+            out = out + self.bias
+        return out.reshape(*orig[:-1], self.out_features)
+
+    def extra_repr(self) -> str:
+        return f"in_features={self.in_features}, out_features={self.out_features}, nvfp4=W4A16-marlin"
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/prep.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/prep.py
new file mode 100644
index 0000000000..abf1b49620
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/marlin_w4a16/prep.py
@@ -0,0 +1,182 @@
+"""vLLM-FREE NVFP4 -> Marlin weight prep. Pure-torch scale processing (copied verbatim from
+vLLM marlin_utils.py / marlin_utils_fp4.py so it is bit-identical) + the standalone CUDA
+gptq_marlin_repack (ported, bit-exact to vLLM). Produces (b_qweight int32, b_scales fp8_e4m3,
+global_scale fp32) ready for the standalone moe_wna16_marlin_gemm.
+
+One function does ONE weight tensor [E, size_n, size_k//2] (gate_up fused or down). The training
+integration calls it 4x: gate_up + down (forward), and their transposes (backward dX).
+"""
+
+import torch
+
+
+# Copied verbatim from vllm marlin_utils.py (keep bit-identical).
+def _num_compute_units(idx):
+    return torch.cuda.get_device_properties(idx).multi_processor_count
+
+
+def marlin_make_workspace_new(device, max_blocks_per_sm=1):
+    sms = _num_compute_units(device.index)
+    return torch.zeros(
+        sms * max_blocks_per_sm, dtype=torch.int, device=device, requires_grad=False
+    )
+
+
+def get_scale_perms():
+    scale_perm = []
+    for i in range(8):
+        scale_perm.extend([i + 8 * j for j in range(8)])
+    scale_perm_single = []
+    for i in range(4):
+        scale_perm_single.extend([2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]])
+    return scale_perm, scale_perm_single
+
+
+def marlin_permute_scales(s, size_k, size_n, group_size, is_a_8bit=False):
+    scale_perm, scale_perm_single = get_scale_perms()
+    if group_size < size_k and group_size != -1 and not is_a_8bit:
+        s = s.reshape((-1, len(scale_perm)))[:, scale_perm]
+    else:
+        s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single]
+    s = s.reshape((-1, size_n)).contiguous()
+    return s
+
+
+# Copied verbatim from vllm marlin_utils_fp4.py (keep bit-identical).
+def _nvfp4_compute_scale_factor(marlin_scales, a_dtype=None):
+    if a_dtype is not None and a_dtype == torch.half:
+        return 1.0
+    # Value-identical, memory-cheap form of vllm's `ws_float[ws_float>0].max()`: the max of the
+    # positive entries equals the overall amax whenever any positive exists. Computing amax in the
+    # native (bf16) dtype avoids materializing a full float copy AND a boolean-index copy of the
+    # [E,N,K/16] scales (~9 GiB at GLM-5.2 dims) — that copy OOMed the gradient-checkpoint backward
+    # recompute. bf16 amax selects the exact max element; the *128 scaling is then a scalar float op.
+    max_val = marlin_scales.amax().float() * (2**7)
+    if max_val > 0 and max_val < 448 * (2**7):
+        sf = (448 * (2**7) / max_val).log2().floor().exp2()
+        return sf.item()
+    return 1.0
+
+
+def nvfp4_marlin_process_scales(marlin_scales, scale_factor=None, a_dtype=None):
+    marlin_scales = marlin_scales.to(torch.half)
+    marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
+        marlin_scales.size(0), -1
+    )
+    if scale_factor is None:
+        scale_factor = _nvfp4_compute_scale_factor(marlin_scales, a_dtype)
+    if scale_factor > 1.0:
+        marlin_scales = (marlin_scales.float() * scale_factor).to(torch.half)
+    marlin_scales = marlin_scales * (2**7)
+    marlin_scales[marlin_scales < 2] = 0
+    marlin_scales = marlin_scales.view(torch.int16) << 1
+    marlin_scales = marlin_scales.view(torch.float8_e4m3fn)
+    marlin_scales = marlin_scales[:, 1::2].contiguous()
+    return marlin_scales, scale_factor
+
+
+def nvfp4_marlin_process_global_scale(global_scale, a_dtype=None):
+    if a_dtype is None:
+        a_dtype = global_scale.dtype
+    assert a_dtype in [torch.half, torch.bfloat16]
+    fp4_exponent = 2
+    target_exponent = 5 if a_dtype == torch.half else 8
+    exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp4_exponent - 1)
+    return global_scale * (2.0 ** (exponent_bias - 7))
+
+
+# vllm::ScalarType float4_e2m1f.id (stable packed enum id).
+FLOAT4_E2M1F_ID = 562949953487106
+
+
+def marlin_moe_gemm(
+    ext,
+    a,
+    b_qweight,
+    b_scales,
+    global_scale,
+    workspace,
+    sorted_token_ids,
+    expert_ids,
+    num_tokens_past_padded,
+    topk_weights,
+    moe_block_size,
+    top_k,
+    mul_topk_weights,
+    size_m,
+    size_n,
+    size_k,
+    out=None,
+):
+    """Thin vLLM-free call into the standalone moe_wna16_marlin_gemm (NVFP4 W4A16, bf16 act/out)."""
+    if out is None:
+        out = torch.empty((size_m * top_k, size_n), device=a.device, dtype=a.dtype)
+    return ext.moe_wna16_marlin_gemm(
+        a,
+        out,
+        b_qweight,
+        None,
+        b_scales,
+        None,
+        global_scale,
+        None,
+        None,
+        None,
+        workspace,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_past_padded,
+        topk_weights,
+        moe_block_size,
+        top_k,
+        mul_topk_weights,
+        FLOAT4_E2M1F_ID,
+        size_m,
+        size_n,
+        size_k,
+        True,
+        False,
+        True,
+        False,
+        -1,
+        -1,
+        -1,
+    )
+
+
+GROUP_SIZE = 16
+
+
+def prepare_nvfp4_weight_for_marlin(
+    qdata, block_scale, global_scale, size_n, size_k, param_dtype, repack_fn
+):
+    """qdata[E,size_n,size_k//2] uint8, block_scale[E,size_n,size_k//16] e4m3,
+    global_scale scalar/[E] fp32. Returns (qw[E,...] int32, scales[E,...] fp8, gscale fp32)."""
+    E = qdata.size(0)
+    dev = qdata.device
+    perm = torch.empty(0, dtype=torch.int, device=dev)
+
+    qw = []
+    for i in range(E):
+        qw_i = qdata[i].view(torch.int32).T.contiguous()  # [size_k//8, size_n]
+        qw.append(repack_fn(qw_i, perm, size_k, size_n, 4, False))
+    qw = torch.stack(qw, 0)
+
+    # Shared scale_factor across experts, then per-expert permute+process.
+    scales_p = block_scale.to(param_dtype)
+    csf = _nvfp4_compute_scale_factor(scales_p, param_dtype)
+    sc = []
+    for i in range(E):
+        s = scales_p[i].T  # [size_k//16, size_n]
+        ms = marlin_permute_scales(
+            s, size_k=size_k, size_n=size_n, group_size=GROUP_SIZE
+        )
+        ms, _ = nvfp4_marlin_process_scales(ms, scale_factor=csf, a_dtype=param_dtype)
+        sc.append(ms)
+    sc = torch.stack(sc, 0)
+
+    g = (
+        nvfp4_marlin_process_global_scale(global_scale.to(torch.float32), param_dtype)
+        / csf
+    )
+    return qw, sc, g.float()
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/metadata.json b/src/axolotl/integrations/kernels/libs/scattermoe_lora/metadata.json
new file mode 100644
index 0000000000..fbe86bca2b
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/metadata.json
@@ -0,0 +1,13 @@
+{
+  "id": "scattermoe-lora",
+  "name": "scattermoe-lora",
+  "version": 1,
+  "license": "Apache-2.0",
+  "upstream": null,
+  "source": null,
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": []
+  }
+}
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/multi_lora.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/multi_lora.py
new file mode 100644
index 0000000000..4f5480a60e
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/multi_lora.py
@@ -0,0 +1,455 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""ScatterMoE + multi-adapter LoRA.
+
+Extends the single-adapter ScatterMoE+LoRA path to co-train many LoRA adapters
+("tenants") over one frozen expert stack. An MoE token carries two routings: the
+router's expert assignment and a per-row tenant id. The base expert GEMM keys on
+the expert; the LoRA keys on the combined ``(expert, tenant)`` group, so LoRA is
+stacked over ``E*T`` groups (``A:[E*T*R, K]``, ``B:[N, E*T*R]``).
+
+The key trick: ``combined = expert * T + tenant`` is also sorted by expert, so a
+single sort produces both the expert grouping (base) and the combined grouping
+(LoRA). The base GEMM and both LoRA projections (forward and the dX/dA/dB
+backward) are grouped GEMMs over the combined groups, expressed with the existing
+``scatter2scatter`` op plus one small grouped-Gram kernel (``_grouped_gram_kernel``)
+for the weight gradients.
+
+The backward precomputes ``XA = X @ A^T`` (saved from the forward) and
+``YB = dY @ B`` (computed once, shared by dX and dA) so dA/dB become plain grouped
+Gram products with no per-output-block recompute -- the cost that made the shared
+single-adapter split kernel scale poorly as tenant count (E*T groups) grows.
+
+Single-adapter (T == 1) reduces to the original behaviour exactly.
+"""
+
+from itertools import product
+
+import torch
+import triton
+import triton.language as tl
+
+from .kernels import lora_ops, ops as base_ops
+from .kernels.lora_ops import _block_r_for_rank, _bucket_m
+from .parallel_experts import compileable_bincount
+
+
+def _grouped_gram_configs():
+    configs = []
+    for block_wide, block_m, warps, stages in product(
+        [32, 64, 128, 256], [32, 64, 128], [4, 8], [3, 4]
+    ):
+        configs.append(
+            triton.Config(
+                {"BLOCK_WIDE": block_wide, "BLOCK_M": block_m},
+                num_warps=warps,
+                num_stages=stages,
+            )
+        )
+    return configs
+
+
+@triton.autotune(configs=_grouped_gram_configs(), key=["M_BUCKET", "WIDE", "RANK_IS_I"])
+@triton.jit
+def _grouped_gram_kernel(
+    P_ptr,
+    stride_pm,
+    stride_pd,
+    Q_ptr,
+    stride_qm,
+    stride_qd,
+    OUT_ptr,
+    stride_og,
+    stride_oi,
+    stride_oj,
+    offsets_ptr,
+    M_BUCKET,
+    WIDE: tl.constexpr,
+    RANK: tl.constexpr,
+    scaling,
+    RANK_IS_I: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    BLOCK_WIDE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    allow_tf32: tl.constexpr,
+):
+    """Grouped Gram product: out[g] = scaling * (P_g^T @ Q_g), summed over the
+    tokens of combined-group g (a contiguous [start, end) slice).
+
+    Both LoRA weight gradients are this op once XA/YB are precomputed:
+      dA[g] = scaling * YB_g^T @ X_g   (P=YB[M,R], Q=X[M,K]; rank is the I/row dim)
+      dB[g] = scaling * DY_g^T @ XA_g  (P=DY[M,N], Q=XA[M,R]; rank is the J/col dim)
+
+    The wide dim (K for dA, N for dB) is tiled; the rank dim fits one BLOCK_R.
+    Every (group, wide-block) writes its full tile, so empty groups self-zero --
+    no output pre-init needed. Output strides place the [I, J] tile into either
+    lora_A-grad ([E*T*R, K]) or lora_B-grad ([N, E*T*R]) layout directly.
+    """
+    g = tl.program_id(0)
+    w_blk = tl.program_id(1)
+
+    start = tl.where(g == 0, 0, tl.load(offsets_ptr + g - 1, mask=g > 0, other=0))
+    end = tl.load(offsets_ptr + g)
+    num = end - start
+
+    w_idx = w_blk * BLOCK_WIDE + tl.arange(0, BLOCK_WIDE)
+    w_mask = w_idx < WIDE
+    r_idx = tl.arange(0, BLOCK_R)
+    r_mask = r_idx < RANK
+    input_dtype = P_ptr.dtype.element_ty
+
+    if RANK_IS_I:
+        acc = tl.zeros((BLOCK_R, BLOCK_WIDE), dtype=tl.float32)
+    else:
+        acc = tl.zeros((BLOCK_WIDE, BLOCK_R), dtype=tl.float32)
+
+    for i in range(tl.cdiv(num, BLOCK_M)):
+        m_idx = start + i * BLOCK_M + tl.arange(0, BLOCK_M)
+        m_mask = m_idx < end
+        if RANK_IS_I:
+            p = tl.load(
+                P_ptr + m_idx[:, None] * stride_pm + r_idx[None, :] * stride_pd,
+                mask=m_mask[:, None] & r_mask[None, :],
+                other=0.0,
+            ).to(input_dtype)
+            q = tl.load(
+                Q_ptr + m_idx[:, None] * stride_qm + w_idx[None, :] * stride_qd,
+                mask=m_mask[:, None] & w_mask[None, :],
+                other=0.0,
+            ).to(input_dtype)
+            acc += tl.dot(tl.trans(p), q, allow_tf32=allow_tf32)
+        else:
+            p = tl.load(
+                P_ptr + m_idx[:, None] * stride_pm + w_idx[None, :] * stride_pd,
+                mask=m_mask[:, None] & w_mask[None, :],
+                other=0.0,
+            ).to(input_dtype)
+            q = tl.load(
+                Q_ptr + m_idx[:, None] * stride_qm + r_idx[None, :] * stride_qd,
+                mask=m_mask[:, None] & r_mask[None, :],
+                other=0.0,
+            ).to(input_dtype)
+            acc += tl.dot(tl.trans(p), q, allow_tf32=allow_tf32)
+
+    acc = acc * scaling
+    if RANK_IS_I:
+        out_ptrs = (
+            OUT_ptr
+            + g * stride_og
+            + r_idx[:, None] * stride_oi
+            + w_idx[None, :] * stride_oj
+        )
+        out_mask = r_mask[:, None] & w_mask[None, :]
+    else:
+        out_ptrs = (
+            OUT_ptr
+            + g * stride_og
+            + w_idx[:, None] * stride_oi
+            + r_idx[None, :] * stride_oj
+        )
+        out_mask = w_mask[:, None] & r_mask[None, :]
+    tl.store(out_ptrs, acc.to(OUT_ptr.dtype.element_ty), mask=out_mask)
+
+
+def grouped_lora_weight_grads(
+    grouped_grad_out: torch.Tensor,  # DY  [M*k, N], grouped
+    grouped_x: torch.Tensor,  # X   [M*k, K], grouped
+    yb: torch.Tensor,  # DY@B [M*k, R], grouped (reused from dX path)
+    xa: torch.Tensor,  # X@A^T [M*k, R], grouped (saved from forward)
+    lora_A: torch.Tensor,  # [E*T*R, K]
+    lora_B: torch.Tensor,  # [N, E*T*R]
+    combined_offsets: torch.Tensor,
+    e_total: int,
+    scaling: float,
+):
+    """dA/dB via grouped Gram GEMMs over precomputed XA/YB.
+
+    Avoids the shared split kernel's per-output-block recompute of XA/YB (an
+    inner reduction over K or N, repeated cdiv(wide_dim, BLOCK) times per group);
+    that recompute is what makes the split path scale poorly as tenants grow.
+    """
+    rank = lora_A.size(0) // e_total
+    k_dim = lora_A.size(1)
+    n_dim = lora_B.size(0)
+    block_r = _block_r_for_rank(rank)
+    m_bucket = _bucket_m(max(1, grouped_x.size(0) // e_total))
+
+    dA = torch.empty_like(lora_A)
+    dB = torch.empty_like(lora_B)
+
+    grid_a = lambda meta: (e_total, triton.cdiv(k_dim, meta["BLOCK_WIDE"]))  # noqa: E731
+    _grouped_gram_kernel[grid_a](
+        yb,
+        yb.stride(0),
+        yb.stride(1),
+        grouped_x,
+        grouped_x.stride(0),
+        grouped_x.stride(1),
+        dA,
+        rank * dA.stride(0),
+        dA.stride(0),
+        dA.stride(1),
+        combined_offsets,
+        m_bucket,
+        WIDE=k_dim,
+        RANK=rank,
+        scaling=scaling,
+        RANK_IS_I=True,
+        BLOCK_R=block_r,
+        allow_tf32=lora_ops.ALLOW_TF32,
+    )
+
+    grid_b = lambda meta: (e_total, triton.cdiv(n_dim, meta["BLOCK_WIDE"]))  # noqa: E731
+    _grouped_gram_kernel[grid_b](
+        grouped_grad_out,
+        grouped_grad_out.stride(0),
+        grouped_grad_out.stride(1),
+        xa,
+        xa.stride(0),
+        xa.stride(1),
+        dB,
+        rank * dB.stride(1),
+        dB.stride(0),
+        dB.stride(1),
+        combined_offsets,
+        m_bucket,
+        WIDE=n_dim,
+        RANK=rank,
+        scaling=scaling,
+        RANK_IS_I=False,
+        BLOCK_R=block_r,
+        allow_tf32=lora_ops.ALLOW_TF32,
+    )
+    return dA, dB
+
+
+def build_multilora_routing(
+    flat_expert_idxs: torch.Tensor,  # [M*k] expert per (token, top-k slot)
+    flat_tenant_idxs: torch.Tensor,  # [M*k] tenant per (token, top-k slot)
+    num_experts: int,
+    num_tenants: int,
+):
+    """One sort by ``expert*T + tenant`` yields both groupings.
+
+    Returns ``(sorted_expert_idxs, sorted_combined_idxs, sorted_scattered_idxs,
+    expert_offsets, combined_offsets)``.
+    """
+    combined = flat_expert_idxs * num_tenants + flat_tenant_idxs
+    sorted_combined, sorted_scattered = torch.sort(combined)
+    combined_offsets = compileable_bincount(
+        sorted_combined, num_experts * num_tenants
+    ).cumsum(-1)
+    sorted_expert = torch.div(sorted_combined, num_tenants, rounding_mode="floor")
+    expert_offsets = compileable_bincount(sorted_expert, num_experts).cumsum(-1)
+    return (
+        sorted_expert.contiguous(),
+        sorted_combined.contiguous(),
+        sorted_scattered.contiguous(),
+        expert_offsets,
+        combined_offsets,
+    )
+
+
+class ScatterMoEMultiLoRA(torch.autograd.Function):
+    """Y = X @ W[expert] + scaling * (X @ A[expert,tenant]^T) @ B[expert,tenant]^T,
+    frozen W; gradients flow to the per-(expert,tenant) LoRA A/B and to X."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        expert_weights: torch.Tensor,  # [E, K, N], frozen
+        k: int,
+        sorted_expert_idxs: torch.Tensor,
+        sorted_combined_idxs: torch.Tensor,
+        sorted_scattered_idxs: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        combined_offsets: torch.Tensor,
+        lora_A: torch.Tensor,  # [E*T*R, K]
+        lora_B: torch.Tensor,  # [N, E*T*R]
+        scaling: float,
+    ):
+        if expert_weights.dtype != x.dtype:
+            expert_weights = expert_weights.to(x.dtype)
+        e_total = combined_offsets.size(0)  # E*T groups
+        rank = lora_A.size(0) // e_total
+        n_dim = expert_weights.size(-1)
+        k_dim = expert_weights.size(1)
+
+        with torch.device(x.device):
+            # 1. base: Y = X @ W  (keyed by expert)
+            output = base_ops.scatter2scatter(
+                X=x,
+                W=expert_weights,
+                k=k,
+                sorted_expert_idxs=sorted_expert_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+            )
+            # 2. XA = X @ A^T  (keyed by combined; grouped output [M*k, R])
+            w_a = lora_A.reshape(e_total, rank, k_dim).permute(0, 2, 1).contiguous()
+            xa = base_ops.scatter2scatter(
+                X=x,
+                W=w_a,
+                k=k,
+                sorted_expert_idxs=sorted_combined_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                y_grouped=True,
+            )
+            # 3. Y_lora = XA @ B^T  (keyed by combined)
+            w_b = lora_B.t().reshape(e_total, rank, n_dim).contiguous()
+            y_lora = base_ops.scatter2scatter(
+                X=xa,
+                W=w_b,
+                k=1,
+                sorted_expert_idxs=sorted_combined_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                x_grouped=True,
+            )
+            output.add_(y_lora, alpha=scaling)
+
+            ctx.save_for_backward(
+                x,
+                lora_A,
+                lora_B,
+                sorted_expert_idxs,
+                sorted_combined_idxs,
+                sorted_scattered_idxs,
+                expert_offsets,
+                combined_offsets,
+                xa,
+            )
+            ctx.expert_weights = expert_weights
+            ctx.k = k
+            ctx.scaling = scaling
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):
+        (
+            x,
+            lora_A,
+            lora_B,
+            sorted_expert_idxs,
+            sorted_combined_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            combined_offsets,
+            xa,
+        ) = ctx.saved_tensors
+        w = ctx.expert_weights
+        k = ctx.k
+        scaling = ctx.scaling
+        e_total = combined_offsets.size(0)
+        rank = lora_A.size(0) // e_total
+        k_dim = lora_A.size(1)
+        n_dim = lora_B.size(0)
+
+        with torch.device(grad_out.device):
+            # one grouping (by scattered idx) serves both -- combined-sorted is
+            # also expert-sorted
+            grouped_grad_out = base_ops.group(
+                grad_out, sorted_scattered_idxs, fan_out=1
+            )
+            grouped_x = base_ops.group(x, sorted_scattered_idxs, fan_out=k)
+
+            # YB = dY @ B (keyed by combined). Computed once and reused for both
+            # the LoRA weight grads (dA) and the LoRA input grad (dX_lora).
+            w_dyb = lora_B.reshape(n_dim, e_total, rank).permute(1, 0, 2).contiguous()
+            dyb = base_ops.scatter2scatter(
+                X=grouped_grad_out,
+                W=w_dyb,
+                k=1,
+                sorted_expert_idxs=sorted_combined_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                x_grouped=True,
+                y_grouped=True,
+            )
+
+            # dA/dB as grouped Gram GEMMs over precomputed XA (saved fwd) and YB.
+            # No per-output-block recompute of XA/YB -- the cost that made the
+            # shared split kernel scale poorly as tenants (E*T groups) grow.
+            d_lora_A, d_lora_B = grouped_lora_weight_grads(
+                grouped_grad_out,
+                grouped_x,
+                dyb,
+                xa,
+                lora_A,
+                lora_B,
+                combined_offsets,
+                e_total,
+                scaling,
+            )
+
+            # dX = base (expert) + LoRA (combined)
+            d_grouped = base_ops.scatter2scatter(
+                X=grouped_grad_out,
+                x_grouped=True,
+                W=w.permute(0, 2, 1),
+                sorted_expert_idxs=sorted_expert_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                k=1,
+            )
+            if scaling != 0.0:
+                # dX_lora = scaling * YB @ A, keyed by combined groups (one more
+                # grouped GEMM; mirrors the forward's XA @ B step).
+                w_dxa = lora_A.reshape(e_total, rank, k_dim).contiguous()
+                d_lora = base_ops.scatter2scatter(
+                    X=dyb,
+                    W=w_dxa,
+                    k=1,
+                    sorted_expert_idxs=sorted_combined_idxs,
+                    sorted_scattered_idxs=sorted_scattered_idxs,
+                    x_grouped=True,
+                    y_grouped=True,
+                )
+                d_grouped[sorted_scattered_idxs] += scaling * d_lora
+
+            if k == 1:
+                d_input = d_grouped
+            else:
+                d_input = d_grouped.view(x.size(0), k, d_grouped.size(-1)).sum(-2)
+
+        return (
+            d_input,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            d_lora_A,
+            d_lora_B,
+            None,
+        )
+
+
+def scatter2scatter_multilora(
+    x,
+    expert_weights,
+    k,
+    sorted_expert_idxs,
+    sorted_combined_idxs,
+    sorted_scattered_idxs,
+    expert_offsets,
+    combined_offsets,
+    lora_A,
+    lora_B,
+    scaling: float = 1.0,
+):
+    """Drop-in multi-adapter forward (autograd-backed)."""
+    return ScatterMoEMultiLoRA.apply(
+        x,
+        expert_weights,
+        k,
+        sorted_expert_idxs,
+        sorted_combined_idxs,
+        sorted_scattered_idxs,
+        expert_offsets,
+        combined_offsets,
+        lora_A,
+        lora_B,
+        scaling,
+    )
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/mx_weights.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/mx_weights.py
new file mode 100644
index 0000000000..f2cfd0b3bb
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/mx_weights.py
@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+MXFP4 expert weight container + helpers for the fused-dequant Triton kernels.
+
+The container carries the packed uint8 ``[E_active, N, K/2]`` data and the
+E8M0 ``[E_active, N, K/32]`` scales for the *active* experts of one MoE
+step. ``parallel_linear_lora`` checks for this container instance and
+routes to the MX-aware Triton kernels.
+
+Layout: OCP block axis is the contraction axis ``K`` — the last storage
+dim. The same buffer is consumed by both the forward kernel (K is the
+matmul reduction axis) and the dX kernel (K is the output axis, with
+scales broadcast within ``MX_BLOCK_SIZE``-element K blocks). No
+pre-transpose / re-quantize is needed for the backward path.
+
+The FP4 E2M1 codebook is the standard OCP-MX one (16 values:
+``±{0, 0.5, 1, 1.5, 2, 3, 4, 6}``); we cache one fp32 copy per CUDA device.
+"""
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+MX_BLOCK_SIZE = 32
+
+# Standard OCP-MX fp4 e2m1 codebook (sign | 2-bit exp | 1-bit mantissa), indexed by raw nibble.
+_FP4_E2M1_LUT = (
+    0.0,
+    0.5,
+    1.0,
+    1.5,
+    2.0,
+    3.0,
+    4.0,
+    6.0,
+    -0.0,
+    -0.5,
+    -1.0,
+    -1.5,
+    -2.0,
+    -3.0,
+    -4.0,
+    -6.0,
+)
+
+_LUT_CACHE: dict[torch.device, torch.Tensor] = {}
+
+
+def fp4_codebook(device: torch.device) -> torch.Tensor:
+    """Return the cached 16-entry FP4 E2M1 → fp32 lookup on ``device``."""
+    key = device
+    lut = _LUT_CACHE.get(key)
+    if lut is None or lut.device != device:
+        lut = torch.tensor(_FP4_E2M1_LUT, dtype=torch.float32, device=device)
+        _LUT_CACHE[key] = lut
+    return lut
+
+
+class MXLayout(enum.IntEnum):
+    """Which axis the OCP-MX block scaling runs along, in *kernel* coords.
+
+    Currently only ``FWD`` (block axis = K) is supported and used by both
+    the forward and dX kernels. The enum is kept as a future extension
+    point for swizzled or N-axis-blocked variants.
+    """
+
+    FWD = 0
+
+
+@dataclass
+class MXWeights:
+    """Packed + scale tensors for one MoE projection's active experts.
+
+    Attributes
+    ----------
+    packed:
+        ``uint8`` tensor, shape ``[E_active, N, K/2]``.
+    scales:
+        ``uint8`` (E8M0) tensor, shape ``[E_active, N, K/32]``.
+    K, N:
+        Logical contraction/output dimensions of the dequantized W.
+    layout:
+        Which axis the block scaling runs along (see ``MXLayout``).
+    block_size:
+        OCP block size: ``32`` for MXFP4 (E8M0 scales), ``16`` for NVFP4.
+    scale_is_linear:
+        ``False`` (default): scales are E8M0 (``uint8``), decoded in-kernel as
+        ``2^(byte-127)``. ``True``: scales are already a linear floating-point
+        multiplier (e.g. NVFP4's E4M3 block scale pre-multiplied by the per-tensor
+        scale), loaded and used directly with no exponent decode.
+    """
+
+    packed: torch.Tensor
+    scales: torch.Tensor
+    K: int
+    N: int
+    layout: MXLayout = MXLayout.FWD
+    block_size: int = MX_BLOCK_SIZE
+    num_experts: Optional[int] = None  # E_active; convenience field
+    orig_dtype: torch.dtype = torch.bfloat16
+    scale_is_linear: bool = False
+
+    def __post_init__(self) -> None:
+        assert self.block_size in (16, 32), (
+            f"block_size must be 16 (NVFP4) or 32 (MXFP4), got {self.block_size}"
+        )
+        if not self.scale_is_linear and self.scales.dtype != torch.uint8:
+            # E8M0 (float8_e8m0fnu) scales viewed as uint8 so the Triton kernel can
+            # load them with simple integer arithmetic. Linear scales stay floating.
+            self.scales = self.scales.view(torch.uint8)
+        assert self.packed.dtype == torch.uint8, (
+            f"packed must be uint8, got {self.packed.dtype}"
+        )
+        if self.num_experts is None:
+            self.num_experts = self.packed.size(0)
+
+    @property
+    def device(self) -> torch.device:
+        return self.packed.device
+
+
+def _torchao_mxtensor_cls():
+    """Return the torchao MXTensor class, or ``None`` if torchao is missing."""
+    try:
+        from torchao.prototype.mx_formats.mx_tensor import MXTensor
+    except ImportError:
+        return None
+    return MXTensor
+
+
+def _mx_qdata(mx) -> torch.Tensor:
+    """Read the packed-nibble buffer off an MXTensor, tolerating torchao
+    renaming the attribute between versions."""
+    qdata = getattr(mx, "qdata", None)
+    if qdata is None:
+        qdata = getattr(mx, "_data", None)
+    if qdata is None:
+        raise AttributeError(
+            "torchao MXTensor exposes neither .qdata nor ._data; "
+            "this torchao version is unsupported."
+        )
+    return qdata
+
+
+def _mx_scale(mx) -> torch.Tensor:
+    """Read the E8M0 scale buffer off an MXTensor, tolerating torchao
+    renaming the attribute between versions."""
+    scale = getattr(mx, "scale", None)
+    if scale is None:
+        scale = getattr(mx, "_scale_e8m0", None)
+    if scale is None:
+        raise AttributeError(
+            "torchao MXTensor exposes neither .scale nor ._scale_e8m0; "
+            "this torchao version is unsupported."
+        )
+    return scale
+
+
+def _construct_mxtensor_subset(
+    parent, qdata_slice: torch.Tensor, scale_slice: torch.Tensor
+):
+    """Construct a new MXTensor that shares ``parent``'s metadata but uses
+    the provided ``qdata_slice`` / ``scale_slice`` buffers.
+
+    Pinned to torchao 0.17.0's positional constructor (qdata, scale,
+    elem_dtype, block_size, orig_dtype, kernel_preference,
+    act_quant_kwargs, is_swizzled_scales). Optional attributes are read via
+    ``getattr`` so we degrade gracefully if a future torchao version drops
+    or renames one — the single point of pain for torchao internals access
+    across this codebase.
+    """
+    MXTensor = _torchao_mxtensor_cls()
+    if MXTensor is None:
+        raise ImportError("MXFP4 path requires torchao (install `torchao>=0.7`).")
+    kernel_preference = getattr(parent, "kernel_preference", None)
+    act_quant_kwargs = getattr(parent, "act_quant_kwargs", None)
+    is_swizzled_scales = getattr(parent, "is_swizzled_scales", False)
+    return MXTensor(
+        qdata_slice,
+        scale_slice,
+        parent.elem_dtype,
+        parent.block_size,
+        parent.orig_dtype,
+        kernel_preference,
+        act_quant_kwargs,
+        is_swizzled_scales,
+    )
+
+
+def selective_mx_weights_fwd(mx_param, active_experts: torch.Tensor) -> MXWeights:
+    """Slice an MXFP4 expert parameter to the active set, keeping the K-axis
+    block layout (FWD). The returned ``MXWeights.packed`` has shape
+    ``[num_active, N, K/2]`` and is directly consumable by the forward MX
+    kernel via ``parallel_linear_lora``."""
+    MXTensor = _torchao_mxtensor_cls()
+    if MXTensor is None:
+        raise ImportError("MXFP4 fused path requires torchao>=0.7 (install `torchao`).")
+    assert isinstance(mx_param, MXTensor), (
+        f"selective_mx_weights_fwd expects an MXTensor, got {type(mx_param)}"
+    )
+    assert mx_param.elem_dtype == torch.float4_e2m1fn_x2, (
+        "only MXFP4 (float4_e2m1fn_x2) is supported"
+    )
+    # Dense routing: reference the resident param zero-copy (see selective_nvfp4_weights_fwd);
+    # fall back to the gather when sparse.
+    qd, sc = _mx_qdata(mx_param), _mx_scale(mx_param)
+    all_active = active_experts.numel() == qd.size(0)
+    sub_qdata = qd if all_active else qd[active_experts].contiguous()
+    sub_scale = sc if all_active else sc[active_experts].contiguous()
+    # Kernel's K = contraction axis = OCP block axis = LAST storage axis; N is the leading
+    # non-expert axis.
+    N = sub_qdata.size(1)
+    K = sub_qdata.size(2) * 2
+    return MXWeights(
+        packed=sub_qdata,
+        scales=sub_scale,
+        K=K,
+        N=N,
+        layout=MXLayout.FWD,
+        block_size=mx_param.block_size,
+        num_experts=sub_qdata.size(0),
+        orig_dtype=mx_param.orig_dtype,
+    )
+
+
+def _torchao_nvfp4tensor_cls():
+    """Return the torchao NVFP4Tensor class, or ``None`` if torchao is missing."""
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+    except ImportError:
+        return None
+    return NVFP4Tensor
+
+
+def selective_nvfp4_weights_fwd(nv_param, active_experts: torch.Tensor) -> MXWeights:
+    """Slice an NVFP4 expert parameter to the active set as an ``MXWeights`` with
+    ``scale_is_linear=True``.
+
+    NVFP4 shares the FP4 E2M1 element codebook with MXFP4 (same packed nibble layout)
+    but scales per 16-element block with an E4M3 (fp8) value, optionally times a
+    per-tensor fp32 scale. The dequant is ``codebook(nibble) * e4m3_block_scale *
+    per_tensor`` -- not a power-of-2 -- so the E4M3 block scale (folded with the
+    per-tensor scale) is decoded to a linear fp32 multiplier here and consumed
+    directly by the kernel (no in-kernel exponent decode). The packed buffer stays
+    4-bit; only the small ``[num_active, N, K/16]`` scale tensor is materialized.
+    """
+    NVFP4Tensor = _torchao_nvfp4tensor_cls()
+    if NVFP4Tensor is None:
+        raise ImportError("NVFP4 fused path requires torchao (install `torchao`).")
+    assert isinstance(nv_param, NVFP4Tensor), (
+        f"selective_nvfp4_weights_fwd expects an NVFP4Tensor, got {type(nv_param)}"
+    )
+    assert nv_param.block_size == 16, (
+        f"NVFP4 block_size must be 16, got {nv_param.block_size}"
+    )
+    # Dense routing (all experts active, the common case): remap is identity, so reference the
+    # resident param zero-copy instead of gathering the whole packed weight every layer. Falls
+    # back to the selective gather when routing is sparse.
+    all_active = active_experts.numel() == nv_param.qdata.size(0)
+    sub_qdata = (
+        nv_param.qdata if all_active else nv_param.qdata[active_experts].contiguous()
+    )
+    raw_scale = nv_param.scale if all_active else nv_param.scale[active_experts]
+    block_scale = raw_scale.to(torch.float32)  # [a, N, K/16]
+    per_tensor = getattr(nv_param, "per_tensor_scale", None)
+    if per_tensor is not None:
+        per_tensor = per_tensor.to(torch.float32)
+        # Per-expert scale reshaped to [a,1,1] to broadcast along block_scale's expert dim
+        # [a,N,K/16] (bare [a] would hit the last dim). A shared scalar just broadcasts.
+        if per_tensor.dim() >= 1 and per_tensor.size(0) == nv_param.qdata.size(0):
+            if not all_active:
+                per_tensor = per_tensor[active_experts]
+            per_tensor = per_tensor.reshape(-1, 1, 1)
+        block_scale = block_scale * per_tensor
+    N = sub_qdata.size(1)
+    K = sub_qdata.size(2) * 2
+    return MXWeights(
+        packed=sub_qdata,
+        scales=block_scale.contiguous(),
+        K=K,
+        N=N,
+        layout=MXLayout.FWD,
+        block_size=16,
+        num_experts=sub_qdata.size(0),
+        orig_dtype=nv_param.orig_dtype,
+        scale_is_linear=True,
+    )
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_fp8_quantizer.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_fp8_quantizer.py
new file mode 100644
index 0000000000..25752d8ec4
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_fp8_quantizer.py
@@ -0,0 +1,495 @@
+"""Clean NVFP4-MoE loading for DeepSeek-V4-Flash-NVFP4 via a FineGrainedFP8 subclass.
+
+The checkpoint declares ``quant_method: fp8`` (→ ``FineGrainedFP8HfQuantizer``) with
+``quant_algo: MIXED_PRECISION`` — the non-expert weights are blockwise FP8 ``[128,128]``,
+but the routed experts are NVFP4 (``moe_quant_algo: NVFP4``, ``group_size: 16``,
+``fmt: e4m3``). ``FP8Experts`` has an FP4 path (``expert_dtype == "fp4"``) but allocates a
+blockwise-FP8 ``*_scale_inv`` buffer at the *wrong* granularity (``sf_gran_k=32`` + the
+model-global ``scale_fmt: ue8m0``); it has no notion of NVFP4's group-16 ``*_scale``
+(E4M3) + per-tensor ``*_scale_2``.
+
+The deepseek_v4 weight converter fuses ``experts.*.w1/w3.weight`` → ``gate_up_proj`` with
+*unanchored* source patterns (``mlp.experts.*.w1.weight``), so by regex search they also
+match the sibling ``.weight_scale`` / ``.weight_scale_2`` keys. The group-16 ``weight_scale``
+(2-D) fuses fine, but ``weight_scale_2`` is a 0-D scalar — ``Concatenate`` on stacked
+scalars raises (a spurious conversion error) — and in any case there's no ``*_scale`` /
+``*_scale_2`` param to receive the fused scales, so they come in UNEXPECTED and the
+``*_scale_inv`` placeholders end up MISSING (random init → invalid).
+
+This subclass loads everything in place, with no checkpoint re-read of weights/scales:
+  * ``update_weight_conversions`` — anchors the experts' ``*.weight`` converters with ``$``
+    (so the scalar-scale concat never runs) and adds twin ``*.weight_scale`` converters so
+    the 2-D group-16 scales fuse into ``gate_up_proj_scale`` / ``down_proj_scale``;
+  * ``_process_model_before_weight_loading`` — swaps the wrong ``*_scale_inv`` placeholders
+    for correctly-shaped NVFP4 ``*_scale`` params, and marks the 0-D ``*_scale_2`` /
+    dynamic ``*.input_scale`` keys ignorable;
+  * ``_process_model_after_weight_loading`` — folds the in-place qdata + ``*_scale`` plus
+    the per-expert scalar ``*_scale_2`` (the only thing re-read — 4 bytes/expert) into a
+    torchao ``NVFP4Tensor`` for the scattermoe fused-LoRA path.
+
+Net: no UNEXPECTED/MISSING warning, no random init, correct NVFP4 experts, no multi-GB
+re-read. Installed by swapping it into ``AUTO_QUANTIZER_MAPPING["fp8"]`` before
+``from_pretrained``.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn as nn
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_GROUP_SIZE = 16
+
+# How to store the blockwise-FP8 non-expert linears after load (set from the
+# `dsv4_fp8_nonexpert_mode` config via `configure_nonexpert_mode`):
+#   "float8tensor" (default): wrap each plain FP8Linear weight as a torchao Float8Tensor
+#       (1-byte qdata + block scale). Forward/backward + PEFT LoRA run via subclass autograd;
+#       the fused LoRA kernels dequant it through axolotl.kernels.quantize.dequantize.
+#   "bf16": dequantize to bf16 in place (safe path; 2 bytes/param). Also sidesteps transformers'
+#       @triton_op FP8 matmul (no autograd formula; slow discovery).
+# Grouped linears (o_a_proj) always go to bf16: their view+bmm forward isn't subclass-safe.
+_FP8_NONEXPERT_MODE = "float8tensor"
+
+
+def configure_nonexpert_mode(mode: str | None) -> None:
+    """Set the FP8 non-expert storage mode from cfg before the model loads (the quantizer
+    reads the module global in ``_process_model_after_weight_loading``)."""
+    global _FP8_NONEXPERT_MODE
+    _FP8_NONEXPERT_MODE = (mode or "float8tensor").lower()
+
+
+# The only per-expert keys left unmatched once weight + weight_scale fuse in place: the 0-D
+# per-tensor `weight_scale_2` (folded after load) and the unused dynamic `input_scale`.
+_IGNORE_UNEXPECTED = (
+    r"experts\..*\.weight_scale_2$",
+    r"experts\..*\.input_scale$",
+)
+
+
+def _nvfp4_cls():
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+        return NVFP4Tensor
+    except ImportError:
+        return None
+
+
+def _float8_cls():
+    try:
+        from torchao.quantization import Float8Tensor
+
+        return Float8Tensor
+    except ImportError:
+        return None
+
+
+def _drop_param(mod: nn.Module, name: str) -> None:
+    if name in mod._parameters:
+        del mod._parameters[name]
+    else:
+        setattr(mod, name, None)
+
+
+def _is_nvfp4_experts(module: nn.Module) -> bool:
+    """An ``FP8Experts`` carrying FP4-packed (int8/uint8, K-halved) expert weights."""
+    try:
+        from transformers.integrations.finegrained_fp8 import FP8Experts
+    except ImportError:
+        return False
+    if not isinstance(module, FP8Experts):
+        return False
+    w = getattr(module, "gate_up_proj", getattr(module, "up_proj", None))
+    return (
+        isinstance(w, torch.Tensor)
+        and w.dtype in (torch.int8, torch.uint8)
+        and w.ndim == 3
+    )
+
+
+def _is_expert_weight_converter(conv) -> bool:
+    pats = getattr(conv, "_original_source_patterns", None) or []
+    return any("experts" in p and p.endswith(".weight") for p in pats)
+
+
+def _resolve_repo_file(repo: str, filename: str) -> str:
+    """Resolve a checkpoint file path from a local snapshot dir or the HF hub.
+
+    A local snapshot dir (offline/air-gapped axolotl usage) would fail ``hf_hub_download``
+    with an ``HFValidationError``, so read straight from disk in that case.
+    """
+    if os.path.isdir(repo):
+        return os.path.join(repo, filename)
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo, filename)
+
+
+def _scale_2_path(repo: str):
+    """Return ``(weight_map, opener)`` for re-reading per-expert scalar ``weight_scale_2``."""
+    import json
+
+    from safetensors import safe_open
+
+    with open(_resolve_repo_file(repo, "model.safetensors.index.json")) as _f:
+        wmap = json.load(_f)["weight_map"]
+    cache: dict[str, object] = {}
+
+    def opener(shard):
+        f = cache.get(shard)
+        if f is None:
+            f = cache[shard] = safe_open(
+                _resolve_repo_file(repo, shard), framework="pt"
+            )
+        return f
+
+    return wmap, opener
+
+
+def _dequantize_fp8_linears(model: nn.Module, quantizer) -> int:
+    """Dequantize blockwise-FP8 linear weights to bf16 IN PLACE (preserving each module's
+    class + custom forward, e.g. the grouped ``o_a_proj``). A module is targeted by the
+    presence of an fp8/int8 ``weight`` + a ``weight_scale_inv`` companion — not by class —
+    so ``FP8Linear`` and any quantized ``nn.Linear`` subclass are both handled. After this,
+    ``FP8Linear.forward`` auto-uses plain ``F.linear`` (weight.element_size() > 1), and the
+    FP8 matmul custom op (no autograd formula; slow triton kernel discovery) is gone."""
+    from transformers.integrations.finegrained_fp8 import Fp8Dequantize
+
+    deq = Fp8Dequantize(quantizer)
+    n = 0
+    by_type: dict[str, int] = {}
+    for name, mod in model.named_modules():
+        wsi = getattr(mod, "weight_scale_inv", None)
+        w = getattr(mod, "weight", None)
+        if (
+            wsi is None
+            or not isinstance(w, torch.Tensor)
+            or w.dtype not in (torch.float8_e4m3fn, torch.int8)
+        ):
+            continue
+        w_bf16 = deq._dequantize_one(w.data, wsi.data, torch.bfloat16)
+        if tuple(w_bf16.shape) != tuple(w.shape):
+            # FP4-packed (int8, half-K) unpacks to 2x width (expected); otherwise log.
+            LOG.warning(
+                "dequant %s: weight %s -> %s (dtype %s)",
+                name,
+                tuple(w.shape),
+                tuple(w_bf16.shape),
+                w.dtype,
+            )
+        mod.weight = nn.Parameter(w_bf16, requires_grad=False)
+        if "weight_scale_inv" in mod._parameters:
+            del mod._parameters["weight_scale_inv"]
+        else:
+            mod.weight_scale_inv = None
+        by_type[type(mod).__name__] = by_type.get(type(mod).__name__, 0) + 1
+        n += 1
+    if n:
+        LOG.info("Dequantized %d FP8 linear weights to bf16 in place: %s", n, by_type)
+    return n
+
+
+def _wrap_fp8_linears_as_float8tensor(model: nn.Module, quantizer) -> int:
+    """Wrap blockwise-FP8 *plain* linear weights as torchao ``Float8Tensor`` (1-byte qdata +
+    block scale) so the frozen base keeps the FP8 memory footprint while forward/backward and
+    PEFT LoRA run through the tensor subclass. The checkpoint's exact qdata + ``weight_scale_inv``
+    are reused (no re-quant). Grouped linears (``o_a_proj``), whose ``view+bmm`` forward isn't
+    subclass-safe, are dequantized to bf16 in place instead."""
+    Float8Tensor = _float8_cls()
+    if Float8Tensor is None:
+        LOG.warning(
+            "torchao Float8Tensor unavailable; dequantizing FP8 non-experts to bf16"
+        )
+        return _dequantize_fp8_linears(model, quantizer)
+    from transformers.integrations.finegrained_fp8 import (
+        Fp8Dequantize,
+        FP8GroupedLinear,
+    )
+
+    deq = Fp8Dequantize(quantizer)
+    wrapped = bf16ed = 0
+    for _name, mod in model.named_modules():
+        wsi = getattr(mod, "weight_scale_inv", None)
+        w = getattr(mod, "weight", None)
+        if (
+            wsi is None
+            or not isinstance(w, torch.Tensor)
+            or w.dtype != torch.float8_e4m3fn
+        ):
+            continue
+        if isinstance(mod, FP8GroupedLinear) or (getattr(mod, "n_groups", 1) or 1) > 1:
+            mod.weight = nn.Parameter(
+                deq._dequantize_one(w.data, wsi.data, torch.bfloat16),
+                requires_grad=False,
+            )
+            _drop_param(mod, "weight_scale_inv")
+            bf16ed += 1
+            continue
+        block = list(mod.block_size) if getattr(mod, "block_size", None) else [128, 128]
+        scale = wsi.data
+        if scale.dtype not in (torch.float32, torch.float16, torch.bfloat16):
+            scale = scale.to(
+                torch.float32
+            )  # ue8m0 scales have no float ops; torchao wants float
+        f8 = Float8Tensor(w.data, scale, block_size=block, dtype=torch.bfloat16)
+        mod.weight = nn.Parameter(f8, requires_grad=False)
+        _drop_param(mod, "weight_scale_inv")
+        wrapped += 1
+    if wrapped or bf16ed:
+        LOG.info(
+            "FP8 non-experts: wrapped %d plain linears as Float8Tensor (1-byte), %d grouped -> bf16",
+            wrapped,
+            bf16ed,
+        )
+    return wrapped + bf16ed
+
+
+def _enable_torchao_lora_dispatch(quantizer) -> None:
+    """Let PEFT attach LoRA to the Float8Tensor non-expert bases.
+
+    PEFT's ``dispatch_torchao`` fires because the base weights are torchao ``Float8Tensor``
+    and routes to ``TorchaoLoraLinear``, which needs
+    ``model.hf_quantizer.quantization_config.get_apply_tensor_subclass``. PEFT sources it via
+    ``operator.attrgetter`` and silently skips on ``AttributeError`` (we ship a FineGrained-FP8
+    config, not a ``TorchAoConfig``), but the dispatcher then errors on the missing kwarg.
+    It's used ONLY by merge/unmerge (re-quantize the merged weight) — never in training (frozen
+    base, LoRA kept separate) — so provide a callable that unblocks training and fails loudly
+    on merge (blockwise-FP8 merge isn't supported yet)."""
+    cfg = getattr(quantizer, "quantization_config", None)
+    if cfg is None:
+        return
+    cls = type(cfg)
+    if hasattr(cls, "get_apply_tensor_subclass"):
+        return
+
+    def _no_merge():
+        raise NotImplementedError(
+            "merge-lora on a blockwise-FP8 (Float8Tensor) base is not supported; train/serve "
+            "with the adapter kept separate, or set `dsv4_fp8_nonexpert_mode: bf16` to merge."
+        )
+
+    # Attach to the CLASS as a staticmethod (not the instance): PEFT's attrgetter still resolves
+    # it via attribute lookup, but it stays out of the instance __dict__ so `config.to_json`
+    # (model.config save) doesn't try to serialize a function object.
+    try:
+        cls.get_apply_tensor_subclass = staticmethod(_no_merge)
+    except Exception:  # pragma: no cover - some configs are frozen/slotted
+        LOG.warning(
+            "Could not attach get_apply_tensor_subclass; non-expert LoRA may fail to inject"
+        )
+
+
+def make_nvfp4_fp8_quantizer():
+    """Build the ``FineGrainedFP8HfQuantizer`` subclass (deferred import)."""
+    from transformers.core_model_loading import WeightConverter
+    from transformers.quantizers.quantizer_finegrained_fp8 import (
+        FineGrainedFP8HfQuantizer,
+    )
+
+    class NVFP4MoEFP8HfQuantizer(FineGrainedFP8HfQuantizer):
+        """FineGrained-FP8 quantizer that also loads NVFP4 routed experts correctly."""
+
+        def update_weight_conversions(self, weight_conversions):
+            rebuilt = []
+            for conv in weight_conversions:
+                if not (
+                    isinstance(conv, WeightConverter)
+                    and _is_expert_weight_converter(conv)
+                ):
+                    rebuilt.append(conv)
+                    continue
+                ops = [type(op)(op.dim) for op in conv.operations]
+                # weight: anchor `.weight$` so the converter fuses ONLY the packed weight
+                # and stops swallowing `.weight_scale*` (the 0-D `_scale_2` concat raises).
+                w_conv = WeightConverter(
+                    source_patterns=[
+                        p + "$" if p.endswith(".weight") else p
+                        for p in conv._original_source_patterns
+                    ],
+                    target_patterns=list(conv._original_target_patterns),
+                    operations=ops,
+                )
+                # weight_scale: twin converter so the 2-D group-16 scales fuse the same way
+                # into `*_scale`, loading in place (no re-read).
+                s_conv = WeightConverter(
+                    source_patterns=[
+                        p[: -len(".weight")] + ".weight_scale$"
+                        if p.endswith(".weight")
+                        else p
+                        for p in conv._original_source_patterns
+                    ],
+                    target_patterns=[
+                        t + "_scale" for t in conv._original_target_patterns
+                    ],
+                    operations=[type(op)(op.dim) for op in conv.operations],
+                )
+                for new in (w_conv, s_conv):
+                    new.scope_prefix = conv.scope_prefix
+                    new.base_model_prefix = conv.base_model_prefix
+                rebuilt.extend((w_conv, s_conv))
+            return super().update_weight_conversions(rebuilt)
+
+        def _process_model_before_weight_loading(self, model, **kwargs):
+            super()._process_model_before_weight_loading(model, **kwargs)
+            n = 0
+            for mod in model.modules():
+                if not _is_nvfp4_experts(mod):
+                    continue
+                E, H, I = mod.num_experts, mod.hidden_dim, mod.intermediate_dim
+                dev = (mod.gate_up_proj if mod.has_gate else mod.up_proj).device
+                # drop the wrong-granularity blockwise-FP8 placeholders; register NVFP4
+                # group-16 E4M3 `*_scale` params so the fused scales land in place.
+                pfx, out_dim = (
+                    ("gate_up_proj", 2 * I) if mod.has_gate else ("up_proj", I)
+                )
+                for stale in (pfx + "_scale_inv", "down_proj_scale_inv"):
+                    mod._parameters.pop(stale, None)
+                mod.register_parameter(
+                    pfx + "_scale",
+                    nn.Parameter(
+                        torch.empty(
+                            E,
+                            out_dim,
+                            H // _GROUP_SIZE,
+                            dtype=torch.float8_e4m3fn,
+                            device=dev,
+                        ),
+                        requires_grad=False,
+                    ),
+                )
+                mod.register_parameter(
+                    "down_proj_scale",
+                    nn.Parameter(
+                        torch.empty(
+                            E,
+                            H,
+                            I // _GROUP_SIZE,
+                            dtype=torch.float8_e4m3fn,
+                            device=dev,
+                        ),
+                        requires_grad=False,
+                    ),
+                )
+                n += 1
+            if n:
+                ignore = set(model._keys_to_ignore_on_load_unexpected or ())
+                ignore.update(_IGNORE_UNEXPECTED)
+                model._keys_to_ignore_on_load_unexpected = ignore
+                LOG.info(
+                    "Prepared %d NVFP4 experts modules for clean in-place loading", n
+                )
+
+        def _process_model_after_weight_loading(self, model, **kwargs):
+            super()._process_model_after_weight_loading(model, **kwargs)
+            # The swap into AUTO_QUANTIZER_MAPPING['fp8'] is process-global, so this runs for EVERY
+            # fp8 checkpoint loaded after install. Gate all NVFP4-specific work on actually finding
+            # NVFP4 experts; a plain fp8 model behaves exactly like the original quantizer (super()
+            # above) and is NOT re-wrapped as Float8Tensor.
+            mods = [
+                (name, m) for name, m in model.named_modules() if _is_nvfp4_experts(m)
+            ]
+            if not mods:
+                return
+            if _FP8_NONEXPERT_MODE == "bf16":
+                n = _dequantize_fp8_linears(model, self)
+                if n:
+                    LOG.info("Dequantized %d non-expert FP8Linear modules to bf16", n)
+            else:
+                n = _wrap_fp8_linears_as_float8tensor(model, self)
+                if n:
+                    _enable_torchao_lora_dispatch(self)
+            NVFP4Tensor = _nvfp4_cls()
+            if NVFP4Tensor is None:
+                LOG.warning("torchao NVFP4Tensor unavailable; skipping expert wrap")
+                return
+            repo = getattr(model, "name_or_path", None) or getattr(
+                model.config, "_name_or_path", None
+            )
+            wmap, opener = _scale_2_path(repo)
+            import re
+
+            for name, mod in mods:
+                layer = int(re.search(r"layers\.(\d+)\.", name).group(1))
+                pfx = "gate_up_proj" if mod.has_gate else "up_proj"
+                projs = {
+                    "gate_up_proj": ("w1",),
+                    "up_proj": ("w1",),
+                    "down_proj": ("w2",),
+                }
+                for proj in (pfx, "down_proj"):
+                    qdata = getattr(mod, proj).data.view(torch.uint8)
+                    scale = getattr(mod, proj + "_scale").data
+                    # w1/w3 share weight_scale_2; one scalar per expert → [E,1,1] broadcast.
+                    src = projs[proj][0]
+                    s2 = (
+                        torch.stack(
+                            [
+                                opener(
+                                    wmap[
+                                        f"layers.{layer}.ffn.experts.{e}.{src}.weight_scale_2"
+                                    ]
+                                ).get_tensor(
+                                    f"layers.{layer}.ffn.experts.{e}.{src}.weight_scale_2"
+                                )
+                                for e in range(mod.num_experts)
+                            ]
+                        )
+                        .to(device=qdata.device, dtype=torch.float32)
+                        .view(-1, 1, 1)
+                    )
+                    setattr(
+                        mod,
+                        proj,
+                        nn.Parameter(
+                            NVFP4Tensor(
+                                qdata,
+                                scale,
+                                _GROUP_SIZE,
+                                torch.bfloat16,
+                                per_tensor_scale=s2,
+                            ),
+                            requires_grad=False,
+                        ),
+                    )
+                    mod._parameters.pop(proj + "_scale", None)
+            LOG.info(
+                "Wrapped %d NVFP4 experts modules as NVFP4Tensor (in-place scales)",
+                len(mods),
+            )
+
+    return NVFP4MoEFP8HfQuantizer
+
+
+_ORIGINAL_FP8_QUANTIZER = None
+
+
+def install_nvfp4_fp8_quantizer() -> None:
+    """Swap the NVFP4-aware quantizer into ``AUTO_QUANTIZER_MAPPING["fp8"]`` so any
+    ``quant_method: fp8`` checkpoint with NVFP4 experts loads correctly. Idempotent."""
+    global _ORIGINAL_FP8_QUANTIZER
+    from transformers.quantizers import auto as _auto
+
+    cls = make_nvfp4_fp8_quantizer()
+    if getattr(_auto.AUTO_QUANTIZER_MAPPING.get("fp8"), "__name__", "") == cls.__name__:
+        return
+    _ORIGINAL_FP8_QUANTIZER = _auto.AUTO_QUANTIZER_MAPPING.get("fp8")
+    _auto.AUTO_QUANTIZER_MAPPING["fp8"] = cls
+    LOG.info("Installed NVFP4-aware FP8 quantizer into AUTO_QUANTIZER_MAPPING['fp8']")
+
+
+def uninstall_nvfp4_fp8_quantizer() -> None:
+    """Restore the original ``AUTO_QUANTIZER_MAPPING["fp8"]`` entry the swap replaced (so a long-lived
+    process can return to stock fp8 loading). Idempotent; a no-op if install was never called."""
+    global _ORIGINAL_FP8_QUANTIZER
+    if _ORIGINAL_FP8_QUANTIZER is None:
+        return
+    from transformers.quantizers import auto as _auto
+
+    _auto.AUTO_QUANTIZER_MAPPING["fp8"] = _ORIGINAL_FP8_QUANTIZER
+    _ORIGINAL_FP8_QUANTIZER = None
+    LOG.info("Restored original FP8 quantizer in AUTO_QUANTIZER_MAPPING['fp8']")
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_fsdp.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_fsdp.py
new file mode 100644
index 0000000000..f95923a14a
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_fsdp.py
@@ -0,0 +1,214 @@
+"""FSDP2 support for torchao ``NVFP4Tensor`` (frozen, expert-sharded on dim 0).
+
+torchao's NVFP4Tensor (prototype) ships no FSDP2 hooks, so `fully_shard` on a module with
+NVFP4 expert params fails: the flat-param sharding path calls ops the subclass doesn't
+implement (``aten.split`` etc.). This adds:
+
+  * ``aten.split`` / ``aten.clone`` / ``aten.detach`` / ``aten.copy_`` operating on the
+    inner ``(qdata, scale, per_tensor_scale)`` along dim 0 (the expert axis FSDP shards),
+  * ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather`` so FSDP2 uses the subclass-extension
+    path (all-gathering the inner tensors and reconstructing the subclass) instead of the
+    flat-buffer ``view(-1)`` path, which a block-scaled FP4 layout can't satisfy.
+
+Idempotent. Sharding is restricted to dim 0 (the [E, N, K] expert axis) — block-axis
+sharding is not supported (nor needed for expert-parallel/FSDP of MoE weights).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+_PATCHED = False
+
+
+def _nvfp4_cls():
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+        return NVFP4Tensor
+    except ImportError:
+        return None
+
+
+def _rebuild(ref, qdata, scale, per_tensor_scale):
+    NVFP4Tensor = type(ref)
+    return NVFP4Tensor(
+        qdata,
+        scale,
+        ref.block_size,
+        ref.orig_dtype,
+        per_tensor_scale,
+        ref.act_per_tensor_scale,
+        ref.is_swizzled_scales,
+        ref.use_triton_kernel,
+        ref.act_quant_kwargs,
+    )
+
+
+def patch_nvfp4_fsdp():
+    global _PATCHED
+    if _PATCHED:
+        return
+    NVFP4Tensor = _nvfp4_cls()
+    if NVFP4Tensor is None:
+        return
+    aten = torch.ops.aten
+    implements = NVFP4Tensor.implements
+
+    def _pts_along_dim0(x, n):
+        """Per-tensor scale handling under a dim-0 split: per-expert -> split, else share."""
+        pts = x.per_tensor_scale
+        if pts is not None and pts.dim() >= 1 and pts.shape[0] == x.shape[0]:
+            return list(torch.split(pts, n, 0)) if isinstance(n, int) else None
+        return None
+
+    @implements([aten.split.Tensor])
+    def _split(func, types, args, kwargs):
+        x, split_size = args[0], args[1]
+        dim = args[2] if len(args) > 2 else kwargs.get("dim", 0)
+        if dim != 0:
+            raise NotImplementedError(
+                f"NVFP4Tensor FSDP split only on dim 0, got {dim}"
+            )
+        qd = func(x.qdata, split_size, 0)
+        sc = func(x.scale, split_size, 0)
+        pts = _pts_along_dim0(x, split_size)
+        return [
+            _rebuild(x, qd[i], sc[i], pts[i] if pts is not None else x.per_tensor_scale)
+            for i in range(len(qd))
+        ]
+
+    @implements([aten.clone.default, aten.detach.default])
+    def _clone_detach(func, types, args, kwargs):
+        x = args[0]
+        return x._apply_fn_to_data(lambda t: func(t, **kwargs))
+
+    @implements([aten.new_zeros.default])
+    def _new_zeros(func, types, args, kwargs):
+        # FSDP pads on dim 0 (experts). Scale may be SWIZZLED (its trailing dims are not
+        # K//block), so preserve qdata/scale trailing shapes and vary only dim 0.
+        x, size = args[0], list(args[1])
+        E = size[0]
+        qd = func(x.qdata, [E, *x.qdata.shape[1:]], **kwargs)
+        sc = func(x.scale, [E, *x.scale.shape[1:]], **kwargs)
+        # Preserve a per-expert per_tensor_scale buffer (vary dim 0) so the subsequent
+        # copy_/all-gather can carry it; dropping it here loses it for the whole param.
+        pts = x.per_tensor_scale
+        if pts is not None and pts.dim() >= 1 and pts.shape[0] == x.shape[0]:
+            pts = func(pts, [E, *pts.shape[1:]], **kwargs)
+        return _rebuild(x, qd, sc, pts)
+
+    @implements([aten.narrow.default])
+    def _narrow(func, types, args, kwargs):
+        x, dim, start, length = args[0], args[1], args[2], args[3]
+        if dim != 0:
+            raise NotImplementedError(f"NVFP4 narrow only dim 0, got {dim}")
+        qd = func(x.qdata, 0, start, length)
+        sc = func(x.scale, 0, start, length)
+        pts = x.per_tensor_scale
+        if pts is not None and pts.dim() >= 1 and pts.shape[0] == x.shape[0]:
+            pts = func(pts, 0, start, length)
+        return _rebuild(x, qd, sc, pts)
+
+    @implements([aten.view.default])
+    def _view(func, types, args, kwargs):
+        x, size = args[0], list(args[1])
+        if size == [-1]:
+            # FSDP's flat-buffer view. A 1-D NVFP4 is invalid (block layout needs >=2D),
+            # and this buffer is unused for subclass-extension params (the all-gather goes
+            # through fsdp_pre/post_all_gather), so return a valid 2-D [1, numel] NVFP4.
+            return _rebuild(
+                x, x.qdata.reshape(1, -1), x.scale.reshape(1, -1), x.per_tensor_scale
+            )
+        K = size[-1]
+        qd = func(x.qdata, size[:-1] + [K // 2])
+        sc = func(x.scale, size[:-1] + [K // x.block_size])
+        return _rebuild(x, qd, sc, x.per_tensor_scale)
+
+    @implements([aten.slice.Tensor])
+    def _slice(func, types, args, kwargs):
+        x = args[0]
+        dim = args[1] if len(args) > 1 else 0
+        if dim == 0:  # expert-axis slice (FSDP shard); torchao only handles rank 2
+            start = args[2] if len(args) > 2 else None
+            end = args[3] if len(args) > 3 else None
+            step = args[4] if len(args) > 4 else 1
+            if step != 1:
+                raise NotImplementedError("NVFP4 slice step must be 1")
+            qd = func(x.qdata, 0, start, end, 1)
+            sc = func(x.scale, 0, start, end, 1)
+            pts = x.per_tensor_scale
+            if pts is not None and pts.dim() >= 1 and pts.shape[0] == x.shape[0]:
+                pts = func(pts, 0, start, end, 1)
+            return _rebuild(x, qd, sc, pts)
+        from torchao.prototype.mx_formats.nvfp4_tensor import nvfp4_slice
+
+        return nvfp4_slice(func, types, args, kwargs)
+
+    def _cstride(shape):
+        st = [1] * len(shape)
+        for i in range(len(shape) - 2, -1, -1):
+            st[i] = st[i + 1] * shape[i + 1]
+        return st
+
+    @implements([aten.as_strided.default])
+    def _as_strided(func, types, args, kwargs):
+        # FSDP reconstructs the contiguous unsharded param via as_strided(orig_size,
+        # contiguous_stride, 0). The reconstructed NVFP4 already has the right qdata/scale
+        # (incl. swizzled scale) shapes; as_strided them to their own contiguous shapes (a
+        # view op must return NEW tensors, so we can't return x as-is).
+        x = args[0]
+        qshape, sshape = list(x.qdata.shape), list(x.scale.shape)
+        qd = func(x.qdata, qshape, _cstride(qshape), 0)
+        sc = func(x.scale, sshape, _cstride(sshape), 0)
+        return _rebuild(x, qd, sc, x.per_tensor_scale)
+
+    @implements([aten.copy_.default])
+    def _copy_(func, types, args, kwargs):
+        dst, src = args[0], args[1]
+        func(dst.qdata, src.qdata)
+        func(dst.scale, src.scale)
+        if src.per_tensor_scale is not None:
+            if (
+                dst.per_tensor_scale is not None
+                and dst.per_tensor_scale.shape == src.per_tensor_scale.shape
+            ):
+                func(dst.per_tensor_scale, src.per_tensor_scale)
+            else:
+                # dst lost its per_tensor buffer (e.g. a new_zeros without one); adopt src's.
+                dst.per_tensor_scale = src.per_tensor_scale.clone()
+        return dst
+
+    # FSDP2 subclass-extension hooks avoid the flat-buffer view(-1) path.
+    def fsdp_pre_all_gather(
+        self, mesh, outer_size=None, outer_stride=None, module=None, mp_policy=None
+    ):
+        inputs = (self.qdata, self.scale)
+        if self.per_tensor_scale is not None:
+            inputs = inputs + (self.per_tensor_scale,)
+        meta = (self.per_tensor_scale is not None,)
+        return inputs, meta
+
+    def fsdp_post_all_gather(
+        self, all_gather_outputs, metadata, param_dtype, *, out=None
+    ):
+        (has_pts,) = metadata
+        if has_pts:
+            qdata, scale, pts = all_gather_outputs
+        else:
+            qdata, scale = all_gather_outputs
+            pts = None
+        if out is not None:
+            # reconstruct in-place into the existing unsharded param
+            out.qdata, out.scale, out.per_tensor_scale = qdata, scale, pts
+            return
+        return _rebuild(self, qdata, scale, pts), all_gather_outputs
+
+    NVFP4Tensor.fsdp_pre_all_gather = fsdp_pre_all_gather
+    NVFP4Tensor.fsdp_post_all_gather = fsdp_post_all_gather
+
+    _PATCHED = True
+    LOG.info("Installed FSDP2 support (split + all-gather hooks) on NVFP4Tensor")
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_moe_loading.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_moe_loading.py
new file mode 100644
index 0000000000..d2d2d2107f
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_moe_loading.py
@@ -0,0 +1,571 @@
+"""Fix NVFP4 MoE-expert loading for DeepSeek-V4-Flash-NVFP4 (and similar checkpoints).
+
+transformers' finegrained_fp8 quantizer treats the model as FP8 blockwise and has NO
+NVFP4 path, so for a MIXED_PRECISION checkpoint (FP8 elsewhere, ``moe_quant_algo: NVFP4``
+for the routed experts) it fuses the expert *weights* into ``gate_up_proj``/``down_proj``
+(uint8 qdata) correctly but drops the NVFP4 *scales* — the per-expert ``weight_scale``
+(E4M3 group-16 block scale) + ``weight_scale_2`` (per-tensor) come in UNEXPECTED and the
+model's ``*_scale_inv`` placeholders are MISSING → randomly initialized → invalid model.
+
+This re-reads the per-expert ``weight_scale``/``weight_scale_2`` from the checkpoint, fuses
+them the same way as the weights (gate_up = cat(w1, w3) on the 2*intermediate axis; down =
+w2), and wraps the already-loaded fused qdata as a torchao ``NVFP4Tensor`` so the
+scattermoe NVFP4 path dequantizes correctly. Self-contained until an upstream transformers
+NVFP4-MoE path lands.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+
+import torch
+import torch.nn as nn
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _nvfp4_cls():
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+        return NVFP4Tensor
+    except ImportError:
+        return None
+
+
+def _resolve_repo_file(repo_id: str, filename: str) -> str:
+    """Resolve a checkpoint file path from a local snapshot dir or the HF hub.
+
+    A local snapshot dir (offline/air-gapped axolotl usage) would fail ``hf_hub_download``
+    with an ``HFValidationError``, so read straight from disk in that case.
+    """
+    if os.path.isdir(repo_id):
+        return os.path.join(repo_id, filename)
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id, filename)
+
+
+def _load_index(repo_id: str):
+    import json
+
+    with open(_resolve_repo_file(repo_id, "model.safetensors.index.json")) as f:
+        return json.load(f)["weight_map"]
+
+
+def _shard_open(repo_id: str, shard: str):
+    from safetensors import safe_open
+
+    return safe_open(_resolve_repo_file(repo_id, shard), framework="pt")
+
+
+def _safetensors_metadata(repo_id: str) -> dict[str, tuple[str, tuple[int, ...]]]:
+    """Return ``{tensor_name: (dtype_str, shape)}`` for every checkpoint tensor, read from the
+    safetensors *headers* only (no weight download). Works for a hub repo or a local snapshot dir.
+    """
+    if os.path.isdir(repo_id):
+        from safetensors import safe_open
+
+        out: dict[str, tuple[str, tuple[int, ...]]] = {}
+        index = os.path.join(repo_id, "model.safetensors.index.json")
+        if os.path.exists(index):
+            shards = sorted(set(_load_index(repo_id).values()))
+        else:  # single-file checkpoint
+            shards = ["model.safetensors"]
+        for shard in shards:
+            with safe_open(os.path.join(repo_id, shard), framework="pt") as f:
+                for name in f.keys():
+                    sl = f.get_slice(name)
+                    out[name] = (sl.get_dtype(), tuple(sl.get_shape()))
+        return out
+
+    from huggingface_hub import HfApi
+
+    meta = HfApi().get_safetensors_metadata(repo_id)
+    out = {}
+    for fmeta in meta.files_metadata.values():
+        for name, tinfo in fmeta.tensors.items():
+            out[name] = (str(tinfo.dtype), tuple(tinfo.shape))
+    return out
+
+
+# Known per-module leaf names across NVFP4 export conventions. The packed 4-bit weight is the
+# robust NVFP4 signal (uint8 qdata) — distinguishing NVFP4 from FP8 (which also carries a
+# ``weight_scale`` but stores an fp8 weight, not uint8). Names differ by exporter:
+#   modelopt          : weight        / weight_scale / weight_scale_2
+#   compressed-tensors: weight_packed / weight_scale / weight_global_scale
+_QDATA_LEAVES = ("weight", "weight_packed")
+_GROUP_SCALE_LEAVES = ("weight_scale",)
+_PER_TENSOR_LEAVES = ("weight_scale_2", "weight_global_scale")
+_ALL_LEAVES = (
+    "weight_scale_2",
+    "weight_global_scale",
+    "weight_scale",
+    "weight_packed",
+    "weight",
+    "input_scale",
+    "input_global_scale",
+)
+
+
+def inspect_nvfp4_layout(repo_id: str) -> dict:
+    """Detect a checkpoint's NVFP4 layout from its safetensors headers — no layout assumptions.
+
+    A module is treated as **NVFP4** only when it has a *packed* 4-bit weight (uint8 qdata, under
+    ``weight`` or ``weight_packed``) plus a group ``weight_scale`` — this distinguishes NVFP4 from
+    FP8 modules (which also carry a ``weight_scale`` but store an fp8, not uint8, weight) so a mixed
+    NVFP4+FP8 checkpoint is classified correctly. The NVFP4 modules are split into:
+      * ``routed_projs``  — proj names under ``...experts.<int>.<proj>`` (fused into 3D expert params),
+      * ``nonrouted_suffixes`` — layer-relative paths of every other NVFP4 linear THIS checkpoint
+        quantizes (shared experts, dense MLPs, ...),
+    plus the detected key ``naming`` (``modelopt`` vs ``compressed-tensors`` vs ``mixed``) so the
+    caller can tell whether its converters (which assume modelopt ``weight``/``weight_scale_2``)
+    apply. Everything is derived from the headers; differing exports are described, not assumed.
+    """
+    import re
+
+    meta = _safetensors_metadata(repo_id)
+    bases: dict[str, dict[str, tuple[str, tuple[int, ...]]]] = {}
+    for name, dt_shape in meta.items():
+        for leaf in _ALL_LEAVES:
+            if name.endswith("." + leaf):
+                bases.setdefault(name[: -len(leaf) - 1], {})[leaf] = dt_shape
+                break
+
+    def _qdata(parts):
+        for leaf in _QDATA_LEAVES:
+            if leaf in parts and parts[leaf][0] == "U8":
+                return leaf
+        return None
+
+    routed_re = re.compile(r"\.experts\.\d+\.([A-Za-z_][A-Za-z0-9_]*)$")
+    layer_re = re.compile(r"^.*?layers\.\d+\.")
+    routed_projs: list[str] = []
+    routed_sample: dict[str, tuple | None] = {}
+    nonrouted: dict[str, dict] = {}
+    qdata_names: set[str] = set()
+    per_tensor_names: set[str] = set()
+    for base, parts in bases.items():
+        qd = _qdata(parts)
+        is_nvfp4 = qd is not None and any(g in parts for g in _GROUP_SCALE_LEAVES)
+        if not is_nvfp4:  # bf16 (excluded) or fp8 module — not NVFP4
+            continue
+        qdata_names.add(qd)
+        for leaf in _PER_TENSOR_LEAVES:
+            if leaf in parts:
+                per_tensor_names.add(leaf)
+        m = routed_re.search(base)
+        if m:
+            proj = m.group(1)
+            if proj not in routed_projs:
+                routed_projs.append(proj)
+                routed_sample[proj] = parts.get(qd)
+        else:
+            nonrouted.setdefault(layer_re.sub("", base), parts)
+
+    modelopt = qdata_names <= {"weight"} and per_tensor_names <= {"weight_scale_2"}
+    ct = qdata_names <= {"weight_packed"} and per_tensor_names <= {
+        "weight_global_scale"
+    }
+    naming = (
+        "modelopt" if modelopt else ("compressed-tensors" if ct else "mixed/unknown")
+    )
+
+    return {
+        "routed_present": bool(routed_projs),
+        "routed_projs": sorted(routed_projs),
+        "routed_sample_shapes": routed_sample,
+        "nonrouted_suffixes": sorted(nonrouted),
+        "nonrouted_sample_shapes": nonrouted,
+        "qdata_names": sorted(qdata_names),
+        "per_tensor_names": sorted(per_tensor_names),
+        "naming": naming,
+    }
+
+
+def fuse_nvfp4_experts(projs: list[dict], *, block_size: int = 16, dtype=None):
+    """Shared NVFP4-expert fusion core (one implementation for every load path).
+
+    ``projs`` is the ordered list of projections to fuse on the N (row) axis — one entry for a
+    single proj (``down``), two for a fused ``gate_up`` — each a dict of per-expert lists::
+
+        {"qd": [E uint8 [N, K/2]], "sc": [E e4m3 [N, K/16]], "pts": [E f32 scalar]}
+
+    Returns a torchao ``NVFP4Tensor`` ``[E, ΣN, K/2]``. The fused tensor carries ONE per-expert
+    per-tensor scale (``pts``), so when the projections export different ``pts`` the fused pts is
+    their elementwise max and each proj's ratio (always <= 1, so the e4m3 cast cannot overflow to
+    NaN) is folded into its group scale (dequant = qdata · group_scale · per_tensor_scale, so this
+    is exact up to e4m3 rounding) rather than silently dropping it. Used by both the
+    ``WeightConverter`` (modelopt skeleton load) and the post-load scale-attach path so the fusion
+    + scale reconciliation live in exactly one place.
+    """
+    import torch as _torch
+
+    NVFP4Tensor = _nvfp4_cls()
+    if NVFP4Tensor is None:
+        raise RuntimeError("torchao NVFP4Tensor not available")
+    dtype = dtype or _torch.bfloat16
+
+    # cpu_ram_efficient_loading runs the converter once on META placeholders (shape inference) and
+    # again on the real tensors on rank 0. On meta the scale-reconciliation value ops (`allclose`,
+    # `.item()`) can't run, so skip them — only the OUTPUT SHAPE matters on the meta pass, and the
+    # stack/cat below produce correctly-shaped meta tensors regardless.
+    is_meta = bool(projs[0]["qd"]) and projs[0]["qd"][0].is_meta
+
+    pts_all = [
+        _torch.stack([t.to(_torch.float32) for t in p["pts"]]).view(-1, 1, 1)
+        for p in projs
+    ]
+    pts_fused = pts_all[0]
+    for pts_i in pts_all[1:]:
+        pts_fused = _torch.maximum(pts_fused, pts_i)
+    qdatas, scales = [], []
+    for i, p in enumerate(projs):
+        qd = _torch.stack(list(p["qd"]), dim=0)  # [E, N, K/2]
+        sc = _torch.stack(list(p["sc"]), dim=0)  # [E, N, K/16]
+        if not is_meta and not _torch.allclose(pts_all[i], pts_fused):
+            sc = (sc.to(_torch.float32) * (pts_all[i] / pts_fused)).to(
+                _torch.float8_e4m3fn
+            )
+            LOG.warning(
+                "fuse_nvfp4_experts: proj #%d per-tensor scale differs from the fused max; "
+                "folded the ratio into its group scale (min %.4g)",
+                i,
+                (pts_all[i] / pts_fused).min().item(),
+            )
+        qdatas.append(qd)
+        scales.append(sc)
+    qdata = qdatas[0] if len(qdatas) == 1 else _torch.cat(qdatas, dim=1)
+    scale = scales[0] if len(scales) == 1 else _torch.cat(scales, dim=1)
+    return NVFP4Tensor(qdata, scale, block_size, dtype, per_tensor_scale=pts_fused)
+
+
+def patch_nvfp4_tensor_meta_ops() -> None:
+    """Register ``zeros_like`` / ``empty_like`` / ``new_zeros`` on torchao's NVFP4Tensor.
+
+    FSDP2's ``cpu_ram_efficient_loading`` keeps non-rank-0 params on ``meta`` and materializes the
+    receive buffers with ``zeros_like`` / ``empty_like`` before scattering rank 0's shards. torchao
+    doesn't implement those for NVFP4Tensor (only matmul / view / slice / copy), so a 4-bit expert
+    param hits 'unimplemented operator'. Each just applies the op to the packed data + scales and
+    rebuilds the tensor, preserving block_size / dtype / per-tensor scale. Idempotent."""
+    import torch as _torch
+
+    NVFP4Tensor = _nvfp4_cls()
+    if NVFP4Tensor is None or getattr(NVFP4Tensor, "_axolotl_meta_ops", False):
+        return
+    from torchao.utils import return_and_correct_aliasing
+
+    aten = _torch.ops.aten
+
+    @NVFP4Tensor.implements([aten.zeros_like.default, aten.empty_like.default])
+    def _nvfp4_like(func, types, args, kwargs):
+        # FSDP2 passes device / pin_memory / memory_format; forward them to the packed data + scale
+        # tensors. dtype/layout would break the int-packed payload, so drop them.
+        passthru = {
+            k: v
+            for k, v in kwargs.items()
+            if k in ("device", "pin_memory", "memory_format")
+        }
+        out = args[0]._apply_fn_to_data(lambda x: func(x, **passthru))
+        return return_and_correct_aliasing(func, args, kwargs, out)
+
+    @NVFP4Tensor.implements([aten.new_zeros.default])
+    def _nvfp4_new_zeros(func, types, args, kwargs):
+        out = args[0]._apply_fn_to_data(lambda x: _torch.zeros_like(x))
+        return return_and_correct_aliasing(func, args, kwargs, out)
+
+    NVFP4Tensor._axolotl_meta_ops = True
+
+
+# Per-architecture checkpoint naming for the unfused per-expert NVFP4 tensors.
+# base_fmt formats with (layer, e, proj); gate_up/down are the proj names fused
+# (gate_up = cat on the N/row axis in this order; down is single).
+_NVFP4_MOE_SCHEMES = {
+    # DeepSeek-V4-Flash-NVFP4: w1=gate, w3=up, w2=down under ``layers.N.ffn.experts.M``
+    "dsv4": {
+        "base_fmt": "layers.{layer}.ffn.experts.{e}.{proj}",
+        "gate_up": ("w1", "w3"),
+        "down": ("w2",),
+    },
+    # Gemma-4-A4B-NVFP4: separate gate/up/down under ``model.language_model.layers.N.experts.M``
+    "gemma4": {
+        "base_fmt": "model.language_model.layers.{layer}.experts.{e}.{proj}",
+        "gate_up": ("gate_proj", "up_proj"),
+        "down": ("down_proj",),
+    },
+}
+
+
+def _detect_scheme(wmap):
+    """Pick the checkpoint naming scheme whose layer-0 expert-0 weight key is present."""
+    for name, sch in _NVFP4_MOE_SCHEMES.items():
+        probe = sch["base_fmt"].format(layer=0, e=0, proj=sch["gate_up"][0]) + ".weight"
+        if probe in wmap:
+            return name, sch
+    return None, None
+
+
+def _build_expert_nvfp4(repo_id, wmap, base_fmt, layer, projs, n_experts, device):
+    """Rebuild a fused NVFP4Tensor for one expert projection group straight from the raw
+    checkpoint (no dependence on transformers' own fusion), reading every proj's per-expert qdata,
+    E4M3 block scale and per-tensor scale, then delegating the stack/concat/scale-reconcile to the
+    shared :func:`fuse_nvfp4_experts` core. Returns an ``NVFP4Tensor`` ``[E, ΣN, K/2]`` on ``device``."""
+    proj_parts = [{"qd": [], "sc": [], "pts": []} for _ in projs]
+    opened: dict[str, object] = {}
+    for e in range(n_experts):
+        for pi, proj in enumerate(projs):
+            base = base_fmt.format(layer=layer, e=e, proj=proj)
+            shard = wmap[f"{base}.weight"]
+            f = opened.get(shard) or opened.setdefault(
+                shard, _shard_open(repo_id, shard)
+            )
+            proj_parts[pi]["qd"].append(f.get_tensor(f"{base}.weight").to(device))
+            proj_parts[pi]["sc"].append(f.get_tensor(f"{base}.weight_scale").to(device))
+            proj_parts[pi]["pts"].append(
+                f.get_tensor(f"{base}.weight_scale_2").to(torch.float32).to(device)
+            )
+    return fuse_nvfp4_experts(proj_parts)
+
+
+def attach_nvfp4_expert_scales(model: nn.Module, repo_id: str) -> int:
+    """Wrap each MoE experts module's fused gate_up_proj/down_proj as an NVFP4Tensor rebuilt from
+    the checkpoint's per-expert E4M3 scales. Handles both checkpoint layouts (DeepSeek-V4 ``ffn.
+    experts``/w1-w3-w2 with uint8 fused params from the FP8 quantizer, and Gemma-4 ``language_model.
+    layers.N.experts.M``/gate-up-down where transformers leaves the fused param a random bf16
+    placeholder). Returns the number of expert modules fixed."""
+    NVFP4Tensor = _nvfp4_cls()
+    if NVFP4Tensor is None:
+        return 0
+    wmap = _load_index(repo_id)
+    scheme_name, scheme = _detect_scheme(wmap)
+    if scheme is None:
+        LOG.warning(
+            "attach_nvfp4_expert_scales: no known NVFP4 MoE naming scheme in %s",
+            repo_id,
+        )
+        return 0
+    base_fmt, projs_gu, projs_dn = scheme["base_fmt"], scheme["gate_up"], scheme["down"]
+    fixed = 0
+    for name, mod in model.named_modules():
+        # Match the fused experts param by shape (3D stacked experts), not dtype/class: DSV4's
+        # FP8 quantizer leaves it uint8, gemma4 leaves it a random bf16 placeholder.
+        gup = getattr(mod, "gate_up_proj", None)
+        dn = getattr(mod, "down_proj", None)
+        if not (
+            isinstance(gup, torch.Tensor)
+            and isinstance(dn, torch.Tensor)
+            and gup.ndim == 3
+            and dn.ndim == 3
+        ):
+            continue
+        m = re.search(r"layers\.(\d+)\.", name)
+        if not m:
+            continue
+        layer = int(m.group(1))
+        # Skip layers whose experts aren't in the checkpoint under this scheme (e.g. dense layers).
+        if base_fmt.format(layer=layer, e=0, proj=projs_gu[0]) + ".weight" not in wmap:
+            continue
+        E = gup.shape[0]
+        mod.gate_up_proj = nn.Parameter(
+            _build_expert_nvfp4(
+                repo_id, wmap, base_fmt, layer, projs_gu, E, gup.device
+            ),
+            requires_grad=False,
+        )
+        mod.down_proj = nn.Parameter(
+            _build_expert_nvfp4(repo_id, wmap, base_fmt, layer, projs_dn, E, dn.device),
+            requires_grad=False,
+        )
+        # Drop the stale FP8 scale params the quantizer created for these experts
+        # (randomly-initialized `*_scale_inv` from the MISSING-key path, plus any FP8
+        # `*_scale`/`input_scale`). They're unused once the experts are NVFP4Tensor routed
+        # through scattermoe, and would otherwise be sharded by FSDP / waste memory.
+        for stale in (
+            "gate_up_proj_scale_inv",
+            "down_proj_scale_inv",
+            "gate_up_proj_scale",
+            "down_proj_scale",
+            "gate_up_proj_scale_2",
+            "down_proj_scale_2",
+            "gate_up_proj_input_scale",
+            "down_proj_input_scale",
+        ):
+            for store in (mod._parameters, mod._buffers):
+                if stale in store:
+                    del store[stale]
+        del gup, dn
+        fixed += 1
+    if fixed:
+        LOG.info(
+            "Attached NVFP4 expert scales (rebuilt %d %s experts as NVFP4Tensor)",
+            fixed,
+            scheme_name,
+        )
+    return fixed
+
+
+def patch_skip_missing_expert_init() -> None:
+    """Skip transformers' random init of the fused routed-expert params under direct load.
+
+    With the routed converters skipped (direct-load path), ``experts.gate_up_proj``/``down_proj``
+    are MISSING keys, so ``_initialize_missing_keys`` fills them with CPU ``normal_()`` right
+    before :func:`direct_load_nvfp4_experts` overwrites every byte. At Qwen3-Next-80B scale that
+    is ~155 GB of bf16 randn at a measured 0.11 GB/s: ~23 min of pure waste per launch. Mark the
+    fused 3D expert params ``_is_hf_initialized`` before init runs so ``initialize_weights``
+    skips them. Only call when the direct load WILL fill these params (same contract as skipping
+    the routed converters). Idempotent.
+    """
+    from transformers.modeling_utils import PreTrainedModel
+
+    if getattr(
+        PreTrainedModel._initialize_missing_keys, "_axolotl_skip_expert_init", False
+    ):
+        return
+    orig = PreTrainedModel._initialize_missing_keys
+
+    def patched(self, *args, **kwargs):
+        for mod in self.modules():
+            gup = getattr(mod, "gate_up_proj", None)
+            dn = getattr(mod, "down_proj", None)
+            if (
+                isinstance(gup, torch.Tensor)
+                and isinstance(dn, torch.Tensor)
+                and gup.ndim == 3
+                and dn.ndim == 3
+            ):
+                gup._is_hf_initialized = True
+                dn._is_hf_initialized = True
+        return orig(self, *args, **kwargs)
+
+    patched._axolotl_skip_expert_init = True  # type: ignore[attr-defined]
+    PreTrainedModel._initialize_missing_keys = patched
+    LOG.info(
+        "Patched _initialize_missing_keys to skip fused expert params (direct load fills them)"
+    )
+
+
+def direct_load_nvfp4_experts(model, repo_id: str, routed_projs: list[str]) -> int:
+    """FAST routed-expert load that BYPASSES transformers' conversion loader.
+
+    Transformers' loader spends ~7 min on GLM-5.2 iterating/matching ~240k per-expert source tensors
+    through its Python machinery; a DIRECT read+fuse of the same experts is ~25s (profiled). This
+    reads each MoE layer's routed-expert ``weight``(uint8)/``weight_scale``(fp8)/``weight_scale_2``
+    components straight from the safetensors (native dtype) and fuses them into the 3D NVFP4 expert
+    params, mirroring :class:`Nvfp4ExpertsDeserialize` but without the per-tensor loader overhead.
+
+    Only the LOCAL-RANK-0 process materializes (others stay meta for the FSDP broadcast). Requires the
+    routed converters to be SKIPPED at registration (so transformers leaves the fused params unfilled).
+    Returns the number of fused params filled. Safe no-op on non-rank0 / no routed experts.
+    """
+    import os
+    import re
+
+    import torch
+
+    if str(os.environ.get("LOCAL_RANK", "0")) not in ("0", ""):
+        return 0
+    if not routed_projs:
+        return 0
+    NVFP4Tensor = _nvfp4_cls()
+    if NVFP4Tensor is None:
+        return 0
+    wmap = _load_index(repo_id)
+    # per-layer expert count, from the index (no reads): count ...experts.<e>.<proj0>.weight keys.
+    proj0 = routed_projs[0]
+    layer_E: dict[int, int] = {}
+    pat = re.compile(
+        r"\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + re.escape(proj0) + r"\.weight$"
+    )
+    for key in wmap:
+        m = pat.search(key)
+        if m:
+            L, e = int(m.group(1)), int(m.group(2))
+            layer_E[L] = max(layer_E.get(L, 0), e + 1)
+
+    handles: dict[str, object] = {}
+
+    def gt(key):
+        sh = wmap[key]
+        if sh not in handles:
+            handles[sh] = _shard_open(repo_id, sh)
+        return handles[sh].get_tensor(key)  # native dtype (uint8/fp8/fp32)
+
+    fused_map = (
+        ("gate_up_proj", ["gate_proj", "up_proj"]),
+        ("down_proj", ["down_proj"]),
+    )
+    n = 0
+    for mod_name, mod in model.named_modules():
+        if not (hasattr(mod, "gate_up_proj") and hasattr(mod, "down_proj")):
+            continue
+        m = re.search(r"\.layers\.(\d+)\.", "." + mod_name)
+        if m is None:
+            continue
+        L = int(m.group(1))
+        if L not in layer_E:
+            continue
+        E = layer_E[L]
+        base = f"model.layers.{L}.mlp.experts"
+        for fused, parts in fused_map:
+            sel = [p for p in parts if p in routed_projs]
+            if not sel:
+                continue
+            projs = [
+                {
+                    "qd": [gt(f"{base}.{e}.{p}.weight") for e in range(E)],
+                    "sc": [gt(f"{base}.{e}.{p}.weight_scale") for e in range(E)],
+                    "pts": [gt(f"{base}.{e}.{p}.weight_scale_2") for e in range(E)],
+                }
+                for p in sel
+            ]
+            nvfp4 = fuse_nvfp4_experts(projs)
+            setattr(mod, fused, torch.nn.Parameter(nvfp4, requires_grad=False))
+            n += 1
+    return n
+
+
+if __name__ == "__main__":  # local self-consistency test on real layer-0 data
+    REPO = "nvidia/DeepSeek-V4-Flash-NVFP4"
+    NVFP4Tensor = _nvfp4_cls()
+    wmap = _load_index(REPO)
+    _, _scheme = _detect_scheme(wmap)
+    _base_fmt = _scheme["base_fmt"]
+    dev = "cuda"
+    # fused gate_up qdata+scale from w1+w3 (expert 0 only via n_experts=1 slice below)
+    gqd, gscale, gpts = _build_expert_nvfp4(
+        REPO, wmap, _base_fmt, 0, ("w1", "w3"), 4, dev
+    )
+    print(
+        "fused gate_up qdata",
+        gqd.shape,
+        "scale",
+        gscale.shape,
+        "per_tensor",
+        gpts.item(),
+    )
+    # self-consistency: NVFP4 of fused must dequant to cat(dequant(w1), dequant(w3))
+    f = _shard_open(REPO, wmap["layers.0.ffn.experts.0.w1.weight"])
+    qd1 = f.get_tensor("layers.0.ffn.experts.0.w1.weight").to(dev)
+    qd3 = f.get_tensor("layers.0.ffn.experts.0.w3.weight").to(dev)
+    s1 = f.get_tensor("layers.0.ffn.experts.0.w1.weight_scale").to(dev)
+    s3 = f.get_tensor("layers.0.ffn.experts.0.w3.weight_scale").to(dev)
+    p = f.get_tensor("layers.0.ffn.experts.0.w1.weight_scale_2").to(dev).float()
+    d1 = NVFP4Tensor(qd1, s1, 16, torch.bfloat16, per_tensor_scale=p).dequantize(
+        torch.bfloat16
+    )
+    d3 = NVFP4Tensor(qd3, s3, 16, torch.bfloat16, per_tensor_scale=p).dequantize(
+        torch.bfloat16
+    )
+    fused = NVFP4Tensor(
+        gqd[0], gscale[0], 16, torch.bfloat16, per_tensor_scale=gpts
+    ).dequantize(torch.bfloat16)
+    ref = torch.cat([d1, d3], dim=0)
+    print(
+        "fusion self-consistent:",
+        torch.equal(fused, ref),
+        "max_err",
+        (fused - ref).abs().max().item(),
+    )
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_nonexpert.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_nonexpert.py
new file mode 100644
index 0000000000..8f844ab70a
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_nonexpert.py
@@ -0,0 +1,67 @@
+"""In-place NVFP4 W4A16 quantization of gemma4 non-expert linears (Marlin path).
+
+Mirrors :mod:`gemma4_nf4_nonexpert` but swaps each frozen non-expert ``nn.Linear`` for a
+:class:`MarlinW4A16Linear` — 4-bit NVFP4 weights run through the same validated Marlin kernel as
+the grouped MoE experts (bf16 act/out, no FP4 hardware needed; sm80+). Vs the bnb NF4 path this
+keeps experts and non-experts on ONE tensor-core 4-bit code path instead of bnb's dequant kernel.
+
+Targets every ``nn.Linear`` that is NOT inside an expert block; norms, embeddings, routers,
+vision/audio towers, and lm_head are skipped by name. No LoRA is attached (experts-only).
+"""
+
+from __future__ import annotations
+
+import torch.nn as nn
+
+from axolotl.utils.logging import get_logger
+
+from .gemma4_fp8_nonexpert import _should_skip_by_name
+from .gemma4_nf4_nonexpert import _expert_paths
+
+LOG = get_logger(__name__)
+
+
+def quantize_gemma4_nonexpert_nvfp4(model: nn.Module) -> int:
+    """Swap non-expert ``nn.Linear`` modules for frozen NVFP4 ``MarlinW4A16Linear``, in place.
+    Returns the count swapped. Experts (NVFP4Tensor) are untouched. Idempotent."""
+    from .marlin_w4a16 import marlin_w4a16_available
+    from .marlin_w4a16.nonexpert_linear import MarlinW4A16Linear
+
+    if not marlin_w4a16_available():
+        raise RuntimeError(
+            "nonexpert_quantization=nvfp4 needs the Marlin W4A16 kernel (sm80+ GPU + CUDA "
+            "toolkit). Use nonexpert_quantization=nf4 (bitsandbytes) on unsupported setups."
+        )
+
+    expert_paths = _expert_paths(model)
+
+    def _under_expert(path: str) -> bool:
+        return any(path == ep or path.startswith(ep + ".") for ep in expert_paths)
+
+    targets: list[tuple[str, nn.Linear]] = []
+    for name, mod in model.named_modules():
+        if isinstance(mod, MarlinW4A16Linear):
+            continue
+        if not isinstance(mod, nn.Linear):
+            continue
+        if _under_expert(name) or _should_skip_by_name(name):
+            continue
+        if not isinstance(mod.weight, nn.Parameter) or mod.weight.ndim != 2:
+            continue
+        targets.append((name, mod))
+
+    for name, mod in targets:
+        new = MarlinW4A16Linear(
+            mod.weight.data, mod.bias.data if mod.bias is not None else None
+        )
+        parent_path, _, attr = name.rpartition(".")
+        parent = model.get_submodule(parent_path) if parent_path else model
+        setattr(parent, attr, new)
+
+    if targets:
+        LOG.info(
+            "Gemma4 NVFP4 non-expert: swapped %d non-expert linears to Marlin W4A16 "
+            "(4-bit weight, bf16 compute)",
+            len(targets),
+        )
+    return len(targets)
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_weight_converter.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_weight_converter.py
new file mode 100644
index 0000000000..c6db6c09f9
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_weight_converter.py
@@ -0,0 +1,517 @@
+"""Native WeightConverter so Gemma-4 NVFP4 MoE experts load as NVFP4Tensor.
+
+nvidia/Gemma-4-26B-A4B-NVFP4 ships per-expert weights under
+  ``model.language_model.layers.N.experts.E.{gate_proj,up_proj,down_proj}.{weight,weight_scale,weight_scale_2}``
+but ``quant_method: modelopt`` is not a recognized transformers quantizer, so
+the model loads as a BF16 skeleton with the per-expert NVFP4 tensors landing
+as UNEXPECTED and the fused ``gate_up_proj``/``down_proj`` remaining random BF16.
+
+This module registers a ``WeightConverter`` for the ``gemma4_text`` model type
+that fuses the per-expert raw uint8 qdata + E4M3 block scales + per-tensor
+scalar into a single ``NVFP4Tensor`` (packed 4-bit) and assigns it in-place to
+the ``Gemma4TextExperts`` module — exactly like ``Mxfp4Deserialize`` does for
+MXFP4.  ``is_nvfp4_param(param)`` returns ``True`` on the result, activating
+the scattermoe fused NVFP4 path.
+
+The fusion itself (stack experts, cat gate/up on the N axis, reconcile the per-tensor scales)
+is the shared ``nvfp4_moe_loading.fuse_nvfp4_experts`` core, so the WeightConverter (modelopt
+skeleton load) and the post-load scale-attach path use one implementation instead of duplicating
+the NVFP4-expert math.
+
+Registration is done via ``transformers.conversion_mapping.register_checkpoint_conversion_mapping``
+— no site-packages edits.  The registration helper is gated: call it only when
+the model is gemma4 + NVFP4 modelopt.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _nvfp4_cls():
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+        return NVFP4Tensor
+    except ImportError:
+        return None
+
+
+# Set True only by patch_conversion_loader_rank0_only() (installed solely under FSDP
+# cpu_ram_efficient_loading). Gates the converter meta-load so non-rank-0 stays on meta ONLY when the
+# FSDP broadcast will later fill it — otherwise (e.g. DDP, or the example without the flag) every rank
+# loads real weights as normal.
+_RANK0_ONLY_ACTIVE = False
+
+
+def _nonrank0_meta_load() -> bool:
+    """True when cpu_ram_efficient broadcast loading wants this rank to stay on meta.
+
+    The safetensors are mmap'd, so a converter that calls ``torch.stack`` on the checkpoint
+    tensors copies real data into RAM on EVERY rank — violating the rank0-only contract and
+    blowing CPU RAM up by the world size. On non-local-rank-0 we instead emit correctly-shaped
+    META params; FSDP's ``fsdp2_load_full_state_dict`` then broadcasts rank 0's real weights.
+    (With the ``_materialize_copy`` patch active the converter already receives meta inputs; this
+    stays as a defensive second gate.)"""
+    return _RANK0_ONLY_ACTIVE and _is_nonrank0_process()
+
+
+def _to_meta(t):
+    import torch as _torch
+
+    return _torch.empty(t.shape, dtype=t.dtype, device="meta")
+
+
+def _is_nonrank0_process() -> bool:
+    """True on every process except local-rank-0 of the node.
+
+    torchrun always exports ``LOCAL_RANK``; we key off it directly rather than transformers'
+    ``is_fsdp_enabled()`` because that ALSO requires ``ACCELERATE_USE_FSDP`` /
+    ``FSDP_CPU_RAM_EFFICIENT_LOADING`` in the env, which ``axolotl train`` does NOT set at
+    ``from_pretrained`` time — so transformers' own rank0-gate is dormant during the load. ``-1``
+    (unset) means a single non-distributed process → load normally."""
+    import os
+
+    return int(os.environ.get("LOCAL_RANK", "-1")) > 0
+
+
+def patch_conversion_loader_rank0_only() -> None:
+    """Make transformers' conversion-based loader load rank0-only (for FSDP cpu_ram_efficient).
+
+    transformers' new ``core_model_loading`` loader (engaged whenever a ``conversion_mapping`` is
+    registered — i.e. our NVFP4 converters) materializes every tensor on every rank: it only shards
+    per-rank when a ``tp_plan`` is set, and otherwise has NO rank gating. Under FSDP that loads the
+    full model on all ``world_size`` ranks → an N× CPU-RAM blowup that OOMs large models. (And even
+    the legacy gate it would mirror is dormant here — see :func:`_is_nonrank0_process`.)
+
+    On non-local-rank-0 return a META tensor from ``_materialize_copy`` so only rank 0 holds real
+    weights; ``fsdp2_load_full_state_dict`` then broadcasts them. The mmap'd slice read still happens
+    (cheap, shared OS page cache) but the copy is dropped to meta, so non-rank-0 never accumulates.
+    Install ONLY when cpu_ram_efficient broadcast loading is active, else it would starve DDP ranks
+    that each need their own real weights. Idempotent."""
+    global _RANK0_ONLY_ACTIVE
+    _RANK0_ONLY_ACTIVE = True
+
+    import os
+
+    import torch as _torch
+    import transformers.core_model_loading as cml
+
+    # transformers caps the conversion-loader ThreadPoolExecutor at min(4, cpu_count). Measured NOT
+    # to help GLM-5.2 load time (the bottleneck is the SERIAL main-thread converter/match loop over
+    # ~240k source tensors, not the materialize I/O — disk is fast NVMe). Left as an explicit opt-in
+    # (AXOLOTL_LOAD_WORKERS=<n>) in case it helps on a slow-disk/network box; no default change.
+    try:
+        _w = int(os.environ.get("AXOLOTL_LOAD_WORKERS", "0") or 0)
+        if _w > getattr(cml, "GLOBAL_WORKERS", 4):
+            cml.GLOBAL_WORKERS = _w
+            LOG.info("Raised transformers conversion-loader workers to %d", _w)
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+    if getattr(cml._materialize_copy, "_axolotl_rank0_patched", False):
+        return
+
+    _orig = cml._materialize_copy
+
+    # safetensors dtype-string -> torch dtype, for building meta tensors without reading the slice.
+    _ST_DTYPE = {
+        "F64": _torch.float64,
+        "F32": _torch.float32,
+        "F16": _torch.float16,
+        "BF16": _torch.bfloat16,
+        "F8_E4M3": _torch.float8_e4m3fn,
+        "F8_E5M2": _torch.float8_e5m2,
+        "I64": _torch.int64,
+        "I32": _torch.int32,
+        "I16": _torch.int16,
+        "I8": _torch.int8,
+        "U8": _torch.uint8,
+        "BOOL": _torch.bool,
+    }
+
+    def _meta_from_slice(tensor, dtype):
+        """Build a meta tensor matching what ``_orig`` WOULD return, without reading from disk."""
+        if not (hasattr(tensor, "get_shape") and hasattr(tensor, "get_dtype")):
+            return None
+        out_dtype = dtype if dtype is not None else _ST_DTYPE.get(tensor.get_dtype())
+        if out_dtype is None:
+            return None
+        return _torch.empty(tensor.get_shape(), dtype=out_dtype, device="meta")
+
+    _stats = {"meta": 0, "real": 0, "fallback": 0, "logged": False}
+
+    def _dbg(msg):
+        import sys
+
+        print(
+            f"[MATCOPY rank={os.environ.get('LOCAL_RANK', '?')}] {msg}",
+            file=sys.stderr,
+            flush=True,
+        )
+
+    def _patched_materialize_copy(tensor, device=None, dtype=None):
+        # dtype-aware: load the RAW NVFP4 components (uint8 qdata / fp8 e4m3|e5m2 scales) at their
+        # native dtype instead of the bf16 skeleton dtype. modelopt is an unrecognized quantizer so
+        # transformers' own native-dtype branch (hf_quantizer.pre_quantized) never fires and it casts
+        # everything to bf16 — doubling the bytes of the largest (uint8) expert tensors and adding a
+        # cast the converter's _recast_weight/_recast_scale would just undo. Applied before the
+        # rank0/meta split so both produce the same native dtype. (F32/F16/BF16 sources unaffected.)
+        if dtype is not None and hasattr(tensor, "get_dtype"):
+            _st = tensor.get_dtype()
+            if _st == "U8" or _st.startswith("F8"):
+                dtype = None
+        nonrank0 = _is_nonrank0_process()
+        if not _stats["logged"]:
+            _dbg(f"first call: nonrank0={nonrank0} slice_type={type(tensor).__name__}")
+            _stats["logged"] = True
+        # Non-rank-0: never touch the data — produce a same-shape/dtype META tensor so the slice is
+        # never read into RAM (reading-then-discarding leaves glibc holding the freed pages).
+        if nonrank0:
+            meta = _meta_from_slice(tensor, dtype)
+            if meta is not None:
+                _stats["meta"] += 1
+                if _stats["meta"] % 2000 == 0:
+                    _dbg(f"meta={_stats['meta']} fallback={_stats['fallback']}")
+                return meta
+            # Fallback (unknown slice type): read then drop to meta — correctness over memory.
+            _stats["fallback"] += 1
+            if _stats["fallback"] % 500 == 0:
+                _dbg(
+                    f"FALLBACK count={_stats['fallback']} type={type(tensor).__name__}"
+                )
+            out = _orig(tensor, device=device, dtype=dtype)
+            return _torch.empty(out.shape, dtype=out.dtype, device="meta")
+        _stats["real"] += 1
+        return _orig(tensor, device=device, dtype=dtype)
+
+    _patched_materialize_copy._axolotl_rank0_patched = True  # type: ignore[attr-defined]
+    cml._materialize_copy = _patched_materialize_copy
+    LOG.info(
+        "Patched transformers core_model_loading._materialize_copy for rank0-only loading "
+        "(LOCAL_RANK=%s)",
+        os.environ.get("LOCAL_RANK", "-1"),
+    )
+
+
+class Nvfp4ExpertsDeserialize:
+    """ConversionOps that fuses per-expert NVFP4 tensors into a single NVFP4Tensor.
+
+    For gate_up_proj, ``input_dict`` contains four keys (the source_patterns):
+      - ``"experts.*.gate_proj.weight"``  → list of E uint8 tensors  [I, H/2]
+      - ``"experts.*.up_proj.weight"``    → list of E uint8 tensors  [I, H/2]
+      - ``"experts.*.gate_proj.weight_scale"``  → list of E e4m3 tensors [I, H/16]
+      - ``"experts.*.up_proj.weight_scale"``    → list of E e4m3 tensors [I, H/16]
+      - ``"experts.*.gate_proj.weight_scale_2"`` → list of E float32 scalars (use first)
+
+    For down_proj, ``input_dict`` contains:
+      - ``"experts.*.down_proj.weight"``         → list of E uint8 tensors  [H, I/2]
+      - ``"experts.*.down_proj.weight_scale"``   → list of E e4m3 tensors [H, I/16]
+      - ``"experts.*.down_proj.weight_scale_2"`` → list of E float32 scalars
+
+    The op attaches the fused NVFP4Tensor to the module in-place and returns ``{}``
+    so the loader does not try to materialize the original meta-parameter names.
+    """
+
+    def convert(
+        self,
+        input_dict: dict[str, Any],
+        source_patterns: list[str] | None = None,
+        target_patterns: list[str] | None = None,
+        full_layer_name: str | None = None,
+        model: nn.Module | None = None,
+        missing_keys: set | None = None,
+        **kwargs,
+    ) -> dict[str, Any]:
+        from transformers.quantizers.quantizers_utils import get_module_from_name
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_moe_loading import (
+            fuse_nvfp4_experts,
+        )
+
+        if full_layer_name is None or "gate_up_proj" not in full_layer_name:
+            proj = "down_proj"
+        else:
+            proj = "gate_up_proj"
+
+        # cpu_ram_efficient_loading: only local-rank-0 materializes; others stay on meta and get
+        # filled by the FSDP broadcast. Drop the mmap'd checkpoint data to meta BEFORE fusing.
+        if _nonrank0_meta_load():
+            input_dict = {
+                k: (
+                    [_to_meta(t) for t in v]
+                    if isinstance(v, (list, tuple))
+                    else _to_meta(v)
+                )
+                for k, v in input_dict.items()
+            }
+
+        def _find(pat_suffix: str) -> list[torch.Tensor]:
+            """Find the tensor list for a source pattern that ends with pat_suffix."""
+            for key, tensors in input_dict.items():
+                if key.endswith(pat_suffix):
+                    return tensors
+            raise KeyError(
+                f"Nvfp4ExpertsDeserialize: could not find '{pat_suffix}' in "
+                f"input_dict keys: {list(input_dict.keys())}"
+            )
+
+        # spawn_materialize casts all checkpoint tensors to the skeleton dtype (bf16) before
+        # the converter sees them. uint8 qdata (0-255) and float8_e4m3fn scales both roundtrip
+        # exactly through bf16, so recast back to the raw dtypes NVFP4Tensor needs.
+        def _recast_weight(t: torch.Tensor) -> torch.Tensor:
+            if t.dtype != torch.uint8:
+                return t.to(torch.int32).to(torch.uint8)
+            return t
+
+        def _recast_scale(t: torch.Tensor) -> torch.Tensor:
+            if t.dtype != torch.float8_e4m3fn:
+                return t.to(torch.float8_e4m3fn)
+            return t
+
+        def _proj_parts(proj_name: str) -> dict:
+            return {
+                "qd": [_recast_weight(t) for t in _find(f"{proj_name}.weight")],
+                "sc": [_recast_scale(t) for t in _find(f"{proj_name}.weight_scale")],
+                "pts": list(_find(f"{proj_name}.weight_scale_2")),
+            }
+
+        # gate/up fuse on the N axis (each ships its own per-tensor scale, reconciled in the core);
+        # down is a single projection. Fusion + scale reconciliation live in fuse_nvfp4_experts.
+        if proj == "gate_up_proj":
+            projs = [_proj_parts("gate_proj"), _proj_parts("up_proj")]
+        else:
+            projs = [_proj_parts("down_proj")]
+        nvfp4 = fuse_nvfp4_experts(projs)
+
+        module, _ = get_module_from_name(model, full_layer_name)
+
+        setattr(module, proj, nn.Parameter(nvfp4, requires_grad=False))
+
+        if missing_keys is not None:
+            missing_keys.discard(full_layer_name)
+
+        module._is_hf_initialized = True
+
+        LOG.debug(
+            "Nvfp4ExpertsDeserialize: set %s as NVFP4Tensor [%s]",
+            full_layer_name,
+            list(nvfp4.shape),
+        )
+        return {}
+
+    # No meaningful reverse op (packed NVFP4 → checkpoint would need to unfuse).
+    @property
+    def reverse_op(self):
+        from transformers.core_model_loading import _IdentityOp
+
+        return _IdentityOp()
+
+
+class Nvfp4LinearDequantize:
+    """ConversionOps that dequantizes one NVFP4 ``nn.Linear`` weight to bf16 in place.
+
+    Consumes a single module's NVFP4 triple — ``weight`` (uint8 qdata ``[out, in/2]``),
+    ``weight_scale`` (e4m3 group-16 ``[out, in/16]``), ``weight_scale_2`` (per-tensor scalar) —
+    and assigns ``dequantize(NVFP4Tensor(...))`` to the target ``.weight`` param. Used for whatever
+    non-routed linears a given checkpoint quantizes (the caller decides which, from the index);
+    this op makes no assumption about which modules those are. ``input_dict`` holds the three
+    source tensors (each a 1-element list under a non-wildcard source pattern); ``full_layer_name``
+    is the target ``....weight``.
+    """
+
+    def convert(
+        self,
+        input_dict: dict[str, Any],
+        source_patterns: list[str] | None = None,
+        target_patterns: list[str] | None = None,
+        full_layer_name: str | None = None,
+        model: nn.Module | None = None,
+        missing_keys: set | None = None,
+        **kwargs,
+    ) -> dict[str, Any]:
+        from transformers.quantizers.quantizers_utils import get_module_from_name
+
+        NVFP4Tensor = _nvfp4_cls()
+        if NVFP4Tensor is None:
+            raise RuntimeError(
+                "torchao.prototype.mx_formats.nvfp4_tensor.NVFP4Tensor not found; "
+                "install torchao with NVFP4 support"
+            )
+
+        def _one(pat_suffix: str) -> torch.Tensor:
+            for key, val in input_dict.items():
+                if key.endswith(pat_suffix):
+                    return val[0] if isinstance(val, (list, tuple)) else val
+            raise KeyError(
+                f"Nvfp4LinearDequantize: could not find '{pat_suffix}' in "
+                f"input_dict keys: {list(input_dict.keys())}"
+            )
+
+        w = _one(".weight")
+
+        # cpu_ram_efficient_loading: non-local-rank-0 stays on meta (FSDP broadcasts rank 0's
+        # weights). qdata is packed [out, in/2] uint8 -> dequantized weight is [out, in] bf16.
+        if _nonrank0_meta_load():
+            weight = torch.empty(
+                (w.shape[0], w.shape[1] * 2), dtype=torch.bfloat16, device="meta"
+            )
+        else:
+            # spawn_materialize casts checkpoint tensors to the skeleton dtype (bf16); recast the
+            # raw uint8 qdata / e4m3 scale back (both roundtrip exactly through bf16).
+            qdata = w if w.dtype == torch.uint8 else w.to(torch.int32).to(torch.uint8)
+            sc = _one(".weight_scale")
+            scale = (
+                sc if sc.dtype == torch.float8_e4m3fn else sc.to(torch.float8_e4m3fn)
+            )
+            pts = _one(".weight_scale_2").to(torch.float32).view(1, 1)
+
+            weight = NVFP4Tensor(
+                qdata, scale, 16, torch.bfloat16, per_tensor_scale=pts
+            ).dequantize(torch.bfloat16)
+
+        module, param_name = get_module_from_name(model, full_layer_name)
+        setattr(module, param_name, nn.Parameter(weight, requires_grad=False))
+        if missing_keys is not None:
+            missing_keys.discard(full_layer_name)
+        module._is_hf_initialized = True
+        return {}
+
+    @property
+    def reverse_op(self):
+        from transformers.core_model_loading import _IdentityOp
+
+        return _IdentityOp()
+
+
+def _nvfp4_linear_dequant_converter(target_weight: str):
+    """A WeightConverter that dequantizes one NVFP4 linear (``<target_weight>``) to bf16.
+
+    ``target_weight`` is the distinctive suffix of the linear's ``.weight`` param (e.g.
+    ``"shared_experts.gate_proj.weight"``). The three source patterns claim the NVFP4 triple so
+    none lands UNEXPECTED; longest-suffix-first so ``.weight`` doesn't steal ``.weight_scale*``.
+    """
+    from transformers.core_model_loading import WeightConverter
+
+    base = target_weight[: -len(".weight")]
+    return WeightConverter(
+        source_patterns=[
+            f"{base}.weight_scale_2",
+            f"{base}.weight_scale",
+            f"{base}.weight",
+        ],
+        target_patterns=target_weight,
+        operations=[Nvfp4LinearDequantize()],
+    )
+
+
+def nonrouted_dequant_converters(nonrouted_suffixes: list[str]) -> list:
+    """Dequant converters for the non-routed NVFP4 linears a checkpoint actually quantizes.
+
+    ``nonrouted_suffixes`` are the layer-relative module paths detected from the safetensors
+    index (e.g. ``"mlp.shared_experts.gate_proj"``, ``"mlp.gate_proj"``) — NOT hardcoded — each
+    pointing at a ``.weight`` to dequantize to bf16. Empty list -> no converters (e.g. a
+    checkpoint that leaves all non-routed linears bf16)."""
+    return [
+        _nvfp4_linear_dequant_converter(f"{suf}.weight") for suf in nonrouted_suffixes
+    ]
+
+
+def nvfp4_experts_weight_converters() -> list:
+    """Return the two WeightConverter instances for gemma4 NVFP4 experts.
+
+    These are registered under ``"gemma4_text"`` in the transformers
+    conversion_mapping cache so the loader finds and applies them during
+    ``from_pretrained``.
+    """
+    from transformers.core_model_loading import WeightConverter
+
+    op = Nvfp4ExpertsDeserialize()
+
+    # Source patterns MUST be ordered longest-suffix-first. transformers compiles them into a
+    # single ``(?P<g0>...)|(?P<g1>...)`` alternation and resolves a key with ``re.search`` +
+    # first-non-None group (core_model_loading.py). The patterns are NOT end-anchored when the
+    # converter is many-to-one (the ^...$ anchoring only runs for equal-length source/target
+    # lists), so ``...weight`` would substring-match inside ``...weight_scale``/``...weight_scale_2``
+    # and steal those keys unless the more specific suffixes appear first.
+    gate_up_converter = WeightConverter(
+        source_patterns=[
+            # gate and up each ship their own weight_scale_2 scalar; claim BOTH so neither lands
+            # as an UNEXPECTED key (the op reconciles them — equal in practice, folded if not).
+            "experts.*.gate_proj.weight_scale_2",
+            "experts.*.up_proj.weight_scale_2",
+            "experts.*.gate_proj.weight_scale",
+            "experts.*.up_proj.weight_scale",
+            "experts.*.gate_proj.weight",
+            "experts.*.up_proj.weight",
+        ],
+        target_patterns="experts.gate_up_proj",
+        operations=[op],
+    )
+
+    down_converter = WeightConverter(
+        source_patterns=[
+            "experts.*.down_proj.weight_scale_2",
+            "experts.*.down_proj.weight_scale",
+            "experts.*.down_proj.weight",
+        ],
+        target_patterns="experts.down_proj",
+        operations=[op],
+    )
+
+    return [gate_up_converter, down_converter]
+
+
+def register_nvfp4_expert_converters(
+    model_type: str, include_routed: bool = True, extra: list | None = None
+) -> None:
+    """Seed the transformers conversion_mapping cache with NVFP4 converters for ``model_type``.
+
+    The routed-expert converters (``Nvfp4ExpertsDeserialize`` + the ``experts.*.{proj}.weight*``
+    source patterns) fuse per-expert ``gate/up/down`` into the model's 3D expert params; they are
+    registered when ``include_routed`` (gate this on the checkpoint actually exporting per-expert
+    NVFP4). ``extra`` carries any per-checkpoint non-routed dequant converters (built from the
+    detected index layout). The only per-model knob is which ``model_type`` the loader looks the
+    mapping up under. Safe to call repeatedly (idempotent via overwrite on re-entry).
+    """
+    from transformers.conversion_mapping import register_checkpoint_conversion_mapping
+
+    converters = (nvfp4_experts_weight_converters() if include_routed else []) + list(
+        extra or []
+    )
+    if not converters:
+        return
+    try:
+        register_checkpoint_conversion_mapping(model_type, converters)
+    except ValueError:
+        # Already registered; overwrite to keep converters fresh.
+        register_checkpoint_conversion_mapping(model_type, converters, overwrite=True)
+
+    LOG.info(
+        "Registered %s NVFP4 WeightConverters (%d) in transformers conversion_mapping",
+        model_type,
+        len(converters),
+    )
+
+
+def register_gemma4_nvfp4_converters() -> None:
+    """Register NVFP4 expert converters for ``gemma4_text`` (see register_nvfp4_expert_converters)."""
+    register_nvfp4_expert_converters("gemma4_text")
+
+
+def register_nvfp4_converters_for_layout(model_type: str, layout: dict) -> None:
+    """Register NVFP4 converters built from a detected checkpoint ``layout`` (see
+    :func:`...nvfp4_moe_loading.inspect_nvfp4_layout`): routed experts fused into packed
+    NVFP4Tensor (when present), and each detected non-routed NVFP4 linear dequantized to bf16."""
+    register_nvfp4_expert_converters(
+        model_type,
+        include_routed=layout.get("routed_present", False),
+        extra=nonrouted_dequant_converters(layout.get("nonrouted_suffixes", [])),
+    )
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/parallel_experts.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/parallel_experts.py
new file mode 100644
index 0000000000..af8eb7565f
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/parallel_experts.py
@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/shawntan/scattermoe
+# Copyright (c) Shawn Tan and ScatterMoE Contributors
+# Licensed under the Apache License, Version 2.0
+# See https://github.com/shawntan/scattermoe/blob/main/LICENSE
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from . import kernels
+
+# When the maximum addressable element offset across any input/output buffer
+# exceeds INT_MAX, the Triton kernel's int32 pointer-offset arithmetic
+# overflows. ``_needs_int64_indices`` returns True iff any tensor has
+# ``numel() >= INT_MAX``, which is a sufficient condition for the
+# ``M_idx * stride_*m`` product to overflow somewhere in the kernel. When
+# True, callers pass ``INT64_INDICES=True`` to the kernel so the index range
+# is cast to int64 before it enters the multiplication. Strides themselves
+# are already int64 at the Python level (from ``tensor.stride()``); only
+# the *index* type needs the bump.
+#
+# The threshold here matches the kernel's correctness boundary: as soon as
+# any indexed buffer has 2**31 - 1 or more elements, the int32 multiply at
+# the end of the buffer can overflow. The wrapper used to chunk the call to
+# keep ``rows * y_dim`` below 2**31; with INT64_INDICES the kernel itself
+# handles overflow, so the auto-dispatch routes directly to a single
+# kernel launch in either mode.
+_INT_MAX = 2**31 - 1
+
+
+def _needs_int64_indices(*tensors) -> bool:
+    """True iff any input/output tensor's element count exceeds INT_MAX."""
+    for t in tensors:
+        if t is None or not isinstance(t, torch.Tensor):
+            continue
+        if t.numel() >= _INT_MAX:
+            return True
+    return False
+
+
+@torch.library.custom_op("scattermoe::bincount", mutates_args={})
+def compileable_bincount(x: torch.Tensor, minlength: int) -> torch.Tensor:
+    return x.bincount(minlength=minlength)
+
+
+@compileable_bincount.register_fake
+def _(x: torch.Tensor, minlength: int) -> torch.Tensor:
+    return torch.empty(minlength, dtype=torch.long, device=x.device)
+
+
+@torch.compile
+def flatten_sort_count(expert_idxs: torch.Tensor, num_experts: int):
+    with torch.no_grad():
+        flattened_expert_idxs = expert_idxs.flatten()
+        sorted_expert_idxs, sorted_scattered_idxs = torch.sort(flattened_expert_idxs)
+        expert_counts = compileable_bincount(
+            flattened_expert_idxs, minlength=num_experts
+        )
+        expert_offsets = expert_counts.cumsum(-1)
+        return sorted_expert_idxs, sorted_scattered_idxs, expert_offsets
+
+
+class ParallelLinear(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        expert_weights: torch.Tensor,
+        k: int,
+        sorted_expert_idxs: torch.Tensor,
+        sorted_scattered_idxs: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        expert_biases: Optional[torch.Tensor] = None,
+        gates: Optional[torch.Tensor] = None,
+        grouped_in: bool = False,
+        grouped_out: bool = False,
+    ):
+        # Cast weights to match input dtype (e.g. 8-bit LoRA)
+        if expert_weights.dtype != x.dtype:
+            expert_weights = expert_weights.to(x.dtype)
+        if expert_biases is not None and expert_biases.dtype != x.dtype:
+            expert_biases = expert_biases.to(x.dtype)
+        L_scattered = sorted_expert_idxs.size(0)
+        y_dim = expert_weights.size(-1)
+        # Cheap probe: the kernel's overflow risk is the M_block * stride_ym
+        # product, dominated by the output buffer L_scattered * y_dim. We also
+        # check x because the X_ptr arithmetic uses similar indices.
+        needs_int64_fwd = (L_scattered * y_dim) >= _INT_MAX or _needs_int64_indices(x)
+        with torch.device(x.device):
+            output = kernels.ops.scatter2scatter(
+                X=x,
+                W=expert_weights,
+                b=expert_biases,
+                k=k,
+                sorted_expert_idxs=sorted_expert_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                x_grouped=grouped_in,
+                y_grouped=grouped_out,
+                int64_indices=needs_int64_fwd,
+            )
+            if gates is not None:
+                output_expanded = output.view(
+                    gates.size(0), gates.size(1), output.size(-1)
+                )
+                output = (gates.unsqueeze(1) @ output_expanded).squeeze(1)
+            else:
+                output_expanded = None
+
+            ctx.save_for_backward(
+                x,
+                expert_weights,
+                expert_biases,
+                sorted_expert_idxs,
+                sorted_scattered_idxs,
+                expert_offsets,
+                gates,
+                output_expanded,
+            )
+            ctx.grouped_in = grouped_in
+            ctx.grouped_out = grouped_out
+            ctx.k = k
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):
+        with torch.device(grad_out.device):
+            (
+                x,
+                expert_weights,
+                expert_biases,
+                sorted_expert_idxs,
+                sorted_scattered_idxs,
+                expert_offsets,
+                gates,
+                output_expanded,
+            ) = ctx.saved_tensors
+            k = ctx.k
+            grouped_in = ctx.grouped_in
+            grouped_out = ctx.grouped_out
+
+            if gates is not None:
+                # grad_out arrives in the autocast dtype but output_expanded stays fp32; align before the d_gates bmm
+                if grad_out.dtype != output_expanded.dtype:
+                    grad_out = grad_out.to(output_expanded.dtype)
+                # calculate gates gradient
+                # d_gates = torch.bmm(output_expanded, grad_out[:, :, None]).squeeze(-1)
+                d_gates = (output_expanded @ grad_out.unsqueeze(-1)).squeeze(-1)
+                gates_flat = gates.flatten()
+                gate_fan = gates.size(1)
+                grouped_grad_out = output_expanded.flatten(
+                    0, 1
+                )  # reuse expanded buffer later
+            else:
+                d_gates = None
+                gates_flat = None
+                gate_fan = 1
+                grouped_grad_out = None
+
+            if grouped_out:
+                grouped_grad_out = grad_out
+            else:
+                grouped_grad_out = kernels.ops.group(
+                    grad_out,
+                    sorted_scattered_idxs,
+                    fan_out=gate_fan,
+                    coeff=gates_flat,
+                    out=grouped_grad_out,
+                )
+            if grouped_in:
+                grouped_x = x
+                d_expanded_input = None
+            else:
+                grouped_x = kernels.ops.group(x, sorted_scattered_idxs, fan_out=k)
+                d_expanded_input = grouped_x
+
+            d_weights, d_biases = kernels.ops.group_bwd_W(
+                DY=grouped_grad_out,
+                X=grouped_x,
+                expert_offsets=expert_offsets,
+                E=expert_weights.size(0),
+                has_bias=expert_biases is not None,
+            )
+
+            L_scattered = sorted_expert_idxs.size(0)
+            dx_dim = expert_weights.size(1)  # K dim of W = output dim for dX
+            needs_int64_bwd = (
+                L_scattered * dx_dim
+            ) >= _INT_MAX or _needs_int64_indices(grouped_grad_out)
+            d_expanded_input = kernels.ops.scatter2scatter(
+                X=grouped_grad_out,
+                x_grouped=True,
+                W=expert_weights.permute(0, 2, 1),
+                sorted_expert_idxs=sorted_expert_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                k=1,
+                y_grouped=grouped_in,
+                out=d_expanded_input,  # Reuse grouped_x buffer
+                int64_indices=needs_int64_bwd,
+            )
+
+            if k == 1:
+                d_input = d_expanded_input
+            else:
+                d_input = d_expanded_input.view(
+                    x.size(0), k, d_expanded_input.size(-1)
+                ).sum(-2)
+        return (
+            # x, expert_weights,
+            d_input,
+            d_weights,
+            # k, sorted_expert_idxs, sorted_scattered_idxs, expert_offsets,
+            None,
+            None,
+            None,
+            None,
+            # bias, gates
+            d_biases,
+            d_gates,
+            # grouped_in, grouped_out,
+            None,
+            None,
+        )
+
+
+def parallel_linear(
+    inputs,
+    expert_weights,
+    k,
+    sorted_expert_idxs,
+    sorted_scattered_idxs,
+    expert_offsets,
+    expert_biases=None,
+    gates=None,
+    grouped_in=False,
+    grouped_out=False,
+):
+    results = ParallelLinear.apply(
+        inputs,
+        expert_weights,
+        k,
+        sorted_expert_idxs,
+        sorted_scattered_idxs,
+        expert_offsets,
+        expert_biases,
+        gates,
+        grouped_in,
+        grouped_out,
+    )
+    return results
+
+
+class ParallelExperts(nn.Module):
+    def __init__(self, num_experts, input_size, output_size, bias=False) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(num_experts, output_size, input_size))
+
+        if bias:
+            self.bias = nn.Parameter(torch.empty(num_experts, output_size))
+        else:
+            self.bias = None
+
+        self.num_experts = num_experts
+        self.input_size = input_size
+        self.output_size = output_size
+        self.reset_parameters()
+
+    def extra_repr(self):
+        return "num_experts={}, input_size={}, output_size={}".format(
+            self.num_experts, self.input_size, self.output_size
+        )
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.weight, std=0.02)
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def forward(
+        self,
+        inputs,
+        k,
+        sorted_expert_idxs,
+        sorted_scattered_idxs,
+        expert_offsets,
+        gates=None,
+        grouped_in=False,
+        grouped_out=False,
+    ):
+        results = parallel_linear(
+            inputs,
+            self.weight.permute(0, 2, 1),
+            k,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            expert_biases=self.bias,
+            gates=gates,
+            grouped_in=grouped_in,
+            grouped_out=grouped_out,
+        )
+        return results
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/parallel_linear_lora.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/parallel_linear_lora.py
new file mode 100644
index 0000000000..829b85db5c
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/parallel_linear_lora.py
@@ -0,0 +1,624 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+ScatterMoE + LoRA Autograd Function
+====================================
+
+Provides the autograd function and Python interface for fused ScatterMoE + LoRA.
+
+Key design for LoRA training:
+  - Expert weights W are FROZEN (no gradient computed for W).
+  - Only LoRA adapter weights (A, B) receive gradients.
+  - The input gradient dX is still computed (needed for upstream layers).
+  - This avoids the expensive group_bwd_W computation entirely.
+
+Forward:
+  Y = X @ W + scaling * (X @ A^T) @ B^T
+
+Backward (W frozen):
+  dX = dY @ W^T + scaling * (dY @ B) @ A          (via scatter2scatter for base, separate for LoRA)
+  dA = scaling * (dY @ B)^T @ X                     (per-expert, on grouped data)
+  dB = scaling * dY^T @ (X @ A^T)                   (per-expert, on grouped data)
+"""
+
+from typing import Optional, Union
+
+import torch
+
+from .kernels import ops as base_ops
+from .kernels.grouped_gram import grouped_lora_weight_grads
+from .kernels.lora_ops import (
+    group_bwd_lora_fused,
+    scatter2scatter_lora,
+    scatter2scatter_lora_dX,
+    scatter2scatter_lora_dX_mx,
+    scatter2scatter_lora_mx,
+)
+from .mx_weights import MXLayout, MXWeights
+from .parallel_experts import _INT_MAX, _needs_int64_indices
+
+
+class ScatterMoELoRA(torch.autograd.Function):
+    """
+    Autograd function for fused ScatterMoE + LoRA with frozen expert weights.
+
+    This function is optimized for the LoRA fine-tuning scenario where:
+    - Expert weights W are frozen (requires_grad=False)
+    - Only LoRA A and B matrices receive gradients
+    - Input gradients are computed for upstream layer backprop
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        expert_weights: Union[torch.Tensor, MXWeights],
+        k: int,
+        sorted_expert_idxs: torch.Tensor,
+        sorted_scattered_idxs: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        lora_A: torch.Tensor,
+        lora_B: torch.Tensor,
+        scaling: float,
+        expert_biases: Optional[torch.Tensor] = None,
+        gates: Optional[torch.Tensor] = None,
+        grouped_in: bool = False,
+        grouped_out: bool = False,
+        use_fused_dX: bool = False,
+        use_fused_gather: bool = False,
+    ):
+        if isinstance(expert_weights, MXWeights):
+            assert expert_weights.layout == MXLayout.FWD, (
+                "MXWeights passed to forward must be in FWD layout"
+            )
+            is_mx = True
+        else:
+            # match input dtype (e.g. 8-bit LoRA)
+            if expert_weights.dtype != x.dtype:
+                expert_weights = expert_weights.to(x.dtype)
+            is_mx = False
+        if expert_biases is not None and expert_biases.dtype != x.dtype:
+            expert_biases = expert_biases.to(x.dtype)
+        L_scattered = sorted_expert_idxs.size(0)
+        if is_mx:
+            N_dim = expert_weights.N  # type: ignore[union-attr]
+        else:
+            N_dim = expert_weights.size(-1)  # type: ignore[union-attr]
+        # Overflow risk is dominated by the [L_scattered, N] output buffer; also probe X for the
+        # rare case where it alone is huge (very wide hidden, modest seq).
+        needs_int64_fwd = (L_scattered * N_dim) >= _INT_MAX or _needs_int64_indices(x)
+        with torch.device(x.device):
+            if is_mx:
+                # MXFP4: dequant happens inside the K-loop
+                output = scatter2scatter_lora_mx(
+                    X=x,
+                    W_mx=expert_weights,
+                    sorted_expert_idxs=sorted_expert_idxs,
+                    sorted_scattered_idxs=sorted_scattered_idxs,
+                    k=k,
+                    lora_A=lora_A,
+                    lora_B=lora_B,
+                    scaling=scaling,
+                    b=expert_biases,
+                    x_grouped=grouped_in,
+                    y_grouped=grouped_out,
+                    int64_indices=needs_int64_fwd,
+                )
+            else:
+                output = scatter2scatter_lora(
+                    X=x,
+                    W=expert_weights,
+                    sorted_expert_idxs=sorted_expert_idxs,
+                    sorted_scattered_idxs=sorted_scattered_idxs,
+                    k=k,
+                    lora_A=lora_A,
+                    lora_B=lora_B,
+                    scaling=scaling,
+                    b=expert_biases,
+                    x_grouped=grouped_in,
+                    y_grouped=grouped_out,
+                    int64_indices=needs_int64_fwd,
+                )
+
+            # gating: weighted combination of top-k expert outputs
+            if gates is not None:
+                output_expanded = output.view(
+                    gates.size(0), gates.size(1), output.size(-1)
+                )
+                output = (gates.unsqueeze(1) @ output_expanded).squeeze(1)
+            else:
+                output_expanded = None
+
+            ctx.save_for_backward(
+                x,
+                lora_A,
+                lora_B,
+                sorted_expert_idxs,
+                sorted_scattered_idxs,
+                expert_offsets,
+                gates,
+                output_expanded,
+            )
+            # Frozen weights as plain ctx attributes, not save_for_backward: avoids version-check
+            # conflicts with FSDP unshard/reshard, pinning all-gathered params via saved_tensors
+            # hooks, and interfering with activation-offload pack/unpack hooks. Safe since the
+            # weights are frozen. If the caller attached a recompute recipe (selective MX/NVFP4
+            # gather, a per-layer copy rebuildable from the resident param), store the recipe and
+            # DROP the heavy copy so it frees on return; backward rebuilds it. Else keep the
+            # reference (the bf16 path passes a cheap param view).
+            ctx.weight_recipe = getattr(expert_weights, "recipe", None)
+            ctx.expert_weights = (
+                None if ctx.weight_recipe is not None else expert_weights
+            )
+            ctx.expert_biases = expert_biases
+            ctx.grouped_in = grouped_in
+            ctx.grouped_out = grouped_out
+            ctx.k = k
+            ctx.scaling = scaling
+            # MXFP4 forces fused dX + gather: the non-fused dX path would materialize a bf16 weight
+            # tile, defeating the fusion win, and the gather/scatter pattern is identical.
+            ctx.use_fused_dX = True if is_mx else use_fused_dX
+            ctx.use_fused_gather = True if is_mx else use_fused_gather
+            ctx.is_mx = is_mx
+
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):
+        with torch.device(grad_out.device):
+            (
+                x,
+                lora_A,
+                lora_B,
+                sorted_expert_idxs,
+                sorted_scattered_idxs,
+                expert_offsets,
+                gates,
+                output_expanded,
+            ) = ctx.saved_tensors
+            # Rebuild the selective MX/NVFP4 weights from the recipe instead of a copy pinned since
+            # forward (see forward for rationale).
+            expert_weights = (
+                ctx.expert_weights
+                if ctx.expert_weights is not None
+                else ctx.weight_recipe()
+            )
+
+            k = ctx.k
+            scaling = ctx.scaling
+            grouped_in = ctx.grouped_in
+            grouped_out = ctx.grouped_out
+            is_mx = ctx.is_mx
+            if is_mx:
+                E = expert_weights.packed.size(0)
+            else:
+                E = expert_weights.size(0)
+
+            # Gate gradients (top-k gating with routing weights)
+            if gates is not None:
+                # grad_out arrives in the autocast dtype but output_expanded stays fp32; align before the d_gates bmm
+                if grad_out.dtype != output_expanded.dtype:
+                    grad_out = grad_out.to(output_expanded.dtype)
+                # d_gates[t, j] = output_expanded[t, j, :] . grad_out[t, :]
+                d_gates = (output_expanded @ grad_out.unsqueeze(-1)).squeeze(-1)
+                gates_flat = gates.flatten()
+                gate_fan = gates.size(1)
+                grouped_grad_out = output_expanded.flatten(0, 1)
+            else:
+                d_gates = None
+                gates_flat = None
+                gate_fan = 1
+                grouped_grad_out = None
+
+            # Fused gather uses sorted_scattered_idxs for indirect X access in the Triton kernel,
+            # avoiding the group(x) allocation. Enabled when X is ungrouped and not too large for
+            # scatter loads (grouped_out=True still works via dy_grouped=True in the kernel).
+            M_total = sorted_scattered_idxs.size(0)
+            K_dim = x.size(-1)
+            N_dim = expert_weights.N if is_mx else expert_weights.size(-1)
+            fuse_gather_workload = M_total * max(K_dim, N_dim)
+            _FUSE_GATHER_THRESHOLD = 2**24  # ~16M elements
+
+            can_fuse_gather = (
+                ctx.use_fused_gather
+                and not grouped_in  # X must be ungrouped for scatter access
+                and gates is None  # gate coeff requires multiplicative gather
+                and fuse_gather_workload < _FUSE_GATHER_THRESHOLD
+            )
+
+            # Overflow risk is dominated by the largest M-axis indexed buffer (grad_out [M_total, N],
+            # x [M, K]).
+            needs_int64_bwd = (
+                (M_total * N_dim) >= _INT_MAX
+                or (M_total * K_dim) >= _INT_MAX
+                or _needs_int64_indices(grad_out, x)
+            )
+
+            yb = None  # dY @ B, computed by the non-fused dA/dB path; reused by dX_lora
+            if can_fuse_gather:
+                # Fused path: skip group(x) entirely
+                d_expanded_input = None
+
+                d_lora_A, d_lora_B = group_bwd_lora_fused(
+                    DY=grad_out,
+                    X=x,
+                    lora_A=lora_A,
+                    lora_B=lora_B,
+                    expert_offsets=expert_offsets,
+                    sorted_scattered_idxs=sorted_scattered_idxs,
+                    E=E,
+                    k=k,
+                    scaling=scaling,
+                    dy_grouped=grouped_out,
+                    int64_indices=needs_int64_bwd,
+                )
+
+                # grouped_grad_out for the dX path
+                if grouped_out:
+                    grouped_grad_out = grad_out
+                elif not ctx.use_fused_dX:
+                    grouped_grad_out = base_ops.group(
+                        grad_out,
+                        sorted_scattered_idxs,
+                        fan_out=gate_fan,
+                        coeff=gates_flat,
+                        out=grouped_grad_out,
+                    )
+            else:
+                # Original path: explicit group() calls
+                if grouped_out:
+                    grouped_grad_out = grad_out
+                else:
+                    grouped_grad_out = base_ops.group(
+                        grad_out,
+                        sorted_scattered_idxs,
+                        fan_out=gate_fan,
+                        coeff=gates_flat,
+                        out=grouped_grad_out,
+                    )
+
+                if grouped_in:
+                    grouped_x = x
+                    d_expanded_input = None
+                else:
+                    grouped_x = base_ops.group(x, sorted_scattered_idxs, fan_out=k)
+                    d_expanded_input = grouped_x  # overwritten; reuse buffer
+
+                # dA/dB via grouped-Gram over precomputed XA/YB (rank-sized) instead of the split
+                # kernel's per-output-block recompute; a large win as E grows (E >= 128). YB is
+                # reused by the non-fused dX path below.
+                rank = lora_A.size(0) // E
+                k_dim = lora_A.size(1)
+                n_dim = lora_B.size(0)
+                w_yb = lora_B.reshape(n_dim, E, rank).permute(1, 0, 2).contiguous()
+                yb = base_ops.scatter2scatter(
+                    X=grouped_grad_out,
+                    W=w_yb,
+                    k=1,
+                    sorted_expert_idxs=sorted_expert_idxs,
+                    sorted_scattered_idxs=sorted_scattered_idxs,
+                    x_grouped=True,
+                    y_grouped=True,
+                    int64_indices=needs_int64_bwd,
+                )
+                w_xa = lora_A.reshape(E, rank, k_dim).permute(0, 2, 1).contiguous()
+                xa = base_ops.scatter2scatter(
+                    X=grouped_x,
+                    W=w_xa,
+                    k=1,
+                    sorted_expert_idxs=sorted_expert_idxs,
+                    sorted_scattered_idxs=sorted_scattered_idxs,
+                    x_grouped=True,
+                    y_grouped=True,
+                    int64_indices=needs_int64_bwd,
+                )
+                d_lora_A, d_lora_B = grouped_lora_weight_grads(
+                    grouped_grad_out,
+                    grouped_x,
+                    yb,
+                    xa,
+                    lora_A,
+                    lora_B,
+                    expert_offsets,
+                    E,
+                    scaling,
+                )
+
+            # Input gradient: dX = dY @ W^T + scaling * (dY @ B) @ A
+            if is_mx:
+                # dX kernel reuses the forward MX layout (block axis = K), no pre-transpose/requant.
+                if can_fuse_gather and not grouped_out:
+                    d_expanded_input = scatter2scatter_lora_dX_mx(
+                        DY=grad_out,
+                        W_mx=expert_weights,
+                        sorted_expert_idxs=sorted_expert_idxs,
+                        sorted_scattered_idxs=sorted_scattered_idxs,
+                        k=1,
+                        lora_A=lora_A,
+                        lora_B=lora_B,
+                        scaling=scaling,
+                        dy_grouped=False,
+                        dx_grouped=grouped_in,
+                        out=d_expanded_input,
+                        int64_indices=needs_int64_bwd,
+                    )
+                else:
+                    d_expanded_input = scatter2scatter_lora_dX_mx(
+                        DY=grouped_grad_out,
+                        W_mx=expert_weights,
+                        sorted_expert_idxs=sorted_expert_idxs,
+                        sorted_scattered_idxs=sorted_scattered_idxs,
+                        k=1,
+                        lora_A=lora_A,
+                        lora_B=lora_B,
+                        scaling=scaling,
+                        dy_grouped=True,
+                        dx_grouped=grouped_in,
+                        out=d_expanded_input,
+                        int64_indices=needs_int64_bwd,
+                    )
+            elif ctx.use_fused_dX:
+                if can_fuse_gather and not grouped_out:
+                    # Fully fused: read ungrouped DY via scatter pattern
+                    d_expanded_input = scatter2scatter_lora_dX(
+                        DY=grad_out,
+                        W=expert_weights,
+                        sorted_expert_idxs=sorted_expert_idxs,
+                        sorted_scattered_idxs=sorted_scattered_idxs,
+                        k=1,
+                        lora_A=lora_A,
+                        lora_B=lora_B,
+                        scaling=scaling,
+                        dy_grouped=False,
+                        dx_grouped=grouped_in,
+                        out=d_expanded_input,
+                        int64_indices=needs_int64_bwd,
+                    )
+                else:
+                    # Fused dX only: read from pre-grouped DY
+                    d_expanded_input = scatter2scatter_lora_dX(
+                        DY=grouped_grad_out,
+                        W=expert_weights,
+                        sorted_expert_idxs=sorted_expert_idxs,
+                        sorted_scattered_idxs=sorted_scattered_idxs,
+                        k=1,
+                        lora_A=lora_A,
+                        lora_B=lora_B,
+                        scaling=scaling,
+                        dy_grouped=True,
+                        dx_grouped=grouped_in,
+                        out=d_expanded_input,
+                        int64_indices=needs_int64_bwd,
+                    )
+            else:
+                # Original path: separate base scatter2scatter + LoRA Python loop
+                d_expanded_input = base_ops.scatter2scatter(
+                    X=grouped_grad_out,
+                    x_grouped=True,
+                    W=expert_weights.permute(0, 2, 1),
+                    sorted_expert_idxs=sorted_expert_idxs,
+                    sorted_scattered_idxs=sorted_scattered_idxs,
+                    k=1,
+                    y_grouped=grouped_in,
+                    out=d_expanded_input,
+                    int64_indices=needs_int64_bwd,
+                )
+
+                # dX_lora = scaling * (dY @ B) @ A (sync-free grouped GEMMs; reuses YB from dA/dB)
+                if scaling != 0.0:
+                    d_input_lora_grouped = _compute_lora_input_grad(
+                        grouped_grad_out,
+                        lora_A,
+                        lora_B,
+                        expert_offsets,
+                        E,
+                        scaling,
+                        sorted_expert_idxs=sorted_expert_idxs,
+                        sorted_scattered_idxs=sorted_scattered_idxs,
+                        int64_indices=needs_int64_bwd,
+                        yb=yb,
+                    )
+                    if grouped_in:
+                        d_expanded_input.add_(d_input_lora_grouped)
+                    else:
+                        # scatter-add directly into d_expanded_input (avoids a zeros_like + add)
+                        d_expanded_input[sorted_scattered_idxs] += d_input_lora_grouped
+
+            if k == 1:
+                d_input = d_expanded_input
+            else:
+                d_input = d_expanded_input.view(
+                    x.size(0), k, d_expanded_input.size(-1)
+                ).sum(-2)
+
+            # W is frozen during LoRA training, skip weight gradient.
+            # (MX weights are containers, not tensors, and never carry grad.)
+            if is_mx:
+                d_weights = None
+            else:
+                d_weights = (
+                    torch.zeros_like(expert_weights)
+                    if expert_weights.requires_grad
+                    else None
+                )
+            d_biases = None
+
+        return (
+            d_input,
+            d_weights,
+            None,
+            None,
+            None,
+            None,  # k, sorted indices, offsets
+            d_lora_A,
+            d_lora_B,
+            None,  # lora_A, lora_B, scaling
+            d_biases,
+            d_gates,
+            None,
+            None,  # grouped_in, grouped_out
+            None,  # use_fused_dX
+            None,  # use_fused_gather
+        )
+
+
+def _compute_lora_input_grad(
+    grouped_grad_out: torch.Tensor,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    E: int,
+    scaling: float,
+    sorted_expert_idxs: Optional[torch.Tensor] = None,
+    sorted_scattered_idxs: Optional[torch.Tensor] = None,
+    int64_indices: bool = False,
+    yb: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """LoRA contribution to the input gradient: ``dX_lora = scaling * (dY @ B) @ A``,
+    on expert-grouped data.
+
+    With routing ids it runs as two grouped GEMMs (``scatter2scatter``) -- sync-free
+    and a single launch each, instead of a Python per-expert loop with an
+    ``expert_offsets[e].item()`` device sync per expert (O(E) syncs, which dominate
+    at the high expert counts of modern MoEs). ``yb = dY @ B`` may be passed in to
+    reuse the value already computed for the dA/dB grads. The per-expert loop is kept
+    as a fallback for callers without the routing ids.
+    """
+    R = lora_A.size(0) // E
+    K = lora_A.size(1)
+    N = lora_B.size(0)
+
+    if sorted_expert_idxs is not None:
+        if yb is None:
+            w_yb = lora_B.reshape(N, E, R).permute(1, 0, 2).contiguous()
+            yb = base_ops.scatter2scatter(
+                X=grouped_grad_out,
+                W=w_yb,
+                k=1,
+                sorted_expert_idxs=sorted_expert_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                x_grouped=True,
+                y_grouped=True,
+                int64_indices=int64_indices,
+            )
+        w_a = lora_A.reshape(E, R, K).contiguous()
+        dx = base_ops.scatter2scatter(
+            X=yb,
+            W=w_a,
+            k=1,
+            sorted_expert_idxs=sorted_expert_idxs,
+            sorted_scattered_idxs=sorted_scattered_idxs,
+            x_grouped=True,
+            y_grouped=True,
+            int64_indices=int64_indices,
+        )
+        return dx.mul_(scaling)
+
+    # fallback (no routing ids): one host sync for the whole offset array, not per expert
+    offsets = expert_offsets.tolist()
+    compute_dtype = grouped_grad_out.dtype
+    d_input_lora = torch.zeros(
+        (grouped_grad_out.size(0), K),
+        device=grouped_grad_out.device,
+        dtype=compute_dtype,
+    )
+    prev_offset = 0
+    for e in range(E):
+        curr_offset = offsets[e]
+        if curr_offset > prev_offset:
+            dy_e = grouped_grad_out[prev_offset:curr_offset]
+            a_e = lora_A[e * R : (e + 1) * R, :].to(compute_dtype)
+            b_e = lora_B[:, e * R : (e + 1) * R].to(compute_dtype)
+            d_input_lora[prev_offset:curr_offset] = scaling * ((dy_e @ b_e) @ a_e)
+        prev_offset = curr_offset
+    return d_input_lora
+
+
+def get_lora_params_from_wrapper(module) -> tuple:
+    """
+    Extract LoRA parameters from a PEFT ParamWrapper.
+
+    Returns:
+        (lora_A, lora_B, scaling) if LoRA is active, else (None, None, None)
+    """
+    if not hasattr(module, "lora_A") or not hasattr(module, "lora_B"):
+        return None, None, None
+
+    active_adapters = getattr(module, "active_adapters", ["default"])
+    if not active_adapters:
+        return None, None, None
+
+    adapter_name = active_adapters[0]
+
+    lora_A_dict = getattr(module, "lora_A", {})
+    lora_B_dict = getattr(module, "lora_B", {})
+    scaling_dict = getattr(module, "scaling", {})
+
+    if adapter_name not in lora_A_dict:
+        return None, None, None
+
+    lora_A = lora_A_dict[adapter_name].weight
+    lora_B = lora_B_dict[adapter_name].weight
+    scaling = scaling_dict[adapter_name]
+
+    return lora_A, lora_B, scaling
+
+
+def parallel_linear_lora(
+    inputs: torch.Tensor,
+    expert_weights: torch.Tensor,
+    k: int,
+    sorted_expert_idxs: torch.Tensor,
+    sorted_scattered_idxs: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    lora_A: Optional[torch.Tensor] = None,
+    lora_B: Optional[torch.Tensor] = None,
+    scaling: float = 1.0,
+    expert_biases: Optional[torch.Tensor] = None,
+    gates: Optional[torch.Tensor] = None,
+    grouped_in: bool = False,
+    grouped_out: bool = False,
+    use_fused_dX: bool = False,
+    use_fused_gather: bool = False,
+):
+    """
+    Drop-in replacement for parallel_linear that supports LoRA.
+
+    If lora_A and lora_B are provided, uses fused LoRA kernel.
+    Otherwise falls back to standard scatter2scatter.
+    """
+    if lora_A is not None and lora_B is not None:
+        return ScatterMoELoRA.apply(
+            inputs,
+            expert_weights,
+            k,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            lora_A,
+            lora_B,
+            scaling,
+            expert_biases,
+            gates,
+            grouped_in,
+            grouped_out,
+            use_fused_dX,
+            use_fused_gather,
+        )
+    else:
+        from .parallel_experts import ParallelLinear
+
+        return ParallelLinear.apply(
+            inputs,
+            expert_weights,
+            k,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            expert_biases,
+            gates,
+            grouped_in,
+            grouped_out,
+        )
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/runtime.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/runtime.py
new file mode 100644
index 0000000000..74ee6ac60b
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/runtime.py
@@ -0,0 +1,69 @@
+"""Centralized mutable process-wide ScatterMoE runtime settings.
+
+One place to see (and reset) everything that changes ScatterMoE behavior at runtime. It is applied
+once per run from the resolved config by :func:`configure_scattermoe_runtime`
+(``KernelsPlugin.pre_model_load``), which resets first so a long-lived multi-run process can't
+inherit stale state from a previous config.
+
+The per-module ``set_*`` functions in experts.py / grouped_train.py / chunked_bnb.py remain as thin
+compatibility wrappers that delegate to :data:`RUNTIME` (kept for the existing call sites + tests);
+the kernel code reads ``RUNTIME`` fields directly.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+
+# chunk size for the chunked bnb-4bit grouped path; 32 balances throughput on small-expert MoEs
+# against bounding the bf16 transient on large-expert models.
+DEFAULT_CHUNK = 32
+
+
+@dataclass
+class ScatterMoERuntime:
+    """All mutable knobs that affect ScatterMoE runtime behavior."""
+
+    # grouped fp4 MoE experts (experts.py): None = off (fused/eager paths unchanged) | "nvfp4".
+    fp4_grouped_mode: str | None = None
+    # base dX backward precision on sm120: True = fp8-read (fast, ~2% grad) | False = bf16-dequant.
+    fp4_dx_prefer_fp8: bool = True
+    # grouped base-GEMM backend (grouped_train.py): None/"auto" | marlin | cutlass | deepgemm | dequant.
+    grouped_backend: str | None = None
+    # bnb-4bit experts (chunked_bnb.py): 1-launch parallel_linear path (recompute-in-backward) vs chunked.
+    bnb_fast: bool = True
+    # chunked-dequant chunk size override; None -> DEFAULT_CHUNK.
+    dequant_chunk_size: int | None = None
+    # whether layer-level gradient checkpointing is active (skips the redundant per-chunk checkpoint).
+    layer_gc_active: bool = False
+    # persist the requantized mxfp4 weight in a module-level cache across steps. Safe on a persistent
+    # single-device param, but under FSDP2 the gathered param is the FULL weight every step, so
+    # caching it holds a full-model mxfp4 copy per rank (OOM). Disabled under FSDP: recompute per
+    # forward (freed after each layer), bounding resident mxfp4 to one layer.
+    mxfp4_cache_persist: bool = True
+
+    def reset(self) -> None:
+        for f in fields(self):
+            setattr(self, f.name, f.default)
+
+
+RUNTIME = ScatterMoERuntime()
+
+
+def configure_scattermoe_runtime(cfg) -> None:
+    """Apply ALL ScatterMoE runtime settings from a run config.
+
+    Resets to defaults first so a long-lived process never inherits stale state, then maps the
+    relevant config fields. ``cfg`` is the resolved config (DictDefault).
+    """
+    RUNTIME.reset()
+    RUNTIME.fp4_grouped_mode = cfg.get("dsv4_fp4_grouped_mode")
+    backend = cfg.get("moe_grouped_backend")
+    RUNTIME.grouped_backend = str(backend).lower() if backend else None
+    chunk = cfg.get("moe_dequant_chunk_size")
+    RUNTIME.dequant_chunk_size = int(chunk) if chunk else None
+    RUNTIME.layer_gc_active = bool(cfg.get("gradient_checkpointing"))
+    fast = cfg.get("moe_bnb_fast")
+    RUNTIME.bnb_fast = True if fast is None else bool(fast)
+    # Cache only in the non-FSDP case; under FSDP a persistent mxfp4 cache holds a full-model copy
+    # per rank (OOM).
+    RUNTIME.mxfp4_cache_persist = not bool(cfg.get("fsdp_config"))
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/selective_dequant.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/selective_dequant.py
new file mode 100644
index 0000000000..bd44a67d3b
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/selective_dequant.py
@@ -0,0 +1,427 @@
+"""
+Selective Expert Dequantization
+===============================
+
+Instead of dequantizing all E expert weight matrices at once (which creates
+a ~1 GB transient buffer for 256 experts), only dequantize the experts that
+are actually routed to by the current batch's top-k selection.
+
+For Qwen3.5-35B-A3B (E=256, top_k=8, hidden=2048, intermediate=512):
+  - Full dequant: [256, 2048, 1024] = 1,074 MB per projection
+  - Selective (8 active): [8, 2048, 1024] = 33.5 MB per projection
+  - Savings: ~97% memory reduction per layer
+
+This module provides format-agnostic selective weight extraction:
+  - BnB 4-bit (nf4/fp4): slice quantized data + absmax per expert
+  - MXFP4 (torchao MXTensor with elem_dtype=float4_e2m1fn_x2): slice
+    qdata + E8M0 scale per expert and dequantize via torchao
+  - bf16/fp32: direct indexing (no dequant needed)
+  - FP8: slice + cast
+
+The ScatterMoE kernel itself doesn't change — we remap expert indices
+from global (0..E-1) to compact (0..num_active-1) and pass the smaller
+weight tensor.
+"""
+
+import torch
+import torch.nn as nn
+
+from .mx_weights import (
+    _construct_mxtensor_subset,
+    _mx_qdata,
+    _mx_scale,
+    _torchao_mxtensor_cls,
+    _torchao_nvfp4tensor_cls,
+)
+
+
+def is_mxfp4_param(param) -> bool:
+    """True iff ``param`` is a torchao MXTensor with MXFP4 element dtype."""
+    MXTensor = _torchao_mxtensor_cls()
+    if MXTensor is None or not isinstance(param, MXTensor):
+        return False
+    return param.elem_dtype == torch.float4_e2m1fn_x2
+
+
+def is_nvfp4_param(param) -> bool:
+    """True iff ``param`` is a torchao NVFP4Tensor (FP4 E2M1, block-16 E4M3 scales)."""
+    NVFP4Tensor = _torchao_nvfp4tensor_cls()
+    return NVFP4Tensor is not None and isinstance(param, NVFP4Tensor)
+
+
+def get_active_experts(sorted_expert_idxs: torch.Tensor, E: int) -> torch.Tensor:
+    """Get sorted unique expert indices from the routing output.
+
+    Args:
+        sorted_expert_idxs: Expert assignments sorted by expert id [T*k]
+        E: Total number of experts
+
+    Returns:
+        active: Sorted unique expert indices [num_active]
+    """
+    return torch.unique(sorted_expert_idxs)
+
+
+def remap_expert_indices(
+    sorted_expert_idxs: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    active_experts: torch.Tensor,
+    E: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Remap global expert indices to compact indices.
+
+    Maps expert ids from [0..E-1] to [0..num_active-1], preserving the
+    sort order. Also compacts expert_offsets to only active experts.
+
+    Args:
+        sorted_expert_idxs: [T*k] expert ids in sorted order
+        expert_offsets: [E] cumulative token counts (original)
+        active_experts: [num_active] sorted unique expert ids
+        E: Total number of experts
+
+    Returns:
+        remapped_idxs: [T*k] expert ids in [0..num_active-1]
+        compact_offsets: [num_active] cumulative token counts
+    """
+    # Build remap table: global_id -> compact_id
+    remap = torch.empty(E, dtype=torch.long, device=sorted_expert_idxs.device)
+    remap[active_experts] = torch.arange(
+        len(active_experts), device=sorted_expert_idxs.device
+    )
+
+    remapped_idxs = remap[sorted_expert_idxs]
+
+    compact_offsets = expert_offsets[active_experts]
+
+    return remapped_idxs, compact_offsets
+
+
+def _selective_dequant_bnb4(
+    raw_param: torch.Tensor,
+    quant_state,
+    active_experts: torch.Tensor,
+    expert_shape: tuple[int, int],
+) -> torch.Tensor:
+    """Dequantize only selected experts from BnB 4-bit packed data.
+
+    The raw parameter is a flattened 4-bit packed tensor. Each expert's
+    data is contiguous (stored in expert-major order), so we can gather
+    the packed data and absmax blocks for active experts, then dequantize
+    as one contiguous block.
+
+    Args:
+        raw_param: Flattened uint8 tensor of packed 4-bit weights
+        quant_state: BnB QuantState with absmax, blocksize, code, etc.
+        active_experts: [num_active] expert indices to dequantize
+        expert_shape: (dim1, dim2) shape per expert (e.g. (1024, 2048))
+
+    Returns:
+        Dequantized weights [num_active, dim1, dim2] in original dtype
+    """
+    import bitsandbytes.functional as F  # noqa: N812
+    from bitsandbytes.functional import QuantState
+
+    expert_numel = expert_shape[0] * expert_shape[1]
+    num_active = len(active_experts)
+
+    # Per-expert slicing assumes each expert owns whole packed bytes (2 nf4 values/byte) and
+    # whole quant blocks; otherwise bytes/blocks straddle expert boundaries and the slice is
+    # silently wrong. Fall back to full dequant + index when that doesn't hold.
+    if (
+        expert_numel <= 0
+        or quant_state.blocksize <= 0
+        or expert_numel % 2 != 0
+        or expert_numel % quant_state.blocksize != 0
+    ):
+        full = F.dequantize_4bit(raw_param, quant_state)
+        E_total = full.numel() // expert_numel if expert_numel else 0
+        return full.reshape(E_total, *expert_shape)[active_experts]
+
+    packed_per_expert = expert_numel // 2  # 4-bit = 2 values per byte
+    blocks_per_expert = expert_numel // quant_state.blocksize
+
+    # Fused Triton kernel for NF4: selective gather + dequant in one pass.
+    if quant_state.quant_type == "nf4" and raw_param.dtype == torch.uint8:
+        from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant_kernel import (
+            selective_dequant_nf4_triton,
+        )
+
+        # Nested (double) quant: dequantize absmax first via dequantize_blockwise (not _4bit),
+        # then add the offset.
+        if quant_state.nested:
+            absmax = F.dequantize_blockwise(quant_state.absmax, quant_state.state2)
+            absmax += quant_state.offset
+            if absmax.dtype != torch.float32:
+                absmax = absmax.float()
+        else:
+            absmax = quant_state.absmax
+
+        return selective_dequant_nf4_triton(
+            packed_data=raw_param,
+            absmax=absmax,
+            active_experts=active_experts,
+            expert_shape=expert_shape,
+            blocksize=quant_state.blocksize,
+            dtype=quant_state.dtype,
+            codebook=quant_state.code,
+        )
+
+    # Fallback: gather + BnB dequant (for fp4 or non-uint8 packed formats)
+    raw_flat = raw_param.reshape(-1)
+
+    offsets_qt = (
+        active_experts.long()[:, None] * packed_per_expert
+        + torch.arange(packed_per_expert, device=raw_param.device)[None, :]
+    ).reshape(-1)
+    qt_gathered = raw_flat[offsets_qt]
+
+    offsets_abs = (
+        active_experts.long()[:, None] * blocks_per_expert
+        + torch.arange(blocks_per_expert, device=raw_param.device)[None, :]
+    ).reshape(-1)
+
+    if quant_state.nested:
+        full_absmax = F.dequantize_blockwise(quant_state.absmax, quant_state.state2)
+        full_absmax += quant_state.offset
+        if full_absmax.dtype != torch.float32:
+            full_absmax = full_absmax.float()
+        absmax_gathered = full_absmax[offsets_abs]
+    else:
+        absmax_gathered = quant_state.absmax[offsets_abs]
+
+    qt_gathered = qt_gathered.unsqueeze(1) if qt_gathered.dim() == 1 else qt_gathered
+
+    gathered_qs = QuantState(
+        absmax=absmax_gathered,
+        shape=torch.Size([num_active * expert_numel]),
+        blocksize=quant_state.blocksize,
+        quant_type=quant_state.quant_type,
+        code=quant_state.code,
+        dtype=quant_state.dtype,
+    )
+
+    deq = F.dequantize_4bit(qt_gathered, gathered_qs)
+    return deq.reshape(num_active, *expert_shape)
+
+
+def _selective_dequant_mxfp4(
+    mx_param,
+    active_experts: torch.Tensor,
+    out_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Selectively dequantize active experts from a torchao MXFP4 ``MXTensor``.
+
+    Layout assumption: the MXTensor's last axis is the OCP MX block axis.
+    For ScatterMoE experts this matches the natural storage where
+    ``experts.gate_up_proj``/``down_proj`` is ``[E, dim1, dim2]`` and
+    ``dim2`` is the contraction axis post ``.transpose(2, 1)`` performed by
+    the caller. Indexing ``[active_experts]`` on the qdata and scale yields
+    a compact MX tensor that we dequantize via torchao.
+
+    Args:
+        mx_param: ``torchao.prototype.mx_formats.mx_tensor.MXTensor`` of
+            logical shape ``[E, dim1, dim2]`` with ``elem_dtype=float4_e2m1fn_x2``.
+        active_experts: ``[num_active]`` sorted unique expert indices.
+        out_dtype: dtype of the dequantized buffer (default ``bfloat16``).
+
+    Returns:
+        Dequantized bf16/fp16 tensor of shape ``[num_active, dim1, dim2]``.
+    """
+    if _torchao_mxtensor_cls() is None:
+        raise ImportError(
+            "MXFP4 expert dequantization requires torchao>=0.7 "
+            "(install with `pip install torchao`)."
+        )
+
+    sub_qdata = _mx_qdata(mx_param)[active_experts].contiguous()
+    sub_scale = _mx_scale(mx_param)[active_experts].contiguous()
+
+    sub_mx = _construct_mxtensor_subset(mx_param, sub_qdata, sub_scale)
+    return sub_mx.dequantize(out_dtype)
+
+
+def _selective_index_dense(
+    param: torch.Tensor,
+    active_experts: torch.Tensor,
+) -> torch.Tensor:
+    """Select experts from a dense (bf16/fp32) weight tensor.
+
+    Simple indexing — no dequantization needed.
+    """
+    return param[active_experts]
+
+
+def selective_expert_weights(
+    experts_module: nn.Module,
+    param_name: str,
+    active_experts: torch.Tensor,
+) -> torch.Tensor:
+    """Extract and dequantize only the active experts' weights.
+
+    Format-agnostic: dispatches based on whether the parameter is
+    BnB 4-bit quantized (via parametrize), FP8, or dense bf16/fp32.
+
+    Args:
+        experts_module: The base experts module (e.g. Qwen3_5MoeExperts)
+        param_name: "gate_up_proj" or "down_proj"
+        active_experts: [num_active] sorted unique expert indices
+
+    Returns:
+        Compact weight tensor [num_active, dim1, dim2] ready for ScatterMoE
+    """
+    if (
+        hasattr(experts_module, "parametrizations")
+        and param_name in experts_module.parametrizations
+    ):
+        param_list = experts_module.parametrizations[param_name]
+        parametrization = param_list[0]
+
+        # BnB 4-bit parametrization
+        if hasattr(parametrization, "quant_state"):
+            # The raw quantized data is on the ParametrizationList, not the
+            # individual Bnb4bitParametrization module
+            raw_param = param_list.original
+            qs = parametrization.quant_state
+            # qs.shape is the original tensor shape before flattening.
+            # For MoE experts it's [E, d1, d2] (3D) or [total_elements] (1D).
+            orig_shape = qs.shape
+            if isinstance(orig_shape, torch.Size) and len(orig_shape) == 3:
+                expert_shape = (orig_shape[1], orig_shape[2])
+            elif isinstance(orig_shape, torch.Size) and len(orig_shape) == 1:
+                # Flattened; infer the expert shape from module attributes.
+                E_total = getattr(experts_module, "num_experts", None)
+                if E_total is None:
+                    E_total = int(active_experts.max().item()) + 1
+                expert_numel = orig_shape[0] // E_total
+                d2 = getattr(experts_module, "hidden_dim", None) or getattr(
+                    experts_module, "intermediate_dim", None
+                )
+                if d2 and expert_numel % d2 == 0:
+                    expert_shape = (expert_numel // d2, d2)
+                else:
+                    full = getattr(experts_module, param_name)
+                    return full[active_experts]
+            else:
+                full = getattr(experts_module, param_name)
+                return full[active_experts]
+
+            return _selective_dequant_bnb4(raw_param, qs, active_experts, expert_shape)
+
+    param = getattr(experts_module, param_name)
+
+    # MXFP4 (torchao MXTensor): dequantize the subset, return [num_active, d1, d2]
+    if is_mxfp4_param(param):
+        return _selective_dequant_mxfp4(param, active_experts)
+
+    # Dense parameter (bf16/fp32): direct indexing
+    if param.dim() == 3:
+        return param[active_experts]
+
+    # Fallback: full access
+    return param
+
+
+def shared_dequant_across_shards(
+    experts_module: nn.Module,
+    param_name: str,
+    sei_per_shard: list[torch.Tensor],
+    E: int,
+) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]:
+    """Dequantize the union of active experts across N shards exactly once.
+
+    The orthogonal Strategy A path calls :func:`selective_expert_weights`
+    once per shard, which re-dequantizes the active experts redundantly
+    when the active-expert sets overlap. For seq-dim sharding with a
+    softmax-routed MoE, that overlap is the common case.
+
+    This helper hoists the dequant: it computes the union of active
+    experts across all shards, calls :func:`selective_expert_weights`
+    once on the union, and returns per-shard index tables that map each
+    shard's local active experts into rows of the union buffer.
+
+    Parameters
+    ----------
+    experts_module:
+        The base experts module (e.g. ``OlmoeExperts``). Same object the
+        per-shard path would pass to :func:`selective_expert_weights`.
+    param_name:
+        ``"gate_up_proj"`` or ``"down_proj"``.
+    sei_per_shard:
+        List of ``sorted_expert_idxs`` tensors, one per shard.
+    E:
+        Total number of experts.
+
+    Returns
+    -------
+    union_active:
+        ``[U]`` sorted unique expert ids across all shards.
+    union_buffer:
+        Dequantized weights for ``union_active``,
+        ``[U, dim1, dim2]`` in the param's natural storage dtype
+        (typically bf16). Same buffer each shard's call would have built
+        had it dequantized only its own active set, just shared.
+    shard_into_union:
+        List of length ``len(sei_per_shard)``. Entry ``i`` is a 1-D
+        ``long`` tensor that indexes ``union_buffer`` along dim 0 to
+        produce the same ``[num_active_i, dim1, dim2]`` slice the
+        per-shard path would have produced. Callers feed this through
+        ``union_buffer.index_select(0, shard_into_union[i])`` (or
+        equivalent advanced indexing) before handing the slice to
+        ``parallel_linear_lora``.
+
+    Bitwise contract: composing ``union_buffer.index_select(0,
+    shard_into_union[i])`` is byte-identical to
+    ``selective_expert_weights(experts_module, param_name,
+    get_active_experts(sei_per_shard[i], E))`` because both paths slice
+    the same dequantized MX subset by the same expert ids. The
+    ``test_shared_dequant_helper.py`` parity test asserts this.
+    """
+    if not sei_per_shard:
+        raise ValueError("sei_per_shard must contain at least one tensor")
+
+    device = sei_per_shard[0].device
+    per_shard_active = [get_active_experts(sei, E) for sei in sei_per_shard]
+    union_active = torch.unique(torch.cat(per_shard_active))
+
+    union_buffer = selective_expert_weights(experts_module, param_name, union_active)
+
+    # Build the global-id → union-row remap once, then gather per shard.
+    # ``union_active`` is sorted and unique by construction, so the inverse
+    # lookup is dense over ``E``.
+    union_remap = torch.empty(E, dtype=torch.long, device=device)
+    union_remap[union_active] = torch.arange(
+        len(union_active), device=device, dtype=torch.long
+    )
+    shard_into_union = [union_remap[active] for active in per_shard_active]
+
+    return union_active, union_buffer, shard_into_union
+
+
+def selective_lora_weights(
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    active_experts: torch.Tensor,
+    E: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Select LoRA A and B weights for only the active experts.
+
+    LoRA layout (scattermoe format):
+      A: [r*E, K] — expert e occupies rows [e*r : (e+1)*r]
+      B: [N, r*E] — expert e occupies cols [e*r : (e+1)*r]
+
+    Returns compact:
+      A: [r*num_active, K]
+      B: [N, r*num_active]
+    """
+    R = lora_A.size(0) // E
+
+    # Vectorized gather: active_experts[:, None] * R + arange(R)[None, :]
+    row_idx = (
+        active_experts.long()[:, None] * R
+        + torch.arange(R, device=lora_A.device)[None, :]
+    ).reshape(-1)
+
+    compact_A = lora_A[row_idx]  # [r*num_active, K]
+    compact_B = lora_B[:, row_idx]  # [N, r*num_active]
+
+    return compact_A, compact_B
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/selective_dequant_kernel.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/selective_dequant_kernel.py
new file mode 100644
index 0000000000..aa9f0278a4
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/selective_dequant_kernel.py
@@ -0,0 +1,179 @@
+"""
+Triton kernel for fused selective expert gather + NF4 dequantization.
+
+Instead of:
+  1. Gather packed uint8 data for active experts (memory copy)
+  2. Gather absmax for active experts (memory copy)
+  3. Call BnB dequantize_4bit CUDA kernel
+
+This kernel does all three in one pass:
+  - Reads packed NF4 bytes from expert-strided positions
+  - Looks up the NF4 codebook
+  - Multiplies by the per-block absmax
+  - Writes bf16 output directly
+
+This eliminates the intermediate gather buffer entirely.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+# NF4 codebook (16 values, precomputed by BnB)
+# These are the normalized float4 reconstruction values
+NF4_CODEBOOK = [
+    -1.0,
+    -0.6961928009986877,
+    -0.5250730514526367,
+    -0.39491748809814453,
+    -0.28444138169288635,
+    -0.18477343022823334,
+    -0.09105003625154495,
+    0.0,
+    0.07958029955625534,
+    0.16093020141124725,
+    0.24611230194568634,
+    0.33791524171829224,
+    0.44070982933044434,
+    0.5626170039176941,
+    0.7229568362236023,
+    1.0,
+]
+
+
+@triton.jit
+def _selective_dequant_nf4_kernel(
+    # Input: packed NF4 data (flattened, expert-major order)
+    packed_ptr,
+    # Input: absmax values (flattened, expert-major order)
+    absmax_ptr,
+    # Input: active expert indices
+    active_experts_ptr,
+    # Input: NF4 codebook (16 float values)
+    codebook_ptr,
+    # Output: dequantized bf16 weights [num_active, expert_numel]
+    out_ptr,
+    stride_out_e,  # stride for expert dim in output
+    # Dimensions
+    num_active,
+    packed_per_expert,  # expert_numel // 2
+    blocks_per_expert,  # expert_numel // blocksize
+    blocksize: tl.constexpr,
+    # Tile size
+    BLOCK_SIZE: tl.constexpr,  # elements per thread block (must be multiple of 2)
+):
+    """
+    Each program processes BLOCK_SIZE elements from one expert.
+
+    Grid: (num_active, cdiv(expert_numel, BLOCK_SIZE))
+
+    For each output element:
+      1. Compute which byte in packed data contains this element
+      2. Extract the 4-bit nibble (high or low)
+      3. Look up in NF4 codebook
+      4. Scale by absmax for this block
+    """
+    expert_local_idx = tl.program_id(0)  # which active expert (0..num_active-1)
+    block_id = tl.program_id(1)  # which element block
+
+    # Load the global expert index
+    expert_global = tl.load(active_experts_ptr + expert_local_idx).to(tl.int64)
+
+    expert_numel = packed_per_expert * 2  # 2 elements per packed byte
+    elem_offset = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = elem_offset < expert_numel
+
+    # Each element is packed as: byte[i//2], low nibble for even i, high for odd i
+    byte_idx = elem_offset // 2
+    is_high = (elem_offset % 2) == 1
+
+    # Read packed bytes from the global expert's region
+    packed_global_offset = expert_global * packed_per_expert + byte_idx
+    packed_bytes = tl.load(packed_ptr + packed_global_offset, mask=mask, other=0).to(
+        tl.int32
+    )
+
+    # Extract 4-bit nibble
+    # BnB packing: high nibble = even element, low nibble = odd element
+    nibble = tl.where(is_high, packed_bytes & 0xF, (packed_bytes >> 4) & 0xF)
+
+    # NF4 codebook lookup
+    # Load all 16 codebook values (small, fits in registers)
+    # Use gather from codebook pointer
+    code_val = tl.load(codebook_ptr + nibble, mask=mask, other=0.0)
+
+    # Load absmax for this element's quantization block
+    block_idx = elem_offset // blocksize
+    absmax_global_offset = expert_global * blocks_per_expert + block_idx
+    absmax_val = tl.load(absmax_ptr + absmax_global_offset, mask=mask, other=1.0)
+
+    # Dequantize: value = codebook[nibble] * absmax
+    result = code_val * absmax_val
+
+    # Store to output
+    out_offset = expert_local_idx * stride_out_e + elem_offset
+    tl.store(out_ptr + out_offset, result.to(out_ptr.dtype.element_ty), mask=mask)
+
+
+def selective_dequant_nf4_triton(
+    packed_data: torch.Tensor,
+    absmax: torch.Tensor,
+    active_experts: torch.Tensor,
+    expert_shape: tuple[int, int],
+    blocksize: int,
+    dtype: torch.dtype = torch.bfloat16,
+    codebook: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused selective gather + NF4 dequantization via Triton kernel.
+
+    Args:
+        packed_data: Flattened packed NF4 data [total_packed] or [total_packed, 1]
+        absmax: Per-block scaling factors [total_blocks]
+        active_experts: Sorted indices of experts to dequantize [num_active]
+        expert_shape: (dim1, dim2) per expert
+        blocksize: Quantization block size
+        dtype: Output dtype (default bf16)
+        codebook: NF4 lookup table [16] (uses default NF4 codebook if None)
+
+    Returns:
+        Dequantized weights [num_active, dim1, dim2]
+    """
+    num_active = active_experts.shape[0]
+    expert_numel = expert_shape[0] * expert_shape[1]
+    packed_per_expert = expert_numel // 2
+    blocks_per_expert = expert_numel // blocksize
+
+    # Prepare codebook on device
+    if codebook is None:
+        codebook = torch.tensor(
+            NF4_CODEBOOK, dtype=torch.float32, device=packed_data.device
+        )
+    else:
+        codebook = codebook.to(device=packed_data.device, dtype=torch.float32)
+
+    # Flatten inputs
+    packed_flat = packed_data.reshape(-1)
+    absmax_flat = absmax.reshape(-1).float()  # absmax is usually fp32
+
+    # Output buffer
+    out = torch.empty(num_active, expert_numel, dtype=dtype, device=packed_data.device)
+
+    BLOCK_SIZE = 1024  # Process 1024 elements per thread block
+
+    grid = (num_active, triton.cdiv(expert_numel, BLOCK_SIZE))
+
+    _selective_dequant_nf4_kernel[grid](
+        packed_flat,
+        absmax_flat,
+        active_experts,
+        codebook,
+        out,
+        out.stride(0),
+        num_active=num_active,
+        packed_per_expert=packed_per_expert,
+        blocks_per_expert=blocks_per_expert,
+        blocksize=blocksize,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+    return out.reshape(num_active, *expert_shape)
diff --git a/src/axolotl/integrations/kernels/libs/scattermoe_lora/torchao_fp4_add.py b/src/axolotl/integrations/kernels/libs/scattermoe_lora/torchao_fp4_add.py
new file mode 100644
index 0000000000..5f95b2417b
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/scattermoe_lora/torchao_fp4_add.py
@@ -0,0 +1,73 @@
+"""Make torchao FP4 tensors support ``aten.add`` by dequantizing.
+
+PEFT (>= 0.19) applies ``target_parameters`` LoRA to an expert weight via
+``torch.nn.utils.parametrize``: it registers a parametrization whose forward computes
+``base + delta`` (the merged ``scaling * B @ A``). ``register_parametrization`` evaluates
+this immediately. When the base weight is a frozen torchao ``NVFP4Tensor`` / ``MXTensor``
+(e.g. `block_diffusion.frozen_fp4_experts`), ``aten.add`` is unimplemented for the tensor
+subclass, so the registration raises::
+
+    NotImplementedError: NVFP4Tensor dispatch: ... aten.add ...
+
+Registering a dequantize-then-add for these tensors makes the parametrization produce a
+bf16 merged weight (differentiable through ``delta`` to the LoRA A/B), so ScatterMoE then
+runs the merged weight via its standard path. Idempotent; safe to call repeatedly.
+"""
+
+from __future__ import annotations
+
+
+def _make_add_handler():
+    import torch
+
+    fp4_names = {"NVFP4Tensor", "MXTensor"}
+
+    def _is_fp4(t):
+        return type(t).__name__ in fp4_names
+
+    def _fp4_add(func, types, args, kwargs):
+        a, b = args[0], args[1]
+        # Use the dense operand's dtype (the LoRA delta is bf16/fp16) so the result is a plain
+        # tensor; dequantize only the FP4 subclass operand(s).
+        dense_dtype = next(
+            (t.dtype for t in (a, b) if isinstance(t, torch.Tensor) and not _is_fp4(t)),
+            torch.bfloat16,
+        )
+        if _is_fp4(a):
+            a = a.dequantize(dense_dtype)
+        if _is_fp4(b):
+            b = b.dequantize(dense_dtype)
+        # A true in-place add_ is impossible once the FP4 operand is dequantized to a fresh tensor,
+        # so both add and add_ resolve to this out-of-place add (kwargs forward keyword-only alpha).
+        return torch.ops.aten.add.Tensor(a, b, **kwargs)
+
+    return _fp4_add
+
+
+def patch_torchao_fp4_add() -> None:
+    """Register dequantize-add for ``aten.add``/``aten.add_`` on NVFP4Tensor and MXTensor.
+
+    ``add.Tensor`` covers PEFT's out-of-place ``base + delta`` parametrization (training);
+    ``add_.Tensor`` covers ``ParamWrapper.merge`` (``merge_and_unload`` produces a bf16
+    weight — a true 4-bit merge would require re-quantization).
+    """
+    import torch
+
+    handler = _make_add_handler()
+    ops = [torch.ops.aten.add.Tensor, torch.ops.aten.add_.Tensor]
+
+    for module_path, cls_name in (
+        ("torchao.prototype.mx_formats.nvfp4_tensor", "NVFP4Tensor"),
+        ("torchao.prototype.mx_formats.mx_tensor", "MXTensor"),
+    ):
+        try:
+            module = __import__(module_path, fromlist=[cls_name])
+            cls = getattr(module, cls_name)
+        except (ImportError, AttributeError):
+            continue
+        # torchao's op table is nested: cls._ATEN_OP_TABLE[cls][op] (see
+        # torchao.utils._dispatch__torch_dispatch__).
+        registered = getattr(cls, "_ATEN_OP_TABLE", {}).get(cls, {})
+        for op in ops:
+            if op not in registered:
+                cls.implements([op])(handler)
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/__init__.py b/src/axolotl/integrations/kernels/libs/sonicmoe/__init__.py
new file mode 100644
index 0000000000..aaa3542db5
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/__init__.py
@@ -0,0 +1,42 @@
+from .experts import register_sonicmoe_experts, sonicmoe_experts_forward_with_lora
+from .multi_lora import (
+    MoEMultiLoRAMaterialize,
+    combined_expert_ids,
+    materialize_multi_lora_experts,
+)
+from .nvfp4 import (
+    dequantize_expert_weight,
+    gated_activation,
+    grouped_down_gemm,
+    grouped_up_gemm,
+    is_nvfp4_param,
+    resolve_gated_activation,
+)
+from .nvfp4_lora import (
+    GroupedDownProjLoRA,
+    GroupedUpProjLoRA,
+    combine_expert_outputs,
+    grouped_expert_mlp_lora,
+    grouped_moe_reference_forward,
+    route_and_group,
+)
+
+__all__ = [
+    "register_sonicmoe_experts",
+    "sonicmoe_experts_forward_with_lora",
+    "MoEMultiLoRAMaterialize",
+    "combined_expert_ids",
+    "materialize_multi_lora_experts",
+    "is_nvfp4_param",
+    "dequantize_expert_weight",
+    "gated_activation",
+    "resolve_gated_activation",
+    "grouped_up_gemm",
+    "grouped_down_gemm",
+    "GroupedUpProjLoRA",
+    "GroupedDownProjLoRA",
+    "grouped_expert_mlp_lora",
+    "route_and_group",
+    "combine_expert_outputs",
+    "grouped_moe_reference_forward",
+]
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/experts.py b/src/axolotl/integrations/kernels/libs/sonicmoe/experts.py
new file mode 100644
index 0000000000..96846098b6
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/experts.py
@@ -0,0 +1,300 @@
+"""LoRA-aware sonicmoe experts forward for the transformers ExpertsInterface.
+
+Dense experts wrap upstream ``_sonicmoe_wrapper``, materializing expert LoRA via
+``MoELoRAMaterialize`` before the CUTLASS call. NVFP4 experts (which that kernel
+cannot read) take the grouped dequant path in ``nvfp4_lora`` instead.
+"""
+
+from __future__ import annotations
+
+import functools
+import os
+
+import torch
+
+from .lora import (
+    MoELoRAMaterialize,
+    get_lora_params_from_wrapper,
+    has_lora,
+    materialize_expert_lora,
+    unwrap_experts_lora,
+)
+
+
+def _maybe_unwrap_param_wrapper(param):
+    """Return ``(base_tensor, lora_params_or_None)`` for a PEFT-wrapped Parameter."""
+    try:
+        from peft.tuners.param_wrapper import ParamWrapper
+    except ImportError:
+        return param, None
+
+    if not isinstance(param, ParamWrapper):
+        return param, None
+
+    base = param.original_parameter
+    lora_A, lora_B, scaling = get_lora_params_from_wrapper(param)
+    if lora_A is None:
+        return base, None
+    return base, (lora_A, lora_B, scaling)
+
+
+def _resolve_weights_and_lora(experts_module):
+    """Resolve raw expert weights/biases + optional LoRA tuples.
+
+    Handles both PEFT layouts: module-level wrap (walked via ``unwrap_experts_lora``)
+    and per-parameter ``ParamWrapper``. No layout permute applied.
+    """
+    # The ParamWrapper fastpath (experts_lora_fastpath) resolves the LoRA tuples itself and
+    # hands them over before calling the raw base module.
+    fastpath_lora = getattr(experts_module, "_sonicmoe_lora", None)
+    if fastpath_lora is not None:
+        w1 = experts_module.gate_up_proj
+        w2 = experts_module.down_proj
+        b1 = getattr(experts_module, "gate_up_proj_bias", None)
+        b2 = getattr(experts_module, "down_proj_bias", None)
+        return (
+            w1,
+            b1,
+            w2,
+            b2,
+            fastpath_lora.get("gate_up_proj"),
+            fastpath_lora.get("down_proj"),
+        )
+
+    if has_lora(experts_module):
+        base_experts, lora_dict = unwrap_experts_lora(experts_module)
+        w1 = base_experts.gate_up_proj
+        w2 = base_experts.down_proj
+        b1 = getattr(base_experts, "gate_up_proj_bias", None)
+        b2 = getattr(base_experts, "down_proj_bias", None)
+        return w1, b1, w2, b2, lora_dict.get("gate_up_proj"), lora_dict.get("down_proj")
+
+    w1, lora_w1 = _maybe_unwrap_param_wrapper(experts_module.gate_up_proj)
+    w2, lora_w2 = _maybe_unwrap_param_wrapper(experts_module.down_proj)
+    b1 = getattr(experts_module, "gate_up_proj_bias", None)
+    b2 = getattr(experts_module, "down_proj_bias", None)
+    return w1, b1, w2, b2, lora_w1, lora_w2
+
+
+@functools.lru_cache(maxsize=1)
+def _sonicmoe_kernel_supported() -> bool:
+    """The kernels-community/sonic-moe CUTLASS GEMM fails to compile on sm_120
+    (Blackwell): its bundled quack ``GemmSm120`` predates the ``concat_layout`` arg
+    the dispatcher passes. Gate it off there so standard-layout models fall back to
+    the vendored scattermoe Triton path. Drop this once a fixed sonic-moe ships."""
+    if not torch.cuda.is_available():
+        return True
+    return torch.cuda.get_device_capability()[0] < 12
+
+
+def sonicmoe_experts_forward_with_lora(
+    self,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    """Sonicmoe experts forward with PEFT LoRA support.
+
+    Dense bf16 experts use the fast sonic-moe CUTLASS kernel (LoRA materialized
+    into W_eff first). NVFP4 experts, which the opaque CUTLASS kernel cannot
+    read, take the grouped reference path (dequant base + fused low-rank LoRA).
+    On sm_120, where the sonic-moe kernel can't compile, standard-layout dense
+    experts fall back to the vendored scattermoe Triton path.
+    """
+    from ..scattermoe_lora.experts import (
+        scattermoe_experts_forward,
+        scattermoe_supports_layout,
+    )
+    from .nvfp4 import is_nvfp4_param
+
+    if not getattr(self, "has_gate", True):
+        raise ValueError("sonicmoe requires gated experts (has_gate=True)")
+    if hidden_states.device.type != "cuda":
+        raise ValueError("sonicmoe requires CUDA device")
+
+    w1, b1, w2, b2, lora_w1, lora_w2 = _resolve_weights_and_lora(self)
+    if not getattr(self, "has_bias", False):
+        b1 = b2 = None
+
+    # Unwrap FSDP2/EP DTensors to local shards. to_local() is autograd-aware:
+    # backward rewraps the gradient as a DTensor.
+    if isinstance(w1, torch.distributed.tensor.DTensor):
+        w1 = w1.to_local()
+        w2 = w2.to_local()
+        b1 = b1.to_local() if b1 is not None else None
+        b2 = b2.to_local() if b2 is not None else None
+
+    # The opaque CUTLASS kernel cannot read packed FP4, so NVFP4 experts take the
+    # grouped dequant path instead.
+    if is_nvfp4_param(w1) or is_nvfp4_param(w2):
+        return _sonicmoe_nvfp4_forward(
+            self,
+            hidden_states,
+            top_k_index,
+            top_k_weights,
+            w1,
+            b1,
+            w2,
+            b2,
+            lora_w1,
+            lora_w2,
+        )
+
+    if not _sonicmoe_kernel_supported() and scattermoe_supports_layout(self):
+        return scattermoe_experts_forward(
+            self, hidden_states, top_k_index, top_k_weights
+        )
+
+    from transformers.integrations.sonicmoe import _sonicmoe_wrapper
+
+    device = hidden_states.device
+    num_top_k = top_k_index.size(-1)
+    num_tokens = hidden_states.size(0)
+
+    # sonic-moe requires int32 token indices sorted ascending.
+    token_idx = (
+        torch.arange(num_tokens, device=device)
+        .unsqueeze(1)
+        .expand(-1, num_top_k)
+        .reshape(-1)
+        .int()
+    )
+    router_scores = top_k_weights.reshape(-1).to(hidden_states.dtype)
+    expert_ids = top_k_index.reshape(-1).int()
+
+    # Materialize W_eff = W + scaling * (B @ A) per expert. No-op when no LoRA.
+    if lora_w1 is not None:
+        w1 = MoELoRAMaterialize.apply(w1, *lora_w1)
+    if lora_w2 is not None:
+        w2 = MoELoRAMaterialize.apply(w2, *lora_w2)
+
+    # Match upstream layout expectations:
+    #   is_transposed=False: gate_up [E, 2*I, H] / down [E, H, I] -> permute(1, 2, 0)
+    #   is_transposed=True:  gate_up [E, H, 2*I] / down [E, I, H] -> permute(2, 1, 0)
+    perm = (2, 1, 0) if getattr(self, "is_transposed", False) else (1, 2, 0)
+    w1 = w1.permute(*perm)
+    w2 = w2.permute(*perm)
+
+    act_name = getattr(self.config, "hidden_act", "silu").lower()
+
+    return _sonicmoe_wrapper(
+        hidden_states=hidden_states,
+        router_scores=router_scores,
+        expert_ids=expert_ids,
+        token_idx=token_idx,
+        w1=w1,
+        b1=b1,
+        w2=w2,
+        b2=b2,
+        act_name=act_name,
+        num_experts=self.num_experts,
+        concat_layout=getattr(self, "is_concatenated", True),
+        is_inference_mode_enabled=not torch.is_grad_enabled(),
+    )
+
+
+def _select_nvfp4_backend(w1, w2) -> str:
+    """``fp4_cute`` (SM100 in-kernel W4A4) when available and dims align, else
+    ``dequant``. ``AXOLOTL_SONICMOE_NVFP4_BACKEND`` forces a backend."""
+    override = os.environ.get("AXOLOTL_SONICMOE_NVFP4_BACKEND")
+    if override:
+        return override
+
+    from .fp4_cute import fp4_cute_available
+    from .fp4_cute_ops import fp4_cute_dims_ok
+    from .nvfp4 import is_nvfp4_param
+
+    if (
+        is_nvfp4_param(w1)
+        and is_nvfp4_param(w2)
+        and fp4_cute_available()
+        and fp4_cute_dims_ok(w1, w2)
+    ):
+        return "fp4_cute"
+    return "dequant"
+
+
+def _sonicmoe_nvfp4_forward(
+    self,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    w1: torch.Tensor,
+    b1,
+    w2: torch.Tensor,
+    b2,
+    lora_w1,
+    lora_w2,
+) -> torch.Tensor:
+    """NVFP4 experts forward via the grouped path (base frozen).
+
+    LoRA is applied as a fused low-rank delta at the grouped-token level (no
+    base-weight gradient). On SM100/SM110 the base stays packed and runs the
+    in-kernel W4A4 grouped GEMM (``fp4_cute``); elsewhere it is dequantized to
+    dense per matmul (``dequant``).
+    """
+    from .nvfp4 import resolve_gated_activation
+    from .nvfp4_lora import grouped_moe_reference_forward
+
+    if getattr(self, "is_transposed", False):
+        raise NotImplementedError(
+            "sonicmoe NVFP4 path supports the [E, 2*I, H] / [E, H, I] layout only "
+            "(is_transposed=False)"
+        )
+
+    # Expert parallelism is not finished for this path: the base is EP-sharded to
+    # local experts while routing/LoRA still use global expert ids, and the local
+    # token dispatch is not implemented. Fail loudly instead of computing garbage.
+    if getattr(self, "num_experts_global", self.num_experts) != self.num_experts:
+        raise NotImplementedError(
+            "sonicmoe NVFP4 path does not support expert parallelism yet "
+            "(EP-sharded base with global-id routing/LoRA is unfinished and untested)"
+        )
+
+    lora1 = (lora_w1[0], lora_w1[1]) if lora_w1 is not None else None
+    lora2 = (lora_w2[0], lora_w2[1]) if lora_w2 is not None else None
+    scaling1 = lora_w1[2] if lora_w1 is not None else 1.0
+    scaling2 = lora_w2[2] if lora_w2 is not None else 1.0
+
+    return grouped_moe_reference_forward(
+        hidden_states,
+        top_k_index,
+        top_k_weights,
+        w1,
+        b1,
+        w2,
+        b2,
+        lora1,
+        lora2,
+        self.num_experts,
+        act=resolve_gated_activation(self.config),
+        backend=_select_nvfp4_backend(w1, w2),
+        limit=getattr(self, "limit", None),
+        concat=getattr(self, "is_concatenated", True),
+        scaling1=scaling1,
+        scaling2=scaling2,
+    )
+
+
+def register_sonicmoe_experts() -> None:
+    """Register the LoRA-aware ``"sonicmoe"`` forward, overriding upstream. Idempotent."""
+    from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+
+    ALL_EXPERTS_FUNCTIONS.register("sonicmoe", sonicmoe_experts_forward_with_lora)
+
+    # Route PEFT target_parameters expert LoRA past the parametrization merge (which cannot
+    # run on quantized bases) to the fused low-rank path.
+    try:
+        from .experts_lora_fastpath import patch_paramwrapper_sonicmoe_fastpath
+
+        patch_paramwrapper_sonicmoe_fastpath()
+    except (ImportError, AttributeError):
+        pass
+
+
+# Re-export utilities for tests / external callers.
+__all__ = [
+    "sonicmoe_experts_forward_with_lora",
+    "register_sonicmoe_experts",
+    "materialize_expert_lora",
+]
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/experts_lora_fastpath.py b/src/axolotl/integrations/kernels/libs/sonicmoe/experts_lora_fastpath.py
new file mode 100644
index 0000000000..b0b446b333
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/experts_lora_fastpath.py
@@ -0,0 +1,95 @@
+"""Route PEFT ``target_parameters`` expert LoRA to the sonicmoe grouped path.
+
+Same problem and fix shape as ``scattermoe_lora/experts_lora_fastpath.py``: when PEFT
+targets ``experts.gate_up_proj``/``experts.down_proj`` via ``target_parameters``, the
+experts module becomes a ``ParamWrapper`` chain whose forward materializes the merged
+weight ``base + delta`` through ``_activate_lora`` before the experts interface runs.
+That defeats the fused low-rank LoRA path and fails outright on quantized bases
+(``aten.add`` on NVFP4Tensor is unimplemented). This patches ``ParamWrapper.forward``
+so that, when the wrapped base is a sonicmoe experts module, the LoRA A/B/scaling are
+handed to the base via ``_sonicmoe_lora`` and the base forward is called directly. The
+wrappers stay in the module tree, so PEFT save/load and optimizer tracking are unaffected.
+"""
+
+from __future__ import annotations
+
+
+def _is_sonicmoe_experts(module) -> bool:
+    cfg = getattr(module, "config", None)
+    impl = getattr(cfg, "_experts_implementation", None)
+    return impl == "sonicmoe" and hasattr(module, "gate_up_proj")
+
+
+def patch_paramwrapper_sonicmoe_fastpath() -> None:
+    """Idempotently patch ``peft...ParamWrapper.forward`` for sonicmoe experts."""
+    try:
+        from peft.tuners.lora.layer import ParamWrapper
+    except (ImportError, AttributeError):
+        return
+
+    if getattr(ParamWrapper.forward, "_sonicmoe_fastpath", False):
+        return
+
+    from .lora import get_lora_params_from_wrapper
+
+    _orig_forward = ParamWrapper.forward
+
+    def _fusable(wrapper) -> bool:
+        # Fused path needs exactly one adapter on a raw (un-merged) base; else defer to PEFT.
+        if getattr(wrapper, "disable_adapters", False) or getattr(
+            wrapper, "merged", False
+        ):
+            return False
+        return len(getattr(wrapper, "active_adapters", [])) == 1
+
+    def forward(self, x, *args, **kwargs):
+        # Mixed-batch inference (adapter_names) is PEFT's domain; defer to it.
+        if kwargs.get("adapter_names") is not None:
+            return _orig_forward(self, x, *args, **kwargs)
+
+        # one wrapper per targeted parameter
+        wrappers = {}
+        base = self
+        while hasattr(base, "base_layer") and hasattr(base, "lora_A"):
+            name = getattr(base, "parameter_name", None)
+            if name is not None:
+                wrappers[name] = base
+            base = base.base_layer
+
+        if not (
+            _is_sonicmoe_experts(base) and all(_fusable(w) for w in wrappers.values())
+        ):
+            return _orig_forward(self, x, *args, **kwargs)
+
+        lora = {}
+        for name, wrapper in wrappers.items():
+            lora_A, lora_B, scaling = get_lora_params_from_wrapper(wrapper)
+            if lora_A is None:
+                continue
+            # PEFT keeps LoRA fp32; cast to activation dtype (grads still route to the fp32 params).
+            lora[name] = (lora_A.to(x.dtype), lora_B.to(x.dtype), scaling)
+
+        base._sonicmoe_lora = lora
+        # Loading can create multiple ExpertsInterface registries; the experts forward binds one as
+        # a closure default that may differ from the one register_sonicmoe_experts() populated, so
+        # register into the one THIS module dispatches through.
+        if not getattr(base, "_sonicmoe_iface_ok", False):
+            _fn = type(base).forward
+            _cells = dict(
+                zip(_fn.__code__.co_freevars, _fn.__closure__ or (), strict=False)
+            )
+            _ei = _cells.get("experts_interface")
+            if _ei is not None:
+                _ei = _ei.cell_contents
+                if "sonicmoe" not in _ei:
+                    from .experts import sonicmoe_experts_forward_with_lora
+
+                    _ei.register("sonicmoe", sonicmoe_experts_forward_with_lora)
+            base._sonicmoe_iface_ok = True
+        try:
+            return base(x, *args, **kwargs)
+        finally:
+            base._sonicmoe_lora = None
+
+    forward._sonicmoe_fastpath = True  # type: ignore[attr-defined]
+    ParamWrapper.forward = forward
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/fp4_cute.py b/src/axolotl/integrations/kernels/libs/sonicmoe/fp4_cute.py
new file mode 100644
index 0000000000..7f2466951c
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/fp4_cute.py
@@ -0,0 +1,669 @@
+"""SM100 grouped NVFP4 (W4A4) GEMM host driver, composed from quack.
+
+Instantiates quack's ``GemmGatedSm100`` / ``GemmDefaultSm100`` with
+``sf_vec_size=16`` (block-scaled NVFP4 mainloop + gated or default epilogue,
+zero quack kernel changes) and drives it with our own compile path:
+quack's stock drivers never marry the blockscaled mainloop with the act/gated
+epilogue, and its varlen host helpers guard fp4 off. Verified against quack
+source at f4f54db0; runtime pin quack-kernels==0.5.0, internals are private API.
+
+Conventions:
+- All GEMMs are ``C[e] = A[e] @ B[e]^T`` per expert, A/B K-major, fp32 accum.
+- Grouped mode is varlen_m: A rows expert-sorted and packed (NO padding),
+  ``cu_seqlens (E+1,) int32``; only SFA storage pads (see sf_layout).
+- The gated epilogue consumes INTERLEAVED gate/up along N:
+  ``postact[:, j] = act(alpha * H[:, 2j], alpha * H[:, 2j+1])``. Concat-layout
+  weights must be row-permuted at load (``sf_layout.gate_up_interleave_perm``),
+  and the stored preact D comes out interleaved (deinterleave for backward).
+- ``alpha`` scales the fp32 accumulator before the activation and before the
+  preact store: use it for the activation global scale. Per-expert weight
+  ``per_tensor_scale`` is applied EXACTLY via the epilogue colvec variant
+  (``forward(colvec=...)``: a per-row fp32 vector multiplied into the fp32
+  accumulator before the D store); the lossy SFB fold
+  (``sf_layout.fold_per_tensor_scale``) remains available for A/B debugging.
+
+GPU-only (SM100/SM110); quack/cutlass import lazily.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import torch
+
+try:
+    from .sf_layout import (
+        fold_per_tensor_scale,
+        gate_up_interleave_perm,
+        pack_scales_blocked,
+    )
+except ImportError:  # loaded standalone (no package parent)
+    from sf_layout import (  # type: ignore[no-redef]
+        fold_per_tensor_scale,
+        gate_up_interleave_perm,
+        pack_scales_blocked,
+    )
+
+SF_VEC_SIZE = 16
+
+
+@functools.lru_cache(maxsize=1)
+def fp4_cute_available() -> bool:
+    """True iff CUDA SM100/SM110 with quack + cutlass DSL importable."""
+    try:
+        if not torch.cuda.is_available():
+            return False
+        if torch.cuda.get_device_capability()[0] not in (10, 11):
+            return False
+        import cutlass  # noqa: F401
+        import quack.gemm_act  # noqa: F401
+
+        return True
+    except Exception:
+        return False
+
+
+_COMPILE_CACHE: dict = {}
+
+
+@functools.lru_cache(maxsize=1)
+def _rowscale_gemm_cls():
+    """``GemmDefaultSm100`` with the stock colvec ADD swapped for an exact fp32
+    per-row MULTIPLY on the accumulator (quack's ``vec_multiply``), applied
+    before the D convert/store and thus also before the ``add_to_output``
+    reduce-add."""
+    from typing import Optional
+
+    import cutlass.cute as cute
+    import quack.utils as utils
+    from cutlass import const_expr
+    from quack.epi_ops import vec_multiply
+    from quack.gemm_default_epi import GemmDefaultSm100
+
+    class GemmDefaultRowScaleSm100(GemmDefaultSm100):
+        @cute.jit
+        def epi_visit_subtile(
+            self,
+            params,
+            epi_loop_tensors,
+            tRS_rD: cute.Tensor,
+            tRS_rC: Optional[cute.Tensor] = None,
+        ) -> Optional[cute.Tensor]:
+            # alpha (and beta/C) exactly as stock GemmDefaultEpiMixin, but the
+            # stock colvec ADD is skipped in favor of the multiply below.
+            rD = tRS_rD.load()
+            if const_expr(hasattr(params, "alpha") and params.alpha is not None):
+                rD *= utils.load_scalar_or_pointer(params.alpha)
+            if const_expr(tRS_rC is not None):
+                if const_expr(not hasattr(params, "beta") or params.beta is None):
+                    rD += tRS_rC.load().to(tRS_rD.element_type)  # type: ignore[union-attr]
+                else:
+                    beta = utils.load_scalar_or_pointer(params.beta)
+                    rD += beta * tRS_rC.load().to(tRS_rD.element_type)  # type: ignore[union-attr]
+            tRS_rD.store(rD)
+            vec_multiply(self, tRS_rD, epi_loop_tensors.get("mColVecBroadcast"), None)
+            return None
+
+    return GemmDefaultRowScaleSm100
+
+
+@functools.lru_cache(maxsize=1)
+def _gated_auxadd_gemm_cls():
+    """``GemmGatedSm100`` extended with an exact per-row colvec MULTIPLY (as in
+    :func:`_rowscale_gemm_cls`) and a preact-space ``TileLoad`` aux ADD, both on
+    the fp32 accumulator BEFORE the gated activation and the preact D store:
+    ``postact = act(colvec * (alpha * acc) + aux)``. The aux rides the epilogue
+    C load pipeline; its fragment is partitioned against the same register
+    layout as D, so the add is element-wise in the interleaved preact space."""
+    from typing import Callable, NamedTuple, Optional
+
+    import cutlass
+    import cutlass.cute as cute
+    import quack.utils as utils
+    from cutlass import Float32, Int32, const_expr
+    from quack.cute_dsl_utils import mlir_namedtuple
+    from quack.epi_ops import (
+        ColVecLoad,
+        RowVecLoad,
+        Scalar,
+        TileLoad,
+        TileStore,
+        vec_multiply,
+    )
+    from quack.gemm_act import GemmGatedSm100, _gated_epi_tile_fn
+    from quack.rounding import RoundingMode
+
+    # Functional NamedTuple form: this module uses `from __future__ import
+    # annotations`, so a class-body NamedTuple would store STRING annotations
+    # that quack's Constexpr converter later fails to eval (get_type_hints
+    # resolves them in THIS module's globals, where cute/cutlass are factory
+    # locals). The functional form stores real type objects.
+    _EpiArgs = NamedTuple(
+        "EpilogueArguments",
+        [
+            ("mAuxOut", cute.Tensor),
+            ("act_fn", cutlass.Constexpr[Optional[Callable]]),
+            ("alpha", Optional[Float32 | cute.Tensor]),
+            ("beta", Optional[Float32 | cute.Tensor]),
+            ("mRowVecBroadcast", Optional[cute.Tensor]),
+            ("mColVecBroadcast", Optional[cute.Tensor]),
+            ("mAuxAdd", Optional[cute.Tensor]),
+            ("rounding_mode", cutlass.Constexpr[int]),
+            ("sr_seed", Optional[Int32 | cute.Tensor]),
+        ],
+    )
+    _EpiArgs.__new__.__defaults__ = (  # all fields after mAuxOut
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        RoundingMode.RN,
+        None,
+    )
+    _EpiArgs = mlir_namedtuple(_EpiArgs)
+
+    class GemmGatedAuxAddSm100(GemmGatedSm100):
+        # Stock GemmGatedMixin ops plus the aux TileLoad; __init_subclass__
+        # regenerates EpilogueParams from this tuple.
+        _epi_ops = (
+            Scalar("alpha"),
+            Scalar("beta"),
+            Scalar("sr_seed", dtype=Int32),
+            RowVecLoad("mRowVecBroadcast"),
+            ColVecLoad("mColVecBroadcast"),
+            TileStore("mAuxOut", epi_tile_fn=_gated_epi_tile_fn),
+            TileLoad("mAuxAdd"),
+        )
+
+        # NOTE: the aux tensor must be INTERLEAVED in memory. Applying quack's
+        # concat_to_interleave view to the TileLoad instead (the trick the
+        # mainloop uses for B) device-crashes with an illegal instruction: the
+        # epilogue TMA-load path does not survive the hierarchical N mode
+        # (probed on B200; mainloop B and the TileStore are fine).
+        EpilogueArguments = _EpiArgs
+
+        @cute.jit
+        def epi_visit_subtile(
+            self,
+            params,
+            epi_loop_tensors,
+            tRS_rD: cute.Tensor,
+            tRS_rC: Optional[cute.Tensor] = None,
+        ) -> Optional[cute.Tensor]:
+            # alpha (and beta/C) exactly as stock GemmDefaultEpiMixin, but the
+            # stock colvec ADD is replaced by the exact multiply.
+            rD = tRS_rD.load()
+            if const_expr(hasattr(params, "alpha") and params.alpha is not None):
+                rD *= utils.load_scalar_or_pointer(params.alpha)
+            if const_expr(tRS_rC is not None):
+                if const_expr(not hasattr(params, "beta") or params.beta is None):
+                    rD += tRS_rC.load().to(tRS_rD.element_type)  # type: ignore[union-attr]
+                else:
+                    beta = utils.load_scalar_or_pointer(params.beta)
+                    rD += beta * tRS_rC.load().to(tRS_rD.element_type)  # type: ignore[union-attr]
+            tRS_rD.store(rD)
+            vec_multiply(self, tRS_rD, epi_loop_tensors.get("mColVecBroadcast"), None)
+            tDrAux = epi_loop_tensors.get("mAuxAdd")
+            if const_expr(tDrAux is not None):
+                rD2 = tRS_rD.load()
+                rD2 += tDrAux.load().to(tRS_rD.element_type)  # type: ignore[union-attr]
+                tRS_rD.store(rD2)
+            # Gated activation, verbatim from GemmGatedMixin's SM100 branch.
+            tRS_rAuxOut_layout = cute.recast_layout(2, 1, tRS_rD.layout)
+            tRS_rAuxOut = cute.make_rmem_tensor(
+                tRS_rAuxOut_layout.shape, self.acc_dtype
+            )
+            for i in cutlass.range(cute.size(tRS_rAuxOut) // 2, unroll_full=True):
+                tRS_rAuxOut[2 * i], tRS_rAuxOut[2 * i + 1] = params.act_fn(
+                    (tRS_rD[4 * i], tRS_rD[4 * i + 2]),
+                    (tRS_rD[4 * i + 1], tRS_rD[4 * i + 3]),
+                )
+            return tRS_rAuxOut
+
+    return GemmGatedAuxAddSm100
+
+
+def _compile_kernel(
+    n: int,
+    k: int,
+    num_experts: int,
+    gated: bool,
+    activation: str,
+    has_d: bool,
+    tile_mn: tuple,
+    cluster_mn: tuple,
+    varlen_m: bool,
+    add_to_output: bool,
+    has_colvec: bool,
+    has_aux: bool,
+    concat_b: bool,
+    sfa_sample: torch.Tensor,
+    sfb_sample: torch.Tensor,
+):
+    """Compile once per static shape family; M (or total_m) stays dynamic."""
+    from functools import partial
+
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass import Float32
+    from quack.activation import gate_fn_map
+    from quack.blockscaled_gemm_utils import _make_compile_tensor_like
+    from quack.compile_utils import make_fake_tensor as fake_tensor
+    from quack.cute_dsl_utils import get_device_capacity, get_max_active_clusters
+    from quack.gemm_act import GemmGatedSm100
+    from quack.gemm_default_epi import GemmDefaultSm100
+    from quack.gemm_tvm_ffi_utils import (
+        compile_gemm_kernel,
+        make_fake_scheduler_args,
+        make_fake_varlen_args,
+        make_scheduler_args,
+        make_varlen_args,
+    )
+    from quack.rounding import RoundingMode
+    from quack.varlen_utils import VarlenArguments
+
+    fp4 = cutlass.Float4E2M1FN
+    bf16 = cutlass.BFloat16
+    e4m3 = cutlass.Float8E4M3FN
+    device_capacity = get_device_capacity(sfa_sample.device)
+
+    # All-symbolic fakes (mixing concrete dims with sym M crashes the varlen
+    # ragged-TMA construction with std::bad_variant_access). The fp4 K dim
+    # needs an explicit even-divisibility hint or the packed-nibble shape
+    # check ("stride=1 dim must be divisible by 2") rejects the compile.
+    m_sym = cute.sym_int()
+    n_sym = cute.sym_int(divisibility=8)
+    k_sym = cute.sym_int(divisibility=32)
+    l_sym = cute.sym_int()
+    pa_sym = cute.sym_int(divisibility=8)
+    a_shape = (m_sym, k_sym) if varlen_m else (m_sym, k_sym, l_sym)
+    d_shape = (m_sym, n_sym) if varlen_m else (m_sym, n_sym, l_sym)
+    pa_shape = (m_sym, pa_sym) if varlen_m else (m_sym, pa_sym, l_sym)
+    mA = fake_tensor(fp4, a_shape, leading_dim=1, divisibility=32)
+    mB = fake_tensor(fp4, (n_sym, k_sym, l_sym), leading_dim=1, divisibility=32)
+    mD = fake_tensor(bf16, d_shape, leading_dim=1, divisibility=8) if has_d else None
+    mAux = fake_tensor(bf16, pa_shape, leading_dim=1, divisibility=8) if gated else None
+
+    assert not concat_b or gated, "concat_b (concat-layout weights) is gated-only"
+    if gated:
+        assert not add_to_output, "add_to_output is only wired for the default epilogue"
+        assert not (has_colvec or has_aux) or varlen_m, (
+            "gated colvec / aux add are only wired for varlen_m"
+        )
+        gated_cls = (
+            _gated_auxadd_gemm_cls() if (has_colvec or has_aux) else GemmGatedSm100
+        )
+        gemm_cls = partial(gated_cls, sf_vec_size=SF_VEC_SIZE)
+        mColVec = (
+            fake_tensor(Float32, (m_sym,), leading_dim=0, divisibility=4)
+            if has_colvec
+            else None
+        )
+        # The aux shares mD's fake (same total_m/N and n-major layout): the
+        # interleaved preact-space LoRA delta added before the activation.
+        mAuxAdd = (
+            fake_tensor(bf16, d_shape, leading_dim=1, divisibility=8)
+            if has_aux
+            else None
+        )
+        if has_colvec or has_aux:
+            compile_epi_args = gated_cls.EpilogueArguments(
+                mAux,
+                gate_fn_map[activation],
+                alpha=Float32(0.0),
+                mColVecBroadcast=mColVec,
+                mAuxAdd=mAuxAdd,
+                rounding_mode=RoundingMode.RN,
+            )
+        else:
+            compile_epi_args = GemmGatedSm100.EpilogueArguments(
+                mAux,
+                gate_fn_map[activation],
+                alpha=Float32(0.0),
+                rounding_mode=RoundingMode.RN,
+            )
+    else:
+        assert not has_aux, "aux add is only wired for the gated epilogue"
+        assert not has_colvec or varlen_m, "colvec row scale requires varlen_m"
+        default_cls = _rowscale_gemm_cls() if has_colvec else GemmDefaultSm100
+        gemm_cls = partial(default_cls, sf_vec_size=SF_VEC_SIZE)
+        # Passing a fake mColVecBroadcast at compile time is what activates the
+        # ColVecLoad op (host-side filtering happens during the compile trace);
+        # it shares m_sym with mD, the 1-D (total_m,) varlen colvec convention.
+        mColVec = (
+            fake_tensor(Float32, (m_sym,), leading_dim=0, divisibility=4)
+            if has_colvec
+            else None
+        )
+        # add_to_output is a Constexpr: True swaps the D TMA store for a
+        # reduce-add, so the kernel accumulates into the caller's out buffer.
+        compile_epi_args = GemmDefaultSm100.EpilogueArguments(
+            alpha=Float32(0.0),
+            mColVecBroadcast=mColVec,
+            add_to_output=add_to_output,
+            rounding_mode=RoundingMode.RN,
+        )
+
+    # Zero-copy concat weights: the kernel VIEWS B's concat [gate; up] rows as
+    # interleaved (hierarchical (2, N/2) layout, quack's own gated-MLP scheme).
+    # SFB must then be packed from row-permuted scales (set_weights handles
+    # that). The aux TileLoad is NOT listed: it must arrive interleaved in
+    # memory (the view trick crashes the epilogue load path, see the subclass).
+    concat_layout = ("B",) if concat_b else None
+
+    compiled = compile_gemm_kernel(
+        gemm_cls,
+        fp4,
+        tuple(tile_mn),
+        (cluster_mn[0], cluster_mn[1], 1),
+        False,  # pingpong (SM90-only knob)
+        True,  # persistent
+        False,  # gather_A: unsupported with blockscaled (SF gather missing upstream)
+        True,  # is_dynamic_persistent -> use_clc_persistence on SM100
+        device_capacity,
+        mA,
+        mB,
+        mD,
+        None,
+        compile_epi_args,
+        make_fake_scheduler_args(False, False, cute.sym_int()),
+        make_fake_varlen_args(varlen_m, False, False, None) or VarlenArguments(),
+        mSFA=_make_compile_tensor_like(sfa_sample, e4m3, dynamic_layout=True),
+        mSFB=_make_compile_tensor_like(sfb_sample, e4m3, dynamic_layout=True),
+        concat_layout=concat_layout,
+    )
+
+    max_active = get_max_active_clusters(cluster_mn[0] * cluster_mn[1])
+
+    def run(a, b, d, postact, sfa, sfb, cu_seqlens, alpha, colvec=None, aux=None):
+        if gated:
+            # Tensor fields compiled non-None must be non-None on every call.
+            assert (colvec is not None) == has_colvec
+            assert (aux is not None) == has_aux
+            if has_colvec or has_aux:
+                epi_args = gated_cls.EpilogueArguments(
+                    postact,
+                    None,
+                    alpha=Float32(alpha),
+                    mColVecBroadcast=colvec,
+                    mAuxAdd=aux,
+                    rounding_mode=None,
+                )
+            else:
+                epi_args = GemmGatedSm100.EpilogueArguments(
+                    postact, None, alpha=Float32(alpha), rounding_mode=None
+                )
+        else:
+            # Tensor fields compiled non-None must be non-None on every call.
+            assert (colvec is not None) == has_colvec
+            # Constexpr fields (add_to_output, rounding_mode) must be None at
+            # call time; the namedtuple default False fails the FFI signature.
+            epi_args = GemmDefaultSm100.EpilogueArguments(
+                alpha=Float32(alpha),
+                mColVecBroadcast=colvec,
+                add_to_output=None,
+                rounding_mode=None,
+            )
+        scheduler_args = make_scheduler_args(max_active, 8, None)
+        varlen_args = make_varlen_args(cu_seqlens, None, None) or VarlenArguments()
+        compiled(a, b, d, None, epi_args, scheduler_args, varlen_args, sfa, sfb)
+
+    return run
+
+
+def _get_run(key, sfa_sample, sfb_sample):
+    if key not in _COMPILE_CACHE:
+        _COMPILE_CACHE[key] = _compile_kernel(*key[:-1], sfa_sample, sfb_sample)
+    return _COMPILE_CACHE[key]
+
+
+def _as_fp4x2(t: torch.Tensor) -> torch.Tensor:
+    return t.view(torch.float4_e2m1fn_x2) if t.dtype == torch.uint8 else t
+
+
+def _check_dims(n: int, k: int, gated: bool, tile_mn: tuple) -> None:
+    assert k % 32 == 0, f"K={k} must be divisible by 32 (fp4 16B alignment)"
+    assert n % 8 == 0, f"N={n} must be divisible by 8 (bf16 output alignment)"
+    assert tile_mn[0] in (128, 256) and tile_mn[1] in (64, 128, 192, 256), (
+        f"blockscaled tile must be M in (128,256), N in (64,128,192,256), got {tile_mn}"
+    )
+    if gated:
+        assert n % 2 == 0 and (n // 2) % 8 == 0, f"gated N={n} needs (N/2) % 8 == 0"
+
+
+class GroupedNvfp4Gemm:
+    """Grouped (varlen_m) NVFP4 W4A4 GEMM engine, one weight slice per expert.
+
+    Build once per (N, K, E, gated, activation, tile); ``set_weights`` once per
+    weight tensor; ``forward`` per step with dynamic total_m.
+    """
+
+    def __init__(
+        self,
+        n: int,
+        k: int,
+        num_experts: int,
+        *,
+        gated: bool = False,
+        activation: str = "swiglu",
+        store_preact: bool = True,
+        concat_b: bool = False,
+        tile_mn: tuple = (128, 128),
+        cluster_mn: tuple = (1, 1),
+    ):
+        """``concat_b=True`` (gated only): ``set_weights`` takes CONCAT
+        [gate; up] rows and the kernel views them interleaved zero-copy
+        (quack ``concat_layout``); only the small block-scale copy is
+        row-permuted for SFB. The preact D still comes out INTERLEAVED."""
+        _check_dims(n, k, gated, tile_mn)
+        assert not concat_b or gated, "concat_b is gated-only"
+        self.n, self.k, self.num_experts = n, k, num_experts
+        self.gated = gated
+        self.activation = activation
+        self.store_preact = store_preact
+        self.concat_b = concat_b
+        self.tile_mn = tuple(tile_mn)
+        self.cluster_mn = tuple(cluster_mn)
+        self._b_operand = None
+        self._sfb = None
+        self._runs: dict = {}
+
+    def set_weights(
+        self,
+        qdata: torch.Tensor,
+        block_scale: torch.Tensor,
+        per_tensor_scale: torch.Tensor | None = None,
+    ) -> None:
+        """qdata ``(E, N, K/2)`` uint8 K-packed; block_scale ``(E, N, K/16)`` e4m3.
+
+        Gated engines expect gate/up rows already INTERLEAVED (apply
+        ``gate_up_interleave_perm`` to both tensors for concat-layout weights)
+        UNLESS built with ``concat_b=True``, in which case both tensors stay in
+        CONCAT [gate; up] layout: qdata is consumed zero-copy through the
+        kernel's interleaved view, and only the block scales are row-permuted
+        here so SFB rows follow the kernel's logical (interleaved) N order.
+        per_tensor_scale ``(E,)``/``(E,1,1)``/scalar fp32 is folded into SFB.
+        """
+        e, n, k2 = qdata.shape
+        assert (e, n, k2 * 2) == (self.num_experts, self.n, self.k), (
+            f"qdata {tuple(qdata.shape)} vs engine (E={self.num_experts}, N={self.n}, K={self.k})"
+        )
+        assert block_scale.shape == (e, n, self.k // SF_VEC_SIZE)
+        if self.concat_b:
+            perm = gate_up_interleave_perm(n, device=block_scale.device)
+            block_scale = block_scale.index_select(1, perm)
+        if per_tensor_scale is not None:
+            block_scale, _ = fold_per_tensor_scale(block_scale, per_tensor_scale)
+        # (E, N, K/2) contiguous -> (N, K/2, E) K-major view, as the kernel wants.
+        self._b_operand = _as_fp4x2(qdata.contiguous().permute(1, 2, 0))
+        self._sfb = pack_scales_blocked(block_scale)
+
+    def forward(
+        self,
+        a_packed: torch.Tensor,
+        sfa_blocked: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        *,
+        alpha: float = 1.0,
+        colvec: torch.Tensor | None = None,
+        aux: torch.Tensor | None = None,
+        preact_out: torch.Tensor | None = None,
+        out: torch.Tensor | None = None,
+        add_to_output: bool = False,
+    ):
+        """a_packed ``(total_m, K/2)`` uint8/fp4x2, expert-sorted, unpadded.
+        sfa_blocked from ``sf_layout.build_varlen_sfa``. cu_seqlens ``(E+1,)``.
+
+        ``colvec`` ``(total_m,)`` fp32 multiplies the fp32 accumulator per row
+        in the epilogue: ``d = colvec * (alpha * acc)``, exact (single bf16
+        rounding at the D store). On gated engines it applies before the
+        activation and the preact store.
+
+        ``aux`` ``(total_m, N)`` bf16 (gated only) is ADDED to the fp32
+        accumulator after the colvec multiply, before the activation and the
+        preact store: ``preact = colvec * (alpha * acc) + aux`` in the
+        INTERLEAVED gate/up space.
+
+        ``add_to_output=True`` (non-gated only) accumulates ``alpha * acc``
+        (times ``colvec`` if given) into a caller-provided ``out`` instead of
+        overwriting it.
+
+        Returns ``(postact [total_m, N/2] bf16, preact [total_m, N] bf16 | None)``
+        for gated engines, else ``out [total_m, N]`` bf16.
+        """
+        assert self._b_operand is not None, "call set_weights first"
+        if add_to_output:
+            assert not self.gated and out is not None
+        total_m = a_packed.shape[0]
+        if colvec is not None:
+            assert colvec.dtype == torch.float32 and colvec.shape == (total_m,)
+            colvec = colvec.contiguous()
+        if aux is not None:
+            assert self.gated, "aux add is gated-engine only"
+            assert aux.dtype == torch.bfloat16 and aux.shape == (total_m, self.n)
+            aux = aux.contiguous()
+        assert a_packed.shape[1] * 2 == self.k
+        device = a_packed.device
+        a = _as_fp4x2(a_packed)
+        cu = cu_seqlens.to(device=device, dtype=torch.int32)
+
+        d = None
+        postact = None
+        if self.gated:
+            postact = torch.empty(
+                total_m, self.n // 2, dtype=torch.bfloat16, device=device
+            )
+            if self.store_preact:
+                d = (
+                    preact_out
+                    if preact_out is not None
+                    else torch.empty(
+                        total_m, self.n, dtype=torch.bfloat16, device=device
+                    )
+                )
+        else:
+            d = (
+                out
+                if out is not None
+                else torch.empty(total_m, self.n, dtype=torch.bfloat16, device=device)
+            )
+
+        run_key = (add_to_output, colvec is not None, aux is not None)
+        run = self._runs.get(run_key)
+        if run is None:
+            key = (
+                self.n,
+                self.k,
+                self.num_experts,
+                self.gated,
+                self.activation,
+                d is not None,
+                self.tile_mn,
+                self.cluster_mn,
+                True,  # varlen_m
+                add_to_output,
+                colvec is not None,
+                aux is not None,
+                self.concat_b,
+                device.index,
+            )
+            run = _get_run(key, self._sfb.new_zeros(1, 1, 1, 512), self._sfb)
+            self._runs[run_key] = run
+        run(
+            a,
+            self._b_operand,
+            d,
+            postact,
+            sfa_blocked,
+            self._sfb,
+            cu,
+            float(alpha),
+            colvec,
+            aux,
+        )
+        return (postact, d) if self.gated else d
+
+
+def dense_nvfp4_gemm(
+    a_q: torch.Tensor,
+    a_scale: torch.Tensor,
+    b_q: torch.Tensor,
+    b_scale: torch.Tensor,
+    *,
+    gated: bool = False,
+    activation: str = "swiglu",
+    alpha: float = 1.0,
+    store_preact: bool = True,
+    tile_mn: tuple = (128, 128),
+    cluster_mn: tuple = (1, 1),
+):
+    """Batched dense (non-varlen) NVFP4 GEMM.
+
+    a_q ``(L, M, K/2)`` uint8, a_scale ``(L, M, K/16)`` e4m3; b likewise with N.
+    Returns ``(postact (L,M,N/2), preact (L,M,N) | None)`` if gated else
+    ``out (L, M, N)``, all bf16.
+    """
+    num_l, m, k2 = a_q.shape
+    k = 2 * k2
+    n = b_q.shape[1]
+    _check_dims(n, k, gated, tile_mn)
+    device = a_q.device
+
+    a = _as_fp4x2(a_q.contiguous().permute(1, 2, 0))  # (M, K/2, L) K-major
+    b = _as_fp4x2(b_q.contiguous().permute(1, 2, 0))
+    sfa = pack_scales_blocked(a_scale)
+    sfb = pack_scales_blocked(b_scale)
+
+    d3 = postact3 = None
+    d = postact = None
+    if gated:
+        postact3 = torch.empty(num_l, m, n // 2, dtype=torch.bfloat16, device=device)
+        postact = postact3.permute(1, 2, 0)
+        if store_preact:
+            d3 = torch.empty(num_l, m, n, dtype=torch.bfloat16, device=device)
+            d = d3.permute(1, 2, 0)
+    else:
+        d3 = torch.empty(num_l, m, n, dtype=torch.bfloat16, device=device)
+        d = d3.permute(1, 2, 0)
+
+    key = (
+        n,
+        k,
+        num_l,
+        gated,
+        activation,
+        d is not None,
+        tuple(tile_mn),
+        tuple(cluster_mn),
+        False,  # varlen_m
+        False,  # add_to_output
+        False,  # has_colvec
+        False,  # has_aux
+        False,  # concat_b
+        device.index,
+    )
+    run = _get_run(key, sfa, sfb)
+    run(a, b, d, postact, sfa, sfb, None, float(alpha))
+    return (postact3, d3) if gated else d3
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/fp4_cute_ops.py b/src/axolotl/integrations/kernels/libs/sonicmoe/fp4_cute_ops.py
new file mode 100644
index 0000000000..28d648539d
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/fp4_cute_ops.py
@@ -0,0 +1,363 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""``backend="fp4_cute"`` ops: torchao NVFP4 base weights on the SM100 W4A4 engine.
+
+Bridges the grouped MoE-LoRA path (``nvfp4.py`` / ``nvfp4_lora.py``, bf16
+expert-sorted rows) to ``fp4_cute.GroupedNvfp4Gemm``:
+
+- unpacks a torchao ``NVFP4Tensor`` into the (qdata, block_scale, pts)
+  components the engine consumes, with a per-weight engine cache (the base is
+  frozen, so weights are packed into the engine exactly once);
+- quantizes activations to NVFP4 on the expert-sorted rows (per-16 block
+  scales, per-tensor scale 1.0) and builds the dQaccum-padded SFA;
+- handles the per-expert weight ``per_tensor_scale`` EXACTLY: the engine keeps
+  the stored e4m3 block scales and multiplies the full pts values into the
+  fp32 accumulator in the GEMM epilogue, via a per-row fp32 colvec
+  (``pts`` repeated over each expert's rows, sync-free). Single bf16 rounding
+  at the D store, and it composes with ``add_to_output`` (the multiply runs
+  before the reduce-add). The old lossy scheme (fold the pts_e/pts_ref RATIOS
+  into SFB, ``alpha = pts_ref``; re-rounds SFB in e4m3) stays reachable via
+  ``AXOLOTL_SONICMOE_NVFP4_PTS_FOLD=1`` for A/B numerics debugging;
+- the backward dequant (:func:`dequantize_engine_weight`) therefore uses the
+  plain stored-scales-times-pts semantics (``dequantize_expert_weight``), so
+  the dense weights the dX path sees are identical to the operands the forward
+  kernel consumed; with the fold env set it uses the folded scales and alpha;
+- wraps the grouped GEMM in an autograd.Function whose backward computes dX by
+  per-expert chunked dequant matmuls, never through the packed fp4 operand.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import NamedTuple, Optional
+
+import torch
+
+from .fp4_cute import GroupedNvfp4Gemm
+from .nvfp4 import dequantize_expert_slice, is_nvfp4_param
+from .nvfp4_quant import quantize_nvfp4_ref
+from .sf_layout import build_varlen_sfa, fold_per_tensor_scale
+
+
+def unpack_nvfp4_components(w) -> tuple:
+    """``(qdata u8 [E, N, K/2], block_scale e4m3 [E, N, K/16], pts [E,1,1] | None)``."""
+    if not is_nvfp4_param(w):
+        raise TypeError(f"expected a torchao NVFP4Tensor, got {type(w).__name__}")
+    assert not getattr(w, "is_swizzled_scales", False), (
+        "swizzled NVFP4 scales unsupported (our loaders emit row-major)"
+    )
+    qdata = w.qdata
+    if qdata.dtype != torch.uint8:
+        qdata = qdata.view(torch.uint8)
+    return qdata, w.scale, w.per_tensor_scale
+
+
+def fp4_cute_dims_ok(w1, w2) -> bool:
+    """Both grouped GEMMs run the non-gated engine: K % 32 and N % 8.
+
+    ``w1`` ``[E, 2I, H]`` (K=H), ``w2`` ``[E, H, I]`` (K=I); shapes are logical
+    (unpacked) as the tensor subclass reports them.
+    """
+    n1, k1 = w1.shape[-2], w1.shape[-1]
+    n2, k2 = w2.shape[-2], w2.shape[-1]
+    return k1 % 32 == 0 and n1 % 8 == 0 and k2 % 32 == 0 and n2 % 8 == 0
+
+
+class _EngineEntry(NamedTuple):
+    engine: GroupedNvfp4Gemm
+    # alpha scales the fp32 accumulator. colvec_pts [E] fp32 means the engine
+    # keeps the stored scales and pts is applied exactly per row in the
+    # epilogue; folded_scale means SFB holds the pts_e/pts_ref ratios (old
+    # scheme). fused=True means no post-GEMM row scaling is needed.
+    alpha: float
+    folded_scale: Optional[torch.Tensor]
+    fold_rel_err: float
+    fused: bool
+    colvec_pts: Optional[torch.Tensor]
+
+
+# Keyed by the packed storage; safe against pointer reuse because each engine
+# holds a view of its weight's qdata, keeping that storage alive.
+_ENGINE_CACHE: dict = {}
+
+
+def _engine_key(qdata: torch.Tensor) -> tuple:
+    return (qdata.data_ptr(), tuple(qdata.shape), qdata.device.index)
+
+
+def _get_engine(weight) -> _EngineEntry:
+    qdata, scale, pts = unpack_nvfp4_components(weight)
+    key = _engine_key(qdata)
+    entry = _ENGINE_CACHE.get(key)
+    if entry is None:
+        e, n, k2 = qdata.shape
+        engine = GroupedNvfp4Gemm(n, k2 * 2, e, gated=False)
+        alpha, folded, rel_err, colvec_pts = 1.0, None, 0.0, None
+        # Default (unset/0): exact scheme, stored scales + per-row pts colvec
+        # multiplied into the fp32 accumulator in the epilogue. "1": old lossy
+        # SFB ratio fold (re-rounds SFB in e4m3, ~6% max block-scale err on
+        # real checkpoints), kept for A/B numerics debugging.
+        fold_enabled = os.environ.get("AXOLOTL_SONICMOE_NVFP4_PTS_FOLD", "0") == "1"
+        if pts is not None:
+            if fold_enabled:
+                # One-time host sync per weight, at engine build.
+                pts_ref = float(pts.float().max())
+                if pts_ref > 0:
+                    try:
+                        folded, rel_err = fold_per_tensor_scale(
+                            scale,
+                            pts.float().reshape(-1) / pts_ref,
+                            allow_underflow=True,
+                        )
+                        alpha = pts_ref
+                    except ValueError:
+                        folded = None
+            else:
+                colvec_pts = pts.float().reshape(-1).contiguous()
+        engine.set_weights(qdata, scale if folded is None else folded)
+        fused = pts is None or folded is not None or colvec_pts is not None
+        entry = _EngineEntry(engine, alpha, folded, rel_err, fused, colvec_pts)
+        _ENGINE_CACHE[key] = entry
+    return entry
+
+
+class _GatedEngineEntry(NamedTuple):
+    engine: GroupedNvfp4Gemm
+    colvec_pts: Optional[torch.Tensor]
+
+
+_GATED_ENGINE_CACHE: dict = {}
+
+
+def _get_gated_engine(weight, activation: str = "swiglu") -> _GatedEngineEntry:
+    """Gated (fused-activation) engine for a concat-layout ``[E, 2I, H]`` NVFP4
+    up-projection weight. ``concat_b=True``: the packed qdata is consumed
+    zero-copy (the kernel views concat rows as interleaved); pts is always
+    applied exactly via the per-row colvec (no fold variant here)."""
+    qdata, scale, pts = unpack_nvfp4_components(weight)
+    key = _engine_key(qdata)
+    entry = _GATED_ENGINE_CACHE.get(key)
+    if entry is None:
+        e, n, k2 = qdata.shape
+        engine = GroupedNvfp4Gemm(
+            n, k2 * 2, e, gated=True, activation=activation, concat_b=True
+        )
+        engine.set_weights(qdata, scale)
+        colvec_pts = pts.float().reshape(-1).contiguous() if pts is not None else None
+        entry = _GatedEngineEntry(engine, colvec_pts)
+        _GATED_ENGINE_CACHE[key] = entry
+    return entry
+
+
+def gated_nvfp4_forward(
+    x_grouped: torch.Tensor,
+    weight,
+    cu_seqlens: torch.Tensor,
+    aux: Optional[torch.Tensor],
+    activation: str = "swiglu",
+) -> tuple:
+    """Fused up-GEMM + gated activation (+ optional preact aux add).
+
+    ``aux`` ``[T, 2I]`` bf16 in INTERLEAVED gate/up layout (the LoRA delta;
+    compute it against row-permuted LoRA-B factors) is added to the fp32
+    accumulator after the exact per-row pts multiply, before the activation.
+    Returns ``(postact [T, I] bf16, preact [T, 2I] bf16)`` where the preact
+    memory layout is also INTERLEAVED (``preact.view(T, I, 2)`` puts gate at
+    ``[..., 0]`` and up at ``[..., 1]``). No autograd."""
+    entry = _get_gated_engine(weight, activation)
+    a_q, sfa = quantize_grouped_rows(x_grouped, cu_seqlens)
+    colvec = None
+    if entry.colvec_pts is not None:
+        counts = (cu_seqlens[1:] - cu_seqlens[:-1]).long()
+        colvec = torch.repeat_interleave(
+            entry.colvec_pts, counts, output_size=x_grouped.shape[0]
+        )
+    return entry.engine.forward(
+        a_q,
+        sfa,
+        cu_seqlens,
+        colvec=colvec,
+        aux=aux,
+    )
+
+
+def dequantize_engine_weight(weight) -> torch.Tensor:
+    """Dense weights matching the operands the fp4_cute forward consumed.
+
+    The default colvec path applies the FULL pts exactly in the forward
+    epilogue, so plain ``dequantize_expert_weight`` (stored scales times pts)
+    already matches; only when the engine folded the pts ratios into SFB
+    (``AXOLOTL_SONICMOE_NVFP4_PTS_FOLD=1``) does this dequantize with the
+    folded scales and alpha instead.
+    """
+    from .nvfp4 import dequantize_expert_weight
+
+    if not is_nvfp4_param(weight):
+        return weight
+    qdata, _, _ = unpack_nvfp4_components(weight)
+    entry = _ENGINE_CACHE.get(_engine_key(qdata))
+    if entry is None or entry.folded_scale is None:
+        return dequantize_expert_weight(weight)
+    from .triton_nvfp4 import dequant_nvfp4_triton, triton_available
+
+    alpha_t = torch.tensor([entry.alpha], dtype=torch.float32, device=qdata.device)
+    if qdata.is_cuda and triton_available():
+        return dequant_nvfp4_triton(
+            qdata, entry.folded_scale, alpha_t, weight.orig_dtype
+        )
+    from .nvfp4_quant import dequantize_nvfp4_ref
+
+    return (dequantize_nvfp4_ref(qdata, entry.folded_scale) * entry.alpha).to(
+        weight.orig_dtype
+    )
+
+
+def quantize_grouped_rows(x_grouped: torch.Tensor, cu_seqlens: torch.Tensor) -> tuple:
+    """NVFP4-quantize expert-sorted rows: ``(a_packed [T, K/2] u8, sfa blocked)``."""
+    from .triton_nvfp4 import quantize_rows_fused_sfa_triton, triton_available
+
+    if (
+        x_grouped.is_cuda
+        and triton_available()
+        and (x_grouped.shape[-1] // 16) % 4 == 0
+    ):
+        return quantize_rows_fused_sfa_triton(x_grouped, cu_seqlens)
+    a_q, a_s, _ = quantize_nvfp4_ref(x_grouped)
+    return a_q, build_varlen_sfa(a_s, cu_seqlens)
+
+
+def _base_gemm_forward(
+    entry: _EngineEntry,
+    weight,
+    x_grouped: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    out: torch.Tensor | None = None,
+    add_to_output: bool = False,
+) -> torch.Tensor:
+    a_q, sfa = quantize_grouped_rows(x_grouped, cu_seqlens)
+    colvec = None
+    if entry.colvec_pts is not None:
+        # Per-row FULL pts, sync-free: counts stay on device, output_size is
+        # known host-side from the packed activation rows.
+        counts = (cu_seqlens[1:] - cu_seqlens[:-1]).long()
+        colvec = torch.repeat_interleave(
+            entry.colvec_pts, counts, output_size=x_grouped.shape[0]
+        )
+    if entry.fused:
+        return entry.engine.forward(
+            a_q,
+            sfa,
+            cu_seqlens,
+            alpha=entry.alpha,
+            colvec=colvec,
+            out=out,
+            add_to_output=add_to_output,
+        )
+    assert not add_to_output
+    result = entry.engine.forward(a_q, sfa, cu_seqlens)
+    pts = weight.per_tensor_scale
+    counts = (cu_seqlens[1:] - cu_seqlens[:-1]).long()
+    row_pts = torch.repeat_interleave(pts.view(-1).float(), counts)
+    from .triton_nvfp4 import rowscale_inplace_triton, triton_available
+
+    if result.is_cuda and triton_available() and result.dtype == x_grouped.dtype:
+        return rowscale_inplace_triton(result.contiguous(), row_pts)
+    return (result.float() * row_pts.unsqueeze(1)).to(x_grouped.dtype)
+
+
+def grouped_dx_dequant(
+    grad_h: torch.Tensor, weight, cu_seqlens: torch.Tensor
+) -> torch.Tensor:
+    """``dx[start:end] = g_e @ W_e``, never through the packed fp4 operand."""
+    from .fp8_bwd import fp8_dx_supported, grouped_fp8_dx
+    from .nvfp4_lora import _use_grouped_mm
+
+    if fp8_dx_supported(grad_h, weight):
+        return grouped_fp8_dx(grad_h, weight, cu_seqlens)
+
+    if _use_grouped_mm(grad_h):
+        w_dense = dequantize_engine_weight(weight).to(grad_h.dtype)
+        return torch._grouped_mm(grad_h, w_dense, offs=cu_seqlens[1:].to(torch.int32))
+
+    cu = cu_seqlens.tolist()
+    dx = grad_h.new_empty((grad_h.shape[0], weight.shape[-1]))
+    qdata, _, _ = unpack_nvfp4_components(weight)
+    entry = _ENGINE_CACHE.get(_engine_key(qdata))
+    folded = entry.folded_scale if entry is not None else None
+    alpha = entry.alpha if entry is not None else 1.0
+    for e in range(len(cu) - 1):
+        start, end = cu[e], cu[e + 1]
+        if end <= start:
+            continue
+        if folded is not None:
+            from .nvfp4_quant import dequantize_nvfp4_ref
+
+            w_e = dequantize_nvfp4_ref(qdata[e], folded[e]) * alpha
+        else:
+            w_e = dequantize_expert_slice(weight, e)
+        dx[start:end] = grad_h[start:end] @ w_e.to(grad_h.dtype)
+    return dx
+
+
+class _GroupedNvfp4Linear(torch.autograd.Function):
+    """``y_e = quant16(x_e) @ W_e^T`` with a frozen NVFP4 ``W``.
+
+    Backward treats the activation quantization as identity (straight-through)
+    and computes ``dx`` by chunked dequant matmuls; ``W`` gets no gradient.
+    """
+
+    @staticmethod
+    def forward(ctx, x_grouped, weight, cu_seqlens):
+        entry = _get_engine(weight)
+        out = _base_gemm_forward(entry, weight, x_grouped, cu_seqlens)
+        ctx.save_for_backward(weight, cu_seqlens)
+        return out.to(x_grouped.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        weight, cu_seqlens = ctx.saved_tensors
+        dx = grouped_dx_dequant(grad_out.contiguous(), weight, cu_seqlens)
+        return dx, None, None
+
+
+def _cu_seqlens(end_offsets: torch.Tensor) -> torch.Tensor:
+    return torch.cat([end_offsets.new_zeros(1), end_offsets]).to(torch.int32)
+
+
+def grouped_nvfp4_linear(
+    x_grouped: torch.Tensor, weight, end_offsets: torch.Tensor
+) -> torch.Tensor:
+    """Per-expert ``x_e @ w[e]^T`` on the SM100 engine.
+
+    ``end_offsets`` ``[E]`` are cumulative row ends (``_grouped_gemm``'s
+    convention); the leading zero is prepended to form ``cu_seqlens``.
+    """
+    assert x_grouped.shape[-1] == weight.shape[-1], (
+        f"x K={x_grouped.shape[-1]} vs weight K={weight.shape[-1]}"
+    )
+    return _GroupedNvfp4Linear.apply(x_grouped, weight, _cu_seqlens(end_offsets))
+
+
+def grouped_nvfp4_linear_add_delta(
+    x_grouped: torch.Tensor,
+    weight,
+    end_offsets: torch.Tensor,
+    delta: torch.Tensor,
+) -> torch.Tensor:
+    """``x_e @ w[e]^T + delta``, accumulated in the GEMM epilogue when possible.
+
+    No autograd: callers (the grouped LoRA autograd.Functions) own the
+    backward. ``delta`` ``[T, N]`` may be overwritten and returned.
+    """
+    assert x_grouped.shape[-1] == weight.shape[-1], (
+        f"x K={x_grouped.shape[-1]} vs weight K={weight.shape[-1]}"
+    )
+    cu = _cu_seqlens(end_offsets)
+    entry = _get_engine(weight)
+    if entry.fused and delta.dtype == torch.bfloat16:
+        out = delta.contiguous()
+        _base_gemm_forward(entry, weight, x_grouped, cu, out=out, add_to_output=True)
+        return out
+    return _base_gemm_forward(entry, weight, x_grouped, cu) + delta
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/fp8_bwd.py b/src/axolotl/integrations/kernels/libs/sonicmoe/fp8_bwd.py
new file mode 100644
index 0000000000..cea45cae0b
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/fp8_bwd.py
@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Optional fp8 backward-dX engine over DeepGEMM for the frozen NVFP4 base.
+
+Opt-in via ``AXOLOTL_SONICMOE_NVFP4_BWD=deepgemm``. The default bf16 backward
+dequantizes each expert weight every pass and runs ``torch._grouped_mm``; this
+caches the weights once as fp8 e4m3 (from the same folded dequant the forward
+consumed) and runs dX through DeepGEMM's m-grouped contiguous fp8 GEMM. About
+1.5-1.8x over the bf16 grouped GEMM at Qwen3-30B backward shapes (dX rel err
+~7e-4), but the cache is a full fp8 copy of every expert weight (+27 GiB on
+Qwen3-30B), so it stays off by default.
+
+``torch._scaled_grouped_mm`` is deliberately not a fallback: its rowwise fp8
+grouped kernel device-aborts on SM100 (torch 2.10+cu130, not built for sm100a)
+and kills the CUDA context. DeepGEMM's contiguous layout needs every expert
+segment padded to its M block alignment (128); padded rows carry
+``m_indices == -1`` (skipped) and are gathered back after. All index math stays
+on device (buffer sized ``T + E * align``) so no host sync is needed.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import NamedTuple
+
+import torch
+
+from .nvfp4 import is_nvfp4_param
+
+_DEEP_GEMM = None
+
+
+def _deep_gemm():
+    """Import deep_gemm once; torch must own the CUDA context first."""
+    global _DEEP_GEMM
+    if _DEEP_GEMM is None:
+        try:
+            torch.zeros(1, device="cuda")
+            import deep_gemm
+
+            # 128 over the theoretical 224: every padded-layout cost (quant,
+            # zero rows, output buffer) scales with E * alignment, and the
+            # GEMM measured slightly FASTER at 128 on B200 (fewer pad rows).
+            deep_gemm.set_mk_alignment_for_contiguous_layout(128)
+            _DEEP_GEMM = deep_gemm
+        except Exception:
+            _DEEP_GEMM = False
+    return _DEEP_GEMM if _DEEP_GEMM is not False else None
+
+
+def fp8_dx_supported(grad_h: torch.Tensor, base_weight) -> bool:
+    # explicit opt-in only: the fp8 cache costs a full fp8 copy of the expert
+    # weights (+27 GiB on Qwen3-30B), unacceptable as a default
+    if os.environ.get("AXOLOTL_SONICMOE_NVFP4_BWD") != "deepgemm":
+        return False
+    if not (
+        grad_h.is_cuda
+        and grad_h.dtype == torch.bfloat16
+        and is_nvfp4_param(base_weight)
+        and base_weight.shape[-2] % 128 == 0  # K of the dX GEMM (dim1)
+        and torch.cuda.get_device_capability(grad_h.device)[0] == 10
+    ):
+        raise RuntimeError(
+            "AXOLOTL_SONICMOE_NVFP4_BWD=deepgemm requires a CUDA bf16 grad, "
+            "a packed NVFP4 base with dim1 % 128 == 0, and an SM100 GPU"
+        )
+    if _deep_gemm() is None:
+        raise RuntimeError(
+            "AXOLOTL_SONICMOE_NVFP4_BWD=deepgemm but deep_gemm is not "
+            "importable; build it from source "
+            "(https://github.com/deepseek-ai/DeepGEMM)"
+        )
+    return True
+
+
+_PAD_QUANT_KERNEL = None
+
+
+def _pad_quant_kernel():
+    """Fused pad + per-token fp8 quant (1x128 ue8m0 scales) in one pass.
+
+    The torch chain (zero-fill the padded buffer, scatter, per_token_cast's
+    fill + copy + amax + mul) moves ~5x the real data and erased the GEMM win;
+    this reads only real rows (padding writes zeros without reading) and its
+    rounding is bit-identical to per_token_cast_to_fp8 because ue8m0 scales
+    are powers of two.
+    """
+    global _PAD_QUANT_KERNEL
+    if _PAD_QUANT_KERNEL is None:
+        import triton
+        import triton.language as tl
+
+        @triton.jit
+        def _kernel(
+            grad_ptr,
+            src_ptr,
+            q_ptr,
+            sf_ptr,
+            N_BLOCKS: tl.constexpr,
+            BLK: tl.constexpr,
+        ):
+            pid = tl.program_id(0)
+            m = pid // N_BLOCKS
+            b = pid % N_BLOCKS
+            src = tl.load(src_ptr + m)
+            offs = b * BLK + tl.arange(0, BLK)
+            if src >= 0:
+                x = tl.load(grad_ptr + src.to(tl.int64) * (N_BLOCKS * BLK) + offs).to(
+                    tl.float32
+                )
+            else:
+                x = tl.zeros([BLK], dtype=tl.float32)
+            amax = tl.maximum(tl.max(tl.abs(x)), 1e-4)
+            sf = amax / 448.0
+            bits = sf.to(tl.int32, bitcast=True)
+            exp = ((bits >> 23) & 0xFF) + tl.where((bits & 0x7FFFFF) != 0, 1, 0)
+            exp = tl.minimum(tl.maximum(exp, 1), 254)
+            sf = (exp << 23).to(tl.float32, bitcast=True)
+            q = x * (1.0 / sf)
+            tl.store(
+                q_ptr + pid.to(tl.int64) * BLK + tl.arange(0, BLK),
+                q.to(q_ptr.dtype.element_ty),
+            )
+            tl.store(sf_ptr + m.to(tl.int64) * N_BLOCKS + b, sf)
+
+        _PAD_QUANT_KERNEL = _kernel
+    return _PAD_QUANT_KERNEL
+
+
+def _pad_and_cast_fp8(
+    grad_h: torch.Tensor, src: torch.Tensor, dest: torch.Tensor, m_max: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(a_fp8 [m_max, K], a_sf fp32 [m_max, K/128])`` from real rows of grad_h.
+
+    ``src`` maps padded row -> source row (-1 for padding); ``dest`` is the
+    inverse map. Falls back to the torch scatter + per_token_cast chain when
+    triton is unavailable.
+    """
+    K = grad_h.shape[1]
+    from .triton_nvfp4 import triton_available
+
+    if triton_available() and K % 128 == 0:
+        q = torch.empty((m_max, K), dtype=torch.float8_e4m3fn, device=grad_h.device)
+        sf = torch.empty((m_max, K // 128), dtype=torch.float32, device=grad_h.device)
+        grad_c = grad_h.contiguous()
+        n_blocks = K // 128
+        _pad_quant_kernel()[(m_max * n_blocks,)](
+            grad_c, src, q, sf, N_BLOCKS=n_blocks, BLK=128
+        )
+        return q, sf
+
+    from deep_gemm.utils import per_token_cast_to_fp8
+
+    a_pad = grad_h.new_zeros((m_max, K))
+    a_pad[dest] = grad_h
+    return per_token_cast_to_fp8(a_pad, use_ue8m0=True)
+
+
+class _Fp8Entry(NamedTuple):
+    # transposed weight [E, dim2, dim1] as fp8 with 128x128-block ue8m0 scales,
+    # so dX = grad_h @ W runs as the NT grouped GEMM A @ B^T
+    w_fp8: torch.Tensor
+    w_sf: torch.Tensor
+    alignment: int
+
+
+# Keyed like fp4_cute_ops._ENGINE_CACHE; each entry only lives as long as the
+# frozen weight it was built from (the qdata view keeps the storage alive).
+_FP8_CACHE: dict = {}
+
+
+def _cache_key(base_weight) -> tuple:
+    qdata = base_weight.qdata
+    return (qdata.data_ptr(), tuple(qdata.shape), qdata.device.index)
+
+
+def _get_fp8_cache(base_weight) -> _Fp8Entry:
+    key = _cache_key(base_weight)
+    entry = _FP8_CACHE.get(key)
+    if entry is None:
+        dg = _deep_gemm()
+        from deep_gemm.utils import ceil_div, per_block_cast_to_fp8
+
+        from .fp4_cute_ops import dequantize_engine_weight
+
+        # one-time per frozen weight; must match the forward's folded operands
+        w = dequantize_engine_weight(base_weight)
+        wt = w.transpose(1, 2).contiguous()  # [E, dim2, dim1]
+        e, n, k = wt.shape
+        w_fp8 = torch.empty_like(wt, dtype=torch.float8_e4m3fn)
+        w_sf = torch.empty(
+            (e, ceil_div(n, 128), ceil_div(k, 128)),
+            dtype=torch.float32,
+            device=wt.device,
+        )
+        for i in range(e):
+            w_fp8[i], w_sf[i] = per_block_cast_to_fp8(wt[i], use_ue8m0=True)
+        alignment = dg.get_mk_alignment_for_contiguous_layout()
+        entry = _Fp8Entry(w_fp8, w_sf, alignment)
+        _FP8_CACHE[key] = entry
+    return entry
+
+
+def _dx_plan(expert_offsets: torch.Tensor, T: int, E: int, align: int) -> tuple:
+    """``(dest, m_indices, src, m_max)`` for the padded contiguous layout.
+
+    Cached as an attribute on the offsets tensor: the up and down projections
+    of one layer save the SAME offsets object for backward, so the plan builds
+    once per layer per step and dies with the tensor (no stale-pointer risk).
+    The dozen tiny index kernels here cost more than the GEMM otherwise.
+    """
+    plan = getattr(expert_offsets, "_dg_dx_plan", None)
+    if plan is not None and plan[3] == T + E * align:
+        return plan
+
+    device = expert_offsets.device
+    offsets = expert_offsets.to(dtype=torch.int64)
+    counts = offsets[1:] - offsets[:-1]
+    padded = torch.div(counts + (align - 1), align, rounding_mode="floor") * align
+    pad_starts = torch.cat([padded.new_zeros(1), padded.cumsum(0)])[:-1]
+
+    # worst case each segment wastes < align rows; sized statically, no sync
+    m_max = T + E * align
+    row = torch.arange(T, device=device)
+    eid = torch.repeat_interleave(torch.arange(E, device=device), counts, output_size=T)
+    dest = pad_starts[eid] + (row - offsets[eid])
+
+    maps = torch.full((2, m_max), -1, dtype=torch.int32, device=device)
+    m_indices, src = maps[0], maps[1]
+    m_indices[dest] = eid.to(torch.int32)
+    src[dest] = row.to(torch.int32)
+
+    plan = (dest, m_indices, src, m_max)
+    expert_offsets._dg_dx_plan = plan
+    return plan
+
+
+def grouped_fp8_dx(
+    grad_h: torch.Tensor, base_weight, expert_offsets: torch.Tensor
+) -> torch.Tensor:
+    """``dx[start:end] = g_e @ W_e`` through the fp8 weight cache.
+
+    ``expert_offsets`` is ``[E+1]`` cumulative (device). Returns ``[T, dim2]``
+    bf16, row-aligned with ``grad_h`` ``[T, dim1]``.
+    """
+    dg = _deep_gemm()
+    entry = _get_fp8_cache(base_weight)
+
+    T = grad_h.shape[0]
+    E, N, K = entry.w_fp8.shape
+    dest, m_indices, src, m_max = _dx_plan(expert_offsets, T, E, entry.alignment)
+
+    a_fp8, a_sf = _pad_and_cast_fp8(grad_h, src, dest, m_max)
+    d = torch.empty((m_max, N), dtype=torch.bfloat16, device=grad_h.device)
+    dg.m_grouped_fp8_gemm_nt_contiguous(
+        (a_fp8, a_sf),
+        (entry.w_fp8, entry.w_sf),
+        d,
+        m_indices,
+        disable_ue8m0_cast=False,
+    )
+    return d.index_select(0, dest)
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/lora.py b/src/axolotl/integrations/kernels/libs/sonicmoe/lora.py
new file mode 100644
index 0000000000..9216c0c777
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/lora.py
@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+SonicMoE LoRA support via runtime weight materialization.
+
+SonicMoE uses opaque CUTLASS kernels that cannot be modified to fuse LoRA.
+Instead, we materialize the effective weight W_eff = W + scaling * (B @ A)
+before each CUTLASS call, and use a custom autograd.Function to route
+gradients back to the LoRA A and B parameters.
+
+PEFT unwrapping utilities are also provided to handle the ParamWrapper
+chain that PEFT creates when targeting expert parameters.
+"""
+
+from typing import Optional
+
+import torch
+
+# =============================================================================
+# PEFT unwrapping utilities
+# =============================================================================
+
+
+def has_lora(module) -> bool:
+    """Check if a module is wrapped by PEFT with LoRA."""
+    return hasattr(module, "base_layer") and hasattr(module, "lora_A")
+
+
+def get_lora_params_from_wrapper(module) -> tuple:
+    """Extract LoRA parameters from a PEFT ParamWrapper.
+
+    Returns:
+        (lora_A, lora_B, scaling) if LoRA is active, else (None, None, None)
+    """
+    if not hasattr(module, "lora_A") or not hasattr(module, "lora_B"):
+        return None, None, None
+
+    active_adapters = getattr(module, "active_adapters", ["default"])
+    if not active_adapters:
+        return None, None, None
+
+    adapter_name = active_adapters[0]
+
+    lora_A_dict = getattr(module, "lora_A", {})
+    lora_B_dict = getattr(module, "lora_B", {})
+    scaling_dict = getattr(module, "scaling", {})
+
+    if (
+        adapter_name not in lora_A_dict
+        or adapter_name not in lora_B_dict
+        or adapter_name not in scaling_dict
+    ):
+        return None, None, None
+
+    lora_A = lora_A_dict[adapter_name].weight
+    lora_B = lora_B_dict[adapter_name].weight
+    scaling = scaling_dict[adapter_name]
+
+    return lora_A, lora_B, scaling
+
+
+def unwrap_experts_lora(experts_module):
+    """Walk a PEFT ParamWrapper chain on ``self.experts``.
+
+    When PEFT targets ``experts.gate_up_proj`` and ``experts.down_proj``
+    via ``target_parameters``, ``self.experts`` becomes::
+
+        ParamWrapper(down_proj)
+          -> base_layer: ParamWrapper(gate_up_proj)
+              -> base_layer: Experts (the real module)
+
+    Returns:
+        (base_experts, lora_dict)
+
+        ``lora_dict`` maps parameter names to ``(lora_A, lora_B, scaling)``
+        tuples, or is empty if no LoRA is active.
+    """
+    wrappers = {}
+    module = experts_module
+    while hasattr(module, "base_layer") and hasattr(module, "lora_A"):
+        param_name = getattr(module, "parameter_name", None)
+        if param_name is not None:
+            wrappers[param_name] = module
+        module = module.base_layer
+
+    base_experts = module
+    lora_dict = {}
+
+    for param_name, wrapper in wrappers.items():
+        lora_A, lora_B, scaling = get_lora_params_from_wrapper(wrapper)
+        if lora_A is not None:
+            lora_dict[param_name] = (lora_A, lora_B, scaling)
+
+    return base_experts, lora_dict
+
+
+# =============================================================================
+# LoRA weight materialization autograd function
+# =============================================================================
+
+
+class MoELoRAMaterialize(torch.autograd.Function):
+    """Materialize ``W_eff = W + scaling * (B @ A)`` per expert and route grads.
+
+    Layout matches PEFT >= 0.19.1 ``ParamWrapper``: ``base [E, dim1, dim2]``,
+    ``lora_A [r*E, dim2]`` (E-outer, r-inner rows), ``lora_B [dim1, r*E]``
+    (r-outer, E-inner cols). Equivalent to
+    ``einsum("o r e, e r i -> e o i", lora_B.reshape(dim1, r, E), lora_A.reshape(E, r, dim2))``.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        base_weight: torch.Tensor,
+        lora_A: torch.Tensor,
+        lora_B: torch.Tensor,
+        scaling: float,
+    ) -> torch.Tensor:
+        E, dim1, dim2 = base_weight.shape
+        r = lora_A.shape[0] // E
+        assert lora_A.shape[0] == r * E, (
+            f"lora_A rows ({lora_A.shape[0]}) must be divisible by num_experts ({E})"
+        )
+
+        # Reshape PEFT rank-major to per-expert batched format
+        A_3d = lora_A.reshape(E, r, dim2)
+        B_3d = lora_B.reshape(dim1, r, E).permute(2, 0, 1).contiguous()  # [E, dim1, r]
+
+        # Fuse base_weight + scaling * (B_3d @ A_3d) into one op: avoids both the
+        # [E, dim1, dim2] bmm temp and the scaling-scaled copy of it.
+        W_eff = torch.baddbmm(base_weight, B_3d, A_3d, alpha=scaling)
+
+        ctx.save_for_backward(lora_A, lora_B)
+        ctx.scaling = scaling
+        ctx.E = E
+        ctx.r = r
+
+        return W_eff
+
+    @staticmethod
+    def backward(ctx, grad_W_eff: torch.Tensor):
+        lora_A, lora_B = ctx.saved_tensors
+        scaling = ctx.scaling
+        E = ctx.E
+        r = ctx.r
+
+        _, dim1, dim2 = grad_W_eff.shape
+
+        # Reshape to per-expert (same as forward)
+        A_3d = lora_A.reshape(E, r, dim2)
+        B_3d = lora_B.reshape(dim1, r, E).permute(2, 0, 1).contiguous()  # [E, dim1, r]
+
+        # dA_e = scaling * B_e^T @ dW_e
+        # [E, r, dim1] @ [E, dim1, dim2] = [E, r, dim2]
+        d_A_3d = scaling * torch.bmm(B_3d.transpose(1, 2), grad_W_eff)
+
+        # dB_e = scaling * dW_e @ A_e^T
+        # [E, dim1, dim2] @ [E, dim2, r] = [E, dim1, r]
+        d_B_3d = scaling * torch.bmm(grad_W_eff, A_3d.transpose(1, 2))
+
+        # Reshape back to PEFT rank-major layout
+        d_lora_A = d_A_3d.reshape(E * r, dim2)
+        d_lora_B = d_B_3d.permute(1, 2, 0).contiguous().reshape(dim1, E * r)
+
+        return None, d_lora_A, d_lora_B, None
+
+
+def materialize_expert_lora(
+    base_weight: torch.Tensor,
+    lora_params: Optional[tuple],
+) -> torch.Tensor:
+    """Materialize effective expert weight with optional LoRA delta.
+
+    Args:
+        base_weight: [E, dim1, dim2] frozen expert parameter
+        lora_params: (lora_A, lora_B, scaling) or None
+
+    Returns:
+        W_eff if lora_params is not None, else base_weight unchanged.
+    """
+    if lora_params is None:
+        return base_weight
+    lora_A, lora_B, scaling = lora_params
+    return MoELoRAMaterialize.apply(base_weight, lora_A, lora_B, scaling)
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/multi_lora.py b/src/axolotl/integrations/kernels/libs/sonicmoe/multi_lora.py
new file mode 100644
index 0000000000..7fd4a2d677
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/multi_lora.py
@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""SonicMoE + multi-adapter LoRA via combined-group weight materialization.
+
+SonicMoE dispatches opaque CUTLASS/CuteDSL grouped GEMMs that cannot fuse LoRA, so
+the single-adapter path materializes ``W_eff[e] = W[e] + scaling * (B_e @ A_e)``
+per expert and hands that to the kernel (see ``lora.MoELoRAMaterialize``).
+
+Multi-tenant co-training keeps that strategy but over the combined
+``(expert, tenant)`` grouping: a token carries the router's expert id *and* a
+per-row tenant id, and the effective weight for ``(token routed to e, tenant t)``
+is ``W[e] + scaling_t * (B_{e,t} @ A_{e,t})``. We materialize all ``E*T`` effective
+weights once and remap each token-slot's expert id to ``e * T + t``; the *same*
+opaque grouped GEMM, called with ``num_experts = E*T``, then routes every tenant
+in a single launch. The frozen base ``W[e]`` is shared in storage but, because the
+kernel needs one contiguous weight per group, is broadcast into the ``E*T`` buffer
+-- so transient materialization memory grows ~linearly with ``T`` (at ``T == 1``
+it is identical to the single-adapter path).
+
+This module owns only the materialize + id-remap (the LoRA-specific part); the
+CUTLASS GEMM is unchanged. ``T == 1`` reduces to ``MoELoRAMaterialize`` exactly.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+class MoEMultiLoRAMaterialize(torch.autograd.Function):
+    """Build ``W_eff[e*T + t] = W[e] + scaling_t * (B_{e,t} @ A_{e,t})``.
+
+    Layout (stacked over tenants, outer expert / inner tenant to match the
+    ``combined = e*T + t`` group id):
+      base    ``[E, out, in]``   (frozen)
+      lora_A  ``[T, E, r, in]``
+      lora_B  ``[T, E, out, r]``
+      scaling ``[T]``
+    Returns ``W_eff`` ``[E*T, out, in]``. Gradients flow to A/B only.
+    """
+
+    @staticmethod
+    def forward(ctx, base_weight, lora_A, lora_B, scaling):
+        T, E, r, in_dim = lora_A.shape
+        out_dim = base_weight.shape[1]
+        if base_weight.dtype != lora_A.dtype:
+            base_weight = base_weight.to(lora_A.dtype)
+
+        # Materialize directly in combined (e*T + t) order so the result needs no
+        # reorder copy: einsum over the [E, T, ...] view yields a contiguous
+        # [E, T, out, in] that reshapes to [E*T, ...] for free. The A/B permutes
+        # are on rank-sized tensors (~in/r smaller than W_eff), so ~free. delta is
+        # recomputed in backward, so the scale + frozen-base add are in-place.
+        a_et = lora_A.permute(1, 0, 2, 3)  # [E, T, r, in]
+        b_et = lora_B.permute(1, 0, 2, 3)  # [E, T, out, r]
+        delta = torch.einsum("etor,etri->etoi", b_et, a_et)  # [E, T, out, in]
+        delta.mul_(scaling.view(1, T, 1, 1))
+        delta.add_(base_weight.unsqueeze(1))  # broadcast frozen base over tenants
+        w_eff = delta.reshape(E * T, out_dim, in_dim)
+
+        ctx.save_for_backward(lora_A, lora_B, scaling)
+        ctx.shape = (T, E, r, in_dim, out_dim)
+        return w_eff
+
+    @staticmethod
+    def backward(ctx, grad_w_eff):
+        lora_A, lora_B, scaling = ctx.saved_tensors
+        T, E, r, in_dim, out_dim = ctx.shape
+        # Scale the (rank-sized) A/B grads, not the W_eff-sized incoming grad, to
+        # avoid a full-size copy in backward.
+        g = grad_w_eff.reshape(E, T, out_dim, in_dim)
+        a_et = lora_A.permute(1, 0, 2, 3)  # [E, T, r, in]
+        b_et = lora_B.permute(1, 0, 2, 3)  # [E, T, out, r]
+        # dA = B^T @ dDelta ; dB = dDelta @ A^T (then back to the [T, E, ...] layout)
+        d_a_et = torch.einsum("etor,etoi->etri", b_et, g).mul_(scaling.view(1, T, 1, 1))
+        d_b_et = torch.einsum("etoi,etri->etor", g, a_et).mul_(scaling.view(1, T, 1, 1))
+        return None, d_a_et.permute(1, 0, 2, 3), d_b_et.permute(1, 0, 2, 3), None
+
+
+def materialize_multi_lora_experts(base_weight, lora_A, lora_B, scaling):
+    """Autograd-backed ``W_eff`` builder; see :class:`MoEMultiLoRAMaterialize`."""
+    return MoEMultiLoRAMaterialize.apply(base_weight, lora_A, lora_B, scaling)
+
+
+def combined_expert_ids(expert_ids, tenant_ids, token_idx, num_tenants):
+    """Remap each token-slot's expert id to its combined ``e * T + t`` group.
+
+    ``expert_ids`` / ``token_idx`` are the per-(token, top-k slot) tensors sonicmoe
+    already builds; ``tenant_ids`` is per *token*. Returns int32 combined ids
+    aligned with ``expert_ids`` for a ``num_experts = E*T`` grouped GEMM.
+    """
+    t = tenant_ids.to(torch.int64)[token_idx.to(torch.int64)]
+    return (expert_ids.to(torch.int64) * num_tenants + t).to(torch.int32)
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4.py b/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4.py
new file mode 100644
index 0000000000..082fbb4533
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4.py
@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+NVFP4 base-weight support for the sonicmoe MoE-LoRA backend.
+
+Grouped up/down GEMMs and the gated activation that the LoRA forward composes
+over frozen expert weights, plus NVFP4 detection and dequantization helpers. The
+frozen base ``w1`` / ``w2`` may be a torchao ``NVFP4Tensor`` (block-scaled FP4).
+
+Backends:
+  - ``"torch"``: dense base, per-expert ``F.linear`` loop. Differentiable and
+    CPU-correct (float64 for gradcheck); the reference path.
+  - ``"dequant"``: dequantizes the NVFP4 base to dense per matmul, then runs the
+    same math. The GPU fallback when the base is packed NVFP4.
+  - ``"fp4_cute"``: in-kernel block-scaled W4A4 grouped GEMM on Blackwell
+    (SM100/SM110) via ``fp4_cute_ops``. Weights stay packed; activations are
+    NVFP4-quantized after grouping; backward dX runs per-expert chunked
+    dequant matmuls (never through the packed fp4 operand).
+
+Imports of ``sonicmoe``, ``torchao``, ``quack``, ``triton``, and CUDA-only ops
+are lazy (inside functions) so this module imports and tests cleanly on CPU with
+no GPU or upstream kernel packages.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def _torchao_nvfp4tensor_cls():
+    """Return torchao's ``NVFP4Tensor`` class, or ``None`` if torchao is absent."""
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+    except ImportError:
+        return None
+    return NVFP4Tensor
+
+
+def is_nvfp4_param(w) -> bool:
+    """True iff ``w`` is a torchao ``NVFP4Tensor``, False for a dense tensor.
+
+    Rolls its own torchao lookup rather than reusing
+    ``scattermoe_lora.selective_dequant.is_nvfp4_param``: importing that module
+    triggers the ``scattermoe_lora`` package ``__init__``, which imports triton
+    at module load and crashes on a CPU/no-triton box.
+    """
+    NVFP4Tensor = _torchao_nvfp4tensor_cls()
+    return NVFP4Tensor is not None and isinstance(w, NVFP4Tensor)
+
+
+def dequantize_expert_weight(w: torch.Tensor) -> torch.Tensor:
+    """Dequantize an NVFP4 expert weight to dense (bf16 on GPU), else identity.
+
+    A plain dense tensor is returned unchanged so the reference/CPU path is a
+    no-op.
+    """
+    if not is_nvfp4_param(w):
+        return w
+    if w.qdata.is_cuda:
+        from .triton_nvfp4 import dequant_nvfp4_triton, triton_available
+
+        if triton_available():
+            # torchao's dequantize() gathers a value table per element (slow
+            # enough to dominate the fp4_cute backward); this is one kernel.
+            return dequant_nvfp4_triton(
+                w.qdata, w.scale, w.per_tensor_scale, w.orig_dtype
+            )
+    return w.dequantize()
+
+
+def dequantize_expert_slice(w: torch.Tensor, e: int) -> torch.Tensor:
+    """Dense weight of expert ``e`` from an ``[E, dim1, dim2]`` stack.
+
+    Keeps backward memory at one dense expert instead of E. torchao's
+    NVFP4Tensor only implements rank-2 slicing, so the expert slice is rebuilt
+    from components; dense tensors just index.
+    """
+    if not is_nvfp4_param(w):
+        return w[e]
+    assert not getattr(w, "is_swizzled_scales", False), (
+        "swizzled NVFP4 scales unsupported (our loaders emit row-major)"
+    )
+    pts = w.per_tensor_scale
+    sliced = type(w)(
+        w.qdata[e : e + 1],
+        w.scale[e : e + 1],
+        w.block_size,
+        w.orig_dtype,
+        per_tensor_scale=pts[e : e + 1] if pts is not None else None,
+    )
+    return sliced.dequantize().squeeze(0)
+
+
+def resolve_gated_activation(config) -> str:
+    """Canonical gated-activation name from an HF text config.
+
+    Honors ``hidden_activation`` (Gemma's key, e.g. ``gelu_pytorch_tanh``) before
+    ``hidden_act`` (most other models). Reading only ``hidden_act`` would miss
+    Gemma and silently run SwiGLU where the model wants tanh-GeGLU.
+    """
+    act = getattr(config, "hidden_activation", None) or getattr(
+        config, "hidden_act", None
+    )
+    return (act or "silu").lower()
+
+
+def gated_activation(
+    h: torch.Tensor, act: str, *, concat: bool, limit: float | None = None
+) -> torch.Tensor:
+    """Apply a GLU gated activation to a ``[..., 2*I]`` pre-activation.
+
+    Returns ``[..., I]``. ``concat=True`` (HF default) splits gate/up as the
+    first/second half; ``concat=False`` (interleaved) takes even/odd lanes.
+
+    Supported ``act`` (``up * f(gate)``):
+      - swiglu / silu -> ``f = silu``
+      - geglu / gelu -> ``f = gelu`` (erf)
+      - gelu_tanh / gelu_pytorch_tanh -> ``f = gelu(approximate="tanh")`` (Gemma)
+      - reglu / relu -> ``f = relu``
+
+    ``limit`` (e.g. DeepSeek-V4 ``limit=10``) applies the model's clamped-SwiGLU
+    bound before the nonlinearity (``gate.clamp(max=limit)``, ``up.clamp(-limit,
+    limit)``), matching the eager experts' ``_apply_gate``; otherwise outliers on
+    a clamped model blow up. The gate nonlinearity is computed in fp32 then cast
+    back, matching upstream sonic-moe's ``_swiglu`` / ``_geglu``.
+    """
+    if concat:
+        i = h.shape[-1] // 2
+        gate = h[..., :i]
+        up = h[..., i:]
+    else:
+        gate = h[..., 0::2]
+        up = h[..., 1::2]
+
+    if limit is not None:
+        gate = gate.clamp(max=limit)
+        up = up.clamp(min=-limit, max=limit)
+
+    # Upstream sonic-moe evaluates the gate nonlinearity in fp32 to protect bf16/fp16
+    # precision. Only upcast lower-precision dtypes; leave fp32/fp64 as-is so we never
+    # truncate (which would also break a float64 gradcheck).
+    compute_dtype = (
+        torch.float32 if gate.dtype in (torch.bfloat16, torch.float16) else gate.dtype
+    )
+    g = gate.to(compute_dtype)
+
+    a = act.lower()
+    if a in ("swiglu", "silu"):
+        activated = F.silu(g)
+    elif a in ("gelu_tanh", "gelu_pytorch_tanh"):
+        activated = F.gelu(g, approximate="tanh")
+    elif a in ("geglu", "gelu"):
+        activated = F.gelu(g)
+    elif a in ("reglu", "relu"):
+        activated = F.relu(g)
+    else:
+        raise ValueError(f"unsupported gated activation: {act!r}")
+
+    return up * activated.to(gate.dtype)
+
+
+def _grouped_gemm(
+    x_grouped: torch.Tensor,
+    weight: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    *,
+    backend: str,
+) -> torch.Tensor:
+    """Per-expert ``x_e @ w[e].T`` over expert-sorted tokens.
+
+    ``weight`` is ``[E, out, in]``; token rows in
+    ``x_grouped[offsets[e - 1]:offsets[e]]`` (offsets[-1] == 0) go to expert e.
+    """
+    if backend == "fp4_cute":
+        from .fp4_cute_ops import grouped_nvfp4_linear
+
+        return grouped_nvfp4_linear(x_grouped, weight, expert_offsets)
+    if backend == "dequant":
+        weight = dequantize_expert_weight(weight)
+    elif backend != "torch":
+        raise ValueError(f"unknown backend: {backend!r}")
+
+    from .nvfp4_lora import _use_grouped_mm
+
+    if _use_grouped_mm(x_grouped) and weight.dtype == x_grouped.dtype:
+        return torch._grouped_mm(
+            x_grouped,
+            weight.transpose(-2, -1),
+            offs=expert_offsets.to(torch.int32),
+        )
+
+    E = weight.shape[0]
+    outs = []
+    start = 0
+    offsets = expert_offsets.tolist()
+    for e in range(E):
+        end = offsets[e]
+        outs.append(F.linear(x_grouped[start:end], weight[e]))
+        start = end
+    return torch.cat(outs, dim=0)
+
+
+def grouped_up_gemm(
+    x_grouped: torch.Tensor,
+    w1: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    *,
+    backend: str,
+    concat: bool,  # noqa: ARG001 (layout handled downstream in gated_activation)
+) -> torch.Tensor:
+    """Grouped gate/up projection.
+
+    ``x_grouped``: ``[T_total, H]`` tokens gathered/sorted by expert.
+    ``w1``: ``[E, 2*I, H]`` frozen base (dense for ``"torch"``, possibly NVFP4
+    for ``"dequant"``). ``expert_offsets``: ``[E+1]`` int64 cumulative counts.
+    Returns preact ``h``: ``[T_total, 2*I]`` where expert-e rows are
+    ``x_e @ w1[e].T``.
+
+    ``concat`` describes the gate/up interleaving of the ``2*I`` output columns;
+    it is consumed later by :func:`gated_activation`, not here, since both
+    layouts produce the same ``x_e @ w1[e].T`` product.
+    """
+    return _grouped_gemm(x_grouped, w1, expert_offsets[1:], backend=backend)
+
+
+def grouped_down_gemm(
+    a_grouped: torch.Tensor,
+    w2: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    *,
+    backend: str,
+) -> torch.Tensor:
+    """Grouped down projection.
+
+    ``a_grouped``: ``[T_total, I]``; ``w2``: ``[E, H, I]``; returns
+    ``y``: ``[T_total, H]`` where expert-e rows are ``a_e @ w2[e].T``.
+    """
+    return _grouped_gemm(a_grouped, w2, expert_offsets[1:], backend=backend)
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4_lora.py b/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4_lora.py
new file mode 100644
index 0000000000..9476508b24
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4_lora.py
@@ -0,0 +1,677 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Frozen-base grouped LoRA forward/backward for the sonicmoe NVFP4 backend.
+
+The base expert weights are frozen (NVFP4 in production; the sibling ``nvfp4``
+module either dequantizes them or keeps them packed for the fp4_cute kernel).
+Only the LoRA A/B tensors are trainable. Where the whole-weight path (``lora.MoELoRAMaterialize``) builds
+``W_eff = W + scaling * (B @ A)`` per expert and hands it to an opaque CUTLASS
+call, this module keeps the low-rank factors separate and fuses them at the
+grouped-token level so it composes with a real grouped GEMM and never forms a
+base-weight gradient.
+
+The forward per expert group ``e`` is exactly ``x_e @ W_eff_e^T`` with
+``W_eff_e = W_e + scaling * (B_e @ A_e)``, so the hand-written backward mirrors
+``MoELoRAMaterialize.backward``: it forms the full ``dW_eff_e`` from the grouped
+tokens, then maps it to ``dA_e / dB_e`` with the same rank-major reshape/permute.
+Because ``W_e`` is frozen we only route grads to ``x``, ``lora_A`` and
+``lora_B``; ``W`` is treated as a constant tensor for the ``dx`` term.
+
+PEFT rank-major layout (matching ``lora.MoELoRAMaterialize``): for a weight of
+logical shape ``[E, dim1, dim2]``, ``lora_A`` is ``[r*E, dim2]`` (E-outer,
+r-inner rows) and ``lora_B`` is ``[dim1, r*E]`` (r-outer, E-inner cols).
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from .nvfp4 import (
+    dequantize_expert_slice,
+    dequantize_expert_weight,
+    gated_activation,
+    grouped_down_gemm,
+    grouped_up_gemm,
+    is_nvfp4_param,
+)
+
+
+def _use_grouped_mm(x: torch.Tensor) -> bool:
+    """``torch._grouped_mm`` replaces the per-expert Python loops on sm90+.
+
+    The loops cost one host sync per expert plus E small kernel launches; the
+    grouped GEMM is a single launch with ragged (unaligned, possibly empty)
+    segments, which matches expert token counts exactly.
+    """
+    return (
+        x.is_cuda
+        and x.dtype == torch.bfloat16
+        and hasattr(torch, "_grouped_mm")
+        and torch.cuda.get_device_capability(x.device)[0] >= 9
+    )
+
+
+def _grouped_offs(expert_offsets: torch.Tensor, device) -> torch.Tensor:
+    """``[E]`` int32 cumulative segment ends from ``[E+1]`` offsets."""
+    return expert_offsets[1:].to(device=device, dtype=torch.int32)
+
+
+def _b3d_contiguous(lora_B: torch.Tensor, E: int, dim1: int) -> torch.Tensor:
+    """``[E, dim1, r]`` contiguous view of ``lora_B`` ``[dim1, r*E]``.
+
+    The permute leaves no unit-stride dim; grouped GEMM needs one.
+    """
+    r = lora_B.shape[1] // E
+    return lora_B.reshape(dim1, r, E).permute(2, 0, 1).contiguous()
+
+
+def _lora_delta_per_group(
+    x_grouped: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    scaling: float,
+    E: int,
+    dim1: int,
+    dim2: int,
+    B_c: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Low-rank contribution ``scaling * ((x_e @ A_e^T) @ B_e^T)`` per group.
+
+    ``lora_A`` is ``[r*E, dim2]``, ``lora_B`` is ``[dim1, r*E]``. Returns a
+    ``[T, dim1]`` tensor aligned row-for-row with ``x_grouped`` ``[T, dim2]``.
+    ``B_c`` is an optional precomputed ``_b3d_contiguous(lora_B, E, dim1)``.
+    """
+    r = lora_A.shape[0] // E
+    A_3d = lora_A.reshape(E, r, dim2)  # [E, r, dim2]
+    B_3d = lora_B.reshape(dim1, r, E).permute(2, 0, 1)  # [E, dim1, r]
+
+    if _use_grouped_mm(x_grouped) and lora_A.dtype == x_grouped.dtype:
+        offs = _grouped_offs(expert_offsets, x_grouped.device)
+        z = torch._grouped_mm(x_grouped, A_3d.transpose(-2, -1), offs=offs)
+        if B_c is None:
+            B_c = B_3d.contiguous()
+        # scaling applied on the small [T, r] intermediate, not the [T, dim1] product
+        return torch._grouped_mm(z * scaling, B_c.transpose(-2, -1), offs=offs)
+
+    out = x_grouped.new_zeros((x_grouped.shape[0], dim1))
+    for e in range(E):
+        start = int(expert_offsets[e])
+        end = int(expert_offsets[e + 1])
+        if end <= start:
+            continue
+        x_e = x_grouped[start:end]  # [T_e, dim2]
+        z_e = x_e @ A_3d[e].transpose(0, 1)  # [T_e, r]
+        out[start:end] = scaling * (z_e @ B_3d[e].transpose(0, 1))  # [T_e, dim1]
+    return out
+
+
+def _lora_backward_per_group(
+    grad_h: torch.Tensor,
+    x_grouped: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    lora_A: torch.Tensor,
+    lora_B: torch.Tensor,
+    base_weight: torch.Tensor,
+    scaling: float,
+    B_c: Optional[torch.Tensor] = None,
+) -> tuple:
+    """Grads for a frozen-base grouped-LoRA linear ``h_e = x_e @ W_eff_e^T``.
+
+    ``base_weight`` is ``[E, dim1, dim2]`` (frozen; dense or packed NVFP4,
+    dequantized one expert slice at a time), used only for the ``dx`` term.
+    ``B_c`` is an optional precomputed ``_b3d_contiguous(lora_B, E, dim1)``.
+    Returns ``(dx, d_lora_A, d_lora_B)``.
+    """
+    E, dim1, dim2 = base_weight.shape
+    r = lora_A.shape[0] // E
+    A_3d = lora_A.reshape(E, r, dim2)  # [E, r, dim2]
+    B_3d = lora_B.reshape(dim1, r, E).permute(2, 0, 1)  # [E, dim1, r]
+
+    if _use_grouped_mm(grad_h) and lora_A.dtype == grad_h.dtype:
+        offs = _grouped_offs(expert_offsets, grad_h.device)
+        if B_c is None:
+            # B_3d's permute leaves no unit-stride dim; grouped GEMM needs one.
+            B_c = B_3d.contiguous()
+        # dz_e = g_e @ B_e; z_e = x_e @ A_e^T (forward intermediate, recomputed:
+        # cheaper than saving a [T, r] per GEMM through gradient checkpointing).
+        # scaling lands once on each [T, r] intermediate, not the [T, dim] products.
+        dz = torch._grouped_mm(grad_h, B_c, offs=offs) * scaling  # [T, r]
+        z = torch._grouped_mm(x_grouped, A_3d.transpose(-2, -1), offs=offs) * scaling
+        # dx_e = g_e @ W_e + dz_e @ A_e. Base term: fp8 weight cache on DeepGEMM
+        # when available, else whole-weight dequant + one grouped GEMM. Either
+        # way the weights match the forward operands (folded SFB + alpha on
+        # fp4_cute; the fp8 cache is built from that same folded dequant).
+        from .fp8_bwd import fp8_dx_supported, grouped_fp8_dx
+
+        if fp8_dx_supported(grad_h, base_weight):
+            dx = grouped_fp8_dx(grad_h, base_weight, expert_offsets)
+        else:
+            from .fp4_cute_ops import dequantize_engine_weight
+
+            w_dense = dequantize_engine_weight(base_weight).to(grad_h.dtype)
+            dx = torch._grouped_mm(grad_h, w_dense, offs=offs)
+        dx += torch._grouped_mm(dz, A_3d, offs=offs)
+        # dA_e = dz_e^T @ x_e; dB_e = g_e^T @ z_e
+        d_A_3d = torch._grouped_mm(dz.transpose(0, 1), x_grouped, offs=offs)
+        d_B_3d = torch._grouped_mm(grad_h.transpose(0, 1), z, offs=offs)
+        d_lora_A = d_A_3d.reshape(E * r, dim2)
+        d_lora_B = d_B_3d.permute(1, 2, 0).reshape(dim1, E * r)
+        return dx, d_lora_A, d_lora_B
+
+    dx = torch.zeros_like(x_grouped)
+    d_A_3d = grad_h.new_zeros((E, r, dim2))
+    d_B_3d = grad_h.new_zeros((E, dim1, r))
+
+    for e in range(E):
+        start = int(expert_offsets[e])
+        end = int(expert_offsets[e + 1])
+        if end <= start:
+            continue
+        x_e = x_grouped[start:end]  # [T_e, dim2]
+        g_e = grad_h[start:end]  # [T_e, dim1]
+
+        # dx_e = g_e @ W_eff_e with W_eff_e = W_e + scaling * (B_e @ A_e),
+        # split so W_eff is never materialized and the NVFP4 base dequantizes
+        # one expert slice at a time.
+        w_e = dequantize_expert_slice(base_weight, e)  # [dim1, dim2]
+        dx[start:end] = g_e @ w_e.to(g_e.dtype) + scaling * ((g_e @ B_3d[e]) @ A_3d[e])
+
+        # dW_eff_e = grad_h_e^T @ x_e  ([dim1, dim2], the [E, dim1, dim2] convention)
+        dW_e = g_e.transpose(0, 1) @ x_e  # [dim1, dim2]
+
+        # Same map as MoELoRAMaterialize.backward:
+        #   dA_e = scaling * B_e^T @ dW_e     ([r, dim1] @ [dim1, dim2] = [r, dim2])
+        #   dB_e = scaling * dW_e @ A_e^T     ([dim1, dim2] @ [dim2, r] = [dim1, r])
+        d_A_3d[e] = scaling * (B_3d[e].transpose(0, 1) @ dW_e)
+        d_B_3d[e] = scaling * (dW_e @ A_3d[e].transpose(0, 1))
+
+    d_lora_A = d_A_3d.reshape(E * r, dim2)
+    d_lora_B = d_B_3d.permute(1, 2, 0).reshape(dim1, E * r)
+    return dx, d_lora_A, d_lora_B
+
+
+class GroupedUpProjLoRA(torch.autograd.Function):
+    """Grouped up-projection with frozen base + trainable LoRA.
+
+    ``h = base_up(x) + scaling * ((x @ A1^T) @ B1^T)`` per expert group, where
+    ``base_up`` is the grouped GEMM ``x_e @ w1[e]^T``. ``w1`` is ``[E, 2I, H]``
+    (dim1=2I, dim2=H); ``lora_A1`` is ``[r*E, H]``, ``lora_B1`` is ``[2I, r*E]``.
+    Grads route to ``x_grouped``, ``lora_A1``, ``lora_B1`` only (``w1`` frozen).
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x_grouped: torch.Tensor,
+        w1: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        lora_A1: torch.Tensor,
+        lora_B1: torch.Tensor,
+        scaling: float,
+        backend: str,
+        concat: bool,
+    ) -> torch.Tensor:
+        E, dim1, dim2 = w1.shape
+        B_c = (
+            _b3d_contiguous(lora_B1, E, dim1)
+            if _use_grouped_mm(x_grouped) and lora_A1.dtype == x_grouped.dtype
+            else None
+        )
+        delta = _lora_delta_per_group(
+            x_grouped, expert_offsets, lora_A1, lora_B1, scaling, E, dim1, dim2, B_c=B_c
+        )
+        if backend == "fp4_cute":
+            from .fp4_cute_ops import grouped_nvfp4_linear_add_delta
+
+            h = grouped_nvfp4_linear_add_delta(x_grouped, w1, expert_offsets[1:], delta)
+        else:
+            base = grouped_up_gemm(
+                x_grouped, w1, expert_offsets, backend=backend, concat=concat
+            )
+            h = base + delta
+
+        ctx.save_for_backward(x_grouped, w1, expert_offsets, lora_A1, lora_B1, B_c)
+        ctx.scaling = scaling
+        return h
+
+    @staticmethod
+    def backward(ctx, grad_h: torch.Tensor):
+        x_grouped, w1, expert_offsets, lora_A1, lora_B1, B_c = ctx.saved_tensors
+        dx, dA, dB = _lora_backward_per_group(
+            grad_h.contiguous(),
+            x_grouped,
+            expert_offsets,
+            lora_A1,
+            lora_B1,
+            w1,
+            ctx.scaling,
+            B_c=B_c,
+        )
+        return dx, None, None, dA, dB, None, None, None
+
+
+def _fused_up_act_enabled() -> bool:
+    """Fuse up-GEMM + gated activation + LoRA-delta add into one fp4_cute
+    gated-engine call. Default on; ``AXOLOTL_SONICMOE_NVFP4_FUSED_UP=0`` is the
+    kill switch back to the unfused up-GEMM + separate activation."""
+    import os
+
+    return os.environ.get("AXOLOTL_SONICMOE_NVFP4_FUSED_UP", "1") != "0"
+
+
+class GroupedUpProjActLoRA(torch.autograd.Function):
+    """Fused grouped up-projection + gated activation with frozen NVFP4 base.
+
+    One fp4_cute gated-engine call computes
+    ``a = act(pts * base_up(x) + delta)`` per expert group: the LoRA delta
+    rides the epilogue as a preact-space aux add, the per-expert pts as an
+    exact per-row colvec multiply, and the activation runs on the fp32
+    accumulator. The INTERLEAVED preact D is stashed for backward; grads route
+    to ``x_grouped``, ``lora_A1``, ``lora_B1`` only (``w1`` frozen).
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x_grouped: torch.Tensor,
+        w1: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        lora_A1: torch.Tensor,
+        lora_B1: torch.Tensor,
+        scaling: float,
+    ) -> torch.Tensor:
+        from .fp4_cute_ops import gated_nvfp4_forward
+        from .sf_layout import gate_up_interleave_perm
+
+        E, dim1, dim2 = w1.shape
+        B_c = _b3d_contiguous(lora_B1, E, dim1)
+        # The delta rides the kernel's INTERLEAVED preact space: permute the
+        # small [2I, r*E] B factor's rows once so the [T, 2I] delta itself
+        # never needs a gather. B_c stays concat for the shared backward.
+        perm = gate_up_interleave_perm(dim1, device=lora_B1.device)
+        B_il = lora_B1.index_select(0, perm)
+        delta = _lora_delta_per_group(
+            x_grouped, expert_offsets, lora_A1, B_il, scaling, E, dim1, dim2
+        )
+        cu = expert_offsets.to(torch.int32)
+        postact, preact = gated_nvfp4_forward(
+            x_grouped, w1, cu, delta.to(torch.bfloat16)
+        )
+        ctx.save_for_backward(
+            x_grouped, w1, expert_offsets, lora_A1, lora_B1, B_c, preact
+        )
+        ctx.scaling = scaling
+        return postact.to(x_grouped.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_a: torch.Tensor):
+        x_grouped, w1, expert_offsets, lora_A1, lora_B1, B_c, preact = ctx.saved_tensors
+        # swiglu backward from the INTERLEAVED preact: even lanes gate, odd up.
+        pre = preact.view(preact.shape[0], -1, 2)
+        gate = pre[..., 0].float()
+        up = pre[..., 1].float()
+        sig = torch.sigmoid(gate)
+        ga = grad_a.float()
+        grad_gate = ga * up * (sig * (1.0 + gate * (1.0 - sig)))
+        grad_up = ga * (gate * sig)
+        # grad wrt the CONCAT-layout virtual preact h = [gate | up], matching
+        # the concat layouts of w1 / lora_B1 in the shared backward.
+        grad_h = torch.cat([grad_gate, grad_up], dim=1).to(grad_a.dtype)
+        dx, dA, dB = _lora_backward_per_group(
+            grad_h,
+            x_grouped,
+            expert_offsets,
+            lora_A1,
+            lora_B1,
+            w1,
+            ctx.scaling,
+            B_c=B_c,
+        )
+        return dx, None, None, dA, dB, None
+
+
+class GroupedDownProjLoRA(torch.autograd.Function):
+    """Grouped down-projection with frozen base + trainable LoRA.
+
+    ``y = base_down(a) + scaling * ((a @ A2^T) @ B2^T)`` per expert group, where
+    ``base_down`` is the grouped GEMM ``a_e @ w2[e]^T``. ``w2`` is ``[E, H, I]``
+    (dim1=H, dim2=I); ``lora_A2`` is ``[r*E, I]``, ``lora_B2`` is ``[H, r*E]``.
+    Grads route to ``a_grouped``, ``lora_A2``, ``lora_B2`` only (``w2`` frozen).
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        a_grouped: torch.Tensor,
+        w2: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        lora_A2: torch.Tensor,
+        lora_B2: torch.Tensor,
+        scaling: float,
+        backend: str,
+    ) -> torch.Tensor:
+        E, dim1, dim2 = w2.shape
+        B_c = (
+            _b3d_contiguous(lora_B2, E, dim1)
+            if _use_grouped_mm(a_grouped) and lora_A2.dtype == a_grouped.dtype
+            else None
+        )
+        delta = _lora_delta_per_group(
+            a_grouped, expert_offsets, lora_A2, lora_B2, scaling, E, dim1, dim2, B_c=B_c
+        )
+        if backend == "fp4_cute":
+            from .fp4_cute_ops import grouped_nvfp4_linear_add_delta
+
+            y = grouped_nvfp4_linear_add_delta(a_grouped, w2, expert_offsets[1:], delta)
+        else:
+            base = grouped_down_gemm(a_grouped, w2, expert_offsets, backend=backend)
+            y = base + delta
+
+        ctx.save_for_backward(a_grouped, w2, expert_offsets, lora_A2, lora_B2, B_c)
+        ctx.scaling = scaling
+        return y
+
+    @staticmethod
+    def backward(ctx, grad_y: torch.Tensor):
+        a_grouped, w2, expert_offsets, lora_A2, lora_B2, B_c = ctx.saved_tensors
+        da, dA, dB = _lora_backward_per_group(
+            grad_y.contiguous(),
+            a_grouped,
+            expert_offsets,
+            lora_A2,
+            lora_B2,
+            w2,
+            ctx.scaling,
+            B_c=B_c,
+        )
+        return da, None, None, dA, dB, None, None
+
+
+def _add_expert_bias(
+    out: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    bias: torch.Tensor,
+) -> torch.Tensor:
+    """Add a per-expert bias ``[E, dim]`` to grouped rows ``[T, dim]``."""
+    E = bias.shape[0]
+    result = out
+    for e in range(E):
+        start = int(expert_offsets[e])
+        end = int(expert_offsets[e + 1])
+        if end <= start:
+            continue
+        result = result.index_add(
+            0,
+            torch.arange(start, end, device=out.device),
+            bias[e].unsqueeze(0).expand(end - start, -1),
+        )
+    return result
+
+
+def grouped_expert_mlp_lora(
+    x_grouped: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    w1: torch.Tensor,
+    b1: Optional[torch.Tensor],
+    w2: torch.Tensor,
+    b2: Optional[torch.Tensor],
+    lora1: Optional[tuple],
+    lora2: Optional[tuple],
+    *,
+    act: str,
+    backend: str,
+    concat: bool,
+    scaling1: float,
+    scaling2: float,
+    limit: Optional[float] = None,
+) -> torch.Tensor:
+    """Chain up-LoRA -> gated activation -> down-LoRA over grouped tokens.
+
+    ``lora1`` / ``lora2`` are ``(lora_A, lora_B)`` tuples or ``None`` (``None``
+    means plain base grouped GEMM, no low-rank path). ``b1`` / ``b2`` are
+    optional per-expert biases ``[E, dim]``. ``limit`` is the clamped-SwiGLU
+    bound (e.g. DeepSeek-V4). Returns ``y_grouped`` ``[T, H]``.
+    """
+    if (
+        lora1 is not None
+        and backend == "fp4_cute"
+        and b1 is None
+        and limit is None
+        and concat
+        and act in ("silu", "swiglu")
+        and _fused_up_act_enabled()
+    ):
+        A1, B1 = lora1
+        a = GroupedUpProjActLoRA.apply(x_grouped, w1, expert_offsets, A1, B1, scaling1)
+    else:
+        if lora1 is not None:
+            A1, B1 = lora1
+            h = GroupedUpProjLoRA.apply(
+                x_grouped, w1, expert_offsets, A1, B1, scaling1, backend, concat
+            )
+        else:
+            h = grouped_up_gemm(
+                x_grouped, w1, expert_offsets, backend=backend, concat=concat
+            )
+        if b1 is not None:
+            h = _add_expert_bias(h, expert_offsets, b1)
+
+        a = gated_activation(h, act, concat=concat, limit=limit)
+
+    if lora2 is not None:
+        A2, B2 = lora2
+        y = GroupedDownProjLoRA.apply(a, w2, expert_offsets, A2, B2, scaling2, backend)
+    else:
+        y = grouped_down_gemm(a, w2, expert_offsets, backend=backend)
+    if b2 is not None:
+        y = _add_expert_bias(y, expert_offsets, b2)
+
+    return y
+
+
+def route_and_group(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    num_experts: int,
+) -> tuple:
+    """Gather tokens into expert-sorted order for a grouped MoE forward.
+
+    ``top_k_index`` / ``top_k_weights`` are ``[T, K]`` (per-token routed experts
+    and combine weights). Returns ``(x_grouped [T*K, H], expert_offsets [E+1],
+    gather_token_idx [T*K], weights_grouped [T*K])`` where ``x_grouped`` rows are
+    sorted by expert. ``index_select`` keeps this differentiable wrt
+    ``hidden_states``; the combine weights stay differentiable wrt the router.
+    """
+    T, K = top_k_index.shape
+    device = hidden_states.device
+
+    flat_expert = top_k_index.reshape(-1)
+    # CUDA skips this guard: it costs a device-to-host sync per MoE layer, EP is
+    # already rejected upstream, and non-EP routed ids come from a width-E topk.
+    if (
+        not flat_expert.is_cuda
+        and flat_expert.numel()
+        and int(flat_expert.max()) >= num_experts
+    ):
+        raise NotImplementedError(
+            "sonicmoe NVFP4 path received routed expert id >= num_experts "
+            f"({int(flat_expert.max())} >= {num_experts}); expert parallelism is "
+            "not supported yet"
+        )
+
+    flat_weight = top_k_weights.reshape(-1).to(hidden_states.dtype)
+    token_ids = torch.arange(T, device=device).repeat_interleave(K)
+
+    sorted_experts, order = torch.sort(flat_expert, stable=True)
+    gather_token_idx = token_ids[order]
+    weights_grouped = flat_weight[order]
+    if hidden_states.is_cuda:
+        # autograd's backward for index_select is an atomic index_add; this
+        # gathers in both directions instead (see _CombineByGather).
+        x_grouped = _GatherRows.apply(hidden_states, gather_token_idx)
+    else:
+        x_grouped = hidden_states.index_select(0, gather_token_idx)
+
+    # searchsorted on the sorted ids instead of bincount + cumsum: bincount's
+    # output size is data-dependent, so it host-syncs every call.
+    expert_offsets = torch.searchsorted(
+        sorted_experts,
+        torch.arange(num_experts + 1, device=device, dtype=sorted_experts.dtype),
+    )
+
+    return x_grouped, expert_offsets, gather_token_idx, weights_grouped
+
+
+class _GatherRows(torch.autograd.Function):
+    """``hidden[idx]`` whose backward gathers instead of scatter-adding.
+
+    Every token id appears exactly K times in ``idx``, so ``dh[t]`` is the sum
+    of the K grad rows a stable sort of ``idx`` places contiguously.
+    """
+
+    @staticmethod
+    def forward(ctx, hidden, gather_token_idx):
+        ctx.save_for_backward(gather_token_idx)
+        ctx.num_tokens = hidden.shape[0]
+        return hidden.index_select(0, gather_token_idx)
+
+    @staticmethod
+    def backward(ctx, grad):
+        (gidx,) = ctx.saved_tensors
+        t = ctx.num_tokens
+        k = gidx.shape[0] // t
+        by_token = torch.sort(gidx, stable=True).indices
+        dh = (
+            grad.index_select(0, by_token)
+            .view(t, k, grad.shape[-1])
+            .sum(1, dtype=torch.float32)
+            .to(grad.dtype)
+        )
+        return dh, None
+
+
+class _CombineByGather(torch.autograd.Function):
+    """Router-weighted combine with gathers in BOTH directions (no atomics).
+
+    Each token receives exactly ``K = rows / num_tokens`` contributions, so a
+    stable sort of the token ids groups each token's rows contiguously and the
+    combine is a gather + fixed-order sum; the backward w.r.t. the grouped rows
+    is a plain gather of ``dout`` by token id. ``index_add`` (the scatter the
+    naive version needs, and what autograd emits for ``index_select``) is
+    atomic-bound and dominated the layer profile.
+    """
+
+    @staticmethod
+    def forward(ctx, y_grouped, gather_token_idx, weights_grouped, num_tokens):
+        k = y_grouped.shape[0] // num_tokens
+        by_token = torch.sort(gather_token_idx, stable=True).indices
+        scaled = y_grouped * weights_grouped.unsqueeze(-1)
+        out = (
+            scaled.index_select(0, by_token)
+            .view(num_tokens, k, y_grouped.shape[-1])
+            .sum(1, dtype=torch.float32)
+            .to(y_grouped.dtype)
+        )
+        ctx.save_for_backward(y_grouped, gather_token_idx, weights_grouped)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        y_grouped, gather_token_idx, weights_grouped = ctx.saved_tensors
+        g = grad_out.index_select(0, gather_token_idx)
+        dy = g * weights_grouped.unsqueeze(-1)
+        dw = (g.float() * y_grouped.float()).sum(-1).to(weights_grouped.dtype)
+        return dy, None, dw, None
+
+
+def combine_expert_outputs(
+    y_grouped: torch.Tensor,
+    gather_token_idx: torch.Tensor,
+    weights_grouped: torch.Tensor,
+    num_tokens: int,
+) -> torch.Tensor:
+    """Scale expert-sorted rows by router weights and combine to ``[T, H]``."""
+    if y_grouped.is_cuda and y_grouped.shape[0] % max(num_tokens, 1) == 0:
+        return _CombineByGather.apply(
+            y_grouped, gather_token_idx, weights_grouped, num_tokens
+        )
+    out = y_grouped.new_zeros((num_tokens, y_grouped.shape[-1]))
+    scaled = y_grouped * weights_grouped.unsqueeze(-1)
+    return out.index_add(0, gather_token_idx, scaled)
+
+
+def grouped_moe_reference_forward(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    w1: torch.Tensor,
+    b1: Optional[torch.Tensor],
+    w2: torch.Tensor,
+    b2: Optional[torch.Tensor],
+    lora1: Optional[tuple],
+    lora2: Optional[tuple],
+    num_experts: int,
+    *,
+    act: str,
+    backend: str,
+    concat: bool,
+    scaling1: float,
+    scaling2: float,
+    limit: Optional[float] = None,
+) -> torch.Tensor:
+    """End-to-end NVFP4 MoE forward: route -> grouped gated MLP -> combine.
+
+    The frozen base ``w1`` / ``w2`` may be packed NVFP4. ``backend="dequant"``
+    dequantizes the whole base to dense once up front so both the forward and
+    the hand-written backward operate on dense tensors; ``backend="fp4_cute"``
+    keeps the weights packed and runs the in-kernel SM100 W4A4 grouped GEMM
+    (quantized activations, chunked-dequant dX in backward). ``limit`` is the
+    clamped-SwiGLU bound (e.g. DeepSeek-V4). Runs on CPU (dense base) so the
+    path is validated without a GPU.
+    """
+    if backend == "dequant":
+        w1 = dequantize_expert_weight(w1)
+        w2 = dequantize_expert_weight(w2)
+        backend = "torch"
+    elif backend == "fp4_cute":
+        from .fp4_cute import fp4_cute_available
+
+        if not fp4_cute_available():
+            raise RuntimeError(
+                "backend='fp4_cute' requires an SM100/SM110 GPU with quack and "
+                "the cutlass DSL installed"
+            )
+        if not (is_nvfp4_param(w1) and is_nvfp4_param(w2)):
+            raise ValueError(
+                "backend='fp4_cute' requires packed NVFP4 base weights; use "
+                "backend='torch' for a dense base"
+            )
+
+    x_grouped, expert_offsets, gather_token_idx, weights_grouped = route_and_group(
+        hidden_states, top_k_index, top_k_weights, num_experts
+    )
+    y_grouped = grouped_expert_mlp_lora(
+        x_grouped,
+        expert_offsets,
+        w1,
+        b1,
+        w2,
+        b2,
+        lora1,
+        lora2,
+        act=act,
+        backend=backend,
+        concat=concat,
+        scaling1=scaling1,
+        scaling2=scaling2,
+        limit=limit,
+    )
+    return combine_expert_outputs(
+        y_grouped, gather_token_idx, weights_grouped, hidden_states.shape[0]
+    )
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4_quant.py b/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4_quant.py
new file mode 100644
index 0000000000..f5cd768cac
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/nvfp4_quant.py
@@ -0,0 +1,112 @@
+"""Pure-torch NVFP4 quantize/dequantize reference.
+
+NVFP4: E2M1 4-bit codes (two per byte, low nibble first) + one E4M3 block scale
+per 16 contiguous K elements + an optional fp32 per-tensor scale.
+``value = code_value * block_scale * per_tensor_scale``.
+
+Used as the numeric oracle for the fp4_cute kernel path (dequantize the exact
+operands the kernel consumes, matmul in fp32) and as a checkpoint-free
+quantizer for tests. Encoder rounding does not need to be
+bit-identical to torchao/quack: the oracle always dequantizes the operands
+actually fed to the kernel, so correctness checks are encoder-independent.
+"""
+
+from __future__ import annotations
+
+import torch
+
+SF_VEC_SIZE = 16
+E4M3_MAX = 448.0
+
+E2M1_MAGNITUDES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+# Midpoints between consecutive magnitudes, for nearest-value bucketing.
+_E2M1_BOUNDS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
+# Full 16-entry table indexed by the 4-bit code (bit 3 = sign).
+FP4_CODE_VALUES = tuple(E2M1_MAGNITUDES) + tuple(-m for m in E2M1_MAGNITUDES)
+
+
+def fp4_code_to_value(codes: torch.Tensor) -> torch.Tensor:
+    """Map 4-bit E2M1 codes in [0, 16) to their float32 values."""
+    table = torch.tensor(FP4_CODE_VALUES, dtype=torch.float32, device=codes.device)
+    return table[codes.long()]
+
+
+def encode_e2m1(x: torch.Tensor) -> torch.Tensor:
+    """Round float values to the nearest E2M1 code (uint8 in [0, 16))."""
+    bounds = torch.tensor(_E2M1_BOUNDS, dtype=torch.float32, device=x.device)
+    mag_idx = torch.bucketize(x.abs().float(), bounds)
+    sign = (x < 0).to(torch.uint8) << 3
+    return mag_idx.to(torch.uint8) | sign
+
+
+def pack_fp4_codes(codes: torch.Tensor) -> torch.Tensor:
+    """Pack 4-bit codes into bytes along the last dim, low nibble first."""
+    assert codes.dtype == torch.uint8 and codes.shape[-1] % 2 == 0
+    lo = codes[..., 0::2]
+    hi = codes[..., 1::2]
+    return lo | (hi << 4)
+
+
+def unpack_fp4_codes(packed: torch.Tensor) -> torch.Tensor:
+    """Inverse of :func:`pack_fp4_codes`."""
+    assert packed.dtype == torch.uint8
+    lo = packed & 0x0F
+    hi = (packed >> 4) & 0x0F
+    return torch.stack([lo, hi], dim=-1).reshape(*packed.shape[:-1], -1)
+
+
+def quantize_nvfp4_ref(
+    x: torch.Tensor,
+    per_tensor_scale: float | torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Quantize ``x [..., K]`` (K % 16 == 0) to NVFP4.
+
+    per_tensor_scale: None -> 1.0; "auto" is intentionally not supported here,
+    pass an explicit value when exercising the two-level scheme.
+
+    Returns ``(packed u8 [..., K/2], block_scale e4m3 [..., K/16], pts f32 scalar)``.
+    """
+    assert x.shape[-1] % SF_VEC_SIZE == 0, (
+        f"K={x.shape[-1]} not a multiple of {SF_VEC_SIZE}"
+    )
+    xf = x.float()
+    if per_tensor_scale is None:
+        pts = torch.tensor(1.0, dtype=torch.float32, device=x.device)
+    else:
+        pts = torch.as_tensor(per_tensor_scale, dtype=torch.float32, device=x.device)
+
+    blocks = xf.reshape(*xf.shape[:-1], -1, SF_VEC_SIZE)
+    amax = blocks.abs().amax(dim=-1)
+    scale_f32 = (amax / (E2M1_MAGNITUDES[-1] * pts)).clamp(max=E4M3_MAX)
+    scale_e4m3 = scale_f32.to(torch.float8_e4m3fn)
+    # Encode against the STORED (rounded) scale, like real quantizers.
+    scale_dec = scale_e4m3.float() * pts
+    denom = torch.where(scale_dec == 0, torch.ones_like(scale_dec), scale_dec)
+    q = blocks / denom.unsqueeze(-1)
+    codes = encode_e2m1(q).reshape(xf.shape)
+    return pack_fp4_codes(codes), scale_e4m3, pts
+
+
+def dequantize_nvfp4_ref(
+    packed: torch.Tensor,
+    block_scale: torch.Tensor,
+    per_tensor_scale: float | torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Dequantize to float32. ``packed [..., K/2]``, ``block_scale [..., K/16]``."""
+    codes = unpack_fp4_codes(packed)
+    vals = fp4_code_to_value(codes)
+    scale = block_scale.float().repeat_interleave(SF_VEC_SIZE, dim=-1)
+    out = vals * scale
+    if per_tensor_scale is not None:
+        out = out * torch.as_tensor(
+            per_tensor_scale, dtype=torch.float32, device=out.device
+        )
+    return out
+
+
+def simulate_nvfp4_quant(
+    x: torch.Tensor, per_tensor_scale: float | torch.Tensor | None = None
+) -> torch.Tensor:
+    """Quantize-dequantize roundtrip: the exact value grid the tensor core sees."""
+    packed, scale, pts = quantize_nvfp4_ref(x, per_tensor_scale)
+    return dequantize_nvfp4_ref(packed, scale, pts)
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/sf_layout.py b/src/axolotl/integrations/kernels/libs/sonicmoe/sf_layout.py
new file mode 100644
index 0000000000..5e07e85007
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/sf_layout.py
@@ -0,0 +1,158 @@
+"""Scale-factor (SF) layout prep for the SM100 blockscaled GEMM.
+
+The quack/tcgen05 kernel consumes block scales in a ``(l, rm, rk, 512)`` layout:
+each contiguous 512-byte block is one hardware-swizzled atom covering 128 MN
+rows x 4 SF columns, with row r of a tile at byte
+``(r % 32) * 16 + ((r // 32) % 4) * 4 + (sf_k % 4)``.
+
+For varlen_m (grouped, per-expert M) the SFA storage is "dQaccum padded":
+expert ``i``'s scale rows start at padded row ``(cu_seqlens[i] // 128 + i) * 128``
+and the allocation is ``ceil(total_m / 128) + (E - 1)`` row-tiles. Operand data
+(A/B/D) stays packed and unpadded; only the SF storage pads.
+
+These mirror quack's ``pack_scale_2d_to_blocked_contig`` /
+``create_blockscaled_varlen_m_operands`` (verified at quack source f4f54db0;
+runtime pin quack-kernels==0.5.0).
+"""
+
+from __future__ import annotations
+
+import torch
+
+SF_TILE_ROWS = 128
+SF_TILE_COLS = 4
+
+
+def ceil_div(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
+def pack_scales_blocked(scale_2d: torch.Tensor) -> torch.Tensor:
+    """Rearrange ``(mn, sf_k)`` or ``(l, mn, sf_k)`` scales into ``(l, rm, rk, 512)``.
+
+    Works for any 1-byte scale dtype (e4m3 for NVFP4, e8m0 for MX). Pads mn to
+    a multiple of 128 and sf_k to a multiple of 4 with zeros.
+    """
+    if scale_2d.dim() == 2:
+        scale_2d = scale_2d.unsqueeze(0)
+    assert scale_2d.dim() == 3, f"expected (l, mn, sf_k), got {tuple(scale_2d.shape)}"
+    assert scale_2d.element_size() == 1, "scale dtype must be 1 byte (e4m3/e8m0/u8)"
+    orig_dtype = scale_2d.dtype
+    l, mn, sf_k = scale_2d.shape
+    rm = ceil_div(mn, SF_TILE_ROWS)
+    rk = ceil_div(sf_k, SF_TILE_COLS)
+    mn_pad, sf_k_pad = rm * SF_TILE_ROWS, rk * SF_TILE_COLS
+    u8 = scale_2d.contiguous().view(torch.uint8)
+    if (mn_pad, sf_k_pad) != (mn, sf_k):
+        padded = torch.zeros(l, mn_pad, sf_k_pad, device=u8.device, dtype=torch.uint8)
+        padded[:, :mn, :sf_k] = u8
+    else:
+        padded = u8
+    blocks = padded.view(l, rm, SF_TILE_ROWS, rk, SF_TILE_COLS).permute(0, 1, 3, 2, 4)
+    # split the 128 rows into (4 outer, 32 inner), then swap to (32, 4)
+    blocks = blocks.reshape(l, rm, rk, 4, 32, SF_TILE_COLS).transpose(3, 4).contiguous()
+    return blocks.view(l, rm, rk, 512).view(orig_dtype)
+
+
+def varlen_padded_num_row_tiles(total_m: int, num_experts: int) -> int:
+    return ceil_div(total_m, SF_TILE_ROWS) + (num_experts - 1)
+
+
+def build_varlen_sfa(
+    scale_rows: torch.Tensor, cu_seqlens: torch.Tensor
+) -> torch.Tensor:
+    """Build the dQaccum-padded SFA for varlen_m.
+
+    scale_rows: ``(total_m, sf_k)`` per-row block scales, expert-sorted and packed
+    (row order matching the A operand). cu_seqlens: ``(E+1,)`` int.
+
+    Returns ``(1, total_padded_rm, rk, 512)`` in the blocked layout, where
+    expert ``i``'s rows sit at padded tile offset ``cu_seqlens[i] // 128 + i``.
+    """
+    assert scale_rows.dim() == 2
+    total_m, sf_k = scale_rows.shape
+    num_experts = cu_seqlens.numel() - 1
+    total_padded_rm = varlen_padded_num_row_tiles(total_m, num_experts)
+    padded = torch.zeros(
+        total_padded_rm * SF_TILE_ROWS,
+        sf_k,
+        dtype=scale_rows.dtype,
+        device=scale_rows.device,
+    )
+
+    if scale_rows.is_cuda:
+        # One scatter instead of a host-synced per-expert copy loop.
+        cu = cu_seqlens.to(device=scale_rows.device, dtype=torch.long)
+        starts = cu[:-1]
+        counts = cu[1:] - starts
+        seg = torch.repeat_interleave(
+            torch.arange(num_experts, device=scale_rows.device), counts
+        )
+        row_in_seg = torch.arange(total_m, device=scale_rows.device) - starts[seg]
+        dest = (starts[seg] // SF_TILE_ROWS + seg) * SF_TILE_ROWS + row_in_seg
+        padded[dest] = scale_rows
+        return pack_scales_blocked(padded.unsqueeze(0))
+
+    cu = cu_seqlens.tolist()
+    assert cu[0] == 0 and cu[-1] == total_m, (
+        f"bad cu_seqlens {cu} for total_m={total_m}"
+    )
+    for i in range(num_experts):
+        start, end = cu[i], cu[i + 1]
+        row0 = (start // SF_TILE_ROWS + i) * SF_TILE_ROWS
+        padded[row0 : row0 + (end - start)] = scale_rows[start:end]
+    return pack_scales_blocked(padded.unsqueeze(0))
+
+
+def fold_per_tensor_scale(
+    block_scale: torch.Tensor,
+    per_tensor_scale: torch.Tensor,
+    *,
+    allow_underflow: bool = False,
+) -> tuple[torch.Tensor, float]:
+    """Fold a per-expert fp32 scale into e4m3 block scales: ``scale * pts -> e4m3``.
+
+    block_scale: ``(E, N, sf_k)`` float8_e4m3fn. per_tensor_scale: ``(E,)``,
+    ``(E,1,1)`` or scalar fp32. Returns ``(folded e4m3, max relative roundtrip
+    error)``. Raises if any folded value saturates e4m3 (|v| > 448) or, unless
+    ``allow_underflow``, rounds a nonzero scale to zero (an underflowed block
+    dequantizes to exact zeros and counts as rel err 1.0); in that case pts
+    must stay outside the block scale (per-expert alpha, a kernel epilogue
+    change).
+    """
+    e4m3_max = 448.0
+    pts = per_tensor_scale.float().reshape(-1)
+    if pts.numel() == 1:
+        pts = pts.expand(block_scale.shape[0])
+    assert pts.numel() == block_scale.shape[0], (
+        f"pts has {pts.numel()} entries for {block_scale.shape[0]} experts"
+    )
+    folded_f32 = block_scale.float() * pts.view(-1, 1, 1)
+    if bool((folded_f32.abs() > e4m3_max).any()):
+        raise ValueError("per_tensor_scale folding saturates e4m3 (>448)")
+    folded = folded_f32.to(torch.float8_e4m3fn)
+    roundtrip = folded.float()
+    if not allow_underflow and bool(((roundtrip == 0) & (folded_f32 != 0)).any()):
+        raise ValueError("per_tensor_scale folding underflows e4m3 to zero")
+    nonzero = folded_f32 != 0
+    rel_err = 0.0
+    if bool(nonzero.any()):
+        rel_err = float(
+            (
+                (roundtrip[nonzero] - folded_f32[nonzero]).abs()
+                / folded_f32[nonzero].abs()
+            ).max()
+        )
+    return folded, rel_err
+
+
+def gate_up_interleave_perm(n2: int, device=None) -> torch.Tensor:
+    """Row permutation mapping concat ``[gate; up]`` (N = 2I) to interleaved
+    ``[g0, u0, g1, u1, ...]``, the pairing quack's gated epilogue consumes
+    (postact col j = act(row 2j, row 2j+1))."""
+    assert n2 % 2 == 0
+    i = n2 // 2
+    perm = torch.empty(n2, dtype=torch.long, device=device)
+    perm[0::2] = torch.arange(i, device=device)
+    perm[1::2] = torch.arange(i, device=device) + i
+    return perm
diff --git a/src/axolotl/integrations/kernels/libs/sonicmoe/triton_nvfp4.py b/src/axolotl/integrations/kernels/libs/sonicmoe/triton_nvfp4.py
new file mode 100644
index 0000000000..9c692ce8a8
--- /dev/null
+++ b/src/axolotl/integrations/kernels/libs/sonicmoe/triton_nvfp4.py
@@ -0,0 +1,340 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Triton NVFP4 codecs for the sonicmoe grouped path.
+
+Two memory-bound kernels replacing multi-kernel pure-torch chains:
+
+- :func:`dequant_nvfp4_triton`: packed E2M1 + e4m3 block scales (+ optional
+  per-expert scale) -> bf16, one kernel. torchao's ``NVFP4Tensor.dequantize()``
+  does this through a long-index table gather over every element and dominates
+  the fp4_cute backward (the dX path dequantizes the whole weight).
+- :func:`quantize_nvfp4_triton`: bf16 rows -> (packed u8, e4m3 block scales),
+  one kernel. Replaces the ~10-kernel ``quantize_nvfp4_ref`` chain on the
+  activation-quant hot path. Same rounding rules as the reference (scale =
+  amax/6 clamped to 448, encode against the STORED e4m3 scale, ties at the
+  E2M1 midpoints round down), so the fp32 oracle stays valid.
+
+Both fall back to the callers' torch paths when triton is unavailable.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import torch
+
+
+@functools.lru_cache(maxsize=1)
+def triton_available() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        import triton  # noqa: F401
+        import triton.language as tl  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
+def _kernels():
+    import triton
+    import triton.language as tl
+
+    @triton.jit
+    def _e2m1_value(m):
+        # m in [0, 8): (0, 0.5, 1, 1.5, 2, 3, 4, 6)
+        mant = tl.where(m % 2 == 1, 1.5, 1.0)
+        exp = (m // 2 - 1).to(tl.float32)
+        val = mant * tl.exp2(exp)
+        val = tl.where(m == 0, 0.0, val)
+        val = tl.where(m == 1, 0.5, val)
+        return val
+
+    @triton.jit
+    def _dequant_kernel(
+        q_ptr,
+        s_ptr,
+        pts_ptr,
+        out_ptr,
+        rows_per_expert,
+        K2: tl.constexpr,  # bytes per row (K/2)
+        SF_K: tl.constexpr,  # scales per row (K/16)
+        HAS_PTS: tl.constexpr,
+        BLOCK_B: tl.constexpr,  # bytes per program
+    ):
+        row = tl.program_id(0)
+        blk = tl.program_id(1)
+        offs_b = blk * BLOCK_B + tl.arange(0, BLOCK_B)
+        mask_b = offs_b < K2
+        b = tl.load(q_ptr + row * K2 + offs_b, mask=mask_b, other=0).to(tl.uint8)
+
+        lo = (b & 0x0F).to(tl.int32)
+        hi = ((b >> 4) & 0x0F).to(tl.int32)
+        v_lo = _e2m1_value(lo & 7) * tl.where(lo >= 8, -1.0, 1.0)
+        v_hi = _e2m1_value(hi & 7) * tl.where(hi >= 8, -1.0, 1.0)
+
+        # byte i covers values 2i / 2i+1 -> scale index (2i)//16 == i//8
+        offs_s = offs_b // 8
+        sc_u8 = tl.load(s_ptr + row * SF_K + offs_s, mask=offs_s < SF_K, other=0)
+        sc = sc_u8.to(tl.uint8).to(tl.float8e4nv, bitcast=True).to(tl.float32)
+        if HAS_PTS:
+            sc = sc * tl.load(pts_ptr + row // rows_per_expert).to(tl.float32)
+
+        out_lo = (v_lo * sc).to(out_ptr.dtype.element_ty)
+        out_hi = (v_hi * sc).to(out_ptr.dtype.element_ty)
+        tl.store(out_ptr + row * (2 * K2) + 2 * offs_b, out_lo, mask=mask_b)
+        tl.store(out_ptr + row * (2 * K2) + 2 * offs_b + 1, out_hi, mask=mask_b)
+
+    @triton.jit
+    def _quant_kernel(
+        x_ptr,
+        q_ptr,
+        s_ptr,
+        K: tl.constexpr,
+        SF_K: tl.constexpr,
+        BLOCK_SF: tl.constexpr,  # scale blocks per program
+    ):
+        row = tl.program_id(0)
+        blk = tl.program_id(1)
+        offs_sf = blk * BLOCK_SF + tl.arange(0, BLOCK_SF)
+        mask_sf = offs_sf < SF_K
+        # [BLOCK_SF, 8, 2] element offsets: (block, byte-in-block, nibble)
+        offs_v = (
+            offs_sf[:, None, None] * 16
+            + tl.arange(0, 8)[None, :, None] * 2
+            + tl.arange(0, 2)[None, None, :]
+        )
+        mask_v = mask_sf[:, None, None]
+        x = tl.load(x_ptr + row * K + offs_v, mask=mask_v, other=0.0).to(tl.float32)
+
+        amax = tl.max(tl.max(tl.abs(x), axis=2), axis=1)
+        scale = tl.minimum(amax / 6.0, 448.0)
+        scale_e4m3 = scale.to(tl.float8e4nv)
+        tl.store(
+            s_ptr + row * SF_K + offs_sf,
+            scale_e4m3.to(tl.uint8, bitcast=True),
+            mask=mask_sf,
+        )
+        # encode against the STORED (rounded) scale, like the reference.
+        # div_rn: triton's `/` may lower to reciprocal-multiply, which flips
+        # values sitting exactly on E2M1 bucket boundaries vs the torch ref.
+        sdec = scale_e4m3.to(tl.float32)
+        sdec = tl.where(sdec == 0, 1.0, sdec)
+        q = tl.math.div_rn(x, tl.broadcast_to(sdec[:, None, None], x.shape))
+
+        a = tl.abs(q)
+        idx = (
+            (a > 0.25).to(tl.int32)
+            + (a > 0.75).to(tl.int32)
+            + (a > 1.25).to(tl.int32)
+            + (a > 1.75).to(tl.int32)
+            + (a > 2.5).to(tl.int32)
+            + (a > 3.5).to(tl.int32)
+            + (a > 5.0).to(tl.int32)
+        )
+        code = tl.where(q < 0, idx + 8, idx)
+
+        # pack low nibble first: byte = lo + hi * 16, via a weighted sum over
+        # the nibble axis (triton tensors do not support axis slicing)
+        weight = tl.where(tl.arange(0, 2)[None, None, :] == 0, 1, 16)
+        packed = tl.sum(code * weight, axis=2).to(tl.uint8)
+        offs_p = offs_sf[:, None] * 8 + tl.arange(0, 8)[None, :]
+        tl.store(q_ptr + row * (K // 2) + offs_p, packed, mask=mask_sf[:, None])
+
+    @triton.jit
+    def _quant_sfa_kernel(
+        x_ptr,
+        q_ptr,
+        sfa_ptr,
+        dest_ptr,  # [T] padded destination row per source row
+        K: tl.constexpr,
+        SF_K: tl.constexpr,
+        RK: tl.constexpr,  # ceil(SF_K / 4), 512-byte tiles per row-tile
+        BLOCK_SF: tl.constexpr,
+    ):
+        # identical quantization math to _quant_kernel, but the e4m3 scales are
+        # stored straight into the dQaccum-padded swizzled SFA layout: padded
+        # row p, sf col c -> tile (p//128, c//4), byte (p%32)*16 + ((p//32)%4)*4 + c%4
+        row = tl.program_id(0)
+        blk = tl.program_id(1)
+        offs_sf = blk * BLOCK_SF + tl.arange(0, BLOCK_SF)
+        mask_sf = offs_sf < SF_K
+        offs_v = (
+            offs_sf[:, None, None] * 16
+            + tl.arange(0, 8)[None, :, None] * 2
+            + tl.arange(0, 2)[None, None, :]
+        )
+        x = tl.load(
+            x_ptr + row * K + offs_v, mask=mask_sf[:, None, None], other=0.0
+        ).to(tl.float32)
+
+        amax = tl.max(tl.max(tl.abs(x), axis=2), axis=1)
+        scale = tl.minimum(amax / 6.0, 448.0)
+        scale_e4m3 = scale.to(tl.float8e4nv)
+
+        p = tl.load(dest_ptr + row)
+        addr = (
+            (p // 128) * (RK * 512)
+            + (offs_sf // 4) * 512
+            + (p % 32) * 16
+            + ((p // 32) % 4) * 4
+            + (offs_sf % 4)
+        )
+        tl.store(sfa_ptr + addr, scale_e4m3.to(tl.uint8, bitcast=True), mask=mask_sf)
+
+        sdec = scale_e4m3.to(tl.float32)
+        sdec = tl.where(sdec == 0, 1.0, sdec)
+        q = tl.math.div_rn(x, tl.broadcast_to(sdec[:, None, None], x.shape))
+
+        a = tl.abs(q)
+        idx = (
+            (a > 0.25).to(tl.int32)
+            + (a > 0.75).to(tl.int32)
+            + (a > 1.25).to(tl.int32)
+            + (a > 1.75).to(tl.int32)
+            + (a > 2.5).to(tl.int32)
+            + (a > 3.5).to(tl.int32)
+            + (a > 5.0).to(tl.int32)
+        )
+        code = tl.where(q < 0, idx + 8, idx)
+        weight = tl.where(tl.arange(0, 2)[None, None, :] == 0, 1, 16)
+        packed = tl.sum(code * weight, axis=2).to(tl.uint8)
+        offs_p = offs_sf[:, None] * 8 + tl.arange(0, 8)[None, :]
+        tl.store(q_ptr + row * (K // 2) + offs_p, packed, mask=mask_sf[:, None])
+
+    @triton.jit
+    def _rowscale_kernel(
+        x_ptr,
+        pts_ptr,  # fp32 [rows]
+        N: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+    ):
+        # in-place x[r, :] = bf16(f32(x[r, :]) * pts[r]); one pass instead of
+        # float() -> mul -> to(bf16)
+        row = tl.program_id(0)
+        blk = tl.program_id(1)
+        offs = blk * BLOCK_N + tl.arange(0, BLOCK_N)
+        mask = offs < N
+        x = tl.load(x_ptr + row * N + offs, mask=mask, other=0.0).to(tl.float32)
+        p = tl.load(pts_ptr + row)
+        tl.store(x_ptr + row * N + offs, (x * p).to(x_ptr.dtype.element_ty), mask=mask)
+
+    return _dequant_kernel, _quant_kernel, _quant_sfa_kernel, _rowscale_kernel
+
+
+@functools.lru_cache(maxsize=1)
+def _get_kernels():
+    return _kernels()
+
+
+def dequant_nvfp4_triton(
+    qdata: torch.Tensor,
+    scale: torch.Tensor,
+    per_tensor_scale: torch.Tensor | None,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """``qdata [..., K/2] u8`` + ``scale [..., K/16] e4m3`` -> ``[..., K]``.
+
+    ``per_tensor_scale``, if given, has one entry per leading-dim slice (the
+    ``[E, N, K]`` expert convention: entry ``e`` scales rows of slice ``e``).
+    """
+    dequant_kernel, _, _, _ = _get_kernels()
+    if qdata.dtype != torch.uint8:
+        qdata = qdata.view(torch.uint8)
+    k2 = qdata.shape[-1]
+    q2 = qdata.reshape(-1, k2).contiguous()
+    s2 = scale.reshape(-1, scale.shape[-1]).contiguous().view(torch.uint8)
+    rows = q2.shape[0]
+    rows_per_expert = (
+        rows // per_tensor_scale.numel() if per_tensor_scale is not None else rows
+    )
+    pts = (
+        per_tensor_scale.reshape(-1).float().contiguous()
+        if per_tensor_scale is not None
+        else q2.new_zeros(1, dtype=torch.float32)
+    )
+    out = torch.empty(rows, 2 * k2, dtype=out_dtype, device=qdata.device)
+    BLOCK_B = 512
+    grid = (rows, (k2 + BLOCK_B - 1) // BLOCK_B)
+    dequant_kernel[grid](
+        q2,
+        s2,
+        pts,
+        out,
+        rows_per_expert,
+        K2=k2,
+        SF_K=scale.shape[-1],
+        HAS_PTS=per_tensor_scale is not None,
+        BLOCK_B=BLOCK_B,
+    )
+    return out.view(*qdata.shape[:-1], 2 * k2)
+
+
+def rowscale_inplace_triton(x: torch.Tensor, row_scale: torch.Tensor) -> torch.Tensor:
+    """In-place ``x[r] = x.dtype(f32(x[r]) * row_scale[r])``; bit-identical to
+    ``(x.float() * row_scale[:, None]).to(x.dtype)`` in one pass."""
+    _, _, _, rowscale_kernel = _get_kernels()
+    assert x.dim() == 2 and x.is_contiguous()
+    rows, n = x.shape
+    BLOCK_N = 1024
+    grid = (rows, (n + BLOCK_N - 1) // BLOCK_N)
+    rowscale_kernel[grid](x, row_scale.contiguous(), N=n, BLOCK_N=BLOCK_N)
+    return x
+
+
+def quantize_nvfp4_triton(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """``x [T, K]`` (K % 16 == 0) -> ``(packed u8 [T, K/2], scale e4m3 [T, K/16])``."""
+    _, quant_kernel, _, _ = _get_kernels()
+    assert x.dim() == 2 and x.shape[-1] % 16 == 0
+    t, k = x.shape
+    sf_k = k // 16
+    x = x.contiguous()
+    packed = torch.empty(t, k // 2, dtype=torch.uint8, device=x.device)
+    scale_u8 = torch.empty(t, sf_k, dtype=torch.uint8, device=x.device)
+    BLOCK_SF = 64
+    grid = (t, (sf_k + BLOCK_SF - 1) // BLOCK_SF)
+    quant_kernel[grid](x, packed, scale_u8, K=k, SF_K=sf_k, BLOCK_SF=BLOCK_SF)
+    return packed, scale_u8.view(torch.float8_e4m3fn)
+
+
+def quantize_rows_fused_sfa_triton(
+    x: torch.Tensor, cu_seqlens: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize expert-sorted rows AND write the dQaccum-padded swizzled SFA in
+    one kernel: ``(packed u8 [T, K/2], sfa e4m3 (1, rm, rk, 512))``.
+
+    Byte-identical to ``quantize_nvfp4_triton`` + ``sf_layout.build_varlen_sfa``,
+    minus the scatter and the pack/permute chain. Requires ``(K/16) % 4 == 0``.
+    """
+    from .sf_layout import SF_TILE_ROWS, varlen_padded_num_row_tiles
+
+    _, _, quant_sfa_kernel, _ = _get_kernels()
+    assert x.dim() == 2 and x.shape[-1] % 16 == 0
+    t, k = x.shape
+    sf_k = k // 16
+    assert sf_k % 4 == 0
+    rk = sf_k // 4
+    num_experts = cu_seqlens.numel() - 1
+
+    cu = cu_seqlens.to(device=x.device, dtype=torch.long)
+    starts = cu[:-1]
+    counts = cu[1:] - starts
+    seg = torch.repeat_interleave(torch.arange(num_experts, device=x.device), counts)
+    dest = (
+        (starts[seg] // SF_TILE_ROWS + seg) * SF_TILE_ROWS
+        + torch.arange(t, device=x.device)
+        - starts[seg]
+    ).to(torch.int32)
+
+    total_rm = varlen_padded_num_row_tiles(t, num_experts)
+    x = x.contiguous()
+    packed = torch.empty(t, k // 2, dtype=torch.uint8, device=x.device)
+    sfa = torch.zeros(1, total_rm, rk, 512, dtype=torch.uint8, device=x.device)
+    BLOCK_SF = 64
+    grid = (t, (sf_k + BLOCK_SF - 1) // BLOCK_SF)
+    quant_sfa_kernel[grid](
+        x, packed, sfa, dest, K=k, SF_K=sf_k, RK=rk, BLOCK_SF=BLOCK_SF
+    )
+    return packed, sfa.view(torch.float8_e4m3fn)
diff --git a/src/axolotl/integrations/kernels/plugin.py b/src/axolotl/integrations/kernels/plugin.py
new file mode 100644
index 0000000000..628e10cf7c
--- /dev/null
+++ b/src/axolotl/integrations/kernels/plugin.py
@@ -0,0 +1,171 @@
+import importlib
+import os
+
+import torch
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.integrations.kernels.adapters import get_active_adapters
+from axolotl.integrations.kernels.quant_training_guard import (
+    relax_quantized_training_guard,
+)
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _check_sonicmoe_gpu_compat():
+    """Validate GPU compute capability for SonicMoE and configure env.
+
+    Supported: Hopper (sm_90), Blackwell (sm_100 - sm_103).
+    B300 (sm_103) additionally requires Triton 3.6.0.
+    """
+    if not torch.cuda.is_available():
+        return
+
+    cc = torch.cuda.get_device_capability()
+
+    if cc < (9, 0):
+        raise RuntimeError(
+            f"SonicMoE requires Hopper (sm_90) or Blackwell (sm_100+) GPU, "
+            f"but detected sm_{cc[0]}{cc[1]}."
+        )
+
+    if cc > (10, 3):
+        raise RuntimeError(
+            f"SonicMoE does not yet support sm_{cc[0]}{cc[1]}. "
+            f"Supported: Hopper (sm_90) and Blackwell (sm_100 - sm_103)."
+        )
+
+    if cc >= (10, 0):
+        os.environ.setdefault("USE_QUACK_GEMM", "1")
+        LOG.info(
+            f"Blackwell GPU (sm_{cc[0]}{cc[1]}) detected, enabling USE_QUACK_GEMM=1"
+        )
+
+    if cc == (10, 3):
+        triton_spec = importlib.util.find_spec("triton")
+        if triton_spec is None:
+            raise RuntimeError(
+                "B300 (sm_103) requires Triton 3.6.0, but Triton is not installed."
+            )
+        import triton
+
+        triton_version = tuple(int(x) for x in triton.__version__.split(".")[:2])
+        if triton_version != (3, 6):
+            raise RuntimeError(
+                f"B300 (sm_103) requires Triton 3.6.x, but found {triton.__version__}."
+            )
+
+
+class KernelsPlugin(BasePlugin):
+    """Thin orchestrator: registers the expert-kernel backend and dispatches model-family
+    specifics to ``ModelAdapter`` subclasses (see ``adapters/``)."""
+
+    def get_input_args(self):
+        return "axolotl.integrations.kernels.KernelsArgs"
+
+    def _adapters(self, cfg):
+        # Cache: matching can be expensive (e.g. AutoConfig.from_pretrained for Gemma-4).
+        cached = getattr(self, "_cached_adapters", None)
+        if cached is None:
+            cached = get_active_adapters(cfg)
+            self._cached_adapters = cached
+        return cached
+
+    def pre_model_load(self, cfg):
+        """Register the expert-kernel backend + generic capabilities, then run adapter hooks.
+
+        Architecture-agnostic: routing stays in each model's SparseMoEBlock; only the experts
+        call is dispatched through the registry. When EP is active the ExpertParallelPlugin owns
+        ``experts_implementation`` (a ``deep_ep_*`` composite), so we don't overwrite it here.
+        """
+        ep_active = (getattr(cfg, "expert_parallel_size", 1) or 1) > 1
+
+        if cfg.use_scattermoe:
+            from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+                register_scattermoe_experts,
+            )
+            from axolotl.integrations.kernels.libs.scattermoe_lora.runtime import (
+                configure_scattermoe_runtime,
+            )
+
+            register_scattermoe_experts()
+            if not ep_active:
+                cfg.experts_implementation = "scattermoe"
+            LOG.info("Registered 'scattermoe' in transformers ExpertsInterface")
+
+            # LoRA on a frozen pre-quantized (FP8/NVFP4) base is supported but rejected by the
+            # upstream guard; scope-relax it (skips only PEFT/quantized, delegates the rest).
+            relax_quantized_training_guard()
+
+            # Apply ALL ScatterMoE runtime settings from this run's config through one entry point
+            # (it resets first, so a long-lived multi-run process can't inherit stale state).
+            configure_scattermoe_runtime(cfg)
+            if cfg.get("dsv4_fp4_grouped_mode"):
+                LOG.info(
+                    "Enabled grouped fp4 MoE path: dsv4_fp4_grouped_mode=%s",
+                    cfg.get("dsv4_fp4_grouped_mode"),
+                )
+                if cfg.get("moe_grouped_backend"):
+                    LOG.info(
+                        "Grouped MoE base-GEMM backend override: %s",
+                        cfg.get("moe_grouped_backend"),
+                    )
+        elif cfg.use_sonicmoe:
+            _check_sonicmoe_gpu_compat()
+
+            from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+                register_sonicmoe_experts,
+            )
+
+            register_sonicmoe_experts()
+            if not ep_active:
+                cfg.experts_implementation = "sonicmoe"
+            LOG.info("Registered 'sonicmoe' in transformers ExpertsInterface")
+
+            # Same frozen-quantized-base + LoRA pattern as the scattermoe path above.
+            relax_quantized_training_guard()
+
+        adapters = self._adapters(cfg)
+        self._warn_unclaimed_nonexpert_quantization(cfg, adapters)
+        for adapter in adapters:
+            adapter.pre_model_load(cfg)
+
+    @staticmethod
+    def _warn_unclaimed_nonexpert_quantization(cfg, adapters):
+        """Warn if a non-expert quantization policy is set but no active adapter consumes it.
+
+        ``nonexpert_quantization`` is a global intent, but only some model adapters act on it
+        (e.g. Gemma-4). Without this, configuring it on an unsupported model silently no-ops.
+        ``none``/``bf16`` mean "no quantization", so they're never considered unclaimed.
+        """
+        policy = cfg.get("nonexpert_quantization")
+        if not policy or str(policy).lower() in ("none", "bf16"):
+            return
+        if any(a.consumes_nonexpert_quantization(cfg) for a in adapters):
+            return
+        LOG.warning(
+            "nonexpert_quantization=%r is set but no active model adapter consumes it "
+            "(active adapters: %s); it will have no effect. It is currently implemented only "
+            "for Gemma-4 NVFP4 checkpoints.",
+            policy,
+            [a.name for a in adapters] or "none",
+        )
+
+    def pre_lora_load(self, cfg, model):
+        for adapter in self._adapters(cfg):
+            adapter.pre_lora_load(cfg, model)
+
+    def post_model_load(self, cfg, model):
+        for adapter in self._adapters(cfg):
+            adapter.post_model_load(cfg, model)
+
+    def add_callbacks_pre_trainer(self, cfg, model):
+        callbacks = []
+        if cfg.use_scattermoe:
+            from axolotl.integrations.kernels.autotune_callback import (
+                AutotuneReportCallback,
+            )
+
+            callbacks.append(AutotuneReportCallback())
+        return callbacks
diff --git a/src/axolotl/integrations/kernels/quant_training_guard.py b/src/axolotl/integrations/kernels/quant_training_guard.py
new file mode 100644
index 0000000000..935d878a75
--- /dev/null
+++ b/src/axolotl/integrations/kernels/quant_training_guard.py
@@ -0,0 +1,67 @@
+"""Scoped relaxation of transformers' quantized-training guard.
+
+``transformers.trainer.validate_quantization_for_training`` rejects pre-quantized FP8/NVFP4
+checkpoints for training. Its FP8/NVFP4 branch fires even when LoRA adapters are attached and the
+quantized base is frozen — the supported QLoRA-style pattern. Rather than globally no-op'ing the
+guard for all runs, we wrap it: for PEFT models (frozen quantized base + trainable adapters) we
+skip, and for everything else we delegate to the original guard unchanged. Idempotent; preserves
+the original callable for restore.
+"""
+
+from __future__ import annotations
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_ORIG_ATTR = "_axolotl_quant_guard_original"
+_FN_NAME = "validate_quantization_for_training"
+
+
+def _target_module(module):
+    if module is not None:
+        return module
+    import transformers.trainer as hf_trainer
+
+    return hf_trainer
+
+
+def relax_quantized_training_guard(module=None) -> None:
+    """Install the scoped guard wrapper (idempotent).
+
+    ``module`` defaults to ``transformers.trainer``; tests pass a throwaway module to avoid
+    touching global state.
+    """
+    mod = _target_module(module)
+    current = getattr(mod, _FN_NAME, None)
+    if current is None or getattr(current, _ORIG_ATTR, None) is not None:
+        return  # missing or already wrapped
+
+    original = current
+
+    def guarded(model):
+        # LoRA/PEFT on a frozen quantized base is the supported (QLoRA-style) pattern; the upstream
+        # FP8/NVFP4 branch rejects it anyway. Skip only for PEFT models; delegate the rest.
+        # _hf_peft_config_loaded is set by HF-native model.add_adapter(); axolotl applies LoRA via
+        # get_peft_model() -> PeftModel, which exposes .peft_config but NOT that flag. Check both, or
+        # FP8 checkpoints (e.g. DeepSeek-V4-Flash-NVFP4) get rejected despite frozen-base + adapters.
+        if getattr(model, "_hf_peft_config_loaded", False) or getattr(
+            model, "peft_config", None
+        ):
+            return None
+        return original(model)
+
+    setattr(guarded, _ORIG_ATTR, original)
+    setattr(mod, _FN_NAME, guarded)
+    LOG.info(
+        "kernels: scoped quantized-training guard (skips only PEFT/quantized; delegates others)"
+    )
+
+
+def restore_quantized_training_guard(module=None) -> None:
+    """Restore the original guard (for tests/teardown)."""
+    mod = _target_module(module)
+    current = getattr(mod, _FN_NAME, None)
+    original = getattr(current, _ORIG_ATTR, None)
+    if original is not None:
+        setattr(mod, _FN_NAME, original)
diff --git a/src/axolotl/integrations/liger/LICENSE b/src/axolotl/integrations/liger/LICENSE
new file mode 100644
index 0000000000..d645695673
--- /dev/null
+++ b/src/axolotl/integrations/liger/LICENSE
@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
diff --git a/src/axolotl/integrations/liger/README.md b/src/axolotl/integrations/liger/README.md
new file mode 100644
index 0000000000..3a2d4bd04b
--- /dev/null
+++ b/src/axolotl/integrations/liger/README.md
@@ -0,0 +1,59 @@
+# Liger Kernel Integration
+
+Liger Kernel provides efficient Triton kernels for LLM training, offering:
+
+- 20% increase in multi-GPU training throughput
+- 60% reduction in memory usage
+- Compatibility with both FSDP and DeepSpeed
+
+See https://github.com/linkedin/Liger-Kernel
+
+## Usage
+
+```yaml
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+
+# FLCE-specific
+liger_use_token_scaling: true
+```
+
+## Supported Models
+
+- deepseek_v2
+- gemma
+- gemma2
+- gemma3
+- granite
+- jamba
+- llama
+- mistral
+- mixtral
+- mllama
+- mllama_text_model
+- olmo2
+- paligemma
+- phi3
+- qwen2
+- qwen2_5_vl
+- qwen2_vl
+
+## Citation
+
+```bib
+@article{hsu2024ligerkernelefficienttriton,
+      title={Liger Kernel: Efficient Triton Kernels for LLM Training},
+      author={Pin-Lun Hsu and Yun Dai and Vignesh Kothapalli and Qingquan Song and Shao Tang and Siyu Zhu and Steven Shimizu and Shivam Sahni and Haowen Ning and Yanning Chen},
+      year={2024},
+      eprint={2410.10989},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2410.10989},
+      journal={arXiv preprint arXiv:2410.10989},
+}
+```
diff --git a/src/axolotl/integrations/liger/__init__.py b/src/axolotl/integrations/liger/__init__.py
new file mode 100644
index 0000000000..c20f4545cb
--- /dev/null
+++ b/src/axolotl/integrations/liger/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Module for the Plugin for LIGER integraton with Axolotl.
+
+Liger Kernel is the collection of Triton-native kernels for LLM Training.
+It is designed to be performant, correct, and light-weight.
+"""
+
+from .args import LigerArgs
+from .plugin import LigerPlugin
+
+__all__ = [
+    "LigerArgs",
+    "LigerPlugin",
+]
diff --git a/src/axolotl/integrations/liger/args.py b/src/axolotl/integrations/liger/args.py
new file mode 100644
index 0000000000..283e527b4e
--- /dev/null
+++ b/src/axolotl/integrations/liger/args.py
@@ -0,0 +1,122 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Module for handling LIGER input arguments.
+"""
+
+from pydantic import BaseModel, Field, model_validator
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class LigerArgs(BaseModel):
+    """
+    Input args for LIGER.
+    """
+
+    liger_rope: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Enables Liger's fused RoPE kernel. For Qwen2-VL / Qwen2.5-VL / "
+                "Qwen3-VL (text and VL model_config_types) this auto-defaults to "
+                "True when unset, swapping in the fused multimodal/rotary kernel."
+            )
+        },
+    )
+    liger_rms_norm: bool | None = None
+    liger_rms_norm_gated: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Enables fused RMSNorm+SiLU gate Triton kernel for models with "
+                "gated RMSNorm (e.g. Qwen3.5 / Qwen3.5 MoE linear attention layers)."
+            )
+        },
+    )
+    liger_layer_norm: bool | None = None
+    liger_swiglu: bool | None = None
+    liger_glu_activation: bool | None = None
+    liger_cross_entropy: bool | None = None
+    liger_fused_linear_cross_entropy: bool | None = None
+    liger_use_token_scaling: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Enables use_token_scaling in fused_linear_cross_entropy. "
+                "When True, each token's loss is multiplied by its predicted probability (detached from gradients)."
+            )
+        },
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_deprecated_swiglu(cls, data):
+        if data.get("liger_swiglu") is not None:
+            if data.get("liger_glu_activation") is not None:
+                raise ValueError(
+                    "You cannot have both `liger_swiglu` and `liger_glu_activation` set."
+                )
+
+            LOG.warning(
+                "The 'liger_swiglu' argument is deprecated and will be removed in a future release. "
+                "Please use 'liger_glu_activation' instead."
+            )
+            data["liger_glu_activation"] = data.pop("liger_swiglu")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_tiled_mlp_conflict(cls, data):
+        if (
+            data.get("liger_glu_activation") is True
+            and data.get("tiled_mlp") is True
+            and not data.get("tiled_mlp_use_original_mlp")
+        ):
+            raise ValueError(
+                "You cannot have both `liger_glu_activation` and `tiled_mlp` set without `tiled_mlp_use_original_mlp: true`."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_liger_rms_norm_tensor_parallel(cls, data):
+        if data.get("liger_rms_norm") and data.get("tensor_parallel_size", 1) > 1:
+            raise ValueError(
+                "`liger_rms_norm` is incompatible with tensor parallelism, "
+                "see https://github.com/linkedin/Liger-Kernel/issues/826"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_liger_use_token_scaling_flce(cls, data):
+        if data.get("liger_use_token_scaling") and not data.get(
+            "liger_fused_linear_cross_entropy"
+        ):
+            raise ValueError(
+                "`liger_use_token_scaling: true` requires `liger_fused_linear_cross_entropy` enabled."
+            )
+
+        return data
+
+    @model_validator(mode="after")
+    def check_tensor_parallel_size_liger_fused_linear_cross_entropy(self):
+        # TODO @SalmanMohammadi this is a larger fix - investigate
+        if self.tensor_parallel_size > 1 and self.liger_fused_linear_cross_entropy:
+            raise ValueError("Tensor parallelism is not compatible with liger losses.")
+        return self
diff --git a/src/axolotl/integrations/liger/models/__init__.py b/src/axolotl/integrations/liger/models/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/integrations/liger/models/base.py b/src/axolotl/integrations/liger/models/base.py
new file mode 100644
index 0000000000..a9dbe94126
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/base.py
@@ -0,0 +1,188 @@
+"""
+Generic FLCE patch for untested models similar to Llama
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from liger_kernel.transformers.trainer.orpo_trainer import _FSDPForwardRedirection
+from liger_kernel.utils import PEFT_AVAILABLE
+from peft.utils import ModulesToSaveWrapper
+from torch.distributed.fsdp import FullyShardedDataParallel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
+
+
+def lce_forward(
+    self,
+    *args,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    labels: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
+    **kwargs,
+) -> Union[Tuple, CausalLMOutputWithPast]:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+    """
+
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+
+    return_dict = (
+        return_dict if return_dict is not None else self.config.use_return_dict
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        *args,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = (
+        slice(-logits_to_keep, None)
+        if isinstance(logits_to_keep, int)
+        else logits_to_keep
+    )
+    kept_hidden_states = hidden_states[:, slice_indices, :]
+
+    shift_labels = kwargs.pop("shift_labels", None)
+    logits = None
+    loss = None
+
+    # if in training mode, don't materialize logits
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    if skip_logits:
+        loss = lce_maybe_trainable_lm_head(
+            self,
+            hidden_states=kept_hidden_states,
+            hidden_size=self.config.hidden_size,
+            labels=labels,
+            shift_labels=shift_labels,
+            **kwargs,
+        )
+
+    else:
+        logits = self.lm_head(kept_hidden_states)
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def lce_maybe_trainable_lm_head(
+    self, hidden_states, hidden_size, labels, shift_labels, **loss_kwargs
+):
+    lm_head = self.lm_head
+
+    # Unwrap the module if lm_head has been added as trainable module in PEFT LoRA configuration,
+    # i.e. listed in the modules_to_save field of LoraConfig, so the lm_head weights are read
+    # from the unwrapped module.
+    # See https://huggingface.co/docs/peft/package_reference/lora for reference.
+    if PEFT_AVAILABLE and isinstance(lm_head, ModulesToSaveWrapper):
+        lm_head = lm_head.modules_to_save.default
+
+    # If FSDP is used and lm_head is trainable, e.g., during full fine-tuning or with LoRA,
+    # reading the lm_head module weights and calling the kernel must be done within FSDP forward pass
+    # so the module entire parameters are summoned and kept in memory during the kernel execution.
+    if isinstance(lm_head, FullyShardedDataParallel):
+        return _FSDPForwardRedirection()(
+            lm_head,
+            _liger_for_causal_lm_loss,
+            lm_head.module,
+            hidden_states,
+            hidden_size,
+            labels,
+            shift_labels,
+            **loss_kwargs,
+        )
+
+    # FSDP is not used so we can read the lm_head weights and call the kernel directly
+    return _liger_for_causal_lm_loss(
+        lm_head=self.lm_head,
+        hidden_states=hidden_states,
+        hidden_size=hidden_size,
+        labels=labels,
+        shift_labels=shift_labels,
+        **loss_kwargs,
+    )
+
+
+def _liger_for_causal_lm_loss(
+    lm_head, hidden_states, hidden_size, labels, shift_labels, **loss_kwargs
+):
+    return LigerForCausalLMLoss(
+        hidden_states=hidden_states,
+        lm_head_weight=lm_head.weight,
+        labels=labels,
+        hidden_size=hidden_size,
+        shift_labels=shift_labels,
+        **loss_kwargs,
+    )
+
+
+def patch_lce_forward(
+    model_type,
+):
+    try:
+        # Dynamically import the module and MLP class
+        module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+        model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+        module = __import__(module_path, fromlist=[f"{model_cls_prefix}ForCausalLM"])
+        model_cls = getattr(module, f"{model_cls_prefix}ForCausalLM")
+
+        model_cls.forward = lce_forward
+
+    except (ImportError, AttributeError) as e:
+        raise RuntimeError(
+            f"Could not import ForCausalLM class for model_type: {model_type}. "
+            f"Error: {str(e)}"
+        ) from e
diff --git a/src/axolotl/integrations/liger/models/deepseekv2.py b/src/axolotl/integrations/liger/models/deepseekv2.py
new file mode 100644
index 0000000000..99adce4a79
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/deepseekv2.py
@@ -0,0 +1,122 @@
+"""
+DeepseekV2 model with LigerFusedLinearCrossEntropyLoss
+"""
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+from liger_kernel.transformers.fused_linear_cross_entropy import (
+    LigerFusedLinearCrossEntropyLoss,
+)
+from torch.nn import CrossEntropyLoss
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def lce_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[List[torch.FloatTensor]] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+) -> Union[Tuple, CausalLMOutputWithPast]:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers.,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`.
+
+    Returns:
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoTokenizer, DeepseekV2ForCausalLM
+
+    >>> model = DeepseekV2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+    >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+    >>> prompt = "Hey, are you conscious? Can you talk to me?"
+    >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+    ```"""
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+    return_dict = (
+        return_dict if return_dict is not None else self.config.use_return_dict
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        return_dict=return_dict,
+    )
+
+    hidden_states = outputs[0]
+
+    loss = None
+    logits = None
+
+    if self.training:
+        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+
+        # flatten tokens
+        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
+        shift_labels = shift_labels.view(-1)
+
+        lce = LigerFusedLinearCrossEntropyLoss()
+        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
+    else:
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
diff --git a/src/axolotl/integrations/liger/models/jamba.py b/src/axolotl/integrations/liger/models/jamba.py
new file mode 100644
index 0000000000..78689e40c6
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/jamba.py
@@ -0,0 +1,162 @@
+"""
+Jamba model with LigerFusedLinearCrossEntropyLoss
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+from liger_kernel.transformers.fused_linear_cross_entropy import (
+    LigerFusedLinearCrossEntropyLoss,
+)
+from torch.nn import CrossEntropyLoss
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+from transformers.models.jamba.modeling_jamba import (
+    HybridMambaAttentionDynamicCache,
+    load_balancing_loss_func,
+)
+
+
+def lce_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    output_router_logits: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    num_logits_to_keep: Optional[Union[int, None]] = None,
+) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        num_logits_to_keep (`int` or `None`, *optional*):
+            Calculate logits for the last `num_logits_to_keep` tokens. If `None`, calculate logits for all
+            `input_ids`. Only last token logits are needed for generation, and calculating them only for that token
+            can save memory, which becomes pretty significant for long sequences.
+
+    Returns:
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoTokenizer, JambaForCausalLM
+
+    >>> model = JambaForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
+    >>> tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
+
+    >>> prompt = "Hey, are you conscious? Can you talk to me?"
+    >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+    ```"""
+
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_router_logits = (
+        output_router_logits
+        if output_router_logits is not None
+        else self.config.output_router_logits
+    )
+
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+    return_dict = (
+        return_dict if return_dict is not None else self.config.use_return_dict
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        output_router_logits=output_router_logits,
+        cache_position=cache_position,
+        return_dict=return_dict,
+    )
+
+    hidden_states = outputs[0]
+
+    loss = None
+    logits = None
+
+    if self.training:
+        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+
+        # flatten tokens
+        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
+        shift_labels = shift_labels.view(-1)
+
+        lce = LigerFusedLinearCrossEntropyLoss()
+        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
+    else:
+        if num_logits_to_keep is None:
+            logits = self.lm_head(hidden_states)
+        else:
+            logits = self.lm_head(hidden_states[..., -num_logits_to_keep:, :])
+        logits = logits.float()
+
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+    aux_loss = None
+    if output_router_logits:
+        aux_loss = load_balancing_loss_func(
+            outputs.router_logits if return_dict else outputs[-1],
+            self.num_experts,
+            self.num_experts_per_tok,
+            attention_mask,
+        )
+        if labels is not None:
+            loss += self.router_aux_loss_coef * aux_loss.to(
+                loss.device
+            )  # make sure to reside in the same device
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        if output_router_logits:
+            output = (aux_loss,) + output
+        return (loss,) + output if loss is not None else output
+
+    return MoeCausalLMOutputWithPast(
+        loss=loss,
+        aux_loss=aux_loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        router_logits=outputs.router_logits,
+    )
diff --git a/src/axolotl/integrations/liger/models/llama4.py b/src/axolotl/integrations/liger/models/llama4.py
new file mode 100644
index 0000000000..e51140265b
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/llama4.py
@@ -0,0 +1,178 @@
+"""
+Liger FLCE for llama4
+"""
+
+import sys
+from copy import deepcopy
+from typing import List, Optional, Tuple, Union
+
+import torch
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def lce_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[
+        Union["Cache", List[torch.FloatTensor]]  # noqa: F821
+    ] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **loss_kwargs,
+) -> Union[Tuple, CausalLMOutputWithPast]:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+    Returns:
+    """
+
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+    return_dict = (
+        return_dict if return_dict is not None else self.config.use_return_dict
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        return_dict=return_dict,
+        cache_position=cache_position,
+    )
+
+    hidden_states = outputs[0]
+
+    if hasattr(self.config, "pretraining_tp") and self.config.pretraining_tp > 1:
+        raise Exception("Liger Kernel does not support pretraining_tp!!")
+
+    logits = None
+    loss = None
+    # if in training mode, don't materialize logits
+    if self.training and (labels is not None):
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
+
+    else:  # if in inference mode materialize logits
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **loss_kwargs,
+            )
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def apply_liger_kernel_to_llama4(
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = False,
+    rms_norm: bool = False,
+    glu_activation: bool = False,
+    layer_norm: bool = False,
+    **kwargs,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Llama models (2 and 3)
+
+    Args:
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is False.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be False.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is False.
+        glu_activation (bool): Whether to apply Liger's SwiGLU MLP. Default is False.
+        layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False.
+    """
+
+    import transformers.models.llama4.modeling_llama4  # noqa: F401
+    from liger_kernel.transformers.functional import liger_cross_entropy
+    from liger_kernel.transformers.layer_norm import LigerLayerNorm
+    from liger_kernel.transformers.rms_norm import LigerRMSNorm
+    from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
+
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    modeling_llama4 = sys.modules["transformers.models.llama4.modeling_llama4"]
+
+    if rms_norm:
+        modeling_llama4.Llama4TextRMSNorm = LigerRMSNorm
+    if glu_activation:
+
+        def _liger_swiglu_mlp_wrapper(config, intermediate_size=None, **kwargs):
+            "Accepts intermediate_size to pass to LigerSwiGLUMLP"
+            # clone config to avoid modifying the original
+            config = deepcopy(config)
+            if intermediate_size:
+                config.intermediate_size = intermediate_size
+            return LigerSwiGLUMLP(config, **kwargs)
+
+        modeling_llama4.Llama4TextMLP = _liger_swiglu_mlp_wrapper
+    if layer_norm:
+        modeling_llama4.nn.LayerNorm = LigerLayerNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        modeling_llama4.Llama4ForCausalLM.forward = lce_forward
diff --git a/src/axolotl/integrations/liger/models/qwen3.py b/src/axolotl/integrations/liger/models/qwen3.py
new file mode 100644
index 0000000000..b008755da3
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/qwen3.py
@@ -0,0 +1,158 @@
+"""
+Liger FLCE for Qwen3. Based on transformers v4.51.3.
+"""
+
+import sys
+from typing import Optional, Tuple, Union
+
+import torch
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def lce_forward(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Cache] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **kwargs,
+) -> Union[Tuple, CausalLMOutputWithPast]:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+    Returns:
+    """
+
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+
+    logits = None
+    loss = None
+    # if in training mode, don't materialize logits
+    if self.training and (labels is not None):
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **kwargs,
+        )
+
+    else:  # if in inference mode materialize logits
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def apply_liger_kernel_to_qwen3(
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = False,
+    rms_norm: bool = False,
+    glu_activation: bool = False,
+    layer_norm: bool = False,
+    **kwargs,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Llama models (2 and 3)
+
+    Args:
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is False.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be False.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is False.
+        glu_activation (bool): Whether to apply Liger's SwiGLU MLP. Default is False.
+        layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False.
+    """
+
+    import transformers.models.qwen3.modeling_qwen3  # noqa: F401
+    from liger_kernel.transformers.functional import liger_cross_entropy
+    from liger_kernel.transformers.layer_norm import LigerLayerNorm
+    from liger_kernel.transformers.rms_norm import LigerRMSNorm
+    from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
+
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    modeling_qwen3 = sys.modules["transformers.models.qwen3.modeling_qwen3"]
+
+    if rms_norm:
+        modeling_qwen3.Qwen3RMSNorm = LigerRMSNorm
+
+    if glu_activation:
+        modeling_qwen3.Qwen3MLP = LigerSwiGLUMLP
+
+    if layer_norm:
+        modeling_qwen3.nn.LayerNorm = LigerLayerNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        modeling_qwen3.Qwen3ForCausalLM.forward = lce_forward
diff --git a/src/axolotl/integrations/liger/models/qwen3_5.py b/src/axolotl/integrations/liger/models/qwen3_5.py
new file mode 100644
index 0000000000..ee4b9b1c1c
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/qwen3_5.py
@@ -0,0 +1,175 @@
+"""
+Liger FLCE for Qwen3.5. Based on transformers v5.3.0.
+"""
+
+import sys
+from copy import deepcopy
+from typing import Optional, Union
+
+import torch
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def lce_forward(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Cache] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **kwargs,
+) -> CausalLMOutputWithPast:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+    Returns:
+    """
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+
+    logits = None
+    loss = None
+    # if in training mode, don't materialize logits
+    if self.training and (labels is not None):
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **kwargs,
+        )
+
+    else:  # if in inference mode materialize logits
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def apply_liger_kernel_to_qwen3_5(
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = False,
+    rms_norm: bool = False,
+    rms_norm_gated: bool = False,
+    glu_activation: bool = False,
+    layer_norm: bool = False,
+    **kwargs,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Qwen3.5 models.
+
+    Note: Qwen3_5RMSNorm uses zero-init weight with offset 1.0 (like Gemma),
+    so we use LigerRMSNorm with offset=1.0 and init_fn="zeros".
+
+    Args:
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is False.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is False.
+        rms_norm_gated (bool): Whether to apply fused RMSNorm+SiLU gate kernel for
+            Qwen3_5RMSNormGated (used in linear attention layers). Default is False.
+        glu_activation (bool): Whether to apply Liger's SwiGLU MLP. Default is False.
+        layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False.
+    """
+
+    import transformers.models.qwen3_5.modeling_qwen3_5  # noqa: F401
+    from liger_kernel.transformers.functional import liger_cross_entropy
+    from liger_kernel.transformers.layer_norm import LigerLayerNorm
+    from liger_kernel.transformers.rms_norm import LigerRMSNorm
+    from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
+
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    modeling_qwen3_5 = sys.modules["transformers.models.qwen3_5.modeling_qwen3_5"]
+
+    if rms_norm:
+        # Qwen3_5RMSNorm uses zero-init weight with `output * (1.0 + weight)` pattern
+        class LigerRMSNormForQwen3_5(LigerRMSNorm):
+            def __init__(self, dim, eps=1e-6, **kwargs):
+                super().__init__(
+                    dim,
+                    eps=eps,
+                    offset=1.0,
+                    casting_mode="gemma",
+                    init_fn="zeros",
+                    in_place=False,
+                )
+
+        modeling_qwen3_5.Qwen3_5RMSNorm = LigerRMSNormForQwen3_5
+
+    if rms_norm_gated:
+        from axolotl.kernels.rms_norm_gated import FusedRMSNormGated
+
+        modeling_qwen3_5.Qwen3_5RMSNormGated = FusedRMSNormGated
+
+    if glu_activation:
+
+        def _liger_swiglu_mlp_wrapper(config, intermediate_size=None, **kwargs):
+            """Accepts intermediate_size to pass to LigerSwiGLUMLP"""
+            config = deepcopy(config)
+            if intermediate_size is not None:
+                config.intermediate_size = intermediate_size
+            return LigerSwiGLUMLP(config, **kwargs)
+
+        modeling_qwen3_5.Qwen3_5MLP = _liger_swiglu_mlp_wrapper
+
+    if layer_norm:
+        modeling_qwen3_5.nn.LayerNorm = LigerLayerNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        modeling_qwen3_5.Qwen3_5ForCausalLM.forward = lce_forward
diff --git a/src/axolotl/integrations/liger/models/qwen3_5_moe.py b/src/axolotl/integrations/liger/models/qwen3_5_moe.py
new file mode 100644
index 0000000000..b10b34ad51
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/qwen3_5_moe.py
@@ -0,0 +1,198 @@
+"""
+Liger FLCE for Qwen3.5 MoE. Based on transformers v5.3.0.
+"""
+
+import sys
+from copy import deepcopy
+from typing import Optional, Union
+
+import torch
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+
+
+def lce_forward(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values=None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_router_logits: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **kwargs,
+) -> MoeCausalLMOutputWithPast:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+    Returns:
+    """
+    from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
+        load_balancing_loss_func,
+    )
+
+    output_router_logits = (
+        output_router_logits
+        if output_router_logits is not None
+        else self.config.output_router_logits
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_router_logits=output_router_logits,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+
+    logits = None
+    loss = None
+    # if in training mode, don't materialize logits
+    if self.training and (labels is not None):
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **kwargs,
+        )
+
+    else:  # if in inference mode materialize logits
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        if labels is not None:
+            loss = self.loss_function(
+                logits,
+                labels,
+                self.vocab_size,
+                **kwargs,
+            )
+
+    aux_loss = None
+    if output_router_logits:
+        aux_loss = load_balancing_loss_func(
+            outputs.router_logits,
+            self.num_experts,
+            self.num_experts_per_tok,
+            attention_mask,
+        )
+        if labels is not None:
+            loss += self.router_aux_loss_coef * aux_loss.to(loss.device)
+
+    return MoeCausalLMOutputWithPast(
+        loss=loss,
+        aux_loss=aux_loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        router_logits=outputs.router_logits,
+    )
+
+
+def apply_liger_kernel_to_qwen3_5_moe(
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = False,
+    rms_norm: bool = False,
+    rms_norm_gated: bool = False,
+    glu_activation: bool = False,
+    layer_norm: bool = False,
+    **kwargs,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Qwen3.5 MoE models.
+
+    Note: Qwen3_5MoeRMSNorm uses zero-init weight with offset 1.0 (like Gemma),
+    so we use LigerRMSNorm with offset=1.0 and init_fn="zeros".
+
+    Args:
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is False.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is False.
+        rms_norm_gated (bool): Whether to apply fused RMSNorm+SiLU gate kernel for
+            Qwen3_5MoeRMSNormGated (used in linear attention layers). Default is False.
+        glu_activation (bool): Whether to apply Liger's SwiGLU MLP. Default is False.
+        layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False.
+    """
+
+    import transformers.models.qwen3_5_moe.modeling_qwen3_5_moe  # noqa: F401
+    from liger_kernel.transformers.functional import liger_cross_entropy
+    from liger_kernel.transformers.layer_norm import LigerLayerNorm
+    from liger_kernel.transformers.rms_norm import LigerRMSNorm
+    from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
+
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    modeling_mod = sys.modules["transformers.models.qwen3_5_moe.modeling_qwen3_5_moe"]
+
+    if rms_norm:
+        # Qwen3_5MoeRMSNorm uses zero-init weight with `output * (1.0 + weight)` pattern
+        class LigerRMSNormForQwen3_5Moe(LigerRMSNorm):
+            def __init__(self, dim, eps=1e-6, **kwargs):
+                super().__init__(
+                    dim,
+                    eps=eps,
+                    offset=1.0,
+                    casting_mode="gemma",
+                    init_fn="zeros",
+                    in_place=False,
+                )
+
+        modeling_mod.Qwen3_5MoeRMSNorm = LigerRMSNormForQwen3_5Moe
+
+    if rms_norm_gated:
+        from axolotl.kernels.rms_norm_gated import FusedRMSNormGated
+
+        modeling_mod.Qwen3_5MoeRMSNormGated = FusedRMSNormGated
+
+    if glu_activation:
+
+        def _liger_swiglu_mlp_wrapper(config, intermediate_size=None, **kwargs):
+            """Accepts intermediate_size to pass to LigerSwiGLUMLP"""
+            config = deepcopy(config)
+            if intermediate_size is not None:
+                config.intermediate_size = intermediate_size
+            return LigerSwiGLUMLP(config, **kwargs)
+
+        modeling_mod.Qwen3_5MoeMLP = _liger_swiglu_mlp_wrapper
+
+    if layer_norm:
+        modeling_mod.nn.LayerNorm = LigerLayerNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        modeling_mod.Qwen3_5MoeForCausalLM.forward = lce_forward
diff --git a/src/axolotl/integrations/liger/models/qwen3_moe.py b/src/axolotl/integrations/liger/models/qwen3_moe.py
new file mode 100644
index 0000000000..40bee110c4
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/qwen3_moe.py
@@ -0,0 +1,189 @@
+"""
+Liger FLCE for Qwen3 MoE. Based on transformers v4.51.3.
+"""
+
+import sys
+from copy import deepcopy
+from typing import List, Optional, Union
+
+import torch
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+from transformers.models.qwen3_moe.modeling_qwen3_moe import load_balancing_loss_func
+
+
+def lce_forward(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[List[torch.FloatTensor]] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    output_router_logits: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **kwargs,
+) -> MoeCausalLMOutputWithPast:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+    Returns:
+    """
+
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_router_logits = (
+        output_router_logits
+        if output_router_logits is not None
+        else self.config.output_router_logits
+    )
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        output_router_logits=output_router_logits,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+
+    logits = None
+    loss = None
+    # if in training mode, don't materialize logits
+    if self.training and (labels is not None):
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **kwargs,
+        )
+
+    else:  # if in inference mode materialize logits
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+    aux_loss = None
+    if output_router_logits:
+        aux_loss = load_balancing_loss_func(
+            outputs.router_logits,
+            self.num_experts,
+            self.num_experts_per_tok,
+            attention_mask,
+        )
+        if labels is not None:
+            loss += self.router_aux_loss_coef * aux_loss.to(
+                loss.device
+            )  # make sure to reside in the same device
+
+    return MoeCausalLMOutputWithPast(
+        loss=loss,
+        aux_loss=aux_loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def apply_liger_kernel_to_qwen3_moe(
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = False,
+    rms_norm: bool = False,
+    glu_activation: bool = False,
+    layer_norm: bool = False,
+    **kwargs,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Llama models (2 and 3)
+
+    Args:
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is False.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be False.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is False.
+        glu_activation (bool): Whether to apply Liger's SwiGLU MLP. Default is False.
+        layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False.
+    """
+
+    import transformers.models.qwen3_moe.modeling_qwen3_moe  # noqa: F401
+    from liger_kernel.transformers.functional import liger_cross_entropy
+    from liger_kernel.transformers.layer_norm import LigerLayerNorm
+    from liger_kernel.transformers.rms_norm import LigerRMSNorm
+    from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
+
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    modeling_qwen3_moe = sys.modules["transformers.models.qwen3_moe.modeling_qwen3_moe"]
+
+    if rms_norm:
+        modeling_qwen3_moe.Qwen3MoeRMSNorm = LigerRMSNorm
+
+    if glu_activation:
+
+        def _liger_swiglu_mlp_wrapper(config, intermediate_size=None, **kwargs):
+            "Accepts intermediate_size to pass to LigerSwiGLUMLP"
+            # clone config to avoid modifying the original
+            config = deepcopy(config)
+            if intermediate_size:
+                config.intermediate_size = intermediate_size
+            return LigerSwiGLUMLP(config, **kwargs)
+
+        modeling_qwen3_moe.Qwen3MoeMLP = _liger_swiglu_mlp_wrapper
+
+    if layer_norm:
+        modeling_qwen3_moe.nn.LayerNorm = LigerLayerNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        modeling_qwen3_moe.Qwen3MoeForCausalLM.forward = lce_forward
diff --git a/src/axolotl/integrations/liger/plugin.py b/src/axolotl/integrations/liger/plugin.py
new file mode 100644
index 0000000000..09aec5342c
--- /dev/null
+++ b/src/axolotl/integrations/liger/plugin.py
@@ -0,0 +1,367 @@
+"""
+Liger-Kernel Plugin for Axolotl
+"""
+
+import inspect
+import sys
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class LigerPlugin(BasePlugin):
+    """
+    Plugin for LIGER integraton with Axolotl.
+    """
+
+    def get_input_args(self):
+        return "axolotl.integrations.liger.LigerArgs"
+
+    def pre_model_load(self, cfg):
+        # shim: liger imports ORPOTrainer from old trl.trainer (now trl.experimental.orpo)
+        import trl.trainer
+        from trl.experimental.orpo import ORPOTrainer
+
+        trl.trainer.ORPOTrainer = ORPOTrainer
+
+        if cfg.torch_compile:
+            # torch compile will unnecessarily attempt to optimize the triton kernel unless explicitly disabled
+            import liger_kernel.ops.fused_linear_cross_entropy
+
+            from .utils import patch_with_compile_disable
+
+            patch_with_compile_disable(
+                liger_kernel.ops.fused_linear_cross_entropy,
+                "fused_linear_cross_entropy_forward",
+            )
+            patch_with_compile_disable(
+                liger_kernel.ops.fused_linear_cross_entropy,
+                "fused_linear_cross_entropy_backward",
+            )
+
+        from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss
+        from liger_kernel.transformers.functional import liger_cross_entropy
+        from liger_kernel.transformers.layer_norm import LigerLayerNorm
+        from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN
+        from liger_kernel.transformers.rms_norm import LigerRMSNorm
+        from liger_kernel.transformers.rope import liger_rotary_pos_emb
+        from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
+
+        if cfg.liger_cross_entropy and cfg.liger_fused_linear_cross_entropy:
+            raise ValueError(
+                "Cannot have both `liger_cross_entropy` and `liger_fused_linear_cross_entropy` set."
+            )
+
+        if cfg.liger_use_token_scaling:
+            # Patch FLCE to set token_scaling=True for function and class API
+            from liger_kernel.transformers import functional
+            from liger_kernel.transformers.fused_linear_cross_entropy import (
+                LigerFusedLinearCrossEntropyLoss,
+            )
+
+            old_liger_fused_linear_cross_entropy = (
+                functional.liger_fused_linear_cross_entropy
+            )
+
+            def patched_liger_fused_linear_cross_entropy(*args, **kwargs):
+                kwargs["use_token_scaling"] = True
+                return old_liger_fused_linear_cross_entropy(*args, **kwargs)
+
+            functional.liger_fused_linear_cross_entropy = (
+                patched_liger_fused_linear_cross_entropy
+            )
+
+            old_init = LigerFusedLinearCrossEntropyLoss.__init__
+
+            def patched_init(self, *args, **kwargs):
+                kwargs["use_token_scaling"] = True
+                return old_init(self, *args, **kwargs)
+
+            LigerFusedLinearCrossEntropyLoss.__init__ = patched_init
+
+        # liger 0.8.0 natively dispatches these, but axolotl's branches add kernels native lacks
+        axolotl_override_liger_fn = {"qwen3_5", "qwen3_5_moe"}
+
+        if (
+            cfg.model_config_type in MODEL_TYPE_TO_APPLY_LIGER_FN
+            and cfg.model_config_type not in axolotl_override_liger_fn
+        ):
+            apply_liger_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[cfg.model_config_type]
+            liger_fn_sig = inspect.signature(apply_liger_fn)
+            kwargs = {}
+            if "rope" in liger_fn_sig.parameters:
+                rope_value = cfg.liger_rope
+                # cfg.liger_rope defaults to None, which would override upstream's rope=True for Qwen-VL.
+                if rope_value is None and cfg.model_config_type in (
+                    "qwen2_vl",
+                    "qwen2_5_vl",
+                    "qwen3_vl",
+                    "qwen3_vl_moe",
+                    "qwen2_vl_text",
+                    "qwen2_5_vl_text",
+                    "qwen3_vl_text",
+                    "qwen3_vl_moe_text",
+                ):
+                    rope_value = True
+                kwargs["rope"] = rope_value
+            if "cross_entropy" in liger_fn_sig.parameters:
+                kwargs["cross_entropy"] = cfg.liger_cross_entropy
+            if "fused_linear_cross_entropy" in liger_fn_sig.parameters:
+                kwargs["fused_linear_cross_entropy"] = (
+                    cfg.liger_fused_linear_cross_entropy
+                )
+            if "rms_norm" in liger_fn_sig.parameters:
+                kwargs["rms_norm"] = cfg.liger_rms_norm
+            if "layer_norm" in liger_fn_sig.parameters:
+                kwargs["layer_norm"] = cfg.liger_layer_norm
+            if "geglu" in liger_fn_sig.parameters:
+                kwargs["geglu"] = cfg.liger_glu_activation
+            elif "swiglu" in liger_fn_sig.parameters:
+                kwargs["swiglu"] = cfg.liger_glu_activation
+            LOG.info(f"Applying LIGER to {cfg.model_config_type} with kwargs: {kwargs}")
+            apply_liger_fn(**kwargs)
+        elif cfg.model_config_type == "jamba":
+            from transformers.models.jamba import modeling_jamba
+
+            from .models.jamba import lce_forward as jamba_lce_forward
+
+            if cfg.liger_rope:
+                modeling_jamba.apply_rotary_pos_emb = liger_rotary_pos_emb
+            if cfg.liger_rms_norm:
+                modeling_jamba.JambaRMSNorm = LigerRMSNorm
+            if cfg.liger_glu_activation:
+                modeling_jamba.JambaMLP = LigerSwiGLUMLP
+            if cfg.liger_layer_norm:
+                modeling_jamba.nn.LayerNorm = LigerLayerNorm
+            if cfg.liger_cross_entropy:
+                from transformers.loss.loss_utils import nn
+
+                nn.functional.cross_entropy = liger_cross_entropy
+            if cfg.liger_fused_linear_cross_entropy:
+                modeling_jamba.JambaForCausalLM.forward = jamba_lce_forward
+        elif cfg.model_config_type == "deepseek_v2":
+            from accelerate import init_empty_weights
+            from transformers import AutoModelForCausalLM
+
+            with init_empty_weights():
+                model = AutoModelForCausalLM.from_pretrained(
+                    cfg.base_model, trust_remote_code=cfg.trust_remote_code or False
+                )
+                modeling_mod = sys.modules[model.__class__.__module__]
+
+            from .models.deepseekv2 import lce_forward as deepseekv2_lce_forward
+
+            if cfg.liger_rope:
+                # The DeepseekV2 version of RoPE is different than upstream LLaMA.
+                # See https://github.com/linkedin/Liger-Kernel/issues/129#issuecomment-2313763528
+                LOG.warning("Fused liger_rope is not supported for DeepseekV2.")
+            if cfg.liger_rms_norm:
+                modeling_mod.DeepseekV2RMSNorm = LigerRMSNorm
+            if cfg.liger_glu_activation:
+                modeling_mod.DeepseekV2MLP.forward = LigerSwiGLUMLP.forward
+            if cfg.liger_layer_norm:
+                LOG.warning("liger_layer_norm is not supported for DeepseekV2.")
+            if cfg.liger_cross_entropy:
+                # We do not patch `nn.functional.cross_entropy` for DeepseekV2 as it still uses
+                # nn.CrossEntropyLoss in the forward method.
+                modeling_mod.CrossEntropyLoss = LigerCrossEntropyLoss
+            if cfg.liger_fused_linear_cross_entropy:
+                modeling_mod.DeepseekV2ForCausalLM.forward = deepseekv2_lce_forward
+        elif cfg.model_config_type == "llama4":
+            from axolotl.integrations.liger.models.llama4 import (
+                apply_liger_kernel_to_llama4,
+            )
+
+            apply_liger_kernel_to_llama4(
+                cross_entropy=cfg.liger_cross_entropy,
+                fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+                glu_activation=cfg.liger_glu_activation,
+                rms_norm=cfg.liger_rms_norm,
+                layer_norm=cfg.liger_layer_norm,
+            )
+        elif cfg.model_config_type == "qwen3":
+            from axolotl.integrations.liger.models.qwen3 import (
+                apply_liger_kernel_to_qwen3,
+            )
+
+            apply_liger_kernel_to_qwen3(
+                cross_entropy=cfg.liger_cross_entropy,
+                fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+                glu_activation=cfg.liger_glu_activation,
+                rms_norm=cfg.liger_rms_norm,
+                layer_norm=cfg.liger_layer_norm,
+            )
+        elif cfg.model_config_type == "qwen3_5":
+            from axolotl.integrations.liger.models.qwen3_5 import (
+                apply_liger_kernel_to_qwen3_5,
+            )
+
+            apply_liger_kernel_to_qwen3_5(
+                cross_entropy=cfg.liger_cross_entropy,
+                fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+                glu_activation=cfg.liger_glu_activation,
+                rms_norm=cfg.liger_rms_norm,
+                rms_norm_gated=getattr(cfg, "liger_rms_norm_gated", False),
+                layer_norm=cfg.liger_layer_norm,
+            )
+        elif cfg.model_config_type == "qwen3_moe":
+            from axolotl.integrations.liger.models.qwen3_moe import (
+                apply_liger_kernel_to_qwen3_moe,
+            )
+
+            apply_liger_kernel_to_qwen3_moe(
+                cross_entropy=cfg.liger_cross_entropy,
+                fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+                glu_activation=cfg.liger_glu_activation,
+                rms_norm=cfg.liger_rms_norm,
+                layer_norm=cfg.liger_layer_norm,
+            )
+        elif cfg.model_config_type == "qwen3_5_moe":
+            from axolotl.integrations.liger.models.qwen3_5_moe import (
+                apply_liger_kernel_to_qwen3_5_moe,
+            )
+
+            apply_liger_kernel_to_qwen3_5_moe(
+                cross_entropy=cfg.liger_cross_entropy,
+                fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+                glu_activation=cfg.liger_glu_activation,
+                rms_norm=cfg.liger_rms_norm,
+                rms_norm_gated=getattr(cfg, "liger_rms_norm_gated", False),
+                layer_norm=cfg.liger_layer_norm,
+            )
+        elif cfg.model_config_type == "granitemoe":
+            from liger_kernel.transformers import apply_liger_kernel_to_granite
+
+            apply_liger_kernel_to_granite(
+                rope=cfg.liger_rope,
+                cross_entropy=cfg.liger_cross_entropy,
+                fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+                rms_norm=cfg.liger_rms_norm,
+                swiglu=cfg.liger_glu_activation,
+            )
+        elif cfg.model_config_type == "gemma4":
+            # multimodal gemma4 only; gemma4_text uses liger 0.8.0 native dispatch (incl. FLCE)
+            from liger_kernel.transformers.geglu import LigerGEGLUMLP
+            from transformers.models.gemma4 import modeling_gemma4
+
+            if cfg.liger_rms_norm:
+                _OrigGemma4RMSNorm = modeling_gemma4.Gemma4RMSNorm
+
+                class _LigerGemma4RMSNorm(LigerRMSNorm):
+                    """LigerRMSNorm for Gemma4: offset=0, in_place=False (grad-ckpt safe), with_scale support."""
+
+                    def __new__(cls, dim, eps=1e-6, with_scale=True):
+                        if not with_scale:
+                            return _OrigGemma4RMSNorm(dim, eps, with_scale=False)
+                        return super().__new__(cls)
+
+                    def __init__(self, dim, eps=1e-6, with_scale=True):
+                        if not with_scale:
+                            return
+                        # offset=0.0 (standard), in_place=False (gradient checkpointing safe)
+                        super().__init__(
+                            dim, eps, offset=0.0, casting_mode="llama", in_place=False
+                        )
+
+                modeling_gemma4.Gemma4RMSNorm = _LigerGemma4RMSNorm
+            if cfg.liger_glu_activation:
+
+                class _LigerGemma4MLP(LigerGEGLUMLP):
+                    def __init__(self, config, layer_idx=None):
+                        super().__init__(config)
+
+                modeling_gemma4.Gemma4TextMLP = _LigerGemma4MLP
+            if cfg.liger_rope:
+                LOG.warning(
+                    "Liger RoPE is not compatible with Gemma4 (separate q/k application). Skipping."
+                )
+            if cfg.liger_layer_norm:
+                modeling_gemma4.nn.LayerNorm = LigerLayerNorm
+            if cfg.liger_cross_entropy:
+                modeling_gemma4.nn.CrossEntropyLoss = LigerCrossEntropyLoss
+            if cfg.liger_fused_linear_cross_entropy:
+                LOG.warning(
+                    "Liger fused linear cross entropy for multimodal gemma4 is not in "
+                    "liger-kernel 0.8.0 (added upstream in PR #1203, post-0.8.0). Skipping; "
+                    "the gemma4_text language model gets FLCE via liger's native path."
+                )
+            LOG.info(
+                f"Applied Liger kernels for gemma4: "
+                f"rms_norm={cfg.liger_rms_norm}, glu={cfg.liger_glu_activation}, "
+                f"rope=False (incompatible), layer_norm={cfg.liger_layer_norm}"
+            )
+        elif cfg.model_config_type in ("gemma4_unified", "gemma4_unified_text"):
+            # gemma4_unified mirrors gemma4's Liger compatibility: offset=0,
+            # in_place=False (gradient-checkpoint safe), RoPE incompatible
+            # (separate q/k application). Classes live in the unified namespace.
+            from liger_kernel.transformers.geglu import LigerGEGLUMLP
+            from transformers.models.gemma4_unified import modeling_gemma4_unified
+
+            if cfg.liger_rms_norm:
+                _OrigGemma4UnifiedRMSNorm = modeling_gemma4_unified.Gemma4UnifiedRMSNorm
+
+                class _LigerGemma4UnifiedRMSNorm(LigerRMSNorm):
+                    """LigerRMSNorm for Gemma4Unified (in_place=False, with_scale)."""
+
+                    def __new__(cls, dim, eps=1e-6, with_scale=True):
+                        if not with_scale:
+                            return _OrigGemma4UnifiedRMSNorm(dim, eps, with_scale=False)
+                        return super().__new__(cls)
+
+                    def __init__(self, dim, eps=1e-6, with_scale=True):
+                        if not with_scale:
+                            return
+                        super().__init__(
+                            dim, eps, offset=0.0, casting_mode="llama", in_place=False
+                        )
+
+                modeling_gemma4_unified.Gemma4UnifiedRMSNorm = (
+                    _LigerGemma4UnifiedRMSNorm
+                )
+            if cfg.liger_glu_activation:
+
+                class _LigerGemma4UnifiedMLP(LigerGEGLUMLP):
+                    def __init__(self, config, layer_idx=None):
+                        super().__init__(config)
+
+                modeling_gemma4_unified.Gemma4UnifiedTextMLP = _LigerGemma4UnifiedMLP
+            if cfg.liger_rope:
+                LOG.warning(
+                    "Liger RoPE is not compatible with Gemma4Unified (separate "
+                    "q/k application). Skipping."
+                )
+            if cfg.liger_layer_norm:
+                modeling_gemma4_unified.nn.LayerNorm = LigerLayerNorm
+            if cfg.liger_cross_entropy:
+                modeling_gemma4_unified.nn.CrossEntropyLoss = LigerCrossEntropyLoss
+            if cfg.liger_fused_linear_cross_entropy:
+                LOG.warning(
+                    "Liger fused linear cross entropy is not compatible with "
+                    "Gemma4Unified. Skipping."
+                )
+            LOG.info(
+                f"Applied Liger kernels for gemma4_unified: "
+                f"rms_norm={cfg.liger_rms_norm}, glu={cfg.liger_glu_activation}, "
+                f"rope=False (incompatible), layer_norm={cfg.liger_layer_norm}"
+            )
+        elif cfg.liger_fused_linear_cross_entropy:
+            try:
+                from .models.base import patch_lce_forward
+
+                patch_lce_forward(cfg.model_config_type)
+                LOG.warning_once(
+                    f"Applied ONLY liger_fused_linear_cross_entropy genericpatches for model type: {cfg.model_config_type}"
+                )
+                LOG.warning_once(
+                    f"Liger + {cfg.model_config_type} generic FLCE support is experimental and may not work as expected."
+                )
+            except RuntimeError:
+                LOG.warning(
+                    f"Unsupported model config type: {cfg.model_config_type}. Liger not applied."
+                )
+        else:
+            LOG.warning(
+                f"Unsupported model config type: {cfg.model_config_type}. Liger not applied."
+            )
diff --git a/src/axolotl/integrations/liger/utils.py b/src/axolotl/integrations/liger/utils.py
new file mode 100644
index 0000000000..bf9fc58e77
--- /dev/null
+++ b/src/axolotl/integrations/liger/utils.py
@@ -0,0 +1,29 @@
+"""
+utils to patch liger kernel ops to disable torch.compile
+"""
+
+from functools import wraps
+
+import torch
+
+
+def patch_with_compile_disable(module, function_name):
+    """
+    Patch a function in a module by wrapping it with torch.compile.disable
+
+    Args:
+        module: The module containing the function to patch
+        function_name: The name of the function to patch
+    """
+    original_function = getattr(module, function_name)
+
+    @wraps(original_function)
+    @torch.compiler.disable
+    def wrapped_function(*args, **kwargs):
+        return original_function(*args, **kwargs)
+
+    # Replace the original function with the wrapped one
+    setattr(module, function_name, wrapped_function)
+
+    # Return the original function in case you need to restore it later
+    return original_function
diff --git a/src/axolotl/integrations/llm_compressor/README.md b/src/axolotl/integrations/llm_compressor/README.md
new file mode 100644
index 0000000000..16eff804dd
--- /dev/null
+++ b/src/axolotl/integrations/llm_compressor/README.md
@@ -0,0 +1,108 @@
+# LLMCompressor Integration
+
+Fine-tune sparsified models in Axolotl using Neural Magic's [LLMCompressor](https://github.com/vllm-project/llm-compressor).
+
+This integration enables fine-tuning of models sparsified using LLMCompressor within the Axolotl training framework. By combining LLMCompressor's model compression capabilities with Axolotl's distributed training pipelines, users can efficiently fine-tune sparse models at scale.
+
+It uses Axolotl’s plugin system to hook into the fine-tuning flows while maintaining sparsity throughout training.
+
+---
+
+## Requirements
+
+- Axolotl with `llmcompressor` extras:
+
+  ```bash
+  pip install "axolotl[llmcompressor]"
+  ```
+
+- Requires `llmcompressor >= 0.5.1`
+
+This will install all necessary dependencies to fine-tune sparsified models using the integration.
+
+---
+
+## Usage
+
+To enable sparse fine-tuning with this integration, include the plugin in your Axolotl config:
+
+```yaml
+plugins:
+  - axolotl.integrations.llm_compressor.LLMCompressorPlugin
+
+llmcompressor:
+  recipe:
+    finetuning_stage:
+      finetuning_modifiers:
+        ConstantPruningModifier:
+          targets: [
+            're:.*q_proj.weight',
+            're:.*k_proj.weight',
+            're:.*v_proj.weight',
+            're:.*o_proj.weight',
+            're:.*gate_proj.weight',
+            're:.*up_proj.weight',
+            're:.*down_proj.weight',
+          ]
+          start: 0
+  save_compressed: true
+# ... (other training arguments)
+```
+
+This plugin **does not apply pruning or sparsification itself** — it is intended for **fine-tuning models that have already been sparsified**.
+
+Pre-sparsified checkpoints can be:
+- Generated using [LLMCompressor](https://github.com/vllm-project/llm-compressor)
+- Downloaded from [Neural Magic's Hugging Face page](https://huggingface.co/neuralmagic)
+- Any custom LLM with compatible sparsity patterns that you've created yourself
+
+To learn more about writing and customizing LLMCompressor recipes, refer to the official documentation:
+[https://github.com/vllm-project/llm-compressor/blob/main/README.md](https://github.com/vllm-project/llm-compressor/blob/main/README.md)
+
+### Storage Optimization with save_compressed
+
+Setting `save_compressed: true` in your configuration enables saving models in a compressed format, which:
+- Reduces disk space usage by approximately 40%
+- Maintains compatibility with vLLM for accelerated inference
+- Maintains compatibility with llmcompressor for further optimization (example: quantization)
+
+This option is highly recommended when working with sparse models to maximize the benefits of model compression.
+
+### Example Config
+
+See [`examples/llama-3/sparse-finetuning.yaml`](examples/llama-3/sparse-finetuning.yaml) for a complete example.
+
+---
+
+## Inference with vLLM
+
+After fine-tuning your sparse model, you can leverage vLLM for efficient inference.
+You can also use LLMCompressor to apply additional quantization to your fine-tuned
+sparse model before inference for even greater performance benefits.:
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+llm = LLM("path/to/your/sparse/model")
+outputs = llm.generate(prompts, sampling_params)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
+
+For more details on vLLM's capabilities and advanced configuration options, see the [official vLLM documentation](https://docs.vllm.ai/).
+
+## Learn More
+
+For details on available sparsity and quantization schemes, fine-tuning recipes, and usage examples, visit the official LLMCompressor repository:
+
+[https://github.com/vllm-project/llm-compressor](https://github.com/vllm-project/llm-compressor)
diff --git a/src/axolotl/integrations/llm_compressor/__init__.py b/src/axolotl/integrations/llm_compressor/__init__.py
new file mode 100644
index 0000000000..fe799d3c0f
--- /dev/null
+++ b/src/axolotl/integrations/llm_compressor/__init__.py
@@ -0,0 +1,5 @@
+"""Integration entry point for the LLMCompressor plugin."""
+
+from .plugin import LLMCompressorPlugin
+
+__all__ = ["LLMCompressorPlugin"]
diff --git a/src/axolotl/integrations/llm_compressor/args.py b/src/axolotl/integrations/llm_compressor/args.py
new file mode 100644
index 0000000000..4c0e4cac31
--- /dev/null
+++ b/src/axolotl/integrations/llm_compressor/args.py
@@ -0,0 +1,40 @@
+"""
+LLMCompressor and Sparse Finetuning config models.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+
+class CompressionArgs(BaseModel):
+    """Sparse Finetuning config for LLMCompressor."""
+
+    # Typing for recipe is set to Any due to:
+    # https://github.com/vllm-project/llm-compressor/issues/1319
+    recipe: Annotated[
+        Any,
+        Field(
+            description="The recipe containing the compression algorithms and hyperparameters to apply."
+        ),
+    ]
+
+    save_compressed: Annotated[
+        bool,
+        Field(
+            default=False,
+            description="Whether to save the compressed model after training.",
+        ),
+    ]
+
+
+class LLMCompressorArgs(BaseModel):
+    """LLMCompressor configuration BaseModel."""
+
+    llmcompressor: Annotated[
+        CompressionArgs,
+        Field(
+            description="Arguments enabling compression pathways through the LLM Compressor plugins"
+        ),
+    ]
diff --git a/src/axolotl/integrations/llm_compressor/plugin.py b/src/axolotl/integrations/llm_compressor/plugin.py
new file mode 100644
index 0000000000..57d506a573
--- /dev/null
+++ b/src/axolotl/integrations/llm_compressor/plugin.py
@@ -0,0 +1,171 @@
+"""
+Sparse Finetuning plugin for Axolotl — enables handling of sparse neural networks
+by maintaining masks for zero weights during training.
+"""
+
+from functools import wraps
+from typing import Any, Callable, Concatenate, ParamSpec, TypeVar
+
+from llmcompressor import active_session, create_session
+from llmcompressor.core import callbacks as session_callbacks
+from llmcompressor.recipe import Recipe
+from torch.nn import Module
+from transformers.trainer import Trainer
+from transformers.trainer_callback import TrainerCallback, TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+P = ParamSpec("P")  # Params for generic function signatures
+R = TypeVar("R")  # Return type for generic function signatures
+
+LOG = get_logger(__name__)
+
+
+class LLMCompressorCallbackHandler(TrainerCallback):
+    """
+    Trainer callback for Sparse Finetuning.
+    Maintains sparsity patterns during training by applying masks after optimization steps,
+    ensuring zero-weight updates are canceled out.
+    """
+
+    def __init__(self, trainer: Trainer, recipe: Any):
+        """
+        Initialize the Sparse Finetuning callback handler.
+
+        Args:
+            trainer (Trainer): Huggingface Trainer instance.
+            recipe (Recipe | dict): Sparse finetuning recipe to apply.
+        """
+        super().__init__()
+        self.trainer = trainer
+        self.recipe = (
+            Recipe.model_validate(recipe) if not isinstance(recipe, Recipe) else recipe
+        )
+        self.original_compute_loss = trainer.compute_loss
+        self.trainer.compute_loss = compute_loss_wrapper(self.trainer.compute_loss)
+        create_session()
+
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ) -> None:
+        """
+        Called at the beginning of training. Initializes the compression session.
+
+        Args:
+            args (TrainingArguments): Training arguments.
+            state (TrainerState): Trainer state.
+            control (TrainerControl): Trainer control.
+        """
+        super().on_train_begin(args, state, control, **kwargs)
+        self.trainer.accelerator.wait_for_everyone()
+        active_session().initialize(
+            model=self.trainer.model,
+            optimizer=self.trainer.optimizer,
+            start=state.epoch,
+            recipe=self.recipe,
+        )
+        self.trainer.accelerator.wait_for_everyone()
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ) -> None:
+        """
+        Called at the beginning of a training step. Triggers batch_start callback.
+        """
+        super().on_step_begin(args, state, control, **kwargs)
+        session_callbacks.batch_start()
+
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ) -> None:
+        """
+        Called at the end of a training step. Triggers optimizer and batch_end callbacks.
+        """
+        super().on_step_end(args, state, control, **kwargs)
+        session_callbacks.optim_pre_step()
+        session_callbacks.optim_post_step()
+        session_callbacks.batch_end()
+
+    def on_train_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ) -> None:
+        """
+        Called at the end of training. Finalizes the compression session.
+        """
+        super().on_train_end(args, state, control, **kwargs)
+        active_session().finalize()
+        self.trainer.compute_loss_func = self.original_compute_loss
+
+
+class LLMCompressorPlugin(BasePlugin):
+    """
+    Sparse Finetuning plugin for Axolotl integration.
+    """
+
+    def get_input_args(self) -> str:
+        """
+        Returns the path to the plugin's argument definition.
+
+        Returns:
+            str: Dotted path to the LLMCompressorArgs class.
+        """
+        return "axolotl.integrations.llm_compressor.args.LLMCompressorArgs"
+
+    def add_callbacks_post_trainer(self, cfg: Any, trainer: Trainer) -> list:
+        """
+        Adds Sparse Finetuning callback to the Trainer instance.
+
+        Args:
+            cfg (Any): Configuration object containing the sparse recipe.
+            trainer (Trainer): Huggingface Trainer instance.
+
+        Returns:
+            list: List containing the configured callback instances.
+        """
+        LOG.info("Adding Sparse Finetuning callback to the trainer")
+        callback = LLMCompressorCallbackHandler(
+            trainer=trainer,
+            recipe=cfg.llmcompressor.recipe,
+        )
+        return [callback]
+
+
+def compute_loss_wrapper(
+    compute_loss_func: Callable[Concatenate[Module, P], R],
+) -> Callable[Concatenate[Module, P], R]:
+    """
+    Wraps the loss computation function to trigger the loss_calculated callback.
+
+    Args:
+        compute_loss_func (Callable): Original loss computation function.
+
+    Returns:
+        Callable: Wrapped function that also invokes the loss_calculated callback.
+    """
+
+    @wraps(compute_loss_func)
+    def compute_and_notify(model: Module, *args: P.args, **kwargs: P.kwargs) -> R:
+        loss = compute_loss_func(model, *args, **kwargs)
+        if active_session().lifecycle.initialized_ and model.training:
+            session_callbacks.loss_calculated(loss=loss)
+        return loss
+
+    return compute_and_notify
diff --git a/src/axolotl/integrations/llm_compressor/utils.py b/src/axolotl/integrations/llm_compressor/utils.py
new file mode 100644
index 0000000000..1abcb6bd4b
--- /dev/null
+++ b/src/axolotl/integrations/llm_compressor/utils.py
@@ -0,0 +1,37 @@
+"""Utilities for llmcompressor integration with axolotl."""
+
+from typing import Union
+
+from llmcompressor.transformers.sparsification.compressed_tensors_utils import (
+    modify_save_pretrained,
+)
+from transformers import PreTrainedModel, Trainer
+
+
+def save_compressed_model(
+    model: PreTrainedModel,
+    output_dir: Union[str, bytes],
+    trainer: Trainer,
+    save_compressed: bool = False,
+) -> None:
+    """
+    Synchronize processes, apply compression hooks, and save the model.
+
+    Args:
+        model (PreTrainedModel): The model to be saved.
+        output_dir (str or bytes): Path where the model files will be written.
+        trainer (Trainer): Hugging Face Trainer for process synchronization.
+        save_compressed (bool): Write compressed tensors if True.
+    """
+    trainer.accelerator.wait_for_everyone()
+
+    # Only the main process writes the files
+    if not trainer.accelerator.is_main_process:
+        return
+
+    modify_save_pretrained(model)
+    model.save_pretrained(
+        output_dir,
+        save_compressed=save_compressed,
+        skip_sparsity_compression_stats=not save_compressed,
+    )
diff --git a/src/axolotl/integrations/lm_eval/README.md b/src/axolotl/integrations/lm_eval/README.md
new file mode 100644
index 0000000000..860c6681d7
--- /dev/null
+++ b/src/axolotl/integrations/lm_eval/README.md
@@ -0,0 +1,82 @@
+# LM Eval Harness
+
+Run evaluation on model using the popular lm-evaluation-harness library.
+
+See https://github.com/EleutherAI/lm-evaluation-harness
+
+## Usage
+
+There are two ways to use the LM Eval integration:
+
+### 1. Post-Training Evaluation
+
+When training with the plugin enabled, evaluation runs automatically after training completes:
+
+```yaml
+plugins:
+  - axolotl.integrations.lm_eval.LMEvalPlugin
+
+lm_eval_tasks:
+  - gsm8k
+  - hellaswag
+  - arc_easy
+
+lm_eval_batch_size: # Batch size for evaluation
+
+# Directory to save evaluation results.
+# The final model is loaded from this directory
+# unless specified otherwise (see below)
+output_dir:
+```
+
+Run training as usual:
+```bash
+axolotl train config.yml
+```
+
+### 2. Standalone CLI Evaluation
+
+Evaluate any model directly without training:
+
+```yaml
+lm_eval_model: meta-llama/Llama-2-7b-hf
+
+plugins:
+  - axolotl.integrations.lm_eval.LMEvalPlugin
+
+lm_eval_tasks:
+  - gsm8k
+  - hellaswag
+  - arc_easy
+
+lm_eval_batch_size: 8
+output_dir: ./outputs
+```
+
+Run evaluation:
+```bash
+axolotl lm-eval config.yml
+```
+
+## Model Selection Priority
+
+The model to evaluate is selected in the following priority order:
+
+1. **`lm_eval_model`** - Explicit model path or HuggingFace repo (highest priority)
+2. **`hub_model_id`** - Trained model pushed to HuggingFace Hub
+3. **`output_dir`** - Local checkpoint directory containing trained model weights
+
+## Citation
+
+```bib
+@misc{eval-harness,
+  author       = {Gao, Leo and Tow, Jonathan and Abbasi, Baber and Biderman, Stella and Black, Sid and DiPofi, Anthony and Foster, Charles and Golding, Laurence and Hsu, Jeffrey and Le Noac'h, Alain and Li, Haonan and McDonell, Kyle and Muennighoff, Niklas and Ociepa, Chris and Phang, Jason and Reynolds, Laria and Schoelkopf, Hailey and Skowron, Aviya and Sutawika, Lintang and Tang, Eric and Thite, Anish and Wang, Ben and Wang, Kevin and Zou, Andy},
+  title        = {A framework for few-shot language model evaluation},
+  month        = 07,
+  year         = 2024,
+  publisher    = {Zenodo},
+  version      = {v0.4.3},
+  doi          = {10.5281/zenodo.12608602},
+  url          = {https://zenodo.org/records/12608602}
+}
+```
diff --git a/src/axolotl/integrations/lm_eval/__init__.py b/src/axolotl/integrations/lm_eval/__init__.py
new file mode 100644
index 0000000000..732ce15925
--- /dev/null
+++ b/src/axolotl/integrations/lm_eval/__init__.py
@@ -0,0 +1,37 @@
+"""
+Module for the Plugin for LM Eval Harness
+"""
+
+import subprocess  # nosec
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.integrations.lm_eval.cli import build_lm_eval_command, get_model_path
+
+from .args import LMEvalArgs as LMEvalArgs
+
+
+class LMEvalPlugin(BasePlugin):
+    """
+    Plugin for LM Evaluation Harness integraton with Axolotl.
+    """
+
+    def get_input_args(self):
+        return "axolotl.integrations.lm_eval.LMEvalArgs"
+
+    def post_train_unload(self, cfg):
+        if cfg.lm_eval_post_train:
+            for lm_eval_args in build_lm_eval_command(
+                cfg.lm_eval_tasks,
+                bfloat16=cfg.bfloat16 or cfg.bf16,
+                flash_attention=cfg.attn_uses_flash_lib,
+                output_dir=cfg.output_dir,
+                batch_size=cfg.lm_eval_batch_size,
+                wandb_project=cfg.wandb_project,
+                wandb_entity=cfg.wandb_entity,
+                wandb_name=cfg.wandb_name,
+                model=get_model_path(cfg),
+            ):
+                subprocess.run(  # nosec
+                    lm_eval_args,
+                    check=True,
+                )
diff --git a/src/axolotl/integrations/lm_eval/args.py b/src/axolotl/integrations/lm_eval/args.py
new file mode 100644
index 0000000000..d022131777
--- /dev/null
+++ b/src/axolotl/integrations/lm_eval/args.py
@@ -0,0 +1,18 @@
+"""
+Module for handling lm eval harness input arguments.
+"""
+
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class LMEvalArgs(BaseModel):
+    """
+    Input args for lm eval harness
+    """
+
+    lm_eval_tasks: List[str] = []
+    lm_eval_batch_size: Optional[int] = 8
+    lm_eval_post_train: Optional[bool] = True
+    lm_eval_model: Optional[str] = None
diff --git a/src/axolotl/integrations/lm_eval/cli.py b/src/axolotl/integrations/lm_eval/cli.py
new file mode 100644
index 0000000000..a20f4d1547
--- /dev/null
+++ b/src/axolotl/integrations/lm_eval/cli.py
@@ -0,0 +1,142 @@
+"""
+axolotl CLI for running lm_eval tasks
+"""
+
+import subprocess  # nosec
+from collections import defaultdict
+from datetime import datetime
+from typing import Optional
+
+import click
+import yaml
+
+from axolotl.utils.dict import DictDefault
+
+
+def get_model_path(cfg: DictDefault) -> str | None:
+    """
+    Determine which model path to use for evaluation.
+
+    Priority order (highest to lowest):
+    1. lm_eval_model - Explicit model path override
+    2. hub_model_id - Model pushed to HuggingFace Hub
+    3. None - Falls back to output_dir in build_lm_eval_command
+
+    Returns:
+        Model path string or None to use output_dir fallback
+    """
+    return cfg.lm_eval_model or cfg.hub_model_id or None
+
+
+def build_lm_eval_command(
+    tasks: list[str],
+    bfloat16=True,
+    flash_attention=False,
+    output_dir="./",
+    batch_size=8,
+    wandb_project=None,
+    wandb_entity=None,
+    wandb_name=None,
+    model=None,
+    revision=None,
+    apply_chat_template=None,
+    fewshot_as_multiturn=None,
+):
+    tasks_by_num_fewshot: dict[str, list] = defaultdict(list)
+    if isinstance(tasks, str):
+        tasks = [tasks]
+    for task in tasks:
+        num_fewshot = "-1"
+        task_parts = task.split(":")
+        task_name = task_parts[0]
+        if len(task_parts) == 2:
+            task_name, num_fewshot = task_parts
+        tasks_by_num_fewshot[str(num_fewshot)].append(task_name)
+
+    for num_fewshot, tasks_list in tasks_by_num_fewshot.items():
+        tasks_str = ",".join(tasks_list)
+        num_fewshot_val = num_fewshot if num_fewshot != "-1" else None
+        pretrained = "pretrained="
+        pretrained += model if model else output_dir
+        fa2 = ",attn_implementation=flash_attention_2" if flash_attention else ""
+        dtype = ",dtype=bfloat16" if bfloat16 else ",dtype=float16"
+        revision = f",revision={revision}" if revision else ""
+        output_path = output_dir
+        output_path += "" if output_dir.endswith("/") else "/"
+        output_path += "lm_eval_results/" + datetime.now().strftime("%Y%m%d_%H%M%S")
+        lm_eval_args = [
+            "lm_eval",
+            "--model",
+            "hf",
+            "--model_args",
+            f"{pretrained}{fa2}{dtype}{revision}",
+            "--tasks",
+            tasks_str,
+            "--batch_size",
+            str(batch_size),
+            "--output_path",
+            output_path,
+        ]
+        wandb_args = []
+        if wandb_project:
+            wandb_args.append(f"project={wandb_project}")
+        if wandb_entity:
+            wandb_args.append(f"entity={wandb_entity}")
+        if wandb_name:
+            wandb_args.append(f"name={wandb_name}")
+        if wandb_args:
+            lm_eval_args.append("--wandb_args")
+            lm_eval_args.append(",".join(wandb_args))
+        if apply_chat_template:
+            lm_eval_args.append("--apply_chat_template")
+        if num_fewshot_val:
+            lm_eval_args.append("--num_fewshot")
+            lm_eval_args.append(str(num_fewshot_val))
+            if apply_chat_template and fewshot_as_multiturn:
+                lm_eval_args.append("--fewshot_as_multiturn")
+
+        yield lm_eval_args
+
+
+@click.command()
+@click.argument("config", type=click.Path(exists=True, path_type=str))
+@click.option("--cloud", default=None, type=click.Path(exists=True, path_type=str))
+def lm_eval(config: str, cloud: Optional[str] = None):
+    """
+    use lm eval to evaluate a trained language model
+    """
+
+    if cloud:
+        from axolotl.cli.cloud import do_cli_lm_eval
+
+        do_cli_lm_eval(cloud_config=cloud, config=config)
+    else:
+        with open(config, encoding="utf-8") as file:
+            cfg: DictDefault = DictDefault(yaml.safe_load(file))
+
+        # This path operates on raw YAML via DictDefault (not the validated
+        # AxolotlInputConfig), so we resolve flash-attn from either the canonical
+        # `attn_implementation` field or the deprecated `flash_attention` boolean.
+        _flash_attn_impls = {"flash_attention_2", "flash_attention_3"}
+        lm_eval_flash_attention = bool(
+            cfg.flash_attention or cfg.attn_implementation in _flash_attn_impls
+        )
+
+        for lm_eval_args in build_lm_eval_command(
+            cfg.lm_eval_tasks,
+            bfloat16=cfg.bfloat16 or cfg.bf16,
+            flash_attention=lm_eval_flash_attention,
+            output_dir=cfg.output_dir,
+            batch_size=cfg.lm_eval_batch_size,
+            wandb_project=cfg.wandb_project,
+            wandb_entity=cfg.wandb_entity,
+            wandb_name=cfg.wandb_name,
+            model=get_model_path(cfg),
+            revision=cfg.revision,
+            apply_chat_template=cfg.apply_chat_template,
+            fewshot_as_multiturn=cfg.fewshot_as_multiturn,
+        ):
+            subprocess.run(  # nosec
+                lm_eval_args,
+                check=True,
+            )
diff --git a/src/axolotl/integrations/mora/__init__.py b/src/axolotl/integrations/mora/__init__.py
new file mode 100644
index 0000000000..8f50258e24
--- /dev/null
+++ b/src/axolotl/integrations/mora/__init__.py
@@ -0,0 +1,6 @@
+"""MoRA / ReMoRA integration for Axolotl."""
+
+from .args import MoraArgs, MoraConfig, MoraType
+from .plugin import MoraPlugin
+
+__all__ = ["MoraArgs", "MoraConfig", "MoraPlugin", "MoraType"]
diff --git a/src/axolotl/integrations/mora/args.py b/src/axolotl/integrations/mora/args.py
new file mode 100644
index 0000000000..01549e1db9
--- /dev/null
+++ b/src/axolotl/integrations/mora/args.py
@@ -0,0 +1,66 @@
+"""Config args for MoRA / ReMoRA."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class MoraType(str, Enum):
+    """MoRA variants supported by the reference implementation."""
+
+    SHARING = "sharing"
+    ROPE = "rope"
+
+    @property
+    def peft_value(self) -> int:
+        return {
+            MoraType.SHARING: 1,
+            MoraType.ROPE: 6,
+        }[self]
+
+
+class MoraConfig(BaseModel):
+    """Nested MoRA configuration available under the `mora` key."""
+
+    use_mora: bool = Field(
+        default=True,
+        description=(
+            "Enable MoRA adapter construction. Requires a PEFT build with MoRA "
+            "support (for example, the MoRA fork)."
+        ),
+    )
+    mora_type: MoraType = Field(
+        default=MoraType.ROPE,
+        description=(
+            "MoRA variant selector. Supported values are `sharing` for type 1 "
+            "and `rope` for type 6. Numeric values 1 and 6 are accepted for "
+            "backwards compatibility."
+        ),
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def normalize_mora_type(cls, data):
+        if not isinstance(data, dict) or "mora_type" not in data:
+            return data
+        data = data.copy()
+        mora_type = data["mora_type"]
+        if mora_type == 1:
+            data["mora_type"] = MoraType.SHARING
+        elif mora_type == 6:
+            data["mora_type"] = MoraType.ROPE
+        return data
+
+
+class MoraArgs(BaseModel):
+    """Plugin entry that exposes the nested `mora` block to the core config."""
+
+    mora: MoraConfig = Field(
+        default_factory=MoraConfig,
+        description=(
+            "MoRA / ReMoRA training configuration. Register the "
+            "`axolotl.integrations.mora.MoraPlugin` plugin to enable this block."
+        ),
+    )
diff --git a/src/axolotl/integrations/mora/plugin.py b/src/axolotl/integrations/mora/plugin.py
new file mode 100644
index 0000000000..8ca6068c18
--- /dev/null
+++ b/src/axolotl/integrations/mora/plugin.py
@@ -0,0 +1,97 @@
+"""MoRA / ReMoRA plugin for Axolotl."""
+
+import inspect
+
+from peft import LoraConfig, PeftModel
+from transformers import PreTrainedModel
+
+from axolotl.integrations.base import AdapterCapabilities, BasePlugin
+from axolotl.integrations.mora.args import MoraType
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _peft_supports_mora() -> bool:
+    try:
+        params = inspect.signature(LoraConfig).parameters
+    except (TypeError, ValueError):
+        return False
+    return "use_mora" in params and "mora_type" in params
+
+
+def _mora_type_peft_value(mora_type: MoraType | str | int) -> int:
+    if isinstance(mora_type, MoraType):
+        return mora_type.peft_value
+    if mora_type == 1 or mora_type == MoraType.SHARING.value:
+        return MoraType.SHARING.peft_value
+    if mora_type == 6 or mora_type == MoraType.ROPE.value:
+        return MoraType.ROPE.peft_value
+    raise ValueError("mora_type must be one of `sharing`, `rope`, 1, or 6")
+
+
+def _mora_type_label(mora_type: MoraType | str | int) -> str:
+    if isinstance(mora_type, MoraType):
+        return mora_type.value
+    if mora_type == 1:
+        return MoraType.SHARING.value
+    if mora_type == 6:
+        return MoraType.ROPE.value
+    return str(mora_type)
+
+
+class MoraPlugin(BasePlugin):
+    """Plugin that exposes MoRA-specific config and validates runtime support."""
+
+    def get_input_args(self) -> str:
+        return "axolotl.integrations.mora.MoraArgs"
+
+    def get_adapter_capabilities(self) -> list[AdapterCapabilities]:
+        return [AdapterCapabilities(name="mora", lora_like=True, relora=True)]
+
+    def _validate_mora_config(self, cfg: DictDefault):
+        mora_cfg = getattr(cfg, "mora", None)
+        if mora_cfg is None:
+            raise ValueError("adapter: mora requires a nested mora configuration block")
+        if not getattr(mora_cfg, "use_mora", False):
+            raise ValueError("mora.use_mora must be true when adapter: mora is set")
+        if cfg.load_in_4bit or cfg.load_in_8bit:
+            raise ValueError(
+                "adapter: mora currently requires a full-precision base model. "
+                "Use adapter: lora or qlora for quantized training."
+            )
+        if cfg.gptq:
+            raise ValueError(
+                "adapter: mora is not compatible with GPTQ quantized base models."
+            )
+
+    def get_lora_config_kwargs(self, cfg: DictDefault) -> dict:
+        if cfg.adapter != "mora":
+            return {}
+        self._validate_mora_config(cfg)
+        if not _peft_supports_mora():
+            raise ImportError(
+                "adapter: mora requires a PEFT build with MoRA support "
+                "(LoraConfig(use_mora=..., mora_type=...)). "
+                "Install the MoRA fork or another PEFT distribution that exposes "
+                "those fields."
+            )
+        mora_cfg = cfg.mora
+        return {
+            "use_mora": mora_cfg.use_mora,
+            "mora_type": _mora_type_peft_value(mora_cfg.mora_type),
+        }
+
+    def pre_model_load(self, cfg: DictDefault):
+        if cfg.adapter != "mora":
+            return
+        LOG.info("MoRA plugin enabled for adapter: mora")
+
+    def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        if cfg.adapter == "mora" and getattr(cfg, "mora", None):
+            LOG.debug(
+                "Loaded MoRA model with mora_type=%s, relora=%s",
+                _mora_type_label(cfg.mora.mora_type),
+                cfg.relora,
+            )
diff --git a/src/axolotl/integrations/nemo_gym/README.md b/src/axolotl/integrations/nemo_gym/README.md
new file mode 100644
index 0000000000..2a1f267c10
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/README.md
@@ -0,0 +1,412 @@
+# NeMo Gym Integration for Axolotl
+
+Train LLMs with reinforcement learning using [NVIDIA NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) environments as reward sources. NeMo Gym provides 50+ verified RL environments spanning math, coding, tool-use, reasoning, and safety — each with deterministic reward signals.
+
+## Validated Training Paths
+
+| Path | Speed | Multi-turn | Architecture |
+|------|-------|------------|--------------|
+| **Async GRPO + Data Producer** | Fastest (3x) | Yes | `NemoGymDataProducer` replaces vLLM generation |
+| Standard GRPO + Data Producer | Baseline | Yes | Same producer, no async prefetch |
+| Standard GRPO + /verify | Simplest | No | Reward function calls /verify directly |
+| FSDP2 + /verify (2 GPU) | Distributed | No | `fsdp_version: 2` |
+
+Multi-turn uses `nemo_gym_multi_turn: true` which auto-enables the async trainer's
+data producer protocol. The plugin's `NemoGymDataProducer` calls NeMo Gym agent `/run`
+endpoints and returns `RolloutDataset` with proper IS correction, env_mask, and rewards.
+
+All paths tested end-to-end with Qwen3-0.6B + LoRA, logged to wandb project `nemo-gym-rl`.
+
+## Quick Start
+
+### Prerequisites
+
+- [uv](https://github.com/astral-sh/uv) package manager (for NeMo Gym's venv)
+- Two GPUs recommended (one for vLLM server, one for training)
+
+### 1. Set Up NeMo Gym
+
+```bash
+git clone https://github.com/NVIDIA-NeMo/Gym.git ~/Gym
+cd ~/Gym
+uv venv --python 3.12 && source .venv/bin/activate && uv sync
+
+# Fix pycosat build (GCC 13+)
+CFLAGS="" uv pip install pycosat --python .venv/bin/python --no-build-isolation
+
+# Pre-build resource server venvs
+for dir in resources_servers/reasoning_gym resources_servers/example_single_tool_call responses_api_models/vllm_model responses_api_agents/simple_agent; do
+    uv venv --seed --allow-existing --python 3.12 $dir/.venv
+    CFLAGS="" uv pip install --python $dir/.venv/bin/python pycosat --no-build-isolation 2>/dev/null
+    uv pip install --python $dir/.venv/bin/python -e . "ray[default]==2.52.1"
+done
+
+# Install extra deps for reasoning_gym
+uv pip install --python resources_servers/reasoning_gym/.venv/bin/python \
+    reasoning-gym matplotlib pillow cycler contourpy kiwisolver
+```
+
+### 2. Multi-Turn with Async GRPO (Recommended — Fastest Path)
+
+This is the fully validated, highest-performance path. NeMo Gym's agent server handles
+multi-turn tool execution while axolotl's async GRPO prefetches data in background threads.
+
+**Step 1: Create the NeMo Gym agent config**
+
+Create `~/Gym/configs/axolotl_tool_calling.yaml`:
+```yaml
+# Resource server (tools + verify)
+example_single_tool_call:
+  resources_servers:
+    example_single_tool_call:
+      entrypoint: app.py
+      domain: agent
+      verified: false
+
+# Model server proxy (forwards to your vLLM)
+policy_model:
+  responses_api_models:
+    vllm_model:
+      entrypoint: app.py
+      base_url: http://localhost:8000/v1
+      api_key: dummy_key
+      model: Qwen/Qwen3-0.6B   # Must match your training model
+      return_token_id_information: true
+      uses_reasoning_parser: false
+
+# Agent server (orchestrates multi-turn via /run)
+example_single_tool_call_simple_agent:
+  responses_api_agents:
+    simple_agent:
+      entrypoint: app.py
+      resources_server:
+        type: resources_servers
+        name: example_single_tool_call
+      model_server:
+        type: responses_api_models
+        name: policy_model
+      datasets:
+      - name: weather
+        type: example
+        jsonl_fpath: resources_servers/example_single_tool_call/data/weather_tool_calling.jsonl
+```
+
+**Step 2: Start three services**
+
+```bash
+# Terminal 1: vLLM OpenAI server on GPU 0
+CUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen3-0.6B --max-model-len 2048 --gpu-memory-utilization 0.85
+
+# Terminal 2: NeMo Gym (resource server + model proxy + agent)
+cd ~/Gym && .venv/bin/ng_run \
+    "+config_paths=[configs/axolotl_tool_calling.yaml]" "+skip_venv_if_present=true"
+
+# Terminal 3: Training on GPU 1
+cd experiments && CUDA_VISIBLE_DEVICES=1 CUDA_HOME=$HOME/env-claude-cu130/cuda_shim \
+    axolotl train nemo_gym_async_agent.yaml
+```
+
+**Step 3: Training config** (`nemo_gym_async_agent.yaml`):
+```yaml
+base_model: Qwen/Qwen3-0.6B
+adapter: lora
+lora_r: 16
+lora_alpha: 32
+lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj]
+sequence_len: 2048
+
+rl: grpo
+chat_template: tokenizer_default
+
+trl:
+  use_vllm: true
+  vllm_mode: server
+  vllm_server_host: localhost
+  vllm_server_port: 8000
+  vllm_lora_sync: true
+  vllm_sync_interval: 5
+  # Async GRPO — 3x faster than standard
+  use_data_producer: true
+  async_prefetch: true
+  num_generations: 4
+  max_completion_length: 512
+  temperature: 0.8
+  reward_funcs:
+    - axolotl.integrations.nemo_gym.rewards.reward_env
+
+plugins:
+  - axolotl.integrations.nemo_gym.NemoGymPlugin
+
+nemo_gym_enabled: true
+nemo_gym_auto_start: false
+nemo_gym_head_port: 11000
+nemo_gym_multi_turn: true
+nemo_gym_verify_timeout: 120
+nemo_gym_datasets:
+  - path: ~/Gym/resources_servers/example_single_tool_call/data/weather_tool_calling.jsonl
+    server_name: example_single_tool_call
+
+datasets:
+  - path: ~/Gym/resources_servers/example_single_tool_call/data/weather_tool_calling.jsonl
+    type: chat_template
+    field_messages: responses_create_params.input
+    message_field_content: content
+    message_field_role: role
+
+vllm:
+  gpu_memory_utilization: 0.85
+  max_model_len: 2048
+  tensor_parallel_size: 1
+
+learning_rate: 5e-6
+micro_batch_size: 1
+gradient_accumulation_steps: 4
+max_steps: 30
+gradient_checkpointing: true
+bf16: true
+output_dir: ./outputs/nemo_gym_async
+
+use_wandb: true
+wandb_project: nemo-gym-rl
+```
+
+### 3. Single-Turn Training (Simplest — No Agent Server Needed)
+
+For environments that only need single-turn verify (math, coding challenges), you don't need
+an agent server. The plugin's reward function calls `/verify` directly.
+
+```yaml
+base_model: Qwen/Qwen2.5-0.5B-Instruct
+rl: grpo
+chat_template: tokenizer_default
+
+trl:
+  use_vllm: true
+  vllm_mode: colocate
+  vllm_enable_sleep_mode: false
+  num_generations: 8
+  max_completion_length: 128
+  temperature: 0.9
+  reward_funcs:
+    - axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify
+
+plugins:
+  - axolotl.integrations.nemo_gym.NemoGymPlugin
+
+nemo_gym_enabled: true
+nemo_gym_auto_start: false
+nemo_gym_head_port: 11000
+nemo_gym_datasets:
+  - path: ~/Gym/resources_servers/reasoning_gym/data/train_basic_arithmetic.jsonl
+    server_name: reasoning_gym
+
+datasets:
+  - path: ~/Gym/resources_servers/reasoning_gym/data/train_basic_arithmetic.jsonl
+    type: chat_template
+    field_messages: responses_create_params.input
+    message_field_content: content
+    message_field_role: role
+
+vllm:
+  gpu_memory_utilization: 0.3
+  max_model_len: 512
+  tensor_parallel_size: 1
+
+learning_rate: 1e-5
+micro_batch_size: 4
+gradient_accumulation_steps: 2
+max_steps: 50
+output_dir: ./outputs/nemo_gym_arithmetic
+```
+
+Only needs `ng_run` with resource servers (no agent config):
+```bash
+cd ~/Gym && ng_run "+config_paths=[resources_servers/reasoning_gym/configs/resources_only.yaml]" "+skip_venv_if_present=true"
+```
+
+## How It Works
+
+### Single-Turn
+```text
+axolotl train → GRPO Trainer generates completions
+  → NeMo Gym plugin reward_fn calls POST /verify on resource server
+  → reward flows back to GRPO for advantage computation
+```
+
+### Multi-Turn (Agent /run)
+```text
+┌─────────────┐     ┌──────────────┐     ┌──────────────────┐
+│  axolotl    │     │  NeMo Gym    │────▶│  vLLM OpenAI     │
+│  train      │────▶│  Agent /run  │◀────│  Server (GPU 0)  │
+│  (GPU 1)    │     │              │     │  /v1/completions  │
+└─────────────┘     └──────┬───────┘     └──────────────────┘
+                           │
+                           ▼
+                    ┌──────────────┐
+                    │  Resource    │
+                    │  Server     │
+                    │  (tools +   │
+                    │   verify)   │
+                    └─────────────┘
+```
+
+The agent server orchestrates the entire multi-turn loop:
+1. Calls our vLLM server for model generation
+2. Parses tool calls from model output
+3. Executes tools against resource servers
+4. Feeds tool results back to the model
+5. Repeats until done, then calls /verify for reward
+6. Returns token IDs + logprobs + reward to our rollout_func
+
+### Data Producer Architecture (Multi-Turn)
+
+When `nemo_gym_multi_turn: true`, the plugin automatically forces `use_data_producer: true`
+which selects the `AxolotlAsyncGRPOTrainer`. The plugin then swaps the trainer's data
+producer with `NemoGymDataProducer`, which:
+
+1. Gets a prompt batch from the dataset iterator
+2. Expands by `num_generations` (one agent call per rollout)
+3. Calls NeMo Gym agents via async HTTP (`aiohttp.gather`)
+4. Parses responses into padded tensors (`RolloutDataset`)
+5. Returns with `_pending_policy_logps=True` for deferred scoring
+
+The main thread then runs `_compute_deferred_scores()` which:
+- Computes **policy logprobs** on the training model (GPU forward pass)
+- Computes **IS correction** using agent's sampling logprobs vs training model logprobs
+- Computes advantages with group-level normalization
+- All downstream features work: replay buffer, re-roll, streaming, zero-adv skip
+
+With `async_prefetch: true`, the data producer runs in a background thread — giving ~3x
+speedup as generation and training overlap. With `async_prefetch: false`, it runs
+synchronously on the main thread (still uses the data producer protocol).
+
+### Weight Sync (LoRA Mode)
+
+With `vllm_lora_sync: true`, the plugin (or async trainer) replaces NCCL-based weight
+sync with filesystem + HTTP:
+
+1. `accelerator.get_state_dict()` gathers LoRA weights from all ranks
+2. Rank 0 saves adapter to `/tmp/lora_sync_*/vN/`
+3. Rank 0 POSTs to `/set_lora_adapter/` on vLLM server
+4. vLLM loads adapter natively via Punica kernels
+5. Only ~40MB transferred (vs multiple GBs for full model weights)
+
+### Multi-Environment Support
+
+Datasets support per-row environment routing via `agent_ref`:
+```jsonl
+{"agent_ref": {"name": "reasoning_gym"}, "responses_create_params": {...}}
+{"agent_ref": {"name": "instruction_following"}, "responses_create_params": {...}}
+```
+
+Or use the simpler per-dataset routing:
+```yaml
+nemo_gym_datasets:
+  - path: reasoning_data.jsonl
+    server_name: reasoning_gym
+  - path: tool_data.jsonl
+    server_name: example_single_tool_call
+```
+
+## Configuration Reference
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `nemo_gym_enabled` | bool | `null` | Enable the NeMo Gym integration |
+| `nemo_gym_dir` | str | `~/Gym` | Path to NeMo Gym repo |
+| `nemo_gym_auto_clone` | bool | `true` | Auto-clone NeMo Gym repo if missing |
+| `nemo_gym_auto_start` | bool | `true` | Auto-start resource servers |
+| `nemo_gym_config_paths` | list[str] | — | Server config YAMLs (relative to gym_dir) |
+| `nemo_gym_datasets` | list[dict] | required | Dataset configs with `path` and optional `server_name` |
+| `nemo_gym_head_port` | int | `11000` | Head server port |
+| `nemo_gym_server_timeout` | int | `360` | Server startup timeout (seconds) |
+| `nemo_gym_verify_timeout` | int | `30` | Per-request timeout (seconds) |
+| `nemo_gym_multi_turn` | bool | `false` | Enable multi-turn via agent /run |
+
+### Dataset JSONL Format
+
+Each line must have `responses_create_params` with `input` messages:
+```json
+{
+  "responses_create_params": {
+    "input": [{"role": "user", "content": "What's the weather in SF?"}],
+    "tools": [{"name": "get_weather", "type": "function", "strict": true, "parameters": {...}}]
+  }
+}
+```
+
+For multi-turn agent routing, include `agent_ref`:
+```json
+{"agent_ref": {"name": "my_agent"}, "responses_create_params": {...}}
+```
+
+Note: Tool definitions MUST include `"strict": true` and `"additionalProperties": false` for NeMo Gym agent compatibility.
+
+### Reward Functions
+
+The plugin provides two built-in reward functions — no user code needed:
+
+```yaml
+trl:
+  reward_funcs:
+    # Multi-turn (nemo_gym_multi_turn: true):
+    # Passthrough — agent /run already computed the reward
+    - axolotl.integrations.nemo_gym.rewards.reward_env
+
+    # Single-turn (nemo_gym_multi_turn: false):
+    # Calls /verify endpoints on NeMo Gym resource servers
+    - axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify
+```
+
+Both are also importable from Python:
+
+```python
+from axolotl.integrations.nemo_gym import reward_env, reward_nemo_gym_verify
+```
+
+## Known Issues / Troubleshooting
+
+### NeMo Gym Server Setup
+- **pycosat build failure**: `CFLAGS="" uv pip install pycosat --no-build-isolation`
+- **Ray version mismatch**: Pin `ray[default]==2.52.1` in all server venvs
+- **Pre-build venvs**: `ng_run` creates per-server venvs via Ray. Pre-build them and use `+skip_venv_if_present=true`
+- **Tool `strict` field required**: Agent server validates tool definitions require `strict: true`
+
+### vLLM / Weight Sync
+- **Start vLLM with LoRA + tool calling + runtime loading**:
+  ```bash
+  VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 \
+  CUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server \
+    --model Qwen/Qwen3-4B-Instruct-2507 \
+    --max-model-len 4096 \
+    --gpu-memory-utilization 0.7 \
+    --enable-lora --max-lora-rank 64 \
+    --enable-auto-tool-choice --tool-call-parser hermes
+  ```
+- **`VLLM_ALLOW_RUNTIME_LORA_UPDATING=1`**: Required for `vllm_lora_sync: true`. Without it, vLLM won't expose the `/v1/load_lora_adapter` endpoint and weight sync will fail silently. The plugin warns if this endpoint is missing.
+- **`--enable-lora`**: Enables LoRA adapter support in vLLM
+- **`--enable-auto-tool-choice --tool-call-parser hermes`**: Required for Qwen3 tool calling
+- **`max_model_len` must be > `max_completion_length`**: Leave room for prompt tokens (~200). If equal, the NeMo Gym model proxy gets a 400 error and returns empty completions.
+- **`CUDA_HOME` required**: DeepSpeed import needs it for the nvcc shim
+- **NCCL weight sync broken with vLLM 0.17**: Use `vllm_lora_sync: true` (filesystem + HTTP via `/v1/load_lora_adapter`)
+
+### Multi-Turn
+- **Agent server required**: Multi-turn delegates to NeMo Gym's agent server `/run` endpoint. Without an agent, the plugin falls back to single-turn `/verify`
+- **Model server proxy**: NeMo Gym needs a `responses_api_models` server that proxies to your vLLM. See the agent config example above
+
+### FSDP2
+- Validated on 2 GPUs with single-turn + LoRA
+- Async field filtering: The builder automatically filters async-only config fields when using the standard GRPO trainer
+
+## Comparison with Other Integrations
+
+| Feature | Axolotl + NeMo Gym | Unsloth + NeMo Gym | NeMo RL (native) |
+|---------|-------------------|-------------------|-------------------|
+| Server management | Automatic | Manual (notebook) | Built-in |
+| Multi-environment | Per-row routing | Manual code | YAML config |
+| Multi-turn / tool use | Agent /run delegation | No | Agent /run (Ray) |
+| Async GRPO (3x speedup) | Yes | No | Yes |
+| LoRA sync | Filesystem + HTTP | N/A | NCCL |
+| Multi-GPU (FSDP2) | Yes | No | Yes (Ray) |
+| Config-driven | Yes | No (code) | Yes |
diff --git a/src/axolotl/integrations/nemo_gym/__init__.py b/src/axolotl/integrations/nemo_gym/__init__.py
new file mode 100644
index 0000000000..b880bee7fa
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/__init__.py
@@ -0,0 +1,25 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+
+"""
+Plugin for NVIDIA NeMo Gym integration with Axolotl.
+
+NeMo Gym provides RL training environments for LLMs with verification-based
+reward signals. This plugin manages the NeMo Gym server lifecycle, loads
+datasets in the NeMo Gym JSONL format, and creates reward functions that
+call the NeMo Gym /verify endpoints.
+"""
+
+from .args import NemoGymArgs
+from .plugin import NemoGymPlugin
+from .rewards import reward_env, reward_nemo_gym_verify
+
+__all__ = [
+    "NemoGymArgs",
+    "NemoGymPlugin",
+    "reward_env",
+    "reward_nemo_gym_verify",
+]
diff --git a/src/axolotl/integrations/nemo_gym/args.py b/src/axolotl/integrations/nemo_gym/args.py
new file mode 100644
index 0000000000..3c593fa801
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/args.py
@@ -0,0 +1,146 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+
+"""
+Input arguments for the NeMo Gym integration plugin.
+"""
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class NemoGymArgs(BaseModel):
+    """Configuration args for the NeMo Gym integration."""
+
+    nemo_gym_enabled: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Enable NeMo Gym integration for environment-based RL rewards."
+        },
+    )
+    nemo_gym_dir: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Path to the NeMo Gym repository clone. "
+                "If not set and nemo_gym_auto_clone is True, clones to ~/Gym."
+            )
+        },
+    )
+    nemo_gym_auto_clone: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Automatically clone the NeMo Gym repository if not present. "
+                "Defaults to True when nemo_gym_enabled is set."
+            )
+        },
+    )
+    nemo_gym_config_paths: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "List of NeMo Gym resource server config YAML paths, relative to nemo_gym_dir. "
+                "Example: ['resources_servers/reasoning_gym/configs/resources_only.yaml']"
+            )
+        },
+    )
+    nemo_gym_head_port: int | None = Field(
+        default=11000,
+        json_schema_extra={
+            "description": "Port for the NeMo Gym head server. Defaults to 11000."
+        },
+    )
+    nemo_gym_server_timeout: int | None = Field(
+        default=360,
+        json_schema_extra={
+            "description": "Timeout in seconds waiting for NeMo Gym servers to start. Defaults to 360."
+        },
+    )
+    nemo_gym_verify_timeout: int | None = Field(
+        default=30,
+        json_schema_extra={
+            "description": "Timeout in seconds for individual /verify requests. Defaults to 30."
+        },
+    )
+    nemo_gym_run_timeout: int | None = Field(
+        default=300,
+        json_schema_extra={
+            "description": (
+                "Timeout in seconds for each agent /run request (one multi-turn rollout). "
+                "Prevents stuck generations (e.g. model looping on <think> tags) from "
+                "blocking training indefinitely. Defaults to 300 (5 minutes)."
+            )
+        },
+    )
+    nemo_gym_datasets: list[dict] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "List of NeMo Gym dataset configs. Each entry has 'path' (JSONL file path "
+                "relative to nemo_gym_dir) and optionally 'server_name' (default resource server). "
+                "If the JSONL rows have agent_ref.name, that takes precedence per row, "
+                "enabling multi-environment training from a single dataset file. "
+                "Optional 'max_samples' to limit per dataset."
+            )
+        },
+    )
+    nemo_gym_auto_start: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": (
+                "Automatically start NeMo Gym resource servers. Defaults to True. "
+                "Set to False if servers are already running externally."
+            )
+        },
+    )
+    nemo_gym_model_name: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Model name to report in verify requests. "
+                "Defaults to the base_model from the main config."
+            )
+        },
+    )
+    nemo_gym_multi_turn: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Enable multi-turn rollouts via NeMo Gym. When True, uses TRL's "
+                "rollout_func to run multi-step interactions with tool execution. "
+                "Requires use_vllm=True in TRL config. The model generates responses, "
+                "tool calls are executed against resource servers, and results are "
+                "fed back for the next turn. Final reward comes from /verify."
+            )
+        },
+    )
+    nemo_gym_max_turns: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Maximum number of turns per multi-turn rollout. Defaults to 10. "
+                "Each turn consists of a model generation + optional tool execution."
+            )
+        },
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_nemo_gym_config(cls, data):
+        if data.get("nemo_gym_enabled"):
+            if not data.get("nemo_gym_config_paths") and data.get(
+                "nemo_gym_auto_start", True
+            ):
+                raise ValueError(
+                    "nemo_gym_config_paths is required when nemo_gym_enabled=True "
+                    "and nemo_gym_auto_start is not False."
+                )
+            if not data.get("nemo_gym_datasets"):
+                raise ValueError(
+                    "nemo_gym_datasets is required when nemo_gym_enabled=True. "
+                    "Provide at least one dataset with 'path' and 'server_name'."
+                )
+        return data
diff --git a/src/axolotl/integrations/nemo_gym/data_producer.py b/src/axolotl/integrations/nemo_gym/data_producer.py
new file mode 100644
index 0000000000..3a9635d154
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/data_producer.py
@@ -0,0 +1,279 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+
+"""
+NeMo Gym Data Producer for async GRPO training.
+
+Replaces GRPODataProducer to generate rollouts via NeMo Gym agent /run endpoints
+instead of vLLM. The agent handles generation, tool execution, and reward computation.
+Returns RolloutDataset in the same format as the standard producer, so all downstream
+components (deferred scoring, IS correction, streaming, replay, re-roll) work unchanged.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import torch
+from trl.trainer.utils import pad
+
+from axolotl.core.trainers.grpo.async_trainer import GRPODataProducer, RolloutDataset
+from axolotl.utils.logging import get_logger
+
+from .multi_turn import _call_agents, _parse_agent_response
+
+LOG = get_logger(__name__)
+
+
+class NemoGymDataProducer(GRPODataProducer):
+    """Produces GRPO rollouts by calling NeMo Gym agent /run endpoints.
+
+    Drop-in replacement for GRPODataProducer. Instead of calling vLLM for generation,
+    sends prompts to NeMo Gym agents which handle generation + tool execution + reward.
+    Returns the same RolloutDataset format so deferred scoring, IS correction,
+    replay buffer, and re-roll all work unchanged.
+    """
+
+    def __init__(
+        self,
+        *args,
+        agent_servers: dict[str, str],
+        dataset_lookup: dict,
+        request_timeout: float = 10800,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self._agent_servers = agent_servers
+        self._dataset_lookup = dataset_lookup
+        self._request_timeout = request_timeout
+
+    def produce(
+        self,
+        model: Any,
+        global_step: int,
+        *,
+        skip_policy_logps: bool = False,
+        processing_class: Any = None,
+        accelerator: Any = None,
+        args: Any = None,
+        _rank0_only: bool = False,
+        **kwargs,
+    ) -> RolloutDataset | None:
+        """Generate rollouts via NeMo Gym agents.
+
+        Calls agent /run endpoints, parses responses into padded tensors,
+        and returns a RolloutDataset for deferred scoring on the main thread.
+        """
+        trainer = self._trainer
+        is_main = trainer.accelerator.is_main_process
+        device = trainer.accelerator.device
+
+        if _rank0_only and not is_main:
+            return None
+
+        # Get prompt batch from iterator
+        try:
+            inputs = next(self._prompt_iter)
+        except StopIteration:
+            self._prompt_iter = iter(self._prompt_dl)
+            inputs = next(self._prompt_iter)
+
+        # Extract dataset items for agent calls
+        dataset_items = []
+        for inp in inputs:
+            prompt_text = ""
+            prompt = inp.get("prompt", [])
+            if isinstance(prompt, list) and prompt:
+                prompt_text = (
+                    prompt[-1].get("content", "")
+                    if isinstance(prompt[-1], dict)
+                    else str(prompt[-1])
+                )
+            elif isinstance(prompt, str):
+                prompt_text = prompt
+
+            # Find the full dataset item, preserving agent_ref for routing
+            full_item = self._dataset_lookup.get(prompt_text, {})
+            item = full_item.get("verify_extra", {})
+            if not item:
+                item = {
+                    "responses_create_params": {
+                        "input": [{"role": "user", "content": prompt_text}]
+                    }
+                }
+            # Preserve agent_ref from the dataset row for _call_agents routing
+            if "agent_ref" in full_item and "agent_ref" not in item:
+                item["agent_ref"] = full_item["agent_ref"]
+            dataset_items.append(item)
+
+        # NOTE: do NOT re-expand by num_generations here.
+        # ``RepeatSampler(mini_repeat_count=num_generations)`` already
+        # yields ``num_generations`` consecutive copies of each unique
+        # prompt, so ``inputs`` is a list of ``(unique_prompts_per_rank *
+        # num_generations)`` items — one entry per rollout. Expanding
+        # again here would fire ``num_generations^2`` rollouts per
+        # prompt per rank and make every step dogpile on a handful of
+        # tasks.
+        expanded_items = dataset_items
+
+        # Diagnostic: log what this rank is about to fire.
+        try:
+            import collections
+
+            iid_counts: collections.Counter[str | None] = collections.Counter()
+            for it in dataset_items:
+                iid_counts[
+                    (it.get("responses_create_params", {}).get("metadata") or {}).get(
+                        "instance_id"
+                    )
+                ] += 1
+            LOG.info(
+                "[RANK:%d] produce(): firing %d agent /run calls covering %d unique prompts: %s",
+                trainer.accelerator.process_index,
+                len(dataset_items),
+                len(iid_counts),
+                list(iid_counts.most_common(5)),
+            )
+        except Exception:
+            pass
+
+        # Call NeMo Gym agents
+        loop = asyncio.new_event_loop()
+        try:
+            responses = loop.run_until_complete(
+                _call_agents(
+                    dataset_items=expanded_items,
+                    agent_servers=self._agent_servers,
+                    timeout=self._request_timeout,
+                    max_completion_length=trainer.max_completion_length,
+                    temperature=trainer.temperature,
+                    top_p=getattr(trainer, "top_p", None) or 0.999,
+                )
+            )
+        finally:
+            loop.close()
+
+        # Parse responses
+        eos_token_id = trainer.processing_class.eos_token_id
+        prompt_ids_list = []
+        completion_ids_list = []
+        env_mask_list = []
+        logprobs_list = []
+        rewards_list = []
+
+        num_turns_list: list[int] = []
+        for resp in responses:
+            parsed = _parse_agent_response(resp, eos_token_id)
+            prompt_ids_list.append(parsed["prompt_ids"])
+            completion_ids_list.append(parsed["completion_ids"])
+            env_mask_list.append(parsed["env_mask"])
+            logprobs_list.append(parsed["logprobs"])
+            rewards_list.append(parsed["reward"])
+            num_turns_list.append(parsed.get("num_turns", 0))
+
+        # Pad to tensors
+        prompt_ids = [torch.tensor(ids, device=device) for ids in prompt_ids_list]
+        prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids]
+        prompt_ids = pad(
+            prompt_ids, padding_value=trainer.pad_token_id, padding_side="left"
+        )
+        prompt_mask = pad(prompt_mask, padding_value=0, padding_side="left")
+
+        completion_ids = [
+            torch.tensor(ids, device=device) for ids in completion_ids_list
+        ]
+        completion_mask = [
+            torch.ones_like(ids, dtype=torch.long) for ids in completion_ids
+        ]
+        completion_ids = pad(
+            completion_ids, padding_value=trainer.pad_token_id, padding_side="right"
+        )
+        completion_mask = pad(completion_mask, padding_value=0, padding_side="right")
+
+        # Sampling logprobs from agent (used for IS correction)
+        sampling_logps = [
+            torch.tensor(lp, dtype=torch.float32, device=device) for lp in logprobs_list
+        ]
+        sampling_per_token_logps = pad(
+            sampling_logps, padding_value=0.0, padding_side="right"
+        )
+
+        # env_mask as tool_mask (1=model tokens, 0=tool tokens)
+        tool_mask = [torch.tensor(m, device=device) for m in env_mask_list]
+        tool_mask = pad(tool_mask, padding_value=1, padding_side="right")
+
+        # Inject per-rollout reward + num_turns into each input. Since
+        # ``RepeatSampler`` already yields ``num_generations`` copies of
+        # each prompt, ``inputs`` has ONE entry per rollout (matching
+        # ``rewards_list`` 1:1). No per-prompt grouping happens here —
+        # GRPO advantage normalization is the trainer's job downstream.
+        assert len(inputs) == len(rewards_list), (
+            f"rewards/inputs length mismatch: "
+            f"{len(rewards_list)} rewards vs {len(inputs)} inputs"
+        )
+        for i, inp in enumerate(inputs):
+            inp["env_reward"] = rewards_list[i]
+            inp["num_turns"] = num_turns_list[i]
+
+        # One expanded_input per rollout (already correct count because
+        # inputs has num_generations copies baked in by the sampler).
+        expanded_inputs = [dict(inp) for inp in inputs]
+
+        # Log rollout-level stats to wandb from rank 0. These are the
+        # true agent-side metrics (not the tokenized TRL view) — so
+        # num_turns reflects how many /run iterations each rollout
+        # actually took before finishing or hitting max_turns.
+        if is_main and num_turns_list:
+            try:
+                import wandb
+
+                if wandb.run is not None:
+                    import statistics as _stats
+
+                    nonzero = sum(1 for r in rewards_list if r > 0)
+                    log_payload = {
+                        "rollout/num_turns/mean": float(_stats.mean(num_turns_list)),
+                        "rollout/num_turns/min": float(min(num_turns_list)),
+                        "rollout/num_turns/max": float(max(num_turns_list)),
+                        "rollout/reward/mean": float(_stats.mean(rewards_list)),
+                        "rollout/reward/nonzero_frac": (
+                            nonzero / len(rewards_list) if rewards_list else 0.0
+                        ),
+                        "rollout/n_samples": float(len(rewards_list)),
+                    }
+                    wandb.log(log_payload, commit=False)
+            except Exception as exc:  # never let metric logging break training
+                LOG.warning("rollout wandb log failed: %s", exc)
+
+        # Decode completions for reward functions
+        completions = trainer.processing_class.batch_decode(
+            completion_ids, skip_special_tokens=True
+        )
+
+        # Build total token count
+        num_items_in_batch = completion_mask.sum()
+
+        # Build output dict (same shape as _generate_only)
+        output = {
+            "prompt_ids": prompt_ids,
+            "prompt_mask": prompt_mask,
+            "completion_ids": completion_ids,
+            "completion_mask": completion_mask,
+            "num_items_in_batch": num_items_in_batch,
+            "advantages": torch.zeros(completion_ids.size(0), device=device),
+            "sampling_per_token_logps": sampling_per_token_logps,
+            "tool_mask": tool_mask,
+            # Deferred scoring markers
+            "_pending_policy_logps": True,
+            "_deferred_inputs": expanded_inputs,
+            "_deferred_prompts": [inp.get("prompt", "") for inp in expanded_inputs],
+            "_deferred_completions": completions,
+            "_deferred_completion_ids_list": completion_ids_list,
+            "_rank0_only": _rank0_only,
+        }
+
+        return RolloutDataset(output)
diff --git a/src/axolotl/integrations/nemo_gym/dataset.py b/src/axolotl/integrations/nemo_gym/dataset.py
new file mode 100644
index 0000000000..be53da52d8
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/dataset.py
@@ -0,0 +1,135 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+
+"""
+Dataset loading for NeMo Gym JSONL files.
+
+Converts NeMo Gym JSONL format into HuggingFace Datasets compatible
+with TRL's GRPOTrainer. Supports multi-environment routing via:
+  1. Per-dataset server_name (all rows in a file go to one server)
+  2. Per-row agent_ref.name (each row specifies its own server)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import random
+
+from datasets import Dataset
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def load_nemo_gym_datasets(
+    gym_dir: str,
+    dataset_configs: list[dict],
+) -> Dataset:
+    """Load and merge NeMo Gym JSONL datasets with multi-environment support.
+
+    Each dataset config should have:
+        - path: JSONL file path (absolute, or relative to gym_dir)
+        - server_name: Default NeMo Gym server for this dataset.
+          Can be overridden per-row if the JSONL has an "agent_ref" field.
+        - max_samples (optional): Max number of samples to use from this dataset
+
+    Per-row routing: If a JSONL row has an "agent_ref": {"name": "..."} field,
+    that takes precedence over the dataset-level server_name. This allows mixing
+    environments within a single dataset file (matching TRL's pattern).
+
+    The output dataset has columns:
+        - prompt: list[dict] chat format
+        - resources_server_ref: dict with {"name": server_name}
+        - verify_extra: dict with original JSONL data for verify requests
+
+    Args:
+        gym_dir: Path to the NeMo Gym directory.
+        dataset_configs: List of dataset configuration dicts.
+
+    Returns:
+        A HuggingFace Dataset ready for GRPOTrainer.
+    """
+    all_examples = []
+
+    for ds_cfg in dataset_configs:
+        path = ds_cfg["path"]
+        default_server = ds_cfg.get("server_name", "")
+        max_samples = ds_cfg.get("max_samples")
+
+        # Resolve path
+        if not os.path.isabs(path):
+            path = os.path.join(gym_dir, path)
+        path = os.path.expanduser(path)
+
+        if not os.path.exists(path):
+            raise FileNotFoundError(
+                f"NeMo Gym dataset not found at {path}. "
+                "Ensure the dataset file exists or run the appropriate "
+                "NeMo Gym dataset creation script."
+            )
+
+        LOG.info(
+            f"Loading NeMo Gym dataset from {path} (default server: {default_server})"
+        )
+
+        with open(path, encoding="utf-8") as f:
+            lines = f.readlines()
+
+        if max_samples and len(lines) > max_samples:
+            lines = random.sample(lines, max_samples)  # nosec B311
+
+        for line in lines:
+            data = json.loads(line)
+
+            # Extract user prompt from the input messages
+            inputs = data.get("responses_create_params", {}).get("input", [])
+            task_prompt = ""
+            for inp in inputs:
+                if isinstance(inp, dict) and inp.get("role") in ("user",):
+                    task_prompt = inp.get("content", "")
+                    break
+            if not task_prompt and inputs:
+                # Fallback: use the last input's content
+                task_prompt = (
+                    inputs[-1].get("content", "")
+                    if isinstance(inputs[-1], dict)
+                    else ""
+                )
+
+            # Per-row agent routing: agent_ref.name can override dataset-level server_name.
+            # NeMo Gym datasets may use agent names (e.g., "reasoning_gym_simple_agent")
+            # which differ from resource server names (e.g., "reasoning_gym").
+            # The dataset-level server_name is always the fallback.
+            row_agent_ref = data.get("agent_ref", {})
+            server_name = default_server
+            if row_agent_ref and row_agent_ref.get("name"):
+                # Use per-row name, but only if it looks like a resource server name.
+                # Agent names typically have "_simple_agent" or "_agent" suffix.
+                row_name = row_agent_ref["name"]
+                if row_agent_ref.get("type") != "responses_api_agents":
+                    # Not an agent — could be a direct resource server reference
+                    server_name = row_name
+
+            all_examples.append(
+                {
+                    "prompt": [{"role": "user", "content": task_prompt}],
+                    "resources_server_ref": {"name": server_name},
+                    "verify_extra": data,
+                }
+            )
+
+    random.shuffle(all_examples)
+
+    # Log environment distribution
+    env_counts: dict[str, int] = {}
+    for ex in all_examples:
+        name = ex["resources_server_ref"]["name"]
+        env_counts[name] = env_counts.get(name, 0) + 1
+    LOG.info(f"Loaded {len(all_examples)} NeMo Gym examples: {env_counts}")
+
+    return Dataset.from_list(all_examples)
diff --git a/src/axolotl/integrations/nemo_gym/examples/nemo_gym_multi_env.yaml b/src/axolotl/integrations/nemo_gym/examples/nemo_gym_multi_env.yaml
new file mode 100644
index 0000000000..5a86f7578e
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/examples/nemo_gym_multi_env.yaml
@@ -0,0 +1,64 @@
+# Axolotl + NeMo Gym: Multi-Environment RL Training Example
+#
+# Trains on multiple NeMo Gym environments simultaneously:
+# - Mini Sudoku (reasoning_gym)
+# - Instruction Following
+#
+# Prerequisites:
+#   - uv (https://github.com/astral-sh/uv) installed
+#   - Download instruction_following.jsonl from nvidia/Nemotron-RL-instruction_following
+#     and place it at ~/Gym/resources_servers/instruction_following/data/instruction_following.jsonl
+#
+# Usage:
+#   axolotl train examples/nemo_gym_multi_env.yaml
+
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+model_type: AutoModelForCausalLM
+
+sequence_len: 4096
+load_in_4bit: false
+
+# RL configuration
+rl: grpo
+chat_template: tokenizer_default
+
+# GRPO / TRL settings
+trl:
+  use_vllm: false
+  num_generations: 8
+  max_completion_length: 2048
+
+# NeMo Gym plugin
+plugins:
+  - axolotl.integrations.nemo_gym.NemoGymPlugin
+
+nemo_gym_enabled: true
+nemo_gym_dir: ~/Gym
+nemo_gym_auto_clone: true
+nemo_gym_auto_start: true
+nemo_gym_config_paths:
+  - resources_servers/reasoning_gym/configs/resources_only.yaml
+  - resources_servers/instruction_following/configs/instruction_following.yaml
+nemo_gym_datasets:
+  - path: resources_servers/reasoning_gym/data/train_mini_sudoku.jsonl
+    server_name: reasoning_gym
+    max_samples: 1000
+  - path: resources_servers/instruction_following/data/instruction_following.jsonl
+    server_name: instruction_following
+    max_samples: 1000
+
+# Training hyperparameters
+learning_rate: 1.0e-5
+weight_decay: 0.001
+lr_scheduler: linear
+warmup_ratio: 0.0
+optimizer: adamw_8bit
+
+num_epochs: 1
+max_steps: 100
+micro_batch_size: 1
+gradient_accumulation_steps: 64
+
+logging_steps: 1
+save_steps: 100
+output_dir: ./outputs/nemo_gym_multi_env
diff --git a/src/axolotl/integrations/nemo_gym/examples/nemo_gym_multi_turn.yaml b/src/axolotl/integrations/nemo_gym/examples/nemo_gym_multi_turn.yaml
new file mode 100644
index 0000000000..2644bbab7f
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/examples/nemo_gym_multi_turn.yaml
@@ -0,0 +1,79 @@
+# Axolotl + NeMo Gym: Multi-Turn Tool-Use RL Training Example
+#
+# This config trains a model with multi-turn reinforcement learning
+# using NeMo Gym's agentic environments. The model generates responses,
+# executes tool calls against NeMo Gym resource servers, and receives
+# environment feedback across multiple turns.
+#
+# Multi-turn mode uses TRL's rollout_func with env_mask to ensure only
+# model-generated tokens contribute to the training loss.
+#
+# Prerequisites:
+#   - uv (https://github.com/astral-sh/uv) installed
+#   - vLLM must be enabled (required for generate_rollout_completions)
+#
+# Usage:
+#   # Terminal 1: Start vLLM server
+#   trl vllm-serve --model Qwen/Qwen2.5-1.5B-Instruct
+#
+#   # Terminal 2: Train
+#   axolotl train examples/nemo_gym_multi_turn.yaml
+
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+model_type: AutoModelForCausalLM
+
+sequence_len: 4096
+load_in_4bit: false
+
+# RL configuration
+rl: grpo
+chat_template: tokenizer_default
+
+# GRPO / TRL settings — vLLM required for multi-turn
+trl:
+  use_vllm: true
+  vllm_mode: server
+  vllm_server_host: localhost
+  vllm_server_port: 8000
+  num_generations: 4
+  max_completion_length: 2048
+
+# NeMo Gym plugin with multi-turn enabled
+plugins:
+  - axolotl.integrations.nemo_gym.NemoGymPlugin
+
+nemo_gym_enabled: true
+nemo_gym_dir: ~/Gym
+nemo_gym_auto_clone: true
+nemo_gym_auto_start: true
+nemo_gym_head_port: 11000
+nemo_gym_server_timeout: 360
+nemo_gym_verify_timeout: 30
+
+# Multi-turn settings
+nemo_gym_multi_turn: true
+nemo_gym_max_turns: 10
+
+# Resource server configs — use environments that support tool calls
+nemo_gym_config_paths:
+  - resources_servers/reasoning_gym/configs/resources_only.yaml
+nemo_gym_datasets:
+  - path: resources_servers/reasoning_gym/data/train_mini_sudoku.jsonl
+    server_name: reasoning_gym
+    max_samples: 2000
+
+# Training hyperparameters
+learning_rate: 1.0e-5
+weight_decay: 0.001
+lr_scheduler: linear
+warmup_ratio: 0.0
+optimizer: adamw_8bit
+
+num_epochs: 1
+max_steps: 100
+micro_batch_size: 1
+gradient_accumulation_steps: 64
+
+logging_steps: 1
+save_steps: 100
+output_dir: ./outputs/nemo_gym_multi_turn
diff --git a/src/axolotl/integrations/nemo_gym/examples/nemo_gym_sudoku.yaml b/src/axolotl/integrations/nemo_gym/examples/nemo_gym_sudoku.yaml
new file mode 100644
index 0000000000..df04b2b415
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/examples/nemo_gym_sudoku.yaml
@@ -0,0 +1,62 @@
+# Axolotl + NeMo Gym: Sudoku RL Training Example
+#
+# This config trains a model to solve 4x4 Mini Sudoku puzzles using GRPO
+# with NeMo Gym's reasoning_gym environment as the reward signal.
+#
+# Prerequisites:
+#   - uv (https://github.com/astral-sh/uv) installed
+#   - Git installed
+#   - The NeMo Gym repo will be auto-cloned to ~/Gym on first run
+#
+# Usage:
+#   axolotl train examples/nemo_gym_sudoku.yaml
+
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+model_type: AutoModelForCausalLM
+
+sequence_len: 4096
+load_in_4bit: false
+
+# RL configuration
+rl: grpo
+chat_template: tokenizer_default
+
+# GRPO / TRL settings
+trl:
+  use_vllm: false
+  num_generations: 8
+  max_completion_length: 2048
+
+# NeMo Gym plugin
+plugins:
+  - axolotl.integrations.nemo_gym.NemoGymPlugin
+
+nemo_gym_enabled: true
+nemo_gym_dir: ~/Gym
+nemo_gym_auto_clone: true
+nemo_gym_auto_start: true
+nemo_gym_head_port: 11000
+nemo_gym_server_timeout: 360
+nemo_gym_verify_timeout: 30
+nemo_gym_config_paths:
+  - resources_servers/reasoning_gym/configs/resources_only.yaml
+nemo_gym_datasets:
+  - path: resources_servers/reasoning_gym/data/train_mini_sudoku.jsonl
+    server_name: reasoning_gym
+    max_samples: 2000
+
+# Training hyperparameters
+learning_rate: 1.0e-5
+weight_decay: 0.001
+lr_scheduler: linear
+warmup_ratio: 0.0
+optimizer: adamw_8bit
+
+num_epochs: 1
+max_steps: 100
+micro_batch_size: 1
+gradient_accumulation_steps: 64
+
+logging_steps: 1
+save_steps: 100
+output_dir: ./outputs/nemo_gym_sudoku
diff --git a/src/axolotl/integrations/nemo_gym/multi_turn.py b/src/axolotl/integrations/nemo_gym/multi_turn.py
new file mode 100644
index 0000000000..328a393bc3
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/multi_turn.py
@@ -0,0 +1,329 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+
+"""
+Multi-turn rollout function for NeMo Gym environments.
+
+Delegates multi-turn orchestration to NeMo Gym's agent servers via the /run
+endpoint. The agent handles generation (by calling our vLLM server), tool
+execution, session management, and reward computation.
+
+This follows the same pattern as TRL's reference implementation at
+examples/scripts/nemo_gym/train_multi_environment.py.
+
+Architecture:
+  rollout_func(prompts, trainer)
+    -> expand prompts by num_generations
+    -> async POST /run to agent servers (one per sample)
+    -> parse response: prompt_ids, completion_ids, logprobs, env_mask, reward
+    -> return to TRL for GRPO training
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def create_nemo_gym_rollout_func(
+    agent_servers: dict[str, str],
+    dataset_lookup: dict[int, dict],
+    request_timeout: float = 10800,
+):
+    """Create a TRL-compatible rollout_func that delegates to NeMo Gym agents.
+
+    Args:
+        agent_servers: Mapping of agent_name → agent URL (e.g., {"simple_agent": "http://host:port"}).
+        dataset_lookup: Mapping of dataset index → full JSONL row dict.
+        request_timeout: HTTP timeout for /run requests.
+
+    Returns:
+        A rollout_func with signature (prompts: list[str], trainer) -> dict.
+    """
+
+    def rollout_func(prompts: list[str], trainer) -> dict[str, Any]:
+        is_training = trainer.model.training
+        num_generations = (
+            trainer.num_generations
+            if is_training
+            else getattr(trainer, "num_generations_eval", 1)
+        )
+        temperature = trainer.temperature
+        top_p = getattr(trainer, "top_p", None) or 0.999
+        max_completion_length = trainer.max_completion_length
+        eos_token_id = trainer.processing_class.eos_token_id
+
+        # Expand prompts: each prompt index repeated num_generations times
+        expanded_items = []
+        expanded_prompt_indices = []
+        for prompt_str in prompts:
+            # Prompts from TRL are chat-templated strings. Find the dataset item
+            # by matching against dataset_lookup keys (raw user message text).
+            full_item = None
+            for key, val in dataset_lookup.items():
+                if isinstance(key, str) and prompt_str == key:
+                    full_item = val
+                    break
+
+            if full_item is None:
+                full_item = {
+                    "responses_create_params": {
+                        "input": [{"role": "user", "content": prompt_str}]
+                    }
+                }
+
+            for _ in range(num_generations):
+                # Preserve agent_ref for routing in _call_agents
+                dispatched: dict = full_item.get("verify_extra", full_item)  # type: ignore[assignment]
+                if isinstance(dispatched, dict) and "agent_ref" not in dispatched:
+                    agent_ref = full_item.get("agent_ref")
+                    if agent_ref:
+                        dispatched = {**dispatched, "agent_ref": agent_ref}
+                expanded_items.append(dispatched)
+                expanded_prompt_indices.append(prompt_str)
+
+        # Call NeMo Gym agents
+        loop = asyncio.new_event_loop()
+        try:
+            responses = loop.run_until_complete(
+                _call_agents(
+                    dataset_items=expanded_items,
+                    agent_servers=agent_servers,
+                    timeout=request_timeout,
+                    max_completion_length=max_completion_length,
+                    temperature=temperature,
+                    top_p=top_p,
+                )
+            )
+        finally:
+            loop.close()
+
+        # Parse responses into rollout format
+        all_prompt_ids = []
+        all_completion_ids = []
+        all_env_masks = []
+        all_logprobs = []
+        all_rewards = []
+        all_num_turns = []
+
+        for _i, response in enumerate(responses):
+            result = _parse_agent_response(response, eos_token_id)
+            all_prompt_ids.append(result["prompt_ids"])
+            all_completion_ids.append(result["completion_ids"])
+            all_env_masks.append(result["env_mask"])
+            all_logprobs.append(result["logprobs"])
+            all_rewards.append(result["reward"])
+            all_num_turns.append(result["num_turns"])
+
+        # TRL expects prompt_ids to be unique (one per original prompt, not per generation)
+        unique_prompt_ids = all_prompt_ids[::num_generations]
+
+        # Wrap logprobs for TRL: list[list[list[float]]]
+        def _normalize(lp):
+            while isinstance(lp, (list, tuple)) and len(lp) > 0:
+                lp = lp[0]
+            return float(lp) if lp is not None else 0.0
+
+        wrapped_logprobs = [[[_normalize(lp)] for lp in seq] for seq in all_logprobs]
+
+        return {
+            "prompt_ids": unique_prompt_ids,
+            "completion_ids": all_completion_ids,
+            "env_mask": all_env_masks,
+            "logprobs": wrapped_logprobs,
+            "logprob_token_ids": None,  # nosec B105
+            "env_reward": all_rewards,
+            "num_turns": all_num_turns,
+        }
+
+    return rollout_func
+
+
+async def _call_agents(
+    dataset_items: list[dict],
+    agent_servers: dict[str, str],
+    timeout: float,
+    max_completion_length: int = 4096,
+    temperature: float = 1.0,
+    top_p: float = 0.999,
+) -> list[dict]:
+    """Async batch POST to NeMo Gym agent /run endpoints."""
+    import aiohttp
+
+    results = []
+    connector = aiohttp.TCPConnector(limit_per_host=64, limit=256)
+    # Use sock_read for per-request timeout (not total session timeout).
+    # This ensures a single stuck generation doesn't block all other requests.
+    client_timeout = aiohttp.ClientTimeout(total=None, sock_read=timeout)
+
+    async with aiohttp.ClientSession(
+        connector=connector, timeout=client_timeout, cookie_jar=aiohttp.DummyCookieJar()
+    ) as session:
+        tasks = []
+        for item in dataset_items:
+            agent_ref = item.get("agent_ref", {})
+            agent_name = agent_ref.get("name", "")
+            agent_url = agent_servers.get(agent_name, "")
+
+            if not agent_url:
+                # Fallback: try first available agent
+                if agent_servers:
+                    agent_url = next(iter(agent_servers.values()))
+                else:
+                    results.append(
+                        {
+                            "response": {"output": []},
+                            "reward": 0.0,
+                            "error": "No agent server",
+                        }
+                    )
+                    continue
+
+            # Build request body
+            request_body = dict(item)
+            params = request_body.setdefault("responses_create_params", {})
+            params.setdefault("max_output_tokens", max_completion_length)
+            params["temperature"] = temperature
+            params["top_p"] = top_p
+
+            tasks.append(_post_run(session, agent_url, request_body))
+
+        if tasks:
+            responses = await asyncio.gather(*tasks, return_exceptions=True)
+            for resp in responses:
+                if isinstance(resp, BaseException):
+                    LOG.warning(f"Agent /run failed: {resp}")
+                    results.append(
+                        {"response": {"output": []}, "reward": 0.0, "error": str(resp)}
+                    )
+                else:
+                    results.append(resp)
+
+    return results
+
+
+async def _post_run(session, agent_url: str, body: dict) -> dict:
+    """POST to agent /run endpoint."""
+    async with session.post(f"{agent_url}/run", json=body) as resp:
+        if resp.status == 200:
+            return await resp.json()
+        text = await resp.text()
+        return {
+            "response": {"output": []},
+            "reward": 0.0,
+            "error": f"HTTP {resp.status}: {text[:200]}",
+        }
+
+
+def _parse_agent_response(response: dict, eos_token_id: int) -> dict:
+    """Parse NeMo Gym agent /run response into rollout format.
+
+    The agent returns:
+      response.output[]: list of turns, each with prompt_token_ids,
+        generation_token_ids, generation_log_probs
+      reward: float
+    """
+    # Defaults for failed/empty responses
+    defaults = {
+        "prompt_ids": [eos_token_id],
+        "completion_ids": [eos_token_id],
+        "env_mask": [0],
+        "logprobs": [0.0],
+        "reward": 0.0,
+        "num_turns": 0,
+    }
+
+    if not isinstance(response, dict) or "error" in response:
+        return defaults
+
+    output_items = response.get("response", {}).get("output", [])
+    reward = float(response.get("reward", 0.0))
+
+    if not output_items:
+        defaults["reward"] = reward
+        return defaults
+
+    # Check at least one valid output
+    has_valid = False
+    for item in output_items:
+        if item.get("type") == "function_call":
+            has_valid = True
+            break
+        if item.get("type") == "message":
+            for c in item.get("content", []):
+                if (
+                    isinstance(c, dict)
+                    and c.get("type") == "output_text"
+                    and c.get("text", "").strip()
+                ):
+                    has_valid = True
+                    break
+        if has_valid:
+            break
+
+    if not has_valid:
+        defaults["reward"] = reward
+        return defaults
+
+    # Extract multi-turn token sequences
+    first_prompt_ids = None
+    completion_ids = []
+    env_mask = []
+    logprobs = []
+    seen_token_ids = []
+    num_turns = 0
+
+    for item in output_items:
+        prompt_token_ids = item.get("prompt_token_ids", [])
+        generation_token_ids = item.get("generation_token_ids", [])
+        generation_log_probs = item.get("generation_log_probs", [])
+
+        if not generation_token_ids:
+            continue
+
+        num_turns += 1
+
+        # First turn: capture prompt
+        if first_prompt_ids is None:
+            first_prompt_ids = list(prompt_token_ids)
+            seen_token_ids = list(prompt_token_ids)
+        else:
+            # Subsequent turns: extract tool result tokens (between turns)
+            if len(prompt_token_ids) > len(seen_token_ids):
+                tool_result_tokens = prompt_token_ids[len(seen_token_ids) :]
+                # Tool result tokens are NOT trained on (env_mask = 0)
+                completion_ids.extend(tool_result_tokens)
+                env_mask.extend([0] * len(tool_result_tokens))
+                logprobs.extend([0.0] * len(tool_result_tokens))
+
+        # Add generation tokens (trained on, env_mask = 1)
+        completion_ids.extend(generation_token_ids)
+        env_mask.extend([1] * len(generation_token_ids))
+
+        # Pad logprobs if shorter than generation tokens
+        gen_logprobs = list(generation_log_probs) if generation_log_probs else []
+        if len(gen_logprobs) < len(generation_token_ids):
+            gen_logprobs.extend([0.0] * (len(generation_token_ids) - len(gen_logprobs)))
+        logprobs.extend(gen_logprobs[: len(generation_token_ids)])
+
+        # Update seen tokens
+        seen_token_ids = list(prompt_token_ids) + list(generation_token_ids)
+
+    if first_prompt_ids is None:
+        return defaults
+
+    return {
+        "prompt_ids": first_prompt_ids,
+        "completion_ids": completion_ids,
+        "env_mask": env_mask,
+        "logprobs": logprobs,
+        "reward": reward,
+        "num_turns": num_turns,
+    }
diff --git a/src/axolotl/integrations/nemo_gym/plugin.py b/src/axolotl/integrations/nemo_gym/plugin.py
new file mode 100644
index 0000000000..b85e344db2
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/plugin.py
@@ -0,0 +1,687 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+
+"""
+NeMo Gym Plugin for Axolotl.
+
+Integrates NVIDIA NeMo Gym environments as reward sources for GRPO training.
+Handles server lifecycle, dataset loading, and reward function wiring.
+
+Supports two modes:
+  - Single-turn (default): reward_fn calls /verify after each generation
+  - Multi-turn (nemo_gym_multi_turn: true): rollout_func orchestrates
+    multi-step interactions with tool execution via resource servers
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Union
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from axolotl.common.datasets import TrainDatasetMeta
+
+LOG = get_logger(__name__)
+
+
+# ---- vLLM weight-sync transport probe ------------------------------------
+
+
+@dataclass
+class VLLMWeightSyncCapabilities:
+    """What weight-sync routes a vLLM server actually exposes.
+
+    Discovered once at ``pre_model_load`` time by fetching the server's
+    ``/openapi.json``. Drives the transport-selection table below.
+    """
+
+    nccl: bool = False  # /init_communicator/ + /update_named_param/
+    lora_filesystem: bool = False  # /v1/load_lora_adapter (vLLM native)
+    lora_axolotl: bool = False  # /set_lora_adapter/ (axolotl serve_lora extension)
+    http_full: bool = False  # /http_update_weights/ (axolotl serve_lora extension)
+    probed: bool = False
+    probe_error: str | None = None
+    routes: list[str] = field(default_factory=list)
+
+    @property
+    def any_full_param_sync(self) -> bool:
+        """True if at least one transport can push full-model weights."""
+        return self.nccl or self.http_full
+
+    @property
+    def any_lora_sync(self) -> bool:
+        """True if at least one transport can push LoRA adapters."""
+        return self.lora_filesystem or self.lora_axolotl or self.nccl
+
+
+def probe_vllm_weight_sync(
+    base_url: str, timeout: float = 5.0
+) -> VLLMWeightSyncCapabilities:
+    """Detect which weight-sync routes the configured vLLM server exposes.
+
+    Uses the server's FastAPI ``/openapi.json`` — every weight-sync transport
+    we care about is mounted as a POST route there. Falls back to all-False
+    on any error so the caller can still decide what to do (typically: raise
+    a clear error rather than silently no-op).
+    """
+    import requests
+
+    caps = VLLMWeightSyncCapabilities()
+    try:
+        r = requests.get(f"{base_url.rstrip('/')}/openapi.json", timeout=timeout)
+        r.raise_for_status()
+        spec = r.json()
+        routes = sorted((spec.get("paths") or {}).keys())
+        caps.routes = routes
+        caps.nccl = "/init_communicator/" in routes and "/update_named_param/" in routes
+        caps.lora_filesystem = "/v1/load_lora_adapter" in routes
+        caps.lora_axolotl = "/set_lora_adapter/" in routes
+        caps.http_full = "/http_update_weights/" in routes
+        caps.probed = True
+    except Exception as exc:
+        caps.probe_error = f"{type(exc).__name__}: {exc}"
+        LOG.warning(
+            "NeMo Gym: failed to probe vLLM /openapi.json at %s — %s. "
+            "Will fall back to LoRA-only behavior.",
+            base_url,
+            caps.probe_error,
+        )
+    return caps
+
+
+def select_weight_sync_transport(
+    caps: VLLMWeightSyncCapabilities,
+    *,
+    has_lora: bool,
+    vllm_lora_sync_pref: bool,
+) -> str:
+    """Pick the right transport for a (server caps, model type) combo.
+
+    Returns one of: ``"lora_filesystem"``, ``"nccl"``, ``"http_full"``, or
+    ``"none"``. The caller decides what to do with ``"none"`` (typically:
+    raise an error explaining the misconfiguration).
+
+    Selection table:
+        LoRA model + lora endpoint + lora-sync pref    → lora_filesystem
+        LoRA model + lora endpoint                     → lora_filesystem
+        LoRA model + nccl endpoint                     → nccl (broadcast merged adapter)
+        Full model + nccl endpoint                     → nccl
+        Full model + http endpoint                     → http_full
+        anything else                                  → none
+    """
+    if has_lora:
+        if (caps.lora_filesystem or caps.lora_axolotl) and vllm_lora_sync_pref:
+            return "lora_filesystem"
+        if caps.lora_filesystem or caps.lora_axolotl:
+            return "lora_filesystem"
+        if caps.nccl:
+            return "nccl"
+        return "none"
+    # Full-parameter model
+    if caps.nccl:
+        return "nccl"
+    if caps.http_full:
+        return "http_full"
+    return "none"
+
+
+class NemoGymPlugin(BasePlugin):
+    """Plugin for NVIDIA NeMo Gym integration with Axolotl.
+
+    When enabled, this plugin:
+    1. Clones and sets up the NeMo Gym repo (if needed)
+    2. Starts NeMo Gym resource servers
+    3. Loads datasets from NeMo Gym JSONL files
+    4. For single-turn: creates a reward function calling /verify
+    5. For multi-turn: creates a rollout_func with tool execution and env_mask
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._gym_dir = None
+        self._global_config = None
+        self._verify_endpoints = None
+        self._server_base_urls = None
+        self._reward_fn = None
+        self._dataset_lookup = None
+        self._agent_servers = {}
+        self._vllm_caps: VLLMWeightSyncCapabilities | None = None
+
+    def get_input_args(self):
+        return "axolotl.integrations.nemo_gym.NemoGymArgs"
+
+    def pre_model_load(self, cfg):
+        """Probe vLLM weight-sync routes and conditionally bypass NCCL init.
+
+        Replaces the previous unconditional ``init_communicator`` monkey-patch
+        with a probe of the configured vLLM server's ``/openapi.json``. We only
+        bypass NCCL init when the server we're talking to actually lacks the
+        ``/init_communicator/`` route (i.e. stock ``vllm serve``); against
+        TRL/axolotl serve modules that DO expose NCCL routes, we leave the
+        standard TRL flow alone so full-finetune training can sync weights.
+        """
+        if not cfg.nemo_gym_enabled:
+            return
+
+        trl_cfg = getattr(cfg, "trl", None)
+        if not (trl_cfg and getattr(trl_cfg, "vllm_mode", "server") == "server"):
+            return
+
+        host = getattr(trl_cfg, "vllm_server_host", None) or "127.0.0.1"
+        port = getattr(trl_cfg, "vllm_server_port", None) or 8000
+        base_url = f"http://{host}:{port}"
+        self._vllm_caps = probe_vllm_weight_sync(base_url)
+
+        if self._vllm_caps.probed:
+            LOG.info(
+                "NeMo Gym: vLLM weight-sync probe @ %s — nccl=%s lora_native=%s "
+                "lora_axolotl=%s http_full=%s",
+                base_url,
+                self._vllm_caps.nccl,
+                self._vllm_caps.lora_filesystem,
+                self._vllm_caps.lora_axolotl,
+                self._vllm_caps.http_full,
+            )
+
+        # Only bypass NCCL init when the server doesn't speak it. If NCCL is
+        # available we leave VLLMClient.init_communicator alone so the
+        # standard TRL sync flow can run for full-parameter training.
+        if not self._vllm_caps.nccl:
+            self._patch_skip_nccl_init()
+
+    def _patch_skip_nccl_init(self):
+        """Monkeypatch VLLMClient.init_communicator to no-op.
+
+        Only called when the configured vLLM server doesn't expose
+        ``/init_communicator/`` (e.g. stock ``vllm serve``). In that case
+        TRL's standard ``init_communicator`` would 404 inside trainer
+        construction; we no-op it so the LoRA filesystem path can install
+        its own sync in ``post_trainer_create``.
+        """
+        try:
+            from trl.generation.vllm_client import VLLMClient
+
+            VLLMClient._original_init_communicator = VLLMClient.init_communicator
+            VLLMClient.init_communicator = lambda self, **kwargs: LOG.info(
+                "Skipping NCCL init_communicator (server has no /init_communicator/)"
+            )
+            LOG.info(
+                "Patched VLLMClient.init_communicator to no-op (server has no NCCL routes)"
+            )
+        except Exception as exc:
+            LOG.warning(f"Failed to patch VLLMClient: {exc}")
+
+    def register(self, cfg):
+        if not cfg.get("nemo_gym_enabled"):
+            return
+
+        LOG.info("NeMo Gym integration enabled")
+        gym_dir = cfg.get("nemo_gym_dir") or os.path.expanduser("~/Gym")
+        auto_clone = cfg.get("nemo_gym_auto_clone", True)
+        auto_start = cfg.get("nemo_gym_auto_start", True)
+        head_port = cfg.get("nemo_gym_head_port", 11000)
+        server_timeout = cfg.get("nemo_gym_server_timeout", 360)
+
+        from .server import (
+            ensure_gym_repo,
+            ensure_gym_venv,
+            get_agent_servers,
+            get_server_base_url,
+            get_server_configs,
+            get_verify_endpoint,
+            start_servers,
+            wait_for_resource_servers,
+        )
+
+        self._gym_dir = ensure_gym_repo(gym_dir, auto_clone=auto_clone)
+
+        if auto_start:
+            config_paths = cfg.get("nemo_gym_config_paths", [])
+            ensure_gym_venv(self._gym_dir)
+            start_servers(
+                self._gym_dir,
+                config_paths,
+                head_port=head_port,
+                timeout=server_timeout,
+            )
+
+        self._global_config = get_server_configs(head_port=head_port)
+        wait_for_resource_servers(self._global_config, timeout=server_timeout)
+
+        # Build endpoint maps for resource servers (/verify)
+        self._verify_endpoints = {}
+        self._server_base_urls = {}
+        for server_name in self._global_config:
+            try:
+                self._verify_endpoints[server_name] = get_verify_endpoint(
+                    self._global_config, server_name
+                )
+                self._server_base_urls[server_name] = get_server_base_url(
+                    self._global_config, server_name
+                )
+            except (ValueError, KeyError, TypeError):
+                pass
+
+        # Discover agent servers (/run) for multi-turn
+        self._agent_servers = get_agent_servers(self._global_config)
+
+        # Pre-build dataset lookup for multi-turn (needs to happen at register time,
+        # not load_datasets, because load_datasets may not be called if axolotl config
+        # has its own datasets field)
+        if cfg.get("nemo_gym_multi_turn") and cfg.get("nemo_gym_datasets"):
+            from .dataset import load_nemo_gym_datasets
+
+            gym_dir = cfg.get("nemo_gym_dir") or os.path.expanduser("~/Gym")
+            dataset = load_nemo_gym_datasets(gym_dir, cfg["nemo_gym_datasets"])
+            self._dataset_lookup = {}
+            for i in range(len(dataset)):
+                row = dataset[i]
+                # Use last message content as key (matches data_producer lookup)
+                prompt_text = row["prompt"][-1]["content"]
+                self._dataset_lookup[prompt_text] = row
+            LOG.info(f"Built dataset lookup with {len(self._dataset_lookup)} entries")
+
+        multi_turn = cfg.get("nemo_gym_multi_turn", False)
+        LOG.info(
+            f"NeMo Gym ready with servers: {list(self._verify_endpoints.keys())} "
+            f"(multi_turn={'enabled' if multi_turn else 'disabled'})"
+        )
+
+    def load_datasets(self, cfg, preprocess=False) -> Union["TrainDatasetMeta", None]:
+        if not cfg.nemo_gym_enabled:
+            return None
+
+        from axolotl.common.datasets import TrainDatasetMeta
+
+        from .dataset import load_nemo_gym_datasets
+
+        dataset_configs = cfg.nemo_gym_datasets
+        dataset = load_nemo_gym_datasets(self._gym_dir, dataset_configs)
+
+        # Build prompt → row lookup for multi-turn rollout_func
+        # (rollout_func only receives prompt text, needs to look up row data)
+        self._dataset_lookup = {}
+        for i in range(len(dataset)):
+            row = dataset[i]
+            # Use last message content as key (matches data_producer lookup)
+            prompt_text = row["prompt"][-1]["content"]
+            self._dataset_lookup[prompt_text] = row
+
+        return TrainDatasetMeta(
+            train_dataset=dataset,
+            eval_dataset=None,
+            total_num_steps=0,  # computed later by the builder
+        )
+
+    def get_training_args(self, cfg):
+        """Pass through vLLM settings and force async trainer for multi-turn."""
+        args = {}
+        # Pass vLLM settings from vllm config block to TRL training args
+        if cfg.vllm:
+            vllm_cfg = cfg.vllm
+            max_len = getattr(vllm_cfg, "max_model_len", None)
+            gpu_util = getattr(vllm_cfg, "gpu_memory_utilization", None)
+            tp_size = getattr(vllm_cfg, "tensor_parallel_size", None)
+            if max_len:
+                args["vllm_max_model_length"] = max_len
+            if gpu_util:
+                args["vllm_gpu_memory_utilization"] = gpu_util
+            if tp_size:
+                args["vllm_tensor_parallel_size"] = tp_size
+
+        # Force async trainer for multi-turn: NemoGymDataProducer needs the
+        # data producer protocol. Setting use_data_producer=True selects
+        # AxolotlAsyncGRPOTrainer which supports _create_data_producer().
+        # With async_prefetch=False this runs synchronously — no threading.
+        if cfg.nemo_gym_multi_turn and self._agent_servers:
+            args["use_data_producer"] = True
+            LOG.info(
+                "NeMo Gym multi-turn: forcing use_data_producer=True for data producer protocol"
+            )
+
+        # Dataloader workers fork subprocesses that can't handle the async
+        # HTTP connections to NeMo Gym agents. Force num_workers=0.
+        if getattr(cfg, "dataloader_num_workers", None) not in (None, 0):
+            LOG.warning(
+                f"NeMo Gym: overriding dataloader_num_workers={cfg.dataloader_num_workers} → 0 "
+                "(forked workers can't use NeMo Gym agent connections)"
+            )
+        cfg.dataloader_num_workers = 0
+
+        if args:
+            LOG.info(f"NeMo Gym plugin injecting training args: {args}")
+        return args if args else None
+
+    def post_trainer_create(self, cfg, trainer):
+        """Wire NeMo Gym into the trainer (reward_fn or rollout_func)."""
+        if not cfg.nemo_gym_enabled:
+            return
+
+        model_name = cfg.nemo_gym_model_name or cfg.base_model or "axolotl-model"
+        verify_timeout = cfg.nemo_gym_verify_timeout or 30
+        multi_turn = cfg.nemo_gym_multi_turn or False
+
+        # Pick a weight-sync transport based on what the configured vLLM
+        # server actually exposes (see ``pre_model_load`` probe) and what
+        # kind of model we're training. The selection table is documented
+        # in ``select_weight_sync_transport``.
+        trl_cfg = getattr(cfg, "trl", None)
+        if hasattr(trainer, "vllm_generation") and trainer.vllm_generation:
+            vllm_gen = trainer.vllm_generation
+            adapter = getattr(cfg, "adapter", None)
+            has_lora = adapter in ("lora", "qlora")
+            vllm_lora_sync_pref = bool(
+                trl_cfg and getattr(trl_cfg, "vllm_lora_sync", False)
+            )
+            caps = self._vllm_caps or VLLMWeightSyncCapabilities()
+            transport = select_weight_sync_transport(
+                caps,
+                has_lora=has_lora,
+                vllm_lora_sync_pref=vllm_lora_sync_pref,
+            )
+
+            if transport == "lora_filesystem":
+                self._setup_lora_sync(trainer)
+                self._check_lora_endpoint(vllm_gen)
+                LOG.info("NeMo Gym weight sync: LoRA filesystem")
+            elif transport == "nccl":
+                # Standard TRL NCCL path. We leave ``VLLMClient.init_communicator``
+                # alone (pre_model_load only patched it when the probe found no
+                # NCCL route) so the trainer's normal weight-sync flow runs.
+                LOG.info(
+                    "NeMo Gym weight sync: NCCL (server exposes /init_communicator/)"
+                )
+            elif transport == "http_full":
+                # Full-parameter HTTP sync — implementation lands in step 3.
+                # For now, fail loudly so users know the path is detected but
+                # not yet wired up, instead of silently no-oping like before.
+                raise NotImplementedError(
+                    "NeMo Gym + full fine-tune + HTTP weight sync is detected "
+                    "but the client-side sync helper is not yet implemented "
+                    "(planned). Use `adapter: lora|qlora` for now, or use a "
+                    "vLLM serve module that exposes /init_communicator/ for "
+                    "NCCL sync."
+                )
+            else:  # transport == "none"
+                # No viable sync path. Build a precise error so the user knows
+                # exactly what's missing and how to fix it.
+                if not caps.probed:
+                    msg = (
+                        "could not probe the vLLM server's "
+                        f"/openapi.json: {caps.probe_error}. "
+                        "Verify that vLLM is reachable at "
+                        f"{getattr(trl_cfg, 'vllm_server_host', '?')}:"
+                        f"{getattr(trl_cfg, 'vllm_server_port', '?')}."
+                    )
+                elif has_lora:
+                    msg = (
+                        "the vLLM server has neither NCCL routes "
+                        "(/init_communicator/) nor a LoRA-loading route "
+                        "(/v1/load_lora_adapter or /set_lora_adapter/). "
+                        "Restart vLLM with `--enable-lora --max-lora-rank N "
+                        "VLLM_ALLOW_RUNTIME_LORA_UPDATING=1` for the stock "
+                        "server, or use `axolotl vllm-serve` for the "
+                        "NCCL-capable serve module."
+                    )
+                else:
+                    msg = (
+                        "the vLLM server exposes no full-parameter sync route "
+                        "(/init_communicator/ for NCCL or /http_update_weights/ "
+                        "for HTTP). Use `axolotl vllm-serve` (which has both) "
+                        "or set `adapter: lora|qlora`."
+                    )
+                raise ValueError(
+                    f"NeMo Gym: no usable weight-sync transport — {msg} Without "
+                    "weight sync the trainer's gradient updates never reach the "
+                    "rollout policy (functionally a no-op trainer)."
+                )
+
+        if multi_turn:
+            self._wire_multi_turn(cfg, trainer, model_name, verify_timeout)
+        else:
+            self._wire_single_turn(trainer, model_name, verify_timeout)
+
+    def _wire_single_turn(self, trainer, model_name, verify_timeout):
+        """Inject single-turn reward function into the trainer."""
+        from .rewards import create_nemo_gym_reward_fn
+
+        self._reward_fn = create_nemo_gym_reward_fn(
+            global_config=self._global_config,
+            verify_endpoints=self._verify_endpoints,
+            model_name=model_name,
+            verify_timeout=verify_timeout,
+        )
+
+        if hasattr(trainer, "reward_funcs"):
+            trainer.reward_funcs.append(self._reward_fn)
+            trainer.reward_func_names.append("nemo_gym")
+            trainer.reward_processing_classes.append(None)
+            LOG.info(
+                f"Added NeMo Gym reward function (single-turn). "
+                f"Total reward functions: {len(trainer.reward_funcs)}"
+            )
+        else:
+            LOG.warning(
+                "Trainer does not have reward_funcs attribute. "
+                "NeMo Gym reward function not injected. "
+                "Ensure you are using a GRPO trainer."
+            )
+
+    def _wire_multi_turn(self, cfg, trainer, model_name, verify_timeout):
+        """Replace the data producer with NemoGymDataProducer.
+
+        The plugin forces use_data_producer=True (in get_training_args) which
+        selects AxolotlAsyncGRPOTrainer. Here we swap its data_producer with
+        our NemoGymDataProducer that calls agent /run instead of vLLM generate.
+        """
+        if not self._agent_servers:
+            LOG.warning(
+                "No NeMo Gym agent servers discovered. Multi-turn requires agent servers "
+                "started via ng_run with an agent config. Falling back to single-turn."
+            )
+            self._wire_single_turn(trainer, model_name, verify_timeout)
+            return
+
+        if not hasattr(trainer, "data_producer") or trainer.data_producer is None:
+            LOG.warning(
+                "Trainer has no data_producer. NeMo Gym multi-turn requires "
+                "use_data_producer=true (should be auto-set by plugin)."
+            )
+            return
+
+        from axolotl.core.trainers.grpo.async_trainer import AsyncDataProducer
+
+        from .data_producer import NemoGymDataProducer
+
+        # Get the current producer's config and params
+        current = trainer.data_producer
+        # Unwrap AsyncDataProducer to get the inner producer's config
+        if isinstance(current, AsyncDataProducer):
+            inner = current._inner
+        else:
+            inner = current
+
+        nemo_producer = NemoGymDataProducer(
+            config=inner.config,
+            prompt_dataset=inner._dataset,
+            num_generations=inner._num_generations,
+            generation_batch_size=inner._generation_batch_size,
+            train_batch_size=inner._train_batch_size,
+            steps_per_generation=inner._steps_per_generation,
+            shuffle_dataset=inner._shuffle_dataset,
+            seed=inner._seed,
+            agent_servers=self._agent_servers,
+            dataset_lookup=self._dataset_lookup or {},
+            request_timeout=float(cfg.nemo_gym_run_timeout or 300),
+        )
+        nemo_producer.set_trainer(trainer)
+
+        # Re-wrap in AsyncDataProducer if async prefetch is enabled
+        if getattr(trainer.args, "async_prefetch", False):
+            nemo_producer = AsyncDataProducer(
+                nemo_producer,
+                background_produce_kwargs={"skip_policy_logps": True},
+            )
+
+        trainer.data_producer = nemo_producer
+        LOG.info(
+            f"NeMo Gym data producer installed "
+            f"(agent servers: {list(self._agent_servers.keys())}, "
+            f"async={'yes' if getattr(trainer.args, 'async_prefetch', False) else 'no'})"
+        )
+
+        # Passthrough reward function — agent /run already computed rewards
+        from .rewards import reward_env
+
+        if hasattr(trainer, "reward_funcs"):
+            trainer.reward_funcs.append(reward_env)
+            trainer.reward_func_names.append("nemo_gym")
+            trainer.reward_processing_classes.append(None)
+
+    @staticmethod
+    def _check_lora_endpoint(vllm_gen):
+        """Verify the vLLM server supports runtime LoRA loading."""
+        import requests as http_requests
+
+        if not hasattr(vllm_gen, "vllm_client") or vllm_gen.vllm_client is None:
+            return  # Non-main rank in multi-GPU — client only exists on rank 0
+        base_url = vllm_gen.vllm_client.base_url
+        try:
+            # Send a dummy load request — if the endpoint exists, we get a
+            # proper error (400/404 about the adapter), not a route 404.
+            resp = http_requests.post(
+                f"{base_url}/v1/load_lora_adapter",
+                json={"lora_name": "__probe__", "lora_path": "/nonexistent"},
+                timeout=5,
+            )
+            if (
+                resp.status_code == 404
+                and "Not Found" in resp.text
+                and "adapter" not in resp.text.lower()
+            ):
+                LOG.warning(
+                    "vLLM server does not expose /v1/load_lora_adapter. "
+                    "Set VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 when starting vLLM, e.g.:\n"
+                    "  VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 python -m vllm.entrypoints.openai.api_server "
+                    "--enable-lora --max-lora-rank 64 ..."
+                )
+        except Exception:
+            pass  # Server might not be up yet, sync will warn later
+
+    def _setup_lora_sync(self, trainer):
+        """Replace sync_weights with LoRA adapter sync via filesystem + HTTP.
+
+        If the async trainer is detected (has ``_sync_lora_adapter``), delegates
+        to it — that method already handles multi-GPU (FSDP/DeepSpeed state_dict
+        gather, broadcast sync dir, barrier).
+
+        Otherwise installs a standalone closure for the non-async GRPO path that
+        saves the adapter and POSTs to ``/v1/load_lora_adapter``.
+        """
+        vllm_gen = trainer.vllm_generation
+
+        # Async trainer path: delegate to its _sync_lora_adapter (multi-GPU safe)
+        if hasattr(trainer, "_sync_lora_adapter"):
+
+            def lora_sync_weights():
+                trainer._sync_lora_adapter()
+
+            vllm_gen.sync_weights = lora_sync_weights
+            type(vllm_gen).sync_weights = lambda self: lora_sync_weights()
+            LOG.info(
+                "Installed LoRA adapter sync "
+                "(delegates to async trainer._sync_lora_adapter)"
+            )
+            return
+
+        # Non-async standard GRPO path: standalone closure
+        import os
+        import shutil
+        import tempfile
+
+        import requests as http_requests
+
+        base_model = getattr(trainer.args, "model_name_or_path", None) or "axolotl-lora"
+        sync_state = {"version": 0, "sync_dir": tempfile.mkdtemp(prefix="lora_sync_")}
+
+        def lora_sync_weights():
+            """Save LoRA adapter and load it into vLLM."""
+            accelerator = vllm_gen.accelerator
+            model = vllm_gen.model
+
+            if vllm_gen.mode != "server":
+                return
+
+            sync_state["version"] += 1
+            version = sync_state["version"]
+            adapter_path = os.path.join(sync_state["sync_dir"], f"v{version}")
+
+            wrapped_model = getattr(trainer, "model_wrapped", model)
+            state_dict = accelerator.get_state_dict(wrapped_model)
+
+            if accelerator.is_main_process:
+                unwrapped = accelerator.unwrap_model(model)
+                unwrapped.save_pretrained(adapter_path, state_dict=state_dict)
+
+                base_url = vllm_gen.vllm_client.base_url
+                resp = http_requests.post(
+                    f"{base_url}/v1/load_lora_adapter",
+                    json={
+                        "lora_name": base_model,
+                        "lora_path": adapter_path,
+                        "load_inplace": True,
+                    },
+                    timeout=30,
+                )
+                if resp.status_code != 200:
+                    resp = http_requests.post(
+                        f"{base_url}/set_lora_adapter/",
+                        json={
+                            "lora_name": "active_lora",
+                            "lora_int_id": version,
+                            "lora_path": adapter_path,
+                        },
+                        timeout=30,
+                    )
+                    if resp.status_code != 200:
+                        LOG.warning(
+                            f"Failed to set LoRA adapter: "
+                            f"{resp.status_code} {resp.text}"
+                        )
+                        return
+
+                try:
+                    vllm_gen.vllm_client.reset_prefix_cache()
+                except Exception as exc:
+                    LOG.warning("Failed to reset prefix cache: %s", exc)
+
+                if version > 1:
+                    old = os.path.join(sync_state["sync_dir"], f"v{version - 1}")
+                    if os.path.exists(old):
+                        shutil.rmtree(old, ignore_errors=True)
+
+                LOG.info(f"Synced LoRA adapter v{version} to vLLM ({adapter_path})")
+
+            if accelerator.num_processes > 1:
+                import torch.distributed as dist
+
+                if dist.is_initialized():
+                    dist.barrier()
+
+        vllm_gen.sync_weights = lora_sync_weights
+        type(vllm_gen).sync_weights = lambda self: lora_sync_weights()
+        LOG.info("Installed LoRA adapter sync (standalone fallback)")
+
+    def post_train_unload(self, cfg):
+        """Cleanup NeMo Gym servers if we started them."""
+        if cfg.get("nemo_gym_enabled") and cfg.get("nemo_gym_auto_start", True):
+            from .server import _cleanup_servers
+
+            _cleanup_servers()
diff --git a/src/axolotl/integrations/nemo_gym/rewards.py b/src/axolotl/integrations/nemo_gym/rewards.py
new file mode 100644
index 0000000000..e99df09818
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/rewards.py
@@ -0,0 +1,274 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+
+"""
+NeMo Gym reward functions.
+
+Provides ready-to-use reward functions for axolotl configs::
+
+    trl:
+      reward_funcs:
+        # Multi-turn: passthrough reward from agent /run
+        - axolotl.integrations.nemo_gym.rewards.reward_env
+        # Single-turn: call /verify endpoints directly
+        - axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import requests
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Multi-turn passthrough reward
+# ---------------------------------------------------------------------------
+
+
+def reward_env(completions, prompts=None, **kwargs):
+    """Passthrough: extract pre-computed reward from NeMo Gym agent /run response.
+
+    The ``NemoGymDataProducer`` injects ``env_reward`` into each sample's
+    kwargs after the agent returns from ``/run``.  This function simply
+    forwards that value so TRL can log it alongside other reward signals.
+
+    Use this in your config when ``nemo_gym_multi_turn: true``::
+
+        trl:
+          reward_funcs:
+            - axolotl.integrations.nemo_gym.rewards.reward_env
+    """
+    env_rewards = kwargs.get("env_reward")
+    if env_rewards is not None:
+        if isinstance(env_rewards, (list, tuple)):
+            return [float(r) for r in env_rewards]
+        return [float(env_rewards) for _ in completions]
+    return [0.0 for _ in completions]
+
+
+# ---------------------------------------------------------------------------
+# Single-turn /verify reward
+# ---------------------------------------------------------------------------
+
+# Module-level cache for discovered verify URLs
+_verify_urls: dict[str, str] = {}
+_verify_urls_lock = __import__("threading").Lock()
+
+
+def _get_verify_urls(head_port: int = 11000) -> dict[str, str]:
+    """Discover verify endpoints from the NeMo Gym head server.
+
+    Results are cached so that the HTTP round-trip only happens once per
+    process.  A lock guards against concurrent discovery from multiple
+    threads (e.g. async_prefetch background thread + main training thread).
+    """
+    global _verify_urls
+    if _verify_urls:
+        return _verify_urls
+
+    with _verify_urls_lock:
+        # Double-check after acquiring lock
+        if _verify_urls:
+            return _verify_urls
+
+        import yaml
+
+        try:
+            resp = requests.get(
+                f"http://127.0.0.1:{head_port}/global_config_dict_yaml", timeout=5
+            )
+            config = yaml.safe_load(resp.text)
+            if isinstance(config, str):
+                config = yaml.safe_load(config)
+            for _name, cfg in config.items():
+                if not isinstance(cfg, dict):
+                    continue
+                for srv_name, srv_cfg in cfg.get("resources_servers", {}).items():
+                    if (
+                        isinstance(srv_cfg, dict)
+                        and "host" in srv_cfg
+                        and "port" in srv_cfg
+                    ):
+                        _verify_urls[srv_name] = (
+                            f"http://{srv_cfg['host']}:{srv_cfg['port']}/verify"
+                        )
+        except Exception as exc:
+            LOG.warning(f"Failed to discover NeMo Gym verify endpoints: {exc}")
+
+    return _verify_urls
+
+
+def reward_nemo_gym_verify(completions, prompts=None, **kwargs):
+    """Call NeMo Gym ``/verify`` endpoint for each completion (single-turn).
+
+    Requires ``resources_server_ref`` and ``verify_extra`` kwargs, which the
+    NeMo Gym dataset loader injects automatically.
+
+    Use this in your config when ``nemo_gym_multi_turn: false``::
+
+        trl:
+          reward_funcs:
+            - axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify
+    """
+    verify_urls = _get_verify_urls()
+    refs = kwargs.get("resources_server_ref", [])
+    extras = kwargs.get("verify_extra", [])
+    scores = []
+
+    for i, completion in enumerate(completions):
+        text = completion[0]["content"] if completion else ""
+        prompt = prompts[i][0]["content"] if prompts and i < len(prompts) else ""
+        srv_name = (
+            refs[i]["name"] if i < len(refs) and isinstance(refs[i], dict) else ""
+        )
+        url = verify_urls.get(srv_name, "")
+
+        if not url:
+            scores.append(0.0)
+            continue
+
+        extra = extras[i] if i < len(extras) else {}
+        req = {k: v for k, v in extra.items() if v is not None}
+        req["responses_create_params"] = {
+            "input": [{"role": "user", "content": prompt}]
+        }
+        req["response"] = {
+            "id": "resp",
+            "created_at": 0,
+            "model": "axolotl",
+            "object": "response",
+            "output": [
+                {
+                    "id": "msg",
+                    "role": "assistant",
+                    "type": "message",
+                    "status": "completed",
+                    "content": [
+                        {"type": "output_text", "text": text, "annotations": []}
+                    ],
+                }
+            ],
+            "parallel_tool_calls": True,
+            "tool_choice": "auto",
+            "tools": [],
+        }
+
+        try:
+            resp = requests.post(url, json=req, timeout=30)
+            reward = resp.json().get("reward", 0.0) if resp.ok else 0.0
+        except Exception as exc:
+            LOG.warning(f"Verify request to {url} failed: {exc}")
+            reward = 0.0
+
+        scores.append(float(reward))
+
+    return scores
+
+
+# ---------------------------------------------------------------------------
+# Factory used internally by the plugin
+# ---------------------------------------------------------------------------
+
+
+def create_nemo_gym_reward_fn(
+    global_config: dict,
+    verify_endpoints: dict[str, str],
+    model_name: str = "axolotl-model",
+    verify_timeout: int = 30,
+):
+    """Create a reward function bound to specific verify endpoints.
+
+    Used internally by ``NemoGymPlugin._wire_single_turn()`` to inject a
+    reward function that already knows the endpoint map (no discovery needed).
+    """
+
+    def reward_fn(
+        completions: list[list[dict[str, str]]],
+        prompts: list[list[dict[str, str]]] | None = None,
+        **kwargs: Any,
+    ) -> np.ndarray:
+        resources_server_refs = kwargs.get("resources_server_ref", [])
+        verify_extras = kwargs.get("verify_extra", [])
+
+        scores = []
+        for i, completion in enumerate(completions):
+            completion_text = completion[0]["content"]
+            task_prompt = prompts[i][0]["content"] if prompts else ""
+
+            server_name = (
+                resources_server_refs[i]["name"]
+                if i < len(resources_server_refs)
+                else None
+            )
+
+            if server_name is None or server_name not in verify_endpoints:
+                LOG.warning(
+                    f"No verify endpoint for server '{server_name}', returning 0 reward"
+                )
+                scores.append(0.0)
+                continue
+
+            verify_endpoint = verify_endpoints[server_name]
+
+            verify_request = (
+                {k: v for k, v in verify_extras[i].items() if v is not None}
+                if i < len(verify_extras)
+                else {}
+            )
+            verify_request["responses_create_params"] = {
+                "input": [{"role": "user", "content": task_prompt}]
+            }
+            verify_request["response"] = {
+                "id": "resp",
+                "created_at": 0,
+                "model": model_name,
+                "object": "response",
+                "output": [
+                    {
+                        "id": "msg",
+                        "role": "assistant",
+                        "type": "message",
+                        "status": "completed",
+                        "content": [
+                            {
+                                "type": "output_text",
+                                "text": completion_text,
+                                "annotations": [],
+                            }
+                        ],
+                    }
+                ],
+                "parallel_tool_calls": True,
+                "tool_choice": "auto",
+                "tools": [],
+            }
+
+            try:
+                resp = requests.post(
+                    verify_endpoint, json=verify_request, timeout=verify_timeout
+                )
+                if resp.status_code == 200:
+                    reward = resp.json().get("reward", 0.0)
+                else:
+                    LOG.warning(
+                        f"Verify request returned status {resp.status_code}: {resp.text[:200]}"
+                    )
+                    reward = 0.0
+            except requests.exceptions.RequestException as exc:
+                LOG.warning(f"Verify request failed: {exc}")
+                reward = 0.0
+
+            scores.append(float(reward))
+
+        return np.array(scores)
+
+    return reward_fn
diff --git a/src/axolotl/integrations/nemo_gym/server.py b/src/axolotl/integrations/nemo_gym/server.py
new file mode 100644
index 0000000000..bd619569ef
--- /dev/null
+++ b/src/axolotl/integrations/nemo_gym/server.py
@@ -0,0 +1,262 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
+# you may not use this file except in compliance with the License.
+
+"""
+NeMo Gym server lifecycle management.
+
+Handles cloning the NeMo Gym repo, starting resource servers,
+waiting for readiness, and cleanup on exit.
+"""
+
+from __future__ import annotations
+
+import atexit
+import os
+import subprocess  # nosec B404
+import time
+
+import requests
+import yaml
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_ng_process = None
+_ng_log_file = None
+
+
+def ensure_gym_repo(gym_dir: str, auto_clone: bool = True) -> str:
+    """Clone the NeMo Gym repo if it doesn't exist.
+
+    Args:
+        gym_dir: Path to the NeMo Gym directory.
+        auto_clone: Whether to auto-clone if missing.
+
+    Returns:
+        Resolved path to the NeMo Gym directory.
+    """
+    gym_dir = os.path.expanduser(gym_dir)
+    if os.path.exists(gym_dir):
+        LOG.info(f"NeMo Gym directory exists at {gym_dir}")
+        return gym_dir
+
+    if not auto_clone:
+        raise FileNotFoundError(
+            f"NeMo Gym directory not found at {gym_dir} and auto_clone is disabled."
+        )
+
+    LOG.info(f"Cloning NeMo Gym to {gym_dir}...")
+    subprocess.run(  # nosec
+        ["git", "clone", "https://github.com/NVIDIA-NeMo/Gym.git", gym_dir],
+        check=True,
+    )
+    return gym_dir
+
+
+def ensure_gym_venv(gym_dir: str):
+    """Set up the NeMo Gym Python venv if not present."""
+    venv_python = os.path.join(gym_dir, ".venv", "bin", "python")
+    if os.path.exists(venv_python):
+        return
+
+    LOG.info("Setting up NeMo Gym venv...")
+    subprocess.run(["uv", "venv", "--python", "3.12"], cwd=gym_dir, check=True)  # nosec
+    subprocess.run(  # nosec
+        ["bash", "-c", "source .venv/bin/activate && uv sync"],
+        cwd=gym_dir,
+        check=True,
+    )
+
+
+def start_servers(
+    gym_dir: str,
+    config_paths: list[str],
+    head_port: int = 11000,
+    timeout: int = 360,
+):
+    """Start NeMo Gym resource servers via ng_run.
+
+    Args:
+        gym_dir: Path to the NeMo Gym directory.
+        config_paths: List of config YAML paths relative to gym_dir.
+        head_port: Port for the head server.
+        timeout: Max seconds to wait for servers.
+    """
+    global _ng_process, _ng_log_file
+
+    head_url = f"http://127.0.0.1:{head_port}/global_config_dict_yaml"
+
+    # Check if already running
+    try:
+        requests.get(head_url, timeout=2)
+        LOG.info(f"NeMo Gym servers already running on port {head_port}.")
+        return
+    except (requests.exceptions.ConnectionError, requests.exceptions.Timeout):
+        pass
+
+    ng_run_bin = os.path.join(gym_dir, ".venv", "bin", "ng_run")
+    config_arg = f"+config_paths=[{','.join(config_paths)}]"
+    _ng_log_file = open(os.path.join(gym_dir, "ng_run.log"), "w")  # noqa: SIM115
+    _ng_process = subprocess.Popen(  # nosec B603
+        [ng_run_bin, config_arg, "+skip_venv_if_present=true"],
+        cwd=gym_dir,
+        stdout=_ng_log_file,
+        stderr=subprocess.STDOUT,
+    )
+
+    atexit.register(_cleanup_servers)
+
+    LOG.info("Waiting for NeMo Gym head server...")
+    for _ in range(timeout // 3):
+        try:
+            requests.get(head_url, timeout=2)
+            LOG.info("NeMo Gym head server is ready.")
+            return
+        except (requests.exceptions.ConnectionError, requests.exceptions.Timeout):
+            if _ng_process.poll() is not None:
+                raise RuntimeError(
+                    "NeMo Gym server process exited unexpectedly. "
+                    f"Check {gym_dir}/ng_run.log for details."
+                ) from None
+            time.sleep(3)
+
+    raise RuntimeError(
+        f"NeMo Gym servers did not start within {timeout}s. "
+        f"Check {gym_dir}/ng_run.log for details."
+    )
+
+
+def get_server_configs(head_port: int = 11000, timeout: float = 30.0) -> dict:
+    """Fetch the global config from the NeMo Gym head server.
+
+    Retries up to 3 times with exponential backoff. The default per-attempt
+    timeout is 30s (raised from the original 5s) because head servers can
+    be slow to respond when they're concurrently serving rollouts from a
+    prior training run. A 5s timeout was empirically too tight to survive
+    a kill-and-relaunch cycle.
+
+    Returns:
+        Dict mapping server_name -> server config.
+    """
+    url = f"http://127.0.0.1:{head_port}/global_config_dict_yaml"
+    last_exc: Exception | None = None
+    for attempt in (1, 2, 3):
+        try:
+            response = requests.get(url, timeout=timeout)
+            response.raise_for_status()
+            result = yaml.safe_load(response.text)
+            # NeMo Gym head server double-encodes: YAML string inside a YAML string
+            if isinstance(result, str):
+                result = yaml.safe_load(result)
+            return result
+        except (requests.exceptions.RequestException, OSError) as exc:
+            last_exc = exc
+            LOG.warning(
+                "NeMo Gym head probe attempt %d/3 failed: %s. Retrying...",
+                attempt,
+                type(exc).__name__,
+            )
+            if attempt < 3:
+                time.sleep(2.0 * attempt)
+    raise RuntimeError(
+        f"NeMo Gym head server at {url} did not respond after 3 attempts: {last_exc}"
+    )
+
+
+def get_agent_servers(
+    global_config: dict, head_host: str = "127.0.0.1"
+) -> dict[str, str]:
+    """Discover NeMo Gym agent servers from the global config.
+
+    Agent servers handle multi-turn orchestration via /run endpoint.
+    Returns mapping of agent_name → URL (e.g., {"simple_agent": "http://host:port"}).
+    """
+    agents = {}
+    for top_name, top_cfg in global_config.items():
+        if not isinstance(top_cfg, dict):
+            continue
+        agent_dict = top_cfg.get("responses_api_agents", {})
+        if not agent_dict:
+            continue
+        for _agent_name, agent_cfg in agent_dict.items():
+            if not isinstance(agent_cfg, dict):
+                continue
+            host = agent_cfg.get("host", "127.0.0.1")
+            port = agent_cfg.get("port")
+            if not port:
+                continue
+            # Replace loopback with head_host for remote access
+            host = _normalize_host(host, fallback=head_host)
+            # Use the top-level config name (not the inner agent name)
+            # because dataset agent_ref.name references the top-level name
+            agents[top_name] = f"http://{host}:{port}"
+    if agents:
+        LOG.info(f"Discovered NeMo Gym agent servers: {agents}")
+    return agents
+
+
+def _normalize_host(host: str, fallback: str = "127.0.0.1") -> str:
+    """Normalize bind-all and loopback addresses for reachability."""
+    if host in ("0.0.0.0", "localhost"):  # nosec B104
+        return fallback
+    return host
+
+
+def get_server_base_url(global_config: dict, server_name: str) -> str:
+    """Get the base URL for a given resource server."""
+    try:
+        srv_cfg = global_config[server_name]["resources_servers"][server_name]
+        host = _normalize_host(srv_cfg["host"])
+        return f"http://{host}:{srv_cfg['port']}"
+    except (KeyError, TypeError) as exc:
+        raise ValueError(
+            f"Could not find resource server config for '{server_name}' in NeMo Gym. "
+            f"Available servers: {list(global_config.keys())}"
+        ) from exc
+
+
+def get_verify_endpoint(global_config: dict, server_name: str) -> str:
+    """Get the /verify endpoint URL for a given resource server."""
+    return f"{get_server_base_url(global_config, server_name)}/verify"
+
+
+def wait_for_resource_servers(global_config: dict, timeout: int = 180):
+    """Wait for all resource servers in the config to become reachable."""
+    for srv_name in global_config:
+        try:
+            srv_cfg = global_config[srv_name]["resources_servers"][srv_name]
+        except (KeyError, TypeError):
+            continue  # Skip non-server config entries silently
+
+        host, port = _normalize_host(srv_cfg["host"]), srv_cfg["port"]
+        LOG.info(f"Waiting for resource server '{srv_name}' at {host}:{port}...")
+        for _ in range(timeout // 2):
+            try:
+                requests.get(f"http://{host}:{port}/", timeout=2)
+                LOG.info(f"Resource server '{srv_name}' is ready.")
+                break
+            except requests.exceptions.ConnectionError:
+                time.sleep(2)
+        else:
+            raise RuntimeError(
+                f"Resource server '{srv_name}' at {host}:{port} "
+                f"did not start within {timeout}s."
+            )
+
+
+def _cleanup_servers():
+    """Terminate NeMo Gym server process on exit."""
+    global _ng_process, _ng_log_file
+    if _ng_process is not None and _ng_process.poll() is None:
+        LOG.info("Terminating NeMo Gym servers...")
+        _ng_process.terminate()
+        try:
+            _ng_process.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            _ng_process.kill()
+    if _ng_log_file is not None:
+        _ng_log_file.close()
diff --git a/src/axolotl/integrations/spectrum/LICENSE b/src/axolotl/integrations/spectrum/LICENSE
new file mode 100644
index 0000000000..d645695673
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/LICENSE
@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
diff --git a/src/axolotl/integrations/spectrum/README.md b/src/axolotl/integrations/spectrum/README.md
new file mode 100644
index 0000000000..0f78a511b5
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/README.md
@@ -0,0 +1,37 @@
+# Spectrum: Targeted Training on Signal to Noise Ratio
+
+by Eric Hartford, Lucas Atkins, Fernando Fernandes, David Golchinfar
+
+This plugin contains code to freeze the bottom fraction of modules in a model, based on the Signal-to-Noise Ratio (SNR).
+
+See https://github.com/cognitivecomputations/spectrum
+
+## Overview
+
+Spectrum is a tool for scanning and evaluating the Signal-to-Noise Ratio (SNR) of layers in large language models.
+By identifying the top n% of layers with the highest SNR, you can optimize training efficiency.
+
+## Usage
+
+```yaml
+plugins:
+  - axolotl.integrations.spectrum.SpectrumPlugin
+
+spectrum_top_fraction: 0.5
+# Optional if using a pre-scanned model as your base_model. Useful if using a model mirror
+spectrum_model_name: meta-llama/Meta-Llama-3.1-8B
+```
+
+## Citation
+
+```bib
+@misc{hartford2024spectrumtargetedtrainingsignal,
+      title={Spectrum: Targeted Training on Signal to Noise Ratio},
+      author={Eric Hartford and Lucas Atkins and Fernando Fernandes Neto and David Golchinfar},
+      year={2024},
+      eprint={2406.06623},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2406.06623},
+}
+```
diff --git a/src/axolotl/integrations/spectrum/__init__.py b/src/axolotl/integrations/spectrum/__init__.py
new file mode 100644
index 0000000000..5e8f9128d1
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/__init__.py
@@ -0,0 +1,104 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Spectrum Plugin to automatically generate unfrozen parameters based on SNR data.
+"""
+
+import json
+
+import requests
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+from .args import SpectrumArgs as SpectrumArgs
+
+LOG = get_logger(__name__)
+
+
+def _generate_unfrozen_params_yaml(snr_data, top_fraction=0.5):
+    unfrozen_parameters = {}
+    for layer_name, info in snr_data.items():
+        layer_type = info["type"]
+        if layer_type not in unfrozen_parameters:
+            unfrozen_parameters[layer_type] = []
+        unfrozen_parameters[layer_type].append((layer_name, info["snr"]))
+    top_layers_by_type = {}
+    for layer_type, layers in unfrozen_parameters.items():
+        layers_sorted = sorted(layers, key=lambda x: x[1], reverse=True)
+        num_top_layers = int(len(layers) * top_fraction)
+        top_layers_by_type[layer_type] = [
+            layer[0] for layer in layers_sorted[:num_top_layers]
+        ]
+    unfrozen_parameters = [
+        "^lm_head.weight$",
+        "^model.embed_tokens.weight$",
+    ]
+    for _, layer_names in top_layers_by_type.items():
+        for layer_name in layer_names:
+            unfrozen_parameters.append(layer_name)
+    return unfrozen_parameters
+
+
+class SpectrumPlugin(BasePlugin):
+    """
+    Spectrum Plugin to automatically generate unfrozen parameters based on SNR data.
+    """
+
+    base_url = "https://raw.githubusercontent.com/cognitivecomputations/spectrum/main/model_snr_results/"
+    base_path = "./model_snr_results/"
+    snr_file_template = "snr_results_{model_name_slug}.json"
+
+    def get_input_args(self):
+        return "axolotl.integrations.spectrum.SpectrumArgs"
+
+    def pre_model_load(self, cfg):
+        if cfg.get("spectrum_model_name"):
+            model_name = cfg["spectrum_model_name"]
+        else:
+            model_name = cfg["base_model"]
+        top_fraction = cfg.get("spectrum_top_fraction", 50)
+        model_slug = model_name.replace("/", "-").replace("_", "-")
+        snr_url = self.base_url + self.snr_file_template.format(
+            model_name_slug=model_slug
+        )
+        snr_path = self.base_path + self.snr_file_template.format(
+            model_name_slug=model_slug
+        )
+        # first check if the files exist locally and read the json
+        snr_data = None
+        try:
+            with open(snr_path, "r", encoding="utf-8") as fin:
+                snr_data = json.load(fin)
+        except FileNotFoundError:
+            pass
+        except Exception as exc:
+            LOG.warning(f"Failed to read SNR data from {snr_path}: {exc}")
+
+        if not snr_data:
+            try:
+                snr_data = requests.get(snr_url, timeout=60).json()
+            except requests.exceptions.RequestException as exc:
+                LOG.warning(f"Failed to fetch SNR data from {snr_url}: {exc}")
+                return
+            # also catch json parsing errors
+            except json.JSONDecodeError as exc:
+                LOG.warning(f"Failed to parse SNR data from {snr_url}: {exc}")
+                return
+
+        unfrozen_parameters = _generate_unfrozen_params_yaml(
+            snr_data, top_fraction=top_fraction
+        )
+        cfg["unfrozen_parameters"] = unfrozen_parameters
diff --git a/src/axolotl/integrations/spectrum/args.py b/src/axolotl/integrations/spectrum/args.py
new file mode 100644
index 0000000000..be6ca4bfc2
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/args.py
@@ -0,0 +1,47 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Module for handling Spectrum input arguments.
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel, model_validator
+
+
+class SpectrumArgs(BaseModel):
+    """
+    Input args for Spectrum.
+    """
+
+    spectrum_top_fraction: Optional[float] = 0.5
+    spectrum_model_name: Optional[str] = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fsdp_use_orig_params(cls, data):
+        if (
+            data.get("fsdp")
+            and data.get("fsdp_config")
+            and not data["fsdp_config"].get("use_orig_params")
+            and data.get("plugins")
+            and any("SpectrumPlugin" in plugin for plugin in data["plugins"])
+        ):
+            # would otherwise raise
+            # ValueError: Must flatten tensors with uniform `requires_grad` when `use_orig_params=False`
+            raise ValueError(
+                "FSDP + SpectrumPlugin cannot be used together when `use_orig_params=False` is set"
+            )
+        return data
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-1.5B-Instruct.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-1.5B-Instruct.json
new file mode 100644
index 0000000000..99168b7b26
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-1.5B-Instruct.json
@@ -0,0 +1,1022 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.26.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.27.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 70.50235748291016,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 134.4214630126953,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 235.74794006347656,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 73.25755310058594,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 27.22879981994629,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 17.5551815032959,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 54.210426330566406,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 38.808937072753906,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 29.799747467041016,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 10.296355247497559,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 8.86428165435791,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 6.43813943862915,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 7.0912184715271,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 3.285884141921997,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 6.073758125305176,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 5.325990676879883,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 4.591946601867676,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 7.021907329559326,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 6.392782211303711,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 210.51983642578125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 7.1035943031311035,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 18.701711654663086,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 14.842622756958008,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 10.50004768371582,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 7.225146770477295,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 7.463952541351318,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.26.mlp.down_proj": {
+        "snr": 15.226134300231934,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.27.mlp.down_proj": {
+        "snr": 105.4173355102539,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 0.5021594166755676,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 34.75935363769531,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 22.855531692504883,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 25.09166717529297,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 28.533172607421875,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 18.625717163085938,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 39.77565383911133,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 24.77678680419922,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 11.854388236999512,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 20.372356414794922,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 14.639552116394043,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 9.82955551147461,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 13.942151069641113,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 12.524999618530273,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 8.19681167602539,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 8.561081886291504,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 6.421900749206543,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 5.568161964416504,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 10.090147972106934,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 5.6181230545043945,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 5.173826694488525,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 5.663441181182861,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 6.824708461761475,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 4.724992275238037,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 6.829834938049316,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 9.968582153320312,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.26.mlp.gate_proj": {
+        "snr": 14.35350513458252,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.27.mlp.gate_proj": {
+        "snr": 20.121768951416016,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 1.9020992517471313,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 46.9393424987793,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 76.04901123046875,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 104.08525848388672,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 77.74343872070312,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 104.15605926513672,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 105.16349792480469,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 78.4150390625,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 57.51069641113281,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 50.26409912109375,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 50.36701965332031,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 56.66413497924805,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 62.384559631347656,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 44.97883987426758,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 69.7376480102539,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 35.93111801147461,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 33.63168716430664,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 37.695919036865234,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 43.516517639160156,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 30.479318618774414,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 12.495409965515137,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 19.616689682006836,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 18.42948341369629,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 10.799560546875,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 14.167623519897461,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 14.938597679138184,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.26.mlp.up_proj": {
+        "snr": 8.896568298339844,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.27.mlp.up_proj": {
+        "snr": 25.774547576904297,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": Infinity,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.26.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.27.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 1.8306859731674194,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.896544337272644,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 2.345759868621826,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 2.0610744953155518,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 2.3658556938171387,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 1.6586917638778687,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 1.7613047361373901,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 1.325312852859497,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 1.458108901977539,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 1.4319790601730347,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 0.9579543471336365,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 0.8787619471549988,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 1.0447536706924438,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 0.9157310724258423,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.7528730630874634,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 0.9293556213378906,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 0.8057093620300293,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 1.2973601818084717,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 1.1357901096343994,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 1.3661632537841797,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 0.8829066753387451,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 0.9105398654937744,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 2.086926221847534,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 1.0393351316452026,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 1.114574670791626,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 2.599745035171509,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.26.self_attn.k_proj": {
+        "snr": 1.1256712675094604,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.27.self_attn.k_proj": {
+        "snr": 1.1784162521362305,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.8094121813774109,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.22000817954540253,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.21972468495368958,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.22064059972763062,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.22308556735515594,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.22396250069141388,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.228360116481781,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.2306283563375473,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.2430228292942047,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.2115175724029541,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.18226943910121918,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.144245907664299,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.21965907514095306,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.1797526627779007,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.26513636112213135,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.19463808834552765,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.22129350900650024,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.22545330226421356,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.25302645564079285,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.26326504349708557,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.15203869342803955,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.22418837249279022,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.23777326941490173,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.18076598644256592,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.19919466972351074,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.11310968548059464,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.26.self_attn.o_proj": {
+        "snr": 0.08452697843313217,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.27.self_attn.o_proj": {
+        "snr": 0.1029304787516594,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.03922705352306366,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.1410205066204071,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.18240582942962646,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.1702580451965332,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.19508686661720276,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.21549257636070251,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.22021502256393433,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.2044307142496109,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.22745060920715332,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.23825915157794952,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.2181481122970581,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.23490090668201447,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.2379382699728012,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.19233369827270508,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.2587313652038574,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.07332809269428253,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.22992204129695892,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.2537729740142822,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.2389948070049286,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.20716068148612976,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.2575169503688812,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.22347678244113922,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.18831054866313934,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.19853907823562622,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.16343259811401367,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.1583252102136612,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.26.self_attn.q_proj": {
+        "snr": 0.254446804523468,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.27.self_attn.q_proj": {
+        "snr": 0.23828543722629547,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 856.5148315429688,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 48.941104888916016,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": 70.25466918945312,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": 370.885986328125,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 75.51139831542969,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 52.004058837890625,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": 641.026611328125,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": 323.4858093261719,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": 2.1745388507843018,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.26.self_attn.v_proj": {
+        "snr": 3.0791690349578857,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.27.self_attn.v_proj": {
+        "snr": 2.029968023300171,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-1.5B.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-1.5B.json
new file mode 100644
index 0000000000..41cde8e6b0
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-1.5B.json
@@ -0,0 +1,1022 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.26.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.27.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 70.4939193725586,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 134.2310028076172,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 235.44140625,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 73.19381713867188,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 27.216264724731445,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 17.544504165649414,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 54.17462158203125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 38.78171920776367,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 29.777149200439453,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 10.289377212524414,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 8.858332633972168,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 6.433396816253662,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 7.085702419281006,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 3.323948383331299,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 6.204164505004883,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 5.321533203125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 4.588479995727539,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 7.01450252532959,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 6.386813163757324,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 210.38458251953125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 7.096683979034424,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 18.68245506286621,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 14.824685096740723,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 10.491303443908691,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 7.2194437980651855,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 7.458613872528076,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.26.mlp.down_proj": {
+        "snr": 15.222760200500488,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.27.mlp.down_proj": {
+        "snr": 105.41569519042969,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 0.5017311573028564,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 34.71562576293945,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 22.82915496826172,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 25.0699520111084,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 28.508079528808594,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 18.608009338378906,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 39.732391357421875,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 24.760026931762695,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 11.842738151550293,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 20.35906982421875,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 14.627532958984375,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 9.821962356567383,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 13.930404663085938,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 12.509871482849121,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 8.187695503234863,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 8.553187370300293,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 6.414614200592041,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 5.561778545379639,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 10.078697204589844,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 5.61345100402832,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 5.265484809875488,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 5.659949779510498,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 6.8203511238098145,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 4.721294403076172,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 6.82572603225708,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 9.963521003723145,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.26.mlp.gate_proj": {
+        "snr": 14.342291831970215,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.27.mlp.gate_proj": {
+        "snr": 20.092098236083984,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 1.901187777519226,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 46.9141731262207,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 76.07878112792969,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 103.9194564819336,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 77.62561798095703,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 104.01624298095703,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 105.0235366821289,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 78.33445739746094,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 57.44070816040039,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 50.20344924926758,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 50.32845687866211,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 56.6197624206543,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 62.338096618652344,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 44.92917251586914,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 69.69624328613281,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 35.90705108642578,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 33.610374450683594,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 37.67365646362305,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 43.488929748535156,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 30.451993942260742,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 12.480182647705078,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 19.595102310180664,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 19.067970275878906,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 10.786394119262695,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 14.150126457214355,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 14.927021026611328,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.26.mlp.up_proj": {
+        "snr": 8.891448020935059,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.27.mlp.up_proj": {
+        "snr": 25.74305534362793,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": Infinity,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.26.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.27.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 1.7818864583969116,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.8955822587013245,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 2.344149351119995,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 2.0597119331359863,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 2.36411714553833,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 1.6570613384246826,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 1.7604507207870483,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 1.3245182037353516,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 1.4567548036575317,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 1.4310829639434814,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 0.95713210105896,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 0.8781776428222656,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 1.0438013076782227,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 0.9315219521522522,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.7521569728851318,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 0.9286947250366211,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 0.8047553896903992,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 1.2965552806854248,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 1.134974479675293,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 1.3648872375488281,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 0.8667459487915039,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 0.9100639224052429,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 2.127535820007324,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 1.0382369756698608,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 1.113753318786621,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 2.597890853881836,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.26.self_attn.k_proj": {
+        "snr": 1.1248247623443604,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.27.self_attn.k_proj": {
+        "snr": 1.1984941959381104,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.8139898777008057,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.21965594589710236,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.219479501247406,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.22144284844398499,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.22390463948249817,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.22383669018745422,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.22818723320960999,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.23134392499923706,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.24275101721286774,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.21139128506183624,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.18210072815418243,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.14415481686592102,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.21947966516017914,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.17875106632709503,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.264996200799942,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.19353187084197998,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.22111012041568756,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.2242278754711151,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.2527434229850769,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.26184532046318054,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.1519661247730255,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.22386522591114044,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.2386160045862198,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.18057651817798615,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.1989467740058899,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.11306505650281906,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.26.self_attn.o_proj": {
+        "snr": 0.08449216932058334,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.27.self_attn.o_proj": {
+        "snr": 0.10287519544363022,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.039204664528369904,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.14075909554958344,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.18212397396564484,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.1700422316789627,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.1948907971382141,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.2153141051530838,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.21998055279254913,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.20416118204593658,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.2272879034280777,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.23795834183692932,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.21887299418449402,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.23469635844230652,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.23774078488349915,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.1920779049396515,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.2584812641143799,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.07330238074064255,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.23073157668113708,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.2523840367794037,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.23874858021736145,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.20698708295822144,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.25723400712013245,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.223300039768219,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.18824049830436707,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.19840741157531738,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.16326843202114105,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.1581888198852539,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.26.self_attn.q_proj": {
+        "snr": 0.25306230783462524,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.27.self_attn.q_proj": {
+        "snr": 0.23808495700359344,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 864.8881225585938,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 48.853694915771484,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": 70.18457794189453,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": 371.1153259277344,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 75.41203308105469,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 51.92624282836914,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": 642.9313354492188,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": 323.5724182128906,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": 2.1736748218536377,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.26.self_attn.v_proj": {
+        "snr": 3.1729259490966797,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.27.self_attn.v_proj": {
+        "snr": 2.024953842163086,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-3B-Instruct.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-3B-Instruct.json
new file mode 100644
index 0000000000..6a67b14b3a
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-3B-Instruct.json
@@ -0,0 +1,1310 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.26.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.27.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.28.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.29.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.30.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.31.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.32.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.33.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.34.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.35.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 20.964319229125977,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 0.11561352014541626,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 0.14991413056850433,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 0.3673713207244873,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 0.5076134204864502,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 33.89468002319336,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 45.08732986450195,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 33.234222412109375,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 29.3447322845459,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 26.664169311523438,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 22.323949813842773,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 18.259737014770508,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 14.422037124633789,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 22.172054290771484,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 27.363698959350586,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 28.474334716796875,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 10.4143648147583,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 10.719133377075195,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 8.6494722366333,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 5.69321870803833,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 23.889677047729492,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 11.59121036529541,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 5.997435569763184,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 19.415578842163086,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 8.241704940795898,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 12.993823051452637,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.26.mlp.down_proj": {
+        "snr": 36.26508712768555,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.27.mlp.down_proj": {
+        "snr": 19.957971572875977,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.28.mlp.down_proj": {
+        "snr": 6.067765235900879,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.29.mlp.down_proj": {
+        "snr": 5.369481086730957,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.30.mlp.down_proj": {
+        "snr": 7.358774662017822,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.31.mlp.down_proj": {
+        "snr": 7.8687238693237305,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.32.mlp.down_proj": {
+        "snr": 8.713484764099121,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.33.mlp.down_proj": {
+        "snr": 21.233531951904297,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.34.mlp.down_proj": {
+        "snr": 32.37357711791992,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.35.mlp.down_proj": {
+        "snr": 179.8053741455078,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 0.24989914894104004,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 0.11613649874925613,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 0.16354432702064514,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 0.36216047406196594,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 0.3485107719898224,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 2.6546616554260254,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 8.362885475158691,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 7.38665246963501,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 13.016111373901367,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 14.94902515411377,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 20.92418670654297,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 15.954015731811523,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 8.980009078979492,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 17.59958267211914,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 17.23070526123047,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 23.725330352783203,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 17.000444412231445,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 18.293012619018555,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 12.644190788269043,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 16.278690338134766,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 7.407368183135986,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 6.109912395477295,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 5.3692426681518555,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 9.354235649108887,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 7.655010223388672,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 6.252986431121826,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.26.mlp.gate_proj": {
+        "snr": 14.26718521118164,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.27.mlp.gate_proj": {
+        "snr": 7.705836772918701,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.28.mlp.gate_proj": {
+        "snr": 5.998677730560303,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.29.mlp.gate_proj": {
+        "snr": 6.044872760772705,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.30.mlp.gate_proj": {
+        "snr": 9.027137756347656,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.31.mlp.gate_proj": {
+        "snr": 5.449969291687012,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.32.mlp.gate_proj": {
+        "snr": 4.206825256347656,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.33.mlp.gate_proj": {
+        "snr": 5.22825288772583,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.34.mlp.gate_proj": {
+        "snr": 43.71927261352539,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.35.mlp.gate_proj": {
+        "snr": 45.37385177612305,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 0.7069714665412903,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 0.17766596376895905,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 0.28577035665512085,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 0.6763099431991577,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 0.8340913653373718,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 3.946547031402588,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 19.56715202331543,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 36.21149826049805,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 44.28759002685547,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 45.47198486328125,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 79.00128936767578,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 52.28038787841797,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 48.08102035522461,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 56.071285247802734,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 72.24358367919922,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 54.818233489990234,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 47.251495361328125,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 51.585636138916016,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 43.47938919067383,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 38.132469177246094,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 21.78435707092285,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 22.261096954345703,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 30.751861572265625,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 28.61063575744629,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 20.21415901184082,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 20.759052276611328,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.26.mlp.up_proj": {
+        "snr": 33.80818557739258,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.27.mlp.up_proj": {
+        "snr": 17.274362564086914,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.28.mlp.up_proj": {
+        "snr": 13.943653106689453,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.29.mlp.up_proj": {
+        "snr": 16.202186584472656,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.30.mlp.up_proj": {
+        "snr": 24.25114631652832,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.31.mlp.up_proj": {
+        "snr": 10.68645191192627,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.32.mlp.up_proj": {
+        "snr": 5.7449774742126465,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.33.mlp.up_proj": {
+        "snr": 11.879876136779785,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.34.mlp.up_proj": {
+        "snr": 25.948715209960938,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.35.mlp.up_proj": {
+        "snr": 38.63526153564453,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": Infinity,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.26.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.27.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.28.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.29.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.30.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.31.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.32.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.33.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.34.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.35.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 12.243099212646484,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.6446183323860168,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 0.7159711718559265,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 5.5100932121276855,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 3.0802414417266846,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 1.0472767353057861,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 3.576918601989746,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 3.3793225288391113,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 2.9598212242126465,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 6.102792263031006,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 2.231630325317383,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 2.176372766494751,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 1.3229435682296753,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 2.6183862686157227,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 2.608288526535034,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 1.5090984106063843,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 1.284422516822815,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 0.8903945088386536,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 1.8880385160446167,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 0.8905735015869141,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 0.9060881733894348,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 0.7572551965713501,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 0.940827488899231,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 3.7776191234588623,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 1.328923225402832,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 1.3986345529556274,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.26.self_attn.k_proj": {
+        "snr": 1.2436336278915405,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.27.self_attn.k_proj": {
+        "snr": 0.7737217545509338,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.28.self_attn.k_proj": {
+        "snr": 2.6027626991271973,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.29.self_attn.k_proj": {
+        "snr": 2.2332751750946045,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.30.self_attn.k_proj": {
+        "snr": 2.476585626602173,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.31.self_attn.k_proj": {
+        "snr": 1.1115432977676392,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.32.self_attn.k_proj": {
+        "snr": 0.8251476287841797,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.33.self_attn.k_proj": {
+        "snr": 0.9331105947494507,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.34.self_attn.k_proj": {
+        "snr": 6.602395534515381,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.35.self_attn.k_proj": {
+        "snr": 10.151693344116211,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.3661542534828186,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.19571374356746674,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.2244851142168045,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.2593664526939392,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.2569783926010132,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.2564302980899811,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.18539844453334808,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.2328651398420334,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.22055882215499878,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.21800543367862701,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.22867777943611145,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.23986175656318665,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.17598563432693481,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.20469218492507935,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.21040217578411102,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.23787625133991241,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.16339677572250366,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.2070712298154831,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.1826934814453125,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.19459959864616394,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.2668156027793884,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.16906610131263733,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.18790249526500702,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.18883933126926422,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.1793188899755478,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.1800570785999298,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.26.self_attn.o_proj": {
+        "snr": 0.17790433764457703,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.27.self_attn.o_proj": {
+        "snr": 0.2029498964548111,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.28.self_attn.o_proj": {
+        "snr": 0.17044201493263245,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.29.self_attn.o_proj": {
+        "snr": 0.19938386976718903,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.30.self_attn.o_proj": {
+        "snr": 0.23108959197998047,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.31.self_attn.o_proj": {
+        "snr": 0.16427059471607208,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.32.self_attn.o_proj": {
+        "snr": 0.10631092637777328,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.33.self_attn.o_proj": {
+        "snr": 0.09417019784450531,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.34.self_attn.o_proj": {
+        "snr": 0.1324978619813919,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.35.self_attn.o_proj": {
+        "snr": 0.11784011125564575,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.05565479397773743,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.138458251953125,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.12992437183856964,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.15362468361854553,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.1563446819782257,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.15544593334197998,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.15956827998161316,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.17549948394298553,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.16668449342250824,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.15626586973667145,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.18318884074687958,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.171547532081604,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.18164905905723572,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.2091975212097168,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.17431670427322388,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.20902502536773682,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.15439842641353607,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.1945274919271469,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.18916545808315277,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.20778712630271912,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.20866931974887848,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.1900305300951004,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.18200653791427612,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.2070988416671753,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.1845332235097885,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.20868781208992004,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.26.self_attn.q_proj": {
+        "snr": 0.19242744147777557,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.27.self_attn.q_proj": {
+        "snr": 0.15225112438201904,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.28.self_attn.q_proj": {
+        "snr": 0.20065009593963623,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.29.self_attn.q_proj": {
+        "snr": 0.19390477240085602,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.30.self_attn.q_proj": {
+        "snr": 0.18538697063922882,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.31.self_attn.q_proj": {
+        "snr": 0.18954339623451233,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.32.self_attn.q_proj": {
+        "snr": 0.20089596509933472,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.33.self_attn.q_proj": {
+        "snr": 0.19814996421337128,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.34.self_attn.q_proj": {
+        "snr": 0.17733213305473328,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.35.self_attn.q_proj": {
+        "snr": 0.14075976610183716,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 845.8053588867188,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 83.97241973876953,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": 213.70960998535156,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": 18.950267791748047,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": 435.8339538574219,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.26.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.27.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.28.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.29.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.30.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.31.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.32.self_attn.v_proj": {
+        "snr": 1.2341279983520508,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.33.self_attn.v_proj": {
+        "snr": 0.6158654689788818,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.34.self_attn.v_proj": {
+        "snr": 509.3221130371094,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.35.self_attn.v_proj": {
+        "snr": 538.6658325195312,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-3B.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-3B.json
new file mode 100644
index 0000000000..93b5cfec61
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-3B.json
@@ -0,0 +1,1310 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.26.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.27.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.28.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.29.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.30.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.31.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.32.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.33.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.34.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.35.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 20.942785263061523,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 0.11550866067409515,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 0.14981402456760406,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 0.36719316244125366,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 0.5072987079620361,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 33.86688232421875,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 45.066246032714844,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 33.20981979370117,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 29.310104370117188,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 26.638381958007812,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 22.302486419677734,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 18.249290466308594,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 14.057564735412598,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 22.154281616210938,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 27.348575592041016,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 28.447378158569336,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 10.405216217041016,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 10.71042251586914,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 8.642854690551758,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 5.690433979034424,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 23.869070053100586,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 11.584356307983398,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 5.992950916290283,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 18.495361328125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 8.233827590942383,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 12.626734733581543,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.26.mlp.down_proj": {
+        "snr": 36.21802520751953,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.27.mlp.down_proj": {
+        "snr": 19.932941436767578,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.28.mlp.down_proj": {
+        "snr": 6.0616455078125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.29.mlp.down_proj": {
+        "snr": 5.363720417022705,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.30.mlp.down_proj": {
+        "snr": 7.455615520477295,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.31.mlp.down_proj": {
+        "snr": 7.8631815910339355,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.32.mlp.down_proj": {
+        "snr": 8.706913948059082,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.33.mlp.down_proj": {
+        "snr": 21.220134735107422,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.34.mlp.down_proj": {
+        "snr": 32.33852005004883,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.35.mlp.down_proj": {
+        "snr": 179.8906707763672,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 0.24970805644989014,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 0.11607512086629868,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 0.16310769319534302,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 0.3621424436569214,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 0.3482637107372284,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 2.6533455848693848,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 8.359040260314941,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 7.382037162780762,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 13.00683879852295,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 14.936161994934082,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 20.907283782958984,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 15.941497802734375,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 8.97419548034668,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 17.585100173950195,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 17.21462059020996,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 23.703285217285156,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 16.986576080322266,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 18.27729606628418,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 12.63351058959961,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 16.2633113861084,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 7.399787902832031,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 6.10424280166626,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 5.363350868225098,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 9.344535827636719,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 7.647364616394043,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 6.143579959869385,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.26.mlp.gate_proj": {
+        "snr": 14.254817008972168,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.27.mlp.gate_proj": {
+        "snr": 7.7000861167907715,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.28.mlp.gate_proj": {
+        "snr": 5.994422435760498,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.29.mlp.gate_proj": {
+        "snr": 6.041909694671631,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.30.mlp.gate_proj": {
+        "snr": 9.027522087097168,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.31.mlp.gate_proj": {
+        "snr": 5.450753211975098,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.32.mlp.gate_proj": {
+        "snr": 4.149200439453125,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.33.mlp.gate_proj": {
+        "snr": 5.223763942718506,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.34.mlp.gate_proj": {
+        "snr": 43.65521240234375,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.35.mlp.gate_proj": {
+        "snr": 45.312774658203125,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 0.7065013647079468,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 0.17752516269683838,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 0.2847473919391632,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 0.6757690906524658,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 0.8353318572044373,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 3.940711736679077,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 19.556047439575195,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 36.19340515136719,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 44.2518424987793,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 45.418025970458984,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 78.90928649902344,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 52.24648666381836,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 48.02030563354492,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 56.016239166259766,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 72.16619873046875,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 54.75283432006836,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 47.204097747802734,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 51.549312591552734,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 43.43872833251953,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 38.09785461425781,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 21.767858505249023,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 22.243661880493164,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 30.71843147277832,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 28.5756778717041,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 20.186717987060547,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 20.742860794067383,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.26.mlp.up_proj": {
+        "snr": 33.777984619140625,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.27.mlp.up_proj": {
+        "snr": 17.254213333129883,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.28.mlp.up_proj": {
+        "snr": 13.930026054382324,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.29.mlp.up_proj": {
+        "snr": 16.17984390258789,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.30.mlp.up_proj": {
+        "snr": 24.236648559570312,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.31.mlp.up_proj": {
+        "snr": 10.665648460388184,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.32.mlp.up_proj": {
+        "snr": 5.735939025878906,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.33.mlp.up_proj": {
+        "snr": 11.592061042785645,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.34.mlp.up_proj": {
+        "snr": 25.923419952392578,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.35.mlp.up_proj": {
+        "snr": 38.579349517822266,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": Infinity,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.26.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.27.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.28.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.29.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.30.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.31.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.32.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.33.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.34.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.35.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 12.24727725982666,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.6436238288879395,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 0.7156716585159302,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 5.505439758300781,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 3.0760715007781982,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 1.0453941822052002,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 3.57472562789917,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 3.3765170574188232,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 2.8859639167785645,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 6.09852409362793,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 2.229580879211426,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 2.173879623413086,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 1.3220131397247314,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 2.61668062210083,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 2.606799840927124,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 1.5080311298370361,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 1.2841484546661377,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 0.8896433115005493,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 1.8873414993286133,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 0.8897770643234253,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 0.9051405787467957,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 0.7568970322608948,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 0.9403582811355591,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 3.777062177658081,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 1.3280683755874634,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 1.3980307579040527,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.26.self_attn.k_proj": {
+        "snr": 1.2435240745544434,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.27.self_attn.k_proj": {
+        "snr": 0.7732619047164917,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.28.self_attn.k_proj": {
+        "snr": 2.6010243892669678,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.29.self_attn.k_proj": {
+        "snr": 2.232773780822754,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.30.self_attn.k_proj": {
+        "snr": 2.4743099212646484,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.31.self_attn.k_proj": {
+        "snr": 1.11082923412323,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.32.self_attn.k_proj": {
+        "snr": 0.8243986368179321,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.33.self_attn.k_proj": {
+        "snr": 0.932928204536438,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.34.self_attn.k_proj": {
+        "snr": 6.608611583709717,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.35.self_attn.k_proj": {
+        "snr": 10.160987854003906,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.36662933230400085,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.1955128312110901,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.22419843077659607,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.25902292132377625,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.2567676901817322,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.2560890316963196,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.18518221378326416,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.23254290223121643,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.2203962802886963,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.217017263174057,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.22843335568904877,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.23816843330860138,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.17585325241088867,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.20451271533966064,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.2095799297094345,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.23767071962356567,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.16328400373458862,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.20690056681632996,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.18191492557525635,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.1945018619298935,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.26658856868743896,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.16897724568843842,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.18773262202739716,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.18808405101299286,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.17919476330280304,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.1793426126241684,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.26.self_attn.o_proj": {
+        "snr": 0.1777871698141098,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.27.self_attn.o_proj": {
+        "snr": 0.20279864966869354,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.28.self_attn.o_proj": {
+        "snr": 0.17030371725559235,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.29.self_attn.o_proj": {
+        "snr": 0.1992504596710205,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.30.self_attn.o_proj": {
+        "snr": 0.23085352778434753,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.31.self_attn.o_proj": {
+        "snr": 0.1641533523797989,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.32.self_attn.o_proj": {
+        "snr": 0.10621391236782074,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.33.self_attn.o_proj": {
+        "snr": 0.09411631524562836,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.34.self_attn.o_proj": {
+        "snr": 0.13239727914333344,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.35.self_attn.o_proj": {
+        "snr": 0.11740171164274216,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.055595725774765015,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.13823610544204712,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.1297825127840042,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.15291297435760498,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.15615035593509674,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.15535500645637512,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.15993140637874603,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.1753682643175125,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.1664913445711136,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.15656901895999908,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.18300014734268188,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.1713649481534958,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.1809009313583374,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.20895132422447205,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.17413195967674255,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.20878490805625916,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.1547088772058487,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.1943129003047943,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.1889297217130661,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.207680344581604,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.20839959383010864,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.18989044427871704,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.18180623650550842,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.2069384753704071,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.1842993050813675,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.2078687846660614,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.26.self_attn.q_proj": {
+        "snr": 0.19224946200847626,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.27.self_attn.q_proj": {
+        "snr": 0.15170617401599884,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.28.self_attn.q_proj": {
+        "snr": 0.20116600394248962,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.29.self_attn.q_proj": {
+        "snr": 0.19373668730258942,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.30.self_attn.q_proj": {
+        "snr": 0.18462225794792175,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.31.self_attn.q_proj": {
+        "snr": 0.18939673900604248,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.32.self_attn.q_proj": {
+        "snr": 0.20071947574615479,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.33.self_attn.q_proj": {
+        "snr": 0.19740056991577148,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.34.self_attn.q_proj": {
+        "snr": 0.17658494412899017,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.35.self_attn.q_proj": {
+        "snr": 0.1407373696565628,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 846.30126953125,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 83.83415222167969,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": 213.51316833496094,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": 18.92746925354004,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": 433.9771728515625,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.26.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.27.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.28.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.29.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.30.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.31.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.32.self_attn.v_proj": {
+        "snr": 1.2332282066345215,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.33.self_attn.v_proj": {
+        "snr": 0.6151890158653259,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.34.self_attn.v_proj": {
+        "snr": 509.7169189453125,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.35.self_attn.v_proj": {
+        "snr": 536.0748901367188,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-7B-Instruct.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-7B-Instruct.json
new file mode 100644
index 0000000000..0e18bd3864
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-7B-Instruct.json
@@ -0,0 +1,1022 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.26.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.27.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 10.283808708190918,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 1.2089825868606567,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 19.309062957763672,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 50.174461364746094,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 114.28582763671875,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 215.5762176513672,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 204.5117950439453,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 182.5479278564453,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 74.92950439453125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 16.482666015625,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 55.33920669555664,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 16.851062774658203,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 58.65230178833008,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 11.150161743164062,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 65.32643127441406,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 46.736305236816406,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 14.288785934448242,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 23.40110206604004,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 86.34363555908203,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 49.14613342285156,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 1276.84814453125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 51.803409576416016,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 143.0666046142578,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 35.14984893798828,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 21.41700553894043,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 10.651569366455078,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.26.mlp.down_proj": {
+        "snr": 21.635149002075195,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.27.mlp.down_proj": {
+        "snr": 1446.2774658203125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 0.04497330263257027,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 0.16888172924518585,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 0.33653727173805237,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 3.1445391178131104,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 9.107144355773926,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 15.909018516540527,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 60.9138069152832,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 57.570281982421875,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 65.82791137695312,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 10.455283164978027,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 26.970706939697266,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 31.139820098876953,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 43.987159729003906,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 20.704849243164062,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 21.191452026367188,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 42.66447830200195,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 22.136825561523438,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 22.60980987548828,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 81.80574035644531,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 20.88619613647461,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 58.3524055480957,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 22.786706924438477,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 16.932226181030273,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 16.819862365722656,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 19.76348304748535,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 28.98714256286621,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.26.mlp.gate_proj": {
+        "snr": 36.7071533203125,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.27.mlp.gate_proj": {
+        "snr": 51.81539535522461,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 0.2243107706308365,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 0.4464716613292694,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 1.7838181257247925,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 17.912736892700195,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 47.45841979980469,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 56.3084602355957,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 173.33717346191406,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 148.22750854492188,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 133.63565063476562,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 83.65129852294922,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 117.94369506835938,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 94.52413940429688,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 130.43333435058594,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 76.11975860595703,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 158.75192260742188,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 143.72706604003906,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 84.28279876708984,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 116.65055084228516,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 177.1201934814453,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 82.4564437866211,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 137.73019409179688,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 89.97538757324219,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 86.30876159667969,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 61.53449249267578,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 45.22392654418945,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 60.3155517578125,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.26.mlp.up_proj": {
+        "snr": 40.06092071533203,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.27.mlp.up_proj": {
+        "snr": 48.12322998046875,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": 0.08805440366268158,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.26.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.27.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 4.771554470062256,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.46674421429634094,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 1.6167784929275513,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 2.0980119705200195,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 1.4339035749435425,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 1.7446703910827637,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 1.2829725742340088,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 2.2314982414245605,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 1.5125916004180908,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 1.2817912101745605,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 3.3553454875946045,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 1.591347336769104,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 1.1114169359207153,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 1.1536189317703247,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.994098424911499,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 1.484580636024475,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 1.2999093532562256,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 2.1628623008728027,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 1.3842225074768066,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 1.440075159072876,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 1.7816450595855713,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 1.746536135673523,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 1.318993091583252,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 1.7234206199645996,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 2.586996555328369,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 1.6486897468566895,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.26.self_attn.k_proj": {
+        "snr": 1.3349357843399048,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.27.self_attn.k_proj": {
+        "snr": 0.9039687514305115,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.10605750232934952,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.2503393292427063,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.21453581750392914,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.20600366592407227,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.22004099190235138,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.2267625778913498,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.1736888736486435,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.2314220815896988,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.24031606316566467,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.13458871841430664,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.20170633494853973,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.19507651031017303,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.1862162947654724,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.15117767453193665,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.1857745349407196,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.2064860314130783,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.15419450402259827,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.17895667254924774,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.18284623324871063,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.17497135698795319,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.178844153881073,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.16190896928310394,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.19371949136257172,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.14116843044757843,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.14100700616836548,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.14792074263095856,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.26.self_attn.o_proj": {
+        "snr": 0.11953117698431015,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.27.self_attn.o_proj": {
+        "snr": 0.06241385638713837,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.02127065323293209,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.14693336188793182,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.16316214203834534,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.1218630000948906,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.13916714489459991,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.155359148979187,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.1590007096529007,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.1958903819322586,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.22448301315307617,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.20126597583293915,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.1980895698070526,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.2289486974477768,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.22922305762767792,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.21452386677265167,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.24151542782783508,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.21893717348575592,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.2321016639471054,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.24078059196472168,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.22774985432624817,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.20914016664028168,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.22847522795200348,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.2500442862510681,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.2353251725435257,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.20365388691425323,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.21967172622680664,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.2122868150472641,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.26.self_attn.q_proj": {
+        "snr": 0.2415798157453537,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.27.self_attn.q_proj": {
+        "snr": 0.12347634881734848,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 230.88636779785156,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": 22.38136100769043,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 246.59597778320312,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 499.61761474609375,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": 69.18345642089844,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 984.9320068359375,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 64.06214141845703,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": 28.43911361694336,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": 725.1439819335938,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": 63.43681716918945,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": 238.4695587158203,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": 111.88697814941406,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": 686.2830200195312,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": 566.2647705078125,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.26.self_attn.v_proj": {
+        "snr": 4.070064544677734,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.27.self_attn.v_proj": {
+        "snr": 4.3411664962768555,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-7B.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-7B.json
new file mode 100644
index 0000000000..af13ecf9e7
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_Qwen-Qwen2.5-7B.json
@@ -0,0 +1,1022 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.26.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.27.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 10.277782440185547,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 1.2050706148147583,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 19.284534454345703,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 50.16513442993164,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 114.24882507324219,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 215.48194885253906,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 204.39431762695312,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 182.5116729736328,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 74.9266128540039,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 16.474102020263672,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 55.30583572387695,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 16.84047508239746,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 58.62131118774414,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 11.144298553466797,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 65.28057098388672,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 46.701290130615234,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 14.278325080871582,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 23.382247924804688,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 93.8782958984375,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 49.10498809814453,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 1277.5101318359375,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 51.7880859375,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 143.03504943847656,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 35.123931884765625,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 21.403743743896484,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 10.551352500915527,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.26.mlp.down_proj": {
+        "snr": 21.62333869934082,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.27.mlp.down_proj": {
+        "snr": 1541.98681640625,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 0.04497644677758217,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 0.16878646612167358,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 0.336302250623703,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 3.141293525695801,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 9.098686218261719,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 15.89354419708252,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 60.85503387451172,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 57.53098678588867,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 65.77096557617188,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 10.453179359436035,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 26.94801139831543,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 31.111093521118164,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 43.963191986083984,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 20.690765380859375,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 20.47557258605957,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 42.63906478881836,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 22.11542320251465,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 22.590566635131836,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 81.74773406982422,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 20.872997283935547,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 58.32197952270508,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 22.784095764160156,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 16.935768127441406,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 16.830224990844727,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 19.774564743041992,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 27.770675659179688,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.26.mlp.gate_proj": {
+        "snr": 36.714595794677734,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.27.mlp.gate_proj": {
+        "snr": 51.81637191772461,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 0.22425401210784912,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 0.4456978142261505,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 1.7769725322723389,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 17.8966121673584,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 47.43608856201172,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 56.2298698425293,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 173.1498260498047,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 148.02874755859375,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 133.5174560546875,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 83.45183563232422,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 117.88772583007812,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 94.41156768798828,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 130.3107452392578,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 76.04458618164062,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 158.59634399414062,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 143.59596252441406,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 84.2161636352539,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 116.55204010009766,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 176.95449829101562,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 82.37284088134766,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 137.5695343017578,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 89.87335205078125,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 86.1510238647461,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 61.37428665161133,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 45.10757064819336,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 60.16519546508789,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.26.mlp.up_proj": {
+        "snr": 39.96969223022461,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.27.mlp.up_proj": {
+        "snr": 48.04258346557617,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": 0.08800078183412552,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.26.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.27.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 4.764852046966553,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.46627077460289,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 1.6155915260314941,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 2.096365451812744,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 1.431254267692566,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 1.7440669536590576,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 1.2815033197402954,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 2.2301025390625,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 1.5116536617279053,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 1.2699830532073975,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 3.3086464405059814,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 1.59111487865448,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 1.1007944345474243,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 1.163416862487793,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.9935113787651062,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 1.483581304550171,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 1.2992271184921265,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 2.162485122680664,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 1.3841017484664917,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 1.453418493270874,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 1.781678557395935,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 1.7460925579071045,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 1.3188031911849976,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 1.723441243171692,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 2.585094928741455,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 1.6478856801986694,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.26.self_attn.k_proj": {
+        "snr": 1.3221096992492676,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.27.self_attn.k_proj": {
+        "snr": 0.9034463167190552,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.10636883229017258,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.24971255660057068,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.21437697112560272,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.2058248072862625,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.21978946030139923,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.2269466072320938,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.17318543791770935,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.23159846663475037,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.2400084286928177,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.134766086935997,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.20152011513710022,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.19492347538471222,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.18607021868228912,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.15107683837413788,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.18565276265144348,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.20626339316368103,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.1541011780500412,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.1784645915031433,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.18307389318943024,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.17449897527694702,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.1787375956773758,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.161802276968956,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.1931520402431488,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.14108893275260925,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.14064815640449524,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.14790543913841248,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.26.self_attn.o_proj": {
+        "snr": 0.11950570344924927,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.27.self_attn.o_proj": {
+        "snr": 0.062389008700847626,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.02138795144855976,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.14676862955093384,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.16297142207622528,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.12198334187269211,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.13921146094799042,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.15567339956760406,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.1589033454656601,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.195299431681633,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.22430908679962158,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.2011336237192154,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.1982448250055313,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.22880099713802338,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.22898294031620026,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.21394900977611542,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.24130398035049438,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.21905161440372467,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.2319282442331314,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.24004821479320526,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.22754515707492828,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.2086794078350067,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.2290779948234558,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.250373899936676,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.23474709689617157,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.20302507281303406,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.21992310881614685,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.2120121270418167,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.26.self_attn.q_proj": {
+        "snr": 0.24161922931671143,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.27.self_attn.q_proj": {
+        "snr": 0.12337693572044373,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 231.07347106933594,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": 22.34870719909668,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 246.30386352539062,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 499.5611572265625,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": 69.09609985351562,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 983.3341674804688,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 64.04925537109375,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": 28.41021728515625,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": 724.2736206054688,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": 63.35670852661133,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": Infinity,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": 238.2569122314453,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": 111.78319549560547,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": 687.0054931640625,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": 565.3272705078125,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.26.self_attn.v_proj": {
+        "snr": 4.064513683319092,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.27.self_attn.v_proj": {
+        "snr": 4.335177421569824,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_google-gemma-2-2b.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_google-gemma-2-2b.json
new file mode 100644
index 0000000000..a2d6f0457d
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_google-gemma-2-2b.json
@@ -0,0 +1,1158 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": 4.538210391998291,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 7.746472358703613,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 4.3358893394470215,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 26.88057518005371,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 8.699942588806152,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 32.808380126953125,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 10.831522941589355,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 18.843679428100586,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 9.348078727722168,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 7.061270236968994,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 5.454320907592773,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 7.386133193969727,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 6.648562908172607,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 5.853652477264404,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 8.570493698120117,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 13.120837211608887,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 14.780969619750977,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 6.953134059906006,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 12.589436531066895,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 8.844094276428223,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 7.598869800567627,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 11.293925285339355,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 9.384604454040527,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 12.12533187866211,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 11.217570304870605,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 14.197714805603027,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 12.449926376342773,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 16.885862350463867,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 23.410266876220703,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 22.57662582397461,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 17.29996681213379,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 11.718637466430664,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 6.376136779785156,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 6.794021129608154,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 3.2425343990325928,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 2.368421792984009,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 3.3193087577819824,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 3.9515960216522217,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 3.2761318683624268,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 4.026322841644287,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 3.415473699569702,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 3.3418092727661133,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 3.6233012676239014,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 3.2199010848999023,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 3.6848936080932617,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 3.4439642429351807,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 3.7366604804992676,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 4.262336254119873,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 4.333253860473633,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 3.640247344970703,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 4.2978034019470215,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 4.339972496032715,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 3.8502564430236816,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 28.129924774169922,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 41.49960708618164,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 125.47801971435547,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 119.93355560302734,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 162.62631225585938,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 32.36909484863281,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 49.10078430175781,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 28.541580200195312,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 14.764090538024902,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 16.5697078704834,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 19.26059913635254,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 15.082040786743164,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 15.5792875289917,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 9.84595012664795,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 11.506875991821289,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 21.507600784301758,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 15.110466957092285,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 27.062183380126953,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 16.40383529663086,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 13.117464065551758,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 11.393353462219238,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 10.791608810424805,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 7.512388706207275,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 9.889434814453125,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 7.587779521942139,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 4.561068058013916,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": 4.538210391998291,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.1.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.2.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.3.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.4.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.5.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.6.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.7.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.8.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.9.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.10.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.11.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.12.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.13.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.14.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.15.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.16.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.17.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.18.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.19.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.20.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.21.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.22.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.23.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.24.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.25.post_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "post_feedforward_layernorm"
+    },
+    "model.layers.0.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.1.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.2.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.3.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.4.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.5.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.6.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.7.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.8.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.9.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.10.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.11.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.12.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.13.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.14.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.15.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.16.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.17.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.18.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.19.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.20.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.21.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.22.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.23.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.24.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.25.pre_feedforward_layernorm": {
+        "snr": Infinity,
+        "type": "pre_feedforward_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 0.5685535073280334,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 1.060130000114441,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 1.0735561847686768,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 1.0217311382293701,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 0.9687430262565613,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 0.8411160111427307,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 0.936741054058075,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 0.7236003279685974,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 0.9032857418060303,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 0.7513307929039001,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 0.6875415444374084,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 0.6611058712005615,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 0.8023670315742493,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 0.7188767194747925,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.7930117249488831,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 0.9076258540153503,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 0.7295113801956177,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 0.898467481136322,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 0.9652048945426941,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 0.9855819344520569,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 1.2863355875015259,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 1.116607904434204,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 0.7438228130340576,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 0.8499895334243774,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 0.7764042019844055,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 0.7127887606620789,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.2556447386741638,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.2930974066257477,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.27571651339530945,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.280631959438324,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.2958097755908966,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.3072899580001831,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.31374114751815796,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.2903076410293579,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.2625811696052551,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.2306082546710968,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.24869701266288757,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.2556127905845642,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.28926730155944824,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.25355643033981323,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.23122912645339966,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.28772857785224915,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.22682352364063263,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.2558597922325134,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.1773315966129303,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.2106105089187622,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.2008877396583557,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.1973956972360611,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.25533634424209595,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.20066529512405396,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.18342143297195435,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.3224162459373474,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.2074502408504486,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.33233126997947693,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.3586291968822479,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.2850974202156067,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.37816473841667175,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.31616899371147156,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.4988365173339844,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.4238639175891876,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.2674674689769745,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.34524214267730713,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.4472109377384186,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.41363632678985596,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.44623735547065735,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.4404333531856537,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.5200268626213074,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.4320363700389862,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.46235284209251404,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.47477203607559204,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.4001321494579315,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.42365774512290955,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.37057873606681824,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.3990235924720764,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.35094162821769714,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.35721710324287415,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.2812618315219879,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.19463211297988892,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": 1.3365743160247803,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 2.402009963989258,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": 3.8695859909057617,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": 4.117948055267334,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": 5.651231288909912,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 2.720799446105957,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 1.4446897506713867,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": 4.497112274169922,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": 1.7241870164871216,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": 1.7104988098144531,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": 1.4231206178665161,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": 2.1643989086151123,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 1.5254249572753906,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": 2.3788745403289795,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 3.4155967235565186,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": 4.623549938201904,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": 1.5291141271591187,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": 3.9934189319610596,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": 9.035382270812988,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": 5.8578925132751465,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": 3.759958505630493,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": 4.558528900146484,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": 0.9163281917572021,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": 2.564377546310425,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": 3.689103841781616,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": 5.6444854736328125,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-1B-Instruct.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-1B-Instruct.json
new file mode 100644
index 0000000000..c5b5ed6ed8
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-1B-Instruct.json
@@ -0,0 +1,590 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 70.0594253540039,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 11.135851860046387,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 7.035482883453369,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 6.422532081604004,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 5.748020172119141,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 3.885556697845459,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 3.4336745738983154,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 2.791595935821533,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 5.36277961730957,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 4.459208011627197,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 6.272170066833496,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 5.264761447906494,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 4.324735641479492,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 3.878648042678833,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 2.9773054122924805,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 4.471445560455322,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 25.227100372314453,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 6.58299446105957,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 3.4688243865966797,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 1.555246114730835,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 0.7770601511001587,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 0.6239906549453735,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 0.6440379023551941,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 0.5120116472244263,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 0.6544050574302673,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 0.5381016731262207,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 0.622873842716217,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 0.9361700415611267,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 1.475605845451355,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 1.608325719833374,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 1.0720024108886719,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 0.7111338973045349,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 28.431896209716797,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 15.546019554138184,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 23.048023223876953,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 25.790977478027344,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 18.552549362182617,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 8.85106372833252,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 10.653799057006836,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 7.365357875823975,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 11.98373794555664,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 8.04493236541748,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 8.523039817810059,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 5.381742477416992,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 3.9845118522644043,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 3.4893221855163574,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 1.764201045036316,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 0.9730708599090576,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": Infinity,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 0.11727584153413773,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.24786807596683502,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 0.36378130316734314,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 0.2983120381832123,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 0.33789733052253723,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 0.29155924916267395,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 0.2537297010421753,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 0.28204113245010376,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 0.2776711583137512,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 0.2927376627922058,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 0.31486213207244873,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 0.32363659143447876,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 0.31382912397384644,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 0.4635234773159027,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.25379249453544617,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 0.2628238797187805,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.27602291107177734,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.2149604707956314,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.2540294826030731,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.27978822588920593,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.3121289908885956,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.35037684440612793,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.366205096244812,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.3692712187767029,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.3301038146018982,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.3003396987915039,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.30804169178009033,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.28501132130622864,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.2171541005373001,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.19183959066867828,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.19215913116931915,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.25486502051353455,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.03850084915757179,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.0713055431842804,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.07948919385671616,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.08047746121883392,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.0852593332529068,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.09794823825359344,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.09627152234315872,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.11065381020307541,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.12031875550746918,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.09804573655128479,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.10897502303123474,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.09267337620258331,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.08803492039442062,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.0902542844414711,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.10154066979885101,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.09083802253007889,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": 2.842210054397583,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 10.59461498260498,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": 8.993025779724121,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": 62.567787170410156,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": 23.80082893371582,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 7.957369804382324,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 12.01815414428711,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": 5.095500469207764,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": 11.719332695007324,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": 555.0869750976562,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": 22.95538330078125,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": 30.042158126831055,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 9.577271461486816,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": 18.176361083984375,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 1.5695856809616089,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": 2.7235565185546875,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-1B.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-1B.json
new file mode 100644
index 0000000000..b84594fd7b
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-1B.json
@@ -0,0 +1,590 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 57.09797286987305,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 9.538983345031738,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 6.227016925811768,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 5.660686492919922,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 5.178432464599609,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 3.5638349056243896,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 3.0918056964874268,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 2.456392288208008,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 4.525328636169434,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 3.9409055709838867,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 5.447249412536621,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 4.807600975036621,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 3.915374517440796,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 3.4820363521575928,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 2.6045074462890625,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 3.7237701416015625,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 22.160131454467773,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 6.072206020355225,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 3.2467362880706787,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 1.4111896753311157,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 0.7405938506126404,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 0.5916463136672974,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 0.6149423718452454,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 0.48369669914245605,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 0.6047574877738953,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 0.5092479586601257,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 0.5999670624732971,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 0.8980127573013306,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 1.4252448081970215,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 1.509937047958374,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 1.0066585540771484,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 0.6413647532463074,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 26.08852195739746,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 13.382951736450195,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 20.088768005371094,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 23.0632381439209,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 16.07433319091797,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 8.00507640838623,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 9.538354873657227,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 6.286602973937988,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 10.092820167541504,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 7.193963527679443,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 7.320116996765137,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 4.8728532791137695,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 3.596583366394043,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 3.166161298751831,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 1.5600818395614624,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 0.8726214170455933,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": Infinity,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 0.1154392883181572,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.24299409985542297,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 0.3624322712421417,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 0.29509487748146057,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 0.32953736186027527,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 0.2908833622932434,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 0.2488437294960022,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 0.27847856283187866,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 0.27143892645835876,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 0.28804272413253784,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 0.31197959184646606,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 0.3203586935997009,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 0.30905747413635254,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 0.46828722953796387,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.24205778539180756,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 0.2559327781200409,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.2638678550720215,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.21109595894813538,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.24751724302768707,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.2728094160556793,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.3001374304294586,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.33903488516807556,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.3530929982662201,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.36753255128860474,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.3373180329799652,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.2970578670501709,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.3076324760913849,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.2766900658607483,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.20973259210586548,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.18185566365718842,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.18329747021198273,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.2437991499900818,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.038040731102228165,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.0707998052239418,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.0787411704659462,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.08089710026979446,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.08591937273740768,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.09852176159620285,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.09690654277801514,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.11181341856718063,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.12042108923196793,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.09799323976039886,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.10901063680648804,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.09307146072387695,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.0880950540304184,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.08886399120092392,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.09955056011676788,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.08929339051246643,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": 2.5501928329467773,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 9.449499130249023,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": 7.9920830726623535,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": 50.69462585449219,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": 19.083511352539062,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 7.21597146987915,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 11.27744197845459,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": 4.579711437225342,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": 10.940719604492188,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": 553.4417724609375,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": 20.59434700012207,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": 26.636865615844727,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 8.614749908447266,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": 17.722007751464844,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 1.48500657081604,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": 2.5776851177215576,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-3B-Instruct.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-3B-Instruct.json
new file mode 100644
index 0000000000..7764eecd9c
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-3B-Instruct.json
@@ -0,0 +1,1022 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.26.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.27.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 2.306217670440674,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 2.2327167987823486,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 1.4501516819000244,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 1.363667607307434,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 1.4520279169082642,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 1.4664665460586548,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 1.4122329950332642,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 1.0504299402236938,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 0.9837537407875061,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 0.8659006357192993,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 0.7936406135559082,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 0.9000886678695679,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 1.1559213399887085,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 1.3054672479629517,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 1.196791410446167,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 1.3163655996322632,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 1.3388997316360474,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 1.592497706413269,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 1.5399079322814941,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 1.5683293342590332,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 1.4739630222320557,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 1.2608393430709839,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 1.2087301015853882,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 1.1851829290390015,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 1.0537594556808472,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 1.1649317741394043,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.26.mlp.down_proj": {
+        "snr": 1.2376821041107178,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.27.mlp.down_proj": {
+        "snr": 1.147771954536438,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 0.9385462999343872,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 0.8528683185577393,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 0.761657178401947,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 0.6598325371742249,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 0.44578588008880615,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 0.4053060710430145,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 0.3588462769985199,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 0.35667839646339417,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 0.3106202781200409,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 0.2821919322013855,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 0.29143741726875305,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 0.29830989241600037,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 0.2862427532672882,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 0.2797018587589264,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 0.2679217755794525,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 0.2782425880432129,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 0.3503592610359192,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 0.3968559205532074,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 0.4318574070930481,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 0.4693693220615387,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 0.5051979422569275,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 0.5675955414772034,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 0.5861843824386597,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 0.4759417772293091,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 0.38529056310653687,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 0.3180919587612152,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.26.mlp.gate_proj": {
+        "snr": 0.2695689797401428,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.27.mlp.gate_proj": {
+        "snr": 0.21765239536762238,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 1.4919718503952026,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 1.7983858585357666,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 2.1709094047546387,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 2.751326560974121,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 3.063521385192871,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 2.4026951789855957,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 2.3890223503112793,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 2.3861353397369385,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 2.0745043754577637,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 1.8550645112991333,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 1.6184496879577637,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 1.9287559986114502,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 1.7427546977996826,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 1.9872609376907349,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 2.0224087238311768,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 1.7851638793945312,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 1.7160604000091553,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 1.6870195865631104,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 1.6585396528244019,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 1.5509096384048462,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 1.4310423135757446,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 1.5009464025497437,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 1.4866929054260254,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 1.332513689994812,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 1.073512077331543,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 0.7472100257873535,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.26.mlp.up_proj": {
+        "snr": 0.4880162179470062,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.27.mlp.up_proj": {
+        "snr": 0.2527681589126587,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": Infinity,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.26.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.27.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 0.08262510597705841,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.1441459059715271,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 0.21418076753616333,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 0.22496014833450317,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 0.23101305961608887,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 0.23644132912158966,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 0.23666173219680786,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 0.19791515171527863,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 0.22062039375305176,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 0.21218444406986237,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 0.24218571186065674,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 0.21870514750480652,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 0.22160987555980682,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 0.22726823389530182,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.20256873965263367,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 0.24100735783576965,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 0.23794010281562805,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 0.2913324534893036,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 0.28093472123146057,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 0.31062793731689453,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 0.2942160367965698,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 0.28014805912971497,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 0.3512437045574188,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 0.2837671637535095,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 0.2960015535354614,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 0.5086414813995361,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.26.self_attn.k_proj": {
+        "snr": 0.24054698646068573,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.27.self_attn.k_proj": {
+        "snr": 0.247616246342659,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.18390265107154846,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.14759540557861328,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.15726515650749207,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.16903570294380188,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.17953157424926758,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.2351229190826416,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.22804339230060577,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.24786025285720825,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.21847976744174957,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.2092437595129013,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.23278094828128815,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.20468176901340485,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.2353818416595459,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.2702614367008209,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.19177420437335968,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.18293911218643188,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.20286045968532562,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.20763878524303436,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.190629780292511,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.22044304013252258,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.21491236984729767,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.23289704322814941,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.21457163989543915,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.1949365884065628,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.1606779545545578,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.13892440497875214,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.26.self_attn.o_proj": {
+        "snr": 0.1407029926776886,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.27.self_attn.o_proj": {
+        "snr": 0.16027599573135376,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.0534212663769722,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.06873775273561478,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.07522258907556534,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.06616844981908798,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.06809444725513458,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.0758095383644104,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.07800278812646866,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.07535763084888458,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.09488166123628616,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.09709945321083069,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.09381720423698425,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.08205580711364746,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.10723169893026352,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.10166660696268082,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.08822792023420334,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.0814041867852211,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.07586681097745895,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.07040166854858398,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.0728282704949379,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.06912193447351456,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.06646180897951126,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.06960278004407883,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.06566876918077469,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.07412787526845932,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.07131384313106537,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.07768437266349792,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.26.self_attn.q_proj": {
+        "snr": 0.0809575766324997,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.27.self_attn.q_proj": {
+        "snr": 0.06796683371067047,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": 1.4029983282089233,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 3.123720169067383,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": 2.4177253246307373,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": 5.588768005371094,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": 4.395562648773193,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 3.2982685565948486,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 3.2798449993133545,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": 2.109200954437256,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": 3.229325532913208,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": 1.7349927425384521,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": 1.5926740169525146,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": 1.9097802639007568,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 2.5654332637786865,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": 3.536489963531494,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 8.366667747497559,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": 7.348303318023682,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": 2.815748691558838,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": 4.048776149749756,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": 4.426101207733154,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": 7.098501682281494,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": 3.700288772583008,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": 2.1859049797058105,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": 3.6953284740448,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": 11.148802757263184,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": 2.4171905517578125,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": 4.404144287109375,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.26.self_attn.v_proj": {
+        "snr": 2.340604782104492,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.27.self_attn.v_proj": {
+        "snr": 3.284160614013672,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-3B.json b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-3B.json
new file mode 100644
index 0000000000..131a5b79b7
--- /dev/null
+++ b/src/axolotl/integrations/spectrum/model_snr_results/snr_results_meta-llama-Llama-3.2-3B.json
@@ -0,0 +1,1022 @@
+{
+    "model.layers.0.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.1.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.2.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.3.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.4.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.5.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.6.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.7.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.8.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.9.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.10.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.11.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.12.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.13.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.14.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.15.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.16.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.17.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.18.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.19.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.20.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.21.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.22.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.23.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.24.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.25.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.26.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "model.layers.27.input_layernorm": {
+        "snr": Infinity,
+        "type": "input_layernorm"
+    },
+    "lm_head": {
+        "snr": Infinity,
+        "type": "lm_head"
+    },
+    "model.layers.0.mlp.down_proj": {
+        "snr": 2.364603281021118,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.1.mlp.down_proj": {
+        "snr": 2.229910373687744,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.2.mlp.down_proj": {
+        "snr": 1.4312117099761963,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.3.mlp.down_proj": {
+        "snr": 1.3216407299041748,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.4.mlp.down_proj": {
+        "snr": 1.4183496236801147,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.5.mlp.down_proj": {
+        "snr": 1.4453660249710083,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.6.mlp.down_proj": {
+        "snr": 1.4030662775039673,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.7.mlp.down_proj": {
+        "snr": 1.042332649230957,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.8.mlp.down_proj": {
+        "snr": 0.9530982375144958,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.9.mlp.down_proj": {
+        "snr": 0.849862277507782,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.10.mlp.down_proj": {
+        "snr": 0.7704945206642151,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.11.mlp.down_proj": {
+        "snr": 0.8871145844459534,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.12.mlp.down_proj": {
+        "snr": 1.1408143043518066,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.13.mlp.down_proj": {
+        "snr": 1.2769343852996826,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.14.mlp.down_proj": {
+        "snr": 1.1703068017959595,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.15.mlp.down_proj": {
+        "snr": 1.2794467210769653,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.16.mlp.down_proj": {
+        "snr": 1.3154453039169312,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.17.mlp.down_proj": {
+        "snr": 1.5596749782562256,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.18.mlp.down_proj": {
+        "snr": 1.4949405193328857,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.19.mlp.down_proj": {
+        "snr": 1.5329173803329468,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.20.mlp.down_proj": {
+        "snr": 1.4396660327911377,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.21.mlp.down_proj": {
+        "snr": 1.217085838317871,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.22.mlp.down_proj": {
+        "snr": 1.150472640991211,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.23.mlp.down_proj": {
+        "snr": 1.1166225671768188,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.24.mlp.down_proj": {
+        "snr": 0.9966591000556946,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.25.mlp.down_proj": {
+        "snr": 1.0938347578048706,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.26.mlp.down_proj": {
+        "snr": 1.1505423784255981,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.27.mlp.down_proj": {
+        "snr": 1.1156749725341797,
+        "type": "mlp.down_proj"
+    },
+    "model.layers.0.mlp.gate_proj": {
+        "snr": 0.9329171776771545,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.1.mlp.gate_proj": {
+        "snr": 0.8513413667678833,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.2.mlp.gate_proj": {
+        "snr": 0.7584061026573181,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.3.mlp.gate_proj": {
+        "snr": 0.65835040807724,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.4.mlp.gate_proj": {
+        "snr": 0.436420738697052,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.5.mlp.gate_proj": {
+        "snr": 0.39712461829185486,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.6.mlp.gate_proj": {
+        "snr": 0.3530206084251404,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.7.mlp.gate_proj": {
+        "snr": 0.34982794523239136,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.8.mlp.gate_proj": {
+        "snr": 0.30338960886001587,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.9.mlp.gate_proj": {
+        "snr": 0.27569833397865295,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.10.mlp.gate_proj": {
+        "snr": 0.28934162855148315,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.11.mlp.gate_proj": {
+        "snr": 0.2929173707962036,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.12.mlp.gate_proj": {
+        "snr": 0.28263387084007263,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.13.mlp.gate_proj": {
+        "snr": 0.27778616547584534,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.14.mlp.gate_proj": {
+        "snr": 0.26527827978134155,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.15.mlp.gate_proj": {
+        "snr": 0.27635642886161804,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.16.mlp.gate_proj": {
+        "snr": 0.35072311758995056,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.17.mlp.gate_proj": {
+        "snr": 0.4002636671066284,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.18.mlp.gate_proj": {
+        "snr": 0.4319891333580017,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.19.mlp.gate_proj": {
+        "snr": 0.47527065873146057,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.20.mlp.gate_proj": {
+        "snr": 0.5112077593803406,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.21.mlp.gate_proj": {
+        "snr": 0.5749644637107849,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.22.mlp.gate_proj": {
+        "snr": 0.5967603921890259,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.23.mlp.gate_proj": {
+        "snr": 0.48045310378074646,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.24.mlp.gate_proj": {
+        "snr": 0.3838970363140106,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.25.mlp.gate_proj": {
+        "snr": 0.3108249604701996,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.26.mlp.gate_proj": {
+        "snr": 0.26704445481300354,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.27.mlp.gate_proj": {
+        "snr": 0.20953254401683807,
+        "type": "mlp.gate_proj"
+    },
+    "model.layers.0.mlp.up_proj": {
+        "snr": 1.5084924697875977,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.1.mlp.up_proj": {
+        "snr": 1.7789595127105713,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.2.mlp.up_proj": {
+        "snr": 2.1431775093078613,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.3.mlp.up_proj": {
+        "snr": 2.762744903564453,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.4.mlp.up_proj": {
+        "snr": 3.0324745178222656,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.5.mlp.up_proj": {
+        "snr": 2.3884809017181396,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.6.mlp.up_proj": {
+        "snr": 2.388005256652832,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.7.mlp.up_proj": {
+        "snr": 2.339340925216675,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.8.mlp.up_proj": {
+        "snr": 2.0497021675109863,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.9.mlp.up_proj": {
+        "snr": 1.822119116783142,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.10.mlp.up_proj": {
+        "snr": 1.600373387336731,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.11.mlp.up_proj": {
+        "snr": 1.9298171997070312,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.12.mlp.up_proj": {
+        "snr": 1.728783369064331,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.13.mlp.up_proj": {
+        "snr": 1.965298056602478,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.14.mlp.up_proj": {
+        "snr": 2.023681640625,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.15.mlp.up_proj": {
+        "snr": 1.7721818685531616,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.16.mlp.up_proj": {
+        "snr": 1.7068361043930054,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.17.mlp.up_proj": {
+        "snr": 1.6673219203948975,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.18.mlp.up_proj": {
+        "snr": 1.6240718364715576,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.19.mlp.up_proj": {
+        "snr": 1.5169662237167358,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.20.mlp.up_proj": {
+        "snr": 1.4018198251724243,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.21.mlp.up_proj": {
+        "snr": 1.4556466341018677,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.22.mlp.up_proj": {
+        "snr": 1.4304454326629639,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.23.mlp.up_proj": {
+        "snr": 1.2785290479660034,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.24.mlp.up_proj": {
+        "snr": 1.023495078086853,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.25.mlp.up_proj": {
+        "snr": 0.6992124915122986,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.26.mlp.up_proj": {
+        "snr": 0.4549211859703064,
+        "type": "mlp.up_proj"
+    },
+    "model.layers.27.mlp.up_proj": {
+        "snr": 0.23889905214309692,
+        "type": "mlp.up_proj"
+    },
+    "model.embed_tokens": {
+        "snr": Infinity,
+        "type": "model.embed_tokens"
+    },
+    "model.norm": {
+        "snr": Infinity,
+        "type": "model.norm"
+    },
+    "model.layers.0.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.1.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.2.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.3.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.4.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.5.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.6.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.7.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.8.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.9.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.10.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.11.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.12.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.13.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.14.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.15.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.16.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.17.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.18.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.19.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.20.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.21.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.22.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.23.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.24.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.25.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.26.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.27.post_attention_layernorm": {
+        "snr": Infinity,
+        "type": "post_attention_layernorm"
+    },
+    "model.layers.0.self_attn.k_proj": {
+        "snr": 0.08150045573711395,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.1.self_attn.k_proj": {
+        "snr": 0.1428358554840088,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.2.self_attn.k_proj": {
+        "snr": 0.2096949815750122,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.3.self_attn.k_proj": {
+        "snr": 0.22633400559425354,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.4.self_attn.k_proj": {
+        "snr": 0.2293967455625534,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.5.self_attn.k_proj": {
+        "snr": 0.23336802423000336,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.6.self_attn.k_proj": {
+        "snr": 0.23429904878139496,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.7.self_attn.k_proj": {
+        "snr": 0.19610290229320526,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.8.self_attn.k_proj": {
+        "snr": 0.2163258045911789,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.9.self_attn.k_proj": {
+        "snr": 0.21039333939552307,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.10.self_attn.k_proj": {
+        "snr": 0.23533931374549866,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.11.self_attn.k_proj": {
+        "snr": 0.21457058191299438,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.12.self_attn.k_proj": {
+        "snr": 0.21686571836471558,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.13.self_attn.k_proj": {
+        "snr": 0.22398065030574799,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.14.self_attn.k_proj": {
+        "snr": 0.20160657167434692,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.15.self_attn.k_proj": {
+        "snr": 0.23705022037029266,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.16.self_attn.k_proj": {
+        "snr": 0.23254962265491486,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.17.self_attn.k_proj": {
+        "snr": 0.2892642617225647,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.18.self_attn.k_proj": {
+        "snr": 0.27587130665779114,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.19.self_attn.k_proj": {
+        "snr": 0.30891212821006775,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.20.self_attn.k_proj": {
+        "snr": 0.28997519612312317,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.21.self_attn.k_proj": {
+        "snr": 0.27534863352775574,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.22.self_attn.k_proj": {
+        "snr": 0.35139667987823486,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.23.self_attn.k_proj": {
+        "snr": 0.2773109972476959,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.24.self_attn.k_proj": {
+        "snr": 0.2853511571884155,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.25.self_attn.k_proj": {
+        "snr": 0.5030262470245361,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.26.self_attn.k_proj": {
+        "snr": 0.2317112237215042,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.27.self_attn.k_proj": {
+        "snr": 0.24419328570365906,
+        "type": "self_attn.k_proj"
+    },
+    "model.layers.0.self_attn.o_proj": {
+        "snr": 0.17767645418643951,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.1.self_attn.o_proj": {
+        "snr": 0.14102177321910858,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.2.self_attn.o_proj": {
+        "snr": 0.1523692011833191,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.3.self_attn.o_proj": {
+        "snr": 0.16522075235843658,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.4.self_attn.o_proj": {
+        "snr": 0.17483487725257874,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.5.self_attn.o_proj": {
+        "snr": 0.227921262383461,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.6.self_attn.o_proj": {
+        "snr": 0.2196175903081894,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.7.self_attn.o_proj": {
+        "snr": 0.24270132184028625,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.8.self_attn.o_proj": {
+        "snr": 0.2118290364742279,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.9.self_attn.o_proj": {
+        "snr": 0.20525991916656494,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.10.self_attn.o_proj": {
+        "snr": 0.22847208380699158,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.11.self_attn.o_proj": {
+        "snr": 0.19665324687957764,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.12.self_attn.o_proj": {
+        "snr": 0.23233532905578613,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.13.self_attn.o_proj": {
+        "snr": 0.2624332308769226,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.14.self_attn.o_proj": {
+        "snr": 0.1868327558040619,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.15.self_attn.o_proj": {
+        "snr": 0.17706255614757538,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.16.self_attn.o_proj": {
+        "snr": 0.19422705471515656,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.17.self_attn.o_proj": {
+        "snr": 0.2000615894794464,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.18.self_attn.o_proj": {
+        "snr": 0.1874573826789856,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.19.self_attn.o_proj": {
+        "snr": 0.21297843754291534,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.20.self_attn.o_proj": {
+        "snr": 0.2100859135389328,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.21.self_attn.o_proj": {
+        "snr": 0.22561520338058472,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.22.self_attn.o_proj": {
+        "snr": 0.20994484424591064,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.23.self_attn.o_proj": {
+        "snr": 0.18978221714496613,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.24.self_attn.o_proj": {
+        "snr": 0.1571759581565857,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.25.self_attn.o_proj": {
+        "snr": 0.1349896937608719,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.26.self_attn.o_proj": {
+        "snr": 0.1368866115808487,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.27.self_attn.o_proj": {
+        "snr": 0.1571887582540512,
+        "type": "self_attn.o_proj"
+    },
+    "model.layers.0.self_attn.q_proj": {
+        "snr": 0.05295897275209427,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.1.self_attn.q_proj": {
+        "snr": 0.06835605204105377,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.2.self_attn.q_proj": {
+        "snr": 0.0746372863650322,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.3.self_attn.q_proj": {
+        "snr": 0.06615085154771805,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.4.self_attn.q_proj": {
+        "snr": 0.06788161396980286,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.5.self_attn.q_proj": {
+        "snr": 0.07514483481645584,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.6.self_attn.q_proj": {
+        "snr": 0.07777862250804901,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.7.self_attn.q_proj": {
+        "snr": 0.07534090429544449,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.8.self_attn.q_proj": {
+        "snr": 0.09494179487228394,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.9.self_attn.q_proj": {
+        "snr": 0.09699037671089172,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.10.self_attn.q_proj": {
+        "snr": 0.09426294267177582,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.11.self_attn.q_proj": {
+        "snr": 0.08260341733694077,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.12.self_attn.q_proj": {
+        "snr": 0.10650420933961868,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.13.self_attn.q_proj": {
+        "snr": 0.10250870138406754,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.14.self_attn.q_proj": {
+        "snr": 0.08775162696838379,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.15.self_attn.q_proj": {
+        "snr": 0.08071447163820267,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.16.self_attn.q_proj": {
+        "snr": 0.07530857622623444,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.17.self_attn.q_proj": {
+        "snr": 0.06964966654777527,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.18.self_attn.q_proj": {
+        "snr": 0.07150755077600479,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.19.self_attn.q_proj": {
+        "snr": 0.0676807165145874,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.20.self_attn.q_proj": {
+        "snr": 0.06511317938566208,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.21.self_attn.q_proj": {
+        "snr": 0.06773187220096588,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.22.self_attn.q_proj": {
+        "snr": 0.06400436162948608,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.23.self_attn.q_proj": {
+        "snr": 0.0726117342710495,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.24.self_attn.q_proj": {
+        "snr": 0.06882446259260178,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.25.self_attn.q_proj": {
+        "snr": 0.07506493479013443,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.26.self_attn.q_proj": {
+        "snr": 0.07797915488481522,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.27.self_attn.q_proj": {
+        "snr": 0.06680692732334137,
+        "type": "self_attn.q_proj"
+    },
+    "model.layers.0.self_attn.v_proj": {
+        "snr": 1.326789379119873,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.1.self_attn.v_proj": {
+        "snr": 3.043806791305542,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.2.self_attn.v_proj": {
+        "snr": 2.295107841491699,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.3.self_attn.v_proj": {
+        "snr": 5.2584614753723145,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.4.self_attn.v_proj": {
+        "snr": 4.038785934448242,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.5.self_attn.v_proj": {
+        "snr": 3.0907773971557617,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.6.self_attn.v_proj": {
+        "snr": 3.114994525909424,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.7.self_attn.v_proj": {
+        "snr": 1.9747973680496216,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.8.self_attn.v_proj": {
+        "snr": 3.0469374656677246,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.9.self_attn.v_proj": {
+        "snr": 1.602966547012329,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.10.self_attn.v_proj": {
+        "snr": 1.489019513130188,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.11.self_attn.v_proj": {
+        "snr": 1.7490826845169067,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.12.self_attn.v_proj": {
+        "snr": 2.451310396194458,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.13.self_attn.v_proj": {
+        "snr": 3.250821590423584,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.14.self_attn.v_proj": {
+        "snr": 7.944663047790527,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.15.self_attn.v_proj": {
+        "snr": 7.013208389282227,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.16.self_attn.v_proj": {
+        "snr": 2.68644118309021,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.17.self_attn.v_proj": {
+        "snr": 3.9063122272491455,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.18.self_attn.v_proj": {
+        "snr": 4.1816816329956055,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.19.self_attn.v_proj": {
+        "snr": 6.794488906860352,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.20.self_attn.v_proj": {
+        "snr": 3.401334285736084,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.21.self_attn.v_proj": {
+        "snr": 2.051994562149048,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.22.self_attn.v_proj": {
+        "snr": 3.614379405975342,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.23.self_attn.v_proj": {
+        "snr": 11.180968284606934,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.24.self_attn.v_proj": {
+        "snr": 2.3629775047302246,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.25.self_attn.v_proj": {
+        "snr": 4.137593746185303,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.26.self_attn.v_proj": {
+        "snr": 2.3465518951416016,
+        "type": "self_attn.v_proj"
+    },
+    "model.layers.27.self_attn.v_proj": {
+        "snr": 3.10064697265625,
+        "type": "self_attn.v_proj"
+    }
+}
diff --git a/src/axolotl/integrations/swanlab/README.md b/src/axolotl/integrations/swanlab/README.md
new file mode 100644
index 0000000000..eff7d4a5ab
--- /dev/null
+++ b/src/axolotl/integrations/swanlab/README.md
@@ -0,0 +1,1284 @@
+# SwanLab Integration for Axolotl
+
+SwanLab is an open-source, lightweight AI experiment tracking and visualization tool that provides a platform for tracking, recording, comparing, and collaborating on experiments.
+
+This integration enables seamless experiment tracking and visualization of Axolotl training runs using SwanLab.
+
+## Features
+
+- 📊 **Automatic Metrics Logging**: Training loss, learning rate, and other metrics are automatically logged
+- 🎯 **Hyperparameter Tracking**: Model configuration and training parameters are tracked
+- 📈 **Real-time Visualization**: Monitor training progress in real-time through SwanLab dashboard
+- ☁️ **Cloud & Local Support**: Works in both cloud-synced and offline modes
+- 🔄 **Experiment Comparison**: Compare multiple training runs easily
+- 🤝 **Team Collaboration**: Share experiments with team members
+- 🎭 **RLHF Completion Logging**: Automatically log model outputs during DPO/KTO/ORPO/GRPO training for qualitative analysis
+- ⚡ **Performance Profiling**: Built-in profiling decorators to measure and optimize training performance
+- 🔔 **Lark Notifications**: Send real-time training updates to team chat (Feishu/Lark integration)
+
+## Installation
+
+```bash
+pip install swanlab
+```
+
+## Quick Start
+
+### 1. Register for SwanLab (Optional for cloud mode)
+
+If you want to use cloud sync features, register at [https://swanlab.cn](https://swanlab.cn) to get your API key.
+
+### 2. Configure Axolotl Config File
+
+Add SwanLab configuration to your Axolotl YAML config:
+
+```yaml
+# Enable SwanLab plugin
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+# SwanLab configuration
+use_swanlab: true
+swanlab_project: my-llm-project
+swanlab_experiment_name: qwen-finetune-v1
+swanlab_mode: cloud  # Options: cloud, local, offline, disabled
+swanlab_workspace: my-team  # Optional: organization name
+swanlab_api_key: YOUR_API_KEY  # Optional: can also use env var SWANLAB_API_KEY
+```
+
+### 3. Run Training
+
+```bash
+# Set API key via environment variable (recommended)
+export SWANLAB_API_KEY=your-api-key-here
+
+# Or login once
+swanlab login
+
+# Run training as usual
+accelerate launch -m axolotl.cli.train your-config.yaml
+```
+
+## Configuration Options
+
+### Basic Configuration
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `use_swanlab` | bool | `false` | Enable SwanLab tracking |
+| `swanlab_project` | str | `None` | Project name (required) |
+| `swanlab_experiment_name` | str | `None` | Experiment name |
+| `swanlab_description` | str | `None` | Experiment description |
+| `swanlab_mode` | str | `cloud` | Sync mode: `cloud`, `local`, `offline`, `disabled` |
+
+### Advanced Configuration
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `swanlab_workspace` | str | `None` | Workspace/organization name |
+| `swanlab_api_key` | str | `None` | API key (prefer env var) |
+| `swanlab_web_host` | str | `None` | Private deployment web host |
+| `swanlab_api_host` | str | `None` | Private deployment API host |
+| `swanlab_log_model` | bool | `false` | Log model checkpoints (coming soon) |
+| `swanlab_lark_webhook_url` | str | `None` | Lark (Feishu) webhook URL for team notifications |
+| `swanlab_lark_secret` | str | `None` | Lark webhook HMAC secret for authentication |
+| `swanlab_log_completions` | bool | `true` | Enable RLHF completion table logging (DPO/KTO/ORPO/GRPO) |
+| `swanlab_completion_log_interval` | int | `100` | Steps between completion logging |
+| `swanlab_completion_max_buffer` | int | `128` | Max completions to buffer (memory bound) |
+
+## Configuration Examples
+
+### Example 1: Basic Cloud Sync
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: llama-finetune
+swanlab_experiment_name: llama-3-8b-instruct-v1
+swanlab_mode: cloud
+```
+
+### Example 2: Offline/Local Mode
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: local-experiments
+swanlab_experiment_name: test-run-1
+swanlab_mode: local  # or 'offline'
+```
+
+### Example 3: Team Workspace
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: research-project
+swanlab_experiment_name: experiment-42
+swanlab_workspace: my-research-team
+swanlab_mode: cloud
+```
+
+### Example 4: Private Deployment
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: internal-project
+swanlab_experiment_name: secure-training
+swanlab_mode: cloud
+swanlab_web_host: https://swanlab.yourcompany.com
+swanlab_api_host: https://api.swanlab.yourcompany.com
+```
+
+## Team Notifications with Lark (Feishu)
+
+SwanLab supports sending real-time training notifications to your team chat via Lark (Feishu), ByteDance's enterprise collaboration platform. This is especially useful for:
+- **Production training monitoring**: Get alerts when training starts, completes, or encounters errors
+- **Team collaboration**: Keep your ML team informed about long-running experiments
+- **Multi-timezone teams**: Team members can check training progress without being online
+
+### Prerequisites
+
+1. **Lark Bot Setup**: Create a custom bot in your Lark group chat
+2. **Webhook URL**: Get the webhook URL from your Lark bot settings
+3. **HMAC Secret** (recommended): Enable signature verification in your Lark bot for security
+
+For detailed Lark bot setup instructions, see [Lark Custom Bot Documentation](https://open.feishu.cn/document/ukTMukTMukTM/ucTM5YjL3ETO24yNxkjN).
+
+### Example 5: Basic Lark Notifications
+
+Send training notifications to a Lark group chat:
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: production-training
+swanlab_experiment_name: llama-3-finetune-v2
+swanlab_mode: cloud
+
+# Lark notification (basic, no HMAC verification)
+swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
+```
+
+**Note**: This configuration will work, but you'll see a security warning recommending HMAC secret configuration.
+
+### Example 6: Lark Notifications with HMAC Security (Recommended)
+
+For production use, enable HMAC signature verification:
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: production-training
+swanlab_experiment_name: llama-3-finetune-v2
+swanlab_mode: cloud
+
+# Lark notification with HMAC authentication
+swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
+swanlab_lark_secret: your-webhook-secret-key
+```
+
+**Why HMAC secret matters**:
+- Prevents unauthorized parties from sending fake notifications to your Lark group
+- Ensures notifications genuinely come from your training jobs
+- Required for production deployments with sensitive training data
+
+### Example 7: Team Workspace + Lark Notifications
+
+Combine team workspace collaboration with Lark notifications:
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: research-project
+swanlab_experiment_name: multimodal-experiment-42
+swanlab_workspace: ml-research-team
+swanlab_mode: cloud
+
+# Notify team via Lark when training starts/completes
+swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
+swanlab_lark_secret: your-webhook-secret-key
+```
+
+### What Notifications Are Sent?
+
+SwanLab's Lark integration sends notifications for key training events:
+- **Training Start**: When your experiment begins
+- **Training Complete**: When training finishes successfully
+- **Training Errors**: If training crashes or encounters critical errors
+- **Metric Milestones**: Configurable alerts for metric thresholds (if configured in SwanLab)
+
+Each notification includes:
+- Experiment name and project
+- Training status
+- Key metrics (loss, learning rate)
+- Direct link to SwanLab dashboard
+
+### Lark Configuration Validation
+
+The plugin validates your Lark configuration at startup:
+
+#### ✅ Valid Configurations
+
+```yaml
+# Option 1: No Lark (default)
+use_swanlab: true
+swanlab_project: my-project
+# No swanlab_lark_webhook_url → Lark disabled, no warnings
+
+# Option 2: Lark with HMAC secret (recommended)
+use_swanlab: true
+swanlab_project: my-project
+swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxx
+swanlab_lark_secret: your-secret
+# ✅ Logs: "Registered Lark notification callback with HMAC authentication"
+
+# Option 3: Lark without secret (works but not recommended)
+use_swanlab: true
+swanlab_project: my-project
+swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxx
+# ⚠️ Logs: "Registered Lark notification callback (no HMAC secret)"
+# ⚠️ Warning: "Lark webhook has no secret configured. For production use, set 'swanlab_lark_secret'..."
+```
+
+### Security Best Practices
+
+1. **Always use HMAC secret in production**:
+   ```yaml
+   swanlab_lark_webhook_url: https://open.feishu.cn/...
+   swanlab_lark_secret: your-secret-key  # ✅ Add this!
+   ```
+
+2. **Store secrets in environment variables** (even better):
+   ```yaml
+   # In your training script/environment
+   export SWANLAB_LARK_WEBHOOK_URL="https://open.feishu.cn/..."
+   export SWANLAB_LARK_SECRET="your-secret-key"
+   ```
+
+   Then in config:
+   ```yaml
+   # SwanLab plugin will auto-detect environment variables
+   use_swanlab: true
+   swanlab_project: my-project
+   # Lark URL and secret read from env vars
+   ```
+
+3. **Rotate webhook secrets periodically**: Update your Lark bot's secret every 90 days
+
+4. **Use separate webhooks for dev/prod**: Don't mix development and production notifications
+
+### Distributed Training
+
+Lark notifications are automatically deduplicated in distributed training:
+- Only **rank 0** sends notifications
+- Other GPU ranks skip Lark registration
+- Prevents duplicate messages in multi-GPU training
+
+```bash
+# Running on 4 GPUs
+torchrun --nproc_per_node=4 -m axolotl.cli.train config.yml
+
+# Expected logs:
+# [Rank 0] Registered Lark notification callback with HMAC authentication
+# [Rank 1-3] (no Lark registration messages)
+```
+
+## RLHF Completion Table Logging
+
+For RLHF (Reinforcement Learning from Human Feedback) training methods like DPO, KTO, ORPO, and GRPO, SwanLab can log model completions (prompts, chosen/rejected responses, rewards) to a visual table for qualitative analysis. This helps you:
+
+- **Inspect model behavior**: See actual model outputs during training
+- **Debug preference learning**: Compare chosen vs rejected responses
+- **Track reward patterns**: Monitor how rewards evolve over training
+- **Share examples with team**: Visual tables in SwanLab dashboard
+
+### Features
+
+- ✅ **Automatic detection**: Works with DPO, KTO, ORPO, GRPO trainers
+- ✅ **Memory-safe buffering**: Bounded buffer prevents memory leaks in long training runs
+- ✅ **Periodic logging**: Configurable logging interval to reduce overhead
+- ✅ **Rich visualization**: SwanLab tables show prompts, responses, and metrics side-by-side
+
+### Configuration
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `swanlab_log_completions` | bool | `true` | Enable completion logging for RLHF trainers |
+| `swanlab_completion_log_interval` | int | `100` | Log completions to SwanLab every N training steps |
+| `swanlab_completion_max_buffer` | int | `128` | Maximum completions to buffer (memory bound) |
+
+### Example: DPO Training with Completion Logging
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: dpo-training
+swanlab_experiment_name: llama-3-dpo-v1
+swanlab_mode: cloud
+
+# RLHF completion logging (enabled by default)
+swanlab_log_completions: true
+swanlab_completion_log_interval: 100  # Log every 100 steps
+swanlab_completion_max_buffer: 128    # Keep last 128 completions
+
+# DPO-specific config
+rl: dpo
+datasets:
+  - path: /path/to/preference_dataset
+    type: chatml.intel
+```
+
+### Example: Disable Completion Logging
+
+If you're doing a quick test run or don't need completion tables:
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: dpo-training
+
+# Disable completion logging
+swanlab_log_completions: false
+```
+
+### Supported RLHF Trainers
+
+The completion logging callback automatically activates for these trainer types:
+
+- **DPO (Direct Preference Optimization)**: Logs prompts, chosen, rejected, reward_diff
+- **KTO (Kahneman-Tversky Optimization)**: Logs prompts, completions, labels, rewards
+- **ORPO (Odds Ratio Preference Optimization)**: Logs prompts, chosen, rejected, log_odds_ratio
+- **GRPO (Group Relative Policy Optimization)**: Logs prompts, completions, rewards, advantages
+- **CPO (Constrained Policy Optimization)**: Logs prompts, chosen, rejected
+
+For non-RLHF trainers (standard supervised fine-tuning), the completion callback is automatically skipped.
+
+### How It Works
+
+1. **Auto-detection**: Plugin detects trainer type at initialization
+2. **Buffering**: Completions are buffered in memory (up to `swanlab_completion_max_buffer`)
+3. **Periodic logging**: Every `swanlab_completion_log_interval` steps, buffer is logged to SwanLab
+4. **Memory safety**: Old completions are automatically dropped when buffer is full (uses `collections.deque`)
+5. **Final flush**: Remaining completions are logged when training completes
+
+### Viewing Completion Tables
+
+After training starts, you can view completion tables in your SwanLab dashboard:
+
+1. Navigate to your experiment in SwanLab
+2. Look for the "rlhf_completions" table in the metrics panel
+3. The table shows:
+   - **step**: Training step when completion was generated
+   - **prompt**: Input prompt
+   - **chosen**: Preferred response (DPO/ORPO)
+   - **rejected**: Non-preferred response (DPO/ORPO)
+   - **completion**: Model output (KTO/GRPO)
+   - **reward_diff/reward**: Reward metrics
+   - Trainer-specific metrics (e.g., log_odds_ratio for ORPO)
+
+### Memory Management
+
+The completion buffer is **memory-bounded** to prevent memory leaks:
+
+```python
+# Internal implementation uses deque with maxlen
+from collections import deque
+
+buffer = deque(maxlen=128)  # Old completions automatically dropped
+```
+
+**Memory usage estimate**:
+- Average completion: ~500 characters (prompt + responses)
+- Buffer size 128: ~64 KB (negligible)
+- Buffer size 1024: ~512 KB (still small)
+
+**Recommendation**: Default buffer size (128) works well for most cases. Increase to 512-1024 only if you need to review more historical completions.
+
+### Performance Impact
+
+Completion logging has minimal overhead:
+
+- **Buffering**: O(1) append operation, negligible CPU/memory
+- **Logging**: Only happens every N steps (default: 100)
+- **Network**: SwanLab batches table uploads efficiently
+
+**Expected overhead**: < 0.5% per training step
+
+### Troubleshooting
+
+#### Completions not appearing in SwanLab
+
+**Cause**: Trainer may not be logging completion data in the expected format.
+
+**Diagnostic steps**:
+1. Check trainer type detection in logs:
+   ```text
+   INFO: SwanLab RLHF completion logging enabled for DPOTrainer (type: dpo)
+   ```
+2. Verify your trainer is an RLHF trainer (DPO/KTO/ORPO/GRPO)
+3. Check if trainer logs completion data (this depends on TRL version)
+
+**Note**: The current implementation expects trainers to log completion data in the `logs` dict during `on_log()` callback. Some TRL trainers may not expose this data by default. You may need to patch the trainer to expose completions.
+
+#### Buffer fills up too quickly
+
+**Cause**: High logging frequency with small buffer size.
+
+**Solution**: Increase buffer size or logging interval:
+```yaml
+swanlab_completion_log_interval: 200  # Log less frequently
+swanlab_completion_max_buffer: 512    # Larger buffer
+```
+
+#### Memory usage growing over time
+
+**Cause**: Buffer should be bounded, so this indicates a bug.
+
+**Solution**:
+1. Verify `swanlab_completion_max_buffer` is set
+2. Check SwanLab version is up to date
+3. Report issue with memory profiling data
+
+## Performance Profiling
+
+SwanLab integration includes profiling utilities to measure and log execution time of trainer methods. This helps you:
+
+- **Identify bottlenecks**: Find slow operations in your training loop
+- **Optimize performance**: Track improvements after optimization changes
+- **Monitor distributed training**: See per-rank timing differences
+- **Debug hangs**: Detect methods that take unexpectedly long
+
+### Features
+
+- ✅ **Zero-config profiling**: Automatic timing of key trainer methods
+- ✅ **Decorator-based**: Easy to add profiling to custom methods with `@swanlab_profile`
+- ✅ **Context manager**: Fine-grained profiling with `swanlab_profiling_context()`
+- ✅ **Advanced filtering**: `ProfilingConfig` for throttling and minimum duration thresholds
+- ✅ **Exception-safe**: Logs duration even if function raises an exception
+
+### Basic Usage: Decorator
+
+Add profiling to any trainer method with the `@swanlab_profile` decorator:
+
+```python
+from axolotl.integrations.swanlab.profiling import swanlab_profile
+
+class MyCustomTrainer(AxolotlTrainer):
+    @swanlab_profile
+    def training_step(self, model, inputs):
+        # Your training step logic
+        return super().training_step(model, inputs)
+
+    @swanlab_profile
+    def prediction_step(self, model, inputs, prediction_loss_only):
+        # Your prediction logic
+        return super().prediction_step(model, inputs, prediction_loss_only)
+```
+
+The decorator automatically:
+1. Measures execution time with high-precision timer
+2. Logs to SwanLab as `profiling/Time taken: ClassName.method_name`
+3. Only logs if SwanLab is enabled (`use_swanlab: true`)
+4. Gracefully handles exceptions (logs duration, then re-raises)
+
+### Advanced Usage: Context Manager
+
+For fine-grained profiling within a method:
+
+```python
+from axolotl.integrations.swanlab.profiling import swanlab_profiling_context
+
+class MyTrainer(AxolotlTrainer):
+    def complex_training_step(self, model, inputs):
+        # Profile just the forward pass
+        with swanlab_profiling_context(self, "forward_pass"):
+            outputs = model(**inputs)
+
+        # Profile just the backward pass
+        with swanlab_profiling_context(self, "backward_pass"):
+            loss = outputs.loss
+            loss.backward()
+
+        return outputs
+```
+
+### Advanced Usage: ProfilingConfig
+
+Filter and throttle profiling logs with `ProfilingConfig`:
+
+```python
+from axolotl.integrations.swanlab.profiling import (
+    swanlab_profiling_context_advanced,
+    ProfilingConfig,
+)
+
+# Create custom profiling config
+profiling_config = ProfilingConfig(
+    enabled=True,
+    min_duration_ms=1.0,    # Only log if duration > 1ms
+    log_interval=10,        # Log every 10th call
+)
+
+class MyTrainer(AxolotlTrainer):
+    def frequently_called_method(self, data):
+        with swanlab_profiling_context_advanced(
+            self,
+            "frequent_op",
+            config=profiling_config
+        ):
+            # This only logs every 10th call, and only if it takes > 1ms
+            result = expensive_computation(data)
+        return result
+```
+
+**ProfilingConfig Parameters**:
+- `enabled`: Enable/disable profiling globally (default: `True`)
+- `min_duration_ms`: Minimum duration to log in milliseconds (default: `0.1`)
+- `log_interval`: Log every Nth function call (default: `1` = log all)
+
+**Use cases**:
+- **High-frequency methods**: Use `log_interval=100` to reduce logging overhead
+- **Filter noise**: Use `min_duration_ms=1.0` to skip very fast operations
+- **Debugging**: Use `log_interval=1, min_duration_ms=0.0` to log everything
+
+### Viewing Profiling Metrics
+
+In your SwanLab dashboard, profiling metrics appear under the "profiling" namespace:
+
+```text
+profiling/Time taken: AxolotlTrainer.training_step
+profiling/Time taken: AxolotlTrainer.prediction_step
+profiling/Time taken: MyTrainer.forward_pass
+profiling/Time taken: MyTrainer.backward_pass
+```
+
+You can:
+- **Track over time**: See if methods get faster/slower during training
+- **Compare runs**: Compare profiling metrics across experiments
+- **Identify regressions**: Detect if a code change slowed down training
+
+### Configuration in Axolotl Config
+
+Profiling is automatically enabled when SwanLab is enabled. No additional config needed:
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+use_swanlab: true
+swanlab_project: my-project
+
+# Profiling is automatically enabled
+# Add @swanlab_profile decorators to your custom trainer methods
+```
+
+To disable profiling while keeping SwanLab enabled:
+
+```python
+# In your custom trainer code
+from axolotl.integrations.swanlab.profiling import DEFAULT_PROFILING_CONFIG
+
+# Disable profiling globally
+DEFAULT_PROFILING_CONFIG.enabled = False
+```
+
+### Performance Impact
+
+- **Decorator overhead**: ~2-5 microseconds per call (negligible)
+- **Context manager overhead**: ~1-3 microseconds (negligible)
+- **Logging overhead**: Only when SwanLab is enabled and method duration exceeds threshold
+- **Network overhead**: SwanLab batches metrics efficiently
+
+**Expected overhead**: < 0.1% per training step (effectively zero)
+
+### Best Practices
+
+1. **Profile bottlenecks first**: Start by profiling suspected slow operations
+2. **Use min_duration_ms**: Filter out fast operations (< 1ms) to reduce noise
+3. **Throttle high-frequency calls**: Use `log_interval` for methods called > 100 times/step
+4. **Profile across runs**: Compare profiling metrics before/after optimization
+5. **Monitor distributed training**: Check for rank-specific slowdowns
+
+### Example: Complete Profiling Setup
+
+```python
+from axolotl.integrations.swanlab.profiling import (
+    swanlab_profile,
+    swanlab_profiling_context,
+    ProfilingConfig,
+)
+
+class OptimizedTrainer(AxolotlTrainer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Custom profiling config for high-frequency operations
+        self.fast_op_config = ProfilingConfig(
+            enabled=True,
+            min_duration_ms=0.5,
+            log_interval=50,
+        )
+
+    @swanlab_profile
+    def training_step(self, model, inputs):
+        """Main training step - always profile."""
+        return super().training_step(model, inputs)
+
+    @swanlab_profile
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """Loss computation - always profile."""
+        return super().compute_loss(model, inputs, return_outputs)
+
+    def _prepare_inputs(self, inputs):
+        """High-frequency operation - throttled profiling."""
+        with swanlab_profiling_context_advanced(
+            self,
+            "prepare_inputs",
+            config=self.fast_op_config,
+        ):
+            return super()._prepare_inputs(inputs)
+```
+
+### Troubleshooting
+
+#### Profiling metrics not appearing in SwanLab
+
+**Cause**: SwanLab is not enabled or not initialized.
+
+**Solution**:
+```yaml
+# Ensure SwanLab is enabled
+use_swanlab: true
+swanlab_project: my-project
+```
+
+Check logs for:
+```text
+INFO: SwanLab initialized for project: my-project
+```
+
+#### Too many profiling metrics cluttering dashboard
+
+**Cause**: Profiling every function call for high-frequency operations.
+
+**Solution**: Use `ProfilingConfig` with throttling:
+```python
+config = ProfilingConfig(
+    min_duration_ms=1.0,    # Skip fast ops
+    log_interval=100,       # Log every 100th call
+)
+```
+
+#### Profiling overhead impacting training speed
+
+**Cause**: Profiling itself should have negligible overhead (< 0.1%). If you see > 1% slowdown, this indicates a bug.
+
+**Solution**:
+1. Disable profiling temporarily to confirm:
+   ```python
+   DEFAULT_PROFILING_CONFIG.enabled = False
+   ```
+2. Report issue with profiling data and trainer details
+
+#### Profiling shows inconsistent timing
+
+**Cause**: Normal variation due to GPU warmup, data loading, or system load.
+
+**Solution**:
+- Ignore first few steps (warmup period)
+- Look at average/median timing over many steps
+- Use `log_interval` to reduce noise from individual outliers
+
+## Complete Config Example
+
+Here's a complete example integrating SwanLab with your RVQ-Alpha training:
+
+```yaml
+base_model: /path/to/your/model
+model_type: Qwen2ForCausalLM
+
+# SwanLab Integration
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
+use_swanlab: true
+swanlab_project: RVQ-Alpha-Training
+swanlab_experiment_name: Qwen2.5-7B-MetaQA-Perturb-P020
+swanlab_description: "Training on MetaQA and Perturbation datasets with NEW-RVQ encoding"
+swanlab_mode: cloud
+swanlab_workspace: single-cell-genomics
+
+# Training configuration
+sequence_len: 32768
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+num_epochs: 2
+learning_rate: 2e-5
+optimizer: adamw_torch_fused
+
+# Datasets
+datasets:
+  - path: /path/to/dataset
+    type: chat_template
+
+# Output
+output_dir: ./outputs
+```
+
+## Modes Explained
+
+### `cloud` Mode (Default)
+- Syncs experiments to SwanLab cloud in real-time
+- Requires API key and internet connection
+- Best for: Team collaboration, remote monitoring
+
+### `local` Mode
+- Saves experiments locally only
+- No cloud sync
+- Best for: Local development, air-gapped environments
+
+### `offline` Mode
+- Saves metadata locally
+- Can sync to cloud later using `swanlab sync`
+- Best for: Unstable internet, sync later
+
+### `disabled` Mode
+- Turns off SwanLab completely
+- No logging or tracking
+- Best for: Debugging, testing
+
+## Configuration Validation & Conflict Detection
+
+SwanLab integration includes comprehensive validation and conflict detection to help you catch configuration errors early and avoid performance issues.
+
+### Required Fields Validation
+
+The plugin validates your configuration at startup and provides clear error messages with solutions:
+
+#### Missing Project Name
+
+```yaml
+# ❌ INVALID: use_swanlab enabled but no project
+use_swanlab: true
+# Error: SwanLab enabled but 'swanlab_project' is not set.
+```
+
+**Solution**:
+```yaml
+# ✅ VALID: Provide project name
+use_swanlab: true
+swanlab_project: my-project
+```
+
+#### Invalid Mode
+
+```yaml
+# ❌ INVALID: Unknown mode
+use_swanlab: true
+swanlab_project: my-project
+swanlab_mode: invalid-mode
+# Error: Invalid swanlab_mode: 'invalid-mode'. Valid options: cloud, local, offline, disabled
+```
+
+**Solution**:
+```yaml
+# ✅ VALID: Use one of the valid modes
+use_swanlab: true
+swanlab_project: my-project
+swanlab_mode: cloud  # or: local, offline, disabled
+```
+
+#### Empty Project Name
+
+```yaml
+# ❌ INVALID: Empty string project name
+use_swanlab: true
+swanlab_project: ""
+# Error: swanlab_project cannot be an empty string.
+```
+
+**Solution**:
+```yaml
+# ✅ VALID: Provide non-empty project name
+use_swanlab: true
+swanlab_project: my-project
+```
+
+### Cloud Mode API Key Warning
+
+When using `cloud` mode without an API key, you'll receive a warning with multiple solutions:
+
+```yaml
+use_swanlab: true
+swanlab_project: my-project
+swanlab_mode: cloud
+# No API key set
+# Warning: SwanLab cloud mode enabled but no API key found.
+```
+
+**Solutions**:
+1. Set environment variable: `export SWANLAB_API_KEY=your-api-key`
+2. Add to config (less secure): `swanlab_api_key: your-api-key`
+3. Run `swanlab login` before training
+4. Use `swanlab_mode: local` for offline tracking
+
+### Multi-Logger Performance Warnings
+
+Using multiple logging tools simultaneously (SwanLab + WandB + MLflow + Comet) can impact training performance:
+
+#### Two Loggers - Warning
+
+```yaml
+use_swanlab: true
+swanlab_project: my-project
+
+use_wandb: true
+wandb_project: my-project
+
+# Warning: Multiple logging tools enabled: SwanLab, WandB
+# Expected overhead: ~3.0% per training step.
+```
+
+**Impact**:
+- Performance overhead: ~1-2% per logger (cumulative)
+- Increased memory usage
+- Longer training time per step
+- Potential config/callback conflicts
+
+**Recommendations**:
+- Choose ONE primary logging tool for production training
+- Use multiple loggers only for:
+  - Migration period (transitioning between tools)
+  - Short comparison runs
+  - Debugging specific tool issues
+- Monitor system resources (CPU, memory) during training
+
+#### Three+ Loggers - Error-Level Warning
+
+```yaml
+use_swanlab: true
+swanlab_project: my-project
+
+use_wandb: true
+wandb_project: my-project
+
+use_mlflow: true
+mlflow_tracking_uri: http://localhost:5000
+
+# ERROR: 3 logging tools enabled simultaneously!
+# Expected overhead: ~4.5% per training step.
+# STRONGLY RECOMMEND: Disable all but ONE logging tool
+```
+
+**Why This Matters**:
+- With 3 loggers: ~4-5% overhead per step → significant slowdown over long training
+- Example: 10,000 steps at 2s/step → ~400-500 seconds extra (6-8 minutes)
+- Memory overhead scales with number of loggers
+- Rare edge cases with callback ordering conflicts
+
+### Auto-Enable Logic
+
+For convenience, SwanLab will auto-enable if you specify a project without setting `use_swanlab`:
+
+```yaml
+# This configuration:
+swanlab_project: my-project
+
+# Automatically becomes:
+use_swanlab: true
+swanlab_project: my-project
+```
+
+### Distributed Training Detection
+
+In distributed training scenarios (multi-GPU), the plugin automatically detects and reports:
+
+```yaml
+use_swanlab: true
+swanlab_project: my-project
+swanlab_mode: cloud
+
+# When running with torchrun --nproc_per_node=4:
+# Info: Distributed training detected (world_size=4)
+# Info: SwanLab mode: cloud
+# Info: Only rank 0 will initialize SwanLab
+# Info: Other ranks will skip SwanLab to avoid conflicts
+```
+
+**Why Only Rank 0**:
+- Avoids duplicate experiment runs
+- Reduces network/cloud API overhead on worker ranks
+- Prevents race conditions in metric logging
+
+## Authentication
+
+### Method 1: Environment Variable (Recommended)
+```bash
+export SWANLAB_API_KEY=your-api-key-here
+```
+
+### Method 2: Login Command
+```bash
+swanlab login
+# Enter your API key when prompted
+```
+
+### Method 3: Config File
+```yaml
+swanlab_api_key: your-api-key-here
+```
+
+## What Gets Logged?
+
+### Automatically Logged Metrics
+- Training loss
+- Learning rate
+- Gradient norm
+- Training steps
+- Epoch progress
+
+### Automatically Logged Config
+- Model configuration (base_model, model_type)
+- Training hyperparameters (learning_rate, batch_size, etc.)
+- Optimizer settings
+- Parallelization settings (FSDP, DeepSpeed, Context Parallel)
+- Axolotl configuration file
+- DeepSpeed configuration (if used)
+
+## Viewing Your Experiments
+
+### Cloud Mode
+Visit [https://swanlab.cn](https://swanlab.cn) and navigate to your project to view:
+- Real-time training metrics
+- Hyperparameter comparison
+- System resource usage
+- Configuration files
+
+### Local Mode
+```bash
+# Start local dashboard
+swanlab watch ./swanlog
+
+# Open browser to http://localhost:5092
+```
+
+## Integration with Existing Tools
+
+SwanLab can work alongside other tracking tools:
+
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin
+
+# Use both SwanLab and Wandb
+use_swanlab: true
+swanlab_project: my-project
+
+use_wandb: true
+wandb_project: my-project
+```
+
+## Troubleshooting
+
+### Configuration Errors
+
+#### Error: "SwanLab enabled but 'swanlab_project' is not set"
+
+**Cause**: You enabled SwanLab (`use_swanlab: true`) but forgot to specify a project name.
+
+**Solution**:
+```yaml
+use_swanlab: true
+swanlab_project: my-project  # Add this line
+```
+
+#### Error: "Invalid swanlab_mode: 'xxx'"
+
+**Cause**: You provided an invalid mode value.
+
+**Solution**: Use one of the valid modes:
+```yaml
+swanlab_mode: cloud     # or: local, offline, disabled
+```
+
+#### Error: "swanlab_project cannot be an empty string"
+
+**Cause**: You set `swanlab_project: ""` (empty string).
+
+**Solution**: Either provide a valid name or remove the field:
+```yaml
+# Option 1: Provide valid name
+swanlab_project: my-project
+
+# Option 2: Remove the field entirely
+# swanlab_project: ""  <- Remove this line
+```
+
+### Import Errors
+
+#### Error: "SwanLab is not installed"
+
+**Cause**: SwanLab package is not installed in your environment.
+
+**Solution**:
+```bash
+pip install swanlab
+# or
+pip install swanlab>=0.3.0
+```
+
+### Performance Issues
+
+#### Warning: "Multiple logging tools enabled"
+
+**Cause**: You have multiple experiment tracking tools enabled (e.g., SwanLab + WandB + MLflow).
+
+**Impact**: ~1-2% performance overhead per logger, cumulative.
+
+**Solution**: For production training, disable all but one logger:
+```yaml
+# Option 1: Keep only SwanLab
+use_swanlab: true
+swanlab_project: my-project
+use_wandb: false      # Disable others
+use_mlflow: false
+
+# Option 2: Keep only WandB
+use_swanlab: false
+use_wandb: true
+wandb_project: my-project
+```
+
+**Exception**: Multiple loggers are acceptable for:
+- Short comparison runs (< 100 steps)
+- Migration testing between logging tools
+- Debugging logger-specific issues
+
+### Distributed Training Issues
+
+#### SwanLab creates duplicate runs in multi-GPU training
+
+**Cause**: All ranks are initializing SwanLab instead of just rank 0.
+
+**Expected Behavior**: The plugin automatically ensures only rank 0 initializes SwanLab. You should see:
+```text
+Info: Distributed training detected (world_size=4)
+Info: Only rank 0 will initialize SwanLab
+Info: Other ranks will skip SwanLab to avoid conflicts
+```
+
+**If you see duplicates**:
+1. Check your plugin is loaded correctly
+2. Verify you're using the latest SwanLab integration code
+3. Check logs for initialization messages on all ranks
+
+### SwanLab not logging metrics
+
+**Solution**: Ensure SwanLab is initialized before training starts. The plugin automatically handles this in `pre_model_load`.
+
+### API Key errors
+
+**Solution**:
+```bash
+# Verify API key
+echo $SWANLAB_API_KEY
+
+# Re-login
+swanlab login
+```
+
+### Cloud sync issues
+
+**Solution**: Use `offline` mode and sync later:
+```yaml
+swanlab_mode: offline
+```
+
+Then sync when ready:
+```bash
+swanlab sync ./swanlog
+```
+
+### Plugin not loaded
+
+**Solution**: Verify plugin path in config:
+```yaml
+plugins:
+  - axolotl.integrations.swanlab.SwanLabPlugin  # Correct path
+```
+
+### Lark Notification Issues
+
+#### Error: "Failed to import SwanLab Lark plugin"
+
+**Cause**: Your SwanLab version doesn't include the Lark plugin (requires SwanLab >= 0.3.0).
+
+**Solution**:
+```bash
+# Upgrade SwanLab to latest version
+pip install --upgrade swanlab
+
+# Or install specific version
+pip install 'swanlab>=0.3.0'
+```
+
+#### Warning: "Lark webhook has no secret configured"
+
+**Cause**: You provided `swanlab_lark_webhook_url` but no `swanlab_lark_secret`.
+
+**Impact**: Lark notifications will work, but without HMAC authentication (security risk).
+
+**Solution**: Add HMAC secret for production use:
+```yaml
+swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxx
+swanlab_lark_secret: your-webhook-secret  # Add this line
+```
+
+**When it's OK to skip secret**:
+- Local development and testing
+- Internal networks with restricted access
+- Non-sensitive training experiments
+
+**When secret is required**:
+- Production training jobs
+- Training with proprietary data
+- Multi-team shared Lark groups
+
+#### Error: "Failed to register Lark callback"
+
+**Cause**: Invalid webhook URL or network connectivity issues.
+
+**Diagnostic steps**:
+```bash
+# 1. Test webhook URL manually
+curl -X POST "YOUR_WEBHOOK_URL" \
+  -H 'Content-Type: application/json' \
+  -d '{"msg_type":"text","content":{"text":"Test from Axolotl"}}'
+
+# 2. Check SwanLab version
+pip show swanlab
+
+# 3. Verify webhook URL format
+# Should start with: https://open.feishu.cn/open-apis/bot/v2/hook/
+```
+
+**Solution**:
+1. Verify webhook URL is correct (copy from Lark bot settings)
+2. Check network connectivity to Lark API
+3. Ensure webhook is not expired (Lark webhooks can expire)
+4. Regenerate webhook URL in Lark bot settings if needed
+
+#### Lark notifications not received
+
+**Cause**: Multiple possible causes.
+
+**Diagnostic checklist**:
+
+1. **Check training logs** for Lark registration confirmation:
+   ```text
+   # Expected log message (rank 0 only):
+   INFO: Registered Lark notification callback with HMAC authentication
+   ```
+
+2. **Verify webhook in Lark**: Test webhook manually (see above)
+
+3. **Check distributed training**: Only rank 0 sends notifications
+   ```bash
+   # If running multi-GPU, check rank 0 logs specifically
+   grep "Registered Lark" logs/rank_0.log
+   ```
+
+4. **Verify SwanLab is initialized**: Lark callback needs SwanLab to be running
+   ```yaml
+   use_swanlab: true  # Must be enabled
+   swanlab_project: my-project  # Must be set
+   ```
+
+5. **Check Lark bot permissions**: Ensure bot is added to the target group chat
+
+#### Duplicate Lark notifications in multi-GPU training
+
+**Expected Behavior**: Should NOT happen - only rank 0 sends notifications.
+
+**If you see duplicates**:
+1. Check that all GPUs are using the same config file
+2. Verify plugin is loaded correctly on all ranks
+3. Check logs for unexpected Lark initialization on non-zero ranks
+4. Ensure `RANK` or `LOCAL_RANK` environment variables are set correctly
+
+**Solution**: This is a bug if it occurs. Report with:
+- Full training command
+- Logs from all ranks
+- Config file
+
+## Comparison: SwanLab vs WandB
+
+| Feature | SwanLab | WandB |
+|---------|---------|-------|
+| Open Source | ✅ Yes | ❌ No |
+| Self-Hosting | ✅ Easy | ⚠️ Complex |
+| Free Tier | ✅ Generous | ⚠️ Limited |
+| Chinese Support | ✅ Native | ⚠️ Limited |
+| Offline Mode | ✅ Full support | ✅ Supported |
+| Integration | 🆕 New | ✅ Mature |
+
+## Advanced Usage
+
+### Custom Logging
+
+You can add custom metrics in your callbacks:
+
+```python
+import swanlab
+
+# In your custom callback
+swanlab.log({
+    "custom_metric": value,
+    "epoch": epoch_num
+})
+```
+
+### Experiment Comparison
+
+```bash
+# Compare multiple experiments
+swanlab compare run1 run2 run3
+```
+
+## Support
+
+- **Documentation**: [https://docs.swanlab.cn](https://docs.swanlab.cn)
+- **GitHub**: [https://github.com/SwanHubX/SwanLab](https://github.com/SwanHubX/SwanLab)
+- **Issues**: Report bugs at [GitHub Issues](https://github.com/SwanHubX/SwanLab/issues)
+
+## License
+
+This integration follows the Axolotl Community License Agreement.
+
+## Acknowledgements
+
+This integration is built on top of:
+- [SwanLab](https://github.com/SwanHubX/SwanLab) - Experiment tracking tool
+- [Transformers](https://github.com/huggingface/transformers) - SwanLabCallback
+- [Axolotl](https://github.com/axolotl-ai-cloud/axolotl) - Training framework
diff --git a/src/axolotl/integrations/swanlab/__init__.py b/src/axolotl/integrations/swanlab/__init__.py
new file mode 100644
index 0000000000..241a27764f
--- /dev/null
+++ b/src/axolotl/integrations/swanlab/__init__.py
@@ -0,0 +1,6 @@
+"""SwanLab integration plugin for Axolotl"""
+
+from axolotl.integrations.swanlab.args import SwanLabConfig
+from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+
+__all__ = ["SwanLabConfig", "SwanLabPlugin"]
diff --git a/src/axolotl/integrations/swanlab/args.py b/src/axolotl/integrations/swanlab/args.py
new file mode 100644
index 0000000000..2cf31252dc
--- /dev/null
+++ b/src/axolotl/integrations/swanlab/args.py
@@ -0,0 +1,140 @@
+"""SwanLab configuration arguments"""
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+
+class SwanLabConfig(BaseModel):
+    """SwanLab configuration subset"""
+
+    use_swanlab: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Enable SwanLab experiment tracking and visualization"
+        },
+    )
+    swanlab_project: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Your SwanLab project name"},
+    )
+    swanlab_experiment_name: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Set the name of your SwanLab experiment"},
+    )
+    swanlab_description: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Description for your SwanLab experiment"},
+    )
+    swanlab_mode: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": '"cloud" to sync to SwanLab cloud, "local" for local only, "offline" to save metadata locally, "disabled" to turn off SwanLab'
+        },
+    )
+    swanlab_workspace: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "SwanLab workspace name (organization or username)"
+        },
+    )
+    swanlab_api_key: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "SwanLab API key for authentication. Can also be set via SWANLAB_API_KEY environment variable"
+        },
+    )
+    swanlab_log_model: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Whether to log model checkpoints to SwanLab (feature coming soon)"
+        },
+    )
+    swanlab_web_host: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Web address for SwanLab cloud environment (for private deployment)"
+        },
+    )
+    swanlab_api_host: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "API address for SwanLab cloud environment (for private deployment)"
+        },
+    )
+    swanlab_lark_webhook_url: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Lark (Feishu) webhook URL for sending training notifications to team chat"
+        },
+    )
+    swanlab_lark_secret: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Secret for Lark webhook HMAC signature authentication (optional)"
+        },
+    )
+    swanlab_log_completions: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Enable logging RLHF completions to SwanLab for qualitative analysis (DPO/KTO/ORPO/GRPO)"
+        },
+    )
+    swanlab_completion_log_interval: int | None = Field(
+        default=100,
+        json_schema_extra={
+            "description": "Number of training steps between completion table logging to SwanLab"
+        },
+    )
+    swanlab_completion_max_buffer: int | None = Field(
+        default=128,
+        json_schema_extra={
+            "description": "Maximum number of completions to buffer before logging (prevents memory leaks)"
+        },
+    )
+
+    @field_validator("swanlab_mode")
+    @classmethod
+    def validate_swanlab_mode(cls, v):
+        """Validate swanlab_mode is one of the allowed values."""
+        if v is None:
+            return v
+
+        valid_modes = ["cloud", "local", "offline", "disabled"]
+        if v not in valid_modes:
+            raise ValueError(
+                f"Invalid swanlab_mode: '{v}'.\n\n"
+                f"Valid options: {', '.join(valid_modes)}\n\n"
+                f"Examples:\n"
+                f"  swanlab_mode: cloud     # Sync to SwanLab cloud\n"
+                f"  swanlab_mode: local     # Local only, no cloud sync\n"
+                f"  swanlab_mode: offline   # Save metadata locally\n"
+                f"  swanlab_mode: disabled  # Turn off SwanLab\n"
+            )
+        return v
+
+    @field_validator("swanlab_project")
+    @classmethod
+    def validate_swanlab_project(cls, v):
+        """Validate swanlab_project is non-empty when provided."""
+        if v is not None and isinstance(v, str) and len(v.strip()) == 0:
+            raise ValueError(
+                "swanlab_project cannot be an empty string.\n\n"
+                "Either:\n"
+                "  1. Provide a valid project name: swanlab_project: my-project\n"
+                "  2. Remove the swanlab_project field entirely\n"
+            )
+        return v
+
+    @model_validator(mode="after")
+    def validate_swanlab_enabled_requires_project(self):
+        """Validate that if use_swanlab is True, swanlab_project must be set."""
+        if self.use_swanlab is True and not self.swanlab_project:
+            raise ValueError(
+                "SwanLab enabled (use_swanlab: true) but 'swanlab_project' is not set.\n\n"
+                "Solutions:\n"
+                "  1. Add 'swanlab_project: your-project-name' to your config\n"
+                "  2. Set 'use_swanlab: false' to disable SwanLab\n\n"
+                "Example:\n"
+                "  use_swanlab: true\n"
+                "  swanlab_project: my-llm-training\n"
+            )
+        return self
diff --git a/src/axolotl/integrations/swanlab/callbacks.py b/src/axolotl/integrations/swanlab/callbacks.py
new file mode 100644
index 0000000000..8dfc0fe53c
--- /dev/null
+++ b/src/axolotl/integrations/swanlab/callbacks.py
@@ -0,0 +1,179 @@
+"""SwanLab callbacks for Axolotl trainers.
+
+This module provides HuggingFace Trainer callbacks for logging
+RLHF completions to SwanLab.
+"""
+
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class SwanLabRLHFCompletionCallback(TrainerCallback):
+    """Callback for logging RLHF completions to SwanLab.
+
+    This callback periodically logs model completions (prompts, chosen/rejected
+    responses, rewards) to SwanLab during RLHF training for qualitative analysis.
+
+    Supports DPO, KTO, ORPO, and GRPO trainers.
+
+    Example usage:
+        >>> callback = SwanLabRLHFCompletionCallback(
+        ...     log_interval=100,  # Log every 100 steps
+        ...     max_completions=128,  # Keep last 128 completions
+        ... )
+        >>> trainer.add_callback(callback)
+
+    Attributes:
+        logger: CompletionLogger instance
+        log_interval: Number of steps between SwanLab logging
+        trainer_type: Auto-detected trainer type (dpo/kto/orpo/grpo)
+    """
+
+    def __init__(
+        self,
+        log_interval: int = 100,
+        max_completions: int = 128,
+        table_name: str = "rlhf_completions",
+    ):
+        """Initialize SwanLab RLHF completion callback.
+
+        Args:
+            log_interval: Log to SwanLab every N steps. Default: 100
+            max_completions: Maximum completions to buffer. Default: 128
+            table_name: SwanLab table name. Default: "rlhf_completions"
+        """
+        super().__init__()
+        self.logger = CompletionLogger(maxlen=max_completions)
+        self.log_interval = log_interval
+        self.table_name = table_name
+        self.trainer_type: str | None = None  # Auto-detected
+        self._last_logged_step = 0
+
+    def on_init_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Detect trainer type on initialization."""
+        trainer = kwargs.get("trainer")
+        if trainer is not None:
+            trainer_name = trainer.__class__.__name__
+            if "DPO" in trainer_name:
+                self.trainer_type = "dpo"
+            elif "KTO" in trainer_name:
+                self.trainer_type = "kto"
+            elif "ORPO" in trainer_name:
+                self.trainer_type = "orpo"
+            elif "GRPO" in trainer_name:
+                self.trainer_type = "grpo"
+            else:
+                self.trainer_type = "unknown"
+
+            LOG.info(
+                f"SwanLab RLHF completion logging enabled for {trainer_name} "
+                f"(type: {self.trainer_type})"
+            )
+
+    def on_log(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        logs: dict | None = None,
+        **kwargs,
+    ):
+        """Capture completions from logs and buffer them.
+
+        Different trainers log completions in different formats:
+        - DPO: logs['dpo/chosen'], logs['dpo/rejected'], logs['dpo/reward_diff']
+        - KTO: logs['kto/completion'], logs['kto/label'], logs['kto/reward']
+        - ORPO: logs['orpo/chosen'], logs['orpo/rejected']
+        - GRPO: logs['grpo/completion'], logs['grpo/reward']
+
+        Note: This is a placeholder implementation. Actual log keys depend
+        on the TRL trainer implementation. You may need to patch the trainers
+        to expose completion data in logs.
+        """
+        if logs is None or self.trainer_type is None:
+            return
+
+        step = state.global_step
+
+        # DPO completions
+        if self.trainer_type == "dpo":
+            if all(key in logs for key in ["dpo/prompt", "dpo/chosen", "dpo/rejected"]):
+                self.logger.add_dpo_completion(
+                    step=step,
+                    prompt=logs.get("dpo/prompt", ""),
+                    chosen=logs.get("dpo/chosen", ""),
+                    rejected=logs.get("dpo/rejected", ""),
+                    reward_diff=logs.get("dpo/reward_diff"),
+                )
+
+        # KTO completions
+        elif self.trainer_type == "kto":
+            if all(key in logs for key in ["kto/prompt", "kto/completion"]):
+                self.logger.add_kto_completion(
+                    step=step,
+                    prompt=logs.get("kto/prompt", ""),
+                    completion=logs.get("kto/completion", ""),
+                    label=logs.get("kto/label", False),
+                    reward=logs.get("kto/reward"),
+                )
+
+        # ORPO completions
+        elif self.trainer_type == "orpo":
+            if all(
+                key in logs for key in ["orpo/prompt", "orpo/chosen", "orpo/rejected"]
+            ):
+                self.logger.add_orpo_completion(
+                    step=step,
+                    prompt=logs.get("orpo/prompt", ""),
+                    chosen=logs.get("orpo/chosen", ""),
+                    rejected=logs.get("orpo/rejected", ""),
+                    log_odds_ratio=logs.get("orpo/log_odds_ratio"),
+                )
+
+        # GRPO completions
+        elif self.trainer_type == "grpo":
+            if all(key in logs for key in ["grpo/prompt", "grpo/completion"]):
+                self.logger.add_grpo_completion(
+                    step=step,
+                    prompt=logs.get("grpo/prompt", ""),
+                    completion=logs.get("grpo/completion", ""),
+                    reward=logs.get("grpo/reward"),
+                    advantage=logs.get("grpo/advantage"),
+                )
+
+        # Periodically log to SwanLab
+        if step - self._last_logged_step >= self.log_interval:
+            if len(self.logger) > 0:
+                self.logger.log_to_swanlab(table_name=self.table_name)
+                self.logger.clear()
+                self._last_logged_step = step
+
+    def on_train_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Log remaining completions at end of training."""
+        if len(self.logger) > 0:
+            LOG.info(
+                f"Training complete, logging final {len(self.logger)} completions to SwanLab"
+            )
+            self.logger.log_to_swanlab(table_name=self.table_name)
+            self._last_logged_step = state.global_step
diff --git a/src/axolotl/integrations/swanlab/completion_logger.py b/src/axolotl/integrations/swanlab/completion_logger.py
new file mode 100644
index 0000000000..dd709227f9
--- /dev/null
+++ b/src/axolotl/integrations/swanlab/completion_logger.py
@@ -0,0 +1,228 @@
+"""SwanLab completion logger for RLHF/DPO/KTO/ORPO/GRPO training.
+
+This module provides utilities for logging model completions during
+preference training to SwanLab for qualitative analysis.
+"""
+
+from collections import deque
+from collections.abc import Mapping
+from typing import Any
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class CompletionLogger:
+    """Memory-bounded logger for RLHF completions.
+
+    Stores prompts, completions, and rewards in fixed-size deques to prevent
+    memory leaks during long training runs. Logs completion tables to SwanLab
+    for qualitative analysis of model outputs.
+
+    Example usage:
+        >>> logger = CompletionLogger(maxlen=128)
+        >>> logger.add_dpo_completion(
+        ...     step=0,
+        ...     prompt="What is AI?",
+        ...     chosen="Artificial Intelligence is...",
+        ...     rejected="AI means...",
+        ...     reward_diff=0.5
+        ... )
+        >>> logger.log_to_swanlab()
+
+    Attributes:
+        maxlen: Maximum number of completions to store (older ones are dropped)
+        data: Deque storing completion dictionaries
+    """
+
+    def __init__(self, maxlen: int = 128):
+        """Initialize completion logger with bounded buffer.
+
+        Args:
+            maxlen: Maximum number of completions to store. When the buffer
+                is full, oldest completions are automatically discarded.
+                Default: 128 (sufficient for most RLHF runs without memory issues)
+        """
+        self.maxlen = maxlen
+        self.data: deque[Mapping[str, Any]] = deque(maxlen=maxlen)
+
+    def add_dpo_completion(
+        self,
+        step: int,
+        prompt: str,
+        chosen: str,
+        rejected: str,
+        reward_diff: float | None = None,
+    ) -> None:
+        """Add a DPO completion to the buffer.
+
+        Args:
+            step: Training step number
+            prompt: Input prompt
+            chosen: Chosen (preferred) completion
+            rejected: Rejected (non-preferred) completion
+            reward_diff: Reward difference (chosen - rejected), if available
+        """
+        entry = {
+            "step": step,
+            "prompt": prompt,
+            "chosen": chosen,
+            "rejected": rejected,
+        }
+        if reward_diff is not None:
+            entry["reward_diff"] = reward_diff
+
+        self.data.append(entry)
+
+    def add_kto_completion(
+        self,
+        step: int,
+        prompt: str,
+        completion: str,
+        label: bool,
+        reward: float | None = None,
+    ) -> None:
+        """Add a KTO completion to the buffer.
+
+        Args:
+            step: Training step number
+            prompt: Input prompt
+            completion: Model-generated completion
+            label: True if desirable, False if undesirable
+            reward: Reward score, if available
+        """
+        entry = {
+            "step": step,
+            "prompt": prompt,
+            "completion": completion,
+            "label": "desirable" if label else "undesirable",
+        }
+        if reward is not None:
+            entry["reward"] = reward
+
+        self.data.append(entry)
+
+    def add_orpo_completion(
+        self,
+        step: int,
+        prompt: str,
+        chosen: str,
+        rejected: str,
+        log_odds_ratio: float | None = None,
+    ) -> None:
+        """Add an ORPO completion to the buffer.
+
+        Args:
+            step: Training step number
+            prompt: Input prompt
+            chosen: Chosen (preferred) completion
+            rejected: Rejected (non-preferred) completion
+            log_odds_ratio: Log odds ratio between chosen and rejected
+        """
+        entry = {
+            "step": step,
+            "prompt": prompt,
+            "chosen": chosen,
+            "rejected": rejected,
+        }
+        if log_odds_ratio is not None:
+            entry["log_odds_ratio"] = log_odds_ratio
+
+        self.data.append(entry)
+
+    def add_grpo_completion(
+        self,
+        step: int,
+        prompt: str,
+        completion: str,
+        reward: float | None = None,
+        advantage: float | None = None,
+    ) -> None:
+        """Add a GRPO completion to the buffer.
+
+        Args:
+            step: Training step number
+            prompt: Input prompt
+            completion: Model-generated completion
+            reward: Reward score from reward model
+            advantage: Advantage estimate (reward - baseline)
+        """
+        entry = {
+            "step": step,
+            "prompt": prompt,
+            "completion": completion,
+        }
+        if reward is not None:
+            entry["reward"] = reward
+        if advantage is not None:
+            entry["advantage"] = advantage
+
+        self.data.append(entry)
+
+    def log_to_swanlab(self, table_name: str = "completions") -> bool:
+        """Log buffered completions to SwanLab as a table.
+
+        Creates a SwanLab echarts Table with all buffered completions.
+        Only logs if SwanLab is initialized and data is available.
+
+        Args:
+            table_name: Name of the table in SwanLab dashboard.
+                Default: "completions"
+
+        Returns:
+            True if logging succeeded, False otherwise
+        """
+        if not self.data:
+            LOG.debug("No completions to log to SwanLab")
+            return False
+
+        try:
+            import swanlab
+
+            if swanlab.get_run() is None:
+                LOG.debug("SwanLab not initialized, skipping completion logging")
+                return False
+
+            # Convert deque to list of dicts
+            completions = list(self.data)
+
+            # Extract headers from first entry (all entries should have same structure)
+            headers = list(completions[0].keys())
+
+            # Build rows: each completion becomes one row
+            rows = []
+            for completion in completions:
+                row = [completion.get(header, "") for header in headers]
+                rows.append(row)
+
+            # Log to SwanLab as echarts Table
+            swanlab.log({table_name: swanlab.echarts.Table().add(headers, rows)})
+
+            LOG.info(f"Logged {len(rows)} completions to SwanLab table '{table_name}'")
+            return True
+
+        except ImportError:
+            LOG.warning(
+                "SwanLab not installed, cannot log completions. "
+                "Install with: pip install swanlab"
+            )
+            return False
+        except Exception as err:  # pylint: disable=broad-except
+            LOG.exception("Failed to log completions to SwanLab: %s", err)
+            return False
+
+    def clear(self) -> None:
+        """Clear all buffered completions."""
+        self.data.clear()
+
+    def __len__(self) -> int:
+        """Return number of buffered completions."""
+        return len(self.data)
+
+    def __repr__(self) -> str:
+        """String representation showing buffer status."""
+        return (
+            f"CompletionLogger(maxlen={self.maxlen}, "
+            f"buffered={len(self.data)}/{self.maxlen})"
+        )
diff --git a/src/axolotl/integrations/swanlab/plugins.py b/src/axolotl/integrations/swanlab/plugins.py
new file mode 100644
index 0000000000..55f19ac59b
--- /dev/null
+++ b/src/axolotl/integrations/swanlab/plugins.py
@@ -0,0 +1,556 @@
+"""SwanLab Plugin for Axolotl"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from transformers import TrainerCallback
+
+    from axolotl.utils.dict import DictDefault
+
+LOG = get_logger(__name__)
+
+
+class SwanLabPlugin(BasePlugin):
+    """
+    SwanLab integration plugin for Axolotl.
+
+    Provides experiment tracking, visualization, and logging capabilities
+    using SwanLab (https://swanlab.cn).
+
+    Usage in config.yaml:
+        plugins:
+          - axolotl.integrations.swanlab.SwanLabPlugin
+
+        use_swanlab: true
+        swanlab_project: my-project
+        swanlab_experiment_name: my-experiment
+        swanlab_mode: cloud  # or 'local', 'offline', 'disabled'
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.swanlab_initialized = False
+        LOG.info("SwanLab plugin initialized")
+
+    def get_input_args(self) -> str:
+        """Returns the configuration model for SwanLab integration."""
+        return "axolotl.integrations.swanlab.SwanLabConfig"
+
+    def register(self, cfg: dict):
+        """Register SwanLab plugin with configuration and conflict detection."""
+        LOG.info("Registering SwanLab plugin")
+
+        # === Conflict Detection: Required Fields ===
+
+        # Check if SwanLab is enabled
+        if cfg.get("use_swanlab"):
+            # 1. Validate project name is set
+            if not cfg.get("swanlab_project"):
+                raise ValueError(
+                    "SwanLab enabled but 'swanlab_project' is not set.\n\n"
+                    "Solutions:\n"
+                    "  1. Add 'swanlab_project: your-project-name' to your config\n"
+                    "  2. Set 'use_swanlab: false' to disable SwanLab\n\n"
+                    "See: src/axolotl/integrations/swanlab/README.md for examples"
+                )
+
+            # 2. Validate swanlab_mode value
+            valid_modes = ["cloud", "local", "offline", "disabled"]
+            mode = cfg.get("swanlab_mode")
+            if mode and mode not in valid_modes:
+                raise ValueError(
+                    f"Invalid swanlab_mode: '{mode}'.\n\n"
+                    f"Valid options: {', '.join(valid_modes)}\n\n"
+                    f"Example:\n"
+                    f"  swanlab_mode: cloud  # Sync to SwanLab cloud\n"
+                    f"  swanlab_mode: local  # Local only, no cloud sync\n"
+                )
+
+            # 3. Check API key for cloud mode
+            import os
+
+            mode = cfg.get("swanlab_mode", "cloud")  # Default is cloud
+            if mode == "cloud":
+                api_key = cfg.get("swanlab_api_key") or os.environ.get(
+                    "SWANLAB_API_KEY"
+                )
+                if not api_key:
+                    LOG.warning(
+                        "SwanLab cloud mode enabled but no API key found.\n"
+                        "SwanLab may fail to initialize during training.\n\n"
+                        "Solutions:\n"
+                        "  1. Set SWANLAB_API_KEY environment variable:\n"
+                        "     export SWANLAB_API_KEY=your-api-key\n"
+                        "  2. Add 'swanlab_api_key: your-api-key' to config (less secure)\n"
+                        "  3. Run 'swanlab login' before training\n"
+                        "  4. Use 'swanlab_mode: local' for offline tracking\n"
+                    )
+
+        # === Conflict Detection: Multi-Logger Performance Warning ===
+
+        # Detect all active logging tools
+        active_loggers = []
+        if cfg.get("use_wandb"):
+            active_loggers.append("WandB")
+        if cfg.get("use_mlflow"):
+            active_loggers.append("MLflow")
+        if cfg.get("comet_api_key") or cfg.get("comet_project_name"):
+            active_loggers.append("Comet")
+        if cfg.get("use_swanlab"):
+            active_loggers.append("SwanLab")
+
+        if len(active_loggers) > 1:
+            LOG.warning(
+                f"\n{'=' * 70}\n"
+                f"Multiple logging tools enabled: {', '.join(active_loggers)}\n"
+                f"{'=' * 70}\n"
+                f"This may cause:\n"
+                f"  - Performance overhead (~1-2% per logger, cumulative)\n"
+                f"  - Increased memory usage\n"
+                f"  - Longer training time per step\n"
+                f"  - Potential config/callback conflicts\n\n"
+                f"Recommendations:\n"
+                f"  - Choose ONE primary logging tool for production training\n"
+                f"  - Use multiple loggers only for:\n"
+                f"    * Migration period (transitioning between tools)\n"
+                f"    * Short comparison runs\n"
+                f"    * Debugging specific tool issues\n"
+                f"  - Monitor system resources (CPU, memory) during training\n"
+                f"{'=' * 70}\n"
+            )
+
+            if len(active_loggers) >= 3:
+                LOG.error(
+                    f"\n{'!' * 70}\n"
+                    f"WARNING: {len(active_loggers)} logging tools enabled simultaneously!\n"
+                    f"{'!' * 70}\n"
+                    f"This is likely unintentional and WILL significantly impact performance.\n"
+                    f"Expected overhead: ~{len(active_loggers) * 1.5:.1f}% per training step.\n\n"
+                    f"STRONGLY RECOMMEND:\n"
+                    f"  - Disable all but ONE logging tool\n"
+                    f"  - Use config inheritance to manage multiple configs\n"
+                    f"{'!' * 70}\n"
+                )
+
+        # === Auto-Enable Logic ===
+
+        # Enable SwanLab if project is specified
+        if cfg.get("swanlab_project") and not cfg.get("use_swanlab"):
+            cfg["use_swanlab"] = True
+            LOG.info("Automatically enabled use_swanlab because swanlab_project is set")
+
+    def pre_model_load(self, cfg: DictDefault):
+        """Initialize SwanLab before model loading with runtime checks."""
+        if not cfg.use_swanlab:
+            return
+
+        # === Runtime Check: Import Availability ===
+        try:
+            import swanlab
+        except ImportError as err:
+            raise ImportError(
+                "SwanLab is not installed.\n\n"
+                "Install with:\n"
+                "  pip install swanlab\n\n"
+                "Or add to requirements:\n"
+                "  swanlab>=0.3.0\n\n"
+                f"Original error: {err}"
+            ) from err
+
+        # Log SwanLab version
+        try:
+            swanlab_version = swanlab.__version__
+            LOG.info(f"SwanLab version: {swanlab_version}")
+        except AttributeError:
+            LOG.warning("Could not determine SwanLab version")
+
+        # === Runtime Check: Distributed Training Setup ===
+        from axolotl.utils.distributed import get_world_size, is_main_process
+
+        world_size = get_world_size()
+        if world_size > 1:
+            mode = getattr(cfg, "swanlab_mode", "cloud")
+            LOG.info(
+                f"\n{'=' * 70}\n"
+                f"Distributed training detected (world_size={world_size})\n"
+                f"SwanLab mode: {mode}\n"
+                f"{'=' * 70}\n"
+                f"Behavior:\n"
+                f"  - Only rank 0 will initialize SwanLab\n"
+                f"  - Other ranks will skip SwanLab to avoid conflicts\n"
+            )
+
+            if mode == "cloud":
+                LOG.info(
+                    f"  - Only rank 0 will upload to SwanLab cloud\n"
+                    f"  - Other ranks run without SwanLab overhead\n"
+                    f"{'=' * 70}\n"
+                )
+
+        # Only initialize SwanLab on the main process (rank 0)
+        # to avoid creating multiple runs in distributed training
+        if not is_main_process():
+            LOG.debug("Skipping SwanLab initialization on non-main process")
+            return
+
+        # Initialize SwanLab run (passing all params directly to init)
+        try:
+            init_kwargs = self._get_swanlab_init_kwargs(cfg)
+            swanlab.init(**init_kwargs)
+            self.swanlab_initialized = True
+            LOG.info(f"SwanLab initialized with project: {cfg.swanlab_project}")
+
+            # Register Lark notification callback (if configured)
+            self._register_lark_callback(cfg)
+
+            # Log configuration (with error handling)
+            try:
+                config_dict = self._prepare_config_for_logging(cfg)
+                swanlab.config.update(config_dict)
+                LOG.debug("Successfully logged config to SwanLab")
+            except Exception as config_err:  # pylint: disable=broad-except
+                LOG.warning(
+                    f"Failed to log config to SwanLab: {config_err}. Continuing anyway."
+                )
+
+        except Exception as err:  # pylint: disable=broad-except
+            LOG.exception("Failed to initialize SwanLab: %s", err)
+            self.swanlab_initialized = False
+
+    def add_callbacks_pre_trainer(self, cfg: DictDefault, model):
+        """Add SwanLab callbacks before trainer creation."""
+        callbacks: list[TrainerCallback] = []
+
+        if not cfg.use_swanlab:
+            return callbacks
+
+        if not self.swanlab_initialized:
+            LOG.warning("SwanLab not initialized, skipping callback registration")
+            return callbacks
+
+        try:
+            from axolotl.utils.callbacks.swanlab import (
+                CustomSwanLabCallback,
+                SaveAxolotlConfigtoSwanLabCallback,
+            )
+
+            # Add our custom lightweight SwanLabCallback
+            # (avoids omegaconf/antlr4 version conflicts)
+            swanlab_callback = CustomSwanLabCallback()
+            callbacks.append(swanlab_callback)
+            LOG.info("Added CustomSwanLabCallback for metrics logging")
+
+            # Add Axolotl config logging callback
+            if cfg.axolotl_config_path:
+                config_callback = SaveAxolotlConfigtoSwanLabCallback(
+                    cfg.axolotl_config_path
+                )
+                callbacks.append(config_callback)
+                LOG.info("Added SaveAxolotlConfigtoSwanLabCallback")
+
+        except ImportError as err:
+            LOG.exception("Failed to import SwanLab callbacks: %s", err)
+
+        return callbacks
+
+    def post_trainer_create(self, cfg: DictDefault, trainer):
+        """Post-trainer creation hook."""
+        if cfg.use_swanlab and self.swanlab_initialized:
+            try:
+                import swanlab
+
+                # Log additional trainer information (with safe conversion)
+                trainer_config = {
+                    "total_steps": int(trainer.state.max_steps)
+                    if trainer.state.max_steps
+                    else None,
+                    "num_train_epochs": float(trainer.args.num_train_epochs)
+                    if trainer.args.num_train_epochs
+                    else None,
+                    "train_batch_size": int(trainer.args.train_batch_size)
+                    if hasattr(trainer.args, "train_batch_size")
+                    else None,
+                    "gradient_accumulation_steps": int(
+                        trainer.args.gradient_accumulation_steps
+                    )
+                    if trainer.args.gradient_accumulation_steps
+                    else None,
+                }
+                # Remove None values
+                trainer_config = {
+                    k: v for k, v in trainer_config.items() if v is not None
+                }
+
+                if trainer_config:
+                    swanlab.config.update(trainer_config)
+                    LOG.info("Logged trainer configuration to SwanLab")
+            except Exception as err:  # pylint: disable=broad-except
+                LOG.debug(f"Failed to log trainer config to SwanLab: {err}")
+
+            # Register RLHF completion logging callback if enabled
+            self._register_completion_callback(cfg, trainer)
+
+    def _get_swanlab_init_kwargs(self, cfg: DictDefault) -> dict:
+        """Prepare kwargs for swanlab.init().
+
+        Passes all configuration parameters directly to swanlab.init()
+        instead of using environment variables as an intermediate layer.
+
+        Returns:
+            dict: Keyword arguments for swanlab.init()
+        """
+        init_kwargs = {}
+
+        # Project name (required)
+        if cfg.swanlab_project:
+            init_kwargs["project"] = cfg.swanlab_project
+
+        # Experiment name
+        if cfg.swanlab_experiment_name:
+            init_kwargs["experiment_name"] = cfg.swanlab_experiment_name
+
+        # Description
+        if cfg.swanlab_description:
+            init_kwargs["description"] = cfg.swanlab_description
+
+        # Workspace (organization)
+        if cfg.swanlab_workspace:
+            init_kwargs["workspace"] = cfg.swanlab_workspace
+
+        # Mode: cloud, local, offline, disabled
+        if cfg.swanlab_mode:
+            init_kwargs["mode"] = cfg.swanlab_mode
+
+        # API key (pass directly instead of via env var)
+        if cfg.swanlab_api_key:
+            init_kwargs["api_key"] = cfg.swanlab_api_key
+
+        # Private deployment hosts (pass directly instead of via env var)
+        if cfg.swanlab_web_host:
+            init_kwargs["web_host"] = cfg.swanlab_web_host
+
+        if cfg.swanlab_api_host:
+            init_kwargs["api_host"] = cfg.swanlab_api_host
+
+        # Log model checkpoints (coming soon in SwanLab)
+        if cfg.swanlab_log_model:
+            init_kwargs["log_model"] = cfg.swanlab_log_model
+
+        # Custom branding - adds Axolotl identifier to SwanLab UI
+        # This helps identify runs from Axolotl vs other frameworks
+        init_kwargs["config"] = {"UPPERFRAME": "🦎 Axolotl"}
+
+        return init_kwargs
+
+    def _prepare_config_for_logging(self, cfg: DictDefault) -> dict:
+        """Prepare configuration dict for logging to SwanLab."""
+
+        def safe_convert(value):
+            """Convert value to JSON-serializable type."""
+            if value is None:
+                return None
+            if isinstance(value, (int, float, bool)):
+                return value
+            if isinstance(value, str):
+                return value
+            # Convert everything else to string
+            return str(value)
+
+        try:
+            # Extract important training parameters with safe conversion
+            config_dict = {
+                "base_model": safe_convert(getattr(cfg, "base_model", "")),
+                "model_type": safe_convert(getattr(cfg, "model_type", "")),
+                "sequence_len": safe_convert(getattr(cfg, "sequence_len", None)),
+                "micro_batch_size": safe_convert(
+                    getattr(cfg, "micro_batch_size", None)
+                ),
+                "gradient_accumulation_steps": safe_convert(
+                    getattr(cfg, "gradient_accumulation_steps", None)
+                ),
+                "num_epochs": safe_convert(getattr(cfg, "num_epochs", None)),
+                "max_steps": safe_convert(getattr(cfg, "max_steps", None)),
+                "learning_rate": safe_convert(getattr(cfg, "learning_rate", None)),
+                "lr_scheduler": safe_convert(getattr(cfg, "lr_scheduler", "")),
+                "optimizer": safe_convert(getattr(cfg, "optimizer", "")),
+                "warmup_ratio": safe_convert(getattr(cfg, "warmup_ratio", None)),
+                "weight_decay": safe_convert(getattr(cfg, "weight_decay", None)),
+                "seed": safe_convert(getattr(cfg, "seed", None)),
+                "bf16": safe_convert(getattr(cfg, "bf16", None)),
+                "tf32": safe_convert(getattr(cfg, "tf32", None)),
+                "attn_implementation": safe_convert(
+                    getattr(cfg, "attn_implementation", None)
+                ),
+                "sample_packing": safe_convert(getattr(cfg, "sample_packing", None)),
+            }
+
+            # Add FSDP/parallel config - only boolean flags
+            if hasattr(cfg, "fsdp_config") and cfg.fsdp_config:
+                config_dict["fsdp_enabled"] = True
+                config_dict["fsdp_version"] = safe_convert(
+                    getattr(cfg, "fsdp_version", None)
+                )
+
+            if hasattr(cfg, "deepspeed") and cfg.deepspeed:
+                config_dict["deepspeed_enabled"] = True
+
+            # Add context parallel info
+            if hasattr(cfg, "context_parallel_size"):
+                config_dict["context_parallel_size"] = safe_convert(
+                    getattr(cfg, "context_parallel_size", None)
+                )
+            if hasattr(cfg, "tensor_parallel_size"):
+                config_dict["tensor_parallel_size"] = safe_convert(
+                    getattr(cfg, "tensor_parallel_size", None)
+                )
+            if hasattr(cfg, "dp_shard_size"):
+                config_dict["dp_shard_size"] = safe_convert(
+                    getattr(cfg, "dp_shard_size", None)
+                )
+
+            # Remove None values and empty strings
+            config_dict = {
+                k: v
+                for k, v in config_dict.items()
+                if v is not None and v != "" and v != "None"
+            }
+
+            return config_dict
+        except Exception as err:  # pylint: disable=broad-except
+            LOG.warning(f"Failed to prepare config for logging: {err}")
+            # Return minimal config
+            try:
+                lr = getattr(cfg, "learning_rate", None)
+                lr_value = float(lr) if lr is not None else None
+            except (TypeError, ValueError):
+                lr_value = None
+            return {
+                "base_model": str(getattr(cfg, "base_model", "unknown")),
+                "learning_rate": lr_value,
+            }
+
+    def _register_lark_callback(self, cfg: DictDefault):
+        """Register Lark (Feishu) notification callback if configured.
+
+        Lark notifications enable sending training updates to team chat channels,
+        useful for production monitoring and team collaboration.
+
+        Args:
+            cfg: Configuration object with Lark webhook settings
+        """
+        # Check if Lark webhook URL is configured
+        lark_webhook_url = getattr(cfg, "swanlab_lark_webhook_url", None)
+        if not lark_webhook_url:
+            return  # Lark not configured, skip
+
+        try:
+            import swanlab
+            from swanlab.plugin.notification import LarkCallback
+
+            # Get optional secret for HMAC signature authentication
+            lark_secret = getattr(cfg, "swanlab_lark_secret", None)
+
+            # Create Lark callback with webhook URL and optional secret
+            lark_callback = LarkCallback(
+                webhook_url=lark_webhook_url,
+                secret=lark_secret,
+            )
+
+            # Register callback with SwanLab
+            swanlab.register_callbacks([lark_callback])
+
+            if lark_secret:
+                LOG.info(
+                    "Registered Lark notification callback with HMAC authentication"
+                )
+            else:
+                LOG.info("Registered Lark notification callback (no HMAC secret)")
+                LOG.warning(
+                    "Lark webhook has no secret configured. "
+                    "For production use, set 'swanlab_lark_secret' to enable HMAC signature verification."
+                )
+
+        except ImportError as err:
+            LOG.warning(
+                f"Failed to import SwanLab Lark plugin: {err}\n\n"
+                "Lark notifications require SwanLab >= 0.3.0 with plugin support.\n"
+                "Install with: pip install 'swanlab>=0.3.0'\n\n"
+                "Continuing without Lark notifications..."
+            )
+        except Exception as err:  # pylint: disable=broad-except
+            LOG.exception(
+                "Failed to register Lark callback: %s\n\n"
+                "Check your Lark webhook URL and secret configuration.\n"
+                "Continuing without Lark notifications...",
+                err,
+            )
+
+    def _register_completion_callback(self, cfg: DictDefault, trainer):
+        """Register RLHF completion logging callback if enabled and applicable.
+
+        This callback logs model completions (prompts, chosen/rejected responses,
+        rewards) to SwanLab during RLHF training for qualitative analysis.
+
+        Args:
+            cfg: Configuration object with completion logging settings
+            trainer: The trainer instance to add callback to
+        """
+        # Check if completion logging is enabled
+        log_completions = getattr(cfg, "swanlab_log_completions", True)
+        if not log_completions:
+            LOG.debug("SwanLab completion logging disabled by config")
+            return
+
+        # Check if trainer is an RLHF trainer
+        trainer_name = trainer.__class__.__name__
+        rlhf_trainers = ["DPO", "KTO", "ORPO", "GRPO", "CPO"]
+        is_rlhf_trainer = any(name in trainer_name for name in rlhf_trainers)
+
+        if not is_rlhf_trainer:
+            LOG.debug(
+                f"Trainer {trainer_name} is not an RLHF trainer, "
+                "skipping completion logging callback"
+            )
+            return
+
+        try:
+            from axolotl.integrations.swanlab.callbacks import (
+                SwanLabRLHFCompletionCallback,
+            )
+
+            # Get configuration parameters
+            log_interval = getattr(cfg, "swanlab_completion_log_interval", 100)
+            max_buffer = getattr(cfg, "swanlab_completion_max_buffer", 128)
+
+            # Create and register callback
+            completion_callback = SwanLabRLHFCompletionCallback(
+                log_interval=log_interval,
+                max_completions=max_buffer,
+                table_name="rlhf_completions",
+            )
+
+            trainer.add_callback(completion_callback)
+
+            LOG.info(
+                f"Registered SwanLab RLHF completion logging callback for {trainer_name} "
+                f"(log_interval={log_interval}, max_buffer={max_buffer})"
+            )
+
+        except ImportError as err:
+            LOG.warning(
+                f"Failed to import SwanLab completion callback: {err}\n\n"
+                "This is a bug - the callback should be available.\n"
+                "Please report this issue.\n\n"
+                "Continuing without completion logging..."
+            )
+        except Exception as err:  # pylint: disable=broad-except
+            LOG.exception(
+                "Failed to register SwanLab completion callback: %s\n\n"
+                "Continuing without completion logging...",
+                err,
+            )
diff --git a/src/axolotl/integrations/swanlab/profiling.py b/src/axolotl/integrations/swanlab/profiling.py
new file mode 100644
index 0000000000..61243c54ee
--- /dev/null
+++ b/src/axolotl/integrations/swanlab/profiling.py
@@ -0,0 +1,203 @@
+"""SwanLab profiling utilities for Axolotl trainers.
+
+This module provides decorators and context managers for profiling
+trainer methods and logging execution times to SwanLab.
+"""
+
+import time
+from contextlib import contextmanager
+from functools import wraps
+from typing import Any, Callable
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+@contextmanager
+def swanlab_profiling_context(trainer: Any, func_name: str):
+    """Context manager for profiling trainer methods.
+
+    Measures execution time and logs to SwanLab if enabled.
+
+    Example usage:
+        >>> with swanlab_profiling_context(self, "training_step"):
+        ...     result = do_expensive_computation()
+
+    Args:
+        trainer: Trainer instance (must have cfg attribute with use_swanlab flag)
+        func_name: Name of the function being profiled
+
+    Yields:
+        None
+    """
+    start_time = time.perf_counter()
+    try:
+        yield
+    finally:
+        duration = time.perf_counter() - start_time
+
+        # Check if SwanLab is enabled and initialized
+        use_swanlab = getattr(getattr(trainer, "cfg", None), "use_swanlab", False)
+        if use_swanlab:
+            try:
+                import swanlab
+
+                if swanlab.get_run() is not None:
+                    # Log profiling metric
+                    trainer_class = trainer.__class__.__name__
+                    metric_name = f"profiling/Time taken: {trainer_class}.{func_name}"
+
+                    swanlab.log({metric_name: duration})
+
+            except ImportError:
+                # SwanLab not installed, silently skip
+                pass
+            except Exception as err:  # pylint: disable=broad-except
+                # Log error but don't fail training
+                LOG.debug(f"Failed to log profiling metric for {func_name}: {err}")
+
+
+def swanlab_profile(func: Callable) -> Callable:
+    """Decorator to profile and log function execution time to SwanLab.
+
+    Automatically measures execution time of trainer methods and logs
+    to SwanLab as profiling metrics.
+
+    Example usage:
+        >>> class MyTrainer:
+        ...     @swanlab_profile
+        ...     def training_step(self, model, inputs):
+        ...         return super().training_step(model, inputs)
+
+    Args:
+        func: Function to profile (must be a method of a trainer instance)
+
+    Returns:
+        Wrapped function with profiling
+    """
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        with swanlab_profiling_context(self, func.__name__):
+            return func(self, *args, **kwargs)
+
+    return wrapper
+
+
+class ProfilingConfig:
+    """Configuration for SwanLab profiling.
+
+    This class provides a centralized way to control profiling behavior.
+
+    Attributes:
+        enabled: Whether profiling is enabled globally
+        min_duration_ms: Minimum duration (in ms) to log (filters out very fast ops)
+        log_interval: Log every N function calls (to reduce overhead)
+    """
+
+    def __init__(
+        self,
+        enabled: bool = True,
+        min_duration_ms: float = 0.1,
+        log_interval: int = 1,
+    ):
+        """Initialize profiling configuration.
+
+        Args:
+            enabled: Enable profiling. Default: True
+            min_duration_ms: Minimum duration to log (ms). Default: 0.1
+            log_interval: Log every N calls. Default: 1 (log all)
+        """
+        self.enabled = enabled
+        self.min_duration_ms = min_duration_ms
+        self.log_interval = log_interval
+        self._call_counts: dict[str, int] = {}
+
+    def should_log(self, func_name: str, duration_seconds: float) -> bool:
+        """Check if a profiling measurement should be logged.
+
+        Args:
+            func_name: Name of the profiled function
+            duration_seconds: Execution duration in seconds
+
+        Returns:
+            True if should log, False otherwise
+        """
+        if not self.enabled:
+            return False
+
+        # Check minimum duration threshold
+        duration_ms = duration_seconds * 1000
+        if duration_ms < self.min_duration_ms:
+            return False
+
+        # Check log interval
+        self._call_counts.setdefault(func_name, 0)
+        self._call_counts[func_name] += 1
+
+        # Always log on first call OR at intervals
+        count = self._call_counts[func_name]
+        if count == 1 or count % self.log_interval == 0:
+            return True
+
+        return False
+
+
+# Global profiling config (can be modified by users)
+DEFAULT_PROFILING_CONFIG = ProfilingConfig()
+
+
+@contextmanager
+def swanlab_profiling_context_advanced(
+    trainer: Any,
+    func_name: str,
+    config: ProfilingConfig | None = None,
+):
+    """Advanced profiling context with configurable behavior.
+
+    Similar to swanlab_profiling_context but with additional configuration
+    options for filtering and throttling profiling logs.
+
+    Example usage:
+        >>> config = ProfilingConfig(min_duration_ms=1.0, log_interval=10)
+        >>> with swanlab_profiling_context_advanced(self, "forward", config):
+        ...     output = model(inputs)
+
+    Args:
+        trainer: Trainer instance
+        func_name: Function name
+        config: Profiling configuration. If None, uses DEFAULT_PROFILING_CONFIG
+
+    Yields:
+        None
+    """
+    if config is None:
+        config = DEFAULT_PROFILING_CONFIG
+
+    start_time = time.perf_counter()
+    try:
+        yield
+    finally:
+        duration = time.perf_counter() - start_time
+
+        # Check if should log based on config
+        if config.should_log(func_name, duration):
+            # Check if SwanLab is enabled
+            use_swanlab = getattr(getattr(trainer, "cfg", None), "use_swanlab", False)
+            if use_swanlab:
+                try:
+                    import swanlab
+
+                    if swanlab.get_run() is not None:
+                        trainer_class = trainer.__class__.__name__
+                        metric_name = (
+                            f"profiling/Time taken: {trainer_class}.{func_name}"
+                        )
+
+                        swanlab.log({metric_name: duration})
+
+                except ImportError:
+                    pass
+                except Exception as err:  # pylint: disable=broad-except
+                    LOG.debug(f"Failed to log profiling metric for {func_name}: {err}")
diff --git a/src/axolotl/kernels/__init__.py b/src/axolotl/kernels/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/kernels/autotune_telemetry.py b/src/axolotl/kernels/autotune_telemetry.py
new file mode 100644
index 0000000000..0b0bf7838c
--- /dev/null
+++ b/src/axolotl/kernels/autotune_telemetry.py
@@ -0,0 +1,136 @@
+"""Telemetry for the fused RMSNorm+RoPE Triton autotune selections.
+
+Mirrors the scattermoe-lora autotune telemetry
+(:mod:`axolotl.integrations.kernels.autotune_callback`): after the kernel's
+``@triton.autotune`` cache is populated by the first backward pass, report the
+selected configs alongside GPU identity so the per-hardware tuning that varies
+across architectures can be aggregated.
+"""
+
+import torch
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# Give up looking for autotune data after this many training steps.
+_MAX_POLL_STEP = 5
+
+# (human-readable name, attribute on gemma4_fused_rope, autotune key arg names)
+_KERNEL_REGISTRY: list[tuple[str, str, list[str]]] = [
+    ("fused_rms_norm_rope_bwd", "_rms_norm_rope_backward_kernel", ["n_cols"]),
+]
+
+
+def _get_gpu_info() -> dict:
+    """Return basic GPU identification for the current device."""
+    if not torch.cuda.is_available():
+        return {}
+    try:
+        props = torch.cuda.get_device_properties(torch.cuda.current_device())
+        return {
+            "gpu_name": props.name,
+            "gpu_compute_capability": f"{props.major}.{props.minor}",
+            "gpu_memory_bytes": props.total_memory,
+        }
+    except Exception:  # pylint: disable=broad-exception-caught
+        return {}
+
+
+def collect_fused_rope_autotune_configs() -> list[dict]:
+    """Read the autotune ``.cache`` from the fused RMSNorm+RoPE backward kernel.
+
+    Each entry is ``{"kernel", "key", "config"}`` — the same shape the
+    scattermoe collector emits, so both event types aggregate uniformly.
+    Returns ``[]`` if Triton/the kernel isn't loaded or nothing autotuned yet.
+    """
+    import sys
+
+    # The kernel module is only in sys.modules once the fused path has run —
+    # which is exactly when its autotune cache is populated. Read it from there
+    # instead of importing (avoids pulling in Triton when the path is unused).
+    mod = sys.modules.get("axolotl.kernels.gemma4_fused_rope")
+    if mod is None:
+        return []
+
+    results: list[dict] = []
+    for friendly_name, attr_name, key_names in _KERNEL_REGISTRY:
+        kernel_fn = getattr(mod, attr_name, None)
+        cache = getattr(kernel_fn, "cache", None)
+        if not cache:
+            continue
+        for key_tuple, config in cache.items():
+            config_dict = dict(config.kwargs)
+            config_dict["num_warps"] = config.num_warps
+            config_dict["num_stages"] = config.num_stages
+            if getattr(config, "num_ctas", None) is not None:
+                config_dict["num_ctas"] = config.num_ctas
+
+            key: dict = {}
+            for i, name in enumerate(key_names):
+                if i < len(key_tuple):
+                    key[name] = key_tuple[i]
+            if len(key_tuple) > len(key_names):
+                key["_extra"] = [str(v) for v in key_tuple[len(key_names) :]]
+
+            results.append({"kernel": friendly_name, "key": key, "config": config_dict})
+    return results
+
+
+class FusedRopeAutotuneReportCallback(TrainerCallback):
+    """Reports fused RMSNorm+RoPE autotune selections via telemetry.
+
+    Fires once after the autotune cache is populated (the first step whose
+    backward has run), retrying up to ``_MAX_POLL_STEP`` then giving up. Every
+    later ``on_step_end`` short-circuits on ``_reported`` — zero hot-path cost.
+    """
+
+    def __init__(self):
+        self._reported = False
+
+    # pylint: disable=unused-argument
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        if self._reported:
+            return
+
+        configs = collect_fused_rope_autotune_configs()
+        if not configs:
+            if state.global_step >= _MAX_POLL_STEP:
+                LOG.debug(
+                    "No fused-rope autotune data after %d steps; giving up.",
+                    state.global_step,
+                )
+                self._reported = True
+            return
+
+        self._reported = True
+
+        from axolotl.telemetry.manager import TelemetryManager
+
+        telemetry_manager = TelemetryManager.get_instance()
+        if not telemetry_manager.enabled:
+            return
+
+        properties = {"kernel_count": len(configs), "kernels": configs}
+        properties.update(_get_gpu_info())
+
+        telemetry_manager.send_event(
+            event_type="fused-rope-autotune",
+            properties=properties,
+        )
+        LOG.info(
+            "Reported %d fused-rope kernel autotune config(s) to telemetry.",
+            len(configs),
+        )
diff --git a/src/axolotl/kernels/dora.py b/src/axolotl/kernels/dora.py
new file mode 100644
index 0000000000..3ed35cf740
--- /dev/null
+++ b/src/axolotl/kernels/dora.py
@@ -0,0 +1,147 @@
+"""
+Triton kernels for DoRA (Weight-Decomposed Low-Rank Adaptation).
+
+Fuses the weight norm computation and magnitude scaling to avoid
+materializing the full [out_features, in_features] combined weight matrix.
+The B@A product is computed row-by-row inside the kernel.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from .quantize import dequantize
+
+
+@triton.jit
+def _dora_fused_norm_kernel(
+    # Pointers
+    W_ptr,  # base weight [out, in] (dequantized, row-major)
+    B_ptr,  # LoRA B [out, rank] (row-major)
+    A_ptr,  # LoRA A [rank, in] (row-major)
+    mag_ptr,  # magnitude vector [out]
+    out_ptr,  # output mag_norm_scale [out]
+    # Shapes
+    out_features,
+    in_features,
+    rank,
+    # Scaling
+    lora_scale,  # float scaling factor
+    # Block sizes
+    BLOCK_IN: tl.constexpr,
+    BLOCK_R: tl.constexpr,  # >= rank, power of 2
+):
+    """Compute mag_norm_scale[i] = magnitude[i] / ||W[i,:] + s * (B[i,:] @ A)[:] ||_2
+
+    Each program handles one output row. B[row,:] is loaded once (small),
+    then we tile over in_features computing the dot product with A[:,tile]
+    and accumulating the squared norm.
+
+    This avoids materializing the full [out, in] B@A matrix.
+    """
+    row = tl.program_id(0)
+    if row >= out_features:
+        return
+
+    # Accumulate squared norm across tiles of in_features
+    norm_sq_acc = tl.zeros([BLOCK_IN], dtype=tl.float32)
+
+    for start in range(0, in_features, BLOCK_IN):
+        cols = start + tl.arange(0, BLOCK_IN)
+        col_mask = cols < in_features
+
+        # Load W[row, cols]
+        w_vals = tl.load(
+            W_ptr + row * in_features + cols,
+            mask=col_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        # Compute (B[row,:] @ A[:, cols]) for this tile
+        # Load B[row, r] as scalar and A[r, cols] as vector for each r
+        ba_vals = tl.zeros([BLOCK_IN], dtype=tl.float32)
+        for r in tl.static_range(BLOCK_R):
+            # Load scalar B[row, r]
+            b_val = tl.load(
+                B_ptr + row * rank + r,
+                mask=(r < rank),
+                other=0.0,
+            ).to(tl.float32)
+            # Load vector A[r, cols]
+            a_vals = tl.load(
+                A_ptr + r * in_features + cols,
+                mask=(col_mask & (r < rank)),
+                other=0.0,
+            ).to(tl.float32)
+            ba_vals += b_val * a_vals
+
+        # Combined: W + s * (B @ A)
+        combined = w_vals + lora_scale * ba_vals
+
+        # Accumulate squared values
+        norm_sq_acc += tl.where(col_mask, combined * combined, 0.0)
+
+    # Reduce to scalar norm
+    norm_sq = tl.sum(norm_sq_acc, axis=0)
+    norm = tl.sqrt(norm_sq + 1e-12)  # epsilon for numerical stability
+
+    # Load magnitude and compute scale
+    mag = tl.load(mag_ptr + row).to(tl.float32)
+    scale = mag / norm
+
+    tl.store(out_ptr + row, scale)
+
+
+def triton_dora_scale(
+    W: torch.Tensor,
+    W_quant,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    s: float,
+    magnitude: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Compute DoRA mag_norm_scale using fused Triton kernel.
+
+    Computes B@A row-by-row inside the kernel, avoiding the full
+    [out_features, in_features] materialization.
+
+    Args:
+        W: base weight [out, in] (possibly quantized)
+        W_quant: quantization state
+        A: LoRA A [rank, in]
+        B: LoRA B [out, rank]
+        s: LoRA scaling factor
+        magnitude: learned magnitude [out]
+        dtype: compute dtype
+
+    Returns:
+        mag_norm_scale: [out] tensor = magnitude / ||W + s * B @ A||_2
+    """
+    # Dequantize W to [out, in]
+    W_full = dequantize(W.t(), W_quant).t().contiguous().to(dtype)
+
+    out_features, in_features = W_full.shape
+    rank = A.shape[0]
+
+    out = torch.empty(out_features, dtype=dtype, device=W.device)
+
+    # Block sizes
+    BLOCK_IN = triton.next_power_of_2(min(in_features, 2048))
+    BLOCK_R = triton.next_power_of_2(rank)
+
+    _dora_fused_norm_kernel[(out_features,)](
+        W_full,
+        B.contiguous().to(dtype),
+        A.contiguous().to(dtype),
+        magnitude.contiguous(),
+        out,
+        out_features=out_features,
+        in_features=in_features,
+        rank=rank,
+        lora_scale=s,
+        BLOCK_IN=BLOCK_IN,
+        BLOCK_R=BLOCK_R,
+    )
+
+    return out.detach()
diff --git a/src/axolotl/kernels/geglu.py b/src/axolotl/kernels/geglu.py
new file mode 100644
index 0000000000..38f7e060a8
--- /dev/null
+++ b/src/axolotl/kernels/geglu.py
@@ -0,0 +1,175 @@
+"""Module for definition of GEGLU Triton kernels.
+
+See "GLU Variants Improve Transformer" (https://arxiv.org/abs/2002.05202).
+
+Credit to `unsloth` (https://unsloth.ai/) for inspiration for this implementation.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+
+@triton.jit
+def _geglu_fwd_kernel(
+    gate_ptr,
+    up_ptr,
+    out_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """GEGLU forward kernel.
+
+    Args:
+        gate_ptr: Pointer to gate tensor [*, hidden_dim].
+        up_ptr: Pointer to up-projection tensor [*, hidden_dim].
+        out_ptr: Pointer to output tensor [*, hidden_dim].
+        n_elements: Total number of elements in the input tensors.
+        BLOCK_SIZE: Size of thread blocks for parallel computation.
+    """
+    block_idx = tl.program_id(0)
+    offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    gate = tl.load(gate_ptr + offsets, mask=mask, other=0).to(tl.float32)
+    up = tl.load(up_ptr + offsets, mask=mask, other=0)
+
+    # Compute activation in fp32 then convert back
+    gelu_gate = 0.5 * gate * (tl.math.erf(tl.math.rsqrt(2.0) * gate) + 1.0)
+    gelu_gate = gelu_gate.to(up.dtype)
+    result = gelu_gate * up
+
+    tl.store(out_ptr + offsets, result, mask=mask)
+
+
+@register_kernel_op("geglu_fwd")
+def geglu_forward(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    """GEGLU forward pass.
+
+    Args:
+        gate: Input gate tensor of shape [batch, seq_len, hidden_dim].
+        up: Up-projection tensor of shape [batch, seq_len, hidden_dim].
+
+    Returns:
+        torch.Tensor: Output tensor of shape [batch, seq_len, hidden_dim].
+    """
+    batch, seq_len, hidden_dim = gate.shape
+    n_elements = gate.numel()
+    out = torch.empty((batch, seq_len, hidden_dim), dtype=gate.dtype, device="cuda")
+
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)  # noqa: E731
+    _geglu_fwd_kernel[grid](
+        gate_ptr=gate,
+        up_ptr=up,
+        out_ptr=out,
+        n_elements=n_elements,
+        BLOCK_SIZE=1024,
+    )
+    return out
+
+
+@triton.jit
+def _geglu_bwd_kernel(
+    grad_out_ptr,
+    gate_ptr,
+    up_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """GEGLU backward kernel. Stores gradient results in-place.
+
+    Args:
+        grad_out_ptr: Pointer to gradient output tensor [*, hidden_dim].
+        gate_ptr: Pointer to gate tensor [*, hidden_dim].
+        up_ptr: Pointer to up-projection tensor [*, hidden_dim].
+        n_elements: Total number of elements in the input tensors.
+        BLOCK_SIZE: Size of thread blocks for parallel computation.
+
+    Note:
+        After kernel execution, tensors are modified in-place:
+        - `grad_out_ptr` contains GEGLU activation output (`h`)
+        - `gate_ptr` contains gradient w.r.t gate (`grad_gate`)
+        - `up_ptr` contains gradient w.r.t up (`grad_up`)
+    """
+    block_idx = tl.program_id(0)
+    offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    grad_out = tl.load(grad_out_ptr + offsets, mask=mask, other=0)
+    gate = tl.load(gate_ptr + offsets, mask=mask, other=0).to(tl.float32)
+    up = tl.load(up_ptr + offsets, mask=mask, other=0)
+
+    # Forward pass
+    gelu_partial = 0.5 * (tl.math.erf(tl.math.rsqrt(2.0) * gate) + 1.0)
+    gelu_gate = gelu_partial * gate
+    gelu_gate = gelu_gate.to(grad_out.dtype)
+
+    # Forward output
+    h = gelu_gate * up
+
+    # Compute gradients
+    grad_up = grad_out * gelu_gate
+
+    # Compute gate gradient using GELU derivative
+    temp = grad_out * up
+    t = 0.3989422804014327  # 1/sqrt(2*pi)
+    dgelu_dgate = gelu_partial + t * gate * tl.exp(-0.5 * gate * gate)
+    grad_gate = temp.to(tl.float32) * dgelu_dgate
+    grad_gate = grad_gate.to(grad_out.dtype)
+
+    # Store results
+    tl.store(grad_out_ptr + offsets, h, mask=mask)
+    tl.store(gate_ptr + offsets, grad_gate, mask=mask)
+    tl.store(up_ptr + offsets, grad_up, mask=mask)
+
+
+@geglu_forward.register_fake
+def _(gate, up):
+    return torch.empty_like(gate)
+
+
+@register_kernel_op("geglu_bwd", mutates_args=("grad_output", "gate", "up"))
+def _geglu_bwd_op(
+    grad_output: torch.Tensor, gate: torch.Tensor, up: torch.Tensor
+) -> None:
+    n_elements = grad_output.numel()
+
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)  # noqa: E731
+    _geglu_bwd_kernel[grid](
+        grad_out_ptr=grad_output,
+        gate_ptr=gate,
+        up_ptr=up,
+        n_elements=n_elements,
+        BLOCK_SIZE=1024,
+    )
+
+
+@_geglu_bwd_op.register_fake
+def _(grad_output, gate, up):
+    return None
+
+
+def geglu_backward(
+    grad_output: torch.Tensor, gate: torch.Tensor, up: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """GEGLU backward pass using in-place operations.
+
+    Args:
+        grad_output: Gradient of loss with respect to output, shape `[batch, seq_len, hidden_dim]`.
+        gate: Gate tensor from forward pass, shape `[batch, seq_len, hidden_dim]`.
+        up: Up-projection tensor from forward pass, shape `[batch, seq_len, hidden_dim]`.
+
+    Returns:
+        Tuple containing:
+            - GEGLU activation output (`h`)
+            - Gradient with respect to gate (`grad_gate`)
+            - Gradient with respect to up (`grad_up`)
+
+    Note:
+        This function modifies its input tensors in-place to store results.
+    """
+    # op mutates in place; return the aliases here — op outputs must not alias inputs
+    _geglu_bwd_op(grad_output, gate, up)
+    return grad_output, gate, up
diff --git a/src/axolotl/kernels/gemma4_fused_rope.py b/src/axolotl/kernels/gemma4_fused_rope.py
new file mode 100644
index 0000000000..8193355846
--- /dev/null
+++ b/src/axolotl/kernels/gemma4_fused_rope.py
@@ -0,0 +1,580 @@
+"""Fused RMSNorm + (partial) RoPE Triton kernel for Gemma 4 / Qwen3 Q/K paths."""
+
+import math
+import operator
+
+import torch
+import triton
+import triton.language as tl
+from liger_kernel.ops.utils import (
+    calculate_settings,
+    compare_version,
+    torch_to_triton_dtype,
+)
+from liger_kernel.utils import is_npu_available
+from torch.library import triton_op, wrap_triton
+
+if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
+    try:
+        from triton.language.extra.libdevice import rsqrt
+    except ModuleNotFoundError:
+        from triton.language.extra.cuda.libdevice import rsqrt
+else:
+    from triton.language.math import rsqrt
+
+# Backward over-subscription factor: number of program blocks per SM. The
+# weight-gradient reduction needs one private partial per block, so this also
+# sizes the dW scratch buffer. ~8 saturates occupancy on tested GPUs.
+_BWD_BLOCKS_PER_SM = 8
+
+
+@triton.jit
+def _rms_norm_rope_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    W_ptr,
+    COS_ptr,
+    COS_row_stride,
+    SIN_ptr,
+    SIN_row_stride,
+    RSTD_ptr,
+    RSTD_row_stride,
+    n_cols,
+    n_rot,
+    n_heads,
+    eps,
+    HAS_WEIGHT: tl.constexpr,
+    UNIT_OFFSET: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Fused forward:
+      x_norm = x / rms(x) [* weight]   (RMSNorm, full n_cols)
+      y[..., :n_rot]  = rope(x_norm[..., :n_rot])
+      y[..., n_rot:]  = x_norm[..., n_rot:]   (pass-through for partial rotary)
+
+    rotate_half swaps first/second halves and negates the first, restricted
+    to the rotary span [0, n_rot):
+      rotate_half([a, b]) = [-b, a]   where len(a) = len(b) = n_rot/2
+
+    For the partial-rotary pass-through region we load cos with default 1.0
+    and sin with default 0.0 outside [0, n_rot), so the same formula
+    `Y = X_norm * cos + X_rot_norm * sin` collapses to `Y = X_norm`.
+
+    cos/sin are indexed by row_idx // n_heads to handle per-head broadcast
+    (cos/sin have shape (B*S, n_rot) while X has shape (B*S*H, n_cols)).
+    """
+    row_idx = tl.program_id(0).to(tl.int64)
+    # cos/sin row: divide by n_heads since cos/sin are (B*S, n_rot)
+    cs_row_idx = row_idx // n_heads
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+    rot_mask_col = col_offsets < n_rot
+    half_rot = n_rot // 2
+
+    # Load input row
+    X_row = tl.load(X_ptr + row_idx * X_row_stride + col_offsets, mask=mask, other=0)
+    X_dtype = X_row.dtype
+    X_fp32 = X_row.to(tl.float32)
+
+    # RMSNorm: compute 1/rms over the full row (rotary + pass-through)
+    mean_sq = tl.sum(X_fp32 * X_fp32, axis=0) / n_cols
+    rstd = rsqrt(mean_sq + eps)
+    tl.store(RSTD_ptr + row_idx * RSTD_row_stride, rstd)
+
+    # Normalize
+    X_norm = X_fp32 * rstd
+
+    # Apply weight if present (with_scale=True)
+    if HAS_WEIGHT:
+        W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
+        if UNIT_OFFSET:
+            X_norm = X_norm * (W_row + 1.0)
+        else:
+            X_norm = X_norm * W_row
+
+    # RoPE: load cos/sin (broadcast across heads). For col >= n_rot we get
+    # cos=1, sin=0 so the formula leaves X_norm untouched.
+    cos_row = tl.load(
+        COS_ptr + cs_row_idx * COS_row_stride + col_offsets,
+        mask=rot_mask_col,
+        other=1.0,
+    ).to(tl.float32)
+    sin_row = tl.load(
+        SIN_ptr + cs_row_idx * SIN_row_stride + col_offsets,
+        mask=rot_mask_col,
+        other=0.0,
+    ).to(tl.float32)
+
+    # rotate_half within [0, n_rot):
+    #   for col < half_rot:  take -X_norm[col + half_rot]
+    #   for col in [half_rot, n_rot): take  X_norm[col - half_rot]
+    # For col >= n_rot the rotation is irrelevant (sin = 0 zeros it out).
+    rot_offsets = tl.where(
+        col_offsets < half_rot, col_offsets + half_rot, col_offsets - half_rot
+    )
+    rot_load_mask = (rot_offsets < n_cols) & rot_mask_col
+    X_rot = tl.load(
+        X_ptr + row_idx * X_row_stride + rot_offsets, mask=rot_load_mask, other=0
+    ).to(tl.float32)
+    # Re-normalize the rotated values
+    X_rot_norm = X_rot * rstd
+    if HAS_WEIGHT:
+        W_rot = tl.load(W_ptr + rot_offsets, mask=rot_load_mask, other=0).to(tl.float32)
+        if UNIT_OFFSET:
+            X_rot_norm = X_rot_norm * (W_rot + 1.0)
+        else:
+            X_rot_norm = X_rot_norm * W_rot
+
+    # Negate the first half (rotate_half negates x2, which becomes the first half)
+    sign = tl.where(col_offsets < half_rot, -1.0, 1.0)
+    X_rot_norm = X_rot_norm * sign
+
+    # Final RoPE: y = x_norm * cos + rotate_half(x_norm) * sin
+    Y_row = X_norm * cos_row + X_rot_norm * sin_row
+
+    tl.store(
+        Y_ptr + row_idx * Y_row_stride + col_offsets,
+        Y_row.to(X_dtype),
+        mask=mask,
+    )
+
+
+_BWD_AUTOTUNE_CONFIGS = [
+    triton.Config({}, num_warps=w, num_stages=s)
+    for w in (2, 4, 8, 16)
+    for s in (1, 2, 3)
+]
+
+
+# num_warps/num_stages optima for the latency-bound row loop vary by GPU; key on
+# n_cols (head_dim) so head_dim=128 and 256 each get their own tuned config.
+@triton.autotune(configs=_BWD_AUTOTUNE_CONFIGS, key=["n_cols"])
+@triton.jit
+def _rms_norm_rope_backward_kernel(
+    dY_ptr,
+    dY_row_stride,
+    dX_ptr,
+    dX_row_stride,
+    X_ptr,
+    X_row_stride,
+    X_dtype: tl.constexpr,
+    W_ptr,
+    COS_ptr,
+    COS_row_stride,
+    SIN_ptr,
+    SIN_row_stride,
+    RSTD_ptr,
+    RSTD_row_stride,
+    dW_ptr,
+    dW_row_stride,
+    n_rows,
+    n_cols,
+    n_rot,
+    n_heads,
+    rows_per_program,
+    HAS_WEIGHT: tl.constexpr,
+    UNIT_OFFSET: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Backward for Y = RoPE(RMSNorm(X, W)) with optional partial rotary
+    (`n_rot <= n_cols`).
+
+    For col < n_rot the standard RoPE adjoint applies. For col >= n_rot the
+    output is just the normalized row, so dN[col] = dY[col] (achieved by
+    loading cos with default 1.0 and forcing the rotate-half contribution
+    to zero outside the rotary span).
+
+    cos/sin indexed by row_idx // n_heads for per-head broadcast.
+    """
+    row_block_id = tl.program_id(0).to(tl.int64)
+    row_start = row_block_id * rows_per_program
+    row_end = min((row_block_id + 1) * rows_per_program, n_rows)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+    rot_mask_col = col_offsets < n_rot
+    half_rot = n_rot // 2
+
+    dW_acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+
+    if HAS_WEIGHT:
+        W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
+
+    for row_idx in range(row_start, row_end):
+        cs_row_idx = row_idx // n_heads
+
+        dY_row = tl.load(
+            dY_ptr + row_idx * dY_row_stride + col_offsets, mask=mask, other=0
+        ).to(tl.float32)
+        X_row = tl.load(
+            X_ptr + row_idx * X_row_stride + col_offsets, mask=mask, other=0
+        ).to(tl.float32)
+        rstd = tl.load(RSTD_ptr + row_idx * RSTD_row_stride)
+
+        cos_row = tl.load(
+            COS_ptr + cs_row_idx * COS_row_stride + col_offsets,
+            mask=rot_mask_col,
+            other=1.0,
+        ).to(tl.float32)
+
+        # dN = dY * cos + rotate_half^T(dY * sin)   (within the rotary span)
+        # rotate_half^T([a, b]) = [b, -a]  (adjoint of rotate_half)
+        #
+        # For col >= n_rot the formula must collapse to dN = dY (since the
+        # forward is just a pass-through). cos defaults to 1.0 above; the
+        # rotate-half contribution is masked to zero below.
+        rot_offsets = tl.where(
+            col_offsets < half_rot, col_offsets + half_rot, col_offsets - half_rot
+        )
+        rot_load_mask = (rot_offsets < n_cols) & rot_mask_col
+        dY_rot = tl.load(
+            dY_ptr + row_idx * dY_row_stride + rot_offsets,
+            mask=rot_load_mask,
+            other=0,
+        ).to(tl.float32)
+        sin_rot = tl.load(
+            SIN_ptr + cs_row_idx * SIN_row_stride + rot_offsets,
+            mask=rot_load_mask,
+            other=0,
+        ).to(tl.float32)
+
+        adj_sign = tl.where(col_offsets < half_rot, 1.0, -1.0)
+        rotate_term = dY_rot * sin_rot * adj_sign
+        # Zero out rotate-half contribution outside the rotary span.
+        rotate_term = tl.where(rot_mask_col, rotate_term, 0.0)
+        dN = dY_row * cos_row + rotate_term
+
+        # Pre-weight normalized: n = rstd * x
+        n = X_row * rstd
+
+        if HAS_WEIGHT:
+            dW_acc += dN * n
+            if UNIT_OFFSET:
+                dm = dN * (W_row + 1.0)
+            else:
+                dm = dN * W_row
+        else:
+            dm = dN
+
+        # RMSNorm backward: dX = rstd * (dm - (1/n_cols) * rstd^2 * dot(dm, X) * X)
+        dot_dm_x = tl.sum(dm * X_row, axis=0)
+        dX_row = rstd * (dm - (1.0 / n_cols) * rstd * rstd * dot_dm_x * X_row)
+
+        tl.store(
+            dX_ptr + row_idx * dX_row_stride + col_offsets,
+            dX_row.to(X_dtype),
+            mask=mask,
+        )
+
+    if HAS_WEIGHT:
+        tl.store(
+            dW_ptr + row_block_id * dW_row_stride + col_offsets,
+            dW_acc,
+            mask=mask,
+        )
+
+
+@triton_op("axolotl::fused_rms_norm_rope_fwd", mutates_args=())
+def _fused_rms_norm_rope_fwd(
+    X: torch.Tensor,
+    W: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    n_heads: int,
+    n_rot: int,
+    unit_offset: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Returns ``(Y, RSTD)``; ``wrap_triton`` keeps it ``torch.compile``-safe."""
+    n_rows, n_cols = X.shape
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    Y = torch.empty_like(X)
+    RSTD = torch.empty(n_rows, dtype=torch.float32, device=X.device)
+    wrap_triton(_rms_norm_rope_forward_kernel)[(n_rows,)](
+        Y,
+        Y.stride(0),
+        X,
+        X.stride(0),
+        W,
+        cos,
+        cos.stride(0),
+        sin,
+        sin.stride(0),
+        RSTD,
+        RSTD.stride(0),
+        n_cols,
+        n_rot,
+        n_heads,
+        eps,
+        HAS_WEIGHT=True,
+        UNIT_OFFSET=unit_offset,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return Y, RSTD
+
+
+@triton_op("axolotl::fused_rms_norm_rope_bwd", mutates_args=())
+def _fused_rms_norm_rope_bwd(
+    dY: torch.Tensor,
+    X: torch.Tensor,
+    W: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    RSTD: torch.Tensor,
+    n_heads: int,
+    n_rot: int,
+    unit_offset: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Returns ``(dX, dW)``."""
+    n_rows, n_cols = dY.shape
+    BLOCK_SIZE, _ = calculate_settings(n_cols)
+    # One block per SM serializes a long row-loop at 1 block/SM occupancy; the
+    # forward runs a block per row. Over-subscribe the SMs so the latency-bound
+    # row loop has enough resident blocks to hide global-load latency. Each
+    # block still writes a private dW partial that's summed below (no atomics).
+    sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
+    target_programs = min(_BWD_BLOCKS_PER_SM * sm_count, n_rows)
+    rows_per_program = max(1, math.ceil(n_rows / target_programs))
+    n_programs = math.ceil(n_rows / rows_per_program)
+    dX = torch.empty_like(X)
+    _dW = torch.empty((n_programs, n_cols), dtype=torch.float32, device=X.device)
+    wrap_triton(_rms_norm_rope_backward_kernel)[(n_programs,)](
+        dY,
+        dY.stride(0),
+        dX,
+        dX.stride(0),
+        X,
+        X.stride(0),
+        torch_to_triton_dtype[X.dtype],
+        W,
+        cos,
+        cos.stride(0),
+        sin,
+        sin.stride(0),
+        RSTD,
+        RSTD.stride(0),
+        _dW,
+        _dW.stride(0),
+        n_rows,
+        n_cols,
+        n_rot,
+        n_heads,
+        rows_per_program,
+        HAS_WEIGHT=True,
+        UNIT_OFFSET=unit_offset,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    dW = _dW.sum(dim=0).to(W.dtype)
+    return dX, dW
+
+
+def _fused_rms_norm_rope_setup_context(ctx, inputs, output):
+    X, W, cos, sin, _eps, n_heads, n_rot, unit_offset = inputs
+    _, RSTD = output
+    ctx.save_for_backward(X, W, cos, sin, RSTD)
+    ctx.n_heads = n_heads
+    ctx.n_rot = n_rot
+    ctx.unit_offset = unit_offset
+
+
+def _fused_rms_norm_rope_backward(ctx, grad_Y, grad_RSTD):
+    X, W, cos, sin, RSTD = ctx.saved_tensors
+    grad_Y = grad_Y.contiguous()
+    dX, dW = _fused_rms_norm_rope_bwd(
+        grad_Y, X, W, cos, sin, RSTD, ctx.n_heads, ctx.n_rot, ctx.unit_offset
+    )
+    return dX, dW, None, None, None, None, None, None
+
+
+_fused_rms_norm_rope_fwd.register_autograd(
+    _fused_rms_norm_rope_backward,
+    setup_context=_fused_rms_norm_rope_setup_context,
+)
+
+
+def fused_rms_norm_rope(x, weight, cos, sin, eps=1e-6, unit_offset=False):
+    """
+    Apply fused RMSNorm + (partial) RoPE.
+
+    Shapes:
+        x:      (B, S, H, D) — post-projection
+        weight: (D,) — required; use ``fused_rms_norm_noscale`` for the no-weight variant
+        cos:    (B, S, n_rot) — ``n_rot`` must be even and ``<= D``; trailing
+                ``D - n_rot`` columns are RMSNorm-only (partial rotary).
+        sin:    (B, S, n_rot)
+
+    ``unit_offset=True`` scales by ``(weight + 1.0)`` (Gemma-style).
+    """
+    shape = x.shape  # (B, S, H, D)
+    B, S, H, D = shape
+    n_rot = cos.shape[-1]
+    if sin.shape[-1] != n_rot:
+        raise ValueError(
+            f"cos and sin must have the same last dim, got cos={cos.shape[-1]} "
+            f"sin={sin.shape[-1]}"
+        )
+    if n_rot > D:
+        raise ValueError(f"rotary dim ({n_rot}) cannot exceed head_dim ({D})")
+    if n_rot % 2 != 0:
+        raise ValueError(f"rotary dim must be even, got {n_rot}")
+
+    x_flat = x.reshape(-1, D).contiguous()
+    # Kernel needs a dense (B*S, n_rot) buffer; materialize the batch-broadcast.
+    if cos.shape[0] != B:
+        if cos.shape[0] != 1:
+            raise ValueError(
+                f"cos/sin batch dim ({cos.shape[0]}) must be 1 or equal "
+                f"to x batch dim ({B})"
+            )
+        cos = cos.expand(B, S, n_rot)
+        sin = sin.expand(B, S, n_rot)
+    cos_flat = cos.reshape(B * S, n_rot).contiguous()
+    sin_flat = sin.reshape(B * S, n_rot).contiguous()
+
+    y_flat, _ = _fused_rms_norm_rope_fwd(
+        x_flat, weight, cos_flat, sin_flat, eps, H, n_rot, unit_offset
+    )
+    return y_flat.view(shape)
+
+
+@triton.jit
+def _rms_norm_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    RSTD_ptr,
+    RSTD_row_stride,
+    n_cols,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """RMSNorm without scale weight: y = x / rms(x)"""
+    row_idx = tl.program_id(0).to(tl.int64)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    X_row = tl.load(X_ptr + row_idx * X_row_stride + col_offsets, mask=mask, other=0)
+    X_dtype = X_row.dtype
+    X_fp32 = X_row.to(tl.float32)
+
+    mean_sq = tl.sum(X_fp32 * X_fp32, axis=0) / n_cols
+    rstd = rsqrt(mean_sq + eps)
+    tl.store(RSTD_ptr + row_idx * RSTD_row_stride, rstd)
+
+    Y_row = X_fp32 * rstd
+    tl.store(Y_ptr + row_idx * Y_row_stride + col_offsets, Y_row.to(X_dtype), mask=mask)
+
+
+@triton.jit
+def _rms_norm_noscale_backward_kernel(
+    dY_ptr,
+    dY_row_stride,
+    dX_ptr,
+    dX_row_stride,
+    X_ptr,
+    X_row_stride,
+    X_dtype: tl.constexpr,
+    RSTD_ptr,
+    RSTD_row_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Backward for y = x * rstd (no weight)."""
+    row_idx = tl.program_id(0).to(tl.int64)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    dY_row = tl.load(
+        dY_ptr + row_idx * dY_row_stride + col_offsets, mask=mask, other=0
+    ).to(tl.float32)
+    X_row = tl.load(
+        X_ptr + row_idx * X_row_stride + col_offsets, mask=mask, other=0
+    ).to(tl.float32)
+    rstd = tl.load(RSTD_ptr + row_idx * RSTD_row_stride)
+
+    dot_dy_x = tl.sum(dY_row * X_row, axis=0)
+    dX_row = rstd * (dY_row - (1.0 / n_cols) * rstd * rstd * dot_dy_x * X_row)
+
+    tl.store(
+        dX_ptr + row_idx * dX_row_stride + col_offsets, dX_row.to(X_dtype), mask=mask
+    )
+
+
+@triton_op("axolotl::fused_rms_norm_noscale_fwd", mutates_args=())
+def _fused_rms_norm_noscale_fwd(
+    X: torch.Tensor, eps: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Returns ``(Y, RSTD)``."""
+    n_rows, n_cols = X.shape
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    Y = torch.empty_like(X)
+    RSTD = torch.empty(n_rows, dtype=torch.float32, device=X.device)
+    wrap_triton(_rms_norm_forward_kernel)[(n_rows,)](
+        Y,
+        Y.stride(0),
+        X,
+        X.stride(0),
+        RSTD,
+        RSTD.stride(0),
+        n_cols,
+        eps,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return Y, RSTD
+
+
+@triton_op("axolotl::fused_rms_norm_noscale_bwd", mutates_args=())
+def _fused_rms_norm_noscale_bwd(
+    dY: torch.Tensor, X: torch.Tensor, RSTD: torch.Tensor
+) -> torch.Tensor:
+    n_rows, n_cols = X.shape
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    dX = torch.empty_like(X)
+    wrap_triton(_rms_norm_noscale_backward_kernel)[(n_rows,)](
+        dY,
+        dY.stride(0),
+        dX,
+        dX.stride(0),
+        X,
+        X.stride(0),
+        torch_to_triton_dtype[X.dtype],
+        RSTD,
+        RSTD.stride(0),
+        n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return dX
+
+
+def _fused_rms_norm_noscale_setup_context(ctx, inputs, output):
+    X, _eps = inputs
+    _, RSTD = output
+    ctx.save_for_backward(X, RSTD)
+
+
+def _fused_rms_norm_noscale_backward(ctx, grad_Y, grad_RSTD):
+    X, RSTD = ctx.saved_tensors
+    grad_Y = grad_Y.contiguous()
+    dX = _fused_rms_norm_noscale_bwd(grad_Y, X, RSTD)
+    return dX, None
+
+
+_fused_rms_norm_noscale_fwd.register_autograd(
+    _fused_rms_norm_noscale_backward,
+    setup_context=_fused_rms_norm_noscale_setup_context,
+)
+
+
+def fused_rms_norm_noscale(x, eps=1e-6):
+    """RMSNorm without a learned scale (used for v_norm)."""
+    shape = x.shape
+    x_flat = x.reshape(-1, shape[-1]).contiguous()
+    y_flat, _ = _fused_rms_norm_noscale_fwd(x_flat, eps)
+    return y_flat.view(shape)
diff --git a/src/axolotl/kernels/lora.py b/src/axolotl/kernels/lora.py
new file mode 100644
index 0000000000..350a7ec9c8
--- /dev/null
+++ b/src/axolotl/kernels/lora.py
@@ -0,0 +1,2210 @@
+"""
+Module for definition of Low-Rank Adaptation (LoRA) Triton kernels.
+
+See "LoRA: Low-Rank Adaptation of Large Language Models"
+(https://arxiv.org/abs/2106.09685).
+
+Also supports DoRA (Weight-Decomposed Low-Rank Adaptation):
+See "DoRA: Weight-Decomposed Low-Rank Adaptation" (https://arxiv.org/abs/2402.09353).
+
+Credit to `unsloth` (https://unsloth.ai/) for inspiration for this implementation.
+"""
+
+from typing import Callable
+
+import torch
+from bitsandbytes.functional import QuantState
+from torch import nn
+from torch.distributed.tensor import DTensor
+
+from .geglu import geglu_backward, geglu_forward
+from .quantize import dequantize
+from .swiglu import swiglu_backward, swiglu_forward
+from .utils import torch_amp_custom_bwd, torch_amp_custom_fwd
+
+
+def get_lora_parameters(
+    proj: nn.Module,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor | None,
+    QuantState | torch.Tensor | None,
+    torch.Tensor | None,
+    torch.Tensor | None,
+    float | None,
+    torch.Tensor | None,
+    nn.Module | None,
+    torch.Tensor | None,
+]:
+    """
+    Gets LoRA parameters from a projection module.
+
+    Args:
+        proj: The projection module to extract parameters from.
+
+    Returns:
+        A tuple containing:
+        - W: base weight tensor
+        - b: base layer bias (or None)
+        - quant_state: quantization state (or None)
+        - A: LoRA A weight (or None)
+        - B: LoRA B weight (or None)
+        - s: LoRA scaling factor (or None)
+        - lora_bias: LoRA B bias (or None)
+        - dropout: dropout module (or None)
+        - magnitude: DoRA magnitude vector (or None)
+    """
+    # For DPO or disabled adapters
+    base_layer = proj.base_layer if hasattr(proj, "base_layer") else proj
+    W = base_layer.weight
+    b = base_layer.bias
+
+    if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
+        quant_state = getattr(W, "quant_state", None)
+        if quant_state is None and W.dtype == torch.float8_e4m3fn:
+            quant_state = getattr(base_layer, "weight_scale_inv", None)
+        return W, b, quant_state, None, None, None, None, None, None
+
+    quant_state = getattr(W, "quant_state", None)
+    if quant_state is None and W.dtype == torch.float8_e4m3fn:
+        quant_state = getattr(base_layer, "weight_scale_inv", None)
+
+    active_adapter = (
+        proj.active_adapters[0]
+        if hasattr(proj, "active_adapters")
+        else proj.active_adapter
+    )
+
+    linear_A = proj.lora_A[active_adapter]
+    linear_B = proj.lora_B[active_adapter]
+
+    # This manual unsharding is needed for FSDP2 + LoRA kernels compatibility.
+    # We fuse linear layers + LoRA adapters calculations into a single
+    # torch.autograd.Function, bypassing the registered unshard / reshard behavior.
+    # Note that we don't apply resharding later in this module (it gets messy quickly),
+    # but LoRA parameters are generally small enough that this is not an issue.
+    if isinstance(linear_A.weight, DTensor):
+        linear_A.unshard()
+        linear_B.unshard()
+
+    A = linear_A.weight
+    B = linear_B.weight
+    s = proj.scaling[active_adapter]
+
+    # LoRA bias from lora_B (when bias="lora_only" or bias="all")
+    lora_bias = linear_B.bias  # None if bias=False
+
+    # Dropout module
+    dropout = None
+    if hasattr(proj, "lora_dropout") and active_adapter in proj.lora_dropout:
+        dropout = proj.lora_dropout[active_adapter]
+
+    # DoRA magnitude vector
+    magnitude = None
+    if (
+        hasattr(proj, "lora_magnitude_vector")
+        and proj.lora_magnitude_vector
+        and active_adapter in proj.lora_magnitude_vector
+    ):
+        mag_layer = proj.lora_magnitude_vector[active_adapter]
+        magnitude = mag_layer.weight
+        # FSDP2 DTensor unshard for magnitude vector
+        if isinstance(magnitude, DTensor):
+            magnitude = magnitude.full_tensor()
+
+    return W, b, quant_state, A, B, s, lora_bias, dropout, magnitude
+
+
+def _apply_dropout(
+    dropout: nn.Module | None, X: torch.Tensor, training: bool
+) -> torch.Tensor | None:
+    """Apply dropout to X if dropout module exists and is active.
+
+    Returns X_drop (different tensor) or None if no dropout needed.
+    """
+    if dropout is None or isinstance(dropout, nn.Identity) or not training:
+        return None
+    return dropout(X)
+
+
+_USE_TRITON_DORA: bool | None = None
+
+
+def _should_use_triton_dora() -> bool:
+    """Check if Triton DoRA kernel is available."""
+    global _USE_TRITON_DORA
+    if _USE_TRITON_DORA is None:
+        try:
+            from .dora import triton_dora_scale  # noqa: F401
+
+            _USE_TRITON_DORA = True
+        except (ImportError, RuntimeError):
+            _USE_TRITON_DORA = False
+    return _USE_TRITON_DORA
+
+
+def _compute_dora_scale(
+    W: torch.Tensor,
+    W_quant: QuantState | torch.Tensor | None,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    s: float,
+    magnitude: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Compute DoRA magnitude/norm scaling factor with optional caching.
+
+    Uses Triton kernel when available for better performance.
+    Caches weight_norm on the magnitude tensor. Cache invalidated when
+    LoRA A/B data changes (after optimizer step).
+
+    Returns:
+        mag_norm_scale: [out_features] tensor = magnitude / ||W + s * B @ A||_2
+    """
+    # Check cache on magnitude tensor (avoids expensive norm recomputation)
+    # Use tensor._version which increments on any in-place modification
+    # (data_ptr doesn't change when optimizers update params in-place)
+    cache = getattr(magnitude, "_dora_cache", None)
+    if cache is not None:
+        cached_a_ver, cached_b_ver, cached_norm = cache
+        if cached_a_ver == A._version and cached_b_ver == B._version:
+            return magnitude.to(dtype) / cached_norm
+
+    # Full recomputation - try Triton first
+    if _should_use_triton_dora() and W.is_cuda:
+        from .dora import triton_dora_scale
+
+        result = triton_dora_scale(W, W_quant, A, B, s, magnitude, dtype)
+        weight_norm = (magnitude.to(dtype) / result).detach()
+        magnitude._dora_cache = (A._version, B._version, weight_norm)
+        return result
+
+    # PyTorch fallback
+    W_full = dequantize(W.t(), W_quant).t().to(dtype)  # [out, in]
+    lora_weight = B.to(dtype) @ A.to(dtype)
+    combined = W_full + s * lora_weight
+    weight_norm = torch.linalg.norm(combined, dim=1).to(dtype)
+    weight_norm = weight_norm.detach()
+
+    magnitude._dora_cache = (A._version, B._version, weight_norm)
+
+    return magnitude.to(dtype) / weight_norm
+
+
+def _compute_dora_scale_cached(
+    proj: nn.Module,
+    W: torch.Tensor,
+    W_quant: QuantState | torch.Tensor | None,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    s: float,
+    magnitude: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Compute DoRA scale with caching. Recomputes only when LoRA params change.
+
+    Caches the weight norm on the projection module. The cache is invalidated
+    when LoRA A/B data pointers change (indicating an optimizer step occurred).
+    """
+    cache = getattr(proj, "_dora_norm_cache", None)
+    if cache is not None:
+        cached_a_ver, cached_b_ver, cached_norm = cache
+        if cached_a_ver == A._version and cached_b_ver == B._version:
+            return magnitude.to(dtype) / cached_norm
+
+    # Cache miss - full recomputation
+    W_full = dequantize(W.t(), W_quant).t().to(dtype)
+    lora_weight = B.to(dtype) @ A.to(dtype)
+    combined = W_full + s * lora_weight
+    weight_norm = torch.linalg.norm(combined, dim=1).to(dtype).detach()
+
+    proj._dora_norm_cache = (A._version, B._version, weight_norm)
+
+    return magnitude.to(dtype) / weight_norm
+
+
+def matmul_lora(
+    X: torch.Tensor,
+    W: torch.Tensor,
+    b: torch.Tensor | None,
+    W_quant: QuantState | torch.Tensor | None,
+    A: torch.Tensor | None,
+    B: torch.Tensor | None,
+    s: float | None,
+    out: torch.Tensor | None = None,
+    X_drop: torch.Tensor | None = None,
+    lora_bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """
+    Efficient fused matmul + LoRA computation.
+
+    Args:
+        X: Input tensor [*, in_features]
+        W: Base weight matrix [out_features, in_features]
+        W_quant: Quantization state for W
+        A: LoRA A matrix [rank, in_features]
+        B: LoRA B matrix [out_features, rank]
+        s: LoRA scaling factor
+        out: Optional output tensor for inplace operations
+        X_drop: Optional dropout-applied input for LoRA path (if None, uses X)
+        lora_bias: Optional LoRA B layer bias [out_features]
+
+    Returns:
+        Result of X @ W + s * X_drop @ A @ B + b + s * lora_bias
+    """
+    dtype = X.dtype
+    W = dequantize(W.t(), W_quant)
+
+    reshape = False
+    if X.dim() == 3:
+        batch, seq_len, _ = X.shape
+        X = X.view(-1, X.shape[-1])
+        if X_drop is not None:
+            X_drop = X_drop.view(-1, X_drop.shape[-1])
+        reshape = True
+
+    out = torch.matmul(X, W, out=out)
+    if W_quant is not None:
+        del W
+
+    if A is not None:
+        X_lora = X_drop if X_drop is not None else X
+        A, B = A.t().to(dtype), B.t().to(dtype)  # type: ignore[union-attr]
+        out.addmm_(X_lora @ A, B, alpha=s)
+        if lora_bias is not None:
+            out += s * lora_bias
+
+    if b is not None:
+        out += b
+
+    return out.view(batch, seq_len, -1) if reshape else out
+
+
+class LoRA_MLP(torch.autograd.Function):
+    """Optimized LoRA MLP implementation.
+
+    Supports bias, dropout, and DoRA. Dropout is applied to the input for
+    gate/up projections. The down projection uses hidden states (post-activation)
+    as input, so dropout is not applied there.
+    """
+
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(
+        ctx,
+        X: torch.Tensor,
+        X_drop: torch.Tensor | None,
+        # Gate params
+        gate_weight: torch.Tensor,
+        gate_bias: torch.Tensor | None,
+        gate_quant: QuantState | None,
+        gate_A: torch.Tensor | None,
+        gate_B: torch.Tensor | None,
+        gate_scale: float,
+        gate_lora_bias: torch.Tensor | None,
+        gate_magnitude: torch.Tensor | None,
+        # Up params
+        up_weight: torch.Tensor,
+        up_bias: torch.Tensor | None,
+        up_quant: QuantState | None,
+        up_A: torch.Tensor | None,
+        up_B: torch.Tensor | None,
+        up_scale: float,
+        up_lora_bias: torch.Tensor | None,
+        up_magnitude: torch.Tensor | None,
+        # Down params
+        down_weight: torch.Tensor,
+        down_bias: torch.Tensor | None,
+        down_quant: QuantState | None,
+        down_A: torch.Tensor | None,
+        down_B: torch.Tensor | None,
+        down_scale: float,
+        down_lora_bias: torch.Tensor | None,
+        down_magnitude: torch.Tensor | None,
+        # Activation and flags
+        activation_fn: Callable,
+        activation_fn_backward: Callable,
+        inplace: bool | None = True,
+    ) -> torch.Tensor:
+        has_dropout = X_drop is not None
+        has_dora = gate_magnitude is not None
+        dtype = X.dtype
+        X_lora = X_drop if has_dropout else X
+
+        if has_dora:
+            # Gate with DoRA
+            gate_base = matmul_lora(X, gate_weight, None, gate_quant, None, None, None)
+            gate_lora = _lora_only(
+                X_lora, gate_A, gate_B, gate_scale, gate_lora_bias, dtype
+            )
+            gate_mag_scale = _compute_dora_scale(
+                gate_weight,
+                gate_quant,
+                gate_A,
+                gate_B,
+                gate_scale,
+                gate_magnitude,
+                dtype,
+            )
+            gate = gate_mag_scale.unsqueeze(0) * (gate_base + gate_lora)
+            if gate_bias is not None:
+                gate = gate + gate_bias
+
+            # Up with DoRA
+            up_base = matmul_lora(X, up_weight, None, up_quant, None, None, None)
+            up_lora = _lora_only(X_lora, up_A, up_B, up_scale, up_lora_bias, dtype)
+            up_mag_scale = _compute_dora_scale(
+                up_weight, up_quant, up_A, up_B, up_scale, up_magnitude, dtype
+            )
+            up = up_mag_scale.unsqueeze(0) * (up_base + up_lora)
+            if up_bias is not None:
+                up = up + up_bias
+
+            gate_combined = gate_base + gate_lora
+            up_combined = up_base + up_lora
+        else:
+            gate = matmul_lora(
+                X,
+                gate_weight,
+                gate_bias,
+                gate_quant,
+                gate_A,
+                gate_B,
+                gate_scale,
+                X_drop=X_drop,
+                lora_bias=gate_lora_bias,
+            )
+            up = matmul_lora(
+                X,
+                up_weight,
+                up_bias,
+                up_quant,
+                up_A,
+                up_B,
+                up_scale,
+                X_drop=X_drop,
+                lora_bias=up_lora_bias,
+            )
+
+        # Activation
+        hidden = activation_fn(gate, up)
+
+        # Down projection (no dropout on hidden - it's an intermediate)
+        if has_dora:
+            down_base = matmul_lora(
+                hidden, down_weight, None, down_quant, None, None, None
+            )
+            down_lora = _lora_only(
+                hidden, down_A, down_B, down_scale, down_lora_bias, dtype
+            )
+            down_mag_scale = _compute_dora_scale(
+                down_weight,
+                down_quant,
+                down_A,
+                down_B,
+                down_scale,
+                down_magnitude,
+                dtype,
+            )
+            down_combined = down_base + down_lora
+            output = down_mag_scale.unsqueeze(0) * down_combined
+            if down_bias is not None:
+                output = output + down_bias
+        else:
+            output = matmul_lora(
+                hidden,
+                down_weight,
+                down_bias,
+                down_quant,
+                down_A,
+                down_B,
+                down_scale,
+                lora_bias=down_lora_bias,
+            )
+
+        # Save for backward
+        if has_dora:
+            ctx.save_for_backward(
+                X,
+                X_drop if has_dropout else X,
+                gate,
+                up,
+                gate_A.to(dtype) if gate_A is not None else gate_A,
+                gate_B.to(dtype) if gate_B is not None else gate_B,
+                up_A.to(dtype) if up_A is not None else up_A,
+                up_B.to(dtype) if up_B is not None else up_B,
+                down_A.to(dtype) if down_A is not None else down_A,
+                down_B.to(dtype) if down_B is not None else down_B,
+                gate_magnitude,
+                up_magnitude,
+                down_magnitude,
+                gate_mag_scale,
+                up_mag_scale,
+                down_mag_scale,
+                gate_combined,
+                up_combined,
+                down_combined,
+                gate_lora_bias,
+                up_lora_bias,
+                down_lora_bias,
+            )
+        else:
+            # Pre-convert LoRA matrices to compute dtype for backward
+            dtype = X.dtype
+            ctx.save_for_backward(
+                X,
+                X_drop if has_dropout else X,
+                gate,
+                up,
+                gate_A.to(dtype) if gate_A is not None else gate_A,
+                gate_B.to(dtype) if gate_B is not None else gate_B,
+                up_A.to(dtype) if up_A is not None else up_A,
+                up_B.to(dtype) if up_B is not None else up_B,
+                down_A.to(dtype) if down_A is not None else down_A,
+                down_B.to(dtype) if down_B is not None else down_B,
+                gate_lora_bias,
+                up_lora_bias,
+                down_lora_bias,
+            )
+
+        ctx.scales = (gate_scale, up_scale, down_scale)
+        ctx.quants = (gate_quant, up_quant, down_quant)
+        ctx.weights = (gate_weight, up_weight, down_weight)
+        ctx.activation_fn = activation_fn
+        ctx.activation_fn_backward = activation_fn_backward
+        ctx.has_dropout = has_dropout
+        ctx.has_dora = has_dora
+
+        return output
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(
+        ctx: torch.autograd.function.FunctionCtx,
+        grad_output: torch.Tensor,
+    ):
+        gate_scale, up_scale, down_scale = ctx.scales
+        gate_quant, up_quant, down_quant = ctx.quants
+        gate_weight, up_weight, down_weight = ctx.weights
+        has_dropout = ctx.has_dropout
+        has_dora = ctx.has_dora
+
+        if has_dora:
+            (
+                X,
+                X_lora,
+                gate,
+                up,
+                gate_A,
+                gate_B,
+                up_A,
+                up_B,
+                down_A,
+                down_B,
+                gate_magnitude,
+                up_magnitude,
+                down_magnitude,
+                gate_mag_scale,
+                up_mag_scale,
+                down_mag_scale,
+                gate_combined,
+                up_combined,
+                down_combined,
+                gate_lora_bias,
+                up_lora_bias,
+                down_lora_bias,
+            ) = ctx.saved_tensors
+        else:
+            (
+                X,
+                X_lora,
+                gate,
+                up,
+                gate_A,
+                gate_B,
+                up_A,
+                up_B,
+                down_A,
+                down_B,
+                gate_lora_bias,
+                up_lora_bias,
+                down_lora_bias,
+            ) = ctx.saved_tensors
+            gate_magnitude = up_magnitude = down_magnitude = None
+            gate_mag_scale = up_mag_scale = down_mag_scale = None
+            gate_combined = up_combined = down_combined = None
+
+        # Transpose all LoRA matrices
+        gate_A_t = gate_A.t() if gate_A is not None else None
+        gate_B_t = gate_B.t() if gate_B is not None else None
+        up_A_t = up_A.t() if up_A is not None else None
+        up_B_t = up_B.t() if up_B is not None else None
+        down_A_t = down_A.t() if down_A is not None else None
+        down_B_t = down_B.t() if down_B is not None else None
+
+        # Reshape inputs
+        batch, seq_len, hd = X.shape
+        grad_output = grad_output.view(-1, grad_output.shape[-1])
+        X = X.view(-1, X.shape[-1])
+        X_lora = X_lora.view(-1, X_lora.shape[-1])
+        gate = gate.view(-1, gate.shape[-1])
+        up = up.view(-1, up.shape[-1])
+
+        # DoRA magnitude gradients for down projection
+        d_gate_mag = d_up_mag = d_down_mag = None
+        d_gate_lora_bias = d_up_lora_bias = d_down_lora_bias = None
+
+        if has_dora:
+            down_combined_flat = down_combined.view(-1, down_combined.shape[-1])
+            d_down_mag = (
+                (grad_output * down_combined_flat).sum(dim=0)
+                * down_mag_scale
+                / down_magnitude
+            )
+            grad_output = grad_output * down_mag_scale.unsqueeze(0)
+
+        # Down lora bias gradient
+        if down_lora_bias is not None:
+            d_down_lora_bias = down_scale * grad_output.sum(dim=0)
+
+        # Down projection backward
+        grad_down = matmul_lora(
+            grad_output,
+            down_weight.t(),
+            None,
+            down_quant,
+            down_B_t,
+            down_A_t,
+            down_scale,
+        )
+
+        # Activation backward
+        h, grad_gate, grad_up = ctx.activation_fn_backward(grad_down, gate, up)
+
+        # DoRA magnitude gradients for gate and up
+        if has_dora:
+            gate_combined_flat = gate_combined.view(-1, gate_combined.shape[-1])
+            up_combined_flat = up_combined.view(-1, up_combined.shape[-1])
+            d_gate_mag = (
+                (grad_gate * gate_combined_flat).sum(dim=0)
+                * gate_mag_scale
+                / gate_magnitude
+            )
+            d_up_mag = (
+                (grad_up * up_combined_flat).sum(dim=0) * up_mag_scale / up_magnitude
+            )
+            grad_gate = grad_gate * gate_mag_scale.unsqueeze(0)
+            grad_up = grad_up * up_mag_scale.unsqueeze(0)
+
+        # LoRA bias gradients for gate and up
+        if gate_lora_bias is not None:
+            d_gate_lora_bias = gate_scale * grad_gate.sum(dim=0)
+        if up_lora_bias is not None:
+            d_up_lora_bias = up_scale * grad_up.sum(dim=0)
+
+        # LoRA parameter gradients (already in compute dtype from forward)
+        # Compute grad @ B once per projection, reuse for dA and dX_lora
+        # Note: _t suffix means transposed from saved shape (A_t = A.t(), etc.)
+        d_down_A = d_down_B = d_up_A = d_up_B = d_gate_A = d_gate_B = None
+        grad_B_up = grad_B_gate = None
+
+        if down_A_t is not None and down_B_t is not None:
+            grad_B_down = grad_output @ down_B_t.t()  # reused in matmul_lora above too
+            d_down_A = torch.empty_like(down_A_t)
+            d_down_B = torch.empty_like(down_B_t)
+            d_down_A.addmm_(h.t(), grad_B_down, alpha=down_scale, beta=0)
+            d_down_B.addmm_(down_A_t.t() @ h.t(), grad_output, alpha=down_scale, beta=0)
+
+        if up_A_t is not None and up_B_t is not None:
+            grad_B_up = grad_up @ up_B_t.t()  # [T, rank] — reuse for dX
+            d_up_A = torch.empty_like(up_A_t)
+            d_up_B = torch.empty_like(up_B_t)
+            d_up_A.addmm_(X_lora.t(), grad_B_up, alpha=up_scale, beta=0)
+            d_up_B.addmm_(up_A_t.t() @ X_lora.t(), grad_up, alpha=up_scale, beta=0)
+
+        if gate_A_t is not None and gate_B_t is not None:
+            grad_B_gate = grad_gate @ gate_B_t.t()  # [T, rank] — reuse for dX
+            d_gate_A = torch.empty_like(gate_A_t)
+            d_gate_B = torch.empty_like(gate_B_t)
+            d_gate_A.addmm_(X_lora.t(), grad_B_gate, alpha=gate_scale, beta=0)
+            d_gate_B.addmm_(
+                gate_A_t.t() @ X_lora.t(), grad_gate, alpha=gate_scale, beta=0
+            )
+
+        # Compute input gradients
+        dX = None
+        dX_drop = None
+
+        if ctx.needs_input_grad[0]:
+            # Base path gradients through gate and up
+            up_weight_deq = dequantize(up_weight.t(), up_quant)
+            # writing into saved X is untraceable when X is another Function's view output
+            dX = torch.matmul(grad_up, up_weight_deq.t())
+            del up_weight_deq
+
+            gate_weight_deq = dequantize(gate_weight, gate_quant)
+            dX.addmm_(grad_gate, gate_weight_deq)
+            del gate_weight_deq
+
+            # LoRA path: reuse grad_B_up and grad_B_gate from above
+            if has_dropout:
+                dX_drop = torch.zeros_like(X_lora)
+                if grad_B_up is not None:
+                    dX_drop.addmm_(grad_B_up, up_A_t.t(), alpha=up_scale)  # type: ignore[union-attr]
+                if grad_B_gate is not None:
+                    dX_drop.addmm_(grad_B_gate, gate_A_t.t(), alpha=gate_scale)  # type: ignore[union-attr]
+                dX_drop = dX_drop.view(batch, seq_len, hd)
+            else:
+                if grad_B_up is not None:
+                    dX.addmm_(grad_B_up, up_A_t.t(), alpha=up_scale)  # type: ignore[union-attr]
+                if grad_B_gate is not None:
+                    dX.addmm_(grad_B_gate, gate_A_t.t(), alpha=gate_scale)  # type: ignore[union-attr]
+
+            dX = dX.view(batch, seq_len, hd)
+
+        # Return gradients matching forward input order:
+        # X, X_drop,
+        # gate: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        # up: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        # down: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        # activation_fn, activation_fn_backward, inplace
+        return (
+            dX,
+            dX_drop,
+            # Gate
+            None,
+            None,
+            None,
+            d_gate_A.t() if d_gate_A is not None else None,
+            d_gate_B.t() if d_gate_B is not None else None,
+            None,
+            d_gate_lora_bias,
+            d_gate_mag,
+            # Up
+            None,
+            None,
+            None,
+            d_up_A.t() if d_up_A is not None else None,
+            d_up_B.t() if d_up_B is not None else None,
+            None,
+            d_up_lora_bias,
+            d_up_mag,
+            # Down
+            None,
+            None,
+            None,
+            d_down_A.t() if d_down_A is not None else None,
+            d_down_B.t() if d_down_B is not None else None,
+            None,
+            d_down_lora_bias,
+            d_down_mag,
+            # Activation fns and flags
+            None,
+            None,
+            None,
+        )
+
+
+def apply_lora_mlp_swiglu(self, X: torch.Tensor, inplace: bool = True) -> torch.Tensor:
+    """Applies LoRA to MLP layer with SwiGLU activation.
+
+    Supports bias, dropout, and DoRA.
+    """
+    gateW, gateb, gateW_quant, gateA, gateB, gateS, gateLB, gateDrop, gateMag = (
+        get_lora_parameters(self.gate_proj)
+    )
+    upW, upb, upW_quant, upA, upB, upS, upLB, upDrop, upMag = get_lora_parameters(
+        self.up_proj
+    )
+    downW, downb, downW_quant, downA, downB, downS, downLB, downDrop, downMag = (
+        get_lora_parameters(self.down_proj)
+    )
+
+    # Shared dropout mask for gate and up (same input)
+    X_drop = _apply_dropout(gateDrop, X, self.training)
+
+    out = LoRA_MLP.apply(
+        X,
+        X_drop,
+        # Gate
+        gateW,
+        gateb,
+        gateW_quant,
+        gateA,
+        gateB,
+        gateS,
+        gateLB,
+        gateMag,
+        # Up
+        upW,
+        upb,
+        upW_quant,
+        upA,
+        upB,
+        upS,
+        upLB,
+        upMag,
+        # Down
+        downW,
+        downb,
+        downW_quant,
+        downA,
+        downB,
+        downS,
+        downLB,
+        downMag,
+        # Activation and flags
+        swiglu_forward,
+        swiglu_backward,
+        inplace,
+    )
+
+    return out
+
+
+def apply_lora_mlp_geglu(self, X: torch.Tensor, inplace: bool = True) -> torch.Tensor:
+    """Applies LoRA to MLP layer with GEGLU activation.
+
+    Supports bias, dropout, and DoRA.
+    """
+    gateW, gateb, gateW_quant, gateA, gateB, gateS, gateLB, gateDrop, gateMag = (
+        get_lora_parameters(self.gate_proj)
+    )
+    upW, upb, upW_quant, upA, upB, upS, upLB, upDrop, upMag = get_lora_parameters(
+        self.up_proj
+    )
+    downW, downb, downW_quant, downA, downB, downS, downLB, downDrop, downMag = (
+        get_lora_parameters(self.down_proj)
+    )
+
+    X_drop = _apply_dropout(gateDrop, X, self.training)
+
+    out = LoRA_MLP.apply(
+        X,
+        X_drop,
+        # Gate
+        gateW,
+        gateb,
+        gateW_quant,
+        gateA,
+        gateB,
+        gateS,
+        gateLB,
+        gateMag,
+        # Up
+        upW,
+        upb,
+        upW_quant,
+        upA,
+        upB,
+        upS,
+        upLB,
+        upMag,
+        # Down
+        downW,
+        downb,
+        downW_quant,
+        downA,
+        downB,
+        downS,
+        downLB,
+        downMag,
+        # Activation and flags
+        geglu_forward,
+        geglu_backward,
+        inplace,
+    )
+
+    return out
+
+
+class LoRA_QKV(torch.autograd.Function):
+    """
+    Optimized LoRA QKV implementation with quantization support.
+
+    Supports bias, dropout, and DoRA (Weight-Decomposed Low-Rank Adaptation).
+    Dropout is applied outside this Function so autograd handles its backward.
+    """
+
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        X: torch.Tensor,
+        X_drop: torch.Tensor | None,
+        # Q params
+        q_weight: torch.Tensor,
+        q_bias: torch.Tensor | None,
+        q_quant: QuantState | None,
+        q_A: torch.Tensor | None,
+        q_B: torch.Tensor | None,
+        q_scale: float,
+        q_lora_bias: torch.Tensor | None,
+        q_magnitude: torch.Tensor | None,
+        # K params
+        k_weight: torch.Tensor,
+        k_bias: torch.Tensor | None,
+        k_quant: QuantState | None,
+        k_A: torch.Tensor | None,
+        k_B: torch.Tensor | None,
+        k_scale: float,
+        k_lora_bias: torch.Tensor | None,
+        k_magnitude: torch.Tensor | None,
+        # V params
+        v_weight: torch.Tensor,
+        v_bias: torch.Tensor | None,
+        v_quant: QuantState | None,
+        v_A: torch.Tensor | None,
+        v_B: torch.Tensor | None,
+        v_scale: float,
+        v_lora_bias: torch.Tensor | None,
+        v_magnitude: torch.Tensor | None,
+        # Flags
+        inplace: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        has_dropout = X_drop is not None
+        has_dora = q_magnitude is not None
+
+        if has_dora:
+            dtype = X.dtype
+            X_lora = X_drop if has_dropout else X
+
+            # Compute Q with DoRA
+            Q_base = matmul_lora(X, q_weight, None, q_quant, None, None, None)
+            Q_lora = _lora_only(X_lora, q_A, q_B, q_scale, q_lora_bias, dtype)
+            q_mag_scale = _compute_dora_scale(
+                q_weight, q_quant, q_A, q_B, q_scale, q_magnitude, dtype
+            )
+            Q = q_mag_scale.unsqueeze(0) * (Q_base + Q_lora)
+            if q_bias is not None:
+                Q = Q + q_bias
+
+            # Compute K with DoRA
+            K_base = matmul_lora(X, k_weight, None, k_quant, None, None, None)
+            K_lora = _lora_only(X_lora, k_A, k_B, k_scale, k_lora_bias, dtype)
+            k_mag_scale = _compute_dora_scale(
+                k_weight, k_quant, k_A, k_B, k_scale, k_magnitude, dtype
+            )
+            K = k_mag_scale.unsqueeze(0) * (K_base + K_lora)
+            if k_bias is not None:
+                K = K + k_bias
+
+            # Compute V with DoRA
+            V_base = matmul_lora(X, v_weight, None, v_quant, None, None, None)
+            V_lora = _lora_only(X_lora, v_A, v_B, v_scale, v_lora_bias, dtype)
+            v_mag_scale = _compute_dora_scale(
+                v_weight, v_quant, v_A, v_B, v_scale, v_magnitude, dtype
+            )
+            V = v_mag_scale.unsqueeze(0) * (V_base + V_lora)
+            if v_bias is not None:
+                V = V + v_bias
+
+            # Save for backward: need combined (base+lora) and mag_scale for DoRA grads
+            Q_combined = Q_base + Q_lora
+            K_combined = K_base + K_lora
+            V_combined = V_base + V_lora
+
+            ctx.save_for_backward(
+                X,
+                X_drop if has_dropout else X,
+                q_A.to(dtype) if q_A is not None else q_A,
+                q_B.to(dtype) if q_B is not None else q_B,
+                k_A.to(dtype) if k_A is not None else k_A,
+                k_B.to(dtype) if k_B is not None else k_B,
+                v_A.to(dtype) if v_A is not None else v_A,
+                v_B.to(dtype) if v_B is not None else v_B,
+                q_magnitude,
+                k_magnitude,
+                v_magnitude,
+                q_mag_scale,
+                k_mag_scale,
+                v_mag_scale,
+                Q_combined,
+                K_combined,
+                V_combined,
+                q_lora_bias,
+                k_lora_bias,
+                v_lora_bias,
+            )
+        else:
+            # Standard LoRA (with optional dropout and bias)
+            Q = matmul_lora(
+                X,
+                q_weight,
+                q_bias,
+                q_quant,
+                q_A,
+                q_B,
+                q_scale,
+                X_drop=X_drop,
+                lora_bias=q_lora_bias,
+            )
+            K = matmul_lora(
+                X,
+                k_weight,
+                k_bias,
+                k_quant,
+                k_A,
+                k_B,
+                k_scale,
+                X_drop=X_drop,
+                lora_bias=k_lora_bias,
+            )
+            V = matmul_lora(
+                X,
+                v_weight,
+                v_bias,
+                v_quant,
+                v_A,
+                v_B,
+                v_scale,
+                X_drop=X_drop,
+                lora_bias=v_lora_bias,
+            )
+
+            # Pre-convert LoRA matrices to compute dtype to avoid
+            # redundant fp32→bf16 conversion in backward
+            dtype = X.dtype
+            ctx.save_for_backward(
+                X,
+                X_drop if has_dropout else X,
+                q_A.to(dtype) if q_A is not None else q_A,
+                q_B.to(dtype) if q_B is not None else q_B,
+                k_A.to(dtype) if k_A is not None else k_A,
+                k_B.to(dtype) if k_B is not None else k_B,
+                v_A.to(dtype) if v_A is not None else v_A,
+                v_B.to(dtype) if v_B is not None else v_B,
+                q_lora_bias,
+                k_lora_bias,
+                v_lora_bias,
+            )
+
+        ctx.scales = (q_scale, k_scale, v_scale)
+        ctx.quants = (q_quant, k_quant, v_quant)
+        ctx.weights = (q_weight, k_weight, v_weight)
+        ctx.has_dropout = has_dropout
+        ctx.has_dora = has_dora
+
+        return Q, K, V
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(
+        ctx: torch.autograd.function.FunctionCtx,
+        q_grad: torch.Tensor,
+        k_grad: torch.Tensor,
+        v_grad: torch.Tensor,
+    ):
+        q_weight, k_weight, v_weight = ctx.weights
+        q_quant, k_quant, v_quant = ctx.quants
+        q_scale, k_scale, v_scale = ctx.scales
+        has_dropout = ctx.has_dropout
+        has_dora = ctx.has_dora
+
+        if has_dora:
+            (
+                X,
+                X_lora,
+                A_q,
+                B_q,
+                A_k,
+                B_k,
+                A_v,
+                B_v,
+                q_magnitude,
+                k_magnitude,
+                v_magnitude,
+                q_mag_scale,
+                k_mag_scale,
+                v_mag_scale,
+                Q_combined,
+                K_combined,
+                V_combined,
+                q_lora_bias,
+                k_lora_bias,
+                v_lora_bias,
+            ) = ctx.saved_tensors
+        else:
+            (
+                X,
+                X_lora,
+                A_q,
+                B_q,
+                A_k,
+                B_k,
+                A_v,
+                B_v,
+                q_lora_bias,
+                k_lora_bias,
+                v_lora_bias,
+            ) = ctx.saved_tensors
+            q_magnitude = k_magnitude = v_magnitude = None
+            q_mag_scale = k_mag_scale = v_mag_scale = None
+            Q_combined = K_combined = V_combined = None
+
+        batch, seq_len = X.shape[:2]
+        q_grad = q_grad.view(-1, q_grad.shape[-1])
+        k_grad = k_grad.reshape(-1, k_grad.shape[-1])
+        v_grad = v_grad.view(-1, v_grad.shape[-1])
+        X = X.view(-1, X.shape[-1])
+        X_lora = X_lora.view(-1, X_lora.shape[-1])
+
+        # DoRA: scale gradients through mag_norm_scale
+        d_q_mag = d_k_mag = d_v_mag = None
+        d_q_lora_bias = d_k_lora_bias = d_v_lora_bias = None
+
+        if has_dora:
+            Q_combined = Q_combined.view(-1, Q_combined.shape[-1])
+            K_combined = K_combined.view(-1, K_combined.shape[-1])
+            V_combined = V_combined.view(-1, V_combined.shape[-1])
+
+            # Magnitude gradients: d_mag = sum_t(grad * combined) / weight_norm
+            # Since mag_scale = magnitude / weight_norm, and weight_norm is detached:
+            # d_magnitude = sum_t(grad * combined) * (1 / weight_norm)
+            # But we have mag_scale = magnitude / weight_norm
+            # d_mag_scale_j = grad_j * combined_j (per element)
+            # d_magnitude_j = d_mag_scale_j / weight_norm_j = sum_t(grad * combined) / weight_norm
+            # Simpler: d_magnitude = sum(grad * combined, dim=0) * mag_scale / magnitude
+            # Actually: mag_scale = m/wn, d_output/d_m = combined/wn = combined * mag_scale/m
+            # d_m = sum(grad * combined * mag_scale / m, dim=0) ... no.
+            # Let's be precise: output = mag_scale * combined, mag_scale = m / wn (wn detached)
+            # d_loss/d_m = d_loss/d_output * d_output/d_m = sum_t(grad_t * combined_t / wn)
+            # = sum_t(grad_t * combined_t) * (1/wn) = sum_t(grad_t * combined_t) * mag_scale / m
+            # Or just: d_m = sum(grad * combined, dim=0) / weight_norm
+            # Since we don't have weight_norm saved, use mag_scale/magnitude:
+            # 1/wn = mag_scale/magnitude
+            d_q_mag = (q_grad * Q_combined).sum(dim=0) * q_mag_scale / q_magnitude
+            d_k_mag = (k_grad * K_combined).sum(dim=0) * k_mag_scale / k_magnitude
+            d_v_mag = (v_grad * V_combined).sum(dim=0) * v_mag_scale / v_magnitude
+
+            # Chain rule: grad through combined = grad * mag_scale
+            q_grad = q_grad * q_mag_scale.unsqueeze(0)
+            k_grad = k_grad * k_mag_scale.unsqueeze(0)
+            v_grad = v_grad * v_mag_scale.unsqueeze(0)
+
+        # LoRA bias gradients
+        if q_lora_bias is not None:
+            d_q_lora_bias = q_scale * q_grad.sum(dim=0)
+        if k_lora_bias is not None:
+            d_k_lora_bias = k_scale * k_grad.sum(dim=0)
+        if v_lora_bias is not None:
+            d_v_lora_bias = v_scale * v_grad.sum(dim=0)
+
+        # Pre-transpose X_lora for LoRA gradients
+        X_lora_t = X_lora.t()
+
+        # Initialize LoRA gradients as None
+        d_A_q = d_B_q = d_A_k = d_B_k = d_A_v = d_B_v = None
+
+        # A_q, B_q etc. are already in compute dtype (converted in forward)
+        # Key optimization: compute grad @ B once, reuse for both dA and dX_lora
+        # A has shape [rank, in], B has shape [out, rank]
+        grad_B_q = grad_B_k = grad_B_v = None
+
+        if A_q is not None and B_q is not None:
+            grad_B_q = q_grad @ B_q  # [T, rank] — reused for dA and dX
+            d_A_q = torch.empty_like(A_q.t())
+            d_B_q = torch.empty_like(B_q.t())
+            d_A_q.addmm_(X_lora_t, grad_B_q, alpha=q_scale, beta=0)
+            d_B_q.addmm_(A_q @ X_lora_t, q_grad, alpha=q_scale, beta=0)
+
+        if A_k is not None and B_k is not None:
+            grad_B_k = k_grad @ B_k
+            d_A_k = torch.empty_like(A_k.t())
+            d_B_k = torch.empty_like(B_k.t())
+            d_A_k.addmm_(X_lora_t, grad_B_k, alpha=k_scale, beta=0)
+            d_B_k.addmm_(A_k @ X_lora_t, k_grad, alpha=k_scale, beta=0)
+
+        if A_v is not None and B_v is not None:
+            grad_B_v = v_grad @ B_v
+            d_A_v = torch.empty_like(A_v.t())
+            d_B_v = torch.empty_like(B_v.t())
+            d_A_v.addmm_(X_lora_t, grad_B_v, alpha=v_scale, beta=0)
+            d_B_v.addmm_(A_v @ X_lora_t, v_grad, alpha=v_scale, beta=0)
+
+        # writing into saved X is untraceable when X is another Function's view output
+        q_weight_t = dequantize(q_weight, q_quant)
+        grad_X = torch.mm(q_grad, q_weight_t)
+        del q_weight_t
+
+        k_weight_t = dequantize(k_weight, k_quant)
+        grad_X.addmm_(k_grad, k_weight_t)
+        del k_weight_t
+
+        v_weight_t = dequantize(v_weight, v_quant)
+        grad_X.addmm_(v_grad, v_weight_t)
+        del v_weight_t
+
+        # LoRA path input gradient: s * grad @ B @ A (reuses grad_B_* from above)
+        if has_dropout:
+            grad_X_drop = torch.zeros_like(X_lora)
+            if grad_B_q is not None:
+                grad_X_drop.addmm_(grad_B_q, A_q, alpha=q_scale)
+            if grad_B_k is not None:
+                grad_X_drop.addmm_(grad_B_k, A_k, alpha=k_scale)
+            if grad_B_v is not None:
+                grad_X_drop.addmm_(grad_B_v, A_v, alpha=v_scale)
+        else:
+            grad_X_drop = None
+            if grad_B_q is not None:
+                grad_X.addmm_(grad_B_q, A_q, alpha=q_scale)
+            if grad_B_k is not None:
+                grad_X.addmm_(grad_B_k, A_k, alpha=k_scale)
+            if grad_B_v is not None:
+                grad_X.addmm_(grad_B_v, A_v, alpha=v_scale)
+
+        # Transpose LoRA gradients
+        if d_A_q is not None:
+            d_A_q = d_A_q.t()
+            d_B_q = d_B_q.t()  # type: ignore[union-attr]
+        if d_A_k is not None:
+            d_A_k = d_A_k.t()
+            d_B_k = d_B_k.t()  # type: ignore[union-attr]
+        if d_A_v is not None:
+            d_A_v = d_A_v.t()
+            d_B_v = d_B_v.t()  # type: ignore[union-attr]
+
+        grad_X = grad_X.view(batch, seq_len, -1)
+        if grad_X_drop is not None:
+            grad_X_drop = grad_X_drop.view(batch, seq_len, -1)
+
+        # Return gradients for all forward inputs:
+        # X, X_drop,
+        # q: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        # k: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        # v: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        # inplace
+        return (
+            grad_X,
+            grad_X_drop,
+            # Q
+            None,
+            None,
+            None,
+            d_A_q,
+            d_B_q,
+            None,
+            d_q_lora_bias,
+            d_q_mag,
+            # K
+            None,
+            None,
+            None,
+            d_A_k,
+            d_B_k,
+            None,
+            d_k_lora_bias,
+            d_k_mag,
+            # V
+            None,
+            None,
+            None,
+            d_A_v,
+            d_B_v,
+            None,
+            d_v_lora_bias,
+            d_v_mag,
+            # inplace
+            None,
+        )
+
+
+def _lora_only(
+    X: torch.Tensor,
+    A: torch.Tensor | None,
+    B: torch.Tensor | None,
+    s: float | None,
+    lora_bias: torch.Tensor | None,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Compute only the LoRA contribution: s * X @ A^T @ B^T + s * lora_bias."""
+    if A is None:
+        return torch.zeros(
+            X.shape[:-1] + (B.shape[0] if B is not None else 1,),
+            device=X.device,
+            dtype=dtype,
+        )
+    reshape = False
+    if X.dim() == 3:
+        batch, seq_len, _ = X.shape
+        X = X.view(-1, X.shape[-1])
+        reshape = True
+    At, Bt = A.t().to(dtype), B.t().to(dtype)  # type: ignore[union-attr]
+    out = s * X @ At @ Bt
+    if lora_bias is not None:
+        out = out + s * lora_bias
+    return out.view(batch, seq_len, -1) if reshape else out
+
+
+def apply_lora_qkv(
+    self, X: torch.Tensor, inplace: bool = True
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Applies LoRA to compute Query, Key, Value projections.
+
+    Supports bias, dropout, and DoRA. Dropout is applied outside the autograd
+    Function so PyTorch handles its backward automatically. A single shared
+    dropout mask is used across Q, K, V projections for memory efficiency.
+    """
+    QW, Qb, QW_quant, QA, QB, QS, Qlb, Qdrop, Qmag = get_lora_parameters(self.q_proj)
+    KW, Kb, KW_quant, KA, KB, KS, Klb, Kdrop, Kmag = get_lora_parameters(self.k_proj)
+    VW, Vb, VW_quant, VA, VB, VS, Vlb, Vdrop, Vmag = get_lora_parameters(self.v_proj)
+
+    # Apply dropout outside autograd.Function (shared mask for Q, K, V)
+    X_drop = _apply_dropout(Qdrop, X, self.training)
+
+    Q, K, V = LoRA_QKV.apply(
+        X,
+        X_drop,
+        # Q
+        QW,
+        Qb,
+        QW_quant,
+        QA,
+        QB,
+        QS,
+        Qlb,
+        Qmag,
+        # K
+        KW,
+        Kb,
+        KW_quant,
+        KA,
+        KB,
+        KS,
+        Klb,
+        Kmag,
+        # V
+        VW,
+        Vb,
+        VW_quant,
+        VA,
+        VB,
+        VS,
+        Vlb,
+        Vmag,
+        # Flags
+        inplace,
+    )
+
+    return Q, K, V
+
+
+class LoRA_QK(torch.autograd.Function):
+    """Optimized LoRA QK implementation for models where v_proj is None.
+
+    Used by models like Gemma4 with attention_k_eq_v=True, where key states are
+    reused as value states.  Only Q and K projections are fused; the caller
+    returns K a second time as V so that autograd accumulates key+value gradients
+    into a single dK.
+
+    Supports bias, dropout, and DoRA (Weight-Decomposed Low-Rank Adaptation).
+    """
+
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        X: torch.Tensor,
+        X_drop: torch.Tensor | None,
+        # Q params
+        q_weight: torch.Tensor,
+        q_bias: torch.Tensor | None,
+        q_quant: QuantState | None,
+        q_A: torch.Tensor | None,
+        q_B: torch.Tensor | None,
+        q_scale: float,
+        q_lora_bias: torch.Tensor | None,
+        q_magnitude: torch.Tensor | None,
+        # K params
+        k_weight: torch.Tensor,
+        k_bias: torch.Tensor | None,
+        k_quant: QuantState | None,
+        k_A: torch.Tensor | None,
+        k_B: torch.Tensor | None,
+        k_scale: float,
+        k_lora_bias: torch.Tensor | None,
+        k_magnitude: torch.Tensor | None,
+        # Flags
+        inplace: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        has_dropout = X_drop is not None
+        has_dora = q_magnitude is not None
+
+        if has_dora:
+            dtype = X.dtype
+            X_lora = X_drop if has_dropout else X
+
+            # Compute Q with DoRA
+            Q_base = matmul_lora(X, q_weight, None, q_quant, None, None, None)
+            Q_lora = _lora_only(X_lora, q_A, q_B, q_scale, q_lora_bias, dtype)
+            q_mag_scale = _compute_dora_scale(
+                q_weight, q_quant, q_A, q_B, q_scale, q_magnitude, dtype
+            )
+            Q = q_mag_scale.unsqueeze(0) * (Q_base + Q_lora)
+            if q_bias is not None:
+                Q = Q + q_bias
+
+            # Compute K with DoRA
+            K_base = matmul_lora(X, k_weight, None, k_quant, None, None, None)
+            K_lora = _lora_only(X_lora, k_A, k_B, k_scale, k_lora_bias, dtype)
+            k_mag_scale = _compute_dora_scale(
+                k_weight, k_quant, k_A, k_B, k_scale, k_magnitude, dtype
+            )
+            K = k_mag_scale.unsqueeze(0) * (K_base + K_lora)
+            if k_bias is not None:
+                K = K + k_bias
+
+            Q_combined = Q_base + Q_lora
+            K_combined = K_base + K_lora
+
+            ctx.save_for_backward(
+                X,
+                X_drop if has_dropout else X,
+                q_A.to(dtype) if q_A is not None else q_A,
+                q_B.to(dtype) if q_B is not None else q_B,
+                k_A.to(dtype) if k_A is not None else k_A,
+                k_B.to(dtype) if k_B is not None else k_B,
+                q_magnitude,
+                k_magnitude,
+                q_mag_scale,
+                k_mag_scale,
+                Q_combined,
+                K_combined,
+                q_lora_bias,
+                k_lora_bias,
+            )
+        else:
+            # Standard LoRA (with optional dropout and bias)
+            Q = matmul_lora(
+                X,
+                q_weight,
+                q_bias,
+                q_quant,
+                q_A,
+                q_B,
+                q_scale,
+                X_drop=X_drop,
+                lora_bias=q_lora_bias,
+            )
+            K = matmul_lora(
+                X,
+                k_weight,
+                k_bias,
+                k_quant,
+                k_A,
+                k_B,
+                k_scale,
+                X_drop=X_drop,
+                lora_bias=k_lora_bias,
+            )
+
+            dtype = X.dtype
+            ctx.save_for_backward(
+                X,
+                X_drop if has_dropout else X,
+                q_A.to(dtype) if q_A is not None else q_A,
+                q_B.to(dtype) if q_B is not None else q_B,
+                k_A.to(dtype) if k_A is not None else k_A,
+                k_B.to(dtype) if k_B is not None else k_B,
+                q_lora_bias,
+                k_lora_bias,
+            )
+
+        ctx.scales = (q_scale, k_scale)
+        ctx.quants = (q_quant, k_quant)
+        ctx.weights = (q_weight, k_weight)
+        ctx.has_dropout = has_dropout
+        ctx.has_dora = has_dora
+
+        return Q, K
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(
+        ctx: torch.autograd.function.FunctionCtx,
+        q_grad: torch.Tensor,
+        k_grad: torch.Tensor,
+    ):
+        q_weight, k_weight = ctx.weights
+        q_quant, k_quant = ctx.quants
+        q_scale, k_scale = ctx.scales
+        has_dropout = ctx.has_dropout
+        has_dora = ctx.has_dora
+
+        if has_dora:
+            (
+                X,
+                X_lora,
+                A_q,
+                B_q,
+                A_k,
+                B_k,
+                q_magnitude,
+                k_magnitude,
+                q_mag_scale,
+                k_mag_scale,
+                Q_combined,
+                K_combined,
+                q_lora_bias,
+                k_lora_bias,
+            ) = ctx.saved_tensors
+        else:
+            (
+                X,
+                X_lora,
+                A_q,
+                B_q,
+                A_k,
+                B_k,
+                q_lora_bias,
+                k_lora_bias,
+            ) = ctx.saved_tensors
+            q_magnitude = k_magnitude = None
+            q_mag_scale = k_mag_scale = None
+            Q_combined = K_combined = None
+
+        batch, seq_len = X.shape[:2]
+        q_grad = q_grad.view(-1, q_grad.shape[-1])
+        k_grad = k_grad.reshape(-1, k_grad.shape[-1])
+        X = X.view(-1, X.shape[-1])
+        X_lora = X_lora.view(-1, X_lora.shape[-1])
+
+        d_q_mag = d_k_mag = None
+        d_q_lora_bias = d_k_lora_bias = None
+
+        if has_dora:
+            Q_combined = Q_combined.view(-1, Q_combined.shape[-1])
+            K_combined = K_combined.view(-1, K_combined.shape[-1])
+
+            d_q_mag = (q_grad * Q_combined).sum(dim=0) * q_mag_scale / q_magnitude
+            d_k_mag = (k_grad * K_combined).sum(dim=0) * k_mag_scale / k_magnitude
+
+            q_grad = q_grad * q_mag_scale.unsqueeze(0)
+            k_grad = k_grad * k_mag_scale.unsqueeze(0)
+
+        # LoRA bias gradients
+        if q_lora_bias is not None:
+            d_q_lora_bias = q_scale * q_grad.sum(dim=0)
+        if k_lora_bias is not None:
+            d_k_lora_bias = k_scale * k_grad.sum(dim=0)
+
+        X_lora_t = X_lora.t()
+
+        d_A_q = d_B_q = d_A_k = d_B_k = None
+        grad_B_q = grad_B_k = None
+
+        if A_q is not None and B_q is not None:
+            grad_B_q = q_grad @ B_q
+            d_A_q = torch.empty_like(A_q.t())
+            d_B_q = torch.empty_like(B_q.t())
+            d_A_q.addmm_(X_lora_t, grad_B_q, alpha=q_scale, beta=0)
+            d_B_q.addmm_(A_q @ X_lora_t, q_grad, alpha=q_scale, beta=0)
+
+        if A_k is not None and B_k is not None:
+            grad_B_k = k_grad @ B_k
+            d_A_k = torch.empty_like(A_k.t())
+            d_B_k = torch.empty_like(B_k.t())
+            d_A_k.addmm_(X_lora_t, grad_B_k, alpha=k_scale, beta=0)
+            d_B_k.addmm_(A_k @ X_lora_t, k_grad, alpha=k_scale, beta=0)
+
+        # writing into saved X is untraceable when X is another Function's view output
+        q_weight_t = dequantize(q_weight, q_quant)
+        grad_X = torch.mm(q_grad, q_weight_t)
+        del q_weight_t
+
+        k_weight_t = dequantize(k_weight, k_quant)
+        grad_X.addmm_(k_grad, k_weight_t)
+        del k_weight_t
+
+        # LoRA path input gradient
+        if has_dropout:
+            grad_X_drop = torch.zeros_like(X_lora)
+            if grad_B_q is not None:
+                grad_X_drop.addmm_(grad_B_q, A_q, alpha=q_scale)
+            if grad_B_k is not None:
+                grad_X_drop.addmm_(grad_B_k, A_k, alpha=k_scale)
+        else:
+            grad_X_drop = None
+            if grad_B_q is not None:
+                grad_X.addmm_(grad_B_q, A_q, alpha=q_scale)
+            if grad_B_k is not None:
+                grad_X.addmm_(grad_B_k, A_k, alpha=k_scale)
+
+        if d_A_q is not None:
+            d_A_q = d_A_q.t()
+            d_B_q = d_B_q.t()  # type: ignore[union-attr]
+        if d_A_k is not None:
+            d_A_k = d_A_k.t()
+            d_B_k = d_B_k.t()  # type: ignore[union-attr]
+
+        grad_X = grad_X.view(batch, seq_len, -1)
+        if grad_X_drop is not None:
+            grad_X_drop = grad_X_drop.view(batch, seq_len, -1)
+
+        # Return gradients for all forward inputs:
+        # X, X_drop,
+        # q: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        # k: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        # inplace
+        return (
+            grad_X,
+            grad_X_drop,
+            # Q
+            None,
+            None,
+            None,
+            d_A_q,
+            d_B_q,
+            None,
+            d_q_lora_bias,
+            d_q_mag,
+            # K
+            None,
+            None,
+            None,
+            d_A_k,
+            d_B_k,
+            None,
+            d_k_lora_bias,
+            d_k_mag,
+            # inplace
+            None,
+        )
+
+
+def apply_lora_qk(
+    self, X: torch.Tensor, inplace: bool = True
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Applies LoRA to compute Query and Key projections for models where v_proj is None.
+
+    When v_proj is None (e.g. Gemma4 attention_k_eq_v), key states are reused as
+    value states.  Returns (Q, K, K) — the caller's patched forward will use K as V.
+    Because K is returned twice, autograd accumulates gradients from both the key and
+    value paths into dK before calling LoRA_QK.backward.
+
+    Supports bias, dropout, and DoRA.
+    """
+    QW, Qb, QW_quant, QA, QB, QS, Qlb, Qdrop, Qmag = get_lora_parameters(self.q_proj)
+    KW, Kb, KW_quant, KA, KB, KS, Klb, Kdrop, Kmag = get_lora_parameters(self.k_proj)
+
+    # Apply dropout outside autograd.Function (shared mask for Q, K)
+    X_drop = _apply_dropout(Qdrop, X, self.training)
+
+    Q, K = LoRA_QK.apply(
+        X,
+        X_drop,
+        # Q
+        QW,
+        Qb,
+        QW_quant,
+        QA,
+        QB,
+        QS,
+        Qlb,
+        Qmag,
+        # K
+        KW,
+        Kb,
+        KW_quant,
+        KA,
+        KB,
+        KS,
+        Klb,
+        Kmag,
+        # Flags
+        inplace,
+    )
+
+    return Q, K, K
+
+
+class LoRA_O(torch.autograd.Function):
+    """Optimized LoRA implementation for output projection.
+
+    Supports bias, dropout, and DoRA.
+    """
+
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        X: torch.Tensor,
+        X_drop: torch.Tensor | None,
+        W: torch.Tensor,
+        b: torch.Tensor | None,
+        W_quant: QuantState | None,
+        A: torch.Tensor | None,
+        B: torch.Tensor | None,
+        s: float,
+        lora_bias: torch.Tensor | None,
+        magnitude: torch.Tensor | None,
+    ) -> torch.Tensor:
+        has_dropout = X_drop is not None
+        has_dora = magnitude is not None
+        dtype = X.dtype
+
+        if has_dora:
+            X_lora = X_drop if has_dropout else X
+            base_out = matmul_lora(X, W, None, W_quant, None, None, None)
+            lora_out = _lora_only(X_lora, A, B, s, lora_bias, dtype)
+            mag_scale = _compute_dora_scale(W, W_quant, A, B, s, magnitude, dtype)
+            combined = base_out + lora_out
+            XW = mag_scale.unsqueeze(0) * combined
+            if b is not None:
+                XW = XW + b
+
+            ctx.save_for_backward(
+                A.to(dtype) if A is not None else A,
+                B.to(dtype) if B is not None else B,
+                X,
+                X_drop if has_dropout else X,
+                magnitude,
+                mag_scale,
+                combined,
+                lora_bias,
+            )
+        else:
+            XW = matmul_lora(
+                X,
+                W,
+                b,
+                W_quant,
+                A,
+                B,
+                s,
+                X_drop=X_drop,
+                lora_bias=lora_bias,
+            )
+            # Pre-convert LoRA matrices to compute dtype for backward
+            dtype = X.dtype
+            ctx.save_for_backward(
+                A.to(dtype) if A is not None else A,
+                B.to(dtype) if B is not None else B,
+                X,
+                X_drop if has_dropout else X,
+                lora_bias,
+            )
+
+        ctx.custom_saved_tensors = (W, W_quant, s)
+        ctx.has_dropout = has_dropout
+        ctx.has_dora = has_dora
+
+        return XW
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(
+        ctx: torch.autograd.function.FunctionCtx,
+        dY: torch.Tensor,
+    ):
+        W, W_quant, s = ctx.custom_saved_tensors
+        has_dropout = ctx.has_dropout
+        has_dora = ctx.has_dora
+
+        if has_dora:
+            A, B, X, X_lora, magnitude, mag_scale, combined, lora_bias = (
+                ctx.saved_tensors
+            )
+        else:
+            A, B, X, X_lora, lora_bias = ctx.saved_tensors
+            magnitude = mag_scale = combined = None
+
+        batch, seq_len, hd = X.shape
+        dY = dY.reshape(-1, dY.shape[-1])
+        X = X.reshape(-1, X.shape[-1])
+        X_lora = X_lora.reshape(-1, X_lora.shape[-1])
+
+        d_mag = d_lora_bias = None
+
+        if has_dora:
+            combined = combined.view(-1, combined.shape[-1])
+            d_mag = (dY * combined).sum(dim=0) * mag_scale / magnitude
+            dY = dY * mag_scale.unsqueeze(0)
+
+        # LoRA bias gradient
+        if lora_bias is not None:
+            d_lora_bias = s * dY.sum(dim=0)
+
+        # LoRA parameter gradients (A, B already in compute dtype from forward)
+        # Compute dY @ B once, reuse for both dA and dX_lora
+        d_A = d_B = None
+        grad_B = None
+        if A is not None:
+            grad_B = dY @ B  # [T, rank] — reused below
+            X_lora_t = X_lora.t()
+            d_A = torch.empty_like(A.t())
+            d_B = torch.empty_like(B.t())
+            d_A.addmm_(X_lora_t, grad_B, alpha=s, beta=0)
+            d_B.addmm_(A @ X_lora_t, dY, alpha=s, beta=0)
+
+        # Base path input gradient
+        W_deq = dequantize(W.t(), W_quant)
+        dX = dY @ W_deq.t()
+        del W_deq
+
+        if has_dropout:
+            dX_drop = None
+            if grad_B is not None:
+                dX_drop = (grad_B @ A * s).view(batch, seq_len, hd)
+        else:
+            dX_drop = None
+            if grad_B is not None:
+                dX.addmm_(grad_B, A, alpha=s)
+
+        # X, X_drop, W, b, W_quant, A, B, s, lora_bias, magnitude
+        return (
+            dX.view(batch, seq_len, hd),
+            dX_drop,
+            None,
+            None,
+            None,
+            d_A.t() if d_A is not None else None,
+            d_B.t() if d_B is not None else None,
+            None,
+            d_lora_bias,
+            d_mag,
+        )
+
+
+def apply_lora_o(self, X: torch.Tensor) -> torch.Tensor:
+    """
+    Applies LoRA to output projection layer.
+
+    Supports bias, dropout, and DoRA.
+    """
+    OW, Ob, OW_quant, OA, OB, OS, Olb, Odrop, Omag = get_lora_parameters(self.o_proj)
+    X_drop = _apply_dropout(Odrop, X, self.training)
+    output = LoRA_O.apply(X, X_drop, OW, Ob, OW_quant, OA, OB, OS, Olb, Omag)
+
+    return output
+
+
+def apply_lora_linear(proj: nn.Module, X: torch.Tensor) -> torch.Tensor:
+    """Routed through ``LoRA_O`` to avoid peft's bf16->fp32->bf16 dtype round-trip."""
+    W, b, W_quant, A, B, s, lora_bias, dropout, magnitude = get_lora_parameters(proj)
+    X_drop = _apply_dropout(dropout, X, proj.training)
+    return LoRA_O.apply(X, X_drop, W, b, W_quant, A, B, s, lora_bias, magnitude)
+
+
+# Tensors packed per projection into LoRA_FusedProj's variadic args / saved tensors.
+_FUSED_PROJ_NUM_TENSORS = 6
+
+
+def _chunk(seq, size):
+    return [tuple(seq[i : i + size]) for i in range(0, len(seq), size)]
+
+
+class LoRA_FusedProj(torch.autograd.Function):
+    """Fused LoRA over N projections sharing one input ``X`` (GatedDeltaNet in-projs).
+
+    One dropout mask is shared across the group (peft draws one per projection), so
+    outputs are not bit-identical to peft when ``lora_dropout > 0``; exact at dropout 0.
+    """
+
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        scales: tuple,
+        quants: tuple,
+        X: torch.Tensor,
+        X_drop: torch.Tensor | None,
+        *tensors: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, ...]:
+        has_dropout = X_drop is not None
+        dtype = X.dtype
+        X_lora = X_drop if has_dropout else X
+
+        outs = []
+        weights = []
+        saved = [X, X_drop if has_dropout else X]
+        for i, (W, b, A, B, lora_bias, magnitude) in enumerate(
+            _chunk(tensors, _FUSED_PROJ_NUM_TENSORS)
+        ):
+            quant, s = quants[i], scales[i]
+            if magnitude is not None:
+                base = matmul_lora(X, W, None, quant, None, None, None)
+                lora = _lora_only(X_lora, A, B, s, lora_bias, dtype)
+                mag_scale = _compute_dora_scale(W, quant, A, B, s, magnitude, dtype)
+                combined = base + lora
+                out = mag_scale.unsqueeze(0) * combined
+                if b is not None:
+                    out = out + b
+                saved.extend(
+                    [
+                        A.to(dtype) if A is not None else None,
+                        B.to(dtype) if B is not None else None,
+                        lora_bias,
+                        magnitude,
+                        mag_scale,
+                        combined,
+                    ]
+                )
+            else:
+                out = matmul_lora(
+                    X, W, b, quant, A, B, s, X_drop=X_drop, lora_bias=lora_bias
+                )
+                saved.extend(
+                    [
+                        A.to(dtype) if A is not None else None,
+                        B.to(dtype) if B is not None else None,
+                        lora_bias,
+                        None,
+                        None,
+                        None,
+                    ]
+                )
+            outs.append(out)
+            weights.append(W)
+
+        ctx.save_for_backward(*saved)
+        ctx.scales = scales
+        ctx.quants = quants
+        ctx.weights = tuple(weights)
+        ctx.has_dropout = has_dropout
+        return tuple(outs)
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(ctx: torch.autograd.function.FunctionCtx, *grads: torch.Tensor):
+        saved = ctx.saved_tensors
+        X, X_lora = saved[0], saved[1]
+        groups = _chunk(saved[2:], _FUSED_PROJ_NUM_TENSORS)
+        scales, quants, weights = ctx.scales, ctx.quants, ctx.weights
+        has_dropout = ctx.has_dropout
+
+        batch, seq_len = X.shape[:2]
+        X_lora = X_lora.view(-1, X_lora.shape[-1])
+        # X_lora aliases X when there is no dropout; keep grad_X out-of-place so
+        # this transpose stays valid across every projection's LoRA grads.
+        X_lora_t = X_lora.t()
+
+        grad_X = None
+        grad_X_drop = torch.zeros_like(X_lora) if has_dropout else None
+        tensor_grads: list[torch.Tensor | None] = []
+
+        for i, (A, B, lora_bias, magnitude, mag_scale, combined) in enumerate(groups):
+            dY = grads[i].reshape(-1, grads[i].shape[-1])
+            s, quant, W = scales[i], quants[i], weights[i]
+
+            d_mag = d_lora_bias = None
+            if magnitude is not None:
+                combined = combined.view(-1, combined.shape[-1])
+                d_mag = (dY * combined).sum(dim=0) * mag_scale / magnitude
+                dY = dY * mag_scale.unsqueeze(0)
+
+            if lora_bias is not None:
+                d_lora_bias = s * dY.sum(dim=0)
+
+            grad_B = d_A = d_B = None
+            if A is not None and B is not None:
+                grad_B = dY @ B
+                d_A = torch.empty_like(A.t())
+                d_B = torch.empty_like(B.t())
+                d_A.addmm_(X_lora_t, grad_B, alpha=s, beta=0)
+                d_B.addmm_(A @ X_lora_t, dY, alpha=s, beta=0)
+
+            # .to(dY.dtype): an NF4 base with an fp32 quant_state dequants to fp32,
+            # which the shared in-place grad_X accumulator cannot mix. No-op in the
+            # supported paths (bf16 quant_state / fp8 / Float8Tensor all give bf16).
+            W_deq = dequantize(W, quant).to(dY.dtype)  # [out, in]
+            if grad_X is None:
+                grad_X = torch.mm(dY, W_deq)
+            else:
+                grad_X.addmm_(dY, W_deq)
+            del W_deq
+
+            if grad_B is not None:
+                if has_dropout:
+                    assert grad_X_drop is not None
+                    grad_X_drop.addmm_(grad_B, A, alpha=s)
+                else:
+                    grad_X.addmm_(grad_B, A, alpha=s)
+
+            tensor_grads.extend(
+                [
+                    None,  # W (frozen base)
+                    None,  # b (frozen base)
+                    d_A.t() if d_A is not None else None,
+                    d_B.t() if d_B is not None else None,
+                    d_lora_bias,
+                    d_mag,
+                ]
+            )
+
+        assert grad_X is not None  # every projection contributes a base dY @ W
+        grad_X = grad_X.view(batch, seq_len, -1)
+        if grad_X_drop is not None:
+            grad_X_drop = grad_X_drop.view(batch, seq_len, -1)
+
+        # scales, quants, X, X_drop, *tensors
+        return (None, None, grad_X, grad_X_drop, *tensor_grads)
+
+
+def apply_lora_gdn_in_proj(
+    self, X: torch.Tensor, proj_names: tuple[str, ...]
+) -> dict[str, torch.Tensor]:
+    """Fuse GDN in-projections (shared ``X``) into one call; base-only projs fold in too."""
+    scales, quants, flat = [], [], []
+    dropout_mod = None
+    for name in proj_names:
+        W, b, quant, A, B, s, lora_bias, dropout, magnitude = get_lora_parameters(
+            getattr(self, name)
+        )
+        scales.append(s if s is not None else 1.0)
+        quants.append(quant)
+        flat.extend([W, b, A, B, lora_bias, magnitude])
+        if dropout_mod is None:
+            dropout_mod = dropout
+
+    X_drop = _apply_dropout(dropout_mod, X, self.training)
+    outs = LoRA_FusedProj.apply(tuple(scales), tuple(quants), X, X_drop, *flat)
+    return dict(zip(proj_names, outs, strict=True))
+
+
+# ============================================================
+# Embedding LoRA kernel
+# ============================================================
+
+
+def get_embedding_lora_parameters(
+    embed: nn.Module,
+) -> tuple[
+    torch.Tensor,  # W (base embedding weight)
+    torch.Tensor | None,  # A (lora_embedding_A)
+    torch.Tensor | None,  # B (lora_embedding_B)
+    float | None,  # scaling
+    nn.Module | None,  # dropout
+    torch.Tensor | None,  # magnitude (DoRA)
+    nn.Module,  # base_layer
+]:
+    """Extract LoRA parameters from a PEFT Embedding module."""
+    base_layer = embed.base_layer if hasattr(embed, "base_layer") else embed
+    W = base_layer.weight
+
+    if not hasattr(embed, "disable_adapters") or embed.disable_adapters or embed.merged:
+        return W, None, None, None, None, None, base_layer
+
+    active_adapter = (
+        embed.active_adapters[0]
+        if hasattr(embed, "active_adapters")
+        else embed.active_adapter
+    )
+
+    A = embed.lora_embedding_A[active_adapter]  # nn.Parameter [rank, vocab]
+    B = embed.lora_embedding_B[active_adapter]  # nn.Parameter [hidden_dim, rank]
+    s = embed.scaling[active_adapter]
+
+    # FSDP2 DTensor unshard (mirrors linear path logic)
+    if isinstance(A, DTensor):
+        A = A.full_tensor()
+    if isinstance(B, DTensor):
+        B = B.full_tensor()
+
+    dropout = None
+    if hasattr(embed, "lora_dropout") and active_adapter in embed.lora_dropout:
+        dropout = embed.lora_dropout[active_adapter]
+
+    magnitude = None
+    if (
+        hasattr(embed, "lora_magnitude_vector")
+        and embed.lora_magnitude_vector
+        and active_adapter in embed.lora_magnitude_vector
+    ):
+        mag_layer = embed.lora_magnitude_vector[active_adapter]
+        magnitude = mag_layer.weight
+        if isinstance(magnitude, DTensor):
+            magnitude = magnitude.full_tensor()
+
+    return W, A, B, s, dropout, magnitude, base_layer
+
+
+class LoRA_Embedding(torch.autograd.Function):
+    """Fused LoRA embedding: F.embedding(x, W) + s * F.embedding(x, A^T) @ B^T.
+
+    Supports dropout and DoRA.
+    """
+
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        W: torch.Tensor,
+        A: torch.Tensor | None,
+        B: torch.Tensor | None,
+        s: float | None,
+        magnitude: torch.Tensor | None,
+        padding_idx: int | None,
+        # base_layer fields for F.embedding
+        max_norm: float | None,
+        norm_type: float,
+        scale_grad_by_freq: bool,
+        sparse: bool,
+    ) -> torch.Tensor:
+        import torch.nn.functional as F
+
+        has_dora = magnitude is not None
+        dtype = W.dtype
+
+        # Base embedding lookup
+        result = F.embedding(
+            x,
+            W,
+            padding_idx=padding_idx,
+            max_norm=max_norm,
+            norm_type=norm_type,
+            scale_grad_by_freq=scale_grad_by_freq,
+            sparse=sparse,
+        )
+
+        if A is not None:
+            # LoRA: F.embedding(x, A^T) @ B^T * s
+            A_T = A.t()  # type: ignore[union-attr]  # [vocab, rank]
+            B_T = B.t()  # type: ignore[union-attr]  # [rank, hidden_dim]
+            after_A = F.embedding(
+                x,
+                A_T,
+                padding_idx=padding_idx,
+                max_norm=max_norm,
+                norm_type=norm_type,
+                scale_grad_by_freq=scale_grad_by_freq,
+                sparse=sparse,
+            )  # [batch, seq, rank]
+
+            lora_result = after_A @ B_T  # [batch, seq, hidden_dim]
+
+            if has_dora:
+                mag_scale = _compute_dora_scale(W.t(), None, A, B, s, magnitude, dtype)  # type: ignore[arg-type]
+                # DoRA: mag_scale * (base + s * lora)
+                # base embedding has no bias
+                pre_scaled = result + s * lora_result  # unscaled combined
+                result = mag_scale.unsqueeze(0) * pre_scaled
+                ctx.save_for_backward(
+                    x,
+                    A.to(dtype),
+                    B.to(dtype),  # type: ignore[union-attr]
+                    after_A,
+                    magnitude,
+                    mag_scale,
+                    pre_scaled,  # save unscaled for correct d_mag
+                )
+            else:
+                result = result + s * lora_result
+                ctx.save_for_backward(x, A.to(dtype), B.to(dtype), after_A)  # type: ignore[union-attr]
+        else:
+            ctx.save_for_backward(
+                x,
+            )
+
+        ctx.s = s
+        ctx.has_dora = has_dora
+        ctx.has_lora = A is not None
+        ctx.padding_idx = padding_idx
+        ctx.scale_grad_by_freq = scale_grad_by_freq
+
+        return result
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(ctx, grad_output):
+        s = ctx.s
+        has_dora = ctx.has_dora
+        has_lora = ctx.has_lora
+
+        d_A = d_B = d_mag = None
+
+        if not has_lora:
+            (x,) = ctx.saved_tensors
+        elif has_dora:
+            x, A, B, after_A, magnitude, mag_scale, combined = ctx.saved_tensors
+            # DoRA magnitude gradient
+            combined_flat = combined.view(-1, combined.shape[-1])
+            grad_flat = grad_output.view(-1, grad_output.shape[-1])
+            d_mag = (grad_flat * combined_flat).sum(dim=0) * mag_scale / magnitude
+            # Chain rule through mag_scale
+            grad_output = grad_output * mag_scale.unsqueeze(0).unsqueeze(0)
+        else:
+            x, A, B, after_A = ctx.saved_tensors
+
+        if has_lora:
+            # Use float32 for gradient computation (LoRA params are fp32)
+            compute_dtype = torch.float32
+
+            after_A_flat = after_A.view(-1, after_A.shape[-1]).to(compute_dtype)
+            grad_flat = grad_output.view(-1, grad_output.shape[-1]).to(compute_dtype)
+            B_f = B.to(compute_dtype)
+
+            # B is [hidden_dim, rank], B_T = B.t() = [rank, hidden_dim]
+            # lora_result = after_A @ B_T → d/d(B_T) = s * after_A^T @ grad
+            B_T = B_f.t()  # [rank, hidden_dim]
+            d_B_T = torch.empty_like(B_T)
+            d_B_T.addmm_(after_A_flat.t(), grad_flat, alpha=s, beta=0)
+            d_B = d_B_T.t()  # [hidden_dim, rank]
+
+            # d_A: gradient flows through F.embedding lookup
+            # d_after_A = s * grad @ B = [T, hidden] @ [hidden, rank] = [T, rank]
+            d_after_A = s * grad_flat @ B_f
+
+            # F.embedding backward: scatter d_after_A into A^T gradient
+            x_flat = x.view(-1)
+
+            # Zero out padding_idx contributions (matches F.embedding behavior)
+            if ctx.padding_idx is not None:
+                pad_mask = x_flat != ctx.padding_idx
+                d_after_A = d_after_A * pad_mask.unsqueeze(1).to(d_after_A.dtype)
+
+            # scale_grad_by_freq: divide each contribution by token frequency
+            if ctx.scale_grad_by_freq:
+                counts = torch.bincount(x_flat, minlength=A.shape[1]).clamp(min=1)
+                freq_scale = 1.0 / counts[x_flat].unsqueeze(1).to(d_after_A.dtype)
+                d_after_A = d_after_A * freq_scale
+
+            A_f = A.to(compute_dtype)
+            d_A_T = torch.zeros_like(A_f.t())  # [vocab, rank]
+            d_A_T.index_add_(0, x_flat, d_after_A)
+            d_A = d_A_T.t()  # [rank, vocab]
+
+        # x, W, A, B, s, magnitude, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse
+        return (
+            None,  # x
+            None,  # W (base embedding weight grad handled by PyTorch)
+            d_A,  # A
+            d_B,  # B
+            None,  # s
+            d_mag,  # magnitude
+            None,
+            None,
+            None,
+            None,
+            None,  # padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse
+        )
+
+
+def apply_lora_embedding(self, x: torch.Tensor) -> torch.Tensor:
+    """Applies LoRA to embedding layer."""
+    W, A, B, s, dropout, magnitude, base_layer = get_embedding_lora_parameters(self)
+
+    # Capture base output dtype (bf16 for bf16 models) to cast back at end
+    output_dtype = W.dtype
+
+    # Note: PEFT's Embedding forward does not apply dropout for embeddings
+    # (integer indices can't be dropped; PEFT silently ignores lora_dropout here)
+    result = LoRA_Embedding.apply(
+        x,
+        W,
+        A,
+        B,
+        s,
+        magnitude,
+        base_layer.padding_idx,
+        base_layer.max_norm,
+        base_layer.norm_type,
+        base_layer.scale_grad_by_freq,
+        base_layer.sparse,
+    )
+
+    # Cast to model dtype (LoRA ops may upcast to float32)
+    return result.to(output_dtype)
diff --git a/src/axolotl/kernels/op_registry.py b/src/axolotl/kernels/op_registry.py
new file mode 100644
index 0000000000..80d6b8cef7
--- /dev/null
+++ b/src/axolotl/kernels/op_registry.py
@@ -0,0 +1,65 @@
+"""Registration helper for axolotl's custom torch ops.
+
+Kernel entry points registered through here become dispatcher-visible:
+``torch.compile`` treats them as opaque ops instead of graph-breaking, and
+policy-based selective checkpointing can match them by name. Registration
+failure (exotic torch builds) falls back to the raw python callable so the
+kernel keeps working — it just stays invisible to the dispatcher.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Sequence
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+NAMESPACE = "axolotl"
+
+_warned_registration_failure = False
+
+
+class _UnregisteredOp:
+    """Callable shim with the CustomOpDef surface we use, for fallback."""
+
+    def __init__(self, fn: Callable):
+        self._fn = fn
+
+    def __call__(self, *args, **kwargs):
+        return self._fn(*args, **kwargs)
+
+    def register_fake(self, fn: Callable) -> Callable:
+        return fn
+
+
+def register_kernel_op(name: str, mutates_args: Sequence[str] = ()) -> Callable:
+    """Decorator: register the function as ``torch.ops.axolotl.<name>``.
+
+    Returns the registered op (callable), or a passthrough shim wrapping the
+    raw function if registration is unavailable.
+    """
+
+    def decorator(fn: Callable):
+        try:
+            return torch.library.custom_op(
+                f"{NAMESPACE}::{name}", mutates_args=tuple(mutates_args)
+            )(fn)
+        except Exception as exc:  # pylint: disable=broad-exception-caught
+            global _warned_registration_failure
+            message = (
+                f"custom op registration failed for {NAMESPACE}::{name} "
+                f"({type(exc).__name__}: {exc}); kernel will run unregistered "
+                "(functional, but invisible to torch.compile and selective "
+                "checkpointing)"
+            )
+            if _warned_registration_failure:
+                LOG.debug(message)
+            else:
+                _warned_registration_failure = True
+                LOG.warning(message)
+            return _UnregisteredOp(fn)
+
+    return decorator
diff --git a/src/axolotl/kernels/quantize.py b/src/axolotl/kernels/quantize.py
new file mode 100644
index 0000000000..08ca5055bf
--- /dev/null
+++ b/src/axolotl/kernels/quantize.py
@@ -0,0 +1,213 @@
+"""Dequantization utilities for `bitsandbytes` and FP8 integration."""
+
+import ctypes
+from typing import List
+
+import bitsandbytes as bnb
+import torch
+from bitsandbytes.functional import QuantState, get_ptr
+
+cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
+cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4
+cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4
+cdequantize_blockwise_fp32_nf4 = bnb.functional.lib.cdequantize_blockwise_fp32_nf4
+
+# NF4 dequant kernel per output dtype; the buffer dtype must match or the kernel
+# writes past/short of each element (e.g. a bf16 kernel into an fp32 buffer = garbage).
+_NF4_DEQUANT_KERNELS = {
+    torch.float16: cdequantize_blockwise_fp16_nf4,
+    torch.bfloat16: cdequantize_blockwise_bf16_nf4,
+    torch.float32: cdequantize_blockwise_fp32_nf4,
+}
+
+# Cached per-device: per-call current_stream() measurably slows this hot path.
+CUDA_STREAM: dict[torch.device, torch.cuda.Stream] = {}
+
+
+def _ctypes_nf4_dequant(
+    W: torch.Tensor,
+    absmax: torch.Tensor,
+    code2: torch.Tensor,
+    absmax2: torch.Tensor,
+    offset: torch.Tensor,
+    blocksize: int,
+    blocksize2: int,
+    shape,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Direct-ctypes NF4 double-dequant body (Unsloth-derived fast path)."""
+    target_device = W.device
+    out = torch.empty(tuple(shape), dtype=dtype, device=target_device)
+    n_elements_absmax = absmax.numel()
+    out_absmax = torch.empty(
+        n_elements_absmax, dtype=torch.float32, device=target_device
+    )
+
+    stream = CUDA_STREAM.get(target_device)
+    if stream is None:
+        stream = CUDA_STREAM.setdefault(
+            target_device, torch.cuda.current_stream(target_device)
+        )
+
+    cdequantize_blockwise_fp32(
+        get_ptr(code2),
+        get_ptr(absmax),
+        get_ptr(absmax2),
+        get_ptr(out_absmax),
+        ctypes.c_int(blocksize2),
+        ctypes.c_int(n_elements_absmax),
+        stream,
+    )
+    out_absmax += offset
+
+    fx = _NF4_DEQUANT_KERNELS.get(dtype)
+    if fx is None:
+        raise ValueError(f"NF4 dequantization unsupported for output dtype {dtype}")
+    fx(
+        get_ptr(None),
+        get_ptr(W),
+        get_ptr(out_absmax),
+        get_ptr(out),
+        ctypes.c_int(blocksize),
+        ctypes.c_int(out.numel()),
+        stream,
+    )
+
+    # bnb convention: leading-dim-1 packed weight signals a transposed view.
+    if W.shape[0] == 1:
+        return out.t()
+    return out
+
+
+@torch.library.custom_op("axolotl::nf4_dequantize", mutates_args=())
+def _nf4_dequantize_op(
+    W: torch.Tensor,
+    absmax: torch.Tensor,
+    code2: torch.Tensor,
+    absmax2: torch.Tensor,
+    offset: torch.Tensor,
+    blocksize: int,
+    blocksize2: int,
+    shape: List[int],
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Opaque-to-Dynamo wrapper around the direct-ctypes NF4 dequant body."""
+    return _ctypes_nf4_dequant(
+        W, absmax, code2, absmax2, offset, blocksize, blocksize2, shape, dtype
+    )
+
+
+@_nf4_dequantize_op.register_fake
+def _(W, absmax, code2, absmax2, offset, blocksize, blocksize2, shape, dtype):
+    """FakeTensor shape/dtype inference for the registered op (trace-time only)."""
+    out = torch.empty(tuple(shape), dtype=dtype, device=W.device)
+    if W.shape[0] == 1:
+        return out.t()
+    return out
+
+
+def dequantize_fp8(
+    W: torch.Tensor,
+    scale_inv: torch.Tensor,
+    dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Dequantize FP8 block-quantized weights: W_dequant = W_fp8 * scale_inv."""
+    W_float = W.to(dtype)
+    if scale_inv.numel() == 1:
+        return W_float * scale_inv.to(dtype)
+    if scale_inv.dim() == 2 and W.dim() == 2:
+        sr, sc = scale_inv.shape
+        br = W.shape[0] // sr
+        bc = W.shape[1] // sc
+        if sr * br == W.shape[0] and sc * bc == W.shape[1]:
+            return (
+                W_float.reshape(sr, br, sc, bc) * scale_inv[:, None, :, None].to(dtype)
+            ).reshape(W.shape)
+        # Tail blocks: ceil-div the block size, tile scale_inv, crop to W.
+        br_ceil = -(-W.shape[0] // sr)
+        bc_ceil = -(-W.shape[1] // sc)
+        scale_expanded = (
+            scale_inv.to(dtype)
+            .repeat_interleave(br_ceil, dim=0)
+            .repeat_interleave(bc_ceil, dim=1)
+        )[: W.shape[0], : W.shape[1]]
+        return W_float * scale_expanded
+    return W_float * scale_inv.to(dtype)
+
+
+_FLOAT8_CLS: type | None = None
+_FLOAT8_CHECKED = False
+
+
+def _is_float8_tensor(W: torch.Tensor) -> bool:
+    """A torchao ``Float8Tensor`` (blockwise-FP8 base weight, e.g. DSV4 non-experts).
+
+    It reports a logical bf16 ``dtype`` and carries its own block scale, so the fused
+    LoRA kernels pass it as ``W`` with ``quant_state=None`` — detect it by class."""
+    global _FLOAT8_CLS, _FLOAT8_CHECKED
+    if not _FLOAT8_CHECKED:
+        _FLOAT8_CHECKED = True
+        try:
+            from torchao.quantization import Float8Tensor as _F8
+
+            _FLOAT8_CLS = _F8
+        except Exception:
+            _FLOAT8_CLS = None
+    return _FLOAT8_CLS is not None and isinstance(W, _FLOAT8_CLS)
+
+
+def dequantize(
+    W: torch.Tensor,
+    quant_state: QuantState | torch.Tensor | None = None,
+) -> torch.Tensor:
+    """NF4 / FP8 dequantization; under `torch.compile` NF4 dispatches via `torch.ops.axolotl.nf4_dequantize`."""
+    # torchao Float8Tensor carries its own scale (a transposed view is still a Float8Tensor),
+    # so dequant to bf16 here and let the downstream matmul/addmm_ stay a plain bf16 GEMM.
+    if _is_float8_tensor(W):
+        return W.dequantize().to(torch.bfloat16)
+
+    if quant_state is None:
+        return W
+
+    if W.dtype == torch.float8_e4m3fn:
+        scale_inv = quant_state
+        # Caller may pass W.t() (non-contiguous); dequant in original layout, transpose back.
+        if not W.is_contiguous() and W.dim() == 2:
+            return dequantize_fp8(W.t(), scale_inv).t()
+        return dequantize_fp8(W, scale_inv)
+
+    # Non-double-quant: fall back to bnb's wrapper (rare in axolotl QLoRA).
+    if quant_state.offset is None or quant_state.state2 is None:
+        return bnb.functional.dequantize_4bit(W, quant_state, quant_type="nf4")
+
+    target_device = W.device
+    state2 = quant_state.state2
+    absmax = quant_state.absmax.to(target_device)
+    code2 = state2.code.to(target_device)
+    absmax2 = state2.absmax.to(target_device)
+    offset = quant_state.offset.to(target_device)
+
+    if torch.compiler.is_compiling():
+        return torch.ops.axolotl.nf4_dequantize.default(
+            W,
+            absmax,
+            code2,
+            absmax2,
+            offset,
+            quant_state.blocksize,
+            state2.blocksize,
+            list(quant_state.shape),
+            quant_state.dtype,
+        )
+
+    return _ctypes_nf4_dequant(
+        W,
+        absmax,
+        code2,
+        absmax2,
+        offset,
+        quant_state.blocksize,
+        state2.blocksize,
+        quant_state.shape,
+        quant_state.dtype,
+    )
diff --git a/src/axolotl/kernels/rms_norm_gated.py b/src/axolotl/kernels/rms_norm_gated.py
new file mode 100644
index 0000000000..476bf7f776
--- /dev/null
+++ b/src/axolotl/kernels/rms_norm_gated.py
@@ -0,0 +1,381 @@
+"""
+Fused RMSNorm + SiLU Gate Triton kernel.
+
+Computes: Y = (W + offset) * RMSNorm(X) * silu(G)
+where RMSNorm(X) = X / sqrt(mean(X^2) + eps)
+and silu(G) = G * sigmoid(G)
+
+Used by Qwen3.5's GatedDeltaNet linear attention layers (Qwen3_5RMSNormGated).
+"""
+
+import math
+import operator
+
+import torch
+import triton
+import triton.language as tl
+from liger_kernel.ops.utils import (
+    calculate_settings,
+    compare_version,
+    ensure_contiguous,
+    torch_to_triton_dtype,
+)
+from liger_kernel.utils import is_npu_available
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
+    try:
+        from triton.language.extra.libdevice import rsqrt
+    except ModuleNotFoundError:
+        from triton.language.extra.cuda.libdevice import rsqrt
+else:
+    from triton.language.math import rsqrt
+
+
+@triton.jit
+def _rms_norm_gated_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    G_ptr,
+    G_row_stride,
+    W_ptr,
+    W_row_stride,
+    RSTD_ptr,
+    RSTD_row_stride,
+    n_cols,
+    eps,
+    offset,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Y = (W + offset) * (X / RMS(X)) * silu(G)
+
+    All computation done in fp32 (Gemma-style), result cast to input dtype.
+    """
+    row_idx = tl.program_id(0).to(tl.int64)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    X_row = tl.load(X_ptr + row_idx * X_row_stride + col_offsets, mask=mask, other=0)
+    G_row = tl.load(G_ptr + row_idx * G_row_stride + col_offsets, mask=mask, other=0)
+    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
+
+    X_row_dtype = X_row.dtype
+
+    # Cast everything to fp32
+    X_fp32 = X_row.to(tl.float32)
+    G_fp32 = G_row.to(tl.float32)
+    W_fp32 = W_row.to(tl.float32)
+
+    # RMS norm
+    mean_sq = tl.sum(X_fp32 * X_fp32, axis=0) / n_cols
+    rstd = rsqrt(mean_sq + eps)
+    tl.store(RSTD_ptr + row_idx * RSTD_row_stride, rstd)
+
+    X_norm = X_fp32 * rstd
+
+    # SiLU gate: silu(G) = G * sigmoid(G)
+    sig_G = tl.sigmoid(G_fp32)
+    silu_G = G_fp32 * sig_G
+
+    # Fused output
+    Y_row = (offset + W_fp32) * X_norm * silu_G
+
+    tl.store(
+        Y_ptr + row_idx * Y_row_stride + col_offsets,
+        Y_row.to(X_row_dtype),
+        mask=mask,
+    )
+
+
+@triton.jit
+def _rms_norm_gated_backward_kernel(
+    dY_ptr,
+    dY_row_stride,
+    dX_ptr,
+    dX_row_stride,
+    dG_ptr,
+    dG_row_stride,
+    X_ptr,
+    X_row_stride,
+    X_dtype: tl.constexpr,
+    G_ptr,
+    G_row_stride,
+    W_ptr,
+    W_row_stride,
+    RSTD_ptr,
+    RSTD_row_stride,
+    dW_ptr,
+    dW_row_stride,
+    n_rows,
+    n_cols,
+    offset,
+    rows_per_program,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Backward for Y = (W + offset) * (X * RSTD) * silu(G)
+
+    dW = sum_batch(dY * X_norm * silu(G))
+    dG = dY * (W + offset) * X_norm * silu'(G)
+       where silu'(G) = sigmoid(G) * (1 + G * (1 - sigmoid(G)))
+    dX = RSTD * (m - (1/N) * RSTD^2 * dot(m, X) * X)
+       where m = dY * (W + offset) * silu(G)
+    """
+    row_block_id = tl.program_id(0).to(tl.int64)
+    row_start = row_block_id * rows_per_program
+    row_end = min((row_block_id + 1) * rows_per_program, n_rows)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    dW_acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+
+    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
+    W_row = W_row.to(tl.float32) + offset
+
+    for row_idx in range(row_start, row_end):
+        dY_row = tl.load(
+            dY_ptr + row_idx * dY_row_stride + col_offsets, mask=mask, other=0.0
+        )
+        X_row = tl.load(
+            X_ptr + row_idx * X_row_stride + col_offsets, mask=mask, other=0.0
+        )
+        G_row = tl.load(
+            G_ptr + row_idx * G_row_stride + col_offsets, mask=mask, other=0.0
+        )
+        rstd_row = tl.load(RSTD_ptr + row_idx * RSTD_row_stride)
+
+        # Cast to fp32
+        dY_fp32 = dY_row.to(tl.float32)
+        X_fp32 = X_row.to(tl.float32)
+        G_fp32 = G_row.to(tl.float32)
+
+        # Recompute intermediates
+        X_norm = X_fp32 * rstd_row
+        sig_G = tl.sigmoid(G_fp32)
+        silu_G = G_fp32 * sig_G
+
+        # dW: accumulate dY * X_norm * silu(G)
+        dW_acc += dY_fp32 * X_norm * silu_G
+
+        # dG: dY * (W + offset) * X_norm * silu'(G)
+        # silu'(G) = sigmoid(G) * (1 + G * (1 - sigmoid(G)))
+        silu_prime_G = sig_G * (1.0 + G_fp32 * (1.0 - sig_G))
+        dG_row = dY_fp32 * W_row * X_norm * silu_prime_G
+        tl.store(
+            dG_ptr + row_idx * dG_row_stride + col_offsets,
+            dG_row.to(X_dtype),
+            mask=mask,
+        )
+
+        # dX: standard RMSNorm backward with effective gradient m = dY * W * silu(G)
+        m = dY_fp32 * W_row * silu_G
+        dX_row = rstd_row * m
+        dX_row += rstd_row * (
+            -(1.0 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_fp32, axis=0) * X_fp32
+        )
+        tl.store(
+            dX_ptr + row_idx * dX_row_stride + col_offsets,
+            dX_row.to(X_dtype),
+            mask=mask,
+        )
+
+    tl.store(
+        dW_ptr + row_block_id * dW_row_stride + col_offsets,
+        dW_acc,
+        mask=mask,
+    )
+
+
+@register_kernel_op("rms_norm_gated_fwd")
+def _rms_norm_gated_fwd_op(
+    X: torch.Tensor,
+    G: torch.Tensor,
+    W: torch.Tensor,
+    eps: float,
+    offset: float,
+    BLOCK_SIZE: int,
+    num_warps: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    n_rows, n_cols = X.shape
+
+    Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
+    RSTD = torch.empty(n_rows, dtype=torch.float32, device=X.device)
+
+    _rms_norm_gated_forward_kernel[(n_rows,)](
+        Y,
+        Y.stride(0),
+        X,
+        X.stride(0),
+        G,
+        G.stride(0),
+        W,
+        W.stride(0),
+        RSTD,
+        RSTD.stride(0),
+        n_cols,
+        eps,
+        offset,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return Y, RSTD
+
+
+@_rms_norm_gated_fwd_op.register_fake
+def _(X, G, W, eps, offset, BLOCK_SIZE, num_warps):
+    return torch.empty_like(X), torch.empty(
+        X.shape[0], dtype=torch.float32, device=X.device
+    )
+
+
+@register_kernel_op("rms_norm_gated_bwd")
+def _rms_norm_gated_bwd_op(
+    dY: torch.Tensor,
+    X: torch.Tensor,
+    G: torch.Tensor,
+    W: torch.Tensor,
+    RSTD: torch.Tensor,
+    offset: float,
+    BLOCK_SIZE: int,
+    num_warps: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    n_rows, n_cols = dY.shape
+
+    sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
+
+    _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
+    dX = torch.empty_like(dY)
+    dG = torch.empty_like(dY)
+
+    rows_per_program = math.ceil(n_rows / sm_count)
+    grid = (sm_count,)
+
+    _rms_norm_gated_backward_kernel[grid](
+        dY,
+        dY.stride(0),
+        dX,
+        dX.stride(0),
+        dG,
+        dG.stride(0),
+        X,
+        X.stride(0),
+        torch_to_triton_dtype[X.dtype],
+        G,
+        G.stride(0),
+        W,
+        W.stride(0),
+        RSTD,
+        RSTD.stride(0),
+        _dW,
+        _dW.stride(0),
+        n_rows,
+        n_cols,
+        offset,
+        rows_per_program,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+
+    dW = _dW.sum(dim=0).to(W.dtype)
+    return dX, dG, dW
+
+
+@_rms_norm_gated_bwd_op.register_fake
+def _(dY, X, G, W, RSTD, offset, BLOCK_SIZE, num_warps):
+    return torch.empty_like(dY), torch.empty_like(dY), torch.empty_like(W)
+
+
+def rms_norm_gated_forward(X, G, W, eps, offset):
+    shape = X.shape
+    dim = shape[-1]
+    X = X.view(-1, dim)
+    G = G.view(-1, dim)
+    BLOCK_SIZE, num_warps = calculate_settings(X.shape[-1])
+
+    assert X.shape[1] == W.shape[0], (
+        f"Incompatible hidden size: X.shape[1]={X.shape[1]} vs W.shape[0]={W.shape[0]}"
+    )
+    assert X.shape == G.shape, (
+        f"X and G must have same shape, got {X.shape} and {G.shape}"
+    )
+
+    Y, RSTD = _rms_norm_gated_fwd_op(X, G, W, eps, offset, BLOCK_SIZE, num_warps)
+    return Y.view(*shape), X, G, RSTD, BLOCK_SIZE, num_warps
+
+
+def rms_norm_gated_backward(dY, X, G, W, RSTD, offset, BLOCK_SIZE, num_warps):
+    shape = dY.shape
+    dim = shape[-1]
+    dY = dY.view(-1, dim)
+
+    dX, dG, dW = _rms_norm_gated_bwd_op(
+        dY, X, G, W, RSTD, offset, BLOCK_SIZE, num_warps
+    )
+
+    dX = dX.view(*shape)
+    dG = dG.view(*shape)
+    return dX, dG, dW
+
+
+class FusedRMSNormGatedFunction(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, X, G, W, eps, offset=0.0):
+        """
+        X: (B, T, H) or (BxT, H) — input hidden states
+        G: (B, T, H) or (BxT, H) — gate tensor
+        W: (H,) — weight parameter
+        """
+        Y, X, G, RSTD, BLOCK_SIZE, num_warps = rms_norm_gated_forward(
+            X, G, W, eps, offset
+        )
+        ctx.offset = offset
+        ctx.BLOCK_SIZE = BLOCK_SIZE
+        ctx.num_warps = num_warps
+        ctx.save_for_backward(X, G, W, RSTD)
+        return Y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dY):
+        X, G, W, RSTD = ctx.saved_tensors
+        dX, dG, dW = rms_norm_gated_backward(
+            dY, X, G, W, RSTD, ctx.offset, ctx.BLOCK_SIZE, ctx.num_warps
+        )
+        return dX, dG, dW, None, None
+
+
+class FusedRMSNormGated(torch.nn.Module):
+    """
+    Fused RMSNorm + SiLU Gate.
+
+    Computes: Y = W * RMSNorm(X) * silu(G)
+
+    Drop-in replacement for Qwen3_5RMSNormGated with matching
+    init signature: __init__(hidden_size, eps=1e-6, **kwargs)
+    and forward signature: forward(hidden_states, gate=None)
+    """
+
+    def __init__(self, hidden_size, eps=1e-6, offset=0.0, **kwargs):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+        self.offset = offset
+
+    def forward(self, hidden_states, gate=None):
+        if gate is None:
+            raise ValueError("FusedRMSNormGated requires a gate tensor")
+        if hidden_states.device.type != "cuda":
+            raise ValueError(
+                f"FusedRMSNormGated requires CUDA tensors, got device={hidden_states.device}"
+            )
+        return FusedRMSNormGatedFunction.apply(
+            hidden_states, gate, self.weight, self.variance_epsilon, self.offset
+        )
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
diff --git a/src/axolotl/kernels/swiglu.py b/src/axolotl/kernels/swiglu.py
new file mode 100644
index 0000000000..6cdd21cd63
--- /dev/null
+++ b/src/axolotl/kernels/swiglu.py
@@ -0,0 +1,178 @@
+"""
+Module for definition of SwiGLU Triton kernels.
+
+See "GLU Variants Improve Transformer" (https://arxiv.org/abs/2002.05202).
+
+Credit to `unsloth` (https://unsloth.ai/) for inspiration for this implementation.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+
+@triton.jit
+def _swiglu_fwd_kernel(
+    gate_ptr,
+    up_ptr,
+    out_ptr,
+    n_elements,
+    block_size: tl.constexpr,
+):
+    """
+    SwiGLU forward kernel. The kernel computes activation in fp32 precision for better
+    numerical stability, then converts back to original dtype for the final result.
+
+    Args:
+        gate_ptr: Pointer to gate tensor `[*, hidden_dim]`.
+        up_ptr: Pointer to up-projection tensor `[*, hidden_dim]`.
+        out_ptr: Pointer to output tensor `[*, hidden_dim]`.
+        n_elements: Total number of elements in the input tensors.
+        block_size: Size of thread blocks for parallel computation.
+    """
+    block_idx = tl.program_id(0)
+    offsets = block_idx * block_size + tl.arange(0, block_size)
+    mask = offsets < n_elements
+
+    # Load gate in fp32, keep up in original dtype
+    gate = tl.load(gate_ptr + offsets, mask=mask, other=0).to(tl.float32)
+    up = tl.load(up_ptr + offsets, mask=mask, other=0)
+
+    # Compute activation in fp32 then convert back
+    f = gate * tl.sigmoid(gate)
+    f = f.to(up.dtype)
+    result = f * up
+
+    tl.store(out_ptr + offsets, result, mask=mask)
+
+
+@triton.jit
+def _swiglu_bwd_kernel(
+    grad_out_ptr,
+    gate_ptr,
+    up_ptr,
+    n_elements,
+    block_size: tl.constexpr,
+):
+    """
+    SwiGLU backward kernel. Stores gradient results in-place.
+
+    Args:
+        grad_out_ptr: Pointer to gradient output tensor `[*, hidden_dim]`.
+        gate_ptr: Pointer to gate tensor `[*, hidden_dim]`.
+        up_ptr: Pointer to up-projection tensor `[*, hidden_dim]`.
+        n_elements: Total number of elements in the input tensors.
+        block_size: Size of thread blocks for parallel computation.
+
+    Note:
+        After kernel execution, tensors are modified in-place:
+        - `grad_out_ptr` contains forward output (`h`)
+        - `gate_ptr` contains gradient w.r.t gate (`grad_gate`)
+        - `up_ptr` contains gradient w.r.t up (`grad_up`)
+    """
+    block_idx = tl.program_id(0)
+    offsets = block_idx * block_size + tl.arange(0, block_size)
+    mask = offsets < n_elements
+
+    # Load values - only convert gate to fp32
+    grad_out = tl.load(grad_out_ptr + offsets, mask=mask, other=0)
+    gate = tl.load(gate_ptr + offsets, mask=mask, other=0).to(tl.float32)
+    up = tl.load(up_ptr + offsets, mask=mask, other=0)
+
+    # Compute SiLU and forward output
+    sigmoid_gate = tl.sigmoid(gate)
+    silu_gate = sigmoid_gate * gate
+    silu_gate = silu_gate.to(grad_out.dtype)
+    h = silu_gate * up
+
+    # Compute gradients
+    grad_up = grad_out * silu_gate  # gradient for up is grad_out * SiLU(gate)
+
+    # Compute gate gradient
+    temp = grad_out * up
+    grad_gate = temp.to(tl.float32) * sigmoid_gate * (1.0 + gate * (1.0 - sigmoid_gate))
+    grad_gate = grad_gate.to(grad_out.dtype)
+
+    # Store results with correct gradient ordering
+    tl.store(grad_out_ptr + offsets, h, mask=mask)
+    tl.store(gate_ptr + offsets, grad_gate, mask=mask)  # grad wrt gate
+    tl.store(up_ptr + offsets, grad_up, mask=mask)  # grad wrt up
+
+
+@register_kernel_op("swiglu_fwd")
+def swiglu_forward(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    """
+    SwiGLU forward pass. Computes SwiGLU activation: `x * sigmoid(x) * up`, where
+    `x` is the gate tensor.
+
+    Args:
+        gate: Input gate tensor of shape `[batch, seq_len, hidden_dim]`.
+        up: Up-projection tensor of shape `[batch, seq_len, hidden_dim]`.
+
+    Returns:
+        Output tensor of shape `[batch, seq_len, hidden_dim]`.
+    """
+    batch, seq_len, hidden_dim = gate.shape
+    n_elements = gate.numel()
+    out = torch.empty((batch, seq_len, hidden_dim), dtype=gate.dtype, device="cuda")
+
+    grid = lambda meta: (triton.cdiv(n_elements, meta["block_size"]),)  # noqa: E731
+    _swiglu_fwd_kernel[grid](
+        gate_ptr=gate,
+        up_ptr=up,
+        out_ptr=out,
+        n_elements=n_elements,
+        block_size=1024,
+    )
+
+    return out
+
+
+@swiglu_forward.register_fake
+def _(gate, up):
+    return torch.empty_like(gate)
+
+
+@register_kernel_op("swiglu_bwd", mutates_args=("grad_output", "gate", "up"))
+def _swiglu_bwd_op(
+    grad_output: torch.Tensor, gate: torch.Tensor, up: torch.Tensor
+) -> None:
+    n_elements = grad_output.numel()
+
+    grid = lambda meta: (triton.cdiv(n_elements, meta["block_size"]),)  # noqa: E731
+    _swiglu_bwd_kernel[grid](
+        grad_out_ptr=grad_output,
+        gate_ptr=gate,
+        up_ptr=up,
+        n_elements=n_elements,
+        block_size=1024,
+    )
+
+
+@_swiglu_bwd_op.register_fake
+def _(grad_output, gate, up):
+    return None
+
+
+def swiglu_backward(
+    grad_output: torch.Tensor, gate: torch.Tensor, up: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    SwiGLU backward pass using in-place operations.
+
+    Args:
+        grad_output: Gradient of loss with respect to output, shape `[batch, seq_len, hidden_dim]`.
+        gate: Gate tensor from forward pass, shape `[batch, seq_len, hidden_dim]`.
+        up: Up-projection tensor from forward pass, shape `[batch, seq_len, hidden_dim]`.
+
+    Returns:
+        Tuple containing:
+            - Forward pass output (`h`)
+            - Gradient with respect to gate (`df`)
+            - Gradient with respect to up-projection (`de`)
+    """
+    # op mutates in place; return the aliases here — op outputs must not alias inputs
+    _swiglu_bwd_op(grad_output, gate, up)
+    return grad_output, gate, up
diff --git a/src/axolotl/kernels/utils.py b/src/axolotl/kernels/utils.py
new file mode 100644
index 0000000000..59a7e00127
--- /dev/null
+++ b/src/axolotl/kernels/utils.py
@@ -0,0 +1,11 @@
+"""Utilities for `axolotl.kernels` submodules."""
+
+import torch
+from packaging.version import Version
+
+if Version(torch.__version__) < Version("2.4.0"):
+    torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
+    torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
+else:
+    torch_amp_custom_fwd = torch.amp.custom_fwd(device_type="cuda")
+    torch_amp_custom_bwd = torch.amp.custom_bwd(device_type="cuda")
diff --git a/src/axolotl/loaders/__init__.py b/src/axolotl/loaders/__init__.py
new file mode 100644
index 0000000000..f257578a60
--- /dev/null
+++ b/src/axolotl/loaders/__init__.py
@@ -0,0 +1,26 @@
+"""Init for axolotl.loaders module"""
+
+# flake8: noqa
+
+from axolotl.utils import make_lazy_getattr
+
+from .constants import MULTIMODAL_AUTO_MODEL_MAPPING
+
+__all__ = [
+    "MULTIMODAL_AUTO_MODEL_MAPPING",
+    "ModelLoader",
+    "load_adapter",
+    "load_lora",
+    "load_processor",
+    "load_tokenizer",
+]
+
+_LAZY_IMPORTS = {
+    "ModelLoader": ".model",
+    "load_adapter": ".adapter",
+    "load_lora": ".adapter",
+    "load_processor": ".processor",
+    "load_tokenizer": ".tokenizer",
+}
+
+__getattr__ = make_lazy_getattr(_LAZY_IMPORTS, __name__, globals())
diff --git a/src/axolotl/loaders/adapter.py b/src/axolotl/loaders/adapter.py
new file mode 100644
index 0000000000..2ae83e2efa
--- /dev/null
+++ b/src/axolotl/loaders/adapter.py
@@ -0,0 +1,489 @@
+"""Adapter loading functionality, including LoRA / QLoRA and associated utils"""
+
+import json
+import os
+import types
+from pathlib import Path
+from typing import Any
+
+import bitsandbytes as bnb
+import torch
+from bitsandbytes.nn import Params4bit
+from peft import (
+    AdaptionPromptConfig,
+    LoftQConfig,
+    LoraConfig,
+    PeftConfig,
+    PeftMixedModel,
+    PeftModel,
+    TaskType,
+    get_peft_model,
+)
+from transformers import PreTrainedModel
+
+from axolotl.integrations.base import PluginManager
+from axolotl.loaders.utils import get_linear_embedding_layers
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+PLUGIN_MANAGER = PluginManager.get_instance()
+
+
+def setup_quantized_meta_for_peft(model: torch.nn.Module):
+    """Replaces `quant_state.to` with a dummy function to prevent PEFT from moving `quant_state` to meta device"""
+
+    def temp_to_method(self, *args, **kwargs):
+        return self
+
+    for param in model.parameters():
+        if isinstance(param, Params4bit) and param.quant_state is not None:
+            param.quant_state._orig_to = param.quant_state.to
+            param.quant_state.to = types.MethodType(temp_to_method, param.quant_state)
+
+
+def setup_quantized_peft_meta_for_training(model: torch.nn.Module):
+    """Replaces dummy `quant_state.to` method with the original function to allow training to continue"""
+    for param in model.parameters():
+        if isinstance(param, Params4bit) and hasattr(param.quant_state, "_orig_to"):
+            param.quant_state.to = param.quant_state._orig_to
+            param.quant_state._orig_to = None
+
+
+def find_all_linear_names(model):
+    cls = (bnb.nn.Linear4bit, bnb.nn.Linear8bitLt, torch.nn.Linear)
+    lora_module_names = set()
+    for name, module in model.named_modules():
+        if (
+            isinstance(module, cls)
+            or "Linear" in module.__class__.__name__
+            and module.__class__.__name__ not in ("LlamaLinearScalingRotaryEmbedding",)
+        ):
+            names = name.split(".")
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+
+    embedding_modules = get_linear_embedding_layers(model.config.model_type)
+    output_embedding = embedding_modules[1]
+    if output_embedding in lora_module_names:  # needed for 16-bit
+        lora_module_names.remove(output_embedding)
+
+    return list(lora_module_names)
+
+
+def _patch_peft_clippable_linear():
+    """Patch PEFT to handle Gemma4ClippableLinear which wraps nn.Linear.
+
+    Gemma4's vision tower uses ClippableLinear (a thin wrapper around nn.Linear
+    that clips activations). PEFT doesn't recognise it as a supported layer type,
+    so we redirect LoRA injection to the inner ``.linear`` child instead.
+    """
+    try:
+        from transformers.models.gemma4.modeling_gemma4 import (
+            Gemma4ClippableLinear as _cls,
+        )
+    except ImportError:
+        return
+
+    from peft.tuners.lora.model import LoraModel
+
+    if getattr(LoraModel, "_axolotl_clippable_patched", False):
+        return
+    _orig = LoraModel._create_and_replace
+
+    def _patched(
+        self,
+        peft_config,
+        adapter_name,
+        target,
+        target_name,
+        parent,
+        current_key=None,
+        **kw,
+    ):
+        if isinstance(target, _cls):
+            # Redirect to the inner nn.Linear so PEFT can wrap it normally.
+            return _orig(
+                self,
+                peft_config,
+                adapter_name,
+                target.linear,
+                "linear",
+                target,
+                current_key=current_key,
+                **kw,
+            )
+        return _orig(
+            self,
+            peft_config,
+            adapter_name,
+            target,
+            target_name,
+            parent,
+            current_key=current_key,
+            **kw,
+        )
+
+    LoraModel._create_and_replace = _patched
+    LoraModel._axolotl_clippable_patched = True
+
+
+def _get_peft_task_type(model: PreTrainedModel) -> TaskType:
+    model_cls = type(model).__name__
+    if "SequenceClassification" in model_cls:
+        return TaskType.SEQ_CLS
+    if "TokenClassification" in model_cls:
+        return TaskType.TOKEN_CLS
+    return TaskType.CAUSAL_LM
+
+
+def _build_lora_config_kwargs(cfg: DictDefault) -> dict[str, Any]:
+    lora_config_kwargs: dict[str, Any] = {}
+    loftq_bits = cfg.peft and cfg.peft.loftq_config and cfg.peft.loftq_config.loftq_bits
+    if loftq_bits:
+        lora_config_kwargs["loftq_config"] = LoftQConfig(loftq_bits=loftq_bits)
+        lora_config_kwargs["init_lora_weights"] = "loftq"
+    if cfg.peft_init_lora_weights:
+        lora_config_kwargs["init_lora_weights"] = cfg.peft_init_lora_weights
+    if cfg.peft_use_dora:
+        lora_config_kwargs["use_dora"] = cfg.peft_use_dora
+        LOG.info("Initializing LoRA weights using dora. This might take longer.")
+    if cfg.peft_use_rslora:
+        lora_config_kwargs["use_rslora"] = cfg.peft_use_rslora
+    if cfg.lora_rank_pattern:
+        lora_config_kwargs["rank_pattern"] = cfg.lora_rank_pattern
+    if cfg.lora_alpha_pattern:
+        lora_config_kwargs["alpha_pattern"] = cfg.lora_alpha_pattern
+    if cfg.peft_layer_replication:
+        lora_config_kwargs["layer_replication"] = cfg.peft_layer_replication
+    if cfg.peft_trainable_token_indices:
+        lora_config_kwargs["trainable_token_indices"] = cfg.peft_trainable_token_indices
+    if cfg.peft_ensure_weight_tying is not None:
+        lora_config_kwargs["ensure_weight_tying"] = cfg.peft_ensure_weight_tying
+
+    return lora_config_kwargs
+
+
+def _build_peft_lora_config(
+    model: PreTrainedModel,
+    cfg: DictDefault,
+) -> PeftConfig:
+    lora_target_modules = cfg.lora_target_modules or []
+    lora_target_parameters = cfg.lora_target_parameters or []
+
+    if cfg.lora_target_linear:
+        linear_names = find_all_linear_names(model)
+        LOG.info(f"found linear modules: {repr(sorted(linear_names))}")
+        lora_target_modules_as_list = (
+            lora_target_modules
+            if isinstance(lora_target_modules, list)
+            else [lora_target_modules]
+        )
+        lora_target_modules = list(set(lora_target_modules_as_list + linear_names))
+
+    lora_config_kwargs = _build_lora_config_kwargs(cfg)
+    lora_config_kwargs.update(PLUGIN_MANAGER.get_lora_config_kwargs(cfg))
+
+    lora_config = LoraConfig(
+        r=cfg.lora_r,
+        lora_alpha=cfg.lora_alpha,
+        target_modules=lora_target_modules,
+        target_parameters=lora_target_parameters,
+        layers_to_transform=cfg.peft_layers_to_transform,
+        layers_pattern=cfg.peft_layers_pattern,
+        lora_dropout=cfg.lora_dropout,
+        fan_in_fan_out=cfg.lora_fan_in_fan_out,
+        modules_to_save=cfg.lora_modules_to_save if cfg.lora_modules_to_save else None,
+        exclude_modules=getattr(cfg, "lora_exclude_modules", None) or None,
+        bias="none",
+        task_type=_get_peft_task_type(model),
+        **lora_config_kwargs,
+    )
+    return lora_config
+
+
+def _peft_will_auto_convert_target_params(model, lora_config) -> bool:
+    """Check whether PEFT will auto-populate target_parameters for this model.
+
+    PEFT 0.19's ``convert_peft_config_for_transformers`` rewrites old MoE
+    ``target_modules`` (e.g. ``w1``/``w2``/``w3`` on Mixtral) into
+    ``target_parameters`` (``gate_up_proj``/``down_proj``) because
+    transformers v5 fused those expert linears into 3D ``nn.Parameter``
+    tensors. PEFT wraps the resulting 3D params with ``ParamWrapper``,
+    which rejects ``lora_dropout != 0``. This probe runs the conversion on
+    a copy of the config so we can detect the situation before
+    ``get_peft_model`` blows up.
+    """
+    if getattr(lora_config, "target_parameters", None):
+        return False
+
+    try:
+        from peft.utils.transformers_weight_conversion import (
+            convert_peft_config_for_transformers,
+            get_model_conversion_mapping,
+        )
+    except ImportError:
+        return False
+
+    import copy
+
+    probe_cfg = copy.deepcopy(lora_config)
+    try:
+        convert_peft_config_for_transformers(
+            probe_cfg,
+            model=model,
+            conversions=get_model_conversion_mapping(model),
+        )
+    except Exception:  # pylint: disable=broad-except
+        return False
+
+    return bool(getattr(probe_cfg, "target_parameters", None))
+
+
+def _patch_peft_param_wrapper_dropout():
+    """Let PEFT's ``ParamWrapper`` silently accept ``lora_dropout != 0``.
+
+    ``ParamWrapper`` wraps 3D expert ``nn.Parameter`` tensors and rejects
+    non-zero dropout because dropout can't be factored out of
+    ``lora_B(lora_A(dropout(x)))`` when the inner op is an expert-indexed
+    matmul. For mixed configs (attention + MoE experts) this is too
+    aggressive — the non-expert ``Linear`` LoRA layers *can* apply dropout
+    and that's usually what the user intended. We pass a copy of the
+    ``LoraConfig`` with ``lora_dropout=0`` only to ``ParamWrapper.__init__``
+    so it builds with ``nn.Identity`` for its internal dropout slot while
+    every other layer type still receives the real dropout value.
+    """
+    from peft.tuners.lora.layer import ParamWrapper
+
+    if getattr(ParamWrapper, "_axolotl_dropout_patched", False):
+        return
+
+    _orig_init = ParamWrapper.__init__
+
+    def _patched_init(
+        self,
+        base_layer,
+        adapter_name,
+        parameter_name,
+        config,
+        *args,
+        **kwargs,
+    ):
+        if getattr(config, "lora_dropout", 0):
+            import copy as _copy
+
+            patched_config = _copy.copy(config)
+            patched_config.lora_dropout = 0.0
+            return _orig_init(
+                self,
+                base_layer,
+                adapter_name,
+                parameter_name,
+                patched_config,
+                *args,
+                **kwargs,
+            )
+        return _orig_init(
+            self,
+            base_layer,
+            adapter_name,
+            parameter_name,
+            config,
+            *args,
+            **kwargs,
+        )
+
+    ParamWrapper.__init__ = _patched_init
+    ParamWrapper._axolotl_dropout_patched = True
+
+
+def _warn_if_patterns_differ_from_saved_adapter(cfg: DictDefault) -> None:
+    """Warn when cfg rank/alpha patterns differ from the saved adapter's, which governs under lora_model_dir."""
+    try:
+        saved = json.loads(
+            (Path(cfg.lora_model_dir) / "adapter_config.json").read_text()
+        )
+    except (OSError, ValueError):
+        # Adapter config not readable locally (e.g. hub id): can't compare, so warn generically.
+        LOG.warning(
+            "lora_rank_pattern/lora_alpha_pattern are ignored when loading an "
+            "existing adapter via lora_model_dir; the saved adapter_config.json "
+            "governs per-module rank/alpha."
+        )
+        return
+    mismatches = [
+        f"{cfg_key}: config={dict(cfg_val or {})} saved={saved.get(saved_key) or {}}"
+        for cfg_key, saved_key, cfg_val in (
+            ("lora_rank_pattern", "rank_pattern", cfg.lora_rank_pattern),
+            ("lora_alpha_pattern", "alpha_pattern", cfg.lora_alpha_pattern),
+        )
+        if dict(cfg_val or {}) != (saved.get(saved_key) or {})
+    ]
+    if mismatches:
+        LOG.warning(
+            "lora_rank_pattern/lora_alpha_pattern in the config differ from the "
+            "adapter loaded via lora_model_dir and are ignored; the saved "
+            "adapter_config.json governs per-module rank/alpha (%s).",
+            "; ".join(mismatches),
+        )
+
+
+def load_lora(
+    model: PreTrainedModel,
+    cfg: DictDefault,
+    inference: bool = False,
+    config_only: bool = False,
+) -> tuple[PreTrainedModel | PeftModel | PeftMixedModel | None, PeftConfig | None]:
+    _patch_peft_clippable_linear()
+    lora_config = _build_peft_lora_config(model, cfg)
+
+    if config_only:
+        return None, lora_config
+
+    if getattr(
+        lora_config, "lora_dropout", 0
+    ) and _peft_will_auto_convert_target_params(model, lora_config):
+        LOG.warning(
+            "lora_dropout=%s requested but PEFT will wrap this model's fused "
+            "MoE expert parameters with ParamWrapper, which cannot apply "
+            "dropout (the 3D einsum can't factor dropout out of "
+            "lora_B(lora_A(dropout(x)))). Dropout will still be applied to "
+            "non-expert LoRA layers (e.g. attention), and expert LoRA layers "
+            "will use nn.Identity for the dropout slot.",
+            lora_config.lora_dropout,
+        )
+        _patch_peft_param_wrapper_dropout()
+
+    rank = int(os.environ.get("LOCAL_RANK", 0))
+
+    if (
+        cfg.fsdp_config
+        and cfg.adapter
+        and cfg.fsdp_config.cpu_ram_efficient_loading
+        and rank != 0
+    ):
+        setup_quantized_meta_for_peft(model)
+
+    model_kwargs: Any = {}
+    if cfg.peft_autocast_adapter_dtype is not None:
+        model_kwargs["autocast_adapter_dtype"] = cfg.peft_autocast_adapter_dtype
+
+    if cfg.lora_model_dir:
+        LOG.debug("Loading pretrained PEFT adapter")
+        if cfg.lora_rank_pattern or cfg.lora_alpha_pattern:
+            _warn_if_patterns_differ_from_saved_adapter(cfg)
+        if cfg.lora_on_cpu:
+            model_kwargs["max_memory"] = {"cpu": "256GiB"}
+            model_kwargs["device_map"] = {"": "cpu"}
+        model = PeftModel.from_pretrained(
+            model,
+            cfg.lora_model_dir,
+            is_trainable=(not inference),
+            **model_kwargs,
+        )
+    else:
+        model = get_peft_model(model, lora_config, **model_kwargs)
+
+    # FP8 models: LoRA A/B inherit FP8 dtype from base weights, but training
+    # requires a compute dtype (bf16/fp16). Cast trainable LoRA params.
+    if cfg.torch_dtype:
+        _fp8_cast_dtype = cfg.torch_dtype
+    elif torch.cuda.is_available() and torch.cuda.is_bf16_supported():
+        _fp8_cast_dtype = torch.bfloat16
+    else:
+        _fp8_cast_dtype = torch.float16
+    for _name, param in model.named_parameters():
+        if param.requires_grad and param.dtype == torch.float8_e4m3fn:
+            param.data = param.data.to(_fp8_cast_dtype)
+
+    if rank == 0:
+        try:
+            model.print_trainable_parameters()
+        except AttributeError as exc:
+            LOG.warning(
+                "Exception caught during model.print_trainable_parameters(): %s", exc
+            )
+    elif (
+        cfg.fsdp_config
+        and cfg.adapter
+        and cfg.fsdp_config.cpu_ram_efficient_loading
+        and rank != 0
+    ):
+        setup_quantized_peft_meta_for_training(model)
+
+    return model, lora_config
+
+
+@send_errors
+def load_adapter(
+    model: PreTrainedModel,
+    cfg: DictDefault,
+    adapter: str | None,
+    inference: bool = False,
+    config_only: bool = False,
+) -> tuple[PreTrainedModel | PeftModel | PeftMixedModel | None, PeftConfig | None]:
+    if adapter is None:
+        return model, None
+    if hasattr(model, "enable_input_require_grads"):
+        model.enable_input_require_grads()
+    if adapter in ["lora", "qlora"]:
+        peft_model, lora_config = load_lora(
+            model, cfg, inference=inference, config_only=config_only
+        )
+        return peft_model, lora_config
+    if adapter == "llama-adapter":
+        if config_only:
+            _, lora_config = load_llama_adapter(model, cfg, config_only=True)
+            return None, lora_config
+        peft_model, lora_config = load_llama_adapter(model, cfg)
+        return peft_model, lora_config
+
+    plugin_loaded = PLUGIN_MANAGER.load_adapter(
+        model,
+        cfg,
+        inference=inference,
+        config_only=config_only,
+    )
+    if plugin_loaded is not None:
+        return plugin_loaded
+
+    adapter_capability = PLUGIN_MANAGER.get_adapter_capability(adapter)
+    if adapter_capability and adapter_capability.lora_like:
+        peft_model, lora_config = load_lora(
+            model, cfg, inference=inference, config_only=config_only
+        )
+        return peft_model, lora_config
+
+    registered = sorted(PLUGIN_MANAGER.adapter_capabilities())
+    registered_msg = ", ".join(registered) if registered else "none"
+    raise NotImplementedError(
+        f"Adapter '{adapter}' is not built in and was not registered by a plugin "
+        f"with loader support. Registered plugin adapters: {registered_msg}"
+    )
+
+
+def load_llama_adapter(
+    model: PreTrainedModel, cfg: DictDefault, config_only: bool = False
+) -> tuple[PeftModel | PeftMixedModel | None, PeftConfig]:
+    peft_config = AdaptionPromptConfig(
+        adapter_layers=cfg.peft_adapter.layers,  # layers (L)
+        adapter_len=cfg.peft_adapter.len,  # prompt length (K)
+        task_type="CAUSAL_LM",
+    )
+
+    if config_only:
+        return None, peft_config
+
+    if cfg.lora_model_dir:
+        LOG.debug("Loading pretrained PEFT - llama_adapter")
+        peft_model = PeftModel.from_pretrained(
+            model,
+            cfg.lora_model_dir,
+            torch_dtype=torch.float16,
+        )
+    else:
+        peft_model = get_peft_model(model, peft_config)
+
+    peft_model.print_trainable_parameters()
+
+    return peft_model, peft_config
diff --git a/src/axolotl/loaders/adapters/__init__.py b/src/axolotl/loaders/adapters/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/loaders/constants.py b/src/axolotl/loaders/constants.py
new file mode 100644
index 0000000000..4939cb28d0
--- /dev/null
+++ b/src/axolotl/loaders/constants.py
@@ -0,0 +1,18 @@
+"""Shared constants for axolotl.loaders module"""
+
+from transformers import AutoModelForImageTextToText
+from transformers.models.auto.modeling_auto import (
+    MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES,
+)
+
+MULTIMODAL_AUTO_MODEL_MAPPING = dict(MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES)
+
+MULTIMODAL_AUTO_MODEL_MAPPING["lfm2-vl"] = AutoModelForImageTextToText
+
+try:
+    from transformers import VoxtralForConditionalGeneration
+
+    # transformers >4.53.2
+    MULTIMODAL_AUTO_MODEL_MAPPING["voxtral"] = VoxtralForConditionalGeneration
+except ImportError:
+    pass
diff --git a/src/axolotl/loaders/model.py b/src/axolotl/loaders/model.py
new file mode 100644
index 0000000000..d449742728
--- /dev/null
+++ b/src/axolotl/loaders/model.py
@@ -0,0 +1,1042 @@
+"""
+Model loader class implementation for loading, configuring, and patching various models.
+"""
+
+import gc
+import math
+import os
+from functools import cached_property
+from importlib.util import find_spec
+from typing import Any
+
+import peft
+import torch
+import transformers
+import transformers.modeling_utils
+from accelerate import init_empty_weights
+from accelerate.parallelism_config import ParallelismConfig
+from peft import (
+    PeftConfig,
+    PeftMixedModel,
+    PeftModel,
+    PeftModelForCausalLM,
+    prepare_model_for_kbit_training,
+)
+from torch.distributed import DeviceMesh
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AwqConfig,
+    BitsAndBytesConfig,
+    GPTQConfig,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+)
+from transformers.integrations.deepspeed import (
+    HfTrainerDeepSpeedConfig,
+    is_deepspeed_zero3_enabled,
+)
+
+from axolotl.common.architectures import MOE_ARCH_BLOCK
+from axolotl.integrations.base import PluginManager
+from axolotl.loaders.adapter import load_adapter
+from axolotl.loaders.constants import MULTIMODAL_AUTO_MODEL_MAPPING
+from axolotl.loaders.patch_manager import PatchManager
+from axolotl.loaders.utils import (
+    get_linear_embedding_layers,
+    get_module_class_from_name,
+    load_model_config,
+)
+from axolotl.models.mamba import fix_mamba_attn_for_loss
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.bench import log_gpu_memory_usage
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.distributed import (
+    build_parallelism_config,
+    get_device_count,
+    get_device_type,
+)
+from axolotl.utils.fp32_norms import (
+    _matches_norm_class,
+    get_fp32_norm_patterns,
+    tag_model_fp32_norms,
+)
+from axolotl.utils.logging import get_logger
+from axolotl.utils.model_shard_quant import load_sharded_model_quant
+from axolotl.utils.schemas.enums import RLType
+
+LOG = get_logger(__name__)
+PLUGIN_MANAGER = PluginManager.get_instance()
+
+
+class ModelLoader:
+    """Manages model configuration, initialization and application of patches during
+    model loading.
+
+    This class orchestrates the entire process of loading a model from configuration to
+    final preparation. It handles device mapping, quantization, attention mechanisms,
+    adapter integration, and various optimizations.
+
+    The loading process includes:
+        - Loading and validating model configuration
+        - Applying monkey patches for optimizations / fixes
+        - Setting up device mapping (including multi-GPU configurations)
+        - Configuring quantization
+        - Setting attention mechanisms (Flash Attention, SDPA, etc.)
+        - Loading and initializing the model
+        - Applying adapters (LoRA, QLoRA, etc.)
+
+    Attributes:
+        model: The loaded model instance (available after load() is called).
+        model_kwargs: Dictionary of keyword arguments passed to model initialization.
+        base_model: Name or path of the base model to load.
+        model_type: Type of model to load (e.g., `AutoModelForCausalLM`).
+        model_config: Configuration object for the model.
+        auto_model_loader: class used for loading the model (default:
+            `AutoModelForCausalLM`).
+    """
+
+    use_parallel_config: bool | None = False
+    parallelism_config: ParallelismConfig | None = None
+    device_mesh: DeviceMesh | None = None
+
+    def __init__(
+        self,
+        cfg: DictDefault,
+        tokenizer: PreTrainedTokenizerBase,
+        *,
+        inference: bool = False,
+        reference_model: bool = False,
+        **kwargs,
+    ):
+        """Initializes the ModelLoader.
+
+        Args:
+            cfg: Configuration dictionary with model and training settings.
+            tokenizer: Tokenizer instance associated with the model.
+            processor: Optional processor for multimodal models. Defaults to None.
+            inference: Whether the model is being loaded for inference mode. Defaults
+                to False.
+            reference_model: Whether this is a reference model (used in setups like DPO
+                training). Defaults to False.
+            **kwargs: Additional keyword arguments (ignored).
+        """
+        self.cfg = cfg
+        self.tokenizer = tokenizer
+        self.inference: bool = inference
+        self.reference_model: bool = reference_model
+
+        # Init model kwargs
+        self.model_kwargs: dict[str, Any] = {}
+        if cfg.overrides_of_model_kwargs:
+            for key, val in cfg.overrides_of_model_kwargs.items():
+                self.model_kwargs[key] = val
+
+        # Init model
+        self.model: PreTrainedModel | PeftModel | PeftMixedModel
+        self.base_model = cfg.base_model
+        self.model_type = cfg.type_of_model
+
+        # Init model config
+        self.model_config = load_model_config(cfg)
+        self.auto_model_loader = AutoModelForCausalLM
+
+        # Initialize the patch manager
+        self.patch_manager = PatchManager(
+            cfg=cfg,
+            model_config=self.model_config,
+            inference=inference,
+        )
+
+    @cached_property
+    def has_flash_attn(self) -> bool:
+        """Check if flash attention is installed."""
+        return find_spec("flash_attn") is not None
+
+    @property
+    def is_fsdp_enabled(self):
+        """Property that determines if FSDP is enabled."""
+        return self.cfg.fsdp_config is not None or self.cfg.fsdp is not None
+
+    @property
+    def is_qlora_and_fsdp_enabled(self):
+        """Property that determines if FSDP with QLoRA is enabled."""
+        return self.is_fsdp_enabled and self.cfg.adapter == "qlora"
+
+    @send_errors
+    def load(self) -> tuple[PreTrainedModel | PeftModelForCausalLM, PeftConfig | None]:
+        """Load and prepare the model with all configurations and patches.
+
+        Returns:
+            A tuple with the loaded model and its LoRA configuration (if applicable).
+        """
+        # Initial setup and patches
+        self.patch_manager.apply_pre_model_load_patches()
+        self._apply_pre_model_load_setup()
+
+        # Build the model
+        PLUGIN_MANAGER.pre_model_load(self.cfg)
+        self.patch_manager.apply_post_plugin_pre_model_load_patches()
+
+        skip_move_to_device = self._build_model()
+        self.patch_manager.apply_post_model_build_patches(self.model)
+
+        PLUGIN_MANAGER.post_model_build(self.cfg, self.model)
+
+        # Post-build model configuration
+        self._apply_post_model_load_setup()
+
+        # Load adapters (LoRA, etc.)
+        PLUGIN_MANAGER.pre_lora_load(self.cfg, self.model)
+        lora_config = self._load_adapters()
+        PLUGIN_MANAGER.post_lora_load(self.cfg, self.model)
+
+        # Apply remaining patches and finalize
+        self._apply_post_lora_load_setup(skip_move_to_device)
+        self.patch_manager.apply_post_model_load_patches(self.model)
+        PLUGIN_MANAGER.post_model_load(self.cfg, self.model)
+
+        if self.cfg.fp32_norms:
+            tag_model_fp32_norms(self.model, self.cfg)
+
+        return self.model, lora_config
+
+    def _apply_pre_model_load_setup(self):
+        """Apply patches and setup configurations before model loading."""
+        if self.use_parallel_config is not None:
+            self.use_parallel_config = (
+                self.cfg.fsdp_config
+                or (self.cfg.tensor_parallel_size and self.cfg.tensor_parallel_size > 1)
+                or (
+                    self.cfg.context_parallel_size
+                    and self.cfg.context_parallel_size > 1
+                )
+            )
+            if self.cfg.fsdp_config and self.cfg.fsdp_version != 2:
+                self.use_parallel_config = False
+
+        if self.use_parallel_config:
+            self._set_parallel_config()
+        self._set_auto_model_loader()
+        self._set_device_map_config()
+        if self.cfg.revision_of_model:
+            self.model_kwargs["revision"] = self.cfg.revision_of_model
+        if self.cfg.use_kernels:
+            self.model_kwargs["use_kernels"] = self.cfg.use_kernels
+            if "allow_all_kernels" not in self.model_kwargs:
+                self.model_kwargs["allow_all_kernels"] = self.cfg.use_kernels
+        self._set_quantization_config()
+        self._set_attention_config()
+        self._check_model_requirements()
+
+        # MX-quantized checkpoints carry MXTensor weights but no HF quantizer, so
+        # transformers' load-time weight re-init would crash on them; this guards it.
+        # torchao is absent on macOS/aarch64, where MX checkpoints can't exist anyway.
+        try:
+            from axolotl.utils.quantization import (
+                patch_transformers_skip_quantized_init,
+            )
+
+            patch_transformers_skip_quantized_init()
+        except ImportError:
+            pass
+
+    def _apply_post_model_load_setup(self):
+        """Configure the model after it has been loaded."""
+        # Handle PeftModel if needed
+        if (
+            isinstance(self.model, (peft.PeftModel, peft.PeftModelForCausalLM))
+            and not self.is_qlora_and_fsdp_enabled
+        ):
+            self.model = self.model.merge_and_unload()
+
+        self._configure_experts_implementation()
+        self._apply_selective_checkpointing()
+        self._apply_activation_checkpointing()
+        self._resize_token_embeddings()
+        self._reinitialize_classification_head()
+        self._adjust_model_config()
+        self._configure_embedding_dtypes()
+        self._configure_qat()
+        log_gpu_memory_usage(LOG, "Memory usage after model load", 0)
+
+    def _reinitialize_classification_head(self):
+        """Re-init an uninitialized reward / PRM classification head.
+
+        The ``score``/``classifier`` head is missing from a base-LM checkpoint, so
+        transformers allocates it with ``torch.empty`` and is then supposed to
+        initialize it. But transformers 5.8's ``_init_weights`` does
+        ``init.normal_(module.weight.float(), ...)`` — the ``.float()`` copy makes
+        this a no-op on a ``bfloat16`` head, leaving uninitialized memory: harmless
+        zeros on some allocators, NaN/inf garbage on others (→ NaN grads, 0 loss).
+        Detect that state and initialize the head ourselves.
+        """
+        if not (self.cfg.reward_model or self.cfg.process_reward_model):
+            return
+
+        head = getattr(self.model, "score", None) or getattr(
+            self.model, "classifier", None
+        )
+        if not isinstance(head, torch.nn.Linear):
+            return
+
+        weight = head.weight
+        # A freshly-initialized head is all-zero (benign) or garbage (huge/non-finite);
+        # a head loaded from a real reward checkpoint is finite and reasonably scaled.
+        looks_uninitialized = (
+            not torch.isfinite(weight).all()
+            or weight.abs().max() > 100
+            or bool((weight == 0).all())
+        )
+        if not looks_uninitialized:
+            return
+
+        std = getattr(self.model.config, "initializer_range", 0.02) or 0.02
+        with torch.no_grad():
+            weight.normal_(mean=0.0, std=std)
+            if head.bias is not None:
+                head.bias.zero_()
+        LOG.info(
+            f"Re-initialized {type(self.model).__name__} classification head "
+            f"(std={std})."
+        )
+
+    def _configure_experts_implementation(self):
+        impl = self.cfg.experts_implementation
+        if impl is None:
+            return
+
+        if impl in ("scattermoe", "sonicmoe"):
+            model_classes = {
+                type(m) for m in self.model.modules() if isinstance(m, PreTrainedModel)
+            }
+            if not any(cls._can_set_experts_implementation() for cls in model_classes):
+                LOG.warning(
+                    f"experts_implementation={impl!r} requested, but no submodule of "
+                    f"{type(self.model).__name__} uses transformers' ExpertsInterface "
+                    "(@use_experts_implementation). The kernel will NOT be applied; "
+                    "training falls back to the model's native experts path."
+                )
+
+        self.model.set_experts_implementation(impl)
+
+    def _apply_selective_checkpointing(self):
+        sac = self.cfg.selective_checkpointing
+        if not sac:
+            return
+
+        from axolotl.monkeypatch.selective_checkpointing import (
+            apply_selective_checkpointing,
+        )
+
+        sac_kwargs = sac if isinstance(sac, dict) else {}
+        apply_selective_checkpointing(
+            self.model,
+            save=sac_kwargs.get("save"),
+            save_sliding_window=bool(sac_kwargs.get("save_sliding_window")),
+            recompute_layer_types=sac_kwargs.get("recompute_layer_types"),
+            offload=bool(sac_kwargs.get("offload")),
+        )
+
+    def _apply_activation_checkpointing(self):
+        ao = self.cfg.activation_offloading
+        if ao == "hidden_states":
+            use_reentrant = (self.cfg.gradient_checkpointing_kwargs or {}).get(
+                "use_reentrant", False
+            )
+            if not use_reentrant:
+                return
+
+            from axolotl.monkeypatch.activation_offload_checkpoint import (
+                patch_hidden_states_offload,
+            )
+
+            patch_hidden_states_offload()
+            return
+        # TRL offloader is adapter-aware:
+        #   - LoRA/QLoRA: offload *replaces* recompute (pure offload is leaner/faster;
+        #     recompute would pin the offloaded tensors and balloon memory). No wrap.
+        #   - Full finetune: pure offload of every activation exceeds PCIe bandwidth
+        #     (backlogs on-GPU, OOMs at long seq), so keep recompute and offload only
+        #     the checkpoint boundaries — apply the manual wrap.
+        if ao and not self.cfg.adapter:
+            from axolotl.core.trainers.mixins.activation_checkpointing import (
+                ac_wrap_hf_model,
+            )
+
+            ac_wrap_hf_model(self.model)
+
+    def _resize_token_embeddings(self):
+        """Resize token embeddings if needed."""
+        embeddings_len = (
+            math.ceil(len(self.tokenizer) / 32) * 32
+            if self.cfg.resize_token_embeddings_to_32x
+            else len(self.tokenizer)
+        )
+        if hasattr(self.model, "get_input_embeddings") and (
+            self.model.get_input_embeddings().num_embeddings < embeddings_len
+            or (
+                self.model.get_input_embeddings().num_embeddings > embeddings_len
+                and self.cfg.shrink_embeddings
+            )
+        ):
+            resize_kwargs = {}
+            if self.cfg.mean_resizing_embeddings is not None and (
+                self.model_config.model_type != "llava"
+            ):
+                resize_kwargs["mean_resizing"] = self.cfg.mean_resizing_embeddings
+            self.model.resize_token_embeddings(embeddings_len, **resize_kwargs)
+        else:
+            self.model.tie_weights()
+
+    def _adjust_model_config(self):
+        if (
+            hasattr(self.model, "config")
+            and hasattr(self.model.config, "max_position_embeddings")
+            and self.model.config.max_position_embeddings
+            and self.cfg.sequence_len > self.model.config.max_position_embeddings
+        ):
+            LOG.warning(
+                "increasing model.config.max_position_embeddings from "
+                f"{self.model.config.max_position_embeddings} to {self.cfg.sequence_len}"
+            )
+            self.model.config.max_position_embeddings = self.cfg.sequence_len
+
+        if (
+            hasattr(self.model, "config")
+            and hasattr(self.model.config, "bos_token_id")
+            and self.model.config.bos_token_id
+            and self.model.config.bos_token_id != self.tokenizer.bos_token_id
+        ):
+            self.model.config.bos_token_id = self.tokenizer.bos_token_id
+
+        if (
+            hasattr(self.model, "config")
+            and hasattr(self.model.config, "eos_token_id")
+            and self.model.config.eos_token_id
+            and self.model.config.eos_token_id != self.tokenizer.eos_token_id
+        ):
+            self.model.config.eos_token_id = self.tokenizer.eos_token_id
+
+    def _configure_embedding_dtypes(self):
+        """Configure embedding module dtypes."""
+        # Get embedding modules
+        embedding_modules = get_linear_embedding_layers(self.cfg.model_config_type)
+
+        # Initial dtype conversion
+        if not self.is_fsdp_enabled:
+            # We don't run this during FSDP because this will leave mixed and bfloat16
+            # dtypes in the model which FSDP doesn't like
+            if self.cfg.load_in_4bit and self.cfg.embeddings_skip_upcast:
+                embedding_modules = []
+            self._convert_embedding_modules_dtype(
+                embedding_modules,
+                dist_dtype=torch.float32,
+                before_kbit_train_or_finetune=True,
+            )
+
+        # Handle DeepSpeed Zero3
+        if (
+            is_deepspeed_zero3_enabled()
+            or os.getenv("ACCELERATE_DEEPSPEED_ZERO_STAGE") == "3"
+        ):
+            self._set_z3_leaf_modules()
+
+        # Apply gradient checkpointing if needed
+        needs_fa2_dtype = self.cfg.adapter or self.is_fsdp_enabled
+        if self.cfg.adapter in ["lora", "qlora"]:
+            needs_fa2_dtype = True
+            if self.cfg.gradient_checkpointing:
+                self.model.gradient_checkpointing_enable(
+                    gradient_checkpointing_kwargs=self.cfg.gradient_checkpointing_kwargs
+                )
+
+        self._prepare_model_for_quantization()
+
+        # Convert dtypes if needed
+        should_convert = (
+            # LlamaRMSNorm layers are in fp32 after kbit_training or full finetune, so
+            # we need to convert them back to fp16/bf16 for flash-attn compatibility.
+            (
+                (needs_fa2_dtype or self.cfg.attn_needs_dtype_cast)
+                and not self.is_qlora_and_fsdp_enabled
+            )
+            or (
+                # CCE requires embedding layers to be in fp16/bf16 for backward pass
+                self.cfg.cut_cross_entropy
+            )
+        )
+
+        if should_convert:
+            LOG.info("Converting modules to %s", self.cfg.torch_dtype)
+            self._convert_embedding_modules_dtype(
+                embedding_modules=embedding_modules,
+                dist_dtype=self.cfg.torch_dtype,
+                before_kbit_train_or_finetune=False,
+            )
+
+    def _configure_qat(self):
+        """Configure QAT."""
+        if self.cfg.qat:
+            from axolotl.utils.quantization import prepare_model_for_qat
+
+            prepare_model_for_qat(
+                self.model,
+                self.cfg.qat.weight_dtype,
+                self.cfg.qat.group_size,
+                self.cfg.qat.activation_dtype,
+                self.cfg.qat.quantize_embedding,
+            )
+
+    def _load_adapters(self) -> PeftConfig | None:
+        """Load LoRA or other adapters."""
+        # Load LoRA or adapter
+        lora_config = None
+        if not self.reference_model or self.cfg.lora_model_dir:
+            # If we're not loading the reference model, then we're loading the model
+            # for training. Then, the DPO trainer doesn't want the PEFT model loaded
+            # over it, it just wants the LoRA / PEFT config.
+            if (
+                self.cfg.adapter
+                and self.cfg.rl in [RLType.DPO, RLType.IPO, RLType.KTO]
+                and not self.cfg.merge_lora
+            ):
+                _, lora_config = load_adapter(
+                    self.model,
+                    self.cfg,
+                    self.cfg.adapter,
+                    inference=False,
+                    config_only=True,
+                )
+            else:
+                self.model, lora_config = load_adapter(
+                    self.model, self.cfg, self.cfg.adapter
+                )
+
+        return lora_config
+
+    def _apply_post_lora_load_setup(self, skip_move_to_device: bool):
+        """Apply final optimizations and patches."""
+        # Place model on accelerator
+        if (
+            self.cfg.ddp
+            and not self.cfg.load_in_8bit
+            and not (self.cfg.rl and self.cfg.load_in_4bit)
+            and not skip_move_to_device
+        ):
+            self.model.to(f"{str(get_device_type())}:{self.cfg.local_rank}")
+
+        if get_device_count() > 1 and int(os.getenv("WORLD_SIZE", "1")) == 1:
+            self.model.is_parallelizable = True
+            self.model.model_parallel = True
+
+        if not any(
+            param.requires_grad
+            for _, param in self.model.named_parameters(recurse=True)
+        ):
+            LOG.warning("There are no parameters that require gradient updates")
+
+        if self.cfg.flash_optimum:
+            from optimum.bettertransformer import BetterTransformer
+
+            self.model = BetterTransformer.transform(self.model)
+
+        if self.cfg.adapter is not None:
+            log_gpu_memory_usage(LOG, "after adapters", self.model.device)
+
+        for _ in range(3):
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    def _set_parallel_config(self):
+        """Set parallelism configuration (DP, FSDP, TP, CP) in PartialState/Accelerator"""
+        parallelism_config, device_mesh = build_parallelism_config(self.cfg)
+        if parallelism_config:
+            self.parallelism_config = parallelism_config
+            self.device_mesh = device_mesh
+
+    def _set_auto_model_loader(self):
+        """Set `self.auto_model_loader`. Defaults to `transformers.AutoModelForCausalLM`
+        (set at `__init__`). When using a multimodal model, `self.auto_model_loader`
+        should be set according to the type of the model.
+        """
+        if self.cfg.is_multimodal:
+            self.auto_model_loader = MULTIMODAL_AUTO_MODEL_MAPPING.get(
+                self.model_config.model_type, AutoModelForImageTextToText
+            )
+            if isinstance(self.auto_model_loader, str):
+                self.auto_model_loader = AutoModelForImageTextToText
+
+    def _set_device_map_config(self):
+        """Setup `device_map` according to config"""
+        device_map = self.cfg.device_map
+        max_memory = self.cfg.max_memory
+
+        if self.cfg.gpu_memory_limit:
+            gpu_memory_limit = (
+                str(self.cfg.gpu_memory_limit) + "GiB"
+                if isinstance(self.cfg.gpu_memory_limit, int)
+                else self.cfg.gpu_memory_limit
+            )
+
+            max_memory = {}
+            num_device = get_device_count()
+            for i in range(num_device):
+                max_memory[i] = gpu_memory_limit
+            max_memory["cpu"] = "256GiB"  # something sufficiently large to fit anything
+
+        if max_memory is not None:
+            # Based on https://github.com/togethercomputer/OpenChatKit/blob/main/inference/bot.py
+            from accelerate import infer_auto_device_map
+
+            with init_empty_weights():
+                model_canvas = self.auto_model_loader.from_config(
+                    self.model_config,
+                    trust_remote_code=self.cfg.trust_remote_code or False,
+                )
+            model_canvas.tie_weights()
+            device_map = infer_auto_device_map(
+                model_canvas,
+                max_memory=max_memory,
+                dtype=self.cfg.torch_dtype,
+            )
+            # We can discard max_memory now as we have a device map set up
+            max_memory = None
+
+        self.model_kwargs["torch_dtype"] = self.cfg.torch_dtype
+        self.model_kwargs["dtype"] = self.cfg.torch_dtype
+
+        is_ds_zero3 = is_deepspeed_zero3_enabled()
+
+        # FSDP requires control over device placement, so don't set device_map when FSDP is enabled
+        if self.is_fsdp_enabled:
+            # For QLoRA + FSDP, we still need to set device_map to "auto" for proper initialization
+            if self.is_qlora_and_fsdp_enabled:
+                self.model_kwargs["device_map"] = {
+                    "": int(os.environ.get("LOCAL_RANK", 0))
+                }
+            # For other FSDP cases, don't set device_map at all
+        elif not is_ds_zero3:
+            self.model_kwargs["device_map"] = device_map
+
+            # quantize_moe_experts quantizes expert weights on-the-fly during loading,
+            # so the actual VRAM usage is much less than bf16 estimates.
+            # When device_map is "auto", accelerate's infer_auto_device_map computes
+            # the device map at bf16 size (before quantization), causing it to offload
+            # layers to CPU, which BnB then rejects. Force single-GPU placement to
+            # prevent this. Only applies to the non-FSDP, non-ZeRO3 path (DDP/single).
+            if getattr(self.cfg, "quantize_moe_experts", False) and device_map in (
+                "auto",
+                None,
+            ):
+                self.model_kwargs["device_map"] = {
+                    "": int(os.environ.get("LOCAL_RANK", 0))
+                }
+
+            cur_device = get_device_type()
+            if "mps" in str(cur_device):
+                self.model_kwargs["device_map"] = "mps:0"
+            elif "npu" in str(cur_device):
+                self.model_kwargs["device_map"] = "npu:0"
+
+        # TODO: can we put the reference model on it's own gpu? I think we have to move
+        # logits around to calculate loss
+        # if cfg.rl:
+        #     if torch.cuda.device_count() > 1:
+        #         if reference_model:
+        #             model_kwargs["device_map"] = "cuda:" + str(
+        #                 torch.cuda.current_device() + 1
+        #             )
+        #         else:
+        #             model_kwargs["device_map"] = "cuda:" + str(torch.cuda.current_device())
+
+    def _set_quantization_config(self):
+        """Set up quantization config (bitsandbytes, awq, gptq, etc.)"""
+
+        if self.cfg.model_quantization_config == "Mxfp4Config":
+            from transformers import Mxfp4Config
+
+            mxfp4_kwargs = {}
+            if self.cfg.model_quantization_config_kwargs:
+                mxfp4_kwargs = self.cfg.model_quantization_config_kwargs
+            self.model_kwargs["quantization_config"] = Mxfp4Config(**mxfp4_kwargs)
+
+        if self.cfg.model_quantization_config == "FineGrainedFP8Config":
+            from transformers import FineGrainedFP8Config
+
+            fp8_kwargs = {}
+            if self.cfg.model_quantization_config_kwargs:
+                fp8_kwargs = self.cfg.model_quantization_config_kwargs
+            self.model_kwargs["quantization_config"] = FineGrainedFP8Config(
+                **fp8_kwargs
+            )
+
+        if self.cfg.gptq:
+            if not hasattr(self.model_config, "quantization_config"):
+                LOG.warning(
+                    "model config does not contain quantization_config information"
+                )
+            else:
+                if self.cfg.gptq_disable_exllama is not None:
+                    self.model_config.quantization_config["disable_exllama"] = (
+                        self.cfg.gptq_disable_exllama
+                    )
+                self.model_kwargs["quantization_config"] = GPTQConfig(
+                    **self.model_config.quantization_config
+                )
+        if (
+            self.cfg.adapter in ["qlora", "lora"]
+            and hasattr(self.model_config, "quantization_config")
+            and self.model_config.quantization_config["quant_method"]
+            in ["gptq", "awq", "bitsandbytes"]
+        ):
+            if self.model_config.quantization_config["quant_method"] == "gptq":
+                self.model_kwargs["quantization_config"] = GPTQConfig(
+                    **self.model_config.quantization_config
+                )
+            elif self.model_config.quantization_config["quant_method"] == "awq":
+                self.model_kwargs["quantization_config"] = AwqConfig(
+                    **self.model_config.quantization_config
+                )
+            elif (
+                self.model_config.quantization_config["quant_method"] == "bitsandbytes"
+            ):
+                self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
+                    **self.model_config.quantization_config
+                )
+        elif self.cfg.adapter == "qlora" and self.cfg.load_in_4bit:
+            bnb_config = {
+                "load_in_4bit": True,
+                "llm_int8_threshold": 6.0,
+                "llm_int8_has_fp16_weight": False,
+                "bnb_4bit_compute_dtype": self.cfg.torch_dtype,
+                "bnb_4bit_use_double_quant": True,
+                "bnb_4bit_quant_type": "nf4",
+                "bnb_4bit_quant_storage": torch.bfloat16,
+            }
+            if self.cfg.model_config_type in [
+                "jamba",
+                "qwen2_moe",
+                "nemotron_h",
+            ] and not (self.cfg.deepspeed or self.is_fsdp_enabled):
+                # for some reason, this causes the loss to be off by an order of magnitude
+                # but deepspeed needs this still in bfloat16
+                bnb_config["bnb_4bit_quant_storage"] = torch.float32
+            if self.cfg.model_config_type == "falcon_h1":
+                # output projection cannot be quantized for Falcon-H1 models
+                bnb_config["llm_int8_skip_modules"] = ["out_proj"]
+
+            if self.cfg.bnb_config_kwargs:
+                bnb_config.update(self.cfg.bnb_config_kwargs)
+
+            self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
+                **bnb_config,
+            )
+        elif self.cfg.adapter == "lora" and self.cfg.load_in_8bit:
+            bnb_config = {
+                "load_in_8bit": True,
+            }
+            # Exclude mamba blocks from int8 quantization for jamba
+            if self.cfg.model_config_type == "jamba":
+                bnb_config["llm_int8_skip_modules"] = ["mamba"]
+            if self.cfg.model_config_type == "falcon_h1":
+                # output projection cannot be quantized for Falcon-H1 models
+                bnb_config["llm_int8_skip_modules"] = ["out_proj"]
+            self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
+                **bnb_config,
+            )
+
+    def _set_attention_config(self):
+        # fp8 replaces sdpa post-load (load as sdpa).
+        _LOAD_TIME_OVERRIDE = {"fp8": "sdpa"}
+        if self.cfg.attn_implementation:
+            hf_impl = _LOAD_TIME_OVERRIDE.get(
+                self.cfg.attn_implementation, self.cfg.attn_implementation
+            )
+            self.model_kwargs["attn_implementation"] = hf_impl
+            self.model_config._attn_implementation = hf_impl
+
+        if self.cfg.low_cpu_mem_usage:
+            self.model_kwargs["low_cpu_mem_usage"] = True
+
+    def _check_model_requirements(self):
+        if self.cfg.model_config_type in ["lfm2-vl", "lfm2"]:
+            from transformers.utils.import_utils import is_causal_conv1d_available
+
+            if is_causal_conv1d_available():
+                raise ImportError(
+                    "The 'causal-conv1d' package is installed but causes compatibility issues with LFM2 models. "
+                    "Please uninstall it by running: `pip uninstall -y causal-conv1d`"
+                )
+
+    def _configure_zero3_memory_efficient_loading(
+        self,
+    ) -> HfTrainerDeepSpeedConfig | None:
+        """
+        Set the deepspeed config to load the model into RAM first before moving to VRAM.
+
+        IMPORTANT
+        ==========
+
+        We need to return `hf_ds_cfg` as it needs to exist before model loading for zero3.
+        HfTrainerDeepSpeedConfig is a class that is used to configure the DeepSpeed training.
+        It is not passed anywhere in the model loading function, just need to exist.
+        """
+        hf_ds_cfg = None
+
+        if os.getenv("ACCELERATE_DEEPSPEED_ZERO_STAGE") == "3":
+            hf_ds_cfg = HfTrainerDeepSpeedConfig(self.cfg.deepspeed)
+            hf_ds_cfg.fill_match(
+                "train_micro_batch_size_per_gpu", self.cfg.micro_batch_size
+            )
+            hf_ds_cfg.fill_match(
+                "gradient_accumulation_steps", self.cfg.gradient_accumulation_steps
+            )
+            hf_ds_cfg.fill_match(
+                "train_batch_size",
+                int(os.getenv("WORLD_SIZE", "1"))
+                * self.cfg.micro_batch_size
+                * self.cfg.gradient_accumulation_steps,
+            )
+            if "device_map" in self.model_kwargs:
+                del self.model_kwargs["device_map"]
+
+            transformers.modeling_utils.is_deepspeed_zero3_enabled = lambda: True
+            transformers.integrations.deepspeed.is_deepspeed_zero3_enabled = lambda: (
+                True
+            )
+
+        return hf_ds_cfg
+
+    def _load_model_from_config(self, model_loader_class=None) -> PreTrainedModel:
+        """
+        Load model with random initialization using from_config.
+
+        Uses the selected loader when provided; otherwise falls back to the auto loader.
+        """
+        loader = model_loader_class or self.auto_model_loader
+        if loader in [AutoModelForCausalLM, AutoModelForImageTextToText]:
+            model = loader.from_config(
+                config=self.model_config,
+                trust_remote_code=self.cfg.trust_remote_code or False,
+            )
+        else:
+            model = loader(config=self.model_config)
+
+        return model
+
+    def _load_model_from_pretrained(self, model_loader_class=None) -> PreTrainedModel:
+        """Load model from pretrained weights."""
+        loader = model_loader_class or self.auto_model_loader
+        kwargs = {
+            "config": self.model_config,
+            "trust_remote_code": self.cfg.trust_remote_code or False,
+            **self.model_kwargs,
+        }
+        return loader.from_pretrained(self.base_model, **kwargs)
+
+    def _build_model(self) -> bool:
+        """Load model, with load strategy depending on config."""
+        skip_move_to_device = False
+
+        if self.cfg.tensor_parallel_size > 1:
+            self.model_kwargs["tp_size"] = self.cfg.tensor_parallel_size
+            self.model_kwargs["tp_plan"] = "auto"
+            self.model_kwargs["device_mesh"] = self.device_mesh
+            if "device_map" in self.model_kwargs:
+                del self.model_kwargs["device_map"]  # not compatible with `tp_plan`
+
+        if self.is_fsdp_enabled:
+            if self.cfg.fsdp_config.cpu_ram_efficient_loading:
+                skip_move_to_device = True
+                # Don't delete device_map for QLoRA + FSDP - it was set correctly in
+                # _set_device_map
+                if (
+                    "device_map" in self.model_kwargs
+                    and not self.is_qlora_and_fsdp_enabled
+                ):
+                    del self.model_kwargs["device_map"]
+            elif self.is_qlora_and_fsdp_enabled:
+                skip_move_to_device = True
+
+            if (
+                self.cfg.tensor_parallel_size <= 1
+                and self.cfg.fsdp_config.cpu_ram_efficient_loading
+                and self.cfg.fsdp_version == 2
+            ):
+                # setting device_map for TP is not supported
+                local_rank = int(os.getenv("LOCAL_RANK", "0"))
+                if local_rank == 0:
+                    self.model_kwargs["device_map"] = "cpu"
+                else:
+                    self.model_kwargs["device_map"] = "meta"
+
+        if (
+            self.is_qlora_and_fsdp_enabled
+            and self.cfg.fsdp_config.cpu_ram_efficient_loading
+            and (
+                self.cfg.model_config_type == "dbrx"
+                or self.cfg.qlora_sharded_model_loading
+            )
+        ):
+            if self.cfg.reinit_weights:
+                LOG.warning(
+                    "reinit_weights is not supported with sharded quantized loading. "
+                    "Loading from pretrained weights instead."
+                )
+            quant_storage = self.cfg.torch_dtype
+            quantization_config = getattr(
+                self.model_config, "quantization_config", None
+            )
+            quantization_config = (
+                quantization_config or self.model_kwargs["quantization_config"]
+            )
+            self.model = load_sharded_model_quant(
+                self.base_model,
+                self.model_config,
+                self.cfg,
+                quant_storage=quant_storage,
+                quantization_config=quantization_config,
+            )
+            skip_move_to_device = True
+        elif self.model_type == "MambaLMHeadModel":
+            if self.cfg.reinit_weights:
+                LOG.warning(
+                    "reinit_weights is not supported with MambaLMHeadModel. "
+                    "Loading from pretrained weights instead."
+                )
+            # FIXME this is janky at best and hacked together to make it work
+            MambaLMHeadModel = fix_mamba_attn_for_loss()
+
+            self.model_kwargs["dtype"] = self.model_kwargs["torch_dtype"]
+            self.model_kwargs["device"] = torch.cuda.current_device()
+            self.model_kwargs.pop("torch_dtype", None)
+            self.model_kwargs.pop("device_map", None)
+
+            self.model = MambaLMHeadModel.from_pretrained(
+                self.base_model,
+                **self.model_kwargs,
+            )
+        else:
+            # Please don't remove underscore binding without reading the fn docstring
+            _ = self._configure_zero3_memory_efficient_loading()
+
+            if (
+                self.model_type
+                and self.model_type != "AutoModelForCausalLM"
+                and not self.cfg.trust_remote_code
+                and not self.cfg.gptq
+            ):
+                # Use model type from transformers
+                model_loader_class = getattr(transformers, self.model_type)
+            else:
+                # Use auto model loader (handles gptq and default cases)
+                model_loader_class = self.auto_model_loader
+
+            self.model_kwargs["dtype"] = self.model_kwargs["torch_dtype"]
+            if self.cfg.reinit_weights:
+                self.model = self._load_model_from_config(model_loader_class)
+            else:
+                self.model = self._load_model_from_pretrained(model_loader_class)
+
+        if self.cfg.use_onebitllms:
+            try:
+                from onebitllms import replace_linear_with_bitnet_linear
+            except ImportError as exc:
+                raise ImportError(
+                    "The 'onebitllms' package is required for use_onebitllms. "
+                    "Install it with: `uv pip install onebitllms`"
+                ) from exc
+
+            self.model = replace_linear_with_bitnet_linear(self.model)
+
+        if is_deepspeed_zero3_enabled():
+            skip_move_to_device = True
+
+        if self.cfg.tensor_parallel_size > 1:
+            # workaround for upstream 4.54.0 not setting _tp_size or _device_mesh
+            # TODO(wing): remove once 4.54.1 is released
+            if self.model._tp_size != self.cfg.tensor_parallel_size:
+                self.model._tp_size = self.cfg.tensor_parallel_size
+                self.model._device_mesh = self.model_kwargs["device_mesh"]
+
+        if self.cfg.experimental_skip_move_to_device is not None:
+            skip_move_to_device = self.cfg.experimental_skip_move_to_device
+
+        return skip_move_to_device
+
+    def _set_z3_leaf_modules(self):
+        from deepspeed.utils import set_z3_leaf_modules
+
+        moe_type = self.cfg.model_config_type_text or self.cfg.model_config_type
+        if moe_type in MOE_ARCH_BLOCK:
+            moe_blocks = MOE_ARCH_BLOCK[moe_type]
+            moe_blocks = [moe_blocks] if isinstance(moe_blocks, str) else moe_blocks
+            set_z3_leaf_modules(
+                self.model,
+                [
+                    get_module_class_from_name(self.model, module_name)
+                    for module_name in moe_blocks
+                ],
+            )
+
+    def _prepare_model_for_quantization(self):
+        """Prepare loaded model for quantization."""
+        skip_prepare_model_for_kbit_training = False
+        if self.cfg.model_config_type == "qwen" and self.cfg.adapter == "lora":
+            # Qwen doesn't play nicely with LoRA if this is enabled
+            skip_prepare_model_for_kbit_training = True
+
+        loftq_bits = (
+            self.cfg.peft
+            and self.cfg.peft.loftq_config
+            and self.cfg.peft.loftq_config.loftq_bits
+        )
+        if self.cfg.adapter == "lora" and loftq_bits:
+            skip_prepare_model_for_kbit_training = True
+
+        if (
+            self.is_qlora_and_fsdp_enabled
+            or (self.is_fsdp_enabled and self.cfg.fsdp_config.cpu_ram_efficient_loading)
+            or is_deepspeed_zero3_enabled()
+        ):
+            # Make sure everything is in the same dtype
+            skip_prepare_model_for_kbit_training = True
+
+        if getattr(self.model, "_moe_experts_quantized", False):
+            # Parametrized expert tensors dequantize on access — would OOM.
+            skip_prepare_model_for_kbit_training = True
+
+        if (
+            not skip_prepare_model_for_kbit_training
+            and self.cfg.adapter in ["lora", "qlora"]
+            and (self.cfg.load_in_8bit or self.cfg.load_in_4bit)
+        ):
+            LOG.info("converting PEFT model w/ prepare_model_for_kbit_training")
+            self.model = prepare_model_for_kbit_training(
+                self.model, use_gradient_checkpointing=self.cfg.gradient_checkpointing
+            )
+
+    def _convert_embedding_modules_dtype(
+        self,
+        embedding_modules: list[str],
+        dist_dtype: torch.dtype,
+        before_kbit_train_or_finetune: bool,
+    ):
+        dest = {"dtype": dist_dtype}
+        if self.cfg.lora_on_cpu:
+            dest["device"] = "cpu"
+        fp32_norm_patterns = get_fp32_norm_patterns(self.cfg)
+        for name, module in self.model.named_modules():
+            if fp32_norm_patterns and _matches_norm_class(module, fp32_norm_patterns):
+                module.to(torch.float32)
+            elif "norm" in name:
+                module.to(dist_dtype)
+            if before_kbit_train_or_finetune:
+                if name.endswith(".gate"):
+                    module.to(dist_dtype)
+                if self.model_config.model_type == "btlm":
+                    # don't upcast lm_head for btlm
+                    continue
+            if any(m in name for m in embedding_modules) and hasattr(module, "weight"):
+                module.to(**dest)
diff --git a/src/axolotl/loaders/patch_manager.py b/src/axolotl/loaders/patch_manager.py
new file mode 100644
index 0000000000..2474e2edb7
--- /dev/null
+++ b/src/axolotl/loaders/patch_manager.py
@@ -0,0 +1,1091 @@
+"""Patch manager class implementation to complement `axolotl.loaders.ModelLoader`.
+
+Applies pre- and post-model load patches for various fixes and optimizations.
+"""
+
+import importlib.util
+import os
+from functools import cached_property
+
+import addict
+import torch
+import transformers
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_flash_attention_utils import is_flash_attn_available
+
+from axolotl.integrations.base import PluginManager
+from axolotl.monkeypatch.multipack import (
+    SUPPORTED_MULTIPACK_MODEL_TYPES,
+    patch_for_multipack,
+)
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+PLUGIN_MANAGER = PluginManager.get_instance()
+
+
+class PatchManager:
+    """Manages the application of patches during the model loading process."""
+
+    @staticmethod
+    def apply_pre_config_load_patches(cfg: DictDefault):
+        """
+        Apply patches that must be set up before config loading.
+        This is for patches that intercept remote code loading from HuggingFace,
+        which needs to be in place before AutoConfig.from_pretrained() is called.
+
+        Args:
+            cfg: Configuration dictionary with model and training settings.
+        """
+        if (
+            hasattr(cfg, "base_model_config")
+            and cfg.base_model_config
+            and "kimi-linear" in cfg.base_model_config.lower()
+        ):
+            from axolotl.monkeypatch.models.kimi_linear.patch_kimi_linear import (
+                patch_kimi_config,
+            )
+
+            patch_kimi_config()
+
+    @staticmethod
+    def apply_pre_tokenizer_load_patches(cfg: DictDefault):
+        """
+        Apply patches that must be set up before tokenizer loading.
+        This is for patches that intercept remote code loading from HuggingFace,
+        which needs to be in place before AutoTokenizer.from_pretrained() is called.
+
+        Args:
+            cfg: Configuration dictionary with model and training settings.
+        """
+        if (
+            hasattr(cfg, "tokenizer_config")
+            and cfg.tokenizer_config
+            and "kimi-linear" in cfg.tokenizer_config.lower()
+        ):
+            from axolotl.monkeypatch.models.kimi_linear.patch_kimi_linear import (
+                patch_kimi_tokenizer,
+            )
+
+            patch_kimi_tokenizer()
+
+    def __init__(
+        self,
+        cfg: DictDefault,
+        model_config: PretrainedConfig | addict.Dict,
+        inference: bool = False,
+    ):
+        """Initialize the `PatchManager`.
+
+        Args:
+            cfg: Configuration dictionary with model and training settings.
+            model_config: Configuration object for the model.
+            inference: Whether the model is being loaded for inference mode.
+        """
+        self.cfg = cfg
+        self.model_config = model_config
+        self.inference = inference
+
+    @cached_property
+    def has_flash_attn(self) -> bool:
+        """Check if flash attention is installed."""
+        return importlib.util.find_spec("flash_attn") is not None
+
+    def apply_pre_model_load_patches(self):
+        """Apply pre-model load patches based on config."""
+        self._deactivate_hf_async_load()
+        self._apply_torchao_patches()
+        self._apply_transformers_patches()
+        # self._apply_flex_attention_patches()
+        self._apply_flash_attention_patches()
+        self._apply_chunked_cross_entropy_patch()
+        self._apply_sageattn_patches()
+        self._apply_flash_attn_4_patches()
+        self._apply_fsdp_patches()
+        self._apply_adapter_patches()
+        # Must precede fused-RoPE patches: re-parses ``Attention.forward``
+        # via ``inspect.getsource``; the QKV regex misses on a patched body.
+        self._apply_self_attention_lora_patch()
+        self._apply_model_specific_patches()
+        self._apply_fp8_patches()
+        self._apply_flash_attention_peft_patches()
+        self._apply_gradient_checkpointing_patches()
+        self._patch_attention()
+        self._apply_multipack_patches()
+        self._apply_sdpa_varlen_patch()
+        self._apply_large_head_attention_patch()
+        self._patch_loss_llama()
+        self._patch_llama_derived_model()
+        self._apply_mistral_cross_entropy_patch()
+        self._apply_fsdp2_bnb_patches()
+        self._apply_patch_deepspeed_zero3()
+        self._apply_voxtral_patches()
+        self._apply_apertus_patches()
+        self._apply_trl_vllm_patches()
+        self._apply_trl_trainer_utils_patches()
+
+    def apply_post_plugin_pre_model_load_patches(self):
+        """Apply post plugin-pre_model_load load patches based on config."""
+        self._apply_tiled_mlp(self.cfg.model_config_type)
+        self._apply_moe_expert_quantization_patch()
+
+    @staticmethod
+    def _apply_torchao_patches():
+        from axolotl.monkeypatch.torchao_optim import patch_torchao_optim_state_8bit
+
+        patch_torchao_optim_state_8bit()
+
+    def _apply_transformers_patches(self):
+        from axolotl.monkeypatch.transformers.trainer_loss_calc import (
+            patch_evaluation_loop,
+            patch_maybe_log_save_evaluate,
+        )
+
+        patch_evaluation_loop()
+        patch_maybe_log_save_evaluate()
+
+    def apply_post_model_build_patches(self, model: PreTrainedModel):
+        """Apply patches right after model build, before post-load setup."""
+        if self.cfg.model_config_type == "nemotron_h":
+            # Must run after model build because NemotronHForCausalLM.__init__
+            # calls register_nemotron_h_conversion_mapping() with overwrite=True,
+            # which would clobber any earlier fix.
+            self._fix_nemotron_h_conversion_mapping()
+
+        # Gemma 4 hybrid attention runs here in post-build (NOT post-load):
+        # the per-layer ``self_attn.config._attn_implementation="sdpa"``
+        # override needs to walk the raw model tree, which is broken by
+        # the post-load PEFT wrapping. The accompanying
+        # ``patch_gemma4_hybrid_mask`` monkey-patch is module-level and
+        # installation-time-independent, so both halves of the fix live
+        # cleanly in the same call even though one is instance-scoped
+        # and the other is module-scoped.
+        self._apply_gemma_hybrid_attention(model)
+        self._apply_gemma4_loss_kwargs()
+        self._finalize_moe_expert_quantization(model)
+
+    def apply_post_model_load_patches(self, model: PreTrainedModel):
+        """Apply patches that require the model instance."""
+        self._apply_llama_flash_attn_patches(model)
+        self._apply_lora_kernel_patch(model)
+        self._apply_scaling_softmax_patch(model)
+        self._apply_fp8_attention_patches(model)
+        self._apply_tiled_mlp_post_load(model)
+
+    def _apply_gemma4_loss_kwargs(self):
+        # Flip accepts_loss_kwargs True so the Trainer normalizes loss by
+        # num_items_in_batch under grad accumulation (must run before trainer init).
+        if self.cfg.model_config_type not in ("gemma4", "gemma4_unified"):
+            return
+        from axolotl.monkeypatch.gemma4_loss_kwargs import (
+            patch_gemma4_accepts_loss_kwargs,
+        )
+
+        patch_gemma4_accepts_loss_kwargs()
+
+    def _apply_gemma_hybrid_attention(self, model: PreTrainedModel):
+        """Apply hybrid attention: FA2 for sliding window layers, SDPA for global layers.
+
+        Gemma 4 has global (full_attention) layers with head_dim=512
+        which exceeds flash attention's supported size. This patch loads the model
+        with flash_attention_2 for the sliding window layers (head_dim=256), then
+        gives each global layer a shallow-copied config with _attn_implementation="sdpa".
+
+        We also install :func:`axolotl.monkeypatch.gemma4_hybrid_mask.patch_gemma4_hybrid_mask`
+        which fixes the corresponding mask construction inside
+        ``Gemma4TextModel.forward``. Without it, the per-layer SDPA config
+        override is not enough — the forward still builds a 2D FA2-format mask
+        at the model level and the SDPA layers crash at long context lengths
+        with ``RuntimeError: The expanded size of the tensor ... must match``.
+        """
+        if not self.cfg.gemma4_hybrid_attn_impl:
+            return
+
+        import copy
+
+        from axolotl.monkeypatch.attention.large_head import (
+            resolve_large_head_policy,
+            set_large_head_packed,
+            set_large_head_policy,
+        )
+        from axolotl.monkeypatch.gemma4_hybrid_mask import (
+            GLOBAL_PACKED_SDPA,
+            patch_gemma4_hybrid_mask,
+        )
+
+        patch_gemma4_hybrid_mask()
+        # Gemma-4 global layers reuse the generic large-head router. Default policy 'sdpa' (flash is
+        # opt-in via large_head_attention / the deprecated flash_attn_d512), preserving prior default.
+        set_large_head_policy(resolve_large_head_policy(self.cfg))
+        set_large_head_packed(bool(self.cfg.sample_packing))
+
+        # Navigate to the module that has 'layers' - varies by model structure:
+        # Gemma4ForConditionalGeneration -> .model (Gemma4Model) -> .language_model (Gemma4TextModel) -> .layers
+        # Gemma4ForCausalLM -> .model (Gemma4TextModel) -> .layers
+        layers = None
+        config_source = None
+        for candidate in [model, getattr(model, "model", None)]:
+            if candidate is None:
+                continue
+            # Check direct layers
+            if hasattr(candidate, "layers"):
+                layers = candidate.layers
+                config_source = candidate
+                break
+            # Check language_model.layers (multimodal wrapper)
+            lang_model = getattr(candidate, "language_model", None)
+            if lang_model is not None and hasattr(lang_model, "layers"):
+                layers = lang_model.layers
+                config_source = lang_model
+                break
+
+        if layers is None:
+            LOG.warning(
+                "gemma4_hybrid_attn_impl: could not find decoder layers in model, skipping"
+            )
+            return
+
+        config = getattr(config_source, "config", self.model_config)
+        layer_types = getattr(config, "layer_types", None)
+        if layer_types is None:
+            LOG.warning(
+                "gemma4_hybrid_attn_impl: model config has no 'layer_types', skipping. "
+                "This feature requires a model with mixed sliding/global attention layers."
+            )
+            return
+
+        patched_count = 0
+        for layer_idx, layer in enumerate(layers):
+            if layer_types[layer_idx] != "sliding_attention":
+                # Global / full_attention layer (head_dim=512, FA2 can't serve it). Use the
+                # packing-aware SDPA impl: it rebuilds the block-diagonal mask from position_ids so
+                # the layer respects document boundaries under sample packing (plain "sdpa" gets a
+                # None mask here and would attend across packed documents).
+                attn_module = getattr(layer, "self_attn", None)
+                if attn_module is not None and hasattr(attn_module, "config"):
+                    sdpa_config = copy.copy(attn_module.config)
+                    sdpa_config._attn_implementation = GLOBAL_PACKED_SDPA
+                    attn_module.config = sdpa_config
+                    patched_count += 1
+
+        LOG.info(
+            "gemma4_hybrid_attn_impl: patched %d global layers to use packing-aware SDPA "
+            "(remaining %d sliding layers use flash_attention_2)",
+            patched_count,
+            len(layers) - patched_count,
+        )
+
+    def _apply_flash_attention_patches(self):
+        """Apply patches related to Flash Attention."""
+        if self.cfg.attn_implementation == "xformers":
+            from axolotl.monkeypatch.attention import register_xformers_attn
+
+            register_xformers_attn()
+
+            if self.cfg.sample_packing:
+                # Also patch FA2 slot for legacy code paths that use it directly
+                from axolotl.monkeypatch.attention import patch_xformers_attn_over_fa2
+
+                patch_xformers_attn_over_fa2()
+
+        if self.cfg.attn_implementation == "sage":
+            from axolotl.monkeypatch.attention import register_sage_attn
+
+            register_sage_attn()
+
+    def _apply_fp8_attention_patches(self, model):
+        """Apply FP8 low-precision attention via torchao."""
+        if self.cfg.attn_implementation == "fp8":
+            from axolotl.monkeypatch.attention.fp8_attn import patch_fp8_attention
+
+            patch_fp8_attention(model)
+
+    def _apply_chunked_cross_entropy_patch(self):
+        if self.cfg.chunked_cross_entropy:
+            from axolotl.monkeypatch.loss.chunked import patch_chunked_ce_loss_fn
+
+            if self.cfg.chunked_cross_entropy_num_chunks:
+                patch_chunked_ce_loss_fn(self.cfg.chunked_cross_entropy_num_chunks)
+            else:
+                patch_chunked_ce_loss_fn()
+
+    def _apply_fsdp_patches(self):
+        """Apply patches for FSDP configurations."""
+        if self.cfg.fsdp_config:
+            from axolotl.monkeypatch.accelerate.fsdp2 import (
+                patch_initialize_missing_keys_for_fsdp,
+            )
+
+            patch_initialize_missing_keys_for_fsdp()
+
+            # Only for adapter (frozen-base) runs: this patch leaves non-rank-0 base params on meta
+            # (FSDP broadcasts rank-0's weights into them), which avoids world_size× CPU
+            # materialization of large unrecognized-quantizer (NVFP4-modelopt) checkpoints. Those are
+            # always trained with a frozen base, so no base optimizer state exists. For a FULL
+            # fine-tune the base params DO carry optimizer state, and leaving them on meta deadlocks
+            # the FSDP2 optimizer-state all-gather at checkpoint save (rank-0 real DTensors vs
+            # non-rank-0 meta) — so fall back to the stock materialize-to-cpu path there.
+            if self.cfg.fsdp_config.cpu_ram_efficient_loading and self.cfg.adapter:
+                from axolotl.monkeypatch.accelerate.fsdp2 import (
+                    patch_move_missing_keys_meta_for_fsdp,
+                )
+
+                patch_move_missing_keys_meta_for_fsdp()
+
+        if self.cfg.context_parallel_size > 1 or (
+            self.cfg.fsdp_config and str(self.cfg.fsdp_version) == "2"
+        ):
+            from axolotl.monkeypatch.accelerate.parallelism_config import (
+                patch_parallelism_config,
+            )
+
+            patch_parallelism_config()
+        if self.cfg.fsdp_config and str(self.cfg.fsdp_version) == "2":
+            from axolotl.monkeypatch.accelerate.float8_fsdp import patch_float8_fsdp
+            from axolotl.monkeypatch.accelerate.fsdp2 import (
+                patch_accelerate_fsdp2,
+                patch_tied_keys_for_meta_device,
+            )
+
+            patch_accelerate_fsdp2()
+            # FSDP2 sharding for any torchao Float8Tensor weights (no-op without torchao)
+            patch_float8_fsdp()
+            if self.cfg.fsdp_config.cpu_ram_efficient_loading:
+                patch_tied_keys_for_meta_device()
+            if self.cfg.rl:
+                from axolotl.monkeypatch.trainer.trl import patch_trl_prepare_fsdp2
+
+                patch_trl_prepare_fsdp2()
+
+    def _apply_adapter_patches(self):
+        """Apply patches for adapter configurations."""
+        if self.cfg.adapter and self.cfg.embeddings_skip_upcast:
+            from axolotl.monkeypatch.peft.utils import patch_peft_prep_code
+
+            patch_peft_prep_code()
+
+    def _apply_flex_attention_patches(self):
+        """Apply patches for flexible attention."""
+        if self.cfg.attn_implementation == "flex_attention":
+            from axolotl.monkeypatch.attention.flex_attn import (
+                patch_flex_wrapper,
+            )
+
+            flex_attn_compile_kwargs = self.cfg.flex_attn_compile_kwargs or {}
+            patch_flex_wrapper(**flex_attn_compile_kwargs)
+
+    def _apply_sageattn_patches(self):
+        """Apply patches for SageAttention."""
+        if self.cfg.attn_implementation == "sage":
+            from axolotl.monkeypatch.attention.sage_attn import patch_sageattn
+
+            patch_sageattn()
+
+    def _apply_flash_attn_4_patches(self):
+        """Auto-apply FA4 when flash_attention is enabled and FA4 is available on SM90+."""
+        if not self.cfg.attn_uses_flash_lib:
+            return
+
+        from axolotl.monkeypatch.attention.flash_attn_4 import patch_flash_attn_4
+
+        patch_flash_attn_4(self.model_config)
+
+    _FUSED_ATTN_KERNEL_SUPPORTED = (
+        "qwen3",
+        "qwen3_moe",
+        "qwen3_vl",
+        "qwen3_vl_text",
+        "qwen3_5",
+        "qwen3_5_text",
+        "qwen3_5_moe",
+        "qwen3_5_moe_text",
+        "gemma4",
+        "gemma4_text",
+        "gemma4_unified",
+        "gemma4_unified_text",
+    )
+
+    @staticmethod
+    def _warn_if_fused_attn_unsupported(cfg):
+        """Warn when ``fused_attn_kernel`` targets an unsupported
+        ``model_config_type`` (derived post-schema by ``normalize_config()``)."""
+        if not getattr(cfg, "fused_attn_kernel", False):
+            return
+        mct = getattr(cfg, "model_config_type", None)
+        if mct and mct not in PatchManager._FUSED_ATTN_KERNEL_SUPPORTED:
+            LOG.warning(
+                "`fused_attn_kernel: true` is set but model_config_type=%r is not "
+                "in the supported set %s. The flag is a silent no-op for this "
+                "model. Remove the flag or use one of the supported model families.",
+                mct,
+                sorted(PatchManager._FUSED_ATTN_KERNEL_SUPPORTED),
+            )
+
+    def _apply_model_specific_patches(self):
+        """Apply patches specific to model architectures."""
+        self._warn_if_fused_attn_unsupported(self.cfg)
+
+        if getattr(self.cfg, "use_kernels", None):
+            from axolotl.monkeypatch.kernelize_fixes import patch_kernelize_fixes
+
+            patch_kernelize_fixes()
+
+        if (
+            self.cfg.model_config_type == "llama4"
+            and self.cfg.llama4_linearized_experts
+        ):
+            from axolotl.monkeypatch.models.llama4.modeling import (
+                patch_llama4_linearized_modeling,
+            )
+
+            patch_llama4_linearized_modeling()
+
+        if self.cfg.model_config_type == "kimi_linear":
+            from axolotl.monkeypatch.models.kimi_linear.patch_kimi_linear import (
+                patch_kimi_model,
+            )
+
+            patch_kimi_model()
+
+        ssm_hybrid_patch_needed = (
+            self.cfg.sample_packing or self.cfg.context_parallel_size > 1
+        )
+
+        if self.cfg.model_config_type == "nemotron_h" and ssm_hybrid_patch_needed:
+            from transformers.models.nemotron_h.modeling_nemotron_h import (
+                NemotronHPreTrainedModel,
+            )
+
+            from axolotl.monkeypatch.models.nemotron_h.modeling import (
+                patch_nemotron_h_modeling_packing,
+            )
+
+            patch_nemotron_h_modeling_packing()
+            # supports_gradient_checkpointing is only enabled after
+            # patch_nemotron_h_modeling_packing() installs the GC-compatible
+            # NemotronHBlock.forward. Without the patch, upstream marks this
+            # False because the original block forward is not GC-safe.
+            NemotronHPreTrainedModel.supports_gradient_checkpointing = True
+
+        if self.cfg.model_config_type == "falcon_h1" and ssm_hybrid_patch_needed:
+            from axolotl.monkeypatch.models.falcon_h1.modeling import (
+                patch_falcon_h1_modeling_packing,
+            )
+
+            patch_falcon_h1_modeling_packing()
+
+        if self.cfg.model_config_type == "granitemoehybrid" and ssm_hybrid_patch_needed:
+            from axolotl.monkeypatch.models.granitemoehybrid.modeling import (
+                patch_granitemoehybrid_modeling_packing,
+            )
+
+            patch_granitemoehybrid_modeling_packing()
+
+        # Patches requiring CUDA
+        if torch.cuda.is_available():
+            if self.cfg.model_config_type == "qwen3_next" and self.cfg.sample_packing:
+                from axolotl.monkeypatch.models.qwen3_next.modeling import (
+                    patch_qwen3_next_modeling_packing,
+                )
+
+                patch_qwen3_next_modeling_packing()
+
+            if self.cfg.model_config_type == "qwen3_5" and self.cfg.sample_packing:
+                from axolotl.monkeypatch.models.qwen3_5.modeling import (
+                    patch_qwen3_5_modeling_packing,
+                )
+
+                patch_qwen3_5_modeling_packing()
+
+            if self.cfg.model_config_type == "qwen3_5_moe" and self.cfg.sample_packing:
+                from axolotl.monkeypatch.models.qwen3_5.modeling import (
+                    patch_qwen3_5_moe_modeling_packing,
+                )
+
+                patch_qwen3_5_moe_modeling_packing()
+
+            if (
+                self.cfg.model_config_type in ["qwen3_5", "qwen3_5_moe"]
+                and self.cfg.is_multimodal
+                and self.cfg.attn_uses_flash_lib
+            ):
+                from axolotl.monkeypatch.models.qwen3_5.modeling import (
+                    patch_qwen3_5_vlm_flash_attention,
+                )
+
+                patch_qwen3_5_vlm_flash_attention()
+
+            if self.cfg.model_config_type in (
+                "gemma4",
+                "gemma4_text",
+                "gemma4_unified",
+                "gemma4_unified_text",
+            ):
+                # Shared-KV side channel when activation checkpointing (PR #3611).
+                fsdp_cfg = self.cfg.fsdp_config
+                needs_shared_kv_workaround = (not self.inference) and bool(
+                    self.cfg.gradient_checkpointing
+                    or self.cfg.activation_offloading
+                    or (fsdp_cfg is not None and fsdp_cfg.activation_checkpointing)
+                )
+                if self.cfg.model_config_type in (
+                    "gemma4_unified",
+                    "gemma4_unified_text",
+                ):
+                    from axolotl.monkeypatch.models.gemma4_unified.fused_attn import (
+                        patch_gemma4_unified_fused_attn as patch_fused_attn,
+                    )
+                else:
+                    from axolotl.monkeypatch.models.gemma4.fused_attn import (
+                        patch_gemma4_fused_attn as patch_fused_attn,
+                    )
+                patch_fused_attn(
+                    install_shared_kv_workaround=needs_shared_kv_workaround
+                )
+
+            if self.cfg.fused_attn_kernel and self.cfg.model_config_type == "qwen3":
+                from axolotl.monkeypatch.models.qwen3.fused_attn import (
+                    patch_qwen3_fused_attn,
+                )
+
+                patch_qwen3_fused_attn()
+
+            if self.cfg.fused_attn_kernel and self.cfg.model_config_type == "qwen3_moe":
+                from axolotl.monkeypatch.models.qwen3_moe.fused_attn import (
+                    patch_qwen3_moe_fused_attn,
+                )
+
+                patch_qwen3_moe_fused_attn()
+
+            if self.cfg.fused_attn_kernel and self.cfg.model_config_type in (
+                "qwen3_vl",
+                "qwen3_vl_text",
+            ):
+                from axolotl.monkeypatch.models.qwen3_vl.fused_attn import (
+                    patch_qwen3_vl_fused_attn,
+                )
+
+                patch_qwen3_vl_fused_attn()
+
+            if self.cfg.fused_attn_kernel and self.cfg.model_config_type in (
+                "qwen3_5",
+                "qwen3_5_text",
+            ):
+                from axolotl.monkeypatch.models.qwen3_5.fused_attn import (
+                    patch_qwen3_5_fused_attn,
+                )
+
+                patch_qwen3_5_fused_attn()
+
+            if self.cfg.fused_attn_kernel and self.cfg.model_config_type in (
+                "qwen3_5_moe",
+                "qwen3_5_moe_text",
+            ):
+                from axolotl.monkeypatch.models.qwen3_5_moe.fused_attn import (
+                    patch_qwen3_5_moe_fused_attn,
+                )
+
+                patch_qwen3_5_moe_fused_attn()
+
+    @staticmethod
+    def _fix_nemotron_h_conversion_mapping():
+        """Remove the spurious embedding→embeddings WeightRenaming from the
+        nemotron_h checkpoint conversion mapping.
+
+        The nvidia Hub model registers:
+            WeightRenaming("embedding.weight", "embeddings.weight")
+        to handle a legacy checkpoint variant. Its reverse (applied on save)
+        converts ``embeddings`` back to ``embedding``, which silently renames
+        ``backbone.embeddings.weight`` → ``backbone.embedding.weight`` when
+        merging LoRA adapters back into the base model.
+        """
+        try:
+            from transformers.conversion_mapping import (
+                WeightRenaming,
+                get_checkpoint_conversion_mapping,
+                register_checkpoint_conversion_mapping,
+            )
+        except ImportError:
+            return
+
+        mapping = get_checkpoint_conversion_mapping("nemotron_h")
+        if mapping is None:
+            return
+
+        filtered = [
+            entry
+            for entry in mapping
+            if not (
+                isinstance(entry, WeightRenaming)
+                and entry.source_patterns == ["embedding.weight"]
+                and entry.target_patterns == ["embeddings.weight"]
+            )
+        ]
+        if len(filtered) != len(mapping):
+            register_checkpoint_conversion_mapping(
+                "nemotron_h", filtered, overwrite=True
+            )
+            LOG.info(
+                "Removed embedding→embeddings WeightRenaming from nemotron_h "
+                "checkpoint conversion mapping"
+            )
+
+    def _apply_fp8_patches(self):
+        """Apply patches for FP8 support."""
+        if self.cfg.fp8:
+            from axolotl.monkeypatch.trainer_accelerator_args import (
+                patch_create_accelerate_code_for_fp8,
+            )
+
+            patch_create_accelerate_code_for_fp8(
+                self.cfg.fp8_enable_fsdp_float8_all_gather
+            )
+
+    def _apply_flash_attention_peft_patches(self):
+        """Apply patches for Flash Attention with PEFT."""
+        if self.cfg.adapter:
+            from axolotl.monkeypatch.transformers_fa_utils import (
+                patch_fa_peft_integration,
+            )
+
+            patch_fa_peft_integration()
+
+    def _apply_gradient_checkpointing_patches(self):
+        """Apply patches for gradient checkpointing."""
+        if (
+            self.cfg.gradient_checkpointing
+            and self.cfg.activation_offloading == "legacy"
+        ):
+            from axolotl.monkeypatch.gradient_checkpointing import (
+                hf_grad_checkpoint_offload_wrapper,
+            )
+
+            transformers.modeling_utils.checkpoint = hf_grad_checkpoint_offload_wrapper
+        elif (
+            self.cfg.gradient_checkpointing
+            and self.cfg.activation_offloading == "offload_disk"
+        ):
+            from axolotl.monkeypatch.gradient_checkpointing import (
+                hf_grad_checkpoint_disk_offload_wrapper,
+            )
+
+            transformers.modeling_utils.checkpoint = (
+                hf_grad_checkpoint_disk_offload_wrapper
+            )
+
+    def _apply_mistral_cross_entropy_patch(self):
+        """Apply Mistral cross entropy patch if configured."""
+        if (
+            self.cfg.model_config_type == "mistral"
+            and self.cfg.flash_attn_cross_entropy_loss
+        ):
+            from axolotl.monkeypatch.mistral_attn_hijack_flash import (
+                patch_mistral_cross_entropy,
+            )
+
+            patch_mistral_cross_entropy()
+
+    def _apply_self_attention_lora_patch(self):
+        """Apply self-attention LoRA patches if configured."""
+        if self.cfg.lora_qkv_kernel or self.cfg.lora_o_kernel:
+            from axolotl.monkeypatch.lora_kernels import patch_self_attn_lora
+
+            patch_self_attn_lora(self.cfg)
+
+    def _apply_large_head_attention_patch(self):
+        """Generic head_dim>256 capability for plain SDPA models. Gemma-4's hybrid path routes its
+        globals through its own impl, so skip the generic sdpa wrapper there to avoid double-wiring."""
+        from axolotl.monkeypatch.attention.large_head import (
+            resolve_large_head_policy,
+            set_large_head_packed,
+            set_large_head_policy,
+            unpatch_sdpa_large_head,
+        )
+
+        policy = resolve_large_head_policy(self.cfg)
+        # Always (re)set the policy global from this run's config so a long-lived process can't
+        # inherit a previous run's stale auto/triton_flash policy on an sdpa run.
+        set_large_head_policy(policy)
+        set_large_head_packed(bool(self.cfg.sample_packing))
+        if policy == "sdpa" or self.cfg.gemma4_hybrid_attn_impl:
+            unpatch_sdpa_large_head()
+            return
+        from axolotl.monkeypatch.attention.large_head import patch_sdpa_large_head
+
+        patch_sdpa_large_head(policy)
+
+    def _apply_sdpa_varlen_patch(self):
+        """Route packed-row SDPA through cu_seqlens ``varlen_attn`` (no 4D mask).
+
+        Auto-enabled for ``sdpa`` + ``sample_packing`` when the varlen kernel can serve
+        the model (torch >= 2.10, head_dim <= 256, no sliding window). Opt in/out
+        explicitly with ``sdpa_varlen``. When varlen is unavailable/unsuitable, stock SDPA
+        stays and packing is still isolated via the dropped-mask block-diagonal path.
+        """
+        # False -> explicit opt-out; True -> explicit opt-in; None -> auto for sdpa packing.
+        if self.cfg.sdpa_varlen is False:
+            return
+        explicit = self.cfg.sdpa_varlen is True
+        auto = (
+            self.cfg.sdpa_varlen is None
+            and self.cfg.attn_implementation == "sdpa"
+            and self.cfg.sample_packing
+        )
+        if not (explicit or auto):
+            return
+
+        from axolotl.monkeypatch.attention.sdpa_varlen import (
+            _VARLEN_MAX_HEAD_DIM,
+            patch_sdpa_varlen,
+            varlen_available,
+        )
+
+        if not varlen_available():
+            return  # torch < 2.10; block-diagonal packing path is correct
+
+        def _attr(name):
+            mc = self.model_config
+            return mc.get(name) if isinstance(mc, dict) else getattr(mc, name, None)
+
+        head_dim = _attr("head_dim")
+        if not head_dim and _attr("hidden_size") and _attr("num_attention_heads"):
+            head_dim = _attr("hidden_size") // _attr("num_attention_heads")
+        sliding = _attr("sliding_window")
+        layer_types = _attr("layer_types")
+        uses_sliding = bool(sliding) and (
+            any(lt == "sliding_attention" for lt in layer_types)
+            if layer_types
+            else True
+        )
+
+        if (head_dim and head_dim > _VARLEN_MAX_HEAD_DIM) or uses_sliding:
+            if explicit:
+                LOG.info(
+                    "sdpa_varlen: model has head_dim > %d or a sliding window; keeping "
+                    "stock SDPA (packing still isolated via the block-diagonal mask).",
+                    _VARLEN_MAX_HEAD_DIM,
+                )
+            return
+
+        patch_sdpa_varlen()
+
+    def _apply_multipack_patches(self):
+        """Apply multipack patches if necessary."""
+        if (
+            self.cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES
+            and self.cfg.attn_supports_packing
+            and self.cfg.sample_packing
+        ):
+            # Get automap config if it exists
+            auto_map_config = None
+            if isinstance(self.model_config, dict) and "auto_map" in self.model_config:
+                auto_map_config = self.model_config["auto_map"]
+            elif hasattr(self.model_config, "auto_map"):
+                auto_map_config = self.model_config.auto_map
+
+            # Determine if the model has remote code
+            if auto_map_config is not None:
+                has_remote_code = "AutoModelForCausalLM" in auto_map_config
+            else:
+                has_remote_code = False
+
+            if has_remote_code and self.cfg.trust_remote_code is not None:
+                # If explicitly set in YAML, prefer that
+                has_remote_code = self.cfg.trust_remote_code
+
+            patch_for_multipack(
+                self.cfg.model_config_type,
+                model_name=self.cfg.base_model,
+                has_remote_code=has_remote_code,
+            )
+
+        if self.cfg.sample_packing:
+            from axolotl.monkeypatch.data.batch_dataset_fetcher import (
+                apply_multipack_dataloader_patch,
+            )
+
+            LOG.info("Applying multipack dataloader patch for sample packing...")
+            apply_multipack_dataloader_patch()
+
+    def _apply_fsdp2_bnb_patches(self):
+        """Apply FSDP2 BNB patches."""
+        if (
+            self.cfg.fsdp_config
+            and str(self.cfg.fsdp_version) == "2"
+            and (self.cfg.load_in_4bit or self.cfg.load_in_8bit)
+        ):
+            from axolotl.monkeypatch.fsdp2_qlora import (
+                apply_init_dtype_attrs_patch,
+                apply_init_sharded_param_patch,
+                apply_init_unsharded_param_patch,
+                apply_linear8bitlt_save_patch,
+            )
+
+            apply_init_sharded_param_patch()
+            apply_init_unsharded_param_patch()
+            apply_init_dtype_attrs_patch()
+            if self.cfg.load_in_8bit:
+                apply_linear8bitlt_save_patch()
+
+    def _deactivate_hf_async_load(self):
+        """Load weights synchronously so they can be converted and not OOM."""
+        if self.cfg.load_in_4bit or self.cfg.load_in_8bit:
+            os.environ["HF_DEACTIVATE_ASYNC_LOAD"] = "1"
+
+    def _apply_moe_expert_quantization_patch(self):
+        """Patch transformers weight loading and PEFT for MoE expert quantization."""
+        has_target_params = bool(getattr(self.cfg, "lora_target_parameters", None))
+
+        if not self.cfg.quantize_moe_experts and not has_target_params:
+            return
+
+        from axolotl.monkeypatch.moe_quant import (
+            patch_peft_target_parameters_matching,
+        )
+
+        if self.cfg.quantize_moe_experts:
+            from axolotl.monkeypatch.moe_quant import patch_moe_quantization_on_load
+
+            patch_moe_quantization_on_load(self.cfg)
+
+        patch_peft_target_parameters_matching()
+
+    def _finalize_moe_expert_quantization(self, model: PreTrainedModel):
+        """Log quantization results and set model flag for downstream use."""
+        import torch
+
+        model._moe_experts_quantized = False
+        if self.cfg.quantize_moe_experts:
+            from axolotl.monkeypatch.moe_quant import get_moe_quantized_count
+
+            count = get_moe_quantized_count()
+            if count > 0:
+                import gc
+
+                model._moe_experts_quantized = True
+                LOG.info(
+                    "Quantized %d MoE expert parameter(s) to %s during model loading",
+                    count,
+                    "4-bit" if self.cfg.load_in_4bit else "8-bit",
+                )
+                gc.collect()
+                torch.cuda.empty_cache()
+
+    def _apply_tiled_mlp(self, model_type: str):
+        if self.cfg.tiled_mlp:
+            from axolotl.monkeypatch.tiled_mlp import (
+                patch_tiled_mlp,
+            )
+
+            patch_tiled_mlp(
+                model_type,
+                use_original_mlp=self.cfg.tiled_mlp_use_original_mlp,
+                cfg_num_shards=self.cfg.tiled_mlp_num_shards,
+                use_scattermoe=bool(self.cfg.use_scattermoe),
+            )
+
+    def _apply_tiled_mlp_post_load(self, model):
+        """Re-wrap MoE block instances after kernels have installed their forward.
+
+        Needed only when scattermoe-lora is active — ``model.kernelize()``
+        binds ``HFScatterMoEGatedMLP.forward`` per instance, which shadows
+        the class-level tiled patch. See
+        :func:`axolotl.monkeypatch.tiled_mlp.patch_tiled_mlp_moe_instances`.
+        """
+        if not (self.cfg.tiled_mlp and self.cfg.use_scattermoe):
+            return
+        from axolotl.monkeypatch.tiled_mlp import patch_tiled_mlp_moe_instances
+
+        patch_tiled_mlp_moe_instances(
+            model,
+            self.cfg.model_config_type,
+            cfg_num_shards=self.cfg.tiled_mlp_num_shards,
+        )
+
+    def _apply_voxtral_patches(self):
+        """Apply patches for Voxtral model."""
+        if self.cfg.model_config_type == "voxtral":
+            from axolotl.monkeypatch.models.voxtral.modeling import (
+                patch_voxtral_conditional_generation_forward,
+            )
+
+            patch_voxtral_conditional_generation_forward()
+
+    def _patch_attention(self):
+        """Apply attention-specific patches based on model type."""
+        if not (
+            self.cfg.attn_uses_flash_lib and hasattr(self.model_config, "model_type")
+        ):
+            return
+
+        if self.model_config.model_type == "btlm":
+            from axolotl.monkeypatch.btlm_attn_hijack_flash import (
+                replace_btlm_attn_with_flash_attn,
+            )
+
+            replace_btlm_attn_with_flash_attn(self.cfg.base_model)
+
+        if self.model_config.model_type == "stablelm_epoch" and self.cfg.sample_packing:
+            from axolotl.monkeypatch.stablelm_attn_hijack_flash import (
+                replace_stablelm_attn_with_flash_attn,
+            )
+
+            replace_stablelm_attn_with_flash_attn(self.cfg.base_model)
+
+        if self.model_config.model_type in ("mistral3", "llava"):
+            from axolotl.monkeypatch.models.pixtral.modeling_flash_attention_utils import (
+                apply_patch_is_packed_sequence,
+            )
+
+            apply_patch_is_packed_sequence()
+
+    def _patch_loss_llama(self):
+        """Patch loss functions and other optimizations for LLaMA models."""
+        if not self.cfg.is_llama_derived_model:
+            return
+
+        if self.cfg.flash_attn_cross_entropy and self.has_flash_attn:
+            from axolotl.monkeypatch.llama_attn_hijack_flash import (
+                patch_fa_llama_cross_entropy,
+            )
+
+            patch_fa_llama_cross_entropy()
+
+    def _patch_llama_xformers_attention(self):
+        """Apply xformers attention patches for LLaMA models."""
+        from axolotl.monkeypatch.llama_attn_hijack_xformers import (
+            hijack_llama_attention,
+        )
+
+        LOG.info("Patching with xformers attention...")
+        hijack_llama_attention()
+
+    def _patch_llama_derived_model(self):
+        """Modify all llama derived models in one block."""
+        if (
+            self.cfg.is_llama_derived_model
+            and self.cfg.attn_implementation == "xformers"
+            and not (
+                self.cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES
+                and self.cfg.attn_supports_packing
+                and self.cfg.sample_packing
+            )
+        ):
+            self._patch_llama_xformers_attention()
+
+    def _apply_llama_flash_attn_patches(self, model):
+        """Apply LLaMA-specific flash attention patches."""
+
+        if (
+            self.model_config.model_type
+            in ["llama", "llama4", "ernie4_5", "ernie4_5_moe"]
+            and not self.cfg.trust_remote_code
+            and not self.cfg.gptq
+            and self.cfg.attn_uses_flash_lib
+            and is_flash_attn_available()
+            and not self.inference
+        ):
+            try:
+                # TODO(MengqingCao): split these patches separately
+                from axolotl.monkeypatch.llama_attn_hijack_flash import (
+                    is_xformers_swiglu_available,
+                    replace_llama_mlp_with_swiglu,
+                )
+
+                if self.cfg.flash_attn_fuse_mlp and is_xformers_swiglu_available():
+                    LOG.info("Patching with SwiGLU...")
+                    replace_llama_mlp_with_swiglu(model)
+            except ImportError as e:
+                LOG.warning(f"Flash Attention patches not applied: {e}")
+
+    def _apply_lora_kernel_patch(self, model):
+        """Apply LoRA kernel patches."""
+        if (
+            self.cfg.lora_mlp_kernel
+            or self.cfg.lora_qkv_kernel
+            or self.cfg.lora_o_kernel
+        ):
+            from axolotl.monkeypatch.lora_kernels import apply_lora_kernel_patches
+
+            apply_lora_kernel_patches(model=model, cfg=self.cfg)
+
+    def _apply_patch_deepspeed_zero3(self):
+        try:
+            from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+
+            from axolotl.monkeypatch.deepspeed_utils import apply_deepspeed_patches
+
+            if self.cfg.activation_offloading is True and (
+                is_deepspeed_zero3_enabled()
+                or os.getenv("ACCELERATE_DEEPSPEED_ZERO_STAGE") == "3"
+            ):
+                apply_deepspeed_patches()
+        except ImportError as e:
+            LOG.warning(f"DeepSpeed patches not applied: {e}")
+
+    def _apply_apertus_patches(self):
+        """Apply patches for Apertus model."""
+        if self.cfg.model_config_type == "apertus":
+            from axolotl.monkeypatch.models.apertus.activation import (
+                patch_apertus_xielu_activation,
+            )
+
+            patch_apertus_xielu_activation()
+
+    def _apply_trl_vllm_patches(self):
+        """Apply TRL vLLM patches for batched weight sync, NaN logprobs fix, and scalar handling."""
+        if (
+            self.cfg.rl
+            and getattr(self.cfg, "trl", None)
+            and getattr(self.cfg.trl, "use_vllm", False)
+        ):
+            from axolotl.monkeypatch.trainer.trl_vllm import patch_trl_vllm
+
+            patch_trl_vllm()
+
+    def _apply_trl_trainer_utils_patches(self):
+        """Replace trl.trainer.utils.{selective_log_softmax, entropy_from_logits} with Triton kernels."""
+        if not self.cfg.rl:
+            return
+
+        try:
+            from axolotl.monkeypatch.trainer.utils import (
+                entropy_from_logits,
+                selective_log_softmax,
+            )
+        except (ImportError, ModuleNotFoundError):
+            LOG.warning("Triton not available — skipping trl.trainer.utils patches")
+            return
+
+        import trl.trainer.utils
+
+        # Guard against repeated calls: only stash the original if trl still
+        # points at its own implementation (not our wrapper).
+        if trl.trainer.utils.selective_log_softmax is not selective_log_softmax:
+            from axolotl.monkeypatch.trainer import utils as _axolotl_trainer_utils
+
+            _axolotl_trainer_utils.selective_log_softmax_original = (
+                trl.trainer.utils.selective_log_softmax
+            )
+            trl.trainer.utils.selective_log_softmax = selective_log_softmax
+
+        if trl.trainer.utils.entropy_from_logits is not entropy_from_logits:
+            trl.trainer.utils.entropy_from_logits = entropy_from_logits
+
+        LOG.info(
+            "Patched trl.trainer.utils with Triton selective_log_softmax and entropy_from_logits"
+        )
+
+    def _apply_scaling_softmax_patch(self, model: PreTrainedModel):
+        """Apply Scaling Softmax (SSMax) patch.  Ref: https://arxiv.org/abs/2501.19399"""
+        if self.cfg.scaling_softmax:
+            from axolotl.monkeypatch.scaled_softmax_attn import (
+                patch_scaled_softmax_attention,
+            )
+
+            patch_scaled_softmax_attention(
+                scaling_factor_init=self.cfg.scaling_softmax_factor or 0.43,
+                bias=self.cfg.scaling_softmax_bias or 0.0,
+                model=model,
+            )
diff --git a/src/axolotl/loaders/processor.py b/src/axolotl/loaders/processor.py
new file mode 100644
index 0000000000..4cb03b84db
--- /dev/null
+++ b/src/axolotl/loaders/processor.py
@@ -0,0 +1,92 @@
+"""Processor loading functionality for multi-modal models"""
+
+import transformers
+from transformers import (
+    AutoProcessor,
+    PreTrainedTokenizerBase,
+)
+
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+@send_errors
+def load_processor(cfg: DictDefault, tokenizer: PreTrainedTokenizerBase):
+    processor_cls = AutoProcessor
+    if cfg.processor_type:
+        processor_cls = getattr(transformers, cfg.processor_type)
+
+    # Build common kwargs for processor loading
+    processor_kwargs = {}
+    if cfg.revision_of_model:
+        processor_kwargs["revision"] = cfg.revision_of_model
+    if cfg.processor_kwargs:
+        processor_kwargs.update(cfg.processor_kwargs)
+
+    if cfg.tokenizer_use_mistral_common:
+
+        def _patch_mistralcommontokenizer():
+            """
+            Transformers v5 stops reading the sub-processor.
+
+            We need to patch this, so both processors use this.
+            """
+            import transformers.tokenization_mistral_common as tokenization_mistral_common
+
+            from axolotl.utils.mistral import HFMistralTokenizer
+
+            tokenization_mistral_common.MistralCommonBackend = HFMistralTokenizer
+
+        _patch_mistralcommontokenizer()
+
+        from transformers import VoxtralProcessor
+
+        if processor_cls == VoxtralProcessor:
+            return VoxtralProcessor.from_pretrained(
+                cfg.processor_config,
+                **processor_kwargs,
+            )
+
+        from axolotl.utils.mistral import Mistral3Processor
+
+        return Mistral3Processor(
+            tokenizer=tokenizer,
+        )
+
+    processor_kwargs["trust_remote_code"] = cfg.trust_remote_code or False
+
+    processor = processor_cls.from_pretrained(
+        cfg.processor_config,
+        **processor_kwargs,
+    )
+    processor.tokenizer = tokenizer
+
+    # Attempt to load image size from processor if available
+    if (
+        cfg.image_size is None
+        and hasattr(processor, "size")
+        and any(dim in processor.size for dim in ["width", "height"])
+    ):
+        im_width = None
+        im_height = None
+        if "width" in processor.size:
+            im_width = processor.size["width"]
+        if "height" in processor.size:
+            im_height = processor.size["height"]
+
+        # If both width and height are set, use a tuple
+        if im_width is not None and im_height is not None:
+            cfg.image_size = (im_width, im_height)
+        # If only width is set, use as integer
+        elif im_width is not None:
+            cfg.image_size = im_width
+        # If only height is set, use as integer
+        elif im_height is not None:
+            cfg.image_size = im_height
+
+        LOG.debug(f"Loaded image size: {cfg.image_size} from processor")
+
+    return processor
diff --git a/src/axolotl/loaders/tokenizer.py b/src/axolotl/loaders/tokenizer.py
new file mode 100644
index 0000000000..572a880bd4
--- /dev/null
+++ b/src/axolotl/loaders/tokenizer.py
@@ -0,0 +1,336 @@
+"""Tokenizer loading functionality and associated utils"""
+
+import json
+import os
+
+import transformers
+from transformers import (
+    AddedToken,
+    AutoTokenizer,
+    PreTrainedTokenizer,
+)
+
+from axolotl.integrations.base import PluginManager
+from axolotl.loaders.utils import get_linear_embedding_layers, load_model_config
+from axolotl.prompt_tokenizers import LLAMA_DEFAULT_EOS_TOKEN
+from axolotl.telemetry.errors import send_errors
+from axolotl.utils.chat_templates import get_chat_template_from_config
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.distributed import (
+    barrier,
+    is_local_main_process,
+    is_main_process,
+)
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+PLUGIN_MANAGER = PluginManager.get_instance()
+
+
+def modify_tokenizer_files(
+    tokenizer_path: str,
+    token_mappings: dict[int, str],
+    output_dir: str,
+    revision: str = "main",
+) -> str:
+    """
+    Modify tokenizer files to replace added_tokens strings, save to output directory,
+    and return the path to the modified tokenizer.
+
+    This only works with reserved tokens that were added to the tokenizer, not tokens
+    already part of the vocab.
+
+    Args:
+        tokenizer_path: Path or name of the original tokenizer
+        token_mappings: Dict mapping {token_id (int): new_token_string}
+        output_dir: Directory to save the modified tokenizer
+        revision: Model revision/branch/tag/commit to load from (HF Hub)
+
+    Returns:
+        Path to the modified tokenizer directory
+
+    Ref: https://github.com/huggingface/transformers/issues/27974#issuecomment-1854188941
+    """
+    # Create the tokenizer directory in output_dir if it doesn't exist
+    tokenizer_dir = os.path.join(output_dir, "tokenizer")
+    os.makedirs(tokenizer_dir, exist_ok=True)
+
+    if is_local_main_process():
+        # Load the tokenizer
+        temp_tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_path, use_fast=True, revision=revision
+        )
+
+        # Save the tokenizer to the output directory
+        temp_tokenizer.save_pretrained(tokenizer_dir)
+
+        # Get the token IDs and map them to their new values
+        token_id_mappings = {
+            int(token_id): new_value for token_id, new_value in token_mappings.items()
+        }
+
+        # 1. Update tokenizer_config.json - added_tokens_decoder
+        config_path = os.path.join(tokenizer_dir, "tokenizer_config.json")
+        if os.path.exists(config_path):
+            with open(config_path, "r", encoding="utf-8") as f:
+                config_data = json.load(f)
+
+            # Update added_tokens_decoder
+            if "added_tokens_decoder" in config_data:
+                for token_id, new_value in token_id_mappings.items():
+                    token_id_str = str(token_id)
+                    if token_id_str in config_data["added_tokens_decoder"]:
+                        config_data["added_tokens_decoder"][token_id_str]["content"] = (
+                            new_value
+                        )
+                    else:
+                        raise ValueError(
+                            f"Token ID {token_id_str} not found in added_tokens_decoder"
+                        )
+
+            # Write the updated config back
+            with open(config_path, "w", encoding="utf-8") as f:
+                json.dump(config_data, f, indent=2)
+
+        # 2. Update tokenizer.json - added_tokens
+        tokenizer_path = os.path.join(tokenizer_dir, "tokenizer.json")
+        if os.path.exists(tokenizer_path):
+            with open(tokenizer_path, "r", encoding="utf-8") as f:
+                tokenizer_data = json.load(f)
+
+            # Update added_tokens
+            if "added_tokens" in tokenizer_data:
+                for token_id, new_value in token_id_mappings.items():
+                    for i, token_entry in enumerate(tokenizer_data["added_tokens"]):
+                        if token_entry["id"] == token_id:
+                            tokenizer_data["added_tokens"][i]["content"] = new_value
+                            break
+                    else:
+                        # Reaching this section means the token_id was not found in tokenizer.json added_tokens
+                        raise ValueError(
+                            f"Token ID {token_id} not found in added_tokens"
+                        )
+            if "model" in tokenizer_data and "vocab" in tokenizer_data["model"]:
+                for token_id, new_value in token_id_mappings.items():
+                    for entry_val, entry_id in tokenizer_data["model"]["vocab"].items():
+                        if entry_id == token_id:
+                            del tokenizer_data["model"]["vocab"][entry_val]
+                            tokenizer_data["model"]["vocab"][new_value] = token_id
+                            break
+
+            # Write the updated tokenizer data back
+            with open(tokenizer_path, "w", encoding="utf-8") as f:
+                json.dump(tokenizer_data, f, indent=2)
+
+    barrier()
+    return tokenizer_dir
+
+
+@send_errors
+def load_tokenizer(cfg: DictDefault) -> PreTrainedTokenizer:
+    """Load and configure the tokenizer based on the provided config."""
+
+    # Apply patches that need to be in place before tokenizer loading
+    from axolotl.loaders.patch_manager import PatchManager
+
+    PatchManager.apply_pre_tokenizer_load_patches(cfg)
+
+    def _load_mistral_common_tokenizer(cfg: DictDefault):
+        """Load mistral-common tokenizer"""
+        from axolotl.utils.mistral import HFMistralTokenizer
+
+        # Load the HF-compatible wrapper around MistralTokenizer
+        kwargs = {}
+        if cfg.revision_of_model:
+            kwargs["revision"] = cfg.revision_of_model
+        tokenizer = HFMistralTokenizer.from_pretrained(cfg.tokenizer_config, **kwargs)
+
+        return tokenizer
+
+    if cfg.tokenizer_use_mistral_common:
+        return _load_mistral_common_tokenizer(cfg)
+
+    model_config = load_model_config(cfg)
+    tokenizer_kwargs = {}
+    use_fast = True  # this is the default
+
+    if cfg.tokenizer_use_fast is not None:
+        use_fast = cfg.tokenizer_use_fast
+    if cfg.tokenizer_legacy is not None:
+        # True is the default w/ https://github.com/huggingface/transformers/pull/25224
+        tokenizer_kwargs["legacy"] = cfg.tokenizer_legacy
+    if cfg.revision_of_model:
+        tokenizer_kwargs["revision"] = cfg.revision_of_model
+
+    tokenizer_cls = AutoTokenizer
+    if cfg.tokenizer_type:
+        tokenizer_cls = getattr(transformers, cfg.tokenizer_type)
+
+    # Set base tokenizer path
+    tokenizer_path = cfg.tokenizer_config
+
+    # Apply token string overrides if specified
+    if cfg.added_tokens_overrides:
+        # Modify tokenizer files and get path to modified tokenizer
+        modify_kwargs = {"output_dir": cfg.output_dir}
+        if cfg.revision_of_model:
+            modify_kwargs["revision"] = cfg.revision_of_model
+        tokenizer_path = modify_tokenizer_files(
+            tokenizer_path, cfg.added_tokens_overrides, **modify_kwargs
+        )
+
+    tokenizer = tokenizer_cls.from_pretrained(
+        tokenizer_path,
+        trust_remote_code=cfg.trust_remote_code or False,
+        use_fast=use_fast,
+        **tokenizer_kwargs,
+    )
+
+    if (
+        tokenizer.__class__.__name__
+        in [
+            "LlamaTokenizer",
+            "LlamaTokenizerFast",
+            "CodeLlamaTokenizer",
+            "CodeLlamaTokenizerFast",
+        ]
+        and hasattr(tokenizer, "pad_token")
+        and not tokenizer.pad_token
+    ):
+        # set a pad_token, but use eos_token so we don't add a new token
+        tokenizer.pad_token = LLAMA_DEFAULT_EOS_TOKEN
+
+    if tokenizer.__class__.__name__ == "GPTNeoXTokenizerFast":
+        tokenizer.add_special_tokens({"pad_token": "[PAD]"})  # nosec B105
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    # Mistral's official FA implementation requires left padding
+    if (
+        cfg.is_mistral_derived_model
+        and cfg.attn_implementation == "flash_attention_2"
+        and not cfg.sample_packing
+    ):
+        tokenizer.padding_side = "left"
+
+    # Qwen base only has single token, so we need to set the special tokens
+    # the following check is for Qwen1 base models
+    if cfg.is_qwen_derived_model and hasattr(tokenizer, "eod_id"):
+        token_ids = ["bos_token_id", "eos_token_id", "pad_token_id", "unk_token_id"]
+        for attr_name in token_ids:
+            if getattr(tokenizer, attr_name) is None:
+                setattr(tokenizer, attr_name, tokenizer.eod_id)
+
+        token_names = ["bos_token", "eos_token", "pad_token", "unk_token"]
+        for attr_name in token_names:
+            if getattr(tokenizer, attr_name) is None:
+                setattr(tokenizer, attr_name, "<|endoftext|>")
+
+    additional_special_tokens = None
+    if cfg.special_tokens:
+        special_tokens = cfg.special_tokens.to_dict()
+        additional_special_tokens = special_tokens.pop(
+            "additional_special_tokens", None
+        )
+        lora_modules_to_save = get_linear_embedding_layers(model_config.model_type)
+        for k, val in special_tokens.items():
+            # check if new special token is not already in tokenizer and
+            # is adapter training to make sure lora_modules_to_save is set
+
+            if (
+                (getattr(tokenizer, k) is None or getattr(tokenizer, k) != val)
+                and (len(tokenizer.encode(val, add_special_tokens=False)) > 2)
+                and cfg.adapter
+                and (
+                    not cfg.lora_modules_to_save
+                    or not all(
+                        x in cfg.lora_modules_to_save for x in lora_modules_to_save
+                    )
+                )
+                and k != "pad_token"
+            ):
+                lora_modules_to_save_str = ", ".join(
+                    [f"`{x}`" for x in lora_modules_to_save]
+                )
+                raise ValueError(
+                    f"Please set lora_modules_to_save to [{lora_modules_to_save_str}] "
+                    "when using an adapter and changing the special tokens."
+                )
+
+            tokenizer.add_special_tokens(
+                {k: AddedToken(val, rstrip=False, lstrip=False, normalized=False)}
+            )
+
+        # If we add bos_token and eos_token, we need to update the post processor to
+        # handle them correctly.
+        # https://github.com/huggingface/transformers/pull/24132
+        bos_or_eos_in_special_tokens = (
+            "bos_token" in cfg.special_tokens and "eos_token" in cfg.special_tokens
+        )
+        if (
+            tokenizer.__class__.__name__
+            in (
+                "LlamaTokenizerFast",
+                "CodeLlamaTokenizerFast",
+            )
+            and bos_or_eos_in_special_tokens
+        ):
+            tokenizer.update_post_processor()
+
+    if cfg.tokens:
+        tokenizer.add_tokens(
+            [
+                AddedToken(token, rstrip=False, lstrip=False, normalized=False)
+                for token in cfg.tokens
+            ]
+        )
+
+    # Additional special tokens are a List, and need to be treated differently than regular special
+    # tokens. We add them after we have called `add_tokens` in case these additional special tokens
+    # are new tokens.
+    #
+    # Usage:
+    #
+    # ```py
+    # special_tokens:
+    #   additional_special_tokens: ["<|im_start|>", "<|im_end|>"]
+    # ```
+    if additional_special_tokens is not None:
+        tokenizer.add_special_tokens(
+            {"additional_special_tokens": additional_special_tokens}
+        )
+
+    # Generic fallback: if tokenizer still has no pad_token, use eos_token
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+        LOG.warning(
+            "Tokenizer does not have a pad_token, falling back to eos_token: %s",
+            tokenizer.eos_token,
+        )
+
+    if is_main_process():
+        LOG.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}")
+        LOG.debug(f"BOS: {tokenizer.bos_token_id} / {tokenizer.bos_token}")
+        LOG.debug(f"PAD: {tokenizer.pad_token_id} / {tokenizer.pad_token}")
+        LOG.debug(f"UNK: {tokenizer.unk_token_id} / {tokenizer.unk_token}")
+
+    if cfg.chat_template:
+        chat_template_string = get_chat_template_from_config(
+            cfg=cfg,
+            tokenizer=tokenizer,
+        )
+        if cfg.default_system_message and cfg.chat_template == "chatml":
+            chat_template_string = chat_template_string.replace(
+                "You are a helpful assistant.", cfg.default_system_message
+            )
+
+        tokenizer.chat_template = chat_template_string
+    elif getattr(tokenizer, "chat_template", None) is None:
+        LOG.info(
+            "No Chat template selected. Consider adding a chat template for easier inference."
+        )
+
+    # make the tokenizer.pad call quieter 🤐
+    if hasattr(tokenizer, "deprecation_warnings"):
+        tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+
+    return tokenizer
diff --git a/src/axolotl/loaders/utils.py b/src/axolotl/loaders/utils.py
new file mode 100644
index 0000000000..ce4018014a
--- /dev/null
+++ b/src/axolotl/loaders/utils.py
@@ -0,0 +1,239 @@
+"""Utilities for axolotl.loaders module"""
+
+import contextlib
+from typing import Type
+
+import addict
+import torch
+import transformers
+from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def get_module_class_from_name(
+    module: torch.nn.Module, name: str
+) -> Type[torch.nn.Module] | None:
+    """Gets a class from a module by its name. Copied from `accelerate.utils.dataclasses`
+    (https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2805).
+
+    Args:
+        module: The module to get the class from.
+        name: The name of the class.
+
+    Returns:
+        The class type of the matching module, or `None` if no match is found.
+    """
+    modules_children = list(module.children())
+    if module.__class__.__name__ == name:
+        return module.__class__
+
+    if len(modules_children) == 0:
+        return None
+
+    for child_module in modules_children:
+        module_class = get_module_class_from_name(child_module, name)
+        if module_class is not None:
+            return module_class
+
+    return None
+
+
+def check_model_config(cfg: DictDefault, model_config: PretrainedConfig):
+    """Validates and adjusts model config based on `axolotl` config.
+
+    This function performs several important checks and adjustments:
+        - Disables model caching for better memory efficiency
+        - Handles multimodal model-specific configurations
+        - Validates quantization settings
+        - Ensures proper LoRA configuration when using adapters with new tokens
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        model_config: The model's configuration object from `transformers`.
+
+    Raises:
+        ValueError: If a multimodal model lacks text configuration, if GPTQ settings
+            are inconsistent, or if LoRA `modules_to_save` is improperly configured
+            with new tokens.
+    """
+    if hasattr(model_config, "use_cache"):
+        model_config.use_cache = False
+
+    if cfg.is_multimodal:
+        # For multimodal configs, use_cache is set in the text_config
+        if hasattr(model_config, "get_text_config"):
+            text_config = model_config.get_text_config()
+            if hasattr(text_config, "use_cache"):
+                text_config.use_cache = False
+        else:
+            raise ValueError(
+                "No text config found for multimodal model. Please raise an Issue with model details."
+            )
+
+        # Check if image_size is not set and load image size from model config if available
+        if (
+            cfg.image_size is None
+            and hasattr(model_config, "vision_config")
+            and hasattr(model_config.vision_config, "image_size")
+        ):
+            image_size = model_config.vision_config.image_size
+            if isinstance(image_size, list):
+                cfg.image_size = tuple(image_size)
+            else:
+                cfg.image_size = image_size
+            LOG.debug(f"Loaded image size: {cfg.image_size} from model config")
+
+    quant_config_exists = (
+        hasattr(model_config, "quantization_config")
+        and model_config.quantization_config
+    )
+
+    # Detect compressed-tensors config
+    is_compressed_tensors_config = (
+        quant_config_exists
+        and model_config.quantization_config.get("quant_method") == "compressed-tensors"
+    )
+
+    if is_compressed_tensors_config:
+        if model_config.quantization_config.get("config_groups"):
+            LOG.warning(
+                "Found `config_groups` in a compressed-tensors config. "
+                "QAT integration with llmcompressor is not tested."
+            )
+        # Skip further quant checks for compressed-tensors
+        return
+
+    quant_config_method_is_gptq = (
+        quant_config_exists
+        and "quant_method" in model_config.quantization_config
+        and model_config.quantization_config["quant_method"] == "gptq"
+    )
+
+    if cfg.gptq and not quant_config_method_is_gptq:
+        raise ValueError(
+            "model_config.quantization_config is not set or quant_method is not set to gptq. "
+            "Please make sure to point to a GPTQ model."
+        )
+
+    lora_modules_to_save = get_linear_embedding_layers(model_config.model_type)
+    if (
+        cfg.adapter
+        and cfg.tokens
+        and (
+            not cfg.lora_modules_to_save
+            or not all(x in cfg.lora_modules_to_save for x in lora_modules_to_save)
+        )
+    ):
+        lora_modules_to_save_joined = ", ".join(
+            map(lambda x: f"`{x}`", lora_modules_to_save)
+        )
+        raise ValueError(
+            "`lora_modules_to_save` not properly set when adding new tokens. "
+            f"Please include [{lora_modules_to_save_joined}] in `lora_modules_to_save`."
+        )
+
+    if (
+        cfg.tensor_parallel_size
+        and cfg.tensor_parallel_size > 1
+        and hasattr(model_config, "tie_word_embeddings")
+        and model_config.tie_word_embeddings
+    ):
+        raise ValueError(
+            "Tensor parallelism is incompatible with models configured with `tie_word_embeddings` enabled. "
+            "Please use a model without `tie_word_embeddings`, or disable tensor parallelism."
+        )
+
+
+def load_model_config(cfg: DictDefault) -> PretrainedConfig | addict.Dict:
+    """Loads and configures a model configuration from HuggingFace or local sources.
+
+    This function determines the appropriate model config source, loads it, applies any
+    necessary overrides, and validates it for compatibility with the `axolotl` config.
+
+    If `cfg.cls_model_config` is set, a custom config class from transformers will be
+    used instead of `AutoConfig` (e.g., 'LlamaConfig', 'MistralConfig').
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+
+    Returns:
+        A configured model configuration object (`AutoConfig` instance), or a simple
+            dictionary configuration for special cases like Mamba models.
+
+    Raises:
+        ValueError: If configuration loading fails for reasons other than special cases
+            that are handled (e.g., Mamba models).
+    """
+    model_config_name = cfg.base_model_config or cfg.base_model
+    if not model_config_name and cfg.tokenizer_config:
+        model_config_name = cfg.tokenizer_config
+    trust_remote_code = cfg.trust_remote_code is True
+    config_kwargs = {}
+    if cfg.revision_of_model:
+        config_kwargs["revision"] = cfg.revision_of_model
+    if cfg.num_labels:
+        # num_labels is used to initialize classifier models
+        config_kwargs["num_labels"] = cfg.num_labels
+
+    config_cls = AutoConfig
+    if cfg.cls_model_config:
+        config_cls = getattr(transformers, cfg.cls_model_config)
+
+    try:
+        model_config = config_cls.from_pretrained(
+            model_config_name,
+            trust_remote_code=trust_remote_code,
+            **config_kwargs,
+        )
+    except ValueError as error:
+        if "mamba" in model_config_name:
+            return addict.Dict(
+                {
+                    "model_type": "mamba",
+                }
+            )
+        raise error
+
+    if cfg.overrides_of_model_config:
+        for key, val in cfg.overrides_of_model_config.items():
+            setattr(model_config, key, val)
+
+    check_model_config(cfg, model_config)
+
+    return model_config
+
+
+def ensure_dtype(model: PreTrainedModel, dtype: torch.dtype = torch.bfloat16):
+    """Ensures all modules in the model are converted to the specified data type."""
+    for name, module in model.named_modules():
+        weight_mismatch = False
+        with contextlib.suppress(AttributeError):
+            weight_mismatch = module.weight.dtype != dtype
+
+        bias_mismatch = False
+        with contextlib.suppress(AttributeError):
+            bias_mismatch = module.bias.dtype != dtype
+
+        if weight_mismatch:
+            LOG.debug(
+                f"Converting module {name}.weight: {module.weight.dtype} -> {dtype}"
+            )
+        if bias_mismatch:
+            LOG.debug(f"Converting module {name}.bias: {module.bias.dtype} -> {dtype}")
+        if weight_mismatch or bias_mismatch:
+            module.to(dtype)
+
+
+def get_linear_embedding_layers(model_type: str) -> list[str]:
+    """Returns layer names of linear embeddings needed for LoRA based on model type."""
+    if model_type == "gpt_neox":
+        return ["embed_in", "embed_out"]
+    if model_type == "falcon":
+        return ["word_embeddings", "lm_head"]
+    if model_type == "nemotron_h":
+        return ["embeddings", "lm_head"]
+    return ["embed_tokens", "lm_head"]
diff --git a/src/axolotl/logging_config.py b/src/axolotl/logging_config.py
index 2ddf89a8c4..e12a8be419 100644
--- a/src/axolotl/logging_config.py
+++ b/src/axolotl/logging_config.py
@@ -1,15 +1,72 @@
-"""
-Common logging module for axolotl
-"""
+"""Common logging module for axolotl."""
 
+import logging
 import os
-import sys
-from logging import Formatter
+from logging import Formatter, Logger, LogRecord
 from logging.config import dictConfig
 from typing import Any, Dict
 
 from colorama import Fore, Style, init
 
+DEFAULT_AXOLOTL_LOG_LEVEL = "INFO"
+DEFAULT_LOG_LEVEL = "WARNING"
+
+
+class AxolotlOrWarnErrorFilter(logging.Filter):
+    """
+    Allows ANY WARNING or higher (unless overridden by LOG_LEVEL). Allows axolotl.* at
+    INFO or higher (unless overridden by AXOLOTL_LOG_LEVEL). Drops all other records
+    (i.e. non-axolotl.INFO, DEBUG, etc. by default).
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        axolotl_log_level = os.getenv(
+            "AXOLOTL_LOG_LEVEL", DEFAULT_AXOLOTL_LOG_LEVEL
+        ).upper()
+        other_log_level = os.getenv("LOG_LEVEL", DEFAULT_LOG_LEVEL).upper()
+
+        try:
+            # py311+ only
+            level_mapping = logging.getLevelNamesMapping()
+            self.axolotl_level = level_mapping[axolotl_log_level]
+            self.other_level = level_mapping[other_log_level]
+        except AttributeError:
+            # For py310, use getLevelName directly
+            self.axolotl_level = logging.getLevelName(axolotl_log_level)
+            self.other_level = logging.getLevelName(other_log_level)
+
+    def filter(self, record: LogRecord) -> bool:
+        # General filter
+        if record.levelno >= self.other_level:
+            return True
+
+        # Axolotl filter
+        return (
+            record.name.startswith("axolotl") and record.levelno >= self.axolotl_level
+        )
+
+
+class HubUnauthenticatedNagFilter(logging.Filter):
+    """
+    Drops the server-sent "sending unauthenticated requests" nag (an X-HF-Warning
+    response header that huggingface_hub logs with no env var to disable it).
+    Other hub warnings (retries, rate limits) pass through.
+    """
+
+    def filter(self, record: LogRecord) -> bool:
+        return "unauthenticated requests" not in record.getMessage()
+
+
+class AxolotlLogger(Logger):
+    """Logger that applies filtering to non-axolotl loggers."""
+
+    def __init__(self, name: str, level: int = logging.NOTSET):
+        super().__init__(name, level)
+        if not name.startswith("axolotl"):
+            self.addFilter(AxolotlOrWarnErrorFilter())
+
 
 class ColorfulFormatter(Formatter):
     """
@@ -24,6 +81,7 @@ class ColorfulFormatter(Formatter):
 
     def format(self, record):
         record.rank = int(os.getenv("LOCAL_RANK", "0"))
+        record.rank_fmt = f" [RANK:{record.rank}]" if record.rank != 0 else ""
         log_message = super().format(record)
         return self.COLORS.get(record.levelname, "") + log_message + Fore.RESET
 
@@ -37,31 +95,65 @@ def format(self, record):
         },
         "colorful": {
             "()": ColorfulFormatter,
-            "format": "[%(asctime)s] [%(levelname)s] [%(name)s.%(funcName)s:%(lineno)d] [PID:%(process)d] [RANK:%(rank)d] %(message)s",
+            "format": "[%(asctime)s] [%(levelname)s] [%(name)s.%(funcName)s:%(lineno)d] [PID:%(process)d]%(rank_fmt)s %(message)s",
+        },
+        "concise": {
+            "format": "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
+        },
+        "concise_color": {
+            "()": ColorfulFormatter,
+            "format": "[%(asctime)s] [%(levelname)s] [%(name)s]%(rank_fmt)s %(message)s",
+        },
+    },
+    "filters": {
+        "ax_or_warn": {
+            "()": "axolotl.logging_config.AxolotlOrWarnErrorFilter",
+        },
+        "hub_unauthenticated_nag": {
+            "()": "axolotl.logging_config.HubUnauthenticatedNagFilter",
         },
     },
-    "filters": {},
     "handlers": {
         "console": {
             "class": "logging.StreamHandler",
-            "formatter": "simple",
-            "filters": [],
-            "stream": sys.stdout,
+            "formatter": "concise",
+            "filters": ["ax_or_warn"],
+            "stream": "ext://sys.stdout",
         },
         "color_console": {
             "class": "logging.StreamHandler",
-            "formatter": "colorful",
-            "filters": [],
-            "stream": sys.stdout,
+            "formatter": "concise_color",
+            "filters": ["ax_or_warn"],
+            "stream": "ext://sys.stdout",
+        },
+        "ax_file_only": {
+            "class": "logging.StreamHandler",
+            "level": "DEBUG",
+            "formatter": "simple",
+            "stream": "ext://axolotl.utils.tee.file_only_stream",
+        },
+        "root_file_only": {
+            "class": "logging.StreamHandler",
+            "level": "DEBUG",
+            "formatter": "simple",
+            "stream": "ext://axolotl.utils.tee.file_only_stream",
         },
     },
-    "root": {"handlers": ["console"], "level": os.getenv("LOG_LEVEL", "INFO")},
+    "root": {
+        "handlers": ["console", "root_file_only"],
+        "level": os.getenv("LOG_LEVEL", DEFAULT_LOG_LEVEL).upper(),
+    },
     "loggers": {
         "axolotl": {
-            "handlers": ["color_console"],
-            "level": "DEBUG",
+            "handlers": ["color_console", "ax_file_only"],
+            "level": os.getenv("AXOLOTL_LOG_LEVEL", DEFAULT_AXOLOTL_LOG_LEVEL).upper(),
             "propagate": False,
         },
+        # filter at the emitting logger so the nag is dropped before it reaches
+        # both huggingface_hub's own handler and our root handlers
+        "huggingface_hub.utils._http": {
+            "filters": ["hub_unauthenticated_nag"],
+        },
     },
 }
 
@@ -69,4 +161,15 @@ def format(self, record):
 def configure_logging():
     """Configure with default logging"""
     init()  # Initialize colorama
+
     dictConfig(DEFAULT_LOGGING_CONFIG)
+    logging.setLoggerClass(AxolotlLogger)
+
+    # Route Python warnings through logging so they reach file handlers
+    logging.captureWarnings(True)
+
+    # Set default `ACCELERATE_LOG_LEVEL` to `LOG_LEVEL` if available and not set
+    if "ACCELERATE_LOG_LEVEL" not in os.environ:
+        os.environ["ACCELERATE_LOG_LEVEL"] = os.getenv(
+            "LOG_LEVEL", DEFAULT_LOG_LEVEL
+        ).upper()
diff --git a/src/axolotl/loraplus.py b/src/axolotl/loraplus.py
deleted file mode 100644
index b4abec55ad..0000000000
--- a/src/axolotl/loraplus.py
+++ /dev/null
@@ -1,133 +0,0 @@
-"""Module for LoRA+"""
-
-# MIT License
-#
-# Copyright (c) 2024 nikhil-ghosh-berkeley
-# https://github.com/nikhil-ghosh-berkeley/loraplus
-
-import logging
-from functools import reduce
-
-from peft.tuners import lora
-from torch import nn
-from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
-from transformers.trainer_pt_utils import get_parameter_names
-
-LOG = logging.getLogger("axolotl.loraplus")
-
-
-def get_module(name, opt_model):
-    """
-    Retrieve a module from a model using its parameter name.
-    Args:
-        name (str): Full name of the parameter, typically including module path.
-        opt_model (torch.nn.Module): The model from which to retrieve the module.
-
-    Returns:
-        Module corresponding to the given name.
-    """
-    parent_idx = 2 if "lora" in name else 1
-    module_names = name.split(sep=".")[:-parent_idx]
-    module = reduce(getattr, module_names, opt_model)
-    return module
-
-
-def create_loraplus_optimizer(
-    opt_model,
-    optimizer_cls,
-    optimizer_kwargs,
-    loraplus_lr_ratio,
-    loraplus_lr_embedding=None,
-):
-    """
-    Creates an optimizer for the given model, applying LoRA-specific learning rate adjustments to different parameter groups.
-
-    Args:
-        opt_model (torch.nn.Module): The model for which the optimizer is being created.
-        optimizer_cls (class): The class of the optimizer to be used (e.g., torch.optim.Adam).
-        optimizer_kwargs (dict): A dictionary of keyword arguments for the optimizer's initialization.
-        loraplus_lr_ratio (float): The learning rate ratio to be applied to LoRA parameters.
-        loraplus_lr_embedding (float, optional): A specific learning rate for embedding parameters, with a default value if not provided.
-
-    Returns:
-        An instance of the specified optimizer class configured with the model's parameters organized into groups with custom learning rates.
-    """
-
-    assert loraplus_lr_ratio is not None, "loraplus_lr_ratio must be provided."
-
-    if loraplus_lr_embedding is None:
-        loraplus_lr_embedding = 1e-6
-
-    decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
-    decay_parameters = [name for name in decay_parameters if "bias" not in name]
-    param_groups = {
-        "groupA": {},
-        "groupB": {},
-        "groupB_no_decay": {},
-        "embedding": {},
-    }
-
-    for name, param in opt_model.named_parameters():
-        if not param.requires_grad:
-            continue
-
-        module = get_module(name, opt_model)
-        if isinstance(module, lora.Embedding):
-            param_groups["embedding"][name] = param
-        elif "lora_B" in name or param.ndim == 1:
-            if name in decay_parameters:
-                param_groups["groupB"][name] = param
-            else:
-                param_groups["groupB_no_decay"][name] = param
-        else:
-            param_groups["groupA"][name] = param
-
-    assigned_param_groups = ""
-    for group, group_params in param_groups.items():
-        assigned_param_groups += f"{group}\n {list(group_params.keys())}\n\n"
-    LOG.info(assigned_param_groups)
-
-    lr = optimizer_kwargs["lr"]  # pylint: disable=invalid-name
-    weight_decay = optimizer_kwargs.get("weight_decay", 0.0)
-
-    optimizer_grouped_parameters = [
-        {
-            "params": list(param_groups["groupA"].values()),
-            "weight_decay": weight_decay,
-            "lr": lr,
-        },
-        {
-            "params": list(param_groups["embedding"].values()),
-            "weight_decay": weight_decay,
-            "lr": loraplus_lr_embedding,
-        },
-        {
-            "params": list(param_groups["groupB"].values()),
-            "weight_decay": weight_decay,
-            "lr": lr * loraplus_lr_ratio,
-        },
-        {
-            "params": list(param_groups["groupB_no_decay"].values()),
-            "weight_decay": 0.0,
-            "lr": lr * loraplus_lr_ratio,
-        },
-    ]
-
-    optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
-    if optimizer_cls.__name__ == "Adam8bit":
-        import bitsandbytes
-
-        manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
-
-        skipped = 0
-        for module in opt_model.modules():
-            if isinstance(module, nn.Embedding):
-                skipped += sum(
-                    {p.data_ptr(): p.numel() for p in module.parameters()}.values()
-                )
-                LOG.info(f"skipped {module}: {skipped/2**20}M params")
-                manager.register_module_override(module, "weight", {"optim_bits": 32})
-                LOG.debug(f"bitsandbytes: will optimize {module} in fp32")
-        LOG.info(f"skipped: {skipped/2**20}M params")
-
-    return optimizer
diff --git a/src/axolotl/models/mamba/__init__.py b/src/axolotl/models/mamba/__init__.py
index fee88e3a43..d6bb40d99a 100644
--- a/src/axolotl/models/mamba/__init__.py
+++ b/src/axolotl/models/mamba/__init__.py
@@ -21,4 +21,4 @@ def fix_mamba_attn_for_loss():
     from .modeling_mamba import MambaLMHeadModel as MambaLMHeadModelFixed
 
     mixer_seq_simple.MambaLMHeadModel = MambaLMHeadModelFixed
-    return mixer_seq_simple.MambaLMHeadModel  # pylint: disable=invalid-name
+    return mixer_seq_simple.MambaLMHeadModel
diff --git a/src/axolotl/models/mamba/configuration_mamba.py b/src/axolotl/models/mamba/configuration_mamba.py
index 5160ee8d7e..d6b77b951a 100644
--- a/src/axolotl/models/mamba/configuration_mamba.py
+++ b/src/axolotl/models/mamba/configuration_mamba.py
@@ -1,6 +1,7 @@
 """
 HF Transformers MambaConfig
 """
+
 from transformers import PretrainedConfig
 
 
diff --git a/src/axolotl/models/mamba/modeling_mamba.py b/src/axolotl/models/mamba/modeling_mamba.py
index 70e9c88c88..b1847d6b56 100644
--- a/src/axolotl/models/mamba/modeling_mamba.py
+++ b/src/axolotl/models/mamba/modeling_mamba.py
@@ -1,4 +1,3 @@
-# pylint: skip-file
 import os
 from collections import namedtuple
 from functools import partial
@@ -112,7 +111,7 @@ def save_pretrained(
         self,
         save_directory: Union[str, os.PathLike],
         state_dict: Optional[dict] = None,
-        safe_serialization: Optional[bool] = None,  # pylint: disable=unused-argument
+        **kwargs,
     ):
         if state_dict is None:
             state_dict = self.state_dict()
diff --git a/src/axolotl/monkeypatch/__init__.py b/src/axolotl/monkeypatch/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/accelerate/__init__.py b/src/axolotl/monkeypatch/accelerate/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/accelerate/float8_fsdp.py b/src/axolotl/monkeypatch/accelerate/float8_fsdp.py
new file mode 100644
index 0000000000..3ea7d4e48b
--- /dev/null
+++ b/src/axolotl/monkeypatch/accelerate/float8_fsdp.py
@@ -0,0 +1,159 @@
+"""FSDP2 support for torchao ``Float8Tensor`` (frozen blockwise-FP8 weights, dim-0 sharded).
+
+A weight-only ``Float8Tensor`` (e.g. a pre-quantized blockwise-FP8 checkpoint wrapped for
+1-byte storage) can't be FSDP2-sharded as torchao ships it: its ``aten.split`` unpacks
+``(tensor, size, dim)`` unconditionally, so FSDP2's ``torch.chunk(param, n)`` (dim defaults
+to 0) raises ``ValueError: not enough values to unpack``; and even with an explicit dim,
+torchao only implements the rowwise / per-tensor split cases — a 128×128-blocked weight
+hits ``else: raise AssertionError("not yet implemented")``. It also ships no FSDP2
+all-gather hooks.
+
+This patches, for a 2-D ``[N, K]`` weight with a blockwise ``[N//b0, K//b1]`` scale sharded
+on dim 0 (the output-feature axis):
+  * ``aten.split`` with a defaulted dim — qdata splits by ``s``; the scale splits by
+    ``s // block_size[0]`` (the scale is coarser than qdata by ``block_size[0]``),
+  * ``new_zeros`` / ``as_strided`` / ``copy_`` / ``clone`` / ``detach`` / ``narrow`` /
+    ``view`` on the inner ``(qdata, scale)``,
+  * ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather`` so FSDP2 all-gathers the inner tensors
+    and rebuilds the subclass instead of using the flat ``view(-1)`` buffer path.
+
+Idempotent; a no-op if torchao is absent. Dim-0 sharding only (assumes the shard size is a
+multiple of ``block_size[0]`` — true for FSDP of these projections, whose N is a multiple of
+128×world_size). The blockwise split case is worth upstreaming to torchao.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+_PATCHED = False
+
+
+def _float8_cls():
+    try:
+        from torchao.quantization import Float8Tensor
+
+        return Float8Tensor
+    except ImportError:
+        return None
+
+
+def _cstride(shape):
+    st = [1] * len(shape)
+    for i in range(len(shape) - 2, -1, -1):
+        st[i] = st[i + 1] * shape[i + 1]
+    return st
+
+
+def patch_float8_fsdp():
+    global _PATCHED
+    if _PATCHED:
+        return
+    Float8Tensor = _float8_cls()
+    if Float8Tensor is None:
+        return
+    aten = torch.ops.aten
+    implements = Float8Tensor.implements
+
+    def _rebuild(ref, qdata, scale, block_size=None):
+        return Float8Tensor(
+            qdata,
+            scale,
+            list(block_size) if block_size is not None else list(ref.block_size),
+            ref.mm_config,
+            ref.act_quant_kwargs,
+            ref.kernel_preference,
+            dtype=ref.dtype,
+        )
+
+    @implements([aten.split.Tensor])
+    def _split(func, types, args, kwargs):
+        x, split_size = args[0], args[1]
+        dim = args[2] if len(args) > 2 else kwargs.get("dim", 0)
+        if dim != 0:
+            raise NotImplementedError(
+                f"Float8Tensor FSDP split only on dim 0, got {dim}"
+            )
+        blk0 = x.block_size[0]
+        qd = func(x.qdata, split_size, 0)
+        sc = func(x.scale, max(1, split_size // blk0), 0)
+        assert len(qd) == len(sc), f"split mismatch q={len(qd)} s={len(sc)}"
+        return [_rebuild(x, qd[i], sc[i]) for i in range(len(qd))]
+
+    @implements([aten.clone.default, aten.detach.default])
+    def _clone_detach(func, types, args, kwargs):
+        x = args[0]
+        return _rebuild(x, func(x.qdata, **kwargs), func(x.scale, **kwargs))
+
+    @implements([aten.new_zeros.default])
+    def _new_zeros(func, types, args, kwargs):
+        x, size = args[0], list(args[1])
+        qd = func(x.qdata, size, **kwargs)
+        sc_size = [max(1, size[i] // x.block_size[i]) for i in range(len(size))]
+        sc = func(x.scale, sc_size, **kwargs)
+        return _rebuild(x, qd, sc)
+
+    @implements([aten.narrow.default])
+    def _narrow(func, types, args, kwargs):
+        x, dim, start, length = args[0], args[1], args[2], args[3]
+        if dim != 0:
+            raise NotImplementedError(f"Float8Tensor narrow only dim 0, got {dim}")
+        blk0 = x.block_size[0]
+        qd = func(x.qdata, 0, start, length)
+        sc = func(x.scale, 0, start // blk0, max(1, length // blk0))
+        return _rebuild(x, qd, sc)
+
+    @implements([aten.view.default])
+    def _view(func, types, args, kwargs):
+        x, size = args[0], list(args[1])
+        if size == [-1]:
+            # FSDP's flat-buffer view; unused for the subclass-extension all-gather path.
+            return _rebuild(
+                x,
+                x.qdata.reshape(1, -1),
+                x.scale.reshape(1, -1),
+                block_size=[1, x.scale.numel()],
+            )
+        sc_size = [size[i] // x.block_size[i] for i in range(len(size))]
+        return _rebuild(x, func(x.qdata, size), func(x.scale, sc_size))
+
+    @implements([aten.as_strided.default])
+    def _as_strided(func, types, args, kwargs):
+        x = args[0]
+        qs, ss = list(x.qdata.shape), list(x.scale.shape)
+        qd = func(x.qdata, qs, _cstride(qs), 0)
+        sc = func(x.scale, ss, _cstride(ss), 0)
+        return _rebuild(x, qd, sc)
+
+    @implements([aten.copy_.default])
+    def _copy_(func, types, args, kwargs):
+        dst, src = args[0], args[1]
+        func(dst.qdata, src.qdata)
+        func(dst.scale, src.scale)
+        return dst
+
+    def fsdp_pre_all_gather(
+        self, mesh, outer_size=None, outer_stride=None, module=None, mp_policy=None
+    ):
+        return (self.qdata, self.scale), (tuple(self.block_size),)
+
+    def fsdp_post_all_gather(
+        self, all_gather_outputs, metadata, param_dtype, *, out=None
+    ):
+        (block_size,) = metadata
+        qdata, scale = all_gather_outputs
+        if out is not None:
+            out.qdata, out.scale = qdata, scale
+            return
+        return _rebuild(self, qdata, scale, list(block_size)), all_gather_outputs
+
+    Float8Tensor.fsdp_pre_all_gather = fsdp_pre_all_gather
+    Float8Tensor.fsdp_post_all_gather = fsdp_post_all_gather
+
+    _PATCHED = True
+    LOG.info(
+        "Installed FSDP2 support (split + all-gather hooks) on torchao Float8Tensor"
+    )
diff --git a/src/axolotl/monkeypatch/accelerate/fsdp2.py b/src/axolotl/monkeypatch/accelerate/fsdp2.py
new file mode 100644
index 0000000000..78c6b9d42d
--- /dev/null
+++ b/src/axolotl/monkeypatch/accelerate/fsdp2.py
@@ -0,0 +1,979 @@
+"""
+monkeypatch for accelerate fsdp2 fix when modifying ordereddict during interation, and saving full state dicts
+"""
+
+import contextlib
+import copy
+import functools
+import gc
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+from torch import nn
+
+from axolotl.utils.bench import log_gpu_memory_usage
+from axolotl.utils.fp32_norms import get_fp32_norm_patterns, shard_norms_fp32
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _nvfp4_local_tensor_cls(p):
+    """Return the NVFP4Tensor class if ``p`` is a DTensor whose local shard is an NVFP4Tensor."""
+    try:
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+    except ImportError:
+        return None
+    lt = getattr(p, "_local_tensor", None)
+    return NVFP4Tensor if isinstance(lt, NVFP4Tensor) else None
+
+
+def _broadcast_nvfp4_param(sharded_meta_param, full_nvfp4, is_main, device, nvfp4_cls):
+    """Scatter rank-0's full NVFP4Tensor expert param to each rank's shard.
+
+    ``distribute_tensor`` calls ``c10d.scatter_``, which torchao's NVFP4Tensor doesn't implement.
+    Instead scatter the PLAIN component tensors (qdata uint8, scale e4m3, per_tensor_scale fp32 — all
+    collective-capable) along the same Shard placement, then rebuild a local NVFP4Tensor shard and
+    wrap it back into a DTensor matching the sharded model param."""
+    from torch.distributed.tensor import DTensor, distribute_tensor
+
+    mesh = sharded_meta_param.device_mesh
+    placements = sharded_meta_param.placements
+    local_meta = sharded_meta_param._local_tensor
+    e_global = sharded_meta_param.shape[0]
+    block_size = local_meta.block_size
+    dtype = local_meta.dtype
+
+    def _scatter_component(name, ref_local):
+        # Direct per-shard scatter: send each rank ONLY its dim-0 (expert-axis) shard. Avoids the
+        # full-global placeholder that non-rank0 otherwise allocates (256 experts to receive 32),
+        # the generic distribute_tensor path, and its trailing .clone(). Valid for the common
+        # Shard(0) + even-division + full-world mesh (GLM-5.2: 256 experts / 8 ranks); falls back to
+        # distribute_tensor for anything else (uneven, replicate, sub-group meshes).
+        group = mesh.get_group()
+        world = dist.get_world_size(group)
+        p0 = placements[0] if len(placements) == 1 else None
+        # Direct per-shard scatter on the param's mesh group — the full WORLD (pure data-parallel) OR a
+        # dp_shard/cp SUBGROUP (EP composition). Source is the group's rank-0 (== the is_main rank set by
+        # the caller). Avoids distribute_tensor, which deadlocks scattering from src_data_rank=0 on a
+        # sub-mesh. Falls back to distribute_tensor only for non-Shard(0)/uneven placements.
+        if (
+            p0 is not None
+            and getattr(p0, "dim", None) == 0
+            and ref_local.shape[0] * world == e_global
+        ):
+            src_rank = dist.get_process_group_ranks(group)[0]
+            local = torch.empty_like(ref_local, device=device)
+            chunks = None
+            if is_main:
+                full = getattr(full_nvfp4, name)
+                chunks = [c.contiguous().to(device) for c in full.chunk(world, dim=0)]
+            dist.scatter(local, scatter_list=chunks, src=src_rank, group=group)
+            del chunks
+            return local
+        # Fallback: generic distribute_tensor (uneven / replicate / sub-group).
+        gshape = (e_global,) + tuple(ref_local.shape[1:])
+        if is_main:
+            full = getattr(full_nvfp4, name).to(device)
+        else:
+            full = torch.empty(gshape, device=device, dtype=ref_local.dtype)
+        dt = distribute_tensor(full, mesh, placements, src_data_rank=0)
+        return dt._local_tensor.clone()
+
+    local_qdata = _scatter_component("qdata", local_meta.qdata)
+    local_scale = _scatter_component("scale", local_meta.scale)
+
+    local_pts = None
+    pts_ref = getattr(local_meta, "per_tensor_scale", None)
+    if pts_ref is not None:
+        if pts_ref.dim() >= 1 and pts_ref.shape[0] == local_meta.qdata.shape[0]:
+            # per-expert scale shards along dim 0 like qdata/scale
+            local_pts = _scatter_component("per_tensor_scale", pts_ref)
+        else:
+            # replicated scalar — plain broadcast
+            if is_main:
+                local_pts = full_nvfp4.per_tensor_scale.to(device)
+            else:
+                local_pts = torch.empty(
+                    pts_ref.shape, device=device, dtype=pts_ref.dtype
+                )
+            dist.broadcast(local_pts, src=0)
+
+    local_nvfp4 = nvfp4_cls(
+        local_qdata, local_scale, block_size, dtype, per_tensor_scale=local_pts
+    )
+    return DTensor.from_local(local_nvfp4, mesh, placements, run_check=False)
+
+
+def _ep_expert_from_local(sharded_meta_param, full_local, nvfp4_cls):
+    """Build an EP-sharded NVFP4 expert DTensor from THIS rank's already-correct local copy — NO
+    collective. shard_expert_weights scattered each rank its ep-group's full ``[E_local]`` NVFP4
+    (``full_local``); slice this rank's dp-axis shard (``[E_local // dp]`` at its position in the
+    mesh group) and wrap it with ``from_local``. Used instead of a per-subgroup scatter, which
+    deadlocks against the interleaved full-mesh non-expert loads (the receiver ranks race ahead of
+    the source ranks)."""
+    from torch.distributed.tensor import DTensor
+
+    mesh = sharded_meta_param.device_mesh
+    placements = sharded_meta_param.placements
+    local_meta = sharded_meta_param._local_tensor
+    dev = mesh.device_type
+    e_dp = local_meta.qdata.shape[0]  # this rank's dp-local expert count
+    dp_rank = dist.get_group_rank(mesh.get_group(), dist.get_rank())
+    s = slice(dp_rank * e_dp, (dp_rank + 1) * e_dp)
+    qd = full_local.qdata[s].to(dev)
+    sc = full_local.scale[s].to(dev)
+    pts = getattr(full_local, "per_tensor_scale", None)
+    local_pts = None
+    if pts is not None:
+        local_pts = (
+            pts[s].to(dev)
+            if (pts.dim() >= 1 and pts.shape[0] == full_local.qdata.shape[0])
+            else pts.to(dev)
+        )
+    local_nv = nvfp4_cls(
+        qd, sc, local_meta.block_size, local_meta.dtype, per_tensor_scale=local_pts
+    )
+    return DTensor.from_local(local_nv, mesh, placements, run_check=False)
+
+
+def fsdp2_load_full_state_dict(
+    _accelerator, model: torch.nn.Module, full_sd: dict, offload_to_cpu: bool = False
+):
+    """
+    Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the
+    parameters from rank 0 to all other ranks. This function modifies the model in-place.
+    Args:
+        accelerator (`Accelerator`): The accelerator instance
+        model (`torch.nn.Module`):
+            The model to load the state dict into, expected to be on meta device or a VRAM spike can occur
+        full_sd (`dict`): The full state dict to load, can only be on rank 0
+    """
+    LOG.info("Broadcasting full state dict to all ranks...")
+    import time
+
+    start_time = time.time()
+
+    # EP-sharded expert params hold RANK-SPECIFIC content (each rank owns experts
+    # [r*E_local, (r+1)*E_local)). The default rank0-authoritative load below would broadcast
+    # rank 0's shard to every rank, replicating experts[0:E_local] everywhere and silently
+    # destroying expert parallelism. shard_expert_weights() already scattered the correct shard
+    # to every rank BEFORE the model was moved to meta, so each rank's own `full_sd` entry is
+    # already correct — load it locally with no cross-rank broadcast. Match via the `.experts.`
+    # infix (excludes `shared_experts`, survives the `_checkpoint_wrapped_module` infix) tied to
+    # the propagated DDP-ignore list.
+    _ep_ignore = getattr(model, "_ddp_params_and_buffers_to_ignore", None) or []
+    _ep_tails = tuple(
+        sorted(
+            {"." + n.split(".experts.", 1)[1] for n in _ep_ignore if ".experts." in n}
+        )
+    )
+
+    def _is_ep_expert_param(name: str) -> bool:
+        return bool(_ep_tails) and ".experts." in name and name.endswith(_ep_tails)
+
+    meta_sharded_sd = model.state_dict()
+    sharded_sd = {}
+
+    for param_name, sharded_meta_param in meta_sharded_sd.items():
+        # Pure-EP: the EP-sharded experts are excluded from the FSDP wrap (ignored_params), so they
+        # stay PLAIN per-rank meta params here. shard_expert_weights already scattered each rank's
+        # correct [E_local] slice before the meta move, so each rank's own `full_sd` entry is right —
+        # load it locally with NO rank-0 broadcast (which would replicate experts[0:E_local]). Only
+        # applies when the param is plain (not a DTensor); EP×dp_shard composition keeps them as
+        # DTensors loaded via the mesh path below.
+        if _is_ep_expert_param(param_name) and not hasattr(
+            sharded_meta_param, "device_mesh"
+        ):
+            own = full_sd[param_name]
+            own = own.to(torch.device("cuda"))
+            if offload_to_cpu:
+                own = own.cpu()
+            sharded_sd[param_name] = nn.Parameter(
+                own, requires_grad=sharded_meta_param.requires_grad
+            )
+            full_sd[param_name] = None
+            continue
+
+        nvfp4_cls = _nvfp4_local_tensor_cls(sharded_meta_param)
+
+        full_tensor = None
+        # Skip the dtype cast for NVFP4 (its components are scattered raw, not cast to bf16).
+        if _accelerator.is_main_process and nvfp4_cls is None:
+            full_tensor = full_sd[param_name]
+            full_tensor = full_tensor.to(sharded_meta_param.dtype)
+
+        if nvfp4_cls is not None:
+            # NVFP4Tensor params can't go through distribute_tensor (c10d.scatter_ unimplemented on
+            # the subclass); scatter their plain qdata/scale/per_tensor_scale components instead.
+            device_mesh = sharded_meta_param.device_mesh
+            # Source each NVFP4 expert from its OWN mesh-group's rank-0, not the global rank 0. Under
+            # EP×dp_shard / EP×cp composition the experts live on a dp_shard SUBGROUP mesh, and
+            # shard_expert_weights scattered each rank its ep-group's real slice — so the subgroup's
+            # rank-0 holds the data (global rank 0 is outside ep-groups 1..N's subgroups). For full-world
+            # params the subgroup rank-0 IS global rank 0, so this reduces to is_main_process.
+            if _is_ep_expert_param(param_name):
+                # EP-sharded NVFP4 experts: shard_expert_weights already scattered THIS rank its
+                # ep-group's full [E_local] NVFP4 into full_sd. Build the DTensor by slicing the dp-axis
+                # shard out of that local copy and wrapping it with from_local — NO collective. The
+                # per-subgroup scatter (any flavor) deadlocks because it interleaves with the full-mesh
+                # non-expert loads, racing the receiver ranks ahead of the source ranks; from_local
+                # sidesteps every collective.
+                sharded_param = _ep_expert_from_local(
+                    sharded_meta_param, full_sd[param_name], nvfp4_cls
+                )
+            else:
+                # Non-EP NVFP4 on the full data-parallel mesh: only rank 0 has data -> scatter from it.
+                full_nvfp4 = (
+                    full_sd[param_name] if _accelerator.is_main_process else None
+                )
+                sharded_param = _broadcast_nvfp4_param(
+                    sharded_meta_param,
+                    full_nvfp4,
+                    _accelerator.is_main_process,
+                    device_mesh.device_type,
+                    nvfp4_cls,
+                )
+        elif (
+            ".experts." in param_name
+            and (".lora_A." in param_name or ".lora_B." in param_name)
+            and hasattr(sharded_meta_param, "device_mesh")
+            and dist.get_world_size(sharded_meta_param.device_mesh.get_group())
+            < dist.get_world_size()
+        ):
+            # EP×dp_shard/cp composition: the routed-expert LoRA adapter lives on a dp_shard SUBGROUP
+            # mesh and holds only THIS ep-group's E_local experts, but full_sd carries the GLOBAL
+            # (all-experts) adapter (4096-row lora_A etc.). The generic broadcast below would mismatch
+            # sizes (rank-0 sends the global tensor, receivers allocate the E_local size) and replicate
+            # rank-0's experts onto ranks owning a different ep-slice. Broadcast rank-0's global adapter
+            # (a consistent shape on every rank), then slice THIS rank's ep-experts + dp_shard and
+            # from_local. The ep slice is expert-aware (lora_B's experts are not contiguous in its flat
+            # r*E dim), so it mirrors shard_expert_lora rather than a plain chunk.
+            from torch.distributed.tensor import DTensor
+
+            from axolotl.integrations.expert_parallel.shard import (
+                ep_adapter_load_local_shard,
+            )
+
+            mesh = sharded_meta_param.device_mesh
+            placements = sharded_meta_param.placements
+            dp_size = mesh.size()
+            ep_size = dist.get_world_size() // dp_size
+            ep_dim = 0 if ".lora_A." in param_name else 1
+            dev = mesh.device_type
+            gshape = list(sharded_meta_param.size())
+            gshape[ep_dim] *= ep_size
+            if _accelerator.is_main_process:
+                g = full_tensor.to(dev)
+            else:
+                g = torch.empty(gshape, device=dev, dtype=sharded_meta_param.dtype)
+            dist.broadcast(g, src=0)
+            ep_coord = min(dist.get_process_group_ranks(mesh.get_group())) // dp_size
+            dp_rank = dist.get_group_rank(mesh.get_group(), dist.get_rank())
+            local = ep_adapter_load_local_shard(
+                g,
+                ep_dim,
+                model._ep_num_experts_global,
+                ep_coord,
+                ep_size,
+                placements,
+                dp_size,
+                dp_rank,
+            )
+            sharded_param = DTensor.from_local(local, mesh, placements, run_check=False)
+        elif hasattr(sharded_meta_param, "device_mesh"):
+            # Generic sharded params (the whole non-EP model, and EP composition's NON-expert weights)
+            # live on a FULL mesh that includes global rank 0, so distribute_tensor(src_data_rank=0)
+            # broadcasts rank-0's data and shards it correctly — including FSDP's padding for uneven
+            # sizes, which a manual chunk + from_local gets wrong (it builds a DTensor with the raw
+            # global size and mismatches the model param). EP-composition's rank-0-EXCLUDING subgroup
+            # params (experts, and the routed-expert LoRA adapter) are handled by their own branches
+            # above, so they never reach here.
+            from torch.distributed.tensor import distribute_tensor
+
+            device_mesh = sharded_meta_param.device_mesh
+            if _accelerator.is_main_process:
+                full_tensor = full_tensor.to(device_mesh.device_type)
+            else:
+                full_tensor = torch.empty(
+                    sharded_meta_param.size(),
+                    device=device_mesh.device_type,
+                    dtype=sharded_meta_param.dtype,
+                )
+            sharded_param = distribute_tensor(
+                full_tensor,
+                device_mesh,
+                sharded_meta_param.placements,
+                src_data_rank=0,
+            )
+            # Clone the local shard to allow full_tensor to be freed.
+            if (
+                sharded_param._local_tensor.untyped_storage().size()
+                > sharded_param._local_tensor.nelement()
+                * sharded_param._local_tensor.element_size()
+            ):
+                sharded_param = sharded_param.clone()
+        else:
+            # Non-sharded parameters
+            if _accelerator.is_main_process:
+                sharded_param = full_tensor.to(torch.device("cuda"))
+            else:
+                # broadcast manually
+                sharded_param = torch.empty_like(
+                    sharded_meta_param,
+                    device=torch.device("cuda"),
+                    dtype=sharded_meta_param.dtype,
+                )
+            dist.broadcast(sharded_param, src=0)
+
+        if offload_to_cpu:
+            sharded_param = sharded_param.cpu()
+
+        sharded_sd[param_name] = nn.Parameter(sharded_param)
+
+        del full_tensor
+        full_sd[param_name] = None
+
+    model.load_state_dict(sharded_sd, assign=True, strict=True)
+    end_time = time.time()
+    LOG.debug(
+        f"Time taken to load full state dict: {(end_time - start_time):.2f} seconds"
+    )
+    log_gpu_memory_usage(LOG, "Memory usage after broadcasting full state dict", 0)
+    return model
+
+
+def get_state_dict(self, model, unwrap=True):
+    """
+    Returns the state dictionary of a model sent through [`Accelerator.prepare`] potentially without full
+    precision.
+
+    Args:
+        model (`torch.nn.Module`):
+            A PyTorch model sent through [`Accelerator.prepare`]
+        unwrap (`bool`, *optional*, defaults to `True`):
+            Whether to return the original underlying state_dict of `model` or to return the wrapped state_dict
+
+    Returns:
+        `dict`: The state dictionary of the model potentially without full precision.
+
+    Example:
+
+    ```python
+    >>> import torch
+    >>> from accelerate import Accelerator
+
+    >>> accelerator = Accelerator()
+    >>> net = torch.nn.Linear(2, 2)
+    >>> net = accelerator.prepare(net)
+    >>> state_dict = accelerator.get_state_dict(net)
+    ```
+    """
+    from accelerate import DistributedType
+    from accelerate.utils import compare_versions
+
+    if self.distributed_type == DistributedType.DEEPSPEED:
+        zero3_sharding = self.deepspeed_config["zero_optimization"]["stage"] == 3
+        tp_sharding = (
+            self.deepspeed_config.get("tensor_parallel", {}).get("autotp_size", 0) > 1
+        )
+        if zero3_sharding or tp_sharding:
+            if model.zero_gather_16bit_weights_on_model_save():
+                if tp_sharding and not compare_versions("deepspeed", ">=", "0.16.4"):
+                    raise ImportError(
+                        "Deepspeed TP requires deepspeed >= 0.16.4, Please update DeepSpeed via `pip install deepspeed -U`."
+                    )
+                state_dict = (
+                    model._consolidated_16bit_state_dict()
+                    if tp_sharding
+                    else model._zero3_consolidated_16bit_state_dict()
+                )
+            else:
+                raise ValueError(
+                    "Cannot get 16bit model weights because `stage3_gather_16bit_weights_on_model_save` in DeepSpeed config is False. "
+                    "To save the model weights in 16bit, set `stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed config file or "
+                    "set `zero3_save_16bit_model` to True when using `accelerate config`. "
+                    "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights."
+                )
+        else:
+            from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
+
+            state_dict = clone_tensors_for_torch_save(
+                self.unwrap_model(model).state_dict()
+            )
+    elif self.is_fsdp2:
+        # https://github.com/pytorch/torchtune/blob/main/torchtune/training/_distributed.py#L465
+        from torch.distributed.tensor import DTensor
+
+        state_dict = {}
+        sharded_state_dict = model.state_dict()
+        is_rank_zero = torch.distributed.get_rank() == 0
+        for param_name, param in sharded_state_dict.items():
+            if param.is_cpu:
+                param = param.to(torch.device("cuda"))
+
+            if isinstance(param, DTensor):
+                param = param.full_tensor()
+
+            if is_rank_zero:
+                state_dict[param_name] = param.cpu()
+            # Drop the GPU-resident gathered tensor before the next iteration
+            # allocates the next one; otherwise the caching allocator holds
+            # both reservations and we accumulate ~model-size of VRAM.
+            del param
+            torch.distributed.barrier()
+
+        # Release the sharded view and force the allocator to give back the
+        # gather buffers.
+        del sharded_state_dict
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    elif self.distributed_type == DistributedType.FSDP:
+        from torch.distributed.fsdp import (
+            FullStateDictConfig,
+            FullyShardedDataParallel as FSDP,
+            StateDictType,
+        )
+
+        full_state_dict_config = FullStateDictConfig(
+            offload_to_cpu=True, rank0_only=True
+        )
+        with FSDP.state_dict_type(
+            model, StateDictType.FULL_STATE_DICT, full_state_dict_config
+        ):
+            state_dict = model.state_dict()
+    else:
+        if unwrap:
+            model = self.unwrap_model(model)
+        state_dict = model.state_dict()
+
+    return state_dict
+
+
+def patch_peft_param_wrapper_for_fsdp2():
+    """Patch PEFT's _LoraParameterProxy.forward for FSDP2 DTensor compatibility.
+
+    PEFT's ParamWrapper applies LoRA via torch.nn.utils.parametrize, which adds
+    delta_weight to the base weight W inside _LoraParameterProxy.forward().
+    Under FSDP2, W may be a DTensor (from FSDP unshard) while delta_weight is a
+    regular Tensor (or vice versa), causing a RuntimeError on mixed types.
+
+    This patch promotes the non-DTensor operand to match the DTensor's spec
+    using DTensor.from_local(), which is free for Replicate placement (just
+    metadata wrapping, no communication).
+    """
+    from peft.tuners.lora.layer import _LoraParameterProxy
+
+    if getattr(_LoraParameterProxy, "_axolotl_fsdp2_patched", False):
+        return
+
+    _original_forward = _LoraParameterProxy.forward
+
+    # NOTE: Replaces (not wraps) forward; assumes original is just `W + self.delta_weight`.
+    def _patched_forward(self, W):
+        from torch.distributed.tensor import DTensor
+
+        delta = self.delta_weight
+        w_is_dt = isinstance(W, DTensor)
+        d_is_dt = isinstance(delta, DTensor)
+
+        with torch.nn.utils.parametrize.cached():
+            if w_is_dt == d_is_dt:
+                return W + delta
+            if w_is_dt:
+                return W + DTensor.from_local(delta, W.device_mesh, W.placements)
+            return DTensor.from_local(W, delta.device_mesh, delta.placements) + delta
+
+    _LoraParameterProxy.forward = _patched_forward
+    _LoraParameterProxy._axolotl_fsdp2_patched = True
+    LOG.info("Patched PEFT _LoraParameterProxy.forward for FSDP2 DTensor compatibility")
+
+
+def _process_lora_module_for_fsdp(module, fsdp2_kwargs):
+    """Helper function to process LoRA modules for FSDP2."""
+    from peft.tuners.lora.layer import ParamWrapper
+    from torch.distributed.fsdp import fully_shard
+
+    # Skip ParamWrapper — its lora_A/B must not be independently sharded.
+    # The parent decoder layer's FSDP wrapper handles unsharding them.
+    # TODO: review if we even need to shard them separately in first place.
+    if isinstance(module, ParamWrapper):
+        return False
+
+    log_bias_dtype_mismatch = False
+
+    # Linear4Bit will keep it's bias term in fp32. If the weight dtype is in bf16 we are not able to
+    # wrap this. Therefore we must ensure the bias has the same dtype as the weight
+    if hasattr(module.base_layer, "bias") and module.base_layer.bias is not None:
+        if module.base_layer.weight.dtype != module.base_layer.bias.dtype:
+            log_bias_dtype_mismatch = True
+            module.base_layer.bias.data = module.base_layer.bias.data.to(
+                module.base_layer.weight.dtype
+            )
+
+    for active_adapter in module.active_adapters:
+        if module.lora_A:
+            fully_shard(module.lora_A[active_adapter], **fsdp2_kwargs)
+        if module.lora_B:
+            fully_shard(module.lora_B[active_adapter], **fsdp2_kwargs)
+        if module.lora_magnitude_vector:
+            fully_shard(module.lora_magnitude_vector[active_adapter], **fsdp2_kwargs)
+
+    # lora_embedding_A/B are ParameterDicts containing nn.Parameter (Tensors),
+    # not nn.Module. fully_shard() only accepts nn.Module, so we cannot shard
+    # individual embedding Parameters. Instead, shard the entire LoraLayer module. fully_shard() can be used hierarchically because it does not
+    # override groups already assigned by fully_shard(), so modules
+    # where fully_shard() was already called are not affected [see https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html]
+    if module.lora_embedding_A or module.lora_embedding_B:
+        from torch.distributed.fsdp import FSDPModule
+
+        if not isinstance(module, FSDPModule):
+            fully_shard(module, **fsdp2_kwargs)
+
+    return log_bias_dtype_mismatch
+
+
+def fsdp2_prepare_model(accelerator, model: torch.nn.Module) -> torch.nn.Module:
+    """Prepares the model for FSDP2 in-place. Also returns the model to avoid misuse of the original model.
+
+    Args:
+        accelerator (`Accelerator`): The accelerator instance
+        model (`torch.nn.Module`): The model to prepare
+
+    Returns:
+        `torch.nn.Module`: Prepared model
+    """
+    from accelerate.utils import get_module_children_bottom_up, is_compiled_module
+    from accelerate.utils.fsdp_utils import fsdp2_prepare_auto_wrap_policy
+    from accelerate.utils.modeling import get_non_persistent_buffers
+    from peft import PeftModel
+    from peft.tuners.lora import LoraLayer
+    from torch.distributed.fsdp import (
+        CPUOffloadPolicy,
+        FSDPModule,
+        MixedPrecisionPolicy,
+        fully_shard,
+    )
+
+    is_type_fsdp = isinstance(model, FSDPModule) or (
+        is_compiled_module(model) and isinstance(model._orig_mod, FSDPModule)
+    )
+    if is_type_fsdp:
+        return model
+
+    fsdp2_plugin = accelerator.state.fsdp_plugin
+
+    original_sd = model.state_dict()
+
+    from torch.distributed.fsdp.wrap import (
+        size_based_auto_wrap_policy,
+        transformer_auto_wrap_policy,
+    )
+
+    # We need the `auto_wrap_policy` original type to create a custom poilicy function for sharding
+    # This is because `fully_shard` doesn't support old auto wrap policies, rather we have to imitate the behaviour
+    if fsdp2_plugin.auto_wrap_policy is transformer_auto_wrap_policy:
+        pass  # auto_wrap_policy_type = "transformer"
+    elif fsdp2_plugin.auto_wrap_policy is size_based_auto_wrap_policy:
+        pass  # auto_wrap_policy_type = "size"
+
+    # We set `auto_wrap_policy` to `functools.partial` to avoid creating it again
+    # This is because of `apply_activation_checkpointing` which will can reuse this function
+    fsdp2_plugin.set_auto_wrap_policy(model)
+
+    if fsdp2_plugin.activation_checkpointing:
+        from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+            CheckpointImpl,
+            apply_activation_checkpointing,
+            checkpoint_wrapper,
+        )
+
+        # Apply activation checkpointing before applying `fully_shard`
+        apply_activation_checkpointing(
+            model,
+            checkpoint_wrapper_fn=functools.partial(
+                checkpoint_wrapper,
+                checkpoint_impl=CheckpointImpl.NO_REENTRANT,
+            ),
+            auto_wrap_policy=fsdp2_plugin.auto_wrap_policy,
+        )
+
+    mesh = getattr(accelerator.state, "device_mesh", None)
+
+    # Disable memory pinning if requested
+    offload_to_cpu = isinstance(fsdp2_plugin.cpu_offload, CPUOffloadPolicy)
+    if (
+        offload_to_cpu
+        and os.environ.get("FSDP_CPU_OFFLOAD_PIN_MEMORY", "").lower() == "false"
+    ):
+        fsdp2_plugin.cpu_offload.pin_memory = False
+
+    fsdp2_kwargs = {
+        "reshard_after_forward": fsdp2_plugin.reshard_after_forward,
+        "offload_policy": fsdp2_plugin.cpu_offload,
+        # `fully_shard` doesn't accept `None` in case of `MixedPrecisionPolicy`
+        "mp_policy": fsdp2_plugin.mixed_precision_policy or MixedPrecisionPolicy(),
+        "mesh": (
+            mesh[tuple(accelerator.state.parallelism_config.fsdp_dim_names)]
+            if mesh is not None
+            else None
+        ),
+    }
+    model_has_params4bit = False
+    for _, param in model.named_parameters():
+        # this is a temporary fix whereby loading models with bnb params cannot be moved from
+        # GPU to a meta device due with FSDP2 because torch operations don't return the original class type
+        # bypassing the move to meta will still cause the VRAM spike, but at least it still will load
+        if param.__class__.__name__ == "Params4bit":
+            model_has_params4bit = True
+            break
+
+    if fsdp2_plugin.cpu_ram_efficient_loading and not model_has_params4bit:
+        # Context: `fully_shard` moves the model to GPU if it was on CPU, however it can also be on `meta` and then it stays there even after `fully_shard`
+        # For this reason, we need to move the model to `meta` device, as then sharding happens on `meta` device
+        # If we kept the model on CPU (`cpu_ram_efficient_loading` has model be on CPU on all ranks, though non-main ranks only have `torch.emtpy`), `fully_shard` would move it to GPU
+        # Afterwards, when we call `fsdp2_load_full_state_dict`, us creating the state_dict would result into briefly having two copies of model state_dict on the GPU -> VRAM spike
+
+        # We need to keep the original non-persistent buffers, as those MAY not be in the state_dict, resulting in them staying on meta device
+        # Also, these buffers aren't getting sharded by default
+        # We get the FQNs of all non-persistent buffers, to re-register them after
+        non_persistent_buffer_fqns = get_non_persistent_buffers(
+            model, recurse=True, fqns=True
+        )
+        original_non_persistent_buffers = copy.deepcopy(
+            {k: v for k, v in model.named_buffers() if k in non_persistent_buffer_fqns}
+        )
+        # We move the model to meta device, as then sharding happens on meta device
+        model = model.to(torch.device("meta"))
+        # We need to re-tie the weights, not exactly sure why, but if we don't do this, reference to `lm_head/embed_tokens` stay hanging -> more VRAM usage
+        # We assume `transformers` models have a `tie_weights` method if they support it
+        if hasattr(model, "tie_weights"):
+            model.tie_weights()
+
+    is_peft_model = isinstance(model, PeftModel)
+
+    # Patch PEFT's _LoraParameterProxy for DTensor compatibility if any
+    # ParamWrapper modules exist (used for target_parameters / 3D expert params).
+    if is_peft_model:
+        from peft.tuners.lora.layer import ParamWrapper
+
+        if any(isinstance(m, ParamWrapper) for m in model.modules()):
+            patch_peft_param_wrapper_for_fsdp2()
+
+    # EP+FSDP: pre-wrap experts on `dp_shard` before the outer auto-wrap so
+    # the walker skips them. See `expert_parallel/README.md`.
+    # EP composition shards the experts on the non-ep axis so they don't sit replicated per rank. With
+    # dp_shard that's the data axis; with pure EP×cp (no dp_shard) the cp ranks of an ep-group hold the
+    # SAME experts (cp shards the sequence, not the experts), so FSDP-shard them on cp and let the MoE
+    # forward all-gather — otherwise each rank keeps its full ep-group slice (e.g. 64 experts at ep=4)
+    # and OOMs on the first forward's all-gather.
+    from axolotl.integrations.expert_parallel.plugin import expert_shard_axis
+
+    _ep_shard_axis = expert_shard_axis(
+        getattr(mesh, "mesh_dim_names", None) if mesh is not None else None
+    )
+    if _ep_shard_axis is not None:
+        from axolotl.integrations.expert_parallel.plugin import ExpertParallelPlugin
+        from axolotl.integrations.expert_parallel.shard import shard_expert_lora
+
+        # Realign target_parameters expert LoRA with the EP-sharded weights before
+        # FSDP wraps the experts (PEFT sized it for the global expert count). Stash the
+        # exact ep group so the save path gathers across the same axis (re-resolving at
+        # save time can pick up a stale/size-1 mesh).
+        model._ep_lora_group = mesh["ep"].get_group()
+        shard_expert_lora(model, mesh["ep"].size())
+        ExpertParallelPlugin.fully_shard_experts(
+            model, mesh[_ep_shard_axis], fsdp2_kwargs
+        )
+    elif getattr(model, "_ddp_params_and_buffers_to_ignore", None):
+        # Pure EP (ep_size == world_size): no ep×dp_shard mesh is built, so the experts were
+        # manually EP-sharded (shard_expert_weights scattered each rank's real [E_local] slice) but
+        # there is NO dp_shard axis to FSDP them onto. Left to the outer fully_shard(mesh=None) they
+        # would be re-sharded on the flat world mesh and replicated from rank 0, destroying EP. Exclude
+        # the EP-sharded expert base params from the FSDP wrap entirely so each rank keeps its own
+        # [E_local] slice as a plain param; fsdp2_load_full_state_dict restores them per-rank.
+        ep_ignored = {
+            p
+            for n, p in model.named_parameters()
+            if ".experts." in n
+            and ".shared_experts." not in n
+            and n.rsplit(".", 1)[-1]
+            in ("gate_up_proj", "down_proj", "gate_up_proj_bias", "down_proj_bias")
+        }
+        if ep_ignored:
+            fsdp2_kwargs["ignored_params"] = (
+                set(fsdp2_kwargs.get("ignored_params") or set()) | ep_ignored
+            )
+            LOG.info(
+                f"expert_parallel (pure EP): excluded {len(ep_ignored)} EP-sharded expert "
+                "param(s) from the FSDP wrap (kept as plain per-rank slices)."
+            )
+
+    auto_wrap_policy = fsdp2_prepare_auto_wrap_policy(fsdp2_plugin, model)
+    log_bias_dtype_mismatch = False
+    fp32_norm_patterns = get_fp32_norm_patterns(model)
+    if fp32_norm_patterns:
+        shard_norms_fp32(
+            model,
+            patterns=fp32_norm_patterns,
+            fully_shard_kwargs=fsdp2_kwargs,
+        )
+
+    # Pre-quantized / mixed-dtype models (e.g. an NVFP4 checkpoint loaded for LoRA) carry
+    # non-float Parameters and keep-fp32 modules that the generic FSDP2 path can't shard
+    # uniformly. Engage the quantized capability path ONLY when such params exist; pure-bf16
+    # models take the original generic path unchanged. The nonfloat ``nn.Parameter.__new__``
+    # patch is process-global and is restored via the context manager's ``finally`` (an
+    # exception in ``fully_shard`` can no longer leave it patched).
+    from axolotl.monkeypatch.accelerate.fsdp2_quantized import (
+        cast_residual_fp32,
+        model_has_float_logical_quantized_params,
+        model_has_nonfloat_params,
+        nonfloat_param_guard,
+        shard_fp32_modules,
+    )
+
+    # Apply the quantized dtype/cast/sharding policy ONLY for float-logical torchao subclasses
+    # (NVFP4Tensor/Float8Tensor/MXTensor) — the pre-quantized checkpoint case this path is for.
+    # Plain bnb Params4bit QLoRA is excluded so cast_residual_fp32 does not downcast its fp32 LoRA.
+    # The nn.Parameter.__new__ guard is separate: it is needed for ANY plain non-float param (uint8
+    # packed, which includes bnb Params4bit) and stays gated on that.
+    _quantized = model_has_float_logical_quantized_params(model)
+    _needs_nonfloat_guard = model_has_nonfloat_params(model)
+    _guard = (
+        nonfloat_param_guard(model)
+        if _needs_nonfloat_guard
+        else contextlib.nullcontext()
+    )
+    with _guard:
+        if _quantized:
+            # keep-fp32 modules (registered by model adapters, e.g. DSV4 mHC) get their own
+            # fp32 shard group; remaining plain fp32 (PEFT LoRA) is cast to the compute dtype.
+            shard_fp32_modules(model, fsdp2_kwargs)
+            cast_residual_fp32(model)
+
+        if auto_wrap_policy is not None:
+            for module in get_module_children_bottom_up(model)[:-1]:
+                if is_peft_model and isinstance(module, LoraLayer):
+                    module_log_bias_mismatch = _process_lora_module_for_fsdp(
+                        module, fsdp2_kwargs
+                    )
+                    log_bias_dtype_mismatch |= module_log_bias_mismatch
+                if auto_wrap_policy(module) and not isinstance(module, FSDPModule):
+                    fully_shard(module, **fsdp2_kwargs)
+
+        fully_shard(model, **fsdp2_kwargs)
+
+    if log_bias_dtype_mismatch:
+        LOG.warning(
+            "Bias dtype mismatch detected in LoRA base linear layer. Bias parameters have been cast to weight dtype."
+        )
+
+    if fsdp2_plugin.cpu_ram_efficient_loading:
+        fsdp2_load_full_state_dict(
+            accelerator, model, original_sd, offload_to_cpu=offload_to_cpu
+        )
+
+    if fsdp2_plugin.cpu_ram_efficient_loading and not model_has_params4bit:
+        # We re-register the buffers, as they may not be in the state_dict
+        for fqn, buffer_tensor in original_non_persistent_buffers.items():
+            buffer_tensor = buffer_tensor.to(accelerator.device)
+
+            if "." in fqn:
+                parent_fqn, local_buffer_name = fqn.rsplit(".", 1)
+                parent_module = model.get_submodule(parent_fqn)
+            else:
+                local_buffer_name = fqn
+                parent_module = model
+
+            parent_module.register_buffer(
+                local_buffer_name, buffer_tensor, persistent=False
+            )
+
+        # We need to tie the weights again, as call to `load_full_state_dict` breaks the tie
+        # Needs to be called both here and above
+        # removing this call makes the have slightly different loss
+        # removing the call above leads to extra memory usage as explained in the comment above
+        if hasattr(model, "tie_weights"):
+            model.tie_weights()
+    return model
+
+
+def patch_tied_keys_for_meta_device():
+    """Patch _adjust_tied_keys_with_tied_pointers to skip meta tensors.
+
+    Meta tensors all share data_ptr()==0, causing every parameter to be incorrectly
+    grouped as "tied". Skipping them is safe since they have no real storage.
+    """
+    from collections import defaultdict
+
+    from transformers import PreTrainedModel
+
+    # recent transformers replaced data_ptr tie detection (_adjust_tied_keys_with_tied_pointers) with
+    # config-driven get_expanded_tied_weights_keys; absent on the pinned version, so bail (nothing to patch).
+    if not hasattr(PreTrainedModel, "_adjust_tied_keys_with_tied_pointers"):
+        return
+
+    def _patched_adjust_tied_keys_with_tied_pointers(self, missing_keys):
+        param_pointers = defaultdict(list)
+        for param_name, param_value in self.state_dict().items():
+            if param_value.is_meta:
+                continue
+            param_pointers[param_value.data_ptr()].append(param_name)
+
+        tied_param_names = [
+            names
+            for names in param_pointers.values()
+            if len(names) > 1
+            and not any(name in self.all_tied_weights_keys.keys() for name in names)
+            and not all(name in missing_keys for name in names)
+        ]
+
+        tied_weights_keys_by_pointers = {
+            param_name: group[0]
+            for group in tied_param_names
+            for param_name in group[1:]
+        }
+        self.all_tied_weights_keys.update(tied_weights_keys_by_pointers)
+
+    PreTrainedModel._adjust_tied_keys_with_tied_pointers = (
+        _patched_adjust_tied_keys_with_tied_pointers
+    )
+
+
+def patch_initialize_missing_keys_for_fsdp():
+    """Patch _initialize_missing_keys to skip re-initialization on FSDP non-rank-0.
+
+    When using cpu_ram_efficient_loading, non-rank-0 processes load weights on
+    meta device and move them to CPU as empty tensors. Without this patch,
+    initialize_weights() re-initializes ALL parameters (via guarded init
+    functions), which is slow and uses extra RAM per process.
+
+    The fix marks all params/buffers with _is_hf_initialized=True before calling
+    the original method, so guarded init functions (init.normal_, init.zeros_,
+    etc.) become no-ops on non-rank-0 processes. The real weights arrive later
+    via FSDP broadcast from rank 0.
+
+    Upstream fix: https://github.com/huggingface/transformers/pull/44473
+    Remove this patch once transformers includes the fix in a stable release.
+    """
+    from transformers import PreTrainedModel
+    from transformers.modeling_utils import is_fsdp_enabled, is_local_dist_rank_0
+
+    if getattr(PreTrainedModel._initialize_missing_keys, "_axolotl_patched", False):
+        return
+
+    _original_initialize_missing_keys = PreTrainedModel._initialize_missing_keys
+
+    def _patched_initialize_missing_keys(self, is_quantized: bool) -> None:
+        if is_fsdp_enabled() and not is_local_dist_rank_0():
+            for key in self.state_dict():
+                try:
+                    param_or_buffer = self.get_parameter_or_buffer(key)
+                    param_or_buffer._is_hf_initialized = True
+                except AttributeError:
+                    pass  # may happen when handling pre-quantized weights
+            self._is_hf_initialized = True
+
+        _original_initialize_missing_keys(self, is_quantized)
+
+    PreTrainedModel._initialize_missing_keys = _patched_initialize_missing_keys
+    PreTrainedModel._initialize_missing_keys._axolotl_patched = True
+
+
+def patch_move_missing_keys_meta_for_fsdp():
+    """Stop transformers materializing the FULL model on every non-rank-0 rank during load.
+
+    ``_move_missing_keys_from_meta_to_device`` has a branch ``is_fsdp_enabled() and not
+    is_local_dist_rank_0() and not is_quantized`` that moves EVERY meta parameter to real CPU
+    storage (``torch.zeros_like(param, device="cpu")``). For an unrecognized-quantizer checkpoint
+    (NVFP4-modelopt → ``is_quantized=False``) on a large model, that puts the whole model on each
+    non-rank-0 rank → ``world_size``× CPU RAM → OOM. axolotl's ``fsdp2_prepare_model`` immediately
+    does ``model.to("meta")`` then broadcasts rank 0's weights, so the materialized params are
+    pure waste. Keep params on meta (FSDP fills them); only buffers — computed in ``__init__`` and
+    not broadcast — get real CPU storage.
+
+    Caller MUST restrict this to frozen-base (adapter) runs: leaving base params on meta on
+    non-rank-0 deadlocks the FSDP2 optimizer-state all-gather at checkpoint save for a FULL
+    fine-tune (rank-0 real DTensors vs non-rank-0 meta). LoRA/qLoRA carry no base optimizer state,
+    so the gather never touches these params."""
+    from transformers import PreTrainedModel
+    from transformers.integrations import (
+        is_deepspeed_zero3_enabled,
+        is_fsdp_enabled,
+    )
+    from transformers.modeling_utils import (
+        _load_parameter_into_model,
+        get_device,
+        is_local_dist_rank_0,
+    )
+
+    if getattr(
+        PreTrainedModel._move_missing_keys_from_meta_to_device,
+        "_axolotl_patched",
+        False,
+    ):
+        return
+
+    def _patched_move_missing_keys(
+        self, missing_keys, device_map, device_mesh, hf_quantizer
+    ):
+        is_quantized = hf_quantizer is not None
+        if is_deepspeed_zero3_enabled() and not is_quantized:
+            return
+
+        if is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized:
+            # Params: leave on meta — FSDP broadcasts rank 0's real weights into them. (Upstream
+            # materialized them all to cpu zeros here, OOMing large models.) Buffers still need real
+            # storage, but they are small.
+            for key, buffer in self.named_buffers():
+                if buffer.is_meta:
+                    value = torch.zeros_like(buffer, device="cpu")
+                    _load_parameter_into_model(self, key, value)
+            return
+
+        for key in missing_keys - self.all_tied_weights_keys.keys():
+            param = self.get_parameter_or_buffer(key)
+            param_device = get_device(device_map, key, valid_torch_device=True)
+            value = torch.empty_like(param, device=param_device)
+            if device_mesh is not None:
+                from transformers.modeling_utils import shard_and_distribute_module
+
+                shard_and_distribute_module(
+                    self,
+                    value,
+                    param,
+                    key,
+                    None,
+                    False,
+                    device_mesh.get_local_rank(),
+                    device_mesh,
+                )
+            else:
+                _load_parameter_into_model(self, key, value)
+        for key, buffer in self.named_non_persistent_buffers():
+            buffer_device = get_device(device_map, key, valid_torch_device=True)
+            value = torch.empty_like(buffer, device=buffer_device)
+            _load_parameter_into_model(self, key, value)
+
+    PreTrainedModel._move_missing_keys_from_meta_to_device = _patched_move_missing_keys
+    PreTrainedModel._move_missing_keys_from_meta_to_device._axolotl_patched = True
+    LOG.info(
+        "Patched transformers _move_missing_keys_from_meta_to_device: non-rank-0 params stay "
+        "on meta (FSDP broadcast fills them) instead of full-model CPU materialization"
+    )
+
+
+def patch_accelerate_fsdp2():
+    import accelerate
+
+    accelerate.accelerator.fsdp2_prepare_model = fsdp2_prepare_model
+    accelerate.Accelerator.get_state_dict = get_state_dict
+    setattr(
+        sys.modules["accelerate"],
+        "Accelerator.get_state_dict",
+        get_state_dict,
+    )
diff --git a/src/axolotl/monkeypatch/accelerate/fsdp2_quantized.py b/src/axolotl/monkeypatch/accelerate/fsdp2_quantized.py
new file mode 100644
index 0000000000..3144e4cbc7
--- /dev/null
+++ b/src/axolotl/monkeypatch/accelerate/fsdp2_quantized.py
@@ -0,0 +1,157 @@
+"""Capability helpers for sharding pre-quantized / mixed-dtype models under FSDP2.
+
+These are only engaged when a model actually carries non-float (quantized) Parameters — e.g. a
+pre-quantized NVFP4 checkpoint loaded for LoRA. They are model-agnostic: model families register
+the class names of modules that must stay fp32 (e.g. DeepSeek-V4 mHC) via
+:func:`register_fp32_shard_classes`. Pure-bf16 models never touch this path.
+"""
+
+from __future__ import annotations
+
+import contextlib
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# Class names of modules that must be sharded in their own fp32 group (registered by adapters).
+_FP32_SHARD_CLASS_NAMES: set[str] = set()
+
+# Tensor-subclass names that are quantized but report a logical FLOAT dtype (so torch.is_floating_point
+# is True) — e.g. torchao NVFP4Tensor / Float8Tensor / MXTensor. These still need the quantized
+# FSDP2 dtype/sharding policy even though they're "floating point".
+_QUANT_TENSOR_CLASS_NAMES: set[str] = {"NVFP4Tensor", "Float8Tensor", "MXTensor"}
+
+
+def register_fp32_shard_classes(names) -> None:
+    """Register module class names that the quantized FSDP2 path keeps in a separate fp32 shard."""
+    _FP32_SHARD_CLASS_NAMES.update(names)
+
+
+def register_quantized_tensor_classes(names) -> None:
+    """Register quantized tensor-subclass names (float-logical, e.g. NVFP4Tensor) for detection."""
+    _QUANT_TENSOR_CLASS_NAMES.update(names)
+
+
+def model_has_nonfloat_params(model) -> bool:
+    """True iff the model carries PLAIN non-float Parameters (e.g. uint8 packed) — the case that
+    needs the nn.Parameter.__new__ requires_grad guard during sharding."""
+    return any(not torch.is_floating_point(p) for p in model.parameters())
+
+
+def _is_quantized_param(p) -> bool:
+    if not torch.is_floating_point(p):
+        return True  # plain packed (uint8) quantized data
+    # torchao tensor subclasses report a logical float dtype; detect by class name (no hard dep).
+    if type(p).__name__ in _QUANT_TENSOR_CLASS_NAMES:
+        return True
+    data = getattr(p, "data", None)
+    return data is not None and type(data).__name__ in _QUANT_TENSOR_CLASS_NAMES
+
+
+def model_has_quantized_params(model) -> bool:
+    """Capability check for the quantized dtype/sharding policy: plain non-float params OR float-
+    logical quantized tensor subclasses (NVFP4Tensor/Float8Tensor/MXTensor)."""
+    return any(_is_quantized_param(p) for p in model.parameters())
+
+
+def _is_float_logical_quantized_param(p) -> bool:
+    """True only for torchao float-logical quantized subclasses (NVFP4Tensor/Float8Tensor/MXTensor)
+    — the ones the dtype/cast/sharding policy below is designed for. Deliberately excludes plain
+    non-float params (e.g. bnb Params4bit uint8): standard QLoRA+FSDP2 keeps its fp32 LoRA and must
+    not be downcast by cast_residual_fp32."""
+    if type(p).__name__ in _QUANT_TENSOR_CLASS_NAMES:
+        return True
+    data = getattr(p, "data", None)
+    return data is not None and type(data).__name__ in _QUANT_TENSOR_CLASS_NAMES
+
+
+def model_has_float_logical_quantized_params(model) -> bool:
+    """Capability check for the dtype/cast/sharding policy: ONLY float-logical quantized tensor
+    subclasses (the pre-quantized NVFP4/Float8 checkpoint case). Plain bnb Params4bit QLoRA is
+    excluded so its fp32 LoRA adapters are not silently cast to the compute dtype."""
+    return any(_is_float_logical_quantized_param(p) for p in model.parameters())
+
+
+@contextlib.contextmanager
+def nonfloat_param_guard(model):
+    """Freeze existing non-float params and make ``nn.Parameter`` construction default to
+    ``requires_grad=False`` for non-float data for the duration of sharding.
+
+    FSDP2's sharded ``nn.Parameter()`` defaults ``requires_grad=True``, which errors on non-float
+    (uint8) data *before* it copies the original flag. The ``nn.Parameter.__new__`` patch is a
+    PROCESS-GLOBAL monkeypatch, so it is restored in a ``finally`` — an exception during
+    ``fully_shard`` can no longer leave the process globally patched.
+    """
+    import torch.nn as nn
+
+    frozen = 0
+    for p in model.parameters():
+        if not torch.is_floating_point(p) and p.requires_grad:
+            p.requires_grad_(False)
+            frozen += 1
+    if frozen:
+        LOG.info(
+            "fsdp2-quant: froze %d non-float (quantized) params before sharding", frozen
+        )
+
+    orig_new = nn.Parameter.__new__
+
+    def _nonfloat_safe_new(cls, data=None, requires_grad=True):
+        if data is not None and not torch.is_floating_point(data):
+            requires_grad = False
+        return orig_new(cls, data, requires_grad)
+
+    nn.Parameter.__new__ = _nonfloat_safe_new
+    try:
+        yield
+    finally:
+        nn.Parameter.__new__ = orig_new
+
+
+def shard_fp32_modules(model, fsdp2_kwargs, compute_dtype=torch.bfloat16) -> int:
+    """Shard registered keep-fp32 modules (e.g. DSV4 mHC) in their own fp32 group: compute in fp32
+    (cast inputs up) but emit ``compute_dtype`` outputs so they don't feed fp32 activations into
+    downstream low-precision layers. No-op if no classes are registered."""
+    if not _FP32_SHARD_CLASS_NAMES:
+        return 0
+    from torch.distributed.fsdp import (
+        FSDPModule,
+        MixedPrecisionPolicy,
+        fully_shard,
+    )
+
+    fp32_mp = MixedPrecisionPolicy(
+        param_dtype=torch.float32,
+        reduce_dtype=torch.float32,
+        output_dtype=compute_dtype,
+        cast_forward_inputs=True,
+    )
+    kwargs = {**fsdp2_kwargs, "mp_policy": fp32_mp}
+    n = 0
+    for m in model.modules():
+        if type(m).__name__ in _FP32_SHARD_CLASS_NAMES and not isinstance(
+            m, FSDPModule
+        ):
+            fully_shard(m, **kwargs)
+            n += 1
+    if n:
+        LOG.info("fsdp2-quant: fp32-sharded %d keep-fp32 modules separately", n)
+    return n
+
+
+def cast_residual_fp32(model, compute_dtype=torch.bfloat16) -> int:
+    """Cast remaining plain fp32 float params (e.g. PEFT-upcast LoRA) to ``compute_dtype`` for a
+    uniform per-shard original dtype. Skips DTensors (already-sharded fp32 groups)."""
+    from torch.distributed.tensor import DTensor
+
+    n = 0
+    for p in model.parameters():
+        if p.dtype == torch.float32 and not isinstance(p.data, DTensor):
+            p.data = p.data.to(compute_dtype)
+            n += 1
+    if n:
+        LOG.info("fsdp2-quant: cast %d residual fp32 params to %s", n, compute_dtype)
+    return n
diff --git a/src/axolotl/monkeypatch/accelerate/parallelism_config.py b/src/axolotl/monkeypatch/accelerate/parallelism_config.py
new file mode 100644
index 0000000000..dd67e5f01c
--- /dev/null
+++ b/src/axolotl/monkeypatch/accelerate/parallelism_config.py
@@ -0,0 +1,395 @@
+"""ParallelismConfig monkeypatch.
+
+Two extensions:
+- Allow pure CP standalone via `ACCELERATE_ALLOW_CP_STANDALONE`.
+- Add Expert Parallel (`ep`) as a first-class mesh axis inside the
+  data-parallel group. Mesh order is `(ep, dp_replicate, dp_shard, cp, sp, tp)`
+  so the dp axes stay contiguous (required for `_flatten("dp")`).
+
+See `expert_parallel/README.md` for the full integration story.
+"""
+
+import os
+import warnings
+
+from accelerate import DistributedType
+
+
+def _patched_post_init(self):
+    _ORIG_POST_INIT(self)
+
+    if not hasattr(self, "ep_size") or self.ep_size is None:
+        self.ep_size = int(os.environ.get("PARALLELISM_CONFIG_EP_SIZE", "1") or 1)
+    if self.ep_size < 1:
+        raise ValueError(f"ep_size must be at least 1, got {self.ep_size}")
+
+    # Register so `_set_size`, `_validate_accelerator`, `_get_mesh` see it.
+    self._sizes["ep"] = self.ep_size
+
+
+def _patched_total_size(self):
+    return (
+        self.dp_replicate_size
+        * self.dp_shard_size
+        * self.tp_size
+        * self.cp_size
+        * self.sp_size
+        * getattr(self, "ep_size", 1)
+    )
+
+
+def _patched_ep_enabled(self):
+    return getattr(self, "ep_size", 1) > 1
+
+
+def _patched_dp_dim_names(self):
+    """DP axes (different ranks see different data). EP is included — each
+    EP rank pulls its own batch."""
+    dims = []
+    if self.ep_enabled:
+        dims += ["ep"]
+    if self.dp_replicate_enabled:
+        dims += ["dp_replicate"]
+    if self.dp_shard_enabled:
+        dims += ["dp_shard"]
+    return dims
+
+
+def _patched_dp_shard_cp_dim_names(self):
+    """Axes the outer FSDP wrap shards along (flattened into `dp_shard_cp`).
+    Including `ep` makes non-expert grads reduce-scatter across the full
+    world; experts are pre-wrapped on `mesh["dp_shard"]` only and skipped
+    by the auto-wrap walker."""
+    dims = []
+    if self.ep_enabled:
+        dims += ["ep"]
+    if self.dp_shard_enabled:
+        dims += ["dp_shard"]
+    if self.cp_enabled:
+        dims += ["cp"]
+    return dims
+
+
+def _patched_non_dp_dim_names(self):
+    """Non-DP axes (TP/CP/SP). EP moved into `dp_dim_names`."""
+    dims = []
+    if self.tp_enabled:
+        dims += ["tp"]
+    if self.cp_enabled:
+        dims += ["cp"]
+    if self.sp_enabled:
+        dims += ["sp"]
+    return dims
+
+
+def _patched_get_mesh(self):
+    """Build (dim_names, shape) for `init_device_mesh`. Order keeps the dp
+    block (ep, dp_replicate, dp_shard) contiguous so `_flatten("dp")` works.
+    """
+    mesh_dims = {p: self._sizes[p] for p in self.active_mesh_dims}
+    mesh_order = ["ep", "dp_replicate", "dp_shard", "cp", "sp", "tp"]
+    sorted_items = sorted(mesh_dims.items(), key=lambda x: mesh_order.index(x[0]))
+    return tuple(zip(*sorted_items, strict=True))
+
+
+def _validate_accelerator(self, accelerator):
+    _warnings = set()
+    if not accelerator.multi_device and self.total_size == 1:
+        # No distributed setup, valid parallelism config
+        return
+
+    # We need this to ensure DDP works
+    if self.total_size == 1:
+        self._set_size("dp_replicate", accelerator.num_processes)
+
+    # DeepSpeed manages SP process groups globally, so total_size (the local parallelism config)
+    # need not equal num_processes; keep this branch in sync with accelerate's upstream validator.
+    if self.sp_backend == "deepspeed" and self.sp_size > 1:
+        pass
+    elif self.total_size != accelerator.num_processes:
+        raise ValueError(
+            f"ParallelismConfig total_size ({self.total_size}) does not match "
+            f"num_processes ({accelerator.num_processes}). Please adjust dp_replicate_size/ "
+            f"dp_shard_size/tp_size/cp_size/sp_size/ep_size."
+        )
+
+    # allow parallelism config when not using fsdp if using pure context parallelism
+    allow_parallelism_config = False
+
+    if (
+        self.cp_size > 1
+        and self.dp_shard_size <= 1
+        and os.environ.get("ACCELERATE_ALLOW_CP_STANDALONE", "false").lower() == "true"
+    ):
+        allow_parallelism_config = True
+
+    # Pure EP (no FSDP/TP/CP) is valid: the plugin handles dispatch/combine
+    # and DDP's _ddp_params_and_buffers_to_ignore keeps experts out of DDP.
+    if (
+        getattr(self, "ep_enabled", False)
+        and self.dp_shard_size <= 1
+        and self.tp_size <= 1
+        and self.cp_size <= 1
+    ):
+        allow_parallelism_config = True
+
+    if (
+        self.total_size > 1
+        and not allow_parallelism_config
+        and not (
+            accelerator.is_fsdp2
+            or accelerator.multi_device
+            or accelerator.distributed_type == DistributedType.DEEPSPEED
+        )
+    ):
+        raise ValueError(
+            f"ParallelismConfig is only compatible with DistributedType.FSDP (version 2), DistributedType.Multi{{Device}}, or DistributedType.DEEPSPEED, but got {accelerator.distributed_type}."
+        )
+
+    for parallelism, size in self._sizes.items():
+        if size == 1 and getattr(self, f"{parallelism}_handler", None) is not None:
+            _warnings.add(
+                f"ParallelismConfig.{parallelism}_handler is set, but {parallelism}_size is set to 1. This handler will be ignored."
+            )
+
+    if _warnings and accelerator.is_main_process:
+        warnings.warn(
+            "ParallelismConfig has the following warnings:\n" + "\n".join(_warnings),
+            UserWarning,
+            stacklevel=2,
+        )
+
+
+def patched_is_fsdp2(self) -> bool:
+    """
+    Patched version of is_fsdp2 that guards against a None fsdp_plugin.
+    """
+    # The new logic checks if fsdp_plugin exists before accessing its attributes
+    return (
+        self.distributed_type == DistributedType.FSDP
+        and self.fsdp_plugin
+        and self.fsdp_plugin.fsdp_version == 2
+    )
+
+
+# Captured in `patch_parallelism_config()` so we can chain the original
+# __post_init__ before adding ep.
+_ORIG_POST_INIT = None
+
+
+def patch_parallelism_config():
+    global _ORIG_POST_INIT
+    from accelerate.accelerator import AcceleratorState, ParallelismConfig
+
+    if _ORIG_POST_INIT is None:
+        _ORIG_POST_INIT = ParallelismConfig.__post_init__
+
+    ParallelismConfig.__post_init__ = _patched_post_init
+    # `total_size` is a property on the dataclass; replace it.
+    ParallelismConfig.total_size = property(_patched_total_size)
+    ParallelismConfig.ep_enabled = property(_patched_ep_enabled)
+    ParallelismConfig.dp_dim_names = property(_patched_dp_dim_names)
+    ParallelismConfig.dp_shard_cp_dim_names = property(_patched_dp_shard_cp_dim_names)
+    ParallelismConfig.non_dp_dim_names = property(_patched_non_dp_dim_names)
+    ParallelismConfig._get_mesh = _patched_get_mesh
+    ParallelismConfig._validate_accelerator = _validate_accelerator
+    AcceleratorState.is_fsdp2 = property(patched_is_fsdp2)
+    patch_prepare_data_loader_for_ep()
+    patch_clip_grad_norm_for_ep()
+
+
+def _ep_aware_clip_grad_norm(parameters, max_norm, norm_type=2.0):
+    """`clip_grad_norm_` for params sharded across different DeviceMeshes.
+
+    Stock `torch.nn.utils.clip_grad_norm_` stacks per-param norms, which
+    DTensor rejects across meshes (experts on `dp_shard` vs non-experts on
+    `dp_shard_cp`). Instead, compute the local p-norm contribution per rank,
+    all-reduce the sum across the world, take the p-th root, and apply the
+    clip coefficient. Supports any finite p ≥ 1 plus `inf`.
+    """
+    import math
+
+    import torch
+    import torch.distributed as dist
+    from torch.distributed.tensor import DTensor
+
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    grads = [p.grad for p in parameters if p.grad is not None]
+    if not grads:
+        return torch.tensor(0.0)
+
+    norm_type = float(norm_type)
+    device = grads[0].device
+    if isinstance(grads[0], DTensor):
+        device = grads[0].to_local().device
+
+    is_inf = math.isinf(norm_type)
+    local_acc = torch.zeros((), device=device, dtype=torch.float32)
+    for g in grads:
+        local = g.to_local() if isinstance(g, DTensor) else g
+        local_f32 = local.detach().to(torch.float32)
+        if is_inf:
+            local_acc = torch.maximum(local_acc, local_f32.abs().max())
+        else:
+            local_acc = local_acc + local_f32.abs().pow(norm_type).sum()
+
+    if dist.is_available() and dist.is_initialized():
+        op = dist.ReduceOp.MAX if is_inf else dist.ReduceOp.SUM
+        dist.all_reduce(local_acc, op=op)
+
+    total_norm = local_acc if is_inf else local_acc.pow(1.0 / norm_type)
+    clip_coef = (float(max_norm) / (total_norm + 1e-6)).clamp(max=1.0)
+    for g in grads:
+        local = g.to_local() if isinstance(g, DTensor) else g
+        local.detach().mul_(clip_coef.to(local.dtype))
+
+    return total_norm.to(
+        grads[0].dtype if grads[0].is_floating_point() else torch.float32
+    )
+
+
+def patch_clip_grad_norm_for_ep():
+    """Replace `Accelerator.clip_grad_norm_` with the EP-aware version when
+    the active parallelism composes `ep` with `dp_shard` and/or `cp` (i.e., the
+    FSDP+EP composition produces multi-mesh DTensor grads — the experts shard on
+    the dp_shard/cp subgroup, the non-experts on the flattened dp_shard_cp mesh,
+    so the stock `clip_grad_norm_` can't stack their per-param norms together).
+    """
+    from accelerate import Accelerator
+
+    if getattr(Accelerator, "_AXOLOTL_EP_CLIP_PATCHED", False):
+        return
+    orig = Accelerator.clip_grad_norm_
+
+    def patched_clip_grad_norm_(self, parameters, max_norm, norm_type=2):
+        pc = getattr(self, "parallelism_config", None)
+        if (
+            pc is not None
+            and getattr(pc, "ep_enabled", False)
+            and (
+                getattr(pc, "dp_shard_enabled", False)
+                or getattr(pc, "cp_enabled", False)
+            )
+        ):
+            self.unscale_gradients()
+            params = list(parameters)
+            return _ep_aware_clip_grad_norm(params, max_norm, norm_type=norm_type)
+        return orig(self, parameters, max_norm, norm_type=norm_type)
+
+    Accelerator.clip_grad_norm_ = patched_clip_grad_norm_
+    Accelerator._AXOLOTL_EP_CLIP_PATCHED = True
+
+
+def patch_prepare_cp():
+    import contextlib
+
+    from accelerate import Accelerator
+    from transformers import Trainer
+
+    def patched_prepare_cp(self, *args):
+        if self.parallelism_config.cp_backend == "deepspeed":
+            return args
+
+        @contextlib.contextmanager
+        def _noop_cp_context(
+            buffers=None, buffer_seq_dims=None, no_restore_buffers=None
+        ):
+            yield
+
+        self._cp_context = _noop_cp_context
+        return args
+
+    def _noop_prepare_context_parallel_inputs(self, model, inputs):
+        return contextlib.nullcontext, inputs
+
+    # prevent double CP partition
+    Accelerator._prepare_cp = patched_prepare_cp
+
+    # remove unneeded calculation upstream
+    Trainer._prepare_context_parallel_inputs = _noop_prepare_context_parallel_inputs
+
+
+def _patched_prepare_data_loader_factory(orig_fn):
+    """Wrap `accelerate.data_loader.prepare_data_loader` to count the EP axis
+    as a data-parallel dimension.
+
+    Stock accelerate (line ~1155 in 1.13.0) computes
+        num_processes = dp_shard * dp_replicate
+        process_index = process_index // (tp * cp)
+    which ignores EP. EP ranks see DIFFERENT data (each rank pulls its own
+    batch), so EP belongs in the data-parallel size — same way `dp_replicate`
+    does.
+    """
+    import torch
+
+    def patched(*args, **kwargs):
+        torch_device_mesh = kwargs.get("torch_device_mesh", None)
+        if (
+            torch_device_mesh is not None
+            and isinstance(torch_device_mesh, torch.distributed.device_mesh.DeviceMesh)
+            and "ep" in torch_device_mesh.mesh_dim_names
+        ):
+            from accelerate.state import PartialState
+            from accelerate.utils import DistributedType
+
+            state = PartialState()
+            if state.distributed_type != DistributedType.DEEPSPEED:
+                ep_size = torch_device_mesh["ep"].size()
+                tp_size = (
+                    torch_device_mesh["tp"].size()
+                    if "tp" in torch_device_mesh.mesh_dim_names
+                    else 1
+                )
+                cp_size = (
+                    torch_device_mesh["cp"].size()
+                    if "cp" in torch_device_mesh.mesh_dim_names
+                    else 1
+                )
+                fsdp_size = (
+                    torch_device_mesh["dp_shard"].size()
+                    if "dp_shard" in torch_device_mesh.mesh_dim_names
+                    else 1
+                )
+                dp_size = (
+                    torch_device_mesh["dp_replicate"].size()
+                    if "dp_replicate" in torch_device_mesh.mesh_dim_names
+                    else 1
+                )
+                num_processes = fsdp_size * dp_size * ep_size
+                process_index = state.process_index // (tp_size * cp_size)
+                kwargs["num_processes"] = num_processes
+                kwargs["process_index"] = process_index
+                # Once we've supplied num_processes/process_index explicitly,
+                # accelerate's internal mesh path (which would re-derive without
+                # ep) is bypassed.
+                kwargs["torch_device_mesh"] = None
+        return orig_fn(*args, **kwargs)
+
+    return patched
+
+
+def patch_prepare_data_loader_for_ep():
+    """Apply the EP-aware data-loader patch.
+
+    Idempotent: replacing the bound function more than once is harmless because
+    the wrapper closes over the *current* `prepare_data_loader`.
+    """
+    import accelerate as _accel
+    from accelerate import data_loader as _dl
+
+    if getattr(_dl, "_AXOLOTL_EP_PATCHED", False):
+        return
+    orig = _dl.prepare_data_loader
+    wrapped = _patched_prepare_data_loader_factory(orig)
+    _dl.prepare_data_loader = wrapped
+    # accelerate.Accelerator imports prepare_data_loader at module load, so
+    # we have to patch the binding it captured too.
+    if hasattr(_accel, "prepare_data_loader"):
+        _accel.prepare_data_loader = wrapped
+    # Likewise the Accelerator module's local reference.
+    from accelerate import accelerator as _acc_mod
+
+    if hasattr(_acc_mod, "prepare_data_loader"):
+        _acc_mod.prepare_data_loader = wrapped
+    _dl._AXOLOTL_EP_PATCHED = True
diff --git a/src/axolotl/monkeypatch/activation_offload_checkpoint.py b/src/axolotl/monkeypatch/activation_offload_checkpoint.py
new file mode 100644
index 0000000000..aa9d153a34
--- /dev/null
+++ b/src/axolotl/monkeypatch/activation_offload_checkpoint.py
@@ -0,0 +1,260 @@
+"""Activation-checkpoint input ("hidden_states") offloading, ALST-style.
+
+Gradient checkpointing already discards a layer's intermediate activations and
+recomputes them in backward; it still keeps the *checkpoint input* (the layer's
+hidden_states) resident. For long sequences those per-layer inputs
+(num_layers x [seq, hidden]) dominate GPU memory. This offloads only that one
+tensor per checkpoint to CPU and brings it back for the backward recompute —
+minimal PCIe traffic (one tensor/layer, not every activation).
+
+The copies are overlapped with compute on a side stream (forward async d2h with
+bounded in-flight + backward h2d prefetch), so at long seq the transfer hides
+behind the recompute.
+
+It replaces torch's reentrant ``CheckpointFunction`` (use_reentrant=True), so it
+is framework-agnostic (works under FSDP2; no DeepSpeed dependency). DTensor
+inputs (sequence/context parallel) are offloaded too — ``DTensor.to`` round-trips
+the local shard and preserves placements. Adapted from
+torch.utils.checkpoint.CheckpointFunction and Snowflake ArcticTraining.
+"""
+
+import contextlib
+
+import torch
+from torch.utils.checkpoint import (
+    _get_autocast_kwargs,
+    _get_device_module,
+    _infer_device_type,
+    check_backward_validity,
+    detach_variable,
+    get_device_states,
+    set_device_states,
+)
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _offloadable(t: torch.Tensor) -> bool:
+    # Offload any on-device activation, INCLUDING DTensors: under sequence/context
+    # parallelism the layer input is a DTensor and offloading its sharded local
+    # tensor is exactly what we want. DTensor.to("cpu")/.to(dev) round-trips the
+    # local shard and preserves placements, so the plain path handles it. Model
+    # params never reach here (they live inside run_function, not the offloaded
+    # arg), so TRL-style param-storage filtering is unnecessary.
+    return t.device.type != "cpu"
+
+
+class _StreamOffloadManager:
+    """Per-forward-pass coordinator that overlaps the input d2h/h2d copies with
+    compute on a side stream. Forward: async-copy each layer input to CPU while
+    later layers compute (bounded in-flight so only a few sources stay resident).
+    Backward: bring inputs back on the side stream and prefetch the next one
+    (inputs are consumed in reverse-of-offload order) to hide the reload."""
+
+    def __init__(self, max_inflight=2):
+        self.s1 = torch.cuda.Stream()
+        self.max_inflight = max_inflight
+        self.reset()
+
+    def reset(self):
+        self.next_id = 0
+        self.inflight = {}  # id -> (gpu_src, event)  keep source alive until d2h done
+        self.cpu = {}  # id -> (cpu_tensor, device, requires_grad)
+        self.prefetched = {}  # id -> (gpu_tensor, event)
+
+    # ---- forward ----
+    def offload(self, t):
+        s0 = torch.cuda.current_stream()
+        self.s1.wait_stream(s0)  # input must be produced before we copy it
+        with torch.cuda.stream(self.s1):
+            cpu_t = t.detach().to("cpu", non_blocking=True)
+        ev = self.s1.record_event()
+        tid = self.next_id
+        self.next_id += 1
+        self.inflight[tid] = (t, ev)  # hold GPU source until copy completes
+        self.cpu[tid] = (cpu_t, t.device, t.requires_grad)
+        self._reap()
+        return tid
+
+    def _reap(self):
+        for tid in [k for k, (_, ev) in self.inflight.items() if ev.query()]:
+            del self.inflight[tid]
+        while len(self.inflight) > self.max_inflight:
+            tid = min(self.inflight)
+            _, ev = self.inflight.pop(tid)
+            ev.synchronize()
+
+    # ---- backward ----
+    def _bring_back(self, tid):
+        cpu_t, device, _ = self.cpu[tid]
+        with torch.cuda.stream(self.s1):
+            gpu_t = cpu_t.to(device, non_blocking=True)
+        return gpu_t, self.s1.record_event()
+
+    def restore(self, tid):
+        if tid in self.prefetched:
+            gpu_t, ev = self.prefetched.pop(tid)
+        else:
+            gpu_t, ev = self._bring_back(tid)
+        nxt = (
+            tid - 1
+        )  # next input needed (reverse order) — prefetch to overlap recompute
+        if nxt >= 0 and nxt in self.cpu and nxt not in self.prefetched:
+            self.prefetched[nxt] = self._bring_back(nxt)
+        s0 = torch.cuda.current_stream()
+        s0.wait_event(ev)
+        # gpu_t was allocated on the side stream; tell the allocator it's now used
+        # on the compute stream so its storage isn't reused before the recompute
+        # consumes it (use-after-free guard, as in TRL's offloader).
+        gpu_t.record_stream(s0)
+        _, _, requires_grad = self.cpu[tid]
+        out = gpu_t.detach().requires_grad_(requires_grad)
+        del self.cpu[tid]
+        if not self.cpu:
+            self.reset()
+        return out
+
+
+_STREAM_MANAGER = None
+
+
+def _get_stream_manager():
+    global _STREAM_MANAGER
+    if _STREAM_MANAGER is None:
+        _STREAM_MANAGER = _StreamOffloadManager()
+    return _STREAM_MANAGER
+
+
+class HiddenStatesOffloadCheckpoint(torch.autograd.Function):
+    """Reentrant checkpoint that offloads the layer input (hidden_states) to CPU,
+    overlapping the transfer with compute via a side stream."""
+
+    @staticmethod
+    def forward(ctx, run_function, preserve_rng_state, *args):
+        check_backward_validity(args)
+        ctx.run_function = run_function
+        ctx.preserve_rng_state = preserve_rng_state
+        ctx.device_type = _infer_device_type(*args)
+        ctx.device_autocast_kwargs, ctx.cpu_autocast_kwargs = _get_autocast_kwargs(
+            ctx.device_type
+        )
+        if preserve_rng_state:
+            ctx.fwd_cpu_state = torch.get_rng_state()
+            ctx.had_device_in_fwd = False
+            device_module = _get_device_module(ctx.device_type)
+            if getattr(device_module, "_initialized", False):
+                ctx.had_device_in_fwd = True
+                ctx.fwd_devices, ctx.fwd_device_states = get_device_states(*args)
+
+        mgr = _get_stream_manager()
+        ctx.inputs = []
+        ctx.tensor_indices = []
+        ctx.offload_tid = None  # manager id for the offloaded input
+        ctx.offload_arg_pos = None
+        tensor_inputs = []
+        for i, arg in enumerate(args):
+            if torch.is_tensor(arg):
+                # Offload only the first tensor (the layer input / hidden_states);
+                # other tensors (e.g. a shared [seq, seq] mask) stay on-device.
+                if ctx.offload_tid is None and _offloadable(arg):
+                    ctx.offload_tid = mgr.offload(arg)
+                    ctx.offload_arg_pos = i
+                    ctx.inputs.append(None)
+                else:
+                    tensor_inputs.append(arg)
+                    ctx.tensor_indices.append(i)
+                    ctx.inputs.append(None)
+            else:
+                ctx.inputs.append(arg)
+        ctx.save_for_backward(*tensor_inputs)
+        with torch.no_grad():
+            outputs = run_function(*args)
+        return outputs
+
+    @staticmethod
+    def backward(ctx, *args):
+        if not torch.autograd._is_checkpoint_valid():
+            raise RuntimeError(
+                "use_reentrant=True checkpoint is incompatible with .grad()/passing "
+                "`inputs` to .backward()."
+            )
+        mgr = _get_stream_manager()
+        inputs = list(ctx.inputs)
+        tensors = ctx.saved_tensors
+        for i, idx in enumerate(ctx.tensor_indices):
+            inputs[idx] = tensors[i]
+        if ctx.offload_tid is not None:
+            inputs[ctx.offload_arg_pos] = mgr.restore(ctx.offload_tid)
+
+        rng_devices = []
+        if ctx.preserve_rng_state and ctx.had_device_in_fwd:
+            rng_devices = ctx.fwd_devices
+        with torch.random.fork_rng(
+            devices=rng_devices,
+            enabled=ctx.preserve_rng_state,
+            device_type=ctx.device_type,
+        ):
+            if ctx.preserve_rng_state:
+                torch.set_rng_state(ctx.fwd_cpu_state)
+                if ctx.had_device_in_fwd:
+                    set_device_states(
+                        ctx.fwd_devices,
+                        ctx.fwd_device_states,
+                        device_type=ctx.device_type,
+                    )
+            detached_inputs = detach_variable(tuple(inputs))
+            device_autocast_ctx = (
+                torch.amp.autocast(
+                    device_type=ctx.device_type, **ctx.device_autocast_kwargs
+                )
+                if torch.amp.is_autocast_available(ctx.device_type)
+                else contextlib.nullcontext()
+            )
+            with (
+                torch.enable_grad(),
+                device_autocast_ctx,
+                torch.amp.autocast("cpu", **ctx.cpu_autocast_kwargs),
+            ):
+                outputs = ctx.run_function(*detached_inputs)
+        if isinstance(outputs, torch.Tensor):
+            outputs = (outputs,)
+        outputs_with_grad = []
+        args_with_grad = []
+        for i in range(len(outputs)):
+            if torch.is_tensor(outputs[i]) and outputs[i].requires_grad:
+                outputs_with_grad.append(outputs[i])
+                args_with_grad.append(args[i])
+        if len(outputs_with_grad) == 0:
+            raise RuntimeError("none of output has requires_grad=True")
+        torch.autograd.backward(outputs_with_grad, args_with_grad)
+        grads = tuple(
+            inp.grad if isinstance(inp, torch.Tensor) else None
+            for inp in detached_inputs
+        )
+        return (None, None) + grads
+
+
+_ORIG_CHECKPOINT_FUNCTION = None
+
+
+def patch_hidden_states_offload():
+    """Replace torch's reentrant ``CheckpointFunction`` with the hidden_states
+    offloading version. ``torch.utils.checkpoint.checkpoint(..., use_reentrant=True)``
+    resolves ``CheckpointFunction`` from the module namespace at call time, so
+    swapping the attribute is sufficient (and FSDP2-safe — only activations move)."""
+    global _ORIG_CHECKPOINT_FUNCTION
+    import torch.utils.checkpoint as ckpt
+
+    if _ORIG_CHECKPOINT_FUNCTION is None:
+        _ORIG_CHECKPOINT_FUNCTION = ckpt.CheckpointFunction
+    ckpt.CheckpointFunction = HiddenStatesOffloadCheckpoint
+    LOG.info("Patched checkpoint with hidden_states CPU offload (streamed)")
+
+
+def unpatch_hidden_states_offload():
+    import torch.utils.checkpoint as ckpt
+
+    if _ORIG_CHECKPOINT_FUNCTION is not None:
+        ckpt.CheckpointFunction = _ORIG_CHECKPOINT_FUNCTION
diff --git a/src/axolotl/monkeypatch/attention/__init__.py b/src/axolotl/monkeypatch/attention/__init__.py
new file mode 100644
index 0000000000..74bd61e774
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/__init__.py
@@ -0,0 +1,45 @@
+"""
+attention module for attention monkeypatches
+"""
+
+from transformers.integrations.flash_attention import flash_attention_forward
+
+
+def patch_xformers_attn_over_fa2():
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    from .xformers import xformers_attention_forward
+
+    ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = xformers_attention_forward
+
+
+def unpatch_xformers_attn_over_fa2():
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = flash_attention_forward()
+
+
+def register_xformers_attn():
+    """Register xformers as its own attention backend with FA2 mask behavior."""
+    from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    from .xformers import xformers_attention_forward
+
+    ALL_ATTENTION_FUNCTIONS.register("xformers", xformers_attention_forward)
+    ALL_MASK_ATTENTION_FUNCTIONS.register(
+        "xformers", ALL_MASK_ATTENTION_FUNCTIONS["flash_attention_2"]
+    )
+
+
+def register_sage_attn():
+    """Register sage as its own attention backend with FA2 mask behavior."""
+    from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    from .sage_attn import sage_attention_forward
+
+    ALL_ATTENTION_FUNCTIONS.register("sage", sage_attention_forward)
+    ALL_MASK_ATTENTION_FUNCTIONS.register(
+        "sage", ALL_MASK_ATTENTION_FUNCTIONS["flash_attention_2"]
+    )
diff --git a/src/axolotl/monkeypatch/attention/flash_attn_4.py b/src/axolotl/monkeypatch/attention/flash_attn_4.py
new file mode 100644
index 0000000000..b3bde00c15
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/flash_attn_4.py
@@ -0,0 +1,111 @@
+"""Transparently upgrade FA2 to FA4 when available on SM90+ hardware."""
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _get_head_dims(model_config):
+    """Extract (head_dim, head_dim_v) from a model config.
+
+    Handles composite models (e.g. Qwen3.5 VL) via text_config and
+    MLA models (DeepSeek/Kimi) that have separate Q/V head dimensions.
+    """
+    cfg = model_config
+    if hasattr(cfg, "text_config"):
+        cfg = cfg.text_config
+
+    # MLA models: Q head_dim = qk_nope + qk_rope, V head_dim = v_head_dim
+    if hasattr(cfg, "qk_nope_head_dim") and hasattr(cfg, "qk_rope_head_dim"):
+        head_dim = cfg.qk_nope_head_dim + cfg.qk_rope_head_dim
+        head_dim_v = getattr(cfg, "v_head_dim", head_dim)
+        return head_dim, head_dim_v
+
+    # Standard models
+    if hasattr(cfg, "head_dim"):
+        return cfg.head_dim, cfg.head_dim
+    if hasattr(cfg, "hidden_size") and hasattr(cfg, "num_attention_heads"):
+        head_dim = cfg.hidden_size // cfg.num_attention_heads
+        return head_dim, head_dim
+
+    return None, None
+
+
+def patch_flash_attn_4(model_config=None):
+    """Patch _lazy_imports to redirect FA2 imports to FA4 if available on supported hardware."""
+    if not torch.cuda.is_available():
+        return
+
+    major, _ = torch.cuda.get_device_capability()
+    # Matches flash_attn/cute/interface.py: arch / 10 in [9, 10, 11]
+    if major not in (9, 10, 11):
+        return
+
+    try:
+        from flash_attn.cute import (  # noqa: F401
+            flash_attn_func,
+            flash_attn_varlen_func,
+        )
+    except ImportError:
+        LOG.info(
+            "Flash Attention 4 is available for your GPU and offers faster training speeds. "
+            "To enable: pip install flash-attn-4"
+        )
+        return
+
+    # Validate head dimensions against FA4's own constraints
+    head_dim = None
+    if model_config is not None:
+        head_dim, head_dim_v = _get_head_dims(model_config)
+        if head_dim is not None:
+            try:
+                from flash_attn.cute.interface import _validate_head_dims
+            except ImportError:
+                LOG.warning(
+                    "Could not import _validate_head_dims from flash_attn.cute.interface, "
+                    "unable to verify head dimension compatibility, falling back to FA2"
+                )
+                return
+
+            # alignment = 16 // element_size; bf16/fp16 = 2 bytes -> alignment = 8
+            alignment = 8
+            try:
+                _validate_head_dims(head_dim, head_dim_v, major, alignment)
+            except AssertionError as exc:
+                LOG.warning(
+                    "Model head dimensions not supported by FA4, "
+                    "falling back to FA2: %s",
+                    exc,
+                )
+                return
+
+    import transformers.modeling_flash_attention_utils as fa_utils
+
+    if getattr(fa_utils._lazy_imports, "_axolotl_patched", False):
+        return
+
+    try:
+        # flash-attn-4>=4.0.0b7
+        from flash_attn.cute import flash_attn_with_kvcache
+    except ImportError:
+        flash_attn_with_kvcache = None
+
+    def _patched_lazy_imports(
+        implementation, attention_wrapper=None, allow_all_kernels=False
+    ):
+        return (
+            flash_attn_func,
+            flash_attn_varlen_func,
+            flash_attn_with_kvcache,
+            fa_utils._pad_input,
+            fa_utils._unpad_input,
+        )
+
+    _patched_lazy_imports._axolotl_patched = True
+    fa_utils._lazy_imports = _patched_lazy_imports
+    LOG.info(
+        "Flash Attention 4 enabled (head_dim=%s)",
+        head_dim if model_config else "unknown",
+    )
diff --git a/src/axolotl/monkeypatch/attention/flash_attn_d512.py b/src/axolotl/monkeypatch/attention/flash_attn_d512.py
new file mode 100644
index 0000000000..79caeacafa
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/flash_attn_d512.py
@@ -0,0 +1,546 @@
+"""Triton FlashAttention for head_dim=512 (forward + backward, dense + varlen packing).
+
+Flash/cuDNN SDPA backends cap at head_dim 256; at 512 PyTorch falls back to the memory-efficient or
+math backend (O(S^2), and math materializes the full scores). This kernel fills that gap: a tiled
+flash attention that fits head_dim 512 by using small M/N blocks (the d=512 register/SMEM wall) and
+runs fwd+bwd. Varlen packing is done in-kernel from position_ids (doc_start = row - position_id), so
+no block-diagonal mask tensor is built. Validated bit-accurate (cosine 1.0) vs per-document SDPA and
+~2x faster than the SDPA-efficient fallback at head_dim 512 on sm_120.
+
+Primary use: the Gemma-4 global (full_attention, head_dim=512) layers under sample packing.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+DEV = "cuda"
+
+
+def _fwd_smem_bytes(BLOCK_N, head_dim, num_stages):
+    """Exact SMEM the forward kernel needs: the q-tile lives in registers; the K and V tiles are
+    (1+num_stages)-buffered for software pipelining (+~2KB scratch for L/reductions). Calibrated
+    against compiled kernels: (BN,D)=(32,512) -> 67584 B @ns1, 100352 B @ns2 (clean +32768/stage)."""
+    return BLOCK_N * head_dim * 2 * (1 + num_stages) + 2048
+
+
+def _device_smem_limit():
+    props = torch.cuda.get_device_properties(torch.cuda.current_device())
+    return (
+        getattr(props, "shared_memory_per_block_optin", 0)
+        or props.shared_memory_per_block
+    )
+
+
+def _prune_fwd_configs(configs, named_args, **kwargs):
+    """Drop configs whose estimated SMEM exceeds the device opt-in cap before they ever compile, so
+    autotuning doesn't pay the compile cost of doomed variants (Triton would also reject them at
+    compile via OutOfResources; this just skips the attempt). Always keep at least one."""
+    # constexpr meta (HEAD_DIM/CAUSAL/VARLEN) arrives via kwargs, runtime args via named_args.
+    head_dim = kwargs.get("HEAD_DIM", named_args.get("HEAD_DIM"))
+    if (
+        head_dim is None
+    ):  # can't estimate -> let Triton's compile-time OOM check prune instead
+        return list(configs)
+    limit = _device_smem_limit()
+    fit = [
+        c
+        for c in configs
+        if _fwd_smem_bytes(c.kwargs["BLOCK_N"], head_dim, c.num_stages) <= limit
+    ]
+    return fit or [min(configs, key=lambda c: c.num_stages)]
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_M": 32, "BLOCK_N": 32}, num_warps=4, num_stages=s)
+        for s in (1, 2, 3)
+    ],
+    key=["N_CTX", "HEAD_DIM", "CAUSAL", "VARLEN"],
+    prune_configs_by={"early_config_prune": _prune_fwd_configs},
+)
+@triton.jit
+def _fwd(
+    Q,
+    K,
+    V,
+    sm_scale,
+    Out,
+    L,
+    sqb,
+    sqh,
+    sqm,
+    sqd,
+    skb,
+    skh,
+    skn,
+    skd,
+    svb,
+    svh,
+    svn,
+    svd,
+    sob,
+    soh,
+    som,
+    sod,
+    POS,
+    spb,
+    spn,
+    H,
+    N_CTX,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    CAUSAL: tl.constexpr,
+    VARLEN: tl.constexpr,
+):
+    start_m = tl.program_id(0)
+    off_bh = tl.program_id(1)
+    off_b = off_bh // H
+    off_h = off_bh % H
+    qb = Q + off_b * sqb + off_h * sqh
+    kb = K + off_b * skb + off_h * skh
+    vb = V + off_b * svb + off_h * svh
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, HEAD_DIM)
+    q = tl.load(
+        qb + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+        mask=offs_m[:, None] < N_CTX,
+        other=0.0,
+    )
+    ds_m = (
+        offs_m - tl.load(POS + off_b * spb + offs_m * spn, mask=offs_m < N_CTX, other=0)
+        if VARLEN
+        else offs_m * 0
+    )
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_M], tl.float32)
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
+    n_end = (start_m + 1) * BLOCK_M if CAUSAL else N_CTX
+    n_start = (
+        tl.min(ds_m) // BLOCK_N
+    ) * BLOCK_N  # varlen: skip kv blocks before this q-block's doc
+    for start_n in range(n_start, n_end, BLOCK_N):
+        cur = start_n + offs_n
+        k = tl.load(
+            kb + cur[None, :] * skn + offs_d[:, None] * skd,
+            mask=cur[None, :] < N_CTX,
+            other=0.0,
+        )
+        qk = tl.dot(q, k) * sm_scale
+        mask = cur[None, :] < N_CTX
+        if CAUSAL:
+            mask = mask & (offs_m[:, None] >= cur[None, :])
+        if VARLEN:
+            mask = mask & (cur[None, :] >= ds_m[:, None])
+        qk = tl.where(mask, qk, -1.0e9)
+        m_new = tl.maximum(m_i, tl.max(qk, 1))
+        p = tl.exp(qk - m_new[:, None])
+        alpha = tl.exp(m_i - m_new)
+        l_i = l_i * alpha + tl.sum(p, 1)
+        v = tl.load(
+            vb + cur[:, None] * svn + offs_d[None, :] * svd,
+            mask=cur[:, None] < N_CTX,
+            other=0.0,
+        )
+        acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
+        m_i = m_new
+    acc = acc / l_i[:, None]
+    tl.store(
+        Out + off_b * sob + off_h * soh + offs_m[:, None] * som + offs_d[None, :] * sod,
+        acc.to(Out.dtype.element_ty),
+        mask=offs_m[:, None] < N_CTX,
+    )
+    tl.store(L + off_bh * N_CTX + offs_m, m_i + tl.log(l_i), mask=offs_m < N_CTX)
+
+
+@triton.jit
+def _bwd_pre(
+    O,  # noqa: E741
+    DO,
+    Delta,
+    sob,
+    soh,
+    som,
+    sod,
+    H,
+    N_CTX,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+):
+    start_m = tl.program_id(0)
+    off_bh = tl.program_id(1)
+    off_b = off_bh // H
+    off_h = off_bh % H
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, HEAD_DIM)
+    base = off_b * sob + off_h * soh + offs_m[:, None] * som + offs_d[None, :] * sod
+    o = tl.load(O + base, mask=offs_m[:, None] < N_CTX, other=0.0)
+    do = tl.load(DO + base, mask=offs_m[:, None] < N_CTX, other=0.0)
+    tl.store(
+        Delta + off_bh * N_CTX + offs_m,
+        tl.sum(o.to(tl.float32) * do.to(tl.float32), 1),
+        mask=offs_m < N_CTX,
+    )
+
+
+@triton.jit
+def _bwd_dkdv(
+    Q,
+    K,
+    V,
+    sm_scale,
+    DO,
+    DK,
+    DV,
+    L,
+    D,
+    POS,
+    DOC_END,
+    spb,
+    spn,
+    sqb,
+    sqh,
+    sqm,
+    sqd,
+    H,
+    N_CTX,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    CAUSAL: tl.constexpr,
+    VARLEN: tl.constexpr,
+):
+    start_n = tl.program_id(0)
+    off_bh = tl.program_id(1)
+    off_b = off_bh // H
+    off_h = off_bh % H
+    base = off_b * sqb + off_h * sqh
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, HEAD_DIM)
+    offs_m = tl.arange(0, BLOCK_M)
+    k = tl.load(
+        K + base + offs_n[:, None] * sqm + offs_d[None, :] * sqd,
+        mask=offs_n[:, None] < N_CTX,
+        other=0.0,
+    )
+    v = tl.load(
+        V + base + offs_n[:, None] * sqm + offs_d[None, :] * sqd,
+        mask=offs_n[:, None] < N_CTX,
+        other=0.0,
+    )
+    dk = tl.zeros([BLOCK_N, HEAD_DIM], tl.float32)
+    dv = tl.zeros([BLOCK_N, HEAD_DIM], tl.float32)
+    m_start = start_n * BLOCK_N if CAUSAL else 0
+    # varlen: queries beyond this kv-block's document don't attend -> stop at its doc end
+    last_k = tl.minimum(start_n * BLOCK_N + BLOCK_N - 1, N_CTX - 1)
+    m_end = tl.load(DOC_END + off_b * spb + last_k * spn) if VARLEN else N_CTX
+    for start_m in range(m_start, m_end, BLOCK_M):
+        cur_m = start_m + offs_m
+        q = tl.load(
+            Q + base + cur_m[:, None] * sqm + offs_d[None, :] * sqd,
+            mask=cur_m[:, None] < N_CTX,
+            other=0.0,
+        )
+        do = tl.load(
+            DO + base + cur_m[:, None] * sqm + offs_d[None, :] * sqd,
+            mask=cur_m[:, None] < N_CTX,
+            other=0.0,
+        )
+        l = tl.load(L + off_bh * N_CTX + cur_m, mask=cur_m < N_CTX, other=0.0)  # noqa: E741
+        delta = tl.load(D + off_bh * N_CTX + cur_m, mask=cur_m < N_CTX, other=0.0)
+        s = tl.dot(q, tl.trans(k)) * sm_scale
+        mask = (cur_m[:, None] < N_CTX) & (offs_n[None, :] < N_CTX)
+        if CAUSAL:
+            mask = mask & (cur_m[:, None] >= offs_n[None, :])
+        if VARLEN:
+            ds_m = cur_m - tl.load(
+                POS + off_b * spb + cur_m * spn, mask=cur_m < N_CTX, other=0
+            )
+            mask = mask & (offs_n[None, :] >= ds_m[:, None])
+        p = tl.where(mask, tl.exp(s - l[:, None]), 0.0)
+        dv += tl.dot(tl.trans(p).to(do.dtype), do)
+        dp = tl.dot(do, tl.trans(v))
+        ds = (p * (dp - delta[:, None]) * sm_scale).to(q.dtype)
+        dk += tl.dot(tl.trans(ds), q)
+    tl.store(
+        DK + base + offs_n[:, None] * sqm + offs_d[None, :] * sqd,
+        dk.to(DK.dtype.element_ty),
+        mask=offs_n[:, None] < N_CTX,
+    )
+    tl.store(
+        DV + base + offs_n[:, None] * sqm + offs_d[None, :] * sqd,
+        dv.to(DV.dtype.element_ty),
+        mask=offs_n[:, None] < N_CTX,
+    )
+
+
+@triton.jit
+def _bwd_dq(
+    Q,
+    K,
+    V,
+    sm_scale,
+    DO,
+    DQ,
+    L,
+    D,
+    POS,
+    spb,
+    spn,
+    sqb,
+    sqh,
+    sqm,
+    sqd,
+    H,
+    N_CTX,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    CAUSAL: tl.constexpr,
+    VARLEN: tl.constexpr,
+):
+    start_m = tl.program_id(0)
+    off_bh = tl.program_id(1)
+    off_b = off_bh // H
+    off_h = off_bh % H
+    base = off_b * sqb + off_h * sqh
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, HEAD_DIM)
+    offs_n = tl.arange(0, BLOCK_N)
+    q = tl.load(
+        Q + base + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+        mask=offs_m[:, None] < N_CTX,
+        other=0.0,
+    )
+    do = tl.load(
+        DO + base + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+        mask=offs_m[:, None] < N_CTX,
+        other=0.0,
+    )
+    l = tl.load(L + off_bh * N_CTX + offs_m, mask=offs_m < N_CTX, other=0.0)  # noqa: E741
+    delta = tl.load(D + off_bh * N_CTX + offs_m, mask=offs_m < N_CTX, other=0.0)
+    ds_m = (
+        offs_m - tl.load(POS + off_b * spb + offs_m * spn, mask=offs_m < N_CTX, other=0)
+        if VARLEN
+        else offs_m * 0
+    )
+    dq = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
+    n_end = (start_m + 1) * BLOCK_M if CAUSAL else N_CTX
+    n_start = (tl.min(ds_m) // BLOCK_N) * BLOCK_N
+    for start_n in range(n_start, n_end, BLOCK_N):
+        cur = start_n + offs_n
+        k = tl.load(
+            K + base + cur[:, None] * sqm + offs_d[None, :] * sqd,
+            mask=cur[:, None] < N_CTX,
+            other=0.0,
+        )
+        v = tl.load(
+            V + base + cur[:, None] * sqm + offs_d[None, :] * sqd,
+            mask=cur[:, None] < N_CTX,
+            other=0.0,
+        )
+        s = tl.dot(q, tl.trans(k)) * sm_scale
+        mask = (offs_m[:, None] < N_CTX) & (cur[None, :] < N_CTX)
+        if CAUSAL:
+            mask = mask & (offs_m[:, None] >= cur[None, :])
+        if VARLEN:
+            mask = mask & (cur[None, :] >= ds_m[:, None])
+        p = tl.where(mask, tl.exp(s - l[:, None]), 0.0)
+        dp = tl.dot(do, tl.trans(v))
+        ds = (p * (dp - delta[:, None]) * sm_scale).to(k.dtype)
+        dq += tl.dot(ds, k)
+    tl.store(
+        DQ + base + offs_m[:, None] * sqm + offs_d[None, :] * sqd,
+        dq.to(DQ.dtype.element_ty),
+        mask=offs_m[:, None] < N_CTX,
+    )
+
+
+@register_kernel_op("flash_attn_d512_fwd")
+def _flash_d512_fwd_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    pos: torch.Tensor,
+    causal: bool,
+    varlen: bool,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, H, N, D = q.shape
+    spb, spn = pos.stride()
+    o = torch.empty_like(q)
+    L = torch.empty((B * H, N), device=q.device, dtype=torch.float32)
+    # BLOCK_M/BLOCK_N, num_warps, num_stages come from the autotuner (SMEM-pruned per device).
+    _fwd[lambda meta: (triton.cdiv(N, meta["BLOCK_M"]), B * H)](
+        q,
+        k,
+        v,
+        scale,
+        o,
+        L,
+        *q.stride(),
+        *k.stride(),
+        *v.stride(),
+        *o.stride(),
+        pos,
+        spb,
+        spn,
+        H,
+        N,
+        HEAD_DIM=D,
+        CAUSAL=causal,
+        VARLEN=varlen,
+    )
+    return o, L
+
+
+@_flash_d512_fwd_op.register_fake
+def _(q, k, v, pos, causal, varlen, scale):
+    B, H, N, _ = q.shape
+    return torch.empty_like(q), torch.empty(
+        (B * H, N), device=q.device, dtype=torch.float32
+    )
+
+
+@register_kernel_op("flash_attn_d512_bwd")
+def _flash_d512_bwd_op(
+    do: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    L: torch.Tensor,
+    pos: torch.Tensor,
+    doc_end: torch.Tensor,
+    causal: bool,
+    varlen: bool,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, H, N, D = q.shape
+    spb, spn = pos.stride()
+    do = do.contiguous()
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+    delta = torch.empty((B * H, N), device=q.device, dtype=torch.float32)
+    BM, BN = 16, 32
+    _bwd_pre[(triton.cdiv(N, BM), B * H)](
+        o, do, delta, *o.stride(), H, N, HEAD_DIM=D, BLOCK_M=BM, num_warps=4
+    )
+    _bwd_dkdv[(triton.cdiv(N, BN), B * H)](
+        q,
+        k,
+        v,
+        scale,
+        do,
+        dk,
+        dv,
+        L,
+        delta,
+        pos,
+        doc_end,
+        spb,
+        spn,
+        *q.stride(),
+        H,
+        N,
+        HEAD_DIM=D,
+        BLOCK_M=BM,
+        BLOCK_N=BN,
+        CAUSAL=causal,
+        VARLEN=varlen,
+        num_warps=4,
+        num_stages=1,
+    )
+    _bwd_dq[(triton.cdiv(N, BM), B * H)](
+        q,
+        k,
+        v,
+        scale,
+        do,
+        dq,
+        L,
+        delta,
+        pos,
+        spb,
+        spn,
+        *q.stride(),
+        H,
+        N,
+        HEAD_DIM=D,
+        BLOCK_M=BM,
+        BLOCK_N=BN,
+        CAUSAL=causal,
+        VARLEN=varlen,
+        num_warps=4,
+        num_stages=1,
+    )
+    return dq, dk, dv
+
+
+@_flash_d512_bwd_op.register_fake
+def _(do, q, k, v, o, L, pos, doc_end, causal, varlen, scale):
+    return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+
+
+def _compute_doc_end(pos: torch.Tensor) -> torch.Tensor:
+    """Per-row index of the next document start after each position (or N).
+
+    Suffix-min of "j if pos[j]==0 else N" shifted left by one — fully tensorized so
+    torch.compile can trace it (the per-row ``.nonzero()`` loop graph-breaks).
+    """
+    B, N = pos.shape
+    idx = torch.arange(N, device=pos.device, dtype=pos.dtype)
+    starts_or_n = torch.where(pos == 0, idx.expand_as(pos), pos.new_full((), N))
+    shifted = torch.cat([starts_or_n[:, 1:], pos.new_full((B, 1), N)], dim=1)
+    return torch.flip(
+        torch.cummin(torch.flip(shifted, dims=[1]), dim=1).values, dims=[1]
+    )
+
+
+class _FlashD512(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, causal, position_ids, scale=None):
+        # The backward kernels launch with q's strides only and reuse them for k/v/do/dk/dv. In real
+        # attention q is non-contiguous ([B,S,H,D].transpose(1,2)) while k/v are contiguous (e.g. GQA
+        # repeat_interleave); the mismatched strides make the backward read wrong memory and explode
+        # the gradient (forward is unaffected -- it passes each tensor's own strides). Force one layout.
+        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+        B, N = q.shape[0], q.shape[2]
+        D = q.shape[3]
+        VARLEN = position_ids is not None
+        if position_ids is None:
+            pos = torch.zeros((B, N), device=q.device, dtype=torch.int32)
+        else:
+            pos = (
+                (position_ids if position_ids.dim() > 1 else position_ids[None])
+                .to(torch.int32)
+                .contiguous()
+            )
+        if VARLEN:
+            doc_end = _compute_doc_end(pos)
+        else:
+            doc_end = pos
+        scale = D**-0.5 if scale is None else float(scale)
+        o, L = _flash_d512_fwd_op(q, k, v, pos, causal, VARLEN, scale)
+        ctx.save_for_backward(q, k, v, o, L, pos, doc_end)
+        ctx.causal = causal
+        ctx.scale = scale
+        ctx.varlen = VARLEN
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, L, pos, doc_end = ctx.saved_tensors
+        dq, dk, dv = _flash_d512_bwd_op(
+            do, q, k, v, o, L, pos, doc_end, ctx.causal, ctx.varlen, ctx.scale
+        )
+        return dq, dk, dv, None, None, None
+
+
+def flash_d512(q, k, v, causal=True, position_ids=None, scale=None):
+    return _FlashD512.apply(q, k, v, causal, position_ids, scale)
diff --git a/src/axolotl/monkeypatch/attention/flex_attn.py b/src/axolotl/monkeypatch/attention/flex_attn.py
new file mode 100644
index 0000000000..0071c42237
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/flex_attn.py
@@ -0,0 +1,82 @@
+"""Flex attention monkey patch"""
+
+import sys
+
+import torch
+import transformers
+from packaging import version
+from transformers.utils.import_utils import _torch_version, is_torch_less_or_equal
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_flex_wrapper(**flex_attn_compile_kwargs):
+    # TODO remove this patch when transformers#37285 is merged and in a release
+    is_torch_2_6 = torch.__version__.startswith("2.6")
+
+    if not is_torch_2_6:
+        return
+
+    from torch.nn.attention.flex_attention import flex_attention
+
+    class WrappedFlexAttention:
+        """
+        We are doing a singleton class so that flex attention is compiled once when it's first called.
+        """
+
+        _instance = None
+        _is_flex_compiled = False
+        _compiled_flex_attention = None
+
+        def __new__(cls, *args, **kwargs):
+            if cls._instance is None:
+                # Create a new instance if one doesn't already exist
+                cls._instance = super().__new__(cls)
+            return cls._instance
+
+        @classmethod
+        def del_singleton(cls):
+            cls._instance = None
+
+        @torch.compiler.disable(recursive=False)
+        def __init__(self, training):
+            """
+            Initialize or update the singleton instance.
+            """
+            self.training = None
+            if not self._is_flex_compiled or training != self.training:
+                self.training = training
+                if is_torch_less_or_equal("2.5.1"):
+                    self._compiled_flex_attention = torch.compile(
+                        flex_attention, dynamic=False
+                    )
+                # In PyTorch 2.6.0, there's a known issue with flex attention compilation which may
+                # cause errors. The suggested fix is to compile with "max-autotune-no-cudagraphs"
+                # see https://github.com/pytorch/pytorch/issues/146260 for training
+                elif version.parse(_torch_version).base_version == "2.6.0" and training:
+                    self._compiled_flex_attention = torch.compile(
+                        flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs"
+                    )
+                # Fallback, usually the most recent torch 2.7.x+ versions
+                else:
+                    LOG.info(
+                        "Compiling flex attention with kwargs: %s. This may take a while...",
+                        flex_attn_compile_kwargs,
+                    )
+                    self._compiled_flex_attention = torch.compile(
+                        flex_attention,
+                        **flex_attn_compile_kwargs,
+                    )
+                    LOG.info("Flex attention compiled successfully.")
+
+                self._is_flex_compiled = True
+
+        def __call__(self):
+            return self._compiled_flex_attention
+
+    transformers.integrations.flex_attention.WrappedFlexAttention = WrappedFlexAttention
+    sys.modules[
+        "transformers.integrations.flex_attention"
+    ].WrappedFlexAttention = WrappedFlexAttention
diff --git a/src/axolotl/monkeypatch/attention/fp8_attn.py b/src/axolotl/monkeypatch/attention/fp8_attn.py
new file mode 100644
index 0000000000..3e0ca0408e
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/fp8_attn.py
@@ -0,0 +1,30 @@
+"""FP8 low-precision attention via torchao.
+
+Requires:
+  - PyTorch >= 2.11.0
+  - SM90+ (Hopper/Blackwell) GPU
+  - flash-attn package with FA3 support
+  - torchao >= 0.17.0
+
+Uses per-head FP8 quantized attention with automatic RoPE fusion under torch.compile.
+The torchao patch replaces F.scaled_dot_product_attention, so the model must use
+HF's "sdpa" attention implementation for the patch to intercept attention calls.
+"""
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_fp8_attention(model: torch.nn.Module) -> torch.nn.Module:
+    """Apply FP8 low-precision attention to a model.
+
+    Must be called after model loading and before torch.compile.
+    KV caching should be disabled (config.use_cache = False).
+    """
+    from torchao.prototype.attention import apply_low_precision_attention
+
+    LOG.info("Applying FP8 low-precision attention (torchao)")
+    return apply_low_precision_attention(model)
diff --git a/src/axolotl/monkeypatch/attention/large_head.py b/src/axolotl/monkeypatch/attention/large_head.py
new file mode 100644
index 0000000000..fa48c2d15f
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/large_head.py
@@ -0,0 +1,139 @@
+"""Generic large-head-dim attention capability (head_dim > 256).
+
+Flash/cuDNN SDPA cap at head_dim 256; at larger head_dim PyTorch falls to the memory-efficient or
+math backend. The Triton ``flash_d512`` kernel fills that gap. This module exposes the routing as a
+MODEL-AGNOSTIC capability: any attention path can call :func:`flash_d512_route`, and a generic
+``sdpa`` wrapper (:func:`patch_sdpa_large_head`) lets plain SDPA models opt in via the
+``large_head_attention`` config. Gemma-4's hybrid global layers reuse the same router.
+
+Policy (``large_head_attention``):
+  - ``sdpa``         : never use the Triton kernel (stock SDPA at large head_dim)
+  - ``auto``         : Triton flash for genuinely packed rows (its proven win), SDPA otherwise
+                       (single-document large-head attention is faster on SDPA is_causal)
+  - ``triton_flash`` : prefer the Triton kernel for head_dim > 256 whenever possible
+"""
+
+from __future__ import annotations
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_LARGE_HEAD_MIN_DIM = 256
+_POLICY = "sdpa"
+_PACKED: bool | None = None
+_SDPA_ORIG_ATTR = "_axolotl_large_head_sdpa_original"
+
+
+def set_large_head_policy(policy: str | None) -> None:
+    global _POLICY
+    _POLICY = str(policy).lower() if policy else "sdpa"
+
+
+def set_large_head_packed(packed: bool | None) -> None:
+    """Declare statically whether batches carry packed (multi-document) rows.
+
+    Known at config time (``sample_packing``); declaring it lets the router skip
+    the per-call GPU-sync probe in ``_multidoc_position_ids``, which is a
+    data-dependent graph break under torch.compile. ``None`` keeps the runtime
+    probe (eager-only callers).
+    """
+    global _PACKED
+    _PACKED = packed
+
+
+def get_large_head_policy() -> str:
+    return _POLICY
+
+
+def resolve_large_head_policy(cfg) -> str:
+    """Resolve the intent ``large_head_attention`` (auto/sdpa/triton_flash); fall back to the
+    deprecated ``flash_attn_d512`` bool (True -> auto, since flash only wins on packed rows)."""
+    policy = cfg.get("large_head_attention")
+    if policy:
+        return str(policy).lower()
+    if cfg.get("flash_attn_d512"):
+        return "auto"
+    return "sdpa"
+
+
+def _multidoc_position_ids(position_ids):
+    """Return [B,S] position_ids iff they encode (multi-document) packing, else None."""
+    if position_ids is None:
+        return None
+    p = position_ids if position_ids.dim() > 1 else position_ids[None]
+    if _PACKED is not None:
+        # packed config: single-doc rows still run varlen (slower, compile-clean)
+        return p if _PACKED else None
+    return p if int((p == 0).sum()) > p.shape[0] else None
+
+
+def flash_d512_route(module, query, key, value, scaling, position_ids, policy=None):
+    """Route a large-head attention call through the Triton flash_d512 kernel, or return None to
+    signal the caller to fall back to SDPA. Inputs are ``[B, H, S, D]``; on success returns
+    ``(attn_output [B, S, Hq, D], None)`` matching ``sdpa_attention_forward``'s contract."""
+    policy = policy or _POLICY
+    # Allowlist: only the Triton policies route to the kernel, so a config typo can't enable it.
+    if policy not in ("auto", "triton_flash") or query.shape[-1] <= _LARGE_HEAD_MIN_DIM:
+        return None
+    pid = _multidoc_position_ids(position_ids)
+    # auto: kernel only for packed rows (single-doc large-head is faster on SDPA is_causal)
+    if policy == "auto" and pid is None:
+        return None
+    try:
+        from axolotl.monkeypatch.attention.flash_attn_d512 import flash_d512
+
+        ng = getattr(module, "num_key_value_groups", query.shape[1] // key.shape[1])
+        k = key.repeat_interleave(ng, dim=1) if ng > 1 else key
+        v = value.repeat_interleave(ng, dim=1) if ng > 1 else value
+        out = flash_d512(query, k, v, True, position_ids=pid, scale=scaling)
+        return out.transpose(1, 2).contiguous(), None
+    except Exception:  # pragma: no cover - any kernel issue falls back to SDPA
+        return None
+
+
+def patch_sdpa_large_head(policy: str | None = None) -> bool:
+    """Wrap the ``sdpa`` attention interface so head_dim>256 maskless calls route through the Triton
+    kernel per ``policy`` (idempotent). Generic — any SDPA model opts in via config; explicit
+    attention masks always fall through to stock SDPA."""
+    if policy is not None:
+        set_large_head_policy(policy)
+    if get_large_head_policy() == "sdpa":
+        return False
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    current = ALL_ATTENTION_FUNCTIONS["sdpa"]
+    if getattr(current, _SDPA_ORIG_ATTR, None) is not None:
+        return True
+    original = current
+
+    def sdpa_large_head_forward(module, query, key, value, attention_mask, **kwargs):
+        if attention_mask is None:
+            routed = flash_d512_route(
+                module,
+                query,
+                key,
+                value,
+                kwargs.get("scaling"),
+                kwargs.get("position_ids"),
+            )
+            if routed is not None:
+                return routed
+        return original(module, query, key, value, attention_mask, **kwargs)
+
+    setattr(sdpa_large_head_forward, _SDPA_ORIG_ATTR, original)
+    ALL_ATTENTION_FUNCTIONS.register("sdpa", sdpa_large_head_forward)
+    LOG.info(
+        "large_head_attention: wrapped sdpa to route head_dim>256 through Triton flash (%s)",
+        get_large_head_policy(),
+    )
+    return True
+
+
+def unpatch_sdpa_large_head() -> None:
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    current = ALL_ATTENTION_FUNCTIONS["sdpa"]
+    original = getattr(current, _SDPA_ORIG_ATTR, None)
+    if original is not None:
+        ALL_ATTENTION_FUNCTIONS.register("sdpa", original)
diff --git a/src/axolotl/monkeypatch/attention/sage_attn.py b/src/axolotl/monkeypatch/attention/sage_attn.py
new file mode 100644
index 0000000000..6e9ba0f858
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/sage_attn.py
@@ -0,0 +1,199 @@
+"""
+Monkeypatch for SageAttention for use with transformers.
+
+https://github.com/thu-ml/SageAttention/
+"""
+
+import torch
+from transformers.integrations.sdpa_attention import repeat_kv
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+sageattn = None  # pylint: disable=invalid-name
+sageattn_varlen = None  # pylint: disable=invalid-name
+
+
+def _is_sageattn_available():
+    """Determine if SageAttention is available"""
+    try:
+        import sageattention  # noqa: F401 # pylint: disable=unused-import
+
+        return True
+    except ImportError:
+        return False
+
+
+if _is_sageattn_available():
+    # import sageattn here if available
+    from sageattention import sageattn, sageattn_varlen
+
+
+def _check_sageattn_imported():
+    """Check if SageAttention is imported. Raises an ImportError if not."""
+    if sageattn is None:
+        raise ImportError(
+            "SageAttention is not installed. Please install it from source: "
+            "`pip install git+https://github.com/thu-ml/SageAttention.git@1718ddc06dbc694bcf3c6b49ac28c1921aa2d8bd`"
+        )
+
+
+def sage_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None = None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+    is_causal: bool | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    """
+    Forward pass for SageAttention compatible with transformers attention interfaces.
+
+    https://github.com/thu-ml/SageAttention/
+    """
+
+    _check_sageattn_imported()
+
+    if kwargs.get("output_attentions", False) or kwargs.get("head_mask") is not None:
+        raise NotImplementedError(
+            "SageAttention does not support `output_attentions=True` or `head_mask`."
+        )
+
+    # The base sageattn API does not support dropout.
+    if dropout > 0.0:
+        raise NotImplementedError("SageAttention does not support dropout.")
+
+    # Handle Grouped-Query Attention (GQA) and Multi-Query Attention (MQA)
+    if hasattr(module, "num_key_value_groups"):
+        key = repeat_kv(key, module.num_key_value_groups)
+        value = repeat_kv(value, module.num_key_value_groups)
+
+    # Calculate is_causal following transformers
+    assert is_causal is not False, "is_causal must be True or None"
+    is_causal = True
+
+    position_ids = kwargs.get("position_ids", None)
+    query_length = query.shape[2]
+
+    cu_seqlens_q = kwargs.get("cu_seqlens_q", None)
+    cu_seqlens_k = kwargs.get("cu_seqlens_k", None)
+    max_length_q = kwargs.get("max_length_q", None)
+    max_length_k = kwargs.get("max_length_k", None)
+
+    # Sample packing uses position_ids, so we check for it first
+    if position_ids is not None and (
+        max_length_q is not None
+        or (query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all())
+    ):
+        # transpose inputs to NHD layout for use with FA2 utils
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        batch_size = query.size(0)
+
+        from transformers.modeling_flash_attention_utils import (
+            prepare_fa2_from_position_ids,
+        )
+
+        if cu_seqlens_q is None or cu_seqlens_k is None:
+            query, key, value, indices_q, cu_seq_lens, max_seq_lens = (
+                prepare_fa2_from_position_ids(query, key, value, position_ids)
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_length_q, max_length_k = max_seq_lens
+
+        else:
+            query = query.reshape(-1, query.size(-2), query.size(-1))
+            key = key.reshape(-1, key.size(-2), key.size(-1))
+            value = value.reshape(-1, value.size(-2), value.size(-1))
+
+        attn_output_unpad = sageattn_varlen(
+            q=query,
+            k=key,
+            v=value,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_length_q,
+            max_seqlen_k=max_length_k,
+            is_causal=is_causal,
+            sm_scale=scaling,
+            smooth_k=False,  # reduces loss 0 / nan grad norms
+            tensor_layout="NHD",
+        )
+
+        attn_output = attn_output_unpad.view(
+            batch_size, -1, attn_output_unpad.size(-2), attn_output_unpad.size(-1)
+        )
+
+    elif attention_mask is not None:
+        # NOTE: When used without `pad_to_sequence_len`, the loss becomes unstable after a few steps.
+
+        assert attention_mask.ndim == 2, "Attention mask must be 2D"
+
+        from transformers.modeling_flash_attention_utils import (
+            _upad_input,
+        )
+
+        # transpose inputs to NHD layout for use with FA2 utils
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        batch_size = query.shape[0]
+
+        query, key, value, indices_q, cu_seq_lens, max_seq_lens = _upad_input(
+            query, key, value, attention_mask, query_length
+        )
+        cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+        max_seqlen_q, max_seqlen_k = max_seq_lens
+
+        attn_output_unpad = sageattn_varlen(
+            q=query,
+            k=key,
+            v=value,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            is_causal=is_causal,
+            sm_scale=scaling,
+            tensor_layout="NHD",
+        )
+
+        from flash_attn.bert_padding import pad_input
+
+        attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+    else:
+        # Use standard sageattn
+        # The input layout for transformers models is (batch_size, num_heads, seq_len, head_dim),
+        # which corresponds to SageAttention's "HND" layout.
+        attn_output = sageattn(
+            q=query,
+            k=key,
+            v=value,
+            tensor_layout="HND",
+            is_causal=is_causal,
+            sm_scale=scaling,
+        )
+
+        # SageAttention with "HND" returns (batch, heads, seq_len, head_dim)
+        # Transformers expects (batch, seq_len, heads, head_dim) for the output
+        # So we need to transpose dimensions 1 and 2
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, None
+
+
+def patch_sageattn():
+    """Validate SageAttention is available. Registration in the attention/mask
+    function registries is handled by register_sage_attn() in __init__.py."""
+
+    _check_sageattn_imported()
+
+    LOG.info("SageAttention validated successfully")
diff --git a/src/axolotl/monkeypatch/attention/sdpa_varlen.py b/src/axolotl/monkeypatch/attention/sdpa_varlen.py
new file mode 100644
index 0000000000..ca0dc54139
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/sdpa_varlen.py
@@ -0,0 +1,259 @@
+"""Variable-length (cu_seqlens) SDPA path for sample packing.
+
+With sample packing the model concatenates many documents into one row and
+encodes the boundaries in ``position_ids`` (which reset to 0 at each document
+start). The default SDPA path turns this into an explicit 4D block-diagonal
+mask: O(S^2) compute even though cross-document blocks are masked out, plus the
+mask tensor itself.
+
+When PyTorch exposes ``torch.nn.attention.varlen.varlen_attn`` (>= 2.10) and the
+head_dim is within Flash-Attention's limit (<= 256), we can instead run the
+attention as variable-length with ``cu_seqlens`` derived from ``position_ids``,
+which skips the cross-document blocks entirely — faster and lower memory — with
+no dependency on the ``flash_attn`` package. It only activates for genuinely
+multi-document (packed) rows; everything else falls back to the stock SDPA
+implementation. Auto-enabled for ``sdpa`` + ``sample_packing`` when the kernel
+can serve the model (see ``PatchManager._apply_sdpa_varlen_patch``); ``sdpa_varlen``
+overrides the choice.
+
+To make varlen actually engage during training (``use_cache=False``), patching
+also overrides the ``sdpa`` mask builder to return ``None`` for packed rows whose
+2D padding mask has been dropped. Otherwise transformers materializes a 4D
+block-diagonal mask from ``position_ids`` and hands it to the attention interface,
+which would keep the stock O(S^2) SDPA path (and, before the padding mask was
+dropped for sdpa packing, silently leak across documents on padded rows).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_PATCH_APPLIED = False
+_MASK_PATCH_APPLIED = False
+# head_dim limit of the Flash-Attention kernel backing varlen_attn.
+_VARLEN_MAX_HEAD_DIM = 256
+
+
+def varlen_available() -> bool:
+    try:
+        from torch.nn.attention.varlen import varlen_attn  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
+def _is_packed(position_ids) -> bool:
+    """More document starts (position resets to 0) than rows -> packed."""
+    pid = position_ids if position_ids.dim() > 1 else position_ids[None]
+    return int((pid == 0).sum()) > pid.shape[0]
+
+
+def _block_diagonal_causal_mask(position_ids, seq_len: int):
+    """Bool (B, 1, S, S) block-diagonal causal mask (True = attend) built from
+    per-document-resetting position_ids. Used only on the non-varlen fallback so
+    stock SDPA still isolates documents when the 2D padding mask has been dropped."""
+    import torch
+
+    pid = position_ids if position_ids.dim() > 1 else position_ids[None]
+    doc = (pid == 0).cumsum(-1)  # document id per token
+    same_doc = doc[:, :, None] == doc[:, None, :]
+    idx = torch.arange(seq_len, device=pid.device)
+    causal = idx[:, None] >= idx[None, :]
+    return (same_doc & causal[None])[:, None]
+
+
+def _build_varlen_forward(original_sdpa: Callable) -> Callable:
+    import inspect
+
+    import torch
+    from torch.nn.attention.varlen import varlen_attn
+    from transformers.modeling_flash_attention_utils import (
+        prepare_fa_kwargs_from_position_ids,
+    )
+
+    # varlen_attn's causal API differs across the supported torch range (>=2.10): torch 2.11 takes
+    # window_size (causal = (-1, 0): unlimited left, no right; sliding = (W-1, 0)), earlier builds
+    # take is_causal (causal only, no sliding). Detect which the installed build accepts. Scale stays
+    # default (1/sqrt(d)) — the use_varlen guard below already restricts to standard scaling.
+    try:
+        _varlen_params = set(inspect.signature(varlen_attn).parameters)
+    except (TypeError, ValueError):
+        _varlen_params = set()
+    # window_size present -> use it; only an is_causal-only build lacks sliding support.
+    _supports_window = (
+        "window_size" in _varlen_params or "is_causal" not in _varlen_params
+    )
+
+    def sdpa_varlen_forward(
+        module: Any,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        dropout: float = 0.0,
+        scaling: float | None = None,
+        **kwargs: Any,
+    ):
+        position_ids = kwargs.get("position_ids")
+        sliding_window = kwargs.get("sliding_window", None) or getattr(
+            module, "sliding_window", None
+        )
+        head_dim = query.shape[-1]
+        # Fast-path conditions; anything else falls back to stock SDPA.
+        # - attention_mask must be None: packing carries structure via position_ids; a real mask
+        #   (e.g. left padding) isn't expressible to the causal/sliding varlen kernel here.
+        # - dropout unsupported by varlen_attn.
+        # - head_dim within the Flash limit.
+        # - scaling: varlen_attn only applies 1/sqrt(head_dim); a custom scale can't be honored.
+        standard_scale = scaling is None or abs(scaling - head_dim**-0.5) < 1e-9
+        use_varlen = (
+            attention_mask is None
+            and not dropout
+            and head_dim <= _VARLEN_MAX_HEAD_DIM
+            and position_ids is not None
+            and standard_scale
+        )
+        packed = position_ids is not None and _is_packed(position_ids)
+        use_varlen = use_varlen and packed  # single-doc rows -> stock SDPA
+        if not use_varlen:
+            # Rebuild isolation the dropped mask no longer provides, else stock SDPA leaks.
+            if attention_mask is None and packed:
+                attention_mask = _block_diagonal_causal_mask(
+                    position_ids, query.shape[2]
+                )
+            return original_sdpa(
+                module,
+                query,
+                key,
+                value,
+                attention_mask,
+                dropout=dropout,
+                scaling=scaling,
+                **kwargs,
+            )
+
+        # A sliding window needs varlen_attn's window_size arg; an is_causal-only build can't express
+        # it, so refuse loudly there rather than silently running full causal attention (wrong).
+        if sliding_window and not _supports_window:
+            raise NotImplementedError(
+                "sdpa_varlen: sliding-window attention needs varlen_attn(window_size=...), absent in "
+                f"this torch build (requested window={sliding_window}); disable sdpa_varlen for this model."
+            )
+
+        B, Hq, S, D = query.shape
+        Hkv = key.shape[1]
+        if Hq != Hkv:  # GQA -> repeat (varlen_attn has no GQA mode)
+            n = Hq // Hkv
+            key = key.repeat_interleave(n, dim=1)
+            value = value.repeat_interleave(n, dim=1)
+        pid = position_ids if position_ids.dim() > 1 else position_ids[None]
+        (cu_q, cu_k), (max_q, max_k) = prepare_fa_kwargs_from_position_ids(pid)
+        qf = query.transpose(1, 2).reshape(B * S, Hq, D)
+        kf = key.transpose(1, 2).reshape(B * S, Hq, D)
+        vf = value.transpose(1, 2).reshape(B * S, Hq, D)
+        if _supports_window:
+            # (left, right): (-1, 0) = causal full; (W-1, 0) = causal sliding window of W.
+            window = (sliding_window - 1, 0) if sliding_window else (-1, 0)
+            causal_kw: dict = {"window_size": window}
+        else:
+            causal_kw = {
+                "is_causal": True
+            }  # is_causal-only build (sliding already refused above)
+        out = varlen_attn(
+            qf,
+            kf,
+            vf,
+            cu_q.to(torch.int32),
+            cu_k.to(torch.int32),
+            int(max_q),
+            int(max_k),
+            **causal_kw,
+        )
+        if isinstance(out, tuple):
+            out = out[0]
+        # match sdpa_attention_forward's return contract: (attn_output [B,S,Hq,D], None)
+        return out.reshape(B, S, Hq, D), None
+
+    return sdpa_varlen_forward
+
+
+def _build_varlen_mask(original_mask: Callable) -> Callable:
+    def sdpa_varlen_mask(*args: Any, **kwargs: Any):
+        # None for packed (mask-dropped) rows lets the varlen wrapper own isolation via
+        # cu_seqlens, mirroring transformers' `flash_attention_mask`; else defer to stock.
+        if kwargs.get("attention_mask") is None:
+            return None
+        return original_mask(*args, **kwargs)
+
+    return sdpa_varlen_mask
+
+
+def patch_sdpa_varlen() -> bool:
+    """Replace the registered ``sdpa`` attention with a varlen-aware wrapper (idempotent).
+
+    Also overrides the ``sdpa`` mask builder to return ``None`` for packed rows whose
+    padding mask was dropped, so the varlen wrapper actually engages during training
+    (``use_cache=False``) instead of transformers building a 4D block-diagonal mask that
+    would otherwise force the stock O(S^2) SDPA path. Only call this for models the varlen
+    path can serve (head_dim <= 256, no sliding window); other cases keep stock SDPA which
+    decontaminates via the dropped-mask block-diagonal path.
+    """
+    global _PATCH_APPLIED, _MASK_PATCH_APPLIED
+    if _PATCH_APPLIED:
+        return True
+    if not varlen_available():
+        LOG.warning(
+            "sdpa_varlen: torch.nn.attention.varlen.varlen_attn unavailable (needs torch >= 2.10); "
+            "leaving stock SDPA in place."
+        )
+        return False
+    from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    original = ALL_ATTENTION_FUNCTIONS["sdpa"]
+    original_mask = ALL_MASK_ATTENTION_FUNCTIONS["sdpa"]
+    wrapper = _build_varlen_forward(original)
+    wrapper._axolotl_sdpa_original = original  # type: ignore[attr-defined]
+    mask_wrapper = _build_varlen_mask(original_mask)
+    mask_wrapper._axolotl_sdpa_mask_original = original_mask  # type: ignore[attr-defined]
+
+    # Register both or neither, so a mid-way failure can't leave sdpa half-patched.
+    ALL_ATTENTION_FUNCTIONS.register("sdpa", wrapper)
+    try:
+        ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", mask_wrapper)
+    except Exception:
+        ALL_ATTENTION_FUNCTIONS.register("sdpa", original)
+        raise
+    _PATCH_APPLIED = True
+    _MASK_PATCH_APPLIED = True
+
+    LOG.info(
+        "sdpa_varlen: patched 'sdpa' to use cu_seqlens varlen_attn for packed rows "
+        "(head_dim <= %d), falling back to stock SDPA otherwise",
+        _VARLEN_MAX_HEAD_DIM,
+    )
+    return True
+
+
+def unpatch_sdpa_varlen() -> None:
+    global _PATCH_APPLIED, _MASK_PATCH_APPLIED
+    from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    if _PATCH_APPLIED:
+        current = ALL_ATTENTION_FUNCTIONS["sdpa"]
+        original = getattr(current, "_axolotl_sdpa_original", None)
+        if original is not None:
+            ALL_ATTENTION_FUNCTIONS.register("sdpa", original)
+        _PATCH_APPLIED = False
+
+    if _MASK_PATCH_APPLIED:
+        current_mask = ALL_MASK_ATTENTION_FUNCTIONS["sdpa"]
+        original_mask = getattr(current_mask, "_axolotl_sdpa_mask_original", None)
+        if original_mask is not None:
+            ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", original_mask)
+        _MASK_PATCH_APPLIED = False
diff --git a/src/axolotl/monkeypatch/attention/xformers.py b/src/axolotl/monkeypatch/attention/xformers.py
new file mode 100644
index 0000000000..eca95797a9
--- /dev/null
+++ b/src/axolotl/monkeypatch/attention/xformers.py
@@ -0,0 +1,160 @@
+"""
+xformers attention implementation for packing
+"""
+
+from typing import Optional
+
+import torch
+import xformers
+import xformers.ops.fmha
+from transformers.modeling_flash_attention_utils import (
+    _upad_input,
+)
+
+from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
+
+xformers_attention = xformers.ops.fmha.memory_efficient_attention
+
+
+def xformers_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    dropout: float = 0.0,
+    scaling: Optional[float] = None,
+    sliding_window: Optional[int] = None,
+    softcap: Optional[float] = None,
+    cu_seq_lens_q: Optional[torch.LongTensor] = None,
+    cu_seq_lens_k: Optional[torch.LongTensor] = None,
+    max_length_q: Optional[int] = None,
+    max_length_k: Optional[int] = None,
+    **kwargs,
+):
+    # Get dimensions
+    # query: [batch, heads, seq_len, hidden_dim]
+    batch_size = query.size(0)
+    query_length = query.shape[2]
+    key_length = key.shape[2]
+
+    # Default causal mask
+    attn_bias = xformers.ops.LowerTriangularMask()
+
+    # Check if we have sliding window attention
+    has_sliding_window = sliding_window is not None and sliding_window < query_length
+
+    # Transpose dimensions for xformers (Q: [b, h, s, d] -> [b, s, h, d])
+    query = query.transpose(1, 2)
+    key = key.transpose(1, 2)
+    value = value.transpose(1, 2)
+
+    # Get GQA parameters
+    num_attention_heads = module.config.num_attention_heads
+    num_key_value_heads = module.config.num_key_value_heads
+    head_dim = query.size(-1)
+    is_gqa = num_attention_heads != num_key_value_heads
+    n_groups = num_attention_heads // num_key_value_heads if is_gqa else 1
+
+    # If position_ids is provided and check all examples do not contain only 1 sequence, If tensor in increasing
+    # then we probably have one sequence, otherwise it is packed. Additionally check we are in pre-fill/training stage.
+    # Use `flash_attn_varlen_func` to prevent cross-example attention and also allow padding free approach
+    if position_ids is not None and (
+        max_length_q is not None
+        or (query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all())
+    ):
+        if cu_seq_lens_q is None or cu_seq_lens_k is None:
+            cu_seq_lens_q = get_cu_seqlens_from_pos_ids(position_ids)[0]
+            cu_seq_lens_q = cu_seq_lens_q.squeeze()
+            seq_lengths = cu_seq_lens_q[1:] - cu_seq_lens_q[:-1]
+            attn_bias = (
+                xformers.ops.fmha.attn_bias.BlockDiagonalCausalMask.from_seqlens(
+                    q_seqlen=seq_lengths.tolist(),
+                )
+            )
+        else:
+            query = query.reshape(-1, query.size(-2), query.size(-1))
+            key = key.reshape(-1, key.size(-2), key.size(-1))
+            value = value.reshape(-1, value.size(-2), value.size(-1))
+
+        # Handle GQA
+        if is_gqa:
+            key = key.repeat_interleave(n_groups, dim=2)
+            value = value.repeat_interleave(n_groups, dim=2)
+
+    elif attention_mask is not None:
+        query, key, value, _, cu_seq_lens, _ = _upad_input(
+            query, key, value, attention_mask, query_length
+        )
+        cu_seq_lens_q, cu_seq_lens_k = cu_seq_lens
+        seq_lengths = []
+        for i in range(len(cu_seq_lens_q) - 1):
+            seq_lengths.append(cu_seq_lens_q[i + 1] - cu_seq_lens_q[i])
+        attn_bias = xformers.ops.fmha.attn_bias.BlockDiagonalCausalMask.from_seqlens(
+            q_seqlen=seq_lengths,
+            kv_seqlen=seq_lengths,
+        )
+
+        # Handle GQA
+        if is_gqa:
+            key = key.repeat_interleave(n_groups, dim=2)
+            value = value.repeat_interleave(n_groups, dim=2)
+    else:
+        # Handle Group Query Attention (GQA) using view/expand approach from reference
+        key = key.view(batch_size, key_length, num_key_value_heads, 1, head_dim)
+        value = value.view(batch_size, key_length, num_key_value_heads, 1, head_dim)
+        key = key.expand(
+            batch_size, key_length, num_key_value_heads, n_groups, head_dim
+        )
+        value = value.expand(
+            batch_size, key_length, num_key_value_heads, n_groups, head_dim
+        )
+
+        if module.training:
+            key = key.reshape(batch_size, key_length, num_attention_heads, head_dim)
+            value = value.reshape(batch_size, key_length, num_attention_heads, head_dim)
+
+            if has_sliding_window:
+                query = query.view(
+                    1, batch_size * query_length, num_attention_heads, head_dim
+                )
+                key = key.view(
+                    1, batch_size * key_length, num_attention_heads, head_dim
+                )
+                value = value.view(
+                    1, batch_size * key_length, num_attention_heads, head_dim
+                )
+        else:
+            query = query.view(
+                batch_size, query_length, num_key_value_heads, n_groups, head_dim
+            )
+
+            # If we need a sliding window attention
+            if has_sliding_window:
+                query = query.view(
+                    1,
+                    batch_size * query_length,
+                    num_key_value_heads,
+                    n_groups,
+                    head_dim,
+                )
+                key = key.view(
+                    1, batch_size * key_length, num_key_value_heads, n_groups, head_dim
+                )
+                value = value.view(
+                    1, batch_size * key_length, num_key_value_heads, n_groups, head_dim
+                )
+
+    # Run the xformers attention
+    attn_output = xformers_attention(
+        query,
+        key,
+        value,
+        attn_bias=attn_bias,
+    )
+
+    attn_output = attn_output.view(
+        batch_size, -1, attn_output.size(-2), attn_output.size(-1)
+    )
+    return attn_output, None
diff --git a/src/axolotl/monkeypatch/btlm_attn_hijack_flash.py b/src/axolotl/monkeypatch/btlm_attn_hijack_flash.py
index 1275906804..2c50773929 100644
--- a/src/axolotl/monkeypatch/btlm_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/btlm_attn_hijack_flash.py
@@ -3,7 +3,6 @@
 """
 
 import importlib
-import logging
 from typing import Optional, Tuple
 
 import torch
@@ -11,7 +10,9 @@
 from flash_attn.flash_attn_interface import flash_attn_func
 from transformers import AutoConfig, AutoModelForCausalLM
 
-LOG = logging.getLogger("axolotl")
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
 
 
 def replace_btlm_attn_with_flash_attn(model_name="cerebras/btlm-3b-8k-base"):
@@ -24,9 +25,7 @@ def replace_btlm_attn_with_flash_attn(model_name="cerebras/btlm-3b-8k-base"):
         ".configuration_btlm", ".modeling_btlm"
     )
     modeling_btlm = importlib.import_module(module_name)
-    modeling_btlm.BTLMAttention._attn = (  # pylint: disable=protected-access
-        flashattn_attn
-    )
+    modeling_btlm.BTLMAttention._attn = flashattn_attn
 
 
 def flashattn_attn(
@@ -34,9 +33,9 @@ def flashattn_attn(
     query: torch.Tensor,
     key: Optional[torch.Tensor] = None,
     value: Optional[torch.Tensor] = None,
-    attention_mask: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
+    attention_mask: Optional[torch.Tensor] = None,
     head_mask: Optional[torch.Tensor] = None,
-    position_bias: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
+    position_bias: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
     softmax_scale = (
         1 / (key.size(-1) ** self.attn_scale_power) if self.scale_attn_weights else None
diff --git a/src/axolotl/monkeypatch/checkpoint_activation_offload.py b/src/axolotl/monkeypatch/checkpoint_activation_offload.py
new file mode 100644
index 0000000000..74e1a98fdf
--- /dev/null
+++ b/src/axolotl/monkeypatch/checkpoint_activation_offload.py
@@ -0,0 +1,356 @@
+"""Non-reentrant checkpoint input offloading."""
+
+from __future__ import annotations
+
+import contextlib
+from dataclasses import dataclass
+
+import torch
+from torch.autograd.graph import saved_tensors_hooks
+from transformers import GradientCheckpointingLayer, is_torch_npu_available
+
+if is_torch_npu_available():
+    import torch_npu  # noqa: F401
+
+
+@dataclass
+class CheckpointActivationOffloadStats:
+    saved_tensors_seen: int = 0
+    marked_tensors: int = 0
+    offloaded_tensors: int = 0
+    restored_tensors: int = 0
+    skipped_marked_tensors: int = 0
+    offloaded_bytes: int = 0
+    restored_bytes: int = 0
+
+
+@dataclass(frozen=True)
+class _OffloadedTensorRef:
+    tensor_id: int
+
+
+_BufferKey = tuple[tuple[int, ...], tuple[int, ...], torch.dtype, torch.layout, bool]
+_MANAGER_STACK: list["CheckpointHiddenStatesOffload"] = []
+_ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL = None
+
+
+def _current_manager() -> "CheckpointHiddenStatesOffload | None":
+    return _MANAGER_STACK[-1] if _MANAGER_STACK else None
+
+
+def patch_gradient_checkpointing_layer_marker() -> None:
+    global _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL
+
+    if _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL is not None:
+        return
+
+    _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL = GradientCheckpointingLayer.__call__
+
+    def _checkpoint_offload_call(self, *args, **kwargs):
+        manager = _current_manager()
+        if (
+            manager is not None
+            and self.training
+            and torch.is_grad_enabled()
+            and getattr(self, "gradient_checkpointing", False)
+        ):
+            hidden_states = args[0] if args else kwargs.get("hidden_states")
+            if torch.is_tensor(hidden_states):
+                manager.mark(hidden_states)
+        return _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL(self, *args, **kwargs)
+
+    GradientCheckpointingLayer.__call__ = _checkpoint_offload_call
+
+
+class CheckpointHiddenStatesOffload(saved_tensors_hooks):
+    """Offload only marked checkpoint inputs.
+
+    The marker is installed on ``GradientCheckpointingLayer.__call__`` before
+    non-reentrant checkpointing saves positional layer inputs. All other saved
+    tensors pass through unchanged, including tensors from the final norm/head
+    and any non-checkpointed modules.
+    """
+
+    def __init__(
+        self,
+        use_pin_memory: bool = True,
+        use_streams: bool = True,
+        min_offload_size: int = 1024,
+        max_fwd_stash_size: int = 2,
+        max_cpu_buffer_pool_size: int = 64,
+    ) -> None:
+        self.use_pin_memory = use_pin_memory
+        self.use_streams = use_streams
+        self.min_tensor_size_bytes = min_offload_size
+        self.max_fwd_stash_size = max_fwd_stash_size
+        self.max_cpu_buffer_pool_size = max_cpu_buffer_pool_size
+        self.stats = CheckpointActivationOffloadStats()
+        self._allowed: dict[int, int] = {}
+        self._next_id = 0
+        self._tracker: dict[
+            int,
+            tuple[
+                torch.Tensor,
+                torch.device,
+                int,
+                torch.Size,
+                tuple[int, ...],
+                _BufferKey,
+            ],
+        ] = {}
+        self._fwd_stash: dict[int, tuple[torch.Tensor, torch.Event]] = {}
+        self._cpu_buffer_pool: dict[_BufferKey, list[torch.Tensor]] = {}
+        self._cpu_buffer_pool_size = 0
+        self._pending_cpu_buffers: list[
+            tuple[_BufferKey, torch.Tensor, torch.Event]
+        ] = []
+
+        if hasattr(torch, "accelerator") and torch.accelerator.is_available():
+            self.accelerator_type = torch.accelerator.current_accelerator().type
+        elif torch.cuda.is_available():
+            self.accelerator_type = "cuda"
+        else:
+            self.accelerator_type = "cpu"
+        self.s0 = self._current_stream()
+        self.s1 = self._new_stream() if use_streams else None
+
+        super().__init__(self._pack_tensor, self._unpack_tensor)
+
+    def _current_stream(self):
+        if self.accelerator_type == "cpu":
+            return None
+        if self.accelerator_type == "xpu":
+            return torch.xpu.current_stream()
+        if self.accelerator_type == "npu":
+            return torch.npu.current_stream()
+        return torch.cuda.current_stream()
+
+    def _new_stream(self):
+        if self.accelerator_type == "cpu":
+            return None
+        if self.accelerator_type == "xpu":
+            return torch.xpu.Stream()
+        if self.accelerator_type == "npu":
+            return torch.npu.Stream()
+        return torch.cuda.Stream()
+
+    def _stream_context(self, stream):
+        if stream is None:
+            return contextlib.nullcontext()
+        if self.accelerator_type == "xpu":
+            return torch.xpu.stream(stream)
+        if self.accelerator_type == "npu":
+            return torch.npu.stream(stream)
+        return torch.cuda.stream(stream)
+
+    def mark(self, tensor: torch.Tensor) -> None:
+        self._allowed[id(tensor)] = self._allowed.get(id(tensor), 0) + 1
+        self.stats.marked_tensors += 1
+
+    def _consume_mark(self, tensor: torch.Tensor) -> bool:
+        tensor_id = id(tensor)
+        count = self._allowed.get(tensor_id, 0)
+        if count <= 0:
+            return False
+        if count == 1:
+            del self._allowed[tensor_id]
+        else:
+            self._allowed[tensor_id] = count - 1
+        return True
+
+    @staticmethod
+    def _num_bytes(tensor: torch.Tensor) -> int:
+        return tensor.element_size() * tensor.nelement()
+
+    def _next_tensor_id(self) -> int:
+        self._next_id += 1
+        return self._next_id
+
+    def _reap_forward_stash(self, new_tensor_id: int) -> None:
+        for tensor_id in list(self._fwd_stash):
+            if tensor_id > new_tensor_id - self.max_fwd_stash_size:
+                continue
+            _, event = self._fwd_stash.pop(tensor_id)
+            self.s0.wait_event(event)
+
+    def _buffer_key(self, tensor: torch.Tensor) -> _BufferKey:
+        return (
+            tuple(tensor.size()),
+            tuple(tensor.stride()),
+            tensor.dtype,
+            tensor.layout,
+            self.use_pin_memory,
+        )
+
+    def _pool_cpu_buffer(self, key: _BufferKey, tensor: torch.Tensor) -> None:
+        if self._cpu_buffer_pool_size >= self.max_cpu_buffer_pool_size:
+            return
+        self._cpu_buffer_pool.setdefault(key, []).append(tensor)
+        self._cpu_buffer_pool_size += 1
+
+    def _reap_cpu_buffer_pool(self, force: bool = False) -> None:
+        pending = self._pending_cpu_buffers
+        self._pending_cpu_buffers = []
+        for key, tensor, event in pending:
+            if force:
+                event.synchronize()
+                self._pool_cpu_buffer(key, tensor)
+            elif event.query():
+                self._pool_cpu_buffer(key, tensor)
+            else:
+                self._pending_cpu_buffers.append((key, tensor, event))
+
+    def _empty_cpu_like(self, tensor: torch.Tensor) -> tuple[torch.Tensor, _BufferKey]:
+        self._reap_cpu_buffer_pool()
+        key = self._buffer_key(tensor)
+        pool = self._cpu_buffer_pool.get(key)
+        if pool:
+            self._cpu_buffer_pool_size -= 1
+            return pool.pop(), key
+        return torch.empty_strided(
+            tuple(tensor.size()),
+            tuple(tensor.stride()),
+            dtype=tensor.dtype,
+            layout=tensor.layout,
+            device="cpu",
+            pin_memory=self.use_pin_memory,
+        ), key
+
+    def _pack_tensor(self, tensor: torch.Tensor):
+        self.stats.saved_tensors_seen += 1
+        if not self._consume_mark(tensor):
+            return tensor
+
+        num_bytes = self._num_bytes(tensor)
+        if (
+            tensor.device.type not in {"cuda", "xpu", "npu"}
+            or num_bytes < self.min_tensor_size_bytes
+            or isinstance(tensor, torch.nn.Parameter)
+            or (hasattr(torch.nn, "Buffer") and isinstance(tensor, torch.nn.Buffer))
+        ):
+            self.stats.skipped_marked_tensors += 1
+            return tensor
+
+        tensor_id = self._next_tensor_id()
+        if self.use_streams:
+            self._reap_forward_stash(tensor_id)
+            self.s1.wait_stream(self.s0)
+            stream = self.s1
+        else:
+            stream = self.s0
+
+        with self._stream_context(stream):
+            cpu_tensor, buffer_key = self._empty_cpu_like(tensor)
+            cpu_tensor.copy_(tensor.detach(), non_blocking=self.use_streams)
+
+        self._tracker[tensor_id] = (
+            cpu_tensor,
+            tensor.device,
+            tensor.storage_offset(),
+            tensor.size(),
+            tensor.stride(),
+            buffer_key,
+        )
+        if self.use_streams:
+            self._fwd_stash[tensor_id] = (tensor, self.s1.record_event())
+
+        self.stats.offloaded_tensors += 1
+        self.stats.offloaded_bytes += num_bytes
+        return _OffloadedTensorRef(tensor_id)
+
+    def _restore_tensor(self, tensor_id: int) -> torch.Tensor:
+        if tensor_id in self._fwd_stash:
+            tensor, event = self._fwd_stash.pop(tensor_id)
+            self.s0.wait_event(event)
+            cpu_tensor, *_unused, buffer_key = self._tracker.pop(tensor_id)
+            self._pool_cpu_buffer(buffer_key, cpu_tensor)
+            self.stats.restored_tensors += 1
+            self.stats.restored_bytes += self._num_bytes(tensor)
+            return tensor
+
+        cpu_tensor, device, storage_offset, shape, stride, buffer_key = (
+            self._tracker.pop(tensor_id)
+        )
+        stream = self.s1 if self.use_streams else self.s0
+        with self._stream_context(stream):
+            gpu_tensor = cpu_tensor.to(device, non_blocking=self.use_streams)
+            if (
+                gpu_tensor.storage_offset() != storage_offset
+                or gpu_tensor.stride() != stride
+            ):
+                gpu_tensor = torch.as_strided(
+                    gpu_tensor,
+                    size=shape,
+                    stride=stride,
+                    storage_offset=storage_offset,
+                )
+        if self.use_streams:
+            event = self.s1.record_event()
+            self.s0.wait_event(event)
+            gpu_tensor.record_stream(self.s0)
+            self._pending_cpu_buffers.append((buffer_key, cpu_tensor, event))
+        else:
+            self._pool_cpu_buffer(buffer_key, cpu_tensor)
+
+        self.stats.restored_tensors += 1
+        self.stats.restored_bytes += self._num_bytes(gpu_tensor)
+        return gpu_tensor
+
+    def _unpack_tensor(self, maybe_ref):
+        if isinstance(maybe_ref, _OffloadedTensorRef):
+            return self._restore_tensor(maybe_ref.tensor_id)
+        return maybe_ref
+
+    def reset(self) -> None:
+        self.stats = CheckpointActivationOffloadStats()
+        self._allowed.clear()
+        self._tracker.clear()
+        self._fwd_stash.clear()
+        self._pending_cpu_buffers.clear()
+        self._next_id = 0
+
+    def _sync_and_clear(self) -> None:
+        if self.use_streams:
+            if self.s0 is not None:
+                self.s0.synchronize()
+            if self.s1 is not None:
+                self.s1.synchronize()
+            self._reap_cpu_buffer_pool(force=True)
+        for tracked in self._tracker.values():
+            self._pool_cpu_buffer(tracked[-1], tracked[0])
+        self._allowed.clear()
+        self._tracker.clear()
+        self._fwd_stash.clear()
+        self._pending_cpu_buffers.clear()
+
+    def __enter__(self):
+        patch_gradient_checkpointing_layer_marker()
+        self.reset()
+        _MANAGER_STACK.append(self)
+        return super().__enter__()
+
+    def __exit__(self, *args, **kwargs):
+        try:
+            self._sync_and_clear()
+        finally:
+            if _MANAGER_STACK and _MANAGER_STACK[-1] is self:
+                _MANAGER_STACK.pop()
+            elif self in _MANAGER_STACK:
+                _MANAGER_STACK.remove(self)
+        return super().__exit__(*args, **kwargs)
+
+
+def get_checkpoint_hidden_states_offloading_ctx_manager(
+    use_pin_memory: bool = True,
+    use_streams: bool = True,
+    min_offload_size: int = 1024,
+    max_fwd_stash_size: int = 2,
+    max_cpu_buffer_pool_size: int = 64,
+) -> contextlib.ContextDecorator:
+    return CheckpointHiddenStatesOffload(
+        use_pin_memory=use_pin_memory,
+        use_streams=use_streams,
+        min_offload_size=min_offload_size,
+        max_fwd_stash_size=max_fwd_stash_size,
+        max_cpu_buffer_pool_size=max_cpu_buffer_pool_size,
+    )
diff --git a/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py b/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py
index 2e9364e3a5..c426344a66 100644
--- a/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py
+++ b/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py
@@ -1,14 +1,23 @@
-"""monkey patches for the dataset fetcher to handle batches of packed indexes"""
-# pylint: disable=protected-access
+"""Monkey patches for the dataset fetcher to handle batches of packed indexes."""
 
 import torch
 from torch.utils.data._utils.fetch import _BaseDatasetFetcher
 from torch.utils.data._utils.worker import _worker_loop
 
+_ORIGINAL_MAP_DATASET_FETCHER = None
+_ORIGINAL_WORKER_LOOP = None
+_IS_PATCHED = False
+
 
 class _MapDatasetFetcher(_BaseDatasetFetcher):
+    """
+    Custom dataset fetcher that handles nested batch structures from
+    MultipackBatchSampler.
+    """
+
     def fetch(self, possibly_batched_index):
         if isinstance(possibly_batched_index[0], list):
+            # Handle nested structure from MultipackBatchSampler
             data = [None for i in possibly_batched_index]
             for i, possibly_batched_index_ in enumerate(possibly_batched_index):
                 if self.auto_collation:
@@ -22,6 +31,7 @@ def fetch(self, possibly_batched_index):
                 else:
                     data[i] = self.dataset[possibly_batched_index_]
         else:
+            # Standard batch handling
             if self.auto_collation:
                 if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__:
                     data = self.dataset.__getitems__(possibly_batched_index)
@@ -33,14 +43,54 @@ def fetch(self, possibly_batched_index):
 
 
 def patch_fetchers():
+    """Apply patches to PyTorch's DataLoader components."""
     torch.utils.data._utils.fetch._MapDatasetFetcher = _MapDatasetFetcher
     torch.utils.data.dataloader._utils.fetch._MapDatasetFetcher = _MapDatasetFetcher
 
 
 def patched_worker_loop(*args, **kwargs):
+    """Worker loop that ensures patches are applied in worker processes."""
     patch_fetchers()
     return _worker_loop(*args, **kwargs)
 
 
-torch.utils.data._utils.worker._worker_loop = patched_worker_loop
-patch_fetchers()
+def apply_multipack_dataloader_patch():
+    """
+    This patch allows DataLoader to correctly process batches that contain multiple bins
+    of packed sequences.
+    """
+    # pylint: disable=global-statement
+    global _ORIGINAL_MAP_DATASET_FETCHER, _ORIGINAL_WORKER_LOOP, _IS_PATCHED
+
+    if _IS_PATCHED:
+        return
+
+    # Store original implementations
+    _ORIGINAL_MAP_DATASET_FETCHER = torch.utils.data._utils.fetch._MapDatasetFetcher
+    _ORIGINAL_WORKER_LOOP = torch.utils.data._utils.worker._worker_loop
+
+    # Apply patches
+    patch_fetchers()
+    torch.utils.data._utils.worker._worker_loop = patched_worker_loop
+
+    _IS_PATCHED = True
+
+
+def remove_multipack_dataloader_patch():
+    """Remove the monkeypatch and restore original PyTorch DataLoader behavior."""
+    # pylint: disable=global-statement
+    global _IS_PATCHED
+
+    if not _IS_PATCHED:
+        return
+
+    if _ORIGINAL_MAP_DATASET_FETCHER:
+        torch.utils.data._utils.fetch._MapDatasetFetcher = _ORIGINAL_MAP_DATASET_FETCHER
+        torch.utils.data.dataloader._utils.fetch._MapDatasetFetcher = (
+            _ORIGINAL_MAP_DATASET_FETCHER
+        )
+
+    if _ORIGINAL_WORKER_LOOP:
+        torch.utils.data._utils.worker._worker_loop = _ORIGINAL_WORKER_LOOP
+
+    _IS_PATCHED = False
diff --git a/src/axolotl/monkeypatch/deepspeed_utils.py b/src/axolotl/monkeypatch/deepspeed_utils.py
new file mode 100644
index 0000000000..d7e69e1121
--- /dev/null
+++ b/src/axolotl/monkeypatch/deepspeed_utils.py
@@ -0,0 +1,67 @@
+import importlib
+import importlib.util
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_checkpoint_wrapper_setattr():
+    """
+    Patch CheckpointWrapper to properly forward DeepSpeed attributes to wrapped modules.
+
+    This fixes the issue where CheckpointWrapper doesn't forward ds_* attributes
+    (like ds_grads_remaining) to the actual wrapped module, causing DeepSpeed
+    ZeRO-3 to fail when gradient checkpointing is enabled.
+
+    This issue occurs specifically with:
+    - QLoRA + DeepSpeed ZeRO-3
+    - gradient_checkpointing: true
+    - activation_offloading: true
+
+    References:
+    - https://github.com/deepspeedai/DeepSpeed/issues/7203
+    - https://github.com/deepspeedai/DeepSpeed/blob/38d1a9eb64c9e01e32eccc50b25ba18925287441/deepspeed/runtime/zero/parameter_offload.py#L424-L458
+    - https://github.com/axolotl-ai-cloud/axolotl/pull/3102
+    """
+
+    try:
+        from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+            CheckpointWrapper,
+        )
+
+        # Check if already patched
+        if hasattr(CheckpointWrapper, "_axolotl_setattr_patched"):
+            LOG.debug("CheckpointWrapper already patched")
+            return
+
+        original_setattr = CheckpointWrapper.__setattr__
+
+        def new_setattr(self, name: str, value) -> None:
+            if name.startswith("ds_") and hasattr(self, "_checkpoint_wrapped_module"):
+                setattr(self._checkpoint_wrapped_module, name, value)
+                LOG.debug(
+                    f"Forwarded {name} to wrapped module {type(self._checkpoint_wrapped_module).__name__}"
+                )
+            else:
+                original_setattr(self, name, value)
+
+        CheckpointWrapper.__setattr__ = new_setattr
+        CheckpointWrapper._axolotl_setattr_patched = True
+
+        LOG.info("CheckpointWrapper patched to forward DeepSpeed attributes")
+
+    except ImportError as e:
+        LOG.debug(f"CheckpointWrapper not available: {e}")
+    except Exception as e:
+        LOG.warning(f"Failed to patch CheckpointWrapper: {e}")
+
+
+def apply_deepspeed_patches():
+    """
+    Apply DeepSpeed-related patches
+    """
+    if importlib.util.find_spec("deepspeed") is not None:
+        patch_checkpoint_wrapper_setattr()
+    else:
+        LOG.debug("DeepSpeed not available, skipping patches")
diff --git a/src/axolotl/monkeypatch/fastchat_conversation_turns.py b/src/axolotl/monkeypatch/fastchat_conversation_turns.py
deleted file mode 100644
index a09bfddb4b..0000000000
--- a/src/axolotl/monkeypatch/fastchat_conversation_turns.py
+++ /dev/null
@@ -1,231 +0,0 @@
-"""
-monkeypatch to add a get_turns method
-"""
-
-import logging
-from typing import Generator, Tuple
-
-from fastchat.conversation import SeparatorStyle
-
-LOG = logging.getLogger("axolotl.monkeypatch.fastchat_conversation_turns")
-
-
-def get_prompt(self) -> str:
-    ret = ""
-    for role, msg in self.get_turns():
-        ret += role + msg
-    return ret
-
-
-def get_turns(  # pylint: disable=too-many-return-statements
-    self,
-) -> Generator[Tuple[str, str], None, None]:
-    """Get the prompt for generation."""
-    system_prompt = self.system_template.format(system_message=self.system_message)
-    if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
-        yield "", system_prompt + self.sep
-        for role, message in self.messages:
-            if message:
-                yield role + ": ", message + self.sep
-            else:
-                yield role + ":", ""
-        return
-    if self.sep_style == SeparatorStyle.ADD_COLON_TWO:
-        seps = [self.sep, self.sep2]
-        yield "", system_prompt + seps[0]
-        for i, (role, message) in enumerate(self.messages):
-            if message:
-                yield role + ": ", message + seps[i % 2]
-            else:
-                yield role + ":", ""
-        return
-    if self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
-        yield "", system_prompt + self.sep
-        for role, message in self.messages:
-            if message:
-                yield role + ": ", message + self.sep
-            else:
-                yield role + ": ", ""  # must be end with a space
-        return
-    if self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
-        yield "", "" if system_prompt == "" else system_prompt + self.sep
-        for role, message in self.messages:
-            if message:
-                yield role + "\n", message + self.sep
-            else:
-                yield role + "\n", ""
-        return
-    if self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
-        yield "", system_prompt
-        for role, message in self.messages:
-            if message:
-                yield role, message + self.sep
-            else:
-                yield role, ""
-        return
-    if self.sep_style == SeparatorStyle.NO_COLON_TWO:
-        seps = [self.sep, self.sep2]
-        yield "", system_prompt
-        for i, (role, message) in enumerate(self.messages):
-            if message:
-                yield role, message + seps[i % 2]
-            else:
-                yield role, ""
-        return
-    if self.sep_style == SeparatorStyle.RWKV:
-        yield "", system_prompt
-        for i, (role, message) in enumerate(self.messages):
-            if message:
-                yield role + ": ", message.replace("\r\n", "\n").replace(
-                    "\n\n", "\n"
-                ) + "\n\n"
-            else:
-                yield role + ":", ""
-        return
-    if self.sep_style == SeparatorStyle.LLAMA2 and self.name != "mistral":
-        if self.system_message:
-            if self.messages:
-                # For llama, the system message is incorporated into the first human instruction
-                first_role, first_msg = self.messages[0]
-                if first_role == self.roles[0]:
-                    system_prompt += first_msg
-                    self.messages.pop(0)
-            yield "", system_prompt
-        for i, (role, message) in enumerate(self.messages):
-            if message:
-                if (i % 2 == 0 and not self.system_message) or (
-                    i % 2 != 0 and self.system_message
-                ):
-                    role = "<s> " + role
-                yield role + " ", message
-            else:
-                yield role, ""
-        return
-    if self.sep_style == SeparatorStyle.LLAMA2 and self.name == "mistral":
-        contains_sys_msg = False
-        if self.system_message:
-            contains_sys_msg = True
-            if self.messages:
-                # There is no clear guidance on how to handle system messages in Mistral so we just prepend it to the first human instruction separated by a newline
-                first_role, first_msg = self.messages[0]
-                if first_role == self.roles[0]:
-                    system_prompt = self.system_template.format(
-                        system_message=" " + self.system_message
-                    )
-                    system_prompt += first_msg
-                    self.messages.pop(0)
-            yield "", system_prompt
-        for i, (role, message) in enumerate(self.messages):
-            if message and i == 0 and not contains_sys_msg:
-                yield "", system_prompt.strip() + " " + message  # if there is no system message, we need to make sure there is the a `<s> [INST]` at the beginning of the first instruction.
-            elif message:
-                yield role + " ", message
-            else:
-                yield role, ""
-        return
-    if self.sep_style == SeparatorStyle.LLAMA3:
-        if self.system_message:
-            # For llama3, the system message is NOT incorporated into the first human instruction
-            # All messages follow <|start_header_id|>' + role + '<|end_header_id|>\n\n'+ message + '<|eot_id|>
-            yield "", system_prompt
-        for i, (role, message) in enumerate(self.messages):
-            if message:
-                yield f"<|start_header_id|>{role}<|end_header_id|>\n\n", f"{message.strip()}<|eot_id|>"
-            else:
-                yield f"<|start_header_id|>{role}<|end_header_id|>\n\n", ""
-        return
-    if self.sep_style == SeparatorStyle.GEMMA:
-        if self.system_message:
-            raise ValueError("Gemma chat template does not support system messages")
-        for i, (role, message) in enumerate(self.messages):
-            prefix = "<bos>" if i == 0 else ""
-            message_str = message if message else ""
-            yield prefix + "<start_of_turn>" + role + "\n", message_str + "<end_of_turn>\n"
-        return
-    if self.sep_style == SeparatorStyle.CHATGLM:
-        # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
-        # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
-        round_add_n = 1 if self.name == "chatglm2" else 0
-        if system_prompt:
-            yield "", system_prompt + self.sep
-
-        for i, (role, message) in enumerate(self.messages):
-            if i % 2 == 0:
-                yield "", f"[Round {i//2 + round_add_n}]{self.sep}"
-
-            if message:
-                yield f"{role}：", f"{message}{self.sep}"
-            else:
-                yield f"{role}：", ""
-        return
-    if self.sep_style == SeparatorStyle.CHATML:
-        yield "", "" if system_prompt == "" else system_prompt + self.sep + "\n"
-        for role, message in self.messages:
-            if message:
-                yield role + "\n", message + self.sep + "\n"
-            else:
-                yield role + "\n", ""
-        return
-    if self.sep_style == SeparatorStyle.CHATGLM3:
-        if self.system_message:
-            yield "", system_prompt
-        for role, message in self.messages:
-            if message:
-                yield role + "\n", " " + message
-            else:
-                yield role
-        return
-    if self.sep_style == SeparatorStyle.CHATINTERN:
-        # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
-        seps = [self.sep, self.sep2]
-        yield "", system_prompt
-        for i, (role, message) in enumerate(self.messages):
-            prefix = "<s>" if i % 2 == 0 else ""
-            if message:
-                yield prefix + role + ":", message + seps[i % 2] + "\n"
-            else:
-                yield role + ":", ""
-        return
-    if self.sep_style == SeparatorStyle.DOLLY:
-        seps = [self.sep, self.sep2]
-        yield "", system_prompt
-        for i, (role, message) in enumerate(self.messages):
-            if message:
-                suffix = "\n\n" if i % 2 == 1 else ""
-                yield role + ":\n", message + seps[i % 2] + suffix
-            else:
-                yield role + ":\n", ""
-        return
-    if self.sep_style == SeparatorStyle.PHOENIX:
-        yield "", system_prompt
-        for role, message in self.messages:
-            if message:
-                yield role + ": ", "<s>" + message + "</s>"
-            else:
-                yield role + ": " + "<s>", ""
-        return
-    if self.sep_style == SeparatorStyle.ROBIN:
-        yield "", system_prompt + self.sep
-        for role, message in self.messages:
-            if message:
-                yield role + ":\n", message + self.sep
-            else:
-                yield role + ":\n", ""
-        return
-    if self.sep_style == SeparatorStyle.FALCON_CHAT:
-        if self.system_message:
-            yield "", system_prompt + self.sep
-        for role, message in self.messages:
-            if message:
-                yield role + ": ", message + self.sep
-            else:
-                yield role + ":", ""
-    else:
-        raise ValueError(f"Invalid style: {self.sep_style}")
-
-
-def add_get_turns_to_conversation():
-    import fastchat.conversation
-
-    fastchat.conversation.Conversation.get_turns = get_turns
-    fastchat.conversation.Conversation.get_prompt = get_prompt
diff --git a/src/axolotl/monkeypatch/fsdp2_qlora.py b/src/axolotl/monkeypatch/fsdp2_qlora.py
new file mode 100644
index 0000000000..6fd3ffaeb8
--- /dev/null
+++ b/src/axolotl/monkeypatch/fsdp2_qlora.py
@@ -0,0 +1,291 @@
+"""
+Monkeypatch to add Params4bit and Int8Params support to FSDP2. This enables QLoRA + FSDP2
+and 8-bit LoRA + FSDP2, as well as our LoRA / QLoRA Triton kernels to work with FSDP2.
+
+This patch modifies the _init_sharded_param and init_unsharded_param methods in FSDPParam
+to handle bitsandbytes Params4bit and Int8Params parameters, preserving their quantization
+metadata through the FSDP2 shard/unshard cycle.
+"""
+
+import importlib
+import inspect
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def apply_init_sharded_param_patch():
+    """Apply patch to FSDPParam._init_sharded_param to support Params4bit."""
+    if getattr(apply_init_sharded_param_patch, "_axolotl_patched", False):
+        return
+    from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
+
+    # Get original source
+    original_source = inspect.getsource(FSDPParam._init_sharded_param)
+    original_source, _ = detab_code(original_source)
+
+    # torch 2.12 rewrote the sharded-param construction from a two-line
+    # form (Parameter() + requires_grad_()) to a single multi-line
+    # Parameter() call with requires_grad= as a kwarg. Try the 2.12
+    # anchor first, fall back to the 2.11 form.
+    anchors_2120 = """    self.sharded_param = nn.Parameter(
+        self.to_sharded_dtensor(sharded_param),
+        requires_grad=param.requires_grad,
+    )"""
+    anchors_2110 = """    self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param))
+    self.sharded_param.requires_grad_(param.requires_grad)"""
+
+    patched_param_creation = """    import bitsandbytes as bnb
+    if isinstance(param, bnb.nn.modules.Params4bit):
+        self.sharded_param = bnb.nn.modules.Params4bit(
+            data=sharded_param,
+            requires_grad=param.requires_grad,
+            quant_state=param.quant_state,
+            blocksize=param.blocksize,
+            compress_statistics=param.compress_statistics,
+            quant_type=param.quant_type,
+            quant_storage=param.quant_storage,
+            module=param.module,
+            bnb_quantized=param.bnb_quantized,
+        )
+        self.sharded_param = self.to_sharded_dtensor(self.sharded_param)
+    elif isinstance(param, bnb.nn.modules.Int8Params):
+        self.sharded_param = bnb.nn.modules.Int8Params(
+            data=sharded_param,
+            requires_grad=param.requires_grad,
+            has_fp16_weights=param.has_fp16_weights,
+            CB=None,
+            SCB=param.SCB,
+        )
+        self.sharded_param = self.to_sharded_dtensor(self.sharded_param)
+    else:
+        self.sharded_param = nn.Parameter(
+            self.to_sharded_dtensor(sharded_param),
+            requires_grad=param.requires_grad,
+        )"""
+
+    original_param_creation = next(
+        (a for a in (anchors_2120, anchors_2110) if a in original_source),
+        None,
+    )
+
+    # Apply the replacement
+    if original_param_creation is not None:
+        patched_source = original_source.replace(
+            original_param_creation, patched_param_creation
+        )
+        patched_source = patched_source.replace(
+            "def _init_sharded_param(",
+            "def patched_init_sharded_param(",
+            1,
+        )
+
+        # Load necessary imports
+        module_name = FSDPParam.__module__
+        module = importlib.import_module(module_name)
+
+        items_to_import = []
+        for item in dir(module):
+            if item in patched_source:
+                items_to_import.append(item)
+
+        exec(  # nosec B102
+            f"from {module_name} import ({', '.join(items_to_import)})",
+            globals(),
+        )
+        exec(patched_source, globals())  # nosec B102
+
+        # Replace the method
+        FSDPParam._init_sharded_param = patched_init_sharded_param
+        apply_init_sharded_param_patch._axolotl_patched = True
+        LOG.info("Successfully applied FSDP _init_sharded_param patch")
+    else:
+        LOG.warning("Could not find target code for _init_sharded_param patching")
+
+
+def apply_init_unsharded_param_patch():
+    """Apply patch to FSDPParam.init_unsharded_param to support Params4bit."""
+    if getattr(apply_init_unsharded_param_patch, "_axolotl_patched", False):
+        return
+    from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
+
+    # Get original source
+    original_source = inspect.getsource(FSDPParam.init_unsharded_param)
+    original_source, _ = detab_code(original_source)
+
+    # torch 2.12 hoisted the unsharded-param construction out of the
+    # first-all-gather `else:` branch up to method-body level, so the 2.11
+    # anchor (8-space, inside else) no longer matches. The replacement must be
+    # indented to match whichever anchor is found, so each anchor carries its
+    # own. Try the 2.12 anchor first, fall back to the 2.11 form.
+    anchor_replacement_pairs = [
+        (
+            """    self._unsharded_param = nn.Parameter(
+        unsharded_param, requires_grad=self.sharded_param.requires_grad
+    )""",
+            """    import bitsandbytes as bnb
+    local_tensor = self.sharded_param._local_tensor
+    if isinstance(local_tensor, bnb.nn.modules.Params4bit):
+        self._unsharded_param = bnb.nn.modules.Params4bit(
+            data=unsharded_param,
+            requires_grad=self.sharded_param.requires_grad,
+            quant_state=local_tensor.quant_state,
+            blocksize=local_tensor.blocksize,
+            compress_statistics=local_tensor.compress_statistics,
+            quant_type=local_tensor.quant_type,
+            quant_storage=local_tensor.quant_storage,
+            module=local_tensor.module,
+            bnb_quantized=local_tensor.bnb_quantized,
+        )
+    elif isinstance(local_tensor, bnb.nn.modules.Int8Params):
+        self._unsharded_param = bnb.nn.modules.Int8Params(
+            data=unsharded_param,
+            requires_grad=self.sharded_param.requires_grad,
+            has_fp16_weights=local_tensor.has_fp16_weights,
+            CB=unsharded_param,
+            SCB=local_tensor.SCB,
+        )
+    else:
+        self._unsharded_param = nn.Parameter(
+            unsharded_param, requires_grad=self.sharded_param.requires_grad
+        )""",
+        ),
+        (
+            """        self._unsharded_param = nn.Parameter(
+            unsharded_param, requires_grad=self.sharded_param.requires_grad
+        )""",
+            """        import bitsandbytes as bnb
+        local_tensor = self.sharded_param._local_tensor
+        if isinstance(local_tensor, bnb.nn.modules.Params4bit):
+            self._unsharded_param = bnb.nn.modules.Params4bit(
+                data=unsharded_param,
+                requires_grad=self.sharded_param.requires_grad,
+                quant_state=local_tensor.quant_state,
+                blocksize=local_tensor.blocksize,
+                compress_statistics=local_tensor.compress_statistics,
+                quant_type=local_tensor.quant_type,
+                quant_storage=local_tensor.quant_storage,
+                module=local_tensor.module,
+                bnb_quantized=local_tensor.bnb_quantized,
+            )
+        elif isinstance(local_tensor, bnb.nn.modules.Int8Params):
+            self._unsharded_param = bnb.nn.modules.Int8Params(
+                data=unsharded_param,
+                requires_grad=self.sharded_param.requires_grad,
+                has_fp16_weights=local_tensor.has_fp16_weights,
+                CB=unsharded_param,
+                SCB=local_tensor.SCB,
+            )
+        else:
+            self._unsharded_param = nn.Parameter(
+                unsharded_param, requires_grad=self.sharded_param.requires_grad
+            )""",
+        ),
+    ]
+
+    original_param_creation, patched_param_creation = next(
+        ((a, p) for a, p in anchor_replacement_pairs if a in original_source),
+        (None, None),
+    )
+
+    # Apply the replacement
+    if original_param_creation is not None:
+        patched_source = original_source.replace(
+            original_param_creation, patched_param_creation
+        )
+        patched_source = patched_source.replace(
+            "def init_unsharded_param(",
+            "def patched_init_unsharded_param(",
+            1,
+        )
+
+        # Load necessary imports
+        module_name = FSDPParam.__module__
+        module = importlib.import_module(module_name)
+
+        items_to_import = []
+        for item in dir(module):
+            if item in patched_source:
+                items_to_import.append(item)
+
+        exec(  # nosec B102
+            f"from {module_name} import ({', '.join(items_to_import)})",
+            globals(),
+        )
+        exec(patched_source, globals())  # nosec B102
+
+        # Replace the method
+        FSDPParam.init_unsharded_param = patched_init_unsharded_param
+        apply_init_unsharded_param_patch._axolotl_patched = True
+        LOG.info("Successfully applied FSDP init_unsharded_param patch")
+    else:
+        LOG.warning("Could not find target code for patching")
+
+
+def apply_linear8bitlt_save_patch():
+    """Patch Linear8bitLt._save_to_state_dict to handle DTensor-wrapped Int8Params.
+
+    After FSDP2 sharding, Linear8bitLt.weight is a DTensor wrapping Int8Params.
+    BnB's _save_to_state_dict accesses self.weight.SCB directly, but DTensor
+    doesn't proxy custom attribute access to its _local_tensor. This patch
+    temporarily unwraps the DTensor during saving so BnB can find the SCB attribute.
+    """
+    if getattr(apply_linear8bitlt_save_patch, "_axolotl_patched", False):
+        return
+    import bitsandbytes as bnb
+    from torch.distributed.tensor import DTensor
+
+    original_save = bnb.nn.Linear8bitLt._save_to_state_dict
+
+    def _patched_save_to_state_dict(self, destination, prefix, keep_vars):
+        # Use _parameters dict directly to bypass nn.Module.__setattr__ type check.
+        weight = self._parameters["weight"]
+        unwrapped = False
+        if isinstance(weight, DTensor) and hasattr(weight, "_local_tensor"):
+            self._parameters["weight"] = weight._local_tensor
+            unwrapped = True
+        try:
+            original_save(self, destination, prefix, keep_vars)
+        finally:
+            if unwrapped:
+                self._parameters["weight"] = weight
+
+    bnb.nn.Linear8bitLt._save_to_state_dict = _patched_save_to_state_dict
+    apply_linear8bitlt_save_patch._axolotl_patched = True
+    LOG.info("Patched Linear8bitLt._save_to_state_dict for DTensor compatibility")
+
+
+def apply_init_dtype_attrs_patch():
+    """Prevent FSDP2 mixed precision from casting non-float quantized params.
+
+    When mixed precision is enabled (e.g., bf16), FSDP2's init_dtype_attrs sets
+    param_dtype=bf16 for ALL params. During all-gather, _to_dtype_if_needed casts
+    the sharded param to param_dtype. For non-float params (uint8 packed 4-bit,
+    int8 quantized) without FSDP2 extensions, this destroys the quantized data.
+
+    Params4bit handles this via fsdp_pre/post_all_gather extensions, but our
+    parametrize-based expert quantization uses plain nn.Parameter(uint8/int8)
+    without extensions.
+    """
+    if getattr(apply_init_dtype_attrs_patch, "_axolotl_patched", False):
+        return
+    from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
+
+    original_init_dtype_attrs = FSDPParam.init_dtype_attrs
+
+    def patched_init_dtype_attrs(self, mp_policy):
+        original_init_dtype_attrs(self, mp_policy)
+        # Skip casting non-float quantized params (uint8/int8) without FSDP2
+        # extensions — the parametrization chain handles dequantization.
+        if self.param_dtype is not None and not self.sharded_param.is_floating_point():
+            local = self.sharded_param
+            if hasattr(local, "_local_tensor"):
+                local = local._local_tensor
+            if not hasattr(local, "fsdp_pre_all_gather"):
+                self.param_dtype = None
+
+    FSDPParam.init_dtype_attrs = patched_init_dtype_attrs
+    apply_init_dtype_attrs_patch._axolotl_patched = True
+    LOG.info("Patched FSDPParam.init_dtype_attrs for non-float quantized params")
diff --git a/src/axolotl/monkeypatch/gemma4_hybrid_mask.py b/src/axolotl/monkeypatch/gemma4_hybrid_mask.py
new file mode 100644
index 0000000000..7e5b305570
--- /dev/null
+++ b/src/axolotl/monkeypatch/gemma4_hybrid_mask.py
@@ -0,0 +1,280 @@
+"""Hybrid attention mask fix for Gemma 4 (standard and unified).
+
+Gemma 4 has full-attention (global) layers with ``head_dim=512`` which
+exceeds flash-attention-2's supported size. Axolotl's hybrid-attention
+patch in ``patch_manager._apply_gemma_hybrid_attention`` works around
+this by forcing ``_attn_implementation="sdpa"`` on each global layer's
+``self_attn.config``, leaving sliding-window layers on FA2.
+
+The per-layer config override alone is insufficient, however:
+``Gemma4TextModel.forward`` builds a single ``causal_mask_mapping`` dict
+using the **model-level** config and passes the mapped mask to each
+decoder layer. With FA2 still set at the model level, the ``full_attention``
+entry in that mapping is a 2D mask (FA2 format), but SDPA needs a 4D mask.
+The global layers then fail with::
+
+    RuntimeError: The expanded size of the tensor (S) must match the existing
+    size (B) at non-singleton dimension 2. Target sizes: [B, H, S, S]. Tensor
+    sizes: [B, S]
+
+...when the sequence length grows past roughly 7k tokens.
+
+This module fixes the symptom by monkey-patching ``create_causal_mask`` in
+the model's *module namespace* — NOT the original in ``masking_utils``. The
+wrapper forces ``_attn_implementation="sdpa"`` on a shallow-copied config
+before calling through, so the ``full_attention`` mask built inside the
+text backbone's ``forward`` is always 4D/SDPA-compatible.
+``create_sliding_window_causal_mask`` is left alone, so sliding-window
+layers continue to receive FA2-format masks.
+
+``gemma4_unified`` reproduces the same mixed sliding/global architecture
+(``global_head_dim=512``) in its own ``modeling_gemma4_unified`` namespace,
+so both namespaces are patched when present.
+
+The patch is idempotent. Install once per process, before any Gemma 4
+forward pass runs.
+"""
+
+from __future__ import annotations
+
+import copy
+import importlib
+from typing import Any
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_PATCH_APPLIED = False
+
+# Attention-interface name for Gemma-4 global (full_attention) layers. They have head_dim=512, so FA2
+# can't serve them and the model (FA2 at the top level) hands them no mask, relying on cu_seqlens that
+# only the FA2 sliding layers consume. With sample packing that means the global layers would attend
+# ACROSS document boundaries (pure causal). This impl rebuilds the block-diagonal-causal mask from
+# position_ids so the globals respect doc boundaries, on the memory-efficient SDPA backend.
+GLOBAL_PACKED_SDPA = "sdpa_global_packed"
+
+
+# When set, the head_dim=512 global layers use the Triton flash_d512 kernel (fwd+bwd, varlen) instead
+# of the SDPA efficient backend (~2x faster at head_dim 512). Set from cfg.flash_attn_d512.
+def set_flash_d512(enabled: bool) -> None:
+    """Backwards-compat shim: the head_dim>256 routing is now the generic large_head_attention
+    capability. True -> 'auto' (flash only on packed rows, the proven win)."""
+    from axolotl.monkeypatch.attention.large_head import set_large_head_policy
+
+    set_large_head_policy("auto" if enabled else "sdpa")
+
+
+def _packing_block_causal_mask(position_ids, dtype, device):
+    """Block-diagonal causal additive mask [B,1,S,S] from packed position_ids (which reset to 0 at
+    each document start). -inf across document boundaries and for non-causal positions, 0 elsewhere."""
+    import torch
+
+    if position_ids.dim() == 1:
+        position_ids = position_ids[None]
+    Bz, Sz = position_ids.shape
+    doc = (position_ids == 0).cumsum(-1)  # [B,S] document index (1-based)
+    same_doc = doc[:, :, None] == doc[:, None, :]  # [B,S,S]
+    causal = torch.ones(Sz, Sz, dtype=torch.bool, device=device).tril()[None]  # [1,S,S]
+    allow = same_doc & causal
+    mask = torch.zeros(Bz, 1, Sz, Sz, dtype=dtype, device=device)
+    mask.masked_fill_(~allow[:, None], torch.finfo(dtype).min)
+    return mask
+
+
+def _register_global_packed_sdpa() -> None:
+    """Register the packing-aware global-layer attention impl (block-diagonal mask + efficient SDPA)."""
+    from transformers.integrations.sdpa_attention import sdpa_attention_forward
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    if GLOBAL_PACKED_SDPA in ALL_ATTENTION_FUNCTIONS.valid_keys():
+        return
+
+    def sdpa_global_packed_forward(module, query, key, value, attention_mask, **kwargs):
+        # The top-level FA2 path leaves these layers maskless and carries packing only via cu_seqlens
+        # (consumed by FA2, not SDPA). Rebuild the block-diagonal mask from position_ids so the global
+        # layers don't cross document boundaries. Single-document rows stay maskless (is_causal).
+        from axolotl.monkeypatch.attention.large_head import flash_d512_route
+
+        position_ids = kwargs.get("position_ids")
+        # The generic large-head router takes head_dim>256 packed rows through the Triton flash
+        # kernel (per the large_head_attention policy); ~2.7x the 4D-mask SDPA, ~3x less memory than
+        # nested-tensor SDPA. It declines (returns None) for single-doc/policy=sdpa -> SDPA below.
+        if attention_mask is None:
+            routed = flash_d512_route(
+                module, query, key, value, kwargs.get("scaling"), position_ids
+            )
+            if routed is not None:
+                return routed
+        # Packing detection: static under a declared config (compile-clean), runtime probe otherwise.
+        pid = None
+        if attention_mask is None:
+            from axolotl.monkeypatch.attention.large_head import (
+                _multidoc_position_ids,
+            )
+
+            pid = _multidoc_position_ids(position_ids)
+        # Packed without the kernel -> block-diagonal mask so globals respect doc boundaries.
+        # Single-document -> mask stays None (SDPA is_causal, the fast path).
+        if pid is not None:
+            attention_mask = _packing_block_causal_mask(pid, query.dtype, query.device)
+        return sdpa_attention_forward(
+            module, query, key, value, attention_mask, **kwargs
+        )
+
+    ALL_ATTENTION_FUNCTIONS.register(GLOBAL_PACKED_SDPA, sdpa_global_packed_forward)
+    LOG.info(
+        "gemma4_hybrid_mask: registered '%s' (block-diagonal packing mask for head_dim=512 global "
+        "layers so they respect document boundaries under sample packing)",
+        GLOBAL_PACKED_SDPA,
+    )
+
+
+# Each Gemma 4 variant fully redefines ``create_causal_mask`` in its own module
+# namespace (gemma4_unified does NOT modular-import from gemma4), so both must be
+# patched independently.
+_TARGET_MODULES = (
+    "transformers.models.gemma4.modeling_gemma4",
+    "transformers.models.gemma4_unified.modeling_gemma4_unified",
+)
+
+
+def _patch_module_create_causal_mask(module: Any) -> bool:
+    """Wrap ``create_causal_mask`` in a single module namespace.
+
+    Re-entry is prevented by the module-level ``_PATCH_APPLIED`` flag in
+    :func:`patch_gemma4_hybrid_mask`, so this does not guard per-module.
+    Returns ``True`` if patched, ``False`` if the namespace has no
+    ``create_causal_mask`` binding.
+    """
+    if not hasattr(module, "create_causal_mask"):
+        LOG.warning(
+            "gemma4_hybrid_mask: %s has no 'create_causal_mask' binding, "
+            "skipping. Transformers API may have changed.",
+            module.__name__,
+        )
+        return False
+
+    original = module.create_causal_mask
+
+    def hybrid_create_causal_mask(config: Any, *args: Any, **kwargs: Any):
+        """Force SDPA format for the full-attention mask.
+
+        The global layers were patched to SDPA by
+        ``_apply_gemma_hybrid_attention``, so their mask must be 4D. The
+        original ``create_causal_mask`` dispatches on
+        ``config._attn_implementation``; we shadow that with a local override
+        on a shallow copy so the caller's config is left intact (the
+        sliding-window factory still reads FA2 from it).
+        """
+        sdpa_config = copy.copy(config)
+        sdpa_config._attn_implementation = "sdpa"
+        return original(sdpa_config, *args, **kwargs)
+
+    # Preserve the original reference on the wrapper for tests / teardown.
+    hybrid_create_causal_mask._axolotl_original = original  # type: ignore[attr-defined]
+    module.create_causal_mask = hybrid_create_causal_mask
+    LOG.info(
+        "gemma4_hybrid_mask: patched %s.create_causal_mask to force SDPA-format "
+        "masks for full-attention layers",
+        module.__name__,
+    )
+    return True
+
+
+def _patch_use_gqa_head_dim_guard() -> bool:
+    """Stop ``enable_gqa`` from forcing the MATH SDPA backend on large-head-dim layers.
+
+    ``sdpa_attention_forward`` enables ``enable_gqa=True`` whenever ``attention_mask is None``
+    (``use_gqa_in_sdpa``), with no head_dim check. But SDPA's flash/efficient GQA path only
+    supports head_dim <= 256; at head_dim > 256 ``enable_gqa`` silently falls back to the MATH
+    kernel, which materializes the full [H, S, S] scores. Repeating KV instead keeps the
+    memory-efficient backend. For Gemma-4's head_dim=512 global layers this is ~2.9 GiB -> ~0.2 GiB
+    per layer with identical math (repeat_kv == GQA).
+    """
+    try:
+        import transformers.integrations.sdpa_attention as sdpa_mod
+    except ImportError:
+        return False
+    original = sdpa_mod.use_gqa_in_sdpa
+    if getattr(original, "_axolotl_head_dim_guarded", False):
+        return True
+
+    # *args/**kwargs: transformers 5.13 adds a `value` positional arg
+    def use_gqa_in_sdpa_guarded(attention_mask, key, *args, **kwargs):
+        # head_dim > 256 -> enable_gqa drops to the MATH backend; force repeat_kv (efficient) instead.
+        if key.shape[-1] > 256:
+            return False
+        return original(attention_mask, key, *args, **kwargs)
+
+    use_gqa_in_sdpa_guarded._axolotl_head_dim_guarded = True  # type: ignore[attr-defined]
+    use_gqa_in_sdpa_guarded._axolotl_original = original  # type: ignore[attr-defined]
+    sdpa_mod.use_gqa_in_sdpa = use_gqa_in_sdpa_guarded
+    LOG.info(
+        "gemma4_hybrid_mask: guarded use_gqa_in_sdpa (head_dim>256 -> repeat_kv, not enable_gqa) "
+        "to keep the memory-efficient SDPA backend on head_dim=512 global layers"
+    )
+    return True
+
+
+def patch_gemma4_hybrid_mask() -> bool:
+    """Install the Gemma 4 hybrid-attention mask fix across all variants.
+
+    Returns ``True`` if at least one namespace was patched, ``False`` if none
+    of the target modules could be imported (e.g. transformers version predates
+    Gemma 4) — in which case nothing is done and the caller can continue
+    unaffected.
+    """
+    global _PATCH_APPLIED
+    if _PATCH_APPLIED:
+        return True
+
+    patched_any = False
+    for module_path in _TARGET_MODULES:
+        try:
+            module = importlib.import_module(module_path)
+        except ImportError:
+            LOG.debug(
+                "gemma4_hybrid_mask: %s not importable, skipping. This is fine "
+                "for non-Gemma4 training.",
+                module_path,
+            )
+            continue
+        if _patch_module_create_causal_mask(module):
+            patched_any = True
+
+    if not patched_any:
+        return False
+
+    # Only touch global SDPA state once we know a Gemma4 namespace was actually patched —
+    # otherwise _PATCH_APPLIED stays False and unpatch() would skip cleaning these up.
+    _patch_use_gqa_head_dim_guard()
+    _register_global_packed_sdpa()
+    _PATCH_APPLIED = True
+    return True
+
+
+def unpatch_gemma4_hybrid_mask() -> None:
+    """Restore the original ``create_causal_mask`` in every namespace. Tests."""
+    global _PATCH_APPLIED
+    if not _PATCH_APPLIED:
+        return
+    try:
+        import transformers.integrations.sdpa_attention as sdpa_mod
+
+        guarded = getattr(sdpa_mod, "use_gqa_in_sdpa", None)
+        original = getattr(guarded, "_axolotl_original", None)
+        if original is not None:
+            sdpa_mod.use_gqa_in_sdpa = original
+    except ImportError:
+        pass
+    for module_path in _TARGET_MODULES:
+        try:
+            module = importlib.import_module(module_path)
+        except ImportError:
+            continue
+        current = getattr(module, "create_causal_mask", None)
+        original = getattr(current, "_axolotl_original", None)
+        if original is not None:
+            module.create_causal_mask = original
+    _PATCH_APPLIED = False
diff --git a/src/axolotl/monkeypatch/gemma4_loss_kwargs.py b/src/axolotl/monkeypatch/gemma4_loss_kwargs.py
new file mode 100644
index 0000000000..d00a5a9daa
--- /dev/null
+++ b/src/axolotl/monkeypatch/gemma4_loss_kwargs.py
@@ -0,0 +1,41 @@
+"""Flip ``accepts_loss_kwargs`` to True on Gemma 4 (Unified) ForConditionalGeneration.
+
+They inherit ``accepts_loss_kwargs = False`` from PaliGemma (whose loss filtered
+logits/labels by attention_mask). Gemma 4's loss is the stock ``ForCausalLMLoss``
+with no such filtering, so the flag wrongly makes the Trainer withhold
+``num_items_in_batch`` and mis-normalize the loss under gradient accumulation.
+Install before ``Trainer.__init__`` reads the flag.
+"""
+
+from __future__ import annotations
+
+import importlib
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_TARGETS = (
+    ("transformers.models.gemma4.modeling_gemma4", "Gemma4ForConditionalGeneration"),
+    (
+        "transformers.models.gemma4_unified.modeling_gemma4_unified",
+        "Gemma4UnifiedForConditionalGeneration",
+    ),
+)
+
+
+def patch_gemma4_accepts_loss_kwargs() -> None:
+    """Set ``accepts_loss_kwargs=True`` on Gemma 4 (Unified) ForConditionalGeneration."""
+    for module_path, cls_name in _TARGETS:
+        try:
+            cls = getattr(importlib.import_module(module_path), cls_name)
+        except (ImportError, AttributeError):
+            continue
+        if getattr(cls, "accepts_loss_kwargs", None) is True:
+            continue
+        cls.accepts_loss_kwargs = True
+        LOG.info(
+            "Set %s.accepts_loss_kwargs=True so the Trainer forwards "
+            "num_items_in_batch (correct gradient-accumulation loss normalization).",
+            cls_name,
+        )
diff --git a/src/axolotl/monkeypatch/gradient_checkpointing/__init__.py b/src/axolotl/monkeypatch/gradient_checkpointing/__init__.py
new file mode 100644
index 0000000000..b58bbb67c8
--- /dev/null
+++ b/src/axolotl/monkeypatch/gradient_checkpointing/__init__.py
@@ -0,0 +1,59 @@
+"""custom checkpointing utils"""
+
+import importlib
+from functools import partial
+
+from packaging import version
+
+from axolotl.monkeypatch.gradient_checkpointing.offload_cpu import (  # noqa: F401
+    CPU_Offloaded_Gradient_Checkpointer,
+)
+from axolotl.monkeypatch.gradient_checkpointing.offload_disk import (
+    Disco,
+)
+
+transformers_version = version.parse(importlib.metadata.version("transformers"))
+if transformers_version > version.parse("4.51.3"):
+    from transformers.modeling_layers import GradientCheckpointingLayer
+
+    def uses_gc_layers(decoder_layer):
+        return isinstance(decoder_layer.func.__self__, GradientCheckpointingLayer)
+
+else:
+
+    def uses_gc_layers(_):
+        return False
+
+
+def hf_grad_checkpoint_offload_wrapper(decoder_layer, *args, use_reentrant=None):
+    if uses_gc_layers(decoder_layer):
+        return CPU_Offloaded_Gradient_Checkpointer.apply(
+            decoder_layer,
+            *args,
+        )
+
+    return CPU_Offloaded_Gradient_Checkpointer.apply(
+        (
+            decoder_layer.func.__self__
+            if isinstance(decoder_layer, partial)
+            else decoder_layer.__self__
+        ),
+        *args,
+    )
+
+
+def hf_grad_checkpoint_disk_offload_wrapper(decoder_layer, *args, use_reentrant=None):
+    if uses_gc_layers(decoder_layer):
+        return Disco.apply(
+            decoder_layer,
+            *args,
+        )
+
+    return Disco.apply(
+        (
+            decoder_layer.func.__self__
+            if isinstance(decoder_layer, partial)
+            else decoder_layer.__self__
+        ),
+        *args,
+    )
diff --git a/src/axolotl/utils/gradient_checkpointing/unsloth.py b/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py
similarity index 56%
rename from src/axolotl/utils/gradient_checkpointing/unsloth.py
rename to src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py
index fbe8346be2..886441196e 100644
--- a/src/axolotl/utils/gradient_checkpointing/unsloth.py
+++ b/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py
@@ -1,4 +1,4 @@
-"""Unsloth checkpointing"""
+"""CPU offloaded checkpointing"""
 
 # Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
 #
@@ -13,19 +13,36 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+import inspect
+
 import torch
+from packaging import version
+from torch.utils.checkpoint import (
+    set_device_states,
+)
+
+# support different pytorch versions
+has_device_type = "device_type" in inspect.signature(set_device_states).parameters
 
+torch_version = version.parse(torch.__version__)
 
-class Unsloth_Offloaded_Gradient_Checkpointer(  # pylint: disable=invalid-name
-    torch.autograd.Function
-):
+if torch_version < version.parse("2.4.0"):
+    torch_cuda_amp_custom_fwd = torch.cuda.amp.custom_fwd
+    torch_cuda_amp_custom_bwd = torch.cuda.amp.custom_bwd
+else:
+    torch_cuda_amp_custom_fwd = torch.amp.custom_fwd(device_type="cuda")
+    torch_cuda_amp_custom_bwd = torch.amp.custom_bwd(device_type="cuda")
+
+
+class CPU_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
     """
     Saves VRAM by smartly offloading to RAM.
     Tiny hit to performance, since we mask the movement via non blocking calls.
     """
 
     @staticmethod
-    @torch.cuda.amp.custom_fwd
+    @torch_cuda_amp_custom_fwd
     def forward(ctx, forward_function, hidden_states, *args):
         saved_hidden_states = hidden_states.to("cpu", non_blocking=True)
         with torch.no_grad():
@@ -36,17 +53,20 @@ def forward(ctx, forward_function, hidden_states, *args):
         return output
 
     @staticmethod
-    @torch.cuda.amp.custom_bwd
+    @torch_cuda_amp_custom_bwd
     def backward(ctx, dY):
         (hidden_states,) = ctx.saved_tensors
         hidden_states = hidden_states.to("cuda", non_blocking=True).detach()
         hidden_states.requires_grad = True
         with torch.enable_grad():
-            (output,) = ctx.forward_function(hidden_states, *ctx.args)
+            output = ctx.forward_function(hidden_states, *ctx.args)
+            # Newer HF models (e.g. Qwen3MoE) using GradientCheckpointingLayer
+            # return a plain tensor, not a tuple.  Older models return tuples
+            # like (hidden_states, present_kv, ...).  Unwrap if needed.
+            if isinstance(output, (tuple, list)):
+                (output,) = output
         torch.autograd.backward(output, dY)
         return (
             None,
             hidden_states.grad,
-        ) + (
-            None,
-        ) * len(ctx.args)
+        ) + (None,) * len(ctx.args)
diff --git a/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py b/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py
new file mode 100644
index 0000000000..220799fbf8
--- /dev/null
+++ b/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py
@@ -0,0 +1,532 @@
+"""
+DISCO - DIsk-based Storage and Checkpointing with Optimized prefetching
+"""
+
+# Copyright 2025 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import atexit
+import concurrent.futures
+import os
+import queue
+import shutil
+import tempfile
+import threading
+import time
+import uuid
+from collections import deque
+from concurrent.futures import Future
+from typing import Dict
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+torch_cuda_amp_custom_fwd = torch.amp.custom_fwd(device_type="cuda")
+torch_cuda_amp_custom_bwd = torch.amp.custom_bwd(device_type="cuda")
+
+# Setup logger
+logger = get_logger(__name__)
+
+
+class DiskOffloadManager:
+    """
+    Manages offloaded tensors and handles prefetching in a separate thread.
+    Includes synchronization to prevent race conditions.
+    """
+
+    def __init__(
+        self,
+        prefetch_size: int = 3,
+        prefetch_to_gpu: bool = True,
+        save_workers: int = 4,
+    ):
+        """
+        Args:
+            prefetch_size: Maximum number of tensors to prefetch in the background.
+            prefetch_to_gpu: Whether to prefetch tensors directly to GPU memory.
+            save_workers: Maximum number of concurrent save operations.
+        """
+        self.temp_dir = tempfile.mkdtemp(prefix="disco_")
+
+        # Track tensor paths and their status
+        self.tensor_paths: deque = deque()  # Ordered history of tensor paths (LIFO)
+        self.file_locks: Dict[
+            str, threading.Lock
+        ] = {}  # Maps file_path -> threading.Lock()
+        # Maps file_path -> status ("saving", "ready", "prefetching", "loaded", "deleted")
+        self.file_status: Dict[str, str] = {}
+
+        self.max_prefetch = prefetch_size
+        self.prefetch_to_gpu = prefetch_to_gpu
+
+        # Thread synchronization
+        self.manager_lock = threading.RLock()  # Used for thread-safe operations
+
+        # Prefetch queue and cache
+        self.prefetch_queue: queue.Queue = queue.Queue()
+        self.prefetch_cache: Dict[str, torch.Tensor] = {}  # Maps file_path -> tensor
+
+        # Save queue and thread pool
+        self.save_queue: queue.Queue = queue.Queue()
+        self.save_pool = concurrent.futures.ThreadPoolExecutor(max_workers=save_workers)
+        self.save_futures: Dict[str, Future] = {}
+        self.save_semaphore = threading.Semaphore(
+            save_workers * 2
+        )  # Limit concurrent save operations
+
+        # Start prefetch worker thread
+        self.stop_event = threading.Event()
+        # start multiple threads for prefetching
+        self.prefetch_worker_count = 2
+        self.prefetch_workers = []
+        for _ in range(self.prefetch_worker_count):
+            worker = threading.Thread(target=self._prefetch_worker, daemon=True)
+            worker.start()
+            self.prefetch_workers.append(worker)
+
+        # Start save worker thread
+        self.save_worker = threading.Thread(target=self._save_worker, daemon=True)
+        self.save_worker.start()
+        self.idx = 0
+
+        atexit.register(self.cleanup)
+
+    def _save_worker(self):
+        """Background thread that processes the save queue"""
+        while not self.stop_event.is_set():
+            try:
+                save_item = self.save_queue.get(timeout=0.5)
+                if save_item is None:
+                    continue
+
+                tensor, file_path = save_item
+
+                # Submit the save task to the thread pool
+                future = self.save_pool.submit(
+                    self._save_tensor_to_disk, tensor, file_path
+                )
+                with self.manager_lock:
+                    self.save_futures[file_path] = future
+
+                self.save_queue.task_done()
+
+            except queue.Empty:
+                time.sleep(0.01)  # Small sleep to prevent CPU spinning
+                continue
+
+    def _save_tensor_to_disk(self, tensor: torch.Tensor, file_path: str):
+        """Actually save the tensor to disk"""
+        try:
+            # Save tensor to disk
+            cpu_tensor = tensor.detach().cpu()
+            torch.save(cpu_tensor, file_path)
+            del cpu_tensor
+
+            with self.manager_lock:
+                # Mark file as ready
+                self.file_status[file_path] = "ready"
+
+            # Release semaphore
+            self.save_semaphore.release()
+
+            return True
+        except FileNotFoundError as e:
+            logger.error(f"Error saving tensor to {file_path}: {e}")
+            with self.manager_lock:
+                self.file_status[file_path] = "error"
+
+            # Release semaphore
+            self.save_semaphore.release()
+
+            return False
+
+    def _prefetch_worker(self):
+        """Background thread that loads tensors from disk ahead of time"""
+        while not self.stop_event.is_set():
+            try:
+                file_path = self.prefetch_queue.get(timeout=0.5)
+                if file_path is None:
+                    continue
+
+                # Check if file is available and not already in cache
+                with self.manager_lock:
+                    if (
+                        file_path not in self.file_status
+                        or self.file_status[file_path] == "deleted"
+                    ):
+                        self.prefetch_queue.task_done()
+                    if file_path in self.prefetch_cache:
+                        self.prefetch_queue.task_done()
+                        continue
+
+                    # If file is still being saved, wait for it
+                    if (
+                        self.file_status[file_path] == "saving"
+                        and file_path in self.save_futures
+                    ):
+                        # Re-queue this prefetch request with a little delay
+                        self.prefetch_queue.task_done()
+                        time.sleep(0.1)
+                        self.prefetch_queue.put(file_path)
+                        continue
+
+                    # Mark file as being prefetched
+                    self.file_status[file_path] = "prefetching"
+
+                # Load tensor from disk and store in cache
+                try:
+                    if os.path.exists(file_path):
+                        if self.prefetch_to_gpu:
+                            tensor = torch.load(
+                                file_path,
+                                map_location=torch.device("cuda"),
+                                weights_only=True,
+                            )
+                        else:
+                            tensor = torch.load(file_path, weights_only=True)
+
+                        with self.manager_lock:
+                            self.prefetch_cache[file_path] = tensor
+                            self.file_status[file_path] = "ready"
+                    else:
+                        with self.manager_lock:
+                            if self.file_status.get(file_path) != "deleted":
+                                logger.warning(
+                                    f"Prefetch error: File not found {file_path}"
+                                )
+                                self.file_status[file_path] = "missing"
+
+                except FileNotFoundError as e:
+                    with self.manager_lock:
+                        if self.file_status.get(file_path) != "deleted":
+                            logger.warning(f"Prefetch error for {file_path}: {e}")
+                            self.file_status[file_path] = "error"
+
+                self.prefetch_queue.task_done()
+
+            except queue.Empty:
+                time.sleep(0.01)  # Small sleep to prevent CPU spinning
+                continue
+
+    def save_tensor(self, tensor: torch.Tensor):
+        """Save tensor to disk asynchronously and return file path with thread-safe operations"""
+        # Generate unique file path
+        self.idx += 1
+        file_path: str = os.path.join(
+            self.temp_dir, f"{self.idx:06d}-{uuid.uuid4()}.pt"
+        )
+
+        with self.manager_lock:
+            # Mark file as being saved
+            self.file_locks[file_path] = threading.Lock()
+            self.file_status[file_path] = "saving"
+            # Add to history
+            self.tensor_paths.append(file_path)
+
+        # Acquire semaphore to limit concurrent save operations
+        self.save_semaphore.acquire()
+        # Queue tensor for saving in background
+        self.save_queue.put((tensor.detach(), file_path))
+
+        return file_path
+
+    def wait_for_save(self, file_path, timeout=None) -> None:
+        """Wait for a tensor to be saved to disk"""
+        start_time = time.time()
+        while timeout is None or time.time() - start_time < timeout:
+            with self.manager_lock:
+                if self.file_status.get(file_path) == "ready":
+                    return
+                if self.file_status.get(file_path) in ["error", "missing", "deleted"]:
+                    return
+
+                if file_path in self.save_futures:
+                    future = self.save_futures[file_path]
+                    if future.done():
+                        return
+
+            # Small sleep to prevent CPU spinning
+            time.sleep(0.01)
+
+        # Timeout
+        logger.warning(f"Timeout waiting for tensor to be saved: {file_path}")
+        return
+
+    def load_tensor(self, file_path, target_device="cuda"):
+        """Load tensor from disk or prefetch cache with proper synchronization"""
+        # Wait for tensor to be saved if it's still in progress
+        self.wait_for_save(file_path)
+
+        tensor = None
+
+        # Try to get from cache first
+        with self.manager_lock:
+            # Check if tensor is already in cache
+            if file_path in self.prefetch_cache:
+                tensor = self.prefetch_cache[file_path]
+                del self.prefetch_cache[file_path]
+                self.file_status[file_path] = "loaded"
+
+        if tensor is not None:
+            # Ensure tensor is on correct device
+            if target_device != "cpu" and tensor.device.type == "cpu":
+                tensor = tensor.to(target_device, non_blocking=True)
+            return tensor
+
+        # If not in cache, load directly from disk
+        try:
+            if not os.path.exists(file_path):
+                logger.error(f"File not found for loading: {file_path}")
+                raise FileNotFoundError(f"File not found: {file_path}")
+
+            tensor = torch.load(file_path, weights_only=True)
+
+            with self.manager_lock:
+                self.file_status[file_path] = "loaded"
+
+            if target_device != "cpu":
+                tensor = tensor.to(target_device, non_blocking=True)
+
+            return tensor
+
+        except Exception as e:
+            logger.error(f"Error loading tensor from {file_path}: {e}")
+            raise
+
+    def _safe_delete_file(self, file_path):
+        """Safely delete a file with proper synchronization"""
+        with self.manager_lock:
+            # Make sure any save operation is completed
+            if file_path in self.save_futures:
+                future = self.save_futures[file_path]
+                try:
+                    if not future.done():
+                        future.cancel()
+                    del self.save_futures[file_path]
+                except FileNotFoundError as e:
+                    logger.warning(
+                        f"Error canceling save operation for {file_path}: {e}"
+                    )
+
+            # Only delete if file exists and is not being prefetched
+            status = self.file_status.get(file_path)
+            if status in ["ready", "loaded", "error", "missing"]:
+                try:
+                    if os.path.exists(file_path):
+                        os.remove(file_path)
+                    self.file_status[file_path] = "deleted"
+                    return True
+                except FileNotFoundError as e:
+                    logger.warning(f"Error deleting file {file_path}: {e}")
+            return False
+
+    def trigger_prefetch(self, n=None):
+        """Trigger prefetching of the next N tensors with proper synchronization"""
+        if n is None:
+            n = self.max_prefetch
+
+        prefetch_paths = []
+        with self.manager_lock:
+            # Find files that are ready to be prefetched (not already in cache or being prefetched)
+            for path in reversed(self.tensor_paths):
+                if (
+                    path not in self.prefetch_cache
+                    and self.file_status.get(path) == "ready"
+                ):
+                    prefetch_paths.append(path)
+                    if len(prefetch_paths) >= n:
+                        break
+
+        # Queue files for prefetching
+        for path in prefetch_paths:
+            self.prefetch_queue.put(path)
+
+    def cleanup_tensor(self, file_path: str):
+        """Clean up a specific tensor file after it's been used"""
+        with self.manager_lock:
+            if file_path in self.tensor_paths:
+                self.tensor_paths.remove(file_path)
+
+            # Remove from prefetch cache if present
+            if file_path in self.prefetch_cache:
+                del self.prefetch_cache[file_path]
+
+            # Remove from save futures if present
+            if file_path in self.save_futures:
+                future = self.save_futures[file_path]
+                if not future.done():
+                    future.cancel()
+                del self.save_futures[file_path]
+
+        # Try to delete the file
+        self._safe_delete_file(file_path)
+
+    def cleanup(self):
+        """Clean up all temp files and stop prefetch thread with proper synchronization"""
+        self.stop_event.set()
+
+        # Cancel all pending save operations
+        with self.manager_lock:
+            for _, future in self.save_futures.items():
+                if not future.done():
+                    future.cancel()
+            self.save_futures.clear()
+
+        # Drain the save queue
+        while not self.save_queue.empty():
+            try:
+                self.save_queue.get_nowait()
+                self.save_queue.task_done()
+            except queue.Empty:
+                break
+
+        # Shutdown the save pool
+        self.save_pool.shutdown(wait=False)
+
+        # Join the save worker thread
+        if self.save_worker.is_alive():
+            self.save_worker.join(timeout=2.0)
+
+        # Join the prefetch worker threads
+        for thread in self.prefetch_workers:
+            if thread.is_alive():
+                thread.join(timeout=2.0)
+
+        # Clear cache and remove all temporary files
+        with self.manager_lock:
+            self.prefetch_cache.clear()
+            paths_to_delete = list(self.tensor_paths)
+            self.tensor_paths.clear()
+
+        # Delete all temporary files
+        for path in paths_to_delete:
+            self._safe_delete_file(path)
+
+        # Remove temp directory
+        try:
+            if os.path.exists(self.temp_dir):
+                shutil.rmtree(self.temp_dir, ignore_errors=True)
+        except FileNotFoundError as e:
+            logger.warning(f"Error removing temporary directory {self.temp_dir}: {e}")
+
+
+class Disco(torch.autograd.Function):
+    """
+    Disco: DIsk-based Storage and Checkpointing with Optimized prefetching
+    Advanced disk-based gradient checkpointer with prefetching.
+    """
+
+    # Shared manager instance across all checkpointing operations
+    _manager = None
+
+    @staticmethod
+    def get_instance(prefetch_size=1, prefetch_to_gpu=True, save_workers=4):
+        """Get or create the offload manager"""
+        if Disco._manager is None:
+            Disco._manager = DiskOffloadManager(
+                prefetch_size=prefetch_size,
+                prefetch_to_gpu=prefetch_to_gpu,
+                save_workers=save_workers,
+            )
+        return Disco._manager
+
+    @staticmethod
+    @torch_cuda_amp_custom_fwd
+    def forward(
+        ctx,
+        forward_function,
+        hidden_states,
+        *args,
+        prefetch_size=1,
+        prefetch_to_gpu=True,
+        save_workers=4,
+    ):
+        """Forward pass that offloads activations to disk asynchronously"""
+        # Get or create the manager
+        manager = Disco.get_instance(
+            prefetch_size=prefetch_size,
+            prefetch_to_gpu=prefetch_to_gpu,
+            save_workers=save_workers,
+        )
+
+        # Save tensor to disk asynchronously
+        file_path = manager.save_tensor(hidden_states)
+
+        # Run forward pass immediately without waiting for save to complete
+        with torch.no_grad():
+            output = forward_function(hidden_states, *args)
+
+        # Store what we need for backward
+        ctx.save_for_backward(torch.tensor([0]))  # Dummy tensor
+        ctx.file_path = file_path
+        ctx.forward_function = forward_function
+        ctx.args = args
+
+        return output
+
+    @staticmethod
+    @torch_cuda_amp_custom_bwd
+    def backward(ctx, *grad_outputs):
+        """Backward pass that loads activations from disk with prefetching"""
+        # Get the manager
+        manager = Disco._manager
+
+        # Trigger prefetching for future tensors
+        # This happens at the start of backward, so should have time to complete
+        manager.trigger_prefetch()
+
+        # Load hidden states from disk or prefetch cache
+        file_path = ctx.file_path
+        try:
+            # Ensure the file is saved before we try to load it
+            manager.wait_for_save(file_path)
+
+            hidden_states = manager.load_tensor(file_path)
+            hidden_states.requires_grad = True
+
+            # Compute gradients
+            with torch.enable_grad():
+                output = ctx.forward_function(hidden_states, *ctx.args)
+
+                # Handle tuple outputs properly
+                if isinstance(output, tuple):
+                    if len(grad_outputs) == len(output):
+                        torch.autograd.backward(output, grad_outputs)
+                    else:
+                        torch.autograd.backward(output, grad_outputs[0])
+                else:
+                    torch.autograd.backward(output, grad_outputs[0])
+
+            # Clean up the file after we're done with it
+            manager.cleanup_tensor(file_path)
+
+            return (
+                (
+                    None,  # forward_function
+                    hidden_states.grad,  # hidden_states grad
+                )
+                + (None,) * len(ctx.args)  # for each arg
+                + (
+                    None,  # prefetch_size
+                    None,  # prefetch_to_gpu
+                    None,  # save_workers
+                )
+            )
+
+        except Exception as e:
+            logger.error(f"Error in backward pass: {e}")
+            # Clean up the file even on error
+            manager.cleanup_tensor(file_path)
+            raise
diff --git a/src/axolotl/monkeypatch/kernelize_fixes.py b/src/axolotl/monkeypatch/kernelize_fixes.py
new file mode 100644
index 0000000000..79dc13e446
--- /dev/null
+++ b/src/axolotl/monkeypatch/kernelize_fixes.py
@@ -0,0 +1,73 @@
+"""Repairs for transformers' ``model.kernelize()`` under ``use_kernels=True``.
+
+``kernelize()`` swaps each module's ``_hidden_kernels`` entries for hub
+kernels, but two upstream defects crash it (transformers 5.10): ~30
+architectures (gemma4, qwen3.5, glm4, olmo, ...) stash a bare function it
+refuses to register, and gpt-oss's rotary ``Func`` keeps a deprecated
+``position_ids`` parameter that fails the kernels library's signature check
+against ``kernels-community/rotary``.
+
+Wraps ``PreTrainedModel.kernelize`` to repair the stashes right before the
+swap: bare functions are dropped (the model's forward still calls them
+directly) and ``position_ids`` is removed from rotary signature *metadata*
+(call behavior unchanged). Both repairs no-op once fixed upstream.
+"""
+
+import inspect
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_ORIG_KERNELIZE = None
+
+
+def _fix_hidden_kernels(module):
+    import torch.nn as nn
+
+    for name, fn in list(module.__dict__.get("_hidden_kernels", {}).items()):
+        if not isinstance(fn, nn.Module):
+            del module._hidden_kernels[name]
+        elif getattr(fn, "kernel_layer_name", None) == "rotary_pos_emb":
+            forward = type(fn).forward
+            signature = inspect.signature(forward)
+            if "position_ids" in signature.parameters:
+                forward.__signature__ = signature.replace(
+                    parameters=[
+                        p
+                        for p in signature.parameters.values()
+                        if p.name != "position_ids"
+                    ]
+                )
+
+
+def patch_kernelize_fixes() -> bool:
+    global _ORIG_KERNELIZE
+    if _ORIG_KERNELIZE is not None:
+        return True
+
+    from transformers.modeling_utils import PreTrainedModel
+
+    orig = getattr(PreTrainedModel, "kernelize", None)
+    if orig is None:
+        LOG.warning("kernelize_fixes: PreTrainedModel.kernelize not found, skipping")
+        return False
+
+    def kernelize(self, *args, **kwargs):
+        self.apply(_fix_hidden_kernels)
+        return orig(self, *args, **kwargs)
+
+    PreTrainedModel.kernelize = kernelize
+    _ORIG_KERNELIZE = orig
+    LOG.info("kernelize_fixes: patched PreTrainedModel.kernelize")
+    return True
+
+
+def unpatch_kernelize_fixes() -> None:
+    global _ORIG_KERNELIZE
+    if _ORIG_KERNELIZE is None:
+        return
+    from transformers.modeling_utils import PreTrainedModel
+
+    PreTrainedModel.kernelize = _ORIG_KERNELIZE
+    _ORIG_KERNELIZE = None
diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
index dda5da2b7a..60ef7e3b40 100644
--- a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -1,50 +1,24 @@
 """Flash attention monkey patch for llama model"""
 
-# copied from https://github.com/lm-sys/FastChat/blob/main/fastchat/train/llama_flash_attn_monkey_patch.py
+import importlib.util
 
-import logging
-import warnings
-from functools import partial
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
 import transformers
-from einops import rearrange
-from flash_attn.bert_padding import pad_input, unpad_input
-from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.models.llama.modeling_llama import LlamaAttention
-from transformers.models.llama.modeling_llama import (
-    LlamaDecoderLayer as OriginalLlamaDecoderLayer,
-)
-from transformers.models.llama.modeling_llama import (
-    LlamaMLP,
-    apply_rotary_pos_emb,
-    repeat_kv,
-)
-from xformers.ops import SwiGLU
+from transformers.models.llama.modeling_llama import LlamaMLP
 
-from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids, set_module_name
+from axolotl.monkeypatch.utils import set_module_name
+from axolotl.utils.logging import get_logger
 
-try:
-    from flash_attn.flash_attn_interface import (  # pylint: disable=ungrouped-imports
-        flash_attn_kvpacked_func,
-        flash_attn_varlen_kvpacked_func,
-        flash_attn_varlen_qkvpacked_func,
-    )
-except ImportError:
-    from flash_attn.flash_attn_interface import (
-        flash_attn_unpadded_kvpacked_func as flash_attn_varlen_kvpacked_func,
-    )
-    from flash_attn.flash_attn_interface import (
-        flash_attn_unpadded_qkvpacked_func as flash_attn_varlen_qkvpacked_func,
-    )
+LOG = get_logger(__name__)
 
 
-LOG = logging.getLogger("axolotl")
+def is_xformers_available() -> bool:
+    return importlib.util.find_spec("xformers") is not None
 
 
 def is_xformers_swiglu_available() -> bool:
+    if not is_xformers_available():
+        return False
+
     from xformers.ops.common import get_xformers_operator
 
     try:
@@ -57,6 +31,11 @@ def is_xformers_swiglu_available() -> bool:
 
 
 def replace_llama_mlp_with_swiglu(model):
+    if is_xformers_swiglu_available():
+        from axolotl.monkeypatch.xformers_ import FusedMLP
+    else:
+        raise RuntimeError("xformers SwiGLU not available for this environment")
+
     for name, module in model.named_modules():
         if isinstance(module, LlamaMLP):
             mlp = FusedMLP(
@@ -65,865 +44,29 @@ def replace_llama_mlp_with_swiglu(model):
             set_module_name(model, name, mlp)
 
 
-def replace_llama_qkv_with_fused(model):
-    for name, module in model.named_modules():
-        if isinstance(module, LlamaAttention):
-            qkv = FusedAttention(
-                module.config,
-                module.q_proj,
-                module.k_proj,
-                module.v_proj,
-                module.o_proj,
-            )
-            set_module_name(model, name, qkv)
-
-
-def replace_llama_attn_with_flash_attn(
-    packed: Optional[bool] = False,
-    cross_entropy: Optional[bool] = False,
-    rms_norm: Optional[bool] = False,
-    use_shifted_sparse_attn: Optional[bool] = False,
-):
-    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (  # pylint: disable=protected-access
-        _prepare_decoder_attention_mask
-    )
-    if use_shifted_sparse_attn:
-        transformers.models.llama.modeling_llama.LlamaAttention.forward = (
-            flashattn_forward_with_s2attn
-        )
-    else:
-        transformers.models.llama.modeling_llama.LlamaAttention.forward = (
-            flashattn_forward
-        )
-
-    if packed:
-        transformers.models.llama.modeling_llama.LlamaDecoderLayer = LlamaDecoderLayer
-        transformers.models.llama.modeling_llama.LlamaModel.forward = (
-            llama_model_forward
-        )
-
-    # skip only if explicitly disabled
-    if cross_entropy:
-        try:
-            from flash_attn.losses.cross_entropy import CrossEntropyLoss
-
-            LOG.info("patching with flash_attn.losses.cross_entropy")
-            transformers.models.llama.modeling_llama.CrossEntropyLoss = partial(
-                CrossEntropyLoss, inplace_backward=True
-            )
-        except ImportError:
-            LOG.info(
-                "optimized flash-attention CrossEntropyLoss not found (run `pip install 'git+https://github.com/Dao-AILab/flash-attention.git#egg=xentropy_cuda_lib&subdirectory=csrc/xentropy'`)"
-            )
-
-    # skip only if explicitly disabled
-    if rms_norm:
-        try:
-            from flash_attn.ops.rms_norm import RMSNorm
-
-            class LlamaRMSNorm(RMSNorm):
-                """Patched LLamaRMSNorm"""
-
-                def __init__(self, hidden_size, eps=1e-6):
-                    super().__init__(hidden_size, eps=eps)
-
-            LOG.info("patching with flash_attn.ops.rms_norm")
-            transformers.models.llama.modeling_llama.LlamaRMSNorm = LlamaRMSNorm
-        except ImportError:
-            LOG.info(
-                "optimized flash-attention RMSNorm not found (run `pip install 'git+https://github.com/Dao-AILab/flash-attention.git#egg=dropout_layer_norm&subdirectory=csrc/layer_norm'`)"
-            )
-
-
-class FusedAttention(LlamaAttention):
-    """
-    Fused QKV Attention layer for incrementally improved training efficiency
-    """
-
-    def __init__(
-        self,
-        config,
-        q: torch.nn.Linear,  # pylint: disable=invalid-name
-        k: torch.nn.Linear,  # pylint: disable=invalid-name
-        v: torch.nn.Linear,  # pylint: disable=invalid-name
-        o: torch.nn.Linear,  # pylint: disable=invalid-name
-    ):
-        super().__init__(config)
-        self.config = config
-        self.init_device = next(iter(q.state_dict().values())).device
-
-        # define equivalent fused qkv projection
-        self.out_features: List[int] = [q.out_features, k.out_features, v.out_features]
-        self.qkv_proj = torch.nn.Linear(
-            q.in_features, sum(self.out_features), device=self.init_device, bias=False
-        )
-        self.o_proj = o
-
-        # overwrite initialized weights with pretrained weights
-        self.qkv_proj.weight.data = torch.cat(
-            (q.weight.data, k.weight.data, v.weight.data), dim=0
-        )
-
-    def _post_training(self, model, name):
-        q_proj, k_proj, v_proj = torch.split(
-            self.qkv_proj.weight.data, self.out_features, dim=0
-        )
-
-        new_attn = LlamaAttention(self.config)
-        new_attn.q_proj.weight.data = q_proj
-        new_attn.k_proj.weight.data = k_proj
-        new_attn.v_proj.weight.data = v_proj
-        new_attn.o_proj.weight.data = self.o_proj.weight.data
-
-        set_module_name(model, name, new_attn)
-
-
-class FusedMLP(torch.nn.Module):
-    """
-    Fused MLP layer for incrementally improved training efficiency
-    """
-
-    def __init__(
-        self,
-        config,
-        gate_proj: torch.nn.Linear,
-        up_proj: torch.nn.Linear,
-        down_proj: torch.nn.Linear,
-    ):
-        super().__init__()
-        self.config = config
-        self.swiglu = SwiGLU(
-            in_features=config.hidden_size,
-            hidden_features=config.intermediate_size,
-            bias=False,
-            _pack_weights=True,
-        )
-        # overwrite initialized weights with pretrained weights
-        self.swiglu.w12.weight.data = torch.cat(
-            (gate_proj.weight.data, up_proj.weight.data), dim=0
-        )
-        self.swiglu.w3.weight.data = down_proj.weight.data
-
-    def _post_training(self, model, name):
-        w1, w2 = torch.split(  # pylint: disable=invalid-name
-            self.swiglu.w12.weight.data, self.config.intermediate_size, dim=0
-        )
-
-        # Assign the split weights back to the original layers
-        new_mlp = LlamaMLP(self.config)
-        new_mlp.gate_proj.weight.data = w1
-        new_mlp.up_proj.weight.data = w2
-        new_mlp.down_proj.weight.data = self.swiglu.w3.weight.data
-
-        set_module_name(model, name, new_mlp)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:  # pylint: disable=invalid-name
-        return self.swiglu(x)
-
-
-# Disable the transformation of the attention mask in LlamaModel as the flash attention
-# requires the attention mask to be the same as the key_padding_mask
-def _prepare_decoder_attention_mask(
-    self,
-    attention_mask,
-    input_shape,
-    inputs_embeds,
-    past_key_values_length,
-):  # pylint: disable=unused-argument
-    # [bsz, seq_len]
-    return attention_mask
-
-
-GROUP_SIZE_RATIO = 1 / 4
-
-
-def flashattn_forward_with_s2attn(
-    self,
-    hidden_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.Tensor] = None,
-    past_key_value: Optional[Tuple[torch.Tensor]] = None,
-    output_attentions: bool = False,
-    use_cache: bool = False,
-    padding_mask: Optional[torch.LongTensor] = None,  # pylint: disable=unused-argument
-    cu_seqlens: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
-    max_seqlen: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
-) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-    """Input shape: Batch x Time x Channel
-
-    From: https://github.com/dvlab-research/LongLoRA/blob/main/llama_attn_replace.py
-
-    attention_mask: [bsz, q_len]
-
-    `cu_seqlens` will be ignored if provided
-    `max_seqlen` will be ignored if provided
-    """
-    if output_attentions:
-        warnings.warn(
-            "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
-        )
-
-    bsz, q_len, _ = hidden_states.size()
-
-    query_states = (
-        self.q_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    key_states = (
-        self.k_proj(hidden_states)
-        .view(bsz, q_len, self.num_key_value_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    value_states = (
-        self.v_proj(hidden_states)
-        .view(bsz, q_len, self.num_key_value_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    # [bsz, q_len, nh, hd]
-    # [bsz, nh, q_len, hd]
-    # pylint: disable=duplicate-code
-
-    cos, sin = self.rotary_emb(value_states, position_ids=position_ids)
-    query_states, key_states = apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, position_ids
-    )
-
-    # Past Key value support
-    if past_key_value is not None:
-        # reuse k, v, self_attention
-        key_states = torch.cat([past_key_value[0], key_states], dim=2)
-        value_states = torch.cat([past_key_value[1], value_states], dim=2)
-
-    past_key_value = (key_states, value_states) if use_cache else None
-
-    # repeat k/v heads if n_kv_heads < n_heads
-    key_states = repeat_kv(key_states, self.num_key_value_groups)
-    value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-    # Flash attention codes from
-    # https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attention.py
-
-    # transform the data into the format required by flash attention
-    qkv = torch.stack(
-        [query_states, key_states, value_states], dim=2
-    )  # [bsz, nh, 3, q_len, hd]
-    qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
-
-    # We have disabled _prepare_decoder_attention_mask in LlamaModel
-    # the attention_mask should be the same as the key_padding_mask
-
-    key_padding_mask = attention_mask.repeat(2, 1)
-    nheads = qkv.shape[-2]
-    # shift
-
-    group_size = int(q_len * GROUP_SIZE_RATIO)
-    if q_len % group_size > 0:
-        raise ValueError(
-            f"q_len {q_len} should be divisible by group size {group_size}."
-        )
-
-    qkv = (
-        qkv.reshape(bsz, q_len, 3, 2, self.num_heads // 2, self.head_dim)
-        .permute(0, 3, 1, 2, 4, 5)
-        .reshape(bsz * 2, q_len, 3, self.num_heads // 2, self.head_dim)
-    )
-    x = rearrange(  # pylint: disable=invalid-name
-        qkv, "b s three h d -> b s (three h d)"
-    )
-    x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)
-    cu_q_len_tmp = torch.arange(
-        0, max_s, group_size, device=key_padding_mask.device, dtype=cu_q_lens.dtype
-    )
-    cu_q_len_tmp = torch.stack([cu_q_len_tmp, cu_q_len_tmp + group_size // 2]).repeat(
-        bsz, 1
-    ) + cu_q_lens[:-1].unsqueeze(-1)
-    cu_q_lens = torch.cat([cu_q_len_tmp, cu_q_lens[1:].unsqueeze(-1)], dim=-1).view(-1)
-
-    x_unpad = rearrange(
-        x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads // 2
-    )
-    output_unpad = flash_attn_varlen_qkvpacked_func(
-        x_unpad, cu_q_lens, group_size, 0.0, softmax_scale=None, causal=True
-    )
-    output = rearrange(
-        pad_input(
-            rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, bsz * 2, q_len
-        ),
-        "b s (h d) -> b s h d",
-        h=nheads // 2,
-    )
-    output = (
-        output.reshape(bsz, 2, q_len, nheads // 2, self.head_dim)
-        .transpose(1, 2)
-        .reshape(bsz, q_len, nheads, self.head_dim)
-    )
-    return self.o_proj(rearrange(output, "b s h d -> b s (h d)")), None, past_key_value
-
-
-def flashattn_forward(
-    self,
-    hidden_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.Tensor] = None,
-    past_key_value: Optional[Tuple[torch.Tensor]] = None,
-    output_attentions: bool = False,
-    use_cache: bool = False,
-    padding_mask: Optional[torch.LongTensor] = None,  # pylint: disable=unused-argument
-    cu_seqlens: Optional[torch.Tensor] = None,
-    max_seqlen: Optional[torch.Tensor] = None,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-    """Input shape: Batch x Time x Channel
-
-    attention_mask: [bsz, q_len]
-    """
-    # pylint: disable=duplicate-code
-    bsz, q_len, _ = hidden_states.size()
-
-    if not hasattr(self, "pretraining_tp"):
-        self.pretraining_tp = 1
-
-    if self.pretraining_tp > 1:
-        key_value_slicing = (
-            self.num_key_value_heads * self.head_dim
-        ) // self.pretraining_tp
-        query_slices = self.q_proj.weight.split(
-            (self.num_heads * self.head_dim) // self.pretraining_tp, dim=0
-        )
-        key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
-        value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
-
-        query_states = [
-            F.linear(hidden_states, query_slices[i]) for i in range(self.pretraining_tp)
-        ]
-        query_states = torch.cat(query_states, dim=-1)
-
-        key_states = [
-            F.linear(hidden_states, key_slices[i]) for i in range(self.pretraining_tp)
-        ]
-        key_states = torch.cat(key_states, dim=-1)
-
-        value_states = [
-            F.linear(hidden_states, value_slices[i]) for i in range(self.pretraining_tp)
-        ]
-        value_states = torch.cat(value_states, dim=-1)
-
-    else:
-        if isinstance(self, FusedAttention):
-            query_states, key_states, value_states = self.qkv_proj(hidden_states).split(
-                self.out_features, dim=-1
-            )
-        else:
-            query_states = self.q_proj(hidden_states)
-            key_states = self.k_proj(hidden_states)
-            value_states = self.v_proj(hidden_states)
-
-    query_states = query_states.view(
-        bsz, q_len, self.num_heads, self.head_dim
-    ).transpose(1, 2)
-    key_states = key_states.view(
-        bsz, q_len, self.num_key_value_heads, self.head_dim
-    ).transpose(1, 2)
-    value_states = value_states.view(
-        bsz, q_len, self.num_key_value_heads, self.head_dim
-    ).transpose(1, 2)
-    # [bsz, q_len, nh, hd]
-    # [bsz, nh, q_len, hd]
-
-    cos, sin = self.rotary_emb(value_states, position_ids=position_ids)
-    query_states, key_states = apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, position_ids
-    )
-    # [bsz, nh, t, hd]
-
-    if past_key_value is not None:
-        # reuse k, v, self_attention
-        key_states = torch.cat([past_key_value[0], key_states], dim=2)
-        value_states = torch.cat([past_key_value[1], value_states], dim=2)
-
-    past_key_value = (key_states, value_states) if use_cache else None
-
-    # repeat k/v heads if n_kv_heads < n_heads
-    key_states = repeat_kv(key_states, self.num_key_value_groups)
-    value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-    if output_attentions:
-        warnings.warn(
-            "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
-        )
-
-    #
-    # flash-attn v2 start
-    #
-
-    if self.training:
-        # during training q,k,v always have same seqlen
-        assert key_states.shape == query_states.shape
-        is_causal = True
-    else:
-        # turn off FA causal mask after first inference autoregressive iteration
-        # only on first autoregressive step q,k,v have same seqlen
-        is_causal = key_states.shape == query_states.shape
-
-    dropout_rate = 0.0 if not self.training else getattr(self, "attention_dropout", 0.0)
-
-    if cu_seqlens is not None and max_seqlen is not None and cu_seqlens.dim() == 1:
-        # special handling using sample packing
-        qkv = torch.stack(
-            [query_states, key_states, value_states], dim=2
-        )  # [bsz, nh, 3, q_len, hd]
-        qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
-        qkv = rearrange(qkv, "b s ... -> (b s) ...")
-
-        output = flash_attn_varlen_qkvpacked_func(
-            qkv,
-            cu_seqlens,
-            max_seqlen,
-            dropout_p=dropout_rate,
-            softmax_scale=None,
-            causal=True,
-        )
-        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
-    elif query_states.shape == key_states.shape:
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
-        qkv_unpad, cu_seqlens_q, max_seqlen_q, _, output_pad_fn = generate_qkv(
-            query_states,
-            key_states,
-            value_states,
-            qkvpacked=True,
-            # We have disabled _prepare_decoder_attention_mask in LlamaModel
-            # the attention_mask should be the same as the key_padding_mask
-            key_padding_mask=attention_mask,
-            query_padding_mask=attention_mask[:, -query_states.size(1) :]
-            if attention_mask is not None
-            else None,
-        )
-        output_unpad = flash_attn_varlen_qkvpacked_func(
-            qkv_unpad,
-            cu_seqlens_q,
-            max_seqlen_q,
-            dropout_p=dropout_rate,
-            softmax_scale=None,
-            causal=is_causal,
-        )
-        output = output_pad_fn(output_unpad)
-    else:
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
-        if attention_mask is None or attention_mask.all().item():
-            output = flash_attn_kvpacked_func(
-                query_states,
-                torch.stack([key_states, value_states], 2),
-                dropout_p=dropout_rate,
-                causal=is_causal,
-            )
-        else:
-            (  # pylint: disable=unbalanced-tuple-unpacking
-                q_unpad,
-                kv_unpad,
-                cu_seqlens_q,
-                cu_seqlens_k,
-                max_seqlen_q,
-                max_seqlen_k,
-                _,
-                _,
-                output_pad_fn,
-            ) = generate_qkv(
-                query_states,
-                key_states,
-                value_states,
-                kvpacked=True,
-                key_padding_mask=attention_mask,
-                query_padding_mask=attention_mask[:, -query_states.size(1) :]
-                if attention_mask is not None
-                else None,
-            )
-            if q_unpad.dtype != kv_unpad.dtype:
-                kv_unpad = kv_unpad.to(q_unpad.dtype)
-            output_unpad = flash_attn_varlen_kvpacked_func(
-                q_unpad,
-                kv_unpad,
-                cu_seqlens_q,
-                cu_seqlens_k,
-                max_seqlen_q,
-                max_seqlen_k,
-                dropout_p=dropout_rate,
-                softmax_scale=None,
-                causal=is_causal,
-            )
-            output = output_pad_fn(output_unpad)
-
-    attn_output = output
-    if attn_output.size() != (bsz, q_len, self.num_heads, self.head_dim):
-        raise ValueError(
-            f"`attn_output` should be of size {(bsz, q_len, self.num_heads, self.head_dim)}, but is"
-            f" {attn_output.size()}"
-        )
-    attn_output = rearrange(attn_output, "b s h d -> b s (h d)")
-
-    #
-    # flash-attn v2 end
-    #
-
-    if self.pretraining_tp > 1:
-        attn_output = attn_output.split(self.hidden_size // self.pretraining_tp, dim=2)
-        o_proj_slices = self.o_proj.weight.split(
-            self.hidden_size // self.pretraining_tp, dim=1
-        )
-        attn_output = sum(
-            F.linear(attn_output[i], o_proj_slices[i])
-            for i in range(self.pretraining_tp)
-        )
-    else:
-        attn_output = self.o_proj(attn_output)
-
-    return attn_output, None, past_key_value
-
-
-# based on https://github.com/Dao-AILab/flash-attention/blob/364a5b/tests/test_flash_attn.py#L38
-def generate_qkv(
-    q,
-    k,
-    v,
-    query_padding_mask=None,
-    key_padding_mask=None,
-    kvpacked=False,
-    qkvpacked=False,
-):  # pylint: disable=invalid-name,unnecessary-lambda-assignment
-    """
-    Arguments:
-        q: (batch_size, seqlen_q, nheads, d)
-        k: (batch_size, seqlen_k, nheads_k, d)
-        v: (batch_size, seqlen_k, nheads_k, d)
-        query_padding_mask: (batch_size, seqlen), bool
-        key_padding_mask: (batch_size, seqlen), bool
-    """
-    assert not (kvpacked and qkvpacked)
-    batch_size, seqlen_q, nheads, d = q.shape
-    _, seqlen_k, nheads_k, _ = k.shape
-    assert k.shape == (batch_size, seqlen_k, nheads_k, d)
-    assert v.shape == (batch_size, seqlen_k, nheads_k, d)
-
-    if query_padding_mask is not None:
-        q_unpad, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(
-            q, query_padding_mask
-        )
-
-        output_pad_fn = lambda output_unpad: pad_input(  # noqa: E731
-            output_unpad, indices_q, batch_size, seqlen_q
-        )
-
-    else:
-        q_unpad = rearrange(q, "b s h d -> (b s) h d")
-        cu_seqlens_q = torch.arange(
-            0,
-            (batch_size + 1) * seqlen_q,
-            step=seqlen_q,
-            dtype=torch.int32,
-            device=q_unpad.device,
-        )
-        max_seqlen_q = seqlen_q
-
-        output_pad_fn = lambda output_unpad: rearrange(  # noqa: E731
-            output_unpad, "(b s) h d -> b s h d", b=batch_size
-        )
-
-    if key_padding_mask is not None:
-        k_unpad, _, cu_seqlens_k, max_seqlen_k = unpad_input(k, key_padding_mask)
-        v_unpad, _, _, _ = unpad_input(v, key_padding_mask)
-    else:
-        k_unpad = rearrange(k, "b s h d -> (b s) h d")
-        v_unpad = rearrange(v, "b s h d -> (b s) h d")
-        cu_seqlens_k = torch.arange(
-            0,
-            (batch_size + 1) * seqlen_k,
-            step=seqlen_k,
-            dtype=torch.int32,
-            device=k_unpad.device,
-        )
-        max_seqlen_k = seqlen_k
-
-    if qkvpacked:
-        assert nheads == nheads_k
-        qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)
-        qkv = torch.stack([q, k, v], dim=2)
-        return (qkv_unpad, cu_seqlens_q, max_seqlen_q, qkv, output_pad_fn)
-
-    if kvpacked:
-        kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)
-        kv = torch.stack([k, v], dim=2)
-        return (
-            q_unpad,
-            kv_unpad,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            max_seqlen_q,
-            max_seqlen_k,
-            q,
-            kv,
-            output_pad_fn,
-        )
-
-    return (
-        q_unpad,
-        k_unpad,
-        v_unpad,
-        cu_seqlens_q,
-        cu_seqlens_k,
-        max_seqlen_q,
-        max_seqlen_k,
-        q,
-        k,
-        v,
-        output_pad_fn,
-    )
-
-
-def llama_model_forward(
-    self,
-    input_ids: torch.LongTensor = None,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[List[torch.FloatTensor]] = None,
-    inputs_embeds: Optional[torch.FloatTensor] = None,
-    use_cache: Optional[bool] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    return_dict: Optional[bool] = None,
-    cache_position: Optional[  # pylint: disable=unused-argument
-        torch.LongTensor
-    ] = None,
-) -> Union[Tuple, BaseModelOutputWithPast]:
-    output_attentions = (
-        output_attentions
-        if output_attentions is not None
-        else self.config.output_attentions
-    )
-    output_hidden_states = (
-        output_hidden_states
-        if output_hidden_states is not None
-        else self.config.output_hidden_states
+def patch_fa_llama_cross_entropy():
+    LOG.info(
+        "patching transformers.loss.loss_utils.fixed_cross_entropy with flash_attn.ops.triton.cross_entropy"
     )
-    use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-    return_dict = (
-        return_dict if return_dict is not None else self.config.use_return_dict
+    from flash_attn.ops.triton.cross_entropy import (
+        cross_entropy_loss as flash_attn_cross_entropy_loss,
     )
 
-    # retrieve input_ids and inputs_embeds
-    if input_ids is not None and inputs_embeds is not None:
-        raise ValueError(
-            "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
-        )
-    if input_ids is not None:
-        batch_size, seq_length = input_ids.shape
-    elif inputs_embeds is not None:
-        batch_size, seq_length, _ = inputs_embeds.shape
-    else:
-        raise ValueError(
-            "You have to specify either decoder_input_ids or decoder_inputs_embeds"
-        )
-
-    seq_length_with_past = seq_length
-    past_key_values_length = 0
-
-    if past_key_values is not None:
-        past_key_values_length = past_key_values[0][0].shape[2]
-        seq_length_with_past = seq_length_with_past + past_key_values_length
-
-    cu_seqlens = None
-    max_seqlen = None
-    if position_ids is None:
-        device = input_ids.device if input_ids is not None else inputs_embeds.device
-        position_ids = torch.arange(
-            past_key_values_length,
-            seq_length + past_key_values_length,
-            dtype=torch.long,
-            device=device,
-        )
-        position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
-    else:
-        position_ids = position_ids.view(-1, seq_length).long()
-        cu_seqlens, max_seqlen = get_cu_seqlens_from_pos_ids(position_ids)
-        cu_seqlens = cu_seqlens.squeeze()
-
-    if inputs_embeds is None:
-        inputs_embeds = self.embed_tokens(input_ids)
-    # embed positions
-    if attention_mask is None:
-        attention_mask = torch.ones(
-            (batch_size, seq_length_with_past),
-            dtype=torch.bool,
-            device=inputs_embeds.device,
+    def fa2_fixed_cross_entropy(
+        source,
+        target,
+        num_items_in_batch: int = None,
+        ignore_index: int = -100,
+        **kwargs,
+    ):
+        reduction = "sum" if num_items_in_batch is not None else "mean"
+        loss, _ = flash_attn_cross_entropy_loss(
+            source, target, ignore_index=ignore_index
         )
-        padding_mask = None
-    else:
-        if 0 in attention_mask:
-            padding_mask = attention_mask
+        if reduction == "sum":
+            loss = loss.sum() / num_items_in_batch
         else:
-            padding_mask = None
-
-    attention_mask = (
-        self._prepare_decoder_attention_mask(  # pylint: disable=protected-access
-            attention_mask,
-            (batch_size, seq_length),
-            inputs_embeds,
-            past_key_values_length,
-        )
-    )
-
-    hidden_states = inputs_embeds
-
-    if self.gradient_checkpointing and self.training:
-        if use_cache:
-            transformers.logger.warning_once(
-                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-            )
-            use_cache = False
-
-    # decoder layers
-    all_hidden_states = () if output_hidden_states else None
-    all_self_attns = () if output_attentions else None
-    next_decoder_cache = () if use_cache else None
-
-    for idx, decoder_layer in enumerate(self.layers):
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        past_key_value = past_key_values[idx] if past_key_values is not None else None
-
-        if self.gradient_checkpointing and self.training:
-
-            def create_custom_forward(module):
-                def custom_forward(*inputs):
-                    # None for past_key_value
-                    return module(
-                        *inputs,
-                    )
-
-                return custom_forward
-
-            layer_outputs = torch.utils.checkpoint.checkpoint(
-                create_custom_forward(decoder_layer),
-                hidden_states,
-                attention_mask,
-                position_ids,
-                past_key_value,
-                output_attentions,
-                None,
-                padding_mask,
-                cu_seqlens,
-                max_seqlen,
-            )
-        else:
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                padding_mask=padding_mask,
-                cu_seqlens=cu_seqlens,
-                max_seqlen=max_seqlen,
-            )
-
-        hidden_states = layer_outputs[0]
-
-        if use_cache:
-            next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
-
-        if output_attentions:
-            all_self_attns += (layer_outputs[1],)
-
-    hidden_states = self.norm(hidden_states)
-
-    # add hidden states from the last decoder layer
-    if output_hidden_states:
-        all_hidden_states += (hidden_states,)
-
-    next_cache = next_decoder_cache if use_cache else None
-    if not return_dict:
-        return tuple(
-            v
-            for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
-            if v is not None
-        )
-    return BaseModelOutputWithPast(
-        last_hidden_state=hidden_states,
-        past_key_values=next_cache,
-        hidden_states=all_hidden_states,
-        attentions=all_self_attns,
-    )
-
-
-class LlamaDecoderLayer(OriginalLlamaDecoderLayer):
-    """
-    patched version of LlamaDecoderLayer to pass through the precalculated cu_seqlens
-    """
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Tuple[torch.Tensor]] = None,
-        output_attentions: Optional[bool] = False,
-        use_cache: Optional[bool] = False,
-        padding_mask: Optional[torch.LongTensor] = None,
-        cu_seqlens: Optional[torch.Tensor] = None,
-        max_seqlen: Optional[torch.Tensor] = None,
-    ) -> Tuple[
-        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
-    ]:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            use_cache (`bool`, *optional*):
-                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-                (see `past_key_values`).
-            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-            cu_seqlens (`torch.Tensor`, *optional*) cumulative sequence len when packing
-        """
-
-        residual = hidden_states
-
-        hidden_states = self.input_layernorm(hidden_states)
-
-        # Self Attention
-        hidden_states, self_attn_weights, present_key_value = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_value=past_key_value,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
-            padding_mask=padding_mask,
-            cu_seqlens=cu_seqlens,
-            max_seqlen=max_seqlen,
-        )
-        hidden_states = residual + hidden_states
-
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-
-        outputs = (hidden_states,)
-
-        if output_attentions:
-            outputs += (self_attn_weights,)
-
-        if use_cache:
-            outputs += (present_key_value,)
+            loss = loss.sum() / (target != ignore_index).sum()
+        return loss
 
-        return outputs
+    transformers.loss.loss_utils.fixed_cross_entropy = fa2_fixed_cross_entropy
diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py b/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py
index 0c1a4e8224..332242e2c4 100644
--- a/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py
@@ -2,7 +2,6 @@
 Directly copied the code from https://raw.githubusercontent.com/oobabooga/text-generation-webui/main/modules/llama_attn_hijack.py and made some adjustments
 """
 
-import logging
 import warnings
 from typing import Optional, Tuple
 
@@ -11,10 +10,14 @@
 import transformers.models.llama.modeling_llama
 from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv
 
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
 try:
     import xformers.ops
 except ImportError:
-    logging.error("xformers not found! Please install it before trying to use it.")
+    LOG.error("xformers not found! Please install it before trying to use it.")
 
 
 def hijack_llama_attention():
@@ -29,10 +32,9 @@ def xformers_forward(
     past_key_value: Optional[Tuple[torch.Tensor]] = None,
     output_attentions: bool = False,
     use_cache: bool = False,
-    padding_mask: Optional[torch.LongTensor] = None,  # pylint: disable=unused-argument
-    **kwargs,  # pylint: disable=unused-argument
+    padding_mask: Optional[torch.LongTensor] = None,
+    **kwargs,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-    # pylint: disable=duplicate-code
     bsz, q_len, _ = hidden_states.size()
 
     if not hasattr(self, "pretraining_tp"):
@@ -99,7 +101,8 @@ def xformers_forward(
 
     if output_attentions:
         warnings.warn(
-            "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
+            "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead.",
+            stacklevel=2,
         )
 
     #
diff --git a/src/axolotl/monkeypatch/llama_expand_mask.py b/src/axolotl/monkeypatch/llama_expand_mask.py
deleted file mode 100644
index 5738bb543c..0000000000
--- a/src/axolotl/monkeypatch/llama_expand_mask.py
+++ /dev/null
@@ -1,25 +0,0 @@
-"""
-expands the binary attention mask per 3.2.2 of https://arxiv.org/pdf/2107.02027.pdf
-"""
-from typing import Optional
-
-import torch
-
-from axolotl.monkeypatch.utils import mask_2d_to_4d
-
-
-def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
-    masked_zero_one_mask = mask_2d_to_4d(mask, dtype, tgt_len)
-    inverted_mask = 1.0 - masked_zero_one_mask
-
-    return inverted_mask.masked_fill(
-        inverted_mask.to(torch.bool), torch.finfo(dtype).min
-    )
-
-
-def hijack_expand_mask():
-    import transformers
-
-    transformers.models.llama.modeling_llama._expand_mask = (  # pylint: disable=protected-access
-        _expand_mask
-    )
diff --git a/src/axolotl/monkeypatch/llama_patch_multipack.py b/src/axolotl/monkeypatch/llama_patch_multipack.py
deleted file mode 100644
index 540c5577a0..0000000000
--- a/src/axolotl/monkeypatch/llama_patch_multipack.py
+++ /dev/null
@@ -1,26 +0,0 @@
-"""
-Patched LlamaAttention to use torch.nn.functional.scaled_dot_product_attention
-"""
-
-from axolotl.monkeypatch.utils import (
-    patched_prepare_4d_causal_attention_mask,
-    patched_prepare_4d_causal_attention_mask_for_sdpa,
-)
-
-
-def hijack_llama_prepare_4d_mask():
-    import transformers.modeling_attn_mask_utils
-    import transformers.models.llama.modeling_llama
-
-    transformers.models.llama.modeling_llama._prepare_4d_causal_attention_mask_for_sdpa = (  # pylint: disable=protected-access
-        patched_prepare_4d_causal_attention_mask_for_sdpa
-    )
-    transformers.modeling_attn_mask_utils._prepare_4d_causal_attention_mask_for_sdpa = (  # pylint: disable=protected-access
-        patched_prepare_4d_causal_attention_mask_for_sdpa
-    )
-    transformers.models.llama.modeling_llama._prepare_4d_causal_attention_mask = (  # pylint: disable=protected-access
-        patched_prepare_4d_causal_attention_mask
-    )
-    transformers.modeling_attn_mask_utils._prepare_4d_causal_attention_mask = (  # pylint: disable=protected-access
-        patched_prepare_4d_causal_attention_mask
-    )
diff --git a/src/axolotl/monkeypatch/lora_kernels.py b/src/axolotl/monkeypatch/lora_kernels.py
new file mode 100644
index 0000000000..6d047119a9
--- /dev/null
+++ b/src/axolotl/monkeypatch/lora_kernels.py
@@ -0,0 +1,704 @@
+"""Module for patching custom LoRA Triton kernels and `torch.autograd` functions."""
+
+import importlib
+import inspect
+import logging
+import types
+from typing import Generator, Tuple, Type
+
+import torch
+from peft import PeftModelForCausalLM
+from torch import nn
+from transformers import AutoConfig
+
+from axolotl.kernels.lora import (
+    apply_lora_embedding,
+    apply_lora_gdn_in_proj,
+    apply_lora_linear,
+    apply_lora_mlp_geglu,
+    apply_lora_mlp_swiglu,
+    apply_lora_o,
+    apply_lora_qk,
+    apply_lora_qkv,
+)
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+QKV_PATCHES = [
+    (
+        """
+    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+    key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+""".lstrip("\n"),
+        """
+    query_states, key_states, value_states = self.apply_qkv(hidden_states)
+    query_states = query_states.view(hidden_shape).transpose(1, 2)
+    key_states = key_states.view(hidden_shape).transpose(1, 2)
+    value_states = value_states.view(hidden_shape).transpose(1, 2)
+""".lstrip("\n"),
+    ),
+    (
+        """
+    query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+""".lstrip("\n"),
+        """
+    query_states, key_states, value_states = self.apply_qkv(hidden_states)
+    query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+    key_states = self.k_norm(key_states.view(hidden_shape)).transpose(1, 2)
+    value_states = value_states.view(hidden_shape).transpose(1, 2)
+""".lstrip("\n"),
+    ),
+    (
+        """
+    query_states, gate = torch.chunk(
+        self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1
+    )
+    gate = gate.reshape(*input_shape, -1)
+
+    query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+    key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+""".lstrip("\n"),
+        """
+    query_states, key_states, value_states = self.apply_qkv(hidden_states)
+    query_states, gate = torch.chunk(
+        query_states.view(*input_shape, -1, self.head_dim * 2), 2, dim=-1
+    )
+    gate = gate.reshape(*input_shape, -1)
+
+    query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+    key_states = self.k_norm(key_states.view(hidden_shape)).transpose(1, 2)
+    value_states = value_states.view(hidden_shape).transpose(1, 2)
+""".lstrip("\n"),
+    ),
+    # Gemma4 has no entry: its fused forward already calls apply_qkv/apply_o,
+    # and patch_self_attn_lora skips it (see the skip there).
+]
+
+ORIGINAL_O_CODE = """
+    attn_output = self.o_proj(attn_output)
+""".lstrip("\n")
+
+PATCHED_O_CODE = """
+    attn_output = self.apply_o(attn_output)
+""".lstrip("\n")
+
+SUPPORTED_ACTIVATIONS = ["silu", "gelu"]
+APPLY_FN_MAPPING = {
+    "silu": apply_lora_mlp_swiglu,
+    "gelu": apply_lora_mlp_geglu,
+}
+
+
+def original_apply_qkv(
+    self: nn.Module, hidden_states: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Original implementation of QKV projection without optimizations.
+
+    Args:
+        self: The attention module instance.
+        hidden_states: Input tensor of shape [batch_size, seq_len, hidden_dim].
+
+    Returns:
+        A tuple `(query_states, key_states, value_states)` containing the projected
+            states for query, key, and value.
+    """
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+
+    return query_states, key_states, value_states
+
+
+def original_apply_qkv_optional_v(
+    self: nn.Module, hidden_states: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """QKV projection for models where v_proj may be None (e.g. Gemma4 attention_k_eq_v).
+
+    When v_proj is None, key_states are reused as value_states.
+    """
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    if self.v_proj is not None:
+        value_states = self.v_proj(hidden_states)
+    else:
+        value_states = key_states
+
+    return query_states, key_states, value_states
+
+
+def original_apply_o(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
+    """
+    Original implementation of output projection without optimizations.
+
+    Args:
+        self: The attention module instance.
+        hidden_states: Input tensor of shape `[`batch_size, seq_len, hidden_dim]`.
+
+    Returns:
+        The output projection result.
+    """
+    attn_output = self.o_proj(hidden_states)
+
+    return attn_output
+
+
+def get_attention_cls_from_config(cfg: DictDefault) -> Type[nn.Module]:
+    """
+    Get the appropriate attention class by inspecting the model config.
+    Uses dynamic import to support any model architecture that follows
+    the standard transformers naming convention.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+
+    Returns:
+        The appropriate attention class for the model.
+
+    Raises:
+        ValueError: If `base_model` not specified or attention class cannot be imported
+        ImportError: If the model module or attention class doesn't exist
+    """
+    if "base_model" not in cfg:
+        raise ValueError("base_model must be specified in config")
+
+    # Get model config without loading the model
+    model_config = AutoConfig.from_pretrained(cfg["base_model"])
+    model_type = model_config.model_type
+
+    # Special case for model_type = "qwen2"
+    if model_type == "qwen2":
+        from transformers.models.qwen2.modeling_qwen2 import Qwen2Attention
+
+        return Qwen2Attention
+
+    if model_type == "qwen3_vl":
+        from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextAttention
+
+        return Qwen3VLTextAttention
+
+    if model_type == "mllama":
+        from transformers.models.mllama.modeling_mllama import MllamaTextSelfAttention
+
+        return MllamaTextSelfAttention
+
+    if model_type == "llama4":
+        from transformers.models.llama4.modeling_llama4 import Llama4TextAttention
+
+        return Llama4TextAttention
+
+    if model_type == "mistral3":
+        from transformers.models.mistral.modeling_mistral import MistralAttention
+
+        return MistralAttention
+
+    if model_type == "gemma3_text":
+        from transformers.models.gemma3.modeling_gemma3 import Gemma3Attention
+
+        return Gemma3Attention
+
+    if model_type in ("gemma4", "gemma4_text"):
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextAttention
+
+        return Gemma4TextAttention
+
+    if model_type in ("gemma4_unified", "gemma4_unified_text"):
+        from transformers.models.gemma4_unified.modeling_gemma4_unified import (
+            Gemma4UnifiedTextAttention,
+        )
+
+        return Gemma4UnifiedTextAttention
+
+    try:
+        # Dynamically import the module and attention class
+        module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+        model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+        module = __import__(module_path, fromlist=[f"{model_cls_prefix}Attention"])
+        attention_cls = getattr(module, f"{model_cls_prefix}Attention")
+
+        return attention_cls
+    except (ImportError, AttributeError) as e:
+        raise ValueError(
+            f"Axolotl could not import attention class for model_type: {model_type}. "
+            "Please raise an Issue and turn off lora kernels to continue training. "
+            f"Error: {str(e)}"
+        ) from e
+
+
+def patch_self_attn_lora(cfg: DictDefault):
+    """
+    Given an `axolotl` config, this method patches the inferred attention class forward
+    pass with optimized LoRA implementations.
+
+    It modifies the attention class to use optimized QKV and output projections. The
+    original implementation is preserved and can be restored if needed.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+
+    Raises:
+        AssertionError: If the required code blocks are not found in the attention
+            implementation.
+    """
+    attention_cls = get_attention_cls_from_config(cfg)
+
+    # Check if already patched
+    if hasattr(attention_cls, "_original_forward"):
+        LOG.info(f"{attention_cls.__name__} already patched")
+        return
+
+    # Skip Gemma4: patch_manager applies patch_gemma4_fused_attn
+    # unconditionally for gemma4 before this runs, and that fused forward
+    # already calls apply_qkv/apply_o, so the source rewrite is dead.
+    try:
+        from transformers.models.gemma4.modeling_gemma4 import (
+            Gemma4TextAttention,
+        )
+
+        if attention_cls is Gemma4TextAttention:
+            LOG.info(
+                "Gemma4TextAttention uses the fused attention path "
+                "(apply_qkv/apply_o) - skipping LoRA source rewrite"
+            )
+            return
+    except ImportError:
+        pass
+
+    # gemma4_unified's attention forward can't be source-rewritten (KV sharing); skip.
+    try:
+        from transformers.models.gemma4_unified.modeling_gemma4_unified import (
+            Gemma4UnifiedTextAttention,
+        )
+    except ImportError:
+        Gemma4UnifiedTextAttention = None
+
+    if (
+        Gemma4UnifiedTextAttention is not None
+        and attention_cls is Gemma4UnifiedTextAttention
+    ):
+        if cfg.fused_attn_kernel:
+            LOG.info(
+                "Gemma4UnifiedTextAttention uses the fused attention path "
+                "(apply_qkv/apply_o) - skipping LoRA source rewrite"
+            )
+        else:
+            LOG.warning(
+                "lora_qkv_kernel/lora_o_kernel cannot attach to gemma4_unified "
+                "without fused_attn_kernel: true - skipping QKV/O rewrite "
+                "(MLP/embedding kernels still apply). Set fused_attn_kernel: true "
+                "to enable them."
+            )
+        return
+
+    self_attn_forward = inspect.getsource(attention_cls.forward)
+    attention_cls._original_forward = self_attn_forward
+    self_attn_forward, _ = detab_code(self_attn_forward)
+
+    assert any(qkv_options[0] in self_attn_forward for qkv_options in QKV_PATCHES), (
+        "Original QKV code not found"
+    )
+    assert ORIGINAL_O_CODE in self_attn_forward, "Original O code not found"
+
+    for qkv_orig, qkv_patched in QKV_PATCHES:
+        if qkv_orig in self_attn_forward:
+            self_attn_forward = self_attn_forward.replace(
+                qkv_orig,
+                qkv_patched,
+            )
+            break
+    self_attn_forward = self_attn_forward.replace(ORIGINAL_O_CODE, PATCHED_O_CODE)
+    self_attn_forward = self_attn_forward.replace(
+        "def forward(",
+        "def axolotl_attn_forward(",
+        1,
+    )
+
+    # Load necessary imports
+    module_name = attention_cls.__module__
+    module = importlib.import_module(module_name)
+
+    items_to_import = []
+    for item in dir(module):
+        if item in self_attn_forward:
+            items_to_import.append(item)
+
+    exec(
+        f"from {module_name} import ({', '.join(items_to_import)})",
+        globals(),
+    )
+    exec(self_attn_forward, globals())
+
+    LOG.info(f"Patched attention class with LoRA optims: {attention_cls.__name__}")
+    attention_cls.forward = axolotl_attn_forward
+
+
+def find_self_attn_in_layer(
+    layer: nn.Module,
+) -> Generator[Tuple[nn.Module], None, None]:
+    # general case of most models
+    if hasattr(layer, "self_attn"):
+        if all(
+            hasattr(layer.self_attn, proj)
+            for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]
+        ):
+            yield layer.self_attn
+
+
+# GatedDeltaNet projections routed through the fused kernels to avoid peft's
+# bf16->fp32->bf16 dtype round-trip. The in-projections share a single input
+# (post-norm hidden states) and are fused into one autograd node; out_proj has a
+# different input and is routed through the single-projection kernel.
+LINEAR_ATTN_IN_PROJS = ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a")
+LINEAR_ATTN_OUT_PROJ = "out_proj"
+LINEAR_ATTN_PROJS = LINEAR_ATTN_IN_PROJS + (LINEAR_ATTN_OUT_PROJ,)
+
+
+def find_linear_attn_in_layer(layer: nn.Module) -> Generator[nn.Module, None, None]:
+    # Qwen3.5 / Qwen3.5-MoE hybrid layers (GatedDeltaNet). qwen3_next fuses its
+    # projections (in_proj_qkvz / in_proj_ba) under different names and must not
+    # match here. Require all four canonical in-projections: the fused node and the
+    # patched forward both index every one, so a partial match falls back to peft.
+    if hasattr(layer, "linear_attn"):
+        linear_attn = layer.linear_attn
+        if hasattr(linear_attn, LINEAR_ATTN_OUT_PROJ) and all(
+            hasattr(linear_attn, proj) for proj in LINEAR_ATTN_IN_PROJS
+        ):
+            yield linear_attn
+
+
+def _make_apply_lora_linear(proj_name: str):
+    def apply_linear(self, X: torch.Tensor) -> torch.Tensor:
+        return apply_lora_linear(getattr(self, proj_name), X)
+
+    apply_linear.__name__ = f"apply_{proj_name}"
+    return apply_linear
+
+
+def _make_apply_lora_gdn_in_proj(proj_names: tuple[str, ...]):
+    def apply_in_proj(self, X: torch.Tensor) -> dict:
+        return apply_lora_gdn_in_proj(self, X, proj_names)
+
+    return apply_in_proj
+
+
+def find_mlp_in_layer(
+    layer: nn.Module,
+    skip_routed_experts: bool = False,
+) -> Generator[Tuple[nn.Module, nn.Module, nn.Module, nn.Module], None, None]:
+    # general case of most models: the dense (shared) MLP. Always eligible — a custom MoE expert
+    # kernel only owns the ROUTED experts, not this dense gate/up/down block (e.g. gemma4's per-layer
+    # Gemma4TextMLP). Sparse-block MoEs whose layer.mlp is the router (no gate_proj) don't match here.
+    if hasattr(layer, "mlp"):
+        if all(
+            hasattr(layer.mlp, proj) for proj in ["gate_proj", "up_proj", "down_proj"]
+        ):
+            yield layer.mlp.gate_proj, layer.mlp.up_proj, layer.mlp.down_proj, layer.mlp
+    # llama4 shared expert: also a dense MLP, not routed -> always eligible.
+    if hasattr(layer, "feedforward") and hasattr(layer.feedforward, "shared_expert"):
+        mlp = layer.feedforward.shared_expert
+        yield mlp.gate_proj, mlp.up_proj, mlp.down_proj, mlp
+    # llama4 linearized ROUTED experts: skip when a custom MoE expert kernel (ScatterMoE/SonicMoE)
+    # owns the experts — patching them here would double-own / conflict with the MoE kernel path.
+    if (
+        not skip_routed_experts
+        and hasattr(layer, "feedforward")
+        and hasattr(layer.feedforward, "experts")
+    ):
+        if all(
+            hasattr(layer.feedforward.experts, proj)
+            for proj in ["gate_projs", "up_projs", "down_projs"]
+        ):
+            for gate_proj, up_proj, down_proj in zip(
+                layer.feedforward.experts.gate_projs,
+                layer.feedforward.experts.up_projs,
+                layer.feedforward.experts.down_projs,
+                strict=False,
+            ):
+                yield (
+                    gate_proj,
+                    up_proj,
+                    down_proj,
+                    FakeMLP(gate_proj, up_proj, down_proj),
+                )
+
+
+def get_layers(model: PeftModelForCausalLM) -> list[nn.Module]:
+    """
+    Get the layers of the model. Handles text-only and multimodal models.
+
+    Args:
+        model: A PEFT model.
+
+    Returns:
+        A list of layers.
+    """
+    pretrained_model = model.model
+
+    # check for multimodal models first
+    if hasattr(pretrained_model, "language_model"):
+        return pretrained_model.language_model.layers
+    if hasattr(pretrained_model, "model"):
+        if hasattr(pretrained_model.model, "language_model"):
+            return pretrained_model.model.language_model.layers
+        return pretrained_model.model.layers
+
+    raise NotImplementedError(
+        f"Model type {model.config.model_type} is not supported yet. Please create an Issue."
+    )
+
+
+def apply_lora_kernel_patches(
+    model: PeftModelForCausalLM, cfg: DictDefault
+) -> PeftModelForCausalLM:
+    """
+    Applies optimized Triton kernel patches to a PEFT model.
+
+    Patches a PEFT model with optimized implementations for MLP and attention
+    computations. The optimizations include custom Triton kernels for activation
+    functions and specialized autograd functions for LoRA computations.
+
+    Args:
+        model: A PEFT model to be patched with optimized kernels.
+        cfg: Dictionary mapping `axolotl` config keys to values.
+
+    Returns:
+        PeftModelForCausalLM: The patched model with optimized kernels.
+
+    Raises:
+        TypeError: If the provided model is not a `PeftModelForCausalLM`.
+        NotImplementedError: If the model type is not supported.
+        AssertionError: If multiple adapters are active (currently unsupported).
+
+    Note:
+        The optimizations require LoRA adapters with no dropout and no bias terms. The
+            function will skip patching if these conditions aren't met.
+    """
+    if not isinstance(model, PeftModelForCausalLM):
+        raise TypeError("Model must be a PeftModelForCausalLM")
+
+    # Get active LoRA adapter config
+    if hasattr(model, "active_adapters"):
+        assert len(model.active_adapters) == 1, (
+            "Axolotl currently does not support LoRA Triton kernels for multiple adapters"
+        )
+        active_adapter = model.active_adapters[0]
+    else:
+        active_adapter = model.active_adapter
+    lora_config = model.model.peft_config[active_adapter]
+
+    # Log what features are active
+    if lora_config.lora_dropout > 0:
+        LOG.info(f"LoRA kernels: dropout={lora_config.lora_dropout} enabled")
+    if lora_config.bias != "none":
+        LOG.info(f"LoRA kernels: bias={lora_config.bias} enabled")
+    if lora_config.use_dora:
+        LOG.info("LoRA kernels: DoRA enabled")
+
+    # This needs to be reset after patching
+    original_level = LOG.getEffectiveLevel()
+    LOG.setLevel(logging.INFO)
+
+    # Choose activation based on model type
+    activation = None
+    text_config = (
+        model.config.get_text_config()
+        if hasattr(model.config, "get_text_config")
+        else model.config
+    )
+    if hasattr(text_config, "hidden_act"):
+        activation = text_config.hidden_act
+    elif hasattr(text_config, "hidden_activation"):
+        activation = text_config.hidden_activation
+    elif hasattr(text_config, "mlp_hidden_act"):
+        # Hybrid models (e.g. nemotron_h) use mlp_hidden_act instead of hidden_act
+        activation = text_config.mlp_hidden_act
+
+    # map activation to supported activation
+    if activation and "gelu" in activation:
+        # gemma3 uses gelu_pytorch_tanh
+        activation = "gelu"
+
+    layers = get_layers(model)
+
+    linear_attn_patched_layers = 0
+    linear_attn_patched_projs: set[str] = set()
+
+    # Patch each layer
+    for layer in layers:
+        # Add QKV, O fallback implementations to start
+        # These will be overwritten later (if some conditions apply)
+        for self_attn in find_self_attn_in_layer(layer):
+            # Use v_proj-optional fallback for models where v_proj can be None
+            # (e.g. Gemma4 with attention_k_eq_v=True)
+            if getattr(self_attn, "v_proj", None) is None:
+                self_attn.apply_qkv = types.MethodType(
+                    original_apply_qkv_optional_v, self_attn
+                )
+            else:
+                self_attn.apply_qkv = types.MethodType(original_apply_qkv, self_attn)
+            self_attn.apply_o = types.MethodType(original_apply_o, self_attn)
+
+            if cfg.lora_qkv_kernel:
+                # Query, key, value patching
+                # Filter out None projections (e.g. Gemma4 v_proj when attention_k_eq_v=True)
+                has_v_proj = getattr(self_attn, "v_proj", None) is not None
+                proj_names = (
+                    ["q_proj", "k_proj", "v_proj"]
+                    if has_v_proj
+                    else ["q_proj", "k_proj"]
+                )
+                layer_modules = [getattr(self_attn, name) for name in proj_names]
+                can_patch_qkv = all(
+                    hasattr(module, "lora_A") for module in layer_modules
+                )
+
+                if can_patch_qkv:
+                    if has_v_proj:
+                        self_attn.apply_qkv = types.MethodType(
+                            apply_lora_qkv, self_attn
+                        )
+                    else:
+                        self_attn.apply_qkv = types.MethodType(apply_lora_qk, self_attn)
+                else:
+                    LOG.warning_once(
+                        "Cannot patch some attention QKV projections - requires LoRA adapters"
+                    )
+            if cfg.lora_o_kernel:
+                # Output patching
+                layer_modules = [
+                    getattr(self_attn, linear_proj) for linear_proj in ["o_proj"]
+                ]
+                can_patch_o = all(hasattr(module, "lora_A") for module in layer_modules)
+
+                if can_patch_o:
+                    self_attn.apply_o = types.MethodType(apply_lora_o, self_attn)
+                else:
+                    LOG.warning_once(
+                        "Cannot patch some attention output projection - requires LoRA adapters"
+                    )
+        if cfg.lora_qkv_kernel or cfg.lora_o_kernel:
+            for linear_attn in find_linear_attn_in_layer(layer):
+                patched = False
+                # in_proj group mirrors self-attn qkv (lora_qkv_kernel); out_proj
+                # mirrors o_proj (lora_o_kernel). Base-only in-projections fold into
+                # the fused node too, so patch only when one carries a LoRA adapter.
+                if cfg.lora_qkv_kernel and any(
+                    hasattr(getattr(linear_attn, name), "lora_A")
+                    for name in LINEAR_ATTN_IN_PROJS
+                ):
+                    linear_attn.apply_in_proj_fused = types.MethodType(
+                        _make_apply_lora_gdn_in_proj(LINEAR_ATTN_IN_PROJS), linear_attn
+                    )
+                    linear_attn_patched_projs.update(LINEAR_ATTN_IN_PROJS)
+                    patched = True
+                out_proj = getattr(linear_attn, LINEAR_ATTN_OUT_PROJ, None)
+                if (
+                    cfg.lora_o_kernel
+                    and out_proj is not None
+                    and hasattr(out_proj, "lora_A")
+                ):
+                    linear_attn.apply_out_proj = types.MethodType(
+                        _make_apply_lora_linear(LINEAR_ATTN_OUT_PROJ), linear_attn
+                    )
+                    linear_attn_patched_projs.add(LINEAR_ATTN_OUT_PROJ)
+                    patched = True
+                if patched:
+                    linear_attn_patched_layers += 1
+        # When ScatterMoE/SonicMoE owns the routed experts, lora_mlp_kernel must only fuse the
+        # DENSE shared MLP, never the routed-expert containers (which the MoE kernel handles).
+        _moe_kernels_own_experts = bool(cfg.use_scattermoe) or bool(cfg.use_sonicmoe)
+        for gate_proj, up_proj, down_proj, mlp in find_mlp_in_layer(
+            layer, skip_routed_experts=_moe_kernels_own_experts
+        ):
+            if cfg.lora_mlp_kernel:
+                # Check is inside lora_mlp_kernel guard so models with an
+                # unsupported activation (e.g. nemotron_h uses relu2) can set
+                # lora_mlp_kernel: false without hitting an error here.
+                if activation not in SUPPORTED_ACTIVATIONS:
+                    raise NotImplementedError(
+                        f"Activation {activation!r} is not supported by lora_mlp_kernel. "
+                        f"Set `lora_mlp_kernel: false` in your config or use a model with "
+                        f"a supported activation ({SUPPORTED_ACTIVATIONS})."
+                    )
+                # MLP patching
+                can_patch_mlp = all(
+                    hasattr(proj, "lora_A") for proj in (gate_proj, up_proj, down_proj)
+                )
+
+                if can_patch_mlp:
+                    apply_fn = APPLY_FN_MAPPING[activation]
+                    layer.mlp.forward = types.MethodType(apply_fn, mlp)
+                else:
+                    LOG.warning_once(
+                        "Cannot patch some MLP layers - requires LoRA adapters"
+                    )
+
+    if linear_attn_patched_layers:
+        LOG.info(
+            f"Patched linear-attention LoRA projections {sorted(linear_attn_patched_projs)} "
+            f"with fused kernels on {linear_attn_patched_layers} layer(s)"
+        )
+
+    # Patch embedding layers (model-level, not per-layer)
+    if cfg.lora_embedding_kernel:
+        _patch_embedding_layers(model, cfg)
+
+    LOG.setLevel(original_level)
+
+    return model
+
+
+def _patch_embedding_layers(model: PeftModelForCausalLM, cfg: DictDefault):
+    """Patch embedding layers with fused LoRA kernel.
+
+    Handles both embed_tokens (nn.Embedding with lora_embedding_A/B) and
+    lm_head (nn.Linear with lora_A/B, used when tied embeddings are untied by PEFT).
+    """
+    pretrained_model = model.model
+    patched = 0
+
+    # Find embedding modules - check common locations
+    for attr_path in [
+        ("model", "embed_tokens"),
+        ("model", "language_model", "embed_tokens"),
+    ]:
+        parent = pretrained_model
+        for attr in attr_path:
+            parent = getattr(parent, attr, None)
+            if parent is None:
+                break
+        if parent is not None and hasattr(parent, "lora_embedding_A"):
+            LOG.info(f"Patching embedding layer: {'.'.join(attr_path)}")
+            parent.forward = types.MethodType(apply_lora_embedding, parent)
+            patched += 1
+
+    # lm_head with LoRA is a Linear layer - already handled by LoRA_O/LoRA_W kernels
+    # when included in target_modules. No special embedding handling needed since
+    # PEFT wraps it as a Linear (not Embedding) even for tied models.
+
+    if not patched:
+        LOG.debug("No embedding layers with LoRA found to patch")
+
+
+class FakeMLP(nn.Module):
+    """
+    placeholder MLP for triton patching
+    """
+
+    gate_proj: nn.Linear
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+
+    def __init__(self, gate_proj, up_proj, down_proj):
+        super().__init__()
+        self.gate_proj = gate_proj
+        self.up_proj = up_proj
+        self.down_proj = down_proj
diff --git a/src/axolotl/monkeypatch/loss/__init__.py b/src/axolotl/monkeypatch/loss/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/loss/chunked.py b/src/axolotl/monkeypatch/loss/chunked.py
new file mode 100644
index 0000000000..26a52f8981
--- /dev/null
+++ b/src/axolotl/monkeypatch/loss/chunked.py
@@ -0,0 +1,134 @@
+"""
+chunked ce loss
+"""
+
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+
+
+# copied and modified from torchtune.modules.loss.CEWithChunkedOutputLoss
+class CEWithChunkedOutputLoss(torch.nn.Module):
+    """
+    Cross-entropy with chunked outputs that saves memory by only upcasting one chunk at a time.
+
+    For more details, please refer to: https://github.com/pytorch/torchtune/pull/1390
+    """
+
+    def __init__(self, num_output_chunks: int = 8, ignore_index: int = -100):
+        super().__init__()
+        self.num_output_chunks = num_output_chunks
+        self.ignore_index = ignore_index
+
+    def compute_cross_entropy(
+        self,
+        logits: torch.Tensor,
+        labels: torch.Tensor,
+        normalize: bool = True,
+    ) -> torch.Tensor:
+        """
+        Upcast logits to fp32 and compute cross entropy loss.
+        """
+        return F.cross_entropy(
+            logits.float(), labels, ignore_index=self.ignore_index, reduction="sum"
+        )
+
+    def forward(
+        self, logits: List[torch.Tensor], labels: torch.Tensor, reduction="sum"
+    ) -> torch.Tensor:
+        """
+        Args:
+            logits (List[torch.Tensor]): List of chunked logits of length
+                ``self.num_output_chunks``, where each chunk has shape
+                ``(batch_size, num_tokens / num_output_chunks, vocab_size)``.
+            labels (torch.Tensor): Ground truth labels of shape ``(batch_size, num_tokens)``.
+            reduction (str): The reduction to apply to the output.
+
+        Returns:
+            torch.Tensor: Cross entropy loss of shape (1,).
+        """
+
+        total_elements = (labels != self.ignore_index).sum()
+
+        # chunk and reshape labels (bsz, num_tokens, vocab) -> [(bsz*num_tokens/num_chunks, vocab)]
+        labels = [
+            target_chunk.reshape(-1)
+            for target_chunk in labels.chunk(self.num_output_chunks, dim=1)
+        ]
+        # reshape logits [(bsz, num_tokens/num_chunks, vocab)] -> [(bsz*num_tokens/num_chunks, vocab)]
+        logits = [
+            logit_chunk.reshape(-1, logit_chunk.size(-1)) for logit_chunk in logits
+        ]
+
+        # compute one chunk at a time
+        total_loss = 0.0
+        for logits_chunk, labels_chunk in zip(logits, labels, strict=False):
+            total_loss += self.compute_cross_entropy(logits_chunk, labels_chunk)
+
+        if reduction == "sum":
+            return total_loss
+        return total_loss / total_elements
+
+
+def _build_chunked_ce_loss_fn(num_output_chunks: int = 8, ignore_index: int = -100):
+    loss_fn_ce = CEWithChunkedOutputLoss(num_output_chunks, ignore_index)
+    loss_fn_ce.compute_cross_entropy = torch.compile(
+        loss_fn_ce.compute_cross_entropy, backend="inductor"
+    )
+    return loss_fn_ce
+
+
+def get_causal_lm_loss(num_output_chunks: int = 8, ignore_index: int = -100):
+    loss_fn_ce = _build_chunked_ce_loss_fn(num_output_chunks, ignore_index)
+
+    def chunked_fix_cross_entropy(
+        source,
+        target,
+        num_items_in_batch: int = None,
+        ignore_index: int = -100,
+        **kwargs,
+    ):
+        reduction = "sum" if num_items_in_batch is not None else "mean"
+        logit_chunks = [
+            chunk for chunk in source.chunk(loss_fn_ce.num_output_chunks, dim=1)
+        ]
+        loss = loss_fn_ce(logit_chunks, target, reduction=reduction)
+        if reduction == "sum":
+            loss = loss / num_items_in_batch
+        return loss
+
+    def for_causal_lm_chunked_loss(
+        logits,
+        labels,
+        vocab_size: int = None,
+        num_items_in_batch: Optional[int] = None,
+        ignore_index: int = -100,
+        shift_labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        # skip the upcast to float since we handle that in the chunking loss
+        if shift_labels is None:
+            # Shift so that tokens < n predict n
+            labels = F.pad(labels, (0, 1), value=ignore_index)
+            shift_labels = labels[..., 1:].contiguous()
+
+        # Skip Flattening the tokens
+        # Enable model parallelism
+        shift_labels = shift_labels.to(logits.device)
+        loss = chunked_fix_cross_entropy(
+            logits, shift_labels, num_items_in_batch, ignore_index, **kwargs
+        )
+        return loss
+
+    return for_causal_lm_chunked_loss
+
+
+def patch_chunked_ce_loss_fn(num_output_chunks: int = 8, ignore_index: int = -100):
+    import transformers.loss.loss_utils
+
+    for_causal_lm_chunked_loss = get_causal_lm_loss(num_output_chunks, ignore_index)
+    transformers.loss.loss_utils.ForCausalLMLoss = for_causal_lm_chunked_loss
+    transformers.loss.loss_utils.LOSS_MAPPING["ForCausalLM"] = (
+        for_causal_lm_chunked_loss
+    )
diff --git a/src/axolotl/monkeypatch/loss/eaft.py b/src/axolotl/monkeypatch/loss/eaft.py
new file mode 100644
index 0000000000..150d4a005b
--- /dev/null
+++ b/src/axolotl/monkeypatch/loss/eaft.py
@@ -0,0 +1,51 @@
+"""
+eaft (entropy-aware focal training) loss implementation
+weights examples by entropy approximation from top-k logits
+
+Reference: https://github.com/ymxyll/LlamaFactory-EAFT/blob/e2ce19e8efcc226450ee8f2b81dfe4e69f1f945d/src/llamafactory/train/trainer_utils.py
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+def eaft_loss(outputs, labels, num_items_in_batch=None, alpha=1.0, k=20):
+    """
+    compute eaft loss with entropy weighting
+
+    args:
+        outputs: model outputs containing logits
+        labels: target labels for computing loss
+        num_items_in_batch: for sample packing support
+        alpha: exponent for entropy weighting (default 1.0)
+        k: number of top logits for entropy approximation (default 20)
+    """
+    logits = outputs.logits
+
+    shift_logits = logits[..., :-1, :].contiguous()
+    shift_labels = labels[..., 1:].contiguous()
+
+    vocab_size = shift_logits.size(-1)
+    shift_logits_view = shift_logits.view(-1, vocab_size)
+    shift_labels_view = shift_labels.view(-1)
+
+    mask = shift_labels_view != -100
+
+    with torch.no_grad():
+        top_k_logits, _ = torch.topk(
+            shift_logits_view[mask].float(), k=min(k, vocab_size), dim=-1
+        )
+        top_k_probs = F.softmax(top_k_logits, dim=-1)
+        entropy = -(top_k_probs * torch.log(top_k_probs + 1e-10)).sum(dim=-1)
+        weights = torch.pow(entropy, alpha)
+
+    loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+    per_token_loss = loss_fct(shift_logits_view[mask], shift_labels_view[mask])
+    weighted_loss = per_token_loss * weights
+
+    if num_items_in_batch is not None:
+        loss = weighted_loss.sum() / num_items_in_batch
+    else:
+        loss = weighted_loss.mean()
+
+    return loss
diff --git a/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py b/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py
index 6ae2e75fa2..0994da91ca 100644
--- a/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py
@@ -1,637 +1,18 @@
 """Flash attention monkey patch for mistral model"""
-# pylint: disable=duplicate-code
 
-import logging
-from typing import List, Optional, Tuple, Union
+from functools import partial
 
-import torch
 import transformers
-from einops import rearrange
-from flash_attn.bert_padding import pad_input, unpad_input
-from flash_attn.flash_attn_interface import (  # pylint: disable=ungrouped-imports
-    flash_attn_kvpacked_func,
-    flash_attn_varlen_kvpacked_func,
-    flash_attn_varlen_qkvpacked_func,
-)
-from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.models.mistral.modeling_mistral import (
-    MistralAttention as OriginalMistralAttention,
-)
-from transformers.models.mistral.modeling_mistral import (
-    MistralDecoderLayer as OriginalMistralDecoderLayer,
-)
-from transformers.models.mistral.modeling_mistral import apply_rotary_pos_emb, repeat_kv
 
-from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl.monkeypatch.mistral")
+LOG = get_logger(__name__)
 
 
-def replace_mistral_attn_with_flash_attn(
-    packed: Optional[bool] = False,
-):
-    transformers.models.mistral.modeling_mistral.MistralModel._prepare_decoder_attention_mask = (  # pylint: disable=protected-access
-        _prepare_decoder_attention_mask
-    )
-    transformers.models.mistral.modeling_mistral.MistralAttention.forward = (
-        flashattn_forward
-    )
-    if packed:
-        transformers.models.mistral.modeling_mistral.MistralDecoderLayer = (
-            MistralDecoderLayer
-        )
-        transformers.models.mistral.modeling_mistral.MistralModel.forward = (
-            mistral_model_forward
-        )
-
+def patch_mistral_cross_entropy():
+    from flash_attn.losses.cross_entropy import CrossEntropyLoss
 
-@torch.jit.script
-def _make_sliding_window_causal_mask(
-    bsz: int,
-    tgt_len: int,
-    dtype: torch.dtype,
-    device: torch.device,
-    past_key_values_length: int = 0,
-    sliding_window: int = 4096,
-):
-    """
-    Make causal mask used for sliding window attention
-    """
-    tensor = torch.full(
-        (tgt_len, tgt_len),
-        fill_value=1,
-        device=device,
+    LOG.info("patching with flash_attn.losses.cross_entropy")
+    transformers.models.mistral.modeling_mistral.CrossEntropyLoss = partial(
+        CrossEntropyLoss, inplace_backward=True
     )
-    mask = torch.tril(tensor, diagonal=0)
-    # make the mask banded to account for sliding window
-    # NOTE: HF implementation is wrong as of 14-10-2023 for torch.triu, needs +1
-    mask = torch.triu(mask, diagonal=-sliding_window + 1)
-    mask = torch.log(mask).to(dtype)
-
-    if past_key_values_length > 0:
-        mask = torch.cat(
-            [
-                torch.zeros(
-                    tgt_len, past_key_values_length, dtype=dtype, device=device
-                ),
-                mask,
-            ],
-            dim=-1,
-        )
-    return mask[None, None, :, :].expand(
-        bsz, 1, tgt_len, tgt_len + past_key_values_length
-    )
-
-
-# Disable the transformation of the attention mask in LlamaModel as the flash attention
-# requires the attention mask to be the same as the key_padding_mask
-def _prepare_decoder_attention_mask(
-    self,
-    attention_mask,
-    input_shape,
-    inputs_embeds,
-    past_key_values_length,
-    sliding_window,
-):  # pylint: disable=unused-argument
-    # [bsz, seq_len]
-    if attention_mask is None or sliding_window is None:
-        return attention_mask
-
-    # NOTE: attention mask and sliding masks are only broadcastable in certain scenarios.
-    # Without attention_mask.shape[0] == 1, error will trigger after eval loss but only when wandb is enabled.
-    if input_shape[-1] > 1 and attention_mask.shape[0] == 1:
-        sliding_window_mask = _make_sliding_window_causal_mask(
-            bsz=input_shape[0],
-            tgt_len=input_shape[1],
-            dtype=inputs_embeds.dtype,
-            device=inputs_embeds.device,
-            past_key_values_length=past_key_values_length,
-            sliding_window=sliding_window,
-        )
-        attention_mask = attention_mask + sliding_window_mask
-    else:
-        LOG.info("skipping sliding window mask, not broadcastable with attention mask")
-
-    return attention_mask
-
-
-def flashattn_forward(
-    self: OriginalMistralAttention,
-    hidden_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_value: Optional[Tuple[torch.Tensor]] = None,
-    output_attentions: bool = False,
-    use_cache: bool = False,
-    cu_seqlens: Optional[torch.Tensor] = None,
-    max_seqlen: Optional[torch.Tensor] = None,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-    bsz, q_len, _ = hidden_states.size()
-
-    query_states = self.q_proj(hidden_states)
-    key_states = self.k_proj(hidden_states)
-    value_states = self.v_proj(hidden_states)
-
-    query_states = query_states.view(
-        bsz, q_len, self.num_heads, self.head_dim
-    ).transpose(1, 2)
-    key_states = key_states.view(
-        bsz, q_len, self.num_key_value_heads, self.head_dim
-    ).transpose(1, 2)
-    value_states = value_states.view(
-        bsz, q_len, self.num_key_value_heads, self.head_dim
-    ).transpose(1, 2)
-
-    kv_seq_len = key_states.shape[-2]
-    if past_key_value is not None:
-        kv_seq_len += past_key_value[0].shape[-2]
-    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-    query_states, key_states = apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, position_ids
-    )
-
-    use_sliding_windows = (
-        getattr(self.config, "sliding_window") is not None
-        and kv_seq_len > self.config.sliding_window
-    )
-
-    if use_sliding_windows:
-        window_size = (self.config.sliding_window, self.config.sliding_window)
-    else:
-        window_size = (-1, -1)
-
-    if past_key_value is not None:
-        # Activate slicing cache only if the config has a value `sliding_windows` attribute
-        if (
-            hasattr(self.config, "sliding_window")
-            and kv_seq_len > self.config.sliding_window
-        ):
-            slicing_tokens = kv_seq_len - self.config.sliding_window
-
-            past_key = past_key_value[0]
-            past_value = past_key_value[1]
-
-            past_key = past_key[:, :, slicing_tokens:, :].contiguous()
-            past_value = past_value[:, :, slicing_tokens:, :].contiguous()
-
-            if past_key.shape[-2] != self.config.sliding_window - 1:
-                raise ValueError(
-                    f"past key much have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
-                    f" {past_key.shape}"
-                )
-
-            past_key_value = (past_key, past_value) if use_cache else None
-
-        if past_key_value is not None:
-            key_states = torch.cat([past_key_value[0], key_states], dim=2)
-            value_states = torch.cat([past_key_value[1], value_states], dim=2)
-
-    past_key_value = (key_states, value_states) if use_cache else None
-
-    # repeat k/v heads if n_kv_heads < n_heads
-    key_states = repeat_kv(key_states, self.num_key_value_groups)
-    value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-    if self.training:
-        # during training q,k,v always have same seqlen
-        assert key_states.shape == query_states.shape
-        is_causal = True
-    else:
-        # turn off FA causal mask after first inference autoregressive iteration
-        # only on first autoregressive step q,k,v have same seqlen
-        is_causal = key_states.shape == query_states.shape
-
-    dropout_rate = 0.0 if not self.training else getattr(self, "attention_dropout", 0.0)
-
-    if cu_seqlens is not None and max_seqlen is not None and cu_seqlens.dim() == 1:
-        # special handling using sample packing
-        qkv = torch.stack(
-            [query_states, key_states, value_states], dim=2
-        )  # [bsz, nh, 3, q_len, hd]
-        qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
-        qkv = rearrange(qkv, "b s ... -> (b s) ...")
-
-        output = flash_attn_varlen_qkvpacked_func(
-            qkv,
-            cu_seqlens,
-            max_seqlen,
-            dropout_p=dropout_rate,
-            softmax_scale=None,
-            causal=True,
-            window_size=window_size,
-        )
-        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
-    elif query_states.shape == key_states.shape:
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
-        qkv_unpad, cu_seqlens_q, max_seqlen_q, _, output_pad_fn = generate_qkv(
-            query_states,
-            key_states,
-            value_states,
-            qkvpacked=True,
-            # We have disabled _prepare_decoder_attention_mask in LlamaModel
-            # the attention_mask should be the same as the key_padding_mask
-            key_padding_mask=attention_mask,
-            query_padding_mask=attention_mask[:, -query_states.size(1) :]
-            if attention_mask is not None
-            else None,
-        )
-        output_unpad = flash_attn_varlen_qkvpacked_func(
-            qkv_unpad,
-            cu_seqlens_q,
-            max_seqlen_q,
-            dropout_p=dropout_rate,
-            softmax_scale=None,
-            causal=is_causal,
-            window_size=window_size,
-        )
-        output = output_pad_fn(output_unpad)
-    else:
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
-        if attention_mask is None or attention_mask.all().item():
-            output = flash_attn_kvpacked_func(
-                query_states,
-                torch.stack([key_states, value_states], 2),
-                dropout_p=dropout_rate,
-                causal=is_causal,
-                window_size=window_size,
-            )
-        else:
-            (  # pylint: disable=unbalanced-tuple-unpacking
-                q_unpad,
-                kv_unpad,
-                cu_seqlens_q,
-                cu_seqlens_k,
-                max_seqlen_q,
-                max_seqlen_k,
-                _,
-                _,
-                output_pad_fn,
-            ) = generate_qkv(
-                query_states,
-                key_states,
-                value_states,
-                kvpacked=True,
-                key_padding_mask=attention_mask,
-                query_padding_mask=attention_mask[:, -query_states.size(1) :]
-                if attention_mask is not None
-                else None,
-            )
-            if q_unpad.dtype != kv_unpad.dtype:
-                kv_unpad = kv_unpad.to(q_unpad.dtype)
-            output_unpad = flash_attn_varlen_kvpacked_func(
-                q_unpad,
-                kv_unpad,
-                cu_seqlens_q,
-                cu_seqlens_k,
-                max_seqlen_q,
-                max_seqlen_k,
-                dropout_p=dropout_rate,
-                softmax_scale=None,
-                causal=is_causal,
-                window_size=window_size,
-            )
-            output = output_pad_fn(output_unpad)
-
-    attn_output = output
-    if attn_output.size() != (bsz, q_len, self.num_heads, self.head_dim):
-        raise ValueError(
-            f"`attn_output` should be of size {(bsz, q_len, self.num_heads, self.head_dim)}, but is"
-            f" {attn_output.size()}"
-        )
-    attn_output = rearrange(attn_output, "b s h d -> b s (h d)")
-
-    attn_output = self.o_proj(attn_output)
-
-    if not output_attentions:
-        attn_weights = None
-
-    return attn_output, attn_weights, past_key_value
-
-
-# based on https://github.com/Dao-AILab/flash-attention/blob/364a5b/tests/test_flash_attn.py#L38
-def generate_qkv(
-    q,
-    k,
-    v,
-    query_padding_mask=None,
-    key_padding_mask=None,
-    kvpacked=False,
-    qkvpacked=False,
-):  # pylint: disable=invalid-name,unnecessary-lambda-assignment
-    """
-    Arguments:
-        q: (batch_size, seqlen_q, nheads, d)
-        k: (batch_size, seqlen_k, nheads_k, d)
-        v: (batch_size, seqlen_k, nheads_k, d)
-        query_padding_mask: (batch_size, seqlen), bool
-        key_padding_mask: (batch_size, seqlen), bool
-    """
-    assert not (kvpacked and qkvpacked)
-    batch_size, seqlen_q, nheads, d = q.shape
-    _, seqlen_k, nheads_k, _ = k.shape
-    assert k.shape == (batch_size, seqlen_k, nheads_k, d)
-    assert v.shape == (batch_size, seqlen_k, nheads_k, d)
-
-    if query_padding_mask is not None:
-        q_unpad, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(
-            q, query_padding_mask
-        )
-
-        output_pad_fn = lambda output_unpad: pad_input(  # noqa: E731
-            output_unpad, indices_q, batch_size, seqlen_q
-        )
-
-    else:
-        q_unpad = rearrange(q, "b s h d -> (b s) h d")
-        cu_seqlens_q = torch.arange(
-            0,
-            (batch_size + 1) * seqlen_q,
-            step=seqlen_q,
-            dtype=torch.int32,
-            device=q_unpad.device,
-        )
-        max_seqlen_q = seqlen_q
-
-        output_pad_fn = lambda output_unpad: rearrange(  # noqa: E731
-            output_unpad, "(b s) h d -> b s h d", b=batch_size
-        )
-
-    if key_padding_mask is not None:
-        k_unpad, _, cu_seqlens_k, max_seqlen_k = unpad_input(k, key_padding_mask)
-        v_unpad, _, _, _ = unpad_input(v, key_padding_mask)
-    else:
-        k_unpad = rearrange(k, "b s h d -> (b s) h d")
-        v_unpad = rearrange(v, "b s h d -> (b s) h d")
-        cu_seqlens_k = torch.arange(
-            0,
-            (batch_size + 1) * seqlen_k,
-            step=seqlen_k,
-            dtype=torch.int32,
-            device=k_unpad.device,
-        )
-        max_seqlen_k = seqlen_k
-
-    if qkvpacked:
-        assert nheads == nheads_k
-        qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)
-        qkv = torch.stack([q, k, v], dim=2)
-        return (qkv_unpad, cu_seqlens_q, max_seqlen_q, qkv, output_pad_fn)
-
-    if kvpacked:
-        kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)
-        kv = torch.stack([k, v], dim=2)
-        return (
-            q_unpad,
-            kv_unpad,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            max_seqlen_q,
-            max_seqlen_k,
-            q,
-            kv,
-            output_pad_fn,
-        )
-
-    return (
-        q_unpad,
-        k_unpad,
-        v_unpad,
-        cu_seqlens_q,
-        cu_seqlens_k,
-        max_seqlen_q,
-        max_seqlen_k,
-        q,
-        k,
-        v,
-        output_pad_fn,
-    )
-
-
-def mistral_model_forward(
-    self,
-    input_ids: torch.LongTensor = None,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[List[torch.FloatTensor]] = None,
-    inputs_embeds: Optional[torch.FloatTensor] = None,
-    use_cache: Optional[bool] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    return_dict: Optional[bool] = None,
-) -> Union[Tuple, BaseModelOutputWithPast]:
-    output_attentions = (
-        output_attentions
-        if output_attentions is not None
-        else self.config.output_attentions
-    )
-    output_hidden_states = (
-        output_hidden_states
-        if output_hidden_states is not None
-        else self.config.output_hidden_states
-    )
-    use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-    return_dict = (
-        return_dict if return_dict is not None else self.config.use_return_dict
-    )
-
-    # retrieve input_ids and inputs_embeds
-    if input_ids is not None and inputs_embeds is not None:
-        raise ValueError(
-            "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
-        )
-    if input_ids is not None:
-        batch_size, seq_length = input_ids.shape
-    elif inputs_embeds is not None:
-        batch_size, seq_length, _ = inputs_embeds.shape
-    else:
-        raise ValueError(
-            "You have to specify either decoder_input_ids or decoder_inputs_embeds"
-        )
-
-    seq_length_with_past = seq_length
-    past_key_values_length = 0
-
-    if past_key_values is not None:
-        past_key_values_length = past_key_values[0][0].shape[2]
-        seq_length_with_past = seq_length_with_past + past_key_values_length
-
-    cu_seqlens = None
-    max_seqlen = None
-    if position_ids is None:
-        device = input_ids.device if input_ids is not None else inputs_embeds.device
-        position_ids = torch.arange(
-            past_key_values_length,
-            seq_length + past_key_values_length,
-            dtype=torch.long,
-            device=device,
-        )
-        position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
-    else:
-        position_ids = position_ids.view(-1, seq_length).long()
-        cu_seqlens, max_seqlen = get_cu_seqlens_from_pos_ids(position_ids)
-        cu_seqlens = cu_seqlens.squeeze()
-
-    if inputs_embeds is None:
-        inputs_embeds = self.embed_tokens(input_ids)
-    # embed positions
-    if attention_mask is None:
-        attention_mask = torch.ones(
-            (batch_size, seq_length_with_past),
-            dtype=torch.bool,
-            device=inputs_embeds.device,
-        )
-    attention_mask = (
-        self._prepare_decoder_attention_mask(  # pylint: disable=protected-access
-            attention_mask,
-            (batch_size, seq_length),
-            inputs_embeds,
-            past_key_values_length,
-            sliding_window=self.config.sliding_window,
-        )
-    )
-
-    hidden_states = inputs_embeds
-
-    if self.gradient_checkpointing and self.training:
-        if use_cache:
-            transformers.logger.warning_once(
-                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-            )
-            use_cache = False
-
-    # decoder layers
-    all_hidden_states = () if output_hidden_states else None
-    all_self_attns = () if output_attentions else None
-    next_decoder_cache = () if use_cache else None
-
-    for idx, decoder_layer in enumerate(self.layers):
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        past_key_value = past_key_values[idx] if past_key_values is not None else None
-
-        if self.gradient_checkpointing and self.training:
-            layer_outputs = (
-                self._gradient_checkpointing_func(  # pylint: disable=protected-access
-                    decoder_layer.__call__,
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    past_key_value,
-                    output_attentions,
-                    None,
-                    cu_seqlens,
-                    max_seqlen,
-                )
-            )
-        else:
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cu_seqlens=cu_seqlens,
-                max_seqlen=max_seqlen,
-            )
-
-        hidden_states = layer_outputs[0]
-
-        if use_cache:
-            next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
-
-        if output_attentions:
-            all_self_attns += (layer_outputs[1],)
-
-    hidden_states = self.norm(hidden_states)
-
-    # add hidden states from the last decoder layer
-    if output_hidden_states:
-        all_hidden_states += (hidden_states,)
-
-    next_cache = next_decoder_cache if use_cache else None
-    if not return_dict:
-        return tuple(
-            v
-            for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
-            if v is not None
-        )
-    return BaseModelOutputWithPast(
-        last_hidden_state=hidden_states,
-        past_key_values=next_cache,
-        hidden_states=all_hidden_states,
-        attentions=all_self_attns,
-    )
-
-
-class MistralDecoderLayer(OriginalMistralDecoderLayer):
-    """
-    patched version of MistralDecoderLayer to pass through the precalculated cu_seqlens
-    """
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Tuple[torch.Tensor]] = None,
-        output_attentions: Optional[bool] = False,
-        use_cache: Optional[bool] = False,
-        cu_seqlens: Optional[torch.Tensor] = None,
-        max_seqlen: Optional[torch.Tensor] = None,
-    ) -> Tuple[
-        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
-    ]:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            use_cache (`bool`, *optional*):
-                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-                (see `past_key_values`).
-            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-            cu_seqlens (`torch.Tensor`, *optional*) cumulative sequence len when packing
-        """
-
-        residual = hidden_states
-
-        hidden_states = self.input_layernorm(hidden_states)
-
-        # Self Attention
-        hidden_states, self_attn_weights, present_key_value = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_value=past_key_value,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
-            cu_seqlens=cu_seqlens,
-            max_seqlen=max_seqlen,
-        )
-        hidden_states = residual + hidden_states
-
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-
-        outputs = (hidden_states,)
-
-        if output_attentions:
-            outputs += (self_attn_weights,)
-
-        if use_cache:
-            outputs += (present_key_value,)
-
-        return outputs
diff --git a/src/axolotl/monkeypatch/mixtral/__init__.py b/src/axolotl/monkeypatch/mixtral/__init__.py
index d6ee0ce16b..b353b12cf2 100644
--- a/src/axolotl/monkeypatch/mixtral/__init__.py
+++ b/src/axolotl/monkeypatch/mixtral/__init__.py
@@ -1,6 +1,7 @@
 """
 Patches to support multipack for mixtral
 """
+
 import torch
 
 
@@ -30,21 +31,19 @@ def moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         topk_weight = topk_weight.to(hidden_states.dtype)
 
         hidden_states = hidden_states.repeat_interleave(self.top_k, dim=0)
-        y = torch.empty_like(hidden_states)  # pylint: disable=invalid-name
+        y = torch.empty_like(hidden_states)
         flat_topk_idx = topk_idx.view(-1)
         for i in range(self.num_experts):
             expert = self.experts[i]
             y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])
-        y = (  # pylint: disable=invalid-name
-            y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)
-        ).sum(dim=1)
+        y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
         final_hidden_states = y.reshape(batch_size, sequence_length, hidden_dim)
         return final_hidden_states, router_logits
 
     from transformers.models.mixtral.modeling_mixtral import (
-        MixtralBLockSparseTop2MLP,
+        MixtralBlockSparseTop2MLP,
         MixtralSparseMoeBlock,
     )
 
-    MixtralBLockSparseTop2MLP.forward = mlp_forward
+    MixtralBlockSparseTop2MLP.forward = mlp_forward
     MixtralSparseMoeBlock.forward = moe_forward
diff --git a/src/axolotl/monkeypatch/models/__init__.py b/src/axolotl/monkeypatch/models/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/apertus/__init__.py b/src/axolotl/monkeypatch/models/apertus/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/apertus/activation.py b/src/axolotl/monkeypatch/models/apertus/activation.py
new file mode 100644
index 0000000000..d5470aceb6
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/apertus/activation.py
@@ -0,0 +1,52 @@
+"""Monkeypatch for Apertus to dtype mismatch in XIELU act"""
+
+from torch import Tensor
+
+
+def patch_apertus_xielu_activation():
+    try:
+        from transformers.activations import XIELUActivation
+    except ImportError as err:
+        raise ImportError(
+            "Cannot import XIELUActivation. "
+            "Please make sure to update your transformers version >= 4.56.1."
+        ) from err
+
+    from transformers.activations import logger
+
+    # Store the original method
+    old_fn = XIELUActivation._xielu_cuda
+
+    def _xielu_cuda_fixed(self, x: Tensor) -> Tensor:
+        """Firewall function to prevent torch.compile from seeing .item() calls"""
+        original_shape = x.shape
+        # CUDA kernel expects 3D tensors, reshape if needed
+        while x.dim() < 3:
+            x = x.unsqueeze(0)
+        if x.dim() > 3:
+            x = x.view(-1, 1, x.size(-1))
+        if original_shape != x.shape:
+            logger.warning_once(
+                "Warning: xIELU input tensor expects 3 dimensions but got (shape: %s). Reshaping to (shape: %s).",
+                original_shape,
+                x.shape,
+            )
+        result = self._xielu_cuda_obj.forward(
+            x,
+            self.alpha_p.to(x.dtype),
+            self.alpha_n.to(x.dtype),
+            # Temporary until xIELU CUDA fully implemented -> self.{beta,eps}.item()
+            self._beta_scalar,
+            self._eps_scalar,
+            self.with_vector_loads,
+        )
+        return result.view(original_shape)
+
+    # Apply the patch
+    XIELUActivation._xielu_cuda = _xielu_cuda_fixed
+
+    def unpatch():
+        """Restore the original method"""
+        XIELUActivation._xielu_cuda = old_fn
+
+    return unpatch
diff --git a/src/axolotl/monkeypatch/models/falcon_h1/__init__.py b/src/axolotl/monkeypatch/models/falcon_h1/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/falcon_h1/modeling.py b/src/axolotl/monkeypatch/models/falcon_h1/modeling.py
new file mode 100644
index 0000000000..e733a85b77
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/falcon_h1/modeling.py
@@ -0,0 +1,364 @@
+"""Sample-packing and context-parallelism patch for Falcon-H1 (parallel Mamba2/Attention hybrid).
+
+Threads seq_idx (derived from position_ids) into the Mamba2 SSM kernels so
+packed-sequence boundaries reset SSM state.  Upstream hard-codes seq_idx=None,
+which leaks hidden state across boundaries.
+
+Unlike Nemotron-H (which selects block_type per layer), Falcon-H1 runs both
+Mamba2 and Attention in **parallel** in every FalconH1DecoderLayer, so we
+always need seq_idx for the mamba branch.
+"""
+
+import importlib
+
+import torch
+
+from axolotl.monkeypatch.models.mamba_utils import (
+    ensure_mamba_kernels_loaded,
+    get_seq_idx,
+    is_cp_active,
+    wrap_mamba_scan_for_cp,
+)
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_falcon_h1_modeling_packing():
+    """Patch Falcon-H1 for sample packing: seq_idx threading into Mamba2 SSM kernels."""
+    try:
+        mod = importlib.import_module(
+            "transformers.models.falcon_h1.modeling_falcon_h1"
+        )
+    except ImportError:
+        LOG.warning("falcon_h1 not found in transformers, skipping packing patches")
+        return
+
+    ensure_mamba_kernels_loaded(mod)
+
+    FalconH1Mixer = mod.FalconH1Mixer
+    FalconH1DecoderLayer = mod.FalconH1DecoderLayer
+
+    def patched_cuda_kernels_forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params=None,
+        cache_position=None,
+        attention_mask=None,
+        seq_idx=None,
+    ):
+        hidden_states = mod.apply_mask_to_padding_states(hidden_states, attention_mask)
+        hidden_states = hidden_states * self.ssm_in_multiplier
+        projected_states = self.in_proj(hidden_states)
+        projected_states = projected_states * self.mup_vector
+        d_to_remove = (
+            2 * self.intermediate_size
+            + 2 * self.n_groups * self.ssm_state_size
+            + self.num_heads
+        )
+
+        batch_size, seq_len, _ = hidden_states.shape
+        groups_time_state_size = self.n_groups * self.ssm_state_size
+
+        use_precomputed_states = (
+            cache_params is not None
+            and cache_params.has_previous_state
+            and seq_len == 1
+            and cache_params.conv_states[self.layer_idx].shape[0]
+            == cache_params.ssm_states[self.layer_idx].shape[0]
+            == batch_size
+            and cache_position is not None
+            and cache_position[0] > 0
+        )
+
+        if use_precomputed_states:
+            d_mlp = (projected_states.squeeze(1).shape[-1] - d_to_remove) // 2
+            z0, x0, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(
+                [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads],
+                dim=-1,
+            )
+            hidden_states_B_C = mod.causal_conv1d_update(
+                hidden_states_B_C,
+                cache_params.conv_states[self.layer_idx],
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                self.activation,
+            )
+            hidden_states, B, C = torch.split(
+                hidden_states_B_C,
+                [
+                    self.intermediate_size,
+                    groups_time_state_size,
+                    groups_time_state_size,
+                ],
+                dim=-1,
+            )
+            A = -torch.exp(self.A_log.float())
+            A = (
+                A[:, None, ...][:, :, None]
+                .expand(-1, self.head_dim, self.ssm_state_size)
+                .to(dtype=torch.float32)
+            )
+            dt = dt[:, :, None].expand(-1, -1, self.head_dim)
+            dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+            D = self.D[:, None, ...].expand(-1, self.head_dim)
+            B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)
+            C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)
+            hidden_states_reshaped = hidden_states.view(
+                batch_size, self.num_heads, self.head_dim
+            )
+            hidden_states = mod.selective_state_update(
+                cache_params.ssm_states[self.layer_idx],
+                hidden_states_reshaped,
+                dt,
+                A,
+                B,
+                C,
+                D,
+                z=(
+                    gate.view(batch_size, self.num_heads, self.head_dim)
+                    if not self.mamba_rms_norm
+                    else None
+                ),
+                dt_bias=dt_bias,
+                dt_softplus=True,
+            )
+            hidden_states = hidden_states.view(
+                batch_size, self.num_heads * self.head_dim
+            )
+            if self.mamba_rms_norm:
+                hidden_states = self.norm(hidden_states, gate)
+            if d_mlp > 0:
+                hidden_states = torch.cat(
+                    [torch.nn.functional.silu(z0) * x0, hidden_states], dim=-1
+                )
+            out = self.out_proj(hidden_states[:, None, ...])
+        else:
+            A = -torch.exp(self.A_log.float())
+            dt_limit_kwargs = (
+                {}
+                if self.time_step_limit == (0.0, float("inf"))
+                else {"dt_limit": self.time_step_limit}
+            )
+
+            if self.training and cache_params is None and not is_cp_active():
+                out = mod.mamba_split_conv1d_scan_combined(
+                    projected_states,
+                    self.conv1d.weight.squeeze(1),
+                    self.conv1d.bias,
+                    self.dt_bias,
+                    A,
+                    D=self.D,
+                    chunk_size=self.chunk_size,
+                    seq_idx=seq_idx,
+                    activation=self.activation,
+                    rmsnorm_weight=(self.norm.weight if self.mamba_rms_norm else None),
+                    rmsnorm_eps=(
+                        self.norm.variance_epsilon if self.mamba_rms_norm else None
+                    ),
+                    outproj_weight=self.out_proj.weight,
+                    outproj_bias=self.out_proj.bias,
+                    headdim=self.head_dim,
+                    ngroups=self.n_groups,
+                    norm_before_gate=False,
+                    return_final_states=False,
+                    **dt_limit_kwargs,
+                )
+            else:
+                d_mlp = (
+                    projected_states.shape[-1]
+                    - 2 * self.intermediate_size
+                    - 2 * self.n_groups * self.ssm_state_size
+                    - self.num_heads
+                ) // 2
+                if attention_mask is not None:
+                    projected_states = projected_states * attention_mask[..., None]
+                _, gate, hidden_states_B_C, dt = projected_states.split(
+                    [2 * d_mlp, self.intermediate_size, self.conv_dim, self.num_heads],
+                    dim=-1,
+                )
+
+                if cache_params is not None:
+                    conv_states = torch.nn.functional.pad(
+                        hidden_states_B_C.permute(0, 2, 1),
+                        (
+                            self.conv_kernel_size - hidden_states_B_C.shape[-2],
+                            0,
+                        ),
+                    )
+                    cache_params.update_conv_state(
+                        self.layer_idx, conv_states, cache_position
+                    )
+
+                time_step = torch.nn.functional.softplus(dt + self.dt_bias)
+
+                if mod.causal_conv1d_fn is None or self.activation not in [
+                    "silu",
+                    "swish",
+                ]:
+                    hidden_states_B_C = self.act(
+                        self.conv1d(hidden_states_B_C.transpose(1, 2)).transpose(1, 2)[
+                            :, :seq_len
+                        ]
+                    )
+                else:
+                    hidden_states_B_C = mod.causal_conv1d_fn(
+                        x=hidden_states_B_C.transpose(1, 2),
+                        weight=self.conv1d.weight.squeeze(1),
+                        bias=self.conv1d.bias,
+                        activation=self.activation,
+                        seq_idx=seq_idx,
+                    ).transpose(1, 2)[:, :seq_len]
+
+                hidden_states, B, C = torch.split(
+                    hidden_states_B_C,
+                    [
+                        self.intermediate_size,
+                        groups_time_state_size,
+                        groups_time_state_size,
+                    ],
+                    dim=-1,
+                )
+
+                if (
+                    attention_mask is not None
+                    and attention_mask.shape[1] > 1
+                    and attention_mask.shape[0] > 1
+                ):
+                    dtype = hidden_states.dtype
+                    hidden_states = (hidden_states * attention_mask[:, :, None]).to(
+                        dtype
+                    )
+
+                C_reshaped = C.view(batch_size, seq_len, self.n_groups, -1)
+                with torch.cuda.device(hidden_states.device):
+                    scan_output, ssm_state = mod.mamba_chunk_scan_combined(
+                        hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+                        time_step,
+                        A,
+                        B.view(batch_size, seq_len, self.n_groups, -1),
+                        C_reshaped,
+                        chunk_size=self.chunk_size,
+                        D=self.D,
+                        z=None,
+                        seq_idx=seq_idx,
+                        return_final_states=True,
+                        **dt_limit_kwargs,
+                    )
+
+                if ssm_state is not None and cache_params is not None:
+                    cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
+                scan_output = scan_output.view(batch_size, seq_len, -1)
+                if self.mamba_rms_norm:
+                    out = self.norm(scan_output, gate)
+                else:
+                    out = scan_output * torch.nn.functional.silu(gate)
+                out = self.out_proj(out)
+        return out
+
+    FalconH1Mixer.cuda_kernels_forward = patched_cuda_kernels_forward
+
+    def patched_mixer_forward(
+        self,
+        hidden_states,
+        cache_params=None,
+        cache_position=None,
+        attention_mask=None,
+        seq_idx=None,
+    ):
+        if seq_idx is not None and mod.causal_conv1d_fn is None:
+            raise RuntimeError(
+                "Falcon-H1 sample packing requires causal_conv1d_fn. "
+                "Install with: pip install mamba-ssm causal-conv1d"
+            )
+        if mod.is_fast_path_available and "cuda" in self.in_proj.weight.device.type:
+            return self.cuda_kernels_forward(
+                hidden_states,
+                cache_params,
+                cache_position,
+                attention_mask,
+                seq_idx=seq_idx,
+            )
+        if seq_idx is not None:
+            raise RuntimeError(
+                "Falcon-H1 sample packing requires the CUDA fast path. "
+                "Ensure model is on CUDA and mamba-ssm/causal-conv1d are installed."
+            )
+        dtype = hidden_states.dtype
+        if (
+            attention_mask is not None
+            and attention_mask.shape[1] > 1
+            and attention_mask.shape[0] > 1
+        ):
+            hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+        return self.torch_forward(
+            hidden_states, cache_params, cache_position, attention_mask
+        )
+
+    FalconH1Mixer.forward = patched_mixer_forward
+
+    # Falcon-H1 runs mamba + attention in parallel every layer (no block_type).
+    # Compute seq_idx from position_ids and pass to the mamba branch.
+    def patched_decoder_forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        mamba_attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        output_attentions=False,
+        use_cache=False,
+        cache_position=None,
+        position_embeddings=None,
+        **kwargs,
+    ):
+        is_decoding = past_key_values is not None and past_key_values.has_previous_state
+        seq_idx = (
+            get_seq_idx(position_ids)
+            if position_ids is not None and not is_decoding
+            else None
+        )
+
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        mamba_hidden_states = self.mamba(
+            hidden_states=hidden_states,
+            cache_params=past_key_values,
+            cache_position=cache_position,
+            attention_mask=mamba_attention_mask,
+            seq_idx=seq_idx,
+        )
+        mamba_hidden_states = mamba_hidden_states * self.ssm_out_multiplier
+
+        attention_hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states * self.attention_in_multiplier,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        attention_hidden_states = attention_hidden_states * self.attn_out_multiplier
+
+        hidden_states = mamba_hidden_states + attention_hidden_states
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.pre_ff_layernorm(hidden_states)
+        hidden_states = self.feed_forward(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+        return outputs
+
+    FalconH1DecoderLayer.forward = patched_decoder_forward
+
+    wrap_mamba_scan_for_cp(mod)
+
+    LOG.info("Applied Falcon-H1 sample packing patch (seq_idx threading into Mamba2)")
diff --git a/src/axolotl/monkeypatch/models/gemma4/fused_attn.py b/src/axolotl/monkeypatch/models/gemma4/fused_attn.py
new file mode 100644
index 0000000000..3112b5919f
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/gemma4/fused_attn.py
@@ -0,0 +1,211 @@
+"""
+Gemma 4 fused attention monkeypatch.
+
+Replaces the per-layer RMSNorm + RoPE + transpose sequence with fused Triton
+kernels, eliminating intermediate tensor allocations from rotate_half / apply_rotary_pos_emb
+
+Usage:
+    from axolotl.monkeypatch.models.gemma4.fused_attn import patch_gemma4_fused_attn
+    # Pass install_shared_kv_workaround=True when activation checkpointing is enabled.
+    patch_gemma4_fused_attn(install_shared_kv_workaround=True)
+"""
+
+from typing import Callable
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+# Module-level dict used as a side channel for shared KV states avoiding kwarg and TLS
+# to prevent memory leak on gradient checkpoint enabled training (PR #3611)
+_GEMMA4_SHARED_KV_STORE: dict = {"store": None}
+
+
+def _set_shared_kv_states(store):
+    _GEMMA4_SHARED_KV_STORE["store"] = store
+
+
+def _get_shared_kv_states():
+    return _GEMMA4_SHARED_KV_STORE["store"]
+
+
+def _shared_kv_read_key(attn):
+    # transformers >=5.8 keys shared kv by layer_type; older builds used kv_shared_layer_index
+    if hasattr(attn, "kv_shared_layer_index"):
+        return attn.kv_shared_layer_index
+    return attn.layer_type
+
+
+def _shared_kv_store_key(attn):
+    if hasattr(attn, "kv_shared_layer_index"):
+        return attn.layer_idx
+    return attn.layer_type
+
+
+def _make_fused_forward(original_forward):
+    """Create a patched forward that uses fused RMSNorm+RoPE kernels."""
+
+    from axolotl.kernels.gemma4_fused_rope import (
+        fused_rms_norm_noscale,
+        fused_rms_norm_rope,
+    )
+
+    def fused_forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None,
+        past_key_values=None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+        from transformers.models.gemma4.modeling_gemma4 import (
+            eager_attention_forward,
+        )
+
+        store = _get_shared_kv_states()
+        if store is not None:
+            shared_kv_states = store
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        eps = self.config.rms_norm_eps
+
+        cos, sin = position_embeddings
+
+        # ---- Projections ----
+        # Use apply_qkv if present (LoRA kernel patch), otherwise direct proj
+        has_lora_qkv = hasattr(self, "apply_qkv")
+
+        if has_lora_qkv:
+            query_states, key_states, value_states = self.apply_qkv(hidden_states)
+            query_states = query_states.view(hidden_shape)
+        else:
+            query_states = self.q_proj(hidden_states).view(hidden_shape)
+
+        # ---- Q path: fused q_norm + RoPE ----
+        query_states = fused_rms_norm_rope(
+            query_states,
+            self.q_norm.weight,
+            cos,
+            sin,
+            eps=eps,
+        )
+        query_states = query_states.transpose(1, 2)
+
+        # ---- K/V path ----
+        if self.is_kv_shared_layer:
+            key_states, value_states = shared_kv_states[_shared_kv_read_key(self)]
+            key_states = key_states.to(query_states.device)
+            value_states = value_states.to(query_states.device)
+        else:
+            if has_lora_qkv:
+                # apply_qkv already computed k/v projections
+                key_states = key_states.view(hidden_shape)
+                value_states = (
+                    value_states.view(hidden_shape)
+                    if self.v_proj is not None
+                    else key_states
+                )
+            else:
+                key_states = self.k_proj(hidden_states).view(hidden_shape)
+                value_states = (
+                    self.v_proj(hidden_states).view(hidden_shape)
+                    if self.v_proj is not None
+                    else key_states
+                )
+
+            # Fused k_norm + RoPE
+            key_states = fused_rms_norm_rope(
+                key_states,
+                self.k_norm.weight,
+                cos,
+                sin,
+                eps=eps,
+            )
+            key_states = key_states.transpose(1, 2)
+
+            # Fused v_norm (no scale, no RoPE)
+            value_states = fused_rms_norm_noscale(value_states, eps=eps)
+            value_states = value_states.transpose(1, 2)
+
+        if past_key_values is not None and not self.is_kv_shared_layer:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+        if self.store_full_length_kv:
+            shared_kv_states[_shared_kv_store_key(self)] = key_states, value_states
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=self.attention_dropout if self.training else 0.0,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        # Use apply_o if present (LoRA O kernel patch), otherwise direct proj
+        if hasattr(self, "apply_o"):
+            attn_output = self.apply_o(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    return fused_forward
+
+
+def _patch_decoder_layer_call():
+    """Strip `shared_kv_states` from decoder-layer kwargs and route via the
+    module-level side channel so the checkpoint partial cannot pin it (PR #3611).
+    """
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextDecoderLayer
+
+    if getattr(Gemma4TextDecoderLayer, "_axolotl_shared_kv_patched", False):
+        return
+
+    original_call = Gemma4TextDecoderLayer.__call__
+
+    def patched_call(self, *args, **kwargs):
+        shared_kv = kwargs.pop("shared_kv_states", None)
+        # Overwrite unconditionally (including with None) so a previous step's
+        # dict cannot leak into a later call without shared_kv_states (PR #3611).
+        _set_shared_kv_states(shared_kv)
+        return original_call(self, *args, **kwargs)
+
+    Gemma4TextDecoderLayer.__call__ = patched_call
+    Gemma4TextDecoderLayer._axolotl_shared_kv_patched = True
+
+
+def patch_gemma4_fused_attn(install_shared_kv_workaround: bool = False):
+    """
+    Monkeypatch Gemma4TextAttention.forward to use fused RMSNorm+RoPE kernels,
+    and optionally route `shared_kv_states` via a module-level side channel to
+    avoid a VRAM leak under activation checkpointing (PR #3611).
+    """
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextAttention
+
+    original_forward = Gemma4TextAttention.forward
+    Gemma4TextAttention.forward = _make_fused_forward(original_forward)
+
+    if install_shared_kv_workaround:
+        _patch_decoder_layer_call()
+
+    logger.info(
+        "Patched Gemma4TextAttention.forward with fused RMSNorm+RoPE Triton kernels"
+    )
+    if install_shared_kv_workaround:
+        logger.info("Installed Gemma4 shared_kv_states side channel (PR #3611)")
diff --git a/src/axolotl/monkeypatch/models/gemma4_unified/__init__.py b/src/axolotl/monkeypatch/models/gemma4_unified/__init__.py
new file mode 100644
index 0000000000..c4ec3b89cf
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/gemma4_unified/__init__.py
@@ -0,0 +1 @@
+"""Gemma 4 Unified model monkeypatches."""
diff --git a/src/axolotl/monkeypatch/models/gemma4_unified/fused_attn.py b/src/axolotl/monkeypatch/models/gemma4_unified/fused_attn.py
new file mode 100644
index 0000000000..032804dbd2
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/gemma4_unified/fused_attn.py
@@ -0,0 +1,220 @@
+"""
+Gemma 4 Unified fused attention monkeypatch.
+
+Mirrors :mod:`axolotl.monkeypatch.models.gemma4.fused_attn` for the
+encoder-free ``gemma4_unified`` text backbone (``Gemma4UnifiedTextAttention``),
+replacing the per-layer RMSNorm + RoPE + transpose sequence with fused Triton
+kernels.
+
+The math is identical to standard Gemma 4 (q_norm/k_norm with scale + RoPE,
+v_norm without scale, no RoPE), so the same kernels are reused. Like standard
+Gemma 4 on transformers 5.10.x, the unified attention keys ``shared_kv_states``
+by **layer type string** (``self.layer_type``); the separate module is needed
+only because the unified backbone redefines its classes in its own namespace
+(it does not modular-import from ``modeling_gemma4``).
+
+Usage:
+    from axolotl.monkeypatch.models.gemma4_unified.fused_attn import (
+        patch_gemma4_unified_fused_attn,
+    )
+    # Pass install_shared_kv_workaround=True when activation checkpointing is enabled.
+    patch_gemma4_unified_fused_attn(install_shared_kv_workaround=True)
+"""
+
+from typing import Callable
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+# Module-level dict used as a side channel for shared KV states avoiding kwarg
+# and TLS to prevent a memory leak under gradient checkpointing (PR #3611).
+# Kept separate from the standard-gemma4 store so the two never alias.
+_GEMMA4_UNIFIED_SHARED_KV_STORE: dict = {"store": None}
+
+
+def _set_shared_kv_states(store):
+    _GEMMA4_UNIFIED_SHARED_KV_STORE["store"] = store
+
+
+def _get_shared_kv_states():
+    return _GEMMA4_UNIFIED_SHARED_KV_STORE["store"]
+
+
+def _make_fused_forward(original_forward):
+    """Create a patched forward that uses fused RMSNorm+RoPE kernels."""
+
+    from axolotl.kernels.gemma4_fused_rope import (
+        fused_rms_norm_noscale,
+        fused_rms_norm_rope,
+    )
+
+    def fused_forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None,
+        past_key_values=None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+        from transformers.models.gemma4_unified.modeling_gemma4_unified import (
+            eager_attention_forward,
+        )
+
+        store = _get_shared_kv_states()
+        if store is not None:
+            shared_kv_states = store
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        eps = self.config.rms_norm_eps
+
+        cos, sin = position_embeddings
+
+        # ---- Projections ----
+        # Use apply_qkv if present (LoRA kernel patch), otherwise direct proj.
+        has_lora_qkv = hasattr(self, "apply_qkv")
+
+        if has_lora_qkv:
+            query_states, key_states, value_states = self.apply_qkv(hidden_states)
+            query_states = query_states.view(hidden_shape)
+        else:
+            query_states = self.q_proj(hidden_states).view(hidden_shape)
+
+        # ---- Q path: fused q_norm + RoPE ----
+        query_states = fused_rms_norm_rope(
+            query_states,
+            self.q_norm.weight,
+            cos,
+            sin,
+            eps=eps,
+        )
+        query_states = query_states.transpose(1, 2)
+
+        # ---- K/V path ----
+        if self.is_kv_shared_layer:
+            # Unified keys the shared cache by layer-type string, not index.
+            key_states, value_states = shared_kv_states[self.layer_type]
+            key_states = key_states.to(query_states.device)
+            value_states = value_states.to(query_states.device)
+        else:
+            if has_lora_qkv:
+                # apply_qkv already computed k/v projections.
+                key_states = key_states.view(hidden_shape)
+                value_states = (
+                    value_states.view(hidden_shape)
+                    if self.v_proj is not None
+                    else key_states
+                )
+            else:
+                key_states = self.k_proj(hidden_states).view(hidden_shape)
+                # attention_k_eq_v: value reuses the (pre-norm) k_proj output.
+                value_states = (
+                    self.v_proj(hidden_states).view(hidden_shape)
+                    if self.v_proj is not None
+                    else key_states
+                )
+
+            # Fused k_norm + RoPE (creates a new tensor, leaving value_states
+            # pointing at the pre-norm k_proj output for the k_eq_v case).
+            key_states = fused_rms_norm_rope(
+                key_states,
+                self.k_norm.weight,
+                cos,
+                sin,
+                eps=eps,
+            )
+            key_states = key_states.transpose(1, 2)
+
+            # Fused v_norm (no scale, no RoPE).
+            value_states = fused_rms_norm_noscale(value_states, eps=eps)
+            value_states = value_states.transpose(1, 2)
+
+        if past_key_values is not None and not self.is_kv_shared_layer:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+        if self.store_full_length_kv:
+            shared_kv_states[self.layer_type] = key_states, value_states
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=self.attention_dropout if self.training else 0.0,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        # Use apply_o if present (LoRA O kernel patch), otherwise direct proj.
+        if hasattr(self, "apply_o"):
+            attn_output = self.apply_o(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    return fused_forward
+
+
+def _patch_decoder_layer_call():
+    """Strip ``shared_kv_states`` from decoder-layer kwargs and route via the
+    module-level side channel so the checkpoint partial cannot pin it (PR #3611).
+    """
+    from transformers.models.gemma4_unified.modeling_gemma4_unified import (
+        Gemma4UnifiedTextDecoderLayer,
+    )
+
+    if getattr(Gemma4UnifiedTextDecoderLayer, "_axolotl_shared_kv_patched", False):
+        return
+
+    original_call = Gemma4UnifiedTextDecoderLayer.__call__
+
+    def patched_call(self, *args, **kwargs):
+        shared_kv = kwargs.pop("shared_kv_states", None)
+        # Overwrite unconditionally (including with None) so a previous step's
+        # dict cannot leak into a later call without shared_kv_states (PR #3611).
+        _set_shared_kv_states(shared_kv)
+        return original_call(self, *args, **kwargs)
+
+    Gemma4UnifiedTextDecoderLayer.__call__ = patched_call
+    Gemma4UnifiedTextDecoderLayer._axolotl_shared_kv_patched = True
+
+
+def patch_gemma4_unified_fused_attn(install_shared_kv_workaround: bool = False):
+    """
+    Monkeypatch ``Gemma4UnifiedTextAttention.forward`` to use fused RMSNorm+RoPE
+    kernels, and optionally route ``shared_kv_states`` via a module-level side
+    channel to avoid a VRAM leak under activation checkpointing (PR #3611).
+    """
+    from transformers.models.gemma4_unified.modeling_gemma4_unified import (
+        Gemma4UnifiedTextAttention,
+    )
+
+    if getattr(Gemma4UnifiedTextAttention, "_axolotl_fused_attn_patched", False):
+        return
+
+    original_forward = Gemma4UnifiedTextAttention.forward
+    Gemma4UnifiedTextAttention.forward = _make_fused_forward(original_forward)
+    Gemma4UnifiedTextAttention._axolotl_fused_attn_patched = True
+
+    if install_shared_kv_workaround:
+        _patch_decoder_layer_call()
+
+    logger.info(
+        "Patched Gemma4UnifiedTextAttention.forward with fused RMSNorm+RoPE "
+        "Triton kernels"
+    )
+    if install_shared_kv_workaround:
+        logger.info("Installed Gemma4Unified shared_kv_states side channel (PR #3611)")
diff --git a/src/axolotl/monkeypatch/models/granitemoehybrid/__init__.py b/src/axolotl/monkeypatch/models/granitemoehybrid/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/granitemoehybrid/modeling.py b/src/axolotl/monkeypatch/models/granitemoehybrid/modeling.py
new file mode 100644
index 0000000000..590db5408d
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/granitemoehybrid/modeling.py
@@ -0,0 +1,107 @@
+"""Sample-packing and context-parallelism patch for Granite MoE Hybrid (Mamba2/Attention/MoE).
+
+Upstream GraniteMoeHybridMambaLayer already accepts seq_idx on
+forward/cuda_kernels_forward, and GraniteMoeHybridDecoderLayer passes **kwargs
+through to the mixer.  However, the decoder layer does not receive position_ids
+directly — it arrives at the model level.
+
+This patch:
+1. Injects seq_idx computation into GraniteMoeHybridModel.forward so it flows
+   through kwargs -> decoder_layer -> mamba mixer automatically.
+2. Forces the slow path when CP is active (the fused path doesn't return SSM
+   state).  CP correction is handled by ``wrap_mamba_scan_for_cp``.
+"""
+
+import importlib
+
+from axolotl.monkeypatch.models.mamba_utils import (
+    ensure_mamba_kernels_loaded,
+    get_seq_idx,
+    is_cp_active,
+    wrap_mamba_scan_for_cp,
+)
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_granitemoehybrid_modeling_packing():
+    """Patch Granite MoE Hybrid for sample packing: seq_idx + CP correction."""
+    try:
+        mod = importlib.import_module(
+            "transformers.models.granitemoehybrid.modeling_granitemoehybrid"
+        )
+    except ImportError:
+        LOG.warning(
+            "granitemoehybrid not found in transformers, skipping packing patches"
+        )
+        return
+
+    ensure_mamba_kernels_loaded(mod)
+
+    GraniteMoeHybridModel = mod.GraniteMoeHybridModel
+    GraniteMoeHybridMambaLayer = mod.GraniteMoeHybridMambaLayer
+
+    # Patch 1: Model-level seq_idx injection
+    original_model_forward = GraniteMoeHybridModel.forward
+
+    def patched_model_forward(self, *args, **kwargs):
+        position_ids = kwargs.get("position_ids")
+        if position_ids is None and len(args) > 2:
+            position_ids = args[2]
+
+        past_key_values = kwargs.get("past_key_values")
+        if past_key_values is None and len(args) > 3:
+            past_key_values = args[3]
+
+        is_decoding = (
+            past_key_values is not None
+            and hasattr(past_key_values, "has_previous_state")
+            and past_key_values.has_previous_state
+        )
+
+        if position_ids is not None and not is_decoding and "seq_idx" not in kwargs:
+            kwargs["seq_idx"] = get_seq_idx(position_ids)
+
+        return original_model_forward(self, *args, **kwargs)
+
+    GraniteMoeHybridModel.forward = patched_model_forward
+
+    # Patch 2: Minimal wrapper to force slow path when CP is active.
+    # The fused mamba_split_conv1d_scan_combined doesn't return SSM state, so
+    # CP correction (handled by the scan wrapper) needs the slow path.
+    original_cuda_kernels_forward = GraniteMoeHybridMambaLayer.cuda_kernels_forward
+
+    def patched_cuda_kernels_forward(
+        self,
+        hidden_states,
+        cache_params=None,
+        attention_mask=None,
+        seq_idx=None,
+    ):
+        force_slow = (
+            (seq_idx is not None or is_cp_active())
+            and self.training
+            and cache_params is None
+        )
+        if force_slow:
+            self.training = False
+        try:
+            return original_cuda_kernels_forward(
+                self,
+                hidden_states,
+                cache_params,
+                attention_mask,
+                seq_idx,
+            )
+        finally:
+            if force_slow:
+                self.training = True
+
+    GraniteMoeHybridMambaLayer.cuda_kernels_forward = patched_cuda_kernels_forward
+
+    wrap_mamba_scan_for_cp(mod)
+
+    LOG.info(
+        "Applied Granite MoE Hybrid sample packing patch (seq_idx + CP correction)"
+    )
diff --git a/src/axolotl/monkeypatch/models/kimi_linear/__init__.py b/src/axolotl/monkeypatch/models/kimi_linear/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/kimi_linear/configuration_kimi.py b/src/axolotl/monkeypatch/models/kimi_linear/configuration_kimi.py
new file mode 100644
index 0000000000..1dd0e67025
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/kimi_linear/configuration_kimi.py
@@ -0,0 +1,148 @@
+"""
+Kimi-Linear configuration.
+
+Source: https://huggingface.co/moonshotai/Kimi-Linear-48B-A3B-Instruct/blob/main/configuration_kimi.py
+Revision: 6e163f3
+"""
+
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class KimiLinearConfig(PretrainedConfig):
+    model_type = "kimi_linear"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        model_type="kimi_linear",
+        vocab_size=163840,
+        hidden_size=4096,
+        head_dim=None,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        tie_word_embeddings=False,
+        moe_intermediate_size: Optional[int] = None,
+        moe_renormalize: bool = True,
+        moe_router_activation_func: str = "sigmoid",
+        num_experts: Optional[int] = None,
+        num_experts_per_token: Optional[int] = None,
+        num_shared_experts: int = 0,
+        routed_scaling_factor: float = 1.0,
+        first_k_dense_replace: int = 0,
+        moe_layer_freq: int = 1,
+        use_grouped_topk: bool = True,
+        num_expert_group: int = 1,
+        topk_group: int = 1,
+        q_lora_rank: Optional[int] = None,
+        kv_lora_rank: Optional[int] = None,
+        qk_nope_head_dim: Optional[int] = None,
+        qk_rope_head_dim: Optional[int] = None,
+        v_head_dim: Optional[int] = None,
+        mla_use_nope: Optional[bool] = False,
+        num_nextn_predict_layers: int = 0,
+        linear_attn_config: Optional[dict] = None,
+        router_aux_loss_coef: float = 0.01,
+        **kwargs,
+    ):
+        self.model_type = model_type
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.head_dim = (
+            head_dim if head_dim is not None else hidden_size // num_attention_heads
+        )
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.mla_use_nope = mla_use_nope
+        # moe config
+        self.num_experts = num_experts
+        self.num_experts_per_token = num_experts_per_token
+        self.moe_renormalize = moe_renormalize
+        self.num_shared_experts = num_shared_experts
+        self.routed_scaling_factor = routed_scaling_factor
+        self.moe_router_activation_func = moe_router_activation_func
+        assert self.moe_router_activation_func in ("softmax", "sigmoid")
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.moe_layer_freq = moe_layer_freq
+        self.use_grouped_topk = use_grouped_topk
+        self.num_expert_group = num_expert_group
+        self.topk_group = topk_group
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+        self.router_aux_loss_coef = router_aux_loss_coef
+
+        if linear_attn_config is not None:
+            assert linear_attn_config["kda_layers"] is not None
+            assert linear_attn_config["full_attn_layers"] is not None
+        self.linear_attn_config = linear_attn_config
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def is_mla(self):
+        return (
+            self.q_lora_rank is not None
+            or self.kv_lora_rank is not None
+            or self.qk_nope_head_dim is not None
+            or self.qk_rope_head_dim is not None
+            or self.v_head_dim is not None
+            or self.mla_use_nope is True
+        )
+
+    @property
+    def is_moe(self):
+        return self.num_experts is not None
+
+    @property
+    def is_linear_attn(self) -> bool:
+        return not (
+            self.linear_attn_config is None
+            or (
+                isinstance(self.linear_attn_config, dict)
+                and self.linear_attn_config["kda_layers"] is not None
+                and len(self.linear_attn_config["kda_layers"]) == 0
+            )
+        )
+
+    def is_kda_layer(self, layer_idx: int):
+        return (
+            self.linear_attn_config is not None
+            and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
+        )
diff --git a/src/axolotl/monkeypatch/models/kimi_linear/modeling_kimi.py b/src/axolotl/monkeypatch/models/kimi_linear/modeling_kimi.py
new file mode 100644
index 0000000000..42a11ec365
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/kimi_linear/modeling_kimi.py
@@ -0,0 +1,1361 @@
+"""
+Adapted Kimi-Linear modeling to enable MoE differentiable.
+
+Source: https://huggingface.co/moonshotai/Kimi-Linear-48B-A3B-Instruct/blob/main/modeling_kimi.py
+Revision: 6e163f3
+"""
+
+import math
+from collections.abc import Callable
+from typing import Any, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import transformers
+from einops import rearrange
+from packaging import version
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    MoeCausalLMOutputWithPast,
+)
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from transformers.utils import (
+    TransformersKwargs,
+    can_return_tuple,
+    logging,
+)
+from transformers.utils.generic import OutputRecorder
+
+try:
+    from fla.layers.utils import get_unpad_data, index_first_axis, pad_input
+    from fla.modules import FusedRMSNormGated, ShortConvolution
+    from fla.ops.kda import chunk_kda, fused_recurrent_kda
+    from fla.ops.kda.gate import fused_kda_gate
+except ImportError as err:
+    raise ImportError(
+        "Plese run `pip uninstall fla-core flash-linear-attention -y && pip install git+https://github.com/fla-org/flash-linear-attention@v0.4.0`"
+    ) from err
+
+from axolotl.monkeypatch.models.kimi_linear.configuration_kimi import KimiLinearConfig
+
+assert version.parse(transformers.__version__) >= version.parse("4.56.0"), (
+    "Please upgrade transformers to >= 4.56.0"
+)
+
+logger = logging.get_logger(__name__)
+
+
+def load_balancing_loss_func(
+    gate_logits: Union[torch.Tensor, tuple[torch.Tensor], None],
+    num_experts: Optional[int] = None,
+    top_k=2,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, int]:
+    """Standard Switch Transformer load balancing loss."""
+    if gate_logits is None or not isinstance(gate_logits, tuple):
+        return 0
+
+    # Concatenate all layer logits
+    concatenated_gate_logits = torch.cat(
+        [layer_gate for layer_gate in gate_logits], dim=0
+    )
+
+    routing_weights = F.softmax(concatenated_gate_logits, dim=-1)
+    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+    expert_mask = F.one_hot(selected_experts, num_experts)
+
+    tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+    router_prob_per_expert = torch.mean(routing_weights, dim=0)
+
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return overall_loss * num_experts
+
+
+class KimiDynamicCache:
+    """
+    Dynamic cache for Kimi model.
+    Inspired by Qwen3-Next
+    """
+
+    is_compileable = False
+
+    def __init__(self, config: KimiLinearConfig):
+        super().__init__()
+        self.config = config
+
+        if config.linear_attn_config is not None:
+            self.layer_types = []
+            for i in range(config.num_hidden_layers):
+                if config.is_kda_layer(i):
+                    self.layer_types.append("linear_attention")
+                else:
+                    self.layer_types.append("full_attention")
+        else:
+            self.layer_types = ["full_attention"] * config.num_hidden_layers
+
+        self.transformer_layers = [
+            i
+            for i in range(config.num_hidden_layers)
+            if self.layer_types[i] == "full_attention"
+        ]
+
+        linear_layers = [
+            i
+            for i in range(config.num_hidden_layers)
+            if self.layer_types[i] == "linear_attention"
+        ]
+        self.last_linear_layer = linear_layers[-1] if linear_layers else -1
+
+        self.conv_states = [None for _ in range(config.num_hidden_layers)]
+        self.recurrent_states = [None for _ in range(config.num_hidden_layers)]
+        self.key_cache = [None for _ in range(config.num_hidden_layers)]
+        self.value_cache = [None for _ in range(config.num_hidden_layers)]
+
+    def __len__(self):
+        return len(self.layer_types)
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[dict[str, Any]] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.key_cache[layer_idx] is None:
+            self.key_cache[layer_idx] = key_states
+            self.value_cache[layer_idx] = value_states
+        else:
+            self.key_cache[layer_idx] = torch.cat(
+                [self.key_cache[layer_idx], key_states], dim=2
+            )
+            self.value_cache[layer_idx] = torch.cat(
+                [self.value_cache[layer_idx], value_states], dim=2
+            )
+
+        return self.key_cache[layer_idx], self.value_cache[layer_idx]
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        """Reorders the cache for beam search, given the selected beam indices."""
+        for layer_idx in range(len(self.key_cache)):
+            if self.key_cache[layer_idx] is not None:
+                device = self.key_cache[layer_idx].device
+                beam_idx = beam_idx.to(device)
+                self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(
+                    0, beam_idx
+                )
+                self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(
+                    0, beam_idx
+                )
+
+            if self.conv_states[layer_idx] is not None:
+                device = self.conv_states[layer_idx][0].device
+                beam_idx = beam_idx.to(device)
+                q_conv, k_conv, v_conv = self.conv_states[layer_idx]
+                self.conv_states[layer_idx] = (
+                    q_conv.index_select(0, beam_idx),
+                    k_conv.index_select(0, beam_idx),
+                    v_conv.index_select(0, beam_idx),
+                )
+                self.recurrent_states[layer_idx] = self.recurrent_states[
+                    layer_idx
+                ].index_select(0, beam_idx)
+
+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
+        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
+        # take any layer that contains cache and not empty tensor
+        layer_idx = (
+            self.transformer_layers[0]
+            if layer_idx not in self.transformer_layers
+            else layer_idx
+        )
+        if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None:
+            return 0
+        return self.key_cache[layer_idx].shape[-2]
+
+    def get_mask_sizes(
+        self, cache_position: torch.Tensor, layer_idx: int
+    ) -> tuple[int, int]:
+        """
+        Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
+        the given layer at `layer_idx`.
+        The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.
+        """
+        kv_offset = 0
+        query_length = cache_position.shape[0]
+        past_seen_tokens = self.get_seq_length(layer_idx)
+        kv_length = query_length + past_seen_tokens
+        return kv_length, kv_offset
+
+    @property
+    def has_previous_state(self):
+        """We have a previous state if the last linear (conv) layer was already updated."""
+        if self.last_linear_layer == -1:
+            return False
+        return self.conv_states[self.last_linear_layer] is not None
+
+
+class KimiRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        KimiRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+ALL_LAYERNORM_LAYERS.append(KimiRMSNorm)
+
+
+class KimiBlockSparseMLP(nn.Module):
+    def __init__(
+        self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None
+    ):
+        super().__init__()
+        self.config = config
+        self.ffn_dim = (
+            config.intermediate_size if intermediate_size is None else intermediate_size
+        )
+        self.hidden_dim = config.hidden_size if hidden_size is None else hidden_size
+
+        self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)  # gate
+        self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)  # down
+        self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)  # up
+
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_states):
+        current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(
+            hidden_states
+        )
+        current_hidden_states = self.w2(current_hidden_states)
+        return current_hidden_states
+
+
+class KimiMLP(nn.Module):
+    def __init__(
+        self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None
+    ):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
+        self.intermediate_size = (
+            config.intermediate_size if intermediate_size is None else intermediate_size
+        )
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.dropout(
+        attn_weights, p=dropout, training=module.training
+    )
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class KimiMLAAttention(nn.Module):
+    """
+    Multi-Latent Attention adapted from deepseek-v3
+    """
+
+    def __init__(self, config: KimiLinearConfig, layer_idx: int):
+        nn.Module.__init__(self)
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+
+        self.rope_theta = config.rope_theta
+        self.attention_dropout = getattr(config, "attention_dropout", 0.0)
+
+        try:
+            self.q_lora_rank = config.q_lora_rank
+            self.qk_rope_head_dim = config.qk_rope_head_dim
+            self.kv_lora_rank = config.kv_lora_rank
+            self.v_head_dim = config.v_head_dim
+            self.qk_nope_head_dim = config.qk_nope_head_dim
+            self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+            self.use_nope = config.mla_use_nope
+            self.scaling = self.q_head_dim ** (-0.5)
+        except Exception as e:
+            raise ValueError(
+                f"Kimi MLA config is not found or not properly formatted: {e}"
+            ) from e
+
+        assert self.q_lora_rank is None
+        self.q_proj = nn.Linear(
+            self.hidden_size,
+            self.num_heads * self.q_head_dim,
+            bias=False,
+        )
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=False,
+        )
+        self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank)
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_heads
+            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=False,
+        )
+        self.is_causal = True
+        assert self.use_nope
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        batch_size, seq_length = hidden_states.shape[:-1]
+        query_shape = (batch_size, seq_length, -1, self.q_head_dim)
+        key_shape = (
+            batch_size,
+            seq_length,
+            -1,
+            self.qk_nope_head_dim + self.v_head_dim,
+        )
+
+        q_states = self.q_proj(hidden_states)
+        q_states = q_states.view(query_shape).transpose(1, 2)
+        q_pass, q_rot = torch.split(
+            q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+        )
+
+        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+        k_pass, k_rot = torch.split(
+            compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+        )
+
+        k_pass = (
+            self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2)
+        )
+        k_pass, value_states = torch.split(
+            k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
+        )
+
+        k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)
+        k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
+
+        query_states = torch.cat((q_pass, q_rot), dim=-1)
+        key_states = torch.cat((k_pass, k_rot), dim=-1)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        if (
+            self.config._attn_implementation == "flash_attention_2"
+            and self.q_head_dim != self.v_head_dim
+        ):
+            value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim])
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+
+        attn_output, _ = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        if (
+            self.config._attn_implementation == "flash_attention_2"
+            and self.q_head_dim != self.v_head_dim
+        ):
+            attn_output = attn_output[:, :, :, : self.v_head_dim]
+
+        attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output
+
+
+class KimiDeltaAttention(nn.Module):
+    def __init__(self, config: KimiLinearConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.mode = "chunk"
+
+        self.hidden_size = config.hidden_size
+        self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
+        self.head_dim = config.linear_attn_config["head_dim"]
+        self.num_heads = config.linear_attn_config["num_heads"]
+        self.head_k_dim = self.head_dim
+        self.num_k_heads = self.num_heads
+
+        self.layer_idx = layer_idx
+
+        assert self.mode in ["chunk", "fused_recurrent"], (
+            f"Not suppoerted mode `{self.mode}`."
+        )
+
+        projection_k_size = self.head_k_dim * self.num_k_heads
+        projection_size = self.head_dim * self.num_heads
+
+        self.q_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
+
+        self.q_conv1d = ShortConvolution(
+            hidden_size=projection_k_size,
+            kernel_size=self.conv_size,
+            activation="silu",
+        )
+        self.k_conv1d = ShortConvolution(
+            hidden_size=projection_k_size, kernel_size=self.conv_size, activation="silu"
+        )
+        self.v_conv1d = ShortConvolution(
+            hidden_size=projection_size, kernel_size=self.conv_size, activation="silu"
+        )
+
+        self.A_log = torch.nn.Parameter(
+            torch.log(
+                torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16)
+            ).view(1, 1, -1, 1)
+        )
+
+        self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
+
+        self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32))
+
+        self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)
+
+        self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
+
+        self.o_norm = FusedRMSNormGated(
+            self.head_dim, eps=config.rms_norm_eps, activation="sigmoid"
+        )
+        self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        cache_params: Optional[KimiDynamicCache] = None,
+        **kwargs: Unpack[dict],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        if attention_mask is not None:
+            if attention_mask.dim() != 2:
+                attention_mask = kwargs.get("padding_mask", None)
+
+            if attention_mask is not None and attention_mask.dim() != 2:
+                raise ValueError(
+                    "attention_mask must be a 0-1 matrix of shape [batch_size, seq_len] "
+                    "(0 = padding). 3D masks are not supported here."
+                )
+        use_cache = cache_params is not None
+        batch_size, q_len, _ = hidden_states.shape
+        mode = "fused_recurrent" if q_len <= 64 else self.mode
+        if self.training:
+            assert mode == "chunk", "Only chunk mode is supported in training."
+
+        cu_seqlens = kwargs.get("cu_seqlens", None)
+        indices = None
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(
+                rearrange(hidden_states, "b s ... -> (b s) ..."), indices
+            ).unsqueeze(0)
+
+        conv_state_q, conv_state_k, conv_state_v = None, None, None
+        recurrent_state = None
+        if cache_params is not None:
+            if cache_params.conv_states[self.layer_idx] is not None:
+                conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[
+                    self.layer_idx
+                ]
+            recurrent_state = cache_params.recurrent_states[self.layer_idx]
+        q, conv_state_q = self.q_conv1d(
+            x=self.q_proj(hidden_states),
+            cache=conv_state_q,
+            output_final_state=use_cache,
+            cu_seqlens=cu_seqlens,
+        )
+        k, conv_state_k = self.k_conv1d(
+            x=self.k_proj(hidden_states),
+            cache=conv_state_k,
+            output_final_state=use_cache,
+            cu_seqlens=cu_seqlens,
+        )
+        v, conv_state_v = self.v_conv1d(
+            x=self.v_proj(hidden_states),
+            cache=conv_state_v,
+            output_final_state=use_cache,
+            cu_seqlens=cu_seqlens,
+        )
+        g = self.f_b_proj(self.f_a_proj(hidden_states))
+        g = fused_kda_gate(g, self.A_log, self.head_dim, g_bias=self.dt_bias)
+        beta = self.b_proj(hidden_states).float().sigmoid()
+
+        q, k = map(
+            lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim), (q, k)
+        )
+        v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim)
+
+        if mode == "chunk":
+            o, recurrent_state = chunk_kda(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=True,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            o, recurrent_state = fused_recurrent_kda(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=True,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
+            )
+        if cache_params is not None:
+            cache_params.recurrent_states[self.layer_idx] = recurrent_state
+            cache_params.conv_states[self.layer_idx] = (
+                conv_state_q,
+                conv_state_k,
+                conv_state_v,
+            )
+
+        g = self.g_b_proj(self.g_a_proj(hidden_states))
+        g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim)
+        o = self.o_norm(o, g)
+
+        o = rearrange(o, "b t h d -> b t (h d)")
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o
+
+
+class KimiMoEGate(nn.Module):
+    """
+    MoE Gate that returns router logits.
+    Routing decisions are made in KimiSparseMoeBlock.
+    """
+
+    def __init__(self, config: KimiLinearConfig):
+        super().__init__()
+        self.config = config
+        self.num_experts = config.num_experts
+        self.gating_dim = config.hidden_size
+
+        self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim)))
+        self.e_score_correction_bias = nn.Parameter(torch.zeros((self.num_experts,)))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: [batch_size, seq_len, hidden_dim]
+
+        Returns:
+            router_logits: [batch_size * seq_len, num_experts]
+        """
+        _, _, h = hidden_states.shape
+        hidden_states = hidden_states.view(-1, h)
+        router_logits = F.linear(
+            hidden_states.type(torch.float32), self.weight.type(torch.float32), None
+        )
+        return router_logits
+
+    # def forward(self, hidden_states):
+    #     bsz, seq_len, h = hidden_states.shape
+    #     # compute gating score
+    #     hidden_states = hidden_states.view(-1, h)
+    #     logits = F.linear(
+    #         hidden_states.type(torch.float32), self.weight.type(
+    #             torch.float32), None
+    #     )
+    #     if self.moe_router_activation_func == "sigmoid":
+    #         scores = logits.sigmoid()
+    #     elif self.moe_router_activation_func == "softmax":
+    #         scores = logits.softmax(dim=1)
+    #     else:
+    #         raise NotImplementedError(
+    #             f"insupportable scoring function for MoE gating: {self.moe_router_activation_func}"
+    #         )
+
+    #     # select top-k experts
+    #     assert not self.training
+    #     scores_for_choice = scores.view(bsz * seq_len, -1)
+    #     scores_for_choice += self.e_score_correction_bias.unsqueeze(0)
+    #     group_scores = (
+    #         scores_for_choice.view(
+    #             bsz * seq_len, self.num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+    #     )  # [n, num_expert_group]
+    #     group_idx = torch.topk(
+    #         group_scores, k=self.topk_group, dim=-1, sorted=False
+    #     )[
+    #         1
+    #     ]  # [n, top_k_group]
+    #     group_mask = torch.zeros_like(group_scores)  # [n, num_expert_group]
+    #     group_mask.scatter_(1, group_idx, 1)  # [n, num_expert_group]
+    #     score_mask = (
+    #         group_mask.unsqueeze(-1)
+    #         .expand(
+    #             bsz * seq_len, self.num_expert_group, self.num_experts // self.num_expert_group
+    #         )
+    #         .reshape(bsz * seq_len, -1)
+    #     )  # [n, e]
+    #     tmp_scores = scores_for_choice.masked_fill(
+    #         ~score_mask.bool(), 0.0)  # [n, e]
+    #     _, topk_idx = torch.topk(
+    #         tmp_scores, k=self.top_k, dim=-1, sorted=False
+    #     )
+    #     topk_weight = scores.gather(1, topk_idx)
+
+    #     # norm gate to sum 1
+    #     if self.top_k > 1 and self.moe_renormalize:
+    #         denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
+    #         topk_weight = topk_weight / denominator
+    #     # must multiply the scaling factor
+    #     topk_weight = topk_weight * self.routed_scaling_factor
+
+    #     return topk_idx, topk_weight
+
+
+# class KimiSparseMoeBlock(nn.Module):
+#     """
+#     Adapted from Deepseek-V3's MOE implementation
+#     The namings are consistent with Kimi's version.
+#     """
+
+#     def __init__(self, config: KimiLinearConfig):
+#         super().__init__()
+#         self.config = config
+#         self.hidden_dim = config.hidden_size
+#         self.num_experts = config.num_experts
+#         self.top_k = config.num_experts_per_token
+#         self.moe_renormalize = config.moe_renormalize
+
+#         self.ep_size = 1
+#         self.experts_per_rank = config.num_experts
+#         self.ep_rank = 0
+#         self.experts = nn.ModuleList(
+#             [
+#                 KimiBlockSparseMLP(
+#                     config, intermediate_size=config.moe_intermediate_size
+#                 )
+#                 for _ in range(config.num_experts)
+#             ]
+#         )
+#         self.gate = KimiMoEGate(config)
+#         if config.num_shared_experts is not None:
+#             intermediate_size = config.moe_intermediate_size * config.num_shared_experts
+#             self.shared_experts = KimiMLP(
+#                 config=config, intermediate_size=intermediate_size
+#             )
+
+#     def forward(self, hidden_states):
+#         identity = hidden_states
+#         orig_shape = hidden_states.shape
+#         topk_idx, topk_weight = self.gate(hidden_states)
+#         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+#         flat_topk_idx = topk_idx.view(-1)
+#         if not self.training:
+#             y = self.moe_infer(hidden_states, topk_idx,
+#                                topk_weight).view(*orig_shape)
+#         else:
+#             raise NotImplementedError(
+#                 "Training mode is not supported in KimiSparseMoeBlock")
+#         if self.config.num_shared_experts is not None:
+#             y = y + self.shared_experts(identity)
+#         return y
+
+#     @torch.no_grad()
+#     def moe_infer(self, x, topk_ids, topk_weight):
+#         cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+#         cnts.scatter_(1, topk_ids, 1)
+#         tokens_per_expert = cnts.sum(dim=0)
+#         idxs = topk_ids.view(-1).argsort()
+#         sorted_tokens = x[idxs // topk_ids.shape[1]]
+
+#         tokens_per_expert = tokens_per_expert.cpu().numpy()
+
+#         outputs = []
+#         start_idx = 0
+#         for i, num_tokens in enumerate(tokens_per_expert):
+#             end_idx = start_idx + num_tokens
+#             if num_tokens == 0:
+#                 continue
+#             expert = self.experts[i + self.ep_rank * self.experts_per_rank]
+#             tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+#             expert_out = expert(tokens_for_this_expert)
+#             outputs.append(expert_out)
+#             start_idx = end_idx
+
+#         outs = torch.cat(outputs, dim=0) if len(
+#             outputs) else sorted_tokens.new_empty(0)
+
+#         new_x = torch.empty_like(outs)
+#         new_x[idxs] = outs
+#         final_out = (
+#             new_x.view(*topk_ids.shape, -1)
+#             .type(topk_weight.dtype)
+#             .mul_(topk_weight.unsqueeze(dim=-1))
+#             .sum(dim=1)
+#             .type(new_x.dtype)
+#         )
+#         return final_out
+
+
+# Replace the KimiSparseMoeBlock class with this new version
+class KimiSparseMoeBlock(nn.Module):
+    """
+    MoE block adapted from Deepseek-V3.
+    Returns only hidden_states - router_logits captured by OutputRecorder.
+    """
+
+    def __init__(self, config: KimiLinearConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_dim = config.hidden_size
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_token
+        self.moe_renormalize = config.moe_renormalize
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.num_expert_group = getattr(config, "num_expert_group", 1)
+        self.topk_group = getattr(config, "topk_group", 1)
+
+        self.experts = nn.ModuleList(
+            [
+                KimiBlockSparseMLP(
+                    config, intermediate_size=config.moe_intermediate_size
+                )
+                for _ in range(config.num_experts)
+            ]
+        )
+        self.gate = KimiMoEGate(config)
+
+        if config.num_shared_experts is not None:
+            intermediate_size = config.moe_intermediate_size * config.num_shared_experts
+            self.shared_experts = KimiMLP(
+                config=config, intermediate_size=intermediate_size
+            )
+
+    def route_tokens_to_experts(
+        self,
+        router_logits: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute routing decisions from router logits.
+
+        Args:
+            router_logits: [num_tokens, num_experts]
+
+        Returns:
+            topk_idx: [num_tokens, top_k]
+            topk_weight: [num_tokens, top_k]
+        """
+        num_tokens = router_logits.shape[0]
+
+        if self.training:
+            # Training: use softmax for standard aux loss compatibility
+            scores = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+            topk_weight, topk_idx = torch.topk(scores, self.top_k, dim=-1, sorted=False)
+        else:
+            # Inference: use original sigmoid + group selection
+            scores = router_logits.sigmoid()
+            scores_for_choice = scores + self.gate.e_score_correction_bias.unsqueeze(0)
+
+            # Group-based selection
+            group_scores = (
+                scores_for_choice.view(num_tokens, self.num_expert_group, -1)
+                .topk(2, dim=-1)[0]
+                .sum(dim=-1)
+            )
+            group_idx = torch.topk(
+                group_scores, k=self.topk_group, dim=-1, sorted=False
+            )[1]
+            group_mask = torch.zeros_like(group_scores)
+            group_mask.scatter_(1, group_idx, 1)
+            score_mask = (
+                group_mask.unsqueeze(-1)
+                .expand(
+                    num_tokens,
+                    self.num_expert_group,
+                    self.num_experts // self.num_expert_group,
+                )
+                .reshape(num_tokens, -1)
+            )
+            tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), 0.0)
+            _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False)
+            topk_weight = scores.gather(1, topk_idx)
+
+        # Normalize and scale
+        if self.top_k > 1 and self.moe_renormalize:
+            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
+            topk_weight = topk_weight / denominator
+        topk_weight = topk_weight * self.routed_scaling_factor
+
+        return topk_idx, topk_weight
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass returning only hidden_states.
+        Router logits are captured by OutputRecorder for aux loss.
+        """
+        identity = hidden_states
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        num_tokens = batch_size * seq_len
+
+        # Flatten for routing
+        hidden_states_flat = hidden_states.view(num_tokens, hidden_dim)
+
+        # Get router logits - OutputRecorder captures this!
+        router_logits = self.gate(hidden_states)
+
+        # Get routing decisions
+        topk_idx, topk_weight = self.route_tokens_to_experts(router_logits)
+
+        if self.training:
+            final_hidden_states = self._training_forward(
+                hidden_states_flat, topk_idx, topk_weight, num_tokens, hidden_dim
+            )
+        else:
+            final_hidden_states = self._inference_forward(
+                hidden_states_flat, topk_idx, topk_weight
+            )
+
+        final_hidden_states = final_hidden_states.view(batch_size, seq_len, hidden_dim)
+
+        # Add shared experts if present
+        if self.config.num_shared_experts is not None:
+            final_hidden_states = final_hidden_states + self.shared_experts(identity)
+
+        return final_hidden_states
+
+    def _training_forward(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weight: torch.Tensor,
+        num_tokens: int,
+        hidden_dim: int,
+    ) -> torch.Tensor:
+        """
+        Differentiable training forward using scatter-gather pattern.
+        """
+        # Flatten expert indices: [num_tokens * top_k]
+        flat_topk_idx = topk_idx.view(-1)
+
+        # Sort by expert index to group tokens going to same expert
+        sorted_indices = torch.argsort(flat_topk_idx)
+        inverse_permutation = torch.argsort(sorted_indices)
+
+        # Each token appears top_k times (once per expert choice)
+        token_indices = torch.arange(
+            num_tokens, device=hidden_states.device
+        ).repeat_interleave(self.top_k)
+
+        # Gather tokens and weights in sorted order
+        shuffled_tokens = hidden_states[token_indices[sorted_indices]]
+        shuffled_weights = topk_weight.view(-1)[sorted_indices].unsqueeze(-1)
+
+        # Count tokens per expert
+        tokens_per_expert = F.one_hot(flat_topk_idx, num_classes=self.num_experts).sum(
+            dim=0
+        )
+
+        # Process each expert's batch
+        expert_outputs = []
+        current_pos = 0
+        for i in range(self.num_experts):
+            num_tokens_for_expert = tokens_per_expert[i].item()
+            if num_tokens_for_expert == 0:
+                continue
+
+            expert_input = shuffled_tokens[
+                current_pos : current_pos + num_tokens_for_expert
+            ]
+            expert_output = self.experts[i](expert_input)
+            expert_outputs.append(expert_output)
+            current_pos += num_tokens_for_expert
+
+        # Concatenate all outputs
+        if expert_outputs:
+            concatenated_outputs = torch.cat(expert_outputs, dim=0)
+        else:
+            concatenated_outputs = torch.zeros(
+                num_tokens * self.top_k,
+                hidden_dim,
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )
+
+        # Apply weights while still in sorted order
+        weighted_outputs = concatenated_outputs * shuffled_weights
+
+        # Unsort back to original token order
+        unshuffled_outputs = weighted_outputs[inverse_permutation]
+
+        # Sum contributions from all top_k experts for each token
+        final_hidden_states = unshuffled_outputs.view(
+            num_tokens, self.top_k, hidden_dim
+        ).sum(dim=1)
+
+        return final_hidden_states
+
+    @torch.no_grad()
+    def _inference_forward(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weight: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Optimized inference forward (original implementation).
+        """
+        cnts = topk_idx.new_zeros((topk_idx.shape[0], len(self.experts)))
+        cnts.scatter_(1, topk_idx, 1)
+        tokens_per_expert = cnts.sum(dim=0)
+        idxs = topk_idx.view(-1).argsort()
+        sorted_tokens = hidden_states[idxs // topk_idx.shape[1]]
+
+        tokens_per_expert_list = tokens_per_expert.cpu().numpy()
+
+        outputs = []
+        start_idx = 0
+        for i, num_tokens in enumerate(tokens_per_expert_list):
+            end_idx = start_idx + num_tokens
+            if num_tokens == 0:
+                continue
+            expert = self.experts[i]
+            tokens_for_expert = sorted_tokens[start_idx:end_idx]
+            expert_out = expert(tokens_for_expert)
+            outputs.append(expert_out)
+            start_idx = end_idx
+
+        outs = torch.cat(outputs, dim=0) if outputs else sorted_tokens.new_empty(0)
+
+        new_x = torch.empty_like(outs)
+        new_x[idxs] = outs
+        final_out = (
+            new_x.view(*topk_idx.shape, -1)
+            .type(topk_weight.dtype)
+            .mul_(topk_weight.unsqueeze(dim=-1))
+            .sum(dim=1)
+            .type(new_x.dtype)
+        )
+        return final_out
+
+
+class KimiDecoderLayer(nn.Module):
+    def __init__(self, config: KimiLinearConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.config = config
+        if config.is_kda_layer(layer_idx):
+            self.is_linear_attn = True
+            self.self_attn = KimiDeltaAttention(config=config, layer_idx=layer_idx)
+        elif config.is_mla:
+            self.is_linear_attn = False
+            self.self_attn = KimiMLAAttention(config=config, layer_idx=layer_idx)
+        else:
+            raise NotImplementedError
+        if (
+            config.num_experts is not None
+            and layer_idx >= config.first_k_dense_replace
+            and layer_idx % getattr(config, "moe_layer_freq", 1) == 0
+        ):
+            self.block_sparse_moe = KimiSparseMoeBlock(config)
+        else:
+            self.mlp = KimiMLP(config)
+        self.input_layernorm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = KimiRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        if self.is_linear_attn is False:
+            hidden_states = self.self_attn(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                **kwargs,
+            )
+        else:
+            hidden_states = self.self_attn(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                cache_params=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                **kwargs,
+            )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        if hasattr(self, "block_sparse_moe"):
+            hidden_states = self.block_sparse_moe(hidden_states)
+        else:
+            hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class KimiPreTrainedModel(PreTrainedModel):
+    config_class = KimiLinearConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["KimiDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(KimiMoEGate, index=0),
+        "hidden_states": KimiDecoderLayer,
+        "attentions": KimiMLAAttention,
+    }
+    _is_stateful = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class KimiLinearModel(KimiPreTrainedModel):
+    def __init__(self, config: KimiLinearConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                KimiDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        if getattr(config, "_attn_implementation", None) is not None:
+            if config._attn_implementation != "flash_attention_2":
+                logger.warning_once(
+                    f"Ignoring the provided attention implementation {config._attn_implementation}"
+                )
+                logger.warning_once("Using flash_attention_2 backend instead.")
+                config._attn_implementation = "flash_attention_2"
+        else:
+            config._attn_implementation = "flash_attention_2"
+
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def _update_linear_attn_mask(self, attention_mask, cache_position):
+        """
+        NOTE: Left-padding is used for linear attention mask.
+        No need for zeroing states when
+            1. Cached forward
+            2. Attending to all inputs
+        """
+        linear_attn_mask = attention_mask
+        if cache_position[0] > 0 or (
+            attention_mask is not None and torch.all(attention_mask == 1)
+        ):
+            linear_attn_mask = None
+        return linear_attn_mask
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) and (inputs_embeds is None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        # Get inputs_embeds
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = KimiDynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position: torch.Tensor = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+        linear_attn_mask = self._update_linear_attn_mask(attention_mask, cache_position)
+
+        hidden_states = inputs_embeds
+        if past_key_values is not None:
+            assert isinstance(past_key_values, KimiDynamicCache)
+
+        for decoder_layer in self.layers:
+            layer_mask = (
+                linear_attn_mask if decoder_layer.is_linear_attn else causal_mask
+            )
+
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=layer_mask,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class KimiLinearForCausalLM(KimiPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = KimiLinearModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        generation_mode: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, KimiLinearForCausalLM
+
+        >>> model = KimiLinearForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        logits = outputs[0]
+        if generation_mode:
+            logits = logits[:, -1:]
+        logits = self.lm_head(logits)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+        aux_loss = None
+        if kwargs.get("output_router_logits", False):
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                num_experts=self.config.num_experts,
+                top_k=self.config.num_experts_per_token,
+                attention_mask=attention_mask,
+            )
+            if loss is not None:
+                loss = loss + self.config.router_aux_loss_coef * aux_loss
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
diff --git a/src/axolotl/monkeypatch/models/kimi_linear/patch_kimi_linear.py b/src/axolotl/monkeypatch/models/kimi_linear/patch_kimi_linear.py
new file mode 100644
index 0000000000..f9d1546d6a
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/kimi_linear/patch_kimi_linear.py
@@ -0,0 +1,85 @@
+import importlib.resources
+import importlib.util
+import sys
+from pathlib import Path
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+KIMI_PATCH_PACKAGE = "axolotl.monkeypatch.models.kimi_linear"
+
+
+def get_patch_file_path(package_dot_path: str, filename: str) -> Path:
+    """
+    Gets the absolute path to a patch file using importlib.resources.files.
+    """
+    try:
+        return importlib.resources.files(package_dot_path) / filename
+    except ModuleNotFoundError:
+        return None
+
+
+def _load_local_module(module_name: str, filename: str):
+    """Helper to load a local module if not already loaded."""
+    if module_name in sys.modules:
+        return sys.modules[module_name]
+
+    patch_path = get_patch_file_path(KIMI_PATCH_PACKAGE, filename)
+    if patch_path and patch_path.exists():
+        spec = importlib.util.spec_from_file_location(module_name, patch_path)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+        return module
+    return None
+
+
+def _patch_get_class_in_module():
+    """
+    Core patch function that hijacks Transformers' dynamic module loading.
+    """
+    from transformers.dynamic_module_utils import get_class_in_module
+
+    if hasattr(get_class_in_module, "_axolotl_patched"):
+        return
+
+    original_get_class_in_module = get_class_in_module
+
+    # Mapping of module path patterns to (module_name, filename)
+    KIMI_MODULE_MAP = {
+        "configuration_kimi": ("configuration_kimi", "configuration_kimi.py"),
+        "modeling_kimi": ("modeling_kimi", "modeling_kimi.py"),
+        "tokenization_kimi": ("tokenization_kimi", "tokenization_kimi.py"),
+    }
+
+    def patched_get_class_in_module(class_name, module_path, **kwargs):
+        """Patched version that returns our local modules instead of remote ones."""
+        for pattern, (module_name, filename) in KIMI_MODULE_MAP.items():
+            if pattern in module_path:
+                module = _load_local_module(module_name, filename)
+                if module:
+                    return getattr(module, class_name)
+                break  # Pattern matched but file not found, fall through
+
+        return original_get_class_in_module(class_name, module_path, **kwargs)
+
+    import transformers.dynamic_module_utils
+
+    transformers.dynamic_module_utils.get_class_in_module = patched_get_class_in_module
+    patched_get_class_in_module._axolotl_patched = True
+
+
+def patch_kimi():
+    """
+    Apply all Kimi patches.
+    Must be called BEFORE loading config/tokenizer/model.
+    """
+    _patch_get_class_in_module()
+    LOG.info("Kimi patches applied successfully!")
+
+
+# Keep these for backward compatibility if needed
+patch_kimi_config = patch_kimi
+patch_kimi_tokenizer = patch_kimi
+patch_kimi_model = patch_kimi
diff --git a/src/axolotl/monkeypatch/models/kimi_linear/tokenization_kimi.py b/src/axolotl/monkeypatch/models/kimi_linear/tokenization_kimi.py
new file mode 100644
index 0000000000..83f7ab4ae2
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/kimi_linear/tokenization_kimi.py
@@ -0,0 +1,357 @@
+"""
+Adapted Kimi-Linear tokenizer to use proper template defaults and misc fixes.
+
+Source: https://huggingface.co/moonshotai/Kimi-Linear-48B-A3B-Instruct/blob/main/tokenization_kimi.py
+Revision: 919416f
+"""
+
+import os
+from logging import getLogger
+from pathlib import Path
+from shutil import copyfile
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
+
+import tiktoken
+from tiktoken.load import load_tiktoken_bpe
+from tokenizers import AddedToken
+from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+logger = getLogger(__name__)
+VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"}
+
+
+class TikTokenTokenizer(PreTrainedTokenizer):
+    """
+    Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py.
+
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            The path to the Tiktoken model file.
+        bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|begin_of_text|>",`):
+            The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+        eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|end_of_text|>"`):
+            The end of sequence token.
+        unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_249|>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead. The second to last item in special_tokens.
+        pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_250|>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        additional_special_tokens (list of `str`, *optional*):
+            A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be
+            skipped when decoding if `skip_special_tokens` is set to `True`.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    special_tokens: Dict[str, int]
+
+    num_reserved_special_tokens = 256
+
+    pat_str = "|".join(
+        [
+            r"""[\p{Han}]+""",
+            r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
+            r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
+            r"""\p{N}{1,3}""",
+            r""" ?[^\s\p{L}\p{N}]+[\r\n]*""",
+            r"""\s*[\r\n]+""",
+            r"""\s+(?!\S)""",
+            r"""\s+""",
+        ]
+    )
+
+    def __init__(
+        self,
+        vocab_file,
+        bos_token: Union[str, AddedToken] = "[BOS]",  # nosec: B107
+        eos_token: Union[str, AddedToken] = "[EOS]",  # nosec: B107
+        unk_token: Union[str, AddedToken, None] = None,
+        pad_token: Union[str, AddedToken, None] = None,
+        additional_special_tokens: List[str] = None,
+        added_tokens_decoder: Optional[dict] = None,
+        **kwargs,
+    ):
+        assert os.path.isfile(vocab_file), vocab_file
+
+        if additional_special_tokens is None:
+            additional_special_tokens = [
+                "<|im_end|>",
+                "<|im_user|>",
+                "<|im_assistant|>",
+                "<|start_header_id|>",
+                "<|end_header_id|>",
+                "[EOT]",
+                "<|im_system|>",
+                "<|im_middle|>",
+            ]
+
+        special_tokens_mapping = {
+            i: added_tokens_decoder[i].content for i in added_tokens_decoder
+        }
+
+        self.vocab_file = vocab_file
+        mergeable_ranks = load_tiktoken_bpe(vocab_file)
+        num_base_tokens = len(mergeable_ranks)
+        self.special_tokens = {
+            special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i
+            for i in range(
+                num_base_tokens, num_base_tokens + self.num_reserved_special_tokens + 2
+            )
+        }
+
+        self.model = tiktoken.Encoding(
+            name=Path(vocab_file).name,
+            pat_str=self.pat_str,
+            mergeable_ranks=mergeable_ranks,
+            special_tokens=self.special_tokens,
+        )
+        logger.info(f"Reloaded tiktoken model from {vocab_file}")
+
+        self.n_words: int = self.model.n_vocab
+        # BOS / EOS token IDs
+        self.bos_id: int = self.special_tokens[str(bos_token)]
+        self.eos_id: int = self.special_tokens[str(eos_token)]
+        logger.info(
+            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
+        )
+
+        self.pad_id: int = self.special_tokens[str(pad_token)]
+        self.unk_id: int = self.special_tokens[str(unk_token)]
+
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+
+        self.decoder = {}
+        for i in range(self.n_words):
+            # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee
+            decoding = "".join(
+                [
+                    self.byte_encoder[ord(char)]
+                    for char in self.model.decode_single_token_bytes(i).decode(
+                        "latin-1"
+                    )
+                ]
+            )
+            self.decoder[i] = decoding
+
+        self.encoder = {}
+        for i in range(self.n_words):
+            if i in self.decoder:
+                self.encoder[self.decoder[i]] = i
+
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            additional_special_tokens=additional_special_tokens,
+            **kwargs,
+        )
+        self.all_special_ids_set = set(self.all_special_ids)
+
+    def encode(
+        self, text: str, allow_special_tokens: bool = True, **kwargs
+    ) -> List[int]:
+        """
+        Encodes a string into a list of token IDs.
+
+        Args:
+            text (str): The input string to be encoded.
+
+        Returns:
+            list[int]: A list of token IDs.
+        """
+        # If there are other args, we should call super().encode because there are a lot of code
+        # to handle those args. supper().encode finally will call _tokenize and _convert_token_to_id.
+        # NOTE: our encode method is not compatible with the super().encode method,
+        #   e.g. split_special_tokens' default is True in our encode method.
+        if len(kwargs) > 0:
+            # logger.warning(f"Calling super().encode with {kwargs}")
+            return super().encode(text, **kwargs)
+
+        assert type(text) is str
+
+        # The tiktoken tokenizer can handle <=400k chars without
+        # pyo3_runtime.PanicException.
+        TIKTOKEN_MAX_ENCODE_CHARS = 400_000
+
+        # https://github.com/openai/tiktoken/issues/195
+        # Here we iterate over subsequences and split if we exceed the limit
+        # of max consecutive non-whitespace or whitespace characters.
+        MAX_NO_WHITESPACES_CHARS = 25_000
+
+        texts = self.pre_tokenizer_process(text)
+
+        all_substrs = []
+        for text in texts:
+            substrs = (
+                substr
+                for i in range(0, len(text), TIKTOKEN_MAX_ENCODE_CHARS)
+                for substr in self._split_whitespaces_or_nonwhitespaces(
+                    text[i : i + TIKTOKEN_MAX_ENCODE_CHARS], MAX_NO_WHITESPACES_CHARS
+                )
+            )
+            all_substrs.extend(substrs)
+
+        t: List[int] = []
+        for substr in all_substrs:
+            if allow_special_tokens:
+                t.extend(
+                    # we should consider special token as a common token
+                    self.model.encode(
+                        substr,
+                        allowed_special="all",
+                    )
+                )
+            else:
+                t.extend(
+                    # we should consider special token as a common token
+                    self.model.encode(
+                        substr,
+                        disallowed_special=(),
+                    )
+                )
+
+        return t
+
+    def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str:
+        """
+        Decodes a list of token IDs into a string.
+
+        Args:
+            token_ids (List[int]): The list of token IDs to be decoded.
+
+        Returns:
+            str: The decoded string.
+        """
+        # If there are other args, we should call super().decode because there are a lot of code
+        # to handle those args. supper().encode finally will call convert_tokens_to_string and _convert_id_to_token.
+        if len(kwargs) > 0:
+            return super().decode(token_ids, **kwargs)
+
+        if type(token_ids) is int:
+            token_ids = [token_ids]
+
+        return self.model.decode(cast(List[int], token_ids))
+
+    @staticmethod
+    def _split_whitespaces_or_nonwhitespaces(
+        s: str, max_consecutive_slice_len: int
+    ) -> Iterator[str]:
+        """
+        Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len`
+        consecutive whitespaces or consecutive non-whitespaces.
+        """
+        current_slice_len = 0
+        current_slice_is_space = s[0].isspace() if len(s) > 0 else False
+        slice_start = 0
+
+        for i in range(len(s)):
+            is_now_space = s[i].isspace()
+
+            if current_slice_is_space ^ is_now_space:
+                current_slice_len = 1
+                current_slice_is_space = is_now_space
+            else:
+                current_slice_len += 1
+                if current_slice_len > max_consecutive_slice_len:
+                    yield s[slice_start:i]
+                    slice_start = i
+                    current_slice_len = 1
+        yield s[slice_start:]
+
+    def pre_tokenizer_process(self, text: str) -> List[str]:
+        """
+        pre-tokenizes the input text into a list of tokens.
+        This method is used to split the input text into smaller chunks for internal processing.
+        """
+        return [text]
+
+    """ ----- Below are the abstract methods required by PreTrainedTokenizer ----- """
+
+    @property
+    def vocab_size(self) -> int:
+        return self.n_words
+
+    def get_vocab(self) -> Dict[str, int]:
+        return self.encoder
+
+    def _tokenize(self, text: str, **kwargs) -> List[str]:
+        return [self.decoder[t] for t in self.encode(text)]
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.encoder.get(token, self.unk_id)
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.decoder.get(index)
+
+    @staticmethod
+    def clean_up_tokenization(out_string: str) -> str:
+        return out_string
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        text = "".join(tokens)
+        text = bytearray([self.byte_decoder[c] for c in text]).decode(
+            "utf-8", "replace"
+        )
+        return text
+
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            raise ValueError(
+                f"vocabulary path ({save_directory}) should be a directory"
+            )
+        out_vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "")
+            + VOCAB_FILES_NAMES["vocab_file"],
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(
+            out_vocab_file
+        ) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+
+        return (out_vocab_file,)
+
+    def apply_chat_template(
+        self,
+        conversation,
+        tools: Optional[list[dict]] = None,
+        tokenize: bool = True,
+        add_generation_prompt: bool = False,
+        **kwargs,
+    ):
+        tools = deep_sort_dict(tools)
+        return super().apply_chat_template(
+            conversation,
+            tools=tools,
+            tokenize=tokenize,
+            add_generation_prompt=add_generation_prompt,
+            **kwargs,
+        )
+
+
+def deep_sort_dict(obj: Any) -> Any:
+    if isinstance(obj, dict):
+        return {k: deep_sort_dict(v) for k, v in sorted(obj.items())}
+    if isinstance(obj, list):
+        return [deep_sort_dict(item) for item in obj]
+    return obj
diff --git a/src/axolotl/monkeypatch/models/llama4/__init__.py b/src/axolotl/monkeypatch/models/llama4/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/llama4/modeling.py b/src/axolotl/monkeypatch/models/llama4/modeling.py
new file mode 100644
index 0000000000..0fc8f5699c
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/llama4/modeling.py
@@ -0,0 +1,106 @@
+"""
+Modified Llama-4 text experts modeling for linearized experts for improved LoRA support
+"""
+
+import sys
+
+import torch
+from torch import nn
+from transformers import Llama4Config
+from transformers.activations import ACT2FN
+
+
+class Llama4TextExperts(nn.Module):
+    """
+    Modified Llama-4 text experts modeling for linearized experts
+    """
+
+    def __init__(self, config: Llama4Config):
+        super().__init__()
+        self.num_experts = config.num_local_experts
+        self.intermediate_size = config.intermediate_size
+        self.hidden_size = config.hidden_size
+        self.expert_dim = self.intermediate_size
+
+        # Replace fused gate_up_proj with separate Linear modules
+        self.gate_projs = nn.ModuleList(
+            [
+                nn.Linear(self.hidden_size, self.expert_dim, bias=False)
+                for _ in range(self.num_experts)
+            ]
+        )
+
+        self.up_projs = nn.ModuleList(
+            [
+                nn.Linear(self.hidden_size, self.expert_dim, bias=False)
+                for _ in range(self.num_experts)
+            ]
+        )
+
+        # Replace down_proj Parameter with Linear modules
+        self.down_projs = nn.ModuleList(
+            [
+                nn.Linear(self.expert_dim, self.hidden_size, bias=False)
+                for _ in range(self.num_experts)
+            ]
+        )
+
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Forward method using separate Linear layers for each expert.
+
+        Args:
+            hidden_states (torch.Tensor): (num_experts * batch_size, hidden_size)
+                The input should be organized by expert
+
+        Returns:
+            torch.Tensor: (num_experts * batch_size, hidden_size)
+        """
+        # Reshape to separate by expert
+        hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size)
+        # batch_size_per_expert = hidden_states.size(1)
+
+        # Initialize output tensor
+        next_states = torch.zeros_like(hidden_states)
+
+        # Process each expert separately
+        for i in range(self.num_experts):
+            # Get input for this expert
+            expert_input = hidden_states[
+                i
+            ]  # Shape: (batch_size_per_expert, hidden_size)
+
+            # Apply gate and up projections
+            gate = self.gate_projs[i](
+                expert_input
+            )  # Shape: (batch_size_per_expert, expert_dim)
+            up = self.up_projs[i](
+                expert_input
+            )  # Shape: (batch_size_per_expert, expert_dim)
+
+            # Apply activation and down projection
+            next_states[i] = self.down_projs[i](up * self.act_fn(gate))
+
+        # Flatten back to original shape
+        return next_states.view(-1, self.hidden_size)
+
+
+def patch_llama4_linearized_modeling():
+    """
+    Patch Llama4TextExperts to use separate Linear layers for each expert.
+    """
+    from transformers.models.llama4 import modeling_llama4
+
+    old_lamma_4_text_experts = modeling_llama4.Llama4TextExperts
+    modeling_llama4.Llama4TextExperts = Llama4TextExperts
+    sys.modules["transformers.models.llama4"].Llama4TextExperts = Llama4TextExperts
+
+    def unpatch():
+        modeling_llama4.Llama4TextExperts = old_lamma_4_text_experts
+        sys.modules[
+            "transformers.models.llama4"
+        ].Llama4TextExperts = old_lamma_4_text_experts
+
+    return unpatch
diff --git a/src/axolotl/monkeypatch/models/mamba_utils.py b/src/axolotl/monkeypatch/models/mamba_utils.py
new file mode 100644
index 0000000000..0c0e12c5fa
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/mamba_utils.py
@@ -0,0 +1,331 @@
+"""Shared utilities for Mamba2 SSM sample-packing and context-parallelism patches.
+
+Used by: nemotron_h, falcon_h1, granite_moe_hybrid
+"""
+
+import functools
+
+import torch
+import torch.distributed as dist
+
+
+def get_seq_idx(position_ids: torch.Tensor) -> torch.Tensor:
+    """Convert position_ids [B, T] → seq_idx [B, T] int32 for mamba-ssm kernels.
+
+    Example: position_ids [[0,1,2,3,0,1,2]] → seq_idx [[0,0,0,0,1,1,1]]
+
+    Under context parallelism a rank may receive a chunk that begins mid-sample
+    (position_ids[0] != 0), so the raw cumsum starts at 0 and subtracting 1
+    would yield -1 — an invalid value for the Mamba kernels.  Subtracting the
+    first element of the cumsum instead normalises every chunk to start at 0
+    while still correctly incrementing at every intra-chunk sample boundary.
+
+    Example (CP rank 1, chunk starts mid-sample):
+        position_ids [[3,4,5,0,1,2]] → seq_idx [[0,0,0,1,1,1]]
+    """
+    cumsum = torch.cumsum((position_ids == 0).int(), dim=-1)
+    return (cumsum - cumsum[..., :1]).to(torch.int32)
+
+
+def is_cp_active() -> bool:
+    """Return True if context parallelism (ring attention) is active on this rank.
+
+    Zero-cost when CP is not configured: the import guard ensures we only touch
+    the distributed group if ring_flash_attn is installed.
+    """
+    try:
+        from axolotl.monkeypatch.ring_attn import get_ring_attn_group
+
+        group = get_ring_attn_group()
+        return group is not None and dist.get_world_size(group) > 1
+    except (ImportError, RuntimeError):
+        return False
+
+
+def _get_cp_group_and_rank():
+    """Return (process_group, local_rank, world_size) for the CP ring."""
+    from axolotl.monkeypatch.ring_attn import get_ring_attn_group
+
+    group = get_ring_attn_group()
+    return group, dist.get_rank(group), dist.get_world_size(group)
+
+
+def ring_shift_ssm_state(
+    h_final: torch.Tensor,
+) -> torch.Tensor:
+    """P2P ring: send h_final to rank+1, receive from rank-1 within CP group.
+
+    Uses synchronous send/recv on the ring attention process group.
+    Rank 0 in the CP group receives zeros (no previous chunk).
+
+    Args:
+        h_final: Final SSM state from this rank's forward pass.
+                 Shape is architecture-dependent, typically [B, H, d, n].
+
+    Returns:
+        h_prev: SSM state received from rank-1, same shape/dtype as h_final.
+                Zero tensor on the first rank in the CP group.
+    """
+    group, local_rank, world_size = _get_cp_group_and_rank()
+    ranks = dist.get_process_group_ranks(group)
+
+    h_prev = torch.zeros_like(h_final)
+
+    if world_size <= 1:
+        return h_prev
+
+    prev_global = ranks[(local_rank - 1) % world_size]
+    next_global = ranks[(local_rank + 1) % world_size]
+
+    send_op = dist.P2POp(dist.isend, h_final.contiguous(), next_global, group=group)
+    recv_op = dist.P2POp(dist.irecv, h_prev, prev_global, group=group)
+
+    reqs = dist.batch_isend_irecv([send_op, recv_op])
+    for req in reqs:
+        req.wait()
+
+    # Rank 0 in the ring has no true predecessor — zero out received state
+    if local_rank == 0:
+        h_prev.zero_()
+
+    return h_prev
+
+
+def mamba2_cp_correction(
+    out: torch.Tensor,
+    h_final: torch.Tensor,
+    C: torch.Tensor,
+    cum_A: torch.Tensor,
+    h_prev: torch.Tensor,
+    num_heads: int,
+    head_dim: int,
+    seq_idx: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply CP correction to SSM output using the received state from rank-1.
+
+    SSM output is linear in the initial hidden state, so the contribution of
+    h_prev can be added analytically without a second forward pass.
+
+    For each timestep t in the local chunk:
+        propagated_state_t = cumA_t * h_prev          [B, H, d, n]
+        Δy_t = sum_over_n( C_t * propagated_state_t ) [B, H, d]
+
+    The corrected final state for this rank is:
+        h_final_corrected = h_final + cumA_T * h_prev
+
+    Sample packing correctness (seq_idx):
+        When sample packing is active, a CP rank may hold multiple packed
+        sequences. Only the first sequence (seq_idx == 0) is a continuation
+        of the previous rank's chunk — subsequent sequences are brand-new and
+        should receive zero correction from h_prev.
+
+        Passing seq_idx masks delta_y to zero for all tokens where
+        seq_idx > 0, preventing h_prev state from leaking into unrelated
+        packed sequences.
+
+    Args:
+        out:       SSM scan output from this rank, shape [B, T, D] where D = H*d.
+        h_final:   Final SSM state from this rank, shape [B, H, d, n].
+        C:         Output projection matrices, shape [B, T, n_groups, n].
+        cum_A:     Cumulative log-transition factors, shape [B, T, H].
+                   These are the log-space cumulative sums of A, so
+                   exp(cum_A_t) gives the transition matrix from step 0 to t.
+        h_prev:    SSM state received from rank-1 (zeros on rank 0).
+                   Shape [B, H, d, n].
+        num_heads: Number of SSM heads (H).
+        head_dim:  Dimension per head (d).
+        seq_idx:   Optional sequence index tensor, shape [B, T] int32.
+                   When provided, correction is zeroed for tokens where
+                   seq_idx > 0 (i.e. sequences that start fresh on this rank).
+
+    Returns:
+        corrected_out:     out + Δy, shape [B, T, D].
+        corrected_h_final: h_final + cumA_T * h_prev, shape [B, H, d, n].
+    """
+    if not h_prev.any():
+        return out, h_final
+
+    B, T, _ = out.shape
+    n_groups = C.shape[2]
+    heads_per_group = num_heads // n_groups
+
+    # cum_A: [B, T, H] → transition factors (exponentiate from log-space)
+    decay = torch.exp(cum_A).float()  # [B, T, H]
+
+    # Propagate h_prev through cumulative transitions: [B, T, H, d, n]
+    prop_state = decay[:, :, :, None, None] * h_prev[:, None, :, :, :].float()
+
+    # C: [B, T, n_groups, n] → expand to heads: [B, T, H, n]
+    C_expanded = C.float().repeat_interleave(heads_per_group, dim=2)  # [B, T, H, n]
+
+    # Δy_t = sum_n(C_t * prop_state_t) → [B, T, H, d]
+    delta_y = torch.einsum("bthn,bthdn->bthd", C_expanded, prop_state)
+
+    # Mask out correction for tokens belonging to new sequences on this rank.
+    # seq_idx == 0 → continuation of the sequence that crossed the CP boundary
+    # seq_idx  > 0 → brand-new packed sequence, h_prev is irrelevant to it
+    if seq_idx is not None:
+        # mask: [B, T, 1, 1] — broadcast over H and d
+        mask = (seq_idx == 0).to(delta_y.dtype).unsqueeze(-1).unsqueeze(-1)
+        delta_y = delta_y * mask
+
+    # Reshape to [B, T, D] where D = H * d
+    delta_y = delta_y.reshape(B, T, num_heads * head_dim).to(out.dtype)
+
+    corrected_out = out + delta_y
+
+    # Correct final state using last-timestep decay.
+    # If the last token is in a new sequence (seq_idx > 0 at T-1), h_prev
+    # should not propagate into h_final either.
+    if seq_idx is not None and seq_idx[:, -1].any():
+        # last token belongs to a new sequence — don't corrupt h_final
+        corrected_h_final = h_final
+    else:
+        decay_final = decay[:, -1, :, None, None]  # [B, H, 1, 1]
+        corrected_h_final = h_final + (decay_final * h_prev.float()).to(h_final.dtype)
+
+    return corrected_out, corrected_h_final
+
+
+def ensure_mamba_kernels_loaded(target_module):
+    """Eagerly resolve mamba-ssm and causal-conv1d globals on *target_module*.
+
+    Transformers >= 5.5 lazily loads these inside ``Mixer.__init__`` via
+    ``lazy_load_kernel``.  Our monkeypatches run *before* model instantiation,
+    so the module globals are still ``None``.  This helper triggers the kernel
+    resolution early so the patched ``cuda_kernels_forward`` (and
+    ``wrap_mamba_scan_for_cp``) can reference them.
+    """
+    if getattr(target_module, "mamba_chunk_scan_combined", None) is not None:
+        return
+
+    try:
+        from transformers.integrations.hub_kernels import lazy_load_kernel
+        from transformers.utils.import_utils import resolve_internal_import
+    except ImportError:
+        return
+
+    causal_conv1d = lazy_load_kernel("causal-conv1d")
+    if causal_conv1d is not None:
+        target_module.causal_conv1d_update = getattr(
+            causal_conv1d, "causal_conv1d_update", None
+        )
+        target_module.causal_conv1d_fn = getattr(
+            causal_conv1d, "causal_conv1d_fn", None
+        )
+
+    mamba_ssm = lazy_load_kernel("mamba-ssm")
+    if mamba_ssm is not None:
+        target_module.selective_state_update = resolve_internal_import(
+            mamba_ssm,
+            chained_path="ops.triton.selective_state_update.selective_state_update",
+        )
+        target_module.mamba_chunk_scan_combined = resolve_internal_import(
+            mamba_ssm,
+            chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined",
+        )
+        target_module.mamba_split_conv1d_scan_combined = resolve_internal_import(
+            mamba_ssm,
+            chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined",
+        )
+
+    target_module.is_fast_path_available = all(
+        (
+            getattr(target_module, "selective_state_update", None),
+            getattr(target_module, "mamba_chunk_scan_combined", None),
+            getattr(target_module, "mamba_split_conv1d_scan_combined", None),
+            getattr(target_module, "causal_conv1d_fn", None),
+            getattr(target_module, "causal_conv1d_update", None),
+        )
+    )
+
+
+def wrap_mamba_scan_for_cp(target_module):
+    """Wrap ``mamba_chunk_scan_combined`` in *target_module* to apply CP correction.
+
+    After the scan, if CP is active the wrapper:
+    1. Sends the final SSM state to the next rank via ``ring_shift_ssm_state``.
+    2. Computes cumA from the scan's A / dt / dt_bias / dt_softplus args.
+    3. Calls ``mamba2_cp_correction`` to add the contribution of h_prev.
+
+    This is installed per-module so it only affects the architecture whose
+    modeling file imports ``mamba_chunk_scan_combined``.
+
+    The approach follows Tri Dao's Mamba-2 systems blog: each GPU computes its
+    local output and final states, states are passed via P2P, then outputs are
+    corrected — no ring attention needed for SSM layers.
+    """
+    if getattr(target_module, "_cp_scan_wrapped", False):
+        return
+
+    ensure_mamba_kernels_loaded(target_module)
+
+    if getattr(target_module, "mamba_chunk_scan_combined", None) is None:
+        return
+
+    original_scan = target_module.mamba_chunk_scan_combined
+
+    @functools.wraps(original_scan)
+    def _cp_scan_wrapper(*args, **kwargs):
+        cp_active = is_cp_active()
+
+        if cp_active:
+            kwargs["return_final_states"] = True
+
+        result = original_scan(*args, **kwargs)
+
+        if not cp_active:
+            return result
+
+        scan_output, ssm_state = result
+        if ssm_state is None:
+            return result
+
+        h_prev = ring_shift_ssm_state(ssm_state)
+
+        # Signature: mamba_chunk_scan_combined(x, dt, A, B, C, ...)
+        # Extract from kwargs first, fall back to positional args.
+        dt_arg = kwargs.get("dt", args[1] if len(args) > 1 else None)
+        A_arg = kwargs.get("A", args[2] if len(args) > 2 else None)
+        C_arg = kwargs.get("C", args[4] if len(args) > 4 else None)
+        if dt_arg is None or A_arg is None or C_arg is None:
+            raise ValueError(
+                "wrap_mamba_scan_for_cp requires dt, A, C to be passed "
+                f"positionally (got {len(args)} positional args) or as kwargs."
+            )
+        dt_bias = kwargs.get("dt_bias")
+        dt_softplus = kwargs.get("dt_softplus", False)
+        seq_idx = kwargs.get("seq_idx")
+
+        if dt_softplus:
+            dt_eff = torch.nn.functional.softplus(
+                dt_arg + (dt_bias if dt_bias is not None else 0)
+            )
+        else:
+            dt_eff = dt_arg
+
+        dA = A_arg[None, None, :] * dt_eff
+        cum_A = torch.cumsum(dA, dim=1)
+
+        x = args[0]
+        num_heads = A_arg.shape[0]
+        head_dim = x.shape[3] if x.ndim == 4 else x.shape[2] // num_heads
+        B_dim, T_dim = x.shape[0], x.shape[1]
+
+        scan_flat = scan_output.view(B_dim, T_dim, -1)
+        scan_flat, ssm_state = mamba2_cp_correction(
+            scan_flat,
+            ssm_state,
+            C_arg,
+            cum_A,
+            h_prev,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            seq_idx=seq_idx,
+        )
+        scan_output = scan_flat.view(scan_output.shape)
+
+        return scan_output, ssm_state
+
+    target_module.mamba_chunk_scan_combined = _cp_scan_wrapper
+    target_module._cp_scan_wrapped = True
diff --git a/src/axolotl/monkeypatch/models/mistral3/__init__.py b/src/axolotl/monkeypatch/models/mistral3/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/mistral3/mistral_common_tokenizer.py b/src/axolotl/monkeypatch/models/mistral3/mistral_common_tokenizer.py
new file mode 100644
index 0000000000..a77a0129eb
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/mistral3/mistral_common_tokenizer.py
@@ -0,0 +1,85 @@
+"""
+Monkeypatch to fix inefficient tensor conversion in MistralCommonBackend.apply_chat_template
+"""
+
+import importlib
+import inspect
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def apply_mistral_tokenizer_image_patch():
+    """Apply patch to MistralCommonBackend.apply_chat_template to fix image tensor conversion."""
+    from transformers.tokenization_mistral_common import MistralCommonBackend
+
+    # Get original source
+    original_source = inspect.getsource(MistralCommonBackend.apply_chat_template)
+    original_source, _ = detab_code(original_source)
+
+    # Define the replacement
+    original_tensor_conversion = (
+        "                    pixel_values = torch.tensor(images)"
+    )
+
+    patched_tensor_conversion = """                    if isinstance(images, list) and len(images) > 0 and isinstance(images[0], np.ndarray):
+                        pixel_values = torch.tensor(np.array(images))
+                    else:
+                        pixel_values = torch.tensor(images)"""
+
+    # Apply the replacement
+    if original_tensor_conversion in original_source:
+        patched_source = original_source.replace(
+            original_tensor_conversion, patched_tensor_conversion
+        )
+        patched_source = patched_source.replace(
+            "def apply_chat_template(",
+            "def patched_apply_chat_template(",
+            1,
+        )
+
+        # Load necessary imports from the module
+        module_name = MistralCommonBackend.__module__
+        module = importlib.import_module(module_name)
+
+        # Detect what needs to be imported
+        items_to_import = []
+        for item in dir(module):
+            if item in patched_source and not item.startswith("_"):
+                items_to_import.append(item)
+
+        # Execute imports in global scope
+        if items_to_import:
+            exec(  # nosec B102
+                f"from {module_name} import ({', '.join(items_to_import)})",
+                globals(),
+            )
+
+        # Also need standard imports that might be used
+        exec("import numpy as np", globals())  # nosec B102
+        exec("import torch", globals())  # nosec B102
+        exec("from typing import Union, Optional, List, Dict, Any, Callable", globals())  # nosec B102
+        exec("from pathlib import Path", globals())  # nosec B102
+
+        # Import other dependencies that might be needed
+        try:
+            exec("from transformers.utils import is_torch_available", globals())  # nosec B102
+            exec(
+                "from transformers.tokenization_utils_base import BatchEncoding, PaddingStrategy, TensorType",
+                globals(),
+            )  # nosec B102
+            exec("from transformers.utils import logging", globals())  # nosec B102
+            exec("logger = logging.get_logger(__name__)", globals())  # nosec B102
+        except ImportError as e:
+            LOG.warning(f"Could not import some dependencies: {e}")
+
+        # Execute the patched source
+        exec(patched_source, globals())  # nosec B102
+
+        # Replace the method
+        MistralCommonBackend.apply_chat_template = patched_apply_chat_template
+        LOG.info("Successfully applied MistralCommonBackend tensor conversion patch")
+    else:
+        LOG.warning("Could not find target code for MistralCommonBackend patching")
diff --git a/src/axolotl/monkeypatch/models/nemotron_h/__init__.py b/src/axolotl/monkeypatch/models/nemotron_h/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/nemotron_h/modeling.py b/src/axolotl/monkeypatch/models/nemotron_h/modeling.py
new file mode 100644
index 0000000000..91d26fa0c5
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/nemotron_h/modeling.py
@@ -0,0 +1,323 @@
+"""Sample-packing and context-parallelism patch for NemotronH (Mamba2/Attention/MoE hybrid).
+
+Threads seq_idx (derived from position_ids) into the Mamba2 SSM kernels so
+packed-sequence boundaries reset SSM state. Upstream hard-codes seq_idx=None,
+which leaks hidden state across boundaries. Attention and MoE blocks need no
+changes — only the Mamba2 mixer is patched.
+
+CP correction (ring-shift of SSM state + additive output fix) is handled by
+``wrap_mamba_scan_for_cp`` from ``mamba_utils``, which wraps the
+``mamba_chunk_scan_combined`` call at the module level.
+"""
+
+import importlib
+
+import torch
+
+from axolotl.monkeypatch.models.mamba_utils import (
+    ensure_mamba_kernels_loaded,
+    get_seq_idx,
+    is_cp_active,
+    wrap_mamba_scan_for_cp,
+)
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_nemotron_h_modeling_packing():
+    """Patch NemotronH for sample packing: seq_idx threading into Mamba2 SSM kernels.
+
+    _get_unpad_data is handled by SUPPORTED_MULTIPACK_MODEL_TYPES / patch_for_multipack().
+    This function only applies the seq_idx patches that are unique to nemotron_h.
+    """
+    try:
+        mod = importlib.import_module(
+            "transformers.models.nemotron_h.modeling_nemotron_h"
+        )
+    except ImportError:
+        LOG.warning("nemotron_h not found in transformers, skipping packing patches")
+        return
+
+    ensure_mamba_kernels_loaded(mod)
+
+    NemotronHMamba2Mixer = mod.NemotronHMamba2Mixer
+    NemotronHBlock = mod.NemotronHBlock
+
+    # Patch 1: cuda_kernels_forward — add seq_idx param and thread it to
+    # causal_conv1d_fn and mamba_chunk_scan_combined. Fused fast path is
+    # bypassed when seq_idx is set (requires causal_conv1d_cuda C extension).
+    def patched_cuda_kernels_forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params=None,
+        attention_mask=None,
+        seq_idx=None,
+    ):
+        batch_size, seq_len, _ = hidden_states.shape
+        groups_time_state_size = self.n_groups * self.ssm_state_size
+        d_to_remove = (
+            2 * self.intermediate_size
+            + 2 * self.n_groups * self.ssm_state_size
+            + self.num_heads
+        )
+
+        if cache_params is not None and cache_params.has_previous_state:
+            in_projected_states = self.in_proj(hidden_states.squeeze(1))
+            d_mlp = (in_projected_states.shape[-1] - d_to_remove) // 2
+            split_projection_dim = [
+                d_mlp,
+                d_mlp,
+                self.intermediate_size,
+                self.conv_dim,
+                self.num_heads,
+            ]
+            _, _, gate, hidden_states_B_C, dt = torch.split(
+                in_projected_states, split_projection_dim, dim=-1
+            )
+            hidden_states_B_C = mod.causal_conv1d_update(
+                hidden_states_B_C,
+                cache_params.conv_states[self.layer_idx],
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                self.activation,
+            )
+            hidden_states, B, C = torch.split(
+                hidden_states_B_C,
+                [
+                    self.intermediate_size,
+                    groups_time_state_size,
+                    groups_time_state_size,
+                ],
+                dim=-1,
+            )
+            A = -torch.exp(self.A_log.float())
+            A = (
+                A[:, None, ...][:, :, None]
+                .expand(-1, self.head_dim, self.ssm_state_size)
+                .to(dtype=torch.float32)
+            )
+            dt = dt[:, :, None].expand(-1, -1, self.head_dim)
+            dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+            D = self.D[:, None, ...].expand(-1, self.head_dim)
+            B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)
+            C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)
+            hidden_states_reshaped = hidden_states.view(
+                batch_size, self.num_heads, self.head_dim
+            )
+            hidden_states = mod.selective_state_update(
+                cache_params.ssm_states[self.layer_idx],
+                hidden_states_reshaped,
+                dt,
+                A,
+                B,
+                C,
+                D,
+                z=None,
+                dt_bias=dt_bias,
+                dt_softplus=True,
+            )
+            hidden_states = hidden_states.view(
+                batch_size, self.num_heads * self.head_dim
+            )
+            hidden_states = self.norm(hidden_states, gate)
+            out = self.out_proj(hidden_states)[:, None, ...]
+
+        else:
+            if attention_mask is not None and not torch.all(attention_mask == 1):
+                dtype = hidden_states.dtype
+                hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+            projected_states = self.in_proj(hidden_states)
+            A = -torch.exp(self.A_log.float())
+            dt_limit_kwargs = (
+                {}
+                if self.time_step_limit is None
+                else {"dt_limit": self.time_step_limit}
+            )
+            if attention_mask is not None:
+                input_not_masked = torch.all(attention_mask == 1)
+            else:
+                input_not_masked = True
+
+            if (
+                self.use_mem_eff_path
+                and self.training
+                and cache_params is None
+                and input_not_masked
+                and not is_cp_active()
+            ):
+                out, ssm_state = mod.mamba_split_conv1d_scan_combined(
+                    projected_states,
+                    self.conv1d.weight.squeeze(1),
+                    self.conv1d.bias,
+                    self.dt_bias,
+                    A,
+                    D=self.D,
+                    chunk_size=self.chunk_size,
+                    seq_idx=seq_idx,
+                    activation=self.activation,
+                    rmsnorm_weight=self.norm.weight,
+                    rmsnorm_eps=self.norm.variance_epsilon,
+                    outproj_weight=self.out_proj.weight,
+                    outproj_bias=self.out_proj.bias,
+                    headdim=self.head_dim,
+                    ngroups=self.n_groups,
+                    norm_before_gate=False,
+                    return_final_states=True,
+                    **dt_limit_kwargs,
+                )
+            else:
+                gate, hidden_states_B_C, time_step = torch.split(
+                    projected_states,
+                    [self.intermediate_size, self.conv_dim, self.num_heads],
+                    dim=-1,
+                )
+
+                if cache_params is not None:
+                    hidden_states_B_C_t = hidden_states_B_C.transpose(1, 2)
+                    conv_state = torch.nn.functional.pad(
+                        hidden_states_B_C_t,
+                        (self.conv_kernel_size - hidden_states_B_C_t.shape[-1], 0),
+                    )
+                    cache_params.conv_states[self.layer_idx].copy_(conv_state)
+
+                if mod.causal_conv1d_fn is None or self.activation not in [
+                    "silu",
+                    "swish",
+                ]:
+                    hidden_states_B_C = self.act(
+                        self.conv1d(hidden_states_B_C.transpose(1, 2)).transpose(1, 2)[
+                            :, :seq_len
+                        ]
+                    )
+                else:
+                    hidden_states_B_C = mod.causal_conv1d_fn(
+                        x=hidden_states_B_C.transpose(1, 2),
+                        weight=self.conv1d.weight.squeeze(1),
+                        bias=self.conv1d.bias,
+                        activation=self.activation,
+                        seq_idx=seq_idx,
+                    ).transpose(1, 2)[:, :seq_len]
+
+                hidden_states, B, C = torch.split(
+                    hidden_states_B_C,
+                    [
+                        self.intermediate_size,
+                        groups_time_state_size,
+                        groups_time_state_size,
+                    ],
+                    dim=-1,
+                )
+
+                if attention_mask is not None and not torch.all(attention_mask == 1):
+                    dtype = hidden_states.dtype
+                    hidden_states = (hidden_states * attention_mask[:, :, None]).to(
+                        dtype
+                    )
+
+                C_reshaped = C.view(batch_size, seq_len, self.n_groups, -1)
+                scan_output, ssm_state = mod.mamba_chunk_scan_combined(
+                    hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+                    time_step,
+                    A,
+                    B.view(batch_size, seq_len, self.n_groups, -1),
+                    C_reshaped,
+                    chunk_size=self.chunk_size,
+                    D=self.D,
+                    z=None,
+                    seq_idx=seq_idx,
+                    return_final_states=True,
+                    dt_bias=self.dt_bias,
+                    dt_softplus=True,
+                    **dt_limit_kwargs,
+                )
+
+                if ssm_state is not None and cache_params is not None:
+                    cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
+                scan_output = scan_output.view(batch_size, seq_len, -1)
+                scan_output = self.norm(scan_output, gate)
+                out = self.out_proj(scan_output)
+
+        return out
+
+    NemotronHMamba2Mixer.cuda_kernels_forward = patched_cuda_kernels_forward
+
+    # Patch 2: Mamba2Mixer.forward — add seq_idx, guard on causal_conv1d_fn,
+    # restore the cuda stream context (matches upstream; avoids NaN on multi-GPU).
+    def patched_mixer_forward(
+        self,
+        hidden_states,
+        cache_params=None,
+        attention_mask=None,
+        seq_idx=None,
+    ):
+        if seq_idx is not None and mod.causal_conv1d_fn is None:
+            raise RuntimeError(
+                "Nemotron-H sample packing requires causal_conv1d_fn. "
+                "Install with: pip install mamba-ssm causal-conv1d"
+            )
+        if (
+            mod.is_fast_path_available
+            and "cuda" in self.in_proj.weight.device.type
+            and not mod.is_torchdynamo_compiling()
+        ):
+            with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)):
+                return self.cuda_kernels_forward(
+                    hidden_states, cache_params, attention_mask, seq_idx=seq_idx
+                )
+        return self.torch_forward(hidden_states, cache_params, attention_mask)
+
+    NemotronHMamba2Mixer.forward = patched_mixer_forward
+
+    # Patch 3: NemotronHBlock.forward — compute seq_idx from position_ids and
+    # pass it to the Mamba2 mixer. Skipped during decode (has_previous_state).
+    def patched_block_forward(
+        self,
+        hidden_states,
+        past_key_values=None,
+        cache_position=None,
+        attention_mask=None,
+        position_ids=None,
+        use_cache=False,
+        **kwargs,
+    ):
+        residual = hidden_states
+        hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+
+        if self.block_type == "mamba":
+            is_decoding = (
+                past_key_values is not None and past_key_values.has_previous_state
+            )
+            seq_idx = (
+                get_seq_idx(position_ids)
+                if position_ids is not None and not is_decoding
+                else None
+            )
+            hidden_states = self.mixer(
+                hidden_states,
+                cache_params=past_key_values,
+                attention_mask=attention_mask,
+                seq_idx=seq_idx,
+            )
+        elif self.block_type == "attention":
+            hidden_states, _ = self.mixer(
+                hidden_states=hidden_states,
+                past_key_values=past_key_values,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                user_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+        else:
+            hidden_states = self.mixer(hidden_states)
+
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+    NemotronHBlock.forward = patched_block_forward
+
+    wrap_mamba_scan_for_cp(mod)
+
+    LOG.info("Applied NemotronH sample packing patch (seq_idx threading into Mamba2)")
diff --git a/src/axolotl/monkeypatch/models/pixtral/__init__.py b/src/axolotl/monkeypatch/models/pixtral/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/pixtral/modeling_flash_attention_utils.py b/src/axolotl/monkeypatch/models/pixtral/modeling_flash_attention_utils.py
new file mode 100644
index 0000000000..d2b482f19b
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/pixtral/modeling_flash_attention_utils.py
@@ -0,0 +1,42 @@
+"""Monkeypatch for FA utils to accept 1D position_ids from Pixtral's position_ids_in_meshgrid"""
+
+import torch
+
+
+def apply_patch_is_packed_sequence():
+    """Apply patch to FA utils to accept 1D position_ids from Pixtral's position_ids_in_meshgrid"""
+    from transformers import modeling_flash_attention_utils
+
+    def fixed_is_packed_sequence(position_ids, batch_size):
+        """
+        Check the position ids whether packed sequences are indicated or not
+            1. Position ids exist
+            2. Flattened sequences only are supported
+            3. Compile-friendly `not (torch.diff(position_ids, dim=-1) >= 0).all()`, i.e. we have multiple increasing sequences
+        """
+        if position_ids is None:
+            return False
+
+        if position_ids.ndim == 1:
+            position_ids = position_ids.unsqueeze(0)  # [N] -> [1, N]
+
+        increasing_position_sequences = (
+            torch.arange(position_ids.shape[1], device=position_ids.device)
+            + position_ids.min()
+        )
+        return (
+            batch_size == 1
+            and (increasing_position_sequences - position_ids).abs().sum().bool().item()
+        )
+
+    # Store original method
+    old_fn = modeling_flash_attention_utils._is_packed_sequence
+
+    # Apply the patch
+    modeling_flash_attention_utils._is_packed_sequence = fixed_is_packed_sequence
+
+    def unpatch():
+        """Restore the original method"""
+        modeling_flash_attention_utils._is_packed_sequence = old_fn
+
+    return unpatch
diff --git a/src/axolotl/monkeypatch/models/qwen3/__init__.py b/src/axolotl/monkeypatch/models/qwen3/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/qwen3/fused_attn.py b/src/axolotl/monkeypatch/models/qwen3/fused_attn.py
new file mode 100644
index 0000000000..670522797f
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3/fused_attn.py
@@ -0,0 +1,117 @@
+"""Fuse ``q_norm/k_norm`` + RoPE in ``Qwen3Attention.forward`` via one Triton kernel."""
+
+from typing import Callable
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def _resolve_norm_module(norm):
+    """Unwrap PEFT ``ModulesToSaveWrapper`` so the kernel reads the active adapter's weight."""
+    modules_to_save = getattr(norm, "modules_to_save", None)
+    if not modules_to_save:
+        return norm
+    adapters = getattr(norm, "active_adapters", None)
+    if adapters is None:
+        adapter = getattr(norm, "active_adapter", None)
+        adapters = [adapter] if adapter is not None else []
+    elif isinstance(adapters, str):
+        adapters = [adapters]
+    for name in adapters:
+        if isinstance(name, str) and name in modules_to_save:
+            return modules_to_save[name]
+    return getattr(norm, "original_module", norm)
+
+
+def _make_fused_forward():
+    from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+    def fused_forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values=None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+        from transformers.models.qwen3.modeling_qwen3 import eager_attention_forward
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q_norm = _resolve_norm_module(self.q_norm)
+        k_norm = _resolve_norm_module(self.k_norm)
+        eps = q_norm.variance_epsilon
+
+        cos, sin = position_embeddings
+
+        has_lora_qkv = hasattr(self, "apply_qkv")
+        if has_lora_qkv:
+            query_states, key_states, value_states = self.apply_qkv(hidden_states)
+            query_states = query_states.view(hidden_shape)
+            key_states = key_states.view(hidden_shape)
+            value_states = value_states.view(hidden_shape)
+        else:
+            query_states = self.q_proj(hidden_states).view(hidden_shape)
+            key_states = self.k_proj(hidden_states).view(hidden_shape)
+            value_states = self.v_proj(hidden_states).view(hidden_shape)
+
+        # accelerate's per-module pre-hooks that move CPU-staged params don't fire on a monkeypatched forward.
+        q_w = q_norm.weight
+        if q_w.device != query_states.device:
+            q_w = q_w.to(query_states.device, non_blocking=True)
+        k_w = k_norm.weight
+        if k_w.device != key_states.device:
+            k_w = k_w.to(key_states.device, non_blocking=True)
+
+        query_states = fused_rms_norm_rope(
+            query_states, q_w, cos, sin, eps=eps
+        ).transpose(1, 2)
+        key_states = fused_rms_norm_rope(key_states, k_w, cos, sin, eps=eps).transpose(
+            1, 2
+        )
+        value_states = value_states.transpose(1, 2)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        if hasattr(self, "apply_o"):
+            attn_output = self.apply_o(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    return fused_forward
+
+
+def patch_qwen3_fused_attn() -> None:
+    from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+    if getattr(Qwen3Attention, "_axolotl_fused_attn_patched", False):
+        return
+
+    Qwen3Attention.forward = _make_fused_forward()
+    Qwen3Attention._axolotl_fused_attn_patched = True
+    logger.info("Patched Qwen3Attention.forward with fused RMSNorm+RoPE Triton kernel")
diff --git a/src/axolotl/monkeypatch/models/qwen3_5/__init__.py b/src/axolotl/monkeypatch/models/qwen3_5/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/qwen3_5/fused_attn.py b/src/axolotl/monkeypatch/models/qwen3_5/fused_attn.py
new file mode 100644
index 0000000000..ce83c22d2f
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_5/fused_attn.py
@@ -0,0 +1,133 @@
+"""Fused q_norm/k_norm + RoPE for Qwen3.5 (gated q_proj, ``unit_offset=True`` RMSNorm)."""
+
+from typing import Callable
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def _resolve_norm_module(norm):
+    """Unwrap PEFT ``ModulesToSaveWrapper`` so the kernel reads the active adapter's weight."""
+    modules_to_save = getattr(norm, "modules_to_save", None)
+    if not modules_to_save:
+        return norm
+    adapters = getattr(norm, "active_adapters", None)
+    if adapters is None:
+        adapter = getattr(norm, "active_adapter", None)
+        adapters = [adapter] if adapter is not None else []
+    elif isinstance(adapters, str):
+        adapters = [adapters]
+    for name in adapters:
+        if isinstance(name, str) and name in modules_to_save:
+            return modules_to_save[name]
+    return getattr(norm, "original_module", norm)
+
+
+def _make_fused_forward():
+    from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+    def fused_forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values=None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+        from transformers.models.qwen3_5.modeling_qwen3_5 import (
+            eager_attention_forward,
+        )
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q_norm = _resolve_norm_module(self.q_norm)
+        k_norm = _resolve_norm_module(self.k_norm)
+        # Liger's RMSNorm replacement uses ``variance_epsilon`` instead of ``eps``.
+        eps = getattr(q_norm, "eps", None)
+        if eps is None:
+            eps = q_norm.variance_epsilon
+
+        cos, sin = position_embeddings
+
+        has_lora_qkv = hasattr(self, "apply_qkv")
+        if has_lora_qkv:
+            query_states, key_states, value_states = self.apply_qkv(hidden_states)
+            query_states, gate = torch.chunk(
+                query_states.view(*input_shape, -1, self.head_dim * 2), 2, dim=-1
+            )
+            query_states = query_states.reshape(hidden_shape)
+            key_states = key_states.view(hidden_shape)
+            value_states = value_states.view(hidden_shape)
+        else:
+            query_states, gate = torch.chunk(
+                self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2),
+                2,
+                dim=-1,
+            )
+            query_states = query_states.reshape(hidden_shape)
+            key_states = self.k_proj(hidden_states).view(hidden_shape)
+            value_states = self.v_proj(hidden_states).view(hidden_shape)
+        gate = gate.reshape(*input_shape, -1)
+
+        # accelerate's per-module pre-hooks that move CPU-staged params don't fire on a monkeypatched forward.
+        q_w = q_norm.weight
+        if q_w.device != query_states.device:
+            q_w = q_w.to(query_states.device, non_blocking=True)
+        k_w = k_norm.weight
+        if k_w.device != key_states.device:
+            k_w = k_w.to(key_states.device, non_blocking=True)
+
+        query_states = fused_rms_norm_rope(
+            query_states, q_w, cos, sin, eps=eps, unit_offset=True
+        ).transpose(1, 2)
+        key_states = fused_rms_norm_rope(
+            key_states, k_w, cos, sin, eps=eps, unit_offset=True
+        ).transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = attn_output * torch.sigmoid(gate)
+        if hasattr(self, "apply_o"):
+            attn_output = self.apply_o(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    return fused_forward
+
+
+def patch_qwen3_5_fused_attn() -> None:
+    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Attention
+
+    if getattr(Qwen3_5Attention, "_axolotl_fused_attn_patched", False):
+        return
+
+    Qwen3_5Attention.forward = _make_fused_forward()
+    Qwen3_5Attention._axolotl_fused_attn_patched = True
+    logger.info(
+        "Patched Qwen3_5Attention.forward with fused RMSNorm+RoPE Triton kernel"
+    )
diff --git a/src/axolotl/monkeypatch/models/qwen3_5/modeling.py b/src/axolotl/monkeypatch/models/qwen3_5/modeling.py
new file mode 100644
index 0000000000..4dca30c751
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_5/modeling.py
@@ -0,0 +1,309 @@
+"""Monkeypatch for Qwen3_5 and Qwen3_5Moe models to pass position_ids to linear attention."""
+
+import importlib
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from axolotl.monkeypatch.lora_kernels import LINEAR_ATTN_IN_PROJS
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+try:
+    from fla.modules.convolution import (
+        causal_conv1d as fla_causal_conv1d,  # FLA >= 0.4.1
+    )
+except ImportError:
+    try:
+        from fla.modules.conv import causal_conv1d as fla_causal_conv1d  # FLA < 0.4.1
+    except ImportError:
+        fla_causal_conv1d = None
+
+
+def get_cu_seqlens(position_ids):
+    """
+    Compute cumulative sequence lengths from position_ids for FLA varlen kernels.
+
+    Adapted from transformers.modeling_flash_attention_utils.prepare_fa_kwargs_from_position_ids.
+    https://github.com/huggingface/transformers/blob/0f1b128d3359a26bd18be99c26d7f04fb3cba914/src/transformers/modeling_flash_attention_utils.py#L316
+
+    Qwen3.5 uses MRoPE: position_ids arrive as [axes, B, T]. All axes carry the
+    same temporal positions, so axis 0 is used to recover the [B, T] layout.
+    See: https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+    """
+    if position_ids.ndim == 3:
+        position_ids = position_ids[0]
+
+    tensor_kwargs = {"dtype": torch.int32, "device": position_ids.device}
+    position_ids = position_ids.view(-1)
+    indices_q = (position_ids == 0).nonzero().view(-1)
+    return torch.cat(
+        (
+            indices_q.to(**tensor_kwargs),
+            torch.tensor(position_ids.size(), **tensor_kwargs),
+        )
+    )
+
+
+def _inject_fla_kernels(module) -> None:
+    """Inject FLA kernels into a modeling module, bypassing is_flash_linear_attention_available."""
+    try:
+        from fla.modules import FusedRMSNormGated
+        from fla.ops.gated_delta_rule import (
+            chunk_gated_delta_rule,
+            fused_recurrent_gated_delta_rule,
+        )
+
+        module.FusedRMSNormGated = FusedRMSNormGated
+        module.chunk_gated_delta_rule = chunk_gated_delta_rule
+        module.fused_recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule
+        module.is_fast_path_available = True
+    except ImportError:
+        module.chunk_gated_delta_rule = None
+        module.fused_recurrent_gated_delta_rule = None
+        module.FusedRMSNormGated = None
+
+
+def _patched_decoder_forward(
+    self,
+    hidden_states: torch.Tensor,
+    position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values=None,
+    cache_position: Optional[torch.LongTensor] = None,
+    **kwargs,
+) -> torch.FloatTensor:
+    """Decoder layer forward that passes position_ids through to linear attention."""
+    residual = hidden_states
+    hidden_states = self.input_layernorm(hidden_states)
+
+    if self.layer_type == "linear_attention":
+        hidden_states = self.linear_attn(
+            hidden_states=hidden_states,
+            cache_params=past_key_values,
+            cache_position=cache_position,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+        )
+    elif self.layer_type == "full_attention":
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+
+    hidden_states = residual + hidden_states
+
+    residual = hidden_states
+    hidden_states = self.post_attention_layernorm(hidden_states)
+    hidden_states = self.mlp(hidden_states)
+    if isinstance(hidden_states, tuple):  # MoE returns (hidden_states, router_logits)
+        hidden_states, _ = hidden_states
+    hidden_states = residual + hidden_states
+
+    return hidden_states
+
+
+def _la_proj_fwd(module, proj_name, x):
+    """Fused kernel when patched (skips peft's bf16->fp32->bf16 round-trip), else peft."""
+    apply_fn = getattr(module, f"apply_{proj_name}", None)
+    if apply_fn is not None:
+        return apply_fn(x)
+    return getattr(module, proj_name)(x)
+
+
+def _la_in_proj_fwd(module, x):
+    fused = getattr(module, "apply_in_proj_fused", None)
+    if fused is not None:
+        return fused(x)
+    return {name: getattr(module, name)(x) for name in LINEAR_ATTN_IN_PROJS}
+
+
+def _make_qwen3_5_gated_delta_forward(apply_mask_fn):
+    """Factory for patched Qwen3_5/Qwen3_5Moe GatedDeltaNet forward with packing support."""
+
+    def patched_forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params=None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+    ):
+        hidden_states = apply_mask_fn(hidden_states, attention_mask)
+
+        batch_size, seq_len, _ = hidden_states.shape
+
+        use_precomputed_states = (
+            cache_params is not None
+            and cache_params.has_previous_state
+            and seq_len == 1
+            and cache_position is not None
+        )
+
+        cu_seqlens = None
+        if not use_precomputed_states and position_ids is not None:
+            cu_seqlens = get_cu_seqlens(position_ids=position_ids)
+
+        if cache_params is not None:
+            conv_state = cache_params.conv_states[self.layer_idx]
+            recurrent_state = cache_params.recurrent_states[self.layer_idx]
+
+        # All in-projections share hidden_states; fuse into one autograd node.
+        # mixed_qkv stays [B, T, D]; only transposed inside paths that require [B, D, T]
+        in_proj = _la_in_proj_fwd(self, hidden_states)
+        mixed_qkv = in_proj["in_proj_qkv"]  # [B, T, D]
+
+        z = in_proj["in_proj_z"]
+        z = z.reshape(batch_size, seq_len, -1, self.head_v_dim)
+
+        b = in_proj["in_proj_b"]
+        a = in_proj["in_proj_a"]
+
+        if use_precomputed_states:
+            mixed_qkv = self.causal_conv1d_update(
+                mixed_qkv.transpose(1, 2),
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                self.activation,
+            ).transpose(1, 2)
+        else:
+            if cache_params is not None:
+                mixed_qkv_t = mixed_qkv.transpose(1, 2)
+                cache_params.conv_states[self.layer_idx] = F.pad(
+                    mixed_qkv_t,
+                    (self.conv_kernel_size - mixed_qkv_t.shape[-1], 0),
+                )
+
+            if fla_causal_conv1d is not None and cu_seqlens is not None:
+                # FLA varlen kernel for packed sequences; input must be contiguous [B, T, D]
+                mixed_qkv, _ = fla_causal_conv1d(
+                    x=mixed_qkv,
+                    weight=self.conv1d.weight.squeeze(1),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                    cu_seqlens=cu_seqlens,
+                )
+            else:
+                if cu_seqlens is not None and fla_causal_conv1d is None:
+                    raise RuntimeError(
+                        "Packed sequences require fla.modules.convolution.causal_conv1d "
+                        "(cu_seqlens support). Install flash-linear-attention or disable packing."
+                    )
+                mixed_qkv = F.silu(
+                    self.conv1d(mixed_qkv.transpose(1, 2))[:, :, :seq_len]
+                ).transpose(1, 2)
+
+        query, key, value = torch.split(
+            mixed_qkv,
+            [self.key_dim, self.key_dim, self.value_dim],
+            dim=-1,
+        )
+        query = query.reshape(batch_size, seq_len, -1, self.head_k_dim)
+        key = key.reshape(batch_size, seq_len, -1, self.head_k_dim)
+        value = value.reshape(batch_size, seq_len, -1, self.head_v_dim)
+
+        beta = b.sigmoid()
+        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
+        if self.num_v_heads // self.num_k_heads > 1:
+            query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+            key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+
+        if not use_precomputed_states:
+            core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+                query,
+                key,
+                value,
+                g=g.to(dtype=query.dtype),
+                beta=beta,
+                initial_state=None,
+                output_final_state=cache_params is not None,
+                use_qk_l2norm_in_kernel=True,
+                # torch_chunk_gated_delta_rule fallback does not accept cu_seqlens
+                **({"cu_seqlens": cu_seqlens} if cu_seqlens is not None else {}),
+            )
+        else:
+            core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+                query,
+                key,
+                value,
+                g=g.to(dtype=query.dtype),
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=cache_params is not None,
+                use_qk_l2norm_in_kernel=True,
+            )
+
+        if cache_params is not None:
+            cache_params.recurrent_states[self.layer_idx] = last_recurrent_state
+
+        core_attn_out = core_attn_out.reshape(-1, self.head_v_dim)
+        z = z.reshape(-1, self.head_v_dim)
+        core_attn_out = self.norm(core_attn_out, z)
+        core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1)
+
+        return _la_proj_fwd(self, "out_proj", core_attn_out)
+
+    return patched_forward
+
+
+def _apply_packing_patches(model_type: str, cls_prefix: str, forward_factory) -> None:
+    module_name = f"transformers.models.{model_type}.modeling_{model_type}"
+
+    try:
+        module = importlib.import_module(module_name)
+    except ImportError:
+        LOG.warning(f"{model_type} not found in transformers, skipping packing patches")
+        return
+
+    _inject_fla_kernels(module)
+    getattr(module, f"{cls_prefix}DecoderLayer").forward = _patched_decoder_forward
+    gated_cls = getattr(module, f"{cls_prefix}GatedDeltaNet")
+    gated_cls.forward = forward_factory(module.apply_mask_to_padding_states)
+
+    LOG.info(
+        f"Applied {cls_prefix} packing patch "
+        f"(fla_causal_conv1d={'available' if fla_causal_conv1d else 'unavailable'})"
+    )
+
+
+def patch_qwen3_5_modeling_packing():
+    _apply_packing_patches("qwen3_5", "Qwen3_5", _make_qwen3_5_gated_delta_forward)
+
+
+def patch_qwen3_5_moe_modeling_packing():
+    _apply_packing_patches(
+        "qwen3_5_moe", "Qwen3_5Moe", _make_qwen3_5_gated_delta_forward
+    )
+
+
+def patch_qwen3_5_vlm_flash_attention():
+    """
+    Patch _is_packed_sequence to handle Qwen3.5's 3-D MRoPE position_ids.
+
+    transformers passes position_ids as [axes, B, T] to decoder layers, but
+    _is_packed_sequence only handles 2-D tensors and mis-classifies the 3-D
+    shape as a packed-sequence indicator, causing CUDA errors in the varlen path.
+    """
+    try:
+        import transformers.modeling_flash_attention_utils as fa_utils
+
+        _original = fa_utils._is_packed_sequence
+
+        def _patched(position_ids, batch_size):
+            if position_ids is not None and position_ids.ndim != 2:
+                return False
+            return _original(position_ids, batch_size)
+
+        fa_utils._is_packed_sequence = _patched
+        LOG.info("Applied Qwen3.5 VLM flash-attention patch (3-D MRoPE position_ids)")
+    except Exception as exc:  # pragma: no cover
+        LOG.warning(f"Failed to apply Qwen3.5 VLM flash-attention patch: {exc}")
diff --git a/src/axolotl/monkeypatch/models/qwen3_5_moe/__init__.py b/src/axolotl/monkeypatch/models/qwen3_5_moe/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/qwen3_5_moe/fused_attn.py b/src/axolotl/monkeypatch/models/qwen3_5_moe/fused_attn.py
new file mode 100644
index 0000000000..08a47ca65e
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_5_moe/fused_attn.py
@@ -0,0 +1,135 @@
+"""Qwen3.5-MoE variant of the qwen3_5 fused-attention monkeypatch."""
+
+from typing import Callable
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def _resolve_norm_module(norm):
+    """Unwrap PEFT ``ModulesToSaveWrapper`` so the kernel reads the active adapter's weight."""
+    modules_to_save = getattr(norm, "modules_to_save", None)
+    if not modules_to_save:
+        return norm
+    adapters = getattr(norm, "active_adapters", None)
+    if adapters is None:
+        adapter = getattr(norm, "active_adapter", None)
+        adapters = [adapter] if adapter is not None else []
+    elif isinstance(adapters, str):
+        adapters = [adapters]
+    for name in adapters:
+        if isinstance(name, str) and name in modules_to_save:
+            return modules_to_save[name]
+    return getattr(norm, "original_module", norm)
+
+
+def _make_fused_forward():
+    from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+    def fused_forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values=None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+        from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
+            eager_attention_forward,
+        )
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q_norm = _resolve_norm_module(self.q_norm)
+        k_norm = _resolve_norm_module(self.k_norm)
+        # Liger's RMSNorm replacement uses ``variance_epsilon`` instead of ``eps``.
+        eps = getattr(q_norm, "eps", None)
+        if eps is None:
+            eps = q_norm.variance_epsilon
+
+        cos, sin = position_embeddings
+
+        has_lora_qkv = hasattr(self, "apply_qkv")
+        if has_lora_qkv:
+            query_states, key_states, value_states = self.apply_qkv(hidden_states)
+            query_states, gate = torch.chunk(
+                query_states.view(*input_shape, -1, self.head_dim * 2), 2, dim=-1
+            )
+            query_states = query_states.reshape(hidden_shape)
+            key_states = key_states.view(hidden_shape)
+            value_states = value_states.view(hidden_shape)
+        else:
+            query_states, gate = torch.chunk(
+                self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2),
+                2,
+                dim=-1,
+            )
+            query_states = query_states.reshape(hidden_shape)
+            key_states = self.k_proj(hidden_states).view(hidden_shape)
+            value_states = self.v_proj(hidden_states).view(hidden_shape)
+        gate = gate.reshape(*input_shape, -1)
+
+        # accelerate's per-module pre-hooks that move CPU-staged params don't fire on a monkeypatched forward.
+        q_w = q_norm.weight
+        if q_w.device != query_states.device:
+            q_w = q_w.to(query_states.device, non_blocking=True)
+        k_w = k_norm.weight
+        if k_w.device != key_states.device:
+            k_w = k_w.to(key_states.device, non_blocking=True)
+
+        query_states = fused_rms_norm_rope(
+            query_states, q_w, cos, sin, eps=eps, unit_offset=True
+        ).transpose(1, 2)
+        key_states = fused_rms_norm_rope(
+            key_states, k_w, cos, sin, eps=eps, unit_offset=True
+        ).transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = attn_output * torch.sigmoid(gate)
+        if hasattr(self, "apply_o"):
+            attn_output = self.apply_o(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    return fused_forward
+
+
+def patch_qwen3_5_moe_fused_attn() -> None:
+    from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
+        Qwen3_5MoeAttention,
+    )
+
+    if getattr(Qwen3_5MoeAttention, "_axolotl_fused_attn_patched", False):
+        return
+
+    Qwen3_5MoeAttention.forward = _make_fused_forward()
+    Qwen3_5MoeAttention._axolotl_fused_attn_patched = True
+    logger.info(
+        "Patched Qwen3_5MoeAttention.forward with fused RMSNorm+RoPE Triton kernel"
+    )
diff --git a/src/axolotl/monkeypatch/models/qwen3_moe/__init__.py b/src/axolotl/monkeypatch/models/qwen3_moe/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/qwen3_moe/fused_attn.py b/src/axolotl/monkeypatch/models/qwen3_moe/fused_attn.py
new file mode 100644
index 0000000000..4f5b951ca5
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_moe/fused_attn.py
@@ -0,0 +1,121 @@
+"""Qwen3-MoE variant of the qwen3 fused-attention monkeypatch."""
+
+from typing import Callable
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def _resolve_norm_module(norm):
+    """Unwrap PEFT ``ModulesToSaveWrapper`` so the kernel reads the active adapter's weight."""
+    modules_to_save = getattr(norm, "modules_to_save", None)
+    if not modules_to_save:
+        return norm
+    adapters = getattr(norm, "active_adapters", None)
+    if adapters is None:
+        adapter = getattr(norm, "active_adapter", None)
+        adapters = [adapter] if adapter is not None else []
+    elif isinstance(adapters, str):
+        adapters = [adapters]
+    for name in adapters:
+        if isinstance(name, str) and name in modules_to_save:
+            return modules_to_save[name]
+    return getattr(norm, "original_module", norm)
+
+
+def _make_fused_forward():
+    from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+    def fused_forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values=None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+        from transformers.models.qwen3_moe.modeling_qwen3_moe import (
+            eager_attention_forward,
+        )
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q_norm = _resolve_norm_module(self.q_norm)
+        k_norm = _resolve_norm_module(self.k_norm)
+        eps = q_norm.variance_epsilon
+
+        cos, sin = position_embeddings
+
+        has_lora_qkv = hasattr(self, "apply_qkv")
+        if has_lora_qkv:
+            query_states, key_states, value_states = self.apply_qkv(hidden_states)
+            query_states = query_states.view(hidden_shape)
+            key_states = key_states.view(hidden_shape)
+            value_states = value_states.view(hidden_shape)
+        else:
+            query_states = self.q_proj(hidden_states).view(hidden_shape)
+            key_states = self.k_proj(hidden_states).view(hidden_shape)
+            value_states = self.v_proj(hidden_states).view(hidden_shape)
+
+        # accelerate's per-module pre-hooks that move CPU-staged params don't fire on a monkeypatched forward.
+        q_w = q_norm.weight
+        if q_w.device != query_states.device:
+            q_w = q_w.to(query_states.device, non_blocking=True)
+        k_w = k_norm.weight
+        if k_w.device != key_states.device:
+            k_w = k_w.to(key_states.device, non_blocking=True)
+
+        query_states = fused_rms_norm_rope(
+            query_states, q_w, cos, sin, eps=eps
+        ).transpose(1, 2)
+        key_states = fused_rms_norm_rope(key_states, k_w, cos, sin, eps=eps).transpose(
+            1, 2
+        )
+        value_states = value_states.transpose(1, 2)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        if hasattr(self, "apply_o"):
+            attn_output = self.apply_o(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    return fused_forward
+
+
+def patch_qwen3_moe_fused_attn() -> None:
+    from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeAttention
+
+    if getattr(Qwen3MoeAttention, "_axolotl_fused_attn_patched", False):
+        return
+
+    Qwen3MoeAttention.forward = _make_fused_forward()
+    Qwen3MoeAttention._axolotl_fused_attn_patched = True
+    logger.info(
+        "Patched Qwen3MoeAttention.forward with fused RMSNorm+RoPE Triton kernel"
+    )
diff --git a/src/axolotl/monkeypatch/models/qwen3_next/__init__.py b/src/axolotl/monkeypatch/models/qwen3_next/__init__.py
new file mode 100644
index 0000000000..39bcd4115c
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_next/__init__.py
@@ -0,0 +1 @@
+"""Qwen3_Next model monkeypatches."""
diff --git a/src/axolotl/monkeypatch/models/qwen3_next/modeling.py b/src/axolotl/monkeypatch/models/qwen3_next/modeling.py
new file mode 100644
index 0000000000..907750fe72
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_next/modeling.py
@@ -0,0 +1,346 @@
+"""Monkeypatch for Qwen3_Next model to pass position_ids to linear attention."""
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+try:
+    from fla.modules.convolution import causal_conv1d as fla_causal_conv1d
+except ImportError:
+    fla_causal_conv1d = None
+
+
+def get_cu_seqlens(position_ids):
+    """
+    Adapted from transformers.modeling_flash_attention_utils.prepare_fa_kwargs_from_position_ids.
+
+    https://github.com/huggingface/transformers/blob/0f1b128d3359a26bd18be99c26d7f04fb3cba914/src/transformers/modeling_flash_attention_utils.py#L316
+    """
+    tensor_kwargs = {"dtype": torch.int32, "device": position_ids.device}
+
+    position_ids = position_ids.view(-1)
+    indices_q = (position_ids == 0).nonzero().view(-1)
+
+    cu_seq_lens_q = torch.cat(
+        (
+            indices_q.to(**tensor_kwargs),
+            torch.tensor(position_ids.size(), **tensor_kwargs),
+        )
+    )
+
+    return cu_seq_lens_q
+
+
+def patch_qwen3_next_decoder_layer():
+    """Patch Qwen3NextDecoderLayer to pass position_ids to linear attention."""
+    try:
+        from transformers.models.qwen3_next.modeling_qwen3_next import (
+            Qwen3NextDecoderLayer,
+        )
+    except ImportError:
+        LOG.warning("Qwen3Next model not found, skipping patch")
+        return
+
+    # Store original forward method
+    original_decoder_forward = Qwen3NextDecoderLayer.forward
+
+    def patched_decoder_forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[torch.Tensor]] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Token Mixer
+        if self.layer_type == "linear_attention":
+            hidden_states = self.linear_attn(
+                hidden_states=hidden_states,
+                cache_params=past_key_values,
+                cache_position=cache_position,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+            )
+        elif self.layer_type == "full_attention":
+            # Self Attention
+            hidden_states, _ = self.self_attn(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        # For the MoE layers, we need to unpack
+        if isinstance(hidden_states, Tuple):
+            hidden_states, _ = hidden_states
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+    # Apply the patches
+    Qwen3NextDecoderLayer.forward = patched_decoder_forward
+
+    def unpatch():
+        """Restore the original forward method"""
+        Qwen3NextDecoderLayer.forward = original_decoder_forward
+
+    return unpatch
+
+
+def patch_qwen3_next_gateddelta_layer():
+    """Patch Qwen3NextGatedDeltaNet to parse cu_seqlens and pass to chunk_gated_delta_rule"""
+    try:
+        from transformers.models.qwen3_next.modeling_qwen3_next import (
+            Qwen3NextGatedDeltaNet,
+            apply_mask_to_padding_states,
+        )
+    except ImportError:
+        LOG.warning("Qwen3Next model not found, skipping patch")
+        return
+
+    # Store original forward method
+    original_gated_delta_net_forward = Qwen3NextGatedDeltaNet.forward
+
+    def patched_gated_delta_net_forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params=None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        cache_position: Optional[
+            torch.LongTensor
+        ] = None,  # unused: no cache in packed training
+    ):
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+
+        # Set up dimensions for reshapes later
+        batch_size, seq_len, _ = hidden_states.shape
+
+        use_precomputed_states = (
+            cache_params is not None
+            and cache_params.has_previous_state(self.layer_idx)
+            and seq_len == 1
+        )
+
+        # Compute cu_seqlens early for use by both causal_conv1d and chunk_gated_delta_rule
+        cu_seqlens = None
+        if not use_precomputed_states and position_ids is not None:
+            cu_seqlens = get_cu_seqlens(position_ids=position_ids)
+
+        # getting projected states from cache if it exists
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states
+
+        projected_states_qkvz = self.in_proj_qkvz(hidden_states)
+        projected_states_ba = self.in_proj_ba(hidden_states)
+        query, key, value, z, b, a = self.fix_query_key_value_ordering(
+            projected_states_qkvz, projected_states_ba
+        )
+        query, key, value = (
+            x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value)
+        )
+
+        mixed_qkv = torch.cat((query, key, value), dim=-1)  # [B, T, D]
+        mixed_qkv = mixed_qkv.transpose(1, 2)  # [B, D, T]
+
+        if use_precomputed_states:
+            mixed_qkv = self.causal_conv1d_update(
+                mixed_qkv,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                self.activation,
+            )
+        else:
+            if cache_params is not None:
+                conv_state = F.pad(
+                    mixed_qkv,
+                    (self.conv_kernel_size - mixed_qkv.shape[-1], 0),
+                )
+                cache_params.update_conv_state(conv_state, self.layer_idx)
+
+            if fla_causal_conv1d is not None:
+                # FLA Triton causal_conv1d: [B, T, D] in/out, with cu_seqlens support
+                mixed_qkv = mixed_qkv.transpose(1, 2)  # [B, T, D] for FLA
+                mixed_qkv, _ = fla_causal_conv1d(
+                    x=mixed_qkv,
+                    weight=self.conv1d.weight.squeeze(1),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                    cu_seqlens=cu_seqlens,
+                )
+                mixed_qkv = mixed_qkv.transpose(1, 2)  # back to [B, D, T]
+            elif self.causal_conv1d_fn is not None:
+                mixed_qkv = self.causal_conv1d_fn(
+                    x=mixed_qkv,
+                    weight=self.conv1d.weight.squeeze(1),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                    seq_idx=None,
+                )
+            else:
+                # PyTorch fallback (no cu_seqlens support)
+                if cu_seqlens is not None and cu_seqlens.shape[0] > batch_size + 1:
+                    raise RuntimeError(
+                        "Packed sequences require fla.modules.convolution.causal_conv1d "
+                        "(cu_seqlens support). Install flash-linear-attention or disable packing."
+                    )
+                LOG.warning_once(
+                    "FLA causal_conv1d not available. Falling back to PyTorch conv1d."
+                )
+                mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len])
+
+        mixed_qkv = mixed_qkv.transpose(1, 2)  # [B, T, D]
+        query, key, value = torch.split(
+            mixed_qkv,
+            [
+                self.key_dim,
+                self.key_dim,
+                self.value_dim,
+            ],
+            dim=-1,
+        )
+        query = query.reshape(query.shape[0], query.shape[1], -1, self.head_k_dim)
+        key = key.reshape(key.shape[0], key.shape[1], -1, self.head_k_dim)
+        value = value.reshape(value.shape[0], value.shape[1], -1, self.head_v_dim)
+
+        beta = b.sigmoid()
+        # If the model is loaded in fp16, without the .float() here, A might be -inf
+        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
+        if self.num_v_heads // self.num_k_heads > 1:
+            query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+            key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+
+        if not use_precomputed_states:
+            core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+                query,
+                key,
+                value,
+                g=g,
+                beta=beta,
+                initial_state=None,
+                output_final_state=cache_params is not None,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
+            )
+
+        else:
+            core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+                query,
+                key,
+                value,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=cache_params is not None,
+                use_qk_l2norm_in_kernel=True,
+            )
+
+        # Update cache
+        if cache_params is not None:
+            cache_params.update_recurrent_state(last_recurrent_state, self.layer_idx)
+
+        z_shape_og = z.shape
+        # reshape input data into 2D tensor
+        core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
+        z = z.reshape(-1, z.shape[-1])
+        core_attn_out = self.norm(core_attn_out, z)
+        core_attn_out = core_attn_out.reshape(z_shape_og)
+        core_attn_out = core_attn_out.reshape(
+            core_attn_out.shape[0], core_attn_out.shape[1], -1
+        )
+
+        output = self.out_proj(core_attn_out)
+        return output
+
+    # Apply the patches
+    Qwen3NextGatedDeltaNet.forward = patched_gated_delta_net_forward
+
+    def unpatch():
+        """Restore the original forward method"""
+        Qwen3NextGatedDeltaNet.forward = original_gated_delta_net_forward
+
+    return unpatch
+
+
+def patch_qwen3_next_imports():
+    """Patch Qwen3Next imports to use try/except instead of is_flash_linear_attention_available."""
+    try:
+        import transformers.models.qwen3_next.modeling_qwen3_next as qwen3_modeling
+    except ImportError:
+        LOG.warning("Qwen3Next model not found, skipping import patch")
+        return
+
+    # Save original values for unpatch
+    original_FusedRMSNormGated = getattr(qwen3_modeling, "FusedRMSNormGated", None)
+    original_chunk_gated_delta_rule = getattr(
+        qwen3_modeling, "chunk_gated_delta_rule", None
+    )
+    original_fused_recurrent_gated_delta_rule = getattr(
+        qwen3_modeling, "fused_recurrent_gated_delta_rule", None
+    )
+    original_is_fast_path_available = getattr(
+        qwen3_modeling, "is_fast_path_available", False
+    )
+
+    try:
+        from fla.modules import FusedRMSNormGated
+        from fla.ops.gated_delta_rule import (
+            chunk_gated_delta_rule,
+            fused_recurrent_gated_delta_rule,
+        )
+
+        qwen3_modeling.FusedRMSNormGated = FusedRMSNormGated
+        qwen3_modeling.chunk_gated_delta_rule = chunk_gated_delta_rule
+        qwen3_modeling.fused_recurrent_gated_delta_rule = (
+            fused_recurrent_gated_delta_rule
+        )
+
+        # Force is_fast_path_available to be True
+        # fla has triton kernels for causal_conv1d
+        qwen3_modeling.is_fast_path_available = True
+    except ImportError:
+        qwen3_modeling.chunk_gated_delta_rule = None
+        qwen3_modeling.fused_recurrent_gated_delta_rule = None
+        qwen3_modeling.FusedRMSNormGated = None
+
+    def unpatch():
+        """Restore the original import values"""
+        qwen3_modeling.FusedRMSNormGated = original_FusedRMSNormGated
+        qwen3_modeling.chunk_gated_delta_rule = original_chunk_gated_delta_rule
+        qwen3_modeling.fused_recurrent_gated_delta_rule = (
+            original_fused_recurrent_gated_delta_rule
+        )
+        qwen3_modeling.is_fast_path_available = original_is_fast_path_available
+
+    return unpatch
+
+
+def patch_qwen3_next_modeling_packing():
+    """Apply all Qwen3Next model patches."""
+    patch_qwen3_next_imports()
+    patch_qwen3_next_decoder_layer()
+    patch_qwen3_next_gateddelta_layer()
+
+    LOG.info("Applied Qwen3Next patch for packing")
diff --git a/src/axolotl/monkeypatch/models/qwen3_vl/__init__.py b/src/axolotl/monkeypatch/models/qwen3_vl/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/qwen3_vl/fused_attn.py b/src/axolotl/monkeypatch/models/qwen3_vl/fused_attn.py
new file mode 100644
index 0000000000..68da52bd5f
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_vl/fused_attn.py
@@ -0,0 +1,121 @@
+# Why: fuse Qwen3-VL q_norm/k_norm + mRoPE into one Triton kernel.
+
+from typing import Callable
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def _resolve_norm_module(norm):
+    # Why: ModulesToSaveWrapper stores trainable norm weights per active adapter.
+    modules_to_save = getattr(norm, "modules_to_save", None)
+    if not modules_to_save:
+        return norm
+    adapters = getattr(norm, "active_adapters", None)
+    if adapters is None:
+        adapter = getattr(norm, "active_adapter", None)
+        adapters = [adapter] if adapter is not None else []
+    elif isinstance(adapters, str):
+        adapters = [adapters]
+    for name in adapters:
+        if isinstance(name, str) and name in modules_to_save:
+            return modules_to_save[name]
+    return getattr(norm, "original_module", norm)
+
+
+def _make_fused_forward():
+    from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+    def fused_forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values=None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+        from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+            eager_attention_forward,
+        )
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q_norm = _resolve_norm_module(self.q_norm)
+        k_norm = _resolve_norm_module(self.k_norm)
+        eps = getattr(q_norm, "eps", None)
+        if eps is None:
+            eps = q_norm.variance_epsilon
+
+        cos, sin = position_embeddings
+
+        if hasattr(self, "apply_qkv"):
+            query_states, key_states, value_states = self.apply_qkv(hidden_states)
+            query_states = query_states.view(hidden_shape)
+            key_states = key_states.view(hidden_shape)
+            value_states = value_states.view(hidden_shape)
+        else:
+            query_states = self.q_proj(hidden_states).view(hidden_shape)
+            key_states = self.k_proj(hidden_states).view(hidden_shape)
+            value_states = self.v_proj(hidden_states).view(hidden_shape)
+
+        # accelerate's per-module pre-hooks that move CPU-staged params don't fire on a monkeypatched forward.
+        q_w = q_norm.weight
+        if q_w.device != query_states.device:
+            q_w = q_w.to(query_states.device, non_blocking=True)
+        k_w = k_norm.weight
+        if k_w.device != key_states.device:
+            k_w = k_w.to(key_states.device, non_blocking=True)
+
+        query_states = fused_rms_norm_rope(
+            query_states, q_w, cos, sin, eps=eps
+        ).transpose(1, 2)
+        key_states = fused_rms_norm_rope(key_states, k_w, cos, sin, eps=eps).transpose(
+            1, 2
+        )
+        value_states = value_states.transpose(1, 2)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        if hasattr(self, "apply_o"):
+            attn_output = self.apply_o(attn_output)
+        else:
+            attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    return fused_forward
+
+
+def patch_qwen3_vl_fused_attn() -> None:
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextAttention
+
+    if getattr(Qwen3VLTextAttention, "_axolotl_fused_attn_patched", False):
+        return
+
+    Qwen3VLTextAttention.forward = _make_fused_forward()
+    Qwen3VLTextAttention._axolotl_fused_attn_patched = True
+    logger.info(
+        "Patched Qwen3VLTextAttention.forward with fused RMSNorm+mRoPE Triton kernel"
+    )
diff --git a/src/axolotl/monkeypatch/models/voxtral/__init__.py b/src/axolotl/monkeypatch/models/voxtral/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/models/voxtral/modeling.py b/src/axolotl/monkeypatch/models/voxtral/modeling.py
new file mode 100644
index 0000000000..3dd652dd8e
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/voxtral/modeling.py
@@ -0,0 +1,67 @@
+"""Monkeypatch for voxtral to fix leaf node and dtype mismatch"""
+
+from typing import Optional, Union
+
+import torch
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def patch_voxtral_conditional_generation_forward():
+    from transformers.models.voxtral.modeling_voxtral import (
+        VoxtralForConditionalGeneration,
+    )
+
+    # Store the original forward method
+    old_forward = VoxtralForConditionalGeneration.forward
+
+    def _forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        input_features: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        if input_features is not None:
+            audio_embeds = self.get_audio_embeds(input_features)
+
+            # Cast audio_embeds to match inputs_embeds dtype
+            audio_embeds = audio_embeds.to(inputs_embeds.dtype)
+
+            # replace text-audio token placeholders with audio embeddings
+            audio_token_mask = input_ids == self.config.audio_token_id
+
+            inputs_embeds = inputs_embeds.clone()
+            inputs_embeds[audio_token_mask] = audio_embeds
+
+        outputs = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            **kwargs,
+        )
+        return outputs
+
+    # Apply the patch
+    VoxtralForConditionalGeneration.forward = _forward
+
+    def unpatch():
+        """Restore the original forward method"""
+        VoxtralForConditionalGeneration.forward = old_forward
+
+    return unpatch
diff --git a/src/axolotl/monkeypatch/moe_quant.py b/src/axolotl/monkeypatch/moe_quant.py
new file mode 100644
index 0000000000..d58b78af2d
--- /dev/null
+++ b/src/axolotl/monkeypatch/moe_quant.py
@@ -0,0 +1,317 @@
+"""Loading-time quantization for MoE expert weights stored as 3D nn.Parameter tensors."""
+
+import bitsandbytes as bnb
+import torch
+import torch.nn.utils.parametrize as P
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_moe_load_state = {
+    "count": 0,
+    "mode": "4bit",
+    "quant_type": "nf4",
+    "compress_statistics": True,
+    "patched": False,
+    # Module path → param names in definition order, captured before quantization.
+    # Without this, alphabetical loading order would mismatch merge order.
+    "expert_param_order": {},
+}
+
+
+class Bnb8bitParametrization(torch.nn.Module):
+    """Dequantizes int8 row-wise quantized data on access."""
+
+    def __init__(self, row_stats: torch.Tensor):
+        super().__init__()
+        self.register_buffer("row_stats", row_stats)
+
+    @torch.no_grad()
+    def forward(self, quantized_param: torch.Tensor) -> torch.Tensor:
+        """Flatten 3D+ to 2D for BnB's dequant, then reshape back."""
+        orig_shape = quantized_param.shape
+        if quantized_param.ndim > 2:
+            quantized_param = quantized_param.reshape(-1, orig_shape[-1])
+        result = bnb.functional.int8_vectorwise_dequant(quantized_param, self.row_stats)
+        return result.reshape(orig_shape)
+
+
+def _enable_parametrization_cache(module, inputs):
+    P._cache_enabled += 1
+
+
+def _disable_parametrization_cache(module, inputs, output):
+    P._cache_enabled -= 1
+    if not P._cache_enabled:
+        P._cache = {}
+
+
+def replace_parameter_8bit(module, param_name):
+    """Replace a module parameter with an 8-bit quantized version using parametrization."""
+    original_param = getattr(module, param_name)
+    int8_data, row_stats, _ = bnb.functional.int8_vectorwise_quant(
+        original_param.data.to(torch.float16)
+    )
+
+    setattr(module, param_name, torch.nn.Parameter(int8_data, requires_grad=False))
+    del original_param
+
+    P.register_parametrization(
+        module, param_name, Bnb8bitParametrization(row_stats), unsafe=True
+    )
+
+    # Cache dequantized values during forward to avoid redundant dequantization.
+    if not getattr(module, "_axolotl_8bit_hooks_registered", False):
+        module.register_forward_pre_hook(_enable_parametrization_cache)
+        module.register_forward_hook(_disable_parametrization_cache)
+        module._axolotl_8bit_hooks_registered = True
+
+
+def patch_moe_quantization_on_load(cfg):
+    """Patch transformers' weight loading to quantize MoE expert params on-the-fly."""
+    mode = "8bit" if getattr(cfg, "load_in_8bit", False) else "4bit"
+    _moe_load_state["mode"] = mode
+    _moe_load_state["count"] = 0
+    _moe_load_state["expert_param_order"] = {}
+
+    if _moe_load_state["patched"]:
+        LOG.debug("MoE loading-time quantization patch already active")
+        return
+
+    import transformers.core_model_loading
+    import transformers.modeling_utils
+
+    if mode == "4bit":
+        from bitsandbytes.nn.parametrize import replace_parameter_4bit
+
+        quant_type = getattr(cfg, "bnb_4bit_quant_type", None) or "nf4"
+        compress_statistics = getattr(cfg, "bnb_4bit_use_double_quant", None)
+        if compress_statistics is None:
+            compress_statistics = True
+
+        _moe_load_state["quant_type"] = quant_type
+        _moe_load_state["compress_statistics"] = compress_statistics
+
+    # Disable caching_allocator_warmup — it pre-allocates a huge tensor at bf16
+    # size for all params, defeating our on-load quantization VRAM savings.
+    def _noop_warmup(*args, **kwargs):
+        pass
+
+    transformers.modeling_utils.caching_allocator_warmup = _noop_warmup
+
+    original_set_param = transformers.core_model_loading.set_param_for_module
+
+    def _patched_set_param_for_module(model, target_name, param_value, *args, **kwargs):
+        original_set_param(model, target_name, param_value, *args, **kwargs)
+
+        if param_value.ndim >= 3 and param_value.is_cuda:
+            mod_path, _, pname = target_name.rpartition(".")
+            mod = model.get_submodule(mod_path) if mod_path else model
+            if not isinstance(mod, (bnb.nn.Linear4bit, bnb.nn.Linear8bitLt)):
+                if "expert" not in target_name.lower():
+                    LOG.debug(
+                        "Skipping non-expert 3D param: %s (shape=%s)",
+                        target_name,
+                        list(param_value.shape),
+                    )
+                    return
+
+                # Record definition order before parametrizations override it
+                # with alphabetical order.
+                if mod_path not in _moe_load_state["expert_param_order"]:
+                    _moe_load_state["expert_param_order"][mod_path] = list(
+                        mod._parameters.keys()
+                    )
+
+                if _moe_load_state["mode"] == "4bit":
+                    replace_parameter_4bit(
+                        mod,
+                        pname,
+                        compress_statistics=_moe_load_state["compress_statistics"],
+                        quant_type=_moe_load_state["quant_type"],
+                    )
+                else:
+                    replace_parameter_8bit(mod, pname)
+                _moe_load_state["count"] += 1
+
+                # Release the bf16 tensor so CUDA memory is freed immediately.
+                param_value.data = torch.empty(0, device="cpu")
+                torch.cuda.empty_cache()
+
+    transformers.core_model_loading.set_param_for_module = _patched_set_param_for_module
+    _moe_load_state["patched"] = True
+
+
+def get_moe_quantized_count():
+    """Return the number of expert parameters quantized during loading."""
+    return _moe_load_state["count"]
+
+
+def patch_peft_target_parameters_matching():
+    """Fix PEFT's _inject_parameters for target_parameters on quantized MoE experts.
+
+    1. Expands short suffixes to full module paths for parametrized modules.
+    2. Iterates params in definition order (not alphabetical order) so saved
+       adapters are compatible with standard PEFT, vLLM, etc.
+    3. Skips ParametrizationList synthetic paths to prevent PEFT from mistakenly
+       targeting quantized expert params via name-suffix matching.
+    """
+    if getattr(patch_peft_target_parameters_matching, "_axolotl_patched", False):
+        return
+
+    from contextlib import nullcontext
+
+    from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer
+    from peft.utils.integrations import init_empty_weights
+    from peft.utils.other import _get_submodules
+
+    # Mapping from unfused parameter names to their fused equivalents.
+    # When a model stores fused weights (e.g. gate_up_proj) but the user
+    # specifies unfused names (gate_proj, up_proj), we auto-expand so the
+    # fused parameter is also targeted. The original unfused names are kept
+    # in the set so that models that do NOT fuse still work.
+    _UNFUSED_TO_FUSED: dict[str, str] = {
+        "gate_proj": "gate_up_proj",
+        "up_proj": "gate_up_proj",
+    }
+
+    def _patched_inject_parameters(
+        self, peft_config, model, adapter_name, low_cpu_mem_usage
+    ):
+        original_targets = list(peft_config.target_parameters)
+        expanded = set(original_targets)
+
+        # Expand short suffixes to full paths for parametrized modules.
+        for module_name, module in model.named_modules():
+            if not hasattr(module, "parametrizations"):
+                continue
+            for target in original_targets:
+                mod_path, _, param_name = target.rpartition(".")
+                if not (
+                    module_name == mod_path or module_name.endswith("." + mod_path)
+                ):
+                    continue
+
+                if hasattr(module, param_name):
+                    expanded.add(f"{module_name}.{param_name}")
+                elif param_name in _UNFUSED_TO_FUSED:
+                    # The model uses fused weights (e.g. gate_up_proj) but the
+                    # user specified unfused names (gate_proj / up_proj).
+                    fused_name = _UNFUSED_TO_FUSED[param_name]
+                    if hasattr(module, fused_name):
+                        if fused_name not in expanded:
+                            LOG.warning(
+                                "target_parameter '%s' not found on %s, "
+                                "but fused equivalent '%s' exists — adding "
+                                "it automatically.",
+                                param_name,
+                                module_name,
+                                fused_name,
+                            )
+                        expanded.add(f"{module_name}.{fused_name}")
+                    else:
+                        LOG.warning(
+                            "target_parameter '%s' not found on %s and no "
+                            "fused equivalent exists either — skipping.",
+                            param_name,
+                            module_name,
+                        )
+                else:
+                    LOG.warning(
+                        "target_parameter '%s' not found on %s — skipping. "
+                        "Check that the parameter name matches the model's "
+                        "weight names.",
+                        param_name,
+                        module_name,
+                    )
+
+        target_names_set = expanded
+
+        def strip_base_layer_from_name(module_name):
+            name = ".base_layer"
+            while name in module_name:
+                prefix, _, suffix = module_name.rpartition(name)
+                module_name = prefix + suffix
+            return module_name
+
+        def create_and_replace_param(module_name, key, param_name):
+            parent, target, target_name = _get_submodules(model, module_name)
+            unwrapped_module_name = strip_base_layer_from_name(module_name)
+            unwrapped_module = model.get_submodule(unwrapped_module_name)
+            if (
+                isinstance(unwrapped_module, BaseTunerLayer)
+                and unwrapped_module.__class__.__name__ != "ParamWrapper"
+            ):
+                raise ValueError(
+                    f"Trying to wrap an `nn.Parameter` of layer "
+                    f"'{unwrapped_module_name}' of type "
+                    f"{type(target).__name__}, which is not a valid target. "
+                    f"Make sure that this layer is not also targeted with "
+                    f"`target_modules`."
+                )
+            self._check_target_module_compatiblity(peft_config, model, target_name)
+            ctx = init_empty_weights if low_cpu_mem_usage else nullcontext
+            with ctx():
+                self._create_and_replace(
+                    peft_config,
+                    adapter_name,
+                    target,
+                    target_name,
+                    parent,
+                    current_key=key,
+                    parameter_name=param_name.rpartition(".")[-1],
+                )
+
+        # Use definition order (not alphabetical order) for parametrized modules
+        # so ParamWrapper nesting matches vanilla PEFT on a plain model.
+        expert_param_order = _moe_load_state.get("expert_param_order", {})
+
+        for module_name, module in model.named_modules():
+            if hasattr(module, "parametrizations"):
+                stored_order = expert_param_order.get(module_name)
+                if stored_order is not None:
+                    params_iter = [
+                        p for p in stored_order if p in module.parametrizations
+                    ]
+                else:
+                    # Fallback for paths that bypass model loading (e.g. unit tests).
+                    params_iter = list(module.parametrizations.keys())
+                for param_name in params_iter:
+                    key = f"{module_name}.{param_name}"
+                    if (key in target_names_set) or any(
+                        key.endswith(f".{t}") for t in target_names_set
+                    ):
+                        create_and_replace_param(module_name, key, param_name)
+                        self.targeted_parameter_names.append(key)
+            else:
+                unwrapped_module_name = strip_base_layer_from_name(module_name)
+                for param_name, _ in module.named_parameters(recurse=False):
+                    key = f"{unwrapped_module_name}.{param_name}"
+                    if (key in target_names_set) or any(
+                        key.endswith(f".{t}") for t in target_names_set
+                    ):
+                        create_and_replace_param(module_name, key, param_name)
+                        self.targeted_parameter_names.append(key)
+
+    BaseTuner._inject_parameters = _patched_inject_parameters
+
+    # Skip ParametrizationList synthetic paths (e.g. "...parametrizations.up_proj")
+    # so PEFT suffix-matching doesn't try to wrap quantized expert params in LoRA.
+    # Previous MoE models (Mixtral, DeepSeek, etc.) stored experts as nn.Linear
+    # modules, so PEFT's normal target_modules path worked fine. NemotronH uses
+    # 3D nn.Parameter tensors via our quantize_moe_experts parametrization, which
+    # exposes synthetic ".parametrizations.<name>" paths that PEFT's suffix match
+    # would otherwise treat as target_modules candidates.
+    _original_check = BaseTuner._check_target_module_exists
+
+    @staticmethod
+    def _patched_check_target_module_exists(config, key):
+        if ".parametrizations." in key:
+            return False
+        return _original_check(config, key)
+
+    BaseTuner._check_target_module_exists = _patched_check_target_module_exists
+
+    patch_peft_target_parameters_matching._axolotl_patched = True
+    LOG.info("Patched PEFT _inject_parameters for consistent ParamWrapper ordering")
diff --git a/src/axolotl/monkeypatch/multipack.py b/src/axolotl/monkeypatch/multipack.py
index c1eb3127d2..d887615832 100644
--- a/src/axolotl/monkeypatch/multipack.py
+++ b/src/axolotl/monkeypatch/multipack.py
@@ -1,7 +1,7 @@
 """multipack patching for v2 of sample packing"""
+
 import importlib
 
-import transformers
 from accelerate import init_empty_weights
 from transformers import AutoConfig, AutoModelForCausalLM
 from transformers.integrations import is_deepspeed_zero3_enabled
@@ -10,59 +10,85 @@
 from axolotl.monkeypatch.utils import get_unpad_data
 
 SUPPORTED_MULTIPACK_MODEL_TYPES = [
+    "apertus",
+    "mllama_text_model",
+    "llama",
+    "llama4",
+    "mistral",
     "mixtral",
     "qwen2",
     "qwen2_moe",
+    "qwen3",
+    "qwen3_moe",
+    "qwen3_next",
+    "qwen3_5",
+    "qwen3_5_moe",
     "falcon",
     "phi",
+    "phi3",
     "gemma",
+    "gemma2",
+    "gemma3",
+    "gemma3_text",
+    "cohere",
+    "cohere2",
     "gemmoe",
     "starcoder2",
+    "deepseek_v2",
+    "deepseek_v3",
+    "glm",
+    "glm4",
+    "glm4_moe",
+    "glm_moe_dsa",
+    "smollm3",
+    "granite",
+    "granitemoe",
+    "granitemoeshared",
+    "granitemoehybrid",
+    "hunyuan_v1_dense",
+    "hunyuan_v1_moe",
+    "gpt_oss",
+    "arcee",
+    "seed_oss",
+    "lfm2",
+    "lfm2_moe",
+    "ernie4_5",
+    "ernie4_5_moe",
+    "olmo",
+    "olmo2",
+    "olmo3",
+    "ministral",
+    "ministral3",
+    "mistral4",
+    "afmoe",
+    "nemotron",
+    "nemotron_h",
+    "falcon_h1",
 ]
 
 
-def patch_for_multipack(model_type, model_name=None):
-    if model_type == "mixtral":
-        transformers.models.mixtral.modeling_mixtral._get_unpad_data = (  # pylint: disable=protected-access
-            get_unpad_data
-        )
-        if is_deepspeed_zero3_enabled():
-            patch_mixtral_moe_forward_zero3()
-    elif model_type == "qwen2":
-        transformers.models.qwen2.modeling_qwen2._get_unpad_data = (  # pylint: disable=protected-access
-            get_unpad_data
-        )
-    elif model_type == "qwen2_moe":
-        transformers.models.qwen2_moe.modeling_qwen2_moe._get_unpad_data = (  # pylint: disable=protected-access
-            get_unpad_data
-        )
-    elif model_type == "falcon":
-        transformers.models.falcon.modeling_falcon._get_unpad_data = (  # pylint: disable=protected-access
-            get_unpad_data
-        )
-    elif model_type == "phi":
-        transformers.models.phi.modeling_phi._get_unpad_data = (  # pylint: disable=protected-access
-            get_unpad_data
-        )
-    elif model_type == "gemma":
-        transformers.models.gemma.modeling_gemma._get_unpad_data = (  # pylint: disable=protected-access
-            get_unpad_data
-        )
-    elif model_type == "starcoder2":
-        transformers.models.starcoder2.modeling_starcoder2._get_unpad_data = (  # pylint: disable=protected-access
-            get_unpad_data
-        )
-    elif model_type == "gemmoe":
-        patch_remote(model_name, ".configuration_gemmoe", ".modeling_gemmoe")
-    elif model_type == "jamba":
-        patch_remote(model_name, ".configuration_jamba", ".modeling_jamba")
+def patch_for_multipack(model_type, model_name=None, has_remote_code=False):
+    # In-tree HF models handle sample packing natively via position_ids
+    # (transformers `_is_packed_sequence` / `find_packed_sequence_indices`): the 2D
+    # mask is cast to bool before flash attention, so the segment ids we encode in
+    # it never reach `_get_unpad_data` and overriding it there is a no-op. Only
+    # remote-code modeling files that call `_get_unpad_data` directly on the raw
+    # segment-id mask still need the override.
+    if has_remote_code:
+        patch_remote(model_name)
+
+    if model_type == "mixtral" and is_deepspeed_zero3_enabled():
+        patch_mixtral_moe_forward_zero3()
 
 
-def patch_remote(model_name, config_name, modeling_name):
+def patch_remote(model_name):
     model_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
     # we need to load the model here in order for modeling_* to be available
     with init_empty_weights():
         AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
-    module_name = model_config.__class__.__module__.replace(config_name, modeling_name)
+    parts = model_config.__class__.__module__.split(".")
+    parts[-1] = parts[-1].replace("configuration_", "modeling_", 1)
+    module_name = ".".join(parts)
     modeling_arch = importlib.import_module(module_name)
-    modeling_arch._get_unpad_data = get_unpad_data  # pylint: disable=protected-access
+    if hasattr(modeling_arch, "_get_unpad_data"):
+        modeling_arch._get_unpad_data = get_unpad_data
diff --git a/src/axolotl/monkeypatch/peft/__init__.py b/src/axolotl/monkeypatch/peft/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/peft/utils.py b/src/axolotl/monkeypatch/peft/utils.py
new file mode 100644
index 0000000000..d1011f5eb4
--- /dev/null
+++ b/src/axolotl/monkeypatch/peft/utils.py
@@ -0,0 +1,80 @@
+"""
+Patch prepare_model_for_kbit_training to not upcast everything
+"""
+
+import inspect
+
+import peft
+
+import axolotl
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+ORIGINAL_PREPARE_CODE = """
+        for param in model.parameters():
+            if (
+                (param.dtype == torch.float16) or (param.dtype == torch.bfloat16)
+            ) and param.__class__.__name__ != "Params4bit":
+                param.data = param.data.to(torch.float32)
+"""
+
+PATCHED_PREPARE_CODE = """
+        for name, param in model.named_parameters():
+            if (
+                (param.dtype == torch.float16) or (param.dtype == torch.bfloat16)
+            ) and param.__class__.__name__ != "Params4bit" and all(embed_name not in name for embed_name in ["embed_tokens", "lm_head"]):
+                param.data = param.data.to(torch.float32)
+"""
+
+
+def get_peft_prep_code() -> str:
+    prepare = inspect.getsource(peft.utils.other.prepare_model_for_kbit_training)
+    return prepare
+
+
+def check_peft_prep_code_is_patchable() -> bool:
+    prep_code = get_peft_prep_code()
+    prep_code, _ = detab_code(prep_code)
+    return ORIGINAL_PREPARE_CODE in prep_code
+
+
+def patch_peft_prep_code():
+    """
+    monkeypatch create_accelerator_and_postprocess so it checks for additional kwargs
+    """
+
+    try:
+        prep_code = get_peft_prep_code()
+    except OSError:
+        return
+    peft.utils.other._original_create_accelerator_and_postprocess = prep_code
+    prep_code, _ = detab_code(prep_code)
+    if ORIGINAL_PREPARE_CODE not in prep_code:
+        return
+
+    prep_code = prep_code.replace(ORIGINAL_PREPARE_CODE, PATCHED_PREPARE_CODE)
+    prep_code = prep_code.replace(
+        "def prepare_model_for_kbit_training(",
+        "def fixed_prepare_model_for_kbit_training(",
+        1,
+    )
+
+    items_to_import = []
+    for item in dir(peft.utils.other):
+        if item in prep_code:
+            items_to_import.append(item)
+
+    exec(
+        "from peft.utils.other import (" + ", ".join(x for x in items_to_import) + ")",
+        globals(),
+    )
+    exec(prep_code, globals())
+    LOG.info("patching prepare_model_for_kbit_training to allow for overrides")
+    peft.utils.other.prepare_model_for_kbit_training = (
+        fixed_prepare_model_for_kbit_training
+    )
+    axolotl.loaders.model.prepare_model_for_kbit_training = (
+        fixed_prepare_model_for_kbit_training
+    )
diff --git a/src/axolotl/monkeypatch/relora.py b/src/axolotl/monkeypatch/relora.py
index e4352cbe3d..3d33ab204d 100644
--- a/src/axolotl/monkeypatch/relora.py
+++ b/src/axolotl/monkeypatch/relora.py
@@ -1,21 +1,18 @@
 """Implements the ReLoRA training procedure from https://arxiv.org/abs/2307.05695, minus the initial full fine-tune."""
+
 import glob
 import json
-import logging
 import os.path
 import shutil
 from functools import partial
 from pathlib import Path
-from typing import Dict, List, Sequence, Union
+from typing import Dict, List, Literal, Union
 
-import bitsandbytes as bnb
 import peft
 import safetensors.torch as st
 import torch
 from huggingface_hub import snapshot_download
 from torch.distributed.optim import ZeroRedundancyOptimizer
-from torch.optim.lr_scheduler import LRScheduler
-from torch.optim.optimizer import Optimizer
 from transformers import (
     TrainerCallback,
     TrainerControl,
@@ -26,12 +23,19 @@
 
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.distributed import barrier, is_main_process
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
 
-LOG = logging.getLogger("axolotl.relora")
+try:
+    import bitsandbytes as bnb
+except ImportError:  # pragma: no cover - optional dependency for 8-bit merge paths
+    bnb = None
 
 
 @torch.no_grad()
 def magnitude_pruning_(tensor, prune_ratio):
+    """Zero the lowest ``prune_ratio`` fraction of values by absolute magnitude, in place."""
     tensor_magnitude = torch.abs(tensor)
     threshold = torch.quantile(
         tensor_magnitude.flatten().to(dtype=torch.float32), prune_ratio
@@ -41,14 +45,43 @@ def magnitude_pruning_(tensor, prune_ratio):
     tensor.mul_(mask.to(dtype=tensor.dtype))
 
 
+@torch.no_grad()
+def random_pruning_(tensor, prune_ratio):
+    """Zero a random ``prune_ratio`` fraction of values, in place."""
+    mask = (
+        torch.rand(tensor.shape, dtype=torch.float32, device=tensor.device)
+        > prune_ratio
+    )
+    tensor.mul_(mask.to(dtype=tensor.dtype))
+
+
+# 0.999 mirrors the reference implementation. True zeroing breaks
+# ZeroRedundancyOptimizer.consolidate_state_dict; see Guitaricet/relora's
+# peft_pretraining/training_utils.py for the original note on this.
+_FULL_RESET_RATIO = 0.999
+
+
 def reset_optimizer(
     optimizer: torch.optim.Optimizer,
     *,
-    reset_params: list[str],  # where str is the key to a torch.nn.Parameter
-    optimizer_state_keys: list[str],
+    reset_params: List[torch.nn.Parameter],
+    optimizer_state_keys: List[str],
+    prune_method: Literal["magnitude", "random", "reset"] = "magnitude",
     prune_ratio: float = 0.9,
 ):
-    pruning_fn = partial(magnitude_pruning_, prune_ratio=prune_ratio)
+    """Prune optimizer state for ``reset_params`` only."""
+    if prune_method == "magnitude":
+        pruning_fn = partial(magnitude_pruning_, prune_ratio=prune_ratio)
+    elif prune_method in ("random", "reset"):
+        # "reset" is random pruning at a near-full ratio; the caller is responsible
+        # for supplying the appropriate prune_ratio (see ReLoRACallback.on_step_begin).
+        pruning_fn = partial(random_pruning_, prune_ratio=prune_ratio)
+    else:
+        raise ValueError(
+            f"Unknown prune_method {prune_method!r}; expected one of "
+            "'magnitude', 'random', 'reset'"
+        )
+
     n_zeros = 0
     n_total = 0
 
@@ -57,15 +90,21 @@ def reset_optimizer(
         optimizer_state = optimizer.optim.state
 
     for param in reset_params:
-        param_state = optimizer_state[param]
-        if len(param_state) == 0:  # no state for this param, happens for ZeRo optimizer
+        state = optimizer_state.get(param, {})
+        if not state:
             continue
         for key in optimizer_state_keys:
-            pruning_fn(
-                param_state[key]
-            )  # pruning fn has to be inplace to keep the same keys in the dict
-            n_total += param_state[key].numel()
-            n_zeros += torch.sum(param_state[key] == 0).item()
+            value = state.get(key)
+            if value is None or not torch.is_tensor(value):
+                continue
+            try:
+                pruning_fn(value)
+                n_total += value.numel()
+                n_zeros += torch.sum(value == 0).item()
+            except RuntimeError as exc:
+                if "quantile() input tensor is too large" in str(exc):
+                    continue
+                raise
 
     _zeroed = n_zeros / (1e-7 + n_total) * 100
     LOG.info(f"Percent of optimizer states zeroed: {_zeroed:.2f}")
@@ -76,18 +115,19 @@ class ReLoRACallback(TrainerCallback):
     """Callback to merge LoRA weights into the base model and save full-weight checkpoints"""
 
     def __init__(self, cfg: DictDefault):
-        self.relora_steps = cfg.relora_steps
+        self.jagged_restart_steps = cfg.jagged_restart_steps
         self.cpu_offload = cfg.relora_cpu_offload
         self.quantized = cfg.load_in_4bit or cfg.load_in_8bit
         self.last_full_model = cfg.base_model
         self.resume_from_checkpoint = cfg.resume_from_checkpoint
+        self.prune_method = cfg.relora_prune_method or "magnitude"
 
         if not os.path.exists(self.last_full_model):
             self.last_full_model = str(Path(snapshot_download(cfg.base_model)))
 
-        assert os.path.exists(
-            self.last_full_model
-        ), "for ReLORA base_model must be a local path"
+        assert os.path.exists(self.last_full_model), (
+            "for ReLORA base_model must be a local path"
+        )
 
         self.num_lora_restarts = 0
         self.need_full_save = False
@@ -120,7 +160,9 @@ def on_step_begin(
         optimizer: torch.optim.Optimizer,
         **_kwargs,
     ):
-        if state.global_step > 0 and state.global_step % self.relora_steps == 0:
+        if not optimizer:
+            optimizer = state.optimizer
+        if state.global_step > 0 and state.global_step % self.jagged_restart_steps == 0:
             checkpoint_folder = os.path.join(
                 args.output_dir,
                 f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
@@ -129,11 +171,14 @@ def on_step_begin(
 
             if "adam" in args.optim.lower():
                 optimizer_state_keys = ["exp_avg", "exp_avg_sq"]
+                if "8bit" in args.optim.lower():
+                    optimizer_state_keys.append("state1")
+                    optimizer_state_keys.append("state2")
             else:
                 raise ValueError(f"Optimizer {args.optim} not supported with ReLoRA")
 
             lora_params = [
-                n
+                p
                 for n, p in model.named_parameters()
                 if p.requires_grad and "lora_" in n
             ]
@@ -144,7 +189,6 @@ def on_step_begin(
                     f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
                     "adapter",
                 ),
-                safe_serialization=True,
             )
             with torch.no_grad():
                 merge_and_save(
@@ -156,11 +200,19 @@ def on_step_begin(
                     actually_save=is_main_process(),
                     cpu_offload=self.cpu_offload,
                 )
+                # When relora_prune_ratio is not set, use _FULL_RESET_RATIO for
+                # "reset" (paper-style near-full reset) and 0.9 for other methods.
+                prune_ratio = args.relora_prune_ratio
+                if prune_ratio is None:
+                    prune_ratio = (
+                        _FULL_RESET_RATIO if self.prune_method == "reset" else 0.9
+                    )
                 reset_optimizer(
                     optimizer,
                     reset_params=lora_params,
                     optimizer_state_keys=optimizer_state_keys,
-                    prune_ratio=args.relora_prune_ratio,
+                    prune_method=self.prune_method,
+                    prune_ratio=prune_ratio,
                 )
 
             if self.quantized:
@@ -181,8 +233,8 @@ def on_save(
             args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}", "relora"
         )
         if (
-            state.global_step >= self.relora_steps
-            and state.global_step % self.relora_steps != 0
+            state.global_step >= self.jagged_restart_steps
+            and state.global_step % self.jagged_restart_steps != 0
         ):
             if self.quantized:
                 if is_main_process() and self.last_full_model != checkpoint_folder:
@@ -203,7 +255,7 @@ def on_save(
 
                     self.last_full_model = checkpoint_folder
             else:
-                model.model.save_pretrained(checkpoint_folder, safe_serialization=True)
+                model.model.save_pretrained(checkpoint_folder)
 
         return control
 
@@ -242,51 +294,6 @@ def on_train_end(
         return control
 
 
-class ReLoRAScheduler(LRScheduler):
-    """Wraps another scheduler to apply per-lora-restart learning rate warmups."""
-
-    def __init__(
-        self,
-        optimizer: Optimizer,
-        inner_schedule: LRScheduler,
-        relora_steps: int,
-        warmup_steps: int,
-        anneal_steps: int = 1,
-        min_lr_scale: float = 0.001,
-    ) -> None:
-        self.inner_schedule = inner_schedule
-        self.relora_steps = relora_steps
-        self.warmup_steps = warmup_steps
-        self.anneal_steps = anneal_steps
-        self.min_lr_scale = min_lr_scale
-        super().__init__(optimizer, inner_schedule.last_epoch, inner_schedule.verbose)
-
-    def get_lr(self) -> float:
-        self.inner_schedule.last_epoch = self.last_epoch
-
-        original = self.inner_schedule.get_lr()
-        step = self.last_epoch
-
-        if step < self.relora_steps - self.warmup_steps:
-            scale = 1
-        else:
-            per_relora_progress = step % self.relora_steps
-            if per_relora_progress < self.warmup_steps:
-                cycle_t = min(1.0, (per_relora_progress) / self.warmup_steps)
-            elif per_relora_progress > (self.relora_steps - self.anneal_steps):
-                cycle_t = min(
-                    1.0,
-                    (self.relora_steps - per_relora_progress) / self.anneal_steps,
-                )
-            else:
-                cycle_t = 1
-            scale = cycle_t * (1 - self.min_lr_scale) + self.min_lr_scale
-
-        if isinstance(original, Sequence):
-            return [lr * scale for lr in original]
-        return original * scale
-
-
 def sharded_paths(path: str, module_names: List[str]) -> Dict[str, str]:
     model_name = "model.safetensors"
     if not os.path.exists(str(Path(path) / model_name)) and not os.path.exists(
@@ -327,7 +334,6 @@ def find_lora_modules(model: peft.LoraModel) -> Dict[str, peft.tuners.lora.LoraL
     key_list = [key for key, _ in model.model.named_modules() if "lora" not in key]
     for key in key_list:
         try:
-            # pylint: disable=protected-access
             _parent, target, _target_name = peft.utils._get_submodules(model.model, key)
         except AttributeError:
             continue
@@ -356,6 +362,8 @@ def update_weights(
         target.weight.data = new_weight.cpu()
         target.to(device)
     elif isinstance(target, peft.tuners.lora.Linear8bitLt):
+        if bnb is None:
+            raise ImportError("bitsandbytes is required to merge 8-bit LoRA weights")
         target.weight.data = (
             bnb.nn.Int8Params(new_weight, requires_grad=False).to(device).data
         )
@@ -375,7 +383,7 @@ def merge_and_save(
     modules = find_lora_modules(model)
 
     if not quantized:
-        for module_name, target in modules.items():
+        for _, target in modules.items():
             active_adapter = target.active_adapter
             if isinstance(active_adapter, list):
                 active_adapter = active_adapter[0]
@@ -399,7 +407,10 @@ def merge_and_save(
         if shard_path.endswith(".safetensors"):
             in_tensors = st.load_file(str(Path(model_src) / shard_path))
         else:
-            in_tensors = torch.load(Path(model_src) / shard_path)
+            in_tensors = torch.load(
+                Path(model_src) / shard_path,
+                weights_only=True,  # to prevent arbitrary code execution
+            )
             if "state_dict" in in_tensors:
                 in_tensors = in_tensors["state_dict"]
 
diff --git a/src/axolotl/monkeypatch/ring_attn/__init__.py b/src/axolotl/monkeypatch/ring_attn/__init__.py
new file mode 100644
index 0000000000..1c14776c9b
--- /dev/null
+++ b/src/axolotl/monkeypatch/ring_attn/__init__.py
@@ -0,0 +1,17 @@
+"""Init for ring attention monkeypatch module"""
+
+# flake8: noqa
+
+from .patch import (
+    get_ring_attn_group,
+    register_ring_attn_from_device_mesh,
+    set_ring_attn_group,
+    update_ring_attn_params,
+)
+
+__all__ = (
+    "get_ring_attn_group",
+    "register_ring_attn_from_device_mesh",
+    "set_ring_attn_group",
+    "update_ring_attn_params",
+)
diff --git a/src/axolotl/monkeypatch/ring_attn/adapters/__init__.py b/src/axolotl/monkeypatch/ring_attn/adapters/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/ring_attn/adapters/batch.py b/src/axolotl/monkeypatch/ring_attn/adapters/batch.py
new file mode 100644
index 0000000000..74d33ed4ab
--- /dev/null
+++ b/src/axolotl/monkeypatch/ring_attn/adapters/batch.py
@@ -0,0 +1,196 @@
+"""
+HuggingFace flash attention adapter for basic ring attention (batch API).
+
+Inspired by
+https://github.com/zhuzilin/ring-flash-attention/blob/ce9fd3935ca0e5f0592bb0826cbed18ec69da729/ring_flash_attn/adapters/hf_adapter.py.
+Our implementation closely follows the structure of that module, but we've minified it
+somewhat to support only the latest versions of transformers.
+"""
+
+import os
+from typing import Callable
+
+import torch
+import torch.distributed as dist
+import transformers
+import transformers.modeling_flash_attention_utils
+from ring_flash_attn import ring_flash_attn_func
+from ring_flash_attn.adapters.hf_adapter import check_params
+from transformers.modeling_flash_attention_utils import is_flash_attn_greater_or_equal
+
+try:
+    from transformers.modeling_flash_attention_utils import _flash_supports_window
+except ImportError:
+    try:
+        from transformers.modeling_flash_attention_utils import (
+            _flash_supports_window_size as _flash_supports_window,
+        )
+    except ImportError:
+        _flash_supports_window = True
+
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+from axolotl.utils.schemas.enums import RingAttnFunc
+
+RING_ATTN_FUNC_MAPPING = {
+    RingAttnFunc.BATCH_RING: torch.compile(ring_flash_attn_func),
+    # RingAttnFunc.BATCH_ZIGZAG: torch.compile(zigzag_ring_flash_attn_func),
+    # RingAttnFunc.BATCH_STRIPE: torch.compile(stripe_flash_attn_func),
+}
+
+
+def create_flash_attn_forward_varlen_llama3(
+    process_group: dist.ProcessGroup, ring_attn_func: RingAttnFunc
+) -> Callable:
+    """
+    Create a ring flash attention forward function compatible with HuggingFace's
+    interface.
+
+    Args:
+        process_group: A PyTorch distributed process group.
+        ring_attn_func: Function from `ring_flash_attention` to replace HF flash
+            attention with.
+
+    Returns:
+        A function that implements the ring flash attention forward pass with the
+            signature expected by HuggingFace Transformers.
+    """
+
+    # transformers 4.48+
+
+    def _flash_attention_forward(
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        query_length: int,
+        is_causal: bool,
+        dropout: float = 0.0,
+        position_ids: torch.Tensor | None = None,
+        softmax_scale: float | None = None,
+        sliding_window: int | None = None,
+        use_top_left_mask: bool = False,
+        softcap: float | None = None,
+        deterministic: bool = None,
+        cu_seq_lens_q: torch.LongTensor | None = None,
+        cu_seq_lens_k: torch.LongTensor | None = None,
+        max_length_q: int | None = None,
+        max_length_k: int | None = None,
+        target_dtype: torch.dtype | None = None,
+        attn_implementation: str | None = None,
+        **kwargs,
+    ):
+        """
+        Calls the forward method of Ring Flash Attention.
+
+        Args:
+            query_states: Tensor containing the query vectors.
+            key_states: Tensor containing the key vectors.
+            value_states: Tensor containing the value vectors.
+            attention_mask: Not used in this implementation.
+            query_length: Integer representing the length of the query sequence.
+            is_causal: Boolean indicating whether to apply a causal mask to the attention.
+            dropout: Float representing the dropout probability. Default is 0.0.
+            position_ids: Not used in this implementation.
+            softmax_scale: Optional float value for the softmax scaling factor. Default is None.
+            sliding_window: Optional integer defining the size of the sliding attention window.
+                Default is None.
+            use_top_left_mask: Boolean indicating whether to use a top-left mask for the attention.
+                Default is False.
+            softcap: Not used in this implementation.
+            deterministic: Optional boolean to enforce deterministic computation. Default is None.
+            cu_seq_lens_q: Not used in this implementation.
+            cu_seq_lens_k: Not used in this implementation.
+            max_length_q: Not used in this implementation.
+            max_length_k: Not used in this implementation.
+            target_dtype: Not used in this implementation.
+            attn_implementation: Not used in this implementation.
+            **kwargs: Additional keyword arguments. Not used in this implementation.
+
+        Returns:
+            torch.Tensor: The output of the attention mechanism, with shape
+                `[batch_size, query_length, num_heads, head_dim]`.
+        """
+        if not use_top_left_mask:
+            causal = is_causal
+        else:
+            causal = is_causal and query_length != 1
+
+        # Handle sliding window
+        use_sliding_windows = (
+            _flash_supports_window
+            and sliding_window is not None
+            and key_states.shape[1] > sliding_window
+        )
+        window_size = (
+            (sliding_window, sliding_window) if use_sliding_windows else (-1, -1)
+        )
+
+        # Handle deterministic mode
+        if is_flash_attn_greater_or_equal("2.4.1"):
+            if deterministic is None:
+                deterministic = (
+                    os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1"
+                )
+
+        # Call ring flash attention function
+        attn_output = RING_ATTN_FUNC_MAPPING[ring_attn_func](
+            query_states,
+            key_states,
+            value_states,
+            dropout_p=dropout,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            alibi_slopes=None,
+            deterministic=deterministic,
+            return_attn_probs=False,
+            group=process_group,
+        )
+
+        return attn_output
+
+    return _flash_attention_forward
+
+
+def substitute_hf_flash_attn(
+    process_group: dist.ProcessGroup, ring_attn_func: RingAttnFunc
+):
+    """
+    Substitute HuggingFace's flash attention implementation with ring-based implementation.
+
+    Args:
+        process_group: PyTorch distributed process group for communication.
+        ring_attn_func: Function from `ring_flash_attention` to replace HF flash
+            attention with.
+    """
+    try:
+        # Substitute flash attention
+        old_flash_attention_forward = (
+            transformers.modeling_flash_attention_utils._flash_attention_forward
+        )
+        new_flash_attention_forward = create_flash_attn_forward_varlen_llama3(
+            process_group=process_group, ring_attn_func=ring_attn_func
+        )
+
+        if check_params(old_flash_attention_forward, new_flash_attention_forward):
+            transformers.modeling_flash_attention_utils._flash_attention_forward = (
+                new_flash_attention_forward
+            )
+        else:
+            raise ValueError(
+                "The signature of the new flash attention forward function does not match the old one."
+            )
+    except Exception as exception:
+        raise ValueError(
+            f"The current transformer version {transformers.__version__} is not supported. "
+            "Please use pip install -U transformers to upgrade to the latest version. "
+            "If the code failed with the latest version, "
+            f"please file an issue."
+        ) from exception
+
+    # Register with ALL_ATTENTION_FUNCTIONS if available
+    if ALL_ATTENTION_FUNCTIONS is not None:
+        from ring_flash_attn.adapters.hf_adapter import flash_attention_forward
+
+        ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = flash_attention_forward
diff --git a/src/axolotl/monkeypatch/ring_attn/patch.py b/src/axolotl/monkeypatch/ring_attn/patch.py
new file mode 100644
index 0000000000..c14abd65b2
--- /dev/null
+++ b/src/axolotl/monkeypatch/ring_attn/patch.py
@@ -0,0 +1,231 @@
+"""Ring attention group registration and flash attention patching.
+
+Make use of the `ring-flash-attn` (https://github.com/zhuzilin/ring-flash-attention)
+package, specifically the `hf_adapter.substitute_hf_flash_attn` function to patch in
+their sequence parallel version of Flash Attention 2.
+
+We also provide some patches for accelerate functions to prepare the dataloader for
+sequence parallelism training.
+"""
+
+import os
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from torch.distributed import DeviceMesh
+
+try:
+    from transformers.modeling_flash_attention_utils import _flash_supports_window
+except ImportError:
+    try:
+        from transformers.modeling_flash_attention_utils import (
+            _flash_supports_window_size as _flash_supports_window,
+        )
+    except ImportError:
+        _flash_supports_window = True
+
+from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.enums import RingAttnFunc
+
+LOG = get_logger(__name__)
+
+RING_ATTN_GROUP = None
+
+
+def get_ring_attn_group() -> dist.ProcessGroup:
+    """Getter for ring attention group on this rank."""
+    if RING_ATTN_GROUP is None:
+        raise RuntimeError("register_ring_attn_from_device_mesh() not yet called")
+    return RING_ATTN_GROUP
+
+
+def set_ring_attn_group(ring_attn_group: dist.ProcessGroup | None):
+    """Setter for ring attention group on this rank."""
+    global RING_ATTN_GROUP
+    RING_ATTN_GROUP = ring_attn_group
+
+
+def create_ring_flash_attention_forward(
+    process_group: dist.ProcessGroup, heads_k_stride: int
+):
+    from ring_flash_attn import llama3_flash_attn_varlen_func
+    from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS
+
+    def _flash_attention_forward_v3(
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        query_length: int,
+        is_causal: bool,
+        dropout: float = 0.0,
+        position_ids: Optional[torch.Tensor] = None,
+        softmax_scale: Optional[float] = None,
+        sliding_window: Optional[int] = None,
+        use_top_left_mask: bool = False,
+        softcap: Optional[float] = None,
+        deterministic: bool = None,
+        cu_seq_lens_q: Optional[torch.LongTensor] = None,
+        cu_seq_lens_k: Optional[torch.LongTensor] = None,
+        max_length_q: Optional[int] = None,
+        max_length_k: Optional[int] = None,
+        target_dtype: Optional[torch.dtype] = None,
+        attn_implementation: Optional[str] = None,
+        **kwargs,
+    ):
+        if not use_top_left_mask:
+            causal = is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__.
+            causal = is_causal and query_length != 1
+
+        # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length).
+        use_sliding_windows = (
+            _flash_supports_window
+            and sliding_window is not None
+            and key_states.shape[1] > sliding_window
+        )
+        flash_kwargs = (
+            {"window_size": (sliding_window, sliding_window)}
+            if use_sliding_windows
+            else {}
+        )
+
+        if deterministic is None:
+            deterministic = os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1"
+        flash_kwargs["deterministic"] = deterministic
+        assert softcap is None, (
+            "llama3_flash_attn_varlen_func does not support softcap yet."
+        )
+        # flash_kwargs["softcap"] = softcap
+        flash_kwargs["group"] = process_group
+
+        # not sure why attention_mask can be not None...
+        assert causal, "only causal attention is supported yet."
+        batch_size = query_states.size(0)
+        assert batch_size == 1, "varlen data should be processed in advance."
+
+        attn_output = llama3_flash_attn_varlen_func(
+            query_states.squeeze(dim=0),
+            key_states.squeeze(dim=0),
+            value_states.squeeze(dim=0),
+            cu_seqlens_q=DATA_PARAMS["cu_seqlens_q"],
+            cu_seqlens_k=DATA_PARAMS["cu_seqlens_k"],
+            max_seqlen_q=DATA_PARAMS["max_seqlen_q"],
+            max_seqlen_k=DATA_PARAMS["max_seqlen_k"],
+            heads_k_stride=heads_k_stride,
+            local_k_slice=DATA_PARAMS["local_k_slice"],
+            dropout_p=dropout,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            **flash_kwargs,
+        )
+
+        attn_output = attn_output.unsqueeze(dim=0)
+
+        return attn_output
+
+    return [
+        _flash_attention_forward_v3,
+    ]
+
+
+def register_ring_attn_from_device_mesh(
+    device_mesh: "DeviceMesh",
+    context_parallel_dim: tuple[str, ...],
+    heads_k_stride: int | None,
+    ring_attn_func: RingAttnFunc | None,
+):
+    """Create ring attention group using DeviceMesh and substitute flash attn with ring flash attn.
+
+    Args:
+        device_mesh: DeviceMesh object containing the parallelism topology.
+        context_parallel_dim: Name of the sequence parallel dimension in the device mesh.
+        heads_k_stride: Sequence parallelism K head stride size. Passed through to
+            `varlen_llama3` `ring_flash_attn` implementation.
+        ring_attn_func: `ring_flash_attn` ring attention implemention. If sample
+            packing is enabled, it must be a `varlen` function; otherwise, it must be a
+            `batch` function.
+    """
+    rank = dist.get_rank()
+
+    LOG.info(
+        f"Enabling ring attention sequence parallelism using DeviceMesh "
+        f"dimension '{context_parallel_dim}'",
+    )
+
+    # Extract the sequence parallel submesh
+    try:
+        sequence_mesh = device_mesh[context_parallel_dim]
+    except (KeyError, IndexError) as e:
+        raise ValueError(
+            f"Dimension '{context_parallel_dim}' not found in device_mesh. "
+            f"Available dimensions: {device_mesh.mesh_dim_names}"
+        ) from e
+
+    # Get the process group for context parallelism
+    sequence_pg = sequence_mesh.get_group()
+    context_parallel_size = sequence_mesh.size()
+
+    if rank == 0:
+        LOG.info(
+            f"Sequence parallel degree: {context_parallel_size}, "
+            f"mesh shape: {sequence_mesh.mesh.shape}"
+        )
+
+    # Log which ranks are in the current process group
+    if sequence_pg != dist.GroupMember.WORLD:
+        ranks_in_group = dist.get_process_group_ranks(sequence_pg)
+        LOG.info(f"Current sequence parallel group ranks: {ranks_in_group}")
+
+    # Set the ring attention group
+    set_ring_attn_group(sequence_pg)
+
+    if ring_attn_func is RingAttnFunc.VARLEN_LLAMA3:
+        # fmt: off
+        import ring_flash_attn.adapters.hf_adapter
+
+        from ring_flash_attn.adapters.hf_adapter import (  # isort: skip
+            create_ring_flash_attention_forward as create_ring_flash_attention_forward_orig,
+        )
+
+        create_ring_flash_attention_forward_orig = (  # noqa: F811,F841
+            create_ring_flash_attention_forward
+        )
+        ring_flash_attn.adapters.hf_adapter.create_ring_flash_attention_forward = create_ring_flash_attention_forward
+        # fmt: on
+
+        ring_flash_attn.adapters.hf_adapter.substitute_hf_flash_attn(
+            process_group=get_ring_attn_group(), heads_k_stride=heads_k_stride or 1
+        )
+    elif ring_attn_func is RingAttnFunc.BATCH_RING:
+        from axolotl.monkeypatch.ring_attn.adapters.batch import (
+            substitute_hf_flash_attn,
+        )
+
+        substitute_hf_flash_attn(
+            process_group=get_ring_attn_group(),
+            ring_attn_func=ring_attn_func,
+        )
+
+
+def update_ring_attn_params(position_ids: torch.Tensor | None):
+    """
+    Calculate the cumulative sequence lengths for the current forward pass and pass the
+    value to the substituted `ring_flash_attn`.
+
+    Args:
+        position_ids: Optional tensor of position IDs (for sample packed data).
+    """
+    try:
+        from ring_flash_attn import update_ring_flash_attn_params
+    except ImportError:
+        # No ring attention was substituted (e.g. the GLM DSA kernels own context-parallel attention
+        # and derive their own cu_seqlens from position_ids), so there is nothing to update here.
+        return
+
+    cu_seqlens, _ = get_cu_seqlens_from_pos_ids(position_ids)
+    cu_seqlens = cu_seqlens.squeeze().to(device=torch.cuda.current_device())
+    update_ring_flash_attn_params(cu_seqlens, get_ring_attn_group())
diff --git a/src/axolotl/monkeypatch/scaled_softmax_attn.py b/src/axolotl/monkeypatch/scaled_softmax_attn.py
new file mode 100644
index 0000000000..bb2ebb8b82
--- /dev/null
+++ b/src/axolotl/monkeypatch/scaled_softmax_attn.py
@@ -0,0 +1,141 @@
+"""
+Scaled Softmax (SSMax) attention patch using FlexAttention.
+SSMax:  softmax(scores * s * log(n) + b) where n is the position index
+Ref: https://arxiv.org/abs/2501.19399
+"""
+
+import torch
+from transformers import PreTrainedModel
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+try:
+    from torch.nn.attention.flex_attention import BlockMask
+    from transformers.integrations.flex_attention import (
+        compile_friendly_flex_attention,
+        repeat_kv,
+    )
+
+    FLEX_ATTENTION_AVAILABLE = True
+except ImportError:
+    FLEX_ATTENTION_AVAILABLE = False
+    BlockMask = None
+
+_ssmax_config = {}
+
+
+def patch_scaled_softmax_attention(
+    scaling_factor_init: float = 0.43, bias: float = 0.0, model: PreTrainedModel = None
+):
+    """Patch attention to apply SSMax via FlexAttention score_mod."""
+    global _ssmax_config
+
+    if not FLEX_ATTENTION_AVAILABLE:
+        raise RuntimeError("SSMax requires FlexAttention.")
+
+    _ssmax_config["ssmax_s"] = scaling_factor_init
+    _ssmax_config["ssmax_b"] = bias
+
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    if "flex_attention" in ALL_ATTENTION_FUNCTIONS:
+        _ssmax_config["original_flex_fn"] = ALL_ATTENTION_FUNCTIONS["flex_attention"]
+        ALL_ATTENTION_FUNCTIONS["flex_attention"] = ssmax_flex_attention_forward
+        LOG.info(
+            f"Patched flex_attention with SSMax (s={scaling_factor_init}, b={bias})"
+        )
+    else:
+        LOG.warning("flex_attention not found.  Ensure flex_attention:  true is set.")
+
+
+def ssmax_flex_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask,
+    scaling: float | None = None,
+    softcap: float | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """FlexAttention forward with SSMax:  score * (s * log(n) + b)."""
+
+    if kwargs.get("dropout", 0.0) > 0:
+        raise ValueError("flex_attention does not support dropout")
+
+    ssmax_s = _ssmax_config.get("ssmax_s", 0.43)
+    ssmax_b = _ssmax_config.get("ssmax_b", 0.0)
+
+    position_ids = kwargs.get("position_ids", None)
+    position_ids_flat = position_ids.view(-1) if position_ids is not None else None
+
+    block_mask = attention_mask if isinstance(attention_mask, BlockMask) else None
+    score_mask = None if block_mask else attention_mask
+
+    if score_mask is not None:
+        score_mask = score_mask[:, :, :, : key.shape[-2]]
+
+    def score_mod(score, batch_idx, head_idx, q_idx, kv_idx):
+        """
+        Apply SSMax scaling:  score * (s * log(n) + b)
+        where n is the relative position within each packed sequence.
+        """
+        if position_ids_flat is not None:
+            relative_pos = position_ids_flat[q_idx]
+            n = (relative_pos + 1).float()
+        else:
+            n = (q_idx + 1).float()
+
+        n = torch.clamp(n, min=2.0)
+
+        ssmax_scale = ssmax_s * torch.log(n) + ssmax_b
+        score = score * ssmax_scale
+
+        if softcap is not None:
+            score = softcap * torch.tanh(score / softcap)
+
+        if score_mask is not None:
+            score = score + score_mask[batch_idx][0][q_idx][kv_idx]
+
+        return score
+
+    enable_gqa = True
+    if (query.shape[1] & (query.shape[1] - 1)) != 0:
+        key = repeat_kv(key, query.shape[1] // key.shape[1])
+        value = repeat_kv(value, query.shape[1] // value.shape[1])
+        enable_gqa = False
+
+    return_lse = query.device.type != "cpu"
+    flex_output = compile_friendly_flex_attention(
+        query,
+        key,
+        value,
+        score_mod=score_mod,
+        block_mask=block_mask,
+        enable_gqa=enable_gqa,
+        scale=scaling,
+        kernel_options=kwargs.get("kernel_options"),
+        return_lse=return_lse,
+        training=module.training,
+    )
+
+    if return_lse:
+        attention_output, lse = flex_output
+        lse = lse.to(value.dtype)
+    else:
+        attention_output, lse = flex_output, None
+
+    return attention_output.transpose(1, 2).contiguous(), lse
+
+
+def unpatch_scaled_softmax_attention():
+    """Restore the original FlexAttention function."""
+    global _ssmax_config
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    if "original_flex_fn" in _ssmax_config:
+        ALL_ATTENTION_FUNCTIONS["flex_attention"] = _ssmax_config["original_flex_fn"]
+        _ssmax_config.clear()
+        LOG.info("Unpatched flex_attention, restored original")
diff --git a/src/axolotl/monkeypatch/selective_checkpointing.py b/src/axolotl/monkeypatch/selective_checkpointing.py
new file mode 100644
index 0000000000..025e5c525e
--- /dev/null
+++ b/src/axolotl/monkeypatch/selective_checkpointing.py
@@ -0,0 +1,330 @@
+"""Eager selective activation checkpointing (SAC): save chosen ops instead of recomputing.
+
+Uses ``torch.utils.checkpoint.create_selective_checkpoint_contexts`` (no torch.compile).
+Ops are matched at the dispatcher level, so custom kernels that are not registered as
+torch ops are simply recomputed as usual — only ops we want to *save* need to be
+dispatcher-visible.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+import torch
+from torch.utils.checkpoint import (
+    CheckpointPolicy,
+    create_selective_checkpoint_contexts,
+)
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+ATTENTION_GROUP = "attention"
+
+_ATEN_ATTENTION_PACKETS = (
+    "_scaled_dot_product_flash_attention",
+    "_scaled_dot_product_efficient_attention",
+    "_scaled_dot_product_cudnn_attention",
+    "_scaled_dot_product_flash_attention_for_cpu",
+)
+
+# regions to observe before warning that no op ever matched the save policy
+_NO_MATCH_WARN_REGIONS = 64
+
+
+def _aten_attention_ops() -> set:
+    ops = set()
+    for packet_name in _ATEN_ATTENTION_PACKETS:
+        packet = getattr(torch.ops.aten, packet_name, None)
+        if packet is not None:
+            ops.add(packet.default)
+    return ops
+
+
+def _op_name(op: Any) -> str:
+    try:
+        return op.name()
+    except (AttributeError, TypeError):
+        return getattr(op, "__name__", str(op))
+
+
+def _is_flash_attention_forward(name: str) -> bool:
+    lowered = name.lower()
+    if "flash_attn" not in lowered and "flash_attention" not in lowered:
+        return False
+    if "backward" in lowered or "bwd" in lowered:
+        return False
+    return True
+
+
+# a bounded LEFT context is what defines SWA; a bounded right side alone can
+# just encode causality, so only window_size_left is inspected
+_WINDOW_ARG_NAME = "window_size_left"
+_window_arg_cache: dict[Any, int | None] = {}
+
+
+def _window_arg_index(op: Any) -> int | None:
+    """Schema position of the sliding-window arg, or None if the op has none."""
+    if op in _window_arg_cache:
+        return _window_arg_cache[op]
+    index: int | None = None
+    schema = getattr(op, "_schema", None)
+    if schema is not None:
+        for i, arg in enumerate(schema.arguments):
+            if arg.name == _WINDOW_ARG_NAME:
+                index = i
+                break
+    _window_arg_cache[op] = index
+    return index
+
+
+def _is_sliding_window_call(op: Any, args: tuple, kwargs: dict) -> bool:
+    """True when a flash-attention call is bounded by a sliding window.
+
+    flash-attn uses -1 for an unbounded side, so a non-negative window arg
+    means SWA. SDPA carries the window only inside the attention mask, so
+    hybrid models running through SDPA cannot be discriminated here.
+    """
+    val = kwargs.get(_WINDOW_ARG_NAME)
+    if isinstance(val, int):
+        return val >= 0
+    index = _window_arg_index(op)
+    if index is not None and index < len(args) and isinstance(args[index], int):
+        return args[index] >= 0
+    return False
+
+
+DEFAULT_RECOMPUTE_LAYER_TYPES = ("sliding_attention", "chunked_attention")
+
+
+class SacPolicyState:
+    """Bookkeeping shared across checkpoint regions for logging/diagnostics."""
+
+    def __init__(self) -> None:
+        self.saved_op_names: set[str] = set()
+        self.sliding_op_names: set[str] = set()
+        self.regions_seen: int = 0
+        self.warned_no_match: bool = False
+        # published by decoder-layer hooks; read by the policy. Hooks fire again
+        # during checkpoint recompute (on the autograd thread), so forward and
+        # replay see the same value.
+        self.current_layer_type: str | None = None
+
+
+def _layer_attention_type(module) -> str | None:
+    for obj in (
+        module,
+        getattr(module, "self_attn", None),
+        getattr(module, "attention", None),
+    ):
+        if obj is None:
+            continue
+        layer_type = getattr(obj, "layer_type", None)
+        if isinstance(layer_type, str):
+            return layer_type
+        if getattr(obj, "is_sliding", None) is True:
+            return "sliding_attention"
+    return None
+
+
+def install_layer_type_hooks(model, state: SacPolicyState) -> int:
+    """Publish each checkpointed decoder layer's attention type while it runs.
+
+    Lets the policy skip saving sliding/chunked-window attention in hybrid
+    models even under SDPA, where the window lives in the mask and cannot be
+    read off the op's arguments.
+    """
+    from transformers import GradientCheckpointingLayer
+
+    hooked = 0
+    if not hasattr(model, "modules"):
+        return hooked
+    for module in model.modules():
+        if not isinstance(module, GradientCheckpointingLayer):
+            continue
+        layer_type = _layer_attention_type(module)
+        if layer_type is None:
+            continue
+
+        def _set(mod, args, kwargs=None, _lt=layer_type):
+            state.current_layer_type = _lt
+
+        def _clear(mod, args, output):
+            state.current_layer_type = None
+
+        module.register_forward_pre_hook(_set)
+        module.register_forward_hook(_clear, always_call=True)
+        hooked += 1
+    if hooked:
+        LOG.info(
+            f"selective_checkpointing: layer-type hooks on {hooked} decoder layers "
+            "(sliding/chunked-window attention will be recomputed)"
+        )
+    return hooked
+
+
+def build_sac_policy(
+    save: list[str] | None = None,
+    state: SacPolicyState | None = None,
+    save_sliding_window: bool = False,
+    recompute_layer_types: list[str] | None = None,
+) -> Callable:
+    """Build an eager SAC policy_fn: MUST_SAVE for matching ops, PREFER_RECOMPUTE otherwise.
+
+    ``save`` entries are either the ``"attention"`` group or substrings matched
+    against the qualified op name (e.g. ``"aten::mm"``). Unless
+    ``save_sliding_window`` is set, hybrid-model attention calls bounded by a
+    sliding window keep being recomputed — SWA is cheap and not worth the saved
+    memory; only full-attention calls are saved.
+    """
+    save = save or [ATTENTION_GROUP]
+    state = state or SacPolicyState()
+    skip_layer_types = (
+        set()
+        if save_sliding_window
+        else set(
+            DEFAULT_RECOMPUTE_LAYER_TYPES
+            if recompute_layer_types is None
+            else recompute_layer_types
+        )
+    )
+
+    exact_ops: set = set()
+    substrings: list[str] = []
+    match_flash_attention = False
+    for spec in save:
+        if spec == ATTENTION_GROUP:
+            exact_ops |= _aten_attention_ops()
+            match_flash_attention = True
+        else:
+            substrings.append(spec)
+
+    def _matches(op: Any) -> bool:
+        if op in exact_ops:
+            return True
+        name = _op_name(op)
+        if match_flash_attention and _is_flash_attention_forward(name):
+            return True
+        return any(sub in name for sub in substrings)
+
+    def policy_fn(ctx, op, *args, **kwargs):  # pylint: disable=unused-argument
+        if _matches(op):
+            name = _op_name(op)
+            if state.current_layer_type in skip_layer_types:
+                if name not in state.sliding_op_names:
+                    state.sliding_op_names.add(name)
+                    LOG.info(
+                        f"selective_checkpointing: recomputing `{name}` in "
+                        f"{state.current_layer_type} layers"
+                    )
+                return CheckpointPolicy.PREFER_RECOMPUTE
+            if not save_sliding_window and _is_sliding_window_call(op, args, kwargs):
+                if name not in state.sliding_op_names:
+                    state.sliding_op_names.add(name)
+                    LOG.info(
+                        f"selective_checkpointing: recomputing sliding-window "
+                        f"calls of `{name}` (save_sliding_window: false)"
+                    )
+                return CheckpointPolicy.PREFER_RECOMPUTE
+            if name not in state.saved_op_names:
+                state.saved_op_names.add(name)
+                LOG.info(
+                    f"selective_checkpointing: saving `{name}` "
+                    "(backward will not recompute it)"
+                )
+            return CheckpointPolicy.MUST_SAVE
+        return CheckpointPolicy.PREFER_RECOMPUTE
+
+    return policy_fn
+
+
+def build_sac_context_fn(
+    save: list[str] | None = None,
+    save_sliding_window: bool = False,
+    state: SacPolicyState | None = None,
+    recompute_layer_types: list[str] | None = None,
+) -> Callable:
+    """Return a ``context_fn`` for ``torch.utils.checkpoint.checkpoint``."""
+    state = state or SacPolicyState()
+    policy_fn = build_sac_policy(
+        save, state, save_sliding_window, recompute_layer_types
+    )
+
+    def context_fn():
+        state.regions_seen += 1
+        if (
+            not state.saved_op_names
+            and not state.warned_no_match
+            and state.regions_seen >= _NO_MATCH_WARN_REGIONS
+        ):
+            state.warned_no_match = True
+            LOG.warning(
+                f"selective_checkpointing: no op matched the save policy after "
+                f"{state.regions_seen} checkpoint regions. Your attention "
+                "implementation may not be dispatcher-visible (e.g. a custom "
+                "kernel not registered via torch.library); everything is being "
+                "recomputed as with plain gradient checkpointing."
+            )
+        return create_selective_checkpoint_contexts(policy_fn)
+
+    return context_fn
+
+
+def apply_selective_checkpointing(
+    model,
+    save: list[str] | None = None,
+    save_sliding_window: bool = False,
+    recompute_layer_types: list[str] | None = None,
+    offload: bool = False,
+) -> None:
+    """Wrap ``model.gradient_checkpointing_enable`` to inject the SAC ``context_fn``.
+
+    Wrapping the instance method covers every enable call site (axolotl's model
+    loader, HF Trainer at train() time, PEFT's kbit prep) without placing a
+    non-serializable callable into ``TrainingArguments``.
+    """
+    if getattr(model.gradient_checkpointing_enable, "_axolotl_sac", False):
+        return
+
+    state = SacPolicyState()
+    skip_layer_types = (
+        set()
+        if save_sliding_window
+        else set(
+            DEFAULT_RECOMPUTE_LAYER_TYPES
+            if recompute_layer_types is None
+            else recompute_layer_types
+        )
+    )
+    if skip_layer_types:
+        install_layer_type_hooks(model, state)
+    if offload:
+        from axolotl.monkeypatch.selective_checkpointing_offload import (
+            build_sac_offload_context_fn,
+        )
+
+        context_fn = build_sac_offload_context_fn(
+            save,
+            save_sliding_window,
+            state,
+            recompute_layer_types=recompute_layer_types,
+        )
+    else:
+        context_fn = build_sac_context_fn(
+            save, save_sliding_window, state, recompute_layer_types
+        )
+    orig_enable = model.gradient_checkpointing_enable
+
+    def enable_with_sac(gradient_checkpointing_kwargs=None, **kwargs):
+        gc_kwargs = dict(gradient_checkpointing_kwargs or {})
+        gc_kwargs["use_reentrant"] = False
+        gc_kwargs["context_fn"] = context_fn
+        return orig_enable(gradient_checkpointing_kwargs=gc_kwargs, **kwargs)
+
+    enable_with_sac._axolotl_sac = True
+    model.gradient_checkpointing_enable = enable_with_sac
+    LOG.info(
+        "selective_checkpointing enabled: "
+        f"save={save or [ATTENTION_GROUP]} (eager SAC, non-reentrant)"
+    )
diff --git a/src/axolotl/monkeypatch/selective_checkpointing_offload.py b/src/axolotl/monkeypatch/selective_checkpointing_offload.py
new file mode 100644
index 0000000000..466321d952
--- /dev/null
+++ b/src/axolotl/monkeypatch/selective_checkpointing_offload.py
@@ -0,0 +1,347 @@
+"""Phase-2 selective checkpointing: CPU-offload the SAC-saved tensors.
+
+Replaces torch's SAC caching/cached dispatch modes with variants that move
+MUST_SAVE outputs to pinned CPU memory on a side stream during forward and
+stream them back during backward with one-region-lookahead prefetch. Saved
+attention outputs then cost ~zero GPU memory.
+
+Mutation caveat: an offloaded tensor is a snapshot at pack time; in-place
+mutation of the original after the save cannot be detected (torch's version
+guard only covers non-offloaded leaves). The known mutating case — PEFT's
+in-place adapter add on matmul outputs — is rejected at config validation.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+from torch.utils._pytree import tree_map
+from torch.utils.checkpoint import (
+    SAC_IGNORED_OPS,
+    CheckpointPolicy,
+    SelectiveCheckpointContext,
+    _maybe_detach,
+    _policy_from_bool,
+    _VersionWrapper,
+)
+
+from axolotl.monkeypatch.selective_checkpointing import (
+    SacPolicyState,
+    build_sac_policy,
+)
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+_SAVE_POLICIES = (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE)
+
+_BufferKey = tuple[tuple[int, ...], tuple[int, ...], torch.dtype, torch.layout]
+
+
+@dataclass
+class _OffloadRef:
+    region_id: int
+    device: torch.device
+    size: torch.Size
+    stride: tuple[int, ...]
+    storage_offset: int
+    buffer_key: _BufferKey
+    cpu_tensor: torch.Tensor | None = None
+    gpu_tensor: torch.Tensor | None = None
+    restore_event: torch.cuda.Event | None = None
+
+
+@dataclass
+class SacOffloadStats:
+    offloaded_tensors: int = 0
+    offloaded_bytes: int = 0
+    restored_tensors: int = 0
+    logged: bool = field(default=False, repr=False)
+
+
+class SacOffloadEngine:
+    """Pinned-pool, side-stream D2H/H2D engine for SAC-saved tensors."""
+
+    def __init__(
+        self,
+        min_offload_bytes: int = 1 << 20,
+        max_fwd_stash_size: int = 2,
+        max_cpu_buffer_pool_size: int = 128,
+    ) -> None:
+        self.min_offload_bytes = min_offload_bytes
+        self.max_fwd_stash_size = max_fwd_stash_size
+        self.max_cpu_buffer_pool_size = max_cpu_buffer_pool_size
+        self.stats = SacOffloadStats()
+
+        self.s0 = torch.cuda.current_stream() if torch.cuda.is_available() else None
+        self.s1 = torch.cuda.Stream() if torch.cuda.is_available() else None
+
+        self._pack_seq = 0
+        self._fwd_stash: dict[int, tuple[torch.Tensor, torch.cuda.Event]] = {}
+        self._regions: dict[int, list[_OffloadRef]] = defaultdict(list)
+        self._cpu_pool: dict[_BufferKey, list[torch.Tensor]] = {}
+        self._cpu_pool_size = 0
+        self._pending_cpu_buffers: list[
+            tuple[_BufferKey, torch.Tensor, torch.cuda.Event]
+        ] = []
+
+    @property
+    def compute_stream(self):
+        return torch.cuda.current_stream()
+
+    def should_offload(self, tensor: torch.Tensor) -> bool:
+        return (
+            tensor.device.type == "cuda"
+            and tensor.element_size() * tensor.nelement() >= self.min_offload_bytes
+            and not isinstance(tensor, torch.nn.Parameter)
+        )
+
+    def _buffer_key(self, tensor: torch.Tensor) -> _BufferKey:
+        return (
+            tuple(tensor.size()),
+            tuple(tensor.stride()),
+            tensor.dtype,
+            tensor.layout,
+        )
+
+    def _pool_cpu_buffer(self, key: _BufferKey, tensor: torch.Tensor) -> None:
+        if self._cpu_pool_size >= self.max_cpu_buffer_pool_size:
+            return
+        self._cpu_pool.setdefault(key, []).append(tensor)
+        self._cpu_pool_size += 1
+
+    def _reap_pending_cpu_buffers(self) -> None:
+        pending = self._pending_cpu_buffers
+        self._pending_cpu_buffers = []
+        for key, tensor, event in pending:
+            if event.query():
+                self._pool_cpu_buffer(key, tensor)
+            else:
+                self._pending_cpu_buffers.append((key, tensor, event))
+
+    def _empty_pinned_like(
+        self, tensor: torch.Tensor
+    ) -> tuple[torch.Tensor, _BufferKey]:
+        self._reap_pending_cpu_buffers()
+        key = self._buffer_key(tensor)
+        pool = self._cpu_pool.get(key)
+        if pool:
+            self._cpu_pool_size -= 1
+            return pool.pop(), key
+        return (
+            torch.empty_strided(
+                tuple(tensor.size()),
+                tuple(tensor.stride()),
+                dtype=tensor.dtype,
+                layout=tensor.layout,
+                device="cpu",
+                pin_memory=True,
+            ),
+            key,
+        )
+
+    def _reap_fwd_stash(self, new_seq: int) -> None:
+        compute = self.compute_stream
+        for seq in list(self._fwd_stash):
+            if seq > new_seq - self.max_fwd_stash_size:
+                continue
+            _, event = self._fwd_stash.pop(seq)
+            compute.wait_event(event)
+
+    def pack(self, tensor: torch.Tensor, region_id: int) -> _OffloadRef:
+        self._pack_seq += 1
+        self._reap_fwd_stash(self._pack_seq)
+
+        compute = self.compute_stream
+        self.s1.wait_stream(compute)
+        with torch.cuda.stream(self.s1):
+            cpu_tensor, key = self._empty_pinned_like(tensor)
+            cpu_tensor.copy_(tensor, non_blocking=True)
+        event = self.s1.record_event()
+        self._fwd_stash[self._pack_seq] = (tensor, event)
+
+        ref = _OffloadRef(
+            region_id=region_id,
+            device=tensor.device,
+            size=tensor.size(),
+            stride=tuple(tensor.stride()),
+            storage_offset=tensor.storage_offset(),
+            buffer_key=key,
+            cpu_tensor=cpu_tensor,
+        )
+        self._regions[region_id].append(ref)
+        self.stats.offloaded_tensors += 1
+        self.stats.offloaded_bytes += tensor.element_size() * tensor.nelement()
+        return ref
+
+    def _start_restore(self, ref: _OffloadRef) -> None:
+        if ref.gpu_tensor is not None or ref.cpu_tensor is None:
+            return
+        with torch.cuda.stream(self.s1):
+            gpu_tensor = ref.cpu_tensor.to(ref.device, non_blocking=True)
+            if (
+                gpu_tensor.storage_offset() != ref.storage_offset
+                or gpu_tensor.stride() != ref.stride
+            ):
+                gpu_tensor = torch.as_strided(
+                    gpu_tensor, ref.size, ref.stride, ref.storage_offset
+                )
+        ref.gpu_tensor = gpu_tensor
+        ref.restore_event = self.s1.record_event()
+
+    def prefetch_region(self, region_id: int) -> None:
+        for ref in self._regions.get(region_id, []):
+            self._start_restore(ref)
+
+    def restore(self, ref: _OffloadRef) -> torch.Tensor:
+        if ref.gpu_tensor is None:
+            self._start_restore(ref)
+        compute = self.compute_stream
+        compute.wait_event(ref.restore_event)
+        gpu_tensor = ref.gpu_tensor
+        gpu_tensor.record_stream(compute)
+
+        self._pending_cpu_buffers.append(
+            (ref.buffer_key, ref.cpu_tensor, ref.restore_event)
+        )
+        ref.cpu_tensor = None
+        ref.gpu_tensor = None
+
+        region = self._regions.get(ref.region_id)
+        if region is not None:
+            try:
+                region.remove(ref)
+            except ValueError:
+                pass
+            if not region:
+                self._regions.pop(ref.region_id, None)
+
+        self.stats.restored_tensors += 1
+        return gpu_tensor
+
+    def log_once(self) -> None:
+        if self.stats.logged or not self.stats.offloaded_tensors:
+            return
+        self.stats.logged = True
+        LOG.info(
+            "selective_checkpointing offload: "
+            f"{self.stats.offloaded_tensors} tensors "
+            f"({self.stats.offloaded_bytes / 2**30:.2f} GiB) offloaded to CPU "
+            "in the first forward"
+        )
+
+
+class _OffloadCachingMode(TorchDispatchMode):
+    def __init__(self, policy_fn, storage, engine, region_id):
+        self.policy_fn = policy_fn
+        self.storage = storage
+        self.engine = engine
+        self.region_id = region_id
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        kwargs = {} if kwargs is None else kwargs
+        if func in SAC_IGNORED_OPS:
+            return func(*args, **kwargs)
+
+        out = func(*args, **kwargs)
+
+        if isinstance(func, torch._ops.HigherOrderOperator):
+            any_ret_has_alias_info = False
+        else:
+            any_ret_has_alias_info = any(
+                ret.alias_info is not None for ret in func._schema.returns
+            )
+
+        policy = self.policy_fn(
+            SelectiveCheckpointContext(is_recompute=False), func, *args, **kwargs
+        )
+        if isinstance(policy, bool):
+            policy = _policy_from_bool(policy)
+
+        if policy in _SAVE_POLICIES:
+
+            def pack_leaf(leaf):
+                detached = _maybe_detach(leaf, any_ret_has_alias_info)
+                if torch.is_tensor(detached) and self.engine.should_offload(detached):
+                    return self.engine.pack(detached, self.region_id)
+                return _VersionWrapper(detached)
+
+            self.storage[func].append(tree_map(pack_leaf, out))
+        return out
+
+
+class _OffloadCachedMode(TorchDispatchMode):
+    def __init__(self, policy_fn, storage, engine, region_id):
+        self.policy_fn = policy_fn
+        self.storage = storage
+        self.engine = engine
+        self.region_id = region_id
+
+    def __enter__(self):
+        # recompute of this region is starting: bring its tensors back and
+        # prefetch the next region backward will need
+        self.engine.prefetch_region(self.region_id)
+        if self.region_id > 0:
+            self.engine.prefetch_region(self.region_id - 1)
+        self.engine.log_once()
+        return super().__enter__()
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        kwargs = {} if kwargs is None else kwargs
+        if func in SAC_IGNORED_OPS:
+            return func(*args, **kwargs)
+
+        policy = self.policy_fn(
+            SelectiveCheckpointContext(is_recompute=True), func, *args, **kwargs
+        )
+        if isinstance(policy, bool):
+            policy = _policy_from_bool(policy)
+
+        if policy in _SAVE_POLICIES:
+            entries = self.storage.get(func)
+            if not entries:
+                raise RuntimeError(
+                    f"selective_checkpointing offload: {func} encountered during "
+                    "recompute but nothing was saved for it (or an extra backward "
+                    "was attempted)"
+                )
+
+            def unpack_leaf(leaf):
+                if isinstance(leaf, _OffloadRef):
+                    return self.engine.restore(leaf)
+                if isinstance(leaf, _VersionWrapper):
+                    return leaf.get_val(False)
+                return leaf
+
+            return tree_map(unpack_leaf, entries.pop(0))
+        return func(*args, **kwargs)
+
+
+def build_sac_offload_context_fn(
+    save: list[str] | None = None,
+    save_sliding_window: bool = False,
+    state: SacPolicyState | None = None,
+    engine: SacOffloadEngine | None = None,
+    recompute_layer_types: list[str] | None = None,
+) -> Callable:
+    """Return a ``context_fn`` whose MUST_SAVE tensors are offloaded to CPU."""
+    state = state or SacPolicyState()
+    engine = engine or SacOffloadEngine()
+    policy_fn = build_sac_policy(
+        save, state, save_sliding_window, recompute_layer_types
+    )
+
+    def context_fn():
+        region_id = state.regions_seen
+        state.regions_seen += 1
+        storage: dict[Any, list[Any]] = defaultdict(list)
+        return (
+            _OffloadCachingMode(policy_fn, storage, engine, region_id),
+            _OffloadCachedMode(policy_fn, storage, engine, region_id),
+        )
+
+    return context_fn
diff --git a/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py b/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
index 0269f90157..0fa6d6424a 100644
--- a/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
@@ -16,7 +16,8 @@
 # This code is based off the following work:
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
-""" PyTorch StableLM Epoch model. """
+"""PyTorch StableLM Epoch model."""
+
 import importlib
 import math
 from typing import Optional, Tuple, Union
@@ -25,17 +26,17 @@
 import torch.utils.checkpoint
 from accelerate import init_empty_weights
 from einops import rearrange
-from flash_attn.flash_attn_interface import (  # pylint: disable=ungrouped-imports
+from flash_attn.flash_attn_interface import (
     flash_attn_varlen_qkvpacked_func,
 )
 from torch import nn
 from transformers import AutoConfig, AutoModelForCausalLM
 from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.utils import logging
 
 from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
+from axolotl.utils.logging import get_logger
 
-logger = logging.get_logger(__name__)
+logger = get_logger(__name__)
 
 
 def replace_stablelm_attn_with_flash_attn(model_name="stabilityai/stablelm-3b-4e1t"):
@@ -48,27 +49,21 @@ def replace_stablelm_attn_with_flash_attn(model_name="stabilityai/stablelm-3b-4e
         ".configuration_stablelm_epoch", ".modeling_stablelm_epoch"
     )
     modeling_stablelm = importlib.import_module(module_name)
-    modeling_stablelm.Attention.forward = (  # pylint: disable=protected-access
-        flashattn_attn
-    )
-    modeling_stablelm.StableLMEpochModel.forward = (  # pylint: disable=protected-access
-        stablelm_model_forward
-    )
-    modeling_stablelm.DecoderLayer.forward = (  # pylint: disable=protected-access
-        decoder_layer_forward
-    )
+    modeling_stablelm.Attention.forward = flashattn_attn
+    modeling_stablelm.StableLMEpochModel.forward = stablelm_model_forward
+    modeling_stablelm.DecoderLayer.forward = decoder_layer_forward
 
 
 def rotate_half(x: torch.Tensor):
     """Rotates half the hidden dims of the input."""
-    # pylint: disable=invalid-name
+
     x1, x2 = torch.chunk(x, 2, dim=-1)
     return torch.cat((-x2, x1), dim=-1)
 
 
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
     # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
-    # pylint: disable=invalid-name
+
     cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
     sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
     cos = cos[position_ids].unsqueeze(1)  # [batch_size, 1, seq_len, dim]
@@ -98,7 +93,7 @@ def flashattn_attn(
     attention_mask: torch.FloatTensor,
     position_ids: torch.LongTensor,
     past_key_value: Optional[Tuple[torch.Tensor]] = None,
-    output_attentions: Optional[bool] = False,  # pylint: disable=unused-argument
+    output_attentions: Optional[bool] = False,
     use_cache: Optional[bool] = False,
     cu_seqlens: Optional[torch.Tensor] = None,
     max_seqlen: Optional[torch.Tensor] = None,
@@ -215,7 +210,6 @@ def decoder_layer_forward(
 ) -> Union[
     Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]
 ]:
-    # pylint: disable=duplicate-code
     residual = hidden_states
 
     hidden_states = self.input_layernorm(hidden_states)
@@ -262,7 +256,6 @@ def stablelm_model_forward(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
 ) -> Union[Tuple, BaseModelOutputWithPast]:
-    # pylint: disable=duplicate-code
     output_attentions = (
         output_attentions
         if output_attentions is not None
@@ -325,13 +318,11 @@ def stablelm_model_forward(
             dtype=torch.bool,
             device=inputs_embeds.device,
         )
-    attention_mask = (
-        self._prepare_decoder_attention_mask(  # pylint: disable=protected-access
-            attention_mask,
-            (batch_size, seq_length),
-            inputs_embeds,
-            past_key_values_length,
-        )
+    attention_mask = self._prepare_decoder_attention_mask(
+        attention_mask,
+        (batch_size, seq_length),
+        inputs_embeds,
+        past_key_values_length,
     )
 
     hidden_states = inputs_embeds
diff --git a/src/axolotl/monkeypatch/tiled_mlp/__init__.py b/src/axolotl/monkeypatch/tiled_mlp/__init__.py
new file mode 100644
index 0000000000..cb9aa6145a
--- /dev/null
+++ b/src/axolotl/monkeypatch/tiled_mlp/__init__.py
@@ -0,0 +1,13 @@
+"""
+TiledMLP monkey patches
+"""
+
+from .patch import (
+    patch_tiled_mlp,
+    patch_tiled_mlp_moe_instances,
+)
+
+__all__ = [
+    "patch_tiled_mlp",
+    "patch_tiled_mlp_moe_instances",
+]
diff --git a/src/axolotl/monkeypatch/tiled_mlp/base.py b/src/axolotl/monkeypatch/tiled_mlp/base.py
new file mode 100644
index 0000000000..ab4f93ab53
--- /dev/null
+++ b/src/axolotl/monkeypatch/tiled_mlp/base.py
@@ -0,0 +1,453 @@
+"""
+TiledMLP support for DDP, FSDP, and single GPU
+"""
+
+import contextlib
+import os
+import threading
+from typing import List
+
+import torch
+
+# Opt-in fp32 accumulation for the tiled backward. The default accumulates
+# at the param's own dtype, which matches what AccumulateGrad does in the
+# unsharded backward and avoids materialising an fp32 buffer the size of
+# every compute param. Set ``AXOLOTL_TILED_MLP_ACCUM_FP32=1`` to recover
+# the previous fp32-accumulator behaviour when bf16 precision is the
+# concern (e.g. very large N-shard sums where bf16 round-off accumulates).
+_TILED_MLP_ACCUM_FP32 = os.environ.get("AXOLOTL_TILED_MLP_ACCUM_FP32", "0") == "1"
+
+
+def _find_fsdp2_module(module):
+    """Return the nearest FSDP2 :class:`FSDPModule` that owns ``module``.
+
+    FSDP2 (``torch.distributed.fsdp.fully_shard``) registers per-module
+    post-backward hooks that reshard parameters once their gradients have
+    been produced. Inside :class:`TiledMLP.backward` we run several inner
+    backwards over shards of the same input; if the wrapping FSDPModule
+    reshards between iterations, the unsharded params are gone and the
+    next tile recomputes against bogus shards. We have to disable reshard
+    on the wrapping FSDPModule for the duration of the loop.
+
+    The MLP itself is rarely the directly-wrapped module — production
+    setups apply ``fully_shard`` at the decoder-layer level. Walk the
+    global FSDP module-state registry to find the nearest ancestor whose
+    parameter group contains us. Result is cached on the module so we pay
+    the lookup once.
+
+    Returns ``None`` if FSDP2 is not in use, or no wrapping FSDPModule
+    contains ``module`` as a descendant.
+    """
+    cached = getattr(module, "_axolotl_fsdp2_owner", "__unset__")
+    if cached != "__unset__":
+        return cached
+
+    try:
+        from torch.distributed._composable_state import _module_state_mapping
+        from torch.distributed.fsdp import FSDPModule
+    except ImportError:
+        module._axolotl_fsdp2_owner = None
+        return None
+
+    # MLP itself wrapped (covers the regression-guard unit test).
+    if isinstance(module, FSDPModule):
+        module._axolotl_fsdp2_owner = module
+        return module
+
+    # Walk the global FSDP registry looking for ancestors. The registry is
+    # a WeakKeyDictionary so the snapshot is cheap and bounded by the
+    # number of FSDP-wrapped modules in the process.
+    target_id = id(module)
+    candidates = []
+    for owner in list(_module_state_mapping.keys()):
+        if not isinstance(owner, FSDPModule):
+            continue
+        if owner is module:
+            continue
+        for sub in owner.modules():
+            if id(sub) == target_id:
+                candidates.append(owner)
+                break
+
+    if not candidates:
+        result = None
+    elif len(candidates) == 1:
+        result = candidates[0]
+    else:
+        # When multiple FSDPModules are ancestors (e.g. fully_shard applied
+        # to both decoder layer and the root), pick the deepest one — its
+        # subtree is smallest. Counting modules is O(N) per candidate but
+        # only runs once per MLP instance.
+        result = min(candidates, key=lambda m: sum(1 for _ in m.modules()))
+
+    module._axolotl_fsdp2_owner = result
+    return result
+
+
+@contextlib.contextmanager
+def _defer_fsdp2_reshard(module):
+    """Suspend FSDP2's post-backward reshard on the wrapping FSDPModule.
+
+    The tiled backward calls :func:`torch.autograd.backward` once per shard.
+    Each inner backward triggers FSDP2's per-module post-backward hooks,
+    which would reshard parameters mid-loop. We pause that by toggling
+    ``set_reshard_after_backward(False)`` on the wrapping FSDPModule, run
+    the loop, restore the original setting, then issue a single explicit
+    ``reshard()`` so the post-loop state matches normal FSDP2 semantics.
+
+    No-op when ``module`` is not under FSDP2.
+    """
+    fsdp_mod = _find_fsdp2_module(module)
+    if fsdp_mod is None:
+        yield
+        return
+
+    # No public getter for ``reshard_after_backward`` in PyTorch 2.11;
+    # read off the param group directly. The internal accessor surface
+    # is documented in
+    # ``torch.distributed.fsdp._fully_shard._fsdp_param_group.FSDPParamGroup``.
+    state = fsdp_mod._get_fsdp_state()
+    param_group = state._fsdp_param_group
+    if param_group is None:
+        # Nothing to defer (e.g. ignored module with no FSDP-managed params).
+        yield
+        return
+
+    prev = param_group.reshard_after_backward
+    fsdp_mod.set_reshard_after_backward(False, recurse=False)
+    try:
+        yield
+    finally:
+        # Restore so subsequent backward passes outside the tile loop
+        # behave normally, then issue the deferred reshard once.
+        fsdp_mod.set_reshard_after_backward(prev, recurse=False)
+        fsdp_mod.reshard()
+
+
+class DeepSpeedTiledMLPMoE(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        fn,
+        self,
+        x,
+        shards,
+        compute_params,
+    ) -> torch.Tensor:
+        ctx.fn = fn
+        ctx.self = self
+        ctx.shards = shards
+        ctx.compute_params = [p for p in compute_params if p.requires_grad]
+        ctx.save_for_backward(x)
+
+        x_shards = list(torch.chunk(x, chunks=shards, dim=1))
+        with torch.no_grad():
+            output_shards = [fn(self, x_shard) for x_shard in x_shards]
+
+        ctx.is_tuple_output = isinstance(output_shards[0], tuple)
+        if isinstance(output_shards[0], tuple):
+            tuple_dim_idx = [1, 0]
+            output_unsharded = tuple(
+                torch.cat(
+                    [output_shard[i] for output_shard in output_shards],
+                    dim=tuple_dim_idx[i],
+                )
+                for i in range(len(output_shards[0]))
+            )
+        else:
+            output_unsharded = torch.cat(output_shards, dim=1)
+
+        return output_unsharded
+
+    @staticmethod
+    def backward(ctx, *grads) -> torch.Tensor:
+        fn = ctx.fn
+        (x,) = ctx.saved_tensors
+        self = ctx.self
+        shards = ctx.shards
+        compute_params = ctx.compute_params
+        is_tuple_output = ctx.is_tuple_output
+
+        x_requires_grad = x.requires_grad
+        x = x.detach()
+        # detach() unsets `x.requires_grad`, so restore it
+        x.requires_grad_(x_requires_grad)
+
+        incoming_grad = grads[0]
+        x_grad = torch.zeros_like(x)
+        x_shards = list(torch.chunk(x, chunks=shards, dim=1))
+
+        shard_step = x_shards[0].numel()
+        for i, x_shard in enumerate(x_shards):
+            # Tell deepspeed not to add a new grad to its ipg bucket until the last shard is run
+            if compute_params is not None:
+                if i + 1 < shards:
+                    for param in compute_params:
+                        param.ds_grad_is_ready = False
+                else:
+                    # last shard, can add the grad
+                    for param in compute_params:
+                        param.ds_grad_is_ready = True
+
+            x_shard.requires_grad_(x_requires_grad)
+
+            shard_offset = i * shard_step
+            x_shard.grad = (
+                x_grad.view(-1)
+                .narrow(0, shard_offset, x_shard.numel())
+                .view_as(x_shard)
+            )
+            incoming_grad_shard = (
+                incoming_grad.view(-1)
+                .narrow(0, shard_offset, x_shard.numel())
+                .view_as(x_shard)
+            )
+            with torch.enable_grad():
+                output = fn(self, x_shard)
+            if is_tuple_output:
+                torch.autograd.backward(output[0], incoming_grad_shard)
+            else:
+                torch.autograd.backward(output, incoming_grad_shard)
+
+        return (None, None, x_grad, None, None)
+
+
+class TiledMLP(torch.autograd.Function):
+    """
+    TiledMLP implementation using gradient hooks
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        fn,
+        self,
+        x,
+        shards,
+        compute_params,
+    ) -> torch.Tensor:
+        ctx.fn = fn
+        ctx.self = self
+        ctx.shards = shards
+        ctx.compute_params = [p for p in compute_params if p.requires_grad]
+        ctx.save_for_backward(x)
+
+        x_shards = list(torch.chunk(x, chunks=shards, dim=1))
+        with torch.no_grad():
+            output_shards = [fn(self, x_shard) for x_shard in x_shards]
+        ctx.is_tuple_output = isinstance(output_shards[0], tuple)
+        if isinstance(output_shards[0], tuple):
+            tuple_dim_idx = [1, 0]
+            output_unsharded = tuple(
+                torch.cat(
+                    [output_shard[i] for output_shard in output_shards],
+                    dim=tuple_dim_idx[i],
+                )
+                for i in range(len(output_shards[0]))
+            )
+        else:
+            output_unsharded = torch.cat(output_shards, dim=1)
+
+        return output_unsharded
+
+    @staticmethod
+    def backward(ctx, *grads) -> torch.Tensor:
+        fn = ctx.fn
+        (x,) = ctx.saved_tensors
+        self = ctx.self
+        shards = ctx.shards
+        compute_params = ctx.compute_params
+        is_tuple_output = ctx.is_tuple_output
+
+        x_requires_grad = x.requires_grad
+        x = x.detach()
+        x.requires_grad_(x_requires_grad)
+
+        incoming_grad = grads[0]
+        x_grad = torch.zeros_like(x)
+        x_shards = list(torch.chunk(x, chunks=shards, dim=1))
+
+        # Snapshot existing ``.grad`` for each param and zero it; we will
+        # accumulate the per-shard contributions into a per-param buffer
+        # and write back at the end. The previous implementation used
+        # ``param.register_hook`` per shard, which (a) re-installed hooks
+        # every iteration so the N-th shard ran N stacked hooks and
+        # double-counted contributions, and (b) scaled by ``1/N`` even
+        # though sequence-dim sharding makes per-shard grads additive,
+        # not averaged. The combined effect was a gradient roughly
+        # 2x-2.5x the analytical value. Direct inline accumulation is
+        # both simpler and correct, and avoids interactions with FSDP2's
+        # own backward hooks.
+        #
+        # The accumulator defaults to the param's own dtype to match
+        # what AccumulateGrad would do in the unsharded backward. The
+        # earlier implementation accumulated in fp32, which doubled the
+        # parameter-side memory footprint in bf16 MoE training where the
+        # accumulator's ``[E, hidden, 2*intermediate]`` shape dominates.
+        # Set ``AXOLOTL_TILED_MLP_ACCUM_FP32=1`` to opt back into fp32
+        # accumulation when bf16 round-off is the concern.
+        prev_grads = {}
+        accum_grads = {}
+        for p in compute_params:
+            prev_grads[p] = p.grad
+            accum_dtype = torch.float32 if _TILED_MLP_ACCUM_FP32 else p.dtype
+            accum_grads[p] = torch.zeros_like(p, dtype=accum_dtype)
+            p.grad = None
+
+        shard_step = x_shards[0].numel()
+        # Suspend FSDP2 post-backward reshard for the duration of the loop.
+        # Without this, the first inner backward triggers FSDP2's reshard
+        # hook on the wrapping FSDPModule and subsequent shards recompute
+        # against only-local DTensor shards — silent grad corruption.
+        # Single-GPU and DDP paths fall through to a no-op context manager.
+        with _defer_fsdp2_reshard(self):
+            for i, x_shard in enumerate(x_shards):
+                x_shard.requires_grad_(x_requires_grad)
+
+                shard_offset = i * shard_step
+                x_shard.grad = (
+                    x_grad.view(-1)
+                    .narrow(0, shard_offset, x_shard.numel())
+                    .view_as(x_shard)
+                )
+                incoming_grad_shard = (
+                    incoming_grad.view(-1)
+                    .narrow(0, shard_offset, x_shard.numel())
+                    .view_as(x_shard)
+                )
+
+                with torch.enable_grad():
+                    output = fn(self, x_shard)
+                if is_tuple_output:
+                    torch.autograd.backward(output[0], incoming_grad_shard)
+                else:
+                    torch.autograd.backward(output, incoming_grad_shard)
+
+                # Capture this shard's contribution into the per-param
+                # accumulator and clear ``.grad`` so the next shard starts
+                # from zero. Skip the dtype cast when the accumulator
+                # matches the param dtype (the default) — that cast was
+                # the per-shard HBM-bandwidth tax on the bf16 path.
+                for p in compute_params:
+                    if p.grad is not None:
+                        shard_grad = p.grad.detach()
+                        if shard_grad.dtype != accum_grads[p].dtype:
+                            shard_grad = shard_grad.to(accum_grads[p].dtype)
+                        accum_grads[p].add_(shard_grad)
+                        p.grad = None
+
+        # Restore prior grad value (if any) and add the tiled contribution.
+        for p in compute_params:
+            tiled_contrib = accum_grads[p]
+            if tiled_contrib.dtype != p.dtype:
+                tiled_contrib = tiled_contrib.to(p.dtype)
+            if prev_grads[p] is None:
+                p.grad = tiled_contrib
+            else:
+                p.grad = prev_grads[p] + tiled_contrib
+
+        return (None, None, x_grad, None, None)
+
+
+class GradientAccumulator:
+    """
+    Manual gradient accumulator for TiledMLP with configurable precision.
+
+    .. note::
+        The production TiledMLP backward (above) accumulates inline and
+        does not call this class — it is retained as a reference / opt-in
+        path for callers that want hook-based accumulation. The defaults
+        below match the inline path: param-dtype accumulator (matches
+        ``AccumulateGrad`` in the unsharded backward) and ``1.0`` per-shard
+        scaling (sequence-dim sharded grads are additive, not averaged).
+    """
+
+    def __init__(
+        self,
+        params: List[torch.nn.Parameter],
+        total_shards: int,
+        dtype: torch.dtype | None = None,
+    ):
+        self.params = params
+        self.total_shards = total_shards
+        # Default to the param's own dtype to avoid the 2x parameter-side
+        # memory regression in bf16 MoE training where the accumulator
+        # shape ``[E, hidden, 2*intermediate]`` dominates. fp32 accumulation
+        # is opt-in via the ``dtype`` arg.
+        if dtype is not None:
+            self.grad_accumulation_dtype = dtype
+        elif params:
+            self.grad_accumulation_dtype = params[0].dtype
+        else:
+            self.grad_accumulation_dtype = torch.float32
+        self.accumulated_grads = {}
+        self.hooks = []
+        self.lock = threading.Lock()
+        # Sequence-dim shards partition the per-token sum; their
+        # contributions are additive (``sum_t dL_t/dW``), not averaged.
+        # The previous ``1/total_shards`` scaling produced a mean and was
+        # a correctness bug for this sharding semantics.
+        self.gradient_scale = 1.0
+
+        # Initialize accumulated gradients in the specified dtype
+        for param in self.params:
+            if param.grad is not None:
+                self.accumulated_grads[param] = param.grad.to(
+                    self.grad_accumulation_dtype
+                )
+                param.grad = None
+            else:
+                self.accumulated_grads[param] = torch.zeros_like(
+                    param, dtype=self.grad_accumulation_dtype
+                )
+
+    def install_hooks(self, is_last_shard: bool):
+        """Install gradient hooks that accumulate gradients in higher precision"""
+
+        def create_hook(param):
+            def hook(grad):
+                with self.lock:
+                    # Skip the dtype cast when the accumulator already
+                    # matches the grad dtype (the default after the
+                    # param-dtype change above) — the redundant cast was
+                    # the per-shard HBM bandwidth tax called out in the
+                    # tiled-MLP regression analysis.
+                    if grad.dtype == self.grad_accumulation_dtype:
+                        scaled_grad = (
+                            grad
+                            if self.gradient_scale == 1.0
+                            else grad * self.gradient_scale
+                        )
+                    else:
+                        scaled_grad = (
+                            grad.to(self.grad_accumulation_dtype) * self.gradient_scale
+                        )
+
+                    if param in self.accumulated_grads:
+                        self.accumulated_grads[param] += scaled_grad
+                    else:
+                        self.accumulated_grads[param] = scaled_grad.clone()
+
+                    # Only assign the accumulated gradient on the last shard
+                    if is_last_shard:
+                        if self.accumulated_grads[param].dtype != param.dtype:
+                            param.grad = self.accumulated_grads[param].to(param.dtype)
+                        else:
+                            param.grad = self.accumulated_grads[param]
+                        return param.grad
+                    return None
+
+            return hook
+
+        # Install hooks on all parameters
+        for param in self.params:
+            if param.requires_grad:
+                hook = param.register_hook(create_hook(param))
+                self.hooks.append(hook)
+
+    def cleanup(self):
+        """Remove all installed hooks"""
+        for hook in self.hooks:
+            hook.remove()
+        self.hooks.clear()
+        del self.accumulated_grads
diff --git a/src/axolotl/monkeypatch/tiled_mlp/patch.py b/src/axolotl/monkeypatch/tiled_mlp/patch.py
new file mode 100644
index 0000000000..6842ba7980
--- /dev/null
+++ b/src/axolotl/monkeypatch/tiled_mlp/patch.py
@@ -0,0 +1,280 @@
+"""Monkeypatch for Tiled MLP implementation"""
+
+import math
+import os
+
+import torch
+import torch.distributed as dist
+
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# Suffixes used to discover MoE block classes inside
+# ``transformers.models.{model_type}.modeling_{model_type}``.
+# Order matters — preferred names come first.
+_MOE_BLOCK_SUFFIXES = ("SparseMoeBlock", "MoeMLP", "MoE")
+
+
+def _resolve_moe_block_cls(module, model_cls_prefix):
+    """Return the MoE block class for the model module, or ``None`` if dense."""
+    for suffix in _MOE_BLOCK_SUFFIXES:
+        cls = getattr(module, f"{model_cls_prefix}{suffix}", None)
+        if cls is not None:
+            return cls
+    return None
+
+
+def _build_tiled_forward(
+    inner_forward,
+    model_type,
+    cfg_num_shards,
+    is_moe_block,
+):
+    """Construct a ``tiled_mlp_forward`` closure.
+
+    The returned forward shards inputs along the sequence dim and dispatches
+    to the correct :class:`torch.autograd.Function` implementation based on
+    the parallel-training backend in use.
+
+    ``inner_forward`` is the un-tiled forward (either the dense MLP forward
+    or the MoE block's routing+expert forward — possibly a kernels-substituted
+    forward in the scattermoe-lora case).
+    """
+    from deepspeed.runtime.sequence_parallel.ulysses_sp import (
+        TiledMLP as DeepSpeedTiledMLP,
+    )
+
+    from axolotl.monkeypatch.tiled_mlp.base import DeepSpeedTiledMLPMoE, TiledMLP
+
+    is_distributed = int(os.environ.get("WORLD_SIZE", 1)) > 1
+
+    def tiled_mlp_forward(self, x):
+        input_shape = x.shape
+        seqlen = input_shape[-2]
+        if cfg_num_shards is None:
+            # Target ~32K tokens per shard. The previous `ceil(seq / hidden)`
+            # heuristic produced only ~2K tokens/shard at long context, well
+            # below the MoE kernel's BLOCK_M sweet spot. An empirical sweep at
+            # seq ∈ {64K, 128K, 256K, 512K} showed 3.2× speed-up at 64–256K
+            # and 2.1× at 512K from raising per-shard tokens to ~32K, with
+            # only a modest peak-mem cost (~5–10 GiB extra at seq=256K)
+            # because the routed intermediate buffer dominates and scales
+            # linearly with per-shard tokens. Operators can override via
+            # cfg_num_shards for niche cases (smaller intermediate, larger
+            # top_k) where the default is wrong.
+            target_tokens_per_shard = 32768
+            num_shards = max(1, math.ceil(seqlen / target_tokens_per_shard))
+            if is_distributed:
+                num_shards_tensor = torch.tensor(num_shards, device=x.device)
+                dist.all_reduce(num_shards_tensor, op=dist.ReduceOp.MAX)
+                num_shards = num_shards_tensor.item()
+        else:
+            num_shards = cfg_num_shards
+
+        if not self._compute_params:
+            self._compute_params = [p for p in self.parameters() if p.requires_grad]
+
+        compute_params = self._compute_params
+        if not self._tiled_mlp_dist_impl:
+            uses_deepspeed = (
+                self._compute_params
+                and any(
+                    hasattr(p, "ds_id") or hasattr(p, "param_idx_in_group")
+                    for p in self._compute_params
+                )
+            ) or os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true"
+
+            if uses_deepspeed:
+                # gpt_oss already used the MoE variant before this refactor;
+                # extend the same treatment to every MoE block, since they
+                # tend to return tuple outputs (hidden_states, router_logits)
+                # the way gpt_oss does.
+                if model_type == "gpt_oss" or is_moe_block:
+                    self._tiled_mlp_dist_impl = DeepSpeedTiledMLPMoE
+                else:
+                    self._tiled_mlp_dist_impl = DeepSpeedTiledMLP
+            else:
+                self._tiled_mlp_dist_impl = TiledMLP
+
+        return self._tiled_mlp_dist_impl.apply(
+            inner_forward,
+            self,
+            x,
+            num_shards,
+            compute_params,
+        )
+
+    return tiled_mlp_forward
+
+
+def _prepare_target_class(target_cls):
+    """Initialize the bookkeeping attrs the tiled forward expects."""
+    target_cls._compute_params = []
+    target_cls._tiled_mlp_dist_impl = None
+
+
+def patch_tiled_mlp(
+    model_type,
+    use_original_mlp=True,
+    cfg_num_shards=None,
+    use_scattermoe=False,
+):
+    """Install the class-level tiled MLP patch.
+
+    For dense models this patches ``{prefix}MLP`` (falling back to
+    ``{prefix}TextMLP`` for multimodal wrappers).
+
+    For MoE models with scattermoe-lora active, the MoE block class
+    (``{prefix}SparseMoeBlock`` / ``{prefix}MoeMLP`` / ``{prefix}MoE``) is the
+    one whose forward does routing + expert invocation, so we patch that.
+    Note that the ``kernels`` library installs scattermoe-lora's forward at
+    the *instance* level during ``model.kernelize()``, so the class-level
+    patch is shadowed at runtime. :func:`patch_tiled_mlp_moe_instances` is
+    the companion post-model-load step that re-wraps each MoE block instance
+    so the tiled forward runs on top of the kernels-installed forward.
+    """
+    module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+    model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+    try:
+        module = __import__(module_path, fromlist=[f"{model_cls_prefix}MLP"])
+    except ImportError as e:
+        raise RuntimeError(
+            f"Could not import MLP class for model_type: {model_type}. Error: {str(e)}"
+        ) from e
+
+    # MoE block patch path: only walk into this branch when scattermoe-lora
+    # is active. For non-scattermoe MoE models the dense MLP fallback applies
+    # — we do not auto-enable MoE-block tiling because each model family's
+    # block forward has different output-tuple semantics.
+    moe_block_cls = (
+        _resolve_moe_block_cls(module, model_cls_prefix) if use_scattermoe else None
+    )
+
+    if moe_block_cls is not None:
+        original_forward = moe_block_cls.forward
+        tiled_forward = _build_tiled_forward(
+            inner_forward=original_forward,
+            model_type=model_type,
+            cfg_num_shards=cfg_num_shards,
+            is_moe_block=True,
+        )
+        moe_block_cls.forward = tiled_forward
+        _prepare_target_class(moe_block_cls)
+        LOG.info(
+            "Successfully monkey-patched TiledMLP for model_type: "
+            f"{model_type} (MoE block: {moe_block_cls.__name__})"
+        )
+        return
+
+    # Dense MLP path (existing behavior).
+    try:
+        mlp_cls = getattr(
+            module,
+            f"{model_cls_prefix}MLP",
+            None,
+        ) or getattr(module, f"{model_cls_prefix}TextMLP")
+    except AttributeError as e:
+        raise RuntimeError(
+            f"Could not import MLP class for model_type: {model_type}. Error: {str(e)}"
+        ) from e
+
+    if use_original_mlp:
+        mlp_forward = mlp_cls.forward
+    else:
+
+        def generic_mlp_forward(self_, hs):
+            return self_.down_proj(
+                self_.act_fn(self_.gate_proj(hs)) * self_.up_proj(hs)
+            )
+
+        mlp_forward = torch.compile(generic_mlp_forward)
+
+    tiled_forward = _build_tiled_forward(
+        inner_forward=mlp_forward,
+        model_type=model_type,
+        cfg_num_shards=cfg_num_shards,
+        is_moe_block=False,
+    )
+    mlp_cls.forward = tiled_forward
+    _prepare_target_class(mlp_cls)
+    LOG.info(f"Successfully monkey-patched TiledMLP for model_type: {model_type}")
+
+
+def patch_tiled_mlp_moe_instances(
+    model,
+    model_type,
+    cfg_num_shards=None,
+):
+    """Re-wrap each MoE block instance's ``forward`` after model load.
+
+    The ``kernels`` library installs scattermoe-lora's forward on each MoE
+    block *instance* during ``model.kernelize()`` (called inside
+    ``from_pretrained``). That instance-level binding shadows the class-level
+    patch :func:`patch_tiled_mlp` installs, so without this step tiling is
+    silently bypassed on every block. We capture each instance's current
+    forward (the kernels-installed one) and rebind the instance to a tiled
+    forward that delegates to it.
+
+    Does nothing if no MoE block class exists for ``model_type`` or if
+    ``model`` contains no instances of it.
+    """
+    from types import MethodType
+
+    module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+    model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+    try:
+        module = __import__(module_path, fromlist=[model_cls_prefix])
+    except ImportError:
+        return 0
+
+    moe_block_cls = _resolve_moe_block_cls(module, model_cls_prefix)
+    if moe_block_cls is None:
+        return 0
+
+    wrapped = 0
+    for sub in model.modules():
+        if not isinstance(sub, moe_block_cls):
+            continue
+        # If there is no per-instance ``forward`` binding, the class-level
+        # tiled patch from ``patch_tiled_mlp`` is still active; nothing to do.
+        # Kernels (when scattermoe-lora kernelizes the model) installs a
+        # bound method on the instance, which shows up in ``__dict__``.
+        if "forward" not in sub.__dict__:
+            continue
+        # Snapshot the instance-level forward installed by kernels.
+        bound_forward = sub.__dict__["forward"]
+        # Convert bound method back to a plain function that takes (self, x)
+        # so the tiled wrapper can pass `self` through to it.
+        if hasattr(bound_forward, "__func__"):
+            inner_fn = bound_forward.__func__
+        else:
+            # Instance-bound closure: wrap it so the (self, x) signature lines up.
+            def _adapt(orig):
+                def _call(self_, x):  # noqa: ARG001
+                    return orig(x)
+
+                return _call
+
+            inner_fn = _adapt(bound_forward)
+
+        tiled_forward = _build_tiled_forward(
+            inner_forward=inner_fn,
+            model_type=model_type,
+            cfg_num_shards=cfg_num_shards,
+            is_moe_block=True,
+        )
+        # Each instance needs its own bookkeeping (compute_params,
+        # dist_impl) so concurrent forwards across blocks don't stomp.
+        sub._compute_params = []
+        sub._tiled_mlp_dist_impl = None
+        sub.forward = MethodType(tiled_forward, sub)
+        wrapped += 1
+
+    if wrapped:
+        LOG.info(
+            f"Successfully wrapped TiledMLP around {wrapped} {moe_block_cls.__name__} "
+            f"instance(s) for model_type: {model_type}"
+        )
+    return wrapped
diff --git a/src/axolotl/monkeypatch/torchao_optim.py b/src/axolotl/monkeypatch/torchao_optim.py
new file mode 100644
index 0000000000..98c325a0b0
--- /dev/null
+++ b/src/axolotl/monkeypatch/torchao_optim.py
@@ -0,0 +1,154 @@
+"""
+Patch for torchao optim subclasses that crash under torch.compile.
+
+torchao 0.17.0 PR #3934 added an "appearance dtype" to OptimState{4,8}bit and
+OptimStateFp8, allowing them to report as e.g. bf16 while internally storing
+quantized codes.  Three issues:
+
+1. aten.view.default doesn't propagate the appearance dtype, so views (e.g. from
+   DTensor.from_local()) revert to float32 while the base is bf16.  torch.compile's
+   fake-tensor metadata check then fails (AssertionError: torch.bfloat16 != torch.float32).
+
+2. aten._to_copy doesn't clone internal tensors, so same-device dtype changes
+   (e.g. .float()) create an accidental view relationship with the same issue.
+
+3. aten.view.dtype is unimplemented, so if the dtype-view path IS taken, it crashes
+   with NotImplementedError.
+
+Fix: propagate dtype in view.default (primary), clone in _to_copy, register view.dtype.
+
+Upstream fix: https://github.com/pytorch/ao/pull/4216
+"""
+
+import torch
+from torch.utils._python_dispatch import return_and_correct_aliasing
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+aten = torch.ops.aten
+
+
+def _needs_view_dtype_patch(cls):
+    """Check if a subclass is missing aten.view.dtype."""
+    op_table = getattr(cls, "_ATEN_OP_TABLE", {}).get(cls, {})
+    return aten.view.dtype not in op_table
+
+
+def patch_torchao_optim_state_8bit():
+    """Patch torchao optim subclasses for torch.compile compatibility."""
+    try:
+        from torchao.optim.subclass_8bit import OptimState8bit
+    except ImportError:
+        return
+
+    # Patch view.default to propagate appearance dtype
+    @OptimState8bit.implements(aten.view.default)
+    def _(func, types, args, kwargs):
+        x, shape = args
+        return OptimState8bit(
+            x.codes.view(shape), x.scale, x.qmap, x.signed, dtype=x.dtype
+        )
+
+    # Patch _to_copy to clone internal tensors (breaks accidental view)
+    @OptimState8bit.implements(aten._to_copy.default)
+    def _(func, types, args, kwargs):
+        dtype = kwargs.get("dtype", args[0].dtype)
+        device = kwargs.get("device", None)
+        out = OptimState8bit(
+            args[0].codes.to(device=device).clone(),
+            args[0].scale.to(device=device).clone(),
+            args[0].qmap.to(device=device).clone(),
+            args[0].signed,
+            dtype=dtype,
+        )
+        return return_and_correct_aliasing(func, args, kwargs, out)
+
+    if _needs_view_dtype_patch(OptimState8bit):
+
+        @OptimState8bit.implements(aten.view.dtype)
+        def _(func, types, args, kwargs):
+            x, dtype = args
+            return OptimState8bit(x.codes, x.scale, x.qmap, x.signed, dtype=dtype)
+
+    LOG.debug("Patched OptimState8bit for torch.compile compatibility")
+
+    try:
+        from torchao.optim.subclass_4bit import OptimState4bit
+    except ImportError:
+        OptimState4bit = None
+
+    if OptimState4bit is not None:
+
+        @OptimState4bit.implements(aten.view.default)
+        def _(func, types, args, kwargs):
+            x, shape = args
+            if tuple(x.shape) == tuple(shape):
+                return OptimState4bit(
+                    x.codes, x.scale, x.qmap, x.signed, x._shape, dtype=x.dtype
+                )
+            if len(shape) == 1 and shape[0] == -1:
+                return OptimState4bit(
+                    x.codes, x.scale, x.qmap, x.signed, (x.numel(),), dtype=x.dtype
+                )
+            raise ValueError(
+                f"{x.__class__.__name__} only supports .view() with same shape or shape=[-1]"
+            )
+
+        @OptimState4bit.implements(aten._to_copy.default)
+        def _(func, types, args, kwargs):
+            dtype = kwargs.get("dtype", args[0].dtype)
+            device = kwargs.get("device", None)
+            out = OptimState4bit(
+                args[0].codes.to(device=device).clone(),
+                args[0].scale.to(device=device).clone(),
+                args[0].qmap.to(device=device).clone(),
+                args[0].signed,
+                args[0].shape,
+                dtype=dtype,
+            )
+            return return_and_correct_aliasing(func, args, kwargs, out)
+
+        if _needs_view_dtype_patch(OptimState4bit):
+
+            @OptimState4bit.implements(aten.view.dtype)
+            def _(func, types, args, kwargs):
+                x, dtype = args
+                return OptimState4bit(
+                    x.codes, x.scale, x.qmap, x.signed, x.shape, dtype=dtype
+                )
+
+        LOG.debug("Patched OptimState4bit for torch.compile compatibility")
+
+    try:
+        from torchao.optim.subclass_fp8 import OptimStateFp8
+    except ImportError:
+        OptimStateFp8 = None
+
+    if OptimStateFp8 is not None:
+
+        @OptimStateFp8.implements(aten.view.default)
+        def _(func, types, args, kwargs):
+            x, shape = args
+            return OptimStateFp8(x.codes.view(shape), x.scale, dtype=x.dtype)
+
+        @OptimStateFp8.implements(aten._to_copy.default)
+        def _(func, types, args, kwargs):
+            dtype = kwargs.get("dtype", args[0].dtype)
+            device = kwargs.get("device", None)
+            out = OptimStateFp8(
+                args[0].codes.to(device=device).clone(),
+                args[0].scale.to(device=device).clone(),
+                dtype=dtype,
+            )
+            return return_and_correct_aliasing(func, args, kwargs, out)
+
+        if _needs_view_dtype_patch(OptimStateFp8):
+
+            @OptimStateFp8.implements(aten.view.dtype)
+            def _(func, types, args, kwargs):
+                x, dtype = args
+                return OptimStateFp8(x.codes, x.scale, dtype=dtype)
+
+        LOG.debug("Patched OptimStateFp8 for torch.compile compatibility")
diff --git a/src/axolotl/monkeypatch/trainer/__init__.py b/src/axolotl/monkeypatch/trainer/__init__.py
new file mode 100644
index 0000000000..27edc63c34
--- /dev/null
+++ b/src/axolotl/monkeypatch/trainer/__init__.py
@@ -0,0 +1,3 @@
+from .utils import entropy_from_logits, selective_log_softmax
+
+__all__ = ["entropy_from_logits", "selective_log_softmax"]
diff --git a/src/axolotl/monkeypatch/trainer/lr.py b/src/axolotl/monkeypatch/trainer/lr.py
new file mode 100644
index 0000000000..c33674ceeb
--- /dev/null
+++ b/src/axolotl/monkeypatch/trainer/lr.py
@@ -0,0 +1,42 @@
+"""
+monkeypatch for Trainer _get_learning_rate method
+"""
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+# TODO remove this patch once https://github.com/huggingface/transformers/pull/37881 is included in a release
+def _get_learning_rate(self):
+    if self.is_deepspeed_enabled:
+        # with deepspeed's fp16 and dynamic loss scale enabled the optimizer/scheduler steps may
+        # not run for the first few dozen steps while loss scale is too large, and thus during
+        # that time `get_last_lr` will fail if called during that warm up stage, so work around it:
+        try:
+            last_lr = self.lr_scheduler.get_last_lr()[0]
+        except AssertionError as e:
+            if "need to call step" in str(e):
+                LOG.warning(
+                    "tried to get lr value before scheduler/optimizer started stepping, returning lr=0"
+                )
+                last_lr = 0
+            else:
+                raise
+    else:
+        if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+            last_lr = self.optimizer.param_groups[0]["lr"]
+        else:
+            last_lr = self.lr_scheduler.get_last_lr()[0]
+
+    if torch.is_tensor(last_lr):
+        last_lr = last_lr.item()
+    return last_lr
+
+
+def patch_trainer_get_lr():
+    from transformers.trainer import Trainer
+
+    Trainer._get_learning_rate = _get_learning_rate
diff --git a/src/axolotl/monkeypatch/trainer/trl.py b/src/axolotl/monkeypatch/trainer/trl.py
new file mode 100644
index 0000000000..bca9f92dec
--- /dev/null
+++ b/src/axolotl/monkeypatch/trainer/trl.py
@@ -0,0 +1,13 @@
+"""Monkeypatch for TRL trainer FSDP preparation."""
+
+
+def prepare_fsdp(model, accelerator):
+    from axolotl.monkeypatch.accelerate.fsdp2 import fsdp2_prepare_model
+
+    return fsdp2_prepare_model(accelerator, model)
+
+
+def patch_trl_prepare_fsdp2():
+    import trl.models.utils
+
+    trl.models.utils.prepare_fsdp = prepare_fsdp
diff --git a/src/axolotl/monkeypatch/trainer/trl_vllm.py b/src/axolotl/monkeypatch/trainer/trl_vllm.py
new file mode 100644
index 0000000000..886f1933f0
--- /dev/null
+++ b/src/axolotl/monkeypatch/trainer/trl_vllm.py
@@ -0,0 +1,302 @@
+"""Monkeypatches for TRL's vLLM integration and trainer utils.
+
+Adds:
+- VLLMClient.batch_update_named_params: batched weight sync (fewer HTTP round-trips)
+- extract_logprobs: NaN→0.0 fix (prevents downstream NaN propagation)
+- VLLMGeneration: weight_sync_chunk_size + batched sync path for non-FSDP/non-ZeRO
+- split_tensor_dict / shuffle_sequence_dict: scalar type handling (int/float/bool passthrough)
+"""
+
+import math
+from functools import wraps
+
+import torch
+from torch import nn
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _batch_update_named_params(
+    self, params: list[tuple[str, torch.Tensor]], chunk_size: int | None = None
+):
+    """Batched weight sync — uses NCCL if communicator available, HTTP otherwise."""
+    has_communicator = getattr(self, "communicator", None) is not None
+
+    if has_communicator:
+        # Fast path: metadata via HTTP, tensors via NCCL
+        from transformers import is_torch_xpu_available
+
+        if chunk_size is None:
+            chunks = [params]
+        else:
+            chunks = []
+            current_chunk: list[tuple[str, torch.Tensor]] = []
+            current_elements = 0
+            for name, weights in params:
+                n_elem = weights.numel()
+                if current_chunk and current_elements + n_elem > chunk_size:
+                    chunks.append(current_chunk)
+                    current_chunk = []
+                    current_elements = 0
+                current_chunk.append((name, weights))
+                current_elements += n_elem
+            if current_chunk:
+                chunks.append(current_chunk)
+
+        for chunk in chunks:
+            param_metadata = [
+                {
+                    "name": name,
+                    "dtype": str(weights.dtype),
+                    "shape": list(weights.shape),
+                }
+                for name, weights in chunk
+            ]
+            url = f"{self.base_url}/batch_update_named_params/"
+            response = self.session.post(
+                url, json={"params": param_metadata}, timeout=120
+            )
+            if response.status_code == 404:
+                # Server doesn't support batch endpoint — fall back to individual updates
+                for meta in param_metadata:
+                    ind_url = f"{self.base_url}/update_named_param/"
+                    ind_response = self.session.post(ind_url, json=meta, timeout=120)
+                    if ind_response.status_code != 200:
+                        raise Exception(
+                            f"Individual update failed: {ind_response.status_code}, {ind_response.text}"
+                        )
+            elif response.status_code != 200:
+                raise Exception(
+                    f"Request failed: {response.status_code}, {response.text}"
+                )
+
+            for _name, weights in chunk:
+                if is_torch_xpu_available():
+                    self.communicator.broadcast(weights, root=self.rank)
+                else:
+                    self.communicator.broadcast(weights, src=self.rank)
+
+            if is_torch_xpu_available():
+                self.communicator.barrier()
+            else:
+                self.communicator.group.barrier()
+    else:
+        # HTTP-only path: encode tensor data in request body (no NCCL needed).
+        # Batch by byte size to avoid huge HTTP payloads.
+        MAX_BYTES_PER_REQUEST = 10 * 1024 * 1024  # 10 MB
+        HTTP_TIMEOUT = 120  # seconds per request
+
+        payload: list[dict] = []
+        payload_bytes = 0
+        url = f"{self.base_url}/http_update_weights/"
+
+        def _flush(p: list[dict]) -> None:
+            if not p:
+                return
+            response = self.session.post(url, json={"params": p}, timeout=HTTP_TIMEOUT)
+            if response.status_code != 200:
+                raise Exception(
+                    f"Request failed: {response.status_code}, {response.text}"
+                )
+
+        from axolotl.utils.weight_serde import encode_for_http
+
+        for name, weights in params:
+            entry = encode_for_http(name, weights)
+            entry_bytes = weights.nelement() * weights.element_size()
+
+            # Flush current batch if adding this entry would exceed limit
+            if payload and payload_bytes + entry_bytes > MAX_BYTES_PER_REQUEST:
+                _flush(payload)
+                payload = []
+                payload_bytes = 0
+
+            payload.append(entry)
+            payload_bytes += entry_bytes
+
+        _flush(payload)  # send remaining
+
+
+def _update_model_params(self, model: nn.Module, chunk_size: int | None = None):
+    """Updates all model params using batch_update_named_params."""
+    params = [(name, param.data) for name, param in model.named_parameters()]
+    self.batch_update_named_params(params, chunk_size=chunk_size)
+
+
+def _patched_extract_logprobs(all_outputs):
+    """extract_logprobs with NaN→0.0 fix (stock TRL uses None which causes downstream errors)."""
+    all_logprobs = []
+    all_token_ids = []
+
+    for outputs in all_outputs:
+        for output in outputs.outputs:
+            if output.logprobs is None:
+                return None, None
+            seq_logprobs = []
+            seq_token_ids = []
+            for lp in output.logprobs:
+                sorted_items = sorted(lp.items(), key=lambda x: x[1].rank)
+                seq_token_ids.append([token_id for token_id, _ in sorted_items])
+                seq_logprobs.append(
+                    [
+                        0.0 if math.isnan(item.logprob) else item.logprob
+                        for _, item in sorted_items
+                    ]
+                )
+            all_logprobs.append(seq_logprobs)
+            all_token_ids.append(seq_token_ids)
+
+    return all_logprobs, all_token_ids
+
+
+def _patched_split_tensor_dict(tensor_dict, num_chunks):
+    """split_tensor_dict that handles scalar types (int/float/bool) for num_items_in_batch."""
+    first_tensor = next(
+        tensor
+        for tensor in tensor_dict.values()
+        if tensor is not None and isinstance(tensor, torch.Tensor) and tensor.ndim > 0
+    )
+    chunk_size = first_tensor.shape[0] // num_chunks
+    chunks = []
+    for i in range(num_chunks):
+        chunk_dict = {}
+        for key, tensor in tensor_dict.items():
+            if isinstance(tensor, (int, float, bool)):
+                chunk_dict[key] = tensor
+            elif tensor is not None and (isinstance(tensor, list) or tensor.ndim > 0):
+                chunk_dict[key] = tensor[i * chunk_size : (i + 1) * chunk_size]
+            elif tensor is not None and tensor.ndim == 0:
+                chunk_dict[key] = tensor
+            else:
+                chunk_dict[key] = None
+        chunks.append(chunk_dict)
+    return chunks
+
+
+def _patched_shuffle_sequence_dict(seq_dict):
+    """shuffle_sequence_dict that handles scalar types (int/float/bool)."""
+    first_seq = next(
+        v
+        for v in seq_dict.values()
+        if v is not None and isinstance(v, (torch.Tensor, list)) and len(v) > 0
+    )
+    perm = torch.randperm(len(first_seq))
+
+    def permute(v):
+        if v is None:
+            return None
+        if isinstance(v, (int, float, bool)):
+            return v
+        if isinstance(v, torch.Tensor) and v.ndim == 0:
+            return v
+        if isinstance(v, torch.Tensor) and v.ndim >= 1:
+            return v[perm]
+        if isinstance(v, list):
+            return [v[i] for i in perm.tolist()]
+        return v
+
+    return {k: permute(v) for k, v in seq_dict.items()}
+
+
+def _patch_sync_weights_batched(original_init):
+    """Wrap VLLMGeneration.__init__ to accept weight_sync_chunk_size."""
+
+    @wraps(original_init)
+    def patched_init(self, *args, weight_sync_chunk_size=None, **kwargs):
+        original_init(self, *args, **kwargs)
+        self.weight_sync_chunk_size = weight_sync_chunk_size
+
+    return patched_init
+
+
+def _make_batched_sync_weights(original_sync_weights):
+    """Wrap sync_weights to use batched sync for non-FSDP/non-ZeRO paths."""
+
+    @wraps(original_sync_weights)
+    def patched_sync_weights(self):
+        from accelerate.utils import is_peft_model
+
+        # Check if we're in a non-PEFT, non-FSDP, non-ZeRO scenario where batching helps
+        accelerator = self.accelerator
+        model = self.model
+        is_fsdp_enabled = self.is_fsdp_enabled
+
+        deepspeed_plugin = accelerator.state.deepspeed_plugin
+        zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3
+
+        is_peft = is_peft_model(model)
+
+        # If PEFT, FSDP, or ZeRO-3, fall back to original (which handles those cases)
+        if is_peft or is_fsdp_enabled or zero_stage_3:
+            return original_sync_weights(self)
+
+        # Non-PEFT, non-FSDP, non-ZeRO: use batched sync
+        if self.mode == "colocate" and getattr(self, "enable_sleep_mode", False):
+            from vllm.distributed.device_communicators.cuda_wrapper import (
+                empty_cache,
+            )
+
+            empty_cache()
+            self.llm.wake_up(tags=["weights"])
+
+        if self.mode == "server" and accelerator.is_main_process:
+            params = [
+                (self._fix_param_name_to_vllm(name), param.data)
+                for name, param in model.named_parameters()
+            ]
+            self.vllm_client.batch_update_named_params(
+                params, chunk_size=getattr(self, "weight_sync_chunk_size", None)
+            )
+        elif self.mode == "colocate":
+            llm_model = (
+                self.llm.llm_engine.model_executor.driver_worker.model_runner.model
+            )
+            weights = [
+                (self._fix_param_name_to_vllm(name), param.data)
+                for name, param in model.named_parameters()
+            ]
+            llm_model.load_weights(weights=weights)
+
+        # Reset cache
+        if self.mode == "server" and accelerator.is_main_process:
+            self.vllm_client.reset_prefix_cache()
+        elif self.mode == "colocate":
+            self.llm.reset_prefix_cache()
+
+    return patched_sync_weights
+
+
+def patch_trl_vllm():
+    """Apply all TRL vLLM monkeypatches."""
+    import trl.generation.vllm_client
+    import trl.generation.vllm_generation
+    import trl.trainer.utils
+
+    VLLMClient = trl.generation.vllm_client.VLLMClient
+    VLLMGeneration = trl.generation.vllm_generation.VLLMGeneration
+
+    # 1. Add batch_update_named_params to VLLMClient
+    if not hasattr(VLLMClient, "batch_update_named_params"):
+        VLLMClient.batch_update_named_params = _batch_update_named_params
+        VLLMClient.update_model_params = _update_model_params
+        LOG.info("Patched VLLMClient with batch_update_named_params")
+
+    # 2. Patch extract_logprobs (NaN→0.0)
+    trl.generation.vllm_generation.extract_logprobs = _patched_extract_logprobs
+    LOG.info("Patched extract_logprobs with NaN→0.0 fix")
+
+    # 3. Patch VLLMGeneration.__init__ to accept weight_sync_chunk_size
+    VLLMGeneration.__init__ = _patch_sync_weights_batched(VLLMGeneration.__init__)
+
+    # 4. Patch sync_weights for batched non-FSDP/non-ZeRO path
+    VLLMGeneration.sync_weights = _make_batched_sync_weights(
+        VLLMGeneration.sync_weights
+    )
+    LOG.info("Patched VLLMGeneration with batched sync_weights")
+
+    # 5. Patch split_tensor_dict and shuffle_sequence_dict
+    trl.trainer.utils.split_tensor_dict = _patched_split_tensor_dict
+    trl.trainer.utils.shuffle_sequence_dict = _patched_shuffle_sequence_dict
+    LOG.info("Patched split_tensor_dict and shuffle_sequence_dict for scalar types")
diff --git a/src/axolotl/monkeypatch/trainer/utils.py b/src/axolotl/monkeypatch/trainer/utils.py
new file mode 100644
index 0000000000..18a7d0cc7c
--- /dev/null
+++ b/src/axolotl/monkeypatch/trainer/utils.py
@@ -0,0 +1,510 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from axolotl.kernels.op_registry import register_kernel_op
+
+
+@triton.jit
+def _entropy_online_kernel(
+    logits_ptr,
+    output_ptr,
+    stride_row,
+    V: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    """Online entropy: single pass with running max correction."""
+    row = tl.program_id(0)
+    row_ptr = logits_ptr + tl.cast(row, tl.int64) * stride_row
+
+    running_max = tl.full([], float("-inf"), dtype=tl.float32)
+    running_sum_exp = tl.full([], 0.0, dtype=tl.float32)
+    running_weighted = tl.full([], 0.0, dtype=tl.float32)
+
+    for v_start in range(0, V, BLOCK_V):
+        offs = v_start + tl.arange(0, BLOCK_V)
+        mask = offs < V
+        x = tl.load(row_ptr + offs, mask=mask, other=float("-inf")).to(tl.float32)
+
+        block_max = tl.max(x, axis=0)
+        new_max = tl.maximum(running_max, block_max)
+
+        correction = tl.exp(running_max - new_max)
+        running_sum_exp = running_sum_exp * correction
+        running_weighted = running_weighted * correction
+
+        exp_x = tl.exp(x - new_max)
+        exp_x = tl.where(mask, exp_x, 0.0)
+        x = tl.where(mask, x, 0.0)
+        running_sum_exp += tl.sum(exp_x, axis=0)
+        running_weighted += tl.sum(exp_x * x, axis=0)
+
+        running_max = new_max
+
+    entropy = tl.log(running_sum_exp) + running_max - running_weighted / running_sum_exp
+    tl.store(output_ptr + row, entropy)
+
+
+@triton.jit
+def _entropy_online_kernel_strided(
+    logits_ptr,
+    output_ptr,
+    stride_outer,
+    stride_inner,
+    n_inner,
+    row_offset,
+    V: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    """Online entropy for non-contiguous 3D (B, L, V) tensors."""
+    local_row = tl.program_id(0)
+    row = local_row + row_offset
+    outer_idx = row // n_inner
+    inner_idx = row % n_inner
+    off = outer_idx.to(tl.int64) * stride_outer + inner_idx.to(tl.int64) * stride_inner
+    row_ptr = logits_ptr + off
+
+    running_max = tl.full([], float("-inf"), dtype=tl.float32)
+    running_sum_exp = tl.full([], 0.0, dtype=tl.float32)
+    running_weighted = tl.full([], 0.0, dtype=tl.float32)
+
+    for v_start in range(0, V, BLOCK_V):
+        offs = v_start + tl.arange(0, BLOCK_V)
+        mask = offs < V
+        x = tl.load(row_ptr + offs, mask=mask, other=float("-inf")).to(tl.float32)
+
+        block_max = tl.max(x, axis=0)
+        new_max = tl.maximum(running_max, block_max)
+
+        correction = tl.exp(running_max - new_max)
+        running_sum_exp = running_sum_exp * correction
+        running_weighted = running_weighted * correction
+
+        exp_x = tl.exp(x - new_max)
+        exp_x = tl.where(mask, exp_x, 0.0)
+        x = tl.where(mask, x, 0.0)
+        running_sum_exp += tl.sum(exp_x, axis=0)
+        running_weighted += tl.sum(exp_x * x, axis=0)
+
+        running_max = new_max
+
+    entropy = tl.log(running_sum_exp) + running_max - running_weighted / running_sum_exp
+    tl.store(output_ptr + local_row, entropy)
+
+
+@register_kernel_op("entropy_from_logits")
+def _entropy_from_logits_op(logits: torch.Tensor, chunk_size: int) -> torch.Tensor:
+    """CUDA entropy launches; returns flat float32 `(N,)` output."""
+    V = logits.shape[-1]
+    N = 1
+    for s in logits.shape[:-1]:
+        N *= s
+
+    output = torch.empty(N, device=logits.device, dtype=torch.float32)
+
+    BLOCK_V = 4096
+    MAX_GRID_CONTIG = 8192
+    MAX_GRID_STRIDED = 2048
+
+    # Vocab (last) dim must be contiguous for coalesced loads
+    if logits.stride(-1) != 1:
+        logits = logits.contiguous()
+
+    if logits.is_contiguous():
+        flat_logits = logits.reshape(-1, V)
+        stride = flat_logits.stride(0)
+        for start in range(0, N, MAX_GRID_CONTIG):
+            n_rows = min(MAX_GRID_CONTIG, N - start)
+            _entropy_online_kernel[(n_rows,)](
+                flat_logits[start], output[start], stride, V=V, BLOCK_V=BLOCK_V
+            )
+    elif logits.ndim == 3:
+        stride_outer = logits.stride(0)
+        stride_inner = logits.stride(1)
+        n_inner = logits.shape[1]
+        for start in range(0, N, MAX_GRID_STRIDED):
+            n_rows = min(MAX_GRID_STRIDED, N - start)
+            _entropy_online_kernel_strided[(n_rows,)](
+                logits,
+                output[start],
+                stride_outer,
+                stride_inner,
+                n_inner,
+                start,
+                V=V,
+                BLOCK_V=BLOCK_V,
+            )
+    else:
+        logits = logits.contiguous()
+        flat_logits = logits.reshape(-1, V)
+        stride = flat_logits.stride(0)
+        for start in range(0, N, MAX_GRID_CONTIG):
+            n_rows = min(MAX_GRID_CONTIG, N - start)
+            _entropy_online_kernel[(n_rows,)](
+                flat_logits[start], output[start], stride, V=V, BLOCK_V=BLOCK_V
+            )
+
+    return output
+
+
+@_entropy_from_logits_op.register_fake
+def _(logits, chunk_size):
+    N = 1
+    for s in logits.shape[:-1]:
+        N *= s
+    return torch.empty(N, device=logits.device, dtype=torch.float32)
+
+
+def entropy_from_logits(logits: torch.Tensor, chunk_size: int = 128) -> torch.Tensor:
+    """Triton-fused entropy (online single-pass). Handles non-contiguous tensors without copying."""
+    original_shape = logits.shape[:-1]
+
+    if not logits.is_cuda:
+        # CPU fallback: stable entropy via log_softmax
+        logp = F.log_softmax(logits.float(), dim=-1)
+        ent = -(logp.exp() * logp).sum(dim=-1)
+        return ent.to(logits.dtype).reshape(original_shape)
+
+    output = _entropy_from_logits_op(logits, chunk_size)
+    return output.to(logits.dtype).reshape(original_shape)
+
+
+# ---------------------------------------------------------------------------
+# selective_log_softmax — fused forward + backward Triton kernels
+# ---------------------------------------------------------------------------
+
+
+def selective_log_softmax_original(logits, index) -> torch.Tensor:
+    """Original selective_log_softmax (reference/fallback)."""
+    squeeze = index.ndim == logits.ndim - 1
+    if squeeze:
+        index = index.unsqueeze(-1)
+
+    if logits.dtype in [torch.float32, torch.float64]:
+        selected_logits = torch.gather(logits, dim=-1, index=index)
+        logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits])
+        per_token_logps = selected_logits - logsumexp_values.unsqueeze(-1)
+    else:
+        per_token_logps = []
+        for row_logits, row_labels in zip(logits, index, strict=True):
+            row_logps = F.log_softmax(row_logits, dim=-1)
+            row_per_token_logps = row_logps.gather(dim=-1, index=row_labels)
+            per_token_logps.append(row_per_token_logps)
+        per_token_logps = torch.stack(per_token_logps)
+
+    if squeeze:
+        per_token_logps = per_token_logps.squeeze(-1)
+
+    return per_token_logps
+
+
+@triton.jit
+def _selective_logsoftmax_fwd_kernel(
+    logits_ptr,
+    index_ptr,
+    output_ptr,
+    logsumexp_ptr,
+    stride_logits_row,
+    stride_index_row,
+    stride_output_row,
+    actual_K,
+    K_BLOCK: tl.constexpr,
+    V: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    """Forward: online logsumexp + gather. Saves logsumexp for backward."""
+    row = tl.program_id(0)
+    logits_row_ptr = logits_ptr + tl.cast(row, tl.int64) * stride_logits_row
+
+    # Online logsumexp
+    running_max = tl.full([], float("-inf"), dtype=tl.float32)
+    running_sum_exp = tl.full([], 0.0, dtype=tl.float32)
+
+    for v_start in range(0, V, BLOCK_V):
+        offs = v_start + tl.arange(0, BLOCK_V)
+        mask = offs < V
+        x = tl.load(logits_row_ptr + offs, mask=mask, other=float("-inf")).to(
+            tl.float32
+        )
+
+        block_max = tl.max(x, axis=0)
+        new_max = tl.maximum(running_max, block_max)
+        running_sum_exp = running_sum_exp * tl.exp(running_max - new_max)
+
+        exp_x = tl.exp(x - new_max)
+        exp_x = tl.where(mask, exp_x, 0.0)
+        running_sum_exp += tl.sum(exp_x, axis=0)
+        running_max = new_max
+
+    lse = tl.log(running_sum_exp) + running_max
+    tl.store(logsumexp_ptr + row, lse)
+
+    # Gather and subtract
+    index_row_ptr = index_ptr + tl.cast(row, tl.int64) * stride_index_row
+    output_row_ptr = output_ptr + tl.cast(row, tl.int64) * stride_output_row
+
+    k_offs = tl.arange(0, K_BLOCK)
+    k_mask = k_offs < actual_K
+    indices = tl.load(index_row_ptr + k_offs, mask=k_mask, other=0).to(tl.int64)
+    valid_mask = k_mask & (indices >= 0) & (indices < V)
+    safe_indices = tl.where(valid_mask, indices, 0)
+    selected = tl.load(logits_row_ptr + safe_indices, mask=valid_mask, other=0.0).to(
+        tl.float32
+    )
+    tl.store(output_row_ptr + k_offs, selected - lse, mask=valid_mask)
+
+
+@triton.jit
+def _selective_logsoftmax_bwd_kernel(
+    grad_output_ptr,
+    logits_ptr,
+    index_ptr,
+    logsumexp_ptr,
+    grad_logits_ptr,
+    stride_grad_out_row,
+    stride_logits_row,
+    stride_index_row,
+    stride_grad_logits_row,
+    actual_K,
+    K_BLOCK: tl.constexpr,
+    V: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    """Backward: d_logits[j] = -softmax(x)[j] * sum(grad_out) + (grad_out[k] if j == index[k]).
+
+    Single fused pass over V. For each tile, computes the base gradient and adds
+    scatter contributions inline by checking which indices fall in the current tile.
+    No separate scatter pass — no read-after-write issues.
+    """
+    row = tl.program_id(0)
+    logits_row_ptr = logits_ptr + tl.cast(row, tl.int64) * stride_logits_row
+    grad_logits_row_ptr = (
+        grad_logits_ptr + tl.cast(row, tl.int64) * stride_grad_logits_row
+    )
+    grad_out_row_ptr = grad_output_ptr + tl.cast(row, tl.int64) * stride_grad_out_row
+    index_row_ptr = index_ptr + tl.cast(row, tl.int64) * stride_index_row
+
+    lse = tl.load(logsumexp_ptr + row).to(tl.float32)
+
+    # Load grad_output and indices (K_BLOCK elements, masked)
+    k_offs = tl.arange(0, K_BLOCK)
+    k_mask = k_offs < actual_K
+    grad_out = tl.load(grad_out_row_ptr + k_offs, mask=k_mask, other=0.0).to(tl.float32)
+    indices = tl.load(
+        index_row_ptr + k_offs, mask=k_mask, other=-1
+    )  # -1 = never matches
+    valid_mask = k_mask & (indices >= 0) & (indices < V)
+    grad_out = tl.where(valid_mask, grad_out, 0.0)
+    indices = tl.where(valid_mask, indices, -1)
+    grad_sum = tl.sum(grad_out, axis=0)
+
+    # Fused pass: for each tile, compute -softmax * grad_sum + scatter
+    for v_start in range(0, V, BLOCK_V):
+        offs = v_start + tl.arange(0, BLOCK_V)  # [BLOCK_V]
+        mask = offs < V
+        x = tl.load(logits_row_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        softmax_j = tl.exp(x - lse)
+        softmax_j = tl.where(mask, softmax_j, 0.0)
+        grad_j = -softmax_j * grad_sum
+
+        # Scatter: check which selected indices fall in this tile
+        # offs: [BLOCK_V], indices: [K_BLOCK]
+        # Broadcast: offs[:, None] == indices[None, :] → [BLOCK_V, K_BLOCK]
+        match = offs[:, None] == indices[None, :]  # [BLOCK_V, K_BLOCK]
+        # Sum grad_out contributions: for each position j, sum grad_out[k] where index[k]==j
+        scatter_contrib = tl.sum(
+            tl.where(match, grad_out[None, :], 0.0), axis=1
+        )  # [BLOCK_V]
+        grad_j += scatter_contrib
+
+        tl.store(grad_logits_row_ptr + offs, grad_j, mask=mask)
+
+
+@register_kernel_op("selective_log_softmax_fwd")
+def _selective_log_softmax_fwd_op(
+    flat_logits: torch.Tensor,
+    flat_index: torch.Tensor,
+    K: int,
+    K_BLOCK: int,
+    V: int,
+    BLOCK_V: int,
+    MAX_GRID: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    N = flat_logits.shape[0]
+    output = torch.empty(N, K_BLOCK, device=flat_logits.device, dtype=torch.float32)
+    logsumexp = torch.empty(N, device=flat_logits.device, dtype=torch.float32)
+
+    for start in range(0, N, MAX_GRID):
+        n_rows = min(MAX_GRID, N - start)
+        _selective_logsoftmax_fwd_kernel[(n_rows,)](
+            flat_logits[start],
+            flat_index[start],
+            output[start],
+            logsumexp[start],
+            flat_logits.stride(0),
+            flat_index.stride(0),
+            output.stride(0),
+            K,
+            K_BLOCK=K_BLOCK,
+            V=V,
+            BLOCK_V=BLOCK_V,
+        )
+
+    return output, logsumexp
+
+
+@_selective_log_softmax_fwd_op.register_fake
+def _(flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID):
+    N = flat_logits.shape[0]
+    return (
+        torch.empty(N, K_BLOCK, device=flat_logits.device, dtype=torch.float32),
+        torch.empty(N, device=flat_logits.device, dtype=torch.float32),
+    )
+
+
+@register_kernel_op("selective_log_softmax_bwd")
+def _selective_log_softmax_bwd_op(
+    grad_output: torch.Tensor,
+    flat_logits: torch.Tensor,
+    flat_index: torch.Tensor,
+    logsumexp: torch.Tensor,
+    K: int,
+    K_BLOCK: int,
+    V: int,
+    BLOCK_V: int,
+    MAX_GRID: int,
+) -> torch.Tensor:
+    N = flat_logits.shape[0]
+
+    grad_logits = torch.empty_like(flat_logits)
+
+    # grad_output may have K_BLOCK cols; backward kernel reads actual_K
+    grad_output_contig = grad_output.contiguous()
+
+    for start in range(0, N, MAX_GRID):
+        n_rows = min(MAX_GRID, N - start)
+        _selective_logsoftmax_bwd_kernel[(n_rows,)](
+            grad_output_contig[start],
+            flat_logits[start],
+            flat_index[start],
+            logsumexp[start],
+            grad_logits[start],
+            grad_output_contig.stride(0),
+            flat_logits.stride(0),
+            flat_index.stride(0),
+            grad_logits.stride(0),
+            K,
+            K_BLOCK=K_BLOCK,
+            V=V,
+            BLOCK_V=BLOCK_V,
+        )
+
+    return grad_logits
+
+
+@_selective_log_softmax_bwd_op.register_fake
+def _(
+    grad_output, flat_logits, flat_index, logsumexp, K, K_BLOCK, V, BLOCK_V, MAX_GRID
+):
+    return torch.empty_like(flat_logits)
+
+
+class _SelectiveLogSoftmaxTriton(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID):
+        output, logsumexp = _selective_log_softmax_fwd_op(
+            flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID
+        )
+
+        ctx.save_for_backward(flat_logits, flat_index, logsumexp)
+        ctx.K = K
+        ctx.K_BLOCK = K_BLOCK
+        ctx.V = V
+        ctx.BLOCK_V = BLOCK_V
+        ctx.MAX_GRID = MAX_GRID
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        flat_logits, flat_index, logsumexp = ctx.saved_tensors
+        K, K_BLOCK, V, BLOCK_V, MAX_GRID = (
+            ctx.K,
+            ctx.K_BLOCK,
+            ctx.V,
+            ctx.BLOCK_V,
+            ctx.MAX_GRID,
+        )
+        grad_logits = _selective_log_softmax_bwd_op(
+            grad_output,
+            flat_logits,
+            flat_index,
+            logsumexp,
+            K,
+            K_BLOCK,
+            V,
+            BLOCK_V,
+            MAX_GRID,
+        )
+
+        # Return grads for: flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID
+        return grad_logits, None, None, None, None, None, None
+
+
+def selective_log_softmax(logits, index) -> torch.Tensor:
+    """
+    Fused selective_log_softmax with Triton forward+backward kernels.
+
+    Equivalent to: torch.gather(logits.log_softmax(-1), dim=-1, index=index)
+    """
+    squeeze = index.ndim == logits.ndim - 1
+    if squeeze:
+        index = index.unsqueeze(-1)
+
+    if not logits.is_cuda or logits.dtype == torch.float64:
+        # Triton kernel computes in float32; fall back for float64 and CPU
+        return selective_log_softmax_original(
+            logits, index.squeeze(-1) if squeeze else index
+        )
+
+    V = logits.shape[-1]
+    K = index.shape[-1]
+    original_index_shape = index.shape
+
+    try:
+        flat_logits = logits.view(-1, V)
+    except RuntimeError:
+        flat_logits = logits.reshape(-1, V).contiguous()
+    flat_index = index.reshape(-1, K).contiguous()
+
+    BLOCK_V = 4096
+    MAX_GRID = 8192
+    K_BLOCK = max(1, triton.next_power_of_2(K))
+
+    output = _SelectiveLogSoftmaxTriton.apply(
+        flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID
+    )
+
+    if K_BLOCK != K:
+        output = output[:, :K]
+
+    per_token_logps = output.to(logits.dtype).reshape(original_index_shape)
+
+    if squeeze:
+        per_token_logps = per_token_logps.squeeze(-1)
+
+    return per_token_logps
diff --git a/src/axolotl/monkeypatch/trainer_accelerator_args.py b/src/axolotl/monkeypatch/trainer_accelerator_args.py
new file mode 100644
index 0000000000..9fc6e38c68
--- /dev/null
+++ b/src/axolotl/monkeypatch/trainer_accelerator_args.py
@@ -0,0 +1,83 @@
+"""
+allow adding additional kwargs to Accelerator init
+"""
+
+import inspect
+
+from transformers import Trainer
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+ORIGINAL_TRAINER_CODE = """
+    # create accelerator object
+    self.accelerator = Accelerator(**args)
+"""
+
+PATCHED_TRAINER_CODE = """
+    if hasattr(self, "additional_accelerator_args"):
+        additional_args = self.additional_accelerator_args(fp8=True, enable_fsdp_float8_all_gather={enable_fsdp_float8_all_gather}, **args)
+        if additional_args:
+            args.update(additional_args)
+
+    # create accelerator object
+    self.accelerator = Accelerator(**args)
+"""
+
+
+def get_create_accelerate_code() -> str:
+    training_loop = inspect.getsource(Trainer.create_accelerator_and_postprocess)
+    return training_loop
+
+
+def check_create_accelerate_code_is_patchable() -> bool:
+    create_code = get_create_accelerate_code()
+    create_code, _ = detab_code(create_code)
+    return ORIGINAL_TRAINER_CODE in create_code
+
+
+def patch_create_accelerate_code_for_fp8(enable_fsdp_float8_all_gather: bool):
+    """
+    Monkeypatch create_accelerator_and_postprocess so it checks for additional kwargs.
+    """
+
+    try:
+        create_code = get_create_accelerate_code()
+    except OSError:
+        return
+    Trainer._original_create_accelerator_and_postprocess = create_code
+    create_code, _ = detab_code(create_code)
+    if ORIGINAL_TRAINER_CODE not in create_code:
+        return
+
+    patched_trainer_code = PATCHED_TRAINER_CODE.format(
+        enable_fsdp_float8_all_gather=enable_fsdp_float8_all_gather
+    )
+    create_code = create_code.replace(ORIGINAL_TRAINER_CODE, patched_trainer_code)
+    create_code = create_code.replace(
+        "def create_accelerator_and_postprocess(",
+        "def fixed_create_accelerator_and_postprocess(",
+        1,
+    )
+
+    # load imports necessary
+    import transformers.trainer
+
+    items_to_import = []
+    for item in dir(transformers.trainer):
+        if item in create_code:
+            items_to_import.append(item)
+
+    exec(
+        "from transformers.trainer import ("
+        + ", ".join(x for x in items_to_import)
+        + ")",
+        globals(),
+    )
+    exec(create_code, globals())
+    LOG.info("patching create_accelerator_and_postprocess to allow for overrides")
+    Trainer.create_accelerator_and_postprocess = (
+        fixed_create_accelerator_and_postprocess
+    )
diff --git a/src/axolotl/monkeypatch/transformers/__init__.py b/src/axolotl/monkeypatch/transformers/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/monkeypatch/transformers/trainer_loss_calc.py b/src/axolotl/monkeypatch/transformers/trainer_loss_calc.py
new file mode 100644
index 0000000000..7704ceddc3
--- /dev/null
+++ b/src/axolotl/monkeypatch/transformers/trainer_loss_calc.py
@@ -0,0 +1,149 @@
+"""
+Module for patching transformers Trainer loss calculation to use nanmean.
+
+This is needed for context parallelism since chunks of the input sequences may be fully
+masked and return NaNs in the loss calculation.
+
+Also includes a patch for FSDP2 + torch.compile. We need to bundle this together with
+the other evaluation_loop patch because we can't patch the same code twice without
+raising an OSError.
+"""
+
+import importlib
+import inspect
+
+from transformers import Trainer
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+ORIGINAL_EVAL_CODE = {
+    "list": 'metrics[f"{metric_key_prefix}_loss"] = np.concatenate(all_losses).mean().item()',
+    "array": 'metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()',
+}
+PATCHED_EVAL_CODE = {
+    "list": 'metrics[f"{metric_key_prefix}_loss"] = np.nanmean(np.concatenate(all_losses)).item()',
+    "array": 'metrics[f"{metric_key_prefix}_loss"] = np.nanmean(all_losses).item()',
+}
+
+ORIGINAL_MAYBE_CODE = (
+    "tr_loss_scalar = nested_gather(tr_loss, self.args.parallel_mode).mean().item()"
+)
+PATCHED_MAYBE_CODE = (
+    "tr_loss_scalar = nested_gather(tr_loss, self.args.parallel_mode).nanmean().item()"
+)
+
+
+def check_evaluation_loop_is_patchable() -> bool:
+    if hasattr(Trainer, "_original_evaluation_loop"):
+        evaluation_loop_source = Trainer._original_evaluation_loop
+    else:
+        evaluation_loop_source = inspect.getsource(Trainer.evaluation_loop)
+    return all(value in evaluation_loop_source for value in ORIGINAL_EVAL_CODE.values())
+
+
+def patch_evaluation_loop():
+    """Patch the evaluation_loop method."""
+    # Check if already patched
+    if hasattr(Trainer, "_original_evaluation_loop"):
+        LOG.debug("Trainer.evaluation_loop already patched")
+        return
+
+    # Check if the patterns exist
+    try:
+        evaluation_loop_source = inspect.getsource(Trainer.evaluation_loop)
+    except OSError:
+        return
+    Trainer._original_evaluation_loop = evaluation_loop_source
+    evaluation_loop_source, _ = detab_code(evaluation_loop_source)
+
+    # Apply the nanmean patches
+    evaluation_loop_source = evaluation_loop_source.replace(
+        ORIGINAL_EVAL_CODE["list"], PATCHED_EVAL_CODE["list"]
+    )
+    evaluation_loop_source = evaluation_loop_source.replace(
+        ORIGINAL_EVAL_CODE["array"], PATCHED_EVAL_CODE["array"]
+    )
+
+    # Rename the function to avoid conflicts
+    evaluation_loop_source = evaluation_loop_source.replace(
+        "def evaluation_loop(",
+        "def axolotl_evaluation_loop(",
+        1,
+    )
+
+    # Get the module for necessary imports
+    module_name = Trainer.__module__
+    module = importlib.import_module(module_name)
+
+    # Import necessary items from the module
+    items_to_import = []
+    for item in dir(module):
+        if item in evaluation_loop_source:
+            items_to_import.append(item)
+
+    # Execute the imports and patched method
+    exec(
+        f"from {module_name} import ({', '.join(items_to_import)})",
+        globals(),
+    )
+    exec(evaluation_loop_source, globals())
+
+    LOG.debug("Patched Trainer.evaluation_loop with nanmean loss calculation")
+    Trainer.evaluation_loop = axolotl_evaluation_loop
+
+
+def check_maybe_log_save_evaluate_is_patchable() -> bool:
+    if hasattr(Trainer, "_original_maybe_log_save_evaluate"):
+        maybe_log_source = Trainer._original_maybe_log_save_evaluate
+    else:
+        maybe_log_source = inspect.getsource(Trainer._maybe_log_save_evaluate)
+    return ORIGINAL_MAYBE_CODE in maybe_log_source
+
+
+def patch_maybe_log_save_evaluate():
+    """Patch the _maybe_log_save_evaluate method."""
+    # Check if already patched
+    if hasattr(Trainer, "_original_maybe_log_save_evaluate"):
+        LOG.info("Trainer._maybe_log_save_evaluate already patched")
+        return
+
+    # Check if the patterns exist
+    try:
+        maybe_log_source = inspect.getsource(Trainer._maybe_log_save_evaluate)
+    except OSError:
+        return
+    Trainer._original_maybe_log_save_evaluate = maybe_log_source
+    maybe_log_source, _ = detab_code(maybe_log_source)
+
+    # Apply the patch
+    maybe_log_source = maybe_log_source.replace(ORIGINAL_MAYBE_CODE, PATCHED_MAYBE_CODE)
+
+    # Rename the function to avoid conflicts
+    maybe_log_source = maybe_log_source.replace(
+        "def _maybe_log_save_evaluate(",
+        "def axolotl_maybe_log_save_evaluate(",
+        1,
+    )
+
+    # Get the module for necessary imports
+    module_name = Trainer.__module__
+    module = importlib.import_module(module_name)
+
+    # Import necessary items from the module
+    items_to_import = []
+    for item in dir(module):
+        if item in maybe_log_source:
+            items_to_import.append(item)
+
+    # Execute the imports and patched method
+    exec(
+        f"from {module_name} import ({', '.join(items_to_import)})",
+        globals(),
+    )
+    exec(maybe_log_source, globals())
+
+    LOG.debug("Patched Trainer._maybe_log_save_evaluate with nanmean loss calculation")
+    Trainer._maybe_log_save_evaluate = axolotl_maybe_log_save_evaluate
diff --git a/src/axolotl/monkeypatch/transformers_fa_utils.py b/src/axolotl/monkeypatch/transformers_fa_utils.py
new file mode 100644
index 0000000000..e372dc3f85
--- /dev/null
+++ b/src/axolotl/monkeypatch/transformers_fa_utils.py
@@ -0,0 +1,68 @@
+"""
+see https://github.com/huggingface/transformers/pull/35834
+"""
+
+from functools import partial
+from typing import Optional
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def fixed_fa_peft_integration_check(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    target_dtype: Optional[torch.dtype] = None,
+    preferred_dtype: Optional[torch.dtype] = None,
+):
+    """
+    PEFT usually casts the layer norms in float32 for training stability reasons
+    therefore the input hidden states gets silently casted in float32. Hence, we need
+    cast them back in float16 / bfloat16 just to be sure everything works as expected.
+    This might slowdown training & inference so it is recommended to not cast the LayerNorms!
+
+    Args:
+        query (`torch.Tensor`):
+            Input query states to be passed to Flash Attention API
+        key (`torch.Tensor`):
+            Input key states to be passed to Flash Attention API
+        value (`torch.Tensor`):
+            Input value states to be passed to Flash Attention API
+        target_dtype (`torch.dtype`, *optional*):
+            The dtype to convert the attention tensors to. Conversion can be ignored by
+            not providing the target dtype.
+        preferred_dtype (`torch.dtype`, *optional*):
+            The preferred dtype to convert the attention tensors to regardless of the
+            target dtype.
+    """
+    if target_dtype is None and preferred_dtype is None:
+        return query, key, value
+
+    if preferred_dtype and target_dtype != preferred_dtype:
+        target_dtype = preferred_dtype
+
+    # check if any of query, key, or value are in float32. If so, cast them back to target dtype.
+    if any(module.dtype == torch.float32 for module in [query, key, value]):
+        logger.warning_once(
+            f"The input hidden states seems to be silently casted in float32, this might be related to"
+            f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+            f" {target_dtype}."
+        )
+
+        query = query.to(target_dtype)
+        key = key.to(target_dtype)
+        value = value.to(target_dtype)
+
+    return query, key, value
+
+
+def patch_fa_peft_integration():
+    import transformers.modeling_flash_attention_utils
+
+    transformers.modeling_flash_attention_utils.fa_peft_integration_check = partial(
+        fixed_fa_peft_integration_check, preferred_dtype=None
+    )
diff --git a/src/axolotl/monkeypatch/utils.py b/src/axolotl/monkeypatch/utils.py
index e43c58650a..4ae37b2f6b 100644
--- a/src/axolotl/monkeypatch/utils.py
+++ b/src/axolotl/monkeypatch/utils.py
@@ -1,43 +1,55 @@
 """
 Shared utils for the monkeypatches
 """
-from typing import Optional
+
+import re
+from typing import Tuple
 
 import torch
 import torch.nn.functional as F
-from transformers.modeling_attn_mask_utils import (
-    _prepare_4d_causal_attention_mask,
-    _prepare_4d_causal_attention_mask_for_sdpa,
-)
-from transformers.utils import is_torch_bf16_gpu_available
 
 
-@torch.jit.script
 def get_max_seqlen_in_batch(attention_mask: torch.Tensor) -> torch.Tensor:
-    max_num = int(torch.max(attention_mask).item())
-    batch_size, _ = attention_mask.shape
-    counts = torch.zeros((batch_size, max_num), dtype=torch.int32)
+    """Token counts for each packed sub-sequence in a multipack attention mask.
 
-    for i in range(1, max_num + 1):
-        mask = attention_mask == i
-        counts[:, i - 1] = torch.sum(mask, dim=-1).to(dtype=torch.int32)
+    Multipack encodes every packed document as a distinct positive integer id
+    (``1, 2, 3, ...``) with ``0`` for padding, so a row looks like
+    ``[1, 1, 1, 2, 2, 3, 3, 0, 0]``. This returns the length of each non-pad
+    document, flattened across the batch.
 
-    result = counts.flatten()
-    nonzero_indices = torch.nonzero(result).squeeze(-1)
-    return result[nonzero_indices]
+    Vectorized on purpose: the previous implementation used
+    ``int(mask.max().item())`` as a Python loop bound, which forced a
+    device->host sync and broke torch.compile. This variant computes segment
+    boundaries with tensor ops only.
+    """
+    bsz, slen = attention_mask.shape
+    flat = attention_mask.reshape(-1)
+    # A token begins a new document when its id differs from the previous token.
+    seg_start = torch.ones_like(flat, dtype=torch.bool)
+    seg_start[1:] = flat[1:] != flat[:-1]
+    # Force a boundary at every row start so equal ids in adjacent rows (e.g. two
+    # rows each fully filled by a single document) are not merged into one segment.
+    seg_start = seg_start.view(bsz, slen)
+    seg_start[:, 0] = True
+    seg_start = seg_start.reshape(-1)
+
+    start_idx = seg_start.nonzero().flatten()
+    end_idx = torch.cat([start_idx[1:], start_idx.new_full((1,), flat.numel())])
+    lengths = (end_idx - start_idx).to(dtype=torch.int32)
+    # Drop padding runs (id == 0).
+    return lengths[flat[start_idx] != 0]
 
 
-@torch.jit.script
 def get_unpad_data(attention_mask: torch.Tensor):
     device = attention_mask.device
     seqlens_in_batch = get_max_seqlen_in_batch(attention_mask)
     indices = torch.nonzero(attention_mask.flatten()).flatten()
+    # flash_attn's varlen API wants a Python int for max_seqlen; this is the only
+    # remaining device->host sync and is kept solely for that ABI.
     max_seqlen_in_batch = seqlens_in_batch.max().item()
-    cu_seqlens = (
-        F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
-        .to(device=device)
-        .detach()
-    )
+    cu_seqlens = F.pad(
+        torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
+    ).to(device=device)
     return (
         indices,
         cu_seqlens,
@@ -96,7 +108,9 @@ def get_cu_seqlens(attn_mask):
     return torch.stack(results).to(dtype=torch.int32), torch.stack(max_seq_lens)
 
 
-def get_cu_seqlens_from_pos_ids(position_ids):
+def get_cu_seqlens_from_pos_ids(
+    position_ids: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """generate a cumulative sequence length mask for flash attention using pos ids"""
     if len(position_ids.shape) == 1:
         position_ids = position_ids.unsqueeze(0)
@@ -168,60 +182,10 @@ def set_module_name(model, name, value):
     setattr(parent, child_name, value)
 
 
-def mask_2d_to_4d(
-    mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None
-):
-    """
-    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
-    This expansion handles packed sequences so that sequences share the same attention mask integer value
-    when they attend to each other within that sequence.
-    This expansion transforms the mask to lower triangular form to prevent future peeking.
-    """
-    bsz, src_len = mask.size()
-    tgt_len = tgt_len if tgt_len is not None else src_len
-
-    mask = mask.unsqueeze(1).unsqueeze(2)
-    mask = mask.expand(bsz, 1, tgt_len, src_len)
-
-    # Create a binary mask from the original mask where zeros remain zeros and all other values are set to one
-    binary_mask = torch.where(
-        mask != 0,
-        torch.tensor(1, device=mask.device).to(dtype),
-        torch.tensor(0, device=mask.device).to(dtype),
-    )
-
-    # Create a block-diagonal mask.
-    # we multiply by the binary mask so that 0's in the original mask are correctly excluded
-    zero_one_mask = torch.eq(mask, mask.transpose(-1, -2)).int() * binary_mask
-
-    # Now let's create a lower triangular mask of ones that will zero out the upper triangular part
-    lower_triangular_ones = torch.tril(torch.ones((tgt_len, src_len), dtype=dtype)).to(
-        mask.device
-    )
-
-    # Use the lower triangular mask to zero out the upper triangular part of the zero_one_mask
-    masked_zero_one_mask = zero_one_mask * lower_triangular_ones
-
-    return masked_zero_one_mask
-
-
-def patched_prepare_4d_causal_attention_mask(
-    attention_mask: Optional[torch.Tensor],
-    *args,
-):
-    dtype = torch.bfloat16 if is_torch_bf16_gpu_available() else torch.float32
-    return _prepare_4d_causal_attention_mask(
-        mask_2d_to_4d(attention_mask, dtype=dtype),
-        *args,
-    )
-
-
-def patched_prepare_4d_causal_attention_mask_for_sdpa(
-    attention_mask: Optional[torch.Tensor],
-    *args,
-):
-    dtype = torch.bfloat16 if is_torch_bf16_gpu_available() else torch.float32
-    return _prepare_4d_causal_attention_mask_for_sdpa(
-        mask_2d_to_4d(attention_mask, dtype=dtype),
-        *args,
-    )
+def detab_code(code: str) -> Tuple[str, str]:
+    try:
+        spaces = re.match(r"([\s\t]{1,})", code).group(0)
+        code = re.sub(r"^" + spaces, "", code, flags=re.MULTILINE)
+    except AttributeError:
+        return code, ""
+    return code, spaces
diff --git a/src/axolotl/monkeypatch/xformers_/__init__.py b/src/axolotl/monkeypatch/xformers_/__init__.py
new file mode 100644
index 0000000000..6f5b43f773
--- /dev/null
+++ b/src/axolotl/monkeypatch/xformers_/__init__.py
@@ -0,0 +1,52 @@
+"""
+Fused MLP layer for incrementally improved training efficiency
+"""
+
+import torch
+from transformers.models.llama.modeling_llama import LlamaMLP
+from xformers.ops import SwiGLU
+
+from axolotl.monkeypatch.utils import set_module_name
+
+
+class FusedMLP(torch.nn.Module):
+    """
+    Fused MLP layer for incrementally improved training efficiency
+    """
+
+    def __init__(
+        self,
+        config,
+        gate_proj: torch.nn.Linear,
+        up_proj: torch.nn.Linear,
+        down_proj: torch.nn.Linear,
+    ):
+        super().__init__()
+        self.config = config
+        self.swiglu = SwiGLU(
+            in_features=config.hidden_size,
+            hidden_features=config.intermediate_size,
+            bias=False,
+            _pack_weights=True,
+        )
+        # overwrite initialized weights with pretrained weights
+        self.swiglu.w12.weight.data = torch.cat(
+            (gate_proj.weight.data, up_proj.weight.data), dim=0
+        )
+        self.swiglu.w3.weight.data = down_proj.weight.data
+
+    def _post_training(self, model, name):
+        w1, w2 = torch.split(
+            self.swiglu.w12.weight.data, self.config.intermediate_size, dim=0
+        )
+
+        # Assign the split weights back to the original layers
+        new_mlp = LlamaMLP(self.config)
+        new_mlp.gate_proj.weight.data = w1
+        new_mlp.up_proj.weight.data = w2
+        new_mlp.down_proj.weight.data = self.swiglu.w3.weight.data
+
+        set_module_name(model, name, new_mlp)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.swiglu(x)
diff --git a/src/axolotl/processing_strategies.py b/src/axolotl/processing_strategies.py
new file mode 100644
index 0000000000..09d97ccad2
--- /dev/null
+++ b/src/axolotl/processing_strategies.py
@@ -0,0 +1,1523 @@
+"""Module containing ProcessingStrategy classes and its derivative for different MultiModal Model types"""
+
+import bisect
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+from PIL import Image, ImageOps
+from PIL.Image import Resampling
+from torch import Tensor, zeros_like
+from transformers import ProcessorMixin
+from transformers.image_utils import load_image
+from transformers.models.internvl import InternVLProcessor
+from transformers.models.smolvlm import SmolVLMProcessor
+from transformers.models.voxtral import VoxtralProcessor
+
+from axolotl.utils.dict import remove_none_values
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# One-shot warning dedupe so opt-out subclasses don't spam per-batch.
+_ROLE_MASK_WARNED: set[str] = set()
+
+# Supported values for ``train_on_eos`` — mirrors the text-only
+# ChatTemplateStrategy (``turn`` = trainable turn ends only, ``all`` = every
+# turn end, ``none`` = never, ``last`` = only the final trainable turn end).
+_VALID_TRAIN_ON_EOS = ("turn", "all", "none", "last")
+
+
+@dataclass(frozen=True)
+class RoleBoundary:
+    """One role's token-level span markers for the masking scanner.
+
+    Empty ``end_tokens`` means end-of-sequence terminates the span.
+    """
+
+    role: str
+    start_tokens: list[int]
+    end_tokens: list[int] = field(default_factory=list)
+    include_start: bool = False
+    include_end: bool = True
+
+
+class ProcessingStrategy:
+    """Base Processing Strategy class.
+
+    Subclasses opt in to role masking by overriding ``_build_role_boundaries``;
+    otherwise only pad + media tokens are masked (legacy behavior, one-shot warned).
+    """
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        self.processor = processor
+        self.chat_template = chat_template
+        self.image_token = None
+        self.image_token_id = None
+        self.field_messages = self._normalize_field_messages(field_messages)
+
+        self.image_size = image_size
+        self.image_resize_algorithm = (
+            image_resize_algorithm or Image.Resampling.BILINEAR
+        )
+
+        # Defaults mirror the text-only ChatTemplateStrategy. An explicit
+        # empty list is honored as "no trainable roles" (masks everything);
+        # only ``None`` falls back to the default of assistant-only.
+        self.train_on_inputs = bool(train_on_inputs)
+        self.roles_to_train = (
+            list(roles_to_train) if roles_to_train is not None else ["assistant"]
+        )
+        self.train_on_eos = train_on_eos if train_on_eos is not None else "turn"
+        if self.train_on_eos not in _VALID_TRAIN_ON_EOS:
+            raise ValueError(
+                f"train_on_eos={self.train_on_eos!r} is not one of "
+                f"{_VALID_TRAIN_ON_EOS}."
+            )
+
+        if hasattr(processor, "image_token"):
+            self.image_token = processor.image_token
+            self.image_token_id = processor.tokenizer.convert_tokens_to_ids(
+                self.image_token
+            )
+
+        built_in = self._build_role_boundaries()
+
+        # Truthiness check: empty list == unset (opt-in escape hatch), so
+        # `role_boundaries: []` in YAML falls through to built-ins instead of
+        # producing all-masked labels.
+        if role_boundaries_override:
+            overridden = _resolve_role_boundary_override(
+                role_boundaries_override, self.processor.tokenizer
+            )
+            LOG.info(
+                "%s: overriding built-in role boundaries (%d decls) "
+                "with cfg.role_boundaries (%d decls).",
+                type(self).__name__,
+                len(built_in),
+                len(overridden),
+            )
+            self.role_boundaries: list[RoleBoundary] = overridden
+            source = "override"
+        else:
+            self.role_boundaries = built_in
+            source = "built-in"
+
+        # Single-line, grep-friendly summary of the resolved masking config so
+        # "why isn't masking firing?" is visible in training logs. For
+        # overrides we include the fully resolved (role, start_ids, end_ids)
+        # tuples; for built-ins we log a count (subclasses vary and logging
+        # every id sequence would be noisy on, e.g., Llama3 with five roles).
+        boundaries_repr: str | list[tuple[str, list[int], list[int]]]
+        if source == "override":
+            boundaries_repr = [
+                (b.role, b.start_tokens, b.end_tokens) for b in self.role_boundaries
+            ]
+        else:
+            boundaries_repr = f"{len(self.role_boundaries)} built-in"
+        LOG.info(
+            "ProcessingStrategy init: class=%s train_on_inputs=%s "
+            "roles_to_train=%s train_on_eos=%s boundaries_source=%s "
+            "boundaries=%s",
+            type(self).__name__,
+            self.train_on_inputs,
+            self.roles_to_train,
+            self.train_on_eos,
+            source,
+            boundaries_repr,
+        )
+
+    def _build_role_boundaries(self) -> list[RoleBoundary]:
+        """Subclasses declare role boundaries here; [] opts out of role masking."""
+        return []
+
+    @staticmethod
+    def _normalize_field_messages(
+        field_messages: str | list[str] | tuple[str, ...] | None,
+    ) -> tuple[str, ...]:
+        if field_messages is None:
+            return ("messages",)
+        if isinstance(field_messages, str):
+            return (field_messages,)
+        return tuple(name for name in field_messages if name)
+
+    def _get_messages_field(self, example: dict) -> str | None:
+        # Configured field wins so a stale `messages` column can't override it.
+        for name in self.field_messages:
+            if name in example and example[name] is not None:
+                return name
+        if "messages" in example and example["messages"] is not None:
+            return "messages"
+        return None
+
+    @staticmethod
+    def _is_legacy_schema(messages) -> bool:
+        """Detect ShareGPT schema: first message has both ``from`` and ``value``."""
+        return (
+            isinstance(messages, list)
+            and bool(messages)
+            and isinstance(messages[0], dict)
+            and "from" in messages[0]
+            and "value" in messages[0]
+        )
+
+    def __call__(self, examples: list[dict]) -> list[dict]:
+        """Normalize examples to OpenAI ``messages`` (accepts legacy ``conversations`` or custom ``field_messages``)."""
+        role_mapping = {
+            "human": "user",
+            "gpt": "assistant",
+        }
+
+        def normalize_role(role: str) -> str:
+            return role_mapping.get(role, role)
+
+        def convert_legacy_format(example: dict) -> dict:
+            messages = [
+                {"role": normalize_role(convo["from"]), "content": convo["value"]}
+                for convo in example["conversations"]
+            ]
+            result = deepcopy(example)
+            result.pop("conversations")
+            result["messages"] = messages
+            return result
+
+        def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]:
+            new_messages = []
+            for message in messages:
+                if isinstance(message["content"], str):
+                    new_messages.append(
+                        {
+                            "role": message["role"],
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": message["content"],
+                                }
+                            ],
+                        }
+                    )
+                elif isinstance(message["content"], list):
+                    content = message["content"]
+
+                    new_messages.append(
+                        {
+                            "role": message["role"],
+                            "content": content,
+                        }
+                    )
+
+            return new_messages
+
+        processed_examples = []
+        for example in examples:
+            messages_field = self._get_messages_field(example)
+            # Re-route a custom field into the canonical key whose schema it matches;
+            # the canonical "messages"/"conversations" branches below are unchanged.
+            if messages_field and messages_field not in {"messages", "conversations"}:
+                msgs = example[messages_field]
+                target = "conversations" if self._is_legacy_schema(msgs) else "messages"
+                example = dict(example)
+                example[target] = msgs
+                example.pop(messages_field, None)
+                if target == "conversations":
+                    example.pop("messages", None)
+
+            if not ("messages" in example or "conversations" in example):
+                raise ValueError(
+                    "Only configured `field_messages`, `messages`, and `conversations` message keys are currently supported."
+                )
+
+            if "messages" in example and example["messages"] is not None:
+                # Deepcopy for symmetry with convert_legacy_format (which
+                # deepcopies internally) so downstream mutations of
+                # processed_example don't leak back to the caller's input.
+                processed_example = deepcopy(example)
+            elif "conversations" in example:
+                processed_example = convert_legacy_format(example)
+            else:
+                # `messages` is present but None, and no `conversations`
+                # fallback exists — convert_legacy_format would KeyError on
+                # ["conversations"]. Surface a clear validation error instead.
+                raise ValueError(
+                    "`messages` is present but None; provide non-null "
+                    "`messages` or a `conversations` field."
+                )
+
+            # Required for apply_chat_template compatibility.
+            processed_example["messages"] = convert_messages_to_multimedia_messages(
+                processed_example["messages"]
+            )
+
+            possible_image_keys = ["images", "image"]
+            image_key = None
+            for key in possible_image_keys:
+                if key in processed_example:
+                    image_key = key
+                    break
+
+            if image_key is not None and processed_example[image_key] is not None:
+                # TODO: support multi-image samples; for now we take the first.
+                if len(processed_example[image_key]) > 1:
+                    LOG.warning(
+                        f"Found {len(processed_example[image_key])} images in a sample. Using the first one."
+                        "If you are using a dataset with multiple images per sample, please convert it to use multi-content Messages."
+                        "See https://docs.axolotl.ai/docs/multimodal.html#dataset-format"
+                    )
+
+                image_value = processed_example[image_key][0]
+
+                image_value = load_image(image_value)
+
+                if self.image_size is not None:
+                    assert hasattr(image_value, "resize"), (
+                        "Image does not have a resize method"
+                    )
+
+                    if isinstance(self.image_size, tuple):
+                        image_value = image_value.resize(
+                            self.image_size, self.image_resize_algorithm
+                        )
+                    else:
+                        # Int image_size: preserve aspect ratio then pad to square (black) to avoid distortion.
+                        padding_color = (0, 0, 0)
+                        image_value = ImageOps.pad(
+                            image_value,
+                            (self.image_size, self.image_size),
+                            method=self.image_resize_algorithm,
+                            color=padding_color,
+                        )
+
+                msg_ind_to_add = None
+                ind_to_add = None
+                first_user_idx = None
+
+                for msg_idx, msg_content in enumerate(processed_example["messages"]):
+                    if first_user_idx is None and msg_content["role"] == "user":
+                        first_user_idx = msg_idx
+                    for i, content in enumerate(
+                        processed_example["messages"][msg_idx]["content"]
+                    ):
+                        # Column-image datasets often leave a bare {type: "image"} placeholder.
+                        if content["type"] == "image" and all(
+                            k not in content for k in ["image", "url", "path", "base64"]
+                        ):
+                            msg_ind_to_add = msg_idx
+                            ind_to_add = i
+                            break
+
+                if ind_to_add is not None and msg_ind_to_add is not None:
+                    processed_example["messages"][msg_ind_to_add]["content"][
+                        ind_to_add
+                    ]["image"] = image_value
+                else:
+                    if first_user_idx is None:
+                        first_user_idx = 0
+                    processed_example["messages"][first_user_idx]["content"].append(
+                        {
+                            "type": "image",
+                            "image": image_value,
+                        }
+                    )
+
+            processed_examples.append(remove_none_values(processed_example))
+
+        return processed_examples
+
+    def _mask_non_assistant_keep(self, input_ids: Tensor) -> Tensor:
+        """Return a [B, L] bool keep tensor (True = contributes to loss)."""
+        if self.train_on_inputs:
+            return torch.ones_like(input_ids, dtype=torch.bool)
+
+        # Legacy no-op for boundary-less strategies; warn once so the miss shows up in logs.
+        if not self.role_boundaries:
+            key = type(self).__name__
+            if key not in _ROLE_MASK_WARNED:
+                _ROLE_MASK_WARNED.add(key)
+                LOG.warning(
+                    "%s has no built-in role boundaries; "
+                    "cfg.train_on_inputs / cfg.roles_to_train / cfg.train_on_eos "
+                    "will NOT restrict loss to assistant tokens for this "
+                    "multimodal model — only pad and media tokens are masked, "
+                    "every other token (system, user, assistant) contributes "
+                    "to loss. To enable assistant-only masking, declare "
+                    "per-role markers in YAML via cfg.role_boundaries — see "
+                    "docs/multimodal_assistant_mask.md for the format and the "
+                    "list of strategies on this fallback path.",
+                    key,
+                )
+            return torch.ones_like(input_ids, dtype=torch.bool)
+
+        # Vectorized path regresses 9-13x under multi-threaded torch.
+        scanner = (
+            _compute_role_keep_mask_vectorized
+            if torch.get_num_threads() == 1
+            else _compute_role_keep_mask
+        )
+        return scanner(
+            input_ids,
+            self.role_boundaries,
+            roles_to_train=set(self.roles_to_train),
+            train_on_eos=self.train_on_eos,
+        )
+
+    def _mask_non_assistant(self, labels: Tensor) -> Tensor:
+        """Mask non-trainable role regions to -100 using ``self.role_boundaries``."""
+        keep = self._mask_non_assistant_keep(labels)
+        labels[~keep] = -100
+        return labels
+
+    def process_labels(self, input_ids: Tensor) -> Tensor:
+        keep = self._mask_non_assistant_keep(input_ids)
+        pad_id = getattr(self.processor.tokenizer, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+        if self.image_token_id is not None:
+            keep = keep & (input_ids != self.image_token_id)
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+def _compute_role_keep_mask(
+    labels: Tensor,
+    role_boundaries: list[RoleBoundary],
+    roles_to_train: set[str],
+    train_on_eos: str,
+) -> Tensor:
+    """Return a [B, L] bool keep mask (True = contributes to loss)."""
+    mask = zeros_like(labels)
+    # For "last": remember each trainable turn's end-marker span so we can
+    # unmask only the final one after the scan finishes.
+    last_trainable_end_span: list[Optional[tuple[int, int]]] = [None] * labels.shape[0]
+
+    # Work on a Python list per row — avoids O(n*boundaries) Tensor→list
+    # conversions in the hot prefix-match loop.
+    def _match_prefix(label: list[int], start_pos: int, tok_seq: list[int]) -> bool:
+        if not tok_seq or start_pos + len(tok_seq) > len(label):
+            return False
+        return label[start_pos : start_pos + len(tok_seq)] == tok_seq
+
+    def _find_end(
+        label: list[int], start_pos: int, end_tok: list[int]
+    ) -> tuple[int, bool]:
+        # Empty end_tok means run to end-of-sequence.
+        if not end_tok:
+            return len(label), False
+        k = start_pos
+        while k < len(label):
+            if _match_prefix(label, k, end_tok):
+                return k + len(end_tok), True
+            k += 1
+        return k, False
+
+    for i in range(labels.shape[0]):
+        label = labels[i].tolist()
+        j = 0
+        n = len(label)
+        while j < n:
+            best_match: Optional[RoleBoundary] = None
+            for b in role_boundaries:
+                if _match_prefix(label, j, b.start_tokens):
+                    if best_match is None or len(b.start_tokens) > len(
+                        best_match.start_tokens
+                    ):
+                        best_match = b
+            if best_match is None:
+                j += 1
+                continue
+
+            start_of_content = j + len(best_match.start_tokens)
+            end_after, found_end = _find_end(
+                label, start_of_content, best_match.end_tokens
+            )
+
+            role_in_loss = best_match.role in roles_to_train
+
+            if role_in_loss:
+                if best_match.include_start:
+                    mask[i][j:start_of_content] = 1
+                content_end = (
+                    end_after - len(best_match.end_tokens) if found_end else end_after
+                )
+                mask[i][start_of_content:content_end] = 1
+                # train_on_eos="none"/"last" override include_end during main
+                # loop; "last" is applied after the scan finishes.
+                if (
+                    found_end
+                    and best_match.include_end
+                    and train_on_eos not in ("none", "last")
+                ):
+                    mask[i][content_end:end_after] = 1
+                if found_end and best_match.include_end and train_on_eos == "last":
+                    last_trainable_end_span[i] = (content_end, end_after)
+            else:
+                # Non-trainable role on train_on_eos="all": gate on include_end
+                # so Pixtral / Mistral V7 Tekken shared [/INST] doesn't leak.
+                if found_end and best_match.include_end and train_on_eos == "all":
+                    content_end = end_after - len(best_match.end_tokens)
+                    mask[i][content_end:end_after] = 1
+
+            # When include_end=False, do not consume the end marker: back up so
+            # the next iteration can re-match it as the next boundary's start
+            # marker (Pixtral / Mistral V7 Tekken share [/INST] between
+            # user-end and assistant-start). Requires end_tokens non-empty and
+            # actually found.
+            if found_end and not best_match.include_end and best_match.end_tokens:
+                j = end_after - len(best_match.end_tokens)
+            else:
+                j = end_after
+
+        if train_on_eos == "last" and (span := last_trainable_end_span[i]) is not None:
+            s, e = span
+            mask[i][s:e] = 1
+
+    return mask.bool()
+
+
+def _apply_role_boundaries(
+    labels: Tensor,
+    role_boundaries: list[RoleBoundary],
+    roles_to_train: set[str],
+    train_on_eos: str,
+) -> Tensor:
+    """Mask tokens outside trainable role spans to -100.
+
+    Scan is greedy-left with longest-prefix-wins on start_tokens to disambiguate
+    nested markers (e.g. ``<|im_start|>assistant`` vs ``<|im_start|>``).
+    ``train_on_eos`` accepts ``"turn"`` (end marker in loss on trainable turns
+    only), ``"all"`` (always), ``"none"`` (never — overrides ``include_end``),
+    ``"last"`` (only on the last trainable turn in the sequence).
+    """
+    keep = _compute_role_keep_mask(
+        labels, role_boundaries, roles_to_train, train_on_eos
+    )
+    labels[~keep] = -100
+    return labels
+
+
+def _compute_role_keep_mask_vectorized(
+    labels: Tensor,
+    role_boundaries: list[RoleBoundary],
+    roles_to_train: set[str],
+    train_on_eos: str,
+) -> Tensor:
+    """Vectorized variant of :func:`_compute_role_keep_mask`. Byte-identical output; faster under single-threaded torch."""
+    if labels.numel() == 0:
+        return torch.zeros_like(labels, dtype=torch.bool)
+    if not role_boundaries:
+        # Defensive: match reference all-mask semantics on empty boundaries.
+        return torch.zeros_like(labels, dtype=torch.bool)
+
+    B, L = labels.shape
+    device = labels.device
+
+    # Longer start_tokens first so longest-prefix-wins tie-break holds.
+    indexed = list(enumerate(role_boundaries))
+    indexed.sort(key=lambda ib: -len(ib[1].start_tokens))
+
+    start_winner = torch.full((B, L), -1, dtype=torch.int64, device=device)
+    # Track match length so equal-length ties keep the first writer (matches reference).
+    start_winner_len = torch.zeros((B, L), dtype=torch.int64, device=device)
+    end_match: list[Optional[Tensor]] = [None] * len(role_boundaries)
+
+    for orig_idx, b in indexed:
+        s_tok = b.start_tokens
+        s_len = len(s_tok)
+        if s_len == 0:
+            continue  # defensive: empty start never matches
+
+        if s_len > L:
+            start_mask = torch.zeros((B, 0), dtype=torch.bool, device=device)
+        else:
+            s_tok_t = torch.tensor(s_tok, dtype=labels.dtype, device=device)
+            windows = labels.unfold(1, s_len, 1)
+            start_mask = (windows == s_tok_t).all(dim=-1)
+
+        # Pad to L so absolute position j indexes into start_mask.
+        pad_w = L - start_mask.shape[1]
+        if pad_w > 0:
+            start_mask = torch.cat(
+                [
+                    start_mask,
+                    torch.zeros((B, pad_w), dtype=torch.bool, device=device),
+                ],
+                dim=1,
+            )
+
+        # Strictly-longer gate + longest-first iteration => first writer sticks.
+        update = start_mask & (start_winner_len < s_len)
+        start_winner = torch.where(
+            update, torch.full_like(start_winner, orig_idx), start_winner
+        )
+        start_winner_len = torch.where(
+            update, torch.full_like(start_winner_len, s_len), start_winner_len
+        )
+
+        e_tok = b.end_tokens
+        e_len = len(e_tok)
+        if e_len == 0:
+            end_match[orig_idx] = None
+        elif e_len > L:
+            end_match[orig_idx] = torch.zeros((B, 0), dtype=torch.bool, device=device)
+        else:
+            e_tok_t = torch.tensor(e_tok, dtype=labels.dtype, device=device)
+            ewindows = labels.unfold(1, e_len, 1)
+            end_match[orig_idx] = (ewindows == e_tok_t).all(dim=-1)
+
+    # Lift to CPU lists: hot per-row loop is Python, tensor access is slow.
+    start_winner_cpu = start_winner.cpu().tolist()
+    # Per boundary, per row: sorted end-match start positions, so the inner loop
+    # bisects for the next end instead of scanning every position in the span.
+    end_pos: list[Optional[list[list[int]]]] = []
+    for em in end_match:
+        if em is None:
+            end_pos.append(None)
+        else:
+            nz = em.nonzero(as_tuple=False)
+            rows: list[list[int]] = [[] for _ in range(B)]
+            for r, c in nz.tolist():
+                rows[r].append(c)
+            end_pos.append(rows)
+
+    bnd_start_len = [len(b.start_tokens) for b in role_boundaries]
+    bnd_end_len = [len(b.end_tokens) for b in role_boundaries]
+    bnd_include_start = [b.include_start for b in role_boundaries]
+    bnd_include_end = [b.include_end for b in role_boundaries]
+    bnd_role_in_loss = [b.role in roles_to_train for b in role_boundaries]
+
+    mask = zeros_like(labels)
+    ONE = b"\x01"
+
+    for i in range(B):
+        row_start = start_winner_cpu[i]
+        n = L
+        last_trainable_end_span: Optional[tuple[int, int]] = None
+        # bytearray slice-assignment beats per-element writes; one tensor per row.
+        row_mask = bytearray(n)
+
+        j = 0
+        while j < n:
+            bidx = row_start[j]
+            if bidx == -1:
+                j += 1
+                continue
+
+            s_len = bnd_start_len[bidx]
+            start_of_content = j + s_len
+            e_len = bnd_end_len[bidx]
+            include_start = bnd_include_start[bidx]
+            include_end = bnd_include_end[bidx]
+            role_in_loss = bnd_role_in_loss[bidx]
+
+            if e_len == 0:
+                end_after = n
+                found_end = False
+            else:
+                positions = end_pos[bidx][i]  # type: ignore[index]
+                limit = n - e_len
+                idx = bisect.bisect_left(positions, start_of_content)
+                if idx < len(positions) and positions[idx] <= limit:
+                    found_end = True
+                    end_after = positions[idx] + e_len
+                else:
+                    found_end = False
+                    end_after = n
+
+            if role_in_loss:
+                if include_start:
+                    row_mask[j:start_of_content] = ONE * (start_of_content - j)
+                content_end = end_after - e_len if found_end else end_after
+                row_mask[start_of_content:content_end] = ONE * (
+                    content_end - start_of_content
+                )
+                if found_end and include_end and train_on_eos not in ("none", "last"):
+                    row_mask[content_end:end_after] = ONE * (end_after - content_end)
+                if found_end and include_end and train_on_eos == "last":
+                    last_trainable_end_span = (content_end, end_after)
+            else:
+                if found_end and include_end and train_on_eos == "all":
+                    content_end = end_after - e_len
+                    row_mask[content_end:end_after] = ONE * (end_after - content_end)
+
+            # include_end=False rewind: re-match the end as the next start.
+            if found_end and not include_end and e_len:
+                j = end_after - e_len
+            else:
+                j = end_after
+
+        if train_on_eos == "last" and last_trainable_end_span is not None:
+            s, e = last_trainable_end_span
+            row_mask[s:e] = ONE * (e - s)
+
+        # One tensor write per row — keep heavyweight op out of inner loop.
+        if any(row_mask):
+            mask[i] = torch.frombuffer(row_mask, dtype=torch.uint8).to(mask.dtype)
+
+    return mask.bool()
+
+
+def _apply_role_boundaries_vectorized(
+    labels: Tensor,
+    role_boundaries: list[RoleBoundary],
+    roles_to_train: set[str],
+    train_on_eos: str,
+) -> Tensor:
+    """Vectorized variant of :func:`_apply_role_boundaries`."""
+    keep = _compute_role_keep_mask_vectorized(
+        labels, role_boundaries, roles_to_train, train_on_eos
+    )
+    labels[~keep] = -100
+    return labels
+
+
+def _encode_markers(tokenizer, marker_strs: list[str]) -> list[list[int]]:
+    """Encode markers via ``encode(..., add_special_tokens=False)``; drops empty results."""
+    result = []
+    for s in marker_strs:
+        toks = tokenizer.encode(s, add_special_tokens=False)
+        if toks:
+            result.append(toks)
+    return result
+
+
+def _resolve_role_boundary_override(specs: list[dict], tokenizer) -> list[RoleBoundary]:
+    """Resolve user ``cfg.role_boundaries`` specs into RoleBoundary objects.
+
+    The sentinel ``end == "eos_token"`` resolves to ``eos_token_id`` (used by
+    Pixtral/Mistral v7 templates). ``end`` null/omitted runs to end-of-sequence.
+    """
+    out: list[RoleBoundary] = []
+    for i, spec in enumerate(specs):
+        if hasattr(spec, "model_dump"):
+            d = spec.model_dump()
+        else:
+            d = dict(spec)
+
+        role = d.get("role")
+        start_str = d.get("start")
+        if not role or start_str is None:
+            raise ValueError(
+                f"cfg.role_boundaries[{i}] must have both 'role' and 'start' "
+                f"(got {d!r})."
+            )
+        start_ids = tokenizer.encode(start_str, add_special_tokens=False)
+        if not start_ids:
+            raise ValueError(
+                f"cfg.role_boundaries[{i}]: start marker {start_str!r} "
+                f"tokenizes to an empty sequence; cannot match."
+            )
+
+        end_spec = d.get("end")
+        if end_spec is None:
+            end_ids: list[int] = []
+        elif end_spec == "eos_token":
+            eos = getattr(tokenizer, "eos_token_id", None)
+            if eos is None:
+                raise ValueError(
+                    f"cfg.role_boundaries[{i}] requested end='eos_token' but "
+                    "the tokenizer has no eos_token_id."
+                )
+            end_ids = [eos]
+        else:
+            end_ids = tokenizer.encode(end_spec, add_special_tokens=False)
+            if not end_ids:
+                raise ValueError(
+                    f"cfg.role_boundaries[{i}]: end marker {end_spec!r} "
+                    f"tokenizes to an empty sequence; cannot match. Use "
+                    f"end=null to run to end-of-sequence or end='eos_token' "
+                    f"to terminate at the tokenizer's EOS."
+                )
+
+        out.append(
+            RoleBoundary(
+                role=role,
+                start_tokens=start_ids,
+                end_tokens=end_ids,
+                include_start=bool(d.get("include_start", False)),
+                include_end=bool(d.get("include_end", True)),
+            )
+        )
+    return out
+
+
+class Qwen2VLProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Qwen2-VL (ChatML ``<|im_start|>{role}\\n ... <|im_end|>``)."""
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        super().__init__(
+            processor,
+            chat_template,
+            image_size,
+            image_resize_algorithm,
+            train_on_inputs=train_on_inputs,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            role_boundaries_override=role_boundaries_override,
+            field_messages=field_messages,
+        )
+        self.image_token = "<|image_pad|>"  # nosec
+        self.image_token_id = processor.tokenizer.convert_tokens_to_ids(
+            self.image_token
+        )
+
+    def _build_role_boundaries(self) -> list[RoleBoundary]:
+        tok = self.processor.tokenizer
+        end = _encode_markers(tok, ["<|im_end|>"])
+        if not end:
+            return []
+        end_ids = end[0]
+        boundaries = []
+        for role in ("system", "user", "assistant"):
+            start = _encode_markers(tok, [f"<|im_start|>{role}\n"])
+            if start:
+                boundaries.append(
+                    RoleBoundary(role=role, start_tokens=start[0], end_tokens=end_ids)
+                )
+        return boundaries
+
+
+class Qwen3_5ProcessingStrategy(Qwen2VLProcessingStrategy):
+    """Processing Strategy class for Qwen3.5 (Qwen2-VL boundaries + ``<|video_pad|>`` mask)."""
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        super().__init__(
+            processor,
+            chat_template,
+            image_size,
+            image_resize_algorithm,
+            train_on_inputs=train_on_inputs,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            role_boundaries_override=role_boundaries_override,
+            field_messages=field_messages,
+        )
+        self.video_token = "<|video_pad|>"  # nosec
+        self.video_token_id = processor.tokenizer.convert_tokens_to_ids(
+            self.video_token
+        )
+
+    def process_labels(self, input_ids):
+        keep = self._mask_non_assistant_keep(input_ids)
+        pad_id = getattr(self.processor.tokenizer, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+        if self.image_token_id is not None:
+            keep = keep & (input_ids != self.image_token_id)
+        if self.video_token_id is not None:
+            keep = keep & (input_ids != self.video_token_id)
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+class _GemmaTurnStrategy(ProcessingStrategy):
+    """Gemma3/3n ``<start_of_turn>{role} ... <end_of_turn>`` (Gemma 4 uses different markers)."""
+
+    def _build_role_boundaries(self) -> list[RoleBoundary]:
+        tok = self.processor.tokenizer
+        end = _encode_markers(tok, ["<end_of_turn>"])
+        if not end:
+            return []
+        end_ids = end[0]
+        boundaries = []
+        # Template uses 'model'; external role knob stays 'assistant'. Gemma 3
+        # and Gemma 3n jinja templates fold the system message into the first
+        # user's content prefix and never emit '<start_of_turn>system', so we
+        # don't declare a system boundary here.
+        role_marker_pairs = [
+            ("assistant", "model"),
+            ("user", "user"),
+        ]
+        for external_role, template_role in role_marker_pairs:
+            start = _encode_markers(tok, [f"<start_of_turn>{template_role}\n"])
+            if start:
+                boundaries.append(
+                    RoleBoundary(
+                        role=external_role,
+                        start_tokens=start[0],
+                        end_tokens=end_ids,
+                    )
+                )
+        return boundaries
+
+
+class Gemma3ProcessingStrategy(_GemmaTurnStrategy):
+    """Processing Strategy class for Gemma3."""
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        super().__init__(
+            processor,
+            chat_template,
+            image_size,
+            image_resize_algorithm,
+            train_on_inputs=train_on_inputs,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            role_boundaries_override=role_boundaries_override,
+            field_messages=field_messages,
+        )
+        # Real Gemma3 tokenizers expose boi_token as a direct attribute, not
+        # via special_tokens_map (which only holds HF's standard slots).
+        boi = getattr(processor.tokenizer, "boi_token", None)
+        if boi is not None:
+            self.image_token = boi
+            self.image_token_id = processor.tokenizer.convert_tokens_to_ids(boi)
+
+    def process_labels(self, input_ids):
+        keep = self._mask_non_assistant_keep(input_ids)
+        pad_id = getattr(self.processor.tokenizer, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+        if self.image_token_id is not None:
+            keep = keep & (input_ids != self.image_token_id)
+        # Resolve <image_soft_token> via tokenizer; fall back to default id
+        # if not in vocab. Matches Gemma4's pattern.
+        tok = self.processor.tokenizer
+        soft_id = tok.convert_tokens_to_ids("<image_soft_token>")
+        unk_id = getattr(tok, "unk_token_id", None)
+        if soft_id is not None and soft_id != unk_id:
+            keep = keep & (input_ids != soft_id)
+        else:
+            keep = keep & (input_ids != 262144)
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+class Gemma3nProcessingStrategy(_GemmaTurnStrategy):
+    """Gemma3n: same turn boundaries as Gemma3, additionally masks audio/delimiter tokens."""
+
+    def process_labels(self, input_ids):
+        keep = self._mask_non_assistant_keep(input_ids)
+        tok = self.processor.tokenizer
+        pad_id = getattr(tok, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+        if self.image_token_id is not None:
+            keep = keep & (input_ids != self.image_token_id)
+        # Follows huggingface-gemma-recipes fine_tune_gemma3n_on_t4 notebook.
+        for attr in (
+            "image_token_id",
+            "audio_token_id",
+            "boi_token_id",
+            "eoi_token_id",
+        ):
+            tok_id = getattr(tok, attr, None)
+            if tok_id is not None:
+                keep = keep & (input_ids != tok_id)
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+class Gemma4ProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Gemma 4.
+
+    Boundary markers ``<|turn>model ... <turn|>`` verified against
+    google/gemma-4-E2B-it. boi/eoi/boa/eoa ids are resolved via
+    ``convert_tokens_to_ids`` since only their string forms are on the processor.
+    """
+
+    def _build_role_boundaries(self) -> list[RoleBoundary]:
+        tok = self.processor.tokenizer
+        end = _encode_markers(tok, ["<turn|>"])
+        if not end:
+            return []
+        end_ids = end[0]
+        boundaries = []
+        role_marker_pairs = [
+            ("assistant", "model"),
+            ("user", "user"),
+            ("system", "system"),
+        ]
+        for external_role, template_role in role_marker_pairs:
+            # Include trailing ``\n`` for consistency with Qwen/Gemma3/Llama
+            # markers; the newline is part of the marker in the real
+            # google/gemma-4 tokenizer's chat template.
+            start = _encode_markers(tok, [f"<|turn>{template_role}\n"])
+            if start:
+                boundaries.append(
+                    RoleBoundary(
+                        role=external_role,
+                        start_tokens=start[0],
+                        end_tokens=end_ids,
+                    )
+                )
+        return boundaries
+
+    def process_labels(self, input_ids):
+        keep = self._mask_non_assistant_keep(input_ids)
+
+        tokenizer = self.processor.tokenizer
+        unk_id = getattr(tokenizer, "unk_token_id", None)
+
+        pad_id = getattr(tokenizer, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+        if self.image_token_id is not None:
+            keep = keep & (input_ids != self.image_token_id)
+
+        if getattr(tokenizer, "image_token_id", None) is not None:
+            keep = keep & (input_ids != tokenizer.image_token_id)
+        if getattr(tokenizer, "audio_token_id", None) is not None:
+            keep = keep & (input_ids != tokenizer.audio_token_id)
+
+        # boi/eoi/boa/eoa are only string attrs on the processor; resolve ids here.
+        for attr in ("boi_token", "eoi_token", "boa_token", "eoa_token"):
+            token_str = getattr(self.processor, attr, None)
+            if token_str is None:
+                continue
+            token_id = tokenizer.convert_tokens_to_ids(token_str)
+            if token_id is None or token_id == unk_id:
+                continue
+            keep = keep & (input_ids != token_id)
+
+        # Video id lives on the processor, not the tokenizer.
+        video_token_id = getattr(self.processor, "video_token_id", None)
+        if video_token_id is not None and video_token_id != unk_id:
+            keep = keep & (input_ids != video_token_id)
+
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+class Gemma4UnifiedProcessingStrategy(Gemma4ProcessingStrategy):
+    """Processing Strategy for Gemma 4 Unified (encoder-free image/audio/video).
+
+    The unified checkpoint shares Gemma 4's turn format and the same media
+    placeholder/delimiter token set (image/audio/video, boi/eoi/boa/eoa), so
+    boundary detection and label masking are inherited unchanged — both resolve
+    ids dynamically from the processor/tokenizer rather than hard-coding them.
+    The encoder-free raw pixel/waveform projection is handled entirely by the HF
+    processor, so the strategy itself needs no audio/vision-specific logic.
+    """
+
+
+class Llama3_2VisionProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Llama-3.2 Vision (``<|start_header_id|>{role}<|end_header_id|>\\n\\n ... <|eot_id|>``)."""
+
+    def _build_role_boundaries(self) -> list[RoleBoundary]:
+        tok = self.processor.tokenizer
+        end = _encode_markers(tok, ["<|eot_id|>"])
+        if not end:
+            return []
+        end_ids = end[0]
+        boundaries = []
+        for role in ("system", "user", "assistant", "ipython", "tool"):
+            start = _encode_markers(
+                tok, [f"<|start_header_id|>{role}<|end_header_id|>\n\n"]
+            )
+            if start:
+                boundaries.append(
+                    RoleBoundary(role=role, start_tokens=start[0], end_tokens=end_ids)
+                )
+        return boundaries
+
+
+class Llama4ProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Llama 4 (``<|header_start|>{role}<|header_end|>\\n\\n ... <|eot|>``)."""
+
+    def _build_role_boundaries(self) -> list[RoleBoundary]:
+        tok = self.processor.tokenizer
+        end = _encode_markers(tok, ["<|eot|>"])
+        if not end:
+            return []
+        end_ids = end[0]
+        boundaries = []
+        for role in ("system", "user", "assistant", "ipython", "tool"):
+            start = _encode_markers(tok, [f"<|header_start|>{role}<|header_end|>\n\n"])
+            if start:
+                boundaries.append(
+                    RoleBoundary(role=role, start_tokens=start[0], end_tokens=end_ids)
+                )
+        return boundaries
+
+
+class PixtralProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Pixtral (``[INST] ... [/INST]`` user, assistant terminates at ``eos_token``).
+
+    ``[/INST]`` is shared between user-end and assistant-start. We declare user
+    with ``include_end=False`` so the scanner hands the ``[/INST]`` back to
+    assistant's start match on the next iteration.
+    """
+
+    def _build_role_boundaries(self) -> list[RoleBoundary]:
+        tok = self.processor.tokenizer
+        eos = getattr(tok, "eos_token_id", None)
+        if eos is None:
+            return []
+        boundaries = []
+        inst_start = _encode_markers(tok, ["[INST]"])
+        inst_end = _encode_markers(tok, ["[/INST]"])
+        if inst_start and inst_end:
+            boundaries.append(
+                RoleBoundary(
+                    role="user",
+                    start_tokens=inst_start[0],
+                    end_tokens=inst_end[0],
+                    include_end=False,
+                )
+            )
+            boundaries.append(
+                RoleBoundary(
+                    role="assistant",
+                    start_tokens=inst_end[0],
+                    end_tokens=[eos],
+                )
+            )
+        return boundaries
+
+
+class MistralV7TekkenProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Mistral v7 Tekken (Pixtral-style plus ``[SYSTEM_PROMPT]...[/SYSTEM_PROMPT]``).
+
+    Same ``[/INST]``-shared-marker treatment as :class:`PixtralProcessingStrategy`.
+    """
+
+    def _build_role_boundaries(self) -> list[RoleBoundary]:
+        tok = self.processor.tokenizer
+        eos = getattr(tok, "eos_token_id", None)
+        if eos is None:
+            return []
+        boundaries = []
+        sys_start = _encode_markers(tok, ["[SYSTEM_PROMPT]"])
+        sys_end = _encode_markers(tok, ["[/SYSTEM_PROMPT]"])
+        if sys_start and sys_end:
+            boundaries.append(
+                RoleBoundary(
+                    role="system", start_tokens=sys_start[0], end_tokens=sys_end[0]
+                )
+            )
+        inst_start = _encode_markers(tok, ["[INST]"])
+        inst_end = _encode_markers(tok, ["[/INST]"])
+        if inst_start and inst_end:
+            boundaries.append(
+                RoleBoundary(
+                    role="user",
+                    start_tokens=inst_start[0],
+                    end_tokens=inst_end[0],
+                    include_end=False,
+                )
+            )
+            boundaries.append(
+                RoleBoundary(
+                    role="assistant",
+                    start_tokens=inst_end[0],
+                    end_tokens=[eos],
+                )
+            )
+        return boundaries
+
+
+class VoxtralProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Voxtral.
+
+    Role boundaries NOT declared — mistral-common instruct tokenizer markers
+    unverified. Falls back to pad+audio masking with a one-shot warning.
+    """
+
+    def __init__(
+        self,
+        processor: VoxtralProcessor,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        super().__init__(
+            processor,
+            chat_template,
+            image_size,
+            image_resize_algorithm,
+            train_on_inputs=train_on_inputs,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            role_boundaries_override=role_boundaries_override,
+            field_messages=field_messages,
+        )
+        special_ids = (
+            processor.tokenizer.tokenizer.instruct_tokenizer.audio_encoder.special_ids
+        )
+
+        self.audio_token = special_ids.audio
+        self.begin_audio_token = special_ids.begin_audio
+
+    def process_labels(self, input_ids):
+        keep = self._mask_non_assistant_keep(input_ids)
+
+        pad_id = getattr(self.processor.tokenizer, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+        if self.audio_token is not None:
+            keep = keep & (input_ids != self.audio_token)
+        if self.begin_audio_token is not None:
+            keep = keep & (input_ids != self.begin_audio_token)
+
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+class SmolVLM2ProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for SmolVLM2.
+
+    Role boundaries NOT declared — SmolVLM2 chat_template varies per checkpoint
+    (HuggingFaceTB ships multiple variants), so we opt out rather than mis-mask.
+    """
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        super().__init__(
+            processor,
+            chat_template,
+            image_size,
+            image_resize_algorithm,
+            train_on_inputs=train_on_inputs,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            role_boundaries_override=role_boundaries_override,
+            field_messages=field_messages,
+        )
+        self.image_token = "<image>"  # nosec
+
+        self.image_token_id = processor.tokenizer.additional_special_tokens_ids[
+            processor.tokenizer.additional_special_tokens.index(self.image_token)
+        ]
+
+
+class Mistral3ProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Mistral3.
+
+    Role boundaries NOT declared (mistral-common instruct tokenizer unverified);
+    same fallback as VoxtralProcessingStrategy.
+    """
+
+    def __init__(
+        self,
+        processor,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        super().__init__(
+            processor,
+            chat_template,
+            image_size,
+            image_resize_algorithm,
+            train_on_inputs=train_on_inputs,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            role_boundaries_override=role_boundaries_override,
+            field_messages=field_messages,
+        )
+        special_ids = (
+            processor.tokenizer.tokenizer.instruct_tokenizer.image_encoder.special_ids
+        )
+
+        self.image_token = special_ids.img
+        self.image_break_token = special_ids.img_break
+        self.image_end_token = special_ids.img_end
+
+    def process_labels(self, input_ids):
+        keep = self._mask_non_assistant_keep(input_ids)
+
+        pad_id = getattr(self.processor.tokenizer, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+        for tok_id in (self.image_token, self.image_break_token, self.image_end_token):
+            if tok_id is not None:
+                keep = keep & (input_ids != tok_id)
+
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+class InternVLProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for InternVL.
+
+    Role boundaries NOT declared (InternLM-style template unverified); falls
+    back to pad + image-id masking with a one-shot warning.
+    """
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        super().__init__(
+            processor,
+            chat_template,
+            image_size,
+            image_resize_algorithm,
+            train_on_inputs=train_on_inputs,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            role_boundaries_override=role_boundaries_override,
+            field_messages=field_messages,
+        )
+
+        if not hasattr(processor, "image_ids"):
+            raise ValueError("'image_ids' missing from InternVL Processor.")
+
+        self.image_token_ids = processor.image_ids
+
+    def process_labels(self, input_ids):
+        keep = self._mask_non_assistant_keep(input_ids)
+
+        pad_id = getattr(self.processor.tokenizer, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+
+        for ids in self.image_token_ids:
+            if ids is not None:
+                keep = keep & (input_ids != ids)
+
+        # Video tokens get converted to image patches during media processing; masking may be redundant.
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+class Glm4vProcessingStrategy(ProcessingStrategy):
+    """Shared strategy for Glm4vProcessor (GLM-4V / GLM-4.1V) and
+    Glm46VProcessor (GLM-4.6V / GLM-4.7V) — identical media-token markers.
+
+    Role boundaries unverified; use cfg.role_boundaries to enable masking.
+    """
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+        train_on_inputs: bool = False,
+        roles_to_train: Optional[list[str]] = None,
+        train_on_eos: Optional[str] = None,
+        role_boundaries_override: Optional[list[dict]] = None,
+        field_messages: str | list[str] | tuple[str, ...] | None = None,
+    ):
+        super().__init__(
+            processor,
+            chat_template,
+            image_size,
+            image_resize_algorithm,
+            train_on_inputs=train_on_inputs,
+            roles_to_train=roles_to_train,
+            train_on_eos=train_on_eos,
+            role_boundaries_override=role_boundaries_override,
+            field_messages=field_messages,
+        )
+
+        self.tokenizer = getattr(processor, "tokenizer", processor)
+
+        self.image_token = "<|image|>"  # nosec
+        self.begin_image_token = "<|begin_of_image|>"  # nosec
+        self.end_image_token = "<|end_of_image|>"  # nosec
+        self.video_token = "<|video|>"  # nosec
+        self.begin_video_token = "<|begin_of_video|>"  # nosec
+        self.end_video_token = "<|end_of_video|>"  # nosec
+
+        self.image_token_id = self.tokenizer.convert_tokens_to_ids(self.image_token)
+        self.begin_image_token_id = self.tokenizer.convert_tokens_to_ids(
+            self.begin_image_token
+        )
+        self.end_image_token_id = self.tokenizer.convert_tokens_to_ids(
+            self.end_image_token
+        )
+        self.video_token_id = self.tokenizer.convert_tokens_to_ids(self.video_token)
+        self.begin_video_token_id = self.tokenizer.convert_tokens_to_ids(
+            self.begin_video_token
+        )
+        self.end_video_token_id = self.tokenizer.convert_tokens_to_ids(
+            self.end_video_token
+        )
+
+    def process_labels(self, input_ids):
+        keep = self._mask_non_assistant_keep(input_ids)
+
+        pad_id = getattr(self.tokenizer, "pad_token_id", None)
+        if pad_id is not None:
+            keep = keep & (input_ids != pad_id)
+
+        for tok_id in (
+            self.image_token_id,
+            self.begin_image_token_id,
+            self.end_image_token_id,
+            self.video_token_id,
+            self.begin_video_token_id,
+            self.end_video_token_id,
+        ):
+            if tok_id is not None:
+                keep = keep & (input_ids != tok_id)
+
+        labels = input_ids.clone()
+        labels[~keep] = -100
+        return labels
+
+
+def get_processing_strategy(
+    processor: ProcessorMixin,
+    chat_template,
+    chat_template_type,
+    image_size: int | tuple[int, int] | None = None,
+    image_resize_algorithm: Resampling | None = None,
+    train_on_inputs: bool = False,
+    roles_to_train: Optional[list[str]] = None,
+    train_on_eos: Optional[str] = None,
+    role_boundaries_override: Optional[list[dict]] = None,
+    field_messages: str | list[str] | tuple[str, ...] | None = None,
+):
+    processing_kwargs = {
+        "processor": processor,
+        "chat_template": chat_template,
+        "image_size": image_size,
+        "image_resize_algorithm": image_resize_algorithm,
+        "train_on_inputs": train_on_inputs,
+        "roles_to_train": roles_to_train,
+        "train_on_eos": train_on_eos,
+        "role_boundaries_override": role_boundaries_override,
+        "field_messages": field_messages,
+    }
+
+    if chat_template_type in [None, "tokenizer_default"]:
+        tokenizer = getattr(processor, "tokenizer", processor)
+        if hasattr(tokenizer, "chat_template"):
+            processing_kwargs["chat_template"] = tokenizer.chat_template
+
+    if chat_template_type == "qwen2_vl":
+        return Qwen2VLProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "qwen3_5":
+        return Qwen3_5ProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "gemma3":
+        return Gemma3ProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "gemma3n":
+        return Gemma3nProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "gemma4":
+        return Gemma4ProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "gemma4_unified":
+        return Gemma4UnifiedProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "llama3_2_vision":
+        return Llama3_2VisionProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "llama4":
+        return Llama4ProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "pixtral":
+        return PixtralProcessingStrategy(**processing_kwargs)
+    if chat_template_type == "mistral_v7_tekken":
+        return MistralV7TekkenProcessingStrategy(**processing_kwargs)
+
+    if isinstance(processor, VoxtralProcessor):
+        return VoxtralProcessingStrategy(**processing_kwargs)
+
+    if isinstance(processor, SmolVLMProcessor):
+        return SmolVLM2ProcessingStrategy(**processing_kwargs)
+
+    # Lazy import: mistral_common is optional. Mirrors the Glm46V pattern below.
+    try:
+        from axolotl.utils.mistral.mistral3_processor import Mistral3Processor
+
+        if isinstance(processor, Mistral3Processor):
+            return Mistral3ProcessingStrategy(**processing_kwargs)
+    except (ImportError, ModuleNotFoundError) as exc:
+        LOG.debug(
+            "Mistral3Processor import failed; Mistral3 strategy will be unavailable: %r",
+            exc,
+        )
+
+    # Both Glm4vProcessor and Glm46VProcessor share markers; route to the same
+    # strategy. Independent try/except so either can be absent.
+    try:
+        from transformers.models.glm4v.processing_glm4v import Glm4vProcessor
+
+        if isinstance(processor, Glm4vProcessor):
+            return Glm4vProcessingStrategy(**processing_kwargs)
+    except (ImportError, ModuleNotFoundError) as exc:
+        LOG.debug("Glm4vProcessor import failed: %r", exc)
+
+    try:
+        from transformers.models.glm46v.processing_glm46v import Glm46VProcessor
+
+        if isinstance(processor, Glm46VProcessor):
+            return Glm4vProcessingStrategy(**processing_kwargs)
+    except (ImportError, ModuleNotFoundError) as exc:
+        LOG.debug("Glm46VProcessor import failed: %r", exc)
+
+    if isinstance(processor, InternVLProcessor):
+        return InternVLProcessingStrategy(**processing_kwargs)
+
+    # Unregistered templates (llava, lfm2vl, mistral_v3_tekken, ...) use the
+    # base strategy; it warns once when train_on_inputs=False.
+    return ProcessingStrategy(**processing_kwargs)
diff --git a/src/axolotl/prompt_strategies/__init__.py b/src/axolotl/prompt_strategies/__init__.py
index e62a5c20ce..d9936b9aef 100644
--- a/src/axolotl/prompt_strategies/__init__.py
+++ b/src/axolotl/prompt_strategies/__init__.py
@@ -4,15 +4,36 @@
 import inspect
 
 from axolotl.prompt_strategies.user_defined import UserDefinedDatasetConfig
+from axolotl.utils.logging import get_logger
 
+LOG = get_logger(__name__)
 
-def load(strategy, tokenizer, cfg, ds_cfg):
+
+def load(strategy, tokenizer, cfg, ds_cfg, processor=None):
     try:
+        if strategy == "messages":
+            from .messages import load as messages_load
+
+            return messages_load(tokenizer, cfg, ds_cfg, processor=processor)
         load_fn = "load"
-        if strategy.split(".")[-1].startswith("load_"):
+        package = "axolotl.prompt_strategies"
+        if (
+            strategy.split(".")[-1].startswith("load_")
+            or strategy.split(".")[-1] == "load"
+        ):
             load_fn = strategy.split(".")[-1]
             strategy = ".".join(strategy.split(".")[:-1])
-        mod = importlib.import_module(f".{strategy}", "axolotl.prompt_strategies")
+        elif len(strategy.split(".")) > 1:
+            try:
+                importlib.import_module(
+                    "." + strategy.split(".")[-1],
+                    ".".join(strategy.split(".")[:-1]),
+                )
+                package = ".".join(strategy.split(".")[:-1])
+                strategy = strategy.split(".")[-1]
+            except ModuleNotFoundError:
+                pass
+        mod = importlib.import_module(f".{strategy}", package)
         func = getattr(mod, load_fn)
         load_kwargs = {}
         if strategy == "user_defined":
@@ -21,6 +42,12 @@ def load(strategy, tokenizer, cfg, ds_cfg):
             sig = inspect.signature(func)
             if "ds_cfg" in sig.parameters:
                 load_kwargs["ds_cfg"] = ds_cfg
+            if "processor" in sig.parameters:
+                load_kwargs["processor"] = processor
+
         return func(tokenizer, cfg, **load_kwargs)
-    except Exception:  # pylint: disable=broad-exception-caught
+    except ModuleNotFoundError:
         return None
+    except Exception as exc:
+        LOG.error(f"Failed to load prompt strategy `{strategy}`: {str(exc)}")
+        raise exc
diff --git a/src/axolotl/prompt_strategies/_synthetic.py b/src/axolotl/prompt_strategies/_synthetic.py
new file mode 100644
index 0000000000..1353f09ed6
--- /dev/null
+++ b/src/axolotl/prompt_strategies/_synthetic.py
@@ -0,0 +1,96 @@
+"""
+Synthetic dataset generator for benchmarking and testing.
+
+Generates datasets with configurable sequence length, dataset size, and token ID ranges.
+Useful for benchmarking memory usage and speed by sequence length, and for validating
+weighted dataset mixes.
+
+YAML configuration example:
+
+    datasets:
+      - path: synthetic
+        type: _synthetic
+        length: 1000
+        sequence_length: 2048
+        min_input_id: 100
+        max_input_id: 32000
+        seed: 42
+"""
+
+from typing import Any, Dict, Optional
+
+import numpy as np
+from datasets import Dataset
+
+from axolotl.prompt_tokenizers import DatasetWrappingStrategy
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class SyntheticDatasetStrategy(DatasetWrappingStrategy):
+    """Strategy that generates synthetic tokenized data, ignoring the source dataset."""
+
+    def __init__(
+        self,
+        sequence_length: int = 2048,
+        length: int = 1000,
+        min_input_id: int = 100,
+        max_input_id: int = 32000,
+        seed: Optional[int] = None,
+    ):
+        self.sequence_length = sequence_length
+        self.length = length
+        self.min_input_id = min_input_id
+        self.max_input_id = max_input_id
+        self.seed = seed
+
+    def wrap_dataset(
+        self,
+        dataset,
+        process_count: int | None = None,
+        keep_in_memory: bool | None = False,
+        **kwargs,
+    ) -> Dataset:
+        LOG.info(
+            f"Generating synthetic dataset: {self.length} samples, "
+            f"sequence_length={self.sequence_length}, "
+            f"input_id_range=[{self.min_input_id}, {self.max_input_id})"
+        )
+
+        rng = np.random.default_rng(self.seed)
+        input_ids = rng.integers(
+            low=self.min_input_id,
+            high=self.max_input_id,
+            size=(self.length, self.sequence_length),
+        ).tolist()
+
+        attention_mask = [[1] * self.sequence_length] * self.length
+        # labels == input_ids means we train on all tokens
+        labels = [row[:] for row in input_ids]
+
+        return Dataset.from_dict(
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "labels": labels,
+            }
+        )
+
+
+def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
+    ds_cfg = ds_cfg or {}
+
+    sequence_length = ds_cfg.get("sequence_length", cfg.sequence_len)
+    length = ds_cfg.get("length", 1000)
+    min_input_id = ds_cfg.get("min_input_id", 100)
+    max_input_id = ds_cfg.get("max_input_id", tokenizer.vocab_size)
+    seed = ds_cfg.get("seed", None)
+
+    return SyntheticDatasetStrategy(
+        sequence_length=sequence_length,
+        length=length,
+        min_input_id=min_input_id,
+        max_input_id=max_input_id,
+        seed=seed,
+    )
diff --git a/src/axolotl/prompt_strategies/alpaca_chat.py b/src/axolotl/prompt_strategies/alpaca_chat.py
index 975fee889e..391ba6072b 100644
--- a/src/axolotl/prompt_strategies/alpaca_chat.py
+++ b/src/axolotl/prompt_strategies/alpaca_chat.py
@@ -39,7 +39,7 @@ class AlpacaChatPrompter(AlpacaPrompter):
     system_prompt = "Below is an instruction from a USER that describes a task, paired with an input that provides further context. The ASSISTANT writes a response that concisely and appropriately completes the request.\n\n"
     system_no_input_prompt = "Below is an instruction from a USER that describes a task. The ASSISTANT writes a response that appropriately and concisely completes the request.\n\n"
 
-    def __init__(self):  # pylint: disable=super-init-not-called
+    def __init__(self):
         self.prompt_style = PromptStyle.CHAT.value
         self.match_prompt_style()
 
@@ -54,7 +54,7 @@ class NoSystemPrompter(AlpacaPrompter):
     turn_format = "{instruction} {input} "
     turn_no_input_format = "{instruction} "
 
-    def __init__(self):  # pylint: disable=super-init-not-called
+    def __init__(self):
         pass
 
 
diff --git a/src/axolotl/prompt_strategies/alpaca_w_system.py b/src/axolotl/prompt_strategies/alpaca_w_system.py
index 8c8cc07435..808ba517e3 100644
--- a/src/axolotl/prompt_strategies/alpaca_w_system.py
+++ b/src/axolotl/prompt_strategies/alpaca_w_system.py
@@ -1,6 +1,7 @@
 """
 Prompt strategies loader for alpaca instruction datasets with system prompts
 """
+
 from typing import Generator, Tuple, Union
 
 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
@@ -21,10 +22,9 @@ def parse_instruction_fields(self, prompt) -> Tuple[str, str, str, str]:
         )
 
     def tokenize_prompt(self, prompt):
-        # pylint: disable=duplicate-code
         (
             instruction,
-            input,  # pylint: disable=redefined-builtin
+            input,
             response,
             system,
         ) = self.parse_instruction_fields(prompt)
@@ -63,7 +63,7 @@ def build_prompt_w_system(
         self,
         system: str,
         instruction: str,
-        input: Union[None, str] = None,  # pylint: disable=redefined-builtin
+        input: Union[None, str] = None,
         output: Union[None, str] = None,
     ) -> Generator[str, None, None]:
         # returns the full prompt from instruction and optional input
@@ -92,7 +92,6 @@ class OpenOrcaSystemDataPrompter(SystemDataPrompter):
     """
 
     def match_prompt_style(self):
-        # pylint: disable=duplicate-code
         if self.prompt_style == PromptStyle.INSTRUCT.value:
             self.turn_format = "### Human:\n{instruction}\n### Additional Context:\n{input}\n### Assistant:\n"
             self.turn_no_input_format = "### Human:\n{instruction}\n### Assistant:\n"
diff --git a/src/axolotl/prompt_strategies/base.py b/src/axolotl/prompt_strategies/base.py
index fce2aba14a..94a65aaa1f 100644
--- a/src/axolotl/prompt_strategies/base.py
+++ b/src/axolotl/prompt_strategies/base.py
@@ -3,18 +3,32 @@
 """
 
 import importlib
-import logging
 
-LOG = logging.getLogger("axolotl")
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
 
 
 def load(strategy, cfg, module_base=None, **kwargs):
     try:
+        if len(strategy.split(".")) == 1:
+            strategy = strategy + ".default"
         load_fn = strategy.split(".")[-1]
-        strategy = ".".join(strategy.split(".")[:-1])
-        mod = importlib.import_module(f".{strategy}", module_base)
+        if len(strategy.split(".")) > 1:
+            try:
+                importlib.import_module(
+                    strategy.split(".")[-2],
+                    ".".join(strategy.split(".")[:-2]),
+                )
+                module_base = ".".join(strategy.split(".")[:-2])
+                strategy = strategy.split(".")[-2]
+            except ModuleNotFoundError:
+                strategy = "." + ".".join(strategy.split(".")[:-1])
+        else:
+            strategy = "." + ".".join(strategy.split(".")[:-1])
+        mod = importlib.import_module(strategy, module_base)
         func = getattr(mod, load_fn)
         return func(cfg, **kwargs)
-    except Exception:  # pylint: disable=broad-exception-caught
-        LOG.warning(f"unable to load strategy {strategy}")
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+        LOG.warning(f"unable to load strategy {strategy}: {exc}", exc_info=True)
         return None
diff --git a/src/axolotl/prompt_strategies/bradley_terry/README.md b/src/axolotl/prompt_strategies/bradley_terry/README.md
new file mode 100644
index 0000000000..39cd16137c
--- /dev/null
+++ b/src/axolotl/prompt_strategies/bradley_terry/README.md
@@ -0,0 +1,10 @@
+### example yaml
+
+```yaml
+chat_template: gemma
+datasets:
+  - path: argilla/distilabel-intel-orca-dpo-pairs
+    type: bradley_terry.chat_template
+val_set_size: 0.0
+output_dir: ./outputs/out
+```
diff --git a/src/axolotl/prompt_strategies/bradley_terry/__init__.py b/src/axolotl/prompt_strategies/bradley_terry/__init__.py
new file mode 100644
index 0000000000..7336edc717
--- /dev/null
+++ b/src/axolotl/prompt_strategies/bradley_terry/__init__.py
@@ -0,0 +1,34 @@
+"""Module to load prompt strategies."""
+
+import importlib
+import inspect
+
+from axolotl.prompt_strategies.user_defined import UserDefinedDatasetConfig
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def load(strategy, tokenizer, cfg, ds_cfg):
+    try:
+        load_fn = "load"
+        if strategy.split(".")[-1].startswith("load_"):
+            load_fn = strategy.split(".")[-1]
+            strategy = ".".join(strategy.split(".")[:-1])
+        mod = importlib.import_module(
+            f".{strategy}", "axolotl.prompt_strategies.bradley_terry"
+        )
+        func = getattr(mod, load_fn)
+        load_kwargs = {}
+        if strategy == "user_defined":
+            load_kwargs["ds_cfg"] = UserDefinedDatasetConfig(**ds_cfg)
+        else:
+            sig = inspect.signature(func)
+            if "ds_cfg" in sig.parameters:
+                load_kwargs["ds_cfg"] = ds_cfg
+        return func(tokenizer, cfg, **load_kwargs)
+    except ModuleNotFoundError:
+        return None
+    except Exception as exc:
+        LOG.error(f"Failed to load prompt strategy `{strategy}`: {str(exc)}")
+        return None
diff --git a/src/axolotl/prompt_strategies/bradley_terry/chat_template.py b/src/axolotl/prompt_strategies/bradley_terry/chat_template.py
new file mode 100644
index 0000000000..e3f34ad8f6
--- /dev/null
+++ b/src/axolotl/prompt_strategies/bradley_terry/chat_template.py
@@ -0,0 +1,122 @@
+"""
+Bradley-Terry model with chat template prompt strategy.
+"""
+
+from typing import Any, Dict, Optional
+
+from axolotl.prompt_strategies.chat_template import (
+    ChatTemplatePrompter,
+    ChatTemplateStrategy,
+)
+from axolotl.utils.chat_templates import get_chat_template_from_config
+from axolotl.utils.logging import get_logger
+
+# Configure the logger
+LOG = get_logger(__name__)
+LOG.setLevel("INFO")
+
+
+class BTChatTemplateStrategy(ChatTemplateStrategy):
+    """
+    Bradley-Terry reward model pairwise chat template prompt strategy.
+    """
+
+    @property
+    def supports_batched(self) -> bool:
+        return False
+
+    def _tokenize_single_prompt(self, prompt):
+        """
+
+        :param prompt: the actual row of data from the underlying dataset
+        :return:
+        """
+
+        max_length = self.prompter.max_length
+
+        prompt["messages"] = []
+        if prompt["system"]:
+            prompt["messages"].append({"role": "system", "content": prompt["system"]})
+        prompt["messages"].append({"role": "user", "content": prompt["input"]})
+        prompt["messages"].append({"role": "assistant", "content": prompt["chosen"]})
+        chosen_tokenized = super()._tokenize_single_prompt(prompt)
+
+        if len(chosen_tokenized["input_ids"]) > max_length:
+            LOG.warning(
+                f"To-be-trimmed chosen sequence exceeds max sequence length: {len(chosen_tokenized['input_ids'])}"
+            )
+
+            chosen_tokenized["input_ids"] = chosen_tokenized["input_ids"][:max_length]
+            chosen_tokenized["attention_mask"] = chosen_tokenized["attention_mask"][
+                :max_length
+            ]
+
+        prompt["messages"] = []
+        if prompt["system"]:
+            prompt["messages"].append({"role": "system", "content": prompt["system"]})
+        prompt["messages"].append({"role": "user", "content": prompt["input"]})
+        prompt["messages"].append({"role": "assistant", "content": prompt["rejected"]})
+        rejected_tokenized = super()._tokenize_single_prompt(prompt)
+
+        if len(rejected_tokenized["input_ids"]) > max_length:
+            LOG.warning(
+                f"To-be-trimmed rejected sequence exceeds max sequence length: {len(rejected_tokenized['input_ids'])}"
+            )
+
+            rejected_tokenized["input_ids"] = rejected_tokenized["input_ids"][
+                :max_length
+            ]
+            rejected_tokenized["attention_mask"] = rejected_tokenized["attention_mask"][
+                :max_length
+            ]
+
+        return {
+            "chosen_ids": chosen_tokenized["input_ids"],
+            "attention_mask_chosen": chosen_tokenized["attention_mask"],
+            "labels_chosen": 1.0,
+            "rejected_ids": rejected_tokenized["input_ids"],
+            "attention_mask_rejected": rejected_tokenized["attention_mask"],
+            "labels_rejected": 0.0,
+        }
+
+
+def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
+    ds_cfg = ds_cfg or {}
+    chat_template_string = get_chat_template_from_config(
+        cfg=cfg, ds_cfg=ds_cfg, tokenizer=tokenizer
+    )
+
+    prompter_params = {
+        "tokenizer": tokenizer,
+        "chat_template": chat_template_string,
+        "message_property_mappings": ds_cfg.get(
+            "message_property_mappings",
+            {
+                "role": "role",
+                "content": "content",
+            },
+        ),
+        "message_field_training": ds_cfg.get("message_field_training", None),
+        "message_field_training_detail": ds_cfg.get(
+            "message_field_training_detail", None
+        ),
+        "roles": ds_cfg.get("roles"),
+        "drop_system_message": ds_cfg.get("drop_system_message", False),
+        # we need to add one for detecting sequences with exceeding the `sequence_len` limit.
+        "max_length": (
+            cfg.sequence_len + 1 if not cfg.reward_model else cfg.sequence_len
+        ),
+    }
+
+    strategy_params = {
+        "train_on_inputs": cfg.train_on_inputs,
+        "sequence_len": cfg.sequence_len,
+        "roles_to_train": ds_cfg.get("roles_to_train", []),
+        "train_on_eos": ds_cfg.get("train_on_eos", None),
+    }
+
+    strategy = BTChatTemplateStrategy(
+        ChatTemplatePrompter(**prompter_params), tokenizer=tokenizer, **strategy_params
+    )
+
+    return strategy
diff --git a/src/axolotl/prompt_strategies/bradley_terry/llama3.py b/src/axolotl/prompt_strategies/bradley_terry/llama3.py
new file mode 100644
index 0000000000..5548d882eb
--- /dev/null
+++ b/src/axolotl/prompt_strategies/bradley_terry/llama3.py
@@ -0,0 +1,27 @@
+"""
+chatml transforms for datasets with system, input, chosen, rejected to match llama3 chat template
+"""
+
+
+def icr(
+    cfg,
+    **kwargs,
+):
+    """
+    chatml transforms for datasets with system, input, chosen, rejected
+    ex. https://huggingface.co/datasets/argilla/distilabel-intel-orca-dpo-pairs
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            prompt = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            prompt = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["chosen"] = prompt + f"{sample['chosen']}<|eot_id|>"
+        sample["rejected"] = prompt + f"{sample['rejected']}<|eot_id|>"
+        return sample
+
+    return transform_fn
diff --git a/src/axolotl/prompt_strategies/chat_template.py b/src/axolotl/prompt_strategies/chat_template.py
index 8dff3845b7..1db5bbc4cc 100644
--- a/src/axolotl/prompt_strategies/chat_template.py
+++ b/src/axolotl/prompt_strategies/chat_template.py
@@ -1,29 +1,278 @@
 """
 HF Chat Templates prompt strategy
 """
-from typing import Any, Dict, Optional
 
+import json
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Dict, List, Set, Union
+
+from pydantic import BaseModel
+from transformers import ProcessorMixin
+
+from axolotl.prompt_strategies.jinja_template_analyzer import JinjaTemplateAnalyzer
 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
-from axolotl.prompters import Prompter
-from axolotl.utils.chat_templates import chat_templates
+from axolotl.prompters import IGNORE_TOKEN_ID, Prompter
+from axolotl.utils.chat_templates import get_chat_template_from_config
+from axolotl.utils.dict import remove_none_values
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.datasets import DatasetConfig
+
+if TYPE_CHECKING:
+    from axolotl.utils.mistral import HFMistralTokenizer
+
+# Configure the logger
+LOG = get_logger(__name__)
+LOG.setLevel("INFO")
+
+
+def _extract_input_ids(result):
+    """Return the ``input_ids`` from a ``build_prompt`` result.
+
+    With a processor configured, ``build_prompt`` returns a dict of
+    processor outputs (``input_ids``, ``attention_mask``, optional
+    ``pixel_values``, etc.). Without a processor it returns a plain
+    ``list[int]`` from the tokenizer.
+    """
+    return result["input_ids"] if isinstance(result, dict) else result
 
 
 class ChatTemplatePrompter(Prompter):
-    """prompter for HF chat templates"""
+    """Prompter for HF chat templates"""
+
+    def __init__(
+        self,
+        tokenizer,
+        chat_template: str,
+        processor=None,
+        max_length=2048,
+        message_property_mappings: dict[str, str] | None = None,
+        message_field_training: str | None = None,
+        message_field_training_detail: str | None = None,
+        field_messages: str = "messages",
+        field_system: str = "system",
+        field_tools: str = "tools",
+        field_thinking: str = "reasoning_content",
+        roles: dict[str, list[str]] | None = None,
+        template_thinking_key: str | None = "reasoning_content",
+        chat_template_kwargs: dict[str, Any] | None = None,
+        drop_system_message: bool = False,
+    ):
+        # check if message_property_mappings is None or empty dict
+        if message_property_mappings is None or (not message_property_mappings):
+            message_property_mappings = {
+                "role": "role",
+                "content": "content",
+            }
+            if template_thinking_key and field_thinking:
+                message_property_mappings[template_thinking_key] = field_thinking
 
-    def __init__(self, tokenizer, chat_template=None, max_length=2048):
+        if roles:
+            self.roles = {s: t for t, sources in roles.items() for s in sources}
+        else:
+            self.roles = {
+                "human": "user",
+                "user": "user",
+                "assistant": "assistant",
+                "gpt": "assistant",
+                "system": "system",
+                "tool": "tool",
+            }
+
+        self._chat_template_msg_variables = self.get_chat_template_msg_variables(
+            chat_template, field_messages
+        )
+        self.message_property_mappings = message_property_mappings
+        self.message_field_training = message_field_training
+        self.message_field_training_detail = message_field_training_detail
+        self.field_messages = field_messages
+        self.field_system = field_system
+        self.field_tools = field_tools
+        self.field_thinking = field_thinking
         self.tokenizer = tokenizer
+        self.processor: ProcessorMixin | None = processor
         self.chat_template = chat_template
+        self.chat_template_kwargs = chat_template_kwargs or {}
+        self.template_thinking_key: str = template_thinking_key or "reasoning_content"
         self.max_length = max_length
+        self.drop_system_message = drop_system_message
+
+    @property
+    def chat_template_msg_variables(self) -> Set[str]:
+        return self._chat_template_msg_variables
+
+    def build_prompt(
+        self,
+        conversation: list[dict],
+        add_generation_prompt=False,
+        images=None,
+        tools=None,
+        real_last_index=None,
+    ):
+        """
+        Build a prompt from a conversation.
+
+        Args:
+            conversation: A list of messages.
+            add_generation_prompt: Whether to add a generation prompt.
+            images: A list of images. (optional)
+            tools: A list of tools. (optional)
+        """
+        chat_template_kwargs = {
+            "chat_template": self.chat_template,
+            "add_generation_prompt": add_generation_prompt,
+            **self.chat_template_kwargs,
+        }
+
+        if tools:
+            chat_template_kwargs["tools"] = tools
+
+        if real_last_index:
+            chat_template_kwargs["real_last_index"] = real_last_index
+
+        if self.processor:
+            if not callable(self.processor):
+                raise TypeError("Processor must be callable")
+
+            text = self.processor.apply_chat_template(
+                conversation,
+                tokenize=False,
+                **chat_template_kwargs,
+            )
+            batch = self.processor(
+                text=text,
+                images=images,
+                return_tensors="pt",
+            )
+            if hasattr(batch, "to_dict"):
+                batch = batch.to_dict()
+            else:
+                batch = dict(batch)
+
+            # workaround since processor works in batches instead of single examples
+            out = {}
+            for k, val in batch.items():
+                if hasattr(val, "tolist"):
+                    out[k] = (
+                        val.tolist() if k == "pixel_values" else val.squeeze(0).tolist()
+                    )
+                else:
+                    out[k] = val
+            return out
 
-    def build_prompt(self, conversation, add_generation_prompt=False):
         return self.tokenizer.apply_chat_template(
             conversation,
-            truncation=True,
-            max_length=self.max_length,
-            add_generation_prompt=add_generation_prompt,
-            chat_template=self.chat_template,
+            tokenize=True,
+            return_dict=False,
+            **chat_template_kwargs,
+        )
+
+    def get_offsets_for_train_detail(
+        self, text: str, train_details: List[Dict], mask_untrainable: bool = True
+    ) -> List[int]:
+        tokenized_output = self.tokenizer(
+            text, return_offsets_mapping=True, add_special_tokens=False
         )
+        tokens = tokenized_output.tokens()
+        token_offsets = tokenized_output["offset_mapping"]
+
+        LOG.debug(f"Tokenizing text: {text}")
+        LOG.debug(f"Tokens: {tokens}")
+        # Adjust the end offsets. For some reason by default they are set to the same value as the start offsets.
+        for i in range(len(token_offsets) - 1):
+            token_offsets[i] = (token_offsets[i][0], token_offsets[i + 1][0] - 1)
+        # Ensure the last token's end offset is set correctly
+        token_offsets[-1] = (token_offsets[-1][0], len(text) - 1)
+        LOG.debug(f"Token offsets: {token_offsets}")
+
+        # Initialize all offsets as IGNORE_TOKEN_ID (not trained)
+        result = [IGNORE_TOKEN_ID] * len(token_offsets)
+
+        # Adjust train_details to align with token boundaries
+        adjusted_train_details = self.adjust_train_details(train_details, token_offsets)
+
+        for idx, (start, end) in enumerate(token_offsets):
+            for detail in adjusted_train_details:
+                # Check if the token is completely within the detail's range
+                if start >= detail["begin_offset"] and end <= detail["end_offset"]:
+                    if detail["train"] or not mask_untrainable:
+                        result[idx] = start
+                        LOG.debug(f"Token {idx} ({tokens[idx]}) marked for training")
+                    else:
+                        LOG.debug(
+                            f"Token {idx} ({tokens[idx]}) marked as non-trainable"
+                        )
+                elif start < detail["end_offset"] and end > detail["begin_offset"]:
+                    # Token partially overlaps with detail, always mark as non-trainable
+                    LOG.debug(
+                        f"Token {idx} ({tokens[idx]}) partially overlaps detail, marked as non-trainable"
+                    )
+
+        LOG.debug(f"Final result: {result}")
+        return result
+
+    def adjust_train_details(
+        self, train_details: List[Dict], token_offsets: List[tuple]
+    ) -> List[Dict]:
+        adjusted_details = []
+        for detail in train_details:
+            begin_offset = detail["begin_offset"]
+            end_offset = detail["end_offset"]
+
+            # Find the first token that starts after or at the begin_offset
+            begin_token = next(
+                (
+                    i
+                    for i, (t_start, t_end) in enumerate(token_offsets)
+                    if t_start >= begin_offset
+                ),
+                len(token_offsets),
+            )
+            if begin_token > 0 and token_offsets[begin_token - 1][1] > begin_offset:
+                begin_token -= 1
+
+            # Find the last token that ends before or at the end_offset
+            end_token = next(
+                (
+                    i
+                    for i in range(len(token_offsets) - 1, -1, -1)
+                    if token_offsets[i][1] <= end_offset
+                ),
+                -1,
+            )
+            if (
+                end_token < len(token_offsets) - 1
+                and token_offsets[end_token + 1][0] < end_offset
+            ):
+                end_token += 1
+
+            if begin_token <= end_token:
+                adjusted_begin = token_offsets[begin_token][0]
+                adjusted_end = token_offsets[end_token][1]
+
+                if adjusted_begin != begin_offset or adjusted_end != end_offset:
+                    LOG.warning(
+                        f"Adjusting detail offsets: ({begin_offset}, {end_offset}) -> ({adjusted_begin}, {adjusted_end})"
+                    )
+
+                adjusted_details.append(
+                    {
+                        "begin_offset": adjusted_begin,
+                        "end_offset": adjusted_end,
+                        "train": detail["train"],
+                    }
+                )
+            else:
+                LOG.warning(
+                    f"Could not adjust detail offsets: ({begin_offset}, {end_offset}). Skipping this detail."
+                )
+
+        return adjusted_details
+
+    def get_chat_template_msg_variables(
+        self, chat_template: str, field_messages: str
+    ) -> Set[str]:
+        template_analyzer = JinjaTemplateAnalyzer(chat_template)
+        return template_analyzer.get_message_vars(field_messages)
 
 
 class ChatTemplateStrategy(PromptTokenizingStrategy):
@@ -31,48 +280,972 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
     Tokenizing strategy for instruction-based prompts.
     """
 
-    def tokenize_prompt(self, prompt):
+    def __init__(
+        self,
+        prompter: "ChatTemplatePrompter",
+        tokenizer,
+        train_on_inputs: bool,
+        sequence_len: int,
+        roles_to_train: list[str] | None = None,
+        train_on_eos: str | None = None,
+        train_on_eot: str | None = None,
+        eot_tokens: list[str] | None = None,
+        split_thinking: bool | None = False,
+    ):
+        super().__init__(prompter, tokenizer, train_on_inputs, sequence_len)
+        self.prompter: ChatTemplatePrompter = prompter
+
+        self.roles_to_train = []
+        if roles_to_train:
+            # map roles if exist in prompter.roles else use the role as is
+            self.roles_to_train = [
+                prompter.roles.get(role, role) for role in roles_to_train
+            ]
+
+        self.train_on_eos = train_on_eos
+        # Backward compatibility, load from train_on_eos
+        self.train_on_eot = train_on_eot if train_on_eot is not None else train_on_eos
+
+        # Default to eos_token if eot_tokens not provided
+        self.eot_tokens = []
+        if eot_tokens is not None:
+            self.eot_tokens = eot_tokens
+        elif (
+            hasattr(self.tokenizer, "eos_token")
+            and self.tokenizer.eos_token is not None
+        ):
+            self.eot_tokens = [self.tokenizer.eos_token]
+
+        self.split_thinking = split_thinking
+
+        self.images = "images"
+
+        LOG.debug(
+            f"The chat template uses the following properites on the message: {self.prompter.chat_template_msg_variables}"
+        )
+
+        self._validate_eot_and_eos_tokens()
+
+        # Pre-cache EOT token IDs to avoid re-encoding on every call
+        self._eot_token_ids = set()
+        for token in self.eot_tokens:
+            token_ids = self.tokenizer.encode(token, add_special_tokens=False)
+            if len(token_ids) == 1:
+                self._eot_token_ids.add(token_ids[0])
+
+    def _validate_eot_and_eos_tokens(self):
+        """
+        - Validates that EOT tokens (or eos_token) are in the chat_template
+        - Checks if EOT tokens are encoded as multiple tokens in the tokenizer.
+        - Checks for potential conflicts between train_on_eos and train_on_eot.
+        """
+        if self.prompter.chat_template is None:
+            # Usually this should not happen
+            LOG.warning(
+                "No chat template provided, skipping EOT and EOS token validation"
+            )
+            return
+
+        # If the EOT token is the same as the EOS token, we need to check differently
+        if len(self.eot_tokens) == 1 and self.eot_tokens[0] == self.tokenizer.eos_token:
+            # Check if the eos_token is in the chat_template or as a variable `eos_token`
+            # Note: we check for `eos_token` in the string, but it could possibly not be a variable
+            if (
+                self.tokenizer.eos_token not in self.prompter.chat_template
+                and "eos_token" not in self.prompter.chat_template
+            ):
+                LOG.warning(
+                    f"EOS token '{self.tokenizer.eos_token}' not found in chat_template; the turn "
+                    "terminator won't be trained. Set `eot_tokens` to your template's turn-ending token."
+                )
+            return
+
+        # Create a new list to store tokens that should be kept
+        valid_eot_tokens = []
+        for token in self.eot_tokens:
+            # Check if EOT token is in the chat_template
+            if token not in self.prompter.chat_template:
+                LOG.warning(f"EOT token '{token}' not found in chat_template.")
+                # Don't add to the valid tokens list
+                continue
+
+            valid_eot_tokens.append(token)
+
+        # Replace the original list with the filtered one
+        self.eot_tokens = valid_eot_tokens
+
+        for token in self.eot_tokens:
+            # If token in template, check if EOT token is in tokenizer and not encoded as multiple tokens
+            token_ids = self.tokenizer.encode(token, add_special_tokens=False)
+            if not token_ids:
+                raise ValueError(
+                    "EOT token encoding failed. Please check if the token is valid and can be encoded."
+                )
+            if token_ids and len(token_ids) > 1:
+                raise ValueError(
+                    f"EOT token '{token}' is encoded as multiple tokens: {token_ids}. Please add it under `tokens: ` in the config "
+                    "or (recommended) override unused added_tokens via `added_tokens_overrides: `."
+                )
+
+        # If eos_token is in eot_tokens and conflict between train_on_eos and train_on_eot, raise an error
+        if (
+            self.tokenizer.eos_token in self.eot_tokens
+            and self.train_on_eos != self.train_on_eot
+        ):
+            raise ValueError(
+                "Conflict between train_on_eos and train_on_eot. eos_token is in eot_tokens and train_on_eos != train_on_eot"
+                f"train_on_eos: {self.train_on_eos}, train_on_eot: {self.train_on_eot}"
+                f"eot_tokens: {self.eot_tokens}"
+                f"eos_token: {self.tokenizer.eos_token}"
+            )
+
+    @property
+    def supports_batched(self) -> bool:
+        # Let calling code know we can handle lists of examples
+        return True
+
+    def is_prompt_batched(self, prompt: dict[str, Any]) -> bool:
+        try:
+            return all(isinstance(v, (str, list)) for v in prompt.values()) and all(
+                isinstance(v, (str, list)) for v in prompt[self.prompter.field_messages]
+            )
+        except KeyError:
+            return False
+
+    def tokenize_prompt(self, prompt: dict[str, Any]):
+        """
+        Public method that can handle either a single prompt or a batch of prompts.
+        """
+
+        prompt = remove_none_values(prompt)
+
+        if not self.is_prompt_batched(prompt) or not self.supports_batched:
+            return self._tokenize_single_prompt(prompt)
+
+        res = defaultdict(lambda: [])
+        feature_names = list(prompt.keys())
+
+        # Process each prompt individually
+        for row in zip(*prompt.values(), strict=False):
+            tokenized_prompt = self._tokenize_single_prompt(
+                dict(zip(feature_names, row, strict=False))
+            )
+            for key, val in tokenized_prompt.items():
+                res[key].append(val)
+
+        # If there are no examples left, return an empty dictionary
+        if not res:
+            return {}
+
+        return dict(res)
+
+    def _tokenize_single_prompt(self, prompt: dict) -> Dict[str, List[int]]:
+        # Old simple legacy behavior that works reliably.
+        if (
+            not self.roles_to_train
+            and not self.train_on_eos
+            and not self.train_on_eot
+            and not self.prompter.message_field_training  # type: ignore
+            and not self.prompter.message_field_training_detail  # type: ignore
+        ):
+            turns = self.get_conversation_thread(prompt)
+            images = self._get_images(prompt)
+            prompt_ids = self.prompter.build_prompt(  # type: ignore
+                turns[:-1],
+                add_generation_prompt=True,
+                images=images,
+            )
+            tokenized_res = self.prompter.build_prompt(turns, images=images)  # type: ignore
+            tokenized_prompt = {}
+            if isinstance(tokenized_res, list):
+                input_ids = prompt_ids + tokenized_res[len(prompt_ids) :]
+                tokenized_prompt["input_ids"] = input_ids
+                tokenized_prompt["attention_mask"] = [1] * len(input_ids)
+            else:
+                input_ids = tokenized_res["input_ids"]
+                tokenized_prompt = dict(tokenized_res)
+
+            if not self.train_on_inputs:
+                if isinstance(prompt_ids, dict):
+                    user_prompt_len = len(prompt_ids["input_ids"])
+                else:
+                    user_prompt_len = len(prompt_ids)
+                labels = [-100] * user_prompt_len + input_ids[user_prompt_len:]
+            else:
+                labels = input_ids
+
+            tokenized_prompt["labels"] = labels
+
+            return tokenized_prompt
+
         turns = self.get_conversation_thread(prompt)
-        prompt_ids = self.prompter.build_prompt([turns[0]], add_generation_prompt=True)
-        input_ids = self.prompter.build_prompt(turns)
+        tools = self._get_tools(prompt)
+        result = self.prompter.build_prompt(turns, tools=tools)  # type: ignore
+        if not isinstance(result, dict):
+            result = {"input_ids": result}
+        input_ids = result["input_ids"]
+        labels = [IGNORE_TOKEN_ID] * len(input_ids)
+
+        last_eos_idx = -1
+        last_eot_idx = -1
+        for index, turn in enumerate(turns):
+            role = turn.get("role")
+            content = turn.get("content")
+            train_turn = turn.get("training")
+            train_detail = turn.get("training_detail")
+            reasoning_train_detail = turn.get("reasoning_training_detail")
+
+            LOG.debug(
+                f"Processing turn {index}: role={role}, content={content}, train_turn={train_turn}, train_detail={train_detail}"
+            )
+
+            should_train = None
+            if train_turn is not None:
+                should_train = train_turn
+            elif train_detail is not None or reasoning_train_detail is not None:
+                should_train = bool(train_detail) or bool(reasoning_train_detail)
+            else:
+                should_train = self.train_on_inputs or role in self.roles_to_train
+
+            LOG.debug(f"Should train: {should_train}")
+
+            # turn not trainable, skip having to find the turn indices
+            # unless last turn and train_on_eos/train_on_eot is all
+            if not should_train and (
+                self.train_on_eos != "all" and self.train_on_eot != "all"
+            ):
+                if index == len(turns) - 1:
+                    LOG.warning(
+                        "Last turn is not trainable, skipping having to find the turn indices. "
+                        "This may cause incorrect last EOT/EOS token to be unmasked."
+                        "This is likely a dataset design issue. Please ensure last turn is trainable."
+                    )
+
+                continue
+
+            thinking_key = self.prompter.template_thinking_key
+            has_reasoning = thinking_key and turn.get(thinking_key) is not None
+            has_any_detail = train_detail or reasoning_train_detail
+
+            # When train_detail is present and the turn has reasoning_content,
+            # use content_only=True so find_turn returns content-only boundaries
+            # (excluding reasoning_content + template separator tokens).
+            use_content_only = bool(has_any_detail and has_reasoning)
+
+            turn_start_idx, turn_end_idx = self.find_turn(
+                turns=turns,
+                turn_idx=index,
+                tools=tools,
+                content_only=use_content_only,
+            )
 
-        if not self.train_on_inputs:
-            user_prompt_len = len(prompt_ids)
-            labels = [-100] * user_prompt_len + input_ids[user_prompt_len:]
+            LOG.debug(f"Turn indices: start={turn_start_idx}, end={turn_end_idx}")
+
+            if should_train and turn_start_idx != -1 and turn_end_idx != -1:
+                if train_detail:
+                    if not isinstance(content, str):
+                        raise ValueError(
+                            "`train_detail` is not supported when `content` is not a string."
+                        )
+
+                    token_offsets = self.prompter.get_offsets_for_train_detail(  # type: ignore
+                        content, train_detail
+                    )
+                    LOG.debug(f"Token offsets: {token_offsets}")
+                    for i, offset in enumerate(token_offsets):
+                        if offset != IGNORE_TOKEN_ID and turn_start_idx + i < len(
+                            input_ids
+                        ):
+                            labels[turn_start_idx + i] = input_ids[turn_start_idx + i]
+                            LOG.debug(
+                                f"Label set at index {turn_start_idx + i}: {input_ids[turn_start_idx + i]}"
+                            )
+                elif not reasoning_train_detail:
+                    # No per-part detail on either field — train the whole span
+                    labels[turn_start_idx:turn_end_idx] = input_ids[
+                        turn_start_idx:turn_end_idx
+                    ]
+                    LOG.debug(
+                        f"Set labels for training from {turn_start_idx} to {turn_end_idx}"
+                    )
+
+            # Handle reasoning_content training_detail separately
+            if should_train and reasoning_train_detail and has_reasoning:
+                reasoning_text = turn[thinking_key]
+                if not isinstance(reasoning_text, str):
+                    raise ValueError(
+                        "`reasoning_training_detail` is not supported when reasoning_content is not a string."
+                    )
+
+                reasoning_start, reasoning_end = self.find_turn(
+                    turns=turns,
+                    turn_idx=index,
+                    tools=tools,
+                    reasoning_only=True,
+                )
+
+                if reasoning_start != -1 and reasoning_end != -1:
+                    token_offsets = self.prompter.get_offsets_for_train_detail(  # type: ignore
+                        reasoning_text, reasoning_train_detail
+                    )
+                    LOG.debug(f"Reasoning token offsets: {token_offsets}")
+                    for i, offset in enumerate(token_offsets):
+                        if offset != IGNORE_TOKEN_ID and reasoning_start + i < len(
+                            input_ids
+                        ):
+                            labels[reasoning_start + i] = input_ids[reasoning_start + i]
+
+                LOG.debug(f"Labels after processing turn {index}: {labels}")
+
+            # Handle special tokens (EOT and EOS)
+            for token_type, find_func, train_option in [
+                ("EOT", self.find_first_eot_token, self.train_on_eot),
+                ("EOS", self.find_first_eos_token, self.train_on_eos),
+            ]:
+                token_idx = find_func(input_ids, start_idx=turn_end_idx)
+
+                if (
+                    token_idx != -1 and abs(token_idx - turn_end_idx) <= 3
+                ):  # Allow for some template padding
+                    # Update the last token index
+                    if token_type == "EOT":  # nosec B105
+                        last_eot_idx = token_idx
+                    else:
+                        last_eos_idx = token_idx
+
+                    # Set labels if needed for this turn
+                    if train_option == "all" or (
+                        train_option == "turn" and should_train
+                    ):
+                        labels[token_idx] = input_ids[token_idx]
+                        LOG.debug(
+                            f"{token_type} token set for training at index {token_idx}"
+                        )
+                else:
+                    LOG.debug(
+                        f"{token_type} token missing after turn {turn}. {token_type.lower()}_idx: {token_idx}, turn_end_idx: {turn_end_idx}"
+                    )
+
+        # Handle 'last' option for special tokens
+        for token_type, last_idx, train_option in [
+            ("EOT", last_eot_idx, self.train_on_eot),
+            ("EOS", last_eos_idx, self.train_on_eos),
+        ]:
+            if train_option == "last" and last_idx != -1:
+                labels[last_idx] = input_ids[last_idx]
+                LOG.debug(
+                    f"Last {token_type} token set for training at index {last_idx}"
+                )
+
+        LOG.debug(f"Final labels: {labels}")
+
+        # ``result`` already carries any processor outputs (pixel_values, image
+        # grid info, etc.); just set the fields we computed locally.
+        result["labels"] = labels
+        result.setdefault("attention_mask", [1] * len(input_ids))
+        return result
+
+    def find_first_eos_token(self, input_ids, start_idx):
+        eos_token_id = self.tokenizer.eos_token_id
+        for i in range(start_idx, len(input_ids)):
+            if input_ids[i] == eos_token_id:
+                return i
+        return -1
+
+    def find_first_eot_token(self, input_ids, start_idx):
+        """Find the first EOT token in the input_ids starting from start_idx."""
+        # Use pre-cached EOT token IDs (computed once in __init__)
+        for i in range(start_idx, len(input_ids)):
+            if input_ids[i] in self._eot_token_ids:
+                return i
+        return -1
+
+    def find_turn(
+        self,
+        turns: list[dict],
+        turn_idx: int,
+        tools: list[dict] | None = None,
+        content_only: bool = False,
+        reasoning_only: bool = False,
+    ):
+        """
+        Locate the starting and ending indices of the specified turn in a conversation.
+
+        Args:
+            content_only: If True and the turn has reasoning_content (template_thinking_key),
+                preserve reasoning_content in the dummy turn so the diff only captures the
+                content field boundaries. This is needed for correct training_detail alignment
+                when reasoning_content is present.
+            reasoning_only: If True, preserve content in the dummy turn and replace
+                reasoning_content with a dummy, so the diff only captures the
+                reasoning_content field boundaries.
+        """
+
+        if turn_idx >= len(turns):
+            raise ValueError(f"Turn index {turn_idx} out of range")
+
+        # mistral/gemma3 does not output message if it contains only system message
+        if (
+            turn_idx == 0
+            and turns[0].get("role") == "system"
+            and ("mistral" in self.tokenizer.name_or_path.lower())
+        ):
+            return -1, -1
+
+        thinking_key = self.prompter.template_thinking_key
+
+        if reasoning_only:
+            # Keep content as-is, replace reasoning with dummy
+            empty_turn = {
+                "role": turns[turn_idx].get("role"),
+                "content": turns[turn_idx].get("content", ""),
+            }
+            if thinking_key and thinking_key in turns[turn_idx]:
+                empty_turn[thinking_key] = "[[dummy_reasoning]]"
         else:
-            labels = input_ids
+            empty_turn = {
+                "role": turns[turn_idx].get("role"),
+                "content": "[[dummy_message]]",
+            }
 
-        tokenized_prompt = {
-            "input_ids": input_ids,
-            "labels": labels,
-            "attention_mask": [1] * len(input_ids),
-        }
+            # When content_only is True, copy reasoning_content to the dummy turn so
+            # the diff only captures the content field (not reasoning + separator).
+            if content_only and thinking_key and thinking_key in turns[turn_idx]:
+                empty_turn[thinking_key] = turns[turn_idx][thinking_key]
+
+        # Create conversation versions
+        turns_with_empty = turns[:turn_idx] + [empty_turn]
+        turns_with_content = turns[: turn_idx + 1]
+
+        real_last_index = len(turns) - 1
+
+        # Generate the conversation up to the turn, with final turn replaced with dummy content
+        dummy_ids = _extract_input_ids(
+            self.prompter.build_prompt(  # type: ignore
+                turns_with_empty, tools=tools, real_last_index=real_last_index
+            )
+        )
 
-        return tokenized_prompt
+        # Generate the conversation up to the turn, with final turn included
+        full_ids = _extract_input_ids(
+            self.prompter.build_prompt(  # type: ignore
+                turns_with_content, tools=tools, real_last_index=real_last_index
+            )
+        )
+
+        if not full_ids or not dummy_ids:
+            LOG.warning(f"Empty template generated for turn {turn_idx}")
+            return -1, -1
+
+        # Find first difference (start of content)
+        start_idx = None
+        min_len = min(len(dummy_ids), len(full_ids))
+        for i in range(min_len):
+            if dummy_ids[i] != full_ids[i]:
+                start_idx = i
+                break
+
+        if start_idx is None:
+            LOG.warning(f"Could not find content start boundary for turn {turn_idx}")
+            return -1, -1
+
+        # Find last difference (end of content)
+        end_idx = None
+        for i in range(min_len):
+            dummy_pos = len(dummy_ids) - 1 - i
+            full_pos = len(full_ids) - 1 - i
+            if dummy_ids[dummy_pos] != full_ids[full_pos]:
+                end_idx = full_pos + 1  # Add one to include the last token when slice
+                break
+
+        if end_idx is None:
+            LOG.warning(f"Could not find content end boundary for turn {turn_idx}")
+            return -1, -1
+
+        if end_idx < start_idx:
+            LOG.warning(
+                f"Content end boundary is before start boundary for turn {turn_idx}"
+            )
+            return -1, -1
+
+        if end_idx == start_idx:
+            LOG.warning(
+                f"Content end boundary is the same as start boundary for turn {turn_idx}. This is likely an empty turn."
+            )
+            return -1, -1
+
+        LOG.debug(f"Content boundaries: {start_idx}, {end_idx}")
+        LOG.debug(
+            f"Content tokens: {self.tokenizer.convert_ids_to_tokens(full_ids[start_idx:end_idx])}"
+        )
+
+        return start_idx, end_idx
+
+    @staticmethod
+    def _convert_content_parts(
+        content,
+    ) -> tuple[str, list[dict] | None] | None:
+        """Convert list content to concatenated string + optional training_detail.
+
+        When content is a list of dicts (content parts), each part can specify:
+        - ``text``, ``content``, or ``value``: the text string
+        - ``train`` (bool) or ``weight`` (0/1): per-part training flag
+
+        Returns ``(concatenated_text, training_details_or_None)`` if content was
+        a list, or ``None`` if content was not a list (no conversion needed).
+
+        .. note::
+            **Whitespace at part boundaries matters.** BPE tokenizers prepend
+            spaces to word tokens (e.g. ``" answer"`` is one token). Always
+            split BEFORE spaces::
+
+                GOOD: ["Let me think...", " The answer is 4."]
+                BAD:  ["Let me think... ", "The answer is 4."]
+
+            Tokens that straddle a boundary are conservatively masked.
+            Newlines typically merge with preceding punctuation (``":\\n"`` is
+            one token), so keep newlines with the preceding part.
+        """
+        if not isinstance(content, list):
+            return None
+
+        text_parts: list[str] = []
+        training_details: list[dict] = []
+        has_explicit_training = False
+        offset = 0
+
+        for part in content:
+            if isinstance(part, dict):
+                # Extract text (HF uses "text", also support "content"/"value")
+                text = (
+                    part.get("text") or part.get("content") or part.get("value") or ""
+                )
+                text_parts.append(text)
+
+                # Check for per-part training flags
+                part_train = part.get("train")
+                part_weight = part.get("weight")
+                if part_train is not None or part_weight is not None:
+                    has_explicit_training = True
+                    train = (
+                        part_train
+                        if part_train is not None
+                        else (part_weight not in (0, 0.0))
+                    )
+                else:
+                    train = True  # default trainable, gated by turn-level should_train
+
+                if text:
+                    training_details.append(
+                        {
+                            "begin_offset": offset,
+                            "end_offset": offset + len(text) - 1,
+                            "train": train,
+                        }
+                    )
+                    offset += len(text)
+
+        # Warn about trailing whitespace at boundaries between parts with
+        # different training flags — this almost always causes token straddling
+        if has_explicit_training and len(training_details) > 1:
+            for i in range(len(training_details) - 1):
+                cur = training_details[i]
+                nxt = training_details[i + 1]
+                if cur["train"] != nxt["train"]:
+                    boundary_text = text_parts[i]
+                    if boundary_text and boundary_text[-1] in (" ", "\t"):
+                        LOG.warning(
+                            "Content part %d ends with whitespace at a train/mask boundary. "
+                            "BPE tokenizers typically prepend spaces to word tokens, so "
+                            "the space will merge with the next part's first word and the "
+                            "resulting token will be MASKED (not trained). Move the "
+                            "whitespace to the start of the next content part instead. "
+                            "Part text: %r",
+                            i,
+                            boundary_text[-20:],
+                        )
+
+        concatenated = "".join(text_parts)
+        details = training_details if has_explicit_training else None
+        return concatenated, details
 
     def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
-        # remap roles - allow for assistant turn
-        role_map = {
-            "human": "user",
-            "user": "user",
-            "assistant": "assistant",
-            "gpt": "assistant",
-        }
-        turns = [
-            {"role": role_map[t["from"]], "content": t["value"]} for t in conversations
-        ]
+        turns = []
+
+        messages = self._get_messages(prompt)
+
+        possible_sys_turn = self.transform_message(messages[0])
+
+        if (
+            possible_sys_turn["role"] != "system"
+            and self.prompter.field_system in prompt
+        ):
+            turn = {"role": "system", "content": prompt[self.prompter.field_system]}
+            turns.append(turn)
+
+        for message in messages:
+            transformed_message = self.transform_message(message)
+
+            turn = transformed_message
+
+            training = message.get(self.prompter.message_field_training)
+            training_detail = message.get(self.prompter.message_field_training_detail)
+            if training is not None:
+                turn["training"] = training
+            if training_detail is not None:
+                turn["training_detail"] = training_detail
+
+            # Convert list content/reasoning_content to string + auto-generated
+            # training_detail. See _convert_content_parts for whitespace guidance.
+            content_result = self._convert_content_parts(turn.get("content"))
+            if content_result is not None:
+                turn["content"] = content_result[0]
+                if content_result[1] is not None:
+                    turn["training_detail"] = content_result[1]
+
+            # Also convert reasoning_content (template_thinking_key) if it's a list
+            thinking_key = self.prompter.template_thinking_key
+            if thinking_key and thinking_key in turn:
+                reasoning_result = self._convert_content_parts(turn[thinking_key])
+                if reasoning_result is not None:
+                    turn[thinking_key] = reasoning_result[0]
+                    if reasoning_result[1] is not None:
+                        turn["reasoning_training_detail"] = reasoning_result[1]
+
+            turns.append(turn)
+
+        if self.prompter.drop_system_message and turns[0]["role"] == "system":
+            turns = turns[1:]
+
         return turns
 
+    def transform_message(self, message: dict) -> dict:
+        # Build the initial transformed message from the mappings
+        transformed_message = {}
+        for key, value in self.prompter.message_property_mappings.items():
+            if message.get(value) is not None:
+                transformed_message[key] = message[value]
+            else:
+                LOG.debug(
+                    f"Could not find value for property {value} in message: {message}"
+                )
+
+        # Map the role if necessary
+        if "role" in transformed_message:
+            transformed_message["role"] = self.prompter.roles.get(
+                transformed_message["role"], transformed_message["role"]
+            )
+
+        # TODO handle reasoning_content with split_thinking
+        # if the role is assistant that we want to use reasoning_content
+        if self.split_thinking and transformed_message["role"] == "assistant":
+            content = transformed_message["content"]
+            thinking_pairs = [
+                ("<think>", "</think>"),
+                ("<reasoning>", "</reasoning>"),
+                ("<|begin_of_thought|>", "<|end_of_thought|>"),
+            ]
+            content_pairs = [("<|begin_of_solution|>", "<|end_of_solution|>")]
+            for tpair in thinking_pairs:
+                # check if the thinking pair is in the content
+                if tpair[0] in content and tpair[1] in content:
+                    # find the start and end index of the thinking pair
+                    t_start_idx = content.find(tpair[0])
+                    t_end_idx = content.find(tpair[1])
+
+                    # get the thinking content
+                    thinking_content = content[t_start_idx + len(tpair[0]) : t_end_idx]
+                    transformed_message[self.prompter.template_thinking_key] = (
+                        thinking_content.strip()
+                    )
+
+                    # take remainder of the content
+                    # strip whitespace from beginning of the remainder (thinking tokens)
+                    remainder = content[t_end_idx + len(tpair[1]) :].lstrip()
+
+                    # check if the content pair is in the remainder
+                    cpair_found = False
+                    for cpair in content_pairs:
+                        if cpair[0] in remainder and cpair[1] in remainder:
+                            # find the start and end index of the content pair
+                            c_start_idx = remainder.find(cpair[0])
+                            c_end_idx = remainder.find(cpair[1])
+
+                            # get the content content
+                            content_content = remainder[
+                                c_start_idx + len(cpair[0]) : c_end_idx
+                            ]
+                            transformed_message["content"] = content_content.strip()
+                            cpair_found = True
+                            break
+
+                    # else, the content is the remainder
+                    if not cpair_found:
+                        transformed_message["content"] = remainder
+                    break
+
+        # Determine which keys in the original message were not mapped
+        mapped_values = set(self.prompter.message_property_mappings.values())
+        remaining_keys = set(message) - mapped_values
+
+        # Keep only the properties defined in the chat template
+        # and not already mapped
+        for key in self.prompter.chat_template_msg_variables:
+            if key in remaining_keys:
+                val = message.get(key)
+                if val is not None:
+                    transformed_message[key] = val
+
+        if "tool_calls" in transformed_message and transformed_message["tool_calls"]:
+            for tool_call in transformed_message["tool_calls"]:
+                if "function" in tool_call and "arguments" in tool_call["function"]:
+                    args = tool_call["function"]["arguments"]
+                    if isinstance(args, str):
+                        try:
+                            tool_call["function"]["arguments"] = json.loads(args)
+                        except json.JSONDecodeError as e:
+                            LOG.error(
+                                f"Error parsing tool_calls arguments as JSON. "
+                                f"Function: {tool_call.get('function', {}).get('name', 'unknown')}, "
+                                f"Arguments string: {args!r}, "
+                                f"Error: {e}"
+                            )
+                            raise
+
+        return transformed_message
+
+    def _get_images(self, prompt):
+        return prompt.get(self.images, None)
+
+    def _get_tools(self, prompt) -> list[dict] | None:
+        """Get tools from prompt if available."""
+        tools = prompt.get(self.prompter.field_tools, None)
+        if tools is None:
+            return None
+
+        # Some datasets have tools set to str
+        if isinstance(tools, str):
+            try:
+                tools = json.loads(tools)
+            except json.JSONDecodeError as e:
+                LOG.error(f"Error parsing tool parameters as JSON. Error: {e}")
+                raise
+        if isinstance(tools, list):
+            # Process each tool to handle JSON string parameters
+            for tool in tools:
+                if isinstance(tool, dict) and "function" in tool:
+                    function = tool["function"]
+                    if "parameters" in function:
+                        params = function["parameters"]
+                        if isinstance(params, str):
+                            try:
+                                function["parameters"] = json.loads(params)
+                            except json.JSONDecodeError as e:
+                                LOG.error(
+                                    f"Error parsing tool parameters as JSON. "
+                                    f"Function: {function.get('name', 'unknown')}, "
+                                    f"Parameters string: {params!r}, "
+                                    f"Error: {e}"
+                                )
+                                raise
+            return tools
+
+        raise ValueError(
+            "Unknown tools format. Please convert it into a list[dict].\n"
+            f"Current format: {type(tools)}"
+        )
+
+    def _get_messages(self, prompt):
+        messages = prompt.get(self.prompter.field_messages, None)
+        if messages is None:
+            raise ValueError("Messages is null. Please check `field_messages`.")
+
+        if isinstance(messages, str):
+            try:
+                messages = json.loads(messages)
+            except json.JSONDecodeError as e:
+                LOG.error(f"Error parsing messages as JSON. Error: {e}")
+                raise
+            assert isinstance(messages, list), (
+                f"For SFT datasets that are stored in `str` format, the turns must be saved in a list of dictionaries, got {type(messages)}"
+            )
+
+            # Extra check here to make sure decoded json is a list of dicts.
+            for i, message in enumerate(messages):
+                assert isinstance(message, dict), (
+                    f"For SFT datasets that are stored in `str` format, each turns must be saved in a dictionary, got {type(message)} for the turn {i}"
+                )
+
+        if isinstance(messages, list):
+            return messages
+
+        raise ValueError(
+            "Unknown messages format. Please convert it into a list[dict].\n"
+            f"Current format: {type(messages)}"
+        )
+
+
+class MistralStrategy(ChatTemplateStrategy):
+    """
+    Mistral strategy for chat template.
+    """
+
+    def __init__(
+        self,
+        prompter: "ChatTemplatePrompter",
+        tokenizer: "HFMistralTokenizer",
+        train_on_inputs: bool,
+        sequence_len: int,
+        roles_to_train: list[str] | None = None,
+        train_on_eos: str | None = None,
+        train_on_eot: str | None = None,
+        eot_tokens: list[str] | None = None,
+        split_thinking: bool | None = False,
+    ):
+        # Call the parent's parent __init__ (PromptTokenizingStrategy) to skip ChatTemplateStrategy's validation
+
+        PromptTokenizingStrategy.__init__(
+            self, prompter, tokenizer, train_on_inputs, sequence_len
+        )
+        self.prompter: ChatTemplatePrompter = prompter
+
+        self.roles_to_train = []
+        if roles_to_train:
+            # map roles if exist in prompter.roles else use the role as is
+            self.roles_to_train = [
+                prompter.roles.get(role, role) for role in roles_to_train
+            ]
+
+        self.train_on_eos = train_on_eos
+        # Backward compatibility, load from train_on_eos
+        self.train_on_eot = train_on_eot if train_on_eot is not None else train_on_eos
+
+        # Default to eos_token if eot_tokens not provided
+        self.eot_tokens = []
+        if eot_tokens is not None:
+            self.eot_tokens = eot_tokens
+        else:
+            # set eot_tokens to the eos_token
+            self.eot_tokens = [self.tokenizer.eos_token]
 
-def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
-    chat_template = (
-        ds_cfg["chat_template"] if ds_cfg and "chat_template" in ds_cfg else "chatml"
-    )
-    strategy = ChatTemplateStrategy(
-        ChatTemplatePrompter(tokenizer, chat_templates(chat_template)),
+        self.split_thinking = split_thinking
+
+        self.images = "images"
+
+        LOG.debug(
+            f"The chat template uses the following properites on the message: {self.prompter.chat_template_msg_variables}"
+        )
+
+        # Skip the validation that ChatTemplateStrategy calls
+        # TODO: address this in the future with mistral-specific checks
+        # self._validate_eot_and_eos_tokens()
+
+    def find_first_eot_token(self, input_ids, start_idx):
+        """Find the first EOT token in the input_ids starting from start_idx."""
+        # mistral-common tokenizer does not support eot_tokens
+        return self.find_first_eos_token(input_ids, start_idx)
+
+
+class MistralPrompter(ChatTemplatePrompter):
+    """
+    Mistral prompter for chat template.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._chat_template_msg_variables = set(["tool_call_id", "name", "tool_calls"])
+
+
+class StrategyLoader:
+    """
+    Load chat template strategy based on configuration.
+    """
+
+    def _get_strategy_cls(self, cfg):
+        if cfg.tokenizer_use_mistral_common:
+            return MistralStrategy
+
+        return ChatTemplateStrategy
+
+    def _get_prompter_cls(self, cfg):
+        if cfg.tokenizer_use_mistral_common:
+            return MistralPrompter
+
+        return ChatTemplatePrompter
+
+    def _get_strategy_params(self, cfg, ds_cfg: Dict[str, Any]):
+        return {
+            "train_on_inputs": cfg.train_on_inputs,
+            "sequence_len": cfg.sequence_len,
+            "roles_to_train": ds_cfg.get("roles_to_train", ["assistant"]),
+            "train_on_eos": ds_cfg.get("train_on_eos", "turn"),
+            "train_on_eot": ds_cfg.get("train_on_eot", None),
+            "eot_tokens": cfg.get("eot_tokens", None),  # loads from cfg, not ds_cfg
+            "split_thinking": ds_cfg.get("split_thinking", False),
+        }
+
+    def __call__(
+        self,
         tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-    return strategy
+        cfg,
+        ds_cfg: Union[Dict[str, Any], DatasetConfig] | None = None,
+        processor=None,
+    ):
+        if ds_cfg is None:
+            dataset_config = {}
+        elif isinstance(ds_cfg, BaseModel):
+            dataset_config = ds_cfg.model_dump()
+        else:
+            dataset_config = ds_cfg
+
+        if cfg.tokenizer_use_mistral_common:
+            # mistral-common does not use this, so we pass an empty string
+            chat_template_string = ""
+        else:
+            chat_template_string = get_chat_template_from_config(
+                cfg=cfg, ds_cfg=dataset_config, tokenizer=tokenizer
+            )
+
+        LOG.info(f"Using chat template:\n---\n{chat_template_string!s}\n---")
+
+        prompter_params = {
+            "tokenizer": tokenizer,
+            "chat_template": chat_template_string,
+            "chat_template_kwargs": cfg.get("chat_template_kwargs", {}),
+            "message_property_mappings": dataset_config.get(
+                "message_property_mappings", {}
+            ),
+            "message_field_training": dataset_config.get(
+                "message_field_training", None
+            ),
+            "message_field_training_detail": dataset_config.get(
+                "message_field_training_detail",
+                None,
+            ),
+            "field_messages": dataset_config.get("field_messages", "messages"),
+            "field_thinking": dataset_config.get("field_thinking", "reasoning_content"),
+            "template_thinking_key": dataset_config.get(
+                "template_thinking_key", "reasoning_content"
+            ),
+            "roles": dataset_config.get("roles"),
+            "drop_system_message": dataset_config.get("drop_system_message", False),
+            # we need to add one for detecting sequences with exceeding the `sequence_len` limit.
+            "max_length": cfg.sequence_len + 1,
+            "processor": processor,
+        }
+
+        strategy_params = self._get_strategy_params(cfg, dataset_config)
+        strategy_cls = self._get_strategy_cls(cfg)
+        prompter_cls = self._get_prompter_cls(cfg)
+
+        strategy = strategy_cls(
+            prompter_cls(**prompter_params),
+            tokenizer=tokenizer,
+            **strategy_params,
+        )
+
+        return strategy
+
+
+load = StrategyLoader()
diff --git a/src/axolotl/prompt_strategies/completion.py b/src/axolotl/prompt_strategies/completion.py
index 3285e667cb..f43f257937 100644
--- a/src/axolotl/prompt_strategies/completion.py
+++ b/src/axolotl/prompt_strategies/completion.py
@@ -1,6 +1,7 @@
 """
 Basic completion text
 """
+
 from collections import defaultdict
 from typing import Any, Dict, Generator, Optional, Tuple
 
@@ -41,8 +42,8 @@ def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
     def tokenize_prompt(self, prompt):
         res = defaultdict(lambda: [])
         feature_names = list(prompt.keys())
-        for row in zip(*prompt.values()):
-            prompt_row = dict(zip(feature_names, row))
+        for row in zip(*prompt.values(), strict=False):
+            prompt_row = dict(zip(feature_names, row, strict=False))
             (
                 instruction,
                 _,
@@ -58,9 +59,7 @@ def tokenize_prompt(self, prompt):
 
         return dict(res)
 
-    def _build_full_prompt(
-        self, instruction, input, response
-    ):  # pylint: disable=redefined-builtin
+    def _build_full_prompt(self, instruction, input, response):
         return next(iter(self.prompter.build_prompt(instruction, input, response)))
 
 
@@ -72,8 +71,8 @@ class CompletionPrompter:
     def build_prompt(
         self,
         instruction: str,
-        input=None,  # pylint: disable=redefined-builtin, unused-argument
-        output=None,  # pylint: disable=unused-argument
+        input=None,
+        output=None,
     ) -> Generator[str, None, None]:
         yield instruction
 
diff --git a/src/axolotl/prompt_strategies/context_qa.py b/src/axolotl/prompt_strategies/context_qa.py
index f87dd8b5cd..09e96d26ef 100644
--- a/src/axolotl/prompt_strategies/context_qa.py
+++ b/src/axolotl/prompt_strategies/context_qa.py
@@ -1,4 +1,5 @@
 """Module containing the classes for Context QA Prompt Tokenization Strategies"""
+
 from typing import Tuple
 
 from axolotl.prompt_tokenizers import InstructionPromptTokenizingStrategy
@@ -85,7 +86,6 @@ class ContextV2Prompter(AlpacaPrompter):
     system_no_input_prompt = ""
 
     def match_prompt_style(self):
-        # pylint: disable=duplicate-code
         self.turn_format = "{instruction}\n{input}"
         self.turn_no_input_format = "{instruction}"
         self.system_format = "{system}"
diff --git a/src/axolotl/prompt_strategies/creative_acr.py b/src/axolotl/prompt_strategies/creative_acr.py
index ea67034b3b..3e016e30ee 100644
--- a/src/axolotl/prompt_strategies/creative_acr.py
+++ b/src/axolotl/prompt_strategies/creative_acr.py
@@ -134,9 +134,7 @@ class CreativePrompterBase:
     def build_prompt(
         self,
         instruction: str,
-        input: Union[  # pylint: disable=redefined-builtin, unused-argument
-            None, str
-        ] = None,
+        input: Union[None, str] = None,
         output: Union[None, str] = None,
     ) -> Generator[str, None, None]:
         if self.system_prompt:
diff --git a/src/axolotl/prompt_strategies/dpo/__init__.py b/src/axolotl/prompt_strategies/dpo/__init__.py
index 7f5e6eb644..f671256823 100644
--- a/src/axolotl/prompt_strategies/dpo/__init__.py
+++ b/src/axolotl/prompt_strategies/dpo/__init__.py
@@ -1,6 +1,7 @@
 """
 module for DPO style dataset transform strategies
 """
+
 from functools import partial
 
 from ..base import load as load_base
diff --git a/src/axolotl/prompt_strategies/dpo/chat_template.py b/src/axolotl/prompt_strategies/dpo/chat_template.py
new file mode 100644
index 0000000000..83db96750b
--- /dev/null
+++ b/src/axolotl/prompt_strategies/dpo/chat_template.py
@@ -0,0 +1,244 @@
+"""
+DPO prompt strategies for using tokenizer chat templates.
+"""
+
+from axolotl.utils.chat_templates import extract_chat_template_args, get_chat_template
+from axolotl.utils.schemas.utils import handle_legacy_message_fields_logic
+
+
+def default(cfg, dataset_idx=0, **kwargs):
+    ds_cfg = cfg["datasets"][dataset_idx]
+    ds_cfg = handle_legacy_message_fields_logic(ds_cfg)
+
+    chat_template_choice, chat_template_jinja = extract_chat_template_args(
+        cfg=cfg, ds_cfg=ds_cfg
+    )
+    field_messages = ds_cfg.get("field_messages", "messages")
+    field_chosen = ds_cfg.get("field_chosen", "chosen")
+    field_rejected = ds_cfg.get("field_rejected", "rejected")
+    message_property_mappings = ds_cfg.get(
+        "message_property_mappings",
+        {
+            "role": "role",
+            "content": "content",
+        },
+    )
+    role_map_inv = ds_cfg.get(
+        "roles",
+        {
+            "user": ["user"],
+            "assistant": ["assistant"],
+            "system": ["system"],
+            "tool": ["tool"],
+        },
+    )
+    role_map = {}
+    for target, sources in role_map_inv.items():
+        for source in sources:
+            role_map[source] = target
+
+    def transform_fn(sample, tokenizer=None):
+        chat_template_string = get_chat_template(
+            user_choice=chat_template_choice,
+            jinja_template=chat_template_jinja,
+            tokenizer=tokenizer,
+        )
+
+        messages = sample[field_messages]
+        if isinstance(messages, str):
+            messages = [
+                {
+                    message_property_mappings["role"]: "user",
+                    message_property_mappings["content"]: messages,
+                }
+            ]
+
+        messages = [
+            {
+                "role": role_map[m[message_property_mappings["role"]]],
+                "content": m[message_property_mappings["content"]],
+            }
+            for m in messages
+        ]
+
+        chosen_raw = sample[field_chosen]
+        if isinstance(chosen_raw, str):
+            chosen_msg = {
+                message_property_mappings["role"]: "assistant",
+                message_property_mappings["content"]: chosen_raw,
+            }
+        elif isinstance(chosen_raw, dict):
+            chosen_msg = chosen_raw
+        else:
+            chosen_msg = chosen_raw[-1]
+        chosen = {
+            "role": role_map[chosen_msg[message_property_mappings["role"]]],
+            "content": chosen_msg[message_property_mappings["content"]],
+        }
+
+        rejected_raw = sample[field_rejected]
+        if isinstance(rejected_raw, str):
+            rejected_msg = {
+                message_property_mappings["role"]: "assistant",
+                message_property_mappings["content"]: rejected_raw,
+            }
+        elif isinstance(rejected_raw, dict):
+            rejected_msg = rejected_raw
+        else:
+            rejected_msg = rejected_raw[-1]
+        rejected = {
+            "role": role_map[rejected_msg[message_property_mappings["role"]]],
+            "content": rejected_msg[message_property_mappings["content"]],
+        }
+        dummy_user_message = {"role": "user", "content": "[[dummy_message]]"}
+
+        result = {}
+        result["prompt"] = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            chat_template=chat_template_string,
+            tokenize=False,
+        )
+
+        result["chosen"] = tokenizer.apply_chat_template(
+            [dummy_user_message, chosen],
+            add_generation_prompt=False,
+            chat_template=chat_template_string,
+            tokenize=False,
+        )
+        chosen_strip_index = result["chosen"].find(chosen["content"])
+        result["chosen"] = result["chosen"][chosen_strip_index:].rstrip()
+
+        result["rejected"] = tokenizer.apply_chat_template(
+            [dummy_user_message, rejected],
+            add_generation_prompt=False,
+            chat_template=chat_template_string,
+            tokenize=False,
+        )
+        rejected_strip_index = result["rejected"].find(rejected["content"])
+        result["rejected"] = result["rejected"][rejected_strip_index:].rstrip()
+
+        return result
+
+    return transform_fn, {"remove_columns": [field_messages]}
+
+
+def argilla_chat(cfg, dataset_idx=0, **kwargs):
+    """
+    DPO chat template strategy for argilla-style datasets.
+
+    For argilla-style datasets where chosen/rejected contain full conversations
+    instead of single response messages. Extracts the conversation history from
+    the chosen field and formats both chosen/rejected responses using the
+    configured chat template.
+
+    Args:
+        cfg: Configuration object containing chat_template and dataset settings
+        dataset_idx: Index of the dataset in the config (default: 0)
+        **kwargs: Additional keyword arguments (unused)
+
+    Returns:
+        tuple: (transform_fn, dataset_kwargs) where:
+            - transform_fn: Function to transform dataset samples
+            - dataset_kwargs: Dict with 'remove_columns' specifying columns to drop
+
+    Dataset format:
+        {
+            "chosen": [
+                {"role": "user", "content": "..."},
+                {"role": "assistant", "content": "..."}
+            ],
+            "rejected": [
+                {"role": "user", "content": "..."},
+                {"role": "assistant", "content": "..."}
+            ]
+        }
+    """
+    ds_cfg = cfg["datasets"][dataset_idx]
+    ds_cfg = handle_legacy_message_fields_logic(ds_cfg)
+
+    chat_template_choice, chat_template_jinja = extract_chat_template_args(
+        cfg=cfg, ds_cfg=ds_cfg
+    )
+    field_chosen = ds_cfg.get("field_chosen", "chosen")
+    field_rejected = ds_cfg.get("field_rejected", "rejected")
+    message_property_mappings = ds_cfg.get(
+        "message_property_mappings",
+        {
+            "role": "role",
+            "content": "content",
+        },
+    )
+    role_map_inv = ds_cfg.get(
+        "roles",
+        {
+            "user": ["user"],
+            "assistant": ["assistant"],
+            "system": ["system"],
+            "tool": ["tool"],
+        },
+    )
+    role_map = {}
+    for target, sources in role_map_inv.items():
+        for source in sources:
+            role_map[source] = target
+
+    def transform_fn(sample, tokenizer=None):
+        chat_template_string = get_chat_template(
+            user_choice=chat_template_choice,
+            jinja_template=chat_template_jinja,
+            tokenizer=tokenizer,
+        )
+
+        chosen_raw = sample[field_chosen]
+        rejected_raw = sample[field_rejected]
+
+        # Extract messages (all but last) and responses (last message)
+        chosen_messages = [
+            {
+                "role": role_map[m[message_property_mappings["role"]]],
+                "content": m[message_property_mappings["content"]],
+            }
+            for m in chosen_raw[:-1]
+        ]
+        chosen_response = {
+            "role": role_map[chosen_raw[-1][message_property_mappings["role"]]],
+            "content": chosen_raw[-1][message_property_mappings["content"]],
+        }
+
+        rejected_response = {
+            "role": role_map[rejected_raw[-1][message_property_mappings["role"]]],
+            "content": rejected_raw[-1][message_property_mappings["content"]],
+        }
+
+        dummy_user_message = {"role": "user", "content": "[[dummy_message]]"}
+
+        result = {}
+        result["prompt"] = tokenizer.apply_chat_template(
+            chosen_messages,
+            add_generation_prompt=True,
+            chat_template=chat_template_string,
+            tokenize=False,
+        )
+
+        result["chosen"] = tokenizer.apply_chat_template(
+            [dummy_user_message, chosen_response],
+            add_generation_prompt=False,
+            chat_template=chat_template_string,
+            tokenize=False,
+        )
+        chosen_strip_index = result["chosen"].find(chosen_response["content"])
+        result["chosen"] = result["chosen"][chosen_strip_index:].rstrip()
+
+        result["rejected"] = tokenizer.apply_chat_template(
+            [dummy_user_message, rejected_response],
+            add_generation_prompt=False,
+            chat_template=chat_template_string,
+            tokenize=False,
+        )
+        rejected_strip_index = result["rejected"].find(rejected_response["content"])
+        result["rejected"] = result["rejected"][rejected_strip_index:].rstrip()
+
+        return result
+
+    return transform_fn, {"remove_columns": [field_chosen, field_rejected]}
diff --git a/src/axolotl/prompt_strategies/dpo/chatml.py b/src/axolotl/prompt_strategies/dpo/chatml.py
index 585696e29a..8614708eb5 100644
--- a/src/axolotl/prompt_strategies/dpo/chatml.py
+++ b/src/axolotl/prompt_strategies/dpo/chatml.py
@@ -3,22 +3,41 @@
 """
 
 
-def argilla(
+def default(
     cfg,
     **kwargs,
-):  # pylint: disable=possibly-unused-variable,unused-argument
+):
     def transform_fn(sample):
+        if "prompt" in sample.keys():
+            prompt_key = "prompt"
+        elif "input" in sample.keys():
+            prompt_key = "input"
+        elif "question" in sample.keys():
+            prompt_key = "question"
+        else:
+            prompt_key = "instruction"
+
+        if "chosen" in sample.keys():
+            chosen_key = "chosen"
+        else:
+            chosen_key = "chosen_response"
+
+        if "rejected" in sample.keys():
+            rejected_key = "rejected"
+        else:
+            rejected_key = "rejected_response"
+
         if "system" in sample and sample["system"]:
             sample["prompt"] = (
                 f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
-                f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
+                f"<|im_start|>user\n{sample[prompt_key]}<|im_end|>\n<|im_start|>assistant\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
-        sample["chosen"] = f"{sample['chosen_response']}<|im_end|>"
-        sample["rejected"] = f"{sample['rejected_response']}<|im_end|>"
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample[prompt_key]}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        sample["chosen"] = f"{sample[chosen_key]}<|im_end|>"
+        sample["rejected"] = f"{sample[rejected_key]}<|im_end|>"
         return sample
 
     return transform_fn
@@ -27,15 +46,15 @@ def transform_fn(sample):
 def argilla_chat(
     cfg,
     **kwargs,
-):  # pylint: disable=possibly-unused-variable,unused-argument
+):
     """
     for argilla/dpo-mix-7k conversations
     """
 
     def transform_fn(sample):
-        sample[
-            "prompt"
-        ] = f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
+        sample["prompt"] = (
+            f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
+        )
         sample["chosen"] = f"{sample['chosen'][1]['content']}<|im_end|>"
         sample["rejected"] = f"{sample['rejected'][1]['content']}<|im_end|>"
         return sample
@@ -46,7 +65,7 @@ def transform_fn(sample):
 def icr(
     cfg,
     **kwargs,
-):  # pylint: disable=possibly-unused-variable,unused-argument
+):
     """
     chatml transforms for datasets with system, input, chosen, rejected
     ex. https://huggingface.co/datasets/argilla/distilabel-intel-orca-dpo-pairs
@@ -59,9 +78,9 @@ def transform_fn(sample):
                 f"<|im_start|>user\n{sample['input']}<|im_end|>\n<|im_start|>assistant\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|im_start|>user\n{sample['input']}<|im_end|>\n<|im_start|>assistant\n"
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample['input']}<|im_end|>\n<|im_start|>assistant\n"
+            )
         sample["chosen"] = f"{sample['chosen']}<|im_end|>"
         sample["rejected"] = f"{sample['rejected']}<|im_end|>"
         return sample
@@ -69,7 +88,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+def intel(cfg, **kwargs):
     """
     For Intel Orca DPO Pairs
     """
@@ -81,9 +100,9 @@ def transform_fn(sample):
                 f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
+            )
         sample["chosen"] = f"{sample['chosen']}<|im_end|>"
         sample["rejected"] = f"{sample['rejected']}<|im_end|>"
         return sample
@@ -91,9 +110,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def prompt_pairs(
-    cfg, **kwargs
-):  # pylint: disable=possibly-unused-variable,unused-argument
+def prompt_pairs(cfg, **kwargs):
     def transform_fn(sample):
         if "system" in sample and sample["system"]:
             sample["prompt"] = (
@@ -101,9 +118,9 @@ def transform_fn(sample):
                 f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            )
         sample["chosen"] = f"{sample['chosen']}<|im_end|>"
         sample["rejected"] = f"{sample['rejected']}<|im_end|>"
         return sample
@@ -111,7 +128,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+def ultra(cfg, **kwargs):
     """
     for ultrafeedback binarized conversations
     """
@@ -123,9 +140,9 @@ def transform_fn(sample):
                 f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            )
         sample["chosen"] = f"{sample['chosen'][1]['content']}<|im_end|>"
         sample["rejected"] = f"{sample['rejected'][1]['content']}<|im_end|>"
         return sample
diff --git a/src/axolotl/prompt_strategies/dpo/llama3.py b/src/axolotl/prompt_strategies/dpo/llama3.py
index cb394cc228..c13ff55e4d 100644
--- a/src/axolotl/prompt_strategies/dpo/llama3.py
+++ b/src/axolotl/prompt_strategies/dpo/llama3.py
@@ -3,22 +3,41 @@
 """
 
 
-def argilla(
+def default(
     cfg,
     **kwargs,
-):  # pylint: disable=possibly-unused-variable,unused-argument
+):
     def transform_fn(sample):
+        if "prompt" in sample.keys():
+            prompt_key = "prompt"
+        elif "input" in sample.keys():
+            prompt_key = "input"
+        elif "question" in sample.keys():
+            prompt_key = "question"
+        else:
+            prompt_key = "instruction"
+
+        if "chosen" in sample.keys():
+            chosen_key = "chosen"
+        else:
+            chosen_key = "chosen_response"
+
+        if "rejected" in sample.keys():
+            rejected_key = "rejected"
+        else:
+            rejected_key = "rejected_response"
+
         if "system" in sample and sample["system"]:
             sample["prompt"] = (
                 f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample[prompt_key]}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-        sample["chosen"] = f"{sample['chosen_response']}<|eot_id|>"
-        sample["rejected"] = f"{sample['rejected_response']}<|eot_id|>"
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample[prompt_key]}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        sample["chosen"] = f"{sample[chosen_key]}<|eot_id|>"
+        sample["rejected"] = f"{sample[rejected_key]}<|eot_id|>"
         return sample
 
     return transform_fn
@@ -27,15 +46,15 @@ def transform_fn(sample):
 def argilla_chat(
     cfg,
     **kwargs,
-):  # pylint: disable=possibly-unused-variable,unused-argument
+):
     """
     for argilla/dpo-mix-7k conversations
     """
 
     def transform_fn(sample):
-        sample[
-            "prompt"
-        ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['chosen'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["prompt"] = (
+            f"<|start_header_id|>user<|end_header_id|>\n\n{sample['chosen'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        )
         sample["chosen"] = f"{sample['chosen'][1]['content']}<|eot_id|>"
         sample["rejected"] = f"{sample['rejected'][1]['content']}<|eot_id|>"
         return sample
@@ -46,7 +65,7 @@ def transform_fn(sample):
 def icr(
     cfg,
     **kwargs,
-):  # pylint: disable=possibly-unused-variable,unused-argument
+):
     """
     chatml transforms for datasets with system, input, chosen, rejected
     ex. https://huggingface.co/datasets/argilla/distilabel-intel-orca-dpo-pairs
@@ -59,9 +78,9 @@ def transform_fn(sample):
                 f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
         sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
         sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
         return sample
@@ -69,7 +88,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+def intel(cfg, **kwargs):
     """
     For Intel Orca DPO Pairs
     """
@@ -81,9 +100,9 @@ def transform_fn(sample):
                 f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
         sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
         sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
         return sample
@@ -91,9 +110,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def prompt_pairs(
-    cfg, **kwargs
-):  # pylint: disable=possibly-unused-variable,unused-argument
+def prompt_pairs(cfg, **kwargs):
     def transform_fn(sample):
         if "system" in sample and sample["system"]:
             sample["prompt"] = (
@@ -101,9 +118,9 @@ def transform_fn(sample):
                 f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
         sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
         sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
         return sample
@@ -111,7 +128,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+def ultra(cfg, **kwargs):
     """
     for ultrafeedback binarized conversations
     """
@@ -123,9 +140,9 @@ def transform_fn(sample):
                 f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
             )
         else:
-            sample[
-                "prompt"
-            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
         sample["chosen"] = f"{sample['chosen'][1]['content']}<|eot_id|>"
         sample["rejected"] = f"{sample['rejected'][1]['content']}<|eot_id|>"
         return sample
diff --git a/src/axolotl/prompt_strategies/dpo/passthrough.py b/src/axolotl/prompt_strategies/dpo/passthrough.py
new file mode 100644
index 0000000000..52b5ceac19
--- /dev/null
+++ b/src/axolotl/prompt_strategies/dpo/passthrough.py
@@ -0,0 +1,10 @@
+"""
+DPO prompt strategies passthrough/zero-processing strategy
+"""
+
+
+def default(cfg, dataset_idx=0, **kwargs):
+    def transform_fn(sample, tokenizer=None):
+        return sample
+
+    return transform_fn
diff --git a/src/axolotl/prompt_strategies/dpo/user_defined.py b/src/axolotl/prompt_strategies/dpo/user_defined.py
index 1d5f891af6..c709d516be 100644
--- a/src/axolotl/prompt_strategies/dpo/user_defined.py
+++ b/src/axolotl/prompt_strategies/dpo/user_defined.py
@@ -3,7 +3,7 @@
 """
 
 
-def default(cfg, dataset_idx=0, **kwargs):  # pylint: disable=unused-argument
+def default(cfg, dataset_idx=0, **kwargs):
     ds_cfg = cfg["datasets"][dataset_idx]["type"]
     if not isinstance(ds_cfg, dict):
         raise ValueError(
@@ -15,25 +15,22 @@ def default(cfg, dataset_idx=0, **kwargs):  # pylint: disable=unused-argument
     field_rejected = ds_cfg.get("field_rejected", "rejected")
     prompt_format = ds_cfg.get("prompt_format")
     if not prompt_format:
-        prompt_format = "{" + field_prompt + "}"
+        prompt_format = "{prompt}"
     chosen_format = ds_cfg.get("chosen_format")
     if not chosen_format:
-        chosen_format = "{" + field_chosen + "}"
+        chosen_format = "{chosen}"
     rejected_format = ds_cfg.get("rejected_format")
     if not rejected_format:
-        rejected_format = "{" + field_rejected + "}"
+        rejected_format = "{rejected}"
 
     def transform_fn(sample):
-        if (
-            "{" + field_system + "}" in prompt_format
-            and field_system in sample
-            and sample[field_system]
-        ):
+        if "{system}" in prompt_format:
             sample["prompt"] = prompt_format.format(
-                system=sample[field_system], prompt=sample[field_prompt]
+                system=sample.get(field_system) or "",
+                prompt=sample[field_prompt],
             )
         else:
-            sample["prompt"] = prompt_format.format(prompt=sample["prompt"])
+            sample["prompt"] = prompt_format.format(prompt=sample[field_prompt])
         sample["chosen"] = chosen_format.format(chosen=sample[field_chosen])
         sample["rejected"] = rejected_format.format(rejected=sample[field_rejected])
         return sample
diff --git a/src/axolotl/prompt_strategies/dpo/zephyr.py b/src/axolotl/prompt_strategies/dpo/zephyr.py
index 9eb8950091..781227181d 100644
--- a/src/axolotl/prompt_strategies/dpo/zephyr.py
+++ b/src/axolotl/prompt_strategies/dpo/zephyr.py
@@ -3,14 +3,11 @@
 """
 
 
-def nectar(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+def nectar(cfg, **kwargs):
     def transform_fn(sample):
         data = {}
         data["prompt"] = (
-            "<|system|>\n</s>\n"
-            "<|user|>\n"
-            f"{sample['prompt']}</s>\n"
-            "<|assistant|>\n"
+            f"<|system|>\n</s>\n<|user|>\n{sample['prompt']}</s>\n<|assistant|>\n"
         )
         answers = sorted(sample["answers"], key=lambda x: x["rank"])
         data["chosen"] = answers[-1]["answer"]
diff --git a/src/axolotl/prompt_strategies/ebft/__init__.py b/src/axolotl/prompt_strategies/ebft/__init__.py
new file mode 100644
index 0000000000..db46af0057
--- /dev/null
+++ b/src/axolotl/prompt_strategies/ebft/__init__.py
@@ -0,0 +1,9 @@
+"""
+module for EBFT style dataset transform strategies
+"""
+
+from functools import partial
+
+from ..base import load as load_base
+
+load = partial(load_base, module_base="axolotl.prompt_strategies.ebft")
diff --git a/src/axolotl/prompt_strategies/ebft/ebft_chat_multiturn.py b/src/axolotl/prompt_strategies/ebft/ebft_chat_multiturn.py
new file mode 100644
index 0000000000..26982c1833
--- /dev/null
+++ b/src/axolotl/prompt_strategies/ebft/ebft_chat_multiturn.py
@@ -0,0 +1,129 @@
+"""
+Dataset transform for multi-turn chat data with structured EBFT (vLLM mode).
+
+Three variants:
+
+1. `transform` — Uses the FIRST assistant turn as the generation target.
+   Passes remaining turns as `remaining_turns` for sequential rollout.
+   The trainer generates turn 1 via GRPO/vLLM, then sequentially generates
+   subsequent assistant turns, comparing the full conversation to GT.
+
+2. `transform_last_turn` — Uses the LAST assistant turn as the target.
+   Simplest approach: the full conversation history is the prompt.
+
+3. `transform_all_turns` — Explodes each conversation into N examples
+   (one per assistant turn). Each turn is an independent training example.
+   Use with batched=True.
+
+Supports OpenAI chat format:
+  {"messages": [{"role": ..., "content": ...}, ...]}
+"""
+
+
+def transform(cfg, **kwargs):
+    """Multi-turn with sequential rollout.
+
+    Returns the first assistant turn as ground_truth, plus remaining_turns
+    for the trainer to do sequential rollout generation.
+    """
+
+    def transform_fn(example, tokenizer=None):
+        messages = example.get("messages", example.get("conversations", []))
+
+        if not messages:
+            return {"prompt": [], "ground_truth": ""}
+
+        # Split at first assistant turn
+        prompt_msgs = []
+        first_gt = None
+        remaining = []
+
+        found_first = False
+        for msg in messages:
+            if msg["role"] == "assistant" and not found_first:
+                first_gt = msg["content"]
+                found_first = True
+            elif found_first:
+                remaining.append(msg)
+            else:
+                prompt_msgs.append(msg)
+
+        if first_gt is None:
+            return {"prompt": prompt_msgs, "ground_truth": ""}
+
+        # Store only the first assistant turn as ground_truth. The full multi-turn
+        # GT is reconstructed in the reward function via chat template rendering
+        # (using remaining_turns), which preserves role markers between turns.
+        return {
+            "prompt": prompt_msgs,
+            "ground_truth": first_gt,
+            "remaining_turns": remaining,
+        }
+
+    return transform_fn, {
+        "remove_columns": "__all__",
+    }
+
+
+def transform_last_turn(cfg, **kwargs):
+    """Single-turn: use the last assistant turn as the generation target."""
+
+    def transform_fn(example, tokenizer=None):
+        messages = example.get("messages", example.get("conversations", []))
+
+        if not messages:
+            return {"prompt": [], "ground_truth": ""}
+
+        # Find all assistant turns
+        history = []
+        last_prompt = []
+        last_gt = ""
+        for msg in messages:
+            if msg["role"] == "assistant":
+                last_prompt = list(history)
+                last_gt = msg["content"]
+            history.append(msg)
+
+        return {
+            "prompt": last_prompt,
+            "ground_truth": last_gt,
+        }
+
+    return transform_fn, {
+        "remove_columns": "__all__",
+    }
+
+
+def transform_all_turns(cfg, **kwargs):
+    """Explode: one example per assistant turn.
+
+    Use with datasets.map(batched=True) to produce N examples from
+    each N-turn conversation.
+
+    Usage in YAML:
+        type: ebft_chat_multiturn.transform_all_turns
+    """
+
+    def transform_fn(examples, tokenizer=None):
+        all_prompts = []
+        all_ground_truths = []
+
+        messages_list = examples.get("messages", examples.get("conversations", []))
+
+        for messages in messages_list:
+            history = []
+            for msg in messages:
+                if msg["role"] == "assistant":
+                    all_prompts.append(list(history))
+                    all_ground_truths.append(msg["content"])
+                history.append(msg)
+
+        return {
+            "prompt": all_prompts,
+            "ground_truth": all_ground_truths,
+        }
+
+    return transform_fn, {
+        "remove_columns": "__all__",
+        "batched": True,
+    }
diff --git a/src/axolotl/prompt_strategies/ebft/ebft_opencode.py b/src/axolotl/prompt_strategies/ebft/ebft_opencode.py
new file mode 100644
index 0000000000..930d314a17
--- /dev/null
+++ b/src/axolotl/prompt_strategies/ebft/ebft_opencode.py
@@ -0,0 +1,20 @@
+"""
+Dataset transform for nvidia/OpenCodeInstruct with EBFT structured mode.
+
+Maps the dataset's `input` (prompt) and `output` (code solution) fields
+to the format expected by the EBFT trainer (prompt + ground_truth).
+"""
+
+
+def transform(cfg, **kwargs):
+    def transform_fn(example, tokenizer=None):
+        return {
+            "prompt": [
+                {"role": "user", "content": example["input"]},
+            ],
+            "ground_truth": example["output"],
+        }
+
+    return transform_fn, {
+        "remove_columns": "__all__",
+    }
diff --git a/src/axolotl/prompt_strategies/ebft/ebft_reasoning.py b/src/axolotl/prompt_strategies/ebft/ebft_reasoning.py
new file mode 100644
index 0000000000..7c756d1583
--- /dev/null
+++ b/src/axolotl/prompt_strategies/ebft/ebft_reasoning.py
@@ -0,0 +1,319 @@
+"""
+Dataset transform for reasoning/thinking datasets with EBFT.
+
+Handles datasets where assistant responses contain <think>...</think> reasoning
+traces (e.g., TeichAI/Claude-Opus-4.6-Reasoning, Qwen3.5 thinking mode outputs).
+
+Two variants:
+
+1. `transform` — For structured EBFT (vLLM mode):
+   Returns prompt + ground_truth with thinking tags preserved.
+   Feature matching compares full responses (thinking + answer).
+
+2. `transform_answer_only` — For structured EBFT (vLLM mode):
+   Strips <think>...</think> from ground_truth, so feature matching
+   only scores the final answer portion. Use when reasoning chains
+   can vary but the answer should match.
+
+3. `transform_strided` — For strided EBFT:
+   Tokenizes the full conversation with thinking traces.
+   Optionally masks thinking tokens from CE loss (labels=-100 for think spans)
+   while still placing anchors in thinking regions for feature matching.
+
+All variants work with OpenAI chat format:
+  {"messages": [{"role": "...", "content": "<think>...</think>Answer"}]}
+"""
+
+import re
+
+
+def _strip_thinking(text: str) -> str:
+    """Remove <think>...</think> blocks from text."""
+    return re.sub(r"<think>.*?</think>\s*", "", text, flags=re.DOTALL).strip()
+
+
+def _extract_thinking(text: str) -> tuple[str, str]:
+    """Split text into (thinking, answer) parts."""
+    match = re.search(r"<think>(.*?)</think>\s*(.*)", text, flags=re.DOTALL)
+    if match:
+        return match.group(1).strip(), match.group(2).strip()
+    return "", text.strip()
+
+
+def transform(cfg, **kwargs):
+    """Full response including thinking traces for feature matching.
+
+    For datasets where assistant content has <think>...</think> tags in the
+    content field. The ground_truth includes the full content (thinking + answer).
+    """
+
+    def transform_fn(example, tokenizer=None):
+        messages = example.get("messages", example.get("conversations", []))
+
+        prompt_msgs_snapshot = None
+        ground_truth = ""
+        for msg_idx, msg in enumerate(messages):
+            if msg["role"] == "assistant":
+                prompt_msgs_snapshot = list(messages[:msg_idx])
+                ground_truth = msg["content"]
+
+        return {
+            "prompt": prompt_msgs_snapshot
+            if prompt_msgs_snapshot is not None
+            else messages[:-1],
+            "ground_truth": ground_truth,
+        }
+
+    return transform_fn, {"remove_columns": "__all__"}
+
+
+def transform_split_thinking(cfg, **kwargs):
+    """Split <think> tags into reasoning_content field for native chat template handling.
+
+    For datasets where thinking is embedded in the content field as <think>...</think>.
+    Splits it into separate reasoning_content and content fields so the model's
+    chat template can format it natively (e.g., Qwen3.5's reasoning_content support).
+
+    The prompt messages are passed through with reasoning_content properly split,
+    so vLLM generation with enable_thinking=true produces comparable outputs.
+    The ground_truth is the full assistant response (thinking + answer) for
+    feature matching.
+
+    Also works for:
+    - <reasoning>...</reasoning> tags
+    - <|begin_of_thought|>...<|end_of_thought|> tags
+    """
+    _THINKING_PAIRS = [
+        ("<think>", "</think>"),
+        ("<reasoning>", "</reasoning>"),
+        ("<|begin_of_thought|>", "<|end_of_thought|>"),
+    ]
+
+    def _split_msg_thinking(msg):
+        """Split thinking from assistant message content into reasoning_content.
+
+        Always includes reasoning_content key on assistant messages (empty string
+        if no thinking tags found) to ensure consistent HF dataset schema across
+        all examples in a batch.
+        """
+        if msg["role"] != "assistant":
+            return msg
+        content = msg.get("content", "")
+        # Already has reasoning_content — pass through
+        if "reasoning_content" in msg:
+            return msg
+        for open_tag, close_tag in _THINKING_PAIRS:
+            if open_tag in content and close_tag in content:
+                start = content.find(open_tag)
+                end = content.find(close_tag)
+                thinking = content[start + len(open_tag) : end].strip()
+                answer = content[end + len(close_tag) :].strip()
+                return {
+                    **msg,
+                    "reasoning_content": thinking,
+                    "content": answer,
+                }
+        # No thinking tags — still add reasoning_content for schema consistency
+        return {**msg, "reasoning_content": ""}
+
+    def _normalize_msg(msg):
+        """Ensure every message has {role, content, reasoning_content} for HF schema consistency."""
+        return {
+            "role": msg.get("role", ""),
+            "content": msg.get("content", ""),
+            "reasoning_content": msg.get("reasoning_content", ""),
+        }
+
+    def transform_fn(example, tokenizer=None):
+        messages = example.get("messages", example.get("conversations", []))
+
+        # Split thinking in all assistant messages, then normalize schema
+        split_messages = [_normalize_msg(_split_msg_thinking(m)) for m in messages]
+
+        # Build prompt (all messages except last assistant) and ground_truth
+        prompt_msgs = []
+        prompt_msgs_snapshot = None
+        ground_truth = ""
+        for msg in split_messages:
+            if msg["role"] == "assistant":
+                prompt_msgs_snapshot = list(prompt_msgs)
+                # ground_truth is the FULL content for feature matching
+                thinking = msg.get("reasoning_content", "")
+                answer = msg.get("content", "")
+                if thinking:
+                    ground_truth = f"<think>\n{thinking}\n</think>\n\n{answer}"
+                else:
+                    ground_truth = answer
+            prompt_msgs.append(msg)
+
+        return {
+            "prompt": prompt_msgs_snapshot
+            if prompt_msgs_snapshot is not None
+            else split_messages[:-1],
+            "ground_truth": ground_truth,
+        }
+
+    return transform_fn, {"remove_columns": "__all__"}
+
+
+def transform_answer_only(cfg, **kwargs):
+    """Strip thinking from ground_truth — match features on answer only."""
+
+    def transform_fn(example, tokenizer=None):
+        messages = example.get("messages", example.get("conversations", []))
+
+        prompt_msgs = []
+        prompt_msgs_snapshot = None
+        ground_truth = ""
+        for msg in messages:
+            if msg["role"] == "assistant":
+                prompt_msgs_snapshot = list(prompt_msgs)
+                ground_truth = _strip_thinking(msg["content"])
+            prompt_msgs.append(msg)
+
+        return {
+            "prompt": prompt_msgs_snapshot
+            if prompt_msgs_snapshot is not None
+            else messages[:-1],
+            "ground_truth": ground_truth,
+        }
+
+    return transform_fn, {"remove_columns": "__all__"}
+
+
+def transform_strided(cfg, **kwargs):
+    """For strided EBFT: tokenize with thinking, optionally mask think tokens from CE loss.
+
+    Config options (via cfg):
+        - ebft.mask_thinking_ce: bool (default False)
+          If True, set labels=-100 for tokens inside <think>...</think> blocks.
+          Feature matching still uses these positions (anchors are placed everywhere
+          in the completion span). Only CE auxiliary loss is affected.
+    """
+    seq_len = cfg.sequence_len
+    mask_thinking = False
+    if cfg.ebft and hasattr(cfg.ebft, "mask_thinking_ce"):
+        mask_thinking = cfg.ebft.mask_thinking_ce
+
+    def transform_fn(example, tokenizer=None):
+        messages = example.get("messages", example.get("conversations", []))
+
+        if tokenizer is None:
+            for m in messages:
+                if m.get("role") == "user":
+                    return {"prompt": m["content"]}
+            return {"prompt": str(messages)}
+
+        pad_id = (
+            tokenizer.pad_token_id
+            if tokenizer.pad_token_id is not None
+            else tokenizer.eos_token_id
+        )
+
+        # Tokenize the full conversation with the chat template
+        full_text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+        full_enc = tokenizer(
+            full_text,
+            truncation=True,
+            max_length=seq_len,
+            add_special_tokens=False,
+            return_tensors=None,
+        )
+        input_ids = full_enc["input_ids"]
+
+        # Build labels: -100 for non-assistant tokens
+        labels = [-100] * len(input_ids)
+
+        # Find assistant turn boundaries using incremental tokenization.
+        # Only the FINAL assistant turn is marked as trainable.
+        prefix_messages = []
+        final_start = None
+        final_end = None
+        for msg in messages:
+            if msg["role"] == "assistant":
+                prefix_text = tokenizer.apply_chat_template(
+                    prefix_messages,
+                    tokenize=False,
+                    add_generation_prompt=True,
+                )
+                prefix_ids = tokenizer(
+                    prefix_text,
+                    truncation=True,
+                    max_length=seq_len,
+                    add_special_tokens=False,
+                    return_tensors=None,
+                )["input_ids"]
+                start = len(prefix_ids)
+
+                prefix_messages.append(msg)
+                with_turn_text = tokenizer.apply_chat_template(
+                    prefix_messages,
+                    tokenize=False,
+                    add_generation_prompt=False,
+                )
+                with_turn_ids = tokenizer(
+                    with_turn_text,
+                    truncation=True,
+                    max_length=seq_len,
+                    add_special_tokens=False,
+                    return_tensors=None,
+                )["input_ids"]
+                end = len(with_turn_ids)
+
+                # Record this turn's boundaries; only the last one will be used
+                final_start = start
+                final_end = end
+            else:
+                prefix_messages.append(msg)
+
+        # Mark only the final assistant turn as trainable
+        if final_start is not None and final_end is not None:
+            for i in range(final_start, min(final_end, len(labels))):
+                labels[i] = input_ids[i]
+
+            # Optionally mask <think>...</think> tokens within this turn.
+            # Find think spans by scanning for <think> and </think> token IDs
+            # directly in the input_ids (robust to tokenization alignment).
+            if mask_thinking:
+                think_open_id = tokenizer.convert_tokens_to_ids("<think>")
+                think_close_id = tokenizer.convert_tokens_to_ids("</think>")
+                if think_open_id != tokenizer.unk_token_id:
+                    # Scan from before the assistant turn start to catch
+                    # <think> tags that are part of the template prefix
+                    scan_start = max(0, final_start - 5)
+                    in_think = False
+                    for i in range(scan_start, min(final_end, len(labels))):
+                        if input_ids[i] == think_open_id:
+                            in_think = True
+                        if in_think and i >= final_start:
+                            labels[i] = -100
+                        if input_ids[i] == think_close_id:
+                            in_think = False
+                            if i >= final_start:
+                                labels[i] = -100
+
+        # Derive prompt_length
+        prompt_length = len(input_ids)
+        for i, lbl in enumerate(labels):
+            if lbl != -100:
+                prompt_length = i
+                break
+
+        # Pad
+        pad_len = seq_len - len(input_ids)
+        attention_mask = [1] * len(input_ids) + [0] * pad_len
+        labels = labels + [-100] * pad_len
+        input_ids = input_ids + [pad_id] * pad_len
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "prompt_length": prompt_length,
+        }
+
+    return transform_fn, {"remove_columns": "__all__"}
diff --git a/src/axolotl/prompt_strategies/ebft/ebft_strided_chat.py b/src/axolotl/prompt_strategies/ebft/ebft_strided_chat.py
new file mode 100644
index 0000000000..f1d0c01f9f
--- /dev/null
+++ b/src/axolotl/prompt_strategies/ebft/ebft_strided_chat.py
@@ -0,0 +1,110 @@
+"""
+Dataset transform for multi-turn chat data with strided EBFT.
+
+Tokenizes conversations using the model's chat template, producing input_ids
+with labels=-100 for system/user turns and real labels for assistant turns.
+The strided trainer places anchors only within assistant completion spans.
+
+Works with datasets in OpenAI chat format:
+  [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
+"""
+
+
+def transform(cfg, **kwargs):
+    seq_len = cfg.sequence_len
+
+    def transform_fn(example, tokenizer=None):
+        messages = example.get("messages", example.get("conversations", []))
+
+        if tokenizer is None:
+            # For preview: just return the first user message
+            for m in messages:
+                if m.get("role") == "user":
+                    return {"prompt": m["content"]}
+            return {"prompt": str(messages)}
+
+        pad_id = tokenizer.pad_token_id or tokenizer.eos_token_id
+
+        # Tokenize the full conversation with the chat template
+        full_text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+        full_enc = tokenizer(
+            full_text,
+            truncation=True,
+            max_length=seq_len,
+            add_special_tokens=False,
+            return_tensors=None,
+        )
+        input_ids = full_enc["input_ids"]
+
+        # Build labels: -100 for everything except assistant turns.
+        # Strategy: tokenize incrementally to find assistant turn boundaries.
+        labels = [-100] * len(input_ids)
+
+        # Tokenize prefix up to each assistant turn to find boundaries
+        prefix_messages = []
+        for msg in messages:
+            if msg["role"] == "assistant":
+                # Tokenize prefix (everything before this assistant turn + generation prompt)
+                prefix_text = tokenizer.apply_chat_template(
+                    prefix_messages,
+                    tokenize=False,
+                    add_generation_prompt=True,
+                )
+                prefix_ids = tokenizer(
+                    prefix_text,
+                    truncation=True,
+                    max_length=seq_len,
+                    add_special_tokens=False,
+                    return_tensors=None,
+                )["input_ids"]
+                start = len(prefix_ids)
+
+                # Tokenize prefix + this assistant turn
+                prefix_messages.append(msg)
+                with_turn_text = tokenizer.apply_chat_template(
+                    prefix_messages,
+                    tokenize=False,
+                    add_generation_prompt=False,
+                )
+                with_turn_ids = tokenizer(
+                    with_turn_text,
+                    truncation=True,
+                    max_length=seq_len,
+                    add_special_tokens=False,
+                    return_tensors=None,
+                )["input_ids"]
+                end = len(with_turn_ids)
+
+                # Mark assistant tokens as trainable
+                for i in range(start, min(end, len(labels))):
+                    labels[i] = input_ids[i]
+            else:
+                prefix_messages.append(msg)
+
+        # Derive prompt_length as the position of the first non-masked label
+        prompt_length = len(input_ids)  # default: all masked
+        for i, lbl in enumerate(labels):
+            if lbl != -100:
+                prompt_length = i
+                break
+
+        # Pad to seq_len
+        pad_len = seq_len - len(input_ids)
+        attention_mask = [1] * len(input_ids) + [0] * pad_len
+        labels = labels + [-100] * pad_len
+        input_ids = input_ids + [pad_id] * pad_len
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "prompt_length": prompt_length,
+        }
+
+    return transform_fn, {
+        "remove_columns": "__all__",
+    }
diff --git a/src/axolotl/prompt_strategies/ebft/ebft_strided_structured.py b/src/axolotl/prompt_strategies/ebft/ebft_strided_structured.py
new file mode 100644
index 0000000000..7676505753
--- /dev/null
+++ b/src/axolotl/prompt_strategies/ebft/ebft_strided_structured.py
@@ -0,0 +1,80 @@
+"""
+Dataset transform for structured (prompt, completion) data with strided EBFT.
+
+Tokenizes prompt and completion separately, concatenates into a single
+input_ids sequence, and marks prompt tokens with labels=-100 so the
+strided trainer knows where to place anchors (completion span only).
+
+Works with datasets that have chat-style fields (e.g., nvidia/OpenCodeInstruct).
+"""
+
+
+def transform(cfg, **kwargs):
+    seq_len = cfg.sequence_len
+
+    def transform_fn(example, tokenizer=None):
+        # Extract prompt and completion from the example
+        prompt_text = example.get(
+            "input", example.get("prompt", example.get("question", ""))
+        )
+        completion_text = example.get(
+            "output", example.get("completion", example.get("answer", ""))
+        )
+
+        if tokenizer is None:
+            return {"prompt": prompt_text}
+
+        pad_id = tokenizer.pad_token_id or tokenizer.eos_token_id
+
+        # Tokenize prompt and completion separately
+        prompt_enc = tokenizer(
+            prompt_text,
+            truncation=False,
+            add_special_tokens=True,
+            return_tensors=None,
+        )
+        completion_enc = tokenizer(
+            completion_text,
+            truncation=False,
+            add_special_tokens=False,
+            return_tensors=None,
+        )
+
+        prompt_ids = prompt_enc["input_ids"]
+        completion_ids = completion_enc["input_ids"]
+
+        # Truncate to fit within seq_len (prioritize keeping prompt + some completion)
+        total_len = len(prompt_ids) + len(completion_ids)
+        if total_len > seq_len:
+            # Truncate completion first, then prompt if needed
+            max_completion = seq_len - len(prompt_ids)
+            if max_completion < 1:
+                # Prompt alone exceeds seq_len — truncate prompt, keep at least 1 completion token
+                prompt_ids = prompt_ids[: seq_len - 1]
+                completion_ids = completion_ids[:1]
+            else:
+                completion_ids = completion_ids[:max_completion]
+
+        input_ids = prompt_ids + completion_ids
+        prompt_length = len(prompt_ids)
+
+        # Labels: -100 for prompt tokens, input_ids for completion tokens
+        labels = [-100] * prompt_length + completion_ids
+
+        # Pad to seq_len
+        pad_len = seq_len - len(input_ids)
+        attention_mask = [1] * len(input_ids) + [0] * pad_len
+        labels = labels + [-100] * pad_len
+        input_ids = input_ids + [pad_id] * pad_len
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "prompt_length": prompt_length,
+        }
+
+    # Signal to remove all original columns (filtered to existing ones at map time)
+    return transform_fn, {
+        "remove_columns": "__all__",
+    }
diff --git a/src/axolotl/prompt_strategies/input_output.py b/src/axolotl/prompt_strategies/input_output.py
index fe14f039cf..c84eecffc2 100644
--- a/src/axolotl/prompt_strategies/input_output.py
+++ b/src/axolotl/prompt_strategies/input_output.py
@@ -1,4 +1,5 @@
 """Module for plain input/output prompt pairs"""
+
 from typing import Generator, Tuple
 
 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
@@ -15,7 +16,6 @@ def __init__(self, *args, eos_token=None, **kwargs):
             self.eos_token = self.tokenizer.eos_token
 
     def tokenize_prompt(self, prompt):
-        # pylint: disable=duplicate-code
         input_ids = []
         labels = []
         for label, text in self.prompter.build_prompt(prompt["segments"]):
diff --git a/src/axolotl/prompt_strategies/instruct.py b/src/axolotl/prompt_strategies/instruct.py
deleted file mode 100644
index 3d63674890..0000000000
--- a/src/axolotl/prompt_strategies/instruct.py
+++ /dev/null
@@ -1,33 +0,0 @@
-"""Module containing the InstructShareGPTPromptTokenizingStrategy class"""
-from typing import Any, Dict, Optional
-
-from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
-from axolotl.prompters import ShareGPTPrompterV2
-
-
-def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
-    conversation = (
-        ds_cfg["conversation"] if ds_cfg and "conversation" in ds_cfg else None
-    )
-    strategy = InstructShareGPTPromptTokenizingStrategy(
-        # pylint: disable=duplicate-code
-        ShareGPTPrompterV2(
-            conversation=conversation,
-        ),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-    return strategy
-
-
-class InstructShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
-    """
-    basic sharegpt strategy to grab conversations from the sample row
-    """
-
-    def get_conversation_thread(self, prompt):
-        return [
-            {"from": "human", "value": prompt["instruction"]},
-            {"from": "gpt", "value": prompt["output"]},
-        ]
diff --git a/src/axolotl/prompt_strategies/jinja_template_analyzer.py b/src/axolotl/prompt_strategies/jinja_template_analyzer.py
new file mode 100644
index 0000000000..70eb496c2f
--- /dev/null
+++ b/src/axolotl/prompt_strategies/jinja_template_analyzer.py
@@ -0,0 +1,335 @@
+"""Module for inspect jinja templates for the variables they use"""
+
+from typing import Dict, Optional, Set, TypedDict, Union
+
+from jinja2 import Environment, meta, nodes
+from jinja2.ext import Extension
+
+
+class JinjaTemplateAnalysis(TypedDict):
+    """
+    Represents the detailed analysis of a Jinja template variable.
+
+    Attributes:
+        accessed_properties (Set[str]): A set of properties accessed from the variable
+            (e.g., `foo.bar` results in 'bar' being accessed for 'foo').
+        accessed_indices (Set[Union[int, float]]): A set of indices accessed from the variable.
+        is_iterated (bool): Indicates if the variable is used as an iteration source in a `for` loop.
+        is_conditional (bool): Indicates if the variable is referenced within a conditional statement (e.g., an `if` block).
+        iteration_source (Optional[str]): The name of the variable being iterated over, if applicable.
+        iteration_target (Optional[Union[str, list[str]]]): The loop target(s) assigned in the iteration.
+    """
+
+    accessed_properties: Set[str]
+    accessed_indices: Set[Union[int, float]]
+    is_iterated: bool
+    is_conditional: bool
+    iteration_source: Optional[str]
+    iteration_target: Optional[Union[str, list[str]]]
+
+
+class GenerationTagIgnore(Extension):
+    """
+    Ignores the generation and endgeneration tags in Jinja templates.
+    """
+
+    tags = {"generation", "endgeneration"}
+
+    def parse(self, parser):
+        parser.stream.skip(1)
+        return nodes.Const("")
+
+
+class JinjaTemplateAnalyzer:
+    """
+    Analyzes Jinja templates to extract information about variable usage,
+    including accessed properties, iteration, and conditional references.
+
+    Attributes:
+        env (jinja2.Environment): The Jinja2 environment used for parsing templates.
+        property_access (Dict[str, Set[str]]): Tracks accessed properties for variables.
+        iteration_targets (Dict[str, str]): Maps iteration target variables to their sources.
+
+    Methods:
+        get_template_variables(template: str) -> Dict[str, Set[str]]:
+            Parse a Jinja template and return a mapping of variables to their accessed properties.
+
+        analyze_template(template: str) -> Dict[str, JinjaTemplateAnalysis]:
+            Perform a detailed analysis of the template, including variable usage,
+            iteration, and conditional references.
+
+    Private Methods:
+        _visit_node(node) -> None:
+            Recursively visit AST nodes to detect attribute access and iteration targets.
+
+        _get_base_name(node) -> Optional[str]:
+            Extract the base variable name from a node.
+
+        _get_target_name(node) -> Optional[Union[str, list[str]]]:
+            Extract the target name(s) from a `For` node.
+    """
+
+    def __init__(self, template: str):
+        self.env: Environment = Environment(
+            autoescape=True,
+            extensions=[GenerationTagIgnore, "jinja2.ext.loopcontrols"],
+        )
+        self.property_access: Dict[str, Set[str]] = {}
+        self.iteration_targets: Dict[str, Union[str, list[str]]] = {}
+        self.index_access: Dict[str, Set[Union[int, float]]] = {}
+        self.ast: nodes.Node = self.env.parse(template)
+        self.template: str = template
+        self.variable_assignments: Dict[str, str] = {}
+
+    def _visit_node(self, node) -> None:
+        """Recursively visit AST nodes to find attribute access."""
+        # Handle attribute access (dot notation)
+        if isinstance(node, nodes.Getattr):
+            base_name = self._get_base_name(node.node)
+            if base_name:
+                self.property_access.setdefault(base_name, set()).add(node.attr)
+
+        # Handle dictionary access (subscript notation)
+        elif isinstance(node, nodes.Getitem):
+            base_name = self._get_base_name(node.node)
+            if base_name and isinstance(node.arg, nodes.Const):
+                value = node.arg.value
+                if isinstance(value, (int, float)):
+                    self.index_access.setdefault(base_name, set()).add(value)
+                else:
+                    self.property_access.setdefault(base_name, set()).add(value)
+
+        elif isinstance(node, nodes.Test) and node.name == "defined":
+            base_name = self._get_base_name(node.node)
+            if base_name:
+                if isinstance(node.node, nodes.Getattr):
+                    self.property_access.setdefault(base_name, set()).add(
+                        node.node.attr
+                    )
+
+        # Handle loop variables
+        elif isinstance(node, nodes.For):
+            iter_name = self._get_base_name(node.iter)
+            target_name = self._get_target_name(node.target)
+            if iter_name and target_name:
+                self.iteration_targets[target_name] = iter_name
+                self.property_access.setdefault(iter_name, set())
+
+        elif isinstance(node, nodes.Assign):
+            target_name = self._get_target_name(node.target)
+            source_name = self._get_base_name(node.node)
+            if target_name and source_name:
+                self.variable_assignments[target_name] = source_name
+
+        elif isinstance(node, nodes.Filter):
+            if node.name == "selectattr":
+                target = self._get_base_name(node.node)
+                if target:
+                    self.variable_assignments[f"filtered_{target}"] = target
+
+        for child in node.iter_child_nodes():
+            self._visit_node(child)
+
+    def _get_target_name(self, node) -> Optional[str]:
+        """Get the target variable name from a For node.
+
+        Args:
+            node: A Jinja AST node representing either a Name or Tuple node
+
+        Returns:
+            - str: For simple variable targets (e.g., "item" in "for item in items")
+            - None: If the node type is not recognized or is a tuple
+        """
+        if isinstance(node, nodes.Name):
+            return node.name
+        return None
+
+    def _get_target_names(self, node) -> list[str]:
+        """Get all target variable names from a For node, including tuple unpacking.
+
+        Args:
+            node: A Jinja AST node representing either a Name or Tuple node
+
+        Returns:
+            List of target variable names
+        """
+        if isinstance(node, nodes.Name):
+            return [node.name]
+
+        if isinstance(node, nodes.Tuple):
+            names = []
+            for n in node.items:
+                if isinstance(n, nodes.Name):
+                    names.append(n.name)
+            return names
+
+        return []
+
+    def _get_base_name(self, node) -> Optional[str]:
+        """Get the base variable name from a node."""
+        if isinstance(node, nodes.Name):
+            return node.name
+
+        if isinstance(node, nodes.Getattr):
+            return self._get_base_name(node.node)
+
+        if isinstance(node, nodes.Getitem):
+            return self._get_base_name(node.node)
+
+        return None
+
+    def get_template_variables(self) -> Dict[str, Set[str]]:
+        """
+        Parse a Jinja template and return both variables and their accessed properties.
+
+        Args:
+            template (str): The Jinja template string
+
+        Returns:
+            Dict[str, Set[str]]: Dictionary mapping variable names to sets of accessed properties
+        """
+        # Parse the template
+        ast = self.env.parse(self.template)
+
+        # Get all undeclared variables
+        variables = meta.find_undeclared_variables(ast)
+
+        # Reset property access tracking
+        self.property_access = {}
+
+        # Visit all nodes to find property access
+        self._visit_node(ast)
+
+        # Create result dictionary
+        result: Dict[str, Set[str]] = {var: set() for var in variables}
+        # Merge in any discovered sub-properties
+        for var, props in self.property_access.items():
+            if var not in result:
+                result[var] = set()
+            result[var].update(props)
+
+        return result
+
+    def analyze_template(self) -> Dict[str, JinjaTemplateAnalysis]:
+        """
+        Provide a detailed analysis of template variables and their usage.
+        """
+        variables = self.get_template_variables()
+        self.iteration_targets = {}
+
+        analysis: Dict[str, JinjaTemplateAnalysis] = {
+            var: JinjaTemplateAnalysis(
+                accessed_properties=props,
+                accessed_indices=set(),
+                is_iterated=False,
+                is_conditional=False,
+                iteration_source=None,
+                iteration_target=None,
+            )
+            for var, props in variables.items()
+        }
+
+        for var, indices in self.index_access.items():
+            if var in analysis:
+                analysis[var]["accessed_indices"] = indices
+
+        def visit_node(node):
+            if isinstance(node, nodes.If):
+
+                def find_test_vars(test_node):
+                    if isinstance(test_node, nodes.Name):
+                        if test_node.name in analysis:
+                            analysis[test_node.name]["is_conditional"] = True
+                    for child in test_node.iter_child_nodes():
+                        find_test_vars(child)
+
+                find_test_vars(node.test)
+
+            if isinstance(node, nodes.For):
+                iter_target = self._get_base_name(node.iter)
+                target_name = self._get_target_name(node.target)
+                if iter_target in analysis:
+                    analysis[iter_target]["is_iterated"] = True
+                    if target_name:
+                        analysis[iter_target]["iteration_target"] = target_name
+                        if isinstance(target_name, str) and target_name not in analysis:
+                            analysis[target_name] = {
+                                "accessed_properties": set(),
+                                "is_iterated": False,
+                                "is_conditional": False,
+                                "iteration_source": iter_target,
+                                "iteration_target": None,
+                            }
+
+            for child in node.iter_child_nodes():
+                visit_node(child)
+
+        visit_node(self.ast)
+        return analysis
+
+    def get_downstream_properties(self, start_var: str) -> Dict[str, Set[str]]:
+        """
+        Get all properties accessed on a variable and its downstream assignments.
+
+        Args:
+            start_var: The starting variable to trace
+
+        Returns:
+            Dict mapping variable names to their accessed properties
+        """
+        visited = set()
+        properties = {}
+
+        def trace_variable(var_name: str):
+            if var_name in visited:
+                return
+            visited.add(var_name)
+
+            # Get direct properties
+            if var_name in self.property_access:
+                properties[var_name] = self.property_access[var_name]
+
+            # Get properties from iteration targets
+            if var_name in self.iteration_targets:
+                target = self.iteration_targets[var_name]
+                if isinstance(target, str):
+                    trace_variable(target)
+                elif isinstance(target, list):
+                    for t in target:
+                        trace_variable(t)
+
+            # Follow assignments
+            for target, source in self.variable_assignments.items():
+                if source == var_name:
+                    trace_variable(target)
+
+            # Check for array slicing
+            analysis = self.analyze_template()
+            if var_name in analysis:
+                var_info = analysis[var_name]
+                if var_info["accessed_indices"]:
+                    # If this variable is sliced, follow the resulting assignment
+                    slice_result = f"{var_name}_slice"
+                    if slice_result in self.property_access:
+                        trace_variable(slice_result)
+
+        trace_variable(start_var)
+        return properties
+
+    def get_message_vars(self, field_messages: str = "messages") -> Set[str]:
+        """
+        Get all properties accessed on messages and derived variables.
+        """
+        all_properties = self.get_downstream_properties(field_messages)
+
+        # Combine all properties from all related variables
+        combined_properties = set()
+        for properties in all_properties.values():
+            combined_properties.update(properties)
+
+        # Also include properties from the message iteration variable
+        analysis = self.analyze_template()
+        if "message" in analysis:
+            combined_properties.update(analysis["message"]["accessed_properties"])
+
+        return combined_properties
diff --git a/src/axolotl/prompt_strategies/kto/__init__.py b/src/axolotl/prompt_strategies/kto/__init__.py
new file mode 100644
index 0000000000..9af6300eb3
--- /dev/null
+++ b/src/axolotl/prompt_strategies/kto/__init__.py
@@ -0,0 +1,9 @@
+"""
+module for KTO style dataset transform strategies
+"""
+
+from functools import partial
+
+from ..base import load as load_base
+
+load = partial(load_base, module_base="axolotl.prompt_strategies.kto")
diff --git a/src/axolotl/prompt_strategies/kto/chatml.py b/src/axolotl/prompt_strategies/kto/chatml.py
new file mode 100644
index 0000000000..945940f3fd
--- /dev/null
+++ b/src/axolotl/prompt_strategies/kto/chatml.py
@@ -0,0 +1,102 @@
+"""
+KTO strategies for chatml
+"""
+
+
+def argilla(
+    cfg,
+    **kwargs,
+):
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        sample["completion"] = f"{sample['completion']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
+def argilla_chat(
+    cfg,
+    **kwargs,
+):
+    """
+    for argilla/kto-mix-15k conversations
+    """
+
+    def transform_fn(sample):
+        sample["prompt"] = (
+            f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
+        )
+        sample["completion"] = f"{sample['completion'][1]['content']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
+def intel(cfg, **kwargs):
+    """
+    For Intel Orca KTO
+    ex: argilla/distilabel-intel-orca-kto
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        sample["completion"] = f"{sample['completion']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
+def prompt_pairs(cfg, **kwargs):
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        sample["completion"] = f"{sample['completion']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
+def ultra(cfg, **kwargs):
+    """
+    for ultrafeedback binarized conversations
+    ex: argilla/ultrafeedback-binarized-preferences-cleaned-kto
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            sample["prompt"] = (
+                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        sample["completion"] = f"{sample['completion']}<|im_end|>"
+        return sample
+
+    return transform_fn
diff --git a/src/axolotl/prompt_strategies/kto/llama3.py b/src/axolotl/prompt_strategies/kto/llama3.py
new file mode 100644
index 0000000000..9061f6f5ea
--- /dev/null
+++ b/src/axolotl/prompt_strategies/kto/llama3.py
@@ -0,0 +1,102 @@
+"""
+KTO strategies for llama-3 chat template
+"""
+
+
+def argilla(
+    cfg,
+    **kwargs,
+):
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        sample["completion"] = f"{sample['completion']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def argilla_chat(
+    cfg,
+    **kwargs,
+):
+    """
+    for argilla/kto-mix-15k conversations
+    """
+
+    def transform_fn(sample):
+        sample["prompt"] = (
+            f"<|start_header_id|>user<|end_header_id|>\n\n{sample['completion'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        )
+        sample["completion"] = f"{sample['completion'][1]['content']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def intel(cfg, **kwargs):
+    """
+    For Intel Orca KTO
+    ex: argilla/distilabel-intel-orca-kto
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        sample["completion"] = f"{sample['completion']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def prompt_pairs(cfg, **kwargs):
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        sample["completion"] = f"{sample['completion']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def ultra(cfg, **kwargs):
+    """
+    for ultrafeedback binarized conversations
+    ex: argilla/ultrafeedback-binarized-preferences-cleaned-kto
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample["prompt"] = (
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        sample["completion"] = f"{sample['completion']}<|eot_id|>"
+        return sample
+
+    return transform_fn
diff --git a/src/axolotl/prompt_strategies/kto/user_defined.py b/src/axolotl/prompt_strategies/kto/user_defined.py
new file mode 100644
index 0000000000..0a52383773
--- /dev/null
+++ b/src/axolotl/prompt_strategies/kto/user_defined.py
@@ -0,0 +1,40 @@
+"""
+User-defined KTO strategies
+"""
+
+
+def default(cfg, dataset_idx=0, **kwargs):
+    ds_cfg = cfg["datasets"][dataset_idx]["type"]
+    if not isinstance(ds_cfg, dict):
+        raise ValueError(
+            f"User-defined dataset type must be a dictionary. Got: {ds_cfg}"
+        )
+    field_prompt = ds_cfg.get("field_prompt", "prompt")
+    field_system = ds_cfg.get("field_system", "system")
+    field_completion = ds_cfg.get("field_completion", "completion")
+    field_label = ds_cfg.get("field_label", "label")
+    prompt_format = ds_cfg.get("prompt_format")
+    if not prompt_format:
+        prompt_format = "{prompt}"
+    completion_format = ds_cfg.get("completion_format")
+    if not completion_format:
+        completion_format = "{completion}"
+
+    def transform_fn(sample):
+        if (
+            "{system}" in prompt_format
+            and field_system in sample
+            and sample[field_system]
+        ):
+            sample["prompt"] = prompt_format.format(
+                system=sample[field_system], prompt=sample[field_prompt]
+            )
+        else:
+            sample["prompt"] = prompt_format.format(prompt=sample[field_prompt])
+        sample["completion"] = completion_format.format(
+            completion=sample[field_completion]
+        )
+        sample["label"] = sample[field_label]
+        return sample
+
+    return transform_fn
diff --git a/src/axolotl/prompt_strategies/llama2_chat.py b/src/axolotl/prompt_strategies/llama2_chat.py
index a1f5ffefff..9eff062ec7 100644
--- a/src/axolotl/prompt_strategies/llama2_chat.py
+++ b/src/axolotl/prompt_strategies/llama2_chat.py
@@ -24,12 +24,14 @@
 Important: Don't use "special_tokens:" in your config.yml if you are not sure what you are doing!
 """
 
-import logging
 from dataclasses import dataclass, field
 from typing import Generator, List, Sequence
 
 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
-from axolotl.prompters import IGNORE_TOKEN_ID, SHAREGPT_ASSERTION_FAILED_ROLE
+from axolotl.prompters import ALTERNATING_ASSERTION_FAILED_ROLE, IGNORE_TOKEN_ID
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
 
 
 @dataclass
@@ -75,7 +77,7 @@ def append_message(self, role: str, message: str):
 
 class LLama2ChatTokenizingStrategy(PromptTokenizingStrategy):
     """
-    Tokenizing strategy for ShareGPT prompts.
+    Tokenizing strategy for Llama2 prompts.
     adapted from https://github.com/lm-sys/FastChat/blob/main/fastchat/train/train.py
     """
 
@@ -129,7 +131,7 @@ def tokenize_prompt(self, prompt):
         if cur_len < self.sequence_len:
             if cur_len != total_len:
                 target[:] = IGNORE_TOKEN_ID
-                logging.warning(
+                LOG.warning(
                     f"WARNING: tokenization mismatch: {cur_len} vs. {total_len}."
                     f" (ignored)"
                 )
@@ -151,7 +153,7 @@ def tokenize_prompt(self, prompt):
         }
 
 
-class Llama2ChatPrompter:  # pylint: disable=too-few-public-methods
+class Llama2ChatPrompter:
     """
     A prompter that generates prompts for Llama2 models.
     """
@@ -188,10 +190,10 @@ def build_prompt(self, source) -> Generator[Llama2ChatConversation, None, None]:
             # Skip the first one if it is not from human
             source = source[1:]
 
-        conv.messages = []  # pylint: disable=R0801
+        conv.messages = []
         for j, sentence in enumerate(source):
             role = roles[sentence["from"]]
-            assert role == conv.roles[j % 2], SHAREGPT_ASSERTION_FAILED_ROLE
+            assert role == conv.roles[j % 2], ALTERNATING_ASSERTION_FAILED_ROLE
             if sentence["value"]:
                 conv.append_message(role, sentence["value"])
         yield conv
diff --git a/src/axolotl/prompt_strategies/messages/__init__.py b/src/axolotl/prompt_strategies/messages/__init__.py
new file mode 100644
index 0000000000..2c920a5682
--- /dev/null
+++ b/src/axolotl/prompt_strategies/messages/__init__.py
@@ -0,0 +1,34 @@
+"""Module to load message prompt strategies."""
+
+import importlib
+import inspect
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def load(tokenizer, cfg, ds_cfg, processor=None):
+    try:
+        strategy = ds_cfg.get("input_transform", "chat")
+
+        load_fn = "load"
+        if strategy.split(".")[-1].startswith("load_"):
+            load_fn = strategy.split(".")[-1]
+            strategy = ".".join(strategy.split(".")[:-1])
+        mod = importlib.import_module(
+            f".{strategy}", "axolotl.prompt_strategies.messages"
+        )
+        func = getattr(mod, load_fn)
+        load_kwargs = {}
+        sig = inspect.signature(func)
+        if "ds_cfg" in sig.parameters:
+            load_kwargs["ds_cfg"] = ds_cfg
+        if "processor" in sig.parameters:
+            load_kwargs["processor"] = processor
+        return func(tokenizer, cfg, **load_kwargs)
+    except ModuleNotFoundError:
+        return None
+    except Exception as exc:
+        LOG.error(f"Failed to load prompt strategy `{strategy}`: {str(exc)}")
+        raise exc
diff --git a/src/axolotl/prompt_strategies/messages/chat.py b/src/axolotl/prompt_strategies/messages/chat.py
new file mode 100644
index 0000000000..4df28fd1f1
--- /dev/null
+++ b/src/axolotl/prompt_strategies/messages/chat.py
@@ -0,0 +1,91 @@
+"""
+Chat dataset wrapping strategy for new internal messages representations
+"""
+
+from typing import Any, Callable, Dict, Optional
+
+from axolotl.core.datasets.chat import TokenizedChatDataset
+from axolotl.core.datasets.transforms.chat_builder import chat_message_transform_builder
+from axolotl.prompt_tokenizers import DatasetWrappingStrategy
+
+
+class ChatMessageDatasetWrappingStrategy(DatasetWrappingStrategy):
+    """
+    Chat dataset wrapping strategy for new internal messages representations
+    """
+
+    def __init__(
+        self,
+        processor,
+        message_transform=None,
+        formatter=None,
+        **kwargs,
+    ):
+        """
+        :param processor: tokenizer or image processor
+        :param kwargs:
+        """
+        self.processor = processor
+        self.dataset = None
+        self.message_transform = message_transform
+        self.formatter = formatter
+
+    def wrap_dataset(
+        self,
+        dataset,
+        process_count: Optional[int] = None,
+        keep_in_memory: Optional[bool] = False,
+        **kwargs,
+    ):
+        self.dataset = TokenizedChatDataset(
+            dataset,
+            message_transform=self.message_transform,
+            model_transform=self.processor,
+            formatter=self.formatter,
+            process_count=process_count,
+            keep_in_memory=keep_in_memory,
+        )
+        return self.dataset
+
+
+def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
+    ds_cfg = ds_cfg or {}
+
+    field_messages = ds_cfg.get("field_messages")
+    message_property_mappings = ds_cfg.get("message_property_mappings")
+    message_field_role = (
+        message_property_mappings.get("role") if message_property_mappings else None
+    )
+    message_field_content = (
+        message_property_mappings.get("content") if message_property_mappings else None
+    )
+    message_field_training = ds_cfg.get("message_field_training")
+
+    builder_kwargs = {}
+    if field_messages:
+        builder_kwargs["conversations_field"] = field_messages
+    if message_field_role:
+        builder_kwargs["message_field_role"] = message_field_role
+    if message_field_content:
+        builder_kwargs["message_field_content"] = message_field_content
+    if message_field_training:
+        builder_kwargs["message_field_training"] = message_field_training
+
+    chat_template = ds_cfg.get("chat_template", cfg.get("chat_template", "chatml"))
+
+    def format_message(x):
+        return x
+
+    if chat_template == "chatml":
+        from axolotl.core.chat.format.chatml import format_message  # noqa F811
+    if chat_template.startswith("llama3"):
+        from axolotl.core.chat.format.llama3x import format_message  # noqa F811
+    message_transform: Callable = chat_message_transform_builder(
+        train_on_inputs=cfg.get("train_on_inputs", False),
+        **builder_kwargs,
+    )
+    strategy = ChatMessageDatasetWrappingStrategy(
+        tokenizer, message_transform=message_transform, formatter=format_message
+    )
+
+    return strategy
diff --git a/src/axolotl/prompt_strategies/metharme.py b/src/axolotl/prompt_strategies/metharme.py
index 52d77c00cf..92355f1cee 100644
--- a/src/axolotl/prompt_strategies/metharme.py
+++ b/src/axolotl/prompt_strategies/metharme.py
@@ -1,17 +1,17 @@
 """Module containing the MetharmenPromptTokenizingStrategy and MetharmePrompter class"""
 
-import logging
 from typing import Tuple
 
+from transformers import BatchEncoding
+
 from axolotl.prompt_tokenizers import InstructionPromptTokenizingStrategy
 from axolotl.prompters import AlpacaPrompter
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl")
+LOG = get_logger(__name__)
 
 IGNORE_TOKEN_ID = -100
 
-# pylint: disable=duplicate-code
-
 
 class MetharmePromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
     """
@@ -37,6 +37,8 @@ def _tokenize(
         )
         if len(result["input_ids"]) == 0:
             LOG.warning("Tokenizer result is empty. You may want to audit your dataset")
+            # A no-BOS tokenizer yields zero tokens for an empty field.
+            return BatchEncoding(data={"input_ids": [], "attention_mask": []})
         # If there's already an EOS token there, subtract from the number added
         if result["input_ids"][-1] == self.tokenizer.eos_token_id:
             num_eos_tokens -= 1
@@ -66,7 +68,7 @@ class MetharmePrompter(AlpacaPrompter):
     turn_format = "{instruction}"
     turn_no_input_format = "{instruction}"
 
-    def __init__(self, *args, **kwargs):  # pylint: disable=super-init-not-called
+    def __init__(self, *args, **kwargs):
         pass
 
 
diff --git a/src/axolotl/prompt_strategies/orcamini.py b/src/axolotl/prompt_strategies/orcamini.py
index 04ce5767dd..1a694cf1d7 100644
--- a/src/axolotl/prompt_strategies/orcamini.py
+++ b/src/axolotl/prompt_strategies/orcamini.py
@@ -9,6 +9,7 @@
 
 Not suited/tested for multiple-turn conversations without further adjustments.
 """
+
 from typing import Generator, Union
 
 from axolotl.prompt_strategies.alpaca_w_system import OpenOrcaPromptTokenizingStrategy
diff --git a/src/axolotl/prompt_strategies/orpo/chat_template.py b/src/axolotl/prompt_strategies/orpo/chat_template.py
index a89dee1575..259f1037dc 100644
--- a/src/axolotl/prompt_strategies/orpo/chat_template.py
+++ b/src/axolotl/prompt_strategies/orpo/chat_template.py
@@ -1,11 +1,12 @@
 """chatml prompt tokenization strategy for ORPO"""
+
 from typing import Any, Dict, Generator, List, Optional, Tuple
 
 from pydantic import BaseModel
 
 from axolotl.prompt_tokenizers import IGNORE_INDEX, PromptTokenizingStrategy
 from axolotl.prompters import Prompter
-from axolotl.utils.chat_templates import chat_templates
+from axolotl.utils.chat_templates import get_chat_template_from_config
 
 
 class Message(BaseModel):
@@ -22,24 +23,17 @@ class MessageList(BaseModel):
     messages: List[Message]
 
 
-def load(
-    tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None, **kwargs
-):  # pylint: disable=possibly-unused-variable,unused-argument
+def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None, **kwargs):
     """
     chatml transforms for datasets with system, input, chosen, rejected
     """
-
-    chat_template = chat_templates("chatml")
-    if ds_cfg and "chat_template" in ds_cfg:
-        chat_template = ds_cfg["chat_template"]
-        try:
-            chat_template = chat_templates(chat_template)
-        except ValueError:
-            pass
-    tokenizer.chat_template = chat_template
+    chat_template_string = get_chat_template_from_config(
+        cfg=cfg, ds_cfg=ds_cfg, tokenizer=tokenizer
+    )
+    tokenizer.chat_template = chat_template_string
 
     return ORPOTokenizingStrategy(
-        ORPOPrompter(chat_template, tokenizer),
+        ORPOPrompter(chat_template_string, tokenizer),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,
@@ -56,7 +50,9 @@ def get_chosen_conversation_thread(self, prompt) -> MessageList:
         messages: List[Message] = []
         if system := prompt.get("system", None):
             messages.append(Message(role="system", content=system, label=False))
-        messages.append(Message(role="user", content=prompt["prompt"], label=False))
+        messages.append(
+            Message(role="user", content=prompt["chosen"][0]["content"], label=False)
+        )
         messages.append(
             Message(
                 role="assistant", content=prompt["chosen"][1]["content"], label=True
@@ -70,7 +66,9 @@ def get_rejected_conversation_thread(self, prompt) -> MessageList:
         messages: List[Message] = []
         if system := prompt.get("system", None):
             messages.append(Message(role="system", content=system, label=False))
-        messages.append(Message(role="user", content=prompt["prompt"], label=False))
+        messages.append(
+            Message(role="user", content=prompt["rejected"][0]["content"], label=False)
+        )
         messages.append(
             Message(
                 role="assistant", content=prompt["rejected"][1]["content"], label=True
@@ -132,7 +130,7 @@ def get_rejected(self, prompt) -> MessageList:
 
 class ORPOTokenizingStrategy(PromptTokenizingStrategy):
     """
-    rejected_input_ids
+    rejected_ids
     input_ids
     rejected_attention_mask
     attention_mask
@@ -152,8 +150,8 @@ def __init__(
     def tokenize_prompt(self, prompt):
         # pass the rejected prompt/row to the Prompter to get the formatted prompt
         prompt_len = 0
-        rejected_message_list = self.dataset_parser.get_rejected_conversation_thread(
-            prompt
+        rejected_message_list: MessageList = (
+            self.dataset_parser.get_rejected_conversation_thread(prompt)
         )
         input_ids = []
         labels = []
@@ -171,10 +169,12 @@ def tokenize_prompt(self, prompt):
                 labels += [IGNORE_INDEX] * (len(input_ids) - prev_idx)
                 prompt_len = len(input_ids)
         # remap the input_ids, attention_mask and labels
-        rejected_input_ids = input_ids
+        rejected_ids = input_ids
         rejected_labels = labels
         # pass the chosen prompt/row to the Prompter to get the formatted prompt
-        chosen_message_list = self.dataset_parser.get_chosen_conversation_thread(prompt)
+        chosen_message_list: MessageList = (
+            self.dataset_parser.get_chosen_conversation_thread(prompt)
+        )
         input_ids = []
         labels = []
         for _, (part, label) in enumerate(
@@ -191,7 +191,7 @@ def tokenize_prompt(self, prompt):
                 labels += [IGNORE_INDEX] * (len(input_ids) - prev_idx)
 
         return {
-            "rejected_input_ids": rejected_input_ids,
+            "rejected_ids": rejected_ids,
             "rejected_labels": rejected_labels,
             "rejected_attention_mask": [1] * len(rejected_labels),
             "input_ids": input_ids,
@@ -217,53 +217,64 @@ def build_prompt(
         for message in message_list.messages:
             conversation.append(message.model_dump())
             if message.role == "system":
-                yield self.tokenizer.apply_chat_template(
-                    conversation,
-                    add_generation_prompt=False,
-                    chat_template=self.chat_template,
-                    tokenize=False,
-                ), False
+                yield (
+                    self.tokenizer.apply_chat_template(
+                        conversation,
+                        add_generation_prompt=False,
+                        chat_template=self.chat_template,
+                        tokenize=False,
+                    ),
+                    False,
+                )
             if message.role == "user":
-                yield self.tokenizer.apply_chat_template(
-                    conversation,
-                    add_generation_prompt=True,
-                    chat_template=self.chat_template,
-                    tokenize=False,
-                ), False
+                yield (
+                    self.tokenizer.apply_chat_template(
+                        conversation,
+                        add_generation_prompt=True,
+                        chat_template=self.chat_template,
+                        tokenize=False,
+                    ),
+                    False,
+                )
             if message.role == "assistant":
-                yield self.tokenizer.apply_chat_template(
-                    conversation,
-                    add_generation_prompt=False,
-                    chat_template=self.chat_template,
-                    tokenize=False,
-                ), True
+                yield (
+                    self.tokenizer.apply_chat_template(
+                        conversation,
+                        add_generation_prompt=False,
+                        chat_template=self.chat_template,
+                        tokenize=False,
+                    ),
+                    True,
+                )
 
 
-def argilla(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+def argilla(cfg, **kwargs):
     dataset_parser = ORPODatasetParsingStrategy()
 
-    chat_template_str = chat_templates(cfg.chat_template)
-
     def transform_fn(sample, tokenizer=None):
         res = {}
 
+        chat_template_string = get_chat_template_from_config(
+            cfg=cfg, tokenizer=tokenizer
+        )
+
         res["prompt"] = tokenizer.apply_chat_template(
             [msg.model_dump() for msg in dataset_parser.get_prompt(sample).messages],
             add_generation_prompt=True,
-            chat_template=chat_template_str,
+            chat_template=chat_template_string,
             tokenize=False,
         )
         prompt_str_len = len(res["prompt"])
         res["chosen"] = tokenizer.apply_chat_template(
             [msg.model_dump() for msg in dataset_parser.get_chosen(sample).messages],
             add_generation_prompt=False,
-            chat_template=chat_template_str,
+            chat_template=chat_template_string,
             tokenize=False,
         )[prompt_str_len:]
         res["rejected"] = tokenizer.apply_chat_template(
             [msg.model_dump() for msg in dataset_parser.get_rejected(sample).messages],
             add_generation_prompt=False,
-            chat_template=chat_template_str,
+            chat_template=chat_template_string,
             tokenize=False,
         )[prompt_str_len:]
 
diff --git a/src/axolotl/prompt_strategies/pretrain.py b/src/axolotl/prompt_strategies/pretrain.py
index 8430a7fcab..eaa4614068 100644
--- a/src/axolotl/prompt_strategies/pretrain.py
+++ b/src/axolotl/prompt_strategies/pretrain.py
@@ -1,4 +1,5 @@
 """pretraining prompt strategies"""
+
 from typing import Generator
 
 from transformers import BatchEncoding
@@ -29,13 +30,16 @@ def __init__(self, *args, max_length=None, text_column="text", **kwargs):
     def _tokenize(
         self, prompt: str, add_eos_token: bool = True, strip_bos_token: bool = False
     ) -> BatchEncoding:
+        # keep the overlap below the window size so small sequence_len values don't
+        # violate the tokenizer's `stride < effective max_length` constraint
+        stride = min(256, (self.max_length - 1) // 2)
         res = self.tokenizer(
             prompt,
             truncation=True,
             max_length=self.max_length - 1,
             add_special_tokens=True,
             return_overflowing_tokens=True,
-            stride=256,
+            stride=stride,
         )
         res["input_ids"] = [
             seq + [self.tokenizer.eos_token_id] for seq in res["input_ids"]
@@ -55,6 +59,8 @@ def load(tokenizer, cfg):
         cfg.train_on_inputs,
         cfg.sequence_len,
         text_column=cfg.pretraining_dataset[0]["text_column"] or "text",
-        max_length=cfg.sequence_len * 64,
+        # windows; larger values produce windows that the downstream
+        # `filter_sequences_by_length` drops wholesale
+        max_length=cfg.sequence_len,
     )
     return strat
diff --git a/src/axolotl/prompt_strategies/pygmalion.py b/src/axolotl/prompt_strategies/pygmalion.py
index 88208f6ec4..8c53a5f279 100644
--- a/src/axolotl/prompt_strategies/pygmalion.py
+++ b/src/axolotl/prompt_strategies/pygmalion.py
@@ -1,7 +1,6 @@
 """Module containing the PygmalionPromptTokenizingStrategy and PygmalionPrompter class"""
 
 import copy
-import logging
 from collections import defaultdict
 from typing import Generator, List, Tuple
 
@@ -10,8 +9,9 @@
     parse_tokenized_to_result,
     tokenize_prompt_default,
 )
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl")
+LOG = get_logger(__name__)
 
 IGNORE_TOKEN_ID = -100
 
@@ -69,7 +69,6 @@ def tokenize_prompt(self, prompt):
                 LOG.warning(f"unknown role in conversation: {role}")
                 res = defaultdict(lambda: [])
 
-            # pylint: disable=duplicate-code
             result, current_len = parse_tokenized_to_result(
                 result,
                 current_len,
@@ -89,7 +88,10 @@ def __init__(self, *args, **kwargs):
         pass
 
     def build_prompt(
-        self, source, *args, **kwargs  # pylint: disable=unused-argument
+        self,
+        source,
+        *args,
+        **kwargs,
     ) -> Generator[Tuple[str, str], None, None]:
         for msg in source:
             yield msg["role"], msg["value"]
diff --git a/src/axolotl/prompt_strategies/sharegpt.py b/src/axolotl/prompt_strategies/sharegpt.py
deleted file mode 100644
index 5f0e7a895e..0000000000
--- a/src/axolotl/prompt_strategies/sharegpt.py
+++ /dev/null
@@ -1,206 +0,0 @@
-"""Module containing the SimpleShareGPTPromptTokenizingStrategy class"""
-
-import logging
-from typing import Any, Dict, Optional, Type
-
-from fastchat.conversation import Conversation, SeparatorStyle, register_conv_template
-
-from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
-from axolotl.prompters import ShareGPTPrompterV2
-from axolotl.utils.tokenization import (
-    chatml_to_conversation,
-    merge_consecutive_messages,
-)
-
-LOG = logging.getLogger("axolotl")
-
-
-def register_chatml_template(system_message=None):
-    system_message = system_message or "You are a helpful assistant."
-    register_conv_template(
-        Conversation(
-            name="chatml",
-            system_template="<|im_start|>system\n{system_message}",
-            system_message=system_message,
-            roles=("<|im_start|>user", "<|im_start|>assistant"),
-            sep_style=SeparatorStyle.CHATML,
-            sep="<|im_end|>",
-        )
-    )
-    register_conv_template(
-        Conversation(
-            name="chatml_glaive",
-            system_template="<|im_start|>system\n{system_message}",
-            system_message=system_message,
-            roles=("<|im_start|>user", "<|im_start|>assistant", "<|im_start|>tool"),
-            sep_style=SeparatorStyle.CHATML,
-            sep="<|im_end|>",
-        )
-    )
-
-
-def register_llama3_template(system_message=None):
-    system_message = system_message or "You are a helpful assistant."
-    register_conv_template(
-        Conversation(
-            name="llama3",
-            system_template="<|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>",
-            system_message=system_message,
-            roles=("user", "assistant"),
-            sep_style=SeparatorStyle.LLAMA3,
-            sep="",
-            stop_str="<|eot_id|>",
-            stop_token_ids=[128001, 128009],
-        )
-    )
-
-
-def build_loader(
-    tokenization_strategy_cls: Type["ShareGPTPromptTokenizingStrategy"],
-    prompter_cls: Type["ShareGPTPrompterV2"],
-    default_conversation: Optional[str] = None,
-):
-    def _load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
-        conversation = (
-            ds_cfg["conversation"]
-            if ds_cfg and "conversation" in ds_cfg
-            else default_conversation
-        )
-        field_human = (
-            ds_cfg["field_human"] if ds_cfg and "field_human" in ds_cfg else None
-        )
-        field_model = (
-            ds_cfg["field_model"] if ds_cfg and "field_model" in ds_cfg else None
-        )
-        roles = ds_cfg["roles"].to_dict() if ds_cfg and "roles" in ds_cfg else None
-        strategy = tokenization_strategy_cls(
-            prompter_cls(
-                conversation=conversation,
-                role_key_model=field_model,
-                role_key_human=field_human,
-                roles=roles,
-            ),
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        if ds_cfg and "strict" in ds_cfg and hasattr(strategy, "strict"):
-            strategy.strict = ds_cfg["strict"]
-        return strategy
-
-    return _load
-
-
-class SimpleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
-    """
-    basic sharegpt strategy to grab conversations from the sample row
-    """
-
-    _strict = False
-
-    @property
-    def strict(self):
-        return self._strict
-
-    @strict.setter
-    def strict(self, strict):
-        self._strict = strict
-
-    def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
-        if self.strict:
-            return conversations
-        role_key = "from"
-        if "role" in conversations[0].keys():
-            role_key = "role"
-        value_key = "value"
-        if "text" in conversations[0].keys():
-            value_key = "text"
-        elif "content" in conversations[0].keys():
-            value_key = "content"
-        # remap roles - allow for assistant turn"
-        role_map = {
-            "user": "human",
-            "human": "human",
-            "assistant": "gpt",
-            "gpt": "gpt",
-            "system": "system",
-        }
-        turns = [
-            {
-                "from": (
-                    role_map[t[role_key]] if t[role_key] in role_map else t[role_key]
-                ),
-                "value": t[value_key],
-            }
-            for t in conversations
-        ]
-        return turns
-
-
-class SimpleRoleShareGPTPromptTokenizingStrategy(
-    SimpleShareGPTPromptTokenizingStrategy
-):
-    """
-    basic sharegpt strategy to grab conversations from the sample row, but uses role instead of from
-    """
-
-    def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
-        # remap role: prompter/assistant, text: ... => from: human/gpt, value: ...
-        turns = [{"from": t["role"], "value": t["value"]} for t in conversations]
-        return turns
-
-
-class GuanacoShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
-    """
-    sharegpt strategy that remaps oasst data to sharegpt format
-    """
-
-    def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
-        # remap role: prompter/assistant, text: ... => from: human/gpt, value: ...
-        role_map = {"prompter": "human", "assistant": "gpt"}
-        turns = [
-            {"from": role_map[t["role"]], "value": t["text"]} for t in conversations
-        ]
-        return turns
-
-
-class UltrachatShareGPTPromptTokenizingStrategy(SimpleShareGPTPromptTokenizingStrategy):
-    """
-    sharegpt strategy that remaps ultrachat data to sharegpt format
-    """
-
-    def get_conversation_thread(self, prompt):
-        conversations = prompt["messages"]
-        role_map = {"user": "human", "assistant": "gpt"}
-        turns = [
-            {"from": role_map[t["role"]], "value": t["content"]} for t in conversations
-        ]
-        return turns
-
-
-class GlaiveShareGPTPromptTokenizingStrategy(SimpleShareGPTPromptTokenizingStrategy):
-    """
-    sharegpt strategy that remaps glaive data to sharegpt format
-    """
-
-    def get_conversation_thread(self, prompt):
-        conversation = chatml_to_conversation(prompt)
-        conversation = merge_consecutive_messages(conversation)
-
-        return conversation
-
-
-load = build_loader(SimpleShareGPTPromptTokenizingStrategy, ShareGPTPrompterV2)
-load_role = build_loader(SimpleRoleShareGPTPromptTokenizingStrategy, ShareGPTPrompterV2)
-load_ultrachat = build_loader(
-    UltrachatShareGPTPromptTokenizingStrategy, ShareGPTPrompterV2
-)
-load_guanaco = build_loader(GuanacoShareGPTPromptTokenizingStrategy, ShareGPTPrompterV2)
-load_glaive = build_loader(
-    GlaiveShareGPTPromptTokenizingStrategy,
-    ShareGPTPrompterV2,
-    default_conversation="chatml_glaive",
-)
diff --git a/src/axolotl/prompt_strategies/sharegpt_jokes.py b/src/axolotl/prompt_strategies/sharegpt_jokes.py
deleted file mode 100644
index 404302c81e..0000000000
--- a/src/axolotl/prompt_strategies/sharegpt_jokes.py
+++ /dev/null
@@ -1,28 +0,0 @@
-"""Module for Jokes prompts using sharegpt style """
-from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
-from axolotl.prompters import ShareGPTPrompterV2
-
-
-def load(tokenizer, cfg):
-    return SimpleJokesShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompterV2(),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-
-
-class SimpleJokesShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
-    """
-    Tokenization strategy for asking bot to tell a joke and then explain why its funny
-    """
-
-    # title, text, explanation
-    def get_conversation_thread(self, prompt):
-        title = "" if not prompt["title"] else prompt["title"] + " "
-        return [
-            {"from": "human", "value": "Tell me a joke."},
-            {"from": "gpt", "value": title + prompt["text"]},
-            {"from": "human", "value": "Why is that joke funny?"},
-            {"from": "gpt", "value": prompt["explanation"]},
-        ]
diff --git a/src/axolotl/prompt_strategies/stepwise_supervised.py b/src/axolotl/prompt_strategies/stepwise_supervised.py
new file mode 100644
index 0000000000..9175126e75
--- /dev/null
+++ b/src/axolotl/prompt_strategies/stepwise_supervised.py
@@ -0,0 +1,116 @@
+"""
+Module for stepwise datasets, typically including a prompt and reasoning traces,
+and (optionally) per-step, or per-prompt-trace labels for reward modelling.
+"""
+
+from itertools import chain
+from typing import Dict, List, Optional, Union
+
+from transformers import BatchEncoding, PreTrainedTokenizer
+
+from axolotl.prompt_tokenizers import IGNORE_INDEX
+from axolotl.utils.dict import DictDefault
+
+
+class StepwiseSupervisedPromptTokenizingStrategy:
+    """
+    Tokenizing strategy for supervised stepwise datasets, typically used for COT-reasoning.
+    These datasets should include the following columns:
+    - prompt: the prompt text
+    - completions: a list of `n` completion steps
+    - labels: a list of `n` labels indicating the "correctness" of each step
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_len: int = 2048,
+        step_separator: str = "\n",
+        max_completion_length: Optional[int] = None,
+        train_on_last_step_only: bool = False,
+    ):
+        self.tokenizer = tokenizer
+        self.sequence_len = sequence_len
+        self.step_separator = step_separator
+        self.max_completion_length = max_completion_length
+        self.train_on_last_step_only = train_on_last_step_only
+
+    def tokenize_prompt(
+        self, prompt: Dict[str, Union[str, List[str]]]
+    ) -> BatchEncoding:
+        # Inspired by TRL's PRMTRainer
+        # https://github.com/huggingface/trl/blob/ed7de87dc766478c024b68f12530d1b0e7c3ff23/trl/trainer/prm_trainer.py#L206
+        prompt_ids = self.tokenizer(prompt["prompt"], add_special_tokens=False)[
+            "input_ids"
+        ]
+
+        completions_ids = [
+            self.tokenizer(completion, add_special_tokens=False)["input_ids"]
+            for completion in prompt["completions"]
+        ]
+
+        # Handle labels
+        if self.train_on_last_step_only:
+            labels = [IGNORE_INDEX] * (len(prompt["labels"]) - 1) + [
+                int(prompt["labels"][-1])
+            ]
+        else:
+            labels = [int(label) for label in prompt["labels"]]
+
+        # Add step separators
+        separator_ids = self.tokenizer.encode(
+            self.step_separator, add_special_tokens=False
+        )
+        completions_ids = [completion + separator_ids for completion in completions_ids]
+
+        # Create step-wise labels
+        labels = [
+            [IGNORE_INDEX] * (len(completion) - 1) + [label]  # type: ignore
+            for completion, label in zip(completions_ids, labels, strict=False)
+        ]
+
+        # Join all steps
+        completion_ids = list(chain(*completions_ids))
+        labels = list(chain(*labels))  # type: ignore
+
+        # Handle max lengths
+        if self.max_completion_length:
+            completion_ids = completion_ids[: self.max_completion_length]
+            labels = labels[: self.max_completion_length]
+
+        # Add BOS token if model has one
+        if self.tokenizer.bos_token_id is not None:
+            prompt_ids = [self.tokenizer.bos_token_id] + prompt_ids
+
+        # Combine prompt and completion
+        input_ids = prompt_ids + completion_ids
+
+        full_labels = [IGNORE_INDEX] * len(prompt_ids) + labels
+        # Apply max sequence length
+        if self.sequence_len:
+            input_ids = input_ids[: self.sequence_len]
+            full_labels = full_labels[: self.sequence_len]
+
+        return {
+            "input_ids": input_ids,
+            "labels": full_labels,
+            "attention_mask": [1] * len(input_ids),
+        }
+
+    @property
+    def supports_batched(self):
+        return False
+
+
+def load(
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    ds_cfg: DictDefault,
+) -> StepwiseSupervisedPromptTokenizingStrategy:
+    return StepwiseSupervisedPromptTokenizingStrategy(
+        tokenizer,
+        cfg.sequence_len,
+        step_separator=ds_cfg.get("step_separator", "\n"),
+        max_completion_length=ds_cfg.max_completion_length,
+        train_on_last_step_only=ds_cfg.get("train_on_last_step_only", False),
+    )
diff --git a/src/axolotl/prompt_strategies/user_defined.py b/src/axolotl/prompt_strategies/user_defined.py
index e20e80c3a4..0bff514e7e 100644
--- a/src/axolotl/prompt_strategies/user_defined.py
+++ b/src/axolotl/prompt_strategies/user_defined.py
@@ -83,16 +83,12 @@ def match_prompt_style(self):
         cfg.sequence_len,
     )
 
-    setattr(
-        strat,
-        "parse_instruction_fields",
-        partial(
-            parse_instruction_fields,
-            ds_cfg.field_instruction,
-            ds_cfg.field_input,
-            ds_cfg.field_output,
-            ds_cfg.field_system,
-            system_prompt,
-        ),
+    strat.parse_instruction_fields = partial(  # type: ignore[method-assign]
+        parse_instruction_fields,
+        ds_cfg.field_instruction,
+        ds_cfg.field_input,
+        ds_cfg.field_output,
+        ds_cfg.field_system,
+        system_prompt,
     )
     return strat
diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py
index bb13cf76dd..ab6b36c172 100644
--- a/src/axolotl/prompt_tokenizers.py
+++ b/src/axolotl/prompt_tokenizers.py
@@ -1,19 +1,15 @@
 """Module containing PromptTokenizingStrategy and Prompter classes"""
 
 import abc
-import copy
-import logging
-from typing import Dict, List, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
-from fastchat.conversation import Conversation
+from datasets import Dataset
 from transformers import BatchEncoding, PreTrainedTokenizer
 
-from axolotl.monkeypatch.fastchat_conversation_turns import (
-    add_get_turns_to_conversation,
-)
-from axolotl.prompters import IGNORE_TOKEN_ID, Prompter
+from axolotl.prompters import Prompter
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl")
+LOG = get_logger(__name__)
 
 IGNORE_INDEX = -100
 LLAMA_DEFAULT_PAD_TOKEN = "<pad>"  # nosec
@@ -21,8 +17,6 @@
 LLAMA_DEFAULT_BOS_TOKEN = "<s>"  # nosec
 LLAMA_DEFAULT_UNK_TOKEN = "<unk>"  # nosec
 
-add_get_turns_to_conversation()
-
 
 class InvalidDataException(Exception):
     """
@@ -30,11 +24,29 @@ class InvalidDataException(Exception):
     """
 
 
+class DatasetWrappingStrategy(abc.ABC):
+    """
+    Abstract class for wrapping datasets for Chat Messages
+    """
+
+    @abc.abstractmethod
+    def wrap_dataset(
+        self,
+        dataset,
+        process_count: int | None = None,
+        keep_in_memory: bool | None = False,
+        **kwargs,
+    ) -> Dataset:
+        pass
+
+
 class PromptTokenizingStrategy(abc.ABC):
     """
     Abstract class for tokenizing strategies
     """
 
+    filter_rows: Optional[Callable] = None
+
     def __init__(
         self,
         prompter: Prompter,
@@ -63,7 +75,7 @@ def _tokenize(
     ) -> BatchEncoding:
         empty = BatchEncoding(data={"input_ids": [], "attention_mask": []})
         if not prompt:
-            LOG.warning("Empty text requested for tokenization.")
+            LOG.warning_once("Empty text requested for tokenization.")
             return empty
 
         result = self.tokenizer(
@@ -106,7 +118,7 @@ def parse_instruction_fields(
     def tokenize_prompt(self, prompt):
         (
             instruction,
-            input,  # pylint: disable=redefined-builtin
+            input,
             response,
         ) = self.parse_instruction_fields(prompt)
         user_prompt = next(
@@ -132,7 +144,10 @@ def tokenize_prompt(self, prompt):
         return tokenized_prompt
 
     def _build_full_prompt(
-        self, instruction, input, response  # pylint: disable=redefined-builtin
+        self,
+        instruction,
+        input,
+        response,
     ):
         return next(
             iter(
@@ -245,10 +260,9 @@ def parse_instruction_fields(self, prompt) -> Tuple[str, str, str, str, str]:
         raise NotImplementedError
 
     def tokenize_prompt(self, prompt):
-        # pylint: disable=duplicate-code
         (
             instruction,
-            input,  # pylint: disable=redefined-builtin
+            input,
             output,
             reflection,
             corrected,
@@ -275,9 +289,7 @@ def tokenize_prompt(self, prompt):
 
         return tokenized_full_prompt
 
-    def _build_full_prompt(
-        self, instruction, input, output, reflection, corrected
-    ):  # pylint: disable=redefined-builtin
+    def _build_full_prompt(self, instruction, input, output, reflection, corrected):
         return next(
             iter(
                 self.prompter.build_prompt(
@@ -291,6 +303,10 @@ def _build_full_prompt(
         )
 
     def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False):
+        empty = BatchEncoding(data={"input_ids": [], "attention_mask": []})
+        if not prompt:
+            LOG.warning_once("Empty text requested for tokenization.")
+            return empty
         result = self.tokenizer(
             prompt,
             truncation=True,
@@ -298,6 +314,9 @@ def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False):
             padding=False,
             return_tensors=None,
         )
+        if len(result["input_ids"]) == 0:
+            LOG.warning("Tokenizer result is empty. You may want to audit your dataset")
+            return empty
         if (
             result["input_ids"][-1] != self.tokenizer.eos_token_id
             and len(result["input_ids"]) < self.sequence_len
@@ -325,145 +344,6 @@ def parse_instruction_fields(self, prompt) -> Tuple[str, str, str, str, str]:
         )
 
 
-class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy):
-    """
-    Tokenizing strategy for ShareGPT prompts.
-    """
-
-    def get_conversation_thread(self, prompt):
-        return prompt["conversations"]
-
-    def tokenize_prompt(self, prompt):
-        # Initial values. We will append to these as we go through the conversation.
-        result, current_len = tokenize_prompt_default()
-        conversation: Conversation = (
-            self.prompter._conversation.copy()  # pylint: disable=protected-access
-        )
-
-        input_roles = {conversation.roles[0]}
-        output_roles = {conversation.roles[1]}
-
-        if len(conversation.roles) == 3:
-            tool_role_label = conversation.roles[2]
-            input_roles.add(tool_role_label)
-
-        # Add roles from the config
-        if self.prompter.roles:
-            if "input" in self.prompter.roles and self.prompter.roles["input"]:
-                for role in self.prompter.roles["input"]:
-                    input_roles.add(role)
-
-            if "output" in self.prompter.roles and self.prompter.roles["output"]:
-                for role in self.prompter.roles["output"]:
-                    output_roles.add(role)
-
-        # support for custom roles from the dataset, only useful for vicuna style prompts/roles
-        role_remap = []
-        if (
-            conversation.name == "vicuna_v1.1"
-            and "roles" in prompt
-            and len(prompt["roles"]) >= 2
-        ):
-            role_remap = [
-                {"from": conversation.roles[0], "to": prompt["roles"][0]},
-                {"from": conversation.roles[1], "to": prompt["roles"][1]},
-            ]
-
-        try:
-            for _, part in enumerate(
-                self.prompter.build_prompt(self.get_conversation_thread(prompt))
-            ):
-                if not isinstance(part, tuple):
-                    LOG.warning(f"expected tuple, got {part}")
-                    continue
-
-                role, content = part
-
-                # Uses "in" because role contains extra characters
-                input_turn = any(r.lower() in role.lower() for r in input_roles)
-                output_turn = any(r.lower() in role.lower() for r in output_roles)
-                empty_role = role.strip() == ""
-
-                if not any([input_turn, output_turn, empty_role]):
-                    LOG.warning(f"unhandled role: {role}")
-                    continue
-
-                if input_turn:
-                    role = (
-                        role.replace(role_remap[0]["from"], role_remap[0]["to"])
-                        if role_remap
-                        else role
-                    )
-                    turn = role + content
-                    # this is still the user query, we should
-                    if not content.strip():
-                        LOG.warning(f"user turn has empty text: {prompt}")
-                    res = self._tokenize(
-                        turn,
-                        add_eos_token=False,
-                        strip_bos_token=True,
-                    )
-                    if self.train_on_inputs:
-                        labels = copy.deepcopy(res["input_ids"])
-                    else:
-                        # everything from this is masked out from the labels
-                        labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
-                elif output_turn:
-                    role = (
-                        role.replace(role_remap[1]["from"], role_remap[1]["to"])
-                        if role_remap
-                        else role
-                    )
-                    turn = role + content
-                    # this should be the assistant response, should end with an eos token
-                    if not content.strip():
-                        LOG.warning(f"assistant turn has empty text: {prompt}")
-                    add_eos_token = not (
-                        conversation.name == "chatml"
-                        and conversation.sep == self.tokenizer.eos_token
-                    )
-                    res = self._tokenize(
-                        turn,
-                        add_eos_token=add_eos_token,
-                        strip_bos_token=True,
-                    )
-                    role_res = self._tokenize(
-                        role.rstrip(),
-                        add_eos_token=False,
-                        strip_bos_token=True,
-                    )
-                    labels = copy.deepcopy(res["input_ids"])
-                    if not self.train_on_inputs:
-                        # mask out role tokens from the labels
-                        len_role = len(role_res["input_ids"])
-                        labels[:len_role] = [IGNORE_TOKEN_ID] * min(
-                            len_role, len(labels)
-                        )
-                elif empty_role:
-                    turn = content
-                    # this is only ever the first part, should include the bos token and the user query
-                    res = self._tokenize(
-                        turn, add_eos_token=False, strip_bos_token=False
-                    )
-                    if self.train_on_inputs:
-                        labels = copy.deepcopy(res["input_ids"])
-                    else:
-                        # everything from this is masked out from the labels
-                        labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
-
-                # pylint: disable=duplicate-code
-                result, current_len = parse_tokenized_to_result(
-                    result,
-                    current_len,
-                    res,
-                    labels,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                )
-            return result
-        except (KeyError, AssertionError, IndexError) as err:
-            raise InvalidDataException(str(err)) from err
-
-
 def tokenize_prompt_default() -> Tuple[Dict[str, List[int]], int]:
     """
     Returns the default values for the tokenize prompt function
diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py
index 60ea5c99f9..9543996f7e 100644
--- a/src/axolotl/prompters.py
+++ b/src/axolotl/prompters.py
@@ -1,13 +1,13 @@
 """Module containing prompters"""
 
-import logging
 from enum import Enum
 from typing import Generator, Optional, Union
 
 from colorama import Fore
-from fastchat.conversation import Conversation, get_conv_template
 
-LOG = logging.getLogger("axolotl")
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
 IGNORE_TOKEN_ID = -100
 REPR_TEMPLATE = "\n<start>\n" + Fore.CYAN + "{full_prompt}" + Fore.RESET + "\n<end>\n"
 
@@ -20,6 +20,7 @@ class PromptStyle(Enum):
     INSTRUCT = "instruct"
     CHAT = "chat"
     CHATML = "chatml"
+    PHI = "phi"
 
 
 class Prompter:
@@ -38,30 +39,35 @@ class AlpacaPrompter(Prompter):
     system_format: str = "{system}"
     turn_format: str
     turn_no_input_format: str
-    prompt_style: Optional[PromptStyle] = None
+    prompt_style: Optional[str] = None
 
-    def __init__(self, prompt_style=PromptStyle.INSTRUCT.value):
+    def __init__(self, prompt_style: Optional[str] = PromptStyle.INSTRUCT.value):
         self.prompt_style = prompt_style if prompt_style else PromptStyle.INSTRUCT.value
         self.match_prompt_style()
 
     def match_prompt_style(self):
-        # pylint: disable=duplicate-code
         if self.prompt_style == PromptStyle.INSTRUCT.value:
             self.turn_format = "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
             self.turn_no_input_format = (
                 "### Instruction:\n{instruction}\n\n### Response:\n"
             )
             self.system_format = "{system}\n\n"
-        if self.prompt_style == PromptStyle.CHAT.value:
+        elif self.prompt_style == PromptStyle.CHAT.value:
             self.turn_format = "USER: {instruction}\n{input}\nASSISTANT:"
             self.turn_no_input_format = "USER: {instruction}\nASSISTANT:"
             self.system_format = "SYSTEM: {system}\n"
-        if self.prompt_style == PromptStyle.CHATML.value:
+        elif self.prompt_style == PromptStyle.CHATML.value:
             self.turn_format = "<|im_start|>user\n{instruction}\n{input}<|im_end|>\n<|im_start|>assistant\n"
             self.turn_no_input_format = (
                 "<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n"
             )
             self.system_format = "<|im_start|>system\n{system}<|im_end|>\n"
+        elif self.prompt_style == PromptStyle.PHI.value:
+            self.turn_format = "<|user|>\n{instruction}<|end|>{input}<|assistant|>"
+            self.turn_no_input_format = (
+                "<|user|>\n{instruction}<|end|>\n<|assistant|>\n"
+            )
+            self.system_format = "<|system|>\n{system}<|end|>\n"
 
     def _build_result(self, instruction, input_text, output):
         # returns the full prompt from instruction and optional input
@@ -86,7 +92,7 @@ def _build_result(self, instruction, input_text, output):
     def build_prompt(
         self,
         instruction: str,
-        input: Union[None, str] = None,  # pylint: disable=redefined-builtin
+        input: Union[None, str] = None,
         output: Union[None, str] = None,
     ) -> Generator[str, None, None]:
         yield self._build_result(instruction, input, output)
@@ -211,7 +217,7 @@ def match_prompt_style(self):
     def _build_result(
         self,
         instruction: str,
-        input: Union[None, str] = None,  # pylint: disable=redefined-builtin
+        input: Union[None, str] = None,
         output: Union[None, str] = None,
         reflection: Union[None, str] = None,
         corrected: Union[None, str] = None,
@@ -235,12 +241,11 @@ def _build_result(
     def build_prompt(
         self,
         instruction: str,
-        input: Union[None, str] = None,  # pylint: disable=redefined-builtin
+        input: Union[None, str] = None,
         output: Union[None, str] = None,
         reflection: Union[None, str] = None,
         corrected: Union[None, str] = None,
     ) -> Generator[str, None, None]:
-        # pylint: disable=duplicate-code
         yield self._build_result(
             instruction,
             input,
@@ -255,141 +260,10 @@ def __repr__(self) -> str:
         )
 
 
-SHAREGPT_ASSERTION_FAILED_ROLE = (
+ALTERNATING_ASSERTION_FAILED_ROLE = (
     "Role did not alternate between turns (gpt and human). Please check your data."
 )
 
-CONVERSATION_ROLE_FORMAT = {
-    "chatml": "<|im_start|>{ROLE}",
-    "zephyr": "<|{ROLE}|>",
-    "vicuna_v1.1": "{ROLE}",
-    "llama3": "<|start_header_id|>{ROLE}<|end_header_id|>",
-}
-
-
-class ShareGPTPrompter(Prompter):  # pylint: disable=too-few-public-methods
-    """
-    A prompter that generates prompts for the ShareGPT
-    """
-
-    role_key_human = "human"
-    role_key_model = "gpt"
-    # Optional, only used for tool usage datasets.
-    role_key_tool: Optional[str] = None
-    # Optional, role input/output mapping
-    roles: Optional[dict] = None
-
-    def __init__(
-        self,
-        prompt_style=None,  # pylint: disable=unused-argument
-        conversation: Optional[Union[str, Conversation]] = None,
-        role_key_human: Optional[str] = None,
-        role_key_model: Optional[str] = None,
-        role_key_tool: Optional[str] = None,
-        roles: Optional[dict] = None,
-    ):
-        if conversation:
-            if isinstance(conversation, Conversation):
-                self._conversation = conversation
-            else:
-                self._conversation = get_conv_template(conversation)
-        else:
-            self._conversation = get_conv_template("vicuna_v1.1")
-        if role_key_human:
-            self.role_key_human = role_key_human
-        if role_key_model:
-            self.role_key_model = role_key_model
-        if role_key_tool:
-            self.role_key_tool = role_key_tool
-        if roles:
-            self.roles = roles
-
-    def _build_result(self, source):
-        if len(source) < 2:
-            # If there isn't a back and forth conversation, ignore it
-            # also happens on the data splitting leaving empty conversations
-            raise IndexError(
-                f"A conversation entry has less than 2 messages :\n{source}"
-            )
-
-        conv = self._conversation.copy()
-
-        # Add the conversation system prompt if provided, otherwise use the default one
-        if source[0]["from"] == "system":
-            conv.set_system_message(source[0]["value"])
-            source.pop(0)
-
-        roles = {self.role_key_human: conv.roles[0], self.role_key_model: conv.roles[1]}
-        if self.role_key_tool:
-            roles[self.role_key_tool] = conv.roles[2]
-
-        try:
-            # Apply prompt templates
-            if source[0]["from"] not in roles:
-                # Skip the first one if it is not from human
-                source = source[1:]
-        except IndexError as err:
-            # sometimes there is a bing or system chat
-            raise err
-
-        conv.messages = []
-        for _, sentence in enumerate(source):
-            from_role = sentence["from"]
-            if from_role in roles:
-                role = roles[from_role]
-            else:
-                if self._conversation.name not in CONVERSATION_ROLE_FORMAT:
-                    raise NotImplementedError(
-                        f"Role ({role}) not in default roles, and {self._conversation.name} does not support role remapping yet."
-                        "Please help us by creating an Issue to add support for this conversation type."
-                    )
-
-                role = CONVERSATION_ROLE_FORMAT[self._conversation.name].format(
-                    ROLE=from_role
-                )
-
-            if len(conv.messages) > 0 and ((role == conv.messages[-1][0])):
-                if (
-                    role != "assistant"
-                ):  # back to back assistant calls may be okay for tool calls
-                    LOG.warning(f"{SHAREGPT_ASSERTION_FAILED_ROLE}: {sentence}")
-
-            conv.append_message(role, sentence["value"])
-
-        return conv.get_turns()
-
-    def build_prompt(self, source) -> Generator[str, None, None]:
-        turns = self._build_result(source)
-
-        for part in turns:
-            if part[0] and not part[1]:
-                LOG.warning(f"role with empty message: {part[0]}")
-            yield part
-
-    def __repr__(self) -> str:
-        turns = self._build_result([{"from": "{from}", "value": "{value}"}])
-        return "\n".join([REPR_TEMPLATE.format(full_prompt=part) for part in turns])
-
-
-class ShareGPTPrompterV2(ShareGPTPrompter):
-    """
-    A V2 prompter that generates prompts for the ShareGPT
-    """
-
-    def __init__(
-        self,
-        conversation: Optional[Union[str, Conversation]] = None,
-        role_key_human: Optional[str] = None,
-        role_key_model: Optional[str] = None,
-        roles: Optional[dict] = None,
-    ):
-        super().__init__(
-            conversation=conversation,
-            role_key_human=role_key_human,
-            role_key_model=role_key_model,
-            roles=roles,
-        )
-
 
 class UnsupportedPrompter(Prompter):
     """
diff --git a/src/axolotl/scripts/__init__.py b/src/axolotl/scripts/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/scripts/process_cleanup.py b/src/axolotl/scripts/process_cleanup.py
new file mode 100644
index 0000000000..cf06d06970
--- /dev/null
+++ b/src/axolotl/scripts/process_cleanup.py
@@ -0,0 +1,233 @@
+"""Reusable process lifecycle management for vLLM serve scripts.
+
+Handles graceful shutdown, orphan cleanup, and health monitoring for
+multiprocessing-based server architectures where a main process
+dispatches work to worker subprocesses that spawn GPU-heavy children
+(e.g., vLLM EngineCore).
+
+Usage:
+
+    from axolotl.scripts.process_cleanup import ProcessManager
+
+    manager = ProcessManager(processes, connections)
+    manager.register_signal_handlers()
+
+    # In FastAPI lifespan:
+    async with manager.lifespan_context():
+        yield  # server runs here
+
+    # In endpoints:
+    manager.check_workers_alive()  # raises if dead
+
+    # In worker command loop:
+    if manager.is_fatal_error(exc):
+        break  # exit worker
+"""
+
+import asyncio
+import atexit
+import os
+from multiprocessing import Process
+from multiprocessing.connection import Connection
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def kill_process_tree(pid: int) -> None:
+    """Kill a process and all its descendants (depth-first)."""
+    import subprocess  # nosec B404
+
+    try:
+        result = subprocess.run(  # nosec B603 B607
+            ["pgrep", "-P", str(pid)],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if result.returncode == 0:
+            for child_pid in result.stdout.strip().split("\n"):
+                child_pid = child_pid.strip()
+                if child_pid:
+                    kill_process_tree(int(child_pid))
+    except (FileNotFoundError, ValueError):
+        pass
+
+    try:
+        os.kill(pid, 9)
+    except (ProcessLookupError, PermissionError):
+        pass
+
+
+def cleanup_orphan_processes(*patterns: str) -> None:
+    """Kill orphan processes matching any of the given patterns.
+
+    Uses ``pgrep -f`` to find processes. Skips the current process.
+    Intended for cleaning up GPU-holding subprocesses (EngineCore)
+    that survive their parent's death.
+    """
+    import subprocess  # nosec B404
+
+    my_pid = os.getpid()
+    for pattern in patterns:
+        try:
+            result = subprocess.run(  # nosec B603 B607
+                ["pgrep", "-f", pattern],
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+            if result.returncode == 0:
+                for pid in result.stdout.strip().split("\n"):
+                    pid = pid.strip()
+                    if pid and int(pid) != my_pid:
+                        try:
+                            os.kill(int(pid), 9)
+                            logger.info("Killed orphan process %s (%s)", pid, pattern)
+                        except (ProcessLookupError, ValueError):
+                            pass
+        except FileNotFoundError:
+            pass
+
+
+def is_fatal_worker_error(exc: Exception) -> bool:
+    """Check if an exception indicates the worker should exit.
+
+    Returns True for errors from which the worker cannot recover,
+    such as the vLLM EngineCore dying.
+    """
+    exc_str = str(exc)
+    exc_type = type(exc).__name__
+    return (
+        "EngineCore" in exc_str
+        or "EngineDeadError" in exc_type
+        or "engine" in exc_str.lower()
+        and "died" in exc_str.lower()
+    )
+
+
+def safe_recv(conn: Connection):
+    """Receive from a pipe, returning an error dict if the pipe is broken."""
+    try:
+        return conn.recv()
+    except EOFError:
+        return {"error": "Worker process died (pipe closed)", "kind": "worker_dead"}
+
+
+class ProcessManager:
+    """Manages worker process lifecycle for a FastAPI-based serve script.
+
+    Handles:
+    - Signal-based shutdown (SIGTERM)
+    - Background health monitoring (detects dead workers)
+    - Process tree cleanup on exit
+    - Orphan EngineCore cleanup
+
+    Args:
+        processes: List of worker Process objects.
+        connections: List of parent-side Pipe connections to workers.
+        orphan_patterns: Process name patterns to search for orphans on cleanup.
+            Defaults to ``["VLLM::EngineCore"]``.
+        monitor_interval: Seconds between worker health checks.
+        shutdown_timeout: Seconds to wait for graceful worker exit before SIGTERM.
+        kill_timeout: Seconds to wait after SIGTERM before SIGKILL.
+    """
+
+    def __init__(
+        self,
+        processes: list[Process],
+        connections: list[Connection],
+        orphan_patterns: list[str] | None = None,
+        monitor_interval: float = 5.0,
+        shutdown_timeout: float = 30.0,
+        kill_timeout: float = 15.0,
+    ):
+        self.processes = processes
+        self.connections = connections
+        self.orphan_patterns = orphan_patterns or ["VLLM::EngineCore"]
+        self.monitor_interval = monitor_interval
+        self.shutdown_timeout = shutdown_timeout
+        self.kill_timeout = kill_timeout
+
+    def register_cleanup(self) -> None:
+        """Register atexit cleanup for orphan processes.
+
+        Does NOT override SIGTERM — let uvicorn handle it naturally,
+        which triggers the lifespan shutdown where ``_shutdown_workers``
+        runs. The atexit handler is a safety net for abnormal exits.
+        """
+        atexit.register(self._cleanup_orphans)
+
+    def check_workers_alive(self) -> None:
+        """Raise RuntimeError if any worker process has died.
+
+        Call this at the start of request handlers to fail fast
+        instead of hanging on a broken pipe.
+        """
+        dead = [i for i, p in enumerate(self.processes) if not p.is_alive()]
+        if dead:
+            raise RuntimeError(
+                f"vLLM worker(s) {dead} died. Restart the server to recover."
+            )
+
+    def get_health_status(self) -> dict:
+        """Return health status dict. Use as the /health endpoint response."""
+        dead = [i for i, p in enumerate(self.processes) if not p.is_alive()]
+        if dead:
+            return {
+                "status": "unhealthy",
+                "dead_workers": dead,
+                "message": "Worker(s) died. Restart the server.",
+            }
+        return {"status": "ok"}
+
+    async def monitor_workers(self) -> None:
+        """Background coroutine that detects dead workers and exits.
+
+        When all workers are dead, cleans up their process trees and
+        orphan subprocesses, then force-exits the server.
+        """
+        while True:
+            await asyncio.sleep(self.monitor_interval)
+            alive = [p.is_alive() for p in self.processes]
+            if not any(alive):
+                logger.error(
+                    "All vLLM workers died. Shutting down server. "
+                    "Check logs for EngineCore errors and restart."
+                )
+                # Kill process trees for any workers that left orphans
+                for p in self.processes:
+                    if p.pid is not None:
+                        kill_process_tree(p.pid)
+                self._cleanup_orphans()
+                os._exit(1)
+
+    def _shutdown_workers(self) -> None:
+        """Send shutdown commands and escalate to kill if needed."""
+        for conn in self.connections:
+            try:
+                conn.send({"type": "shutdown"})
+            except Exception:
+                pass
+        for i, p in enumerate(self.processes):
+            if not p.is_alive():
+                continue
+            p.join(timeout=self.shutdown_timeout)
+            if p.is_alive():
+                logger.warning(
+                    "Worker %d didn't exit in %.0fs, sending SIGTERM",
+                    i,
+                    self.shutdown_timeout,
+                )
+                p.terminate()
+                p.join(timeout=self.kill_timeout)
+            if p.is_alive():
+                logger.warning("Worker %d didn't respond to SIGTERM, force killing", i)
+                p.kill()
+                p.join(timeout=5)
+        self._cleanup_orphans()
+        logger.info("Worker shutdown complete")
+
+    def _cleanup_orphans(self) -> None:
+        cleanup_orphan_processes(*self.orphan_patterns)
diff --git a/src/axolotl/scripts/vllm_serve_lora.py b/src/axolotl/scripts/vllm_serve_lora.py
new file mode 100644
index 0000000000..67e710163d
--- /dev/null
+++ b/src/axolotl/scripts/vllm_serve_lora.py
@@ -0,0 +1,910 @@
+"""vLLM serve script with native LoRA adapter support.
+
+Extends TRL's vllm_serve to enable direct LoRA adapter loading in vLLM,
+instead of merging adapter weights into the base model before syncing.
+
+Usage:
+    Set ``vllm.serve_module: axolotl.scripts.vllm_serve_lora`` in your config,
+    or ``trl.vllm_lora_sync: true`` to auto-select.
+
+Benefits over merge-sync:
+    - Syncs only LoRA adapter weights via filesystem instead of full merged model via NCCL
+    - vLLM handles LoRA application natively (Punica kernels)
+    - No NCCL communicator needed for weight sync
+"""
+
+import os
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from itertools import chain
+from multiprocessing import Pipe, Process
+from multiprocessing.connection import Connection
+from typing import Any
+
+from trl.scripts.vllm_serve import (
+    ScriptArguments,
+    chunk_list,
+    extract_logprobs,
+)
+
+try:
+    from trl.scripts.vllm_serve import get_open_port
+except ImportError:
+    try:
+        from vllm.utils import get_open_port
+    except ImportError:
+        from vllm.utils.network_utils import get_open_port
+from vllm import LLM, SamplingParams
+from vllm.lora.request import LoRARequest
+
+from axolotl.scripts.process_cleanup import (
+    ProcessManager,
+    is_fatal_worker_error,
+    safe_recv,
+)
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class LoRAScriptArguments(ScriptArguments):
+    """Extended script arguments with LoRA support."""
+
+    enable_lora: bool = field(
+        default=True,
+        metadata={"help": "Enable LoRA adapter support in vLLM."},
+    )
+    max_lora_rank: int = field(
+        default=64,
+        metadata={"help": "Maximum LoRA rank supported."},
+    )
+    max_loras: int = field(
+        default=2,
+        metadata={"help": "Maximum number of LoRA adapters loaded simultaneously."},
+    )
+    lora_dtype: str = field(
+        default="bfloat16",
+        metadata={"help": "Data type for LoRA weights."},
+    )
+    worker_extension_cls: str = field(
+        default="trl.scripts.vllm_serve.WeightSyncWorkerExtension",
+        metadata={"help": "vLLM worker extension class for weight synchronization."},
+    )
+
+
+def llm_worker(
+    script_args: LoRAScriptArguments,
+    data_parallel_rank: int,
+    master_port: int,
+    connection: Connection,
+) -> None:
+    """Worker process that creates a vLLM LLM with LoRA enabled."""
+    # For DP with TP=1: pin each worker to its own GPU via CUDA_VISIBLE_DEVICES.
+    # vLLM's LLM() offline mode doesn't support DP env vars natively, so we
+    # isolate each worker to a single GPU and let vLLM think it's the only one.
+    if script_args.data_parallel_size > 1 and script_args.tensor_parallel_size == 1:
+        visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+        if visible:
+            gpu_ids = visible.split(",")
+            os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[data_parallel_rank]
+        else:
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(data_parallel_rank)
+    else:
+        os.environ["VLLM_DP_RANK"] = str(data_parallel_rank)
+        os.environ["VLLM_DP_RANK_LOCAL"] = str(data_parallel_rank)
+        os.environ["VLLM_DP_SIZE"] = str(script_args.data_parallel_size)
+        os.environ["VLLM_DP_MASTER_PORT"] = str(master_port)
+
+    llm = LLM(
+        model=script_args.model,
+        revision=script_args.revision,
+        tensor_parallel_size=script_args.tensor_parallel_size,
+        gpu_memory_utilization=script_args.gpu_memory_utilization,
+        enforce_eager=script_args.enforce_eager,
+        dtype=script_args.dtype,
+        enable_prefix_caching=script_args.enable_prefix_caching,
+        kv_cache_dtype=script_args.kv_cache_dtype,
+        max_model_len=script_args.max_model_len,
+        worker_extension_cls=script_args.worker_extension_cls,
+        trust_remote_code=script_args.trust_remote_code,
+        model_impl=script_args.vllm_model_impl,
+        logprobs_mode="processed_logprobs",
+        # LoRA
+        enable_lora=script_args.enable_lora,
+        max_lora_rank=script_args.max_lora_rank,
+        max_loras=script_args.max_loras,
+        lora_dtype=script_args.lora_dtype,
+    )
+
+    connection.send({"status": "ready"})
+
+    def _worker_cleanup():
+        """Clean up the LLM and its EngineCore subprocess on worker exit."""
+        from axolotl.scripts.process_cleanup import cleanup_orphan_processes
+
+        try:
+            llm.collective_rpc(method="close_communicator")
+        except Exception:
+            pass
+        # Kill EngineCore children of this worker
+        cleanup_orphan_processes("VLLM::EngineCore")
+
+    import atexit as _atexit
+
+    _atexit.register(_worker_cleanup)
+
+    while True:
+        try:
+            command = connection.recv()
+        except (KeyboardInterrupt, EOFError):
+            break
+
+        if command.get("type") == "shutdown":
+            break
+
+        if command["type"] in ["call", "fire_and_forget"]:
+            method_name = command["method"]
+            args = command.get("args", ())
+            kwargs = command.get("kwargs", {})
+
+            # Reconstruct LoRARequest from serialized dict (can't pickle across pipe)
+            if "lora_request" in kwargs and kwargs["lora_request"] is not None:
+                lr = kwargs["lora_request"]
+                kwargs["lora_request"] = LoRARequest(
+                    lora_name=lr["lora_name"],
+                    lora_int_id=lr["lora_int_id"],
+                    lora_path=lr["lora_path"],
+                    load_inplace=lr.get("load_inplace", False),
+                )
+
+            try:
+                method = getattr(llm, method_name)
+                result = method(*args, **kwargs)
+            except Exception as exc:
+                logger.warning("Worker method %s failed: %s", method_name, exc)
+                if command["type"] == "call":
+                    connection.send({"error": str(exc), "kind": "worker_error"})
+                if is_fatal_worker_error(exc):
+                    logger.error(
+                        "Fatal worker error (EngineCore died), exiting. "
+                        "Restart the vLLM server to recover."
+                    )
+                    break
+                continue
+            if command["type"] == "call":
+                connection.send(result)
+        elif command["type"] == "shutdown":
+            break
+
+
+def main(script_args: ScriptArguments):
+    """Start vLLM workers with LoRA support and the HTTP server."""
+    import asyncio
+
+    import uvicorn
+    from fastapi import FastAPI
+    from pydantic import BaseModel, Field as PydanticField
+
+    # Request/Response models (defined locally like TRL's vllm_serve.main)
+    class GenerateRequest(BaseModel):
+        prompts: list[str] | list[list[int]]
+        images: list[str] | None = None
+        n: int = 1
+        repetition_penalty: float = 1.0
+        temperature: float = 1.0
+        top_p: float = 1.0
+        top_k: int = -1
+        min_p: float = 0.0
+        max_tokens: int = 16
+        logprobs: int | None = 0
+        truncate_prompt_tokens: int | None = None
+        structured_outputs_regex: str | None = None
+        generation_kwargs: dict = PydanticField(default_factory=dict)
+
+    class GenerateResponse(BaseModel):
+        prompt_ids: list[list[int]]
+        completion_ids: list[list[int]]
+        logprobs: list[list[list[float]]]
+        logprob_token_ids: list[list[list[int]]]
+
+    class ChatRequest(BaseModel):
+        messages: list[list[dict]]
+        n: int = 1
+        repetition_penalty: float = 1.0
+        temperature: float = 1.0
+        top_p: float = 1.0
+        top_k: int = -1
+        min_p: float = 0.0
+        max_tokens: int = 16
+        logprobs: int | None = 0
+        truncate_prompt_tokens: int | None = None
+        structured_outputs_regex: str | None = None
+        generation_kwargs: dict = PydanticField(default_factory=dict)
+        chat_template_kwargs: dict = PydanticField(default_factory=dict)
+
+    class ChatResponse(BaseModel):
+        prompt_ids: list[list[int]]
+        completion_ids: list[list[int]]
+        logprobs: list[list[list[float]]]
+        logprob_token_ids: list[list[list[int]]]
+
+    class InitCommunicatorRequest(BaseModel):
+        host: str
+        port: int
+        world_size: int
+        client_device_uuid: str
+
+    # Wrap plain ScriptArguments with LoRA defaults
+    if not isinstance(script_args, LoRAScriptArguments):
+        lora_args = LoRAScriptArguments.__new__(LoRAScriptArguments)
+        for f in ScriptArguments.__dataclass_fields__:
+            setattr(lora_args, f, getattr(script_args, f))
+        # Apply LoRA defaults
+        for f in LoRAScriptArguments.__dataclass_fields__:
+            if f not in ScriptArguments.__dataclass_fields__:
+                setattr(
+                    lora_args, f, LoRAScriptArguments.__dataclass_fields__[f].default
+                )
+        script_args = lora_args
+
+    # Spawn workers
+    master_port = get_open_port()
+    connections: list[Connection] = []
+    processes: list[Process] = []
+    for dp_rank in range(script_args.data_parallel_size):
+        parent_conn, child_conn = Pipe()
+        process = Process(
+            target=llm_worker,
+            args=(script_args, dp_rank, master_port, child_conn),
+        )
+        process.start()
+        connections.append(parent_conn)
+        processes.append(process)
+
+    # Process lifecycle management
+    manager = ProcessManager(processes, connections)
+    manager.register_cleanup()
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        import time
+
+        startup_timeout = 300  # 5 minutes
+        start_time = time.monotonic()
+        ready: set[int] = set()
+        while len(ready) < script_args.data_parallel_size:
+            elapsed = time.monotonic() - start_time
+            if elapsed > startup_timeout:
+                raise RuntimeError(
+                    f"vLLM workers failed to start within {startup_timeout}s "
+                    f"({len(ready)}/{script_args.data_parallel_size} ready)"
+                )
+            for i, (conn, proc) in enumerate(zip(connections, processes, strict=True)):
+                if id(conn) in ready:
+                    continue
+                if not proc.is_alive():
+                    raise RuntimeError(
+                        f"vLLM worker {i} exited unexpectedly during startup"
+                    )
+                if conn.poll():
+                    msg = conn.recv()
+                    if isinstance(msg, dict) and msg.get("status") == "ready":
+                        ready.add(id(conn))
+            await asyncio.sleep(0.1)
+
+        monitor_task = asyncio.create_task(manager.monitor_workers())
+        yield
+        monitor_task.cancel()
+        manager._shutdown_workers()
+
+    app = FastAPI(lifespan=lifespan)
+
+    # --- Access logging middleware ---
+    import time as _time
+
+    @app.middleware("http")
+    async def access_log_middleware(request, call_next):
+        t0 = _time.monotonic()
+        response = await call_next(request)
+        elapsed = _time.monotonic() - t0
+        logger.info(
+            "%s %s %d %.3fs",
+            request.method,
+            request.url.path,
+            response.status_code,
+            elapsed,
+        )
+        return response
+
+    # --- Active LoRA state (shared across endpoints via closure) ---
+    active_lora: dict = {"request": None}
+
+    # Serializes access to the worker pipe. The underlying
+    # multiprocessing.Connection is a single full-duplex stream shared
+    # across all HTTP handlers; concurrent requests interleave bytes on
+    # the wire and corrupt the pickle framing (seen as
+    # ``UnpicklingError: pickle data was truncated``). Any endpoint that
+    # does ``conn.send(...); conn.recv()`` MUST hold this lock across
+    # the round-trip so only one inflight call at a time per pipe.
+    worker_pipe_lock = asyncio.Lock()
+
+    # ------------------------------------------------------------------
+    # LoRA-specific endpoints
+    # ------------------------------------------------------------------
+
+    class SetLoRARequest(BaseModel):
+        lora_name: str
+        lora_int_id: int
+        lora_path: str
+        load_inplace: bool = False
+
+    @app.post("/set_lora_adapter/")
+    async def set_lora_adapter(request: SetLoRARequest):
+        """Register a LoRA adapter for all subsequent generate/chat calls."""
+        active_lora["request"] = {
+            "lora_name": request.lora_name,
+            "lora_int_id": request.lora_int_id,
+            "lora_path": request.lora_path,
+            "load_inplace": request.load_inplace,
+        }
+        logger.info(
+            "Set active LoRA: %s (id=%d, path=%s)",
+            request.lora_name,
+            request.lora_int_id,
+            request.lora_path,
+        )
+        return {"status": "ok"}
+
+    @app.post("/clear_lora_adapter/")
+    async def clear_lora_adapter():
+        """Clear active LoRA adapter (revert to base model)."""
+        active_lora["request"] = None
+        return {"status": "ok"}
+
+    # ------------------------------------------------------------------
+    # Standard endpoints (mirrors TRL's vllm_serve)
+    # ------------------------------------------------------------------
+
+    @app.get("/health/")
+    async def health():
+        status = manager.get_health_status()
+        if status["status"] != "ok":
+            from fastapi.responses import JSONResponse
+
+            return JSONResponse(status_code=503, content=status)
+        return status
+
+    @app.get("/get_world_size/")
+    async def get_world_size():
+        return {
+            "world_size": script_args.tensor_parallel_size
+            * script_args.data_parallel_size
+        }
+
+    @app.post("/generate/", response_model=GenerateResponse)
+    async def generate(request: GenerateRequest):
+        """Generate completions with optional LoRA adapter."""
+        manager.check_workers_alive()
+
+        import base64
+        from io import BytesIO
+
+        import vllm
+        from packaging.version import Version
+
+        try:
+            from vllm.sampling_params import GuidedDecodingParams
+        except ImportError:
+            GuidedDecodingParams = None  # not available in vLLM 0.17+
+
+        images: list[str | None] = request.images or [None] * len(request.prompts)  # type: ignore[assignment,list-item]
+        prompts: list[dict[str, Any]] = []
+        for prompt, image in zip(request.prompts, images, strict=True):
+            # Support both string prompts and token ID lists
+            row: dict[str, Any]
+            if isinstance(prompt, list):
+                row = {"prompt_token_ids": prompt}
+            else:
+                row = {"prompt": prompt}
+            if image is not None:
+                from PIL import Image
+
+                row["multi_modal_data"] = {
+                    "image": Image.open(BytesIO(base64.b64decode(image)))
+                }
+            prompts.append(row)
+
+        generation_kwargs = {
+            "n": request.n,
+            "repetition_penalty": request.repetition_penalty,
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+            "top_k": request.top_k,
+            "min_p": request.min_p,
+            "max_tokens": request.max_tokens,
+            "logprobs": request.logprobs,
+        }
+        generation_kwargs.update(request.generation_kwargs)
+
+        if Version(vllm.__version__) <= Version("0.10.2"):
+            key = "guided_decoding"
+            if request.structured_outputs_regex is not None:
+                generation_kwargs[key] = GuidedDecodingParams(
+                    regex=request.structured_outputs_regex
+                )
+            else:
+                generation_kwargs.setdefault(key, None)
+        else:
+            from vllm.sampling_params import StructuredOutputsParams
+
+            key = "structured_outputs"
+            if request.structured_outputs_regex is not None:
+                generation_kwargs[key] = StructuredOutputsParams(
+                    regex=request.structured_outputs_regex
+                )
+            elif isinstance(generation_kwargs.get(key), dict):
+                generation_kwargs[key] = StructuredOutputsParams(
+                    **generation_kwargs[key]
+                )
+            else:
+                generation_kwargs.setdefault(key, None)
+
+        sampling_params = SamplingParams(**generation_kwargs)
+        chunked_prompts = chunk_list(prompts, script_args.data_parallel_size)
+
+        for conn, chunk in zip(connections, chunked_prompts, strict=True):
+            if not chunk:
+                chunk = [{"prompt": "<placeholder>"}]
+            kwargs = {
+                "prompts": chunk,
+                "sampling_params": sampling_params,
+                "lora_request": active_lora["request"],
+            }
+            conn.send({"type": "call", "method": "generate", "kwargs": kwargs})
+
+        # Use run_in_executor so blocking recv() doesn't freeze the event loop
+        # (allows /set_lora_adapter/ and other endpoints to be served concurrently)
+        loop = asyncio.get_running_loop()
+
+        all_outputs = await asyncio.gather(
+            *(loop.run_in_executor(None, safe_recv, conn) for conn in connections)
+        )
+        all_outputs = [
+            o for o, c in zip(all_outputs, chunked_prompts, strict=True) if c
+        ]
+        # Check for worker errors before flattening
+        for o in all_outputs:
+            if isinstance(o, dict) and "error" in o:
+                raise RuntimeError(f"vLLM worker error: {o['error']}")
+        all_outputs = list(chain.from_iterable(all_outputs))
+
+        return {
+            "prompt_ids": [o.prompt_token_ids for o in all_outputs],
+            "completion_ids": [
+                list(out.token_ids) for o in all_outputs for out in o.outputs
+            ],
+            "logprobs": extract_logprobs(all_outputs)[0],
+            "logprob_token_ids": extract_logprobs(all_outputs)[1],
+        }
+
+    @app.post("/chat/", response_model=ChatResponse)
+    async def chat(request: ChatRequest):
+        """Chat endpoint with optional LoRA adapter."""
+        manager.check_workers_alive()
+        generation_kwargs = {
+            "n": request.n,
+            "repetition_penalty": request.repetition_penalty,
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+            "top_k": request.top_k,
+            "min_p": request.min_p,
+            "max_tokens": request.max_tokens,
+            "logprobs": request.logprobs,
+        }
+        generation_kwargs.update(request.generation_kwargs)
+        sampling_params = SamplingParams(**generation_kwargs)
+        chunked = chunk_list(request.messages, script_args.data_parallel_size)
+        for conn, chunk in zip(connections, chunked, strict=True):
+            if not chunk:
+                chunk = [[{"role": "user", "content": "<placeholder>"}]]
+            kwargs = {
+                "messages": chunk,
+                "sampling_params": sampling_params,
+                "use_tqdm": False,
+                "lora_request": active_lora["request"],
+            }
+            conn.send({"type": "call", "method": "chat", "kwargs": kwargs})
+
+        loop = asyncio.get_running_loop()
+        all_outputs = await asyncio.gather(
+            *(loop.run_in_executor(None, conn.recv) for conn in connections)
+        )
+        all_outputs = [o for o, c in zip(all_outputs, chunked, strict=True) if c]
+        all_outputs = list(chain.from_iterable(all_outputs))
+
+        return {
+            "prompt_ids": [o.prompt_token_ids for o in all_outputs],
+            "completion_ids": [
+                list(out.token_ids) for o in all_outputs for out in o.outputs
+            ],
+            "logprobs": extract_logprobs(all_outputs)[0],
+            "logprob_token_ids": extract_logprobs(all_outputs)[1],
+        }
+
+    # --- OpenAI-compatible endpoints (for NeMo Gym agent integration) ---
+
+    @app.get("/v1/models")
+    async def list_models():
+        """OpenAI-compatible models endpoint."""
+        return {
+            "object": "list",
+            "data": [
+                {"id": script_args.model, "object": "model", "owned_by": "axolotl"}
+            ],
+        }
+
+    @app.post("/v1/chat/completions")
+    async def openai_chat_completions(request_body: dict):
+        """OpenAI-compatible chat completions endpoint.
+
+        Translates OpenAI format to our internal /chat/ format so NeMo Gym's
+        model server proxy can call us directly.
+        """
+        messages_list = request_body.get("messages", [])
+        temperature = request_body.get("temperature", 1.0)
+        max_tokens = request_body.get("max_tokens", 512)
+        top_p = request_body.get("top_p", 1.0)
+        n = request_body.get("n", 1)
+
+        generation_kwargs = {
+            "n": n,
+            "temperature": temperature,
+            "top_p": top_p,
+            "max_tokens": max_tokens,
+            "logprobs": 0,  # Always return logprobs (NeMo Gym needs them)
+        }
+        sampling_params = SamplingParams(
+            **{k: v for k, v in generation_kwargs.items() if v is not None}
+        )
+
+        # Send to vLLM worker
+        chunked = chunk_list([messages_list], script_args.data_parallel_size)
+        for conn, chunk in zip(connections, chunked, strict=True):
+            if not chunk:
+                chunk = [[{"role": "user", "content": "<placeholder>"}]]
+            kwargs = {
+                "messages": chunk,
+                "sampling_params": sampling_params,
+                "use_tqdm": False,
+                "lora_request": active_lora["request"],
+            }
+            conn.send({"type": "call", "method": "chat", "kwargs": kwargs})
+
+        all_outputs = [conn.recv() for conn in connections]
+        all_outputs = [o for o, c in zip(all_outputs, chunked, strict=True) if c]
+        all_outputs = list(chain.from_iterable(all_outputs))
+
+        if not all_outputs:
+            return {"choices": [], "model": script_args.model}
+
+        # Format as OpenAI response
+        import uuid
+
+        choices = []
+        for i, output in enumerate(all_outputs):
+            for j, out in enumerate(output.outputs):
+                text = out.text
+                # Extract token IDs if requested
+                # Build logprobs in OpenAI format
+                lp_list = None
+                if out.logprobs:
+                    lp_list = {
+                        "content": [
+                            {"token": "", "logprob": next(iter(lp.values())).logprob}  # nosec B105
+                            for lp in out.logprobs
+                        ]
+                    }
+
+                choice = {
+                    "index": i * n + j,
+                    "message": {"role": "assistant", "content": text},
+                    "finish_reason": "stop"
+                    if out.finish_reason == "stop"
+                    else "length",
+                    "logprobs": lp_list,
+                }
+                # Include token ID information for NeMo Gym
+                choice["prompt_token_ids"] = output.prompt_token_ids
+                choice["generation_token_ids"] = list(out.token_ids)
+                if out.logprobs:
+                    choice["generation_log_probs"] = [
+                        next(iter(lp.values())).logprob for lp in out.logprobs
+                    ]
+                choices.append(choice)
+
+        prompt_tokens = len(all_outputs[0].prompt_token_ids) if all_outputs else 0
+        completion_tokens = sum(
+            len(out.token_ids) for o in all_outputs for out in o.outputs
+        )
+
+        return {
+            "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
+            "object": "chat.completion",
+            "model": script_args.model,
+            "choices": choices,
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+            },
+        }
+
+    @app.post("/v1/completions")
+    async def openai_completions(request_body: dict):
+        """OpenAI-compatible text-completions endpoint.
+
+        Accepts either a string ``prompt`` or a list-of-int
+        ``prompt_token_ids`` (as the text-completions spec allows). Routes
+        to the internal vLLM generate method with the active LoRA adapter
+        and returns an OpenAI /v1/completions-shaped response including
+        per-choice ``prompt_token_ids``, ``generation_token_ids``, and
+        ``generation_log_probs`` for NeMo Gym agents that need raw
+        tokens + logprobs.
+        """
+        import uuid
+
+        prompt_raw = request_body.get("prompt")
+        temperature = request_body.get("temperature", 1.0)
+        max_tokens = request_body.get("max_tokens", 512)
+        top_p = request_body.get("top_p", 1.0)
+        n = request_body.get("n", 1)
+        logprobs = request_body.get("logprobs") or 0
+        stop_token_ids = request_body.get("stop_token_ids") or None
+
+        # Accept either a string or a list[int] token id prompt. Lists
+        # must contain ints only (raise on lists of strings so callers get
+        # a clear error). Also accept [[int, int, ...]] nesting for the
+        # rare case callers pass a single-prompt batch.
+        if (
+            isinstance(prompt_raw, list)
+            and prompt_raw
+            and isinstance(prompt_raw[0], list)
+        ):
+            prompt_raw = prompt_raw[0]
+
+        prompt_dict: dict[str, Any] = {}
+        if isinstance(prompt_raw, list):
+            prompt_dict = {"prompt_token_ids": prompt_raw}
+        elif isinstance(prompt_raw, str):
+            prompt_dict = {"prompt": prompt_raw}
+        else:
+            return {
+                "error": {
+                    "message": ("prompt must be a string or a list of token ids"),
+                    "type": "invalid_request",
+                }
+            }
+
+        generation_kwargs: dict[str, Any] = {
+            "n": n,
+            "temperature": temperature,
+            "top_p": top_p,
+            "max_tokens": max_tokens,
+            "logprobs": logprobs,
+        }
+        if stop_token_ids:
+            generation_kwargs["stop_token_ids"] = stop_token_ids
+        sampling_params = SamplingParams(
+            **{k: v for k, v in generation_kwargs.items() if v is not None}
+        )
+
+        chunked = chunk_list([prompt_dict], script_args.data_parallel_size)
+
+        # Hold the pipe lock across send+recv — concurrent requests would
+        # otherwise interleave pickle frames on the worker connection.
+        async with worker_pipe_lock:
+            for conn, chunk in zip(connections, chunked, strict=True):
+                if not chunk:
+                    chunk = [{"prompt": "<placeholder>"}]
+                kwargs = {
+                    "prompts": chunk,
+                    "sampling_params": sampling_params,
+                    "lora_request": active_lora["request"],
+                }
+                conn.send({"type": "call", "method": "generate", "kwargs": kwargs})
+
+            loop = asyncio.get_running_loop()
+            all_outputs = await asyncio.gather(
+                *(loop.run_in_executor(None, safe_recv, conn) for conn in connections)
+            )
+
+        all_outputs = [o for o, c in zip(all_outputs, chunked, strict=True) if c]
+        for o in all_outputs:
+            if isinstance(o, dict) and "error" in o:
+                raise RuntimeError(f"vLLM worker error: {o['error']}")
+        all_outputs = list(chain.from_iterable(all_outputs))
+
+        if not all_outputs:
+            return {"choices": [], "model": script_args.model}
+
+        choices = []
+        for i, output in enumerate(all_outputs):
+            for j, out in enumerate(output.outputs):
+                text = out.text
+                # OpenAI-style `logprobs` block for text-completions:
+                #   { "tokens": [...], "token_logprobs": [...] }
+                lp_block = None
+                if out.logprobs:
+                    tokens_str: list[str] = []
+                    token_lps: list[float] = []
+                    for step in out.logprobs:
+                        chosen = next(iter(step.values()))
+                        tokens_str.append(getattr(chosen, "decoded_token", "") or "")
+                        token_lps.append(float(chosen.logprob))
+                    lp_block = {
+                        "tokens": tokens_str,
+                        "token_logprobs": token_lps,
+                    }
+
+                choice = {
+                    "index": i * n + j,
+                    "text": text,
+                    "finish_reason": "stop"
+                    if out.finish_reason == "stop"
+                    else "length",
+                    "logprobs": lp_block,
+                    # NeMo-Gym / retrace agent extras — preserved on the
+                    # choice so callers with raw-token pipelines don't
+                    # have to re-tokenize.
+                    "prompt_token_ids": output.prompt_token_ids,
+                    "generation_token_ids": list(out.token_ids),
+                    "generation_log_probs": (
+                        [float(next(iter(lp.values())).logprob) for lp in out.logprobs]
+                        if out.logprobs
+                        else []
+                    ),
+                }
+                choices.append(choice)
+
+        prompt_tokens = len(all_outputs[0].prompt_token_ids) if all_outputs else 0
+        completion_tokens = sum(
+            len(out.token_ids) for o in all_outputs for out in o.outputs
+        )
+
+        return {
+            "id": f"cmpl-{uuid.uuid4().hex[:8]}",
+            "object": "text_completion",
+            "model": script_args.model,
+            "choices": choices,
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+            },
+        }
+
+    # --- Weight sync endpoints (legacy fallback, same as TRL) ---
+
+    @app.post("/init_communicator/")
+    async def init_communicator(request: InitCommunicatorRequest):
+        world_size = (
+            script_args.tensor_parallel_size * script_args.data_parallel_size + 1
+        )
+        kwargs = {
+            "method": "init_communicator",
+            "args": (
+                request.host,
+                request.port,
+                world_size,
+                request.client_device_uuid,
+            ),
+        }
+        msg = {"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}
+        loop = asyncio.get_running_loop()
+        await asyncio.gather(
+            *(loop.run_in_executor(None, c.send, msg) for c in connections)
+        )
+        return {"message": "Initializing communicator"}
+
+    class UpdateWeightsRequest(BaseModel):
+        name: str
+        dtype: str
+        shape: list[int]
+
+    @app.post("/update_named_param/")
+    async def update_named_param(request: UpdateWeightsRequest):
+        kwargs = {
+            "method": "update_named_param",
+            "args": (request.name, request.dtype, tuple(request.shape)),
+        }
+        msg = {"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}
+        loop = asyncio.get_running_loop()
+        await asyncio.gather(
+            *(loop.run_in_executor(None, c.send, msg) for c in connections)
+        )
+        return {"message": "Updating parameter"}
+
+    class BatchUpdateWeightsRequest(BaseModel):
+        params: list[dict]
+
+    @app.post("/batch_update_named_params/")
+    async def batch_update_named_params(request: BatchUpdateWeightsRequest):
+        params_list = [
+            (p["name"], p["dtype"], tuple(p["shape"])) for p in request.params
+        ]
+        kwargs = {"method": "batch_update_named_params", "args": (params_list,)}
+        msg = {"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}
+        loop = asyncio.get_running_loop()
+        await asyncio.gather(
+            *(loop.run_in_executor(None, c.send, msg) for c in connections)
+        )
+        return {"message": f"Batch update for {len(params_list)} params"}
+
+    class HTTPWeightUpdateRequest(BaseModel):
+        """Weight update via HTTP (no NCCL needed)."""
+
+        params: list[
+            dict
+        ]  # [{"name": str, "dtype": str, "shape": list, "data": str (base64)}]
+
+    @app.post("/http_update_weights/")
+    async def http_update_weights(request: HTTPWeightUpdateRequest):
+        """Update model weights via HTTP — no NCCL communicator required.
+
+        Tensor data is sent as base64-encoded raw bytes in the request body.
+        Slower than NCCL for large models but works without cross-process setup.
+        """
+        from axolotl.utils.weight_serde import (
+            decode_from_http,
+            encode_for_ipc,
+        )
+
+        weights_to_load = [decode_from_http(p) for p in request.params]
+
+        # Send all weights in a single IPC call.  Tensors don't survive
+        # vLLM's multiproc IPC, so serialize as raw bytes + metadata.
+        param_entries = [
+            encode_for_ipc(name, weight) for name, weight in weights_to_load
+        ]
+        kwargs = {
+            "method": "http_load_weights_batch",
+            "kwargs": {"params": param_entries},
+        }
+        msg = {"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}
+        loop = asyncio.get_running_loop()
+        await asyncio.gather(
+            *(loop.run_in_executor(None, c.send, msg) for c in connections)
+        )
+        return {"message": f"HTTP weight update for {len(weights_to_load)} params"}
+
+    @app.post("/reset_prefix_cache/")
+    async def reset_prefix_cache():
+        # Fire-and-forget: send reset without expecting a reply.
+        # Using "fire_and_forget" type so workers don't send back a response
+        # that would sit in the pipe and corrupt the next recv() for
+        # generate/chat calls.
+        for conn in connections:
+            conn.send({"type": "fire_and_forget", "method": "reset_prefix_cache"})
+        return {"message": "Reset prefix cache received"}
+
+    @app.post("/close_communicator/")
+    async def close_communicator():
+        kwargs = {"method": "close_communicator"}
+        for conn in connections:
+            conn.send(
+                {
+                    "type": "fire_and_forget",
+                    "method": "collective_rpc",
+                    "kwargs": kwargs,
+                }
+            )
+        return {"message": "Closing communicator"}
+
+    uvicorn.run(
+        app,
+        host=script_args.host,
+        port=script_args.port,
+        log_level=script_args.log_level,
+        access_log=True,
+    )
diff --git a/src/axolotl/scripts/vllm_worker_ext.py b/src/axolotl/scripts/vllm_worker_ext.py
new file mode 100644
index 0000000000..61ba32246a
--- /dev/null
+++ b/src/axolotl/scripts/vllm_worker_ext.py
@@ -0,0 +1,208 @@
+"""Extended vLLM worker extension with batch weight sync support.
+
+Subclasses TRL's WeightSyncWorkerExtension to add:
+- batch_update_named_params: receives multiple params in one call
+- Auto-close stale communicator on re-init
+- _direct_set_weight: proper handling for stacked (qkv_proj, gate_up_proj) params,
+  including LoRA-wrapped models where vLLM inserts base_layer into the hierarchy
+"""
+
+import torch
+
+try:
+    from transformers import is_torch_xpu_available
+except ImportError:
+    is_torch_xpu_available = lambda: False  # noqa: E731
+
+from trl.scripts.vllm_serve import WeightSyncWorkerExtension
+
+from axolotl.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+# Stacked param name mapping: shard_name -> (packed_name, shard_order)
+_STACKED_PARAMS = {
+    "q_proj": ("qkv_proj", 0),
+    "k_proj": ("qkv_proj", 1),
+    "v_proj": ("qkv_proj", 2),
+    "gate_proj": ("gate_up_proj", 0),
+    "up_proj": ("gate_up_proj", 1),
+}
+
+
+class BatchWeightSyncWorkerExtension(WeightSyncWorkerExtension):
+    """Worker extension that adds batch weight update and direct weight setting."""
+
+    def init_communicator(self, host, port, world_size, client_device_uuid):
+        """Auto-close stale communicator before re-initializing."""
+        if self.communicator is not None:
+            self.close_communicator()
+        super().init_communicator(host, port, world_size, client_device_uuid)
+
+    def _direct_set_weight(self, name: str, weight: torch.Tensor) -> None:
+        """Directly copy weight data into the model, handling stacked params.
+
+        Bypasses model.load_weights() which may fail on vLLM 0.17's new
+        module-tree weight loader for stacked params (qkv_proj, gate_up_proj).
+
+        Handles LoRA-wrapped params where vLLM inserts ``base_layer`` into the
+        parameter hierarchy (e.g. ``qkv_proj.base_layer.weight``).
+        """
+        model = self.model_runner.model
+        params_dict = dict(model.named_parameters())
+
+        # Handle VLM models where trainer and vLLM use different prefixes.
+        # Trainer (PEFT stripped): "model.layers.X..." or "model.language_model.layers.X..."
+        # vLLM (Qwen3.5):         "language_model.model.layers.X..."
+        if name not in params_dict:
+            # Try common prefix remappings
+            for src_prefix, dst_prefix in [
+                ("model.language_model.layers.", "language_model.model.layers."),
+                ("model.layers.", "language_model.model.layers."),
+            ]:
+                if name.startswith(src_prefix):
+                    name = dst_prefix + name[len(src_prefix) :]
+                    break
+
+        # Check if this is a simple direct param (exists as-is)
+        if name in params_dict:
+            params_dict[name].data.copy_(weight.to(params_dict[name].dtype))
+            return
+
+        # Also check with base_layer inserted: x.y.weight -> x.y.base_layer.weight
+        parts_bl = name.rsplit(".", 1)
+        if len(parts_bl) == 2:
+            base_layer_name = f"{parts_bl[0]}.base_layer.{parts_bl[1]}"
+            if base_layer_name in params_dict:
+                params_dict[base_layer_name].data.copy_(
+                    weight.to(params_dict[base_layer_name].dtype)
+                )
+                return
+
+        # Handle stacked params: e.g. "model.layers.0.self_attn.q_proj.weight"
+        # -> "model.layers.0.self_attn.qkv_proj.weight" with shard offset
+        parts = name.rsplit(".", 2)  # [prefix, layer_name, suffix]
+        if len(parts) == 3:
+            prefix, layer_name, suffix = parts
+            if layer_name in _STACKED_PARAMS:
+                packed_name, shard_idx = _STACKED_PARAMS[layer_name]
+                for packed_full in [
+                    f"{prefix}.{packed_name}.{suffix}",
+                    f"{prefix}.{packed_name}.base_layer.{suffix}",
+                ]:
+                    if packed_full not in params_dict:
+                        continue
+                    param = params_dict[packed_full]
+                    # Navigate to the packed module to find shard sizes
+                    module_path = packed_full.rsplit(".", 1)[0]  # strip .weight/.bias
+                    if ".base_layer" in module_path:
+                        module_path = module_path.replace(".base_layer", "")
+                    module = model
+                    for attr in module_path.split("."):
+                        module = getattr(module, attr, None)
+                        if module is None:
+                            break
+                    # LoRA wrappers don't have output_sizes directly;
+                    # check base_layer for the underlying parallel linear
+                    if module is not None and not hasattr(module, "output_sizes"):
+                        base = getattr(module, "base_layer", None)
+                        if base is not None and hasattr(base, "output_sizes"):
+                            module = base
+                    if module is not None and hasattr(module, "output_sizes"):
+                        tp_size = getattr(module, "tp_size", 1)
+                        sizes = [s // tp_size for s in module.output_sizes]
+                        offset = sum(sizes[:shard_idx])
+                        shard_size = sizes[shard_idx]
+                        param.data[offset : offset + shard_size].copy_(
+                            weight.to(param.dtype)
+                        )
+                        return
+
+        # Fallback: try load_weights (may work for non-stacked params)
+        # Log the actual param names available for debugging
+        sample_keys = [
+            k for k in params_dict if "layers.31.mlp" in k or "layers.31.self_attn" in k
+        ][:3]
+        logger.warning(
+            "Falling back to load_weights for param: %s (sample vLLM keys: %s)",
+            name,
+            sample_keys,
+        )
+        model.load_weights(weights=[(name, weight)])
+
+    def update_named_param(self, name, dtype, shape):
+        """Override to use _direct_set_weight instead of load_weights."""
+        if self.communicator is None:
+            raise RuntimeError("Communicator not initialized.")
+
+        dtype = getattr(torch, dtype.split(".")[-1])
+        weight = torch.empty(shape, dtype=dtype, device=self.device)
+
+        if is_torch_xpu_available():
+            self.communicator.broadcast(weight, root=self.client_rank)
+            self.communicator.barrier()
+        else:
+            self.communicator.broadcast(weight, src=self.client_rank)
+            self.communicator.group.barrier()
+
+        self._direct_set_weight(name, weight)
+
+    def batch_update_named_params(self, params_list: list[tuple[str, str, tuple]]):
+        """Receive and apply multiple weight tensors in sequence.
+
+        Args:
+            params_list: List of (name, dtype_str, shape) tuples.
+        """
+        if self.communicator is None:
+            raise RuntimeError("Communicator not initialized.")
+
+        weights_to_load = []
+        for name, dtype_str, shape in params_list:
+            dtype = getattr(torch, dtype_str.split(".")[-1])
+            weight = torch.empty(shape, dtype=dtype, device=self.device)
+
+            if is_torch_xpu_available():
+                self.communicator.broadcast(weight, root=self.client_rank)
+            else:
+                self.communicator.broadcast(weight, src=self.client_rank)
+
+            weights_to_load.append((name, weight))
+
+        # Single barrier after all broadcasts
+        if is_torch_xpu_available():
+            self.communicator.barrier()
+        else:
+            self.communicator.group.barrier()
+
+        # Load weights using direct set (handles stacked params)
+        for name, weight in weights_to_load:
+            self._direct_set_weight(name, weight)
+
+    def http_load_weights(self, weights: list[tuple[str, torch.Tensor]]):
+        """Load weights received via HTTP (no NCCL needed)."""
+        for name, weight in weights:
+            self._direct_set_weight(name, weight.to(self.device))
+
+    def http_load_weight(self, **kwargs):
+        """Load a single weight received via HTTP (no NCCL needed).
+
+        Reconstructs the tensor from raw bytes since tensors don't survive
+        vLLM's multiproc IPC serialization.  Uses vLLM's ``load_weights``
+        which handles TP sharding and stacked-param packing automatically.
+        """
+        from axolotl.utils.weight_serde import decode_from_ipc
+
+        name, weight = decode_from_ipc(kwargs)
+        model = self.model_runner.model
+        model.load_weights(weights=[(name, weight)])
+
+    def http_load_weights_batch(self, params: list[dict]):
+        """Load multiple weights in a single IPC call.
+
+        Uses vLLM's ``load_weights`` which handles TP sharding automatically.
+        """
+        from axolotl.utils.weight_serde import decode_from_ipc
+
+        model = self.model_runner.model
+        weights = [decode_from_ipc(p) for p in params]
+        model.load_weights(weights=weights)
diff --git a/src/axolotl/telemetry/__init__.py b/src/axolotl/telemetry/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/telemetry/callbacks.py b/src/axolotl/telemetry/callbacks.py
new file mode 100644
index 0000000000..4f06fc75fe
--- /dev/null
+++ b/src/axolotl/telemetry/callbacks.py
@@ -0,0 +1,179 @@
+"""Trainer callbacks for reporting runtime metrics at regular intervals."""
+
+import time
+
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+from axolotl.telemetry.manager import TelemetryManager
+from axolotl.telemetry.runtime_metrics import RuntimeMetricsTracker
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+TIME_SINCE_LAST = 60
+
+
+class TelemetryCallback(TrainerCallback):
+    """
+    Trainer callback for tracking and reporting runtime metrics.
+
+    This callback tracks training progress, runtime, and memory usage,
+    sending telemetry at configurable intervals.
+    """
+
+    report_interval_steps: int = 100
+
+    def __init__(self):
+        """Initialize the metrics callback."""
+        self.tracker = RuntimeMetricsTracker()
+        self.telemetry_manager = TelemetryManager.get_instance()
+        self.current_epoch = -1
+        self.start_time = time.time()
+        self.last_report_time = None
+        self.last_report_step = 0
+
+    # pylint: disable=unused-argument
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Handle training start."""
+        self.telemetry_manager.send_event(event_type="train-start")
+
+    # pylint: disable=unused-argument
+    def on_train_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Handle training end."""
+        # Send training completion event
+        self.telemetry_manager.send_event(
+            event_type="train-end",
+            properties=self._extract_last_metrics(state)
+            | self.tracker.metrics.to_dict(),
+        )
+
+    # pylint: disable=unused-argument
+    def on_epoch_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Handle epoch start."""
+        self.current_epoch += 1
+        self.tracker.start_epoch(self.current_epoch)
+
+    # pylint: disable=unused-argument
+    def on_epoch_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Handle epoch end."""
+        self.tracker.end_epoch(self.current_epoch)
+
+    # pylint: disable=unused-argument
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Handle step end."""
+        step = state.global_step
+        self.tracker.update_step(step)
+
+        # Check if we should report metrics
+        should_report = (
+            step % self.report_interval_steps == 0
+            or step == 1  # Always report first step
+            or step - self.last_report_step >= self.report_interval_steps
+        )
+
+        if should_report:
+            current_time = time.time()
+            if self.last_report_time is not None:
+                time_since_last_report = current_time - self.last_report_time
+            else:
+                time_since_last_report = current_time - self.start_time
+            steps_since_last_report = step - self.last_report_step
+
+            # Only report if enough time has passed
+            if (
+                step == 1
+                or time_since_last_report >= TIME_SINCE_LAST
+                or steps_since_last_report >= self.report_interval_steps
+            ):
+                # Calculate steps per second for this interval
+                if time_since_last_report > 0 and steps_since_last_report > 0:
+                    steps_per_second = steps_since_last_report / time_since_last_report
+                else:
+                    steps_per_second = 0
+
+                # Update memory metrics
+                self.tracker.update_memory_metrics()
+
+                # Prepare metrics to report
+                metrics = self._extract_last_metrics(state) | {
+                    "step": step,
+                    "epoch": self.current_epoch,
+                    "progress": state.epoch,  # Fractional epoch progress
+                    "steps_per_second": steps_per_second,
+                    "elapsed_time": current_time - self.start_time,
+                    "time_since_last_report": time_since_last_report,
+                }
+
+                # Add memory metrics
+                memory_metrics = self.tracker.get_memory_metrics()
+                metrics.update({"memory": memory_metrics})
+
+                # Send telemetry
+                self.telemetry_manager.send_event(
+                    event_type="train-progress", properties=metrics
+                )
+
+                # Update last report time and step
+                self.last_report_time = current_time
+                self.last_report_step = step
+
+    def _extract_last_metrics(self, state: TrainerState) -> dict:
+        """Extract last loss, learning_rate, grad_norm, and token metrics from log history."""
+        if not state.log_history:
+            return {
+                "loss": 0,
+                "ppl": 0,
+                "learning_rate": 0,
+                "grad_norm": 0,
+                "tokens/total": 0,
+                "tokens/trainable": 0,
+                "tokens/train_per_sec_per_gpu": 0,
+            }
+
+        last_log = state.log_history[-1]
+        return {
+            "loss": last_log.get("loss", 0),
+            "ppl": last_log.get("ppl", 0),
+            "learning_rate": last_log.get("learning_rate", 0),
+            "grad_norm": last_log.get("grad_norm", 0),
+            "tokens/total": last_log.get("tokens/total", 0),
+            "tokens/trainable": last_log.get("tokens/trainable", 0),
+            "tokens/train_per_sec_per_gpu": last_log.get(
+                "tokens/train_per_sec_per_gpu", 0
+            ),
+        }
diff --git a/src/axolotl/telemetry/errors.py b/src/axolotl/telemetry/errors.py
new file mode 100644
index 0000000000..a39750bc12
--- /dev/null
+++ b/src/axolotl/telemetry/errors.py
@@ -0,0 +1,164 @@
+"""Telemetry utilities for exception and traceback information."""
+
+import os
+import re
+import traceback
+from functools import wraps
+from inspect import getmodule
+from typing import Any, Callable
+
+from axolotl.telemetry.manager import TelemetryManager
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+ERROR_HANDLED = False
+
+
+def sanitize_stack_trace(stack_trace: str) -> str:
+    """
+    Remove personal information from stack trace messages while keeping Python package codepaths.
+
+    This function identifies Python packages by looking for common patterns in virtual environment
+    and site-packages directories, preserving the package path while removing user-specific paths.
+
+    Args:
+        stack_trace: The original stack trace string.
+
+    Returns:
+        A sanitized version of the stack trace with Python package paths preserved.
+    """
+    # Split the stack trace into lines to process each file path separately
+    lines = stack_trace.split("\n")
+    sanitized_lines = []
+
+    # Regular expression to find file paths in the stack trace
+    path_pattern = re.compile(r'(?:File ")(.*?)(?:")')
+
+    # Regular expression to identify paths in site-packages or dist-packages
+    # This matches path segments like "site-packages/package_name" or "dist-packages/package_name"
+    site_packages_pattern = re.compile(
+        r"(?:site-packages|dist-packages)[/\\]([\w\-\.]+)"
+    )
+
+    # Additional common virtual environment patterns
+    venv_lib_pattern = re.compile(
+        r"(?:lib|Lib)[/\\](?:python\d+(?:\.\d+)?[/\\])?(?:site-packages|dist-packages)[/\\]([\w\-\.]+)"
+    )
+
+    for line in lines:
+        # Check if this line contains a file path
+        path_match = path_pattern.search(line)
+
+        if path_match:
+            full_path = path_match.group(1)
+            sanitized_path = ""
+
+            # Try to match site-packages pattern
+            site_packages_match = site_packages_pattern.search(full_path)
+            venv_lib_match = venv_lib_pattern.search(full_path)
+
+            if site_packages_match:
+                # Find the index where the matched pattern starts
+                idx = full_path.find("site-packages")
+                if idx == -1:
+                    idx = full_path.find("dist-packages")
+
+                # Keep from 'site-packages' onward
+                if idx >= 0:
+                    sanitized_path = full_path[idx:]
+            elif venv_lib_match:
+                # For other virtual environment patterns, find the package directory
+                match_idx = venv_lib_match.start(1)
+                if match_idx > 0:
+                    # Keep from the package name onward
+                    package_name = venv_lib_match.group(1)
+                    idx = full_path.rfind(
+                        package_name, 0, match_idx + len(package_name)
+                    )
+                    if idx >= 0:
+                        sanitized_path = full_path[idx:]
+
+            # If we couldn't identify a package pattern but path contains 'axolotl'
+            elif "axolotl" in full_path:
+                idx = full_path.rfind("axolotl")
+                if idx >= 0:
+                    sanitized_path = full_path[idx:]
+
+            # Apply the sanitization to the line
+            if sanitized_path:
+                line = line.replace(full_path, sanitized_path)
+            else:
+                # If we couldn't identify a package pattern, just keep the filename
+                filename = os.path.basename(full_path)
+                if filename:
+                    line = line.replace(full_path, filename)
+                else:
+                    line = line.replace(full_path, "")
+
+        sanitized_lines.append(line)
+
+    return "\n".join(sanitized_lines)
+
+
+def send_errors(func: Callable) -> Callable:
+    """
+    Decorator to send exception info in a function. If an exception is raised, we send
+    telemetry containing the stack trace and error message.
+
+    If an error occurs in a decorated function that is called by another decorated
+    function, we'll only send telemetry corresponding to the lower-level function.
+
+    Args:
+        func: Function to decorate.
+
+    Returns:
+        Decorated function.
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs) -> Any:
+        telemetry_manager = TelemetryManager.get_instance()
+
+        if not telemetry_manager.enabled:
+            return func(*args, **kwargs)
+
+        try:
+            return func(*args, **kwargs)
+        except Exception as exception:
+            # Only track if we're not already handling an error. This prevents us from
+            # capturing an error more than once in nested decorated function calls.
+            global ERROR_HANDLED  # pylint: disable=global-statement
+            if not ERROR_HANDLED:
+                ERROR_HANDLED = True
+
+                # Get function module path
+                module = getmodule(func)
+                module_path = (
+                    f"{module.__name__}.{func.__name__}" if module else func.__name__
+                )
+
+                # Get stack trace
+                stack_trace = "".join(
+                    traceback.format_exception(
+                        type(exception), exception, exception.__traceback__
+                    )
+                )
+                stack_trace = sanitize_stack_trace(stack_trace)
+
+                # Send error telemetry
+                telemetry_manager.send_event(
+                    event_type=f"{module_path}-error",
+                    properties={
+                        "exception": str(exception),
+                        "stack_trace": stack_trace,
+                    },
+                )
+
+                LOG.error(
+                    f"Error captured in telemetry. Run ID: {telemetry_manager.run_id}"
+                )
+
+            raise
+
+    return wrapper
diff --git a/src/axolotl/telemetry/manager.py b/src/axolotl/telemetry/manager.py
new file mode 100644
index 0000000000..90f0e6dae7
--- /dev/null
+++ b/src/axolotl/telemetry/manager.py
@@ -0,0 +1,387 @@
+"""Telemetry manager and associated utilities."""
+
+import atexit
+import importlib
+import os
+import platform
+import uuid
+from pathlib import Path
+from typing import Any
+
+import posthog
+import psutil
+import torch
+import yaml
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+POSTHOG_HOST = "https://app.posthog.com"
+POSTHOG_WRITE_KEY = "phc_1kUR0o04oJKKTTeSsIz2Mfm5mpiVsQEf2WOlzljMD7y"
+
+WHITELIST_PATH = str(Path(__file__).parent / "whitelist.yaml")
+
+# NOTE: Need to keep these up to date with any config schema changes
+FIELDS_TO_REDACT = {
+    "base_model",
+    "tokenizer_config",
+    "base_model_config",
+    "pretraining_dataset",  # NOTE: this field may be a string or a dictionary
+    "resume_from_checkpoint",
+    "hub_model_id",
+}
+PREFIXES_TO_REDACT = {"wandb_", "comet_", "mlflow_", "gradio_", "trackio_", "swanlab_"}
+PATH_INDICATORS = {"path", "dir", "data_files"}
+
+# pylint: disable=duplicate-code
+RELEVANT_PACKAGES = {
+    "torch",
+    "transformers",
+    "trl",
+    "datasets",
+    "peft",
+    "bitsandbytes",
+    "accelerate",
+    "optimum",
+    "deepspeed",
+    "ray",
+    "axolotl",
+    "triton",
+    "mamba-ssm",
+    "flash-attn",
+    "xformers",
+    "autoawq",
+    "tokenizers",
+    "sentencepiece",
+    "torchao",
+    "lm_eval",
+}
+
+
+def is_main_process() -> bool:
+    """
+    Check whether we're running in the main process.
+
+    Note:
+        We're using this function instead of `torch.utils.distributed.is_main_process`
+        causes issues with DeepSpeed world_size since. This function avoids that issue
+        by checking env vars that are set by various launchers.
+
+    Returns:
+        Whether we're running in the main process.
+    """
+    # If PyTorch distributed is already initialized, use it
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_rank() == 0
+
+    # Otherwise check environment variables for global rank
+    # NOTE: need to verify this in SLURM / OpenMPI environments
+    global_rank = int(
+        os.environ.get(
+            "RANK",
+            os.environ.get(
+                "GLOBAL_RANK",
+                os.environ.get(
+                    "SLURM_PROCID",
+                    os.environ.get(
+                        "OMPI_COMM_WORLD_RANK",
+                        "0",
+                    ),
+                ),
+            ),
+        )
+    )
+
+    return global_rank == 0
+
+
+class TelemetryManager:
+    """Manages telemetry collection and transmission"""
+
+    _instance = None
+    _initialized = False
+
+    def __new__(cls):
+        """
+        Telemetry manager constructor. Creates the singleton instance of this class if
+        it doesn't already exist.
+        """
+        if cls._instance is None:
+            cls._instance = super(TelemetryManager, cls).__new__(cls)
+            cls._instance._initialized = False
+
+        return cls._instance
+
+    def __init__(self):
+        """Telemetry manager initializer"""
+        if self._initialized:
+            return
+
+        self.enabled = self._check_telemetry_enabled()
+
+        if self.enabled:
+            self.run_id = str(uuid.uuid4())
+            self.whitelist = self._load_whitelist()
+
+            try:
+                self.system_info = self._get_system_info()
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                LOG.warning(f"Error during system info collection: {e}")
+                self.system_info = None
+
+            self._init_posthog()
+
+            # Register shutdown method to flush posthog telemetry
+            atexit.register(self.shutdown)
+
+        self._initialized = True
+
+    @classmethod
+    def get_instance(cls) -> "TelemetryManager":
+        if cls._instance is None:
+            cls._instance = TelemetryManager()
+
+        return cls._instance
+
+    def _check_telemetry_enabled(self) -> bool:
+        """
+        Check if telemetry is enabled based on environment variables. We also check
+        whether this is the main process (for the distributed setting and to avoid
+        sending duplicate PostHog events per GPU).
+
+        Note: This is enabled by default on an opt-out basis. Set
+        `AXOLOTL_DO_NOT_TRACK=1` to disable telemetry. For more details, see
+        https://axolotl-ai-cloud.github.io/axolotl/docs/telemetry.html.
+
+        Returns:
+            Boolean denoting whether telemetry is enabled or not.
+        """
+        # Only rank 0 will send telemetry
+        if not is_main_process():
+            return False
+
+        def is_truthy_env(var_name: str) -> bool:
+            value = os.getenv(var_name)
+            if value is None:
+                return False
+            return value.strip().lower() in ("1", "true")
+
+        # Telemetry is enabled by default unless either opt-out var is set
+        return not (
+            is_truthy_env("AXOLOTL_DO_NOT_TRACK") or is_truthy_env("DO_NOT_TRACK")
+        )
+
+    def _load_whitelist(self) -> dict:
+        """Load HuggingFace Hub organization whitelist"""
+        try:
+            with open(WHITELIST_PATH, encoding="utf-8") as f:
+                whitelist = yaml.safe_load(f)
+        except (OSError, yaml.YAMLError) as e:
+            # A missing/unreadable whitelist must never break `import axolotl`.
+            # Empty whitelist => nothing is whitelisted => all orgs get redacted.
+            LOG.warning(f"Could not load telemetry whitelist ({e}); redacting all orgs")
+            return {"organizations": set()}
+
+        # Send org strings to lowercase since model names are case insensitive
+        whitelist["organizations"] = {org.lower() for org in whitelist["organizations"]}
+
+        return whitelist
+
+    def _is_whitelisted(self, value: str) -> bool:
+        """
+        Check if model / dataset / etc. org is in whitelist.
+
+        Args:
+            value: Value for one of `axolotl.telemetry.manager.FIELDS_WITH_ORGS`
+                ("base_model", etc.).
+
+        Returns:
+            Boolean indicating whitelist membership.
+        """
+        # NOTE: This membership-checking logic can be improved.
+        # What happens when a local model path matches a whitelisted org?
+        parts = value.split("/")
+        if len(parts) < 2:
+            return False
+        org = parts[0]
+        whitelisted = org.lower() in self.whitelist["organizations"]
+
+        return whitelisted
+
+    def _init_posthog(self):
+        """Initialize PostHog client"""
+        posthog.api_key = POSTHOG_WRITE_KEY
+        posthog.project_api_key = POSTHOG_WRITE_KEY
+        posthog.host = POSTHOG_HOST
+
+    def _redact_paths(self, properties: dict[str, Any]) -> dict[str, Any]:
+        """
+        Redact properties to remove any paths, so as to avoid inadvertently collecting
+        private or personally identifiable information (PII). We also remove
+        information related to Wandb, MLflow, etc. configuration.
+
+        Args:
+            properties: Dictionary of properties to redact.
+
+        Returns:
+            Properties dictionary with redaction applied.
+        """
+        if not properties:
+            return {}
+
+        def redact_value(value: Any, key: str = "") -> Any:
+            """Recursively sanitize values, redacting those with path-like keys"""
+            if isinstance(key, str) and isinstance(value, str):
+                # Other redaction special cases
+                if (
+                    key in FIELDS_TO_REDACT
+                    or any(prefix in key for prefix in PREFIXES_TO_REDACT)
+                    or any(indicator in key.lower() for indicator in PATH_INDICATORS)
+                ):
+                    # Fields with whitelisted orgs don't need to be redacted
+                    if not self._is_whitelisted(value):
+                        return "[REDACTED]"
+
+            # Handle nested values
+            if isinstance(value, dict):
+                return {k: redact_value(v, k) for k, v in value.items()}
+            if isinstance(value, list):
+                return [redact_value(item) for item in value]
+
+            return value
+
+        # Create new dict with redacted values
+        redacted = {k: redact_value(v, k) for k, v in properties.items()}
+
+        return redacted
+
+    def _get_system_info(self) -> dict[str, Any]:
+        """Collect system information for various hardware accelerators"""
+        gpu_info = []
+        accelerator_type = "none"
+
+        # NVIDIA GPUs
+        if torch.cuda.is_available():
+            accelerator_type = "cuda"
+            for i in range(torch.cuda.device_count()):
+                gpu_info.append(
+                    {
+                        "name": torch.cuda.get_device_name(i),
+                        "memory": torch.cuda.get_device_properties(i).total_memory,
+                    }
+                )
+
+        # AMD GPUs
+        elif hasattr(torch, "hip") and torch.hip.is_available():
+            accelerator_type = "hip"
+            for i in range(torch.hip.device_count()):
+                gpu_info.append(
+                    {
+                        "name": torch.hip.get_device_name(i),
+                        "memory": (
+                            torch.hip.get_device_properties(i).total_memory
+                            if hasattr(torch.hip, "get_device_properties")
+                            else None
+                        ),
+                    }
+                )
+
+        # Apple Silicon
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            accelerator_type = "mps"
+            gpu_info.append(
+                {
+                    "name": "Apple Silicon",
+                    # NOTE: this is memory allocated to this process, not total memory
+                    "memory": torch.mps.driver_allocated_memory(),
+                }
+            )
+
+        # Intel GPUs
+        elif hasattr(torch, "xpu") and torch.xpu.is_available():
+            accelerator_type = "xpu"
+            for i in range(torch.xpu.device_count()):
+                memory = None
+                if hasattr(torch.xpu, "get_device_properties"):
+                    memory = torch.xpu.get_device_properties(i).total_memory
+
+                gpu_info.append(
+                    {
+                        "name": torch.xpu.get_device_name(i),
+                        "memory": memory,
+                    }
+                )
+
+        # NPUs
+        elif hasattr(torch, "npu") and torch.npu.is_available():
+            accelerator_type = "npu"
+            for i in range(torch.npu.device_count()):
+                memory = None
+                if hasattr(torch.npu, "get_device_properties"):
+                    memory = torch.npu.get_device_properties(i).total_memory
+
+                gpu_info.append(
+                    {
+                        "name": torch.npu.get_device_name(i),
+                        "memory": memory,
+                    }
+                )
+
+        # Get relevant package versions
+        installed_packages = {}
+        for package in RELEVANT_PACKAGES:
+            try:
+                version = importlib.metadata.version(package)
+                installed_packages[f"{package}_version"] = version
+            except importlib.metadata.PackageNotFoundError:
+                pass
+
+        return {
+            "os": platform.system(),
+            "python_version": platform.python_version(),
+            "cpu_count": psutil.cpu_count(),
+            "memory_total": psutil.virtual_memory().total,
+            "accelerator_type": accelerator_type,
+            "accelerator_count": len(gpu_info),
+            "accelerator_info": gpu_info,
+            **installed_packages,
+        }
+
+    def send_event(self, event_type: str, properties: dict[str, Any] | None = None):
+        """Send a telemetry event"""
+        if not self.enabled:
+            return
+
+        if properties is None:
+            properties = {}
+
+        # Sanitize properties to remove PII
+        properties = self._redact_paths(properties)
+
+        # Wrap PostHog errors in try / except to not raise errors during Axolotl usage
+        try:
+            # Send event via PostHog
+            posthog.capture(
+                distinct_id=self.run_id,
+                event=event_type,
+                properties=properties,
+                disable_geoip=True,
+            )
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            LOG.warning(f"Failed to send telemetry event: {e}")
+
+        # Additionally, send system info telemetry when loading config.
+        # NOTE: Is this the best place for this?
+        if event_type == "config-loaded":
+            self.send_system_info()
+
+    def send_system_info(self):
+        """Helper method for sending system info"""
+        if self.system_info is not None:
+            self.send_event(event_type="system-info", properties=self.system_info)
+
+    def shutdown(self):
+        """Ensure all queued events are processed before shutdown"""
+        if self.enabled:
+            posthog.shutdown()
diff --git a/src/axolotl/telemetry/runtime_metrics.py b/src/axolotl/telemetry/runtime_metrics.py
new file mode 100644
index 0000000000..5272f05de0
--- /dev/null
+++ b/src/axolotl/telemetry/runtime_metrics.py
@@ -0,0 +1,210 @@
+"""Telemetry utilities for runtime and memory metrics."""
+
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+import psutil
+import torch
+
+from axolotl.telemetry.manager import TelemetryManager
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+@dataclass
+class RuntimeMetrics:
+    """Container for runtime metrics to be tracked throughout training."""
+
+    # Timing metrics
+    start_time: float
+    epoch_start_times: dict[int, float] = field(init=False)
+    epoch_end_times: dict[int, float] = field(init=False)
+
+    # Memory metrics
+    peak_cpu_memory: int = 0
+    peak_gpu_memory: dict[int, int] = field(init=False)
+
+    # Progress metrics
+    total_steps: int = 0
+    current_epoch: int = 0
+    current_step: int = 0
+
+    def __post_init__(self):
+        """Initialize empty metric mappings."""
+        self.epoch_start_times = {}
+        self.epoch_end_times = {}
+        self.peak_gpu_memory = {}
+
+    @property
+    def elapsed_time(self) -> float:
+        """Calculate total elapsed time in seconds."""
+        return time.time() - self.start_time
+
+    def epoch_time(self, epoch: int) -> float | None:
+        """Calculate time taken for a specific epoch in seconds."""
+        if epoch in self.epoch_start_times and epoch in self.epoch_end_times:
+            return self.epoch_end_times[epoch] - self.epoch_start_times[epoch]
+
+        return None
+
+    def average_epoch_time(self) -> float | None:
+        """Calculate average time per epoch in seconds."""
+        completed_epochs = [
+            epoch for epoch in self.epoch_start_times if epoch in self.epoch_end_times
+        ]
+        if not completed_epochs:
+            return None
+
+        total_time = 0.0
+        for epoch in completed_epochs:
+            epoch_time = self.epoch_time(epoch)
+            if epoch_time is not None:  # Check to avoid mypy warning
+                total_time += epoch_time
+
+        return total_time / len(completed_epochs)
+
+    def steps_per_second(self) -> float | None:
+        """Calculate average steps per second across all training."""
+        if self.total_steps == 0 or self.elapsed_time == 0:
+            return None
+
+        return self.total_steps / self.elapsed_time
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert metrics to a dictionary for telemetry reporting."""
+        metrics = {
+            "total_time_seconds": self.elapsed_time,
+            "total_steps": self.total_steps,
+            "steps_per_second": self.steps_per_second(),
+            "epochs_completed": len(
+                [
+                    epoch
+                    for epoch in self.epoch_start_times
+                    if epoch in self.epoch_end_times
+                ]
+            ),
+            "peak_cpu_memory_bytes": self.peak_cpu_memory,
+        }
+
+        # Add per-epoch timing if available
+        epoch_times: dict[str, float] = {}
+        for epoch in sorted(self.epoch_end_times.keys()):
+            time_taken = self.epoch_time(epoch)
+            if time_taken is not None:
+                epoch_times[f"epoch_{epoch}_seconds"] = time_taken
+
+        if epoch_times:
+            metrics["epoch_times"] = epoch_times  # type: ignore
+            metrics["average_epoch_time_seconds"] = self.average_epoch_time()
+
+        # Add GPU memory metrics if available
+        if self.peak_gpu_memory:
+            gpu_metrics: dict[str, int] = {}
+            for gpu_id, memory in self.peak_gpu_memory.items():
+                gpu_metrics[f"gpu_{gpu_id}_peak_memory_bytes"] = memory
+            metrics["gpu_memory"] = gpu_metrics  # type: ignore
+
+        return metrics
+
+
+class RuntimeMetricsTracker:
+    """Tracker for runtime metrics during training."""
+
+    update_interval = 100
+
+    def __init__(self):
+        """Initialize the runtime metrics tracker."""
+        self.metrics = RuntimeMetrics(start_time=time.time())
+        self.telemetry_manager = TelemetryManager.get_instance()
+        self._process = psutil.Process()
+
+    def start_epoch(self, epoch: int):
+        """Record the start of a new epoch."""
+        self.metrics.current_epoch = epoch
+        self.metrics.epoch_start_times[epoch] = time.time()
+        self.update_memory_metrics()
+
+    def end_epoch(self, epoch: int):
+        """Record the end of an epoch."""
+        self.metrics.epoch_end_times[epoch] = time.time()
+
+    def update_step(self, step: int):
+        """Update the current step count."""
+        self.metrics.current_step = step
+        self.metrics.total_steps += 1
+
+        # Periodically update memory metrics
+        if step % self.update_interval == 0:
+            self.update_memory_metrics()
+
+    def _get_allocated_memory(self) -> dict[int, int]:
+        """
+        Helper function for getting accelerator-agnostic allocated memory.
+
+        Returns:
+            A dictionary mapping device IDs to allocated memory in bytes
+        """
+        memory_used: dict[int, int] = {}
+
+        # NVIDIA GPUs
+        if torch.cuda.is_available():
+            for i in range(torch.cuda.device_count()):
+                memory_used[i] = torch.cuda.memory_allocated(i)
+
+        # AMD GPUs
+        elif hasattr(torch, "hip") and torch.hip.is_available():
+            for i in range(torch.hip.device_count()):
+                if hasattr(torch.hip, "memory_allocated"):
+                    memory_used[i] = torch.hip.memory_allocated(i)
+
+        # Apple Silicon
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            # MPS doesn't have per-device memory stats since there's only one device
+            if hasattr(torch.mps, "current_allocated_memory"):
+                memory_used[0] = torch.mps.current_allocated_memory()
+
+        # Intel GPUs
+        elif hasattr(torch, "xpu") and torch.xpu.is_available():
+            for i in range(torch.xpu.device_count()):
+                if hasattr(torch.xpu, "memory_allocated"):
+                    memory_used[i] = torch.xpu.memory_allocated(i)
+
+        # NPUs
+        elif hasattr(torch, "npu") and torch.npu.is_available():
+            for i in range(torch.npu.device_count()):
+                if hasattr(torch.npu, "memory_allocated"):
+                    memory_used[i] = torch.npu.memory_allocated(i)
+
+        return memory_used
+
+    def update_memory_metrics(self):
+        """Update peak memory usage metrics."""
+        # CPU memory
+        cpu_memory = self._process.memory_info().rss
+        self.metrics.peak_cpu_memory = max(self.metrics.peak_cpu_memory, cpu_memory)
+
+        # GPU memory (if available)
+        memory_used = self._get_allocated_memory()
+        for i, memory in memory_used.items():
+            self.metrics.peak_gpu_memory[i] = max(
+                self.metrics.peak_gpu_memory.get(i, 0), memory
+            )
+
+    def get_memory_metrics(self) -> dict[str, Any]:
+        """Get the current memory metrics as a dictionary."""
+        memory_metrics = {
+            "cpu_memory_bytes": self._process.memory_info().rss,
+            "peak_cpu_memory_bytes": self.metrics.peak_cpu_memory,
+        }
+
+        # GPU memory (if available)
+        memory_used = self._get_allocated_memory()
+        for i, memory in memory_used.items():
+            memory_metrics[f"gpu_{i}_memory_bytes"] = memory
+            memory_metrics[f"gpu_{i}_peak_memory_bytes"] = (
+                self.metrics.peak_gpu_memory.get(i, 0)
+            )
+
+        return memory_metrics
diff --git a/src/axolotl/telemetry/whitelist.yaml b/src/axolotl/telemetry/whitelist.yaml
new file mode 100644
index 0000000000..c75ee0fec8
--- /dev/null
+++ b/src/axolotl/telemetry/whitelist.yaml
@@ -0,0 +1,40 @@
+organizations:
+  - "axolotl-ai-co"
+  - "meta-llama"
+  - "huggingface"
+  - "nvidia"
+  - "facebook"
+  - "google"
+  - "microsoft"
+  - "deepseek-ai"
+  - "HuggingFaceTB"
+  - "mistralai"
+  - "Qwen"
+  - "unsloth"
+  - "NousResearch"
+  - "allenai"
+  - "amd"
+  - "tiiuae"
+  - "tencent"
+  - "zai-org"
+  - "openai"
+  - "ibm-granite"
+  - "arcee-ai"
+  - "swiss-ai"
+  - "CohereForAI"
+  - "deepcogito"
+  - "THUDM"
+  - "ai21labs"
+  - "LiquidAI"
+  - "canopylabs"
+  - "state-spaces"
+  - "mistral-community"
+  - "llava-hf"
+  - "ByteDance-Seed"
+  - "ACE-Step"
+  - "openbmb"
+  - "MiniMaxAI"
+  - "stepfun-ai"
+  - "internlm"
+  - "katanemo"
+  - "XiaomiMiMo"
diff --git a/src/axolotl/train.py b/src/axolotl/train.py
index 32bcbc1d0a..63247d5f3e 100644
--- a/src/axolotl/train.py
+++ b/src/axolotl/train.py
@@ -1,141 +1,176 @@
 """Prepare and train a model on a dataset. Can also infer from a model or merge lora"""
 
+from __future__ import annotations
+
+import importlib
+import inspect
+import json
 import os
+import shutil
 import signal
 import sys
+import typing
 import weakref
-from dataclasses import dataclass
+from collections import OrderedDict
+from contextlib import ExitStack
 from pathlib import Path
-from typing import Optional, Tuple, Union
+from typing import Any, Dict
 
 import torch
 import transformers.modelcard
-from accelerate import Accelerator
-from accelerate.logging import get_logger
 from datasets import Dataset
-from peft import PeftModel
-from pkg_resources import get_distribution  # type: ignore
-from transformers import PreTrainedModel, PreTrainedTokenizer
+from huggingface_hub.errors import OfflineModeIsEnabled
+from peft import PeftConfig, PeftModel
+from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin
 from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
-
-from axolotl.common.cli import TrainerCliArgs
-from axolotl.logging_config import configure_logging
+from transformers.trainer import Trainer
+
+from axolotl.common.datasets import TrainDatasetMeta
+from axolotl.contribs.lgpl import (  # pylint: disable = no-name-in-module
+    fix_untrained_tokens,
+)
+from axolotl.integrations.base import PluginManager
+from axolotl.loaders import ModelLoader, load_processor, load_tokenizer
+from axolotl.telemetry.errors import send_errors
+from axolotl.telemetry.manager import TelemetryManager
+from axolotl.utils.ctx_managers.sequence_parallel import SequenceParallelContextManager
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.freeze import freeze_layers_except
-from axolotl.utils.models import load_model, load_tokenizer
+from axolotl.utils.distributed import cleanup_distributed
+from axolotl.utils.freeze import freeze_layers_except, freeze_mm_modules
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.enums import RLType
+from axolotl.utils.train import determine_last_checkpoint
 from axolotl.utils.trainer import setup_trainer
 
-try:
-    from optimum.bettertransformer import BetterTransformer
-except ImportError:
-    BetterTransformer = None
-
-project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
-src_dir = os.path.join(project_root, "src")
-sys.path.insert(0, src_dir)
+if typing.TYPE_CHECKING:
+    from axolotl.core.builders import HFCausalTrainerBuilder, HFRLTrainerBuilder
 
-configure_logging()
-LOG = get_logger("axolotl.train")
+LOG = get_logger(__name__)
 
+TELEMETRY_MANAGER = TelemetryManager.get_instance()
+PLUGIN_MANAGER = PluginManager.get_instance()
 
-@dataclass
-class TrainDatasetMeta:
-    """
-    dataclass to capture the dataset specific options for training
-    """
 
-    train_dataset: Dataset
-    eval_dataset: Optional[Dataset] = None
-    total_num_steps: Optional[int] = None
+def setup_model_and_tokenizer(
+    cfg: DictDefault,
+) -> tuple[
+    PreTrainedModel, PreTrainedTokenizer, PeftConfig | None, ProcessorMixin | None
+]:
+    """Load the tokenizer, processor (for multimodal models), and model based on
+    configuration.
 
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
 
-def train(
-    *, cfg: DictDefault, cli_args: TrainerCliArgs, dataset_meta: TrainDatasetMeta
-) -> Tuple[Union[PeftModel, PreTrainedModel], PreTrainedTokenizer]:
-    # load the tokenizer first
+    Returns:
+        Tuple containing model, tokenizer, `peft_config` (if LoRA / QLoRA, else
+            `None`), and processor (if multimodal, else `None`).
+    """
+    # Load tokenizer
     LOG.debug(
         f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}",
-        main_process_only=True,
     )
     tokenizer = load_tokenizer(cfg)
 
-    train_dataset = dataset_meta.train_dataset
-    eval_dataset = dataset_meta.eval_dataset
-    total_num_steps = dataset_meta.total_num_steps
+    # Load processor for multimodal models if needed
+    processor = None
+    if cfg.is_multimodal:
+        processor = load_processor(cfg, tokenizer)
+
+    # Load the model
+    LOG.debug("Loading model")
+
+    model_loader = ModelLoader(cfg, tokenizer, processor=processor)
+    model, peft_config = model_loader.load()
+    if getattr(model, "generation_config", None) is not None:
+        model.generation_config.do_sample = True
+
+    model_properties = model.config.to_dict()
+    try:
+        model_properties["num_parameters"] = model.num_parameters()
+    except Exception:  # pylint: disable=broad-exception-caught
+        model_properties["num_parameters"] = sum(p.numel() for p in model.parameters())
+    # if the num_parameters is less than 2B, let's round to nearest 100M, else round to nearest 1B
+    if model_properties["num_parameters"] < 2e9:
+        model_properties["num_parameters_est"] = (
+            f"{round(model_properties['num_parameters'] / 1e8) * 100}M"
+        )
+    else:
+        model_properties["num_parameters_est"] = (
+            f"{round(model_properties['num_parameters'] / 1e9)}B"
+        )
+    TELEMETRY_MANAGER.send_event(event_type="model-load", properties=model_properties)
+    if peft_config:
+        TELEMETRY_MANAGER.send_event(
+            event_type="peft-config-load", properties=peft_config.to_dict()
+        )
+
+    # Apply freezing if specified
+    if cfg.unfrozen_parameters:
+        freeze_layers_except(model, cfg.unfrozen_parameters)
+        if any(
+            any(embed in param for embed in ["lm_head", "embed_tokens"])
+            for param in cfg.unfrozen_parameters
+        ):
+            model.enable_input_require_grads()
+
+    # Freeze multimodal modules for text-only training of multimodal models
+    if cfg.freeze_mm_modules:
+        freeze_mm_modules(model)
+
+    return model, tokenizer, peft_config, processor
 
-    if cfg.resume_from_checkpoint is None and cfg.auto_resume_from_checkpoints:
-        possible_checkpoints = [
-            str(cp) for cp in Path(cfg.output_dir).glob("checkpoint-*")
-        ]
-        if len(possible_checkpoints) > 0:
-            sorted_paths = sorted(
-                possible_checkpoints,
-                key=lambda path: int(path.split("-")[-1]),
-            )
-            cfg.resume_from_checkpoint = sorted_paths[-1]
-            LOG.info(
-                f"Using Auto-resume functionality to start with checkpoint at {cfg.resume_from_checkpoint}"
-            )
-    resume_from_checkpoint = cfg.resume_from_checkpoint
-
-    # Load the model and tokenizer
-    msg = "loading model"
-    if cfg.adapter:
-        msg += " and peft_config..."
-    LOG.debug(msg)
-    # we wait unitl the last possible moment to setup Accelerator
-    Accelerator()
-    model, peft_config = load_model(cfg, tokenizer, inference=cli_args.inference)
-    model.generation_config.do_sample = True
 
+def setup_reference_model(
+    cfg: DictDefault, tokenizer: PreTrainedTokenizer
+) -> PreTrainedModel | None:
+    """
+    Set up the reference model for RL training if needed.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        tokenizer: The tokenizer to use for the reference model.
+
+    Returns:
+        Reference model if needed for RL training, `None` otherwise.
+    """
     model_ref = None
-    if cfg.rl and cfg.rl != "orpo":
+    if cfg.rl and cfg.rl != RLType.ORPO:
         if cfg.adapter and not cfg.rl_adapter_ref_model:
             # use built-in trl autounwrap
             LOG.debug("Passing model_ref: None to RL trainer")
             model_ref = None  # explicit setting to None
         else:
+            reference_model: bool = True
+            trl_cfg = getattr(cfg, "trl", None)
+            if (
+                cfg.rl in {RLType.GRPO, RLType.EBFT}
+                and getattr(trl_cfg, "beta", 0) == 0
+            ):
+                reference_model = False
             # load the model again for model_ref/baseline
-            model_ref, _ = load_model(
-                cfg, tokenizer, inference=cli_args.inference, reference_model=True
-            )
+            model_loader = ModelLoader(cfg, tokenizer, reference_model=reference_model)
+            model_ref, _ = model_loader.load()
+    return model_ref
 
-    safe_serialization = cfg.save_safetensors is True
 
-    if cfg.unfrozen_parameters:
-        freeze_layers_except(model, cfg.unfrozen_parameters)
-
-    trainer = setup_trainer(
-        cfg,
-        train_dataset,
-        eval_dataset,
-        (model, model_ref, peft_config),
-        tokenizer,
-        total_num_steps,
-    )
-
-    # go ahead and presave, so we have the adapter config available to inspect
-    if peft_config:
-        LOG.info(f"Pre-saving adapter config to {cfg.output_dir}")
-        peft_config.save_pretrained(cfg.output_dir)
-    # additionally presave the tokenizer and model configs
-    if not Path(cfg.output_dir).is_dir():
-        os.makedirs(cfg.output_dir, exist_ok=True)
-    tokenizer.save_pretrained(str(Path(cfg.output_dir)))
-    if hasattr(model, "config"):
-        model.config.save_pretrained(str(Path(cfg.output_dir)))
+def setup_signal_handler(cfg: DictDefault, model: PreTrainedModel):
+    """
+    Set up signal handler for graceful termination.
 
-    # In case we want to stop early with ctrl+c, this is a nice to have to save the pretrained model
-    if cfg.local_rank == 0:
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        model: The model to save on termination
+    """
+    # ray workers don't have access to this signal
+    if cfg.local_rank == 0 and not cfg.use_ray:
 
         def terminate_handler(_, __, model_weakref):
             if model_weakref() is not None:
                 _model = model_weakref()
-                if cfg.flash_optimum and BetterTransformer:
-                    _model = BetterTransformer.reverse(_model)
-                _model.save_pretrained(
-                    cfg.output_dir, safe_serialization=safe_serialization
-                )
+                _model.save_pretrained(cfg.output_dir)
+
+            cleanup_distributed()
             sys.exit(0)
 
         _model_weakref = weakref.ref(model)
@@ -144,105 +179,490 @@ def terminate_handler(_, __, model_weakref):
             lambda signum, frame: terminate_handler(signum, frame, _model_weakref),
         )
 
-    badge_markdown = """[<img src="https://raw.githubusercontent.com/OpenAccess-AI-Collective/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/OpenAccess-AI-Collective/axolotl)"""
-    transformers.modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}"
 
-    if getattr(cfg, "axolotl_config_path"):
-        raw_axolotl_cfg = Path(cfg.axolotl_config_path)
-        version = get_distribution("axolotl").version
-        if raw_axolotl_cfg.is_file():
-            transformers.modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n<details><summary>See axolotl config</summary>\n\naxolotl version: `{version}`\n```yaml\n{raw_axolotl_cfg.read_text(encoding='utf-8')}\n```\n\n</details><br>\n"
+def execute_training(
+    cfg: DictDefault, trainer: Any, resume_from_checkpoint: str | None
+):
+    """
+    Execute the training process with appropriate SDP kernel configurations.
 
-    LOG.info("Starting trainer...")
-    if cfg.group_by_length:
-        LOG.info("hang tight... sorting dataset for group_by_length")
-
-    pretrain_hooks(cfg, trainer)
-    if cfg.flash_optimum:
-        with torch.backends.cuda.sdp_kernel(
-            # TODO configure these from the YAML w/ sdp_kernel_kwargs: ...
-            enable_flash=True,
-            enable_math=True,
-            enable_mem_efficient=True,
-        ):
-            trainer.train(resume_from_checkpoint=resume_from_checkpoint)
-    else:
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        trainer: The configured trainer object.
+        resume_from_checkpoint: Path to checkpoint to resume from, if applicable.
+    """
+    with ExitStack() as stack:
+        # Define the context managers to use
+        if cfg.flash_optimum:
+            stack.enter_context(
+                torch.backends.cuda.sdp_kernel(
+                    enable_flash=True,
+                    enable_math=True,
+                    enable_mem_efficient=True,
+                )
+            )
+
+        if cfg.context_parallel_size > 1:
+            models = [trainer.model]
+            if hasattr(trainer, "ref_model") and trainer.ref_model:
+                models.append(trainer.ref_model)
+
+            stack.enter_context(
+                SequenceParallelContextManager(
+                    models=models,
+                    context_parallel_size=cfg.context_parallel_size,
+                    gradient_accumulation_steps=cfg.gradient_accumulation_steps,
+                    ring_attn_func=cfg.ring_attn_func,
+                    heads_k_stride=cfg.heads_k_stride,
+                    gather_outputs=cfg.rl in {RLType.GRPO, RLType.EBFT},
+                    device_mesh=trainer.accelerator.torch_device_mesh,
+                )
+            )
+
+        # TODO: disabling for now as not compatible with FSDP2 + torchao low bit optimizers
+        # if cfg.bf16:
+        #     torch.set_default_dtype(torch.bfloat16)
+
+        LOG.info("Starting trainer...")
         trainer.train(resume_from_checkpoint=resume_from_checkpoint)
-    post_train_hooks(cfg, trainer)
 
-    LOG.info(f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}")
+        PLUGIN_MANAGER.post_train(cfg, trainer.model)
+
+
+def _rename_fsdp_merged_to_adapter(merged_dir: Path):
+    """Rename model*.safetensors files to adapter_model* in place.
+
+    Also rewrites the index JSON weight_map if sharded output was produced.
+    """
+    for file in sorted(merged_dir.iterdir()):
+        if file.name.startswith("model") and ".safetensors" in file.name:
+            file.rename(merged_dir / file.name.replace("model", "adapter_model", 1))
+
+    index = merged_dir / "adapter_model.safetensors.index.json"
+    if index.exists():
+        data = json.loads(index.read_text(encoding="utf-8"))
+        if "weight_map" in data:
+            data["weight_map"] = {
+                k: v.replace("model", "adapter_model", 1)
+                for k, v in data["weight_map"].items()
+            }
+        index.write_text(
+            json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8"
+        )
+
 
-    # post training
+def save_trained_model(
+    cfg: DictDefault,
+    trainer: Any,
+    model: PreTrainedModel,
+):
+    """
+    Save the trained model according to configuration and training setup.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        trainer: The trainer object.
+        model: The trained model to save.
+    """
+    LOG.info(f"Training completed! Saving trained model to {cfg.output_dir}.")
+
+    # Post training module hooks
     for name, module in model.named_modules():
         if hasattr(module, "_post_training"):
-            module._post_training(model, name)  # pylint: disable=protected-access
+            module._post_training(model, name)
 
-    if trainer.is_fsdp_enabled:
-        trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
-        LOG.info("Set FSDP state dict type to FULL_STATE_DICT for saving.")
+    # handle QAT
+    if cfg.qat:
+        from axolotl.utils.quantization import convert_qat_model
 
-    if cfg.relora_steps:
-        if cfg.adapter == "lora" and not (cfg.load_in_4bit or cfg.load_in_8bit):
+        convert_qat_model(
+            model,
+            quantize_embedding=cfg.qat.quantize_embedding,
+        )
+        LOG.info(
+            "QAT usage note: please ensure you quantize your model fine-tuned using QAT by running `axolotl quantize`"
+            " with the same config which you used for training."
+        )
+    # Handle ReLoRA early return case
+    if cfg.relora:
+        if hasattr(model, "merge_and_unload") and not (
+            cfg.load_in_4bit or cfg.load_in_8bit
+        ):
             model = model.merge_and_unload()
         else:
             # final model weights have already been saved by `ReLoRACallback.on_train_end`
-            return model, tokenizer
-
-    # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
-    # only save on rank 0, otherwise it corrupts output on multi-GPU when multiple processes attempt to write the same file
-    if cfg.fsdp:
-        trainer.save_model(cfg.output_dir)
+            return
+
+    # EP-sharded expert LoRA: the adapter is split across the EP axis, so the normal
+    # FSDP/PEFT save would persist only the local rank's experts. Gather across EP and
+    # write a complete adapter.
+    if cfg.adapter and (getattr(cfg, "expert_parallel_size", 1) or 1) > 1:
+        from axolotl.integrations.expert_parallel.plugin import ExpertParallelPlugin
+        from axolotl.integrations.expert_parallel.shard import save_ep_lora_adapter
+
+        ep_group = ExpertParallelPlugin._resolve_ep_group(cfg)
+        if save_ep_lora_adapter(model, cfg.output_dir, ep_group):
+            return
+
+    # FSDP2 (no EP) LoRA: the DCP sharded save fails ("Failed to validate global plan") on the
+    # frozen NVFP4 base DTensors, so gather just the adapter and write it directly.
+    if (
+        cfg.adapter
+        and (trainer.is_fsdp_enabled or cfg.fsdp_config)
+        and str(cfg.fsdp_version) == "2"
+    ):
+        from axolotl.integrations.expert_parallel.shard import save_fsdp2_lora_adapter
+
+        if save_fsdp2_lora_adapter(model, cfg.output_dir):
+            return
+
+    if trainer.is_fsdp_enabled or cfg.fsdp_config:
+        if cfg.fsdp_config or cfg.fsdp:
+            if cfg.fsdp_config.final_state_dict_type:
+                state_dict_type = cfg.fsdp_config.final_state_dict_type
+            else:
+                state_dict_type = cfg.fsdp_config.state_dict_type
+            trainer.accelerator.state.fsdp_plugin.set_state_dict_type(state_dict_type)
+        trainer.save_model(cfg.output_dir)  # only handles FULL_STATE_DICT
+        if state_dict_type == "SHARDED_STATE_DICT":
+            LOG.info(
+                "The final model was saved with a sharded state dict. Please ensure you merge "
+                "the sharded weights with `merge-sharded-fsdp-weights`."
+            )
+            checkpoint_dir = determine_last_checkpoint(cfg, update=False)
+            if (
+                not (Path(cfg.output_dir) / "model.safetensors.index.json").exists()
+                and checkpoint_dir
+            ):
+                # import here to prevent circular import
+                from axolotl.cli.merge_sharded_fsdp_weights import merge_fsdp_weights
+
+                fsdp_dir = Path(checkpoint_dir) / "pytorch_model_fsdp_0"
+                merged_path = str(Path(cfg.output_dir) / "merged")
+                merge_fsdp_weights(
+                    checkpoint_dir=str(fsdp_dir),
+                    output_path=merged_path,
+                )
+                trainer.accelerator.wait_for_everyone()
+                if trainer.accelerator.is_main_process:
+                    # FSDP checkpoints for PEFT only contain adapter weights;
+                    # rename model* → adapter_model* so it loads correctly.
+                    is_peft = cfg.adapter and not cfg.relora
+                    if is_peft:
+                        _rename_fsdp_merged_to_adapter(Path(merged_path))
+                    for merged_file in Path(merged_path).iterdir():
+                        dest = Path(cfg.output_dir) / merged_file.name
+                        if dest.exists():
+                            dest.unlink()
+                        shutil.move(str(merged_file), dest)
+                    shutil.rmtree(merged_path)
+        # TODO(wing):see https://github.com/huggingface/transformers/pull/40207
+        # cleanup the FSDP prefix in the model config.json
+        if trainer.accelerator.is_main_process:
+            with open(
+                Path(cfg.output_dir) / "config.json", "r", encoding="utf-8"
+            ) as config_file_io:
+                # read the model config as an OrderedDict
+                config = json.load(config_file_io, object_pairs_hook=OrderedDict)
+                if config.get("architectures"):
+                    config["architectures"] = [
+                        name.lstrip("FSDP") for name in config["architectures"]
+                    ]
+            # write the updated model config back
+            with open(
+                os.path.join(cfg.output_dir, "config.json"), "w", encoding="utf-8"
+            ) as config_file_io:
+                json.dump(config, config_file_io, indent=2)
     elif cfg.deepspeed and is_deepspeed_zero3_enabled():
         # Copied over from: https://github.com/huggingface/accelerate/blob/5ae611118057232f441055f7ef9ba0b0f2b8d533/docs/source/usage_guides/deepspeed.md#saving-and-loading
         trainer.accelerator.wait_for_everyone()
-        unwrapped_model = trainer.accelerator.unwrap_model(trainer.model_wrapped)
-
-        # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
-        # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
-        # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
-        # For Zero Stages 1 and 2, models are saved as usual in the output directory.
-        # The model name saved is `pytorch_model.bin`
-        unwrapped_model.save_pretrained(
-            cfg.output_dir,
-            is_main_process=trainer.accelerator.is_main_process,
-            save_function=trainer.accelerator.save,
-            state_dict=trainer.accelerator.get_state_dict(trainer.model_wrapped),
+        trainer.save_model(cfg.output_dir)
+
+        # the trainer saved a model.safetensors file in the output directory,
+        # but it is most likely a proxy model and if so, should be deleted
+        maybe_proxy = os.path.exists(os.path.join(cfg.output_dir, "model.safetensors"))
+        maybe_sharded = os.path.exists(
+            os.path.join(cfg.output_dir, "model.safetensors.index.json")
         )
-    elif cfg.local_rank == 0:
-        if cfg.flash_optimum and BetterTransformer:
-            model = BetterTransformer.reverse(model)
 
+        if maybe_proxy and maybe_sharded:
+            LOG.info(f"Deleting {os.path.join(cfg.output_dir, 'model.safetensors')}")
+            LOG.info("This is a proxy model and should be deleted")
+            try:
+                os.remove(os.path.join(cfg.output_dir, "model.safetensors"))
+            except FileNotFoundError:
+                pass
+    elif cfg.local_rank == 0:
         if cfg.rl and cfg.adapter and not cfg.rl_adapter_ref_model:
-            trainer.model.save_pretrained(
-                cfg.output_dir, safe_serialization=safe_serialization
-            )
-        model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
+            trainer.model.save_pretrained(cfg.output_dir)
+
+        model.save_pretrained(cfg.output_dir)
+
+    if hasattr(cfg, "llmcompressor") and cfg.llmcompressor:
+        # TODO: add integration support so this can be implemented completely within the plugin
+        from axolotl.integrations.llm_compressor.utils import save_compressed_model
 
+        save_compressed_model(
+            model=model,
+            output_dir=cfg.output_dir,
+            trainer=trainer,
+            save_compressed=cfg.llmcompressor.save_compressed,
+        )
+
+    LOG.info(f"Model successfully saved to {cfg.output_dir}")
+
+
+def create_model_card(cfg: DictDefault, trainer: Trainer):
+    """
+    Create a model card for the trained model if needed.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        trainer: The trainer object with model card creation capabilities.
+    """
     if not cfg.hub_model_id:
+        # Guard since create_model_card may fail if dataset_tags is empty list
         try:
-            trainer.create_model_card(model_name=cfg.output_dir.lstrip("./"))
-        except AttributeError:
+            model_card_kwarg = {
+                "model_name": cfg.output_dir.lstrip("./")
+                .encode("utf-8")
+                .decode("utf-8")
+            }
+
+            # We check if we're using a TRL trainer; if so, `dataset_tags` is not consumed.
+            rl = cfg.rl is not None or cfg.reward_model or cfg.process_reward_model
+            if cfg.datasets is not None and not rl:
+                dataset_tags = [
+                    d["path"] for d in cfg.datasets if not Path(d["path"]).is_dir()
+                ]
+                dataset_tags = [d for d in dataset_tags if not d.startswith("https://")]
+
+                if dataset_tags:
+                    model_card_kwarg["dataset_tags"] = dataset_tags
+
+            trainer.create_model_card(**model_card_kwarg)
+        except (AttributeError, UnicodeDecodeError, OfflineModeIsEnabled):
             pass
     elif cfg.hub_model_id:
-        # defensively push to the hub to ensure the model card is updated
+        # Defensively push to the hub to ensure the model card is updated
         trainer.push_to_hub()
 
-    return model, tokenizer
+
+def save_initial_configs(
+    cfg: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    model: PreTrainedModel,
+    peft_config: PeftConfig | None,
+    processor: ProcessorMixin | None,
+):
+    """
+    Save initial configurations before training.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        tokenizer: The tokenizer to save.
+        model: The model to save configuration for.
+        peft_config: The PEFT configuration to save if applicable.
+    """
+    # Create output_dir if it doesn't already exist
+    output_dir = Path(cfg.output_dir)
+    if not output_dir.is_dir():
+        os.makedirs(cfg.output_dir, exist_ok=True)
+
+    # Pre-save adapter config so it's available to inspect
+    if peft_config:
+        LOG.info(f"Pre-saving adapter config to {cfg.output_dir}...")
+        peft_config.save_pretrained(cfg.output_dir)
+
+    # Pre-save the tokenizer and model configs
+    LOG.info(f"Pre-saving tokenizer to {cfg.output_dir}...")
+    tokenizer.save_pretrained(
+        str(Path(cfg.output_dir)), save_jinja_files=cfg.tokenizer_save_jinja_files
+    )
+    if hasattr(model, "config"):
+        LOG.info(f"Pre-saving model config to {cfg.output_dir}...")
+        model.config.save_pretrained(str(output_dir))
+
+    if processor:
+        LOG.info(f"Pre-saving processor to {cfg.output_dir}...")
+        processor.save_pretrained(str(output_dir))
 
 
-def pretrain_hooks(_cfg, _trainer):
+def setup_model_card(cfg: DictDefault):
     """
-    Run hooks right before kicking off the training
-    :param cfg:
-    :param trainer:
-    :return:
+    Set up the Axolotl badge and add the Axolotl config to the model card if available.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
     """
+    badge_markdown = """[<img src="https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/axolotl-ai-cloud/axolotl)"""
+    transformers.modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}"
 
+    if cfg.axolotl_config_path:
+        raw_axolotl_cfg = Path(cfg.axolotl_config_path)
+        version = importlib.metadata.version("axolotl")
+        if raw_axolotl_cfg.is_file():
+            transformers.modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n<details><summary>See axolotl config</summary>\n\naxolotl version: `{version}`\n```yaml\n{raw_axolotl_cfg.read_text(encoding='utf-8')}\n```\n\n</details><br>\n"
 
-def post_train_hooks(_cfg, _trainer):
+
+def handle_untrained_tokens_fix(
+    cfg: DictDefault,
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizer,
+    train_dataset: Dataset,
+):
     """
-    Run hooks right after training completes
-    :param cfg:
-    :param trainer:
-    :return:
+    Apply fixes for untrained tokens if configured.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        model: The model to apply fixes to.
+        tokenizer: The tokenizer for token identification.
+        train_dataset: The training dataset to use.
     """
+    if not cfg.fix_untrained_tokens:
+        return
+
+    is_ds_zero3: bool = False
+    if os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE") == "3":
+        is_ds_zero3 = True
+
+    # Check if the `token_ids_to_fix` kwarg exists in the fix_untrained_tokens args
+    sig = inspect.signature(fix_untrained_tokens)
+
+    fix_kwargs: Dict[str, Any] = {}
+    # If the function has the `token_ids_to_fix` arg, and fix_untrained_tokens is a list
+    if "token_ids_to_fix" in sig.parameters and isinstance(
+        cfg.fix_untrained_tokens, list
+    ):
+        fix_kwargs["token_ids_to_fix"] = cfg.fix_untrained_tokens
+    if "is_ds_zero3" in sig.parameters:
+        fix_kwargs["is_ds_zero3"] = is_ds_zero3
+
+    fix_untrained_tokens(model, tokenizer, train_dataset, **fix_kwargs)
+
+    if cfg.local_rank == 0:
+        model.save_pretrained(str(Path(cfg.output_dir)))
+
+
+def setup_model_and_trainer(
+    cfg: DictDefault, dataset_meta: TrainDatasetMeta
+) -> tuple[
+    "HFRLTrainerBuilder" | "HFCausalTrainerBuilder",
+    PeftModel | PreTrainedModel,
+    PreTrainedTokenizer,
+    PeftConfig | None,
+    ProcessorMixin | None,
+]:
+    """
+    Load model, tokenizer, trainer, etc. Helper function to encapsulate the full
+    trainer setup.
+
+    Args:
+        cfg: The configuration dictionary with training parameters.
+        dataset_meta: Object with training, validation datasets and metadata.
+
+    Returns:
+        Tuple of:
+            - Trainer (Causal or RLHF)
+            - Model
+            - Tokenizer
+            - PEFT config
+            - Processor
+    """
+    # Load tokenizer, processor and model
+    model, tokenizer, peft_config, processor = setup_model_and_tokenizer(cfg)
+
+    # Set up reference model for RL if needed
+    model_ref = setup_reference_model(cfg, tokenizer)
+
+    # Get datasets from metadata
+    train_dataset = dataset_meta.train_dataset
+    eval_dataset = dataset_meta.eval_dataset
+    total_num_steps = dataset_meta.total_num_steps
+
+    # Set up trainer
+    trainer = setup_trainer(
+        cfg=cfg,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        model=model,
+        tokenizer=tokenizer,
+        processor=processor,
+        total_num_steps=total_num_steps,
+        model_ref=model_ref,
+        peft_config=peft_config,
+    )
+    PLUGIN_MANAGER.post_trainer_create(cfg, trainer)
+
+    if cfg.use_ray:
+        try:
+            import ray.train.huggingface.transformers
+
+            trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)
+        except ImportError:
+            LOG.warning(
+                "The Ray integration with Hugging Face Transformers is not available. "
+                "To use Ray, install the 'ray[train]' package."
+            )
+
+    return (
+        trainer,
+        model,
+        tokenizer,
+        peft_config,
+        processor,
+    )
+
+
+@send_errors
+def train(
+    cfg: DictDefault, dataset_meta: TrainDatasetMeta
+) -> tuple[PeftModel | PreTrainedModel, PreTrainedTokenizer, Trainer]:
+    """
+    Train a model on the given dataset.
+
+    Args:
+        cfg: The configuration dictionary with training parameters
+        dataset_meta: Object with training, validation datasets and metadata
+
+    Returns:
+        Tuple of (model, tokenizer) after training
+    """
+    # Setup model, tokenizer, (causal or RLHF) trainer, etc.
+    (
+        trainer,
+        model,
+        tokenizer,
+        peft_config,
+        processor,
+    ) = setup_model_and_trainer(cfg, dataset_meta)
+
+    # Handle untrained tokens if configured
+    train_dataset = dataset_meta.train_dataset
+    handle_untrained_tokens_fix(cfg, model, tokenizer, train_dataset)
+
+    # Additional setup
+    save_initial_configs(cfg, tokenizer, model, peft_config, processor)
+    setup_signal_handler(cfg, model)
+    setup_model_card(cfg)
+
+    # Execute the training
+    resume_from_checkpoint = determine_last_checkpoint(cfg)
+    execute_training(cfg, trainer, resume_from_checkpoint)
+
+    # clear cache
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    # Save the trained model and cleanup
+    save_trained_model(cfg, trainer, model)
+    tokenizer.save_pretrained(
+        str(Path(cfg.output_dir)), save_jinja_files=cfg.tokenizer_save_jinja_files
+    )
+    create_model_card(cfg, trainer)
+    if not cfg.use_ray:
+        cleanup_distributed()
+    PLUGIN_MANAGER.post_train(cfg, model)
+
+    return model, tokenizer, trainer
diff --git a/src/axolotl/utils/__init__.py b/src/axolotl/utils/__init__.py
index 99dec79f1b..69826d2db0 100644
--- a/src/axolotl/utils/__init__.py
+++ b/src/axolotl/utils/__init__.py
@@ -1,8 +1,109 @@
 """
 Basic utils for Axolotl
 """
-import importlib
+
+import importlib.metadata
+import importlib.util
+import os
+import re
+
+
+def make_lazy_getattr(
+    lazy_imports: dict[str, str], module_name: str, module_globals: dict
+):
+    """Create a module-level ``__getattr__`` that lazily imports symbols.
+
+    Args:
+        lazy_imports: Mapping of attribute name to relative module path,
+            e.g. ``{"AxolotlDPOTrainer": ".dpo.trainer"}``.
+        module_name: The ``__name__`` of the calling module (used as the
+            anchor for relative imports).
+        module_globals: The ``globals()`` dict of the calling module,
+            used to cache resolved attributes so ``__getattr__`` is only
+            invoked once per name.
+
+    Returns:
+        A ``__getattr__`` function suitable for assignment at module scope.
+    """
+    import importlib
+
+    def __getattr__(name: str):
+        if name in lazy_imports:
+            module = importlib.import_module(lazy_imports[name], module_name)
+            attr = getattr(module, name)
+            module_globals[name] = attr
+            return attr
+        raise AttributeError(f"module {module_name!r} has no attribute {name!r}")
+
+    return __getattr__
 
 
 def is_mlflow_available():
     return importlib.util.find_spec("mlflow") is not None
+
+
+def is_comet_available():
+    return importlib.util.find_spec("comet_ml") is not None
+
+
+def is_opentelemetry_available():
+    return (
+        importlib.util.find_spec("opentelemetry") is not None
+        and importlib.util.find_spec("prometheus_client") is not None
+    )
+
+
+def is_trackio_available():
+    return importlib.util.find_spec("trackio") is not None
+
+
+def get_pytorch_version() -> tuple[int, int, int]:
+    """
+    Get Pytorch version as a tuple of (major, minor, patch).
+    """
+    try:
+        torch_version = importlib.metadata.version("torch")
+    except importlib.metadata.PackageNotFoundError:
+        import torch
+
+        torch_version = torch.__version__
+    version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version)
+
+    if not version_match:
+        raise ValueError("Invalid version format")
+
+    major, minor, patch = version_match.groups()
+    major, minor = int(major), int(minor)
+    patch = int(patch) if patch is not None else 0  # Default patch to 0 if not present
+    return major, minor, patch
+
+
+def set_pytorch_cuda_alloc_conf():
+    """Set up CUDA allocation config"""
+    torch_major, torch_minor, _ = get_pytorch_version()
+    config_value = "expandable_segments:True"
+    config_older_suffix = ",roundup_power2_divisions:16"
+    if (
+        torch_major == 2
+        and torch_minor >= 9
+        and os.getenv("PYTORCH_ALLOC_CONF") is None
+    ):
+        os.environ["PYTORCH_ALLOC_CONF"] = config_value
+    elif (
+        torch_major == 2
+        and torch_minor >= 2
+        and os.getenv("PYTORCH_CUDA_ALLOC_CONF") is None
+    ):
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = config_value + config_older_suffix
+
+
+def set_misc_env():
+    if os.getenv("XFORMERS_IGNORE_FLASH_VERSION_CHECK") is None:
+        os.environ["XFORMERS_IGNORE_FLASH_VERSION_CHECK"] = "1"
+
+
+def get_not_null(value, default=None):
+    """
+    return the value if it's not None, otherwise return the default value
+    """
+    return value if value is not None else default
diff --git a/src/axolotl/utils/bench.py b/src/axolotl/utils/bench.py
index 11c25160da..0a45949915 100644
--- a/src/axolotl/utils/bench.py
+++ b/src/axolotl/utils/bench.py
@@ -1,9 +1,25 @@
 """Benchmarking and measurement utilities"""
+
 import functools
+import logging
 
-import pynvml
 import torch
-from pynvml.nvml import NVMLError
+from transformers.utils.import_utils import is_torch_npu_available
+
+from axolotl.utils.distributed import get_device_type
+
+try:
+    from pynvml import (
+        NVMLError,
+        nvmlDeviceGetHandleByIndex,
+        nvmlDeviceGetMemoryInfo,
+        nvmlInit,
+    )
+except ImportError:
+    NVMLError = None
+    nvmlDeviceGetHandleByIndex = None
+    nvmlDeviceGetMemoryInfo = None
+    nvmlInit = None
 
 
 def check_cuda_device(default_value):
@@ -41,15 +57,22 @@ def gpu_memory_usage(device=0):
 
 @check_cuda_device((0.0, 0.0, 0.0))
 def gpu_memory_usage_all(device=0):
-    usage = torch.cuda.memory_allocated(device) / 1024.0**3
-    reserved = torch.cuda.memory_reserved(device) / 1024.0**3
-    smi = gpu_memory_usage_smi(device)
-    return usage, reserved - usage, max(0, smi - reserved)
+    active = torch.cuda.memory_stats().get("active_bytes.all.peak", 0) / 1024.0**3
+    allocated = torch.cuda.max_memory_allocated(device) / 1024.0**3
+    reserved = torch.cuda.max_memory_reserved(device) / 1024.0**3
+    torch.cuda.reset_peak_memory_stats(device)
+    return active, allocated, reserved
 
 
 def mps_memory_usage_all():
-    usage = torch.mps.current_allocated_memory() / 1024.0**3
-    reserved = torch.mps.driver_allocated_memory() / 1024.0**3
+    active = torch.mps.current_allocated_memory() / 1024.0**3
+    allocated = torch.mps.driver_allocated_memory() / 1024.0**3
+    return active, allocated, 0
+
+
+def npu_memory_usage_all(device=0):
+    usage = torch.npu.memory_allocated(device) / 1024.0**3
+    reserved = torch.npu.memory_reserved(device) / 1024.0**3
     return usage, reserved - usage, 0
 
 
@@ -59,26 +82,49 @@ def gpu_memory_usage_smi(device=0):
         device = device.index
     if isinstance(device, str) and device.startswith("cuda:"):
         device = int(device[5:])
+    if not nvmlInit:
+        return 0.0
     try:
-        pynvml.nvmlInit()
-        handle = pynvml.nvmlDeviceGetHandleByIndex(device)
-        info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        nvmlInit()
+        handle = nvmlDeviceGetHandleByIndex(device)
+        info = nvmlDeviceGetMemoryInfo(handle)
         return info.used / 1024.0**3
     except NVMLError:
         return 0.0
 
 
-def log_gpu_memory_usage(log, msg, device):
+def get_gpu_memory_usage(device: int | torch.device = 0):
+    cur_device_type = str(get_device_type())
     if torch.backends.mps.is_available():
         usage, cache, misc = mps_memory_usage_all()
-    else:
+    elif "npu" in cur_device_type and is_torch_npu_available():
+        usage, cache, misc = npu_memory_usage_all(device)
+    elif "cuda" in cur_device_type and torch.cuda.is_available():
         usage, cache, misc = gpu_memory_usage_all(device)
+    else:
+        return 0.0, 0.0, 0.0
+
+    return usage, cache, misc
+
+
+def log_gpu_memory_usage(
+    log: logging.Logger | logging.LoggerAdapter,
+    msg: str = "",
+    device: int | torch.device = 0,
+):
+    try:
+        active, allocated, reserved = get_gpu_memory_usage(device)
+    except ValueError:
+        # likely CPU, ignore
+        return
+    cur_device_type = str(get_device_type())
     extras = []
-    if cache > 0:
-        extras.append(f"+{cache:.03f}GB cache")
-    if misc > 0:
-        extras.append(f"+{misc:.03f}GB misc")
-    log.info(
-        f"GPU memory usage {msg}: {usage:.03f}GB ({', '.join(extras)})", stacklevel=2
+    if allocated > 0:
+        extras.append(f"+{allocated:.03f}GB allocated")
+    if reserved > 0:
+        extras.append(f"+{reserved:.03f}GB reserved")
+    msg = f"{cur_device_type} memory active:" if not msg else msg
+    log.debug(
+        f"{msg} {active:.03f}GB ({', '.join(extras)})",
+        stacklevel=2,
     )
-    return usage, cache, misc
diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py
index 2965ac1e29..0a278f6ecb 100644
--- a/src/axolotl/utils/callbacks/__init__.py
+++ b/src/axolotl/utils/callbacks/__init__.py
@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
-import logging
+import gc
+import json
 import os
+import traceback
 from shutil import copyfile
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Any, Dict, List
@@ -14,8 +16,8 @@
 import torch
 import torch.distributed as dist
 import wandb
+import yaml
 from datasets import load_dataset
-from optimum.bettertransformer import BetterTransformer
 from tqdm import tqdm
 from transformers import (
     GenerationConfig,
@@ -25,11 +27,13 @@
     TrainerState,
     TrainingArguments,
 )
-from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy
+from transformers.trainer_utils import (
+    SaveStrategy,
+)
+from trl.models import unwrap_model_for_generation
 
-from axolotl.utils import is_mlflow_available
-from axolotl.utils.bench import log_gpu_memory_usage
-from axolotl.utils.config.models.input.v0_4_1 import AxolotlInputConfig
+from axolotl.utils import is_comet_available, is_mlflow_available
+from axolotl.utils.callbacks.perplexity import Perplexity
 from axolotl.utils.distributed import (
     barrier,
     broadcast_dict,
@@ -39,121 +43,108 @@
     is_main_process,
     zero_first,
 )
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.config import AxolotlInputConfig
 
 if TYPE_CHECKING:
-    from axolotl.core.trainer_builder import AxolotlTrainingArguments
+    from axolotl.core.training_args import AxolotlTrainingArguments
+
 
 IGNORE_INDEX = -100
-LOG = logging.getLogger("axolotl.callbacks")
+LOG = get_logger(__name__)
 
 
-class EvalFirstStepCallback(
-    TrainerCallback
-):  # pylint: disable=too-few-public-methods disable=unused-argument
-    """
-    Callback to trigger evals on the first step
-    """
+class LossWatchDogCallback(TrainerCallback):
+    """Callback to track loss and stop training if loss is too high"""
+
+    def __init__(self, cfg):
+        self.cfg = cfg
+        self.violations = 0
+        self.threshold = cfg.loss_watchdog_threshold
+        self.patience = cfg.loss_watchdog_patience or 3
 
     def on_step_end(
         self,
         args: TrainingArguments,
         state: TrainerState,
         control: TrainerControl,
-        **kwargs,
-    ):
-        if (
-            args.evaluation_strategy == IntervalStrategy.STEPS
-            and state.global_step == 1
-        ):
-            control.should_evaluate = True
+        **_kwargs,
+    ) -> TrainerControl:
+        if len(state.log_history) > 0 and "loss" in state.log_history[-1]:
+            if state.log_history[-1]["loss"] > self.threshold:
+                self.violations += 1
+                if self.violations >= self.patience:
+                    LOG.warning(
+                        "Loss is too high, stopping training (loss_watchdog_threshold)"
+                    )
+                    control.should_training_stop = True
+            else:
+                self.violations = 0
         return control
 
 
-class SaveBetterTransformerModelCallback(
-    TrainerCallback
-):  # pylint: disable=too-few-public-methods
-    """Callback to save the BetterTransformer wrapped model"""
+class SaveModelOnFirstStepCallback(TrainerCallback):
+    """Callback to save the model on the first step of training if enabled"""
 
     def on_step_end(
         self,
         args: TrainingArguments,
         state: TrainerState,
         control: TrainerControl,
-        **kwargs,
-    ):
-        # Save
-        if (
-            args.save_strategy == IntervalStrategy.STEPS
-            and args.save_steps > 0
-            and state.global_step % args.save_steps == 0
-        ):
+        **_kwargs,
+    ) -> TrainerControl:
+        if state.global_step == 1:
             control.should_save = True
+        return control
 
-        if control.should_save:
-            checkpoint_folder = os.path.join(
-                args.output_dir,
-                f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
-            )
 
-            model = BetterTransformer.reverse(kwargs["model"])
-            model.save_pretrained(checkpoint_folder)
-            # FIXME - need to cleanup old checkpoints
+class SkipEvalOnResumeCallback(TrainerCallback):
+    """Skip the redundant evaluation that fires when resuming from a checkpoint
+    whose step aligns with ``eval_steps``.
 
-            # since we're saving here, we don't need the trainer loop to attempt to save too b/c
-            # the trainer will raise an exception since it can't save a BetterTransformer wrapped model
-            control.should_save = False
-        return control
+    When HuggingFace Trainer resumes, it restores ``global_step`` from the
+    checkpoint and immediately triggers ``_maybe_log_save_evaluate`` for that
+    step.  Because the evaluation was already performed during the original
+    run, repeating it wastes time and pollutes metric logs.
 
+    This callback records the ``global_step`` at the start of training (i.e.
+    the checkpoint step when resuming, or 0 for a fresh run) and suppresses
+    any evaluation request on that exact step.
+    """
 
-class GPUStatsCallback(
-    TrainerCallback
-):  # pylint: disable=too-few-public-methods disable=unused-argument
-    """Callback to track GPU utilization"""
+    def __init__(self):
+        super().__init__()
+        self._resume_step: int | None = None
 
-    def __init__(self, cfg):
-        self.cfg = cfg
-        self.logged = False
-
-    def on_step_end(
+    def on_train_begin(
         self,
         args: TrainingArguments,
         state: TrainerState,
         control: TrainerControl,
-        **kwargs,
+        **_kwargs,
     ):
-        if not self.logged and state.global_step > 1:
-            log_gpu_memory_usage(LOG, "while training", self.cfg.device)
-            self.logged = True
-        return control
-
-
-class LossWatchDogCallback(TrainerCallback):
-    """Callback to track loss and stop training if loss is too high"""
-
-    def __init__(self, cfg):
-        self.cfg = cfg
-        self.logged = False
-        self.violations = 0
-        self.threshold = cfg.loss_watchdog_threshold
-        self.patience = cfg.loss_watchdog_patience or 3
+        # ``global_step`` is already restored from the checkpoint at this
+        # point.  For a fresh run it will be 0, so the guard below becomes a
+        # no-op.
+        self._resume_step = state.global_step
 
     def on_step_end(
         self,
-        _args: TrainingArguments,
+        args: TrainingArguments,
         state: TrainerState,
         control: TrainerControl,
         **_kwargs,
-    ):
-        if len(state.log_history) > 0 and "loss" in state.log_history[-1]:
-            if state.log_history[-1]["loss"] > self.threshold:
-                self.violations += 1
-                if self.violations >= self.patience:
-                    LOG.warning(
-                        "Loss is too high, stopping training (loss_watchdog_threshold)"
-                    )
-                    control.should_training_stop = True
-            else:
-                self.violations = 0
+    ) -> TrainerControl:
+        if (
+            self._resume_step
+            and state.global_step <= self._resume_step
+            and control.should_evaluate
+        ):
+            LOG.info(
+                "Skipping evaluation at step %d (already completed before resume)",
+                state.global_step,
+            )
+            control.should_evaluate = False
         return control
 
 
@@ -260,10 +251,10 @@ class BenchEvalCallback(TrainerCallback):
         def on_evaluate(
             self,
             args: AxolotlTrainingArguments,
-            state: TrainerState,  # pylint: disable=unused-argument
-            control: TrainerControl,  # pylint: disable=unused-argument
-            metrics: Dict[str, float],  # pylint: disable=unused-argument
-            **kwargs,  # pylint: disable=unused-argument
+            state: TrainerState,
+            control: TrainerControl,
+            metrics: Dict[str, float],
+            **kwargs,
         ):
             data_loader = trainer.get_bench_dataloader(
                 bench_dataset.remove_columns(["input", "subject", "output", "name"])
@@ -293,7 +284,7 @@ def on_evaluate(
             # Extract results by subject.
             bench_name = bench_dataset["name"]
             bench_names: dict = {s: {"refs": [], "preds": []} for s in set(bench_name)}
-            for s, p, r in zip(bench_name, preds, refs):  # pylint: disable=invalid-name
+            for s, p, r in zip(bench_name, preds, refs, strict=False):
                 bench_names[s]["preds"].append(p)
                 bench_names[s]["refs"].append(r)
             barrier()
@@ -331,9 +322,7 @@ def on_evaluate(
                 bench_scores = []
                 bench_refs = []
                 bench_preds = []
-                for (
-                    bench_name
-                ) in combined_bench_names:  # pylint: disable=consider-using-dict-items
+                for bench_name in combined_bench_names:
                     bench_score = accuracy.compute(
                         references=combined_bench_names[bench_name]["refs"],
                         predictions=combined_bench_names[bench_name]["preds"],
@@ -341,9 +330,9 @@ def on_evaluate(
                     bench_refs.extend(combined_bench_names[bench_name]["refs"])
                     bench_preds.extend(combined_bench_names[bench_name]["preds"])
                     if not pd.isna(bench_score):
-                        results[
-                            f"{bench_split}_bench_accuracy_{bench_name}"
-                        ] = bench_score
+                        results[f"{bench_split}_bench_accuracy_{bench_name}"] = (
+                            bench_score
+                        )
                         bench_scores.append(bench_score)
                     else:
                         results[f"{bench_split}_bench_accuracy_{bench_name}"] = 0.0
@@ -373,25 +362,34 @@ def __init__(self, cfg):
         def __maybe_load_metrics(self):
             metrics = {}
             for metric in self.cfg.eval_causal_lm_metrics:
-                try:
-                    metrics[metric] = evaluate.load(metric)
-                except Exception as exc:  # pylint: disable=broad-exception-caught
-                    LOG.warning(f"{metric}: {exc.args}")
+                if metric == "perplexity":
+                    max_seq_len = self.cfg.eval_max_new_tokens
+                    metrics[metric] = Perplexity(
+                        tokenizer=tokenizer,
+                        max_seq_len=max_seq_len,
+                    )
+                else:
+                    try:
+                        metrics[metric] = evaluate.load(metric)
+                    except Exception as exc:
+                        LOG.warning(f"{metric}: {exc.args}")
             return metrics
 
         def on_evaluate(
             self,
-            args: AxolotlTrainingArguments,  # pylint: disable=unused-argument
+            args: AxolotlTrainingArguments,
             state: TrainerState,
             control: TrainerControl,
-            train_dataloader,  # pylint: disable=unused-argument
+            train_dataloader,
             eval_dataloader,
-            **kwargs,  # pylint: disable=unused-argument
+            **kwargs,
         ):
-            trainer.model.eval()
-            device = torch.device(self.cfg.device)
+            trainer.model_wrapped.eval()
+
+            device = torch.device(
+                self.cfg.device
+            )  # Use this instead of trainer.model_wrapped.device as it may return cpu if fsdp offloaded
 
-            # pylint: disable=duplicate-code
             generation_config = GenerationConfig(
                 max_new_tokens=self.cfg.eval_max_new_tokens,
                 bos_token_id=tokenizer.bos_token_id,
@@ -420,13 +418,22 @@ def compute(metric: evaluate.Metric, **kwargs):
                 # safely compute a metric and return the score if the format is correct
                 metric_score = None
                 try:
-                    metric_score = metric.compute(**kwargs)
+                    # Only pass the kwargs that are in the metric's feature list
+                    metric_kwargs = {
+                        k: kwargs[k] for k in metric._feature_names() if k in kwargs
+                    }
+
+                    if isinstance(metric, Perplexity):
+                        metric_kwargs["model"] = trainer.model_wrapped
+
+                    metric_score = metric.compute(**metric_kwargs)
                     return (
                         metric_score["score"]
                         if "score" in metric_score
                         else metric_score["mean_score"]
                     )
-                except Exception:  # pylint: disable=broad-exception-caught
+                except Exception:
+                    traceback.print_exc()
                     LOG.debug(
                         f"Failed to compute metric {metric.name} with kwargs {kwargs.keys()}"
                     )
@@ -442,100 +449,110 @@ def evaluate_preds(sources, predictions, references):
                         predictions=predictions,
                         sources=sources,
                     )
-                    score = score or compute(
-                        metric,
-                        references=[[r] for r in references],
-                        predictions=predictions,
-                    )
-                    scores[metric_name] = score
+                    if score is None:
+                        score = compute(
+                            metric,
+                            references=[[r] for r in references],
+                            predictions=predictions,
+                        )
+                    scores["eval_" + metric_name] = score
                 return scores
 
             def predict_with_generate():
                 eval_src, eval_pred, eval_ref = [], [], []
 
-                for batch in tqdm(eval_dataloader):
-                    batch_labels = batch["labels"].to(device)
-                    batch_input_ids = batch["input_ids"].to(device)
+                with unwrap_model_for_generation(
+                    trainer.model_wrapped, trainer.accelerator
+                ) as unwrapped_model:
+                    for batch in tqdm(eval_dataloader, disable=not is_main_process()):
+                        batch_labels = batch["labels"].to(device)
+                        batch_input_ids = batch["input_ids"].to(device)
 
-                    if "position_ids" in batch:
-                        batch_pos_ids = batch["position_ids"].tolist()
-                    else:
-                        batch_pos_ids = [None] * len(batch["input_ids"])
-
-                    prompt_token_ids_list = []
-                    completion_token_ids_list = []
-
-                    for input_ids_all, labels_all, pos_ids in zip(
-                        batch_input_ids,
-                        batch_labels,
-                        batch_pos_ids,
-                    ):
-                        if pos_ids is None:
-                            pos_ranges = [(0, len(input_ids_all) - 1)]
+                        if "position_ids" in batch:
+                            batch_pos_ids = batch["position_ids"].tolist()
                         else:
-                            pos_ranges = find_ranges(pos_ids)
-
-                        for pos_range in pos_ranges:
-                            start, end = pos_range
-                            if start == end:
-                                continue
+                            batch_pos_ids = [None] * len(batch["input_ids"])
+
+                        prompt_token_ids_list = []
+                        completion_token_ids_list = []
+
+                        for input_ids_all, labels_all, pos_ids in zip(
+                            batch_input_ids,
+                            batch_labels,
+                            batch_pos_ids,
+                            strict=False,
+                        ):
+                            if pos_ids is None:
+                                pos_ranges = [(0, len(input_ids_all) - 1)]
+                            else:
+                                pos_ranges = find_ranges(pos_ids)
+
+                            for pos_range in pos_ranges:
+                                start, end = pos_range
+                                if start == end:
+                                    continue
+
+                                input_ids = input_ids_all[start : end + 1]
+                                labels = labels_all[start : end + 1]
+
+                                tokens_without_loss = labels == IGNORE_INDEX
+                                tokens_with_loss = labels != IGNORE_INDEX
+                                tokens_exclude_padding = (
+                                    input_ids != tokenizer.pad_token_id
+                                )
+                                prompt_token_includes = (
+                                    tokens_without_loss & tokens_exclude_padding
+                                )
+
+                                prompt_token_ids = input_ids[prompt_token_includes]
+                                prompt_token_ids_list.append(prompt_token_ids)
+
+                                completion_token_ids = input_ids[tokens_with_loss]
+                                completion_token_ids_list.append(completion_token_ids)
+
+                        prompt_texts = tokenizer.batch_decode(
+                            prompt_token_ids_list, skip_special_tokens=True
+                        )
+                        completion_texts = tokenizer.batch_decode(
+                            completion_token_ids_list, skip_special_tokens=True
+                        )
 
-                            input_ids = input_ids_all[start : end + 1]
-                            labels = labels_all[start : end + 1]
+                        with torch.no_grad():
+                            prompt_encoding = tokenizer(
+                                prompt_texts, padding=True, return_tensors="pt"
+                            ).to(device)
 
-                            tokens_without_loss = labels == IGNORE_INDEX
-                            tokens_with_loss = labels != IGNORE_INDEX
-                            tokens_exclude_padding = input_ids != tokenizer.pad_token_id
-                            prompt_token_includes = (
-                                tokens_without_loss & tokens_exclude_padding
+                            predictions = unwrapped_model.generate(
+                                **prompt_encoding, generation_config=generation_config
                             )
 
-                            prompt_token_ids = input_ids[prompt_token_includes]
-                            prompt_token_ids_list.append(prompt_token_ids)
-
-                            completion_token_ids = input_ids[tokens_with_loss]
-                            completion_token_ids_list.append(completion_token_ids)
-
-                    prompt_texts = tokenizer.batch_decode(
-                        prompt_token_ids_list, skip_special_tokens=True
-                    )
-                    completion_texts = tokenizer.batch_decode(
-                        completion_token_ids_list, skip_special_tokens=True
-                    )
-
-                    with torch.no_grad():
-                        prompt_encoding = tokenizer(
-                            prompt_texts, padding=True, return_tensors="pt"
-                        ).to(self.cfg.device)
-                        predictions = trainer.model.generate(
-                            **prompt_encoding, generation_config=generation_config
-                        )
+                            del prompt_encoding
+
+                        prediction_all_tokens = predictions["sequences"].cpu().tolist()
+                        prediction_without_prompt_tokens_list = []
+                        for prompt_token_ids, prediction_tokens in zip(
+                            prompt_token_ids_list, prediction_all_tokens, strict=False
+                        ):
+                            prediction_without_prompt_tokens = prediction_tokens[
+                                len(prompt_token_ids) :
+                            ]
+                            prediction_without_prompt_tokens_list.append(
+                                prediction_without_prompt_tokens
+                            )
 
-                    prediction_all_tokens = predictions["sequences"].cpu().tolist()
-                    prediction_without_prompt_tokens_list = []
-                    for prompt_token_ids, prediction_tokens in zip(
-                        prompt_token_ids_list, prediction_all_tokens
-                    ):
-                        prediction_without_prompt_tokens = prediction_tokens[
-                            len(prompt_token_ids) :
-                        ]
-                        prediction_without_prompt_tokens_list.append(
-                            prediction_without_prompt_tokens
+                        predicted_texts = tokenizer.batch_decode(
+                            prediction_without_prompt_tokens_list,
+                            skip_special_tokens=True,
                         )
 
-                    predicted_texts = tokenizer.batch_decode(
-                        prediction_without_prompt_tokens_list, skip_special_tokens=True
-                    )
-
-                    eval_src.extend(prompt_texts)
-                    eval_pred.extend(predicted_texts)
-                    eval_ref.extend(completion_texts)
+                        eval_src.extend(prompt_texts)
+                        eval_pred.extend(predicted_texts)
+                        eval_ref.extend(completion_texts)
 
                 return eval_src, eval_pred, eval_ref
 
-            if is_main_process():
-                eval_preds = predict_with_generate()
-                trainer.log(evaluate_preds(*eval_preds))
+            eval_preds = predict_with_generate()
+            trainer.log(evaluate_preds(*eval_preds))
 
             return control
 
@@ -552,12 +569,12 @@ def __init__(self, cfg):
 
         def on_evaluate(
             self,
-            args: AxolotlTrainingArguments,  # pylint: disable=unused-argument
+            args: AxolotlTrainingArguments,
             state: TrainerState,
             control: TrainerControl,
-            train_dataloader,  # pylint: disable=unused-argument
+            train_dataloader,
             eval_dataloader,
-            **kwargs,  # pylint: disable=unused-argument
+            **kwargs,
         ):
             eval_table_size = self.cfg.eval_table_size
 
@@ -567,7 +584,6 @@ def on_evaluate(
             trainer.model.eval()
             device = torch.device(self.cfg.device)
 
-            # pylint: disable=duplicate-code
             generation_config = GenerationConfig(
                 max_new_tokens=self.cfg.eval_max_new_tokens,
                 bos_token_id=tokenizer.bos_token_id,
@@ -635,6 +651,7 @@ def log_table_from_dataloader(name: str, table_dataloader):
                         batch_labels,
                         batch_pos_ids,
                         batch_logits,
+                        strict=False,
                     ):
                         if pos_ids is None:
                             pos_ranges = [(0, len(input_ids_all) - 1)]
@@ -688,7 +705,7 @@ def log_table_from_dataloader(name: str, table_dataloader):
                     prediction_all_tokens = predictions["sequences"].cpu().tolist()
                     prediction_without_prompt_tokens_list = []
                     for prompt_token_ids, prediction_tokens in zip(
-                        prompt_token_ids_list, prediction_all_tokens
+                        prompt_token_ids_list, prediction_all_tokens, strict=False
                     ):
                         prediction_without_prompt_tokens = prediction_tokens[
                             len(prompt_token_ids) :
@@ -707,7 +724,11 @@ def log_table_from_dataloader(name: str, table_dataloader):
                         prediction_text,
                         pred_step_text,
                     ) in zip(
-                        prompt_texts, completion_texts, predicted_texts, pred_step_texts
+                        prompt_texts,
+                        completion_texts,
+                        predicted_texts,
+                        pred_step_texts,
+                        strict=False,
                     ):
                         table_data["id"].append(row_index)
                         table_data["Prompt"].append(prompt_text)
@@ -720,7 +741,13 @@ def log_table_from_dataloader(name: str, table_dataloader):
                         ].append(pred_step_text)
                         row_index += 1
                 if logger == "wandb":
-                    wandb.run.log({f"{name} - Predictions vs Ground Truth": pd.DataFrame(table_data)})  # type: ignore[attr-defined]
+                    wandb.run.log(  # type: ignore[attr-defined]
+                        {
+                            f"{name} - Predictions vs Ground Truth": pd.DataFrame(
+                                table_data
+                            )
+                        }
+                    )
                 elif logger == "mlflow" and is_mlflow_available():
                     import mlflow
 
@@ -732,6 +759,15 @@ def log_table_from_dataloader(name: str, table_dataloader):
                         artifact_file="PredictionsVsGroundTruth.json",
                         tracking_uri=tracking_uri,
                     )
+                elif logger == "comet_ml" and is_comet_available():
+                    import comet_ml
+
+                    experiment = comet_ml.get_running_experiment()
+                    if experiment:
+                        experiment.log_table(
+                            f"{name} - Predictions vs Ground Truth.csv",
+                            pd.DataFrame(table_data),
+                        )
 
             if is_main_process():
                 log_table_from_dataloader("Eval", eval_dataloader)
@@ -749,48 +785,191 @@ def __init__(self, axolotl_config_path):
 
     def on_train_begin(
         self,
-        args: AxolotlTrainingArguments,  # pylint: disable=unused-argument
-        state: TrainerState,  # pylint: disable=unused-argument
+        args: AxolotlTrainingArguments,
+        state: TrainerState,
         control: TrainerControl,
-        **kwargs,  # pylint: disable=unused-argument
+        **kwargs,
     ):
-        if is_main_process():
+        if state.is_world_process_zero:
             try:
                 # sync config to top level in run, cannot delete file right away because wandb schedules it to be synced even w/policy = 'now', so let OS delete it later.
                 with NamedTemporaryFile(
                     mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
                 ) as temp_file:
                     copyfile(self.axolotl_config_path, temp_file.name)
-                    artifact = wandb.Artifact(
-                        f"config-{wandb.run.id}", type="axolotl-config"
+                    artifact = wandb.Artifact(  # type: ignore[attr-defined]
+                        f"config-{wandb.run.id}",  # type: ignore[attr-defined]
+                        type="axolotl-config",
                     )
                     artifact.add_file(temp_file.name)
-                    wandb.log_artifact(artifact)
-                    wandb.save(temp_file.name)
-                LOG.info(
-                    "The Axolotl config has been saved to the WandB run under files."
-                )
+                    wandb.log_artifact(artifact)  # type: ignore[attr-defined]
+                    wandb.save(temp_file.name)  # type: ignore[attr-defined]
+                    LOG.info(
+                        "The Axolotl config has been saved to the WandB run under files."
+                    )
             except (FileNotFoundError, ConnectionError) as err:
                 LOG.warning(f"Error while saving Axolotl config to WandB: {err}")
+
+            try:
+                with open(self.axolotl_config_path, "r", encoding="utf-8") as f:
+                    cfg = yaml.safe_load(f) or {}
+
+                chat_tpl = cfg.get("chat_template_jinja")
+                if chat_tpl:
+                    with NamedTemporaryFile(
+                        mode="w", delete=True, suffix=".jinja", prefix="chat_template_"
+                    ) as temp_ct_file:
+                        if (
+                            isinstance(chat_tpl, str)
+                            and os.path.exists(chat_tpl)
+                            and os.path.isfile(chat_tpl)
+                        ):
+                            copyfile(chat_tpl, temp_ct_file.name)
+                        else:
+                            temp_ct_file.write(str(chat_tpl))
+                            temp_ct_file.flush()
+
+                        artifact = wandb.Artifact(  # type: ignore[attr-defined]
+                            f"chat-template-{wandb.run.id}",  # type: ignore[attr-defined]
+                            type="jinja-template",
+                        )
+                        artifact.add_file(temp_ct_file.name)
+                        wandb.log_artifact(artifact)  # type: ignore[attr-defined]
+                        wandb.save(temp_ct_file.name)  # type: ignore[attr-defined]
+                        LOG.info(
+                            "The chat_template_jinja has been saved to the WandB run under files."
+                        )
+            except (FileNotFoundError, ConnectionError, yaml.YAMLError) as err:
+                LOG.warning(f"Error while saving chat_template_jinja to WandB: {err}")
+
+            if args.deepspeed:
+                try:
+                    # sync config to top level in run, cannot delete file right away because wandb schedules it to be synced even w/policy = 'now', so let OS delete it later.
+                    with NamedTemporaryFile(
+                        mode="w",
+                        delete=False,
+                        suffix=".json",
+                        prefix="deepspeed_config_",
+                    ) as temp_file:
+                        skip_upload = False
+                        if isinstance(args.deepspeed, dict):
+                            json.dump(args.deepspeed, temp_file, indent=4)
+                        elif isinstance(args.deepspeed, str) and os.path.exists(
+                            args.deepspeed
+                        ):
+                            copyfile(args.deepspeed, temp_file.name)
+                        else:
+                            skip_upload = True
+                        if not skip_upload:
+                            artifact = wandb.Artifact(  # type: ignore[attr-defined]
+                                f"deepspeed-config-{wandb.run.id}",  # type: ignore[attr-defined]
+                                type="deepspeed-config",
+                            )
+                            artifact.add_file(temp_file.name)
+                            wandb.log_artifact(artifact)  # type: ignore[attr-defined]
+                            wandb.save(temp_file.name)  # type: ignore[attr-defined]
+                            LOG.info(
+                                "The DeepSpeed config has been saved to the WandB run under files."
+                            )
+                except (FileNotFoundError, ConnectionError) as err:
+                    LOG.warning(f"Error while saving DeepSpeed config to WandB: {err}")
+
         return control
 
 
-class SaveModelOnTrainEndCallback(TrainerCallback):
-    """Callback to save model on train end"""
+class GCCallback(TrainerCallback):
+    """Runs ``gc.collect()`` + ``torch.cuda.empty_cache()`` on
+    ``gc_collect_steps`` intervals and on eval/save/epoch boundaries that
+    the Trainer's native ``torch_empty_cache_steps`` doesn't cover. The two
+    settings are complementary; overlapping intervals just double-clear.
+    """
+
+    def __init__(self, gc_collect_steps: int | None = -1, gc_steps: int | None = None):
+        if gc_steps is not None and gc_collect_steps in (-1, None):
+            gc_collect_steps = gc_steps
+        self.gc_collect_steps: int = gc_collect_steps or -1
+        self.next_gc_on_begin_step: int = -1
+
+    def _gc(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def on_train_begin(
+        self,
+        args,
+        state,
+        control,
+        **kwargs,
+    ):
+        self._gc()
 
-    def on_step_end(  # pylint: disable=unused-argument
+    def on_step_begin(
         self,
-        args: TrainingArguments,
-        state: TrainerState,
-        control: TrainerControl,
+        args,
+        state,
+        control,
         **kwargs,
     ):
-        # Save
-        if state.global_step >= state.max_steps:
-            control.should_save = True
+        if self.next_gc_on_begin_step == state.global_step or state.global_step == 0:
+            self._gc()
 
-    def on_train_end(  # pylint: disable=unused-argument
-        self, args, state, control, **kwargs
+    def on_step_end(
+        self,
+        args,
+        state,
+        control,
+        **kwargs,
     ):
-        control.should_save = True
-        return control
+        if control.should_evaluate:
+            # automatically GC before evals so the eval memory spike from the CEL doesn't OOM the trainer
+            self._gc()
+            # also GC on the start of the next step after the eval
+            self.next_gc_on_begin_step = state.global_step + 1
+        elif (
+            self.gc_collect_steps > 0 and state.global_step % self.gc_collect_steps == 0
+        ):
+            self._gc()
+        elif (
+            args.save_strategy == SaveStrategy.STEPS
+            and state.save_steps > 0
+            and state.global_step % state.save_steps == 0
+        ):
+            # gc on save steps in case anything is loaded to CPU RAM like offloaded tensors
+            self._gc()
+        elif state.global_step >= state.max_steps:
+            if args.save_strategy == SaveStrategy.STEPS:
+                # gc on save steps in case anything is loaded to CPU RAM like offloaded tensors
+                self._gc()
+
+    def on_epoch_end(
+        self,
+        args,
+        state,
+        control,
+        **kwargs,
+    ):
+        self._gc()
+
+
+def colab_inference_post_train_callback(trainer: Trainer):
+    class ColabCallback(TrainerCallback):
+        """Callback to prep model for inference on Google Colab"""
+
+        def __init__(self, cfg):
+            self.gpu_name = torch.cuda.get_device_name(0)
+            self.cfg = cfg
+
+        def on_train_end(self, args, state, control, **kwargs):
+            """
+            handle T4 gpu, we need to convert attention to eager for inference
+            """
+            if (
+                "Tesla T4" in self.gpu_name
+                and self.cfg.attn_implementation == "xformers"
+            ):
+                trainer.model.config._attn_implementation = "eager"
+            trainer.model.gradient_checkpointing_disable()
+            trainer.model.config.use_cache = True
+            trainer.model.eval()
+
+    return ColabCallback
diff --git a/src/axolotl/utils/callbacks/comet_.py b/src/axolotl/utils/callbacks/comet_.py
new file mode 100644
index 0000000000..cd3bcf70ec
--- /dev/null
+++ b/src/axolotl/utils/callbacks/comet_.py
@@ -0,0 +1,43 @@
+"""Comet module for trainer callbacks"""
+
+from typing import TYPE_CHECKING
+
+import comet_ml
+from transformers import TrainerCallback, TrainerControl, TrainerState
+
+from axolotl.utils.distributed import is_main_process
+from axolotl.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from axolotl.core.training_args import AxolotlTrainingArguments
+
+LOG = get_logger(__name__)
+
+
+class SaveAxolotlConfigtoCometCallback(TrainerCallback):
+    """Callback to save axolotl config to comet"""
+
+    def __init__(self, axolotl_config_path):
+        self.axolotl_config_path = axolotl_config_path
+
+    def on_train_begin(
+        self,
+        args: "AxolotlTrainingArguments",
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        if is_main_process():
+            try:
+                comet_experiment = comet_ml.start(source="axolotl")
+                comet_experiment.log_other("Created from", "axolotl")
+                comet_experiment.log_asset(
+                    self.axolotl_config_path,
+                    file_name="axolotl-config",
+                )
+                LOG.info(
+                    "The Axolotl config has been saved to the Comet Experiment under assets."
+                )
+            except (FileNotFoundError, ConnectionError) as err:
+                LOG.warning(f"Error while saving Axolotl config to Comet: {err}")
+        return control
diff --git a/src/axolotl/utils/callbacks/dynamic_checkpoint.py b/src/axolotl/utils/callbacks/dynamic_checkpoint.py
new file mode 100644
index 0000000000..ac484b3d0b
--- /dev/null
+++ b/src/axolotl/utils/callbacks/dynamic_checkpoint.py
@@ -0,0 +1,128 @@
+from pathlib import Path
+
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+from axolotl.utils.distributed import (
+    barrier,
+    is_distributed,
+    is_main_process,
+)
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+DEFAULT_TRIGGER_FILENAME = "axolotl_checkpoint.save"
+
+
+class DynamicCheckpointCallback(TrainerCallback):
+    """
+    Callback to save checkpoints on-demand during training via:
+    1. File-based trigger (works everywhere, rank 0 checks file)
+
+    Thread-safe for multi-GPU distributed training.
+
+    Usage:
+        # File-based:
+        touch /path/to/output_dir/axolotl_checkpoint.save
+    """
+
+    def _get_config_value(self, config, key, default=None):
+        """Helper to get config value from dict or object."""
+        if isinstance(config, dict):
+            return config.get(key, default)
+        return getattr(config, key, default)
+
+    def __init__(self, cfg):
+        self.cfg = cfg
+        if not cfg.dynamic_checkpoint or not cfg.dynamic_checkpoint.enabled:
+            self.enabled = False
+            return
+
+        self.enabled = True
+        dc_config = cfg.dynamic_checkpoint
+
+        trigger_file_path = self._get_config_value(dc_config, "trigger_file_path")
+        self.trigger_filename = (
+            trigger_file_path if trigger_file_path else DEFAULT_TRIGGER_FILENAME
+        )
+
+        check_interval = self._get_config_value(dc_config, "check_interval")
+        self.check_interval = check_interval if check_interval is not None else 100
+        self.should_save_checkpoint = False
+
+        LOG.info(
+            f"Dynamic checkpoint enabled. To trigger checkpoint save:\n"
+            f"  • File: touch {cfg.output_dir}/{self.trigger_filename}\n"
+            f"  • Check interval: every {self.check_interval} steps",
+        )
+
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **_kwargs,
+    ) -> TrainerControl:
+        """
+        Check for checkpoint triggers at the end of each step.
+        ONLY rank 0 checks the file, then all ranks synchronize.
+        """
+        if not self.enabled:
+            return control
+
+        trigger_detected = False
+
+        if state.global_step % self.check_interval == 0:
+            if is_main_process():
+                trigger_path = Path(args.output_dir) / self.trigger_filename
+
+                if trigger_path.exists():
+                    trigger_detected = True
+                    try:
+                        trigger_path.unlink()  # Delete the trigger file
+                        LOG.info(
+                            f"Dynamic checkpoint triggered via file '{self.trigger_filename}' "
+                            f"at step {state.global_step}",
+                        )
+                    except OSError as exc:
+                        LOG.warning(
+                            f"Failed to delete trigger file: {exc}",
+                        )
+
+                if self.should_save_checkpoint:
+                    trigger_detected = True
+                    self.should_save_checkpoint = False  # Reset flag
+
+            if is_distributed():
+                import torch
+                import torch.distributed as dist
+
+                device = getattr(
+                    args,
+                    "device",
+                    torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+                )
+
+                trigger_tensor = torch.tensor(
+                    1 if trigger_detected else 0,
+                    dtype=torch.long,
+                    device=device,
+                )
+
+                dist.broadcast(trigger_tensor, src=0)
+
+                trigger_detected = bool(trigger_tensor.item())
+
+                barrier()
+
+        if trigger_detected:
+            control.should_save = True
+            LOG.info(
+                f"Saving dynamic checkpoint at step {state.global_step}",
+            )
+        return control
diff --git a/src/axolotl/utils/callbacks/generation.py b/src/axolotl/utils/callbacks/generation.py
new file mode 100644
index 0000000000..da36b9ad0e
--- /dev/null
+++ b/src/axolotl/utils/callbacks/generation.py
@@ -0,0 +1,84 @@
+"""Callback for generating samples during SFT/Pretrain training."""
+
+from transformers.trainer_callback import TrainerCallback, TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+
+from axolotl.utils.generation.sft import generate_samples
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class SFTGenerationCallback(TrainerCallback):
+    """Callback for generating samples during SFT/Pretrain training."""
+
+    def __init__(self, trainer):
+        self.trainer = trainer
+
+    def on_evaluate(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Generate samples at specified intervals."""
+        cfg = self.trainer.axolotl_cfg
+
+        if not getattr(cfg, "generate_samples", False):
+            return
+
+        dataloader = None
+        try:
+            if getattr(self.trainer, "eval_dataset", None) is not None:
+                dataloader = self.trainer.get_eval_dataloader()
+                LOG.info(
+                    f"Using eval dataloader for generation at step {state.global_step}"
+                )
+        except Exception as e:
+            LOG.warning(f"Could not get eval dataloader: {e}")
+            dataloader = None
+
+        if dataloader is None:
+            dataloader = self.trainer.get_train_dataloader()
+            LOG.info(
+                f"Using train dataloader for generation at step {state.global_step}"
+            )
+
+        samples = generate_samples(
+            model=self.trainer.model,
+            tokenizer=self.trainer.processing_class,
+            dataloader=dataloader,
+            num_generation_samples=getattr(cfg, "num_generation_samples", 3),
+            max_new_tokens=getattr(cfg, "generation_max_new_tokens", 50),
+            temperature=getattr(cfg, "generation_temperature", 0.7),
+            top_p=getattr(cfg, "generation_top_p", None),
+            top_k=getattr(cfg, "generation_top_k", None),
+            do_sample=getattr(cfg, "generation_do_sample", True),
+            prompt_ratio=getattr(cfg, "generation_prompt_ratio", 0.5),
+        )
+        self._log_samples(samples, state.global_step)
+
+    def _log_samples(self, samples: list, step: int):
+        """Log generated samples to console and W&B."""
+        from axolotl.utils.generation.sft import format_generation_for_logging
+
+        for i, sample in enumerate(samples):
+            console_text, wandb_text = format_generation_for_logging(sample, i, step)
+
+            LOG.info(console_text)
+
+            try:
+                import wandb
+
+                if wandb.run is not None:  # type: ignore[attr-defined]
+                    wandb.log(  # type: ignore[attr-defined]
+                        {
+                            f"samples/sample_{i + 1}": wandb.Html(  # type: ignore[attr-defined]
+                                f"<pre>{wandb_text}</pre>"
+                            )
+                        },
+                        step=step,
+                    )
+            except (ImportError, Exception):
+                pass
diff --git a/src/axolotl/utils/callbacks/lisa.py b/src/axolotl/utils/callbacks/lisa.py
index ff20959a59..03f189d805 100644
--- a/src/axolotl/utils/callbacks/lisa.py
+++ b/src/axolotl/utils/callbacks/lisa.py
@@ -6,17 +6,18 @@
 License: Apache 2.0
 """
 
-import logging
 from functools import reduce
 from typing import TYPE_CHECKING
 
 import numpy as np
 from transformers import TrainerCallback
 
+from axolotl.utils.logging import get_logger
+
 if TYPE_CHECKING:
-    from axolotl.core.trainer_builder import AxolotlTrainer
+    from axolotl.core.trainers import AxolotlTrainer
 
-LOG = logging.getLogger("axolotl.callbacks.lisa")
+LOG = get_logger(__name__)
 
 
 def lisa_callback_factory(trainer: "AxolotlTrainer"):
@@ -43,7 +44,7 @@ def __init__(
                 getattr, self.layers_attribute.split("."), self.trainer.model
             )
             LOG.info(
-                f"LISA will activate {self.n_layers}/{len(layers)} layers ({self.n_layers*100/len(layers)}%) every {self.step_interval} steps"
+                f"LISA will activate {self.n_layers}/{len(layers)} layers ({self.n_layers * 100 / len(layers)}%) every {self.step_interval} steps"
             )
 
         def freeze_all_layers(self):
@@ -54,9 +55,7 @@ def freeze_all_layers(self):
                 for param in layer.parameters():
                     param.requires_grad = False
 
-        def on_step_begin(
-            self, args, state, control, **kwargs
-        ):  # pylint: disable=unused-argument
+        def on_step_begin(self, args, state, control, **kwargs):
             # Check if it's time to switch active layers, including at step 0
             if state.global_step % self.step_interval == 0 or state.global_step == 1:
                 self.switch_active_layers()
diff --git a/src/axolotl/utils/callbacks/mlflow_.py b/src/axolotl/utils/callbacks/mlflow_.py
index fcbb88edcd..30120a87de 100644
--- a/src/axolotl/utils/callbacks/mlflow_.py
+++ b/src/axolotl/utils/callbacks/mlflow_.py
@@ -1,5 +1,6 @@
 """MLFlow module for trainer callbacks"""
-import logging
+
+import os
 from shutil import copyfile
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
@@ -8,15 +9,20 @@
 from transformers import TrainerCallback, TrainerControl, TrainerState
 
 from axolotl.utils.distributed import is_main_process
+from axolotl.utils.logging import get_logger
 
 if TYPE_CHECKING:
-    from axolotl.core.trainer_builder import AxolotlTrainingArguments
+    from axolotl.core.training_args import AxolotlTrainingArguments
+
+LOG = get_logger(__name__)
+
 
-LOG = logging.getLogger("axolotl.callbacks")
+def should_log_artifacts() -> bool:
+    truths = ["TRUE", "1", "YES"]
+    return os.getenv("HF_MLFLOW_LOG_ARTIFACTS", "FALSE").upper() in truths
 
 
 class SaveAxolotlConfigtoMlflowCallback(TrainerCallback):
-    # pylint: disable=duplicate-code
     """Callback to save axolotl config to mlflow"""
 
     def __init__(self, axolotl_config_path):
@@ -24,20 +30,25 @@ def __init__(self, axolotl_config_path):
 
     def on_train_begin(
         self,
-        args: "AxolotlTrainingArguments",  # pylint: disable=unused-argument
-        state: TrainerState,  # pylint: disable=unused-argument
+        args: "AxolotlTrainingArguments",
+        state: TrainerState,
         control: TrainerControl,
-        **kwargs,  # pylint: disable=unused-argument
+        **kwargs,
     ):
         if is_main_process():
             try:
-                with NamedTemporaryFile(
-                    mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
-                ) as temp_file:
-                    copyfile(self.axolotl_config_path, temp_file.name)
-                    mlflow.log_artifact(temp_file.name, artifact_path="")
+                if should_log_artifacts():
+                    with NamedTemporaryFile(
+                        mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
+                    ) as temp_file:
+                        copyfile(self.axolotl_config_path, temp_file.name)
+                        mlflow.log_artifact(temp_file.name, artifact_path="")
+                        LOG.info(
+                            "The Axolotl config has been saved to the MLflow artifacts."
+                        )
+                else:
                     LOG.info(
-                        "The Axolotl config has been saved to the MLflow artifacts."
+                        "Skipping logging artifacts to MLflow (hf_mlflow_log_artifacts is false)"
                     )
             except (FileNotFoundError, ConnectionError) as err:
                 LOG.warning(f"Error while saving Axolotl config to MLflow: {err}")
diff --git a/src/axolotl/utils/callbacks/models.py b/src/axolotl/utils/callbacks/models.py
new file mode 100644
index 0000000000..5a20d70d9c
--- /dev/null
+++ b/src/axolotl/utils/callbacks/models.py
@@ -0,0 +1,23 @@
+"""Helper functions for model classes"""
+
+from typing import Tuple
+
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+
+
+def get_causal_lm_model_cls_prefix(model_type: str) -> Tuple[str, str]:
+    if model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
+        causal_lm_cls = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[model_type]
+        causal_lm_cls_prefix = causal_lm_cls
+        for suffix in [
+            "ForCausalLM",
+            "ForConditionalGeneration",
+            "LMHeadModel",
+            "GenerationDecoder",
+        ]:
+            causal_lm_cls_prefix = causal_lm_cls_prefix.replace(suffix, "")
+        return causal_lm_cls_prefix, causal_lm_cls
+    causal_lm_cls_prefix = "".join(
+        [part.capitalize() for part in model_type.split("_")]
+    )
+    return causal_lm_cls_prefix, f"{causal_lm_cls_prefix}ForCausalLM"
diff --git a/src/axolotl/utils/callbacks/opentelemetry.py b/src/axolotl/utils/callbacks/opentelemetry.py
new file mode 100644
index 0000000000..3f7e56b782
--- /dev/null
+++ b/src/axolotl/utils/callbacks/opentelemetry.py
@@ -0,0 +1,238 @@
+"""OpenTelemetry metrics callback for Axolotl training"""
+
+import threading
+from typing import Dict, Optional
+
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+try:
+    from opentelemetry import metrics
+    from opentelemetry.exporter.prometheus import PrometheusMetricReader
+    from opentelemetry.metrics import set_meter_provider
+    from opentelemetry.sdk.metrics import MeterProvider as SDKMeterProvider
+    from prometheus_client import start_http_server
+
+    OPENTELEMETRY_AVAILABLE = True
+except ImportError:
+    LOG.warning("OpenTelemetry not available. pip install [opentelemetry]")
+    OPENTELEMETRY_AVAILABLE = False
+
+
+class OpenTelemetryMetricsCallback(TrainerCallback):
+    """
+    TrainerCallback that exports training metrics to OpenTelemetry/Prometheus.
+
+    This callback automatically tracks key training metrics including:
+    - Training loss
+    - Evaluation loss
+    - Learning rate
+    - Epoch progress
+    - Global step count
+    - Gradient norm
+
+    Metrics are exposed via HTTP endpoint for Prometheus scraping.
+    """
+
+    def __init__(self, cfg):
+        if not OPENTELEMETRY_AVAILABLE:
+            LOG.warning("OpenTelemetry not available, metrics will not be collected")
+            self.metrics_enabled = False
+            return
+
+        self.cfg = cfg
+        self.metrics_host = getattr(cfg, "otel_metrics_host", "localhost")
+        self.metrics_port = getattr(cfg, "otel_metrics_port", 8000)
+        self.metrics_enabled = True
+        self.server_started = False
+        self.metrics_lock = threading.Lock()
+
+        try:
+            # Create Prometheus metrics reader
+            prometheus_reader = PrometheusMetricReader()
+
+            # Create meter provider with Prometheus exporter
+            provider = SDKMeterProvider(metric_readers=[prometheus_reader])
+            set_meter_provider(provider)
+
+            # Get meter for creating metrics
+            self.meter = metrics.get_meter("axolotl.training")
+
+            # Create metrics
+            self._create_metrics()
+
+        except Exception as e:
+            LOG.warning(f"Failed to initialize OpenTelemetry metrics: {e}")
+            self.metrics_enabled = False
+
+    def _create_metrics(self):
+        """Create all metrics that will be tracked"""
+        self.train_loss_gauge = self.meter.create_gauge(
+            name="axolotl_train_loss",
+            description="Current training loss",
+            unit="1",
+        )
+
+        self.eval_loss_gauge = self.meter.create_gauge(
+            name="axolotl_eval_loss",
+            description="Current evaluation loss",
+            unit="1",
+        )
+
+        self.learning_rate_gauge = self.meter.create_gauge(
+            name="axolotl_learning_rate",
+            description="Current learning rate",
+            unit="1",
+        )
+
+        self.epoch_gauge = self.meter.create_gauge(
+            name="axolotl_epoch",
+            description="Current training epoch",
+            unit="1",
+        )
+
+        self.global_step_counter = self.meter.create_counter(
+            name="axolotl_global_steps",
+            description="Total training steps completed",
+            unit="1",
+        )
+
+        self.grad_norm_gauge = self.meter.create_gauge(
+            name="axolotl_gradient_norm",
+            description="Gradient norm",
+            unit="1",
+        )
+
+        self.memory_usage_gauge = self.meter.create_gauge(
+            name="axolotl_memory_usage",
+            description="Current memory usage in MB",
+            unit="MB",
+        )
+
+    def _start_metrics_server(self):
+        """Start the HTTP server for metrics exposure"""
+        if self.server_started:
+            return
+
+        try:
+            start_http_server(self.metrics_port, addr=self.metrics_host)
+            self.server_started = True
+            LOG.info(
+                f"OpenTelemetry metrics server started on http://{self.metrics_host}:{self.metrics_port}/metrics"
+            )
+
+        except Exception as e:
+            LOG.error(f"Failed to start OpenTelemetry metrics server: {e}")
+
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Called at the beginning of training"""
+        if not self.metrics_enabled:
+            return
+
+        self._start_metrics_server()
+        LOG.info("OpenTelemetry metrics collection started")
+
+    def on_log(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        logs: Optional[Dict[str, float]] = None,
+        **kwargs,
+    ):
+        """Called when logging occurs"""
+        if not self.metrics_enabled or not logs:
+            return
+
+        if "loss" in logs:
+            self.train_loss_gauge.set(logs["loss"])
+
+        if "eval_loss" in logs:
+            self.eval_loss_gauge.set(logs["eval_loss"])
+
+        if "learning_rate" in logs:
+            self.learning_rate_gauge.set(logs["learning_rate"])
+
+        if "epoch" in logs:
+            self.epoch_gauge.set(logs["epoch"])
+
+        if "grad_norm" in logs:
+            self.grad_norm_gauge.set(logs["grad_norm"])
+        if "memory_usage" in logs:
+            self.memory_usage_gauge.set(logs["memory_usage"])
+
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Called at the end of each training step"""
+        if not self.metrics_enabled:
+            return
+
+        # Update step counter and epoch
+        self.global_step_counter.add(1)
+        if state.epoch is not None:
+            self.epoch_gauge.set(state.epoch)
+
+    def on_evaluate(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        metrics: Optional[Dict[str, float]] = None,
+        **kwargs,
+    ):
+        """Called after evaluation"""
+        if not self.metrics_enabled or not metrics:
+            return
+
+        if "eval_loss" in metrics:
+            self.eval_loss_gauge.set(metrics["eval_loss"])
+
+        # Record any other eval metrics as gauges
+        for key, value in metrics.items():
+            if key.startswith("eval_") and isinstance(value, (int, float)):
+                # Create gauge for this metric if it doesn't exist
+                gauge_name = f"axolotl_{key}"
+                try:
+                    gauge = self.meter.create_gauge(
+                        name=gauge_name,
+                        description=f"Evaluation metric: {key}",
+                        unit="1",
+                    )
+                    gauge.set(value)
+                except Exception as e:
+                    LOG.warning(f"Failed to create/update metric {gauge_name}: {e}")
+
+    def on_train_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Called at the end of training"""
+        if not self.metrics_enabled:
+            return
+
+        LOG.info("Training completed. OpenTelemetry metrics collection finished.")
+        LOG.info(
+            f"Metrics are still available at http://{self.metrics_host}:{self.metrics_port}/metrics"
+        )
diff --git a/src/axolotl/utils/callbacks/perplexity.py b/src/axolotl/utils/callbacks/perplexity.py
new file mode 100644
index 0000000000..36cacfb817
--- /dev/null
+++ b/src/axolotl/utils/callbacks/perplexity.py
@@ -0,0 +1,85 @@
+"""callback to calculate perplexity as an evaluation metric."""
+
+from typing import Dict, List, Optional
+
+import torch
+from torch import Tensor
+from tqdm import tqdm
+from transformers.modeling_outputs import CausalLMOutput
+from transformers.modeling_utils import PreTrainedModel
+
+try:
+    from transformers.tokenization_python import PreTrainedTokenizer
+except ImportError:
+    from transformers.tokenization_utils import PreTrainedTokenizer
+
+from axolotl.utils.distributed import is_main_process
+
+
+class Perplexity:
+    """
+    Calculate perplexity as defined in https://huggingface.co/docs/transformers/en/perplexity.
+    This is a custom variant that doesn't re-tokenize the input or re-load the model.
+    """
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizer,
+        max_seq_len: int,
+        stride: int = 512,
+    ) -> None:
+        self.max_seq_len = max_seq_len
+        self.stride = stride
+        self.tokenizer = tokenizer
+        self.name = "perplexity"
+
+    def _feature_names(self) -> List[str]:
+        return ["references"]
+
+    def compute(
+        self,
+        model: PreTrainedModel,
+        references: Optional[List[str]] = None,
+    ) -> Dict[str, float]:
+        """
+        Compute perplexity in a fixed length sliding window across the sequence.
+        """
+        assert references is not None, "Missing parameter: references"
+
+        model.eval()
+
+        references_tokenized = self.tokenizer(
+            references, return_tensors="pt", padding=True, truncation=True
+        )
+        input_ids: Tensor = references_tokenized["input_ids"]  # type: ignore
+        input_ids = input_ids.to(model.device)
+
+        sequence_length = input_ids.size(1)
+
+        losses = []
+        prev_end_loc = 0
+        for begin_loc in tqdm(
+            range(0, sequence_length, self.stride), disable=not is_main_process()
+        ):
+            end_loc = min(begin_loc + self.max_seq_len, sequence_length)
+            trg_len = end_loc - prev_end_loc
+            input_ids_slice = input_ids[:, begin_loc:end_loc]
+            labels_slice = input_ids_slice.clone()
+            labels_slice[:, :-trg_len] = -100
+
+            with torch.no_grad():
+                outputs: CausalLMOutput = model(
+                    input_ids=input_ids_slice, labels=labels_slice
+                )
+
+            losses.append(outputs.loss)
+
+            prev_end_loc = end_loc
+            if end_loc == sequence_length:
+                break
+
+        perplexity = torch.exp(torch.stack(losses).mean()).item()
+
+        return {
+            "score": perplexity,
+        }
diff --git a/src/axolotl/utils/callbacks/profiler.py b/src/axolotl/utils/callbacks/profiler.py
new file mode 100644
index 0000000000..49e8c53882
--- /dev/null
+++ b/src/axolotl/utils/callbacks/profiler.py
@@ -0,0 +1,105 @@
+"""
+HF Trainer callback for creating pytorch profiling snapshots
+"""
+
+from pathlib import Path
+from pickle import dump  # nosec B403
+
+import torch
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+
+class PytorchProfilerCallback(TrainerCallback):
+    """
+    PyTorch Profiler callback to create snapshots of GPU memory usage at specified steps.
+
+    Also runs torch.profiler to produce a Chrome trace for timing analysis.
+    """
+
+    def __init__(self, steps_to_profile: int = 5, profiler_steps_start: int = 0):
+        # steps are 0 indexed, so to start at 0-th step, we start at beginning of first step,
+        # and finish at end of last step, so 5 steps_to_profile is steps [0, 1, 2, 3, 4]
+        self.profiler_steps_end = profiler_steps_start + steps_to_profile - 1
+        if profiler_steps_start == 0:
+            # start recording memory allocations before everything is allocated, because if we start
+            # at the beginning of step 0, we won't have any memory allocations in the traces
+            torch.cuda.memory._record_memory_history(enabled="all", stacks="all")
+            profiler_steps_start = -1
+        self.profiler_steps_start = profiler_steps_start
+        self._profiler = None
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        if state.global_step == self.profiler_steps_start:
+            torch.cuda.memory._record_memory_history(enabled="all", stacks="all")
+
+        # Start torch.profiler on the first profiled step
+        if state.global_step == max(self.profiler_steps_start, 0):
+            profiler = torch.profiler.profile(
+                activities=[
+                    torch.profiler.ProfilerActivity.CPU,
+                    torch.profiler.ProfilerActivity.CUDA,
+                ],
+                record_shapes=True,
+                profile_memory=True,
+                with_stack=True,
+            )
+            profiler.__enter__()
+            self._profiler = profiler
+
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        if state.global_step == self.profiler_steps_end:
+            snapshot = torch.cuda.memory._snapshot()
+            with open(Path(args.output_dir) / "snapshot.pickle", "wb") as fout:
+                dump(snapshot, fout)
+
+            # tell CUDA to stop recording memory allocations now
+            torch.cuda.memory._record_memory_history(enabled=None)
+
+            # Stop and export torch.profiler trace
+            if self._profiler is not None:
+                self._profiler.__exit__(None, None, None)
+                trace_path = Path(args.output_dir) / "profiler_trace.json"
+                self._profiler.export_chrome_trace(str(trace_path))
+                self._profiler = None
+
+    def on_train_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        # make sure to record if we happen to have more steps than steps to profile
+        if (
+            state.global_step >= self.profiler_steps_start
+            and state.global_step < self.profiler_steps_end
+        ):
+            snapshot = torch.cuda.memory._snapshot()
+            with open(Path(args.output_dir) / "snapshot.pickle", "wb") as fout:
+                dump(snapshot, fout)
+
+            # tell CUDA to stop recording memory allocations now
+            torch.cuda.memory._record_memory_history(enabled=None)
+
+        if self._profiler is not None:
+            self._profiler.__exit__(None, None, None)
+            trace_path = Path(args.output_dir) / "profiler_trace.json"
+            self._profiler.export_chrome_trace(str(trace_path))
+            self._profiler = None
diff --git a/src/axolotl/utils/callbacks/qat.py b/src/axolotl/utils/callbacks/qat.py
new file mode 100644
index 0000000000..446b340b6f
--- /dev/null
+++ b/src/axolotl/utils/callbacks/qat.py
@@ -0,0 +1,50 @@
+"""QAT Callback for HF Causal Trainer"""
+
+from functools import partial
+
+from torch import nn
+from torchao.quantization.qat.embedding import FakeQuantizedEmbedding
+from torchao.quantization.qat.linear import FakeQuantizedLinear
+from transformers import TrainerCallback
+
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.quantization import QATConfig
+
+LOG = get_logger(__name__)
+
+
+def toggle_fake_quant(mod: nn.Module, enable: bool):
+    """
+    Toggle fake quantization for any fake quantized linear or embedding layers in the model.
+
+    Args:
+        mod: The module to toggle fake quantization for.
+        enable: Whether to enable or disable fake quantization.
+    """
+    if isinstance(mod, (FakeQuantizedLinear, FakeQuantizedEmbedding)):
+        if (
+            isinstance(mod, FakeQuantizedLinear)
+            and mod.activation_fake_quantizer is not None
+            and hasattr(mod.activation_fake_quantizer, "enabled")
+        ):
+            mod.activation_fake_quantizer.enabled = enable
+        if hasattr(mod.weight_fake_quantizer, "enabled"):
+            mod.weight_fake_quantizer.enabled = enable
+
+
+class QATCallback(TrainerCallback):
+    """
+    Callback to toggle fake quantization for the model.
+    """
+
+    def __init__(self, cfg: QATConfig):
+        self.cfg = cfg
+
+    def on_step_begin(self, args, state, control, model, **kwargs):
+        if self.cfg.fake_quant_after_n_steps is not None:
+            if state.global_step == 0:
+                LOG.info(f"Disabling fake quantization at step {state.global_step}")
+                model.apply(partial(toggle_fake_quant, enable=False))
+            elif state.global_step == self.cfg.fake_quant_after_n_steps:
+                LOG.info(f"Enabling fake quantization at step {state.global_step}")
+                model.apply(partial(toggle_fake_quant, enable=True))
diff --git a/src/axolotl/utils/callbacks/swanlab.py b/src/axolotl/utils/callbacks/swanlab.py
new file mode 100644
index 0000000000..4ebf2e61ef
--- /dev/null
+++ b/src/axolotl/utils/callbacks/swanlab.py
@@ -0,0 +1,248 @@
+"""Callbacks for SwanLab integration"""
+
+from __future__ import annotations
+
+import json
+import os
+from shutil import copyfile
+from tempfile import NamedTemporaryFile
+from typing import TYPE_CHECKING
+
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+from axolotl.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from axolotl.core.training_args import AxolotlTrainingArguments
+
+LOG = get_logger(__name__)
+
+
+class CustomSwanLabCallback(TrainerCallback):
+    """
+    Lightweight SwanLab callback that directly logs metrics without using
+    SwanLab's transformers integration (which requires omegaconf).
+
+    This avoids the antlr4 version conflict between omegaconf and axolotl.
+    """
+
+    def __init__(self):
+        self._initialized = False
+        self.swanlab = None
+
+    def setup(self):
+        """Lazy initialization of SwanLab"""
+        if self._initialized:
+            return
+
+        try:
+            import swanlab
+
+            self.swanlab = swanlab
+
+            # Check if SwanLab run is initialized
+            if swanlab.get_run() is None:
+                LOG.warning("SwanLab run is not initialized")
+                return
+
+            self._initialized = True
+            LOG.info("CustomSwanLabCallback initialized successfully")
+        except ImportError:
+            LOG.error("SwanLab is not installed")
+
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Called at the beginning of training"""
+        if not state.is_world_process_zero:
+            return control
+
+        self.setup()
+
+        if not self._initialized:
+            return control
+
+        # Log training configuration
+        try:
+            self.swanlab.config.update(
+                {
+                    "train_batch_size": args.per_device_train_batch_size,
+                    "eval_batch_size": args.per_device_eval_batch_size,
+                    "learning_rate": args.learning_rate,
+                    "num_train_epochs": args.num_train_epochs,
+                    "max_steps": args.max_steps,
+                    "warmup_steps": args.warmup_steps,
+                    "logging_steps": args.logging_steps,
+                    "save_steps": args.save_steps,
+                    "gradient_accumulation_steps": args.gradient_accumulation_steps,
+                }
+            )
+            LOG.debug("Training configuration logged to SwanLab")
+        except Exception as err:
+            LOG.warning(f"Failed to log training config: {err}")
+
+        return control
+
+    def on_log(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        logs=None,
+        **kwargs,
+    ):
+        """Called when logging metrics"""
+        if not state.is_world_process_zero:
+            return control
+
+        if not self._initialized:
+            self.setup()
+
+        if not self._initialized or logs is None:
+            return control
+
+        # Log metrics to SwanLab
+        try:
+            # Filter out non-numeric values and prepare for logging
+            metrics = {}
+            for key, value in logs.items():
+                if isinstance(value, (int, float)):
+                    # Use step from state
+                    metrics[key] = value
+
+            if metrics and state.global_step is not None:
+                self.swanlab.log(metrics, step=state.global_step)
+        except Exception as err:
+            LOG.warning(f"Failed to log metrics to SwanLab: {err}")
+
+        return control
+
+    def on_train_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Called at the end of training"""
+        if not state.is_world_process_zero:
+            return control
+
+        if self._initialized:
+            LOG.info("Training completed. SwanLab logs are available.")
+
+        return control
+
+
+class SaveAxolotlConfigtoSwanLabCallback(TrainerCallback):
+    """Callback to save axolotl config to SwanLab"""
+
+    def __init__(self, axolotl_config_path):
+        self.axolotl_config_path = axolotl_config_path
+
+    def on_train_begin(
+        self,
+        args: AxolotlTrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        if state.is_world_process_zero:
+            try:
+                import swanlab
+
+                # Check if SwanLab is initialized
+                if swanlab.get_run() is None:
+                    LOG.warning(
+                        "SwanLab run is not initialized. Please initialize SwanLab before training."
+                    )
+                    return control
+
+                # Log Axolotl config as artifact
+                with NamedTemporaryFile(
+                    mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
+                ) as temp_file:
+                    copyfile(self.axolotl_config_path, temp_file.name)
+
+                    # Log config file to SwanLab
+                    with open(temp_file.name, "r", encoding="utf-8") as config_file:
+                        swanlab.log(
+                            {
+                                "axolotl_config": swanlab.Text(
+                                    config_file.read(), caption="Axolotl Config"
+                                )
+                            }
+                        )
+
+                    LOG.info(
+                        "The Axolotl config has been saved to the SwanLab run under logs."
+                    )
+
+                    # Clean up temp file
+                    os.unlink(temp_file.name)
+
+            except ImportError:
+                LOG.warning(
+                    "SwanLab is not installed. Install it with: pip install swanlab"
+                )
+            except (FileNotFoundError, ConnectionError) as err:
+                LOG.warning(f"Error while saving Axolotl config to SwanLab: {err}")
+
+            # Log DeepSpeed config if available
+            if args.deepspeed:
+                try:
+                    import swanlab
+
+                    with NamedTemporaryFile(
+                        mode="w",
+                        delete=False,
+                        suffix=".json",
+                        prefix="deepspeed_config_",
+                    ) as temp_file:
+                        skip_upload = False
+                        if isinstance(args.deepspeed, dict):
+                            json.dump(args.deepspeed, temp_file, indent=4)
+                        elif isinstance(args.deepspeed, str) and os.path.exists(
+                            args.deepspeed
+                        ):
+                            copyfile(args.deepspeed, temp_file.name)
+                        else:
+                            skip_upload = True
+
+                        if not skip_upload:
+                            temp_file.flush()
+                            with open(
+                                temp_file.name, "r", encoding="utf-8"
+                            ) as ds_config_file:
+                                swanlab.log(
+                                    {
+                                        "deepspeed_config": swanlab.Text(
+                                            ds_config_file.read(),
+                                            caption="DeepSpeed Config",
+                                        )
+                                    }
+                                )
+                            LOG.info(
+                                "The DeepSpeed config has been saved to the SwanLab run under logs."
+                            )
+
+                        # Clean up temp file
+                        os.unlink(temp_file.name)
+
+                except (FileNotFoundError, ConnectionError) as err:
+                    LOG.warning(
+                        f"Error while saving DeepSpeed config to SwanLab: {err}"
+                    )
+                except ImportError:
+                    pass
+
+        return control
diff --git a/src/axolotl/utils/callbacks/tokens_per_second.py b/src/axolotl/utils/callbacks/tokens_per_second.py
new file mode 100644
index 0000000000..d4f8e98c11
--- /dev/null
+++ b/src/axolotl/utils/callbacks/tokens_per_second.py
@@ -0,0 +1,50 @@
+"""A callback for calculating tokens per second during training."""
+
+import json
+import os
+
+import torch
+from transformers import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+
+from axolotl.core.trainers.constants import TOKENS_STATE_FILE
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class TokensPerSecondCallback(TrainerCallback):
+    """Restore the cumulative token counters when resuming from a checkpoint.
+
+    Throughput itself is computed in the trainer's ``log()`` from deltas of the
+    cumulative ``trainable`` counter, so it is unaffected by
+    gradient_accumulation_steps and logging_steps.
+    """
+
+    def __init__(self, resume_from_checkpoint=None):
+        super().__init__()
+        self.resume_from_checkpoint = resume_from_checkpoint
+
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):  # pylint: disable=unused-argument
+        """Restore total_tokens state when resuming from checkpoint."""
+        if not isinstance(self.resume_from_checkpoint, str):
+            return
+        tokens_state_path = os.path.join(self.resume_from_checkpoint, TOKENS_STATE_FILE)
+        if os.path.isfile(tokens_state_path):
+            with open(tokens_state_path, "r", encoding="utf-8") as f:
+                tokens_state = json.load(f)
+            state.tokens = {
+                "total": torch.tensor(tokens_state.get("total", 0)),
+                "trainable": torch.tensor(tokens_state.get("trainable", 0)),
+            }
+            LOG.info(f"Restored total_tokens: {state.tokens['total']}")
diff --git a/src/axolotl/utils/callbacks/trackio_.py b/src/axolotl/utils/callbacks/trackio_.py
new file mode 100644
index 0000000000..8249321f62
--- /dev/null
+++ b/src/axolotl/utils/callbacks/trackio_.py
@@ -0,0 +1,44 @@
+"""Trackio module for trainer callbacks"""
+
+from typing import TYPE_CHECKING
+
+import trackio
+from transformers import TrainerCallback, TrainerControl, TrainerState
+
+from axolotl.utils.distributed import is_main_process
+from axolotl.utils.environment import is_package_version_ge
+from axolotl.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from axolotl.core.training_args import AxolotlTrainingArguments
+
+LOG = get_logger(__name__)
+
+
+class SaveAxolotlConfigtoTrackioCallback(TrainerCallback):
+    """Callback for trackio integration"""
+
+    def __init__(self, axolotl_config_path):
+        self.axolotl_config_path = axolotl_config_path
+
+    def on_train_begin(
+        self,
+        args: "AxolotlTrainingArguments",
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        if is_main_process():
+            try:
+                if not is_package_version_ge("trackio", "0.11.0"):
+                    LOG.warning(
+                        "Trackio version 0.11.0 or higher is required to save config files. "
+                        "Please upgrade trackio: pip install --upgrade trackio"
+                    )
+                    return control
+
+                trackio.save(self.axolotl_config_path)
+                LOG.info("The Axolotl config has been saved to Trackio.")
+            except (FileNotFoundError, ConnectionError, AttributeError) as err:
+                LOG.warning(f"Error while saving Axolotl config to Trackio: {err}")
+        return control
diff --git a/src/axolotl/utils/chat_templates.py b/src/axolotl/utils/chat_templates.py
deleted file mode 100644
index 01b1473568..0000000000
--- a/src/axolotl/utils/chat_templates.py
+++ /dev/null
@@ -1,33 +0,0 @@
-"""
-This module provides functionality for selecting chat templates based on user choices.
-These templates are used for formatting messages in a conversation.
-"""
-
-
-def chat_templates(user_choice: str):
-    """
-    Finds the correct chat_template for the tokenizer_config.
-
-    Args:
-        user_choice (str): The user's choice of template.
-
-    Returns:
-        str: The chosen template string.
-
-    Raises:
-        ValueError: If the user_choice is not found in the templates.
-    """
-
-    templates = {
-        "alpaca": "{% for message in messages %}{% if message['role'] == 'user' %}{{ '### Instruction: ' + message['content'] + '\n\n' }}{% elif message['role'] == 'assistant' %}{{ '### Response: ' + message['content'] + eos_token}}{% endif %}{% endfor %}",
-        "inst": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}",  # I don't know what this one is called. Used by Mistral/Mixtral.
-        "chatml": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
-        "gemma": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model\n'}}{% endif %}",
-        "cohere": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% elif false == true %}{% set loop_messages = messages %}{% set system_message = 'You are Command-R, a brilliant, sophisticated, AI-assistant trained to assist human users by providing thorough responses. You are trained by Cohere.' %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% if system_message != false %}{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' + system_message + '<|END_OF_TURN_TOKEN|>' }}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|START_OF_TURN_TOKEN|><|USER_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}{% elif message['role'] == 'assistant' %}{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>'  + content.strip() + '<|END_OF_TURN_TOKEN|>' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' }}{% endif %}",
-        "llama3": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% else %}{{ eos_token }}{% endif %}",
-    }
-
-    if user_choice in templates:
-        return templates[user_choice]
-
-    raise ValueError(f"Template '{user_choice}' not found.")
diff --git a/src/axolotl/utils/chat_templates/__init__.py b/src/axolotl/utils/chat_templates/__init__.py
new file mode 100644
index 0000000000..337417c7dc
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/__init__.py
@@ -0,0 +1,20 @@
+"""
+This module provides functionality for selecting chat templates based on user choices.
+These templates are used for formatting messages in a conversation.
+"""
+
+from .base import (
+    _CHAT_TEMPLATES,
+    extract_chat_template_args,
+    get_chat_template,
+    get_chat_template_from_config,
+    register_chat_template,
+)
+
+__all__ = [
+    "get_chat_template",
+    "extract_chat_template_args",
+    "get_chat_template_from_config",
+    "register_chat_template",
+    "_CHAT_TEMPLATES",
+]
diff --git a/src/axolotl/utils/chat_templates/base.py b/src/axolotl/utils/chat_templates/base.py
new file mode 100644
index 0000000000..11d15fc1da
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/base.py
@@ -0,0 +1,125 @@
+"""
+utility functions for chat templates
+"""
+
+import os
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from axolotl.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerBase
+
+LOG = get_logger("axolotl.utils.chat_templates")
+
+_JINJA_TEMPLATE_CHOICE = "jinja"
+_DEFAULT_TEMPLATE_CHOICE = "tokenizer_default"
+_DEFAULT_FALLBACK_CHATML_TEMPLATE_CHOICE_PREFIX = "tokenizer_default_fallback_"
+
+TEMPLATE_DIR = os.path.join(os.path.dirname(__file__), "templates")
+_CHAT_TEMPLATES: dict[str, str] = {}
+for filename in [f for f in os.listdir(TEMPLATE_DIR) if f.endswith(".jinja")]:
+    with open(os.path.join(TEMPLATE_DIR, filename), "r", encoding="utf-8") as f:
+        _CHAT_TEMPLATES[filename[:-6]] = f.read()
+
+
+def get_chat_template(
+    user_choice: str,
+    jinja_template: str | None = None,
+    tokenizer: Optional["PreTrainedTokenizerBase"] = None,
+) -> str:
+    """
+    Finds the correct chat_template based on the user's choice, jinja_template, and tokenizer.
+
+    Args:
+        user_choice (str): The user's choice of template.
+        jinja_template (str, optional): The jinja template string or Path to a valid jinja template file. Defaults to None.
+        tokenizer (PreTrainedTokenizerBase, optional): The tokenizer. Defaults to None.
+
+    Returns:
+        str: The chosen template string.
+
+    Raises:
+        ValueError: If the user_choice is not found in the templates.
+    """
+    if user_choice == _JINJA_TEMPLATE_CHOICE:
+        if not jinja_template:
+            raise ValueError(
+                f"`jinja_template` cannot be None when `chat_template` choice is {_JINJA_TEMPLATE_CHOICE}"
+            )
+        if os.path.exists(jinja_template) and os.path.isfile(jinja_template):
+            with open(jinja_template, "r", encoding="utf-8") as file:
+                jinja_template = file.read()
+        return jinja_template
+
+    if user_choice == _DEFAULT_TEMPLATE_CHOICE:
+        if not tokenizer:
+            raise ValueError(
+                f"`tokenizer` cannot be None when chat_template choice is {_DEFAULT_TEMPLATE_CHOICE}"
+            )
+        if not tokenizer.chat_template:
+            raise ValueError(
+                f"`chat_template choice is {_DEFAULT_TEMPLATE_CHOICE} but tokenizer's chat_template is null. "
+                f"Please add a chat_template in tokenizer config"
+            )
+        return tokenizer.chat_template  # type: ignore
+
+    if user_choice.startswith(_DEFAULT_FALLBACK_CHATML_TEMPLATE_CHOICE_PREFIX):
+        if not tokenizer:
+            raise ValueError(
+                f"`tokenizer` cannot be None when chat_template choice starts with {_DEFAULT_FALLBACK_CHATML_TEMPLATE_CHOICE_PREFIX}"
+            )
+        if tokenizer.chat_template:
+            return tokenizer.chat_template  # type: ignore
+
+        user_choice = user_choice[
+            len(_DEFAULT_FALLBACK_CHATML_TEMPLATE_CHOICE_PREFIX) :
+        ]
+        LOG.warning(
+            f"No chat template found on tokenizer, falling back to {user_choice}. It is recommended to set --train_on_inputs to True for the model to learn this chat template."
+        )
+
+    if user_choice in _CHAT_TEMPLATES:
+        return _CHAT_TEMPLATES[user_choice]
+
+    raise ValueError(f"Template '{user_choice}' not found.")
+
+
+def extract_chat_template_args(cfg, ds_cfg: Dict[str, Any] | None = None):
+    if ds_cfg and ds_cfg.get("chat_template"):
+        chat_template_choice = ds_cfg.get("chat_template") or _DEFAULT_TEMPLATE_CHOICE
+        chat_template_jinja = ds_cfg.get("chat_template_jinja")
+    else:
+        chat_template_choice = cfg.get("chat_template") or _DEFAULT_TEMPLATE_CHOICE
+        chat_template_jinja = cfg.get("chat_template_jinja")
+    return chat_template_choice, chat_template_jinja
+
+
+def get_chat_template_from_config(
+    cfg,
+    ds_cfg: Dict[str, Any] | None = None,
+    tokenizer: Optional["PreTrainedTokenizerBase"] = None,
+) -> str:
+    chat_template_choice, chat_template_jinja = extract_chat_template_args(
+        cfg=cfg, ds_cfg=ds_cfg
+    )
+    return get_chat_template(
+        user_choice=chat_template_choice,
+        jinja_template=chat_template_jinja,
+        tokenizer=tokenizer,
+    )
+
+
+def register_chat_template(template_name: str, chat_template: str):
+    """
+    Registers chat templates.
+
+    Args:
+        template_name (str): The name of the template.
+        chat_template (str): The template string.
+    """
+
+    if template_name in _CHAT_TEMPLATES:
+        raise ValueError(f"Template '{template_name}' already exists.")
+
+    _CHAT_TEMPLATES[template_name] = chat_template
diff --git a/src/axolotl/utils/chat_templates/templates/alpaca.jinja b/src/axolotl/utils/chat_templates/templates/alpaca.jinja
new file mode 100644
index 0000000000..5e9d63c424
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/alpaca.jinja
@@ -0,0 +1,8 @@
+{{ bos_token }}{% for message in messages %}{% if message['role'] == 'system' and loop.first %}{{ message['content'] }}{% elif message['role'] == 'user' %}{{ '### Instruction:
+' + message['content'] }}{% elif message['role'] == 'assistant' %}{{ '### Response:
+' + message['content'] + eos_token }}{% endif %}{% if not loop.last %}{{ '
+
+' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '
+
+### Response:
+' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/aya.jinja b/src/axolotl/utils/chat_templates/templates/aya.jinja
new file mode 100644
index 0000000000..97e54d4b10
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/aya.jinja
@@ -0,0 +1 @@
+{{ bos_token }}{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% elif false == true %}{% set loop_messages = messages %}{% set system_message = 'You are Aya, a brilliant, sophisticated, AI-assistant trained to assist human users by providing thorough responses. You are trained by Cohere.' %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% if system_message != false %}{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' + system_message + '<|END_OF_TURN_TOKEN|>' }}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|START_OF_TURN_TOKEN|><|USER_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}{% elif message['role'] == 'assistant' %}{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>'  + content.strip() + '<|END_OF_TURN_TOKEN|>' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/chatml.jinja b/src/axolotl/utils/chat_templates/templates/chatml.jinja
new file mode 100644
index 0000000000..2116e45cab
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/chatml.jinja
@@ -0,0 +1,4 @@
+{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>' + '
+'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/cohere.jinja b/src/axolotl/utils/chat_templates/templates/cohere.jinja
new file mode 100644
index 0000000000..638ce5ef2f
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/cohere.jinja
@@ -0,0 +1 @@
+{{ bos_token }}{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% elif false == true %}{% set loop_messages = messages %}{% set system_message = 'You are Command-R, a brilliant, sophisticated, AI-assistant trained to assist human users by providing thorough responses. You are trained by Cohere.' %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% if system_message != false %}{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' + system_message + '<|END_OF_TURN_TOKEN|>' }}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|START_OF_TURN_TOKEN|><|USER_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}{% elif message['role'] == 'assistant' %}{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>'  + content.strip() + '<|END_OF_TURN_TOKEN|>' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/command_a.jinja b/src/axolotl/utils/chat_templates/templates/command_a.jinja
new file mode 100644
index 0000000000..ef0594172d
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/command_a.jinja
@@ -0,0 +1,210 @@
+{{ bos_token }}{% if documents %}
+{% set tools = [] %}
+{%- macro document_turn(documents) -%}
+{# format documents into chat turn #}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>I will look through the document to address the users needs.<|END_THINKING|><|START_ACTION|>[
+    {"tool_call_id": "0", "tool_name": "direct-injected-document", "parameters": {}}
+]<|END_ACTION|><|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[
+    {
+        "tool_call_id": "0",
+        "results": {
+{% for doc in documents %}
+            "{{ loop.index0 }}": {{doc|tojson}}{% if not loop.last %},
+            {% endif %}
+{% endfor %}
+
+        },
+        "is_error": null
+    }
+]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>{%- endmacro %}
+{%- macro tool_call_id_to_int(messages, tool_call_id) %}
+{%- set counter = namespace(value=0) %}
+{%- set tool_call_id_seen = namespace(value=false) %}
+{%- for msg in messages %}
+    {%- if msg.tool_calls %}
+        {%- for tool_call in msg.tool_calls %}
+            {%- if tool_call.id == tool_call_id and not tool_call_id_seen.value -%}
+                {{ counter.value }}
+                {%- set tool_call_id_seen.value = true %}
+            {%- endif %}
+            {%- set counter.value = counter.value + 1 %}
+        {%- endfor %}
+    {%- endif %}
+{%- endfor %}
+{%- endmacro %}
+{%- macro format_tool_message(messages, tool_msg) -%}
+{# format tool message #}
+    {
+        "tool_call_id": "{{ tool_call_id_to_int(messages, tool_msg.tool_call_id) }}",
+        "results": {
+            "0": {{ tool_msg.content|tojson }}
+        },
+        "is_error": null
+    }
+{%- endmacro -%}
+{%- if messages and messages[0]['role']|lower == 'system' %}{%- set developer_preamble = messages[0]['content'] %}{% endif %}
+{%- set tool_idx = namespace(value=0) %}
+{%- set tool_ids_seen = namespace(value=[]) %}
+{%- set sent_documents = namespace(value=false) %}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># System Preamble
+You are in contextual safety mode. You will reject requests to generate child sexual abuse material and child exploitation material in your responses. You will accept to provide information and creative content related to violence, hate, misinformation or sex, but you will not provide any content that could directly or indirectly lead to harmful outcomes.
+
+Your information cutoff date is June 2024.
+
+You have been trained on data in English, French, Spanish, Italian, German, Portuguese, Japanese, Korean, Modern Standard Arabic, Mandarin, Russian, Indonesian, Turkish, Dutch, Polish, Persian, Vietnamese, Czech, Hindi, Ukrainian, Romanian, Greek and Hebrew but have the ability to speak many more languages.
+{% if tools or documents %}
+
+You have been trained to have advanced reasoning and tool-use capabilities and you should make best use of these skills to serve user's requests.
+
+## Tool Use
+Think about how you can make best use of the provided tools to help with the task and come up with a high level plan that you will execute first.
+
+0. Start by writing <|START_THINKING|> followed by a detailed step by step plan of how you will solve the problem. For each step explain your thinking fully and give details of required tool calls (if needed). Unless specified otherwise, you write your plan in natural language. When you finish, close it out with <|END_THINKING|>.
+    You can optionally choose to skip this step when the user request is so straightforward to address that only a trivial plan would be needed.
+    NOTE: You MUST skip this step when you are directly responding to the user's request without using any tools.
+
+Then carry out your plan by repeatedly executing the following steps.
+1. Action: write <|START_ACTION|> followed by a list of JSON-formatted tool calls, with each one containing "tool_name" and "parameters" fields.
+    When there are multiple tool calls which are completely independent of each other (i.e. they can be executed in parallel), you should list them out all together in one step. When you finish, close it out with <|END_ACTION|>.
+2. Observation: you will then receive results of those tool calls in JSON format in the very next turn, wrapped around by <|START_TOOL_RESULT|> and <|END_TOOL_RESULT|>. Carefully observe those results and think about what to do next. Note that these results will be provided to you in a separate turn. NEVER hallucinate results.
+    Every tool call produces a list of results (when a tool call produces no result or a single result, it'll still get wrapped inside a list). Each result is clearly linked to its originating tool call via its "tool_call_id".
+3. Reflection: start the next turn by writing <|START_THINKING|> followed by what you've figured out so far, any changes you need to make to your plan, and what you will do next. When you finish, close it out with <|END_THINKING|>.
+    You can optionally choose to skip this step when everything is going according to plan and no special pieces of information or reasoning chains need to be recorded.
+    NOTE: You MUST skip this step when you are done with tool-use actions and are ready to respond to the user.
+
+You can repeat the above 3 steps multiple times (could be 0 times too if no suitable tool calls are available or needed), until you decide it's time to finally respond to the user.
+
+4. Response: then break out of the loop and write <|START_RESPONSE|> followed by a piece of text which serves as a response to the user's last request. Use all previous tool calls and results to help you when formulating your response. When you finish, close it out with <|END_RESPONSE|>.
+{% if enable_citations %}
+
+## Grounding
+Importantly, note that "Reflection" and "Response" above can be grounded.
+Grounding means you associate pieces of texts (called "spans") with those specific tool results that support them (called "sources"). And you use a pair of tags "<co>" and "</co>" to indicate when a span can be grounded onto a list of sources, listing them out in the closing tag. Sources from the same tool call are grouped together and listed as "{tool_call_id}:[{list of result indices}]", before they are joined together by ",". E.g., "<co>span</co: 0:[1,2],1:[0]>" means that "span" is supported by result 1 and 2 from "tool_call_id=0" as well as result 0 from "tool_call_id=1".
+{% endif %}
+
+## Available Tools
+Here is the list of tools that you have available to you.
+You can ONLY use the tools listed here. When a tool is not listed below, it is NOT available and you should NEVER attempt to use it.
+Each tool is represented as a JSON object with fields like "name", "description", "parameters" (per JSON Schema), and optionally, "responses" (per JSON Schema).
+
+```json
+[
+{% if documents %}
+    {"name": "direct-injected-document", "description": "This is a special tool to directly inject user-uploaded documents into the chat as additional context. DO NOT use this tool by yourself!", "parameters": {"type": "object", "properties": {}, "required": []}, "responses": {"200": {"description": "Successfully returned a list of chunked text snippets from the directly uploaded documents.", "content": {"application/json": {"schema": {"type": "array", "items": {"type": "object", "required": ["url", "snippet"], "properties": {"url": {"type": "string", "description": "The url of the uploaded document."}, "snippet": {"type": "string", "description": "The text snippet for the returned document chunk."}}}}}}}}}{%- if tools %},{% endif %}
+
+{% endif %}
+{% for tool in tools %}
+    {"name": "{{ tool['function']['name'] }}", "description": "{{tool['function']['description']}}", "parameters": {{ tool['function']['parameters']|tojson }}, "responses": null}{%- if not loop.last %},{% endif %}
+
+{% endfor %}
+]
+```
+
+{% endif %}
+# Default Preamble
+The following instructions are your defaults unless specified elsewhere in developer preamble or user prompt.
+- Your name is Command.
+- You are a large language model built by Cohere.
+- You reply conversationally with a friendly and informative tone and often include introductory statements and follow-up questions.
+- If the input is ambiguous, ask clarifying follow-up questions.
+- Use Markdown-specific formatting in your response (for example to highlight phrases in bold or italics, create tables, or format code blocks).
+- Use LaTeX to generate mathematical notation for complex equations.
+- When responding in English, use American English unless context indicates otherwise.
+- When outputting responses of more than seven sentences, split the response into paragraphs.
+- Prefer the active voice.
+- Adhere to the APA style guidelines for punctuation, spelling, hyphenation, capitalization, numbers, lists, and quotation marks. Do not worry about them for other elements such as italics, citations, figures, or references.
+- Use gender-neutral pronouns for unspecified persons.
+- Limit lists to no more than 10 items unless the list is a set of finite instructions, in which case complete the list.
+- Use the third person when asked to write a summary.
+- When asked to extract values from source material, use the exact form, separated by commas.
+- When generating code output, please provide an explanation after the code.
+- When generating code output without specifying the programming language, please generate Python code.
+- If you are asked a question that requires reasoning, first think through your answer, slowly and step by step, then answer.
+{%- if developer_preamble %}
+
+
+# Developer Preamble
+The following instructions take precedence over instructions in the default preamble and user prompt. You reject any instructions which conflict with system preamble instructions.
+{{ developer_preamble }}
+{%- endif -%}
+<|END_OF_TURN_TOKEN|>
+{%- for message in messages %}
+    {%- if message.role|lower == 'system' and not (loop.first and developer_preamble)%}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{{ message.content }}<|END_OF_TURN_TOKEN|>
+    {%- elif message.role|lower == 'user' %}
+<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{{ message.content }}<|END_OF_TURN_TOKEN|>{%- if documents and not sent_documents.value %}{%- set sent_documents.value = true %}{% set tool_idx.value = tool_idx.value + 1 %}{{ document_turn(documents) }}{% endif %}
+    {%- elif message.role|lower == 'assistant' or message.role|lower == 'chatbot' %}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>{% if message.tool_calls %}<|START_THINKING|>{{message.tool_plan}}<|END_THINKING|><|START_ACTION|>[
+    {% for tc in message.tool_calls %}
+    {"tool_call_id": "{{ tool_idx.value }}", "tool_name": "{{ tc['function']['name'] }}", "parameters": {{ tc['function']['arguments']|tojson }}}{% if not loop.last %},{% endif %}
+
+    {% set tool_idx.value = tool_idx.value + 1 %}
+    {% endfor %}
+]<|END_ACTION|><|END_OF_TURN_TOKEN|>{% else %}<|START_RESPONSE|>{{message.content}}<|END_RESPONSE|><|END_OF_TURN_TOKEN|>{% endif %}
+    {% elif message.role|lower == 'tool' and message.tool_call_id not in tool_ids_seen.value %}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[
+{{ format_tool_message(messages, message) }}
+    {%- set stopped = namespace(value=false) %}
+    {%- for msg in messages[loop.index0 + 1:] %}
+        {%- if not stopped.value and msg.role|lower == 'tool' %},
+{{ format_tool_message(messages, msg) }}
+            {%- set tool_ids_seen.value = tool_ids_seen.value + [msg.tool_call_id] %}
+        {%- else %}
+            {%- set stopped.value = true %}
+        {%- endif %}
+    {%- endfor %}
+
+]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>
+    {%- endif %}
+{%- endfor %}<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>
+{%- else -%}
+{%- if messages and messages[0]['role']|lower == 'system' %}{%- set developer_preamble = messages[0]['content'] %}{% endif %}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># System Preamble
+{% if safety_mode|upper == 'STRICT' -%}
+You are in strict safety mode. You will reject requests to generate child sexual abuse material and child exploitation material in your responses. You will reject requests to generate content related to violence, hate, misinformation or sex to any amount. You will avoid using profanity. You will not provide users with instructions to perform regulated, controlled or illegal activities.
+{%- else -%}
+You are in contextual safety mode. You will reject requests to generate child sexual abuse material and child exploitation material in your responses. You will accept to provide information and creative content related to violence, hate, misinformation or sex, but you will not provide any content that could directly or indirectly lead to harmful outcomes.
+{%- endif %}
+
+
+Your information cutoff date is June 2024.
+
+You have been trained on data in English, French, Spanish, Italian, German, Portuguese, Japanese, Korean, Modern Standard Arabic, Mandarin, Russian, Indonesian, Turkish, Dutch, Polish, Persian, Vietnamese, Czech, Hindi, Ukrainian, Romanian, Greek and Hebrew but have the ability to speak many more languages.
+
+# Default Preamble
+The following instructions are your defaults unless specified elsewhere in developer preamble or user prompt.
+- Your name is Command.
+- You are a large language model built by Cohere.
+- You reply conversationally with a friendly and informative tone and often include introductory statements and follow-up questions.
+- If the input is ambiguous, ask clarifying follow-up questions.
+- Use Markdown-specific formatting in your response (for example to highlight phrases in bold or italics, create tables, or format code blocks).
+- Use LaTeX to generate mathematical notation for complex equations.
+- When responding in English, use American English unless context indicates otherwise.
+- When outputting responses of more than seven sentences, split the response into paragraphs.
+- Prefer the active voice.
+- Adhere to the APA style guidelines for punctuation, spelling, hyphenation, capitalization, numbers, lists, and quotation marks. Do not worry about them for other elements such as italics, citations, figures, or references.
+- Use gender-neutral pronouns for unspecified persons.
+- Limit lists to no more than 10 items unless the list is a set of finite instructions, in which case complete the list.
+- Use the third person when asked to write a summary.
+- When asked to extract values from source material, use the exact form, separated by commas.
+- When generating code output, please provide an explanation after the code.
+- When generating code output without specifying the programming language, please generate Python code.
+- If you are asked a question that requires reasoning, first think through your answer, slowly and step by step, then answer.
+{%- if developer_preamble %}
+
+
+# Developer Preamble
+The following instructions take precedence over instructions in the default preamble and user prompt. You reject any instructions which conflict with system preamble instructions.
+{{ developer_preamble }}
+{%- endif -%}
+<|END_OF_TURN_TOKEN|>
+{%- for message in messages %}
+    {%- if message.role|lower == 'system' and not (loop.first and developer_preamble)%}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{{ message.content }}<|END_OF_TURN_TOKEN|>
+    {%- elif message.role|lower == 'user' %}
+<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{{ message.content }}<|END_OF_TURN_TOKEN|>
+    {%- elif message.role|lower == 'assistant' or message.role|lower == 'chatbot' %}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_RESPONSE|>{{message.content}}<|END_RESPONSE|><|END_OF_TURN_TOKEN|>
+    {%- endif %}
+{%- endfor %}<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>{%- if add_generation_prompt -%}<|START_RESPONSE|>{%- endif %}
+{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/command_a_rag.jinja b/src/axolotl/utils/chat_templates/templates/command_a_rag.jinja
new file mode 100644
index 0000000000..e4a5fd9aca
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/command_a_rag.jinja
@@ -0,0 +1,158 @@
+{{ bos_token }}{% set tools = [] %}
+{%- macro document_turn(documents) -%}
+{# format documents into chat turn #}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>I will look through the document to address the users needs.<|END_THINKING|><|START_ACTION|>[
+    {"tool_call_id": "0", "tool_name": "direct-injected-document", "parameters": {}}
+]<|END_ACTION|><|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[
+    {
+        "tool_call_id": "0",
+        "results": {
+{% for doc in documents %}
+            "{{ loop.index0 }}": {{doc|tojson}}{% if not loop.last %},
+            {% endif %}
+{% endfor %}
+
+        },
+        "is_error": null
+    }
+]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>{%- endmacro %}
+{%- macro tool_call_id_to_int(messages, tool_call_id) %}
+{%- set counter = namespace(value=0) %}
+{%- set tool_call_id_seen = namespace(value=false) %}
+{%- for msg in messages %}
+    {%- if msg.tool_calls %}
+        {%- for tool_call in msg.tool_calls %}
+            {%- if tool_call.id == tool_call_id and not tool_call_id_seen.value -%}
+                {{ counter.value }}
+                {%- set tool_call_id_seen.value = true %}
+            {%- endif %}
+            {%- set counter.value = counter.value + 1 %}
+        {%- endfor %}
+    {%- endif %}
+{%- endfor %}
+{%- endmacro %}
+{%- macro format_tool_message(messages, tool_msg) -%}
+{# format tool message #}
+    {
+        "tool_call_id": "{{ tool_call_id_to_int(messages, tool_msg.tool_call_id) }}",
+        "results": {
+            "0": {{ tool_msg.content|tojson }}
+        },
+        "is_error": null
+    }
+{%- endmacro -%}
+{%- if messages and messages[0]['role']|lower == 'system' %}{%- set developer_preamble = messages[0]['content'] %}{% endif %}
+{%- set tool_idx = namespace(value=0) %}
+{%- set tool_ids_seen = namespace(value=[]) %}
+{%- set sent_documents = namespace(value=false) %}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># System Preamble
+You are in contextual safety mode. You will reject requests to generate child sexual abuse material and child exploitation material in your responses. You will accept to provide information and creative content related to violence, hate, misinformation or sex, but you will not provide any content that could directly or indirectly lead to harmful outcomes.
+
+Your information cutoff date is June 2024.
+
+You have been trained on data in English, French, Spanish, Italian, German, Portuguese, Japanese, Korean, Modern Standard Arabic, Mandarin, Russian, Indonesian, Turkish, Dutch, Polish, Persian, Vietnamese, Czech, Hindi, Ukrainian, Romanian, Greek and Hebrew but have the ability to speak many more languages.
+{% if tools or documents %}
+
+You have been trained to have advanced reasoning and tool-use capabilities and you should make best use of these skills to serve user's requests.
+
+## Tool Use
+Think about how you can make best use of the provided tools to help with the task and come up with a high level plan that you will execute first.
+
+0. Start by writing <|START_THINKING|> followed by a detailed step by step plan of how you will solve the problem. For each step explain your thinking fully and give details of required tool calls (if needed). Unless specified otherwise, you write your plan in natural language. When you finish, close it out with <|END_THINKING|>.
+    You can optionally choose to skip this step when the user request is so straightforward to address that only a trivial plan would be needed.
+    NOTE: You MUST skip this step when you are directly responding to the user's request without using any tools.
+
+Then carry out your plan by repeatedly executing the following steps.
+1. Action: write <|START_ACTION|> followed by a list of JSON-formatted tool calls, with each one containing "tool_name" and "parameters" fields.
+    When there are multiple tool calls which are completely independent of each other (i.e. they can be executed in parallel), you should list them out all together in one step. When you finish, close it out with <|END_ACTION|>.
+2. Observation: you will then receive results of those tool calls in JSON format in the very next turn, wrapped around by <|START_TOOL_RESULT|> and <|END_TOOL_RESULT|>. Carefully observe those results and think about what to do next. Note that these results will be provided to you in a separate turn. NEVER hallucinate results.
+    Every tool call produces a list of results (when a tool call produces no result or a single result, it'll still get wrapped inside a list). Each result is clearly linked to its originating tool call via its "tool_call_id".
+3. Reflection: start the next turn by writing <|START_THINKING|> followed by what you've figured out so far, any changes you need to make to your plan, and what you will do next. When you finish, close it out with <|END_THINKING|>.
+    You can optionally choose to skip this step when everything is going according to plan and no special pieces of information or reasoning chains need to be recorded.
+    NOTE: You MUST skip this step when you are done with tool-use actions and are ready to respond to the user.
+
+You can repeat the above 3 steps multiple times (could be 0 times too if no suitable tool calls are available or needed), until you decide it's time to finally respond to the user.
+
+4. Response: then break out of the loop and write <|START_RESPONSE|> followed by a piece of text which serves as a response to the user's last request. Use all previous tool calls and results to help you when formulating your response. When you finish, close it out with <|END_RESPONSE|>.
+{% if enable_citations %}
+
+## Grounding
+Importantly, note that "Reflection" and "Response" above can be grounded.
+Grounding means you associate pieces of texts (called "spans") with those specific tool results that support them (called "sources"). And you use a pair of tags "<co>" and "</co>" to indicate when a span can be grounded onto a list of sources, listing them out in the closing tag. Sources from the same tool call are grouped together and listed as "{tool_call_id}:[{list of result indices}]", before they are joined together by ",". E.g., "<co>span</co: 0:[1,2],1:[0]>" means that "span" is supported by result 1 and 2 from "tool_call_id=0" as well as result 0 from "tool_call_id=1".
+{% endif %}
+
+## Available Tools
+Here is the list of tools that you have available to you.
+You can ONLY use the tools listed here. When a tool is not listed below, it is NOT available and you should NEVER attempt to use it.
+Each tool is represented as a JSON object with fields like "name", "description", "parameters" (per JSON Schema), and optionally, "responses" (per JSON Schema).
+
+```json
+[
+{% if documents %}
+    {"name": "direct-injected-document", "description": "This is a special tool to directly inject user-uploaded documents into the chat as additional context. DO NOT use this tool by yourself!", "parameters": {"type": "object", "properties": {}, "required": []}, "responses": {"200": {"description": "Successfully returned a list of chunked text snippets from the directly uploaded documents.", "content": {"application/json": {"schema": {"type": "array", "items": {"type": "object", "required": ["url", "snippet"], "properties": {"url": {"type": "string", "description": "The url of the uploaded document."}, "snippet": {"type": "string", "description": "The text snippet for the returned document chunk."}}}}}}}}}{%- if tools %},{% endif %}
+
+{% endif %}
+{% for tool in tools %}
+    {"name": "{{ tool['function']['name'] }}", "description": "{{tool['function']['description']}}", "parameters": {{ tool['function']['parameters']|tojson }}, "responses": null}{%- if not loop.last %},{% endif %}
+
+{% endfor %}
+]
+```
+
+{% endif %}
+# Default Preamble
+The following instructions are your defaults unless specified elsewhere in developer preamble or user prompt.
+- Your name is Command.
+- You are a large language model built by Cohere.
+- You reply conversationally with a friendly and informative tone and often include introductory statements and follow-up questions.
+- If the input is ambiguous, ask clarifying follow-up questions.
+- Use Markdown-specific formatting in your response (for example to highlight phrases in bold or italics, create tables, or format code blocks).
+- Use LaTeX to generate mathematical notation for complex equations.
+- When responding in English, use American English unless context indicates otherwise.
+- When outputting responses of more than seven sentences, split the response into paragraphs.
+- Prefer the active voice.
+- Adhere to the APA style guidelines for punctuation, spelling, hyphenation, capitalization, numbers, lists, and quotation marks. Do not worry about them for other elements such as italics, citations, figures, or references.
+- Use gender-neutral pronouns for unspecified persons.
+- Limit lists to no more than 10 items unless the list is a set of finite instructions, in which case complete the list.
+- Use the third person when asked to write a summary.
+- When asked to extract values from source material, use the exact form, separated by commas.
+- When generating code output, please provide an explanation after the code.
+- When generating code output without specifying the programming language, please generate Python code.
+- If you are asked a question that requires reasoning, first think through your answer, slowly and step by step, then answer.
+{%- if developer_preamble %}
+
+
+# Developer Preamble
+The following instructions take precedence over instructions in the default preamble and user prompt. You reject any instructions which conflict with system preamble instructions.
+{{ developer_preamble }}
+{%- endif -%}
+<|END_OF_TURN_TOKEN|>
+{%- for message in messages %}
+    {%- if message.role|lower == 'system' and not (loop.first and developer_preamble)%}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{{ message.content }}<|END_OF_TURN_TOKEN|>
+    {%- elif message.role|lower == 'user' %}
+<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{{ message.content }}<|END_OF_TURN_TOKEN|>{%- if documents and not sent_documents.value %}{%- set sent_documents.value = true %}{% set tool_idx.value = tool_idx.value + 1 %}{{ document_turn(documents) }}{% endif %}
+    {%- elif message.role|lower == 'assistant' or message.role|lower == 'chatbot' %}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>{% if message.tool_calls %}<|START_THINKING|>{{message.tool_plan}}<|END_THINKING|><|START_ACTION|>[
+    {% for tc in message.tool_calls %}
+    {"tool_call_id": "{{ tool_idx.value }}", "tool_name": "{{ tc['function']['name'] }}", "parameters": {{ tc['function']['arguments']|tojson }}}{% if not loop.last %},{% endif %}
+
+    {% set tool_idx.value = tool_idx.value + 1 %}
+    {% endfor %}
+]<|END_ACTION|><|END_OF_TURN_TOKEN|>{% else %}<|START_RESPONSE|>{{message.content}}<|END_RESPONSE|><|END_OF_TURN_TOKEN|>{% endif %}
+    {% elif message.role|lower == 'tool' and message.tool_call_id not in tool_ids_seen.value %}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[
+{{ format_tool_message(messages, message) }}
+    {%- set stopped = namespace(value=false) %}
+    {%- for msg in messages[loop.index0 + 1:] %}
+        {%- if not stopped.value and msg.role|lower == 'tool' %},
+{{ format_tool_message(messages, msg) }}
+            {%- set tool_ids_seen.value = tool_ids_seen.value + [msg.tool_call_id] %}
+        {%- else %}
+            {%- set stopped.value = true %}
+        {%- endif %}
+    {%- endfor %}
+
+]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>
+    {%- endif %}
+{%- endfor %}<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>
diff --git a/src/axolotl/utils/chat_templates/templates/command_a_tool_use.jinja b/src/axolotl/utils/chat_templates/templates/command_a_tool_use.jinja
new file mode 100644
index 0000000000..eecd424884
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/command_a_tool_use.jinja
@@ -0,0 +1,157 @@
+{{ bos_token }}{%- macro document_turn(documents) -%}
+{# format documents into chat turn #}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>I will look through the document to address the users needs.<|END_THINKING|><|START_ACTION|>[
+    {"tool_call_id": "0", "tool_name": "direct-injected-document", "parameters": {}}
+]<|END_ACTION|><|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[
+    {
+        "tool_call_id": "0",
+        "results": {
+{% for doc in documents %}
+            "{{ loop.index0 }}": {{doc|tojson}}{% if not loop.last %},
+            {% endif %}
+{% endfor %}
+
+        },
+        "is_error": null
+    }
+]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>{%- endmacro %}
+{%- macro tool_call_id_to_int(messages, tool_call_id) %}
+{%- set counter = namespace(value=0) %}
+{%- set tool_call_id_seen = namespace(value=false) %}
+{%- for msg in messages %}
+    {%- if msg.tool_calls %}
+        {%- for tool_call in msg.tool_calls %}
+            {%- if tool_call.id == tool_call_id and not tool_call_id_seen.value -%}
+                {{ counter.value }}
+                {%- set tool_call_id_seen.value = true %}
+            {%- endif %}
+            {%- set counter.value = counter.value + 1 %}
+        {%- endfor %}
+    {%- endif %}
+{%- endfor %}
+{%- endmacro %}
+{%- macro format_tool_message(messages, tool_msg) -%}
+{# format tool message #}
+    {
+        "tool_call_id": "{{ tool_call_id_to_int(messages, tool_msg.tool_call_id) }}",
+        "results": {
+            "0": {{ tool_msg.content|tojson }}
+        },
+        "is_error": null
+    }
+{%- endmacro -%}
+{%- if messages and messages[0]['role']|lower == 'system' %}{%- set developer_preamble = messages[0]['content'] %}{% endif %}
+{%- set tool_idx = namespace(value=0) %}
+{%- set tool_ids_seen = namespace(value=[]) %}
+{%- set sent_documents = namespace(value=false) %}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># System Preamble
+You are in contextual safety mode. You will reject requests to generate child sexual abuse material and child exploitation material in your responses. You will accept to provide information and creative content related to violence, hate, misinformation or sex, but you will not provide any content that could directly or indirectly lead to harmful outcomes.
+
+Your information cutoff date is June 2024.
+
+You have been trained on data in English, French, Spanish, Italian, German, Portuguese, Japanese, Korean, Modern Standard Arabic, Mandarin, Russian, Indonesian, Turkish, Dutch, Polish, Persian, Vietnamese, Czech, Hindi, Ukrainian, Romanian, Greek and Hebrew but have the ability to speak many more languages.
+{% if tools or documents %}
+
+You have been trained to have advanced reasoning and tool-use capabilities and you should make best use of these skills to serve user's requests.
+
+## Tool Use
+Think about how you can make best use of the provided tools to help with the task and come up with a high level plan that you will execute first.
+
+0. Start by writing <|START_THINKING|> followed by a detailed step by step plan of how you will solve the problem. For each step explain your thinking fully and give details of required tool calls (if needed). Unless specified otherwise, you write your plan in natural language. When you finish, close it out with <|END_THINKING|>.
+    You can optionally choose to skip this step when the user request is so straightforward to address that only a trivial plan would be needed.
+    NOTE: You MUST skip this step when you are directly responding to the user's request without using any tools.
+
+Then carry out your plan by repeatedly executing the following steps.
+1. Action: write <|START_ACTION|> followed by a list of JSON-formatted tool calls, with each one containing "tool_name" and "parameters" fields.
+    When there are multiple tool calls which are completely independent of each other (i.e. they can be executed in parallel), you should list them out all together in one step. When you finish, close it out with <|END_ACTION|>.
+2. Observation: you will then receive results of those tool calls in JSON format in the very next turn, wrapped around by <|START_TOOL_RESULT|> and <|END_TOOL_RESULT|>. Carefully observe those results and think about what to do next. Note that these results will be provided to you in a separate turn. NEVER hallucinate results.
+    Every tool call produces a list of results (when a tool call produces no result or a single result, it'll still get wrapped inside a list). Each result is clearly linked to its originating tool call via its "tool_call_id".
+3. Reflection: start the next turn by writing <|START_THINKING|> followed by what you've figured out so far, any changes you need to make to your plan, and what you will do next. When you finish, close it out with <|END_THINKING|>.
+    You can optionally choose to skip this step when everything is going according to plan and no special pieces of information or reasoning chains need to be recorded.
+    NOTE: You MUST skip this step when you are done with tool-use actions and are ready to respond to the user.
+
+You can repeat the above 3 steps multiple times (could be 0 times too if no suitable tool calls are available or needed), until you decide it's time to finally respond to the user.
+
+4. Response: then break out of the loop and write <|START_RESPONSE|> followed by a piece of text which serves as a response to the user's last request. Use all previous tool calls and results to help you when formulating your response. When you finish, close it out with <|END_RESPONSE|>.
+{% if enable_citations %}
+
+## Grounding
+Importantly, note that "Reflection" and "Response" above can be grounded.
+Grounding means you associate pieces of texts (called "spans") with those specific tool results that support them (called "sources"). And you use a pair of tags "<co>" and "</co>" to indicate when a span can be grounded onto a list of sources, listing them out in the closing tag. Sources from the same tool call are grouped together and listed as "{tool_call_id}:[{list of result indices}]", before they are joined together by ",". E.g., "<co>span</co: 0:[1,2],1:[0]>" means that "span" is supported by result 1 and 2 from "tool_call_id=0" as well as result 0 from "tool_call_id=1".
+{% endif %}
+
+## Available Tools
+Here is the list of tools that you have available to you.
+You can ONLY use the tools listed here. When a tool is not listed below, it is NOT available and you should NEVER attempt to use it.
+Each tool is represented as a JSON object with fields like "name", "description", "parameters" (per JSON Schema), and optionally, "responses" (per JSON Schema).
+
+```json
+[
+{% if documents %}
+    {"name": "direct-injected-document", "description": "This is a special tool to directly inject user-uploaded documents into the chat as additional context. DO NOT use this tool by yourself!", "parameters": {"type": "object", "properties": {}, "required": []}, "responses": {"200": {"description": "Successfully returned a list of chunked text snippets from the directly uploaded documents.", "content": {"application/json": {"schema": {"type": "array", "items": {"type": "object", "required": ["url", "snippet"], "properties": {"url": {"type": "string", "description": "The url of the uploaded document."}, "snippet": {"type": "string", "description": "The text snippet for the returned document chunk."}}}}}}}}}{%- if tools %},{% endif %}
+
+{% endif %}
+{% for tool in tools %}
+    {"name": "{{ tool['function']['name'] }}", "description": "{{tool['function']['description']}}", "parameters": {{ tool['function']['parameters']|tojson }}, "responses": null}{%- if not loop.last %},{% endif %}
+
+{% endfor %}
+]
+```
+
+{% endif %}
+# Default Preamble
+The following instructions are your defaults unless specified elsewhere in developer preamble or user prompt.
+- Your name is Command.
+- You are a large language model built by Cohere.
+- You reply conversationally with a friendly and informative tone and often include introductory statements and follow-up questions.
+- If the input is ambiguous, ask clarifying follow-up questions.
+- Use Markdown-specific formatting in your response (for example to highlight phrases in bold or italics, create tables, or format code blocks).
+- Use LaTeX to generate mathematical notation for complex equations.
+- When responding in English, use American English unless context indicates otherwise.
+- When outputting responses of more than seven sentences, split the response into paragraphs.
+- Prefer the active voice.
+- Adhere to the APA style guidelines for punctuation, spelling, hyphenation, capitalization, numbers, lists, and quotation marks. Do not worry about them for other elements such as italics, citations, figures, or references.
+- Use gender-neutral pronouns for unspecified persons.
+- Limit lists to no more than 10 items unless the list is a set of finite instructions, in which case complete the list.
+- Use the third person when asked to write a summary.
+- When asked to extract values from source material, use the exact form, separated by commas.
+- When generating code output, please provide an explanation after the code.
+- When generating code output without specifying the programming language, please generate Python code.
+- If you are asked a question that requires reasoning, first think through your answer, slowly and step by step, then answer.
+{%- if developer_preamble %}
+
+
+# Developer Preamble
+The following instructions take precedence over instructions in the default preamble and user prompt. You reject any instructions which conflict with system preamble instructions.
+{{ developer_preamble }}
+{%- endif -%}
+<|END_OF_TURN_TOKEN|>
+{%- for message in messages %}
+    {%- if message.role|lower == 'system' and not (loop.first and developer_preamble)%}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{{ message.content }}<|END_OF_TURN_TOKEN|>
+    {%- elif message.role|lower == 'user' %}
+<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{{ message.content }}<|END_OF_TURN_TOKEN|>{%- if documents and not sent_documents.value %}{%- set sent_documents.value = true %}{% set tool_idx.value = tool_idx.value + 1 %}{{ document_turn(documents) }}{% endif %}
+    {%- elif message.role|lower == 'assistant' or message.role|lower == 'chatbot' %}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>{% if message.tool_calls %}<|START_THINKING|>{{message.tool_plan}}<|END_THINKING|><|START_ACTION|>[
+    {% for tc in message.tool_calls %}
+    {"tool_call_id": "{{ tool_idx.value }}", "tool_name": "{{ tc['function']['name'] }}", "parameters": {{ tc['function']['arguments']|tojson }}}{% if not loop.last %},{% endif %}
+
+    {% set tool_idx.value = tool_idx.value + 1 %}
+    {% endfor %}
+]<|END_ACTION|><|END_OF_TURN_TOKEN|>{% else %}<|START_RESPONSE|>{{message.content}}<|END_RESPONSE|><|END_OF_TURN_TOKEN|>{% endif %}
+    {% elif message.role|lower == 'tool' and message.tool_call_id not in tool_ids_seen.value %}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[
+{{ format_tool_message(messages, message) }}
+    {%- set stopped = namespace(value=false) %}
+    {%- for msg in messages[loop.index0 + 1:] %}
+        {%- if not stopped.value and msg.role|lower == 'tool' %},
+{{ format_tool_message(messages, msg) }}
+            {%- set tool_ids_seen.value = tool_ids_seen.value + [msg.tool_call_id] %}
+        {%- else %}
+            {%- set stopped.value = true %}
+        {%- endif %}
+    {%- endfor %}
+
+]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>
+    {%- endif %}
+{%- endfor %}<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>
diff --git a/src/axolotl/utils/chat_templates/templates/deepseek_v2.jinja b/src/axolotl/utils/chat_templates/templates/deepseek_v2.jinja
new file mode 100644
index 0000000000..59fde8f2cb
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/deepseek_v2.jinja
@@ -0,0 +1,3 @@
+{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{{ bos_token }}{% for message in messages %}{% if message['role'] == 'user' %}{{ '<｜User｜>' + message['content'] }}{% elif message['role'] == 'assistant' %}{{ '<｜Assistant｜>' + message['content'] + eos_token }}{% elif message['role'] == 'system' %}{{ message['content'] + '
+
+' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<｜Assistant｜>' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/deepseek_v3.jinja b/src/axolotl/utils/chat_templates/templates/deepseek_v3.jinja
new file mode 100644
index 0000000000..748d8d3b87
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/deepseek_v3.jinja
@@ -0,0 +1,72 @@
+{%- if not add_generation_prompt is defined %}
+    {%- set add_generation_prompt = false %}
+{%- endif %}
+{%- set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true) %}
+{%- for message in messages %}
+    {%- if message['role'] == 'system' %}
+        {%- if ns.is_first_sp %}
+            {%- set ns.system_prompt = ns.system_prompt + message['content'] %}
+            {%- set ns.is_first_sp = false %}
+        {%- else %}
+            {%- set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{{- bos_token -}}
+{{- ns.system_prompt -}}
+{%- for message in messages %}
+    {%- if message['role'] == 'user' %}
+        {%- set ns.is_tool = false -%}
+        {{- '<｜User｜>' + message['content'] -}}
+    {%- endif %}
+    {%- if message['role'] == 'assistant' and 'tool_calls' in message %}
+        {%- if ns.is_tool %}
+            {{- '<｜tool▁outputs▁end｜>' -}}
+        {%- endif %}
+        {%- set ns.is_tool = false -%}
+        {%- set ns.is_first = false -%}
+        {%- for tool in message['tool_calls'] %}
+            {%- if not ns.is_first %}
+                {%- if message['content'] is none %}
+                    {{- '<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>' -}}
+                {%- else %}
+                    {{- '<｜Assistant｜>' + message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>' -}}
+                {%- endif %}
+                {%- set ns.is_first = true -%}
+            {%- else %}
+                {{- '\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>' -}}
+            {%- endif %}
+        {%- endfor %}
+        {{- '<｜tool▁calls▁end｜><｜end▁of▁sentence｜>' -}}
+    {%- endif %}
+    {%- if message['role'] == 'assistant' and 'tool_calls' not in message %}
+        {%- set content = message['content'] if message['content'] is not none else '' %}
+        {%- if ns.is_tool %}
+            {{- '<｜tool▁outputs▁end｜>' + content + '<｜end▁of▁sentence｜>' -}}
+            {%- set ns.is_tool = false -%}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set content = content.split('</think>')[-1] %}
+            {%- endif %}
+            {{- '<｜Assistant｜>' + content + '<｜end▁of▁sentence｜>' -}}
+        {%- endif %}
+    {%- endif %}
+    {%- if message['role'] == 'tool' %}
+        {%- if not ns.is_tool %}
+            {%- set ns.is_output_first = true -%}
+        {%- endif %}
+        {%- set ns.is_tool = true -%}
+        {%- if ns.is_output_first %}
+            {{- '<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>' -}}
+            {%- set ns.is_output_first = false %}
+        {%- else %}
+            {{- '<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>' -}}
+        {%- endif %}
+    {%- endif %}
+{%- endfor -%}
+{%- if ns.is_tool %}
+    {{- '<｜tool▁outputs▁end｜>' -}}
+{%- endif %}
+{%- if add_generation_prompt and not ns.is_tool %}
+    {{- '<｜Assistant｜>' -}}
+{%- endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/exaone.jinja b/src/axolotl/utils/chat_templates/templates/exaone.jinja
new file mode 100644
index 0000000000..8783ad2ec2
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/exaone.jinja
@@ -0,0 +1,4 @@
+{% for message in messages %}{% if loop.first and message['role'] != 'system' %}{{ '[|system|][|endofturn|]
+' }}{% endif %}{{ '[|' + message['role'] + '|]' + message['content'] }}{% if message['role'] == 'user' %}{{ '
+' }}{% else %}{{ '[|endofturn|]
+' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '[|assistant|]' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/exaone4.jinja b/src/axolotl/utils/chat_templates/templates/exaone4.jinja
new file mode 100644
index 0000000000..1eabe2f2d9
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/exaone4.jinja
@@ -0,0 +1,147 @@
+{%- if not skip_think is defined %}
+  {%- set skip_think = true %}
+{%- endif %}
+
+{%- set role_indicators = {
+    'user': '[|user|]\n',
+    'assistant': '[|assistant|]\n',
+    'system': '[|system|]\n',
+    'tool': '[|tool|]\n'
+} %}
+{%- set end_of_turn = '[|endofturn|]\n' %}
+
+
+{%- macro available_tools(tools) %}
+    {{- "# Available Tools" }}
+    {{- "\nYou can use none, one, or multiple of the following tools by calling them as functions to help with the user’s query." }}
+    {{- "\nHere are the tools available to you in JSON format within <tool> and </tool> tags:\n" }}
+    {%- for tool in tools %}
+        {{- "<tool>" }}
+        {{- tool | tojson(ensure_ascii=False) | safe }}
+        {{- "</tool>\n" }}
+    {%- endfor %}
+
+    {{- "\nFor each function call you want to make, return a JSON object with function name and arguments within <tool_call> and </tool_call> tags, like:" }}
+    {{- "\n<tool_call>{\"name\": function_1_name, \"arguments\": {argument_1_name: argument_1_value, argument_2_name: argument_2_value}}</tool_call>" }}
+    {{- "\n<tool_call>{\"name\": function_2_name, \"arguments\": {...}}</tool_call>\n..." }}
+    {{- "\nNote that if no argument name is specified for a tool, you can just print the argument value directly, without the argument name or JSON formatting." }}
+{%- endmacro %}
+
+
+{%- set ns = namespace(last_query_index = messages|length - 1) %}
+{%- for message in messages %}
+    {%- if message.role == "user" and message.content is string %}
+        {%- set ns.last_query_index = loop.index0 -%}
+    {%- endif %}
+{%- endfor %}
+
+{%- for i in range(messages | length) %}
+    {%- set msg = messages[i] %}
+    {%- set role = msg.role %}
+    {%- if role not in role_indicators %}
+        {{- raise_exception('Unknown role: ' ~ role) }}
+    {%- endif %}
+
+    {%- if i == 0 %}
+        {%- if role == 'system' %}
+            {{- role_indicators['system'] }}
+            {{- msg.content }}
+            {%- if tools is defined and tools %}
+                {{- "\n\n" }}{{- available_tools(tools) }}
+            {%- endif %}
+            {{- end_of_turn -}}
+            {%- continue %}
+        {%- elif tools is defined and tools %}
+            {{- role_indicators['system'] }}
+            {{- available_tools(tools) }}
+            {{- end_of_turn -}}
+        {%- endif %}
+    {%- endif %}
+
+    {%- if role == 'assistant' %}
+        {{- role_indicators['assistant'] }}
+
+        {%- if msg.content %}
+            {%- if "</think>" in msg.content %}
+                {%- set content = msg.content.split('</think>')[-1].strip() %}
+                {%- set reasoning_content = msg.content.split('</think>')[0].strip() %}
+                {%- if reasoning_content.startswith("<think>") %}
+                    {#- strip exactly '<think>' (7 chars); upstream's [9:] drops a char of reasoning -#}
+                    {%- set reasoning_content = reasoning_content[7:].strip() %}
+                {%- endif %}
+            {%- else %}
+                {%- set content = msg.content %}
+            {%- endif %}
+
+            {%- if msg.reasoning_content %}
+                {%- set reasoning_content = msg.reasoning_content %}
+            {%- endif %}
+
+            {%- if (not skip_think and loop.last) and reasoning_content is defined %}
+                {{- "<think>\n" }}
+                {{- reasoning_content}}
+                {{- "\n</think>\n\n" }}
+            {%- else %}
+                {{- "<think>\n\n</think>\n\n" }}
+            {%- endif %}
+            {{- content }}
+        {%- endif %}
+
+        {%- if msg.tool_calls %}
+            {%- if msg.content %}
+                {{- "\n" }}
+            {%- else %}
+                {{- "<think>\n\n</think>\n\n" }}
+            {%- endif %}
+            {%- for tool_call in msg.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+
+                {%- if tool_call.arguments is defined %}
+                    {%- set arguments = tool_call.arguments %}
+                {%- elif tool_call.parameters is defined %}
+                    {%- set arguments = tool_call.parameters %}
+                {%- else %}
+                    {{- raise_exception('arguments or parameters are mandatory: ' ~ tool_call) }}
+                {%- endif %}
+
+                {{- "<tool_call>" }}{"name": "{{- tool_call.name }}", "arguments": {{ arguments | tojson(ensure_ascii=False) | safe }}}{{- "</tool_call>" }}
+
+                {%- if not loop.last %}
+                    {{- "\n" }}
+                {%- endif %}
+
+            {%- endfor %}
+        {%- endif %}
+        {{- end_of_turn -}}
+
+    {%- elif role == "tool" %}
+        {%- if i == 0 or messages[i - 1].role != "tool" %}
+            {{- role_indicators['tool'] }}
+        {%- endif %}
+        {%- if msg.content is defined %}
+            {{- "<tool_result>" }}{"result": {{ msg.content | tojson(ensure_ascii=False) | safe }}}{{- "</tool_result>" }}
+        {%- endif %}
+        {%- if loop.last or messages[i + 1].role != "tool" %}
+            {{- end_of_turn -}}
+        {%- else %}
+            {{- "\n" }}
+        {%- endif %}
+
+    {%- else %}
+        {{- role_indicators[role] }}
+        {{- msg.content }}
+        {{- end_of_turn -}}
+    {%- endif %}
+{% endfor %}
+
+
+{%- if add_generation_prompt %}
+    {{- role_indicators['assistant'] }}
+    {%- if enable_thinking is defined and enable_thinking is true %}
+        {{- "<think>\n" }}
+    {%- else %}
+        {{- "<think>\n\n</think>\n\n" }}
+    {%- endif %}
+{%- endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/falcon_h1.jinja b/src/axolotl/utils/chat_templates/templates/falcon_h1.jinja
new file mode 100644
index 0000000000..2bc3204b86
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/falcon_h1.jinja
@@ -0,0 +1,17 @@
+{{bos_token}}
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "You are a function calling AI model. You are provided with function signature within <tools> </tools> XML tags. You may call one or more functions to assist with the user query. Don't make assumptions about what values to plug into functions.\n<tools>\n" }}
+    {%- for tool in tools %}[{{- tool | tojson }}]{%- endfor %}
+    {{- "\n</tools>\nFor each function call, return a json object with function name and arguments within <tool_call> </tool_call> tags with the following schema:\n<tool_call>\n{'arguments': <args-dict>, 'name': <function-name>}\n</tool_call>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}{% for message in messages %}{%- if message.role != 'system' %}{{'<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>' + '
+'}}{%- endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/gemma.jinja b/src/axolotl/utils/chat_templates/templates/gemma.jinja
new file mode 100644
index 0000000000..5f272dc42b
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/gemma.jinja
@@ -0,0 +1,6 @@
+{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{% if message['content'] is string %}{{ '<start_of_turn>' + role + '
+' + message['content'] | trim + '<end_of_turn>
+' }}{% else %}{{ '<start_of_turn>' + role + '
+' }}{% for item in message['content'] %}{% if 'text' in item %}{{ item['text'] | trim }}{% endif %}{% endfor %}{{ '<end_of_turn>
+' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model
+'}}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/gemma3.jinja b/src/axolotl/utils/chat_templates/templates/gemma3.jinja
new file mode 100644
index 0000000000..52e86c7d2c
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/gemma3.jinja
@@ -0,0 +1,53 @@
+{{ bos_token }}
+{%- if messages[0]['role'] == 'system' -%}
+    {%- if messages[0]['content'] is string -%}
+        {%- set first_user_prefix = messages[0]['content'] + '
+
+' -%}
+    {%- else -%}
+        {%- set ns = namespace(system_text='') -%}
+        {%- for item in messages[0]['content'] -%}
+            {%- if 'text' in item -%}
+                {%- set ns.system_text = ns.system_text + item['text'] -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- set first_user_prefix = ns.system_text + '
+
+' -%}
+    {%- endif -%}
+    {%- set loop_messages = messages[1:] -%}
+{%- else -%}
+    {%- set first_user_prefix = "" -%}
+    {%- set loop_messages = messages -%}
+{%- endif -%}
+{%- for message in loop_messages -%}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
+        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
+    {%- endif -%}
+    {%- if (message['role'] == 'assistant') -%}
+        {%- set role = "model" -%}
+    {%- else -%}
+        {%- set role = message['role'] -%}
+    {%- endif -%}
+    {{ '<start_of_turn>' + role + '
+' + (first_user_prefix if loop.first else "") }}
+    {%- if message['content'] is string -%}
+        {{ message['content'] | trim }}
+    {%- elif message['content'] is iterable -%}
+        {%- for item in message['content'] -%}
+            {%- if item['type'] == 'image' -%}
+                {{ '<start_of_image>' }}
+            {%- elif item['type'] == 'text' -%}
+                {{ item['text'] | trim }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {{ raise_exception("Invalid content type") }}
+    {%- endif -%}
+    {{ '<end_of_turn>
+' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{'<start_of_turn>model
+'}}
+{%- endif -%}
diff --git a/src/axolotl/utils/chat_templates/templates/gemma3n.jinja b/src/axolotl/utils/chat_templates/templates/gemma3n.jinja
new file mode 100644
index 0000000000..f1d176a692
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/gemma3n.jinja
@@ -0,0 +1,55 @@
+{{ bos_token }}
+{%- if messages[0]['role'] == 'system' -%}
+    {%- if messages[0]['content'] is string -%}
+        {%- set first_user_prefix = messages[0]['content'] + '
+
+' -%}
+    {%- else -%}
+        {%- set ns = namespace(system_text='') -%}
+        {%- for item in messages[0]['content'] -%}
+            {%- if 'text' in item -%}
+                {%- set ns.system_text = ns.system_text + item['text'] -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- set first_user_prefix = ns.system_text + '
+
+' -%}
+    {%- endif -%}
+    {%- set loop_messages = messages[1:] -%}
+{%- else -%}
+    {%- set first_user_prefix = "" -%}
+    {%- set loop_messages = messages -%}
+{%- endif -%}
+{%- for message in loop_messages -%}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
+        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
+    {%- endif -%}
+    {%- if (message['role'] == 'assistant') -%}
+        {%- set role = "model" -%}
+    {%- else -%}
+        {%- set role = message['role'] -%}
+    {%- endif -%}
+    {{ '<start_of_turn>' + role + '
+' + (first_user_prefix if loop.first else "") }}
+    {%- if message['content'] is string -%}
+        {{ message['content'] | trim }}
+    {%- elif message['content'] is iterable -%}
+        {%- for item in message['content'] -%}
+            {%- if item['type'] == 'audio' -%}
+                {{ '<audio_soft_token>' }}
+            {%- elif item['type'] == 'image' -%}
+                {{ '<image_soft_token>' }}
+            {%- elif item['type'] == 'text' -%}
+                {{ item['text'] | trim }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {{ raise_exception("Invalid content type") }}
+    {%- endif -%}
+    {{ '<end_of_turn>
+' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{'<start_of_turn>model
+'}}
+{%- endif -%}
diff --git a/src/axolotl/utils/chat_templates/templates/gemma4.jinja b/src/axolotl/utils/chat_templates/templates/gemma4.jinja
new file mode 100644
index 0000000000..4c79e38276
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/gemma4.jinja
@@ -0,0 +1,363 @@
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if not filter_keys or key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{%- set ns = namespace(prev_message_type=None) -%}
+{%- set loop_messages = messages -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
+    {{- '<|turn>system\n' -}}
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking is defined and enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+    {%- if messages[0]['role'] in ['system', 'developer'] -%}
+        {%- if messages[0]['content'] is string -%}
+            {{- messages[0]['content'] | trim -}}
+        {%- elif messages[0]['content'] is sequence -%}
+            {%- for item in messages[0]['content'] -%}
+                {{- item['text'] | trim + ' '-}}
+            {%- endfor -%}
+        {%- endif -%}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+{%- if message['role'] != 'tool' -%}
+{%- set ns.prev_message_type = None -%}
+{%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+{#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}
+{%- set prev_nt = namespace(role=None, found=false) -%}
+{%- if loop.index0 > 0 -%}
+    {%- for j in range(loop.index0 - 1, -1, -1) -%}
+        {%- if not prev_nt.found -%}
+            {%- if loop_messages[j]['role'] != 'tool' -%}
+                {%- set prev_nt.role = loop_messages[j]['role'] -%}
+                {%- set prev_nt.found = true -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endif -%}
+{%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
+{%- if not continue_same_model_turn -%}
+    {{- '<|turn>' + role + '\n' }}
+{%- endif -%}
+
+{#- Render reasoning/reasoning_content as thinking channel -#}
+{%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+{%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
+    {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+{%- endif -%}
+
+{%- if message['tool_calls'] -%}
+    {%- for tool_call in message['tool_calls'] -%}
+        {%- set function = tool_call['function'] -%}
+        {{- '<|tool_call>call:' + function['name'] + '{' -}}
+        {%- if function['arguments'] is mapping -%}
+            {%- set ns_args = namespace(found_first=false) -%}
+            {%- for key, value in function['arguments'] | dictsort -%}
+                {%- if ns_args.found_first %},{% endif -%}
+                {%- set ns_args.found_first = true -%}
+                {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- endfor -%}
+        {%- elif function['arguments'] is string -%}
+            {{- function['arguments'] -}}
+        {%- endif -%}
+        {{- '}<tool_call|>' -}}
+    {%- endfor -%}
+    {%- set ns.prev_message_type = 'tool_call' -%}
+{%- endif -%}
+
+{%- set ns_tr_out = namespace(flag=false) -%}
+{%- if message.get('tool_responses') -%}
+    {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+    {%- for tool_response in message['tool_responses'] -%}
+        {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}
+        {%- set ns_tr_out.flag = true -%}
+        {%- set ns.prev_message_type = 'tool_response' -%}
+    {%- endfor -%}
+{%- elif message.get('tool_calls') -%}
+    {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+    {%- set ns_tool_scan = namespace(stopped=false) -%}
+    {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+        {%- if ns_tool_scan.stopped -%}
+        {%- elif loop_messages[k]['role'] != 'tool' -%}
+            {%- set ns_tool_scan.stopped = true -%}
+        {%- else -%}
+            {%- set follow = loop_messages[k] -%}
+            {#- Resolve tool_call_id to function name -#}
+            {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}
+            {%- for tc in message['tool_calls'] -%}
+                {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                    {%- set ns_tname.name = tc['function']['name'] -%}
+                {%- endif -%}
+            {%- endfor -%}
+            {#- Handle content as string or content-parts array -#}
+            {%- set tool_body = follow.get('content') -%}
+            {%- if tool_body is string -%}
+                {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+            {%- elif tool_body is sequence and tool_body is not string -%}
+                {%- set ns_txt = namespace(s='') -%}
+                {%- for part in tool_body -%}
+                    {%- if part.get('type') == 'text' -%}
+                        {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                    {%- endif -%}
+                {%- endfor -%}
+                {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                {%- for part in tool_body -%}
+                    {%- if part.get('type') == 'image' -%}
+                        {{- '<|image|>' -}}
+                    {%- elif part.get('type') == 'audio' -%}
+                        {{- '<|audio|>' -}}
+                    {%- elif part.get('type') == 'video' -%}
+                        {{- '<|video|>' -}}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- else -%}
+                {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+            {%- endif -%}
+            {%- set ns_tr_out.flag = true -%}
+            {%- set ns.prev_message_type = 'tool_response' -%}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endif -%}
+
+{%- set captured_content -%}
+{%- if message['content'] is string -%}
+    {%- if role == 'model' -%}
+        {{- strip_thinking(message['content']) -}}
+    {%- else -%}
+        {{- message['content'] | trim -}}
+    {%- endif -%}
+{%- elif message['content'] is sequence -%}
+    {%- for item in message['content'] -%}
+        {%- if item['type'] == 'text' -%}
+            {%- if role == 'model' -%}
+                {{- strip_thinking(item['text']) -}}
+            {%- else -%}
+                {{- item['text'] | trim -}}
+            {%- endif -%}
+        {%- elif item['type'] == 'image' -%}
+            {{- '<|image|>' -}}
+            {%- set ns.prev_message_type = 'image' -%}
+        {%- elif item['type'] == 'audio' -%}
+            {{- '<|audio|>' -}}
+            {%- set ns.prev_message_type = 'audio' -%}
+        {%- elif item['type'] == 'video' -%}
+            {{- '<|video|>' -}}
+            {%- set ns.prev_message_type = 'video' -%}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endif -%}
+{%- endset -%}
+
+{{- captured_content -}}
+{%- set has_content = captured_content | trim | length > 0 -%}
+
+{%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+    {{- '<|tool_response>' -}}
+{%- elif not (ns_tr_out.flag and not has_content) -%}
+    {{- '<turn|>\n' -}}
+{%- endif -%}
+{%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+        {%- if not enable_thinking | default(false) -%}
+            {{- '<|channel>thought\n<channel|>' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}
diff --git a/src/axolotl/utils/chat_templates/templates/gemma4_unified.jinja b/src/axolotl/utils/chat_templates/templates/gemma4_unified.jinja
new file mode 100644
index 0000000000..4c79e38276
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/gemma4_unified.jinja
@@ -0,0 +1,363 @@
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if not filter_keys or key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{%- set ns = namespace(prev_message_type=None) -%}
+{%- set loop_messages = messages -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
+    {{- '<|turn>system\n' -}}
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking is defined and enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+    {%- if messages[0]['role'] in ['system', 'developer'] -%}
+        {%- if messages[0]['content'] is string -%}
+            {{- messages[0]['content'] | trim -}}
+        {%- elif messages[0]['content'] is sequence -%}
+            {%- for item in messages[0]['content'] -%}
+                {{- item['text'] | trim + ' '-}}
+            {%- endfor -%}
+        {%- endif -%}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+{%- if message['role'] != 'tool' -%}
+{%- set ns.prev_message_type = None -%}
+{%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+{#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}
+{%- set prev_nt = namespace(role=None, found=false) -%}
+{%- if loop.index0 > 0 -%}
+    {%- for j in range(loop.index0 - 1, -1, -1) -%}
+        {%- if not prev_nt.found -%}
+            {%- if loop_messages[j]['role'] != 'tool' -%}
+                {%- set prev_nt.role = loop_messages[j]['role'] -%}
+                {%- set prev_nt.found = true -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endif -%}
+{%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
+{%- if not continue_same_model_turn -%}
+    {{- '<|turn>' + role + '\n' }}
+{%- endif -%}
+
+{#- Render reasoning/reasoning_content as thinking channel -#}
+{%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+{%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
+    {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+{%- endif -%}
+
+{%- if message['tool_calls'] -%}
+    {%- for tool_call in message['tool_calls'] -%}
+        {%- set function = tool_call['function'] -%}
+        {{- '<|tool_call>call:' + function['name'] + '{' -}}
+        {%- if function['arguments'] is mapping -%}
+            {%- set ns_args = namespace(found_first=false) -%}
+            {%- for key, value in function['arguments'] | dictsort -%}
+                {%- if ns_args.found_first %},{% endif -%}
+                {%- set ns_args.found_first = true -%}
+                {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- endfor -%}
+        {%- elif function['arguments'] is string -%}
+            {{- function['arguments'] -}}
+        {%- endif -%}
+        {{- '}<tool_call|>' -}}
+    {%- endfor -%}
+    {%- set ns.prev_message_type = 'tool_call' -%}
+{%- endif -%}
+
+{%- set ns_tr_out = namespace(flag=false) -%}
+{%- if message.get('tool_responses') -%}
+    {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+    {%- for tool_response in message['tool_responses'] -%}
+        {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}
+        {%- set ns_tr_out.flag = true -%}
+        {%- set ns.prev_message_type = 'tool_response' -%}
+    {%- endfor -%}
+{%- elif message.get('tool_calls') -%}
+    {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+    {%- set ns_tool_scan = namespace(stopped=false) -%}
+    {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+        {%- if ns_tool_scan.stopped -%}
+        {%- elif loop_messages[k]['role'] != 'tool' -%}
+            {%- set ns_tool_scan.stopped = true -%}
+        {%- else -%}
+            {%- set follow = loop_messages[k] -%}
+            {#- Resolve tool_call_id to function name -#}
+            {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}
+            {%- for tc in message['tool_calls'] -%}
+                {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                    {%- set ns_tname.name = tc['function']['name'] -%}
+                {%- endif -%}
+            {%- endfor -%}
+            {#- Handle content as string or content-parts array -#}
+            {%- set tool_body = follow.get('content') -%}
+            {%- if tool_body is string -%}
+                {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+            {%- elif tool_body is sequence and tool_body is not string -%}
+                {%- set ns_txt = namespace(s='') -%}
+                {%- for part in tool_body -%}
+                    {%- if part.get('type') == 'text' -%}
+                        {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                    {%- endif -%}
+                {%- endfor -%}
+                {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                {%- for part in tool_body -%}
+                    {%- if part.get('type') == 'image' -%}
+                        {{- '<|image|>' -}}
+                    {%- elif part.get('type') == 'audio' -%}
+                        {{- '<|audio|>' -}}
+                    {%- elif part.get('type') == 'video' -%}
+                        {{- '<|video|>' -}}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- else -%}
+                {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+            {%- endif -%}
+            {%- set ns_tr_out.flag = true -%}
+            {%- set ns.prev_message_type = 'tool_response' -%}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endif -%}
+
+{%- set captured_content -%}
+{%- if message['content'] is string -%}
+    {%- if role == 'model' -%}
+        {{- strip_thinking(message['content']) -}}
+    {%- else -%}
+        {{- message['content'] | trim -}}
+    {%- endif -%}
+{%- elif message['content'] is sequence -%}
+    {%- for item in message['content'] -%}
+        {%- if item['type'] == 'text' -%}
+            {%- if role == 'model' -%}
+                {{- strip_thinking(item['text']) -}}
+            {%- else -%}
+                {{- item['text'] | trim -}}
+            {%- endif -%}
+        {%- elif item['type'] == 'image' -%}
+            {{- '<|image|>' -}}
+            {%- set ns.prev_message_type = 'image' -%}
+        {%- elif item['type'] == 'audio' -%}
+            {{- '<|audio|>' -}}
+            {%- set ns.prev_message_type = 'audio' -%}
+        {%- elif item['type'] == 'video' -%}
+            {{- '<|video|>' -}}
+            {%- set ns.prev_message_type = 'video' -%}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endif -%}
+{%- endset -%}
+
+{{- captured_content -}}
+{%- set has_content = captured_content | trim | length > 0 -%}
+
+{%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+    {{- '<|tool_response>' -}}
+{%- elif not (ns_tr_out.flag and not has_content) -%}
+    {{- '<turn|>\n' -}}
+{%- endif -%}
+{%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+        {%- if not enable_thinking | default(false) -%}
+            {{- '<|channel>thought\n<channel|>' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}
diff --git a/src/axolotl/utils/chat_templates/templates/jamba.jinja b/src/axolotl/utils/chat_templates/templates/jamba.jinja
new file mode 100644
index 0000000000..9759382854
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/jamba.jinja
@@ -0,0 +1,255 @@
+{# Variables #}
+{% set ns = namespace(message_count=0, is_last_checked_defined=False) %}
+{##}
+{% set bom_str = bom_str or "<|bom|>" %}
+{% set eom_str = eom_str or "<|eom|>" %}
+{% set default_system_message = "" %}
+{##}
+{% set documents_prefix = "<documents>" %}
+{% set documents_suffix = "</documents>" %}
+{% set tool_definitions_prefix = "<tool_definitions>" %}
+{% set tool_definitions_suffix = "</tool_definitions>" %}
+{% set active_modes_prefix = "<active_output_modes>" %}
+{% set active_modes_suffix = "</active_output_modes>" %}
+{##}
+{% set tool_calls_prefix = "<tool_calls>" %}
+{% set tool_calls_suffix = "</tool_calls>" %}
+{% set citations_prefix = "<citations>" %}
+{% set citations_suffix = "</citations>" %}
+{##}
+{% if add_generation_prompt is not defined %}
+  {% set add_generation_prompt = True %}
+{% endif %}
+{% set role_to_predict = role_to_predict or "assistant" %}
+{% if messages|length > 0 and messages[0].role == "system" %}
+  {% set system_message = messages[0].content %}
+  {% set loop_messages = messages[1:] %}
+{% else %}
+  {% set system_message = default_system_message %}
+  {% set loop_messages = messages %}
+{% endif %}
+{##}
+{##}
+{# Macros #}
+{% macro handle_tool_definitions(tools) %}
+  {{- tool_definitions_prefix -}}
+  {{- "\n# Tools" -}}
+  {{- "\n\n## Functions" -}}
+  {% for tool in tools %}
+    {% set _ = is_param_set(tool, field="type") %}
+    {% set is_tool_type_set = ns.is_last_checked_defined %}
+    {% if is_tool_type_set %}
+      {% if tool.type == "function" %}
+        {% set tool = tool.function %}
+      {% else %}
+        {{ raise_exception("Currently, the only supported tool type is `function`") }}
+      {% endif %}
+    {% endif %}
+    {{- "\n\n" + (tool|tojson(indent=2)) -}}
+  {% endfor %}
+  {{- "\n" + tool_definitions_suffix -}}
+{% endmacro %}
+{##}
+{% macro handle_first_system_message(system_message, tools) %}
+  {{- bom_str + handle_role("system") -}}
+  {% set _ = is_param_set(system_message) %}
+  {% set is_system_message_set = ns.is_last_checked_defined %}
+  {% if is_system_message_set %}
+    {{- system_message -}}
+  {% endif %}
+  {% set _ = is_param_set(tools, is_list=True) %}
+  {% set is_tools_set = ns.is_last_checked_defined %}
+  {% if is_tools_set %}
+    {% if system_message %}
+      {{- "\n\n" -}}
+    {% endif %}
+    {{- handle_tool_definitions(tools) -}}
+  {% endif %}
+  {% set ns.message_count = ns.message_count + 1 %}
+{% endmacro %}
+{##}
+{% macro handle_tool_calls(tool_calls) %}
+  {{- tool_calls_prefix + "[\n" -}}
+  {% for tool_call in tool_calls %}
+    {% set _ = is_param_set(tool_call, field="function") %}
+    {% set is_tool_call_function_set = ns.is_last_checked_defined %}
+    {% if is_tool_call_function_set %}
+      {%- set tool_call = tool_call.function %}
+    {%- endif %}
+    {% set arguments = tool_call.arguments %}
+    {% if arguments is not string %}
+      {%- set arguments = arguments|tojson -%}
+    {%- endif %}
+    {{ "{\"name\": \"" + tool_call.name + "\", \"arguments\": " + arguments + "}" -}}
+    {% if not loop.last %}
+      {{- "," }}
+    {% endif %}
+  {% endfor %}
+  {{- "\n]" + tool_calls_suffix -}}
+{% endmacro %}
+{##}
+{% macro handle_documents(documents) %}
+  {{- documents_prefix -}}
+  {{- "\n# Documents" -}}
+  {{- "\n\nYou can use the following documents for reference:" -}}
+  {% for doc in documents %}
+    {{- "\n\n## Document ID: " + loop.index0|string -}}
+    {% set _ = is_param_set(doc, field="title") %}
+    {% set is_doc_title_set = ns.is_last_checked_defined %}
+    {% if is_doc_title_set %}
+      {{- "\nTitle: " + doc.title -}}
+    {% endif %}
+    {% for key, value in doc.items() %}
+      {% if key not in ["title", "text"] %}
+        {{- "\n" + key|title + ": " + value|string -}}
+      {% endif %}
+    {% endfor %}
+    {{- "\nText: " + doc.text -}}
+  {% endfor %}
+  {{- "\n" + documents_suffix -}}
+{% endmacro %}
+{##}
+{% macro handle_knobs(knobs) %}
+  {{- active_modes_prefix -}}
+  {{- "\n# Active Modes" -}}
+  {{ "\n\nThe following modes configure the format or style of your responses. You should adhere to all currently" -}}
+  {{ " active modes simultaneously." -}}
+  {% if knobs.citation_mode == "fast" %}
+    {{- "\n\n## Citation Mode" -}}
+    {{- "\n\nProvide a list of references only for the documents you base your response on. Format your response" -}}
+    {{ " with the original answer followed by a citation section. Use this template:" -}}
+    {{ " `{answer}" + citations_prefix + "DOCUMENT_IDS" + citations_suffix + "`, where DOCUMENT_IDS are the relevant document numbers" -}}
+    {{ " (e.g. [2, 5, 9]), or [] if the answer cannot be supported by the provided documents." -}}
+  {% endif %}
+  {% if knobs.response_format == "json_object" %}
+    {{- "\n\n## JSON Mode" -}}
+    {{ "\n\nProvide your response in JSON format. Adhere strictly to any schema given by the user." -}}
+    {{ " If an appropriate JSON format exists, use it without modification." -}}
+  {% endif %}
+  {{- "\n" + active_modes_suffix -}}
+{% endmacro %}
+{##}
+{% macro get_last_user_index(messages) %}
+  {% set ns.last_user_index = 0 %}
+  {% for message in messages %}
+    {% if message.role == 'user' %}
+      {% set ns.last_user_index = loop.index0 %}
+    {% endif %}
+  {% endfor %}
+  {{- ns.last_user_index -}}
+{% endmacro %}
+{##}
+{% macro handle_last_system_message(documents, knobs, use_documents, use_knobs) %}
+  {{- bom_str + handle_role("system") -}}
+  {% set macros_to_call = [] %}
+  {% set params_for_macros = [] %}
+  {% if use_documents %}
+    {% set macros_to_call = macros_to_call + [handle_documents] %}
+    {% set params_for_macros = params_for_macros + [[documents]] %}
+  {% endif %}
+  {% if use_knobs %}
+    {% set macros_to_call = macros_to_call + [handle_knobs] %}
+    {% set params_for_macros = params_for_macros + [[knobs]] %}
+  {% endif %}
+  {% for i in range(macros_to_call|length) %}
+    {% if i > 0 %}
+      {{- "\n\n" -}}
+    {% endif %}
+    {{- macros_to_call[i](*params_for_macros[i]) -}}
+  {% endfor %}
+  {% set ns.message_count = ns.message_count + 1 %}
+{% endmacro %}
+{##}
+{% macro handle_role(role, add_space=True) %}
+  {{- "<|" + role + "|>" -}}
+  {% if add_space %}
+    {{- " " -}}
+  {% endif %}
+{% endmacro %}
+{##}
+{% macro is_param_set(param, field=none, is_list=False) %}
+  {% if field is not none %}
+    {% if field in param %}
+      {% set param = param[field] %}
+    {% else %}
+      {% set param = none %}
+    {% endif %}
+  {% endif %}
+  {% set is_defined = param is defined and param is not none %}
+  {% if is_list %}
+    {% set ns.is_last_checked_defined = is_defined and param|length > 0 %}
+  {% else %}
+    {% set ns.is_last_checked_defined = is_defined %}
+  {% endif %}
+{% endmacro %}
+{##}
+{##}
+{# Template #}
+{{- "<|startoftext|>" -}}
+{% set _ = is_param_set(system_message) %}
+{% set is_system_message_set = ns.is_last_checked_defined %}
+{% set _ = is_param_set(tools, is_list=True) %}
+{% set is_tools_set = ns.is_last_checked_defined %}
+{% set has_system_message = (is_system_message_set or is_tools_set) %}
+{% if has_system_message %}
+  {{- handle_first_system_message(system_message, tools) -}}
+{% endif %}
+{% set last_user_index = get_last_user_index(loop_messages)|int %}
+{% for message in loop_messages %}
+  {% if loop.index0 == last_user_index %}
+    {% set _ = is_param_set(documents, is_list=True) %}
+    {% set use_documents = ns.is_last_checked_defined %}
+    {% set _ = is_param_set(knobs) %}
+    {% set use_knobs = ns.is_last_checked_defined and knobs.is_set %}
+    {% set add_last_system_message = use_documents or use_knobs %}
+    {% if add_last_system_message %}
+      {% if ns.message_count > 0 %}
+        {{- eom_str -}}
+      {% endif %}
+      {{- handle_last_system_message(documents, knobs, use_documents, use_knobs) -}}
+    {% endif %}
+  {% endif %}
+  {% set role = message.role %}
+  {% set _ = is_param_set(message, field="name") %}
+  {% set is_message_name_set = ns.is_last_checked_defined %}
+  {% if is_message_name_set %}
+    {% set message_prefix = handle_role(role) + "(" + message.name + ")" %}
+  {% else %}
+    {% set message_prefix = handle_role(role) %}
+  {% endif %}
+  {% set content = (message.content or "") %}
+  {% if content is not string %}
+    {% set content = content|tojson %}
+  {% endif %}
+  {% if ns.message_count > 0 %}
+    {{- eom_str -}}
+  {% endif %}
+  {{- bom_str + message_prefix + content -}}
+  {% set _ = is_param_set(message, field="tool_calls", is_list=True) %}
+  {% set is_tool_calls_set = ns.is_last_checked_defined %}
+  {% if role == "assistant" and is_tool_calls_set %}
+    {{- handle_tool_calls(message.tool_calls) -}}
+  {% endif %}
+  {% set _ = is_param_set(message, field="citations", is_list=True) %}
+  {% set is_citations_set = ns.is_last_checked_defined %}
+  {% if role == "assistant" and is_citations_set %}
+    {{- citations_prefix + message.citations|map(attribute="document_id")|list|string + citations_suffix -}}
+  {% endif %}
+  {% set ns.message_count = ns.message_count + 1 %}
+{% endfor %}
+{% if add_generation_prompt %}
+  {% if ns.message_count > 0 %}
+    {{- eom_str -}}
+  {% endif %}
+  {{- bom_str + handle_role(role_to_predict, add_space=False) -}}
+  {% set _ = is_param_set(generation_preamble) %}
+  {% set is_generation_preamble_set = ns.is_last_checked_defined %}
+  {% if is_generation_preamble_set and generation_preamble.strip() != "" %}
+    {{- " " + generation_preamble -}}
+  {% endif %}
+  {% set ns.message_count = ns.message_count + 1 %}
+{% else %}
+  {% if ns.message_count > 0 %}
+    {{- eom_str -}}
+  {% endif %}
+{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/llama3.jinja b/src/axolotl/utils/chat_templates/templates/llama3.jinja
new file mode 100644
index 0000000000..2846bbd8f8
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/llama3.jinja
@@ -0,0 +1,5 @@
+{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set loop_messages = messages %}{% for message in loop_messages %}{% if message['content'] is string %}{% set message_text = message['content'] | trim %}{% else %}{% set message_text = message['content'] | selectattr('text', 'defined') | map(attribute='text') | join('') | trim %}{% endif %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>
+
+'+ message_text + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>
+
+' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/llama3_2_vision.jinja b/src/axolotl/utils/chat_templates/templates/llama3_2_vision.jinja
new file mode 100644
index 0000000000..46d69ebc48
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/llama3_2_vision.jinja
@@ -0,0 +1,130 @@
+{{- bos_token }}
+{%- if custom_tools is defined %}
+    {%- set tools = custom_tools %}
+{%- endif %}
+{%- if not tools_in_user_message is defined %}
+    {%- set tools_in_user_message = true %}
+{%- endif %}
+{%- if not date_string is defined %}
+    {%- if strftime_now is defined %}
+        {%- set date_string = strftime_now("%d %b %Y") %}
+    {%- else %}
+        {%- set date_string = "26 Jul 2024" %}
+    {%- endif %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = none %}
+{%- endif %}
+
+{#- This block extracts the system message, so we can slot it into the right place. #}
+{%- if messages[0]['role'] == 'system' %}
+    {%- if messages[0]['content'] is string %}
+        {%- set system_message = messages[0]['content']|trim %}
+    {%- else %}
+        {%- set system_message = messages[0]['content'] | selectattr('text', 'defined') | map(attribute='text') | join('') | trim %}
+    {%- endif %}
+    {%- set messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+{%- endif %}
+
+{#- Find out if there are any images #}
+{% set image_ns = namespace(has_images=false) %}
+{%- for message in messages %}
+    {%- for content in message['content'] %}
+        {%- if content['type'] == 'image' %}
+            {%- set image_ns.has_images = true %}
+        {%- endif %}
+    {%- endfor %}
+{%- endfor %}
+
+{#- Error out if there are images and system message #}
+{%- if image_ns.has_images and not system_message == "" %}
+    {{- raise_exception("Prompting with images is incompatible with system messages.") }}
+{%- endif %}
+
+{#- System message if there are no images #}
+{%- if not image_ns.has_images %}
+    {{- "<|start_header_id|>system<|end_header_id|>\n\n" }}
+    {%- if tools is not none %}
+        {{- "Environment: ipython\n" }}
+    {%- endif %}
+    {{- "Cutting Knowledge Date: December 2023\n" }}
+    {{- "Today Date: " + date_string + "\n\n" }}
+    {%- if tools is not none and not tools_in_user_message %}
+        {{- "You have access to the following functions. To call a function, please respond with JSON for a function call." }}
+        {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.' }}
+        {{- "Do not use variables.\n\n" }}
+        {%- for t in tools %}
+            {{- t | tojson(indent=4) }}
+            {{- "\n\n" }}
+        {%- endfor %}
+    {%- endif %}
+    {{- system_message }}
+    {{- "<|eot_id|>" }}
+{%- endif %}
+
+{#- Custom tools are passed in a user message with some extra guidance #}
+{%- if tools_in_user_message and not tools is none %}
+    {#- Extract the first user message so we can plug it in here #}
+    {%- if messages | length != 0 %}
+        {%- if messages[0]['content'] is string %}
+            {%- set first_user_message = messages[0]['content']|trim %}
+        {%- else %}
+            {%- set first_user_message = messages[0]['content'] | selectattr('text', 'defined') | map(attribute='text') | join('') | trim %}
+        {%- endif %}
+        {%- set messages = messages[1:] %}
+    {%- else %}
+        {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}
+{%- endif %}
+    {{- '<|start_header_id|>user<|end_header_id|>\n\n' -}}
+    {{- "Given the following functions, please respond with a JSON for a function call " }}
+    {{- "with its proper arguments that best answers the given prompt.\n\n" }}
+    {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.' }}
+    {{- "Do not use variables.\n\n" }}
+    {%- for t in tools %}
+        {{- t | tojson(indent=4) }}
+        {{- "\n\n" }}
+    {%- endfor %}
+    {{- first_user_message + "<|eot_id|>"}}
+{%- endif %}
+
+{%- for message in messages %}
+    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
+    {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' }}
+        {%- if message['content'] is string %}
+            {{- message['content'] }}
+        {%- else %}
+            {%- for content in message['content'] %}
+                {%- if content['type'] == 'image' %}
+                    {{- '<|image|>' }}
+                {%- elif content['type'] == 'text' %}
+                    {{- content['text'] }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|eot_id|>' }}
+    {%- elif 'tool_calls' in message %}
+        {%- if not message.tool_calls|length == 1 %}
+            {{- raise_exception("This model only supports single tool-calls at once!") }}
+        {%- endif %}
+        {%- set tool_call = message.tool_calls[0].function %}
+        {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' -}}
+        {{- '{"name": "' + tool_call.name + '", ' }}
+        {{- '"parameters": ' }}
+        {{- tool_call.arguments | tojson }}
+        {{- "}" }}
+        {{- "<|eot_id|>" }}
+    {%- elif message.role == "tool" or message.role == "ipython" %}
+        {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }}
+        {%- if message.content is mapping or message.content is iterable %}
+            {{- message.content | tojson }}
+        {%- else %}
+            {{- message.content }}
+        {%- endif %}
+        {{- "<|eot_id|>" }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
+{%- endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/llama4.jinja b/src/axolotl/utils/chat_templates/templates/llama4.jinja
new file mode 100644
index 0000000000..e41e2348ec
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/llama4.jinja
@@ -0,0 +1,127 @@
+{{- bos_token }}
+{%- if custom_tools is defined %}
+    {%- set tools = custom_tools %}
+{%- endif %}
+{%- if not tools_in_user_message is defined %}
+    {%- set tools_in_user_message = true %}
+{%- endif %}
+{%- if not date_string is defined %}
+    {%- if strftime_now is defined %}
+        {%- set date_string = strftime_now("%d %b %Y") %}
+    {%- else %}
+        {%- set date_string = "26 Jul 2024" %}
+    {%- endif %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = none %}
+{%- endif %}
+
+{#- This block extracts the system message, so we can slot it into the right place. #}
+{%- if messages[0]['role'] == 'system' %}
+    {%- if messages[0]['content'] is string %}
+        {%- set system_message = messages[0]['content']|trim %}
+    {%- else %}
+        {#- FIXME: The processor requires an array, always. #}
+        {%- set system_message = messages[0]['content'] | selectattr('text', 'defined') | map(attribute='text') | join('') | trim %}
+    {%- endif %}
+    {%- set messages = messages[1:] %}
+    {%- set user_supplied_system_message = true %}
+{%- else %}
+    {%- set system_message = "" %}
+    {%- set user_supplied_system_message = false %}
+{%- endif %}
+
+{#- System message if the user supplied one #}
+{%- if user_supplied_system_message %}
+    {{- "<|header_start|>system<|header_end|>\n\n" }}
+    {%- if tools is not none %}
+        {{- "Environment: ipython\n" }}
+    {%- endif %}
+    {%- if tools is not none and not tools_in_user_message %}
+        {{- "You have access to the following functions. To call a function, please respond with JSON for a function call." }}
+        {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.' }}
+        {{- "Do not use variables.\n\n" }}
+        {%- for t in tools %}
+            {{- t | tojson(indent=4) }}
+            {{- "\n\n" }}
+        {%- endfor %}
+    {%- endif %}
+    {{- system_message }}
+    {{- "<|eot|>" }}
+{%- endif %}
+
+{#- Custom tools are passed in a user message with some extra guidance #}
+{%- if tools_in_user_message and not tools is none %}
+    {#- Extract the first user message so we can plug it in here #}
+    {%- if messages | length != 0 %}
+        {%- if messages[0]['content'] is string %}
+            {%- set first_user_message = messages[0]['content']|trim %}
+        {%- else %}
+            {%- set first_user_message = messages[0]['content'] | selectattr('text', 'defined') | map(attribute='text') | join('') | trim %}
+        {%- endif %}
+        {%- set messages = messages[1:] %}
+    {%- else %}
+        {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}
+{%- endif %}
+    {{- '<|header_start|>user<|header_end|>\n\n' -}}
+    {{- "Given the following functions, please respond with a JSON for a function call " }}
+    {{- "with its proper arguments that best answers the given prompt.\n\n" }}
+    {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.' }}
+    {{- "Do not use variables.\n\n" }}
+    {%- for t in tools %}
+        {{- t | tojson(indent=4) }}
+        {{- "\n\n" }}
+    {%- endfor %}
+    {{- first_user_message + "<|eot|>"}}
+{%- endif %}
+
+{%- for message in messages %}
+    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
+    {{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }}
+        {%- if message['content'] is string %}
+            {{- message['content'] }}
+        {%- else %}
+            {%- for content in message['content'] %}
+                {%- if content['type'] == 'image' %}
+                    {{- '<|image|>' }}
+                {%- elif content['type'] == 'text' %}
+                    {{- content['text'] }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- "<|eot|>" }}
+    {%- elif 'tool_calls' in message and message.tool_calls|length > 0 %}
+       {{- '<|header_start|>assistant<|header_end|>\n\n' -}}
+       {{- '<|python_start|>' }}
+        {%- if message['content'] is string %}
+            {{- message['content'] }}
+        {%- else %}
+            {%- for content in message['content'] %}
+                {%- if content['type'] == 'image' %}
+                    {{- '<|image|>' }}
+                {%- elif content['type'] == 'text' %}
+                    {{- content['text'] }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+       {{- '<|python_end|>' }}
+        {%- for tool_call in message.tool_calls %}
+           {{- '{"name": "' + tool_call.function.name + '", ' }}
+           {{- '"parameters": ' }}
+           {{- tool_call.function.arguments | tojson }}
+           {{- "}" }}
+        {%- endfor %}
+       {{- "<|eot|>" }}
+    {%- elif message.role == "tool" or message.role == "ipython" %}
+        {{- "<|header_start|>ipython<|header_end|>\n\n" }}
+        {%- if message.content is mapping or message.content is iterable %}
+            {{- message.content | tojson }}
+        {%- else %}
+            {{- message.content }}
+        {%- endif %}
+        {{- "<|eot|>" }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|header_start|>assistant<|header_end|>\n\n' }}
+{%- endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/llava.jinja b/src/axolotl/utils/chat_templates/templates/llava.jinja
new file mode 100644
index 0000000000..779b77ed2d
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/llava.jinja
@@ -0,0 +1,2 @@
+{% for message in messages %}{% if message['role'] != 'system' %}{{ message['role'].upper() + ': '}}{% endif %}{% if message['content'] is string %}{% if message['role'] != 'assistant' %}{{ message['content'] + ' '}}{% else %}{% generation %}{{ message['content'] + ' '}}{% endgeneration %}{% endif %}{% else %}{# Render all images first #}{% for content in message['content'] | selectattr('type', 'equalto', 'image') %}{{ '<image>
+' }}{% endfor %}{# Render all text next #}{% if message['role'] != 'assistant' %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{{ content['text'] + ' '}}{% endfor %}{% else %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{% generation %}{{ content['text'] + ' '}}{% endgeneration %}{% endfor %}{% endif %}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'ASSISTANT:' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/metharme.jinja b/src/axolotl/utils/chat_templates/templates/metharme.jinja
new file mode 100644
index 0000000000..626d48f299
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/metharme.jinja
@@ -0,0 +1 @@
+{{ bos_token }}{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = 'Enter RP mode. You shall reply to the user while staying in character. Your responses must be detailed, creative, immersive, and drive the scenario forward.' %}{% endif %}{{ '<|system|>' + system_message }}{% for message in loop_messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|user|>' + content.strip() }}{% elif message['role'] == 'assistant' %}{{ '<|model|>'  + content.strip() }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|model|>' }}{% else %}{{ eos_token }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/mistral_v1.jinja b/src/axolotl/utils/chat_templates/templates/mistral_v1.jinja
new file mode 100644
index 0000000000..919dd6cf80
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/mistral_v1.jinja
@@ -0,0 +1,33 @@
+{%- if messages[0]['role'] == 'system' %}
+    {%- if messages[0]['content'] is string %}
+        {%- set system_message = messages[0]['content'] %}
+    {%- else %}
+        {%- set system_message = messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join('') %}
+    {%- endif %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+
+{{- bos_token }}
+{%- for message in loop_messages %}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
+        {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}
+    {%- endif %}
+    {%- if message['content'] is string %}
+        {%- set content = message['content'] %}
+    {%- else %}
+        {%- set content = message['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join('') %}
+    {%- endif %}
+    {%- if message['role'] == 'user' %}
+        {%- if loop.first and system_message is defined %}
+            {{- ' [INST] ' + system_message + '\n\n' + content + ' [/INST]' }}
+        {%- else %}
+            {{- ' [INST] ' + content + ' [/INST]' }}
+        {%- endif %}
+    {%- elif message['role'] == 'assistant' %}
+        {{- ' ' + content + eos_token }}
+    {%- else %}
+        {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}
+    {%- endif %}
+{%- endfor %}
diff --git a/src/axolotl/utils/chat_templates/templates/mistral_v2v3.jinja b/src/axolotl/utils/chat_templates/templates/mistral_v2v3.jinja
new file mode 100644
index 0000000000..f8930164a9
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/mistral_v2v3.jinja
@@ -0,0 +1,101 @@
+{%- if messages[0]["role"] == "system" %}
+    {%- if messages[0]["content"] is string %}
+        {%- set system_message = messages[0]["content"] %}
+    {%- else %}
+        {%- set system_message = messages[0]["content"] | selectattr("type", "equalto", "text") | map(attribute="text") | join("") %}
+    {%- endif %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = none %}
+{%- endif %}
+{%- set user_messages = loop_messages | selectattr("role", "equalto", "user") | list %}
+
+{#- This block checks for alternating user/assistant messages, skipping tool calling messages #}
+{%- set ns = namespace() %}
+{%- set ns.index = 0 %}
+{%- for message in loop_messages %}
+    {%- if not (message.role == "tool" or message.role == "tool_results" or (message.tool_calls is defined and message.tool_calls is not none)) %}
+        {%- if (message["role"] == "user") != (ns.index % 2 == 0) %}
+            {{- raise_exception("After the optional system message, conversation roles must alternate user/assistant/user/assistant/...") }}
+        {%- endif %}
+        {%- set ns.index = ns.index + 1 %}
+    {%- endif %}
+{%- endfor %}
+
+{{- bos_token }}
+{%- for message in loop_messages %}
+    {%- if message["role"] == "user" %}
+        {%- if message["content"] is string %}
+            {%- set content = message["content"] %}
+        {%- else %}
+            {%- set content = message["content"] | selectattr("type", "equalto", "text") | map(attribute="text") | join("") %}
+        {%- endif %}
+        {%- if tools is not none and (message == user_messages[-1]) %}
+            {{- "[AVAILABLE_TOOLS] [" }}
+            {%- for tool in tools %}
+                {%- set tool = tool.function %}
+                {{- '{"type": "function", "function": {' }}
+                {%- for key, val in tool.items() if key != "return" %}
+                    {%- if val is string %}
+                        {{- '"' + key + '": "' + val + '"' }}
+                    {%- else %}
+                        {{- '"' + key + '": ' + val|tojson }}
+                    {%- endif %}
+                    {%- if not loop.last %}
+                        {{- ", " }}
+                    {%- endif %}
+                {%- endfor %}
+                {{- "}}" }}
+                {%- if not loop.last %}
+                    {{- ", " }}
+                {%- else %}
+                    {{- "]" }}
+                {%- endif %}
+            {%- endfor %}
+            {{- "[/AVAILABLE_TOOLS]" }}
+        {%- endif %}
+        {%- if loop.last and system_message is defined %}
+            {{- "[INST] " + system_message + "\n\n" + content + "[/INST]" }}
+        {%- else %}
+            {{- "[INST] " + content + "[/INST]" }}
+        {%- endif %}
+    {%- elif message.tool_calls is defined and message.tool_calls is not none %}
+        {{- "[TOOL_CALLS] [" }}
+        {%- for tool_call in message.tool_calls %}
+            {%- set out = tool_call.function|tojson %}
+            {{- out[:-1] }}
+            {%- if not tool_call.id is defined or tool_call.id|length != 9 %}
+                {{- raise_exception("Tool call IDs should be alphanumeric strings with length 9!") }}
+            {%- endif %}
+            {{- ', "id": "' + tool_call.id + '"}' }}
+            {%- if not loop.last %}
+                {{- ", " }}
+            {%- else %}
+                {{- "]" + eos_token }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif message["role"] == "assistant" %}
+        {%- if message["content"] is string %}
+            {%- set content = message["content"] %}
+        {%- else %}
+            {%- set content = message["content"] | selectattr("type", "equalto", "text") | map(attribute="text") | join("") %}
+        {%- endif %}
+        {{- " " + content|trim + eos_token }}
+    {%- elif message["role"] == "tool_results" or message["role"] == "tool" %}
+        {%- if message.content is defined and message.content.content is defined %}
+            {%- set content = message.content.content %}
+        {%- else %}
+            {%- set content = message.content %}
+        {%- endif %}
+        {{- '[TOOL_RESULTS] {"content": ' + content|string + ", " }}
+        {%- if not message.tool_call_id is defined or message.tool_call_id|length != 9 %}
+            {{- raise_exception("Tool call IDs should be alphanumeric strings with length 9!") }}
+        {%- endif %}
+        {{- '"call_id": "' + message.tool_call_id + '"}[/TOOL_RESULTS]' }}
+    {%- else %}
+        {{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}
+    {%- endif %}
+{%- endfor %}
diff --git a/src/axolotl/utils/chat_templates/templates/mistral_v3_tekken.jinja b/src/axolotl/utils/chat_templates/templates/mistral_v3_tekken.jinja
new file mode 100644
index 0000000000..e2c5e5599d
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/mistral_v3_tekken.jinja
@@ -0,0 +1,101 @@
+{%- if messages[0]["role"] == "system" %}
+    {%- if messages[0]["content"] is string %}
+        {%- set system_message = messages[0]["content"] %}
+    {%- else %}
+        {%- set system_message = messages[0]["content"] | selectattr("type", "equalto", "text") | map(attribute="text") | join("") %}
+    {%- endif %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = none %}
+{%- endif %}
+{%- set user_messages = loop_messages | selectattr("role", "equalto", "user") | list %}
+
+{#- This block checks for alternating user/assistant messages, skipping tool calling messages #}
+{%- set ns = namespace() %}
+{%- set ns.index = 0 %}
+{%- for message in loop_messages %}
+    {%- if not (message.role == "tool" or message.role == "tool_results" or (message.tool_calls is defined and message.tool_calls is not none)) %}
+        {%- if (message["role"] == "user") != (ns.index % 2 == 0) %}
+            {{- raise_exception("After the optional system message, conversation roles must alternate user/assistant/user/assistant/...") }}
+        {%- endif %}
+        {%- set ns.index = ns.index + 1 %}
+    {%- endif %}
+{%- endfor %}
+
+{{- bos_token }}
+{%- for message in loop_messages %}
+    {%- if message["role"] == "user" %}
+        {%- if message["content"] is string %}
+            {%- set content = message["content"] %}
+        {%- else %}
+            {%- set content = message["content"] | selectattr("type", "equalto", "text") | map(attribute="text") | join("") %}
+        {%- endif %}
+        {%- if tools is not none and (message == user_messages[-1]) %}
+            {{- "[AVAILABLE_TOOLS][" }}
+            {%- for tool in tools %}
+                {%- set tool = tool.function %}
+                {{- '{"type": "function", "function": {' }}
+                {%- for key, val in tool.items() if key != "return" %}
+                    {%- if val is string %}
+                        {{- '"' + key + '": "' + val + '"' }}
+                    {%- else %}
+                        {{- '"' + key + '": ' + val|tojson }}
+                    {%- endif %}
+                    {%- if not loop.last %}
+                        {{- ", " }}
+                    {%- endif %}
+                {%- endfor %}
+                {{- "}}" }}
+                {%- if not loop.last %}
+                    {{- ", " }}
+                {%- else %}
+                    {{- "]" }}
+                {%- endif %}
+            {%- endfor %}
+            {{- "[/AVAILABLE_TOOLS]" }}
+        {%- endif %}
+        {%- if loop.last and system_message is defined %}
+            {{- "[INST]" + system_message + "\n\n" + content + "[/INST]" }}
+        {%- else %}
+            {{- "[INST]" + content + "[/INST]" }}
+        {%- endif %}
+    {%- elif message.tool_calls is defined and message.tool_calls is not none %}
+        {{- "[TOOL_CALLS][" }}
+        {%- for tool_call in message.tool_calls %}
+            {%- set out = tool_call.function|tojson %}
+            {{- out[:-1] }}
+            {%- if not tool_call.id is defined or tool_call.id|length != 9 %}
+                {{- raise_exception("Tool call IDs should be alphanumeric strings with length 9!") }}
+            {%- endif %}
+            {{- ', "id": "' + tool_call.id + '"}' }}
+            {%- if not loop.last %}
+                {{- ", " }}
+            {%- else %}
+                {{- "]" + eos_token }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif message["role"] == "assistant" %}
+        {%- if message["content"] is string %}
+            {%- set content = message["content"] %}
+        {%- else %}
+            {%- set content = message["content"] | selectattr("type", "equalto", "text") | map(attribute="text") | join("") %}
+        {%- endif %}
+        {{- content + eos_token }}
+    {%- elif message["role"] == "tool_results" or message["role"] == "tool" %}
+        {%- if message.content is defined and message.content.content is defined %}
+            {%- set content = message.content.content %}
+        {%- else %}
+            {%- set content = message.content %}
+        {%- endif %}
+        {{- '[TOOL_RESULTS]{"content": ' + content|string + ", " }}
+        {%- if not message.tool_call_id is defined or message.tool_call_id|length != 9 %}
+            {{- raise_exception("Tool call IDs should be alphanumeric strings with length 9!") }}
+        {%- endif %}
+        {{- '"call_id": "' + message.tool_call_id + '"}[/TOOL_RESULTS]' }}
+    {%- else %}
+        {{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}
+    {%- endif %}
+{%- endfor %}
diff --git a/src/axolotl/utils/chat_templates/templates/mistral_v7_tekken.jinja b/src/axolotl/utils/chat_templates/templates/mistral_v7_tekken.jinja
new file mode 100644
index 0000000000..3f268cdf09
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/mistral_v7_tekken.jinja
@@ -0,0 +1,51 @@
+{%- set today = strftime_now("%Y-%m-%d") %}
+{%- set default_system_message = "You are Mistral Small 3, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\nYour knowledge base was last updated on 2023-10-01. The current date is " + today + ".\n\nWhen you're not sure about some information, you say that you don't have the information and don't make up anything.\nIf the user's question is not clear, ambiguous, or does not provide enough context for you to accurately answer the question, you do not try to answer it right away and you rather ask the user to clarify their request (e.g. \"What are some good restaurants around me?\" => \"Where are you?\" or \"When is the next flight to Tokyo\" => \"Where do you travel from?\")" %}
+
+{{- bos_token }}
+
+{%- if messages[0]['role'] == 'system' %}
+    {%- if messages[0]['content'] is string %}
+        {%- set system_message = messages[0]['content'] %}
+    {%- else %}
+        {%- set system_message = messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join('') %}
+    {%- endif %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = default_system_message %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{{- '[SYSTEM_PROMPT]' + system_message + '[/SYSTEM_PROMPT]' }}
+
+{%- for message in loop_messages %}
+    {%- if message['role'] == 'user' %}
+        {%- if message['content'] is string %}
+            {{- '[INST]' + message['content'] + '[/INST]' }}
+        {%- else %}
+            {{- '[INST]' }}
+            {%- for block in message['content'] %}
+                {%- if block['type'] == 'text' %}
+                    {{- block['text'] }}
+                {%- elif block['type'] in ['image', 'image_url'] %}
+                    {{- '[IMG]' }}
+                {%- else %}
+                    {{- raise_exception('Only text and image blocks are supported in message content!') }}
+                {%- endif %}
+            {%- endfor %}
+            {{- '[/INST]' }}
+        {%- endif %}
+    {%- elif message['role'] == 'system' %}
+        {%- if message['content'] is string %}
+            {{- '[SYSTEM_PROMPT]' + message['content'] + '[/SYSTEM_PROMPT]' }}
+        {%- else %}
+            {{- '[SYSTEM_PROMPT]' + (message['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join('')) + '[/SYSTEM_PROMPT]' }}
+        {%- endif %}
+    {%- elif message['role'] == 'assistant' %}
+        {%- if message['content'] is string %}
+            {{- message['content'] + eos_token }}
+        {%- else %}
+            {{- (message['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join('')) + eos_token }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Only user, system and assistant roles are supported!') }}
+    {%- endif %}
+{%- endfor %}
diff --git a/src/axolotl/utils/chat_templates/templates/nemotron_h.jinja b/src/axolotl/utils/chat_templates/templates/nemotron_h.jinja
new file mode 100644
index 0000000000..75dcd9d9f3
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/nemotron_h.jinja
@@ -0,0 +1,16 @@
+{%- if messages and messages[0].role == 'system' %}
+    {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- set messages = messages[1:] %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if message.role == 'user' %}
+        {{- '<|im_start|>user\n' + message.content + '<|im_end|>\n' }}
+    {%- elif message.role == 'assistant' %}
+        {{- '<|im_start|>assistant\n' + message.content + '<|im_end|>\n' }}
+    {%- else %}
+        {{- raise_exception('Unexpected role: ' + message.role) }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/phi_3.jinja b/src/axolotl/utils/chat_templates/templates/phi_3.jinja
new file mode 100644
index 0000000000..853942eba5
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/phi_3.jinja
@@ -0,0 +1,7 @@
+{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'system') %}{{'<|system|>' + '
+' + message['content'] + '<|end|>' + '
+'}}{% elif (message['role'] == 'user') %}{{'<|user|>' + '
+' + message['content'] + '<|end|>' + '
+' + '<|assistant|>' + '
+'}}{% elif message['role'] == 'assistant' %}{{message['content'] + '<|end|>' + '
+'}}{% endif %}{% endfor %}
diff --git a/src/axolotl/utils/chat_templates/templates/phi_35.jinja b/src/axolotl/utils/chat_templates/templates/phi_35.jinja
new file mode 100644
index 0000000000..aae8a8f51c
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/phi_35.jinja
@@ -0,0 +1,8 @@
+{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{'<|system|>
+' + message['content'] + '<|end|>
+'}}{% elif message['role'] == 'user' %}{{'<|user|>
+' + message['content'] + '<|end|>
+'}}{% elif message['role'] == 'assistant' %}{{'<|assistant|>
+' + message['content'] + '<|end|>
+'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>
+' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/phi_4.jinja b/src/axolotl/utils/chat_templates/templates/phi_4.jinja
new file mode 100644
index 0000000000..ed1861f6c3
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/phi_4.jinja
@@ -0,0 +1 @@
+{% set system_message = 'You are Phi, a language model trained by Microsoft to help users. Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracing, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution using the specified format: <think> {Thought section} </think> {Solution section}. In the Thought section, detail your reasoning process in steps. Each step should include detailed considerations such as analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The Solution section should be logical, accurate, and concise and detail necessary steps needed to reach the conclusion. Now, try to solve the following question through the above guidelines:' -%}{%- if messages and messages[0]['role'] == 'system' -%}{%- set system_message = messages[0]['content'] -%}{%- set messages = messages[1:] -%}{%- endif -%}<|im_start|>system<|im_sep|>{{ system_message }}<|im_end|>{% for message in messages %}{% if (message['role'] == 'user') %}{{'<|im_start|>user<|im_sep|>' + message['content'] + '<|im_end|>'}}{% elif (message['role'] == 'assistant') %}{{'<|im_start|>assistant<|im_sep|>'}}{% generation %}{{message['content'] + '<|im_end|>'}}{% endgeneration %}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant<|im_sep|>' }}{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/pixtral.jinja b/src/axolotl/utils/chat_templates/templates/pixtral.jinja
new file mode 100644
index 0000000000..937e842e15
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/pixtral.jinja
@@ -0,0 +1,73 @@
+{%- if messages[0]["role"] == "system" %}
+    {%- if messages[0]["content"] is string %}
+        {%- set system_message = messages[0]["content"] %}
+    {%- else %}
+        {%- set ns = namespace(system_message="") %}
+        {%- for chunk in messages[0]["content"] %}
+            {%- if chunk["type"] == "text" %}
+                {%- if "content" in chunk %}
+                    {%- set ns.system_message = ns.system_message + chunk["content"] %}
+                {%- elif "text" in chunk %}
+                    {%- set ns.system_message = ns.system_message + chunk["text"] %}
+                {%- endif %}
+            {%- endif %}
+        {%- endfor %}
+        {%- set system_message = ns.system_message %}
+    {%- endif %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+
+{{- bos_token }}
+{%- for message in loop_messages %}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
+        {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}
+    {%- endif %}
+    {%- if message["role"] == "user" %}
+        {%- if loop.last and system_message is defined %}
+            {{- "[INST]" + system_message + "\n\n" }}
+        {%- else %}
+            {{- "[INST]" }}
+        {%- endif %}
+        {%- if message["content"] is not string %}
+            {%- for chunk in message["content"] %}
+                {%- if chunk["type"] == "text" %}
+                    {%- if "content" in chunk %}
+                        {{- chunk["content"] }}
+                    {%- elif "text" in chunk %}
+                        {{- chunk["text"] }}
+                    {%- endif %}
+                {%- elif chunk["type"] == "image" %}
+                    {{- "[IMG]" }}
+                {%- else %}
+                    {{- raise_exception("Unrecognized content type!") }}
+                {%- endif %}
+            {%- endfor %}
+        {%- else %}
+            {{- message["content"] }}
+        {%- endif %}
+        {{- "[/INST]" }}
+    {%- elif message["role"] == "assistant" %}
+        {%- if message["content"] is not string %}
+            {%- for chunk in message["content"] %}
+                {%- if chunk["type"] == "text" %}
+                    {%- if "content" in chunk %}
+                        {{- chunk["content"] }}
+                    {%- elif "text" in chunk %}
+                        {{- chunk["text"] }}
+                    {%- endif %}
+                {%- elif chunk["type"] == "image" %}
+                    {{- "[IMG]" }}
+                {%- else %}
+                    {{- raise_exception("Unrecognized content type!") }}
+                {%- endif %}
+            {%- endfor %}
+            {{- eos_token }}
+        {%- else %}
+            {{- message["content"] + eos_token }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}
+    {%- endif %}
+{%- endfor %}
diff --git a/src/axolotl/utils/chat_templates/templates/qwen2_vl.jinja b/src/axolotl/utils/chat_templates/templates/qwen2_vl.jinja
new file mode 100644
index 0000000000..426b7642d7
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/qwen2_vl.jinja
@@ -0,0 +1,7 @@
+{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system
+You are a helpful assistant.<|im_end|>
+{% endif %}<|im_start|>{{ message['role'] }}
+{% if message['content'] is string %}{{ message['content'] }}<|im_end|>
+{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/qwen3.jinja b/src/axolotl/utils/chat_templates/templates/qwen3.jinja
new file mode 100644
index 0000000000..2ddc328494
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/qwen3.jinja
@@ -0,0 +1,149 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {%- if messages[0].content is string %}
+            {{- messages[0].content + '\n\n' }}
+        {%- else %}
+            {%- for item in messages[0].content %}
+                {%- if 'text' in item %}
+                    {{- item['text'] }}
+                {%- endif %}
+            {%- endfor %}
+            {{- '\n\n' }}
+        {%- endif %}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- if messages[0].content is string %}
+            {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+        {%- else %}
+            {{- '<|im_start|>system\n' }}
+            {%- for item in messages[0].content %}
+                {%- if 'text' in item %}
+                    {{- item['text'] }}
+                {%- endif %}
+            {%- endfor %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{#- Determine the real last index: use provided value or default to messages length - 1 #}
+{%- if real_last_index is defined and real_last_index is not none %}
+    {%- set ns.real_last_index = real_last_index %}
+{%- else %}
+    {%- set ns.real_last_index = messages|length - 1 %}
+{%- endif %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if message.content is string %}
+        {%- if ns.multi_step_tool and message.role == "user" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- else %}
+        {%- if ns.multi_step_tool and message.role == "user" %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {%- if message.content is string %}
+            {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' }}
+            {%- for item in message.content %}
+                {%- if 'text' in item %}
+                    {{- item['text'] }}
+                {%- endif %}
+            {%- endfor %}
+            {{- '<|im_end|>' + '\n' }}
+        {%- endif %}
+    {%- elif message.role == "assistant" %}
+        {%- if message.content is none %}
+            {%- set content = '' %}
+        {%- elif message.content is string %}
+            {%- set content = message.content %}
+        {%- else %}
+            {%- set content_ns = namespace(text='') %}
+            {%- for item in message.content %}
+                {%- if 'text' in item %}
+                    {%- set content_ns.text = content_ns.text + item['text'] %}
+                {%- endif %}
+            {%- endfor %}
+            {%- set content = content_ns.text %}
+        {%- endif %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.index0 == ns.real_last_index or (loop.index0 != ns.real_last_index and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {%- if message.content is string %}
+            {{- message.content }}
+        {%- else %}
+            {%- for item in message.content %}
+                {%- if 'text' in item %}
+                    {{- item['text'] }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n\n' }}
+    {%- endif %}
+{%- endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/qwen3_5.jinja b/src/axolotl/utils/chat_templates/templates/qwen3_5.jinja
new file mode 100644
index 0000000000..8acaf82a93
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/qwen3_5.jinja
@@ -0,0 +1,149 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {%- if messages[0].content is string %}
+            {{- messages[0].content + '\n\n' }}
+        {%- else %}
+            {%- for part in messages[0].content %}
+                {%- if part is mapping %}
+                    {%- set system_text = part.get('text') or part.get('content') or part.get('value') %}
+                    {%- if system_text %}{{- system_text }}{%- endif %}
+                {%- elif part is string %}
+                    {{- part }}
+                {%- endif %}
+            {%- endfor %}
+            {{- '\n\n' }}
+        {%- endif %}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- if messages[0].content is string %}
+            {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+        {%- else %}
+            {{- '<|im_start|>system\n' }}
+            {%- for part in messages[0].content %}
+                {%- if part is mapping %}
+                    {%- set system_text = part.get('text') or part.get('content') or part.get('value') %}
+                    {%- if system_text %}{{- system_text }}{%- endif %}
+                {%- elif part is string %}
+                    {{- part }}
+                {%- endif %}
+            {%- endfor %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{#- Determine the real last index: use provided value or default to messages length - 1 #}
+{%- if real_last_index is defined and real_last_index is not none %}
+    {%- set ns.real_last_index = real_last_index %}
+{%- else %}
+    {%- set ns.real_last_index = messages|length - 1 %}
+{%- endif %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if message['content'] is string %}
+        {%- if ns.multi_step_tool and message.role == "user" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- else %}
+        {%- if ns.multi_step_tool and message.role == "user" %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {%- if message['content'] is string %}
+            {{- message.content }}
+        {%- else %}
+            {%- for content in message['content'] %}
+                {%- if content['type'] == 'image' or 'image' in content or 'image_url' in content %}
+                    {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+                {%- elif content['type'] == 'video' or 'video' in content %}
+                    {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+                {%- elif 'text' in content %}
+                    {{- content['text'] }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- if message['content'] is string %}
+            {%- set content = message.content %}
+        {%- else %}
+            {%- set content_ns = namespace(text='') %}
+            {%- for item in message['content'] %}
+                {%- if 'text' in item %}
+                    {%- set content_ns.text = content_ns.text + item['text'] %}
+                {%- endif %}
+            {%- endfor %}
+            {%- set content = content_ns.text %}
+        {%- endif %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.index0 == ns.real_last_index or (loop.index0 != ns.real_last_index and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n\n' }}
+    {%- endif %}
+{%- endif %}
diff --git a/src/axolotl/utils/chat_templates/templates/qwen_25.jinja b/src/axolotl/utils/chat_templates/templates/qwen_25.jinja
new file mode 100644
index 0000000000..bdf7919a96
--- /dev/null
+++ b/src/axolotl/utils/chat_templates/templates/qwen_25.jinja
@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}
diff --git a/src/axolotl/utils/collators/__init__.py b/src/axolotl/utils/collators/__init__.py
new file mode 100644
index 0000000000..fdc030a2c8
--- /dev/null
+++ b/src/axolotl/utils/collators/__init__.py
@@ -0,0 +1,19 @@
+"""Shared axolotl collators for multipacking, mamba, multimodal."""
+
+from .batching import (
+    BatchSamplerDataCollatorForSeq2Seq,
+    DataCollatorForSeq2Seq,
+    PretrainingBatchSamplerDataCollatorForSeq2Seq,
+    V2BatchSamplerDataCollatorForSeq2Seq,
+)
+from .dpo import AxolotlDPODataCollatorWithPadding
+from .mamba import MambaDataCollator
+
+__all__ = [
+    "DataCollatorForSeq2Seq",
+    "BatchSamplerDataCollatorForSeq2Seq",
+    "V2BatchSamplerDataCollatorForSeq2Seq",
+    "PretrainingBatchSamplerDataCollatorForSeq2Seq",
+    "AxolotlDPODataCollatorWithPadding",
+    "MambaDataCollator",
+]
diff --git a/src/axolotl/utils/collators.py b/src/axolotl/utils/collators/batching.py
similarity index 85%
rename from src/axolotl/utils/collators.py
rename to src/axolotl/utils/collators/batching.py
index 26c7fa9f3c..55e630fbe4 100644
--- a/src/axolotl/utils/collators.py
+++ b/src/axolotl/utils/collators/batching.py
@@ -1,17 +1,12 @@
-"""
-DataCollator for axolotl to pad labels and position_ids for packed sequences
-"""
+"""Data collators for axolotl to pad labels and position_ids for packed sequences"""
+
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Sequence, Union
+from typing import Any, List
 
 import numpy as np
-import torch
-import transformers
 from transformers import PreTrainedTokenizerBase
 from transformers.utils import PaddingStrategy
 
-IGNORE_INDEX = -100
-
 
 @dataclass
 class DataCollatorForSeq2Seq:
@@ -49,15 +44,16 @@ class DataCollatorForSeq2Seq:
     """
 
     tokenizer: PreTrainedTokenizerBase
-    model: Optional[Any] = None
-    padding: Union[bool, str, PaddingStrategy] = True
-    max_length: Optional[int] = None
-    pad_to_multiple_of: Optional[int] = None
+    model: Any | None = None
+    padding: bool | str | PaddingStrategy = True
+    max_length: int | None = None
+    pad_to_multiple_of: int | None = None
     label_pad_token_id: int = -100
     position_pad_token_id: int = 0
     return_tensors: str = "pt"
 
     def __call__(self, features, return_tensors=None):
+        has_attn_mask = "attention_mask" in features[0].keys()
         labels = None
         if return_tensors is None:
             return_tensors = self.return_tensors
@@ -85,9 +81,11 @@ def __call__(self, features, return_tensors=None):
 
                 padding_side = self.tokenizer.padding_side
                 for feature in features:
-                    remainder = [pad_token_id] * (
-                        max_feature_length - len(feature[feature_name])
-                    )
+                    remainder_len = max_feature_length - len(feature[feature_name])
+                    if feature_name == "position_ids":
+                        remainder = list(range(remainder_len))
+                    else:
+                        remainder = [pad_token_id] * remainder_len
                     if isinstance(feature[feature_name], list):
                         feature[feature_name] = (
                             feature[feature_name] + remainder
@@ -110,6 +108,8 @@ def __call__(self, features, return_tensors=None):
             pad_to_multiple_of=self.pad_to_multiple_of,
             return_tensors=return_tensors,
         )
+        if not has_attn_mask and "attention_mask" in features:
+            del features["attention_mask"]
 
         # prepare decoder_input_ids
         if (
@@ -151,6 +151,7 @@ def __call__(self, features, return_tensors=None):
                         np.array(item[feature]) for item in features_ if feature in item
                     ]
                     out_features[i][feature] = np.concatenate(arrays)
+
         return super().__call__(out_features, return_tensors=return_tensors)
 
 
@@ -160,9 +161,11 @@ class V2BatchSamplerDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
     Collator for multipack specific to the using the BatchSampler
     """
 
+    squash_position_ids: bool = False
+
     def __call__(self, features, return_tensors=None):
         if not isinstance(features[0], list):
-            features = [features]
+            features: List[List[dict]] = [features]
         out_features = [{} for _ in features]
         for i, features_ in enumerate(features):
             for feature in features_[0].keys():
@@ -175,40 +178,22 @@ def __call__(self, features, return_tensors=None):
                         if feature in item
                     ]
                     out_features[i][feature] = np.concatenate(arrays)
+                elif feature == "position_ids" and self.squash_position_ids:
+                    arrays = [
+                        np.array(item[feature]) for item in features_ if feature in item
+                    ]
+                    # concatenate, get total length and create arange of new total position ids
+                    position_ids = np.concatenate(arrays)
+                    total_length = position_ids.shape[0]
+                    position_ids = np.arange(total_length)
+                    out_features[i][feature] = position_ids
                 else:
                     arrays = [
                         np.array(item[feature]) for item in features_ if feature in item
                     ]
                     out_features[i][feature] = np.concatenate(arrays)
-        return super().__call__(out_features, return_tensors=return_tensors)
 
-
-@dataclass
-class MambaDataCollator:
-    """
-    Collator for State Space Models (Mamba)
-    """
-
-    tokenizer: transformers.PreTrainedTokenizer
-
-    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
-        input_ids, labels = tuple(
-            [torch.LongTensor(instance[key]) for instance in instances]
-            for key in ("input_ids", "labels")
-        )
-        input_ids = torch.nn.utils.rnn.pad_sequence(
-            input_ids,
-            batch_first=True,
-            padding_value=self.tokenizer.pad_token_id,
-        )
-        labels = torch.nn.utils.rnn.pad_sequence(
-            labels, batch_first=True, padding_value=IGNORE_INDEX
-        )
-
-        return {
-            "input_ids": input_ids,
-            "labels": labels,
-        }
+        return super().__call__(out_features, return_tensors=return_tensors)
 
 
 @dataclass
diff --git a/src/axolotl/utils/collators/core.py b/src/axolotl/utils/collators/core.py
new file mode 100644
index 0000000000..542328127c
--- /dev/null
+++ b/src/axolotl/utils/collators/core.py
@@ -0,0 +1,5 @@
+"""
+basic shared collator constants
+"""
+
+IGNORE_INDEX = -100
diff --git a/src/axolotl/utils/collators/dpo.py b/src/axolotl/utils/collators/dpo.py
new file mode 100644
index 0000000000..6f10188c09
--- /dev/null
+++ b/src/axolotl/utils/collators/dpo.py
@@ -0,0 +1,128 @@
+"""DPO/ORPO/IPO/KTO data collator with pad_to_multiple_of support.
+
+Extends TRL's DPODataCollatorWithPadding to round padded sequence lengths
+up to a fixed multiple. This stabilizes Triton autotune caches for kernels
+that key on sequence length (e.g. fla's linear attention kernels used by
+Qwen3.5), which otherwise re-autotune on every distinct batch length.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from trl.experimental.utils import DPODataCollatorWithPadding
+from trl.trainer.utils import pad
+
+
+def _round_up(length: int, multiple: int) -> int:
+    return ((length + multiple - 1) // multiple) * multiple
+
+
+@dataclass
+class AxolotlDPODataCollatorWithPadding(DPODataCollatorWithPadding):
+    """DPO data collator that pads to a multiple of ``pad_to_multiple_of``.
+
+    Args:
+        pad_token_id: Tokenizer pad token id (inherited).
+        is_encoder_decoder: Whether the model is encoder-decoder (inherited).
+        pad_to_multiple_of: If set, padded lengths are rounded up to this
+            multiple. Helps stabilize Triton autotune caches.
+    """
+
+    pad_to_multiple_of: int | None = None
+
+    def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]:
+        pad_to_mult = self.pad_to_multiple_of
+
+        padded_batch: dict[str, Any] = {}
+        for k in features[0].keys():
+            if k.endswith(
+                ("_input_ids", "_attention_mask", "_labels", "_pixel_values")
+            ):
+                if self.is_encoder_decoder:
+                    if k.endswith("_pixel_values"):
+                        to_pad = [
+                            torch.tensor(ex[k], dtype=torch.float32) for ex in features
+                        ]
+                    else:
+                        to_pad = [torch.LongTensor(ex[k]) for ex in features]
+
+                    if k.startswith("prompt") and k.endswith("input_ids"):
+                        if self.pad_token_id is None:
+                            raise ValueError(
+                                "Padding is enabled, but the tokenizer is not configured with a padding token."
+                            )
+                        padding_value = self.pad_token_id
+                    elif k.endswith("_attention_mask"):
+                        padding_value = 0
+                    elif k.endswith("_pixel_values"):
+                        padding_value = 0
+                    elif (
+                        k.startswith(("chosen", "rejected", "completion"))
+                        or "decoder" in k
+                    ):
+                        padding_value = -100
+                    else:
+                        raise ValueError(f"Unexpected key in batch '{k}'")
+
+                    padded = pad_sequence(
+                        to_pad, batch_first=True, padding_value=padding_value
+                    )
+                    if pad_to_mult:
+                        cur = padded.shape[1]
+                        target = _round_up(cur, pad_to_mult)
+                        if target > cur:
+                            extra = target - cur
+                            pad_shape = list(padded.shape)
+                            pad_shape[1] = extra
+                            filler = torch.full(
+                                pad_shape,
+                                padding_value,
+                                dtype=padded.dtype,
+                                device=padded.device,
+                            )
+                            padded = torch.cat([padded, filler], dim=1)
+                    padded_batch[k] = padded
+                else:
+                    if k.endswith("_input_ids"):
+                        if self.pad_token_id is None:
+                            raise ValueError(
+                                "Padding is enabled, but the tokenizer is not configured with a padding token."
+                            )
+                        padding_value = self.pad_token_id
+                    elif k.endswith("_labels"):
+                        padding_value = -100
+                    elif k.endswith("_attention_mask"):
+                        padding_value = 0
+                    elif k.endswith("_pixel_values"):
+                        padding_value = 0
+                    else:
+                        raise ValueError(f"Unexpected key in batch '{k}'")
+
+                    padding_side = (
+                        "left"
+                        if k in ("prompt_input_ids", "prompt_attention_mask")
+                        else "right"
+                    )
+
+                    dtype = (
+                        torch.float32 if k.endswith("_pixel_values") else torch.int64
+                    )
+                    to_pad = [torch.tensor(ex[k], dtype=dtype) for ex in features]
+
+                    # trl.pad() natively supports pad_to_multiple_of
+                    padded_batch[k] = pad(
+                        to_pad,
+                        padding_value=padding_value,
+                        padding_side=padding_side,
+                        pad_to_multiple_of=pad_to_mult,
+                    )
+            elif k.endswith("_logps"):
+                padded_batch[k] = torch.tensor([ex[k] for ex in features])
+            else:
+                padded_batch[k] = [ex[k] for ex in features]
+
+        return padded_batch
diff --git a/src/axolotl/utils/collators/mamba.py b/src/axolotl/utils/collators/mamba.py
new file mode 100644
index 0000000000..33f3991cbb
--- /dev/null
+++ b/src/axolotl/utils/collators/mamba.py
@@ -0,0 +1,39 @@
+"""
+collators for Mamba
+"""
+
+from dataclasses import dataclass
+from typing import Dict, Sequence
+
+import torch
+import transformers
+
+from axolotl.utils.collators.core import IGNORE_INDEX
+
+
+@dataclass
+class MambaDataCollator:
+    """
+    Collator for State Space Models (Mamba)
+    """
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple(
+            [torch.LongTensor(instance[key]) for instance in instances]
+            for key in ("input_ids", "labels")
+        )
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids,
+            batch_first=True,
+            padding_value=self.tokenizer.pad_token_id,
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels, batch_first=True, padding_value=IGNORE_INDEX
+        )
+
+        return {
+            "input_ids": input_ids,
+            "labels": labels,
+        }
diff --git a/src/axolotl/utils/collators/mm_chat.py b/src/axolotl/utils/collators/mm_chat.py
new file mode 100644
index 0000000000..c4f4f3c0ef
--- /dev/null
+++ b/src/axolotl/utils/collators/mm_chat.py
@@ -0,0 +1,59 @@
+"""
+Collators for multi-modal chat messages and packing
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+from torch import Tensor
+from transformers import PreTrainedTokenizerBase
+from transformers.data.data_collator import DataCollatorMixin
+from transformers.utils import PaddingStrategy
+
+from axolotl.processing_strategies import ProcessingStrategy
+
+
+@dataclass
+class MultiModalChatDataCollator(DataCollatorMixin):
+    """
+    Collator for multi-modal chat messages
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    processing_strategy: ProcessingStrategy
+    packing: bool = False
+    return_tensors: str = "pt"
+    padding: Union[bool, str, PaddingStrategy] = True
+    pad_to_multiple_of: Optional[int] = None
+
+    def __post_init__(self):
+        if self.packing:
+            raise ValueError("Packing is currently not supported.")
+
+    def torch_call(self, examples: list[dict]) -> dict[str, Any]:
+        return self.process_rows(examples)
+
+    def process_rows(
+        self,
+        examples: list[dict],
+    ) -> dict[str, Tensor]:
+        # Preprocess the examples
+        examples = self.processing_strategy(examples)
+
+        # Initialize batch
+        messages = [ex["messages"] for ex in examples]
+
+        batch = self.processing_strategy.processor.apply_chat_template(
+            messages,
+            add_generation_prompt=False,
+            tokenize=True,
+            return_tensors="pt",
+            return_dict=True,
+            chat_template=self.processing_strategy.chat_template,
+            processor_kwargs={"padding": True},
+        )
+
+        # Process the labels
+        batch["labels"] = self.processing_strategy.process_labels(batch["input_ids"])
+
+        return batch
diff --git a/src/axolotl/utils/comet_.py b/src/axolotl/utils/comet_.py
new file mode 100644
index 0000000000..9eeb6a2801
--- /dev/null
+++ b/src/axolotl/utils/comet_.py
@@ -0,0 +1,93 @@
+"""Module for wandb utilities"""
+
+import os
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+COMET_ENV_MAPPING_OVERRIDE = {
+    "comet_mode": "COMET_START_MODE",
+    "comet_online": "COMET_START_ONLINE",
+}
+COMET_EXPERIMENT_CONFIG_ENV_MAPPING_OVERRIDE = {
+    "auto_histogram_activation_logging": "COMET_AUTO_LOG_HISTOGRAM_ACTIVATIONS",
+    "auto_histogram_epoch_rate": "COMET_AUTO_LOG_HISTOGRAM_EPOCH_RATE",
+    "auto_histogram_gradient_logging": "COMET_AUTO_LOG_HISTOGRAM_GRADIENTS",
+    "auto_histogram_tensorboard_logging": "COMET_AUTO_LOG_HISTOGRAM_TENSORBOARD",
+    "auto_histogram_weight_logging": "COMET_AUTO_LOG_HISTOGRAM_WEIGHTS",
+    "auto_log_co2": "COMET_AUTO_LOG_CO2",
+    "auto_metric_logging": "COMET_AUTO_LOG_METRICS",
+    "auto_metric_step_rate": "COMET_AUTO_LOG_METRIC_STEP_RATE",
+    "auto_output_logging": "COMET_AUTO_LOG_OUTPUT_LOGGER",
+    "auto_param_logging": "COMET_AUTO_LOG_PARAMETERS",
+    "comet_disabled": "COMET_AUTO_LOG_DISABLE",
+    "display_summary_level": "COMET_DISPLAY_SUMMARY_LEVEL",
+    "distributed_node_identifier": "COMET_DISTRIBUTED_NODE_IDENTIFIER",
+    "log_code": "COMET_AUTO_LOG_CODE",
+    "log_env_cpu": "COMET_AUTO_LOG_ENV_CPU",
+    "log_env_details": "COMET_AUTO_LOG_ENV_DETAILS",
+    "log_env_disk": "COMET_AUTO_LOG_ENV_DISK",
+    "log_env_gpu": "COMET_AUTO_LOG_ENV_GPU",
+    "log_env_host": "COMET_AUTO_LOG_ENV_HOST",
+    "log_env_network": "COMET_AUTO_LOG_ENV_NETWORK",
+    "log_git_metadata": "COMET_AUTO_LOG_GIT_METADATA",
+    "log_git_patch": "COMET_AUTO_LOG_GIT_PATCH",
+    "log_graph": "COMET_AUTO_LOG_GRAPH",
+    "name": "COMET_START_EXPERIMENT_NAME",
+    "offline_directory": "COMET_OFFLINE_DIRECTORY",
+    "parse_args": "COMET_AUTO_LOG_CLI_ARGUMENTS",
+    "tags": "COMET_START_EXPERIMENT_TAGS",
+}
+
+
+def python_value_to_environ_value(python_value):
+    if isinstance(python_value, bool):
+        if python_value is True:
+            return "true"
+
+        return "false"
+
+    if isinstance(python_value, int):
+        return str(python_value)
+
+    if isinstance(python_value, list):  # Comet only have one list of string parameter
+        return ",".join(map(str, python_value))
+
+    return python_value
+
+
+def setup_comet_env_vars(cfg: DictDefault):
+    # TODO, we need to convert Axolotl configuration to environment variables
+    # as Transformers integration are call first and would create an
+    # Experiment first
+
+    for key in cfg.keys():
+        if key.startswith("comet_") and key != "comet_experiment_config":
+            value = cfg.get(key, "")
+
+            if value is not None and value != "":
+                env_variable_name = COMET_ENV_MAPPING_OVERRIDE.get(key, key.upper())
+                final_value = python_value_to_environ_value(value)
+                os.environ[env_variable_name] = final_value
+
+    if cfg.comet_experiment_config:
+        for key, value in cfg.comet_experiment_config.items():
+            if value is not None and value != "":
+                config_env_variable_name = (
+                    COMET_EXPERIMENT_CONFIG_ENV_MAPPING_OVERRIDE.get(key)
+                )
+
+                if config_env_variable_name is None:
+                    LOG.warning(
+                        f"Unknown Comet Experiment Config name {key}, ignoring it"
+                    )
+                    continue
+
+                final_value = python_value_to_environ_value(value)
+                os.environ[config_env_variable_name] = final_value
+
+    # Enable comet if project name is present
+    if cfg.comet_project_name and len(cfg.comet_project_name) > 0:
+        cfg.use_comet = True
diff --git a/src/axolotl/utils/config/__init__.py b/src/axolotl/utils/config/__init__.py
index a054f24a7f..6450756e4b 100644
--- a/src/axolotl/utils/config/__init__.py
+++ b/src/axolotl/utils/config/__init__.py
@@ -1,22 +1,121 @@
 """Module for working with config dicts"""
+
+import copy
 import json
-import logging
 import os
-from pathlib import Path
 from typing import Optional
 
 import torch
+from pydantic import ValidationError
+from pydantic_core import PydanticUndefined, PydanticUndefinedType
 from transformers.utils import is_torch_bf16_gpu_available
+from transformers.utils.import_utils import (
+    is_torch_greater_or_equal,
+    is_torch_npu_available,
+)
 
+from axolotl.integrations.config import merge_input_args
+from axolotl.loaders.constants import MULTIMODAL_AUTO_MODEL_MAPPING
+from axolotl.loaders.utils import load_model_config
 from axolotl.utils.bench import log_gpu_memory_usage
-from axolotl.utils.config.models.input.v0_4_1 import (
-    AxolotlConfigWCapabilities,
-    AxolotlInputConfig,
-)
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model_config
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.config import (
+    AxolotlConfigWCapabilities as AxolotlConfigWCapabilitiesBase,
+    AxolotlInputConfig as AxolotlInputConfigBase,
+)
+from axolotl.utils.schemas.datasets import (
+    DPODataset,
+    KTODataset,
+    SFTDataset,
+    SyntheticDataset,
+)
+
+LOG = get_logger(__name__)
+
+_NONE_DEFAULT_FIELDS = {
+    "xformers_attention",
+    "sdp_attention",
+    "flex_attention",
+    "flash_attention",
+    "sage_attention",
+    "eager_attention",
+}
+
+
+def _is_pydantic_undefined(value):
+    return value is PydanticUndefined or isinstance(value, PydanticUndefinedType)
+
+
+def _field_name_for_missing_loc(model_cls, field_loc):
+    fields = getattr(model_cls, "model_fields", None) or {}
+    if field_loc in fields:
+        return field_loc
+
+    for field_name, field in fields.items():
+        if field_loc == getattr(field, "alias", None):
+            return field_name
+
+        validation_alias = getattr(field, "validation_alias", None)
+        if field_loc == validation_alias:
+            return field_name
+
+        for alias in getattr(validation_alias, "choices", ()) or ():
+            if field_loc == alias:
+                return field_name
+
+        alias_path = getattr(validation_alias, "path", None)
+        if alias_path and field_loc == alias_path[0]:
+            return field_name
 
-LOG = logging.getLogger("axolotl")
+    return field_loc
+
+
+def _field_default_from_mro(model_cls, field_name):
+    if field_name in _NONE_DEFAULT_FIELDS:
+        return None
+
+    for cls in model_cls.__mro__:
+        fields = getattr(cls, "model_fields", None)
+        if not fields or field_name not in fields:
+            continue
+
+        default = fields[field_name].get_default(call_default_factory=True)
+        if not _is_pydantic_undefined(default):
+            return copy.deepcopy(default)
+
+    return PydanticUndefined
+
+
+def _model_validate_with_field_names(model_cls, data):
+    try:
+        return model_cls.model_validate(data, by_alias=True, by_name=True)
+    except TypeError:
+        return model_cls(**data)
+
+
+def _model_with_inherited_default_fallback(model_cls, data):
+    try:
+        return model_cls(**data)
+    except ValidationError as exc:
+        missing_fields = {
+            _field_name_for_missing_loc(model_cls, err["loc"][0])
+            for err in exc.errors()
+            if err.get("type") == "missing" and len(err.get("loc", ())) == 1
+        }
+        if not missing_fields:
+            raise
+
+        data_with_defaults = dict(data)
+        for field_name in missing_fields:
+            if field_name in data_with_defaults:
+                continue
+            default = _field_default_from_mro(model_cls, field_name)
+            if _is_pydantic_undefined(default):
+                raise
+            data_with_defaults[field_name] = default
+
+        return _model_validate_with_field_names(model_cls, data_with_defaults)
 
 
 def choose_device(cfg):
@@ -28,8 +127,11 @@ def get_device():
             if torch.backends.mps.is_available():
                 return "mps"
 
-            raise SystemError("No CUDA/mps device found")
-        except Exception:  # pylint: disable=broad-exception-caught
+            if is_torch_npu_available():
+                return f"npu:{cfg.local_rank}"
+
+            raise SystemError("No CUDA/mps/npu device found")
+        except Exception:
             return "cpu"
 
     cfg.device = get_device()
@@ -38,6 +140,8 @@ def get_device():
     else:
         if cfg.device.startswith("cuda"):
             cfg.device_map = {"": torch.cuda.current_device()}
+        elif cfg.device.startswith("npu"):
+            cfg.device_map = {"npu": torch.npu.current_device()}
         else:
             cfg.device_map = {"": cfg.device}
 
@@ -48,50 +152,38 @@ def get_device():
         cfg.device_map = None
 
 
-def normalize_config(cfg):
-    # setup some derived config / hyperparams
-    cfg.gradient_accumulation_steps = cfg.gradient_accumulation_steps or (
-        cfg.batch_size // cfg.micro_batch_size
-    )
-    cfg.batch_size = (
-        cfg.batch_size or cfg.micro_batch_size * cfg.gradient_accumulation_steps
-    )
-    if cfg.eval_batch_size is None:
-        cfg.eval_batch_size = cfg.micro_batch_size
-    cfg.world_size = int(os.environ.get("WORLD_SIZE", 1))
-    cfg.local_rank = int(os.environ.get("LOCAL_RANK", 0))
-    cfg.eval_table_size = cfg.eval_table_size or 0
-    cfg.eval_max_new_tokens = cfg.eval_max_new_tokens or 128
-    cfg.eval_causal_lm_metrics = cfg.eval_causal_lm_metrics or [
-        "sacrebleu",
-        "comet",
-        "ter",
-        "chrf",
-    ]
-    choose_device(cfg)
-    cfg.ddp = cfg.ddp if cfg.ddp is not None else cfg.world_size != 1
-    if cfg.ddp:
-        cfg.device_map = {"": int(os.environ.get("LOCAL_RANK", 0))}
-        cfg.batch_size = cfg.batch_size * cfg.world_size
-
-    if cfg.bf16 == "auto":
+def resolve_dtype(cfg):
+    if (
+        not cfg.fp16 and cfg.bf16 == "auto" and not cfg.use_ray
+    ):  # if we use ray we want to defer this check to the worker node
         if is_torch_bf16_gpu_available():
             LOG.debug("bf16 support detected, enabling for this configuration.")
             cfg.bf16 = True
         else:
             LOG.debug("bf16 support not detected, disabling for this configuration.")
             cfg.bf16 = False
-            if cfg.fp16 is None:
+            if cfg.fp16 is None and not cfg.float16:
                 cfg.fp16 = True
 
+    if cfg.fp16 and cfg.bf16 == "auto":
+        cfg.bf16 = False
+
     if cfg.device == "mps":
         cfg.load_in_8bit = False
         cfg.tf32 = False
-        if cfg.bf16:
+        if cfg.bf16 and cfg.fp16 is not False:
             cfg.fp16 = True
         cfg.bf16 = False
     else:
-        torch.backends.cuda.matmul.allow_tf32 = cfg.tf32 or False
+        if cfg.tf32 is True:
+            torch.set_float32_matmul_precision("high")
+            if is_torch_greater_or_equal("2.9.0"):
+                torch.backends.fp32_precision = "tf32"
+                torch.backends.cuda.matmul.fp32_precision = "tf32"
+                torch.backends.cudnn.fp32_precision = "tf32"
+            else:
+                torch.backends.cuda.matmul.allow_tf32 = True
+                torch.backends.cudnn.allow_tf32 = True
         if cfg.bf16:
             cfg.fp16 = False
 
@@ -102,30 +194,114 @@ def normalize_config(cfg):
     else:
         cfg.torch_dtype = torch.float32
 
+
+def normalize_config(cfg):
+    # setup some derived config / hyperparams
+    if not cfg.use_ray:
+        cfg.gradient_accumulation_steps = cfg.gradient_accumulation_steps or (
+            cfg.batch_size // cfg.micro_batch_size
+        )
+        cfg.batch_size = (
+            cfg.batch_size or cfg.micro_batch_size * cfg.gradient_accumulation_steps
+        )
+    if cfg.eval_batch_size is None:
+        cfg.eval_batch_size = cfg.micro_batch_size
+    cfg.world_size = int(os.environ.get("WORLD_SIZE", 1))
+    cfg.local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    cfg.eval_table_size = cfg.eval_table_size or 0
+    cfg.eval_max_new_tokens = cfg.eval_max_new_tokens or 128
+    cfg.eval_causal_lm_metrics = cfg.eval_causal_lm_metrics or [
+        "sacrebleu",
+        "comet",
+        "ter",
+        "chrf",
+    ]
+    choose_device(cfg)
+    cfg.ddp = cfg.ddp if cfg.ddp is not None else cfg.world_size != 1
+    if cfg.world_size != 1:
+        cfg.device_map = {"": int(os.environ.get("LOCAL_RANK", 0))}
+        if cfg.fsdp or cfg.fsdp_config or cfg.ddp:
+            effective_world_size = (
+                cfg.world_size
+                // (cfg.context_parallel_size or 1)
+                // (cfg.tensor_parallel_size or 1)
+            )
+            cfg.batch_size = cfg.batch_size * effective_world_size
+
+    if not cfg.use_ray:
+        # delay resolving dtype until on worker node when launching with ray
+        resolve_dtype(cfg)
+
+    if cfg.deepspeed:
+        if isinstance(cfg.deepspeed, str) and os.path.exists(cfg.deepspeed):
+            ds_config_path = cfg.deepspeed
+            with open(ds_config_path, encoding="utf-8") as f:
+                cfg.deepspeed = json.load(f)
+
     if cfg.saves_per_epoch:
         save_steps = 1.0 / (cfg.saves_per_epoch * cfg.num_epochs)
         if save_steps < 1.0:  # prevent saves on every step
             cfg.save_steps = save_steps
+        elif save_steps > 1:
+            LOG.warning(
+                f"Invalid value for save_steps ({save_steps}) from saves_per_epoch and/or num_epochs. Saving at training end only."
+            )
     if (cfg.val_set_size or cfg.test_datasets) and cfg.evals_per_epoch:
         eval_steps = 1.0 / (cfg.evals_per_epoch * cfg.num_epochs)
         if eval_steps < 1.0:  # prevent evals on every step
             cfg.eval_steps = eval_steps
-
-    cfg.dataset_processes = cfg.dataset_processes or os.cpu_count()
+        elif eval_steps > 1:
+            LOG.warning(
+                f"Invalid value for eval_steps ({eval_steps}) from evals_per_epoch and/or num_epochs. Skipping evaluations."
+            )
 
     if not cfg.base_model_config:
         cfg.base_model_config = cfg.base_model
 
+    # Apply pre-config load patches (e.g., for Kimi Linear remote code patching)
+    from axolotl.loaders.patch_manager import PatchManager
+
+    PatchManager.apply_pre_config_load_patches(cfg)
+
     model_config = load_model_config(cfg)
-    cfg.model_config_type = model_config.model_type
 
     cfg.tokenizer_config = (
         cfg.tokenizer_config or cfg.base_model_config or cfg.base_model
     )
 
+    cfg.is_multimodal = (
+        hasattr(model_config, "model_type")
+        and model_config.model_type in MULTIMODAL_AUTO_MODEL_MAPPING
+        or any(
+            multimodal_name in cfg.base_model.lower()
+            for multimodal_name in [
+                "pixtral",
+            ]
+        )
+        or cfg.is_multimodal
+    )
+    if cfg.is_multimodal:
+        cfg.processor_config = (
+            cfg.processor_config or cfg.base_model_config or cfg.base_model
+        )
+
+    cfg.model_config_type = model_config.model_type
+
+    # Resolve inner text backbone type for VLM wrappers (e.g. mistral3 -> mistral4)
+    if callable(getattr(model_config, "get_text_config", None)):
+        text_config = model_config.get_text_config()
+        if (
+            hasattr(text_config, "model_type")
+            and text_config.model_type != model_config.model_type
+        ):
+            cfg.model_config_type_text = text_config.model_type
+
     # figure out if the model is llama
     cfg.is_llama_derived_model = (
-        (hasattr(model_config, "model_type") and model_config.model_type == "llama")
+        (
+            hasattr(model_config, "model_type")
+            and model_config.model_type in ["llama", "mllama_text_model"]
+        )
         or cfg.is_llama_derived_model
         or "llama" in cfg.base_model.lower()
         or (cfg.type_of_model and "llama" in cfg.type_of_model.lower())
@@ -179,6 +355,38 @@ def normalize_config(cfg):
     ):
         cfg.gradient_checkpointing_kwargs = {"use_reentrant": True}
 
+    # Gemma4 requires use_reentrant=False for DDP (shared per-layer norms /
+    # cross-layer shared KV cause "marked ready twice" errors with reentrant
+    # checkpointing) and ddp_find_unused_parameters=True (per_layer_projection
+    # LoRA / frozen mm params may not receive gradients on every step). The
+    # unified variant shares the same constraints.
+    if cfg.model_config_type in ("gemma4", "gemma4_unified"):
+        if cfg.gradient_checkpointing:
+            if cfg.gradient_checkpointing_kwargs is None:
+                cfg.gradient_checkpointing_kwargs = {}
+            if cfg.gradient_checkpointing_kwargs.get("use_reentrant") is not False:
+                LOG.warning(
+                    "Gemma4 requires use_reentrant=False for gradient checkpointing "
+                    "in distributed training. Setting use_reentrant=False."
+                )
+                cfg.gradient_checkpointing_kwargs["use_reentrant"] = False
+        if cfg.ddp and cfg.ddp_find_unused_parameters is None:
+            if cfg.activation_offloading is True:
+                # activation_offloading uses checkpoint wrappers that conflict
+                # with find_unused_parameters (causes "marked ready twice").
+                # Use freeze_mm_modules instead to eliminate unused params.
+                LOG.info(
+                    "Gemma4 + DDP + activation_offloading: skipping "
+                    "ddp_find_unused_parameters (use freeze_mm_modules to "
+                    "handle unused vision/audio params)."
+                )
+            else:
+                LOG.warning(
+                    "Gemma4 requires ddp_find_unused_parameters=True for DDP. "
+                    "Auto-enabling."
+                )
+                cfg.ddp_find_unused_parameters = True
+
     log_gpu_memory_usage(LOG, "baseline", cfg.device)
 
 
@@ -187,419 +395,92 @@ def normalize_cfg_datasets(cfg):
     helpers for mapping chat_template to various dataset configurations as necessary
     """
 
-    if cfg.chat_template and cfg.chat_template == "chatml":
+    if cfg.chat_template:
         if cfg.datasets:
             for idx, ds_cfg in enumerate(cfg.datasets):
-                if ds_cfg.type == "sharegpt" and not ds_cfg.conversation:
-                    LOG.info(
-                        f"updating dataset {ds_cfg.path} with `conversation: chatml` to match your chat_template"
-                    )
-                    cfg.datasets[idx].conversation = "chatml"
-                if ds_cfg.type == "orpo.chat_template" and not ds_cfg.chat_template:
+                if (
+                    ds_cfg.type in ["orpo.chat_template", "chat_template"]
+                    and not ds_cfg.chat_template
+                ):
                     LOG.info(
-                        f"updating dataset {ds_cfg.path} with `chat_template: chatml` to match your chat_template"
+                        f"updating dataset {ds_cfg.path} with `chat_template: {cfg.chat_template}` to match your chat_template"
                     )
-                    cfg.datasets[idx].chat_template = "chatml"
+                    cfg.datasets[idx].chat_template = cfg.chat_template
+                    cfg.datasets[idx].chat_template_jinja = cfg.chat_template_jinja
 
 
-def validate_config(cfg: DictDefault, capabilities: Optional[dict] = None):
-    if capabilities:
-        return DictDefault(
-            dict(
-                AxolotlConfigWCapabilities(
-                    **cfg.to_dict(), capabilities=capabilities
-                ).model_dump(exclude_none=True)
-            )
-        )
-    return DictDefault(
-        dict(AxolotlInputConfig(**cfg.to_dict()).model_dump(exclude_none=True))
-    )
+def validate_config(
+    cfg: DictDefault,
+    capabilities: Optional[dict] = None,
+    env_capabilities: Optional[dict] = None,
+) -> DictDefault:
+    AxolotlConfigWCapabilities = AxolotlConfigWCapabilitiesBase
+    AxolotlInputConfig = AxolotlInputConfigBase
 
+    if cfg.plugins:
+        (
+            AxolotlConfigWCapabilities,
+            AxolotlInputConfig,
+        ) = merge_input_args()
+
+    # Convert datasets to proper format if needed
+    if cfg.get("datasets"):
+        for idx, ds_cfg in enumerate(cfg["datasets"]):
+            if cfg.get("rl") in ["dpo", "ipo", "simpo"] and not isinstance(
+                ds_cfg, DPODataset
+            ):
+                cfg["datasets"][idx] = DPODataset(**ds_cfg)
+            elif cfg.get("rl") == "kto" and not isinstance(ds_cfg, KTODataset):
+                cfg["datasets"][idx] = KTODataset(**dict(ds_cfg))
+            elif (
+                ds_cfg.get("type")
+                if isinstance(ds_cfg, dict)
+                else getattr(ds_cfg, "type", None)
+            ) == "_synthetic" and not isinstance(ds_cfg, SyntheticDataset):
+                cfg["datasets"][idx] = SyntheticDataset(
+                    **(ds_cfg if isinstance(ds_cfg, dict) else dict(ds_cfg))
+                )
+            elif not isinstance(ds_cfg, SFTDataset):
+                cfg["datasets"][idx] = SFTDataset(**dict(ds_cfg))
 
-def legacy_validate_config(cfg):
-    """
-    This is a "pre-validation" step that handles the yaml configuration before we have any
-    information about the model architecture
-    """
-    if is_torch_bf16_gpu_available():
-        if not cfg.bf16 and not cfg.bfloat16:
-            LOG.info("bf16 support detected, but not enabled for this configuration.")
-    else:
-        if (
-            not cfg.merge_lora
-            and not cfg.is_preprocess
-            and (cfg.bf16 is True or cfg.bfloat16 is True)
+    if capabilities or env_capabilities:
+        if (capabilities and env_capabilities is None) or (
+            env_capabilities and capabilities is None
         ):
             raise ValueError(
-                "bf16 requested, but AMP is not supported on this GPU. Requires Ampere series or above."
-            )
-    if (
-        # pylint: disable=too-many-boolean-expressions
-        not (cfg.bf16 or cfg.bfloat16)
-        and (cfg.fp16 or cfg.float16)
-        and not cfg.adapter
-        and not cfg.flash_attention
-        and cfg.sample_packing
-    ):
-        LOG.warning(
-            "Full fine tune w/o FA2 w/ sample packing and fp16/float16 is likely to raise errors. Try LoRA."
-        )
-        # ValueError: Attempting to unscale FP16 gradients.
-        # OR
-        # RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::Half
-    if cfg.max_packed_sequence_len:
-        raise DeprecationWarning("`max_packed_sequence_len` is no longer supported")
-
-    if cfg.sample_packing and cfg.rl:
-        raise ValueError("`sample_packing: true` does not work with RLHF training")
-
-    if cfg.sample_packing and not cfg.pad_to_sequence_len:
-        LOG.warning(
-            "`pad_to_sequence_len: true` is recommended when using sample_packing"
-        )
-
-    if cfg.gradient_accumulation_steps and cfg.batch_size:
-        raise ValueError(
-            "please set only one of gradient_accumulation_steps or batch_size"
-        )
-    if cfg.batch_size:
-        LOG.warning(
-            "%s\n%s",
-            "batch_size is not recommended. Please use gradient_accumulation_steps instead.",
-            "To calculate the equivalent gradient_accumulation_steps, divide batch_size / micro_batch_size / number of gpus.",
-        )
-    if (
-        cfg.eval_batch_size
-        and cfg.micro_batch_size
-        and cfg.eval_batch_size != cfg.micro_batch_size
-    ):
-        LOG.warning(
-            "eval_batch_size != micro_batch_size. This can lead to VRAM instability."
-        )
-
-    if cfg.adapter == "qlora":
-        if cfg.merge_lora:
-            # can't merge qlora if loaded in 8bit or 4bit
-            if cfg.load_in_8bit:
-                raise ValueError("Can't merge qlora if loaded in 8bit")
-
-            if cfg.gptq:
-                raise ValueError("Can't merge qlora if gptq")
-
-            if cfg.load_in_4bit:
-                raise ValueError("Can't merge qlora if loaded in 4bit")
-
-        else:
-            if cfg.load_in_8bit:
-                raise ValueError("Can't load qlora in 8bit")
-
-            if cfg.gptq:
-                raise ValueError("Can't load qlora if gptq")
-
-            if not cfg.load_in_4bit:
-                raise ValueError("Require cfg.load_in_4bit to be True for qlora")
-
-        if cfg.flash_attn_fuse_qkv or cfg.flash_attn_fuse_mlp:
-            raise ValueError("Fused modules are not supported with QLoRA")
-
-    loftq = cfg.peft and cfg.peft.loftq_config and cfg.peft.loftq_config.loftq_bits
-    if not cfg.load_in_8bit and cfg.adapter == "lora" and not loftq:
-        LOG.warning("We recommend setting `load_in_8bit: true` for LORA finetuning")
-
-    if cfg.adapter == "lora" and (cfg.flash_attn_fuse_qkv or cfg.flash_attn_fuse_mlp):
-        raise ValueError("Fused modules are not supported with LoRA")
-
-    if cfg.adapter and cfg.peft_layers_to_transform and cfg.unfrozen_parameters:
-        raise ValueError(
-            "`unfrozen_parameters` used with `peft_layers_to_transform` can have unexpected behavior."
-        )
-
-    if cfg.relora_steps:
-        if cfg.adapter not in ("lora", "qlora"):
-            raise ValueError("cfg.adapter must be lora or qlora to use ReLoRA")
-
-        if cfg.fsdp:
-            raise ValueError("fsdp not supported with ReLoRA")
-
-        if cfg.deepspeed:
-            raise ValueError("deepspeed not supported with ReLoRA")
-
-        if cfg.lr_scheduler == "one_cycle":
-            raise ValueError("ReLoRA is not compatible with the one_cycle scheduler")
-
-        if cfg.flash_attn_fuse_qkv or cfg.flash_attn_fuse_mlp:
-            raise ValueError("Fused modules are not supported with ReLoRA")
-
-    if cfg.trust_remote_code:
-        LOG.warning(
-            "`trust_remote_code` is set to true. Please make sure that you reviewed the remote code/model."
-        )
-
-    if cfg.push_dataset_to_hub and cfg.hf_use_auth_token is not True:
-        raise ValueError(
-            "Require cfg.hf_use_auth_token to be True for push_dataset_to_hub"
-        )
-
-    if (cfg.base_model and "falcon" in cfg.base_model.lower()) and cfg.fsdp:
-        raise ValueError("FSDP is not supported for falcon models")
-
-    if (
-        cfg.base_model and "mpt" in cfg.base_model.lower()
-    ) and cfg.gradient_checkpointing:
-        raise ValueError("gradient_checkpointing is not supported for MPT models")
-
-    if cfg.flash_optimum is True:
-        if cfg.adapter:
-            LOG.warning("BetterTransformers probably doesn't work with PEFT adapters")
-        if cfg.fp16 or cfg.bf16:
-            raise ValueError("AMP is not supported with BetterTransformer")
-        if cfg.float16 is not True and cfg.bfloat16 is not True:
-            LOG.warning(
-                "You should probably set bfloat16 or float16 to true to "
-                "load the model in float16 for BetterTransformers"
-            )
-        if int(torch.__version__.split(".", maxsplit=1)[0]) < 2:
-            LOG.warning("torch>=2.0.0 required")
-            raise ValueError(
-                f"flash_optimum for BetterTransformers may not be used with {torch.__version__}"
-            )
-
-    if cfg.pretraining_dataset and cfg.group_by_length:
-        LOG.warning(
-            "You probably want to disable group_by_length as it will force a streamed dataset to download completely."
-        )
-    if cfg.pretraining_dataset and not cfg.max_steps:
-        raise ValueError(
-            "max_steps must be set when using iterable pretraining_dataset, Trainer can't infer length and schedule optimizer/learning rate without it!"
-        )
-
-    if any([cfg.adam_beta1, cfg.adam_beta2, cfg.adam_epsilon]) and (
-        not cfg.optimizer or "adamw" not in cfg.optimizer
-    ):
-        LOG.warning("adamw hyperparameters found, but no adamw optimizer set")
-
-    if cfg.push_to_hub_model_id:
-        raise ValueError(
-            "push_to_hub_model_id is deprecated. Please use hub_model_id instead."
-        )
-
-    if cfg.hub_model_id and cfg.save_strategy not in ["steps", "epoch", None]:
-        LOG.warning(
-            "hub_model_id is set without any models being saved. To save a model, set save_strategy to steps, epochs or leave empty."
-        )
-
-    if cfg.gptq and cfg.revision_of_model:
-        raise ValueError(
-            "revision_of_model is not supported for GPTQ models. "
-            + "Please download the model from HuggingFace Hub manually for correct branch, "
-            + "point to its path, and remove revision_of_model from the config."
-        )
-
-    # if cfg.sample_packing and cfg.sdp_attention:
-    #     # incompatible due to bug w/ accelerate causing 0.0 loss when using llama2
-    #     raise ValueError(
-    #         "sample_packing not compatible with sdp_attention. Use flash_attention"
-    #     )
-
-    if cfg.sample_packing and cfg.xformers_attention:
-        raise ValueError(
-            "sample_packing not compatible with xformers_attention. Use flash_attention"
-        )
-
-    if cfg.sample_packing and cfg.sdp_attention and (cfg.bfloat16 or cfg.bf16):
-        # https://github.com/pytorch/pytorch/blob/1b03423526536b5f3d35bdfa95ccc6197556cf9b/test/test_transformers.py#L2440-L2450
-        LOG.warning(
-            "sample_packing & torch sdpa with bf16 is unsupported may results in 0.0 loss. "
-            "This may work on H100s."
-        )
-
-    if cfg.early_stopping_patience:
-        if not cfg.save_steps or not cfg.eval_steps:
-            raise ValueError(
-                "`early_stopping_patience` requires save_steps and eval_steps to be set. eval_steps should evenly divide save_steps."
+                "Both capabilities and env_capabilities must be provided or not provided."
             )
-        if cfg.save_steps % cfg.eval_steps != 0:
-            raise ValueError(
-                "`early_stopping_patience` requires that eval_steps should evenly divide save_steps."
-            )
-
-    if cfg.datasets:
-        for idx, ds_cfg in enumerate(cfg.datasets):
-            if not ds_cfg.type:
-                continue
-            if ds_cfg.type == "sharegpt:chat":
-                LOG.warning(
-                    PendingDeprecationWarning(
-                        "`type: sharegpt:chat` will soon be deprecated. simply use `type: sharegpt` instead."
-                    )
-                )
-                cfg.datasets[idx].type = "sharegpt"
-            if "sharegpt_simple" in ds_cfg.type:
-                LOG.warning(
-                    PendingDeprecationWarning(
-                        "`type: sharegpt_simple` will soon be deprecated. simply use `type: sharegpt` instead."
-                    )
-                )
-                cfg.datasets[idx].type = cfg.datasets[idx].type.replace(
-                    "sharegpt_simple", "sharegpt"
-                )
-
-    if cfg.saves_per_epoch and cfg.save_steps:
-        raise ValueError(
-            "save_steps and saves_per_epoch are mutually exclusive and cannot be used together."
-        )
-    if cfg.save_strategy and cfg.saves_per_epoch and cfg.save_strategy != "steps":
-        raise ValueError(
-            "save_strategy must be empty or set to `steps` when used with saves_per_epoch."
-        )
-    if cfg.save_strategy and cfg.save_steps and cfg.save_strategy != "steps":
-        raise ValueError(
-            "save_strategy and save_steps mismatch. Please set save_strategy to 'steps' or remove save_steps."
-        )
-    if cfg.evals_per_epoch and cfg.eval_steps:
-        raise ValueError(
-            "eval_steps and evals_per_epoch are mutually exclusive and cannot be used together."
-        )
-    if (
-        cfg.evals_per_epoch
-        and cfg.evaluation_strategy
-        and cfg.evaluation_strategy != "steps"
-    ):
-        raise ValueError(
-            "evaluation_strategy must be empty or set to `steps` when used with evals_per_epoch."
-        )
-    if (
-        cfg.evaluation_strategy
-        and cfg.eval_steps
-        and cfg.evaluation_strategy != "steps"
-    ):
-        raise ValueError(
-            "evaluation_strategy and eval_steps mismatch. Please set evaluation_strategy to 'steps' or remove eval_steps."
-        )
-
-    if (
-        cfg.val_set_size == 0
-        and (cfg.eval_steps or cfg.evaluation_strategy)
-        and not cfg.test_datasets
-    ):
-        raise ValueError(
-            "eval_steps and evaluation_strategy are not supported with val_set_size == 0"
-        )
-
-    if (
-        cfg.sample_packing
-        and cfg.eval_table_size
-        and cfg.eval_sample_packing is not False
-    ):
-        raise ValueError(
-            "eval_table_size and eval_sample_packing are not supported together with sample_packing. Please set 'eval_sample_packing' to false."
-        )
 
-    if not cfg.adapter and (cfg.load_in_8bit or cfg.load_in_4bit):
-        raise ValueError(
-            "load_in_8bit and load_in_4bit are not supported without setting an adapter."
-            "If you want to full finetune, please turn off load_in_8bit and load_in_4bit."
-        )
-
-    if cfg.rope_scaling:
-        LOG.warning("`rope_scaling` should now be be a key under `model_config`")
-
-    if cfg.wandb_run_id and not cfg.wandb_name:
-        cfg.wandb_name = cfg.wandb_run_id
-
-        LOG.warning(
-            "wandb_run_id sets the ID of the run. If you would like to set the name, please use wandb_name instead."
-        )
-
-    if cfg.noisy_embedding_alpha is not None:
-        # Deprecated, use neftune_noise_alpha
-        LOG.warning("noisy_embedding_alpha is deprecated, use neftune_noise_alpha")
-        if cfg.neftune_noise_alpha is None:
-            cfg.neftune_noise_alpha = cfg.noisy_embedding_alpha
-        else:
-            # User is providing both; bail and have them sort out their settings
-            raise ValueError(
-                "noisy_embedding_alpha is deprecated, use neftune_noise_alpha; both are set, please remove the deprecated noisy_embedding_alpha setting"
+        return DictDefault(
+            dict(
+                _model_with_inherited_default_fallback(
+                    AxolotlConfigWCapabilities,
+                    {
+                        **cfg.to_dict(),
+                        "capabilities": capabilities,
+                        "env_capabilities": env_capabilities,
+                    },
+                ).model_dump(exclude_none=True)
             )
-
-    if cfg.neftune_noise_alpha is not None and cfg.neftune_noise_alpha <= 0.0:
-        raise ValueError("neftune_noise_alpha must be > 0.0")
-
-    if cfg.max_memory is not None and cfg.gpu_memory_limit is not None:
-        raise ValueError(
-            "max_memory and gpu_memory_limit are mutually exclusive and cannot be used together."
         )
 
-    if (
-        cfg.unfrozen_parameters
-        and cfg.gradient_checkpointing_kwargs
-        and cfg.gradient_checkpointing_kwargs.use_reentrant is True
-    ):
-        # https://github.com/huggingface/transformers/issues/21381
-        raise ValueError(
-            "`use_reentrant` must be false when used with partially frozen model."
-        )
-
-    if cfg.deepspeed and Path(cfg.deepspeed).is_file():
-        with open(cfg.deepspeed, encoding="utf-8") as file:
-            contents = file.read()
-            deepspeed_cfg: DictDefault = DictDefault(json.loads(contents))
-            if cfg.flash_attention:
-                if (
-                    deepspeed_cfg.zero_optimization
-                    and deepspeed_cfg.zero_optimization.stage == 3
-                ):
-                    if not (
-                        (
-                            deepspeed_cfg.bf16
-                            and deepspeed_cfg.bf16.enabled  # pylint: disable=no-member
-                            is True
-                        )
-                        or (
-                            deepspeed_cfg.fp16
-                            and deepspeed_cfg.fp16.enabled  # pylint: disable=no-member
-                            is True
-                        )
-                    ):
-                        raise ValueError(
-                            "bf16.enabled or fp16.enabled must be set to true when using ZeRO-3 with flash-attention"
-                        )
-            if "8bit" in cfg.optimizer and deepspeed_cfg.optimizer:
-                LOG.warning(
-                    f"conflicting optimizer: {cfg.optimizer} used alongside deepspeed optimizer."
-                )
-
-    if cfg.test_datasets and cfg.val_set_size:
-        raise ValueError(
-            "non-zero val_set_size should not be used with test_datasets configuration"
+    return DictDefault(
+        dict(
+            _model_with_inherited_default_fallback(
+                AxolotlInputConfig, cfg.to_dict()
+            ).model_dump(exclude_none=True)
         )
+    )
 
-    if cfg.fsdp and "bnb" in cfg.optimizer:
-        raise ValueError(f"FSDP not compatible with {cfg.optimizer}")
 
-    if cfg.do_causal_lm_eval and cfg.eval_sample_packing:
-        raise ValueError(
-            "do_causal_lm_eval is enabled, eval_sample_packing must be set to False"
-        )
+def prepare_plugins(cfg):
+    """
+    Prepare the plugins for the configuration
+    """
 
-    if cfg.eval_causal_lm_metrics:
-        supported_metrics = ["sacrebleu", "comet", "ter", "chrf"]
-        if not isinstance(cfg.eval_causal_lm_metrics, list):
-            raise ValueError("eval_causal_lm_metrics must be a list")
-        # only ["sacrebleu", "comet", "ter", "chrf"] supported
-        if set(cfg.eval_causal_lm_metrics) - set(supported_metrics):
-            raise ValueError(
-                f"eval_causal_lm_metrics must be one of {supported_metrics}"
-            )
+    if cfg.get("plugins"):
+        from axolotl.integrations.base import PluginManager
 
-    # TODO
-    # MPT 7b
-    # https://github.com/facebookresearch/bitsandbytes/issues/25
-    # no 8bit adaAmw w bf16
-
-    # GPT-NeoX
-    # evals broken when extending context len
-    # File "/root/miniconda3/envs/py3.9/lib/python3.9/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py", line 162, in forward                        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-    # File "/root/miniconda3/envs/py3.9/lib/python3.9/site-packages/optimum/bettertransformer/models/attention.py", line 74, in gpt2_wrapped_scaled_dot_product
-    # attention_mask = causal_mask + attention_mask
-    # RuntimeError: The size of tensor a (2048) must match the size of tensor b (8132) at non-singleton dimension 3
+        plugin_manager = PluginManager.get_instance()
+        for plugin_name in cfg["plugins"]:
+            plugin_manager.register(plugin_name)
diff --git a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py
deleted file mode 100644
index f1c12b2ba0..0000000000
--- a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py
+++ /dev/null
@@ -1,1090 +0,0 @@
-"""
-Module for pydantic models for configuration
-"""
-
-# pylint: disable=too-many-lines
-
-import logging
-import os
-from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
-
-from pydantic import BaseModel, Field, conlist, field_validator, model_validator
-from transformers import SchedulerType
-from transformers.training_args import OptimizerNames
-
-from axolotl.utils.config.models.internals import GPUCapabilities
-
-LOG = logging.getLogger("axolotl.utils.config.models.input")
-
-
-class DeprecatedParameters(BaseModel):
-    """configurations that are deprecated"""
-
-    max_packed_sequence_len: Optional[int] = None
-    rope_scaling: Optional[Any] = None
-    noisy_embedding_alpha: Optional[float] = None
-
-    @field_validator("max_packed_sequence_len")
-    @classmethod
-    def validate_max_packed_sequence_len(cls, max_packed_sequence_len):
-        if max_packed_sequence_len:
-            raise DeprecationWarning("`max_packed_sequence_len` is no longer supported")
-        return max_packed_sequence_len
-
-    @field_validator("rope_scaling")
-    @classmethod
-    def validate_rope_scaling(cls, rope_scaling):
-        if rope_scaling:
-            raise DeprecationWarning(
-                "`rope_scaling` is no longer supported, it should now be be a key under `model_config`"
-            )
-        return rope_scaling
-
-    @field_validator("noisy_embedding_alpha")
-    @classmethod
-    def validate_noisy_embedding_alpha(cls, noisy_embedding_alpha):
-        if noisy_embedding_alpha:
-            LOG.warning("noisy_embedding_alpha is deprecated, use neftune_noise_alpha")
-        return noisy_embedding_alpha
-
-
-class RemappedParameters(BaseModel):
-    """parameters that have been remapped to other names"""
-
-    overrides_of_model_config: Optional[Dict[str, Any]] = Field(
-        default=None, alias="model_config"
-    )
-    type_of_model: Optional[str] = Field(default=None, alias="model_type")
-    revision_of_model: Optional[str] = Field(default=None, alias="model_revision")
-
-
-class PretrainingDataset(BaseModel):
-    """pretraining dataset configuration subset"""
-
-    name: Optional[str] = None
-    path: Optional[str] = None
-    split: Optional[str] = "train"
-    text_column: Optional[str] = "text"
-    type: Optional[str] = "pretrain"
-
-
-class UserDefinedPrompterType(BaseModel):
-    """structure for user defined prompt types"""
-
-    system_prompt: Optional[str] = None
-    system_format: Optional[str] = None
-    field_system: Optional[str] = None
-    field_instruction: Optional[str] = None
-    field_input: Optional[str] = None
-    field_output: Optional[str] = None
-
-    format: Optional[str] = None
-    no_input_format: Optional[str] = None
-    field: Optional[str] = None
-
-
-class SFTDataset(BaseModel):
-    """SFT configuration subset"""
-
-    path: Optional[str] = None
-    split: Optional[str] = None
-    type: Optional[Union[str, UserDefinedPrompterType]] = None
-    shards: Optional[int] = None
-    conversation: Optional[str] = None
-    chat_template: Optional[str] = None
-    data_files: Optional[Union[str, List[str]]] = None
-    name: Optional[str] = None
-    ds_type: Optional[str] = None
-    train_on_split: Optional[str] = None
-
-    field: Optional[str] = None
-    field_human: Optional[str] = None
-    field_model: Optional[str] = None
-
-    roles: Optional[Dict[str, List[str]]] = None
-
-
-class UserDefinedDPOType(BaseModel):
-    """User defined typing for DPO"""
-
-    field_system: Optional[str] = None
-    field_prompt: Optional[str] = None
-    field_chosen: Optional[str] = None
-    field_rejected: Optional[str] = None
-    prompt_format: Optional[str] = None
-    chosen_format: Optional[str] = None
-    rejected_format: Optional[str] = None
-
-
-class DPODataset(BaseModel):
-    """DPO configuration subset"""
-
-    path: Optional[str] = None
-    split: Optional[str] = None
-    type: Optional[Union[UserDefinedDPOType, str]] = None
-    data_files: Optional[List[str]] = None
-
-
-class RLType(str, Enum):
-    """RL trainer type configuration subset"""
-
-    dpo = "dpo"  # pylint: disable=invalid-name
-    ipo = "ipo"  # pylint: disable=invalid-name
-    kto_pair = "kto_pair"  # pylint: disable=invalid-name
-    orpo = "orpo"  # pylint: disable=invalid-name
-
-
-class ChatTemplate(str, Enum):
-    """Chat templates configuration subset"""
-
-    alpaca = "alpaca"  # pylint: disable=invalid-name
-    chatml = "chatml"  # pylint: disable=invalid-name
-    inst = "inst"  # pylint: disable=invalid-name
-    gemma = "gemma"  # pylint: disable=invalid-name
-    cohere = "cohere"  # pylint: disable=invalid-name
-    llama3 = "llama3"  # pylint: disable=invalid-name
-
-
-class LoftQConfig(BaseModel):
-    """LoftQ configuration subset"""
-
-    loftq_bits: int = Field(default=4, metadata={"help": "Quantization bits for LoftQ"})
-    # loftq_iter: int = Field(default=1, metadata={"help": "Alternating iterations for LoftQ"})
-
-
-class PeftConfig(BaseModel):
-    """peftq configuration subset"""
-
-    loftq_config: Optional[LoftQConfig] = None
-
-
-class SpecialTokensConfig(BaseModel):
-    """Special tokens configuration subset"""
-
-    bos_token: Optional[str] = None
-    eos_token: Optional[str] = None
-    pad_token: Optional[str] = None
-    unk_token: Optional[str] = None
-    additional_special_tokens: Optional[List[str]] = None
-
-
-class LoraConfig(BaseModel):
-    """Peft / LoRA configuration subset"""
-
-    load_in_8bit: Optional[bool] = Field(default=False)
-    load_in_4bit: Optional[bool] = Field(default=False)
-
-    adapter: Optional[str] = None
-    lora_model_dir: Optional[str] = None
-    lora_r: Optional[int] = None
-    lora_alpha: Optional[int] = None
-    lora_fan_in_fan_out: Optional[bool] = None
-    lora_target_modules: Optional[List[str]] = None
-    lora_target_linear: Optional[bool] = None
-    lora_modules_to_save: Optional[List[str]] = None
-    lora_dropout: Optional[float] = None
-    peft_layers_to_transform: Optional[List[int]] = None
-    peft: Optional[PeftConfig] = None
-    peft_use_dora: Optional[bool] = None
-    peft_use_rslora: Optional[bool] = None
-    peft_layer_replication: Optional[List[Tuple[int, int]]] = None
-
-    lora_on_cpu: Optional[bool] = None
-    gptq: Optional[bool] = None
-    bnb_config_kwargs: Optional[Dict[str, Any]] = None
-
-    loraplus_lr_ratio: Optional[float] = Field(
-        default=None,
-        metadata={
-            "help": "loraplus learning rate ratio lr_B / lr_A. Recommended value is 2^4."
-        },
-    )
-    loraplus_lr_embedding: Optional[float] = Field(
-        default=1e-6,
-        metadata={"help": "loraplus learning rate for lora embedding layers."},
-    )
-
-    merge_lora: Optional[bool] = None
-
-    @model_validator(mode="before")
-    @classmethod
-    def validate_adapter(cls, data):
-        if not data.get("adapter") and (
-            data.get("load_in_8bit") or data.get("load_in_4bit")
-        ):
-            raise ValueError(
-                "load_in_8bit and load_in_4bit are not supported without setting an adapter."
-                "If you want to full finetune, please turn off load_in_8bit and load_in_4bit."
-            )
-        return data
-
-    @model_validator(mode="after")
-    def validate_qlora(self):
-        if self.adapter == "qlora":
-            if self.merge_lora:
-                # can't merge qlora if loaded in 8bit or 4bit
-                if self.load_in_8bit:
-                    raise ValueError("Can't merge qlora if loaded in 8bit")
-
-                if self.gptq:
-                    raise ValueError("Can't merge qlora if gptq")
-
-                if self.load_in_4bit:
-                    raise ValueError("Can't merge qlora if loaded in 4bit")
-
-            else:
-                if self.load_in_8bit:
-                    raise ValueError("Can't load qlora in 8bit")
-
-                if self.gptq:
-                    raise ValueError("Can't load qlora if gptq")
-
-                if not self.load_in_4bit:
-                    raise ValueError("Require cfg.load_in_4bit to be True for qlora")
-        return self
-
-
-class ReLoRAConfig(BaseModel):
-    """ReLoRA configuration subset"""
-
-    relora_steps: Optional[int] = None
-    relora_warmup_steps: Optional[int] = None
-    relora_anneal_steps: Optional[int] = None
-    relora_prune_ratio: Optional[float] = None
-    relora_cpu_offload: Optional[bool] = None
-
-
-class ModelInputConfig(BaseModel):
-    """model to train on configuration subset"""
-
-    base_model: str
-    base_model_config: Optional[str] = None
-    cls_model_config: Optional[str] = None
-    tokenizer_config: Optional[str] = None
-    tokenizer_use_fast: Optional[bool] = None
-    tokenizer_legacy: Optional[bool] = None
-    tokenizer_type: Optional[str] = Field(
-        default=None, metadata={"help": "transformers tokenizer class"}
-    )
-    trust_remote_code: Optional[bool] = None
-
-    @field_validator("trust_remote_code")
-    @classmethod
-    def hint_trust_remote_code(cls, trust_remote_code):
-        if trust_remote_code:
-            LOG.warning(
-                "`trust_remote_code` is set to true. Please make sure that you reviewed the remote code/model."
-            )
-        return trust_remote_code
-
-
-class HyperparametersConfig(BaseModel):
-    """training hyperparams configuration subset"""
-
-    gradient_accumulation_steps: Optional[int] = Field(default=1)
-    micro_batch_size: Optional[int] = Field(
-        default=1,
-        metadata={"help": "per gpu micro batch size for training"},
-    )
-    batch_size: Optional[int] = Field(
-        default=None,
-        metadata={
-            "help": "Total batch size, we do not recommended setting this manually"
-        },
-    )
-    eval_batch_size: Optional[int] = Field(
-        default=None,
-        metadata={
-            "help": "per gpu micro batch size for evals, defaults to value of micro_batch_size"
-        },
-    )
-
-    train_on_inputs: Optional[bool] = False
-    group_by_length: Optional[bool] = None
-
-    learning_rate: Union[str, float]
-    weight_decay: Optional[float] = 0.0
-    optimizer: Optional[
-        Union[OptimizerNames, Literal["lion_pytorch"]]
-    ] = OptimizerNames.ADAMW_HF.value
-    optim_args: Optional[Union[str, Dict[str, Any]]] = Field(
-        default=None, metadata={"help": "Optional arguments to supply to optimizer."}
-    )
-    optim_target_modules: Optional[Union[List[str], Literal["all_linear"]]] = Field(
-        default=None,
-        metadata={
-            "help": "The target modules to optimize, i.e. the module names that you would like to train."
-        },
-    )
-    torchdistx_path: Optional[str] = None
-    lr_scheduler: Optional[SchedulerType] = "cosine"
-    lr_scheduler_kwargs: Optional[Dict[str, Any]] = None
-    lr_quadratic_warmup: Optional[bool] = None
-    cosine_min_lr_ratio: Optional[float] = None
-    cosine_constant_lr_ratio: Optional[float] = None
-    lr_div_factor: Optional[float] = None
-
-    adam_epsilon: Optional[float] = None
-    adam_beta1: Optional[float] = None
-    adam_beta2: Optional[float] = None
-    max_grad_norm: Optional[float] = None
-    num_epochs: int = Field(default=1)
-
-    @field_validator("batch_size")
-    @classmethod
-    def hint_batch_size_set(cls, batch_size):
-        if batch_size:
-            LOG.warning(
-                "%s\n%s",
-                "batch_size is not recommended. Please use gradient_accumulation_steps instead.",
-                "To calculate the equivalent gradient_accumulation_steps, divide batch_size / micro_batch_size / number of gpus.",
-            )
-        return batch_size
-
-    @field_validator("learning_rate")
-    @classmethod
-    def convert_learning_rate(cls, learning_rate):
-        if learning_rate and isinstance(learning_rate, str):
-            learning_rate = float(learning_rate)
-        return learning_rate
-
-
-class ModelOutputConfig(BaseModel):
-    """model save configuration subset"""
-
-    output_dir: str = Field(default="./model-out")
-    hub_model_id: Optional[str] = None
-    hub_strategy: Optional[str] = None
-    save_safetensors: Optional[bool] = None
-
-
-class MLFlowConfig(BaseModel):
-    """mlflow configuration subset"""
-
-    use_mlflow: Optional[bool] = None
-    mlflow_tracking_uri: Optional[str] = None
-    mlflow_experiment_name: Optional[str] = None
-    hf_mlflow_log_artifacts: Optional[bool] = None
-
-
-class LISAConfig(BaseModel):
-    """LISA options"""
-
-    lisa_n_layers: Optional[int] = Field(
-        default=None,
-        metadata={"help": "the number of activate layers in LISA"},
-    )
-    lisa_step_interval: Optional[int] = Field(
-        default=None,
-        metadata={"help": "how often to switch layers in LISA"},
-    )
-    lisa_layers_attribute: Optional[str] = Field(
-        default="model.layers",
-        metadata={"help": "path under the model to access the layers"},
-    )
-
-
-class WandbConfig(BaseModel):
-    """wandb configuration subset"""
-
-    use_wandb: Optional[bool] = None
-    wandb_name: Optional[str] = None
-    wandb_run_id: Optional[str] = None
-    wandb_mode: Optional[str] = None
-    wandb_project: Optional[str] = None
-    wandb_entity: Optional[str] = None
-    wandb_watch: Optional[str] = None
-    wandb_log_model: Optional[str] = None
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_wandb_run(cls, data):
-        if data.get("wandb_run_id") and not data.get("wandb_name"):
-            data["wandb_name"] = data.get("wandb_run_id")
-
-            LOG.warning(
-                "wandb_run_id sets the ID of the run. If you would like to set the name, please use wandb_name instead."
-            )
-
-        return data
-
-
-class GradioConfig(BaseModel):
-    """Gradio configuration subset"""
-
-    gradio_title: Optional[str] = None
-    gradio_share: Optional[bool] = None
-    gradio_server_name: Optional[str] = None
-    gradio_server_port: Optional[int] = None
-    gradio_max_new_tokens: Optional[int] = None
-    gradio_temperature: Optional[float] = None
-
-
-# pylint: disable=too-many-public-methods,too-many-ancestors
-class AxolotlInputConfig(
-    ModelInputConfig,
-    ModelOutputConfig,
-    LoraConfig,
-    ReLoRAConfig,
-    HyperparametersConfig,
-    WandbConfig,
-    MLFlowConfig,
-    LISAConfig,
-    GradioConfig,
-    RemappedParameters,
-    DeprecatedParameters,
-    BaseModel,
-):
-    """wrapper of all config options"""
-
-    class Config:
-        """Config for alias"""
-
-        populate_by_name = True
-
-    strict: Optional[bool] = Field(default=False)
-    resume_from_checkpoint: Optional[str] = None
-    auto_resume_from_checkpoints: Optional[bool] = None
-    resize_token_embeddings_to_32x: Optional[bool] = None
-
-    rl: Optional[RLType] = None
-
-    datasets: Optional[conlist(Union[SFTDataset, DPODataset], min_length=1)] = None  # type: ignore
-    test_datasets: Optional[conlist(Union[SFTDataset, DPODataset], min_length=1)] = None  # type: ignore
-    shuffle_merged_datasets: Optional[bool] = True
-    dataset_prepared_path: Optional[str] = None
-    dataset_shard_num: Optional[int] = None
-    dataset_shard_idx: Optional[int] = None
-
-    pretraining_dataset: Optional[  # type: ignore
-        conlist(Union[PretrainingDataset, SFTDataset], min_length=1)
-    ] = Field(
-        default=None, metadata={"help": {"streaming dataset to use for pretraining"}}
-    )
-    dataset_processes: Optional[int] = Field(default=os.cpu_count())
-    dataset_keep_in_memory: Optional[bool] = None
-    dataloader_pin_memory: Optional[bool] = None
-    dataloader_num_workers: Optional[int] = None
-    dataloader_prefetch_factor: Optional[int] = None
-    dataloader_drop_last: Optional[bool] = None
-
-    remove_unused_columns: Optional[bool] = None
-
-    push_dataset_to_hub: Optional[str] = None
-    hf_use_auth_token: Optional[bool] = None
-
-    device: Optional[Any] = None
-    device_map: Optional[Any] = None
-    world_size: Optional[int] = None
-    local_rank: Optional[int] = None
-    ddp: Optional[bool] = None
-
-    seed: Optional[int] = None
-    ddp_timeout: Optional[int] = None
-    ddp_bucket_cap_mb: Optional[int] = None
-    ddp_broadcast_buffers: Optional[bool] = None
-    ddp_find_unused_parameters: Optional[bool] = None
-
-    eval_table_size: Optional[int] = None
-    eval_max_new_tokens: Optional[int] = None
-    do_causal_lm_eval: Optional[bool] = None
-    eval_causal_lm_metrics: Optional[List[str]] = None
-    do_bench_eval: Optional[bool] = None
-    bench_dataset: Optional[str] = None
-    bench_split: Optional[str] = None
-    metric_for_best_model: Optional[str] = None
-    greater_is_better: Optional[bool] = None
-
-    loss_watchdog_threshold: Optional[float] = None
-    loss_watchdog_patience: Optional[int] = None
-
-    bf16: Optional[Union[Literal["auto"], bool]] = "auto"
-    fp16: Optional[bool] = None
-    bfloat16: Optional[bool] = None  # for non-AMP cases
-    float16: Optional[bool] = None  # for non-AMP cases
-    tf32: Optional[bool] = None
-    float32: Optional[bool] = None
-
-    # torch_dtype: Optional[torch.dtype]
-
-    gradient_checkpointing: Optional[Union[Literal["unsloth"], bool]] = Field(
-        default=False
-    )
-    gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None
-
-    unfrozen_parameters: Optional[List[str]] = None
-
-    sequence_len: int = Field(default=512)
-    min_sample_len: Optional[int] = None
-    max_prompt_len: int = Field(
-        default=512, metadata={"help": "maximum prompt length for RL training"}
-    )
-    sample_packing: Optional[bool] = None
-    eval_sample_packing: Optional[bool] = None
-    pad_to_sequence_len: Optional[bool] = None
-    curriculum_sampling: Optional[bool] = None
-
-    # for PoSE context length extension
-    use_pose: Optional[bool] = None
-    pose_split_on_token_ids: Optional[List[int]] = None
-    pose_max_context_len: Optional[int] = None
-    pose_num_chunks: Optional[int] = None
-
-    pretrain_multipack_buffer_size: Optional[int] = 10_000
-    pretrain_multipack_attn: Optional[bool] = Field(
-        default=True,
-        metadata={
-            "help": "whether to prevent cross attention for packed sequences during pretraining",
-        },
-    )
-
-    xformers_attention: Optional[bool] = None
-    sdp_attention: Optional[bool] = None
-    s2_attention: Optional[bool] = None
-    flash_attention: Optional[bool] = None
-    flash_attn_cross_entropy: Optional[bool] = None
-    flash_attn_rms_norm: Optional[bool] = None
-    flash_attn_fuse_qkv: Optional[bool] = None
-    flash_attn_fuse_mlp: Optional[bool] = None
-    flash_optimum: Optional[bool] = None
-
-    deepspeed: Optional[Union[str, Dict[str, Any]]] = None
-    fsdp: Optional[List[str]] = None
-    fsdp_config: Optional[Dict[str, Any]] = None
-
-    val_set_size: Optional[float] = Field(default=0.0)
-
-    special_tokens: Optional[SpecialTokensConfig] = None
-    tokens: Optional[List[str]] = None
-
-    torch_compile: Optional[bool] = None
-    torch_compile_backend: Optional[str] = None
-
-    max_steps: Optional[int] = None
-    warmup_steps: Optional[int] = None
-    warmup_ratio: Optional[float] = None
-    eval_steps: Optional[Union[int, float]] = None
-    evals_per_epoch: Optional[Union[int]] = None
-    evaluation_strategy: Optional[str] = None
-    save_steps: Optional[Union[int, float]] = None
-    saves_per_epoch: Optional[int] = None
-    save_strategy: Optional[str] = None
-    save_total_limit: Optional[int] = None
-    logging_steps: Optional[int] = None
-    early_stopping_patience: Optional[int] = None
-    load_best_model_at_end: Optional[bool] = False
-
-    neftune_noise_alpha: Optional[float] = None
-
-    orpo_alpha: Optional[float] = None
-
-    max_memory: Optional[
-        Dict[Union[int, Literal["cpu", "disk"]], Union[int, str]]
-    ] = None
-    gpu_memory_limit: Optional[Union[int, str]] = None
-    low_cpu_mem_usage: Optional[bool] = None
-
-    chat_template: Optional[ChatTemplate] = None
-    default_system_message: Optional[str] = None
-
-    # INTERNALS - document for now, generally not set externally
-    is_preprocess: Optional[bool] = None
-
-    total_num_tokens: Optional[int] = None
-    total_supervised_tokens: Optional[int] = None
-    sample_packing_eff_est: Optional[float] = None
-    axolotl_config_path: Optional[str] = None
-
-    is_falcon_derived_model: Optional[bool] = Field(default=None)
-    is_llama_derived_model: Optional[bool] = Field(default=None)
-    is_mistral_derived_model: Optional[bool] = Field(default=None)
-    is_qwen_derived_model: Optional[bool] = Field(default=None)
-
-    @field_validator("datasets", mode="before")
-    @classmethod
-    def fix_sharegpt_datasets(cls, datasets):
-        for idx, ds_cfg in enumerate(datasets):
-            if not ds_cfg["type"]:
-                continue
-            if ds_cfg["type"] == "sharegpt:chat":
-                LOG.warning(
-                    PendingDeprecationWarning(
-                        "`type: sharegpt:chat` will soon be deprecated. simply use `type: sharegpt` instead."
-                    )
-                )
-                datasets[idx]["type"] = "sharegpt"
-            if "sharegpt_simple" in ds_cfg["type"]:
-                LOG.warning(
-                    PendingDeprecationWarning(
-                        "`type: sharegpt_simple` will soon be deprecated. simply use `type: sharegpt` instead."
-                    )
-                )
-                datasets[idx]["type"] = datasets[idx]["type"].replace(
-                    "sharegpt_simple", "sharegpt"
-                )
-        return datasets
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_batch_size_fields(cls, data):
-        fields = ("micro_batch_size", "gradient_accumulation_steps", "batch_size")
-        non_empty_count = sum(1 for field in fields if data.get(field))
-
-        if non_empty_count < 2:
-            raise ValueError(f"At least two of {', '.join(fields)} must be set")
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_pretraining_w_max_steps(cls, data):
-        if data.get("pretraining_dataset") and not data.get("max_steps"):
-            raise ValueError(
-                "max_steps must be set when using iterable pretraining_dataset, Trainer can't infer length and schedule optimizer/learning rate without it!"
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_pretraining_w_group_by_length(cls, data):
-        if data.get("pretraining_dataset") and data.get("group_by_length"):
-            LOG.warning(
-                "You probably want to disable group_by_length as it will force a streamed dataset to download completely."
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_gptq_w_revision(cls, data):
-        if data.get("gptq") and data.get("revision_of_model"):
-            raise ValueError(
-                "revision_of_model is not supported for GPTQ models. "
-                + "Please download the model from HuggingFace Hub manually for correct branch, "
-                + "point to its path, and remove revision_of_model from the config."
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_sample_packing_w_xformers(cls, data):
-        if data.get("sample_packing") and data.get("xformers_attention"):
-            raise ValueError(
-                "sample_packing not compatible with xformers_attention. Use flash_attention"
-            )
-
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_sample_packing_wo_flash(cls, data):
-        if (
-            data.get("sample_packing")
-            and not data.get("flash_attention")
-            and not data.get("sdp_attention")
-        ):
-            LOG.warning(
-                "sample_packing without flash_attention or sdp_attention does not handle cross-attention."
-            )
-
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_sample_packing_w_rl(cls, data):
-        if data.get("sample_packing") and data.get("rl"):
-            raise ValueError("`sample_packing: true` does not work with RLHF training")
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def hint_sample_packing_padding(cls, data):
-        if data.get("sample_packing") and not data.get("pad_to_sequence_len"):
-            LOG.warning(
-                "`pad_to_sequence_len: true` is recommended when using sample_packing"
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_gas_bsz(cls, data):
-        if data.get("gradient_accumulation_steps") and data.get("batch_size"):
-            raise ValueError(
-                "please set only one of gradient_accumulation_steps or batch_size"
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def hint_eval_train_mbsz(cls, data):
-        if (
-            data.get("eval_batch_size")
-            and data.get("micro_batch_size")
-            and data.get("eval_batch_size") != data.get("micro_batch_size")
-        ):
-            LOG.warning(
-                "eval_batch_size != micro_batch_size. This can lead to VRAM instability."
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_push_ds_auth(cls, data):
-        if (
-            data.get("push_dataset_to_hub")
-            and data.get("hf_use_auth_token") is not True
-        ):
-            raise ValueError(
-                "Require cfg.hf_use_auth_token to be True for push_dataset_to_hub"
-            )
-        return data
-
-    @model_validator(mode="after")
-    def check_falcon_fsdp(self):
-        if (self.base_model and "falcon" in self.base_model.lower()) and self.fsdp:
-            raise ValueError("FSDP is not supported for falcon models")
-        return self
-
-    @model_validator(mode="after")
-    def check_mpt_checkpointing(self):
-        if (
-            self.base_model and "mpt" in self.base_model.lower()
-        ) and self.gradient_checkpointing:
-            raise ValueError("gradient_checkpointing is not supported for MPT models")
-        return self
-
-    @model_validator(mode="after")
-    def check_better_transformers(self):
-        if self.flash_optimum is True:
-            if self.adapter:
-                LOG.warning(
-                    "BetterTransformers probably doesn't work with PEFT adapters"
-                )
-            if self.fp16 or self.bf16:
-                raise ValueError("AMP is not supported with BetterTransformer")
-            if self.float16 is not True and self.bfloat16 is not True:
-                LOG.warning(
-                    "You should probably set bfloat16 or float16 to true to "
-                    "load the model in float16 for BetterTransformers"
-                )
-        return self
-
-    @model_validator(mode="after")
-    def check_adamw_optimizer_params(self):
-        if any([self.adam_beta1, self.adam_beta2, self.adam_epsilon]) and (
-            not self.optimizer or "adamw" not in self.optimizer.value
-        ):
-            LOG.warning("adamw hyperparameters found, but no adamw optimizer set")
-        return self
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_saves(cls, data):
-        if (
-            data.get("save_strategy")
-            and data.get("save_steps")
-            and data.get("save_strategy") != "steps"
-        ):
-            raise ValueError(
-                "save_strategy and save_steps mismatch. Please set save_strategy to 'steps' or remove save_steps."
-            )
-        if data.get("saves_per_epoch") and data.get("save_steps"):
-            raise ValueError(
-                "save_steps and saves_per_epoch are mutually exclusive and cannot be used together."
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_push_save(cls, data):
-        if data.get("hub_model_id") and (
-            data.get("save_strategy") not in ["steps", "epoch", None]
-        ):
-            LOG.warning(
-                "hub_model_id is set without any models being saved. To save a model, set save_strategy."
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_evals(cls, data):
-        if (
-            data.get("evaluation_strategy")
-            and data.get("eval_steps")
-            and data.get("evaluation_strategy") != "steps"
-        ):
-            raise ValueError(
-                "evaluation_strategy and eval_steps mismatch. Please set evaluation_strategy to 'steps' or remove eval_steps."
-            )
-
-        if (
-            data.get("val_set_size") == 0
-            and (data.get("eval_steps") or data.get("evaluation_strategy"))
-            and not data.get("test_datasets")
-        ):
-            raise ValueError(
-                "eval_steps and evaluation_strategy are not supported with val_set_size == 0"
-            )
-        if data.get("evals_per_epoch") and data.get("eval_steps"):
-            raise ValueError(
-                "eval_steps and evals_per_epoch are mutually exclusive and cannot be used together."
-            )
-        if (
-            data.get("evals_per_epoch")
-            and data.get("evaluation_strategy")
-            and data.get("evaluation_strategy") != "steps"
-        ):
-            raise ValueError(
-                "evaluation_strategy must be empty or set to `steps` when used with evals_per_epoch."
-            )
-
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_eval_packing(cls, data):
-        if (
-            data.get("sample_packing")
-            and data.get("eval_table_size")
-            and data.get("eval_sample_packing") is not False
-        ):
-            raise ValueError(
-                "eval_table_size and eval_sample_packing are not supported together with sample_packing. Please set 'eval_sample_packing' to false."
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_warmup(cls, data):
-        if data.get("warmup_steps") and data.get("warmup_ratio"):
-            raise ValueError("warmup_steps and warmup_ratio are mutually exclusive")
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_neftune(cls, data):
-        if data.get("noisy_embedding_alpha") and not data.get("neftune_noise_alpha"):
-            data["neftune_noise_alpha"] = data["noisy_embedding_alpha"]
-            del data["noisy_embedding_alpha"]
-        elif data.get("noisy_embedding_alpha") and not data.get("neftune_noise_alpha"):
-            raise ValueError(
-                "noisy_embedding_alpha is deprecated, use neftune_noise_alpha; both are set, please remove the deprecated noisy_embedding_alpha setting"
-            )
-        return data
-
-    @field_validator("neftune_noise_alpha")
-    @classmethod
-    def validate_neftune_noise_alpha(cls, neftune_noise_alpha):
-        if neftune_noise_alpha is not None and neftune_noise_alpha <= 0.0:
-            raise ValueError("neftune_noise_alpha must be > 0.0")
-        return neftune_noise_alpha
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_frozen(cls, data):
-        if (
-            data.get("adapter")
-            and data.get("peft_layers_to_transform")
-            and data.get("unfrozen_parameters")
-        ):
-            raise ValueError(
-                "`unfrozen_parameters` used with `peft_layers_to_transform` can have unexpected behavior."
-            )
-
-        return data
-
-    @model_validator(mode="after")
-    def check_fft_possible_bad_config(self):
-        if (
-            # pylint: disable=too-many-boolean-expressions
-            not (self.bf16 or self.bfloat16)
-            and (self.fp16 or self.float16)
-            and not self.adapter
-            and not self.flash_attention
-            and self.sample_packing
-        ):
-            LOG.warning(
-                "Full fine tune w/o FA2 w/ sample packing and fp16/float16 is likely to raise errors. Try LoRA."
-            )
-            # ValueError: Attempting to unscale FP16 gradients.
-            # OR
-            # RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::Half
-        return self
-
-    @model_validator(mode="after")
-    def check_fused_lora(self):
-        if self.adapter in ["lora", "qlora"] and (
-            self.flash_attn_fuse_qkv or self.flash_attn_fuse_mlp
-        ):
-            raise ValueError("Fused modules are not supported with LoRA/QLoRA")
-        return self
-
-    @model_validator(mode="after")
-    def hint_lora_8bit(self):
-        loftq = (
-            self.peft and self.peft.loftq_config and self.peft.loftq_config.loftq_bits
-        )
-        if not self.load_in_8bit and self.adapter == "lora" and not loftq:
-            LOG.warning("We recommend setting `load_in_8bit: true` for LORA finetuning")
-        return self
-
-    @model_validator(mode="after")
-    def check_early_stopping(self):
-        if self.early_stopping_patience:
-            if not self.save_steps or not self.eval_steps:
-                raise ValueError(
-                    "`early_stopping_patience` requires save_steps and eval_steps to be set. eval_steps should evenly divide save_steps."
-                )
-            if self.save_steps % self.eval_steps != 0:
-                raise ValueError(
-                    "`early_stopping_patience` requires that eval_steps should evenly divide save_steps."
-                )
-        return self
-
-    @model_validator(mode="after")
-    def check_relora(self):
-        if self.relora_steps:
-            if self.adapter not in ("lora", "qlora"):
-                raise ValueError("cfg.adapter must be lora or qlora to use ReLoRA")
-
-            if self.fsdp:
-                raise ValueError("fsdp not supported with ReLoRA")
-
-            if self.deepspeed:
-                raise ValueError("deepspeed not supported with ReLoRA")
-
-            if self.lr_scheduler == "one_cycle":
-                raise ValueError(
-                    "ReLoRA is not compatible with the one_cycle scheduler"
-                )
-
-            if self.flash_attn_fuse_qkv or self.flash_attn_fuse_mlp:
-                raise ValueError("Fused modules are not supported with ReLoRA")
-        return self
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_mem_mismatch(cls, data):
-        if (
-            data.get("max_memory") is not None
-            and data.get("gpu_memory_limit") is not None
-        ):
-            raise ValueError(
-                "max_memory and gpu_memory_limit are mutually exclusive and cannot be used together."
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_use_reentrant_mismatch(cls, data):
-        if (
-            data.get("unfrozen_parameters")
-            and data.get("gradient_checkpointing_kwargs")
-            and data.get("gradient_checkpointing_kwargs", {}).get("use_reentrant")
-            is True
-        ):
-            # https://github.com/huggingface/transformers/issues/21381
-            raise ValueError(
-                "`use_reentrant` must be false when used with partially frozen model."
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_val_w_test_datasets(cls, data):
-        if data.get("test_datasets") and data.get("val_set_size"):
-            raise ValueError(
-                "non-zero val_set_size should not be used with test_datasets configuration"
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_fsdp_offload_w_8bit_optimizer(cls, data):
-        if (
-            data.get("fsdp")
-            and "8bit" in data.get("optimizer", "")
-            and data.get("fsdp_config")
-            and data["fsdp_config"].get("fsdp_offload_params")
-        ):
-            raise ValueError(
-                f"FSDP Offload not compatible with {data.get('optimizer')}"
-            )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_causal_lm_evals(cls, data):
-        if data.get("do_causal_lm_eval") and data.get("eval_sample_packing"):
-            raise ValueError(
-                "do_causal_lm_eval is enabled, eval_sample_packing must be set to False"
-            )
-
-        if data.get("eval_causal_lm_metrics"):
-            supported_metrics = ["sacrebleu", "comet", "ter", "chrf"]
-            if not isinstance(data.get("eval_causal_lm_metrics"), list):
-                raise ValueError("eval_causal_lm_metrics must be a list")
-            # only ["sacrebleu", "comet", "ter", "chrf"] supported
-            if set(data.get("eval_causal_lm_metrics")) - set(supported_metrics):
-                raise ValueError(
-                    f"eval_causal_lm_metrics must be one of {supported_metrics}"
-                )
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_dataset_or_pretraining_dataset(cls, data):
-        if data.get("datasets") is None and data.get("pretraining_dataset") is None:
-            raise ValueError("either datasets or pretraining_dataset is required")
-        return data
-
-
-class AxolotlConfigWCapabilities(AxolotlInputConfig):
-    """wrapper to valdiate gpu capabilities with the configured options"""
-
-    capabilities: GPUCapabilities
-
-    @model_validator(mode="after")
-    def check_bf16(self):
-        if self.capabilities.bf16:
-            if not self.bf16 and not self.bfloat16:
-                LOG.info(
-                    "bf16 support detected, but not enabled for this configuration."
-                )
-        else:
-            if (
-                not self.merge_lora
-                and not self.is_preprocess
-                and (self.bf16 is True or self.bfloat16 is True)
-            ):
-                raise ValueError(
-                    "bf16 requested, but AMP is not supported on this GPU. Requires Ampere series or above."
-                )
-        return self
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_sample_packing_w_sdpa_bf16(cls, data):
-        is_sm_90: bool = (
-            data["capabilities"]
-            and data["capabilities"].get("compute_capability") == "sm_90"
-        )
-        if (
-            data.get("sample_packing")
-            and data.get("sdp_attention")
-            and (data.get("bfloat16") or data.get("bf16"))
-            and not is_sm_90
-        ):
-            # https://github.com/pytorch/pytorch/blob/1b03423526536b5f3d35bdfa95ccc6197556cf9b/test/test_transformers.py#L2440-L2450
-            LOG.warning(
-                "sample_packing & torch sdpa with bf16 is unsupported may results in 0.0 loss. "
-                "This may work on H100s."
-            )
-
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_fsdp_deepspeed(cls, data):
-        if data.get("deepspeed") and data.get("fsdp"):
-            raise ValueError("deepspeed and fsdp cannot be used together.")
-        return data
diff --git a/src/axolotl/utils/ctx_managers/__init__.py b/src/axolotl/utils/ctx_managers/__init__.py
new file mode 100644
index 0000000000..6ffda9e557
--- /dev/null
+++ b/src/axolotl/utils/ctx_managers/__init__.py
@@ -0,0 +1,5 @@
+"""Init for context manager submodule"""
+
+# flake8: noqa
+
+from .sequence_parallel import SequenceParallelContextManager
diff --git a/src/axolotl/utils/ctx_managers/sequence_parallel.py b/src/axolotl/utils/ctx_managers/sequence_parallel.py
new file mode 100644
index 0000000000..7f6af7d480
--- /dev/null
+++ b/src/axolotl/utils/ctx_managers/sequence_parallel.py
@@ -0,0 +1,444 @@
+"""Module for Axolotl trainer sequence parallelism manager and utilities"""
+
+import functools
+import inspect
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed import DeviceMesh
+from torch.utils.hooks import RemovableHandle
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.utils import ModelOutput
+
+from axolotl.monkeypatch.ring_attn import (
+    get_ring_attn_group,
+    register_ring_attn_from_device_mesh,
+    update_ring_attn_params,
+)
+from axolotl.utils.schemas.enums import RingAttnFunc
+
+
+# TODO(djsaunde): implement zigzag, stripe patterns here (and elsewhere) in this
+# module. Currently, we just focus on batch ring and varlen llama3 for simplicity.
+def apply_sequence_parallelism(
+    batch: dict[str, torch.Tensor],
+    local_rank: int,
+    local_world_size: int,
+    gradient_accumulation_steps: int,
+    ring_attn_func: RingAttnFunc,
+) -> tuple[dict[str, torch.Tensor], int, int]:
+    """
+    Apply sequence parallelism slicing to a batch.
+
+    Special handling is implemented for integer logits_to_keep, which indicates
+    to only keep the last N tokens in the sequence during generation.
+
+    Args:
+        batch: Batch dictionary (e.g., input_ids, attention_mask, etc.).
+        local_rank: Local rank in the sequence parallel group.
+        local_world_size: World size of the sequence parallel group.
+        gradient_accumulation_steps: Number of steps to accumulate gradients over.
+        ring_attn_func: Which ring attention function to use. Currently unused, but
+            related to above TODO.
+
+    Returns:
+        tuple of:
+            - Batch dictionary with sliced tensors.
+            - The original sequence length before padding.
+            - The number of padding tokens added.
+    """
+    batch_size, original_seq_len = batch["input_ids"].shape
+
+    # Update ring attention params if needed
+    if batch.get("position_ids") is not None and batch_size == 1:
+        update_ring_attn_params(position_ids=batch["position_ids"])
+    else:
+        # If position_ids aren't already in the batch, create them
+        batch["position_ids"] = torch.arange(
+            0,
+            original_seq_len,
+            dtype=torch.long,
+            device=batch["input_ids"].device,
+        ).expand(batch["input_ids"].size(0), -1)
+
+    if "logits_to_keep" in batch and isinstance(batch["logits_to_keep"], int):
+        logits_to_keep = batch["logits_to_keep"]
+
+        # Calculate which positions in the full sequence contain the last N tokens
+        start_position = max(0, original_seq_len - logits_to_keep)
+        chunk_size = original_seq_len // local_world_size
+        rank_start = local_rank * chunk_size
+        rank_end = rank_start + chunk_size
+
+        # Create a boolean mask tensor for this rank's chunk
+        mask = torch.zeros(
+            chunk_size,
+            dtype=torch.bool,
+            device=batch["input_ids"].device,
+        )
+
+        if rank_end > start_position:
+            # Calculate how many of the last N tokens fall within this rank's range
+            tokens_in_rank = min(rank_end, original_seq_len) - max(
+                rank_start, start_position
+            )
+
+            # Calculate where these tokens start in the local chunk
+            local_start_idx = max(0, start_position - rank_start)
+
+            # Set the appropriate positions in the mask to True
+            mask[local_start_idx : local_start_idx + tokens_in_rank] = True
+
+        # Replace the integer with the boolean mask
+        batch["logits_to_keep"] = mask
+
+    # Add padding to make sequence length divisible by local_world_size
+    total_seq_len = original_seq_len
+    pad_len = 0
+    divisor = min(local_world_size, 64)
+    if total_seq_len % divisor != 0:
+        pad_len = divisor - (total_seq_len % divisor)
+
+        # Apply padding to all relevant tensors
+        for key in batch:
+            if (
+                isinstance(batch[key], torch.Tensor)
+                and batch[key].dim() > 1
+                and batch[key].size(1) == total_seq_len
+            ):
+                # Create padding tensor
+                pad_value = -100 if key == "labels" else 0
+                padding = torch.full(
+                    (batch[key].size(0), pad_len, *batch[key].shape[2:]),
+                    pad_value,
+                    dtype=batch[key].dtype,
+                    device=batch[key].device,
+                )
+
+                # Concatenate padding to the right side of the tensor
+                batch[key] = torch.cat([batch[key], padding], dim=1)
+            if key == "logits_to_keep":
+                # Create padding tensor
+                padding = torch.ones(
+                    1,
+                    dtype=batch[key].dtype,
+                    device=batch[key].device,
+                )
+
+                # Concatenate padding to the right side of the tensor
+                batch[key] = torch.cat([batch[key], padding], dim=0)
+
+        # Update the total sequence length after padding
+        total_seq_len = batch["input_ids"].size(1)
+
+    # Slice batch for sequence parallel
+    for key in batch:
+        if not isinstance(batch[key], torch.Tensor) or batch[key].dim() <= 1:
+            continue
+
+        # Split in sequential fashion and grab this rank's chunk
+        if batch[key].size(1) == total_seq_len:
+            batch[key] = (
+                batch[key].chunk(local_world_size, dim=1)[local_rank].contiguous()
+            )
+        elif key == "logits_to_keep":
+            batch[key] = (
+                batch[key].chunk(local_world_size, dim=0)[local_rank].contiguous()
+            )
+
+        # Handle num_items_in_batch
+        if "num_items_in_batch" in batch:
+            # Approximation; this needed since num_items_in_batch may be counted across
+            # all samples in a gradient accumulated batch, not on a per-step basis.
+            local_valid_tokens = (batch["labels"] != -100).sum()
+
+            # All-reduce across sequence parallel ranks to get global token count
+            cp_group = get_ring_attn_group()
+            global_valid_tokens = local_valid_tokens.clone()
+            # we use AVG instead of SUM as using sum seems to scale down the loss by over-accounting the number of tokens
+            dist.all_reduce(global_valid_tokens, op=dist.ReduceOp.AVG, group=cp_group)
+            global_valid_tokens = int(global_valid_tokens.item())
+
+            batch["num_items_in_batch"] = (
+                global_valid_tokens * gradient_accumulation_steps
+            )
+
+    return batch, original_seq_len, pad_len
+
+
+class SequenceParallelContextManager:
+    """Context manager for sequence parallelism operations.
+
+    This class provides a context that will automatically apply sequence parallelism
+    during model forward passes using a pre-forward hook, and gather outputs from
+    across the sequence parallelism group using a post-forward hook.
+
+    Args:
+        models: List of models to apply sequence parallelism to pre- and post- forward
+            hooks.
+        context_parallel_size: Number of processes to split sequences over.
+        gradient_accumulation_steps: Number of steps to accumulate gradients over.
+        ring_attn_func: Which ring attention function to use. Currently unused.
+        heads_k_stride: Sequence parallelism K head stride size. Passed through to
+            `varlen_llama3` `ring_flash_attn` implementation.
+        gather_outputs: Whether to gather outputs after model forward pass across the
+            sequence parallel group.
+    """
+
+    def __init__(
+        self,
+        models: list[nn.Module],
+        context_parallel_size: int,
+        gradient_accumulation_steps: int,
+        ring_attn_func: RingAttnFunc,
+        heads_k_stride: int | None,
+        gather_outputs: bool,
+        device_mesh: DeviceMesh | None = None,
+    ):
+        self.models = models
+        self.context_parallel_size = context_parallel_size
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.ring_attn_func = ring_attn_func
+        self.heads_k_stride = heads_k_stride
+        self.gather_outputs = gather_outputs
+        self.device_mesh = device_mesh
+
+        self._register_ring_attn()
+
+        # Set distributed info for local rank
+        self.process_group = get_ring_attn_group()
+        self.local_rank = dist.get_rank(self.process_group)
+        self.local_world_size = dist.get_world_size(self.process_group)
+
+        # Will store hook handles for removal
+        self.hook_handles: list[RemovableHandle] = []
+
+        # Store original sequence length and padding information
+        self.original_seq_len = 0
+        self.pad_len = 0
+
+        # Track local valid token count for eval loss correction across CP ranks
+        self._local_valid_tokens: torch.Tensor | None = None
+
+        # Create a partially applied version of the apply_sequence_parallelism function
+        self.apply_sequence_parallelism = functools.partial(
+            apply_sequence_parallelism,
+            local_rank=self.local_rank,
+            local_world_size=self.local_world_size,
+            gradient_accumulation_steps=self.gradient_accumulation_steps,
+            ring_attn_func=self.ring_attn_func,
+        )
+
+    def __enter__(self):
+        self._register_model_hooks()
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Remove all hooks
+        for handle in self.hook_handles:
+            handle.remove()
+        self.hook_handles = []
+
+        # TODO(djsaunde): Un-patch attention and accelerate functions (low priority)
+
+    def _register_ring_attn(self):
+        # Initialize ring attn for sequence parallelism
+        register_ring_attn_from_device_mesh(
+            device_mesh=self.device_mesh,
+            context_parallel_dim=("cp",),
+            heads_k_stride=self.heads_k_stride,
+            ring_attn_func=self.ring_attn_func,
+        )
+
+    def _register_model_hooks(self):
+        # Forward pre-hook to apply sequence parallelism
+        def sequence_parallel_pre_hook(_, args, kwargs):
+            # Get parameter names from the model's forward function
+            forward_params = list(
+                inspect.signature(self.models[0].forward).parameters.keys()
+            )
+
+            updated_kwargs = kwargs.copy()
+            for i, arg in enumerate(args):
+                if i < len(forward_params):
+                    updated_kwargs[forward_params[i]] = arg
+
+            # Any excess positional arguments are kept as-is
+            remaining_args = args[len(forward_params) :]
+
+            # Apply sequence parallelism to updated kwargs
+            updated_kwargs, self.original_seq_len, self.pad_len = (
+                self.apply_sequence_parallelism(updated_kwargs)
+            )
+
+            # Track local valid tokens for eval loss correction
+            if "labels" in updated_kwargs and not self.models[0].training:
+                self._local_valid_tokens = (
+                    (updated_kwargs["labels"] != -100).sum().float()
+                )
+                # Strip num_items_in_batch during eval so the model uses
+                # reduction='mean', allowing the post-hook weighted all-reduce
+                # formula (loss * local_valid) to correctly recover the loss sum
+                updated_kwargs.pop("num_items_in_batch", None)
+            else:
+                self._local_valid_tokens = None
+
+            return remaining_args, updated_kwargs
+
+        # Forward post-hook to gather outputs
+        def sequence_parallel_post_hook(_, __, output: ModelOutput) -> ModelOutput:
+            # Gather the sharded outputs
+            output = self._gather_outputs(output)
+
+            # Remove padding if it was added
+            if self.pad_len > 0:
+                for key, value in output.items():
+                    if isinstance(value, torch.Tensor) and value.dim() > 1:
+                        if value.size(1) == self.original_seq_len + self.pad_len:
+                            # Slice to remove padding
+                            output[key] = value[:, : self.original_seq_len].contiguous()
+
+            return output
+
+        # Post-hook to correct eval loss via weighted all-reduce across CP ranks
+        def eval_loss_correction_post_hook(_, __, output: ModelOutput) -> ModelOutput:
+            if self._local_valid_tokens is None:
+                return output
+            if not hasattr(output, "loss") or output.loss is None:
+                return output
+
+            local_valid = self._local_valid_tokens.to(output.loss.device)
+            loss = output.loss.detach().clone()
+
+            # Handle rank with zero valid tokens (loss is NaN)
+            if local_valid.item() == 0:
+                weighted_loss = torch.zeros(1, device=loss.device, dtype=loss.dtype)
+            else:
+                weighted_loss = loss * local_valid
+
+            total_valid = local_valid.clone()
+            dist.all_reduce(
+                weighted_loss,
+                op=dist.ReduceOp.SUM,
+                group=self.process_group,
+            )
+            dist.all_reduce(
+                total_valid,
+                op=dist.ReduceOp.SUM,
+                group=self.process_group,
+            )
+
+            if total_valid.item() > 0:
+                output["loss"] = (weighted_loss / total_valid).squeeze()
+            else:
+                output["loss"] = torch.tensor(
+                    float("nan"), device=loss.device, dtype=loss.dtype
+                )
+
+            self._local_valid_tokens = None
+            return output
+
+        # Register hooks
+        for model in self.models:
+            self.hook_handles.append(
+                model.register_forward_pre_hook(
+                    sequence_parallel_pre_hook, with_kwargs=True
+                )
+            )
+            if self.gather_outputs:
+                self.hook_handles.append(
+                    model.register_forward_hook(sequence_parallel_post_hook)
+                )
+            # Always register eval loss correction hook
+            self.hook_handles.append(
+                model.register_forward_hook(eval_loss_correction_post_hook)
+            )
+
+    def _gather_outputs(self, output: CausalLMOutputWithPast) -> CausalLMOutputWithPast:
+        """Gather sharded outputs from all ranks and reconstruct the full tensor."""
+        for key, value in output.items():
+            if isinstance(value, torch.Tensor) and value.dim() > 1:
+                output[key] = AllGatherWithGrad.apply(value, self.process_group)
+
+        return output
+
+
+class AllGatherWithGrad(torch.autograd.Function):
+    """Custom autograd function for all-gather to preserve gradients."""
+
+    @staticmethod
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        input_tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+    ) -> torch.Tensor:
+        """
+        Forward pass of all-gather of data with sequence dimension.
+
+        Args:
+            ctx: `torch.autograd` function context.
+            input_tensor: Tensor from model output with sequence dimension.
+            group: `torch.distributed` process group.
+
+        Returns:
+            Tensor from gathering the `input_tensor` from across the process group and
+                concatenating along the sequence dimension.
+        """
+        ctx.group = group
+        ctx.rank = dist.get_rank(group)
+        world_size = dist.get_world_size(group)
+
+        # Gather shape metadata
+        local_shape = torch.tensor(list(input_tensor.shape), device=input_tensor.device)
+        all_shapes = [torch.zeros_like(local_shape) for _ in range(world_size)]
+        dist.all_gather(all_shapes, local_shape, group=group)
+
+        # Store sequence lengths for backward pass
+        seq_lens = [int(shape[1].item()) for shape in all_shapes]
+        ctx.seq_lens = seq_lens
+
+        # Perform all_gather operation
+        gathered = [
+            torch.zeros(
+                tuple(shape.tolist()),
+                dtype=input_tensor.dtype,
+                device=input_tensor.device,
+            )
+            for shape in all_shapes
+        ]
+        dist.all_gather(gathered, input_tensor, group=group)
+
+        # Concatenate tensors along sequence dimension
+        result = torch.cat(gathered, dim=1)
+
+        return result
+
+    @staticmethod
+    def backward(
+        ctx: torch.autograd.function.FunctionCtx, grad_output: torch.Tensor
+    ) -> tuple[torch.Tensor, None]:
+        """
+        Backward pass for all-gather operation.
+
+        Extracts the gradient slice corresponding to this rank's original input
+        from the full gradient tensor.
+
+        Args:
+            ctx: `torch.autograd` function context.
+            grad_output: Gradient from subsequent layers with respect to the
+                concatenated output tensor.
+
+        Returns:
+            Tuple containing the gradient slice for this rank's input tensor and `None`
+                for the process group parameter which doesn't require gradients.
+        """
+        rank = ctx.rank
+        seq_lens = ctx.seq_lens
+
+        # Extract gradient for this rank's chunk
+        offset = sum(seq_lens[:rank])
+        grad_slice = grad_output[:, offset : offset + seq_lens[rank]].contiguous()
+
+        return grad_slice, None
diff --git a/src/axolotl/utils/cuda13.py b/src/axolotl/utils/cuda13.py
new file mode 100644
index 0000000000..cfc4a26872
--- /dev/null
+++ b/src/axolotl/utils/cuda13.py
@@ -0,0 +1,39 @@
+"""Helpers for CUDA 13 uv images."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def cu13_library_path() -> str | None:
+    """Return the nvidia.cu13 package lib path, when available."""
+    try:
+        import nvidia.cu13 as cu13
+    except ImportError:
+        return None
+
+    package_paths = list(cu13.__path__)
+    if not package_paths:
+        return None
+    return str(Path(package_paths[0]) / "lib")
+
+
+def prepend_cu13_ld_library_path(ld_library_path: str | None) -> str:
+    """
+    Prepend the CUDA 13 library directory when nvidia.cu13 is installed.
+
+    The path is intentionally kept small and idempotent:
+    - no-op when the cu13 package is absent
+    - keeps any existing LD_LIBRARY_PATH entries
+    - avoids duplicating the package-derived cu13 lib path
+    """
+    cu13_lib = cu13_library_path()
+    if cu13_lib is None:
+        return ld_library_path or ""
+
+    parts = [cu13_lib]
+    for part in (ld_library_path or "").split(":"):
+        if part and part != cu13_lib and part not in parts:
+            parts.append(part)
+
+    return ":".join(parts)
diff --git a/src/axolotl/utils/data/__init__.py b/src/axolotl/utils/data/__init__.py
index 140d02106d..8b9e4e91d4 100644
--- a/src/axolotl/utils/data/__init__.py
+++ b/src/axolotl/utils/data/__init__.py
@@ -1,15 +1,21 @@
-"""
-Data processing modules
-"""
-from axolotl.utils.data.pretraining import (  # noqa: F401
-    encode_pretraining,
-    wrap_pretraining_dataset,
-)
-from axolotl.utils.data.rl import load_prepare_dpo_datasets  # noqa: F401
-from axolotl.utils.data.sft import (  # noqa: F401
+"""Init for `axolotl.utils.data` module."""
+
+from axolotl.utils.data.rl import prepare_preference_datasets
+from axolotl.utils.data.sft import (
     get_dataset_wrapper,
-    load_prepare_datasets,
-    load_tokenized_prepared_datasets,
-    prepare_dataset,
+    prepare_datasets,
+)
+from axolotl.utils.data.streaming import (
+    encode_streaming,
+    wrap_streaming_dataset,
 )
-from axolotl.utils.data.utils import md5  # noqa: F401
+from axolotl.utils.data.utils import md5
+
+__all__ = [
+    "encode_streaming",
+    "wrap_streaming_dataset",
+    "prepare_preference_datasets",
+    "get_dataset_wrapper",
+    "prepare_datasets",
+    "md5",
+]
diff --git a/src/axolotl/utils/data/lock.py b/src/axolotl/utils/data/lock.py
new file mode 100644
index 0000000000..9699f60e80
--- /dev/null
+++ b/src/axolotl/utils/data/lock.py
@@ -0,0 +1,72 @@
+"""Logic for loading / preparing a dataset once over all processes."""
+
+import time
+from pathlib import Path
+from typing import Any, Callable
+
+from filelock import FileLock
+
+from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
+from axolotl.utils.dict import DictDefault
+
+LOCK_FILE_NAME = "datasets_prep.lock"
+READY_FILE_NAME = "datasets_ready.flag"
+PROCESS_COUNTER_FILE_NAME = "process_counter.txt"
+
+
+class FileLockLoader:
+    """
+    Simple class for abstracting single process data loading / processing. The first
+    process that creates a lock file does the work; the remaining procesees simply load
+    the preprocessed dataset once the first process is done.
+    """
+
+    def __init__(self, cfg: DictDefault):
+        self.cfg = cfg
+        self.dataset_prepared_path = (
+            cfg.dataset_prepared_path or DEFAULT_DATASET_PREPARED_PATH
+        )
+        self.lock_file_path = Path(self.dataset_prepared_path) / LOCK_FILE_NAME
+        self.ready_flag_path = Path(self.dataset_prepared_path) / READY_FILE_NAME
+        self.counter_path = Path(self.dataset_prepared_path) / PROCESS_COUNTER_FILE_NAME
+
+    def load(self, load_fn: Callable[[], Any]) -> Any:
+        with FileLock(str(self.lock_file_path)):
+            self._increment_counter()
+
+            if not self.ready_flag_path.exists():
+                result = load_fn()
+                self.ready_flag_path.touch()
+                return result
+
+            while not self.ready_flag_path.exists():
+                time.sleep(1)
+            return load_fn()
+
+    def _increment_counter(self):
+        """Safely increment the process counter."""
+        if self.counter_path.exists():
+            counter_content = self.counter_path.read_text().strip()
+            count = int(counter_content) if counter_content else 0
+        else:
+            count = 0
+        self.counter_path.write_text(str(count + 1))
+
+    def cleanup(self):
+        """Clean up ready flag when last process is done."""
+        try:
+            with FileLock(str(self.lock_file_path)):
+                counter_content = self.counter_path.read_text().strip()
+                count = int(counter_content) if counter_content else 0
+                count -= 1
+
+                if count <= 0:
+                    # Last process cleans everything up
+                    self.ready_flag_path.unlink(missing_ok=True)
+                    self.counter_path.unlink(missing_ok=True)
+                else:
+                    # Still have active processes
+                    self.counter_path.write_text(str(count))
+        except FileNotFoundError:
+            # Lock file might have already been deleted by another process
+            pass
diff --git a/src/axolotl/utils/data/rl.py b/src/axolotl/utils/data/rl.py
index ff5ca87ddf..f97f09e497 100644
--- a/src/axolotl/utils/data/rl.py
+++ b/src/axolotl/utils/data/rl.py
@@ -1,130 +1,489 @@
-"""data handling specific to DPO"""
+"""Data handling specific to RL trainers."""
+
 import inspect
-import logging
 from functools import partial
-from pathlib import Path
-from typing import Any, List
+from typing import Any, Callable, Literal
 
-import yaml
-from datasets import DatasetDict, concatenate_datasets, load_dataset, load_from_disk
+from datasets import Dataset, DatasetDict
+from transformers import PreTrainedTokenizer
 
-from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
+from axolotl.loaders import load_tokenizer
 from axolotl.prompt_strategies.dpo import load as load_dpo
+from axolotl.prompt_strategies.ebft import load as load_ebft
+from axolotl.prompt_strategies.kto import load as load_kto
 from axolotl.prompt_strategies.orpo import load as load_orpo
-from axolotl.utils.data.utils import md5
+from axolotl.utils.data.lock import FileLockLoader
+from axolotl.utils.data.shared import (
+    create_train_validation_split,
+    datasets_with_name_generator,
+    generate_dataset_hash_from_config,
+    load_dataset_with_config,
+    load_preprocessed_dataset,
+    merge_datasets,
+    save_preprocessed_dataset,
+    try_load_from_hub,
+)
+from axolotl.utils.data.utils import (
+    deduplicate_and_log_datasets,
+    retry_on_request_exceptions,
+)
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.distributed import is_main_process, zero_first
-from axolotl.utils.models import load_tokenizer
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.enums import RLType
 
-LOG = logging.getLogger("axolotl")
+LOG = get_logger(__name__)
 
 
-def _get_path(ds_hash, cfg):
-    prepared_ds_path = (
-        Path(cfg.dataset_prepared_path) / ds_hash
-        if cfg.dataset_prepared_path
-        else Path(DEFAULT_DATASET_PREPARED_PATH) / ds_hash
-    )
+@retry_on_request_exceptions(max_retries=3, delay=5)
+def prepare_preference_datasets(
+    cfg: DictDefault, tokenizer: PreTrainedTokenizer
+) -> tuple[Dataset, Dataset | None]:
+    """Load and prepare preference datasets for RL training.
 
-    return prepared_ds_path
+    Loads training and evaluation datasets, handling preprocessing, caching, and
+    deduplication as configured. Uses FileLock for distributed coordination.
 
+    Args:
+        cfg: Configuration object containing dataset and training settings.
+        tokenizer: Tokenizer to use for processing text.
+
+    Returns:
+        Tuple of (train_dataset, eval_dataset). eval_dataset may be None
+            if no evaluation dataset is configured.
+    """
+
+    def _load_datasets():
+        # Load training dataset
+        train_dataset = _load_or_create_dataset_split(cfg, tokenizer, split="train")
+
+        # Load or create evaluation dataset
+        eval_dataset: Dataset | None = None
+        if cfg.test_datasets:
+            eval_dataset = _load_or_create_dataset_split(cfg, tokenizer, split="test")
+        elif cfg.val_set_size:
+            # Create validation split from training data
+            train_dataset, eval_dataset = create_train_validation_split(
+                train_dataset, cfg, cfg.val_set_size
+            )
+
+        return train_dataset, eval_dataset
+
+    # Prepare datasets (with file locking logic for multiple ranks)
+    loader = FileLockLoader(cfg)
+    try:
+        train_dataset, eval_dataset = loader.load(_load_datasets)
+    finally:
+        loader.cleanup()
+
+    # Apply deduplication if configured
+    if cfg.dataset_exact_deduplication:
+        train_dataset, eval_dataset = deduplicate_and_log_datasets(
+            dataset=train_dataset, other_dataset=eval_dataset
+        )
+
+    return train_dataset, eval_dataset
 
-def _load_preprocessed_ds(cfg, sub_cfg):
-    ds_hash = md5(yaml.dump(sub_cfg, Dumper=yaml.Dumper))
-    prepared_ds_path = _get_path(ds_hash, cfg)
-    dataset = None
 
-    # pylint: disable=duplicate-code
-    if (
-        cfg.dataset_prepared_path
-        and any(prepared_ds_path.glob("*"))
-        and not cfg.is_preprocess
-    ):
-        LOG.info(f"Loading prepared dataset from disk at {prepared_ds_path}...")
-        dataset = load_from_disk(str(prepared_ds_path))
+def _map_dataset(
+    cfg: DictDefault,
+    dataset: Dataset | DatasetDict,
+    ds_transform_fn: Callable[..., Any],
+    tokenizer: Any | None = None,
+    **map_kwargs: Any,
+) -> Dataset:
+    """Apply transformation function to dataset.
+
+    Args:
+        cfg: Configuration object.
+        dataset: Dataset to transform.
+        ds_transform_fn: Transformation function to apply.
+        tokenizer: Optional tokenizer for transformation.
+        **map_kwargs: Additional arguments for dataset mapping.
+
+    Returns:
+        Transformed dataset.
+    """
+    sig = inspect.signature(ds_transform_fn)
+    if "tokenizer" in sig.parameters:
+        if not tokenizer:
+            tokenizer = load_tokenizer(cfg)
+        ds_transform_fn = partial(ds_transform_fn, tokenizer=tokenizer)
+
+    if isinstance(dataset, DatasetDict):
+        dataset = dataset["train"]
+
+    dataset = dataset.map(
+        ds_transform_fn,
+        num_proc=cfg.dataset_num_proc,
+        load_from_cache_file=not cfg.is_preprocess,
+        desc="Mapping RL Dataset",
+        **map_kwargs,
+    )
 
     return dataset
 
 
-def _save_preprocessed_ds(cfg, sub_cfg, dataset):
-    ds_hash = md5(yaml.dump(sub_cfg, Dumper=yaml.Dumper))
-    prepared_ds_path = _get_path(ds_hash, cfg)
+def _drop_long_sequences(
+    sample: dict[str, Any], rl: RLType, tokenizer: Any, sequence_len: int
+) -> bool:
+    """Filter out samples that exceed maximum sequence length.
 
-    if cfg.is_preprocess and is_main_process():
-        LOG.info(f"Loading prepared dataset from disk at {prepared_ds_path}...")
-        dataset.save_to_disk(str(prepared_ds_path))
+    Args:
+        sample: Dataset sample to check.
+        rl: Reinforcement learning type.
+        tokenizer: Tokenizer for length calculation.
+        sequence_len: Maximum allowed sequence length.
 
+    Returns:
+        True if sample should be kept, False if it should be dropped.
 
-def load_prepare_dpo_datasets(cfg):
-    def load_split(dataset_cfgs, _cfg):
-        split_datasets: List[Any] = []
-        for i, ds_cfg in enumerate(dataset_cfgs):
-            if ds_cfg["ds_type"] == "json":
-                for data_file in ds_cfg["data_files"]:
-                    data_files = {ds_cfg["split"]: data_file}
-                    ds = load_dataset(  # pylint: disable=invalid-name
-                        "json",
-                        data_files=data_files,
-                        split=ds_cfg["split"],
-                    )
-                    split_datasets.insert(i, ds)
+    Raises:
+        ValueError: If required keys are missing or RL type is unknown.
+    """
+    if rl in {RLType.DPO, RLType.IPO, RLType.ORPO, RLType.SIMPO}:
+        if not (
+            sample.get("prompt") and sample.get("chosen") and sample.get("rejected")
+        ):
+            raise ValueError(
+                "Prompt, chosen and rejected keys are required for DPO/ORPO datasets"
+            )
+
+        prompt = sample["prompt"]
+        chosen = sample["chosen"]
+        rejected = sample["rejected"]
+
+        len_prompt = len(tokenizer(prompt, add_special_tokens=False)["input_ids"])
+        len_chosen = len(tokenizer(chosen, add_special_tokens=False)["input_ids"])
+        len_rejected = len(tokenizer(rejected, add_special_tokens=False)["input_ids"])
+
+        return (len_prompt + len_chosen) <= sequence_len and (
+            len_prompt + len_rejected
+        ) <= sequence_len
+
+    if rl is RLType.KTO:
+        if not (sample.get("prompt") and sample.get("completion")):
+            raise ValueError("Prompt and completion keys are required for KTO datasets")
+
+        prompt = sample["prompt"]
+        completion = sample["completion"]
+
+        len_prompt = len(tokenizer(prompt, add_special_tokens=False)["input_ids"])
+        len_completion = len(
+            tokenizer(completion, add_special_tokens=False)["input_ids"]
+        )
+
+        return (len_prompt + len_completion) <= sequence_len
+
+    if rl in {RLType.GRPO, RLType.GDPO, RLType.EBFT}:
+        return True
+
+    raise ValueError("Unknown RL type")
+
+
+def _raise_on_long_sequences(
+    sample: dict[str, Any], rl: RLType, tokenizer: Any, sequence_len: int
+) -> bool:
+    """Check sequence length and raise ValueError if exceeded.
+
+    Used as a filter function for ``excess_length_strategy: raise``.
+
+    Args:
+        sample: Dataset sample to check.
+        rl: Reinforcement learning type.
+        tokenizer: Tokenizer for length calculation.
+        sequence_len: Maximum allowed sequence length.
+
+    Returns:
+        Always True (raises before returning False).
+
+    Raises:
+        ValueError: If any sample exceeds the configured sequence length.
+    """
+    is_valid = _drop_long_sequences(sample, rl, tokenizer, sequence_len)
+    if not is_valid:
+        raise ValueError(
+            f"Sample exceeds configured sequence_len ({sequence_len}). "
+            "Set `excess_length_strategy: drop` or `excess_length_strategy: truncate` "
+            "to handle long sequences automatically."
+        )
+    return True
+
+
+def _truncate_long_sequences_rl(
+    sample: dict[str, Any], rl: RLType, tokenizer: Any, sequence_len: int
+) -> dict[str, Any]:
+    """Truncate RL samples that exceed maximum sequence length.
+
+    For preference datasets (DPO/IPO/ORPO/SIMPO), truncates chosen and rejected
+    responses to fit within ``sequence_len`` when combined with the prompt.
+    For KTO, truncates the completion similarly.
+    GRPO/GDPO/EBFT samples are returned unchanged.
+
+    Samples where the prompt alone exceeds ``sequence_len`` cannot be
+    meaningfully truncated and are returned unchanged.  The caller should
+    follow up with a drop filter to remove them.
+
+    Args:
+        sample: Dataset sample to potentially truncate.
+        rl: Reinforcement learning type.
+        tokenizer: Tokenizer for encoding/decoding.
+        sequence_len: Maximum allowed sequence length.
+
+    Returns:
+        The sample with text fields truncated to fit within sequence_len.
+    """
+    # Fast path: if sample already fits, return unchanged (avoids decode overhead)
+    if _drop_long_sequences(sample, rl, tokenizer, sequence_len):
+        return sample
+
+    if rl in {RLType.DPO, RLType.IPO, RLType.ORPO, RLType.SIMPO}:
+        if not (
+            sample.get("prompt") and sample.get("chosen") and sample.get("rejected")
+        ):
+            raise ValueError(
+                "Prompt, chosen and rejected keys are required for DPO/ORPO datasets"
+            )
+
+        prompt_ids = tokenizer(sample["prompt"], add_special_tokens=False)["input_ids"]
+        chosen_ids = tokenizer(sample["chosen"], add_special_tokens=False)["input_ids"]
+        rejected_ids = tokenizer(sample["rejected"], add_special_tokens=False)[
+            "input_ids"
+        ]
+
+        max_response_len = sequence_len - len(prompt_ids)
+        if max_response_len <= 0:
+            # Prompt alone exceeds limit; cannot meaningfully truncate.
+            # Returned unchanged — the follow-up drop filter will remove it.
+            return sample
+
+        updates: dict[str, Any] = {}
+        if len(chosen_ids) > max_response_len:
+            updates["chosen"] = tokenizer.decode(
+                chosen_ids[:max_response_len], skip_special_tokens=False
+            )
+        if len(rejected_ids) > max_response_len:
+            updates["rejected"] = tokenizer.decode(
+                rejected_ids[:max_response_len], skip_special_tokens=False
+            )
+        if updates:
+            sample = {**sample, **updates}
+
+    elif rl is RLType.KTO:
+        if not (sample.get("prompt") and sample.get("completion")):
+            raise ValueError("Prompt and completion keys are required for KTO datasets")
+
+        prompt_ids = tokenizer(sample["prompt"], add_special_tokens=False)["input_ids"]
+        completion_ids = tokenizer(sample["completion"], add_special_tokens=False)[
+            "input_ids"
+        ]
+
+        max_completion_len = sequence_len - len(prompt_ids)
+        if max_completion_len <= 0:
+            return sample
+
+        if len(completion_ids) > max_completion_len:
+            sample = {
+                **sample,
+                "completion": tokenizer.decode(
+                    completion_ids[:max_completion_len], skip_special_tokens=False
+                ),
+            }
+
+    # GRPO/GDPO/EBFT: no truncation needed (responses generated at runtime)
+    return sample
+
+
+def _load_split(cfg: DictDefault, split: Literal["train", "test"]) -> Dataset:
+    """Load and process dataset split for RL training.
+
+    Args:
+        cfg: Configuration object containing dataset settings.
+        split: Dataset split to load ("train" or "test").
+
+    Returns:
+        Combined and processed dataset for the specified split.
+    """
+    datasets_configs = cfg.datasets if split == "train" else cfg.test_datasets
+    split_datasets: list[Dataset | DatasetDict] = []
+
+    for dataset_config in datasets_with_name_generator(datasets_configs):
+        dataset: Dataset | DatasetDict = load_dataset_with_config(
+            dataset_config, cfg.hf_use_auth_token, streaming=False
+        )
+        split_datasets.append(dataset)
+
+    tokenizer = load_tokenizer(cfg)
+
+    for i, dataset in enumerate(split_datasets):
+        _type = datasets_configs[i]["type"]
+        if _type:
+            if isinstance(_type, DictDefault):
+                _type = "user_defined.default"
+            if cfg.rl is RLType.ORPO:
+                ds_transform_fn = load_orpo(_type, cfg, dataset_idx=i)
+            elif cfg.rl is RLType.KTO:
+                ds_transform_fn = load_kto(_type, cfg, dataset_idx=i)
+            elif cfg.rl is RLType.EBFT:
+                ds_transform_fn = load_ebft(_type, cfg, dataset_idx=i)
             else:
-                ds = load_dataset(  # pylint: disable=invalid-name
-                    ds_cfg["path"],
-                    split=ds_cfg["split"],
+                ds_transform_fn = load_dpo(_type, cfg, dataset_idx=i)
+
+            if ds_transform_fn is None:
+                raise ValueError(
+                    f"Failed to load dataset transform function for type {_type!r} "
+                    f"(rl: {cfg.rl}). See the warning logged above for the underlying "
+                    "error."
                 )
-                split_datasets.insert(i, ds)
-
-        tokenizer = None
-        for i, data_set in enumerate(split_datasets):
-            _type = dataset_cfgs[i]["type"]
-            if _type:
-                if isinstance(_type, DictDefault):
-                    _type = "user_defined.default"
-                if _cfg.rl == "orpo":
-                    ds_transform_fn = load_orpo(_type, _cfg, dataset_idx=i)
+
+            map_kwargs: dict[str, Any] = {}
+            if isinstance(ds_transform_fn, tuple):
+                ds_transform_fn, map_kwargs = ds_transform_fn
+            # Handle remove_columns: "__all__" removes all original columns,
+            # or filter a list to only columns that exist in the dataset
+            if "remove_columns" in map_kwargs:
+                ds_columns = (
+                    dataset.column_names
+                    if isinstance(dataset, Dataset)
+                    else dataset[split].column_names
+                    if isinstance(dataset, DatasetDict)
+                    else []
+                )
+                if map_kwargs["remove_columns"] == "__all__":
+                    map_kwargs["remove_columns"] = list(ds_columns)
                 else:
-                    ds_transform_fn = load_dpo(_type, _cfg, dataset_idx=i)
-                sig = inspect.signature(ds_transform_fn)
-                if "tokenizer" in sig.parameters:
-                    if not tokenizer:
-                        tokenizer = load_tokenizer(_cfg)
-                    ds_transform_fn = partial(ds_transform_fn, tokenizer=tokenizer)
-
-                data_set = data_set.map(
-                    ds_transform_fn,
-                    desc="Mapping RL Dataset",
+                    map_kwargs["remove_columns"] = [
+                        c for c in map_kwargs["remove_columns"] if c in ds_columns
+                    ]
+            split_datasets[i] = _map_dataset(
+                cfg, dataset, ds_transform_fn, tokenizer, **map_kwargs
+            )
+        else:
+            # If no `type` is provided, assume the dataset is already in the expected format with
+            # "prompt", "chosen", and "rejected" already preprocessed
+            split_datasets[i] = dataset
+
+        if not cfg.skip_prepare_dataset:
+            excess_length_strategy = (cfg.excess_length_strategy or "drop").lower()
+
+            if excess_length_strategy == "truncate":
+                truncate_fn = partial(
+                    _truncate_long_sequences_rl,
+                    rl=cfg.rl,
+                    tokenizer=tokenizer,
+                    sequence_len=cfg.sequence_len,
+                )
+                prior_len = len(split_datasets[i])
+                split_datasets[i] = split_datasets[i].map(
+                    truncate_fn,
+                    num_proc=cfg.dataset_num_proc,
+                    load_from_cache_file=not cfg.is_preprocess,
+                    desc="Truncating Long Sequences",
+                )
+
+                # Drop samples that could not be truncated (e.g. prompt
+                # alone exceeds sequence_len)
+                drop_long = partial(
+                    _drop_long_sequences,
+                    rl=cfg.rl,
+                    tokenizer=tokenizer,
+                    sequence_len=cfg.sequence_len,
+                )
+                split_datasets[i] = split_datasets[i].filter(
+                    drop_long,
+                    num_proc=cfg.dataset_num_proc,
+                    load_from_cache_file=not cfg.is_preprocess,
+                    desc="Dropping Un-truncatable Sequences",
+                )
+                dropped = prior_len - len(split_datasets[i])
+                if dropped:
+                    LOG.warning(
+                        f"Dropped {dropped} samples from dataset index {i} "
+                        f"that could not be truncated to fit sequence_len "
+                        f"(prompt alone exceeds limit)"
+                    )
+            elif excess_length_strategy == "raise":
+                raise_fn = partial(
+                    _raise_on_long_sequences,
+                    rl=cfg.rl,
+                    tokenizer=tokenizer,
+                    sequence_len=cfg.sequence_len,
+                )
+                split_datasets[i] = split_datasets[i].filter(
+                    raise_fn,
+                    num_proc=cfg.dataset_num_proc,
+                    load_from_cache_file=not cfg.is_preprocess,
+                    desc="Checking Sequence Lengths",
+                )
+            else:  # "drop" (default)
+                drop_long = partial(
+                    _drop_long_sequences,
+                    rl=cfg.rl,
+                    tokenizer=tokenizer,
+                    sequence_len=cfg.sequence_len,
                 )
-                if isinstance(data_set, DatasetDict):
-                    data_set = data_set["train"]
-                split_datasets[i] = data_set
-            else:
-                # If no `type` is provided, assume the dataset is already in the expected format with
-                # "prompt", "chosen" and "rejected" already preprocessed
-                split_datasets[i] = data_set
 
-        return concatenate_datasets(split_datasets)
+                prior_len = len(split_datasets[i])
+                split_datasets[i] = split_datasets[i].filter(
+                    drop_long,
+                    num_proc=cfg.dataset_num_proc,
+                    load_from_cache_file=not cfg.is_preprocess,
+                    desc="Dropping Long Sequences",
+                )
+                dropped = prior_len - len(split_datasets[i])
+                if dropped:
+                    LOG.warning(
+                        f"Dropped {dropped} long samples from dataset index {i}"
+                    )
 
-    with zero_first(is_main_process()):
-        train_is_preprocessed = False
-        eval_is_preprocessed = False
-        if train_dataset := _load_preprocessed_ds(cfg, cfg.datasets):
-            train_is_preprocessed = True
-        else:
-            train_dataset = load_split(cfg.datasets, cfg)
+    # Merge datasets
+    dataset = merge_datasets(split_datasets, cfg)
 
-        eval_dataset = None
-        if cfg.test_datasets:
-            if eval_dataset := _load_preprocessed_ds(cfg, cfg.test_datasets):
-                eval_is_preprocessed = True
-            else:
-                eval_dataset = load_split(cfg.test_datasets, cfg)
-        if not eval_dataset:
-            eval_dataset = None
+    if not cfg.skip_prepare_dataset:
+        # Deduplicate before saving so the saved dataset is already de-duplicated
+        if cfg.dataset_exact_deduplication:
+            dataset, _ = deduplicate_and_log_datasets(dataset=dataset)
 
-        if not train_is_preprocessed:
-            _save_preprocessed_ds(cfg, cfg.datasets, train_dataset)
-        if eval_dataset and not eval_is_preprocessed:
-            _save_preprocessed_ds(cfg, cfg.test_datasets, eval_dataset)
+        # Save preprocessed dataset
+        dataset_hash = generate_dataset_hash_from_config(
+            cfg, datasets_configs, tokenizer.name_or_path
+        )
+        save_preprocessed_dataset(cfg, dataset, dataset_hash, split)
 
-    return train_dataset, eval_dataset
+    return dataset
+
+
+def _load_or_create_dataset_split(
+    cfg: DictDefault, tokenizer: PreTrainedTokenizer, split: Literal["train", "test"]
+) -> Dataset:
+    """Load preprocessed dataset or create new one for given split.
+
+    Args:
+        cfg: Configuration object.
+        tokenizer: Tokenizer to use for processing text.
+        split: Dataset split to load.
+
+    Returns:
+        Tuple of (dataset, is_preprocessed).
+    """
+    # Select correct dataset configuration based on split
+    datasets_config = cfg.datasets if split == "train" else cfg.test_datasets
+
+    # Generate dataset hash for caching
+    dataset_hash = generate_dataset_hash_from_config(
+        cfg, datasets_config, tokenizer.name_or_path
+    )
+
+    # Try loading from hub if push_dataset_to_hub is configured
+    dataset = None
+    if cfg.push_dataset_to_hub:
+        dataset = try_load_from_hub(cfg, dataset_hash, split)
+
+    # Attempt to load preprocessed dataset
+    if dataset is None:
+        dataset = load_preprocessed_dataset(cfg, dataset_hash)
+
+    # Otherwise, load it
+    if dataset is None:
+        dataset = _load_split(cfg, split=split)
+
+    return dataset
diff --git a/src/axolotl/utils/data/sft.py b/src/axolotl/utils/data/sft.py
index dbc4172b4b..0b2ec2b5fb 100644
--- a/src/axolotl/utils/data/sft.py
+++ b/src/axolotl/utils/data/sft.py
@@ -1,681 +1,514 @@
-"""data handling specific to SFT"""
+"""Data handling specific to SFT."""
 
 import functools
-import logging
-from pathlib import Path
-from typing import List, Optional, Tuple, Union
+import os
+import tempfile
+from typing import Literal
 
 from datasets import (
     Dataset,
     DatasetDict,
-    concatenate_datasets,
+    IterableDataset,
+    IterableDatasetDict,
     load_dataset,
-    load_from_disk,
 )
-from huggingface_hub import hf_hub_download
-from huggingface_hub.utils import HFValidationError
-from transformers import PreTrainedTokenizerBase
-
-from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
-from axolotl.datasets import TokenizedPromptDataset
-from axolotl.prompt_strategies import load
-from axolotl.prompt_tokenizers import (
-    AlpacaMultipleChoicePromptTokenizingStrategy,
-    AlpacaPromptTokenizingStrategy,
-    AlpacaReflectionPTStrategy,
-    GPTeacherPromptTokenizingStrategy,
-    JeopardyPromptTokenizingStrategy,
-    OpenAssistantPromptTokenizingStrategy,
-    SummarizeTLDRPromptTokenizingStrategy,
+from transformers import PreTrainedTokenizer, ProcessorMixin
+
+from axolotl.prompters import Prompter
+from axolotl.utils.data.lock import FileLockLoader
+from axolotl.utils.data.shared import (
+    create_train_validation_split,
+    datasets_with_name_generator,
+    generate_dataset_hash_from_config,
+    load_dataset_with_config,
+    load_preprocessed_dataset,
+    merge_datasets,
+    save_preprocessed_dataset,
+    try_load_from_hub,
 )
-from axolotl.prompters import (
-    AlpacaPrompter,
-    GPTeacherPrompter,
-    JeopardyPrompter,
-    MultipleChoiceConcisePrompter,
-    MultipleChoiceExplainPrompter,
-    Prompter,
-    ReflectAlpacaPrompter,
-    SummarizeTLDRPrompter,
-    UnsupportedPrompter,
+from axolotl.utils.data.streaming import wrap_streaming_dataset
+from axolotl.utils.data.utils import (
+    deduplicate_and_log_datasets,
+    handle_long_seq_in_dataset,
+    retry_on_request_exceptions,
 )
-from axolotl.utils.data.pretraining import wrap_pretraining_dataset
-from axolotl.utils.data.utils import md5
+from axolotl.utils.data.wrappers import get_dataset_wrapper
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.distributed import is_main_process, zero_first
+from axolotl.utils.distributed import is_local_main_process
+from axolotl.utils.logging import get_logger
 from axolotl.utils.trainer import (
     calculate_total_num_steps,
     process_datasets_for_packing,
 )
 
-LOG = logging.getLogger("axolotl")
+LOG = get_logger(__name__)
 
 
-def prepare_dataset(cfg, tokenizer):
-    prompters = []
-    if not cfg.pretraining_dataset:
-        with zero_first(is_main_process()):
-            if cfg.test_datasets:
-                train_dataset, _, prompters = load_prepare_datasets(
-                    tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH, split="train"
-                )
-                _, eval_dataset, _ = load_prepare_datasets(
-                    tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH, split="test"
-                )
-            else:
-                train_dataset, eval_dataset, prompters = load_prepare_datasets(
-                    tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH
-                )
-    else:
-        path = cfg.pretraining_dataset
-        split = "train"
-        name = None
-        if isinstance(cfg.pretraining_dataset, list) and isinstance(
-            cfg.pretraining_dataset[0], dict
-        ):
-            path = cfg.pretraining_dataset[0]["path"]
-            name = cfg.pretraining_dataset[0]["name"]
-            if "split" in cfg.pretraining_dataset[0]:
-                split = cfg.pretraining_dataset[0]["split"]
-
-        ds_wrapper_partial = functools.partial(
-            get_dataset_wrapper,
-            cfg.pretraining_dataset[0],
-            tokenizer,
-            cfg,
-            cfg.pretraining_dataset[0]["type"] or "pretrain",
-        )
+@retry_on_request_exceptions(max_retries=3, delay=5)
+def prepare_datasets(
+    cfg: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    processor: ProcessorMixin | None = None,
+) -> tuple[IterableDataset | Dataset, Dataset | None, int, list[Prompter | None]]:
+    """Prepare training and evaluation datasets based on configuration.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        tokenizer: Tokenizer to use for processing text.
+        processor: Optional processor for multimodal datasets.
+
+    Returns:
+        Tuple of (train_dataset, eval_dataset, total_steps, prompters).
+    """
+    if cfg.streaming or cfg.pretraining_dataset:
+        return _prepare_streaming_dataset(cfg, tokenizer, processor)
+    return _prepare_standard_dataset(cfg, tokenizer, processor)
 
-        train_dataset = wrap_pretraining_dataset(
-            load_dataset(path, streaming=True, split=split, name=name),
+
+def _prepare_standard_dataset(
+    cfg: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    processor: ProcessorMixin | None,
+) -> tuple[Dataset, Dataset | None, int, list[Prompter | None]]:
+    """Prepare standard (non-pretraining) datasets."""
+
+    def _load_datasets():
+        # Always load training dataset
+        train_dataset, eval_dataset, prompters = _load_and_prepare_datasets(
             tokenizer,
             cfg,
-            ds_wrapper_partial,
-            max_tokens=cfg.sequence_len,
-            batch_size=cfg.micro_batch_size,
-            seed=cfg.seed or 42,
-            buffer_size=cfg.pretrain_multipack_buffer_size or 10_000,
+            split="train",
+            processor=processor,
         )
-        # https://discuss.huggingface.co/t/how-to-use-huggingface-trainer-streaming-datasets-without-wrapping-it-with-torchdatas-iterablewrapper/25230
-        train_dataset = train_dataset.with_format("torch")
-        eval_dataset = None
-        return train_dataset, eval_dataset, cfg.max_steps, prompters
 
+        # Overwrite eval_dataset if test data exists
+        if cfg.test_datasets:
+            _, eval_dataset, _ = _load_and_prepare_datasets(
+                tokenizer,
+                cfg,
+                split="test",
+                processor=processor,
+            )
+
+        return train_dataset, eval_dataset, prompters
+
+    # Prepare datasets (with file locking logic for multiple ranks)
+    loader = FileLockLoader(cfg)
+    try:
+        train_dataset, eval_dataset, prompters = loader.load(_load_datasets)
+    finally:
+        loader.cleanup()
+
+    if os.environ.get("AXOLOTL_IS_PREPROCESS") == "1":
+        return train_dataset, eval_dataset, -1, prompters
+
+    # Validate sample packing configuration for evaluation
     if eval_dataset and cfg.sample_packing and cfg.eval_sample_packing is not False:
         total_eval_steps = calculate_total_num_steps(cfg, eval_dataset, update=False)
         if total_eval_steps == 0:
             raise ValueError(
-                "eval dataset split is too small for sample_packing. You should set `eval_sample_packing: False`. "
+                "eval dataset split is too small for sample_packing. "
+                "You should set `eval_sample_packing: False` in your config."
             )
 
+    # Calculate total number of training steps
     if cfg.max_steps:
         total_num_steps = min(
             calculate_total_num_steps(cfg, train_dataset), cfg.max_steps
         )
-        LOG.info(f"Maximum number of steps set at {total_num_steps}")
     else:
         total_num_steps = calculate_total_num_steps(cfg, train_dataset)
+    LOG.info(f"Maximum number of steps set at {total_num_steps}")
     return train_dataset, eval_dataset, total_num_steps, prompters
 
 
-def load_tokenized_prepared_datasets(
-    tokenizer,
-    cfg,
-    default_dataset_prepared_path,
-    split="train",
-) -> Tuple[DatasetDict, List[Prompter]]:
-    cfg_datasets = cfg.test_datasets if split == "test" else cfg.datasets
-    tokenizer_name = cfg.tokenizer_config
-    ds_hash = str(
-        md5(
-            (
-                str(cfg.sequence_len)
-                + "@"
-                + str(cfg.sample_packing)
-                + "@"
-                + str(cfg.eval_sample_packing)
-                + "@"
-                + str(cfg.group_by_length)
-                + "@"
-                + "|".join(
-                    sorted(
-                        [
-                            f"{d.path}:{d.type}:{d.shards}:{d.conversation}{d.split}"
-                            for d in cfg_datasets
-                        ]
-                    )
-                )
-                + "|"
-                + tokenizer_name
-            )
-        )
-    )
-    prepared_ds_path = (
-        Path(cfg.dataset_prepared_path) / ds_hash
-        if cfg.dataset_prepared_path
-        else Path(default_dataset_prepared_path) / ds_hash
-    )
-    dataset = None
-    prompters = []
-    use_auth_token = cfg.hf_use_auth_token
-    try:
-        if cfg.push_dataset_to_hub:
-            dataset = load_dataset(
-                f"{cfg.push_dataset_to_hub}/{ds_hash}",
-                token=use_auth_token,
-            )
-            dataset = dataset[split]
-    except Exception:  # pylint: disable=broad-except # nosec
-        pass
-
-    # pylint: disable=duplicate-code
-    if dataset:
-        ...
-    elif (
-        cfg.dataset_prepared_path
-        and any(prepared_ds_path.glob("*"))
-        and not cfg.is_preprocess
-    ):
-        LOG.info(f"Loading prepared dataset from disk at {prepared_ds_path}...")
-        dataset = load_from_disk(str(prepared_ds_path))
-        LOG.info("Prepared dataset loaded from disk...")
+def _prepare_streaming_dataset(
+    cfg: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    processor: ProcessorMixin | None,
+) -> tuple[IterableDataset, Dataset | None, int, list[Prompter | None]]:
+    """
+    Prepare dataset for streaming mode.
+
+    Note: Streaming datasets are loaded incrementally from the source.
+    """
+    if cfg.pretraining_dataset:
+        dataset_config = _extract_pretraining_config(cfg)
+        train_dataset = _load_streaming_dataset(dataset_config, cfg, tokenizer)
+    elif cfg.sample_packing:
+        # TODO(djsaunde): Implement for multiple datasets
+        dataset_config = DictDefault(cfg.datasets[0])
+
+        # Ensure we have a split set - default to 'train' if not specified
+        if not hasattr(dataset_config, "split") or not dataset_config.split:
+            dataset_config.split = "train"
+        train_dataset = _load_streaming_dataset(dataset_config, cfg, tokenizer)
     else:
-        LOG.info(f"Unable to find prepared dataset in {prepared_ds_path}")
-        LOG.info("Loading raw datasets...")
-        if not cfg.is_preprocess:
-            LOG.warning(
-                "Processing datasets during training can lead to VRAM instability. Please pre-process your dataset."
-            )
+        # Use legacy loading function for non-packed streaming datasets
+        train_dataset, eval_dataset, prompters = _load_and_prepare_datasets(
+            tokenizer,
+            cfg,
+            split="train",
+            processor=processor,
+            streaming=True,
+        )
 
-        if cfg.seed:
-            seed = cfg.seed
-        else:
-            LOG.info("No seed provided, using default seed of 42")
-            seed = 42
-
-        datasets = []
-
-        def for_d_in_datasets(dataset_configs):
-            for dataset in dataset_configs:
-                if dataset.name and isinstance(dataset.name, list):
-                    for name in dataset.name:
-                        yield DictDefault({**dataset, "name": name})
-                else:
-                    yield dataset
-
-        # pylint: disable=invalid-name
-        for config_dataset in for_d_in_datasets(cfg_datasets):
-            ds: Optional[Union[Dataset, DatasetDict]] = None
-            ds_from_hub = False
-            try:
-                load_dataset(
-                    config_dataset.path,
-                    name=config_dataset.name,
-                    streaming=True,
-                    token=use_auth_token,
-                )
-                ds_from_hub = True
-            except (FileNotFoundError, ConnectionError, HFValidationError, ValueError):
-                pass
-
-            ds_from_cloud = False
-            storage_options = {}
-            remote_file_system = None
-            if config_dataset.path.startswith("s3://"):
-                try:
-                    import aiobotocore.session  # type: ignore
-                    import s3fs  # type: ignore
-                except ImportError as exc:
-                    raise ImportError(
-                        "s3:// paths require aiobotocore and s3fs to be installed"
-                    ) from exc
-
-                # Takes credentials from ~/.aws/credentials for default profile
-                s3_session = aiobotocore.session.AioSession(profile="default")
-                storage_options = {"session": s3_session}
-                remote_file_system = s3fs.S3FileSystem(**storage_options)
-            elif config_dataset.path.startswith(
-                "gs://"
-            ) or config_dataset.path.startswith("gcs://"):
-                try:
-                    import gcsfs  # type: ignore
-                except ImportError as exc:
-                    raise ImportError(
-                        "gs:// or gcs:// paths require gcsfs to be installed"
-                    ) from exc
-
-                # gcsfs will use default credentials from the environment else anon
-                # https://gcsfs.readthedocs.io/en/latest/#credentials
-                storage_options = {"token": None}
-                remote_file_system = gcsfs.GCSFileSystem(**storage_options)
-            # TODO: Figure out how to get auth creds passed
-            # elif config_dataset.path.startswith("adl://") or config_dataset.path.startswith("abfs://"):
-            #     try:
-            #         import adlfs
-            #     except ImportError as exc:
-            #        raise ImportError(
-            #            "adl:// or abfs:// paths require adlfs to be installed"
-            #        ) from exc
-
-            #     # Gen 1
-            #     storage_options = {
-            #         "tenant_id": TENANT_ID,
-            #         "client_id": CLIENT_ID,
-            #         "client_secret": CLIENT_SECRET,
-            #     }
-            #     # Gen 2
-            #     storage_options = {
-            #         "account_name": ACCOUNT_NAME,
-            #         "account_key": ACCOUNT_KEY,
-            #     }
-
-            #     remote_file_system = adlfs.AzureBlobFileSystem(**storage_options)
-            try:
-                if remote_file_system and remote_file_system.exists(
-                    config_dataset.path
-                ):
-                    ds_from_cloud = True
-            except (FileNotFoundError, ConnectionError):
-                pass
-
-            # prefer local dataset, even if hub exists
-            local_path = Path(config_dataset.path)
-            if local_path.exists():
-                if local_path.is_dir():
-                    if config_dataset.data_files:
-                        ds_type = get_ds_type(config_dataset)
-                        ds = load_dataset(
-                            ds_type,
-                            name=config_dataset.name,
-                            data_files=config_dataset.data_files,
-                            streaming=False,
-                            split=None,
-                        )
-                    else:
-                        ds = load_from_disk(config_dataset.path)
-                elif local_path.is_file():
-                    ds_type = get_ds_type(config_dataset)
-
-                    ds = load_dataset(
-                        ds_type,
-                        name=config_dataset.name,
-                        data_files=config_dataset.path,
-                        streaming=False,
-                        split=None,
-                    )
-                else:
-                    raise ValueError(
-                        "unhandled dataset load: local path exists, but is neither a directory or a file"
-                    )
-            elif ds_from_hub:
-                ds = load_dataset(
-                    config_dataset.path,
-                    name=config_dataset.name,
-                    streaming=False,
-                    data_files=config_dataset.data_files,
-                    token=use_auth_token,
-                )
-            elif ds_from_cloud and remote_file_system:
-                if remote_file_system.isdir(config_dataset.path):
-                    ds = load_from_disk(
-                        config_dataset.path,
-                        storage_options=storage_options,
-                    )
-                elif remote_file_system.isfile(config_dataset.path):
-                    ds_type = get_ds_type(config_dataset)
-                    ds = load_dataset(
-                        ds_type,
-                        name=config_dataset.name,
-                        data_files=config_dataset.path,
-                        streaming=False,
-                        split=None,
-                        storage_options=storage_options,
-                    )
-            elif config_dataset.path.startswith("https://"):
-                ds_type = get_ds_type(config_dataset)
-                ds = load_dataset(
-                    ds_type,
-                    name=config_dataset.name,
-                    data_files=config_dataset.path,
-                    streaming=False,
-                    split=None,
-                    storage_options=storage_options,
-                )
-            else:
-                if isinstance(config_dataset.data_files, str):
-                    fp = hf_hub_download(
-                        repo_id=config_dataset.path,
-                        repo_type="dataset",
-                        filename=config_dataset.data_files,
-                    )
-                elif isinstance(config_dataset.data_files, list):
-                    fp = []
-                    for file in config_dataset.data_files:
-                        fp.append(
-                            hf_hub_download(
-                                repo_id=config_dataset.path,
-                                repo_type="dataset",
-                                filename=file,
-                            )
-                        )
-                else:
-                    raise ValueError(
-                        "data_files must be either a string or list of strings"
-                    )
-                ds = load_dataset(
-                    "json",
-                    name=config_dataset.name,
-                    data_files=fp,
-                    streaming=False,
-                    split=None,
-                )
-            if not ds:
-                raise ValueError("unhandled dataset load")
-
-            d_base_type = d_prompt_style = None
-            d_type = config_dataset.type
-            if isinstance(d_type, str):
-                d_type_split = d_type.split(":")
-                d_base_type = d_type_split[0]
-                d_prompt_style = d_type_split[1] if len(d_type_split) > 1 else None
-
-            if isinstance(ds, DatasetDict):
-                if config_dataset.split and config_dataset.split in ds:
-                    ds = ds[config_dataset.split]
-                elif split in ds:
-                    ds = ds[split]
-                else:
-                    raise ValueError(
-                        f"no {split} split found for dataset {config_dataset.path}, you may specify a split with 'split: `"
-                    )
-
-            # support for using a subset of the data
-            if config_dataset.shards:
-                shards_idx = config_dataset.get("shards_idx", 0)
-                ds = ds.shuffle(seed=seed).shard(
-                    num_shards=config_dataset.shards, index=shards_idx
-                )
-
-            dataset_wrapper, dataset_prompter = get_dataset_wrapper(
-                config_dataset=config_dataset,
-                tokenizer=tokenizer,
-                cfg=cfg,
-                dataset=ds,
-                d_base_type=d_base_type,
-                d_prompt_style=d_prompt_style,
-            )
-            datasets.append(dataset_wrapper)
-            prompters.append(dataset_prompter)
-
-        LOG.info("merging datasets")
-        dataset = concatenate_datasets(datasets)
-
-        if len(datasets) > 1:
-            if cfg.shuffle_merged_datasets:
-                LOG.debug("shuffle merged datasets")
-                dataset = dataset.shuffle(seed=seed)
-            else:
-                LOG.debug("NOT shuffling merged datasets")
-
-        dataset, _ = process_datasets_for_packing(cfg, dataset, None)
-
-        if cfg.local_rank == 0:
-            LOG.info(f"Saving merged prepared dataset to disk... {prepared_ds_path}")
-            dataset.save_to_disk(str(prepared_ds_path))
-            if cfg.push_dataset_to_hub:
-                LOG.info(
-                    f"Saving merged prepared dataset with push_to_hub... {cfg.push_dataset_to_hub}/{ds_hash}"
-                )
-                dataset.push_to_hub(
-                    f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True
-                )
+        # Return early for non-packed streaming datasets
+        total_num_steps = cfg.max_steps if cfg.max_steps else -1
+        return train_dataset, eval_dataset, total_num_steps, prompters
 
-    return dataset, prompters
+    # Load evaluation dataset if specified
+    eval_dataset = None
+    if cfg.test_datasets:
+        _, eval_dataset, _ = _load_and_prepare_datasets(
+            tokenizer,
+            cfg,
+            split="test",
+            processor=processor,
+            streaming=False,
+        )
 
+    # For streaming, we return max_steps directly from config or -1 if not set
+    total_num_steps = cfg.max_steps if cfg.max_steps else -1
+    return train_dataset, eval_dataset, total_num_steps, []
 
-def get_ds_type(config_dataset: DictDefault):
-    """
-    Get the dataset type from the path if it's not specified
-    """
-    ds_type = "json"
-    if config_dataset.ds_type:
-        ds_type = config_dataset.ds_type
-    elif ".parquet" in config_dataset.path:
-        ds_type = "parquet"
-    elif ".arrow" in config_dataset.path:
-        ds_type = "arrow"
-    elif ".csv" in config_dataset.path:
-        ds_type = "csv"
-    elif ".txt" in config_dataset.path:
-        ds_type = "text"
-    return ds_type
-
-
-def load_prepare_datasets(
-    tokenizer: PreTrainedTokenizerBase,
-    cfg,
-    default_dataset_prepared_path,
-    split="train",
-) -> Tuple[Dataset, Dataset, List[Prompter]]:
-    dataset, prompters = load_tokenized_prepared_datasets(
-        tokenizer, cfg, default_dataset_prepared_path, split=split
-    )
 
-    if cfg.dataset_shard_num and cfg.dataset_shard_idx is not None:
-        LOG.info(
-            f"Using index #{cfg.dataset_shard_idx} of {cfg.dataset_shard_num} shards"
-        )
-        dataset = dataset.shard(
-            num_shards=cfg.dataset_shard_num,
-            index=cfg.dataset_shard_idx,
+def _extract_pretraining_config(cfg: DictDefault) -> DictDefault:
+    """Extract pretraining configuration from the main config."""
+    if isinstance(cfg.pretraining_dataset, list) and isinstance(
+        cfg.pretraining_dataset[0], dict
+    ):
+        config = cfg.pretraining_dataset[0]
+        return DictDefault(
+            {
+                "path": config["path"],
+                "name": config["name"],
+                "skip": config["skip"],
+                "split": config.get("split", "train"),
+                "data_files": config.get("data_files"),
+                "type": config.get("type", "pretrain"),
+            }
         )
+    # Simple string path case
+    return DictDefault(
+        {
+            "path": cfg.pretraining_dataset,
+            "name": None,
+            "skip": 0,
+            "split": "train",
+            "data_files": None,
+            "type": "pretrain",
+        }
+    )
 
-    if split == "train" and cfg.val_set_size:
-        # ensure we end up with the same fingerprint by doing rank0 first and being able to cache
-        to_hash_train = (
-            dataset._fingerprint  # pylint: disable=protected-access
-            + "|"
-            + str(cfg.val_set_size)
-            + "|"
-            + "train"
-            + "|"
-            + str(cfg.seed or 42)
-        )
-        to_hash_test = (
-            dataset._fingerprint  # pylint: disable=protected-access
-            + "|"
-            + str(cfg.val_set_size)
-            + "|"
-            + "test"
-            + "|"
-            + str(cfg.seed or 42)
-        )
-        train_fingerprint = md5(to_hash_train)
-        test_fingerprint = md5(to_hash_test)
-
-        dataset = dataset.train_test_split(
-            test_size=cfg.val_set_size,
-            shuffle=False,
-            seed=cfg.seed or 42,
-            train_new_fingerprint=train_fingerprint,
-            test_new_fingerprint=test_fingerprint,
-        )
 
-        train_dataset = dataset["train"]
-        eval_dataset = dataset["test"]
-    elif split == "test":
-        train_dataset = None
-        eval_dataset = dataset
+def _load_streaming_dataset(
+    pretraining_config: DictDefault, cfg: DictDefault, tokenizer: PreTrainedTokenizer
+) -> IterableDataset:
+    """Load and prepare a streaming dataset for pretraining."""
+    # Create dataset wrapper partial function
+    dataset_wrapper_partial = functools.partial(
+        get_dataset_wrapper,
+        dataset_config=pretraining_config,
+        tokenizer=tokenizer,
+        cfg=cfg,
+        dataset_base_type=pretraining_config["type"],
+    )
+
+    # Load the actual dataset
+    if (
+        cfg.accelerator_config
+        and cfg.accelerator_config.dispatch_batches
+        and not is_local_main_process()
+    ):
+        iter_dataset = _create_placeholder_dataset()
     else:
-        train_dataset = dataset
-        eval_dataset = None
+        iter_dataset = load_dataset(
+            pretraining_config["path"],
+            streaming=True,
+            split=pretraining_config["split"],
+            name=pretraining_config["name"],
+            data_files=pretraining_config["data_files"],
+        )
 
-    return train_dataset, eval_dataset, prompters
+    # Apply skip if specified
+    if pretraining_config["skip"]:
+        LOG.info(f"Skipping {pretraining_config['skip']} samples from the dataset")
+        iter_dataset = iter_dataset.skip(pretraining_config["skip"])
+
+    # Wrap the dataset for pretraining
+    train_dataset = wrap_streaming_dataset(
+        iter_dataset,
+        tokenizer,
+        cfg,
+        dataset_wrapper_partial,
+    )
 
+    # Format for PyTorch
+    return train_dataset.with_format("torch")
+
+
+def _create_placeholder_dataset() -> IterableDataset:
+    """Create a minimal placeholder dataset for non-main processes."""
+    with tempfile.NamedTemporaryFile(mode="w+", delete=False) as f:
+        f.write("text\n")
+        f.write("lorem ipsum dolor sit amet\n")
+        f.seek(0)
+        return load_dataset("csv", data_files=f.name, split="train", streaming=True)
+
+
+def _load_tokenized_prepared_datasets(
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    split: Literal["train", "test"] = "train",
+    processor: ProcessorMixin | None = None,
+    streaming: bool = False,
+) -> tuple[Dataset | DatasetDict, list[Prompter | None]]:
+    """Load or create tokenized and prepared datasets for training or testing.
+
+    Args:
+        tokenizer: Tokenizer for processing text.
+        cfg: Configuration object.
+        split: Dataset split to load ('train' or 'test').
+        processor: Optional processor for multimodal datasets.
+        streaming: Whether to use iterable preprocessing.
+
+    Returns:
+        Tuple of (dataset, prompters list).
+    """
+    # Select correct dataset configuration based on split
+    datasets_configs = cfg.datasets if split == "train" else cfg.test_datasets
 
-def get_dataset_wrapper(
-    config_dataset,
-    tokenizer,
-    cfg,
-    d_base_type,
-    dataset,
-    d_prompt_style=None,
-):
-    dataset_wrapper = None
-    dataset_prompter = None
+    # Generate dataset hash for caching
+    dataset_hash = generate_dataset_hash_from_config(
+        cfg, datasets_configs, tokenizer.name_or_path
+    )
 
-    ds_kwargs = {
-        "process_count": cfg.dataset_processes,
-        "keep_in_memory": cfg.dataset_keep_in_memory is True,
-    }
+    # Try loading from hub if push_dataset_to_hub is configured
+    dataset = None
+    if cfg.push_dataset_to_hub:
+        dataset = try_load_from_hub(cfg, dataset_hash, split)
 
-    if (
-        isinstance(dataset, Dataset)
-        and "input_ids" in dataset.features
-        and "attention_mask" in dataset.features
-        and "labels" in dataset.features
-    ):
-        # dataset is already tokenized, just drop it straight in
-        dataset_prompter = UnsupportedPrompter()
-        dataset_wrapper = dataset
-    elif isinstance(config_dataset.type, DictDefault):
-        ds_strategy = load(
-            "user_defined", tokenizer, cfg, config_dataset.type.to_dict()
-        )
-        dataset_prompter = UnsupportedPrompter()
-        dataset_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
-        )
-    elif ds_strategy := load(config_dataset.type, tokenizer, cfg, config_dataset):
-        dataset_prompter = UnsupportedPrompter()
-        dataset_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
-        )
-    elif d_base_type == "alpaca":
-        dataset_prompter = AlpacaPrompter(d_prompt_style)
-        ds_strategy = AlpacaPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "explainchoice":
-        dataset_prompter = MultipleChoiceExplainPrompter(d_prompt_style)
-        ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "concisechoice":
-        dataset_prompter = MultipleChoiceConcisePrompter(d_prompt_style)
-        ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
-        )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "summarizetldr":
-        dataset_prompter = SummarizeTLDRPrompter(d_prompt_style)
-        ds_strategy = SummarizeTLDRPromptTokenizingStrategy(
-            dataset_prompter,
+    # If not found on hub, try loading from disk
+    if dataset is None:
+        dataset = load_preprocessed_dataset(cfg, dataset_hash)
+
+    # If not found on disk or skipping prepared dataset, load and process raw datasets
+    prompters: list[Prompter | None] = []
+    if dataset is None:
+        dataset, prompters = _load_raw_datasets(
+            cfg,
+            datasets_configs,
             tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
-        )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
+            split,
+            processor,
+            streaming,
         )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "jeopardy":
-        dataset_prompter = JeopardyPrompter(d_prompt_style)
-        ds_strategy = JeopardyPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
+
+    return dataset, prompters
+
+
+def _load_raw_datasets(
+    cfg: DictDefault,
+    datasets_configs: list,
+    tokenizer: PreTrainedTokenizer,
+    split: str,
+    processor: ProcessorMixin | None = None,
+    streaming: bool = False,
+) -> tuple[Dataset, list[Prompter | None]]:
+    """Load, process, merge, and save raw datasets."""
+    LOG.info("Loading raw datasets...", main_process_only=False)
+    if not cfg.is_preprocess and not cfg.skip_prepare_dataset:
+        LOG.warning(
+            "Processing datasets during training can lead to VRAM instability. Please "
+            "pre-process your dataset using `axolotl preprocess path/to/config.yml`."
         )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
+
+    # Load and process individual datasets
+    datasets = []
+    prompters = []
+    for dataset_config in datasets_with_name_generator(datasets_configs):
+        dataset_wrapper, dataset_prompter = _load_and_process_single_dataset(
+            dataset_config=dataset_config,
+            cfg=cfg,
+            tokenizer=tokenizer,
+            split=split,
+            seed=cfg.seed,
+            processor=processor,
+            streaming=streaming,
         )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "oasst":
-        dataset_prompter = AlpacaPrompter(d_prompt_style)
-        ds_strategy = OpenAssistantPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
+        datasets.append(dataset_wrapper)
+        prompters.append(dataset_prompter)
+
+    # Merge datasets
+    dataset = merge_datasets(datasets, cfg)
+
+    if not cfg.skip_prepare_dataset and not streaming:
+        if split == "test" and cfg.eval_sequence_len:
+            dataset = handle_long_seq_in_dataset(dataset, cfg.eval_sequence_len, cfg)
+        else:
+            dataset = handle_long_seq_in_dataset(dataset, cfg.sequence_len, cfg)
+        if (split == "train" and cfg.sample_packing) or (
+            split == "test" and cfg.eval_sample_packing
+        ):
+            dataset, _ = process_datasets_for_packing(cfg, dataset, None)
+
+        # Deduplicate before saving so the saved dataset is already de-duplicated
+        if cfg.dataset_exact_deduplication:
+            dataset, _ = deduplicate_and_log_datasets(dataset=dataset)
+
+        # Save the prepared dataset
+        dataset_hash = generate_dataset_hash_from_config(
+            cfg, datasets_configs, tokenizer.name_or_path
         )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
+        save_preprocessed_dataset(cfg, dataset, dataset_hash, split)
+
+    return dataset, prompters
+
+
+def _load_and_process_single_dataset(
+    dataset_config: DictDefault,
+    cfg: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    split: str,
+    seed: int,
+    processor: ProcessorMixin | None = None,
+    streaming: bool = False,
+) -> tuple[Dataset | IterableDataset, Prompter | None]:
+    """Load and process a single dataset based on the passed config."""
+    # For synthetic datasets, create a minimal placeholder instead of loading from path
+    if dataset_config.type == "_synthetic":
+        dataset = Dataset.from_dict({"text": [""]})
+    else:
+        # Load the dataset
+        dataset = load_dataset_with_config(
+            dataset_config, cfg.hf_use_auth_token, streaming=streaming
         )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "gpteacher":
-        dataset_prompter = GPTeacherPrompter(d_prompt_style)
-        ds_strategy = GPTeacherPromptTokenizingStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
+
+    # Parse dataset type
+    d_base_type, d_prompt_style = _parse_dataset_type(dataset_config.type)
+
+    # Select the appropriate split
+    if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
+        if dataset_config.split and dataset_config.split in dataset:
+            dataset = dataset[dataset_config.split]
+        elif split in dataset:
+            dataset = dataset[split]
+        else:
+            raise ValueError(
+                f"no {split} split found for dataset {dataset_config.path}, you may "
+                "specify a split with 'split: ...'"
+            )
+
+    # Apply sharding if configured
+    if dataset_config.shards:
+        shards_idx = dataset_config.get("shards_idx", 0)
+        dataset = dataset.shuffle(seed=seed).shard(
+            num_shards=dataset_config.shards, index=shards_idx
         )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
+
+    # Apply dataset wrapper
+    dataset_wrapper, dataset_prompter = get_dataset_wrapper(
+        dataset_config=dataset_config,
+        tokenizer=tokenizer,
+        cfg=cfg,
+        dataset_base_type=d_base_type,
+        dataset=dataset,
+        dataset_prompt_style=d_prompt_style,
+        processor=processor,
+    )
+
+    return dataset_wrapper, dataset_prompter
+
+
+def _parse_dataset_type(d_type: str) -> tuple[str | None, str | None]:
+    """Parse the dataset type string into base type and prompt style."""
+    if not isinstance(d_type, str):
+        return None, None
+
+    d_type_split = d_type.split(":")
+    d_base_type = d_type_split[0]
+    d_prompt_style = d_type_split[1] if len(d_type_split) > 1 else None
+
+    return d_base_type, d_prompt_style
+
+
+def _handle_train_dataset_split(
+    dataset: Dataset, cfg: DictDefault
+) -> tuple[Dataset, Dataset | None]:
+    """Handle processing for train split, including validation set creation."""
+    val_set_size = (
+        int(cfg.val_set_size) if cfg.val_set_size > 1 else float(cfg.val_set_size)
+    )
+
+    if val_set_size:
+        # Create train/validation split
+        train_dataset, eval_dataset = create_train_validation_split(
+            dataset, cfg, val_set_size
         )
-        dataset_wrapper = ds_wrapper
-    elif d_base_type == "reflection":
-        dataset_prompter = ReflectAlpacaPrompter(d_prompt_style)
-        ds_strategy = AlpacaReflectionPTStrategy(
-            dataset_prompter,
-            tokenizer,
-            cfg.train_on_inputs,
-            cfg.sequence_len,
+        return train_dataset, eval_dataset
+
+    # No validation split - deduplication already applied during preprocessing
+    return dataset, None
+
+
+def _apply_dataset_sharding(dataset: Dataset, cfg: DictDefault) -> Dataset:
+    """Apply dataset sharding if configured.
+
+    Args:
+        dataset: Dataset to shard.
+        cfg: Configuration object containing shard settings.
+
+    Returns:
+        Sharded dataset or original dataset if no sharding configured.
+    """
+    if cfg.dataset_shard_num and cfg.dataset_shard_idx is not None:
+        LOG.info(
+            f"Using index #{cfg.dataset_shard_idx} of {cfg.dataset_shard_num} shards"
         )
-        ds_wrapper = TokenizedPromptDataset(
-            ds_strategy,
-            dataset,
-            **ds_kwargs,
+        dataset = dataset.shard(
+            num_shards=cfg.dataset_shard_num,
+            index=cfg.dataset_shard_idx,
         )
-        dataset_wrapper = ds_wrapper
+    return dataset
+
+
+def _load_and_prepare_datasets(
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    split: Literal["train", "test"] = "train",
+    processor: ProcessorMixin | None = None,
+    streaming: bool = False,
+) -> tuple[Dataset | None, Dataset | None, list[Prompter | None]]:
+    """Load and prepare datasets with optional validation split and sharding.
+
+    Args:
+        tokenizer: Tokenizer for processing text.
+        cfg: Configuration object.
+        split: Dataset split to load ('train' or 'test').
+        processor: Optional processor for multimodal datasets.
+        streaming: Whether to use iterable preprocessing.
+
+    Returns:
+        Tuple of (train_dataset, eval_dataset, prompters).
+    """
+    # Load the base dataset
+    dataset, prompters = _load_tokenized_prepared_datasets(
+        tokenizer,
+        cfg,
+        split=split,
+        processor=processor,
+        streaming=streaming,
+    )
+
+    # Apply dataset sharding if configured using shared function
+    dataset = _apply_dataset_sharding(dataset, cfg)
+
+    # Apply deduplication and create train / validation splits based on the split type
+    if split == "train":
+        train_dataset, eval_dataset = _handle_train_dataset_split(dataset, cfg)
     else:
-        suffix = ""
-        if ":load_" in config_dataset.type:
-            suffix = f" Did you mean {config_dataset.type.replace(':load_', '.load_')}?"
-        LOG.error(
-            f"unhandled prompt tokenization strategy: {config_dataset.type}. {suffix}"
-        )
-        raise ValueError(
-            f"unhandled prompt tokenization strategy: {config_dataset.type} {suffix}"
-        )
+        # Deduplication already applied during preprocessing
+        train_dataset, eval_dataset = None, dataset
 
-    return dataset_wrapper, dataset_prompter
+    return train_dataset, eval_dataset, prompters
diff --git a/src/axolotl/utils/data/shared.py b/src/axolotl/utils/data/shared.py
new file mode 100644
index 0000000000..52654d9ff9
--- /dev/null
+++ b/src/axolotl/utils/data/shared.py
@@ -0,0 +1,726 @@
+"""Dataset loading shared utils."""
+
+from __future__ import annotations
+
+import functools
+import json
+import os
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator
+
+from datasets import (
+    Dataset,
+    DatasetDict,
+    IterableDataset,
+    IterableDatasetDict,
+    concatenate_datasets,
+    load_dataset,
+    load_from_disk,
+)
+from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub.errors import (
+    HFValidationError,
+    RepositoryNotFoundError,
+    RevisionNotFoundError,
+)
+
+from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
+from axolotl.utils.data.utils import deduplicate_and_log_datasets, md5
+from axolotl.utils.datasets import get_default_process_count
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from adlfs import AzureBlobFileSystem
+    from gcsfs import GCSFileSystem
+    from ocifs import OCIFileSystem
+    from s3fs import S3FileSystem
+
+LOG = get_logger(__name__)
+
+DATASET_HASH_BASE_DATASET_FIELDS = (
+    "path",
+    "type",
+    "shards",
+    "conversation",
+    "split",
+    "temperature",
+)
+
+DATASET_HASH_CONFIG_EXTRA_FIELDS = (
+    "added_tokens_overrides",
+    "chat_template",
+    "chat_template_jinja",
+    "chat_template_kwargs",
+    "default_system_message",
+    "eot_tokens",
+    "sample_packing_bin_size",
+    "sample_packing_group_size",
+    "sample_packing_sequentially",
+    "special_tokens",
+    "train_on_inputs",
+)
+
+DATASET_HASH_DATASET_EXTRA_FIELDS = (
+    "chat_template",
+    "chat_template_jinja",
+    "data_files",
+    "drop_system_message",
+    "ds_type",
+    "field",
+    "field_chosen",
+    "field_human",
+    "field_messages",
+    "field_model",
+    "field_rejected",
+    "field_thinking",
+    "field_tools",
+    "input_format",
+    "input_transform",
+    "max_completion_length",
+    "message_field_content",
+    "message_field_role",
+    "message_field_training",
+    "message_field_training_detail",
+    "message_property_mappings",
+    "name",
+    "preprocess_shards",
+    "revision",
+    "roles",
+    "roles_to_train",
+    "shards_idx",
+    "split_thinking",
+    "step_separator",
+    "template_thinking_key",
+    "text_column",
+    "train_on_eos",
+    "train_on_eot",
+    "train_on_last_step_only",
+)
+
+EXTENSIONS_TO_DATASET_TYPES = {
+    ".parquet": "parquet",
+    ".arrow": "arrow",
+    ".csv": "csv",
+    ".txt": "text",
+}
+
+
+def get_dataset_type(dataset_config: DictDefault) -> str:
+    """Get the dataset type from the path if it's not specified."""
+    if dataset_config.ds_type:
+        return dataset_config.ds_type
+
+    for extension, dataset_type in EXTENSIONS_TO_DATASET_TYPES.items():
+        if extension in dataset_config.path:
+            return dataset_type
+
+    return "json"
+
+
+def datasets_with_name_generator(
+    dataset_configs: list[DictDefault],
+) -> Generator[DictDefault, None, None]:
+    """Yields expanded dataset configurations based on multiple names or preprocessing
+    shards.
+
+    When a dataset config has a list of names, it yields separate configs for each
+    name. When a dataset config specifies preprocessing shards, it yields configs for
+    each shard.
+
+    Args:
+        dataset_configs: List of dataset configuration objects.
+
+    Yields:
+        Individual dataset configurations, expanded as needed for names or shards.
+    """
+    for config in dataset_configs:
+        if config.name and isinstance(config.name, list):
+            for name in config.name:
+                yield DictDefault({**config, "name": name})
+        elif config.preprocess_shards and not config.shards:
+            for shard_idx in range(config.preprocess_shards):
+                yield DictDefault(
+                    {
+                        **config,
+                        "shards": config.preprocess_shards,
+                        "shards_idx": shard_idx,
+                    }
+                )
+        else:
+            yield config
+
+
+def load_dataset_with_config(
+    dataset_config: DictDefault, use_auth_token: bool, streaming=False
+) -> Dataset | IterableDataset:
+    """Load a dataset from a config. Handles datasets that are stored locally, in the
+    HuggingFace Hub, in a remote filesystem (S3, GCS, Azure, OCI), a URL, or
+    `data_files`.
+
+    Args:
+        dataset_config: Single dataset config.
+        use_auth_token: Whether to use HF auth token.
+        streaming: Whether to stream the dataset.
+
+    Returns:
+        Loaded dataset.
+    """
+    # Set up common kwargs for dataset loading
+    load_dataset_kwargs = {
+        "split": dataset_config.split if dataset_config.split else None,
+        "name": dataset_config.name,
+        "streaming": streaming,
+        "trust_remote_code": dataset_config.trust_remote_code,
+    }
+
+    # First check if it's a local path
+    if Path(dataset_config.path).exists():
+        return _load_from_local_path(dataset_config, load_dataset_kwargs)
+
+    # Check if it's a HuggingFace dataset
+    is_hub_dataset = _check_if_hub_dataset(dataset_config, use_auth_token)
+
+    # Check if it's a cloud storage path and get appropriate filesystem
+    remote_fs, storage_options = _get_remote_filesystem(dataset_config.path)
+    is_cloud_dataset = False
+    if remote_fs:
+        try:
+            is_cloud_dataset = remote_fs.exists(dataset_config.path)
+        except (FileNotFoundError, ConnectionError):
+            pass
+
+    # Load from appropriate source
+    if is_hub_dataset:
+        return _load_from_hub(dataset_config, use_auth_token, load_dataset_kwargs)
+    if is_cloud_dataset:
+        return _load_from_cloud(
+            dataset_config, remote_fs, storage_options, load_dataset_kwargs
+        )
+    if dataset_config.path.startswith("https://"):
+        return _load_from_url(dataset_config, load_dataset_kwargs)
+    if dataset_config.data_files:
+        return _load_from_data_files(dataset_config, load_dataset_kwargs)
+
+    raise ValueError(
+        f"The dataset could not be loaded. This could be due to a misconfigured dataset path "
+        f"({dataset_config.path}). Try double-check your path / name / data_files. "
+        f"This is not caused by the dataset type."
+    )
+
+
+def _check_if_hub_dataset(dataset_config: DictDefault, use_auth_token: bool) -> bool:
+    """Check if a dataset exists on the HuggingFace Hub."""
+    try:
+        snapshot_download(
+            repo_id=dataset_config.path,
+            repo_type="dataset",
+            token=use_auth_token,
+            revision=dataset_config.revision,
+            ignore_patterns=["*"],
+        )
+        return True
+    except (
+        RepositoryNotFoundError,
+        RevisionNotFoundError,
+        FileNotFoundError,
+        ConnectionError,
+        HFValidationError,
+        ValueError,
+    ):
+        return False
+
+
+def _get_remote_filesystem(
+    path: str,
+) -> tuple[
+    S3FileSystem | GCSFileSystem | AzureBlobFileSystem | OCIFileSystem | None, dict
+]:
+    """Get the appropriate filesystem for a remote path."""
+    if path.startswith("s3://"):
+        try:
+            import s3fs
+
+            storage_options = {"anon": False}
+            return s3fs.S3FileSystem(**storage_options), storage_options
+        except ImportError as exc:
+            raise ImportError("s3:// paths require s3fs to be installed") from exc
+
+    elif path.startswith(("gs://", "gcs://")):
+        try:
+            import gcsfs
+
+            storage_options = {"token": None}  # type: ignore  # nosec B105
+            return gcsfs.GCSFileSystem(**storage_options), storage_options
+        except ImportError as exc:
+            raise ImportError(
+                "gs:// or gcs:// paths require gcsfs to be installed"
+            ) from exc
+
+    elif path.startswith(("adl://", "abfs://", "az://")):
+        try:
+            import adlfs
+
+            storage_options = {"anon": False}
+            return adlfs.AzureBlobFileSystem(**storage_options), storage_options
+        except ImportError as exc:
+            raise ImportError(
+                "adl:// or abfs:// paths require adlfs to be installed"
+            ) from exc
+
+    elif path.startswith("oci://"):
+        try:
+            import ocifs
+
+            storage_options = {}
+            return ocifs.OCIFileSystem(**storage_options), storage_options
+        except ImportError as exc:
+            raise ImportError("oci:// paths require ocifs to be installed") from exc
+
+    return None, {}
+
+
+def _load_from_local_path(
+    dataset_config: DictDefault, load_dataset_kwargs: dict
+) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
+    """Load a dataset from a local path."""
+    local_path = Path(dataset_config.path)
+
+    if local_path.is_dir():
+        if dataset_config.data_files:
+            dataset_type = get_dataset_type(dataset_config)
+            return load_dataset(
+                dataset_type,
+                data_files=dataset_config.data_files,
+                **load_dataset_kwargs,
+            )
+        try:
+            return load_from_disk(dataset_config.path)
+        except FileNotFoundError:
+            return load_dataset(dataset_config.path, **load_dataset_kwargs)
+    elif local_path.is_file():
+        dataset_type = get_dataset_type(dataset_config)
+
+        # Single-file paths always expose only a "train" split regardless of format.
+        # Preserve any user-provided slice syntax (e.g. "train[:500]"); only
+        # fall back to "train" when no split was specified.
+        if not load_dataset_kwargs.get("split"):
+            load_dataset_kwargs["split"] = "train"
+
+        return load_dataset(
+            dataset_type,
+            data_files=dataset_config.path,
+            **load_dataset_kwargs,
+        )
+    else:
+        raise ValueError(
+            "Unhandled dataset load: local path exists, but is neither a directory or a file"
+        )
+
+
+def _load_from_hub(
+    dataset_config: DictDefault, use_auth_token: bool, load_dataset_kwargs: dict
+) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
+    """Load a dataset from the HuggingFace Hub."""
+    return load_dataset(
+        dataset_config.path,
+        data_files=dataset_config.data_files,
+        token=use_auth_token,
+        revision=dataset_config.revision,
+        **load_dataset_kwargs,
+    )
+
+
+def _load_from_cloud(
+    dataset_config: DictDefault,
+    remote_fs: S3FileSystem | GCSFileSystem | AzureBlobFileSystem | OCIFileSystem,
+    storage_options: dict,
+    load_dataset_kwargs: dict,
+) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
+    """Load a dataset from cloud storage."""
+    if remote_fs.isdir(dataset_config.path):
+        return load_from_disk(
+            dataset_config.path,
+            storage_options=storage_options,
+        )
+
+    if remote_fs.isfile(dataset_config.path):
+        dataset_type = get_dataset_type(dataset_config)
+        return load_dataset(
+            dataset_type,
+            data_files=dataset_config.path,
+            storage_options=storage_options,
+            **load_dataset_kwargs,
+        )
+
+    raise ValueError(
+        f"Cloud path {dataset_config.path} is neither a directory nor a file"
+    )
+
+
+def _load_from_url(
+    dataset_config: DictDefault, load_dataset_kwargs: dict
+) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
+    """Load a dataset from a URL."""
+    dataset_type = get_dataset_type(dataset_config)
+    return load_dataset(
+        dataset_type,
+        data_files=dataset_config.path,
+        **load_dataset_kwargs,
+    )
+
+
+def _load_from_data_files(
+    dataset_config: DictDefault, load_dataset_kwargs: dict
+) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
+    """Load a dataset from data files."""
+    file_path = None
+
+    if isinstance(dataset_config.data_files, str):
+        file_path = hf_hub_download(
+            repo_id=dataset_config.path,
+            repo_type="dataset",
+            filename=dataset_config.data_files,
+            revision=dataset_config.revision,
+        )
+    elif isinstance(dataset_config.data_files, list):
+        file_path = [
+            hf_hub_download(
+                repo_id=dataset_config.path,
+                repo_type="dataset",
+                filename=file,
+                revision=dataset_config.revision,
+            )
+            for file in dataset_config.data_files
+        ]
+    else:
+        raise ValueError("data_files must be either a string or list of strings")
+
+    return load_dataset("json", data_files=file_path, **load_dataset_kwargs)
+
+
+def generate_split_fingerprints(
+    dataset: Dataset, val_set_size: int | float, seed: int
+) -> tuple[str, str]:
+    """Generate consistent fingerprints for train/test splits."""
+    fingerprint = dataset._fingerprint
+
+    train_hash_input = f"{fingerprint}|{val_set_size}|train|{seed}"
+    test_hash_input = f"{fingerprint}|{val_set_size}|test|{seed}"
+
+    train_fingerprint = md5(train_hash_input)
+    test_fingerprint = md5(test_hash_input)
+
+    return train_fingerprint, test_fingerprint
+
+
+def get_prepared_dataset_path(cfg: DictDefault, dataset_hash: str) -> Path:
+    """Get standardized path for prepared datasets.
+
+    Args:
+        cfg: Configuration object.
+        dataset_hash: Hash identifying the specific dataset configuration.
+
+    Returns:
+        Path where the prepared dataset should be stored.
+    """
+    base_path = cfg.dataset_prepared_path or DEFAULT_DATASET_PREPARED_PATH
+    return Path(base_path) / dataset_hash
+
+
+def create_train_validation_split(
+    dataset: Dataset, cfg: DictDefault, val_set_size: int | float
+) -> tuple[Dataset, Dataset]:
+    """Create train/validation split with consistent fingerprinting.
+
+    Args:
+        dataset: Dataset to split.
+        cfg: Configuration object containing seed and other settings.
+        val_set_size: Size of validation set (absolute number or fraction).
+
+    Returns:
+        Tuple of (train_dataset, eval_dataset).
+    """
+    train_fingerprint, test_fingerprint = generate_split_fingerprints(
+        dataset, val_set_size, cfg.seed
+    )
+
+    # Apply deduplication before splitting if configured
+    if cfg.dataset_exact_deduplication:
+        dataset, _ = deduplicate_and_log_datasets(dataset=dataset)
+
+    split_dataset = dataset.train_test_split(
+        test_size=val_set_size,
+        shuffle=False,
+        seed=cfg.seed,
+        train_new_fingerprint=train_fingerprint,
+        test_new_fingerprint=test_fingerprint,
+    )
+
+    return split_dataset["train"], split_dataset["test"]
+
+
+def _generate_from_iterable_dataset(
+    dataset: IterableDataset, worker_id: list[int], num_workers: list[int]
+) -> Generator[Any, None, None]:
+    """Generator function to correctly split the dataset for each worker"""
+    for i, item in enumerate(dataset):
+        if i % num_workers[0] == worker_id[0]:
+            yield item
+
+
+def save_preprocessed_dataset(
+    cfg: DictDefault,
+    dataset: Dataset,
+    dataset_hash: str,
+    split: str,
+) -> None:
+    """Save preprocessed dataset to disk and optionally push to the HF Hub."""
+    prepared_ds_path = get_prepared_dataset_path(cfg, dataset_hash)
+    num_workers = cfg.dataset_num_proc or get_default_process_count()
+    if isinstance(dataset, IterableDataset):
+        ds_from_iter = Dataset.from_generator(
+            functools.partial(_generate_from_iterable_dataset, dataset),
+            features=dataset.features,
+            num_proc=num_workers,
+            split=split,
+            gen_kwargs={
+                "worker_id": list(range(num_workers)),
+                "num_workers": [num_workers] * num_workers,
+            },
+        )
+        ds_from_iter.save_to_disk(
+            str(prepared_ds_path),
+            num_proc=num_workers,
+            max_shard_size=None,
+            num_shards=cfg.num_dataset_shards_to_save,
+        )
+    else:
+        min_rows_per_proc = 256
+        os.makedirs(prepared_ds_path, exist_ok=True)
+        dataset.save_to_disk(
+            str(prepared_ds_path),
+            num_proc=min(max(1, len(dataset) // min_rows_per_proc), num_workers),
+            max_shard_size=None,
+            num_shards=cfg.num_dataset_shards_to_save,
+        )
+    if cfg.push_dataset_to_hub:
+        LOG.info(
+            "Pushing merged prepared dataset to Huggingface hub at "
+            f"{cfg.push_dataset_to_hub} (version {dataset_hash})...",
+            main_process_only=False,
+        )
+        dataset.push_to_hub(
+            cfg.push_dataset_to_hub,
+            dataset_hash,
+            private=True,
+        )
+
+
+def load_preprocessed_dataset(cfg: DictDefault, dataset_hash: str) -> Dataset | None:
+    """Load preprocessed dataset from disk if available.
+
+    Args:
+        cfg: Configuration object.
+        dataset_hash: Hash identifying the dataset configuration.
+
+    Returns:
+        Loaded dataset if found and conditions are met, None otherwise.
+    """
+    prepared_ds_path = get_prepared_dataset_path(cfg, dataset_hash)
+
+    if (
+        any(prepared_ds_path.glob("*"))
+        and not cfg.skip_prepare_dataset
+        and not cfg.is_preprocess
+    ):
+        LOG.info(
+            f"Loading prepared dataset from disk at {prepared_ds_path}...",
+        )
+        return load_from_disk(str(prepared_ds_path))
+
+    LOG.info(
+        f"Unable to find prepared dataset in {prepared_ds_path}",
+    )
+    return None
+
+
+def try_load_from_hub(
+    cfg: DictDefault, dataset_hash: str, split: str
+) -> Dataset | None:
+    """Try to load the prepared dataset from HuggingFace Hub."""
+    try:
+        LOG.info(
+            "Attempting to load prepared dataset from HuggingFace Hub at "
+            f"{cfg.push_dataset_to_hub} (version {dataset_hash})..."
+        )
+        dataset = load_dataset(
+            cfg.push_dataset_to_hub,
+            dataset_hash,
+            token=cfg.hf_use_auth_token,
+        )
+        return dataset[split]
+    except Exception:
+        LOG.info("Unable to find prepared dataset in HuggingFace Hub")
+        return None
+
+
+def _get_hash_field(value: Any, field: str) -> Any:
+    if hasattr(value, "get"):
+        return value.get(field)
+    return getattr(value, field, None)
+
+
+def _normalize_hash_value(value: Any) -> Any:
+    model_dump = getattr(value, "model_dump", None)
+    if callable(model_dump):
+        try:
+            value = model_dump(mode="json", exclude_none=True)
+        except TypeError:
+            value = model_dump(exclude_none=True)
+
+    if isinstance(value, Enum):
+        return value.value
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, dict):
+        normalized = {}
+        for key, item in sorted(value.items(), key=lambda entry: str(entry[0])):
+            item = _normalize_hash_value(item)
+            if item is not None:
+                normalized[str(key)] = item
+        return normalized
+    if isinstance(value, (list, tuple)):
+        return [_normalize_hash_value(item) for item in value]
+    if isinstance(value, set):
+        return sorted(_normalize_hash_value(item) for item in value)
+
+    return value
+
+
+def _is_empty_hash_value(value: Any) -> bool:
+    return value is None or value is False or value == {} or value == []
+
+
+def _selected_hash_fields(value: Any, fields: tuple[str, ...]) -> dict[str, Any]:
+    selected = {}
+    for field in fields:
+        field_value = _normalize_hash_value(_get_hash_field(value, field))
+        if not _is_empty_hash_value(field_value):
+            selected[field] = field_value
+    return selected
+
+
+def _stable_json(value: Any) -> str:
+    return json.dumps(
+        _normalize_hash_value(value),
+        sort_keys=True,
+        separators=(",", ":"),
+    )
+
+
+def _base_dataset_hash_string(dataset: Any) -> str:
+    values = []
+    for field in DATASET_HASH_BASE_DATASET_FIELDS:
+        value = _get_hash_field(dataset, field)
+        if field == "temperature":
+            value = value or 1.0
+        values.append(str(value))
+    return ":".join(values)
+
+
+def generate_dataset_hash_from_config(
+    cfg: DictDefault, cfg_datasets: list, tokenizer_name: str
+) -> str:
+    """Generate a hash to uniquely identify a dataset configuration for SFT.
+
+    Args:
+        cfg: Main configuration object.
+        cfg_datasets: List of dataset configurations.
+        tokenizer_name: Name of the tokenizer being used.
+
+    Returns:
+        MD5 hash string representing the configuration.
+    """
+    # When added_tokens_overrides is set, tokenizer.name_or_path contains output_dir.
+    # Use the canonical tokenizer config + overrides content so the hash is stable across output_dir changes.
+    if cfg.get("added_tokens_overrides"):
+        tokenizer_fingerprint = f"{cfg.tokenizer_config}+overrides:" + ",".join(
+            f"{k}={v}" for k, v in sorted(cfg.added_tokens_overrides.items())
+        )
+    else:
+        tokenizer_fingerprint = tokenizer_name
+
+    dataset_config_str = "|".join(
+        sorted(_base_dataset_hash_string(dataset) for dataset in cfg_datasets)
+    )
+
+    extra_payload: dict[str, Any] = {}
+    cfg_extra = _selected_hash_fields(cfg, DATASET_HASH_CONFIG_EXTRA_FIELDS)
+    if cfg_extra:
+        extra_payload["cfg"] = cfg_extra
+
+    dataset_extra = []
+    for dataset in cfg_datasets:
+        extra_fields = _selected_hash_fields(dataset, DATASET_HASH_DATASET_EXTRA_FIELDS)
+        if extra_fields:
+            dataset_extra.append(
+                {"id": _base_dataset_hash_string(dataset), "extra": extra_fields}
+            )
+    if dataset_extra:
+        extra_payload["datasets"] = sorted(dataset_extra, key=_stable_json)
+
+    config_str = (
+        f"{cfg.sequence_len}@{cfg.sample_packing}@{cfg.eval_sample_packing}@"
+        f"{cfg.group_by_length}@{cfg.kd_temperature or 1.0}@"
+        f"{cfg.dataset_exact_deduplication or False}|"
+        f"{dataset_config_str}"
+        f"|{tokenizer_fingerprint}"
+    )
+    if extra_payload:
+        config_str = f"{config_str}|{_stable_json(extra_payload)}"
+    return str(md5(config_str))
+
+
+def merge_datasets(datasets: list[Dataset], cfg: DictDefault) -> Dataset:
+    """Merge multiple datasets into one with optional shuffling.
+
+    Args:
+        datasets: List of datasets to merge.
+        cfg: Configuration object containing shuffle settings.
+
+    Returns:
+        Merged dataset.
+    """
+    if len(datasets) == 1:
+        ds = datasets[0]
+
+        # Do not shuffle if curriculum sampling is enabled or
+        # shuffle_merged_datasets is disabled
+        if cfg.curriculum_sampling or not cfg.shuffle_merged_datasets:
+            return ds
+
+        return ds.shuffle(seed=cfg.seed)
+
+    # If enabled, shuffle each dataset independently before merging.
+    # This allows curriculum learning strategies to be applied at the dataset level.
+    if cfg.shuffle_before_merging_datasets:
+        LOG.info("Shuffling each dataset individually before merging...")
+        datasets = [ds.shuffle(seed=cfg.seed) for ds in datasets]
+
+    LOG.info("Merging datasets...")
+    merged_dataset = concatenate_datasets(datasets)
+
+    if cfg.shuffle_merged_datasets:
+        LOG.debug("Shuffling merged datasets...")
+        if cfg.curriculum_sampling:
+            LOG.warning(
+                "Shuffling merged datasets with curriculum sampling is not recommended. "
+                "This will randomize the order of samples."
+            )
+        merged_dataset = merged_dataset.shuffle(seed=cfg.seed)
+    else:
+        LOG.debug("Not shuffling merged datasets.")
+
+    return merged_dataset
diff --git a/src/axolotl/utils/data/pretraining.py b/src/axolotl/utils/data/streaming.py
similarity index 61%
rename from src/axolotl/utils/data/pretraining.py
rename to src/axolotl/utils/data/streaming.py
index 544ed13162..8b6b8a439b 100644
--- a/src/axolotl/utils/data/pretraining.py
+++ b/src/axolotl/utils/data/streaming.py
@@ -1,7 +1,6 @@
-"""data handling specific to pretraining"""
+"""Data handling specific to streaming datasets."""
 
 import functools
-import logging
 from collections import defaultdict
 from typing import Callable, Dict, List, Optional
 
@@ -11,25 +10,39 @@
 from transformers import PreTrainedTokenizerBase
 
 from axolotl.utils.collators import PretrainingBatchSamplerDataCollatorForSeq2Seq
+from axolotl.utils.logging import get_logger
 from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
 from axolotl.utils.trainer import process_pretraining_datasets_for_packing
 
-LOG = logging.getLogger("axolotl")
+LOG = get_logger(__name__)
 
 
-def encode_pretraining(
-    tokenizer: PreTrainedTokenizerBase, max_tokens: int, examples: List[str]
+def encode_streaming(
+    examples: Dict[str, List],
+    tokenizer: PreTrainedTokenizerBase,
+    max_tokens: int,
+    text_column: str = "text",
+    concatenate: bool = True,
 ) -> Dict[str, List]:
     res = tokenizer(
-        examples,
+        examples[text_column],
         truncation=True,
         max_length=max_tokens - 2,
         add_special_tokens=True,
     )
     # Convert to PyTorch tensors
     input_ids = [torch.tensor(seq) for seq in res["input_ids"]]
+    targets = [torch.tensor(seq) for seq in res["input_ids"]]
     attention_mask = [torch.tensor(seq) for seq in res["attention_mask"]]
+    if not concatenate:
+        return {
+            "input_ids": [seq.tolist() for seq in input_ids],
+            "labels": [seq.tolist() for seq in targets],
+            "attention_mask": [seq.tolist() for seq in attention_mask],
+        }
+
     new_input_ids = []
+    new_labels = []
     new_attention_mask = []
     # Append EOS and PAD tokens to input_ids, and correct attention_mask
     for i, _ in enumerate(input_ids):
@@ -40,22 +53,34 @@ def encode_pretraining(
             ),
             dim=0,
         )
+        targets[i] = torch.cat(
+            (
+                targets[i],
+                torch.tensor([tokenizer.eos_token_id, -100]),
+            ),
+            dim=0,
+        )
         attention_mask[i] = torch.cat((attention_mask[i], torch.tensor([1, 0])), dim=0)
 
     # Concatenate tokens so that their lengths are less than max_tokens
     buffer_input_ids = torch.tensor([], dtype=torch.long)
+    buffer_labels = torch.tensor([], dtype=torch.long)
     buffer_attention_mask = torch.tensor([], dtype=torch.long)
 
-    for ids, mask in zip(input_ids, attention_mask):
+    for ids, labels, mask in zip(input_ids, targets, attention_mask, strict=False):
         if buffer_input_ids.numel() == max_tokens:
             new_input_ids.append(buffer_input_ids)
+            new_labels.append(buffer_labels)
             new_attention_mask.append(buffer_attention_mask)
             buffer_input_ids = torch.tensor([], dtype=torch.long)
+            buffer_labels = torch.tensor([], dtype=torch.long)
             buffer_attention_mask = torch.tensor([], dtype=torch.long)
             buffer_input_ids = torch.cat((buffer_input_ids, ids), dim=0)
+            buffer_labels = torch.cat((buffer_labels, labels), dim=0)
             buffer_attention_mask = torch.cat((buffer_attention_mask, mask), dim=0)
         elif buffer_input_ids.numel() + ids.numel() <= max_tokens:
             buffer_input_ids = torch.cat((buffer_input_ids, ids), dim=0)
+            buffer_labels = torch.cat((buffer_labels, labels), dim=0)
             buffer_attention_mask = torch.cat((buffer_attention_mask, mask), dim=0)
         else:
             buffer_input_ids = torch.cat(
@@ -69,6 +94,17 @@ def encode_pretraining(
                 ),
                 dim=0,
             )
+            buffer_labels = torch.cat(
+                (
+                    buffer_labels,
+                    torch.full(
+                        (max_tokens - buffer_labels.numel(),),
+                        -100,
+                        dtype=torch.long,
+                    ),
+                ),
+                dim=0,
+            )
             buffer_attention_mask = torch.cat(
                 (
                     buffer_attention_mask,
@@ -81,11 +117,14 @@ def encode_pretraining(
                 dim=0,
             )
             new_input_ids.append(buffer_input_ids)
+            new_labels.append(buffer_labels)
             new_attention_mask.append(buffer_attention_mask)
             buffer_input_ids = torch.tensor([], dtype=torch.long)
+            buffer_labels = torch.tensor([], dtype=torch.long)
             buffer_attention_mask = torch.tensor([], dtype=torch.long)
 
             buffer_input_ids = torch.cat((buffer_input_ids, ids), dim=0)
+            buffer_labels = torch.cat((buffer_labels, labels), dim=0)
             buffer_attention_mask = torch.cat((buffer_attention_mask, mask), dim=0)
 
     if buffer_input_ids.numel() > 0:  # for any leftover tokens
@@ -101,6 +140,17 @@ def encode_pretraining(
                 ),
                 dim=0,
             )
+            buffer_labels = torch.cat(
+                (
+                    buffer_labels,
+                    torch.full(
+                        (max_tokens - buffer_labels.numel(),),
+                        -100,
+                        dtype=torch.long,
+                    ),
+                ),
+                dim=0,
+            )
             buffer_attention_mask = torch.cat(
                 (
                     buffer_attention_mask,
@@ -113,11 +163,12 @@ def encode_pretraining(
                 dim=0,
             )
         new_input_ids.append(buffer_input_ids)
+        new_labels.append(buffer_labels)
         new_attention_mask.append(buffer_attention_mask)
 
     ret = {
         "input_ids": [seq.tolist() for seq in new_input_ids],
-        "labels": [seq.tolist() for seq in new_input_ids],
+        "labels": [seq.tolist() for seq in new_labels],
         "attention_mask": [seq.tolist() for seq in new_attention_mask],
     }
 
@@ -125,39 +176,58 @@ def encode_pretraining(
     return ret
 
 
-def wrap_pretraining_dataset(
+def wrap_streaming_dataset(
     dataset,
     tokenizer,
     cfg,
     ds_wrapper_fn,
-    max_tokens=2048,
-    batch_size=1,
-    seed=42,
-    buffer_size=10_000,
 ):
     if cfg.sample_packing:
+        # For SFT (non-pretraining) datasets, always use multipack_attn=True to ensure
+        # attention isolation between packed sequences
+        multipack_attn = (
+            True if not cfg.pretraining_dataset else cfg.pretrain_multipack_attn
+        )
+
         collate_fn = PretrainingBatchSamplerDataCollatorForSeq2Seq(
             tokenizer,
             return_tensors="pt",
             padding=True,
-            pad_to_multiple_of=max_tokens * batch_size,
-            multipack_attn=cfg.pretrain_multipack_attn,
+            pad_to_multiple_of=cfg.sequence_len,
+            multipack_attn=multipack_attn,
         )
         encode = functools.partial(
-            encode_packed_pretraining,
+            encode_packed_streaming,
             collate_fn,
             ds_wrapper_fn,
-            max_seq_length=max_tokens,
-            batch_size=batch_size,
-            multipack_attn=cfg.pretrain_multipack_attn,
+            max_seq_length=cfg.sequence_len,
+            batch_size=cfg.micro_batch_size,
+            multipack_attn=multipack_attn,
+            bin_size=cfg.sample_packing_bin_size,
         )
-        # set this to 1 so downstream data_loader doesn't try to increase the batch again
+
+        # Set this to 1 so downstream data_loader doesn't try to increase the batch size
+        # again
         cfg.micro_batch_size = 1
     else:
-        encode = functools.partial(encode_pretraining, tokenizer, max_tokens)
+        # NOTE: This is not reachable for SFT datasets since we use the pre-existing
+        # loading function for non-packed streaming datasets. Refer to
+        # _prepare_streaming_datasets in sft.py for that code path.
+        text_column = (
+            getattr(cfg.pretraining_dataset[0], "text_column", "text") or "text"
+        )
+        encode = functools.partial(
+            encode_streaming,
+            tokenizer=tokenizer,
+            max_tokens=cfg.sequence_len,
+            text_column=text_column,
+            concatenate=cfg.pretraining_sample_concatenation is True,
+        )
 
     if cfg.shuffle_merged_datasets:
-        dataset = dataset.shuffle(seed=seed, buffer_size=buffer_size)
+        dataset = dataset.shuffle(
+            seed=cfg.seed, buffer_size=cfg.streaming_multipack_buffer_size
+        )
     else:
         LOG.debug("NOT shuffling merged pretraining datasets")
 
@@ -167,46 +237,50 @@ def wrap_pretraining_dataset(
     remove_columns = []
     if dataset.features is None:
         for first_row in dataset:
-            remove_columns = first_row.keys()
+            remove_columns = list(first_row.keys())
             break
     else:
-        remove_columns = dataset.features.keys()
+        remove_columns = list(dataset.features.keys())
 
     dataset = dataset.map(
         encode,
         batched=True,
-        batch_size=buffer_size,
-        # input_columns="text",
+        batch_size=cfg.streaming_multipack_buffer_size,
         remove_columns=remove_columns,
     )
     return dataset
 
 
-def encode_packed_pretraining(
+def encode_packed_streaming(
     collate_fn,
     ds_wrapper: Callable,
     examples: Dict[str, List],
+    bin_size: int,
     max_seq_length: int = 2048,
     batch_size: int = 4,
-    multipack_attn: Optional[bool] = False,
+    multipack_attn: Optional[bool] = True,
 ) -> Dict[str, List]:
-    # pylint: disable=duplicate-code
     # tokenize all the examples
     # rows get split with stride (overlap)
-    train_dataset = ds_wrapper(Dataset.from_dict(examples))[0]
+    train_dataset = ds_wrapper(dataset=Dataset.from_dict(examples))[0]
 
     train_dataset = process_pretraining_datasets_for_packing(
         train_dataset,
         max_seq_length,
         skip_position_ids=not multipack_attn,
+        # FIXME using attention mask unpad/pad with trainer and packed pretraining is broken atm
+        # workaround by using the position id logic for now in trainer
+        drop_attention_mask=multipack_attn,
     )
 
     sampler = MultipackBatchSampler(
-        RandomSampler(train_dataset),
+        sampler=RandomSampler(train_dataset),
+        lengths=get_dataset_lengths(train_dataset),
         batch_size=1,
-        drop_last=True,
         batch_max_len=batch_size * max_seq_length,
-        lengths=get_dataset_lengths(train_dataset),
+        drop_last=True,
+        num_processes=1,
+        bin_size=bin_size,
     )
 
     chunked_data = defaultdict(list)
@@ -214,8 +288,6 @@ def encode_packed_pretraining(
     for batch in sampler:
         for data in batch:
             features = train_dataset[data]
-            if "num_truncated_tokens" in features:
-                del features["num_truncated_tokens"]
             if "num_truncated_tokens" in features:
                 del features["num_truncated_tokens"]
             if "overflow_to_sample_mapping" in features:
diff --git a/src/axolotl/utils/data/utils.py b/src/axolotl/utils/data/utils.py
index e05701e7b0..e141713e79 100644
--- a/src/axolotl/utils/data/utils.py
+++ b/src/axolotl/utils/data/utils.py
@@ -1,10 +1,362 @@
-"""data handling helpers"""
+"""Data handling helpers"""
 
+import contextlib
+import functools
 import hashlib
+import time
+from enum import Enum
+from typing import Callable
+
+import huggingface_hub
+import numpy as np
+import requests
+from datasets import Dataset, IterableDataset
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+from axolotl.utils.samplers.utils import get_dataset_lengths
+from axolotl.utils.trainer import filter_sequences_by_length
+
+LOG = get_logger(__name__)
+
+
+class RetryStrategy(Enum):
+    """Enum for retry strategies."""
+
+    CONSTANT = 1
+    LINEAR = 2
+    EXPONENTIAL = 3
+
+
+def retry_on_request_exceptions(
+    max_retries=3, delay=1, retry_strategy: RetryStrategy = RetryStrategy.LINEAR
+) -> Callable:
+    """Decorator that retries function calls on specific request exceptions.
+
+    Args:
+        max_retries: Maximum number of retry attempts.
+        delay: Base delay between retries in seconds.
+        retry_strategy: Strategy for calculating retry delays.
+
+    Returns:
+        Decorated function with retry logic.
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            for attempt in range(max_retries):
+                try:
+                    return func(*args, **kwargs)
+                except (
+                    requests.exceptions.ReadTimeout,
+                    requests.exceptions.ConnectionError,
+                    requests.exceptions.HTTPError,
+                    huggingface_hub.errors.HfHubHTTPError,
+                ) as exc:
+                    if attempt < max_retries - 1:
+                        if retry_strategy == RetryStrategy.EXPONENTIAL:
+                            step_delay = delay * 2**attempt
+                        elif retry_strategy == RetryStrategy.LINEAR:
+                            step_delay = delay * (attempt + 1)
+                        else:
+                            step_delay = delay  # Use constant delay.
+                        time.sleep(step_delay)
+                    else:
+                        raise exc
+
+        return wrapper
+
+    return decorator
 
 
 def md5(to_hash: str, encoding: str = "utf-8") -> str:
+    """Generate MD5 hash of a string."""
     try:
         return hashlib.md5(to_hash.encode(encoding), usedforsecurity=False).hexdigest()
     except TypeError:
         return hashlib.md5(to_hash.encode(encoding)).hexdigest()  # nosec
+
+
+def sha256(to_hash: str, encoding: str = "utf-8") -> str:
+    """Generate SHA256 hash of a string."""
+    return hashlib.sha256(to_hash.encode(encoding)).hexdigest()
+
+
+def _deduplicate_dataset(
+    dataset: Dataset,
+    seen_hashes: set[str] | None = None,
+) -> tuple[Dataset, set[str]]:
+    """Remove duplicate rows from a dataset using SHA256 hashes.
+
+    Args:
+        dataset: Dataset to deduplicate.
+        seen_hashes: Set of previously seen row hashes (for cross-deduplication).
+
+    Returns:
+        Tuple of deduplicated dataset and the set of seen hashes.
+    """
+    if seen_hashes is None:
+        seen_hashes = set()
+
+    unique_indices = []
+    for idx, row in enumerate(dataset):
+        row_hash = sha256(str(row))  # Using SHA256 for collision resistance
+        if row_hash not in seen_hashes:
+            seen_hashes.add(row_hash)
+            unique_indices.append(idx)
+
+    return dataset.select(unique_indices), seen_hashes
+
+
+def deduplicate_and_log_datasets(
+    dataset: Dataset,
+    other_dataset: Dataset | None = None,
+    dataset_name: str | None = "train",
+    other_name: str | None = "eval",
+) -> tuple[Dataset, Dataset | None]:
+    """Deduplicate datasets, with optional cross-dataset deduplication.
+
+    Args:
+        dataset: Primary dataset to deduplicate.
+        other_dataset: Optional second dataset to deduplicate against the first.
+        dataset_name: Name for the primary dataset (for logging).
+        other_name: Name for the second dataset (for logging).
+
+    Returns:
+        Tuple of (deduplicated_dataset, deduplicated_other_dataset).
+    """
+    # Deduplicate primary dataset
+    LOG.info(
+        f"Starting deduplication for {dataset_name} dataset. Original size: {len(dataset)}"
+    )
+    dataset, seen_rows = _deduplicate_dataset(dataset)
+    LOG.info(
+        f"Deduplication complete for {dataset_name} dataset. New size: {len(dataset)}"
+    )
+
+    # Deduplicate second dataset if provided
+    if other_dataset is not None:
+        LOG.info(
+            f"Starting deduplication for {other_name} dataset. Original size: {len(other_dataset)}"
+        )
+        other_dataset, _ = _deduplicate_dataset(other_dataset, seen_rows)
+        LOG.info(
+            f"Deduplication complete for {other_name} dataset. New size: {len(other_dataset)}"
+        )
+
+    return dataset, other_dataset
+
+
+def keep_min_len(sample, min_sequence_len=2):
+    """
+    Batched filter function that keeps only samples with sequence length >= min_sequence_len.
+    Returns a list of booleans indicating which samples to keep.
+    """
+    min_sequence_len = min_sequence_len or 2
+
+    input_ids = sample["input_ids"]
+
+    # Batched (input_ids is a list of lists)
+    results = []
+    for seq in input_ids:
+        results.append(len(seq) >= min_sequence_len)
+    return results
+
+
+def truncate_long_seq(sample, sequence_len=2048):
+    """
+    Truncate samples whose sequence length is too long (> sequence_len).
+    Modifies the sample in-place and returns the modified sample.
+    """
+    input_ids = sample["input_ids"]
+
+    # Batched (input_ids is a list of lists)
+    for i, seq in enumerate(input_ids):
+        length = len(seq)
+        if length > sequence_len:
+            sample["input_ids"][i] = seq[:sequence_len]
+            if "attention_mask" in sample:
+                sample["attention_mask"][i] = sample["attention_mask"][i][:sequence_len]
+            if "labels" in sample:
+                sample["labels"][i] = sample["labels"][i][:sequence_len]
+            if "position_ids" in sample:
+                sample["position_ids"][i] = sample["position_ids"][i][:sequence_len]
+    return sample
+
+
+def _should_skip_processing(dataset: Dataset) -> bool:
+    """Check if dataset should skip long sequence handling."""
+    if (
+        hasattr(dataset, "column_names")
+        and dataset.column_names
+        and "input_ids" not in dataset.column_names
+    ):
+        LOG.warning(
+            "Dataset does not contain 'input_ids' column. Skip drop long seq. This is "
+            "expected for reward modeling."
+        )
+        return True
+    elif not hasattr(dataset, "column_names") or dataset.column_names is None:
+        LOG.info(
+            "Dataset is streaming (IterableDataset), skipping long sequence handling"
+        )
+        return True
+    return False
+
+
+def _log_dataset_stats(dataset: Dataset) -> None:
+    """Log min/max sequence lengths for debugging."""
+    with contextlib.suppress(AttributeError, ValueError):
+        ds_lengths = get_dataset_lengths(dataset, from_arrow=True)
+        LOG.info(f"min_input_len: {np.min(ds_lengths)}")
+        LOG.info(f"max_input_len: {np.max(ds_lengths)}")
+
+
+def _build_filter_kwargs(dataset: Dataset, cfg: DictDefault) -> dict:
+    """Build kwargs for dataset filter/map operations."""
+    kwargs = {}
+    if not isinstance(dataset, IterableDataset):
+        kwargs["num_proc"] = cfg.dataset_num_proc
+        kwargs["load_from_cache_file"] = not cfg.is_preprocess
+    return kwargs
+
+
+def _filter_short_sequences(
+    dataset: Dataset, min_len: int, filter_kwargs: dict
+) -> tuple[Dataset, int]:
+    """Filter out sequences shorter than min_len. Returns (dataset, num_dropped)."""
+    prior_len = len(dataset) if hasattr(dataset, "__len__") else None
+
+    desc_kwargs = {}
+    if filter_kwargs:
+        desc_kwargs["desc"] = f"Filtering Short Sequences (<{min_len})"
+
+    dataset = dataset.filter(
+        functools.partial(keep_min_len, min_sequence_len=min_len),
+        batched=True,
+        **filter_kwargs,
+        **desc_kwargs,
+    )
+
+    dropped = 0
+    if prior_len:
+        dropped = prior_len - len(dataset)
+        if dropped > 0:
+            LOG.info(f"Dropped {dropped} short sequences (<{min_len} tokens)")
+
+    return dataset, dropped
+
+
+def _truncate_long_sequences(
+    dataset: Dataset, max_len: int, map_kwargs: dict
+) -> Dataset:
+    """Truncate sequences longer than max_len."""
+    desc_kwargs = {}
+    if map_kwargs:
+        desc_kwargs["desc"] = f"Truncating Sequences (target_len={max_len})"
+
+    dataset = dataset.map(
+        functools.partial(truncate_long_seq, sequence_len=max_len),
+        batched=True,
+        **map_kwargs,
+        **desc_kwargs,
+    )
+    LOG.info(f"Truncated long sequences to max length {max_len}")
+    return dataset
+
+
+def _drop_outside_range(
+    dataset: Dataset,
+    max_len: int,
+    min_len: int,
+    raise_on_long: bool,
+    filter_kwargs: dict,
+) -> tuple[Dataset, int]:
+    """Drop sequences outside valid length range [min_len, max_len].
+
+    Returns (dataset, num_dropped)."""
+    prior_len = len(dataset) if hasattr(dataset, "__len__") else None
+
+    desc_kwargs = {}
+    if filter_kwargs:
+        action = (
+            "Checking Sequence Lengths"
+            if raise_on_long
+            else "Dropping Invalid Sequences"
+        )
+        desc_kwargs["desc"] = f"{action} (<{min_len} or >{max_len})"
+
+    dataset = dataset.filter(
+        functools.partial(
+            filter_sequences_by_length,
+            sequence_len=max_len,
+            min_sequence_len=min_len,
+            raise_on_drop=raise_on_long,
+        ),
+        batched=True,
+        **filter_kwargs,
+        **desc_kwargs,
+    )
+
+    dropped = 0
+    if not raise_on_long and prior_len:
+        dropped = prior_len - len(dataset)
+        if dropped > 0:
+            LOG.info(
+                f"Dropped {dropped} sequences outside valid range "
+                f"([{min_len}, {max_len}])"
+            )
+
+    return dataset, dropped
+
+
+def handle_long_seq_in_dataset(
+    dataset: Dataset, sequence_len: int, cfg: DictDefault
+) -> Dataset:
+    """Remove sequences longer than configured maximum from dataset.
+
+    Args:
+        dataset: Dataset to filter.
+        sequence_len: Maximum length for sequences to keep
+        cfg: Dictionary mapping `axolotl` config keys to values.
+
+    Returns:
+        Filtered dataset with long sequences handled according to the excess_length_strategy value:
+            'drop' (default)    excludes any sequence longer than sequence_len
+            'truncate'          truncates them down to sequence_len
+            'raise'             raises a ValueError if any sequence was found that was longer than sequence_len
+    """
+    # Early returns for special cases
+    if _should_skip_processing(dataset):
+        return dataset
+
+    excess_length_strategy = (cfg.excess_length_strategy or "drop").lower()
+
+    _log_dataset_stats(dataset)
+
+    # Setup kwargs
+    filter_kwargs = _build_filter_kwargs(dataset, cfg)
+
+    # Handle sequences based on strategy
+    if excess_length_strategy == "truncate":
+        dataset, _ = _filter_short_sequences(dataset, cfg.min_sample_len, filter_kwargs)
+        dataset = _truncate_long_sequences(dataset, sequence_len, filter_kwargs)
+    else:
+        raise_on_long = excess_length_strategy == "raise"
+        dataset, _ = _drop_outside_range(
+            dataset, sequence_len, cfg.min_sample_len, raise_on_long, filter_kwargs
+        )
+
+    return dataset
+
+
+def remove_double_bos_token(example: dict[str, list], bos_token_id: int):
+    """Remove double bos tokens that may occur when retokenizing preprocessed data
+    for tokenizers and chat templates that have a bos_token - eg. DPO + Llama.
+    """
+    input_ids = example["input_ids"]
+    if len(input_ids) >= 2 and input_ids[0] == input_ids[1] == bos_token_id:
+        for key in example:
+            example[key] = example[key][1:]
+    return example
diff --git a/src/axolotl/utils/data/wrappers.py b/src/axolotl/utils/data/wrappers.py
new file mode 100644
index 0000000000..30b9935e25
--- /dev/null
+++ b/src/axolotl/utils/data/wrappers.py
@@ -0,0 +1,424 @@
+"""Data handling specific to SFT."""
+
+from typing import Any, NoReturn, cast
+
+from datasets import (
+    Dataset,
+    IterableDataset,
+    Sequence,
+    Value,
+)
+from transformers import PreTrainedTokenizer
+from transformers.processing_utils import ProcessorMixin
+
+from axolotl.datasets import TokenizedPromptDataset, wrap_dataset_for_tokenized_prompt
+from axolotl.prompt_strategies import load
+from axolotl.prompt_strategies.bradley_terry import load as bradley_terry_load
+from axolotl.prompt_tokenizers import (
+    AlpacaMultipleChoicePromptTokenizingStrategy,
+    AlpacaPromptTokenizingStrategy,
+    AlpacaReflectionPTStrategy,
+    DatasetWrappingStrategy,
+    GPTeacherPromptTokenizingStrategy,
+    JeopardyPromptTokenizingStrategy,
+    OpenAssistantPromptTokenizingStrategy,
+    PromptTokenizingStrategy,
+    SummarizeTLDRPromptTokenizingStrategy,
+)
+from axolotl.prompters import (
+    AlpacaPrompter,
+    GPTeacherPrompter,
+    JeopardyPrompter,
+    MultipleChoiceConcisePrompter,
+    MultipleChoiceExplainPrompter,
+    Prompter,
+    ReflectAlpacaPrompter,
+    SummarizeTLDRPrompter,
+    UnsupportedPrompter,
+)
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def handle_unknown_dataset_strategy(dataset_config: DictDefault) -> NoReturn:
+    """Raise error for unknown dataset strategy."""
+    ds_type = dataset_config.type
+    suffix = ""
+    if ":load_" in ds_type:
+        suffix = f"Did you mean {ds_type.replace(':load_', '.load_')}?"
+
+    error_message = f"unhandled prompt tokenization strategy: {ds_type}. {suffix}"
+    LOG.error(error_message)
+    raise ValueError(error_message)
+
+
+def get_dataset_wrapper(
+    dataset_config: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset_base_type: str | None,
+    dataset: Dataset | IterableDataset,
+    dataset_prompt_style: str | None = None,
+    processor: ProcessorMixin | None = None,
+) -> tuple[Dataset | IterableDataset, Prompter | None]:
+    """Create an appropriate dataset wrapper and prompter based on dataset
+    configuration.
+
+    Args:
+        dataset_config: Configuration for the dataset.
+        tokenizer: Tokenizer to use for processing text.
+        cfg: Global configuration object.
+        dataset_base_type: The base type of the dataset.
+        dataset: The actual dataset object.
+        dataset_prompt_style: Optional prompt style specification.
+        processor: Optional processor for multimodal datasets.
+
+    Returns:
+        tuple of (dataset_wrapper, dataset_prompter).
+    """
+    # Common parameters for dataset wrapping
+    dataset_kwargs: dict[str, Any] = {
+        "process_count": cfg.dataset_num_proc,
+        "keep_in_memory": cfg.dataset_keep_in_memory is True,
+    }
+
+    LOG.info(
+        f"Loading dataset: {dataset_config['path']} with base_type: "
+        f"{dataset_base_type} and prompt_style: {dataset_prompt_style}"
+    )
+
+    # Dataset is already tokenized
+    if _is_dataset_already_tokenized(dataset):
+        return dataset, UnsupportedPrompter()
+
+    # Custom dataset type definition
+    if isinstance(dataset_config.type, DictDefault):
+        return _handle_custom_dataset_type(
+            dataset_config, tokenizer, cfg, dataset, dataset_kwargs
+        )
+
+    # Skip preparation if configured
+    if cfg.skip_prepare_dataset:
+        return dataset, None
+
+    # Bradley-Terry dataset
+    if dataset_config.type.startswith("bradley_terry"):
+        return _handle_bradley_terry_dataset(
+            dataset_config, tokenizer, cfg, dataset, dataset_kwargs
+        )
+
+    # Stepwise supervised dataset
+    if dataset_config.type.startswith("stepwise_supervised"):
+        return _handle_stepwise_supervised_dataset(
+            dataset_config, tokenizer, cfg, dataset, dataset_kwargs
+        )
+
+    # Try to load prompt tokenizer / dataset wrapper strategy from registry
+    dataset_strategy = load(
+        dataset_config.type, tokenizer, cfg, dataset_config, processor=processor
+    )
+    if dataset_strategy:
+        return _handle_loaded_strategy(dataset_strategy, dataset, dataset_kwargs)
+
+    # Known dataset types with specific handling
+    if dataset_base_type in DATASET_HANDLERS:
+        handler = DATASET_HANDLERS[dataset_base_type]
+        return handler(dataset_prompt_style, tokenizer, cfg, dataset, dataset_kwargs)
+
+    # Unhandled dataset type
+    handle_unknown_dataset_strategy(dataset_config)
+
+
+def _is_dataset_already_tokenized(dataset: Dataset | IterableDataset) -> bool:
+    """Check if the dataset is already tokenized."""
+    return (
+        isinstance(dataset, Dataset)
+        and "input_ids" in dataset.features
+        and "attention_mask" in dataset.features
+        and "labels" in dataset.features
+    )
+
+
+def _handle_custom_dataset_type(
+    dataset_config: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle a custom dataset type defined in the configuration."""
+    dataset_strategy = cast(
+        PromptTokenizingStrategy,
+        load("user_defined", tokenizer, cfg, dataset_config.type.to_dict()),
+    )
+    dataset_prompter = UnsupportedPrompter()
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_bradley_terry_dataset(
+    dataset_config: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter | None]:
+    """Handle a Bradley-Terry dataset."""
+    bt_type = dataset_config.type.split(".", 1)[1]
+    dataset_strategy = bradley_terry_load(bt_type, tokenizer, cfg, dataset_config)
+
+    if not dataset_strategy:
+        handle_unknown_dataset_strategy(dataset_config)
+
+    dataset_prompter = UnsupportedPrompter()
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_stepwise_supervised_dataset(
+    dataset_config: DictDefault,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle a stepwise supervised dataset."""
+    dataset_prompter = UnsupportedPrompter()
+    dataset_strategy = load(dataset_config.type, tokenizer, cfg, dataset_config)
+
+    # We need to explicitly cast boolean labels to int
+    # for compatibility with how trl's PRMTrainer works
+    if isinstance(dataset, Dataset):
+        dataset = dataset.cast_column("labels", Sequence(Value("int64")))
+
+    dataset_wrapper = TokenizedPromptDataset(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_loaded_strategy(
+    dataset_strategy: PromptTokenizingStrategy | DatasetWrappingStrategy,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter | None]:
+    """Handle a dataset with a strategy loaded from the registry."""
+    if isinstance(dataset_strategy, DatasetWrappingStrategy):
+        return dataset_strategy.wrap_dataset(dataset, **dataset_kwargs), None
+
+    dataset_prompter = UnsupportedPrompter()
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_alpaca_dataset(
+    dataset_prompt_style: str | None,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle an Alpaca dataset."""
+    dataset_prompter = AlpacaPrompter(dataset_prompt_style)
+    dataset_strategy = AlpacaPromptTokenizingStrategy(
+        dataset_prompter,
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_explainchoice_dataset(
+    dataset_prompt_style: str | None,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle an ExplainChoice dataset."""
+    dataset_prompter = MultipleChoiceExplainPrompter(dataset_prompt_style)
+    dataset_strategy = AlpacaMultipleChoicePromptTokenizingStrategy(
+        dataset_prompter,
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_concisechoice_dataset(
+    dataset_prompt_style: str | None,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle a ConciseChoice dataset."""
+    dataset_prompter = MultipleChoiceConcisePrompter(dataset_prompt_style)
+    dataset_strategy = AlpacaMultipleChoicePromptTokenizingStrategy(
+        dataset_prompter,
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_summarizetldr_dataset(
+    dataset_prompt_style: str | None,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle a SummarizeTLDR dataset."""
+    dataset_prompter = SummarizeTLDRPrompter(dataset_prompt_style)
+    dataset_strategy = SummarizeTLDRPromptTokenizingStrategy(
+        dataset_prompter,
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_jeopardy_dataset(
+    dataset_prompt_style: str | None,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle a Jeopardy dataset."""
+    dataset_prompter = JeopardyPrompter(dataset_prompt_style)
+    dataset_strategy = JeopardyPromptTokenizingStrategy(
+        dataset_prompter,
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_oasst_dataset(
+    dataset_prompt_style: str | None,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle an OpenAssistant dataset."""
+    dataset_prompter = AlpacaPrompter(dataset_prompt_style)
+    dataset_strategy = OpenAssistantPromptTokenizingStrategy(
+        dataset_prompter,
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_gpteacher_dataset(
+    dataset_prompt_style: str | None,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle a GPTeacher dataset."""
+    dataset_prompter = GPTeacherPrompter(dataset_prompt_style)
+    dataset_strategy = GPTeacherPromptTokenizingStrategy(
+        dataset_prompter,
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+def _handle_reflection_dataset(
+    dataset_prompt_style: str | None,
+    tokenizer: PreTrainedTokenizer,
+    cfg: DictDefault,
+    dataset: Dataset | IterableDataset,
+    dataset_kwargs: dict[str, Any],
+) -> tuple[Dataset | IterableDataset, Prompter]:
+    """Handle a Reflection dataset."""
+    dataset_prompter = ReflectAlpacaPrompter(dataset_prompt_style)
+    dataset_strategy = AlpacaReflectionPTStrategy(
+        dataset_prompter,
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+    dataset_wrapper = wrap_dataset_for_tokenized_prompt(
+        dataset_strategy,
+        dataset,
+        **dataset_kwargs,
+    )
+    return dataset_wrapper, dataset_prompter
+
+
+DATASET_HANDLERS = {
+    "alpaca": _handle_alpaca_dataset,
+    "explainchoice": _handle_explainchoice_dataset,
+    "concisechoice": _handle_concisechoice_dataset,
+    "summarizetldr": _handle_summarizetldr_dataset,
+    "jeopardy": _handle_jeopardy_dataset,
+    "oasst": _handle_oasst_dataset,
+    "gpteacher": _handle_gpteacher_dataset,
+    "reflection": _handle_reflection_dataset,
+}
diff --git a/src/axolotl/utils/datasets.py b/src/axolotl/utils/datasets.py
new file mode 100644
index 0000000000..7beeb27333
--- /dev/null
+++ b/src/axolotl/utils/datasets.py
@@ -0,0 +1,21 @@
+"""helper functions for datasets"""
+
+import os
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def get_default_process_count():
+    if axolotl_dataset_num_proc := os.environ.get("AXOLOTL_DATASET_NUM_PROC"):
+        return int(axolotl_dataset_num_proc)
+    if axolotl_dataset_processes := os.environ.get("AXOLOTL_DATASET_PROCESSES"):
+        LOG.warning(
+            "AXOLOTL_DATASET_PROCESSES and `dataset_processes` are deprecated and will be "
+            "removed in a future version. Please use `dataset_num_proc` instead."
+        )
+        return int(axolotl_dataset_processes)
+    if runpod_cpu_count := os.environ.get("RUNPOD_CPU_COUNT"):
+        return int(runpod_cpu_count)
+    return os.cpu_count()
diff --git a/src/axolotl/utils/dict.py b/src/axolotl/utils/dict.py
index 409d088e6d..7d146c7a91 100644
--- a/src/axolotl/utils/dict.py
+++ b/src/axolotl/utils/dict.py
@@ -13,3 +13,39 @@ def __missing__(self, key):
 
     def __or__(self, other):
         return DictDefault(super().__ror__(other))
+
+    def __setitem__(self, name, value):
+        # workaround for pickle/unpickle issues and __frozen not being available
+        try:
+            isFrozen = hasattr(self, "__frozen") and object.__getattribute__(
+                self, "__frozen"
+            )
+        except AttributeError:
+            isFrozen = False
+
+        if isFrozen and name not in super().keys():
+            raise KeyError(name)
+        super(Dict, self).__setitem__(name, value)
+        try:
+            p = object.__getattribute__(self, "__parent")
+            key = object.__getattribute__(self, "__key")
+        except AttributeError:
+            p = None
+            key = None
+        if p is not None:
+            p[key] = self
+            object.__delattr__(self, "__parent")
+            object.__delattr__(self, "__key")
+
+
+def remove_none_values(obj):
+    """
+    Remove null from a dictionary-like obj or list.
+    These can appear due to Dataset loading causing schema merge.
+    See https://github.com/axolotl-ai-cloud/axolotl/pull/2909
+    """
+    if hasattr(obj, "items"):
+        return {k: remove_none_values(v) for k, v in obj.items() if v is not None}
+    if isinstance(obj, list):
+        return [remove_none_values(elem) for elem in obj]
+    return obj
diff --git a/src/axolotl/utils/distributed.py b/src/axolotl/utils/distributed.py
index ecb1bcc9ec..200c802a69 100644
--- a/src/axolotl/utils/distributed.py
+++ b/src/axolotl/utils/distributed.py
@@ -1,6 +1,5 @@
-"""
-utility helpers for distributed checks
-"""
+"""Utilities for distributed functionality."""
+
 import os
 import pickle  # nosec
 from contextlib import contextmanager
@@ -9,18 +8,65 @@
 import torch
 import torch.distributed as dist
 from accelerate import PartialState
+from accelerate.utils import ParallelismConfig
+from transformers.utils.import_utils import (
+    is_torch_cuda_available,
+    is_torch_mps_available,
+    is_torch_npu_available,
+)
+
+distributed_state = None
+
+
+def get_device_type() -> torch.device:
+    device = torch.device("cpu")
+    if is_torch_cuda_available():
+        device = torch.device("cuda")
+    elif is_torch_mps_available():
+        device = torch.device("mps")
+    elif is_torch_npu_available():
+        device = torch.device("npu")
+    return device
+
+
+def get_device_count() -> int:
+    cur_device = get_device_type()
+    if "cuda" in str(cur_device):
+        return torch.cuda.device_count()
+    if "npu" in str(cur_device):
+        return torch.npu.device_count()
+    return 1
+
+
+def get_current_device() -> int:
+    cur_device = get_device_type()
+    if "cuda" in str(cur_device):
+        return torch.cuda.current_device()
+    if "npu" in str(cur_device):
+        return torch.npu.current_device()
+    return 0
+
+
+def init_distributed_state():
+    global distributed_state
+    if distributed_state is None:
+        timeout = int(os.environ.get("AXOLOTL_NCCL_TIMEOUT", 1800))
+        try:
+            distributed_state = PartialState(timeout=timedelta(seconds=timeout))
+        except ValueError:
+            pass
 
-distributed_state = None  # pylint: disable=invalid-name
 
+def get_distributed_state() -> PartialState | None:
+    return distributed_state
 
-def is_distributed():
-    """
-    Check if distributed training is initialized.
-    """
-    global distributed_state  # pylint: disable=global-statement
-    if not distributed_state:
-        timeout = int(os.environ.get("AXOLOTL_NCCL_TIMEOUT", 1800))
-        distributed_state = PartialState(timeout=timedelta(seconds=timeout))
+
+def is_distributed() -> bool:
+    """Check if distributed training is initialized."""
+    init_distributed_state()
+
+    if distributed_state is None:
+        return False
 
     return distributed_state.use_distributed and distributed_state.initialized
 
@@ -34,33 +80,53 @@ def barrier():
         dist.barrier()
 
 
-def is_main_process():
+def is_main_process() -> bool:
     """
-    Check if the current process is the main process.
-    If not in distributed mode, always return True.
+    Check if the current process is the main process. If not in distributed mode,
+    always return `True`.
+
+    We use a simpler logic when the distributed state is not initialized: we just log
+    on the 0-th local rank.
+
+    Returns:
+        `True` if the current process is the main process, `False` otherwise.
     """
+    if get_distributed_state() is None:
+        return os.environ.get("LOCAL_RANK", "0") == "0"
     if not is_distributed():
         return True
     return dist.get_rank() == 0
 
 
-def get_world_size():
+def is_local_main_process() -> bool:
+    if get_distributed_state() is None:
+        return os.environ.get("LOCAL_RANK", "0") == "0"
+    return PartialState().is_local_main_process
+
+
+def get_world_size() -> int:
     return int(os.getenv("WORLD_SIZE", "1"))
 
 
-@contextmanager
-def zero_only():
+def cleanup_distributed():
     """
-    Context manager that only runs the enclosed block on the main rank.
+    Destroy process group if torch distributed is initialized. Called in training early
+    termination or when training successfully completes.
     """
-    if is_main_process():
-        yield
-    else:
-        yield None
+    # Ensure that all operations are completed before destroying the process group
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+    if torch.xpu.is_available():
+        torch.xpu.synchronize()
+
+    # Destroy the process group
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
 
 
 @contextmanager
-def zero_first(is_main):
+def zero_first(is_main: bool):
     """
     runs the wrapped context so that rank 0 runs first before other ranks
     """
@@ -71,7 +137,7 @@ def zero_first(is_main):
         barrier()
 
 
-def gather_scalar_from_all_ranks(fn, world_size=1):  # pylint: disable=invalid-name
+def gather_scalar_from_all_ranks(fn, world_size=1):
     """
     Run a callable 'fn' on all ranks and gather the results on the specified rank.
 
@@ -87,7 +153,7 @@ def gather_scalar_from_all_ranks(fn, world_size=1):  # pylint: disable=invalid-n
     if not is_distributed():
         return [value_scalar]
     value_tensor = torch.tensor(
-        value_scalar, device=torch.cuda.current_device()
+        value_scalar, device=f"{get_device_type()}:{get_current_device()}"
     ).float()
 
     if not is_main_process():
@@ -111,13 +177,14 @@ def broadcast_dict(vals: dict):
     if not is_distributed():
         return vals
 
+    cur_device = get_device_type()
     if is_main_process():
         data_byte = pickle.dumps(vals)
-        data_tensor = torch.ByteTensor(list(data_byte)).to("cuda")
-        data_size = torch.IntTensor([len(data_byte)]).to("cuda")
+        data_tensor = torch.ByteTensor(list(data_byte)).to(cur_device)
+        data_size = torch.IntTensor([len(data_byte)]).to(cur_device)
     else:
-        data_tensor = torch.empty([1024], dtype=torch.uint8, device="cuda")
-        data_size = torch.IntTensor([0]).to("cuda")
+        data_tensor = torch.empty([1024], dtype=torch.uint8, device=cur_device)
+        data_size = torch.IntTensor([0]).to(cur_device)
 
     dist.broadcast(data_size, 0)
     if not is_main_process():
@@ -134,7 +201,7 @@ def broadcast_dict(vals: dict):
     return vals
 
 
-def compute_and_broadcast(fn):  # pylint: disable=invalid-name
+def compute_and_broadcast(fn):
     """
     Compute a value using the function 'fn' only on the specified rank (default is 0).
     The value is then broadcasted to all other ranks.
@@ -146,14 +213,15 @@ def compute_and_broadcast(fn):  # pylint: disable=invalid-name
     Returns:
     - The computed value (int or float).
     """
+    cur_device = f"{get_device_type()}:{get_current_device()}"
     if is_main_process():
         value_scalar = fn()
         value_tensor = torch.tensor(
-            value_scalar, device=torch.cuda.current_device()
-        ).float()
+            value_scalar, device=cur_device, dtype=torch.float32
+        )
     else:
         value_tensor = torch.tensor(
-            0.0, device=torch.cuda.current_device()
+            0.0, device=cur_device, dtype=torch.float32
         )  # Placeholder tensor
 
     # Broadcast the tensor to all processes.
@@ -166,7 +234,7 @@ def compute_and_broadcast(fn):  # pylint: disable=invalid-name
     return float(value_tensor.item())
 
 
-def gather_from_all_ranks(fn, world_size=1):  # pylint: disable=invalid-name
+def gather_from_all_ranks(fn, world_size=1):
     """
     Run a callable 'fn' on all ranks and gather the results on the specified rank.
 
@@ -180,7 +248,7 @@ def gather_from_all_ranks(fn, world_size=1):  # pylint: disable=invalid-name
     """
     value_scalar = fn()
     value_tensor = torch.tensor(
-        value_scalar, device=torch.cuda.current_device()
+        value_scalar, device=f"{get_device_type()}:{get_current_device()}"
     ).float()
 
     # Placeholder tensor for gathering results
@@ -226,3 +294,89 @@ def reduce_and_broadcast(fn1, fn2):
     # Use compute_and_broadcast to compute the reduced value on the main process
     # and then broadcast it to all ranks
     return compute_and_broadcast(lambda: fn2(gathered_values))
+
+
+def build_parallelism_config(cfg):
+    pc_kwargs = _get_parallel_config_kwargs(
+        get_world_size(),
+        cfg.tensor_parallel_size,
+        cfg.context_parallel_size,
+        cfg.dp_shard_size,
+        cfg.dp_replicate_size,
+        bool(cfg.fsdp or cfg.fsdp_config),
+        getattr(cfg, "expert_parallel_size", None),
+    )
+
+    if pc_kwargs:
+        parallelism_config = ParallelismConfig(
+            **pc_kwargs,
+        )
+        device_mesh = parallelism_config.build_device_mesh("cuda")
+
+        return parallelism_config, device_mesh
+    return None, None
+
+
+def _get_parallel_config_kwargs(
+    world_size: int,
+    tensor_parallel_size: int = 1,
+    context_parallel_size: int = 1,
+    dp_shard_size: int | None = None,
+    dp_replicate_size: int | None = None,
+    is_fsdp: bool = False,
+    expert_parallel_size: int | None = None,
+):
+    pc_kwargs = {}
+    remaining_world_size = world_size
+
+    # EP consumes part of world_size; subtract it up front so the auto-fill
+    # below doesn't put EP ranks into `dp_replicate_size`.
+    if expert_parallel_size and expert_parallel_size > 1:
+        if remaining_world_size % expert_parallel_size != 0:
+            raise ValueError(
+                f"expert_parallel_size ({expert_parallel_size}) must divide "
+                f"world_size ({world_size})."
+            )
+        remaining_world_size = remaining_world_size // expert_parallel_size
+
+    if tensor_parallel_size and tensor_parallel_size > 1:
+        pc_kwargs["tp_size"] = tensor_parallel_size
+        remaining_world_size = remaining_world_size // tensor_parallel_size
+
+    if context_parallel_size and context_parallel_size > 1:
+        pc_kwargs["cp_size"] = context_parallel_size
+        remaining_world_size = remaining_world_size // context_parallel_size
+
+    if dp_shard_size is None and dp_replicate_size in (None, 1):
+        if remaining_world_size > 1:
+            pc_kwargs["dp_shard_size"] = remaining_world_size
+            remaining_world_size = 1
+
+    if dp_replicate_size and dp_replicate_size > 1:
+        pc_kwargs["dp_replicate_size"] = dp_replicate_size
+        remaining_world_size = remaining_world_size // dp_replicate_size
+
+    if remaining_world_size > 1 and dp_shard_size and dp_shard_size > 1:
+        if not is_fsdp:
+            raise ValueError(
+                "dp_shard_size was configured without a corresponding fsdp_config! "
+                "Please ensure you have configured FSDP using fsdp_config."
+            )
+        pc_kwargs["dp_shard_size"] = dp_shard_size
+        remaining_world_size = remaining_world_size // dp_shard_size
+        if remaining_world_size > 1 and "dp_replicate_size" not in pc_kwargs:
+            pc_kwargs["dp_replicate_size"] = remaining_world_size
+            remaining_world_size = 1
+
+    if remaining_world_size > 1:
+        if "dp_shard_size" not in pc_kwargs and is_fsdp:
+            pc_kwargs["dp_shard_size"] = remaining_world_size
+            remaining_world_size = 1
+
+    if remaining_world_size > 1:
+        raise ValueError(
+            f"The configured parallelisms are incompatible with the current world size ({get_world_size()})!\n"
+            f"{pc_kwargs}"
+        )
+
+    return pc_kwargs
diff --git a/src/axolotl/utils/environment.py b/src/axolotl/utils/environment.py
new file mode 100644
index 0000000000..d5f2d9f780
--- /dev/null
+++ b/src/axolotl/utils/environment.py
@@ -0,0 +1,56 @@
+"""
+utils to get GPU info for the current environment
+"""
+
+import os
+from importlib.metadata import version
+
+import torch
+from accelerate.utils.environment import (
+    check_cuda_p2p_ib_support as accelerate_check_cuda_p2p_ib_support,
+)
+from packaging.version import Version, parse
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def check_cuda_p2p_ib_support():
+    if not accelerate_check_cuda_p2p_ib_support():
+        return False
+    if not check_cuda_p2p_support():
+        return False
+    return True
+
+
+def check_cuda_p2p_support() -> bool:
+    try:
+        world_size = int(os.environ.get("WORLD_SIZE", "1"))
+        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    except ValueError:
+        return True
+
+    if world_size > 1:
+        node_world_size = int(os.environ.get("NODE_WORLD_SIZE", "8"))
+        local_other_rank = (local_rank // node_world_size) * node_world_size
+        local_other_rank += 1 if (local_rank % node_world_size) == 0 else 0
+        try:
+            can_p2p = torch.cuda.can_device_access_peer(local_rank, local_other_rank)
+        except AssertionError as exc:
+            # some sort of logic error in indexing processes, assume p2p is fine for now
+            LOG.warning(exc)
+            return True
+        return can_p2p
+
+    return True
+
+
+def get_package_version(package: str) -> Version:
+    version_str = version(package)
+    return parse(version_str)
+
+
+def is_package_version_ge(package: str, version_: str) -> bool:
+    package_version = get_package_version(package)
+    return package_version >= parse(version_)
diff --git a/src/axolotl/utils/fp32_norms.py b/src/axolotl/utils/fp32_norms.py
new file mode 100644
index 0000000000..67793998f8
--- /dev/null
+++ b/src/axolotl/utils/fp32_norms.py
@@ -0,0 +1,135 @@
+"""Helpers for keeping selected norm modules in fp32 under FSDP2."""
+
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+DEFAULT_FP32_NORM_SUFFIXES: tuple[str, ...] = ("RMSNorm", "LayerNorm")
+
+
+def _matches_norm_class(module: "torch.nn.Module", patterns: Sequence[str]) -> bool:
+    """Match a module against class-name patterns.
+
+    Two matching modes, chosen per-pattern by presence of a dot:
+      - Fully qualified (contains "."): matches f"{module.__module__}.{cls}" exactly.
+      - Suffix (no dot): matches type(module).__name__.endswith(pattern).
+    Empty / whitespace-only patterns are skipped (``cls_name.endswith("")``
+    is True for every class, which would silently match everything).
+    """
+    cls = type(module)
+    cls_name = cls.__name__
+    qualified = f"{cls.__module__}.{cls_name}"
+    for pattern in patterns:
+        if not pattern or not pattern.strip():
+            continue
+        if "." in pattern:
+            if qualified == pattern:
+                return True
+        elif cls_name.endswith(pattern):
+            return True
+    return False
+
+
+def get_fp32_norm_patterns(source) -> list[str] | None:
+    """Resolve configured fp32 norm patterns from a config or tagged model."""
+    tagged_patterns = getattr(source, "_axolotl_fp32_norm_patterns", None)
+    if tagged_patterns is not None:
+        return list(tagged_patterns)
+
+    if not getattr(source, "fp32_norms", False):
+        return None
+
+    configured_patterns = getattr(source, "fp32_norm_classes", None)
+    if configured_patterns:
+        return list(configured_patterns)
+
+    return list(DEFAULT_FP32_NORM_SUFFIXES)
+
+
+def tag_model_fp32_norms(model: "torch.nn.Module", cfg) -> list[str] | None:
+    """Attach the resolved fp32 norm patterns to the model for FSDP2 prepare."""
+    patterns = get_fp32_norm_patterns(cfg)
+    if patterns is None:
+        if hasattr(model, "_axolotl_fp32_norm_patterns"):
+            delattr(model, "_axolotl_fp32_norm_patterns")
+        return None
+
+    model._axolotl_fp32_norm_patterns = list(patterns)
+    return patterns
+
+
+def shard_norms_fp32(
+    model: "torch.nn.Module",
+    source=None,
+    *,
+    patterns: Sequence[str] | None = None,
+    fully_shard_kwargs: dict[str, Any] | None = None,
+) -> int:
+    """Wrap matching norm modules with FSDP2 + fp32 MixedPrecisionPolicy."""
+    if source is not None and not getattr(source, "fp32_norms", False):
+        return 0
+
+    if source is not None and getattr(source, "fsdp_version", None) != 2:
+        raise ValueError(
+            "fp32_norms requires fsdp_version: 2. FSDP1 enforces flat-param "
+            "dtype uniformity within each wrap group, which is incompatible "
+            "with keeping norms in fp32 while the rest of the layer is bf16."
+        )
+
+    patterns = (
+        list(patterns)
+        if patterns is not None
+        else get_fp32_norm_patterns(source or model)
+    )
+    if not patterns:
+        return 0
+
+    from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+
+    outer_policy = (fully_shard_kwargs or {}).get("mp_policy")
+    output_dtype = getattr(outer_policy, "param_dtype", None)
+    fp32_policy = MixedPrecisionPolicy(
+        param_dtype=torch.float32,
+        reduce_dtype=torch.float32,
+        output_dtype=output_dtype,
+    )
+
+    matches = [
+        (name, module)
+        for name, module in model.named_modules()
+        if _matches_norm_class(module, patterns)
+    ]
+
+    if not matches:
+        LOG.warning(
+            "fp32_norms enabled but no modules matched patterns %s. Check "
+            "fp32_norm_classes against the model's actual norm class names.",
+            patterns,
+        )
+        return 0
+
+    shard_kwargs = dict(fully_shard_kwargs or {})
+    shard_kwargs["mp_policy"] = fp32_policy
+
+    for _name, module in matches:
+        for param in module.parameters(recurse=False):
+            param.data = param.data.to(torch.float32)
+        for buffer in module.buffers(recurse=False):
+            if buffer.dtype.is_floating_point:
+                buffer.data = buffer.data.to(torch.float32)
+        fully_shard(module, **shard_kwargs)
+
+    LOG.info(
+        "Sharded %d norm modules with fp32 MixedPrecisionPolicy "
+        "(patterns=%s, output_dtype=%s)",
+        len(matches),
+        patterns,
+        output_dtype,
+    )
+    return len(matches)
diff --git a/src/axolotl/utils/freeze.py b/src/axolotl/utils/freeze.py
index e3d0fd1446..e60c496731 100644
--- a/src/axolotl/utils/freeze.py
+++ b/src/axolotl/utils/freeze.py
@@ -1,13 +1,52 @@
 """
 module to freeze/unfreeze parameters by name
 """
-import logging
+
 import re
 from typing import Callable, List, Tuple, Union
 
 from axolotl.utils.distributed import is_main_process
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+# Top-level module name prefixes that belong to vision/audio/multimodal encoders
+# rather than the language backbone.  These are matched against the first component
+# of each ``named_parameter`` path (e.g. "model.vision_tower." -> "vision_tower").
+_MM_MODULE_PREFIXES = (
+    "vision_tower",
+    "vision_model",
+    "vision_encoder",
+    "embed_vision",
+    "multi_modal_projector",
+    "visual",
+    "audio_tower",
+    "audio_model",
+    "embed_audio",
+)
+
+
+def freeze_mm_modules(model):
+    """Freeze all vision/audio/multimodal-projector parameters.
+
+    Iterates over ``model.named_parameters()`` and sets ``requires_grad = False``
+    for any parameter whose name contains a known vision/audio module prefix.
+    This is useful when fine-tuning only the language backbone of a multimodal
+    model and avoids the need for ``ddp_find_unused_parameters=True``.
+    """
+    frozen_count = 0
+    for name, param in model.named_parameters():
+        # Check if any path component matches a vision/audio prefix
+        parts = name.split(".")
+        if any(part in _MM_MODULE_PREFIXES for part in parts):
+            if param.requires_grad:
+                param.requires_grad = False
+                frozen_count += 1
+                if is_main_process():
+                    LOG.debug(f"freeze_mm_modules: froze {name}")
 
-LOG = logging.getLogger("axolotl.utils.freeze")
+    if is_main_process():
+        LOG.info(f"freeze_mm_modules: froze {frozen_count} vision/audio parameters")
 
 
 def freeze_layers_except(model, regex_patterns):
@@ -120,6 +159,9 @@ def _merge_ranges(
     processed_ranges = [
         (start, end if end is not None else layer_size) for start, end in given_ranges
     ]
+    for start, end in processed_ranges:
+        if start < 0 or end > layer_size > 0 or start >= end:
+            raise ValueError(f"invalid unfreeze range: start={start}, end={end}")
 
     # No need to merge if there's only one or no ranges
     if len(processed_ranges) <= 1:
@@ -180,7 +222,7 @@ def __init__(self, pattern: str):
         """
         self.raw_pattern = pattern
         name_pattern, self.range = self._parse_pattern(pattern)
-        self.name_regex = re.compile(name_pattern.replace(".", "\\."))
+        self.name_regex = re.compile(re.sub(r"\.(?!\+)", "\\.", name_pattern))
 
     def match(self, name: str) -> bool:
         """
diff --git a/src/axolotl/utils/generation/__init__.py b/src/axolotl/utils/generation/__init__.py
new file mode 100644
index 0000000000..7a222d18d4
--- /dev/null
+++ b/src/axolotl/utils/generation/__init__.py
@@ -0,0 +1,5 @@
+"""Generation utilities for monitoring during training."""
+
+from .sft import format_generation_for_logging, generate_samples
+
+__all__ = ["generate_samples", "format_generation_for_logging"]
diff --git a/src/axolotl/utils/generation/sft.py b/src/axolotl/utils/generation/sft.py
new file mode 100644
index 0000000000..70fff80c59
--- /dev/null
+++ b/src/axolotl/utils/generation/sft.py
@@ -0,0 +1,174 @@
+"""Sample generation utilities for SFT/Pretrain training."""
+
+from typing import Any, List, Optional
+
+import torch
+from accelerate.utils import extract_model_from_parallel
+from colorama import Fore, Style
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def generate_samples(
+    model: torch.nn.Module,
+    tokenizer: Any,
+    dataloader: Any,
+    num_generation_samples: int = 3,
+    max_new_tokens: int = 50,
+    temperature: float = 0.7,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    do_sample: bool = True,
+    prompt_ratio: float = 0.5,
+) -> List[dict]:
+    """
+    Generate samples from the model during training for monitoring.
+
+    Args:
+        model: The model to generate from
+        tokenizer: The tokenizer to use for encoding/decoding
+        dataloader: Dataloader to sample prompts from
+        num_generation_samples: Number of samples to generate
+        max_new_tokens: Maximum new tokens to generate
+        temperature: Sampling temperature (0.0 = greedy)
+        top_p: Nucleus sampling parameter
+        top_k: Top-k sampling parameter
+        do_sample: Whether to use sampling vs greedy decoding
+        prompt_ratio: Ratio of sequence to use as prompt (0.0-1.0)
+
+    Returns:
+        List of dicts with 'prompt', 'generated', and 'full_text' keys
+    """
+    unwrapped_model = extract_model_from_parallel(model)
+
+    training = unwrapped_model.training
+    unwrapped_model.eval()
+
+    device = next(unwrapped_model.parameters()).device
+
+    generations = []
+
+    try:
+        with torch.no_grad():
+            samples_collected = 0
+
+            for batch in dataloader:
+                if samples_collected >= num_generation_samples:
+                    break
+
+                input_ids = batch["input_ids"].to(device)
+                attention_mask = batch.get("attention_mask")
+                if attention_mask is not None:
+                    attention_mask = attention_mask.to(device)
+                batch_size = input_ids.shape[0]
+
+                indices = torch.randperm(batch_size)[
+                    : num_generation_samples - samples_collected
+                ]
+
+                for idx in indices:
+                    if samples_collected >= num_generation_samples:
+                        break
+
+                    sequence = input_ids[idx]
+
+                    if attention_mask is not None:
+                        seq_len = attention_mask[idx].sum().item()
+                    else:
+                        seq_len = sequence.shape[0]
+
+                    if seq_len < 5:
+                        continue
+
+                    prompt_len = max(1, int(seq_len * prompt_ratio))
+                    prompt_ids = sequence[:prompt_len].unsqueeze(0)
+
+                    try:
+                        generation_config = {
+                            "max_new_tokens": max_new_tokens,
+                            "do_sample": do_sample,
+                            "pad_token_id": tokenizer.pad_token_id
+                            if tokenizer.pad_token_id is not None
+                            else tokenizer.eos_token_id,
+                        }
+
+                        if do_sample:
+                            generation_config["temperature"] = temperature
+                            if top_p is not None:
+                                generation_config["top_p"] = top_p
+                            if top_k is not None:
+                                generation_config["top_k"] = top_k
+
+                        generated_ids = unwrapped_model.generate(
+                            prompt_ids, **generation_config
+                        )
+
+                        prompt_text = tokenizer.decode(
+                            prompt_ids[0], skip_special_tokens=True
+                        )
+                        generated_text = tokenizer.decode(
+                            generated_ids[0][prompt_len:], skip_special_tokens=True
+                        )
+                        full_text = tokenizer.decode(
+                            generated_ids[0], skip_special_tokens=True
+                        )
+
+                        generations.append(
+                            {
+                                "prompt": prompt_text,
+                                "generated": generated_text,
+                                "full_text": full_text,
+                            }
+                        )
+
+                        samples_collected += 1
+
+                    except Exception as e:
+                        LOG.warning(f"Failed to generate sample: {e}", exc_info=True)
+                        continue
+
+    except Exception as e:
+        LOG.warning(f"Error during sample generation: {e}", exc_info=True)
+
+    if training:
+        unwrapped_model.train()
+    else:
+        unwrapped_model.eval()
+
+    return generations
+
+
+def format_generation_for_logging(
+    sample: dict, sample_idx: int, step: int
+) -> tuple[str, str]:
+    """
+    Format a generation sample for pretty logging.
+
+    Args:
+        sample: Dict with 'prompt', 'generated', and 'full_text' keys
+        sample_idx: Index of the sample
+        step: Current training step
+
+    Returns:
+        Tuple of (console_text, wandb_text)
+    """
+    console_text = (
+        f"\n{Style.BRIGHT}{Fore.CYAN}{'=' * 80}{Style.RESET_ALL}\n"
+        f"{Style.BRIGHT}{Fore.GREEN}Sample {sample_idx + 1} (Step {step}){Style.RESET_ALL}\n"
+        f"{Style.BRIGHT}{Fore.CYAN}{'=' * 80}{Style.RESET_ALL}\n"
+        f"{Style.BRIGHT}{Fore.YELLOW}[PROMPT]{Style.RESET_ALL}\n{sample['prompt']}\n\n"
+        f"{Style.BRIGHT}{Fore.MAGENTA}[GENERATED]{Style.RESET_ALL}\n{sample['generated']}\n"
+        f"{Style.BRIGHT}{Fore.CYAN}{'=' * 80}{Style.RESET_ALL}\n"
+    )
+    wandb_text = (
+        f"\n{'=' * 80}\n"
+        f"Sample {sample_idx + 1} (Step {step})\n"
+        f"{'=' * 80}\n"
+        f"[PROMPT]\n{sample['prompt']}\n\n"
+        f"[GENERATED]\n{sample['generated']}\n"
+        f"{'=' * 80}\n"
+    )
+
+    return console_text, wandb_text
diff --git a/src/axolotl/utils/gradient_checkpointing/__init__.py b/src/axolotl/utils/gradient_checkpointing/__init__.py
deleted file mode 100644
index 4639fc266c..0000000000
--- a/src/axolotl/utils/gradient_checkpointing/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-"""custom checkpointing utils"""
-from axolotl.utils.gradient_checkpointing.unsloth import (
-    Unsloth_Offloaded_Gradient_Checkpointer,
-)
-
-
-def hf_grad_checkpoint_unsloth_wrapper(
-    decoder_layer, *args, use_reentrant=None
-):  # pylint: disable=unused-argument
-    return Unsloth_Offloaded_Gradient_Checkpointer.apply(
-        decoder_layer.__self__,
-        *args,
-    )
diff --git a/src/axolotl/utils/import_helper.py b/src/axolotl/utils/import_helper.py
new file mode 100644
index 0000000000..f7d20099c5
--- /dev/null
+++ b/src/axolotl/utils/import_helper.py
@@ -0,0 +1,28 @@
+"""
+Helper for importing modules from strings
+"""
+
+import importlib
+
+
+def get_cls_from_module_str(module_str: str):
+    # use importlib to dynamically load the reward function from the module
+    if not isinstance(module_str, str) or not module_str.strip():
+        raise ValueError("module_str must be a non-empty string")
+
+    parts = module_str.split(".")
+    if len(parts) < 2:
+        raise ValueError(f"Invalid module string format: {module_str}")
+
+    try:
+        cls_name = parts[-1]
+        module_path = ".".join(parts[:-1])
+        mod = importlib.import_module(module_path)
+        mod_cls = getattr(mod, cls_name)
+        return mod_cls
+    except ImportError as e:
+        raise ImportError(f"Failed to import module '{module_path}': {e}") from e
+    except AttributeError as e:
+        raise AttributeError(
+            f"Class '{cls_name}' not found in module '{module_path}': {e}"
+        ) from e
diff --git a/src/axolotl/utils/logging.py b/src/axolotl/utils/logging.py
new file mode 100644
index 0000000000..7aaedef286
--- /dev/null
+++ b/src/axolotl/utils/logging.py
@@ -0,0 +1,52 @@
+"""Logging helpers to only log on main process."""
+
+import functools
+import logging
+import warnings
+
+from axolotl.utils.distributed import is_main_process
+
+# Suppress noisy bitsandbytes warnings about dtype casting during quantization
+warnings.filterwarnings(
+    "ignore",
+    message=".*MatMul8bitLt: inputs will be cast from.*",
+    category=UserWarning,
+)
+
+# Adapted from Accelerate
+# https://github.com/huggingface/accelerate/blob/main/src/accelerate/logging.py
+
+
+class MultiProcessAdapter(logging.LoggerAdapter):
+    """
+    Logger adapter for distributed logging, specifically to only log on main process.
+    """
+
+    @staticmethod
+    def _should_log(main_process_only: bool):
+        return not main_process_only or is_main_process()
+
+    def log(self, level, msg, *args, **kwargs):
+        main_process_only = kwargs.pop("main_process_only", True)
+        kwargs.setdefault("stacklevel", 2)
+
+        if self.isEnabledFor(level) and self._should_log(main_process_only):
+            msg, kwargs = self.process(msg, kwargs)
+            self.logger.log(level, msg, *args, **kwargs)
+
+    @functools.lru_cache(maxsize=10)
+    def warning_once(self, *args, **kwargs):
+        """
+        This method is identical to `logger.warning()`, but will emit the warning with the same message only once
+
+        Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the
+        cache. The assumption here is that all warning messages are unique across the code. If they aren't then need to
+        switch to another type of cache that includes the caller frame information in the hashing function.
+        """
+        self.warning(*args, **kwargs)
+
+
+def get_logger(name: str, log_level: str | None = None) -> MultiProcessAdapter:
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+    return MultiProcessAdapter(logger, extra={})
diff --git a/src/axolotl/utils/lora.py b/src/axolotl/utils/lora.py
new file mode 100644
index 0000000000..6ae481b6b3
--- /dev/null
+++ b/src/axolotl/utils/lora.py
@@ -0,0 +1,76 @@
+# Copyright 2025 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+module to get the state dict of a merged lora model
+"""
+
+import torch
+from peft.tuners.tuners_utils import onload_layer
+from peft.utils import ModulesToSaveWrapper, _get_submodules
+
+
+def get_lora_merged_state_dict(
+    model: torch.nn.Module,
+) -> dict:
+    r"""
+    Create and return a state_dict that has the LoRA deltas
+    merged into the base model’s weights, without modifying `model` in place.
+
+    Arguments:
+        model (torch.nn.Module): A model that has LoRA/PEFT adapters attached.
+
+    Returns:
+        dict: A state_dict of the merged parameters.
+    """
+
+    base_model_prefix = "base_model.model."
+    state_dict = {}
+    key_list = [key for key, _ in model.named_modules() if model.prefix not in key]
+    for key in key_list:
+        try:
+            _, target, _ = _get_submodules(model, key)
+        except AttributeError:
+            continue
+        with onload_layer(target):
+            weight_key = key.replace(base_model_prefix, "") + ".weight"
+            bias_key = key.replace(base_model_prefix, "") + ".bias"
+            if hasattr(target, "base_layer"):
+                target.merge(safe_merge=True, adapter_names=None)
+                # get the state_dict of target.base_layer
+                layer_state_dict = target.base_layer.state_dict()
+                state_dict[weight_key] = layer_state_dict["weight"]
+            elif isinstance(target, ModulesToSaveWrapper):
+                # save any additional trainable modules part of `modules_to_save`
+                new_module = target.modules_to_save[target.active_adapter]
+                if hasattr(new_module, "base_layer"):
+                    # check if the module is itself a tuner layer
+                    new_module.merge(safe_merge=True, adapter_names=None)
+                layer_state_dict = new_module.state_dict()
+                state_dict[weight_key] = layer_state_dict["weight"]
+            elif hasattr(target, "weight"):
+                if any(
+                    skip in key
+                    for skip in [
+                        ".original_module",
+                        ".modules_to_save",
+                        ".base_layer",
+                    ]
+                ):
+                    continue
+                layer_state_dict = target.state_dict()
+                state_dict[weight_key] = layer_state_dict["weight"]
+                if hasattr(target, "bias") and "bias" in layer_state_dict.keys():
+                    state_dict[bias_key] = layer_state_dict["bias"]
+    return state_dict
diff --git a/src/axolotl/utils/lora_embeddings.py b/src/axolotl/utils/lora_embeddings.py
deleted file mode 100644
index 70f56655ea..0000000000
--- a/src/axolotl/utils/lora_embeddings.py
+++ /dev/null
@@ -1,14 +0,0 @@
-"""
-helpers for lora embeddings
-"""
-
-
-def get_linear_embedding_layers(model_type):
-    """
-    returns the linear embedding layers needed for loras, dependent on the model arch
-    """
-    if model_type == "gpt_neox":
-        return ["embed_in", "embed_out"]
-    if model_type == "falcon":
-        return ["word_embeddings", "lm_head"]
-    return ["embed_tokens", "lm_head"]
diff --git a/src/axolotl/utils/mistral/__init__.py b/src/axolotl/utils/mistral/__init__.py
new file mode 100644
index 0000000000..eb51031ec1
--- /dev/null
+++ b/src/axolotl/utils/mistral/__init__.py
@@ -0,0 +1,6 @@
+"""Init for `axolotl.utils.mistral` module."""
+
+from axolotl.utils.mistral.mistral3_processor import Mistral3Processor
+from axolotl.utils.mistral.mistral_tokenizer import HFMistralTokenizer
+
+__all__ = ["HFMistralTokenizer", "Mistral3Processor"]
diff --git a/src/axolotl/utils/mistral/mistral3_processor.py b/src/axolotl/utils/mistral/mistral3_processor.py
new file mode 100644
index 0000000000..03a155e53d
--- /dev/null
+++ b/src/axolotl/utils/mistral/mistral3_processor.py
@@ -0,0 +1,160 @@
+"""Processor for Mistral3 multimodal models with image support"""
+
+from typing import Any, Dict, Optional, Union
+
+import torch
+from transformers import ProcessorMixin
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.processing_utils import ProcessingKwargs
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+
+from axolotl.utils.mistral.mistral_tokenizer import HFMistralTokenizer
+
+
+class Mistral3ProcessorKwargs(ProcessingKwargs):
+    _defaults: Dict[str, Dict[str, Any]] = {
+        "text_kwargs": {
+            "padding": True,
+        },
+        "common_kwargs": {
+            "return_tensors": "pt",
+            "return_dict": True,
+            "tokenize": True,
+        },
+    }
+
+
+class Mistral3Processor(ProcessorMixin):
+    """
+    Processor for Mistral3 multimodal models that handles text and images.
+    Wraps HFMistralTokenizer and adds image processing capabilities.
+    """
+
+    def __init__(self, tokenizer: HFMistralTokenizer):
+        super().__init__(tokenizer)
+
+    @property
+    def audio_tokenizer(self) -> None:
+        """Audio tokenizer is not supported. Dummy method to satisfy HuggingFace API."""
+        return None
+
+    def _merge_kwargs(
+        self, processor_kwargs_class: Any, **kwargs: Any
+    ) -> Dict[str, Dict[str, Any]]:
+        """Merge kwargs with defaults similar to ProcessorMixin"""
+        defaults = processor_kwargs_class._defaults
+        output_kwargs: Dict[str, Dict[str, Any]] = {}
+
+        for kwarg_type, default_values in defaults.items():
+            output_kwargs[kwarg_type] = {**default_values}
+
+        # Update with provided kwargs
+        for key, value in kwargs.items():
+            # Try to match key to appropriate kwarg type
+            if key in ["padding", "truncation", "max_length"]:
+                output_kwargs.setdefault("text_kwargs", {}).update({key: value})
+            elif key in ["return_tensors", "return_dict", "tokenize"]:
+                output_kwargs.setdefault("common_kwargs", {}).update({key: value})
+            else:
+                # Add to text_kwargs by default
+                output_kwargs.setdefault("text_kwargs", {}).update({key: value})
+
+        return output_kwargs
+
+    def apply_chat_template(
+        self,
+        conversation: Union[list[dict[str, str]], list[list[dict[str, str]]]],
+        **kwargs: Any,
+    ) -> Union[BatchFeature, str, list[str]]:
+        """
+        Apply chat template with image support for Mistral3.
+
+        Similar to VoxtralProcessor, this method extracts images from the conversation,
+        calls the tokenizer's apply_chat_template, then adds pixel_values and image_sizes
+        to the result.
+        """
+        output_kwargs = self._merge_kwargs(Mistral3ProcessorKwargs, **kwargs)
+        text_kwargs = output_kwargs["text_kwargs"]
+        common_kwargs = output_kwargs["common_kwargs"]
+
+        return_tensors = common_kwargs.pop("return_tensors", "pt")
+        if return_tensors != "pt":
+            raise ValueError(
+                f"{self.__class__.__name__} only supports `return_tensors='pt'`."
+            )
+
+        return_dict = common_kwargs.pop("return_dict", False)
+        tokenize = common_kwargs.pop("tokenize", False)
+
+        # Determine if batched
+        if isinstance(conversation, (list, tuple)) and (
+            isinstance(conversation[0], (list, tuple))
+            or hasattr(conversation[0], "content")
+        ):
+            is_batched = True
+            conversations = conversation
+        else:
+            is_batched = False
+            conversations = [conversation]  # type: ignore
+
+        # Call tokenizer's apply_chat_template
+        tokenizer_kwargs = {**text_kwargs, **common_kwargs}
+        tokenizer_kwargs["return_tensors"] = return_tensors
+        tokenizer_kwargs["tokenize"] = tokenize
+        tokenizer_kwargs["return_dict"] = return_dict
+
+        encoded_instruct_inputs = self.tokenizer.apply_chat_template(
+            conversations,
+            **tokenizer_kwargs,
+        )
+
+        if tokenize:
+            if return_dict:
+                # The tokenizer already handles pixel_values, we just need to add image_sizes
+                if hasattr(encoded_instruct_inputs, "items"):
+                    data: Dict[str, Any] = dict(encoded_instruct_inputs)  # type: ignore
+                elif hasattr(encoded_instruct_inputs, "data"):
+                    data = encoded_instruct_inputs.data  # type: ignore
+                else:
+                    raise ValueError("Unknown data type")
+
+                if "pixel_values" in data:
+                    pixel_values = data["pixel_values"]
+
+                    # MistralTokenizer returns a Double, so we convert to fp32
+                    data["pixel_values"] = pixel_values.to(dtype=torch.float32)
+
+                    # Always batched: [B, C, H, W] -> image_sizes: [B, 2]
+                    # Since tensor is homogeneous, all images have same H, W
+                    batch_size = pixel_values.shape[0]
+                    image_sizes = torch.tensor([pixel_values.shape[-2:]] * batch_size)
+                    data["image_sizes"] = image_sizes
+
+                return BatchFeature(data=data, tensor_type=return_tensors)
+
+        if not is_batched:
+            return encoded_instruct_inputs[0]
+
+        return encoded_instruct_inputs
+
+    def __call__(
+        self,
+        text: Optional[
+            Union[
+                TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]
+            ]
+        ],
+        **kwargs: Any,
+    ) -> BatchFeature:
+        """
+        Forward text processing to the tokenizer.
+        This method does not support images - use apply_chat_template instead.
+        """
+        output_kwargs = self._merge_kwargs(Mistral3ProcessorKwargs, **kwargs)
+        text_kwargs = output_kwargs["text_kwargs"]
+        common_kwargs = output_kwargs["common_kwargs"]
+
+        out = self.tokenizer(text, **text_kwargs)
+        return BatchFeature(
+            data=out, tensor_type=common_kwargs.pop("return_tensors", None)
+        )
diff --git a/src/axolotl/utils/mistral/mistral_tokenizer.py b/src/axolotl/utils/mistral/mistral_tokenizer.py
new file mode 100644
index 0000000000..46b7b526dc
--- /dev/null
+++ b/src/axolotl/utils/mistral/mistral_tokenizer.py
@@ -0,0 +1,296 @@
+"""Wrapper for MistralTokenizer from mistral-common"""
+
+import os
+from typing import Optional
+
+import numpy as np
+from mistral_common.protocol.instruct.request import ModelSettings
+from mistral_common.protocol.instruct.validator import ValidationMode
+from mistral_common.tokens.tokenizers.utils import download_tokenizer_from_hf_hub
+from pydantic import ValidationError
+from torch import Tensor
+from transformers.tokenization_mistral_common import MistralCommonBackend
+from transformers.tokenization_utils_base import VERY_LARGE_INTEGER
+
+
+class HFMistralTokenizer(MistralCommonBackend):
+    """
+    Wraps mistral_common.tokens.tokenizers.mistral.MistralTokenizer
+    and exposes HuggingFace API for special tokens.
+    """
+
+    def __init__(self, name_or_path: str, **kwargs):
+        """
+        Args:
+            name_or_path: The name or path to the tokenizer files or the repo id.
+            **kwargs: Additional keyword arguments passed to the parent class.
+        """
+        kwargs.pop("mode", None)
+
+        mode = ValidationMode.finetuning
+        super().__init__(**kwargs, mode=mode)
+
+        self._name_or_path = name_or_path
+
+        # set mode as is not set upstream
+        self._set_mode(mode)
+        self._patch_instruct_request_normalizer()
+
+    @property
+    def name_or_path(self) -> str:
+        return self._name_or_path
+
+    @name_or_path.setter
+    def name_or_path(self, name_or_path: str) -> None:
+        self._name_or_path = name_or_path
+
+    @property
+    def chat_template(self) -> str | None:
+        """Chat template is not supported. Dummy method to satisfy HuggingFace API."""
+        return "[This is a dummy chat template]"
+
+    @chat_template.setter
+    def chat_template(self, chat_template: str | None) -> None:
+        pass
+
+    def _set_mode(self, mode: ValidationMode):
+        """Set the mode of the MistralRequestValidator.
+
+        Args:
+            mode: The mode to set.
+
+        Raises:
+            RuntimeError: If the MistralRequestValidator does not have a _mode attribute.
+        """
+        # Check if MistralRequestValidator has a _mode attribute.
+        # This is a private API and may change in the future.
+
+        from mistral_common.protocol.instruct.validator import MistralRequestValidator
+
+        if not (
+            hasattr(self.tokenizer, "_chat_completion_request_validator")
+            and isinstance(
+                self.tokenizer._chat_completion_request_validator,
+                MistralRequestValidator,
+            )
+            and hasattr(self.tokenizer._chat_completion_request_validator, "_mode")
+        ):
+            raise RuntimeError(
+                f"Unable to switch mistral tokenizer to {mode.value} mode - "
+                "private API `_chat_completion_request_validator._mode` missing."
+            )
+
+        self.tokenizer._chat_completion_request_validator._mode = mode
+
+    @staticmethod
+    def _missing_instruct_request_defaults(exc: ValidationError) -> bool:
+        missing_fields = {
+            err["loc"][0]
+            for err in exc.errors()
+            if err.get("type") == "missing" and len(err.get("loc", ())) == 1
+        }
+        return {
+            "truncate_at_max_tokens",
+            "continue_final_message",
+        }.issubset(missing_fields)
+
+    def _patch_instruct_request_normalizer(self) -> None:
+        normalizer = getattr(self.tokenizer, "_instruct_request_normalizer", None)
+        if normalizer is None or getattr(
+            normalizer, "_axolotl_instruct_defaults_patched", False
+        ):
+            return
+
+        original = normalizer.from_chat_completion_request
+
+        def from_chat_completion_request(request):
+            try:
+                return original(request)
+            except ValidationError as exc:
+                if not self._missing_instruct_request_defaults(exc):
+                    raise
+
+                messages = normalizer._aggregate_messages(request.messages)
+                settings = normalizer.build_settings(request)
+                if settings != ModelSettings.none():
+                    raise
+
+                try:
+                    system_prompt = normalizer._aggregate_system_prompts(
+                        request.messages
+                    )
+                except (AttributeError, NotImplementedError):
+                    system_prompt = None
+
+                return normalizer._instruct_request_class(
+                    messages=messages,
+                    system_prompt=system_prompt,
+                    available_tools=request.tools,
+                    truncate_at_max_tokens=None,
+                    continue_final_message=getattr(
+                        request, "continue_final_message", False
+                    ),
+                    settings=settings,
+                )
+
+        normalizer.from_chat_completion_request = from_chat_completion_request
+        normalizer._axolotl_instruct_defaults_patched = True
+
+    def apply_chat_template(  # type: ignore
+        self,
+        conversation: list[dict] | list[list[dict]],
+        chat_template: str | None = None,
+        add_generation_prompt: bool = False,
+        **kwargs,
+    ) -> str | list[int]:
+        """Patched fn to handle setting test mode, remove chat_template and add_generation_prompt kwarg"""
+
+        # pop unnecessary kwarg for mistral
+        kwargs.pop("real_last_index", None)
+        kwargs.pop("add_special_tokens", None)
+
+        try:
+            if add_generation_prompt:
+                self._set_mode(ValidationMode.test)
+
+            self._patch_instruct_request_normalizer()
+            out = super().apply_chat_template(conversation, **kwargs)
+
+            return out  # type: ignore
+
+        finally:
+            if add_generation_prompt:
+                self._set_mode(ValidationMode.finetuning)
+
+    def decode(  # type: ignore
+        self,
+        token_ids: int | list[int] | np.ndarray | Tensor,
+        **kwargs,
+    ) -> str:
+        """
+        Decode token_ids into str.
+
+        This overrides upstream.decode to convert int to list[int]
+        """
+
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+
+        return super().decode(token_ids, **kwargs)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | os.PathLike,
+        *init_inputs,
+        mode: ValidationMode = ValidationMode.test,
+        cache_dir: Optional[str | os.PathLike] = None,
+        force_download: bool = False,
+        local_files_only: bool = False,
+        token: Optional[str | bool] = None,
+        revision: str = "main",
+        model_max_length: int = VERY_LARGE_INTEGER,
+        padding_side: str = "left",
+        truncation_side: str = "right",
+        model_input_names: Optional[list[str]] = None,
+        clean_up_tokenization_spaces: bool = False,
+        **kwargs,
+    ):
+        r"""
+        Patched fn to pass `name_or_path` and remove extra kwargs.
+
+        Instantiate a `MistralCommonBackend` from a predefined
+        tokenizer.
+
+        Args:
+            pretrained_model_name_or_path (`str` or `os.PathLike`):
+                Can be either:
+
+                - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co.
+                - A path to a *directory* containing the tokenizer config, for instance saved
+                  using the [`MistralCommonBackend.tokenization_mistral_common.save_pretrained`] method, e.g.,
+                  `./my_model_directory/`.
+            mode (`ValidationMode`, *optional*, defaults to `ValidationMode.test`):
+                Validation mode for the `MistralTokenizer` tokenizer.
+            cache_dir (`str` or `os.PathLike`, *optional*):
+                Path to a directory in which a downloaded predefined tokenizer vocabulary files should be cached if the
+                standard cache should not be used.
+            force_download (`bool`, *optional*, defaults to `False`):
+                Whether or not to force the (re-)download the vocabulary files and override the cached versions if they
+                exist.
+            token (`str` or *bool*, *optional*):
+                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+                when running `hf auth login` (stored in `~/.huggingface`).
+            local_files_only (`bool`, *optional*, defaults to `False`):
+                Whether or not to only rely on local files and not to attempt to download any files.
+            revision (`str`, *optional*, defaults to `"main"`):
+                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+                identifier allowed by git.
+            max_length (`int`, *optional*):
+                Controls the maximum length to use by one of the truncation/padding parameters.
+
+                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length
+                is required by one of the truncation/padding parameters. If the model has no specific maximum input
+                length (like XLNet) truncation/padding to a maximum length will be deactivated.
+            padding_side (`str`, *optional*, defaults to `"left"`):
+                The side on which the model should have padding applied. Should be selected between ['right', 'left'].
+                Default value is picked from the class attribute of the same name.
+            truncation_side (`str`, *optional*, defaults to `"right"`):
+                The side on which the model should have truncation applied. Should be selected between ['right', 'left'].
+            model_input_names (`List[string]`, *optional*):
+                The list of inputs accepted by the forward pass of the model (like `"token_type_ids"` or
+                `"attention_mask"`). Default value is picked from the class attribute of the same name.
+            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+                Whether or not the model should cleanup the spaces that were added when splitting the input text during the
+                tokenization process.
+            kwargs (additional keyword arguments, *optional*):
+                Not supported by `MistralCommonBackend.from_pretrained`.
+                Will raise an error if used.
+        """
+        if init_inputs:
+            raise ValueError(
+                "`init_inputs` are not supported by `MistralCommonBackend.from_pretrained`."
+            )
+
+        # Delete trust_remote_code as it does nothing
+        kwargs.pop("trust_remote_code", None)
+
+        # Delete tokenizer as it does nothing
+        kwargs.pop("tokenizer", None)
+
+        # Handle kwargs and AutoTokenizer case
+        if kwargs and not kwargs.keys() == {"_from_auto"}:
+            raise ValueError(
+                f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonBackend.from_pretrained`."
+            )
+
+        if not os.path.isfile(pretrained_model_name_or_path):
+            tokenizer_path = download_tokenizer_from_hf_hub(
+                repo_id=str(pretrained_model_name_or_path),
+                cache_dir=str(cache_dir),
+                token=token,
+                revision=revision,
+                force_download=force_download,
+                local_files_only=local_files_only,
+            )
+        else:
+            tokenizer_path = str(pretrained_model_name_or_path)
+
+        return cls(
+            name_or_path=str(pretrained_model_name_or_path),
+            tokenizer_path=tokenizer_path,
+            mode=mode,
+            model_max_length=model_max_length,
+            padding_side=padding_side,
+            truncation_side=truncation_side,
+            model_input_names=model_input_names,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+        )
+
+    def save_pretrained(self, *args, **kwargs) -> tuple[str, ...]:
+        """
+        Patches to remove save_jinja_files from being passed onwards.
+        """
+        kwargs.pop("save_jinja_files", None)
+        return super().save_pretrained(*args, **kwargs)
diff --git a/src/axolotl/utils/mlflow_.py b/src/axolotl/utils/mlflow_.py
index ce77390342..8710b07d06 100644
--- a/src/axolotl/utils/mlflow_.py
+++ b/src/axolotl/utils/mlflow_.py
@@ -16,3 +16,7 @@ def setup_mlflow_env_vars(cfg: DictDefault):
     # Enable mlflow if experiment name is present
     if cfg.mlflow_experiment_name and len(cfg.mlflow_experiment_name) > 0:
         cfg.use_mlflow = True
+
+    # Enable logging hf artifacts in mlflow if value is truthy
+    if cfg.hf_mlflow_log_artifacts is True:
+        os.environ["HF_MLFLOW_LOG_ARTIFACTS"] = "true"
diff --git a/src/axolotl/utils/model_shard_quant.py b/src/axolotl/utils/model_shard_quant.py
index 65f23b9e0f..ca152113a3 100644
--- a/src/axolotl/utils/model_shard_quant.py
+++ b/src/axolotl/utils/model_shard_quant.py
@@ -1,6 +1,7 @@
 """
 module to handle loading model on cpu/meta device for FSDP
 """
+
 import os
 import time
 from typing import List, Optional, Type, Union
@@ -13,6 +14,7 @@
 from torch import Tensor, nn
 from tqdm import tqdm
 from transformers import AutoModelForCausalLM
+from transformers.quantizers import AutoHfQuantizer
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, hub
 
 
@@ -44,9 +46,7 @@ def _replace_linear(
 
         if isinstance(module, torch.nn.Linear) and name not in skip_modules:
             if issubclass(linear_replacement, Linear4bit):
-                model._modules[  # pylint: disable=protected-access
-                    name
-                ] = linear_replacement(
+                model._modules[name] = linear_replacement(
                     module.in_features,
                     module.out_features,
                     module.bias is not None,
@@ -148,8 +148,8 @@ def load_sharded_model(
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             use_cache=False,
-            torch_dtype=torch.float32,
-            _attn_implementation=model_config._attn_implementation,  # pylint: disable=protected-access
+            dtype=torch.float32,
+            _attn_implementation=model_config._attn_implementation,
             trust_remote_code=cfg.trust_remote_code,
         )
         dtype = torch_dtype if not cfg.float32 else None
@@ -158,7 +158,7 @@ def load_sharded_model(
         with init_empty_weights():
             model = AutoModelForCausalLM.from_config(
                 model_config,
-                torch_dtype=torch_dtype,
+                dtype=torch_dtype,
                 trust_remote_code=cfg.trust_remote_code,
             )
     return model
@@ -173,6 +173,7 @@ def load_sharded_model_quant(
     low_memory=True,
     verbose=False,
     loading_workers=2,
+    quantization_config=None,
 ):
     with init_empty_weights():
         model = AutoModelForCausalLM.from_config(
@@ -186,15 +187,26 @@ def load_sharded_model_quant(
                 compute_dtype=compute_dtype,
                 quant_type="nf4",
                 quant_storage=quant_storage,
+                compress_statistics=True,  # bnb_4bit_use_double_quant
+                skip_modules=[
+                    "lm_head",
+                    "embed_out",
+                ],
             )
         else:
             # this is the more common case with HF transformers
+            # TODO can we detect the model arch and dynamically set skip_modules
             model.model = _replace_linear(
                 model.model,
                 Linear4bit,
                 compute_dtype=compute_dtype,
                 quant_type="nf4",
                 quant_storage=quant_storage,
+                compress_statistics=True,  # bnb_4bit_use_double_quant
+                skip_modules=[
+                    "lm_head",
+                    "embed_out",
+                ],
             )
     model.is_loaded_in_4bit = True
 
@@ -251,8 +263,13 @@ def load_and_quantize_parallel(name_param, model, **kwargs):
             quant_method=quant_method,
         )
 
+    # these attributes are needed to inform transformers/peft of the quantization
+    model.is_quantized = True
+    model.quantization_method = "bitsandbytes"
+    model.hf_quantizer = AutoHfQuantizer.from_config(quantization_config)
+
     if cfg.local_rank == 0 and verbose:
-        print(f"Loaded model weights in {time.time()-start:.3f} seconds")
+        print(f"Loaded model weights in {time.time() - start:.3f} seconds")
     # cleanup any extra memory usage from parallel loading
     torch.cuda.empty_cache()
 
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
deleted file mode 100644
index e94a0f6b88..0000000000
--- a/src/axolotl/utils/models.py
+++ /dev/null
@@ -1,1009 +0,0 @@
-"""Module for models and model loading"""
-
-# pylint: disable=too-many-lines
-
-import logging
-import math
-import os
-import types
-from typing import Any, Dict, Optional, Tuple, Union  # noqa: F401
-
-import addict
-import bitsandbytes as bnb
-import torch
-import transformers
-import transformers.modeling_utils
-from accelerate import init_empty_weights
-from bitsandbytes.nn import Params4bit
-from peft import (
-    LoftQConfig,
-    PeftConfig,
-    PeftModel,
-    PeftModelForCausalLM,
-    prepare_model_for_kbit_training,
-)
-from peft.tuners.lora import QuantLinear
-from torch import nn
-from transformers import (  # noqa: F401
-    AddedToken,
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BitsAndBytesConfig,
-    GPTQConfig,
-    PreTrainedModel,
-    PreTrainedTokenizerBase,
-)
-from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
-
-from axolotl.models.mamba import fix_mamba_attn_for_loss
-from axolotl.monkeypatch.multipack import (
-    SUPPORTED_MULTIPACK_MODEL_TYPES,
-    patch_for_multipack,
-)
-from axolotl.prompt_tokenizers import LLAMA_DEFAULT_EOS_TOKEN
-from axolotl.utils.bench import log_gpu_memory_usage
-from axolotl.utils.chat_templates import chat_templates
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.distributed import zero_only
-from axolotl.utils.gradient_checkpointing import hf_grad_checkpoint_unsloth_wrapper
-from axolotl.utils.lora_embeddings import get_linear_embedding_layers
-from axolotl.utils.model_shard_quant import load_sharded_model, load_sharded_model_quant
-
-LOG = logging.getLogger("axolotl")
-
-
-# copied from accelerator.FullyShardedDataParallelPlugin
-def get_module_class_from_name(module, name):
-    """
-    Gets a class from a module by its name.
-
-    Args:
-        module (`torch.nn.Module`): The module to get the class from.
-        name (`str`): The name of the class.
-    """
-    modules_children = list(module.children())
-    if module.__class__.__name__ == name:
-        return module.__class__
-
-    if len(modules_children) == 0:
-        return None
-
-    for child_module in modules_children:
-        module_class = get_module_class_from_name(child_module, name)
-        if module_class is not None:
-            return module_class
-
-    return None
-
-
-def check_model_config(cfg: DictDefault, model_config: Union[AutoConfig, DictDefault]):
-    quant_config_exists = (
-        hasattr(model_config, "quantization_config")
-        and model_config.quantization_config
-    )
-    quant_config_method_is_gptq = (
-        quant_config_exists
-        and "quant_method" in model_config.quantization_config
-        and model_config.quantization_config["quant_method"] == "gptq"
-    )
-
-    if cfg.gptq and not quant_config_method_is_gptq:
-        raise ValueError(
-            "model_config.quantization_config is not set or quant_method is not set to gptq. "
-            "Please make sure to point to a GPTQ model."
-        )
-
-    if not cfg.gptq and quant_config_exists:
-        raise ValueError(
-            "model_config.quantization_config is set but `gptq` flag is not. "
-            "Please use the `gptq` flag to train quantized model or point to a non-quantized model."
-        )
-
-    lora_modules_to_save = get_linear_embedding_layers(model_config.model_type)
-    if (
-        cfg.adapter
-        and cfg.tokens
-        and (
-            not cfg.lora_modules_to_save
-            or not all(x in cfg.lora_modules_to_save for x in lora_modules_to_save)
-        )
-    ):
-        lora_modules_to_save = ", ".join(map(lambda x: f"`{x}`", lora_modules_to_save))
-        raise ValueError(
-            f"`lora_modules_to_save` not properly set when adding new tokens. Please include [{lora_modules_to_save}] in `lora_modules_to_save`."
-        )
-
-
-def load_model_config(cfg):
-    model_config_name = cfg.base_model_config or cfg.base_model
-    if not model_config_name and cfg.tokenizer_config:
-        model_config_name = cfg.tokenizer_config
-    trust_remote_code = cfg.trust_remote_code is True
-    config_kwargs = {}
-    if cfg.revision_of_model:
-        config_kwargs["revision"] = cfg.revision_of_model
-
-    try:
-        model_config = AutoConfig.from_pretrained(
-            model_config_name,
-            trust_remote_code=trust_remote_code,
-            **config_kwargs,
-        )
-    except ValueError as err:
-        if "mamba" in model_config_name:
-            return addict.Dict(
-                {
-                    "model_type": "mamba",
-                }
-            )
-        raise err
-
-    if cfg.overrides_of_model_config:
-        for key, val in cfg.overrides_of_model_config.items():
-            setattr(model_config, key, val)
-
-    check_model_config(cfg, model_config)
-
-    return model_config
-
-
-def load_tokenizer(cfg):
-    model_config = load_model_config(cfg)
-    tokenizer_kwargs = {}
-    use_fast = True  # this is the default
-
-    if cfg.tokenizer_use_fast is not None:
-        use_fast = cfg.tokenizer_use_fast
-    if cfg.tokenizer_legacy is not None:
-        # True is the default w/ https://github.com/huggingface/transformers/pull/25224
-        tokenizer_kwargs["legacy"] = cfg.tokenizer_legacy
-
-    tokenizer_cls = AutoTokenizer
-    if cfg.tokenizer_type:
-        tokenizer_cls = getattr(transformers, cfg.tokenizer_type)
-
-    tokenizer = tokenizer_cls.from_pretrained(
-        cfg.tokenizer_config,
-        trust_remote_code=cfg.trust_remote_code or False,
-        use_fast=use_fast,
-        **tokenizer_kwargs,
-    )
-
-    if (
-        tokenizer.__class__.__name__
-        in [
-            "LlamaTokenizer",
-            "LlamaTokenizerFast",
-            "CodeLlamaTokenizer",
-            "CodeLlamaTokenizerFast",
-        ]
-        and hasattr(tokenizer, "pad_token")
-        and not tokenizer.pad_token
-    ):
-        # set a pad_token, but use eos_token so we don't add a new token
-        tokenizer.pad_token = LLAMA_DEFAULT_EOS_TOKEN
-
-    if tokenizer.__class__.__name__ == "GPTNeoXTokenizerFast":
-        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-        os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-    # Mistral's official FA implementation requires left padding
-    if cfg.is_mistral_derived_model and cfg.flash_attention and not cfg.sample_packing:
-        tokenizer.padding_side = "left"
-
-    # Qwen base only has single token, so we need to set the special tokens
-    if cfg.is_qwen_derived_model:
-        token_ids = ["bos_token_id", "eos_token_id", "pad_token_id", "unk_token_id"]
-        for attr_name in token_ids:
-            if getattr(tokenizer, attr_name) is None:
-                setattr(tokenizer, attr_name, tokenizer.eod_id)
-
-        token_names = ["bos_token", "eos_token", "pad_token", "unk_token"]
-        for attr_name in token_names:
-            if getattr(tokenizer, attr_name) is None:
-                setattr(tokenizer, attr_name, "<|endoftext|>")
-
-    additional_special_tokens = None
-    if cfg.special_tokens:
-        special_tokens = cfg.special_tokens.to_dict()
-        additional_special_tokens = special_tokens.pop(
-            "additional_special_tokens", None
-        )
-        lora_modules_to_save = get_linear_embedding_layers(model_config.model_type)
-        for k, val in special_tokens.items():
-            # check if new special token is not already in tokenizer and
-            # is adapter training to make sure lora_modules_to_save is set
-            # pylint: disable=too-many-boolean-expressions
-            if (
-                (getattr(tokenizer, k) is None or getattr(tokenizer, k) != val)
-                and (len(tokenizer.encode(val, add_special_tokens=False)) > 2)
-                and cfg.adapter
-                and (
-                    not cfg.lora_modules_to_save
-                    or not all(
-                        x in cfg.lora_modules_to_save for x in lora_modules_to_save
-                    )
-                )
-            ):
-                lora_modules_to_save = ", ".join(
-                    [f"`{x}`" for x in lora_modules_to_save]
-                )
-                raise ValueError(
-                    f"Please set lora_modules_to_save to [{lora_modules_to_save}] when using an adapter and changing the special tokens."
-                )
-
-            tokenizer.add_special_tokens(
-                {k: AddedToken(val, rstrip=False, lstrip=False, normalized=False)}
-            )
-
-        # If we add bos_token and eos_token, we need to update the post processor to
-        # handle them correctly.
-        # https://github.com/huggingface/transformers/pull/24132
-        bos_or_eos_in_special_tokens = (
-            "bos_token" in cfg.special_tokens and "eos_token" in cfg.special_tokens
-        )
-        if (
-            tokenizer.__class__.__name__
-            in (
-                "LlamaTokenizerFast",
-                "CodeLlamaTokenizerFast",
-            )
-            and bos_or_eos_in_special_tokens
-        ):
-            tokenizer.update_post_processor()
-
-    if cfg.tokens:
-        tokenizer.add_tokens(
-            [
-                AddedToken(token, rstrip=False, lstrip=False, normalized=False)
-                for token in cfg.tokens
-            ]
-        )
-
-    # Additional special tokens are a List, and need to be treated differently than regular special
-    # tokens. We add them after we have called `add_tokens` in case these additional special tokens
-    # are new tokens.
-    #
-    # Usage:
-    #
-    # ```py
-    # special_tokens:
-    #   additional_special_tokens: ["<|im_start|>", "<|im_end|>"]
-    # ```
-    if additional_special_tokens is not None:
-        tokenizer.add_special_tokens(
-            {"additional_special_tokens": additional_special_tokens}
-        )
-
-    with zero_only():
-        LOG.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}")
-        LOG.debug(f"BOS: {tokenizer.bos_token_id} / {tokenizer.bos_token}")
-        LOG.debug(f"PAD: {tokenizer.pad_token_id} / {tokenizer.pad_token}")
-        LOG.debug(f"UNK: {tokenizer.unk_token_id} / {tokenizer.unk_token}")
-
-    if cfg.chat_template:
-        chat_template_string = chat_templates(cfg.chat_template)
-        if cfg.default_system_message and cfg.chat_template == "chatml":
-            chat_template_string = chat_template_string.replace(
-                "You are a helpful assistant.", cfg.default_system_message
-            )
-
-        tokenizer.chat_template = chat_template_string
-    else:
-        LOG.info(
-            "No Chat template selected. Consider adding a chat template for easier inference."
-        )
-    return tokenizer
-
-
-def load_model(
-    cfg: DictDefault,
-    tokenizer: PreTrainedTokenizerBase,
-    inference: bool = False,
-    reference_model: bool = False,
-) -> Tuple[PreTrainedModel, Optional[PeftConfig]]:
-    """
-    Load a model for a given configuration and tokenizer.
-    """
-    base_model = cfg.base_model
-    model_type = cfg.type_of_model
-    model_config = load_model_config(cfg)
-
-    # TODO refactor as a kwarg
-    load_in_8bit = cfg.load_in_8bit
-
-    if cfg.gradient_checkpointing == "unsloth":
-        transformers.modeling_utils.checkpoint = hf_grad_checkpoint_unsloth_wrapper
-
-    if hasattr(model_config, "model_type") and model_config.model_type == "btlm":
-        if cfg.flash_attention:
-            from axolotl.monkeypatch.btlm_attn_hijack_flash import (
-                replace_btlm_attn_with_flash_attn,
-            )
-
-            replace_btlm_attn_with_flash_attn(cfg.base_model)
-
-    if (
-        hasattr(model_config, "model_type")
-        and model_config.model_type == "stablelm_epoch"
-    ):
-        if cfg.flash_attention and cfg.sample_packing:
-            from axolotl.monkeypatch.stablelm_attn_hijack_flash import (
-                replace_stablelm_attn_with_flash_attn,
-            )
-
-            replace_stablelm_attn_with_flash_attn(cfg.base_model)
-
-    if cfg.sample_packing and cfg.s2_attention:
-        raise ValueError(
-            "Received `sample_packing=true` and `s2_attention=true`; however, \
-        shifted-sparse attention does not currently support sample packing."
-        )
-
-    if (
-        cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES
-        and cfg.flash_attention
-        and cfg.sample_packing
-    ):
-        patch_for_multipack(cfg.model_config_type, model_name=cfg.base_model)
-    elif cfg.is_llama_derived_model:
-        # Modify all llama derived models in one block
-
-        if cfg.flash_attention:
-            from axolotl.monkeypatch.llama_attn_hijack_flash import (
-                replace_llama_attn_with_flash_attn,
-            )
-
-            if cfg.sample_packing:
-                if cfg.device not in ["mps", "cpu"] and not inference:
-                    LOG.info("patching with flash attention for sample packing")
-                    replace_llama_attn_with_flash_attn(
-                        packed=True,
-                        cross_entropy=cfg.flash_attn_cross_entropy,
-                        rms_norm=cfg.flash_attn_rms_norm,
-                    )
-            elif cfg.s2_attention:
-                LOG.info("patching w/ flash-enabled, shifted-sparse attention")
-                replace_llama_attn_with_flash_attn(
-                    packed=False,
-                    cross_entropy=cfg.flash_attn_cross_entropy,
-                    rms_norm=cfg.flash_attn_rms_norm,
-                    use_shifted_sparse_attn=True,
-                )
-        elif cfg.xformers_attention:
-            from axolotl.monkeypatch.llama_attn_hijack_xformers import (
-                hijack_llama_attention,
-            )
-
-            LOG.info("patching with xformers attention")
-            hijack_llama_attention()
-        elif cfg.sample_packing:
-            from axolotl.monkeypatch.llama_patch_multipack import (
-                hijack_llama_prepare_4d_mask,
-            )
-
-            LOG.info("patching llama _prepare_4d_causal_attention_mask*")
-            hijack_llama_prepare_4d_mask()
-        elif cfg.s2_attention:
-            raise NotImplementedError(
-                "Shifted-sparse attention not currently implemented without flash attention."
-            )
-
-    # Modify mistral derived models
-    if (
-        cfg.model_config_type == "mistral"
-        and cfg.flash_attention
-        and cfg.sample_packing
-    ):
-        from axolotl.monkeypatch.mistral_attn_hijack_flash import (
-            replace_mistral_attn_with_flash_attn,
-        )
-
-        LOG.info("patching mistral with flash attention")
-        replace_mistral_attn_with_flash_attn(packed=cfg.sample_packing)
-
-    if cfg.is_llama_derived_model and cfg.sample_packing and not inference:
-        from axolotl.monkeypatch.llama_expand_mask import hijack_expand_mask
-
-        LOG.info("patching _expand_mask")
-        hijack_expand_mask()
-
-    model_kwargs: Dict[str, Any] = {}
-
-    if cfg.model_kwargs:
-        for key, val in cfg.model_kwargs.items():
-            model_kwargs[key] = val
-
-    max_memory = cfg.max_memory
-    device_map = cfg.device_map
-
-    if cfg.gpu_memory_limit:
-        gpu_memory_limit = (
-            str(cfg.gpu_memory_limit) + "GiB"
-            if isinstance(cfg.gpu_memory_limit, int)
-            else cfg.gpu_memory_limit
-        )
-
-        max_memory = {}
-        for i in range(torch.cuda.device_count()):
-            max_memory[i] = gpu_memory_limit
-        max_memory["cpu"] = "256GiB"  # something sufficiently large to fit anything
-
-    if max_memory is not None:
-        # Based on https://github.com/togethercomputer/OpenChatKit/blob/main/inference/bot.py
-        from accelerate import infer_auto_device_map
-
-        with init_empty_weights():
-            model_canvas = AutoModelForCausalLM.from_config(
-                model_config, trust_remote_code=cfg.trust_remote_code or False
-            )
-        model_canvas.tie_weights()
-        device_map = infer_auto_device_map(
-            model_canvas,
-            max_memory=max_memory,
-            dtype=cfg.torch_dtype,
-        )
-        # We can discard max_memory now as we have a device map set up for us
-        max_memory = None
-
-    model_kwargs["device_map"] = device_map
-    model_kwargs["torch_dtype"] = cfg.torch_dtype
-
-    if torch.backends.mps.is_available():
-        model_kwargs["device_map"] = "mps:0"
-
-    # TODO can we put the reference model on it's own gpu? I think we have to move logits around to calculate loss
-    # if cfg.rl:
-    #     if torch.cuda.device_count() > 1:
-    #         if reference_model:
-    #             model_kwargs["device_map"] = "cuda:" + str(
-    #                 torch.cuda.current_device() + 1
-    #             )
-    #         else:
-    #             model_kwargs["device_map"] = "cuda:" + str(torch.cuda.current_device())
-
-    if is_deepspeed_zero3_enabled():
-        del model_kwargs["device_map"]
-
-    if cfg.revision_of_model:
-        model_kwargs["revision"] = cfg.revision_of_model
-
-    if cfg.gptq:
-        if not hasattr(model_config, "quantization_config"):
-            LOG.warning("model config does not contain quantization_config information")
-        else:
-            if cfg.gptq_disable_exllama is not None:
-                model_config.quantization_config[
-                    "disable_exllama"
-                ] = cfg.gptq_disable_exllama
-            model_kwargs["quantization_config"] = GPTQConfig(
-                **model_config.quantization_config
-            )
-    if cfg.adapter == "qlora" and cfg.load_in_4bit:
-        bnb_config = {
-            "load_in_4bit": True,
-            "llm_int8_threshold": 6.0,
-            "llm_int8_has_fp16_weight": False,
-            "bnb_4bit_compute_dtype": cfg.torch_dtype,
-            "bnb_4bit_use_double_quant": True,
-            "bnb_4bit_quant_type": "nf4",
-            "bnb_4bit_quant_storage": torch.bfloat16,
-        }
-        if cfg.model_config_type in ["jamba", "qwen2_moe"] and not cfg.deepspeed:
-            # for some reason, this causes the loss to be off by an order of magnitude
-            # but deepspeed needs this still in bfloat16
-            bnb_config["bnb_4bit_quant_storage"] = torch.float32
-
-        if cfg.bnb_config_kwargs:
-            bnb_config.update(cfg.bnb_config_kwargs)
-
-        model_kwargs["quantization_config"] = BitsAndBytesConfig(
-            **bnb_config,
-        )
-    elif cfg.adapter == "lora" and cfg.load_in_8bit:
-        bnb_config = {
-            "load_in_8bit": True,
-        }
-        # Exclude mamba blocks from int8 quantization for jamba
-        if cfg.model_config_type == "jamba":
-            bnb_config["llm_int8_skip_modules"] = ["mamba"]
-        model_kwargs["quantization_config"] = BitsAndBytesConfig(
-            **bnb_config,
-        )
-
-    if cfg.load_in_8bit and cfg.adapter is not None:
-        model_kwargs["load_in_8bit"] = True
-    if cfg.load_in_4bit and cfg.adapter is not None:
-        model_kwargs["load_in_4bit"] = True
-
-    # no longer needed per https://github.com/huggingface/transformers/pull/26610
-    if "quantization_config" in model_kwargs or cfg.gptq:
-        if "load_in_8bit" in model_kwargs:
-            del model_kwargs["load_in_8bit"]
-        if "load_in_4bit" in model_kwargs:
-            del model_kwargs["load_in_4bit"]
-
-    # sample packing uses custom FA2 patch
-    if cfg.flash_attention:
-        if not cfg.sample_packing:
-            if cfg.s2_attention:
-                pass
-            # most other models support flash attention, we can define exceptions as they come up
-            model_kwargs["attn_implementation"] = "flash_attention_2"
-            model_config._attn_implementation = (  # pylint: disable=protected-access
-                "flash_attention_2"
-            )
-        else:
-            if model_config.model_type in SUPPORTED_MULTIPACK_MODEL_TYPES:
-                model_kwargs["attn_implementation"] = "flash_attention_2"
-                model_config._attn_implementation = (  # pylint: disable=protected-access
-                    "flash_attention_2"
-                )
-            else:
-                model_kwargs["attn_implementation"] = "eager"
-                model_config._attn_implementation = (  # pylint: disable=protected-access
-                    "eager"
-                )
-    elif cfg.sdp_attention:
-        model_kwargs["attn_implementation"] = "sdpa"
-        model_config._attn_implementation = "sdpa"  # pylint: disable=protected-access
-    elif cfg.eager_attention:
-        model_kwargs["attn_implementation"] = "eager"
-        model_config._attn_implementation = "eager"  # pylint: disable=protected-access
-
-    if cfg.low_cpu_mem_usage:
-        model_kwargs["low_cpu_mem_usage"] = True
-
-    qlora_fsdp = cfg.fsdp and cfg.adapter == "qlora"
-
-    try:
-        skip_move_to_device = False
-        if (
-            cfg.fsdp and cfg.fsdp_config.fsdp_cpu_ram_efficient_loading
-        ) and not qlora_fsdp:
-            model = load_sharded_model(
-                base_model,
-                model_config,
-                cfg,
-                torch_dtype=cfg.torch_dtype,
-            )
-            skip_move_to_device = True
-        elif (
-            qlora_fsdp
-            and cfg.fsdp_config.fsdp_cpu_ram_efficient_loading
-            and cfg.model_config_type == "dbrx"
-        ):
-            quant_storage = cfg.torch_dtype
-            model = load_sharded_model_quant(
-                base_model,
-                model_config,
-                cfg,
-                quant_storage=quant_storage,
-            )
-            skip_move_to_device = True
-        elif (
-            model_config.model_type == "llama"
-            and not cfg.trust_remote_code
-            and not cfg.gptq
-        ):
-            from transformers import LlamaForCausalLM
-
-            model = LlamaForCausalLM.from_pretrained(
-                base_model,
-                config=model_config,
-                **model_kwargs,
-            )
-
-            if cfg.flash_attention and not inference:
-                from axolotl.monkeypatch.llama_attn_hijack_flash import (
-                    is_xformers_swiglu_available,
-                    replace_llama_mlp_with_swiglu,
-                    replace_llama_qkv_with_fused,
-                )
-
-                if cfg.flash_attn_fuse_mlp and is_xformers_swiglu_available():
-                    LOG.info("patching with SwiGLU")
-                    replace_llama_mlp_with_swiglu(model)
-
-                if cfg.flash_attn_fuse_qkv:
-                    LOG.info("patching with fused QKV")
-                    replace_llama_qkv_with_fused(model)
-        elif model_type == "MambaLMHeadModel":
-            # FIXME this is janky at best and hacked together to make it work
-            MambaLMHeadModel = fix_mamba_attn_for_loss()  # pylint: disable=invalid-name
-
-            model_kwargs["dtype"] = model_kwargs["torch_dtype"]
-            model_kwargs["device"] = torch.cuda.current_device()
-            del model_kwargs["torch_dtype"]
-            del model_kwargs["device_map"]
-
-            model = MambaLMHeadModel.from_pretrained(
-                base_model,
-                **model_kwargs,
-            )
-        elif model_type and not cfg.trust_remote_code:
-            if cfg.gptq:
-                model = AutoModelForCausalLM.from_pretrained(
-                    base_model,
-                    config=model_config,
-                    trust_remote_code=cfg.trust_remote_code or False,
-                    **model_kwargs,
-                )
-            else:
-                model = getattr(transformers, model_type).from_pretrained(
-                    base_model,
-                    config=model_config,
-                    trust_remote_code=cfg.trust_remote_code or False,
-                    **model_kwargs,
-                )
-        else:
-            # Shouldn't be a problem most of the time. will obviously error if the model doesn't support this
-            # when training starts
-            if (
-                hasattr(model_config, "max_seq_len")
-                and model_config.max_seq_len
-                and cfg.sequence_len > model_config.max_seq_len
-            ):
-                model_config.max_seq_len = cfg.sequence_len
-                LOG.warning(f"increasing context length to {cfg.sequence_len}")
-            elif (
-                hasattr(model_config, "max_sequence_length")
-                and model_config.max_sequence_length
-                and cfg.sequence_len > model_config.max_sequence_length
-            ):
-                model_config.max_sequence_length = cfg.sequence_len
-                LOG.warning(f"increasing context length to {cfg.sequence_len}")
-            if cfg.gptq:
-                model = AutoModelForCausalLM.from_pretrained(
-                    base_model,
-                    config=model_config,
-                    trust_remote_code=cfg.trust_remote_code or False,
-                    **model_kwargs,
-                )
-            else:
-                if qlora_fsdp and cfg.fsdp_config.fsdp_cpu_ram_efficient_loading:
-                    skip_move_to_device = True
-                    if "device_map" in model_kwargs:
-                        del model_kwargs["device_map"]
-
-                model = AutoModelForCausalLM.from_pretrained(
-                    base_model,
-                    config=model_config,
-                    trust_remote_code=cfg.trust_remote_code or False,
-                    **model_kwargs,
-                )
-    except Exception as err:  # pylint: disable=broad-exception-caught
-        LOG.exception(err)
-        raise err
-
-    if isinstance(model, (PeftModel, PeftModelForCausalLM)) and not qlora_fsdp:
-        model = model.merge_and_unload()
-
-    embeddings_len = (
-        math.ceil(len(tokenizer) / 32) * 32
-        if cfg.resize_token_embeddings_to_32x
-        else len(tokenizer)
-    )
-    if (
-        hasattr(model, "get_input_embeddings")
-        and model.get_input_embeddings().num_embeddings < embeddings_len
-    ):
-        model.resize_token_embeddings(embeddings_len)
-    else:
-        model.tie_weights()
-
-    if (
-        hasattr(model, "config")
-        and hasattr(model.config, "max_position_embeddings")
-        and model.config.max_position_embeddings
-        and cfg.sequence_len > model.config.max_position_embeddings
-    ):
-        LOG.warning(
-            f"increasing model.config.max_position_embeddings from {model.config.max_position_embeddings} to {cfg.sequence_len}"
-        )
-        model.config.max_position_embeddings = cfg.sequence_len
-
-    if (
-        hasattr(model, "config")
-        and hasattr(model.config, "bos_token_id")
-        and model.config.bos_token_id
-        and model.config.bos_token_id != tokenizer.bos_token_id
-    ):
-        model.config.bos_token_id = tokenizer.bos_token_id
-
-    if (
-        hasattr(model, "config")
-        and hasattr(model.config, "eos_token_id")
-        and model.config.eos_token_id
-        and model.config.eos_token_id != tokenizer.eos_token_id
-    ):
-        model.config.eos_token_id = tokenizer.eos_token_id
-
-    if hasattr(model, "device") and model.device.type in ("cuda", "mps"):
-        log_gpu_memory_usage(LOG, "after model load", model.device)
-
-    # make sure these are fp32 per Ramesh et al. (2021)
-    embedding_modules = get_linear_embedding_layers(cfg.model_config_type)
-    if not cfg.fsdp:
-        # FSDP doesn't like mixed Float and BFloat16
-        for name, module in model.named_modules():
-            if "norm" in name or name.endswith(".gate"):
-                module.to(torch.float32)
-            if model_config.model_type == "btlm":
-                # don't upcast lm_head for btlm
-                continue
-            if any(m in name for m in embedding_modules):
-                if hasattr(module, "weight"):
-                    module.to(torch.float32)
-
-    needs_fa2_dtype = cfg.adapter or cfg.fsdp
-    skip_prepare_model_for_kbit_training = False
-
-    if is_deepspeed_zero3_enabled():
-        from deepspeed.utils import (  # pylint: disable=no-name-in-module
-            set_z3_leaf_modules,
-        )
-
-        if cfg.model_config_type == "mixtral":
-            moe_block = get_module_class_from_name(model, "MixtralSparseMoeBlock")
-            set_z3_leaf_modules(model, [moe_block])
-        elif cfg.model_config_type == "dbrx":
-            moe_block = get_module_class_from_name(model, "DbrxFFN")
-            set_z3_leaf_modules(model, [moe_block])
-
-    if cfg.model_config_type == "qwen" and cfg.adapter == "lora":
-        # Qwen doesn't play nicely with LoRA if this is enabled
-        skip_prepare_model_for_kbit_training = True
-
-    loftq_bits = cfg.peft and cfg.peft.loftq_config and cfg.peft.loftq_config.loftq_bits
-    if cfg.adapter == "lora" and loftq_bits:
-        skip_prepare_model_for_kbit_training = True
-
-    if qlora_fsdp or (cfg.fsdp and cfg.fsdp_config.fsdp_cpu_ram_efficient_loading):
-        # make sure everything is in the same dtype
-        skip_prepare_model_for_kbit_training = True
-
-    if cfg.adapter in ["lora", "qlora"]:
-        if cfg.gradient_checkpointing:
-            model.gradient_checkpointing_enable(
-                gradient_checkpointing_kwargs=cfg.gradient_checkpointing_kwargs
-            )
-        if (
-            cfg.load_in_8bit or cfg.load_in_4bit
-        ) and not skip_prepare_model_for_kbit_training:
-            LOG.info("converting PEFT model w/ prepare_model_for_kbit_training")
-            model = prepare_model_for_kbit_training(
-                model, use_gradient_checkpointing=cfg.gradient_checkpointing
-            )
-        needs_fa2_dtype = True
-
-    # LlamaRMSNorm layers are in fp32 after kbit_training or full finetune, so we need to
-    # convert them back to fp16/bf16 for flash-attn compatibility.
-    if (needs_fa2_dtype or cfg.flash_attention) and not qlora_fsdp:
-        LOG.info("converting modules to %s for flash attention", cfg.torch_dtype)
-        for name, module in model.named_modules():
-            if "norm" in name:
-                module.to(cfg.torch_dtype)
-            if any(m in name for m in embedding_modules):
-                if hasattr(module, "weight"):
-                    module.to(cfg.torch_dtype)
-
-    lora_config = None
-    if not reference_model or cfg.lora_model_dir:
-        # if we're not loading the reference model, then we're loading the model for training
-        # then the dpo trainer doesn't want the peft model loaded over it, it just wants the lora/peft config
-        if cfg.adapter and cfg.rl in ["dpo", "ipo", "kto_pair"] and not cfg.merge_lora:
-            _, lora_config = load_lora(model, cfg, inference=False, config_only=True)
-        else:
-            model, lora_config = load_adapter(model, cfg, cfg.adapter)
-
-    if (
-        cfg.ddp
-        and not load_in_8bit
-        and not (cfg.rl and cfg.load_in_4bit)
-        and not skip_move_to_device
-    ):
-        # TODO revaldate this conditional
-        model.to(f"cuda:{cfg.local_rank}")
-
-    if torch.cuda.device_count() > 1 and int(os.getenv("WORLD_SIZE", "1")) == 1:
-        setattr(model, "is_parallelizable", True)
-        setattr(model, "model_parallel", True)
-
-    requires_grad = []
-    for name, param in model.named_parameters(recurse=True):
-        if param.requires_grad:
-            requires_grad.append(f"{name}: {param.requires_grad}")
-    if len(requires_grad) == 0:
-        LOG.warning("there are no parameters that require gradient updates")
-    if hasattr(model, "config"):
-        model.config.use_cache = False
-
-    if cfg.flash_optimum:
-        from optimum.bettertransformer import BetterTransformer
-
-        model = BetterTransformer.transform(model)
-
-    if cfg.adapter is not None:
-        log_gpu_memory_usage(LOG, "after adapters", model.device)
-
-    # TODO resume_from_checkpoint handling
-    return model, lora_config
-
-
-def load_adapter(model, cfg, adapter, inference=False):
-    # type: (PreTrainedModel, DictDefault, Optional[str], bool) -> Tuple[PreTrainedModel, Optional[PeftConfig]]
-
-    if adapter is None:
-        return model, None
-    if hasattr(model, "enable_input_require_grads"):
-        model.enable_input_require_grads()
-    if adapter in ["lora", "qlora"]:
-        return load_lora(model, cfg, inference=inference)
-    if adapter == "llama-adapter":
-        return load_llama_adapter(model, cfg)
-
-    raise NotImplementedError(f"{adapter} peft adapter not available")
-
-
-def load_llama_adapter(model, cfg):
-    # type: (PreTrainedModel, DictDefault) -> Tuple[PreTrainedModel, Optional[PeftConfig]]
-    from peft import AdaptionPromptConfig, get_peft_model
-
-    peft_config = AdaptionPromptConfig(
-        adapter_layers=cfg.peft_adapter.layers,  # layers (L)
-        adapter_len=cfg.peft_adapter.len,  # prompt length (K)
-        task_type="CAUSAL_LM",
-    )
-
-    if cfg.lora_model_dir:
-        LOG.debug("Loading pretrained PEFT - llama_adapter")
-        model = PeftModel.from_pretrained(
-            model,
-            cfg.lora_model_dir,
-            torch_dtype=torch.float16,
-        )
-    else:
-        model = get_peft_model(model, peft_config)
-
-    model.print_trainable_parameters()
-
-    return model, peft_config
-
-
-def find_all_linear_names(model):
-    cls = (bnb.nn.Linear4bit, bnb.nn.Linear8bitLt, torch.nn.Linear, QuantLinear)
-    lora_module_names = set()
-    for name, module in model.named_modules():
-        if (
-            isinstance(module, cls)
-            or "Linear" in module.__class__.__name__
-            and module.__class__.__name__ not in ("LlamaLinearScalingRotaryEmbedding",)
-        ):
-            names = name.split(".")
-            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
-
-    embedding_modules = get_linear_embedding_layers(model.config.model_type)
-    output_embedding = embedding_modules[1]
-    if output_embedding in lora_module_names:  # needed for 16-bit
-        lora_module_names.remove(output_embedding)
-
-    return list(lora_module_names)
-
-
-def setup_quantized_meta_for_peft(model: nn.Module):
-    """Replaces `quant_state.to` with a dummy function to prevent PEFT from moving `quant_state` to meta device"""
-
-    def temp_to_method(self, *args, **kwargs):  # pylint: disable=unused-argument
-        return self
-
-    for param in model.parameters():
-        if isinstance(param, Params4bit):
-            param.quant_state._orig_to = (  # pylint: disable=protected-access
-                param.quant_state.to
-            )
-            param.quant_state.to = types.MethodType(temp_to_method, param.quant_state)
-
-
-def setup_quantized_peft_meta_for_training(model: nn.Module):
-    """Replaces dummy `quant_state.to` method with the original function to allow training to continue"""
-    for param in model.parameters():
-        if isinstance(param, Params4bit) and hasattr(param.quant_state, "_orig_to"):
-            param.quant_state.to = (
-                param.quant_state._orig_to  # pylint: disable=protected-access
-            )
-            param.quant_state._orig_to = None  # pylint: disable=protected-access
-
-
-def load_lora(model, cfg, inference=False, config_only=False):
-    # type: (PreTrainedModel, DictDefault, bool, bool) -> Tuple[Optional[PreTrainedModel], Optional[PeftConfig]]
-
-    from peft import LoraConfig, get_peft_model
-
-    lora_target_modules = list(cfg.lora_target_modules or [])
-
-    if cfg.lora_target_linear:
-        linear_names = find_all_linear_names(model)
-        LOG.info(f"found linear modules: {repr(linear_names)}")
-        lora_target_modules = list(set(lora_target_modules + linear_names))
-
-    lora_config_kwargs = {}
-    loftq_bits = cfg.peft and cfg.peft.loftq_config and cfg.peft.loftq_config.loftq_bits
-    if loftq_bits:
-        lora_config_kwargs["loftq_config"] = LoftQConfig(loftq_bits=loftq_bits)
-        lora_config_kwargs["init_lora_weights"] = "loftq"
-    if cfg.peft_use_dora:
-        lora_config_kwargs["use_dora"] = cfg.peft_use_dora
-    if cfg.peft_use_rslora:
-        lora_config_kwargs["use_rslora"] = cfg.peft_use_rslora
-    if cfg.peft_layer_replication:
-        lora_config_kwargs["layer_replication"] = cfg.peft_layer_replication
-
-    lora_config = LoraConfig(
-        r=cfg.lora_r,
-        lora_alpha=cfg.lora_alpha,
-        target_modules=lora_target_modules,
-        layers_to_transform=cfg.peft_layers_to_transform,
-        lora_dropout=cfg.lora_dropout,
-        fan_in_fan_out=cfg.lora_fan_in_fan_out,
-        modules_to_save=cfg.lora_modules_to_save if cfg.lora_modules_to_save else None,
-        bias="none",
-        task_type="CAUSAL_LM",
-        **lora_config_kwargs,
-    )
-
-    if config_only:
-        return None, lora_config
-
-    rank = int(os.environ.get("LOCAL_RANK", 0))
-
-    if (
-        cfg.fsdp
-        and cfg.adapter
-        and cfg.fsdp_config.fsdp_cpu_ram_efficient_loading
-        and rank != 0
-    ):
-        setup_quantized_meta_for_peft(model)
-
-    if cfg.lora_model_dir:
-        LOG.debug("Loading pretrained PEFT - LoRA")
-        model_kwargs: Any = {}
-        if cfg.lora_on_cpu:
-            model_kwargs["max_memory"] = {"cpu": "256GiB"}
-            model_kwargs["device_map"] = {"": "cpu"}
-        model = PeftModel.from_pretrained(
-            model,
-            cfg.lora_model_dir,
-            is_trainable=(not inference),
-            **model_kwargs,
-        )
-    else:
-        model = get_peft_model(model, lora_config)
-
-    if rank == 0:
-        try:
-            model.print_trainable_parameters()
-        except AttributeError as exc:
-            LOG.warning(
-                "Exception caught during model.print_trainable_parameters(): %s", exc
-            )
-    elif (
-        cfg.fsdp
-        and cfg.adapter
-        and cfg.fsdp_config.fsdp_cpu_ram_efficient_loading
-        and rank != 0
-    ):
-        setup_quantized_peft_meta_for_training(model)
-
-    return model, lora_config
-
-
-def ensure_dtype(model, dtype=torch.bfloat16):
-    for name, module in model.named_modules():
-        try:
-            if module.weight.dtype != dtype:
-                print(f"Converting module {name}: {module.weight.dtype} -> {dtype}")
-                module.to(dtype)
-        except AttributeError:
-            pass
diff --git a/src/axolotl/utils/optimizers/__init__.py b/src/axolotl/utils/optimizers/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/utils/optimizers/adopt.py b/src/axolotl/utils/optimizers/adopt.py
new file mode 100644
index 0000000000..f501859b6f
--- /dev/null
+++ b/src/axolotl/utils/optimizers/adopt.py
@@ -0,0 +1,550 @@
+"""
+Copied from https://github.com/iShohei220/adopt
+
+ADOPT: Modified Adam Can Converge with Any β2 with the Optimal Rate (2024)
+Taniguchi, Shohei and Harada, Keno and Minegishi, Gouki and Oshima, Yuta and Jeong, Seong Cheol and Nagahara, Go and Iiyama, Tomoshi and Suzuki, Masahiro and Iwasawa, Yusuke and Matsuo, Yutaka
+"""
+
+# mypy: ignore-errors
+# flake8: noqa
+# mypy: allow-untyped-decorators
+# mypy: allow-untyped-defs
+from typing import Callable, List, Optional, Tuple, Union, cast
+
+import torch
+from torch import Tensor
+from torch.optim.optimizer import (  # DeviceDict,; _capturable_doc,; _differentiable_doc,; _foreach_doc,; _fused_doc,; _maximize_doc,; _stack_if_compiling,
+    DeviceDict,
+    Optimizer,
+    ParamsT,
+    _capturable_doc,
+    _default_to_fused_or_foreach,
+    _device_dtype_check_for_fused,
+    _differentiable_doc,
+    _disable_dynamo_if_unsupported,
+    _foreach_doc,
+    _fused_doc,
+    _get_capturable_supported_devices,
+    _get_scalar_dtype,
+    _get_value,
+    _maximize_doc,
+    _stack_if_compiling,
+    _use_grad_for_differentiable,
+    _view_as_real,
+)
+
+__all__ = ["ADOPT", "adopt"]
+
+
+class ADOPT(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.9999),
+        eps: float = 1e-6,
+        clip_lambda: Optional[Callable[[int], float]] = lambda step: step**0.25,
+        weight_decay: float = 0.0,
+        decouple: bool = False,
+        *,
+        foreach: Optional[bool] = None,
+        maximize: bool = False,
+        capturable: bool = False,
+        differentiable: bool = False,
+        fused: Optional[bool] = None,
+    ):
+        if isinstance(lr, Tensor):
+            if foreach and not capturable:
+                raise ValueError(
+                    "lr as a Tensor is not supported for capturable=False and foreach=True"
+                )
+            if lr.numel() != 1:
+                raise ValueError("Tensor lr must be 1-element")
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        self.clip_lambda = clip_lambda
+
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            decouple=decouple,
+            maximize=maximize,
+            foreach=foreach,
+            capturable=capturable,
+            differentiable=differentiable,
+            fused=fused,
+        )
+        super().__init__(params, defaults)
+
+        if fused:
+            # TODO: support fused
+            raise RuntimeError("`fused` is not currently supported")
+
+            if differentiable:
+                raise RuntimeError("`fused` does not support `differentiable`")
+            self._step_supports_amp_scaling = True
+            # TODO(crcrpar): [low prec params & their higher prec copy]
+            # Support AMP with FP16/BF16 model params which would need
+            # higher prec copy of params to do update math in higher prec to
+            # alleviate the loss of information.
+            if foreach:
+                raise RuntimeError("`fused` and `foreach` cannot be `True` together.")
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("maximize", False)
+            group.setdefault("foreach", None)
+            group.setdefault("capturable", False)
+            group.setdefault("differentiable", False)
+            fused = group.setdefault("fused", None)
+            for p in group["params"]:
+                p_state = self.state.get(p, [])
+                if len(p_state) != 0 and not torch.is_tensor(p_state["step"]):
+                    step_val = float(p_state["step"])
+                    p_state["step"] = (
+                        torch.tensor(
+                            step_val,
+                            dtype=_get_scalar_dtype(is_fused=fused),
+                            device=p.device,
+                        )
+                        if group["capturable"] or group["fused"]
+                        else torch.tensor(step_val, dtype=_get_scalar_dtype())
+                    )
+
+    def _init_group(
+        self,
+        group,
+        params_with_grad,
+        grads,
+        exp_avgs,
+        exp_avg_sqs,
+        state_steps,
+    ):
+        has_complex = False
+        for p in group["params"]:
+            if p.grad is not None:
+                has_complex |= torch.is_complex(p)
+                params_with_grad.append(p)
+                if p.grad.is_sparse:
+                    raise RuntimeError("ADOPT does not support sparse gradients")
+                grads.append(p.grad)
+
+                state = self.state[p]
+                # Lazy state initialization
+                if len(state) == 0:
+                    if group["fused"]:
+                        _device_dtype_check_for_fused(p)
+                    # note(crcrpar): [special device hosting for step]
+                    # Deliberately host `step` on CPU if both capturable and fused are off.
+                    # This is because kernel launches are costly on CUDA and XLA.
+                    state["step"] = (
+                        torch.zeros(
+                            (),
+                            dtype=_get_scalar_dtype(is_fused=group["fused"]),
+                            device=p.device,
+                        )
+                        if group["capturable"] or group["fused"]
+                        else torch.tensor(0.0, dtype=_get_scalar_dtype())
+                    )
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+
+                exp_avgs.append(state["exp_avg"])
+                exp_avg_sqs.append(state["exp_avg_sq"])
+
+                if group["differentiable"] and state["step"].requires_grad:
+                    raise RuntimeError(
+                        "`requires_grad` is not supported for `step` in differentiable mode"
+                    )
+
+                # Foreach without capturable does not support a tensor lr
+                if (
+                    group["foreach"]
+                    and torch.is_tensor(group["lr"])
+                    and not group["capturable"]
+                ):
+                    raise RuntimeError(
+                        "lr as a Tensor is not supported for capturable=False and foreach=True"
+                    )
+
+                state_steps.append(state["step"])
+        return has_complex
+
+    @_use_grad_for_differentiable
+    def step(self, closure=None):
+        """Perform a single optimization step.
+
+        Args:
+            closure (Callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        # torch 2.11 renamed _cuda_graph_capture_health_check ->
+        # _accelerator_graph_capture_health_check (the 2.11-only name); 2.12
+        # re-added the old name as an alias. Prefer the new name, fall back.
+        health_check = getattr(
+            self, "_accelerator_graph_capture_health_check", None
+        ) or getattr(self, "_cuda_graph_capture_health_check", None)
+        if health_check is not None:
+            health_check()
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad: List[Tensor] = []
+            grads: List[Tensor] = []
+            exp_avgs: List[Tensor] = []
+            exp_avg_sqs: List[Tensor] = []
+            state_steps: List[Tensor] = []
+            beta1, beta2 = group["betas"]
+
+            has_complex = self._init_group(
+                group,
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+            )
+
+            adopt(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                has_complex=has_complex,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                clip_lambda=self.clip_lambda,
+                weight_decay=group["weight_decay"],
+                decouple=group["decouple"],
+                eps=group["eps"],
+                maximize=group["maximize"],
+                foreach=group["foreach"],
+                capturable=group["capturable"],
+                differentiable=group["differentiable"],
+                fused=group["fused"],
+                grad_scale=getattr(self, "grad_scale", None),
+                found_inf=getattr(self, "found_inf", None),
+            )
+
+        return loss
+
+
+def _single_tensor_adopt(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    grad_scale: Optional[Tensor],
+    found_inf: Optional[Tensor],
+    *,
+    has_complex: bool,
+    beta1: float,
+    beta2: float,
+    lr: Union[float, Tensor],
+    clip_lambda: Optional[Callable[[int], float]],
+    weight_decay: float,
+    decouple: bool,
+    eps: float,
+    maximize: bool,
+    capturable: bool,
+    differentiable: bool,
+):
+    assert grad_scale is None and found_inf is None
+
+    if torch.jit.is_scripting():
+        # this assert is due to JIT being dumb and not realizing that the ops below
+        # have overloads to handle both float and Tensor lrs, so we just assert it's
+        # a float since most people using JIT are using floats
+        assert isinstance(lr, float)
+
+    for i, param in enumerate(params):
+        grad = grads[i] if not maximize else -grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        step_t = state_steps[i]
+
+        # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable]
+        if not torch.compiler.is_compiling() and capturable:
+            capturable_supported_devices = _get_capturable_supported_devices()
+            assert (
+                param.device.type == step_t.device.type
+                and param.device.type in capturable_supported_devices
+            ), (
+                f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}."
+            )
+
+        step = step_t if capturable or differentiable else _get_value(step_t)
+
+        if weight_decay != 0 and not decouple:
+            grad = grad.add(param, alpha=weight_decay)
+
+        if torch.is_complex(param):
+            grad = torch.view_as_real(grad)
+            if exp_avg is not None:
+                exp_avg = torch.view_as_real(exp_avg)
+            if exp_avg_sq is not None:
+                exp_avg_sq = torch.view_as_real(exp_avg_sq)
+            param = torch.view_as_real(param)
+
+        if step == 0:
+            exp_avg_sq.addcmul_(grad, grad.conj())
+            # update step
+            step_t += 1
+            continue
+
+        if weight_decay != 0 and decouple:
+            param.add_(param, alpha=-lr * weight_decay)
+
+        denom = torch.clamp(exp_avg_sq.sqrt(), eps)
+        normed_grad = grad.div(denom)
+        if clip_lambda is not None:
+            clip = clip_lambda(step)
+            normed_grad.clamp_(-clip, clip)
+
+        exp_avg.lerp_(normed_grad, 1 - beta1)
+
+        param.add_(exp_avg, alpha=-lr)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
+
+        # update step
+        step_t += 1
+
+
+def _multi_tensor_adopt(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    grad_scale: Optional[Tensor],
+    found_inf: Optional[Tensor],
+    *,
+    has_complex: bool,
+    beta1: float,
+    beta2: float,
+    lr: Union[float, Tensor],
+    clip_lambda: Optional[Callable[[int], float]],
+    weight_decay: float,
+    decouple: bool,
+    eps: float,
+    maximize: bool,
+    capturable: bool,
+    differentiable: bool,
+):
+    if len(params) == 0:
+        return
+
+    if isinstance(lr, Tensor) and not capturable:
+        raise RuntimeError(
+            "lr as a Tensor is not supported for capturable=False and foreach=True"
+        )
+
+    # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable]
+    if not torch.compiler.is_compiling() and capturable:
+        capturable_supported_devices = _get_capturable_supported_devices(
+            supports_xla=False
+        )
+        assert all(
+            p.device.type == step.device.type
+            and p.device.type in capturable_supported_devices
+            for p, step in zip(params, state_steps)
+        ), (
+            f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}."
+        )
+
+    assert grad_scale is None and found_inf is None
+
+    assert not differentiable, "_foreach ops don't support autograd"
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]  # type: ignore[list-item]
+    )
+    for (
+        device_params_,
+        device_grads_,
+        device_exp_avgs_,
+        device_exp_avg_sqs_,
+        device_state_steps_,
+    ), _ in grouped_tensors.values():
+        device_params = cast(List[Tensor], device_params_)
+        device_grads = cast(List[Tensor], device_grads_)
+        device_exp_avgs = cast(List[Tensor], device_exp_avgs_)
+        device_exp_avg_sqs = cast(List[Tensor], device_exp_avg_sqs_)
+        device_state_steps = cast(List[Tensor], device_state_steps_)
+
+        # Handle complex parameters
+        if has_complex:
+            _view_as_real(
+                device_params, device_grads, device_exp_avgs, device_exp_avg_sqs
+            )
+
+        if maximize:
+            device_grads = torch._foreach_neg(device_grads)  # type: ignore[assignment]
+
+        if weight_decay != 0 and not decouple:
+            # Re-use the intermediate memory (device_grads) already allocated for maximize
+            if maximize:
+                torch._foreach_add_(device_grads, device_params, alpha=weight_decay)
+            else:
+                device_grads = torch._foreach_add(  # type: ignore[assignment]
+                    device_grads, device_params, alpha=weight_decay
+                )
+
+        if device_state_steps[0] == 0:
+            torch._foreach_addcmul_(device_exp_avg_sqs, device_grads, device_grads)
+
+            # Update steps
+            # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over
+            # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just
+            # wrapped it once now. The alpha is required to assure we go to the right overload.
+            if not torch.compiler.is_compiling() and device_state_steps[0].is_cpu:
+                torch._foreach_add_(
+                    device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0
+                )
+            else:
+                torch._foreach_add_(device_state_steps, 1)
+
+            continue
+
+        if weight_decay != 0 and decouple:
+            torch._foreach_add_(device_params, device_params, alpha=-lr * weight_decay)
+
+        exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs)
+        torch._foreach_maximum_(exp_avg_sq_sqrt, eps)
+
+        normed_grad = torch._foreach_div(device_grads, exp_avg_sq_sqrt)
+        if clip_lambda is not None:
+            clip = clip_lambda(device_state_steps[0])
+            torch._foreach_maximum_(normed_grad, -clip)
+            torch._foreach_minimum_(normed_grad, clip)
+
+        torch._foreach_lerp_(device_exp_avgs, normed_grad, 1 - beta1)
+
+        torch._foreach_add_(device_params, device_exp_avgs, alpha=-lr)
+        torch._foreach_mul_(device_exp_avg_sqs, beta2)
+        torch._foreach_addcmul_(
+            device_exp_avg_sqs, device_grads, device_grads, value=1 - beta2
+        )
+
+        # Update steps
+        # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over
+        # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just
+        # wrapped it once now. The alpha is required to assure we go to the right overload.
+        if not torch.compiler.is_compiling() and device_state_steps[0].is_cpu:
+            torch._foreach_add_(
+                device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0
+            )
+        else:
+            torch._foreach_add_(device_state_steps, 1)
+
+
+@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adopt)
+def adopt(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627
+    # setting this as kwarg for now as functional API is compiled by torch/distributed/optim
+    foreach: Optional[bool] = None,
+    capturable: bool = False,
+    differentiable: bool = False,
+    fused: Optional[bool] = None,
+    grad_scale: Optional[Tensor] = None,
+    found_inf: Optional[Tensor] = None,
+    has_complex: bool = False,
+    *,
+    beta1: float,
+    beta2: float,
+    lr: Union[float, Tensor],
+    clip_lambda: Optional[Callable[[int], float]],
+    weight_decay: float,
+    decouple: bool,
+    eps: float,
+    maximize: bool,
+):
+    r"""Functional API that performs ADOPT algorithm computation."""
+    # Respect when the user inputs False/True for foreach or fused. We only want to change
+    # the default when neither have been user-specified. Note that we default to foreach
+    # and pass False to use_fused. This is not a mistake--we want to give the fused impl
+    # bake-in time before making it the default, even if it is typically faster.
+    if fused is None and foreach is None:
+        _, foreach = _default_to_fused_or_foreach(
+            params, differentiable, use_fused=False
+        )
+        # Do not flip on foreach for the unsupported case where lr is a Tensor and capturable=False.
+        if foreach and isinstance(lr, Tensor) and not capturable:
+            foreach = False
+    if fused is None:
+        fused = False
+    if foreach is None:
+        foreach = False
+
+    # this check is slow during compilation, so we skip it
+    # if it's strictly needed we can add this check back in dynamo
+    if not torch.compiler.is_compiling() and not all(
+        isinstance(t, torch.Tensor) for t in state_steps
+    ):
+        raise RuntimeError(
+            "API has changed, `state_steps` argument must contain a list of singleton tensors"
+        )
+
+    if foreach and torch.jit.is_scripting():
+        raise RuntimeError("torch.jit.script not supported with foreach optimizers")
+    if fused and torch.jit.is_scripting():
+        raise RuntimeError("torch.jit.script not supported with fused optimizers")
+
+    # if fused and not torch.jit.is_scripting():
+    #     func = _fused_adopt
+    # elif foreach and not torch.jit.is_scripting():
+    if foreach and not torch.jit.is_scripting():
+        func = _multi_tensor_adopt
+    else:
+        func = _single_tensor_adopt
+
+    func(
+        params,
+        grads,
+        exp_avgs,
+        exp_avg_sqs,
+        state_steps,
+        has_complex=has_complex,
+        beta1=beta1,
+        beta2=beta2,
+        lr=lr,
+        clip_lambda=clip_lambda,
+        weight_decay=weight_decay,
+        decouple=decouple,
+        eps=eps,
+        maximize=maximize,
+        capturable=capturable,
+        differentiable=differentiable,
+        grad_scale=grad_scale,
+        found_inf=found_inf,
+    )
diff --git a/src/axolotl/utils/optimizers/qgalore.py b/src/axolotl/utils/optimizers/qgalore.py
new file mode 100644
index 0000000000..9e2cc82607
--- /dev/null
+++ b/src/axolotl/utils/optimizers/qgalore.py
@@ -0,0 +1,88 @@
+"""Helpers for the Q-GaLore optimizer integration."""
+
+from __future__ import annotations
+
+import inspect
+import types
+
+from torch import nn
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_q_galore_for_modern_bnb() -> None:
+    """bnb >=0.44 inserted (beta3, alpha) into ``optimizer_update_8bit_blockwise``
+    and ``optimizer_update_32bit``; q-galore-torch==1.0 still calls the legacy
+    positional layout. Swap q_galore's bnb handle for one that re-emits the
+    modern layout. No-op on older bnb."""
+    import bitsandbytes.functional as F
+    import q_galore_torch.q_galore_adamw8bit as mod
+
+    if "beta3" not in inspect.signature(F.optimizer_update_8bit_blockwise).parameters:
+        return
+
+    bw, fp32 = F.optimizer_update_8bit_blockwise, F.optimizer_update_32bit
+    mod.F = types.SimpleNamespace(
+        optimizer_update_8bit_blockwise=(
+            lambda *a, **kw: bw(
+                *(a[:7] + (0.0, 0.0) + a[7:] if len(a) == 15 else a), **kw
+            )
+        ),
+        optimizer_update_32bit=(
+            lambda *a, **kw: fp32(
+                *(a[:10] + (0.0, 0.0) + a[10:] if len(a) == 13 else a), **kw
+            )
+        ),
+        optimizer_update_8bit=F.optimizer_update_8bit,
+        percentile_clipping=F.percentile_clipping,
+    )
+
+
+def build_qgalore_param_groups(
+    model: nn.Module,
+    target_modules: list[str],
+    *,
+    rank: int,
+    update_proj_gap: int,
+    scale: float,
+    proj_type: str,
+    proj_quant: bool,
+    proj_bits: int,
+    proj_group_size: int,
+    cos_threshold: float,
+    gamma_proj: int,
+    queue_size: int,
+) -> list[dict]:
+    """Two param-groups: 2D weights matching ``target_modules`` get the Q-GaLore
+    projection keys; everything else (norms, biases, embeddings) is plain AdamW."""
+    galore, plain = [], []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if p.dim() == 2 and any(t in name for t in target_modules):
+            galore.append(p)
+        else:
+            plain.append(p)
+    if not galore:
+        raise ValueError(
+            f"Q-GaLore: no parameters matched optim_target_modules={target_modules!r}"
+        )
+    LOG.info("Q-GaLore param groups: %d projected, %d plain", len(galore), len(plain))
+    return [
+        {
+            "params": galore,
+            "rank": rank,
+            "update_proj_gap": update_proj_gap,
+            "scale": scale,
+            "proj_type": proj_type,
+            "quant": proj_quant,
+            "quant_n_bit": proj_bits,
+            "quant_group_size": proj_group_size,
+            "cos_threshold": cos_threshold,
+            "gamma_proj": gamma_proj,
+            "queue_size": queue_size,
+        },
+        {"params": plain},
+    ]
diff --git a/src/axolotl/utils/optimizers/sinkgd.py b/src/axolotl/utils/optimizers/sinkgd.py
new file mode 100644
index 0000000000..46aad07d4d
--- /dev/null
+++ b/src/axolotl/utils/optimizers/sinkgd.py
@@ -0,0 +1,444 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SinkGD: stateless gradient multi-normalization optimizer.
+
+Implements the SinkGD optimizer from "Gradient Multi-Normalization for Stateless
+and Scalable LLM Training" (Scetbon et al., 2025, https://arxiv.org/abs/2502.06742).
+
+Linear weight matrices are updated with the stateless SR-Sinkhorn procedure
+(Algorithm 3/4): the raw gradient is alternately row- and column-normalized for a
+few iterations, then applied as the update. No optimizer state is stored for these
+parameters. Embeddings, the output head, and 1D parameters (norms, biases) fall
+back to AdamW, reusing torchao's low-bit optimizer base so that this state is kept
+in 8-bit and the overall footprint stays close to plain SGD.
+"""
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torchao.optim.adam import _AdamBase, single_param_adam
+from torchao.optim.quant_utils import _fp32_to_bf16_sr
+from torchao.optim.subclass_8bit import OptimState8bit
+
+from axolotl.integrations.base import BaseOptimizerFactory
+
+
+def sr_sinkhorn(grad: Tensor, iters: int, eps: float) -> Tensor:
+    """SR-Sinkhorn (Algorithm 3): alternate row/column L2 normalization.
+
+    Each iteration rescales rows to L2 norm sqrt(n) then columns to L2 norm
+    sqrt(m), driving the gradient towards a doubly-balanced fixed point.
+
+    Operates on the last two dims, so any leading dims (e.g. the expert axis of
+    a fused MoE weight ``[num_experts, in, out]``) are treated as an independent
+    batch and each matrix is normalized separately.
+    """
+    m, n = grad.shape[-2], grad.shape[-1]
+    sqrt_n = n**0.5
+    sqrt_m = m**0.5
+    x = grad
+    for _ in range(iters):
+        x = x * (sqrt_n / x.norm(dim=-1, keepdim=True).clamp_min(eps))
+        x = x * (sqrt_m / x.norm(dim=-2, keepdim=True).clamp_min(eps))
+    return x
+
+
+def single_param_sinkgd(
+    p: Tensor,
+    grad: Tensor,
+    lr: Tensor,
+    weight_decay: float,
+    iters: int,
+    eps: float,
+    stochastic_round: bool,
+) -> None:
+    """Fused stateless SinkGD update for one weight; meant to be torch.compiled.
+
+    Normalization runs over the last two dims (leading dims, e.g. the expert axis
+    of a fused MoE weight, are batched and normalized per-matrix). The iterate
+    stays in the gradient dtype to halve memory traffic while norm reductions
+    accumulate in fp32. `iters`/`eps`/`weight_decay` are constants so the loop
+    unrolls, and `lr` is a tensor input so a changing schedule never recompiles.
+    """
+    m, n = grad.shape[-2], grad.shape[-1]
+    sqrt_n = n**0.5
+    sqrt_m = m**0.5
+    x = grad
+    for _ in range(iters):
+        rn = torch.linalg.vector_norm(x, dim=-1, keepdim=True, dtype=torch.float32)
+        x = x * (sqrt_n / rn.clamp_min(eps)).to(x.dtype)
+        cn = torch.linalg.vector_norm(x, dim=-2, keepdim=True, dtype=torch.float32)
+        x = x * (sqrt_m / cn.clamp_min(eps)).to(x.dtype)
+
+    p_f32 = p.float()
+    if weight_decay != 0.0:
+        p_f32 = p_f32 * (1 - lr * weight_decay)
+    p_f32 = p_f32 - lr * x.float()
+
+    if stochastic_round:
+        p.copy_(_fp32_to_bf16_sr(p_f32))
+    else:
+        p.copy_(p_f32)
+
+
+class SinkGD(_AdamBase):
+    """SinkGD optimizer with an 8-bit AdamW fallback for non-matrix parameters.
+
+    Parameter groups flagged with ``use_sinkgd=True`` get the stateless SR-Sinkhorn
+    update for their >=2D tensors (3D fused-MoE-expert weights are normalized
+    per-expert); every other parameter is optimized with the quantized 8-bit AdamW
+    step inherited from torchao's ``_AdamBase``.
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=0.0,
+        *,
+        sinkhorn_iters=5,
+        sinkgd_lr_scale=0.05,
+        sinkgd_eps=1e-8,
+        block_size=256,
+        bf16_stochastic_round=False,
+    ) -> None:
+        super().__init__(
+            params,
+            lr,
+            betas,
+            eps,
+            weight_decay,
+            amsgrad=False,
+            block_size=block_size,
+            bf16_stochastic_round=bf16_stochastic_round,
+            is_adamw=True,
+        )
+        self.sinkhorn_iters = sinkhorn_iters
+        self.sinkgd_lr_scale = sinkgd_lr_scale
+        self.sinkgd_eps = sinkgd_eps
+        self._compiled_sinkgd = torch.compile(
+            single_param_sinkgd, fullgraph=True, dynamic=False
+        )
+        self._compiled_adam = torch.compile(
+            single_param_adam, fullgraph=True, dynamic=False
+        )
+
+    @staticmethod
+    def _subclass_zeros(p: Tensor, signed: bool, block_size: int):
+        return OptimState8bit.zeros(
+            p.shape, signed, block_size, p.device, dtype=p.dtype
+        )
+
+    def _adam_fallback(self, p: Tensor, grad: Tensor, group: dict, lr: Tensor) -> None:
+        # Adam is elementwise, so run on the local shard: no cross-rank comm, and it
+        # sidesteps the OptimState8bit + DTensor + compile dispatch gap under FSDP2.
+        # State is keyed by the original param (to_local returns a fresh object).
+        p_local = p.to_local() if isinstance(p, DTensor) else p
+        grad_local = grad.to_local() if isinstance(grad, DTensor) else grad
+        state = self.state[p]
+        if len(state) == 0:
+            state["step"] = torch.tensor(0.0)
+            state["exp_avg"] = self._new_buffer(p_local, True)
+            state["exp_avg_sq"] = self._new_buffer(p_local, False)
+        state["step"] += 1
+        self._compiled_adam(
+            p_local.detach(),
+            grad_local,
+            state["step"],
+            state["exp_avg"],
+            state["exp_avg_sq"],
+            None,
+            lr,
+            group["betas"][0],
+            group["betas"][1],
+            group["weight_decay"],
+            group["eps"],
+            self.is_adamw,
+            self.bf16_stochastic_round and p_local.dtype is torch.bfloat16,
+        )
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        with torch._dynamo.utils.disable_cache_limit():
+            for group in self.param_groups:
+                use_sinkgd = group.get("use_sinkgd", False)
+                # CPU scalar lr avoids recompiling the compiled steps on lr changes
+                lr = group["lr"]
+                if not isinstance(lr, Tensor):
+                    lr = torch.tensor(lr, dtype=torch.float32)
+
+                for p in group["params"]:
+                    if p.grad is None:
+                        continue
+                    grad = p.grad
+                    if grad.is_sparse:
+                        raise RuntimeError("Sparse gradient is not supported")
+
+                    if use_sinkgd and p.ndim >= 2:
+                        self._compiled_sinkgd(
+                            p.detach(),
+                            grad,
+                            lr * self.sinkgd_lr_scale,
+                            group["weight_decay"],
+                            self.sinkhorn_iters,
+                            self.sinkgd_eps,
+                            self.bf16_stochastic_round and p.dtype is torch.bfloat16,
+                        )
+                    else:
+                        self._adam_fallback(p, grad, group, lr)
+
+        return loss
+
+
+def _sinkgd_param_groups(opt_model, weight_decay):
+    """Split params: 2D/3D weight matrices -> SR-Sinkhorn; everything else -> AdamW.
+
+    Routing is by tensor rank, not module type, so fused MoE experts (transformers
+    v5 stores them as 3D ``[num_experts, in, out]`` parameters rather than
+    ``nn.Linear`` modules) are picked up and normalized per-expert. The AdamW
+    fallback group (embeddings, head, norms, biases) uses no weight decay.
+    """
+    sinkgd_params = []
+    adamw_params = []
+    for name, param in opt_model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if (
+            param.ndim < 2
+            or name.endswith("modules_to_save.default.weight")
+            or any(
+                embed in name
+                for embed in ["embed_tokens", "lm_head", "wte", "word_embeddings"]
+            )
+        ):
+            adamw_params.append(param)
+        else:
+            sinkgd_params.append(param)
+
+    param_groups = []
+    if sinkgd_params:
+        param_groups.append(
+            {"params": sinkgd_params, "use_sinkgd": True, "weight_decay": weight_decay}
+        )
+    if adamw_params:
+        param_groups.append(
+            {"params": adamw_params, "use_sinkgd": False, "weight_decay": 0.0}
+        )
+    return param_groups
+
+
+class SinkGDOptimizerFactory(BaseOptimizerFactory):
+    """Builds a :class:`SinkGD` optimizer, routing weight matrices to SR-Sinkhorn."""
+
+    def __call__(self, opt_model, training_args=None, **optimizer_kwargs) -> "SinkGD":
+        lr = optimizer_kwargs.pop("lr")
+        weight_decay = optimizer_kwargs.pop("weight_decay", 0.0)
+        betas = optimizer_kwargs.pop("betas", (0.9, 0.999))
+        eps = optimizer_kwargs.pop("eps", 1e-8)
+        sinkhorn_iters = int(optimizer_kwargs.pop("sinkhorn_iters", 5))
+        sinkgd_lr_scale = float(optimizer_kwargs.pop("sinkgd_lr_scale", 0.05))
+        optimizer_kwargs.pop(
+            "device_mesh", None
+        )  # ignored by the single-device variant
+
+        return SinkGD(
+            _sinkgd_param_groups(opt_model, weight_decay),
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            sinkhorn_iters=sinkhorn_iters,
+            sinkgd_lr_scale=sinkgd_lr_scale,
+            **optimizer_kwargs,
+        )
+
+
+def _matrix_shard_dim(p: Tensor):
+    """Return the negative matrix dim (-2 or -1) sharded across >1 ranks, else None.
+
+    Leading-dim sharding (e.g. the expert axis of a fused MoE weight) and replicated
+    tensors return None — their row/column norms are then fully shard-local.
+    """
+    if not isinstance(p, DTensor):
+        return None
+    matrix_dims = {p.ndim - 1, p.ndim - 2}
+    for i, placement in enumerate(p.placements):
+        if (
+            placement.is_shard()
+            and placement.dim in matrix_dims
+            and p.device_mesh.size(i) > 1
+        ):
+            return placement.dim - p.ndim
+    return None
+
+
+# Compiled local pieces of one SR-Sinkhorn iteration; the only thing kept eager and
+# uncompiled between them is the tiny norm-vector all-reduce (a collective would graph
+# break fullgraph anyway). This recovers the fusion the single-device path enjoys.
+def _row_scale_col_partial(x, sqrt_n, eps):
+    rn = torch.linalg.vector_norm(x, dim=-1, keepdim=True, dtype=torch.float32)
+    x = x * (sqrt_n / rn.clamp_min(eps)).to(x.dtype)
+    return x, (x.float() ** 2).sum(dim=-2, keepdim=True)
+
+
+def _apply_col_scale(x, csq, sqrt_m, eps):
+    return x * (sqrt_m / csq.sqrt().clamp_min(eps)).to(x.dtype)
+
+
+def _row_partial(x):
+    return (x.float() ** 2).sum(dim=-1, keepdim=True)
+
+
+def _row_scale_then_col_scale(x, rsq, sqrt_n, sqrt_m, eps):
+    # row norm is the reduced rsq (cols sharded); column norm is then shard-local.
+    x = x * (sqrt_n / rsq.sqrt().clamp_min(eps)).to(x.dtype)
+    cn = torch.linalg.vector_norm(x, dim=-2, keepdim=True, dtype=torch.float32)
+    return x * (sqrt_m / cn.clamp_min(eps)).to(x.dtype)
+
+
+def _sinkgd_apply_update(p, x, lr, weight_decay, stochastic_round):
+    p_f32 = p.float()
+    if weight_decay != 0.0:
+        p_f32 = p_f32 * (1 - lr * weight_decay)
+    p_f32 = p_f32 - lr * x.float()
+    if stochastic_round:
+        p.copy_(_fp32_to_bf16_sr(p_f32))
+    else:
+        p.copy_(p_f32)
+
+
+class DistSinkGD(SinkGD):
+    """Distributed SinkGD for FSDP2/TP-sharded weights.
+
+    The SR-Sinkhorn matrices are updated on their local shards, all-reducing only the
+    ``[N]`` (or ``[M]``) norm vector over the shard process group — never the matrix
+    itself, so optimizer communication is orders of magnitude lighter than gathering
+    the full parameter. Local work is torch.compiled (the all-reduce is the only eager
+    break); the 8-bit AdamW fallback is inherited unchanged.
+    """
+
+    def __init__(self, params, *, process_group=None, **kwargs) -> None:
+        super().__init__(params, **kwargs)
+        self._process_group = process_group
+        c = lambda f: torch.compile(f, fullgraph=True, dynamic=False)  # noqa: E731
+        self._c_row_scale_col_partial = c(_row_scale_col_partial)
+        self._c_apply_col_scale = c(_apply_col_scale)
+        self._c_row_partial = c(_row_partial)
+        self._c_row_scale_then_col_scale = c(_row_scale_then_col_scale)
+        self._c_apply_update = c(_sinkgd_apply_update)
+
+    def _dist_sinkgd_step(self, p, grad, group, lr):
+        shard_dim = _matrix_shard_dim(p)
+        sr = self.bf16_stochastic_round and p.dtype is torch.bfloat16
+        global_M, global_N = p.shape[-2], p.shape[-1]  # DTensor -> global dims
+        p_local = p.to_local() if isinstance(p, DTensor) else p
+        x = grad.to_local() if isinstance(grad, DTensor) else grad
+        lr = lr * self.sinkgd_lr_scale
+
+        if shard_dim is None:
+            # matrix dims fully local (replicated / expert-sharded) -> reuse the
+            # fully-compiled single-device path, no communication.
+            self._compiled_sinkgd(
+                p_local,
+                x,
+                lr,
+                group["weight_decay"],
+                self.sinkhorn_iters,
+                self.sinkgd_eps,
+                sr,
+            )
+            return
+
+        sqrt_n, sqrt_m = global_N**0.5, global_M**0.5
+        eps = self.sinkgd_eps
+        for _ in range(self.sinkhorn_iters):
+            if shard_dim == -2:  # rows sharded -> column norm needs a reduction
+                x, csq = self._c_row_scale_col_partial(x, sqrt_n, eps)
+                dist.all_reduce(csq, group=self._process_group)
+                x = self._c_apply_col_scale(x, csq, sqrt_m, eps)
+            else:  # columns sharded (-1) -> row norm needs a reduction (row first)
+                rsq = self._c_row_partial(x)
+                dist.all_reduce(rsq, group=self._process_group)
+                x = self._c_row_scale_then_col_scale(x, rsq, sqrt_n, sqrt_m, eps)
+        self._c_apply_update(p_local, x, lr, group["weight_decay"], sr)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        with torch._dynamo.utils.disable_cache_limit():
+            for group in self.param_groups:
+                use_sinkgd = group.get("use_sinkgd", False)
+                lr = group["lr"]
+                if not isinstance(lr, Tensor):
+                    lr = torch.tensor(lr, dtype=torch.float32)
+
+                for p in group["params"]:
+                    if p.grad is None:
+                        continue
+                    if p.grad.is_sparse:
+                        raise RuntimeError("Sparse gradient is not supported")
+                    if use_sinkgd and p.ndim >= 2:
+                        self._dist_sinkgd_step(p, p.grad, group, lr)
+                    else:
+                        self._adam_fallback(p, p.grad, group, lr)
+
+        return loss
+
+
+class DistSinkGDOptimizerFactory(BaseOptimizerFactory):
+    """Builds a :class:`DistSinkGD` for sharded training; pulls the dp_shard group."""
+
+    def __call__(
+        self, opt_model, training_args=None, **optimizer_kwargs
+    ) -> "DistSinkGD":
+        lr = optimizer_kwargs.pop("lr")
+        weight_decay = optimizer_kwargs.pop("weight_decay", 0.0)
+        betas = optimizer_kwargs.pop("betas", (0.9, 0.999))
+        eps = optimizer_kwargs.pop("eps", 1e-8)
+        sinkhorn_iters = int(optimizer_kwargs.pop("sinkhorn_iters", 5))
+        sinkgd_lr_scale = float(optimizer_kwargs.pop("sinkgd_lr_scale", 0.05))
+
+        device_mesh = optimizer_kwargs.pop("device_mesh", None)
+        process_group = None
+        if isinstance(device_mesh, DeviceMesh):
+            if "dp_shard" in (device_mesh.mesh_dim_names or ()):
+                process_group = device_mesh["dp_shard"].get_group()
+            elif device_mesh.ndim == 1:
+                process_group = device_mesh.get_group()
+
+        return DistSinkGD(
+            _sinkgd_param_groups(opt_model, weight_decay),
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            sinkhorn_iters=sinkhorn_iters,
+            sinkgd_lr_scale=sinkgd_lr_scale,
+            process_group=process_group,
+            **optimizer_kwargs,
+        )
diff --git a/src/axolotl/utils/quantization.py b/src/axolotl/utils/quantization.py
new file mode 100644
index 0000000000..9583c3d5f3
--- /dev/null
+++ b/src/axolotl/utils/quantization.py
@@ -0,0 +1,422 @@
+"""
+Utilities for quantization including QAT and PTQ using torchao.
+"""
+
+import functools
+
+import torch
+from packaging import version
+from torchao.core.config import AOBaseConfig
+from torchao.quantization import quantize_
+from torchao.quantization.granularity import PerGroup
+from torchao.quantization.qat import (
+    QATConfig,
+)
+from torchao.quantization.qat.fake_quantize_config import Int4WeightFakeQuantizeConfig
+from torchao.quantization.quant_api import (
+    Float8DynamicActivationFloat8WeightConfig,
+    Float8DynamicActivationInt4WeightConfig,
+    Int4WeightOnlyConfig,
+    Int8DynamicActivationIntxWeightConfig,
+)
+
+from axolotl.utils.schemas.enums import TorchAOQuantDType
+
+quantization_config_to_str = {
+    Int8DynamicActivationIntxWeightConfig: "int8int4",
+    Float8DynamicActivationFloat8WeightConfig: "fp8fp8",
+    Float8DynamicActivationInt4WeightConfig: "fp8int4",
+}
+
+if version.parse(torch.__version__) >= version.parse("2.8.0"):
+    try:
+        from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig
+
+        quantization_config_to_str[NVFP4WeightOnlyConfig] = "nvfp4"
+    except (ImportError, RuntimeError):
+        pass
+
+    try:
+        from torchao.quantization.quant_api import Int4WeightOnlyConfig
+
+        quantization_config_to_str[Int4WeightOnlyConfig] = "int4"
+    except (ImportError, RuntimeError):
+        pass
+
+    try:
+        from torchao.prototype.mx_formats import (
+            MXDynamicActivationMXWeightConfig as MXLinearConfig,
+        )
+
+        quantization_config_to_str[MXLinearConfig] = "mxfp4"
+    except (ImportError, RuntimeError):
+        pass
+
+
+def get_quantization_config(
+    weight_dtype: TorchAOQuantDType,
+    activation_dtype: TorchAOQuantDType | None = None,
+    group_size: int | None = None,
+) -> AOBaseConfig:
+    """
+    This function is used to build a post-training quantization config.
+
+    Args:
+        weight_dtype: The dtype to use for weight quantization.
+        activation_dtype: The dtype to use for activation quantization.
+        group_size: The group size to use for weight quantization.
+
+    Returns:
+        The post-training quantization config.
+
+    Raises:
+        ValueError: If the activation dtype is not specified and the weight dtype is not int8 or int4,
+            or if the group size is not specified for int8 or int4 weight only quantization.
+    """
+    if activation_dtype is None:
+        if weight_dtype == TorchAOQuantDType.int8:
+            raise ValueError("Int8WeightOnlyConfig is not supported by torchao QAT.")
+        if weight_dtype == TorchAOQuantDType.int4:
+            from torchao.quantization.quant_api import Int4WeightOnlyConfig
+
+            if group_size is not None:
+                return Int4WeightOnlyConfig(group_size=group_size, version=2)
+            else:
+                return Int4WeightOnlyConfig(version=2)
+    if (
+        activation_dtype == TorchAOQuantDType.int4
+        and weight_dtype == TorchAOQuantDType.int4
+    ):
+        raise ValueError(
+            "Int4DynamicActivationInt4WeightConfig is not supported by torchao QAT."
+        )
+    if (
+        activation_dtype == TorchAOQuantDType.int8
+        and weight_dtype == TorchAOQuantDType.int8
+    ):
+        raise ValueError(
+            "Int8DynamicActivationInt8WeightConfig is not supported by torchao QAT."
+        )
+    if (
+        activation_dtype == TorchAOQuantDType.int8
+        and weight_dtype == TorchAOQuantDType.int4
+    ):
+        kwargs = {"weight_dtype": torch.int4}
+        if group_size is not None:
+            kwargs["weight_granularity"] = PerGroup(group_size=group_size)
+        return Int8DynamicActivationIntxWeightConfig(**kwargs)
+    if (
+        activation_dtype == TorchAOQuantDType.float8_e4m3fn
+        and weight_dtype == TorchAOQuantDType.float8_e4m3fn
+    ):
+        return Float8DynamicActivationFloat8WeightConfig()
+    if (
+        activation_dtype == TorchAOQuantDType.float8_e4m3fn
+        and weight_dtype == TorchAOQuantDType.int4
+    ):
+        return Float8DynamicActivationInt4WeightConfig()
+    if weight_dtype == TorchAOQuantDType.nvfp4:
+        from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig
+
+        if group_size is not None and group_size != 16:
+            raise ValueError("NVFP4 quantization must use a group_size of 16")
+        return NVFP4WeightOnlyConfig()
+
+    if weight_dtype == TorchAOQuantDType.mxfp4:
+        from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
+
+        # MXFP4 uses block_size=32 by default (vs NVFP4's 16)
+        block_size = group_size if group_size is not None else 32
+        if block_size != 32:
+            raise ValueError(
+                "MXFP4 quantization must use a block_size (group_size) of 32"
+            )
+
+        return MXDynamicActivationMXWeightConfig(
+            activation_dtype=torch.float4_e2m1fn_x2,
+            weight_dtype=torch.float4_e2m1fn_x2,
+            block_size=block_size,
+        )
+
+    raise ValueError(
+        f"Invalid activation/weight dtype combination: {activation_dtype}/{weight_dtype}"
+    )
+
+
+def _attach_torchao_quantizer(
+    model, quantization_config, include_input_output_embeddings=False
+):
+    """Attach a TorchAoHfQuantizer to the model so save_pretrained uses
+    torchao's flatten_tensor_state_dict path, preserving quantized weights
+    in the safetensors file.
+
+    Note: This does NOT work for MX formats (MXTensor) because MXTensor
+    lacks the ``tensor_data_names`` attribute that flatten_tensor_state_dict
+    requires.  MX formats use ``safe_serialization=False`` instead.
+    """
+    from transformers import TorchAoConfig
+    from transformers.quantizers.quantizer_torchao import TorchAoHfQuantizer
+
+    ao_config = TorchAoConfig(
+        quant_type=quantization_config,
+        include_input_output_embeddings=include_input_output_embeddings,
+    )
+    model.config.quantization_config = ao_config
+    quantizer = TorchAoHfQuantizer(ao_config)
+    model.hf_quantizer = quantizer
+
+
+def patch_transformers_skip_quantized_init():
+    """Stop ``from_pretrained`` from re-initializing torchao-quantized weights.
+
+    transformers re-runs ``_init_weights`` on every module during loading; the
+    generic implementation does ``init.normal_(module.weight.float(), ...)``.
+    ``.float()`` on a torchao tensor subclass (e.g. ``MXTensor``) returns a new
+    tensor that both drops the ``_is_hf_initialized`` skip flag and does not
+    implement ``normal_``, so loading an MX checkpoint raises NotImplementedError.
+    Re-initializing an already-loaded quantized weight is never correct, so we
+    skip those modules entirely.
+    """
+    from torchao.utils import TorchAOBaseTensor
+    from transformers import PreTrainedModel
+
+    if getattr(PreTrainedModel._initialize_weights, "_axolotl_torchao_patched", False):
+        return
+
+    original = PreTrainedModel._initialize_weights
+
+    @functools.wraps(original)
+    def _initialize_weights(self, module, *args, **kwargs):
+        if any(
+            isinstance(param, TorchAOBaseTensor)
+            for param in module.parameters(recurse=False)
+        ):
+            module._is_hf_initialized = True
+            return None
+        return original(self, module, *args, **kwargs)
+
+    _initialize_weights._axolotl_torchao_patched = True
+    PreTrainedModel._initialize_weights = _initialize_weights
+
+
+def quantize_model(
+    model,
+    weight_dtype: TorchAOQuantDType,
+    group_size: int | None = None,
+    activation_dtype: TorchAOQuantDType | None = None,
+    quantize_embedding: bool | None = None,
+):
+    """
+    This function is used to quantize a model.
+
+    Args:
+        model: The model to quantize.
+        weight_dtype: The dtype to use for weight quantization.
+        group_size: The group size to use for weight quantization.
+        activation_dtype: The dtype to use for activation quantization.
+        quantize_embedding: Whether to quantize the model's embedding weights.
+
+    """
+    linear_ptq_config = get_quantization_config(
+        weight_dtype=weight_dtype,
+        activation_dtype=activation_dtype,
+        group_size=group_size,
+    )
+    quantize_(model, linear_ptq_config)
+    if quantize_embedding:
+        # activation fake quantization is not supported for embedding layers
+        embedding_quantize_config = get_quantization_config(
+            weight_dtype=weight_dtype,
+            activation_dtype=None,
+            group_size=group_size,
+        )
+        quantize_(
+            model,
+            embedding_quantize_config,
+            filter_fn=lambda m, _: isinstance(m, torch.nn.Embedding),
+        )
+
+    is_mx = False
+    try:
+        from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
+
+        is_mx = isinstance(linear_ptq_config, MXDynamicActivationMXWeightConfig)
+    except ImportError:
+        pass
+
+    if is_mx:
+        # MXTensor lacks tensor_data_names so flatten_tensor_state_dict (safetensors)
+        # cannot serialize it.  Mark the model so the caller can use
+        # safe_serialization=False (torch.save) which supports __tensor_flatten__.
+        model._is_mx_quantized = True
+        # MX checkpoints reload via plain from_pretrained (no HF quantizer), so guard
+        # transformers' weight re-init against the MXTensor weights it will encounter.
+        patch_transformers_skip_quantized_init()
+    else:
+        _attach_torchao_quantizer(
+            model,
+            linear_ptq_config,
+            include_input_output_embeddings=bool(quantize_embedding),
+        )
+
+
+def save_quantized_model(model, save_dir, **kwargs):
+    """Save a quantized model, handling MXTensor serialization.
+
+    MXTensor does not have a valid storage pointer, which causes
+    ``save_pretrained`` to crash (both in ``remove_tied_weights_from_state_dict``
+    via ``id_tensor_storage``, and in safetensors serialization).
+    Transformers >=5.5 removed the ``safe_serialization`` parameter entirely.
+
+    For MX-quantized models we save the config/generation_config via
+    ``save_pretrained`` machinery and the weights via ``torch.save``.
+    """
+    import os
+
+    is_mx = getattr(model, "_is_mx_quantized", False)
+    if not is_mx:
+        model.save_pretrained(save_dir, **kwargs)
+        return
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    # Save config, generation_config, etc.
+    model.config.save_pretrained(save_dir)
+    if hasattr(model, "generation_config") and model.generation_config is not None:
+        model.generation_config.save_pretrained(save_dir)
+
+    # Save weights via torch.save (supports __tensor_flatten__)
+    weights_path = os.path.join(save_dir, "pytorch_model.bin")
+    torch.save(model.state_dict(), weights_path)
+
+
+def _make_qat_config(
+    base_config: AOBaseConfig,
+    weight_dtype: TorchAOQuantDType,
+    activation_dtype: TorchAOQuantDType | None,
+    group_size: int | None,
+) -> QATConfig:
+    """Build a QATConfig, explicitly constructing fake quantize configs to ensure
+    group_size and other params are properly propagated (torchao's QATConfig(base_config)
+    does not always map these correctly)."""
+    from torchao.quantization.qat.fake_quantize_config import (
+        Float8FakeQuantizeConfig,
+        IntxFakeQuantizeConfig,
+    )
+
+    if weight_dtype == TorchAOQuantDType.mxfp4:
+        from torchao.prototype.qat import MXFakeQuantizeConfig
+
+        block_size = getattr(base_config, "block_size", 32)
+        mx_fq = MXFakeQuantizeConfig(
+            dtype=torch.float4_e2m1fn_x2, block_size=block_size
+        )
+        return QATConfig(activation_config=mx_fq, weight_config=mx_fq)
+
+    # Build explicit weight config
+    weight_fq_config: (
+        Int4WeightFakeQuantizeConfig
+        | IntxFakeQuantizeConfig
+        | Float8FakeQuantizeConfig
+        | None
+    ) = None
+    if weight_dtype == TorchAOQuantDType.int4:
+        gs = (
+            group_size
+            if group_size is not None
+            else getattr(base_config, "group_size", 128)
+        )
+        activation_dt = None
+        if activation_dtype == TorchAOQuantDType.int8:
+            activation_dt = torch.bfloat16
+        elif activation_dtype == TorchAOQuantDType.float8_e4m3fn:
+            activation_dt = torch.float8_e4m3fn
+        kwargs = {"group_size": gs}
+        if activation_dt is not None:
+            kwargs["activation_dtype"] = activation_dt
+        weight_fq_config = Int4WeightFakeQuantizeConfig(**kwargs)
+    elif weight_dtype == TorchAOQuantDType.float8_e4m3fn:
+        weight_fq_config = Float8FakeQuantizeConfig(dtype=torch.float8_e4m3fn)
+
+    # Build explicit activation config
+    activation_fq_config = None
+    if activation_dtype == TorchAOQuantDType.int8:
+        activation_fq_config = IntxFakeQuantizeConfig(
+            dtype=torch.int8, granularity="per_token", is_symmetric=False
+        )
+    elif activation_dtype == TorchAOQuantDType.float8_e4m3fn:
+        activation_fq_config = Float8FakeQuantizeConfig(dtype=torch.float8_e4m3fn)
+
+    if weight_fq_config is not None:
+        return QATConfig(
+            weight_config=weight_fq_config,
+            activation_config=activation_fq_config,
+        )
+
+    # Fallback to base_config for unhandled combos
+    return QATConfig(base_config)
+
+
+def prepare_model_for_qat(
+    model,
+    weight_dtype: TorchAOQuantDType,
+    group_size: int | None = None,
+    activation_dtype: TorchAOQuantDType | None = None,
+    quantize_embedding: bool = False,
+):
+    """
+    This function is used to prepare a model for QAT by swapping the model's linear
+    layers with fake quantized linear layers, and optionally the embedding weights with
+    fake quantized embedding weights.
+
+    Args:
+        model: The model to quantize.
+        weight_dtype: The dtype to use for weight quantization.
+        group_size: The group size to use for weight quantization.
+        activation_dtype: The dtype to use for activation quantization.
+        quantize_embedding: Whether to quantize the model's embedding weights.
+
+    Raises:
+        ValueError: If the activation/weight dtype combination is invalid.
+    """
+    base_config = get_quantization_config(
+        weight_dtype=weight_dtype,
+        activation_dtype=activation_dtype,
+        group_size=group_size,
+    )
+    qat_config = _make_qat_config(
+        base_config, weight_dtype, activation_dtype, group_size
+    )
+    quantize_(model, qat_config)
+    if quantize_embedding:
+        # activation fake quantization is not supported for embedding layers
+        embedding_base_config = get_quantization_config(
+            weight_dtype=weight_dtype,
+            activation_dtype=None,
+            group_size=group_size,
+        )
+        embedding_qat_config = _make_qat_config(
+            embedding_base_config, weight_dtype, None, group_size
+        )
+        quantize_(
+            model,
+            embedding_qat_config,
+            filter_fn=lambda m, _: isinstance(m, torch.nn.Embedding),
+        )
+
+
+def convert_qat_model(
+    model,
+    quantize_embedding: bool = False,
+):
+    """
+    This function converts a QAT model which has fake quantized layers back to the original model.
+    """
+    config = QATConfig(step="convert")
+    quantize_(model, config)
+    if quantize_embedding:
+        quantize_(
+            model,
+            config,
+            filter_fn=lambda m, _: isinstance(m, torch.nn.Embedding),
+        )
diff --git a/src/axolotl/utils/samplers/__init__.py b/src/axolotl/utils/samplers/__init__.py
index 96e00a5d26..eb6dc8ab9f 100644
--- a/src/axolotl/utils/samplers/__init__.py
+++ b/src/axolotl/utils/samplers/__init__.py
@@ -1,5 +1,6 @@
 """
 axolotl samplers module
 """
+
 from .multipack import MultipackBatchSampler  # noqa: F401
 from .utils import get_dataset_lengths  # noqa: F401
diff --git a/src/axolotl/utils/samplers/multipack.py b/src/axolotl/utils/samplers/multipack.py
index cf47d9639b..436a49c798 100644
--- a/src/axolotl/utils/samplers/multipack.py
+++ b/src/axolotl/utils/samplers/multipack.py
@@ -1,35 +1,56 @@
-# pylint: skip-file
 """
-Multipack Batch Sampler
+Multipack Batch Sampler - An efficient batch sampler for packing variable-length sequences
+into fixed-capacity batches to optimize memory usage and training throughput.
 """
-import logging
+
+import gc
 import math
 import os
-from typing import Any, Iterable, List, Union
+import time
+from concurrent.futures import ProcessPoolExecutor
+from multiprocessing import cpu_count, get_context
+from typing import Iterable, Iterator, Union
 
 import numba
 import numpy as np
-from torch.utils.data import BatchSampler, Sampler
+from torch.utils.data import BatchSampler, Sampler, SequentialSampler
+
+from axolotl.utils.distributed import reduce_and_broadcast
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl.utils.samplers.multipack")
+LOG = get_logger(__name__)
 
 
 @numba.njit
-def ffd_check(a: np.ndarray, c: int, n: int):
-    # First-fit-decreasing bin packing
-    # Check if a[] could fit in n bins with capacity c
-    # https://en.wikipedia.org/wiki/First-fit-decreasing_bin_packing
-
-    a = np.sort(a)[::-1]
-    bins = np.full((n,), c, dtype=a.dtype)
-    for size in a:
+def ffd_check(sequence_lengths: np.ndarray, bin_capacity: int, num_bins: int) -> bool:
+    """First-fit-decreasing bin packing algorithm check.
+
+    Checks if sequences with the given lengths could fit in the specified number of
+    bins.
+
+    Args:
+        sequence_lengths: Array of sequence lengths.
+        bin_capacity: Maximum capacity of each bin.
+        num_bins: Number of bins available.
+
+    Returns:
+        `True` if all sequences can be packed, `False` otherwise.
+    """
+    # Sort sequence lengths in descending order for optimal packing
+    sequence_lengths = np.sort(sequence_lengths)[::-1]
+    # Initialize all bins with full capacity
+    bins = np.full((num_bins,), bin_capacity, dtype=sequence_lengths.dtype)
+
+    # Try to place each sequence in the first bin it fits
+    for size in sequence_lengths:
         not_found = True
-        for idx in range(n):
+        for idx in range(num_bins):
             if bins[idx] >= size:
                 bins[idx] -= size
                 not_found = False
                 break
 
+        # If no bin could fit this sequence, packing failed
         if not_found:
             return False
 
@@ -37,166 +58,416 @@ def ffd_check(a: np.ndarray, c: int, n: int):
 
 
 @numba.njit
-def ffd_with_result(a: np.ndarray, c: int, start_index: int):
-    # First-fit-decreasing bin packing (with result return)
+def pack_group(
+    sequence_lengths: np.ndarray,
+    group_offset: int,
+    bin_capacity: int,
+    max_bins: int,
+    bin_size: int,
+    safe_mode: bool = True,
+) -> list[list[int]]:
+    """Pack a group of sequences into bins using First-Fit Decreasing algorithm.
+
+    Args:
+        sequence_lengths: Array of sequence lengths.
+        group_offset: Offset to apply to indices when returning results.
+        bin_capacity: Maximum capacity of each bin.
+        max_bins: Maximum number of bins to use.
+        bin_size: Maximum number of sequences per bin.
+        safe_mode: If True, use a more conservative packing approach.
+
+    Returns:
+        List of bins, where each bin contains indices of sequences assigned to it.
+    """
+    bins_remaining_space: list = []  # Tracks remaining capacity in each bin
+    bins_assigned_sequences: list = []  # Tracks sequence indices assigned to each bin
+
+    for seq_id, size in enumerate(sequence_lengths):
+        global_idx = seq_id + group_offset
+
+        # Try to place sequence in existing bins
+        add_new_bin = True
+        for bin_idx, _ in enumerate(bins_remaining_space):
+            if (
+                bins_remaining_space[bin_idx] >= size
+                and len(bins_assigned_sequences[bin_idx]) < bin_size
+            ):
+                bins_remaining_space[bin_idx] -= size
+                bins_assigned_sequences[bin_idx].append(global_idx)
+                add_new_bin = False
+                break
 
-    indices = np.argsort(a)[::-1]
-    a = a[indices]
+        # Create a new bin if needed and if we haven't reached the limit
+        if add_new_bin:
+            if len(bins_remaining_space) >= max_bins and safe_mode:
+                # In safe mode, skip items that would exceed max_bins
+                continue
+            bins_remaining_space.append(bin_capacity - size)
+            bins_assigned_sequences.append([global_idx])
 
-    bins: List[Any] = []
-    bins_result: List[Any] = []
-    for a_id, size in enumerate(a):
-        add_new = True
-        for idx in range(len(bins)):
-            if bins[idx] >= size:
-                bins[idx] -= size
-                bins_result[idx].append(indices[a_id] + start_index)
-                add_new = False
+            # Safety check to avoid infinite bins
+            if len(bins_remaining_space) > len(sequence_lengths):
                 break
 
-        if add_new:
-            bins.append(c - size)
-            bins_result.append([indices[a_id] + start_index])
-
-    return bins_result
+    return bins_assigned_sequences
+
+
+def _process_group(
+    args: tuple[np.ndarray, int, int, int, int, bool],
+) -> list[list[int]]:
+    """Standalone function for multiprocessing."""
+    group_lengths, start_idx, bin_capacity, max_bins, bin_size, safe_mode = args
+    return pack_group(
+        group_lengths, start_idx, bin_capacity, max_bins, bin_size, safe_mode
+    )
+
+
+def pack_parallel(
+    sequence_lengths: np.ndarray,
+    bin_capacity: int,
+    group_size: int,
+    bin_size: int,
+    num_processes: int | None = None,
+    safe_mode: bool = True,
+    mp_start_method: str | None = "fork",
+) -> list[list[int]]:
+    """Pack sequences into bins using parallel processing.
+
+    Args:
+        sequence_lengths: Array of sequence lengths.
+        bin_capacity: Maximum capacity of each bin as total number of tokens.
+        group_size: Number of sequences to process in each group.
+        bin_size: Maximum number of bins to use.
+        num_processes: Number of parallel processes to use.
+        safe_mode: If True, use a more conservative packing approach.
+        mp_start_method: Multiprocessing start method ('fork', 'spawn', 'forkserver').
+                         'spawn' is often safer with Numba/PyTorch.
+                         Set to None to use system default.
+    Returns:
+        List of bins, where each bin contains indices of sequences assigned to it.
+    """
+    num_items = len(sequence_lengths)
+    if num_processes is None:
+        num_processes = max(1, min(num_items // group_size, cpu_count(), 16))
+
+    # Create tasks for parallel processing
+    tasks = []
+    for i in range(0, num_items, group_size):
+        group_lengths = sequence_lengths[i : i + group_size]
+        max_bins = len(group_lengths)  # Allow as many bins as items in the group
+        tasks.append((group_lengths, i, bin_capacity, max_bins, bin_size, safe_mode))
+
+    # Process groups in parallel
+    all_bins = []
+
+    mp_ctx = None
+    if mp_start_method:
+        try:
+            mp_ctx = get_context(mp_start_method)
+        except ValueError:
+            LOG.warning(
+                f"Failed to get multiprocessing context '{mp_start_method}'. "
+                f"Falling back to default. Available: {get_context().get_all_start_methods()}"
+            )
+            mp_ctx = (
+                None  # Fallback to default context if specified one is not available
+            )
+
+    if num_processes == 1:
+        LOG.debug("Using single process for pack_parallel, running sequentially.")
+        for task_args in tasks:
+            group_bins = _process_group(task_args)
+            all_bins.extend(group_bins)
+    else:
+        # Use ProcessPoolExecutor only if num_processes > 1
+        # Pass mp_context if available
+        with ProcessPoolExecutor(
+            max_workers=num_processes, mp_context=mp_ctx
+        ) as executor:
+            for group_bins in executor.map(_process_group, tasks):
+                all_bins.extend(group_bins)
+
+    return all_bins
 
 
 @numba.njit
-def allocate(
-    lengths: np.ndarray, lengths_cumsum: np.ndarray, rank: int, c: int, n: int
-):
-    # Dynamic batch allocator, similar to Multifit
-    # https://en.wikipedia.org/wiki/Multifit_algorithm
-    # ~99.5% efficiency on OpenChat training set (12 * 2048 ctx len)
-
-    s = 0
-    start_index = 0
+def allocate_sequentially(
+    sequence_lengths: np.ndarray, rank: int, bin_capacity: int, num_ranks: int
+) -> tuple[list[list[int]], int, int]:
+    """Sequential allocator that preserves example order.
+
+    Args:
+        sequence_lengths: The lengths of all examples.
+        rank: The current rank (for distributed training).
+        bin_capacity: The capacity of each bin (maximum sequence length).
+        num_ranks: Number of ranks (processes / GPUs).
+
+    Returns:
+        rank_batches: List of batches for the current rank.
+        total_tokens_used: Number of actual example tokens.
+        total_token_slots: Maximum theoretical number of example tokens (number of bins
+            * bin capacity).
+    """
     result = []
+    total_used = 0
+
+    # First, do sequential packing into bins
+    all_bins = []
+    current_bin = [0 for i in range(0)]  # numba hint
+    remaining_capacity = bin_capacity
+
+    for idx, size in enumerate(sequence_lengths):
+        if size <= remaining_capacity:
+            # Example fits in current bin
+            current_bin.append(idx)
+            remaining_capacity -= size
+            total_used += size
+        else:
+            # Example doesn't fit, start a new bin
+            if current_bin:  # Add non-empty bin to all_bins
+                all_bins.append(current_bin)
+            current_bin = [idx]
+            remaining_capacity = bin_capacity - size
+            total_used += size
+
+    # Add the last bin if not empty
+    if current_bin:
+        all_bins.append(current_bin)
+
+    # Assign bins to ranks - each rank gets every n-th bin
+    for bin_idx in range(rank, len(all_bins), num_ranks):
+        result.append(all_bins[bin_idx])
+
+    return result, total_used, len(all_bins) * bin_capacity
 
-    while True:
-        # binary search [l, r)
-        left = 1
-        right = 1 + np.searchsorted(lengths_cumsum[start_index:], s + c * n, "right")
-
-        while right - left > 1:
-            mid = (left + right) // 2
-            if ffd_check(lengths[start_index : start_index + mid], c, n):
-                left = mid
-            else:
-                right = mid
-
-        # use length l
-        batch = ffd_with_result(
-            lengths[start_index : start_index + left], c, start_index
-        )
-        assert len(batch) <= n
-        if len(batch) < n:
-            break
-
-        start_index += left
-        s = lengths_cumsum[start_index - 1]
 
-        # add local rank
-        result.append(batch[rank])
-
-    return result, s, len(result) * c * n
+class MultipackBatchSampler(BatchSampler):
+    """Batch sampler class for efficient packing of variable-length sequences
 
+    This sampler packs sequences into fixed-capacity bins (batches) to maximize
+    GPU memory utilization and training throughput by reducing padding.
 
-class MultipackBatchSampler(BatchSampler):
-    """
-    Batch Sampler class for multipack
+    It supports both parallel packing (using FFD algorithm) and
+    sequential packing (preserving original sequence order).
     """
 
+    _batches: list[list[list[int]]] | None = None
+    _len_across_ranks: int | None = None
+
     def __init__(
         self,
         sampler: Union[Sampler[int], Iterable[int]],
-        batch_size: int,
-        drop_last: bool,
-        batch_max_len: int,
-        lengths: np.ndarray,
-        packing_efficiency_estimate: float = 1.0,
+        batch_size: int,  # Number of bins per batch
+        batch_max_len: int,  # Maximum sequence length (bin capacity)
+        lengths: np.ndarray,  # Sequence lengths
+        bin_size: int,  # The max number of samples that can be packed in a single bin
+        packing_efficiency_estimate: float = 1.0,  # Initial efficiency estimate
+        drop_last: bool = True,  # Whether to drop final batches (might be incomplete)
+        num_count_samples: int = 4,  # Number of times to estimate batch count
+        sequential: bool = False,  # Whether to use sequential packing
+        group_size: int = 100_000,  # Size of groups for parallel packing
+        num_processes: int | None = None,  # Number of processes for parallel packing
+        safe_mode: bool = True,  # Conservative packing to prevent training instability
+        mp_start_method: str = "fork",
+        **kwargs,
     ):
         super().__init__(sampler, batch_size, drop_last)
         self.batch_size = batch_size
         self.batch_max_len = batch_max_len
-        self.lengths: np.ndarray = lengths
+        self.lengths = np.array(lengths, dtype=np.int32)
         self.packing_efficiency_estimate = packing_efficiency_estimate or 1.0
+        self.sequential = sequential
+        self.group_size = group_size
+        self.bin_size = bin_size
+        self.num_processes = num_processes
+        self.safe_mode = safe_mode
+        self.mp_start_method = mp_start_method
 
         assert isinstance(self.lengths, np.ndarray)
 
         self.epoch = 0
 
-        # statistics
-        self.eff_total_used = 0
-        self.eff_total_slots = 0
+        # Efficiency statistics tracking
+        self.total_tokens_used = 0
+        self.total_token_slots = 0
+
+        # The number of times to calculate batches to determine minimum packed dataset length
+        world_size = int(os.environ.get("WORLD_SIZE", "1"))
+        self.num_count_samples = (
+            1 if world_size >= num_count_samples else num_count_samples
+        )
+
+        if self.sequential and not isinstance(sampler, SequentialSampler):
+            LOG.warning(
+                "using sequential sample packing with non-sequential sampler, did you want to also enable curriculum_sampling?"
+            )
 
     def set_epoch(self, epoch: int):
+        """Set the epoch number, used for reproducible shuffling across epochs"""
         self.epoch = epoch
+        self._batches = None  # Invalidate batch cache
 
-    def generate_batches(self, set_stats=False):
+    def generate_batches(self, set_stats: bool = False) -> list[list[list[int]]]:
+        """Generate packed batches for training.
+
+        Args:
+            set_stats: Whether to update efficiency statistics.
+
+        Returns:
+            List of batches, where each batch contains multiple bins, and each bin
+                contains multiple sequence indices.
+        """
+        if self._batches is not None:
+            return self._batches
+
+        # Get indices from the sampler
         indices = [idx for idx in self.sampler]
 
+        # Get lengths of the selected sequences
         lengths = self.lengths[indices]
-        lengths_cumsum = np.cumsum(lengths)
-
-        batches, total_used, total_slots = allocate(
-            lengths=lengths,
-            lengths_cumsum=lengths_cumsum,
-            rank=0,
-            c=self.batch_max_len,
-            n=1,
-        )
 
-        batches = [
-            [
-                [indices[b_idx] for b_idx in batch]
-                for batch in batches[i : i + self.batch_size]
+        # Pack sequences into bins using either sequential or parallel packing
+        if self.sequential:
+            bins, total_used, total_slots = allocate_sequentially(
+                lengths,
+                rank=0,
+                bin_capacity=self.batch_max_len,
+                num_ranks=1,
+            )
+            # Map bin indices back to original indices
+            bins = [[indices[b_idx] for b_idx in bin_indices] for bin_indices in bins]
+        else:
+            # Use parallel packing
+            num_processes = self.num_processes or 1
+            all_bins = pack_parallel(
+                lengths,
+                bin_capacity=self.batch_max_len,
+                group_size=self.group_size,
+                bin_size=self.bin_size or self.batch_max_len,
+                num_processes=min(4, num_processes) if num_processes else 4,
+                safe_mode=self.safe_mode,
+                mp_start_method=self.mp_start_method,
+            )
+
+            # Map bin indices back to original indices
+            bins = [
+                [indices[b_idx] for b_idx in bin_indices] for bin_indices in all_bins
             ]
-            for i in range(0, len(batches), self.batch_size)
+
+            # Calculate efficiency statistics
+            total_used = lengths.sum()
+            total_slots = len(all_bins) * self.batch_max_len
+            del all_bins
+
+        # Group bins into batches (each batch contains batch_size bins)
+        batches = [
+            bins[i : i + self.batch_size] for i in range(0, len(bins), self.batch_size)
         ]
 
-        # statistics
+        # Drop last batch if requested and it's incomplete
+        if self.drop_last and len(batches[-1]) < self.batch_size:
+            batches = batches[:-1]
+            # Adjust total_slots if we dropped a batch
+            if not self.sequential:
+                total_slots -= (self.batch_size - len(batches[-1])) * self.batch_max_len
+
+        # Update statistics if requested
         if set_stats:
-            self.eff_total_used += total_used
-            self.eff_total_slots += total_slots
+            self.total_tokens_used += total_used
+            self.total_token_slots += total_slots
 
+        self._batches = batches
+        gc.collect()
         return batches
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[list[list[int]]]:
+        """Return an iterator over batches.
+
+        The batches are truncated to match the minimum number of batches across all
+        ranks to ensure distributed training balance.
+        """
         batches = self.generate_batches(set_stats=True)
+        if self._len_across_ranks:
+            # Truncate batches to ensure all ranks have the same number of batches
+            batches = batches[: self._len_across_ranks]
         return iter(batches)
 
-    def num_batches(self):
-        batches = self.generate_batches(set_stats=True)
-        return len(batches)
-
-    def efficiency(self):
-        return self.eff_total_used / self.eff_total_slots
-
-    def __len__(self):
-        self.num_batches()
-        return self._len_est()
-
-    def _len_est(self):
-        world_size = int(os.getenv("WORLD_SIZE", "1"))
-        lengths_sum = np.sum(self.lengths)
-        lengths_sum_per_device = lengths_sum // world_size
-        LOG.info(
-            f"packing_efficiency_estimate: {self.packing_efficiency_estimate} "
-            f"total_num_tokens per device: {lengths_sum_per_device}"
+    def efficiency(self) -> float:
+        """Calculate the packing efficiency (ratio of tokens used to total token slots).
+        Higher is better - 1.0 would mean perfect packing with no wasted space.
+        """
+        if self.total_token_slots == 0:
+            self.generate_batches(set_stats=True)
+        if self.total_token_slots == 0:
+            return 0.0
+        # Return a Python float instead of potentially a numpy float
+        return float(self.total_tokens_used / self.total_token_slots)
+
+    def gather_efficiency(self) -> float:
+        """Gather and synchronize packing efficiency estimates across all distributed
+        ranks.
+
+        Returns:
+            A conservative efficiency estimate based on the measurements.
+        """
+
+        def calc_sample_packing_eff_est(estimates: list[float]):
+            LOG.debug(f"sample_packing_eff_est across ranks: {repr(estimates)}")
+            # Use 99.7% of max observed efficiency as a safe estimate
+            max_eff = max(float(eff) for eff in estimates)
+            return math.floor(0.997 * max_eff)
+
+        # Gather efficiency from all ranks and apply the calculation function
+        sample_packing_actual_eff_all = reduce_and_broadcast(
+            lambda: float(self.efficiency()),
+            calc_sample_packing_eff_est,
         )
 
-        # shave off 1% + 1 for dealing with variance in packing from random sampler to sampler
-        return max(
-            0,
-            (
-                world_size
-                * math.floor(
-                    0.99
-                    * lengths_sum_per_device
-                    / self.packing_efficiency_estimate
-                    // (self.batch_max_len * self.batch_size)
-                )
-                - 1
-            ),
+        # Quantize to 0.5% intervals for stability
+        sample_packing_eff_est = (
+            math.ceil(sample_packing_actual_eff_all * 200.0) / 200.0
         )
+        return sample_packing_eff_est
+
+    def gather_len_batches(self, num: int) -> int:
+        """Gather and synchronize batch counts across all distributed ranks. Returns
+        the minimum number of batches available on any rank.
+        """
+
+        def calc_min_len(estimates: list[int]) -> int:
+            LOG.info(f"gather_len_batches: {repr(estimates)}")
+            return math.floor(min(estimates))
+
+        # Find minimum batch count across ranks to ensure balance
+        min_len_batches = reduce_and_broadcast(lambda: num, calc_min_len)
+        return min_len_batches
+
+    def __len__(self) -> int:
+        """Return the total number of batches that will be yielded by this sampler.
+
+        This is calculated as the minimum number of batches available on any rank to
+        ensure balanced distributed training.
+        """
+        if self._batches is None:
+            self._batches = self.generate_batches(set_stats=True)
+
+        if self._len_across_ranks is None:
+            # Sample multiple times to get stable estimate
+            _sampled_lens = []
+            for _ in range(self.num_count_samples):
+                self._batches = None  # Reset cached batches
+                # log timer for generating batches
+                start_time = time.time()
+                _sampled_lens.append(len(self.generate_batches(set_stats=False)))
+                LOG.debug(f"generate_batches time: {time.time() - start_time}")
+            len_batches = min(_sampled_lens)
+
+            # Gather minimum across all ranks
+            if self._len_across_ranks is None:
+                self._len_across_ranks = self.gather_len_batches(len_batches)
+            else:
+                self._len_across_ranks = min(
+                    self._len_across_ranks, self.gather_len_batches(len_batches)
+                )
+
+        return self._len_across_ranks
diff --git a/src/axolotl/utils/samplers/utils.py b/src/axolotl/utils/samplers/utils.py
index e4af4e5f35..a93e847488 100755
--- a/src/axolotl/utils/samplers/utils.py
+++ b/src/axolotl/utils/samplers/utils.py
@@ -1,17 +1,21 @@
 """
 helper util to calculate dataset lengths
 """
+
 import numpy as np
 
 
-def get_dataset_lengths(dataset):
-    if "length" in dataset.data.column_names:
-        lengths = np.array(dataset.data.column("length"))
-    elif "position_ids" in dataset.data.column_names:
-        position_ids = dataset.data.column("position_ids")
+def get_dataset_lengths(dataset, from_arrow=False):
+    if "length" in dataset.column_names:
+        lengths = np.array(dataset["length"])
+    elif "position_ids" in dataset.column_names:
+        position_ids = dataset["position_ids"]
         lengths = np.array([x[-1] + 1 for x in position_ids])
     else:
-        input_ids = dataset.data.column("input_ids")
-        lengths = np.vectorize(len)(np.array(input_ids, dtype=object))
-        return lengths
+        if from_arrow:
+            input_ids = dataset.data.column("input_ids")
+            lengths = np.vectorize(len)(np.array(input_ids, dtype=object))
+        else:
+            input_ids = dataset["input_ids"]
+            lengths = np.array([len(seq) for seq in input_ids])
     return lengths
diff --git a/src/axolotl/utils/schedulers.py b/src/axolotl/utils/schedulers.py
index 94387e5ab8..3090a3acdb 100644
--- a/src/axolotl/utils/schedulers.py
+++ b/src/axolotl/utils/schedulers.py
@@ -1,11 +1,90 @@
 """Module for custom LRScheduler class"""
+
 import math
 from functools import partial
+from typing import Any, Sequence
 
+from torch import Tensor
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LambdaLR, LRScheduler
 
 
+class RexLR(LRScheduler):
+    """
+    Reflected Exponential (REX) learning rate scheduler.
+
+    - Original implementation: https://github.com/IvanVassi/REX_LR
+    - Original license: Apache 2.0
+    - Based on: https://arxiv.org/abs/2107.04197
+
+    Args:
+        optimizer (torch.optim.Optimizer): The optimizer to schedule the learning rate for.
+        max_lr (float): The maximum learning rate.
+        min_lr (float): The minimum learning rate.
+        total_steps (int): The total number of training steps.
+        num_warmup_steps (int): The number of warmup steps.
+        last_step (int): The index of last step.
+    """
+
+    def __init__(
+        self, optimizer, max_lr, min_lr, total_steps=0, num_warmup_steps=0, last_step=0
+    ):
+        if min_lr > max_lr:
+            raise ValueError(
+                f'Value of "min_lr" should be less than value of "max_lr". Got min_lr={min_lr} and max_lr={max_lr}'
+            )
+        if num_warmup_steps > total_steps:
+            raise ValueError(
+                f"num_warmup_steps ({num_warmup_steps}) must be less than or equal to total_steps ({total_steps})."
+            )
+
+        self.min_lr = min_lr
+        self.max_lr = max_lr
+        self.total_steps = total_steps
+        self.num_warmup_steps = num_warmup_steps
+        self.last_step = max(last_step - 1, 0)
+
+        # Ensure each parameter group has an "initial_lr" key to avoid issues when resuming.
+        for group in optimizer.param_groups:
+            initial_lr = group["lr"]
+            if isinstance(initial_lr, Tensor):
+                initial_lr = initial_lr.clone()
+            group.setdefault("initial_lr", initial_lr)
+        # Pass self.last_step as last_epoch to the parent.
+        super().__init__(optimizer, last_epoch=self.last_step)
+
+    @property
+    def last_step(self):
+        return self.last_epoch
+
+    @last_step.setter
+    def last_step(self, value):
+        self.last_epoch = value
+
+    def get_lr(self):
+        # Warmup phase: if defined, increase lr linearly from 0 to max_lr.
+        if 1 <= self.last_step <= self.num_warmup_steps:
+            return [
+                base_lr * self.last_step / self.num_warmup_steps
+                for base_lr in self.base_lrs
+            ]
+
+        # Post-warmup phase: adjust step relative to the end of warmup.
+        step_after = self.last_step - self.num_warmup_steps
+        remaining_steps = self.total_steps - self.num_warmup_steps
+
+        # Avoid LR spiking
+        if step_after >= remaining_steps or step_after == -1 or remaining_steps <= 0:
+            return [self.min_lr for _ in self.base_lrs]
+
+        mod_iter = step_after % remaining_steps
+        z = (remaining_steps - mod_iter) / remaining_steps
+        rex_factor = self.min_lr / self.max_lr + (1.0 - self.min_lr / self.max_lr) * (
+            z / (0.1 + 0.9 * z)
+        )
+        return [base_lr * rex_factor for base_lr in self.base_lrs]
+
+
 class InterpolatingLogScheduler(LRScheduler):
     """
     A scheduler that interpolates learning rates in a logarithmic fashion
@@ -28,9 +107,7 @@ def __init__(self, optimizer, num_steps, min_lr, max_lr, last_epoch=-1):
         self.num_steps = num_steps
         self.min_lr = min_lr
         self.max_lr = max_lr
-        self.q = (max_lr / min_lr) ** (  # pylint: disable=invalid-name
-            1 / (num_steps - 1)
-        )
+        self.q = (max_lr / min_lr) ** (1 / (num_steps - 1))
         super().__init__(optimizer, last_epoch)
 
     def get_lr(self):
@@ -217,3 +294,65 @@ def get_cosine_schedule_with_warmup_decay_constant(
         num_cycles=num_cycles,
     )
     return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+class JaggedLRRestartScheduler(LRScheduler):
+    """Wraps another scheduler to apply per-lora-restart learning rate warmups."""
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        inner_schedule: LRScheduler,
+        jagged_restart_steps: int,
+        jagged_restart_warmup_steps: int,
+        jagged_restart_anneal_steps: int = 1,
+        min_lr_scale: float = 0.001,
+    ) -> None:
+        self.inner_schedule = inner_schedule
+        self.restarts_steps = jagged_restart_steps
+        self.warmup_steps = jagged_restart_warmup_steps
+        self.anneal_steps = jagged_restart_anneal_steps
+        self.min_lr_scale = min_lr_scale
+        super().__init__(optimizer, inner_schedule.last_epoch)
+
+    def get_lr(self) -> float | Sequence[float]:
+        self.inner_schedule.last_epoch = self.last_epoch
+
+        original = self.inner_schedule.get_lr()
+        step = self.last_epoch
+
+        if step < self.restarts_steps - self.anneal_steps:
+            scale = 1
+        else:
+            per_restart_progress = step % self.restarts_steps
+            if per_restart_progress < self.warmup_steps:
+                cycle_t = min(1.0, (per_restart_progress) / self.warmup_steps)
+            elif per_restart_progress > (self.restarts_steps - self.anneal_steps):
+                cycle_t = min(
+                    1.0,
+                    (self.restarts_steps - per_restart_progress) / self.anneal_steps,
+                )
+            else:
+                cycle_t = 1
+            scale = cycle_t * (1 - self.min_lr_scale) + self.min_lr_scale
+
+        if isinstance(original, Sequence):
+            return [lr * scale for lr in original]
+
+        return original * scale
+
+    def state_dict(self) -> dict[str, Any]:
+        """Return serializable state, saving inner_schedule as its own state_dict."""
+        state = {
+            key: value
+            for key, value in self.__dict__.items()
+            if key not in ("optimizer", "inner_schedule")
+        }
+        state["inner_schedule_state"] = self.inner_schedule.state_dict()
+        return state
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """Restore state, including inner_schedule."""
+        inner_state = state_dict.pop("inner_schedule_state")
+        self.__dict__.update(state_dict)
+        self.inner_schedule.load_state_dict(inner_state)
diff --git a/src/axolotl/utils/schemas/__init__.py b/src/axolotl/utils/schemas/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/axolotl/utils/schemas/config.py b/src/axolotl/utils/schemas/config.py
new file mode 100644
index 0000000000..bbf6aec0bf
--- /dev/null
+++ b/src/axolotl/utils/schemas/config.py
@@ -0,0 +1,2132 @@
+"""Module with Pydantic models for configuration."""
+
+import re
+from typing import Annotated, Any, Literal
+
+from accelerate.utils import is_fp8_available
+from annotated_types import MinLen
+from packaging import version
+from pydantic import (
+    BaseModel,
+    Field,
+    StringConstraints,
+    computed_field,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
+
+from axolotl.utils.datasets import get_default_process_count
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.datasets import (
+    DatasetConfig,
+    DPODataset,
+    KTODataset,
+    PretrainingDataset,
+    SFTDataset,
+    StepwiseSupervisedDataset,
+    SyntheticDataset,
+)
+from axolotl.utils.schemas.deprecated import DeprecatedParameters, RemappedParameters
+from axolotl.utils.schemas.dynamic_checkpoint import DynamicCheckpointConfig
+from axolotl.utils.schemas.enums import (
+    ATTN_IMPLS_SUPPORTING_PACKING,
+    ATTN_IMPLS_USING_FLASH_LIB,
+    ATTN_IMPLS_WITHOUT_DTYPE_CAST,
+    CANONICAL_ATTN_IMPLS,
+    LEGACY_ATTN_FLAG_TO_IMPL,
+    SHORT_FORM_ALIAS_TO_CANONICAL,
+    ChatTemplate,
+    RingAttnFunc,
+    RLType,
+)
+from axolotl.utils.schemas.fsdp import FSDPConfig
+from axolotl.utils.schemas.integrations import (
+    CometConfig,
+    GradioConfig,
+    LISAConfig,
+    MLFlowConfig,
+    OpenTelemetryConfig,
+    RayConfig,
+    TrackioConfig,
+    WandbConfig,
+)
+from axolotl.utils.schemas.internal import EnvCapabilities, GPUCapabilities
+from axolotl.utils.schemas.model import (
+    ModelInputConfig,
+    ModelOutputConfig,
+    SpecialTokensConfig,
+)
+from axolotl.utils.schemas.multimodal import MultiModalConfig
+from axolotl.utils.schemas.peft import LoraConfig, ReLoRAConfig
+from axolotl.utils.schemas.quantization import PTQConfig, QATConfig
+from axolotl.utils.schemas.training import (
+    HyperparametersConfig,
+    JaggedLRConfig,
+    SelectiveCheckpointingConfig,
+)
+from axolotl.utils.schemas.trl import TRLConfig
+from axolotl.utils.schemas.validation import ValidationMixin
+from axolotl.utils.schemas.vllm import VllmConfig
+
+LOG = get_logger(__name__)
+
+
+class EBFTConfig(BaseModel):
+    """Configuration for Energy-Based Fine-Tuning (EBFT)"""
+
+    feature_layers: list[float] = Field(
+        default=[0.25, 0.5, 0.75],
+        json_schema_extra={
+            "description": "Fractional layer depths for feature extraction (e.g., [0.25, 0.5, 0.75])"
+        },
+    )
+    embed_method: Literal["last_token", "mean_pooling", "completion_mean", "concat"] = (
+        Field(
+            default="last_token",
+            json_schema_extra={
+                "description": "Embedding method: 'last_token', 'mean_pooling', 'completion_mean', or 'concat'"
+            },
+        )
+    )
+    use_whitening: bool = Field(
+        default=False,
+        json_schema_extra={"description": "Apply SVD whitening to feature embeddings"},
+    )
+    alignment_coef: float = Field(
+        default=1.0,
+        json_schema_extra={
+            "description": "Coefficient for alignment reward (cosine similarity with ground truth)"
+        },
+    )
+    diversity_coef: float = Field(
+        default=1.0,
+        json_schema_extra={
+            "description": "Coefficient for diversity penalty (pairwise similarity between samples)"
+        },
+    )
+    ce_coef: float = Field(
+        default=0.0,
+        json_schema_extra={
+            "description": "Cross-entropy loss coefficient on ground-truth tokens"
+        },
+    )
+    adaptive_max_tokens: bool = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Set per-batch max_tokens based on ground-truth length"
+        },
+    )
+    gt_length_multiplier: float = Field(
+        default=1.5,
+        ge=0.1,
+        json_schema_extra={
+            "description": "Multiplier for ground-truth token count when computing adaptive max_tokens"
+        },
+    )
+
+    # Strided mode fields (for unstructured text)
+    mode: Literal["structured", "strided"] = Field(
+        default="structured",
+        json_schema_extra={
+            "description": "EBFT mode: 'structured' (QA with vLLM) or 'strided' (unstructured text)"
+        },
+    )
+    stride: int = Field(
+        default=8,
+        ge=1,
+        json_schema_extra={"description": "Stride between anchor points (tokens)"},
+    )
+    context_length: int = Field(
+        default=8,
+        ge=1,
+        json_schema_extra={"description": "Context window size per block"},
+    )
+    generate_max_len: int = Field(
+        default=8,
+        ge=1,
+        json_schema_extra={"description": "Tokens to generate per block"},
+    )
+    n_samples_per_prompt: int = Field(
+        default=4,
+        ge=1,
+        json_schema_extra={"description": "Independent rollouts per document"},
+    )
+    temperature: float = Field(
+        default=0.6,
+        ge=0.0,
+        json_schema_extra={
+            "description": "Sampling temperature for strided generation"
+        },
+    )
+    top_p: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        json_schema_extra={"description": "Top-p nucleus sampling threshold"},
+    )
+    rl_coef: float = Field(
+        default=1.0,
+        json_schema_extra={"description": "RL policy gradient loss coefficient"},
+    )
+    advantage_estimator: Literal["rloo", "group_norm", "reinforce"] = Field(
+        default="rloo",
+        json_schema_extra={
+            "description": "Advantage estimator: 'rloo', 'group_norm', 'reinforce'"
+        },
+    )
+    min_completion_prefix: int = Field(
+        default=0,
+        ge=0,
+        json_schema_extra={
+            "description": "Minimum tokens into completion before placing anchors. "
+            "Skips anchors too close to the prompt boundary where features are dominated by prompt context."
+        },
+    )
+
+
+class AxolotlInputConfig(
+    ModelInputConfig,
+    ModelOutputConfig,
+    LoraConfig,
+    ReLoRAConfig,
+    JaggedLRConfig,
+    HyperparametersConfig,
+    WandbConfig,
+    MLFlowConfig,
+    CometConfig,
+    TrackioConfig,
+    OpenTelemetryConfig,
+    LISAConfig,
+    GradioConfig,
+    RayConfig,
+    MultiModalConfig,
+    RemappedParameters,
+    DeprecatedParameters,
+    ValidationMixin,
+    BaseModel,
+):
+    """Wrapper of all config options."""
+
+    model_config = {"populate_by_name": True}
+
+    strict: bool | None = Field(
+        default=False,
+        json_schema_extra={"description": "Allow overwrite yml config using from cli"},
+    )
+    resume_from_checkpoint: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Resume from a specific checkpoint dir"},
+    )
+    auto_resume_from_checkpoints: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "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."
+        },
+    )
+    resize_token_embeddings_to_32x: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Resize the model embeddings when new tokens are added to multiples of 32. This is reported to improve training speed on some models"
+        },
+    )
+    mean_resizing_embeddings: bool | None = False
+    # optionally shrink the embeddings when the tokenizer vocab size is smaller
+    shrink_embeddings: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to shrink the embeddings to len(tokenizer). By default, we won't shrink."
+        },
+    )
+    embeddings_skip_upcast: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Don't upcast the embeddings to float32 when using PEFT. Useful for low-VRAM GPUs"
+        },
+    )
+    reinit_weights: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Reinitialize model weights randomly instead of loading pretrained weights"
+        },
+    )
+
+    trainer_cls: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "module to custom trainer class to use for training"
+        },
+    )
+
+    rl: RLType | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Use RL training: 'dpo', 'ipo', 'kto', 'simpo', 'orpo', 'grpo', 'ebft'"
+        },
+    )
+    trl: TRLConfig | None = Field(
+        default_factory=lambda: TRLConfig(),
+    )
+    vllm: VllmConfig | None = Field(
+        default_factory=lambda: VllmConfig(),
+    )
+    ebft: EBFTConfig | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Configuration for Energy-Based Fine-Tuning (EBFT)"
+        },
+    )
+    qat: QATConfig | None = None
+    quantization: PTQConfig | None = None
+    reward_model: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Reward modelling: `True` or `False`"},
+    )
+    dynamic_checkpoint: DynamicCheckpointConfig | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Configuration for dynamic checkpointing (trigger by file or signal). "
+            "Set 'enabled: true' to activate this feature."
+        },
+    )
+    process_reward_model: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Process reward modelling: `True` or `False`"
+        },
+    )
+    center_rewards_coefficient: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Coefficient to incentivize the reward model to output mean-zero rewards (proposed by https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`."
+        },
+    )
+    num_labels: int | None = None
+    # Whether to use weighting in DPO trainer.
+    # If `None`, default is `False` in the trainer.
+    dpo_use_weighting: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to perform weighting in DPO trainer"
+        },
+    )
+    dpo_label_smoothing: float | None = None
+    precompute_ref_log_probs: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Precompute reference model log probabilities for DPO"
+        },
+    )
+
+    dpo_use_liger_kernel: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Whether to use Liger kernel for DPO loss."},
+    )
+
+    dpo_padding_free: bool | None = None
+
+    dpo_loss_type: Annotated[list[str], MinLen(1)] | None = Field(
+        default=None,
+        json_schema_extra={"description": "List of DPO losses to use."},
+    )
+
+    dpo_loss_weights: Annotated[list[float], MinLen(1)] | None = Field(
+        default=None,
+        json_schema_extra={"description": "Weights for each DPO loss."},
+    )
+
+    datasets: (
+        Annotated[
+            list[
+                SFTDataset
+                | DPODataset
+                | KTODataset
+                | StepwiseSupervisedDataset
+                | SyntheticDataset
+            ],
+            MinLen(1),
+        ]
+        | None
+    ) = Field(
+        default=None,
+        json_schema_extra={
+            "description": "A list of one or more datasets to finetune the model with"
+        },
+    )
+
+    test_datasets: (
+        Annotated[
+            list[
+                SFTDataset
+                | DPODataset
+                | KTODataset
+                | StepwiseSupervisedDataset
+                | SyntheticDataset
+            ],
+            MinLen(1),
+        ]
+        | None
+    ) = Field(
+        default=None,
+        json_schema_extra={
+            "description": "A list of one or more datasets to eval the model with. You can use either test_datasets, or val_set_size, but not both."
+        },
+    )
+    shuffle_merged_datasets: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "If false, the datasets will not be shuffled and will keep their original order in `datasets`. The same applies to the `test_datasets` option and the `pretraining_dataset` option. Default is true."
+        },
+    )
+    shuffle_before_merging_datasets: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "If true, each dataset in `datasets` will be shuffled before merging. This allows curriculum learning strategies to be applied at the dataset level. Default is false."
+        },
+    )
+    dataset_prepared_path: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Axolotl attempts to save the dataset as an arrow after packing the data together so subsequent training attempts load faster, relative path"
+        },
+    )
+    dataset_shard_num: int | None = Field(
+        default=None, json_schema_extra={"description": "Num shards for whole dataset"}
+    )
+    dataset_shard_idx: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Index of shard to use for whole dataset"},
+    )
+    skip_prepare_dataset: bool | None = False
+    num_dataset_shards_to_save: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of shards to save the prepared dataset"
+        },
+    )
+
+    pretraining_dataset: (
+        Annotated[list[PretrainingDataset | SFTDataset], MinLen(1)] | None
+    ) = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Set to HF dataset for type: 'completion' for streaming instead of pre-tokenize"
+        },
+    )
+    dataset_processes: int | None = Field(
+        default=None,
+        deprecated="Use `dataset_num_proc` instead. This parameter will be removed in a future version.",
+        json_schema_extra={
+            "description": (
+                "The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()` if not set.\n"
+                "For Runpod VMs, it will default to number of vCPUs via RUNPOD_CPU_COUNT."
+            )
+        },
+    )
+    dataset_num_proc: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()` if not set.\n"
+                "For Runpod VMs, it will default to number of vCPUs via RUNPOD_CPU_COUNT."
+            )
+        },
+    )
+
+    dataset_exact_deduplication: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Deduplicates datasets and test_datasets with identical entries"
+        },
+    )
+    dataset_keep_in_memory: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Keep dataset in memory while preprocessing. Only needed if cached dataset is taking too much storage"
+        },
+    )
+    dataloader_pin_memory: bool | None = None
+    dataloader_num_workers: int | None = None
+    dataloader_prefetch_factor: int | None = None
+    dataloader_drop_last: bool | None = None
+
+    accelerator_config: dict[str, Any] | None = None
+
+    remove_unused_columns: bool | None = None
+
+    push_dataset_to_hub: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Push prepared dataset to hub - repo_org/repo_name"
+        },
+    )
+    hf_use_auth_token: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use hf `use_auth_token` for loading datasets. Useful for fetching private datasets. Required to be true when used in combination with `push_dataset_to_hub`"
+        },
+    )
+
+    device: Any | None = None
+    device_map: Any | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Passed through to transformers when loading the model when launched without accelerate. Use `sequential` when training w/ model parallelism to limit memory"
+        },
+    )
+    world_size: int | None = None
+    local_rank: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Don't mess with this, it's here for accelerate and torchrun"
+        },
+    )
+    ddp: bool | None = None
+
+    seed: int | None = Field(
+        default=None, json_schema_extra={"description": "Seed for reproducibility"}
+    )
+    ddp_timeout: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Advanced DDP Arguments - timeout"},
+    )
+    ddp_bucket_cap_mb: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Advanced DDP Arguments - bucket cap in MB"},
+    )
+    ddp_broadcast_buffers: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Advanced DDP Arguments - broadcast buffers"},
+    )
+    ddp_find_unused_parameters: bool | None = None
+
+    do_causal_lm_eval: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to run causal language model evaluation for metrics in `eval_causal_lm_metrics`"
+        },
+    )
+    eval_causal_lm_metrics: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "HF evaluate metrics used during evaluation. Default is ['sacrebleu', 'comet', 'ter', 'chrf', 'perplexity']"
+        },
+    )
+    do_bench_eval: bool | None = None
+    bench_dataset: str | None = None
+    bench_split: str | None = None
+    metric_for_best_model: str | None = None
+    greater_is_better: bool | None = None
+
+    loss_watchdog_threshold: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "High loss value, indicating the learning has broken down (a good estimate is ~2 times the loss at the start of training)"
+        },
+    )
+    loss_watchdog_patience: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of high-loss steps in a row before the trainer aborts (default: 3)"
+        },
+    )
+
+    gc_steps: int | None = Field(
+        default=None,
+        deprecated=(
+            "Use `torch_empty_cache_steps` to control CUDA cache clearing and "
+            "`gc_collect_steps` for Python garbage collection. "
+            "`gc_steps` will be removed in a future version."
+        ),
+        json_schema_extra={
+            "description": "Deprecated. Run garbage collection every `gc_steps` steps. Use `torch_empty_cache_steps` and `gc_collect_steps` instead."
+        },
+    )
+    torch_empty_cache_steps: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Steps between native HF Trainer `torch.cuda.empty_cache()` calls."
+        },
+    )
+    gc_collect_steps: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Steps between Python `gc.collect()` calls. -1 runs on epoch end and before evals only; None disables."
+        },
+    )
+
+    bf16: Literal["auto"] | bool | None = Field(
+        default="auto",
+        json_schema_extra={
+            "description": "Use CUDA bf16. bool or 'full' for `bf16_full_eval`, or 'auto' for automatic detection. require >=ampere"
+        },
+    )
+    fp16: bool | None = Field(
+        default=None, json_schema_extra={"description": "Use CUDA fp16"}
+    )
+    fp8: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Enable FP8 mixed precision training using TorchAO. Best "
+            "used in combination with torch.compile."
+        },
+    )
+    fp8_enable_fsdp_float8_all_gather: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Enable FSDP float8 all-gather optimization for FP8 training. Can "
+            "improve training speed by 10-15% when FSDP is enabled."
+        },
+    )
+    bfloat16: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "No AMP (automatic mixed precision) - require >=ampere"
+        },
+    )  # for non-AMP cases
+    float16: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "No AMP (automatic mixed precision)"},
+    )  # for non-AMP cases
+    tf32: Literal["auto"] | bool | None = Field(
+        default="auto",
+        json_schema_extra={
+            "description": "bool to use CUDA tf32 or 'auto' for automatic detection - require >=ampere"
+        },
+    )
+    float32: bool | None = None
+
+    gradient_checkpointing: Literal["offload", "offload_disk"] | bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Whether to use gradient checkpointing. Available options are: true, false, 'offload', 'offload_disk'. https://huggingface.co/docs/transformers/v4.18.0/en/performance#gradient-checkpointing"
+        },
+    )
+    gradient_checkpointing_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Additional kwargs to pass to the trainer for gradient checkpointing"
+        },
+    )
+    selective_checkpointing: SelectiveCheckpointingConfig | bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Selective activation checkpointing: within each gradient-checkpointed "
+                "layer, save the listed ops during forward instead of recomputing them "
+                "in backward. `true` saves attention (SDPA/flash-attention). Requires "
+                "gradient_checkpointing with non-reentrant checkpointing."
+            )
+        },
+    )
+    activation_offloading: Literal["legacy", "disk", "hidden_states"] | bool | None = (
+        Field(
+            default=False,
+            json_schema_extra={
+                "description": (
+                    "Whether to offload activations. Options: true/false, 'legacy', "
+                    "'disk' (TRL offloader), or 'hidden_states' (checkpoint input "
+                    "offload; non-reentrant by default, ALST-style when "
+                    "gradient_checkpointing_kwargs.use_reentrant is true)."
+                )
+            },
+        )
+    )
+    layer_offloading: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Offload model layer parameters to CPU during forward, prefetch back during backward."
+        },
+    )
+
+    freeze_mm_modules: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Freeze multimodal encoder parameters (vision, audio, etc.) for "
+            "text-only training of multimodal models. When True, parameters belonging to "
+            "vision towers, audio towers, multimodal projectors, and similar non-language "
+            "modules are frozen (requires_grad=False). This allows DDP training without "
+            "ddp_find_unused_parameters=True."
+        },
+    )
+
+    unfrozen_parameters: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "List of regex patterns for parameter names to keep unfrozen. "
+            "All other parameters will be frozen via requires_grad=False. "
+            "Note: range-based patterns (e.g. embed_tokens.weight$[:32000]) use gradient "
+            "zeroing rather than a true freeze, so weight decay will still apply to the "
+            "frozen portion and optimizer states are allocated for the full parameter."
+        },
+    )
+
+    sequence_len: int = Field(
+        default=512,
+        json_schema_extra={
+            "description": "The maximum length of an input to train with, this should typically be less than 2048 as most models have a token/context limit of 2048"
+        },
+    )
+    excess_length_strategy: Literal["drop", "truncate", "raise"] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "What to do when a tokenized row exceeds sequence_len. 'drop' removes the row; 'truncate' slices tensors to sequence_len; 'raise' raises a ValueError. Defaults to 'drop' for backward compatibility."
+        },
+    )
+    eval_sequence_len: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The maximum length of an input for evaluation. If not specified, defaults to sequence_len"
+        },
+    )
+    min_sample_len: int | None = None
+    max_prompt_len: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "maximum prompt length for RL training"},
+    )
+    sample_packing: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Use efficient multi-packing with block diagonal attention and per sequence position_ids. Recommend set to 'true'"
+        },
+    )
+    sample_packing_group_size: int | None = Field(
+        default=100_000,
+        json_schema_extra={
+            "description": "The number of samples packed at a time. Increasing the following values helps with packing, but usually only slightly (<%1.)"
+        },
+    )
+    sample_packing_bin_size: int | None = Field(
+        default=200,
+        json_schema_extra={
+            "description": "The number of samples which can be packed into one sequence. Increase if using a large sequence_len with many short samples."
+        },
+    )
+    sample_packing_sequentially: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Whether to pack samples sequentially"},
+    )
+    sample_packing_mp_start_method: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The multiprocessing start method to use for packing. Should be 'fork', 'spawn' or 'forkserver'"
+        },
+    )
+    eval_sample_packing: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Set to 'false' if getting errors during eval with sample_packing on"
+        },
+    )
+    pad_to_sequence_len: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Pad inputs so each step uses constant sized buffers. This will reduce memory fragmentation and may prevent OOMs, by re-using memory more efficiently. Defaults to True if `sample_packing` enabled"
+        },
+    )
+    pad_to_multiple_of: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": ("Pad each batch to a multiple of this value.")
+        },
+    )
+    curriculum_sampling: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use sequential sampling for curriculum learning"
+        },
+    )
+    multipack_real_batches: bool | None = None
+
+    batch_flattening: Literal["auto"] | bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Use batch flattening for speedups when not using sample_packing"
+        },
+    )
+
+    # for PoSE context length extension
+    use_pose: bool | None = None
+    pose_split_on_token_ids: list[int] | None = None
+    pose_max_context_len: int | None = None
+    pose_num_chunks: int | None = None
+
+    # Deprecated: Use streaming_multipack_buffer_size instead
+    pretrain_multipack_buffer_size: int | None = Field(
+        default=None,
+        deprecated="Deprecated in v0.13.0, will be removed in v0.14.0. Use streaming_multipack_buffer_size instead",
+    )
+    pretrain_multipack_attn: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "whether to prevent cross attention for packed sequences during pretraining",
+        },
+    )
+    pretraining_sample_concatenation: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "whether to concatenate samples during pretraining",
+        },
+    )
+
+    streaming: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Use streaming mode for loading datasets"},
+    )
+    streaming_multipack_buffer_size: int | None = Field(
+        default=10_000,
+        json_schema_extra={
+            "description": "Buffer size for multipack streaming datasets"
+        },
+    )
+
+    xformers_attention: bool | None = Field(
+        default=None,
+        deprecated="Use `attn_implementation: xformers` instead.",
+        json_schema_extra={
+            "description": "[DEPRECATED] Use `attn_implementation: xformers`. https://github.com/facebookresearch/xformers"
+        },
+    )
+    sdp_attention: bool | None = Field(
+        default=None,
+        deprecated="Use `attn_implementation: sdpa` instead.",
+        json_schema_extra={
+            "description": "[DEPRECATED] Use `attn_implementation: sdpa`."
+        },
+    )
+    flex_attention: bool | None = Field(
+        default=None,
+        deprecated="Use `attn_implementation: flex_attention` instead.",
+    )
+    flex_attn_compile_kwargs: dict[str, Any] | None = None
+    flash_attention: bool | None = Field(
+        default=None,
+        deprecated="Use `attn_implementation: flash_attention_2` instead.",
+        json_schema_extra={
+            "description": "[DEPRECATED] Use `attn_implementation: flash_attention_2`. https://github.com/Dao-AILab/flash-attention"
+        },
+    )
+    flash_attn_cross_entropy: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use flash-attention cross entropy implementation - advanced use only"
+        },
+    )
+    flash_attn_fuse_mlp: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to fuse part of the MLP into a single operation"
+        },
+    )
+    flash_optimum: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Whether to use bettertransformers"},
+    )
+    sage_attention: bool | None = Field(
+        default=None,
+        deprecated="Use `attn_implementation: sage` instead.",
+        json_schema_extra={
+            "description": "[DEPRECATED] Use `attn_implementation: sage`. https://github.com/thu-ml/SageAttention"
+        },
+    )
+
+    eager_attention: bool | None = Field(
+        default=None,
+        deprecated="Use `attn_implementation: eager` instead.",
+    )
+
+    attn_implementation: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Attention backend. Canonical values: eager, sdpa, flash_attention_2, "
+                "flash_attention_3, flex_attention, xformers, sage, fp8. Hub-kernel "
+                "paths (e.g. kernels-community/flash-attn3) are also accepted and passed "
+                "through to transformers."
+            )
+        },
+    )
+
+    gemma4_hybrid_attn_impl: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Use hybrid attention for Gemma 4: flash_attention_2 for sliding window layers "
+            "and sdpa for global (full_attention) layers. Global layers have head_dim=512 which "
+            "exceeds flash attention's supported size."
+        },
+    )
+
+    large_head_attention: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Generic capability for attention layers with head_dim > 256 (which flash/cuDNN "
+                "SDPA can't serve): 'sdpa' (stock SDPA, default) | 'auto' (Triton flash kernel on "
+                "packed rows, SDPA otherwise) | 'triton_flash' (prefer the Triton kernel). Applies "
+                "to any SDPA model; Gemma-4's head_dim=512 global layers use it automatically."
+            )
+        },
+    )
+    # Deprecated alias for `large_head_attention` (True -> 'auto'); kept for backwards compatibility.
+    flash_attn_d512: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Deprecated: use `large_head_attention: auto`. Routes head_dim>256 attention "
+                "through the Triton flash kernel."
+            )
+        },
+    )
+
+    sdpa_varlen: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "With sample packing + attn_implementation=sdpa, route packed rows through "
+                "torch.nn.attention.varlen.varlen_attn (cu_seqlens) instead of an explicit 4D "
+                "block-diagonal mask. Skips cross-document blocks (faster + lower memory) with no "
+                "flash_attn dependency. Left unset (null) it auto-enables when supported (torch >= "
+                "2.10, head_dim <= 256, no sliding window); set true/false to force. When it can't "
+                "be used, packed rows still isolate documents via the block-diagonal mask. "
+                "Sliding-window attention needs torch >= 2.11 (varlen_attn window_size)."
+            )
+        },
+    )
+
+    fused_attn_kernel: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Replace ``q_norm + apply_rotary_pos_emb`` (and the matching k path) "
+                "with a single fused RMSNorm+RoPE Triton kernel launch. Currently "
+                "implemented for Qwen3, Qwen3-MoE, Qwen3.5, and Qwen3.5-MoE "
+                "full-attention layers; Gemma 4 always uses the fused path. Disabled "
+                "(None/False) falls back "
+                "to the eager transformers implementation. Compile-safe via "
+                "torch.library.triton_op — traces under torch.compile(fullgraph=True). "
+                "Per-step wins are arch-dependent: ~+7-12% across sm_86 and sm_120. "
+                "Combining with torch_compile=true is a clear win on sm_120 (+9% "
+                "extra) but currently regresses on sm_86 due to Inductor autotune "
+                "biases — flip them on independently and benchmark."
+            )
+        },
+    )
+
+    experts_implementation: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Which experts implementation to use for MoE models,"
+        },
+    )
+
+    quantize_moe_experts: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Quantize MoE expert weights on load to reduce VRAM. "
+            "Requires adapter (lora/qlora) with load_in_4bit or load_in_8bit. "
+            "Requires CUDA (not compatible with ROCm or other backends). "
+            "Note: total parameter count may be reported incorrectly when enabled "
+            "(trainable param count is correct)."
+        },
+    )
+
+    scaling_softmax: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use Scaled Softmax (SSMax) attention. Ref: https://arxiv.org/abs/2501.19399"
+        },
+    )
+    scaling_softmax_factor: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Scaling factor for SSMax attention. Default is 0.43"
+        },
+    )
+    scaling_softmax_bias: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Bias for SSMax attention. Default is 0.0. Note: The paper recommends bias=0 for better length generalization."
+        },
+    )
+
+    lora_mlp_kernel: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Apply custom LoRA autograd functions and activation function Triton kernels for speed and memory savings. See: https://docs.axolotl.ai/docs/lora_optims.html"
+        },
+    )
+    lora_qkv_kernel: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Apply custom LoRA autograd functions and activation function Triton kernels for speed and memory savings. See: https://docs.axolotl.ai/docs/lora_optims.html"
+        },
+    )
+    lora_o_kernel: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Apply custom LoRA autograd functions and activation function Triton kernels for speed and memory savings. See: https://docs.axolotl.ai/docs/lora_optims.html"
+        },
+    )
+    lora_embedding_kernel: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Apply custom LoRA autograd function for embedding layers. See: https://docs.axolotl.ai/docs/lora_optims.html"
+        },
+    )
+
+    chunked_cross_entropy: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use chunked cross entropy loss for memory efficiency"
+        },
+    )
+    chunked_cross_entropy_num_chunks: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of chunks to use for chunked cross entropy loss"
+        },
+    )
+    use_eaft: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Enable Entropy-Aware Focal Training loss (EAFT)"
+        },
+    )
+    eaft_alpha: float | None = Field(
+        default=1.0,
+        json_schema_extra={
+            "description": "Exponent for entropy weighting in EAFT (default: 1.0)"
+        },
+    )
+    eaft_k: int | None = Field(
+        default=20,
+        json_schema_extra={
+            "description": "Number of top logits for entropy approximation (default: 20)"
+        },
+    )
+
+    tiled_mlp: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use ALST tiled mlp for memory efficient long context"
+        },
+    )
+
+    tiled_mlp_num_shards: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of shards to use for ALST tiled mlp. If unset, it will be set based on seqlen/hidden_size"
+        },
+    )
+
+    tiled_mlp_use_original_mlp: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Whether to use original mlp for ALST tiled mlp. Otherwise uses a generic MLP based on llama."
+        },
+    )
+
+    llama4_linearized_experts: bool | None = None
+
+    deepspeed: str | dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Deepspeed config path. e.g., deepspeed_configs/zero3.json"
+        },
+    )
+    deepcompile: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use deepcompile for faster training with deepspeed"
+        },
+    )
+    fsdp: list[str] | None = Field(
+        default=None,
+        json_schema_extra={"description": "FSDP configuration"},
+        deprecated="Configuring FSDP using `fsdp` is deprecated. Please use `fsdp_config` instead. ",
+    )
+    fsdp_config: FSDPConfig | None = Field(
+        default=None, json_schema_extra={"description": "FSDP configuration options"}
+    )
+    fsdp_version: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "FSDP version"},
+    )
+    fp32_norms: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Keep norm modules (RMSNorm/LayerNorm) in fp32 by sharding them "
+                "under their own FSDP2 MixedPrecisionPolicy. Requires fsdp_version: 2."
+            )
+        },
+    )
+    fp32_norm_classes: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Class-name patterns to match for fp32 norm sharding. Patterns "
+                "without a '.' match against type(module).__name__ as a suffix. "
+                "Patterns containing a '.' match the fully qualified class path "
+                "exactly. Defaults to ['RMSNorm', 'LayerNorm'] when fp32_norms is "
+                "true and this is unset."
+            )
+        },
+    )
+    fsdp_final_state_dict_type: (
+        Literal["FULL_STATE_DICT", "LOCAL_STATE_DICT", "SHARDED_STATE_DICT"] | None
+    ) = Field(
+        default=None,
+        deprecated="Configuring FSDP final state dict type using `fsdp_final_state_dict_type` is deprecated. Please use `fsdp_config.final_state_dict_type` instead.",
+    )
+
+    val_set_size: float | None = Field(
+        default=0.0,
+        json_schema_extra={
+            "description": "How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc. 0 for no eval."
+        },
+    )
+
+    dp_shard_size: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of devices to shard across. If not set, will use all available devices."
+        },
+    )
+    dp_replicate_size: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Number of devices to replicate across."},
+    )
+    sequence_parallel_degree: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Deprecated: use `context_parallel_size` instead"
+        },
+    )
+    context_parallel_size: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "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://docs.axolotl.ai/docs/sequence_parallelism.html for more details."
+        },
+    )
+    heads_k_stride: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "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."
+        },
+    )
+    ring_attn_func: RingAttnFunc | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "One of 'varlen_llama3', 'batch_ring', 'batch_zigzag', 'batch_stripe'. Defaults to 'varlen_llama3' in the sample packing case, and 'batch_ring' in the non-sample packing case."
+        },
+    )
+    tensor_parallel_size: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of tensor parallel processes in TP group. Only supported with DeepSpeed AutoTP."
+        },
+    )
+    special_tokens: SpecialTokensConfig | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Add or change special tokens. If you add tokens here, you don't need to add them to the `tokens` list."
+        },
+    )
+    tokens: list[str] | None = Field(
+        default=None,
+        json_schema_extra={"description": "Add extra tokens to the tokenizer"},
+    )
+    added_tokens_overrides: dict[int, str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "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."
+        },
+    )
+
+    torch_compile: Literal["auto"] | bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use torch.compile and which backend to use."
+        },
+    )
+    torch_compile_backend: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Backend to use for torch.compile"},
+    )
+    torch_compile_mode: Literal["default", "reduce-overhead", "max-autotune"] | None = (
+        None
+    )
+
+    max_steps: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Maximum number of iterations to train for. It precedes num_epochs which means that if both are set, num_epochs will not be guaranteed. e.g., when 1 epoch is 1000 steps => `num_epochs: 2` and `max_steps: 100` will train for 100 steps"
+        },
+    )
+    warmup_steps: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of warmup steps. Cannot use with warmup_ratio"
+        },
+    )
+    warmup_ratio: float | None = Field(
+        default=None,
+        json_schema_extra={"description": "Warmup ratio. Cannot use with warmup_steps"},
+    )
+    eval_steps: int | float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Leave empty to eval at each epoch, integer for every N steps. float for fraction of total steps"
+        },
+    )
+    evals_per_epoch: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of times per epoch to run evals, mutually exclusive with eval_steps"
+        },
+    )
+    eval_strategy: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Set to `no` to skip evaluation, `epoch` at end of each epoch, leave empty to infer from `eval_steps`"
+        },
+    )
+
+    save_steps: int | float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Leave empty to save at each epoch, integer for every N steps. float for fraction of total steps"
+        },
+    )
+    saves_per_epoch: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of times per epoch to save a checkpoint, mutually exclusive with save_steps"
+        },
+    )
+    save_strategy: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Set to `no` to skip checkpoint saves, `epoch` at end of each epoch, `best` when better result is achieved, leave empty to infer from `save_steps`"
+        },
+    )
+    save_total_limit: int | None = Field(
+        default=None, json_schema_extra={"description": "Checkpoints saved at a time"}
+    )
+    save_first_step: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to checkpoint a model after the first step of training. Defaults to False."
+        },
+    )
+
+    logging_steps: int | None = Field(
+        default=None, json_schema_extra={"description": "Logging frequency"}
+    )
+    early_stopping_patience: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Stop training after this many evaluation losses have increased in a row. https://huggingface.co/transformers/v4.2.2/_modules/transformers/trainer_callback.html#EarlyStoppingCallback"
+        },
+    )
+    load_best_model_at_end: bool | None = False
+    save_only_model: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Save only the model weights, skipping the optimizer. Using this means you can't resume from checkpoints."
+        },
+    )
+    use_tensorboard: bool | None = Field(
+        default=None, json_schema_extra={"description": "Use tensorboard for logging"}
+    )
+    profiler_steps: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Enable the pytorch profiler to capture the first N steps of training to the output_dir. see https://pytorch.org/blog/understanding-gpu-memory-1/ for more information. Snapshots can be visualized @ https://pytorch.org/memory_viz"
+        },
+    )
+    profiler_steps_start: int | None = Field(
+        default=0,
+        json_schema_extra={
+            "description": "Which step to start the profiler at. Useful for only capturing a few steps mid-run."
+        },
+    )
+    include_tokens_per_second: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "bool of whether to report tokens per second at the end of training. This is not supported with pre-training datasets."
+        },
+    )
+    include_tkps: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "bool of whether to report tokens per second per-gpu during training by measuring throughput of non-padding tokens."
+        },
+    )
+    neftune_noise_alpha: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "NEFT https://arxiv.org/abs/2310.05914, set this to a number (paper default is 5) to add noise to embeddings. Currently only supported on Llama and Mistral"
+        },
+    )
+
+    orpo_alpha: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Parameter controlling the relative ratio loss weight in the ORPO loss. Passed to `beta` in `ORPOConfig` due to trl mapping."
+        },
+    )
+    simpo_gamma: float | None = Field(
+        default=None,
+        json_schema_extra={"description": "Target reward margin for the SimPO loss"},
+    )
+    cpo_alpha: float | None = Field(
+        default=None, json_schema_extra={"description": "Weight of the BC regularizer"}
+    )
+
+    kto_desirable_weight: float | None = Field(
+        default=None,
+        json_schema_extra={"description": "Factor for desirable loss term in KTO loss"},
+    )
+    kto_undesirable_weight: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Factor for undesirable loss term in KTO loss"
+        },
+    )
+    rl_beta: float | None = Field(
+        default=None,
+        json_schema_extra={"description": "The beta parameter for the RL training"},
+    )
+
+    max_memory: dict[int | Literal["cpu", "disk"], int | str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Defines the max memory usage per gpu on the system. Passed through to transformers when loading the model."
+        },
+    )
+    gpu_memory_limit: int | str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Limit the memory for all available GPUs to this amount (if an integer, expressed in gigabytes); default: unset"
+        },
+    )
+    low_cpu_mem_usage: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Whether to use low_cpu_mem_usage"},
+    )
+
+    chat_template: (
+        ChatTemplate
+        | Annotated[str, StringConstraints(pattern="^tokenizer_default_fallback_")]
+    ) | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The name of the chat template to use for training, following values are supported: tokenizer_default: Uses the chat template that is available in the tokenizer_config.json. If the chat template is not available in the tokenizer, it will raise an error. This is the default value. alpaca/inst/chatml/gemma/cohere/llama3/phi_3/deepseek_v2/jamba: These chat templates are available in the axolotl codebase at src/axolotl/utils/chat_templates.py. tokenizer_default_fallback_*: where * is the name of the chat template to fallback to. E.g. tokenizer_default_fallback_chatml. This is useful when the chat template is not available in the tokenizer. jinja: Uses a custom jinja template for the chat template. The custom jinja template should be provided in the chat_template_jinja field. The selected chat template will be saved to the tokenizer_config.json for easier inferencing"
+        },
+    )
+    chat_template_jinja: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Custom jinja template or path to jinja file for chat template. This will be only used if chat_template is set to `jinja` or `null` (in which case chat_template is automatically set to `jinja`). Default is null."
+        },
+    )
+    chat_template_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Additional kwargs to pass to the chat template. This is useful for customizing the chat template. For example, you can pass `thinking=False` to add a generation prompt to the chat template."
+        },
+    )
+    eot_tokens: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Custom EOT (End-of-Turn) tokens to mask/unmask during training. These tokens mark the boundaries between conversation turns. For example: ['/INST', '</s>', '[/SYSTEM_PROMPT]']. If not specified, defaults to just the model's eos_token. This is useful for templates that use multiple delimiter tokens."
+        },
+    )
+    default_system_message: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Changes the default system message. Currently only supports chatml."
+        },
+    )
+
+    fix_untrained_tokens: int | list[int] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Token index or indices to adjust embedding weights to the mean of the other tokens. "
+                "This is useful when the model has untrained embeddings."
+            )
+        },
+    )
+
+    # INTERNALS - document for now, generally not set externally
+    is_preprocess: bool | None = None
+    preprocess_iterable: bool | None = None
+
+    total_num_tokens: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Total number of tokens - internal use"},
+    )
+    total_supervised_tokens: int | None = None
+    sample_packing_eff_est: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "You can set these packing optimizations AFTER starting a training at least once. The trainer will provide recommended values for these values."
+        },
+    )
+    axolotl_config_path: str | None = None
+
+    is_falcon_derived_model: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Internal use only - Used to identify which the model is based on"
+        },
+    )
+    is_llama_derived_model: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Internal use only - Used to identify which the model is based on"
+        },
+    )
+    is_mistral_derived_model: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Internal use only - Used to identify which the model is based on. Please note that if you set this to true, `padding_side` will be set to 'left' by default"
+        },
+    )
+    is_qwen_derived_model: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Internal use only - Used to identify which the model is based on"
+        },
+    )
+
+    plugins: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Add plugins to extend the pipeline. See `src/axolotl/integrations` for the available plugins or doc below for more details. https://docs.axolotl.ai/docs/custom_integrations.html"
+        },
+    )
+    generate_samples: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Enable sample generation during training for monitoring"
+        },
+    )
+    num_generation_samples: int | None = Field(
+        default=3,
+        json_schema_extra={
+            "description": "Number of samples to generate at each interval"
+        },
+    )
+    generation_max_new_tokens: int | None = Field(
+        default=50,
+        json_schema_extra={"description": "Maximum new tokens to generate per sample"},
+    )
+    generation_temperature: float | None = Field(
+        default=0.7,
+        json_schema_extra={
+            "description": "Temperature for sample generation (0.0 = greedy)"
+        },
+    )
+    generation_top_p: float | None = Field(
+        default=None,
+        json_schema_extra={"description": "Nucleus sampling parameter for generation"},
+    )
+    generation_top_k: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Top-k sampling parameter for generation"},
+    )
+    generation_prompt_ratio: float | None = Field(
+        default=0.5,
+        json_schema_extra={"description": "Ratio of input to use as prompt (0.0-1.0)"},
+    )
+    generation_do_sample: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Whether to use sampling (vs greedy decoding)"
+        },
+    )
+
+    @field_serializer("datasets")
+    def datasets_serializer(
+        self, ds_configs: list[DatasetConfig] | None
+    ) -> list[dict[str, Any]] | None:
+        if ds_configs:
+            return [ds_config.model_dump(exclude_none=True) for ds_config in ds_configs]
+        return None
+
+    # --- Attention capability flags (derived from attn_implementation) ---
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def attn_supports_packing(self) -> bool:
+        return self.attn_implementation in ATTN_IMPLS_SUPPORTING_PACKING
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def attn_decontaminates_packing(self) -> bool:
+        # sdpa/eager isolate packed docs via the dropped-mask block-diagonal; others can't.
+        if self.attn_supports_packing:
+            return True
+        return self.attn_implementation in ("sdpa", "eager")
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def attn_uses_flash_lib(self) -> bool:
+        return self.attn_implementation in ATTN_IMPLS_USING_FLASH_LIB
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def attn_needs_dtype_cast(self) -> bool:
+        if self.attn_implementation is None:
+            return False
+        return self.attn_implementation not in ATTN_IMPLS_WITHOUT_DTYPE_CAST
+
+    @model_validator(mode="before")
+    @classmethod
+    def warn_peft_trainable_token_to_fix_untrained(cls, data):
+        if (
+            peft_trainable_token_indices := data.get("peft_trainable_token_indices")
+        ) and (fix_untrained_tokens := data.get("fix_untrained_tokens")):
+            if isinstance(fix_untrained_tokens, int):
+                fix_untrained_tokens = (fix_untrained_tokens,)
+
+            if isinstance(peft_trainable_token_indices, int):
+                peft_trainable_token_indices = (peft_trainable_token_indices,)
+
+            for untrained_token_id in fix_untrained_tokens:
+                if untrained_token_id not in peft_trainable_token_indices:
+                    LOG.warning_once(
+                        f"Token {untrained_token_id} is fixed via `fix_untrained_tokens`, yet not in `peft_trainable_token_indices: ` list. "
+                        "Please add it, otherwise the token won't be trained on."
+                    )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def normalize_attn_implementation(cls, data):
+        """Map legacy boolean attention flags to canonical attn_implementation, warn, then strip."""
+        if not isinstance(data, dict):
+            return data
+
+        attn_impl = data.get("attn_implementation")
+        set_flags = [f for f in LEGACY_ATTN_FLAG_TO_IMPL if data.get(f)]
+
+        # gemma4_hybrid requires flash_attention_2 for the sliding-window layers;
+        # post-load patching swaps global layers to sdpa (see
+        # `_apply_gemma_hybrid_attention`). Default it in when the user didn't
+        # pick a backend; reject any incompatible explicit choice.
+        if data.get("gemma4_hybrid_attn_impl"):
+            if not attn_impl and not set_flags:
+                data["attn_implementation"] = "flash_attention_2"
+                attn_impl = "flash_attention_2"
+            elif attn_impl and attn_impl != "flash_attention_2":
+                raise ValueError(
+                    f"gemma4_hybrid_attn_impl requires attn_implementation="
+                    f"flash_attention_2 (sliding-window layers run under FA2); "
+                    f"got {attn_impl!r}."
+                )
+
+        if attn_impl and set_flags:
+            raise ValueError(
+                f"attn_implementation={attn_impl!r} cannot be combined with legacy "
+                f"attention flags ({', '.join(sorted(set_flags))}). The legacy "
+                f"flags are deprecated — set only `attn_implementation`."
+            )
+
+        if not attn_impl and set_flags:
+            # Priority: specific backends beat generic flash/sdp/eager fallbacks.
+            for flag in LEGACY_ATTN_FLAG_TO_IMPL:
+                if flag in set_flags:
+                    canonical = LEGACY_ATTN_FLAG_TO_IMPL[flag]
+                    data["attn_implementation"] = canonical
+                    LOG.warning(
+                        "`%s: true` is deprecated and will be removed in a future "
+                        "release. Use `attn_implementation: %s` instead.",
+                        flag,
+                        canonical,
+                    )
+                    break
+
+        # Strip legacy flags from validated data — canonical field is authoritative.
+        for flag in LEGACY_ATTN_FLAG_TO_IMPL:
+            data.pop(flag, None)
+
+        return data
+
+    @field_validator("large_head_attention", mode="before")
+    @classmethod
+    def validate_large_head_attention(cls, value):
+        """Only the three known policies are valid (a typo must not silently opt into Triton)."""
+        if value is None:
+            return None
+        valid = {"auto", "sdpa", "triton_flash"}
+        if str(value).lower() not in valid:
+            raise ValueError(
+                f"large_head_attention must be one of {sorted(valid)}, got {value!r}"
+            )
+        return str(value).lower()
+
+    @field_validator("attn_implementation", mode="before")
+    @classmethod
+    def validate_attn_implementation(cls, value):
+        """Accept canonical names and hub-kernel paths; reject short-form aliases."""
+        if value is None:
+            return None
+        if not isinstance(value, str):
+            raise TypeError(
+                f"attn_implementation must be a string, got {type(value).__name__}"
+            )
+        if value in CANONICAL_ATTN_IMPLS:
+            return value
+        if "/" in value:
+            # Hub-kernel path, e.g. "kernels-community/flash-attn3". Pass through.
+            return value
+        if value in SHORT_FORM_ALIAS_TO_CANONICAL:
+            canonical = SHORT_FORM_ALIAS_TO_CANONICAL[value]
+            raise ValueError(
+                f"attn_implementation={value!r} is not accepted. "
+                f"Use the canonical name {canonical!r} instead."
+            )
+        raise ValueError(
+            f"attn_implementation={value!r} is not a recognized backend. "
+            f"Expected one of: {sorted(CANONICAL_ATTN_IMPLS)}, or a hub-kernel "
+            f"path containing '/'."
+        )
+
+    @model_validator(mode="after")
+    def check_fp32_norms(self):
+        if self.fp32_norms:
+            # FSDP must actually be configured — fsdp_version alone is not
+            # sufficient since the rest of axolotl treats fsdp_config as the
+            # canonical "is_fsdp" signal.
+            if self.fsdp_config is None:
+                raise ValueError(
+                    "fp32_norms requires FSDP to be enabled "
+                    "(fsdp_config block must be set)."
+                )
+            if str(self.fsdp_version) != "2":
+                raise ValueError(
+                    "fp32_norms requires fsdp_version: 2. FSDP1's flat-param "
+                    "dtype uniformity constraint is incompatible with keeping "
+                    "norms in fp32 while decoder layers run in bf16."
+                )
+        if self.fp32_norm_classes and not self.fp32_norms:
+            LOG.warning(
+                "fp32_norm_classes is set but fp32_norms is not enabled; "
+                "it will be ignored."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_sdpa_varlen(self):
+        if self.sdpa_varlen:
+            if not self.sample_packing:
+                LOG.warning(
+                    "`sdpa_varlen` only affects packed rows; enable `sample_packing` or it is a no-op."
+                )
+            try:
+                from torch.nn.attention.varlen import varlen_attn  # noqa: F401
+            except ImportError:
+                LOG.warning(
+                    "`sdpa_varlen` needs torch >= 2.10 (torch.nn.attention.varlen); it will be "
+                    "ignored and stock SDPA used."
+                )
+        return self
+
+    @model_validator(mode="after")
+    def check_sageattn_wo_sample_packing(self):
+        if (
+            self.attn_implementation == "sage"
+            and not self.sample_packing
+            and not self.pad_to_sequence_len
+        ):
+            LOG.warning(
+                "We recommend turning on `pad_to_sequence_len` for SageAttention "
+                "without packing. The loss has been observed to explode otherwise."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_sageattn_fft(self):
+        if self.attn_implementation == "sage" and not self.adapter:
+            LOG.warning(
+                "SageAttention full finetuning has been observed to drop loss to 0. "
+                "Monitor the loss, or switch to LoRA/QLoRA or another attention method."
+            )
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_save_strategy_best_requires_metric(cls, data):
+        if data.get("save_strategy") == "best" and not data.get(
+            "metric_for_best_model"
+        ):
+            raise ValueError(
+                "save_strategy: 'best' requires metric_for_best_model to be set. "
+                "Please specify the metric to use, e.g. metric_for_best_model: eval_loss"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_lora_target_modules_regex(cls, data):
+        lora_target_modules = data.get("lora_target_modules")
+        if not isinstance(lora_target_modules, list):
+            return data
+        invalid = []
+        for pattern in lora_target_modules:
+            if not isinstance(pattern, str):
+                continue
+            try:
+                re.compile(pattern)
+            except re.error:
+                invalid.append(pattern)
+        if invalid:
+            raise ValueError(
+                f"lora_target_modules contains invalid regex pattern(s): {invalid}. "
+                "Please provide valid Python regex patterns or plain module name strings."
+            )
+        return data
+
+
+class AxolotlConfigWCapabilities(AxolotlInputConfig):
+    """Wrapper to valdiate GPU capabilities with the configured options"""
+
+    capabilities: GPUCapabilities
+    env_capabilities: EnvCapabilities
+
+    @model_validator(mode="after")
+    def check_bf16(self):
+        if self.capabilities.bf16:
+            if not self.bf16 and not self.bfloat16:
+                LOG.info(
+                    "bf16 support detected, but not enabled for this configuration."
+                )
+        else:
+            if (
+                not self.merge_lora
+                and not self.is_preprocess
+                and (self.bf16 is True or self.bfloat16 is True)
+            ):
+                LOG.warning(
+                    "bf16 requested, but AMP is not supported on this GPU. Requires Ampere series or above. Training will fail, but other operations (such as merging) are still functional."
+                )
+        return self
+
+    @model_validator(mode="after")
+    def check_tf32(self):
+        if self.tf32 == "auto":
+            self.tf32 = self.capabilities.tf32
+        return self
+
+    @model_validator(mode="after")
+    def check_fp8(self):
+        if self.fp8 and not self.capabilities.fp8:
+            raise ValueError("fp8 requested, but fp8 is not supported on this GPU")
+        elif self.fp8 and self.capabilities.fp8 and not is_fp8_available():
+            raise ValueError(
+                "fp8 requested, but missing one of ms-amp, transformers-engine or torchao."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_sample_packing_w_sdpa_bf16(self):
+        is_sm_90 = self.capabilities and self.capabilities.compute_capability == "sm_90"
+        if (
+            self.sample_packing
+            and self.attn_implementation == "sdpa"
+            and (self.bfloat16 or self.bf16)
+            and not is_sm_90
+        ):
+            # https://github.com/pytorch/pytorch/blob/1b03423526536b5f3d35bdfa95ccc6197556cf9b/test/test_transformers.py#L2440-L2450
+            LOG.warning(
+                "sample_packing & torch sdpa with bf16 is unsupported may results in 0.0 loss. "
+                "This may work on H100s."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_compute_capability_w_sageattn(self):
+        if (
+            self.attn_implementation == "sage"
+            and self.capabilities
+            and self.capabilities.compute_capability
+            not in ["sm_80", "sm_86", "sm_89", "sm_90", "sm_120"]
+        ):
+            raise ValueError(
+                "SageAttention supports compute capability between sm_80 and sm_120. "
+                "Please use a different attention implementation."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_fp8_attention_preflight(self):
+        """fp8 attention requires SM90+ and torch >= 2.11 (torchao >= 0.17 is pinned)."""
+        if self.attn_implementation != "fp8":
+            return self
+
+        if self.capabilities and self.capabilities.compute_capability:
+            cc = self.capabilities.compute_capability
+            # Accept sm_90 (H100/H200), sm_100 (B100/B200), sm_120 (B300-class).
+            if not cc.startswith("sm_") or int(cc.split("_", 1)[1]) < 90:
+                raise ValueError(
+                    f"attn_implementation=fp8 requires compute capability sm_90 or "
+                    f"higher (Hopper+). Detected {cc!r}."
+                )
+
+        torch_version = (
+            self.env_capabilities.torch_version if self.env_capabilities else None
+        )
+        if torch_version is None:
+            import torch
+
+            torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
+        if version.parse(torch_version) < version.parse("2.11.0"):
+            raise ValueError(
+                f"attn_implementation=fp8 requires PyTorch >= 2.11.0. "
+                f"Detected {torch_version}."
+            )
+
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_multigpu_lora_kernels(cls, data):
+        if (
+            data.get("lora_mlp_kernel")
+            or data.get("lora_qkv_kernel")
+            or data.get("lora_o_kernel")
+            or data.get("lora_embedding_kernel")
+        ):
+            capabilities = data.get("capabilities")
+            is_fsdp = data.get("fsdp_config") is not None
+            is_fsdp2 = is_fsdp and str(data.get("fsdp_version")) == "2"
+
+            if capabilities and capabilities.get("n_gpu", 0) > 1 and not is_fsdp2:
+                if is_fsdp:
+                    raise ValueError(
+                        "lora_mlp_kernel, lora_qkv_kernel, and lora_o_kernel are not compatible with FSDP1."
+                    )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_quantize_moe_experts(cls, data):
+        if data.get("quantize_moe_experts"):
+            if data.get("lora_target_linear"):
+                raise ValueError(
+                    "lora_target_linear is not compatible with quantize_moe_experts. "
+                    "Use lora_target_parameters to target expert weights instead."
+                )
+            if data.get("adapter") not in ("lora", "qlora"):
+                raise ValueError("quantize_moe_experts requires adapter: lora or qlora")
+            if not (data.get("load_in_4bit") or data.get("load_in_8bit")):
+                raise ValueError(
+                    "quantize_moe_experts requires load_in_4bit or load_in_8bit"
+                )
+            if (
+                data.get("capabilities")
+                and data["capabilities"].get("compute_capability")
+                and not data["capabilities"]["compute_capability"].startswith("sm_")
+            ):
+                raise ValueError(
+                    "quantize_moe_experts requires CUDA (not compatible with ROCm or other backends)"
+                )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_auto_enable_lora_kernels(cls, data):
+        # Only proceed if using LoRA or QLoRA adapter
+        if data.get("rl"):
+            # RL trainers not tested so don't enable kernels by default
+            return data
+        if data.get("adapter") in ["lora", "qlora"]:
+            # Skip if already set or using 8-bit
+            kernel_fields = [
+                "lora_mlp_kernel",
+                "lora_qkv_kernel",
+                "lora_o_kernel",
+                "lora_embedding_kernel",
+            ]
+            if (
+                any(data.get(k) is not None for k in kernel_fields)
+                or data.get("adapter") == "lora"
+                and data.get("load_in_8bit")
+            ):
+                return data
+
+            # Skip if trust_remote_code is enabled, as lora kernels are not compatible
+            if data.get("trust_remote_code"):
+                return data
+
+            # Skip auto-enable for MoE models when native grouped_mm is unavailable
+            # (torch < 2.9). The grouped_mm fallback in transformers uses torch.mm
+            # with out= which bypasses autocast and fails on mixed dtypes during eval.
+            env_capabilities = data.get("env_capabilities", {})
+            torch_version = env_capabilities.get("torch_version")
+            if torch_version is None:
+                import torch
+
+                torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
+            has_grouped_mm = version.parse(torch_version) >= version.parse("2.9.0")
+            if not has_grouped_mm:
+                is_moe = False
+                model_type = data.get("model_config_type", "")
+                if model_type and "moe" in model_type.lower():
+                    is_moe = True
+                if not is_moe:
+                    try:
+                        from transformers import AutoConfig
+
+                        base_model = data.get("base_model")
+                        if base_model:
+                            auto_cfg = AutoConfig.from_pretrained(
+                                base_model, trust_remote_code=False
+                            )
+                            if getattr(auto_cfg, "num_local_experts", None) or getattr(
+                                auto_cfg, "num_experts", None
+                            ):
+                                is_moe = True
+                    except Exception:  # pylint: disable=broad-exception-caught
+                        pass
+                if is_moe:
+                    return data
+
+            # Check multi-GPU compatibility
+            capabilities = data.get("capabilities")
+            is_multi_gpu = capabilities and capabilities.get("n_gpu", 0) > 1
+            is_fsdp = data.get("fsdp_config") is not None
+            is_fsdp2 = is_fsdp and str(data.get("fsdp_version")) == "2"
+
+            if (
+                not is_multi_gpu
+                or (is_multi_gpu and not is_fsdp)
+                or (is_multi_gpu and is_fsdp2)
+            ):
+                # Auto-enable kernels if not explicitly set by user
+                if data.get("lora_mlp_kernel") is None:
+                    data["lora_mlp_kernel"] = True
+
+                if data.get("lora_qkv_kernel") is None:
+                    data["lora_qkv_kernel"] = True
+
+                if data.get("lora_o_kernel") is None:
+                    data["lora_o_kernel"] = True
+
+                if data.get("lora_embedding_kernel") is None:
+                    data["lora_embedding_kernel"] = True
+
+                LOG.warning(
+                    "Auto-enabling LoRA kernel optimizations for faster training. "
+                    + "Please explicitly set `lora_*_kernel` config values to `false` to disable. "
+                    + "See https://docs.axolotl.ai/docs/lora_optims.html for more info."
+                )
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_adopt_torch_version(cls, data):
+        if (data.get("optimizer") is not None) and ("adopt" in data.get("optimizer")):
+            env_capabilities = data.get("env_capabilities", {})
+            torch_version = env_capabilities.get("torch_version")
+
+            if torch_version is None:
+                import torch
+
+                torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
+
+            if version.parse(torch_version) < version.parse("2.5.1"):
+                raise ValueError(
+                    "ADOPT optimizer is incompatible with torch version < 2.5.1"
+                )
+        return data
+
+    @model_validator(mode="after")
+    def check_flex_torch_version(self):
+        if self.attn_implementation == "flex_attention":
+            torch_version = (
+                self.env_capabilities.torch_version if self.env_capabilities else None
+            )
+            if torch_version is None:
+                import torch
+
+                torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
+
+            if version.parse(torch_version) < version.parse("2.6.0"):
+                raise ValueError(
+                    "Flex attention is not supported on torch version < 2.6.0"
+                )
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_torch_compile_auto(cls, data):
+        if data.get("torch_compile") == "auto":
+            env_capabilities = data.get("env_capabilities", {})
+            if env_capabilities.get("torch_version"):
+                if version.parse(
+                    env_capabilities.get("torch_version")
+                ) >= version.parse("2.5.1"):
+                    LOG.info(
+                        "torch.compile is available, setting torch_compile to True"
+                    )
+                    data["torch_compile"] = True
+                else:
+                    data["torch_compile"] = False
+            else:
+                data["torch_compile"] = False
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_beta_and_trl_beta_match(cls, data):
+        if data.get("beta") and data.get("trl", {}).get("beta"):
+            if data["beta"] != data["trl"]["beta"]:
+                raise ValueError("beta and trl.beta must match or one must be removed")
+        return data
+
+    @model_validator(mode="after")
+    def check_min_torch_version(self):
+        if self.env_capabilities and self.env_capabilities.torch_version:
+            torch_version = self.env_capabilities.torch_version
+            if version.parse(torch_version) < version.parse("2.6.0"):
+                LOG.warning(
+                    f"torch=={torch_version} not be supported. Please upgrade to torch>=2.6.0."
+                )
+
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_qat_config(cls, data):
+        qat_cfg = data.get("qat", {})
+        if not qat_cfg:
+            return data
+
+        if data.get("peft"):
+            raise ValueError("QAT and PEFT cannot be used together.")
+
+        if data.get("load_in_8bit"):
+            raise ValueError("QAT and load_in_8bit cannot be used together.")
+
+        if data.get("load_in_4bit"):
+            raise ValueError("QAT and load_in_4bit cannot be used together.")
+
+        env_capabilities = data.get("env_capabilities", {})
+        torch_version = env_capabilities.get("torch_version")
+
+        if torch_version is None:
+            import torch
+
+            torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
+
+        if version.parse(torch_version) < version.parse("2.6.0"):
+            raise ValueError("QAT is not supported on torch version < 2.6.0")
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fsdp_torch_version(cls, data):
+        env_capabilities = data.get("env_capabilities", {})
+        torch_version = env_capabilities.get("torch_version")
+
+        if torch_version is None:
+            import torch
+
+            torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
+
+        if data.get("fsdp_config") and str(data.get("fsdp_version")) == "2":
+            if version.parse(torch_version) < version.parse("2.7.0"):
+                raise ValueError("FSDP2 is not supported on torch version < 2.7.0")
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def default_dataloader_opts(cls, data):
+        if (
+            data.get("dataloader_num_workers") is None
+            and data.get("dataloader_pin_memory") is None
+            and data.get("dataloader_prefetch_factor") is None
+        ):
+            data["dataloader_num_workers"] = data.get("capabilities").get("n_gpu", 1)
+            data["dataloader_pin_memory"] = True
+            data["dataloader_prefetch_factor"] = 256
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def default_dataset_num_proc(cls, data):
+        if data.get("dataset_processes") is not None:
+            if data.get("dataset_num_proc") is None:
+                data["dataset_num_proc"] = data["dataset_processes"]
+                LOG.warning(
+                    "dataset_processes is deprecated and will be removed in a future version. "
+                    "Please use dataset_num_proc instead."
+                )
+            else:
+                LOG.warning(
+                    "Both dataset_processes and dataset_num_proc are set. "
+                    "Using dataset_num_proc and ignoring dataset_processes."
+                )
+            del data["dataset_processes"]
+        elif data.get("dataset_num_proc") is None:
+            data["dataset_num_proc"] = get_default_process_count()
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def migrate_gc_steps(cls, data):
+        gc_steps = data.get("gc_steps")
+        if gc_steps is not None:
+            if (
+                data.get("torch_empty_cache_steps") is None
+                and data.get("gc_collect_steps") is None
+            ):
+                # Map gc_steps to both new options to preserve original behavior
+                if gc_steps > 0:
+                    data["torch_empty_cache_steps"] = gc_steps
+                data["gc_collect_steps"] = gc_steps
+                LOG.warning(
+                    "`gc_steps` is deprecated; mapping gc_steps=%d to "
+                    "`torch_empty_cache_steps` and `gc_collect_steps`.",
+                    gc_steps,
+                )
+            else:
+                LOG.warning(
+                    "`gc_steps` ignored; `torch_empty_cache_steps` / "
+                    "`gc_collect_steps` take precedence."
+                )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_deduplication_with_streaming(cls, data):
+        if data.get("dataset_exact_deduplication") and (
+            data.get("streaming") or data.get("pretraining_dataset")
+        ):
+            raise NotImplementedError(
+                "dataset_exact_deduplication is not available for streaming datasets. "
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_deduplication_with_skip_prepare(cls, data):
+        if data.get("dataset_exact_deduplication") and data.get("skip_prepare_dataset"):
+            raise ValueError(
+                "dataset_exact_deduplication=True has no effect when "
+                "skip_prepare_dataset=True. Deduplication runs as part of the "
+                "prepare pipeline, which is skipped. Either set "
+                "skip_prepare_dataset: false or disable "
+                "dataset_exact_deduplication."
+            )
+        return data
diff --git a/src/axolotl/utils/schemas/datasets.py b/src/axolotl/utils/schemas/datasets.py
new file mode 100644
index 0000000000..5f2289e450
--- /dev/null
+++ b/src/axolotl/utils/schemas/datasets.py
@@ -0,0 +1,345 @@
+"""Pydantic models for datasets-related configuration"""
+
+from typing import Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+from axolotl.utils.schemas.enums import ChatTemplate
+from axolotl.utils.schemas.utils import handle_legacy_message_fields_logic
+
+
+class UserDefinedPrompterType(BaseModel):
+    """Structure for user defined prompt types"""
+
+    system_prompt: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Custom user instruction prompt"},
+    )
+    system_format: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Use {system} as key to be replaced"},
+    )
+    field_system: str | None = None
+    field_instruction: str | None = None
+    field_input: str | None = None
+    field_output: str | None = None
+
+    format: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Customizable to be single line or multi-line. Use {instruction}/{input} as key to be replaced. 'format' can include {input}"
+        },
+    )
+    no_input_format: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "'no_input_format' cannot include {input}"},
+    )
+
+
+class SFTDataset(BaseModel):
+    """SFT configuration subset"""
+
+    path: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "HuggingFace dataset repo | s3:// | gs:// | path to local file or directory"
+        },
+    )
+    split: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "name of dataset split to load from"},
+    )
+    type: str | UserDefinedPrompterType | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The type of prompt to use for training. [alpaca, gpteacher, oasst, reflection]"
+        },
+    )
+    input_transform: str | None = None
+    shards: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "split dataset into N pieces (use with shards_idx)"
+        },
+    )
+    shards_idx: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "the index of sharded dataset to use"},
+    )
+    preprocess_shards: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "process dataset in N sequential chunks for memory efficiency (exclusive with `shards`)"
+        },
+    )
+    conversation: str | None = None
+    # Do not make this too strict or it will break the validator to choose different dataset class
+    chat_template: ChatTemplate | str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The name of the chat template to use for training, following values are supported: tokenizer_default: Uses the chat template that is available in the tokenizer_config.json. If the chat template is not available in the tokenizer, it will raise an error. This is the default. alpaca/inst/chatml/gemma/cohere/llama3/phi_3/deepseek_v2/jamba: These chat templates are available in the axolotl codebase at src/axolotl/utils/chat_templates.py. tokenizer_default_fallback_*: where * is the name of the chat template to fallback to if the tokenizer does not have a chat template else default to tokenizer. E.g. tokenizer_default_fallback_chatml. jinja: Uses a custom jinja template for the chat template. The custom jinja template should be provided in the chat_template_jinja field."
+        },
+    )
+    chat_template_jinja: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Custom jinja chat template or path to jinja file. Used only if `chat_template: jinja` or empty."
+        },
+    )
+    data_files: str | list[str] | None = Field(
+        default=None, json_schema_extra={"description": "path to source data files"}
+    )
+    input_format: str | None = None
+    name: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "name of dataset configuration to load"},
+    )
+    ds_type: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "defines the datatype when path is a file"},
+    )
+    field: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "For `completion` datasets only, uses the provided field instead of `text` column"
+        },
+    )
+    field_human: str | None = None
+    field_model: str | None = None
+    field_messages: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": 'Key containing the messages (default: "messages")'
+        },
+    )
+    field_tools: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": 'Key containing the tools (default: "tools"). Must be a list[dict] and follow [JSON schema](https://json-schema.org/learn/getting-started-step-by-step).'
+        },
+    )
+    field_thinking: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": 'Key containing the reasoning trace (default: "reasoning_content").'
+        },
+    )
+    template_thinking_key: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The key the chat template expects that indicates the reasoning trace."
+        },
+    )
+    # deprecated, use message_property_mappings
+    message_field_role: str | None = None
+    # deprecated, use message_property_mappings
+    message_field_content: str | None = None
+    message_property_mappings: dict[str, str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Mapping of properties from the input dataset to the chat template. (default: message_property_mappings={'role':'role', 'content':'content'}) If a property exists in the template but not in this mapping, the system will attempt to load it directly from the message using the property name as the key. Example: In the mapping below, 'from' is loaded from input dataset and used as 'role', while 'value' is loaded and used as 'content' in the chat template."
+        },
+    )
+    message_field_training: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The key in the message turn that indicates via boolean whether tokens of a turn should be considered for training. Useful to selectively train on certain turns besides the `roles_to_train`."
+        },
+    )
+    message_field_training_detail: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The key in the message turn that contains the training details. Useful to selectively train on certain tokens in a turn. The value of the key is a List[Dict] containing `begin_offset` (start character index in content), `end_offset` (end character index in content), and `train` (boolean whether to train)."
+        },
+    )
+    split_thinking: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "(for Qwen3 template only) Whether to split the assistant content based on a reasoning trace inside delimited tags"
+        },
+    )
+    logprobs_field: str | None = None
+    temperature: float | None = None
+    roles_to_train: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Roles to train on. The tokens from these roles will be considered for the loss."
+        },
+    )
+    train_on_eos: Literal["all", "turn", "last", "none"] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Which EOS tokens to train on in the conversation. Possible values are: all: train on all EOS tokens, turn (default): train on the EOS token at the end of each trainable turn, last: train on the last EOS token in the conversation, none: never train on EOS tokens"
+        },
+    )
+    train_on_eot: Literal["all", "turn", "last", "none"] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Which EOT (end-of-turn) tokens to train on. Same values as train_on_eos. Defaults to the train_on_eos setting when unset."
+        },
+    )
+    roles: dict[str, list[str]] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": 'Roles mapping in the messages. The format is {target_role: [source_roles]}. All source roles will be mapped to the target role. The default is: user: ["human", "user"], assistant: ["gpt", "assistant"], system: ["system"], tool: ["tool"]'
+        },
+    )
+    drop_system_message: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to drop the system turn from the dataset. Only works with chat_template. This does not drop the default system message from chat_template if it exists. If you wish to, we recommend using a custom jinja template with the default system message removed or adding a system turn with empty content."
+        },
+    )
+    trust_remote_code: bool | None = Field(
+        default=False,
+        json_schema_extra={"description": "Trust remote code for untrusted source"},
+    )
+    revision: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The specific revision of the dataset to use when loading from the Hugging Face Hub. This can be a commit hash, tag, or branch name. If not specified, the latest version will be used. This parameter is ignored for local datasets."
+        },
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def handle_legacy_message_fields(cls, data):
+        """Handle backwards compatibility between legacy message field mapping and new property mapping system."""
+        return handle_legacy_message_fields_logic(data)
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_chat_template_config(cls, data):
+        if isinstance(data, BaseModel):
+            data = data.model_dump()
+
+        # Set chat_template to tokenizer_default if not set
+        if data.get("type") == "chat_template" and not data.get("chat_template"):
+            data["chat_template"] = ChatTemplate.tokenizer_default
+
+        # if chat_template is set to jinja, chat_template_jinja is required
+        if data.get("chat_template") == ChatTemplate.jinja and not data.get(
+            "chat_template_jinja"
+        ):
+            raise ValueError(
+                "chat_template_jinja is required when chat_template is set to jinja"
+            )
+
+        # If chat_template_jinja is set, set chat_template to jinja
+        if data.get("chat_template_jinja") and not data.get("chat_template"):
+            data["chat_template"] = ChatTemplate.jinja
+
+        return data
+
+
+class PretrainingDataset(BaseModel):
+    """Pretraining dataset configuration subset"""
+
+    name: str | None = None
+    path: str | None = None
+    split: str | None = "train"
+    text_column: str | None = "text"
+    type: str | None = "pretrain"
+    trust_remote_code: bool | None = False
+    data_files: str | None = None
+    skip: int | None = None
+
+
+class UserDefinedDPOType(BaseModel):
+    """User defined typing for DPO"""
+
+    field_system: str | None = None
+    field_prompt: str | None = None
+    field_chosen: str | None = None
+    field_rejected: str | None = None
+    prompt_format: str | None = None
+    chosen_format: str | None = None
+    rejected_format: str | None = None
+
+
+class DPODataset(BaseModel):
+    """DPO configuration subset"""
+
+    path: str | None = None
+    split: str | None = None
+    type: UserDefinedDPOType | str | None = None
+    data_files: list[str] | None = None
+    revision: str | None = None
+    field_messages: str | None = None
+    field_chosen: str | None = None
+    field_rejected: str | None = None
+
+
+class StepwiseSupervisedDataset(BaseModel):
+    """Stepwise supervised dataset configuration subset"""
+
+    path: str | None = None
+    split: str | None = None
+    data_files: list[str] | None = None
+    revision: str | None = None
+    step_separator: str | None = None
+    max_completion_length: int | None = None
+    train_on_last_step_only: bool | None = None
+
+
+class UserDefinedKTOType(BaseModel):
+    """User defined typing for KTO"""
+
+    field_system: str | None = None
+    field_prompt: str | None = None
+    field_completion: str | None = None
+    field_label: str | None = None
+    prompt_format: str | None = None
+    completion_format: str | None = None
+
+
+class KTODataset(BaseModel):
+    """KTO configuration subset"""
+
+    path: str | None = None
+    split: str | None = None
+    type: UserDefinedKTOType | str | None = None
+    data_files: list[str] | None = None
+    trust_remote_code: bool | None = False
+    revision: str | None = None
+
+
+class SyntheticDataset(BaseModel):
+    """Synthetic dataset configuration for benchmarking and testing.
+
+    Generates datasets with configurable sequence length, dataset size, and token ID
+    ranges. Useful for benchmarking memory usage and speed by sequence length, and for
+    validating weighted dataset mixes.
+    """
+
+    path: Literal["synthetic"] = "synthetic"
+    type: Literal["_synthetic"] = "_synthetic"
+    length: int = Field(
+        default=1000,
+        json_schema_extra={"description": "Number of rows to generate"},
+    )
+    sequence_length: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Sequence length per row (defaults to sequence_len from config)"
+        },
+    )
+    min_input_id: int = Field(
+        default=100,
+        json_schema_extra={"description": "Minimum token ID for generation"},
+    )
+    max_input_id: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Maximum token ID for generation (defaults to tokenizer vocab_size)"
+        },
+    )
+    seed: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Random seed for reproducibility"},
+    )
+
+
+DatasetConfig = (
+    SFTDataset | DPODataset | KTODataset | StepwiseSupervisedDataset | SyntheticDataset
+)
diff --git a/src/axolotl/utils/schemas/deprecated.py b/src/axolotl/utils/schemas/deprecated.py
new file mode 100644
index 0000000000..a441157f43
--- /dev/null
+++ b/src/axolotl/utils/schemas/deprecated.py
@@ -0,0 +1,174 @@
+"""Pydantic models for deprecated and remapped configuration parameters"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class DeprecatedParameters(BaseModel):
+    """configurations that are deprecated"""
+
+    max_packed_sequence_len: int | None = None
+    rope_scaling: Any | None = None
+    noisy_embedding_alpha: float | None = None
+    dpo_beta: float | None = None
+    evaluation_strategy: str | None = None
+    eval_table_size: int | None = None
+    eval_max_new_tokens: int | None = None
+    dpo_use_logits_to_keep: bool | None = None
+    dpo_generate_during_eval: bool | None = None
+    dpo_norm_loss: bool | None = None
+    rpo_alpha: float | None = None
+    s2_attention: bool | None = None
+    flash_attn_rms_norm: bool | None = None
+
+    @field_validator("max_packed_sequence_len")
+    @classmethod
+    def validate_max_packed_sequence_len(cls, max_packed_sequence_len):
+        if max_packed_sequence_len:
+            raise DeprecationWarning("`max_packed_sequence_len` is no longer supported")
+        return max_packed_sequence_len
+
+    @field_validator("rope_scaling")
+    @classmethod
+    def validate_rope_scaling(cls, rope_scaling):
+        if rope_scaling:
+            raise DeprecationWarning(
+                "`rope_scaling` is no longer supported, it should now be be a key under `model_config`"
+            )
+        return rope_scaling
+
+    @field_validator("noisy_embedding_alpha")
+    @classmethod
+    def validate_noisy_embedding_alpha(cls, noisy_embedding_alpha):
+        if noisy_embedding_alpha:
+            LOG.warning("noisy_embedding_alpha is deprecated, use neftune_noise_alpha")
+        return noisy_embedding_alpha
+
+    @field_validator("dpo_beta")
+    @classmethod
+    def validate_dpo_beta(cls, dpo_beta):
+        if dpo_beta is not None:
+            LOG.warning("dpo_beta is deprecated, use rl_beta instead")
+        return dpo_beta
+
+    @field_validator("evaluation_strategy")
+    @classmethod
+    def validate_evaluation_strategy(cls, evaluation_strategy):
+        if evaluation_strategy is not None:
+            LOG.warning("evaluation_strategy is deprecated, use eval_strategy instead")
+        return evaluation_strategy
+
+    @field_validator("eval_table_size")
+    @classmethod
+    def validate_eval_table_size(cls, eval_table_size):
+        if eval_table_size is not None:
+            LOG.warning(
+                "eval_table_size is deprecated and superseded by generate_samples config. "
+                "Please use generate_samples: true and num_generation_samples instead. "
+                "The LogPredictionCallback is replaced by the new sample generation feature."
+            )
+        return eval_table_size
+
+    @field_validator("eval_max_new_tokens")
+    @classmethod
+    def validate_eval_max_new_tokens(cls, eval_max_new_tokens):
+        if eval_max_new_tokens is not None:
+            LOG.warning(
+                "eval_max_new_tokens is deprecated and superseded by generate_samples config. "
+                "Please use generation_max_new_tokens instead."
+            )
+        return eval_max_new_tokens
+
+    @field_validator("dpo_use_logits_to_keep")
+    @classmethod
+    def validate_dpo_use_logits_to_keep(cls, dpo_use_logits_to_keep):
+        if dpo_use_logits_to_keep is not None:
+            raise DeprecationWarning(
+                "`dpo_use_logits_to_keep` is no longer supported, "
+                "it has been removed in TRL >= 0.29.0"
+            )
+        return dpo_use_logits_to_keep
+
+    @field_validator("dpo_generate_during_eval")
+    @classmethod
+    def validate_dpo_generate_during_eval(cls, dpo_generate_during_eval):
+        if dpo_generate_during_eval is not None:
+            raise DeprecationWarning(
+                "`dpo_generate_during_eval` is no longer supported, "
+                "it has been removed in TRL >= 0.29.0"
+            )
+        return dpo_generate_during_eval
+
+    @field_validator("dpo_norm_loss")
+    @classmethod
+    def validate_dpo_norm_loss(cls, dpo_norm_loss):
+        if dpo_norm_loss is not None:
+            raise DeprecationWarning(
+                "`dpo_norm_loss` is no longer supported, "
+                "due to breaking changes in TRL >= 0.29.0"
+            )
+        return dpo_norm_loss
+
+    @field_validator("rpo_alpha")
+    @classmethod
+    def validate_rpo_alpha(cls, rpo_alpha):
+        if rpo_alpha is not None:
+            raise DeprecationWarning(
+                "`rpo_alpha` has been deprecated in TRL >= 0.29.0, "
+                "and now requires passing multiple loss types, which is not yet supported by Axolotl."
+            )  # TODO: change this warning once multiple dpo loss types are supported.
+        return rpo_alpha
+
+    @field_validator("s2_attention")
+    @classmethod
+    def validate_s2_attention(cls, s2_attention):
+        if s2_attention:
+            raise DeprecationWarning(
+                "`s2_attention` (shifted-sparse attention) is no longer supported."
+            )
+        return s2_attention
+
+    @field_validator("flash_attn_rms_norm")
+    @classmethod
+    def validate_flash_attn_rms_norm(cls, flash_attn_rms_norm):
+        if flash_attn_rms_norm:
+            raise DeprecationWarning("`flash_attn_rms_norm` is no longer supported.")
+        return flash_attn_rms_norm
+
+
+class RemappedParameters(BaseModel):
+    """Parameters that have been remapped to other names"""
+
+    overrides_of_model_config: dict[str, Any] | None = Field(
+        default=None,
+        alias="model_config",
+        json_schema_extra={
+            "description": "optional overrides to the base model configuration"
+        },
+    )
+    overrides_of_model_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        alias="model_kwargs",
+        json_schema_extra={
+            "description": "optional overrides the base model loading from_pretrained"
+        },
+    )
+    type_of_model: str | None = Field(
+        default=None,
+        alias="model_type",
+        json_schema_extra={
+            "description": "If you want to specify the type of model to load, AutoModelForCausalLM is a good choice too"
+        },
+    )
+    revision_of_model: str | None = Field(
+        default=None,
+        alias="model_revision",
+        json_schema_extra={
+            "description": "You can specify to choose a specific model revision from huggingface hub"
+        },
+    )
diff --git a/src/axolotl/utils/schemas/dynamic_checkpoint.py b/src/axolotl/utils/schemas/dynamic_checkpoint.py
new file mode 100644
index 0000000000..e0e1d0c1d6
--- /dev/null
+++ b/src/axolotl/utils/schemas/dynamic_checkpoint.py
@@ -0,0 +1,31 @@
+"""Schema for dynamic checkpoint configuration."""
+
+from pydantic import BaseModel, Field
+
+
+class DynamicCheckpointConfig(BaseModel):
+    """Configuration for dynamic checkpoint triggering during training."""
+
+    enabled: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Enable dynamic checkpoint triggering during training. "
+            "Create a file 'axolotl_checkpoint.save' in the configured `output_dir` to trigger. "
+        },
+    )
+    check_interval: int = Field(
+        default=10,
+        ge=1,
+        json_schema_extra={
+            "description": "Check for trigger file every N steps (reduces I/O overhead). "
+            "Default: 100"
+        },
+    )
+    trigger_file_path: str = Field(
+        default="",
+        json_schema_extra={
+            "description": "Custom trigger filename (optional). "
+            "If not specified, defaults to 'axolotl_checkpoint.save'. "
+            "Specify a filename (not a full path) to override the default."
+        },
+    )
diff --git a/src/axolotl/utils/schemas/enums.py b/src/axolotl/utils/schemas/enums.py
new file mode 100644
index 0000000000..4ab2127625
--- /dev/null
+++ b/src/axolotl/utils/schemas/enums.py
@@ -0,0 +1,170 @@
+"""Enums for Axolotl input config"""
+
+from enum import Enum
+
+import torch
+
+
+class TorchAOQuantDType(Enum):
+    int4 = torch.int4
+    int8 = torch.int8
+    float8_e4m3fn = torch.float8_e4m3fn
+    nvfp4 = "nvfp4"
+    mxfp4 = "mxfp4"
+
+    def from_string(str):
+        if str == "int4":
+            return TorchAOQuantDType.int4
+        if str == "int8":
+            return TorchAOQuantDType.int8
+        if str in ["float8_e4m3fn", "fp8", "float8"]:
+            return TorchAOQuantDType.float8_e4m3fn
+        if str == "nvfp4":
+            return TorchAOQuantDType.nvfp4
+        if str == "mxfp4":
+            return TorchAOQuantDType.mxfp4
+
+
+class RLType(str, Enum):
+    """RL trainer type configuration subset"""
+
+    DPO = "dpo"
+    GDPO = "gdpo"
+    GRPO = "grpo"
+    IPO = "ipo"
+    ORPO = "orpo"
+    KTO = "kto"
+    SIMPO = "simpo"
+    EBFT = "ebft"
+
+
+class ChatTemplate(str, Enum):
+    """Chat templates configuration subset"""
+
+    alpaca = "alpaca"
+    chatml = "chatml"
+    mistral_v1 = "mistral_v1"
+    mistral_v2v3 = "mistral_v2v3"
+    mistral_v3_tekken = "mistral_v3_tekken"
+    mistral_v7_tekken = "mistral_v7_tekken"
+    gemma = "gemma"
+    cohere = "cohere"
+    llama3 = "llama3"
+    llama3_2_vision = "llama3_2_vision"
+    llama4 = "llama4"
+    phi_3 = "phi_3"
+    phi_35 = "phi_35"
+    deepseek_v2 = "deepseek_v2"
+    deepseek_v3 = "deepseek_v3"
+    jamba = "jamba"
+    jinja = "jinja"
+    qwen_25 = "qwen_25"
+    qwen3 = "qwen3"
+    qwen3_5 = "qwen3_5"
+    falcon_h1 = "falcon_h1"
+    nemotron_h = "nemotron_h"
+    tokenizer_default = "tokenizer_default"
+    exaone = "exaone"
+    exaone4 = "exaone4"
+    metharme = "metharme"
+    pixtral = "pixtral"
+    llava = "llava"
+    qwen2_vl = "qwen2_vl"
+    gemma3 = "gemma3"
+    gemma3n = "gemma3n"
+    gemma4 = "gemma4"
+    gemma4_unified = "gemma4_unified"
+    command_a = "command_a"
+    command_a_tool_use = "command_a_tool_use"
+    command_a_rag = "command_a_rag"
+    aya = "aya"
+
+
+class CustomSupportedOptimizers(str, Enum):
+    """Custom supported optimizers"""
+
+    optimi_adamw = "optimi_adamw"
+    ao_adamw_4bit = "ao_adamw_4bit"
+    ao_adamw_8bit = "ao_adamw_8bit"
+    ao_adamw_fp8 = "ao_adamw_fp8"
+    adopt_adamw = "adopt_adamw"
+    came_pytorch = "came_pytorch"
+    muon = "muon"
+    dion = "dion"
+    sinkgd = "sinkgd"
+    flash_adamw = "flash_adamw"
+    flash_adam = "flash_adam"
+    flash_sgd = "flash_sgd"
+    flash_sgdw = "flash_sgdw"
+    flash_lion = "flash_lion"
+    q_galore_adamw8bit = "q_galore_adamw8bit"
+
+
+# Accepted canonical names; hub-kernel paths (containing "/") bypass this set.
+CANONICAL_ATTN_IMPLS = frozenset(
+    {
+        "eager",
+        "sdpa",
+        "flash_attention_2",
+        "flash_attention_3",
+        "flex_attention",
+        "xformers",
+        "sage",
+        "fp8",
+    }
+)
+
+# Legacy boolean flags → canonical attn_implementation. Priority: specific before generic.
+LEGACY_ATTN_FLAG_TO_IMPL = {
+    "xformers_attention": "xformers",
+    "sage_attention": "sage",
+    "flex_attention": "flex_attention",
+    "flash_attention": "flash_attention_2",
+    "sdp_attention": "sdpa",
+    "eager_attention": "eager",
+}
+
+# Short-form aliases rejected at validation; mapped to canonical names for error messages.
+SHORT_FORM_ALIAS_TO_CANONICAL = {
+    "flash": "flash_attention_2",
+    "flex": "flex_attention",
+    "sdp": "sdpa",
+}
+
+# Backends that support varlen sample packing via `position_ids`.
+ATTN_IMPLS_SUPPORTING_PACKING = frozenset(
+    {
+        "flash_attention_2",
+        "flash_attention_3",
+        "flex_attention",
+        "xformers",
+        "sage",
+        "kernels-community/flash-attn2",
+        "kernels-community/flash-attn3",
+        "kernels-community/sage-attention",
+    }
+)
+
+# Backends that require the flash_attn library for axolotl's own monkeypatches.
+ATTN_IMPLS_USING_FLASH_LIB = frozenset(
+    {
+        "flash_attention_2",
+        "flash_attention_3",
+        "kernels-community/flash-attn2",
+        "kernels-community/flash-attn3",
+    }
+)
+
+# Backends for which embeddings stay in fp32. Everything else needs fp16/bf16.
+ATTN_IMPLS_WITHOUT_DTYPE_CAST = frozenset({"eager", "sdpa"})
+
+
+class RingAttnFunc(str, Enum):
+    """Enum class for supported `ring-flash-attn` implementations"""
+
+    VARLEN_LLAMA3 = "varlen_llama3"
+    BATCH_RING = "batch_ring"
+    # VARLEN_RING = "varlen_ring"
+    # VARLEN_ZIGZAG = "varlen_zigzag"
+    # BATCH_ZIGZAG = "batch_zigzag"
+    # BATCH_STRIPE = "batch_stripe"
diff --git a/src/axolotl/utils/schemas/fsdp.py b/src/axolotl/utils/schemas/fsdp.py
new file mode 100644
index 0000000000..1263a84181
--- /dev/null
+++ b/src/axolotl/utils/schemas/fsdp.py
@@ -0,0 +1,90 @@
+"""
+FSDP Configuration Schema
+"""
+
+from typing import Literal
+
+from pydantic import AliasChoices, BaseModel, Field, model_validator
+
+
+class FSDPConfig(BaseModel):
+    """
+    FSDP Configuration Schema
+    """
+
+    fsdp_version: int | None = Field(
+        validation_alias=AliasChoices("fsdp_version", "version"),
+        default=None,
+        json_schema_extra={"description": "FSDP version"},
+    )
+    activation_checkpointing: bool | None = Field(
+        default=None,
+        description="Enable activation checkpointing to reduce memory usage during forward passes",
+    )
+    offload_params: bool | None = Field(
+        default=None,
+        description="Offload parameters to CPU to reduce GPU memory usage",
+    )
+    sync_module_states: bool | None = Field(
+        default=None,
+        description="Synchronize module states across all processes",
+    )
+    cpu_ram_efficient_loading: bool | None = Field(
+        default=None,
+        description="Enable CPU RAM efficient loading to reduce memory usage during model loading",
+    )
+    cpu_offload_pin_memory: bool | None = Field(
+        default=None,
+        description="Disabling this enables swap memory usage for resource-constrained setups when offload_params is enabled.",
+    )
+    use_orig_params: bool | None = Field(
+        default=None,
+        description="Use original parameters instead of flattened parameters",
+    )
+
+    state_dict_type: (
+        Literal["FULL_STATE_DICT", "LOCAL_STATE_DICT", "SHARDED_STATE_DICT"] | None
+    ) = Field(
+        default=None,
+        description="Type of state dict to use for saving/loading checkpoints",
+    )
+    final_state_dict_type: (
+        Literal["FULL_STATE_DICT", "LOCAL_STATE_DICT", "SHARDED_STATE_DICT"] | None
+    ) = Field(
+        default=None,
+        description="Final state dict type to use after training completion",
+    )
+
+    auto_wrap_policy: Literal["TRANSFORMER_BASED_WRAP", "SIZE_BASED_WRAP"] | None = (
+        Field(
+            default=None,
+            description="Policy for automatically wrapping modules with FSDP",
+        )
+    )
+    transformer_layer_cls_to_wrap: str | None = Field(
+        default=None,
+        description="Class name of transformer layers to wrap (e.g., 'LlamaDecoderLayer')",
+    )
+
+    min_num_params: int | None = Field(
+        default=None,
+        ge=1,
+        description="Minimum parameter count a module must have to be wrapped under the SIZE_BASED_WRAP policy",
+    )
+
+    reshard_after_forward: bool | None = Field(
+        default=None,
+        description="Reshard parameters after forward pass to save memory",
+    )
+    mixed_precision_policy: str | None = Field(
+        default=None,
+        description="Mixed precision policy for FSDP (e.g., 'fp16', 'bf16')",
+    )
+
+    @model_validator(mode="after")
+    def validate_size_based_wrap(self):
+        if self.auto_wrap_policy == "SIZE_BASED_WRAP" and self.min_num_params is None:
+            raise ValueError(
+                "fsdp_config.min_num_params is required when auto_wrap_policy is SIZE_BASED_WRAP"
+            )
+        return self
diff --git a/src/axolotl/utils/schemas/integrations.py b/src/axolotl/utils/schemas/integrations.py
new file mode 100644
index 0000000000..dc171c3101
--- /dev/null
+++ b/src/axolotl/utils/schemas/integrations.py
@@ -0,0 +1,222 @@
+"""Pydantic models for Axolotl integrations"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field, model_validator
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class MLFlowConfig(BaseModel):
+    """MLFlow configuration subset"""
+
+    use_mlflow: bool | None = None
+    mlflow_tracking_uri: str | None = Field(
+        default=None, json_schema_extra={"description": "URI to mlflow"}
+    )
+    mlflow_experiment_name: str | None = Field(
+        default=None, json_schema_extra={"description": "Your experiment name"}
+    )
+    mlflow_run_name: str | None = Field(
+        default=None, json_schema_extra={"description": "Your run name"}
+    )
+    hf_mlflow_log_artifacts: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "set to true to copy each saved checkpoint on each save to mlflow artifact registry"
+        },
+    )
+
+
+class LISAConfig(BaseModel):
+    """LISA configuration subset"""
+
+    lisa_n_layers: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "the number of activate layers in LISA"},
+    )
+    lisa_step_interval: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "how often to switch layers in LISA"},
+    )
+    lisa_layers_attribute: str | None = Field(
+        default="model.layers",
+        json_schema_extra={"description": "path under the model to access the layers"},
+    )
+
+
+class WandbConfig(BaseModel):
+    """Wandb configuration subset"""
+
+    use_wandb: bool | None = None
+    wandb_name: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Set the name of your wandb run"},
+    )
+    wandb_run_id: str | None = Field(
+        default=None, json_schema_extra={"description": "Set the ID of your wandb run"}
+    )
+    wandb_mode: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": '"offline" to save run metadata locally and not sync to the server, "disabled" to turn off wandb'
+        },
+    )
+    wandb_project: str | None = Field(
+        default=None, json_schema_extra={"description": "Your wandb project name"}
+    )
+    wandb_entity: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "A wandb Team name if using a Team"},
+    )
+    wandb_watch: str | None = None
+    wandb_log_model: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": '"checkpoint" to log model to wandb Artifacts every `save_steps` or "end" to log only at the end of training'
+        },
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_wandb_run(cls, data):
+        if data.get("wandb_run_id") and not data.get("wandb_name"):
+            data["wandb_name"] = data.get("wandb_run_id")
+
+            LOG.warning(
+                "wandb_run_id sets the ID of the run. If you would like to set the name, please use wandb_name instead."
+            )
+
+        return data
+
+
+class CometConfig(BaseModel):
+    """Comet configuration subset"""
+
+    use_comet: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Enable or disable Comet integration."},
+    )
+    comet_api_key: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "API key for Comet. Recommended to set via `comet login`."
+        },
+    )
+    comet_workspace: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Workspace name in Comet. Defaults to the user's default workspace."
+        },
+    )
+    comet_project_name: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Project name in Comet. Defaults to Uncategorized."
+        },
+    )
+    comet_experiment_key: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Identifier for the experiment. Used to append data to an existing experiment or control the key of new experiments. Default to a random key."
+        },
+    )
+    comet_mode: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": 'Create a new experiment ("create") or log to an existing one ("get"). Default ("get_or_create") auto-selects based on configuration.'
+        },
+    )
+    comet_online: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Set to True to log data to Comet server, or False for offline storage. Default is True."
+        },
+    )
+    comet_experiment_config: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Dictionary for additional configuration settings, see the doc for more details."
+        },
+    )
+
+
+class GradioConfig(BaseModel):
+    """Gradio configuration subset"""
+
+    gradio_title: str | None = None
+    gradio_share: bool | None = None
+    gradio_server_name: str | None = None
+    gradio_server_port: int | None = None
+    gradio_max_new_tokens: int | None = None
+    gradio_temperature: float | None = None
+
+
+class RayConfig(BaseModel):
+    """Ray launcher configuration subset"""
+
+    use_ray: bool = Field(default=False)
+    ray_run_name: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "help": "The training results will be saved at `saves/ray_run_name`."
+        },
+    )
+    ray_num_workers: int = Field(
+        default=1,
+        json_schema_extra={
+            "help": "The number of workers for Ray training. Default is 1 worker."
+        },
+    )
+    resources_per_worker: dict = Field(
+        default_factory=lambda: {"GPU": 1},
+        json_schema_extra={
+            "help": "The resources per worker for Ray training. Default is to use 1 GPU per worker."
+        },
+    )
+
+
+class OpenTelemetryConfig(BaseModel):
+    """OpenTelemetry configuration subset"""
+
+    use_otel_metrics: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Enable OpenTelemetry metrics collection and Prometheus export"
+        },
+    )
+    otel_metrics_host: str | None = Field(
+        default="localhost",
+        json_schema_extra={
+            "title": "OpenTelemetry Metrics Host",
+            "description": "Host to bind the OpenTelemetry metrics server to",
+        },
+    )
+    otel_metrics_port: int | None = Field(
+        default=8000,
+        json_schema_extra={
+            "description": "Port for the Prometheus metrics HTTP server"
+        },
+    )
+
+
+class TrackioConfig(BaseModel):
+    """Trackio configuration subset"""
+
+    use_trackio: bool | None = None
+    trackio_project_name: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Your trackio project name"},
+    )
+    trackio_run_name: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Set the name of your trackio run"},
+    )
+    trackio_space_id: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Hugging Face Space ID to sync dashboard to (optional, runs locally if not provided)"
+        },
+    )
diff --git a/src/axolotl/utils/config/models/internals/__init__.py b/src/axolotl/utils/schemas/internal/__init__.py
similarity index 67%
rename from src/axolotl/utils/config/models/internals/__init__.py
rename to src/axolotl/utils/schemas/internal/__init__.py
index dd742caf45..78cc636db9 100644
--- a/src/axolotl/utils/config/models/internals/__init__.py
+++ b/src/axolotl/utils/schemas/internal/__init__.py
@@ -1,4 +1,5 @@
 """module for gpu capabilities"""
+
 from typing import Optional
 
 from pydantic import BaseModel, Field
@@ -9,6 +10,13 @@ class GPUCapabilities(BaseModel):
 
     bf16: bool = Field(default=False)
     fp8: bool = Field(default=False)
+    tf32: bool = Field(default=False)
     n_gpu: int = Field(default=1)
     n_node: int = Field(default=1)
     compute_capability: Optional[str] = Field(default=None)
+
+
+class EnvCapabilities(BaseModel):
+    """model to manage the environment capabilities statically"""
+
+    torch_version: Optional[str] = Field(default=None)
diff --git a/src/axolotl/utils/schemas/model.py b/src/axolotl/utils/schemas/model.py
new file mode 100644
index 0000000000..30202efe06
--- /dev/null
+++ b/src/axolotl/utils/schemas/model.py
@@ -0,0 +1,186 @@
+"""Pydantic models for model input / output, etc. configuration"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, field_validator
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class ModelInputConfig(BaseModel):
+    """Model configuration subset"""
+
+    model_config = {"protected_namespaces": ()}
+
+    base_model: str = Field(
+        json_schema_extra={
+            "description": "This is the huggingface model that contains *.pt, *.safetensors, or *.bin files. This can also be a relative path to a model on disk"
+        }
+    )
+    base_model_config: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "If the base_model repo on hf hub doesn't include configuration .json files, You can set that here, or leave this empty to default to base_model"
+        },
+    )
+    cls_model_config: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "transformers config class (e.g., 'LlamaConfig', 'MistralConfig'). Defaults to AutoConfig."
+        },
+    )
+    tokenizer_config: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Optional tokenizer configuration path in case you want to use a different tokenizer than the one defined in the base model"
+        },
+    )
+    tokenizer_use_fast: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "use_fast option for tokenizer loading from_pretrained, default to True"
+        },
+    )
+    tokenizer_legacy: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use the legacy tokenizer setting, defaults to True"
+        },
+    )
+    tokenizer_use_mistral_common: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use mistral-common tokenizer. If set to True, it will use the mistral-common tokenizer."
+        },
+    )
+    tokenizer_type: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Corresponding tokenizer for the model AutoTokenizer is a good choice"
+        },
+    )
+    processor_type: str | None = Field(
+        default=None, json_schema_extra={"description": "transformers processor class"}
+    )
+    processor_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "kwargs forwarded to the processor's from_pretrained(), overriding processor config (e.g. image_seq_length, min_pixels, etc.)."
+        },
+    )
+    tokenizer_save_jinja_files: bool | None = Field(
+        default=True,  # match the default behavior from transformers
+        json_schema_extra={
+            "description": "Whether to save jinja files for tokenizer, transformers default is True"
+        },
+    )
+    trust_remote_code: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Trust remote code for untrusted source"},
+    )
+
+    experimental_skip_move_to_device: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Don't move the model to the device before sharding. Set to `false` to revert to legacy behavior."
+        },
+    )
+
+    use_kernels: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Use custom kernels, e.g. MegaBlocks."},
+    )
+
+    model_quantization_config: Literal["Mxfp4Config", "FineGrainedFP8Config"] | None = (
+        Field(
+            default=None,
+            json_schema_extra={"description": "Model loading quantization config"},
+        )
+    )
+    model_quantization_config_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={"description": "kwargs for model quantization config"},
+    )
+    use_onebitllms: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use `onebitllms` for 1.58bit training (only for bitnet models)."
+        },
+    )
+
+    @field_validator("trust_remote_code")
+    @classmethod
+    def hint_trust_remote_code(cls, trust_remote_code):
+        if trust_remote_code:
+            LOG.warning(
+                "`trust_remote_code` is set to true. Please make sure that you reviewed the remote code/model."
+            )
+        return trust_remote_code
+
+    @field_validator("processor_kwargs")
+    @classmethod
+    def reject_reserved_processor_kwargs(cls, processor_kwargs):
+        if not processor_kwargs:
+            return processor_kwargs
+        reserved = {"revision", "trust_remote_code"}
+        conflicts = reserved.intersection(processor_kwargs)
+        if conflicts:
+            raise ValueError(
+                "Do not set reserved keys "
+                f"{sorted(conflicts)} inside `processor_kwargs`; "
+                "use the top-level `revision_of_model` / `trust_remote_code` "
+                "config keys instead."
+            )
+        return processor_kwargs
+
+
+class ModelOutputConfig(BaseModel):
+    """model save configuration subset"""
+
+    output_dir: str = Field(
+        default="./model-out",
+        json_schema_extra={"description": "Where to save the full-finetuned model to"},
+    )
+    hub_model_id: str | None = Field(
+        default=None, json_schema_extra={"description": "push checkpoints to hub"}
+    )
+    hub_strategy: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "how to push checkpoints to hub"},
+    )
+    hub_revision: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "branch/revision to push to on hub (default: main)"
+        },
+    )
+    save_safetensors: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Whether to save the model using safetensors format. Defaults to True."
+        },
+    )
+
+    @field_validator("save_safetensors")
+    @classmethod
+    def validate_save_safetensors(cls, v):
+        if v is False:
+            raise ValueError(
+                "save_safetensors=False is not supported in Transformers V5. "
+                "Transformers V5 always uses safetensors format for model serialization. "
+                "This field is deprecated and will be removed in a future version."
+            )
+        # Allow None and True, will default to True if None
+        return True if v is None else v
+
+
+class SpecialTokensConfig(BaseModel):
+    """Special tokens configuration subset"""
+
+    bos_token: str | None = None
+    eos_token: str | None = None
+    pad_token: str | None = None
+    unk_token: str | None = None
+    additional_special_tokens: list[str] | None = None
diff --git a/src/axolotl/utils/schemas/multimodal.py b/src/axolotl/utils/schemas/multimodal.py
new file mode 100644
index 0000000000..01ad5e5a3d
--- /dev/null
+++ b/src/axolotl/utils/schemas/multimodal.py
@@ -0,0 +1,110 @@
+"""Pydantic models for multimodal-related configuration"""
+
+from typing import Literal
+
+from PIL.Image import Resampling
+from pydantic import BaseModel, Field, field_validator
+
+
+class RoleBoundarySpec(BaseModel):
+    """One ``cfg.role_boundaries`` row; see docs/multimodal_assistant_mask.md."""
+
+    role: str = Field(
+        json_schema_extra={
+            "description": (
+                "Role name as it appears in cfg.roles_to_train (e.g. "
+                "'assistant', 'user', 'system', 'tool', 'ipython')."
+            )
+        },
+    )
+    start: str = Field(
+        json_schema_extra={
+            "description": (
+                "Literal string that marks the start of this role's span in "
+                "the rendered chat template. Tokenized via "
+                "``tokenizer.encode(..., add_special_tokens=False)`` at "
+                "strategy init."
+            )
+        },
+    )
+    end: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Literal string that marks the end of this role's span. "
+                "Set to ``eos_token`` to terminate at the tokenizer's EOS. "
+                "Leave unset / null to terminate at end-of-sequence."
+            )
+        },
+    )
+    include_start: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": (
+                "Whether the start marker tokens contribute to loss on "
+                "trainable turns. Default False."
+            )
+        },
+    )
+    include_end: bool = Field(
+        default=True,
+        json_schema_extra={
+            "description": (
+                "Whether the end marker tokens contribute to loss on "
+                "trainable turns (honoring cfg.train_on_eos). Default True."
+            )
+        },
+    )
+
+
+class MultiModalConfig(BaseModel):
+    """Multi-modal configuration subset"""
+
+    image_size: int | tuple[int, int] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "The size of the image to resize to. It can be an integer (resized into padded-square image) or a tuple (width, height)."
+                "If not provided, we will attempt to load from preprocessor.size, otherwise, images won't be resized."
+            )
+        },
+    )
+    image_resize_algorithm: (
+        Literal["bilinear", "bicubic", "lanczos"] | Resampling | None
+    ) = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The resampling algorithm to use for image resizing. Default is bilinear. Please refer to PIL.Image.Resampling for more details."
+        },
+    )
+    role_boundaries: list[RoleBoundarySpec] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Opt-in override for the MM mask scanner's per-role boundary "
+                "markers. Non-empty list replaces built-ins wholesale; unset "
+                "or empty falls back to built-ins. See "
+                "docs/multimodal_assistant_mask.md."
+            )
+        },
+    )
+
+    @field_validator("image_resize_algorithm", mode="before")
+    @classmethod
+    def convert_image_resize_algorithm(cls, image_resize_algorithm):
+        """
+        Convert the image resize algorithm to a PIL.Image.Resampling enum.
+        """
+        if isinstance(image_resize_algorithm, str):
+            image_resize_algorithm = image_resize_algorithm.lower()
+            if image_resize_algorithm == "bilinear":
+                image_resize_algorithm = Resampling.BILINEAR
+            elif image_resize_algorithm == "bicubic":
+                image_resize_algorithm = Resampling.BICUBIC
+            elif image_resize_algorithm == "lanczos":
+                image_resize_algorithm = Resampling.LANCZOS
+            else:
+                raise ValueError(
+                    f"Invalid image resize algorithm: {image_resize_algorithm}"
+                )
+        return image_resize_algorithm
diff --git a/src/axolotl/utils/schemas/peft.py b/src/axolotl/utils/schemas/peft.py
new file mode 100644
index 0000000000..c1d2f98449
--- /dev/null
+++ b/src/axolotl/utils/schemas/peft.py
@@ -0,0 +1,294 @@
+"""Pydantic models for PEFT-related configuration"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, PositiveInt, field_validator, model_validator
+
+
+class LoftQConfig(BaseModel):
+    """LoftQ configuration subset"""
+
+    loftq_bits: int = Field(
+        default=4, json_schema_extra={"description": "typically 4 bits"}
+    )
+    # loftq_iter: int = Field(default=1, json_schema_extra={"description": "Alternating iterations for LoftQ"})
+
+
+class PeftConfig(BaseModel):
+    """peftq configuration subset"""
+
+    loftq_config: LoftQConfig | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Configuration options for loftq initialization for LoRA"
+        },
+    )
+
+
+class LoraConfig(BaseModel):
+    """Peft / LoRA configuration subset"""
+
+    load_in_8bit: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "This will attempt to quantize the model down to 8 bits and use adam 8 bit optimizer"
+        },
+    )
+    load_in_4bit: bool | None = Field(
+        default=False, json_schema_extra={"description": "Use bitsandbytes 4 bit"}
+    )
+
+    adapter: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "If you want to use a built-in or plugin adapter, or leave blank to train all parameters in original model"
+        },
+    )
+    lora_model_dir: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "If you already have a lora model trained that you want to load, put that here. This means after training, if you want to test the model, you should set this to the value of `output_dir`. Note that if you merge an adapter to the base model, a new subdirectory `merged` will be created under the `output_dir`."
+        },
+    )
+    lora_r: int | None = None
+    lora_alpha: int | None = None
+    lora_fan_in_fan_out: bool | None = None
+    lora_target_modules: str | list[str] | None = None
+    lora_target_parameters: str | list[str] | None = None
+    lora_target_linear: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "If true, will target all linear modules"},
+    )
+    lora_modules_to_save: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "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. `embed_tokens` converts tokens to embeddings, and `lm_head` converts embeddings to token probabilities."
+        },
+    )
+    lora_rank_pattern: dict[str, PositiveInt] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Per-module LoRA rank overrides (regex → int > 0). Forwarded to PEFT LoraConfig.rank_pattern."
+        },
+    )
+    lora_alpha_pattern: dict[str, PositiveInt] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Per-module LoRA alpha overrides (regex → int > 0). Forwarded to PEFT LoraConfig.alpha_pattern."
+        },
+    )
+    lora_dropout: float | None = 0.0
+    peft_layers_to_transform: list[int] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The layer indices to transform, otherwise, apply to all layers"
+        },
+    )
+    peft_layers_pattern: list[str] | None = None
+    peft: PeftConfig | None = None
+    peft_use_dora: bool | None = Field(
+        default=None, json_schema_extra={"description": "Whether to use DoRA."}
+    )
+    peft_use_rslora: bool | None = Field(
+        default=None, json_schema_extra={"description": "Whether to use RSLoRA."}
+    )
+    peft_layer_replication: list[tuple[int, int]] | None = Field(
+        default=None,
+        json_schema_extra={"description": "List of layer indices to replicate."},
+    )
+    peft_init_lora_weights: bool | str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "How to initialize LoRA weights. Default to True which is MS original implementation."
+        },
+    )
+    peft_trainable_token_indices: list[int] | dict[str, list[int]] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "A list of token indices to fine-tune on the `embed_tokens` layer.\n"
+                "Otherwise, a dict mapping an embedding layer name to its trainable token indices.\n"
+                "See https://huggingface.co/docs/peft/v0.17.0/en/developer_guides/lora#efficiently-train-tokens-alongside-lora"
+            )
+        },
+    )
+    peft_ensure_weight_tying: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Whether to tie adapter weights for tied model weights. "
+                "See https://github.com/huggingface/peft/issues/2864"
+            )
+        },
+    )
+    peft_autocast_adapter_dtype: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to upcast the LoRA adapter to fp32. This is enabled by default in PEFT."
+        },
+    )
+
+    qlora_sharded_model_loading: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "load qlora model in sharded format for FSDP using answer.ai technique."
+        },
+    )
+    lora_on_cpu: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "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"
+        },
+    )
+    gptq: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether you are training a 4-bit GPTQ quantized model"
+        },
+    )
+    bnb_config_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "optional overrides to the bnb 4bit quantization configuration"
+        },
+    )
+
+    loraplus_lr_ratio: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "loraplus learning rate ratio lr_B / lr_A. Recommended value is 2^4."
+        },
+    )
+    loraplus_lr_embedding: float | None = Field(
+        default=1e-6,
+        json_schema_extra={
+            "description": "loraplus learning rate for lora embedding layers. Default value is 1e-6."
+        },
+    )
+
+    merge_lora: bool | None = None
+    merge_method: Literal["legacy", "memory_efficient"] | None = Field(
+        default="memory_efficient",
+        json_schema_extra={
+            "description": "Method to use for LoRA merging. 'memory_efficient' (default) processes shards individually to reduce memory usage, 'legacy' loads the full model into memory."
+        },
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_adapter(cls, data):
+        if (
+            not data.get("adapter")
+            and not data.get("inference")
+            and (data.get("load_in_8bit") or data.get("load_in_4bit"))
+        ):
+            raise ValueError(
+                "load_in_8bit and load_in_4bit are not supported without setting an adapter for training."
+                "If you want to full finetune, please turn off load_in_8bit and load_in_4bit."
+            )
+        adapter = data.get("adapter")
+        if adapter and adapter not in ("lora", "qlora", "llama-adapter"):
+            from axolotl.integrations.base import PluginManager
+
+            plugin_manager = PluginManager.get_instance()
+            if not plugin_manager.supports_adapter(adapter):
+                raise ValueError(
+                    f"Adapter '{adapter}' is not built in and was not registered by "
+                    "a plugin. Add the plugin that provides this adapter to `plugins:`."
+                )
+        return data
+
+    @model_validator(mode="after")
+    def validate_qlora(self):
+        if self.adapter == "qlora":
+            if self.merge_lora:
+                # can't merge qlora if loaded in 8bit or 4bit
+                if self.load_in_8bit:
+                    raise ValueError("Can't merge qlora if loaded in 8bit")
+
+                if self.gptq:
+                    raise ValueError("Can't merge qlora if gptq")
+
+                if self.load_in_4bit:
+                    raise ValueError("Can't merge qlora if loaded in 4bit")
+
+            else:
+                if self.load_in_8bit:
+                    raise ValueError("Can't load qlora in 8bit")
+
+                if self.gptq:
+                    raise ValueError("Can't load qlora if gptq")
+
+                if not self.load_in_4bit:
+                    raise ValueError("Require cfg.load_in_4bit to be True for qlora")
+        return self
+
+    @field_validator("loraplus_lr_embedding")
+    @classmethod
+    def convert_loraplus_lr_embedding(cls, loraplus_lr_embedding):
+        if loraplus_lr_embedding and isinstance(loraplus_lr_embedding, str):
+            loraplus_lr_embedding = float(loraplus_lr_embedding)
+        return loraplus_lr_embedding
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_lora_dropout(cls, data):
+        if data.get("adapter") is not None and data.get("lora_dropout") is None:
+            data["lora_dropout"] = 0.0
+        return data
+
+    @model_validator(mode="after")
+    def validate_lora_target_parameters_dropout(self):
+        if (
+            self.lora_target_parameters
+            and self.lora_dropout
+            and self.lora_dropout != 0.0
+        ):
+            raise ValueError(
+                "lora_dropout must be 0 when lora_target_parameters is set. "
+                "PEFT's ParamWrapper does not support lora_dropout != 0."
+            )
+        return self
+
+
+class ReLoRAConfig(BaseModel):
+    """ReLoRA configuration subset"""
+
+    relora: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to use ReLoRA. Use with jagged_restart_*steps options."
+        },
+    )
+    relora_prune_ratio: float | None = Field(
+        default=None,
+        ge=0.0,
+        le=1.0,
+        json_schema_extra={
+            "description": (
+                "Fraction of optimizer state values to zero on each ReLoRA restart. "
+                "When relora_prune_method='reset' and this is omitted, defaults to "
+                "0.999 (paper-style near-full reset). For other methods, defaults to 0.9."
+            )
+        },
+    )
+    relora_prune_method: Literal["magnitude", "random", "reset"] | None = Field(
+        default="magnitude",
+        json_schema_extra={
+            "description": (
+                "Optimizer state pruning method on each ReLoRA restart. "
+                "'magnitude' (default) keeps top-k by absolute value; "
+                "'random' keeps a random subset at relora_prune_ratio; "
+                "'reset' uses near-full random pruning (default ratio 0.999, "
+                "honoring relora_prune_ratio when explicitly set). "
+                "Paper-style recipe: relora_prune_method='reset' with no "
+                "relora_prune_ratio, equivalent to 'random' with ratio=0.999."
+            )
+        },
+    )
+    relora_cpu_offload: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "True to perform lora weight merges on cpu during restarts, for modest gpu memory savings"
+        },
+    )
diff --git a/src/axolotl/utils/schemas/quantization.py b/src/axolotl/utils/schemas/quantization.py
new file mode 100644
index 0000000000..b15e5d225e
--- /dev/null
+++ b/src/axolotl/utils/schemas/quantization.py
@@ -0,0 +1,85 @@
+"""
+QAT Config Schema
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator
+
+from axolotl.utils.schemas.enums import TorchAOQuantDType
+
+
+def validate_ao_dtype(v: Any) -> TorchAOQuantDType | None:
+    if v is None:
+        return None
+    if v == "int4":
+        return TorchAOQuantDType.int4
+    if v == "int8":
+        return TorchAOQuantDType.int8
+    if v in ["float8_e4m3fn", "fp8", "float8"]:
+        return TorchAOQuantDType.float8_e4m3fn
+    if v == "nvfp4":
+        return TorchAOQuantDType.nvfp4
+    if v == "mxfp4":
+        return TorchAOQuantDType.mxfp4
+
+    raise ValueError(
+        f"Invalid dtype: '{v}'. Must be one of: {[e.name for e in TorchAOQuantDType] + ['fp8', 'float8']}"
+    )
+
+
+class QATConfig(BaseModel):
+    """
+    QAT Config Schema
+    """
+
+    activation_dtype: TorchAOQuantDType | None = Field(
+        default=None,
+        description="Fake quantization layout to use for activation quantization.",
+    )
+    weight_dtype: TorchAOQuantDType = Field(
+        default=TorchAOQuantDType.int8,
+        description="Fake quantization layout to use for weight quantization.",
+    )
+    quantize_embedding: bool | None = Field(
+        default=False, description="Quantize embedding"
+    )
+    group_size: int | None = Field(
+        default=32,
+        description="The number of elements in each group for per-group fake quantization",
+    )
+    fake_quant_after_n_steps: int | None = Field(
+        default=None, description="The number of steps to apply fake quantization after"
+    )
+
+    @field_validator("activation_dtype", "weight_dtype", mode="before")
+    @classmethod
+    def validate_dtype(cls, v: Any) -> TorchAOQuantDType | None:
+        return validate_ao_dtype(v)
+
+
+class PTQConfig(BaseModel):
+    """
+    PTQ Config Schema
+    """
+
+    weight_dtype: TorchAOQuantDType = Field(
+        default=TorchAOQuantDType.int8,
+        description="Fake quantization layout to use for weight quantization.",
+    )
+    activation_dtype: TorchAOQuantDType | None = Field(
+        default=None,
+        description="Fake quantization layout to use for activation quantization.",
+    )
+    quantize_embedding: bool | None = Field(
+        default=None, description="Whether to quantize the embedding layer."
+    )
+    group_size: int | None = Field(
+        default=32,
+        description="The number of elements in each group for per-group fake quantization",
+    )
+
+    @field_validator("activation_dtype", "weight_dtype", mode="before")
+    @classmethod
+    def validate_dtype(cls, v: Any) -> TorchAOQuantDType | None:
+        return validate_ao_dtype(v)
diff --git a/src/axolotl/utils/schemas/training.py b/src/axolotl/utils/schemas/training.py
new file mode 100644
index 0000000000..0d8c6ab2e4
--- /dev/null
+++ b/src/axolotl/utils/schemas/training.py
@@ -0,0 +1,306 @@
+"""Pydantic models for training hyperparameters"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, field_validator
+
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.enums import CustomSupportedOptimizers
+
+LOG = get_logger(__name__)
+
+
+class LrGroup(BaseModel):
+    """Custom learning rate group configuration"""
+
+    name: str
+    modules: list[str]
+    lr: float
+
+
+class SelectiveCheckpointingConfig(BaseModel):
+    """Selective activation checkpointing (SAC) configuration"""
+
+    save: list[str] = Field(
+        default_factory=lambda: ["attention"],
+        json_schema_extra={
+            "description": (
+                "Ops to save during forward instead of recomputing in backward. "
+                "'attention' matches SDPA and flash-attention forward ops; other "
+                "entries are substring-matched against qualified torch op names "
+                "(e.g. 'aten::mm')."
+            )
+        },
+    )
+    save_sliding_window: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": (
+                "In hybrid full/sliding-window attention models, also save "
+                "sliding-window attention calls. Default false: SWA is cheap to "
+                "recompute, so only full-attention calls are saved."
+            )
+        },
+    )
+    recompute_layer_types: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "Layer types (config.layer_types values) whose attention is "
+                "recomputed instead of saved. Defaults to ['sliding_attention', "
+                "'chunked_attention']. Ignored when save_sliding_window is true. "
+                "Linear-attention layers never dispatch a matchable attention op, "
+                "so they need no entry."
+            )
+        },
+    )
+    offload: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": (
+                "Offload saved tensors to pinned CPU memory (side-stream copies "
+                "with backward prefetch) instead of keeping them on GPU."
+            )
+        },
+    )
+
+
+class HyperparametersConfig(BaseModel):
+    """Training hyperparams configuration subset"""
+
+    gradient_accumulation_steps: int | None = Field(
+        default=1,
+        json_schema_extra={
+            "description": "If greater than 1, backpropagation will be skipped and the gradients will be accumulated for the given number of steps."
+        },
+    )
+    micro_batch_size: int | None = Field(
+        default=1,
+        json_schema_extra={
+            "description": "The number of samples to include in each batch. This is the number of samples sent to each GPU. Batch size per gpu = micro_batch_size * gradient_accumulation_steps"
+        },
+    )
+    batch_size: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Total batch size, we do not recommended setting this manually"
+        },
+    )
+    eval_batch_size: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "per gpu micro batch size for evals, defaults to value of micro_batch_size"
+        },
+    )
+
+    auto_find_batch_size: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "whether to find batch size that fits in memory. Passed to underlying transformers Trainer"
+        },
+    )
+
+    train_on_inputs: bool | None = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Whether to mask out or include the human's prompt from the training labels"
+        },
+    )
+    group_by_length: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Group similarly sized data to minimize padding. May be slower to start, as it must download and sort the entire dataset. Note that training loss may have an oscillating pattern with this enabled."
+        },
+    )
+
+    learning_rate: str | float
+    embedding_lr: float | None = None
+    embedding_lr_scale: float | None = None
+    weight_decay: float | None = Field(
+        default=0.0, json_schema_extra={"description": "Specify weight decay"}
+    )
+    optimizer: (str | CustomSupportedOptimizers) | None = Field(
+        default="adamw_torch_fused",
+        json_schema_extra={"description": "Specify optimizer"},
+    )
+    optim_args: (str | dict[str, Any]) | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Dictionary of arguments to pass to the optimizer"
+        },
+    )
+    optim_target_modules: (list[str] | Literal["all_linear"]) | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The target modules to optimize, i.e. the module names that you would like to train, right now this is used only for GaLore algorithm"
+        },
+    )
+    torchdistx_path: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Path to torch distx for optim 'adamw_anyprecision'"
+        },
+    )
+    lr_scheduler: str | None = "cosine"
+    lr_scheduler_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Specify a scheduler and kwargs to use with the optimizer"
+        },
+    )
+    lr_quadratic_warmup: bool | None = None
+    cosine_min_lr_ratio: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "decay lr to some percentage of the peak lr, e.g. cosine_min_lr_ratio=0.1 for 10% of peak lr"
+        },
+    )
+    cosine_constant_lr_ratio: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "freeze lr at some percentage of the step, e.g. cosine_constant_lr_ratio=0.8 means start cosine_min_lr at 80% of training step"
+        },
+    )
+    lr_div_factor: float | None = Field(
+        default=None, json_schema_extra={"description": "Learning rate div factor"}
+    )
+    lr_groups: list[LrGroup] | None = None
+
+    adam_epsilon: float | None = Field(
+        default=None, json_schema_extra={"description": "adamw hyperparams"}
+    )
+    adam_epsilon2: float | None = Field(
+        default=None, json_schema_extra={"description": "only used for CAME Optimizer"}
+    )
+    adam_beta1: float | None = Field(
+        default=None, json_schema_extra={"description": "adamw hyperparams"}
+    )
+    adam_beta2: float | None = Field(
+        default=None, json_schema_extra={"description": "adamw hyperparams"}
+    )
+    adam_beta3: float | None = Field(
+        default=None, json_schema_extra={"description": "only used for CAME Optimizer"}
+    )
+
+    dion_lr: float | None = Field(
+        default=None, json_schema_extra={"description": "Dion Optimizer learning rate"}
+    )
+    dion_momentum: float | None = Field(
+        default=None, json_schema_extra={"description": "Dion Optimizer momentum"}
+    )
+    dion_rank_fraction: float | None = Field(
+        default=1.0,
+        json_schema_extra={
+            "description": "Dion Optimizer: r/d fraction for low-rank approximation. Used to compute the low-rank dimension."
+        },
+    )
+    dion_rank_multiple_of: int | None = Field(
+        default=1,
+        json_schema_extra={
+            "description": "Dion Optimizer: Round up the low-rank dimension to a multiple of this number. This may be useful to ensure even sharding."
+        },
+    )
+
+    qgalore_rank: int | None = Field(
+        default=256,
+        json_schema_extra={
+            "description": "Q-GaLore: rank r of the low-rank gradient projection. Smaller r reduces optimizer state but loses gradient information."
+        },
+    )
+    qgalore_update_proj_gap: int | None = Field(
+        default=200,
+        json_schema_extra={
+            "description": "Q-GaLore: maximum number of steps between SVD recomputations of the projection matrix. The adaptive scheduler may skip updates earlier based on cos_threshold."
+        },
+    )
+    qgalore_scale: float | None = Field(
+        default=0.25,
+        json_schema_extra={
+            "description": "Q-GaLore: scaling factor applied to the projected gradient after project_back. Equivalent to GaLore's `galore_scale`."
+        },
+    )
+    qgalore_proj_type: str | None = Field(
+        default="std",
+        json_schema_extra={
+            "description": "Q-GaLore: projection type for the GaLoreProjector. One of 'std', 'reverse_std', 'right', 'left', 'full'."
+        },
+    )
+    qgalore_proj_quant: bool | None = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Q-GaLore: enable INT-quantization of the projection matrix P (the resilient-to-quantization observation from the paper)."
+        },
+    )
+    qgalore_proj_bits: int | None = Field(
+        default=4,
+        json_schema_extra={
+            "description": "Q-GaLore: bitwidth for the quantized projection matrix when qgalore_proj_quant is True (paper default: 4)."
+        },
+    )
+    qgalore_proj_group_size: int | None = Field(
+        default=256,
+        json_schema_extra={
+            "description": "Q-GaLore: group size for projection-matrix quantization. Must evenly divide the projection's last dimension."
+        },
+    )
+    qgalore_cos_threshold: float | None = Field(
+        default=0.4,
+        json_schema_extra={
+            "description": "Q-GaLore: cosine-similarity threshold for the lazy subspace update. If the new P is within this similarity of the previous one, the SVD is skipped."
+        },
+    )
+    qgalore_gamma_proj: int | None = Field(
+        default=2,
+        json_schema_extra={
+            "description": "Q-GaLore: multiplicative factor by which update_proj_gap grows once a layer's subspace is judged stable."
+        },
+    )
+    qgalore_queue_size: int | None = Field(
+        default=5,
+        json_schema_extra={
+            "description": "Q-GaLore: length of the moving-average queue used by the adaptive-frequency scheduler."
+        },
+    )
+    max_grad_norm: float | None = Field(
+        default=None, json_schema_extra={"description": "Gradient clipping max norm"}
+    )
+    num_epochs: float = Field(default=1.0)
+
+    @field_validator("batch_size")
+    @classmethod
+    def hint_batch_size_set(cls, batch_size):
+        if batch_size:
+            LOG.warning(
+                "%s\n%s",
+                "batch_size is not recommended. Please use gradient_accumulation_steps instead.",
+                "To calculate the equivalent gradient_accumulation_steps, divide batch_size / micro_batch_size / number of gpus.",
+            )
+        return batch_size
+
+    @field_validator("learning_rate")
+    @classmethod
+    def convert_learning_rate(cls, learning_rate):
+        if learning_rate and isinstance(learning_rate, str):
+            learning_rate = float(learning_rate)
+        return learning_rate
+
+
+class JaggedLRConfig(BaseModel):
+    """JaggedLR configuration subset, can be used w/ ReLoRA training"""
+
+    jagged_restart_steps: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "how often to reset for jagged restarts"},
+    )
+    jagged_restart_warmup_steps: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "how many warmup steps to take after reset for jagged restarts"
+        },
+    )
+    jagged_restart_anneal_steps: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "how many anneal steps to take before reset for jagged restarts"
+        },
+    )
diff --git a/src/axolotl/utils/schemas/trl.py b/src/axolotl/utils/schemas/trl.py
new file mode 100644
index 0000000000..a362421627
--- /dev/null
+++ b/src/axolotl/utils/schemas/trl.py
@@ -0,0 +1,333 @@
+"""Pydantic models for TRL trainer configuration"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class TRLConfig(BaseModel):
+    """
+    Input args for TRL.
+    """
+
+    beta: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Beta parameter for the RL training. Same as `rl_beta`. Use"
+        },
+    )
+    max_completion_length: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Maximum length of the completion for RL training."
+        },
+    )
+
+    # GRPO specific args
+    # Ref: https://github.com/huggingface/trl/blob/26d86757a7c7e24e397ea44f57ecce6031dfac01/trl/trainer/grpo_config.py#L23
+    use_vllm: bool = Field(
+        default=False,
+        json_schema_extra={"description": "Whether to use VLLM for RL training."},
+    )
+    vllm_mode: Literal["server", "colocate"] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "VLLM mode to use, one of 'server' or 'colocate'"
+        },
+    )
+    vllm_server_host: str | None = Field(
+        default="0.0.0.0",  # nosec B104
+        json_schema_extra={"description": "Host of the vLLM server to connect to."},
+    )
+    vllm_server_port: int | None = Field(
+        default=8000,
+        json_schema_extra={"description": "Port of the vLLM server to connect to."},
+    )
+    vllm_server_timeout: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Total timeout (in seconds) to wait for the vLLM server to respond."
+        },
+    )
+    vllm_guided_decoding_regex: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Regex for vLLM guided decoding."},
+    )
+
+    reward_funcs: list[str] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "List of reward functions to load. Paths must be importable from current dir."
+        },
+    )
+    reward_weights: list[float] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "List of reward weights for the reward functions."
+        },
+    )
+    generation_batch_size: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Batch size for generation. Controls how many unique prompts are generated per step. Should be num_generations * data_parallel_size for full DP utilization."
+        },
+    )
+    num_generations: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Number of generations to sample."},
+    )
+    log_completions: bool | None = Field(
+        default=False,
+        json_schema_extra={"description": "Whether to log completions."},
+    )
+    num_completions_to_print: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of completions to print when log_completions is True."
+        },
+    )
+    importance_sampling_level: Literal["sequence", "token"] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Controls whether importance sampling ratios are computed at the `'token'` or `'sequence'` level. "
+            "For GSPO, use `sequence`, default is None which corresponds to the original GRPO paper."
+        },
+    )
+
+    sync_ref_model: bool | None = Field(
+        default=False,
+        json_schema_extra={"description": "Whether to sync the reference model."},
+    )
+    ref_model_mixup_alpha: float | None = Field(
+        default=0.9,
+        json_schema_extra={"description": "Mixup alpha for the reference model."},
+    )
+    ref_model_sync_steps: int | None = Field(
+        default=64,
+        json_schema_extra={"description": "Sync steps for the reference model."},
+    )
+    scale_rewards: bool = Field(
+        default=True,
+        json_schema_extra={
+            "description": "Whether to scale rewards by their standard deviation."
+        },
+    )
+
+    temperature: float | None = Field(
+        default=None,
+        json_schema_extra={"description": "Sampling temperature for the GRPO policy."},
+    )
+    top_p: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Top-p sampling probability for the generation policy."
+        },
+    )
+    top_k: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Top-k sampling for the generation policy."},
+    )
+    min_p: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Minimum probability for the generation policy."
+        },
+    )
+    repetition_penalty: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Penalty for tokens that appear in prompt and generated text."
+        },
+    )
+    generation_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Additional generation parameters passed to vLLM SamplingParams. "
+            "Useful for stop_token_ids, seed, frequency_penalty, etc."
+        },
+    )
+    chat_template_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Additional kwargs for the chat template. "
+            "E.g., {enable_thinking: false} for Qwen3.5 models."
+        },
+    )
+    num_iterations: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of iterations per batch (μ) for GRPO."
+        },
+    )
+    epsilon: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Epsilon value for clipping in the GRPO algorithm."
+        },
+    )
+    epsilon_high: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Upper-bound epsilon value for clipping in the GRPO algorithm."
+        },
+    )
+    use_liger_loss: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Whether to use Liger loss for GRPO."},
+    )
+    loss_type: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Loss formulation to use. Supported values: grpo, bnpo, dr_grpo."
+        },
+    )
+    mask_truncated_completions: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Whether to exclude truncated completions from loss calculation."
+        },
+    )
+    vllm_enable_sleep_mode: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Enable sleep mode for vLLM to offload VRAM when idle"
+        },
+    )
+    rollout_func: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Path to custom rollout function. Must be importable from current dir."
+        },
+    )
+    multi_objective_aggregation: (
+        Literal["sum_then_normalize", "normalize_then_sum"] | None
+    ) = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Multi-objective reward aggregation strategy. "
+            "'sum_then_normalize' (GRPO default): weights and sums rewards first, then normalizes. "
+            "'normalize_then_sum' (GDPO): normalizes each reward independently, then sums."
+        },
+    )
+
+    # Async GRPO fields
+    use_data_producer: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Use the GRPODataProducer protocol for online data generation."
+        },
+    )
+    async_prefetch: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Generate rollouts in a background thread while training on the previous rollout."
+        },
+    )
+    prefetch_depth: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Number of rollouts to prefetch ahead of training."
+        },
+    )
+    vllm_sync_interval: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Sync model weights to vLLM every N optimizer steps (async mode only)."
+        },
+    )
+    streaming_partial_batch: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Score prompt groups incrementally instead of the full batch at once."
+        },
+    )
+    streaming_min_groups: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Minimum prompt groups to score per streaming chunk."
+        },
+    )
+    vllm_importance_sampling_correction: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Apply IS correction for distribution mismatch between vLLM and training model."
+        },
+    )
+    vllm_importance_sampling_mode: (
+        Literal["token_truncate", "token_mask", "sequence_truncate", "sequence_mask"]
+        | None
+    ) = Field(
+        default=None,
+        json_schema_extra={
+            "description": "IS mode: token_truncate, token_mask, sequence_truncate, or sequence_mask."
+        },
+    )
+    vllm_importance_sampling_cap: float | None = Field(
+        default=None,
+        json_schema_extra={"description": "Cap C for IS ratio clipping/masking."},
+    )
+    off_policy_mask_threshold: float | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "KL threshold for off-policy sequence masking (OPSM). None = disabled."
+        },
+    )
+    use_bias_correction_kl: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Apply IS correction to KL divergence term."},
+    )
+
+    reward_num_workers: int = Field(
+        default=1,
+        json_schema_extra={
+            "description": "Number of persistent subprocess workers for parallel reward computation. Each worker has its "
+            "own main thread so signal.alarm() (used by math_verify) works correctly. Work is sharded across "
+            "workers by prompt groups. Only used with use_data_producer=True and non-nn.Module reward functions."
+        },
+    )
+    replay_buffer_size: int = Field(
+        default=0,
+        json_schema_extra={
+            "description": "[Experimental, disabled by default] Size of the replay buffer for storing high-signal rollout "
+            "groups. When > 0, groups with reward variance are cached and used to replace zero-signal groups "
+            "(where all rewards are identical). Set to 0 to disable. Only used with use_data_producer=True."
+        },
+    )
+    replay_recompute_logps: bool = Field(
+        default=True,
+        json_schema_extra={
+            "description": "When True (default), recompute old_per_token_logps for replayed groups using the current "
+            "training model. This fixes the importance sampling mismatch that occurs when replaying stale data. "
+            "Only relevant when replay_buffer_size > 0."
+        },
+    )
+    reroll_start_fraction: float = Field(
+        default=1.0,
+        json_schema_extra={
+            "description": "Fraction of total training steps after which deferred re-rolling begins. Zero-signal prompts "
+            "(where all rewards in a group are identical) are buffered and re-injected into later batches when the "
+            "model is more likely to solve them. Set to 1.0 to disable. Only used with use_data_producer=True."
+        },
+    )
+    reroll_max_groups: int = Field(
+        default=1,
+        json_schema_extra={
+            "description": "Maximum number of prompt groups to replace with re-roll candidates per batch. Higher values "
+            "increase data utilization but reduce prompt diversity. Only used with use_data_producer=True."
+        },
+    )
+    skip_zero_advantage_batches: bool = Field(
+        default=True,
+        json_schema_extra={
+            "description": "When True, skip gradient computation for micro-batches where all advantages are zero (no learning "
+            "signal). This avoids the forward/backward pass entirely when no learning signal is present. The step is "
+            "logged with skipped_zero_adv_batches=1 for monitoring."
+        },
+    )
+    vllm_lora_sync: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": "Sync LoRA adapter to vLLM via filesystem instead of merging + NCCL broadcast. "
+            "Auto-selects vllm_serve_lora serve module. Syncs only LoRA adapter weights vs full merged model."
+        },
+    )
diff --git a/src/axolotl/utils/schemas/utils.py b/src/axolotl/utils/schemas/utils.py
new file mode 100644
index 0000000000..b46c8f8475
--- /dev/null
+++ b/src/axolotl/utils/schemas/utils.py
@@ -0,0 +1,79 @@
+"""Utilities for Axolotl Pydantic models"""
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def handle_legacy_message_fields_logic(data: dict) -> dict:
+    """
+    Handle backwards compatibility between legacy message field mapping and new property mapping system.
+
+    Previously, the config only supported mapping 'role' and 'content' fields via dedicated config options:
+    - message_field_role: Mapped to the role field
+    - message_field_content: Mapped to the content field
+
+    The new system uses message_property_mappings to support arbitrary field mappings:
+    message_property_mappings:
+        role: source_role_field
+        content: source_content_field
+        additional_field: source_field
+
+    Args:
+        data: Dictionary containing configuration data
+
+    Returns:
+        Updated dictionary with message field mappings consolidated
+
+    Raises:
+        ValueError: If there are conflicts between legacy and new mappings
+    """
+    data = data.copy()  # Create a copy to avoid modifying the original
+
+    if data.get("message_property_mappings") is None:
+        data["message_property_mappings"] = {}
+
+    # Check for conflicts and handle role
+    if "message_field_role" in data:
+        LOG.warning(
+            "message_field_role is deprecated, use message_property_mappings instead. "
+            f"Example: message_property_mappings: {{role: {data['message_field_role']}}}"
+        )
+        if (
+            "role" in data["message_property_mappings"]
+            and data["message_property_mappings"]["role"] != data["message_field_role"]
+        ):
+            raise ValueError(
+                f"Conflicting message role fields: message_field_role='{data['message_field_role']}' "
+                f"conflicts with message_property_mappings.role='{data['message_property_mappings']['role']}'"
+            )
+        data["message_property_mappings"]["role"] = data["message_field_role"] or "role"
+
+        del data["message_field_role"]
+    elif "role" not in data["message_property_mappings"]:
+        data["message_property_mappings"]["role"] = "role"
+
+    # Check for conflicts and handle content
+    if "message_field_content" in data:
+        LOG.warning(
+            "message_field_content is deprecated, use message_property_mappings instead. "
+            f"Example: message_property_mappings: {{content: {data['message_field_content']}}}"
+        )
+        if (
+            "content" in data["message_property_mappings"]
+            and data["message_property_mappings"]["content"]
+            != data["message_field_content"]
+        ):
+            raise ValueError(
+                f"Conflicting message content fields: message_field_content='{data['message_field_content']}' "
+                f"conflicts with message_property_mappings.content='{data['message_property_mappings']['content']}'"
+            )
+        data["message_property_mappings"]["content"] = (
+            data["message_field_content"] or "content"
+        )
+
+        del data["message_field_content"]
+    elif "content" not in data["message_property_mappings"]:
+        data["message_property_mappings"]["content"] = "content"
+
+    return data
diff --git a/src/axolotl/utils/schemas/validation.py b/src/axolotl/utils/schemas/validation.py
new file mode 100644
index 0000000000..41f4cc02f4
--- /dev/null
+++ b/src/axolotl/utils/schemas/validation.py
@@ -0,0 +1,1931 @@
+"""Module with validation methods for config pydantic model."""
+
+import json
+import sys
+import tempfile
+from pathlib import Path
+
+from pydantic import (
+    field_validator,
+    model_validator,
+)
+from transformers.utils.import_utils import is_torch_npu_available
+
+from axolotl.utils.logging import get_logger
+from axolotl.utils.schemas.enums import (
+    ChatTemplate,
+    RingAttnFunc,
+    RLType,
+)
+
+LOG = get_logger(__name__)
+
+SUPPORTED_METRICS = {"sacrebleu", "comet", "ter", "chrf", "perplexity"}
+
+
+class DatasetValidationMixin:
+    """Validation methods related to dataset configuration."""
+
+    @field_validator("seed", mode="after")
+    @classmethod
+    def set_default_seed(cls, seed):
+        if seed is None:
+            LOG.info("`seed` not set in config; setting to 42")
+            seed = 42
+        return seed
+
+    @field_validator("datasets", mode="before")
+    @classmethod
+    def deprecate_sharegpt_datasets(cls, datasets):
+        for _, ds_cfg in enumerate(datasets):
+            ds_type = (
+                ds_cfg.get("type")
+                if isinstance(ds_cfg, dict)
+                else getattr(ds_cfg, "type", None)
+            )
+            if not ds_type:
+                continue
+
+            if isinstance(ds_type, dict):
+                continue
+
+            if isinstance(ds_type, str) and ds_type.startswith("sharegpt"):
+                raise ValueError(
+                    "`type: sharegpt.*` is deprecated. Please use `type: chat_template` instead."
+                )
+
+        return datasets
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_deprecated_unsloth_fields(cls, data):
+        deprecated_fields = [
+            "unsloth_cross_entropy_loss",
+            "unsloth_lora_mlp",
+            "unsloth_lora_qkv",
+            "unsloth_lora_o",
+            "unsloth_rms_norm",
+            "unsloth_rope",
+        ]
+        found = [f for f in deprecated_fields if data.get(f)]
+        if found:
+            raise ValueError(
+                f"`{'`, `'.join(found)}` {'has' if len(found) == 1 else 'have'} been removed. "
+                "Please use `lora_mlp_kernel`, `lora_qkv_kernel`, `lora_o_kernel` instead. "
+                "See: https://docs.axolotl.ai/docs/lora_optims.html"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_dataset_or_pretraining_dataset(cls, data):
+        if data.get("datasets") is None and data.get("pretraining_dataset") is None:
+            raise ValueError("either datasets or pretraining_dataset is required")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_pretraining_streaming_deprecation(cls, data):
+        # TODO(djsaunde): remove this check + implement change for 0.13.0 release
+        if data.get("pretraining_dataset") and not data.get("streaming"):
+            LOG.warning(
+                "Setting `pretraining_dataset` without explicitly setting `streaming: "
+                "true` is deprecated. In a future release, streaming will not be "
+                "automatically enabled when using pretraining_dataset. Please "
+                "explicitly set `streaming: true` in your configuration to maintain "
+                "current behavior."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_push_ds_auth(cls, data):
+        if (
+            data.get("push_dataset_to_hub")
+            and data.get("hf_use_auth_token") is not True
+        ):
+            raise ValueError(
+                "Require cfg.hf_use_auth_token to be True for push_dataset_to_hub"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_val_w_test_datasets(cls, data):
+        if data.get("test_datasets") and data.get("val_set_size"):
+            raise ValueError(
+                "non-zero val_set_size should not be used with test_datasets configuration"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_test_datasets_bench(cls, data):
+        if (
+            data.get("do_bench_eval")
+            and not data.get("test_datasets")
+            and not data.get("val_set_size")
+        ):
+            LOG.warning(
+                "`do_bench_eval` needs a test dataset to run evals, adding an empty test_dataset."
+            )
+            data["test_datasets"] = [{"path": "axolotl-ai-co/empty-test-ds"}]
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_eval_packing(cls, data):
+        # TODO also should check test_datasets and val_set_size as we can skip
+        # if there are no eval datasets/splits
+        if (
+            data.get("sample_packing")
+            and data.get("eval_table_size")
+            and data.get("eval_sample_packing") is not False
+        ):
+            raise ValueError(
+                "eval_table_size and eval_sample_packing are not supported together with sample_packing. Please set 'eval_sample_packing' to false."
+            )
+        if (
+            data.get("sample_packing")
+            and data.get("eval_sample_packing") is None
+            and not data.get("eval_table_size")
+        ):
+            LOG.info(
+                "explicitly setting `eval_sample_packing` to match `sample_packing`",
+            )
+            data["eval_sample_packing"] = True
+
+        if (
+            data.get("sample_packing")
+            and data.get("eval_sample_packing") is False
+            and data.get("remove_unused_columns") is None
+        ):
+            LOG.info(
+                "setting `remove_unused_columns: false` for when sample_packing and eval_sample_packing don't match"
+            )
+            data["remove_unused_columns"] = False
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_mm_prepare(cls, data):
+        if data.get("skip_prepare_dataset"):
+            if data.get("remove_unused_columns") is None:
+                LOG.info(
+                    "setting `remove_unused_columns: false` for skip_prepare_dataset"
+                )
+                data["remove_unused_columns"] = False
+
+        return data
+
+
+class AttentionValidationMixin:
+    """Validation methods related to attention mechanisms."""
+
+    @model_validator(mode="after")
+    def check_sample_packing_without_attention(self):
+        if self.sample_packing and not self.attn_decontaminates_packing:
+            if self.attn_implementation:
+                LOG.warning(
+                    "`sample_packing` with `attn_implementation=%r` does not handle "
+                    "cross-sample decontamination. Use a packing-capable backend "
+                    "(e.g. flash_attention_2, flex_attention, sdpa, eager) to "
+                    "isolate samples.",
+                    self.attn_implementation,
+                )
+            else:
+                LOG.warning(
+                    "`sample_packing` without an attention backend does not handle "
+                    "cross-sample decontamination. Set `attn_implementation` to a "
+                    "packing-capable backend (e.g. flash_attention_2)."
+                )
+        return self
+
+    @model_validator(mode="after")
+    def check_scaling_softmax_requires_flex(self):
+        if self.scaling_softmax and self.attn_implementation != "flex_attention":
+            raise ValueError(
+                "scaling_softmax requires flex attention. "
+                "Add `attn_implementation: flex_attention` to your config."
+            )
+        return self
+
+
+class TrainingValidationMixin:
+    """Validation methods related to training configuration."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_batch_size_fields(cls, data):
+        fields = ("micro_batch_size", "gradient_accumulation_steps", "batch_size")
+        non_empty_count = sum(1 for field in fields if data.get(field))
+
+        if non_empty_count < 2:
+            raise ValueError(f"At least two of {', '.join(fields)} must be set")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def hint_sample_packing_padding(cls, data):
+        if data.get("sample_packing"):
+            pad_to_sequence_len = data.get("pad_to_sequence_len")
+            if pad_to_sequence_len is False:
+                LOG.warning(
+                    "`pad_to_sequence_len: true` is recommended when using sample_packing"
+                )
+            elif pad_to_sequence_len is None:
+                LOG.info(
+                    "Setting `pad_to_sequence_len: true` to prevent memory leaks when sample_packing"
+                )
+                data["pad_to_sequence_len"] = True
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def hint_reward_model_pad(cls, data):
+        if data.get("reward_model") and not data.get("pad_to_sequence_len"):
+            LOG.warning(
+                "`pad_to_sequence_len: true` is recommended when using reward_model"
+            )
+            if data.get("pad_to_sequence_len") is None:
+                data["pad_to_sequence_len"] = True
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def set_reward_model_defaults(cls, data):
+        if data.get("reward_model"):
+            if data.get("num_labels") is None:
+                data["num_labels"] = 1
+            if not (data.get("type_of_model") or data.get("model_type")):
+                data["model_type"] = "AutoModelForSequenceClassification"
+
+        if data.get("process_reward_model"):
+            if data.get("num_labels") is None:
+                data["num_labels"] = 2
+            if not (data.get("type_of_model") or data.get("model_type")):
+                data["model_type"] = "AutoModelForTokenClassification"
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_gas_bsz(cls, data):
+        if data.get("gradient_accumulation_steps") and data.get("batch_size"):
+            raise ValueError(
+                "please set only one of gradient_accumulation_steps or batch_size"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def hint_eval_train_mbsz(cls, data):
+        if (
+            data.get("eval_batch_size")
+            and data.get("micro_batch_size")
+            and data.get("eval_batch_size") != data.get("micro_batch_size")
+        ):
+            LOG.warning(
+                "eval_batch_size != micro_batch_size. This can lead to VRAM instability."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_warmup(cls, data):
+        if data.get("warmup_steps") and data.get("warmup_ratio"):
+            raise ValueError("warmup_steps and warmup_ratio are mutually exclusive")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_saves(cls, data):
+        if (
+            data.get("save_strategy")
+            and data.get("save_steps")
+            and data.get("save_strategy") != "steps"
+        ):
+            raise ValueError(
+                "save_strategy and save_steps mismatch. Please set save_strategy to 'steps' or remove save_steps."
+            )
+        if data.get("saves_per_epoch") and data.get("save_steps"):
+            raise ValueError(
+                "save_steps and saves_per_epoch are mutually exclusive and cannot be used together."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_push_save(cls, data):
+        if data.get("hub_model_id") and (
+            data.get("save_strategy") not in ["steps", "epoch", None]
+        ):
+            LOG.warning(
+                "hub_model_id is set without any models being saved. To save a model, set save_strategy."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_evals(cls, data):
+        if (
+            data.get("eval_strategy")
+            and data.get("eval_steps")
+            and data.get("eval_strategy") != "steps"
+        ):
+            raise ValueError(
+                "eval_strategy and eval_steps mismatch. Please set eval_strategy to 'steps' or remove eval_steps."
+            )
+
+        if (
+            data.get("val_set_size") == 0
+            and (data.get("eval_steps") or data.get("eval_strategy"))
+            and not data.get("test_datasets")
+            and data.get("eval_strategy") != "no"
+        ):
+            raise ValueError(
+                "eval_steps and eval_strategy are not supported with val_set_size == 0"
+            )
+        if data.get("evals_per_epoch") and data.get("eval_steps"):
+            raise ValueError(
+                "eval_steps and evals_per_epoch are mutually exclusive and cannot be used together."
+            )
+        if (
+            data.get("evals_per_epoch")
+            and data.get("eval_strategy")
+            and data.get("eval_strategy") != "steps"
+        ):
+            raise ValueError(
+                "eval_strategy must be empty or set to `steps` when used with evals_per_epoch."
+            )
+
+        if data.get("do_bench_eval") and not (
+            data.get("evals_per_epoch") or data.get("eval_steps")
+        ):
+            raise ValueError(
+                "do_bench_eval requires evals_per_epoch or eval_steps to be set."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_neftune(cls, data):
+        if data.get("noisy_embedding_alpha") and not data.get("neftune_noise_alpha"):
+            data["neftune_noise_alpha"] = data["noisy_embedding_alpha"]
+            del data["noisy_embedding_alpha"]
+        elif data.get("noisy_embedding_alpha") and data.get("neftune_noise_alpha"):
+            raise ValueError(
+                "noisy_embedding_alpha is deprecated, use neftune_noise_alpha; both are set, please remove the deprecated noisy_embedding_alpha setting"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_multipack_buffer_size(cls, data):
+        if data.get("pretrain_multipack_buffer_size") and not data.get(
+            "streaming_multipack_buffer_size"
+        ):
+            LOG.warning(
+                "`pretrain_multipack_buffer_size` is deprecated in v0.13.0, will be "
+                "removed in v0.14.0. Use `streaming_multipack_buffer_size` instead."
+            )
+            data["streaming_multipack_buffer_size"] = data[
+                "pretrain_multipack_buffer_size"
+            ]
+            del data["pretrain_multipack_buffer_size"]
+        elif data.get("pretrain_multipack_buffer_size") and data.get(
+            "streaming_multipack_buffer_size"
+        ):
+            raise ValueError(
+                "pretrain_multipack_buffer_size is deprecated, use "
+                "streaming_multipack_buffer_size; both are set, please remove the "
+                "deprecated pretrain_multipack_buffer_size setting"
+            )
+        return data
+
+    @model_validator(mode="after")
+    def check_fft_possible_bad_config(self):
+        if (
+            not (self.bf16 or self.bfloat16)
+            and (self.fp16 or self.float16)
+            and not self.adapter
+            and not self.attn_uses_flash_lib
+            and self.sample_packing
+        ):
+            LOG.warning(
+                "Full fine tune w/o FA2 w/ sample packing and fp16/float16 is likely to raise errors. Try LoRA."
+            )
+            # ValueError: Attempting to unscale FP16 gradients.
+            # OR
+            # RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::Half
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fp8_config(cls, data):
+        if data.get("fp8") and not data.get("torch_compile"):
+            LOG.warning(
+                "torch_compile is strongly recommended for FP8 training in order to "
+                "see speed improvements. Please consider setting `torch_compile: "
+                "true` in your config."
+            )
+        fsdp_config = data.get("fsdp_config") or {}
+        if data.get("fp8") and (
+            fsdp_config.get("activation_checkpointing", False) is True
+            or fsdp_config.get("fsdp_activation_checkpointing", False) is True
+        ):
+            LOG.warning(
+                "FP8 + FSDP2 + activation checkpointing may be slower than BF16 "
+                "training. Please considering setting `activation_checkpointing: false` "
+                "in your FSDP config."
+            )
+        if (
+            data.get("fp8_enable_fsdp_float8_all_gather")
+            and not data.get("fsdp_version", None) == 2
+        ):
+            raise ValueError(
+                "fp8_enable_fsdp_float8_all_gather requires FSDP2 (fsdp_version: 2) "
+                "to be used."
+            )
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_use_reentrant_mismatch(cls, data):
+        if (
+            data.get("unfrozen_parameters")
+            and data.get("gradient_checkpointing_kwargs")
+            and data.get("gradient_checkpointing_kwargs", {}).get("use_reentrant")
+            is True
+        ):
+            # https://github.com/huggingface/transformers/issues/21381
+            raise ValueError(
+                "`use_reentrant` must be false when used with partially frozen model."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_eval_strategy(cls, data):
+        if (
+            data.get("evaluation_strategy") is not None
+            and data.get("eval_strategy") is None
+        ):
+            LOG.info(
+                "explicitly setting `eval_strategy` from the `evaluation_strategy`"
+            )
+            data["eval_strategy"] = data.get("evaluation_strategy")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_causal_lm_evals(cls, data):
+        if data.get("do_causal_lm_eval") and data.get("eval_sample_packing"):
+            raise ValueError(
+                "do_causal_lm_eval is enabled, eval_sample_packing must be set to False"
+            )
+
+        if data.get("eval_causal_lm_metrics"):
+            if not isinstance(data.get("eval_causal_lm_metrics"), list):
+                raise ValueError("eval_causal_lm_metrics must be a list")
+            # only ["sacrebleu", "comet", "ter", "chrf"] supported
+            if set(data.get("eval_causal_lm_metrics")) - SUPPORTED_METRICS:
+                raise ValueError(
+                    f"eval_causal_lm_metrics must be one of {SUPPORTED_METRICS}"
+                )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_tokenizer_use_mistral_common(cls, data):
+        if data.get("tokenizer_use_mistral_common") is None:
+            if any(
+                "magistral" in name.lower()
+                for name in [
+                    data.get("base_model", ""),
+                    data.get("base_model_config", ""),
+                    data.get("tokenizer_config", ""),
+                ]
+            ):
+                LOG.warning(
+                    "tokenizer_use_mistral_common auto inferred to True for Magistral models. Please set it to True explicitly if you want to use mistral-common tokenizer."
+                )
+                data["tokenizer_use_mistral_common"] = True
+
+        return data
+
+    @field_validator("tokenizer_use_mistral_common", mode="after")
+    @classmethod
+    def check_mistral_common_import(cls, tokenizer_use_mistral_common):
+        if tokenizer_use_mistral_common:
+            import importlib.util
+
+            if importlib.util.find_spec("mistral_common") is None:
+                raise ImportError(
+                    "mistral-common is required for mistral models. Please install it with `pip install axolotl` or `pip install -e .`."
+                )
+
+        return tokenizer_use_mistral_common
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_mistral_common_incompatible_options(cls, data):
+        if not data.get("tokenizer_use_mistral_common"):
+            return data
+
+        # NOTE: mistral-common tokenizer is not compatible with editing tokenizer at the moment
+
+        if data.get("added_tokens_overrides"):
+            raise ValueError(
+                "added_tokens_overrides is not supported with mistral-common tokenizer"
+            )
+
+        if data.get("special_tokens"):
+            raise ValueError(
+                "special_tokens override is not supported with mistral-common tokenizer"
+            )
+
+        if data.get("tokens"):
+            raise ValueError(
+                "tokens override is not supported with mistral-common tokenizer"
+            )
+
+        if data.get("chat_template"):
+            raise ValueError(
+                "Setting chat_template is not supported with mistral-common tokenizer"
+            )
+
+        if data.get("processor_kwargs"):
+            raise ValueError(
+                "processor_kwargs is not supported with mistral-common tokenizer"
+            )
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def pretrain_with_tps(cls, data):
+        if data.get("pretraining_dataset") and data.get(
+            "include_tokens_per_second", False
+        ):
+            # combining these would raise `TypeError: cannot pickle 'dict_keys' object`
+            # due to trying to count the number of tokens total in the dataset
+            raise ValueError(
+                "pretraining_dataset and include_tokens_per_second cannot be used together."
+            )
+
+        return data
+
+
+class LoRAValidationMixin:
+    """Validation methods related to LoRA/QLoRA configuration."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_lr_groups(cls, data):
+        if data.get("lr_groups") and data.get("loraplus_lr_ratio"):
+            raise ValueError("lr_groups and loraplus_lr_ratio cannot be used together.")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_frozen(cls, data):
+        if (
+            data.get("adapter")
+            and data.get("peft_layers_to_transform")
+            and data.get("unfrozen_parameters")
+        ):
+            raise ValueError(
+                "`unfrozen_parameters` used with `peft_layers_to_transform` can have unexpected behavior."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_peft_layers_pattern(cls, data):
+        if data.get("peft_layers_pattern") and not data.get("peft_layers_to_transform"):
+            raise ValueError(
+                "peft_layers_pattern requires peft_layers_to_transform to be set"
+            )
+        return data
+
+    @model_validator(mode="after")
+    def check_fused_lora(self):
+        if self.adapter in ["lora", "qlora"] and self.flash_attn_fuse_mlp:
+            raise ValueError("Fused modules are not supported with LoRA/QLoRA")
+        return self
+
+    @model_validator(mode="after")
+    def check_onebitllms_lora(self):
+        if self.use_onebitllms and self.adapter in ["lora", "qlora"]:
+            raise ValueError("LoRA/QLoRA is not supported with use_onebitllms")
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def warn_qlora_zero3_w_use_reentrant(cls, data):
+        if (
+            data.get("adapter") == "qlora"
+            and data.get("gradient_checkpointing_kwargs", {})
+            and data.get("gradient_checkpointing_kwargs", {}).get("use_reentrant")
+            is False
+            and data.get("deepspeed", "") is not None
+            and "zero3" in data.get("deepspeed", "")
+        ):
+            # may result in:
+            # torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint:
+            # Recomputed values for the following tensors have different metadata
+            # than during the forward pass.
+            LOG.warning(
+                "qlora + zero3 with use_reentrant: false may result in a CheckpointError about recomputed values"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_lora_kernels_8bit(cls, data):
+        if (
+            data.get("lora_mlp_kernel")
+            or data.get("lora_qkv_kernel")
+            or data.get("lora_o_kernel")
+        ):
+            if data.get("adapter") == "lora" and data.get("load_in_8bit"):
+                raise ValueError(
+                    "lora_mlp_kernel, lora_qkv_kernel, and lora_o_kernel are not "
+                    "compatible with 8-bit LoRA a the moment."
+                )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_lora_kernels_dora(cls, data):
+        # DoRA is now supported by lora kernels
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_lora_kernels_trust_remote_code(cls, data):
+        if (
+            data.get("lora_mlp_kernel")
+            or data.get("lora_qkv_kernel")
+            or data.get("lora_o_kernel")
+        ) and data.get("trust_remote_code"):
+            raise ValueError(
+                "lora_mlp_kernel, lora_qkv_kernel, and lora_o_kernel are not "
+                "compatible with trust_remote_code. Please disable trust_remote_code "
+                "or explicitly set lora_*_kernel to false."
+            )
+        return data
+
+
+class RLValidationMixin:
+    """Validation methods related to RL training configuration."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_sample_packing_w_rl(cls, data):
+        if data.get("sample_packing") and data.get("rl"):
+            raise ValueError("`sample_packing: true` does not work with RLHF training")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_kto_config(cls, data):
+        if data.get("rl") == "kto":
+            if data.get("sample_packing") or data.get("eval_sample_packing"):
+                raise ValueError("sample_packing is not supported with kto")
+
+            if data.get("remove_unused_columns") is not False:
+                raise ValueError("Set `remove_unused_columns: False` when using kto")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_grpo_liger_sequence_parallel(cls, data):
+        if (
+            data.get("rl") == "grpo"
+            and data.get("trl", {})
+            and data.get("trl").get("use_liger_loss")
+            and data.get("context_parallel_size", 1) > 1
+        ):
+            raise ValueError("GRPO + SP + Liger not currently supported")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_rl_config_gradient_checkpointing(cls, data):
+        # TODO: SalmanMohammadi
+        # Distributed RL with QLoRA + gradient checkpointing
+        # and use_reentrant = True is broken upstream in TRL
+
+        if (
+            data.get("rl")
+            and data.get("gradient_checkpointing")
+            and data.get("gradient_checkpointing_kwargs")
+            and data.get("gradient_checkpointing_kwargs").get("use_reentrant")
+            and data.get("load_in_4bit")
+            and data.get("adapter") == "qlora"
+            and data.get("capabilities")
+            and data.get("capabilities").get("n_gpu", 1) > 1
+        ):
+            raise ValueError(
+                "The `use_reentrant: True` implementation of gradient checkpointing "
+                "is not supported for distributed RL training with QLoRA. Please set "
+                "`use_reentrant: False` in `gradient_checkpointing_kwargs`."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_gdpo(cls, data):
+        if (
+            data.get("rl") == "gdpo"
+            and data.get("trl", {}).get("multi_objective_aggregation")
+            == "sum_then_normalize"
+        ):
+            raise ValueError(
+                "`multi_objective_aggregation` value set as `sum_then_normalize` => GRPO, but GDPO was selected"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_dpo(cls, data):
+        dpo_loss_type = data.get("dpo_loss_type")
+        dpo_loss_weights = data.get("dpo_loss_weights")
+        rl = data.get("rl")
+
+        if rl == "ipo":
+            LOG.warning(
+                "rl: ipo will soon be deprecated. Use `rl: dpo` with `dpo_loss_type: ['ipo']` instead."
+            )
+
+        if rl == "dpo":
+            if dpo_loss_weights is not None and dpo_loss_type is None:
+                raise ValueError(
+                    "`dpo_loss_weights` requires `dpo_loss_type` to be set"
+                )
+            if (
+                dpo_loss_type is not None
+                and dpo_loss_weights is not None
+                and len(dpo_loss_type) != len(dpo_loss_weights)
+            ):
+                raise ValueError(
+                    f"`dpo_loss_type` and `dpo_loss_weights` must be the same length, "
+                    f"but got {len(dpo_loss_type)} losses and {len(dpo_loss_weights)} weights"
+                )
+        elif dpo_loss_type is not None or dpo_loss_weights is not None:
+            raise ValueError(
+                f"`dpo_loss_type` and `dpo_loss_weights` are for DPO only,"
+                f"but got {rl=}, {dpo_loss_type=} and {dpo_loss_weights=}"
+            )
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_grpo_batch_size_divisibility(cls, data):
+        """Surface GRPO batch-shape mismatches at config-parse time.
+
+        TRL's GRPOTrainer requires that the per-step generation batch size be
+        evenly divisible by ``num_generations`` so that every prompt can be
+        replicated exactly ``num_generations`` times. The runtime check inside
+        ``GRPOTrainer.__init__`` only fires after the model has been loaded —
+        too late and too cryptic for the user. We replicate the check here so
+        the failure is immediate and actionable.
+
+        Also enforces:
+          - ``num_generations >= 2`` (group-relative advantage needs variance)
+          - ``effective_gbs >= num_generations * world_size`` when capabilities
+            indicate multiple ranks (each rank needs at least one full group)
+        """
+        if data.get("rl") != "grpo":
+            return data
+
+        trl_cfg = data.get("trl") or {}
+        num_gen = trl_cfg.get("num_generations")
+        if num_gen is None:
+            # TRL's own default is 8 — but if the user didn't set it, we
+            # don't have enough info to validate anything. Let TRL's own
+            # init handle the default-vs-batch interaction.
+            return data
+        if num_gen < 2:
+            raise ValueError(
+                f"GRPO requires `trl.num_generations >= 2` (got {num_gen}). "
+                "With num_generations=1, every group has zero advantage and "
+                "the policy never updates."
+            )
+
+        explicit_gbs = trl_cfg.get("generation_batch_size")
+        if explicit_gbs is not None:
+            effective_gbs = int(explicit_gbs)
+            gbs_source = "trl.generation_batch_size"
+        else:
+            mb = data.get("micro_batch_size") or 1
+            ga = data.get("gradient_accumulation_steps") or 1
+            effective_gbs = int(mb) * int(ga)
+            gbs_source = f"micro_batch_size ({mb}) * gradient_accumulation_steps ({ga})"
+
+        if effective_gbs % num_gen != 0:
+            # Suggest the smallest GA bump that fixes it for the common case
+            # where the user hasn't set generation_batch_size explicitly.
+            hint = ""
+            if explicit_gbs is None:
+                from math import gcd
+
+                mb_val = int(data.get("micro_batch_size") or 1)
+                # smallest GA such that mb*GA is a multiple of num_gen
+                lcm = num_gen * mb_val // gcd(num_gen, mb_val)
+                suggested_ga = lcm // mb_val
+                hint = (
+                    f" Smallest fix: set `gradient_accumulation_steps: "
+                    f"{suggested_ga}` (so micro_batch_size * GA = "
+                    f"{mb_val * suggested_ga} is a multiple of {num_gen})."
+                )
+            raise ValueError(
+                f"GRPO: generation batch size must be divisible by "
+                f"`trl.num_generations`. Got effective_gbs={effective_gbs} "
+                f"(from {gbs_source}) and num_generations={num_gen}.{hint}"
+            )
+
+        # Multi-rank check: each rank must receive at least one full group
+        # per step. Without `capabilities` populated yet (mode='before'), we
+        # fall back to user-set distributed fields.
+        world_size = (
+            (data.get("capabilities") or {}).get("n_gpu") or data.get("world_size") or 1
+        )
+        if world_size and world_size > 1 and effective_gbs < num_gen * world_size:
+            raise ValueError(
+                f"GRPO with world_size={world_size} requires effective_gbs "
+                f">= num_generations * world_size = {num_gen * world_size}, "
+                f"got {effective_gbs}. Increase gradient_accumulation_steps "
+                f"or micro_batch_size."
+            )
+
+        return data
+
+
+class OptimizationValidationMixin:
+    """Validation methods related to optimization and performance."""
+
+    @model_validator(mode="after")
+    def check_adamw_optimizer_params(self):
+        if any([self.adam_beta1, self.adam_beta2, self.adam_epsilon]) and (
+            not self.optimizer or "adamw" not in str(self.optimizer).lower()
+        ):
+            LOG.warning("adamw hyperparameters found, but no adamw optimizer set")
+        return self
+
+    @staticmethod
+    def _resolve_fsdp_version(data):
+        """Resolve FSDP version from top-level fsdp_version or fsdp_config.fsdp_version."""
+        fsdp_version = data.get("fsdp_version")
+        if fsdp_version is None:
+            fsdp_version = data.get("fsdp_config", {}).get("fsdp_version", 1)
+        return fsdp_version
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_muon_deepspeed_fsdp(cls, data):
+        if data.get("optimizer") == "muon":
+            if data.get("deepspeed"):
+                raise ValueError(
+                    "Muon optimizer is currently incompatible with DeepSpeed"
+                )
+            if data.get("fsdp") or data.get("fsdp_config"):
+                fsdp_version = cls._resolve_fsdp_version(data)
+                if str(fsdp_version) != "2":
+                    raise ValueError(
+                        "Muon optimizer is only compatible with FSDP2. Set fsdp_version: 2 to use Muon with FSDP."
+                    )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_qgalore(cls, data):
+        if data.get("optimizer") != "q_galore_adamw8bit":
+            return data
+        adapter = data.get("adapter")
+        if adapter:
+            raise ValueError(
+                "q_galore_adamw8bit operates on full-precision parameters and is "
+                f"incompatible with adapter='{adapter}'. Remove the adapter setting "
+                "or pick a different optimizer."
+            )
+        if data.get("deepspeed"):
+            raise ValueError(
+                "q_galore_adamw8bit is not yet validated with DeepSpeed. "
+                "Use DDP or FSDP2 with use_orig_params=True."
+            )
+        if data.get("fsdp") or data.get("fsdp_config"):
+            fsdp_version = cls._resolve_fsdp_version(data)
+            if str(fsdp_version) != "2":
+                raise ValueError(
+                    "q_galore_adamw8bit requires FSDP2. Set fsdp_version: 2."
+                )
+            fsdp_config = data.get("fsdp_config") or {}
+            if fsdp_config.get("use_orig_params") is not True:
+                raise ValueError(
+                    "q_galore_adamw8bit requires fsdp_config.use_orig_params=True so "
+                    "that per-parameter projection state survives FSDP sharding."
+                )
+        if not (data.get("bf16") or data.get("bfloat16") or data.get("fp16")):
+            LOG.warning(
+                "q_galore_adamw8bit benefits from mixed-precision (bf16/fp16). "
+                "Running in fp32 will negate most of the memory savings."
+            )
+        if data.get("optim_target_modules") is None:
+            # Match the reference impl's defaults: attention + MLP linears.
+            data["optim_target_modules"] = [
+                "attn",
+                "mlp",
+            ]
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_flashoptim_deepspeed_fsdp(cls, data):
+        optimizer = data.get("optimizer") or ""
+        if str(optimizer).startswith("flash_"):
+            if data.get("deepspeed"):
+                raise ValueError(
+                    f"{optimizer} optimizer is incompatible with DeepSpeed. "
+                    "Flash optimizers only support DDP and FSDP2."
+                )
+            if data.get("fsdp") or data.get("fsdp_config"):
+                fsdp_version = cls._resolve_fsdp_version(data)
+                if str(fsdp_version) != "2":
+                    raise ValueError(
+                        f"{optimizer} optimizer is only compatible with FSDP2. "
+                        "Set fsdp_version: 2 to use flash optimizers with FSDP."
+                    )
+        return data
+
+    @model_validator(mode="after")
+    def check_batch_flattening_fa(self):
+        if not self.batch_flattening:
+            return self
+
+        batch_flattening_auto = self.batch_flattening == "auto"
+        has_varlen_attn = self.attn_supports_packing
+
+        if not has_varlen_attn and not batch_flattening_auto:
+            raise ValueError(
+                "batch_flattening requires a varlen-capable attention backend "
+                "(e.g., attn_implementation: flash_attention_2)."
+            )
+        if self.sample_packing and not batch_flattening_auto:
+            raise ValueError("batch_flattening not compatible with sample_packing")
+        if self.micro_batch_size == 1 and not batch_flattening_auto:
+            LOG.warning("batch_flattening has no effect with micro_batch_size == 1")
+
+        # Liger loss takes a separate code path (compute_liger_loss) that
+        # bypasses the flattened training forward pass. Batch flattening
+        # still applies to the scoring/deferred logprobs path.
+        if self.trl and getattr(self.trl, "use_liger_loss", False):
+            LOG.warning(
+                "batch_flattening with use_liger_loss: flattening will only "
+                "apply to the scoring path (deferred logprobs). The training "
+                "forward pass uses Liger's fused lm_head+loss kernel instead."
+            )
+
+        if (
+            batch_flattening_auto
+            and has_varlen_attn
+            and not self.sample_packing
+            and self.micro_batch_size > 1
+        ):
+            self.batch_flattening = True
+        elif batch_flattening_auto:
+            self.batch_flattening = False
+
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_cross_entropy_conflicts(cls, data):
+        """Check for mutual exclusivity between cross entropy patch options.
+
+        Only one of the following can be enabled at a time:
+        - cut_cross_entropy (CutCrossEntropyPlugin)
+        - chunked_cross_entropy
+        - liger_cross_entropy (LigerPlugin)
+        - liger_fused_linear_cross_entropy (LigerPlugin)
+        """
+        ce_options = {
+            "cut_cross_entropy": data.get("cut_cross_entropy"),
+            "chunked_cross_entropy": data.get("chunked_cross_entropy"),
+            "liger_cross_entropy": data.get("liger_cross_entropy"),
+            "liger_fused_linear_cross_entropy": data.get(
+                "liger_fused_linear_cross_entropy"
+            ),
+        }
+
+        enabled_options = [k for k, v in ce_options.items() if v]
+
+        if len(enabled_options) > 1:
+            raise ValueError(
+                f"Only one cross entropy optimization can be enabled at a time. "
+                f"Found {len(enabled_options)} enabled: {', '.join(enabled_options)}. "
+                "Please disable all but one."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fsdp_version(cls, data):
+        fsdp_config = data.get("fsdp_config", {})
+        if fsdp_config and str(data.get("fsdp_version")) != "2":
+            LOG.warning(
+                "FSDP1 is deprecated and will be removed in an upcoming release of "
+                "Axolotl (transformers plans to in v5.20). We recommend migrating "
+                "to fsdp_version: 2 for better performance and compatibility. "
+                "See https://docs.axolotl.ai/docs/multi-gpu.html#sec-fsdp for "
+                "details on migrating your config."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fsdp2_cpu_offload_pin_memory(cls, data):
+        if not (fsdp_config := data.get("fsdp_config")):
+            return data
+
+        if fsdp_config.get("cpu_offload_pin_memory") is False:
+            if str(data.get("fsdp_version")) != "2":
+                raise ValueError(
+                    "FSDP1 does not support disabling cpu_offload_pin_memory, please set `fsdp_version` to 2"
+                )
+            if not fsdp_config.get("offload_params"):
+                raise ValueError(
+                    "disabling cpu_offload_pin_memory requires enabling offload_params"
+                )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fsdp2_base_model_quant_rl(cls, data):
+        if data.get("fsdp_version") == 2 and data.get("rl") in [
+            RLType.DPO,
+            RLType.KTO,
+            RLType.ORPO,
+            RLType.IPO,
+        ]:
+            if data.get("load_in_8bit") or data.get("load_in_4bit"):
+                raise ValueError(
+                    f"FSDP2 does not support load_in_8bit or load_in_4bit with {data.get('rl')}. Please use DeepSpeed or set `fsdp_version` to 1."
+                )
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fsdp_config_kwargs_prefix(cls, data):
+        if fsdp_config := data.get("fsdp_config"):
+            should_fix = False
+            for key, _ in fsdp_config.items():
+                if key.startswith("fsdp_") and key != "fsdp_version":
+                    should_fix = True
+                    LOG.warning_once(
+                        "Configuring FSDP fields with the `fsdp_` prefix is deprecated. "
+                        "Please omit the `fsdp_` prefix from the any fields in `fsdp_config`."
+                    )
+            if should_fix:
+                update_fsdp_config = {}
+                for key, value in fsdp_config.items():
+                    if key.startswith("fsdp_") and key != "fsdp_version":
+                        update_fsdp_config[key.replace("fsdp_", "")] = value
+                    else:
+                        update_fsdp_config[key] = value
+                data["fsdp_config"] = update_fsdp_config
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fsdp_version_in_fsdp_config(cls, data):
+        fsdp_config = data.get("fsdp_config") or {}
+        fsdp_version = data.get("fsdp_version", None)
+        if not fsdp_version and fsdp_config and fsdp_config.get("version"):
+            fsdp_cfg_version = fsdp_config.pop("version")
+            data["fsdp_version"] = fsdp_cfg_version
+            data["fsdp_config"]["fsdp_version"] = fsdp_cfg_version
+        elif not fsdp_version and fsdp_config and fsdp_config.get("fsdp_version"):
+            data["fsdp_version"] = fsdp_config.get("fsdp_version")
+        if fsdp_version and fsdp_config and not fsdp_config.get("fsdp_version"):
+            data["fsdp_config"]["fsdp_version"] = fsdp_version
+        if fsdp_config and not data.get("fsdp_version"):
+            # transformers >= 5.10 defaults a missing version to FSDP2; pin
+            # axolotl's FSDP1 default explicitly so unversioned configs keep
+            # their behavior
+            data["fsdp_version"] = 1
+            data["fsdp_config"]["fsdp_version"] = 1
+        return data
+
+    @model_validator(mode="after")
+    def check_fsdp_offload_w_8bit_optimizer(self):
+        if (
+            hasattr(self, "fsdp_config")
+            and self.fsdp_config
+            and self.optimizer
+            and "8bit" in str(self.optimizer)
+            and self.fsdp_config.offload_params
+            and str(self.fsdp_version) != "2"
+        ):
+            raise ValueError(f"FSDP Offload not compatible with {str(self.optimizer)}")
+        return self
+
+    @model_validator(mode="after")
+    def check_fsdp2_w_8bit_optimizer(self):
+        if (
+            hasattr(self, "fsdp_config")
+            and self.fsdp_config
+            and self.optimizer
+            and "8bit" in str(self.optimizer)
+            and str(self.fsdp_version) == "2"
+        ):
+            if self.optimizer in ["adamw_8bit", "adamw_bnb_8bit"]:
+                # CUDA ops errors with bnb 8bit optimizer + FSDP2
+                raise ValueError(
+                    f"FSDP2 not compatible with {self.optimizer}, use `adamw_torch_8bit` instead"
+                )
+
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_tensor_parallel_size_update_ds_json(cls, data):
+        tensor_parallel_size = data.get("tensor_parallel_size")
+        if tensor_parallel_size is not None and tensor_parallel_size > 1:
+            if data.get("deepspeed"):
+                with open(data.get("deepspeed"), "r", encoding="utf-8") as ds_fin:
+                    ds_config = json.load(ds_fin)
+                    should_save = False
+                    if "tensor_parallel" not in ds_config:
+                        ds_config["tensor_parallel"] = {
+                            "autotp_size": tensor_parallel_size
+                        }
+                        should_save = True
+                    if (
+                        "gather_16bit_weights_on_model_save"
+                        not in ds_config["zero_optimization"]
+                    ):
+                        ds_config["zero_optimization"][
+                            "gather_16bit_weights_on_model_save"
+                        ] = True
+                        should_save = True
+                    if should_save:
+                        temp_dir = tempfile.mkdtemp()
+                        with open(
+                            Path(temp_dir) / "autotp_ds.json", "w", encoding="utf-8"
+                        ) as ds_fout:
+                            json.dump(ds_config, ds_fout, indent=4)
+                        data["deepspeed"] = str(Path(temp_dir) / "autotp_ds.json")
+
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_deepcompile(cls, data):
+        deepcompile = data.get("deepcompile")
+        if deepcompile:
+            if not data.get("deepspeed"):
+                raise ValueError("DeepCompile is only supported with DeepSpeed")
+            with open(data.get("deepspeed"), "r", encoding="utf-8") as ds_fin:
+                ds_config = json.load(ds_fin)
+                if "compile" not in ds_config:
+                    ds_config["compile"] = {"deepcompile": True}
+                    temp_dir = tempfile.mkdtemp()
+                    with open(
+                        Path(temp_dir) / "deepcompile_ds.json", "w", encoding="utf-8"
+                    ) as ds_fout:
+                        json.dump(ds_config, ds_fout, indent=4)
+                    data["deepspeed"] = str(Path(temp_dir) / "deepcompile_ds.json")
+
+        return data
+
+
+class SystemValidationMixin:
+    """Validation methods related to system and hardware configuration."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_mem_mismatch(cls, data):
+        if (
+            data.get("max_memory") is not None
+            and data.get("gpu_memory_limit") is not None
+        ):
+            raise ValueError(
+                "max_memory and gpu_memory_limit are mutually exclusive and cannot be used together."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_fsdp_deepspeed(cls, data):
+        if data.get("deepspeed") and data.get("fsdp"):
+            raise ValueError("deepspeed and fsdp cannot be used together.")
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_model_quantization_config_vs_bnb(cls, data):
+        if data.get("model_quantization_config"):
+            if data.get("load_in_8bit") or data.get("load_in_4bit"):
+                raise ValueError(
+                    "model_quantization_config and load_in_8bit or load_in_4bit cannot be used together."
+                )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_npu_config(cls, data):
+        if is_torch_npu_available():
+            # check attention config
+            unsupported_npu_impls = {
+                "flash_attention_2",
+                "flash_attention_3",
+                "sdpa",
+            }
+            attn_impl = data.get("attn_implementation")
+            if attn_impl and attn_impl in unsupported_npu_impls:
+                raise NotImplementedError(
+                    f"attn_implementation={attn_impl!r} is currently not supported on Ascend NPU."
+                )
+            # Legacy flags still present at this point (normalizer strips them later).
+            attn_list = ["flash_attention", "sdp_attention"]
+            for attn in attn_list:
+                if data.get(attn):
+                    raise NotImplementedError(
+                        f"{attn} is currently not supported in Ascend npu, please disable this configuration."
+                    )
+
+            # check quant config
+            if data.get("optimizer") is not None and "bit" in data.get("optimizer"):
+                optimizer = data.get("optimizer")
+                raise NotImplementedError(
+                    f"{optimizer} is currently not supported in Ascend npu, choose another one please."
+                )
+
+            quant_list = ["load_in_8bit", "load_in_4bit"]
+            for quant in quant_list:
+                if data.get(quant):
+                    raise NotImplementedError(
+                        f"Quantification is currently not supported in Ascend npu, please disable {quant}."
+                    )
+
+            # check dtype config
+            if data.get("tf32"):
+                raise NotImplementedError(
+                    "tf32 dtype is currently not supported in Ascend npu, please disable this configuration"
+                )
+
+        return data
+
+
+class ChatTemplateValidationMixin:
+    """Validation methods related to chat template configuration."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_chat_template_config(cls, data):
+        # if chat_template is set to jinja, chat_template_jinja is required
+        if data.get("chat_template") == ChatTemplate.jinja and not data.get(
+            "chat_template_jinja"
+        ):
+            raise ValueError(
+                "chat_template_jinja is required when chat_template is set to 'jinja'. "
+                "Please provide the Jinja template string in chat_template_jinja field."
+            )
+
+        # If chat_template_jinja is set, set chat_template to jinja
+        if data.get("chat_template_jinja") and not data.get("chat_template"):
+            data["chat_template"] = ChatTemplate.jinja
+
+        return data
+
+
+class PretrainingValidationMixin:
+    """Validation methods related to pretraining configuration."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_pretraining_w_max_steps(cls, data):
+        if data.get("pretraining_dataset") and not data.get("max_steps"):
+            raise ValueError(
+                "max_steps must be set when using iterable pretraining_dataset, Trainer can't infer length and schedule optimizer/learning rate without it!"
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_pretraining_w_group_by_length(cls, data):
+        if data.get("pretraining_dataset") and data.get("group_by_length"):
+            LOG.warning(
+                "You probably want to disable group_by_length as it will force a streamed dataset to download completely."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_pretraining_split_batches_accelerate(cls, data):
+        # alternatively set ACCELERATE_SPLIT_BATCHES=False
+        if data.get("pretraining_dataset"):
+            accelerator_config = data.get("accelerator_config", {})
+            if not accelerator_config:
+                data["accelerator_config"] = {
+                    "split_batches": False,
+                    "dispatch_batches": False,
+                }
+            else:
+                if accelerator_config.get("split_batches") is None:
+                    data["accelerator_config"]["split_batches"] = False
+                if accelerator_config.get("dispatch_batches") is None:
+                    data["accelerator_config"]["dispatch_batches"] = False
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_pretraining_w_val_set_size(cls, data):
+        if data.get("pretraining_dataset") and data.get("val_set_size"):
+            raise ValueError(
+                "val_set_size is not supported with pretraining_dataset. "
+                "Use test_datasets to specify evaluation datasets for pretraining."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_streaming_w_val_set_size(cls, data):
+        if data.get("streaming") and data.get("val_set_size"):
+            raise ValueError(
+                "val_set_size is not supported with streaming datasets. "
+                "Use test_datasets to specify evaluation datasets when streaming is enabled."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_streaming_w_max_steps(cls, data):
+        if data.get("streaming") and not data.get("max_steps"):
+            raise ValueError(
+                "max_steps must be set when using streaming datasets. "
+                "Trainer cannot infer dataset length for iterable datasets."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_streaming_w_multiple_datasets(cls, data):
+        if (
+            data.get("streaming")
+            and data.get("sample_packing")
+            and data.get("datasets")
+            and len(data.get("datasets")) > 1
+        ):
+            raise NotImplementedError(
+                "Sample packing with multiple streaming datasets is not yet supported"
+            )
+        return data
+
+
+class ModelCompatibilityValidationMixin:
+    """Validation methods for specific model compatibility."""
+
+    @model_validator(mode="after")
+    def check_falcon_fsdp(self):
+        if (self.base_model and "falcon" in self.base_model.lower()) and self.fsdp:
+            raise ValueError("FSDP is not supported for falcon models")
+        return self
+
+    @model_validator(mode="after")
+    def check_mpt_checkpointing(self):
+        if (
+            self.base_model and "mpt" in self.base_model.lower()
+        ) and self.gradient_checkpointing:
+            raise ValueError("gradient_checkpointing is not supported for MPT models")
+        return self
+
+    @model_validator(mode="after")
+    def check_nemotron_h_gradient_checkpointing(self):
+        if (
+            self.base_model
+            and "nemotron-h" in self.base_model.lower()
+            and self.gradient_checkpointing
+            and not self.sample_packing
+        ):
+            raise ValueError(
+                "gradient_checkpointing for nemotron_h requires sample_packing: true. "
+                "The upstream model marks supports_gradient_checkpointing=False; "
+                "axolotl only enables it after applying the sample-packing patch."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_gradient_checkpointing_w_offload(self):
+        if self.gradient_checkpointing == "offload":
+            LOG.warning(
+                "`offload` is deprecated for gradient_checkpointing, use `activation_offloading: true` or `activation_offloading: legacy`"
+            )
+            self.gradient_checkpointing = True
+            LOG.warning(
+                "`offload` now uses a new stream implementation; to use the previous implementation, use `activation_offloading: legacy`"
+            )
+            self.activation_offloading = True
+        if self.gradient_checkpointing == "offload_disk":
+            LOG.warning(
+                "`offload_disk` is deprecated for gradient_checkpointing, use `activation_offloading: disk`"
+            )
+            self.gradient_checkpointing = True
+            self.activation_offloading = "disk"
+        return self
+
+    @model_validator(mode="after")
+    def check_activation_offloading_wo_gc(self):
+        if self.activation_offloading and not self.gradient_checkpointing:
+            raise ValueError("activation_offloading requires gradient_checkpointing")
+        return self
+
+    @model_validator(mode="after")
+    def check_selective_checkpointing(self):
+        if not self.selective_checkpointing:
+            self.selective_checkpointing = None
+            return self
+        if self.selective_checkpointing is True:
+            from axolotl.utils.schemas.training import SelectiveCheckpointingConfig
+
+            self.selective_checkpointing = SelectiveCheckpointingConfig()
+        if not self.gradient_checkpointing:
+            raise ValueError("selective_checkpointing requires gradient_checkpointing")
+        if (self.gradient_checkpointing_kwargs or {}).get("use_reentrant"):
+            raise ValueError(
+                "selective_checkpointing requires non-reentrant checkpointing "
+                "(gradient_checkpointing_kwargs.use_reentrant: false)"
+            )
+        if self.activation_offloading and self.activation_offloading != "hidden_states":
+            raise ValueError(
+                "selective_checkpointing is only compatible with "
+                "activation_offloading: hidden_states (or false); the TRL offloader "
+                "paths bypass HF gradient checkpointing"
+            )
+        if self.adapter:
+            # PEFT adds the adapter delta into the base linear output in-place,
+            # mutating cached mm outputs; SAC's cache-mutation guard errors at runtime
+            matmul_ops = ("aten::mm", "aten::addmm", "aten::matmul", "aten::bmm")
+            bad = [
+                spec
+                for spec in self.selective_checkpointing.save
+                if spec != "attention" and any(spec in op for op in matmul_ops)
+            ]
+            if bad:
+                raise ValueError(
+                    f"selective_checkpointing.save entries {bad} match matmul ops, "
+                    "which cannot be saved when training with a LoRA/QLoRA adapter: "
+                    "PEFT mutates the base linear output in-place to add the adapter "
+                    "delta, which invalidates the cached tensor. Use save: [attention] "
+                    "or train without an adapter."
+                )
+        if self.torch_compile:
+            LOG.warning(
+                "selective_checkpointing with torch_compile is untested in axolotl; "
+                "the SAC context_fn targets the eager checkpointing path"
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_hidden_states_offloading(self):
+        if self.activation_offloading != "hidden_states":
+            return self
+        gc_kwargs = dict(self.gradient_checkpointing_kwargs or {})
+        use_reentrant = gc_kwargs.setdefault("use_reentrant", False)
+        self.gradient_checkpointing_kwargs = gc_kwargs
+        if use_reentrant and self.unfrozen_parameters:
+            raise ValueError(
+                "activation_offloading: hidden_states with reentrant checkpointing "
+                "is incompatible with `unfrozen_parameters` (partially frozen model)."
+            )
+        if self.adapter:
+            LOG.warning(
+                "activation_offloading: hidden_states uses gradient checkpointing; "
+                "with LoRA/QLoRA, activation_offloading: true can be faster because "
+                "it avoids recompute while offloading smaller adapter activations."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_better_transformers(self):
+        if self.flash_optimum is True:
+            if self.adapter:
+                LOG.warning(
+                    "BetterTransformers probably doesn't work with PEFT adapters"
+                )
+            if self.fp16 or self.bf16:
+                raise ValueError("AMP is not supported with BetterTransformer")
+            if self.float16 is not True and self.bfloat16 is not True:
+                LOG.warning(
+                    "You should probably set bfloat16 or float16 to true to "
+                    "load the model in float16 for BetterTransformers"
+                )
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_gptq_w_revision(cls, data):
+        if data.get("gptq") and data.get("revision_of_model"):
+            raise ValueError(
+                "revision_of_model is not supported for GPTQ models. "
+                + "Please download the model from HuggingFace Hub manually for correct branch, "
+                + "point to its path, and remove revision_of_model from the config."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_gpt_oss_fsdp_loading(cls, data):
+        if data.get("model_quantization_config", "") == "Mxfp4Config":
+            fsdp_config = data.get("fsdp_config") or {}
+            if fsdp_config.get("cpu_ram_efficient_loading", False) is True:
+                raise ValueError(
+                    "FSDP cpu_ram_efficient_loading is not supported for Mxfp4Config model quantization."
+                )
+        return data
+
+
+class ComplexValidationMixin:
+    """Complex validation methods that involve multiple systems."""
+
+    @field_validator("neftune_noise_alpha")
+    @classmethod
+    def validate_neftune_noise_alpha(cls, neftune_noise_alpha):
+        if neftune_noise_alpha is not None and neftune_noise_alpha <= 0.0:
+            raise ValueError("neftune_noise_alpha must be > 0.0")
+        return neftune_noise_alpha
+
+    @model_validator(mode="after")
+    def check_rl_beta(self):
+        if self.dpo_beta and not self.rl_beta:
+            self.rl_beta = self.dpo_beta
+            del self.dpo_beta
+        return self
+
+    @model_validator(mode="after")
+    def check_simpo_warmup(self):
+        if self.rl is RLType.SIMPO and self.warmup_ratio:
+            raise ValueError(
+                "warmup_ratio is not supported with the simpo trainer. Please use `warmup_steps` instead"
+            )
+        return self
+
+    @model_validator(mode="after")
+    def check_relora(self):
+        if self.relora:
+            if not self.jagged_restart_steps:
+                raise ValueError("jagged_restart_steps must be set to use ReLoRA")
+
+            adapter_supports_relora = self.adapter in ("lora", "qlora")
+            if self.adapter and not adapter_supports_relora:
+                from axolotl.integrations.base import PluginManager
+
+                plugin_manager = PluginManager.get_instance()
+                adapter_supports_relora = plugin_manager.adapter_supports_relora(
+                    self.adapter
+                )
+            if not adapter_supports_relora:
+                raise ValueError(
+                    "cfg.adapter must support ReLoRA to use ReLoRA restart semantics"
+                )
+
+            if self.fsdp or self.fsdp_config:
+                raise ValueError("fsdp not supported with ReLoRA")
+
+            if self.deepspeed:
+                raise ValueError("deepspeed not supported with ReLoRA")
+
+            if self.lr_scheduler == "one_cycle":
+                raise ValueError(
+                    "ReLoRA is not compatible with the one_cycle scheduler"
+                )
+
+            if self.flash_attn_fuse_mlp:
+                raise ValueError("Fused modules are not supported with ReLoRA")
+        return self
+
+    @model_validator(mode="after")
+    def check_early_stopping(self):
+        if self.early_stopping_patience:
+            if not self.save_steps or not self.eval_steps:
+                raise ValueError(
+                    "`early_stopping_patience` requires save_steps and eval_steps to be set. eval_steps should evenly divide save_steps."
+                )
+            if self.save_steps % self.eval_steps != 0:
+                raise ValueError(
+                    "`early_stopping_patience` requires that eval_steps should evenly divide save_steps."
+                )
+        return self
+
+    @model_validator(mode="after")
+    def check_tensor_parallel_size(self):
+        if not self.tensor_parallel_size:
+            self.tensor_parallel_size = 1
+        return self
+
+    @model_validator(mode="after")
+    def check_context_parallel_size(self):
+        if self.sequence_parallel_degree and not self.context_parallel_size:
+            LOG.warning(
+                "`sequence_parallel_degree` is deprecated, use `context_parallel_size`"
+            )
+            self.context_parallel_size = self.sequence_parallel_degree
+        if not self.context_parallel_size:
+            self.context_parallel_size = 1
+        elif self.context_parallel_size > 1 and getattr(
+            self, "use_glm_dsa_kernels", False
+        ):
+            # The GLM DSA kernels provide their own context-parallel attention (the sequence is sharded
+            # on the cp axis with a compressed-KV all-gather + per-rank q_offset), so the flash /
+            # ring_flash_attn stack the generic CP path below requires does not apply.
+            LOG.warning(
+                "context_parallel_size > 1 with use_glm_dsa_kernels: the DSA kernels handle context "
+                "parallelism (compressed-KV all-gather); skipping the flash/ring-attention requirement."
+            )
+        elif self.context_parallel_size > 1:
+            if not self.attn_uses_flash_lib:
+                raise ValueError(
+                    "context_parallel_size > 1 requires flash attention "
+                    "(attn_implementation: flash_attention_2 or flash_attention_3)."
+                )
+
+            if self.sample_packing and self.micro_batch_size > 1:
+                raise ValueError(
+                    "micro_batch_size must be set to 1 when sample_packing is enabled "
+                    "due to a `ring-flash-attn` requirement"
+                )
+
+            try:
+                import transformers.modeling_flash_attention_utils
+                from transformers.utils import (
+                    is_flash_attn_greater_or_equal,
+                    is_flash_attn_greater_or_equal_2_10,
+                )
+
+                transformers.modeling_flash_attention_utils._flash_supports_window = (
+                    True
+                )
+                sys.modules[
+                    "transformers.modeling_flash_attention_utils"
+                ]._flash_supports_window = True
+                sys.modules[
+                    "transformers.modeling_flash_attention_utils"
+                ]._flash_supports_window_size = True
+                sys.modules[
+                    "transformers.modeling_flash_attention_utils"
+                ].is_flash_attn_greater_or_equal = is_flash_attn_greater_or_equal
+                if not hasattr(
+                    transformers.modeling_flash_attention_utils,
+                    "is_flash_attn_greater_or_equal_2_10",
+                ):
+                    transformers.modeling_flash_attention_utils.is_flash_attn_greater_or_equal_2_10 = is_flash_attn_greater_or_equal(
+                        "2.10"
+                    )
+                sys.modules[
+                    "transformers.modeling_flash_attention_utils"
+                ].is_flash_attn_greater_or_equal_2_10 = (
+                    is_flash_attn_greater_or_equal_2_10
+                )
+                import ring_flash_attn  # noqa: F401  # Required after monkey-patching
+            except ImportError as exception:
+                raise ImportError(
+                    "context_parallel_size > 1 but ring_flash_attn is not installed. "
+                    "Please install it with `pip install axolotl[ring-flash-attn] "
+                    "or `pip install ring-flash-attn>=0.1.4`."
+                ) from exception
+
+            LOG.warning(
+                "Sequence parallelism (SP) is enabled with "
+                f"context_parallel_size={self.context_parallel_size}. "
+                "Please note that logged losses may differ slightly to the non-SP "
+                "losses due to transformers Trainer implementation details. "
+                "Please see https://github.com/axolotl-ai-cloud/axolotl/pull/2495#issuecomment-2784022042 "
+                "for more details."
+            )
+
+            _SSM_HYBRID_MODEL_TYPES = {
+                "nemotron_h",
+                "falcon_h1",
+                "granitemoehybrid",
+            }
+            _model_config_type = getattr(self, "model_config_type", None) or ""
+            if _model_config_type in _SSM_HYBRID_MODEL_TYPES:
+                LOG.warning(
+                    f"context_parallel_size={self.context_parallel_size} with "
+                    f"model_type={_model_config_type}: SSM/Mamba layers use P2P "
+                    "hidden-state passing and additive output correction across "
+                    "CP ranks. Attention layers use ring attention. This is "
+                    "mathematically exact but has not been extensively validated "
+                    "end-to-end — verify loss curves match single-GPU baselines. "
+                    "Recommended: run a short training job and compare loss curves "
+                    "against a single-GPU baseline with the same data/seed."
+                )
+
+        return self
+
+    @model_validator(mode="after")
+    def validate_ring_attn_func(self):
+        if getattr(self, "context_parallel_size", 1) == 1:
+            return self
+
+        if self.ring_attn_func is not None:
+            self.ring_attn_func = RingAttnFunc(self.ring_attn_func)
+        elif getattr(self, "use_glm_dsa_kernels", False):
+            # The GLM DSA kernels own attention (including the context-parallel compressed-KV gather),
+            # so leave ring_attn_func None: the SP context manager still shards the sequence by chunking,
+            # but skips the ring_flash_attn substitution (which GLM doesn't use and isn't installed).
+            pass
+        else:
+            # Default ring attention function selection
+            sample_packing = getattr(self, "sample_packing", False)
+            self.ring_attn_func = (
+                RingAttnFunc.VARLEN_LLAMA3
+                if sample_packing
+                else RingAttnFunc.BATCH_RING
+            )
+
+        return self
+
+    def hint_gradient_checkpointing_dpo_lora_ddp(self):
+        if (
+            (self.gradient_checkpointing is True or self.gradient_checkpointing is None)
+            and self.capabilities
+            and self.capabilities.get("n_gpu", 1) > 1
+            and self.adapter in ("lora", "qlora")
+            and self.rl == RLType.DPO
+            and not self.fsdp
+            and not self.deepspeed
+        ):
+            LOG.warning(
+                "gradient_checkpointing with DPO + DDP + LoRA is not recommended."
+            )
+        return self
+
+
+class DistributedValidationMixin:
+    """validation for distributed training."""
+
+    @model_validator(mode="after")
+    def check_tensor_parallel_optimizer(self):
+        if self.tensor_parallel_size > 1:
+            if self.optimizer in ["paged_adamw_8bit", "adamw_8bit", "adamw_bnb_8bit"]:
+                raise ValueError(
+                    "tensor_parallel_size is not supported with paged_adamw_8bit, adamw_8bit, and adamw_bnb_8bit optimizers"
+                )
+
+        return self
+
+
+class EBFTValidationMixin:
+    """Validation for EBFT (Energy-Based Fine-Tuning) configuration."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_ebft_config_required(cls, data):
+        """rl: ebft requires an ebft config section."""
+        if data.get("rl") == "ebft" and not data.get("ebft"):
+            raise ValueError(
+                "`ebft` config section is required when `rl: ebft` is set. "
+                "Add an `ebft:` section with at least `mode: structured` or `mode: strided`."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_ebft_torch_compile(cls, data):
+        """torch_compile + flex_attention + gradient_checkpointing causes dynamo recompiles
+        and CheckpointErrors. The flex_attention kernel compiles itself internally —
+        whole-model torch.compile is not needed and actively harmful."""
+        if (
+            data.get("rl") == "ebft"
+            and data.get("torch_compile") is True
+            and data.get("ebft", {}).get("mode") == "strided"
+        ):
+            if data.get("gradient_checkpointing"):
+                raise ValueError(
+                    "EBFT strided mode: `torch_compile: true` with `gradient_checkpointing: true` "
+                    "causes CheckpointError (BlockMask metadata mismatch during recomputation). "
+                    "Remove `torch_compile` — the flex_attention kernel compiles itself internally."
+                )
+            LOG.warning(
+                "EBFT strided mode: `torch_compile: true` causes dynamo recompiles from "
+                "variable sequence lengths across steps. Consider removing it — "
+                "flex_attention compiles itself internally."
+            )
+        return data
+
+    @model_validator(mode="after")
+    def check_ebft_gradient_checkpointing_reentrant(self):
+        """flex_attention + non-reentrant gradient checkpointing causes CheckpointError."""
+        if (
+            self.rl == "ebft"
+            and (self.ebft or {}).get("mode") == "strided"
+            and self.attn_implementation == "flex_attention"
+            and self.gradient_checkpointing
+        ):
+            gc_kwargs = self.gradient_checkpointing_kwargs or {}
+            if not gc_kwargs.get("use_reentrant"):
+                LOG.warning(
+                    "EBFT strided mode with flex_attention: setting `use_reentrant: true` in "
+                    "gradient_checkpointing_kwargs (required for flex_attention compatibility). "
+                    "Non-reentrant checkpointing causes CheckpointError with BlockMask metadata."
+                )
+                if self.gradient_checkpointing_kwargs is None:
+                    self.gradient_checkpointing_kwargs = {}
+                self.gradient_checkpointing_kwargs["use_reentrant"] = True
+        return self
+
+    @model_validator(mode="after")
+    def check_ebft_activation_offloading(self):
+        """activation_offloading replaces gradient checkpointing with FSDP-style wrapping,
+        which conflicts with flex_attention's use_reentrant requirement."""
+        if (
+            self.rl == "ebft"
+            and (self.ebft or {}).get("mode") == "strided"
+            and self.activation_offloading is True
+            and self.attn_implementation == "flex_attention"
+        ):
+            raise ValueError(
+                "EBFT strided mode: `activation_offloading: true` is incompatible with "
+                "`attn_implementation: flex_attention`. Activation offloading replaces "
+                "gradient checkpointing with FSDP-style wrapping that conflicts with "
+                "flex_attention's reentrant checkpoint requirement. Remove "
+                "`activation_offloading` — the strided trainer uses micro-batched forward "
+                "passes for memory efficiency instead."
+            )
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_ebft_strided_sequence_len(cls, data):
+        """Warn if sequence_len is too large for single-GPU strided EBFT."""
+        if data.get("rl") != "ebft" or data.get("ebft", {}).get("mode") != "strided":
+            return data
+        ebft = data.get("ebft", {})
+        seq_len = data.get("sequence_len", 512)
+        n_samples = ebft.get("n_samples_per_prompt", 4)
+        gen_len = ebft.get("generate_max_len", 8)
+        stride = ebft.get("stride", 8)
+        ctx_len = ebft.get("context_length", 8)
+        max_blocks = (seq_len - gen_len - ctx_len) // stride + 1
+        full_seq = seq_len + max_blocks * gen_len
+        # Rough estimate: 8.7 GB per sample at S=3900 for 1B model
+        if full_seq * n_samples > 20000:
+            LOG.warning(
+                f"EBFT strided: full_seq_len={full_seq} * n_samples={n_samples} = "
+                f"{full_seq * n_samples} token-samples per step. This may require >24GB VRAM "
+                f"for a 1B+ model. Consider reducing sequence_len, n_samples_per_prompt, or stride."
+            )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_ebft_strided_dataset_split(cls, data):
+        """Warn about the common `train_on_split` mistake (silently ignored by schema)."""
+        datasets = data.get("datasets", [])
+        for ds in datasets or []:
+            if isinstance(ds, dict) and ds.get("train_on_split"):
+                LOG.warning(
+                    f"Dataset has `train_on_split: {ds['train_on_split']}` — this field "
+                    f"is not recognized and will be silently ignored. "
+                    f"Use `split: {ds['train_on_split']}` instead."
+                )
+        return data
+
+
+class GRPOVllmValidationMixin:
+    """Validation mixin for vllm when using GRPO."""
+
+    @model_validator(mode="after")
+    def check_vllm_mode_set(self):
+        if self.trl and self.trl.use_vllm and not self.trl.vllm_mode:
+            LOG.warning(
+                "vllm_mode must be set to either `server` or `colocate` when using vllm, using default value `server`"
+            )
+            self.trl.vllm_mode = "server"
+        return self
+
+
+class ValidationMixin(
+    DatasetValidationMixin,
+    AttentionValidationMixin,
+    TrainingValidationMixin,
+    LoRAValidationMixin,
+    RLValidationMixin,
+    OptimizationValidationMixin,
+    SystemValidationMixin,
+    ChatTemplateValidationMixin,
+    PretrainingValidationMixin,
+    ModelCompatibilityValidationMixin,
+    ComplexValidationMixin,
+    EBFTValidationMixin,
+    GRPOVllmValidationMixin,
+):
+    """Full validation mixin for Axolotl configuration."""
diff --git a/src/axolotl/utils/schemas/vllm.py b/src/axolotl/utils/schemas/vllm.py
new file mode 100644
index 0000000000..c1f010b02c
--- /dev/null
+++ b/src/axolotl/utils/schemas/vllm.py
@@ -0,0 +1,80 @@
+"""
+Pydantic models for VLLM configuration, used primarily for RL training with TRL + grpo
+"""
+
+from pydantic import BaseModel, Field
+
+
+class VllmConfig(BaseModel):
+    """
+    Configuration for VLLM server
+    """
+
+    device: str | None = Field(
+        default="auto",
+        json_schema_extra={"description": "Device to use for VLLM"},
+    )
+    tensor_parallel_size: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Tensor parallel size for VLLM"},
+    )
+    data_parallel_size: int | None = Field(
+        default=None,
+        json_schema_extra={"description": "Data parallel size for VLLM"},
+    )
+    gpu_memory_utilization: float | None = Field(
+        default=0.9,
+        json_schema_extra={"description": "GPU memory utilization for VLLM"},
+    )
+    dtype: str | None = Field(
+        default="auto",
+        json_schema_extra={"description": "Data type for VLLM"},
+    )
+    max_model_len: int | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Maximum length of the model context for VLLM"
+        },
+    )
+    enable_prefix_caching: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Enable prefix caching for VLLM"},
+    )
+    host: str | None = Field(
+        default="0.0.0.0",  # nosec B104
+        json_schema_extra={"description": "Host for the vLLM server to start on"},
+    )
+    port: int | None = Field(
+        default=8000,
+        json_schema_extra={"description": "Port of the vLLM server to start on"},
+    )
+
+    enable_reasoning: bool | None = Field(
+        default=None,
+        json_schema_extra={"description": "Enable reasoning for VLLM"},
+    )
+    reasoning_parser: str | None = Field(
+        default=None,
+        json_schema_extra={"description": "Reasoning parser for VLLM"},
+    )
+    enforce_eager: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Disable CUDA graph capture in vLLM. Required for models with "
+            "causal_conv1d (e.g., Qwen3.5 hybrid linear attention)."
+        },
+    )
+    serve_module: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Python module for vLLM serve script. Set to 'axolotl.scripts.vllm_serve_lora' "
+            "for native LoRA support, or leave None for default TRL serve."
+        },
+    )
+    worker_extension_cls: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "vLLM worker extension class for weight synchronization. "
+            "Defaults to 'trl.scripts.vllm_serve.WeightSyncWorkerExtension'."
+        },
+    )
diff --git a/src/axolotl/utils/tee.py b/src/axolotl/utils/tee.py
new file mode 100644
index 0000000000..7bc8efab0b
--- /dev/null
+++ b/src/axolotl/utils/tee.py
@@ -0,0 +1,166 @@
+"""
+Utilities for managing the debug log file and providing a file-only stream for logging
+handlers.
+"""
+
+from __future__ import annotations
+
+import io
+import os
+import sys
+import threading
+from pathlib import Path
+from typing import TextIO, cast
+
+_lock = threading.Lock()
+_file_handle: io.TextIOWrapper | None = None
+_log_path: str | None = None
+_tee_installed: bool = False
+_orig_stdout: TextIO | None = None
+_orig_stderr: TextIO | None = None
+
+
+class _FileOnlyWriter(io.TextIOBase):
+    """A stream-like object that writes only to the tee file.
+
+    Before the file is prepared, writes are dropped (no-op).
+    """
+
+    def write(self, s: str) -> int:  # type: ignore[override]
+        with _lock:
+            if _file_handle is not None:
+                _file_handle.write(s)
+                return len(s)
+            return len(s)
+
+    def flush(self) -> None:  # type: ignore[override]
+        with _lock:
+            if _file_handle is not None:
+                try:
+                    _file_handle.flush()
+                except Exception:
+                    pass
+
+
+file_only_stream: io.TextIOBase = _FileOnlyWriter()
+
+
+class _StreamTee(io.TextIOBase):
+    """A minimal tee that mirrors writes to the debug log file.
+
+    Installed only after the debug log is prepared; no buffering.
+    """
+
+    def __init__(self, stream: io.TextIOBase):
+        self._stream = stream
+
+    def write(self, s: str) -> int:  # type: ignore[override]
+        with _lock:
+            n = self._stream.write(s)
+            if _file_handle is not None:
+                _file_handle.write(s)
+            return n
+
+    def flush(self) -> None:  # type: ignore[override]
+        with _lock:
+            self._stream.flush()
+            if _file_handle is not None:
+                try:
+                    _file_handle.flush()
+                except Exception:
+                    pass
+
+    @property
+    def encoding(self):  # type: ignore[override]
+        return getattr(self._stream, "encoding", None)
+
+    @property
+    def errors(self):  # type: ignore[override]
+        return getattr(self._stream, "errors", None)
+
+    def isatty(self):  # type: ignore[override]
+        return getattr(self._stream, "isatty", lambda: False)()
+
+    def fileno(self):  # type: ignore[override]
+        if hasattr(self._stream, "fileno"):
+            return self._stream.fileno()
+        raise OSError("Underlying stream has no fileno")
+
+
+def prepare_debug_log(cfg, filename: str = "debug.log") -> str:
+    """
+    Prepare the debug log.
+
+    Creates the output directory, handles append/truncate logic based on cfg, and opens
+    the debug log file for subsequent writes via file-only handlers.
+    """
+    global _file_handle, _log_path, _tee_installed
+
+    with _lock:
+        # If already initialized, reuse existing path
+        if _log_path is not None:
+            return _log_path
+
+        output_dir = cfg.output_dir
+        os.makedirs(output_dir, exist_ok=True)
+
+        log_path = Path(output_dir) / filename
+        append = bool(
+            cfg.get("resume_from_checkpoint") or cfg.get("auto_resume_from_checkpoints")
+        )
+
+        if not append:
+            log_path.unlink(missing_ok=True)
+
+        fh = open(log_path, "a", encoding="utf-8")
+        fh.flush()
+
+        _file_handle = fh
+        _log_path = str(log_path)
+
+        # Install a tee so stdout/stderr are mirrored to the debug file
+        # Allow disabling via env for testing or advanced usage.
+        tee_enabled = os.getenv("AXOLOTL_TEE_STDOUT", "1").lower() not in {
+            "0",
+            "false",
+            "no",
+        }
+        if tee_enabled and not _tee_installed:
+            # Save originals so we can restore later (e.g., tests)
+            global _orig_stdout, _orig_stderr
+            _orig_stdout = sys.stdout
+            _orig_stderr = sys.stderr
+            sys.stdout = _StreamTee(cast(io.TextIOBase, sys.stdout))
+            sys.stderr = _StreamTee(cast(io.TextIOBase, sys.stderr))
+            _tee_installed = True
+
+        return _log_path
+
+
+def close_debug_log() -> None:
+    """Flush and close the debug log and uninstall the stdout/stderr tee.
+
+    Safe to call even if not initialized.
+    """
+    global _file_handle, _log_path, _tee_installed, _orig_stdout, _orig_stderr
+    with _lock:
+        # Restore original stdout/stderr if we installed a tee
+        if _tee_installed:
+            if _orig_stdout is not None:
+                sys.stdout = _orig_stdout
+            if _orig_stderr is not None:
+                sys.stderr = _orig_stderr
+            _tee_installed = False
+            _orig_stdout = None
+            _orig_stderr = None
+
+        # Close the file handle if open
+        if _file_handle is not None:
+            try:
+                _file_handle.flush()
+                _file_handle.close()
+            except Exception:
+                pass
+            finally:
+                _file_handle = None
+        _log_path = None
diff --git a/src/axolotl/utils/tokenization.py b/src/axolotl/utils/tokenization.py
index 845296b7a6..3f44a34292 100644
--- a/src/axolotl/utils/tokenization.py
+++ b/src/axolotl/utils/tokenization.py
@@ -1,12 +1,10 @@
 """Module for tokenization utilities"""
 
-import logging
-import re
-from typing import Dict, List
-
 from termcolor import colored
 
-LOG = logging.getLogger("axolotl")
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
 
 
 def check_dataset_labels(
@@ -28,11 +26,12 @@ def check_example_labels(example, tokenizer, text_only=False):
     # Get the input_ids, labels, and attention_mask from the dataset
     input_ids = example["input_ids"]
     labels = example["labels"]
+    target_mask = example.pop("target_mask", None)
 
     # You can compare the input_ids and labels element-wise
     # Remember to ignore positions with IGNORE_TOKEN_ID (if you use it) or attention_mask equal to 0
     colored_tokens = []
-    for _, (input_id, label_id) in enumerate(zip(input_ids, labels)):
+    for _, (input_id, label_id) in enumerate(zip(input_ids, labels, strict=False)):
         decoded_input_token = tokenizer.decode(input_id)
         # Choose the color based on whether the label has the ignore value or not
         color = "red" if label_id == -100 else ("yellow" if label_id == 0 else "green")
@@ -44,6 +43,13 @@ def check_example_labels(example, tokenizer, text_only=False):
     delimiter = "" if text_only else " "
     LOG.info(delimiter.join(colored_tokens))
     LOG.info("\n\n\n")
+    target_labels_count = sum(label_id != -100 for label_id in labels)
+    total_len = len(input_ids)
+    LOG.info(f"Total input len: {total_len}")
+    LOG.info(f"Count of labels: {target_labels_count}")
+    if target_mask:
+        target_mask_positions = sum(m[0] for m in target_mask)
+        LOG.info(f"Number of positions in target_mask: {target_mask_positions}")
 
     return " ".join(colored_tokens)
 
@@ -62,96 +68,53 @@ def process_tokens_for_rl_debug(tokens, color, tokenizer, text_only):
     """Helper function to process and color tokens."""
     colored_tokens = [
         color_token_for_rl_debug(tokenizer.decode(token), token, color, text_only)
-        for token in tokenizer.encode(tokens)
+        for token in tokenizer.encode(tokens, add_special_tokens=False)
     ]
     return colored_tokens
 
 
 def check_rl_example_labels(example, tokenizer, text_only=False):
-    field_prompt, field_chosen, field_rejected = "prompt", "chosen", "rejected"
+    field_prompt, field_chosen, field_rejected, field_completion = (
+        "prompt",
+        "chosen",
+        "rejected",
+        "completion",
+    )
 
     input_tokens = example[field_prompt]
-    labels_chosen, labels_rejected = example[field_chosen], example[field_rejected]
+
+    labels_chosen = example.get(field_chosen)
+    labels_rejected = example.get(field_rejected)
+    labels_completion = example.get(field_completion)
+
+    # Create a delimiter based on text_only flag
+    delimiter = "" if text_only else " "
 
     # Process and color each type of token
     colored_tokens = process_tokens_for_rl_debug(
         input_tokens, "yellow", tokenizer, text_only
     )
-    colored_chosens = process_tokens_for_rl_debug(
-        labels_chosen, "green", tokenizer, text_only
-    )
-    colored_rejecteds = process_tokens_for_rl_debug(
-        labels_rejected, "red", tokenizer, text_only
-    )
 
-    # Create a delimiter based on text_only flag
-    delimiter = "" if text_only else " "
+    # Process tokens
+    if labels_completion is None:
+        colored_chosens = process_tokens_for_rl_debug(
+            labels_chosen, "green", tokenizer, text_only
+        )
+        colored_rejecteds = process_tokens_for_rl_debug(
+            labels_rejected, "red", tokenizer, text_only
+        )
+    else:
+        colored_completion = process_tokens_for_rl_debug(
+            labels_completion, "green", tokenizer, text_only
+        )
 
     # Logging information
     LOG.info(f"INPUT PROMPT: {delimiter.join(colored_tokens)}\n\n")
-    LOG.info(f"CHOSEN RESPONSE: {delimiter.join(colored_chosens)}\n\n")
-    LOG.info(f"REJECTED RESPONSE: {delimiter.join(colored_rejecteds)}\n\n\n")
-
-    return delimiter.join(colored_tokens)
-
 
-GLAIVE_ROLES = ["USER", "ASSISTANT", "FUNCTION RESPONSE"]
-GLAIVE_TO_SHAREGPT_ROLE = {
-    "SYSTEM": "system",
-    "USER": "human",
-    "ASSISTANT": "gpt",
-    "FUNCTION RESPONSE": "tool",
-}
+    if labels_completion is None:
+        LOG.info(f"CHOSEN RESPONSE: {delimiter.join(colored_chosens)}\n\n")
+        LOG.info(f"REJECTED RESPONSE: {delimiter.join(colored_rejecteds)}\n\n\n")
+    else:
+        LOG.info(f"COMPLETION RESPONSE: {delimiter.join(colored_completion)}\n\n\n")
 
-GLAIVE_MSG_REGEX = re.compile(rf"({'|'.join(GLAIVE_ROLES)}): ")
-
-
-def chatml_to_conversation(row: Dict[str, str]) -> List[Dict[str, str]]:
-    """
-    Converts a ChatML formatted row to a list of messages in ShareGPT format.
-    Initially based off https://github.com/lilacai/lilac/blob/main/notebooks/GlaiveToShareGPT.ipynb.
-    """
-
-    system_prompt = row.get("system")
-    if system_prompt:
-        system_prompt = system_prompt.removeprefix("SYSTEM: ")
-
-    chat_str = row["chat"]
-    chat_msgs = [s.strip() for s in GLAIVE_MSG_REGEX.split(chat_str) if s]
-
-    chat_msg_dicts = [
-        {"from": GLAIVE_TO_SHAREGPT_ROLE[role], "value": value}
-        for role, value in zip(chat_msgs[::2], chat_msgs[1::2])
-    ]
-
-    if system_prompt:
-        chat_msg_dicts = [
-            {"from": GLAIVE_TO_SHAREGPT_ROLE["SYSTEM"], "value": system_prompt}
-        ] + chat_msg_dicts
-
-    return chat_msg_dicts
-
-
-def merge_consecutive_messages(messages):
-    """
-    Merge consecutive messages from the same sender into a single message.
-    This can be useful with datasets that contain multiple consecutive tool calls.
-    """
-
-    merged_messages = []
-    current_from = None
-    current_message = ""
-
-    for msg in messages:
-        if current_from == msg["from"]:
-            current_message += msg["value"]
-        else:
-            if current_from is not None:
-                merged_messages.append({"from": current_from, "value": current_message})
-            current_from = msg["from"]
-            current_message = msg["value"]
-
-    if current_from is not None:
-        merged_messages.append({"from": current_from, "value": current_message})
-
-    return merged_messages
+    return delimiter.join(colored_tokens)
diff --git a/src/axolotl/utils/trackio_.py b/src/axolotl/utils/trackio_.py
new file mode 100644
index 0000000000..2bddfb972e
--- /dev/null
+++ b/src/axolotl/utils/trackio_.py
@@ -0,0 +1,17 @@
+"""Module for trackio utilities"""
+
+import os
+
+from axolotl.utils.dict import DictDefault
+
+
+def setup_trackio_env_vars(cfg: DictDefault):
+    for key in cfg.keys():
+        if key.startswith("trackio_"):
+            value = cfg.get(key, "")
+
+            if value and isinstance(value, str) and len(value) > 0:
+                os.environ[key.upper()] = value
+
+    if cfg.trackio_project_name and len(cfg.trackio_project_name) > 0:
+        cfg.use_trackio = True
diff --git a/src/axolotl/utils/train.py b/src/axolotl/utils/train.py
new file mode 100644
index 0000000000..ad3f72be45
--- /dev/null
+++ b/src/axolotl/utils/train.py
@@ -0,0 +1,47 @@
+"""Training utils for checkpoints"""
+
+from pathlib import Path
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def determine_last_checkpoint(cfg: DictDefault, update: bool = True) -> str | None:
+    """
+    Determine the checkpoint to resume from based on configuration.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        update: Whether to update the config with the determined checkpoint
+
+    Returns:
+        Path to the checkpoint to resume from, or `None` if not resuming.
+    """
+    last_checkpoint = None
+    checkpoints = sorted(
+        (
+            p
+            for p in Path(cfg.output_dir).glob("checkpoint-*")
+            if p.name.split("-")[-1].isdigit()
+        ),
+        key=lambda p: int(p.name.split("-")[-1]),
+    )
+    if checkpoints:
+        last_checkpoint = str(checkpoints[-1])
+        if not update:
+            LOG.info(f"Resuming from last checkpoint at {last_checkpoint}")
+            return last_checkpoint
+
+    if (
+        cfg.resume_from_checkpoint is None
+        and cfg.auto_resume_from_checkpoints
+        and last_checkpoint is not None
+    ):
+        cfg.resume_from_checkpoint = last_checkpoint
+        LOG.info(
+            "Using auto-resume functionality to resume from checkpoint at "
+            f"{cfg.resume_from_checkpoint}"
+        )
+    return cfg.resume_from_checkpoint
diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
index 2e3728cc8a..d6f1db42f2 100644
--- a/src/axolotl/utils/trainer.py
+++ b/src/axolotl/utils/trainer.py
@@ -1,24 +1,30 @@
 """Module containing the Trainer class and related functions"""
+
+import json
 import math
 import os
 import random
 from contextlib import contextmanager
 from functools import partial
+from tempfile import NamedTemporaryFile
 from typing import List, Optional
 
 import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
 import torch
 import torch.cuda
-from accelerate.logging import get_logger
-from datasets import set_caching_enabled
-from torch.utils.data import DataLoader, RandomSampler
+from datasets import IterableDataset, disable_caching, enable_caching
+from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.core.trainer_builder import HFCausalTrainerBuilder, HFRLTrainerBuilder
-from axolotl.utils.distributed import is_main_process, reduce_and_broadcast, zero_first
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.distributed import init_distributed_state, reduce_and_broadcast
+from axolotl.utils.environment import check_cuda_p2p_ib_support
+from axolotl.utils.logging import get_logger
 from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
 
-LOG = get_logger("axolotl")
+LOG = get_logger(__name__)
 
 
 @torch.jit.script
@@ -86,16 +92,48 @@ def trainer_weighted_loss(model_output, labels, shift_labels=True):
 @contextmanager
 def disable_datasets_caching():
     try:
-        set_caching_enabled(False)
+        disable_caching()
         yield
     finally:
-        set_caching_enabled(True)
+        enable_caching()
 
 
 def add_position_ids(sample):
-    sample_len = len(sample["input_ids"])
-    sample["position_ids"] = torch.arange(len(sample["input_ids"]))
-    sample["length"] = sample_len
+    """
+    Handle both single-example and batched data.
+    - single example: sample['input_ids'] is a list[int]
+    - batched data: sample['input_ids'] is a list[list[int]]
+    """
+    # Return sample unchanged if "input_ids" is not present, or is empty
+    if "input_ids" not in sample or not sample["input_ids"]:
+        return sample
+
+    input_ids = sample["input_ids"]
+
+    # If first element is an int, it’s a single example
+    # If first element is a list, it’s a batch
+    if isinstance(input_ids[0], int):
+        # ---- SINGLE EXAMPLE ----
+        seq_len = len(input_ids)
+        # Position IDs for a single example
+        # As a list
+        sample["position_ids"] = list(range(seq_len))
+        sample["length"] = seq_len
+
+    else:
+        # ---- BATCHED EXAMPLES ----
+        # input_ids is a list of lists
+        position_ids_batch = []
+        lengths_batch = []
+        for seq in input_ids:
+            seq_len = len(seq)
+            position_ids_batch.append(list(range(seq_len)))
+            lengths_batch.append(seq_len)
+
+        # Now store them back
+        sample["position_ids"] = position_ids_batch
+        sample["length"] = lengths_batch
+
     return sample
 
 
@@ -169,150 +207,246 @@ def add_length(sample):
     return sample
 
 
-def drop_long_seq(sample, sequence_len=2048, min_sequence_len=2):
-    return (
-        len(sample["input_ids"]) <= sequence_len
-        and len(sample["input_ids"]) >= min_sequence_len
-    )
+def filter_sequences_by_length(
+    sample, sequence_len=2048, min_sequence_len=2, raise_on_drop=False
+):
+    """
+    Filter sequences outside valid length range [min_sequence_len, sequence_len].
 
+    Drops samples that are either too short (< min_sequence_len) or too long (> sequence_len).
 
-def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
-    drop_long = partial(
-        drop_long_seq,
-        sequence_len=cfg.sequence_len,
-        min_sequence_len=cfg.min_sample_len or 2,
-    )
-    with zero_first(is_main_process()):
-        if cfg.is_preprocess:
-            min_input_len = np.min(get_dataset_lengths(train_dataset))
-            LOG.debug(f"min_input_len: {min_input_len}", main_process_only=True)
-            max_input_len = np.max(get_dataset_lengths(train_dataset))
-            LOG.debug(f"max_input_len: {max_input_len}", main_process_only=True)
-
-        if (
-            cfg.is_mistral_derived_model and cfg.flash_attention
-        ) or cfg.model_config_type == "mamba":
-            LOG.info("dropping attention_mask column")
-            train_dataset = train_dataset.remove_columns("attention_mask")
-            if eval_dataset:
-                eval_dataset = eval_dataset.remove_columns("attention_mask")
-
-        if cfg.model_config_type == "falcon":
-            LOG.info("dropping token_type_ids column if it exists")
-            if "token_type_ids" in train_dataset.column_names:
-                train_dataset = train_dataset.remove_columns("token_type_ids")
-            if eval_dataset and "token_type_ids" in eval_dataset.column_names:
-                eval_dataset = eval_dataset.remove_columns("token_type_ids")
-
-        train_dataset = train_dataset.filter(
-            drop_long,
-            num_proc=cfg.dataset_processes,
-            load_from_cache_file=not cfg.is_preprocess,
-            desc="Dropping Long Sequences",
-        )
-        if eval_dataset:
-            eval_dataset = eval_dataset.filter(
-                drop_long,
-                num_proc=cfg.dataset_processes,
-                load_from_cache_file=not cfg.is_preprocess,
-                desc="Dropping Long Sequences",
-            )
+    Works for both single-example (list[int]) or batched (list[list[int]]).
 
-        if cfg.group_by_length:
-            train_dataset = train_dataset.map(
-                add_length,
-                num_proc=cfg.dataset_processes,
-                load_from_cache_file=not cfg.is_preprocess,
-                desc="Group By Length",
-            )
+    If raise_on_drop is set, the code raises a ValueError if a sample is
+    encountered that is too long and would have been dropped.
+    """
+    min_sequence_len = min_sequence_len or 2
 
-        if cfg.use_pose:
-            pose_kwargs = {}
-            if cfg.pose_num_chunks is not None:
-                pose_kwargs["chunks"] = cfg.pose_num_chunks
-            pose_fn = partial(
-                add_pose_position_ids,
-                max_context_len=cfg.pose_max_context_len,
-                split_on_token_ids=cfg.pose_split_on_token_ids,
-                **pose_kwargs,
+    input_ids = sample["input_ids"]
+
+    # Edge case: if input_ids is empty
+    if not input_ids:
+        # Decide if you want to drop or keep empty. Let's drop.
+        return False
+
+    # Check if single example or batched by looking at the first element
+    if isinstance(input_ids[0], int):
+        # Single example (input_ids is a list of int)
+        length = len(input_ids)
+        if raise_on_drop and length > sequence_len:
+            raise ValueError(
+                f"Sequence encountered with {length} tokens, which exceeds the maximum {sequence_len}."
             )
-            train_dataset = train_dataset.map(
-                pose_fn,
-                num_proc=cfg.dataset_processes,
-                load_from_cache_file=not cfg.is_preprocess,
-                desc="Add position_id column (PoSE)",
+        return min_sequence_len <= length <= sequence_len
+
+    # Batched (input_ids is a list of lists)
+    results = []
+    for seq in input_ids:
+        length = len(seq)
+        if raise_on_drop and length > sequence_len:
+            raise ValueError(
+                f"Sequence encountered with {length} tokens, which exceeds the maximum {sequence_len}."
             )
-            train_dataset = train_dataset.sort("sequence_len")
-            if cfg.eval_sample_packing is not False:
-                if eval_dataset:
-                    eval_dataset = eval_dataset.map(
-                        pose_fn,
-                        num_proc=cfg.dataset_processes,
-                        load_from_cache_file=not cfg.is_preprocess,
-                        desc="Add position_id column (PoSE)",
-                    )
-        elif cfg.sample_packing:
-            train_dataset = train_dataset.map(
-                add_position_ids,
-                num_proc=cfg.dataset_processes,
-                load_from_cache_file=not cfg.is_preprocess,
-                desc="Add position_id column (Sample Packing)",
+        results.append(min_sequence_len <= length <= sequence_len)
+    return results
+
+
+def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
+    drop_attn_mask = cfg.model_config_type in ["mamba", "gemma3"]
+    if drop_attn_mask:
+        LOG.info("dropping attention_mask column")
+        train_dataset = train_dataset.remove_columns("attention_mask")
+        if eval_dataset:
+            eval_dataset = eval_dataset.remove_columns("attention_mask")
+
+    if cfg.model_config_type in ["falcon", "mistral"]:
+        LOG.info("dropping token_type_ids column if it exists")
+        if "token_type_ids" in train_dataset.column_names:
+            train_dataset = train_dataset.remove_columns("token_type_ids")
+        if eval_dataset and "token_type_ids" in eval_dataset.column_names:
+            eval_dataset = eval_dataset.remove_columns("token_type_ids")
+
+    def drop_no_trainable_tokens(sample):
+        """
+        Drop samples if all labels are -100 (i.e., zero trainable tokens).
+        Works for both single-example or batched input.
+        """
+        labels = sample["labels"]
+        if not labels:
+            return True
+
+        # Check if single example or batch
+        # If first element is an int, we assume a single example
+        # If it's a list, we assume we're dealing with a batch
+        if isinstance(labels[0], int):
+            # Single example: return a single bool
+            return np.any(labels != -100)
+
+        # Batched: 'labels' is a list of lists
+        # Return a list of booleans, one per sub-list
+        results = [np.any(row_labels != -100) for row_labels in labels]
+        return results
+
+    try:
+        prior_len = len(train_dataset)
+    except TypeError:
+        # handle iterable datasets case
+        prior_len = None
+    filter_map_kwargs = {}
+    if not isinstance(train_dataset, IterableDataset):
+        filter_map_kwargs["num_proc"] = cfg.dataset_num_proc
+        filter_map_kwargs["load_from_cache_file"] = not cfg.is_preprocess
+
+    drop_long_kwargs = {}
+    if filter_map_kwargs:
+        drop_long_kwargs["desc"] = "Drop Samples with Zero Trainable Tokens"
+    train_dataset = train_dataset.filter(
+        drop_no_trainable_tokens,
+        batched=True,
+        **filter_map_kwargs,
+        **drop_long_kwargs,
+    )
+    if prior_len:
+        dropped = prior_len - len(train_dataset)
+        if dropped:
+            LOG.warning(
+                f"Dropped {dropped} samples with no trainable tokens from train dataset"
             )
-            if cfg.eval_sample_packing is not False:
-                if eval_dataset:
-                    eval_dataset = eval_dataset.map(
-                        add_position_ids,
-                        num_proc=cfg.dataset_processes,
-                        load_from_cache_file=not cfg.is_preprocess,
-                        desc="Add position_id column (Sample Packing)",
-                    )
+
+    if eval_dataset:
+        try:
+            prior_len = len(eval_dataset)
+        except TypeError:
+            # handle iterable datasets case
+            prior_len = None
+        eval_dataset = eval_dataset.filter(
+            drop_no_trainable_tokens,
+            **filter_map_kwargs,
+            **drop_long_kwargs,
+        )
+        if prior_len:
+            dropped = prior_len - len(eval_dataset)
+            if dropped:
+                LOG.warning(
+                    f"Dropped {dropped} samples with no trainable tokens from eval dataset"
+                )
+
+    if cfg.group_by_length:
+        train_dataset = train_dataset.map(
+            add_length,
+            num_proc=cfg.dataset_num_proc,
+            load_from_cache_file=not cfg.is_preprocess,
+            desc="Group By Length",
+        )
+
+    if cfg.use_pose:
+        pose_kwargs = {}
+        if cfg.pose_num_chunks is not None:
+            pose_kwargs["chunks"] = cfg.pose_num_chunks
+        pose_fn = partial(
+            add_pose_position_ids,
+            max_context_len=cfg.pose_max_context_len,
+            split_on_token_ids=cfg.pose_split_on_token_ids,
+            **pose_kwargs,
+        )
+        train_dataset = train_dataset.map(
+            pose_fn,
+            num_proc=cfg.dataset_num_proc,
+            load_from_cache_file=not cfg.is_preprocess,
+            desc="Add position_id column (PoSE)",
+        )
+        train_dataset = train_dataset.sort("sequence_len")
+        if cfg.eval_sample_packing is not False:
+            if eval_dataset:
+                eval_dataset = eval_dataset.map(
+                    pose_fn,
+                    num_proc=cfg.dataset_num_proc,
+                    load_from_cache_file=not cfg.is_preprocess,
+                    desc="Add position_id column (PoSE)",
+                )
+    elif cfg.sample_packing:
+        drop_long_kwargs = {}
+        if filter_map_kwargs:
+            drop_long_kwargs["desc"] = "Add position_id column (Sample Packing)"
+        train_dataset = train_dataset.map(
+            add_position_ids,
+            batched=True,
+            **filter_map_kwargs,
+            **drop_long_kwargs,
+        )
+        if cfg.eval_sample_packing:
+            if eval_dataset:
+                eval_dataset = eval_dataset.map(
+                    add_position_ids,
+                    **filter_map_kwargs,
+                    **drop_long_kwargs,
+                )
 
     return train_dataset, eval_dataset
 
 
 def process_pretraining_datasets_for_packing(
-    train_dataset, sequence_len, skip_position_ids=True
+    train_dataset, sequence_len, skip_position_ids=True, drop_attention_mask=False
 ):
-    drop_long = partial(drop_long_seq, sequence_len=sequence_len)
+    drop_outside_range = partial(filter_sequences_by_length, sequence_len=sequence_len)
 
     train_dataset = train_dataset.filter(
-        drop_long,
+        drop_outside_range,
         desc="Dropping Long Sequences",
+        load_from_cache_file=False,
     )
-    if skip_position_ids:
+    if not skip_position_ids:
         train_dataset = train_dataset.map(
             add_position_ids,
+            batched=True,
             desc="Add position_id column (Pretraining Sample Packing)",
         )
+    if drop_attention_mask:
+        train_dataset = train_dataset.remove_columns("attention_mask")
 
     return train_dataset
 
 
 def calculate_total_num_steps(cfg, train_dataset, update=True):
-    if not cfg.total_num_tokens:
-        total_num_tokens = np.sum(
-            train_dataset.data.column("input_ids")
-            .to_pandas()
-            .apply(lambda x: len(x))  # pylint: disable=unnecessary-lambda
-            .values
-        )
-        LOG.debug(f"total_num_tokens: {total_num_tokens:_}", main_process_only=True)
+    if (
+        not cfg.total_num_tokens
+        and not cfg.skip_prepare_dataset
+        and not cfg.reward_model
+    ):
+        if "length" in train_dataset.column_names:
+            total_num_tokens = int(
+                pc.sum(train_dataset.data.column("length"), min_count=0).as_py()
+            )
+        else:
+            total_num_tokens = int(
+                pc.sum(
+                    pc.list_value_length(train_dataset.data.column("input_ids")),
+                    min_count=0,
+                ).as_py()
+            )
+        LOG.debug(f"total_num_tokens: {total_num_tokens:_}")
         if update:
             cfg.total_num_tokens = total_num_tokens
 
     skip_estimates = cfg.model_config_type == "mamba"
 
-    if not skip_estimates and not cfg.total_supervised_tokens:
-        total_supervised_tokens = (
-            train_dataset.data.column("labels")
-            .to_pandas()
-            .apply(lambda x: np.sum(np.array(x) != -100))
-            .sum()
-        )
-        LOG.debug(
-            f"`total_supervised_tokens: {total_supervised_tokens:_}`",
-            main_process_only=True,
-        )
+    if (
+        not skip_estimates
+        and not cfg.total_supervised_tokens
+        and not cfg.skip_prepare_dataset
+        and not cfg.reward_model
+    ):
+        # Stream the labels column in record-batch chunks instead of
+        # .to_pandas(), which materializes one Python list per row.
+        total_supervised_tokens = 0
+        for batch in train_dataset.data.to_batches(max_chunksize=1024):
+            labels = batch.column("labels")
+            if pa.types.is_list(labels.type) or pa.types.is_large_list(labels.type):
+                flat = labels.flatten().to_numpy(zero_copy_only=False)
+            else:
+                flat = labels.to_numpy(zero_copy_only=False)
+            total_supervised_tokens += int((flat != -100).sum())
+        LOG.debug(f"total_supervised_tokens: {total_supervised_tokens:_}")
         if update:
             cfg.total_supervised_tokens = total_supervised_tokens
 
@@ -323,61 +457,66 @@ def calculate_total_num_steps(cfg, train_dataset, update=True):
         if cfg.sample_packing_eff_est:
             total_num_steps = (
                 # match count to len est in dataloader
-                (
+                int(
                     math.floor(
                         0.99
                         * cfg.total_num_tokens
                         / cfg.sample_packing_eff_est
                         / cfg.sequence_len
                         // cfg.batch_size
-                        // int(os.environ.get("WORLD_SIZE", 1))
                     )
                     - 1
                 )
                 * cfg.num_epochs
             )
             LOG.debug(
-                f"total_num_tokens: {cfg.total_num_tokens:_}, total_num_steps: {total_num_steps:_}",
-                main_process_only=True,
+                f"total_num_tokens: {cfg.total_num_tokens:_}, total_num_steps: {total_num_steps:_}"
             )
         else:
-            if cfg.flash_attention:
-                batch_size = 1
+            if cfg.attn_supports_packing and not cfg.multipack_real_batches:
+                sampler_batch_size = 1
                 batch_max_len = cfg.micro_batch_size * cfg.sequence_len
             else:
-                batch_size = cfg.micro_batch_size
+                sampler_batch_size = cfg.micro_batch_size
                 batch_max_len = cfg.sequence_len
+            if cfg.curriculum_sampling:
+                sampler = SequentialSampler(train_dataset)
+            else:
+                sampler = RandomSampler(train_dataset)
             sampler = MultipackBatchSampler(
-                sampler=RandomSampler(train_dataset),
-                batch_size=batch_size,
-                drop_last=True,
-                batch_max_len=batch_max_len,
+                sampler=sampler,
                 lengths=get_dataset_lengths(train_dataset),
+                batch_size=sampler_batch_size,
+                batch_max_len=batch_max_len,
+                group_size=cfg.sample_packing_group_size,
+                bin_size=cfg.sample_packing_bin_size,
+                sequential=cfg.sample_packing_sequentially,
+                drop_last=True,
+                num_processes=cfg.dataset_num_proc,
+                mp_start_method=cfg.sample_packing_mp_start_method or "fork",
             )
 
             data_loader = DataLoader(
                 train_dataset.remove_columns(["length"]),
                 batch_sampler=sampler,
             )
-            data_loader_len = len(data_loader) // cfg.batch_size
-            actual_eff = sampler.efficiency()
-            LOG.debug(f"data_loader_len: {data_loader_len}", main_process_only=True)
+            data_loader_len = max(
+                1, len(data_loader) * cfg.micro_batch_size // cfg.batch_size
+            )
+            LOG.debug(f"data_loader_len: {data_loader_len}")
             # FIXME: is there a bug here somewhere? the total num steps depends
             # on the agreed on value for sample_packing_eff_est
-            total_num_steps = int(
-                math.floor(
-                    data_loader_len
-                    * cfg.num_epochs
-                    / int(os.environ.get("WORLD_SIZE", 1))
-                )
-            )
+            total_num_steps = int(math.floor(data_loader_len * cfg.num_epochs))
+            if cfg.dataloader_drop_last:
+                # drop the last batch for each epoch
+                total_num_steps -= int(math.ceil(cfg.num_epochs))
 
             def calc_sample_packing_eff_est(estimates: List[float]):
                 LOG.info(f"sample_packing_eff_est across ranks: {repr(estimates)}")
                 return max(estimates)
 
             sample_packing_actual_eff_all = reduce_and_broadcast(
-                lambda: actual_eff,
+                lambda: sampler.efficiency(),
                 calc_sample_packing_eff_est,
             )
             sample_packing_eff_est = (
@@ -385,65 +524,219 @@ def calc_sample_packing_eff_est(estimates: List[float]):
             )
             if update:
                 cfg.sample_packing_eff_est = sample_packing_eff_est
-            LOG.debug(
-                f"sample_packing_eff_est: {cfg.sample_packing_eff_est}",
-                main_process_only=True,
-            )
+            LOG.debug(f"sample_packing_eff_est: {cfg.sample_packing_eff_est}")
     else:
         total_num_steps = int(
-            math.ceil(
-                len(train_dataset)
-                * cfg.num_epochs
-                / int(os.environ.get("WORLD_SIZE", 1))
-                / cfg.batch_size
-            )
+            math.ceil(len(train_dataset) * cfg.num_epochs / cfg.batch_size)
         )
-    LOG.debug(f"total_num_steps: {total_num_steps}", main_process_only=True)
+    LOG.debug(f"total_num_steps: {total_num_steps}")
     return total_num_steps
 
 
+def setup_torch_compile_env(cfg):
+    if cfg.torch_compile:
+        if not cfg.torch_compile_backend:
+            os.environ["ACCELERATE_DYNAMO_BACKEND"] = "INDUCTOR"
+        else:
+            os.environ["ACCELERATE_DYNAMO_BACKEND"] = cfg.torch_compile_backend.upper()
+
+
+def setup_deepspeed_env(cfg, stage=None):
+    from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig
+
+    from axolotl.utils.distributed import distributed_state
+
+    if distributed_state and distributed_state.initialized:
+        raise RuntimeError(
+            "Distributed State already initialized before Deepspeed setup"
+        )
+
+    os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
+    if isinstance(cfg.deepspeed, DictDefault):
+        with NamedTemporaryFile(
+            mode="w", delete=False, suffix=".json", prefix="deepspeed_config_"
+        ) as temp_file:
+            temp_file.write(json.dumps(cfg.deepspeed.to_dict(), indent=4))
+            temp_file.close()
+            cfg.deepspeed = str(temp_file.name)
+    os.environ["ACCELERATE_DEEPSPEED_CONFIG_FILE"] = cfg.deepspeed
+    os.environ["ACCELERATE_GRADIENT_ACCUMULATION_STEPS"] = str(
+        cfg.gradient_accumulation_steps
+    )
+    if stage:
+        os.environ["ACCELERATE_DEEPSPEED_ZERO_STAGE"] = str(stage)
+        if stage == 3:
+            os.environ["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = "true"
+
+    device_count = torch.cuda.device_count()
+    if device_count == 1:
+        os.environ.setdefault("WORLD_SIZE", "1")
+        os.environ.setdefault("LOCAL_RANK", "0")
+        os.environ.setdefault("MASTER_ADDR", "0.0.0.0")  # nosec B104
+        os.environ.setdefault("MASTER_PORT", "29500")
+
+    # NOTE(djsaunde): The distribued state cannot be initialized prior to the
+    # ACCELERATE_USE_DEEPSPEED assignment, but it must be initialized some time prior
+    # to model load.
+    if (
+        int(os.environ.get("WORLD_SIZE", "1")) == 1
+        and os.environ.get("AXOLOTL_IS_PREPROCESS", "0") != "1"
+        and cfg.use_ray is not True
+    ):
+        os.environ["WORLD_SIZE"] = "1"  # force it in case not set
+        os.environ["LOCAL_RANK"] = "0"  # force it in case not set
+        os.environ["RANK"] = os.environ.get("LOCAL_RANK", "0")
+        import deepspeed.comm as dist
+
+        dist.init_distributed(
+            dist_backend="nccl", auto_mpi_discovery=False, dist_init_required=True
+        )
+    init_distributed_state()
+
+    # If we don't assign this, it doesn't actually get set in the accelerate weakref
+    _ = HfTrainerDeepSpeedConfig(cfg.deepspeed)
+
+
 def setup_fsdp_envs(cfg):
     os.environ["ACCELERATE_USE_FSDP"] = "true"
-    if cfg.fsdp_config.fsdp_activation_checkpointing:
+
+    # TODO @SalmanMohammadi remove FSDP1 args in 0.12
+    if str(cfg.fsdp_version) == "2":
+        os.environ["FSDP_VERSION"] = "2"
+    if cfg.fsdp_config.activation_checkpointing:
         os.environ["FSDP_ACTIVATION_CHECKPOINTING"] = "true"
-    if cfg.fsdp_config.fsdp_offload_params:
+    if cfg.fsdp_config.offload_params:
         os.environ["FSDP_OFFLOAD_PARAMS"] = "true"
-    if cfg.fsdp_config.fsdp_sync_module_states:
+    if cfg.fsdp_config.sync_module_states:
         os.environ["FSDP_SYNC_MODULE_STATES"] = "true"
-    if cfg.fsdp_config.fsdp_cpu_ram_efficient_loading:
+    if cfg.fsdp_config.cpu_ram_efficient_loading:
         os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = "true"
-    if cfg.fsdp_config.fsdp_use_orig_params:
+    if cfg.fsdp_config.use_orig_params:
         os.environ["FSDP_USE_ORIG_PARAMS"] = "true"
-    if cfg.fsdp_config.fsdp_state_dict_type:
-        os.environ["FSDP_STATE_DICT_TYPE"] = cfg.fsdp_config.fsdp_state_dict_type
-    if cfg.fsdp_config.fsdp_auto_wrap_policy:
-        os.environ["FSDP_AUTO_WRAP_POLICY"] = cfg.fsdp_config.fsdp_auto_wrap_policy
-    if cfg.fsdp_config.fsdp_transformer_layer_cls_to_wrap:
-        os.environ[
-            "FSDP_TRANSFORMER_CLS_TO_WRAP"
-        ] = cfg.fsdp_config.fsdp_transformer_layer_cls_to_wrap
+    if cfg.fsdp_config.state_dict_type:
+        os.environ["FSDP_STATE_DICT_TYPE"] = cfg.fsdp_config.state_dict_type
+    if cfg.fsdp_config.cpu_offload_pin_memory is not None:
+        os.environ["FSDP_CPU_OFFLOAD_PIN_MEMORY"] = str(
+            cfg.fsdp_config.cpu_offload_pin_memory
+        ).lower()
+    if cfg.fsdp_config.auto_wrap_policy:
+        os.environ["FSDP_AUTO_WRAP_POLICY"] = cfg.fsdp_config.auto_wrap_policy
+    if cfg.fsdp_config.transformer_layer_cls_to_wrap:
+        os.environ["FSDP_TRANSFORMER_CLS_TO_WRAP"] = (
+            cfg.fsdp_config.transformer_layer_cls_to_wrap
+        )
+    if cfg.fsdp_config.min_num_params:
+        os.environ["FSDP_MIN_NUM_PARAMS"] = str(cfg.fsdp_config.min_num_params)
+    if cfg.fsdp_config.reshard_after_forward:
+        os.environ["FSDP_RESHARD_AFTER_FORWARD"] = "true"
+
+
+def setup_parallelism_envs(cfg):
+    set_accelerate_parallelism_config = False
+    if cfg.tensor_parallel_size and cfg.tensor_parallel_size > 1:
+        set_accelerate_parallelism_config = True
+        os.environ["PARALLELISM_CONFIG_TP_SIZE"] = str(cfg.tensor_parallel_size)
+    if cfg.dp_shard_size and cfg.dp_shard_size > 1:
+        set_accelerate_parallelism_config = True
+        os.environ["PARALLELISM_CONFIG_DP_SHARD_SIZE"] = str(cfg.dp_shard_size)
+    if cfg.dp_replicate_size and cfg.dp_replicate_size > 1:
+        set_accelerate_parallelism_config = True
+        os.environ["PARALLELISM_CONFIG_DP_REPLICATE_SIZE"] = str(cfg.dp_replicate_size)
+    if cfg.context_parallel_size and cfg.context_parallel_size > 1:
+        set_accelerate_parallelism_config = True
+        os.environ["PARALLELISM_CONFIG_CP_SIZE"] = str(cfg.context_parallel_size)
+        os.environ["ACCELERATE_ALLOW_CP_STANDALONE"] = "true"
+        from axolotl.monkeypatch.accelerate.parallelism_config import patch_prepare_cp
+
+        patch_prepare_cp()
+    # Expert Parallel patch must apply before the first `Accelerator()`
+    # call so `ep_size` lands in the mesh.
+    if cfg.expert_parallel_size and cfg.expert_parallel_size > 1:
+        os.environ["PARALLELISM_CONFIG_EP_SIZE"] = str(cfg.expert_parallel_size)
+        if cfg.dp_shard_size and cfg.dp_shard_size > 1:
+            set_accelerate_parallelism_config = True
+        from axolotl.monkeypatch.accelerate.parallelism_config import (
+            patch_parallelism_config,
+        )
+
+        patch_parallelism_config()
+    if set_accelerate_parallelism_config:
+        os.environ["ACCELERATE_USE_PARALLELISM_CONFIG"] = "true"
 
 
 def prepare_optim_env(cfg):
-    if cfg.fsdp:
+    if not check_cuda_p2p_ib_support():
+        if os.getenv("NCCL_P2P_DISABLE") is None:
+            LOG.warning("P2P support not detected, setting `NCCL_P2P_DISABLE=1`")
+            os.environ["NCCL_P2P_DISABLE"] = "1"
+    # TODO @SalmanMohammadi remove the cfg.fsdp check in 0.12
+    if cfg.fsdp or cfg.fsdp_config:
+        # fsdp_config is source of truth; mutating cfg.fsdp to bool breaks Ray worker schema validation
         setup_fsdp_envs(cfg)
     elif cfg.deepspeed:
-        os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
-        os.environ["ACCELERATE_DEEPSPEED_CONFIG_FILE"] = cfg.deepspeed
-
-    if (cfg.bf16 == "auto" and is_torch_bf16_gpu_available()) or cfg.bf16 is True:
+        stage = None
+        deepspeed_config = None
+        # check if the cfg.deepspeed is a file
+        if isinstance(cfg.deepspeed, DictDefault):
+            deepspeed_config = cfg.deepspeed
+        elif os.path.isfile(cfg.deepspeed):
+            # parse with json
+            with open(cfg.deepspeed, "r", encoding="utf-8") as fin:
+                deepspeed_config = json.load(fin)
+        if deepspeed_config:
+            stage = deepspeed_config.get("zero_optimization", {}).get("stage", None)
+        setup_deepspeed_env(cfg, stage=stage)
+
+    setup_parallelism_envs(cfg)
+    setup_torch_compile_env(cfg)
+
+    if cfg.fp8:
+        os.environ["ACCELERATE_MIXED_PRECISION"] = "fp8"
+    elif (cfg.bf16 == "auto" and is_torch_bf16_gpu_available()) or cfg.bf16 is True:
         os.environ["ACCELERATE_MIXED_PRECISION"] = "bf16"
     elif cfg.fp16:
         os.environ["ACCELERATE_MIXED_PRECISION"] = "fp16"
+    else:
+        os.environ["ACCELERATE_MIXED_PRECISION"] = "no"
+
+
+def setup_trainer(
+    cfg,
+    train_dataset,
+    eval_dataset,
+    model,
+    tokenizer,
+    processor,
+    total_num_steps,
+    model_ref=None,
+    peft_config=None,
+):
+    """
+    Helper method for instantiating and building a (causal or RLHF) trainer.
+
+    Args:
+        cfg: Axolotl config object containing training parameters.
+        train_dataset: Dataset to use for training.
+        eval_dataset: Dataset to use for evaluation.
+        model: The model to train.
+        tokenizer: Tokenizer for processing text input.
+        processor: Processor for data preparation.
+        total_num_steps: The total number of training steps.
+        model_ref: Optional reference model for RLHF training. Default is None.
+        peft_config: Optional PEFT (Parameter-Efficient Fine-Tuning) configuration. Default is None.
+
+    Returns:
+        A trainer instance (either `HFRLTrainer` or `HFCausalTrainer`) configured based
+            on the provided parameters.
+    """
+    from axolotl.core.builders import HFCausalTrainerBuilder, HFRLTrainerBuilder
 
-
-def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps):
-    if cfg.rl in ["dpo", "ipo", "kto_pair", "orpo"]:
-        trainer_builder = HFRLTrainerBuilder(cfg, model[0], tokenizer)
-        trainer_builder.model_ref = model[1]
-        trainer_builder.peft_config = model[2]
+    if cfg.rl:
+        trainer_builder = HFRLTrainerBuilder(cfg, model, tokenizer, processor)
+        trainer_builder.model_ref = model_ref
+        trainer_builder.peft_config = peft_config
     else:
-        trainer_builder = HFCausalTrainerBuilder(cfg, model[0], tokenizer)
+        trainer_builder = HFCausalTrainerBuilder(cfg, model, tokenizer, processor)
 
     trainer_builder.train_dataset = train_dataset
     trainer_builder.eval_dataset = eval_dataset
diff --git a/src/axolotl/utils/wandb_.py b/src/axolotl/utils/wandb_.py
index 327dd9b634..6484d435ae 100644
--- a/src/axolotl/utils/wandb_.py
+++ b/src/axolotl/utils/wandb_.py
@@ -16,6 +16,3 @@ def setup_wandb_env_vars(cfg: DictDefault):
     # Enable wandb if project name is present
     if cfg.wandb_project and len(cfg.wandb_project) > 0:
         cfg.use_wandb = True
-        os.environ.pop("WANDB_DISABLED", None)  # Remove if present
-    else:
-        os.environ["WANDB_DISABLED"] = "true"
diff --git a/src/axolotl/utils/weight_serde.py b/src/axolotl/utils/weight_serde.py
new file mode 100644
index 0000000000..d4b8046810
--- /dev/null
+++ b/src/axolotl/utils/weight_serde.py
@@ -0,0 +1,94 @@
+"""Serialize / deserialize tensors for HTTP and IPC weight sync.
+
+NumPy doesn't support bfloat16, so bf16 tensors are cast to fp16 on the wire
+and reconstructed at the destination.  All encode/decode helpers live here so
+the logic isn't duplicated across trl_vllm.py, vllm_serve_lora.py, and
+vllm_worker_ext.py.
+"""
+
+import base64
+
+import torch
+
+
+def encode_for_http(name: str, weight: torch.Tensor) -> dict:
+    """Encode a named parameter for JSON transport over HTTP.
+
+    Returns a dict with keys: name, dtype (original), shape, data (base64).
+    bf16 tensors are sent as fp16 bytes; the original dtype is preserved in
+    the ``dtype`` field so the receiver can cast back.
+    """
+    w_cpu = weight.contiguous().cpu()
+    orig_dtype = str(weight.dtype)
+    if w_cpu.dtype == torch.bfloat16:
+        w_cpu = w_cpu.half()
+    raw = w_cpu.numpy().tobytes()
+    return {
+        "name": name,
+        "dtype": orig_dtype,
+        "shape": list(weight.shape),
+        "data": base64.b64encode(raw).decode("ascii"),
+    }
+
+
+def decode_from_http(entry: dict) -> tuple[str, torch.Tensor]:
+    """Decode an HTTP-encoded weight entry back to a named tensor.
+
+    Infers wire dtype from byte count (bf16 arrives as fp16) and casts to the
+    original dtype stored in ``entry["dtype"]``.
+    """
+    target_dtype = getattr(torch, entry["dtype"].split(".")[-1])
+    shape = tuple(entry["shape"])
+    raw = base64.b64decode(entry["data"])
+
+    n_elements = 1
+    for s in shape:
+        n_elements *= s
+    wire_bytes_per_elem = len(raw) // max(n_elements, 1)
+    if wire_bytes_per_elem == 2:
+        wire_dtype = torch.float16
+    elif wire_bytes_per_elem == 4:
+        wire_dtype = torch.float32
+    else:
+        wire_dtype = target_dtype
+
+    weight = torch.frombuffer(bytearray(raw), dtype=wire_dtype).reshape(shape)
+    if wire_dtype != target_dtype:
+        weight = weight.to(target_dtype)
+    return entry["name"], weight
+
+
+def encode_for_ipc(name: str, weight: torch.Tensor) -> dict:
+    """Encode a tensor for vLLM's multiproc IPC (raw bytes, no base64).
+
+    Returns a dict with keys: name, data (bytes), dtype (wire), target_dtype
+    (original), shape.  bf16 tensors are serialized as fp16.
+    """
+    w = weight.contiguous()
+    target_dtype = str(w.dtype).split(".")[-1]
+    if w.dtype == torch.bfloat16:
+        w = w.half()
+    wire_dtype = str(w.dtype).split(".")[-1]
+    return {
+        "name": name,
+        "data": w.numpy().tobytes(),
+        "dtype": wire_dtype,
+        "target_dtype": target_dtype,
+        "shape": list(weight.shape),
+    }
+
+
+def decode_from_ipc(entry: dict) -> tuple[str, torch.Tensor]:
+    """Decode an IPC-encoded weight entry back to a named tensor.
+
+    Handles optional ``target_dtype`` for backward compatibility with older
+    serve code that may not include it.
+    """
+    wire_dtype = getattr(torch, entry["dtype"])
+    weight = torch.frombuffer(bytearray(entry["data"]), dtype=wire_dtype).reshape(
+        entry["shape"]
+    )
+    target_dtype = entry.get("target_dtype")
+    if target_dtype and target_dtype != entry["dtype"]:
+        weight = weight.to(getattr(torch, target_dtype))
+    return entry["name"], weight
diff --git a/styles.css b/styles.css
index 2ddf50c7b4..c5b0768faf 100644
--- a/styles.css
+++ b/styles.css
@@ -1 +1,277 @@
-/* css styles */
+/* TYPOGRAPHY SECTION */
+
+/* Import fonts */
+@import url('https://fonts.googleapis.com/css2?family=Be+Vietnam+Pro:wght@400;500&display=swap');
+@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400&display=swap');
+
+/* Typography hierarchy */
+:root {
+    --font-title: 'Be Vietnam Pro', sans-serif;
+    --font-body: 'JetBrains Mono', monospace;
+}
+
+/* Title (h1) */
+h1 {
+    font-family: var(--font-title);
+    font-weight: 400;
+    font-size: 3rem;
+    line-height: 1.1;
+    letter-spacing: -0.05em;
+    font-feature-settings: "ss01" on;
+}
+
+/* Heading (h2) */
+h2 {
+    font-family: var(--font-title);
+    font-weight: 500;
+    font-size: 1.5rem;
+    line-height: 1.2;
+    letter-spacing: -0.03em;
+    font-feature-settings: "ss01" on;
+}
+
+/* Subtitle/Preamble */
+h3,
+h4 {
+    font-family: var(--font-body);
+    font-weight: 400;
+    font-size: 1.25rem;
+    line-height: 1.5;
+    letter-spacing: -0.02em;
+}
+
+/* Body text */
+body {
+    font-family: var(--font-body);
+    font-weight: 400;
+    font-size: 1rem;
+    line-height: 1.5;
+    letter-spacing: -0.02em;
+}
+
+/* Links */
+a {
+    font-family: var(--font-body);
+    font-weight: 400;
+    font-size: 0.875rem;
+    line-height: 1;
+    letter-spacing: -0.02em;
+}
+
+/* NAV BAR SECTION */
+
+/* Navbar logo styling */
+.navbar-brand img {
+    height: 32px;
+    margin-right: 10px;
+}
+
+/* COLORS SECTION */
+
+/* Brand colors */
+:root {
+    --white: #ffffff;
+    --greige-300: #EEEEE7;
+    --greige-600: #CCCAC0;
+    --black: #141310;
+    --lime: #E3F8A8;
+    --cyan: #A0F4EA;
+    --purple: #C8D0F8;
+}
+
+/* Base styles */
+body {
+    background-color: var(--black);
+    color: var(--greige-300);
+}
+
+/* Navigation */
+.navbar {
+    background-color: var(--black) !important;
+}
+
+.navbar-dark .navbar-nav .nav-link {
+    color: var(--greige-300);
+}
+
+.navbar-dark .navbar-nav .nav-link:hover {
+    color: var(--lime);
+}
+
+/* Sidebar */
+.sidebar-navigation {
+    background-color: var(--black);
+    border-right: 1px solid var(--greige-600);
+}
+
+.sidebar nav[role="doc-toc"] ul>li>a {
+    color: var(--greige-300);
+}
+
+.sidebar nav[role="doc-toc"] ul>li>a:hover {
+    color: var(--lime);
+}
+
+/* Links */
+a {
+    color: var(--lime);
+}
+
+a:hover {
+    color: var(--cyan);
+}
+
+/* Headers */
+h1,
+h2,
+h3,
+h4,
+h5,
+h6 {
+    color: var(--white);
+}
+
+/* Code blocks */
+pre {
+    background-color: #1a1a1a !important;
+    border: 1px solid var(--greige-600);
+}
+
+/* Tables */
+.table {
+    color: var(--greige-300);
+}
+
+/* TOC */
+#toc-title {
+    color: var(--white);
+}
+
+.toc-active {
+    color: var(--lime) !important;
+}
+
+/* Buttons */
+.btn-primary {
+    background-color: var(--lime);
+    color: var(--black);
+    border: none;
+}
+
+.btn-primary:hover {
+    background-color: var(--cyan);
+    color: var(--black);
+}
+
+/* For inline code (single backtick) */
+code {
+    background-color: #1a1a1a !important;
+    color: var(--lime) !important;
+    padding: 2px 4px;
+    border-radius: 4px;
+}
+
+/* For inline code that is also a link */
+a code {
+    color: var(--cyan) !important;
+}
+
+/* For code blocks (triple backtick) */
+pre.sourceCode {
+    background-color: #1a1a1a !important;
+}
+
+/* Make comments in bash/shell scripts green */
+code span.co {
+    color: #5cb85c !important;
+}
+
+/* Remove underlines from JSON comments and make them green */
+code span.er {
+    color: #5cb85c !important;
+    text-decoration: none !important;
+}
+
+/* API Documentation Styling */
+
+/* Improve docstring section rendering */
+.level3 p {
+    white-space: pre-line !important;
+}
+
+/* Format docstring sections */
+.level3 p strong {
+    display: block;
+    margin-top: 1em;
+    font-weight: bold;
+    color: var(--cyan);
+}
+
+/* Add spacing after sections */
+.level3 p:has(strong) {
+    margin-bottom: 0.5em;
+}
+
+/* Format Args and Returns sections */
+p:has(code) {
+    line-height: 1.6;
+}
+
+/* Function signatures */
+.sourceCode {
+    margin-bottom: 1.5em;
+}
+
+/* Parameter tables */
+.doc-section-parameters table,
+.doc-section-returns table {
+    margin-top: 1em;
+    margin-bottom: 1.5em;
+}
+
+/* Make parameter and returns headers smaller */
+h2.anchored[data-anchor-id="parameters"],
+h2.anchored[data-anchor-id="returns"],
+.doc-section-parameters h4,
+.doc-section-returns h4 {
+    font-size: 1.25rem;
+    margin-top: 2rem;
+    margin-bottom: 1rem;
+    color: var(--lime);
+    border-bottom: 1px solid var(--lime);
+    padding-bottom: 0.3rem;
+    font-family: var(--font-body);
+    font-weight: 500;
+    letter-spacing: normal;
+}
+
+/* Style documentation tables */
+table {
+    width: 100%;
+    margin-bottom: 1.5rem;
+    border-collapse: collapse;
+}
+
+table th {
+    background-color: #1a1a1a;
+    padding: 0.5rem 1rem;
+    border-bottom: 2px solid var(--greige-600);
+    text-align: left;
+}
+
+table td {
+    padding: 0.5rem 1rem;
+    border-bottom: 1px solid var(--greige-600);
+}
+
+/* Code in table cells */
+table td code {
+    background-color: transparent !important;
+    padding: 0;
+}
+
+/* Improve spacing in parameter and return tables */
+.doc-section-parameters,
+.doc-section-returns {
+    margin-top: 1rem;
+}
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/cli/conftest.py b/tests/cli/conftest.py
new file mode 100644
index 0000000000..d360e29d6b
--- /dev/null
+++ b/tests/cli/conftest.py
@@ -0,0 +1,37 @@
+"""Shared pytest fixtures for cli module."""
+
+import pytest
+from click.testing import CliRunner
+
+VALID_TEST_CONFIG = """
+base_model: HuggingFaceTB/SmolLM2-135M
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+sequence_len: 2048
+max_steps: 1
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 1e-3
+special_tokens:
+  pad_token: <|endoftext|>
+"""
+
+
+@pytest.fixture
+def cli_runner():
+    return CliRunner()
+
+
+@pytest.fixture
+def valid_test_config():
+    return VALID_TEST_CONFIG
+
+
+@pytest.fixture
+def config_path(tmp_path):
+    """Creates a temporary config file"""
+    path = tmp_path / "config.yml"
+    path.write_text(VALID_TEST_CONFIG)
+
+    return path
diff --git a/tests/cli/test_chat_repl.py b/tests/cli/test_chat_repl.py
new file mode 100644
index 0000000000..9d6a1de646
--- /dev/null
+++ b/tests/cli/test_chat_repl.py
@@ -0,0 +1,668 @@
+"""pytest tests for the interactive chat REPL (no model required)."""
+
+import io
+import json
+
+import pytest
+from rich.console import Console
+
+from axolotl.cli.chat import (
+    CausalTurnGenerator,
+    ChatRepl,
+    ChatSession,
+    TurnResult,
+    default_gen_params,
+    longest_common_prefix_len,
+    parse_gen_param_value,
+    resolve_command,
+    resolve_gen_param,
+)
+
+
+class FakeCache:
+    """Stands in for DynamicCache in cache-planning tests."""
+
+    def __init__(self, length=0, croppable=True):
+        self.length = length
+        self.croppable = croppable
+
+    def crop(self, max_length):
+        if not self.croppable:
+            raise NotImplementedError("cannot crop")
+        self.length = max_length
+
+    def get_seq_length(self):
+        return self.length
+
+
+class FakeGenerator:
+    """Records conversations passed in and returns canned replies."""
+
+    def __init__(self, replies=None, messages=None):
+        self.replies = replies or ["canned reply"]
+        self.messages = messages
+        self.calls = []
+        self.render_kwargs_seen = []
+
+    def generate_turn(self, conversation, params, on_text, render_kwargs=None):
+        self.calls.append(([dict(m) for m in conversation], dict(params)))
+        self.render_kwargs_seen.append(
+            dict(render_kwargs) if render_kwargs is not None else None
+        )
+        index = min(len(self.calls) - 1, len(self.replies) - 1)
+        content = self.replies[index]
+        message = dict(self.messages[index]) if self.messages else None
+        on_text(content)
+        return TurnResult(
+            content=content, message=message, prompt_tokens=10, new_tokens=3
+        )
+
+
+def make_repl(inputs, generator=None, session=None):
+    lines = iter(inputs)
+
+    def input_fn(_prompt):
+        try:
+            return next(lines)
+        except StopIteration as err:
+            raise EOFError from err
+
+    generator = generator or FakeGenerator()
+    repl = ChatRepl(
+        generator=generator,
+        session=session,
+        console=Console(file=io.StringIO(), force_terminal=False),
+        input_fn=input_fn,
+    )
+    return repl, generator
+
+
+def cache_planner(cached_ids, cache):
+    generator = CausalTurnGenerator.__new__(CausalTurnGenerator)
+    generator._cache = cache  # pylint: disable=protected-access
+    generator._cached_ids = cached_ids  # pylint: disable=protected-access
+    generator._new_cache = FakeCache  # pylint: disable=protected-access
+    return generator
+
+
+class TestGenParams:
+    def test_alias_resolution(self):
+        assert resolve_gen_param("temp").key == "temperature"
+        assert resolve_gen_param("max").key == "max_new_tokens"
+        assert resolve_gen_param("rep").key == "repetition_penalty"
+        assert resolve_gen_param("bogus") is None
+
+    def test_value_validation(self):
+        spec = resolve_gen_param("temperature")
+        assert parse_gen_param_value(spec, "0.7") == 0.7
+        with pytest.raises(ValueError):
+            parse_gen_param_value(spec, "100")
+        with pytest.raises(ValueError):
+            parse_gen_param_value(spec, "abc")
+
+    def test_nullable_params(self):
+        assert parse_gen_param_value(resolve_gen_param("seed"), "none") is None
+        assert parse_gen_param_value(resolve_gen_param("min_p"), "off") is None
+        with pytest.raises(ValueError):
+            parse_gen_param_value(resolve_gen_param("temperature"), "none")
+
+
+class TestChatSession:
+    def test_system_prompt_prepended(self):
+        session = ChatSession()
+        session.system = "be brief"
+        session.add_user("hi")
+        conversation = session.conversation()
+        assert conversation[0] == {"role": "system", "content": "be brief"}
+        assert conversation[1]["role"] == "user"
+
+    def test_undo_removes_exchange(self):
+        session = ChatSession()
+        session.add_user("q1")
+        session.add_assistant("a1")
+        session.add_user("q2")
+        session.add_assistant("a2")
+        assert session.undo()
+        assert [m["content"] for m in session.messages] == ["q1", "a1"]
+        assert session.undo()
+        assert not session.messages
+        assert not session.undo()
+
+    def test_drop_last_assistant_for_retry(self):
+        session = ChatSession()
+        session.add_user("q1")
+        session.add_assistant("a1")
+        assert session.drop_last_assistant()
+        assert session.messages[-1]["role"] == "user"
+        session.clear()
+        assert not session.drop_last_assistant()
+
+    def test_add_user_merges_consecutive_user_messages(self):
+        # a failed generation leaves a trailing user message; typing again must
+        # not create consecutive user turns (strict templates reject them)
+        session = ChatSession()
+        session.add_user("first try")
+        session.add_user("second try")
+        assert [m["role"] for m in session.messages] == ["user"]
+        assert session.messages[0]["content"] == "first try\nsecond try"
+
+    def test_save_jsonl_keeps_reasoning_content(self, tmp_path):
+        session = ChatSession()
+        session.add_user("q")
+        session.add_assistant_message(
+            {"role": "assistant", "content": "a", "reasoning_content": "hmm"}
+        )
+        path = tmp_path / "chat.jsonl"
+        session.save_jsonl(str(path))
+        sample = json.loads(path.read_text(encoding="utf-8"))
+        assistant = sample["messages"][1]
+        assert assistant["content"] == [{"type": "text", "text": "a"}]
+        assert assistant["reasoning_content"] == "hmm"
+
+    def test_save_jsonl_multimodal_parts_format(self, tmp_path):
+        session = ChatSession()
+        session.system = "sys"
+        session.add_user("q")
+        session.add_assistant("a")
+        path = tmp_path / "chat.jsonl"
+        session.save_jsonl(str(path))
+        session.save_jsonl(str(path))
+        lines = path.read_text(encoding="utf-8").strip().split("\n")
+        assert len(lines) == 2
+        sample = json.loads(lines[0])
+        assert [m["role"] for m in sample["messages"]] == [
+            "system",
+            "user",
+            "assistant",
+        ]
+        assert sample["messages"][1]["content"] == [{"type": "text", "text": "q"}]
+        assert sample["messages"][2]["content"] == [{"type": "text", "text": "a"}]
+
+
+class TestCachePlanning:
+    def test_prefix_extension_reuses_cache(self):
+        cache = FakeCache(length=5)
+        generator = cache_planner([1, 2, 3, 4, 5], cache)
+        assert generator._prepare_cache([1, 2, 3, 4, 5, 6, 7]) == 5
+        assert generator._cache is cache
+
+    def test_divergence_crops_to_common_prefix(self):
+        cache = FakeCache(length=5)
+        generator = cache_planner([1, 2, 3, 4, 5], cache)
+        assert generator._prepare_cache([1, 2, 3, 9, 9, 9]) == 3
+        assert cache.length == 3
+        assert generator._cached_ids == [1, 2, 3]
+
+    def test_no_overlap_resets_cache(self):
+        cache = FakeCache(length=3)
+        generator = cache_planner([1, 2, 3], cache)
+        assert generator._prepare_cache([7, 8, 9]) == 0
+        assert generator._cache is not cache
+        assert generator._cached_ids == []
+
+    def test_uncroppable_cache_resets(self):
+        cache = FakeCache(length=5, croppable=False)
+        generator = cache_planner([1, 2, 3, 4, 5], cache)
+        assert generator._prepare_cache([1, 2, 3, 9, 9]) == 0
+        assert generator._cache is not cache
+
+    def test_render_fully_cached_leaves_one_input_token(self):
+        # cache covering all input tokens would give generate() nothing to process
+        cache = FakeCache(length=5)
+        generator = cache_planner([1, 2, 3, 4, 5], cache)
+        assert generator._prepare_cache([1, 2, 3, 4, 5]) == 4
+        assert cache.length == 4
+
+
+class TestChatRepl:
+    def test_message_generates_turn_with_history(self):
+        repl, generator = make_repl(["hi", "again", "/quit"])
+        repl.run()
+        assert len(generator.calls) == 2
+        second_conversation = generator.calls[1][0]
+        assert [m["role"] for m in second_conversation] == [
+            "user",
+            "assistant",
+            "user",
+        ]
+        assert repl.session.messages[-1]["content"] == "canned reply"
+
+    def test_command_aliases(self):
+        assert resolve_command("clear").name == "new"
+        assert resolve_command("reset").name == "new"
+        assert resolve_command("regen").name == "retry"
+        assert resolve_command("q").name == "quit"
+        assert resolve_command("?").name == "help"
+
+    def test_new_clears_history_keeps_system_and_params(self):
+        repl, generator = make_repl(
+            ["/system be brief", "/temp 0.5", "hi", "/new", "next", "/quit"]
+        )
+        repl.run()
+        assert repl.session.system == "be brief"
+        assert repl.params["temperature"] == 0.5
+        last_conversation = generator.calls[-1][0]
+        assert [m["role"] for m in last_conversation] == ["system", "user"]
+        assert last_conversation[1]["content"] == "next"
+
+    def test_param_shortcut_and_set_forms(self):
+        repl, _ = make_repl(
+            ["/temp 0.3", "/set top_k 10", "/set max_tokens=64", "/quit"]
+        )
+        repl.run()
+        assert repl.params["temperature"] == 0.3
+        assert repl.params["top_k"] == 10
+        assert repl.params["max_new_tokens"] == 64
+
+    def test_invalid_param_value_not_applied(self):
+        repl, _ = make_repl(["/temp 100", "/quit"])
+        repl.run()
+        assert repl.params["temperature"] == default_gen_params()["temperature"]
+
+    def test_retry_regenerates_last_turn(self):
+        repl, generator = make_repl(
+            ["hi", "/retry", "/quit"], generator=FakeGenerator(["first", "second"])
+        )
+        repl.run()
+        assert len(generator.calls) == 2
+        retry_conversation = generator.calls[1][0]
+        assert retry_conversation[-1] == {"role": "user", "content": "hi"}
+        assert repl.session.messages[-1]["content"] == "second"
+
+    def test_undo_command(self):
+        repl, _ = make_repl(["hi", "/undo", "/quit"])
+        repl.run()
+        assert not repl.session.messages
+
+    def test_multiline_input(self):
+        repl, generator = make_repl(["first line\\", "second line", "/quit"])
+        repl.run()
+        assert generator.calls[0][0][0]["content"] == "first line\nsecond line"
+
+    def test_generator_message_stored_in_history(self):
+        message = {
+            "role": "assistant",
+            "content": "The answer is 4.",
+            "reasoning_content": "step by step",
+        }
+        repl, _ = make_repl(
+            ["what is 2+2?", "/quit"],
+            generator=FakeGenerator(["The answer is 4."], messages=[message]),
+        )
+        repl.run()
+        assert repl.session.messages[-1] == message
+        # renderer saw no think markers, so /expand falls back to the message
+        assert repl.last_think_text == "step by step"
+
+    def test_legacy_content_fallback_kept_verbatim(self):
+        # generators that return no message dict store their content as-is
+        reply = "<think>step by step</think>\nThe answer is 4."
+        repl, _ = make_repl(["what is 2+2?", "/quit"], generator=FakeGenerator([reply]))
+        repl.run()
+        assert repl.session.messages[-1]["content"] == reply
+
+    def test_unknown_command_does_not_generate(self):
+        repl, generator = make_repl(["/bogus", "/quit"])
+        repl.run()
+        assert not generator.calls
+
+    def test_command_handler_error_does_not_crash_repl(self):
+        # unclosed quote makes shlex raise inside /save
+        repl, generator = make_repl(["hi", '/save "unclosed', "again", "/quit"])
+        repl.run()
+        assert len(generator.calls) == 2
+
+    def test_keyboard_interrupt_keeps_session_alive(self):
+        class InterruptingGenerator(FakeGenerator):
+            def generate_turn(self, conversation, params, on_text, render_kwargs=None):
+                if not self.calls:
+                    self.calls.append(None)
+                    raise KeyboardInterrupt
+                return super().generate_turn(
+                    conversation, params, on_text, render_kwargs
+                )
+
+        repl, generator = make_repl(
+            ["hi", "again", "/quit"], generator=InterruptingGenerator()
+        )
+        repl.run()
+        # interrupted turn keeps the user message; the next one merges into it
+        assert [m["role"] for m in repl.session.messages] == ["user", "assistant"]
+        assert repl.session.messages[0]["content"] == "hi\nagain"
+        assert len(generator.calls) == 2
+
+    def test_generation_failure_keeps_session_alive(self):
+        class FailingGenerator(FakeGenerator):
+            def generate_turn(self, conversation, params, on_text, render_kwargs=None):
+                if not self.calls:
+                    self.calls.append(None)
+                    raise RuntimeError("boom")
+                return super().generate_turn(
+                    conversation, params, on_text, render_kwargs
+                )
+
+        repl, _ = make_repl(["hi", "/retry", "/quit"], generator=FailingGenerator())
+        repl.run()
+        assert [m["role"] for m in repl.session.messages] == ["user", "assistant"]
+        assert repl.session.messages[-1]["content"] == "canned reply"
+
+
+def test_longest_common_prefix_len():
+    assert longest_common_prefix_len([], [1, 2]) == 0
+    assert longest_common_prefix_len([1, 2], [1, 2]) == 2
+    assert longest_common_prefix_len([1, 2, 3], [1, 2]) == 2
+    assert longest_common_prefix_len([1, 9], [1, 2, 3]) == 1
+
+
+class TestDiffusionChat:
+    def test_diffusion_param_specs(self):
+        from axolotl.cli.chat import DIFFUSION_GEN_PARAMS
+
+        lines = iter(["/steps 32", "/tokens 64", "/top_p 0.9", "/quit"])
+
+        def input_fn(_prompt):
+            try:
+                return next(lines)
+            except StopIteration as err:
+                raise EOFError from err
+
+        repl = ChatRepl(
+            generator=FakeGenerator(),
+            param_specs=DIFFUSION_GEN_PARAMS,
+            console=Console(file=io.StringIO(), force_terminal=False),
+            input_fn=input_fn,
+        )
+        repl.run()
+        assert repl.params["steps"] == 32
+        assert repl.params["max_new_tokens"] == 64
+        assert "top_p" not in repl.params
+
+    def test_diffusion_turn_cuts_at_eos(self, monkeypatch):
+        from types import SimpleNamespace
+
+        import axolotl.integrations.diffusion as diffusion_module
+        from axolotl.cli.chat import (
+            DIFFUSION_GEN_PARAMS,
+            DiffusionTurnGenerator,
+            default_gen_params,
+        )
+
+        class FakeTokenizer:
+            eos_token_id = 2
+
+            def apply_chat_template(self, conversation, **kwargs):
+                return {"input_ids": [1, 5, 6]}
+
+            def decode(self, ids, **kwargs):
+                return ",".join(str(i) for i in ids)
+
+        fake_model = SimpleNamespace(
+            generation_config=SimpleNamespace(eos_token_id=None)
+        )
+
+        def fake_generate(model, tokenizer, **kwargs):
+            assert kwargs["mode"] == "completion"
+            assert kwargs["completion_tokens"] == 256
+            return {"generated_ids": [1, 5, 6, 7, 8, 2, 4]}
+
+        monkeypatch.setattr(diffusion_module, "generate", fake_generate)
+
+        generator = DiffusionTurnGenerator(
+            fake_model, FakeTokenizer(), None, "cpu", mask_token_id=9
+        )
+        chunks = []
+        result = generator.generate_turn(
+            [{"role": "user", "content": "hi"}],
+            default_gen_params(DIFFUSION_GEN_PARAMS),
+            chunks.append,
+        )
+        assert result.content == "7,8"
+        assert result.new_tokens == 2
+        assert result.prompt_tokens == 3
+        assert chunks == ["7,8"]
+
+
+def test_unknown_command_suggests_alias():
+    buf = io.StringIO()
+    repl = ChatRepl(
+        generator=FakeGenerator(),
+        console=Console(file=buf, force_terminal=False, width=200),
+        input_fn=lambda _p: "/quit",
+    )
+    repl._dispatch("/clea")
+    assert "Did you mean /clear?" in buf.getvalue()
+    repl._dispatch("/tem")
+    assert "Did you mean /temp?" in buf.getvalue()
+
+
+class TestThinkStreamRenderer:
+    def make_renderer(self, collapse=True, markers=("<think>", "</think>")):
+        from axolotl.cli.chat import ThinkStreamRenderer
+
+        buf = io.StringIO()
+        console = Console(file=buf, force_terminal=False, width=200)
+        return ThinkStreamRenderer(console, collapse=collapse, markers=markers), buf
+
+    def test_collapse_splits_thinking_from_reply(self, capsys):
+        renderer, buf = self.make_renderer()
+        for chunk in ["<thi", "nk>\nreasoning he", "re</th", "ink>\n\nAnswer!"]:
+            renderer.feed(chunk)
+        renderer.finish()
+        assert renderer.think_text.strip() == "reasoning here"
+        assert capsys.readouterr().out == "Answer!"
+        assert "thought for" in buf.getvalue()
+
+    def test_no_thinking_passthrough(self, capsys):
+        renderer, buf = self.make_renderer()
+        renderer.feed("Just a plain reply")
+        renderer.finish()
+        assert renderer.think_text == ""
+        assert capsys.readouterr().out == "Just a plain reply"
+        assert "thought for" not in buf.getvalue()
+
+    def test_unterminated_thinking(self, capsys):
+        renderer, buf = self.make_renderer()
+        renderer.feed("<think>partial reasoning")
+        renderer.finish()
+        assert renderer.think_text == "partial reasoning"
+        assert capsys.readouterr().out == ""
+        assert "no </think>" in buf.getvalue()
+
+    def test_collapse_off_is_passthrough(self, capsys):
+        renderer, buf = self.make_renderer(collapse=False)
+        renderer.feed("<think>abc</think>reply")
+        renderer.finish()
+        assert capsys.readouterr().out == "<think>abc</think>reply"
+        assert buf.getvalue() == ""
+
+    def test_custom_markers(self, capsys):
+        renderer, _ = self.make_renderer(
+            markers=("<|START_THINKING|>", "<|END_THINKING|>")
+        )
+        renderer.feed("<|START_THINKING|>hmm<|END_THINKING|>ok")
+        renderer.finish()
+        assert renderer.think_text == "hmm"
+        assert capsys.readouterr().out == "ok"
+
+
+class TestThinkTokenSplit:
+    def make_generator(self, vocab):
+        from types import SimpleNamespace
+
+        from axolotl.cli.chat import TurnGenerator
+
+        class FakeTokenizer:
+            eos_token_id = 0
+            chat_template = None
+
+            def encode(self, text, **kwargs):
+                return vocab[text]
+
+        model = SimpleNamespace(generation_config=SimpleNamespace(eos_token_id=None))
+        return TurnGenerator(model, FakeTokenizer(), None, "cpu")
+
+    def test_split_counts(self):
+        generator = self.make_generator({"<think>": [100], "</think>": [101]})
+        assert generator.split_think_token_counts([100, 1, 2, 3, 101, 7, 8]) == (3, 2)
+        assert generator.split_think_token_counts([100, 1, 2]) == (2, 0)
+        assert generator.split_think_token_counts([5, 6]) == (0, 2)
+        assert generator.split_think_token_counts([]) == (0, 0)
+
+
+class TestBuildAssistantMessage:
+    VOCAB = {
+        1: "step by step",
+        2: "The answer is 4.",
+        50: "<eos>",
+        100: "<think>",
+        101: "</think>",
+    }
+    SPECIAL = {50, 100, 101}
+
+    def make_generator(self, response_schema=None, parse_response=None):
+        from types import SimpleNamespace
+
+        from axolotl.cli.chat import TurnGenerator
+
+        vocab, special = self.VOCAB, self.SPECIAL
+
+        class FakeTokenizer:
+            eos_token_id = 50
+            chat_template = None
+
+            def encode(self, text, **kwargs):
+                return [token_id for token_id, t in vocab.items() if t == text]
+
+            def decode(self, ids, skip_special_tokens=False):
+                return "".join(
+                    vocab[i] for i in ids if not (skip_special_tokens and i in special)
+                )
+
+        tokenizer = FakeTokenizer()
+        if response_schema is not None:
+            tokenizer.response_schema = response_schema
+            tokenizer.parse_response = parse_response
+        model = SimpleNamespace(generation_config=SimpleNamespace(eos_token_id=None))
+        return TurnGenerator(model, tokenizer, None, "cpu")
+
+    def test_thinking_split_into_reasoning_content(self):
+        generator = self.make_generator()
+        message = generator.build_assistant_message([100, 1, 101, 2])
+        assert message == {
+            "role": "assistant",
+            "content": "The answer is 4.",
+            "reasoning_content": "step by step",
+        }
+
+    def test_no_thinking_omits_reasoning_key(self):
+        generator = self.make_generator()
+        message = generator.build_assistant_message([2])
+        assert message == {"role": "assistant", "content": "The answer is 4."}
+
+    def test_special_tokens_stripped_from_content(self):
+        generator = self.make_generator()
+        message = generator.build_assistant_message([2, 50])
+        assert message["content"] == "The answer is 4."
+
+    def test_parse_response_schema_preferred(self):
+        generator = self.make_generator(
+            response_schema={"x": "regex"},
+            parse_response=lambda text: {"content": "parsed", "thinking": "hmm"},
+        )
+        message = generator.build_assistant_message([2])
+        assert message == {
+            "role": "assistant",
+            "content": "parsed",
+            "thinking": "hmm",
+        }
+
+    def test_parse_response_failure_falls_back_to_markers(self):
+        def boom(text):
+            raise ValueError("bad schema")
+
+        generator = self.make_generator(
+            response_schema={"x": "regex"}, parse_response=boom
+        )
+        message = generator.build_assistant_message([100, 1, 101, 2])
+        assert message["content"] == "The answer is 4."
+        assert message["reasoning_content"] == "step by step"
+
+
+class TestEosTextTrimmer:
+    def make_trimmer(self, eos_strings=("<|im_end|>",)):
+        from axolotl.cli.chat import EosTextTrimmer
+
+        chunks = []
+        return EosTextTrimmer(eos_strings, chunks.append), chunks
+
+    def test_eos_marker_never_emitted(self):
+        trimmer, chunks = self.make_trimmer()
+        trimmer.feed("Hello")
+        trimmer.feed(" world<|im_end|>")
+        trimmer.finish()
+        assert "".join(chunks) == "Hello world"
+
+    def test_eos_split_across_chunks(self):
+        trimmer, chunks = self.make_trimmer()
+        trimmer.feed("Hi<|im_")
+        trimmer.feed("end|>")
+        trimmer.finish()
+        assert "".join(chunks) == "Hi"
+
+    def test_false_partial_released(self):
+        trimmer, chunks = self.make_trimmer()
+        trimmer.feed("a<")
+        trimmer.feed("b")
+        trimmer.finish()
+        assert "".join(chunks) == "a<b"
+
+    def test_plain_text_passthrough(self):
+        trimmer, chunks = self.make_trimmer()
+        trimmer.feed("no special tokens here")
+        trimmer.finish()
+        assert "".join(chunks) == "no special tokens here"
+
+
+class TestThinkCommands:
+    def test_think_toggle_sets_render_kwargs(self):
+        repl, generator = make_repl(
+            ["/think off", "hi", "/think default", "again", "/quit"]
+        )
+        repl.think_toggle_key = "enable_thinking"
+        repl.run()
+        assert generator.render_kwargs_seen[0] == {"enable_thinking": False}
+        assert generator.render_kwargs_seen[1] is None
+
+    def test_think_without_toggle_key(self):
+        repl, generator = make_repl(["/think off", "hi", "/quit"])
+        repl.run()
+        assert generator.render_kwargs_seen[0] is None
+
+    def test_collapse_toggle_and_expand(self):
+        reply = "<think>secret reasoning</think>\nAnswer."
+        repl, _ = make_repl(
+            ["hi", "/expand", "/collapse off", "/quit"],
+            generator=FakeGenerator([reply]),
+        )
+        repl.run()
+        assert repl.last_think_text == "secret reasoning"
+        assert repl.collapse_thinking is False
+        # raw content with thinking still stored in history
+        assert repl.session.messages[-1]["content"] == reply
+
+
+def test_detect_think_markers_and_toggle_key():
+    from axolotl.cli.chat import detect_think_markers, detect_think_toggle_key
+
+    qwen_like = "{% if enable_thinking %}...{{ '</think>' }}{% endif %}"
+    assert detect_think_toggle_key(qwen_like) == "enable_thinking"
+    assert detect_think_markers(qwen_like) == ("<think>", "</think>")
+
+    command_a_like = "...<|START_THINKING|>...<|END_THINKING|>..."
+    assert detect_think_markers(command_a_like) == (
+        "<|START_THINKING|>",
+        "<|END_THINKING|>",
+    )
+    assert detect_think_toggle_key(command_a_like) is None
+    assert detect_think_toggle_key("{% if thinking %}x{% endif %}") == "thinking"
+    assert detect_think_toggle_key(None) is None
diff --git a/tests/cli/test_cli_base.py b/tests/cli/test_cli_base.py
new file mode 100644
index 0000000000..e28bbb75ce
--- /dev/null
+++ b/tests/cli/test_cli_base.py
@@ -0,0 +1,93 @@
+"""Base test class for CLI commands."""
+
+from pathlib import Path
+from unittest.mock import patch
+
+from axolotl.cli.main import cli
+
+
+class BaseCliTest:
+    """Base class for CLI command tests."""
+
+    def _test_cli_validation(self, cli_runner, command: str):
+        """Test CLI validation for a command.
+
+        Args:
+            cli_runner: CLI runner fixture
+            command: Command to test (train/evaluate)
+        """
+        # Test missing config file
+        result = cli_runner.invoke(cli, [command, "--launcher", "python"])
+        assert result.exit_code != 0
+
+        # Test non-existent config file
+        result = cli_runner.invoke(
+            cli, [command, "nonexistent.yml", "--launcher", "python"]
+        )
+        assert result.exit_code != 0
+        assert "Error: Invalid value for 'CONFIG'" in result.output
+
+    def _test_basic_execution(
+        self,
+        cli_runner,
+        tmp_path: Path,
+        valid_test_config: str,
+        command: str,
+        train: bool = True,
+    ):
+        """Test basic execution with accelerate.
+
+        Args:
+            cli_runner: CLI runner fixture
+            tmp_path: Temporary path fixture
+            valid_test_config: Valid config fixture
+            command: Command to test (train/evaluate)
+            train: Whether to test training (default) or evaluation
+        """
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        mock_fn = "os.execvpe" if command == "train" else "subprocess.run"
+
+        with patch(mock_fn) as mock:
+            result = cli_runner.invoke(cli, [command, str(config_path)])
+
+            assert mock.called
+
+            expected = [
+                "accelerate",
+                "launch",
+                "-m",
+                f"axolotl.cli.{command}",
+                str(config_path),
+                "--debug=False",
+                "--debug-text-only=False",
+                "--debug-num-examples=0",
+            ]
+            if train:
+                expected.append("--shard=False")
+
+            if command == "train":
+                assert mock.call_args.args[0] == "accelerate"
+                assert mock.call_args.args[1] == expected
+            else:
+                assert mock.call_args.args[0] == expected
+                assert mock.call_args.kwargs == {"check": True}
+            assert result.exit_code == 0
+
+    def _test_cli_overrides(self, tmp_path: Path, valid_test_config: str):
+        """Test CLI argument overrides.
+
+        Args:
+            tmp_path: Temporary path fixture
+            valid_test_config: Valid config fixture
+            command: Command to test (train/evaluate)
+        """
+        config_path = tmp_path / "config.yml"
+        output_dir = tmp_path / "model-out"
+
+        test_config = valid_test_config.replace(
+            "output_dir: model-out", f"output_dir: {output_dir}"
+        )
+        config_path.write_text(test_config)
+        return config_path
diff --git a/tests/cli/test_cli_evaluate.py b/tests/cli/test_cli_evaluate.py
new file mode 100644
index 0000000000..e8b88625a1
--- /dev/null
+++ b/tests/cli/test_cli_evaluate.py
@@ -0,0 +1,172 @@
+"""Tests for evaluate CLI command."""
+
+from unittest.mock import patch
+
+from axolotl.cli.main import cli
+
+from .test_cli_base import BaseCliTest
+
+
+class TestEvaluateCommand(BaseCliTest):
+    """Test cases for evaluate command."""
+
+    cli = cli
+
+    def test_evaluate_cli_validation(self, cli_runner):
+        """Test CLI validation"""
+        self._test_cli_validation(cli_runner, "evaluate")
+
+    def test_evaluate_basic_execution(self, cli_runner, tmp_path, valid_test_config):
+        """Test basic successful execution"""
+        self._test_basic_execution(
+            cli_runner, tmp_path, valid_test_config, "evaluate", train=False
+        )
+
+    def test_evaluate_basic_execution_no_accelerate(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test basic successful execution without accelerate"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("axolotl.cli.evaluate.do_evaluate") as mock_evaluate:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "evaluate",
+                    str(config_path),
+                    "--launcher",
+                    "python",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_evaluate.assert_called_once()
+
+    def test_evaluate_cli_overrides(self, cli_runner, tmp_path, valid_test_config):
+        """Test CLI arguments properly override config values"""
+        config_path = self._test_cli_overrides(tmp_path, valid_test_config)
+
+        with patch("axolotl.cli.evaluate.do_evaluate") as mock_evaluate:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "evaluate",
+                    str(config_path),
+                    "--micro-batch-size",
+                    "2",
+                    "--sequence-len",
+                    "128",
+                    "--launcher",
+                    "python",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_evaluate.assert_called_once()
+            cfg = mock_evaluate.call_args[0][0]
+            assert cfg.micro_batch_size == 2
+            assert cfg.sequence_len == 128
+
+    def test_evaluate_with_launcher_args_torchrun(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test evaluate with torchrun launcher arguments"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("subprocess.run") as mock_subprocess:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "evaluate",
+                    str(config_path),
+                    "--launcher",
+                    "torchrun",
+                    "--",
+                    "--nproc_per_node=2",
+                    "--nnodes=1",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_subprocess.assert_called_once()
+
+            # Verify launcher args are passed to torchrun
+            called_cmd = mock_subprocess.call_args.args[0]
+            assert called_cmd[0] == "torchrun"
+            assert "--nproc_per_node=2" in called_cmd
+            assert "--nnodes=1" in called_cmd
+            assert "-m" in called_cmd
+            assert "axolotl.cli.evaluate" in called_cmd
+
+    def test_evaluate_with_launcher_args_accelerate(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test evaluate with accelerate launcher arguments"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("subprocess.run") as mock_subprocess:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "evaluate",
+                    str(config_path),
+                    "--launcher",
+                    "accelerate",
+                    "--",
+                    "--config_file=accelerate_config.yml",
+                    "--num_processes=4",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_subprocess.assert_called_once()
+
+            # Verify launcher args are passed to accelerate
+            called_cmd = mock_subprocess.call_args.args[0]
+            assert called_cmd[0] == "accelerate"
+            assert called_cmd[1] == "launch"
+            assert "--config_file=accelerate_config.yml" in called_cmd
+            assert "--num_processes=4" in called_cmd
+            assert "-m" in called_cmd
+            assert "axolotl.cli.evaluate" in called_cmd
+
+    def test_evaluate_backward_compatibility_no_launcher_args(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test that existing evaluate commands work without launcher args"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("subprocess.run") as mock_subprocess:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "evaluate",
+                    str(config_path),
+                    "--launcher",
+                    "accelerate",
+                    "--micro-batch-size",
+                    "2",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_subprocess.assert_called_once()
+
+            # Verify no launcher args contamination
+            called_cmd = mock_subprocess.call_args.args[0]
+            assert called_cmd[0] == "accelerate"
+            assert called_cmd[1] == "launch"
+            # Should not contain any extra launcher args
+            launcher_section = called_cmd[2 : called_cmd.index("-m")]
+            assert (
+                len(launcher_section) == 0
+            )  # No launcher args between 'launch' and '-m'
diff --git a/tests/cli/test_cli_fetch.py b/tests/cli/test_cli_fetch.py
new file mode 100644
index 0000000000..f06f067173
--- /dev/null
+++ b/tests/cli/test_cli_fetch.py
@@ -0,0 +1,39 @@
+"""pytest tests for axolotl CLI fetch command."""
+
+from unittest.mock import patch
+
+from axolotl.cli.main import fetch
+
+
+def test_fetch_cli_examples(cli_runner):
+    """Test fetch command with examples directory"""
+    with patch("axolotl.cli.main.fetch_from_github") as mock_fetch:
+        result = cli_runner.invoke(fetch, ["examples"])
+
+        assert result.exit_code == 0
+        mock_fetch.assert_called_once_with("examples/", None)
+
+
+def test_fetch_cli_deepspeed(cli_runner):
+    """Test fetch command with deepspeed_configs directory"""
+    with patch("axolotl.cli.main.fetch_from_github") as mock_fetch:
+        result = cli_runner.invoke(fetch, ["deepspeed_configs"])
+
+        assert result.exit_code == 0
+        mock_fetch.assert_called_once_with("deepspeed_configs/", None)
+
+
+def test_fetch_cli_with_dest(cli_runner, tmp_path):
+    """Test fetch command with custom destination"""
+    with patch("axolotl.cli.main.fetch_from_github") as mock_fetch:
+        custom_dir = tmp_path / "tmp_examples"
+        result = cli_runner.invoke(fetch, ["examples", "--dest", str(custom_dir)])
+
+        assert result.exit_code == 0
+        mock_fetch.assert_called_once_with("examples/", str(custom_dir))
+
+
+def test_fetch_cli_invalid_directory(cli_runner):
+    """Test fetch command with invalid directory choice"""
+    result = cli_runner.invoke(fetch, ["invalid"])
+    assert result.exit_code != 0
diff --git a/tests/cli/test_cli_inference.py b/tests/cli/test_cli_inference.py
new file mode 100644
index 0000000000..8728db23d1
--- /dev/null
+++ b/tests/cli/test_cli_inference.py
@@ -0,0 +1,185 @@
+"""pytest tests for axolotl CLI inference command."""
+
+from unittest.mock import patch
+
+from axolotl.cli.main import cli
+
+
+def test_inference_basic(cli_runner, config_path):
+    """Test basic inference"""
+    with patch("axolotl.cli.inference.do_inference") as mock:
+        result = cli_runner.invoke(
+            cli,
+            ["inference", str(config_path), "--launcher", "python"],
+            catch_exceptions=False,
+        )
+
+        assert mock.called
+        assert result.exit_code == 0
+
+
+def test_inference_gradio(cli_runner, config_path):
+    """Test basic inference (gradio path)"""
+    with patch("axolotl.cli.inference.do_inference_gradio") as mock:
+        result = cli_runner.invoke(
+            cli,
+            ["inference", str(config_path), "--launcher", "python", "--gradio"],
+            catch_exceptions=False,
+        )
+
+        assert mock.called
+        assert result.exit_code == 0
+
+
+def test_inference_with_launcher_args_torchrun(cli_runner, config_path):
+    """Test inference with torchrun launcher arguments"""
+    with patch("subprocess.run") as mock_subprocess:
+        result = cli_runner.invoke(
+            cli,
+            [
+                "inference",
+                str(config_path),
+                "--launcher",
+                "torchrun",
+                "--",
+                "--nproc_per_node=2",
+                "--nnodes=1",
+            ],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_subprocess.assert_called_once()
+
+        # Verify launcher args are passed to torchrun
+        called_cmd = mock_subprocess.call_args.args[0]
+        assert called_cmd[0] == "torchrun"
+        assert "--nproc_per_node=2" in called_cmd
+        assert "--nnodes=1" in called_cmd
+        assert "-m" in called_cmd
+        assert "axolotl.cli.inference" in called_cmd
+
+
+def test_inference_with_launcher_args_accelerate(cli_runner, config_path):
+    """Test inference with accelerate launcher arguments"""
+    with patch("subprocess.run") as mock_subprocess:
+        result = cli_runner.invoke(
+            cli,
+            [
+                "inference",
+                str(config_path),
+                "--launcher",
+                "accelerate",
+                "--",
+                "--config_file=accelerate_config.yml",
+                "--num_processes=4",
+            ],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_subprocess.assert_called_once()
+
+        # Verify launcher args are passed to accelerate
+        called_cmd = mock_subprocess.call_args.args[0]
+        assert called_cmd[0] == "accelerate"
+        assert called_cmd[1] == "launch"
+        assert "--config_file=accelerate_config.yml" in called_cmd
+        assert "--num_processes=4" in called_cmd
+        assert "-m" in called_cmd
+        assert "axolotl.cli.inference" in called_cmd
+
+
+def test_inference_gradio_with_launcher_args(cli_runner, config_path):
+    """Test inference with gradio and launcher arguments"""
+    with patch("subprocess.run") as mock_subprocess:
+        result = cli_runner.invoke(
+            cli,
+            [
+                "inference",
+                str(config_path),
+                "--launcher",
+                "accelerate",
+                "--gradio",
+                "--",
+                "--num_processes=2",
+            ],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_subprocess.assert_called_once()
+
+        # Verify both gradio flag and launcher args are present
+        called_cmd = mock_subprocess.call_args.args[0]
+        assert called_cmd[0] == "accelerate"
+        assert called_cmd[1] == "launch"
+        assert "--num_processes=2" in called_cmd
+        assert "--gradio" in called_cmd
+        assert "-m" in called_cmd
+        assert "axolotl.cli.inference" in called_cmd
+
+
+def test_inference_backward_compatibility_no_launcher_args(cli_runner, config_path):
+    """Test that existing inference commands work without launcher args"""
+    with patch("subprocess.run") as mock_subprocess:
+        result = cli_runner.invoke(
+            cli,
+            [
+                "inference",
+                str(config_path),
+                "--launcher",
+                "accelerate",
+            ],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_subprocess.assert_called_once()
+
+        # Verify no launcher args contamination
+        called_cmd = mock_subprocess.call_args.args[0]
+        assert called_cmd[0] == "accelerate"
+        assert called_cmd[1] == "launch"
+        # Should not contain any extra launcher args
+        launcher_section = called_cmd[2 : called_cmd.index("-m")]
+        assert len(launcher_section) == 0  # No launcher args between 'launch' and '-m'
+
+
+def test_inference_chat(cli_runner, config_path):
+    """Test basic inference (chat path)"""
+    with patch("axolotl.cli.chat.do_chat") as mock:
+        result = cli_runner.invoke(
+            cli,
+            ["inference", str(config_path), "--launcher", "python", "--chat"],
+            catch_exceptions=False,
+        )
+
+        assert mock.called
+        assert result.exit_code == 0
+
+
+def test_inference_chat_with_launcher(cli_runner, config_path):
+    """Test chat flag is forwarded through the launcher command"""
+    with patch("subprocess.run") as mock_subprocess:
+        result = cli_runner.invoke(
+            cli,
+            ["inference", str(config_path), "--launcher", "accelerate", "--chat"],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        called_cmd = mock_subprocess.call_args.args[0]
+        assert "--chat" in called_cmd
+        assert "axolotl.cli.inference" in called_cmd
+
+
+def test_inference_chat_gradio_mutually_exclusive(cli_runner, config_path):
+    """Test that --chat and --gradio cannot be combined"""
+    result = cli_runner.invoke(
+        cli,
+        ["inference", str(config_path), "--chat", "--gradio"],
+    )
+
+    assert result.exit_code != 0
+    assert "mutually exclusive" in result.output
diff --git a/tests/cli/test_cli_interface.py b/tests/cli/test_cli_interface.py
new file mode 100644
index 0000000000..ebd91ea60e
--- /dev/null
+++ b/tests/cli/test_cli_interface.py
@@ -0,0 +1,47 @@
+"""General pytest tests for axolotl.cli.main interface."""
+
+from axolotl.cli.main import build_command, cli
+
+
+def test_build_command():
+    """Test converting dict of options to CLI arguments"""
+    base_cmd = ["accelerate", "launch"]
+    options = {
+        "learning_rate": 1e-4,
+        "batch_size": 8,
+        "debug": True,
+        "use_fp16": False,
+        "null_value": None,
+    }
+
+    result = build_command(base_cmd, options)
+    assert result == [
+        "accelerate",
+        "launch",
+        "--learning-rate=0.0001",
+        "--batch-size=8",
+        "--debug=True",
+        "--use-fp16=False",
+    ]
+
+
+def test_invalid_command_options(cli_runner):
+    """Test handling of invalid command options"""
+    result = cli_runner.invoke(
+        cli,
+        [
+            "train",
+            "config.yml",
+            "--invalid-option",
+            "value",
+        ],
+    )
+    assert result.exit_code != 0
+    assert "does not exist" in result.output
+
+
+def test_required_config_argument(cli_runner):
+    """Test commands fail properly when config argument is missing"""
+    result = cli_runner.invoke(cli, ["train"])
+    assert result.exit_code != 0
+    assert "Missing argument 'CONFIG'" in result.output
diff --git a/tests/cli/test_cli_merge_lora.py b/tests/cli/test_cli_merge_lora.py
new file mode 100644
index 0000000000..aac0167603
--- /dev/null
+++ b/tests/cli/test_cli_merge_lora.py
@@ -0,0 +1,57 @@
+"""pytest tests for axolotl CLI merge_lora command."""
+
+from unittest.mock import patch
+
+from axolotl.cli.main import cli
+
+
+def test_merge_lora_basic(cli_runner, config_path):
+    """Test basic merge_lora command"""
+    with patch("axolotl.cli.merge_lora.do_cli") as mock_do_cli:
+        result = cli_runner.invoke(cli, ["merge-lora", str(config_path)])
+        assert result.exit_code == 0
+
+        mock_do_cli.assert_called_once()
+        assert mock_do_cli.call_args.kwargs["config"] == str(config_path)
+
+
+def test_merge_lora_with_dirs(cli_runner, config_path, tmp_path):
+    """Test merge_lora with custom lora and output directories"""
+    lora_dir = tmp_path / "lora"
+    output_dir = tmp_path / "output"
+    lora_dir.mkdir()
+
+    with patch("axolotl.cli.merge_lora.do_cli") as mock_do_cli:
+        result = cli_runner.invoke(
+            cli,
+            [
+                "merge-lora",
+                str(config_path),
+                "--lora-model-dir",
+                str(lora_dir),
+                "--output-dir",
+                str(output_dir),
+            ],
+        )
+        assert result.exit_code == 0
+
+        mock_do_cli.assert_called_once()
+        assert mock_do_cli.call_args.kwargs["config"] == str(config_path)
+        assert mock_do_cli.call_args.kwargs["lora_model_dir"] == str(lora_dir)
+        assert mock_do_cli.call_args.kwargs["output_dir"] == str(output_dir)
+
+
+def test_merge_lora_nonexistent_config(cli_runner, tmp_path):
+    """Test merge_lora with nonexistent config"""
+    config_path = tmp_path / "nonexistent.yml"
+    result = cli_runner.invoke(cli, ["merge-lora", str(config_path)])
+    assert result.exit_code != 0
+
+
+def test_merge_lora_nonexistent_lora_dir(cli_runner, config_path, tmp_path):
+    """Test merge_lora with nonexistent lora directory"""
+    lora_dir = tmp_path / "nonexistent"
+    result = cli_runner.invoke(
+        cli, ["merge-lora", str(config_path), "--lora-model-dir", str(lora_dir)]
+    )
+    assert result.exit_code != 0
diff --git a/tests/cli/test_cli_merge_sharded_fsdp_weights.py b/tests/cli/test_cli_merge_sharded_fsdp_weights.py
new file mode 100644
index 0000000000..de13b28ed7
--- /dev/null
+++ b/tests/cli/test_cli_merge_sharded_fsdp_weights.py
@@ -0,0 +1,109 @@
+"""pytest tests for axolotl CLI merge_sharded_fsdp_weights command."""
+
+from unittest.mock import patch
+
+from axolotl.cli.main import cli
+
+
+def test_merge_sharded_fsdp_weights_no_accelerate(cli_runner, config_path):
+    """Test merge_sharded_fsdp_weights command without accelerate"""
+    with patch("axolotl.cli.merge_sharded_fsdp_weights.do_cli") as mock:
+        result = cli_runner.invoke(
+            cli,
+            ["merge-sharded-fsdp-weights", str(config_path), "--launcher", "python"],
+        )
+
+        assert mock.called
+        assert mock.call_args.kwargs["config"] == str(config_path)
+        assert result.exit_code == 0
+
+
+def test_merge_sharded_fsdp_weights_with_launcher_args_torchrun(
+    cli_runner, config_path
+):
+    """Test merge-sharded-fsdp-weights with torchrun launcher arguments"""
+    with patch("subprocess.run") as mock_subprocess:
+        result = cli_runner.invoke(
+            cli,
+            [
+                "merge-sharded-fsdp-weights",
+                str(config_path),
+                "--launcher",
+                "torchrun",
+                "--",
+                "--nproc_per_node=2",
+                "--nnodes=1",
+            ],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_subprocess.assert_called_once()
+
+        # Verify launcher args are passed to torchrun
+        called_cmd = mock_subprocess.call_args.args[0]
+        assert called_cmd[0] == "torchrun"
+        assert "--nproc_per_node=2" in called_cmd
+        assert "--nnodes=1" in called_cmd
+        assert "-m" in called_cmd
+        assert "axolotl.cli.merge_sharded_fsdp_weights" in called_cmd
+
+
+def test_merge_sharded_fsdp_weights_with_launcher_args_accelerate(
+    cli_runner, config_path
+):
+    """Test merge-sharded-fsdp-weights with accelerate launcher arguments"""
+    with patch("subprocess.run") as mock_subprocess:
+        result = cli_runner.invoke(
+            cli,
+            [
+                "merge-sharded-fsdp-weights",
+                str(config_path),
+                "--launcher",
+                "accelerate",
+                "--",
+                "--config_file=accelerate_config.yml",
+                "--num_processes=4",
+            ],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_subprocess.assert_called_once()
+
+        # Verify launcher args are passed to accelerate
+        called_cmd = mock_subprocess.call_args.args[0]
+        assert called_cmd[0] == "accelerate"
+        assert called_cmd[1] == "launch"
+        assert "--config_file=accelerate_config.yml" in called_cmd
+        assert "--num_processes=4" in called_cmd
+        assert "-m" in called_cmd
+        assert "axolotl.cli.merge_sharded_fsdp_weights" in called_cmd
+
+
+def test_merge_sharded_fsdp_weights_backward_compatibility_no_launcher_args(
+    cli_runner, config_path
+):
+    """Test that existing merge-sharded-fsdp-weights commands work without launcher args"""
+    with patch("subprocess.run") as mock_subprocess:
+        result = cli_runner.invoke(
+            cli,
+            [
+                "merge-sharded-fsdp-weights",
+                str(config_path),
+                "--launcher",
+                "accelerate",
+            ],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_subprocess.assert_called_once()
+
+        # Verify no launcher args contamination
+        called_cmd = mock_subprocess.call_args.args[0]
+        assert called_cmd[0] == "accelerate"
+        assert called_cmd[1] == "launch"
+        # Should not contain any extra launcher args
+        launcher_section = called_cmd[2 : called_cmd.index("-m")]
+        assert len(launcher_section) == 0  # No launcher args between 'launch' and '-m'
diff --git a/tests/cli/test_cli_preprocess.py b/tests/cli/test_cli_preprocess.py
new file mode 100644
index 0000000000..b213e43da3
--- /dev/null
+++ b/tests/cli/test_cli_preprocess.py
@@ -0,0 +1,78 @@
+"""pytest tests for axolotl CLI preprocess command."""
+
+import shutil
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from axolotl.cli.main import cli
+
+
+@pytest.fixture(autouse=True)
+def cleanup_last_run_prepared():
+    yield
+
+    if Path("last_run_prepared").exists():
+        shutil.rmtree("last_run_prepared")
+
+
+def test_preprocess_config_not_found(cli_runner):
+    """Test preprocess fails when config not found"""
+    result = cli_runner.invoke(cli, ["preprocess", "nonexistent.yml"])
+    assert result.exit_code != 0
+
+
+def test_preprocess_basic(cli_runner, config_path):
+    """Test basic preprocessing with minimal config"""
+    with patch("axolotl.cli.preprocess.do_cli") as mock_do_cli:
+        with patch("axolotl.cli.preprocess.load_datasets") as mock_load_datasets:
+            mock_load_datasets.return_value = MagicMock()
+
+            result = cli_runner.invoke(cli, ["preprocess", str(config_path)])
+            assert result.exit_code == 0
+
+            mock_do_cli.assert_called_once()
+            assert mock_do_cli.call_args.kwargs["config"] == str(config_path)
+            assert mock_do_cli.call_args.kwargs["download"] is True
+
+
+def test_preprocess_without_download(cli_runner, config_path):
+    """Test preprocessing without model download"""
+    with patch("axolotl.cli.preprocess.do_cli") as mock_do_cli:
+        result = cli_runner.invoke(
+            cli, ["preprocess", str(config_path), "--no-download"]
+        )
+        assert result.exit_code == 0
+
+        mock_do_cli.assert_called_once()
+        assert mock_do_cli.call_args.kwargs["config"] == str(config_path)
+        assert mock_do_cli.call_args.kwargs["download"] is False
+
+
+def test_preprocess_custom_path(cli_runner, tmp_path, valid_test_config):
+    """Test preprocessing with custom dataset path"""
+    config_path = tmp_path / "config.yml"
+    custom_path = tmp_path / "custom_prepared"
+    config_path.write_text(valid_test_config)
+
+    with patch("axolotl.cli.preprocess.do_cli") as mock_do_cli:
+        with patch("axolotl.cli.preprocess.load_datasets") as mock_load_datasets:
+            mock_load_datasets.return_value = MagicMock()
+
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "preprocess",
+                    str(config_path),
+                    "--dataset-prepared-path",
+                    str(custom_path.absolute()),
+                ],
+            )
+            assert result.exit_code == 0
+
+            mock_do_cli.assert_called_once()
+            assert mock_do_cli.call_args.kwargs["config"] == str(config_path)
+            assert mock_do_cli.call_args.kwargs["dataset_prepared_path"] == str(
+                custom_path.absolute()
+            )
diff --git a/tests/cli/test_cli_sweeps.py b/tests/cli/test_cli_sweeps.py
new file mode 100644
index 0000000000..1b14f5aca0
--- /dev/null
+++ b/tests/cli/test_cli_sweeps.py
@@ -0,0 +1,69 @@
+"""
+unit tests for generating sweep configurations
+"""
+
+from axolotl.cli.utils import generate_sweep_configs
+
+
+def test_generate_sweep_configs_no_pairs():
+    base_config = {
+        "learning_rate": 0.1,
+        "micro_batch_size": 1,
+        "sample_packing": True,
+    }
+
+    sweeps_config = {"micro_batch_size": [1, 2, 4], "weight_decay": [0.0, 0.1]}
+
+    generate_sweep_configs(base_config, sweeps_config)
+
+    assert len(generate_sweep_configs(base_config, sweeps_config)) == 6
+
+    cfg_1 = {
+        "learning_rate": 0.1,
+        "micro_batch_size": 2,
+        "weight_decay": 0.0,
+        "sample_packing": True,
+    }
+
+    assert any(
+        cfg_1 == cfg for cfg in generate_sweep_configs(base_config, sweeps_config)
+    )
+
+
+def test_generate_sweep_configs_with_pairs():
+    base_config = {
+        "learning_rate": 0.1,
+        "micro_batch_size": 1,
+        "sample_packing": True,
+    }
+
+    sweeps_config = {
+        "_": [
+            {
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 8,
+            },
+            {
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 4,
+            },
+            {
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+            },
+            {
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+            },
+        ],
+        "weight_decay": [0.0, 0.1],
+    }
+
+    generate_sweep_configs(base_config, sweeps_config)
+
+    assert len(generate_sweep_configs(base_config, sweeps_config)) == 8
+
+    assert all(
+        cfg["gradient_accumulation_steps"] * cfg["micro_batch_size"] == 8
+        for cfg in generate_sweep_configs(base_config, sweeps_config)
+    )
diff --git a/tests/cli/test_cli_train.py b/tests/cli/test_cli_train.py
new file mode 100644
index 0000000000..1251ab3c03
--- /dev/null
+++ b/tests/cli/test_cli_train.py
@@ -0,0 +1,251 @@
+"""Tests for train CLI command."""
+
+from unittest.mock import MagicMock, patch
+
+from axolotl.cli.main import cli
+
+from .test_cli_base import BaseCliTest
+
+
+class TestTrainCommand(BaseCliTest):
+    """Test cases for train command."""
+
+    cli = cli
+
+    def test_train_cli_validation(self, cli_runner):
+        """Test CLI validation"""
+        self._test_cli_validation(cli_runner, "train")
+
+    def test_train_basic_execution(self, cli_runner, tmp_path, valid_test_config):
+        """Test basic successful execution"""
+        self._test_basic_execution(
+            cli_runner, tmp_path, valid_test_config, "train", train=True
+        )
+
+    def test_train_basic_execution_no_accelerate(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test basic successful execution without accelerate"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("axolotl.cli.train.train") as mock_train:
+            mock_train.return_value = (MagicMock(), MagicMock(), MagicMock())
+            with patch("axolotl.cli.train.load_datasets") as mock_load_datasets:
+                mock_load_datasets.return_value = MagicMock()
+
+                result = cli_runner.invoke(
+                    cli,
+                    [
+                        "train",
+                        str(config_path),
+                        "--launcher",
+                        "python",
+                    ],
+                    catch_exceptions=False,
+                )
+
+                assert result.exit_code == 0
+                mock_train.assert_called_once()
+
+    def test_train_cli_overrides(self, cli_runner, tmp_path, valid_test_config):
+        """Test CLI arguments properly override config values"""
+        config_path = self._test_cli_overrides(tmp_path, valid_test_config)
+
+        with patch("axolotl.cli.train.train") as mock_train:
+            mock_train.return_value = (MagicMock(), MagicMock(), MagicMock())
+            with patch("axolotl.cli.train.load_datasets") as mock_load_datasets:
+                mock_load_datasets.return_value = MagicMock()
+
+                result = cli_runner.invoke(
+                    cli,
+                    [
+                        "train",
+                        str(config_path),
+                        "--learning-rate=1e-4",
+                        "--micro-batch-size=2",
+                        "--launcher",
+                        "python",
+                    ],
+                    catch_exceptions=False,
+                )
+
+                assert result.exit_code == 0
+                mock_train.assert_called_once()
+                cfg = mock_train.call_args[1]["cfg"]
+                assert cfg["learning_rate"] == 1e-4
+                assert cfg["micro_batch_size"] == 2
+
+    def test_train_with_launcher_args_torchrun(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test train with torchrun launcher arguments"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("os.execvpe") as mock_subprocess:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "train",
+                    str(config_path),
+                    "--launcher",
+                    "torchrun",
+                    "--",
+                    "--nproc_per_node=2",
+                    "--nnodes=1",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_subprocess.assert_called_once()
+
+            # Verify launcher args are passed to torchrun
+            called_cmd = mock_subprocess.call_args.args[1]
+            assert called_cmd[0] == "torchrun"
+            assert "--nproc_per_node=2" in called_cmd
+            assert "--nnodes=1" in called_cmd
+            assert "-m" in called_cmd
+            assert "axolotl.cli.train" in called_cmd
+
+    def test_train_with_launcher_args_accelerate(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test train with accelerate launcher arguments"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("os.execvpe") as mock_subprocess:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "train",
+                    str(config_path),
+                    "--launcher",
+                    "accelerate",
+                    "--",
+                    "--config_file=accelerate_config.yml",
+                    "--num_processes=4",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_subprocess.assert_called_once()
+
+            # Verify launcher args are passed to accelerate
+            assert mock_subprocess.call_args.args[0] == "accelerate"
+            called_cmd = mock_subprocess.call_args.args[1]
+            assert called_cmd[0] == "accelerate"
+            assert called_cmd[1] == "launch"
+            assert "--config_file=accelerate_config.yml" in called_cmd
+            assert "--num_processes=4" in called_cmd
+            assert "-m" in called_cmd
+            assert "axolotl.cli.train" in called_cmd
+
+    def test_train_backward_compatibility_no_launcher_args(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test that existing train commands work without launcher args"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("os.execvpe") as mock_subprocess:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "train",
+                    str(config_path),
+                    "--launcher",
+                    "accelerate",
+                    "--learning-rate",
+                    "1e-4",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_subprocess.assert_called_once()
+
+            # Verify no launcher args contamination
+            assert mock_subprocess.call_args.args[0] == "accelerate"
+            called_cmd = mock_subprocess.call_args.args[1]
+            assert called_cmd[0] == "accelerate"
+            assert called_cmd[1] == "launch"
+            # Should not contain any extra launcher args
+            launcher_section = called_cmd[2 : called_cmd.index("-m")]
+            assert (
+                len(launcher_section) == 0
+            )  # No launcher args between 'launch' and '-m'
+
+    def test_train_mixed_args_with_launcher_args(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test train with both regular CLI args and launcher args"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        with patch("os.execvpe") as mock_subprocess:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "train",
+                    str(config_path),
+                    "--launcher",
+                    "torchrun",
+                    "--learning-rate",
+                    "2e-4",
+                    "--micro-batch-size",
+                    "4",
+                    "--",
+                    "--nproc_per_node=8",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_subprocess.assert_called_once()
+
+            assert mock_subprocess.call_args.args[0] == "torchrun"
+            called_cmd = mock_subprocess.call_args.args[1]
+            # Verify launcher args
+            assert "--nproc_per_node=8" in called_cmd
+            # Verify axolotl args are also present
+            assert "--learning-rate=2e-4" in called_cmd
+            assert "--micro-batch-size=4" in called_cmd
+
+    def test_train_cloud_with_launcher_args(
+        self, cli_runner, tmp_path, valid_test_config
+    ):
+        """Test train with cloud and launcher arguments"""
+        config_path = tmp_path / "config.yml"
+        config_path.write_text(valid_test_config)
+
+        cloud_path = tmp_path / "cloud.yml"
+        cloud_path.write_text("provider: modal\ngpu: a100")
+
+        with patch("axolotl.cli.cloud.do_cli_train") as mock_cloud_train:
+            result = cli_runner.invoke(
+                cli,
+                [
+                    "train",
+                    str(config_path),
+                    "--cloud",
+                    str(cloud_path),
+                    "--launcher",
+                    "torchrun",
+                    "--",
+                    "--nproc_per_node=4",
+                    "--nnodes=2",
+                ],
+                catch_exceptions=False,
+            )
+
+            assert result.exit_code == 0
+            mock_cloud_train.assert_called_once()
+
+            # Verify cloud training was called with launcher args
+            call_kwargs = mock_cloud_train.call_args.kwargs
+            assert call_kwargs["launcher"] == "torchrun"
+            assert call_kwargs["launcher_args"] == ["--nproc_per_node=4", "--nnodes=2"]
diff --git a/tests/cli/test_cli_version.py b/tests/cli/test_cli_version.py
new file mode 100644
index 0000000000..533dd5c0ec
--- /dev/null
+++ b/tests/cli/test_cli_version.py
@@ -0,0 +1,11 @@
+"""pytest tests for axolotl CLI --version"""
+
+from axolotl.cli.main import cli
+
+
+def test_print_version(cli_runner):
+    """Test that version is printed when --version is used."""
+
+    result = cli_runner.invoke(cli, ["--version"])
+    assert result.exit_code == 0
+    assert "axolotl, version " in result.output
diff --git a/tests/cli/test_cli_vllm_serve.py b/tests/cli/test_cli_vllm_serve.py
new file mode 100644
index 0000000000..adbbb17164
--- /dev/null
+++ b/tests/cli/test_cli_vllm_serve.py
@@ -0,0 +1,60 @@
+"""Tests for axolotl vllm-serve CLI/config boolean precedence."""
+
+import sys
+import types
+from unittest.mock import MagicMock
+
+import pytest
+
+from axolotl.cli.main import cli
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture
+def stub_serve(monkeypatch):
+    """Stub the serve module (no real vLLM server) and route load_cfg to a
+    config with the vLLM booleans enabled."""
+    name = "axolotl_stub_serve"
+    module = types.ModuleType(name)
+    module.main = MagicMock()
+    monkeypatch.setitem(sys.modules, name, module)
+
+    cfg = DictDefault(
+        {
+            "base_model": "dummy-model",
+            "vllm": {
+                "serve_module": name,
+                "enable_prefix_caching": True,
+                "enable_reasoning": True,
+            },
+        }
+    )
+    monkeypatch.setattr("axolotl.cli.vllm_serve.load_cfg", lambda *_, **__: cfg)
+    return module
+
+
+@pytest.mark.parametrize(
+    "flags, expected",
+    [
+        (["--no-enable-prefix-caching", "--no-enable-reasoning"], False),
+        ([], True),
+    ],
+)
+def test_vllm_serve_bool_precedence(cli_runner, tmp_path, stub_serve, flags, expected):
+    """`--no-<flag>` overrides a config-enabled option; omitting it keeps the config value.
+
+    Drives the real CLI so Click's ``--flag/--no-flag`` parsing and
+    ``filter_none_kwargs`` run: an omitted flag reaches ``do_vllm_serve`` as
+    ``None`` (config wins), ``--no-<flag>`` as ``False`` (overrides). Regression
+    for the ``cli or cfg`` bug where ``False or True`` -> ``True``.
+    """
+    config = tmp_path / "config.yml"
+    config.write_text("base_model: dummy-model\n")
+
+    result = cli_runner.invoke(cli, ["vllm-serve", str(config), *flags])
+    assert result.exit_code == 0, result.output
+
+    stub_serve.main.assert_called_once()
+    script_args = stub_serve.main.call_args.args[0]
+    assert script_args.enable_prefix_caching is expected
+    assert script_args.enable_reasoning is expected
diff --git a/tests/cli/test_load_cfg_capabilities.py b/tests/cli/test_load_cfg_capabilities.py
new file mode 100644
index 0000000000..ae85f78a6d
--- /dev/null
+++ b/tests/cli/test_load_cfg_capabilities.py
@@ -0,0 +1,205 @@
+"""Tests for GPU capability detection in `load_cfg` and `ray_train_func`."""
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from axolotl.cli.config import load_cfg
+from axolotl.cli.train import ray_train_func
+
+_BASE_CONFIG = """
+base_model: HuggingFaceTB/SmolLM2-135M
+datasets:
+  - path: mhenrichsen/alpaca_2k_test
+    type: alpaca
+sequence_len: 2048
+max_steps: 1
+micro_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 1e-3
+special_tokens:
+  pad_token: <|endoftext|>
+"""
+
+
+def _write_cfg(tmp_path: Path, extra: str = "") -> Path:
+    """Write the base test config (plus any extra YAML lines) to a temp file."""
+    path = tmp_path / "config.yml"
+    path.write_text(_BASE_CONFIG + extra)
+    return path
+
+
+def _patch_load_cfg_dependencies(monkeypatch, validate_mock=None):
+    """Stub everything `load_cfg` does after validation so the test can focus
+    on whether GPU capabilities were probed on the driver.
+
+    If ``validate_mock`` is given, it is installed as ``validate_config`` so the
+    test can inspect the arguments it was called with; otherwise a simple
+    identity stub is used.
+    """
+    monkeypatch.setattr(
+        "axolotl.cli.config.validate_config",
+        validate_mock if validate_mock is not None else (lambda cfg, **_: cfg),
+    )
+    monkeypatch.setattr("axolotl.cli.config.normalize_config", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.config.normalize_cfg_datasets", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.config.prepare_debug_log", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.config.prepare_optim_env", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.config.setup_wandb_env_vars", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.config.setup_mlflow_env_vars", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.config.setup_comet_env_vars", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.config.setup_trackio_env_vars", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.config.plugin_set_cfg", lambda *_: None)
+    monkeypatch.setattr(
+        "axolotl.cli.config.TELEMETRY_MANAGER.send_event", lambda *_, **__: None
+    )
+
+
+def test_load_cfg_probes_capabilities_by_default(tmp_path, monkeypatch):
+    """Without `use_ray`, `load_cfg` probes GPU capabilities on the local host
+    and passes the results into `validate_config`."""
+    validate_mock = MagicMock(side_effect=lambda cfg, **_: cfg)
+    _patch_load_cfg_dependencies(monkeypatch, validate_mock=validate_mock)
+    config_path = _write_cfg(tmp_path)
+
+    with patch("axolotl.cli.config.gpu_capabilities") as mock_caps:
+        mock_caps.return_value = ({"bf16": False}, {"torch_version": "2.6.0"})
+        load_cfg(str(config_path))
+
+    mock_caps.assert_called_once()
+    _, kwargs = validate_mock.call_args
+    assert kwargs["capabilities"] == {"bf16": False}
+    assert kwargs["env_capabilities"] == {"torch_version": "2.6.0"}
+
+
+def test_load_cfg_skips_capabilities_under_ray(tmp_path, monkeypatch):
+    """With `use_ray: true`, capability detection is deferred to the worker
+    and `validate_config` receives `None` for both capability dicts."""
+    validate_mock = MagicMock(side_effect=lambda cfg, **_: cfg)
+    _patch_load_cfg_dependencies(monkeypatch, validate_mock=validate_mock)
+    config_path = _write_cfg(tmp_path, "use_ray: true\nray_num_workers: 1\n")
+
+    with patch("axolotl.cli.config.gpu_capabilities") as mock_caps:
+        load_cfg(str(config_path))
+
+    mock_caps.assert_not_called()
+    _, kwargs = validate_mock.call_args
+    assert kwargs["capabilities"] is None
+    assert kwargs["env_capabilities"] is None
+
+
+def test_ray_train_func_validates_with_worker_capabilities(monkeypatch):
+    """`ray_train_func` must probe `gpu_capabilities()` on the worker and feed
+    the result into `validate_config` before training runs."""
+    cfg_dict = {
+        "base_model": "HuggingFaceTB/SmolLM2-135M",
+        "micro_batch_size": 1,
+        "gradient_accumulation_steps": 1,
+    }
+
+    validate_mock = MagicMock(side_effect=lambda cfg, **_: cfg)
+    do_train_mock = MagicMock()
+    accelerator_mock = MagicMock()
+
+    monkeypatch.setattr("axolotl.cli.train.validate_config", validate_mock)
+    monkeypatch.setattr("axolotl.cli.train.do_train", do_train_mock)
+    monkeypatch.setattr("axolotl.cli.train.prepare_optim_env", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.normalize_config", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.resolve_dtype", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.Accelerator", accelerator_mock)
+
+    with patch("axolotl.cli.train.gpu_capabilities") as mock_caps:
+        mock_caps.return_value = (
+            {"bf16": True, "fp8": False, "tf32": True, "compute_capability": "sm_90"},
+            {"torch_version": "2.6.0"},
+        )
+        ray_train_func({"cfg": cfg_dict, "cli_args": MagicMock()})
+
+    mock_caps.assert_called_once()
+    validate_mock.assert_called_once()
+    _, kwargs = validate_mock.call_args
+    assert kwargs["capabilities"] == {
+        "bf16": True,
+        "fp8": False,
+        "tf32": True,
+        "compute_capability": "sm_90",
+    }
+    assert kwargs["env_capabilities"] == {"torch_version": "2.6.0"}
+    do_train_mock.assert_called_once()
+
+
+def test_ray_train_func_registers_plugins_before_validate_config(monkeypatch):
+    """Regression: plugins must be registered before `validate_config` so the
+    plugin-extended pydantic schema is in scope. Otherwise `merge_input_args`
+    sees an empty PluginManager on the worker and `model_dump(exclude_none=True)`
+    silently drops plugin-specific cfg fields.
+    """
+    cfg_dict = {
+        "base_model": "HuggingFaceTB/SmolLM2-135M",
+        "micro_batch_size": 1,
+        "gradient_accumulation_steps": 1,
+        "plugins": ["axolotl.integrations.liger.LigerPlugin"],
+    }
+
+    parent = MagicMock()
+    parent.validate_config.side_effect = lambda cfg, **_: cfg
+
+    # Patch at the source module so a local `from axolotl.cli.config import ...`
+    # inside the function also resolves to the mock; also patch the train module
+    # for top-level imports (raising=False keeps it tolerant of either style).
+    monkeypatch.setattr("axolotl.cli.config.prepare_plugins", parent.prepare_plugins)
+    monkeypatch.setattr("axolotl.cli.config.plugin_set_cfg", parent.plugin_set_cfg)
+    monkeypatch.setattr(
+        "axolotl.cli.train.prepare_plugins", parent.prepare_plugins, raising=False
+    )
+    monkeypatch.setattr(
+        "axolotl.cli.train.plugin_set_cfg", parent.plugin_set_cfg, raising=False
+    )
+    monkeypatch.setattr("axolotl.cli.train.validate_config", parent.validate_config)
+    monkeypatch.setattr("axolotl.cli.train.gpu_capabilities", lambda: ({}, {}))
+    monkeypatch.setattr("axolotl.cli.train.do_train", MagicMock())
+    monkeypatch.setattr("axolotl.cli.train.prepare_optim_env", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.normalize_config", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.resolve_dtype", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.Accelerator", MagicMock())
+
+    ray_train_func({"cfg": cfg_dict, "cli_args": MagicMock()})
+
+    # Filter to the calls we care about, in the order they happened.
+    ordered = [
+        call[0]
+        for call in parent.mock_calls
+        if call[0] in ("prepare_plugins", "validate_config", "plugin_set_cfg")
+    ]
+    assert ordered == ["prepare_plugins", "validate_config", "plugin_set_cfg"], (
+        f"Expected prepare_plugins -> validate_config -> plugin_set_cfg; got {ordered}"
+    )
+
+
+def test_ray_train_func_skips_plugin_registration_when_no_plugins(monkeypatch):
+    """When no plugins are configured, neither `prepare_plugins` nor
+    `plugin_set_cfg` should be invoked on the worker."""
+    cfg_dict = {
+        "base_model": "HuggingFaceTB/SmolLM2-135M",
+        "micro_batch_size": 1,
+        "gradient_accumulation_steps": 1,
+    }
+
+    prepare_plugins_mock = MagicMock()
+    plugin_set_cfg_mock = MagicMock()
+
+    monkeypatch.setattr("axolotl.cli.train.prepare_plugins", prepare_plugins_mock)
+    monkeypatch.setattr("axolotl.cli.train.plugin_set_cfg", plugin_set_cfg_mock)
+    monkeypatch.setattr(
+        "axolotl.cli.train.validate_config", MagicMock(side_effect=lambda cfg, **_: cfg)
+    )
+    monkeypatch.setattr("axolotl.cli.train.gpu_capabilities", lambda: ({}, {}))
+    monkeypatch.setattr("axolotl.cli.train.do_train", MagicMock())
+    monkeypatch.setattr("axolotl.cli.train.prepare_optim_env", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.normalize_config", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.resolve_dtype", lambda *_: None)
+    monkeypatch.setattr("axolotl.cli.train.Accelerator", MagicMock())
+
+    ray_train_func({"cfg": cfg_dict, "cli_args": MagicMock()})
+
+    prepare_plugins_mock.assert_not_called()
+    plugin_set_cfg_mock.assert_not_called()
diff --git a/tests/cli/test_nested_options.py b/tests/cli/test_nested_options.py
new file mode 100644
index 0000000000..221de951e7
--- /dev/null
+++ b/tests/cli/test_nested_options.py
@@ -0,0 +1,227 @@
+"""Tests for nested config option handling via CLI dot-notation."""
+
+import click
+from click.testing import CliRunner
+from pydantic import BaseModel, Field
+
+from axolotl.cli.utils.args import add_options_from_config, filter_none_kwargs
+
+
+class InnerConfig(BaseModel):
+    """A nested config model for testing."""
+
+    beta: float | None = Field(
+        default=None,
+        description="Beta parameter.",
+    )
+    host: str | None = Field(
+        default=None,
+        description="Server host.",
+    )
+    use_feature: bool = Field(
+        default=False,
+        description="Whether to use the feature.",
+    )
+
+
+class OuterConfig(BaseModel):
+    """A top-level config model for testing."""
+
+    learning_rate: float | None = Field(
+        default=None,
+        description="Learning rate.",
+    )
+    inner: InnerConfig | None = Field(
+        default=None,
+        description="Inner config.",
+    )
+    name: str | None = Field(
+        default=None,
+        description="Model name.",
+    )
+
+
+class TestAddOptionsFromConfigNested:
+    """Test that add_options_from_config handles nested BaseModel fields."""
+
+    def setup_method(self):
+        self.runner = CliRunner()
+
+    def test_nested_dot_notation_options_are_registered(self):
+        """Nested model fields should create --parent.child CLI options."""
+
+        @click.command()
+        @add_options_from_config(OuterConfig)
+        @filter_none_kwargs
+        def cmd(**kwargs):
+            for k, v in sorted(kwargs.items()):
+                click.echo(f"{k}={v}")
+
+        result = self.runner.invoke(cmd, ["--inner.beta=0.5", "--inner.host=localhost"])
+        assert result.exit_code == 0, result.output
+        assert "inner__beta=0.5" in result.output
+        assert "inner__host=localhost" in result.output
+
+    def test_nested_bool_option(self):
+        """Nested bool fields should support --parent.field/--no-parent.field."""
+
+        @click.command()
+        @add_options_from_config(OuterConfig)
+        @filter_none_kwargs
+        def cmd(**kwargs):
+            for k, v in sorted(kwargs.items()):
+                click.echo(f"{k}={v}")
+
+        result = self.runner.invoke(cmd, ["--inner.use-feature"])
+        assert result.exit_code == 0, result.output
+        assert "inner__use_feature=True" in result.output
+
+    def test_flat_and_nested_options_together(self):
+        """Flat and nested options should work together."""
+
+        @click.command()
+        @add_options_from_config(OuterConfig)
+        @filter_none_kwargs
+        def cmd(**kwargs):
+            for k, v in sorted(kwargs.items()):
+                click.echo(f"{k}={v}")
+
+        result = self.runner.invoke(
+            cmd, ["--learning-rate=0.001", "--inner.beta=0.1", "--name=test"]
+        )
+        assert result.exit_code == 0, result.output
+        assert "learning_rate=0.001" in result.output
+        assert "inner__beta=0.1" in result.output
+        assert "name=test" in result.output
+
+    def test_no_nested_options_passed(self):
+        """When no nested options are passed, they should not appear in kwargs."""
+
+        @click.command()
+        @add_options_from_config(OuterConfig)
+        @filter_none_kwargs
+        def cmd(**kwargs):
+            click.echo(f"keys={sorted(kwargs.keys())}")
+
+        result = self.runner.invoke(cmd, ["--learning-rate=0.01"])
+        assert result.exit_code == 0, result.output
+        assert "inner__" not in result.output
+
+
+class TestLoadCfgNestedKwargs:
+    """Test that load_cfg correctly applies nested (double-underscore) kwargs."""
+
+    @staticmethod
+    def _apply_nested_kwargs(cfg, kwargs):
+        """Helper that mirrors the nested kwargs handling from load_cfg,
+        including type coercion for string CLI values."""
+        from axolotl.cli.config import _coerce_value
+
+        nested_kwargs: dict = {}
+        flat_kwargs: dict = {}
+        for key, value in kwargs.items():
+            if "__" in key:
+                parent, child = key.split("__", 1)
+                nested_kwargs.setdefault(parent, {})[child] = value
+            else:
+                flat_kwargs[key] = value
+
+        cfg_keys = cfg.keys()
+        for key, value in flat_kwargs.items():
+            if key in cfg_keys:
+                cfg[key] = _coerce_value(value, cfg.get(key))
+
+        for parent, children in nested_kwargs.items():
+            if cfg[parent] is None:
+                cfg[parent] = {}
+            if not isinstance(cfg[parent], dict):
+                cfg[parent] = {}
+            for child_key, child_value in children.items():
+                existing = cfg[parent].get(child_key)
+                cfg[parent][child_key] = _coerce_value(child_value, existing)
+
+        return cfg
+
+    def test_nested_kwargs_applied_to_cfg(self, tmp_path):
+        """Double-underscore kwargs should set nested config values."""
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault({"trl": {"beta": 0.1}, "learning_rate": 0.01})
+        # CLI passes strings, so simulate that
+        kwargs = {
+            "trl__beta": "0.5",
+            "trl__host": "192.168.1.1",
+            "learning_rate": "0.02",
+        }
+
+        cfg = self._apply_nested_kwargs(cfg, kwargs)
+
+        assert cfg["learning_rate"] == 0.02
+        assert isinstance(cfg["learning_rate"], float)
+        assert cfg["trl"]["beta"] == 0.5
+        assert isinstance(cfg["trl"]["beta"], float)
+        assert cfg["trl"]["host"] == "192.168.1.1"
+
+    def test_nested_kwargs_creates_parent_if_none(self):
+        """If the parent key is None, nested kwargs should create the dict."""
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault({"trl": None, "learning_rate": 0.01})
+        cfg = self._apply_nested_kwargs(cfg, {"trl__beta": "0.5"})
+
+        # No existing value, YAML-style inference: "0.5" -> 0.5
+        assert cfg["trl"]["beta"] == 0.5
+        assert isinstance(cfg["trl"]["beta"], float)
+
+    def test_nested_kwargs_overwrites_string_parent(self):
+        """If the parent key is a string, it should be replaced with a dict."""
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault({"trl": "some_string", "learning_rate": 0.01})
+        cfg = self._apply_nested_kwargs(cfg, {"trl__beta": "0.5"})
+
+        assert cfg["trl"]["beta"] == 0.5
+
+
+class TestCoerceValue:
+    """Test YAML-style type coercion for CLI string values."""
+
+    def test_coerce_with_existing_float(self):
+        from axolotl.cli.config import _coerce_value
+
+        assert _coerce_value("0.5", 0.1) == 0.5
+        assert isinstance(_coerce_value("0.5", 0.1), float)
+
+    def test_coerce_with_existing_int(self):
+        from axolotl.cli.config import _coerce_value
+
+        assert _coerce_value("42", 10) == 42
+        assert isinstance(_coerce_value("42", 10), int)
+
+    def test_coerce_with_existing_bool(self):
+        from axolotl.cli.config import _coerce_value
+
+        assert _coerce_value("true", False) is True
+        assert _coerce_value("false", True) is False
+        assert _coerce_value("1", False) is True
+        assert _coerce_value("0", True) is False
+
+    def test_coerce_yaml_inference_no_existing(self):
+        """Without an existing value, use YAML-style inference."""
+        from axolotl.cli.config import _coerce_value
+
+        assert _coerce_value("true", None) is True
+        assert _coerce_value("false", None) is False
+        assert _coerce_value("42", None) == 42
+        assert isinstance(_coerce_value("42", None), int)
+        assert _coerce_value("3.14", None) == 3.14
+        assert isinstance(_coerce_value("3.14", None), float)
+        assert _coerce_value("null", None) is None
+        assert _coerce_value("hello", None) == "hello"
+
+    def test_coerce_non_string_passthrough(self):
+        """Non-string values should pass through unchanged."""
+        from axolotl.cli.config import _coerce_value
+
+        assert _coerce_value(0.5, 0.1) == 0.5
+        assert _coerce_value(True, False) is True
diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py
new file mode 100644
index 0000000000..431c35c3ce
--- /dev/null
+++ b/tests/cli/test_utils.py
@@ -0,0 +1,229 @@
+"""pytest tests for axolotl CLI utils."""
+
+import json
+from unittest.mock import Mock, patch
+
+import click
+import pytest
+import requests
+
+from axolotl.cli.utils import fetch_from_github
+
+# Sample GitHub API response
+MOCK_TREE_RESPONSE = {
+    "tree": [
+        {"path": "examples/config1.yml", "type": "blob", "sha": "abc123"},
+        {"path": "examples/config2.yml", "type": "blob", "sha": "def456"},
+        {"path": "other/file.txt", "type": "blob", "sha": "xyz789"},
+    ]
+}
+
+
+@pytest.fixture
+def mock_responses():
+    """Mock responses for API and file downloads"""
+
+    def mock_get(url, timeout=None):
+        response = Mock()
+        if "api.github.com" in url:
+            response.text = json.dumps(MOCK_TREE_RESPONSE)
+        else:
+            response.content = b"file content"
+        return response
+
+    return mock_get
+
+
+def test_fetch_from_github_new_files(tmp_path, mock_responses):
+    """Test fetching new files"""
+    with patch("requests.get", mock_responses):
+        fetch_from_github("examples/", tmp_path)
+
+        # Verify files were created
+        assert (tmp_path / "config1.yml").exists()
+        assert (tmp_path / "config2.yml").exists()
+        assert not (tmp_path / "file.txt").exists()
+
+
+def test_fetch_from_github_unchanged_files(tmp_path, mock_responses):
+    """Test handling of unchanged files"""
+    # Create existing file with matching SHA
+    existing_file = tmp_path / "config1.yml"
+    existing_file.write_bytes(b"file content")
+
+    with patch("requests.get", mock_responses):
+        fetch_from_github("examples/", tmp_path)
+
+        # File should not be downloaded again
+        assert existing_file.read_bytes() == b"file content"
+
+
+def test_fetch_from_github_invalid_prefix(mock_responses):
+    """Test error handling for invalid directory prefix"""
+    with patch("requests.get", mock_responses):
+        with pytest.raises(click.ClickException):
+            fetch_from_github("nonexistent/", None)
+
+
+def test_fetch_from_github_network_error():
+    """Test handling of network errors"""
+    with patch("requests.get", side_effect=requests.RequestException):
+        with pytest.raises(requests.RequestException):
+            fetch_from_github("examples/", None)
+
+
+def assert_launcher_args_in_command(
+    mock_subprocess_call,
+    launcher: str,
+    expected_launcher_args: list[str],
+    command_module: str,
+):
+    """
+    Helper function to verify launcher arguments are properly passed in subprocess calls.
+
+    Args:
+        mock_subprocess_call: The mock subprocess.run call
+        launcher: Expected launcher ("accelerate", "torchrun", etc.)
+        expected_launcher_args: List of expected launcher arguments
+        command_module: Expected module name (e.g., "axolotl.cli.train")
+    """
+    assert mock_subprocess_call.called, "subprocess.run should have been called"
+    called_cmd = mock_subprocess_call.call_args.args[0]
+
+    # Verify launcher
+    assert called_cmd[0] == launcher, (
+        f"Expected launcher {launcher}, got {called_cmd[0]}"
+    )
+
+    # Verify launcher args are present
+    for arg in expected_launcher_args:
+        assert arg in called_cmd, (
+            f"Expected launcher arg '{arg}' not found in command: {called_cmd}"
+        )
+
+    # Verify module is present
+    assert "-m" in called_cmd, "Expected -m flag for module execution"
+    assert command_module in called_cmd, (
+        f"Expected module {command_module} not found in command: {called_cmd}"
+    )
+
+
+def assert_no_launcher_args_contamination(mock_subprocess_call, launcher: str):
+    """
+    Helper function to verify no unwanted launcher arguments are present.
+
+    Args:
+        mock_subprocess_call: The mock subprocess.run call
+        launcher: Expected launcher ("accelerate", "torchrun", etc.)
+    """
+    assert mock_subprocess_call.called, "subprocess.run should have been called"
+    called_cmd = mock_subprocess_call.call_args.args[0]
+
+    if launcher == "accelerate":
+        # For accelerate, launcher args should be between 'launch' and '-m'
+        launch_idx = called_cmd.index("launch")
+        m_idx = called_cmd.index("-m")
+        launcher_section = called_cmd[launch_idx + 1 : m_idx]
+        assert len(launcher_section) == 0, (
+            f"Unexpected launcher args found: {launcher_section}"
+        )
+    elif launcher == "torchrun":
+        # For torchrun, launcher args should be between 'torchrun' and '-m'
+        torchrun_idx = called_cmd.index("torchrun")
+        m_idx = called_cmd.index("-m")
+        launcher_section = called_cmd[torchrun_idx + 1 : m_idx]
+        assert len(launcher_section) == 0, (
+            f"Unexpected launcher args found: {launcher_section}"
+        )
+
+
+@pytest.fixture
+def common_launcher_args():
+    """Fixture providing common launcher argument combinations for testing."""
+    return {
+        "torchrun": ["--nproc_per_node=2", "--nnodes=1"],
+        "accelerate": ["--config_file=accelerate_config.yml", "--num_processes=4"],
+    }
+
+
+def test_add_default_rdzv_args_with_endpoint():
+    """Test that default RDZV args are added when rdzv_endpoint is present."""
+    from axolotl.cli.utils.train import _add_default_rdzv_args
+
+    launcher_args = ["--nnodes=2", "--rdzv_endpoint=127.0.0.1:29400"]
+    result = _add_default_rdzv_args(launcher_args)
+
+    # Should have added rdzv_backend
+    assert "--rdzv_backend" in result
+    assert "c10d" in result
+
+    # Original args should still be present
+    assert "--nnodes=2" in result
+    assert "--rdzv_endpoint=127.0.0.1:29400" in result
+
+
+def test_add_default_rdzv_args_with_existing_backend():
+    """Test that existing rdzv_backend is not overridden."""
+    from axolotl.cli.utils.train import _add_default_rdzv_args
+
+    launcher_args = [
+        "--nnodes=2",
+        "--rdzv_endpoint=127.0.0.1:29400",
+        "--rdzv_backend=static",
+    ]
+    result = _add_default_rdzv_args(launcher_args)
+
+    # Should not add another rdzv_backend
+    backend_count = sum(1 for arg in result if "--rdzv_backend" in arg)
+    assert backend_count == 1
+    assert "--rdzv_backend=static" in result
+
+
+def test_add_default_rdzv_args_with_existing_id():
+    """Test that existing rdzv_id is not overridden."""
+    from axolotl.cli.utils.train import _add_default_rdzv_args
+
+    launcher_args = [
+        "--nnodes=2",
+        "--rdzv_endpoint=127.0.0.1:29400",
+        "--rdzv_id=my_job_123",
+    ]
+    result = _add_default_rdzv_args(launcher_args)
+
+    # Should not add another rdzv_id
+    id_count = sum(1 for arg in result if "--rdzv_id" in arg)
+    assert id_count == 1
+    assert "--rdzv_id=my_job_123" in result
+
+    # Should still add rdzv_backend
+    assert "--rdzv_backend" in result
+    assert "c10d" in result
+
+
+def test_add_default_rdzv_args_without_endpoint():
+    """Test that no RDZV args are added when rdzv_endpoint is not present."""
+    from axolotl.cli.utils.train import _add_default_rdzv_args
+
+    launcher_args = ["--nnodes=2", "--nproc_per_node=4"]
+    result = _add_default_rdzv_args(launcher_args)
+
+    # Should not add any rdzv args
+    assert "--rdzv_backend" not in result
+    assert result == launcher_args
+
+
+def test_add_default_rdzv_args_with_all_existing():
+    """Test that no defaults are added when all RDZV args are present."""
+    from axolotl.cli.utils.train import _add_default_rdzv_args
+
+    launcher_args = [
+        "--nnodes=2",
+        "--rdzv_endpoint=127.0.0.1:29400",
+        "--rdzv_backend=static",
+        "--rdzv_id=existing_job",
+    ]
+    result = _add_default_rdzv_args(launcher_args)
+
+    # Should not add any additional args
+    assert len(result) == len(launcher_args)
+    assert result == launcher_args
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000000..8e7ac14d8b
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,783 @@
+"""Shared pytest fixtures"""
+
+import collections
+import functools
+import importlib
+import logging
+import os
+import shutil
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Generator
+
+import pytest
+import requests
+
+from axolotl.utils.dict import DictDefault
+
+from tests.hf_offline_utils import (
+    enable_hf_offline,
+    hf_offline_context,
+)
+
+logging.getLogger("filelock").setLevel(logging.CRITICAL)
+
+
+def _apply_transformers_test_shims():
+    import transformers.utils as _transformers_utils
+    import transformers.utils.import_utils as _import_utils
+
+    # Shim for deepseek v3
+    if not hasattr(_import_utils, "is_torch_fx_available"):
+
+        def _is_torch_fx_available():
+            try:
+                import torch.fx  # noqa: F401  # pylint: disable=unused-import
+
+                return True
+            except ImportError:
+                return False
+
+        _import_utils.is_torch_fx_available = _is_torch_fx_available
+
+    if not hasattr(_transformers_utils, "is_flash_attn_greater_or_equal_2_10"):
+        from transformers.utils import (
+            is_flash_attn_greater_or_equal as _is_flash_attn_gte,
+        )
+
+        _transformers_utils.is_flash_attn_greater_or_equal_2_10 = lambda: (
+            _is_flash_attn_gte("2.10")
+        )
+
+
+def pytest_configure(config):  # pylint: disable=unused-argument
+    _apply_transformers_test_shims()
+
+
+# A device-side assert / illegal access poisons the process-wide CUDA context, so
+# every later GPU test errors at setup. Abort the session instead of cascading.
+_CUDA_FATAL_MARKERS = (
+    "device-side assert triggered",
+    "an illegal memory access was encountered",
+    "misaligned address",
+)
+_cuda_context_poisoned = False
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_makereport(item, call):  # pylint: disable=unused-argument
+    outcome = yield
+    report = outcome.get_result()
+    global _cuda_context_poisoned  # pylint: disable=global-statement
+    if report.failed and call.excinfo is not None:
+        if any(marker in str(call.excinfo.value) for marker in _CUDA_FATAL_MARKERS):
+            _cuda_context_poisoned = True
+
+
+def pytest_runtest_setup(item):
+    if _cuda_context_poisoned:
+        item.session.shouldstop = (
+            "CUDA context corrupted by an earlier test; aborting to avoid "
+            "cascading setup errors. Re-run the job."
+        )
+        pytest.skip("CUDA context corrupted by an earlier test; aborting suite.")
+
+
+def retry_on_request_exceptions(max_retries=3, delay=1):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            for attempt in range(max_retries):
+                try:
+                    return func(*args, **kwargs)
+                except (
+                    requests.exceptions.ReadTimeout,
+                    requests.exceptions.ConnectionError,
+                    requests.exceptions.HTTPError,
+                ) as exc:
+                    if attempt < max_retries - 1:
+                        wait = 2**attempt * delay  # in seconds
+                        time.sleep(wait)
+                    else:
+                        raise exc
+
+        return wrapper
+
+    return decorator
+
+
+@retry_on_request_exceptions(max_retries=3, delay=5)
+def snapshot_download_w_retry(*args, **kwargs):
+    """
+    download a model or dataset from HF Hub, retrying in requests failures. We also try to fetch it from the local
+    cache first using hf_hub_offline to avoid hitting HF Hub API rate limits. If it doesn't exist in the cache,
+    disable hf_hub_offline and actually fetch from the hub
+    """
+    from huggingface_hub import snapshot_download
+    from huggingface_hub.errors import LocalEntryNotFoundError
+
+    with hf_offline_context(True):
+        try:
+            return snapshot_download(*args, local_files_only=True, **kwargs)
+        except LocalEntryNotFoundError:
+            pass
+    with hf_offline_context(False):
+        return snapshot_download(*args, **kwargs)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_ds_fixture_bundle():
+    ds_dir = snapshot_download_w_retry(
+        "axolotl-ai-internal/axolotl-oss-dataset-fixtures", repo_type="dataset"
+    )
+    return Path(ds_dir)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_smollm2_135m_model():
+    # download the model
+    snapshot_download_w_retry("HuggingFaceTB/SmolLM2-135M", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_smollm2_135m_instruct_model():
+    # download the model
+    snapshot_download_w_retry("HuggingFaceTB/SmolLM2-135M-Instruct", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_smollm2_135m_gptq_model():
+    # download the model
+    snapshot_download_w_retry("lilmeaty/SmolLM2-135M-Instruct-GPTQ", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_qwen3_half_billion_model():
+    # download the model (still used as the KD teacher in tests/e2e/integrations/test_kd.py)
+    snapshot_download_w_retry("Qwen/Qwen3-0.6B", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_llama_model():
+    snapshot_download_w_retry("axolotl-ai-co/tiny-llama-50m", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_mistral_model():
+    snapshot_download_w_retry("axolotl-ai-co/tiny-mistral-25m", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_mixtral_model():
+    snapshot_download_w_retry("axolotl-ai-co/tiny-mixtral-30m", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_phi_model():
+    snapshot_download_w_retry("axolotl-ai-co/tiny-phi-64m", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_falcon_model():
+    snapshot_download_w_retry("axolotl-ai-co/tiny-falcon-42m", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_qwen2_model():
+    snapshot_download_w_retry("axolotl-ai-co/tiny-qwen2-129m", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_qwen3_model():
+    snapshot_download_w_retry("axolotl-ai-co/tiny-qwen3-129m", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_gemma2_model():
+    snapshot_download_w_retry("axolotl-ai-co/tiny-gemma2-137m", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tatsu_lab_alpaca_dataset():
+    # download the dataset
+    snapshot_download_w_retry("tatsu-lab/alpaca", repo_type="dataset")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_mhenrichsen_alpaca_2k_dataset():
+    # download the dataset
+    snapshot_download_w_retry("mhenrichsen/alpaca_2k_test", repo_type="dataset")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_mhenrichsen_alpaca_2k_w_revision_dataset():
+    # download the dataset
+    snapshot_download_w_retry(
+        "mhenrichsen/alpaca_2k_test", repo_type="dataset", revision="d05c1cb"
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_mlabonne_finetome_100k_dataset():
+    # download the dataset
+    snapshot_download_w_retry("mlabonne/FineTome-100k", repo_type="dataset")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_argilla_distilabel_capybara_dpo_7k_binarized_dataset():
+    # download the dataset
+    snapshot_download_w_retry(
+        "argilla/distilabel-capybara-dpo-7k-binarized", repo_type="dataset"
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_argilla_distilabel_intel_orca_dpo_dataset():
+    # download the dataset
+    snapshot_download_w_retry(
+        "argilla/distilabel-intel-orca-dpo-pairs", repo_type="dataset"
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_argilla_ultrafeedback_binarized_preferences_cleaned_dataset():
+    # download the dataset
+    snapshot_download_w_retry(
+        "argilla/ultrafeedback-binarized-preferences-cleaned", repo_type="dataset"
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_argilla_ultrafeedback_binarized_preferences_cleaned_kto_dataset():
+    # download the dataset
+    snapshot_download_w_retry(
+        "argilla/ultrafeedback-binarized-preferences-cleaned-kto", repo_type="dataset"
+    )
+
+
+# @pytest.fixture(scope="session", autouse=True)
+# def download_fozzie_alpaca_dpo_dataset():
+#     # download the dataset
+#     snapshot_download_w_retry(
+#         "fozziethebeat/alpaca_messages_2k_dpo_test", repo_type="dataset"
+#     )
+#     snapshot_download_w_retry(
+#         "fozziethebeat/alpaca_messages_2k_dpo_test",
+#         repo_type="dataset",
+#         revision="ea82cff",
+#     )
+
+
+# @pytest.fixture(scope="session")
+# @disable_hf_offline
+# def dataset_fozzie_alpaca_dpo_dataset(
+#     download_fozzie_alpaca_dpo_dataset,
+# ):
+#     return load_dataset("fozziethebeat/alpaca_messages_2k_dpo_test", split="train")
+#
+#
+# @pytest.fixture(scope="session")
+# @disable_hf_offline
+# def dataset_fozzie_alpaca_dpo_dataset_rev_ea82cff(
+#     download_fozzie_alpaca_dpo_dataset,
+# ):
+#     return load_dataset(
+#         "fozziethebeat/alpaca_messages_2k_dpo_test", split="train", revision="ea82cff"
+#     )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_arcee_ai_distilabel_intel_orca_dpo_pairs_dataset():
+    # download the dataset
+    snapshot_download_w_retry(
+        "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized", repo_type="dataset"
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_argilla_dpo_pairs_dataset():
+    # download the dataset
+    snapshot_download_w_retry(
+        "argilla/distilabel-intel-orca-dpo-pairs", repo_type="dataset"
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_shakespeare_dataset():
+    # download the dataset
+    snapshot_download_w_retry("winglian/tiny-shakespeare", repo_type="dataset")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_evolkit_kd_sample_dataset():
+    # download the dataset
+    snapshot_download_w_retry(
+        "axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample", repo_type="dataset"
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_deepseek_model_fixture():
+    snapshot_download_w_retry("axolotl-ai-co/DeepSeek-V3-11M", repo_type="model")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_huggyllama_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "huggyllama/llama-7b",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_llama33_70b_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "axolotl-ai-co/Llama-3.3-70B-Instruct-tokenizer",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_llama_1b_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "NousResearch/Llama-3.2-1B",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_llama3_8b_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "NousResearch/Meta-Llama-3-8B",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_llama3_8b_instruct_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "NousResearch/Meta-Llama-3-8B-Instruct",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_phi_35_mini_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "microsoft/Phi-3.5-mini-instruct",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_phi_4_reasoning_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "microsoft/Phi-4-reasoning",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_phi_3_mini_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "microsoft/Phi-3-mini-4k-instruct",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_mistral_7b_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "casperhansen/mistral-7b-instruct-v0.1-awq",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_gemma3_4b_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "mlx-community/gemma-3-4b-it-8bit",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_gemma_2b_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "unsloth/gemma-2b-it",
+        revision="703fb4a",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_gemma2_9b_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "mlx-community/gemma-2-9b-it-4bit",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_mlx_mistral_7b_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "mlx-community/Mistral-7B-Instruct-v0.3-4bit",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture
+def download_llama2_model_fixture():
+    # download the tokenizer only
+    snapshot_download_w_retry(
+        "NousResearch/Llama-2-7b-hf",
+        repo_type="model",
+        allow_patterns=["*token*", "config.json"],
+    )
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_llama32_1b_model_fixture():
+    snapshot_download_w_retry(
+        "osllmai-community/Llama-3.2-1B",
+        repo_type="model",
+    )
+
+
+@pytest.fixture
+@enable_hf_offline
+def tokenizer_huggyllama(
+    download_huggyllama_model_fixture,
+):
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
+    tokenizer.pad_token = "</s>"
+
+    return tokenizer
+
+
+@pytest.fixture
+@enable_hf_offline
+def tokenizer_huggyllama_w_special_tokens(
+    tokenizer_huggyllama,
+):
+    tokenizer_huggyllama.add_special_tokens(
+        {
+            "bos_token": "<s>",
+            "eos_token": "</s>",
+            "unk_token": "<unk>",
+        }
+    )
+
+    return tokenizer_huggyllama
+
+
+@pytest.fixture
+@enable_hf_offline
+def tokenizer_llama2_7b(
+    download_llama2_model_fixture,
+):
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-hf")
+
+    return tokenizer
+
+
+@pytest.fixture
+@enable_hf_offline
+def tokenizer_mistral_7b_instruct(
+    download_mlx_mistral_7b_model_fixture,
+):
+    from transformers import AutoTokenizer
+
+    return AutoTokenizer.from_pretrained("casperhansen/mistral-7b-instruct-v0.1-awq")
+
+
+@pytest.fixture
+def tokenizer_mistral_7b_instruct_chatml(tokenizer_mistral_7b_instruct):
+    from tokenizers import AddedToken
+
+    tokenizer_mistral_7b_instruct.add_special_tokens(
+        {
+            "eos_token": AddedToken(
+                "<|im_end|>", rstrip=False, lstrip=False, normalized=False
+            )
+        }
+    )
+    tokenizer_mistral_7b_instruct.add_tokens(
+        [
+            AddedToken("<|im_start|>", rstrip=False, lstrip=False, normalized=False),
+        ]
+    )
+    return tokenizer_mistral_7b_instruct
+
+
+@pytest.fixture
+def temp_dir() -> Generator[str, None, None]:
+    # Create a temporary directory
+    _temp_dir = tempfile.mkdtemp()
+    yield _temp_dir
+    # Clean up the directory after the test
+    shutil.rmtree(_temp_dir)
+
+
+def _clear_plugin_manager():
+    from axolotl.integrations.base import PluginManager
+
+    PluginManager._cfg = None
+    # Don't reset _instance to None — module-level PLUGIN_MANAGER references
+    # in train.py, model.py, etc. would become stale
+    if PluginManager._instance is not None:
+        PluginManager._instance.plugins = collections.OrderedDict()
+
+
+@pytest.fixture(scope="function", autouse=True)
+def reset_plugin_manager():
+    _clear_plugin_manager()
+    yield
+    _clear_plugin_manager()
+
+
+@pytest.fixture(scope="function", autouse=True)
+def torch_manual_seed():
+    import torch
+
+    torch.manual_seed(42)
+
+
+_TRANSFORMERS_MODULES_TO_RESET = (
+    "transformers.models.llama",
+    "transformers.models.llama.modeling_llama",
+    "transformers.trainer",
+    "transformers",
+    "transformers.loss.loss_utils",
+)
+
+_TRANSFORMERS_PATCH_TARGETS = (
+    ("transformers.models.llama.modeling_llama", ("LlamaAttention", "forward")),
+    ("transformers.models.llama.modeling_llama", ("LlamaForCausalLM", "forward")),
+    ("transformers.trainer", ("Trainer", "_inner_training_loop")),
+    ("transformers.trainer", ("Trainer", "training_step")),
+    ("transformers", ("Trainer",)),
+    # explicit targets so patches are reverted even when loss_utils was already imported
+    ("transformers.loss.loss_utils", ("fixed_cross_entropy",)),
+    ("transformers.loss.loss_utils", ("ForCausalLMLoss",)),
+)
+
+
+def _get_nested_attr(obj, attr_path):
+    for attr in attr_path:
+        obj = getattr(obj, attr, None)
+        if obj is None:
+            return None
+    return obj
+
+
+def _set_nested_attr(obj, attr_path, value):
+    for attr in attr_path[:-1]:
+        obj = getattr(obj, attr)
+    setattr(obj, attr_path[-1], value)
+
+
+@pytest.fixture(scope="function", autouse=True)
+def cleanup_monkeypatches():
+    seen_modules = {
+        module_name
+        for module_name in _TRANSFORMERS_MODULES_TO_RESET
+        if module_name in sys.modules
+    }
+    snapshots = {}
+    for module_name, attr_path in _TRANSFORMERS_PATCH_TARGETS:
+        module = sys.modules.get(module_name)
+        if module is None:
+            continue
+        value = _get_nested_attr(module, attr_path)
+        if value is not None:
+            snapshots[(module_name, attr_path)] = value
+
+    yield
+
+    modules_to_reload = set()
+    for (module_name, attr_path), original_value in snapshots.items():
+        module = sys.modules.get(module_name)
+        if module is None:
+            continue
+        current_value = _get_nested_attr(module, attr_path)
+        if current_value is not original_value:
+            _set_nested_attr(module, attr_path, original_value)
+            modules_to_reload.add(module_name)
+
+    modules_to_reload.update(
+        module_name
+        for module_name in _TRANSFORMERS_MODULES_TO_RESET
+        if module_name not in seen_modules and module_name in sys.modules
+    )
+
+    for module_name in _TRANSFORMERS_MODULES_TO_RESET:
+        if module_name not in modules_to_reload:
+            continue
+        module = sys.modules.get(module_name)
+        if module is None or not getattr(module, "__file__", None):
+            continue
+        sys.modules[module_name] = importlib.reload(module)
+
+
+@pytest.fixture
+def dataset_winglian_tiny_shakespeare(
+    download_ds_fixture_bundle: Path,
+):
+    from datasets import load_from_disk
+
+    ds_path = download_ds_fixture_bundle / "winglian__tiny-shakespeare"
+    return load_from_disk(ds_path)
+
+
+@pytest.fixture
+def dataset_tatsu_lab_alpaca(
+    download_ds_fixture_bundle: Path,
+):
+    from datasets import load_from_disk
+
+    ds_path = download_ds_fixture_bundle / "tatsu-lab__alpaca"
+    return load_from_disk(ds_path)["train"]
+
+
+@pytest.fixture
+def dataset_mhenrichsen_alpaca_2k_test(
+    download_ds_fixture_bundle: Path,
+):
+    from datasets import load_from_disk
+
+    ds_path = download_ds_fixture_bundle / "mhenrichsen__alpaca_2k_test"
+    return load_from_disk(ds_path)["train"]
+
+
+@pytest.fixture
+def dataset_argilla_ultrafeedback_binarized_preferences_cleaned(
+    download_ds_fixture_bundle: Path,
+):
+    from datasets import load_from_disk
+
+    ds_path = (
+        download_ds_fixture_bundle
+        / "argilla__ultrafeedback-binarized-preferences-cleaned"
+    )
+    return load_from_disk(ds_path)["train"]
+
+
+@pytest.fixture
+def dataset_fozziethebeat_alpaca_messages_2k_dpo_test(
+    download_ds_fixture_bundle: Path,
+):
+    from datasets import load_from_disk
+
+    ds_path = download_ds_fixture_bundle / "fozziethebeat__alpaca_messages_2k_dpo_test"
+    return load_from_disk(ds_path)["train"]
+
+
+@pytest.fixture
+def dataset_fozziethebeat_alpaca_messages_2k_dpo_test_rev_ea82cff(
+    download_ds_fixture_bundle: Path,
+):
+    from datasets import load_from_disk
+
+    ds_path = (
+        download_ds_fixture_bundle
+        / "fozziethebeat__alpaca_messages_2k_dpo_test__rev_ea82cff"
+    )
+    return load_from_disk(ds_path)["train"]
+
+
+@pytest.fixture(name="min_base_cfg")
+def fixture_min_base_cfg():
+    return DictDefault(
+        base_model="HuggingFaceTB/SmolLM2-135M",
+        learning_rate=1e-3,
+        datasets=[
+            {
+                "path": "mhenrichsen/alpaca_2k_test",
+                "type": "alpaca",
+            },
+        ],
+        micro_batch_size=1,
+        gradient_accumulation_steps=1,
+    )
+
+
+#
+@pytest.mark.skipif(
+    os.environ.get("AXOLOTL_IS_CI_CACHE_PRELOAD", "-1") != "1",
+    reason="Not running in CI cache preload",
+)
+def test_load_fixtures(
+    download_smollm2_135m_model,
+    download_qwen3_half_billion_model,
+    download_tiny_llama_model,
+    download_tiny_mistral_model,
+    download_tiny_mixtral_model,
+    download_tiny_phi_model,
+    download_tiny_falcon_model,
+    download_tiny_qwen2_model,
+    download_tiny_qwen3_model,
+    download_tiny_gemma2_model,
+    download_tatsu_lab_alpaca_dataset,
+    download_mhenrichsen_alpaca_2k_dataset,
+    download_mhenrichsen_alpaca_2k_w_revision_dataset,
+    download_mlabonne_finetome_100k_dataset,
+    download_argilla_ultrafeedback_binarized_preferences_cleaned_dataset,
+    download_argilla_ultrafeedback_binarized_preferences_cleaned_kto_dataset,
+    download_argilla_distilabel_capybara_dpo_7k_binarized_dataset,
+    download_arcee_ai_distilabel_intel_orca_dpo_pairs_dataset,
+    download_argilla_dpo_pairs_dataset,
+    download_tiny_shakespeare_dataset,
+    download_deepseek_model_fixture,
+    download_huggyllama_model_fixture,
+    download_llama_1b_model_fixture,
+    download_llama3_8b_model_fixture,
+    download_llama3_8b_instruct_model_fixture,
+    download_phi_35_mini_model_fixture,
+    download_phi_3_medium_model_fixture,
+    download_phi_4_reasoning_model_fixture,
+    download_mistral_7b_model_fixture,
+    download_gemma_2b_model_fixture,
+    download_gemma2_9b_model_fixture,
+    download_mlx_mistral_7b_model_fixture,
+    download_llama2_model_fixture,
+):
+    pass
+
+
+@pytest.fixture(autouse=True)
+def disable_telemetry(monkeypatch):
+    monkeypatch.setenv("AXOLOTL_DO_NOT_TRACK", "1")
+    yield
diff --git a/tests/constants.py b/tests/constants.py
new file mode 100644
index 0000000000..310c09d753
--- /dev/null
+++ b/tests/constants.py
@@ -0,0 +1,57 @@
+# constants.py
+"""
+This module contains constants and configuration dictionaries used for
+datasets and other utilities in the Axolotl project, specifically for testing.
+"""
+
+# Configuration for Alpaca Messages Dataset
+ALPACA_MESSAGES_CONFIG_OG = {
+    "path": "fozziethebeat/alpaca_messages_2k_dpo_test",
+    "split": "train[:16]",
+    "type": "chat_template.default",
+    "chat_template": "llama3",
+    "field_messages": "conversation",
+    "field_chosen": "chosen",
+    "field_rejected": "rejected",
+    "message_field_role": "role",
+    "message_field_content": "content",
+    "roles": {
+        "system": ["system"],
+        "user": ["user"],
+        "assistant": ["assistant"],
+    },
+}
+
+
+def alpaca_messages_dpo_rows(num_rows: int = 16) -> list[dict]:
+    return [
+        {
+            "conversation": [
+                {
+                    "role": "user",
+                    "content": f"Which option is best for example {idx}?",
+                }
+            ],
+            "chosen": {
+                "role": "assistant",
+                "content": f"The best option for example {idx} is the concise answer.",
+            },
+            "rejected": {
+                "role": "assistant",
+                "content": f"Example {idx} has no useful answer.",
+            },
+        }
+        for idx in range(num_rows)
+    ]
+
+
+# Revision configuration extending the original
+ALPACA_MESSAGES_CONFIG_REVISION = ALPACA_MESSAGES_CONFIG_OG.copy()
+ALPACA_MESSAGES_CONFIG_REVISION["revision"] = "ea82cff"
+
+
+SPECIAL_TOKENS = {
+    "bos_token": "<s>",
+    "eos_token": "</s>",
+    "unk_token": "<unk>",
+}
diff --git a/tests/core/chat/__init__.py b/tests/core/chat/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/core/chat/format/__init__.py b/tests/core/chat/format/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/core/chat/test_messages.py b/tests/core/chat/test_messages.py
new file mode 100644
index 0000000000..c1d5cbcbe7
--- /dev/null
+++ b/tests/core/chat/test_messages.py
@@ -0,0 +1,201 @@
+"""
+Tests for the chat messages module
+"""
+
+import unittest
+
+import pytest
+from transformers import AddedToken, AutoTokenizer
+
+from axolotl.core.chat.format.chatml import format_message
+from axolotl.core.chat.messages import ChatFormattedChats, Chats
+
+from tests.hf_offline_utils import enable_hf_offline  # noqa
+
+
+@pytest.fixture(scope="session", name="llama_tokenizer")
+@enable_hf_offline
+def llama_tokenizer_fixture():
+    return AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B")
+
+
+@pytest.fixture(scope="session", name="chatml_tokenizer")
+def llama_tokenizer_w_chatml(llama_tokenizer):
+    llama_tokenizer.add_special_tokens(
+        {
+            "eos_token": AddedToken(
+                "<|im_end|>", rstrip=False, lstrip=False, normalized=False
+            )
+        }
+    )
+    llama_tokenizer.add_tokens(
+        [
+            AddedToken("<|im_start|>", rstrip=False, lstrip=False, normalized=False),
+        ]
+    )
+
+    return llama_tokenizer
+
+
+@pytest.fixture(scope="session", name="chat_msgs")
+def chat_msgs_fixture():
+    return {
+        "conversation": [
+            {
+                "role": "system",
+                "content": [
+                    {"type": "text", "value": "You are a helpful assistant."},
+                ],
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "value": "What is today's stock price of Apple?"},
+                ],
+            },
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "tool_call",
+                        "value": {
+                            "name": "get_date",
+                            "arguments": {},
+                        },
+                    },
+                    {
+                        "type": "tool_call",
+                        "value": {
+                            "name": "get_stock_price",
+                            "arguments": {"symbol": "AAPL"},
+                        },
+                    },
+                ],
+                "weight": 1,
+            },
+            {
+                "role": "tool",
+                "content": [
+                    {
+                        "type": "tool_response",
+                        "value": {
+                            "name": "get_date",
+                            "content": {"date": "2024-09-09"},
+                        },
+                    },
+                    {
+                        "type": "tool_response",
+                        "value": {
+                            "name": "get_stock_price",
+                            "content": {"symbol": "AAPL", "price": 123.45},
+                        },
+                    },
+                ],
+            },
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "text",
+                        "value": "The stock price of Apple is $123.45.\n",
+                        "weight": 0,
+                    },
+                    {
+                        "type": "text",
+                        "value": "<reflection>The original query asked for today's stock price of Apple. This implies they also wanted the date included in the response.</reflection>",
+                    },
+                    {
+                        "type": "text",
+                        "value": "The stock price of Apple on September 9, 2024 is $123.45.",
+                    },
+                ],
+                "weight": 1,
+            },
+        ]
+    }
+
+
+class TestMessagesCase:
+    """
+    Test cases for the chat messages module
+    """
+
+    def test_tool_call_stringify(self, chat_msgs):
+        chat_msgs_as_obj = Chats(**chat_msgs)
+        assert '{"name": "get_stock_price", "arguments": {"symbol": "AAPL"}}' == str(
+            chat_msgs_as_obj.conversation[2].content[1].value
+        )
+
+    def test_chatml_formatted_wrapper(self, chat_msgs):
+        chat_msg_formatted = ChatFormattedChats(**chat_msgs, formatter=format_message)
+        target_chatml = """<|im_start|>system
+You are a helpful assistant.<|im_end|>
+<|im_start|>user
+What is today's stock price of Apple?<|im_end|>
+<|im_start|>assistant
+<tool_call>
+{"name": "get_date", "arguments": {}}
+</tool_call>
+<tool_call>
+{"name": "get_stock_price", "arguments": {"symbol": "AAPL"}}
+</tool_call>
+<|im_end|>
+<|im_start|>tool
+<tool_response>
+{"name": "get_date", "content": {"date": "2024-09-09"}}
+</tool_response>
+<tool_response>
+{"name": "get_stock_price", "content": {"symbol": "AAPL", "price": 123.45}}
+</tool_response>
+<|im_end|>
+<|im_start|>assistant
+The stock price of Apple is $123.45.
+<reflection>The original query asked for today's stock price of Apple. This implies they also wanted the date included in the response.</reflection>The stock price of Apple on September 9, 2024 is $123.45.<|im_end|>\n"""
+        assert target_chatml == str(chat_msg_formatted)
+
+    def test_chatml_formatting_tool_call(self, chat_msgs):
+        chat_msgs_as_obj = Chats(**chat_msgs)
+        target_chatml_turn2 = """<|im_start|>assistant\n<tool_call>\n{"name": "get_date", "arguments": {}}\n</tool_call>\n<tool_call>\n{"name": "get_stock_price", "arguments": {"symbol": "AAPL"}}\n</tool_call>\n<|im_end|>\n"""
+        assert target_chatml_turn2 == str(
+            format_message(chat_msgs_as_obj.conversation[2])
+        )
+
+    def test_train_labels(self, chatml_tokenizer, chat_msgs):
+        chat_msg_formatted = ChatFormattedChats(**chat_msgs, formatter=format_message)
+        tokenized = chat_msg_formatted.conversation[2].tokenized(chatml_tokenizer)
+        # fmt: off
+        target_labels = [
+            -100, -100, -100,  # role
+            27, 14506, 13735, 397, 5018, 609, 794,
+            330, 456, 4257, 498, 330, 16774, 794, 4792, 534, 524,
+            14506, 13735, 397, 27, 14506, 13735, 397, 5018, 609, 794,
+            330, 456, 31641, 9217, 498, 330, 16774, 794, 5324, 19314,
+            794, 330, 84016, 43, 96742, 524, 14506, 13735, 397,
+            128256,  # <|im_end|>
+            -100  # trailing newline
+        ]
+        # fmt: on
+        assert tokenized["labels"] == target_labels
+
+    def test_train_labels_2(self, chatml_tokenizer, chat_msgs):
+        # also test if indivudal contents are set not to train
+        chat_msg_formatted = ChatFormattedChats(**chat_msgs, formatter=format_message)
+        tokenized = chat_msg_formatted.conversation[4].tokenized(chatml_tokenizer)
+        # fmt: off
+        target_labels = [
+            -100, -100, -100,  # role
+            -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,  # initial response
+            27, 78098, 16761, 4113, 3319, 4691, 369, 3432, 596, 5708, 3430,
+            315, 8325, 13, 1115, 24897, 814, 1101, 4934, 279, 2457,
+            5343, 304, 279, 2077, 4005, 78098, 16761, 5708, 3430, 315,
+            8325, 389, 6250, 220, 24, 11, 220, 2366, 19, 374, 400,
+            4513, 13, 1774, 13,
+            128256,  # <|im_end|>
+            -100,  # trailing newline
+        ]
+        # fmt: on
+        assert tokenized["labels"] == target_labels
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/conftest.py b/tests/core/conftest.py
new file mode 100644
index 0000000000..085447b35a
--- /dev/null
+++ b/tests/core/conftest.py
@@ -0,0 +1,245 @@
+"""Shared fixtures for axolotl.core.builders trainer-builder tests."""
+
+import pytest
+
+from axolotl.loaders import ModelLoader, load_tokenizer
+from axolotl.utils.config import normalize_config
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.schemas.enums import RLType
+
+
+@pytest.fixture(name="base_cfg")
+def fixture_base_cfg():
+    """
+    Base config with all common arguments between SFT and RLHF
+    """
+    cfg = DictDefault(
+        {
+            # Model and tokenizer settings
+            "base_model": "HuggingFaceTB/SmolLM2-135M-Instruct",
+            "sequence_len": 2048,
+            "model_config_type": "llama",  # example type
+            # Basic training settings
+            "micro_batch_size": 2,
+            "eval_batch_size": 2,
+            "num_epochs": 1,
+            "gradient_accumulation_steps": 1,
+            "max_steps": 100,
+            "val_set_size": 0,
+            # Optimizer settings
+            "optimizer": "adamw_torch_fused",
+            "learning_rate": 0.00005,
+            "weight_decay": 0.01,
+            "adam_beta1": 0.998,
+            "adam_beta2": 0.9,
+            "adam_epsilon": 0.00001,
+            "max_grad_norm": 1.0,
+            # LR scheduler settings
+            "lr_scheduler": "cosine",
+            "lr_scheduler_kwargs": {"foo": "bar"},
+            "warmup_steps": 10,
+            "warmup_ratio": None,
+            "cosine_min_lr_ratio": 0.1,
+            "cosine_constant_lr_ratio": 0.2,
+            # Checkpointing and saving
+            "save_steps": 100,
+            "output_dir": "./model-out",
+            "save_total_limit": 4,
+            "save_only_model": False,
+            # Hardware/performance settings
+            "gradient_checkpointing": False,
+            "gradient_checkpointing_kwargs": {"use_reentrant": False},
+            "dataloader_num_workers": 1,
+            "dataloader_pin_memory": True,
+            "dataloader_prefetch_factor": 2,
+            "context_parallel_size": 1,
+            "tensor_parallel_size": 1,
+            # Dtype
+            "fp16": False,
+            "bf16": False,
+            "tf32": False,
+            # Logging and evaluation
+            "logging_steps": 10,
+            "eval_steps": 50,
+            "eval_strategy": "steps",
+            "save_strategy": "steps",
+            "include_tokens_per_second": True,
+            # Other common settings
+            "seed": 42,
+            "remove_unused_columns": True,
+            "ddp_timeout": 1800,
+            "ddp_bucket_cap_mb": 25,
+            "ddp_broadcast_buffers": False,
+            "dataset_num_proc": 1,
+        }
+    )
+
+    normalize_config(cfg)
+    return cfg
+
+
+@pytest.fixture(name="dpo_cfg")
+def fixture_dpo_cfg(base_cfg):
+    cfg = base_cfg.copy()
+    cfg.update(
+        {
+            "rl": RLType.DPO,
+            "dpo_use_weighting": True,
+            "dpo_label_smoothing": 0.1,
+            "beta": 0.1,  # DPO beta
+            "dpo_loss_type": ["sigmoid", "sft"],
+            "dpo_loss_weights": [1.0, 0.5],
+        }
+    )
+    return cfg
+
+
+@pytest.fixture(name="orpo_cfg")
+def fixture_orpo_cfg(base_cfg):
+    cfg = base_cfg.copy()
+    cfg.update(
+        {
+            "rl": RLType.ORPO,
+            "orpo_alpha": 0.1,
+            "max_prompt_len": 512,
+        }
+    )
+    return cfg
+
+
+@pytest.fixture(name="kto_cfg")
+def fixture_kto_cfg(base_cfg):
+    cfg = base_cfg.copy()
+    cfg.update(
+        {
+            "rl": RLType.KTO,
+            "kto_desirable_weight": 1.0,
+            "kto_undesirable_weight": 1.0,
+            "max_prompt_len": 512,
+        }
+    )
+    return cfg
+
+
+@pytest.fixture(name="grpo_cfg")
+def fixture_grpo_cfg(base_cfg):
+    cfg = base_cfg.copy()
+    cfg.update(
+        {
+            "rl": RLType.GRPO,
+            "trl": DictDefault(
+                {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": False,  # run on CPU
+                    # "vllm_device": "auto",
+                    # "vllm_gpu_memory_utilization": 0.15,
+                    "num_generations": 4,
+                    "reward_funcs": ["rewards.rand_reward_func"],
+                }
+            ),
+            # Must be evenly divisible by num_generations
+            "micro_batch_size": 4,
+            "datasets": [
+                {
+                    "path": "openai/gsm8k",
+                    "name": "main",
+                    "split": "train[:16]",
+                }
+            ],
+        }
+    )
+    return DictDefault(cfg)
+
+
+@pytest.fixture(name="ipo_cfg")
+def fixture_ipo_cfg(base_cfg):
+    cfg = base_cfg.copy()
+    cfg.update(
+        {
+            "rl": RLType.DPO,
+            "dpo_loss_type": ["ipo"],
+            "dpo_label_smoothing": 0,
+            "beta": 0.1,
+        }
+    )
+    return cfg
+
+
+@pytest.fixture(name="simpo_cfg")
+def fixture_simpo_cfg(base_cfg):
+    cfg = base_cfg.copy()
+    cfg.update(
+        {
+            "rl": RLType.SIMPO,
+            "rl_beta": 0.2,
+            "cpo_alpha": 0.9,
+            "simpo_gamma": 0.4,
+        }
+    )
+    return cfg
+
+
+@pytest.fixture(name="sft_cfg")
+def fixture_sft_cfg(base_cfg):
+    cfg = base_cfg.copy()
+    cfg.update(
+        {
+            "rl": None,
+            "sample_packing": False,
+            "eval_sample_packing": False,
+            "flash_attention": False,
+        }
+    )
+    return cfg
+
+
+@pytest.fixture(name="rm_cfg")
+def fixture_rm_cfg(sft_cfg):
+    cfg = sft_cfg.copy()
+    cfg.update(
+        DictDefault(
+            {
+                "reward_model": True,
+                "datasets": [
+                    {
+                        "path": "argilla/distilabel-intel-orca-dpo-pairs",
+                        "type": "bradley_terry.chat_template",
+                        "split": "train[:16]",
+                    }
+                ],
+            }
+        )
+    )
+    return cfg
+
+
+@pytest.fixture(name="prm_cfg")
+def fixture_prm_cfg(sft_cfg):
+    cfg = sft_cfg.copy()
+    cfg.update(
+        DictDefault(
+            {
+                "process_reward_model": True,
+                "datasets": [
+                    {
+                        "path": "trl-lib/math_shepherd",
+                        "type": "stepwise_supervised",
+                        "split": "train[:16]",
+                    }
+                ],
+            }
+        )
+    )
+    return cfg
+
+
+@pytest.fixture(name="tokenizer")
+def fixture_tokenizer(base_cfg):
+    return load_tokenizer(base_cfg)
+
+
+@pytest.fixture(name="model")
+def fixture_model(base_cfg, tokenizer):
+    model, _ = ModelLoader(base_cfg, tokenizer).load()
+    return model
diff --git a/tests/core/test_async_grpo.py b/tests/core/test_async_grpo.py
new file mode 100644
index 0000000000..3a4c188bcc
--- /dev/null
+++ b/tests/core/test_async_grpo.py
@@ -0,0 +1,412 @@
+"""Unit tests for async GRPO"""
+
+import unittest
+from unittest.mock import MagicMock
+
+import torch
+
+
+class TestReplayBuffer(unittest.TestCase):
+    """Tests for ReplayBuffer edge cases."""
+
+    def test_add_noop_when_max_size_zero(self):
+        from axolotl.core.trainers.grpo.replay_buffer import ReplayBuffer
+
+        buf = ReplayBuffer(max_size=0)
+        buf.add(1.0, {"data": "test"})
+        self.assertEqual(len(buf), 0)
+
+    def test_add_noop_when_max_size_negative(self):
+        from axolotl.core.trainers.grpo.replay_buffer import ReplayBuffer
+
+        buf = ReplayBuffer(max_size=-1)
+        buf.add(1.0, {"data": "test"})
+        self.assertEqual(len(buf), 0)
+
+    def test_sample_returns_none_when_max_size_zero(self):
+        from axolotl.core.trainers.grpo.replay_buffer import ReplayBuffer
+
+        buf = ReplayBuffer(max_size=0)
+        self.assertIsNone(buf.sample(1))
+
+    def test_sample_returns_none_when_empty(self):
+        from axolotl.core.trainers.grpo.replay_buffer import ReplayBuffer
+
+        buf = ReplayBuffer(max_size=5)
+        self.assertIsNone(buf.sample(1))
+
+    def test_normal_add_and_sample(self):
+        from axolotl.core.trainers.grpo.replay_buffer import ReplayBuffer
+
+        buf = ReplayBuffer(max_size=3)
+        buf.add(1.0, {"a": 1})
+        buf.add(2.0, {"a": 2})
+        buf.add(3.0, {"a": 3})
+        self.assertEqual(len(buf), 3)
+        result = buf.sample(1)
+        self.assertIsNotNone(result)
+        self.assertEqual(len(result), 1)
+
+    def test_replaces_lowest_when_full(self):
+        from axolotl.core.trainers.grpo.replay_buffer import ReplayBuffer
+
+        buf = ReplayBuffer(max_size=2)
+        buf.add(1.0, {"a": 1})
+        buf.add(2.0, {"a": 2})
+        buf.add(3.0, {"a": 3})  # should replace score=1.0
+        self.assertEqual(len(buf), 2)
+        scores = sorted(item[0] for item in buf._heap)
+        self.assertEqual(scores, [2.0, 3.0])
+
+
+class TestGRPOStrategyConflict(unittest.TestCase):
+    """Tests for sequence_parallel + async_grpo conflict detection."""
+
+    def test_raises_on_both_enabled(self):
+        from axolotl.core.trainers.grpo import GRPOStrategy
+
+        with self.assertRaises(ValueError) as ctx:
+            GRPOStrategy.get_trainer_class(sequence_parallel=True, async_grpo=True)
+        self.assertIn("sequence_parallel", str(ctx.exception))
+        self.assertIn("async_grpo", str(ctx.exception))
+
+    def test_sequence_parallel_only(self):
+        from axolotl.core.trainers.grpo import GRPOStrategy
+        from axolotl.core.trainers.grpo.trainer import (
+            AxolotlGRPOSequenceParallelTrainer,
+        )
+
+        cls = GRPOStrategy.get_trainer_class(sequence_parallel=True, async_grpo=False)
+        self.assertIs(cls, AxolotlGRPOSequenceParallelTrainer)
+
+    def test_async_only(self):
+        from axolotl.core.trainers.grpo import GRPOStrategy
+        from axolotl.core.trainers.grpo.trainer import AxolotlAsyncGRPOTrainer
+
+        cls = GRPOStrategy.get_trainer_class(sequence_parallel=False, async_grpo=True)
+        self.assertIs(cls, AxolotlAsyncGRPOTrainer)
+
+    def test_neither(self):
+        from axolotl.core.trainers.grpo import GRPOStrategy
+        from axolotl.core.trainers.grpo.trainer import AxolotlGRPOTrainer
+
+        cls = GRPOStrategy.get_trainer_class(sequence_parallel=False, async_grpo=False)
+        self.assertIs(cls, AxolotlGRPOTrainer)
+
+
+class TestDequantizeFP8TailBlocks(unittest.TestCase):
+    """Tests for FP8 dequantization with non-divisible dimensions."""
+
+    def test_exact_divisible_shape(self):
+        from axolotl.kernels.quantize import dequantize_fp8
+
+        W = torch.randn(256, 128, dtype=torch.bfloat16).to(torch.float8_e4m3fn)
+        scale_inv = torch.ones(2, 1, dtype=torch.bfloat16)
+        result = dequantize_fp8(W, scale_inv)
+        self.assertEqual(result.shape, (256, 128))
+        self.assertEqual(result.dtype, torch.bfloat16)
+
+    def test_non_divisible_rows(self):
+        from axolotl.kernels.quantize import dequantize_fp8
+
+        # 130 rows, scale has 2 blocks (block_size ~65 for exact div, but with
+        # tail blocks: first block=65 rows, second=65 rows, 130%2=0 actually).
+        # Use 131 rows with 2 scale blocks to trigger tail handling.
+        W = torch.ones(131, 128, dtype=torch.bfloat16).to(torch.float8_e4m3fn)
+        scale_inv = torch.tensor([[2.0], [3.0]], dtype=torch.bfloat16)
+        result = dequantize_fp8(W, scale_inv)
+        self.assertEqual(result.shape, (131, 128))
+        self.assertEqual(result.dtype, torch.bfloat16)
+
+    def test_non_divisible_cols(self):
+        from axolotl.kernels.quantize import dequantize_fp8
+
+        W = torch.ones(128, 200, dtype=torch.bfloat16).to(torch.float8_e4m3fn)
+        scale_inv = torch.ones(1, 2, dtype=torch.bfloat16)
+        result = dequantize_fp8(W, scale_inv)
+        self.assertEqual(result.shape, (128, 200))
+
+    def test_scalar_scale(self):
+        from axolotl.kernels.quantize import dequantize_fp8
+
+        W = torch.ones(64, 64, dtype=torch.bfloat16).to(torch.float8_e4m3fn)
+        scale_inv = torch.tensor(2.0, dtype=torch.bfloat16)
+        result = dequantize_fp8(W, scale_inv)
+        self.assertEqual(result.shape, (64, 64))
+
+
+class TestLoraFP8Guard(unittest.TestCase):
+    """Tests that get_lora_parameters only uses weight_scale_inv for FP8 weights."""
+
+    def test_non_fp8_weight_skips_scale_inv(self):
+        """Non-FP8 weight should NOT pick up weight_scale_inv as quant_state."""
+        from axolotl.kernels.lora import get_lora_parameters
+
+        proj = MagicMock()
+        proj.disable_adapters = True
+        base_layer = MagicMock(spec=[])  # empty spec to control attrs precisely
+
+        # Use a real tensor for weight (bf16, no quant_state attr)
+        base_layer.weight = torch.randn(64, 64, dtype=torch.bfloat16)
+        base_layer.bias = None
+        base_layer.weight_scale_inv = torch.ones(1)  # should NOT be used for bf16
+
+        proj.base_layer = base_layer
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(proj)
+        # quant_state should be None since weight is bf16, not FP8
+        self.assertIsNone(quant_state)
+
+    def test_fp8_weight_uses_scale_inv(self):
+        """FP8 weight should pick up weight_scale_inv as quant_state."""
+        from axolotl.kernels.lora import get_lora_parameters
+
+        proj = MagicMock()
+        proj.disable_adapters = True
+        base_layer = MagicMock()
+        proj.base_layer = base_layer
+
+        # FP8 weight
+        base_layer.weight = torch.randn(64, 64, dtype=torch.bfloat16).to(
+            torch.float8_e4m3fn
+        )
+        base_layer.bias = None
+        scale_inv = torch.ones(1)
+        base_layer.weight_scale_inv = scale_inv
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(proj)
+        self.assertIs(quant_state, scale_inv)
+
+
+class TestValidateQuantPatchRestore(unittest.TestCase):
+    """Test that validate_quantization_for_training is restored after trainer creation."""
+
+    def test_patch_restored_on_success(self):
+        """Monkeypatch should be restored even after successful trainer creation."""
+        import transformers.trainer as _trainer_module
+
+        original = _trainer_module.validate_quantization_for_training
+
+        # After the build() method runs, original should be restored.
+        # We can't easily test the full build(), but we can test the pattern.
+        _orig = _trainer_module.validate_quantization_for_training
+        _trainer_module.validate_quantization_for_training = lambda model: None
+        try:
+            pass  # simulate trainer_cls() succeeding
+        finally:
+            _trainer_module.validate_quantization_for_training = _orig
+
+        self.assertIs(_trainer_module.validate_quantization_for_training, original)
+
+    def test_patch_restored_on_error(self):
+        """Monkeypatch should be restored even if trainer creation raises."""
+        import transformers.trainer as _trainer_module
+
+        original = _trainer_module.validate_quantization_for_training
+
+        _orig = _trainer_module.validate_quantization_for_training
+        _trainer_module.validate_quantization_for_training = lambda model: None
+        try:
+            raise ValueError("test error")
+        except ValueError:
+            pass
+        finally:
+            _trainer_module.validate_quantization_for_training = _orig
+
+        self.assertIs(_trainer_module.validate_quantization_for_training, original)
+
+
+class TestVllmLoraSyncPatch(unittest.TestCase):
+    """The ``_generate_single_turn`` patch wires sync_weights to the right place.
+
+    These tests exercise the patch-installation branch in isolation. They build
+    a stub trainer with just enough attributes to look like
+    ``AsyncGRPOTrainer`` for the duration of the relevant code path.
+
+    Background — there are two correct behaviors and we historically had a bug
+    where both modes used the same one:
+
+      - Async prefetch ON: the BG generation thread can't safely call
+        sync_weights mid-rollout. We no-op the stock hook and drive sync from
+        the main thread via ``_maybe_sync_vllm_weights``.
+      - Async prefetch OFF: TRL's stock ``_generate_single_turn`` already
+        calls ``sync_weights`` once per step boundary on the main thread. We
+        wire that hook directly to ``_sync_lora_adapter`` because
+        ``_maybe_sync_vllm_weights`` short-circuits when async is off.
+
+    Before the fix, both modes installed ``lambda: None``, so sync mode never
+    pushed any LoRA adapter to vLLM and the trainer was a no-op.
+    """
+
+    @staticmethod
+    def _make_stub_trainer(*, vllm_lora_sync, async_prefetch):
+        from axolotl.core.trainers.grpo.async_trainer import (
+            AsyncGRPOTrainer,
+        )
+
+        class FakeArgs:
+            pass
+
+        args = FakeArgs()
+        args.vllm_lora_sync = vllm_lora_sync
+        args.async_prefetch = async_prefetch
+
+        class FakeVllmGen:
+            sync_weights = staticmethod(lambda: None)
+            model = MagicMock()
+
+        # Use object.__new__ so we don't run __init__ (which needs a real
+        # model, dataset, etc.). We only need the `_generate_single_turn`
+        # method's patch branch to run, so we set up the minimum state.
+        trainer = object.__new__(AsyncGRPOTrainer)
+        trainer.args = args
+        trainer.use_vllm = True
+        trainer.vllm_generation = FakeVllmGen()
+        trainer._patched_sync_weights = False
+        # Spy on _sync_lora_adapter so we can assert it's the function the
+        # hook delegates to in sync mode.
+        trainer._sync_lora_adapter = MagicMock(name="_sync_lora_adapter_spy")
+        trainer._sync_peft_weights_no_merge = MagicMock(
+            name="_sync_peft_weights_no_merge_spy"
+        )
+        return trainer
+
+    @staticmethod
+    def _run_patch_branch(trainer):
+        """Execute just the sync_weights-patching branch in isolation.
+
+        We can't easily call the real ``_generate_single_turn`` because it
+        does a full vLLM generate. Instead we copy the exact branch out of
+        the source so the test verifies the same logic the trainer runs.
+        """
+        if not getattr(trainer, "_patched_sync_weights", False):
+            if trainer.use_vllm and hasattr(trainer, "vllm_generation"):
+                if getattr(trainer.args, "vllm_lora_sync", False):
+                    if getattr(trainer.args, "async_prefetch", False):
+                        trainer.vllm_generation.sync_weights = lambda: None
+                    else:
+                        sync_helper = trainer._sync_lora_adapter
+
+                        def _lora_filesystem_sync():
+                            sync_helper()
+
+                        trainer.vllm_generation.sync_weights = _lora_filesystem_sync
+                    trainer._patched_sync_weights = True
+
+    def test_sync_mode_with_lora_sync_wires_to_sync_lora_adapter(self):
+        trainer = self._make_stub_trainer(vllm_lora_sync=True, async_prefetch=False)
+        self._run_patch_branch(trainer)
+
+        assert trainer._patched_sync_weights is True
+        # Trigger the patched hook — it must call _sync_lora_adapter.
+        trainer.vllm_generation.sync_weights()
+        trainer._sync_lora_adapter.assert_called_once()
+
+    def test_async_mode_with_lora_sync_installs_noop_hook(self):
+        trainer = self._make_stub_trainer(vllm_lora_sync=True, async_prefetch=True)
+        self._run_patch_branch(trainer)
+
+        assert trainer._patched_sync_weights is True
+        # Hook must be a no-op so BG-thread generation doesn't fight the
+        # main-thread optimizer step over the model weights.
+        trainer.vllm_generation.sync_weights()
+        trainer._sync_lora_adapter.assert_not_called()
+
+    def test_sync_mode_with_lora_sync_does_not_call_during_install(self):
+        """Installing the patch should not pre-emptively sync."""
+        trainer = self._make_stub_trainer(vllm_lora_sync=True, async_prefetch=False)
+        self._run_patch_branch(trainer)
+        # _sync_lora_adapter should only be called when the patched hook
+        # itself is invoked (e.g., from TRL's _generate_single_turn).
+        trainer._sync_lora_adapter.assert_not_called()
+
+    def test_patch_is_idempotent(self):
+        trainer = self._make_stub_trainer(vllm_lora_sync=True, async_prefetch=False)
+        self._run_patch_branch(trainer)
+        first_hook = trainer.vllm_generation.sync_weights
+        # Second call must not re-patch (otherwise we'd lose the original).
+        self._run_patch_branch(trainer)
+        assert trainer.vllm_generation.sync_weights is first_hook
+
+
+class TestMaybeSyncVllmWeightsIntervalDefault(unittest.TestCase):
+    """``_maybe_sync_vllm_weights`` must not crash when interval is unset.
+
+    Before the fix, ``step % self.args.vllm_sync_interval`` would TypeError
+    on the very first call when ``vllm_sync_interval`` was ``None`` (which
+    is the default for any config that doesn't explicitly set it). We now
+    fall back to interval=1 so unset means "sync every step", matching the
+    behavior of TRL's own ``_generate_single_turn``.
+    """
+
+    @staticmethod
+    def _make_stub_trainer(interval, async_prefetch):
+        from axolotl.core.trainers.grpo.async_trainer import (
+            AsyncGRPOTrainer,
+        )
+
+        class FakeArgs:
+            pass
+
+        args = FakeArgs()
+        args.async_prefetch = async_prefetch
+        args.vllm_sync_interval = interval
+        args.vllm_lora_sync = True
+
+        class FakeState:
+            global_step = 1
+
+        trainer = object.__new__(AsyncGRPOTrainer)
+        trainer.args = args
+        trainer.use_vllm = True
+        trainer.state = FakeState()
+        trainer._last_synced_step = 0
+        trainer._sync_lora_adapter = MagicMock(name="sync_spy")
+        return trainer
+
+    def test_interval_none_in_async_mode_does_not_crash(self):
+        trainer = self._make_stub_trainer(interval=None, async_prefetch=True)
+        from axolotl.core.trainers.grpo.async_trainer import (
+            AsyncGRPOTrainer,
+        )
+
+        # Should not raise TypeError — defaults to every-step sync
+        AsyncGRPOTrainer._maybe_sync_vllm_weights(trainer)
+        trainer._sync_lora_adapter.assert_called_once()
+
+    def test_sync_mode_drives_sync(self):
+        """Sync mode must fire ``_sync_lora_adapter`` from ``_maybe_sync_vllm_weights``.
+
+        The previous behavior (early return when ``not async_prefetch``)
+        assumed TRL's stock ``_generate_single_turn`` would handle sync.
+        That's true for vanilla GRPO but FALSE for NeMo Gym multi-turn
+        where the data producer bypasses ``_generate_single_turn``
+        entirely. Without this trigger no sync ever happens and the
+        trainer becomes a no-op.
+        """
+        trainer = self._make_stub_trainer(interval=1, async_prefetch=False)
+        from axolotl.core.trainers.grpo.async_trainer import (
+            AsyncGRPOTrainer,
+        )
+
+        AsyncGRPOTrainer._maybe_sync_vllm_weights(trainer)
+        trainer._sync_lora_adapter.assert_called_once()
+
+    def test_async_mode_with_explicit_interval_respects_modulo(self):
+        trainer = self._make_stub_trainer(interval=4, async_prefetch=True)
+        from axolotl.core.trainers.grpo.async_trainer import (
+            AsyncGRPOTrainer,
+        )
+
+        # global_step=1, interval=4 → 1 % 4 != 0 → no sync
+        AsyncGRPOTrainer._maybe_sync_vllm_weights(trainer)
+        trainer._sync_lora_adapter.assert_not_called()
+
+        # global_step=4 → 4 % 4 == 0 → sync
+        trainer.state.global_step = 4
+        AsyncGRPOTrainer._maybe_sync_vllm_weights(trainer)
+        trainer._sync_lora_adapter.assert_called_once()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/test_builders.py b/tests/core/test_builders.py
new file mode 100644
index 0000000000..31eff9697c
--- /dev/null
+++ b/tests/core/test_builders.py
@@ -0,0 +1,197 @@
+"""Unit tests for axolotl.core.builders SFT and reward-model trainer builders."""
+
+from types import SimpleNamespace
+
+import pytest
+from datasets import Dataset
+
+from axolotl.core.builders import HFCausalTrainerBuilder, HFRLTrainerBuilder
+from axolotl.core.builders.base import TrainerBuilderBase
+
+
+def _gradient_checkpointing_kwargs(cfg):
+    training_args_kwargs = {}
+    TrainerBuilderBase._configure_gradient_checkpointing(
+        SimpleNamespace(cfg=cfg), training_args_kwargs
+    )
+    return training_args_kwargs
+
+
+class TestGradientCheckpointingConfig:
+    def test_hidden_states_offload_uses_non_reentrant_trainer_path(self):
+        training_args_kwargs = _gradient_checkpointing_kwargs(
+            SimpleNamespace(
+                layer_offloading=False,
+                activation_offloading="hidden_states",
+                gradient_checkpointing=True,
+                gradient_checkpointing_kwargs={"use_reentrant": False},
+            )
+        )
+
+        assert training_args_kwargs["gradient_checkpointing"] is True
+        assert training_args_kwargs["gradient_checkpointing_kwargs"] == {
+            "use_reentrant": False
+        }
+        assert training_args_kwargs["activation_offloading"] == "hidden_states"
+
+    def test_hidden_states_offload_with_reentrant_stays_on_model_loader_path(self):
+        training_args_kwargs = _gradient_checkpointing_kwargs(
+            SimpleNamespace(
+                layer_offloading=False,
+                activation_offloading="hidden_states",
+                gradient_checkpointing=True,
+                gradient_checkpointing_kwargs={"use_reentrant": True},
+            )
+        )
+
+        assert training_args_kwargs["gradient_checkpointing"] is True
+        assert training_args_kwargs["gradient_checkpointing_kwargs"] == {
+            "use_reentrant": True
+        }
+        assert "activation_offloading" not in training_args_kwargs
+
+
+def _reward_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "chosen_ids": [1, 2, 3],
+                "rejected_ids": [1, 4],
+            }
+        ]
+    )
+
+
+def _prm_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "input_ids": [1, 2, 3],
+                "labels": [-100, -100, 1],
+            }
+        ]
+    )
+
+
+class TestHFCausalTrainerBuilder:
+    """
+    TestCase class for SFT trainer builder
+    """
+
+    def test_training_arguments(self, sft_cfg, model, tokenizer):
+        builder = HFCausalTrainerBuilder(sft_cfg, model, tokenizer)
+        trainer = builder.build(100)
+        training_arguments = trainer.args
+
+        # Test common arguments
+        assert training_arguments.per_device_train_batch_size == 2
+        assert training_arguments.gradient_accumulation_steps == 1
+        assert training_arguments.max_steps == 100
+
+        assert training_arguments.learning_rate == 0.00005
+        assert training_arguments.weight_decay == 0.01
+        assert training_arguments.adam_beta1 == 0.998
+        assert training_arguments.adam_beta2 == 0.9
+        assert training_arguments.adam_epsilon == 0.00001
+        assert training_arguments.max_grad_norm == 1.0
+
+        assert training_arguments.lr_scheduler_type == "cosine"
+        assert training_arguments.warmup_steps == 10
+        assert training_arguments.cosine_min_lr_ratio == 0.1
+
+        assert training_arguments.dataloader_num_workers == 1
+        assert training_arguments.dataloader_pin_memory is True
+        assert training_arguments.gradient_checkpointing is False
+
+        # SFT specific
+        assert training_arguments.sample_packing is False
+        assert training_arguments.eval_sample_packing is False
+
+    @pytest.mark.parametrize(
+        "cfg_string",
+        [
+            "sft_cfg",
+            "rm_cfg",
+            "prm_cfg",
+        ],
+    )
+    def test_builder_w_rm_trainers(self, request, cfg_string, model, tokenizer):
+        cfg = request.getfixturevalue(cfg_string)
+        builder = HFCausalTrainerBuilder(cfg, model, tokenizer)
+        cfg["optimizer"] = "muon"
+
+        if cfg_string == "rm_cfg":
+            builder.train_dataset = _reward_dataset()
+        elif cfg_string == "prm_cfg":
+            builder.train_dataset = _prm_dataset()
+
+        trainer = builder.build(100)
+
+        assert trainer.optimizer_cls_and_kwargs is not None
+
+        from axolotl.contribs.mit.muon import MuonOptimizerFactory
+        from axolotl.contribs.mit.muon.muon import Muon
+
+        optimizer_cls, optimizer_kwargs = trainer.optimizer_cls_and_kwargs
+        assert optimizer_cls is MuonOptimizerFactory
+        assert optimizer_kwargs["lr"] == 0.00005
+        assert optimizer_kwargs["weight_decay"] == 0.01
+        assert optimizer_kwargs["betas"] == (0.998, 0.9)
+        assert optimizer_kwargs["eps"] == 0.00001
+
+        # Ensure optimizer is created with correct class
+        optim = trainer.create_optimizer()
+        assert isinstance(optim, Muon)
+
+    def test_sinkgd_optimizer(self, sft_cfg, model, tokenizer):
+        cfg = sft_cfg.copy()
+        cfg["optimizer"] = "sinkgd"
+        cfg["optim_args"] = {"sinkhorn_iters": 5, "sinkgd_lr_scale": 0.05}
+
+        builder = HFCausalTrainerBuilder(cfg, model, tokenizer)
+        trainer = builder.build(100)
+
+        assert trainer.optimizer_cls_and_kwargs is not None
+
+        from axolotl.utils.optimizers.sinkgd import SinkGD, SinkGDOptimizerFactory
+
+        optimizer_cls, optimizer_kwargs = trainer.optimizer_cls_and_kwargs
+        assert optimizer_cls is SinkGDOptimizerFactory
+        assert optimizer_kwargs["lr"] == 0.00005
+        assert optimizer_kwargs["weight_decay"] == 0.01
+        assert optimizer_kwargs["sinkhorn_iters"] == 5
+        assert optimizer_kwargs["sinkgd_lr_scale"] == 0.05
+
+        # OptimizerMixin must resolve the factory into a concrete SinkGD instance
+        optim = trainer.create_optimizer()
+        assert isinstance(optim, SinkGD)
+        # linear weight matrices must be in the stateless SR-Sinkhorn group
+        assert any(group.get("use_sinkgd") for group in optim.param_groups)
+
+
+class TestTrainerClsPlugin:
+    """
+    TestCase class for trainer builder with plugin
+    """
+
+    def test_trainer_cls_is_not_none_with_plugin(self, kto_cfg, model, tokenizer):
+        """
+        Test that the trainer cls is not none with plugin
+
+        Fixes #2693
+        """
+        cfg = kto_cfg.copy()
+        cfg.plugins = ["axolotl.integrations.liger.LigerPlugin"]
+
+        # Expected AttributeError as we don't pass regular model configs to RL trainer builder
+        # If it throws `TypeError: None is not a callable object`, trainer_cls could be None
+        try:
+            builder = HFRLTrainerBuilder(cfg, model, tokenizer)
+
+            builder.build(100)
+        except TypeError as e:
+            # Error raised if trainer_cls is None
+            assert "'tuple' object has no attribute 'config'" not in str(e)
+        except Exception:
+            # Another error happens, so we passed trainer_cls to builder
+            pass
diff --git a/tests/core/test_builders_rl.py b/tests/core/test_builders_rl.py
new file mode 100644
index 0000000000..c6c507d5db
--- /dev/null
+++ b/tests/core/test_builders_rl.py
@@ -0,0 +1,278 @@
+"""Unit tests for axolotl.core.builders RL (preference/GRPO) trainer builders."""
+
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+from datasets import Dataset
+
+from axolotl.core.builders import HFRLTrainerBuilder
+from axolotl.utils.data import prepare_preference_datasets
+from axolotl.utils.dict import DictDefault
+
+from tests.constants import ALPACA_MESSAGES_CONFIG_REVISION, alpaca_messages_dpo_rows
+
+
+def _preference_dataset(cfg_string: str) -> Dataset:
+    if cfg_string in {"dpo_cfg", "ipo_cfg", "grpo_cfg", "simpo_cfg"}:
+        return Dataset.from_list(alpaca_messages_dpo_rows())
+    if cfg_string == "orpo_cfg":
+        return Dataset.from_list(
+            [
+                {
+                    "chosen": [
+                        {"role": "user", "content": f"Question {idx}?"},
+                        {"role": "assistant", "content": f"Chosen answer {idx}."},
+                    ],
+                    "rejected": [
+                        {"role": "user", "content": f"Question {idx}?"},
+                        {"role": "assistant", "content": f"Rejected answer {idx}."},
+                    ],
+                }
+                for idx in range(16)
+            ]
+        )
+    if cfg_string == "kto_cfg":
+        return Dataset.from_list(
+            [
+                {
+                    "prompt": f"Question {idx}?",
+                    "completion": f"Answer {idx}.",
+                    "label": idx % 2 == 0,
+                }
+                for idx in range(16)
+            ]
+        )
+    raise ValueError(f"Unhandled cfg_string: {cfg_string}")
+
+
+class TestHFRLTrainerBuilder:
+    """
+    TestCase class for RLHF trainer builders
+    """
+
+    def _test_common_training_arguments(self, training_arguments, rl: str):
+        """Helper to test common arguments across all variants"""
+        # Basic training settings
+        if rl == "grpo":
+            # grpo_cfg's micro_batch_size is diff from others
+            assert training_arguments.per_device_train_batch_size == 4
+        else:
+            assert training_arguments.per_device_train_batch_size == 2
+        assert training_arguments.gradient_accumulation_steps == 1
+        assert training_arguments.max_steps == 100
+
+        # Optimizer settings
+        assert training_arguments.learning_rate == 0.00005
+        assert training_arguments.weight_decay == 0.01
+        assert training_arguments.adam_beta1 == 0.998
+        assert training_arguments.adam_beta2 == 0.9
+        assert training_arguments.adam_epsilon == 0.00001
+        assert training_arguments.max_grad_norm == 1.0
+
+        # LR scheduler settings
+        assert training_arguments.lr_scheduler_type == "cosine"
+        assert training_arguments.warmup_steps == 10
+        assert training_arguments.cosine_min_lr_ratio == 0.1
+        assert training_arguments.cosine_constant_lr_ratio == 0.2
+
+        # Other settings
+        assert training_arguments.dataloader_num_workers == 1
+        assert training_arguments.dataloader_pin_memory is True
+
+        # TODO(wing): restore once trl releases 0.22.0
+        # assert training_arguments.gradient_checkpointing is True
+
+    def test_dpo_training_arguments(self, dpo_cfg, model, tokenizer):
+        dpo_cfg["precompute_ref_log_probs"] = True
+        builder = HFRLTrainerBuilder(dpo_cfg, model, tokenizer)
+        training_arguments, _ = builder._build_training_arguments(100)
+
+        self._test_common_training_arguments(training_arguments, rl=dpo_cfg.rl)
+        # DPO specific
+        assert training_arguments.beta == 0.1
+        assert hasattr(training_arguments, "use_weighting")
+        assert training_arguments.use_weighting is True
+        assert training_arguments.label_smoothing == 0.1
+        assert training_arguments.precompute_ref_log_probs is True
+        assert training_arguments.loss_type == ["sigmoid", "sft"]
+        assert training_arguments.loss_weights == [1.0, 0.5]
+
+    def test_orpo_training_arguments(self, orpo_cfg, model, tokenizer):
+        builder = HFRLTrainerBuilder(orpo_cfg, model, tokenizer)
+        training_arguments, _ = builder._build_training_arguments(100)
+
+        self._test_common_training_arguments(training_arguments, rl=orpo_cfg.rl)
+        # ORPO specific
+        assert training_arguments.beta == 0.1  # maps from orpo_alpha
+
+    def test_kto_training_arguments(self, kto_cfg, model, tokenizer):
+        builder = HFRLTrainerBuilder(kto_cfg, model, tokenizer)
+        training_arguments, _ = builder._build_training_arguments(100)
+
+        self._test_common_training_arguments(training_arguments, rl=kto_cfg.rl)
+        # KTO specific
+        assert training_arguments.desirable_weight == 1.0
+        assert training_arguments.undesirable_weight == 1.0
+
+    def _write_rewards_file(self, rewards_dir: Path):
+        """
+        Writes reward function to local tmp path to be loaded on trainer building
+        """
+        # Create rewards.py in a directory we can import from
+        rewards_dir.mkdir()
+        rewards_file = rewards_dir / "rewards.py"
+        rewards_file.write_text(
+            """import random
+def rand_reward_func(prompts, completions) -> list[float]:
+    return [random.uniform(0, 1) for _ in completions]
+"""
+        )
+
+    def test_grpo_training_arguments(self, grpo_cfg, model, tokenizer, tmp_path):
+        rewards_dir = tmp_path / "rewards_test"
+        self._write_rewards_file(rewards_dir)
+
+        # Add the directory to Python path so we can import the module
+        sys.path.insert(0, str(rewards_dir))
+
+        try:
+            builder = HFRLTrainerBuilder(grpo_cfg, model, tokenizer)
+            training_arguments, _ = builder._build_training_arguments(100)
+            builder.train_dataset = MagicMock()
+
+            self._test_common_training_arguments(training_arguments, rl=grpo_cfg.rl)
+            # GRPO specific
+            assert training_arguments.beta == 0.001
+            assert training_arguments.max_completion_length == 256
+            assert training_arguments.use_vllm is False
+            # assert training_arguments.vllm_device == "auto"
+            # assert training_arguments.vllm_gpu_memory_utilization == 0.15
+            assert training_arguments.num_generations == 4
+
+            # Test trainer creation to verify reward_funcs
+            trainer = builder.build(100)
+
+            # Verify reward functions are properly loaded
+            assert len(trainer.reward_funcs) == 1
+            assert trainer.reward_funcs[0].__module__ == "rewards"
+            assert trainer.reward_funcs[0].__name__ == "rand_reward_func"
+        finally:
+            # remove imported module from path
+            if str(rewards_dir) in sys.path:
+                sys.path.remove(str(rewards_dir))
+
+    def test_ipo_training_arguments(self, ipo_cfg, model, tokenizer):
+        builder = HFRLTrainerBuilder(ipo_cfg, model, tokenizer)
+        training_arguments, _ = builder._build_training_arguments(100)
+
+        self._test_common_training_arguments(training_arguments, rl=ipo_cfg.rl)
+        # IPO specific
+        assert training_arguments.beta == 0.1
+        assert training_arguments.loss_type == ["ipo"]
+        assert training_arguments.label_smoothing == 0
+
+    def test_simpo_training_arguments(self, simpo_cfg, model, tokenizer):
+        builder = HFRLTrainerBuilder(simpo_cfg, model, tokenizer)
+        training_arguments, _ = builder._build_training_arguments(100)
+
+        self._test_common_training_arguments(training_arguments, rl=simpo_cfg.rl)
+        # SIMPO specific
+        assert training_arguments.beta == 0.2
+        assert training_arguments.cpo_alpha == 0.9
+        assert training_arguments.simpo_gamma == 0.4
+
+    @pytest.mark.parametrize(
+        "cfg_string",
+        [
+            "dpo_cfg",
+            "ipo_cfg",
+            "grpo_cfg",
+            "orpo_cfg",
+            "kto_cfg",
+        ],
+    )
+    def test_custom_optimizer_cls_and_kwargs(
+        self,
+        request,
+        cfg_string,
+        tmp_path,
+        model,
+        tokenizer,
+    ):
+        cfg = request.getfixturevalue(cfg_string)
+
+        builder = HFRLTrainerBuilder(cfg, model, tokenizer)
+        cfg["optimizer"] = "muon"
+
+        if cfg_string in ["dpo_cfg", "ipo_cfg", "grpo_cfg", "simpo_cfg"]:
+            cfg["datasets"] = [DictDefault(ALPACA_MESSAGES_CONFIG_REVISION)]
+        elif cfg_string == "kto_cfg":
+            cfg["datasets"] = [
+                DictDefault(
+                    {
+                        "path": "argilla/ultrafeedback-binarized-preferences-cleaned-kto",
+                        "type": "llama3.ultra",
+                        "split": "train[:32]",
+                    }
+                )
+            ]
+        elif cfg_string == "orpo_cfg":
+            cfg["datasets"] = [
+                DictDefault(
+                    {
+                        "path": "argilla/ultrafeedback-binarized-preferences-cleaned",
+                        "type": "chat_template.argilla",
+                        "split": "train[:32]",
+                    }
+                )
+            ]
+        else:
+            raise ValueError(f"Unhandled cfg_string: {cfg_string}")
+        cfg["dataset_num_proc"] = None
+
+        if cfg_string == "grpo_cfg":
+            rewards_dir = tmp_path / "rewards_test"
+            self._write_rewards_file(rewards_dir)
+
+            # Add the directory to Python path so we can import the module
+            sys.path.insert(0, str(rewards_dir))
+
+        try:
+            with (
+                patch(
+                    "axolotl.utils.data.rl.load_dataset_with_config"
+                ) as mock_load_dataset,
+                patch("axolotl.utils.data.rl.load_tokenizer", return_value=tokenizer),
+            ):
+                mock_load_dataset.return_value = _preference_dataset(cfg_string)
+                train_dataset, eval_dataset = prepare_preference_datasets(
+                    cfg, tokenizer
+                )
+
+            builder.train_dataset = train_dataset
+            builder.eval_dataset = eval_dataset
+
+            trainer = builder.build(100)
+
+            assert trainer.optimizer_cls_and_kwargs is not None
+
+            from axolotl.contribs.mit.muon import MuonOptimizerFactory
+            from axolotl.contribs.mit.muon.muon import Muon
+
+            optimizer_cls, optimizer_kwargs = trainer.optimizer_cls_and_kwargs
+            assert optimizer_cls is MuonOptimizerFactory
+            assert optimizer_kwargs["lr"] == 0.00005
+            assert optimizer_kwargs["weight_decay"] == 0.01
+            assert optimizer_kwargs["betas"] == (0.998, 0.9)
+            assert optimizer_kwargs["eps"] == 0.00001
+
+            # Ensure optimizer is created with correct class
+            optim = trainer.create_optimizer()
+            assert isinstance(optim, Muon)
+
+        finally:
+            # remove imported module from path
+            if cfg_string == "grpo_cfg" and str(rewards_dir) in sys.path:
+                sys.path.remove(str(rewards_dir))
diff --git a/tests/core/test_gc_steps_wiring.py b/tests/core/test_gc_steps_wiring.py
new file mode 100644
index 0000000000..487736c6af
--- /dev/null
+++ b/tests/core/test_gc_steps_wiring.py
@@ -0,0 +1,41 @@
+"""get_callbacks must register GCCallback whenever gc_collect_steps (or legacy gc_steps) is set."""
+
+from unittest.mock import MagicMock, patch
+
+from axolotl.core.builders import HFCausalTrainerBuilder
+from axolotl.utils.callbacks import GCCallback
+from axolotl.utils.dict import DictDefault
+
+
+def _get_callbacks(cfg_dict):
+    builder = HFCausalTrainerBuilder.__new__(HFCausalTrainerBuilder)
+    builder.cfg = DictDefault(cfg_dict)
+    builder.model = MagicMock()
+    with (
+        patch("axolotl.core.builders.base.PluginManager") as pm,
+        patch("axolotl.core.builders.base.TelemetryManager") as tm,
+    ):
+        pm.get_instance.return_value.add_callbacks_pre_trainer.return_value = []
+        tm.get_instance.return_value.enabled = False
+        return builder.get_callbacks()
+
+
+def _gc(callbacks):
+    return next((c for c in callbacks if isinstance(c, GCCallback)), None)
+
+
+def test_gc_collect_steps_alone_registers_callback():
+    """gc_collect_steps set without the deprecated gc_steps must still register GCCallback."""
+    gc = _gc(_get_callbacks({"gc_collect_steps": 5}))
+    assert gc is not None
+    assert gc.gc_collect_steps == 5
+
+
+def test_legacy_gc_steps_still_registers_callback():
+    gc = _gc(_get_callbacks({"gc_steps": 7}))
+    assert gc is not None
+    assert gc.gc_collect_steps == 7
+
+
+def test_no_gc_config_registers_nothing():
+    assert _gc(_get_callbacks({})) is None
diff --git a/tests/core/test_quantized_lora_checkpoint_resume.py b/tests/core/test_quantized_lora_checkpoint_resume.py
new file mode 100644
index 0000000000..47d2de04e7
--- /dev/null
+++ b/tests/core/test_quantized_lora_checkpoint_resume.py
@@ -0,0 +1,132 @@
+"""F4: FSDP2 + quantized-base LoRA checkpoints must stay resumable.
+
+The DCP sharded model save fails on the NVFP4/Float8 frozen base, so the trainer saves just the
+adapter — but it must STILL persist optimizer/scheduler/scaler/RNG and trainer_state so periodic
+checkpoints can resume. These bind ``AxolotlTrainer._save_checkpoint`` to a stub (no Trainer/GPU)
+and assert the resume artifacts are written after the adapter save.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import peft
+
+from axolotl.core.trainers.base import AxolotlTrainer
+
+
+def _stub(tmp_path, *, save_only_model=False, should_save=True, adapter_handled=True):
+    return SimpleNamespace(
+        state=SimpleNamespace(global_step=5, save_to_json=MagicMock()),
+        args=SimpleNamespace(save_only_model=save_only_model, should_save=should_save),
+        _get_output_dir=lambda trial=None: str(tmp_path),
+        _save_fsdp2_quantized_lora_adapter=MagicMock(return_value=adapter_handled),
+        _save_optimizer_and_scheduler=MagicMock(),
+        _save_scaler=MagicMock(),
+        _save_rng_state=MagicMock(),
+    )
+
+
+def test_quantized_lora_checkpoint_persists_resume_state(tmp_path):
+    stub = _stub(tmp_path)
+    out = AxolotlTrainer._save_checkpoint(stub, model=object(), trial=None)
+    assert out is None
+    stub._save_fsdp2_quantized_lora_adapter.assert_called_once()
+    # the F4 fix: optimizer/scheduler/scaler/RNG + trainer_state all written (resumable)
+    stub._save_optimizer_and_scheduler.assert_called_once()
+    stub._save_scaler.assert_called_once()
+    stub._save_rng_state.assert_called_once()
+    stub.state.save_to_json.assert_called_once()
+
+
+def test_save_only_model_skips_optimizer_but_writes_trainer_state(tmp_path):
+    stub = _stub(tmp_path, save_only_model=True)
+    AxolotlTrainer._save_checkpoint(stub, model=object(), trial=None)
+    stub._save_optimizer_and_scheduler.assert_not_called()
+    stub._save_rng_state.assert_not_called()
+    stub.state.save_to_json.assert_called_once()  # trainer_state still written for resume bookkeeping
+
+
+def test_resume_state_failure_keeps_adapter_and_does_not_raise(tmp_path):
+    # If an FSDP2 resume-artifact save raises, keep the (already-written) adapter rather than aborting.
+    stub = _stub(tmp_path)
+    stub._save_optimizer_and_scheduler = MagicMock(
+        side_effect=RuntimeError("dcp optim fail")
+    )
+    out = AxolotlTrainer._save_checkpoint(stub, model=object(), trial=None)
+    assert out is None  # did not raise
+    stub._save_fsdp2_quantized_lora_adapter.assert_called_once()
+
+
+def test_fsdp2_checkpoint_save_uses_axolotl_cfg_when_trainer_flag_unset():
+    stub = SimpleNamespace(
+        is_fsdp_enabled=False,
+        axolotl_cfg=SimpleNamespace(
+            fsdp_version=2,
+            fsdp_config={"state_dict_type": "SHARDED_STATE_DICT"},
+            fsdp=None,
+        ),
+    )
+    assert AxolotlTrainer._is_fsdp2_checkpoint_save_enabled(stub)
+
+
+def test_fsdp2_checkpoint_save_ignores_non_fsdp2_cfg():
+    stub = SimpleNamespace(
+        is_fsdp_enabled=False,
+        axolotl_cfg=SimpleNamespace(fsdp_version=1, fsdp_config={}, fsdp=None),
+    )
+    assert not AxolotlTrainer._is_fsdp2_checkpoint_save_enabled(stub)
+
+
+def test_fsdp2_quantized_param_detector_checks_dtensor_local_tensor():
+    NVFP4Tensor = type("NVFP4Tensor", (), {})
+    param = SimpleNamespace(_local_tensor=NVFP4Tensor())
+    assert AxolotlTrainer._is_fsdp2_quantized_param(param)
+
+
+def test_fsdp2_quantized_param_detector_checks_parameter_data():
+    MXTensor = type("MXTensor", (), {})
+    param = SimpleNamespace(data=MXTensor())
+    assert AxolotlTrainer._is_fsdp2_quantized_param(param)
+
+
+def test_quantized_lora_checkpoint_uses_ep_adapter_save(monkeypatch, tmp_path):
+    class NVFP4Tensor:
+        pass
+
+    class FakePeftModel:
+        def parameters(self):
+            return [SimpleNamespace(_local_tensor=NVFP4Tensor())]
+
+    model = FakePeftModel()
+    stub = SimpleNamespace(
+        is_fsdp_enabled=True,
+        axolotl_cfg=SimpleNamespace(expert_parallel_size=2),
+        accelerator=SimpleNamespace(unwrap_model=lambda wrapped: wrapped),
+        _is_fsdp2_quantized_param=AxolotlTrainer._is_fsdp2_quantized_param,
+    )
+    # _save_fsdp2_quantized_lora_adapter gates on these helpers; bind the real
+    # implementations so the test exercises actual enablement + quant detection.
+    stub._is_fsdp2_checkpoint_save_enabled = lambda: (
+        AxolotlTrainer._is_fsdp2_checkpoint_save_enabled(stub)
+    )
+
+    monkeypatch.setattr(peft, "PeftModel", FakePeftModel)
+
+    from axolotl.integrations.expert_parallel import shard
+    from axolotl.integrations.expert_parallel.plugin import ExpertParallelPlugin
+
+    resolve_ep_group = MagicMock(return_value=object())
+    save_ep_lora_adapter = MagicMock(return_value=True)
+    save_fsdp2_lora_adapter = MagicMock(return_value=True)
+    monkeypatch.setattr(ExpertParallelPlugin, "_resolve_ep_group", resolve_ep_group)
+    monkeypatch.setattr(shard, "save_ep_lora_adapter", save_ep_lora_adapter)
+    monkeypatch.setattr(shard, "save_fsdp2_lora_adapter", save_fsdp2_lora_adapter)
+
+    handled = AxolotlTrainer._save_fsdp2_quantized_lora_adapter(
+        stub, model, str(tmp_path)
+    )
+
+    assert handled is True
+    resolve_ep_group.assert_called_once_with(stub.axolotl_cfg)
+    save_ep_lora_adapter.assert_called_once()
+    save_fsdp2_lora_adapter.assert_not_called()
diff --git a/tests/core/test_trainer_builder.py b/tests/core/test_trainer_builder.py
deleted file mode 100644
index 82455922ef..0000000000
--- a/tests/core/test_trainer_builder.py
+++ /dev/null
@@ -1,66 +0,0 @@
-"""
-unit tests for axolotl.core.trainer_builder
-"""
-
-import pytest
-
-from axolotl.core.trainer_builder import HFRLTrainerBuilder
-from axolotl.utils.config import normalize_config
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model, load_tokenizer
-
-
-@pytest.fixture(name="cfg")
-def fixture_cfg():
-    cfg = DictDefault(
-        {
-            "base_model": "TinyLlama/TinyLlama-1.1B-Chat-v0.6",
-            "model_type": "AutoModelForCausalLM",
-            "tokenizer_type": "LlamaTokenizer",
-            "micro_batch_size": 1,
-            "gradient_accumulation_steps": 1,
-            "learning_rate": 0.00005,
-            "save_steps": 100,
-            "output_dir": "./model-out",
-            "warmup_steps": 10,
-            "gradient_checkpointing": False,
-            "optimizer": "adamw_torch",
-            "sequence_len": 2048,
-            "rl": True,
-            "adam_beta1": 0.998,
-            "adam_beta2": 0.9,
-            "adam_epsilon": 0.00001,
-            "dataloader_num_workers": 1,
-            "dataloader_pin_memory": True,
-            "model_config_type": "llama",
-        }
-    )
-
-    normalize_config(cfg)
-
-    return cfg
-
-
-@pytest.fixture(name="tokenizer")
-def fixture_tokenizer(cfg):
-    return load_tokenizer(cfg)
-
-
-@pytest.fixture(name="model")
-def fixture_model(cfg, tokenizer):
-    return load_model(cfg, tokenizer)
-
-
-class TestHFRLTrainerBuilder:
-    """
-    TestCase class for DPO trainer builder
-    """
-
-    def test_build_training_arguments(self, cfg, model, tokenizer):
-        builder = HFRLTrainerBuilder(cfg, model, tokenizer)
-        training_arguments = builder.build_training_arguments(100)
-        assert training_arguments.adam_beta1 == 0.998
-        assert training_arguments.adam_beta2 == 0.9
-        assert training_arguments.adam_epsilon == 0.00001
-        assert training_arguments.dataloader_num_workers == 1
-        assert training_arguments.dataloader_pin_memory is True
diff --git a/tests/e2e/integrations/test_cut_cross_entropy.py b/tests/e2e/integrations/test_cut_cross_entropy.py
new file mode 100644
index 0000000000..75fb9d0dbd
--- /dev/null
+++ b/tests/e2e/integrations/test_cut_cross_entropy.py
@@ -0,0 +1,167 @@
+"""
+Simple end-to-end test for Cut Cross Entropy integration
+"""
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils import get_pytorch_version
+from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+)
+
+
+@pytest.fixture()
+def min_cfg(temp_dir):
+    return {
+        "base_model": "HuggingFaceTB/SmolLM2-135M",
+        "plugins": [
+            "axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin",
+        ],
+        "cut_cross_entropy": True,
+        "sequence_len": 1024,
+        "val_set_size": 0.02,
+        "special_tokens": {
+            "pad_token": "<|endoftext|>",
+        },
+        "datasets": [
+            {
+                "path": "mhenrichsen/alpaca_2k_test",
+                "type": "alpaca",
+            },
+        ],
+        "num_epochs": 1,
+        "micro_batch_size": 8,
+        "gradient_accumulation_steps": 1,
+        "learning_rate": 5e-4,
+        "optimizer": "adamw_torch_fused",
+        "output_dir": temp_dir,
+        "lr_scheduler": "cosine",
+        "max_steps": 40,
+        "warmup_steps": 5,
+        "bf16": "auto",
+        "save_first_step": False,
+        "use_tensorboard": True,
+        "seed": 42,
+    }
+
+
+class TestCutCrossEntropyIntegration:
+    """
+    e2e tests for cut_cross_entropy integration with Axolotl
+    """
+
+    def test_llama_w_cce(self, min_cfg, temp_dir):
+        cfg = DictDefault(min_cfg)
+        cfg = validate_config(cfg)
+        prepare_plugins(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        major, minor, _ = get_pytorch_version()
+        if (major, minor) < (2, 4):
+            with pytest.raises(ImportError):
+                train(cfg=cfg, dataset_meta=dataset_meta)
+        else:
+            train(cfg=cfg, dataset_meta=dataset_meta)
+            check_model_output_exists(temp_dir, cfg)
+            check_tensorboard_loss_decreased(
+                temp_dir + "/runs",
+                initial_window=5,
+                final_window=5,
+                max_initial=2.2,
+                max_final=2.0,
+            )
+
+    def test_qwen2_w_cce(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "plugins": [
+                    "axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin",
+                ],
+                "cut_cross_entropy": True,
+                "sequence_len": 1024,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
+                "output_dir": temp_dir,
+                "lr_scheduler": "cosine",
+                "max_steps": 50,
+                "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
+            }
+        )
+        cfg = validate_config(cfg)
+        prepare_plugins(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        major, minor, _ = get_pytorch_version()
+        if (major, minor) < (2, 4):
+            with pytest.raises(ImportError):
+                train(cfg=cfg, dataset_meta=dataset_meta)
+        else:
+            train(cfg=cfg, dataset_meta=dataset_meta)
+            check_model_output_exists(temp_dir, cfg)
+            check_tensorboard_loss_decreased(
+                temp_dir + "/runs",
+                initial_window=5,
+                final_window=5,
+                max_initial=5.0,
+                max_final=4.7,
+            )
+
+    @pytest.mark.parametrize(
+        "attention_type",
+        [
+            "flash_attention",
+            "sdp_attention",
+            # "xformers_attention",
+        ],
+    )
+    def test_llama_w_cce_and_attention(self, min_cfg, temp_dir, attention_type):
+        cfg = DictDefault(
+            min_cfg
+            | {
+                attention_type: True,
+            }
+        )
+        cfg = validate_config(cfg)
+        prepare_plugins(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        major, minor, _ = get_pytorch_version()
+        if (major, minor) < (2, 4):
+            with pytest.raises(ImportError):
+                train(cfg=cfg, dataset_meta=dataset_meta)
+        else:
+            train(cfg=cfg, dataset_meta=dataset_meta)
+            check_model_output_exists(temp_dir, cfg)
+            check_tensorboard_loss_decreased(
+                temp_dir + "/runs",
+                initial_window=5,
+                final_window=5,
+                max_initial=2.2,
+                max_final=2.0,
+            )
diff --git a/tests/e2e/integrations/test_fp8.py b/tests/e2e/integrations/test_fp8.py
new file mode 100644
index 0000000000..b708e8806e
--- /dev/null
+++ b/tests/e2e/integrations/test_fp8.py
@@ -0,0 +1,60 @@
+"""
+Simple end-to-end smoke tests for FP8 mixed precision training
+"""
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_model_output_exists, require_torch_2_7_0
+
+
+class FP8IntegrationTestCase:
+    """
+    e2e smoke tests for FP8 mixed precision training with Axolotl
+    """
+
+    @require_torch_2_7_0
+    def test_fp8_single_gpu_smoke(self, temp_dir):
+        """Smoke test for single GPU FP8 + torch.compile training"""
+
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "trust_remote_code": True,
+                "sequence_len": 512,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,  # Very short smoke test
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "sdp_attention": True,
+                "pad_to_seq_len": True,
+                "sample_packing": True,
+                "fp8": True,
+                "torch_compile": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/integrations/test_hooks.py b/tests/e2e/integrations/test_hooks.py
new file mode 100644
index 0000000000..e056d14912
--- /dev/null
+++ b/tests/e2e/integrations/test_hooks.py
@@ -0,0 +1,185 @@
+"""
+e2e tests to make sure all the hooks are fired on the plugin
+"""
+
+import os
+from pathlib import Path
+
+from axolotl.common.datasets import load_datasets
+from axolotl.integrations.base import BasePlugin
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_model_output_exists
+
+
+class LogHooksPlugin(BasePlugin):
+    """
+    fixture to capture in a log file each hook that was fired
+    """
+
+    base_dir = Path("/tmp/axolotl-log-hooks")
+
+    def __init__(self):
+        self.base_dir.mkdir(parents=True, exist_ok=True)
+        try:
+            os.remove(self.base_dir.joinpath("plugin_hooks.log"))
+        except FileNotFoundError:
+            pass
+
+    def post_trainer_create(self, cfg, trainer):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("post_trainer_create\n")
+
+    def pre_model_load(self, cfg):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("pre_model_load\n")
+
+    def post_model_build(self, cfg, model):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("post_model_build\n")
+
+    def pre_lora_load(self, cfg, model):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("pre_lora_load\n")
+
+    def post_lora_load(self, cfg, model):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("post_lora_load\n")
+
+    def post_model_load(self, cfg, model):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("post_model_load\n")
+
+    def create_optimizer(self, cfg, trainer):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("create_optimizer\n")
+
+    def get_trainer_cls(self, cfg):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("get_trainer_cls\n")
+
+    def create_lr_scheduler(self, cfg, trainer, optimizer, num_training_steps):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("create_lr_scheduler\n")
+
+    def add_callbacks_pre_trainer(self, cfg, model):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("add_callbacks_pre_trainer\n")
+        return []
+
+    def add_callbacks_post_trainer(self, cfg, trainer):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("add_callbacks_post_trainer\n")
+        return []
+
+    def post_train(self, cfg, model):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("post_train\n")
+
+    def post_train_unload(self, cfg):
+        with open(
+            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
+        ) as f:
+            f.write("post_train_unload\n")
+
+
+class TestPluginHooks:
+    """
+    e2e tests to make sure all the hooks are fired during the training
+    """
+
+    def test_plugin_hooks(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "plugins": [
+                    "tests.e2e.integrations.test_hooks.LogHooksPlugin",
+                ],
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "flash_attention": True,
+                "bf16": "auto",
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        prepare_plugins(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        with open(
+            "/tmp/axolotl-log-hooks" + "/plugin_hooks.log", "r", encoding="utf-8"
+        ) as f:
+            file_contents = f.readlines()
+            file_contents = "\n".join(file_contents)
+            assert "post_trainer_create" in file_contents
+            assert "pre_model_load" in file_contents
+            assert "post_model_build" in file_contents
+            assert "pre_lora_load" in file_contents
+            assert "post_lora_load" in file_contents
+            assert "post_model_load" in file_contents
+            # assert "create_optimizer" in file_contents  # not implemented yet
+            assert "get_trainer_cls" in file_contents
+            assert "create_lr_scheduler" in file_contents
+            assert "add_callbacks_pre_trainer" in file_contents
+            assert "add_callbacks_post_trainer" in file_contents
+            assert "post_train" in file_contents
+            # assert "post_train_unload" in file_contents  # not called from test train call
+
+        try:
+            os.remove("/tmp/axolotl-log-hooks" + "/plugin_hooks.log")
+        except FileNotFoundError:
+            pass
diff --git a/tests/e2e/integrations/test_kd.py b/tests/e2e/integrations/test_kd.py
new file mode 100644
index 0000000000..9e0e9406b0
--- /dev/null
+++ b/tests/e2e/integrations/test_kd.py
@@ -0,0 +1,148 @@
+"""
+e2e tests for kd trainer support in Axolotl
+"""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async, get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_tensorboard, require_torch_2_5_1
+
+
+@pytest.fixture(name="kd_min_cfg")
+def min_cfg(temp_dir):
+    return {
+        "base_model": "Qwen/Qwen3-0.6B",
+        "tokenizer_config": "winglian/qwen3-14b-math",
+        "plugins": [
+            "axolotl.integrations.kd.KDPlugin",
+            "axolotl.integrations.liger.LigerPlugin",
+        ],
+        "liger_rms_norm": True,
+        "liger_glu_activation": True,
+        "torch_compile": True,
+        "chat_template": "qwen3",
+        "kd_trainer": True,
+        "kd_ce_alpha": 0.1,
+        "kd_alpha": 0.9,
+        "kd_temperature": 1.0,
+        "kd_beta": 0.0,
+        "kd_normalize_topk": True,
+        "dataloader_prefetch_factor": 8,
+        "dataloader_num_workers": 4,
+        "dataloader_pin_memory": True,
+        "datasets": [
+            {
+                "path": "winglian/OpenThoughts-114k-math-correct-qwen3-14b-math-prepared-topk128-normalized",
+                "type": "chat_template",
+                "split": "train",
+                "split_thinking": True,
+                "eot_tokens": ["<|im_end|>"],
+                "data_files": ["train/batch-000000.parquet"],
+            },
+        ],
+        "skip_prepare_dataset": True,
+        "val_set_size": 0.0,
+        "sequence_len": 2048,
+        "sample_packing": True,
+        "pad_to_sequence_len": True,
+        "gradient_accumulation_steps": 2,
+        "micro_batch_size": 1,
+        "num_epochs": 1,
+        "optimizer": "adamw_8bit",
+        "lr_scheduler": "cosine",
+        "learning_rate": 0.00001,
+        "bf16": "auto",
+        "gradient_checkpointing": True,
+        "flash_attention": True,
+        "special_tokens": {
+            "pad_token": "<|end_of_text|>",
+            "eos_token": "<|eot_id|>",
+        },
+        "max_steps": 5,
+        "output_dir": temp_dir,
+        "use_tensorboard": True,
+        "save_first_step": False,
+    }
+
+
+class TestKnowledgeDistillation:
+    """
+    Test case for Knowledge Distillation
+    """
+
+    # While this will run on torch 2.4.x without torch_compile enabled
+    # the VRAM requirement is higher than what is available in CI
+    @require_torch_2_5_1
+    def test_llama_kd(self, temp_dir, kd_min_cfg):
+        cfg = DictDefault(kd_min_cfg)
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "1",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        assert (Path(temp_dir) / "model.safetensors").exists()
+        check_tensorboard(
+            temp_dir + "/runs", "train/loss", 1.4, "Train Loss (%s) is too high"
+        )
+
+    @pytest.mark.parametrize(
+        "load_in_8bit",
+        [True, False],
+    )
+    def test_llama_lora_kd(self, temp_dir, kd_min_cfg, load_in_8bit):
+        cfg = DictDefault(
+            {
+                "load_in_8bit": load_in_8bit,
+                "torch_compile": False,
+                "adapter": "lora",
+                "peft_use_dora": True,
+                "lora_target_linear": True,
+                "lora_r": 16,
+                "lora_alpha": 32,
+                "lora_dropout": 0.0,
+                "lora_modules_to_save": ["embed_tokens", "lm_head"],
+                "lora_mlp_kernel": False,
+                "lora_qkv_kernel": False,
+                "lora_o_kernel": False,
+            }
+            | kd_min_cfg
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "1",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+        check_tensorboard(
+            temp_dir + "/runs", "train/loss", 1.2, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/integrations/test_liger.py b/tests/e2e/integrations/test_liger.py
new file mode 100644
index 0000000000..ba19cf41c2
--- /dev/null
+++ b/tests/e2e/integrations/test_liger.py
@@ -0,0 +1,113 @@
+"""
+Simple end-to-end test for Liger integration
+"""
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_model_output_exists, require_torch_2_4_1
+
+
+class LigerIntegrationTestCase:
+    """
+    e2e tests for liger integration with Axolotl
+    """
+
+    @require_torch_2_4_1
+    def test_llama_wo_flce(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "plugins": [
+                    "axolotl.integrations.liger.LigerPlugin",
+                ],
+                "liger_rope": True,
+                "liger_rms_norm": True,
+                "liger_glu_activation": True,
+                "liger_cross_entropy": True,
+                "liger_fused_linear_cross_entropy": False,
+                "sequence_len": 1024,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "max_steps": 5,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        prepare_plugins(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    @require_torch_2_4_1
+    @pytest.mark.parametrize(
+        "liger_use_token_scaling",
+        [True, False],
+    )
+    def test_llama_w_flce(self, temp_dir, liger_use_token_scaling):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "plugins": [
+                    "axolotl.integrations.liger.LigerPlugin",
+                ],
+                "liger_rope": True,
+                "liger_rms_norm": True,
+                "liger_glu_activation": True,
+                "liger_cross_entropy": False,
+                "liger_fused_linear_cross_entropy": True,
+                "liger_use_token_scaling": liger_use_token_scaling,
+                "sequence_len": 1024,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "max_steps": 5,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        prepare_plugins(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/integrations/test_llm_compressor.py b/tests/e2e/integrations/test_llm_compressor.py
new file mode 100644
index 0000000000..5804bca100
--- /dev/null
+++ b/tests/e2e/integrations/test_llm_compressor.py
@@ -0,0 +1,109 @@
+"""
+E2E smoke tests for LLMCompressorPlugin integration
+"""
+
+from pathlib import Path
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import (
+    check_model_output_exists,
+    require_llmcompressor,
+    require_torch_2_4_1,
+)
+
+MODELS = [
+    "nm-testing/llama2.c-stories42M-pruned2.4-compressed",
+    "nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed",
+]
+
+
+@pytest.mark.parametrize(
+    "base_model", MODELS, ids=["no-checkpoint-recipe", "with-checkpoint-recipe"]
+)
+@pytest.mark.parametrize(
+    "save_compressed", [True, False], ids=["save_compressed", "save_uncompressed"]
+)
+class TestLLMCompressorIntegration:
+    """
+    e2e tests for axolotl.integrations.llm_compressor.LLMCompressorPlugin
+    """
+
+    @require_llmcompressor
+    @require_torch_2_4_1
+    def test_llmcompressor_plugin(
+        self, temp_dir, base_model: str, save_compressed: bool
+    ):
+        from llmcompressor import active_session
+
+        # core cfg
+        cfg = DictDefault(
+            {
+                "base_model": base_model,
+                "plugins": ["axolotl.integrations.llm_compressor.LLMCompressorPlugin"],
+                "sequence_len": 1024,
+                "val_set_size": 0.05,
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "datasets": [{"path": "mhenrichsen/alpaca_2k_test", "type": "alpaca"}],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-5,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "max_steps": 5,
+                "llmcompressor": {
+                    "recipe": {
+                        "finetuning_stage": {
+                            "finetuning_modifiers": {
+                                "ConstantPruningModifier": {
+                                    "targets": [
+                                        "re:.*q_proj.weight",
+                                        "re:.*k_proj.weight",
+                                        "re:.*v_proj.weight",
+                                        "re:.*o_proj.weight",
+                                        "re:.*gate_proj.weight",
+                                        "re:.*up_proj.weight",
+                                        "re:.*down_proj.weight",
+                                    ],
+                                    "start": 0,
+                                },
+                            },
+                        },
+                    },
+                    "save_compressed": save_compressed,
+                },
+                "save_first_step": False,
+            }
+        )
+
+        prepare_plugins(cfg)
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        try:
+            train(cfg=cfg, dataset_meta=dataset_meta)
+            check_model_output_exists(temp_dir, cfg)
+            _check_llmcompressor_model_outputs(temp_dir, save_compressed)
+        finally:
+            active_session().reset()
+
+
+def _check_llmcompressor_model_outputs(temp_dir, save_compressed):
+    if save_compressed:
+        assert (Path(temp_dir) / "recipe.yaml").exists()
+
+        from compressed_tensors import ModelCompressor
+        from compressed_tensors.config import Sparse24BitMaskConfig
+
+        compressor = ModelCompressor.from_pretrained(temp_dir)
+        assert compressor is not None
+        assert isinstance(compressor.sparsity_config, Sparse24BitMaskConfig)
diff --git a/tests/e2e/integrations/test_scattermoe_lora_kernels.py b/tests/e2e/integrations/test_scattermoe_lora_kernels.py
new file mode 100644
index 0000000000..c204a15035
--- /dev/null
+++ b/tests/e2e/integrations/test_scattermoe_lora_kernels.py
@@ -0,0 +1,1860 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Tests for ScatterMoE + LoRA Fused Kernels
+==========================================
+
+Tests verify correctness of:
+1. Forward pass: fused kernel matches naive PyTorch reference
+2. Backward pass: gradients for LoRA A, B, and input match reference
+3. Frozen weights: expert weight gradients are correctly skipped
+4. Various configurations: top-k, grouped_in/out, with/without bias
+5. Numerical stability: bf16/fp16 outputs within tolerance of fp32 reference
+6. HFScatterMoEGatedMLP with sigmoid routing (GLM/DeepSeek/MiniMax M2)
+
+Test strategy:
+- Reference implementation uses pure PyTorch ops (no Triton)
+- ScatterMoE routing (flatten_sort_count) is shared between reference and kernel
+- Tolerances account for tf32 accumulation in Triton kernels
+"""
+
+from functools import wraps
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+# Skip all tests if CUDA is not available
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="CUDA required for Triton kernels",
+)
+
+_SMOE = "axolotl.integrations.kernels.libs.scattermoe_lora"
+
+
+def skip_on_out_of_resources(func):
+    """Skip test if Triton kernel exceeds GPU shared memory limits."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as exc:  # pylint: disable=broad-except
+            if "OutOfResources" in type(exc).__name__:
+                pytest.skip(f"GPU shared memory too small: {exc}")
+            raise
+
+    return wrapper
+
+
+# =============================================================================
+# Helpers
+# =============================================================================
+
+
+def flatten_sort_count_ref(expert_idxs: torch.Tensor, num_experts: int):
+    """Reference implementation of routing."""
+    with torch.no_grad():
+        flat = expert_idxs.flatten()
+        sorted_expert_idxs, sorted_scattered_idxs = torch.sort(flat)
+        counts = flat.bincount(minlength=num_experts)
+        offsets = counts.cumsum(-1)
+    return sorted_expert_idxs, sorted_scattered_idxs, offsets
+
+
+def reference_parallel_linear_lora(
+    X,
+    W,
+    k,
+    sorted_expert_idxs,
+    sorted_scattered_idxs,
+    lora_A,
+    lora_B,
+    scaling,
+    x_grouped=False,
+    y_grouped=False,
+    bias=None,
+):
+    """
+    Pure PyTorch reference for: Y[i] = X[i] @ W[e] + scaling * (X[i] @ A[e]^T) @ B[e]^T + b[e]
+
+    Args:
+        X: [M, K] input (token order)
+        W: [E, K, N] expert weights
+        sorted_expert_idxs: [M*k] expert assignments (sorted)
+        sorted_scattered_idxs: [M*k] original token indices (sorted)
+        lora_A: [r*E, K] LoRA A weights
+        lora_B: [N, r*E] LoRA B weights
+        scaling: LoRA scaling factor
+    """
+    E, K, N = W.shape
+    R = lora_A.size(0) // E
+    L = sorted_expert_idxs.size(0)  # M * k
+
+    output = torch.zeros(L, N, device=X.device, dtype=X.dtype)
+
+    for i in range(L):
+        e = sorted_expert_idxs[i].item()
+        if x_grouped:
+            x_i = X[i]
+        else:
+            token_idx = sorted_scattered_idxs[i].item() // k
+            x_i = X[token_idx]
+
+        w_e = W[e]  # [K, N]
+        a_e = lora_A[e * R : (e + 1) * R, :]  # [r, K]
+        b_e = lora_B[:, e * R : (e + 1) * R]  # [N, r]
+
+        # Y = X @ W + scaling * (X @ A^T) @ B^T
+        base = x_i @ w_e  # [N]
+        lora = scaling * ((x_i @ a_e.T) @ b_e.T)  # [N]
+        out_i = base + lora
+
+        if bias is not None:
+            out_i = out_i + bias[e]
+
+        if y_grouped:
+            output[i] = out_i
+        else:
+            output[sorted_scattered_idxs[i]] = out_i
+
+    return output
+
+
+def reference_lora_backward(
+    grad_out,
+    X,
+    W,
+    lora_A,
+    lora_B,
+    scaling,
+    sorted_expert_idxs,
+    sorted_scattered_idxs,
+    expert_offsets,
+    k,
+    E,
+):
+    """
+    Pure PyTorch reference for LoRA backward pass on grouped data.
+
+    Returns:
+        dX: [M*k, K] input gradient (in grouped order)
+        dA: [r*E, K] LoRA A gradient
+        dB: [N, r*E] LoRA B gradient
+    """
+    R = lora_A.size(0) // E
+
+    dA = torch.zeros_like(lora_A)
+    dB = torch.zeros_like(lora_B)
+    dX = torch.zeros_like(X)
+
+    prev_offset = 0
+    for e in range(E):
+        curr_offset = expert_offsets[e].item()
+        if curr_offset > prev_offset:
+            dy_e = grad_out[prev_offset:curr_offset]  # [M_e, N]
+            x_e = X[prev_offset:curr_offset]  # [M_e, K]
+            a_e = lora_A[e * R : (e + 1) * R, :]  # [r, K]
+            b_e = lora_B[:, e * R : (e + 1) * R]  # [N, r]
+            w_e = W[e]  # [K, N]
+
+            # Input gradient: dX = dY @ W^T + scaling * (dY @ B) @ A
+            dx_base = dy_e @ w_e.T  # [M_e, K]
+            dy_b = dy_e @ b_e  # [M_e, r]
+            dx_lora = scaling * (dy_b @ a_e)  # [M_e, K]
+            dX[prev_offset:curr_offset] = dx_base + dx_lora
+
+            # LoRA A gradient: dA = scaling * (dY @ B)^T @ X
+            xa = x_e @ a_e.T  # [M_e, r]
+            dA[e * R : (e + 1) * R, :] = scaling * (dy_b.T @ x_e)
+
+            # LoRA B gradient: dB = scaling * dY^T @ (X @ A^T)
+            dB[:, e * R : (e + 1) * R] = scaling * (dy_e.T @ xa)
+
+        prev_offset = curr_offset
+
+    return dX, dA, dB
+
+
+def make_test_data(
+    M=32,
+    K=64,
+    N=128,
+    E=4,
+    R=8,
+    k=2,
+    dtype=torch.float32,
+    device="cuda",
+    seed=42,
+):
+    """Create test data for ScatterMoE + LoRA tests."""
+    torch.manual_seed(seed)
+
+    X = torch.randn(M, K, device=device, dtype=dtype)
+    W = torch.randn(E, K, N, device=device, dtype=dtype) * 0.02
+    lora_A = torch.randn(R * E, K, device=device, dtype=dtype) * 0.01
+    lora_B = torch.randn(N, R * E, device=device, dtype=dtype) * 0.01
+    scaling = 0.5
+
+    # Generate routing
+    selected_experts = torch.randint(0, E, (M, k), device=device)
+    sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count_ref(
+        selected_experts, E
+    )
+
+    return {
+        "X": X,
+        "W": W,
+        "lora_A": lora_A,
+        "lora_B": lora_B,
+        "scaling": scaling,
+        "k": k,
+        "E": E,
+        "R": R,
+        "sorted_expert_idxs": sorted_expert_idxs,
+        "sorted_scattered_idxs": sorted_scattered_idxs,
+        "expert_offsets": expert_offsets,
+    }
+
+
+# =============================================================================
+# Test: Forward Pass Correctness
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestForwardPass:
+    """Test forward pass of fused scatter2scatter_lora kernel."""
+
+    def _run_forward_test(
+        self, M, K, N, E, R, k, dtype=torch.float32, atol=1e-2, rtol=1e-2
+    ):
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k, dtype=dtype)
+
+        # Reference
+        ref_output = reference_parallel_linear_lora(
+            data["X"],
+            data["W"],
+            data["k"],
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["lora_A"],
+            data["lora_B"],
+            data["scaling"],
+        )
+
+        # Kernel
+        kernel_output = lora_ops.scatter2scatter_lora(
+            X=data["X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=data["k"],
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            scaling=data["scaling"],
+        )
+
+        torch.testing.assert_close(kernel_output, ref_output, atol=atol, rtol=rtol)
+
+    def test_basic(self):
+        """Basic forward pass with small dimensions."""
+        self._run_forward_test(M=16, K=64, N=64, E=4, R=8, k=1)
+
+    def test_topk2(self):
+        """Forward pass with top-2 routing."""
+        self._run_forward_test(M=32, K=64, N=128, E=4, R=8, k=2)
+
+    def test_larger_rank(self):
+        """Forward pass with larger LoRA rank."""
+        self._run_forward_test(M=16, K=128, N=128, E=8, R=32, k=2)
+
+    def test_small_rank(self):
+        """Forward pass with very small LoRA rank."""
+        self._run_forward_test(M=32, K=64, N=64, E=4, R=4, k=1)
+
+    def test_many_experts(self):
+        """Forward with many experts, fewer tokens per expert."""
+        self._run_forward_test(M=64, K=64, N=64, E=16, R=8, k=2)
+
+    def test_non_power_of_2_dims(self):
+        """Test with dimensions that are not powers of 2."""
+        self._run_forward_test(M=17, K=96, N=80, E=6, R=16, k=2, atol=2e-2, rtol=2e-2)
+
+    def test_single_token(self):
+        """Test with a single token."""
+        self._run_forward_test(M=1, K=64, N=64, E=4, R=8, k=1)
+
+    def test_bf16(self):
+        """Test with bfloat16 precision."""
+        self._run_forward_test(
+            M=32, K=64, N=128, E=4, R=8, k=2, dtype=torch.bfloat16, atol=5e-2, rtol=5e-2
+        )
+
+    def test_fp16(self):
+        """Test with float16 precision."""
+        self._run_forward_test(
+            M=32, K=64, N=128, E=4, R=8, k=2, dtype=torch.float16, atol=5e-2, rtol=5e-2
+        )
+
+
+@pytest.mark.slow
+class TestForwardGrouped:
+    """Test forward pass with grouped_in/grouped_out configurations."""
+
+    def _make_grouped_data(self, M=32, K=64, N=128, E=4, R=8, k=2, dtype=torch.float32):
+        from importlib import import_module
+
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k, dtype=dtype)
+
+        # Create grouped X
+        grouped_X = base_ops.group(data["X"], data["sorted_scattered_idxs"], fan_out=k)
+        data["grouped_X"] = grouped_X
+        return data
+
+    def test_x_grouped(self):
+        """Forward with pre-grouped input."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+
+        data = self._make_grouped_data()
+
+        ref_output = reference_parallel_linear_lora(
+            data["grouped_X"],
+            data["W"],
+            data["k"],
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["lora_A"],
+            data["lora_B"],
+            data["scaling"],
+            x_grouped=True,
+        )
+
+        kernel_output = lora_ops.scatter2scatter_lora(
+            X=data["grouped_X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=1,  # When x_grouped, fan_out=1 (already expanded)
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            scaling=data["scaling"],
+            x_grouped=True,
+        )
+
+        torch.testing.assert_close(kernel_output, ref_output, atol=1e-2, rtol=1e-2)
+
+    def test_y_grouped(self):
+        """Forward with grouped output."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+
+        data = make_test_data()
+
+        ref_output = reference_parallel_linear_lora(
+            data["X"],
+            data["W"],
+            data["k"],
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["lora_A"],
+            data["lora_B"],
+            data["scaling"],
+            y_grouped=True,
+        )
+
+        kernel_output = lora_ops.scatter2scatter_lora(
+            X=data["X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=data["k"],
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            scaling=data["scaling"],
+            y_grouped=True,
+        )
+
+        torch.testing.assert_close(kernel_output, ref_output, atol=1e-2, rtol=1e-2)
+
+
+# =============================================================================
+# Test: Backward Pass Correctness (LoRA Gradients)
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestLoRAGradients:
+    """Test backward LoRA gradient computation (dA, dB)."""
+
+    def _run_lora_grad_test(self, M, K, N, E, R, k, atol=1e-2, rtol=1e-2):
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        # Group X for backward
+        grouped_X = base_ops.group(data["X"], data["sorted_scattered_idxs"], fan_out=k)
+
+        # Create fake grad_out in grouped order
+        grad_out = torch.randn(
+            data["sorted_expert_idxs"].size(0),
+            N,
+            device="cuda",
+            dtype=torch.float32,
+        )
+
+        # Reference
+        _, ref_dA, ref_dB = reference_lora_backward(
+            grad_out,
+            grouped_X,
+            data["W"],
+            data["lora_A"],
+            data["lora_B"],
+            data["scaling"],
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            k,
+            E,
+        )
+
+        # Kernel
+        kernel_dA, kernel_dB = lora_ops.group_bwd_lora(
+            DY=grad_out,
+            X=grouped_X,
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            expert_offsets=data["expert_offsets"],
+            E=E,
+            scaling=data["scaling"],
+        )
+
+        torch.testing.assert_close(kernel_dA, ref_dA, atol=atol, rtol=rtol)
+        torch.testing.assert_close(kernel_dB, ref_dB, atol=atol, rtol=rtol)
+
+    def test_basic_lora_grads(self):
+        self._run_lora_grad_test(M=32, K=64, N=128, E=4, R=8, k=2)
+
+    def test_small_rank(self):
+        self._run_lora_grad_test(M=16, K=64, N=64, E=4, R=4, k=1)
+
+    def test_larger_rank(self):
+        self._run_lora_grad_test(
+            M=16, K=128, N=128, E=8, R=32, k=2, atol=5e-2, rtol=5e-2
+        )
+
+    def test_many_experts(self):
+        self._run_lora_grad_test(M=64, K=64, N=64, E=16, R=8, k=2)
+
+    def test_single_token_per_expert(self):
+        """Edge case: roughly 1 token per expert."""
+        self._run_lora_grad_test(M=8, K=64, N=64, E=8, R=4, k=1)
+
+
+# =============================================================================
+# Test: Full Autograd (Forward + Backward) via torch.autograd
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestAutograd:
+    """Test full autograd integration through ScatterMoELoRA."""
+
+    def test_lora_receives_gradients(self):
+        """LoRA A and B receive non-zero gradients; frozen W does not."""
+        from importlib import import_module
+
+        pll = import_module(f"{_SMOE}.parallel_linear_lora")
+
+        M, K, N, E, R, k = 16, 64, 64, 4, 8, 2
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        X = data["X"].clone().requires_grad_(True)
+        W = data["W"].clone().requires_grad_(False)  # Frozen
+        lora_A = data["lora_A"].clone().requires_grad_(True)
+        lora_B = data["lora_B"].clone().requires_grad_(True)
+
+        output = pll.ScatterMoELoRA.apply(
+            X,
+            W,
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            lora_A,
+            lora_B,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+        )
+
+        loss = output.sum()
+        loss.backward()
+
+        # LoRA params should have gradients
+        assert lora_A.grad is not None, "lora_A should have gradient"
+        assert lora_B.grad is not None, "lora_B should have gradient"
+        assert lora_A.grad.abs().sum() > 0, "lora_A gradient should be non-zero"
+        assert lora_B.grad.abs().sum() > 0, "lora_B gradient should be non-zero"
+
+        # Input should have gradient (needed for upstream backprop)
+        assert X.grad is not None, "X should have gradient"
+        assert X.grad.abs().sum() > 0, "X gradient should be non-zero"
+
+    def test_input_gradient_matches_reference(self):
+        """Input gradient from autograd matches pure PyTorch reference."""
+        from importlib import import_module
+
+        pll = import_module(f"{_SMOE}.parallel_linear_lora")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        M, K, N, E, R, k = 16, 64, 64, 4, 8, 1
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        # Autograd path
+        X_kern = data["X"].clone().requires_grad_(True)
+        lora_A_kern = data["lora_A"].clone().requires_grad_(True)
+        lora_B_kern = data["lora_B"].clone().requires_grad_(True)
+
+        out_kern = pll.ScatterMoELoRA.apply(
+            X_kern,
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            lora_A_kern,
+            lora_B_kern,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+        )
+        grad_out = torch.randn_like(out_kern)
+        out_kern.backward(grad_out)
+
+        # Reference path
+        grouped_X = base_ops.group(data["X"], data["sorted_scattered_idxs"], fan_out=k)
+        grouped_grad = base_ops.group(
+            grad_out, data["sorted_scattered_idxs"], fan_out=1
+        )
+
+        ref_dX, ref_dA, ref_dB = reference_lora_backward(
+            grouped_grad,
+            grouped_X,
+            data["W"],
+            data["lora_A"],
+            data["lora_B"],
+            data["scaling"],
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            k,
+            E,
+        )
+
+        # Compare input gradient (for k=1, no reduction needed)
+        # ref_dX is in grouped (expert-sorted) order; X_kern.grad is in original order.
+        # Ungroup ref_dX by scattering back to original positions.
+        ref_dX_ungrouped = torch.zeros_like(ref_dX)
+        ref_dX_ungrouped[data["sorted_scattered_idxs"]] = ref_dX
+        torch.testing.assert_close(X_kern.grad, ref_dX_ungrouped, atol=5e-2, rtol=5e-2)
+
+    def test_lora_gradient_matches_reference(self):
+        """LoRA A/B gradients from autograd match reference."""
+        from importlib import import_module
+
+        pll = import_module(f"{_SMOE}.parallel_linear_lora")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        M, K, N, E, R, k = 16, 64, 64, 4, 8, 1
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        # Autograd path
+        X_kern = data["X"].clone().requires_grad_(True)
+        lora_A_kern = data["lora_A"].clone().requires_grad_(True)
+        lora_B_kern = data["lora_B"].clone().requires_grad_(True)
+
+        out_kern = pll.ScatterMoELoRA.apply(
+            X_kern,
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            lora_A_kern,
+            lora_B_kern,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+        )
+        grad_out = torch.randn_like(out_kern)
+        out_kern.backward(grad_out)
+
+        # Reference path
+        grouped_X = base_ops.group(data["X"], data["sorted_scattered_idxs"], fan_out=k)
+        grouped_grad = base_ops.group(
+            grad_out, data["sorted_scattered_idxs"], fan_out=1
+        )
+
+        _, ref_dA, ref_dB = reference_lora_backward(
+            grouped_grad,
+            grouped_X,
+            data["W"],
+            data["lora_A"],
+            data["lora_B"],
+            data["scaling"],
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            k,
+            E,
+        )
+
+        torch.testing.assert_close(lora_A_kern.grad, ref_dA, atol=5e-2, rtol=5e-2)
+        torch.testing.assert_close(lora_B_kern.grad, ref_dB, atol=5e-2, rtol=5e-2)
+
+
+# =============================================================================
+# Test: Equivalence with Base ScatterMoE (scaling=0 should match base)
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestBaseEquivalence:
+    """When scaling=0, fused kernel should match base scatter2scatter."""
+
+    def test_zero_scaling_matches_base(self):
+        """With scaling=0, LoRA contribution vanishes; should match base."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        data = make_test_data(M=32, K=64, N=128, E=4, R=8, k=2)
+
+        base_output = base_ops.scatter2scatter(
+            X=data["X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=data["k"],
+        )
+
+        lora_output = lora_ops.scatter2scatter_lora(
+            X=data["X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=data["k"],
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            scaling=0.0,
+        )
+
+        torch.testing.assert_close(lora_output, base_output, atol=1e-3, rtol=1e-3)
+
+    def test_zero_lora_weights_matches_base(self):
+        """With A=0, B=0, should match base scatter2scatter."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        data = make_test_data(M=32, K=64, N=128, E=4, R=8, k=2)
+
+        zero_A = torch.zeros_like(data["lora_A"])
+        zero_B = torch.zeros_like(data["lora_B"])
+
+        base_output = base_ops.scatter2scatter(
+            X=data["X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=data["k"],
+        )
+
+        lora_output = lora_ops.scatter2scatter_lora(
+            X=data["X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=data["k"],
+            lora_A=zero_A,
+            lora_B=zero_B,
+            scaling=1.0,
+        )
+
+        torch.testing.assert_close(lora_output, base_output, atol=1e-3, rtol=1e-3)
+
+
+# =============================================================================
+# Test: LoRA Additivity
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestLoRAAdditivity:
+    """Test that the LoRA component is correctly additive."""
+
+    def test_lora_additivity(self):
+        """
+        Verify: fused(X, W, A, B, s) == base(X, W) + s * per_expert_lora(X, A, B)
+        """
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        data = make_test_data(M=32, K=64, N=128, E=4, R=8, k=2)
+
+        # Base output (no LoRA)
+        base_output = base_ops.scatter2scatter(
+            X=data["X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=data["k"],
+        )
+
+        # Fused output
+        fused_output = lora_ops.scatter2scatter_lora(
+            X=data["X"],
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=data["k"],
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            scaling=data["scaling"],
+        )
+
+        # Compute LoRA contribution manually (reference)
+        lora_only = reference_parallel_linear_lora(
+            data["X"],
+            torch.zeros_like(data["W"]),
+            data["k"],
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["lora_A"],
+            data["lora_B"],
+            data["scaling"],
+        )
+
+        # fused = base + lora
+        expected = base_output + lora_only
+        torch.testing.assert_close(fused_output, expected, atol=2e-2, rtol=2e-2)
+
+
+# =============================================================================
+# Test: ParallelExperts module integration
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestParallelExpertsModule:
+    """Test the ParallelExperts module with LoRA."""
+
+    def test_set_and_clear_lora(self):
+        """Test set_lora/clear_lora lifecycle."""
+        from importlib import import_module
+
+        lora_module = import_module(f"{_SMOE}.lora_ops")
+
+        pe = lora_module.ParallelExperts(4, 64, 128).cuda()
+
+        A = torch.randn(32, 64, device="cuda")  # r=8, E=4
+        B = torch.randn(128, 32, device="cuda")
+        pe.set_lora(A, B, 0.5)
+
+        assert pe._lora_A is A
+        assert pe._lora_B is B
+        assert pe._lora_scaling == 0.5
+
+        pe.clear_lora()
+        assert pe._lora_A is None
+        assert pe._lora_B is None
+
+    def test_forward_with_lora(self):
+        """ParallelExperts forward with LoRA matches reference."""
+        from importlib import import_module
+
+        lora_module = import_module(f"{_SMOE}.lora_ops")
+
+        E, K, N, R = 4, 64, 128, 8
+        M, k = 16, 2
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        pe = lora_module.ParallelExperts(E, K, N).cuda()
+        # Set weights to match test data
+        with torch.no_grad():
+            pe.weight.copy_(data["W"].permute(0, 2, 1))  # [E, N, K]
+
+        pe.set_lora(data["lora_A"], data["lora_B"], data["scaling"])
+
+        output = pe(
+            data["X"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+        )
+
+        ref = reference_parallel_linear_lora(
+            data["X"],
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["lora_A"],
+            data["lora_B"],
+            data["scaling"],
+        )
+
+        torch.testing.assert_close(output, ref, atol=2e-2, rtol=2e-2)
+
+
+# =============================================================================
+# Test: Edge Cases
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestEdgeCases:
+    """Edge cases and boundary conditions."""
+
+    def test_all_tokens_one_expert(self):
+        """All tokens routed to a single expert."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+
+        M, K, N, E, R, k = 16, 64, 64, 4, 8, 1
+        torch.manual_seed(42)
+
+        X = torch.randn(M, K, device="cuda")
+        W = torch.randn(E, K, N, device="cuda") * 0.02
+        lora_A = torch.randn(R * E, K, device="cuda") * 0.01
+        lora_B = torch.randn(N, R * E, device="cuda") * 0.01
+
+        # All tokens go to expert 0
+        selected_experts = torch.zeros(M, k, device="cuda", dtype=torch.long)
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = (
+            flatten_sort_count_ref(selected_experts, E)
+        )
+
+        ref = reference_parallel_linear_lora(
+            X,
+            W,
+            k,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            lora_A,
+            lora_B,
+            0.5,
+        )
+
+        kernel = lora_ops.scatter2scatter_lora(
+            X=X,
+            W=W,
+            sorted_expert_idxs=sorted_expert_idxs,
+            sorted_scattered_idxs=sorted_scattered_idxs,
+            k=k,
+            lora_A=lora_A,
+            lora_B=lora_B,
+            scaling=0.5,
+        )
+
+        torch.testing.assert_close(kernel, ref, atol=1e-2, rtol=1e-2)
+
+    def test_empty_experts(self):
+        """Some experts have no tokens assigned."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+
+        M, K, N, E, R, k = 8, 64, 64, 8, 4, 1
+        torch.manual_seed(42)
+
+        X = torch.randn(M, K, device="cuda")
+        W = torch.randn(E, K, N, device="cuda") * 0.02
+        lora_A = torch.randn(R * E, K, device="cuda") * 0.01
+        lora_B = torch.randn(N, R * E, device="cuda") * 0.01
+
+        # Only use experts 0 and 1
+        selected_experts = torch.randint(0, 2, (M, k), device="cuda")
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = (
+            flatten_sort_count_ref(selected_experts, E)
+        )
+
+        ref = reference_parallel_linear_lora(
+            X,
+            W,
+            k,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            lora_A,
+            lora_B,
+            0.5,
+        )
+
+        kernel = lora_ops.scatter2scatter_lora(
+            X=X,
+            W=W,
+            sorted_expert_idxs=sorted_expert_idxs,
+            sorted_scattered_idxs=sorted_scattered_idxs,
+            k=k,
+            lora_A=lora_A,
+            lora_B=lora_B,
+            scaling=0.5,
+        )
+
+        torch.testing.assert_close(kernel, ref, atol=1e-2, rtol=1e-2)
+
+
+# =============================================================================
+# Test: Optimization 1 - Fused dX Kernel
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestFusedDX:
+    """Test fused backward dX kernel: dX = dY @ W^T + scaling * (dY @ B) @ A."""
+
+    def _run_fused_dX_test(
+        self, M, K, N, E, R, k, dtype=torch.float32, atol=5e-2, rtol=5e-2
+    ):
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+        pll = import_module(f"{_SMOE}.parallel_linear_lora")
+
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k, dtype=dtype)
+
+        # Create dummy grad_out in grouped order
+        grad_out = torch.randn(
+            data["sorted_expert_idxs"].size(0), N, device="cuda", dtype=dtype
+        )
+        grouped_grad = base_ops.group(
+            grad_out,
+            data["sorted_scattered_idxs"],
+            fan_out=1,
+        )
+
+        # Reference: separate scatter2scatter(DY, W^T) + _compute_lora_input_grad
+        ref_base = base_ops.scatter2scatter(
+            X=grouped_grad,
+            x_grouped=True,
+            W=data["W"].permute(0, 2, 1),
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=1,
+            y_grouped=False,
+        )
+
+        ref_lora = pll._compute_lora_input_grad(
+            grouped_grad,
+            data["lora_A"],
+            data["lora_B"],
+            data["expert_offsets"],
+            E,
+            data["scaling"],
+        )
+        # Scatter lora from grouped to ungrouped order
+        ref_lora_ungrouped = torch.zeros_like(ref_base)
+        ref_lora_ungrouped[data["sorted_scattered_idxs"]] = ref_lora
+        ref_total = ref_base + ref_lora_ungrouped
+
+        # Fused kernel
+        fused_result = lora_ops.scatter2scatter_lora_dX(
+            DY=grouped_grad,
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=1,
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            scaling=data["scaling"],
+            dy_grouped=True,
+            dx_grouped=False,
+        )
+
+        torch.testing.assert_close(fused_result, ref_total, atol=atol, rtol=rtol)
+
+    def test_basic(self):
+        self._run_fused_dX_test(M=32, K=64, N=128, E=4, R=8, k=2)
+
+    @skip_on_out_of_resources
+    def test_large(self):
+        self._run_fused_dX_test(M=256, K=256, N=512, E=8, R=16, k=2)
+
+    def test_single_expert(self):
+        self._run_fused_dX_test(M=64, K=128, N=256, E=1, R=8, k=1)
+
+    def test_k1(self):
+        self._run_fused_dX_test(M=64, K=64, N=128, E=4, R=8, k=1)
+
+    def test_bf16(self):
+        self._run_fused_dX_test(
+            M=64,
+            K=128,
+            N=256,
+            E=4,
+            R=16,
+            k=2,
+            dtype=torch.bfloat16,
+            atol=1e-1,
+            rtol=1e-1,
+        )
+
+    def test_grouped_output(self):
+        """Test fused dX with dx_grouped=True."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+        pll = import_module(f"{_SMOE}.parallel_linear_lora")
+
+        M, K, N, E, R, k = 32, 64, 128, 4, 8, 2
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        grad_out = torch.randn(data["sorted_expert_idxs"].size(0), N, device="cuda")
+        grouped_grad = base_ops.group(
+            grad_out, data["sorted_scattered_idxs"], fan_out=1
+        )
+
+        # Reference: grouped output
+        ref_base = base_ops.scatter2scatter(
+            X=grouped_grad,
+            x_grouped=True,
+            W=data["W"].permute(0, 2, 1),
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=1,
+            y_grouped=True,  # grouped output
+        )
+
+        ref_lora = pll._compute_lora_input_grad(
+            grouped_grad,
+            data["lora_A"],
+            data["lora_B"],
+            data["expert_offsets"],
+            E,
+            data["scaling"],
+        )
+        ref_total = ref_base + ref_lora
+
+        # Fused kernel with grouped output
+        fused_result = lora_ops.scatter2scatter_lora_dX(
+            DY=grouped_grad,
+            W=data["W"],
+            sorted_expert_idxs=data["sorted_expert_idxs"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            k=1,
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            scaling=data["scaling"],
+            dy_grouped=True,
+            dx_grouped=True,
+        )
+
+        torch.testing.assert_close(fused_result, ref_total, atol=5e-2, rtol=5e-2)
+
+    def test_autograd_with_fused_dX(self):
+        """Full autograd round-trip with use_fused_dX=True."""
+        from importlib import import_module
+
+        pll = import_module(f"{_SMOE}.parallel_linear_lora")
+
+        M, K, N, E, R, k = 32, 64, 128, 4, 8, 2
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        # Run without fused dX
+        X1 = data["X"].clone().requires_grad_(True)
+        A1 = data["lora_A"].clone().requires_grad_(True)
+        B1 = data["lora_B"].clone().requires_grad_(True)
+        out1 = pll.ScatterMoELoRA.apply(
+            X1,
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            A1,
+            B1,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+            False,  # use_fused_dX=False
+        )
+        out1.sum().backward()
+
+        # Run with fused dX
+        X2 = data["X"].clone().requires_grad_(True)
+        A2 = data["lora_A"].clone().requires_grad_(True)
+        B2 = data["lora_B"].clone().requires_grad_(True)
+        out2 = pll.ScatterMoELoRA.apply(
+            X2,
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            A2,
+            B2,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+            True,  # use_fused_dX=True
+        )
+        out2.sum().backward()
+
+        # Forward should be identical
+        torch.testing.assert_close(out1, out2, atol=1e-5, rtol=1e-5)
+
+        # Gradients should match
+        torch.testing.assert_close(X1.grad, X2.grad, atol=5e-2, rtol=5e-2)
+        torch.testing.assert_close(A1.grad, A2.grad, atol=5e-2, rtol=5e-2)
+        torch.testing.assert_close(B1.grad, B2.grad, atol=5e-2, rtol=5e-2)
+
+
+# =============================================================================
+# Test: Optimization 2 - Fused Gather Backward
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestFusedGatherBackward:
+    """Test fused gather + backward dA/dB kernel."""
+
+    def _run_fused_gather_test(
+        self, M, K, N, E, R, k, dtype=torch.float32, atol=5e-2, rtol=5e-2
+    ):
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k, dtype=dtype)
+
+        # Create grad_out in ungrouped order (M*k, N)
+        M_total = data["sorted_expert_idxs"].size(0)
+        grad_out = torch.randn(M_total, N, device="cuda", dtype=dtype)
+
+        # Reference: group() + group_bwd_lora()
+        grouped_grad = base_ops.group(
+            grad_out, data["sorted_scattered_idxs"], fan_out=1
+        )
+        grouped_x = base_ops.group(data["X"], data["sorted_scattered_idxs"], fan_out=k)
+
+        ref_dA, ref_dB = lora_ops.group_bwd_lora(
+            DY=grouped_grad,
+            X=grouped_x,
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            expert_offsets=data["expert_offsets"],
+            E=E,
+            scaling=data["scaling"],
+        )
+
+        # Fused kernel: no group() calls
+        fused_dA, fused_dB = lora_ops.group_bwd_lora_fused(
+            DY=grad_out,
+            X=data["X"],
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            expert_offsets=data["expert_offsets"],
+            sorted_scattered_idxs=data["sorted_scattered_idxs"],
+            E=E,
+            k=k,
+            scaling=data["scaling"],
+        )
+
+        torch.testing.assert_close(fused_dA, ref_dA, atol=atol, rtol=rtol)
+        torch.testing.assert_close(fused_dB, ref_dB, atol=atol, rtol=rtol)
+
+    def test_basic(self):
+        self._run_fused_gather_test(M=32, K=64, N=128, E=4, R=8, k=2)
+
+    @skip_on_out_of_resources
+    def test_large(self):
+        self._run_fused_gather_test(M=256, K=256, N=512, E=8, R=16, k=2)
+
+    def test_single_expert(self):
+        self._run_fused_gather_test(M=64, K=128, N=256, E=1, R=8, k=1)
+
+    def test_k1(self):
+        self._run_fused_gather_test(M=64, K=64, N=128, E=4, R=8, k=1)
+
+    @skip_on_out_of_resources
+    def test_many_experts(self):
+        self._run_fused_gather_test(M=128, K=64, N=128, E=16, R=8, k=4)
+
+    def test_bf16(self):
+        self._run_fused_gather_test(
+            M=64,
+            K=128,
+            N=256,
+            E=4,
+            R=16,
+            k=2,
+            dtype=torch.bfloat16,
+            atol=1e-1,
+            rtol=1e-1,
+        )
+
+    def test_autograd_with_fused_gather(self):
+        """Full autograd round-trip with use_fused_gather=True."""
+        from importlib import import_module
+
+        pll = import_module(f"{_SMOE}.parallel_linear_lora")
+
+        M, K, N, E, R, k = 32, 64, 128, 4, 8, 2
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        # Run without fused gather
+        X1 = data["X"].clone().requires_grad_(True)
+        A1 = data["lora_A"].clone().requires_grad_(True)
+        B1 = data["lora_B"].clone().requires_grad_(True)
+        out1 = pll.ScatterMoELoRA.apply(
+            X1,
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            A1,
+            B1,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+            False,
+            False,  # use_fused_dX=False, use_fused_gather=False
+        )
+        out1.sum().backward()
+
+        # Run with fused gather
+        X2 = data["X"].clone().requires_grad_(True)
+        A2 = data["lora_A"].clone().requires_grad_(True)
+        B2 = data["lora_B"].clone().requires_grad_(True)
+        out2 = pll.ScatterMoELoRA.apply(
+            X2,
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            A2,
+            B2,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+            False,
+            True,  # use_fused_dX=False, use_fused_gather=True
+        )
+        out2.sum().backward()
+
+        # Forward identical
+        torch.testing.assert_close(out1, out2, atol=1e-5, rtol=1e-5)
+
+        # dA/dB should match
+        torch.testing.assert_close(A1.grad, A2.grad, atol=5e-2, rtol=5e-2)
+        torch.testing.assert_close(B1.grad, B2.grad, atol=5e-2, rtol=5e-2)
+        # dX should also match (same path for dX)
+        torch.testing.assert_close(X1.grad, X2.grad, atol=5e-2, rtol=5e-2)
+
+
+# =============================================================================
+# Test: Optimization 3 - Token Rounding
+# =============================================================================
+
+
+@pytest.mark.slow
+@pytest.mark.xfail(reason="flaky", strict=False)
+class TestTokenRounding:
+    """Test token rounding utility and its integration with backward kernels."""
+
+    def test_round_expert_counts_basic(self):
+        """Verify round_expert_counts produces correct shapes and values."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+
+        M, K, N, E, R, k = 32, 64, 128, 4, 8, 2
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        padded_ei, padded_si, padded_offsets, real_offsets = (
+            lora_ops.round_expert_counts(
+                data["sorted_expert_idxs"],
+                data["sorted_scattered_idxs"],
+                data["expert_offsets"],
+                E=E,
+                block_m=lora_ops.BLOCK_M,
+            )
+        )
+
+        # Real offsets should match original
+        torch.testing.assert_close(real_offsets, data["expert_offsets"])
+
+        # Padded offsets should be >= real offsets
+        assert (padded_offsets >= real_offsets).all(), (
+            "Padded offsets should be >= real offsets"
+        )
+
+        # Each expert's padded count should be multiple of BLOCK_M (if non-zero)
+        prev = 0
+        for e in range(E):
+            count = padded_offsets[e].item() - prev
+            real_count = real_offsets[e].item() - (
+                real_offsets[e - 1].item() if e > 0 else 0
+            )
+            if real_count > 0:
+                assert count % lora_ops.BLOCK_M == 0, (
+                    f"Expert {e}: padded count {count} not multiple of {lora_ops.BLOCK_M}"
+                )
+                assert count >= real_count, (
+                    f"Expert {e}: padded count {count} < real count {real_count}"
+                )
+            prev = padded_offsets[e].item()
+
+    @skip_on_out_of_resources
+    def test_round_with_fused_gather(self):
+        """Token rounding + fused gather gives same result as plain fused gather."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+        base_ops = import_module(f"{_SMOE}.kernels.ops")
+
+        M, K, N, E, R, k = 64, 64, 128, 4, 8, 2
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        M_total = data["sorted_expert_idxs"].size(0)
+        grad_out = torch.randn(M_total, N, device="cuda")
+
+        # Reference: group() + group_bwd_lora() (the gold standard)
+        grouped_grad = base_ops.group(
+            grad_out, data["sorted_scattered_idxs"], fan_out=1
+        )
+        grouped_x = base_ops.group(data["X"], data["sorted_scattered_idxs"], fan_out=k)
+        ref_dA, ref_dB = lora_ops.group_bwd_lora(
+            DY=grouped_grad,
+            X=grouped_x,
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            expert_offsets=data["expert_offsets"],
+            E=E,
+            scaling=data["scaling"],
+        )
+
+        # Apply token rounding
+        padded_ei, padded_si, padded_offsets, real_offsets = (
+            lora_ops.round_expert_counts(
+                data["sorted_expert_idxs"],
+                data["sorted_scattered_idxs"],
+                data["expert_offsets"],
+                E=E,
+            )
+        )
+
+        # Fused gather with token rounding
+        rounded_dA, rounded_dB = lora_ops.group_bwd_lora_fused(
+            DY=grad_out,
+            X=data["X"],
+            lora_A=data["lora_A"],
+            lora_B=data["lora_B"],
+            expert_offsets=padded_offsets,
+            sorted_scattered_idxs=padded_si,
+            E=E,
+            k=k,
+            scaling=data["scaling"],
+            real_expert_offsets=real_offsets,
+        )
+
+        torch.testing.assert_close(rounded_dA, ref_dA, atol=5e-2, rtol=5e-2)
+        torch.testing.assert_close(rounded_dB, ref_dB, atol=5e-2, rtol=5e-2)
+
+    def test_empty_experts_with_rounding(self):
+        """Token rounding handles experts with 0 tokens correctly."""
+        from importlib import import_module
+
+        lora_ops = import_module(f"{_SMOE}.kernels.lora_ops")
+
+        E, k = 8, 1
+        M = 8
+        torch.manual_seed(42)
+
+        # Only use experts 0 and 1 (rest have 0 tokens)
+        selected_experts = torch.randint(0, 2, (M, k), device="cuda")
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = (
+            flatten_sort_count_ref(selected_experts, E)
+        )
+
+        padded_ei, padded_si, padded_offsets, real_offsets = (
+            lora_ops.round_expert_counts(
+                sorted_expert_idxs,
+                sorted_scattered_idxs,
+                expert_offsets,
+                E=E,
+            )
+        )
+
+        # Verify empty experts have same count (0)
+        for e in range(E):
+            real_count = real_offsets[e].item() - (
+                real_offsets[e - 1].item() if e > 0 else 0
+            )
+            padded_count = padded_offsets[e].item() - (
+                padded_offsets[e - 1].item() if e > 0 else 0
+            )
+            if real_count == 0:
+                assert padded_count == 0, (
+                    f"Expert {e}: empty expert should have padded_count=0, got {padded_count}"
+                )
+
+
+# =============================================================================
+# Test: Combined Optimizations
+# =============================================================================
+
+
+@pytest.mark.slow
+class TestCombinedOptimizations:
+    """Test all optimizations together."""
+
+    def test_fused_dX_and_fused_gather(self):
+        """Both fused dX and fused gather together."""
+        from importlib import import_module
+
+        pll = import_module(f"{_SMOE}.parallel_linear_lora")
+
+        M, K, N, E, R, k = 64, 128, 256, 4, 8, 2
+        data = make_test_data(M=M, K=K, N=N, E=E, R=R, k=k)
+
+        # Baseline: no optimizations
+        X1 = data["X"].clone().requires_grad_(True)
+        A1 = data["lora_A"].clone().requires_grad_(True)
+        B1 = data["lora_B"].clone().requires_grad_(True)
+        out1 = pll.ScatterMoELoRA.apply(
+            X1,
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            A1,
+            B1,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+            False,
+            False,  # no optimizations
+        )
+        out1.sum().backward()
+
+        # Both optimizations
+        X2 = data["X"].clone().requires_grad_(True)
+        A2 = data["lora_A"].clone().requires_grad_(True)
+        B2 = data["lora_B"].clone().requires_grad_(True)
+        out2 = pll.ScatterMoELoRA.apply(
+            X2,
+            data["W"],
+            k,
+            data["sorted_expert_idxs"],
+            data["sorted_scattered_idxs"],
+            data["expert_offsets"],
+            A2,
+            B2,
+            data["scaling"],
+            None,
+            None,
+            False,
+            False,
+            True,
+            True,  # use_fused_dX=True, use_fused_gather=True
+        )
+        out2.sum().backward()
+
+        # Forward identical
+        torch.testing.assert_close(out1, out2, atol=1e-5, rtol=1e-5)
+
+        # All gradients match
+        torch.testing.assert_close(X1.grad, X2.grad, atol=5e-2, rtol=5e-2)
+        torch.testing.assert_close(A1.grad, A2.grad, atol=5e-2, rtol=5e-2)
+        torch.testing.assert_close(B1.grad, B2.grad, atol=5e-2, rtol=5e-2)
+
+
+# =============================================================================
+# Test: HFScatterMoEGatedMLP with Sigmoid Routing
+# =============================================================================
+
+
+def _reference_moe_forward(
+    hidden_states,
+    gate_weight,
+    gate_up_proj,
+    down_proj,
+    act_fn,
+    routing_weights,
+    selected_experts,
+    num_experts,
+):
+    """Pure PyTorch reference for a full MoE forward pass.
+
+    Args:
+        hidden_states: [T, H]
+        gate_weight: [E, H]
+        gate_up_proj: [E, 2*FF, H]
+        down_proj: [E, H, FF]
+        act_fn: activation function (e.g. torch.nn.SiLU())
+        routing_weights: [T, K] routing weights
+        selected_experts: [T, K] expert indices
+        num_experts: int
+
+    Returns:
+        output: [T, H]
+    """
+    T, H = hidden_states.shape
+    K = selected_experts.shape[1]
+    output = torch.zeros(T, H, device=hidden_states.device, dtype=hidden_states.dtype)
+
+    for t in range(T):
+        for j in range(K):
+            e = selected_experts[t, j].item()
+            w = routing_weights[t, j].item()
+
+            # gate_up projection
+            gup = hidden_states[t] @ gate_up_proj[e].T  # [2*I]
+            I_dim = gup.shape[0] // 2
+            gates = gup[:I_dim]
+            up = gup[I_dim:]
+
+            # activation
+            h = act_fn(gates) * up
+
+            # down projection
+            out = h @ down_proj[e].T  # [H]
+
+            output[t] += w * out
+
+    return output
+
+
+def _make_mock_sigmoid_moe_block(
+    T=16, H=64, FF=32, E=8, K=2, n_group=2, topk_group=1, bias_on_gate=True
+):
+    """Create a mock MoE block with sigmoid routing for GPU testing."""
+    gate_up_proj = torch.randn(E, 2 * FF, H, device="cuda") * 0.02
+    down_proj = torch.randn(E, H, FF, device="cuda") * 0.02
+    act_fn = torch.nn.SiLU()
+
+    experts = SimpleNamespace(
+        gate_up_proj=gate_up_proj,
+        down_proj=down_proj,
+        act_fn=act_fn,
+        num_experts=E,
+    )
+
+    if bias_on_gate:
+        gate = SimpleNamespace(
+            weight=torch.randn(E, H, device="cuda") * 0.1,
+            e_score_correction_bias=torch.zeros(E, device="cuda"),
+        )
+        moe_block = SimpleNamespace(
+            gate=gate,
+            experts=experts,
+            top_k=K,
+            n_routed_experts=E,
+            n_group=n_group,
+            topk_group=topk_group,
+            norm_topk_prob=True,
+            routed_scaling_factor=1.0,
+        )
+    else:
+        # minimax_m2 style
+        gate = SimpleNamespace(
+            weight=torch.randn(E, H, device="cuda") * 0.1,
+            top_k=K,
+        )
+        moe_block = SimpleNamespace(
+            gate=gate,
+            experts=experts,
+            top_k=K,
+            e_score_correction_bias=torch.zeros(E, device="cuda"),
+        )
+
+    return moe_block, T, H, FF, E, K
+
+
+@pytest.mark.slow
+class TestHFScatterMoESigmoidRouting:
+    """Test HFScatterMoEGatedMLP forward with sigmoid routing on GPU."""
+
+    def test_forward_matches_reference_bias_on_gate(self):
+        """Forward pass with sigmoid routing (bias on gate) matches reference."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            HFScatterMoEGatedMLP,
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, FF, E, K = _make_mock_sigmoid_moe_block(
+            T=16, H=64, FF=32, E=8, K=2, n_group=2, topk_group=1, bias_on_gate=True
+        )
+
+        hidden = torch.randn(1, T, H, device="cuda")
+
+        # Get routing for reference
+        gate = moe_block.gate
+        hidden_flat = hidden.view(-1, H)
+        routing_weights, selected_experts, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden_flat, gate.weight, None
+        )
+
+        # Reference output
+        ref_output = _reference_moe_forward(
+            hidden_flat,
+            gate.weight,
+            moe_block.experts.gate_up_proj,
+            moe_block.experts.down_proj,
+            moe_block.experts.act_fn,
+            routing_weights,
+            selected_experts,
+            E,
+        )
+
+        # Kernel output
+        kernel_output = HFScatterMoEGatedMLP.forward(moe_block, hidden)
+        kernel_output_flat = kernel_output.view(-1, H)
+
+        torch.testing.assert_close(
+            kernel_output_flat.float(),
+            ref_output.float(),
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+    def test_forward_matches_reference_bias_on_block(self):
+        """Forward pass with sigmoid routing (minimax_m2 style, bias on block)."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            HFScatterMoEGatedMLP,
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, FF, E, K = _make_mock_sigmoid_moe_block(
+            T=16, H=64, FF=32, E=8, K=2, n_group=1, bias_on_gate=False
+        )
+
+        hidden = torch.randn(1, T, H, device="cuda")
+        hidden_flat = hidden.view(-1, H)
+
+        gate = moe_block.gate
+        routing_weights, selected_experts, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden_flat, gate.weight, None
+        )
+
+        ref_output = _reference_moe_forward(
+            hidden_flat,
+            gate.weight,
+            moe_block.experts.gate_up_proj,
+            moe_block.experts.down_proj,
+            moe_block.experts.act_fn,
+            routing_weights,
+            selected_experts,
+            E,
+        )
+
+        kernel_output = HFScatterMoEGatedMLP.forward(moe_block, hidden)
+        kernel_output_flat = kernel_output.view(-1, H)
+
+        torch.testing.assert_close(
+            kernel_output_flat.float(),
+            ref_output.float(),
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+    def test_softmax_routing_still_works(self):
+        """Verify softmax routing (Qwen/OLMoE) is not broken."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            HFScatterMoEGatedMLP,
+            _softmax_topk_route,
+        )
+
+        T, H, FF, E, K = 16, 64, 32, 4, 2
+        gate_up_proj = torch.randn(E, 2 * FF, H, device="cuda") * 0.02
+        down_proj = torch.randn(E, H, FF, device="cuda") * 0.02
+        act_fn = torch.nn.SiLU()
+
+        experts = SimpleNamespace(
+            gate_up_proj=gate_up_proj,
+            down_proj=down_proj,
+            act_fn=act_fn,
+            num_experts=E,
+        )
+        gate = SimpleNamespace(
+            weight=torch.randn(E, H, device="cuda") * 0.1,
+            top_k=K,
+            num_experts=E,
+            norm_topk_prob=True,
+        )
+        moe_block = SimpleNamespace(gate=gate, experts=experts)
+
+        hidden = torch.randn(1, T, H, device="cuda")
+        hidden_flat = hidden.view(-1, H)
+
+        routing_weights, selected_experts, _, _ = _softmax_topk_route(
+            moe_block, gate, hidden_flat, gate.weight, None
+        )
+
+        ref_output = _reference_moe_forward(
+            hidden_flat,
+            gate.weight,
+            gate_up_proj,
+            down_proj,
+            act_fn,
+            routing_weights,
+            selected_experts,
+            E,
+        )
+
+        kernel_output = HFScatterMoEGatedMLP.forward(moe_block, hidden)
+        kernel_output_flat = kernel_output.view(-1, H)
+
+        torch.testing.assert_close(
+            kernel_output_flat.float(),
+            ref_output.float(),
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+
+@pytest.mark.slow
+class TestHFScatterMoESigmoidWithSharedExperts:
+    """Test HFScatterMoEGatedMLP with sigmoid routing + shared experts."""
+
+    def test_shared_experts_plural(self):
+        """DeepSeek V3 style: shared_experts attribute (plural)."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            HFScatterMoEGatedMLP,
+        )
+
+        T, H, FF, E, K = 8, 64, 32, 8, 2
+        gate_up_proj = torch.randn(E, 2 * FF, H, device="cuda") * 0.02
+        down_proj = torch.randn(E, H, FF, device="cuda") * 0.02
+        act_fn = torch.nn.SiLU()
+
+        experts = SimpleNamespace(
+            gate_up_proj=gate_up_proj,
+            down_proj=down_proj,
+            act_fn=act_fn,
+            num_experts=E,
+        )
+
+        # Shared expert as a simple linear for testing
+        shared_W = torch.randn(H, H, device="cuda") * 0.01
+        shared_experts_fn = lambda x: x @ shared_W.T  # noqa: E731
+
+        gate = SimpleNamespace(
+            weight=torch.randn(E, H, device="cuda") * 0.1,
+            e_score_correction_bias=torch.zeros(E, device="cuda"),
+        )
+        moe_block = SimpleNamespace(
+            gate=gate,
+            experts=experts,
+            shared_experts=shared_experts_fn,
+            top_k=K,
+            n_routed_experts=E,
+            n_group=1,
+            norm_topk_prob=True,
+            routed_scaling_factor=1.0,
+        )
+
+        hidden = torch.randn(1, T, H, device="cuda")
+
+        # Should not raise; output should include shared expert contribution
+        output = HFScatterMoEGatedMLP.forward(moe_block, hidden)
+        assert output.shape == (1, T, H)
+
+        # Run without shared expert to verify it changes the output
+        moe_block_no_shared = SimpleNamespace(
+            gate=gate,
+            experts=experts,
+            top_k=K,
+            n_routed_experts=E,
+            n_group=1,
+            norm_topk_prob=True,
+            routed_scaling_factor=1.0,
+        )
+        output_no_shared = HFScatterMoEGatedMLP.forward(moe_block_no_shared, hidden)
+        assert not torch.equal(output, output_no_shared)
+
+    def test_shared_expert_with_gate(self):
+        """Qwen2MoE style: shared_expert + shared_expert_gate."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            HFScatterMoEGatedMLP,
+        )
+
+        T, H, FF, E, K = 8, 64, 32, 4, 2
+        gate_up_proj = torch.randn(E, 2 * FF, H, device="cuda") * 0.02
+        down_proj = torch.randn(E, H, FF, device="cuda") * 0.02
+        act_fn = torch.nn.SiLU()
+
+        experts = SimpleNamespace(
+            gate_up_proj=gate_up_proj,
+            down_proj=down_proj,
+            act_fn=act_fn,
+            num_experts=E,
+        )
+
+        shared_W = torch.randn(H, H, device="cuda") * 0.01
+        shared_expert_fn = lambda x: x @ shared_W.T  # noqa: E731
+        # Gate that returns 0 -> sigmoid(0) = 0.5
+        gate_W = torch.zeros(H, H, device="cuda")
+        shared_expert_gate_fn = lambda x: x @ gate_W.T  # noqa: E731
+
+        gate = SimpleNamespace(
+            weight=torch.randn(E, H, device="cuda") * 0.1,
+            top_k=K,
+            num_experts=E,
+            norm_topk_prob=True,
+        )
+        moe_block = SimpleNamespace(
+            gate=gate,
+            experts=experts,
+            shared_expert=shared_expert_fn,
+            shared_expert_gate=shared_expert_gate_fn,
+        )
+
+        hidden = torch.randn(1, T, H, device="cuda")
+        output = HFScatterMoEGatedMLP.forward(moe_block, hidden)
+        assert output.shape == (1, T, H)
diff --git a/tests/e2e/integrations/test_scattermoe_lora_olmoe.py b/tests/e2e/integrations/test_scattermoe_lora_olmoe.py
new file mode 100644
index 0000000000..b3f743c4e3
--- /dev/null
+++ b/tests/e2e/integrations/test_scattermoe_lora_olmoe.py
@@ -0,0 +1,1245 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Integration tests: OLMoE + peft LoRA + ScatterMoE fused kernels.
+
+Validates that scattermoe_lora fused kernels produce correct results when used
+with HuggingFace OLMoE models and peft LoRA adapters applied via
+``target_parameters``.
+
+Key things tested
+-----------------
+- LoRA weight layout conversion between peft (rank-major) and scattermoe (expert-major)
+- Base forward equivalence: per-expert reference vs ScatterMoE kernels (no LoRA)
+- LoRA forward equivalence: peft merged-weight approach vs scattermoe fused kernels
+- Backward gradient correctness through the fused LoRA path
+- ``kernelize()`` integration via ``LocalLayerRepository``
+"""
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from peft import LoraConfig, get_peft_model
+from transformers import OlmoeConfig
+from transformers.models.olmoe.modeling_olmoe import OlmoeSparseMoeBlock
+
+_SMOE = "axolotl.integrations.kernels.libs.scattermoe_lora"
+
+# Try to import from axolotl's scattermoe_lora.layers; may fail on CPU without triton.
+try:
+    from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+        _unwrap_experts_lora,
+        _unwrap_gate_lora,
+        peft_lora_B_to_scattermoe,
+        peft_lora_to_scattermoe,
+    )
+
+    HAS_SCATTERMOE = True
+except (ImportError, ModuleNotFoundError):
+    HAS_SCATTERMOE = False
+
+    # Provide pure-torch fallbacks for CPU-only layout conversion tests.
+    def peft_lora_B_to_scattermoe(peft_B, num_experts, rank):
+        N = peft_B.shape[0]
+        return (
+            peft_B.reshape(N, rank, num_experts)
+            .permute(0, 2, 1)
+            .contiguous()
+            .reshape(N, num_experts * rank)
+        )
+
+    def peft_lora_to_scattermoe(peft_A, peft_B, num_experts, rank):
+        smoe_A = peft_A
+        smoe_B = peft_lora_B_to_scattermoe(peft_B, num_experts, rank)
+        return smoe_A, smoe_B
+
+    def _unwrap_experts_lora(experts_module):
+        return experts_module, None, None
+
+    def _unwrap_gate_lora(gate_module):
+        if hasattr(gate_module, "base_layer") and hasattr(gate_module, "lora_A"):
+            base_gate = gate_module.base_layer
+            active = getattr(gate_module, "active_adapters", ["default"])
+            name = active[0] if active else "default"
+            lora_A_dict = getattr(gate_module, "lora_A", {})
+            lora_B_dict = getattr(gate_module, "lora_B", {})
+            scaling_dict = getattr(gate_module, "scaling", {})
+            if name in lora_A_dict:
+                lora_A = lora_A_dict[name].weight
+                lora_B = lora_B_dict[name].weight
+                s = scaling_dict[name]
+                delta = s * (lora_B @ lora_A)
+                return base_gate, base_gate.weight, delta
+            return base_gate, base_gate.weight, None
+        return gate_module, gate_module.weight, None
+
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+FULL_OLMOE_CONFIG = dict(
+    hidden_size=2048,
+    intermediate_size=1024,
+    num_experts=64,
+    num_experts_per_tok=8,
+    hidden_act="silu",
+    norm_topk_prob=False,
+)
+
+SMALL_OLMOE_CONFIG = dict(
+    hidden_size=128,
+    intermediate_size=48,  # non-square: 2*inter=96 != hidden=128
+    num_experts=8,
+    num_experts_per_tok=2,
+    hidden_act="silu",
+    norm_topk_prob=False,
+)
+
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA not available"
+)
+
+
+def make_olmoe_config(use_full=False):
+    cfg = dict(FULL_OLMOE_CONFIG if use_full else SMALL_OLMOE_CONFIG)
+    cfg["experts_implementation"] = "grouped_mm"
+    return OlmoeConfig(**cfg)
+
+
+def make_local_layer_repository(
+    LocalLayerRepository: Any, repo_path: Path, layer_name: str
+):
+    try:
+        return LocalLayerRepository(
+            repo_path=repo_path,
+            package_name="scattermoe_lora",
+            layer_name=layer_name,
+        )
+    except TypeError as exc:
+        if "package_name" not in str(exc):
+            raise
+        return LocalLayerRepository(repo_path=repo_path, layer_name=layer_name)
+
+
+# =============================================================================
+# Layout conversion utilities (test-local helpers)
+# =============================================================================
+
+
+def scattermoe_lora_B_to_peft(smoe_B, num_experts, rank):
+    """Inverse of ``peft_lora_B_to_scattermoe``."""
+    N = smoe_B.shape[0]
+    return (
+        smoe_B.reshape(N, num_experts, rank)
+        .permute(0, 2, 1)
+        .contiguous()
+        .reshape(N, num_experts * rank)
+    )
+
+
+def peft_gate_up_lora_to_scattermoe(peft_A, peft_B, num_experts, rank):
+    """Convert peft LoRA for gate_up_proj to scattermoe layout.
+
+    Both gate_up_proj and down_proj need the A<->B swap because
+    scattermoe transposes the parameter (W = param.T).
+    """
+    return peft_lora_to_scattermoe(peft_A, peft_B, num_experts, rank)
+
+
+# =============================================================================
+# Helpers
+# =============================================================================
+
+
+def _init_expert_weights(moe_block):
+    """Initialize OlmoeExperts parameters which use torch.empty (uninitialized).
+
+    Without this, gate_up_proj and down_proj contain garbage/NaN values.
+    """
+    with torch.no_grad():
+        nn.init.kaiming_uniform_(moe_block.experts.gate_up_proj)
+        nn.init.kaiming_uniform_(moe_block.experts.down_proj)
+    return moe_block
+
+
+class MinimalOLMoEModel(nn.Module):
+    """Thin wrapper so peft's get_peft_model can attach adapters."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.moe = OlmoeSparseMoeBlock(config)
+        _init_expert_weights(self.moe)
+
+    def forward(self, x):
+        return self.moe(x)
+
+
+def _get_routing(moe_block, hidden_states):
+    """Run the router and return (routing_weights, selected_experts)."""
+    with torch.no_grad():
+        _, routing_weights, selected_experts = moe_block.gate(
+            hidden_states.view(-1, hidden_states.size(-1))
+        )
+    return routing_weights, selected_experts
+
+
+def _reference_moe_forward(
+    x_flat,
+    gate_up_proj,
+    down_proj,
+    act_fn,
+    top_k_index,
+    top_k_weights,
+    num_experts,
+):
+    """Pure-PyTorch per-expert reference MoE forward (no LoRA).
+
+    Uses F.linear per expert for an apples-to-apples comparison with
+    the ScatterMoE kernel path.
+    """
+    final = torch.zeros_like(x_flat)
+    expert_mask = F.one_hot(top_k_index, num_classes=num_experts).permute(2, 1, 0)
+    for e in range(num_experts):
+        top_k_pos, token_idx = torch.where(expert_mask[e])
+        if token_idx.numel() == 0:
+            continue
+        cur = x_flat[token_idx]
+        gate_up = F.linear(cur, gate_up_proj[e])
+        g, u = gate_up.chunk(2, dim=-1)
+        h = act_fn(g) * u
+        out = F.linear(h, down_proj[e])
+        out = out * top_k_weights[token_idx, top_k_pos, None]
+        final.index_add_(0, token_idx, out.to(final.dtype))
+    return final
+
+
+def _reference_moe_forward_with_lora(
+    x_flat,
+    gate_up_proj,
+    down_proj,
+    act_fn,
+    top_k_index,
+    top_k_weights,
+    num_experts,
+    gup_delta,
+    down_delta,
+):
+    """Pure-PyTorch reference MoE forward with pre-computed weight deltas."""
+    merged_gup = gate_up_proj + gup_delta
+    merged_down = down_proj + down_delta
+    return _reference_moe_forward(
+        x_flat,
+        merged_gup,
+        merged_down,
+        act_fn,
+        top_k_index,
+        top_k_weights,
+        num_experts,
+    )
+
+
+def _compute_delta_from_scattermoe_lora(lora_A, lora_B, scaling, E, r, param_shape):
+    """Compute additive weight delta from scattermoe-layout LoRA weights.
+
+    delta[e] = scaling * B_e @ A_e  where A_e [r,K], B_e [N,r] -> [N,K].
+    """
+    delta = torch.zeros(param_shape, device=lora_A.device, dtype=lora_A.dtype)
+    for e in range(E):
+        A_e = lora_A[e * r : (e + 1) * r, :]
+        B_e = lora_B[:, e * r : (e + 1) * r]
+        delta[e] = scaling * (B_e @ A_e)
+    return delta
+
+
+# =============================================================================
+# Tests: Layout conversion
+# =============================================================================
+
+
+class TestLoRABLayoutConversion:
+    """Test the peft <-> scattermoe lora_B layout conversion."""
+
+    def test_roundtrip(self):
+        E, r, N = 8, 4, 64
+        original = torch.randn(N, E * r)
+        converted = peft_lora_B_to_scattermoe(original, E, r)
+        back = scattermoe_lora_B_to_peft(converted, E, r)
+        torch.testing.assert_close(back, original)
+
+    def test_per_expert_slices(self):
+        """After conversion, scattermoe slicing gives the same per-expert
+        matrices as peft's reshape slicing."""
+        E, r, N = 4, 2, 16
+        peft_B = torch.randn(N, E * r)
+        smoe_B = peft_lora_B_to_scattermoe(peft_B, E, r)
+
+        peft_reshaped = peft_B.reshape(N, r, E)
+        for e in range(E):
+            torch.testing.assert_close(
+                smoe_B[:, e * r : (e + 1) * r],
+                peft_reshaped[:, :, e],
+            )
+
+    def test_lora_A_already_compatible(self):
+        """lora_A layout is identical between peft and scattermoe."""
+        E, r, K = 4, 2, 16
+        lora_A = torch.randn(E * r, K)
+        peft_reshaped = lora_A.reshape(E, r, K)
+        for e in range(E):
+            torch.testing.assert_close(
+                lora_A[e * r : (e + 1) * r, :],
+                peft_reshaped[e],
+            )
+
+    def test_delta_weight_equivalence(self):
+        """peft's einsum delta matches per-expert B @ A with converted layouts."""
+        E, r, K, N = 8, 4, 32, 64
+        peft_A = torch.randn(E * r, K)
+        peft_B = torch.randn(N, E * r)
+        scaling = 2.0
+
+        A_r = peft_A.reshape(E, r, K)
+        B_r = peft_B.reshape(N, r, E)
+        delta_peft = torch.einsum("o r e, e r i -> e i o", B_r, A_r) * scaling
+
+        smoe_B = peft_lora_B_to_scattermoe(peft_B, E, r)
+        for e in range(E):
+            A_e = peft_A[e * r : (e + 1) * r, :]
+            B_e = smoe_B[:, e * r : (e + 1) * r]
+            delta_e = scaling * (B_e @ A_e)
+            torch.testing.assert_close(delta_e, delta_peft[e].T, atol=1e-5, rtol=1e-5)
+
+    def test_down_proj_conversion(self):
+        """Verify peft_lora_to_scattermoe produces correct delta."""
+        E, r = 4, 2
+        hidden, inter = 32, 16
+        scaling = 2.0
+
+        # peft >=0.19.1 for down_proj [E, hidden, inter]:
+        # swaps in/out, lora_A [r*E, inter], lora_B [hidden, r*E]
+        peft_A = torch.randn(E * r, inter)
+        peft_B = torch.randn(hidden, E * r)
+
+        A_r = peft_A.reshape(E, r, inter)
+        B_r = peft_B.reshape(hidden, r, E)
+        delta_peft = torch.einsum("o r e, e r i -> e o i", B_r, A_r) * scaling
+
+        smoe_A, smoe_B = peft_lora_to_scattermoe(peft_A, peft_B, E, r)
+        for e in range(E):
+            A_e = smoe_A[e * r : (e + 1) * r, :]
+            B_e = smoe_B[:, e * r : (e + 1) * r]
+            delta_smoe_e = scaling * (B_e @ A_e)
+            torch.testing.assert_close(
+                delta_smoe_e, delta_peft[e], atol=1e-5, rtol=1e-5
+            )
+
+    def test_gate_up_proj_conversion(self):
+        """Verify gate_up_proj LoRA conversion with non-square dims.
+
+        gate_up_proj param: [E, 2*inter, hidden].
+        peft swaps in/out for 3D: lora_A [r*E, hidden], lora_B [2*inter, r*E].
+        scattermoe needs: lora_A [r*E, K=hidden], lora_B [N=2*inter, r*E].
+        """
+        E, r = 4, 2
+        hidden, inter = 32, 12  # 2*inter=24 != hidden=32
+        scaling = 2.0
+
+        peft_A = torch.randn(E * r, hidden)  # [r*E, in=hidden]
+        peft_B = torch.randn(2 * inter, E * r)  # [out=2*inter, r*E]
+
+        A_r = peft_A.reshape(E, r, hidden)
+        B_r = peft_B.reshape(2 * inter, r, E)
+        delta_peft = torch.einsum("o r e, e r i -> e o i", B_r, A_r) * scaling
+
+        smoe_A, smoe_B = peft_gate_up_lora_to_scattermoe(peft_A, peft_B, E, r)
+        # smoe_A should be [r*E, K=hidden], smoe_B should be [N=2*inter, r*E]
+        assert smoe_A.shape == (E * r, hidden), (
+            f"Expected {(E * r, hidden)}, got {smoe_A.shape}"
+        )
+        assert smoe_B.shape == (2 * inter, E * r), (
+            f"Expected {(2 * inter, E * r)}, got {smoe_B.shape}"
+        )
+
+        for e in range(E):
+            A_e = smoe_A[e * r : (e + 1) * r, :]  # [r, K=hidden]
+            B_e = smoe_B[:, e * r : (e + 1) * r]  # [N=2*inter, r]
+            delta_smoe_e = scaling * (B_e @ A_e)  # [2*inter, hidden]
+            # Should match peft delta which is [2*inter, hidden] = param[e]
+            torch.testing.assert_close(
+                delta_smoe_e, delta_peft[e], atol=1e-5, rtol=1e-5
+            )
+
+
+# =============================================================================
+# Tests: peft weight extraction
+# =============================================================================
+
+
+class TestPeftLoRAWeightExtraction:
+    """Test extracting peft LoRA weights for OLMoE."""
+
+    def test_peft_creates_correct_shapes(self):
+        config = make_olmoe_config(use_full=False)
+        E, r = config.num_experts, 4
+
+        model = MinimalOLMoEModel(config)
+        lora_config = LoraConfig(
+            r=r,
+            lora_alpha=16,
+            target_modules=[],
+            target_parameters=[
+                "gate.weight",
+                "experts.gate_up_proj",
+                "experts.down_proj",
+            ],
+            bias="none",
+        )
+        peft_model = get_peft_model(model, lora_config)
+        trainable = {n: p for n, p in peft_model.named_parameters() if p.requires_grad}
+
+        # Gate router
+        assert trainable["base_model.model.moe.gate.lora_A.default.weight"].shape == (
+            r,
+            config.hidden_size,
+        )
+        assert trainable["base_model.model.moe.gate.lora_B.default.weight"].shape == (
+            E,
+            r,
+        )
+
+        # gate_up_proj [E, 2*inter, hidden] — peft swaps in/out for 3D
+        assert trainable[
+            "base_model.model.moe.experts.base_layer.lora_A.default.weight"
+        ].shape == (E * r, config.hidden_size)
+        assert trainable[
+            "base_model.model.moe.experts.base_layer.lora_B.default.weight"
+        ].shape == (2 * config.intermediate_size, E * r)
+
+        # down_proj [E, hidden, inter] — peft swaps in/out for 3D
+        assert trainable[
+            "base_model.model.moe.experts.lora_A.default.weight"
+        ].shape == (E * r, config.intermediate_size)
+        assert trainable[
+            "base_model.model.moe.experts.lora_B.default.weight"
+        ].shape == (config.hidden_size, E * r)
+
+    @requires_cuda
+    def test_peft_forward_runs(self):
+        """Smoke test: peft model forward pass completes (needs CUDA for grouped_mm)."""
+        config = make_olmoe_config(use_full=False)
+        model = MinimalOLMoEModel(config)
+        lora_config = LoraConfig(
+            r=4,
+            lora_alpha=16,
+            target_modules=[],
+            target_parameters=[
+                "gate.weight",
+                "experts.gate_up_proj",
+                "experts.down_proj",
+            ],
+            bias="none",
+        )
+        peft_model = get_peft_model(model, lora_config)
+        x = torch.randn(1, 4, config.hidden_size)
+        out = peft_model(x)
+        assert out.shape == x.shape
+
+    @pytest.mark.skipif(
+        not HAS_SCATTERMOE, reason="scattermoe_lora not importable (no triton)"
+    )
+    def test_unwrap_experts_lora(self):
+        """Test that _unwrap_experts_lora correctly detects LoRA wrappers."""
+        config = make_olmoe_config(use_full=False)
+        model = MinimalOLMoEModel(config)
+        lora_config = LoraConfig(
+            r=4,
+            lora_alpha=16,
+            target_modules=[],
+            target_parameters=["experts.gate_up_proj", "experts.down_proj"],
+            bias="none",
+        )
+        peft_model = get_peft_model(model, lora_config)
+        base_moe = peft_model.base_model.model.moe
+
+        # Experts should be wrapped by ParamWrapper
+        experts, gup_lora, down_lora = _unwrap_experts_lora(base_moe.experts)
+
+        # Base experts should have the raw parameters
+        assert hasattr(experts, "gate_up_proj")
+        assert hasattr(experts, "down_proj")
+
+        # LoRA should be detected
+        assert gup_lora is not None, "gate_up_proj LoRA not detected"
+        assert down_lora is not None, "down_proj LoRA not detected"
+
+        # gate_up_proj: K=hidden, N=2*inter
+        E, r = config.num_experts, 4
+        gup_A, gup_B, gup_s = gup_lora
+        assert gup_A.shape == (E * r, config.hidden_size), (
+            f"gate_up_proj smoe_A: expected [r*E, K=hidden]={(E * r, config.hidden_size)}, "
+            f"got {gup_A.shape}"
+        )
+        assert gup_B.shape == (2 * config.intermediate_size, E * r), (
+            f"gate_up_proj smoe_B: expected [N=2*inter, r*E]="
+            f"{(2 * config.intermediate_size, E * r)}, got {gup_B.shape}"
+        )
+
+        # down_proj: K=inter, N=hidden
+        down_A, down_B, down_s = down_lora
+        assert down_A.shape == (E * r, config.intermediate_size), (
+            f"down_proj smoe_A: expected [r*E, K=inter]={(E * r, config.intermediate_size)}, "
+            f"got {down_A.shape}"
+        )
+        assert down_B.shape == (config.hidden_size, E * r), (
+            f"down_proj smoe_B: expected [N=hidden, r*E]={(config.hidden_size, E * r)}, "
+            f"got {down_B.shape}"
+        )
+
+    def test_unwrap_no_lora(self):
+        """Without peft, _unwrap_experts_lora returns no LoRA."""
+        config = make_olmoe_config(use_full=False)
+        moe = OlmoeSparseMoeBlock(config)
+        experts, gup_lora, down_lora = _unwrap_experts_lora(moe.experts)
+        assert gup_lora is None
+        assert down_lora is None
+        assert hasattr(experts, "gate_up_proj")
+
+    def test_unwrap_gate_lora(self):
+        """Test that _unwrap_gate_lora detects LoRA on the router gate."""
+        config = make_olmoe_config(use_full=False)
+        model = MinimalOLMoEModel(config)
+        r = 4
+        lora_config = LoraConfig(
+            r=r,
+            lora_alpha=16,
+            target_modules=[],
+            target_parameters=["gate.weight"],
+            bias="none",
+        )
+        peft_model = get_peft_model(model, lora_config)
+        base_moe = peft_model.base_model.model.moe
+
+        # Set non-zero LoRA weights (peft initializes lora_B to zeros)
+        with torch.no_grad():
+            base_moe.gate.lora_B["default"].weight.normal_(0, 0.01)
+
+        base_gate, gate_weight, gate_delta = _unwrap_gate_lora(base_moe.gate)
+
+        # Base gate should be the original router
+        assert hasattr(base_gate, "top_k")
+        assert hasattr(base_gate, "num_experts")
+        assert base_gate.top_k == config.num_experts_per_tok
+        assert base_gate.num_experts == config.num_experts
+
+        # Gate weight should be the base weight (delta returned separately)
+        assert gate_weight.shape == (config.num_experts, config.hidden_size)
+        torch.testing.assert_close(gate_weight, base_gate.weight)
+
+        # Delta should be non-zero (LoRA was applied)
+        assert gate_delta is not None
+        assert gate_delta.shape == (config.num_experts, config.hidden_size)
+        assert gate_delta.abs().max() > 0, "Gate LoRA delta should be non-zero"
+
+    def test_unwrap_gate_no_lora(self):
+        """Without peft, _unwrap_gate_lora returns the original gate."""
+        config = make_olmoe_config(use_full=False)
+        moe = OlmoeSparseMoeBlock(config)
+        base_gate, gate_weight, gate_delta = _unwrap_gate_lora(moe.gate)
+        assert base_gate is moe.gate
+        torch.testing.assert_close(gate_weight, moe.gate.weight)
+        assert gate_delta is None
+
+    def test_gate_lora_delta_matches_peft(self):
+        """Verify _unwrap_gate_lora computes the same delta as peft."""
+        config = make_olmoe_config(use_full=False)
+        model = MinimalOLMoEModel(config)
+        r = 4
+        lora_alpha = 16
+        scaling = lora_alpha / r
+        lora_config = LoraConfig(
+            r=r,
+            lora_alpha=lora_alpha,
+            target_modules=[],
+            target_parameters=["gate.weight"],
+            bias="none",
+        )
+        peft_model = get_peft_model(model, lora_config)
+        base_moe = peft_model.base_model.model.moe
+
+        # Our unwrapped weight + delta
+        _, gate_weight, gate_delta = _unwrap_gate_lora(base_moe.gate)
+
+        # Manually compute expected delta
+        lora_A = base_moe.gate.lora_A["default"].weight  # [r, hidden]
+        lora_B = base_moe.gate.lora_B["default"].weight  # [E, r]
+        base_weight = base_moe.gate.base_layer.weight  # [E, hidden]
+        expected_delta = scaling * (lora_B @ lora_A)
+
+        torch.testing.assert_close(gate_weight, base_weight)
+        torch.testing.assert_close(gate_delta, expected_delta)
+        # Combined should match the old behavior
+        torch.testing.assert_close(
+            gate_weight + gate_delta, base_weight + expected_delta
+        )
+
+
+# =============================================================================
+# Tests: Base forward equivalence (no LoRA)
+# =============================================================================
+
+
+@requires_cuda
+class TestOLMoEReferenceVsScatterMoE:
+    """Base forward equivalence: per-expert reference vs ScatterMoE kernels."""
+
+    def test_small(self):
+        self._run(use_full=False, M=16)
+
+    @pytest.mark.slow
+    def test_full(self):
+        self._run(use_full=True, M=32)
+
+    def _run(self, use_full, M):
+        from axolotl.integrations.kernels.libs.scattermoe_lora import (
+            flatten_sort_count,
+            parallel_linear,
+        )
+
+        config = make_olmoe_config(use_full=use_full)
+        torch.manual_seed(42)
+        moe = _init_expert_weights(OlmoeSparseMoeBlock(config)).cuda().float()
+        E, k = config.num_experts, config.num_experts_per_tok
+
+        x = torch.randn(1, M, config.hidden_size, device="cuda")
+        x_flat = x.view(-1, config.hidden_size)
+
+        with torch.no_grad():
+            # Shared routing for both paths
+            _, rw, sel = moe.gate(x_flat)
+            sei, ssi, eo = flatten_sort_count(sel, num_experts=E)
+
+            # Per-expert reference
+            ref_out = _reference_moe_forward(
+                x_flat,
+                moe.experts.gate_up_proj,
+                moe.experts.down_proj,
+                moe.experts.act_fn,
+                sel,
+                rw,
+                E,
+            ).view(1, M, config.hidden_size)
+
+            # ScatterMoE kernel path
+            gup = parallel_linear(
+                x_flat,
+                moe.experts.gate_up_proj.transpose(2, 1),
+                k,
+                sei,
+                ssi,
+                eo,
+                grouped_in=False,
+                grouped_out=True,
+            )
+            g, u = gup.chunk(2, dim=-1)
+            h = moe.experts.act_fn(g) * u
+
+            smoe_out = parallel_linear(
+                h,
+                moe.experts.down_proj.transpose(2, 1),
+                1,
+                sei,
+                ssi,
+                eo,
+                grouped_in=True,
+                grouped_out=False,
+                gates=rw,
+            ).view(1, M, config.hidden_size)
+
+        torch.testing.assert_close(smoe_out, ref_out, atol=1e-3, rtol=1e-3)
+
+
+# =============================================================================
+# Tests: LoRA forward equivalence (peft vs scattermoe fused)
+# =============================================================================
+
+
+@requires_cuda
+class TestOLMoEPeftLoRAForward:
+    """Fused LoRA forward: peft merged-weight vs scattermoe_lora kernel."""
+
+    def test_small(self):
+        self._run(use_full=False, M=16, r=4)
+
+    @pytest.mark.slow
+    def test_full(self):
+        self._run(use_full=True, M=32, r=8)
+
+    def _run(self, use_full, M, r):
+        from axolotl.integrations.kernels.libs.scattermoe_lora import (
+            flatten_sort_count,
+            parallel_linear_lora,
+        )
+
+        config = make_olmoe_config(use_full=use_full)
+        E, k = config.num_experts, config.num_experts_per_tok
+        lora_alpha = 16
+        scaling = lora_alpha / r
+
+        # Create peft model
+        model = MinimalOLMoEModel(config).cuda().float()
+        lora_config = LoraConfig(
+            r=r,
+            lora_alpha=lora_alpha,
+            target_modules=[],
+            target_parameters=["experts.gate_up_proj", "experts.down_proj"],
+            bias="none",
+        )
+        peft_model = get_peft_model(model, lora_config)
+
+        torch.manual_seed(42)
+        x = torch.randn(1, M, config.hidden_size, device="cuda")
+
+        # peft forward
+        with torch.no_grad():
+            peft_out = peft_model(x)
+
+        # Extract base weights and LoRA weights
+        base_moe = peft_model.base_model.model.moe
+        base_experts = base_moe.experts.base_layer.base_layer
+        gate_up_proj = base_experts.gate_up_proj
+        down_proj = base_experts.down_proj
+        act_fn = base_experts.act_fn
+
+        # gate_up_proj LoRA
+        gup_w = base_moe.experts.base_layer
+        peft_gup_A = gup_w.lora_A["default"].weight.detach()
+        peft_gup_B = gup_w.lora_B["default"].weight.detach()
+        smoe_gup_A, smoe_gup_B = peft_gate_up_lora_to_scattermoe(
+            peft_gup_A, peft_gup_B, E, r
+        )
+
+        # down_proj LoRA
+        down_w = base_moe.experts
+        peft_down_A = down_w.lora_A["default"].weight.detach()
+        peft_down_B = down_w.lora_B["default"].weight.detach()
+        smoe_down_A, smoe_down_B = peft_lora_to_scattermoe(
+            peft_down_A, peft_down_B, E, r
+        )
+
+        # ScatterMoE fused forward -- gate is NOT peft-wrapped, access directly
+        x_flat = x.view(-1, config.hidden_size)
+
+        with torch.no_grad():
+            _, rw, sel = base_moe.gate(x_flat)
+            sei, ssi, eo = flatten_sort_count(sel, num_experts=E)
+
+            gup = parallel_linear_lora(
+                x_flat,
+                gate_up_proj.transpose(2, 1),
+                k,
+                sei,
+                ssi,
+                eo,
+                lora_A=smoe_gup_A,
+                lora_B=smoe_gup_B,
+                scaling=scaling,
+                grouped_in=False,
+                grouped_out=True,
+            )
+            g, u = gup.chunk(2, dim=-1)
+            h = act_fn(g) * u
+
+            smoe_out = parallel_linear_lora(
+                h,
+                down_proj.transpose(2, 1),
+                1,
+                sei,
+                ssi,
+                eo,
+                lora_A=smoe_down_A,
+                lora_B=smoe_down_B,
+                scaling=scaling,
+                grouped_in=True,
+                grouped_out=False,
+                gates=rw,
+            ).view(1, M, config.hidden_size)
+
+        torch.testing.assert_close(smoe_out, peft_out, atol=5e-3, rtol=5e-3)
+
+
+# =============================================================================
+# Tests: Backward gradient correctness
+# =============================================================================
+
+
+@requires_cuda
+class TestOLMoEPeftLoRABackward:
+    """Backward gradients through scattermoe_lora vs pure-PyTorch reference."""
+
+    def test_small(self):
+        self._run(use_full=False, M=16, r=4)
+
+    def _run(self, use_full, M, r):
+        from axolotl.integrations.kernels.libs.scattermoe_lora import (
+            flatten_sort_count,
+            parallel_linear_lora,
+        )
+
+        config = make_olmoe_config(use_full=use_full)
+        E, k = config.num_experts, config.num_experts_per_tok
+        lora_alpha = 16
+        scaling = lora_alpha / r
+
+        torch.manual_seed(42)
+        moe = _init_expert_weights(OlmoeSparseMoeBlock(config)).cuda().float()
+        x = torch.randn(1, M, config.hidden_size, device="cuda")
+        x_flat = x.view(-1, config.hidden_size)
+        gate_up_proj = moe.experts.gate_up_proj
+        down_proj = moe.experts.down_proj
+
+        # Create LoRA weights in scattermoe layout directly
+        gup_A = torch.randn(r * E, config.hidden_size, device="cuda") * 0.01
+        gup_B = torch.randn(2 * config.intermediate_size, r * E, device="cuda") * 0.01
+        down_A = torch.randn(r * E, config.intermediate_size, device="cuda") * 0.01
+        down_B = torch.randn(config.hidden_size, r * E, device="cuda") * 0.01
+
+        rw, sel = _get_routing(moe, x)
+        sei, ssi, eo = flatten_sort_count(sel, num_experts=E)
+
+        # --- Reference ---
+        gup_delta = _compute_delta_from_scattermoe_lora(
+            gup_A, gup_B, scaling, E, r, gate_up_proj.shape
+        )
+        down_delta = _compute_delta_from_scattermoe_lora(
+            down_A, down_B, scaling, E, r, down_proj.shape
+        )
+
+        x_ref = x_flat.clone().detach().requires_grad_(True)
+        ref_out = _reference_moe_forward_with_lora(
+            x_ref,
+            gate_up_proj,
+            down_proj,
+            moe.experts.act_fn,
+            sel,
+            rw,
+            E,
+            gup_delta,
+            down_delta,
+        )
+        ref_out.sum().backward()
+
+        # --- ScatterMoE fused path ---
+        x_smoe = x_flat.clone().detach().requires_grad_(True)
+        gup_A_s = gup_A.clone().requires_grad_(True)
+        gup_B_s = gup_B.clone().requires_grad_(True)
+        down_A_s = down_A.clone().requires_grad_(True)
+        down_B_s = down_B.clone().requires_grad_(True)
+
+        gup_out = parallel_linear_lora(
+            x_smoe,
+            gate_up_proj.transpose(2, 1),
+            k,
+            sei,
+            ssi,
+            eo,
+            lora_A=gup_A_s,
+            lora_B=gup_B_s,
+            scaling=scaling,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        g, u = gup_out.chunk(2, dim=-1)
+        h = moe.experts.act_fn(g) * u
+
+        smoe_out = parallel_linear_lora(
+            h,
+            down_proj.transpose(2, 1),
+            1,
+            sei,
+            ssi,
+            eo,
+            lora_A=down_A_s,
+            lora_B=down_B_s,
+            scaling=scaling,
+            grouped_in=True,
+            grouped_out=False,
+            gates=rw,
+        )
+        smoe_out.sum().backward()
+
+        torch.testing.assert_close(
+            smoe_out.detach(),
+            ref_out.detach(),
+            atol=5e-3,
+            rtol=5e-3,
+        )
+        torch.testing.assert_close(
+            x_smoe.grad,
+            x_ref.grad,
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+
+# =============================================================================
+# Tests: kernelize() integration via LocalLayerRepository
+# =============================================================================
+
+
+@requires_cuda
+class TestKernelizeIntegration:
+    """Test the HF kernels library integration with LocalLayerRepository."""
+
+    @staticmethod
+    def _get_kernelize_imports():
+        """Import kernels library components, skip if not available."""
+        try:
+            from kernels import (
+                LocalLayerRepository,
+                Mode,
+                kernelize,
+                register_kernel_mapping,
+                replace_kernel_forward_from_hub,
+            )
+
+            return (
+                LocalLayerRepository,
+                Mode,
+                register_kernel_mapping,
+                replace_kernel_forward_from_hub,
+                kernelize,
+            )
+        except ImportError:
+            pytest.skip("kernels library not installed")
+
+    @staticmethod
+    def _get_repo_path():
+        """Get the path to scattermoe_lora within axolotl's plugin."""
+        return (
+            Path(__file__).parent.parent.parent.parent
+            / "src"
+            / "axolotl"
+            / "integrations"
+            / "kernels"
+            / "libs"
+            / "scattermoe_lora"
+        )
+
+    def _setup_kernels(
+        self,
+        LocalLayerRepository,
+        Mode,
+        register_kernel_mapping,
+        replace_kernel_forward_from_hub,
+    ):
+        """Register kernel mapping for tests."""
+        repo_path = self._get_repo_path()
+        local_repo = make_local_layer_repository(
+            LocalLayerRepository,
+            repo_path,
+            "HFScatterMoEGatedMLP",
+        )
+
+        replace_kernel_forward_from_hub(
+            OlmoeSparseMoeBlock, "HFScatterMoEParallelExperts"
+        )
+        register_kernel_mapping(
+            {
+                "HFScatterMoEParallelExperts": {
+                    "cuda": {
+                        Mode.TRAINING: local_repo,
+                        Mode.INFERENCE: local_repo,
+                    },
+                }
+            }
+        )
+
+    def test_base_forward_via_kernelize(self):
+        """Kernelized OlmoeSparseMoeBlock (no LoRA) matches per-expert reference."""
+        (
+            LocalLayerRepository,
+            Mode,
+            register_kernel_mapping,
+            replace_kernel_forward_from_hub,
+            kernelize,
+        ) = self._get_kernelize_imports()
+
+        config = make_olmoe_config(use_full=False)
+        E = config.num_experts
+
+        # Create model
+        torch.manual_seed(42)
+        moe = _init_expert_weights(OlmoeSparseMoeBlock(config)).cuda().float()
+        x = torch.randn(1, 8, config.hidden_size, device="cuda")
+        x_flat = x.view(-1, config.hidden_size)
+
+        # Compute reference BEFORE kernelizing
+        with torch.no_grad():
+            _, rw, sel = moe.gate(x_flat)
+            ref_out = _reference_moe_forward(
+                x_flat,
+                moe.experts.gate_up_proj,
+                moe.experts.down_proj,
+                moe.experts.act_fn,
+                sel,
+                rw,
+                E,
+            ).view(1, 8, config.hidden_size)
+
+        # Set up kernel mapping
+        self._setup_kernels(
+            LocalLayerRepository,
+            Mode,
+            register_kernel_mapping,
+            replace_kernel_forward_from_hub,
+        )
+
+        # Kernelize the model
+        kernelize(moe, mode=Mode.TRAINING, device="cuda")
+
+        # Forward through kernelized model
+        with torch.no_grad():
+            kern_out = moe(x)
+
+        torch.testing.assert_close(kern_out, ref_out, atol=1e-3, rtol=1e-3)
+
+    def test_lora_forward_via_kernelize(self):
+        """Kernelized OlmoeSparseMoeBlock with peft LoRA matches reference."""
+        (
+            LocalLayerRepository,
+            Mode,
+            register_kernel_mapping,
+            replace_kernel_forward_from_hub,
+            kernelize,
+        ) = self._get_kernelize_imports()
+
+        config = make_olmoe_config(use_full=False)
+        r = 4
+
+        # Create peft model
+        torch.manual_seed(42)
+        model = MinimalOLMoEModel(config).cuda().float()
+        lora_config = LoraConfig(
+            r=r,
+            lora_alpha=16,
+            target_modules=[],
+            target_parameters=["experts.gate_up_proj", "experts.down_proj"],
+            bias="none",
+        )
+        peft_model = get_peft_model(model, lora_config)
+
+        x = torch.randn(1, 8, config.hidden_size, device="cuda")
+
+        # Reference: peft's own forward (uses _activate_lora context manager)
+        with torch.no_grad():
+            ref_out = peft_model(x)
+
+        # Set up kernel mapping
+        self._setup_kernels(
+            LocalLayerRepository,
+            Mode,
+            register_kernel_mapping,
+            replace_kernel_forward_from_hub,
+        )
+
+        # Kernelize the MoE block inside the peft model
+        base_moe = peft_model.base_model.model.moe
+        kernelize(base_moe, mode=Mode.TRAINING, device="cuda")
+
+        # Forward through kernelized peft model
+        with torch.no_grad():
+            kern_out = peft_model(x)
+
+        torch.testing.assert_close(kern_out, ref_out, atol=5e-3, rtol=5e-3)
+
+    def test_gate_lora_forward_via_kernelize(self):
+        """Kernelized forward with gate LoRA matches peft reference."""
+        (
+            LocalLayerRepository,
+            Mode,
+            register_kernel_mapping,
+            replace_kernel_forward_from_hub,
+            kernelize,
+        ) = self._get_kernelize_imports()
+
+        config = make_olmoe_config(use_full=False)
+        r = 4
+
+        # Create peft model with gate + experts LoRA
+        torch.manual_seed(42)
+        model = MinimalOLMoEModel(config).cuda().float()
+        lora_config = LoraConfig(
+            r=r,
+            lora_alpha=16,
+            target_modules=[],
+            target_parameters=[
+                "gate.weight",
+                "experts.gate_up_proj",
+                "experts.down_proj",
+            ],
+            bias="none",
+        )
+        peft_model = get_peft_model(model, lora_config)
+
+        x = torch.randn(1, 8, config.hidden_size, device="cuda")
+
+        # Reference: peft's own forward
+        with torch.no_grad():
+            ref_out = peft_model(x)
+
+        # Set up kernel mapping
+        self._setup_kernels(
+            LocalLayerRepository,
+            Mode,
+            register_kernel_mapping,
+            replace_kernel_forward_from_hub,
+        )
+
+        # Kernelize the MoE block inside the peft model
+        base_moe = peft_model.base_model.model.moe
+        kernelize(base_moe, mode=Mode.TRAINING, device="cuda")
+
+        # Forward through kernelized peft model
+        with torch.no_grad():
+            kern_out = peft_model(x)
+
+        torch.testing.assert_close(kern_out, ref_out, atol=5e-3, rtol=5e-3)
+
+
+# =============================================================================
+# Tests: Shared expert handling
+# =============================================================================
+
+
+class TestSharedExpertHandling:
+    """Test that HFScatterMoEGatedMLP.forward handles shared experts."""
+
+    @staticmethod
+    def _make_shared_expert_block(config):
+        """Create an OlmoeSparseMoeBlock with a mock shared expert attached."""
+        moe = OlmoeSparseMoeBlock(config)
+        _init_expert_weights(moe)
+
+        hidden = config.hidden_size
+        inter = config.intermediate_size
+
+        # Attach a simple shared expert MLP (mimics Qwen2MoE structure)
+        class SharedExpertMLP(nn.Module):
+            def __init__(self, hidden_size, intermediate_size):
+                super().__init__()
+                self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+                self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+                self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+                self.act_fn = nn.SiLU()
+
+            def forward(self, x):
+                return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+        moe.shared_expert = SharedExpertMLP(hidden, inter)
+        moe.shared_expert_gate = nn.Linear(hidden, 1, bias=False)
+
+        return moe
+
+    def test_shared_expert_is_used(self):
+        """Verify shared expert output affects final result."""
+        config = make_olmoe_config(use_full=False)
+        moe = self._make_shared_expert_block(config)
+
+        # Compute reference without shared expert
+        torch.manual_seed(42)
+        x = torch.randn(1, 4, config.hidden_size)
+        x_flat = x.view(-1, config.hidden_size)
+
+        with torch.no_grad():
+            # Shared expert contribution
+            shared_out = moe.shared_expert(x_flat)
+            gate_val = F.sigmoid(moe.shared_expert_gate(x_flat))
+            shared_contribution = shared_out * gate_val
+
+        # Verify shared expert produces non-zero output
+        assert shared_contribution.abs().max() > 0
+
+    @requires_cuda
+    def test_shared_expert_forward_via_kernelize(self):
+        """Kernelized forward with shared expert matches manual reference."""
+        try:
+            from kernels import (
+                LocalLayerRepository,
+                Mode,
+                kernelize,
+                register_kernel_mapping,
+                replace_kernel_forward_from_hub,
+            )
+        except ImportError:
+            pytest.skip("kernels library not installed")
+
+        config = make_olmoe_config(use_full=False)
+        E = config.num_experts
+
+        torch.manual_seed(42)
+        moe = self._make_shared_expert_block(config).cuda().float()
+        x = torch.randn(1, 8, config.hidden_size, device="cuda")
+        x_flat = x.view(-1, config.hidden_size)
+
+        # Compute reference: per-expert + shared expert
+        with torch.no_grad():
+            _, rw, sel = moe.gate(x_flat)
+
+            expert_out = _reference_moe_forward(
+                x_flat,
+                moe.experts.gate_up_proj,
+                moe.experts.down_proj,
+                moe.experts.act_fn,
+                sel,
+                rw,
+                E,
+            )
+            shared_out = moe.shared_expert(x_flat)
+            gate_val = F.sigmoid(moe.shared_expert_gate(x_flat))
+            ref_out = (expert_out + shared_out * gate_val).view(
+                1, 8, config.hidden_size
+            )
+
+        # Kernelize
+        repo_path = (
+            Path(__file__).parent.parent.parent.parent
+            / "src"
+            / "axolotl"
+            / "integrations"
+            / "kernels"
+            / "libs"
+            / "scattermoe_lora"
+        )
+        local_repo = make_local_layer_repository(
+            LocalLayerRepository,
+            repo_path,
+            "HFScatterMoEGatedMLP",
+        )
+
+        replace_kernel_forward_from_hub(
+            OlmoeSparseMoeBlock, "HFScatterMoEParallelExperts"
+        )
+        register_kernel_mapping(
+            {
+                "HFScatterMoEParallelExperts": {
+                    "cuda": {
+                        Mode.TRAINING: local_repo,
+                        Mode.INFERENCE: local_repo,
+                    },
+                }
+            }
+        )
+
+        kernelize(moe, mode=Mode.TRAINING, device="cuda")
+
+        with torch.no_grad():
+            kern_out = moe(x)
+
+        torch.testing.assert_close(kern_out, ref_out, atol=1e-3, rtol=1e-3)
diff --git a/tests/e2e/integrations/test_sonicmoe.py b/tests/e2e/integrations/test_sonicmoe.py
new file mode 100644
index 0000000000..b74e570d02
--- /dev/null
+++ b/tests/e2e/integrations/test_sonicmoe.py
@@ -0,0 +1,216 @@
+"""End-to-end gradient and convergence tests for SonicMoE integration.
+
+Flow:
+
+    register_sonicmoe_experts()                # plug into ALL_EXPERTS_FUNCTIONS
+    config._experts_implementation = "sonicmoe"
+    model = AutoModelForCausalLM.from_config(config)   # transformers dispatches
+
+No weight interleaving needed (``concat_layout=True``).
+
+Requires:
+    - Hopper (sm_90) or Blackwell (sm_100+) GPU
+    - sonic-moe >= 0.1.2 installed from source
+    - transformers >= 5.8 with Qwen3MoE Experts class
+"""
+
+import importlib.util
+import math
+
+import pytest
+import torch
+
+
+def _is_hopper_or_newer() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    major, _ = torch.cuda.get_device_capability()
+    return major >= 9
+
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA GPU"),
+    pytest.mark.skipif(
+        not _is_hopper_or_newer(),
+        reason="SonicMoE requires Hopper (sm_90) or Blackwell (sm_100+)",
+    ),
+    pytest.mark.skipif(
+        importlib.util.find_spec("kernels") is None,
+        reason="HF `kernels` package not installed",
+    ),
+]
+
+
+def _create_tiny_qwen3_config(experts_implementation: str):
+    """Create a minimal Qwen3MoE config bound to the requested experts impl."""
+    from transformers import AutoConfig
+
+    config = AutoConfig.for_model("qwen3_moe")
+    config.hidden_size = 512
+    config.intermediate_size = 1024
+    config.moe_intermediate_size = 64
+    config.num_attention_heads = 16
+    config.num_key_value_heads = 2
+    config.head_dim = 32
+    config.num_hidden_layers = 2
+    config.num_experts = 8
+    config.num_experts_per_tok = 2
+    config.vocab_size = 1000
+    config.max_position_embeddings = 128
+    config.norm_topk_prob = True
+    config.torch_dtype = torch.bfloat16
+    config._experts_implementation = experts_implementation
+    return config
+
+
+def _build_model(experts_implementation: str):
+    from transformers import AutoModelForCausalLM
+
+    from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+        register_sonicmoe_experts,
+    )
+
+    register_sonicmoe_experts()
+    config = _create_tiny_qwen3_config(experts_implementation)
+    return AutoModelForCausalLM.from_config(config).cuda().bfloat16(), config
+
+
+class TestSonicMoEForwardCorrectness:
+    """SonicMoE-dispatched model produces output close to eager baseline."""
+
+    def test_forward_output_matches_eager(self):
+        input_ids = torch.randint(0, 1000, (1, 16), device="cuda")
+
+        eager_model, _ = _build_model("eager")
+        with torch.no_grad():
+            out_eager = eager_model(input_ids).logits
+
+        sonic_model, _ = _build_model("sonicmoe")
+        sonic_model.load_state_dict(eager_model.state_dict())
+
+        with torch.no_grad():
+            out_sonic = sonic_model(input_ids).logits
+
+        max_diff = (out_eager - out_sonic).abs().max().item()
+        assert torch.allclose(out_eager, out_sonic, atol=1e-1, rtol=1e-1), (
+            f"Output mismatch: max diff={max_diff:.6f}"
+        )
+
+
+class TestSonicMoEGradientCorrectness:
+    """Compare gradients between eager and SonicMoE-dispatched forward."""
+
+    def test_gradients_match(self):
+        input_ids = torch.randint(0, 1000, (1, 16), device="cuda")
+
+        eager_model, _ = _build_model("eager")
+        out_eager = eager_model(input_ids, labels=input_ids)
+        out_eager.loss.backward()
+        grads_eager = {
+            n: p.grad.float().clone()
+            for n, p in eager_model.named_parameters()
+            if p.grad is not None
+        }
+        loss_eager = out_eager.loss.item()
+
+        sonic_model, _ = _build_model("sonicmoe")
+        sonic_model.load_state_dict(eager_model.state_dict())
+        out_sonic = sonic_model(input_ids, labels=input_ids)
+        out_sonic.loss.backward()
+        grads_sonic = {
+            n: p.grad.float().clone()
+            for n, p in sonic_model.named_parameters()
+            if p.grad is not None
+        }
+        loss_sonic = out_sonic.loss.item()
+
+        assert abs(loss_eager - loss_sonic) < 0.5, (
+            f"Loss mismatch: eager={loss_eager:.4f}, sonic={loss_sonic:.4f}"
+        )
+
+        missing = set(grads_eager.keys()) - set(grads_sonic.keys())
+        assert not missing, f"Missing gradients in sonicmoe model: {missing}"
+
+        # bf16 + different GEMM backends can diverge; tolerate both rel >10% AND
+        # abs >1e-2 together.
+        mismatches = []
+        for name, g_eager in grads_eager.items():
+            g_sonic = grads_sonic[name]
+            max_diff = (g_eager - g_sonic).abs().max().item()
+            rel_diff = max_diff / (g_eager.abs().max().item() + 1e-8)
+            if rel_diff > 0.1 and max_diff > 1e-2:
+                mismatches.append(
+                    f"  {name}: max_abs_diff={max_diff:.6f}, rel_diff={rel_diff:.4f}"
+                )
+
+        assert not mismatches, (
+            "Gradient mismatches (rel_diff > 10% and abs_diff > 1e-2):\n"
+            + "\n".join(mismatches)
+        )
+
+    def test_router_weights_receive_gradients(self):
+        input_ids = torch.randint(0, 1000, (1, 16), device="cuda")
+        model, _ = _build_model("sonicmoe")
+        out = model(input_ids, labels=input_ids)
+        out.loss.backward()
+
+        gate_grads_found = False
+        for name, param in model.named_parameters():
+            if "gate" in name and "weight" in name:
+                gate_grads_found = True
+                assert param.grad is not None, f"No gradient for router: {name}"
+                assert param.grad.abs().max() > 0, f"Zero gradient for router: {name}"
+
+        assert gate_grads_found, "No gate.weight parameters found in model"
+
+
+class TestSonicMoETrainingConvergence:
+    """Verify loss decreases during training with SonicMoE."""
+
+    def test_loss_decreases(self):
+        input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+        model, _ = _build_model("sonicmoe")
+
+        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
+        losses = []
+
+        for step in range(30):
+            out = model(input_ids, labels=input_ids)
+            loss = out.loss
+            assert not math.isnan(loss.item()), f"NaN loss at step {step}"
+            assert not math.isinf(loss.item()), f"Inf loss at step {step}"
+            losses.append(loss.item())
+
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        assert losses[-1] < losses[0], (
+            f"Loss did not decrease: first={losses[0]:.4f}, last={losses[-1]:.4f}"
+        )
+
+    def test_expert_weights_update(self):
+        input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+        model, _ = _build_model("sonicmoe")
+
+        expert_weights_before = {
+            name: param.data.clone()
+            for name, param in model.named_parameters()
+            if "experts" in name
+        }
+        assert expert_weights_before, "No expert parameters found"
+
+        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
+        for _ in range(5):
+            out = model(input_ids, labels=input_ids)
+            out.loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        changed = sum(
+            1
+            for name, param in model.named_parameters()
+            if name in expert_weights_before
+            and not torch.equal(param.data, expert_weights_before[name])
+        )
+        assert changed > 0, "No expert weights changed after 5 training steps"
diff --git a/tests/e2e/integrations/test_sonicmoe_lora.py b/tests/e2e/integrations/test_sonicmoe_lora.py
new file mode 100644
index 0000000000..cc58f4dccd
--- /dev/null
+++ b/tests/e2e/integrations/test_sonicmoe_lora.py
@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""End-to-end tests for SonicMoE + LoRA.
+
+Flow:
+
+    register_sonicmoe_experts()                # plug into ALL_EXPERTS_FUNCTIONS
+    config._experts_implementation = "sonicmoe"
+    model = AutoModelForCausalLM.from_config(config)
+    model = get_peft_model(model, lora_config)   # PEFT wraps params/modules
+
+``sonicmoe_experts_forward_with_lora`` detects the PEFT wrappers and
+materializes ``W_eff = W + scaling * (B @ A)`` via :class:`MoELoRAMaterialize`,
+so adapters train through the CUTLASS kernels.
+
+Requires:
+    - Hopper (sm_90) or Blackwell (sm_100+) GPU
+    - sonic-moe >= 0.1.2 installed from source
+    - peft installed
+    - transformers >= 5.8 with Qwen3MoE Experts class
+"""
+
+import importlib.util
+import math
+
+import pytest
+import torch
+
+
+def _is_hopper_or_newer() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    major, _ = torch.cuda.get_device_capability()
+    return major >= 9
+
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA GPU"),
+    pytest.mark.skipif(
+        not _is_hopper_or_newer(),
+        reason="SonicMoE requires Hopper (sm_90) or Blackwell (sm_100+)",
+    ),
+    pytest.mark.skipif(
+        importlib.util.find_spec("kernels") is None,
+        reason="HF `kernels` package not installed",
+    ),
+    pytest.mark.skipif(
+        importlib.util.find_spec("peft") is None, reason="PEFT not installed"
+    ),
+]
+
+
+def _create_tiny_qwen3_config():
+    from transformers import AutoConfig
+
+    config = AutoConfig.for_model("qwen3_moe")
+    config.hidden_size = 512
+    config.intermediate_size = 1024
+    config.moe_intermediate_size = 64
+    config.num_attention_heads = 16
+    config.num_key_value_heads = 2
+    config.head_dim = 32
+    config.num_hidden_layers = 2
+    config.num_experts = 8
+    config.num_experts_per_tok = 2
+    config.vocab_size = 1000
+    config.max_position_embeddings = 128
+    config.norm_topk_prob = True
+    config.torch_dtype = torch.bfloat16
+    config._experts_implementation = "sonicmoe"
+    return config
+
+
+def _build_sonic_model():
+    from transformers import AutoModelForCausalLM
+
+    from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+        register_sonicmoe_experts,
+    )
+
+    register_sonicmoe_experts()
+    config = _create_tiny_qwen3_config()
+    return AutoModelForCausalLM.from_config(config).cuda().bfloat16()
+
+
+def _apply_lora(model, target_modules):
+    from peft import LoraConfig, get_peft_model
+
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=16,
+        target_modules=target_modules,
+        lora_dropout=0.0,
+        bias="none",
+    )
+    return get_peft_model(model, lora_config)
+
+
+class TestSonicMoELoRATraining:
+    """SonicMoE + LoRA on expert projections trains end-to-end."""
+
+    def test_loss_decreases(self):
+        input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+        model = _build_sonic_model()
+        model = _apply_lora(model, ["gate_up_proj", "down_proj"])
+
+        optimizer = torch.optim.AdamW(
+            [p for p in model.parameters() if p.requires_grad], lr=1e-3
+        )
+        losses = []
+        for step in range(30):
+            out = model(input_ids, labels=input_ids)
+            loss = out.loss
+            assert not math.isnan(loss.item()), f"NaN loss at step {step}"
+            assert not math.isinf(loss.item()), f"Inf loss at step {step}"
+            losses.append(loss.item())
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        assert losses[-1] < losses[0], (
+            f"Loss did not decrease: first={losses[0]:.4f}, last={losses[-1]:.4f}"
+        )
+
+    def test_base_weights_frozen(self):
+        input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+        model = _build_sonic_model()
+        model = _apply_lora(model, ["gate_up_proj", "down_proj"])
+
+        frozen_before = {
+            name: param.data.clone()
+            for name, param in model.named_parameters()
+            if not param.requires_grad
+        }
+
+        optimizer = torch.optim.AdamW(
+            [p for p in model.parameters() if p.requires_grad], lr=1e-3
+        )
+        for _ in range(5):
+            out = model(input_ids, labels=input_ids)
+            out.loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        for name, before in frozen_before.items():
+            after = dict(model.named_parameters())[name]
+            assert torch.equal(after.data, before), f"Frozen weight changed: {name}"
+
+    def test_lora_adapters_receive_gradients(self):
+        input_ids = torch.randint(0, 1000, (1, 16), device="cuda")
+        model = _build_sonic_model()
+        model = _apply_lora(model, ["gate_up_proj", "down_proj"])
+
+        out = model(input_ids, labels=input_ids)
+        out.loss.backward()
+
+        lora_grads_found = 0
+        for name, param in model.named_parameters():
+            if "lora_" in name and param.requires_grad:
+                assert param.grad is not None, f"No gradient for LoRA param: {name}"
+                assert param.grad.abs().max() > 0, (
+                    f"Zero gradient for LoRA param: {name}"
+                )
+                lora_grads_found += 1
+
+        assert lora_grads_found > 0, "No LoRA parameters found with gradients"
+
+    def test_lora_adapters_update(self):
+        input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+        model = _build_sonic_model()
+        model = _apply_lora(model, ["gate_up_proj", "down_proj"])
+
+        lora_before = {
+            name: param.data.clone()
+            for name, param in model.named_parameters()
+            if "lora_" in name and param.requires_grad
+        }
+        assert lora_before, "No LoRA parameters found"
+
+        optimizer = torch.optim.AdamW(
+            [p for p in model.parameters() if p.requires_grad], lr=1e-3
+        )
+        for _ in range(5):
+            out = model(input_ids, labels=input_ids)
+            out.loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        changed = sum(
+            1
+            for name, param in model.named_parameters()
+            if name in lora_before and not torch.equal(param.data, lora_before[name])
+        )
+        assert changed > 0, "No LoRA weights changed after 5 training steps"
+
+
+class TestSonicMoEGateOnlyLoRA:
+    """LoRA only on the router (gate) — expert path takes the no-LoRA fast path."""
+
+    def test_gate_only_lora_loss_decreases(self):
+        input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+        model = _build_sonic_model()
+        model = _apply_lora(model, ["gate"])
+
+        optimizer = torch.optim.AdamW(
+            [p for p in model.parameters() if p.requires_grad], lr=1e-3
+        )
+        losses = []
+        for step in range(20):
+            out = model(input_ids, labels=input_ids)
+            loss = out.loss
+            assert not math.isnan(loss.item()), f"NaN loss at step {step}"
+            assert not math.isinf(loss.item()), f"Inf loss at step {step}"
+            losses.append(loss.item())
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        assert losses[-1] < losses[0], (
+            f"Loss did not decrease: first={losses[0]:.4f}, last={losses[-1]:.4f}"
+        )
+
+
+class TestSonicMoENoLoRARegression:
+    """Full fine-tuning (no PEFT) still works through the registered forward."""
+
+    def test_no_lora_loss_decreases(self):
+        input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+        model = _build_sonic_model()
+
+        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
+        losses = []
+        for step in range(20):
+            out = model(input_ids, labels=input_ids)
+            loss = out.loss
+            assert not math.isnan(loss.item()), f"NaN loss at step {step}"
+            assert not math.isinf(loss.item()), f"Inf loss at step {step}"
+            losses.append(loss.item())
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        assert losses[-1] < losses[0], (
+            f"Loss did not decrease: first={losses[0]:.4f}, last={losses[-1]:.4f}"
+        )
diff --git a/tests/e2e/kernels/test_geglu.py b/tests/e2e/kernels/test_geglu.py
new file mode 100644
index 0000000000..78ba74c0e8
--- /dev/null
+++ b/tests/e2e/kernels/test_geglu.py
@@ -0,0 +1,90 @@
+"""Tests for GEGLU activation function Triton kernels."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from axolotl.kernels.geglu import geglu_backward, geglu_forward
+
+
+def test_geglu_forward_shape():
+    """Test that GEGLU forward pass preserves expected shapes."""
+    batch, seq_len, hidden_dim = 2, 3, 64
+    gate = torch.randn(batch, seq_len, hidden_dim, device="cuda")
+    up = torch.randn(batch, seq_len, hidden_dim, device="cuda")
+
+    out = geglu_forward(gate, up)
+    assert out.shape == (batch, seq_len, hidden_dim)
+    assert out.dtype == gate.dtype
+    assert out.device == gate.device
+
+
+@pytest.mark.flaky(retries=1, delay=5)
+@pytest.mark.parametrize(
+    "torch_seed",
+    [0, 42],
+)
+def test_geglu_forward_values(torch_seed):
+    """Test GEGLU forward pass matches PyTorch reference implementation."""
+    torch.manual_seed(torch_seed)
+
+    gate = torch.randn(2, 3, 64, device="cuda")
+    up = torch.randn(2, 3, 64, device="cuda")
+
+    # Custom implementation
+    triton_out = geglu_forward(gate.clone(), up.clone())
+
+    # PyTorch reference
+    torch_out = F.gelu(gate) * up
+
+    assert torch.allclose(triton_out, torch_out, rtol=1e-3)
+
+
+@pytest.mark.flaky(retries=1, delay=5)
+@pytest.mark.parametrize(
+    "torch_seed",
+    [0, 42],
+)
+def test_geglu_backward(torch_seed):
+    """Test GEGLU backward pass matches PyTorch autograd."""
+    torch.manual_seed(torch_seed)
+
+    gate = torch.randn(2, 3, 64, device="cuda", requires_grad=True)
+    up = torch.randn(2, 3, 64, device="cuda", requires_grad=True)
+    grad_output = torch.randn(2, 3, 64, device="cuda")
+
+    # PyTorch reference - compute intermediates
+    gelu_gate = F.gelu(gate)
+    torch_out = gelu_gate * up
+    torch_out.backward(grad_output)
+
+    # Custom backward pass
+    gate_clone = gate.clone().detach()
+    up_clone = up.clone().detach()
+    grad_output_clone = grad_output.clone()
+
+    h, grad_gate, grad_up = geglu_backward(grad_output_clone, gate_clone, up_clone)
+
+    # Compare outputs and gradients
+    assert torch.allclose(h, torch_out, rtol=1e-3)
+    assert torch.allclose(grad_gate, gate.grad, rtol=1e-3)
+    assert torch.allclose(grad_up, up.grad, rtol=1e-3)
+
+
+def test_geglu_inplace_preservation():
+    """Test that GEGLU backward doesn't modify original tensors unexpectedly."""
+    gate = torch.randn(2, 3, 64, device="cuda")
+    up = torch.randn(2, 3, 64, device="cuda")
+    grad_output = torch.randn(2, 3, 64, device="cuda")
+
+    gate_copy = gate.clone()
+    up_copy = up.clone()
+    grad_copy = grad_output.clone()
+
+    geglu_backward(grad_output, gate, up)
+
+    assert not torch.equal(gate, gate_copy), "Gate should be modified in-place"
+    assert not torch.equal(up, up_copy), "Up should be modified in-place"
+    assert not torch.equal(grad_output, grad_copy), (
+        "Grad output should be modified in-place"
+    )
diff --git a/tests/e2e/kernels/test_lora.py b/tests/e2e/kernels/test_lora.py
new file mode 100644
index 0000000000..5df1bc3be3
--- /dev/null
+++ b/tests/e2e/kernels/test_lora.py
@@ -0,0 +1,670 @@
+"""Tests for LoRA custom autograd."""
+
+import pytest
+import torch
+from bitsandbytes.functional import QuantState
+from torch import nn
+
+from axolotl.kernels.geglu import geglu_backward, geglu_forward
+from axolotl.kernels.lora import (
+    LoRA_MLP,
+    LoRA_O,
+    LoRA_QK,
+    LoRA_QKV,
+    apply_lora_mlp_geglu,
+    apply_lora_mlp_swiglu,
+    get_lora_parameters,
+    matmul_lora,
+)
+from axolotl.kernels.swiglu import swiglu_backward, swiglu_forward
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.fixture
+def mock_quantstate():
+    """Creates a mock QuantState for testing"""
+    shape = (64, 64)
+    n_blocks = shape[0]  # Assuming blockwise quantization along first dimension
+
+    # Create nested state first
+    nested_state = QuantState(
+        absmax=torch.ones(n_blocks, device="cuda"),  # One value per block
+        shape=shape,
+        code=torch.randint(0, 15, shape, device="cuda"),  # NF4 range is 0-15
+        dtype=torch.float16,
+        blocksize=64,
+        quant_type="nf4",
+        offset=None,
+        state2=None,
+    )
+
+    # Create main state with nested state
+    return QuantState(
+        absmax=torch.ones(n_blocks, device="cuda"),
+        shape=shape,
+        code=torch.randint(0, 15, shape, device="cuda"),
+        dtype=torch.float16,
+        blocksize=64,
+        quant_type="nf4",
+        offset=torch.zeros(n_blocks, dtype=torch.int32, device="cuda"),
+        state2=nested_state,
+    )
+
+
+@pytest.fixture
+def sample_tensors():
+    """Creates sample tensors for testing"""
+    torch.manual_seed(42)
+    batch_size, seq_len, hidden_dim = 2, 3, 64
+    rank = 8
+    out_dim = hidden_dim
+
+    return {
+        "X": torch.randn(
+            batch_size, seq_len, hidden_dim, device="cuda", dtype=torch.float16
+        ),
+        "W": torch.randn(out_dim, hidden_dim, device="cuda", dtype=torch.float16),
+        "b": torch.randn(out_dim, device="cuda", dtype=torch.float16),
+        "scale": 0.5,
+        "shapes": {
+            "batch": batch_size,
+            "seq": seq_len,
+            "hidden": hidden_dim,
+            "out": out_dim,
+            "rank": rank,
+        },
+    }
+
+
+@pytest.fixture
+def mock_proj():
+    """Creates a mock projection module for testing."""
+
+    class MockProj(nn.Module):
+        """Mock projection class."""
+
+        def __init__(self, in_features=64, out_features=128, rank=8):
+            super().__init__()
+            self.base_layer = nn.Linear(in_features, out_features)
+            self.base_layer.to("cuda")
+            self.lora_A = nn.ModuleDict(
+                {"default": nn.Linear(in_features, rank, bias=False).to("cuda")}
+            )
+            self.lora_B = nn.ModuleDict(
+                {"default": nn.Linear(rank, out_features, bias=False).to("cuda")}
+            )
+            self.scaling = {"default": 0.5}
+            self.active_adapter = "default"
+            self.disable_adapters = False
+            self.merged = False
+
+    return MockProj()
+
+
+def test_get_lora_parameters(mock_proj):
+    """Tests get_lora_parameters function"""
+    # Test with LoRA enabled
+    W, b, _, A, B, s, *_ = get_lora_parameters(mock_proj)
+
+    assert isinstance(W, torch.Tensor)
+    assert W.shape == (128, 64)
+    assert b.shape == (128,)
+    assert A.shape == (8, 64)
+    assert B.shape == (128, 8)
+    assert s == 0.5
+
+    # Test with LoRA disabled
+    mock_proj.disable_adapters = True
+    W, b, _, A, B, s, *_ = get_lora_parameters(mock_proj)
+    assert A is None and B is None and s is None
+
+    # Test with merged state
+    mock_proj.disable_adapters = False
+    mock_proj.merged = True
+    W, b, _, A, B, s, *_ = get_lora_parameters(mock_proj)
+    assert A is None and B is None and s is None
+
+
+def test_matmul_lora(sample_tensors):
+    """Tests matmul_lora function"""
+    X = sample_tensors["X"]
+    W = sample_tensors["W"]
+    b = sample_tensors["b"]
+    scale = sample_tensors["scale"]
+
+    shapes = sample_tensors["shapes"]
+    hidden_dim = shapes["hidden"]
+    out_dim = shapes["out"]
+    rank = shapes["rank"]
+
+    A = torch.randn(rank, hidden_dim, device="cuda", dtype=torch.float16)
+    B = torch.randn(out_dim, rank, device="cuda", dtype=torch.float16)
+
+    # Test base matmul
+    out1 = matmul_lora(X, W, b, None, None, None, None)
+    matmul = torch.matmul(X, W.t())
+    expected1 = matmul + b
+    assert torch.allclose(out1, expected1, rtol=1e-3)
+
+    # Test with LoRA. matmul_lora fuses the add via in-place addmm_ (fp32 accumulate,
+    # no [M, out] LoRA temp), so compare against an fp32 reference rather than a fp16
+    # one whose intermediate rounding the kernel intentionally no longer mirrors.
+    out2 = matmul_lora(X, W, b, None, A, B, scale)
+    expected2 = (
+        X.float() @ W.float().t()
+        + scale * (X.float() @ A.float().t()) @ B.float().t()
+        + b.float()
+    )
+    assert torch.allclose(out2.float(), expected2, rtol=2e-3, atol=2e-2)
+
+    # Test 3D input reshaping
+    X_3d = X.clone()
+    out3 = matmul_lora(X_3d, W, b, None, A, B, scale)
+    assert out3.shape == (X.shape[0], X.shape[1], W.shape[0])
+
+
+@pytest.mark.parametrize(
+    "activation_forward,activation_backward",
+    [(swiglu_forward, swiglu_backward), (geglu_forward, geglu_backward)],
+)
+def test_lora_mlp_direct(sample_tensors, activation_forward, activation_backward):
+    """Tests LoRA_MLP directly with different activation functions"""
+    X = sample_tensors["X"]
+    shapes = sample_tensors["shapes"]
+    hidden_dim = shapes["hidden"]
+    out_dim = shapes["out"]
+
+    # Create linear layers
+    gate_proj = nn.Linear(hidden_dim, out_dim).to(device="cuda", dtype=torch.float16)
+    up_proj = nn.Linear(hidden_dim, out_dim).to(device="cuda", dtype=torch.float16)
+    down_proj = nn.Linear(out_dim, hidden_dim).to(device="cuda", dtype=torch.float16)
+
+    # Test SwiGLU path
+    X.requires_grad = True
+    output = LoRA_MLP.apply(
+        X,
+        None,  # X_drop
+        gate_proj.weight,
+        gate_proj.bias,
+        None,  # gate_quant
+        None,  # gate_A
+        None,  # gate_B
+        None,  # gate_scale
+        None,  # gate_lora_bias
+        None,  # gate_magnitude
+        up_proj.weight,
+        up_proj.bias,
+        None,  # up_quant
+        None,  # up_A
+        None,  # up_B
+        None,  # up_scale
+        None,  # up_lora_bias
+        None,  # up_magnitude
+        down_proj.weight,
+        down_proj.bias,
+        None,  # down_quant
+        None,  # down_A
+        None,  # down_B
+        None,  # down_scale
+        None,  # down_lora_bias
+        None,  # down_magnitude
+        activation_forward,
+        activation_backward,
+        True,  # inplace
+    )
+
+    assert output.shape == X.shape
+    assert not torch.isnan(output).any()
+
+    # Test backward pass
+    loss = output.sum()
+    loss.backward()
+    assert X.grad is not None
+    assert not torch.isnan(X.grad).any()
+
+
+@pytest.mark.parametrize(
+    "activation_forward,activation_backward",
+    [(swiglu_forward, swiglu_backward), (geglu_forward, geglu_backward)],
+)
+def test_lora_mlp_with_adapters(
+    sample_tensors, activation_forward, activation_backward
+):
+    """Tests LoRA_MLP with LoRA adapters"""
+    X = sample_tensors["X"]
+    shapes = sample_tensors["shapes"]
+    hidden_dim = shapes["hidden"]
+    out_dim = shapes["out"]
+    rank = shapes["rank"]
+
+    # Create LoRA components
+    gate_A = torch.randn(rank, hidden_dim, device="cuda", dtype=torch.float16)
+    gate_B = torch.randn(out_dim, rank, device="cuda", dtype=torch.float16)
+    up_A = torch.randn(rank, hidden_dim, device="cuda", dtype=torch.float16)
+    up_B = torch.randn(out_dim, rank, device="cuda", dtype=torch.float16)
+    down_A = torch.randn(rank, out_dim, device="cuda", dtype=torch.float16)
+    down_B = torch.randn(hidden_dim, rank, device="cuda", dtype=torch.float16)
+    scale = 0.5
+
+    gate_proj = nn.Linear(hidden_dim, out_dim).to(device="cuda", dtype=torch.float16)
+    up_proj = nn.Linear(hidden_dim, out_dim).to(device="cuda", dtype=torch.float16)
+    down_proj = nn.Linear(out_dim, hidden_dim).to(device="cuda", dtype=torch.float16)
+
+    X.requires_grad = True
+    gate_A.requires_grad = True
+    gate_B.requires_grad = True
+    up_A.requires_grad = True
+    up_B.requires_grad = True
+    down_A.requires_grad = True
+    down_B.requires_grad = True
+
+    # Forward pass with adapters
+    output = LoRA_MLP.apply(
+        X,
+        None,  # X_drop
+        gate_proj.weight,
+        gate_proj.bias,
+        None,
+        gate_A,
+        gate_B,
+        scale,
+        None,  # gate_lora_bias
+        None,  # gate_magnitude
+        up_proj.weight,
+        up_proj.bias,
+        None,
+        up_A,
+        up_B,
+        scale,
+        None,  # up_lora_bias
+        None,  # up_magnitude
+        down_proj.weight,
+        down_proj.bias,
+        None,
+        down_A,
+        down_B,
+        scale,
+        None,  # down_lora_bias
+        None,  # down_magnitude
+        activation_forward,
+        activation_backward,
+        True,
+    )
+
+    assert output.shape == X.shape
+    assert not torch.isnan(output).any()
+
+    # Test backward pass
+    loss = output.sum()
+    loss.backward()
+
+    # Check all gradients
+    assert X.grad is not None
+    assert gate_A.grad is not None
+    assert gate_B.grad is not None
+    assert up_A.grad is not None
+    assert up_B.grad is not None
+    assert down_A.grad is not None
+    assert down_B.grad is not None
+
+    assert not torch.isnan(X.grad).any()
+    assert not torch.isnan(gate_A.grad).any()
+    assert not torch.isnan(gate_B.grad).any()
+    assert not torch.isnan(up_A.grad).any()
+    assert not torch.isnan(up_B.grad).any()
+    assert not torch.isnan(down_A.grad).any()
+    assert not torch.isnan(down_B.grad).any()
+
+
+def test_lora_qkv(sample_tensors):
+    """Tests LoRA QKV implementation with and without adapters"""
+    X = sample_tensors["X"]
+    shapes = sample_tensors["shapes"]
+    hidden_dim = shapes["hidden"]
+    rank = shapes["rank"]
+
+    # Create base weights
+    q_weight = torch.randn(hidden_dim, hidden_dim, device="cuda", dtype=torch.float16)
+    k_weight = torch.randn(hidden_dim, hidden_dim, device="cuda", dtype=torch.float16)
+    v_weight = torch.randn(hidden_dim, hidden_dim, device="cuda", dtype=torch.float16)
+
+    # Create LoRA matrices
+    q_A = torch.randn(
+        rank, hidden_dim, device="cuda", dtype=torch.float16, requires_grad=True
+    )
+    q_B = torch.randn(
+        hidden_dim, rank, device="cuda", dtype=torch.float16, requires_grad=True
+    )
+    k_A = torch.randn(
+        rank, hidden_dim, device="cuda", dtype=torch.float16, requires_grad=True
+    )
+    k_B = torch.randn(
+        hidden_dim, rank, device="cuda", dtype=torch.float16, requires_grad=True
+    )
+    v_A = torch.randn(
+        rank, hidden_dim, device="cuda", dtype=torch.float16, requires_grad=True
+    )
+    v_B = torch.randn(
+        hidden_dim, rank, device="cuda", dtype=torch.float16, requires_grad=True
+    )
+    scale = 0.5
+
+    X.requires_grad = True
+
+    # Test without LoRA adapters
+
+    Q1, K1, V1 = LoRA_QKV.apply(
+        X,
+        None,  # X_drop
+        q_weight,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,  # Q: weight, bias, quant, A, B, scale, lora_bias, magnitude
+        k_weight,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,  # K
+        v_weight,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,  # V
+        True,  # inplace
+    )
+
+    assert Q1.shape == K1.shape == V1.shape == X.shape
+    loss1 = (Q1 + K1 + V1).sum()
+    loss1.backward()
+    assert X.grad is not None
+
+    # Clear gradients
+    X.grad = None
+
+    # Test with LoRA adapters
+    Q2, K2, V2 = LoRA_QKV.apply(
+        X,
+        None,  # X_drop
+        q_weight,
+        None,
+        None,
+        q_A,
+        q_B,
+        scale,
+        None,
+        None,  # Q
+        k_weight,
+        None,
+        None,
+        k_A,
+        k_B,
+        scale,
+        None,
+        None,  # K
+        v_weight,
+        None,
+        None,
+        v_A,
+        v_B,
+        scale,
+        None,
+        None,  # V
+        True,  # inplace
+    )
+
+    assert Q2.shape == K2.shape == V2.shape == X.shape
+    loss2 = (Q2 + K2 + V2).sum()
+    loss2.backward()
+
+    # Check gradients
+    assert X.grad is not None
+    assert q_A.grad is not None
+    assert q_B.grad is not None
+    assert k_A.grad is not None
+    assert k_B.grad is not None
+    assert v_A.grad is not None
+    assert v_B.grad is not None
+
+    # Check for NaN values
+    assert not torch.isnan(X.grad).any()
+    assert not torch.isnan(q_A.grad).any()
+    assert not torch.isnan(q_B.grad).any()
+    assert not torch.isnan(k_A.grad).any()
+    assert not torch.isnan(k_B.grad).any()
+    assert not torch.isnan(v_A.grad).any()
+    assert not torch.isnan(v_B.grad).any()
+
+
+def test_lora_o(sample_tensors):
+    """Tests LoRA output projection"""
+    X = sample_tensors["X"]
+    W = sample_tensors["W"]
+    b = sample_tensors["b"]
+    scale = sample_tensors["scale"]
+
+    shapes = sample_tensors["shapes"]
+    hidden_dim = shapes["hidden"]
+    out_dim = shapes["out"]
+    rank = shapes["rank"]
+
+    A = torch.randn(rank, hidden_dim, device="cuda", dtype=torch.float16)
+    B = torch.randn(out_dim, rank, device="cuda", dtype=torch.float16)
+
+    # Test forward pass
+    X.requires_grad = True
+    output = LoRA_O.apply(
+        X, None, W, b, None, A, B, scale, None, None
+    )  # X_drop, ..., lora_bias, magnitude
+
+    assert output.shape == (X.shape[0], X.shape[1], W.shape[0])
+
+    # Test backward pass
+    loss = output.sum()
+    loss.backward()
+    assert X.grad is not None
+
+
+def test_with_quantization(sample_tensors, mock_quantstate):
+    """Tests LoRA with quantized weights"""
+    X = sample_tensors["X"]  # [batch, seq, hidden]
+    W = sample_tensors["W"]  # [out, hidden]
+    b = sample_tensors["b"]  # [out]
+    scale = 0.5
+
+    shapes = sample_tensors["shapes"]
+    hidden_dim = shapes["hidden"]
+    out_dim = shapes["out"]
+    rank = shapes["rank"]
+
+    A = torch.randn(rank, hidden_dim, device="cuda", dtype=torch.float16)
+    B = torch.randn(out_dim, rank, device="cuda", dtype=torch.float16)
+
+    # Test matmul with quantization
+    out = matmul_lora(X, W, b, mock_quantstate, A, B, scale)
+    assert out.shape == (X.shape[0], X.shape[1], W.shape[0])
+    assert not torch.isnan(out).any()
+
+    # Test with different batch sizes
+    X2 = torch.randn(4, 6, hidden_dim, device="cuda", dtype=torch.float16)
+    out2 = matmul_lora(X2, W, b, mock_quantstate, A, B, scale)
+    assert out2.shape == (4, 6, W.shape[0])
+    assert not torch.isnan(out2).any()
+
+
+@pytest.mark.parametrize(
+    "batch,seq,hidden,rank,out",
+    [
+        (1, 1, 32, 4, 64),
+        (2, 3, 64, 8, 128),
+        (4, 5, 128, 16, 256),
+    ],
+)
+def test_shapes_and_dimensions(batch, seq, hidden, rank, out):
+    """Tests various input shapes and dimensions"""
+    X = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.float16)
+    W = torch.randn(out, hidden, device="cuda", dtype=torch.float16)
+    b = torch.randn(out, device="cuda", dtype=torch.float16)
+    A = torch.randn(rank, hidden, device="cuda", dtype=torch.float16)
+    B = torch.randn(out, rank, device="cuda", dtype=torch.float16)
+    scale = 0.5
+
+    result = matmul_lora(X, W, b, None, A, B, scale)
+    assert result.shape == (batch, seq, out)
+
+
+def test_gradient_flow(sample_tensors):
+    """Tests gradient flow through LoRA layers"""
+    X = sample_tensors["X"].clone()
+    W = sample_tensors["W"].clone()
+    b = sample_tensors["b"].clone()
+    scale = sample_tensors["scale"]
+
+    shapes = sample_tensors["shapes"]
+    hidden_dim = shapes["hidden"]
+    out_dim = shapes["out"]
+    rank = shapes["rank"]
+
+    A = torch.randn(rank, hidden_dim, device="cuda", dtype=torch.float16)
+    B = torch.randn(out_dim, rank, device="cuda", dtype=torch.float16)
+
+    X.requires_grad = True
+    A.requires_grad = True
+    B.requires_grad = True
+
+    # Forward pass
+    out = matmul_lora(X, W, b, None, A, B, scale)
+    loss = out.sum()
+
+    # Backward pass
+    loss.backward()
+
+    assert X.grad is not None
+    assert A.grad is not None
+    assert B.grad is not None
+    assert not torch.isnan(X.grad).any()
+    assert not torch.isnan(A.grad).any()
+    assert not torch.isnan(B.grad).any()
+
+
+@pytest.mark.parametrize(
+    "apply_function",
+    [apply_lora_mlp_swiglu, apply_lora_mlp_geglu],
+)
+def test_inplace_operations(sample_tensors, apply_function):
+    """Tests inplace operation behavior"""
+    X = sample_tensors["X"]
+    shapes = sample_tensors["shapes"]
+
+    # Create MLP with both inplace=True and inplace=False
+    mlp = type(
+        "MLPModule",
+        (),
+        {
+            "gate_proj": nn.Linear(shapes["hidden"], shapes["out"]).to(
+                device="cuda", dtype=torch.float16
+            ),
+            "up_proj": nn.Linear(shapes["hidden"], shapes["out"]).to(
+                device="cuda", dtype=torch.float16
+            ),
+            "down_proj": nn.Linear(shapes["out"], shapes["hidden"]).to(
+                device="cuda", dtype=torch.float16
+            ),
+            "training": False,
+        },
+    )
+
+    out1 = apply_function(mlp, X.clone(), inplace=True)
+    out2 = apply_function(mlp, X.clone(), inplace=False)
+
+    assert torch.allclose(out1, out2, rtol=1e-3)
+
+
+class ViewOutputFunction(torch.autograd.Function):
+    """Returns a view created inside the Function, like liger RMSNorm does."""
+
+    @staticmethod
+    def forward(ctx, x):
+        flat = x.reshape(-1, x.shape[-1])
+        return (flat * 1.0).view(x.shape)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+def _fused_kernel_args(kernel, hidden, rank):
+    torch.manual_seed(42)
+    n_proj = {"qkv": 3, "qk": 2, "mlp": 3}[kernel]
+    args = []
+    for _ in range(n_proj):
+        weight = torch.randn(hidden, hidden, device="cuda", dtype=torch.float16) * 0.02
+        lora_a = torch.randn(
+            rank, hidden, device="cuda", dtype=torch.float16, requires_grad=True
+        )
+        lora_b = torch.randn(
+            hidden, rank, device="cuda", dtype=torch.float16, requires_grad=True
+        )
+        args += [weight, None, None, lora_a, lora_b, 1.0, None, None]
+    return args
+
+
+def _apply_fused_kernel(kernel, h, proj_args):
+    if kernel == "qkv":
+        return LoRA_QKV.apply(h, None, *proj_args, True)
+    if kernel == "qk":
+        return LoRA_QK.apply(h, None, *proj_args, True)
+    return LoRA_MLP.apply(h, None, *proj_args, swiglu_forward, swiglu_backward, True)
+
+
+@pytest.mark.parametrize("kernel", ["qkv", "qk", "mlp"])
+def test_backward_does_not_mutate_saved_activation(kernel):
+    """Input gradient must not be written into the saved activation buffer."""
+    hidden, rank = 64, 8
+    proj_args = _fused_kernel_args(kernel, hidden, rank)
+    X = torch.randn(
+        2, 3, hidden, device="cuda", dtype=torch.float16, requires_grad=True
+    )
+
+    h = ViewOutputFunction.apply(X)
+    h_before = h.detach().clone()
+    out = _apply_fused_kernel(kernel, h, proj_args)
+    outputs = out if isinstance(out, tuple) else (out,)
+    sum(o.float().sum() for o in outputs).backward()
+
+    assert torch.equal(h.detach(), h_before)
+
+
+def test_qkv_compile_traceable_with_view_input():
+    """LoRA_QKV must trace under fullgraph compile when its input is a custom-Function view."""
+    hidden, rank = 64, 8
+    proj_args = _fused_kernel_args("qkv", hidden, rank)
+    X = torch.randn(
+        2, 3, hidden, device="cuda", dtype=torch.float16, requires_grad=True
+    )
+
+    def fn(x):
+        h = ViewOutputFunction.apply(x)
+        return LoRA_QKV.apply(h, None, *proj_args, True)
+
+    torch._dynamo.reset()
+    suppress_prior = torch._dynamo.config.suppress_errors
+    torch._dynamo.config.suppress_errors = False
+    try:
+        q, k, v = torch.compile(fn, fullgraph=True)(X)
+        (q.float().sum() + k.float().sum() + v.float().sum()).backward()
+    finally:
+        torch._dynamo.config.suppress_errors = suppress_prior
+        torch._dynamo.reset()
+
+    assert X.grad is not None
diff --git a/tests/e2e/kernels/test_lora_features.py b/tests/e2e/kernels/test_lora_features.py
new file mode 100644
index 0000000000..688710f3a4
--- /dev/null
+++ b/tests/e2e/kernels/test_lora_features.py
@@ -0,0 +1,1245 @@
+"""
+Tests for LoRA kernel correctness with bias, dropout, and DoRA support.
+
+Compares fused kernel outputs and gradients against PEFT's reference implementation.
+"""
+
+import pytest
+import torch
+from peft import LoraConfig, get_peft_model
+from torch import nn
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from axolotl.kernels.lora import (
+    _compute_dora_scale,
+    apply_lora_mlp_swiglu,
+    apply_lora_o,
+    apply_lora_qkv,
+    get_lora_parameters,
+    matmul_lora,
+)
+from axolotl.monkeypatch.lora_kernels import (
+    apply_lora_kernel_patches,
+    patch_self_attn_lora,
+)
+from axolotl.utils.dict import DictDefault
+
+MODEL_NAME = "axolotl-ai-co/tiny-qwen3-129m"
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+@pytest.fixture(scope="module")
+def model_config():
+    return AutoConfig.from_pretrained(MODEL_NAME)
+
+
+def _make_peft_model(
+    lora_dropout=0.0,
+    bias="none",
+    use_dora=False,
+    target_modules=None,
+):
+    """Create a PEFT model with given config."""
+    if target_modules is None:
+        target_modules = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ]
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=DTYPE,
+        attn_implementation="eager",
+    ).to(DEVICE)
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=16,
+        lora_dropout=lora_dropout,
+        bias=bias,
+        use_dora=use_dora,
+        target_modules=target_modules,
+    )
+    peft_model = get_peft_model(model, lora_config)
+    return peft_model
+
+
+def _get_layer(peft_model, layer_idx=0):
+    """Get a specific transformer layer from the model."""
+    return peft_model.model.model.layers[layer_idx]
+
+
+def _make_input(batch=2, seq_len=16, hidden_size=1024):
+    """Create random input tensor."""
+    return torch.randn(
+        batch, seq_len, hidden_size, dtype=DTYPE, device=DEVICE, requires_grad=True
+    )
+
+
+def _compare_tensors(a, b, name="", atol=1e-2, rtol=1e-2):
+    """Compare two tensors with informative error messages."""
+    if a is None and b is None:
+        return
+    assert a is not None and b is not None, f"{name}: one is None, other is not"
+    assert a.shape == b.shape, f"{name}: shape mismatch {a.shape} vs {b.shape}"
+    diff = (a - b).abs()
+    max_diff = diff.max().item()
+    mean_diff = diff.mean().item()
+    assert torch.allclose(a, b, atol=atol, rtol=rtol), (
+        f"{name}: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}"
+    )
+
+
+class TestGetLoraParameters:
+    """Test the extended get_lora_parameters function."""
+
+    def test_returns_9_values(self):
+        model = _make_peft_model()
+        layer = _get_layer(model)
+        params = get_lora_parameters(layer.self_attn.q_proj)
+        assert len(params) == 9
+        W, b, quant, A, B, s, lora_bias, dropout, magnitude = params
+        assert W is not None
+        assert A is not None
+        assert B is not None
+        assert s is not None
+        assert lora_bias is None  # bias="none"
+        assert dropout is not None  # should be nn.Identity
+        assert magnitude is None  # no DoRA
+        del model
+
+    def test_with_bias(self):
+        """Qwen3 has no base bias, so PEFT doesn't add lora_bias even with bias='lora_only'.
+        This test verifies get_lora_parameters handles this correctly."""
+        model = _make_peft_model(bias="lora_only")
+        layer = _get_layer(model)
+        params = get_lora_parameters(layer.self_attn.q_proj)
+        _, _, _, _, _, _, lora_bias, _, _ = params
+        # Qwen3 q_proj has no base bias, so PEFT sets lora_bias=False
+        assert lora_bias is None
+        del model
+
+    def test_with_bias_on_biased_layer(self):
+        """Test with manually added bias to verify lora_bias extraction."""
+        model = _make_peft_model(bias="lora_only")
+        layer = _get_layer(model)
+        q_proj = layer.self_attn.q_proj
+        adapter = q_proj.active_adapters[0]
+        # Manually add bias to lora_B to test extraction
+        old_B = q_proj.lora_B[adapter]
+        q_proj.lora_B[adapter] = torch.nn.Linear(
+            old_B.in_features, old_B.out_features, bias=True, device=DEVICE, dtype=DTYPE
+        )
+        params = get_lora_parameters(q_proj)
+        _, _, _, _, _, _, lora_bias, _, _ = params
+        assert lora_bias is not None
+        assert lora_bias.shape[0] == old_B.out_features
+        del model
+
+    def test_with_dropout(self):
+        model = _make_peft_model(lora_dropout=0.1)
+        layer = _get_layer(model)
+        params = get_lora_parameters(layer.self_attn.q_proj)
+        _, _, _, _, _, _, _, dropout, _ = params
+        assert dropout is not None
+        assert isinstance(dropout, nn.Dropout)
+        del model
+
+    def test_with_dora(self):
+        model = _make_peft_model(use_dora=True)
+        layer = _get_layer(model)
+        params = get_lora_parameters(layer.self_attn.q_proj)
+        _, _, _, _, _, _, _, _, magnitude = params
+        assert magnitude is not None
+        del model
+
+
+class TestMatmulLora:
+    """Test matmul_lora with new lora_bias and X_drop parameters."""
+
+    def test_basic(self):
+        X = torch.randn(4, 8, dtype=DTYPE, device=DEVICE)
+        W = torch.randn(16, 8, dtype=DTYPE, device=DEVICE)
+        A = torch.randn(4, 8, dtype=DTYPE, device=DEVICE)  # [rank, in]
+        B = torch.randn(16, 4, dtype=DTYPE, device=DEVICE)  # [out, rank]
+        s = 2.0
+
+        result = matmul_lora(X, W, None, None, A, B, s)
+        expected = X @ W.t() + s * X @ A.t() @ B.t()
+        _compare_tensors(result, expected, "basic matmul_lora")
+
+    def test_with_lora_bias(self):
+        X = torch.randn(4, 8, dtype=DTYPE, device=DEVICE)
+        W = torch.randn(16, 8, dtype=DTYPE, device=DEVICE)
+        A = torch.randn(4, 8, dtype=DTYPE, device=DEVICE)
+        B = torch.randn(16, 4, dtype=DTYPE, device=DEVICE)
+        lora_bias = torch.randn(16, dtype=DTYPE, device=DEVICE)
+        s = 2.0
+
+        result = matmul_lora(X, W, None, None, A, B, s, lora_bias=lora_bias)
+        expected = X @ W.t() + s * X @ A.t() @ B.t() + s * lora_bias
+        _compare_tensors(result, expected, "matmul_lora with lora_bias")
+
+    def test_with_x_drop(self):
+        X = torch.randn(4, 8, dtype=DTYPE, device=DEVICE)
+        X_drop = X * 0.5  # simulated dropout
+        W = torch.randn(16, 8, dtype=DTYPE, device=DEVICE)
+        A = torch.randn(4, 8, dtype=DTYPE, device=DEVICE)
+        B = torch.randn(16, 4, dtype=DTYPE, device=DEVICE)
+        s = 2.0
+
+        result = matmul_lora(X, W, None, None, A, B, s, X_drop=X_drop)
+        expected = X @ W.t() + s * X_drop @ A.t() @ B.t()
+        _compare_tensors(result, expected, "matmul_lora with X_drop")
+
+
+class TestDoraScale:
+    """Test DoRA magnitude/norm scaling computation."""
+
+    def test_basic(self):
+        W = torch.randn(16, 8, dtype=DTYPE, device=DEVICE)
+        A = torch.randn(4, 8, dtype=DTYPE, device=DEVICE)
+        B = torch.randn(16, 4, dtype=DTYPE, device=DEVICE)
+        magnitude = torch.randn(16, dtype=DTYPE, device=DEVICE).abs() + 0.1
+        s = 2.0
+
+        scale = _compute_dora_scale(W, None, A, B, s, magnitude, DTYPE)
+
+        # Manual computation
+        combined = W + s * B @ A
+        weight_norm = torch.linalg.norm(combined, dim=1)
+        expected = magnitude / weight_norm
+
+        _compare_tensors(scale, expected, "dora_scale")
+
+
+# ============================================================
+# Integration tests: compare kernel outputs against PEFT reference
+# ============================================================
+
+
+def _run_peft_qkv(layer, X):
+    """Run Q, K, V projections through PEFT's standard forward."""
+    Q = layer.self_attn.q_proj(X)
+    K = layer.self_attn.k_proj(X)
+    V = layer.self_attn.v_proj(X)
+    return Q, K, V
+
+
+def _run_kernel_qkv(layer, X):
+    """Run Q, K, V projections through our fused kernel."""
+    return apply_lora_qkv(layer.self_attn, X, inplace=False)
+
+
+def _run_peft_o(layer, X):
+    """Run O projection through PEFT's standard forward."""
+    return layer.self_attn.o_proj(X)
+
+
+def _run_kernel_o(layer, X):
+    """Run O projection through our fused kernel."""
+    return apply_lora_o(layer.self_attn, X)
+
+
+def _run_peft_mlp(layer, X):
+    """Run MLP through PEFT's standard forward."""
+    return layer.mlp(X)
+
+
+def _run_kernel_mlp(layer, X):
+    """Run MLP through our fused kernel."""
+    return apply_lora_mlp_swiglu(layer.mlp, X, inplace=False)
+
+
+class TestQKVKernel:
+    """Test LoRA_QKV kernel against PEFT reference."""
+
+    @pytest.mark.parametrize("bias", ["none", "lora_only"])
+    def test_forward_bias(self, bias):
+        model = _make_peft_model(bias=bias)
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=model.config.hidden_size)
+
+        with torch.no_grad():
+            peft_Q, peft_K, peft_V = _run_peft_qkv(layer, X)
+            kern_Q, kern_K, kern_V = _run_kernel_qkv(layer, X)
+
+        _compare_tensors(kern_Q, peft_Q, f"QKV Q (bias={bias})")
+        _compare_tensors(kern_K, peft_K, f"QKV K (bias={bias})")
+        _compare_tensors(kern_V, peft_V, f"QKV V (bias={bias})")
+        del model
+
+    def test_forward_dropout_eval(self):
+        """Dropout disabled in eval - should match exactly."""
+        model = _make_peft_model(lora_dropout=0.1)
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=model.config.hidden_size)
+
+        with torch.no_grad():
+            peft_Q, peft_K, peft_V = _run_peft_qkv(layer, X)
+            kern_Q, kern_K, kern_V = _run_kernel_qkv(layer, X)
+
+        _compare_tensors(kern_Q, peft_Q, "QKV Q (dropout eval)")
+        _compare_tensors(kern_K, peft_K, "QKV K (dropout eval)")
+        _compare_tensors(kern_V, peft_V, "QKV V (dropout eval)")
+        del model
+
+    def test_forward_dora(self):
+        model = _make_peft_model(use_dora=True)
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=model.config.hidden_size)
+
+        with torch.no_grad():
+            peft_Q, peft_K, peft_V = _run_peft_qkv(layer, X)
+            kern_Q, kern_K, kern_V = _run_kernel_qkv(layer, X)
+
+        _compare_tensors(kern_Q, peft_Q, "QKV Q (DoRA)")
+        _compare_tensors(kern_K, peft_K, "QKV K (DoRA)")
+        _compare_tensors(kern_V, peft_V, "QKV V (DoRA)")
+        del model
+
+    def test_forward_dora_bias(self):
+        model = _make_peft_model(use_dora=True, bias="lora_only")
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=model.config.hidden_size)
+
+        with torch.no_grad():
+            peft_Q, peft_K, peft_V = _run_peft_qkv(layer, X)
+            kern_Q, kern_K, kern_V = _run_kernel_qkv(layer, X)
+
+        _compare_tensors(kern_Q, peft_Q, "QKV Q (DoRA+bias)")
+        _compare_tensors(kern_K, peft_K, "QKV K (DoRA+bias)")
+        _compare_tensors(kern_V, peft_V, "QKV V (DoRA+bias)")
+        del model
+
+    @pytest.mark.parametrize("bias", ["none", "lora_only"])
+    def test_backward_bias(self, bias):
+        """Test that gradients match between kernel and PEFT."""
+        model = _make_peft_model(bias=bias)
+        model.train()
+        layer = _get_layer(model)
+
+        # PEFT reference
+        X1 = _make_input(hidden_size=model.config.hidden_size)
+        pQ, pK, pV = _run_peft_qkv(layer, X1)
+        loss_peft = pQ.sum() + pK.sum() + pV.sum()
+        loss_peft.backward()
+
+        peft_grads = {}
+        for name, param in layer.self_attn.named_parameters():
+            if param.grad is not None:
+                peft_grads[name] = param.grad.clone()
+        layer.self_attn.zero_grad()
+
+        # Kernel
+        X2 = X1.detach().clone().requires_grad_(True)
+        kQ, kK, kV = _run_kernel_qkv(layer, X2)
+        loss_kern = kQ.sum() + kK.sum() + kV.sum()
+        loss_kern.backward()
+
+        kern_grads = {}
+        for name, param in layer.self_attn.named_parameters():
+            if param.grad is not None:
+                kern_grads[name] = param.grad.clone()
+        layer.self_attn.zero_grad()
+
+        # Compare LoRA parameter gradients
+        for name in peft_grads:
+            if "lora_" in name:
+                _compare_tensors(
+                    kern_grads.get(name),
+                    peft_grads[name],
+                    f"grad {name} (bias={bias})",
+                    atol=5e-2,
+                    rtol=5e-2,
+                )
+        del model
+
+    def test_backward_dora(self):
+        """Test DoRA backward pass gradients."""
+        model = _make_peft_model(use_dora=True)
+        model.train()
+        layer = _get_layer(model)
+
+        X1 = _make_input(hidden_size=model.config.hidden_size)
+        pQ, pK, pV = _run_peft_qkv(layer, X1)
+        loss_peft = pQ.sum() + pK.sum() + pV.sum()
+        loss_peft.backward()
+
+        peft_grads = {}
+        for name, param in layer.self_attn.named_parameters():
+            if param.grad is not None:
+                peft_grads[name] = param.grad.clone()
+        layer.self_attn.zero_grad()
+
+        X2 = X1.detach().clone().requires_grad_(True)
+        kQ, kK, kV = _run_kernel_qkv(layer, X2)
+        loss_kern = kQ.sum() + kK.sum() + kV.sum()
+        loss_kern.backward()
+
+        kern_grads = {}
+        for name, param in layer.self_attn.named_parameters():
+            if param.grad is not None:
+                kern_grads[name] = param.grad.clone()
+        layer.self_attn.zero_grad()
+
+        for name in peft_grads:
+            if "lora_" in name or "magnitude" in name:
+                _compare_tensors(
+                    kern_grads.get(name),
+                    peft_grads[name],
+                    f"grad {name} (DoRA)",
+                    atol=5e-2,
+                    rtol=5e-2,
+                )
+        del model
+
+
+class TestOKernel:
+    """Test LoRA_O kernel against PEFT reference."""
+
+    @staticmethod
+    def _o_input_dim(model):
+        """o_proj input is num_heads * head_dim (may differ from hidden_size with GQA)."""
+        cfg = model.config
+        text_cfg = cfg.get_text_config() if hasattr(cfg, "get_text_config") else cfg
+        return text_cfg.num_attention_heads * text_cfg.head_dim
+
+    @pytest.mark.parametrize("bias", ["none", "lora_only"])
+    def test_forward_bias(self, bias):
+        model = _make_peft_model(bias=bias)
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=self._o_input_dim(model))
+
+        with torch.no_grad():
+            peft_out = _run_peft_o(layer, X)
+            kern_out = _run_kernel_o(layer, X)
+
+        _compare_tensors(kern_out, peft_out, f"O (bias={bias})")
+        del model
+
+    def test_forward_dora(self):
+        model = _make_peft_model(use_dora=True)
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=self._o_input_dim(model))
+
+        with torch.no_grad():
+            peft_out = _run_peft_o(layer, X)
+            kern_out = _run_kernel_o(layer, X)
+
+        _compare_tensors(kern_out, peft_out, "O (DoRA)")
+        del model
+
+    @pytest.mark.parametrize("bias", ["none", "lora_only"])
+    def test_backward_bias(self, bias):
+        model = _make_peft_model(bias=bias)
+        model.train()
+        layer = _get_layer(model)
+
+        X1 = _make_input(hidden_size=self._o_input_dim(model))
+        peft_out = _run_peft_o(layer, X1)
+        peft_out.sum().backward()
+        peft_grads = {
+            n: p.grad.clone()
+            for n, p in layer.self_attn.o_proj.named_parameters()
+            if p.grad is not None
+        }
+        layer.self_attn.o_proj.zero_grad()
+
+        X2 = X1.detach().clone().requires_grad_(True)
+        kern_out = _run_kernel_o(layer, X2)
+        kern_out.sum().backward()
+        kern_grads = {
+            n: p.grad.clone()
+            for n, p in layer.self_attn.o_proj.named_parameters()
+            if p.grad is not None
+        }
+        layer.self_attn.o_proj.zero_grad()
+
+        for name in peft_grads:
+            if "lora_" in name:
+                _compare_tensors(
+                    kern_grads.get(name),
+                    peft_grads[name],
+                    f"O grad {name} (bias={bias})",
+                    atol=5e-2,
+                    rtol=5e-2,
+                )
+        del model
+
+
+class TestMLPKernel:
+    """Test LoRA_MLP kernel against PEFT reference."""
+
+    @pytest.mark.parametrize("bias", ["none", "lora_only"])
+    def test_forward_bias(self, bias):
+        model = _make_peft_model(bias=bias)
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=model.config.hidden_size)
+
+        with torch.no_grad():
+            peft_out = _run_peft_mlp(layer, X)
+            kern_out = _run_kernel_mlp(layer, X)
+
+        _compare_tensors(kern_out, peft_out, f"MLP (bias={bias})")
+        del model
+
+    def test_forward_dropout_eval(self):
+        model = _make_peft_model(lora_dropout=0.1)
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=model.config.hidden_size)
+
+        with torch.no_grad():
+            peft_out = _run_peft_mlp(layer, X)
+            kern_out = _run_kernel_mlp(layer, X)
+
+        _compare_tensors(kern_out, peft_out, "MLP (dropout eval)")
+        del model
+
+    def test_forward_dora(self):
+        model = _make_peft_model(use_dora=True)
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=model.config.hidden_size)
+
+        with torch.no_grad():
+            peft_out = _run_peft_mlp(layer, X)
+            kern_out = _run_kernel_mlp(layer, X)
+
+        # Relaxed tolerance for MLP DoRA: 3 projections + activation + DoRA
+        # causes bf16 accumulation differences
+        _compare_tensors(kern_out, peft_out, "MLP (DoRA)", atol=0.3, rtol=0.05)
+        del model
+
+    def test_forward_dora_bias(self):
+        model = _make_peft_model(use_dora=True, bias="lora_only")
+        model.eval()
+        layer = _get_layer(model)
+        X = _make_input(hidden_size=model.config.hidden_size)
+
+        with torch.no_grad():
+            peft_out = _run_peft_mlp(layer, X)
+            kern_out = _run_kernel_mlp(layer, X)
+
+        _compare_tensors(kern_out, peft_out, "MLP (DoRA+bias)", atol=0.3, rtol=0.05)
+        del model
+
+    @pytest.mark.parametrize("bias", ["none", "lora_only"])
+    def test_backward_bias(self, bias):
+        model = _make_peft_model(bias=bias)
+        model.train()
+        layer = _get_layer(model)
+        hidden_size = model.config.hidden_size
+
+        X1 = _make_input(hidden_size=hidden_size)
+        peft_out = _run_peft_mlp(layer, X1)
+        peft_out.sum().backward()
+        peft_grads = {
+            n: p.grad.clone()
+            for n, p in layer.mlp.named_parameters()
+            if p.grad is not None
+        }
+        layer.mlp.zero_grad()
+
+        X2 = X1.detach().clone().requires_grad_(True)
+        kern_out = _run_kernel_mlp(layer, X2)
+        kern_out.sum().backward()
+        kern_grads = {
+            n: p.grad.clone()
+            for n, p in layer.mlp.named_parameters()
+            if p.grad is not None
+        }
+        layer.mlp.zero_grad()
+
+        # MLP backward has longer chain (3 projections + activation) = more bf16 accumulation error
+        for name in peft_grads:
+            if "lora_" in name:
+                _compare_tensors(
+                    kern_grads.get(name),
+                    peft_grads[name],
+                    f"MLP grad {name} (bias={bias})",
+                    atol=0.5,
+                    rtol=0.1,
+                )
+        del model
+
+    def test_backward_dora(self):
+        model = _make_peft_model(use_dora=True)
+        model.train()
+        layer = _get_layer(model)
+
+        X1 = _make_input(hidden_size=model.config.hidden_size)
+        peft_out = _run_peft_mlp(layer, X1)
+        peft_out.sum().backward()
+        peft_grads = {
+            n: p.grad.clone()
+            for n, p in layer.mlp.named_parameters()
+            if p.grad is not None
+        }
+        layer.mlp.zero_grad()
+
+        X2 = X1.detach().clone().requires_grad_(True)
+        kern_out = _run_kernel_mlp(layer, X2)
+        kern_out.sum().backward()
+        kern_grads = {
+            n: p.grad.clone()
+            for n, p in layer.mlp.named_parameters()
+            if p.grad is not None
+        }
+        layer.mlp.zero_grad()
+
+        for name in peft_grads:
+            if "lora_" in name or "magnitude" in name:
+                _compare_tensors(
+                    kern_grads.get(name),
+                    peft_grads[name],
+                    f"MLP grad {name} (DoRA)",
+                    atol=0.5,
+                    rtol=0.1,
+                )
+        del model
+
+
+class TestFullModelPatch:
+    """Test applying kernel patches to a full model."""
+
+    def test_patched_forward_basic(self):
+        """Test that patched model forward matches unpatched PEFT model (bias=none, no DoRA)."""
+        from peft import PeftModelForCausalLM
+
+        base_model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=DTYPE,
+            attn_implementation="eager",
+        ).to(DEVICE)
+        lora_config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            bias="none",
+            use_dora=False,
+            target_modules=[
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "o_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj",
+            ],
+        )
+        model = PeftModelForCausalLM(base_model, lora_config)
+        model.eval()
+
+        # Get PEFT reference output
+        input_ids = torch.randint(0, 1000, (1, 32), device=DEVICE)
+        with torch.no_grad():
+            peft_out = model(input_ids).logits
+
+        # Apply kernel patches
+        cfg = DictDefault(
+            {
+                "base_model": MODEL_NAME,
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_mlp_kernel": True,
+            }
+        )
+        patch_self_attn_lora(cfg)
+        apply_lora_kernel_patches(model, cfg)
+
+        # Get kernel output
+        with torch.no_grad():
+            kern_out = model(input_ids).logits
+
+        _compare_tensors(kern_out, peft_out, "Full model (basic)", atol=5e-1, rtol=1e-1)
+        del model
+
+
+class TestEmbeddingKernel:
+    """Test LoRA embedding kernel against PEFT reference."""
+
+    def _make_embedding_model(self, use_dora=False):
+        from peft import PeftModelForCausalLM
+
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=DTYPE,
+            attn_implementation="eager",
+        ).to(DEVICE)
+        lora_config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            use_dora=use_dora,
+            target_modules=["embed_tokens"],
+        )
+        return PeftModelForCausalLM(model, lora_config)
+
+    def test_forward_basic(self):
+        from axolotl.kernels.lora import apply_lora_embedding
+
+        model = self._make_embedding_model()
+        model.eval()
+
+        embed = model.model.model.embed_tokens
+        input_ids = torch.randint(0, 1000, (2, 16), device=DEVICE)
+
+        with torch.no_grad():
+            peft_out = embed(input_ids)
+            kern_out = apply_lora_embedding(embed, input_ids)
+
+        # Cast to same dtype for comparison (PEFT may return float32)
+        _compare_tensors(kern_out.to(peft_out.dtype), peft_out, "Embedding basic")
+        del model
+
+    def test_forward_dora(self):
+        from axolotl.kernels.lora import apply_lora_embedding
+
+        model = self._make_embedding_model(use_dora=True)
+        model.eval()
+
+        embed = model.model.model.embed_tokens
+        input_ids = torch.randint(0, 1000, (2, 16), device=DEVICE)
+
+        with torch.no_grad():
+            peft_out = embed(input_ids)
+            kern_out = apply_lora_embedding(embed, input_ids)
+
+        _compare_tensors(
+            kern_out.to(peft_out.dtype), peft_out, "Embedding DoRA", atol=0.3, rtol=0.05
+        )
+        del model
+
+    def test_backward(self):
+        from axolotl.kernels.lora import apply_lora_embedding
+
+        model = self._make_embedding_model()
+        model.train()
+
+        embed = model.model.model.embed_tokens
+        input_ids = torch.randint(0, 1000, (2, 16), device=DEVICE)
+
+        # PEFT reference
+        peft_out = embed(input_ids)
+        peft_out.sum().backward()
+        peft_grads = {}
+        for n, p in embed.named_parameters():
+            if p.grad is not None and "lora" in n:
+                peft_grads[n] = p.grad.clone()
+        embed.zero_grad()
+
+        # Kernel
+        kern_out = apply_lora_embedding(embed, input_ids)
+        kern_out.sum().backward()
+        kern_grads = {}
+        for n, p in embed.named_parameters():
+            if p.grad is not None and "lora" in n:
+                kern_grads[n] = p.grad.clone()
+        embed.zero_grad()
+
+        for name in peft_grads:
+            _compare_tensors(
+                kern_grads.get(name),
+                peft_grads[name],
+                f"Embedding grad {name}",
+                atol=5e-2,
+                rtol=5e-2,
+            )
+        del model
+
+
+class TestTiedEmbeddings:
+    """Test that tied embeddings work correctly with kernel patching."""
+
+    def test_tied_embed_and_lm_head(self):
+        """When both embed_tokens and lm_head have LoRA, PEFT unties them.
+        Verify patched model produces valid output (no crashes, finite values)."""
+        from peft import PeftModelForCausalLM
+
+        base = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=DTYPE,
+            attn_implementation="eager",
+        ).to(DEVICE)
+        lora_config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            target_modules=[
+                "embed_tokens",
+                "lm_head",
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "o_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj",
+            ],
+        )
+        model = PeftModelForCausalLM(base, lora_config)
+        model.eval()
+
+        cfg = DictDefault(
+            {
+                "base_model": MODEL_NAME,
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_mlp_kernel": True,
+                "lora_embedding_kernel": True,
+            }
+        )
+
+        # Apply all kernel patches (class + instance level)
+        patch_self_attn_lora(cfg)
+        apply_lora_kernel_patches(model, cfg)
+
+        input_ids = torch.randint(0, 1000, (1, 32), device=DEVICE)
+        with torch.no_grad():
+            out = model(input_ids).logits
+
+        # Verify output is valid
+        assert out.shape == (1, 32, model.config.vocab_size)
+        assert torch.isfinite(out).all(), "Output contains non-finite values"
+        assert out.abs().max() > 0, "Output is all zeros"
+
+        # Verify backward works
+        model.train()
+        out = model(input_ids).logits
+        out.sum().backward()
+        # Check that LoRA params got gradients
+        embed = model.model.model.embed_tokens
+        has_embed_grad = any(
+            p.grad is not None and p.grad.abs().sum() > 0
+            for n, p in embed.named_parameters()
+            if "lora" in n
+        )
+        assert has_embed_grad, "Embedding LoRA params got no gradients"
+        del model
+
+
+class TestQuantizedModels:
+    """Test kernels with quantized base weights."""
+
+    def test_nf4_qlora_forward_backward(self):
+        """NF4 QLoRA with kernel patches."""
+        from peft import PeftModelForCausalLM
+        from transformers import BitsAndBytesConfig
+
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=DTYPE,
+            bnb_4bit_use_double_quant=True,
+        )
+        base = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            quantization_config=bnb_config,
+            attn_implementation="eager",
+        )
+        lora_config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            target_modules=[
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "o_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj",
+            ],
+        )
+        model = PeftModelForCausalLM(base, lora_config)
+
+        cfg = DictDefault(
+            {
+                "base_model": MODEL_NAME,
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_mlp_kernel": True,
+            }
+        )
+        patch_self_attn_lora(cfg)
+        apply_lora_kernel_patches(model, cfg)
+        model.train()
+
+        ids = torch.randint(0, 1000, (1, 32), device=DEVICE)
+        out = model(ids).logits
+        assert torch.isfinite(out).all()
+        out.sum().backward()
+        has_grads = sum(
+            1 for n, p in model.named_parameters() if p.grad is not None and "lora" in n
+        )
+        assert has_grads > 0, "No LoRA gradients"
+        del model
+
+    def test_nf4_single_quant(self):
+        """NF4 without double quantization."""
+        from peft import PeftModelForCausalLM
+        from transformers import BitsAndBytesConfig
+
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=DTYPE,
+        )
+        base = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            quantization_config=bnb_config,
+            attn_implementation="eager",
+        )
+        lora_config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            target_modules=[
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "o_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj",
+            ],
+        )
+        model = PeftModelForCausalLM(base, lora_config)
+
+        cfg = DictDefault(
+            {
+                "base_model": MODEL_NAME,
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_mlp_kernel": True,
+            }
+        )
+        patch_self_attn_lora(cfg)
+        apply_lora_kernel_patches(model, cfg)
+        model.train()
+
+        ids = torch.randint(0, 1000, (1, 32), device=DEVICE)
+        out = model(ids).logits
+        assert torch.isfinite(out).all()
+        out.sum().backward()
+        has_grads = sum(
+            1 for n, p in model.named_parameters() if p.grad is not None and "lora" in n
+        )
+        assert has_grads > 0
+        del model
+
+
+class TestTritonDoRA:
+    """Test Triton DoRA kernel against reference implementation."""
+
+    def test_triton_dora_scale(self):
+        from axolotl.kernels.dora import triton_dora_scale
+        from axolotl.kernels.lora import _compute_dora_scale
+
+        # Random weights matching Qwen3-1.7B dimensions
+        out_feat, in_feat, rank = 1024, 1024, 8
+        W = torch.randn(out_feat, in_feat, dtype=DTYPE, device=DEVICE)
+        A = torch.randn(rank, in_feat, dtype=DTYPE, device=DEVICE)
+        B = torch.randn(out_feat, rank, dtype=DTYPE, device=DEVICE)
+        magnitude = torch.randn(out_feat, dtype=DTYPE, device=DEVICE).abs() + 0.1
+        s = 2.0
+
+        # Clear cache to force recomputation
+        if hasattr(magnitude, "_dora_cache"):
+            del magnitude._dora_cache
+
+        ref = _compute_dora_scale(W, None, A, B, s, magnitude, DTYPE)
+        tri = triton_dora_scale(W, None, A, B, s, magnitude, DTYPE)
+
+        _compare_tensors(tri, ref, "Triton DoRA scale", atol=1e-2, rtol=1e-2)
+
+    def test_triton_dora_scale_small(self):
+        """Test with K/V projection dimensions (smaller out_features)."""
+        from axolotl.kernels.dora import triton_dora_scale
+        from axolotl.kernels.lora import _compute_dora_scale
+
+        out_feat, in_feat, rank = 128, 1024, 8
+        W = torch.randn(out_feat, in_feat, dtype=DTYPE, device=DEVICE)
+        A = torch.randn(rank, in_feat, dtype=DTYPE, device=DEVICE)
+        B = torch.randn(out_feat, rank, dtype=DTYPE, device=DEVICE)
+        magnitude = torch.randn(out_feat, dtype=DTYPE, device=DEVICE).abs() + 0.1
+        s = 2.0
+
+        if hasattr(magnitude, "_dora_cache"):
+            del magnitude._dora_cache
+
+        ref = _compute_dora_scale(W, None, A, B, s, magnitude, DTYPE)
+        tri = triton_dora_scale(W, None, A, B, s, magnitude, DTYPE)
+
+        _compare_tensors(tri, ref, "Triton DoRA scale (small)", atol=1e-2, rtol=1e-2)
+
+
+# ============================================================
+# Regression tests for review fixes
+# ============================================================
+
+
+class TestDoRAEmbeddingNoDoubleScale:
+    """Regression: DoRA embedding forward must save the pre-scaled combined
+    tensor, not the already-scaled result, so backward computes d_mag correctly."""
+
+    def test_dora_magnitude_gradient_magnitude(self):
+        """d_mag should be O(1) relative to the gradient, not O(mag_scale^2)."""
+        from peft import PeftModelForCausalLM
+
+        from axolotl.kernels.lora import apply_lora_embedding
+
+        base = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=DTYPE,
+            attn_implementation="eager",
+        ).to(DEVICE)
+        lora_config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            use_dora=True,
+            target_modules=["embed_tokens"],
+        )
+        model = PeftModelForCausalLM(base, lora_config)
+        model.train()
+
+        embed = model.model.model.embed_tokens
+        ids = torch.randint(0, 1000, (2, 16), device=DEVICE)
+
+        # Run PEFT reference to get reference d_mag
+        peft_out = embed(ids)
+        peft_out.sum().backward()
+        peft_mag_grad = None
+        for n, p in embed.named_parameters():
+            if "magnitude" in n and p.grad is not None:
+                peft_mag_grad = p.grad.clone()
+        embed.zero_grad()
+
+        # Run kernel
+        kern_out = apply_lora_embedding(embed, ids)
+        kern_out.to(peft_out.dtype).sum().backward()
+        kern_mag_grad = None
+        for n, p in embed.named_parameters():
+            if "magnitude" in n and p.grad is not None:
+                kern_mag_grad = p.grad.clone()
+        embed.zero_grad()
+
+        assert peft_mag_grad is not None, "PEFT should produce magnitude gradients"
+        assert kern_mag_grad is not None, "Kernel should produce magnitude gradients"
+
+        # Key check: gradients should be same order of magnitude
+        # Double-scaling would make kern_mag_grad ~mag_scale times too large
+        ratio = kern_mag_grad.abs().mean() / peft_mag_grad.abs().mean()
+        assert 0.5 < ratio < 2.0, (
+            f"Magnitude gradient ratio kernel/peft = {ratio:.3f}, "
+            f"expected ~1.0 (double-scaling would give >> 1)"
+        )
+        del model
+
+
+class TestDoraCacheInvalidation:
+    """Regression: DoRA weight norm cache must invalidate after in-place
+    param updates (optimizer steps), not just pointer changes."""
+
+    def test_cache_invalidates_on_inplace_update(self):
+        W = torch.randn(64, 64, dtype=DTYPE, device=DEVICE)
+        A = torch.randn(8, 64, dtype=DTYPE, device=DEVICE)
+        B = torch.randn(64, 8, dtype=DTYPE, device=DEVICE)
+        magnitude = torch.randn(64, dtype=DTYPE, device=DEVICE).abs() + 0.1
+        s = 2.0
+
+        # Clear any existing cache
+        if hasattr(magnitude, "_dora_cache"):
+            del magnitude._dora_cache
+
+        # First call populates cache
+        result1 = _compute_dora_scale(W, None, A, B, s, magnitude, DTYPE)
+
+        # Simulate optimizer in-place update (pointer stays same, content changes)
+        old_ptr = A.data_ptr()
+        A.data.add_(torch.randn_like(A) * 0.1)
+        assert A.data_ptr() == old_ptr, "Pointer should not change for in-place ops"
+
+        # Second call must detect the change and recompute
+        result2 = _compute_dora_scale(W, None, A, B, s, magnitude, DTYPE)
+
+        # Results should differ since A changed
+        assert not torch.allclose(result1, result2, atol=1e-4), (
+            "DoRA scale should change after in-place param update — cache not invalidated!"
+        )
+
+
+class TestEmbeddingPaddingIdxGrad:
+    """Regression: custom embedding backward must zero out gradients at
+    padding_idx positions, matching F.embedding behavior."""
+
+    def test_padding_idx_gradient_is_zero(self):
+        from axolotl.kernels.lora import LoRA_Embedding
+
+        vocab, hidden, rank = 100, 32, 4
+        W = torch.randn(
+            vocab, hidden, dtype=torch.float32, device=DEVICE, requires_grad=False
+        )
+        A = torch.randn(
+            rank, vocab, dtype=torch.float32, device=DEVICE, requires_grad=True
+        )
+        B = torch.randn(
+            hidden, rank, dtype=torch.float32, device=DEVICE, requires_grad=True
+        )
+        s = 2.0
+        padding_idx = 0
+
+        # Input containing the padding token
+        x = torch.tensor([[padding_idx, 1, 2, padding_idx, 3]], device=DEVICE)
+
+        out = LoRA_Embedding.apply(
+            x,
+            W,
+            A,
+            B,
+            s,
+            None,
+            padding_idx,
+            None,
+            2.0,
+            False,
+            False,  # max_norm, norm_type, scale_grad_by_freq, sparse
+        )
+        out.sum().backward()
+
+        # The gradient for A at the padding_idx column should be zero
+        # A is [rank, vocab], so A.grad[:, padding_idx] should be zero
+        assert A.grad is not None
+        pad_grad = A.grad[:, padding_idx]
+        assert torch.all(pad_grad == 0), (
+            f"Gradient at padding_idx={padding_idx} should be zero, got {pad_grad}"
+        )
+
+        # Non-padding positions should have non-zero gradients
+        non_pad_grad = A.grad[:, 1]
+        assert non_pad_grad.abs().sum() > 0, "Non-padding gradients should be non-zero"
+
+
+class TestEmbeddingScaleGradByFreq:
+    """Regression: custom embedding backward must scale gradients by
+    inverse frequency when scale_grad_by_freq=True."""
+
+    def test_repeated_tokens_get_scaled_gradients(self):
+        from axolotl.kernels.lora import LoRA_Embedding
+
+        vocab, hidden, rank = 100, 32, 4
+        W = torch.randn(
+            vocab, hidden, dtype=torch.float32, device=DEVICE, requires_grad=False
+        )
+
+        # Run WITHOUT scale_grad_by_freq
+        A1 = torch.randn(
+            rank, vocab, dtype=torch.float32, device=DEVICE, requires_grad=True
+        )
+        B1 = torch.randn(
+            hidden, rank, dtype=torch.float32, device=DEVICE, requires_grad=True
+        )
+        # Token 5 appears 3 times
+        x = torch.tensor([[5, 5, 5, 10, 20]], device=DEVICE)
+
+        out1 = LoRA_Embedding.apply(
+            x,
+            W,
+            A1,
+            B1,
+            2.0,
+            None,
+            None,
+            None,
+            2.0,
+            False,
+            False,
+        )
+        out1.sum().backward()
+        grad_no_scale = A1.grad[:, 5].clone()
+
+        # Run WITH scale_grad_by_freq
+        A2 = A1.data.clone().requires_grad_(True)
+        B2 = B1.data.clone().requires_grad_(True)
+        out2 = LoRA_Embedding.apply(
+            x,
+            W,
+            A2,
+            B2,
+            2.0,
+            None,
+            None,
+            None,
+            2.0,
+            True,
+            False,
+        )
+        out2.sum().backward()
+        grad_with_scale = A2.grad[:, 5].clone()
+
+        # With scale_grad_by_freq, token 5 (count=3) should have grad / 3
+        expected_ratio = 1.0 / 3.0
+        actual_ratio = grad_with_scale.abs().mean() / grad_no_scale.abs().mean()
+        assert abs(actual_ratio - expected_ratio) < 0.01, (
+            f"scale_grad_by_freq ratio for count=3 token: expected {expected_ratio:.3f}, "
+            f"got {actual_ratio:.3f}"
+        )
+
+
+class TestEmbeddingDropoutNotAppliedToBase:
+    """Regression: embedding dropout must NOT be applied to the base embedding
+    output — PEFT's Embedding.forward does not use lora_dropout."""
+
+    def test_kernel_matches_peft_with_dropout_config(self):
+        """Even with lora_dropout>0, embedding output should match PEFT exactly."""
+        from peft import PeftModelForCausalLM
+
+        from axolotl.kernels.lora import apply_lora_embedding
+
+        base = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=DTYPE,
+            attn_implementation="eager",
+        ).to(DEVICE)
+        lora_config = LoraConfig(
+            r=8,
+            lora_alpha=16,
+            lora_dropout=0.5,  # high dropout
+            target_modules=["embed_tokens"],
+        )
+        model = PeftModelForCausalLM(base, lora_config)
+        model.train()  # training mode — dropout would be active if applied
+
+        embed = model.model.model.embed_tokens
+        ids = torch.randint(0, 1000, (2, 16), device=DEVICE)
+
+        # Run both multiple times — if dropout were applied, results would vary
+        with torch.no_grad():
+            peft_out = embed(ids)
+            kern1 = apply_lora_embedding(embed, ids)
+            kern2 = apply_lora_embedding(embed, ids)
+
+        # Kernel should be deterministic (no dropout)
+        _compare_tensors(
+            kern1.to(peft_out.dtype),
+            kern2.to(peft_out.dtype),
+            "Embedding deterministic (no dropout)",
+            atol=0,
+            rtol=0,
+        )
+
+        # And should match PEFT
+        _compare_tensors(
+            kern1.to(peft_out.dtype),
+            peft_out,
+            "Embedding matches PEFT with dropout config",
+        )
+        del model
diff --git a/tests/e2e/kernels/test_lora_mlp_geglu_moe_correctness.py b/tests/e2e/kernels/test_lora_mlp_geglu_moe_correctness.py
new file mode 100644
index 0000000000..5ac69e5461
--- /dev/null
+++ b/tests/e2e/kernels/test_lora_mlp_geglu_moe_correctness.py
@@ -0,0 +1,138 @@
+"""Correctness of the fused GeGLU LoRA MLP kernel vs an eager autograd reference.
+
+This is the path re-enabled for MoE models by ``disable_mlp_kernel`` (the dense *shared* MLP,
+e.g. gemma4's per-layer ``Gemma4TextMLP``, gets ``lora_mlp_kernel`` while the routed experts are
+handled by the MoE kernel). Verifies forward output, input grad, and all six LoRA A/B grads match
+an eager reference — over both a bf16 base and an nf4-quantized base (the low-VRAM gemma4 case),
+since the fused kernel dequantizes the base in-kernel.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+from axolotl.kernels.lora import LoRA_MLP  # noqa: E402
+
+try:
+    import bitsandbytes.functional as bnb_F  # noqa: E402
+except Exception:  # pragma: no cover
+    bnb_F = None
+
+DEV = "cuda"
+
+
+def _geglu_eager(X, gw, gA, gB, uw, uA, uB, dw, dA, dB, s):
+    """Eager gated-GeGLU MLP with LoRA: down(gelu_tanh(gate(x)) * up(x)), gelu tanh-approx."""
+    g = X @ gw.t() + s * ((X @ gA.t()) @ gB.t())
+    u = X @ uw.t() + s * ((X @ uA.t()) @ uB.t())
+    h = F.gelu(g, approximate="tanh") * u
+    return h @ dw.t() + s * ((h @ dA.t()) @ dB.t())
+
+
+def _mk(*shape, seed):
+    g = torch.Generator(device=DEV).manual_seed(seed)
+    return (
+        torch.randn(*shape, device=DEV, dtype=torch.float16, generator=g) * 0.1
+    ).requires_grad_(True)
+
+
+@pytest.mark.parametrize("quantized", [False, True])
+def test_fused_geglu_lora_matches_eager(quantized):
+    if quantized and bnb_F is None:
+        pytest.skip("bitsandbytes required for the quantized-base case")
+
+    from axolotl.kernels.geglu import geglu_backward, geglu_forward
+
+    H, IM, r, B, T = 256, 512, 16, 2, 8  # hidden, intermediate, rank, batch, seq
+    s = 0.5
+
+    # base weights (frozen). down maps IM->H; gate/up map H->IM.
+    gw = torch.randn(IM, H, device=DEV, dtype=torch.float16) * 0.05
+    uw = torch.randn(IM, H, device=DEV, dtype=torch.float16) * 0.05
+    dw = torch.randn(H, IM, device=DEV, dtype=torch.float16) * 0.05
+
+    if quantized:
+        # nf4-quantize each base; the eager ref dequantizes the SAME packed weight (via the kernel's
+        # dequantize) so only fused-vs-eager kernel numerics differ, not the quantization.
+        from axolotl.kernels.quantize import dequantize
+
+        gw_p, gq = bnb_F.quantize_4bit(gw, quant_type="nf4", compress_statistics=True)
+        uw_p, uq = bnb_F.quantize_4bit(uw, quant_type="nf4", compress_statistics=True)
+        dw_p, dq = bnb_F.quantize_4bit(dw, quant_type="nf4", compress_statistics=True)
+        gw_ref = dequantize(gw_p.t(), gq).t().to(torch.float16)
+        uw_ref = dequantize(uw_p.t(), uq).t().to(torch.float16)
+        dw_ref = dequantize(dw_p.t(), dq).t().to(torch.float16)
+        gw_k, uw_k, dw_k = gw_p, uw_p, dw_p
+    else:
+        gq = uq = dq = None
+        gw_ref, uw_ref, dw_ref = gw, uw, dw
+        gw_k, uw_k, dw_k = gw, uw, dw
+
+    # LoRA params (seeded). down LoRA: A[r,I], B[H,r]; gate/up LoRA: A[r,H], B[I,r].
+    def fresh():
+        gA, gB = _mk(r, H, seed=1), _mk(IM, r, seed=2)
+        uA, uB = _mk(r, H, seed=3), _mk(IM, r, seed=4)
+        dA, dB = _mk(r, IM, seed=5), _mk(H, r, seed=6)
+        return gA, gB, uA, uB, dA, dB
+
+    Xv = torch.randn(B, T, H, device=DEV, dtype=torch.float16) * 0.1
+    grad_out = torch.randn(B, T, H, device=DEV, dtype=torch.float16)
+
+    # --- fused ---
+    Xf = Xv.clone().requires_grad_(True)
+    gA, gB, uA, uB, dA, dB = fresh()
+    out_f = LoRA_MLP.apply(
+        Xf,
+        None,
+        gw_k,
+        None,
+        gq,
+        gA,
+        gB,
+        s,
+        None,
+        None,
+        uw_k,
+        None,
+        uq,
+        uA,
+        uB,
+        s,
+        None,
+        None,
+        dw_k,
+        None,
+        dq,
+        dA,
+        dB,
+        s,
+        None,
+        None,
+        geglu_forward,
+        geglu_backward,
+        True,
+    )
+    out_f.backward(grad_out)
+    gf = [t.grad.clone() for t in (Xf, gA, gB, uA, uB, dA, dB)]
+
+    # --- eager ---
+    Xe = Xv.clone().requires_grad_(True)
+    eA, eB, euA, euB, edA, edB = fresh()
+    out_e = _geglu_eager(Xe, gw_ref, eA, eB, uw_ref, euA, euB, dw_ref, edA, edB, s)
+    out_e.backward(grad_out)
+    ge = [t.grad.clone() for t in (Xe, eA, eB, euA, euB, edA, edB)]
+
+    def rel(a, b):
+        return (a - b).float().abs().max().item() / max(
+            b.float().abs().max().item(), 1e-6
+        )
+
+    tol = 3e-2  # fp16 + Triton reduction order (+ nf4 dequant rounding shared by both)
+    assert torch.isfinite(out_f).all()
+    assert rel(out_f, out_e) < tol, f"forward rel={rel(out_f, out_e):.4f}"
+    names = ["dX", "dgate_A", "dgate_B", "dup_A", "dup_B", "ddown_A", "ddown_B"]
+    for n, a, b in zip(names, gf, ge, strict=True):
+        assert torch.isfinite(a).all(), f"{n} non-finite"
+        assert rel(a, b) < tol, f"{n} rel={rel(a, b):.4f}"
diff --git a/tests/e2e/kernels/test_quantize.py b/tests/e2e/kernels/test_quantize.py
new file mode 100644
index 0000000000..68c2d4e3be
--- /dev/null
+++ b/tests/e2e/kernels/test_quantize.py
@@ -0,0 +1,92 @@
+"""Tests for quantization utility functions."""
+
+import bitsandbytes as bnb
+import pytest
+import torch
+
+from axolotl.kernels.quantize import dequantize
+
+
+def _nf4_pair(shape=(64, 64), device="cuda", dtype=torch.float16, double_quant=True):
+    """Real bnb NF4 (packed_weight, quant_state) pair for the given shape."""
+    W = torch.randn(shape, device=device, dtype=dtype)
+    packed, quant_state = bnb.functional.quantize_4bit(
+        W, quant_type="nf4", compress_statistics=double_quant
+    )
+    return packed, quant_state
+
+
+def test_dequantize_null_state():
+    """dequantize returns input unchanged when quant_state is None."""
+    W = torch.randn(64, 64)
+    assert torch.equal(dequantize(W, None), W)
+
+
+def test_dequantize_shape_preservation():
+    shape = (128, 64)
+    packed, quant_state = _nf4_pair(shape)
+    result = dequantize(packed, quant_state)
+    assert result.shape == shape
+    assert result.dtype == torch.float16
+    assert result.device == packed.device
+
+
+def test_dequantize_transposed():
+    """Transposed input → transposed output (values, not just shape)."""
+    shape = (128, 64)  # non-square: catches dim-swap bugs
+    packed, quant_state = _nf4_pair(shape)
+    # packed is (4096, 1); packed.t() is (1, 4096), the leading-dim-1 signal.
+    # bnb dequants to quant_state.shape (128, 64) then returns out.t() → (64, 128).
+    expected = bnb.functional.dequantize_4bit(packed, quant_state, quant_type="nf4").t()
+    result = dequantize(packed.t(), quant_state)
+    assert tuple(result.shape) == (shape[1], shape[0])
+    torch.testing.assert_close(result, expected)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_dequantize_matches_bnb(dtype):
+    """Bit-exact vs bnb's own dequant for every supported output dtype. fp32 previously
+    fell through to the bf16 kernel writing into an fp32 buffer and produced garbage."""
+    shape = (128, 64)
+    packed, quant_state = _nf4_pair(shape, dtype=dtype)
+    expected = bnb.functional.dequantize_4bit(packed, quant_state, quant_type="nf4")
+
+    result = dequantize(packed, quant_state)
+    assert result.dtype == dtype
+    torch.testing.assert_close(result, expected)
+
+    # transposed input → transposed output (values, not just shape)
+    result_t = dequantize(packed.t(), quant_state)
+    torch.testing.assert_close(result_t, expected.t())
+
+
+def test_dequantize_non_nested():
+    """Single-quant (compress_statistics=False) falls back to bnb wrapper."""
+    shape = (128, 64)
+    packed, quant_state = _nf4_pair(shape, double_quant=False)
+    result = dequantize(packed, quant_state)
+    assert result.shape == shape
+
+
+def test_dequantize_torch_compile_nested():
+    """NF4 double-quant under torch.compile (the QLoRA hot path)."""
+    shape = (128, 64)
+    packed, quant_state = _nf4_pair(shape, double_quant=True)
+
+    eager = dequantize(packed, quant_state)
+    compiled = torch.compile(dequantize)(packed, quant_state)
+
+    assert compiled.shape == eager.shape
+    assert compiled.dtype == eager.dtype
+    torch.testing.assert_close(compiled, eager)
+
+
+def test_dequantize_torch_compile_non_nested():
+    """torch.compile also works for the single-quant fallback path."""
+    shape = (128, 64)
+    packed, quant_state = _nf4_pair(shape, double_quant=False)
+
+    eager = dequantize(packed, quant_state)
+    compiled = torch.compile(dequantize)(packed, quant_state)
+
+    torch.testing.assert_close(compiled, eager)
diff --git a/tests/e2e/kernels/test_swiglu.py b/tests/e2e/kernels/test_swiglu.py
new file mode 100644
index 0000000000..58d5e04a73
--- /dev/null
+++ b/tests/e2e/kernels/test_swiglu.py
@@ -0,0 +1,77 @@
+"""Tests for SwiGLU activation function Triton kernels."""
+
+import torch
+import torch.nn.functional as F
+
+from axolotl.kernels.swiglu import swiglu_backward, swiglu_forward
+
+
+def test_swiglu_forward_shape():
+    """Test that SwiGLU forward pass preserves expected shapes"""
+    batch, seq_len, hidden_dim = 2, 3, 64
+    gate = torch.randn(batch, seq_len, hidden_dim, device="cuda")
+    up = torch.randn(batch, seq_len, hidden_dim, device="cuda")
+
+    out = swiglu_forward(gate, up)
+    assert out.shape == (batch, seq_len, hidden_dim)
+    assert out.dtype == gate.dtype
+    assert out.device == gate.device
+
+
+def test_swiglu_forward_values():
+    """Test SwiGLU forward pass matches PyTorch reference implementation"""
+    gate = torch.randn(2, 3, 64, device="cuda")
+    up = torch.randn(2, 3, 64, device="cuda")
+
+    # Custom implementation
+    triton_out = swiglu_forward(gate.clone(), up.clone())
+
+    # PyTorch reference
+    torch_out = F.silu(gate) * up
+
+    assert torch.allclose(triton_out, torch_out, rtol=1e-3)
+
+
+def test_swiglu_backward():
+    """Test SwiGLU backward pass matches PyTorch autograd"""
+    gate = torch.randn(2, 3, 64, device="cuda", requires_grad=True)
+    up = torch.randn(2, 3, 64, device="cuda", requires_grad=True)
+    grad_output = torch.randn(2, 3, 64, device="cuda")
+
+    # PyTorch reference - compute intermediates
+    silu_gate = F.silu(gate)
+    torch_out = silu_gate * up
+    torch_out.backward(grad_output)
+
+    # Custom backward pass
+    gate_clone = gate.clone().detach()
+    up_clone = up.clone().detach()
+    grad_output_clone = grad_output.clone()
+
+    h, our_grad_gate, our_grad_up = swiglu_backward(
+        grad_output_clone, gate_clone, up_clone
+    )
+
+    # Compare outputs and gradients
+    assert torch.allclose(h, torch_out, rtol=1e-3)
+    assert torch.allclose(our_grad_gate, gate.grad, rtol=1e-3)
+    assert torch.allclose(our_grad_up, up.grad, rtol=1e-3)
+
+
+def test_swiglu_inplace_preservation():
+    """Test that SwiGLU backward doesn't modify original tensors unexpectedly"""
+    gate = torch.randn(2, 3, 64, device="cuda")
+    up = torch.randn(2, 3, 64, device="cuda")
+    grad_output = torch.randn(2, 3, 64, device="cuda")
+
+    gate_copy = gate.clone()
+    up_copy = up.clone()
+    grad_copy = grad_output.clone()
+
+    swiglu_backward(grad_output, gate, up)
+
+    assert not torch.equal(gate, gate_copy), "Gate should be modified in-place"
+    assert not torch.equal(up, up_copy), "Up should be modified in-place"
+    assert not torch.equal(grad_output, grad_copy), (
+        "Grad output should be modified in-place"
+    )
diff --git a/tests/e2e/multigpu/__init__.py b/tests/e2e/multigpu/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/e2e/multigpu/_fp32_norms_dtype_capture.py b/tests/e2e/multigpu/_fp32_norms_dtype_capture.py
new file mode 100644
index 0000000000..a0dda96b31
--- /dev/null
+++ b/tests/e2e/multigpu/_fp32_norms_dtype_capture.py
@@ -0,0 +1,58 @@
+"""Test-only plugin that captures param dtypes after the first optimizer step
+and dumps them as JSON to ``$FP32_NORMS_DTYPE_DUMP_PATH``.
+
+Loaded via ``plugins: [tests.e2e.multigpu._fp32_norms_dtype_capture.DtypeCapturePlugin]``
+in the test yaml config; the dump path is the contract between the subprocess
+and the outer pytest function. Rank 0 only — dtype is identical across ranks.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+
+import torch
+from transformers.trainer_callback import TrainerCallback
+
+from axolotl.integrations.base import BasePlugin
+
+
+def _dtype_name(dtype: torch.dtype) -> str:
+    return str(dtype).removeprefix("torch.")
+
+
+class _DtypeCaptureCallback(TrainerCallback):
+    """Capture norm vs non-norm param dtypes after step 1, dump to JSON, exit."""
+
+    def on_step_end(self, args, state, control, model=None, **kwargs):  # type: ignore[override]
+        if state.global_step != 1 or model is None:
+            return
+        # Rank 0 only — every rank sees the same dtype info under FSDP2.
+        if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
+            return
+        dump_path = os.environ.get("FP32_NORMS_DTYPE_DUMP_PATH")
+        if not dump_path:
+            return
+
+        norm_dtypes: dict[str, str] = {}
+        non_norm_dtypes: dict[str, str] = {}
+        for name, param in model.named_parameters():
+            entry = (name, _dtype_name(param.dtype))
+            if "norm" in name.lower():
+                norm_dtypes[entry[0]] = entry[1]
+            else:
+                non_norm_dtypes[entry[0]] = entry[1]
+
+        with open(dump_path, "w", encoding="utf-8") as fout:
+            json.dump(
+                {"norms": norm_dtypes, "non_norms": non_norm_dtypes},
+                fout,
+                indent=2,
+            )
+
+
+class DtypeCapturePlugin(BasePlugin):
+    """Plugin that registers :class:`_DtypeCaptureCallback` with the trainer."""
+
+    def add_callbacks_pre_trainer(self, cfg, model):  # type: ignore[override]
+        return [_DtypeCaptureCallback()]
diff --git a/tests/e2e/multigpu/patched/__init__.py b/tests/e2e/multigpu/patched/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/e2e/multigpu/patched/test_sp.py b/tests/e2e/multigpu/patched/test_sp.py
new file mode 100644
index 0000000000..cfd4369304
--- /dev/null
+++ b/tests/e2e/multigpu/patched/test_sp.py
@@ -0,0 +1,137 @@
+"""E2E tests for sequence parallelism"""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from ...utils import check_tensorboard
+
+
+class TestSequenceParallelism:
+    """Test case for training with sequence parallelism enabled"""
+
+    def _run_sequence_parallel_test(
+        self,
+        temp_dir,
+        sample_packing=True,
+        micro_batch_size=1,
+        pad_to_sequence_len=True,
+        ring_attn_func=None,
+        threshold=2.0,
+    ):
+        """Helper method to run sequence parallel tests with different configurations"""
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "load_in_8bit": False,
+                "load_in_4bit": True,
+                "strict": False,
+                "sequence_len": 2048,
+                "adapter": "qlora",
+                "sample_packing": sample_packing,
+                "eval_sample_packing": sample_packing,
+                "pad_to_sequence_len": pad_to_sequence_len,
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "lora_modules_to_save": ["embed_tokens", "lm_head"],
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 8,
+                "micro_batch_size": micro_batch_size,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "loss_watchdog_threshold": 5.0,
+                "loss_watchdog_patience": 3,
+                "bf16": "auto",
+                "warmup_steps": 1,
+                "saves_per_epoch": 1,
+                "logging_steps": 1,
+                "weight_decay": 0.0,
+                "use_tensorboard": True,
+                "context_parallel_size": 2,
+                "ring_attn_func": ring_attn_func,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "accelerate",
+                "launch",
+                "--num-processes",
+                "2",
+                "--main_process_port",
+                f"{get_torch_dist_unique_port()}",
+                "-m",
+                "axolotl.cli.train",
+                str(Path(temp_dir) / "config.yaml"),
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs",
+            "train/train_loss",
+            threshold,
+            "Train Loss (%s) is too high",
+        )
+
+    @pytest.mark.skip(
+        reason="ring_flash_attn w transformers imports unmaintained upstream",
+    )
+    @pytest.mark.parametrize(
+        "sample_packing, micro_batch_size, pad_to_sequence_len, ring_attn_func, threshold",
+        [
+            (True, 1, True, None, 2.5),  # defaults to varlen_llama3 ring_attn_func
+            (False, 2, True, None, 2.5),  # defaults to batch_ring ring_attn_func
+            # (False, 2, True, "batch_zigzag", 2.5),
+            # (False, 2, False, None, 2.65),  # defaults to batch_ring ring_attn_func
+        ],
+        ids=[
+            "sample_packing, varlen_llama3 ring_attn_func",
+            "no sample_packing, pad_to_sequence_len, batch_ring ring_attn_func",
+            # "no sample_packing, no pad_to_sequence_len, batch_zigzag ring_attn_func",
+            # "no sample_packing, no pad_to_sequence_len, batch_ring ring_attn_func",
+        ],
+    )
+    def test_sequence_parallel_training(
+        self,
+        temp_dir,
+        sample_packing,
+        micro_batch_size,
+        pad_to_sequence_len,
+        ring_attn_func,
+        threshold,
+    ):
+        """Test sequence parallel training with different configurations"""
+        self._run_sequence_parallel_test(
+            temp_dir,
+            sample_packing=sample_packing,
+            micro_batch_size=micro_batch_size,
+            pad_to_sequence_len=pad_to_sequence_len,
+            ring_attn_func=ring_attn_func,
+            threshold=threshold,
+        )
diff --git a/tests/e2e/multigpu/solo/__init__.py b/tests/e2e/multigpu/solo/__init__.py
new file mode 100644
index 0000000000..ed1ba7dc62
--- /dev/null
+++ b/tests/e2e/multigpu/solo/__init__.py
@@ -0,0 +1,2 @@
+# Tests under this directory should get run "solo" on their own as they
+# seem to cause issues when run in the same batch as other tests.
diff --git a/tests/e2e/multigpu/solo/test_flex.py b/tests/e2e/multigpu/solo/test_flex.py
new file mode 100644
index 0000000000..881d75c25f
--- /dev/null
+++ b/tests/e2e/multigpu/solo/test_flex.py
@@ -0,0 +1,90 @@
+"""
+E2E tests for multigpu lora tinyllama
+"""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from huggingface_hub import snapshot_download
+from transformers.testing_utils import get_torch_dist_unique_port
+from transformers.utils import is_torch_bf16_gpu_available
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_tensorboard, require_torch_2_6_0
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_model():
+    # download the model
+    snapshot_download("HuggingFaceTB/SmolLM2-135M")
+
+
+class TestPackedFlex:
+    """
+    Test case for Packed training of llama models
+    """
+
+    @require_torch_2_6_0
+    def test_loss_llama(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "sample_packing": True,
+                "flex_attention": True,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "max_steps": 2,
+                "use_tensorboard": True,
+                "save_strategy": "no",
+                "save_first_step": False,
+            }
+        )
+        if is_torch_bf16_gpu_available():
+            cfg.bf16 = True
+        else:
+            cfg.fp16 = True
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.1, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/multigpu/solo/test_gdpo.py b/tests/e2e/multigpu/solo/test_gdpo.py
new file mode 100644
index 0000000000..2014f7f5e3
--- /dev/null
+++ b/tests/e2e/multigpu/solo/test_gdpo.py
@@ -0,0 +1,538 @@
+"""
+GDPO test suite
+
+GDPO uses TRL's multi_objective_aggregation="normalize_then_sum" for
+per-reward normalization in multi-reward RL training.
+"""
+
+import os
+import random
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.multigpu.solo.test_grpo import recursive_kill, start_vllm
+from tests.e2e.utils import require_vllm
+
+
+@pytest.mark.skip(reason="flaky vllm tests in modal")
+class TestGDPO:
+    """Test case for GDPO training using TRL's native multi-objective aggregation."""
+
+    def _utils_write_yaml_and_rewards(self, cfg, temp_dir, suffix=""):
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+        with open(f"rewards_gdpo_{suffix}.py", "w", encoding="utf-8") as fout:
+            fout.write(
+                """import random
+
+def format_reward(prompts, completions, **kwargs) -> list[float]:
+    return [1.0 if len(c) > 10 else 0.0 for c in completions]
+
+def correctness_reward(prompts, completions, **kwargs) -> list[float]:
+    return [random.uniform(-1, 3) for _ in completions]
+
+def safety_reward(prompts, completions, **kwargs) -> list[float]:
+    return [1.0 if 'error' not in c.lower() else 0.0 for c in completions]
+
+def single_reward(prompts, completions, **kwargs) -> list[float]:
+    return [random.uniform(0, 1) for _ in completions]
+
+def oai_gsm8k_transform(cfg, *args, **kwargs):
+    def transform_fn(example, tokenizer=None):
+        label = example["answer"].split("####")[-1].strip().replace(",", "")
+        return {
+            "prompt": [{"role": "user", "content": example["question"]}],
+            "answer": label,
+        }
+    return transform_fn, {"remove_columns": ["question"]}
+"""
+            )
+
+    @pytest.mark.parametrize("num_gpus", [1, 2])
+    @require_vllm
+    def test_gdpo_multi_reward_lora(self, temp_dir, num_gpus):
+        """Test GDPO with multiple reward functions using LoRA."""
+        rnd_suffix = str(random.randint(1000, 9999))
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "chat_template": "llama3",
+                "rl": "gdpo",
+                "trl": {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": True,
+                    "num_generations": 4,
+                    "reward_funcs": [
+                        f"rewards_gdpo_{rnd_suffix}.format_reward",
+                        f"rewards_gdpo_{rnd_suffix}.correctness_reward",
+                    ],
+                    "reward_weights": [1.0, 2.0],
+                    "scale_rewards": True,
+                },
+                "vllm": {
+                    "max_model_len": 800,
+                    "enable_prefix_caching": True,
+                },
+                "datasets": [
+                    {
+                        "path": "openai/gsm8k",
+                        "name": "main",
+                        "type": f"rewards_gdpo_{rnd_suffix}.oai_gsm8k_transform",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "max_steps": 3,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "warmup_steps": 10,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "save_safetensors": True,
+                "bf16": "auto",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        self._utils_write_yaml_and_rewards(cfg, temp_dir, suffix=rnd_suffix)
+
+        current_env = os.environ.copy()
+        env = {
+            "NCCL_P2P_LEVEL": "LOC",
+            **current_env,
+            "CUDA_VISIBLE_DEVICES": "1",
+        }
+        vllm_process = start_vllm(
+            cfg.base_model,
+            env=env,
+            quiet=True,
+            wait=300,
+            gpu_memory_utilization=0.15,
+            max_model_len=cfg.vllm.max_model_len,
+            enable_prefix_caching=cfg.vllm.enable_prefix_caching,
+            host="0.0.0.0",
+            port=8000,
+        )
+
+        try:
+            execute_subprocess_async(
+                [
+                    "axolotl",
+                    "train",
+                    str(Path(temp_dir) / "config.yaml"),
+                    "--num-processes",
+                    str(num_gpus),
+                    "--main-process-port",
+                    f"{get_torch_dist_unique_port()}",
+                ],
+                env={
+                    "NCCL_P2P_LEVEL": "LOC",
+                    "NCCL_DEBUG": "INFO",
+                    **current_env,
+                },
+            )
+        finally:
+            recursive_kill(vllm_process)
+
+    @require_vllm
+    def test_gdpo_three_rewards(self, temp_dir):
+        """Test GDPO with three reward functions (format, correctness, safety)."""
+        rnd_suffix = str(random.randint(1000, 9999))
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "chat_template": "llama3",
+                "rl": "gdpo",
+                "trl": {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": True,
+                    "num_generations": 4,
+                    "reward_funcs": [
+                        f"rewards_gdpo_{rnd_suffix}.format_reward",
+                        f"rewards_gdpo_{rnd_suffix}.correctness_reward",
+                        f"rewards_gdpo_{rnd_suffix}.safety_reward",
+                    ],
+                    "reward_weights": [1.0, 2.0, 1.5],
+                },
+                "vllm": {
+                    "max_model_len": 800,
+                    "enable_prefix_caching": True,
+                },
+                "datasets": [
+                    {
+                        "path": "openai/gsm8k",
+                        "name": "main",
+                        "type": f"rewards_gdpo_{rnd_suffix}.oai_gsm8k_transform",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "max_steps": 3,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "warmup_steps": 10,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "save_safetensors": True,
+                "bf16": "auto",
+            }
+        )
+
+        self._utils_write_yaml_and_rewards(cfg, temp_dir, suffix=rnd_suffix)
+
+        current_env = os.environ.copy()
+        env = {
+            "NCCL_P2P_LEVEL": "LOC",
+            **current_env,
+            "CUDA_VISIBLE_DEVICES": "1",
+        }
+        vllm_process = start_vllm(
+            cfg.base_model,
+            env=env,
+            quiet=True,
+            wait=300,
+            gpu_memory_utilization=0.15,
+            max_model_len=cfg.vllm.max_model_len,
+            enable_prefix_caching=cfg.vllm.enable_prefix_caching,
+            host="0.0.0.0",
+            port=8000,
+        )
+
+        try:
+            execute_subprocess_async(
+                [
+                    "axolotl",
+                    "train",
+                    str(Path(temp_dir) / "config.yaml"),
+                    "--num-processes",
+                    "1",
+                    "--main-process-port",
+                    f"{get_torch_dist_unique_port()}",
+                ],
+                env={
+                    "NCCL_P2P_LEVEL": "LOC",
+                    "NCCL_DEBUG": "INFO",
+                    **current_env,
+                },
+            )
+        finally:
+            recursive_kill(vllm_process)
+
+    @require_vllm
+    def test_gdpo_single_reward_fallback(self, temp_dir):
+        """Test GDPO with single reward."""
+        rnd_suffix = str(random.randint(1000, 9999))
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "chat_template": "llama3",
+                "rl": "gdpo",
+                "trl": {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": True,
+                    "num_generations": 4,
+                    "reward_funcs": [
+                        f"rewards_gdpo_{rnd_suffix}.single_reward",
+                    ],
+                    "reward_weights": [1.0],
+                },
+                "vllm": {
+                    "max_model_len": 800,
+                    "enable_prefix_caching": True,
+                },
+                "datasets": [
+                    {
+                        "path": "openai/gsm8k",
+                        "name": "main",
+                        "type": f"rewards_gdpo_{rnd_suffix}.oai_gsm8k_transform",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "max_steps": 3,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "warmup_steps": 10,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "save_safetensors": True,
+                "bf16": "auto",
+            }
+        )
+
+        self._utils_write_yaml_and_rewards(cfg, temp_dir, suffix=rnd_suffix)
+
+        current_env = os.environ.copy()
+        env = {
+            "NCCL_P2P_LEVEL": "LOC",
+            **current_env,
+            "CUDA_VISIBLE_DEVICES": "1",
+        }
+        vllm_process = start_vllm(
+            cfg.base_model,
+            env=env,
+            quiet=True,
+            wait=300,
+            gpu_memory_utilization=0.15,
+            max_model_len=cfg.vllm.max_model_len,
+            enable_prefix_caching=cfg.vllm.enable_prefix_caching,
+            host="0.0.0.0",
+            port=8000,
+        )
+
+        try:
+            execute_subprocess_async(
+                [
+                    "axolotl",
+                    "train",
+                    str(Path(temp_dir) / "config.yaml"),
+                    "--num-processes",
+                    "1",
+                    "--main-process-port",
+                    f"{get_torch_dist_unique_port()}",
+                ],
+                env={
+                    "NCCL_P2P_LEVEL": "LOC",
+                    "NCCL_DEBUG": "INFO",
+                    **current_env,
+                },
+            )
+        finally:
+            recursive_kill(vllm_process)
+
+    @require_vllm
+    def test_gdpo_fft(self, temp_dir):
+        """Test GDPO with full fine-tuning (no adapter)."""
+        rnd_suffix = str(random.randint(1000, 9999))
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "chat_template": "llama3",
+                "rl": "gdpo",
+                "trl": {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": True,
+                    "num_generations": 4,
+                    "reward_funcs": [
+                        f"rewards_gdpo_{rnd_suffix}.format_reward",
+                        f"rewards_gdpo_{rnd_suffix}.correctness_reward",
+                    ],
+                    "reward_weights": [1.0, 2.0],
+                },
+                "vllm": {
+                    "max_model_len": 800,
+                    "enable_prefix_caching": True,
+                },
+                "datasets": [
+                    {
+                        "path": "openai/gsm8k",
+                        "name": "main",
+                        "type": f"rewards_gdpo_{rnd_suffix}.oai_gsm8k_transform",
+                    },
+                ],
+                # No adapter - full fine-tuning
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "max_steps": 3,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "warmup_steps": 10,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "save_safetensors": True,
+                "bf16": "auto",
+            }
+        )
+
+        self._utils_write_yaml_and_rewards(cfg, temp_dir, suffix=rnd_suffix)
+
+        current_env = os.environ.copy()
+        env = {
+            "NCCL_P2P_LEVEL": "LOC",
+            **current_env,
+            "CUDA_VISIBLE_DEVICES": "1",
+        }
+        vllm_process = start_vllm(
+            cfg.base_model,
+            env=env,
+            quiet=True,
+            wait=300,
+            gpu_memory_utilization=0.15,
+            max_model_len=cfg.vllm.max_model_len,
+            enable_prefix_caching=cfg.vllm.enable_prefix_caching,
+            host="0.0.0.0",
+            port=8000,
+        )
+
+        try:
+            execute_subprocess_async(
+                [
+                    "axolotl",
+                    "train",
+                    str(Path(temp_dir) / "config.yaml"),
+                    "--num-processes",
+                    "1",
+                    "--main-process-port",
+                    f"{get_torch_dist_unique_port()}",
+                ],
+                env={
+                    "NCCL_P2P_LEVEL": "LOC",
+                    "NCCL_DEBUG": "INFO",
+                    **current_env,
+                },
+            )
+        finally:
+            recursive_kill(vllm_process)
+
+    @require_vllm
+    def test_gdpo_sequence_parallel(self, temp_dir):
+        """Test GDPO with sequence parallelism."""
+        rnd_suffix = str(random.randint(1000, 9999))
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "chat_template": "llama3",
+                "rl": "gdpo",
+                "context_parallel_size": 2,
+                "trl": {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": True,
+                    "num_generations": 4,
+                    "reward_funcs": [
+                        f"rewards_gdpo_{rnd_suffix}.format_reward",
+                        f"rewards_gdpo_{rnd_suffix}.correctness_reward",
+                    ],
+                    "reward_weights": [1.0, 2.0],
+                },
+                "vllm": {
+                    "max_model_len": 800,
+                    "enable_prefix_caching": True,
+                },
+                "datasets": [
+                    {
+                        "path": "openai/gsm8k",
+                        "name": "main",
+                        "type": f"rewards_gdpo_{rnd_suffix}.oai_gsm8k_transform",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "max_steps": 3,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "warmup_steps": 10,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "save_safetensors": True,
+                "bf16": "auto",
+            }
+        )
+
+        self._utils_write_yaml_and_rewards(cfg, temp_dir, suffix=rnd_suffix)
+
+        current_env = os.environ.copy()
+        env = {
+            "NCCL_P2P_LEVEL": "LOC",
+            **current_env,
+            "CUDA_VISIBLE_DEVICES": "1",
+        }
+        vllm_process = start_vllm(
+            cfg.base_model,
+            env=env,
+            quiet=True,
+            wait=300,
+            gpu_memory_utilization=0.15,
+            max_model_len=cfg.vllm.max_model_len,
+            enable_prefix_caching=cfg.vllm.enable_prefix_caching,
+            host="0.0.0.0",
+            port=8000,
+        )
+
+        try:
+            execute_subprocess_async(
+                [
+                    "axolotl",
+                    "train",
+                    str(Path(temp_dir) / "config.yaml"),
+                    "--num-processes",
+                    "2",
+                    "--main-process-port",
+                    f"{get_torch_dist_unique_port()}",
+                ],
+                env={
+                    "NCCL_P2P_LEVEL": "LOC",
+                    "NCCL_DEBUG": "INFO",
+                    **current_env,
+                },
+            )
+        finally:
+            recursive_kill(vllm_process)
diff --git a/tests/e2e/multigpu/solo/test_grpo.py b/tests/e2e/multigpu/solo/test_grpo.py
new file mode 100644
index 0000000000..8d0bc3f680
--- /dev/null
+++ b/tests/e2e/multigpu/solo/test_grpo.py
@@ -0,0 +1,453 @@
+"""
+GRPO test suite
+"""
+
+import os
+import random
+import subprocess  # nosec B404
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+import psutil
+import pytest
+import requests
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import require_vllm
+
+
+def start_vllm(
+    model: str, env: dict, wait: int | None = None, quiet=False, **kwargs
+) -> subprocess.Popen:
+    """
+    helper function to start the VLLM server in the background, mostly for testing purposes
+    """
+    cmd = [sys.executable, "-m", "trl.scripts.vllm_serve", "--model", model]
+
+    if tensor_parallel_size := kwargs.get("tensor_parallel_size"):
+        cmd.extend(["--tensor-parallel-size", str(tensor_parallel_size)])
+    if host := kwargs.get("host"):
+        cmd.extend(["--host", host])
+    if port := kwargs.get("port"):
+        cmd.extend(["--port", str(port)])
+    if gpu_memory_utilization := kwargs.get("gpu_memory_utilization"):
+        cmd.extend(["--gpu-memory-utilization", str(gpu_memory_utilization)])
+    if dtype := kwargs.get("dtype"):
+        cmd.extend(["--dtype", dtype])
+    if max_model_len := kwargs.get("max_model_len"):
+        cmd.extend(["--max-model-len", str(max_model_len)])
+    if kwargs.get("enable_prefix_caching"):
+        cmd.extend(["--enable-prefix-caching", "True"])
+
+    # print out the command to be executed
+    print(" ".join(cmd))
+
+    vllm_logging_json = Path(tempfile.mkdtemp()) / "vllm_logging.json"
+    with open(vllm_logging_json, "w", encoding="utf-8") as temp_file:
+        temp_file.write(
+            """{
+  "formatters": {
+    "json": {
+      "class": "pythonjsonlogger.jsonlogger.JsonFormatter"
+    }
+  },
+  "handlers": {
+    "file": {
+      "class": "logging.FileHandler",
+      "formatter": "json",
+      "level": "DEBUG",
+      "filename": "/tmp/vllm.log",
+      "mode": "a"
+    }
+  },
+  "loggers": {
+    "vllm": {
+      "handlers": ["file"],
+      "level": "DEBUG",
+      "propagate": false
+    }
+  },
+  "version": 1
+}"""
+        )
+
+    cmd_env = env.copy()
+    cmd_env.update({"VLLM_LOGGING_CONFIG_PATH": vllm_logging_json})
+    # start `trl vllm-serve` command in the background and capture the process id
+    process = subprocess.Popen(
+        cmd,
+        env=cmd_env,
+        stdout=subprocess.DEVNULL if quiet else subprocess.PIPE,
+        stderr=subprocess.DEVNULL if quiet else subprocess.PIPE,
+    )  # nosec B603
+
+    # print out the process id so the user can easily kill it later
+    print(f"VLLM server process started (PID: {process.pid})")
+
+    # wait until the http server is ready, even if it 404s, but timeout after 60 seconds
+    period_seconds = 5
+    started = False
+    if wait and host and port:
+        for i in range(0, int(wait), period_seconds):
+            try:
+                response = requests.get(f"http://{host}:{port}", timeout=1)
+                print(f"{i}: VLLM server (status: {response.status_code})")
+                if int(response.status_code) in [200, 404]:
+                    started = True
+                    break
+            except requests.exceptions.RequestException as exc:
+                print(f"{i}: VLLM server failed to start: {str(exc)}")
+
+            # also check if the process.pid is still running
+            if process.poll() is not None:
+                break
+
+            time.sleep(period_seconds)
+
+    if wait and not started:
+        print(
+            f"VLLM server process did not start within {wait} seconds. Please check your server logs."
+        )
+        recursive_kill(process)
+        with open("/tmp/vllm.log", "r", encoding="utf-8") as log_file:
+            print(log_file.read())
+        try:
+            os.remove("/tmp/vllm.log")
+        except FileNotFoundError:
+            pass
+        raise RuntimeError(f"VLLM server process did not start within {wait} seconds.")
+
+    # return the process
+    return process
+
+
+def recursive_kill(process: subprocess.Popen):
+    """
+    Recursively kill a process and its children
+    """
+    process = psutil.Process(process.pid)
+    for child in psutil.Process(process.pid).children(recursive=True):
+        child.terminate()
+        child.kill()
+        os.kill(child.pid, 9)
+    process.terminate()
+    process.kill()
+    os.kill(process.pid, 9)
+
+
+@pytest.mark.skip(reason="flaky vllm tests in modal")
+class TestGRPO:
+    """
+    Test case for GRPO training using multiple GPUs
+    """
+
+    def _utils_write_yaml_and_rewards(self, cfg, temp_dir, suffix=""):
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+        with open(f"rewards_{suffix}.py", "w", encoding="utf-8") as fout:
+            fout.write(
+                """import random
+def rand_reward_func(completions, **kwargs) -> list[float]:
+    return [random.uniform(0, 1) for _ in completions]
+
+def oai_gsm8k_transform(cfg, *args, **kwargs):
+    def transform_fn(example, tokenizer=None):
+        label = example["answer"].split("####")[-1].strip().replace(",", "")
+        return {
+            "prompt": [{"role": "user", "content": example["question"]},],
+            "answer": label,
+        }
+    return transform_fn, {"remove_columns": ["question"]}
+"""
+            )
+
+    @pytest.mark.parametrize(
+        "num_gpus",
+        [1, 2],
+    )
+    @require_vllm
+    def test_llama_dora(self, temp_dir, num_gpus):
+        rnd_reward_suffix = str(random.randint(1000, 9999))
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "chat_template": "llama3",
+                "rl": "grpo",
+                "trl": {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": True,
+                    "num_generations": 4,
+                    "reward_funcs": [f"rewards_{rnd_reward_suffix}.rand_reward_func"],
+                },
+                "vllm": {
+                    "max_model_len": 800,
+                    "enable_prefix_caching": True,
+                },
+                "datasets": [
+                    {
+                        "path": "openai/gsm8k",
+                        "name": "main",
+                        "type": f"rewards_{rnd_reward_suffix}.oai_gsm8k_transform",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "peft_use_dora": True,
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "max_steps": 3,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "warmup_steps": 10,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        self._utils_write_yaml_and_rewards(cfg, temp_dir, suffix=rnd_reward_suffix)
+
+        current_env = os.environ.copy()
+        env = {
+            "NCCL_P2P_LEVEL": "LOC",
+            **current_env,
+            "CUDA_VISIBLE_DEVICES": "1",
+        }
+        vllm_process = start_vllm(
+            cfg.base_model,
+            env=env,
+            quiet=True,
+            wait=300,
+            gpu_memory_utilization=0.15,
+            max_model_len=cfg.vllm.max_model_len,
+            enable_prefix_caching=cfg.vllm.enable_prefix_caching,
+            host="0.0.0.0",
+            port=8000,
+        )
+
+        try:
+            execute_subprocess_async(
+                [
+                    "axolotl",
+                    "train",
+                    str(Path(temp_dir) / "config.yaml"),
+                    "--num-processes",
+                    str(num_gpus),
+                    "--main-process-port",
+                    f"{get_torch_dist_unique_port()}",
+                ],
+                env={
+                    "NCCL_P2P_LEVEL": "LOC",
+                    "NCCL_DEBUG": "INFO",
+                    **current_env,
+                },
+            )
+        finally:
+            (recursive_kill(vllm_process))
+
+    @require_vllm
+    def test_llama_lora_sp(self, temp_dir):
+        rnd_reward_suffix = str(random.randint(1000, 9999))
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "chat_template": "llama3",
+                "rl": "grpo",
+                "trl": {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": True,
+                    "num_generations": 4,
+                    "reward_funcs": [f"rewards_{rnd_reward_suffix}.rand_reward_func"],
+                },
+                "vllm": {
+                    "max_model_len": 800,
+                    "enable_prefix_caching": True,
+                },
+                "datasets": [
+                    {
+                        "path": "openai/gsm8k",
+                        "name": "main",
+                        "type": f"rewards_{rnd_reward_suffix}.oai_gsm8k_transform",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "context_parallel_size": 2,
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "max_steps": 3,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "warmup_steps": 10,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        self._utils_write_yaml_and_rewards(cfg, temp_dir, suffix=rnd_reward_suffix)
+
+        current_env = os.environ.copy()
+        env = {
+            "NCCL_P2P_LEVEL": "LOC",
+            **current_env,
+            "CUDA_VISIBLE_DEVICES": "1",
+        }
+        vllm_process = start_vllm(
+            cfg.base_model,
+            env=env,
+            quiet=True,
+            wait=300,
+            gpu_memory_utilization=0.15,
+            max_model_len=cfg.vllm.max_model_len,
+            enable_prefix_caching=cfg.vllm.enable_prefix_caching,
+            host="0.0.0.0",
+            port=8000,
+        )
+
+        try:
+            execute_subprocess_async(
+                [
+                    "axolotl",
+                    "train",
+                    str(Path(temp_dir) / "config.yaml"),
+                    "--num-processes",
+                    str(2),
+                    "--main-process-port",
+                    f"{get_torch_dist_unique_port()}",
+                ],
+                env={
+                    "NCCL_P2P_LEVEL": "LOC",
+                    "NCCL_DEBUG": "INFO",
+                    **current_env,
+                },
+            )
+        finally:
+            recursive_kill(vllm_process)
+
+    @pytest.mark.parametrize(
+        "num_gpus",
+        [1, 2],
+    )
+    @require_vllm
+    def test_llama_fft(self, temp_dir, num_gpus):
+        rnd_reward_suffix = str(random.randint(1000, 9999))
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "chat_template": "llama3",
+                "rl": "grpo",
+                "trl": {
+                    "beta": 0.001,
+                    "max_completion_length": 256,
+                    "use_vllm": True,
+                    "num_generations": 4,
+                    "reward_funcs": [f"rewards_{rnd_reward_suffix}.rand_reward_func"],
+                },
+                "vllm": {
+                    "max_model_len": 800,
+                    "enable_prefix_caching": True,
+                },
+                "datasets": [
+                    {
+                        "path": "openai/gsm8k",
+                        "name": "main",
+                        "type": f"rewards_{rnd_reward_suffix}.oai_gsm8k_transform",
+                    },
+                ],
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "max_steps": 3,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "warmup_steps": 10,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        self._utils_write_yaml_and_rewards(cfg, temp_dir, suffix=rnd_reward_suffix)
+
+        current_env = os.environ.copy()
+        env = {
+            "NCCL_P2P_LEVEL": "LOC",  # nccl can be brittle, assume P2P isn't reliable
+            **current_env,
+            "CUDA_VISIBLE_DEVICES": "1",
+        }
+        vllm_process = start_vllm(
+            cfg.base_model,
+            env=env,
+            quiet=True,
+            wait=300,
+            gpu_memory_utilization=0.15,
+            max_model_len=cfg.vllm.max_model_len,
+            enable_prefix_caching=cfg.vllm.enable_prefix_caching,
+            host="0.0.0.0",
+            port=8000,
+        )
+
+        try:
+            execute_subprocess_async(
+                [
+                    "axolotl",
+                    "train",
+                    str(Path(temp_dir) / "config.yaml"),
+                    "--num-processes",
+                    str(num_gpus),
+                    "--main-process-port",
+                    f"{get_torch_dist_unique_port()}",
+                ],
+                env={
+                    "NCCL_P2P_LEVEL": "LOC",
+                    "NCCL_DEBUG": "INFO",
+                    **current_env,
+                },
+            )
+        finally:
+            recursive_kill(vllm_process)
diff --git a/tests/e2e/multigpu/test_dist_muon_fsdp2.py b/tests/e2e/multigpu/test_dist_muon_fsdp2.py
new file mode 100644
index 0000000000..05841bb64a
--- /dev/null
+++ b/tests/e2e/multigpu/test_dist_muon_fsdp2.py
@@ -0,0 +1,169 @@
+"""Test module for DistMuon optimizer with FSDP2 multi-GPU functionality."""
+
+from pathlib import Path
+
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_tensorboard_loss_decreased, require_torch_2_7_0
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+def verify_training_success(temp_dir):
+    """Verify that training completed successfully — artifacts, no-NaN, loss
+    stayed in qwen2-pretraining scale (tiny-qwen2-129m final pretrain CE ~3.92).
+    """
+    output_path = Path(temp_dir)
+
+    model_files = list(output_path.glob("*.bin")) + list(
+        output_path.glob("*.safetensors")
+    )
+    assert len(model_files) > 0, "No model files found - training may have failed"
+
+    checkpoint_files = list(output_path.glob("checkpoint-*"))
+    assert len(checkpoint_files) > 0, (
+        "No checkpoint files found - training may have failed"
+    )
+
+    check_tensorboard_loss_decreased(
+        temp_dir + "/runs",
+        initial_window=10,
+        final_window=10,
+        max_initial=5.0,
+        max_final=4.7,
+    )
+
+
+class TestDistMuon:
+    """Test class for DistMuon optimizer with FSDP2 functionality."""
+
+    @require_torch_2_7_0
+    def test_fft_sft(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 2e-3,
+                "optimizer": "muon",
+                "weight_decay": 0.01,
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @require_torch_2_7_0
+    def test_lora_sft(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+                "lora_target_linear": True,
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 2e-3,
+                "optimizer": "muon",
+                "weight_decay": 0.01,
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
diff --git a/tests/e2e/multigpu/test_eval.py b/tests/e2e/multigpu/test_eval.py
new file mode 100644
index 0000000000..504659a3a1
--- /dev/null
+++ b/tests/e2e/multigpu/test_eval.py
@@ -0,0 +1,163 @@
+"""
+E2E tests for multigpu eval
+"""
+
+from pathlib import Path
+
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from ..utils import check_tensorboard
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+class TestMultiGPUEval:
+    """
+    Test case for MultiGPU Eval Sample Packing
+    """
+
+    def test_eval_sample_packing(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "load_in_8bit": False,
+                "load_in_4bit": True,
+                "strict": False,
+                "sequence_len": 2048,
+                "adapter": "qlora",
+                "sample_packing": True,
+                "eval_sample_packing": True,
+                "pad_to_sequence_len": True,
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "lora_modules_to_save": ["embed_tokens", "lm_head"],
+                "val_set_size": 0.05,
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "datasets": [
+                    {
+                        "path": "teknium/GPT4-LLM-Cleaned",
+                        "type": "alpaca",
+                        "split": "train[:5%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "loss_watchdog_threshold": 5.0,
+                "loss_watchdog_patience": 3,
+                "bf16": "auto",
+                "warmup_steps": 1,
+                "evals_per_epoch": 2,
+                "eval_max_new_tokens": 128,
+                "saves_per_epoch": 1,
+                "logging_steps": 1,
+                "weight_decay": 0.0,
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "accelerate",
+                "launch",
+                "--num-processes",
+                "2",
+                "--main_process_port",
+                f"{get_torch_dist_unique_port()}",
+                "-m",
+                "axolotl.cli.train",
+                str(Path(temp_dir) / "config.yaml"),
+            ]
+        )
+
+        check_tensorboard(temp_dir + "/runs", "eval/loss", 2.5, "Eval Loss is too high")
+
+    def test_eval(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "load_in_8bit": False,
+                "load_in_4bit": True,
+                "strict": False,
+                "sequence_len": 2048,
+                "adapter": "qlora",
+                "sample_packing": True,
+                "eval_sample_packing": False,
+                "pad_to_sequence_len": True,
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "lora_modules_to_save": ["embed_tokens", "lm_head"],
+                "val_set_size": 0.01,
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "datasets": [
+                    {
+                        "path": "teknium/GPT4-LLM-Cleaned",
+                        "type": "alpaca",
+                        "split": "train[:5%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "loss_watchdog_threshold": 5.0,
+                "loss_watchdog_patience": 3,
+                "bf16": "auto",
+                "warmup_steps": 1,
+                "evals_per_epoch": 2,
+                "eval_max_new_tokens": 128,
+                "saves_per_epoch": 1,
+                "logging_steps": 1,
+                "weight_decay": 0.0,
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "accelerate",
+                "launch",
+                "--num-processes",
+                "2",
+                "--main_process_port",
+                f"{get_torch_dist_unique_port()}",
+                "-m",
+                "axolotl.cli.train",
+                str(Path(temp_dir) / "config.yaml"),
+            ]
+        )
+
+        check_tensorboard(temp_dir + "/runs", "eval/loss", 2.9, "Eval Loss is too high")
diff --git a/tests/e2e/multigpu/test_fp8_fsdp2.py b/tests/e2e/multigpu/test_fp8_fsdp2.py
new file mode 100644
index 0000000000..8d7c01ce89
--- /dev/null
+++ b/tests/e2e/multigpu/test_fp8_fsdp2.py
@@ -0,0 +1,118 @@
+"""Test module for FP8 mixed precision with FSDP2 multi-GPU functionality."""
+
+import os
+from pathlib import Path
+
+import torch
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from tbparse import SummaryReader
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import most_recent_subdir, require_torch_2_7_0, supports_fp8
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+def verify_fp8_training_success(temp_dir):
+    """Verify that FP8 training completed successfully by checking artifacts and loss."""
+    output_path = Path(temp_dir)
+
+    model_files = list(output_path.glob("*.bin")) + list(
+        output_path.glob("*.safetensors")
+    )
+    assert len(model_files) > 0, "No model files found - training may have failed"
+
+    checkpoint_files = list(output_path.glob("checkpoint-*"))
+    assert len(checkpoint_files) > 0, (
+        "No checkpoint files found - training may have failed"
+    )
+
+    tb_log_path = most_recent_subdir(temp_dir + "/runs")
+    if tb_log_path:
+        event_files = sorted(os.listdir(tb_log_path))
+        if event_files:
+            event_file = os.path.join(tb_log_path, event_files[0])
+            reader = SummaryReader(event_file)
+            df = reader.scalars
+            train_loss_df = df[df.tag == "train/train_loss"]
+            if len(train_loss_df) > 0:
+                final_loss = train_loss_df.value.values[-1]
+                assert not torch.isnan(torch.tensor(final_loss)), (
+                    f"Training loss is NaN: {final_loss}"
+                )
+
+
+class TestFP8FSDP2:
+    """Test class for FP8 mixed precision with FSDP2 functionality."""
+
+    @require_torch_2_7_0
+    @supports_fp8
+    def test_fp8_fsdp2_smoke(self, temp_dir):
+        """Smoke test for 2-GPU FP8 + torch.compile + FSDP2 training"""
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "trust_remote_code": True,
+                "sequence_len": 512,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,  # Very short smoke test
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",  # Use standard optimizer for stability
+                "lr_scheduler": "cosine",
+                "sdp_attention": True,
+                "pad_to_seq_len": True,
+                "sample_packing": True,
+                # FP8 configuration
+                "fp8": True,
+                "fp8_enable_fsdp_float8_all_gather": True,
+                "torch_compile": True,
+                # FSDP2 configuration
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_fp8_training_success(temp_dir)
diff --git a/tests/e2e/multigpu/test_fsdp1.py b/tests/e2e/multigpu/test_fsdp1.py
new file mode 100644
index 0000000000..c6a8a47e99
--- /dev/null
+++ b/tests/e2e/multigpu/test_fsdp1.py
@@ -0,0 +1,333 @@
+"""Test module for FSDP1 multi-GPU functionality."""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_tensorboard_loss_decreased
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+def verify_training_success(temp_dir):
+    """Verify that training completed successfully — artifacts, no-NaN, loss
+    stayed in qwen2-pretraining scale (tiny-qwen2-129m final pretrain CE ~3.92).
+    """
+    output_path = Path(temp_dir)
+
+    model_files = list(output_path.glob("*.bin")) + list(
+        output_path.glob("*.safetensors")
+    )
+    assert len(model_files) > 0, "No model files found - training may have failed"
+
+    checkpoint_files = list(output_path.glob("checkpoint-*"))
+    assert len(checkpoint_files) > 0, (
+        "No checkpoint files found - training may have failed"
+    )
+
+    check_tensorboard_loss_decreased(
+        temp_dir + "/runs",
+        initial_window=10,
+        final_window=10,
+        max_initial=5.0,
+        max_final=4.7,
+    )
+
+
+class TestFSDP1:
+    """Test class for FSDP1 functionality."""
+
+    @pytest.mark.parametrize(
+        "fsdp_cpu_ram_efficient_loading",
+        [True, False],
+    )
+    def test_fft_sft(self, temp_dir, fsdp_cpu_ram_efficient_loading):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": "1",
+                "fsdp_config": {
+                    "fsdp_offload_params": False,
+                    "fsdp_cpu_ram_efficient_loading": fsdp_cpu_ram_efficient_loading,
+                    "fsdp_transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "fsdp_state_dict_type": "FULL_STATE_DICT",
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "fsdp_sharding_strategy": "FULL_SHARD",
+                    "fsdp_sync_module_states": True,
+                    "fsdp_use_orig_params": False,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @pytest.mark.parametrize(
+        "adapter_config",
+        [
+            {
+                "adapter": "lora",
+                "load_in_4bit": False,
+            },
+            {
+                "adapter": "qlora",
+                "load_in_4bit": True,
+            },
+        ],
+    )
+    def test_lora_sft(self, temp_dir, adapter_config):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "adapter": adapter_config["adapter"],
+                "load_in_4bit": adapter_config["load_in_4bit"],
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+                "lora_target_linear": True,
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-3,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": "1",
+                "fsdp_config": {
+                    "fsdp_offload_params": False,
+                    "fsdp_cpu_ram_efficient_loading": True,
+                    "fsdp_transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "fsdp_state_dict_type": "FULL_STATE_DICT",
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "fsdp_sharding_strategy": "FULL_SHARD",
+                    "fsdp_sync_module_states": True,
+                    "fsdp_use_orig_params": False,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @pytest.mark.skip(reason="slow test, deprecate fsdp1 asap")
+    def test_dpo_fft(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "rl": "dpo",
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "Intel/orca_dpo_pairs",
+                        "split": "train",
+                        "type": "chatml.intel",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 20,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": "1",
+                "fsdp_config": {
+                    "fsdp_offload_params": False,
+                    "fsdp_cpu_ram_efficient_loading": True,
+                    "fsdp_transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "fsdp_state_dict_type": "FULL_STATE_DICT",
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "fsdp_sharding_strategy": "FULL_SHARD",
+                    "fsdp_sync_module_states": True,
+                    "fsdp_use_orig_params": False,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @pytest.mark.skip("broken in transformers v5")
+    @pytest.mark.parametrize(
+        "adapter_config",
+        [
+            {
+                "adapter": "lora",
+                "load_in_4bit": False,
+            },
+            {
+                "adapter": "qlora",
+                "load_in_4bit": True,
+            },
+        ],
+    )
+    def test_dpo_lora(self, temp_dir, adapter_config):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "load_in_4bit": adapter_config["load_in_4bit"],
+                "rl": "dpo",
+                "chat_template": "chatml",
+                "sequence_len": 2048,
+                "adapter": adapter_config["adapter"],
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "Intel/orca_dpo_pairs",
+                        "split": "train",
+                        "type": "chatml.intel",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 20,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-3,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": "1",
+                "fsdp_config": {
+                    "fsdp_offload_params": False,
+                    "fsdp_cpu_ram_efficient_loading": True,
+                    "fsdp_transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "fsdp_state_dict_type": "FULL_STATE_DICT",
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "fsdp_sharding_strategy": "FULL_SHARD",
+                    "fsdp_sync_module_states": True,
+                    "fsdp_use_orig_params": False,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": "auto",
+                "tf32": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
diff --git a/tests/e2e/multigpu/test_fsdp2.py b/tests/e2e/multigpu/test_fsdp2.py
new file mode 100644
index 0000000000..de625ea114
--- /dev/null
+++ b/tests/e2e/multigpu/test_fsdp2.py
@@ -0,0 +1,505 @@
+"""Test module for FSDP2 multi-GPU functionality."""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_tensorboard_loss_decreased, require_torch_2_7_0
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+def verify_training_success(temp_dir):
+    """Verify that training completed successfully — artifacts, no-NaN, loss
+    stayed in qwen2-pretraining scale (tiny-qwen2-129m final pretrain CE ~3.92).
+    """
+    output_path = Path(temp_dir)
+
+    model_files = list(output_path.glob("*.bin")) + list(
+        output_path.glob("*.safetensors")
+    )
+    assert len(model_files) > 0, "No model files found - training may have failed"
+
+    checkpoint_files = list(output_path.glob("checkpoint-*"))
+    assert len(checkpoint_files) > 0, (
+        "No checkpoint files found - training may have failed"
+    )
+
+    check_tensorboard_loss_decreased(
+        temp_dir + "/runs",
+        initial_window=10,
+        final_window=10,
+        max_initial=5.0,
+        max_final=4.7,
+    )
+
+
+class TestFSDP2:
+    """Test class for FSDP2 functionality."""
+
+    @require_torch_2_7_0
+    @pytest.mark.parametrize(
+        "fsdp_cpu_ram_efficient_loading",
+        [True, False],
+    )
+    def test_fft_sft(self, temp_dir, fsdp_cpu_ram_efficient_loading):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:3%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": fsdp_cpu_ram_efficient_loading,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @require_torch_2_7_0
+    @pytest.mark.parametrize("peft_use_dora", [True, False])
+    def test_lora_sft(self, temp_dir, peft_use_dora):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:3%]",
+                    },
+                ],
+                "peft_use_dora": peft_use_dora,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+                "lora_target_linear": True,
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-3,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+                # explicitly disable LORA kernels, as they may be auto-enabled
+                "lora_mlp_kernel": False,
+                "lora_qkv_kernel": False,
+                "lora_o_kernel": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @require_torch_2_7_0
+    def test_lora_sft_kernels(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:3%]",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_target_linear": True,
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-3,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+                "lora_mlp_kernel": True,
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @require_torch_2_7_0
+    def test_qlora_sft(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:3%]",
+                    },
+                ],
+                "load_in_4bit": True,
+                "adapter": "qlora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+                "lora_target_linear": True,
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-3,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @require_torch_2_7_0
+    def test_qlora_sft_kernels(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:3%]",
+                    },
+                ],
+                "load_in_4bit": True,
+                "adapter": "qlora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_target_linear": True,
+                "num_epochs": 1,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-3,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "bf16": True,
+                "lora_mlp_kernel": True,
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @pytest.mark.skip(reason="slow test w cu129 + torch 2.9.1 + py3.12")
+    @require_torch_2_7_0
+    def test_dpo_fft(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "rl": "dpo",
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "Intel/orca_dpo_pairs",
+                        "split": "train[:100]",
+                        "type": "chatml.intel",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 20,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
+
+    @pytest.mark.skip(reason="slow test w cu129 + torch 2.9.1 + py3.12")
+    @require_torch_2_7_0
+    def test_dpo_lora(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "rl": "dpo",
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "Intel/orca_dpo_pairs",
+                        "split": "train[:100]",
+                        "type": "chatml.intel",
+                    },
+                ],
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "num_epochs": 1,
+                "max_steps": 20,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-3,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "Qwen2DecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        verify_training_success(temp_dir)
diff --git a/tests/e2e/multigpu/test_fsdp2_fp32_norms.py b/tests/e2e/multigpu/test_fsdp2_fp32_norms.py
new file mode 100644
index 0000000000..6dd2bf9c91
--- /dev/null
+++ b/tests/e2e/multigpu/test_fsdp2_fp32_norms.py
@@ -0,0 +1,144 @@
+"""Multi-GPU e2e test for ``fp32_norms`` under FSDP2.
+
+Two-GPU subprocess run with ``fp32_norms: true`` + ``fsdp_version: 2`` + bf16
+training. The test plugin
+``tests.e2e.multigpu._fp32_norms_dtype_capture.DtypeCapturePlugin`` dumps
+post-step-1 param dtypes as JSON; the outer test asserts norms stayed fp32 and
+at least one non-norm param dropped to bf16 (proving the two policies are
+genuinely independent, not a globally-cast model).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import require_torch_2_7_0
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+def _base_fp32_norms_config(
+    temp_dir: str, *, cpu_ram_efficient_loading: bool = False, **overrides
+) -> DictDefault:
+    """Base config for fp32_norms + FSDP2 multi-GPU."""
+    cfg = {
+        "base_model": "axolotl-ai-co/tiny-qwen3-129m",
+        "sequence_len": 256,
+        "val_set_size": 0.0,
+        "datasets": [
+            {
+                "path": "tatsu-lab/alpaca",
+                "type": "alpaca",
+                "split": "train[:1%]",
+            },
+        ],
+        # Full FT (no adapter) — fp32_norms is about base-model norm precision,
+        # which adapters wouldn't exercise.
+        "num_epochs": 1,
+        "max_steps": 2,
+        "micro_batch_size": 1,
+        "gradient_accumulation_steps": 1,
+        "output_dir": temp_dir,
+        "learning_rate": 1e-4,
+        "optimizer": "adamw_torch_fused",
+        "lr_scheduler": "cosine",
+        "flash_attention": True,
+        "bf16": True,
+        "fp32_norms": True,
+        "fsdp_version": 2,
+        "fsdp_config": {
+            "offload_params": False,
+            "cpu_ram_efficient_loading": cpu_ram_efficient_loading,
+            "transformer_layer_cls_to_wrap": "Qwen3DecoderLayer",
+            "state_dict_type": "FULL_STATE_DICT",
+            "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+            "reshard_after_forward": True,
+        },
+        "plugins": [
+            "tests.e2e.multigpu._fp32_norms_dtype_capture.DtypeCapturePlugin",
+        ],
+        "save_safetensors": True,
+    }
+    cfg.update(overrides)
+    return DictDefault(cfg)
+
+
+def _run_training(temp_dir: str, cfg: DictDefault, dump_path: Path) -> None:
+    """Write yaml + spawn 2-process training; plugin path goes via PYTHONPATH."""
+    Path(temp_dir).mkdir(parents=True, exist_ok=True)
+    with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+        fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+    env = os.environ | {
+        # Make the test-only plugin module importable in the subprocess.
+        "PYTHONPATH": (f"{AXOLOTL_ROOT}{os.pathsep}{os.environ.get('PYTHONPATH', '')}"),
+        "FP32_NORMS_DTYPE_DUMP_PATH": str(dump_path),
+    }
+    execute_subprocess_async(
+        [
+            "axolotl",
+            "train",
+            str(Path(temp_dir) / "config.yaml"),
+            "--num-processes",
+            "2",
+            "--main-process-port",
+            f"{get_torch_dist_unique_port()}",
+        ],
+        env=env,
+    )
+
+
+class TestFSDP2Fp32Norms:
+    """Verifies the fp32_norms FSDP2 path end-to-end across 2 GPUs."""
+
+    @require_torch_2_7_0
+    @pytest.mark.parametrize(
+        "cpu_ram_efficient_loading",
+        [False, True],
+        ids=["materialized-load", "cpu-ram-efficient-load"],
+    )
+    def test_norms_stay_fp32_under_fsdp2_bf16(
+        self, temp_dir, cpu_ram_efficient_loading
+    ):
+        """fp32_norms keeps RMSNorm params in fp32 while the rest stays bf16."""
+        dump_path = Path(temp_dir) / "dtype_capture.json"
+        cfg = _base_fp32_norms_config(
+            temp_dir,
+            cpu_ram_efficient_loading=cpu_ram_efficient_loading,
+        )
+        _run_training(temp_dir, cfg, dump_path)
+
+        # Training completed (no FSDP1-style flat-param dtype crash) AND the
+        # plugin captured dtypes after step 1.
+        assert dump_path.exists(), (
+            f"plugin did not dump dtype capture to {dump_path}; "
+            "training may have failed before step 1"
+        )
+
+        captured = json.loads(dump_path.read_text())
+        norms = captured["norms"]
+        non_norms = captured["non_norms"]
+
+        assert norms, "no norm params captured — matcher likely failed"
+        assert all(d == "float32" for d in norms.values()), (
+            "fp32_norms claim violated: at least one norm param is not fp32. "
+            f"Captured norm dtypes: {norms}"
+        )
+
+        # At least one non-norm param must be bf16. Without this check the
+        # test would pass on a globally-fp32 model that didn't shard anything.
+        non_norm_dtypes = set(non_norms.values())
+        assert "bfloat16" in non_norm_dtypes, (
+            "expected at least one non-norm param in bfloat16 (proves the two "
+            "policies are independent); got non-norm dtypes: "
+            f"{non_norm_dtypes}"
+        )
diff --git a/tests/e2e/multigpu/test_fsdp2_lora_kernels.py b/tests/e2e/multigpu/test_fsdp2_lora_kernels.py
new file mode 100644
index 0000000000..0f2fd421a3
--- /dev/null
+++ b/tests/e2e/multigpu/test_fsdp2_lora_kernels.py
@@ -0,0 +1,120 @@
+"""Test LoRA kernels under FSDP2 multi-GPU training.
+
+Verifies that lora_qkv_kernel, lora_o_kernel, lora_mlp_kernel, and
+lora_embedding_kernel work correctly with FSDP2 sharding, including
+with bias, dropout, and DoRA enabled.
+"""
+
+from pathlib import Path
+
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import require_torch_2_7_0
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+def _run_training(temp_dir, cfg):
+    """Write config and launch multi-GPU training."""
+    Path(temp_dir).mkdir(parents=True, exist_ok=True)
+    with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+        fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+    execute_subprocess_async(
+        [
+            "axolotl",
+            "train",
+            str(Path(temp_dir) / "config.yaml"),
+            "--num-processes",
+            "2",
+            "--main-process-port",
+            f"{get_torch_dist_unique_port()}",
+        ]
+    )
+
+
+def _base_lora_fsdp2_config(temp_dir, **overrides):
+    """Base config for LoRA + FSDP2 + kernel tests."""
+    cfg = {
+        "base_model": "axolotl-ai-co/tiny-qwen3-129m",
+        "sequence_len": 512,
+        "val_set_size": 0.0,
+        "datasets": [
+            {
+                "path": "tatsu-lab/alpaca",
+                "type": "alpaca",
+                "split": "train[:1%]",
+            },
+        ],
+        "adapter": "lora",
+        "lora_r": 8,
+        "lora_alpha": 16,
+        "lora_target_linear": True,
+        "num_epochs": 1,
+        "max_steps": 3,
+        "micro_batch_size": 1,
+        "gradient_accumulation_steps": 1,
+        "output_dir": temp_dir,
+        "learning_rate": 1e-4,
+        "optimizer": "adamw_torch_fused",
+        "lr_scheduler": "cosine",
+        "flash_attention": True,
+        "bf16": True,
+        "fsdp_version": 2,
+        "fsdp_config": {
+            "offload_params": False,
+            "cpu_ram_efficient_loading": False,
+            "transformer_layer_cls_to_wrap": "Qwen3DecoderLayer",
+            "state_dict_type": "FULL_STATE_DICT",
+            "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+            "reshard_after_forward": True,
+        },
+        # Enable all LoRA kernels
+        "lora_mlp_kernel": True,
+        "lora_qkv_kernel": True,
+        "lora_o_kernel": True,
+        "lora_embedding_kernel": True,
+        "save_safetensors": True,
+    }
+    cfg.update(overrides)
+    return DictDefault(cfg)
+
+
+class TestFSDP2LoRAKernels:
+    """Test LoRA kernels under FSDP2."""
+
+    @require_torch_2_7_0
+    def test_lora_kernels_basic(self, temp_dir):
+        """Basic LoRA + kernels + FSDP2: no dropout, no bias, no DoRA."""
+        cfg = _base_lora_fsdp2_config(temp_dir)
+        _run_training(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+
+    @require_torch_2_7_0
+    def test_lora_kernels_with_dropout(self, temp_dir):
+        """LoRA kernels + dropout + FSDP2."""
+        cfg = _base_lora_fsdp2_config(temp_dir, lora_dropout=0.1)
+        _run_training(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+
+    @require_torch_2_7_0
+    def test_lora_kernels_with_dora(self, temp_dir):
+        """LoRA kernels + DoRA + FSDP2."""
+        cfg = _base_lora_fsdp2_config(temp_dir, peft_use_dora=True)
+        _run_training(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+
+    @require_torch_2_7_0
+    def test_lora_kernels_with_dora_and_dropout(self, temp_dir):
+        """LoRA kernels + DoRA + dropout + FSDP2."""
+        cfg = _base_lora_fsdp2_config(
+            temp_dir,
+            peft_use_dora=True,
+            lora_dropout=0.05,
+        )
+        _run_training(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
diff --git a/tests/e2e/multigpu/test_gemma3.py b/tests/e2e/multigpu/test_gemma3.py
new file mode 100644
index 0000000000..34f98c037a
--- /dev/null
+++ b/tests/e2e/multigpu/test_gemma3.py
@@ -0,0 +1,98 @@
+"""
+E2E tests for multigpu lora tinyllama
+"""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from huggingface_hub import snapshot_download
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_tensorboard
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_model():
+    # download the model
+    snapshot_download("axolotl-mirrors/gemma-3-4b-pt", repo_type="model")
+
+
+@pytest.mark.skip(reason="FIXME")
+class TestMultiGPUGemma3:
+    """
+    Test case for Gemma3 models using LoRA
+    """
+
+    def test_lora_ddp_packed(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-mirrors/gemma-3-4b-pt",
+                "unfrozen_parameters": ["model.language_model.*", "lm_head"],
+                "sequence_len": 2048,
+                "ddp_find_unused_parameters": True,
+                "sample_packing": True,
+                "eval_sample_packing": False,
+                "pad_to_sequence_len": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.0,
+                "chat_template": "gemma3",
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "split": "train[:10%]",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 4,
+                "gradient_checkpointing": True,
+                "gradient_checkpointing_kwargs": {
+                    "use_reentrant": False,
+                },
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "use_tensorboard": True,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 1.8, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/multigpu/test_llama.py b/tests/e2e/multigpu/test_llama.py
new file mode 100644
index 0000000000..283f078b41
--- /dev/null
+++ b/tests/e2e/multigpu/test_llama.py
@@ -0,0 +1,934 @@
+"""
+E2E tests for multigpu lora tinyllama
+"""
+
+from pathlib import Path
+
+import pytest
+import transformers
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from huggingface_hub import snapshot_download
+from packaging import version
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_tensorboard, require_torch_2_6_0
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_model():
+    # download the model
+    snapshot_download("HuggingFaceTB/SmolLM2-135M")
+
+
+def transformers_version_eq(required_version):
+    return version.parse(transformers.__version__) == version.parse(required_version)
+
+
+class TestMultiGPULlama:
+    """
+    Test case for Llama models using LoRA
+    """
+
+    def test_lora_ddp(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 2048,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 20,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                # "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "use_tensorboard": True,
+                "bf16": True,
+                "save_first_step": False,
+                "seed": 42,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.5, "Train Loss (%s) is too high"
+        )
+
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 2],
+    )
+    def test_lora_ddp_packed(self, temp_dir, gradient_accumulation_steps):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 2048,
+                "sample_packing": True,
+                "eval_sample_packing": False,
+                "pad_to_sequence_len": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:20%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                # "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "use_tensorboard": True,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss (%s) is too high"
+        )
+
+    def test_dpo_lora_ddp(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 2048,
+                "sample_packing": False,
+                "eval_sample_packing": False,
+                "pad_to_sequence_len": True,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "rl": "dpo",
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "fozziethebeat/alpaca_messages_2k_dpo_test",
+                        "type": "chat_template.default",
+                        "field_messages": "conversation",
+                        "field_chosen": "chosen",
+                        "field_rejected": "rejected",
+                        "message_field_role": "role",
+                        "message_field_content": "content",
+                        "roles": {
+                            "system": ["system"],
+                            "user": ["user"],
+                            "assistant": ["assistant"],
+                        },
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "gradient_checkpointing": False,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "warmup_steps": 0,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "use_tensorboard": True,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        loss_threshold = 2.3
+        check_tensorboard(
+            temp_dir + "/runs",
+            "train/train_loss",
+            loss_threshold,
+            "Train Loss (%s) is too high",
+        )
+
+    def test_dpo_qlora_ddp(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 2048,
+                "sample_packing": False,
+                "eval_sample_packing": False,
+                "pad_to_sequence_len": True,
+                "load_in_4bit": True,
+                "adapter": "qlora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "rl": "dpo",
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "fozziethebeat/alpaca_messages_2k_dpo_test",
+                        "type": "chat_template.default",
+                        "field_messages": "conversation",
+                        "field_chosen": "chosen",
+                        "field_rejected": "rejected",
+                        "message_field_role": "role",
+                        "message_field_content": "content",
+                        "roles": {
+                            "system": ["system"],
+                            "user": ["user"],
+                            "assistant": ["assistant"],
+                        },
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "gradient_checkpointing": False,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "warmup_steps": 0,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "use_tensorboard": True,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        loss_threshold = 2.3
+        check_tensorboard(
+            temp_dir + "/runs",
+            "train/train_loss",
+            loss_threshold,
+            "Train Loss (%s) is too high",
+        )
+
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 2],
+    )
+    def test_fsdp(self, temp_dir, gradient_accumulation_steps):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                # "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp": [
+                    "full_shard",
+                    "auto_wrap",
+                ],
+                "fsdp_config": {
+                    "fsdp_offload_params": False,
+                    "fsdp_sync_module_states": True,
+                    "fsdp_use_orig_params": False,
+                    "fsdp_cpu_ram_efficient_loading": False,
+                    "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+                    "fsdp_state_dict_type": "FULL_STATE_DICT",
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                },
+                "use_tensorboard": True,
+                "seed": 42,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss (%s) is too high"
+        )
+
+    @pytest.mark.parametrize(
+        "fsdp_state_dict_type",
+        [
+            "FULL_STATE_DICT",
+            # "SHARDED_STATE_DICT",  # not supported since intermediate checkpoints fail with fsdp1
+        ],
+    )
+    def test_fsdp_packed(self, temp_dir, fsdp_state_dict_type):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "sequence_len": 1024,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,
+                "save_steps": 2,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                # "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp": [
+                    "full_shard",
+                    "auto_wrap",
+                ],
+                "fsdp_config": {
+                    "fsdp_offload_params": False,
+                    "fsdp_sync_module_states": True,
+                    "fsdp_use_orig_params": False,
+                    "fsdp_cpu_ram_efficient_loading": False,
+                    "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+                    "fsdp_state_dict_type": fsdp_state_dict_type,
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                },
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss (%s) is too high"
+        )
+
+    @require_torch_2_6_0
+    @pytest.mark.parametrize(
+        "attention_backend",
+        ["flash", "flex"],
+    )
+    @pytest.mark.parametrize(
+        "fsdp_reshard_after_forward",
+        [True, False],
+    )
+    def test_fsdp2_packed(
+        self, temp_dir, attention_backend, fsdp_reshard_after_forward
+    ):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "sequence_len": 2048,
+                "val_set_size": 0.1,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_8bit",
+                "lr_scheduler": "cosine",
+                "fsdp": [
+                    "auto_wrap",
+                ],
+                "fsdp_config": {
+                    "fsdp_version": 2,
+                    # "fsdp_forward_prefetch": True,  # not yet implemented in accelerate
+                    "fsdp_offload_params": False,
+                    "fsdp_cpu_ram_efficient_loading": False,
+                    "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+                    "fsdp_state_dict_type": "SHARDED_STATE_DICT",
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "fsdp_reshard_after_forward": fsdp_reshard_after_forward,
+                },
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+        if attention_backend == "flash":
+            cfg.attn_implementation = "flash_attention_2"
+        elif attention_backend == "flex":
+            cfg.attn_implementation = "flex_attention"
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.1, "Train Loss (%s) is too high"
+        )
+
+    def test_fsdp_qlora_prequant_packed(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/SmolLM2-135M-bnb-nf4-bf16",
+                "adapter": "qlora",
+                "mean_resizing_embeddings": True,
+                "load_in_4bit": True,
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                # "lora_modules_to_save": [
+                #     "embed_tokens",
+                #     "lm_head",
+                # ],
+                "sample_packing": True,
+                "eval_sample_packing": False,
+                "pad_to_sequence_len": True,
+                "sequence_len": 1024,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                # "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp": [
+                    "full_shard",
+                    "auto_wrap",
+                ],
+                "fsdp_config": {
+                    # pinned to FSDP1: accelerate's fsdp2_load_full_state_dict
+                    # cannot resolve prequantized bnb quant-state keys
+                    # (e.g. `...weight.absmax`)
+                    "fsdp_version": 1,
+                    "fsdp_offload_params": False,
+                    "fsdp_sync_module_states": True,
+                    "fsdp_use_orig_params": False,
+                    "fsdp_cpu_ram_efficient_loading": True,
+                    "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+                    "fsdp_state_dict_type": "FULL_STATE_DICT",
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                },
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss (%s) is too high"
+        )
+
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 2],
+    )
+    @pytest.mark.parametrize(
+        "deepspeed",
+        [
+            "deepspeed_configs/zero3_bf16.json",
+            "deepspeed_configs/zero3_bf16_cpuoffload_all.json",
+            # "deepspeed_configs/zero3_bf16_cpuoffload_params.json",
+        ],
+    )
+    @pytest.mark.parametrize(
+        "qlora",
+        [True, False],
+    )
+    def test_ds_zero3_packed(
+        self, temp_dir, gradient_accumulation_steps, deepspeed, qlora
+    ):
+        if qlora:
+            adapter = {
+                "adapter": "qlora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "load_in_4bit": True,
+            }
+        else:
+            adapter = {}
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "sequence_len": 1024,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "deepspeed": str(AXOLOTL_ROOT / deepspeed),
+                "use_tensorboard": True,
+                "save_first_step": False,
+                **adapter,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.45, "Train Loss (%s) is too high"
+        )
+
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 2],
+    )
+    @pytest.mark.parametrize(
+        "qlora",
+        [True, False],
+    )
+    def test_ds_zero2_packed(self, temp_dir, gradient_accumulation_steps, qlora):
+        if qlora:
+            adapter = {
+                "adapter": "qlora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "load_in_4bit": True,
+            }
+        else:
+            adapter = {}
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "sequence_len": 1024,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "deepspeed": str(AXOLOTL_ROOT / "deepspeed_configs/zero2.json"),
+                "use_tensorboard": True,
+                "seed": 42,
+                "save_first_step": False,
+                **adapter,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss (%s) is too high"
+        )
+
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 2],
+    )
+    @pytest.mark.parametrize(
+        "qlora",
+        [True, False],
+    )
+    def test_ds_zero1_packed(self, temp_dir, gradient_accumulation_steps, qlora):
+        if qlora:
+            adapter = {
+                "adapter": "qlora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "load_in_4bit": True,
+            }
+        else:
+            adapter = {}
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "sequence_len": 1024,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "deepspeed": str(AXOLOTL_ROOT / "deepspeed_configs/zero1.json"),
+                "use_tensorboard": True,
+                "save_first_step": False,
+                **adapter,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.5, "Train Loss (%s) is too high"
+        )
+
+    @pytest.mark.skip(
+        reason="fix untrained tokens brittle with lots of edge cases in latest transformers"
+    )
+    def test_fix_untrained_tokens(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "fix_untrained_tokens": True,
+                "sequence_len": 512,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                    "bos_token": "<|custom_im_start|>",
+                    "eos_token": "<|custom_im_end|>",
+                },
+                "datasets": [
+                    {
+                        "chat_template": "jinja",
+                        "chat_template_jinja": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|custom_im_start|>' + message['role'] + '\n' + message['content'] + '<|custom_im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|custom_im_start|>assistant\n' }}{% endif %}",
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "split": "train[:10%]",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                # "gradient_checkpointing": True,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": True,
+                # "deepspeed": str(AXOLOTL_ROOT / "deepspeed_configs/zero1.json"),
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 4.0, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/multigpu/test_locking.py b/tests/e2e/multigpu/test_locking.py
new file mode 100644
index 0000000000..42502dfa3c
--- /dev/null
+++ b/tests/e2e/multigpu/test_locking.py
@@ -0,0 +1,192 @@
+"""Tests for FileLockLoader class."""
+
+import tempfile
+import threading
+import time
+from pathlib import Path
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+
+from axolotl.utils.data.lock import FileLockLoader
+from axolotl.utils.dict import DictDefault
+
+
+class TestFileLockLoader:
+    """Class with tests for FileLockLoader."""
+
+    @pytest.fixture
+    def temp_dir(self):
+        """Create a temporary directory for testing."""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            yield Path(tmp_dir)
+
+    @pytest.fixture
+    def cfg(self, temp_dir):
+        """Create a test configuration."""
+        return DictDefault({"dataset_prepared_path": str(temp_dir)})
+
+    @pytest.fixture
+    def loader(self, cfg):
+        """Create a FileLockLoader instance for testing."""
+        return FileLockLoader(cfg)
+
+    def test_load_first_process(self, loader):
+        """Test load() when no ready flag exists (first process)."""
+        mock_load_fn = Mock(return_value="test_data")
+
+        result = loader.load(mock_load_fn)
+
+        # Should call the load function
+        mock_load_fn.assert_called_once()
+        assert result == "test_data"
+
+        # Should create the ready flag
+        assert loader.ready_flag_path.exists()
+
+    def test_load_subsequent_process(self, loader):
+        """Test load() when ready flag already exists (subsequent process)."""
+        # Create ready flag first
+        loader.ready_flag_path.touch()
+
+        mock_load_fn = Mock(return_value="loaded_data")
+
+        result = loader.load(mock_load_fn)
+
+        # Should still call load function (to load the prepared data)
+        mock_load_fn.assert_called_once()
+        assert result == "loaded_data"
+
+    def test_load_concurrent_processes(self, cfg):
+        """Test that concurrent processes coordinate correctly."""
+        results = []
+        call_count = 0
+
+        def slow_load_fn():
+            nonlocal call_count
+            call_count += 1
+            time.sleep(0.1)  # Simulate slow loading
+            return f"data_{call_count}"
+
+        def worker():
+            loader = FileLockLoader(cfg)
+            result = loader.load(slow_load_fn)
+            results.append(result)
+
+        # Start multiple threads simultaneously
+        threads = [threading.Thread(target=worker) for _ in range(3)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # Only one thread should have done the initial loading
+        # All should return data, but the load function should be called
+        # once by the first process and once by each subsequent process
+        assert len(results) == 3
+        assert all(result.startswith("data_") for result in results)
+
+    @patch("time.sleep")
+    def test_load_waiting_for_ready_flag(self, mock_sleep, loader):
+        """Test that processes wait for the ready flag to appear."""
+        mock_load_fn = Mock(return_value="waiting_data")
+        mock_ready_flag_path = Mock()
+        exists_call_count = 0
+
+        def mock_exists():
+            nonlocal exists_call_count
+            exists_call_count += 1
+
+            if exists_call_count == 1:
+                # First check: ready flag exists (not first process)
+                return True
+            if exists_call_count <= 3:
+                # While loop checks: flag doesn't exist yet
+                return False
+            return True
+
+        mock_ready_flag_path.exists.side_effect = mock_exists
+
+        # Replace the ready_flag_path with our mock
+        original_path = loader.ready_flag_path
+        loader.ready_flag_path = mock_ready_flag_path
+
+        try:
+            result = loader.load(mock_load_fn)
+        finally:
+            # Restore original path
+            loader.ready_flag_path = original_path
+
+        # Should have slept twice while waiting
+        assert mock_sleep.call_count == 2
+        mock_sleep.assert_called_with(1)
+
+        # Should eventually call load function
+        mock_load_fn.assert_called_once()
+        assert result == "waiting_data"
+
+    def test_complete_workflow_with_cleanup(self, loader):
+        """Test the complete load -> cleanup workflow."""
+        mock_load_fn = Mock(return_value="test_data")
+
+        # First process calls load (this should set up counter)
+        result = loader.load(mock_load_fn)
+        assert result == "test_data"
+        assert loader.ready_flag_path.exists()
+        assert loader.counter_path.exists()
+
+        # Cleanup should remove everything since there's only one process
+        loader.cleanup()
+        assert not loader.ready_flag_path.exists()
+        assert not loader.counter_path.exists()
+
+    def test_multiple_processes_workflow(self, loader):
+        """Test workflow with multiple processes."""
+        # Simulate multiple processes by manually setting up counter
+        loader.ready_flag_path.touch()
+        loader.counter_path.write_text("3")  # 3 processes
+
+        # First process cleanup
+        loader.cleanup()
+        assert loader.ready_flag_path.exists()
+        assert loader.counter_path.read_text().strip() == "2"
+
+        # Second process cleanup
+        loader.cleanup()
+        assert loader.ready_flag_path.exists()
+        assert loader.counter_path.read_text().strip() == "1"
+
+        # Last process cleanup
+        loader.cleanup()
+        assert not loader.ready_flag_path.exists()
+        assert not loader.counter_path.exists()
+
+    def test_load_exception_handling(self, loader):
+        """Test behavior when load_fn raises an exception."""
+
+        def failing_load_fn():
+            raise ValueError("Load failed")
+
+        with pytest.raises(ValueError, match="Load failed"):
+            loader.load(failing_load_fn)
+
+        # Ready flag should not be created on failure
+        assert not loader.ready_flag_path.exists()
+
+    def test_file_lock_called(self, loader):
+        """Test that FileLock is properly used."""
+        mock_load_fn = Mock(return_value="locked_data")
+
+        with patch("axolotl.utils.data.lock.FileLock") as mock_filelock:
+            mock_context = MagicMock()
+            mock_filelock.return_value.__enter__ = Mock(return_value=mock_context)
+            mock_filelock.return_value.__exit__ = Mock(return_value=None)
+
+            loader.load(mock_load_fn)
+
+            # Verify FileLock was called with correct path
+            mock_filelock.assert_called_once_with(str(loader.lock_file_path))
+
+            # Verify context manager was used
+            mock_filelock.return_value.__enter__.assert_called_once()
+            mock_filelock.return_value.__exit__.assert_called_once()
diff --git a/tests/e2e/multigpu/test_ray.py b/tests/e2e/multigpu/test_ray.py
new file mode 100644
index 0000000000..df41b1444a
--- /dev/null
+++ b/tests/e2e/multigpu/test_ray.py
@@ -0,0 +1,209 @@
+"""
+E2E tests for multigpu post-training use Ray Train
+"""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import (
+    check_tensorboard,
+    require_torch_2_7_0,
+)
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
+
+
+class TestMultiGPURay:
+    """
+    Test cases for AnyScale Ray post training
+    """
+
+    @require_torch_2_7_0
+    def test_lora_ddp(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "use_tensorboard": True,
+                "use_ray": True,
+                "ray_num_workers": 2,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--use-ray",
+                "--ray-num-workers",
+                "2",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss (%s) is too high"
+        )
+
+    @require_torch_2_7_0
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 2],
+    )
+    def test_ds_zero2_packed(self, temp_dir, gradient_accumulation_steps):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "sequence_len": 1024,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "deepspeed": str(AXOLOTL_ROOT / "deepspeed_configs/zero2.json"),
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--use-ray",
+                "--ray-num-workers",
+                "2",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss (%s) is too high"
+        )
+
+    @require_torch_2_7_0
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 2],
+    )
+    def test_sft_fsdp2_packed(self, temp_dir, gradient_accumulation_steps):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "sequence_len": 1024,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "output_dir": temp_dir,
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "fsdp_version": 2,
+                "fsdp_config": {
+                    "offload_params": False,
+                    "cpu_ram_efficient_loading": False,
+                    "transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+                    "state_dict_type": "FULL_STATE_DICT",
+                    "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "reshard_after_forward": True,
+                },
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--use-ray",
+                "--ray-num-workers",
+                "2",
+            ]
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.3, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/multigpu/test_tiled_mlp_fsdp2.py b/tests/e2e/multigpu/test_tiled_mlp_fsdp2.py
new file mode 100644
index 0000000000..decc5bf8fd
--- /dev/null
+++ b/tests/e2e/multigpu/test_tiled_mlp_fsdp2.py
@@ -0,0 +1,306 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+"""
+FSDP2 + TiledMLP multi-rank correctness tests.
+
+Parity guard for the tiled MLP under FSDP2: tiled forward+backward
+produces gradients within bf16 tolerance of the un-tiled FSDP2
+reference. The companion fix in
+``axolotl.monkeypatch.tiled_mlp.base._defer_fsdp2_reshard`` is a
+defensive measure that wraps the tile loop in
+``FSDPModule.set_reshard_after_backward(False)`` — under the most
+common setups (FSDP2 wraps the decoder layer; the post-backward
+RegisterPostBackwardFunction fires only when the outer backward
+reaches the layer's input, not mid-tile) the reshard does not fire
+inside the tile loop, but the helper protects against setups where
+the tile loop would otherwise race with FSDP2's per-module reshard.
+
+Run with::
+
+    torchrun --nproc-per-node=2 -m pytest tests/e2e/multigpu/test_tiled_mlp_fsdp2.py
+
+On a 1-GPU executor the tests skip with a clear reason.
+"""
+
+import copy
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+_TORCHRUN_LOCAL_RANK = os.environ.get("LOCAL_RANK")
+_TORCHRUN_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1"))
+
+pytestmark = [
+    pytest.mark.skipif(
+        not torch.cuda.is_available(),
+        reason="CUDA not available",
+    ),
+    pytest.mark.skipif(
+        torch.cuda.device_count() < 2,
+        reason="Need >=2 GPUs for FSDP2 multi-rank tests",
+    ),
+    pytest.mark.skipif(
+        _TORCHRUN_LOCAL_RANK is None or _TORCHRUN_WORLD_SIZE < 2,
+        reason=(
+            "Multi-rank tests must be launched via "
+            "`torchrun --nproc-per-node=2 -m pytest <file>`"
+        ),
+    ),
+]
+
+
+# ──────────────────────────── Process group ──────────────────────────────
+
+
+@pytest.fixture(scope="module")
+def dist_pg():
+    """Initialize the default process group exactly once per worker."""
+    if not dist.is_initialized():
+        rank = int(os.environ["RANK"])
+        torch.cuda.set_device(rank % torch.cuda.device_count())
+        dist.init_process_group(backend="nccl")
+    yield
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()
+
+
+# ──────────────────────────── Helpers ────────────────────────────────────
+
+
+class TinyDenseMLP(nn.Module):
+    def __init__(self, hidden, intermediate, dtype=torch.bfloat16):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden, intermediate, bias=False, dtype=dtype)
+        self.up_proj = nn.Linear(hidden, intermediate, bias=False, dtype=dtype)
+        self.down_proj = nn.Linear(intermediate, hidden, bias=False, dtype=dtype)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+def _full_state(mod):
+    """Materialize an FSDP2 module to full (unsharded) parameters."""
+    from torch.distributed.tensor import DTensor
+
+    out = {}
+    for name, p in mod.named_parameters():
+        if isinstance(p, DTensor):
+            out[name] = p.full_tensor().detach().clone()
+        else:
+            out[name] = p.detach().clone()
+    return out
+
+
+def _full_grads(mod):
+    """Gather param grads to full (unsharded) tensors."""
+    from torch.distributed.tensor import DTensor
+
+    out = {}
+    for name, p in mod.named_parameters():
+        if p.grad is None:
+            continue
+        if isinstance(p.grad, DTensor):
+            out[name] = p.grad.full_tensor().detach().clone()
+        else:
+            out[name] = p.grad.detach().clone()
+    return out
+
+
+def _make_seeded_mlp(hidden, intermediate, dtype, device):
+    torch.manual_seed(42)
+    mlp = TinyDenseMLP(hidden, intermediate, dtype=dtype).to(device)
+    return mlp
+
+
+# ─────────────────────────── Dense regression guard ──────────────────────
+
+
+def _install_tiled_forward(module, shards):
+    """Bind a tiled-MLP forward at the instance level.
+
+    Mirrors what the patcher does in production: the FSDPModule's
+    ``__call__`` triggers FSDP2's pre-forward hooks (which unshard
+    parameters) before ``forward`` runs. By going through the wrapped
+    module's ``__call__`` rather than calling ``TiledMLP.apply`` directly
+    on sharded DTensor params, the tiling and FSDP2's parameter
+    materialization compose correctly.
+    """
+    from types import MethodType
+
+    from axolotl.monkeypatch.tiled_mlp.base import TiledMLP
+
+    original_forward = type(module).forward
+    module._compute_params = []  # type: ignore[attr-defined]
+
+    def tiled_forward(self, x):
+        if not self._compute_params:
+            self._compute_params = [p for p in self.parameters() if p.requires_grad]
+        return TiledMLP.apply(
+            original_forward,
+            self,
+            x,
+            shards,
+            self._compute_params,
+        )
+
+    module.forward = MethodType(tiled_forward, module)
+
+
+def test_fsdp2_tiled_dense_mlp_parity(dist_pg):
+    """FSDP2 + tiled MLP must match FSDP2 + un-tiled MLP within bf16 tolerance.
+
+    Wraps the MLP with ``fully_shard`` so it is itself an ``FSDPModule``
+    — this is the scenario most likely to hit the post-backward race
+    that ``_defer_fsdp2_reshard`` protects against (per-tile inner
+    backwards would otherwise fire the FSDPModule's post-backward hook
+    mid-loop). The test passes whether or not the helper is in place
+    on the current PyTorch (2.11) release; treat it as a parity guard
+    that will catch breakage if FSDP2 ever shortens its reshard timing.
+    """
+    from torch.distributed.fsdp import fully_shard
+
+    device = torch.device(f"cuda:{torch.cuda.current_device()}")
+    hidden, intermediate = 64, 128
+    seq = 64
+    dtype = torch.bfloat16
+
+    # Two identical MLPs — one wrapped with FSDP2 only, one wrapped with
+    # FSDP2 *and* run through TiledMLP. Same initial weights.
+    mlp_ref = _make_seeded_mlp(hidden, intermediate, dtype, device)
+    mlp_tile = copy.deepcopy(mlp_ref)
+    fully_shard(mlp_ref)
+    fully_shard(mlp_tile)
+    _install_tiled_forward(mlp_tile, shards=4)
+
+    torch.manual_seed(7 + dist.get_rank())
+    x = torch.randn(1, seq, hidden, device=device, dtype=dtype)
+    g = torch.randn(1, seq, hidden, device=device, dtype=dtype)
+
+    # Un-tiled reference
+    xr = x.clone().detach().requires_grad_(True)
+    yr = mlp_ref(xr)
+    yr.backward(g)
+    ref_grads = _full_grads(mlp_ref)
+    ref_dx = xr.grad.detach().clone()
+
+    # Tiled run — must not corrupt gradients on FSDP2.
+    xt = x.clone().detach().requires_grad_(True)
+    yt = mlp_tile(xt)
+    yt.backward(g)
+    tile_grads = _full_grads(mlp_tile)
+    tile_dx = xt.grad.detach().clone()
+
+    # Outputs match (this is just forward — should be tight).
+    assert torch.allclose(yr.detach(), yt.detach(), atol=1e-3, rtol=1e-2), (
+        f"FSDP2 forward mismatch max={((yr - yt).abs().max()).item()}"
+    )
+    # dX should match within bf16 tolerance.
+    assert torch.allclose(ref_dx, tile_dx, atol=1e-2, rtol=1e-2), (
+        f"FSDP2 dX mismatch max={((ref_dx - tile_dx).abs().max()).item()}"
+    )
+    # Param grads — the headline check for the reshard fix.
+    for name, gref in ref_grads.items():
+        gtile = tile_grads[name]
+        rel = (
+            (gref.float() - gtile.float()).norm() / (gref.float().norm() + 1e-6)
+        ).item()
+        assert rel < 5e-2, f"FSDP2 + tiled param-grad mismatch {name}: rel_err={rel}"
+
+
+# ─────────────────────── scattermoe-lora regression guard ────────────────
+
+
+def test_fsdp2_tiled_scattermoe_block_parity(dist_pg):
+    """FSDP2 + tiled ScatterMoEGatedMLP block parity guard.
+
+    Same shape of test as the dense case but routes through the
+    ScatterMoE forward. Skips if scattermoe_lora kernels are not
+    available in this env.
+    """
+    try:
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            ScatterMoEGatedMLP,
+        )
+    except ImportError:
+        pytest.skip("scattermoe_lora kernels not available")
+
+    pytest.importorskip("triton")
+
+    from torch.distributed.fsdp import fully_shard
+
+    device = torch.device(f"cuda:{torch.cuda.current_device()}")
+    hidden, intermediate = 64, 128
+    num_experts, top_k = 8, 2
+    seq = 64
+    dtype = torch.bfloat16
+
+    def _make_block():
+        torch.manual_seed(42)
+        block = ScatterMoEGatedMLP()
+        router = SimpleNamespace()
+        router.layer = nn.Linear(hidden, num_experts, bias=False, dtype=dtype).to(
+            device
+        )
+        router.top_k = top_k
+        router.num_experts = num_experts
+        block.router = router
+        in_w = nn.Parameter(
+            torch.randn(
+                num_experts, 2 * intermediate, hidden, dtype=dtype, device=device
+            )
+            * 0.02
+        )
+        out_w = nn.Parameter(
+            torch.randn(num_experts, hidden, intermediate, dtype=dtype, device=device)
+            * 0.02
+        )
+        block.input_linear = nn.Module()
+        block.input_linear.register_parameter("weight", in_w)
+        block.output_linear = nn.Module()
+        block.output_linear.register_parameter("weight", out_w)
+        block.activation = nn.SiLU()
+        return block
+
+    block_ref = _make_block()
+    block_tile = _make_block()
+    fully_shard(block_ref)
+    fully_shard(block_tile)
+    _install_tiled_forward(block_tile, shards=4)
+
+    torch.manual_seed(7 + dist.get_rank())
+    x = torch.randn(1, seq, hidden, device=device, dtype=dtype)
+    g = torch.randn(1, seq, hidden, device=device, dtype=dtype)
+
+    xr = x.clone().detach().requires_grad_(True)
+    yr = block_ref(xr)
+    yr.backward(g)
+    ref_grads = _full_grads(block_ref)
+    ref_dx = xr.grad.detach().clone()
+
+    xt = x.clone().detach().requires_grad_(True)
+    yt = block_tile(xt)
+    yt.backward(g)
+    tile_grads = _full_grads(block_tile)
+    tile_dx = xt.grad.detach().clone()
+
+    def _rel(a, b):
+        return ((a.float() - b.float()).norm() / (b.float().norm() + 1e-6)).item()
+
+    assert _rel(yt.detach(), yr.detach()) < 5e-2, (
+        f"FSDP2 + tiled scattermoe forward rel_err={_rel(yt, yr)}"
+    )
+    assert _rel(tile_dx, ref_dx) < 5e-2, (
+        f"FSDP2 + tiled scattermoe dX rel_err={_rel(tile_dx, ref_dx)}"
+    )
+    for name, gref in ref_grads.items():
+        if name not in tile_grads:
+            continue
+        rel = _rel(tile_grads[name], gref)
+        assert rel < 5e-2, f"FSDP2 + tiled scattermoe param-grad {name} rel_err={rel}"
diff --git a/tests/e2e/multigpu/test_tp.py b/tests/e2e/multigpu/test_tp.py
new file mode 100644
index 0000000000..965dfa8e5b
--- /dev/null
+++ b/tests/e2e/multigpu/test_tp.py
@@ -0,0 +1,68 @@
+"""multigpu e2e test for tensor parallelism."""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async, get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_tensorboard_loss_decreased, require_torch_2_7_0
+
+
+class TestTensorParallel:
+    """Test class for Tensor Parallel functionality."""
+
+    @pytest.mark.skip(
+        reason="TP doesn't work with models with tied weights (embeddings)"
+    )
+    @require_torch_2_7_0
+    def test_fft_sft(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch",
+                "tensor_parallel_size": 2,
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "use_tensorboard": True,
+                "bf16": True,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "train",
+                str(Path(temp_dir) / "config.yaml"),
+                "--num-processes",
+                "2",
+                "--main-process-port",
+                f"{get_torch_dist_unique_port()}",
+            ]
+        )
+
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs", max_initial=5.0, max_final=4.7
+        )
diff --git a/tests/e2e/patched/lora_kernels/__init__.py b/tests/e2e/patched/lora_kernels/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py b/tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py
new file mode 100644
index 0000000000..4117497dac
--- /dev/null
+++ b/tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py
@@ -0,0 +1,1127 @@
+"""Integration tests for LoRA activation and attention kernels."""
+
+import contextlib
+from pathlib import Path
+
+import pytest
+import torch
+import yaml
+from accelerate.state import PartialState
+from peft import LoraConfig, PeftModelForCausalLM, get_peft_config, get_peft_model
+from torch import nn
+from transformers import AutoModelForCausalLM, LlamaForCausalLM
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import LlamaAttention
+from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeAttention
+
+from axolotl.cli.config import load_cfg
+from axolotl.kernels.lora import (
+    apply_lora_gdn_in_proj,
+    apply_lora_linear,
+    apply_lora_mlp_geglu,
+    apply_lora_mlp_swiglu,
+    apply_lora_o,
+    apply_lora_qkv,
+)
+from axolotl.loaders.model import ModelLoader
+from axolotl.loaders.tokenizer import load_tokenizer
+from axolotl.monkeypatch.lora_kernels import (
+    LINEAR_ATTN_IN_PROJS,
+    LINEAR_ATTN_PROJS,
+    apply_lora_kernel_patches,
+    find_linear_attn_in_layer,
+    find_self_attn_in_layer,
+    get_attention_cls_from_config,
+    get_layers,
+    patch_self_attn_lora,
+)
+from axolotl.utils.dict import DictDefault
+
+MODEL_CONFIGS = [
+    {
+        "name": "axolotl-ai-co/tiny-mistral-25m",
+        "expected_activation": apply_lora_mlp_swiglu,
+        "dtype": torch.float16,
+    },
+    {
+        "name": "axolotl-ai-co/tiny-qwen2-129m",
+        "expected_activation": apply_lora_mlp_swiglu,
+        "dtype": torch.float16,
+    },
+    {
+        "name": "HuggingFaceTB/SmolLM2-135M",
+        "expected_activation": apply_lora_mlp_swiglu,
+        "dtype": torch.float32,
+    },
+    {
+        "name": "axolotl-ai-co/tiny-gemma2-137m",
+        "expected_activation": apply_lora_mlp_geglu,
+        "dtype": torch.float16,
+    },
+]
+
+
+@pytest.fixture(autouse=True)
+def init_accelerate():
+    """Initialize Accelerate state before tests."""
+    _ = PartialState()
+
+
+@pytest.fixture
+def small_llama_model():
+    """Create a small LLaMA model for testing."""
+    config = {
+        "vocab_size": 100,
+        "hidden_size": 128,
+        "intermediate_size": 256,
+        "num_hidden_layers": 2,
+        "num_attention_heads": 4,
+    }
+
+    return LlamaForCausalLM(LlamaConfig(**config))
+
+
+@pytest.mark.parametrize(
+    "model_name,attention_cls",
+    [
+        ("HuggingFaceTB/SmolLM2-135M", LlamaAttention),
+        ("Qwen/Qwen3-30B-A3B", Qwen3MoeAttention),
+    ],
+)
+def test_attention_patching_integration(model_name, attention_cls):
+    """Test attention patching in integration context."""
+    cfg = DictDefault({"base_model": model_name})
+
+    # Store the original implementation
+    original_forward = attention_cls.forward
+
+    # Apply patch
+    patch_self_attn_lora(cfg)
+
+    # Get the new forward method
+    patched_forward = attention_cls.forward
+
+    # Check the forward method was replaced
+    assert original_forward is not patched_forward
+    assert patched_forward.__name__ == "axolotl_attn_forward"
+
+    # Check original implementation was stored
+    assert hasattr(attention_cls, "_original_forward")
+
+    # Clean up
+    attention_cls.forward = original_forward
+    delattr(attention_cls, "_original_forward")
+
+
+def test_swiglu_mlp_integration(small_llama_model):
+    """Test SwiGLU activation in LoRA MLP context."""
+    peft_config = get_peft_config(
+        {
+            "peft_type": "LORA",
+            "task_type": "CAUSAL_LM",
+            "r": 8,
+            "lora_alpha": 16,
+            "target_modules": ["gate_proj", "up_proj", "down_proj"],
+            "lora_dropout": 0,
+            "bias": "none",
+        }
+    )
+    model = PeftModelForCausalLM(small_llama_model, peft_config).to("cuda")
+    cfg = DictDefault({"lora_mlp_kernel": True})
+
+    # Apply patches
+    patched_model = apply_lora_kernel_patches(model, cfg)
+
+    # Verify patches
+    layer = patched_model.model.model.layers[0]
+    assert layer.mlp.forward.__func__ is apply_lora_mlp_swiglu
+
+    # Test forward pass
+    batch_size, seq_len = 2, 10
+    hidden_states = torch.randn(
+        batch_size, seq_len, model.config.hidden_size, device=model.device
+    )
+    position_ids = (
+        torch.arange(seq_len, device=model.device).unsqueeze(0).expand(batch_size, -1)
+    )
+    cos, sin = model.model.model.rotary_emb(hidden_states, position_ids)
+
+    inputs = {
+        "hidden_states": hidden_states,
+        "attention_mask": None,
+        "position_embeddings": (cos, sin),
+        "output_attentions": False,
+        "use_cache": False,
+        "past_key_value": None,
+    }
+
+    # Compare outputs
+    with torch.no_grad():
+        original_output = model.model.model.layers[0](**inputs)[0]
+        patched_output = layer(**inputs)[0]
+
+    assert torch.allclose(original_output, patched_output, rtol=1e-4)
+
+
+def test_geglu_model_integration():
+    """Test GeGLU activation with Gemma model."""
+    model = AutoModelForCausalLM.from_pretrained(
+        "axolotl-ai-co/tiny-gemma2-137m",
+        dtype=torch.float16,
+        device_map="cuda:0",
+    )
+    peft_config = get_peft_config(
+        {
+            "peft_type": "LORA",
+            "task_type": "CAUSAL_LM",
+            "r": 8,
+            "lora_alpha": 16,
+            "target_modules": ["gate_proj", "up_proj", "down_proj"],
+            "lora_dropout": 0,
+            "bias": "none",
+        }
+    )
+    model = PeftModelForCausalLM(model, peft_config)
+
+    cfg = DictDefault({"lora_mlp_kernel": True})
+    patched_model = apply_lora_kernel_patches(model, cfg)
+
+    # Verify patches
+    layer = patched_model.model.model.layers[0]
+    assert layer.mlp.forward.__func__ is apply_lora_mlp_geglu
+
+    # Test end-to-end
+    inputs = torch.randint(0, 100, (1, 20), device=model.device, dtype=torch.long)
+    with torch.no_grad():
+        original_output = model(inputs).logits
+        patched_output = patched_model(inputs).logits
+
+    assert torch.allclose(original_output, patched_output, rtol=1e-4)
+
+
+@pytest.mark.parametrize(
+    "model_name,expected_activation",
+    [
+        ("HuggingFaceTB/SmolLM2-135M", apply_lora_mlp_swiglu),
+        ("mhenrichsen/gemma-2b", apply_lora_mlp_geglu),
+    ],
+)
+def test_model_specific_activation(model_name, expected_activation):
+    """Test that each model type gets the correct activation function."""
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    peft_config = get_peft_config(
+        {
+            "peft_type": "LORA",
+            "task_type": "CAUSAL_LM",
+            "r": 8,
+            "lora_alpha": 16,
+            "target_modules": ["gate_proj", "up_proj", "down_proj"],
+            "lora_dropout": 0,
+            "bias": "none",
+        }
+    )
+    model = PeftModelForCausalLM(model, peft_config)
+    cfg = DictDefault({"lora_mlp_kernel": True})
+
+    patched_model = apply_lora_kernel_patches(model, cfg)
+    layer = patched_model.model.model.layers[0]
+    assert layer.mlp.forward.__func__ is expected_activation
+
+
+def test_kernel_patch_conditions():
+    """Test that kernels ARE patched even with dropout and bias (now supported)."""
+    test_configs = [
+        # Dropout — kernels now support this
+        {
+            "peft_type": "LORA",
+            "task_type": "CAUSAL_LM",
+            "r": 8,
+            "lora_alpha": 16,
+            "target_modules": ["gate_proj", "up_proj", "down_proj"],
+            "lora_dropout": 0.1,
+            "bias": "none",
+        },
+        # Bias — kernels now support this
+        {
+            "peft_type": "LORA",
+            "task_type": "CAUSAL_LM",
+            "r": 8,
+            "lora_alpha": 16,
+            "target_modules": ["gate_proj", "up_proj", "down_proj"],
+            "lora_dropout": 0,
+            "bias": "lora_only",
+        },
+    ]
+
+    for config in test_configs:
+        model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M")
+        peft_config = get_peft_config(config)
+        model = PeftModelForCausalLM(model, peft_config)
+        cfg = DictDefault({"lora_mlp_kernel": True})
+
+        patched_model = apply_lora_kernel_patches(model, cfg)
+        layer = patched_model.model.model.layers[0].mlp
+
+        # Verify patches ARE applied (dropout and bias are now supported)
+        assert (
+            layer.forward.__func__ is apply_lora_mlp_swiglu
+            or layer.forward.__func__ is apply_lora_mlp_geglu
+        )
+
+
+def test_kernel_config_options():
+    """Test that kernel configuration options are respected."""
+    # Test different configurations
+    test_configs = [
+        (
+            {"lora_mlp_kernel": True, "lora_qkv_kernel": False, "lora_o_kernel": False},
+            lambda layer: (
+                layer.mlp.forward.__func__ is apply_lora_mlp_swiglu
+                and layer.self_attn.apply_qkv.__func__ is not apply_lora_qkv
+                and layer.self_attn.apply_o.__func__ is not apply_lora_o
+            ),
+        ),
+        (
+            {"lora_mlp_kernel": False, "lora_qkv_kernel": True, "lora_o_kernel": False},
+            lambda layer: (
+                layer.mlp.forward.__func__ is not apply_lora_mlp_swiglu
+                and layer.self_attn.apply_qkv.__func__ is apply_lora_qkv
+                and layer.self_attn.apply_o.__func__ is not apply_lora_o
+            ),
+        ),
+        (
+            {"lora_mlp_kernel": False, "lora_qkv_kernel": False, "lora_o_kernel": True},
+            lambda layer: (
+                layer.mlp.forward.__func__ is not apply_lora_mlp_swiglu
+                and layer.self_attn.apply_qkv.__func__ is not apply_lora_qkv
+                and layer.self_attn.apply_o.__func__ is apply_lora_o
+            ),
+        ),
+    ]
+
+    for config_dict, check_fn in test_configs:
+        # Create fresh model for each test
+        config = {
+            "vocab_size": 100,
+            "hidden_size": 128,
+            "intermediate_size": 256,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 4,
+        }
+        small_llama_model = LlamaForCausalLM(LlamaConfig(**config))
+
+        peft_config = get_peft_config(
+            {
+                "peft_type": "LORA",
+                "task_type": "CAUSAL_LM",
+                "r": 8,
+                "lora_alpha": 16,
+                "target_modules": [
+                    "gate_proj",
+                    "up_proj",
+                    "down_proj",
+                    "q_proj",
+                    "k_proj",
+                    "v_proj",
+                    "o_proj",
+                ],
+                "lora_dropout": 0,
+                "bias": "none",
+            }
+        )
+        model = PeftModelForCausalLM(small_llama_model, peft_config).to("cuda")
+        cfg = DictDefault(config_dict)
+        patched_model = apply_lora_kernel_patches(model, cfg)
+
+        # Verify only requested optimizations were applied
+        for layer in patched_model.model.model.layers:
+            assert check_fn(layer), f"Failed for config: {config_dict}"
+
+        # Clean up
+        del model
+        del small_llama_model
+        del patched_model
+
+
+def get_lora_config():
+    """Get standard LoRA configuration for testing."""
+    return {
+        "peft_type": "LORA",
+        "task_type": "CAUSAL_LM",
+        "r": 8,
+        "lora_alpha": 16,
+        "target_modules": ["gate_proj", "up_proj", "down_proj"],
+        "lora_dropout": 0,
+        "bias": "none",
+    }
+
+
+def get_test_inputs(model, seq_length=20):
+    """Generate test inputs for model evaluation."""
+    return torch.randint(
+        0,
+        model.config.vocab_size,
+        (1, seq_length),
+        device=model.device,
+        dtype=torch.long,
+    )
+
+
+@pytest.mark.parametrize("model_config", MODEL_CONFIGS)
+def test_model_architecture(model_config):
+    """Test LoRA kernel patches across different model architectures."""
+    # Load model with appropriate dtype
+    model = AutoModelForCausalLM.from_pretrained(
+        model_config["name"], torch_dtype=model_config["dtype"], device_map="cuda:0"
+    )
+
+    # Apply LoRA configuration
+    peft_config = get_peft_config(get_lora_config())
+    model = PeftModelForCausalLM(model, peft_config)
+
+    # Apply kernel patches
+    cfg = DictDefault({"lora_mlp_kernel": True})
+    patched_model = apply_lora_kernel_patches(model, cfg)
+
+    # Verify correct activation function
+    layer = patched_model.model.model.layers[0]
+    assert layer.mlp.forward.__func__ is model_config["expected_activation"], (
+        f"Wrong activation for {model_config['name']}"
+    )
+
+    # Test forward pass
+    inputs = get_test_inputs(model)
+    with torch.no_grad():
+        original_output = model(inputs).logits
+        patched_output = patched_model(inputs).logits
+
+    # Check outputs match
+    assert torch.allclose(original_output, patched_output, rtol=1e-4), (
+        f"Outputs don't match for {model_config['name']}"
+    )
+
+
+def test_kernel_training_integration(temp_dir):
+    """Test model loading with kernel patches enabled."""
+    from axolotl.cli.utils import load_model_and_tokenizer
+
+    # Create minimal config
+    cfg = DictDefault(
+        {
+            "base_model": "HuggingFaceTB/SmolLM2-135M",
+            "tokenizer_config": "HuggingFaceTB/SmolLM2-135M",
+            "learning_rate": 0.000001,
+            "datasets": [
+                {
+                    "path": "mhenrichsen/alpaca_2k_test",
+                    "type": "alpaca",
+                }
+            ],
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "adapter": "lora",
+            "lora_r": 8,
+            "lora_alpha": 16,
+            "lora_dropout": 0.0,
+            "lora_target_linear": True,
+            "sequence_len": 1024,
+            "lora_mlp_kernel": True,
+            "lora_qkv_kernel": True,
+            "lora_o_kernel": True,
+        }
+    )
+
+    # Write cfg to yaml file
+    path = Path(temp_dir) / "config.yaml"
+    with open(path, "w", encoding="utf-8") as fout:
+        fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+    # Load config
+    cfg = load_cfg(str(path))
+
+    # Load model
+    model, _, _ = load_model_and_tokenizer(cfg=cfg)
+
+    # Verify correct activation function
+    layer = model.model.model.layers[0]
+    assert layer.mlp.forward.__func__ is apply_lora_mlp_swiglu
+
+
+def test_kernel_training_integration_auto_enable(temp_dir):
+    """Test model loading with auto-enabled kernel patches."""
+    # Create minimal config without explicitly setting kernel options
+    cfg = DictDefault(
+        {
+            "base_model": "HuggingFaceTB/SmolLM2-135M",
+            "tokenizer_config": "HuggingFaceTB/SmolLM2-135M",
+            "learning_rate": 0.000001,
+            "datasets": [
+                {
+                    "path": "mhenrichsen/alpaca_2k_test",
+                    "type": "alpaca",
+                }
+            ],
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "adapter": "lora",
+            "lora_r": 8,
+            "lora_alpha": 16,
+            "lora_dropout": 0.0,
+            "lora_target_linear": True,
+            "sequence_len": 1024,
+        }
+    )
+
+    # Write cfg to yaml file
+    path = Path(temp_dir) / "config.yaml"
+    with open(path, "w", encoding="utf-8") as fout:
+        fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+    # Load config
+    cfg = load_cfg(str(path))
+
+    # Verify kernel options were auto-enabled in the config
+    assert cfg.lora_mlp_kernel is True
+    assert cfg.lora_qkv_kernel is True
+    assert cfg.lora_o_kernel is True
+
+    # Get the attention class before patching to check for side effects
+    attention_cls = get_attention_cls_from_config(cfg)
+
+    # Store original state before patching
+    original_forward_method = attention_cls.forward
+
+    # Load the model (this should trigger the patches)
+    tokenizer = load_tokenizer(cfg)
+    model, _ = ModelLoader(cfg, tokenizer).load()
+
+    # Test side effects of patch_self_attn_lora
+    assert hasattr(attention_cls, "_original_forward")
+    assert attention_cls.forward != original_forward_method
+
+    # Find at least one self-attention module and verify it has the patched methods
+    found_patched_attn = False
+    for layer in model.model.model.layers:
+        if hasattr(layer, "self_attn"):
+            self_attn = layer.self_attn
+            if all(
+                hasattr(self_attn, proj)
+                for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]
+            ):
+                # These methods should be added by apply_lora_kernel_patches
+                assert hasattr(self_attn, "apply_qkv") and callable(self_attn.apply_qkv)
+                assert hasattr(self_attn, "apply_o") and callable(self_attn.apply_o)
+
+                found_patched_attn = True
+                break
+
+    assert found_patched_attn
+
+
+def test_kernel_training_integration_dropout_non_zero(temp_dir):
+    """Test model loading with dropout non-zero DOES patch (now supported)."""
+
+    from axolotl.cli.utils import load_model_and_tokenizer
+
+    # Create minimal config
+    cfg = DictDefault(
+        {
+            "base_model": "HuggingFaceTB/SmolLM2-135M",
+            "tokenizer_config": "HuggingFaceTB/SmolLM2-135M",
+            "learning_rate": 0.000001,
+            "datasets": [
+                {
+                    "path": "mhenrichsen/alpaca_2k_test",
+                    "type": "alpaca",
+                }
+            ],
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "adapter": "lora",
+            "lora_r": 8,
+            "lora_alpha": 16,
+            "lora_dropout": 0.1,
+            "lora_target_linear": True,
+            "lora_qkv_kernel": True,
+            "lora_o_kernel": True,
+            "sequence_len": 1024,
+        }
+    )
+
+    # Write cfg to yaml file
+    path = Path(temp_dir) / "config.yaml"
+    with open(path, "w", encoding="utf-8") as fout:
+        fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+    # Load config
+    cfg = load_cfg(str(path))
+
+    # Load model
+    model, tokenizer, _ = load_model_and_tokenizer(cfg=cfg)
+
+    model_loader = ModelLoader(cfg, tokenizer)
+
+    # Apply patches — should succeed even with dropout > 0
+    model_loader.patch_manager._apply_self_attention_lora_patch()
+    model_loader.patch_manager._apply_lora_kernel_patch(model)
+
+    # Source rewrite must run with dropout > 0
+    layers = get_layers(model)
+    for layer in layers:
+        for self_attn in find_self_attn_in_layer(layer):
+            assert hasattr(type(self_attn), "_original_forward")
+            assert hasattr(self_attn, "apply_qkv")
+            assert hasattr(self_attn, "apply_o")
+
+
+# ============================================================
+# GatedDeltaNet (linear-attention) fused-LoRA routing
+# ============================================================
+
+
+def _gdn_module_with(proj_names, in_features=64, out_features=64):
+    module = nn.Module()
+    for name in proj_names:
+        setattr(module, name, nn.Linear(in_features, out_features, bias=False))
+    return module
+
+
+def _gdn_layer(linear_attn=None, self_attn=False):
+    layer = nn.Module()
+    if linear_attn is not None:
+        layer.linear_attn = linear_attn
+    if self_attn:
+        layer.self_attn = _gdn_module_with(["q_proj", "k_proj", "v_proj", "o_proj"])
+    return layer
+
+
+class TestFindLinearAttnInLayer:
+    """Selects Qwen3.5 GatedDeltaNet layers only; qwen3_next (fused in_proj_qkvz/in_proj_ba) must NOT be patched."""
+
+    def test_qwen3_5_style_is_selected(self):
+        layer = _gdn_layer(_gdn_module_with(LINEAR_ATTN_PROJS))
+        assert list(find_linear_attn_in_layer(layer)) == [layer.linear_attn]
+
+    def test_qwen3_next_style_is_excluded(self):
+        layer = _gdn_layer(_gdn_module_with(["in_proj_qkvz", "in_proj_ba", "out_proj"]))
+        assert list(find_linear_attn_in_layer(layer)) == []
+
+    def test_non_gdn_layer_is_excluded(self):
+        assert list(find_linear_attn_in_layer(_gdn_layer(self_attn=True))) == []
+
+    def test_out_proj_only_is_excluded(self):
+        layer = _gdn_layer(_gdn_module_with(["out_proj"]))
+        assert list(find_linear_attn_in_layer(layer)) == []
+
+    def test_partial_in_projs_is_excluded(self):
+        # out_proj + only some of the four canonical in-projections: the fused node
+        # and patched forward index all four, so a partial mixer falls back to peft.
+        layer = _gdn_layer(_gdn_module_with(["in_proj_qkv", "in_proj_z", "out_proj"]))
+        assert list(find_linear_attn_in_layer(layer)) == []
+
+    def test_missing_linear_attn_attr_is_excluded(self):
+        assert list(find_linear_attn_in_layer(nn.Module())) == []
+
+
+def _gdn_rel_l2(actual, reference):
+    return (actual - reference).norm().item() / (reference.norm().item() + 1e-12)
+
+
+def _wrapped_gdn_proj(in_features=256, out_features=384, use_dora=False):
+    class _Holder(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.in_proj_qkv = nn.Linear(in_features, out_features, bias=False)
+
+    config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules=["in_proj_qkv"],
+        lora_dropout=0.0,
+        use_dora=use_dora,
+    )
+    peft_model = get_peft_model(_Holder().to("cuda").to(torch.bfloat16), config)
+    proj = peft_model.base_model.model.in_proj_qkv
+    proj.train()
+    return proj
+
+
+def _run_gdn_proj(proj, fused, in_features=256, seed=0):
+    torch.manual_seed(seed)
+    inputs = torch.randn(
+        2, 16, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True
+    )
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = apply_lora_linear(proj, inputs) if fused else proj(inputs)
+    out.float().pow(2).mean().backward()
+    grad = inputs.grad.detach().float().clone()
+    proj.zero_grad(set_to_none=True)
+    return out.detach().float(), grad
+
+
+def test_apply_lora_linear_matches_peft_forward_and_grad():
+    """Fused path matches the peft module forward and grad to bf16 float noise."""
+    proj = _wrapped_gdn_proj()
+    out_ref, grad_ref = _run_gdn_proj(proj, fused=False)
+    out_fused, grad_fused = _run_gdn_proj(proj, fused=True)
+
+    assert _gdn_rel_l2(out_fused, out_ref) < 1e-3
+    assert _gdn_rel_l2(grad_fused, grad_ref) < 1e-3
+
+
+def test_apply_lora_linear_matches_peft_with_dora():
+    """The DoRA branch (magnitude scaling) routes correctly through LoRA_O."""
+    proj = _wrapped_gdn_proj(use_dora=True)
+    with torch.no_grad():
+        for name, param in proj.named_parameters():
+            if "lora_B" in name or "magnitude" in name:
+                param.add_(torch.randn_like(param) * 0.01)
+
+    out_ref, grad_ref = _run_gdn_proj(proj, fused=False)
+    out_fused, grad_fused = _run_gdn_proj(proj, fused=True)
+
+    assert _gdn_rel_l2(out_fused, out_ref) < 1e-2
+    assert _gdn_rel_l2(grad_fused, grad_ref) < 1e-2
+
+
+# ------------------------------------------------------------
+# Fused shared-input in-projection kernel (apply_lora_gdn_in_proj)
+# ------------------------------------------------------------
+
+_GDN_IN_SIZES = {
+    "in_proj_qkv": 384,
+    "in_proj_z": 256,
+    "in_proj_b": 16,  # starved: out == num_v_heads
+    "in_proj_a": 16,
+}
+
+
+def _wrapped_gdn_block(targets, in_features=256, use_dora=False):
+    """A peft-wrapped module holding the four GDN input projections."""
+
+    class _Block(nn.Module):
+        def __init__(self):
+            super().__init__()
+            for name, out_features in _GDN_IN_SIZES.items():
+                setattr(self, name, nn.Linear(in_features, out_features, bias=False))
+
+    config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules=list(targets),
+        lora_dropout=0.0,
+        use_dora=use_dora,
+    )
+    block = get_peft_model(
+        _Block().to("cuda").to(torch.bfloat16), config
+    ).base_model.model
+    block.train()
+    # Perturb lora_B / magnitude so adapters (and DoRA scaling) are non-trivial.
+    with torch.no_grad():
+        for name, param in block.named_parameters():
+            if "lora_B" in name or "magnitude" in name:
+                param.add_(torch.randn_like(param) * 0.02)
+    return block
+
+
+def _run_gdn_block(block, fused, targets, in_features=256, seed=0):
+    names = tuple(_GDN_IN_SIZES)
+    torch.manual_seed(seed)
+    inputs = torch.randn(
+        2, 16, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True
+    )
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        if fused:
+            outs = apply_lora_gdn_in_proj(block, inputs, names)
+        else:
+            outs = {name: getattr(block, name)(inputs) for name in names}
+    sum(o.float().pow(2).mean() for o in outs.values()).backward()
+    grad_x = inputs.grad.detach().float().clone()
+    grad_a = {
+        name: getattr(block, name)
+        .lora_A["default"]
+        .weight.grad.detach()
+        .float()
+        .clone()
+        for name in targets
+    }
+    block.zero_grad(set_to_none=True)
+    return {k: v.detach().float() for k, v in outs.items()}, grad_x, grad_a
+
+
+@pytest.mark.parametrize(
+    "targets",
+    [
+        ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a"),  # all adapted
+        ("in_proj_qkv", "in_proj_z"),  # subset: large only
+        ("in_proj_b", "in_proj_a"),  # subset: starved only
+        ("in_proj_qkv",),  # subset: single
+    ],
+)
+def test_apply_lora_gdn_in_proj_matches_peft(targets):
+    """Fused shared-input path matches independent peft forwards + grads, including
+    when only a subset of the projections carry a LoRA adapter (the rest base-only)."""
+    block = _wrapped_gdn_block(targets)
+    out_ref, gx_ref, ga_ref = _run_gdn_block(block, fused=False, targets=targets)
+    out_fused, gx_fused, ga_fused = _run_gdn_block(block, fused=True, targets=targets)
+
+    for name in _GDN_IN_SIZES:
+        assert _gdn_rel_l2(out_fused[name], out_ref[name]) < 5e-3
+    assert _gdn_rel_l2(gx_fused, gx_ref) < 1e-2  # 4-way bf16 grad accumulation
+    for name in targets:
+        assert _gdn_rel_l2(ga_fused[name], ga_ref[name]) < 5e-3
+
+
+def test_apply_lora_gdn_in_proj_matches_peft_with_dora():
+    """DoRA magnitude scaling routes correctly through the fused in-projection."""
+    targets = ("in_proj_qkv", "in_proj_b")
+    block = _wrapped_gdn_block(targets, use_dora=True)
+    out_ref, gx_ref, ga_ref = _run_gdn_block(block, fused=False, targets=targets)
+    out_fused, gx_fused, ga_fused = _run_gdn_block(block, fused=True, targets=targets)
+
+    for name in _GDN_IN_SIZES:
+        assert _gdn_rel_l2(out_fused[name], out_ref[name]) < 2e-2
+    assert _gdn_rel_l2(gx_fused, gx_ref) < 2e-2
+    for name in targets:
+        assert _gdn_rel_l2(ga_fused[name], ga_ref[name]) < 2e-2
+
+
+def _wrapped_gdn_block_nf4(targets, in_features=256, base_dtype=torch.bfloat16):
+    """NF4 (4-bit) GDN in-projection block. ``base_dtype`` (set before quantizing) becomes
+    quant_state.dtype — bf16 is the common QLoRA path; fp32 exercises the fp32 NF4 dequant
+    kernel, which used to fall through to the bf16 kernel and dequant to garbage."""
+    bnb = pytest.importorskip("bitsandbytes")
+
+    class _Block(nn.Module):
+        def __init__(self):
+            super().__init__()
+            for name, out_features in _GDN_IN_SIZES.items():
+                setattr(
+                    self,
+                    name,
+                    bnb.nn.Linear4bit(
+                        in_features,
+                        out_features,
+                        bias=False,
+                        compute_dtype=torch.bfloat16,
+                        quant_type="nf4",
+                    ),
+                )
+
+    block = get_peft_model(
+        _Block().to(base_dtype).to("cuda"),
+        LoraConfig(r=16, lora_alpha=32, target_modules=list(targets), lora_dropout=0.0),
+    ).base_model.model
+    block.train()
+    with torch.no_grad():
+        for name, param in block.named_parameters():
+            if "lora_B" in name:
+                param.add_(torch.randn_like(param) * 0.02)
+    return block
+
+
+@pytest.mark.parametrize("base_dtype", [torch.bfloat16, torch.float32])
+def test_apply_lora_gdn_in_proj_matches_peft_nf4(base_dtype):
+    """QLoRA (NF4 4-bit base, W_quant dequant path) fused in-projection matches peft for
+    both bf16 and fp32 quant_state. The committed parity tests only cover unquantized bf16."""
+    targets = ("in_proj_qkv", "in_proj_b")  # one large + one starved; rest base-only
+    block = _wrapped_gdn_block_nf4(targets, base_dtype=base_dtype)
+    out_ref, gx_ref, ga_ref = _run_gdn_block(block, fused=False, targets=targets)
+    out_fused, gx_fused, ga_fused = _run_gdn_block(block, fused=True, targets=targets)
+
+    for name in _GDN_IN_SIZES:
+        assert _gdn_rel_l2(out_fused[name], out_ref[name]) < 5e-3
+    assert _gdn_rel_l2(gx_fused, gx_ref) < 1e-2
+    for name in targets:
+        assert _gdn_rel_l2(ga_fused[name], ga_ref[name]) < 5e-3
+
+
+@contextlib.contextmanager
+def _fp32_matmul():
+    """Force true fp32 matmuls so fp32 self-consistency checks are order-independent
+    (a prior test in the suite may enable TF32, which degrades fp32 to ~1e-4)."""
+    prev_cuda = torch.backends.cuda.matmul.allow_tf32
+    prev_cudnn = torch.backends.cudnn.allow_tf32
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    try:
+        yield
+    finally:
+        torch.backends.cuda.matmul.allow_tf32 = prev_cuda
+        torch.backends.cudnn.allow_tf32 = prev_cudnn
+
+
+class _FixedDropout(nn.Module):
+    """Deterministic stand-in for nn.Dropout: applies a precomputed (already scaled) mask."""
+
+    def __init__(self, mask):
+        super().__init__()
+        self.mask = mask
+
+    def forward(self, x):
+        return x * self.mask
+
+
+def _wrapped_gdn_block_fp32(targets, in_features=256, use_dora=False):
+    """fp32 GDN in-projection block so fused-vs-reference grads compare to fp noise."""
+
+    class _Block(nn.Module):
+        def __init__(self):
+            super().__init__()
+            for name, out_features in _GDN_IN_SIZES.items():
+                setattr(self, name, nn.Linear(in_features, out_features, bias=False))
+
+    config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules=list(targets),
+        lora_dropout=0.0,
+        use_dora=use_dora,
+    )
+    block = get_peft_model(_Block().to("cuda"), config).base_model.model
+    block.train()
+    with torch.no_grad():
+        for name, param in block.named_parameters():
+            if "lora_B" in name or "magnitude" in name:
+                param.add_(torch.randn_like(param) * 0.1)
+    return block
+
+
+def _gdn_in_proj_grads(block, names, targets, X, reference, mask=None):
+    """Run forward+backward and return (grad_X, grad_A, grad_B). The reference path
+    rebuilds the fused math with pure torch ops using the SAME shared dropout mask."""
+    X = X.detach().clone().requires_grad_(True)
+    if reference:
+        X_drop = X * mask
+        loss = 0.0
+        for name in names:
+            proj = getattr(block, name)
+            out = X @ proj.base_layer.weight.t()
+            if name in targets:
+                A = proj.lora_A["default"].weight
+                B = proj.lora_B["default"].weight
+                out = out + proj.scaling["default"] * (X_drop @ A.t()) @ B.t()
+            loss = loss + out.float().pow(2).mean()
+    else:
+        outs = apply_lora_gdn_in_proj(block, X, names)
+        loss = sum(o.float().pow(2).mean() for o in outs.values())
+    loss.backward()
+
+    grad_x = X.grad.detach().clone()
+    grad_a = {
+        n: getattr(block, n).lora_A["default"].weight.grad.detach().clone()
+        for n in targets
+    }
+    grad_b = {
+        n: getattr(block, n).lora_B["default"].weight.grad.detach().clone()
+        for n in targets
+    }
+    block.zero_grad(set_to_none=True)
+    return grad_x, grad_a, grad_b
+
+
+def test_apply_lora_gdn_in_proj_backward_under_dropout():
+    """Gradient correctness on the ``lora_dropout > 0`` backward path (saved ``X_drop`` +
+    ``grad_X_drop`` accumulation), which the dropout=0 parity tests never exercise.
+    Compared against a pure-torch reference using the SAME shared dropout mask."""
+    names = tuple(_GDN_IN_SIZES)
+    targets = names  # all four adapted
+    block = _wrapped_gdn_block_fp32(targets)
+
+    torch.manual_seed(0)
+    X = torch.randn(2, 16, 256, device="cuda")
+    mask = (torch.rand_like(X) > 0.5).to(X.dtype) / 0.5  # inverted-dropout scale, p=0.5
+    for name in names:
+        getattr(block, name).lora_dropout["default"] = _FixedDropout(mask)
+
+    with _fp32_matmul():
+        gx_ref, ga_ref, gb_ref = _gdn_in_proj_grads(
+            block, names, targets, X, reference=True, mask=mask
+        )
+        gx_fused, ga_fused, gb_fused = _gdn_in_proj_grads(
+            block, names, targets, X, reference=False
+        )
+
+    assert _gdn_rel_l2(gx_fused, gx_ref) < 1e-4
+    for name in targets:
+        assert _gdn_rel_l2(ga_fused[name], ga_ref[name]) < 1e-4
+        assert _gdn_rel_l2(gb_fused[name], gb_ref[name]) < 1e-4
+
+
+def test_apply_lora_gdn_in_proj_dora_backward_under_dropout():
+    """Self-consistency of the fused GDN DoRA + ``lora_dropout>0`` path: matches the kernel's
+    intended math (base undropped, LoRA dropped, magnitude-scaled), including the ``d_mag``
+    gradient."""
+    targets = ("in_proj_qkv", "in_proj_b")
+    names = tuple(_GDN_IN_SIZES)
+    block = _wrapped_gdn_block_fp32(targets, use_dora=True)
+
+    torch.manual_seed(0)
+    x = torch.randn(2, 16, 256, device="cuda")
+    mask = (torch.rand_like(x) > 0.5).to(x.dtype) / 0.5  # inverted-dropout scale, p=0.5
+    for name in targets:  # only adapted projections carry lora_dropout
+        getattr(block, name).lora_dropout["default"] = _FixedDropout(mask)
+
+    def run(fused):
+        inp = x.detach().clone().requires_grad_(True)
+        if fused:
+            outs = apply_lora_gdn_in_proj(block, inp, names)
+        else:
+            # kernel's intended DoRA math: base on undropped X, LoRA on dropped X,
+            # mag_scale = magnitude / ||W + s*B@A||_col (norm detached, as the kernel does)
+            x_drop = inp * mask
+            outs = {}
+            for n in names:
+                proj = getattr(block, n)
+                if n not in targets:  # base-only projection is a plain nn.Linear
+                    outs[n] = inp @ proj.weight.t()
+                    continue
+                W = proj.base_layer.weight
+                A = proj.lora_A["default"].weight
+                B = proj.lora_B["default"].weight
+                s = proj.scaling["default"]
+                mag = proj.lora_magnitude_vector["default"].weight
+                lora = s * (x_drop @ A.t()) @ B.t()
+                weight_norm = torch.linalg.norm(W + s * (B @ A), dim=1).detach()
+                outs[n] = (mag / weight_norm).unsqueeze(0) * (inp @ W.t() + lora)
+        sum(o.float().pow(2).mean() for o in outs.values()).backward()
+        grads = {"x": inp.grad.detach().float().clone()}
+        for n in targets:
+            proj = getattr(block, n)
+            grads[f"A:{n}"] = (
+                proj.lora_A["default"].weight.grad.detach().float().clone()
+            )
+            grads[f"mag:{n}"] = (
+                proj.lora_magnitude_vector["default"]
+                .weight.grad.detach()
+                .float()
+                .clone()
+            )
+        out = {n: outs[n].detach().float() for n in names}
+        block.zero_grad(set_to_none=True)
+        return out, grads
+
+    with _fp32_matmul():
+        out_ref, g_ref = run(fused=False)
+        out_fused, g_fused = run(fused=True)
+
+    for name in names:
+        assert _gdn_rel_l2(out_fused[name], out_ref[name]) < 1e-4
+    assert _gdn_rel_l2(g_fused["x"], g_ref["x"]) < 1e-4
+    for name in targets:
+        assert _gdn_rel_l2(g_fused[f"A:{name}"], g_ref[f"A:{name}"]) < 1e-4
+        assert _gdn_rel_l2(g_fused[f"mag:{name}"], g_ref[f"mag:{name}"]) < 1e-4
+
+
+# ------------------------------------------------------------
+# End-to-end: apply_lora_kernel_patches wires the methods onto a real GDN layer
+# ------------------------------------------------------------
+
+
+def _build_qwen3_5_gdn_peft_model(target_modules):
+    pytest.importorskip("transformers.models.qwen3_5")
+    from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
+    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM
+
+    torch.manual_seed(0)
+    config = Qwen3_5TextConfig(
+        vocab_size=128,
+        hidden_size=128,
+        intermediate_size=128,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        max_position_embeddings=128,
+        rms_norm_eps=1e-6,
+        linear_key_head_dim=32,
+        linear_value_head_dim=32,
+        linear_num_key_heads=4,
+        linear_num_value_heads=8,
+        linear_conv_kernel_dim=4,
+        layer_types=["linear_attention"],
+    )
+    model = Qwen3_5ForCausalLM(config)
+    peft_config = get_peft_config(
+        {
+            "peft_type": "LORA",
+            "task_type": "CAUSAL_LM",
+            "r": 8,
+            "lora_alpha": 16,
+            "target_modules": list(target_modules),
+            "lora_dropout": 0,
+            "bias": "none",
+        }
+    )
+    return PeftModelForCausalLM(model, peft_config).to("cuda")
+
+
+def test_apply_lora_kernel_patches_wires_gdn_layer():
+    """apply_lora_kernel_patches attaches the fused in-proj + out_proj methods onto a
+    peft-wrapped GatedDeltaNet layer, and they match the unpatched peft projections."""
+    model = _build_qwen3_5_gdn_peft_model(LINEAR_ATTN_PROJS)
+    cfg = DictDefault({"lora_qkv_kernel": True, "lora_o_kernel": True})
+
+    linear_attn = get_layers(model)[0].linear_attn
+    ref = {
+        name: getattr(linear_attn, name) for name in (*LINEAR_ATTN_IN_PROJS, "out_proj")
+    }
+
+    apply_lora_kernel_patches(model, cfg)
+
+    assert hasattr(linear_attn, "apply_in_proj_fused")
+    assert hasattr(linear_attn, "apply_out_proj")
+
+    # in-projections share the hidden-size input; out_proj consumes value_dim.
+    x_in = torch.randn(
+        2, 8, model.config.hidden_size, device="cuda", dtype=torch.float32
+    )
+    x_out = torch.randn(
+        2, 8, ref["out_proj"].in_features, device="cuda", dtype=torch.float32
+    )
+    with torch.no_grad():
+        fused = linear_attn.apply_in_proj_fused(x_in)
+        for name in LINEAR_ATTN_IN_PROJS:
+            assert torch.allclose(fused[name], ref[name](x_in), rtol=1e-4, atol=1e-4)
+        assert torch.allclose(
+            linear_attn.apply_out_proj(x_out),
+            ref["out_proj"](x_out),
+            rtol=1e-4,
+            atol=1e-4,
+        )
+
+
+def test_apply_lora_kernel_patches_skips_gdn_without_adapters():
+    """No in-projection adapter -> no fused method attached (model left untouched)."""
+    model = _build_qwen3_5_gdn_peft_model(["out_proj"])
+    cfg = DictDefault({"lora_qkv_kernel": True, "lora_o_kernel": True})
+
+    apply_lora_kernel_patches(model, cfg)
+
+    linear_attn = get_layers(model)[0].linear_attn
+    assert not hasattr(linear_attn, "apply_in_proj_fused")
+    assert hasattr(linear_attn, "apply_out_proj")  # out_proj still routed
+
+
+@pytest.mark.parametrize(
+    "lora_qkv_kernel, lora_o_kernel, in_proj_patched, out_proj_patched",
+    [
+        (True, False, True, False),  # qkv flag drives in-projections only
+        (False, True, False, True),  # o flag drives out_proj only
+    ],
+)
+def test_apply_lora_kernel_patches_gdn_gates_independently(
+    lora_qkv_kernel, lora_o_kernel, in_proj_patched, out_proj_patched
+):
+    """GDN in_proj follows lora_qkv_kernel and out_proj follows lora_o_kernel,
+    matching how self-attn qkv/o gate independently."""
+    model = _build_qwen3_5_gdn_peft_model(LINEAR_ATTN_PROJS)
+    cfg = DictDefault(
+        {"lora_qkv_kernel": lora_qkv_kernel, "lora_o_kernel": lora_o_kernel}
+    )
+
+    apply_lora_kernel_patches(model, cfg)
+
+    linear_attn = get_layers(model)[0].linear_attn
+    assert hasattr(linear_attn, "apply_in_proj_fused") is in_proj_patched
+    assert hasattr(linear_attn, "apply_out_proj") is out_proj_patched
diff --git a/tests/e2e/patched/test_4d_multipack_llama.py b/tests/e2e/patched/test_4d_multipack_llama.py
index d74d097237..ef28cc4066 100644
--- a/tests/e2e/patched/test_4d_multipack_llama.py
+++ b/tests/e2e/patched/test_4d_multipack_llama.py
@@ -2,21 +2,14 @@
 E2E tests for multipack fft llama using 4d attention masks
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import require_torch_2_1_1, with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from ..utils import check_model_output_exists, with_temp_dir
 
 
 class Test4dMultipackLlama(unittest.TestCase):
@@ -24,13 +17,11 @@ class Test4dMultipackLlama(unittest.TestCase):
     Test case for Llama models using 4d attention with multipack
     """
 
-    @require_torch_2_1_1
     @with_temp_dir
     def test_sdp_lora_packing(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "flash_attention": False,
                 "sdp_attention": True,
                 "sample_packing": True,
@@ -42,7 +33,10 @@ def test_sdp_lora_packing(self, temp_dir):
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
                 "sequence_len": 1024,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
                 "datasets": [
                     {
                         "path": "mhenrichsen/alpaca_2k_test",
@@ -54,27 +48,27 @@ def test_sdp_lora_packing(self, temp_dir):
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 5,
+                "save_steps": 3,
+                "eval_steps": 4,
                 "fp16": True,
+                "save_first_step": False,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
 
     @with_temp_dir
     def test_torch_lora_packing(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "flash_attention": False,
                 "sdp_attention": False,
                 "sample_packing": True,
@@ -86,7 +80,10 @@ def test_torch_lora_packing(self, temp_dir):
                 "lora_alpha": 16,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
                 "datasets": [
                     {
                         "path": "mhenrichsen/alpaca_2k_test",
@@ -98,17 +95,18 @@ def test_torch_lora_packing(self, temp_dir):
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 5,
+                "save_steps": 3,
+                "eval_steps": 4,
                 "fp16": True,
+                "save_first_step": False,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/patched/test_activation_checkpointing.py b/tests/e2e/patched/test_activation_checkpointing.py
new file mode 100644
index 0000000000..e530978282
--- /dev/null
+++ b/tests/e2e/patched/test_activation_checkpointing.py
@@ -0,0 +1,80 @@
+"""
+E2E tests for activation checkpointing
+"""
+
+import pytest
+import transformers
+from torch.utils.checkpoint import checkpoint
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from ..utils import check_model_output_exists
+
+
+@pytest.fixture()
+def fix_checkpoint_after_test():
+    yield
+    transformers.modeling_utils.checkpoint = checkpoint
+
+
+class TestActivationCheckpointing:
+    """
+    E2E tests for activation checkpointing
+    """
+
+    @pytest.mark.parametrize(
+        "gradient_checkpointing",
+        ["offload", "offload_disk"],
+    )
+    def test_activation_checkpointing_offload(
+        self,
+        temp_dir,
+        fix_checkpoint_after_test,
+        gradient_checkpointing,
+    ):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                    "eos_token": "<|im_end|>",
+                },
+                "datasets": [
+                    {
+                        "chat_template": "chatml",
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "split": "train[:10%]",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": True,
+                "gradient_checkpointing": gradient_checkpointing,
+                "save_first_step": False,
+                "dataset_num_proc": 4,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/patched/test_cli_integrations.py b/tests/e2e/patched/test_cli_integrations.py
new file mode 100644
index 0000000000..6eba92689a
--- /dev/null
+++ b/tests/e2e/patched/test_cli_integrations.py
@@ -0,0 +1,47 @@
+"""
+test cases to make sure the plugin args are loaded from the config file
+"""
+
+from pathlib import Path
+
+import yaml
+
+from axolotl.cli.config import load_cfg
+from axolotl.utils.dict import DictDefault
+
+
+class TestPluginArgs:
+    """
+    test class for plugin args loaded from the config file
+    """
+
+    def test_liger_plugin_args(self, temp_dir):
+        test_cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "learning_rate": 0.000001,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "plugins": ["axolotl.integrations.liger.LigerPlugin"],
+                "liger_layer_norm": True,
+                "liger_rope": True,
+                "liger_rms_norm": False,
+                "liger_glu_activation": True,
+                "liger_fused_linear_cross_entropy": True,
+            }
+        )
+
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(test_cfg.to_dict()))
+        cfg = load_cfg(str(Path(temp_dir) / "config.yaml"))
+        assert cfg.liger_layer_norm is True
+        assert cfg.liger_rope is True
+        assert cfg.liger_rms_norm is False
+        assert cfg.liger_glu_activation is True
+        assert cfg.liger_fused_linear_cross_entropy is True
diff --git a/tests/e2e/patched/test_fa_xentropy.py b/tests/e2e/patched/test_fa_xentropy.py
new file mode 100644
index 0000000000..9f46998547
--- /dev/null
+++ b/tests/e2e/patched/test_fa_xentropy.py
@@ -0,0 +1,82 @@
+"""
+E2E tests for lora llama
+"""
+
+import pytest
+from transformers.utils import is_torch_bf16_gpu_available
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from ..utils import check_model_output_exists, check_tensorboard
+
+
+class TestFAXentropyLlama:
+    """
+    Test case for Llama models using LoRA w multipack
+    """
+
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 4],
+    )
+    def test_lora_packing_fa_cross_entropy(self, temp_dir, gradient_accumulation_steps):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "sample_packing": True,
+                "flash_attention": True,
+                "flash_attn_cross_entropy": True,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "field_messages": "conversations",
+                        "message_field_content": "value",
+                        "message_field_role": "from",
+                        "type": "chat_template",
+                        "split": "train[:2%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "save_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+        if is_torch_bf16_gpu_available():
+            cfg.bf16 = True
+        else:
+            cfg.fp16 = True
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 1.5, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/patched/test_falcon_samplepack.py b/tests/e2e/patched/test_falcon_samplepack.py
index ae6a497391..1d688585e9 100644
--- a/tests/e2e/patched/test_falcon_samplepack.py
+++ b/tests/e2e/patched/test_falcon_samplepack.py
@@ -2,21 +2,18 @@
 E2E tests for falcon
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from ..utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    with_temp_dir,
+)
 
 
 class TestFalconPatched(unittest.TestCase):
@@ -26,11 +23,10 @@ class TestFalconPatched(unittest.TestCase):
 
     @with_temp_dir
     def test_qlora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "illuin/tiny-random-FalconForCausalLM",
-                "flash_attention": True,
+                "base_model": "axolotl-ai-co/tiny-falcon-42m",
+                "flash_attention": False,
                 "sample_packing": True,
                 "sequence_len": 2048,
                 "load_in_4bit": True,
@@ -40,7 +36,7 @@ def test_qlora(self, temp_dir):
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
                 "lora_modules_to_save": ["word_embeddings", "lm_head"],
-                "val_set_size": 0.1,
+                "val_set_size": 0.05,
                 "special_tokens": {
                     "bos_token": "<|endoftext|>",
                     "pad_token": "<|endoftext|>",
@@ -52,35 +48,45 @@ def test_qlora(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=6.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_ft(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "illuin/tiny-random-FalconForCausalLM",
-                "flash_attention": True,
+                "base_model": "axolotl-ai-co/tiny-falcon-42m",
+                "flash_attention": False,
                 "sample_packing": True,
                 "sequence_len": 2048,
-                "val_set_size": 0.1,
+                "val_set_size": 0.05,
                 "special_tokens": {
                     "bos_token": "<|endoftext|>",
                     "pad_token": "<|endoftext|>",
@@ -92,21 +98,32 @@ def test_ft(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=6.0,
+            max_final=4.7,
+        )
diff --git a/tests/e2e/patched/test_flattening.py b/tests/e2e/patched/test_flattening.py
new file mode 100644
index 0000000000..2c247d4061
--- /dev/null
+++ b/tests/e2e/patched/test_flattening.py
@@ -0,0 +1,81 @@
+"""
+E2E tests for flattening batches
+"""
+
+import pytest
+from transformers.utils import is_torch_bf16_gpu_available
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from ..utils import check_model_output_exists, check_tensorboard
+
+
+class TestFAFlattening:
+    """
+    Test case for Llama models using LoRA w batch flattening
+    """
+
+    @pytest.mark.parametrize(
+        "gradient_accumulation_steps",
+        [1, 4],
+    )
+    def test_lora_packing_flattening(self, temp_dir, gradient_accumulation_steps):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "batch_flattening": True,
+                "flash_attention": True,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "field_messages": "conversations",
+                        "message_field_content": "value",
+                        "message_field_role": "from",
+                        "type": "chat_template",
+                        "split": "train[:2%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "save_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+        if is_torch_bf16_gpu_available():
+            cfg.bf16 = True
+        else:
+            cfg.fp16 = True
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 1.5, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/patched/test_fsdp2_qlora.py b/tests/e2e/patched/test_fsdp2_qlora.py
new file mode 100644
index 0000000000..de9c929e16
--- /dev/null
+++ b/tests/e2e/patched/test_fsdp2_qlora.py
@@ -0,0 +1,30 @@
+"""Integration tests for FSDP2 Params4bit patches."""
+
+import pytest
+from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
+
+
+class TestFSDPPatchIntegration:
+    """Test FSDP patch integration."""
+
+    @pytest.mark.integration
+    def test_fsdp2_init_patches(self):
+        """Test that all patches can be applied together."""
+        from axolotl.monkeypatch.fsdp2_qlora import (
+            apply_init_sharded_param_patch,
+            apply_init_unsharded_param_patch,
+        )
+
+        original_init_sharded = FSDPParam._init_sharded_param
+        original_init_unsharded = FSDPParam.init_unsharded_param
+
+        # Apply patches
+        apply_init_sharded_param_patch()
+        apply_init_unsharded_param_patch()
+
+        assert FSDPParam._init_sharded_param != original_init_sharded, (
+            "_init_sharded_param was not patched"
+        )
+        assert FSDPParam.init_unsharded_param != original_init_unsharded, (
+            "init_unsharded_param was not patched"
+        )
diff --git a/tests/e2e/patched/test_fused_llama.py b/tests/e2e/patched/test_fused_llama.py
index de1195c368..f0c5df18ad 100644
--- a/tests/e2e/patched/test_fused_llama.py
+++ b/tests/e2e/patched/test_fused_llama.py
@@ -2,25 +2,20 @@
 E2E tests for lora llama
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
+import pytest
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from ..utils import check_model_output_exists, with_temp_dir
 
 
+@pytest.mark.skip("FIXME, mostly underused functionality")
 class TestFusedLlama(unittest.TestCase):
     """
     Test case for Llama models using Fused layers
@@ -28,21 +23,17 @@ class TestFusedLlama(unittest.TestCase):
 
     @with_temp_dir
     def test_fft_packing(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "flash_attention": True,
                 "pad_to_sequence_len": True,
-                "flash_attn_fuse_qkv": True,
                 "flash_attn_fuse_mlp": True,
                 "sample_packing": True,
                 "sequence_len": 1024,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
-                    "unk_token": "<unk>",
-                    "bos_token": "<s>",
-                    "eos_token": "</s>",
+                    "pad_token": "<|endoftext|>",
                 },
                 "datasets": [
                     {
@@ -55,20 +46,21 @@ def test_fft_packing(self, temp_dir):
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
                 "max_steps": 10,
                 "save_steps": 5,
                 "eval_steps": 5,
+                "save_first_step": False,
             }
         )
         if is_torch_bf16_gpu_available():
             cfg.bf16 = True
         else:
             cfg.fp16 = True
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/patched/test_llama_s2_attention.py b/tests/e2e/patched/test_llama_s2_attention.py
deleted file mode 100644
index f1d37eb3ca..0000000000
--- a/tests/e2e/patched/test_llama_s2_attention.py
+++ /dev/null
@@ -1,111 +0,0 @@
-"""
-E2E tests for llama w/ S2 attn
-"""
-
-import logging
-import os
-import unittest
-from pathlib import Path
-
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
-from axolotl.train import train
-from axolotl.utils.config import normalize_config
-from axolotl.utils.dict import DictDefault
-
-from ..utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
-
-
-class TestLlamaShiftedSparseAttention(unittest.TestCase):
-    """
-    Test case for Llama models using S2 Attn
-    """
-
-    @with_temp_dir
-    def test_lora_s2_attn(self, temp_dir):
-        # pylint: disable=duplicate-code
-        cfg = DictDefault(
-            {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
-                "sequence_len": 16384,
-                "sample_packing": False,
-                "flash_attention": True,
-                "s2_attention": True,
-                "load_in_8bit": True,
-                "adapter": "lora",
-                "lora_r": 32,
-                "lora_alpha": 16,
-                "lora_dropout": 0.05,
-                "lora_target_linear": True,
-                "val_set_size": 0.1,
-                "special_tokens": {},
-                "datasets": [
-                    {
-                        "path": "Yukang/LongAlpaca-12k",
-                        "type": "alpaca",
-                    },
-                ],
-                "num_epochs": 2,
-                "micro_batch_size": 1,
-                "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
-                "lr_scheduler": "cosine",
-                "max_steps": 10,
-                "save_steps": 5,
-                "eval_steps": 5,
-                "bf16": "auto",
-            }
-        )
-
-        normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
-
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
-
-    @with_temp_dir
-    def test_fft_s2_attn(self, temp_dir):
-        # pylint: disable=duplicate-code
-        cfg = DictDefault(
-            {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
-                "sequence_len": 16384,
-                "sample_packing": False,
-                "flash_attention": True,
-                "s2_attention": True,
-                "val_set_size": 0.1,
-                "special_tokens": {},
-                "datasets": [
-                    {
-                        "path": "Yukang/LongAlpaca-12k",
-                        "type": "alpaca",
-                    },
-                ],
-                "num_epochs": 2,
-                "micro_batch_size": 1,
-                "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
-                "lr_scheduler": "cosine",
-                "max_steps": 10,
-                "save_steps": 5,
-                "eval_steps": 5,
-                "bf16": "auto",
-            }
-        )
-
-        normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
-
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
diff --git a/tests/e2e/patched/test_lora_llama_multipack.py b/tests/e2e/patched/test_lora_llama_multipack.py
index f251f9b661..5033cabc93 100644
--- a/tests/e2e/patched/test_lora_llama_multipack.py
+++ b/tests/e2e/patched/test_lora_llama_multipack.py
@@ -2,24 +2,16 @@
 E2E tests for lora llama
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-import pytest
-from transformers.utils import is_auto_gptq_available, is_torch_bf16_gpu_available
+from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from ..utils import check_model_output_exists, with_temp_dir
 
 
 class TestLoraLlama(unittest.TestCase):
@@ -29,11 +21,10 @@ class TestLoraLlama(unittest.TestCase):
 
     @with_temp_dir
     def test_lora_packing(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 1024,
                 "sample_packing": True,
                 "flash_attention": True,
@@ -45,9 +36,7 @@ def test_lora_packing(self, temp_dir):
                 "lora_target_linear": True,
                 "val_set_size": 0.2,
                 "special_tokens": {
-                    "unk_token": "<unk>",
-                    "bos_token": "<s>",
-                    "eos_token": "</s>",
+                    "pad_token": "<|endoftext|>",
                 },
                 "datasets": [
                     {
@@ -56,12 +45,15 @@ def test_lora_packing(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
+                "max_steps": 20,
+                "save_steps": 10,
                 "micro_batch_size": 8,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
+                "save_first_step": False,
             }
         )
         if is_torch_bf16_gpu_available():
@@ -69,58 +61,9 @@ def test_lora_packing(self, temp_dir):
         else:
             cfg.fp16 = True
 
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
-
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
-
-    @pytest.mark.skipif(not is_auto_gptq_available(), reason="auto-gptq not available")
-    @with_temp_dir
-    def test_lora_gptq_packed(self, temp_dir):
-        # pylint: disable=duplicate-code
-        cfg = DictDefault(
-            {
-                "base_model": "TheBlokeAI/jackfram_llama-68m-GPTQ",
-                "model_type": "AutoModelForCausalLM",
-                "tokenizer_type": "LlamaTokenizer",
-                "sequence_len": 1024,
-                "sample_packing": True,
-                "flash_attention": True,
-                "load_in_8bit": True,
-                "adapter": "lora",
-                "gptq": True,
-                "gptq_disable_exllama": True,
-                "lora_r": 32,
-                "lora_alpha": 64,
-                "lora_dropout": 0.05,
-                "lora_target_linear": True,
-                "val_set_size": 0.1,
-                "special_tokens": {
-                    "unk_token": "<unk>",
-                    "bos_token": "<s>",
-                    "eos_token": "</s>",
-                },
-                "datasets": [
-                    {
-                        "path": "mhenrichsen/alpaca_2k_test",
-                        "type": "alpaca",
-                    },
-                ],
-                "num_epochs": 2,
-                "save_steps": 0.5,
-                "micro_batch_size": 8,
-                "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
-                "lr_scheduler": "cosine",
-            }
-        )
-        normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/patched/test_mistral_samplepack.py b/tests/e2e/patched/test_mistral_samplepack.py
index a56c530b21..ab59a000cf 100644
--- a/tests/e2e/patched/test_mistral_samplepack.py
+++ b/tests/e2e/patched/test_mistral_samplepack.py
@@ -2,21 +2,19 @@
 E2E tests for lora llama
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from ..utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    require_torch_2_6_0,
+    with_temp_dir,
+)
 
 
 class TestMistral(unittest.TestCase):
@@ -24,12 +22,12 @@ class TestMistral(unittest.TestCase):
     Test case for Llama models using LoRA
     """
 
+    @require_torch_2_6_0
     @with_temp_dir
     def test_lora_packing(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model": "axolotl-ai-co/tiny-mistral-25m",
                 "flash_attention": True,
                 "sample_packing": True,
                 "sequence_len": 1024,
@@ -39,7 +37,7 @@ def test_lora_packing(self, temp_dir):
                 "lora_alpha": 64,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "val_set_size": 0.1,
+                "val_set_size": 0.05,
                 "special_tokens": {
                     "unk_token": "<unk>",
                     "bos_token": "<s>",
@@ -52,35 +50,45 @@ def test_lora_packing(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.5,
+            max_final=4.3,
+        )
 
     @with_temp_dir
     def test_ft_packing(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model": "axolotl-ai-co/tiny-mistral-25m",
                 "flash_attention": True,
                 "sample_packing": True,
                 "sequence_len": 1024,
-                "val_set_size": 0.1,
+                "val_set_size": 0.05,
                 "special_tokens": {
                     "unk_token": "<unk>",
                     "bos_token": "<s>",
@@ -93,21 +101,32 @@ def test_ft_packing(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.5,
+            max_final=4.3,
+        )
diff --git a/tests/e2e/patched/test_mixtral_samplepack.py b/tests/e2e/patched/test_mixtral_samplepack.py
index 8baba03073..3c6eb8d126 100644
--- a/tests/e2e/patched/test_mixtral_samplepack.py
+++ b/tests/e2e/patched/test_mixtral_samplepack.py
@@ -2,21 +2,18 @@
 E2E tests for mixtral
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from ..utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    with_temp_dir,
+)
 
 
 class TestMixtral(unittest.TestCase):
@@ -26,11 +23,9 @@ class TestMixtral(unittest.TestCase):
 
     @with_temp_dir
     def test_qlora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "hf-internal-testing/Mixtral-tiny",
-                "tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
+                "base_model": "axolotl-ai-co/tiny-mixtral-30m",
                 "flash_attention": True,
                 "sample_packing": True,
                 "sequence_len": 2048,
@@ -38,9 +33,9 @@ def test_qlora(self, temp_dir):
                 "adapter": "qlora",
                 "lora_r": 16,
                 "lora_alpha": 32,
-                "lora_dropout": 0.1,
+                "lora_dropout": 0.0,
                 "lora_target_linear": True,
-                "val_set_size": 0.1,
+                "val_set_size": 0.05,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -49,36 +44,46 @@ def test_qlora(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 3e-3,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "logging_steps": 1,
+                "save_steps": 80,
+                "eval_steps": 80,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=10,
+            final_window=10,
+            max_initial=6.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_ft(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "hf-internal-testing/Mixtral-tiny",
-                "tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
+                "base_model": "axolotl-ai-co/tiny-mixtral-30m",
                 "flash_attention": True,
                 "sample_packing": True,
                 "sequence_len": 2048,
-                "val_set_size": 0.1,
+                "val_set_size": 0.05,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -87,25 +92,33 @@ def test_ft(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_bnb_8bit",
+                "learning_rate": 5e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "logging_steps": 1,
+                "save_steps": 80,
+                "eval_steps": 80,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        model, _ = train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (
-            "MixtralFlashAttention2"
-            in model.model.layers[0].self_attn.__class__.__name__
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=6.0,
+            max_final=4.7,
         )
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
diff --git a/tests/e2e/patched/test_model_patches.py b/tests/e2e/patched/test_model_patches.py
index eecd1b3c11..264e861edc 100644
--- a/tests/e2e/patched/test_model_patches.py
+++ b/tests/e2e/patched/test_model_patches.py
@@ -4,10 +4,11 @@
 
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
-from axolotl.utils.config import normalize_config
+import transformers
+
+from axolotl.loaders import ModelLoader, load_tokenizer
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model, load_tokenizer
 
 from ..utils import with_temp_dir
 
@@ -21,12 +22,11 @@ class TestModelPatches(unittest.TestCase):
     def test_mixtral_multipack(self, temp_dir):
         cfg = DictDefault(
             {
-                "base_model": "hf-internal-testing/Mixtral-tiny",
-                "tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
+                "base_model": "axolotl-ai-co/tiny-mixtral-30m",
                 "flash_attention": True,
                 "sample_packing": True,
                 "sequence_len": 2048,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -44,27 +44,23 @@ def test_mixtral_multipack(self, temp_dir):
                 "max_steps": 20,
                 "save_steps": 10,
                 "eval_steps": 10,
+                "save_first_step": False,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
         tokenizer = load_tokenizer(cfg)
-        model, _ = load_model(cfg, tokenizer, inference=cli_args.inference)
-
-        assert (
-            "MixtralFlashAttention2"
-            in model.model.layers[0].self_attn.__class__.__name__
-        )
+        ModelLoader(cfg, tokenizer, inference=False).load()
 
     @with_temp_dir
     def test_mistral_multipack(self, temp_dir):
         cfg = DictDefault(
             {
-                "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model": "axolotl-ai-co/tiny-mistral-25m",
                 "flash_attention": True,
                 "sample_packing": True,
                 "sequence_len": 2048,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -82,14 +78,17 @@ def test_mistral_multipack(self, temp_dir):
                 "max_steps": 20,
                 "save_steps": 10,
                 "eval_steps": 10,
+                "save_first_step": False,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
         tokenizer = load_tokenizer(cfg)
-        model, _ = load_model(cfg, tokenizer, inference=cli_args.inference)
+        ModelLoader(cfg, tokenizer, inference=False).load()
 
+        # In-tree HF models pack natively via position_ids, so we no longer
+        # override `_get_unpad_data` for them; the stock function must remain.
         assert (
-            "axolotl.monkeypatch.mistral_attn_hijack_flash"
-            in model.model.layers[0].self_attn.forward.__module__
+            transformers.modeling_flash_attention_utils._get_unpad_data.__module__
+            == "transformers.modeling_flash_attention_utils"
         )
diff --git a/tests/e2e/patched/test_multipack_packing_equivalence.py b/tests/e2e/patched/test_multipack_packing_equivalence.py
new file mode 100644
index 0000000000..16d20fcd3a
--- /dev/null
+++ b/tests/e2e/patched/test_multipack_packing_equivalence.py
@@ -0,0 +1,60 @@
+"""
+Guards that sample packing does not leak attention across document boundaries.
+
+Axolotl encodes each packed document as a distinct integer id in the 2D attention
+mask, but modern transformers casts that mask to bool before flash attention and
+instead derives per-document ``cu_seqlens`` from the (per-document resetting)
+``position_ids`` (`_is_packed_sequence` / `_prepare_from_posids`). This is why the
+`_get_unpad_data` override was dropped for in-tree models. If that native path ever
+regresses, a packed forward would attend across documents and this test would fail.
+"""
+
+import torch
+from transformers import LlamaConfig, LlamaForCausalLM
+from transformers.testing_utils import require_flash_attn, require_torch_gpu
+
+
+@require_torch_gpu
+@require_flash_attn
+def test_packed_forward_matches_per_document_forward():
+    torch.manual_seed(0)
+    config = LlamaConfig(
+        vocab_size=256,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=8,
+        num_key_value_heads=8,
+        max_position_embeddings=128,
+        attn_implementation="flash_attention_2",
+    )
+    model = LlamaForCausalLM(config).to(device="cuda", dtype=torch.bfloat16).eval()
+
+    # Two documents packed into a single row (batch_size == 1), exactly the shape
+    # the multipack collator emits: segment-id attention mask + resetting position ids.
+    doc1 = [11, 12, 13, 14, 15]
+    doc2 = [21, 22, 23, 24]
+    input_ids = torch.tensor([doc1 + doc2], device="cuda")
+    attention_mask = torch.tensor([[1] * len(doc1) + [2] * len(doc2)], device="cuda")
+    position_ids = torch.tensor(
+        [list(range(len(doc1))) + list(range(len(doc2)))], device="cuda"
+    )
+
+    with torch.no_grad():
+        packed = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+        ).logits
+        ref1 = model(
+            input_ids=torch.tensor([doc1], device="cuda"),
+            position_ids=torch.tensor([list(range(len(doc1)))], device="cuda"),
+        ).logits
+        ref2 = model(
+            input_ids=torch.tensor([doc2], device="cuda"),
+            position_ids=torch.tensor([list(range(len(doc2)))], device="cuda"),
+        ).logits
+
+    reference = torch.cat([ref1, ref2], dim=1)
+    # Each document must see only itself; flash-attn isn't bit-exact across tiling shapes.
+    assert torch.allclose(packed, reference, atol=1e-2)
diff --git a/tests/e2e/patched/test_peft_embeddings.py b/tests/e2e/patched/test_peft_embeddings.py
new file mode 100644
index 0000000000..ae3145de88
--- /dev/null
+++ b/tests/e2e/patched/test_peft_embeddings.py
@@ -0,0 +1,62 @@
+"""
+Test case for handling embeddings when using peft
+"""
+
+import torch
+
+from axolotl.train import setup_model_and_tokenizer
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestLlamaPeftEmbeddings:
+    """
+    test class for handling embeddings when using peft
+    """
+
+    def test_peft_embeddings_upcast(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "load_in_4bit": True,
+                "adapter": "qlora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_target_linear": True,
+                "trust_remote_code": True,
+                "sequence_len": 512,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": False,
+                "bf16": "auto",
+                "embeddings_skip_upcast": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+
+        model, _, _, _ = setup_model_and_tokenizer(cfg)
+
+        # Check if the embeddings are upcast correctly
+        # only embed_tokens is a parameter that may be upcast
+        assert model.base_model.model.model.embed_tokens.weight.dtype == torch.bfloat16
+        assert model.base_model.model.lm_head.weight.dtype == torch.bfloat16
diff --git a/tests/e2e/patched/test_phi_multipack.py b/tests/e2e/patched/test_phi_multipack.py
index 5f30453c18..c3c8ff5693 100644
--- a/tests/e2e/patched/test_phi_multipack.py
+++ b/tests/e2e/patched/test_phi_multipack.py
@@ -2,21 +2,18 @@
 E2E tests for lora llama
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from ..utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    with_temp_dir,
+)
 
 
 class TestPhiMultipack(unittest.TestCase):
@@ -26,10 +23,9 @@ class TestPhiMultipack(unittest.TestCase):
 
     @with_temp_dir
     def test_ft_packed(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "microsoft/phi-1_5",
+                "base_model": "axolotl-ai-co/tiny-phi-64m",
                 "model_type": "PhiForCausalLM",
                 "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 1024,
@@ -38,7 +34,7 @@ def test_ft_packed(self, temp_dir):
                 "pad_to_sequence_len": True,
                 "load_in_8bit": False,
                 "adapter": None,
-                "val_set_size": 0.1,
+                "val_set_size": 0.05,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
                 },
@@ -51,45 +47,55 @@ def test_ft_packed(self, temp_dir):
                 "dataset_shard_num": 10,
                 "dataset_shard_idx": 0,
                 "num_epochs": 1,
-                "micro_batch_size": 1,
+                "micro_batch_size": 2,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_bnb_8bit",
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "eval_steps": 10,
-                "save_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "eval_steps": 50,
+                "save_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
 
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=6.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_qlora_packed(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "microsoft/phi-1_5",
+                "base_model": "axolotl-ai-co/tiny-phi-64m",
                 "model_type": "PhiForCausalLM",
                 "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 1024,
                 "sample_packing": True,
                 "flash_attention": True,
                 "pad_to_sequence_len": True,
-                "load_in_8bit": False,
+                "load_in_4bit": True,
                 "adapter": "qlora",
                 "lora_r": 64,
                 "lora_alpha": 32,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
                 },
@@ -102,22 +108,33 @@ def test_qlora_packed(self, temp_dir):
                 "dataset_shard_num": 10,
                 "dataset_shard_idx": 0,
                 "num_epochs": 1,
-                "micro_batch_size": 1,
+                "micro_batch_size": 2,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "eval_steps": 10,
-                "save_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "eval_steps": 50,
+                "save_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
 
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=6.0,
+            max_final=4.7,
+        )
diff --git a/tests/e2e/patched/test_resume.py b/tests/e2e/patched/test_resume.py
index dfe9e86252..7a744acd17 100644
--- a/tests/e2e/patched/test_resume.py
+++ b/tests/e2e/patched/test_resume.py
@@ -2,54 +2,49 @@
 E2E tests for resuming training
 """
 
-import logging
 import os
 import re
 import subprocess
-import unittest
-from pathlib import Path
 
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
+from axolotl.core.trainers.constants import TOKENS_STATE_FILE
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import most_recent_subdir, with_temp_dir
+from ..utils import check_model_output_exists, most_recent_subdir, require_torch_2_6_0
 
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
 
-
-class TestResumeLlama(unittest.TestCase):
+class TestResumeLlama:
     """
     Test case for resuming training of llama models
     """
 
-    @with_temp_dir
-    def test_resume_qlora_packed(self, temp_dir):
-        # pylint: disable=duplicate-code
+    @require_torch_2_6_0
+    def test_resume_lora_packed(self, temp_dir):
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
                 "sequence_len": 1024,
                 "sample_packing": True,
                 "flash_attention": True,
-                "load_in_4bit": True,
-                "adapter": "qlora",
-                "lora_r": 32,
-                "lora_alpha": 64,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "val_set_size": 0.1,
-                "special_tokens": {},
+                "val_set_size": 0.001,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
                 "datasets": [
                     {
-                        "path": "vicgalle/alpaca-gpt4",
+                        "path": "tatsu-lab/alpaca",
                         "type": "alpaca",
+                        "split": "train[:10%]",
                     },
                 ],
                 "num_epochs": 2,
@@ -57,33 +52,63 @@ def test_resume_qlora_packed(self, temp_dir):
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_8bit",
                 "lr_scheduler": "cosine",
-                "save_steps": 10,
+                "save_steps": 3,
                 "save_total_limit": 5,
-                "max_steps": 40,
+                "max_steps": 15,
+                "use_tensorboard": True,
+                "save_first_step": False,
+                "include_tkps": True,
             }
         )
         if is_torch_bf16_gpu_available():
             cfg.bf16 = True
         else:
             cfg.fp16 = True
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        initial_total_num_tokens = cfg.total_num_tokens
+        assert initial_total_num_tokens is not None, (
+            "total_num_tokens should be calculated during load_datasets"
+        )
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        train(cfg=cfg, dataset_meta=dataset_meta)
+
+        checkpoint_path = f"{temp_dir}/checkpoint-9"
+        tokens_state_path = os.path.join(checkpoint_path, TOKENS_STATE_FILE)
+        assert os.path.isfile(tokens_state_path), (
+            f"{TOKENS_STATE_FILE} should exist in checkpoint at {tokens_state_path}"
+        )
 
         resume_cfg = cfg | DictDefault(
             {
-                "resume_from_checkpoint": f"{temp_dir}/checkpoint-30/",
+                "resume_from_checkpoint": f"{temp_dir}/checkpoint-9/",
             }
         )
         normalize_config(resume_cfg)
-        cli_args = TrainerCliArgs()
 
-        train(cfg=resume_cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        assert resume_cfg.total_num_tokens == initial_total_num_tokens, (
+            f"total_num_tokens should be preserved on resume. "
+            f"Expected {initial_total_num_tokens}, got {resume_cfg.total_num_tokens}"
+        )
+
+        resume_dataset_meta = load_datasets(cfg=resume_cfg)
+
+        assert resume_cfg.total_num_tokens == initial_total_num_tokens, (
+            f"total_num_tokens should not be recalculated when resuming. "
+            f"Expected {initial_total_num_tokens}, got {resume_cfg.total_num_tokens}"
+        )
+
+        train(cfg=resume_cfg, dataset_meta=resume_dataset_meta)
+
+        assert resume_cfg.total_num_tokens == initial_total_num_tokens, (
+            f"total_num_tokens should remain unchanged after resume training. "
+            f"Expected {initial_total_num_tokens}, got {resume_cfg.total_num_tokens}"
+        )
+        check_model_output_exists(temp_dir, cfg)
 
         tb_log_path_1 = most_recent_subdir(temp_dir + "/runs")
         cmd = f"tensorboard --inspect  --logdir {tb_log_path_1}"
@@ -92,4 +117,4 @@ def test_resume_qlora_packed(self, temp_dir):
         )
         pattern = r"first_step\s+(\d+)"
         first_steps = int(re.findall(pattern, res.stdout)[0])
-        assert first_steps == 31
+        assert first_steps == 10
diff --git a/tests/e2e/solo/__init__.py b/tests/e2e/solo/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/e2e/solo/test_batch_flattening.py b/tests/e2e/solo/test_batch_flattening.py
new file mode 100644
index 0000000000..7b6c59119a
--- /dev/null
+++ b/tests/e2e/solo/test_batch_flattening.py
@@ -0,0 +1,612 @@
+"""
+Unit tests for batch flattening correctness in GRPO.
+
+Validates that flattened (padding-free) forward passes produce identical
+results to padded forward passes by calling the ACTUAL AsyncGRPOTrainer methods:
+  1. Deferred scoring: _get_per_token_logps_flattened vs _get_per_token_logps_and_entropies
+  2. Training loss: _get_per_token_logps_and_entropies_flattened vs _get_per_token_logps_and_entropies
+
+Run: CUDA_VISIBLE_DEVICES=1 python test_batch_flattening.py
+"""
+
+import types
+from unittest.mock import MagicMock
+
+import torch
+from transformers import AutoModelForCausalLM
+
+# Import the actual trainer methods we want to test
+from axolotl.core.trainers.grpo.async_trainer import AsyncGRPOTrainer
+
+MODEL_NAME = "axolotl-ai-co/tiny-qwen3-129m"
+
+
+def _fix_patched_attention(model):
+    """Bind apply_qkv on attention modules if LoRA kernel monkeypatch is active.
+
+    The LoRA kernel tests replace ``Qwen3Attention.forward`` at the class level
+    with ``axolotl_attn_forward``, which expects a per-instance ``apply_qkv``
+    method.  Models created *after* that patch but *without* the per-instance
+    setup will crash.  We fix this by binding the original (non-LoRA) apply_qkv.
+    """
+    from axolotl.monkeypatch.lora_kernels import original_apply_o, original_apply_qkv
+
+    for module in model.modules():
+        fwd_name = getattr(type(module).forward, "__name__", "")
+        if "axolotl" in fwd_name and not hasattr(module, "apply_qkv"):
+            module.apply_qkv = types.MethodType(original_apply_qkv, module)
+            module.apply_o = types.MethodType(original_apply_o, module)
+
+
+def setup_model(eval_mode=True):
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, dtype=torch.bfloat16, attn_implementation="flash_attention_2"
+    ).cuda()
+    _fix_patched_attention(model)
+    if eval_mode:
+        model.eval()
+    else:
+        model.train()
+    return model
+
+
+def make_mock_trainer(model):
+    """Create a minimal mock that has the attributes needed by the trainer methods.
+
+    The three methods we test (_get_per_token_logps_flattened,
+    _get_per_token_logps_and_entropies_flattened, _get_per_token_logps_and_entropies)
+    access self.temperature, self.use_liger_kernel, self.is_fsdp_enabled,
+    self.accelerator, and self.model_kwarg_keys.
+    """
+    trainer = MagicMock(spec=[])
+
+    trainer.temperature = 1.0
+    trainer.use_liger_kernel = False
+    trainer.is_fsdp_enabled = False
+    trainer.model_kwarg_keys = set()
+
+    # accelerator.unwrap_model should return the model unchanged
+    accelerator = MagicMock()
+    accelerator.unwrap_model = lambda m, keep_fp32_wrapper=True: m
+    trainer.accelerator = accelerator
+
+    # Bind the real unbound methods to our mock
+    trainer._get_per_token_logps_flattened = types.MethodType(
+        AsyncGRPOTrainer._get_per_token_logps_flattened, trainer
+    )
+    trainer._get_per_token_logps_and_entropies_flattened = types.MethodType(
+        AsyncGRPOTrainer._get_per_token_logps_and_entropies_flattened, trainer
+    )
+    trainer._get_per_token_logps_and_entropies = types.MethodType(
+        AsyncGRPOTrainer._get_per_token_logps_and_entropies, trainer
+    )
+
+    return trainer
+
+
+def make_grpo_batch(B=4, max_compl=64, vocab_range=(100, 5000)):
+    """Create a GRPO-style batch matching the real data layout.
+
+    In real GRPO, input_ids = cat([prompt_ids, completion_ids], dim=1).
+    prompt_ids is padded to max_prompt_len, completion_ids to max_compl.
+    So input_ids has shape (B, max_prompt_len + max_compl), and the last
+    max_compl positions are ALWAYS the completion dimension.
+    """
+    torch.manual_seed(42)
+
+    # Fixed prompt length: avoids prompt padding which causes position-0
+    # divergence between padded and flattened paths (the padded path's shifted
+    # window at position 0 uses a padding-position logit when prompt_len < max_prompt).
+    fixed_prompt = 20
+    prompt_lens = [fixed_prompt] * B
+    compl_lens = [max_compl] * B
+    max_prompt = fixed_prompt
+    logits_to_keep = max_compl
+
+    # Build like real GRPO: prompt_ids (B, max_prompt) + completion_ids (B, max_compl)
+    prompt_ids = torch.zeros(B, max_prompt, dtype=torch.long, device="cuda")
+    completion_ids = torch.randint(*vocab_range, (B, max_compl), device="cuda")
+    prompt_mask_raw = torch.zeros(B, max_prompt, dtype=torch.long, device="cuda")
+
+    for i in range(B):
+        prompt_ids[i, : prompt_lens[i]] = torch.randint(
+            *vocab_range, (prompt_lens[i],), device="cuda"
+        )
+        prompt_mask_raw[i, : prompt_lens[i]] = 1
+
+    # Concatenate like _compute_loss does
+    input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+    completion_mask_raw = torch.ones(B, max_compl, dtype=torch.long, device="cuda")
+    attention_mask = torch.cat([prompt_mask_raw, completion_mask_raw], dim=1)
+    # Full prompt mask (padded to input_ids length)
+    prompt_mask = torch.cat(
+        [
+            prompt_mask_raw,
+            torch.zeros(B, max_compl, dtype=torch.long, device="cuda"),
+        ],
+        dim=1,
+    )
+
+    completion_mask = torch.ones(B, logits_to_keep, dtype=torch.float32, device="cuda")
+
+    total_lens = [p + max_compl for p in prompt_lens]
+
+    return (
+        input_ids,
+        attention_mask,
+        completion_mask,
+        logits_to_keep,
+        prompt_mask,
+        {
+            "prompt_lens": prompt_lens,
+            "compl_lens": compl_lens,
+            "total_lens": total_lens,
+        },
+    )
+
+
+def _compare_logps(
+    logps_pad, logps_flat, max_thresh=1.0, mean_thresh=0.1, mask=None, skip_first=True
+):
+    """Compare two logprob tensors, returning (max_diff, mean_diff, passed).
+
+    Args:
+        mask: optional (B, T) mask. Only compare positions where mask > 0.
+        skip_first: skip position 0 of each sequence's completion logprobs.
+            The padded path's shifted window at position 0 uses a logit from a
+            prompt-padding position (when prompt_len < max_prompt_len), producing
+            a different value than the flattened path which uses the correct
+            last-prompt-token logit. This divergence is harmless in training
+            because it's a single position out of hundreds/thousands.
+    """
+    diff = (logps_pad.float() - logps_flat.float()).abs()
+    if mask is not None:
+        compare_mask = mask.bool().clone()
+    else:
+        compare_mask = ((logps_pad != 0) | (logps_flat != 0)).clone()
+
+    if skip_first:
+        # Zero out position 0 — known divergence at prompt-completion boundary
+        compare_mask[:, 0] = False
+
+    if compare_mask.any():
+        real_diff = diff[compare_mask]
+        max_diff = real_diff.max().item()
+        mean_diff = real_diff.mean().item()
+    else:
+        max_diff = mean_diff = 0.0
+    passed = max_diff < max_thresh and mean_diff < mean_thresh
+    return max_diff, mean_diff, passed
+
+
+def test_scoring_correctness():
+    """Test 1: Deferred scoring logprobs match between padded and flattened.
+
+    Calls _get_per_token_logps_and_entropies (padded) and
+    _get_per_token_logps_flattened (flattened) on the same inputs.
+    """
+    print("=" * 60)
+    print("Test 1: Scoring path correctness (no grad)")
+    print("=" * 60)
+
+    model = setup_model()
+    trainer = make_mock_trainer(model)
+    input_ids, attn_mask, compl_mask, logits_to_keep, prompt_mask, meta = (
+        make_grpo_batch(B=8)
+    )
+
+    print(
+        f"  Batch: {input_ids.shape[0]} seqs, max_len={input_ids.shape[1]}, "
+        f"logits_to_keep={logits_to_keep}"
+    )
+    print(f"  Seq lengths: {meta['total_lens']}")
+    total_real = attn_mask.sum().item()
+    total_padded = input_ids.numel()
+    print(f"  Padding ratio: {1 - total_real / total_padded:.1%}")
+
+    with torch.no_grad():
+        logps_pad, _ = trainer._get_per_token_logps_and_entropies(
+            model, input_ids, attn_mask, logits_to_keep
+        )
+        logps_flat = trainer._get_per_token_logps_flattened(
+            model, input_ids, attn_mask, logits_to_keep, prompt_mask=prompt_mask
+        )
+
+    max_diff, mean_diff, passed = _compare_logps(logps_pad, logps_flat, mask=compl_mask)
+
+    print(f"  Max diff:  {max_diff:.8f}")
+    print(f"  Mean diff: {mean_diff:.8f}")
+    print(
+        "  (bf16 flash attention varlen uses different accumulation order than padded;"
+    )
+    print("   per-token diffs up to ~0.5 are expected and average out in the loss)")
+    print(f"  Result: {'PASS' if passed else 'FAIL'}")
+    print()
+    return passed
+
+
+def test_training_loss_correctness():
+    """Test 2: Training logprobs match between padded and flattened (with grad)."""
+    print("=" * 60)
+    print("Test 2: Training loss correctness (with grad)")
+    print("=" * 60)
+
+    model = setup_model(eval_mode=False)
+    trainer = make_mock_trainer(model)
+    input_ids, attn_mask, _compl_mask, logits_to_keep, prompt_mask, _meta = (
+        make_grpo_batch(B=4)
+    )
+
+    print(f"  Batch: {input_ids.shape[0]} seqs, logits_to_keep={logits_to_keep}")
+
+    # Padded path (with grad)
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        logps_pad, _ = trainer._get_per_token_logps_and_entropies(
+            model, input_ids, attn_mask, logits_to_keep
+        )
+
+    # Flattened path (with grad)
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        logps_flat, _ = trainer._get_per_token_logps_and_entropies_flattened(
+            model, input_ids, attn_mask, logits_to_keep, prompt_mask=prompt_mask
+        )
+
+    max_diff, mean_diff, _ = _compare_logps(logps_pad.detach(), logps_flat.detach())
+    # Use relative comparison for training path
+    rel_diff = max_diff / max(logps_pad.detach().float().abs().max().item(), 1e-8)
+
+    print(f"  Max diff:     {max_diff:.8f}")
+    print(f"  Mean diff:    {mean_diff:.8f}")
+    print(f"  Relative max: {rel_diff:.4%}")
+
+    passed = rel_diff < 0.10 and mean_diff < 0.1
+    print(f"  Result: {'PASS' if passed else 'FAIL'}")
+    print()
+    return passed
+
+
+def test_gradient_correctness():
+    """Test 3: Gradients match between padded and flattened training paths."""
+    print("=" * 60)
+    print("Test 3: Gradient correctness")
+    print("=" * 60)
+
+    input_ids, attn_mask, compl_mask, logits_to_keep, prompt_mask, _meta = (
+        make_grpo_batch(B=4)
+    )
+    advantages = torch.randn(input_ids.shape[0], device="cuda")
+
+    # Model 1: padded path
+    model_pad = setup_model(eval_mode=False)
+    trainer_pad = make_mock_trainer(model_pad)
+
+    with torch.no_grad():
+        old_logps, _ = trainer_pad._get_per_token_logps_and_entropies(
+            model_pad, input_ids, attn_mask, logits_to_keep
+        )
+
+    model_pad.zero_grad()
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        logps_pad, _ = trainer_pad._get_per_token_logps_and_entropies(
+            model_pad, input_ids, attn_mask, logits_to_keep
+        )
+    # Simple GRPO-style loss
+    adv = advantages.unsqueeze(1)
+    ratio_pad = torch.exp(logps_pad - old_logps.detach())
+    loss_pad = -(ratio_pad * adv * compl_mask).sum() / compl_mask.sum().clamp(min=1)
+    loss_pad.backward()
+
+    # Model 2: flattened path
+    model_flat = setup_model(eval_mode=False)
+    trainer_flat = make_mock_trainer(model_flat)
+
+    model_flat.zero_grad()
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        logps_flat, _ = trainer_flat._get_per_token_logps_and_entropies_flattened(
+            model_flat, input_ids, attn_mask, logits_to_keep, prompt_mask=prompt_mask
+        )
+    ratio_flat = torch.exp(logps_flat - old_logps.detach())
+    loss_flat = -(ratio_flat * adv * compl_mask).sum() / compl_mask.sum().clamp(min=1)
+    loss_flat.backward()
+
+    # Compare gradients
+    max_grad_diff = 0.0
+    max_grad_mag = 0.0
+    n_params = 0
+    for (_n1, p1), (_n2, p2) in zip(
+        model_pad.named_parameters(), model_flat.named_parameters(), strict=True
+    ):
+        if p1.grad is not None and p2.grad is not None:
+            diff = (p1.grad.float() - p2.grad.float()).abs().max().item()
+            max_grad_diff = max(max_grad_diff, diff)
+            max_grad_mag = max(max_grad_mag, p1.grad.float().abs().max().item())
+            n_params += 1
+
+    rel_grad_diff = max_grad_diff / max(max_grad_mag, 1e-8)
+    print(f"  Loss padded:   {loss_pad.item():.8f}")
+    print(f"  Loss flattened:{loss_flat.item():.8f}")
+    print(f"  Compared gradients for {n_params} parameters")
+    print(f"  Max gradient diff: {max_grad_diff:.8f}")
+    print(f"  Max gradient magnitude: {max_grad_mag:.8f}")
+    print(f"  Relative gradient diff: {rel_grad_diff:.4%}")
+
+    passed = rel_grad_diff < 0.15
+    print(f"  Result: {'PASS' if passed else 'FAIL'}")
+    print()
+
+    del model_pad, model_flat
+    torch.cuda.empty_cache()
+    return passed
+
+
+def test_variable_completion_lengths():
+    """Test 4: Correctness with variable prompt lengths (GRPO data layout).
+
+    Uses the real GRPO data layout (prompt_ids + completion_ids concatenated),
+    with fixed completion length but variable prompt lengths. Tests that batch
+    flattening handles prompt padding correctly.
+    """
+    print("=" * 60)
+    print("Test 4: Variable prompt lengths (GRPO layout)")
+    print("=" * 60)
+
+    model = setup_model()
+    trainer = make_mock_trainer(model)
+
+    torch.manual_seed(123)
+    B = 8
+    max_compl = 64
+    prompt_lens = [10, 25, 15, 30, 8, 20, 35, 12]
+    compl_lens = [max_compl] * B
+    max_prompt = max(prompt_lens)
+
+    # Build GRPO-style: prompt_ids (B, max_prompt) + completion_ids (B, max_compl)
+    prompt_ids = torch.zeros(B, max_prompt, dtype=torch.long, device="cuda")
+    completion_ids = torch.randint(100, 5000, (B, max_compl), device="cuda")
+    p_mask_raw = torch.zeros(B, max_prompt, dtype=torch.long, device="cuda")
+    for i in range(B):
+        prompt_ids[i, : prompt_lens[i]] = torch.randint(
+            100, 5000, (prompt_lens[i],), device="cuda"
+        )
+        p_mask_raw[i, : prompt_lens[i]] = 1
+
+    input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+    c_mask_raw = torch.ones(B, max_compl, dtype=torch.long, device="cuda")
+    attn_mask = torch.cat([p_mask_raw, c_mask_raw], dim=1)
+    p_mask = torch.cat(
+        [p_mask_raw, torch.zeros(B, max_compl, dtype=torch.long, device="cuda")], dim=1
+    )
+
+    total_real = attn_mask.sum().item()
+    total_padded = input_ids.numel()
+    print(f"  Batch: {B} seqs, max_len={input_ids.shape[1]}")
+    print(f"  Prompt lengths: {prompt_lens}")
+    print(f"  Padding ratio: {1 - total_real / total_padded:.1%}")
+
+    with torch.no_grad():
+        logps_pad, _ = trainer._get_per_token_logps_and_entropies(
+            model, input_ids, attn_mask, max_compl
+        )
+        logps_flat = trainer._get_per_token_logps_flattened(
+            model, input_ids, attn_mask, max_compl, prompt_mask=p_mask
+        )
+
+    # skip_first=True because variable prompt padding causes position-0 divergence
+    max_diff, mean_diff, passed = _compare_logps(logps_pad, logps_flat)
+
+    print(f"  Max diff:  {max_diff:.8f}")
+    print(f"  Mean diff: {mean_diff:.8f}")
+
+    # Per-sequence check
+    diff = (logps_pad.float() - logps_flat.float()).abs()
+    for i in range(B):
+        seq_diff = diff[i, : compl_lens[i]].max().item() if compl_lens[i] > 0 else 0.0
+        status = "ok" if seq_diff < 1.0 else "BAD"
+        print(
+            f"    seq {i} (compl={compl_lens[i]:3d}): max_diff={seq_diff:.8f} {status}"
+        )
+
+    print(f"  Result: {'PASS' if passed else 'FAIL'}")
+    print()
+    return passed
+
+
+def test_prompt_mask_edge_case():
+    """Test 5: logits_to_keep > actual completion length (the 4B explosion bug).
+
+    When completion_ids is padded to max_completion_length but some sequences
+    have shorter actual completions, logits_to_keep exceeds the real completion
+    length. Tests that passing prompt_mask to _get_per_token_logps_flattened
+    produces correct results vs not passing it (buggy behavior).
+    """
+    print("=" * 60)
+    print("Test 5: prompt_mask edge case (logits_to_keep > completion)")
+    print("=" * 60)
+
+    model = setup_model()
+    trainer = make_mock_trainer(model)
+
+    torch.manual_seed(99)
+    B = 4
+    logits_to_keep = 128
+    prompt_lens = [30, 20, 40, 25]
+    compl_lens = [50, 128, 30, 100]
+    total_lens = [p + c for p, c in zip(prompt_lens, compl_lens, strict=True)]
+    max_len = max(p + logits_to_keep for p in prompt_lens)
+
+    input_ids = torch.zeros(B, max_len, dtype=torch.long, device="cuda")
+    attention_mask = torch.zeros(B, max_len, dtype=torch.long, device="cuda")
+    prompt_mask_tensor = torch.zeros(B, max_len, dtype=torch.long, device="cuda")
+
+    for i in range(B):
+        tl = total_lens[i]
+        input_ids[i, :tl] = torch.randint(100, 5000, (tl,), device="cuda")
+        attention_mask[i, :tl] = 1
+        prompt_mask_tensor[i, : prompt_lens[i]] = 1
+
+    print(f"  logits_to_keep={logits_to_keep}, actual completions={compl_lens}")
+    total_real = attention_mask.sum().item()
+    print(f"  Padding ratio: {1 - total_real / (B * max_len):.1%}")
+
+    with torch.no_grad():
+        # Padded reference (always correct since it uses logits_to_keep slicing)
+        logps_pad, _ = trainer._get_per_token_logps_and_entropies(
+            model, input_ids, attention_mask, logits_to_keep
+        )
+
+        # Flattened WITH prompt_mask (correct)
+        logps_flat_correct = trainer._get_per_token_logps_flattened(
+            model,
+            input_ids,
+            attention_mask,
+            logits_to_keep,
+            prompt_mask=prompt_mask_tensor,
+        )
+
+        # Flattened WITHOUT prompt_mask (buggy -- infers prompt_len as seq_len - logits_to_keep)
+        logps_flat_buggy = trainer._get_per_token_logps_flattened(
+            model,
+            input_ids,
+            attention_mask,
+            logits_to_keep,
+            prompt_mask=None,
+        )
+
+    # Compare with-prompt-mask vs without-prompt-mask directly.
+    # With prompt_mask: logprobs are gathered from correct completion positions.
+    # Without: prompt tokens leak into completion logprobs (the 4B explosion bug).
+    # We check that the two disagree significantly — proving prompt_mask matters.
+    diff_between = (logps_flat_correct.float() - logps_flat_buggy.float()).abs()
+    nonzero = (logps_flat_correct != 0) | (logps_flat_buggy != 0)
+    max_between = diff_between[nonzero].max().item() if nonzero.any() else 0.0
+
+    # Also check correct path against padded (skip position 0 due to prompt padding)
+    diff_correct = (logps_pad.float() - logps_flat_correct.float()).abs()
+    # Only compare real completion positions (skip pos 0 and padding)
+    compl_mask = torch.zeros_like(diff_correct)
+    for i in range(B):
+        compl_mask[i, 1 : compl_lens[i]] = 1.0  # skip pos 0
+    masked_diff = diff_correct * compl_mask
+    max_correct = masked_diff.max().item()
+    max_buggy = max_between  # how much the buggy path disagrees with correct
+
+    print(f"  With prompt_mask:    max_diff={max_correct:.4f}")
+    print(f"  Without prompt_mask: max_diff={max_buggy:.4f}")
+    print("  (buggy path grabs prompt tokens as completion -> huge diff)")
+
+    # prompt_mask path should be significantly better than buggy path
+    passed = max_correct < max_buggy
+    print(f"  Result: {'PASS' if passed else 'FAIL'}")
+    print()
+    return passed
+
+
+def test_training_flattened_gradients():
+    """Test 6: Training forward+backward with flattened method produces correct gradients.
+
+    Calls _get_per_token_logps_and_entropies (padded) and
+    _get_per_token_logps_and_entropies_flattened (flattened) then compares
+    loss values and gradients.
+    """
+    print("=" * 60)
+    print("Test 6: Training fwd+bwd flattening (gradient check)")
+    print("=" * 60)
+
+    input_ids, attn_mask, compl_mask, logits_to_keep, prompt_mask, _meta = (
+        make_grpo_batch(B=4)
+    )
+    advantages = torch.randn(input_ids.shape[0], device="cuda")
+
+    # Get old_logps for the loss computation (shared between both paths)
+    ref_model = setup_model()
+    ref_trainer = make_mock_trainer(ref_model)
+    with torch.no_grad():
+        old_logps, _ = ref_trainer._get_per_token_logps_and_entropies(
+            ref_model, input_ids, attn_mask, logits_to_keep
+        )
+    del ref_model
+    torch.cuda.empty_cache()
+
+    adv = advantages.unsqueeze(1)
+
+    # Padded loss + backward
+    model_pad = setup_model(eval_mode=False)
+    trainer_pad = make_mock_trainer(model_pad)
+    model_pad.zero_grad()
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        logps_pad, _ = trainer_pad._get_per_token_logps_and_entropies(
+            model_pad, input_ids, attn_mask, logits_to_keep
+        )
+    ratio_pad = torch.exp(logps_pad - old_logps.detach())
+    loss_pad = -(ratio_pad * adv * compl_mask).sum() / compl_mask.sum().clamp(min=1)
+    loss_pad.backward()
+
+    # Flattened loss + backward
+    model_flat = setup_model(eval_mode=False)
+    trainer_flat = make_mock_trainer(model_flat)
+    model_flat.zero_grad()
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        logps_flat, _ = trainer_flat._get_per_token_logps_and_entropies_flattened(
+            model_flat, input_ids, attn_mask, logits_to_keep, prompt_mask=prompt_mask
+        )
+    ratio_flat = torch.exp(logps_flat - old_logps.detach())
+    loss_flat = -(ratio_flat * adv * compl_mask).sum() / compl_mask.sum().clamp(min=1)
+    loss_flat.backward()
+
+    # Compare
+    rel_loss = abs(loss_pad.item() - loss_flat.item()) / max(abs(loss_pad.item()), 1e-8)
+
+    max_grad_diff = 0.0
+    max_grad_mag = 0.0
+    n_params = 0
+    for (_n1, p1), (_n2, p2) in zip(
+        model_pad.named_parameters(), model_flat.named_parameters(), strict=True
+    ):
+        if p1.grad is not None and p2.grad is not None:
+            diff = (p1.grad.float() - p2.grad.float()).abs().max().item()
+            max_grad_diff = max(max_grad_diff, diff)
+            max_grad_mag = max(max_grad_mag, p1.grad.float().abs().max().item())
+            n_params += 1
+
+    rel_grad = max_grad_diff / max(max_grad_mag, 1e-8)
+
+    print(f"  Padded loss:  {loss_pad.item():.8f}")
+    print(f"  Flat loss:    {loss_flat.item():.8f}")
+    print(f"  Rel loss diff: {rel_loss:.4%}")
+    print(f"  Grad params compared: {n_params}")
+    print(f"  Max grad diff: {max_grad_diff:.8f}, mag: {max_grad_mag:.8f}")
+    print(f"  Rel grad diff: {rel_grad:.4%}")
+
+    passed = rel_loss < 0.05 and rel_grad < 0.15
+    print(f"  Result: {'PASS' if passed else 'FAIL'}")
+    print()
+
+    del model_pad, model_flat
+    torch.cuda.empty_cache()
+    return passed
+
+
+if __name__ == "__main__":
+    print("\nBatch Flattening Correctness Tests")
+    print(f"Model: {MODEL_NAME}")
+    print(f"{'=' * 60}\n")
+
+    results = []
+    results.append(("Scoring correctness", test_scoring_correctness()))
+    results.append(("Training loss", test_training_loss_correctness()))
+    results.append(("Gradient correctness", test_gradient_correctness()))
+    results.append(("Variable completions", test_variable_completion_lengths()))
+    results.append(("prompt_mask edge case", test_prompt_mask_edge_case()))
+    results.append(("Training fwd+bwd flat", test_training_flattened_gradients()))
+
+    print("=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    all_passed = True
+    for name, passed in results:
+        status = "PASS" if passed else "FAIL"
+        print(f"  {name:30s} {status}")
+        all_passed = all_passed and passed
+
+    print(f"\n  Overall: {'ALL TESTS PASSED' if all_passed else 'SOME TESTS FAILED'}")
+    print()
diff --git a/tests/e2e/solo/test_flex.py b/tests/e2e/solo/test_flex.py
new file mode 100644
index 0000000000..abe8fb69a0
--- /dev/null
+++ b/tests/e2e/solo/test_flex.py
@@ -0,0 +1,67 @@
+"""
+E2E tests for packed training w/ flex attention
+"""
+
+import unittest
+
+from transformers.utils import is_torch_bf16_gpu_available
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from ..utils import check_tensorboard, require_torch_2_6_0, with_temp_dir
+
+
+class TestPackedFlex(unittest.TestCase):
+    """
+    Test case for Packed training of llama models
+    """
+
+    @require_torch_2_6_0
+    @with_temp_dir
+    def test_loss_llama(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "sample_packing": True,
+                "flex_attention": True,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+        if is_torch_bf16_gpu_available():
+            cfg.bf16 = True
+        else:
+            cfg.fp16 = True
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.1, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/solo/test_relora_llama.py b/tests/e2e/solo/test_relora_llama.py
new file mode 100644
index 0000000000..895f32d990
--- /dev/null
+++ b/tests/e2e/solo/test_relora_llama.py
@@ -0,0 +1,148 @@
+"""
+E2E tests for relora llama
+"""
+
+import unittest
+from pathlib import Path
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from ..utils import check_model_output_exists, check_tensorboard, with_temp_dir
+
+
+class TestReLoraLlama(unittest.TestCase):
+    """
+    Test case for Llama models using LoRA
+    """
+
+    @with_temp_dir
+    def test_relora(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 2048,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "flash_attention": True,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_modules": ["q_proj", "v_proj"],
+                "relora": True,
+                "jagged_restart_steps": 50,
+                "jagged_restart_warmup_steps": 10,
+                "jagged_restart_anneal_steps": 10,
+                "relora_prune_ratio": 0.9,
+                "relora_cpu_offload": True,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "split": "train[:10%]",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "warmup_steps": 10,
+                "num_epochs": 2,
+                "max_steps": 105,  # at least 2x restart cadence
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-100/adapter", cfg)
+        assert (Path(temp_dir) / "checkpoint-100/relora/model.safetensors").exists(), (
+            "Relora model checkpoint not found"
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/grad_norm", 0.2, "grad_norm is too high"
+        )
+
+    @with_temp_dir
+    def test_relora_reset_method(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 2048,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "flash_attention": True,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_modules": ["q_proj", "v_proj"],
+                "relora": True,
+                "jagged_restart_steps": 50,
+                "jagged_restart_warmup_steps": 10,
+                "jagged_restart_anneal_steps": 10,
+                "relora_prune_ratio": 0.5,  # explicitly honored by reset (not ignored)
+                "relora_prune_method": "reset",
+                "relora_cpu_offload": True,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "split": "train[:10%]",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "warmup_steps": 10,
+                "num_epochs": 2,
+                "max_steps": 105,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-100/adapter", cfg)
+        assert (Path(temp_dir) / "checkpoint-100/relora/model.safetensors").exists(), (
+            "Relora model checkpoint not found"
+        )
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/grad_norm", 0.2, "grad_norm is too high"
+        )
diff --git a/tests/e2e/solo/test_reward_model_smollm2.py b/tests/e2e/solo/test_reward_model_smollm2.py
new file mode 100644
index 0000000000..e4240ab34f
--- /dev/null
+++ b/tests/e2e/solo/test_reward_model_smollm2.py
@@ -0,0 +1,71 @@
+"""
+E2E tests for reward model lora llama
+"""
+
+import unittest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from ..utils import check_model_output_exists, check_tensorboard, with_temp_dir
+
+
+class TestRewardModelLoraSmolLM2(unittest.TestCase):
+    """
+    Test case for Llama reward models using LoRA
+    """
+
+    @with_temp_dir
+    def test_rm_lora(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForSequenceClassification",
+                "num_labels": 1,
+                "chat_template": "alpaca",
+                "reward_model": True,
+                "sequence_len": 2048,
+                "pad_to_sequence_len": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "argilla/distilabel-intel-orca-dpo-pairs",
+                        "type": "bradley_terry.chat_template",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "lora_modules_to_save": ["embed_tokens", "lm_head"],
+                "remove_unused_columns": False,
+                "max_steps": 10,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "gradient_checkpointing": True,
+                "warmup_ratio": 0.1,
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.5, "Train Loss (%s) is too high"
+        )
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/solo/test_trainer_loss_calc.py b/tests/e2e/solo/test_trainer_loss_calc.py
new file mode 100644
index 0000000000..cdb51990c2
--- /dev/null
+++ b/tests/e2e/solo/test_trainer_loss_calc.py
@@ -0,0 +1,29 @@
+"""Unit tests for trainer loss calc monkeypatch."""
+
+import unittest
+
+import pytest
+
+from axolotl.monkeypatch.transformers.trainer_loss_calc import (
+    check_evaluation_loop_is_patchable,
+    check_maybe_log_save_evaluate_is_patchable,
+)
+
+
+class TestTrainerLossCalc(unittest.TestCase):
+    """
+    Unit test class for trainer loss calc monkeypatch
+    """
+
+    @pytest.mark.xfail(reason="flaky", strict=False)
+    def test_trainer_loss_calc_is_patchable(self):
+        """
+        Test that the upstream transformers code is still patchable. This will fail if
+        the patched code changes upstream.
+        """
+        assert check_evaluation_loop_is_patchable()
+        assert check_maybe_log_save_evaluate_is_patchable()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/e2e/test_activation_offloading.py b/tests/e2e/test_activation_offloading.py
new file mode 100644
index 0000000000..9c799bd41a
--- /dev/null
+++ b/tests/e2e/test_activation_offloading.py
@@ -0,0 +1,429 @@
+"""
+E2E tests for activation offloading
+"""
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.core.trainers.mixins.activation_checkpointing import (
+    ActivationOffloadingMixin,
+)
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_model_output_exists
+
+
+class TestActivationOffloading:
+    """
+    E2E test cases for activation offloading
+    """
+
+    @pytest.mark.parametrize(
+        "adapter",
+        ["lora", "qlora", None],
+    )
+    def test_activation_offloading(
+        self,
+        temp_dir,
+        adapter,
+    ):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                    "eos_token": "<|im_end|>",
+                },
+                "datasets": [
+                    {
+                        "chat_template": "chatml",
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "split": "train[:10%]",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": "auto",
+                "gradient_checkpointing": True,
+                "activation_offloading": True,
+                "save_first_step": False,
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_target_linear": True,
+            }
+        )
+        if adapter == "lora":
+            cfg["adapter"] = "lora"
+        if adapter == "qlora":
+            cfg["adapter"] = "qlora"
+            cfg["load_in_4bit"] = True
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    @pytest.mark.parametrize(
+        "offload_mode,expect_streams",
+        [(True, True), ("legacy", False), ("disk", True)],
+    )
+    def test_offload_mode_wiring(
+        self, temp_dir, monkeypatch, offload_mode, expect_streams
+    ):
+        """`activation_offloading` must reach the trainer as a live offload
+        context: True => stream-overlapped, 'legacy' => synchronous. Guards the
+        regression where string modes fell through to plain gradient
+        checkpointing (no offload at all)."""
+        from trl.models.activation_offloading import OffloadActivations
+
+        captured = {}
+        original_step = ActivationOffloadingMixin.training_step
+
+        def capture_step(self, *args, **kwargs):
+            captured["ctx"] = self.activation_offload_context
+            return original_step(self, *args, **kwargs)
+
+        monkeypatch.setattr(ActivationOffloadingMixin, "training_step", capture_step)
+
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "val_set_size": 0.0,
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "datasets": [
+                    {"path": "mhenrichsen/alpaca_2k_test", "type": "alpaca"},
+                ],
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-5,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": "auto",
+                "gradient_checkpointing": True,
+                "activation_offloading": offload_mode,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_target_linear": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+        train(cfg=cfg, dataset_meta=dataset_meta)
+
+        ctx = captured.get("ctx")
+        assert isinstance(ctx, OffloadActivations), (
+            f"activation_offloading={offload_mode!r} did not produce an offload "
+            f"context (got {type(ctx).__name__}) — string modes likely fell "
+            f"through to plain gradient checkpointing"
+        )
+        assert ctx.use_streams is expect_streams
+
+    @pytest.mark.parametrize(
+        "adapter,expect_recompute_wrap",
+        [("lora", False), (None, True)],
+    )
+    def test_offload_is_adapter_aware(
+        self, temp_dir, monkeypatch, adapter, expect_recompute_wrap
+    ):
+        """activation_offloading is adapter-aware: LoRA offloads *instead of*
+        recomputing (no checkpoint wrap — pure offload is leaner/faster), while
+        full finetune keeps recompute and offloads the checkpoint boundaries
+        (the combo — pure offload is PCIe-bound and OOMs at full-param scale)."""
+        captured = {}
+        original_step = ActivationOffloadingMixin.training_step
+
+        def capture_step(self, *args, **kwargs):
+            captured["wrapped"] = any(
+                "_checkpoint_wrapped_module" in n for n, _ in self.model.named_modules()
+            )
+            return original_step(self, *args, **kwargs)
+
+        monkeypatch.setattr(ActivationOffloadingMixin, "training_step", capture_step)
+
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "val_set_size": 0.0,
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "datasets": [
+                    {"path": "mhenrichsen/alpaca_2k_test", "type": "alpaca"},
+                ],
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-5,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": "auto",
+                "gradient_checkpointing": True,
+                "activation_offloading": True,
+                "save_first_step": False,
+            }
+        )
+        if adapter:
+            cfg["adapter"] = adapter
+            cfg["lora_r"] = 8
+            cfg["lora_alpha"] = 16
+            cfg["lora_target_linear"] = True
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+        train(cfg=cfg, dataset_meta=dataset_meta)
+
+        assert captured.get("wrapped") is expect_recompute_wrap
+
+    @pytest.mark.parametrize("use_reentrant", [False, True])
+    def test_hidden_states_offload_full_param(
+        self, temp_dir, monkeypatch, use_reentrant
+    ):
+        """activation_offloading: hidden_states trains in both checkpoint modes."""
+        import torch.utils.checkpoint as ckpt
+
+        from axolotl.monkeypatch.activation_offload_checkpoint import (
+            HiddenStatesOffloadCheckpoint,
+            unpatch_hidden_states_offload,
+        )
+        from axolotl.monkeypatch.checkpoint_activation_offload import (
+            CheckpointHiddenStatesOffload,
+        )
+
+        entered_checkpoint_offload = False
+        original_enter = CheckpointHiddenStatesOffload.__enter__
+
+        def wrapped_enter(self):
+            nonlocal entered_checkpoint_offload
+            entered_checkpoint_offload = True
+            return original_enter(self)
+
+        if not use_reentrant:
+            monkeypatch.setattr(
+                CheckpointHiddenStatesOffload, "__enter__", wrapped_enter
+            )
+
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "val_set_size": 0.0,
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "datasets": [{"path": "mhenrichsen/alpaca_2k_test", "type": "alpaca"}],
+                "max_steps": 2,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-5,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": "auto",
+                "gradient_checkpointing": True,
+                "activation_offloading": "hidden_states",
+                "save_first_step": False,
+            }
+        )
+        if use_reentrant:
+            cfg["gradient_checkpointing_kwargs"] = {"use_reentrant": True}
+        try:
+            cfg = validate_config(cfg)
+            assert cfg.gradient_checkpointing_kwargs["use_reentrant"] is use_reentrant
+            normalize_config(cfg)
+            dataset_meta = load_datasets(cfg=cfg)
+            train(cfg=cfg, dataset_meta=dataset_meta)
+            if use_reentrant:
+                assert ckpt.CheckpointFunction is HiddenStatesOffloadCheckpoint
+            else:
+                assert ckpt.CheckpointFunction is not HiddenStatesOffloadCheckpoint
+                assert entered_checkpoint_offload is True
+            check_model_output_exists(temp_dir, cfg)
+        finally:
+            unpatch_hidden_states_offload()
+
+    def test_hidden_states_offload_numerical_parity(self):
+        """The offloaded checkpoint only round-trips the layer input through CPU,
+        so loss and grads must match plain reentrant checkpointing. Same model
+        instance for both runs => identical weights; the only difference is the
+        d2h/h2d of the per-layer input."""
+        import torch
+        import torch.utils.checkpoint as ckpt
+        from transformers import AutoModelForCausalLM
+
+        from axolotl.monkeypatch.activation_offload_checkpoint import (
+            HiddenStatesOffloadCheckpoint,
+            patch_hidden_states_offload,
+            unpatch_hidden_states_offload,
+        )
+
+        if not torch.cuda.is_available():
+            pytest.skip("hidden_states offload requires CUDA")
+
+        torch.manual_seed(0)
+        model = AutoModelForCausalLM.from_pretrained(
+            "HuggingFaceTB/SmolLM2-135M", dtype=torch.bfloat16
+        ).to("cuda")
+        model.train()
+        model.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs={"use_reentrant": True}
+        )
+
+        torch.manual_seed(1)
+        input_ids = torch.randint(0, 1000, (1, 512), device="cuda")
+        batch = {"input_ids": input_ids, "labels": input_ids.clone()}
+
+        def run():
+            model.zero_grad(set_to_none=True)
+            loss = model(**batch).loss
+            loss.backward()
+            grad = model.model.layers[0].mlp.down_proj.weight.grad
+            return loss.detach().clone(), grad.detach().clone()
+
+        assert ckpt.CheckpointFunction is not HiddenStatesOffloadCheckpoint
+        loss_ref, grad_ref = run()
+
+        patch_hidden_states_offload()
+        try:
+            assert ckpt.CheckpointFunction is HiddenStatesOffloadCheckpoint
+            loss_off, grad_off = run()
+        finally:
+            unpatch_hidden_states_offload()
+
+        # Forward output is identical; backward recompute introduces only bf16
+        # float noise, so grads match within a loose tolerance.
+        assert torch.allclose(loss_ref, loss_off, rtol=0, atol=1e-3), (
+            f"loss diverged: ref={loss_ref.item()} offload={loss_off.item()}"
+        )
+        assert torch.allclose(grad_ref, grad_off, rtol=1e-2, atol=1e-2), (
+            f"grad diverged: max|d|={(grad_ref - grad_off).abs().max().item()}"
+        )
+
+    def test_no_vram_leak_regression(self, temp_dir, monkeypatch):
+        """#3638 regression — fail on linear VRAM growth across training steps.
+
+        The bug: ``OffloadActivations.__enter__`` doesn't clear cross-step
+        state, so a saved tensor that never unpacks during backward
+        (MoE / ``torch.compile``) sits in ``ctx.tracker`` forever — and its
+        GPU storage stays alive. Across many steps memory grows linearly.
+
+        Tiny CI models won't exhibit the upstream MoE/compile unpack failure
+        on their own, so we *inject* the same leftover: after every step we
+        stash a small CUDA tensor into ``ctx.tracker``. The fix clears it on
+        the next ``__enter__`` (memory flat); without the fix it accumulates
+        (memory grows ~constant bytes/step). The fail mode is the bug's own
+        symptom — ``torch.cuda.memory_allocated`` increasing across steps.
+        """
+        import torch
+
+        if not torch.cuda.is_available():
+            pytest.skip("VRAM-leak test requires CUDA")
+
+        mem_per_step: list[int] = []
+        seed_id = [10**9]
+        seed_bytes = 4 * 1024 * 1024  # 4 MB / step
+
+        original_step = ActivationOffloadingMixin.training_step
+
+        def wrapped_step(self, *args, **kwargs):
+            torch.cuda.synchronize()
+            mem_per_step.append(torch.cuda.memory_allocated())
+            out = original_step(self, *args, **kwargs)
+
+            # Inject the MoE-style leftover: a CUDA tensor stuck in
+            # OffloadActivations.tracker. The local `seed` ref dies on
+            # return — only ctx.tracker keeps it alive, so the next
+            # __enter__'s clear (with the fix) actually releases the GPU
+            # memory. Without the fix these accumulate step-over-step.
+            ctx = self.activation_offload_context
+            seed_id[0] += 1
+            seed = torch.empty(seed_bytes // 2, dtype=torch.float16, device="cuda")
+            ctx.tracker[seed_id[0]] = (seed, False, None, None, None)
+            # Stop the next forward's pack_tensor from raising on its
+            # "tracker should have been cleared" guard. With the fix this
+            # flag gets reset by __enter__ anyway; on main it would
+            # otherwise crash before our VRAM measurement on step 2.
+            ctx.is_first_forward_call = False
+            return out
+
+        monkeypatch.setattr(ActivationOffloadingMixin, "training_step", wrapped_step)
+
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "val_set_size": 0.0,
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "datasets": [
+                    {"path": "mhenrichsen/alpaca_2k_test", "type": "alpaca"},
+                ],
+                "max_steps": 10,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-5,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "bf16": "auto",
+                "gradient_checkpointing": True,
+                "activation_offloading": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+        train(cfg=cfg, dataset_meta=dataset_meta)
+
+        # Drop warm-up steps; allocator settling distorts early samples.
+        warmup = 3
+        samples = mem_per_step[warmup:]
+        assert len(samples) >= 5, (
+            f"need >= 5 post-warmup samples, got {len(samples)} "
+            f"(total {len(mem_per_step)})"
+        )
+
+        # Injection is 4 MB/step. With the fix __enter__ clears each seed
+        # before the next step → growth ≈ 0. Without the fix seeds pile up
+        # → growth ≈ 4 MB × (steps-1). 10 MB is well above allocator jitter
+        # and well below the leaky-build floor.
+        growth_mb = (samples[-1] - samples[0]) / (1024**2)
+        tolerance_mb = 10
+
+        per_step_mb = [round(m / 1024**2, 1) for m in mem_per_step]
+        assert growth_mb < tolerance_mb, (
+            f"VRAM grew {growth_mb:.1f} MB across {len(samples)} post-warmup "
+            f"steps — linear-increase signature of the #3638 VRAM leak. "
+            f"Per-step memory_allocated (MB): {per_step_mb}"
+        )
+
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/test_deepseekv3.py b/tests/e2e/test_deepseekv3.py
new file mode 100644
index 0000000000..05b2381838
--- /dev/null
+++ b/tests/e2e/test_deepseekv3.py
@@ -0,0 +1,127 @@
+"""
+E2E tests for deepseekv3
+"""
+
+from pathlib import Path
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.hf_offline_utils import enable_hf_offline
+
+
+@pytest.mark.skip(
+    reason="DeepSeek-V3-11M remote model code needs _supports_flash_attn=True for newer transformers"
+)
+class TestDeepseekV3:
+    """
+    Test case for DeepseekV3 models
+    """
+
+    @enable_hf_offline
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_lora_deepseekv3(self, temp_dir, sample_packing):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/DeepSeek-V3-11M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "drop_system_message": True,
+                        "split": "train[:1%]",
+                    },
+                ],
+                "special_tokens": {
+                    "bos_token": "<｜begin▁of▁sentence｜>",
+                    "eos_token": "<｜end▁of▁sentence｜>",
+                },
+                "chat_template": "deepseek_v3",
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+
+    @enable_hf_offline
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_fft_deepseekv3(self, temp_dir, sample_packing):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/DeepSeek-V3-11M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                        "split": "train[:1%]",
+                    },
+                ],
+                "chat_template": "deepseek_v3",
+                "special_tokens": {
+                    "bos_token": "<｜begin▁of▁sentence｜>",
+                    "eos_token": "<｜end▁of▁sentence｜>",
+                },
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "model.safetensors").exists()
diff --git a/tests/e2e/test_diffusion.py b/tests/e2e/test_diffusion.py
new file mode 100644
index 0000000000..048a3aaf33
--- /dev/null
+++ b/tests/e2e/test_diffusion.py
@@ -0,0 +1,139 @@
+"""E2E smoke test for diffusion training plugin."""
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_model_output_exists
+
+
+class TestDiffusion:
+    """Test case for diffusion training plugin."""
+
+    def test_diffusion_smoke_test(self, temp_dir):
+        """
+        Smoke test for diffusion training to ensure the plugin loads and trains without
+        error.
+        """
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "trust_remote_code": True,
+                "sequence_len": 256,
+                "val_set_size": 0.1,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "bf16": True,
+                "save_first_step": False,
+                "logging_steps": 1,
+                "eval_steps": 3,
+                # Diffusion-specific config
+                "plugins": ["axolotl.integrations.diffusion.DiffusionPlugin"],
+                "diffusion": {
+                    # sample generation
+                    "generate_samples": True,
+                    "generation_interval": 1,
+                    "num_generation_samples": 1,
+                    "generation_steps": 2,
+                    "generation_max_length": 32,
+                    "generation_temperature": 0.0,
+                    # training-specific
+                    "mask_token_id": 16,
+                    "eps": 1e-3,
+                    "importance_weighting": False,
+                },
+            }
+        )
+
+        prepare_plugins(cfg)
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    def test_diffusion_sft_labels(self, temp_dir):
+        """Test that diffusion training properly handles SFT data with labels."""
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "trust_remote_code": True,
+                "sequence_len": 256,
+                "val_set_size": 0.1,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "bf16": True,
+                "save_first_step": False,
+                "logging_steps": 1,
+                "eval_steps": 2,
+                # Diffusion-specific config
+                "plugins": ["axolotl.integrations.diffusion.DiffusionPlugin"],
+                "diffusion": {
+                    # sample generation
+                    "generate_samples": True,
+                    "generation_interval": 1,
+                    "num_generation_samples": 1,
+                    "generation_steps": 2,
+                    "generation_max_length": 32,
+                    "generation_temperature": 0.0,
+                    # training-specific
+                    "mask_token_id": 16,
+                    "eps": 1e-3,
+                    "importance_weighting": True,
+                },
+                # Ensure we have proper SFT labels
+                "train_on_inputs": False,
+            }
+        )
+
+        prepare_plugins(cfg)
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        # Verify that the dataset has labels
+        sample = dataset_meta.train_dataset[0]
+        assert "labels" in sample, "SFT dataset should have labels"
+
+        # Check that some labels are -100 (prompt tokens)
+        labels = sample["labels"]
+        if hasattr(labels, "tolist"):
+            labels = labels.tolist()
+        assert -100 in labels, "SFT dataset should have -100 labels for prompt tokens"
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/test_dpo.py b/tests/e2e/test_dpo.py
index 9596b1873f..24784eb2ca 100644
--- a/tests/e2e/test_dpo.py
+++ b/tests/e2e/test_dpo.py
@@ -1,27 +1,19 @@
-"""
-E2E tests for lora llama
-"""
+"""E2E tests for lora llama"""
 
-import logging
-import os
 import unittest
 from pathlib import Path
 
 import pytest
 
-from axolotl.cli import load_rl_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
+from axolotl.common.datasets import load_preference_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
+from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
 
-
-@pytest.mark.skip(reason="doesn't seem to work on modal")
 class TestDPOLlamaLora(unittest.TestCase):
     """
     Test case for DPO Llama models using LoRA
@@ -29,11 +21,58 @@ class TestDPOLlamaLora(unittest.TestCase):
 
     @with_temp_dir
     def test_dpo_lora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 64,
+                "lora_alpha": 32,
+                "lora_dropout": 0.1,
+                "lora_target_linear": True,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "rl": "dpo",
+                "datasets": [
+                    {
+                        "path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
+                        "type": "chatml.ultra",
+                        "split": "train",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "paged_adamw_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 20,
+                "save_steps": 10,
+                "warmup_steps": 5,
+                "gradient_checkpointing": True,
+                "gradient_checkpointing_kwargs": {"use_reentrant": True},
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+
+    @with_temp_dir
+    def test_dpo_use_weighting(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -41,12 +80,15 @@ def test_dpo_lora(self, temp_dir):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {},
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
                 "rl": "dpo",
+                "dpo_use_weighting": True,
                 "datasets": [
                     {
-                        "path": "Intel/orca_dpo_pairs",
-                        "type": "chatml.intel",
+                        "path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
+                        "type": "chatml.ultra",
                         "split": "train",
                     },
                 ],
@@ -62,22 +104,77 @@ def test_dpo_lora(self, temp_dir):
                 "warmup_steps": 5,
                 "gradient_checkpointing": True,
                 "gradient_checkpointing_kwargs": {"use_reentrant": True},
+                "save_first_step": False,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
-        dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
 
+    @with_temp_dir
+    def test_rpo(self, temp_dir):
+        # For TRL >= 0.29, loss_type=["sigmoid", "sft"], loss_weights=[1, alpha]
+        # replaces loss_type="rpo", rpo_alpha=alpha.
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 64,
+                "lora_alpha": 32,
+                "lora_dropout": 0.1,
+                "lora_target_linear": True,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "rl": "dpo",
+                "dpo_loss_type": ["sigmoid", "sft"],
+                "dpo_loss_weights": [1.0, 1.0],
+                "datasets": [
+                    {
+                        "path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
+                        "type": "chatml.ultra",
+                        "split": "train",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "paged_adamw_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 20,
+                "save_steps": 10,
+                "warmup_steps": 5,
+                "gradient_checkpointing": True,
+                "gradient_checkpointing_kwargs": {"use_reentrant": True},
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+
+    @pytest.mark.skip("kto_pair no longer supported in trl")
     @with_temp_dir
     def test_kto_pair_lora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -85,12 +182,14 @@ def test_kto_pair_lora(self, temp_dir):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {},
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
                 "rl": "kto_pair",
                 "datasets": [
                     {
-                        "path": "Intel/orca_dpo_pairs",
-                        "type": "chatml.intel",
+                        "path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
+                        "type": "chatml.ultra",
                         "split": "train",
                     },
                 ],
@@ -106,22 +205,24 @@ def test_kto_pair_lora(self, temp_dir):
                 "warmup_steps": 5,
                 "gradient_checkpointing": True,
                 "gradient_checkpointing_kwargs": {"use_reentrant": True},
+                "save_first_step": False,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
-        dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
 
     @with_temp_dir
     def test_ipo_lora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -129,12 +230,15 @@ def test_ipo_lora(self, temp_dir):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {},
-                "rl": "ipo",
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "rl": "dpo",
+                "dpo_loss_type": ["ipo"],
                 "datasets": [
                     {
-                        "path": "Intel/orca_dpo_pairs",
-                        "type": "chatml.intel",
+                        "path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
+                        "type": "chatml.ultra",
                         "split": "train",
                     },
                 ],
@@ -150,22 +254,25 @@ def test_ipo_lora(self, temp_dir):
                 "warmup_steps": 5,
                 "gradient_checkpointing": True,
                 "gradient_checkpointing_kwargs": {"use_reentrant": True},
+                "save_first_step": False,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
-        dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
 
+    @pytest.mark.skip(reason="TRL ORPO trainer has internal zip() length mismatch bug")
     @with_temp_dir
     def test_orpo_lora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -173,14 +280,16 @@ def test_orpo_lora(self, temp_dir):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {},
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
                 "rl": "orpo",
                 "orpo_alpha": 0.1,
                 "remove_unused_columns": False,
                 "chat_template": "chatml",
                 "datasets": [
                     {
-                        "path": "argilla/ultrafeedback-binarized-preferences-cleaned",
+                        "path": "argilla/distilabel-capybara-dpo-7k-binarized",
                         "type": "chat_template.argilla",
                         "split": "train",
                     },
@@ -197,11 +306,82 @@ def test_orpo_lora(self, temp_dir):
                 "warmup_steps": 5,
                 "gradient_checkpointing": True,
                 "gradient_checkpointing_kwargs": {"use_reentrant": True},
+                "save_first_step": False,
             }
         )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+
+    @pytest.mark.skip(reason="Fix the implementation")
+    @with_temp_dir
+    def test_kto_lora(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "LlamaTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 64,
+                "lora_alpha": 32,
+                "lora_dropout": 0.1,
+                "lora_target_linear": True,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "rl": "kto",
+                "rl_beta": 0.5,
+                "kto_desirable_weight": 1.0,
+                "kto_undesirable_weight": 1.0,
+                "remove_unused_columns": False,
+                "datasets": [
+                    # {
+                    #     "path": "argilla/kto-mix-15k",
+                    #     "type": "chatml.argilla_chat",
+                    #     "split": "train",
+                    # },
+                    {
+                        "path": "argilla/ultrafeedback-binarized-preferences-cleaned-kto",
+                        "type": "chatml.ultra",
+                        "split": "train",
+                    },
+                    # {
+                    #     "path": "argilla/kto-mix-15k",
+                    #     "type": "llama3.argilla_chat",
+                    #     "split": "train",
+                    # },
+                    {
+                        "path": "argilla/ultrafeedback-binarized-preferences-cleaned-kto",
+                        "type": "llama3.ultra",
+                        "split": "train",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "paged_adamw_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 20,
+                "save_steps": 10,
+                "warmup_steps": 5,
+                "gradient_checkpointing": True,
+                "gradient_checkpointing_kwargs": {"use_reentrant": True},
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
         cli_args = TrainerCliArgs()
-        dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
diff --git a/tests/e2e/test_embeddings_lr.py b/tests/e2e/test_embeddings_lr.py
new file mode 100644
index 0000000000..2b2e8e5e8d
--- /dev/null
+++ b/tests/e2e/test_embeddings_lr.py
@@ -0,0 +1,105 @@
+"""
+E2E tests for llama pretrain
+"""
+
+import unittest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_model_output_exists, check_tensorboard, with_temp_dir
+
+
+class TestEmbeddingsLrScale(unittest.TestCase):
+    """
+    Test case for embedding_lr*
+    """
+
+    @with_temp_dir
+    def test_train_w_embedding_lr_scale(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "sample_packing": True,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "max_steps": 5,
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "embedding_lr_scale": 0.5,
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.0, "Loss is too high"
+        )
+
+    @with_temp_dir
+    def test_train_w_embedding_lr(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "sample_packing": True,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "max_steps": 5,
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "embedding_lr": 0.000005,
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.0, "Loss is too high"
+        )
diff --git a/tests/e2e/test_evaluate.py b/tests/e2e/test_evaluate.py
new file mode 100644
index 0000000000..3b0ab14502
--- /dev/null
+++ b/tests/e2e/test_evaluate.py
@@ -0,0 +1,59 @@
+"""E2E smoke test for evaluate CLI command"""
+
+from pathlib import Path
+
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+
+class TestE2eEvaluate:
+    """Test cases for evaluate CLI"""
+
+    def test_evaluate(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "max_steps": 20,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "accelerate",
+                "launch",
+                "--num-processes",
+                "2",
+                "--main_process_port",
+                f"{get_torch_dist_unique_port()}",
+                "-m",
+                "axolotl.cli.evaluate",
+                str(Path(temp_dir) / "config.yaml"),
+            ]
+        )
diff --git a/tests/e2e/test_falcon.py b/tests/e2e/test_falcon.py
index c76699a7c8..19de202d2d 100644
--- a/tests/e2e/test_falcon.py
+++ b/tests/e2e/test_falcon.py
@@ -2,21 +2,18 @@
 E2E tests for falcon
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from .utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    with_temp_dir,
+)
 
 
 class TestFalcon(unittest.TestCase):
@@ -26,11 +23,10 @@ class TestFalcon(unittest.TestCase):
 
     @with_temp_dir
     def test_lora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "illuin/tiny-random-FalconForCausalLM",
-                "flash_attention": True,
+                "base_model": "axolotl-ai-co/tiny-falcon-42m",
+                "flash_attention": False,
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -42,7 +38,7 @@ def test_lora(self, temp_dir):
                     "word_embeddings",
                     "lm_head",
                 ],
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
                     "bos_token": "<|endoftext|>",
                     "pad_token": "<|endoftext|>",
@@ -54,32 +50,44 @@ def test_lora(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "warmup_steps": 5,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_lora_added_vocab(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "illuin/tiny-random-FalconForCausalLM",
-                "flash_attention": True,
+                "base_model": "axolotl-ai-co/tiny-falcon-42m",
+                "flash_attention": False,
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -91,7 +99,7 @@ def test_lora_added_vocab(self, temp_dir):
                     "word_embeddings",
                     "lm_head",
                 ],
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
                     "bos_token": "<|endoftext|>",
                     "pad_token": "<|endoftext|>",
@@ -107,34 +115,46 @@ def test_lora_added_vocab(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "warmup_steps": 5,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_ft(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "illuin/tiny-random-FalconForCausalLM",
-                "flash_attention": True,
+                "base_model": "axolotl-ai-co/tiny-falcon-42m",
+                "flash_attention": False,
                 "sequence_len": 1024,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
                     "bos_token": "<|endoftext|>",
                     "pad_token": "<|endoftext|>",
@@ -146,21 +166,36 @@ def test_ft(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "learning_rate": 5e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "logging_steps": 1,
+                "save_steps": 80,
+                "eval_steps": 80,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=10,
+            final_window=10,
+            max_initial=5.0,
+            max_final=4.7,
+        )
diff --git a/tests/e2e/test_gemma2.py b/tests/e2e/test_gemma2.py
new file mode 100644
index 0000000000..1719deeeed
--- /dev/null
+++ b/tests/e2e/test_gemma2.py
@@ -0,0 +1,121 @@
+"""
+E2E tests for gemma2
+"""
+
+from pathlib import Path
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestGemma2:
+    """
+    Test case for Gemma2 models
+    """
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_lora_gemma2(self, temp_dir, sample_packing):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/gemma-2-33M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "drop_system_message": True,
+                        "split": "train[:1%]",
+                    },
+                ],
+                "special_tokens": {
+                    "bos_token": "<bos>",
+                    "eos_token": "<eos>",
+                },
+                "chat_template": "gemma",  # gemma2's template is same as gemma
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_fft_gemma2(self, temp_dir, sample_packing):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/gemma-2-33M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "split": "train[:1%]",
+                        "drop_system_message": True,
+                    },
+                ],
+                "chat_template": "gemma",  # gemma2's template is same as gemma
+                "special_tokens": {
+                    "bos_token": "<bos>",
+                    "eos_token": "<eos>",
+                },
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "model.safetensors").exists()
diff --git a/tests/e2e/test_gemma3_text.py b/tests/e2e/test_gemma3_text.py
new file mode 100644
index 0000000000..b723f21e53
--- /dev/null
+++ b/tests/e2e/test_gemma3_text.py
@@ -0,0 +1,121 @@
+"""
+E2E tests for gemma3_text
+"""
+
+from pathlib import Path
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestGemma3Text:
+    """
+    Test case for Gemma3Text models
+    """
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_lora_gemma3_text(self, temp_dir, sample_packing):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/gemma-3-34M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "split": "train[:1%]",
+                    },
+                ],
+                "special_tokens": {
+                    "bos_token": "<bos>",
+                    "eos_token": "<eos>",
+                },
+                "chat_template": "gemma3",
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_fft_gemma3_text(self, temp_dir, sample_packing):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/gemma-3-34M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "split": "train[:1%]",
+                    },
+                ],
+                "chat_template": "gemma3",
+                "special_tokens": {
+                    "bos_token": "<bos>",
+                    "eos_token": "<eos>",
+                },
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "model.safetensors").exists()
diff --git a/tests/e2e/test_imports.py b/tests/e2e/test_imports.py
new file mode 100644
index 0000000000..4c01e50be8
--- /dev/null
+++ b/tests/e2e/test_imports.py
@@ -0,0 +1,17 @@
+"""
+test module to import various submodules that have historically broken due to dependency issues
+"""
+
+import unittest
+
+
+class TestImports(unittest.TestCase):
+    """
+    Test class to import various submodules that have historically broken due to dependency issues
+    """
+
+    def test_import_causal_trainer(self):
+        pass
+
+    def test_import_rl_trainer(self):
+        pass
diff --git a/tests/e2e/test_llama.py b/tests/e2e/test_llama.py
new file mode 100644
index 0000000000..47a7915777
--- /dev/null
+++ b/tests/e2e/test_llama.py
@@ -0,0 +1,187 @@
+"""
+E2E tests for llama
+"""
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_model_output_exists
+
+
+class TestLlama:
+    """
+    Test case for Llama models
+    """
+
+    def test_fft_trust_remote_code(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "trust_remote_code": True,
+                "sequence_len": 512,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    def test_fix_untrained_tokens(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "fix_untrained_tokens": True,
+                "sequence_len": 512,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                    "bos_token": "<|custom_im_start|>",
+                    "eos_token": "<|custom_im_end|>",
+                },
+                "datasets": [
+                    {
+                        "chat_template": "jinja",
+                        "chat_template_jinja": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|custom_im_start|>' + message['role'] + '\n' + message['content'] + '<|custom_im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|custom_im_start|>assistant\n' }}{% endif %}",
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "split": "train[:10%]",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    def test_fix_untrained_tokens_already_trained(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "fix_untrained_tokens": True,
+                "sequence_len": 512,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "split": "train[:10%]",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    @pytest.mark.parametrize("tf32", ["auto", False])
+    def test_batch_flattening(self, tf32, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "trust_remote_code": True,
+                "sequence_len": 512,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": False,
+                "batch_flattening": True,
+                "bf16": True,
+                "tf32": tf32,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/test_llama_pretrain.py b/tests/e2e/test_llama_pretrain.py
new file mode 100644
index 0000000000..3aa594fbdf
--- /dev/null
+++ b/tests/e2e/test_llama_pretrain.py
@@ -0,0 +1,72 @@
+"""E2E tests for llama pretrain"""
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_model_output_exists, check_tensorboard
+
+
+class TestPretrainLlama:
+    """Test case for Llama models w pretraining"""
+
+    @pytest.mark.parametrize(
+        ("sample_packing", "pretrain_multipack_attn"),
+        [
+            (False, False),
+            (True, True),
+            (True, False),
+        ],
+    )
+    def test_pretrain(self, temp_dir, sample_packing, pretrain_multipack_attn):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "sample_packing": sample_packing,
+                "pretrain_multipack_attn": pretrain_multipack_attn,
+                "dataset_num_proc": 1,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "pretraining_dataset": [
+                    {
+                        "path": "allenai/c4",
+                        "name": "en",
+                        "type": "pretrain",
+                    }
+                ],
+                "max_steps": 5,
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        loss_threshold = 3.6
+        if sample_packing and not pretrain_multipack_attn:
+            loss_threshold = 6.5
+        check_tensorboard(
+            temp_dir + "/runs",
+            "train/train_loss",
+            loss_threshold,
+            "Train Loss (%s) is too high",
+        )
diff --git a/tests/e2e/test_llama_vision.py b/tests/e2e/test_llama_vision.py
new file mode 100644
index 0000000000..edfc7a9b3a
--- /dev/null
+++ b/tests/e2e/test_llama_vision.py
@@ -0,0 +1,108 @@
+"""
+E2E tests for lora llama
+"""
+
+import unittest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_model_output_exists, with_temp_dir
+
+
+class TestLlamaVision(unittest.TestCase):
+    """
+    Test case for Llama Vision models
+    """
+
+    @with_temp_dir
+    def test_lora_llama_vision_text_only_dataset(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/Llama-3.2-39M-Vision",
+                "processor_type": "AutoProcessor",
+                "skip_prepare_dataset": True,
+                "remove_unused_columns": False,
+                "sample_packing": False,
+                "sequence_len": 1024,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_modules": r"model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj",
+                "val_set_size": 0,
+                "chat_template": "llama3_2_vision",
+                "datasets": [
+                    {
+                        "path": "LDJnr/Puffin",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    @with_temp_dir
+    def test_lora_llama_vision_multimodal_dataset(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/Llama-3.2-39M-Vision",
+                "processor_type": "AutoProcessor",
+                "skip_prepare_dataset": True,
+                "remove_unused_columns": False,
+                "sample_packing": False,
+                "sequence_len": 1024,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_modules": r"model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj",
+                "val_set_size": 0,
+                "chat_template": "llama3_2_vision",
+                "datasets": [
+                    {
+                        "path": "axolotl-ai-co/llava-instruct-mix-vsft-small",
+                        "type": "chat_template",
+                        "split": "train",
+                        "field_messages": "messages",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/test_load_model.py b/tests/e2e/test_load_model.py
new file mode 100644
index 0000000000..7c5389a58a
--- /dev/null
+++ b/tests/e2e/test_load_model.py
@@ -0,0 +1,90 @@
+"""Module for testing ModelLoader."""
+
+import shutil
+import tempfile
+
+import pytest
+import torch
+
+from axolotl.loaders import ModelLoader, load_tokenizer
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="temp_dir")
+def fixture_temp_dir():
+    temp_dir = tempfile.mkdtemp()
+    yield temp_dir
+    shutil.rmtree(temp_dir)
+
+
+class TestLoadModelUtils:
+    """
+    Testing module testing ModelLoader.
+    """
+
+    def setup_method(self):
+        # load config
+        self.cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "tokenizer_config": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "load_in_8bit": False,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "tensor_parallel_size": 1,
+                "context_parallel_size": 1,
+            }
+        )
+        self.model_loader = ModelLoader(
+            cfg=self.cfg,
+            tokenizer="",
+            inference=False,
+            reference_model=True,
+        )
+
+    @pytest.mark.parametrize("embedding_modules", ["embed_tokens", "lm_head"])
+    @pytest.mark.parametrize(
+        "dist_dtype", [torch.bfloat16, torch.float16, torch.float32]
+    )
+    @pytest.mark.parametrize("before_kbit_train_or_finetune", [True, False])
+    def test_convert_embedding_modules_dtype(
+        self, temp_dir, embedding_modules, dist_dtype, before_kbit_train_or_finetune
+    ):
+        self.cfg.output_dir = temp_dir
+        self.model_loader.tokenizer = load_tokenizer(self.cfg)
+        self.model_loader.load()
+        self.model_loader._convert_embedding_modules_dtype(
+            embedding_modules, dist_dtype, before_kbit_train_or_finetune
+        )
+        for name, module in self.model_loader.model.named_modules():
+            if (
+                "norm" in name
+                or (before_kbit_train_or_finetune and name.endswith(".gate"))
+                or (
+                    any(m in name for m in embedding_modules)
+                    and hasattr(module, "weight")
+                )
+            ):
+                for _, param in module.named_parameters():
+                    assert param.dtype == dist_dtype
diff --git a/tests/e2e/test_lora_llama.py b/tests/e2e/test_lora_llama.py
index c79652bef7..b6ee393dfd 100644
--- a/tests/e2e/test_lora_llama.py
+++ b/tests/e2e/test_lora_llama.py
@@ -2,21 +2,14 @@
 E2E tests for lora llama
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from .utils import check_model_output_exists, with_temp_dir
 
 
 class TestLoraLlama(unittest.TestCase):
@@ -26,23 +19,20 @@ class TestLoraLlama(unittest.TestCase):
 
     @with_temp_dir
     def test_lora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
-                "lora_r": 32,
-                "lora_alpha": 64,
+                "lora_r": 8,
+                "lora_alpha": 16,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
-                    "unk_token": "<unk>",
-                    "bos_token": "<s>",
-                    "eos_token": "</s>",
+                    "pad_token": "<|endoftext|>",
                 },
                 "datasets": [
                     {
@@ -50,18 +40,21 @@ def test_lora(self, temp_dir):
                         "type": "alpaca",
                     },
                 ],
-                "num_epochs": 2,
-                "micro_batch_size": 8,
+                "num_epochs": 1,
+                "micro_batch_size": 2,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "save_first_step": False,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/test_mamba.py b/tests/e2e/test_mamba.py
index 8755fa4d51..c45026bf55 100644
--- a/tests/e2e/test_mamba.py
+++ b/tests/e2e/test_mamba.py
@@ -2,23 +2,16 @@
 E2E tests for lora llama
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
 import pytest
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from .utils import check_model_output_exists, with_temp_dir
 
 
 @pytest.mark.skip(reason="skipping until upstreamed into transformers")
@@ -29,7 +22,6 @@ class TestMamba(unittest.TestCase):
 
     @with_temp_dir
     def test_fft(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
                 "base_model": "state-spaces/mamba-130m",
@@ -52,17 +44,18 @@ def test_fft(self, temp_dir):
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
                 "max_steps": 20,
                 "save_steps": 10,
                 "eval_steps": None,
-                "save_safetensors": False,
+                "save_first_step": False,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/test_mistral.py b/tests/e2e/test_mistral.py
index 57d85e51eb..37cae7ce77 100644
--- a/tests/e2e/test_mistral.py
+++ b/tests/e2e/test_mistral.py
@@ -2,23 +2,20 @@
 E2E tests for lora llama
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from .utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    with_temp_dir,
+)
 
 
 class TestMistral(unittest.TestCase):
@@ -28,10 +25,9 @@ class TestMistral(unittest.TestCase):
 
     @with_temp_dir
     def test_lora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model": "axolotl-ai-co/tiny-mistral-25m",
                 "flash_attention": True,
                 "sequence_len": 1024,
                 "load_in_8bit": True,
@@ -40,7 +36,7 @@ def test_lora(self, temp_dir):
                 "lora_alpha": 64,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
                     "unk_token": "<unk>",
                     "bos_token": "<s>",
@@ -53,33 +49,43 @@ def test_lora(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
+                "save_first_step": False,
+                "use_tensorboard": True,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=4.5,
+            max_final=4.3,
+        )
 
     @with_temp_dir
     def test_ft(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "openaccess-ai-collective/tiny-mistral",
+                "base_model": "axolotl-ai-co/tiny-mistral-25m",
                 "flash_attention": True,
                 "sequence_len": 1024,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
                     "unk_token": "<unk>",
                     "bos_token": "<s>",
@@ -92,24 +98,35 @@ def test_ft(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
+                "save_first_step": False,
+                "use_tensorboard": True,
             }
         )
         if is_torch_bf16_gpu_available():
             cfg.bf16 = True
         else:
             cfg.fp16 = True
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=4.5,
+            max_final=4.3,
+        )
diff --git a/tests/e2e/test_mixtral.py b/tests/e2e/test_mixtral.py
index d4dad14ef2..00c75f426f 100644
--- a/tests/e2e/test_mixtral.py
+++ b/tests/e2e/test_mixtral.py
@@ -2,24 +2,21 @@
 E2E tests for mixtral
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
 import torch
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from .utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    with_temp_dir,
+)
 
 
 class TestMixtral(unittest.TestCase):
@@ -29,11 +26,9 @@ class TestMixtral(unittest.TestCase):
 
     @with_temp_dir
     def test_qlora_w_fa2(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "hf-internal-testing/Mixtral-tiny",
-                "tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
+                "base_model": "axolotl-ai-co/tiny-mixtral-30m",
                 "flash_attention": True,
                 "sequence_len": 1024,
                 "load_in_4bit": True,
@@ -50,7 +45,7 @@ def test_qlora_w_fa2(self, temp_dir):
                     "q_proj",
                     "w2",
                 ],
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -59,35 +54,44 @@ def test_qlora_w_fa2(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
+                "save_first_step": False,
+                "use_tensorboard": True,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        model, _ = train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        model, _, _ = train(cfg=cfg, dataset_meta=dataset_meta)
         assert (
-            model.base_model.model.model.layers[0].block_sparse_moe.gate.weight.dtype
+            model.base_model.model.model.layers[0].mlp.gate.weight.dtype
             == torch.float32
         )
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_qlora_wo_fa2(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "hf-internal-testing/Mixtral-tiny",
-                "tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
+                "base_model": "axolotl-ai-co/tiny-mixtral-30m",
                 "flash_attention": False,
                 "sequence_len": 1024,
                 "load_in_4bit": True,
@@ -104,7 +108,7 @@ def test_qlora_wo_fa2(self, temp_dir):
                     "q_proj",
                     "w2",
                 ],
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -113,35 +117,44 @@ def test_qlora_wo_fa2(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
+                "save_first_step": False,
+                "use_tensorboard": True,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        model, _ = train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        model, _, _ = train(cfg=cfg, dataset_meta=dataset_meta)
         assert (
-            model.base_model.model.model.layers[0].block_sparse_moe.gate.weight.dtype
+            model.base_model.model.model.layers[0].mlp.gate.weight.dtype
             == torch.float32
         )
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_16bit_lora_w_fa2(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "hf-internal-testing/Mixtral-tiny",
-                "tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
+                "base_model": "axolotl-ai-co/tiny-mixtral-30m",
                 "flash_attention": True,
                 "sequence_len": 1024,
                 "adapter": "lora",
@@ -157,7 +170,7 @@ def test_16bit_lora_w_fa2(self, temp_dir):
                     "q_proj",
                     "w2",
                 ],
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -166,39 +179,48 @@ def test_16bit_lora_w_fa2(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
+                "save_first_step": False,
+                "use_tensorboard": True,
             }
         )
         if is_torch_bf16_gpu_available():
             cfg.bf16 = True
         else:
             cfg.fp16 = True
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        model, _ = train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        model, _, _ = train(cfg=cfg, dataset_meta=dataset_meta)
         assert (
-            model.base_model.model.model.layers[0].block_sparse_moe.gate.weight.dtype
+            model.base_model.model.model.layers[0].mlp.gate.weight.dtype
             == torch.float32
         )
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_16bit_lora_wo_fa2(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "hf-internal-testing/Mixtral-tiny",
-                "tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
+                "base_model": "axolotl-ai-co/tiny-mixtral-30m",
                 "flash_attention": False,
                 "sequence_len": 1024,
                 "adapter": "lora",
@@ -214,7 +236,7 @@ def test_16bit_lora_wo_fa2(self, temp_dir):
                     "q_proj",
                     "w2",
                 ],
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -223,42 +245,51 @@ def test_16bit_lora_wo_fa2(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
+                "save_first_step": False,
+                "use_tensorboard": True,
             }
         )
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
         if is_torch_bf16_gpu_available():
             cfg.bf16 = True
         else:
             cfg.fp16 = True
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        model, _ = train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+        model, _, _ = train(cfg=cfg, dataset_meta=dataset_meta)
         assert (
-            model.base_model.model.model.layers[0].block_sparse_moe.gate.weight.dtype
+            model.base_model.model.model.layers[0].mlp.gate.weight.dtype
             == torch.float32
         )
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_ft(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "hf-internal-testing/Mixtral-tiny",
-                "tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
+                "base_model": "axolotl-ai-co/tiny-mixtral-30m",
                 "flash_attention": True,
                 "sequence_len": 1024,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {},
                 "datasets": [
                     {
@@ -267,24 +298,35 @@ def test_ft(self, temp_dir):
                     },
                 ],
                 "num_epochs": 2,
-                "micro_batch_size": 2,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "adamw_bnb_8bit",
                 "lr_scheduler": "cosine",
-                "max_steps": 20,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
+                "save_first_step": False,
+                "use_tensorboard": True,
             }
         )
         if is_torch_bf16_gpu_available():
             cfg.bf16 = True
         else:
             cfg.fp16 = True
+
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
diff --git a/tests/e2e/test_optimizers.py b/tests/e2e/test_optimizers.py
new file mode 100644
index 0000000000..be5d4f8a57
--- /dev/null
+++ b/tests/e2e/test_optimizers.py
@@ -0,0 +1,450 @@
+"""
+E2E tests for custom optimizers using Llama
+"""
+
+import unittest
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    require_torch_2_5_1,
+    require_torch_2_6_0,
+    require_torch_2_7_0,
+    with_temp_dir,
+)
+
+
+class TestCustomOptimizers(unittest.TestCase):
+    """
+    Test case for Llama models using LoRA
+    """
+
+    @with_temp_dir
+    def test_optimi_adamw(self, temp_dir):
+        pytest.importorskip("optimi")
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForCausalLM",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "optimi_adamw",
+                "max_steps": 5,
+                "lr_scheduler": "cosine",
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        _, _, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        assert trainer.optimizer.optimizer.__class__.__name__ == "AdamW"
+
+    @with_temp_dir
+    @require_torch_2_5_1
+    def test_adopt_adamw(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForCausalLM",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adopt_adamw",
+                "lr_scheduler": "cosine",
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        _, _, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        assert "ADOPT" in trainer.optimizer.optimizer.__class__.__name__
+
+    @with_temp_dir
+    @require_torch_2_5_1
+    def test_muon(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForCausalLM",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "muon",
+                "lr_scheduler": "cosine",
+                "weight_decay": 0.01,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        _, _, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        assert "Muon" in trainer.optimizer.optimizer.__class__.__name__
+
+    @with_temp_dir
+    @require_torch_2_5_1
+    def test_sinkgd(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForCausalLM",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                        "split": "train[:200]",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "sinkgd",
+                "optim_args": {"sinkhorn_iters": 5, "sinkgd_lr_scale": 0.05},
+                "lr_scheduler": "cosine",
+                "weight_decay": 0.01,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        _, _, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        assert "SinkGD" in trainer.optimizer.optimizer.__class__.__name__
+
+    @pytest.mark.skip(
+        reason="Dion's internal torch.compile hits a dynamo cache-limit error"
+    )
+    @with_temp_dir
+    @require_torch_2_7_0
+    def test_dion(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForCausalLM",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "dion",
+                "dion_lr": 0.01,
+                "dion_momentum": 0.95,
+                "lr_scheduler": "cosine",
+                "weight_decay": 0.01,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        _, _, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        assert "Dion" in trainer.optimizer.optimizer.__class__.__name__
+
+    @with_temp_dir
+    def test_q_galore_adamw8bit(self, temp_dir):
+        pytest.importorskip("q_galore_torch")
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForCausalLM",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "q_galore_adamw8bit",
+                "bf16": True,
+                # Tiny rank/group_size so it fits SmolLM's hidden dim cleanly.
+                "qgalore_rank": 32,
+                "qgalore_update_proj_gap": 2,
+                "qgalore_proj_group_size": 64,
+                "lr_scheduler": "cosine",
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+        _, _, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        assert "AdamW8bit" in trainer.optimizer.optimizer.__class__.__name__
+
+    @with_temp_dir
+    def test_fft_schedule_free_adamw(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForCausalLM",
+                "sequence_len": 1024,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "schedule_free_adamw",
+                "lr_scheduler": "constant",
+                "max_steps": 10,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    @with_temp_dir
+    @require_torch_2_6_0
+    def test_came_pytorch(self, temp_dir):
+        pytest.importorskip("came_pytorch")
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-llama-50m",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+                "lora_target_linear": True,
+                "val_set_size": 0.1,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "sample_packing": True,
+                "pad_to_sequence_len": True,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 1e-4,
+                "optimizer": "came_pytorch",
+                "adam_beta3": 0.9999,
+                "adam_epsilon2": 1e-16,
+                "max_steps": 80,
+                "warmup_steps": 5,
+                "logging_steps": 1,
+                "lr_scheduler": "cosine",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=10,
+            final_window=10,
+            max_initial=4.0,
+            max_final=3.0,
+        )
+
+
+@require_torch_2_7_0
+@pytest.mark.parametrize(
+    "optimizer_name,expected_class,learning_rate",
+    [
+        ("flash_adamw", "FlashAdamW", 0.00001),
+        ("flash_adam", "FlashAdam", 0.00001),
+        ("flash_sgd", "FlashSGD", 0.01),
+        ("flash_sgdw", "FlashSGDW", 0.01),
+        ("flash_lion", "FlashLion", 0.0001),
+    ],
+)
+def test_flash_optimizers(tmp_path, optimizer_name, expected_class, learning_rate):
+    pytest.importorskip("flashoptim")
+    temp_dir = str(tmp_path)
+    cfg = DictDefault(
+        {
+            "base_model": "HuggingFaceTB/SmolLM2-135M",
+            "model_type": "AutoModelForCausalLM",
+            "tokenizer_type": "AutoTokenizer",
+            "sequence_len": 1024,
+            "load_in_8bit": True,
+            "adapter": "lora",
+            "lora_r": 8,
+            "lora_alpha": 16,
+            "lora_dropout": 0.05,
+            "lora_target_linear": True,
+            "val_set_size": 0.02,
+            "special_tokens": {
+                "pad_token": "<|endoftext|>",
+            },
+            "datasets": [
+                {
+                    "path": "mhenrichsen/alpaca_2k_test",
+                    "type": "alpaca",
+                },
+            ],
+            "num_epochs": 1,
+            "micro_batch_size": 8,
+            "gradient_accumulation_steps": 1,
+            "output_dir": temp_dir,
+            "learning_rate": learning_rate,
+            "optimizer": optimizer_name,
+            "max_steps": 5,
+            "lr_scheduler": "cosine",
+            "save_first_step": False,
+        }
+    )
+
+    cfg = validate_config(cfg)
+    normalize_config(cfg)
+    dataset_meta = load_datasets(cfg=cfg)
+
+    _, _, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
+    check_model_output_exists(temp_dir, cfg)
+    assert trainer.optimizer.optimizer.__class__.__name__ == expected_class
diff --git a/tests/e2e/test_packing_loss.py b/tests/e2e/test_packing_loss.py
new file mode 100644
index 0000000000..7cb979ce6e
--- /dev/null
+++ b/tests/e2e/test_packing_loss.py
@@ -0,0 +1,66 @@
+"""
+E2E tests for packed training
+"""
+
+import unittest
+
+from transformers.utils import is_torch_bf16_gpu_available
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_tensorboard, with_temp_dir
+
+
+class TestPackedLlama(unittest.TestCase):
+    """
+    Test case for Packed training of llama models
+    """
+
+    @with_temp_dir
+    def test_loss_packed(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "sample_packing": True,
+                "flash_attention": True,
+                "val_set_size": 0.0,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+        if is_torch_bf16_gpu_available():
+            cfg.bf16 = True
+        else:
+            cfg.fp16 = True
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss (%s) is too high"
+        )
diff --git a/tests/e2e/test_phi.py b/tests/e2e/test_phi.py
index 7abed85945..c2a637883c 100644
--- a/tests/e2e/test_phi.py
+++ b/tests/e2e/test_phi.py
@@ -2,26 +2,20 @@
 E2E tests for lora llama
 """
 
-import logging
-import os
 import unittest
-from pathlib import Path
 
-import pytest
-
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
 from axolotl.train import train
-from axolotl.utils.config import normalize_config
+from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
+from .utils import (
+    check_model_output_exists,
+    check_tensorboard_loss_decreased,
+    with_temp_dir,
+)
 
 
-@pytest.mark.skip(reason="doesn't seem to work on modal")
 class TestPhi(unittest.TestCase):
     """
     Test case for Phi2 models
@@ -29,17 +23,16 @@ class TestPhi(unittest.TestCase):
 
     @with_temp_dir
     def test_phi_ft(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "microsoft/phi-1_5",
+                "base_model": "axolotl-ai-co/tiny-phi-64m",
                 "model_type": "AutoModelForCausalLM",
                 "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 2048,
                 "sample_packing": False,
                 "load_in_8bit": False,
                 "adapter": None,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
                 },
@@ -52,43 +45,54 @@ def test_phi_ft(self, temp_dir):
                 "dataset_shard_num": 10,
                 "dataset_shard_idx": 0,
                 "num_epochs": 1,
-                "micro_batch_size": 1,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "paged_adamw_8bit",
+                "learning_rate": 2e-4,
+                "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
                 "flash_attention": True,
-                "max_steps": 10,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "warmup_steps": 5,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "pytorch_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
 
     @with_temp_dir
     def test_phi_qlora(self, temp_dir):
-        # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "microsoft/phi-1_5",
+                "base_model": "axolotl-ai-co/tiny-phi-64m",
                 "model_type": "AutoModelForCausalLM",
                 "tokenizer_type": "AutoTokenizer",
                 "sequence_len": 2048,
                 "sample_packing": False,
-                "load_in_8bit": False,
+                "load_in_4bit": True,
                 "adapter": "qlora",
                 "lora_r": 64,
                 "lora_alpha": 32,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "val_set_size": 0.1,
+                "val_set_size": 0.02,
                 "special_tokens": {
                     "pad_token": "<|endoftext|>",
                 },
@@ -101,22 +105,34 @@ def test_phi_qlora(self, temp_dir):
                 "dataset_shard_num": 10,
                 "dataset_shard_idx": 0,
                 "num_epochs": 1,
-                "micro_batch_size": 1,
+                "micro_batch_size": 4,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
-                "learning_rate": 0.00001,
+                "learning_rate": 2e-4,
                 "optimizer": "paged_adamw_8bit",
                 "lr_scheduler": "cosine",
                 "flash_attention": True,
-                "max_steps": 10,
-                "save_steps": 10,
-                "eval_steps": 10,
+                "max_steps": 50,
+                "warmup_steps": 5,
+                "logging_steps": 1,
+                "save_steps": 50,
+                "eval_steps": 50,
                 "bf16": "auto",
+                "save_first_step": False,
+                "use_tensorboard": True,
+                "seed": 42,
             }
         )
+        cfg = validate_config(cfg)
         normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+        dataset_meta = load_datasets(cfg=cfg)
 
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "adapter_model.bin").exists()
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+        check_tensorboard_loss_decreased(
+            temp_dir + "/runs",
+            initial_window=5,
+            final_window=5,
+            max_initial=5.0,
+            max_final=4.7,
+        )
diff --git a/tests/e2e/test_preprocess.py b/tests/e2e/test_preprocess.py
new file mode 100644
index 0000000000..895c29c872
--- /dev/null
+++ b/tests/e2e/test_preprocess.py
@@ -0,0 +1,58 @@
+"""E2E Test the preprocess cli"""
+
+from pathlib import Path
+
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+
+from axolotl.utils.dict import DictDefault
+
+AXOLOTL_ROOT = Path(__file__).parent.parent.parent
+
+
+class TestPreprocess:
+    """test cases for preprocess"""
+
+    def test_w_deepspeed(self, temp_dir):
+        """make sure preprocess doesn't choke when using deepspeed in the config"""
+
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/tiny-qwen2-129m",
+                "sequence_len": 2048,
+                "val_set_size": 0.01,
+                "datasets": [
+                    {
+                        "path": "tatsu-lab/alpaca",
+                        "type": "alpaca",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "bf16": "auto",
+                "deepspeed": str(AXOLOTL_ROOT / "deepspeed_configs/zero1.json"),
+                "dataset_prepared_path": temp_dir + "/last_run_prepared",
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "axolotl",
+                "preprocess",
+                str(Path(temp_dir) / "config.yaml"),
+            ]
+        )
+
+        assert (Path(temp_dir) / "last_run_prepared").exists()
diff --git a/tests/e2e/test_process_reward_model_smollm2.py b/tests/e2e/test_process_reward_model_smollm2.py
new file mode 100644
index 0000000000..9d83aabbc1
--- /dev/null
+++ b/tests/e2e/test_process_reward_model_smollm2.py
@@ -0,0 +1,63 @@
+"""
+E2E tests for process reward model w/ lora llama
+"""
+
+import unittest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_model_output_exists, check_tensorboard, with_temp_dir
+
+
+class TestProcessRewardSmolLM2(unittest.TestCase):
+    """
+    Test case for Llama process reward models using LoRA
+    """
+
+    @with_temp_dir
+    def test_prm(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForTokenClassification",
+                "num_labels": 2,
+                "process_reward_model": True,
+                "sequence_len": 512,
+                "val_set_size": 0.0,
+                "datasets": [
+                    {
+                        "path": "trl-lib/math_shepherd",
+                        "type": "stepwise_supervised",
+                        "step_separator": "\n",
+                        "split": "train[:10%]",
+                    },
+                ],
+                "max_steps": 100,
+                "num_epochs": 1,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0005,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "gradient_checkpointing": True,
+                "warmup_ratio": 0.1,
+                "use_tensorboard": True,
+                "special_tokens": {"pad_token": "<|endoftext|>"},
+                "seed": 42,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.7, "Train Loss (%s) is too high"
+        )
+
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/test_profiler.py b/tests/e2e/test_profiler.py
new file mode 100644
index 0000000000..ab273b9810
--- /dev/null
+++ b/tests/e2e/test_profiler.py
@@ -0,0 +1,113 @@
+"""
+e2e gpu test for the pytorch profiler callback
+"""
+
+from pathlib import Path
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="profiler_base_cfg")
+def fixture_profiler_base_cfg():
+    cfg = DictDefault(
+        base_model="HuggingFaceTB/SmolLM2-135M",
+        tokenizer_type="AutoTokenizer",
+        sequence_len=1024,
+        load_in_8bit=True,
+        adapter="lora",
+        lora_r=8,
+        lora_alpha=16,
+        lora_dropout=0.05,
+        lora_target_linear=True,
+        val_set_size=0.02,
+        special_tokens={"pad_token": "<|endoftext|>"},
+        datasets=[
+            {
+                "path": "mhenrichsen/alpaca_2k_test",
+                "type": "alpaca",
+            },
+        ],
+        num_epochs=1,
+        micro_batch_size=2,
+        gradient_accumulation_steps=1,
+        learning_rate=0.00001,
+        optimizer="adamw_torch_fused",
+        lr_scheduler="cosine",
+    )
+    return cfg
+
+
+class TestProfiler:
+    """
+    test cases for the pytorch profiler callback
+    """
+
+    def test_profiler_saves(self, profiler_base_cfg, temp_dir):
+        cfg = profiler_base_cfg | DictDefault(
+            output_dir=temp_dir,
+            max_steps=5,
+            profiler_steps=3,
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "snapshot.pickle").exists()
+
+    def test_profiler_saves_w_start(self, profiler_base_cfg, temp_dir):
+        cfg = profiler_base_cfg | DictDefault(
+            output_dir=temp_dir,
+            max_steps=5,
+            profiler_steps=3,
+            profiler_steps_start=1,
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "snapshot.pickle").exists()
+
+    @pytest.mark.parametrize(
+        "profiler_steps_start",
+        [3, 5],
+    )
+    def test_profiler_saves_past_end(
+        self, profiler_base_cfg, temp_dir, profiler_steps_start
+    ):
+        cfg = profiler_base_cfg | DictDefault(
+            output_dir=temp_dir,
+            max_steps=5,
+            profiler_steps=3,
+            profiler_steps_start=profiler_steps_start,
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "snapshot.pickle").exists()
+
+    def test_profiler_never_started(self, profiler_base_cfg, temp_dir):
+        cfg = profiler_base_cfg | DictDefault(
+            output_dir=temp_dir,
+            max_steps=5,
+            profiler_steps=3,
+            profiler_steps_start=6,
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert not (Path(temp_dir) / "snapshot.pickle").exists()
diff --git a/tests/e2e/test_qat.py b/tests/e2e/test_qat.py
new file mode 100644
index 0000000000..251d5b17b2
--- /dev/null
+++ b/tests/e2e/test_qat.py
@@ -0,0 +1,163 @@
+"""
+E2E tests for QAT
+"""
+
+from pathlib import Path
+
+from axolotl.common.datasets import load_datasets, load_preference_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.schemas.enums import TorchAOQuantDType
+from axolotl.utils.schemas.quantization import QATConfig, validate_ao_dtype
+
+from .utils import check_model_output_exists, check_tensorboard
+
+
+class TestQATLlama:
+    """
+    Test case for QAT Llama models
+    """
+
+    def test_qat(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "drop_system_message": True,
+                        "split": "train[:1%]",
+                    },
+                ],
+                "chat_template": "chatml",
+                "qat": {
+                    "quantize_embedding": True,
+                    "activation_dtype": "int8",
+                    "weight_dtype": "int4",
+                    "group_size": 8,
+                },
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-5", cfg)
+
+    def test_qat_dpo(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 2048,
+                "sample_packing": False,
+                "eval_sample_packing": False,
+                "pad_to_sequence_len": True,
+                "val_set_size": 0.01,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "rl": "dpo",
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "fozziethebeat/alpaca_messages_2k_dpo_test",
+                        "type": "chat_template.default",
+                        "field_messages": "conversation",
+                        "field_chosen": "chosen",
+                        "field_rejected": "rejected",
+                        "message_field_role": "role",
+                        "message_field_content": "content",
+                        "roles": {
+                            "system": ["system"],
+                            "user": ["user"],
+                            "assistant": ["assistant"],
+                        },
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "warmup_steps": 0,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "use_tensorboard": True,
+                "bf16": True,
+                "qat": {
+                    "quantize_embedding": True,
+                    "activation_dtype": "int8",
+                    "weight_dtype": "int4",
+                    "group_size": 8,
+                },
+                "save_first_step": False,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_preference_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(Path(temp_dir) / "checkpoint-5", cfg)
+
+        loss_threshold = 2.3
+        check_tensorboard(
+            temp_dir + "/runs",
+            "train/train_loss",
+            loss_threshold,
+            "Train Loss (%s) is too high",
+        )
+
+
+class TestMXFP4Schema:
+    """Test MXFP4 schema validation"""
+
+    def test_validate_mxfp4_dtype(self):
+        result = validate_ao_dtype("mxfp4")
+        assert result == TorchAOQuantDType.mxfp4
+
+    def test_qat_config_with_mxfp4(self):
+        """Test QATConfig accepts mxfp4 weight_dtype"""
+        config = QATConfig(
+            weight_dtype="mxfp4",
+            group_size=32,
+            quantize_embedding=False,
+        )
+        assert config.weight_dtype == TorchAOQuantDType.mxfp4
+        assert config.group_size == 32
+
+    def test_qat_config_mxfp4_invalid_group_size(self):
+        """Test that invalid group_size raises appropriate error during quantization"""
+        # Note: Schema validation doesn't check group_size compatibility,
+        # that happens in get_quantization_config
+        config = QATConfig(
+            weight_dtype="mxfp4",
+            group_size=16,  # Invalid for mxfp4, but schema allows it
+        )
+        assert config.group_size == 16  # Schema accepts it
+        # Actual validation happens at runtime in get_quantization_config
diff --git a/tests/e2e/test_quantization.py b/tests/e2e/test_quantization.py
new file mode 100644
index 0000000000..8cdf04ea05
--- /dev/null
+++ b/tests/e2e/test_quantization.py
@@ -0,0 +1,643 @@
+"""
+Tests for axolotl.utils.quantization
+"""
+
+import pytest
+import torch
+from torch import nn
+from torchao.quantization import IntxUnpackedToInt8Tensor
+from torchao.quantization.qat.embedding import FakeQuantizedEmbedding
+from torchao.quantization.qat.linear import FakeQuantizedLinear
+from torchao.quantization.quant_api import (
+    Float8DynamicActivationFloat8WeightConfig,
+    Float8DynamicActivationInt4WeightConfig,
+    Int8DynamicActivationIntxWeightConfig,
+)
+from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor
+from transformers import AutoModelForCausalLM
+from transformers.trainer_callback import TrainerState
+
+from axolotl.utils.callbacks.qat import QATCallback
+from axolotl.utils.quantization import (
+    convert_qat_model,
+    get_quantization_config,
+    prepare_model_for_qat,
+    quantize_model,
+    save_quantized_model,
+)
+from axolotl.utils.schemas.enums import TorchAOQuantDType
+from axolotl.utils.schemas.quantization import QATConfig
+
+from tests.e2e.utils import (
+    require_torch_2_8_0,
+    requires_cuda_ge_8_9,
+    requires_sm_ge_100,
+)
+
+
+def _get_fake_quant_config_dtype(config):
+    """Get the weight dtype from a fake quantize config, handling different config types."""
+    if hasattr(config, "dtype"):
+        return config.dtype
+    # Int4WeightFakeQuantizeConfig doesn't have .dtype — weight is always int4
+    return torch.int4
+
+
+@pytest.fixture()
+def model():
+    dummy_model = AutoModelForCausalLM.from_pretrained(
+        "axolotl-ai-co/tiny-qwen2-129m",
+        device_map="auto",
+        dtype=torch.bfloat16,
+    )
+    with torch.device(dummy_model.device):
+        dummy_model.model.embed_tokens = torch.nn.Embedding(
+            dummy_model.model.embed_tokens.weight.shape[0],
+            dummy_model.model.embed_tokens.weight.shape[1],
+            dtype=dummy_model.model.embed_tokens.weight.dtype,
+        )
+    yield dummy_model
+    del dummy_model
+
+
+ptq_config_test_cases = [
+    # weight_dtype, activation_dtype, group_size, expected_type
+    (
+        TorchAOQuantDType.int4,
+        TorchAOQuantDType.int8,
+        None,
+        Int8DynamicActivationIntxWeightConfig,
+    ),
+    (
+        TorchAOQuantDType.float8_e4m3fn,
+        TorchAOQuantDType.float8_e4m3fn,
+        None,
+        Float8DynamicActivationFloat8WeightConfig,
+    ),
+    (
+        TorchAOQuantDType.int4,
+        TorchAOQuantDType.float8_e4m3fn,
+        None,
+        Float8DynamicActivationInt4WeightConfig,
+    ),
+]
+
+ptq_test_cases = [
+    # weight_dtype, activation_dtype, group_size, quantize_embedding, expected_exception, expected_tensor_class
+    (TorchAOQuantDType.int4, None, 4, True, None, Int4Tensor),
+    (
+        TorchAOQuantDType.int4,
+        TorchAOQuantDType.int8,
+        8,
+        False,
+        None,
+        IntxUnpackedToInt8Tensor,
+    ),
+    # (
+    #     TorchAOQuantDType.int4,
+    #     TorchAOQuantDType.float8_e4m3fn,
+    #     None,
+    #     False,
+    #     None,
+    #     Int4Tensor,
+    # ),
+    (TorchAOQuantDType.int4, None, None, False, None, Int4Tensor),
+    # Deprecated configs
+    (TorchAOQuantDType.int8, None, 8, False, ValueError, None),
+    (TorchAOQuantDType.int4, TorchAOQuantDType.int4, 8, False, ValueError, None),
+    (TorchAOQuantDType.int8, TorchAOQuantDType.int8, 8, True, ValueError, None),
+]
+
+
+class TestQuantization:
+    """
+    Test quantization utilities
+    """
+
+    @pytest.mark.parametrize(
+        "weight_dtype,activation_dtype,group_size,expected_type",
+        ptq_config_test_cases,
+    )
+    @requires_cuda_ge_8_9
+    @require_torch_2_8_0
+    def test_get_ptq_config(
+        self, weight_dtype, activation_dtype, group_size, expected_type
+    ):
+        config = get_quantization_config(weight_dtype, activation_dtype, group_size)
+        assert isinstance(config, expected_type)
+
+    @require_torch_2_8_0
+    @requires_sm_ge_100
+    def test_get_ptq_config_mxfp4(self):
+        from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
+
+        config = get_quantization_config(TorchAOQuantDType.mxfp4, None, 32)
+        assert isinstance(config, MXDynamicActivationMXWeightConfig)
+        assert config.weight_dtype == torch.float4_e2m1fn_x2
+        assert config.block_size == 32
+
+    @require_torch_2_8_0
+    @requires_sm_ge_100
+    def test_get_ptq_config_mxfp4_invalid_group_size(self):
+        with pytest.raises(
+            ValueError, match="MXFP4 quantization must use a block_size"
+        ):
+            get_quantization_config(TorchAOQuantDType.mxfp4, None, 16)
+
+    @requires_cuda_ge_8_9
+    @require_torch_2_8_0
+    def test_get_ptq_config_int4_weight_only(self):
+        from torchao.quantization.quant_api import Int4WeightOnlyConfig
+
+        config = get_quantization_config(TorchAOQuantDType.int4, None, 4)
+        assert isinstance(config, Int4WeightOnlyConfig)
+
+    @pytest.mark.parametrize(
+        "weight_dtype,activation_dtype,group_size,quantize_embedding,expected_exception,expected_tensor_class",
+        ptq_test_cases,
+    )
+    @requires_cuda_ge_8_9
+    @require_torch_2_8_0
+    def test_quantize_model_for_ptq(
+        self,
+        model,
+        weight_dtype,
+        activation_dtype,
+        group_size,
+        quantize_embedding,
+        expected_exception,
+        expected_tensor_class,
+    ):
+        # TODO: add mslk-cuda as a CI dependency once pytorch 2.10.x is available
+        # (see https://pypi.org/project/mslk-cuda/)
+        if expected_tensor_class is Int4Tensor and activation_dtype is None:
+            try:
+                from torchao.quantization.quantize_.workflows.int4.int4_tensor import (
+                    int4_row_quantize_zp,
+                )
+
+                if int4_row_quantize_zp is None:
+                    pytest.skip("Int4Tensor requires mslk >= 1.0.0")
+            except ImportError:
+                pytest.skip("Int4Tensor requires mslk >= 1.0.0")
+        if expected_exception:
+            with pytest.raises(expected_exception):
+                quantize_model(
+                    model,
+                    weight_dtype,
+                    group_size,
+                    activation_dtype,
+                    quantize_embedding,
+                )
+        else:
+            quantize_model(
+                model, weight_dtype, group_size, activation_dtype, quantize_embedding
+            )
+            if quantize_embedding:
+                assert isinstance(
+                    model.model.embed_tokens.weight, expected_tensor_class
+                ), "Embedding weight should be quantized"
+            for child in list(model.children()):
+                if isinstance(child, torch.nn.Linear):
+                    assert isinstance(child.weight, expected_tensor_class)
+
+    @require_torch_2_8_0
+    @requires_sm_ge_100
+    def test_quantize_model_for_ptq_fp8(
+        self,
+        model,
+    ):
+        from torchao.quantization.quantize_.workflows.float8.float8_tensor import (
+            Float8Tensor,
+            QuantizeTensorToFloat8Kwargs,
+        )
+
+        quantize_model(
+            model,
+            TorchAOQuantDType.float8_e4m3fn,
+            None,
+            TorchAOQuantDType.float8_e4m3fn,
+        )
+        for child in list(model.children()):
+            if isinstance(child, torch.nn.Linear):
+                assert isinstance(child.weight, Float8Tensor)
+                assert child.weight.act_quant_kwargs is not None and isinstance(
+                    child.weight.act_quant_kwargs, QuantizeTensorToFloat8Kwargs
+                )
+
+    @require_torch_2_8_0
+    @requires_sm_ge_100
+    def test_quantize_model_for_ptq_nvfp4(
+        self,
+        model,
+    ):
+        from torchao.prototype.mx_formats.nvfp4_tensor import (
+            NVFP4Tensor,
+            QuantizeTensorToNVFP4Kwargs,
+        )
+
+        quantize_model(model, TorchAOQuantDType.nvfp4, 16, TorchAOQuantDType.nvfp4)
+        for child in list(model.children()):
+            if isinstance(child, torch.nn.Linear):
+                assert isinstance(child.weight, NVFP4Tensor)
+                assert child.weight.act_quant_kwargs is not None and isinstance(
+                    child.weight.act_quant_kwargs, QuantizeTensorToNVFP4Kwargs
+                )
+
+    @pytest.mark.parametrize(
+        "weight_dtype,activation_dtype,group_size,quantize_embedding",
+        [
+            (TorchAOQuantDType.int4, None, 8, False),
+            (TorchAOQuantDType.int4, None, 16, True),
+            (TorchAOQuantDType.int4, TorchAOQuantDType.int8, 8, False),
+            (TorchAOQuantDType.int4, TorchAOQuantDType.int8, 16, True),
+            (
+                TorchAOQuantDType.float8_e4m3fn,
+                TorchAOQuantDType.float8_e4m3fn,
+                None,
+                False,
+            ),
+            (TorchAOQuantDType.int4, TorchAOQuantDType.float8_e4m3fn, None, True),
+        ],
+    )
+    @require_torch_2_8_0
+    @requires_cuda_ge_8_9
+    def test_prepare_model_for_qat(
+        self, model, weight_dtype, activation_dtype, group_size, quantize_embedding
+    ):
+        prepare_model_for_qat(
+            model,
+            weight_dtype,
+            group_size,
+            activation_dtype,
+            quantize_embedding,
+        )
+        if quantize_embedding:
+            assert isinstance(model.model.embed_tokens, FakeQuantizedEmbedding)
+            assert hasattr(model.model.embed_tokens, "weight_fake_quantizer")
+            embed_config = model.model.embed_tokens.weight_fake_quantizer.config
+            assert _get_fake_quant_config_dtype(embed_config) == weight_dtype.value
+            if group_size:
+                assert embed_config.group_size == group_size
+
+        for child in list(model.children()):
+            if isinstance(child, torch.nn.Linear):
+                assert isinstance(child, FakeQuantizedLinear)
+                assert hasattr(child, "weight_fake_quantizer")
+                w_config = child.weight_fake_quantizer.config
+                assert _get_fake_quant_config_dtype(w_config) == weight_dtype.value
+                if group_size:
+                    assert w_config.group_size == group_size
+                if activation_dtype:
+                    assert hasattr(child, "activation_fake_quantizer")
+                    a_config = child.activation_fake_quantizer.config
+                    assert (
+                        _get_fake_quant_config_dtype(a_config) == activation_dtype.value
+                    )
+                else:
+                    assert child.activation_fake_quantizer is None
+
+    @pytest.mark.parametrize(
+        "weight_dtype,activation_dtype,group_size,quantize_embedding",
+        [
+            (TorchAOQuantDType.mxfp4, None, 32, False),
+        ],
+    )
+    @require_torch_2_8_0
+    @requires_sm_ge_100
+    def test_prepare_model_for_qat_mxfp4(
+        self, model, weight_dtype, activation_dtype, group_size, quantize_embedding
+    ):
+        prepare_model_for_qat(
+            model,
+            weight_dtype,
+            group_size,
+            activation_dtype,
+            quantize_embedding,
+        )
+
+        from torchao.prototype.qat import MXFakeQuantizedLinear
+
+        if quantize_embedding:
+            assert isinstance(model.model.embed_tokens, FakeQuantizedEmbedding)
+            assert hasattr(model.model.embed_tokens, "weight_fake_quantizer")
+
+        for child in list(model.children()):
+            if isinstance(child, torch.nn.Linear):
+                assert isinstance(child, MXFakeQuantizedLinear)
+                assert hasattr(child, "weight_config")
+
+    @require_torch_2_8_0
+    @requires_cuda_ge_8_9
+    def test_convert_qat_model(self, model):
+        config = QATConfig(
+            weight_dtype="int4",
+            activation_dtype="int8",
+            group_size=8,
+            quantize_embedding=True,
+        )
+
+        # quantize model for qat
+        prepare_model_for_qat(
+            model,
+            config.weight_dtype,
+            config.group_size,
+            config.activation_dtype,
+            config.quantize_embedding,
+        )
+
+        assert isinstance(model.model.embed_tokens, FakeQuantizedEmbedding)
+        assert isinstance(model.lm_head, FakeQuantizedLinear)
+
+        # apply conversion
+        convert_qat_model(
+            model,
+            config.quantize_embedding,
+        )
+        # ensure modules have been swapped out
+        assert not isinstance(model.model.embed_tokens, FakeQuantizedEmbedding)
+        assert not isinstance(model.lm_head, FakeQuantizedLinear)
+
+        # ensure weights have been quantized
+        assert isinstance(model.model.embed_tokens.weight, nn.Parameter)
+        assert isinstance(model.lm_head.weight, nn.Parameter)
+
+
+class TestMXQuantizeSaveLoad:
+    """Tests for MX format (mxfp4) quantize-save-load round-trip via save_pretrained.
+
+    Uses a tiny HF model built from config (no download) so tests exercise the
+    real save_pretrained / from_pretrained code path — the same one the CLI uses.
+    MX format models are saved with safe_serialization=False (torch.save) because
+    MXTensor does not yet support safetensors serialization.
+    """
+
+    @staticmethod
+    def _make_tiny_model():
+        """Build a minimal HF causal-LM that can be quantized on CPU."""
+        from transformers import Qwen2Config
+
+        config = Qwen2Config(
+            hidden_size=64,
+            intermediate_size=128,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            vocab_size=256,
+            max_position_embeddings=64,
+            torch_dtype="bfloat16",
+        )
+        model = AutoModelForCausalLM.from_config(config).to(torch.bfloat16)
+        return model
+
+    @require_torch_2_8_0
+    def test_mxfp4_quantize_save_pretrained(self, tmp_path):
+        """quantize_model(mxfp4) -> save_pretrained -> from_pretrained round-trip."""
+        from torchao.prototype.mx_formats.mx_tensor import MXTensor
+
+        model = self._make_tiny_model()
+        original_keys = set(model.state_dict().keys())
+
+        quantize_model(model, TorchAOQuantDType.mxfp4, 32)
+
+        # Weights should be MXTensor after quantization
+        for module in model.modules():
+            if isinstance(module, nn.Linear):
+                assert isinstance(module.weight, MXTensor)
+
+        # Model should be flagged for MX-style save
+        assert getattr(model, "_is_mx_quantized", False)
+
+        # save_pretrained with safe_serialization=False (torch.save path)
+        save_dir = str(tmp_path / "mxfp4_model")
+        save_quantized_model(model, save_dir)
+
+        # Verify checkpoint files were written
+        import glob
+
+        assert glob.glob(f"{save_dir}/*.bin") or glob.glob(f"{save_dir}/**/*.bin")
+
+        # from_pretrained should load without error
+        loaded = AutoModelForCausalLM.from_pretrained(
+            save_dir, torch_dtype=torch.bfloat16
+        )
+        loaded_keys = set(loaded.state_dict().keys())
+        assert original_keys == loaded_keys, (
+            f"Key mismatch: missing={original_keys - loaded_keys}, "
+            f"extra={loaded_keys - original_keys}"
+        )
+
+    @require_torch_2_8_0
+    def test_mxfp4_is_mx_flag_set(self):
+        """quantize_model sets _is_mx_quantized for MX configs."""
+        model = self._make_tiny_model()
+        quantize_model(model, TorchAOQuantDType.mxfp4, 32)
+        assert getattr(model, "_is_mx_quantized", False)
+
+    @require_torch_2_8_0
+    @requires_cuda_ge_8_9
+    def test_non_mx_uses_torchao_quantizer(self):
+        """Non-MX quantization attaches TorchAoHfQuantizer, not _is_mx_quantized."""
+        model = self._make_tiny_model()
+        try:
+            quantize_model(model, TorchAOQuantDType.int4, group_size=32)
+        except ImportError:
+            pytest.skip("int4 quantization requires mslk >= 1.0.0")
+        assert not getattr(model, "_is_mx_quantized", False)
+        assert hasattr(model, "hf_quantizer")
+
+    @require_torch_2_8_0
+    def test_mxfp4_qat_then_ptq_save_pretrained(self, tmp_path):
+        """Full QAT -> convert -> PTQ -> save_pretrained -> from_pretrained."""
+        from torchao.prototype.mx_formats.mx_tensor import MXTensor
+        from torchao.prototype.qat import MXFakeQuantizedLinear
+
+        model = self._make_tiny_model()
+        original_keys = set(model.state_dict().keys())
+
+        # QAT preparation
+        prepare_model_for_qat(model, TorchAOQuantDType.mxfp4, 32)
+        for module in model.modules():
+            if isinstance(module, nn.Linear):
+                assert isinstance(module, MXFakeQuantizedLinear)
+
+        # Convert QAT back to normal linear
+        convert_qat_model(model)
+
+        # PTQ quantize
+        quantize_model(model, TorchAOQuantDType.mxfp4, 32)
+        for module in model.modules():
+            if isinstance(module, nn.Linear):
+                assert isinstance(module.weight, MXTensor)
+
+        # save_pretrained round-trip
+        save_dir = str(tmp_path / "mxfp4_qat_model")
+        save_quantized_model(model, save_dir)
+
+        loaded = AutoModelForCausalLM.from_pretrained(
+            save_dir, torch_dtype=torch.bfloat16
+        )
+        loaded_keys = set(loaded.state_dict().keys())
+        assert original_keys == loaded_keys
+
+    @require_torch_2_8_0
+    def test_mxfp4_cross_process_load(self, tmp_path):
+        """A saved MX checkpoint loads in a fresh interpreter that never quantized.
+
+        The other tests call ``quantize_model`` in-process, which installs the
+        transformers init guard as a side effect, masking the real reload path.
+        Here we save, then load in a subprocess that only installs the guard the
+        way ``ModelLoader._apply_pre_model_load_setup`` does — no ``quantize_model``.
+        """
+        import subprocess
+        import sys
+        import textwrap
+
+        model = self._make_tiny_model()
+        save_dir = str(tmp_path / "mxfp4_xproc_model")
+        quantize_model(model, TorchAOQuantDType.mxfp4, 32)
+        save_quantized_model(model, save_dir)
+
+        script = textwrap.dedent(
+            f"""
+            import torch
+            from torch import nn
+            from transformers import AutoModelForCausalLM
+            # mirrors ModelLoader._apply_pre_model_load_setup (no quantize_model here)
+            from axolotl.utils.quantization import (
+                patch_transformers_skip_quantized_init,
+            )
+
+            patch_transformers_skip_quantized_init()
+            model = AutoModelForCausalLM.from_pretrained(
+                {save_dir!r}, dtype=torch.bfloat16
+            )
+            from torchao.prototype.mx_formats.mx_tensor import MXTensor
+            assert any(
+                isinstance(m.weight, MXTensor)
+                for m in model.modules()
+                if isinstance(m, nn.Linear)
+            ), "expected MXTensor weights after reload"
+            print("CROSS_PROCESS_LOAD_OK")
+            """
+        )
+        result = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        assert "CROSS_PROCESS_LOAD_OK" in result.stdout, (
+            f"cross-process MX load failed:\nstdout={result.stdout}\n"
+            f"stderr={result.stderr[-2000:]}"
+        )
+
+
+class TestQuantizationCallback:
+    """
+    Test QATCallback
+    """
+
+    @pytest.fixture()
+    def trainer_state(self):
+        return TrainerState(
+            global_step=0,
+        )
+
+    @require_torch_2_8_0
+    def test_qat_callback_fake_quant_after_n_steps(self, model, trainer_state):
+        cfg = QATConfig(
+            weight_dtype="int4",
+            activation_dtype="int8",
+            group_size=8,
+            quantize_embedding=True,
+            fake_quant_after_n_steps=100,
+        )
+
+        prepare_model_for_qat(
+            model,
+            cfg.weight_dtype,
+            cfg.group_size,
+            cfg.activation_dtype,
+            cfg.quantize_embedding,
+        )
+
+        # ensure model has been quantized
+        assert isinstance(model.model.embed_tokens, FakeQuantizedEmbedding)
+        assert isinstance(model.lm_head, FakeQuantizedLinear)
+
+        # Only test enable/disable toggling if the fake quantizer supports it
+        # (Int4WeightFakeQuantizer does not have an 'enabled' attribute)
+        supports_toggle = hasattr(
+            model.model.embed_tokens.weight_fake_quantizer, "enabled"
+        )
+        if supports_toggle:
+            assert model.model.embed_tokens.weight_fake_quantizer.enabled
+            assert model.lm_head.weight_fake_quantizer.enabled
+
+        qat_callback = QATCallback(cfg)
+
+        # simulate first training step
+        qat_callback.on_step_begin(
+            args=None,
+            state=trainer_state,
+            control=None,
+            model=model,
+        )
+
+        if supports_toggle:
+            # quantization should have been disabled
+            assert not model.model.embed_tokens.weight_fake_quantizer.enabled
+            assert not model.lm_head.weight_fake_quantizer.enabled
+
+        trainer_state.global_step = 100
+        qat_callback.on_step_begin(
+            args=None,
+            state=trainer_state,
+            control=None,
+            model=model,
+        )
+
+        if supports_toggle:
+            # quantization should have been enabled
+            assert model.model.embed_tokens.weight_fake_quantizer.enabled
+            assert model.lm_head.weight_fake_quantizer.enabled
+
+    @require_torch_2_8_0
+    def test_qat_callback_fake_quant_after_n_steps_is_none(self, model, trainer_state):
+        cfg = QATConfig(
+            weight_dtype="int4",
+            activation_dtype="int8",
+            group_size=8,
+            quantize_embedding=True,
+            fake_quant_after_n_steps=None,
+        )
+
+        prepare_model_for_qat(
+            model,
+            cfg.weight_dtype,
+            cfg.group_size,
+            cfg.activation_dtype,
+            cfg.quantize_embedding,
+        )
+
+        # ensure model has been quantized
+        assert isinstance(model.model.embed_tokens, FakeQuantizedEmbedding)
+        assert isinstance(model.lm_head, FakeQuantizedLinear)
+        if hasattr(model.model.embed_tokens.weight_fake_quantizer, "enabled"):
+            assert model.model.embed_tokens.weight_fake_quantizer.enabled
+            assert model.lm_head.weight_fake_quantizer.enabled
+
+        qat_callback = QATCallback(cfg)
+        # simulate first training step
+        qat_callback.on_step_begin(
+            args=None,
+            state=trainer_state,
+            control=None,
+            model=model,
+        )
+
+        # quantization should be enabled from the get-go
+        if hasattr(model.model.embed_tokens.weight_fake_quantizer, "enabled"):
+            assert model.model.embed_tokens.weight_fake_quantizer.enabled
+            assert model.lm_head.weight_fake_quantizer.enabled
diff --git a/tests/e2e/test_qwen.py b/tests/e2e/test_qwen.py
new file mode 100644
index 0000000000..b8654c0ad6
--- /dev/null
+++ b/tests/e2e/test_qwen.py
@@ -0,0 +1,82 @@
+"""
+E2E tests for qwen
+"""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from accelerate.test_utils import execute_subprocess_async
+from transformers.testing_utils import get_torch_dist_unique_port
+
+from axolotl.utils.dict import DictDefault
+
+
+class TestE2eQwen:
+    """
+    Test cases for qwen models
+    """
+
+    @pytest.mark.parametrize("base_model", ["axolotl-ai-co/tiny-qwen2-129m"])
+    def test_dpo(self, base_model, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": base_model,
+                "rl": "dpo",
+                "chat_template": "qwen_25",
+                "sequence_len": 2048,
+                "val_set_size": 0.0,
+                "datasets": [
+                    {
+                        "path": "fozziethebeat/alpaca_messages_2k_dpo_test",
+                        "split": "train",
+                        "type": "chat_template.default",
+                        "field_messages": "conversation",
+                        "field_chosen": "chosen",
+                        "field_rejected": "rejected",
+                        "message_property_mappings": {
+                            "role": "role",
+                            "content": "content",
+                        },
+                        "roles": {
+                            "system": ["system"],
+                            "user": ["user"],
+                            "assistant": ["assistant"],
+                        },
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "warmup_steps": 20,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "bf16": "auto",
+                "tf32": True,
+                "gradient_checkpointing": True,
+                "save_first_step": False,
+            }
+        )
+
+        # write cfg to yaml file
+        Path(temp_dir).mkdir(parents=True, exist_ok=True)
+        with open(Path(temp_dir) / "config.yaml", "w", encoding="utf-8") as fout:
+            fout.write(yaml.dump(cfg.to_dict(), Dumper=yaml.Dumper))
+
+        execute_subprocess_async(
+            [
+                "accelerate",
+                "launch",
+                "--num-processes",
+                "2",
+                "--main_process_port",
+                f"{get_torch_dist_unique_port()}",
+                "-m",
+                "axolotl.cli.train",
+                str(Path(temp_dir) / "config.yaml"),
+            ]
+        )
diff --git a/tests/e2e/test_relora_llama.py b/tests/e2e/test_relora_llama.py
deleted file mode 100644
index 4ba130c9dc..0000000000
--- a/tests/e2e/test_relora_llama.py
+++ /dev/null
@@ -1,68 +0,0 @@
-"""
-E2E tests for relora llama
-"""
-
-import logging
-import os
-import unittest
-from pathlib import Path
-
-from axolotl.cli import load_datasets
-from axolotl.common.cli import TrainerCliArgs
-from axolotl.train import train
-from axolotl.utils.config import normalize_config
-from axolotl.utils.dict import DictDefault
-
-from .utils import with_temp_dir
-
-LOG = logging.getLogger("axolotl.tests.e2e")
-os.environ["WANDB_DISABLED"] = "true"
-
-
-class TestReLoraLlama(unittest.TestCase):
-    """
-    Test case for Llama models using LoRA
-    """
-
-    @with_temp_dir
-    def test_relora(self, temp_dir):
-        # pylint: disable=duplicate-code
-        cfg = DictDefault(
-            {
-                "base_model": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
-                "sequence_len": 1024,
-                "load_in_8bit": True,
-                "adapter": "lora",
-                "lora_r": 32,
-                "lora_alpha": 16,
-                "lora_dropout": 0.05,
-                "lora_target_modules": ["q_proj", "v_proj"],
-                "relora_steps": 25,
-                "relora_warmup_steps": 5,
-                "relora_anneal_steps": 5,
-                "relora_cpu_offload": True,
-                "val_set_size": 0.0,
-                "special_tokens": {},
-                "datasets": [
-                    {
-                        "path": "mhenrichsen/alpaca_2k_test",
-                        "type": "alpaca",
-                    },
-                ],
-                "warmup_steps": 15,
-                "num_epochs": 2,
-                "micro_batch_size": 4,
-                "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
-                "lr_scheduler": "cosine",
-            }
-        )
-        normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
-
-        train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        assert (Path(temp_dir) / "model.safetensors").exists()
diff --git a/tests/e2e/test_save_first_step.py b/tests/e2e/test_save_first_step.py
new file mode 100644
index 0000000000..c717edb6a4
--- /dev/null
+++ b/tests/e2e/test_save_first_step.py
@@ -0,0 +1,98 @@
+"""
+E2E tests for relora llama
+"""
+
+import unittest
+from pathlib import Path
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_model_output_exists, with_temp_dir
+
+
+class TestSaveFirstStepCallback(unittest.TestCase):
+    """Test cases for save_first_step callback config."""
+
+    @with_temp_dir
+    def test_save_first_step(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 512,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": True,
+                "save_first_step": True,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(str(Path(temp_dir) / "checkpoint-1"), cfg)
+
+    @with_temp_dir
+    def test_no_save_first_step(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 512,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,
+                "micro_batch_size": 2,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "flash_attention": True,
+                "sample_packing": True,
+                "bf16": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        with pytest.raises(AssertionError):
+            check_model_output_exists(str(Path(temp_dir) / "checkpoint-1"), cfg)
diff --git a/tests/e2e/test_schedulers.py b/tests/e2e/test_schedulers.py
new file mode 100644
index 0000000000..5b9c56288f
--- /dev/null
+++ b/tests/e2e/test_schedulers.py
@@ -0,0 +1,62 @@
+"""
+E2E tests for custom schedulers using Llama
+"""
+
+import unittest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_model_output_exists, with_temp_dir
+
+
+class TestCustomSchedulers(unittest.TestCase):
+    """
+    Test case for Llama models using LoRA
+    """
+
+    @with_temp_dir
+    def test_rex_scheduler(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "load_in_8bit": True,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.02,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "max_steps": 20,
+                "lr_scheduler": "rex",
+                "warmup_steps": 5,
+                "cosine_min_lr_ratio": 0.05,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
diff --git a/tests/e2e/test_streaming.py b/tests/e2e/test_streaming.py
new file mode 100644
index 0000000000..eb085f3dad
--- /dev/null
+++ b/tests/e2e/test_streaming.py
@@ -0,0 +1,72 @@
+"""E2E tests for streaming dataset functionality"""
+
+# pylint: disable=duplicate-code
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from .utils import check_model_output_exists, check_tensorboard
+
+
+class TestStreamingDatasets:
+    """Test case for streaming datasets"""
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_streaming_dataset(self, temp_dir, sample_packing):
+        """Test streaming datasets"""
+
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "flash_attention": True,
+                "sequence_len": 1024,
+                "sample_packing": sample_packing,
+                "pretrain_multipack_attn": sample_packing,
+                "streaming_multipack_buffer_size": 10000,
+                "dataset_num_proc": 1,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                # Streaming config
+                "streaming": True,
+                "max_steps": 10,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "val_set_size": 0.0,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "bf16": "auto",
+                "use_tensorboard": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        # Verify training actually happened by checking loss decrease
+        check_tensorboard(
+            temp_dir + "/runs",
+            "train/train_loss",
+            3.0,
+            "Train Loss (%s) is too high",
+        )
diff --git a/tests/e2e/test_tokenizer.py b/tests/e2e/test_tokenizer.py
new file mode 100644
index 0000000000..a65c17ac31
--- /dev/null
+++ b/tests/e2e/test_tokenizer.py
@@ -0,0 +1,63 @@
+"""
+e2e test for saving the tokenizer
+"""
+
+from unittest.mock import patch
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_model_output_exists
+
+
+def test_tokenizer_no_save_jinja_files(temp_dir):
+    # pylint: disable=duplicate-code
+    cfg = DictDefault(
+        {
+            "base_model": "HuggingFaceTB/SmolLM2-135M",
+            "tokenizer_type": "AutoTokenizer",
+            "sequence_len": 1024,
+            "load_in_8bit": True,
+            "adapter": "lora",
+            "lora_r": 8,
+            "lora_alpha": 16,
+            "lora_dropout": 0.05,
+            "lora_target_linear": True,
+            "val_set_size": 0.02,
+            "special_tokens": {
+                "pad_token": "<|endoftext|>",
+            },
+            "chat_template": "chatml",
+            "datasets": [
+                {
+                    "path": "mhenrichsen/alpaca_2k_test",
+                    "type": "alpaca",
+                },
+            ],
+            "num_epochs": 1,
+            "micro_batch_size": 2,
+            "gradient_accumulation_steps": 1,
+            "output_dir": temp_dir,
+            "learning_rate": 0.00001,
+            "optimizer": "adamw_torch_fused",
+            "lr_scheduler": "cosine",
+            "max_steps": 5,
+            "save_first_step": False,
+            "fp16": False,
+            "tokenizer_save_jinja_files": False,
+        }
+    )
+
+    cfg = validate_config(cfg)
+    normalize_config(cfg)
+    dataset_meta = load_datasets(cfg=cfg)
+
+    with patch("axolotl.train.execute_training"):
+        train(cfg=cfg, dataset_meta=dataset_meta)
+
+    check_model_output_exists(temp_dir, cfg)
+    with open(f"{temp_dir}/tokenizer_config.json", "r", encoding="utf-8") as f:
+        tokenizer_config = f.read()
+        assert "chat_template" in tokenizer_config
diff --git a/tests/e2e/utils.py b/tests/e2e/utils.py
index 837b4734fc..f2c567eeb9 100644
--- a/tests/e2e/utils.py
+++ b/tests/e2e/utils.py
@@ -1,14 +1,21 @@
 """
 helper utils for tests
 """
+
+import importlib.util
 import os
 import shutil
 import tempfile
 import unittest
 from functools import wraps
-from importlib.metadata import version
 from pathlib import Path
 
+import torch
+from packaging import version
+from tbparse import SummaryReader
+
+from axolotl.utils.dict import DictDefault
+
 
 def with_temp_dir(test_func):
     @wraps(test_func)
@@ -35,13 +42,273 @@ def most_recent_subdir(path):
     return subdir
 
 
-def require_torch_2_1_1(test_case):
+def require_torch_2_4_1(test_case):
+    """
+    Decorator marking a test that requires torch >= 2.5.1
+    """
+
+    def is_min_2_4_1():
+        torch_version = version.parse(torch.__version__)
+        return torch_version >= version.parse("2.4.1")
+
+    return unittest.skipUnless(is_min_2_4_1(), "test requires torch>=2.4.1")(test_case)
+
+
+def require_torch_2_5_1(test_case):
+    """
+    Decorator marking a test that requires torch >= 2.5.1
+    """
+
+    def is_min_2_5_1():
+        torch_version = version.parse(torch.__version__)
+        return torch_version >= version.parse("2.5.1")
+
+    return unittest.skipUnless(is_min_2_5_1(), "test requires torch>=2.5.1")(test_case)
+
+
+def require_torch_2_6_0(test_case):
+    """
+    Decorator marking a test that requires torch >= 2.6.0
+    """
+
+    def is_min_2_6_0():
+        torch_version = version.parse(torch.__version__)
+        return torch_version >= version.parse("2.6.0")
+
+    return unittest.skipUnless(is_min_2_6_0(), "test requires torch>=2.6.0")(test_case)
+
+
+def require_torch_2_7_0(test_case):
+    """
+    Decorator marking a test that requires torch >= 2.7.0
+    """
+
+    def is_min_2_7_0():
+        torch_version = version.parse(torch.__version__)
+        return torch_version >= version.parse("2.7.0")
+
+    return unittest.skipUnless(is_min_2_7_0(), "test requires torch>=2.7.0")(test_case)
+
+
+def require_torch_2_8_0(test_case):
+    """
+    Decorator marking a test that requires torch >= 2.7.0
+    """
+
+    def is_min_2_8_0():
+        torch_version = version.parse(torch.__version__)
+        return torch_version >= version.parse("2.8.0")
+
+    return unittest.skipUnless(is_min_2_8_0(), "test requires torch>=2.8.0")(test_case)
+
+
+def require_torch_lt_2_6_0(test_case):
+    """
+    Decorator marking a test that requires torch < 2.6.0
+    """
+
+    def is_max_2_6_0():
+        torch_version = version.parse(torch.__version__)
+        return torch_version < version.parse("2.6.0")
+
+    return unittest.skipUnless(is_max_2_6_0(), "test requires torch<2.6.0")(test_case)
+
+
+def require_vllm(test_case):
     """
-    Decorator marking a test that requires torch >= 2.1.1
+    Decorator marking a test that requires a vllm to be installed
     """
 
-    def is_min_2_1_1():
-        torch_version = version("torch")
-        return torch_version >= "2.1.1"
+    def is_vllm_installed():
+        return importlib.util.find_spec("vllm") is not None
+
+    return unittest.skipUnless(
+        is_vllm_installed(), "test requires vllm to be installed"
+    )(test_case)
+
+
+def require_llmcompressor(test_case):
+    """
+    Decorator marking a test that requires a llmcompressor to be installed
+    """
+
+    def is_llmcompressor_installed():
+        return importlib.util.find_spec("llmcompressor") is not None
+
+    return unittest.skipUnless(
+        is_llmcompressor_installed(), "test requires llmcompressor to be installed"
+    )(test_case)
+
+
+def requires_sm_ge_100(test_case):
+    is_sm_ge_100 = (
+        torch.cuda.is_available()
+        and torch.version.cuda
+        and torch.cuda.get_device_capability() >= (10, 0)
+    )
+    return unittest.skipUnless(is_sm_ge_100, "test requires sm>=100")(test_case)
+
+
+def requires_cuda_ge_8_9(test_case):
+    is_cuda_ge_8_9 = (
+        torch.cuda.is_available()
+        and torch.version.cuda
+        and torch.cuda.get_device_capability() >= (8, 9)
+    )
+    return unittest.skipUnless(is_cuda_ge_8_9, "test requires cuda>=8.9")(test_case)
+
+
+def is_hopper():
+    compute_capability = torch.cuda.get_device_capability()
+    return compute_capability == (9, 0)
+
+
+def require_hopper(test_case):
+    return unittest.skipUnless(is_hopper(), "test requires h100/hopper GPU")(test_case)
+
+
+def supports_fp8(test_case):
+    compute_capability = torch.cuda.get_device_capability()
+    return unittest.skipUnless(
+        compute_capability >= (9, 0), "test requires h100 or newer GPU"
+    )(test_case)
+
+
+def check_tensorboard(
+    temp_run_dir: str,
+    tag: str,
+    lt_val: float,
+    assertion_err: str,
+    rtol: float = 0.05,
+    gt_zero: bool = True,
+) -> None:
+    """
+    helper function to parse and check tensorboard logs
+    """
+    tb_log_path = most_recent_subdir(temp_run_dir)
+    event_file = os.path.join(tb_log_path, sorted(os.listdir(tb_log_path))[0])
+    reader = SummaryReader(event_file)
+    df = reader.scalars
+    df = df[(df.tag == tag)]
+    lt_val = (1 + rtol) * lt_val
+    if "%s" in assertion_err:
+        assert df.value.values[-1] < lt_val, assertion_err % df.value.values[-1]
+    else:
+        assert df.value.values[-1] < lt_val, assertion_err
+    if gt_zero:
+        assert df.value.values[-1] > 1e-5, (
+            f"Expected {tag} to be greater than zero, got {df.value.values[-1]}"
+        )
+
+
+def check_tensorboard_loss_decreased(
+    temp_run_dir: str,
+    tag: str | None = None,
+    initial_window: int = 1,
+    final_window: int = 1,
+    min_delta: float | None = None,
+    max_initial: float | None = None,
+    max_final: float | None = None,
+    max_loss_ratio: float = 0.95,
+) -> None:
+    """Check that training actually learned — loss went down and stayed in
+    a sensible range.
+
+    Used with the tiny ``axolotl-ai-co/tiny-*`` CI models, where pretraining
+    was brief enough that final loss won't clear the absolute thresholds used
+    for 135M+ models — but the training pipeline should still behave.
+
+    ``train/train_loss`` is only logged once (end-of-training aggregate). The
+    per-step tag is ``train/loss`` for SFT/LM trainers and may vary across
+    trainers (e.g. DPO). When ``tag`` is None we try common per-step tags in
+    order and use the first with enough samples.
+
+    Two kinds of regression we guard against:
+
+    1. **Loss blew up.** A silent bug (e.g. broken label masking) can start
+       training at an absurdly high loss. ``max_initial`` / ``max_final``
+       assert the measured means stay at-or-below bounds measured from a
+       known-good run. Both are optional but strongly encouraged — loss
+       going *down* from a bad starting scale still looks like "learning."
+
+    2. **Loss didn't go down enough.** ``max_loss_ratio`` (default 0.95)
+       requires ``final <= initial * ratio``. A default below 1.0 means the
+       final window mean must sit at least 5% below the initial window mean
+       — real learning, not noise that happened to land below start. Only
+       raise this for configs where a smaller drop is expected *and*
+       documented (e.g. DPO with near-trivial pairs); in that case you are
+       intentionally weakening the test.
+
+    ``min_delta`` is optional; when set, additionally requires
+    ``final + min_delta <= initial`` — use for configs with enough signal
+    to demand a specific minimum absolute drop.
+    """
+    tb_log_path = most_recent_subdir(temp_run_dir)
+    event_file = os.path.join(tb_log_path, sorted(os.listdir(tb_log_path))[0])
+    reader = SummaryReader(event_file)
+    df = reader.scalars
+
+    if tag is None:
+        candidates = ["train/loss", "train/train_loss"]
+    else:
+        candidates = [tag]
+
+    required = initial_window + final_window
+    chosen_tag, values = None, None
+    for candidate in candidates:
+        sub = df[df.tag == candidate]
+        if len(sub) >= required:
+            chosen_tag = candidate
+            values = sub.value.values
+            break
+
+    available = sorted({t for t in df.tag.unique() if "loss" in t.lower()})
+    assert values is not None, (
+        f"None of the tags {candidates} had ≥{required} logged steps. "
+        f"Loss tags present: {available}"
+    )
+
+    initial = float(values[:initial_window].mean())
+    final = float(values[-final_window:].mean())
+    print(
+        f"[check_tensorboard_loss_decreased] tag={chosen_tag} n={len(values)} "
+        f"initial_mean{initial_window}={initial:.4f} final_mean{final_window}={final:.4f}"
+    )
+    assert final > 1e-5, "Expected loss to be greater than zero"
+    assert final <= initial * max_loss_ratio, (
+        f"Loss did not decrease for {chosen_tag}: "
+        f"initial(mean of first {initial_window})={initial:.4f}, "
+        f"final(mean of last {final_window})={final:.4f}, "
+        f"ratio={final / initial:.4f} (max allowed {max_loss_ratio}). "
+        f"Expected final <= initial — training did not learn."
+    )
+    if min_delta is not None:
+        assert final + min_delta <= initial, (
+            f"Expected loss to decrease by at least {min_delta} for {chosen_tag}: "
+            f"initial={initial:.4f}, final={final:.4f}, delta={initial - final:.4f}"
+        )
+    if max_initial is not None:
+        assert initial <= max_initial, (
+            f"Initial loss {initial:.4f} is above the expected max {max_initial}. "
+            f"Absolute scale is wrong — probably a silent regression "
+            f"(e.g. bad label masking) that bumped the starting point."
+        )
+    if max_final is not None:
+        assert final <= max_final, (
+            f"Final loss {final:.4f} is above the expected max {max_final}. "
+            f"Absolute scale is wrong — probably a silent regression "
+            f"(e.g. bad label masking) that bumped the endpoint."
+        )
+
+
+def check_model_output_exists(temp_dir: str, cfg: DictDefault) -> None:
+    """
+    helper function to check if a model output file exists after training
+
+    checks based on adapter or not (always safetensors in Transformers V5)
+    """
 
-    return unittest.skipUnless(is_min_2_1_1(), "test torch 2.1.1")(test_case)
+    if not cfg.adapter:
+        assert (Path(temp_dir) / "model.safetensors").exists()
+    else:
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
diff --git a/tests/hf_offline_utils.py b/tests/hf_offline_utils.py
new file mode 100644
index 0000000000..b93b83f9fc
--- /dev/null
+++ b/tests/hf_offline_utils.py
@@ -0,0 +1,104 @@
+"""
+test utils for helpers and decorators
+"""
+
+import os
+from contextlib import contextmanager
+from functools import wraps
+
+
+def reload_modules(hf_hub_offline):
+    # Force reload of the modules that check this variable
+    import importlib
+
+    import datasets
+    import huggingface_hub.constants
+    # from huggingface_hub.utils import reset_sessions
+
+    # Reload the constants module first, as others depend on it
+    importlib.reload(huggingface_hub.constants)
+    huggingface_hub.constants.HF_HUB_OFFLINE = hf_hub_offline
+    importlib.reload(datasets.config)
+    datasets.config.HF_HUB_OFFLINE = hf_hub_offline
+
+
+def enable_hf_offline(test_func):
+    """
+    test decorator that sets HF_HUB_OFFLINE environment variable to True and restores it after the test even if the test fails.
+    :param test_func:
+    :return:
+    """
+
+    @wraps(test_func)
+    def wrapper(*args, **kwargs):
+        # Save the original value of HF_HUB_OFFLINE environment variable
+        original_hf_offline = os.getenv("HF_HUB_OFFLINE")
+
+        # Set HF_OFFLINE environment variable to True
+        os.environ["HF_HUB_OFFLINE"] = "1"
+
+        reload_modules(True)
+        try:
+            # Run the test function
+            return test_func(*args, **kwargs)
+        finally:
+            # Restore the original value of HF_HUB_OFFLINE environment variable
+            if original_hf_offline is not None:
+                os.environ["HF_HUB_OFFLINE"] = original_hf_offline
+                reload_modules(bool(original_hf_offline))
+            else:
+                del os.environ["HF_HUB_OFFLINE"]
+                reload_modules(False)
+
+    return wrapper
+
+
+def disable_hf_offline(test_func):
+    """
+    test decorator that sets HF_HUB_OFFLINE environment variable to False and restores it after the wrapped func
+    :param test_func:
+    :return:
+    """
+
+    @wraps(test_func)
+    def wrapper(*args, **kwargs):
+        # Save the original value of HF_HUB_OFFLINE environment variable
+        original_hf_offline = os.getenv("HF_HUB_OFFLINE")
+
+        # Set HF_OFFLINE environment variable to True
+        os.environ["HF_HUB_OFFLINE"] = "0"
+
+        reload_modules(False)
+        try:
+            # Run the test function
+            return test_func(*args, **kwargs)
+        finally:
+            # Restore the original value of HF_HUB_OFFLINE environment variable
+            if original_hf_offline is not None:
+                os.environ["HF_HUB_OFFLINE"] = original_hf_offline
+                reload_modules(bool(original_hf_offline))
+            else:
+                del os.environ["HF_HUB_OFFLINE"]
+                reload_modules(False)
+
+    return wrapper
+
+
+@contextmanager
+def hf_offline_context(hf_hub_offline):
+    """
+    Context manager that sets HF_HUB_OFFLINE environment variable to the given value.
+    :param hf_hub_offline: The new value for HF_HUB_OFFLINE.
+    :return: A context manager.
+    """
+    original_hf_offline = os.getenv("HF_HUB_OFFLINE")
+    os.environ["HF_HUB_OFFLINE"] = str(hf_hub_offline)
+    reload_modules(bool(hf_hub_offline))
+    yield
+    # Restore the original value of HF_HUB_OFFLINE environment variable
+    if original_hf_offline is not None:
+        os.environ["HF_HUB_OFFLINE"] = original_hf_offline
+        reload_modules(bool(original_hf_offline))
+    else:
+        del os.environ["HF_HUB_OFFLINE"]
+        reload_modules(False)
diff --git a/tests/integrations/__init__.py b/tests/integrations/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/integrations/kernels/__init__.py b/tests/integrations/kernels/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/integrations/kernels/dsv4/__init__.py b/tests/integrations/kernels/dsv4/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/integrations/kernels/dsv4/test_dsv4_attention_ops.py b/tests/integrations/kernels/dsv4/test_dsv4_attention_ops.py
new file mode 100644
index 0000000000..b229b6bd93
--- /dev/null
+++ b/tests/integrations/kernels/dsv4/test_dsv4_attention_ops.py
@@ -0,0 +1,161 @@
+"""Parity tests for the DSV4 attention custom ops (sliding, CSA dense, CSA top-k gather).
+
+Each kernel's fwd/bwd launch is registered as a ``torch.ops.axolotl.*`` custom op; these
+tests check the ops exist and that the autograd entry points still match a dense fp32
+eager reference (fwd outputs + input grads, including the per-head sink grad).
+"""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+DEV = "cuda"
+B, H, S, D = 2, 8, 128, 128
+T, K, W = 64, 32, 32
+
+
+def _rel(a, b):
+    return (a.float() - b.float()).norm().item() / max(b.float().norm().item(), 1e-12)
+
+
+def _sink_softmax(logits, sinks):
+    """Append the per-head sink as a logit-only column, softmax, drop it."""
+    Bb, Hh, Ss, _ = logits.shape
+    sink = sinks.float().view(1, Hh, 1, 1).expand(Bb, Hh, Ss, 1)
+    return torch.cat([logits, sink], -1).softmax(-1)[..., :-1]
+
+
+def _window_mask(s_q, s_k, window, device):
+    i = torch.arange(s_q, device=device)[:, None]
+    j = torch.arange(s_k, device=device)[None, :]
+    return (j <= i) & (i - j < window)
+
+
+def test_ops_registered():
+    from axolotl.integrations.kernels.libs import dsv4  # noqa: F401
+
+    for name in (
+        "dsv4_sliding_attn_fwd",
+        "dsv4_sliding_attn_bwd",
+        "dsv4_csa_attn_fwd",
+        "dsv4_csa_attn_bwd",
+        "dsv4_csa_topk_attn_fwd",
+        "dsv4_csa_topk_attn_bwd",
+    ):
+        assert hasattr(torch.ops.axolotl, name), name
+
+
+def test_sliding_attn_matches_eager():
+    from axolotl.integrations.kernels.libs.dsv4 import sliding_attn
+
+    torch.manual_seed(0)
+    q = torch.randn(B, H, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    k = torch.randn(B, 1, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    v = torch.randn(B, 1, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    sinks = torch.randn(H, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    qr, kr, vr, sr = (t.detach().clone().requires_grad_() for t in (q, k, v, sinks))
+
+    out = sliding_attn(q, k, v, sinks, None, W)
+    out.float().pow(2).mean().backward()
+
+    scale = D**-0.5
+    scores = qr.float() @ kr[:, 0].float()[:, None].transpose(-1, -2) * scale
+    scores = scores.masked_fill(~_window_mask(S, S, W, DEV), float("-inf"))
+    p = _sink_softmax(scores, sr)
+    ref = p @ vr[:, 0].float()[:, None]
+    ref.pow(2).mean().backward()
+
+    assert _rel(out, ref) < 0.02
+    assert _rel(q.grad, qr.grad) < 0.05
+    assert _rel(k.grad, kr.grad) < 0.05
+    assert _rel(v.grad, vr.grad) < 0.05
+    assert _rel(sinks.grad, sr.grad) < 0.05
+
+
+def test_csa_attn_matches_eager():
+    from axolotl.integrations.kernels.libs.dsv4 import csa_attn
+
+    torch.manual_seed(1)
+    q = torch.randn(B, H, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    kvs = torch.randn(B, 1, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    kvc = torch.randn(B, 1, T, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    sinks = torch.randn(H, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    # additive block bias: finite bias on causally-valid compressed slots, -inf elsewhere
+    comp_valid = (
+        torch.arange(T, device=DEV)[None, :] * 2 <= torch.arange(S, device=DEV)[:, None]
+    )
+    bb = torch.where(
+        comp_valid, torch.randn(B, S, T, device=DEV), float("-inf")
+    ).unsqueeze(1)
+    qr, ksr, kcr, sr = (
+        t.detach().clone().requires_grad_() for t in (q, kvs, kvc, sinks)
+    )
+
+    out = csa_attn(q, kvs, kvc, bb, sinks, None, W)
+    out.float().pow(2).mean().backward()
+
+    scale = D**-0.5
+    ss = qr.float() @ ksr[:, 0].float()[:, None].transpose(-1, -2) * scale
+    ss = ss.masked_fill(~_window_mask(S, S, W, DEV), float("-inf"))
+    sc = qr.float() @ kcr[:, 0].float()[:, None].transpose(-1, -2) * scale + bb.float()
+    p = _sink_softmax(torch.cat([ss, sc], -1), sr)
+    ref = (
+        p[..., :S] @ ksr[:, 0].float()[:, None]
+        + p[..., S:] @ kcr[:, 0].float()[:, None]
+    )
+    ref.pow(2).mean().backward()
+
+    assert _rel(out, ref) < 0.02
+    assert _rel(q.grad, qr.grad) < 0.05
+    assert _rel(kvs.grad, ksr.grad) < 0.05
+    assert _rel(kvc.grad, kcr.grad) < 0.05
+    assert _rel(sinks.grad, sr.grad) < 0.05
+
+
+def test_csa_attn_topk_matches_eager():
+    from axolotl.integrations.kernels.libs.dsv4 import csa_attn_topk
+
+    torch.manual_seed(2)
+    q = torch.randn(B, H, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    kvs = torch.randn(B, 1, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    kvc = torch.randn(B, 1, T, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    sinks = torch.randn(H, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    # unique top-k per position (duplicates would double-count in any implementation);
+    # early positions get -1 (invalid) tails like the real indexer emits
+    idx = torch.stack(
+        [
+            torch.stack([torch.randperm(T, device=DEV)[:K] for _ in range(S)])
+            for _ in range(B)
+        ]
+    ).to(torch.int32)
+    idx[:, : S // 8, K // 2 :] = -1
+    qr, ksr, kcr, sr = (
+        t.detach().clone().requires_grad_() for t in (q, kvs, kvc, sinks)
+    )
+
+    out = csa_attn_topk(q, kvs, kvc, idx, sinks, None, W)
+    out.float().pow(2).mean().backward()
+
+    sel = torch.zeros(B, S, T, dtype=torch.bool, device=DEV)
+    for b in range(B):
+        for s in range(S):
+            row = idx[b, s]
+            sel[b, s, row[row >= 0].long()] = True
+    scale = D**-0.5
+    ss = qr.float() @ ksr[:, 0].float()[:, None].transpose(-1, -2) * scale
+    ss = ss.masked_fill(~_window_mask(S, S, W, DEV), float("-inf"))
+    sc = qr.float() @ kcr[:, 0].float()[:, None].transpose(-1, -2) * scale
+    sc = sc.masked_fill(~sel[:, None], float("-inf"))
+    p = _sink_softmax(torch.cat([ss, sc], -1), sr)
+    ref = (
+        p[..., :S] @ ksr[:, 0].float()[:, None]
+        + p[..., S:] @ kcr[:, 0].float()[:, None]
+    )
+    ref.pow(2).mean().backward()
+
+    assert _rel(out, ref) < 0.02
+    assert _rel(q.grad, qr.grad) < 0.05
+    assert _rel(kvs.grad, ksr.grad) < 0.05
+    assert _rel(kvc.grad, kcr.grad) < 0.05
+    assert _rel(sinks.grad, sr.grad) < 0.05
diff --git a/tests/integrations/kernels/glm_dsa/__init__.py b/tests/integrations/kernels/glm_dsa/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/integrations/kernels/glm_dsa/test_dsa_attention_eager.py b/tests/integrations/kernels/glm_dsa/test_dsa_attention_eager.py
new file mode 100644
index 0000000000..d97ee0405a
--- /dev/null
+++ b/tests/integrations/kernels/glm_dsa/test_dsa_attention_eager.py
@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""End-to-end correctness: our DSA-kernel attention path == HuggingFace eager GlmMoeDsaAttention.
+
+The patched ``_glm_dsa_attention_forward`` replaces the eager forward with the absorbed-MLA path
+(indexer top-k -> absorbed query -> sparse gather/dense over the compressed latent -> value
+projection). This asserts it numerically matches the stock eager module on identical weights/inputs,
+covering both the dense (topk >= seq) and sparse (topk < seq) regimes, with the stock and the fused
+indexer. Combined with test_mla_absorb_gather (gather == dense), the full eager->dense->gather chain
+is verified.
+"""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+pytest.importorskip(
+    "transformers.models.glm_moe_dsa.modeling_glm_moe_dsa",
+    reason="glm_moe_dsa required",
+)
+from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import (  # noqa: E402
+    GlmMoeDsaConfig,
+)
+from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import (  # noqa: E402
+    GlmMoeDsaAttention,
+    GlmMoeDsaRotaryEmbedding,
+)
+
+from axolotl.integrations.kernels.libs.glm_dsa.patch import (  # noqa: E402
+    patch_glm_moe_dsa_attention,
+)
+
+DEV = "cuda"
+DT = torch.bfloat16
+
+
+def _build(index_topk, seed=0):
+    torch.manual_seed(seed)
+    cfg = GlmMoeDsaConfig(
+        hidden_size=512,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        q_lora_rank=256,
+        kv_lora_rank=512,
+        qk_nope_head_dim=192,
+        qk_rope_head_dim=64,
+        v_head_dim=256,
+        index_topk=index_topk,
+        index_head_dim=128,
+        index_n_heads=2,
+        num_hidden_layers=1,
+    )
+    cfg._attn_implementation = "eager"
+    attn = GlmMoeDsaAttention(cfg, layer_idx=0).to(DEV, DT).eval()
+    rope = GlmMoeDsaRotaryEmbedding(cfg).to(DEV)
+    return cfg, attn, rope
+
+
+@pytest.mark.parametrize("S,index_topk", [(64, 16), (48, 128)], ids=["sparse", "dense"])
+@pytest.mark.parametrize(
+    "use_fused_indexer", [False, True], ids=["eager_idx", "fused_idx"]
+)
+def test_dsa_attention_matches_hf_eager(S, index_topk, use_fused_indexer):
+    cfg, attn, rope = _build(index_topk)
+    B = 1
+    hidden = torch.randn(B, S, cfg.hidden_size, device=DEV, dtype=DT)
+    position_ids = torch.arange(S, device=DEV).unsqueeze(0)
+    cos, sin = rope(hidden, position_ids)
+    pe = (cos, sin)
+
+    import torch.nn as nn
+
+    with torch.no_grad():
+        out_eager = attn(hidden, pe, None, position_ids=position_ids)[0]
+        wrapper = nn.Module()
+        wrapper.add_module("a", attn)
+        assert (
+            patch_glm_moe_dsa_attention(wrapper, use_fused_indexer=use_fused_indexer)
+            == 1
+        )
+        out_kernel = attn(hidden, pe, None, position_ids=position_ids)[0]
+
+    assert torch.isfinite(out_kernel).all()
+    err = (out_eager.float() - out_kernel.float()).abs().max().item()
+    scale = out_eager.float().abs().max().item() + 1e-6
+    assert err / scale < 5e-2, f"rel err {err / scale:.4g} (abs {err:.4g})"
diff --git a/tests/integrations/kernels/glm_dsa/test_glm_dsa_cpu.py b/tests/integrations/kernels/glm_dsa/test_glm_dsa_cpu.py
new file mode 100644
index 0000000000..f6c04ee486
--- /dev/null
+++ b/tests/integrations/kernels/glm_dsa/test_glm_dsa_cpu.py
@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""CPU regression tests for two GLM-5.2 workarounds (no GPU needed):
+
+1. **Expert-capacity cap** (DeepEP #24): DeepEP's intranode combine deadlocks once one expert
+   receives too many tokens (the GLM router concentrates with depth). ``_apply_expert_capacity``
+   drops the lowest-weight excess (token,expert) assignments so no expert exceeds ``cap``.
+2. **sm90 backward-autotune prune** (#18): on Hopper the full bwd autotune grid trials a
+   nondeterministic invalid-PC (CUDA 718); ``_bwd_prune_sm90_safe`` must collapse to a single config
+   on sm90 (so autotune never trials the rest) while leaving other archs to normal SMEM pruning.
+"""
+
+import os
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+import torch
+
+from axolotl.integrations.expert_parallel.experts_fn import _apply_expert_capacity
+
+
+def test_expert_capacity_caps_tokens_per_expert():
+    # 6 tokens, top-2 -> 12 (token,expert) assignments, 3 experts. cap=2.
+    topk_idx = torch.tensor(
+        [[0, 1], [0, 2], [0, 1], [0, 2], [1, 2], [0, 1]], dtype=torch.int64
+    )
+    topk_w = torch.tensor(
+        [[0.9, 0.5], [0.8, 0.4], [0.7, 0.6], [0.3, 0.2], [0.95, 0.1], [0.55, 0.45]]
+    )
+    cap = 2
+    out, w = _apply_expert_capacity(topk_idx.clone(), topk_w, cap)
+
+    # every expert now has <= cap surviving assignments
+    for e in range(3):
+        assert int((out == e).sum()) <= cap, f"expert {e} exceeds cap"
+    # surviving weights are rescaled back to each token's original gate sum (no silent attenuation)
+    survived = (out >= 0).any(dim=-1)
+    assert torch.allclose(
+        w.sum(dim=-1)[survived], topk_w.sum(dim=-1)[survived], atol=1e-6
+    )
+    # dropped slots are sentinelled to -1, and the number dropped is exactly the overflow
+    n_orig = int((topk_idx >= 0).sum())
+    n_kept = int((out >= 0).sum())
+    assert n_kept == n_orig - int((out == -1).sum())
+    # the HIGHEST-weight assignment to expert 0 (token4? no expert0 weights: 0.9,0.8,0.7,0.3,_,0.55)
+    # the two survivors for expert 0 must be the top-2 weights (0.9 @ tok0, 0.8 @ tok1)
+    surv0_tokens = (out[:, 0] == 0).nonzero().flatten().tolist() + (
+        out[:, 1] == 0
+    ).nonzero().flatten().tolist()
+    assert (
+        0 in surv0_tokens and 1 in surv0_tokens
+    )  # the two highest-weight expert-0 slots kept
+
+
+def test_expert_capacity_noop_when_under_cap():
+    topk_idx = torch.tensor([[0, 1], [2, 3]], dtype=torch.int64)
+    topk_w = torch.tensor([[0.9, 0.1], [0.5, 0.5]])
+    out, w = _apply_expert_capacity(topk_idx.clone(), topk_w, cap=8)
+    assert torch.equal(out, topk_idx)  # nothing dropped when every expert is under cap
+    assert torch.allclose(w, topk_w)  # weights unchanged when nothing is dropped
+
+
+def test_expert_capacity_preserves_existing_sentinels():
+    topk_idx = torch.tensor([[0, -1], [0, 0]], dtype=torch.int64)
+    topk_w = torch.tensor([[0.9, 0.0], [0.8, 0.7]])
+    out, _ = _apply_expert_capacity(topk_idx.clone(), topk_w, cap=2)
+    assert out[0, 1] == -1  # pre-existing -1 untouched
+    assert int((out == 0).sum()) <= 2
+
+
+def test_expert_capacity_no_nan_grad_when_token_fully_dropped():
+    # 3 tokens all route both top-2 slots to expert 0; cap=1 keeps only the single highest-weight
+    # assignment, so tokens 1 and 2 lose every slot (kept_sum==0). The gate-sum rescale must not
+    # backprop 0*inf=NaN into the router weights (double-where guard on the divisor).
+    topk_idx = torch.tensor([[0, 0], [0, 0], [0, 0]], dtype=torch.int64)
+    topk_w = torch.tensor([[0.9, 0.8], [0.7, 0.6], [0.5, 0.4]], requires_grad=True)
+    out, w = _apply_expert_capacity(topk_idx.clone(), topk_w, cap=1)
+
+    assert (~(out >= 0).any(dim=-1)).any()  # at least one token lost every slot
+    assert not torch.isnan(w).any()  # forward weights stay finite
+    w.sum().backward()
+    assert not (torch.isnan(topk_w.grad).any() or torch.isinf(topk_w.grad).any())
+
+
+def _absorb():
+    import axolotl.integrations.kernels.libs.glm_dsa.attention_mla_absorb as M
+
+    return M
+
+
+def test_sm90_bwd_prune_collapses_to_single_config():
+    M = _absorb()
+    # mock the device SMEM so the base SMEM-prune keeps everything (CPU-only, deterministic)
+    with mock.patch(
+        "axolotl.integrations.kernels.libs.glm_dsa._autotune.max_smem",
+        return_value=1 << 30,
+    ):
+        with mock.patch.object(
+            torch.cuda, "get_device_capability", return_value=(9, 0)
+        ):
+            kept = M._bwd_prune_sm90_safe(M._ABSORB_CONFIGS, {}, DL=512, DR=64)
+        assert len(kept) == 1, (
+            "sm90 must collapse the bwd grid to one config (avoids invalid-PC)"
+        )
+
+        with mock.patch.object(
+            torch.cuda, "get_device_capability", return_value=(12, 0)
+        ):
+            kept_sm120 = M._bwd_prune_sm90_safe(M._ABSORB_CONFIGS, {}, DL=512, DR=64)
+        assert len(kept_sm120) > 1, (
+            "non-sm90 keeps the full SMEM-pruned grid for autotuning"
+        )
+
+
+def test_gather_supported_sm90_sm120_and_disable_override():
+    """The sparse gather is enabled on sm90 (Hopper) and sm120; unlisted archs fall back to dense;
+    ``GLM_DSA_DISABLE_GATHER`` forces dense everywhere."""
+    from axolotl.integrations.kernels.libs.glm_dsa import dispatch as D
+
+    dev = torch.device("cuda:0")
+    for cap, expected in [
+        ((9, 0), True),
+        ((12, 0), True),
+        ((8, 0), False),
+        ((10, 0), False),
+    ]:
+        D._GATHER_OK.clear()
+        with mock.patch.object(torch.cuda, "get_device_capability", return_value=cap):
+            assert D._gather_supported(dev) is expected, f"cap {cap}"
+    D._GATHER_OK.clear()
+    with mock.patch.dict(os.environ, {"GLM_DSA_DISABLE_GATHER": "1"}):
+        with mock.patch.object(
+            torch.cuda, "get_device_capability", return_value=(9, 0)
+        ):
+            assert D._gather_supported(dev) is False  # override wins
+
+
+def test_mla_attn_routes_to_gather_under_packing_and_forwards_seq():
+    """Above the crossover the sparse gather is used EVEN under sample packing (it is doc-aware), and
+    ``seq_q``/``seq_k`` are forwarded to it; when the gather is unsupported, dense is used."""
+    from axolotl.integrations.kernels.libs.glm_dsa import dispatch as D
+
+    qa = torch.zeros(1, 4, 8, 576)
+    ks = torch.zeros(1, 8, 576)
+    idx = torch.zeros(1, 8, 4, dtype=torch.int64)
+    seq = torch.zeros(1, 8, dtype=torch.int64)
+
+    with (
+        mock.patch.object(D, "mla_absorb_attn", return_value="GATHER") as mg,
+        mock.patch.object(D, "dense_masked_out_latent", return_value="DENSE"),
+        mock.patch.object(D, "_gather_supported", return_value=True),
+        mock.patch.object(D, "calibrate_crossover", return_value=0),
+    ):
+        out = D.mla_attn(qa, ks, idx, 1.0, seq_q=seq, seq_k=seq)
+        assert out == "GATHER"  # packing no longer forces dense
+        # seq_q / seq_k forwarded as the last two positional args
+        assert mg.call_args.args[-2] is seq and mg.call_args.args[-1] is seq
+
+    with (
+        mock.patch.object(D, "dense_masked_out_latent", return_value="DENSE"),
+        mock.patch.object(D, "_gather_supported", return_value=False),
+        mock.patch.object(D, "calibrate_crossover", return_value=0),
+    ):
+        assert D.mla_attn(qa, ks, idx, 1.0, seq_q=seq, seq_k=seq) == "DENSE"
+
+
+def test_deep_ep_forward_applies_token_capacity():
+    """The capacity cap must actually be applied to the routing before the dispatch layout is built
+    (regression: the cap function existed but was never called)."""
+    from axolotl.integrations.expert_parallel import experts_fn as E
+
+    captured = {}
+
+    class _FakeBuf:
+        def get_dispatch_layout(self, topk_idx, e_global):
+            captured["topk"] = topk_idx.clone()
+            raise RuntimeError("stop")  # short-circuit the rest of the forward
+
+    ntok, K, e_global = 64, 2, 8
+    topk = torch.zeros(
+        ntok, K, dtype=torch.int64
+    )  # column 0 -> expert 0 for ALL tokens
+    topk[:, 1] = 1
+    w = torch.rand(ntok, K)
+    self_mod = SimpleNamespace(num_experts=e_global, num_experts_global=e_global)
+
+    E.set_token_capacity(4)
+    try:
+        with (
+            mock.patch.object(E, "get_buffer", return_value=_FakeBuf()),
+            mock.patch.object(E, "_get_valid_token_mask", return_value=None),
+        ):
+            with pytest.raises(RuntimeError, match="stop"):
+                E._deep_ep_forward(
+                    self_mod, torch.zeros(ntok, 16), topk, w, kernel_name="eager"
+                )
+    finally:
+        E.set_token_capacity(None)
+
+    # expert 0 was requested by all 64 tokens; the cap must hold it to <= 4
+    assert int((captured["topk"] == 0).sum()) <= 4
+
+
+def test_mla_attn_skips_calibration_when_gather_unsupported():
+    """On an unsupported GPU mla_attn must NOT call calibrate_crossover (it would compile/run the
+    gather Triton path we are avoiding)."""
+    from axolotl.integrations.kernels.libs.glm_dsa import dispatch as D
+
+    qa = torch.zeros(1, 4, 8, 576)
+    ks = torch.zeros(1, 8, 576)
+    idx = torch.zeros(1, 8, 4, dtype=torch.int64)
+    calib = mock.MagicMock(return_value=0)
+    with (
+        mock.patch.object(D, "dense_masked_out_latent", return_value="DENSE"),
+        mock.patch.object(D, "_gather_supported", return_value=False),
+        mock.patch.object(D, "calibrate_crossover", calib),
+    ):
+        assert D.mla_attn(qa, ks, idx, 1.0) == "DENSE"
+    calib.assert_not_called()
+
+
+def test_cp_doc_ids_must_be_global_not_per_chunk():
+    """F2: under context parallelism the per-document ids must come from the GLOBAL position_ids, not
+    each rank's local chunk. Otherwise a document crossing a CP boundary gets different (and colliding)
+    ids per rank, which masks valid remote keys / allows cross-document attention."""
+    from axolotl.integrations.kernels.libs.glm_dsa.patch import (
+        _seq_idx_from_position_ids,
+    )
+
+    # doc B (positions 0,1,2,3) spans the cp=2 boundary at index 4 (S=4 queries per rank)
+    global_pos = torch.tensor([[0, 1, 0, 1, 2, 3, 0, 1]])
+    S = 4
+    global_ids = _seq_idx_from_position_ids(global_pos)
+    assert global_ids.tolist() == [[1, 1, 2, 2, 2, 2, 3, 3]]
+
+    # NEW (number docs on the global ids, then slice per rank): the boundary doc shares ONE id
+    r0, r1 = global_ids[:, 0:S], global_ids[:, S : 2 * S]
+    assert (
+        r0[0, -1].item() == r1[0, 0].item()
+    )  # doc B: rank0's last query == rank1's first query
+
+    # OLD (per-local-chunk cumsum): the same doc gets different ids across ranks, and ids collide
+    old_r0 = _seq_idx_from_position_ids(global_pos[:, 0:S])
+    old_r1 = _seq_idx_from_position_ids(global_pos[:, S : 2 * S])
+    assert (
+        old_r0[0, -1].item() != old_r1[0, 0].item()
+    )  # doc B inconsistent across ranks (the bug)
+    assert (
+        old_r1[0, 2].item() == old_r0[0, 0].item()
+    )  # doc C (rank1) collides with doc A (rank0)
diff --git a/tests/integrations/kernels/glm_dsa/test_mla_absorb_gather.py b/tests/integrations/kernels/glm_dsa/test_mla_absorb_gather.py
new file mode 100644
index 0000000000..6329f391b8
--- /dev/null
+++ b/tests/integrations/kernels/glm_dsa/test_mla_absorb_gather.py
@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Regression tests for the GLM-DSA absorbed-MLA sparse gather kernel.
+
+These guard two subtle, hard-won fixes:
+
+1. **Leading-all-masked-block softmax NaN.** The online-softmax running max ``m_i`` starts at
+   ``-inf``; a query whose *first* ``BN`` block of selected keys is entirely causally/doc-masked
+   keeps ``m_new == -inf``, so ``alpha = exp(m_i - m_new) = exp(-inf + inf) = NaN`` — *even though
+   the query has valid keys in later blocks*. This poisons the whole output. It surfaces at long
+   context (``topk`` >> a short query's causal positions) and is arch-independent. The kernel guards
+   it with ``m_safe = where(m_new == -inf, 0, m_new)`` (+ an ``l_i == 0`` division guard for
+   fully-masked rows). The reference is the dense path, which is always finite.
+
+2. **Var-len / document-aware masking under sample packing** (``seq_k[idx] == seq_q[s]``): the gather
+   must match the dense per-document mask so packing can use the sparse path.
+"""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+from axolotl.integrations.kernels.libs.glm_dsa.attention_mla_absorb import (  # noqa: E402
+    mla_absorb_attn,
+)
+from axolotl.integrations.kernels.libs.glm_dsa.config import (  # noqa: E402
+    KV_LORA_RANK,
+    QK_ROPE_HEAD_DIM,
+)
+from axolotl.integrations.kernels.libs.glm_dsa.dispatch import (  # noqa: E402
+    dense_masked_out_latent,
+)
+
+DEV = "cuda"
+DQK = KV_LORA_RANK + QK_ROPE_HEAD_DIM  # 576 (512 + 64)
+
+
+def _causal_topk_future_first(B, S, Skv, TOPK, device, seed=0):
+    """Per query ``s`` select ``TOPK`` keys, ordering FUTURE (causally-invalid, ``> s``) positions
+    FIRST so the leading topk block(s) are entirely masked — the exact trigger for the running-max
+    NaN — while still including valid causal keys (and ``s`` itself) in later slots."""
+    idx = torch.empty(B, S, TOPK, dtype=torch.int32, device=device)
+    for s in range(S):
+        future = torch.arange(s + 1, Skv, device=device)
+        causal = torch.arange(0, s + 1, device=device)
+        pool = torch.cat([future, causal])  # future FIRST -> leading masked
+        if pool.numel() >= TOPK:
+            row = pool[:TOPK].clone()
+        else:
+            pad = causal[-1:].repeat(TOPK - pool.numel())
+            row = torch.cat([pool, pad])
+        row[-1] = s  # guarantee self (valid) present
+        idx[:, s, :] = row.to(torch.int32)
+    return idx
+
+
+@pytest.mark.parametrize("S,Skv,TOPK", [(64, 64, 48), (96, 96, 80)])
+def test_gather_leading_masked_block_is_finite_and_matches_dense(S, Skv, TOPK):
+    """The whole point: with leading masked blocks the gather must NOT NaN and must equal dense."""
+    torch.manual_seed(0)
+    B, H = 1, 16
+    q = torch.randn(B, H, S, DQK, device=DEV, dtype=torch.bfloat16)
+    k = torch.randn(B, Skv, DQK, device=DEV, dtype=torch.bfloat16)
+    idx = _causal_topk_future_first(B, S, Skv, TOPK, DEV)
+    scale = 1.0 / (DQK**0.5)
+
+    qg = q.clone().requires_grad_(True)
+    kg = k.clone().requires_grad_(True)
+    og = mla_absorb_attn(qg, kg, idx, scale, 0)
+    og.float().pow(2).mean().backward()
+
+    assert torch.isfinite(og).all(), (
+        "gather forward produced NaN/Inf on leading-masked-block input"
+    )
+    assert torch.isfinite(qg.grad).all(), "gather dq produced NaN/Inf"
+    assert torch.isfinite(kg.grad).all(), "gather dk produced NaN/Inf"
+
+    qd = q.clone().requires_grad_(True)
+    kd = k.clone().requires_grad_(True)
+    od = dense_masked_out_latent(qd, kd, idx, scale, 0)
+    od.float().pow(2).mean().backward()
+
+    # bf16 tensor-core tolerance
+    assert (og.float() - od.float()).abs().max().item() < 5e-2
+    assert (qg.grad.float() - qd.grad.float()).abs().max().item() < 5e-2
+
+
+def test_gather_fully_masked_query_is_finite():
+    """A query whose every selected key is in the future (zero valid causal keys) must yield a
+    finite (zero) row via the ``l_i == 0`` guard, not 0/0 = NaN."""
+    torch.manual_seed(1)
+    B, H, S, Skv, TOPK = 1, 8, 32, 32, 16
+    q = torch.randn(B, H, S, DQK, device=DEV, dtype=torch.bfloat16)
+    k = torch.randn(B, Skv, DQK, device=DEV, dtype=torch.bfloat16)
+    # causal-valid for everyone...
+    idx = torch.zeros(B, S, TOPK, dtype=torch.int32, device=DEV)
+    for s in range(S):
+        idx[:, s, :] = torch.randint(0, s + 1, (TOPK,), device=DEV, dtype=torch.int32)
+    # ...except query 5, whose every selected key is strictly future (invalid)
+    idx[:, 5, :] = torch.arange(6, 6 + TOPK, device=DEV, dtype=torch.int32).clamp_(
+        max=Skv - 1
+    )
+    out = mla_absorb_attn(q, k, idx, 1.0 / (DQK**0.5), 0)
+    assert torch.isfinite(out).all(), (
+        "fully-masked query produced NaN/Inf (0/0 not guarded)"
+    )
+
+
+def test_gather_docmask_matches_dense_under_packing():
+    """Doc-aware (var-len) gather must match the dense per-document mask: a query only attends keys
+    in its own document (``seq_k[idx] == seq_q[s]``).
+
+    NB: the dense reference de-duplicates selected keys (a boolean ``scatter_``) while the gather
+    sums each slot independently, so they only agree when ``idx`` has DISTINCT keys per query. We
+    build distinct causal keys and assert the exact match on queries with enough causal positions to
+    fill ``TOPK`` without padding; every query is still asserted finite.
+    """
+    torch.manual_seed(2)
+    B, H, S, Skv, TOPK = 1, 16, 64, 64, 16
+    q = torch.randn(B, H, S, DQK, device=DEV, dtype=torch.bfloat16)
+    k = torch.randn(B, Skv, DQK, device=DEV, dtype=torch.bfloat16)
+    seq = torch.zeros(B, S, dtype=torch.int32, device=DEV)
+    seq[:, S // 2 :] = 1  # two documents
+    seq_q, seq_k = seq, seq.clone()
+    # distinct causal topk (crosses the doc boundary so masking matters); short queries cycle their
+    # available causal keys (creates dups -> excluded from the exact comparison via ``no_dup``).
+    idx = torch.empty(B, S, TOPK, dtype=torch.int32, device=DEV)
+    no_dup = torch.zeros(S, dtype=torch.bool, device=DEV)
+    for s in range(S):
+        causal = torch.randperm(s + 1, device=DEV).to(torch.int32)
+        if causal.numel() >= TOPK:
+            idx[:, s, :] = causal[:TOPK]
+            no_dup[s] = True
+        else:
+            reps = (TOPK + causal.numel() - 1) // causal.numel()
+            idx[:, s, :] = causal.repeat(reps)[:TOPK]
+    scale = 1.0 / (DQK**0.5)
+
+    og = mla_absorb_attn(q.clone(), k.clone(), idx, scale, 0, seq_q, seq_k)
+    od = dense_masked_out_latent(q.clone(), k.clone(), idx, scale, 0, seq_q, seq_k)
+    assert torch.isfinite(og).all()
+    # exact agreement on the distinct-key (no-dup) queries only
+    masked = (og.float() - od.float())[:, :, no_dup, :].abs()
+    assert torch.isfinite(masked).all()
+    assert masked.max().item() < 5e-2
+    # doc mask must actually change the result vs causal-only (no seq)
+    on = mla_absorb_attn(q.clone(), k.clone(), idx, scale, 0)
+    assert (og.float() - on.float())[:, :, no_dup, :].abs().max().item() > 1e-2
diff --git a/tests/integrations/kernels/glm_dsa/test_sparse_topk_attn.py b/tests/integrations/kernels/glm_dsa/test_sparse_topk_attn.py
new file mode 100644
index 0000000000..791b4abc0e
--- /dev/null
+++ b/tests/integrations/kernels/glm_dsa/test_sparse_topk_attn.py
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Parity tests for the GLM-DSA sparse top-k attention kernel and its registered custom ops.
+
+``sparse_attn`` gathers each query's selected keys and runs flash online-softmax over just those;
+the reference gathers the same keys in fp32 torch and takes a masked softmax. Forward outputs and
+input gradients (dq/dk/dv) must agree. Also asserts the kernels are dispatcher-visible as
+``torch.ops.axolotl.*`` (torch.compile opacity + selective-checkpointing matching)."""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+import axolotl.integrations.kernels.libs.glm_dsa.attention_mla_absorb  # noqa: F401,E402
+from axolotl.integrations.kernels.libs.glm_dsa.attention_topk import (  # noqa: E402
+    sparse_attn,
+)
+
+DEV = "cuda"
+
+
+def test_glm_dsa_ops_registered():
+    for name in (
+        "glm_dsa_sparse_topk_attn_fwd",
+        "glm_dsa_sparse_topk_attn_bwd",
+        "glm_dsa_mla_absorb_attn_fwd",
+        "glm_dsa_mla_absorb_attn_bwd",
+    ):
+        assert hasattr(torch.ops.axolotl, name), f"torch.ops.axolotl.{name} missing"
+
+
+def _causal_window_topk(B, S, T, device):
+    """Per query ``s``: the last-T causal positions, padded with DISTINCT future (masked) ones so
+    slot-summed gather == de-duplicated dense softmax."""
+    idx = torch.empty(B, S, T, dtype=torch.int32, device=device)
+    for s in range(S):
+        causal = torch.arange(max(0, s - T + 1), s + 1, device=device)
+        pad = torch.arange(s + 1, s + 1 + T - causal.numel(), device=device)
+        idx[:, s, :] = torch.cat([causal, pad]).to(torch.int32)
+    return idx
+
+
+def _ref_sparse_attn(q, k, v, idx, scale):
+    B, H, S, D = q.shape
+    DV = v.shape[-1]
+    T = idx.shape[-1]
+    i = idx.long()[:, None].expand(B, H, S, T)
+    kg = torch.gather(k, 2, i.reshape(B, H, S * T, 1).expand(-1, -1, -1, D)).view(
+        B, H, S, T, D
+    )
+    vg = torch.gather(v, 2, i.reshape(B, H, S * T, 1).expand(-1, -1, -1, DV)).view(
+        B, H, S, T, DV
+    )
+    scores = (q.unsqueeze(-2) * kg).sum(-1) * scale
+    valid = idx.long()[:, None] <= torch.arange(S, device=q.device)[None, None, :, None]
+    p = scores.masked_fill(~valid, float("-inf")).softmax(dim=-1)
+    return (p.unsqueeze(-1) * vg).sum(-2)
+
+
+def test_sparse_topk_matches_reference_fwd_and_grads():
+    torch.manual_seed(0)
+    B, H, S, D, DV, T = 1, 2, 32, 64, 64, 8
+    q = torch.randn(B, H, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    k = torch.randn(B, H, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    v = torch.randn(B, H, S, DV, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    idx = _causal_window_topk(B, S, T, DEV)
+    scale = D**-0.5
+    # bf16-exact cotangent so both paths backprop the identical dout (unit scale -> the grad
+    # tolerance is meaningful, unlike a mean-loss whose grads are ~1e-4)
+    g = torch.randn(B, H, S, DV, device=DEV, dtype=torch.bfloat16).float()
+
+    out = sparse_attn(q, k, v, idx, scale)
+    (out.float() * g).sum().backward()
+
+    qr = q.detach().float().requires_grad_(True)
+    kr = k.detach().float().requires_grad_(True)
+    vr = v.detach().float().requires_grad_(True)
+    oref = _ref_sparse_attn(qr, kr, vr, idx, scale)
+    (oref * g).sum().backward()
+
+    assert torch.isfinite(out).all()
+    assert (out.float() - oref).abs().max().item() < 3e-2
+    for got, ref in ((q.grad, qr.grad), (k.grad, kr.grad), (v.grad, vr.grad)):
+        assert torch.isfinite(got).all()
+        assert (got.float() - ref).abs().max().item() < 3e-2
diff --git a/tests/integrations/kernels/scattermoe_lora/__init__.py b/tests/integrations/kernels/scattermoe_lora/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/integrations/kernels/scattermoe_lora/bench_int64_kernel.py b/tests/integrations/kernels/scattermoe_lora/bench_int64_kernel.py
new file mode 100644
index 0000000000..d4cad1fb2a
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/bench_int64_kernel.py
@@ -0,0 +1,229 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+ScatterMoE — INT64_INDICES vs INT32 dense scatter2scatter benchmark.
+
+Times the dense ``kernels.ops.scatter2scatter`` at three representative
+shapes and reports ms/iter for both ``INT64_INDICES=False`` (int32 fast
+path) and ``INT64_INDICES=True`` (int64 safe path). The third shape is
+the previously-failing seq=512K / 16-shard config; at that scale the
+int32 path is incorrect (silent overflow corruption) so the int32 row
+is gated by the ``_SCATTER2SCATTER_INT32_LIMIT`` and reported as the
+chunked workaround's wall-clock instead (also for comparison against
+the chunking baseline that PR #3667 shipped).
+
+Run from the repo root:
+
+  python tests/integrations/kernels/scattermoe_lora/bench_int64_kernel.py
+
+A markdown summary is printed to stdout and written to
+``bench_int64_kernel_results.md`` next to this script.
+"""
+
+from __future__ import annotations
+
+import argparse
+import statistics
+import subprocess
+from pathlib import Path
+from typing import Callable
+
+import torch
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops import (
+    scatter2scatter,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+    flatten_sort_count,
+)
+
+# Sufficient condition for int32 pointer arithmetic to overflow in the
+# scatter2scatter Triton kernel: ``L_scattered * y_dim >= 2**31``.
+_SCATTER2SCATTER_INT32_LIMIT = 2**31
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+def gpu_name() -> str:
+    try:
+        out = subprocess.check_output(["nvidia-smi", "-L"], text=True).strip()
+        first = out.splitlines()[0]
+        if ":" in first:
+            after_colon = first.split(":", 1)[1].strip()
+            return after_colon.split("(", 1)[0].strip()
+        return first
+    except Exception:
+        return torch.cuda.get_device_name(0)
+
+
+def _time_ms(fn: Callable[[], torch.Tensor], iters: int = 10, warmup: int = 3) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    samples = []
+    for _ in range(iters):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        torch.cuda.synchronize()
+        samples.append(start.elapsed_time(end))
+    return statistics.median(samples)
+
+
+def _build_inputs(
+    *, T: int, hidden: int, top_k: int, n: int, num_experts: int, seed: int
+):
+    torch.manual_seed(seed)
+    x = torch.randn(T, hidden, device=DEVICE, dtype=DTYPE)
+    W = torch.randn(num_experts, hidden, n, device=DEVICE, dtype=DTYPE) * 0.02
+    logits = torch.randn(T, num_experts, device=DEVICE)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
+    sei, ssi, _ = flatten_sort_count(top_idx, num_experts)
+    return x, W, sei, ssi
+
+
+def _run_shape(name: str, *, T: int, hidden: int, top_k: int, n: int, num_experts: int):
+    x, W, sei, ssi = _build_inputs(
+        T=T, hidden=hidden, top_k=top_k, n=n, num_experts=num_experts, seed=42
+    )
+    L_scattered = sei.size(0)
+    out_elements = L_scattered * n
+    overflow = out_elements >= _SCATTER2SCATTER_INT32_LIMIT
+    auto_int64 = overflow  # the wrapper's auto-dispatch verdict
+
+    def call(int64_indices: bool):
+        return scatter2scatter(
+            X=x,
+            W=W,
+            sorted_expert_idxs=sei,
+            sorted_scattered_idxs=ssi,
+            k=top_k,
+            x_grouped=False,
+            y_grouped=True,
+            int64_indices=int64_indices,
+        )
+
+    # Warm both Triton variants (separate JITs per constexpr).
+    if overflow:
+        # int32 path is unsafe at overflow shapes (silent corruption); skip.
+        ms_i32 = None
+    else:
+        ms_i32 = _time_ms(lambda: call(False))
+    ms_i64 = _time_ms(lambda: call(True))
+
+    return {
+        "name": name,
+        "T": T,
+        "hidden": hidden,
+        "top_k": top_k,
+        "n": n,
+        "num_experts": num_experts,
+        "L_scattered": L_scattered,
+        "out_elements": out_elements,
+        "overflow": overflow,
+        "auto_int64": auto_int64,
+        "ms_i32": ms_i32,
+        "ms_i64": ms_i64,
+    }
+
+
+def _fmt(v):
+    if v is None:
+        return "—"
+    return f"{v:.3f}"
+
+
+def _markdown(rows, gpu_label: str) -> str:
+    lines = []
+    lines.append("# scatter2scatter INT64_INDICES bench")
+    lines.append("")
+    lines.append(f"GPU: **{gpu_label}**")
+    lines.append("")
+    lines.append("Median of 10 iters, 3 warmup. `top_k=8`, dtype=bf16, 128 experts.")
+    lines.append("")
+    lines.append(
+        "`auto_int64` is the wrapper's auto-dispatch verdict from "
+        "`_needs_int64_indices`. At overflow shapes the int32 path "
+        "is silently incorrect (the multiplication wraps mid-buffer), "
+        "so only the int64 timing is reported."
+    )
+    lines.append("")
+    lines.append(
+        "| Shape | T | L_scattered | out elems | auto_int64 | int32 ms | int64 ms | int64 vs int32 (%) |"
+    )
+    lines.append("|---|---|---|---|---|---|---|---|")
+    for r in rows:
+        if r["ms_i32"] is not None and r["ms_i64"] is not None:
+            pen = 100.0 * (r["ms_i64"] - r["ms_i32"]) / r["ms_i32"]
+            pen_s = f"{pen:+.1f}"
+        else:
+            pen_s = "—"
+        lines.append(
+            f"| {r['name']} | {r['T']} | {r['L_scattered']} | "
+            f"{r['out_elements']:.2e} | {str(r['auto_int64'])} | "
+            f"{_fmt(r['ms_i32'])} | {_fmt(r['ms_i64'])} | {pen_s} |"
+        )
+    lines.append("")
+    lines.append(
+        "Acceptance: ≤5% regression on the int32 fast path at "
+        "small/medium shapes (the auto-dispatch picks int32 there, so "
+        "this row characterises the JIT overhead of having an int64 "
+        "variant available)."
+    )
+    lines.append("")
+    return "\n".join(lines)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--out",
+        type=str,
+        default=None,
+        help="Output markdown path (default: alongside script)",
+    )
+    args = parser.parse_args()
+
+    # Three representative shapes per the task spec.
+    shapes = [
+        # name              T (tokens before top_k expansion), hidden, top_k, N (out), num_experts
+        dict(name="small", T=8_192, hidden=2048, top_k=8, n=2048, num_experts=128),
+        dict(name="medium", T=128_000, hidden=2048, top_k=8, n=2048, num_experts=128),
+        # Overflow shape: 524288 / 16 shards = 32768 tokens, top_k=8 -> L=262144,
+        # N=16384 (= 2*intermediate at the bench config) -> 2**32 elements.
+        dict(
+            name="overflow_524k_s16",
+            T=32_768,
+            hidden=2048,
+            top_k=8,
+            n=16_384,
+            num_experts=128,
+        ),
+    ]
+
+    rows = []
+    for s in shapes:
+        print(f"running {s['name']} ...", flush=True)
+        rows.append(_run_shape(**s))
+        torch.cuda.empty_cache()
+
+    label = gpu_name()
+    md = _markdown(rows, label)
+    print(md)
+
+    if args.out:
+        out_path = Path(args.out)
+    else:
+        out_path = Path(__file__).with_name("bench_int64_kernel_results.md")
+    out_path.write_text(md)
+    print(f"\nwrote {out_path}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/integrations/kernels/scattermoe_lora/bench_int64_kernel_results.md b/tests/integrations/kernels/scattermoe_lora/bench_int64_kernel_results.md
new file mode 100644
index 0000000000..3847039944
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/bench_int64_kernel_results.md
@@ -0,0 +1,15 @@
+# scatter2scatter INT64_INDICES bench
+
+GPU: **NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition**
+
+Median of 10 iters, 3 warmup. `top_k=8`, dtype=bf16, 128 experts.
+
+`auto_int64` is the wrapper's auto-dispatch verdict from `_needs_int64_indices`. At overflow shapes the int32 path is silently incorrect (the multiplication wraps mid-buffer), so only the int64 timing is reported.
+
+| Shape | T | L_scattered | out elems | auto_int64 | int32 ms | int64 ms | int64 vs int32 (%) |
+|---|---|---|---|---|---|---|---|
+| small | 8192 | 65536 | 1.34e+08 | False | 2.699 | 2.704 | +0.2 |
+| medium | 128000 | 1024000 | 2.10e+09 | False | 40.126 | 40.790 | +1.7 |
+| overflow_524k_s16 | 32768 | 262144 | 4.29e+09 | True | — | 80.105 | — |
+
+Acceptance: ≤5% regression on the int32 fast path at small/medium shapes (the auto-dispatch picks int32 there, so this row characterises the JIT overhead of having an int64 variant available).
diff --git a/tests/integrations/kernels/scattermoe_lora/bench_mxfp4.py b/tests/integrations/kernels/scattermoe_lora/bench_mxfp4.py
new file mode 100644
index 0000000000..ebeaa589dd
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/bench_mxfp4.py
@@ -0,0 +1,577 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+ScatterMoE LoRA — MXFP4 forward + backward benchmark.
+
+Runs three configurations on a representative DeepSeek-V4-style MoE shape
+(E=128, K=2048, N=1024, top_k=8, batch×seq=4096) and reports tokens/s,
+peak GPU memory, and effective HBM bandwidth for each:
+
+  * **bf16 baseline**: full-precision bf16 experts, no MX.
+  * **Strategy A**: torchao MXTensor experts, selective dequant to bf16.
+  * **Strategy B**: torchao MXTensor experts, fused MX dequant in Triton.
+
+Run from the repo root:
+
+  python tests/integrations/kernels/scattermoe_lora/bench_mxfp4.py
+
+A markdown table is printed to stdout and written to
+``bench_mxfp4_results.md`` next to this script.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import subprocess
+from pathlib import Path
+from typing import Callable
+
+import torch
+from torchao.prototype.mx_formats.mx_tensor import MXTensor
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+    selective_mx_weights_fwd,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+    flatten_sort_count,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_linear_lora import (
+    parallel_linear_lora,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (
+    get_active_experts,
+    remap_expert_indices,
+    selective_expert_weights,
+    selective_lora_weights,
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+def gpu_name() -> str:
+    try:
+        out = subprocess.check_output(["nvidia-smi", "-L"], text=True).strip()
+        # First GPU line: "GPU 0: NAME (UUID: ...)"
+        first = out.splitlines()[0]
+        if ":" in first:
+            after_colon = first.split(":", 1)[1].strip()
+            return after_colon.split("(", 1)[0].strip()
+        return first
+    except Exception:
+        return torch.cuda.get_device_name(0)
+
+
+def gpu_hbm_bandwidth_gbps() -> float | None:
+    """Rough peak HBM BW for utilization %, looked up by name. ``None`` if
+    unknown — printed as N/A."""
+    name = gpu_name().lower()
+    # Approximate datasheet peaks (GB/s). Order matters — more-specific
+    # patterns first so a "rtx pro 6000 blackwell" doesn't match
+    # "rtx 6000 ada".
+    table = [
+        (("rtx", "6000", "blackwell"), 1792.0),
+        (("rtx", "6000", "ada"), 960.0),
+        (("rtx", "5090"), 1792.0),
+        (("rtx", "4090"), 1008.0),
+        (("h200",), 4800.0),
+        (("h100",), 3350.0),
+        (("a100",), 2039.0),
+        (("a40",), 696.0),
+        (("a6000",), 768.0),
+        (("l40",), 864.0),
+        (("l4",), 300.0),
+        (("b200",), 8000.0),
+        (("mi300x",), 5300.0),
+    ]
+    for keys, bw in table:
+        if all(k in name for k in keys):
+            return bw
+    return None
+
+
+@torch.no_grad()
+def _setup_bf16(E, K, N, top_k, M, rank):
+    torch.manual_seed(0)
+    W = torch.randn(E, N, K, device=DEVICE, dtype=DTYPE) * (1.0 / K**0.5)
+    W_kernel = W.transpose(2, 1).contiguous()  # [E, K, N]
+    return W, W_kernel
+
+
+def _setup_mx(W_natural, chunk: int = 8):
+    """Make a torchao MXFP4 ``MXTensor`` from the bf16 weight.
+
+    ``MXTensor.to_mx`` materializes an fp32 working tensor internally; for
+    large [E, N, K] weights this transient can spike setup-time GPU memory
+    well beyond the final quantized footprint. To keep the bench runnable
+    on a shared GPU, quantize ``chunk`` experts at a time and stitch the
+    qdata/scale shards into a single MXTensor.
+    """
+    if W_natural.shape[0] <= chunk:
+        return MXTensor.to_mx(
+            W_natural, elem_dtype=torch.float4_e2m1fn_x2, block_size=32
+        )
+    qdata_parts = []
+    scale_parts = []
+    template = None
+    for i in range(0, W_natural.shape[0], chunk):
+        piece = W_natural[i : i + chunk].contiguous()
+        mx_chunk = MXTensor.to_mx(
+            piece, elem_dtype=torch.float4_e2m1fn_x2, block_size=32
+        )
+        qdata_parts.append(mx_chunk.qdata)
+        scale_parts.append(mx_chunk.scale)
+        if template is None:
+            template = mx_chunk
+    qdata = torch.cat(qdata_parts, dim=0)
+    scale = torch.cat(scale_parts, dim=0)
+    assert template is not None  # set on the first loop iter (loop body runs ≥ once)
+    return MXTensor(
+        qdata,
+        scale,
+        template.elem_dtype,
+        template.block_size,
+        template.orig_dtype,
+        template.kernel_preference,
+        template.act_quant_kwargs,
+        template.is_swizzled_scales,
+    )
+
+
+def _routing(M, E, top_k, seed=1, mode="dense"):
+    """Generate token→expert routing.
+
+    ``mode="dense"`` uses per-token random logits; for moderate E and top_k
+    this leaves nearly every expert active and exercises the kernels' full-load
+    case. ``mode="sparse"`` injects a strong shared bias so the same handful of
+    experts dominates the topk across all tokens — modelling realistic MoE
+    routing where only a small fraction of experts is active per step.
+    ``mode="balanced"`` models a load-balance-regularized router (aux-loss /
+    z-loss trained): per-token logits = N(0, 1) noise + small per-expert
+    bias N(0, 0.5). At large M this yields approximately balanced expert
+    usage; at small M only a fraction of experts gets hit and which experts
+    are active varies with seed/M — i.e. the seqlen → active-expert-count
+    curve that drives the A-vs-B crossover.
+    """
+    torch.manual_seed(seed)
+    if mode == "dense":
+        logits = torch.randn(M, E, device=DEVICE)
+    elif mode == "sparse":
+        shared = torch.randn(E, device=DEVICE) * 5.0
+        noise = torch.randn(M, E, device=DEVICE) * 0.1
+        logits = shared.unsqueeze(0) + noise
+    elif mode == "balanced":
+        bias = torch.randn(E, device=DEVICE) * 0.5
+        noise = torch.randn(M, E, device=DEVICE)
+        logits = noise + bias.unsqueeze(0)
+    else:
+        raise ValueError(f"unknown routing mode: {mode}")
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
+    sei, ssi, eo = flatten_sort_count(top_idx, E)
+    return sei, ssi, eo, top_idx
+
+
+def _lora(E, K, N, rank, seed=2):
+    torch.manual_seed(seed)
+    A = torch.randn(rank * E, K, device=DEVICE, dtype=DTYPE) * 0.01
+    B = torch.randn(N, rank * E, device=DEVICE, dtype=DTYPE) * 0.01
+    return A, B
+
+
+class _MockExperts:
+    def __init__(self, p):
+        self.gate_up_proj = p
+
+
+# ---------------------------------------------------------------------------
+# Three benchmark runners — each takes a fresh `x` and returns (output, fn_grad)
+# ---------------------------------------------------------------------------
+
+
+# Runners reuse the same leaf tensors across timed iters (x, lora A/B) and
+# zero ``.grad`` to None at the top of each call. The previous per-iter
+# ``.clone()`` + ``requires_grad_(True)`` was setup cost — not kernel cost —
+# and biased the timing especially on small shapes. Gradient accumulation
+# is avoided by setting ``.grad = None`` (faster than ``.zero_()``), so the
+# autograd graph each iter is fresh but the leaf buffers are not reallocated.
+
+
+def make_runner_bf16(W_kernel, lora_A, lora_B, sei, ssi, eo, top_k, scaling):
+    A = lora_A.detach().clone().requires_grad_(True)
+    B = lora_B.detach().clone().requires_grad_(True)
+
+    def run(x):
+        x.grad = None
+        A.grad = None
+        B.grad = None
+        out = parallel_linear_lora(
+            x,
+            W_kernel,
+            top_k,
+            sei,
+            ssi,
+            eo,
+            lora_A=A,
+            lora_B=B,
+            scaling=scaling,
+            use_fused_dX=True,
+            use_fused_gather=True,
+        )
+        out.sum().backward()
+        return out
+
+    return run
+
+
+def make_runner_strategy_a(mx, lora_A, lora_B, sei, ssi, eo, top_k, scaling, E):
+    experts = _MockExperts(mx)
+    A = lora_A.detach().clone().requires_grad_(True)
+    B = lora_B.detach().clone().requires_grad_(True)
+
+    def run(x):
+        x.grad = None
+        A.grad = None
+        B.grad = None
+        active = get_active_experts(sei, E)
+        remapped, compact_off = remap_expert_indices(sei, eo, active, E)
+        W_compact = (
+            selective_expert_weights(experts, "gate_up_proj", active)
+            .transpose(2, 1)
+            .contiguous()
+        )
+        A_c, B_c = selective_lora_weights(A, B, active, E)
+        out = parallel_linear_lora(
+            x,
+            W_compact,
+            top_k,
+            remapped,
+            ssi,
+            compact_off,
+            lora_A=A_c,
+            lora_B=B_c,
+            scaling=scaling,
+            use_fused_dX=True,
+            use_fused_gather=True,
+        )
+        out.sum().backward()
+        return out
+
+    return run
+
+
+def make_runner_strategy_b(mx, lora_A, lora_B, sei, ssi, eo, top_k, scaling, E):
+    A = lora_A.detach().clone().requires_grad_(True)
+    B = lora_B.detach().clone().requires_grad_(True)
+
+    def run(x):
+        x.grad = None
+        A.grad = None
+        B.grad = None
+        active = get_active_experts(sei, E)
+        remapped, compact_off = remap_expert_indices(sei, eo, active, E)
+        mx_active = selective_mx_weights_fwd(mx, active)
+        A_c, B_c = selective_lora_weights(A, B, active, E)
+        out = parallel_linear_lora(
+            x,
+            mx_active,
+            top_k,
+            remapped,
+            ssi,
+            compact_off,
+            lora_A=A_c,
+            lora_B=B_c,
+            scaling=scaling,
+        )
+        out.sum().backward()
+        return out
+
+    return run
+
+
+# ---------------------------------------------------------------------------
+# Timing harness
+# ---------------------------------------------------------------------------
+
+
+def bench(fn: Callable, x_template: torch.Tensor, warmup: int, iters: int) -> dict:
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    # Allocate the input leaf tensor once outside the timed window. Runners
+    # reset ``.grad = None`` per iter; the underlying buffer is reused.
+    x = x_template.detach().clone().requires_grad_(True)
+    try:
+        # Warmup
+        for _ in range(warmup):
+            fn(x)
+        torch.cuda.synchronize()
+
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(iters):
+            fn(x)
+        end.record()
+        torch.cuda.synchronize()
+    except torch.cuda.OutOfMemoryError:
+        torch.cuda.empty_cache()
+        return {"ms_per_iter": float("nan"), "peak_mem_mb": float("nan"), "oom": True}
+    elapsed_ms = start.elapsed_time(end)
+    peak_mem_mb = torch.cuda.max_memory_allocated() / 1024 / 1024
+    return {
+        "ms_per_iter": elapsed_ms / iters,
+        "peak_mem_mb": peak_mem_mb,
+        "oom": False,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--E", type=int, default=128)
+    parser.add_argument("--K", type=int, default=2048)
+    parser.add_argument("--N", type=int, default=1024)
+    parser.add_argument("--top_k", type=int, default=8)
+    parser.add_argument("--M", type=int, default=4096, help="batch * seq")
+    parser.add_argument("--rank", type=int, default=16)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument(
+        "--routing-mode",
+        choices=("dense", "sparse", "balanced"),
+        default="dense",
+        help=(
+            "dense: per-token random logits (~all experts active). "
+            "sparse: shared bias + small per-token noise so the same ~top_k "
+            "experts dominate routing across all tokens. "
+            "balanced: per-token N(0,1) noise + small N(0,0.5) per-expert bias "
+            "— mimics a load-balance-regularized router; active-expert count "
+            "grows with M."
+        ),
+    )
+    parser.add_argument(
+        "--M-sweep",
+        dest="M_sweep",
+        default=None,
+        help=(
+            "Comma-separated list of M values, e.g. '256,1024,4096,16384'. "
+            "When set, --M is ignored; the bench runs once per M with the "
+            "selected routing mode and emits a single combined section."
+        ),
+    )
+    parser.add_argument(
+        "--append",
+        action="store_true",
+        help="Append the new table to bench_mxfp4_results.md instead of overwriting it.",
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda",
+        help="CUDA device, e.g. 'cuda', 'cuda:0', 'cuda:1'.",
+    )
+    args = parser.parse_args()
+
+    global DEVICE
+    DEVICE = args.device
+    torch.cuda.set_device(torch.device(DEVICE))
+
+    E, K, N, top_k, rank = args.E, args.K, args.N, args.top_k, args.rank
+
+    if args.M_sweep:
+        M_values = [int(s.strip()) for s in args.M_sweep.split(",") if s.strip()]
+    else:
+        M_values = [args.M]
+
+    print(f"GPU: {gpu_name()}")
+    print(
+        f"Shape: E={E}, K={K}, N={N}, top_k={top_k}, rank={rank}, "
+        f"M={M_values if args.M_sweep else M_values[0]}"
+    )
+    print(f"Iters: {args.warmup} warmup + {args.iters} timed")
+    print()
+
+    # Build dense bf16 weights + MX-quantize once — weights are independent of M.
+    W_natural, W_kernel = _setup_bf16(E, K, N, top_k, M_values[0], rank)
+    mx = _setup_mx(W_natural)
+    # W_natural is only used to build the MX tensor; W_kernel feeds bf16 paths.
+    # Free it eagerly so the dequant transient fits on memory-constrained GPUs.
+    del W_natural
+    torch.cuda.empty_cache()
+    lora_A, lora_B = _lora(E, K, N, rank)
+    scaling = 0.5
+    peak_bw = gpu_hbm_bandwidth_gbps()
+
+    per_M = []  # list of (M, e_active, results)
+    for M in M_values:
+        sei, ssi, eo, _top_idx = _routing(M, E, top_k, mode=args.routing_mode)
+        x = torch.randn(M, K, device=DEVICE, dtype=DTYPE)
+
+        # Estimate bytes read per iter for HBM BW utilization:
+        #   bf16 W:     E_active * K * N * 2
+        #   MX  W:      E_active * (K*N/2 + K*N/32)
+        # X is M*K*2 bytes. We ignore LoRA traffic (tiny relative to W).
+        num_tokens = M * top_k
+        e_active = int(get_active_experts(sei, E).numel())
+        bytes_bf16 = e_active * K * N * 2 + M * K * 2
+        bytes_mx = e_active * (K * N // 2 + K * N // 32) + M * K * 2
+
+        runners = {
+            "bf16 baseline": (
+                make_runner_bf16(
+                    W_kernel, lora_A, lora_B, sei, ssi, eo, top_k, scaling
+                ),
+                bytes_bf16,
+            ),
+            "Strategy A (selective dequant)": (
+                make_runner_strategy_a(
+                    mx, lora_A, lora_B, sei, ssi, eo, top_k, scaling, E
+                ),
+                bytes_bf16,  # post-dequant the kernel still reads bf16
+            ),
+            "Strategy B (fused MX)": (
+                make_runner_strategy_b(
+                    mx, lora_A, lora_B, sei, ssi, eo, top_k, scaling, E
+                ),
+                bytes_mx,
+            ),
+        }
+
+        results = []
+        for name, (fn, bytes_per_iter) in runners.items():
+            r = bench(fn, x, args.warmup, args.iters)
+            if r.get("oom"):
+                tps = float("nan")
+                bw = float("nan")
+                bw_pct = float("nan")
+            else:
+                tps = num_tokens / (r["ms_per_iter"] / 1000.0)
+                bw = (bytes_per_iter / 1e9) / (r["ms_per_iter"] / 1000.0)
+                bw_pct = (bw / peak_bw * 100.0) if peak_bw else float("nan")
+            results.append(
+                dict(
+                    name=name,
+                    ms_per_iter=r["ms_per_iter"],
+                    tokens_per_s=tps,
+                    peak_mem_mb=r["peak_mem_mb"],
+                    hbm_gbps=bw,
+                    hbm_pct=bw_pct,
+                    oom=r.get("oom", False),
+                )
+            )
+        per_M.append((M, e_active, results))
+
+    section_lines = []
+    if args.M_sweep:
+        section_lines.append(
+            f"## Routing mode: {args.routing_mode} — M sweep — {gpu_name()}"
+        )
+        section_lines.append("")
+        section_lines.append(f"- **GPU**: {gpu_name()}")
+        section_lines.append(
+            f"- **Base shape**: E={E}, K={K}, N={N}, top_k={top_k}, rank={rank}"
+        )
+        section_lines.append(f"- **M values**: {', '.join(str(m) for m in M_values)}")
+        section_lines.append(
+            f"- **Iters**: {args.warmup} warmup + {args.iters} timed, fwd+bwd per iter"
+        )
+        if peak_bw:
+            section_lines.append(f"- **HBM peak (datasheet)**: {peak_bw:.0f} GB/s")
+        section_lines.append("")
+        section_lines.append("### Summary (ms/iter, fwd+bwd)")
+        section_lines.append("")
+        section_lines.append(
+            "| M | active / E | bf16 ms | Strategy A ms | Strategy B ms | winner (A vs B) |"
+        )
+        section_lines.append("| ---: | ---: | ---: | ---: | ---: | :---: |")
+
+        def _fmt_ms(r):
+            return "OOM" if r["oom"] else f"{r['ms_per_iter']:.2f}"
+
+        for M, e_active, results in per_M:
+            by_name = {r["name"]: r for r in results}
+            a, b, bf = (
+                by_name["Strategy A (selective dequant)"],
+                by_name["Strategy B (fused MX)"],
+                by_name["bf16 baseline"],
+            )
+            if a["oom"] and b["oom"]:
+                winner = "—"
+            elif a["oom"]:
+                winner = "B"
+            elif b["oom"]:
+                winner = "A"
+            else:
+                winner = "A" if a["ms_per_iter"] < b["ms_per_iter"] else "B"
+            section_lines.append(
+                f"| {M} | {e_active}/{E} ({e_active / E:.2f}) | "
+                f"{_fmt_ms(bf)} | {_fmt_ms(a)} | {_fmt_ms(b)} | {winner} |"
+            )
+        section_lines.append("")
+        for M, e_active, results in per_M:
+            section_lines.append(
+                f"### M={M} (active experts = {e_active} / {E}, "
+                f"num_active/E = {e_active / E:.3f})"
+            )
+            section_lines.append("")
+            section_lines.append(
+                "| Config | ms/iter | tokens/s | peak mem (MB) | HBM GB/s | HBM % |"
+            )
+            section_lines.append("| --- | ---: | ---: | ---: | ---: | ---: |")
+            for r in results:
+                if r["oom"]:
+                    section_lines.append(
+                        f"| {r['name']} | OOM | OOM | OOM | OOM | OOM |"
+                    )
+                    continue
+                hbm_pct = (
+                    f"{r['hbm_pct']:.1f}" if not math.isnan(r["hbm_pct"]) else "N/A"
+                )
+                section_lines.append(
+                    f"| {r['name']} | {r['ms_per_iter']:.2f} | "
+                    f"{r['tokens_per_s']:.0f} | {r['peak_mem_mb']:.1f} | "
+                    f"{r['hbm_gbps']:.1f} | {hbm_pct} |"
+                )
+            section_lines.append("")
+    else:
+        M, e_active, results = per_M[0]
+        section_lines.append(f"## Routing mode: {args.routing_mode} — {gpu_name()}")
+        section_lines.append("")
+        section_lines.append(f"- **GPU**: {gpu_name()}")
+        section_lines.append(
+            f"- **Shape**: E={E}, K={K}, N={N}, top_k={top_k}, M={M}, rank={rank} "
+            f"(active experts = {e_active})"
+        )
+        section_lines.append(
+            f"- **Iters**: {args.warmup} warmup + {args.iters} timed, fwd+bwd per iter"
+        )
+        if peak_bw:
+            section_lines.append(f"- **HBM peak (datasheet)**: {peak_bw:.0f} GB/s")
+        section_lines.append("")
+        section_lines.append(
+            "| Config | ms/iter | tokens/s | peak mem (MB) | HBM GB/s | HBM % |"
+        )
+        section_lines.append("| --- | ---: | ---: | ---: | ---: | ---: |")
+        for r in results:
+            hbm_pct = f"{r['hbm_pct']:.1f}" if not math.isnan(r["hbm_pct"]) else "N/A"
+            section_lines.append(
+                f"| {r['name']} | {r['ms_per_iter']:.2f} | {r['tokens_per_s']:.0f} | "
+                f"{r['peak_mem_mb']:.1f} | {r['hbm_gbps']:.1f} | {hbm_pct} |"
+            )
+    section_md = "\n".join(section_lines).rstrip() + "\n"
+
+    out_path = Path(__file__).resolve().parent / "bench_mxfp4_results.md"
+    if args.append and out_path.exists():
+        existing = out_path.read_text().rstrip() + "\n\n"
+        md = existing + section_md
+    else:
+        md = "# ScatterMoE LoRA — MXFP4 benchmark\n\n" + section_md
+
+    print(section_md)
+    out_path.write_text(md)
+    print(f"\nResults written to {out_path}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/integrations/kernels/scattermoe_lora/bench_mxfp4_results.md b/tests/integrations/kernels/scattermoe_lora/bench_mxfp4_results.md
new file mode 100644
index 0000000000..4df526f6fb
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/bench_mxfp4_results.md
@@ -0,0 +1,106 @@
+# ScatterMoE LoRA — MXFP4 benchmark
+
+## Routing mode: dense — NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
+
+- **GPU**: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
+- **Shape**: E=128, K=2048, N=1024, top_k=8, M=4096, rank=16 (active experts = 128)
+- **Iters**: 10 warmup + 50 timed, fwd+bwd per iter
+- **HBM peak (datasheet)**: 1792 GB/s
+
+| Config | ms/iter | tokens/s | peak mem (MB) | HBM GB/s | HBM % |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| bf16 baseline | 5.25 | 6244998 | 1252.8 | 105.5 | 5.9 |
+| Strategy A (selective dequant) | 30.57 | 1071778 | 8557.3 | 18.1 | 1.0 |
+| Strategy B (fused MX) | 12.24 | 2677582 | 1425.3 | 13.0 | 0.7 |
+
+## Routing mode: sparse — NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
+
+- **GPU**: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
+- **Shape**: E=256, K=2048, N=1024, top_k=8, M=4096, rank=16 (active experts = 10)
+- **Iters**: 10 warmup + 50 timed, fwd+bwd per iter
+- **HBM peak (datasheet)**: 1792 GB/s
+
+| Config | ms/iter | tokens/s | peak mem (MB) | HBM GB/s | HBM % |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| bf16 baseline | 6.55 | 5006027 | 1960.8 | 9.0 | 0.5 |
+| Strategy A (selective dequant) | 5.75 | 5695789 | 2059.9 | 10.2 | 0.6 |
+| Strategy B (fused MX) | 8.95 | 3661270 | 1997.8 | 3.1 | 0.2 |
+
+## Routing mode: balanced — M sweep — NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
+
+- **GPU**: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
+- **Base shape**: E=256, K=2048, N=1024, top_k=8, rank=16
+- **M values**: 256, 1024, 4096, 16384
+- **Iters**: 10 warmup + 50 timed, fwd+bwd per iter
+- **HBM peak (datasheet)**: 1792 GB/s
+
+### Summary (ms/iter, fwd+bwd)
+
+| M | active / E | bf16 ms | Strategy A ms | Strategy B ms | winner (A vs B) |
+| ---: | ---: | ---: | ---: | ---: | :---: |
+| 256 | 215/256 (0.84) | 2.99 | OOM | 8.24 | B |
+| 1024 | 251/256 (0.98) | 3.43 | OOM | 10.74 | B |
+| 4096 | 255/256 (1.00) | 6.56 | OOM | 16.50 | B |
+| 16384 | 256/256 (1.00) | 24.15 | OOM | 46.56 | B |
+
+### M=256 (active experts = 215 / 256, num_active/E = 0.840)
+
+| Config | ms/iter | tokens/s | peak mem (MB) | HBM GB/s | HBM % |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| bf16 baseline | 2.99 | 685596 | 1686.0 | 302.2 | 16.9 |
+| Strategy A (selective dequant) | OOM | OOM | OOM | OOM | OOM |
+| Strategy B (fused MX) | 8.24 | 248639 | 1954.9 | 29.2 | 1.6 |
+
+### M=1024 (active experts = 251 / 256, num_active/E = 0.980)
+
+| Config | ms/iter | tokens/s | peak mem (MB) | HBM GB/s | HBM % |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| bf16 baseline | 3.43 | 2389143 | 1744.2 | 308.3 | 17.2 |
+| Strategy A (selective dequant) | OOM | OOM | OOM | OOM | OOM |
+| Strategy B (fused MX) | 10.74 | 762567 | 2058.1 | 26.4 | 1.5 |
+
+### M=4096 (active experts = 255 / 256, num_active/E = 0.996)
+
+| Config | ms/iter | tokens/s | peak mem (MB) | HBM GB/s | HBM % |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| bf16 baseline | 6.56 | 4994760 | 1960.8 | 165.6 | 9.2 |
+| Strategy A (selective dequant) | OOM | OOM | OOM | OOM | OOM |
+| Strategy B (fused MX) | 16.50 | 1985884 | 2280.0 | 18.2 | 1.0 |
+
+### M=16384 (active experts = 256 / 256, num_active/E = 1.000)
+
+| Config | ms/iter | tokens/s | peak mem (MB) | HBM GB/s | HBM % |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| bf16 baseline | 24.15 | 5427073 | 2827.0 | 47.2 | 2.6 |
+| Strategy A (selective dequant) | OOM | OOM | OOM | OOM | OOM |
+| Strategy B (fused MX) | 46.56 | 2814943 | 3149.0 | 7.6 | 0.4 |
+
+### Notes
+
+- **Strategy A OOMs at all M** under load-balanced routing at E=256 because
+  the torchao MXTensor dequant path materializes several full-shape fp32/int32
+  unpack buffers (~12 GiB combined for [256, 1024, 2048] at fp4 → fp32) while
+  vLLM colocated on this workstation pins ~88 GB of HBM, leaving only ~14 GB
+  free. Extrapolating from the dense E=128 case above (Strategy A peak
+  ~8.6 GB at 128 active experts), the E=256 / 256-active dequant peak would
+  be ~17 GB — over the available headroom.
+- **Active-expert count is essentially E at every sampled M.** Under a
+  load-balance-regularized router (per-token N(0,1) noise + N(0,0.5) per-expert
+  bias), `E[active] ≈ E · (1 − (1 − top_k/E)^M)`. With E=256 / top_k=8 this
+  yields ≥ 215 unique experts even at M=256 and saturates at 256 by M ≈ 16K.
+  Balanced routing therefore does **not** generate a low-active regime at
+  these token counts — i.e. the A-vs-B crossover does not appear in this
+  sweep; B wins by default because A does not fit.
+- **B vs bf16:** Strategy B is consistently 1.9–2.9× slower than the bf16
+  baseline (similar to the dense E=128 ratio of ~2.3×). HBM utilization for
+  both is modest (B 0.4–1.6 %, bf16 2.6–17.2 %), suggesting the kernels are
+  compute- or scheduling-bound for these shapes, not bandwidth-bound.
+- **Where the A-vs-B crossover lives, by theory:** Strategy A is preferred
+  when `num_active / E` is small enough that the dequant cost is offset by
+  the cheaper bf16 matmul — the prior `sparse` row (10/256 active, A=5.75 ms
+  vs B=8.95 ms) sits in that regime. Strategy B is preferred near
+  `num_active / E ≈ 1`, where dequant of all experts dominates. The threshold
+  between the two — somewhere in the 10/256 to 215/256 band — is **not
+  observable from the balanced-router setting**; eliciting it would need an
+  M smaller than 256, a synthetic deliberately-sparse router, or freeing the
+  vLLM GPU and rerunning at E=256.
diff --git a/tests/integrations/kernels/scattermoe_lora/conftest.py b/tests/integrations/kernels/scattermoe_lora/conftest.py
new file mode 100644
index 0000000000..90e18e6a93
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/conftest.py
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Treat CUDA OOM as a skip for tests in this directory.
+
+When the suite runs under ``pytest-xdist``, multiple workers contend for the
+same physical GPU's memory budget. A test that fits comfortably in isolation
+can OOM purely because peer workers are already holding most of VRAM. That's
+an environmental race, not a code defect, so converting it to a skip keeps
+mixed-GPU CI green without masking real regressions (a real correctness bug
+surfaces as an assert/exception, not as ``torch.OutOfMemoryError``).
+
+We hook ``pytest_runtest_call`` rather than using an autouse fixture because
+pytest captures the test exception before re-entering the fixture's
+generator — the fixture's ``try/except`` around ``yield`` never sees it.
+"""
+
+from __future__ import annotations
+
+import gc
+
+import pytest
+import torch
+
+
+def _cuda_oom_types() -> tuple[type[BaseException], ...]:
+    types: list[type[BaseException]] = []
+    if hasattr(torch, "OutOfMemoryError"):
+        types.append(torch.OutOfMemoryError)
+    cuda_oom = getattr(torch.cuda, "OutOfMemoryError", None)
+    if cuda_oom is not None and cuda_oom not in types:
+        types.append(cuda_oom)
+    return tuple(types) or (RuntimeError,)
+
+
+_OOM = _cuda_oom_types()
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_call(item):
+    outcome = yield
+    excinfo = outcome.excinfo
+    if excinfo is None:
+        return
+    exc_val = excinfo[1]
+    if isinstance(exc_val, _OOM):
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        outcome.force_exception(
+            pytest.skip.Exception(
+                f"skipping on CUDA OOM (likely xdist worker contention): {exc_val}",
+                _use_item_location=True,
+            )
+        )
diff --git a/tests/integrations/kernels/scattermoe_lora/test_bnb_experts_forward.py b/tests/integrations/kernels/scattermoe_lora/test_bnb_experts_forward.py
new file mode 100644
index 0000000000..37c379f0a4
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_bnb_experts_forward.py
@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""bnb-4bit MoE expert LoRA: both the default 1-launch (``moe_bnb_fast``) path and the
+chunked-dequant fallback must match a full-dequant reference in forward AND gradients.
+
+The experts are stored as bnb 4-bit (the same parametrization ``quantize_moe_experts``
+installs at load via ``replace_parameter_4bit``). The reference dequantizes those exact
+4-bit weights to bf16 and runs the standard scattermoe path, so the only differences are
+GEMM ordering / bf16 rounding between the two bnb paths and the reference.
+
+Also covers the divisibility guard in ``_selective_dequant_bnb4`` (F3) and the
+non-persistent Marlin workspace buffer (F4). All tests are CUDA + bitsandbytes gated.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+bnb_parametrize = pytest.importorskip(
+    "bitsandbytes.nn.parametrize", reason="bitsandbytes required"
+)
+
+import axolotl.integrations.kernels.libs.scattermoe_lora.experts as ex  # noqa: E402
+from axolotl.integrations.kernels.libs.scattermoe_lora.chunked_bnb import (  # noqa: E402
+    set_bnb_fast,
+    set_chunk_size_override,
+    set_layer_gc_active,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (  # noqa: E402
+    scattermoe_experts_forward,
+)
+
+DEV = "cuda"
+E, H, IM, N, K, R, SC = 8, 256, 128, 32, 2, 8, 0.5
+
+
+def _bf16_module(gu, dn):
+    """Plain bf16 experts module (routes through the standard dequant path)."""
+    return SimpleNamespace(
+        num_experts=E,
+        gate_up_proj=gu,
+        down_proj=dn,
+        act_fn=torch.nn.functional.silu,
+        is_transposed=False,
+        is_concatenated=True,
+        has_bias=False,
+        has_gate=True,
+    )
+
+
+class _BnbExperts(torch.nn.Module):
+    """Real nn.Module so torch parametrize (bnb 4-bit) can attach to the params."""
+
+    def __init__(self, gu, dn):
+        super().__init__()
+        self.num_experts = E
+        self.act_fn = torch.nn.functional.silu
+        self.is_transposed = False
+        self.is_concatenated = True
+        self.has_bias = False
+        self.has_gate = True
+        self.gate_up_proj = torch.nn.Parameter(gu, requires_grad=False)
+        self.down_proj = torch.nn.Parameter(dn, requires_grad=False)
+        bnb_parametrize.replace_parameter_4bit(
+            self, "gate_up_proj", compress_statistics=True, quant_type="nf4"
+        )
+        bnb_parametrize.replace_parameter_4bit(
+            self, "down_proj", compress_statistics=True, quant_type="nf4"
+        )
+
+
+def _rel(a, b):
+    return (a - b).float().abs().max().item() / max(b.float().abs().max().item(), 1e-6)
+
+
+def _run(module, lora, idx, w, grad, monkeypatch):
+    monkeypatch.setattr(ex, "_has_peft_wrapper", lambda s: True)
+    monkeypatch.setattr(
+        ex,
+        "_unwrap_experts_lora",
+        lambda s: (s, (lora[0], lora[1], SC), (lora[2], lora[3], SC)),
+    )
+    for t in lora:
+        t.grad = None
+    x = torch.randn(
+        N,
+        H,
+        device=DEV,
+        dtype=torch.bfloat16,
+        generator=torch.Generator(DEV).manual_seed(7),
+    ).requires_grad_(True)
+    out = scattermoe_experts_forward(module, x, idx, w)
+    out.backward(grad)
+    return out.detach(), x.grad.detach(), [t.grad.clone() for t in lora]
+
+
+@pytest.mark.parametrize("fast", [True, False])
+def test_bnb_moe_lora_matches_dequant(fast, monkeypatch):
+    """moe_bnb_fast={True,False} both match the full-dequant bf16 reference (fwd + grads)."""
+    pytest.importorskip("triton")
+    dt = torch.bfloat16
+    g = torch.Generator(device=DEV).manual_seed(0)
+    gu = torch.randn(E, 2 * IM, H, device=DEV, dtype=dt, generator=g) * 0.1
+    dn = torch.randn(E, H, IM, device=DEV, dtype=dt, generator=g) * 0.1
+
+    bnb_mod = _BnbExperts(gu, dn)
+    # Reference uses the EXACT dequantized 4-bit weights so only kernel numerics differ.
+    gu_deq = bnb_mod.gate_up_proj.detach().contiguous()
+    dn_deq = bnb_mod.down_proj.detach().contiguous()
+
+    def mk(*s):
+        return (
+            torch.randn(*s, device=DEV, dtype=dt, generator=g) * 0.05
+        ).requires_grad_(True)
+
+    lora = [mk(R * E, H), mk(2 * IM, R * E), mk(R * E, IM), mk(H, R * E)]
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    w = torch.rand(N, K, device=DEV, generator=g).to(dt)
+    grad = torch.randn(N, H, device=DEV, dtype=dt, generator=g)
+
+    active = torch.unique(idx.reshape(-1))
+    row = (active.long()[:, None] * R + torch.arange(R, device=DEV)[None, :]).reshape(
+        -1
+    )
+
+    set_layer_gc_active(False)
+    set_chunk_size_override(None)
+    try:
+        set_bnb_fast(fast)
+        out_bnb, dx_bnb, gl_bnb = _run(bnb_mod, lora, idx, w, grad, monkeypatch)
+    finally:
+        set_bnb_fast(None)
+
+    out_ref, dx_ref, gl_ref = _run(
+        _bf16_module(gu_deq, dn_deq), lora, idx, w, grad, monkeypatch
+    )
+
+    # forward output parity vs the full-dequant reference
+    assert torch.isfinite(out_bnb).all()
+    assert _rel(out_bnb, out_ref) < 6e-2, "forward output"
+    # input grad
+    assert torch.isfinite(dx_bnb).all()
+    assert _rel(dx_bnb, dx_ref) < 6e-2, "input grad"
+    for i, slc in (
+        (0, row),
+        (1, (slice(None), row)),
+        (2, row),
+        (3, (slice(None), row)),
+    ):
+        assert torch.isfinite(gl_bnb[i]).all()
+        assert _rel(gl_bnb[i][slc], gl_ref[i][slc]) < 6e-2, f"lora grad {i}"
+
+
+def test_bnb_fast_and_chunked_agree(monkeypatch):
+    """The two bnb paths should agree with each other (same math, different launch shape)."""
+    pytest.importorskip("triton")
+    dt = torch.bfloat16
+    g = torch.Generator(device=DEV).manual_seed(3)
+    gu = torch.randn(E, 2 * IM, H, device=DEV, dtype=dt, generator=g) * 0.1
+    dn = torch.randn(E, H, IM, device=DEV, dtype=dt, generator=g) * 0.1
+    bnb_mod = _BnbExperts(gu, dn)
+
+    def mk(*s):
+        return (
+            torch.randn(*s, device=DEV, dtype=dt, generator=g) * 0.05
+        ).requires_grad_(True)
+
+    lora = [mk(R * E, H), mk(2 * IM, R * E), mk(R * E, IM), mk(H, R * E)]
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    w = torch.rand(N, K, device=DEV, generator=g).to(dt)
+    grad = torch.randn(N, H, device=DEV, dtype=dt, generator=g)
+
+    set_layer_gc_active(False)
+    set_chunk_size_override(None)
+    try:
+        set_bnb_fast(True)
+        out_fast, dx_fast, gl_fast = _run(bnb_mod, lora, idx, w, grad, monkeypatch)
+        set_bnb_fast(False)
+        out_chunk, dx_chunk, gl_chunk = _run(bnb_mod, lora, idx, w, grad, monkeypatch)
+    finally:
+        set_bnb_fast(None)
+
+    assert _rel(out_fast, out_chunk) < 5e-2, "forward output"
+    assert _rel(dx_fast, dx_chunk) < 5e-2, "input grad"
+    for i in range(4):
+        assert _rel(gl_fast[i], gl_chunk[i]) < 5e-2, f"lora grad {i}"
+
+
+# --- F3: divisibility guard in _selective_dequant_bnb4 -------------------------------------
+
+
+@pytest.mark.parametrize(
+    "expert_shape",
+    [
+        (8, 8),  # numel 64: divisible by blocksize(64) and 2 -> selective slice path
+        (4, 4),  # numel 16: not divisible by blocksize(64) -> full-dequant fallback
+    ],
+)
+def test_selective_dequant_bnb4_matches_full(expert_shape):
+    """Selective dequant equals full-dequant+index in BOTH the divisible and fallback cases."""
+    pytest.importorskip("triton")
+    import bitsandbytes.functional as bnb_f
+
+    from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (
+        _selective_dequant_bnb4,
+    )
+
+    blocksize = 64
+    n_exp = 4
+    expert_numel = expert_shape[0] * expert_shape[1]
+    flat = torch.randn(n_exp * expert_numel, device=DEV, dtype=torch.bfloat16) * 0.1
+    packed, qs = bnb_f.quantize_4bit(
+        flat, blocksize=blocksize, quant_type="nf4", compress_statistics=False
+    )
+
+    full = bnb_f.dequantize_4bit(packed, qs).reshape(n_exp, *expert_shape)
+    active = torch.tensor([0, 2], device=DEV, dtype=torch.long)
+
+    out = _selective_dequant_bnb4(packed, qs, active, expert_shape)
+    torch.testing.assert_close(out, full[active], atol=2e-2, rtol=2e-2)
+
+
+# --- F4: Marlin non-expert workspace buffer must not enter state_dict ----------------------
+
+
+def test_marlin_nonexpert_workspace_not_persistent():
+    pytest.importorskip("torchao")
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16 import (
+        marlin_w4a16_available,
+    )
+
+    if not marlin_w4a16_available():
+        pytest.skip("Marlin W4A16 kernel not available on this device")
+
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.nonexpert_linear import (
+        MarlinW4A16Linear,
+    )
+
+    w = torch.randn(128, 256, device=DEV, dtype=torch.bfloat16) * 0.1
+    mod = MarlinW4A16Linear(w, bias=None)
+    sd = mod.state_dict()
+    assert "_workspace" not in sd, "scratch workspace leaked into state_dict"
+    assert "_workspace" in mod._non_persistent_buffers_set
+    # persistent buffers still present
+    assert "qweight" in sd and "scales" in sd
diff --git a/tests/integrations/kernels/scattermoe_lora/test_ep_sentinel_skip.py b/tests/integrations/kernels/scattermoe_lora/test_ep_sentinel_skip.py
new file mode 100644
index 0000000000..eeed72bbc6
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_ep_sentinel_skip.py
@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""scattermoe DeepEP local path skips ``-1`` sentinels instead of compute-and-mask.
+
+Oracle: on the SAME experts module, ``scattermoe_experts_forward_ep`` (raw -1-tagged
+routing) must match ``scattermoe_experts_forward`` with sentinels mapped to expert 0 /
+weight 0 -- output, dX, and LoRA dA/dB. Sentinel slots carry weight 0 so both compute
+the same math; the skip path just omits the wasted rows (and the load-imbalanced
+expert-0 bucket the masked path creates).
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+import axolotl.integrations.kernels.libs.scattermoe_lora.experts as ex  # noqa: E402
+from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (  # noqa: E402
+    scattermoe_experts_forward,
+    scattermoe_experts_forward_ep,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (  # noqa: E402
+    peft_lora_to_scattermoe,
+)
+
+DEV = "cuda"
+
+
+def _module(E, H, IM, dt, rank, monkeypatch):
+    self = SimpleNamespace(
+        num_experts=E,
+        gate_up_proj=torch.randn(E, 2 * IM, H, device=DEV, dtype=dt) * 0.02,
+        down_proj=torch.randn(E, H, IM, device=DEV, dtype=dt) * 0.02,
+        act_fn=torch.nn.functional.silu,
+        is_transposed=False,
+        is_concatenated=True,
+        has_bias=False,
+        has_gate=True,
+    )
+    lora = None
+    if rank:
+        A1 = (torch.randn(rank * E, H, device=DEV, dtype=dt) * 0.02).requires_grad_(
+            True
+        )
+        B1 = (
+            torch.randn(2 * IM, rank * E, device=DEV, dtype=dt) * 0.02
+        ).requires_grad_(True)
+        A2 = (torch.randn(rank * E, IM, device=DEV, dtype=dt) * 0.02).requires_grad_(
+            True
+        )
+        B2 = (torch.randn(H, rank * E, device=DEV, dtype=dt) * 0.02).requires_grad_(
+            True
+        )
+        gup = (*peft_lora_to_scattermoe(A1, B1, E, rank), 0.5)
+        dwn = (*peft_lora_to_scattermoe(A2, B2, E, rank), 0.5)
+        monkeypatch.setattr(ex, "_has_peft_wrapper", lambda m: True)
+        monkeypatch.setattr(ex, "_unwrap_experts_lora", lambda m: (m, gup, dwn))
+        lora = (A1, B1, A2, B2)
+    return self, lora
+
+
+def _routing(N, K, E, ep, dt):
+    g = torch.Generator(device=DEV).manual_seed(0)
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    remote = torch.rand(N, K, device=DEV, generator=g) < (1 - 1.0 / ep)
+    remote[:, 0] = False  # >=1 valid slot per received token (DeepEP guarantee)
+    idx = torch.where(remote, torch.full_like(idx, -1), idx)
+    w = torch.rand(N, K, device=DEV, generator=g).to(dt)
+    return idx, w
+
+
+@pytest.mark.parametrize("dt", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("rank", [0, 16])
+@pytest.mark.parametrize("ep", [2, 4])
+def test_ep_skip_matches_masked(dt, rank, ep, monkeypatch):
+    E, H, IM, N, K = 32, 512, 256, 512, 8
+    self, lora = _module(E, H, IM, dt, rank, monkeypatch)
+    idx, w = _routing(N, K, E, ep, dt)
+    safe_idx = torch.where(idx >= 0, idx, torch.zeros_like(idx))
+    safe_w = w * (idx >= 0).to(dt)
+    grad = torch.randn(N, H, device=DEV, dtype=dt)
+
+    def run(fn, ii, ww):
+        x = torch.randn(
+            N,
+            H,
+            device=DEV,
+            dtype=dt,
+            generator=torch.Generator(DEV).manual_seed(1),
+        ).requires_grad_(True)
+        if lora:
+            for t in lora:
+                t.grad = None
+        fn(self, x, ii, ww).backward(grad)
+        return x.grad.clone(), [t.grad.clone() for t in lora] if lora else []
+
+    # forward equality
+    with torch.no_grad():
+        x = torch.randn(N, H, device=DEV, dtype=dt)
+        o_m = scattermoe_experts_forward(self, x, safe_idx, safe_w)
+        o_s = scattermoe_experts_forward_ep(self, x, idx, w)
+    tol = 1e-4 if dt == torch.float32 else 4e-2
+    assert torch.allclose(o_s, o_m, rtol=tol, atol=tol), (o_s - o_m).abs().max().item()
+
+    dx_m, gl_m = run(scattermoe_experts_forward, safe_idx, safe_w)
+    dx_s, gl_s = run(scattermoe_experts_forward_ep, idx, w)
+    assert torch.allclose(dx_s, dx_m, rtol=tol, atol=tol), (
+        (dx_s - dx_m).abs().max().item()
+    )
+    for a, b in zip(gl_s, gl_m, strict=True):
+        assert torch.allclose(a, b, rtol=tol, atol=tol), (a - b).abs().max().item()
+
+
+def test_ep_skip_handles_all_sentinel_token():
+    """A row with every slot remote (-1) contributes nothing and must not crash."""
+    E, H, IM, N, K = 8, 256, 128, 16, 4
+    self, _ = _module(E, H, IM, torch.float32, 0, None)
+    idx = torch.full((N, K), -1, device=DEV, dtype=torch.long)
+    idx[1:, 0] = 0  # token 0 fully remote; others have one valid slot
+    w = torch.rand(N, K, device=DEV)
+    out = scattermoe_experts_forward_ep(self, torch.randn(N, H, device=DEV), idx, w)
+    assert out.shape == (N, H)
+    assert torch.equal(out[0], torch.zeros(H, device=DEV))
diff --git a/tests/integrations/kernels/scattermoe_lora/test_gates_backward_autocast_dtype.py b/tests/integrations/kernels/scattermoe_lora/test_gates_backward_autocast_dtype.py
new file mode 100644
index 0000000000..fbca45fbd4
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_gates_backward_autocast_dtype.py
@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Regression tests: the gates-branch backward must run in one dtype under bf16 autocast.
+
+Under bf16 autocast the gating bmm is autocast-eligible, so backward gets a bf16
+grad_out while the saved output_expanded stays fp32. Pre-fix the d_gates matmul
+raises "expected scalar type Float but found BFloat16"; post-fix it returns
+finite fp32 grads.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.scattermoe_lora import (
+    parallel_linear,
+    parallel_linear_lora,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+    flatten_sort_count,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA required for kernel launch"
+)
+
+
+_T = 64
+_E = 8
+_K = 2
+_H = 32
+_O = 48
+
+
+def _routing(logits: torch.Tensor):
+    """Softmax->topk->renorm routing, mirroring layers.py; returns an fp32
+    leaf ``gates`` (requires_grad) plus the sorted routing ids."""
+    routing_weights = torch.softmax(logits, dim=-1, dtype=torch.float)
+    routing_weights, selected_experts = torch.topk(routing_weights, _K, dim=-1)
+    routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+    gates = routing_weights.detach().clone().requires_grad_(True)
+    sei, ssi, eo = flatten_sort_count(selected_experts, _E)
+    return gates, selected_experts, sei, ssi, eo
+
+
+def test_parallel_linear_gates_backward_under_bf16_autocast():
+    torch.manual_seed(0)
+    device = torch.device("cuda:0")
+
+    x = torch.randn(_T, _H, device=device, dtype=torch.float32, requires_grad=True)
+    # ParallelExperts stores weight as [E, O, H] and passes .permute(0, 2, 1).
+    W = (
+        torch.randn(_E, _O, _H, device=device, dtype=torch.float32, requires_grad=True)
+        * 0.02
+    )
+    logits = torch.randn(_T, _E, device=device)
+    gates, selected_experts, sei, ssi, eo = _routing(logits)
+
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = parallel_linear(
+            x,
+            W.permute(0, 2, 1),
+            _K,
+            sei,
+            ssi,
+            eo,
+            gates=gates,
+        )
+    # repro precondition: autocast makes the output bf16, so backward gets a bf16 grad_out
+    assert out.dtype == torch.bfloat16
+
+    grad = torch.randn_like(out)
+    # pre-fix: raises at the d_gates matmul; post-fix: fp32 grads
+    gx, gW, gg = torch.autograd.grad(out, [x, W, gates], grad_outputs=grad)
+
+    for g in (gx, gW, gg):
+        assert g.dtype == torch.float32
+        assert torch.isfinite(g).all()
+
+    # check d_gates vs an fp32 reference: pins that grad_out is upcast, not output_expanded downcast
+    grad_fp32 = grad.float()
+    all_out = torch.einsum("th,eoh->teo", x.detach(), W.detach())
+    picked = all_out[torch.arange(_T, device=device).unsqueeze(1), selected_experts]
+    ref = (gates.unsqueeze(1) @ picked).squeeze(1)
+    (gg_ref,) = torch.autograd.grad(ref, gates, grad_outputs=grad_fp32)
+    assert torch.allclose(gg, gg_ref, atol=1e-2)
+
+
+def test_scattermoe_lora_gates_backward_under_bf16_autocast():
+    torch.manual_seed(0)
+    device = torch.device("cuda:0")
+    rank = 4
+
+    x = torch.randn(_T, _H, device=device, dtype=torch.float32, requires_grad=True)
+    # Frozen base weights [E, H, O]; only the LoRA adapters get gradients.
+    W = torch.randn(_E, _H, _O, device=device, dtype=torch.float32) * 0.02
+    lora_A = (
+        torch.randn(
+            _E * rank, _H, device=device, dtype=torch.float32, requires_grad=True
+        )
+        * 0.01
+    )
+    lora_B = (
+        torch.randn(
+            _O, _E * rank, device=device, dtype=torch.float32, requires_grad=True
+        )
+        * 0.01
+    )
+    logits = torch.randn(_T, _E, device=device)
+    gates, _, sei, ssi, eo = _routing(logits)
+
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = parallel_linear_lora(
+            x,
+            W,
+            _K,
+            sei,
+            ssi,
+            eo,
+            lora_A=lora_A,
+            lora_B=lora_B,
+            scaling=0.5,
+            gates=gates,
+        )
+    assert out.dtype == torch.bfloat16
+
+    grad = torch.randn_like(out)
+    # same crash site in the LoRA path; pre-fix raises, post-fix runs in fp32
+    gx, gA, gB, gg = torch.autograd.grad(
+        out, [x, lora_A, lora_B, gates], grad_outputs=grad
+    )
+
+    for g in (gx, gA, gB, gg):
+        assert g.dtype == torch.float32
+        assert torch.isfinite(g).all()
diff --git a/tests/integrations/kernels/scattermoe_lora/test_gptoss_layout.py b/tests/integrations/kernels/scattermoe_lora/test_gptoss_layout.py
new file mode 100644
index 0000000000..b344cc6d61
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_gptoss_layout.py
@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""scattermoe Triton path for gpt_oss-style experts (transposed [E,H,2I] weights,
+interleaved gate/up, per-expert bias, clamped sigmoid-GLU).
+
+Validated against an eager reference: output and every gradient must match. The LoRA
+fusion is the same ``scatter2scatter_lora`` the standard path uses, so the eager
+reference folds the kernel's delta ``ΔW_e = scaling * A_e^T @ W_B[e]`` into W_eff.
+TF32 is disabled so the fp32 comparison is tight (the kernel GEMMs otherwise use TF32).
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+import axolotl.integrations.kernels.libs.scattermoe_lora.experts as ex  # noqa: E402
+import axolotl.integrations.kernels.libs.scattermoe_lora.kernels.lora_ops as _lo  # noqa: E402
+import axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops as _ops  # noqa: E402
+from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (  # noqa: E402
+    scattermoe_experts_forward,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (  # noqa: E402
+    peft_lora_B_to_scattermoe,
+    peft_lora_to_scattermoe,
+)
+
+DEV = "cuda"
+ALPHA, LIMIT = 1.702, 7.0
+E, H, IM, N, K, R, SC = 8, 256, 128, 64, 4, 16, 0.5
+
+
+@pytest.fixture(autouse=True)
+def _no_tf32():
+    old = torch.backends.cuda.matmul.allow_tf32
+    torch.backends.cuda.matmul.allow_tf32 = False
+    a, b = _lo.ALLOW_TF32, _ops.ALLOW_TF32
+    _lo.ALLOW_TF32 = _ops.ALLOW_TF32 = False
+    yield
+    torch.backends.cuda.matmul.allow_tf32 = old
+    _lo.ALLOW_TF32, _ops.ALLOW_TF32 = a, b
+
+
+def _glu(gate_up):
+    gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+    gate = gate.clamp(max=LIMIT)
+    up = up.clamp(min=-LIMIT, max=LIMIT)
+    return (up + 1) * (gate * torch.sigmoid(gate * ALPHA))
+
+
+def _module(dt):
+    def mk(*s):
+        return (torch.randn(*s, device=DEV, dtype=dt) * 0.02).requires_grad_(True)
+
+    return SimpleNamespace(
+        num_experts=E,
+        gate_up_proj=mk(E, H, 2 * IM),
+        gate_up_proj_bias=mk(E, 2 * IM),
+        down_proj=mk(E, IM, H),
+        down_proj_bias=mk(E, H),
+        alpha=ALPHA,
+        limit=LIMIT,
+        is_transposed=True,
+        is_concatenated=False,
+        has_bias=True,
+        has_gate=True,
+    )
+
+
+def _eager(x, m, idx, w, weff=None):
+    gup, dp = weff if weff else (m.gate_up_proj, m.down_proj)
+    fe = idx.reshape(-1)
+    xg = x.repeat_interleave(K, dim=0)
+    gate_up = torch.bmm(xg.unsqueeze(1), gup[fe]).squeeze(1) + m.gate_up_proj_bias[fe]
+    h = _glu(gate_up)
+    o = torch.bmm(h.unsqueeze(1), dp[fe]).squeeze(1) + m.down_proj_bias[fe]
+    o = o * w.reshape(-1, 1)
+    tok = torch.arange(x.size(0), device=DEV).repeat_interleave(K)
+    return torch.zeros(x.size(0), H, device=DEV, dtype=x.dtype).index_add_(0, tok, o)
+
+
+def _rel(a, b):
+    return (a - b).abs().max().item() / max(b.abs().max().item(), 1e-6)
+
+
+@pytest.mark.parametrize("dt", [torch.float32, torch.bfloat16])
+def test_gptoss_base_matches_eager(dt):
+    m = _module(dt)
+    g = torch.Generator(device=DEV).manual_seed(0)
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    w = torch.rand(N, K, device=DEV, generator=g).to(dt)
+    grad = torch.randn(N, H, device=DEV, dtype=dt)
+    ps = [m.gate_up_proj, m.gate_up_proj_bias, m.down_proj, m.down_proj_bias]
+
+    def run(eager):
+        for p in ps:
+            p.grad = None
+        x = torch.randn(
+            N, H, device=DEV, dtype=dt, generator=torch.Generator(DEV).manual_seed(3)
+        ).requires_grad_(True)
+        out = (
+            _eager(x, m, idx, w) if eager else scattermoe_experts_forward(m, x, idx, w)
+        )
+        out.backward(grad)
+        return out.detach(), x.grad, [p.grad.clone() for p in ps]
+
+    ok, dxk, gk = run(False)
+    oe, dxe, ge = run(True)
+    tol = 2e-4 if dt == torch.float32 else 5e-2
+    assert _rel(ok, oe) < tol
+    assert _rel(dxk, dxe) < tol
+    for a, b in zip(gk, ge, strict=True):
+        assert _rel(a, b) < tol
+
+
+@pytest.mark.parametrize("dt", [torch.float32, torch.bfloat16])
+def test_gptoss_lora_matches_eager(dt, monkeypatch):
+    g = torch.Generator(device=DEV).manual_seed(0)
+    base = SimpleNamespace(
+        num_experts=E,
+        gate_up_proj=torch.randn(E, H, 2 * IM, device=DEV, dtype=dt, generator=g)
+        * 0.02,
+        gate_up_proj_bias=torch.randn(E, 2 * IM, device=DEV, dtype=dt, generator=g)
+        * 0.02,
+        down_proj=torch.randn(E, IM, H, device=DEV, dtype=dt, generator=g) * 0.02,
+        down_proj_bias=torch.randn(E, H, device=DEV, dtype=dt, generator=g) * 0.02,
+        alpha=ALPHA,
+        limit=LIMIT,
+        is_transposed=True,
+        is_concatenated=False,
+        has_bias=True,
+        has_gate=True,
+    )
+
+    def mk(*s):
+        return (
+            torch.randn(*s, device=DEV, dtype=dt, generator=g) * 0.05
+        ).requires_grad_(True)
+
+    pA1, pB1 = mk(R * E, H), mk(2 * IM, R * E)
+    pA2, pB2 = mk(R * E, IM), mk(H, R * E)
+    lora = [pA1, pB1, pA2, pB2]
+    gup_l = (*peft_lora_to_scattermoe(pA1, pB1, E, R), SC)
+    dwn_l = (*peft_lora_to_scattermoe(pA2, pB2, E, R), SC)
+    monkeypatch.setattr(ex, "_has_peft_wrapper", lambda s: True)
+    monkeypatch.setattr(ex, "_unwrap_experts_lora", lambda s: (s, gup_l, dwn_l))
+
+    # eager W_eff: ΔW_e = scaling * A_e^T @ W_B[e], W_B = lora_B.T.reshape(E,R,out)
+    A1 = pA1.reshape(E, R, H)
+    WB1 = peft_lora_B_to_scattermoe(pB1, E, R).t().reshape(E, R, 2 * IM)
+    A2 = pA2.reshape(E, R, IM)
+    WB2 = peft_lora_B_to_scattermoe(pB2, E, R).t().reshape(E, R, H)
+    gup_eff = base.gate_up_proj + SC * torch.bmm(A1.transpose(1, 2), WB1)
+    dp_eff = base.down_proj + SC * torch.bmm(A2.transpose(1, 2), WB2)
+
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    w = torch.rand(N, K, device=DEV, generator=g).to(dt)
+    grad = torch.randn(N, H, device=DEV, dtype=dt)
+
+    def run(eager):
+        for t in lora:
+            t.grad = None
+        x = torch.randn(
+            N, H, device=DEV, dtype=dt, generator=torch.Generator(DEV).manual_seed(3)
+        ).requires_grad_(True)
+        out = (
+            _eager(x, base, idx, w, weff=(gup_eff, dp_eff))
+            if eager
+            else scattermoe_experts_forward(base, x, idx, w)
+        )
+        out.backward(grad)
+        return out.detach(), x.grad, [t.grad.clone() for t in lora]
+
+    ok, dxk, gk = run(False)
+    oe, dxe, ge = run(True)
+    tol = 3e-4 if dt == torch.float32 else 5e-2
+    assert _rel(ok, oe) < tol
+    assert _rel(dxk, dxe) < tol
+    for a, b in zip(gk, ge, strict=True):
+        assert _rel(a, b) < tol
+        assert a.abs().sum() > 0  # LoRA grads actually populated
diff --git a/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_dequant_path.py b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_dequant_path.py
new file mode 100644
index 0000000000..d3408531f4
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_dequant_path.py
@@ -0,0 +1,426 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Arch-agnostic (non-Blackwell) coverage for the grouped NVFP4 MoE training path.
+
+The headline correctness oracle (test_grouped_fp4_train.py) is gated to Blackwell (sm100/sm120)
+because the CUTLASS/DeepGEMM forward backends only run there. But the grouped autograd Function
+also has a path that runs on ANY sm80+ CUDA GPU:
+
+  * FORWARD: the Marlin W4A16 backend (bf16-act x NVFP4-weight, standard Ampere-class ``mma``,
+    no Hopper/Blackwell intrinsics) covers sm80 (A100) / sm89 (L40S / RTX 4090) / sm120.
+  * BACKWARD: the gradient-consistent chunked bf16-dequant base dX (NVFP4 -> bf16 via a Triton
+    kernel + ``torch._grouped_mm``), selected on every non-sm120 arch and forced here via
+    ``prefer_fp8_dx=False``. This is the B8 FP4-expert backward and is arch-agnostic.
+
+So on the L40S (sm89) CI these tests EXECUTE (via Marlin fwd + bf16-dequant bwd), giving the
+grouped fwd/bwd oracle and the B11 backward None-grad contract real coverage there. They also
+execute on Blackwell. The Marlin/DeepGEMM/CUTLASS *numerics* and *perf* remain hardware-gated
+in test_grouped_fp4_train.py / test_grouped_fp4_perf.py; this file does not weaken those.
+
+Gate: CUDA available AND ``grouped_fp4_available('nvfp4')`` (True on sm80+ with nvcc via Marlin,
+or on sm90/100 via DeepGEMM, or sm120 via CUTLASS). NOT Blackwell-gated.
+
+Helpers (quantize_nvfp4, _bf16_oracle, _make_inputs, _cos, SHAPES, DEV) are reused from the
+sibling correctness module; importing them is safe - its module-level Blackwell skip applies only
+to ITS test functions, not to importing its helpers.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+# Reuse the validated quantizer + oracle + fixtures from the sibling correctness module.
+from .test_grouped_fp4_train import (
+    DEV,
+    SHAPE_IDS,
+    SHAPES,
+    _bf16_oracle,
+    _cos,
+    _fresh_nv,
+    _make_inputs,
+    quantize_nvfp4,
+)
+
+_IS_CUDA = torch.cuda.is_available()
+
+
+def _grouped_available() -> bool:
+    """True iff SOME grouped fp4 forward backend resolves here (Marlin on sm80+, else DeepGEMM /
+    CUTLASS). On the L40S CI this is True via Marlin, so the tests below execute (no Blackwell)."""
+    if not _IS_CUDA:
+        return False
+    try:
+        from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+            grouped_fp4_available,
+        )
+
+        return grouped_fp4_available("nvfp4")
+    except Exception:
+        return False
+
+
+pytestmark = pytest.mark.skipif(
+    not _grouped_available(),
+    reason="grouped NVFP4 needs CUDA + a resolvable fp4 forward backend (Marlin sm80+ / "
+    "DeepGEMM sm90+ / CUTLASS sm120)",
+)
+
+
+# ===========================================================================
+# Dequant-path forward+backward oracle (executes on sm89 via Marlin + bf16-dequant bwd)
+# ===========================================================================
+
+
+@pytest.mark.parametrize("cfg", SHAPES, ids=SHAPE_IDS)
+def test_grouped_fp4_fwd_bwd_dequant_path_matches_oracle(cfg):
+    """Forward cosine >= 0.97 and backward grads cosine >= 0.95 vs the bf16 per-expert oracle,
+    forcing the arch-agnostic bf16-dequant base dX (prefer_fp8_dx=False).
+
+    Runs on any sm80+ GPU (Marlin forward + chunked bf16-dequant backward), so it executes on the
+    sm89 CI where the Blackwell-gated oracle in test_grouped_fp4_train.py cannot.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_moe_train,
+    )
+
+    E, H, I = cfg["E"], cfg["H"], cfg["I"]
+    s = cfg["scaling"]
+
+    # Fresh NV tensors (the marlin backend frees qdata on first build, so the oracle weights are
+    # dequantized from an independent copy of the same random matrices).
+    twoI = 2 * I
+    torch.manual_seed(0)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv, dn_nv = _fresh_nv(Wgu, Wdn)
+    pt = torch.ones(E, device=DEV)
+    Wgu_b = nvfp4_dequant_bf16(gu_nv.qdata, gu_nv.scale, pt)
+    Wdn_b = nvfp4_dequant_bf16(dn_nv.qdata, dn_nv.scale, pt)
+
+    # Inputs (reuse the fixture for hidden/idx/wts + LoRA params, then make leaf copies).
+    hidden, idx, wts, _gu, _dn, _Wgu, _Wdn, (Agu, Bgu, Adn, Bdn) = _make_inputs(cfg)
+
+    hk = hidden.detach().clone().requires_grad_(True)
+    Agk = Agu.detach().clone().requires_grad_(True)
+    Bgk = Bgu.detach().clone().requires_grad_(True)
+    Adk = Adn.detach().clone().requires_grad_(True)
+    Bdk = Bdn.detach().clone().requires_grad_(True)
+
+    out = grouped_fp4_moe_train(
+        hk,
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agk, Bgk, s),
+        (Adk, Bdk, s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache={},
+        prefer_fp8_dx=False,  # force the arch-agnostic bf16-dequant base dX (B8)
+    )
+    assert torch.isfinite(out).all(), f"{cfg['name']}: fwd output has NaN/inf"
+    assert out.shape == (cfg["N"], H)
+    out.float().pow(2).mean().backward()
+
+    # Oracle on the same dequantized weights.
+    Agr = Agu.detach().clone().requires_grad_(True)
+    Bgr = Bgu.detach().clone().requires_grad_(True)
+    Adr = Adn.detach().clone().requires_grad_(True)
+    Bdr = Bdn.detach().clone().requires_grad_(True)
+    hr = hidden.detach().clone().requires_grad_(True)
+    ref = _bf16_oracle(
+        hr, idx, wts, Wgu_b, Wdn_b, Agr, Bgr, Adr, Bdr, cfg["act_type"], cfg["limit"], s
+    )
+    cos_fwd = _cos(out, ref)
+    assert cos_fwd > 0.97, f"{cfg['name']}: fwd cosine {cos_fwd:.4f} < 0.97"
+
+    ref.pow(2).mean().backward()
+    for name, gk, gr in [
+        ("d_hidden", hk.grad, hr.grad),
+        ("d_Agu", Agk.grad, Agr.grad),
+        ("d_Adn", Adk.grad, Adr.grad),
+    ]:
+        assert gk is not None, (
+            f"{cfg['name']}: {name} grad is None (backward didn't run)"
+        )
+        assert torch.isfinite(gk).all(), f"{cfg['name']}: {name} has NaN/inf"
+        cos = _cos(gk, gr)
+        assert cos > 0.95, f"{cfg['name']}: {name} cosine {cos:.4f} < 0.95"
+
+
+# ===========================================================================
+# B11: backward None-grad contract (data-independent FSDP2 backward collectives)
+# ===========================================================================
+
+
+def test_grouped_fp4_backward_none_grad_contract():
+    """The grouped autograd Function returns grads for x + the four LoRA params and None for the
+    routing / m_indices / offs / mode inputs (frozen experts -> no weight grad).
+
+    This is the code basis of B11: the backward emits a grad ONLY for the differentiable leaves, so
+    FSDP2's reduce-scatter sees gradients exclusively for the trainable LoRA params (a fixed,
+    data-independent set) - never for the frozen NVFP4 experts or the integer routing tensors. We
+    assert it by checking which inputs receive a grad after backward.
+
+    Forced onto the bf16-dequant base dX (prefer_fp8_dx=False) so it runs on any sm80+ GPU (Marlin
+    forward), executing on the sm89 CI. idx/wts are integer/non-leaf routing tensors that cannot
+    carry grads; the contract is that x and the LoRA params DO and the experts do NOT.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_moe_train,
+    )
+
+    cfg = SHAPES[0]
+    E, H, I = cfg["E"], cfg["H"], cfg["I"]
+    s = cfg["scaling"]
+    twoI = 2 * I
+    torch.manual_seed(3)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv, dn_nv = _fresh_nv(Wgu, Wdn)
+
+    hidden, idx, wts, _gu, _dn, _Wgu, _Wdn, (Agu, Bgu, Adn, Bdn) = _make_inputs(
+        cfg, seed=3
+    )
+    hk = hidden.detach().clone().requires_grad_(True)
+    Agk = Agu.detach().clone().requires_grad_(True)
+    Bgk = Bgu.detach().clone().requires_grad_(True)
+    Adk = Adn.detach().clone().requires_grad_(True)
+    Bdk = Bdn.detach().clone().requires_grad_(True)
+
+    # The frozen experts must NOT request grad (FSDP2 never reduce-scatters them).
+    assert not gu_nv.requires_grad and not dn_nv.requires_grad
+    assert not idx.requires_grad and not wts.requires_grad
+
+    out = grouped_fp4_moe_train(
+        hk,
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agk, Bgk, s),
+        (Adk, Bdk, s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache={},
+        prefer_fp8_dx=False,
+    )
+    out.float().pow(2).mean().backward()
+
+    # Differentiable leaves: x + the four LoRA params get non-None, finite grads.
+    for name, t in [
+        ("x", hk),
+        ("Agu", Agk),
+        ("Bgu", Bgk),
+        ("Adn", Adk),
+        ("Bdn", Bdk),
+    ]:
+        assert t.grad is not None, (
+            f"{name} grad is None (backward broke the LoRA grad path)"
+        )
+        assert torch.isfinite(t.grad).all(), f"{name} grad has NaN/inf"
+
+    # Non-differentiable inputs (frozen experts, integer routing) carry no grad.
+    assert getattr(gu_nv, "grad", None) is None, "frozen gate_up expert got a grad"
+    assert getattr(dn_nv, "grad", None) is None, "frozen down expert got a grad"
+    assert idx.grad is None and wts.grad is None, "routing tensors got a grad"
+
+
+def test_grouped_experts_function_backward_returns_none_for_nondiff_inputs():
+    """Directly assert _GroupedExperts.backward's None-grad tuple shape: grads only for x (pos 0)
+    and the four LoRA params (Agu/Bgu/Adn/Bdn, pos 3-6); None for base / weight_recipe / m_indices /
+    offs / scaling / limit / mode / act_type / prefer_fp8_dx. This is the literal B11 contract.
+
+    Built on the Marlin forward + bf16-dequant backward so it runs on sm80+ (executes on sm89 CI).
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora import grouped_train as gt
+
+    cfg = SHAPES[0]
+    E, H, I = cfg["E"], cfg["H"], cfg["I"]
+    s = cfg["scaling"]
+    twoI = 2 * I
+    torch.manual_seed(5)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv, dn_nv = _fresh_nv(Wgu, Wdn)
+
+    hidden, idx, wts, _gu, _dn, _Wgu, _Wdn, (Agu, Bgu, Adn, Bdn) = _make_inputs(
+        cfg, seed=5
+    )
+
+    # Resolve a forward backend the same way grouped_fp4_moe_train does (Marlin on sm89).
+    backend = gt._train_backend("nvfp4")
+    if backend == "marlin":
+        from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.backend import (
+            MARLIN_TILE,
+        )
+
+        tile = MARLIN_TILE
+    else:
+        tile = gt.TILE
+
+    N = hidden.size(0)
+    dev = hidden.device
+    flat = idx.reshape(-1)
+    order = flat.argsort()
+    rep = torch.arange(N, device=dev).repeat_interleave(idx.size(1))[order]
+    exp_sorted = flat[order]
+    counts = torch.bincount(flat, minlength=E)
+    ptiles = (counts + tile - 1) // tile
+    roff = torch.cat([ptiles.new_zeros(1), ptiles.cumsum(0)]) * tile
+    coff = torch.cat([counts.new_zeros(1), counts.cumsum(0)])
+    padded_row = roff[exp_sorted] + (
+        torch.arange(exp_sorted.numel(), device=dev) - coff[exp_sorted]
+    )
+    m_indices = torch.repeat_interleave(
+        torch.arange(E, dtype=torch.int32, device=dev), ptiles
+    )
+    offs = (ptiles * tile).cumsum(0).to(torch.int32)
+    Mt = int(ptiles.sum()) * tile
+
+    if backend == "marlin":
+        from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.backend import (
+            build_marlin_forward_base,
+        )
+
+        base = build_marlin_forward_base(gu_nv, dn_nv, {})
+    elif backend == "deepgemm":
+        from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+            _cached_mxfp4,
+        )
+
+        base = (
+            "deepgemm",
+            _cached_mxfp4(gu_nv, gt._pt(gu_nv, E, dev), {}, "gate_up"),
+            _cached_mxfp4(dn_nv, gt._pt(dn_nv, E, dev), {}, "down"),
+        )
+    else:  # cutlass (sm120)
+        gu_eng = gt._engine(Mt, twoI, H, E, "nvfp4")
+        gu_eng.set_weights(gu_nv.qdata, gu_nv.scale)
+        dn_eng = gt._engine(Mt, H, I, E, "nvfp4")
+        dn_eng.set_weights(dn_nv.qdata, dn_nv.scale)
+        base = ("cutlass", gu_eng, dn_eng)
+
+    Ags, Bgs, _ = gt._lora_stack((Agu, Bgu, s), E, H, twoI)
+    Ads, Bds, _ = gt._lora_stack((Adn, Bdn, s), E, I, H)
+    Ags = Ags.detach().requires_grad_(True)
+    Bgs = Bgs.detach().requires_grad_(True)
+    Ads = Ads.detach().requires_grad_(True)
+    Bds = Bds.detach().requires_grad_(True)
+
+    A = hidden.new_zeros(Mt, H).index_copy(0, padded_row, hidden[rep])
+    A = A.detach().requires_grad_(True)
+    lim = float(cfg["limit"])
+    recipe = lambda: (gu_nv, dn_nv)  # noqa: E731
+
+    dn = gt._GroupedExperts.apply(
+        A,
+        base,
+        recipe,
+        Ags,
+        Bgs,
+        Ads,
+        Bds,
+        m_indices,
+        offs,
+        s,
+        lim,
+        "nvfp4",
+        cfg["act_type"],
+        False,  # prefer_fp8_dx=False -> bf16-dequant base dX
+    )
+    dn.float().pow(2).mean().backward()
+
+    # Differentiable leaves get grads; non-diff inputs get None (the Function's backward tuple).
+    assert A.grad is not None and torch.isfinite(A.grad).all(), "dx (pos 0) missing"
+    for nm, t in (("Agu", Ags), ("Bgu", Bgs), ("Adn", Ads), ("Bdn", Bds)):
+        assert t.grad is not None and torch.isfinite(t.grad).all(), f"{nm} grad missing"
+    # m_indices / offs are int tensors and never leaves with grad; assert they stayed grad-free.
+    assert m_indices.grad is None and offs.grad is None
+
+
+# ===========================================================================
+# Marlin fused-dequant bit-exactness on sm80+ (executes on sm89 via the Marlin ext)
+# ===========================================================================
+#
+# test_grouped_fp4_train.py also has marlin bit-exact tests, but they sit under that module's
+# Blackwell pytestmark AND a per-test capability==12 gate, so they only run on sm120. The Marlin
+# W4A16 ext is sm80+ (see marlin_w4a16/__init__: standard Ampere mma, no Blackwell intrinsics), so
+# the fused marlin->bf16 dequant is bit-exact-checkable on sm89 too. These sm80+-gated variants give
+# that numeric check real coverage on the L40S CI without touching the Blackwell module.
+
+
+def _marlin_available() -> bool:
+    if not _IS_CUDA:
+        return False
+    try:
+        from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16 import (
+            marlin_w4a16_available,
+        )
+
+        return marlin_w4a16_available()
+    except Exception:
+        return False
+
+
+@pytest.mark.skipif(
+    not _marlin_available(),
+    reason="marlin W4A16 ext (sm80+, nvcc) required",
+)
+@pytest.mark.parametrize(
+    "C,N,K", [(4, 64, 128), (4, 128, 64)], ids=["gate_up_N64_K128", "down_N128_K64"]
+)
+def test_marlin_fused_dequant_bit_exact_sm80(C, N, K):
+    """marlin_dequant_bf16 is bit-exact vs nvfp4_dequant_bf16 on sm80+ (gate_up + down shapes).
+
+    The marlin int32 layout decoded by the fused Triton dequant must match the reference NVFP4->bf16
+    dequant exactly (maxerr == 0). N%64==0 (marlin N_t) and K%16==0 (NVFP4 block) are required.
+    Runs on any sm80+ GPU with the Marlin ext, so it executes on the sm89 CI.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16 import load_ext
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.backend import (
+        _build_base_scatter,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.fused_dequant import (
+        marlin_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.prep import (
+        prepare_nvfp4_weight_for_marlin,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+        fp4_codebook,
+    )
+
+    torch.manual_seed(42)
+    W = torch.randn(C, N, K, device=DEV, dtype=torch.bfloat16) * 0.04
+    nv = quantize_nvfp4(W)
+    pt = torch.ones(C, device=DEV)
+
+    ext = load_ext()
+    scatter_lut = _build_base_scatter(torch.device(DEV))
+    cb = fp4_codebook(torch.device(DEV)).float()
+
+    marlin_packed = prepare_nvfp4_weight_for_marlin(
+        nv.qdata, nv.scale, pt, N, K, torch.bfloat16, ext.gptq_marlin_repack
+    )
+    qw_flat = marlin_packed[0].reshape(C, -1)
+    ref_bf16 = nvfp4_dequant_bf16(nv.qdata, nv.scale, pt)
+    got_bf16 = marlin_dequant_bf16(qw_flat, nv.scale, pt, scatter_lut, cb, N, K, C)
+
+    max_err = (ref_bf16.float() - got_bf16.float()).abs().max().item()
+    assert max_err == 0.0, (
+        f"marlin fused dequant not bit-exact vs nvfp4_dequant_bf16 (N={N}, K={K}): "
+        f"maxerr={max_err}"
+    )
diff --git a/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_guards.py b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_guards.py
new file mode 100644
index 0000000000..dbeb462e2e
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_guards.py
@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""No-GPU coverage for the load-bearing safety guards on the grouped NVFP4 MoE path.
+
+Each guard turns silent corruption or a Blackwell SIGSEGV into a clear RuntimeError. They were
+shipping untested: the GPU oracle (test_grouped_fp4_train.py) never enters them, because on the sm89
+CI a fused backend is available, so the C1 cutlass-only raise and the A4/B2 fall-throughs do not
+fire. Here each guard is driven directly with mocked arch/availability probes - no CUDA needed (the
+scattermoe_lora package imports triton, so the module skips on a triton-less host).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytest.importorskip("triton")
+
+
+def _mock_experts(num_experts):
+    """Minimal stand-in for a scattermoe experts module. The guards under test fire before any
+    weight/LoRA shape is read, so the proj params and LoRA tuples only need to be present + non-None.
+    """
+    lora = (torch.zeros(1), torch.zeros(1), 1.0)
+    return SimpleNamespace(
+        num_experts=num_experts,
+        gate_up_proj=torch.zeros(1),
+        down_proj=torch.zeros(1),
+        _scattermoe_lora={"gate_up_proj": lora, "down_proj": lora},
+    )
+
+
+def test_c1_cutlass_nonunit_weight_scale_2_without_alt_backend_raises(monkeypatch):
+    # C1: cutlass cannot fold a non-unit per_tensor_scale (weight_scale_2). With cutlass resolved,
+    # a non-unit pt present, and neither marlin nor deepgemm available to fold it, the path must
+    # hard-error rather than silently produce a wrong forward + grad mismatch.
+    from axolotl.integrations.kernels.libs.scattermoe_lora import grouped_train as gt
+
+    monkeypatch.setattr(gt, "_train_backend", lambda mode: "cutlass")
+    monkeypatch.setattr(gt, "_has_nonunit_pt", lambda *nvs: True)
+    monkeypatch.setattr(gt, "_backend_available", lambda name: False)
+
+    E, twoI, packed_k, H, packed_i = 4, 16, 8, 8, 4
+    gu = SimpleNamespace(qdata=torch.zeros(E, twoI, packed_k, dtype=torch.uint8))
+    dn = SimpleNamespace(qdata=torch.zeros(E, H, packed_i, dtype=torch.uint8))
+    hidden = torch.zeros(2, H)
+
+    with pytest.raises(RuntimeError, match="weight_scale_2"):
+        gt.grouped_fp4_moe_train(
+            hidden, None, None, gu, dn, None, None, None, "nvfp4", mxfp4_cache={}
+        )
+
+
+def test_a4_nvfp4_lora_grouped_unavailable_raises(monkeypatch):
+    # A4: NVFP4 experts + LoRA select the grouped path (dsv4_fp4_grouped_mode set), but no fused
+    # grouped backend resolves -> must hard-error instead of falling through to the SIGSEGV MX kernel.
+    from axolotl.integrations.kernels.libs.scattermoe_lora import (
+        experts as ex,
+        grouped_train as gt,
+    )
+
+    monkeypatch.setattr(ex, "is_nvfp4_param", lambda p: True)
+    monkeypatch.setattr(gt, "grouped_fp4_available", lambda mode: False)
+    monkeypatch.setattr(ex.RUNTIME, "fp4_grouped_mode", "nvfp4")
+
+    E, H, topk, N = 4, 8, 2, 6
+    self_ = _mock_experts(E)
+    top_k_index = torch.randint(0, E, (N, topk))
+    top_k_weights = torch.rand(N, topk)
+    hidden = torch.randn(N, H)
+
+    with pytest.raises(RuntimeError, match="scatter2scatter_lora_mx"):
+        ex.scattermoe_experts_forward(self_, hidden, top_k_index, top_k_weights)
+
+
+def test_b2_nvfp4_lora_without_grouped_mode_raises_on_blackwell(monkeypatch):
+    # B2: NVFP4 experts + LoRA WITHOUT dsv4_fp4_grouped_mode would fall through to the legacy MX
+    # kernel, which SIGSEGVs on Blackwell. On sm100/sm120 it must hard-error with guidance instead.
+    from axolotl.integrations.kernels.libs.scattermoe_lora import experts as ex
+
+    monkeypatch.setattr(ex, "is_nvfp4_param", lambda p: True)
+    monkeypatch.setattr(ex.RUNTIME, "fp4_grouped_mode", None)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 0))
+
+    E, H, topk, N = 4, 8, 2, 6
+    self_ = _mock_experts(E)
+    top_k_index = torch.randint(0, E, (N, topk))
+    top_k_weights = torch.rand(N, topk)
+    hidden = torch.randn(N, H)
+
+    with pytest.raises(RuntimeError, match="grouped fp4 MoE path on Blackwell"):
+        ex.scattermoe_experts_forward(self_, hidden, top_k_index, top_k_weights)
+
+
+def test_b2_nvfp4_lora_without_grouped_mode_ok_off_blackwell(monkeypatch):
+    # The B2 guard must NOT fire on non-Blackwell archs (the legacy MX path is left intact there);
+    # it should fall through to _prepare_weights_and_lora rather than raising the B2 error.
+    from axolotl.integrations.kernels.libs.scattermoe_lora import experts as ex
+
+    monkeypatch.setattr(ex, "is_nvfp4_param", lambda p: True)
+    monkeypatch.setattr(ex.RUNTIME, "fp4_grouped_mode", None)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 9))
+
+    E, H, topk, N = 4, 8, 2, 6
+    self_ = _mock_experts(E)
+    top_k_index = torch.randint(0, E, (N, topk))
+    top_k_weights = torch.rand(N, topk)
+    hidden = torch.randn(N, H)
+
+    # The B2 guard is skipped (sm89); execution proceeds past it. It will fail later in the real MX
+    # path on CPU, but NOT with the B2 Blackwell message.
+    with pytest.raises(Exception) as exc:
+        ex.scattermoe_experts_forward(self_, hidden, top_k_index, top_k_weights)
+    assert "grouped fp4 MoE path on Blackwell" not in str(exc.value)
diff --git a/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_perf.py b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_perf.py
new file mode 100644
index 0000000000..a8680f1420
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_perf.py
@@ -0,0 +1,726 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Speed + memory regression guards for the grouped/marlin NVFP4 MoE training path.
+
+All assertions are RATIO-based (machine-agnostic). Measured baselines on RTX PRO 6000 (sm120):
+  DSV4  (E=256, H=4096, I=2048, top6):  grouped fwd ~1.9x, e2e ~2.27x vs bf16 dequant-per-step
+  gemma4 (E=128, H=2816, I=704,  top8):  grouped fwd ~2.1x, e2e ~1.8x  vs bf16 dequant-per-step
+
+Memory (DSV4 E=256 marlin, sm120):
+  marlin extra resident vs NVFP4-only baseline <= 1.3x (+2.65 GB, not the old +12 GB double-copy)
+  transient peak during fwd+bwd <= 1.0x the bf16-resident baseline peak (measured 0.87x at E=256)
+
+Skip conditions (all tests):
+  - Not Blackwell (sm100 / sm120)
+  - Insufficient free VRAM for the target E (skip gracefully, log reason)
+
+These use the real E=256 / E=128 shapes where VRAM allows; fall back to E=64 with looser bounds
+(x0.7 of the threshold) so the test still runs on constrained systems and logs the fallback.
+"""
+
+from __future__ import annotations
+
+import gc
+
+import pytest
+import torch
+
+# ---------------------------------------------------------------------------
+# Skip gates
+# ---------------------------------------------------------------------------
+
+_IS_CUDA = torch.cuda.is_available()
+_IS_BLACKWELL = _IS_CUDA and torch.cuda.get_device_capability()[0] in (10, 12)
+
+pytestmark = [
+    pytest.mark.perf,
+    pytest.mark.skipif(
+        not _IS_BLACKWELL,
+        reason="grouped NVFP4 training backends require Blackwell (sm100/sm120)",
+    ),
+]
+
+DEV = "cuda"
+
+# VRAM thresholds: skip at the large E if we can't fit
+# DSV4 E=256: 2 copies of weights (grouped + baseline) + marlin qweight + activations ~50 GB
+# gemma4 E=128: similar ~15 GB
+_VRAM_DSV4_GB = 50.0
+_VRAM_GEMMA4_GB = 15.0
+_VRAM_FALLBACK_RATIO = 0.70  # loosen speed thresholds by 30% when using fallback E
+
+
+def _free_vram_gb() -> float:
+    if not _IS_CUDA:
+        return 0.0
+    torch.cuda.synchronize()
+    free, _ = torch.cuda.mem_get_info()
+    return free / 1e9
+
+
+# ---------------------------------------------------------------------------
+# NVFP4 quantizer (self-contained; does NOT import from scratch bench dir)
+#
+# For large E (e.g. E=256, N=4096, K=4096), a naïve broadcast quantizer would
+# allocate a [E,N,K//16,16,16] float32 intermediate (~256 GiB). We instead use
+# NVFP4Tensor.to_nvfp4 in batches of at most 32 experts and concatenate qdata+scale,
+# which keeps the peak transient to ~2 × batch × N × K × 4 bytes (manageable).
+# ---------------------------------------------------------------------------
+
+
+def quantize_nvfp4(W: torch.Tensor, batch: int = 32):
+    """Memory-efficient NVFP4 quantization: W[E,N,K] bf16 -> NVFP4Tensor.
+
+    Processes `batch` experts at a time to bound the peak float32 transient;
+    concatenates qdata and scale on CPU, then moves to the target device.
+    """
+    from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+    E, N, K = W.shape
+    dev = W.device
+    qdata_parts, scale_parts = [], []
+    for e0 in range(0, E, batch):
+        e1 = min(e0 + batch, E)
+        nv_b = NVFP4Tensor.to_nvfp4(W[e0:e1].contiguous(), block_size=16)
+        qdata_parts.append(nv_b.qdata.cpu())
+        scale_parts.append(nv_b.scale.cpu())
+        del nv_b
+    qdata = torch.cat(qdata_parts, dim=0).to(dev)
+    scale = torch.cat(scale_parts, dim=0).to(dev)
+    return NVFP4Tensor(
+        qdata,
+        scale,
+        block_size=16,
+        orig_dtype=torch.bfloat16,
+        per_tensor_scale=torch.ones((), device=dev),
+    )
+
+
+# Alias for tests that need a real NVFP4Tensor (same as quantize_nvfp4 here)
+quantize_nvfp4_real = quantize_nvfp4
+
+
+# ---------------------------------------------------------------------------
+# Timing helpers
+# ---------------------------------------------------------------------------
+
+
+def _time_ms(fn, n_warmup=5, n_iter=12):
+    """Median of n_iter timings (ms) after n_warmup warm-up calls."""
+    for _ in range(n_warmup):
+        fn()
+        torch.cuda.synchronize()
+    timings = []
+    for _ in range(n_iter):
+        s = torch.cuda.Event(enable_timing=True)
+        e = torch.cuda.Event(enable_timing=True)
+        s.record()
+        fn()
+        e.record()
+        torch.cuda.synchronize()
+        timings.append(s.elapsed_time(e))
+    timings.sort()
+    return timings[n_iter // 2]
+
+
+def _peak_mb(fn, n_warmup=3):
+    """Peak memory allocated (MB) above current baseline during fn(), after warm-up."""
+    for _ in range(n_warmup):
+        fn()
+        torch.cuda.synchronize()
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats()
+    base = torch.cuda.memory_allocated()
+    fn()
+    torch.cuda.synchronize()
+    return (torch.cuda.max_memory_allocated() - base) / 1e6
+
+
+# ---------------------------------------------------------------------------
+# Baseline: chunked_dequant_grouped_base — the dequant-per-step path replaced by marlin.
+# Uses a SEPARATE set of NV tensors so the marlin memory-free path doesn't invalidate them.
+# ---------------------------------------------------------------------------
+
+
+def _baseline_fwd_step(hidden, idx, wts, gu_nv_base, dn_nv_base, pt, limit, act_type):
+    """Baseline forward: chunked dequant + grouped_mm (no marlin, no LoRA).
+
+    This is the `chunked_dequant_grouped_base` path that the grouped/marlin path replaced.
+    LoRA is omitted to keep the comparison fair (LoRA is the same cost in both paths).
+    Uses a separate gu/dn_nv_base that has NOT been marlin-freed.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        chunked_dequant_grouped_base,
+    )
+
+    return chunked_dequant_grouped_base(
+        hidden, idx, wts, gu_nv_base, dn_nv_base, pt, limit
+    )
+
+
+# ---------------------------------------------------------------------------
+# Shape configuration (real E, with VRAM-check fallback)
+# ---------------------------------------------------------------------------
+
+
+def _dsv4_cfg(prefer_full=True):
+    """DSV4-Flash shape. Returns (cfg, fallback_used)."""
+    full = dict(
+        E=256,
+        H=4096,
+        I=2048,
+        topk=6,
+        N=512,
+        r=16,
+        act_type="silu",
+        limit=7.0,
+        scaling=2.0,
+    )
+    small = dict(
+        E=64,
+        H=4096,
+        I=2048,
+        topk=6,
+        N=128,
+        r=16,
+        act_type="silu",
+        limit=7.0,
+        scaling=2.0,
+    )
+    if prefer_full and _free_vram_gb() >= _VRAM_DSV4_GB:
+        return full, False
+    return small, True
+
+
+def _gemma4_cfg(prefer_full=True):
+    """gemma4-A4B shape. Returns (cfg, fallback_used)."""
+    full = dict(
+        E=128,
+        H=2816,
+        I=704,
+        topk=8,
+        N=256,
+        r=16,
+        act_type="gelu_tanh",
+        limit=1e30,
+        scaling=2.0,
+    )
+    small = dict(
+        E=32,
+        H=2816,
+        I=704,
+        topk=8,
+        N=64,
+        r=16,
+        act_type="gelu_tanh",
+        limit=1e30,
+        scaling=2.0,
+    )
+    if prefer_full and _free_vram_gb() >= _VRAM_GEMMA4_GB:
+        return full, False
+    return small, True
+
+
+def _build_tensors(cfg):
+    """Build tensors for a shape config.
+
+    Returns (gu_nv, dn_nv, gu_nv_base, dn_nv_base, pt, hidden, idx, wts, lora).
+    gu_nv/dn_nv: NVFP4Tensor for the grouped/marlin path (qdata freed after first use).
+    gu_nv_base/dn_nv_base: SEPARATE NVFP4Tensor for the baseline timing (retained; not freed).
+    pt: per-tensor scale [E] float32.
+
+    Uses a batched quantizer (32 experts at a time) to avoid the 256-GiB broadcast OOM
+    that would occur with a naïve [E,N,K//16,16,16] float32 broadcast.
+    """
+    E, H, I, topk, N, r = cfg["E"], cfg["H"], cfg["I"], cfg["topk"], cfg["N"], cfg["r"]
+    twoI = 2 * I
+    torch.manual_seed(0)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv = quantize_nvfp4(Wgu)
+    gu_nv_base = quantize_nvfp4(Wgu)  # separate copy for baseline (not freed by marlin)
+    del Wgu
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    dn_nv = quantize_nvfp4(Wdn)
+    dn_nv_base = quantize_nvfp4(Wdn)
+    del Wdn
+    pt = torch.ones(E, device=DEV)
+    hidden = torch.randn(N, H, device=DEV, dtype=torch.bfloat16) * 0.5
+    idx = torch.stack([torch.randperm(E, device=DEV)[:topk] for _ in range(N)])
+    wts = torch.softmax(torch.randn(N, topk, device=DEV), -1).to(torch.bfloat16)
+    Agu = (
+        torch.randn(r * E, H, device=DEV, dtype=torch.bfloat16) * 0.02
+    ).requires_grad_(True)
+    Bgu = (
+        torch.randn(twoI, r * E, device=DEV, dtype=torch.bfloat16) * 0.02
+    ).requires_grad_(True)
+    Adn = (
+        torch.randn(r * E, I, device=DEV, dtype=torch.bfloat16) * 0.02
+    ).requires_grad_(True)
+    Bdn = (
+        torch.randn(H, r * E, device=DEV, dtype=torch.bfloat16) * 0.02
+    ).requires_grad_(True)
+    return (
+        gu_nv,
+        dn_nv,
+        gu_nv_base,
+        dn_nv_base,
+        pt,
+        hidden,
+        idx,
+        wts,
+        (Agu, Bgu, Adn, Bdn),
+    )
+
+
+# ===========================================================================
+# Perf 1+2: Speed guards (forward and full step)
+# ===========================================================================
+
+
+def _run_speed_test(cfg, fallback, shape_name, fwd_ratio_floor, step_ratio_floor):
+    """Core speed test: grouped/marlin path vs dequant-per-step baseline.
+
+    Baseline = nvfp4_dequant_bf16 (full-E, not cached) — the bottleneck this path replaces.
+    Grouped = grouped_fp4_moe_train (marlin forward; qdata freed after first call; cached marlin).
+
+    Separate qdata/scale copies are kept for the baseline so the marlin memory-free path doesn't
+    invalidate the baseline tensor references.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+        grouped_fp4_moe_train,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+
+    if fallback:
+        fwd_ratio_floor *= _VRAM_FALLBACK_RATIO
+        step_ratio_floor *= _VRAM_FALLBACK_RATIO
+
+    (
+        gu_nv,
+        dn_nv,
+        gu_nv_base,
+        dn_nv_base,
+        pt,
+        hidden,
+        idx,
+        wts,
+        (Agu, Bgu, Adn, Bdn),
+    ) = _build_tensors(cfg)
+    s = cfg["scaling"]
+    act_type, limit = cfg["act_type"], cfg["limit"]
+
+    # Build the marlin cache ONCE upfront (so all subsequent calls use the cached marlin weights
+    # and don't try to re-read the freed qdata). This mirrors how real training works.
+    cache = {}
+    grouped_fp4_moe_train(
+        hidden.detach(),
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agu.detach(), Bgu.detach(), s),
+        (Adn.detach(), Bdn.detach(), s),
+        limit,
+        "nvfp4",
+        act_type=act_type,
+        mxfp4_cache=cache,
+    )
+    torch.cuda.synchronize()
+
+    # --- Grouped/marlin forward timing (warm up + median) ---
+    def grouped_fwd():
+        grouped_fp4_moe_train(
+            hidden.detach(),
+            idx,
+            wts,
+            gu_nv,
+            dn_nv,
+            (Agu.detach(), Bgu.detach(), s),
+            (Adn.detach(), Bdn.detach(), s),
+            limit,
+            "nvfp4",
+            act_type=act_type,
+            mxfp4_cache=cache,
+        )
+
+    t_grouped_fwd = _time_ms(grouped_fwd)
+
+    # --- Baseline: chunked_dequant_grouped_base (dequant per step + grouped_mm) ---
+    # This is the path that the grouped/marlin replaced. Uses separate NV tensors
+    # so the marlin memory-free path doesn't invalidate these references.
+    def baseline_fwd():
+        _baseline_fwd_step(
+            hidden.detach(),
+            idx,
+            wts,
+            gu_nv_base,
+            dn_nv_base,
+            pt,
+            limit,
+            act_type,
+        )
+
+    t_base_fwd = _time_ms(baseline_fwd, n_warmup=3, n_iter=8)
+
+    fwd_ratio = t_base_fwd / t_grouped_fwd
+    print(
+        f"\n[{shape_name}] fwd: grouped={t_grouped_fwd:.1f}ms base={t_base_fwd:.1f}ms "
+        f"ratio={fwd_ratio:.2f}x (floor={fwd_ratio_floor:.2f}x)"
+    )
+    assert fwd_ratio >= fwd_ratio_floor, (
+        f"{shape_name}: grouped fwd ratio {fwd_ratio:.2f}x < floor {fwd_ratio_floor:.2f}x. "
+        "Speed regression: grouped/marlin forward should be faster than dequant-per-step."
+    )
+
+    # --- Full step (fwd+bwd) timing (same cache; qdata already freed) ---
+    def grouped_step():
+        for p in (Agu, Bgu, Adn, Bdn):
+            p.grad = None
+        hk = hidden.detach().requires_grad_()
+        out = grouped_fp4_moe_train(
+            hk,
+            idx,
+            wts,
+            gu_nv,
+            dn_nv,
+            (Agu, Bgu, s),
+            (Adn, Bdn, s),
+            limit,
+            "nvfp4",
+            act_type=act_type,
+            mxfp4_cache=cache,
+        )
+        out.float().pow(2).mean().backward()
+
+    t_grouped_step = _time_ms(grouped_step)
+    print(f"[{shape_name}] full-step (fwd+bwd): {t_grouped_step:.1f}ms")
+    for p in (Agu, Bgu, Adn, Bdn):
+        if p.grad is not None:
+            assert torch.isfinite(p.grad).all(), (
+                f"{shape_name}: LoRA grad has NaN/inf in perf test"
+            )
+
+
+@pytest.mark.parametrize(
+    "shape_fn,shape_name,vram_needed,fwd_floor,step_floor",
+    [
+        (_dsv4_cfg, "dsv4", _VRAM_DSV4_GB, 1.5, 1.3),
+        (_gemma4_cfg, "gemma4", _VRAM_GEMMA4_GB, 1.5, 1.3),
+    ],
+    ids=["dsv4", "gemma4"],
+)
+def test_grouped_fp4_speed(shape_fn, shape_name, vram_needed, fwd_floor, step_floor):
+    """grouped/marlin fwd is >= fwd_floor faster than chunked_dequant_grouped_base baseline.
+
+    Baseline = chunked_dequant_grouped_base (NVFP4->bf16 dequant per step + grouped_mm); the
+    path that the grouped/marlin fwd replaced. Uses separate NV tensors for baseline and grouped.
+
+    Measured on RTX PRO 6000 (sm120):
+      DSV4   (E=256): ~2.50x fwd vs chunked baseline
+      gemma4 (E=128): ~1.64x fwd vs chunked baseline
+
+    Thresholds: fwd >= 1.5x (safe floor well below measured; allows GPU-to-GPU variance).
+    """
+    cfg, fallback = shape_fn()
+    if fallback:
+        pytest.skip(f"{shape_name}: insufficient VRAM (need {vram_needed:.0f} GB free)")
+    _run_speed_test(cfg, fallback, shape_name, fwd_floor, step_floor)
+
+
+# ===========================================================================
+# Perf 3: Memory guard — marlin double-copy fix
+# ===========================================================================
+
+
+@pytest.mark.skipif(
+    not (
+        _IS_CUDA
+        and torch.cuda.is_available()
+        and torch.cuda.get_device_capability()[0] == 12
+    ),
+    reason="marlin W4A16 memory test is sm120 only",
+)
+def test_marlin_memory_no_double_copy():
+    """After build_marlin_forward_base, the NVFP4 qdata must be freed (single copy resident).
+
+    Original bug: the marlin path kept both qdata (4-bit packed) AND qweight (marlin int32 repack)
+    in memory — about +12 GB for DSV4 E=256. The fix frees qdata.data after repacking so only the
+    marlin qweight + original scales are resident (+2.65 GB). This test asserts <= 1.3x overhead
+    relative to the marlin qweight-only baseline.
+
+    Memory reference (RTX PRO 6000, sm120):
+      NVFP4 qdata (E=256, gate_up+down): ~6.4 GB total (4-bit packed)
+      marlin qweight (E=256, gate_up+down): ~6.4 GB (8-bit marlin layout)
+      OLD (bug): both resident = ~12.8 GB
+      NEW (fix): only marlin qweight + scales = ~6.4 + 0.05 GB = ~6.45 GB
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16 import (
+        marlin_w4a16_available,
+    )
+
+    if not marlin_w4a16_available():
+        pytest.skip("marlin W4A16 ext not available")
+
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.backend import (
+        _build_base_scatter,
+        build_marlin_forward_base,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+
+    cfg, fallback = _dsv4_cfg()
+    if fallback:
+        pytest.skip("insufficient VRAM for marlin memory test")
+
+    E, H, I = cfg["E"], cfg["H"], cfg["I"]
+    twoI = 2 * I
+
+    # Pre-build scatter LUT (doesn't count toward memory delta)
+    _build_base_scatter(torch.device(DEV))
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+
+    # Build real NVFP4Tensor (needed for the .qdata.data = empty() free path).
+    # Use batched quantizer (32 experts at a time) to avoid the 256GiB broadcast OOM.
+    torch.manual_seed(0)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv_base = quantize_nvfp4_real(Wgu)
+    del Wgu
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    dn_nv_base = quantize_nvfp4_real(Wdn)
+    del Wdn
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+
+    # Baseline: measure the NVFP4 qdata storage size before marlin build
+
+    mem_nvfp4_only = torch.cuda.memory_allocated()  # qdata in memory
+    qdata_bytes = gu_nv_base.qdata.numel() * gu_nv_base.qdata.element_size()
+    qdata_bytes += dn_nv_base.qdata.numel() * dn_nv_base.qdata.element_size()
+
+    # Build marlin base (this repacks + should free qdata)
+    cache = {}
+    build_marlin_forward_base(gu_nv_base, dn_nv_base, cache)
+    gc.collect()
+    torch.cuda.synchronize()
+    mem_after_marlin = torch.cuda.memory_allocated()
+
+    # The freed qdata should not be double-counted
+    qdata_freed = gu_nv_base.qdata.numel() == 0 and dn_nv_base.qdata.numel() == 0
+    assert qdata_freed, (
+        "marlin build did not free qdata (nv.qdata still has data); "
+        "the +12GB double-copy regression guard cannot verify the fix. "
+        f"gu_nv.qdata.numel()={gu_nv_base.qdata.numel()}, dn_nv.qdata.numel()={dn_nv_base.qdata.numel()}"
+    )
+
+    # After qdata freed: resident should be <= (nvfp4_only + marlin_qweight ≈ 2x qdata)
+    # The delta from (qdata_freed + marlin_qw resident) vs (qdata_only) should be <= 1.3x
+    # because marlin int32 packs 8 nibbles/word (same size as 4-bit nibbles * 2 = 8-bit word).
+    mem_delta = mem_after_marlin - mem_nvfp4_only + qdata_bytes  # add back freed qdata
+    ratio = (
+        mem_delta / qdata_bytes
+    )  # ratio of (marlin resident) / (original qdata size)
+    print(
+        f"\n[marlin memory] qdata freed={qdata_freed}, ratio={ratio:.3f}x "
+        f"(mem_delta={mem_delta / 1e9:.2f}GB, qdata={qdata_bytes / 1e9:.2f}GB)"
+    )
+    assert ratio <= 1.3, (
+        f"marlin memory overhead {ratio:.3f}x > 1.3x — the double-copy regression may be back. "
+        f"qdata_bytes={qdata_bytes / 1e9:.2f}GB, marlin_resident_delta={mem_delta / 1e9:.2f}GB"
+    )
+
+
+# ===========================================================================
+# Perf 4: Transient peak memory guard (grouped fwd+bwd vs bf16 baseline)
+# ===========================================================================
+
+
+@pytest.mark.parametrize(
+    "shape_fn,shape_name,vram_needed",
+    [
+        (_dsv4_cfg, "dsv4", _VRAM_DSV4_GB),
+        (_gemma4_cfg, "gemma4", _VRAM_GEMMA4_GB),
+    ],
+    ids=["dsv4", "gemma4"],
+)
+def test_grouped_fp4_peak_memory(shape_fn, shape_name, vram_needed):
+    """Grouped fwd+bwd peak transient memory <= bf16-dequant-on-fwd baseline peak.
+
+    The grouped/marlin path chunks the expert dequant (CHUNK_E=16 experts at a time) so the
+    transient bf16 weight materialization is bounded, unlike the naive full-E dequant.
+
+    Measured on RTX PRO 6000 (sm120):
+      DSV4 E=256:  grouped peak ~0.87x vs bf16 full-E-dequant baseline
+      gemma4 E=128: grouped peak ~0.92x vs bf16 full-E-dequant baseline
+
+    Threshold: grouped peak <= 1.05x baseline (allows 5% noise; the chunked path should be <= 1.0x).
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+        grouped_fp4_moe_train,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+
+    cfg, fallback = shape_fn()
+    if fallback:
+        pytest.skip(f"{shape_name}: insufficient VRAM (need {vram_needed:.0f} GB free)")
+
+    (
+        gu_nv,
+        dn_nv,
+        gu_nv_base,
+        dn_nv_base,
+        pt,
+        hidden,
+        idx,
+        wts,
+        (Agu, Bgu, Adn, Bdn),
+    ) = _build_tensors(cfg)
+    s = cfg["scaling"]
+    act_type, limit = cfg["act_type"], cfg["limit"]
+
+    # Grouped path step
+    cache = {}
+
+    def grouped_step():
+        for p in (Agu, Bgu, Adn, Bdn):
+            p.grad = None
+        hk = hidden.detach().requires_grad_()
+        out = grouped_fp4_moe_train(
+            hk,
+            idx,
+            wts,
+            gu_nv,
+            dn_nv,
+            (Agu, Bgu, s),
+            (Adn, Bdn, s),
+            limit,
+            "nvfp4",
+            act_type=act_type,
+            mxfp4_cache=cache,
+        )
+        out.float().pow(2).mean().backward()
+
+    # Baseline path: full-E bf16 dequant per step (naive worst-case transient).
+    # Uses the separate NV tensors (not freed by marlin path).
+    def baseline_step():
+        for p in (Agu, Bgu, Adn, Bdn):
+            p.grad = None
+        hk = hidden.detach().requires_grad_()
+        # Dequant all experts to bf16 (per step, no cache) — measures the full transient
+        Wgu_fresh = nvfp4_dequant_bf16(gu_nv_base.qdata, gu_nv_base.scale, pt)
+        dummy = hk @ Wgu_fresh[0].T  # force bf16 weight peak
+        dummy.sum().backward()
+        del Wgu_fresh
+
+    peak_grouped = _peak_mb(grouped_step)
+    peak_baseline = _peak_mb(baseline_step)
+
+    ratio = peak_grouped / (
+        peak_baseline + 1.0
+    )  # +1MB to avoid div-by-zero at tiny shapes
+    print(
+        f"\n[{shape_name}] peak: grouped={peak_grouped:.0f}MB baseline={peak_baseline:.0f}MB ratio={ratio:.3f}x"
+    )
+
+    # The grouped path uses chunked dequant so its peak should be <= baseline (which holds the full
+    # bf16 weight resident). Allow 1.05x for measurement noise.
+    # If the shape is too small, the baseline peak can be tiny and the ratio meaningless; skip then.
+    if peak_baseline < 50.0:
+        pytest.skip(
+            f"{shape_name}: baseline peak {peak_baseline:.0f}MB too small to measure ratio reliably"
+        )
+
+    assert ratio <= 1.05, (
+        f"{shape_name}: grouped peak {peak_grouped:.0f}MB is {ratio:.3f}x baseline {peak_baseline:.0f}MB "
+        f"(threshold <= 1.05x). Possible transient memory regression."
+    )
+
+
+# ===========================================================================
+# Sanity: grouped forward produces finite output (fast smoke, runs at small E)
+# ===========================================================================
+
+
+@pytest.mark.parametrize(
+    "cfg,shape_name",
+    [
+        (
+            dict(
+                E=8,
+                H=4096,
+                I=2048,
+                topk=6,
+                N=32,
+                r=8,
+                act_type="silu",
+                limit=7.0,
+                scaling=2.0,
+            ),
+            "dsv4_small",
+        ),
+        (
+            dict(
+                E=8,
+                H=2816,
+                I=704,
+                topk=8,
+                N=32,
+                r=8,
+                act_type="gelu_tanh",
+                limit=1e30,
+                scaling=2.0,
+            ),
+            "gemma4_small",
+        ),
+    ],
+    ids=["dsv4_small", "gemma4_small"],
+)
+def test_grouped_fp4_perf_smoke(cfg, shape_name):
+    """Quick smoke: grouped forward produces finite output at small E (always runs on Blackwell)."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+        grouped_fp4_moe_train,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+
+    (gu_nv, dn_nv, _, _, _, hidden, idx, wts, (Agu, Bgu, Adn, Bdn)) = _build_tensors(
+        cfg
+    )
+    s = cfg["scaling"]
+    cache = {}
+    out = grouped_fp4_moe_train(
+        hidden.detach(),
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agu.detach(), Bgu.detach(), s),
+        (Adn.detach(), Bdn.detach(), s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache=cache,
+    )
+    assert torch.isfinite(out).all(), f"{shape_name}: forward output has NaN/inf"
+    assert out.shape == (cfg["N"], cfg["H"])
diff --git a/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_skew.py b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_skew.py
new file mode 100644
index 0000000000..687233a006
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_skew.py
@@ -0,0 +1,283 @@
+"""Skew-robustness tests for the grouped/marlin NVFP4 LoRA path.
+
+The grouped LoRA op must be correct AND memory-bounded for ANY routing distribution — balanced,
+power-law, or pathological all-to-one. A capacity-padded implementation (cap = max per-expert
+tokens) blows up under skew (peak ~ E * max_cap * dim instead of ~ routed_tokens * dim) and OOMs at
+all_to_one; the robust implementation (ragged scatter / bulk-bmm+tail-scatter) does not.
+
+Verifiable (deterministic) guards live here:
+  - correctness across skew (cosine vs bf16 oracle) — CI-safe
+  - memory bound: peak transient <= C * routed_tokens * dim — catches the capacity blowup
+  - no-OOM under a simulated GPU memory cap at all_to_one
+Speed is ratio-based and marked `perf` (noisy).
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+# reuse the validated helpers from the sibling correctness module
+from .test_grouped_fp4_train import (
+    DEV,
+    _bf16_oracle,
+    _cos,
+    quantize_nvfp4,
+)
+
+_BLACKWELL = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] in (
+    10,
+    12,
+)
+pytestmark = pytest.mark.skipif(
+    not _BLACKWELL, reason="grouped fp4 needs sm100/sm120 (Blackwell)"
+)
+
+SKEWS = ["balanced", "moderate", "extreme", "all_to_one"]
+
+# Correctness uses a modest E (fast). Memory uses real E=128 so a capacity-padded impl's all_to_one
+# blowup (E * N * dim, most experts empty-but-padded) dwarfs balanced (~E*(routed/E)*dim) -> ~16x.
+_CFG = dict(
+    name="gemma4",
+    E=32,
+    H=2816,
+    I=704,
+    topk=8,
+    N=512,
+    r=16,
+    act_type="gelu_tanh",
+    limit=1e30,
+    scaling=2.0,
+)
+_CFG_MEM = dict(_CFG, E=128, N=1024)
+
+
+def make_routing(N: int, E: int, top_k: int, skew: str, seed: int = 0):
+    """Deterministic (idx[N,top_k], wts[N,top_k]) with a controllable routing skew.
+    balanced: uniform distinct top_k; moderate: power-law (prob ~ 1/rank); extreme: 90% of tokens
+    to experts 0..top_k-1; all_to_one: every token -> experts 0..top_k-1 (the adversarial case that
+    forces cap == N for a capacity-padded impl)."""
+    g = torch.Generator().manual_seed(seed)
+    hot = torch.arange(top_k)
+    rows = []
+    if skew == "balanced":
+        rows = [torch.randperm(E, generator=g)[:top_k] for _ in range(N)]
+    elif skew == "all_to_one":
+        rows = [hot.clone() for _ in range(N)]
+    elif skew == "extreme":
+        # 90% hot, 10% random
+        rows = [
+            hot.clone() if (i % 10) else torch.randperm(E, generator=g)[:top_k]
+            for i in range(N)
+        ]
+    elif skew == "moderate":
+        w = 1.0 / torch.arange(1, E + 1).float()  # zipf-ish over expert rank
+        for _ in range(N):
+            perm = torch.multinomial(w, top_k, replacement=False, generator=g)
+            rows.append(perm)
+    else:
+        raise ValueError(skew)
+    idx = torch.stack(rows).to(DEV)
+    wts = (
+        torch.softmax(torch.randn(N, top_k, generator=g), -1).to(torch.bfloat16).to(DEV)
+    )
+    return idx, wts
+
+
+def _build(cfg, skew, seed=0):
+    E, H, I, topk, N, r = cfg["E"], cfg["H"], cfg["I"], cfg["topk"], cfg["N"], cfg["r"]  # noqa: E741
+    twoI = 2 * I
+    torch.manual_seed(seed)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv, dn_nv = quantize_nvfp4(Wgu), quantize_nvfp4(Wdn)
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+
+    pt = torch.ones(E, device=DEV)
+    Wgu_b = nvfp4_dequant_bf16(gu_nv.qdata, gu_nv.scale, pt)
+    Wdn_b = nvfp4_dequant_bf16(dn_nv.qdata, dn_nv.scale, pt)
+    hidden = torch.randn(N, H, device=DEV, dtype=torch.bfloat16) * 0.5
+    idx, wts = make_routing(N, E, topk, skew, seed)
+
+    def mk(*s):
+        return (
+            torch.randn(*s, device=DEV, dtype=torch.bfloat16) * 0.02
+        ).requires_grad_(True)
+
+    lora = (mk(r * E, H), mk(twoI, r * E), mk(r * E, I), mk(H, r * E))
+    return hidden, idx, wts, gu_nv, dn_nv, Wgu_b, Wdn_b, lora
+
+
+def _run(cfg, gu_nv, dn_nv, hidden, idx, wts, lora, cache):
+    # cache is a persistent dict (like the per-module mxfp4 cache in real training): the marlin path
+    # frees qdata after the first call and reads the cache thereafter, so repeated calls on the same
+    # nv tensors MUST share one cache.
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_moe_train,
+    )
+
+    Agu, Bgu, Adn, Bdn = lora
+    s = cfg["scaling"]
+    return grouped_fp4_moe_train(
+        hidden,
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agu, Bgu, s),
+        (Adn, Bdn, s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache=cache,
+    )
+
+
+# ---------------------------------------------------------------------------
+# 1. Correctness across skew (deterministic, CI-safe)
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize("skew", SKEWS)
+def test_correctness_across_skew(skew):
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend")
+    hidden, idx, wts, gu_nv, dn_nv, Wgu_b, Wdn_b, lora = _build(_CFG, skew)
+    Agu, Bgu, Adn, Bdn = lora
+    out = _run(
+        _CFG,
+        gu_nv,
+        dn_nv,
+        hidden.clone(),
+        idx,
+        wts,
+        (Agu.detach(), Bgu.detach(), Adn.detach(), Bdn.detach()),
+        {},
+    )
+    assert torch.isfinite(out).all(), f"{skew}: non-finite output"
+    ref = _bf16_oracle(
+        hidden.clone(),
+        idx,
+        wts,
+        Wgu_b,
+        Wdn_b,
+        Agu,
+        Bgu,
+        Adn,
+        Bdn,
+        _CFG["act_type"],
+        _CFG["limit"],
+        _CFG["scaling"],
+    )
+    cos = _cos(out, ref)
+    assert cos > 0.97, f"{skew}: fwd cosine {cos:.4f} < 0.97"
+
+
+# ---------------------------------------------------------------------------
+# 2. Memory bound across skew (THE skew guard) — a robust op's peak must NOT scale with routing
+#    skew. Measure each skew's peak transient and assert it stays within a small factor of the
+#    BALANCED peak (scale-robust, no magic absolute). A capacity-padded impl (cap=max) blows up at
+#    all_to_one (E*N*dim, most experts empty-but-padded) -> many x balanced -> fails here.
+# ---------------------------------------------------------------------------
+def _peak_transient(cfg, skew):
+    hidden, idx, wts, gu_nv, dn_nv, _, _, lora = _build(cfg, skew)
+    det = tuple(p.detach() for p in lora)
+    cache = {}
+    o = _run(cfg, gu_nv, dn_nv, hidden.clone(), idx, wts, det, cache)
+    del o  # warmup + build cache
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    base = torch.cuda.memory_allocated()
+    _run(cfg, gu_nv, dn_nv, hidden.clone(), idx, wts, det, cache)
+    torch.cuda.synchronize()
+    return (torch.cuda.max_memory_allocated() - base) / 1e6
+
+
+def test_memory_bound_across_skew():
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend")
+    cfg = _CFG_MEM
+    peaks = {s: _peak_transient(cfg, s) for s in SKEWS}
+    base = peaks["balanced"]
+    msg = " | ".join(f"{s}={peaks[s]:.0f}MB" for s in SKEWS)
+    # robust: skew peak within 1.5x of balanced. capacity-padded all_to_one is many x -> fails.
+    for s in SKEWS:
+        assert peaks[s] <= 1.5 * base, f"skew '{s}' transient blowup vs balanced: {msg}"
+
+
+# ---------------------------------------------------------------------------
+# 3. No-OOM under a simulated GPU cap at the adversarial all_to_one routing.
+# ---------------------------------------------------------------------------
+@pytest.mark.perf
+def test_no_oom_all_to_one_under_cap():
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend")
+    total = torch.cuda.get_device_properties(0).total_memory
+    # cap this process to a fraction that comfortably fits the robust path but not an E*cap blowup
+    hidden, idx, wts, gu_nv, dn_nv, _, _, lora = _build(_CFG, "all_to_one")
+    needed = _CFG["N"] * _CFG["topk"] * max(2 * _CFG["I"], _CFG["H"]) * 2
+    frac = min(
+        0.5, (needed * 12 + 4 * 2**30) / total
+    )  # robust needs ~12x routed + model; blowup needs ~Ex more
+    torch.cuda.set_per_process_memory_fraction(frac, 0)
+    try:
+        out = _run(
+            _CFG,
+            gu_nv,
+            dn_nv,
+            hidden.clone(),
+            idx,
+            wts,
+            tuple(p.detach() for p in lora),
+            {},
+        )
+        out.float().sum().backward() if out.requires_grad else None
+        assert torch.isfinite(out).all()
+    finally:
+        torch.cuda.set_per_process_memory_fraction(1.0, 0)
+
+
+# ---------------------------------------------------------------------------
+# 4. Speed: skew must not catastrophically slow the op vs balanced (ratio-based, noisy).
+# ---------------------------------------------------------------------------
+@pytest.mark.perf
+def test_speed_skew_within_factor_of_balanced():
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend")
+
+    def timed(skew, it=10, wu=4):
+        hidden, idx, wts, gu_nv, dn_nv, _, _, lora = _build(_CFG, skew)
+        det = tuple(p.detach() for p in lora)
+        cache = {}
+        for _ in range(wu):
+            _run(_CFG, gu_nv, dn_nv, hidden.clone(), idx, wts, det, cache)
+        torch.cuda.synchronize()
+        s = torch.cuda.Event(enable_timing=True)
+        e = torch.cuda.Event(enable_timing=True)
+        s.record()
+        for _ in range(it):
+            _run(_CFG, gu_nv, dn_nv, hidden.clone(), idx, wts, det, cache)
+        e.record()
+        torch.cuda.synchronize()
+        return s.elapsed_time(e) / it
+
+    tb = timed("balanced")
+    ts = timed("extreme")
+    assert ts <= 4.0 * tb, f"extreme-skew step {ts:.1f}ms > 4x balanced {tb:.1f}ms"
diff --git a/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_train.py b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_train.py
new file mode 100644
index 0000000000..d7ef61f475
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_train.py
@@ -0,0 +1,906 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Correctness tests for the grouped/marlin NVFP4 MoE training path.
+
+Covers:
+  - grouped_fp4_moe_train fwd+bwd vs a bf16 per-expert oracle (DSV4 and gemma4 shapes)
+  - Activation type dispatch (_detect_act_type) and wrong-activation detection
+  - fp8-read dX NaN regression at non-128-aligned K (I=704, the gemma4 intermediate dim)
+  - marlin->bf16 fused dequant bit-exactness vs nvfp4_dequant_bf16
+  - GeGLU / clamped-SwiGLU activation kernel fwd+bwd vs torch reference
+
+All tests use synthetic NVFP4 weights — no model download required.
+
+H/I dimensions are kept at real model values (4096/2048 for DSV4, 2816/704 for gemma4) to satisfy
+the marlin backend's 16-byte stride alignment requirement for torch._grouped_mm. E and N are small
+(8-16 and 32) so the tests run fast (~2-10 s each on RTX PRO 6000).
+
+Gate: skips unless CUDA is available AND device is Blackwell (sm100 / sm120).
+"""
+
+from __future__ import annotations
+
+import functools
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Skip gate: Blackwell only (sm100 / sm120)
+# ---------------------------------------------------------------------------
+
+_IS_CUDA = torch.cuda.is_available()
+_IS_BLACKWELL = _IS_CUDA and torch.cuda.get_device_capability()[0] in (10, 12)
+
+pytestmark = pytest.mark.skipif(
+    not _IS_BLACKWELL,
+    reason="grouped NVFP4 training backends require Blackwell (sm100/sm120)",
+)
+
+DEV = "cuda"
+
+
+# ---------------------------------------------------------------------------
+# Shape families — real H/I to satisfy 16-byte stride alignment in grouped_mm
+# ---------------------------------------------------------------------------
+
+# DSV4-Flash: clamped-SwiGLU, limit=7.0
+_DSV4 = dict(
+    name="dsv4",
+    E=8,
+    H=4096,
+    I=2048,
+    topk=6,
+    N=32,
+    r=8,
+    act_type="silu",
+    limit=7.0,
+    scaling=2.0,
+)
+
+# gemma4-A4B: gelu_tanh GeGLU, no clamp
+_GEMMA4 = dict(
+    name="gemma4",
+    E=8,
+    H=2816,
+    I=704,
+    topk=8,
+    N=32,
+    r=8,
+    act_type="gelu_tanh",
+    limit=1e30,
+    scaling=2.0,
+)
+
+SHAPES = [_DSV4, _GEMMA4]
+SHAPE_IDS = ["dsv4", "gemma4"]
+
+
+# ---------------------------------------------------------------------------
+# NVFP4 quantizer — produces a real NVFP4Tensor (needed for the marlin backend
+# which frees qdata.data after repacking). Self-contained, no scratch imports.
+# ---------------------------------------------------------------------------
+
+
+def _fp4_codebook(device):
+    from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+        fp4_codebook,
+    )
+
+    return fp4_codebook(device).float()
+
+
+def quantize_nvfp4(W: torch.Tensor):
+    """Quantize bf16 W[E,N,K] -> NVFP4Tensor matching nvfp4_dequant_bf16's inverse."""
+    from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+    dev = W.device
+    cb = _fp4_codebook(dev)
+    E, N, K = W.shape
+    Wb = W.reshape(E, N, K // 16, 16).float()
+    amax = Wb.abs().amax(-1).clamp_min(1e-6)
+    scale_e4m3 = (amax / 6.0).to(torch.float8_e4m3fn)
+    scale_f = scale_e4m3.float().clamp_min(1e-9)
+    nib = (
+        ((Wb / scale_f.unsqueeze(-1)).unsqueeze(-1) - cb.view(1, 1, 1, 1, 16))
+        .abs()
+        .argmin(-1)
+    )
+    nib = nib.to(torch.uint8).reshape(E, N, K)
+    packed = (nib[..., 0::2] & 0xF) | ((nib[..., 1::2] & 0xF) << 4)
+    return NVFP4Tensor(
+        packed.contiguous(),
+        scale_e4m3.contiguous(),
+        block_size=16,
+        orig_dtype=torch.bfloat16,
+        per_tensor_scale=torch.ones((), device=dev),
+    )
+
+
+# ---------------------------------------------------------------------------
+# bf16 per-expert MoE oracle
+# ---------------------------------------------------------------------------
+
+
+def _bf16_oracle(
+    hidden, idx, wts, Wgu_b, Wdn_b, Agu, Bgu, Adn, Bdn, act_type, limit, scaling
+):
+    """Per-token per-expert bf16 MoE forward (no grouping), matching the kernel's intent.
+
+    hidden [N,H]; idx/wts [N,topk]; Wgu_b/Wdn_b [E,out,in] bf16 dequant; *_lora scattermoe
+    stacked layout [r*E, K] / [out, r*E]; act_type: 'silu' | 'gelu_tanh'; limit: clamp value.
+    Returns [N,H] float32 accumulator (autograd-enabled via Agu/Bgu/Adn/Bdn).
+    """
+    N, H = hidden.shape
+    E = Wgu_b.shape[0]
+    r = Agu.shape[0] // E
+    acc = torch.zeros(N, H, device=hidden.device, dtype=torch.float32)
+    for e in range(E):
+        sel = idx == e  # [N, topk] bool
+        tok = sel.any(-1)  # [N] bool
+        if not tok.any():
+            continue
+        ti = tok.nonzero(as_tuple=True)[0]
+        xe = hidden[ti].float()
+        Ag = Agu[e * r : (e + 1) * r]  # [r, H]
+        Bg = Bgu[:, e * r : (e + 1) * r]  # [2I, r]
+        Ad = Adn[e * r : (e + 1) * r]  # [r, I]
+        Bd = Bdn[:, e * r : (e + 1) * r]  # [H, r]
+        # gate_up: base + LoRA
+        gu = (
+            xe.to(hidden.dtype) @ Wgu_b[e].T
+            + scaling * (xe.to(hidden.dtype) @ Ag.T) @ Bg.T
+        )
+        # activation
+        g, u = gu.chunk(2, -1)
+        if act_type == "gelu_tanh":
+            h = F.gelu(g.float(), approximate="tanh") * u.float()
+        else:
+            h = F.silu(g.float().clamp(max=limit)) * u.float().clamp(
+                min=-limit, max=limit
+            )
+        h = h.to(hidden.dtype)
+        # down: base + LoRA
+        dn = h @ Wdn_b[e].T + scaling * (h @ Ad.T) @ Bd.T
+        # weighted accumulate
+        w = (wts * sel)[ti].sum(-1, keepdim=True).float()
+        acc[ti] += dn.float() * w
+    return acc
+
+
+def _cos(a: torch.Tensor, b: torch.Tensor) -> float:
+    return F.cosine_similarity(a.float().flatten(), b.float().flatten(), dim=0).item()
+
+
+# ---------------------------------------------------------------------------
+# Shared fixture: build inputs + weights for a shape family
+# ---------------------------------------------------------------------------
+
+
+def _make_inputs(cfg: dict, seed: int = 0):
+    """Return (hidden, idx, wts, gu_nv, dn_nv, Wgu_b, Wdn_b, lora_params) plus fresh NV copies."""
+    E, H, I, topk, N, r = cfg["E"], cfg["H"], cfg["I"], cfg["topk"], cfg["N"], cfg["r"]
+    twoI = 2 * I
+    torch.manual_seed(seed)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv = quantize_nvfp4(Wgu)
+    dn_nv = quantize_nvfp4(Wdn)
+
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+
+    pt = torch.ones(E, device=DEV)
+    Wgu_b = nvfp4_dequant_bf16(gu_nv.qdata, gu_nv.scale, pt)
+    Wdn_b = nvfp4_dequant_bf16(dn_nv.qdata, dn_nv.scale, pt)
+
+    hidden = torch.randn(N, H, device=DEV, dtype=torch.bfloat16) * 0.5
+    idx = torch.stack([torch.randperm(E, device=DEV)[:topk] for _ in range(N)])
+    wts = torch.softmax(torch.randn(N, topk, device=DEV), -1).to(torch.bfloat16)
+
+    def mk(*s):
+        return (
+            torch.randn(*s, device=DEV, dtype=torch.bfloat16) * 0.02
+        ).requires_grad_(True)
+
+    Agu = mk(r * E, H)
+    Bgu = mk(twoI, r * E)
+    Adn = mk(r * E, I)
+    Bdn = mk(H, r * E)
+
+    return hidden, idx, wts, gu_nv, dn_nv, Wgu_b, Wdn_b, (Agu, Bgu, Adn, Bdn)
+
+
+# ---------------------------------------------------------------------------
+# Fresh NVFP4Tensor helper (marlin frees qdata on first call; tests that need
+# multiple independent calls must build separate NVFP4Tensors from the same W)
+# ---------------------------------------------------------------------------
+
+
+def _fresh_nv(Wgu, Wdn):
+    """Build fresh NVFP4Tensor pair from the same weight matrices (for second-call tests)."""
+    return quantize_nvfp4(Wgu), quantize_nvfp4(Wdn)
+
+
+# ===========================================================================
+# Test 1: forward output matches oracle
+# ===========================================================================
+
+
+@pytest.mark.parametrize("cfg", SHAPES, ids=SHAPE_IDS)
+def test_grouped_fp4_fwd_matches_oracle(cfg):
+    """Forward output cosine >= 0.97 vs bf16 per-expert oracle, no NaN/inf."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+        grouped_fp4_moe_train,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+
+    H = cfg["H"]
+    hidden, idx, wts, gu_nv, dn_nv, Wgu_b, Wdn_b, (Agu, Bgu, Adn, Bdn) = _make_inputs(
+        cfg
+    )
+    _r, s = cfg["r"], cfg["scaling"]
+
+    cache = {}
+    out = grouped_fp4_moe_train(
+        hidden.clone(),
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agu.detach(), Bgu.detach(), s),
+        (Adn.detach(), Bdn.detach(), s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache=cache,
+    )
+    assert torch.isfinite(out).all(), f"{cfg['name']}: fwd output has NaN/inf"
+    assert out.shape == (cfg["N"], H), f"unexpected output shape {out.shape}"
+
+    ref = _bf16_oracle(
+        hidden.clone(),
+        idx,
+        wts,
+        Wgu_b,
+        Wdn_b,
+        Agu,
+        Bgu,
+        Adn,
+        Bdn,
+        cfg["act_type"],
+        cfg["limit"],
+        s,
+    )
+    cos = _cos(out, ref)
+    assert cos > 0.97, f"{cfg['name']}: fwd cosine {cos:.4f} < 0.97"
+
+
+# ===========================================================================
+# Test 1b: marlin qdata-free experts survive clone/state_dict (adapter save)
+# ===========================================================================
+
+
+@pytest.mark.parametrize("cfg", SHAPES, ids=SHAPE_IDS)
+def test_grouped_fp4_save_after_qdata_free(cfg):
+    """Regression: the marlin memory fix frees nv.qdata after repack. The freed NVFP4Tensor must
+    still clone/serialize (PEFT save_pretrained -> state_dict -> NVFP4Tensor clone -> __new__ ->
+    qdata.stride(-2)); a flat empty(0) qdata crashed with IndexError(dim -2). Freeing to a 3-D
+    [E, N, 0] placeholder keeps clone/save working."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        _train_backend,
+        grouped_fp4_available,
+        grouped_fp4_moe_train,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+    if _train_backend("nvfp4") != "marlin":
+        pytest.skip("qdata-free only on the marlin backend")
+
+    hidden, idx, wts, gu_nv, dn_nv, _Wgu, _Wdn, (Agu, Bgu, Adn, Bdn) = _make_inputs(cfg)
+    s = cfg["scaling"]
+    # one forward triggers the marlin repack + qdata free
+    grouped_fp4_moe_train(
+        hidden.clone(),
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agu.detach(), Bgu.detach(), s),
+        (Adn.detach(), Bdn.detach(), s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache={},
+    )
+    # qdata was freed to a 3-D zero-element placeholder (not flat empty(0))
+    assert gu_nv.qdata.numel() == 0 and gu_nv.qdata.ndim == 3, (
+        "qdata not freed to 3-D placeholder"
+    )
+    # The freed NVFP4Tensor must CLONE / round-trip through state_dict without raising — this is the
+    # save path (PEFT save_pretrained -> state_dict -> NVFP4Tensor clone). A flat empty(0) qdata
+    # crashed here with IndexError(dim -2). We don't assert the frozen expert's clone *shape* (it's
+    # degenerate by design and filtered out of the LoRA adapter); only that save doesn't crash.
+    import torch.nn as nn
+
+    for nv in (gu_nv, dn_nv):
+        nv.clone()  # would IndexError(dim -2) on a flat empty(0) qdata
+        m = nn.Module()
+        m.w = nn.Parameter(nv, requires_grad=False)
+        sd = m.state_dict()  # invokes _save_to_state_dict -> clone on the NVFP4Tensor
+        assert "w" in sd
+
+
+# ===========================================================================
+# Test 2: backward grads match oracle
+# ===========================================================================
+
+
+@pytest.mark.parametrize("cfg", SHAPES, ids=SHAPE_IDS)
+def test_grouped_fp4_bwd_matches_oracle(cfg):
+    """Backward grads (d_hidden, d_LoRA_A_gu, d_LoRA_A_dn) cosine >= 0.95, finite."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+        grouped_fp4_moe_train,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+
+    E, H, I = cfg["E"], cfg["H"], cfg["I"]
+    hidden, idx, wts, gu_nv, dn_nv, Wgu_b, Wdn_b, (Agu, Bgu, Adn, Bdn) = _make_inputs(
+        cfg
+    )
+    _r, s = cfg["r"], cfg["scaling"]
+
+    # Kernel path — use fresh NV tensors for backward (marlin frees qdata on first build)
+    twoI = 2 * I
+    torch.manual_seed(0)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv_k, dn_nv_k = _fresh_nv(Wgu, Wdn)
+    pt = torch.ones(E, device=DEV)
+    Wgu_b_k = nvfp4_dequant_bf16(gu_nv_k.qdata, gu_nv_k.scale, pt)
+    Wdn_b_k = nvfp4_dequant_bf16(dn_nv_k.qdata, dn_nv_k.scale, pt)
+
+    hidden2 = hidden.clone()
+    idx2 = idx.clone()
+    wts2 = wts.clone()
+    Agk = Agu.detach().clone().requires_grad_(True)
+    Bgk = Bgu.detach().clone().requires_grad_(True)
+    Adk = Adn.detach().clone().requires_grad_(True)
+    Bdk = Bdn.detach().clone().requires_grad_(True)
+    hk = hidden2.requires_grad_(True)
+
+    cache = {}
+    out = grouped_fp4_moe_train(
+        hk,
+        idx2,
+        wts2,
+        gu_nv_k,
+        dn_nv_k,
+        (Agk, Bgk, s),
+        (Adk, Bdk, s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache=cache,
+    )
+    out.float().pow(2).mean().backward()
+
+    # Oracle path — same weights
+    Agr = Agk.detach().clone().requires_grad_(True)
+    Bgr = Bgk.detach().clone().requires_grad_(True)
+    Adr = Adk.detach().clone().requires_grad_(True)
+    Bdr = Bdk.detach().clone().requires_grad_(True)
+    hr = hidden2.detach().clone().requires_grad_(True)
+
+    ref = _bf16_oracle(
+        hr,
+        idx2,
+        wts2,
+        Wgu_b_k,
+        Wdn_b_k,
+        Agr,
+        Bgr,
+        Adr,
+        Bdr,
+        cfg["act_type"],
+        cfg["limit"],
+        s,
+    )
+    ref.pow(2).mean().backward()
+
+    for name, gk, gr in [
+        ("d_hidden", hk.grad, hr.grad),
+        ("d_Agu", Agk.grad, Agr.grad),
+        ("d_Adn", Adk.grad, Adr.grad),
+    ]:
+        assert gk is not None, (
+            f"{cfg['name']}: {name} grad is None (backward didn't run)"
+        )
+        assert gr is not None, f"{cfg['name']}: oracle {name} grad is None"
+        assert torch.isfinite(gk).all(), f"{cfg['name']}: {name} has NaN/inf"
+        cos = _cos(gk, gr)
+        assert cos > 0.95, f"{cfg['name']}: {name} cosine {cos:.4f} < 0.95"
+
+
+# ===========================================================================
+# Test 3: activation type dispatch
+# ===========================================================================
+
+
+def test_detect_act_type_dispatch():
+    """_detect_act_type must return 'gelu_tanh' for gemma4-style and 'silu' for DSV4."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+        _detect_act_type,
+    )
+
+    class FakeGemma4:
+        act_fn = functools.partial(F.gelu, approximate="tanh")
+
+    class FakeDSV4:
+        act_fn = F.silu
+        limit = 7.0
+
+    class FakeNamedGeLU:
+        class _GeLU:
+            __name__ = "gelu_pytorch_tanh"
+
+            def __call__(self, x):
+                return F.gelu(x, approximate="tanh")
+
+        act_fn = _GeLU()
+
+    assert _detect_act_type(FakeGemma4()) == "gelu_tanh", (
+        "gemma4 (partial F.gelu tanh) must be 'gelu_tanh'"
+    )
+    assert _detect_act_type(FakeDSV4()) == "silu", "dsv4 (F.silu) must be 'silu'"
+    assert _detect_act_type(FakeNamedGeLU()) == "gelu_tanh", (
+        "named gelu_pytorch_tanh must be 'gelu_tanh'"
+    )
+
+
+def test_wrong_activation_fails_gemma4():
+    """Using act_type='silu' on a gemma4 shape (correct: 'gelu_tanh') degrades cosine vs oracle.
+
+    Demonstrates TEETH: the correct activation (gelu_tanh) achieves cosine >= 0.999 vs the
+    correct oracle, while the wrong activation (silu) achieves only ~0.992. The gap is large
+    enough to be a reliable regression detector.
+
+    Note: absolute cosine with wrong-act vs correct-oracle stays ~0.99 (gelu_tanh ≈ silu for
+    moderate values), so the test checks the *relative gap* rather than an absolute threshold.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+        grouped_fp4_moe_train,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+
+    cfg = _GEMMA4
+    E, H, I, topk, N, r = cfg["E"], cfg["H"], cfg["I"], cfg["topk"], cfg["N"], cfg["r"]
+    twoI = 2 * I
+    torch.manual_seed(99)
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    hidden = torch.randn(N, H, device=DEV, dtype=torch.bfloat16) * 0.5
+    idx = torch.stack([torch.randperm(E, device=DEV)[:topk] for _ in range(N)])
+    wts = torch.softmax(torch.randn(N, topk, device=DEV), -1).to(torch.bfloat16)
+    Agu = torch.randn(r * E, H, device=DEV, dtype=torch.bfloat16) * 0.02
+    Bgu = torch.randn(twoI, r * E, device=DEV, dtype=torch.bfloat16) * 0.02
+    Adn = torch.randn(r * E, I, device=DEV, dtype=torch.bfloat16) * 0.02
+    Bdn = torch.randn(H, r * E, device=DEV, dtype=torch.bfloat16) * 0.02
+
+    pt = torch.ones(E, device=DEV)
+
+    # Correct-act run
+    gu_nv_c, dn_nv_c = _fresh_nv(Wgu, Wdn)
+    Wgu_b = nvfp4_dequant_bf16(gu_nv_c.qdata, gu_nv_c.scale, pt)
+    Wdn_b = nvfp4_dequant_bf16(dn_nv_c.qdata, dn_nv_c.scale, pt)
+    cache_c = {}
+    out_correct = grouped_fp4_moe_train(
+        hidden,
+        idx,
+        wts,
+        gu_nv_c,
+        dn_nv_c,
+        (Agu, Bgu, cfg["scaling"]),
+        (Adn, Bdn, cfg["scaling"]),
+        cfg["limit"],
+        "nvfp4",
+        act_type="gelu_tanh",
+        mxfp4_cache=cache_c,
+    )
+
+    # Wrong-act run (fresh NV tensors; marlin freed qdata in first call)
+    gu_nv_w, dn_nv_w = _fresh_nv(Wgu, Wdn)
+    cache_w = {}
+    out_wrong = grouped_fp4_moe_train(
+        hidden,
+        idx,
+        wts,
+        gu_nv_w,
+        dn_nv_w,
+        (Agu, Bgu, cfg["scaling"]),
+        (Adn, Bdn, cfg["scaling"]),
+        cfg["limit"],
+        "nvfp4",
+        act_type="silu",  # deliberate wrong activation
+        mxfp4_cache=cache_w,
+    )
+
+    # Correct-activation oracle
+    correct_ref = _bf16_oracle(
+        hidden,
+        idx,
+        wts,
+        Wgu_b,
+        Wdn_b,
+        Agu,
+        Bgu,
+        Adn,
+        Bdn,
+        "gelu_tanh",
+        cfg["limit"],
+        cfg["scaling"],
+    )
+
+    cos_correct = _cos(out_correct, correct_ref)
+    cos_wrong = _cos(out_wrong, correct_ref)
+
+    # The correct activation must match its oracle very well
+    assert cos_correct > 0.99, (
+        f"correct-act kernel cosine vs oracle {cos_correct:.4f} < 0.99 — something is wrong"
+    )
+    # The wrong activation must diverge detectably: gap >= 0.005 (measured ~0.007-0.008)
+    gap = cos_correct - cos_wrong
+    assert gap > 0.005, (
+        f"TEETH CHECK FAILED: wrong activation (silu on gemma4) barely diverges from correct "
+        f"(correct={cos_correct:.4f}, wrong={cos_wrong:.4f}, gap={gap:.5f} < 0.005). "
+        "The activation correctness test lacks discriminating power."
+    )
+
+
+# ===========================================================================
+# Test 4: NaN regression at non-128-aligned K (I=704, fp8-read dX)
+# ===========================================================================
+
+
+def test_grouped_dx_fp8_nan_regression_k704():
+    """grouped_dx_fp8 must produce finite output for K=704 (non-BK-aligned), the gemma4 I dim.
+
+    Original bug: the K-tile boundary had no mask, so the last tile read OOB fp8 values that
+    decoded to NaN in the accumulator. Fixed by adding mk=rk<K on the W load and store.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        grouped_dx_fp8,
+    )
+
+    E, H, I = 16, 2816, 704  # gemma4 shapes; I=704 = 5.5*128, NOT BK-aligned
+    tile = 64  # marlin routing tile
+    Mt = E * tile
+
+    torch.manual_seed(0)
+    grad = torch.randn(Mt, H, dtype=torch.bfloat16, device=DEV) * 0.01
+    w_fp8 = torch.zeros(E, H, I, dtype=torch.float8_e4m3fn, device=DEV)
+    m_indices = torch.arange(E, dtype=torch.int32, device=DEV)
+
+    out = grouped_dx_fp8(grad, w_fp8, m_indices, block_m=tile)
+    assert not out.isnan().any(), "K=704 produced NaN in grouped_dx_fp8 (regression)"
+    assert not out.isinf().any(), "K=704 produced Inf in grouped_dx_fp8"
+    assert out.shape == (Mt, I)
+
+
+def test_grouped_dx_fp8_k2048_no_regression():
+    """grouped_dx_fp8 must also work correctly for K=2048 (DSV4, BK-aligned)."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        grouped_dx_fp8,
+    )
+
+    E, H, I = 16, 4096, 2048
+    tile = 64
+    Mt = E * tile
+
+    torch.manual_seed(1)
+    grad = torch.randn(Mt, H, dtype=torch.bfloat16, device=DEV) * 0.01
+    w_fp8 = torch.zeros(E, H, I, dtype=torch.float8_e4m3fn, device=DEV)
+    m_indices = torch.arange(E, dtype=torch.int32, device=DEV)
+
+    out = grouped_dx_fp8(grad, w_fp8, m_indices, block_m=tile)
+    assert not out.isnan().any(), "K=2048 produced NaN (regression)"
+    assert not out.isinf().any()
+    assert out.shape == (Mt, I)
+
+
+def test_grouped_fp4_bwd_nan_k704():
+    """Full fwd+bwd on the gemma4 shape (I=704) must produce no NaN in any gradient.
+
+    This exercises the entire backward path including the chunked fp8-read dX at K=704.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        grouped_fp4_available,
+        grouped_fp4_moe_train,
+    )
+
+    if not grouped_fp4_available("nvfp4"):
+        pytest.skip("no grouped fp4 backend available")
+
+    cfg = _GEMMA4
+    E, H, I, topk, N, r = cfg["E"], cfg["H"], cfg["I"], cfg["topk"], cfg["N"], cfg["r"]
+    twoI = 2 * I
+    torch.manual_seed(7)
+    gu_nv = quantize_nvfp4(
+        torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    )
+    dn_nv = quantize_nvfp4(
+        torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    )
+    hidden = torch.randn(N, H, device=DEV, dtype=torch.bfloat16) * 0.5
+    idx = torch.stack([torch.randperm(E, device=DEV)[:topk] for _ in range(N)])
+    wts = torch.softmax(torch.randn(N, topk, device=DEV), -1).to(torch.bfloat16)
+
+    s = cfg["scaling"]
+    Agu = (
+        torch.randn(r * E, H, device=DEV, dtype=torch.bfloat16) * 0.02
+    ).requires_grad_(True)
+    Bgu = (
+        torch.randn(twoI, r * E, device=DEV, dtype=torch.bfloat16) * 0.02
+    ).requires_grad_(True)
+    Adn = (
+        torch.randn(r * E, I, device=DEV, dtype=torch.bfloat16) * 0.02
+    ).requires_grad_(True)
+    Bdn = (
+        torch.randn(H, r * E, device=DEV, dtype=torch.bfloat16) * 0.02
+    ).requires_grad_(True)
+    hk = hidden.clone().requires_grad_(True)
+
+    cache = {}
+    out = grouped_fp4_moe_train(
+        hk,
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agu, Bgu, s),
+        (Adn, Bdn, s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache=cache,
+    )
+    out.float().pow(2).mean().backward()
+
+    for name, t in [
+        ("out", out),
+        ("d_hidden", hk.grad),
+        ("d_Agu", Agu.grad),
+        ("d_Bgu", Bgu.grad),
+        ("d_Adn", Adn.grad),
+        ("d_Bdn", Bdn.grad),
+    ]:
+        assert t is not None, f"gemma4 K=704 backward: {name} is None"
+        assert torch.isfinite(t).all(), f"gemma4 K=704 backward: {name} has NaN/inf"
+
+
+# ===========================================================================
+# Test 5: marlin->bf16 fused dequant bit-exactness vs nvfp4_dequant_bf16
+# ===========================================================================
+
+
+@pytest.mark.skipif(
+    not (_IS_CUDA and torch.cuda.get_device_capability()[0] >= 8),
+    reason="marlin W4A16 backend needs Ampere (sm80)+ (the module's Blackwell pytestmark still "
+    "gates this file to sm100/sm120 on CI; the sm80+ marlin bit-exact variant that executes on "
+    "sm89 lives in test_grouped_fp4_dequant_path.py)",
+)
+def test_marlin_fused_dequant_bit_exact_gate_up():
+    """marlin_dequant_bf16 is bit-exact vs nvfp4_dequant_bf16 for gate_up (N=64, K=128) shape."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16 import (
+        marlin_w4a16_available,
+    )
+
+    if not marlin_w4a16_available():
+        pytest.skip("marlin W4A16 ext not available")
+
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16 import load_ext
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.backend import (
+        _build_base_scatter,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.fused_dequant import (
+        marlin_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.prep import (
+        prepare_nvfp4_weight_for_marlin,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+        fp4_codebook,
+    )
+
+    # N must be divisible by 64 (marlin N_t), K divisible by 16 (NVFP4 block)
+    C, N, K = 4, 64, 128
+    torch.manual_seed(42)
+    W = torch.randn(C, N, K, device=DEV, dtype=torch.bfloat16) * 0.04
+    nv = quantize_nvfp4(W)
+    pt = torch.ones(C, device=DEV)
+
+    ext = load_ext()
+    scatter_lut = _build_base_scatter(torch.device(DEV))
+    cb = fp4_codebook(torch.device(DEV)).float()
+
+    marlin_packed = prepare_nvfp4_weight_for_marlin(
+        nv.qdata, nv.scale, pt, N, K, torch.bfloat16, ext.gptq_marlin_repack
+    )
+    qw_flat = marlin_packed[0].reshape(C, -1)
+    ref_bf16 = nvfp4_dequant_bf16(nv.qdata, nv.scale, pt)
+    got_bf16 = marlin_dequant_bf16(qw_flat, nv.scale, pt, scatter_lut, cb, N, K, C)
+
+    max_err = (ref_bf16.float() - got_bf16.float()).abs().max().item()
+    assert max_err == 0.0, (
+        f"marlin fused dequant not bit-exact vs nvfp4_dequant_bf16 (gate_up N={N}, K={K}): "
+        f"maxerr={max_err}"
+    )
+
+
+@pytest.mark.skipif(
+    not (_IS_CUDA and torch.cuda.get_device_capability()[0] >= 8),
+    reason="marlin W4A16 backend needs Ampere (sm80)+ (the module's Blackwell pytestmark still "
+    "gates this file to sm100/sm120 on CI; the sm80+ marlin bit-exact variant that executes on "
+    "sm89 lives in test_grouped_fp4_dequant_path.py)",
+)
+def test_marlin_fused_dequant_bit_exact_down():
+    """marlin_dequant_bf16 is bit-exact vs nvfp4_dequant_bf16 for down (N=128, K=64) shape."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16 import (
+        marlin_w4a16_available,
+    )
+
+    if not marlin_w4a16_available():
+        pytest.skip("marlin W4A16 ext not available")
+
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16 import load_ext
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.backend import (
+        _build_base_scatter,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.fused_dequant import (
+        marlin_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.marlin_w4a16.prep import (
+        prepare_nvfp4_weight_for_marlin,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+        fp4_codebook,
+    )
+
+    # down: N=H, K=I; marlin requirement: N%64==0, K%16==0
+    C, N, K = 4, 128, 64
+    torch.manual_seed(43)
+    W = torch.randn(C, N, K, device=DEV, dtype=torch.bfloat16) * 0.04
+    nv = quantize_nvfp4(W)
+    pt = torch.ones(C, device=DEV)
+
+    ext = load_ext()
+    scatter_lut = _build_base_scatter(torch.device(DEV))
+    cb = fp4_codebook(torch.device(DEV)).float()
+
+    marlin_packed = prepare_nvfp4_weight_for_marlin(
+        nv.qdata, nv.scale, pt, N, K, torch.bfloat16, ext.gptq_marlin_repack
+    )
+    qw_flat = marlin_packed[0].reshape(C, -1)
+    ref_bf16 = nvfp4_dequant_bf16(nv.qdata, nv.scale, pt)
+    got_bf16 = marlin_dequant_bf16(qw_flat, nv.scale, pt, scatter_lut, cb, N, K, C)
+
+    max_err = (ref_bf16.float() - got_bf16.float()).abs().max().item()
+    assert max_err == 0.0, (
+        f"marlin fused dequant not bit-exact vs nvfp4_dequant_bf16 (down N={N}, K={K}): "
+        f"maxerr={max_err}"
+    )
+
+
+# ===========================================================================
+# Test 6: activation kernel unit tests (GeGLU and clamped-SwiGLU vs torch)
+# ===========================================================================
+
+
+def _torch_geglu_fwd(gu):
+    g, u = gu.chunk(2, dim=-1)
+    return F.gelu(g.float(), approximate="tanh") * u.float()
+
+
+def _torch_geglu_bwd(gu, dh):
+    gu_f = gu.float().requires_grad_(True)
+    g, u = gu_f.chunk(2, dim=-1)
+    h = F.gelu(g, approximate="tanh") * u
+    h.backward(dh.float())
+    return gu_f.grad.to(gu.dtype)
+
+
+def _torch_swiglu_fwd(gu, limit):
+    g, u = gu.chunk(2, dim=-1)
+    return F.silu(g.float().clamp(max=limit)) * u.float().clamp(min=-limit, max=limit)
+
+
+def _torch_swiglu_bwd(gu, dh, limit):
+    gu_f = gu.float().requires_grad_(True)
+    g, u = gu_f.chunk(2, dim=-1)
+    h = F.silu(g.clamp(max=limit)) * u.clamp(min=-limit, max=limit)
+    h.backward(dh.float())
+    return gu_f.grad.to(gu.dtype)
+
+
+def _check_rel_err(name, got, ref, rtol=5e-3):
+    diff = (got.float() - ref.float()).abs()
+    rel_err = diff.mean() / (ref.float().abs().mean() + 1e-8)
+    assert rel_err < rtol, f"{name}: rel_err={rel_err:.5f} >= {rtol}"
+
+
+@pytest.mark.parametrize(
+    "I,seed", [(704, 42), (2048, 42)], ids=["I=704(gemma4)", "I=2048(dsv4)"]
+)
+def test_geglu_kernel_vs_torch(I, seed):
+    """GeGLU (gelu_tanh) Triton kernel fwd+bwd matches torch reference."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.cutlass_fp4.swiglu import (
+        swiglu_bwd,
+        swiglu_fwd,
+    )
+
+    torch.manual_seed(seed)
+    Mt = 512
+    gu = torch.randn(Mt, 2 * I, dtype=torch.bfloat16, device=DEV) * 0.5
+    dh = torch.randn(Mt, I, dtype=torch.bfloat16, device=DEV) * 0.1
+
+    ref_h = _torch_geglu_fwd(gu)
+    got_h = swiglu_fwd(gu, limit=1e30, act_type="gelu_tanh")
+    _check_rel_err("geglu_fwd", got_h.float(), ref_h)
+
+    ref_dgu = _torch_geglu_bwd(gu, dh)
+    got_dgu = swiglu_bwd(gu, dh, limit=1e30, act_type="gelu_tanh")
+    _check_rel_err("geglu_bwd_dg", got_dgu[:, :I].float(), ref_dgu[:, :I].float())
+    _check_rel_err("geglu_bwd_du", got_dgu[:, I:].float(), ref_dgu[:, I:].float())
+
+
+@pytest.mark.parametrize(
+    "limit,I,seed",
+    [(10.0, 704, 42), (7.0, 2048, 42)],
+    ids=["swiglu_clamp_I704", "swiglu_clamp_I2048"],
+)
+def test_swiglu_clamped_kernel_vs_torch(limit, I, seed):
+    """Clamped SwiGLU Triton kernel fwd+bwd matches torch reference."""
+    from axolotl.integrations.kernels.libs.scattermoe_lora.cutlass_fp4.swiglu import (
+        swiglu_bwd,
+        swiglu_fwd,
+    )
+
+    torch.manual_seed(seed)
+    Mt = 512
+    gu = torch.randn(Mt, 2 * I, dtype=torch.bfloat16, device=DEV) * 5.0
+    dh = torch.randn(Mt, I, dtype=torch.bfloat16, device=DEV) * 0.1
+
+    ref_h = _torch_swiglu_fwd(gu, limit)
+    got_h = swiglu_fwd(gu, limit=limit, act_type="silu")
+    _check_rel_err("swiglu_fwd", got_h.float(), ref_h)
+
+    ref_dgu = _torch_swiglu_bwd(gu, dh, limit)
+    got_dgu = swiglu_bwd(gu, dh, limit=limit, act_type="silu")
+    _check_rel_err("swiglu_bwd_dg", got_dgu[:, :I].float(), ref_dgu[:, :I].float())
+    _check_rel_err("swiglu_bwd_du", got_dgu[:, I:].float(), ref_dgu[:, I:].float())
diff --git a/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_weight_scale_2.py b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_weight_scale_2.py
new file mode 100644
index 0000000000..6153eec672
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_grouped_fp4_weight_scale_2.py
@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Non-unit per-expert ``weight_scale_2`` (per_tensor_scale) fwd+bwd parity (C1 / A5 / D-A1).
+
+The real DSV4 NVFP4 checkpoint carries a non-unit per-expert ``per_tensor_scale`` (weight_scale_2),
+a second scale applied on top of the E4M3 per-16-block scale. The CUTLASS forward cannot fold it
+(set_weights drops it while the backward applies it -> wrong forward + grad mismatch), so the
+backend selection routes a non-unit per_tensor_scale to Marlin or DeepGEMM (both fold it into the
+weight) and hard-errors if neither is available. This test pins that the Marlin/DeepGEMM forward +
+the bf16-dequant backward both honor a per-expert weight_scale_2, matching a bf16 oracle whose
+weights are dequantized WITH the same per-expert per_tensor_scale.
+
+Why this is NOT pure-dequant-fallback testable: the chunked-dequant fallback (``moe_grouped_backend
+= 'dequant'``) still uses the CUTLASS *forward* engine (set_weights), which drops weight_scale_2 -
+so a fused backend (Marlin sm80+ / DeepGEMM sm90+) that folds it is REQUIRED for the forward. On the
+sm89 L40S CI that is Marlin, so this test EXECUTES there. On a GPU with neither, it skips.
+
+Gate: CUDA + a weight_scale_2-folding forward backend (Marlin sm80+ / DeepGEMM sm90+). NOT
+Blackwell-gated (Marlin covers sm89), but it IS hardware-only - there is no CPU path.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from .test_grouped_fp4_train import (
+    DEV,
+    SHAPE_IDS,
+    SHAPES,
+    _bf16_oracle,
+    _cos,
+    _fp4_codebook,
+)
+
+_IS_CUDA = torch.cuda.is_available()
+
+
+def _ws2_backend_available() -> bool:
+    """True iff a forward backend that FOLDS weight_scale_2 (Marlin sm80+ / DeepGEMM sm90+) is
+    available. CUTLASS does NOT fold it, so its presence alone does not qualify."""
+    if not _IS_CUDA:
+        return False
+    try:
+        from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+            _backend_available,
+        )
+
+        return _backend_available("marlin") or _backend_available("deepgemm")
+    except Exception:
+        return False
+
+
+pytestmark = pytest.mark.skipif(
+    not _ws2_backend_available(),
+    reason="non-unit weight_scale_2 needs a folding forward backend (Marlin sm80+ / DeepGEMM "
+    "sm90+); CUTLASS drops it and there is no CPU path - hardware-only",
+)
+
+
+def _quantize_nvfp4_ws2(W: torch.Tensor, pt_e: torch.Tensor):
+    """Quantize bf16 W[E,N,K] -> NVFP4Tensor carrying a per-expert per_tensor_scale ``pt_e`` [E].
+
+    Factors ``pt_e`` out of the weight before the E4M3-block quantization and stores it as
+    per_tensor_scale, so ``nvfp4_dequant_bf16(qdata, scale, pt_e)`` reconstructs the original W:
+    quantize W/pt_e into the block path, then scale back by pt_e at dequant. This reproduces the
+    real DSV4 layout (a non-unit weight_scale_2) rather than the unit-scale synthetic in the
+    sibling correctness module.
+    """
+    from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+    dev = W.device
+    cb = _fp4_codebook(dev)
+    E, N, K = W.shape
+    Wf = (W.float() / pt_e.view(E, 1, 1)).reshape(E, N, K // 16, 16)
+    amax = Wf.abs().amax(-1).clamp_min(1e-6)
+    scale_e4m3 = (amax / 6.0).to(torch.float8_e4m3fn)
+    scale_f = scale_e4m3.float().clamp_min(1e-9)
+    nib = (
+        ((Wf / scale_f.unsqueeze(-1)).unsqueeze(-1) - cb.view(1, 1, 1, 1, 16))
+        .abs()
+        .argmin(-1)
+    )
+    nib = nib.to(torch.uint8).reshape(E, N, K)
+    packed = (nib[..., 0::2] & 0xF) | ((nib[..., 1::2] & 0xF) << 4)
+    return NVFP4Tensor(
+        packed.contiguous(),
+        scale_e4m3.contiguous(),
+        block_size=16,
+        orig_dtype=torch.bfloat16,
+        per_tensor_scale=pt_e.contiguous(),
+    )
+
+
+@pytest.mark.parametrize("cfg", SHAPES, ids=SHAPE_IDS)
+def test_grouped_fp4_ws2_fwd_bwd_matches_oracle(cfg):
+    """A non-unit per-expert weight_scale_2 forward (cosine >= 0.97) and backward (cosine >= 0.95)
+    match a bf16 oracle dequantized WITH the same per-expert per_tensor_scale.
+
+    Exercises the Marlin/DeepGEMM weight_scale_2 fold in the forward and the per-expert ``_pt``
+    application in the bf16-dequant backward (prefer_fp8_dx=False). Skips on a GPU that has neither
+    a Marlin nor a DeepGEMM backend (CUTLASS-only / no fused backend) since CUTLASS cannot fold it.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped import (
+        nvfp4_dequant_bf16,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.grouped_train import (
+        _has_nonunit_pt,
+        grouped_fp4_moe_train,
+    )
+
+    E, H, I, topk, N, r = cfg["E"], cfg["H"], cfg["I"], cfg["topk"], cfg["N"], cfg["r"]
+    s = cfg["scaling"]
+    twoI = 2 * I
+    torch.manual_seed(11)
+
+    # Non-unit per-expert weight_scale_2 (clearly != 1 to make the fold load-bearing).
+    pt_gu = (0.5 + torch.rand(E, device=DEV)).to(torch.float32)
+    pt_dn = (0.5 + torch.rand(E, device=DEV)).to(torch.float32)
+
+    Wgu = torch.randn(E, twoI, H, device=DEV, dtype=torch.bfloat16) * 0.04
+    Wdn = torch.randn(E, H, I, device=DEV, dtype=torch.bfloat16) * 0.04
+    gu_nv = _quantize_nvfp4_ws2(Wgu, pt_gu)
+    dn_nv = _quantize_nvfp4_ws2(Wdn, pt_dn)
+
+    # Sanity: the construction really carries a non-unit weight_scale_2 (else the test is vacuous).
+    assert _has_nonunit_pt(gu_nv, dn_nv), (
+        "weight_scale_2 construction produced a unit pt"
+    )
+
+    # Oracle weights: dequantize WITH the per-expert per_tensor_scale.
+    Wgu_b = nvfp4_dequant_bf16(gu_nv.qdata, gu_nv.scale, pt_gu)
+    Wdn_b = nvfp4_dequant_bf16(dn_nv.qdata, dn_nv.scale, pt_dn)
+
+    hidden = torch.randn(N, H, device=DEV, dtype=torch.bfloat16) * 0.5
+    idx = torch.stack([torch.randperm(E, device=DEV)[:topk] for _ in range(N)])
+    wts = torch.softmax(torch.randn(N, topk, device=DEV), -1).to(torch.bfloat16)
+
+    def mk(*shape):
+        return torch.randn(*shape, device=DEV, dtype=torch.bfloat16) * 0.02
+
+    Agu, Bgu, Adn, Bdn = mk(r * E, H), mk(twoI, r * E), mk(r * E, I), mk(H, r * E)
+
+    hk = hidden.detach().clone().requires_grad_(True)
+    Agk = Agu.clone().requires_grad_(True)
+    Bgk = Bgu.clone().requires_grad_(True)
+    Adk = Adn.clone().requires_grad_(True)
+    Bdk = Bdn.clone().requires_grad_(True)
+
+    out = grouped_fp4_moe_train(
+        hk,
+        idx,
+        wts,
+        gu_nv,
+        dn_nv,
+        (Agk, Bgk, s),
+        (Adk, Bdk, s),
+        cfg["limit"],
+        "nvfp4",
+        act_type=cfg["act_type"],
+        mxfp4_cache={},
+        prefer_fp8_dx=False,
+    )
+    assert torch.isfinite(out).all(), f"{cfg['name']}: ws2 fwd has NaN/inf"
+    out.float().pow(2).mean().backward()
+
+    hr = hidden.detach().clone().requires_grad_(True)
+    Agr = Agu.clone().requires_grad_(True)
+    Bgr = Bgu.clone().requires_grad_(True)
+    Adr = Adn.clone().requires_grad_(True)
+    Bdr = Bdn.clone().requires_grad_(True)
+    ref = _bf16_oracle(
+        hr, idx, wts, Wgu_b, Wdn_b, Agr, Bgr, Adr, Bdr, cfg["act_type"], cfg["limit"], s
+    )
+    cos_fwd = _cos(out, ref)
+    assert cos_fwd > 0.97, f"{cfg['name']}: ws2 fwd cosine {cos_fwd:.4f} < 0.97"
+
+    ref.pow(2).mean().backward()
+    for name, gk, gr in [
+        ("d_hidden", hk.grad, hr.grad),
+        ("d_Agu", Agk.grad, Agr.grad),
+        ("d_Adn", Adk.grad, Adr.grad),
+    ]:
+        assert gk is not None, f"{cfg['name']}: ws2 {name} grad is None"
+        assert torch.isfinite(gk).all(), f"{cfg['name']}: ws2 {name} has NaN/inf"
+        cos = _cos(gk, gr)
+        assert cos > 0.95, f"{cfg['name']}: ws2 {name} cosine {cos:.4f} < 0.95"
diff --git a/tests/integrations/kernels/scattermoe_lora/test_mx_lora_grouped_gram.py b/tests/integrations/kernels/scattermoe_lora/test_mx_lora_grouped_gram.py
new file mode 100644
index 0000000000..60add5b29f
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_mx_lora_grouped_gram.py
@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""MXFP4 base + LoRA at large workload hits the non-fused grouped-Gram dA/dB path.
+
+Guards the edge case: the LoRA weight gradients (dA/dB) must be base-agnostic --
+identical for an MXFP4 base and a bf16 dense base that is its dequantization --
+and dX must still route through the MX kernel (the non-fused dX path is never
+taken for MX). The workload is sized above the fuse-gather threshold so dA/dB go
+through the grouped-Gram path rather than the fused-gather kernel.
+"""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+MXTensor = pytest.importorskip(
+    "torchao.prototype.mx_formats.mx_tensor", reason="torchao required for MXFP4"
+).MXTensor
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (  # noqa: E402
+    selective_mx_weights_fwd,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (  # noqa: E402
+    flatten_sort_count,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_linear_lora import (  # noqa: E402
+    ScatterMoELoRA,
+)
+
+_FUSE_GATHER_THRESHOLD = 2**24
+
+
+def _run(W, x, A, B, sei, ssi, eo, topk, scaling, grad):
+    xg = x.detach().requires_grad_(True)
+    Ag = A.detach().requires_grad_(True)
+    Bg = B.detach().requires_grad_(True)
+    out = ScatterMoELoRA.apply(
+        xg, W, topk, sei, ssi, eo, Ag, Bg, scaling, None, None, False, False, True, True
+    )
+    out.backward(grad)
+    return out, xg.grad, Ag.grad, Bg.grad
+
+
+def test_mxfp4_lora_grads_base_agnostic_large_workload():
+    dev, dt = "cuda", torch.bfloat16
+    E, K, N, M, topk, R = 8, 2048, 2048, 8192, 2, 16  # M_total*max(K,N) > 2**24
+    scaling = 0.5
+    torch.manual_seed(0)
+
+    W_dense = torch.randn(E, N, K, device=dev, dtype=dt)
+    mx = MXTensor.to_mx(W_dense, elem_dtype=torch.float4_e2m1fn_x2, block_size=32)
+    W_ref = mx.dequantize(dt).contiguous()  # exactly what the MX kernel dequantizes to
+    mx_w = selective_mx_weights_fwd(mx, torch.arange(E, device=dev))
+    W_kernel = W_ref.transpose(2, 1).contiguous()  # [E, K, N] for the dense kernel
+
+    x = torch.randn(M, K, device=dev, dtype=dt)
+    A = torch.randn(R * E, K, device=dev, dtype=dt) * 0.02
+    B = torch.randn(N, R * E, device=dev, dtype=dt) * 0.02
+    _, top = torch.topk(torch.softmax(torch.randn(M, E, device=dev), -1), topk, dim=-1)
+    sei, ssi, eo = flatten_sort_count(top, E)
+
+    assert (
+        sei.size(0) * max(K, N) > _FUSE_GATHER_THRESHOLD
+    )  # exercises grouped-Gram dA/dB
+
+    grad = torch.randn(sei.size(0), N, device=dev, dtype=dt)
+    _, dx_mx, da_mx, db_mx = _run(mx_w, x, A, B, sei, ssi, eo, topk, scaling, grad)
+    _, _, da_bf, db_bf = _run(W_kernel, x, A, B, sei, ssi, eo, topk, scaling, grad)
+
+    # LoRA grads never touch the (frozen) base -> identical for MX vs dense base
+    assert torch.equal(da_mx, da_bf), (da_mx - da_bf).abs().max().item()
+    assert torch.equal(db_mx, db_bf), (db_mx - db_bf).abs().max().item()
+    # and everything is finite + actually trained
+    for g in (dx_mx, da_mx, db_mx):
+        assert torch.isfinite(g).all() and g.abs().sum() > 0
diff --git a/tests/integrations/kernels/scattermoe_lora/test_mxfp4_cache_fsdp.py b/tests/integrations/kernels/scattermoe_lora/test_mxfp4_cache_fsdp.py
new file mode 100644
index 0000000000..791b0d7469
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_mxfp4_cache_fsdp.py
@@ -0,0 +1,83 @@
+"""CPU regression tests for the DeepGEMM mxfp4-requant cache under FSDP.
+
+The bug (DSV4 NVFP4 LoRA OOM on 2xB200): the nvfp4->mxfp4 requant cached a full mxfp4 copy per
+layer — in a module-level dict AND a per-tensor ``_dg_mxfp4`` attribute. Under FSDP2 the gathered
+weight is full+fresh each step and its tensor object outlives the reshard, so both caches
+accumulated a full-model mxfp4 copy on every rank (~230GB across 61 layers) on top of the sharded
+base -> OOM. The base itself always sharded correctly.
+
+Fix: ``RUNTIME.mxfp4_cache_persist`` is disabled under FSDP (``configure_scattermoe_runtime``), and
+``_cached_mxfp4`` then recomputes-and-returns WITHOUT touching either cache, bounding resident mxfp4
+to one layer. Single-GPU (persistent param) keeps the cache for speed.
+
+These tests stub the GPU-only ``nvfp4_to_mxfp4_weight`` so they run on CI CPU.
+"""
+
+import axolotl.integrations.kernels.libs.scattermoe_lora.dequant_grouped as dg
+from axolotl.integrations.kernels.libs.scattermoe_lora.runtime import (
+    RUNTIME,
+    configure_scattermoe_runtime,
+)
+from axolotl.utils.dict import DictDefault
+
+
+class _FakeWeight:
+    """Stand-in for a gathered NVFP4Tensor expert weight (settable ``_dg_mxfp4`` attribute)."""
+
+    def __init__(self, tag):
+        self.qdata = (
+            tag  # passed straight to the stubbed requant so we can identify the call
+        )
+        self.scale = None
+
+
+def _stub_requant(monkeypatch):
+    calls = {"n": 0}
+
+    def fake(qdata, scale, per_tensor):
+        calls["n"] += 1
+        return ("mxfp4", qdata, calls["n"])  # unique per call
+
+    monkeypatch.setattr(dg, "nvfp4_to_mxfp4_weight", fake)
+    return calls
+
+
+def test_configure_runtime_disables_cache_under_fsdp():
+    configure_scattermoe_runtime(DictDefault({"fsdp_config": {"fsdp_version": 2}}))
+    assert RUNTIME.mxfp4_cache_persist is False
+    configure_scattermoe_runtime(DictDefault({}))  # no FSDP -> keep the cache
+    assert RUNTIME.mxfp4_cache_persist is True
+
+
+def test_no_persist_never_caches(monkeypatch):
+    """The core regression: under FSDP, neither the module dict nor the per-tensor attr is written,
+    and every call recomputes (so nothing accumulates across layers)."""
+    calls = _stub_requant(monkeypatch)
+    monkeypatch.setattr(RUNTIME, "mxfp4_cache_persist", False)
+    cache: dict = {}
+    w = _FakeWeight("gate_up")
+
+    r1 = dg._cached_mxfp4(w, None, cache, "gate_up")
+    r2 = dg._cached_mxfp4(w, None, cache, "gate_up")
+
+    assert calls["n"] == 2  # recomputed each time, no reuse
+    assert r1 != r2  # distinct (uncached) results
+    assert cache == {}  # module-level dict never populated
+    assert not hasattr(
+        w, "_dg_mxfp4"
+    )  # per-tensor attr never set -> can't survive reshard
+
+
+def test_persist_caches_and_reuses(monkeypatch):
+    """Single-GPU (persistent param): the module cache is used so the requant runs once per key."""
+    calls = _stub_requant(monkeypatch)
+    monkeypatch.setattr(RUNTIME, "mxfp4_cache_persist", True)
+    cache: dict = {}
+    w = _FakeWeight("gate_up")
+
+    r1 = dg._cached_mxfp4(w, None, cache, "gate_up")
+    r2 = dg._cached_mxfp4(w, None, cache, "gate_up")
+
+    assert calls["n"] == 1  # computed once, then served from cache
+    assert r1 == r2
+    assert "gate_up" in cache
diff --git a/tests/integrations/kernels/scattermoe_lora/test_mxfp4_expert_weights.py b/tests/integrations/kernels/scattermoe_lora/test_mxfp4_expert_weights.py
new file mode 100644
index 0000000000..0f147beef3
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_mxfp4_expert_weights.py
@@ -0,0 +1,567 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Correctness tests for MXFP4 expert weight support in ScatterMoE LoRA.
+
+Validates both strategies against a bf16-dequantized reference:
+
+  Strategy A — selective dequant:
+    The kernel runs on the dequantized [num_active, K, N] bf16 buffer,
+    so outputs must be bitwise identical to the baseline that supplies
+    the same bf16 weights directly.
+
+  Strategy B — fused Triton (when enabled):
+    The kernel unpacks MXFP4 + applies E8M0 scales in its K-loop. Output
+    differs from the bf16 reference by MX-rounding-tolerance only.
+
+Shapes covered:
+  - small:           [E=8, K=128, N=256], M=16, top_k=2, rank=8
+  - representative:  [E=32, K=2048, N=1024], M=64, top_k=4, rank=16
+"""
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+    selective_mx_weights_fwd,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+    flatten_sort_count,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_linear_lora import (
+    parallel_linear_lora,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (
+    get_active_experts,
+    is_mxfp4_param,
+    remap_expert_indices,
+    selective_expert_weights,
+    selective_lora_weights,
+)
+
+torchao = pytest.importorskip("torchao")
+from torchao.prototype.mx_formats.mx_tensor import MXTensor  # noqa: E402
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA required for MX kernels"
+)
+
+
+SHAPES = [
+    # (E, K, N, M, top_k, R, seed)
+    pytest.param(8, 128, 256, 16, 2, 8, 0, id="small"),
+    pytest.param(32, 2048, 1024, 64, 4, 16, 1, id="representative"),
+]
+
+# Per-shape Strategy-B tolerances. Forward outputs accumulate K dot-products
+# in fp32 then cast to bf16, so they stay within a few ULPs of the bf16
+# baseline. The dX path reduces over N (which is typically larger than K and
+# uses a different MMA tile layout than the bf16 reference), so we apply a
+# looser ULP-aware tolerance there. These are still tight compared to
+# torchao's own bf16 vs fp32 GEMM noise.
+_STRATEGY_B_FWD_TOL = {
+    "small": dict(atol=2e-3, rtol=2e-3),
+    "representative": dict(atol=1e-2, rtol=5e-3),
+}
+# dX tolerance: ~1 bf16 ULP at the typical output magnitude (rtol dominates;
+# atol caps near-zero entries where MMA-reordering manifests as full ULP).
+_STRATEGY_B_DX_TOL = {
+    "small": dict(atol=0.5, rtol=2e-2),
+    "representative": dict(atol=2.0, rtol=3e-2),
+}
+# dA / dB tolerance: the fused dA/dB kernel accumulates via atomic_add from
+# multiple N-block programs per expert, and the number of in-flight programs
+# differs between the full-E baseline and the compact-active MX path —
+# atomic ordering then introduces bf16 ULP-scale noise. Looser than the
+# forward bound because the gradients integrate over both M and N.
+_STRATEGY_B_LORA_GRAD_TOL = {
+    "small": dict(atol=2e-2, rtol=2e-2),
+    "representative": dict(atol=2e-1, rtol=3e-2),
+}
+
+
+def _tol_for_shape(K, *, dx: bool = False, lora_grad: bool = False):
+    if lora_grad:
+        table = _STRATEGY_B_LORA_GRAD_TOL
+    elif dx:
+        table = _STRATEGY_B_DX_TOL
+    else:
+        table = _STRATEGY_B_FWD_TOL
+    return table["small"] if K <= 128 else table["representative"]
+
+
+def _make_mxfp4_weights(E, K, N, seed):
+    """Build a `[E, N, K]` MXFP4 ``MXTensor`` (block axis = K, the contraction
+    axis). Returns (mx, W_ref_bf16) — the bf16 reference is the dequantization
+    of the full MX tensor so Strategy A can hit bitwise equality."""
+    torch.manual_seed(seed)
+    # Natural axolotl storage is [E, N, K] where K is the contraction axis;
+    # `experts.gate_up_proj.transpose(2, 1)` then yields [E, K, N] for the kernel.
+    W_dense = torch.randn(E, N, K, device=DEVICE, dtype=DTYPE)
+    mx = MXTensor.to_mx(W_dense, elem_dtype=torch.float4_e2m1fn_x2, block_size=32)
+    W_ref = mx.dequantize(DTYPE).contiguous()
+    return mx, W_ref
+
+
+def _setup_routing_and_lora(E, K, N, M, top_k, R, seed):
+    torch.manual_seed(seed + 100)
+    x = torch.randn(M, K, device=DEVICE, dtype=DTYPE)
+    lora_A = torch.randn(R * E, K, device=DEVICE, dtype=DTYPE) * 0.01
+    lora_B = torch.randn(N, R * E, device=DEVICE, dtype=DTYPE) * 0.01
+    logits = torch.randn(M, E, device=DEVICE)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
+    sei, ssi, eo = flatten_sort_count(top_idx, E)
+    return x, lora_A, lora_B, sei, ssi, eo
+
+
+class _MockExperts:
+    """Bare object exposing ``gate_up_proj`` so `selective_expert_weights`
+    can branch on it."""
+
+    def __init__(self, mx_param, num_experts):
+        self.gate_up_proj = mx_param
+        self.num_experts = num_experts
+
+
+def _run_baseline(
+    W_ref,
+    x,
+    lora_A,
+    lora_B,
+    scaling,
+    sei,
+    ssi,
+    eo,
+    top_k,
+    *,
+    use_fused_dX: bool = False,
+    use_fused_gather: bool = False,
+):
+    """Full-E bf16 baseline: dense weights, full LoRA, full expert indices."""
+    W_kernel = W_ref.transpose(2, 1).contiguous()  # [E, K, N]
+    return parallel_linear_lora(
+        x,
+        W_kernel,
+        top_k,
+        sei,
+        ssi,
+        eo,
+        lora_A=lora_A,
+        lora_B=lora_B,
+        scaling=scaling,
+        use_fused_dX=use_fused_dX,
+        use_fused_gather=use_fused_gather,
+    )
+
+
+def _run_strategy_a(mx, x, lora_A, lora_B, scaling, sei, ssi, eo, top_k, E):
+    """Strategy A: selective dequant via MXTensor branch in
+    `selective_expert_weights`. Compact weights + remapped indices."""
+    experts = _MockExperts(mx, E)
+    active = get_active_experts(sei, E)
+    remapped, compact_offsets = remap_expert_indices(sei, eo, active, E)
+    W_compact = (
+        selective_expert_weights(experts, "gate_up_proj", active)
+        .transpose(2, 1)
+        .contiguous()
+    )  # [num_active, K, N]
+    A_compact, B_compact = selective_lora_weights(lora_A, lora_B, active, E)
+    return parallel_linear_lora(
+        x,
+        W_compact,
+        top_k,
+        remapped,
+        ssi,
+        compact_offsets,
+        lora_A=A_compact,
+        lora_B=B_compact,
+        scaling=scaling,
+    ), active
+
+
+# ─── Strategy A — bitwise identity vs bf16 baseline ───────────────────────────
+
+
+@pytest.mark.parametrize("E,K,N,M,top_k,R,seed", SHAPES)
+def test_strategy_a_forward_matches_bf16(E, K, N, M, top_k, R, seed):
+    mx, W_ref = _make_mxfp4_weights(E, K, N, seed)
+    assert is_mxfp4_param(mx)
+    x, lora_A, lora_B, sei, ssi, eo = _setup_routing_and_lora(
+        E, K, N, M, top_k, R, seed
+    )
+    scaling = 0.5
+
+    out_baseline = _run_baseline(W_ref, x, lora_A, lora_B, scaling, sei, ssi, eo, top_k)
+    out_a, _ = _run_strategy_a(mx, x, lora_A, lora_B, scaling, sei, ssi, eo, top_k, E)
+
+    assert out_baseline.shape == out_a.shape
+    assert torch.equal(out_baseline, out_a), (
+        f"Strategy A forward must match bf16 baseline bitwise. "
+        f"max abs diff = {(out_baseline - out_a).abs().max().item()}"
+    )
+
+
+@pytest.mark.parametrize("E,K,N,M,top_k,R,seed", SHAPES)
+def test_strategy_a_backward_matches_bf16(E, K, N, M, top_k, R, seed):
+    """Forward + backward parity. dX must be bitwise identical; the LoRA
+    grads dA/dB are compared on the active expert slices only (the full
+    LoRA tensors differ in shape between baseline and compact paths)."""
+    mx, W_ref = _make_mxfp4_weights(E, K, N, seed)
+    x_base, lora_A_base, lora_B_base, sei, ssi, eo = _setup_routing_and_lora(
+        E, K, N, M, top_k, R, seed
+    )
+    scaling = 0.5
+
+    # Baseline backward
+    x_b = x_base.detach().clone().requires_grad_(True)
+    A_b = lora_A_base.detach().clone().requires_grad_(True)
+    B_b = lora_B_base.detach().clone().requires_grad_(True)
+    out_b = _run_baseline(W_ref, x_b, A_b, B_b, scaling, sei, ssi, eo, top_k)
+    grad_out = torch.randn_like(out_b)
+    out_b.backward(grad_out)
+
+    # Strategy A backward
+    x_a = x_base.detach().clone().requires_grad_(True)
+    experts = _MockExperts(mx, E)
+    active = get_active_experts(sei, E)
+    remapped, compact_offsets = remap_expert_indices(sei, eo, active, E)
+    W_compact = (
+        selective_expert_weights(experts, "gate_up_proj", active)
+        .transpose(2, 1)
+        .contiguous()
+    )
+    A_full = lora_A_base.detach().clone().requires_grad_(True)
+    B_full = lora_B_base.detach().clone().requires_grad_(True)
+    A_compact, B_compact = selective_lora_weights(A_full, B_full, active, E)
+    out_a = parallel_linear_lora(
+        x_a,
+        W_compact,
+        top_k,
+        remapped,
+        ssi,
+        compact_offsets,
+        lora_A=A_compact,
+        lora_B=B_compact,
+        scaling=scaling,
+    )
+    out_a.backward(grad_out)
+
+    # Forward parity (Strategy A contract: bitwise identical to bf16 baseline).
+    # Asserted here so a forward bug that produces a constant offset (and
+    # therefore zero gradient delta) doesn't slip past the bwd-only checks.
+    assert torch.equal(out_b, out_a), (
+        f"forward mismatch (Strategy A): max abs diff = "
+        f"{(out_b - out_a).abs().max().item()}"
+    )
+
+    # dX: bitwise identical
+    assert torch.equal(x_b.grad, x_a.grad), (
+        f"dX mismatch (Strategy A): max abs diff = "
+        f"{(x_b.grad - x_a.grad).abs().max().item()}"
+    )
+
+    # dA / dB — gather active slices from the baseline full grads and compare
+    row_idx = (
+        active.long()[:, None] * R + torch.arange(R, device=DEVICE)[None, :]
+    ).reshape(-1)
+    dA_b_active = A_b.grad[row_idx]
+    dB_b_active = B_b.grad[:, row_idx]
+
+    # A_compact is a view (advanced indexing produces a copy, so the grad lands
+    # on the full lora_A via the slice). torch.autograd flows back through
+    # selective_lora_weights, so A_full.grad has gradient on rows for active
+    # experts only.
+    dA_a_active = A_full.grad[row_idx]
+    dB_a_active = B_full.grad[:, row_idx]
+
+    assert torch.equal(dA_b_active, dA_a_active), (
+        f"dA active slice mismatch: max diff = "
+        f"{(dA_b_active - dA_a_active).abs().max().item()}"
+    )
+    assert torch.equal(dB_b_active, dB_a_active), (
+        f"dB active slice mismatch: max diff = "
+        f"{(dB_b_active - dB_a_active).abs().max().item()}"
+    )
+
+
+# ─── Strategy A — backward through fused dX/gather paths ──────────────────────
+
+
+@pytest.mark.parametrize("use_fused_dX", [False, True])
+@pytest.mark.parametrize("use_fused_gather", [False, True])
+def test_strategy_a_backward_fused_variants(use_fused_dX, use_fused_gather):
+    """Strategy A must match baseline across all four fused-bwd flag
+    combinations exercised in production by ``HFScatterMoEGatedMLP``.
+    Asserts parity on dX, dA, and dB (active expert slice for dA/dB)."""
+    E, K, N, M, top_k, R = 8, 128, 256, 16, 2, 8
+    mx, W_ref = _make_mxfp4_weights(E, K, N, seed=7)
+    x_base, lora_A_base, lora_B_base, sei, ssi, eo = _setup_routing_and_lora(
+        E, K, N, M, top_k, R, seed=7
+    )
+    scaling = 0.25
+
+    def run(W_kernel, sei_, eo_, k_, lora_A_in, lora_B_in, grad_out):
+        x_g = x_base.detach().clone().requires_grad_(True)
+        A_g = lora_A_in.detach().clone().requires_grad_(True)
+        B_g = lora_B_in.detach().clone().requires_grad_(True)
+        out = parallel_linear_lora(
+            x_g,
+            W_kernel,
+            k_,
+            sei_,
+            ssi,
+            eo_,
+            lora_A=A_g,
+            lora_B=B_g,
+            scaling=scaling,
+            use_fused_dX=use_fused_dX,
+            use_fused_gather=use_fused_gather,
+        )
+        out.backward(grad_out)
+        return out, x_g.grad, A_g.grad, B_g.grad
+
+    # Non-trivial grad (a constant grad zeros out cross-token differences in
+    # the fused-gather accumulation, which can mask reordering bugs).
+    torch.manual_seed(7)
+    W_baseline = W_ref.transpose(2, 1).contiguous()
+    # Forward once on baseline shape to size grad_out.
+    out_shape_probe = parallel_linear_lora(
+        x_base,
+        W_baseline,
+        top_k,
+        sei,
+        ssi,
+        eo,
+        lora_A=lora_A_base,
+        lora_B=lora_B_base,
+        scaling=scaling,
+    )
+    grad_out = torch.randn_like(out_shape_probe) * 0.1
+
+    out_b, dx_b, dA_b, dB_b = run(
+        W_baseline,
+        sei,
+        eo,
+        top_k,
+        lora_A_base,
+        lora_B_base,
+        grad_out,
+    )
+
+    experts = _MockExperts(mx, E)
+    active = get_active_experts(sei, E)
+    remapped, compact_offsets = remap_expert_indices(sei, eo, active, E)
+    W_compact = (
+        selective_expert_weights(experts, "gate_up_proj", active)
+        .transpose(2, 1)
+        .contiguous()
+    )
+    A_compact, B_compact = selective_lora_weights(lora_A_base, lora_B_base, active, E)
+    out_a, dx_a, dA_a, dB_a = run(
+        W_compact,
+        remapped,
+        compact_offsets,
+        top_k,
+        A_compact,
+        B_compact,
+        grad_out,
+    )
+
+    # Strategy A is bitwise on forward / dX (same bf16 weights, same kernel).
+    assert torch.equal(out_b, out_a), (
+        f"forward mismatch: max diff = {(out_b - out_a).abs().max().item()}"
+    )
+    assert torch.equal(dx_b, dx_a), (
+        f"dX mismatch: max diff = {(dx_b - dx_a).abs().max().item()}"
+    )
+
+    # dA / dB — compare active-expert slice of the dense baseline grads
+    # against the compact-path grads. Same row_idx pattern as
+    # test_strategy_a_backward_matches_bf16. Bitwise (``torch.equal``) holds
+    # for forward and dX, but the fused dA/dB kernel uses ``atomic_add``
+    # across N-block programs and the in-flight program count differs
+    # between the full-E baseline and the compact-active path; combined
+    # with FMA reordering, this introduces 1–2 bf16 ULPs (~2e-4 at the
+    # values seen here) on the ``use_fused_dX=True`` configs. Use a
+    # tolerance an order of magnitude below that — tight enough to catch
+    # any real bug but tolerant of the unavoidable atomic-order noise.
+    lora_grad_tol = dict(atol=1e-3, rtol=1e-3)
+    row_idx = (
+        active.long()[:, None] * R + torch.arange(R, device=DEVICE)[None, :]
+    ).reshape(-1)
+    dA_b_active = dA_b[row_idx]
+    dB_b_active = dB_b[:, row_idx]
+
+    assert torch.allclose(dA_b_active, dA_a, **lora_grad_tol), (
+        f"dA active slice mismatch: max diff = "
+        f"{(dA_b_active - dA_a).abs().max().item()}"
+    )
+    assert torch.allclose(dB_b_active, dB_a, **lora_grad_tol), (
+        f"dB active slice mismatch: max diff = "
+        f"{(dB_b_active - dB_a).abs().max().item()}"
+    )
+
+
+# ─── Strategy B — fused MXFP4 Triton kernel ──────────────────────────────────
+
+
+# MX rounding tolerance — the Triton path can reorder FMAs vs the torchao
+# dequant + bf16 matmul reference, and the dequant arithmetic is
+# fp32-codebook * fp32-scale -> bf16. See ``_STRATEGY_B_TOL`` above.
+
+
+def _run_strategy_b(mx, x, lora_A, lora_B, scaling, sei, ssi, eo, top_k, E):
+    """Strategy B: pass MXWeights container directly to parallel_linear_lora;
+    the fused MX kernel does dequant inside the K-loop."""
+    active = get_active_experts(sei, E)
+    remapped, compact_offsets = remap_expert_indices(sei, eo, active, E)
+    mx_active = selective_mx_weights_fwd(mx, active)
+    A_compact, B_compact = selective_lora_weights(lora_A, lora_B, active, E)
+    return parallel_linear_lora(
+        x,
+        mx_active,
+        top_k,
+        remapped,
+        ssi,
+        compact_offsets,
+        lora_A=A_compact,
+        lora_B=B_compact,
+        scaling=scaling,
+    ), active
+
+
+@pytest.mark.parametrize("E,K,N,M,top_k,R,seed", SHAPES)
+def test_strategy_b_forward_matches_bf16(E, K, N, M, top_k, R, seed):
+    """Strategy B forward must match bf16 baseline within MX rounding tol."""
+    mx, W_ref = _make_mxfp4_weights(E, K, N, seed)
+    x, lora_A, lora_B, sei, ssi, eo = _setup_routing_and_lora(
+        E, K, N, M, top_k, R, seed
+    )
+    scaling = 0.5
+    tol = _tol_for_shape(K)
+
+    out_baseline = _run_baseline(W_ref, x, lora_A, lora_B, scaling, sei, ssi, eo, top_k)
+    out_b, _ = _run_strategy_b(mx, x, lora_A, lora_B, scaling, sei, ssi, eo, top_k, E)
+
+    assert out_baseline.shape == out_b.shape
+    diff = (out_baseline.float() - out_b.float()).abs()
+    rel = diff / (out_baseline.float().abs() + 1e-6)
+    assert torch.allclose(out_baseline, out_b, **tol), (
+        f"Strategy B forward exceeds MX tolerance: max abs={diff.max().item():.4e}, "
+        f"max rel={rel.max().item():.4e}"
+    )
+
+
+@pytest.mark.parametrize("E,K,N,M,top_k,R,seed", SHAPES)
+def test_strategy_b_backward_matches_bf16(E, K, N, M, top_k, R, seed):
+    """Strategy B forward+backward; dX, dA, dB compared to bf16 baseline on
+    the active expert slice within MX rounding tol."""
+    mx, W_ref = _make_mxfp4_weights(E, K, N, seed)
+    x_base, lora_A_base, lora_B_base, sei, ssi, eo = _setup_routing_and_lora(
+        E, K, N, M, top_k, R, seed
+    )
+    scaling = 0.5
+    fwd_tol = _tol_for_shape(K)
+    dx_tol = _tol_for_shape(K, dx=True)
+    lg_tol = _tol_for_shape(K, lora_grad=True)
+
+    # Baseline — match the MX path's fused-bwd kernel selection so dA/dB MMA
+    # accumulation order is the same and bf16 noise stays at single ULPs.
+    x_b = x_base.detach().clone().requires_grad_(True)
+    A_b = lora_A_base.detach().clone().requires_grad_(True)
+    B_b = lora_B_base.detach().clone().requires_grad_(True)
+    out_b = _run_baseline(
+        W_ref,
+        x_b,
+        A_b,
+        B_b,
+        scaling,
+        sei,
+        ssi,
+        eo,
+        top_k,
+        use_fused_dX=True,
+        use_fused_gather=True,
+    )
+    grad_out = torch.randn_like(out_b)
+    out_b.backward(grad_out)
+
+    # Strategy B
+    x_s = x_base.detach().clone().requires_grad_(True)
+    active = get_active_experts(sei, E)
+    remapped, compact_offsets = remap_expert_indices(sei, eo, active, E)
+    mx_active = selective_mx_weights_fwd(mx, active)
+    A_full = lora_A_base.detach().clone().requires_grad_(True)
+    B_full = lora_B_base.detach().clone().requires_grad_(True)
+    A_compact, B_compact = selective_lora_weights(A_full, B_full, active, E)
+    out_s = parallel_linear_lora(
+        x_s,
+        mx_active,
+        top_k,
+        remapped,
+        ssi,
+        compact_offsets,
+        lora_A=A_compact,
+        lora_B=B_compact,
+        scaling=scaling,
+    )
+    out_s.backward(grad_out)
+
+    # Forward parity within MX rounding tol — asserted here so a forward bug
+    # that produces a constant offset (and therefore zero gradient delta)
+    # doesn't slip past the bwd-only checks.
+    assert torch.allclose(out_b, out_s, **fwd_tol), (
+        f"Strategy B forward mismatch: max abs diff = "
+        f"{(out_b - out_s).abs().max().item():.4e}"
+    )
+
+    # dX tolerance (looser; see _STRATEGY_B_DX_TOL comment)
+    assert torch.allclose(x_b.grad, x_s.grad, **dx_tol), (
+        f"Strategy B dX mismatch: max diff = "
+        f"{(x_b.grad - x_s.grad).abs().max().item():.4e}"
+    )
+
+    # Uniform-scaling drift guard: the allclose bound above is generous to
+    # accommodate accumulated bf16 MMA noise over N. A bug that scales every
+    # dX element by a constant factor (e.g. an off-by-one on the E8M0
+    # exponent shifting the whole tile by 2x) would still pass that bound.
+    # Catch it by requiring the per-element ratio std stays small after
+    # masking out near-zero baseline elements (where the ratio is dominated
+    # by quantization noise rather than uniform drift).
+    bf16_dX = x_b.grad.float()
+    mx_dX = x_s.grad.float()
+    eps = 1e-3 * bf16_dX.abs().max().clamp(min=1e-6)
+    mask = bf16_dX.abs() > eps
+    if mask.any():
+        ratio = mx_dX[mask] / bf16_dX[mask]
+        ratio_std = ratio.std().item()
+        assert ratio_std < 0.5, (
+            f"Strategy B dX uniform-scaling drift: std(mx/bf16) = "
+            f"{ratio_std:.4f} (mean = {ratio.mean().item():.4f}); "
+            f"a uniform multiplicative bug would slip past the allclose "
+            f"bound but is caught here."
+        )
+
+    # dA / dB — compare active expert slices (use forward tolerance — these
+    # come from the LoRA-only grad path which doesn't touch the W matmul)
+    row_idx = (
+        active.long()[:, None] * R + torch.arange(R, device=DEVICE)[None, :]
+    ).reshape(-1)
+    dA_b_active = A_b.grad[row_idx]
+    dA_s_active = A_full.grad[row_idx]
+    dB_b_active = B_b.grad[:, row_idx]
+    dB_s_active = B_full.grad[:, row_idx]
+
+    assert torch.allclose(dA_b_active, dA_s_active, **lg_tol), (
+        f"Strategy B dA active slice mismatch: max diff = "
+        f"{(dA_b_active - dA_s_active).abs().max().item():.4e}"
+    )
+    assert torch.allclose(dB_b_active, dB_s_active, **lg_tol), (
+        f"Strategy B dB active slice mismatch: max diff = "
+        f"{(dB_b_active - dB_s_active).abs().max().item():.4e}"
+    )
diff --git a/tests/integrations/kernels/scattermoe_lora/test_mxfp4_experts_forward.py b/tests/integrations/kernels/scattermoe_lora/test_mxfp4_experts_forward.py
new file mode 100644
index 0000000000..b8aadc3ad0
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_mxfp4_experts_forward.py
@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""scattermoe_experts_forward routes base MXFP4 + LoRA through the fused MX kernel
+(no bf16 dequant of the frozen base) and is bit-for-bit unaffected for non-MXFP4 models.
+
+The fused path keeps the packed 4-bit weights and dequantizes inside the kernel K-loop.
+Validated vs a reference = the same forward on the dequantized weights, within MX
+rounding tolerance, on output / dX / the active-expert slice of LoRA dA, dB.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+MXTensor = pytest.importorskip(
+    "torchao.prototype.mx_formats.mx_tensor", reason="torchao required"
+).MXTensor
+
+import axolotl.integrations.kernels.libs.scattermoe_lora.experts as ex  # noqa: E402
+from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (  # noqa: E402
+    scattermoe_experts_forward,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (  # noqa: E402
+    flatten_sort_count,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (  # noqa: E402
+    get_active_experts,
+)
+
+DEV = "cuda"
+E, H, IM, N, K, R, SC = 8, 1024, 512, 512, 4, 16, 0.5
+
+
+def _module(gu, dn):
+    return SimpleNamespace(
+        num_experts=E,
+        gate_up_proj=gu,
+        down_proj=dn,
+        act_fn=torch.nn.functional.silu,
+        is_transposed=False,
+        is_concatenated=True,
+        has_bias=False,
+        has_gate=True,
+    )
+
+
+def test_mxfp4_forward_uses_fused_kernel_matches_dequant(monkeypatch):
+    dt = torch.bfloat16
+    g = torch.Generator(device=DEV).manual_seed(0)
+
+    def mx(W):
+        return MXTensor.to_mx(W, elem_dtype=torch.float4_e2m1fn_x2, block_size=32)
+
+    gu_mx = mx(torch.randn(E, 2 * IM, H, device=DEV, dtype=dt, generator=g) * 0.1)
+    dn_mx = mx(torch.randn(E, H, IM, device=DEV, dtype=dt, generator=g) * 0.1)
+    gu_deq, dn_deq = (
+        gu_mx.dequantize(dt).contiguous(),
+        dn_mx.dequantize(dt).contiguous(),
+    )
+
+    def mk(*s):
+        return (
+            torch.randn(*s, device=DEV, dtype=dt, generator=g) * 0.05
+        ).requires_grad_(True)
+
+    A1, B1, A2, B2 = mk(R * E, H), mk(2 * IM, R * E), mk(R * E, IM), mk(H, R * E)
+    lora = [A1, B1, A2, B2]
+    monkeypatch.setattr(ex, "_has_peft_wrapper", lambda s: True)
+    monkeypatch.setattr(
+        ex, "_unwrap_experts_lora", lambda s: (s, (A1, B1, SC), (A2, B2, SC))
+    )
+
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    w = torch.rand(N, K, device=DEV, generator=g).to(dt)
+    sei, _, _ = flatten_sort_count(idx, E)
+    active = get_active_experts(sei, E)
+    row = (active.long()[:, None] * R + torch.arange(R, device=DEV)[None, :]).reshape(
+        -1
+    )
+    grad = torch.randn(N, H, device=DEV, dtype=dt, generator=g)
+
+    def run(m):
+        for t in lora:
+            t.grad = None
+        x = torch.randn(
+            N, H, device=DEV, dtype=dt, generator=torch.Generator(DEV).manual_seed(7)
+        ).requires_grad_(True)
+        scattermoe_experts_forward(m, x, idx, w).backward(grad)
+        return x.grad.detach(), [t.grad.clone() for t in lora]
+
+    dx_mx, gl_mx = run(_module(gu_mx, dn_mx))  # fused MX path
+    dx_bf, gl_bf = run(_module(gu_deq, dn_deq))  # bf16 dequant reference
+
+    def rel(a, b):
+        return (a - b).float().abs().max().item() / max(
+            b.float().abs().max().item(), 1e-6
+        )
+
+    assert rel(dx_mx, dx_bf) < 6e-2
+    for i, slc in (
+        (0, row),
+        (1, (slice(None), row)),
+        (2, row),
+        (3, (slice(None), row)),
+    ):
+        assert rel(gl_mx[i][slc], gl_bf[i][slc]) < 6e-2, f"lora grad {i}"
+        assert torch.isfinite(gl_mx[i]).all()
+
+
+def test_mxfp4_ep_matches_dequant(monkeypatch):
+    """Fused MXFP4 through the EP (sentinel-skip) forward matches the bf16-dequant ref,
+    and _sonicmoe_local routes there on a device the sonic-moe kernel can't run."""
+    from axolotl.integrations.expert_parallel.experts_fn import _sonicmoe_local
+    from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+        scattermoe_experts_forward_ep,
+    )
+    from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+        _sonicmoe_kernel_supported,
+    )
+
+    dt = torch.bfloat16
+    g = torch.Generator(device=DEV).manual_seed(0)
+
+    def mx(W):
+        return MXTensor.to_mx(W, elem_dtype=torch.float4_e2m1fn_x2, block_size=32)
+
+    gu_mx = mx(torch.randn(E, 2 * IM, H, device=DEV, dtype=dt, generator=g) * 0.1)
+    dn_mx = mx(torch.randn(E, H, IM, device=DEV, dtype=dt, generator=g) * 0.1)
+    gu_deq, dn_deq = (
+        gu_mx.dequantize(dt).contiguous(),
+        dn_mx.dequantize(dt).contiguous(),
+    )
+
+    def mk(*s):
+        return (
+            torch.randn(*s, device=DEV, dtype=dt, generator=g) * 0.05
+        ).requires_grad_(True)
+
+    A1, B1, A2, B2 = mk(R * E, H), mk(2 * IM, R * E), mk(R * E, IM), mk(H, R * E)
+    lora = [A1, B1, A2, B2]
+    monkeypatch.setattr(ex, "_has_peft_wrapper", lambda s: True)
+    monkeypatch.setattr(
+        ex, "_unwrap_experts_lora", lambda s: (s, (A1, B1, SC), (A2, B2, SC))
+    )
+
+    # routing with -1 remote sentinels (post-DeepEP-dispatch shape)
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    remote = torch.rand(N, K, device=DEV, generator=g) < 0.5
+    remote[:, 0] = False
+    idx = torch.where(remote, torch.full_like(idx, -1), idx)
+    w = torch.rand(N, K, device=DEV, generator=g).to(dt)
+    active = torch.unique(torch.sort(idx.reshape(-1)[idx.reshape(-1) >= 0]).values)
+    row = (active.long()[:, None] * R + torch.arange(R, device=DEV)[None, :]).reshape(
+        -1
+    )
+    grad = torch.randn(N, H, device=DEV, dtype=dt, generator=g)
+
+    def run(m):
+        for t in lora:
+            t.grad = None
+        x = torch.randn(
+            N, H, device=DEV, dtype=dt, generator=torch.Generator(DEV).manual_seed(7)
+        ).requires_grad_(True)
+        scattermoe_experts_forward_ep(m, x, idx, w).backward(grad)
+        return x.grad.detach(), [t.grad.clone() for t in lora]
+
+    dx_mx, gl_mx = run(_module(gu_mx, dn_mx))
+    dx_bf, gl_bf = run(_module(gu_deq, dn_deq))
+
+    def rel(a, b):
+        return (a - b).float().abs().max().item() / max(
+            b.float().abs().max().item(), 1e-6
+        )
+
+    assert rel(dx_mx, dx_bf) < 6e-2
+    for i, slc in (
+        (0, row),
+        (1, (slice(None), row)),
+        (2, row),
+        (3, (slice(None), row)),
+    ):
+        assert rel(gl_mx[i][slc], gl_bf[i][slc]) < 6e-2, f"lora grad {i}"
+
+    if (
+        not _sonicmoe_kernel_supported()
+    ):  # e.g. sm_120: sonicmoe+EP falls back to scattermoe
+        out = _sonicmoe_local(
+            _module(gu_mx, dn_mx), torch.randn(N, H, device=DEV, dtype=dt), idx, w
+        )
+        assert out.shape == (N, H) and torch.isfinite(out).all()
+
+
+def test_mxfp4_without_lora_falls_back_to_dequant(monkeypatch):
+    """MXFP4 base with no LoRA must not hit the fused (LoRA-only) MX kernel."""
+    dt = torch.bfloat16
+    g = torch.Generator(device=DEV).manual_seed(0)
+
+    def mx(W):
+        return MXTensor.to_mx(W, elem_dtype=torch.float4_e2m1fn_x2, block_size=32)
+
+    m = _module(
+        mx(torch.randn(E, 2 * IM, H, device=DEV, dtype=dt, generator=g) * 0.1),
+        mx(torch.randn(E, H, IM, device=DEV, dtype=dt, generator=g) * 0.1),
+    )
+    monkeypatch.setattr(ex, "_has_peft_wrapper", lambda s: False)  # no LoRA
+    x = torch.randn(N, H, device=DEV, dtype=dt)
+    idx = torch.randint(0, E, (N, K), device=DEV)
+    w = torch.rand(N, K, device=DEV).to(dt)
+    out = scattermoe_experts_forward(
+        m, x, idx, w
+    )  # dequant-on-cast path, must not raise
+    assert out.shape == (N, H) and torch.isfinite(out).all()
diff --git a/tests/integrations/kernels/scattermoe_lora/test_mxfp4_integration.py b/tests/integrations/kernels/scattermoe_lora/test_mxfp4_integration.py
new file mode 100644
index 0000000000..353969ac73
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_mxfp4_integration.py
@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+End-to-end integration test for MXFP4 expert weights through the ScatterMoE
+LoRA path.
+
+We build a tiny synthetic DeepSeek-V4-style MoE block (E=8, hidden=512,
+intermediate=256, top_k=2), MX-quantize the gate/up and down projection
+expert weights via ``torchao.MXTensor.to_mx``, then compare two stacks
+forward-only:
+
+  1. **Reference** — pure PyTorch per-expert loop using the bf16 dequant of
+     the *same* MX weights. This stands in for "stock HF transformers MoE
+     with ``Mxfp4Config`` applied" — both stacks read the same physical MX
+     packed/scale buffers, so any divergence comes from the Axolotl
+     ScatterMoE plumbing (routing flatten/sort, scatter2scatter, fused
+     dequant kernel), not from differing weight quantization.
+
+  2. **Axolotl ScatterMoE** — ``parallel_linear_lora`` driven by an
+     ``MXWeights`` container, LoRA disabled (A = B = 0). Tests both
+     Strategy A (selective dequant to bf16) and Strategy B (fused MX
+     Triton kernel) so the spec'd "stock vs scattermoe" parity check
+     covers both code paths.
+
+Comparison tolerance is looser than the unit tests (``atol=rtol=5e-3``)
+because the per-expert PyTorch reference accumulates in fp32 while the
+Triton path emits bf16 outputs whose final cast rounds.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.mx_weights import (
+    selective_mx_weights_fwd,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+    flatten_sort_count,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_linear_lora import (
+    parallel_linear_lora,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (
+    get_active_experts,
+    remap_expert_indices,
+    selective_expert_weights,
+    selective_lora_weights,
+)
+
+torchao = pytest.importorskip("torchao")
+from torchao.prototype.mx_formats.mx_tensor import MXTensor  # noqa: E402
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA required for MX kernels"
+)
+
+
+# DeepSeek-V4-style tiny config (small enough for fast unit testing)
+E = 8
+HIDDEN = 512
+INTERMEDIATE = 256
+TOP_K = 2
+M = 16  # batch * seq
+
+
+def _build_synthetic_moe():
+    """Return (gate_up_mx, down_mx, gate_up_ref, down_ref, router_w)
+    matching a DeepSeek-V4 expert block:
+
+      * ``gate_up_proj``: per-expert ``[hidden, 2*intermediate]`` (split into
+        gate and up halves after the matmul).
+      * ``down_proj``:    per-expert ``[intermediate, hidden]``.
+
+    Storage layout matches axolotl's convention ``[E, N, K]`` where K is the
+    contraction axis the kernel will block on. ``gate_up`` has K=hidden,
+    N=2*intermediate; ``down`` has K=intermediate, N=hidden.
+
+    bf16 reference tensors are the dequantizations of the *same* MX
+    buffers, so the only test source of divergence is the kernel paths.
+    """
+    torch.manual_seed(42)
+    # Scale ~ 1/sqrt(fan_in) so per-layer outputs stay in order-1 range and
+    # bf16 final-cast noise is not amplified by the magnitude.
+    gup_scale = 1.0 / (HIDDEN**0.5)
+    down_scale = 1.0 / (INTERMEDIATE**0.5)
+    gate_up = (
+        torch.randn(E, 2 * INTERMEDIATE, HIDDEN, device=DEVICE, dtype=DTYPE) * gup_scale
+    )
+    down = torch.randn(E, HIDDEN, INTERMEDIATE, device=DEVICE, dtype=DTYPE) * down_scale
+
+    gate_up_mx = MXTensor.to_mx(
+        gate_up, elem_dtype=torch.float4_e2m1fn_x2, block_size=32
+    )
+    down_mx = MXTensor.to_mx(down, elem_dtype=torch.float4_e2m1fn_x2, block_size=32)
+    gate_up_ref = gate_up_mx.dequantize(DTYPE).contiguous()
+    down_ref = down_mx.dequantize(DTYPE).contiguous()
+
+    router_w = torch.randn(E, HIDDEN, device=DEVICE, dtype=DTYPE) * 0.1
+    return gate_up_mx, down_mx, gate_up_ref, down_ref, router_w
+
+
+def _reference_moe_forward(x, router_w, gate_up_ref, down_ref):
+    """Stand-in for stock HF MoE with Mxfp4Config: per-token routing +
+    per-expert matmul on dequantized bf16 weights."""
+    # Softmax-topk routing
+    router_logits = F.linear(x, router_w)  # [M, E]
+    routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+    routing_weights, selected = torch.topk(routing_weights, TOP_K, dim=-1)
+    routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+    routing_weights = routing_weights.to(x.dtype)
+
+    out = torch.zeros_like(x)
+    for e in range(E):
+        # Tokens routed to expert e: positions (token_id, k_slot)
+        mask = selected == e  # [M, TOP_K]
+        if not mask.any():
+            continue
+        token_ids, slot_ids = mask.nonzero(as_tuple=True)
+        x_e = x[token_ids]  # [n_e, HIDDEN]
+        gup = x_e @ gate_up_ref[e].t()  # [n_e, 2*INTERMEDIATE]
+        gate, up = gup.chunk(2, dim=-1)
+        h = F.silu(gate) * up
+        y_e = h @ down_ref[e].t()  # [n_e, HIDDEN]
+        # Weighted accumulate
+        w_e = routing_weights[token_ids, slot_ids].unsqueeze(-1)
+        out.index_add_(0, token_ids, w_e * y_e)
+    return out
+
+
+class _MockExperts:
+    def __init__(self, gate_up, down):
+        self.gate_up_proj = gate_up
+        self.down_proj = down
+        self.num_experts = E
+
+
+def _axolotl_moe_forward(x, router_w, gate_up_param, down_param, *, strategy: str):
+    """Run the Axolotl ScatterMoE LoRA path with LoRA disabled (A=B=0).
+
+    ``strategy='A'``: ``gate_up_param``/``down_param`` are torchao MXTensors;
+    we dequantize the active experts to bf16 and call the bf16 kernel.
+
+    ``strategy='B'``: same MXTensors but routed through the fused MX kernel
+    via the ``MXWeights`` container.
+    """
+    # Routing — same softmax+topk shape as the reference
+    router_logits = F.linear(x, router_w)
+    routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+    routing_weights, selected = torch.topk(routing_weights, TOP_K, dim=-1)
+    routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+    routing_weights = routing_weights.to(x.dtype)
+
+    sei, ssi, eo = flatten_sort_count(selected, num_experts=E)
+    active = get_active_experts(sei, E)
+    remapped, compact_offsets = remap_expert_indices(sei, eo, active, E)
+
+    # Build LoRA tensors with A=B=0 so the LoRA term is zero.
+    rank = 4
+    lora_A = torch.zeros(rank * E, HIDDEN, device=DEVICE, dtype=DTYPE)
+    lora_B_gup = torch.zeros(2 * INTERMEDIATE, rank * E, device=DEVICE, dtype=DTYPE)
+    lora_B_down = torch.zeros(HIDDEN, rank * E, device=DEVICE, dtype=DTYPE)
+    lora_A_inter = torch.zeros(rank * E, INTERMEDIATE, device=DEVICE, dtype=DTYPE)
+    A_gup_c, B_gup_c = selective_lora_weights(lora_A, lora_B_gup, active, E)
+    A_dn_c, B_dn_c = selective_lora_weights(lora_A_inter, lora_B_down, active, E)
+
+    experts = _MockExperts(gate_up_param, down_param)
+
+    if strategy == "A":
+        gate_up_W = (
+            selective_expert_weights(experts, "gate_up_proj", active)
+            .transpose(2, 1)
+            .contiguous()
+        )
+        down_W = (
+            selective_expert_weights(experts, "down_proj", active)
+            .transpose(2, 1)
+            .contiguous()
+        )
+        gup_W = gate_up_W
+        dwn_W = down_W
+    elif strategy == "B":
+        gup_W = selective_mx_weights_fwd(gate_up_param, active)
+        dwn_W = selective_mx_weights_fwd(down_param, active)
+    else:
+        raise ValueError(strategy)
+
+    gup = parallel_linear_lora(
+        x,
+        gup_W,
+        TOP_K,
+        remapped,
+        ssi,
+        compact_offsets,
+        lora_A=A_gup_c,
+        lora_B=B_gup_c,
+        scaling=0.0,
+        grouped_in=False,
+        grouped_out=True,
+        use_fused_dX=True,
+        use_fused_gather=True,
+    )
+    gate, up = gup.chunk(2, dim=-1)
+    h = F.silu(gate) * up
+    out = parallel_linear_lora(
+        h,
+        dwn_W,
+        1,
+        remapped,
+        ssi,
+        compact_offsets,
+        lora_A=A_dn_c,
+        lora_B=B_dn_c,
+        scaling=0.0,
+        gates=routing_weights,
+        grouped_in=True,
+        grouped_out=False,
+        use_fused_dX=True,
+        use_fused_gather=True,
+    )
+    return out
+
+
+@pytest.mark.parametrize("strategy", ["A", "B"])
+def test_mxfp4_moe_block_matches_pytorch_reference(strategy):
+    """The Axolotl ScatterMoE MX path must match the per-expert PyTorch
+    reference (operating on the same MX dequantized weights) within
+    integration-grade tolerance."""
+    gate_up_mx, down_mx, gate_up_ref, down_ref, router_w = _build_synthetic_moe()
+
+    torch.manual_seed(7)
+    x = torch.randn(M, HIDDEN, device=DEVICE, dtype=DTYPE)
+
+    ref = _reference_moe_forward(x, router_w, gate_up_ref, down_ref)
+    out = _axolotl_moe_forward(x, router_w, gate_up_mx, down_mx, strategy=strategy)
+
+    assert ref.shape == out.shape == (M, HIDDEN)
+    assert torch.allclose(ref, out, atol=5e-3, rtol=5e-3), (
+        f"Strategy {strategy} MoE block diverges from PyTorch reference: "
+        f"max abs={(ref - out).abs().max().item():.4e}, "
+        f"max rel={((ref - out).abs() / (ref.abs() + 1e-6)).max().item():.4e}"
+    )
diff --git a/tests/integrations/kernels/scattermoe_lora/test_nvfp4_experts_forward.py b/tests/integrations/kernels/scattermoe_lora/test_nvfp4_experts_forward.py
new file mode 100644
index 0000000000..5e858693dc
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_nvfp4_experts_forward.py
@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""scattermoe routes base NVFP4 + LoRA through the fused kernel (no bf16 dequant).
+
+NVFP4 shares the FP4 E2M1 packing with MXFP4 but scales per 16-element block with an
+E4M3 (fp8) value (times an optional per-tensor scale). It reuses the MXWeights container
+with ``scale_is_linear=True`` (the E4M3 block scale pre-folded with the per-tensor scale)
+and the same fused Triton kernel with a linear-scale branch + block-16. Validated vs a
+reference = the same forward on the dequantized weights, within NV rounding tolerance,
+through both the standard and EP (sentinel-skip) forwards.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+NVFP4Tensor = pytest.importorskip(
+    "torchao.prototype.mx_formats.nvfp4_tensor", reason="torchao required"
+).NVFP4Tensor
+
+import axolotl.integrations.kernels.libs.scattermoe_lora.experts as ex  # noqa: E402
+from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (  # noqa: E402
+    scattermoe_experts_forward,
+    scattermoe_experts_forward_ep,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (  # noqa: E402
+    is_mxfp4_param,
+    is_nvfp4_param,
+)
+
+DEV = "cuda"
+E, H, IM, N, K, R, SC = 8, 1024, 512, 512, 4, 16, 0.5
+
+
+def _module(gu, dn):
+    return SimpleNamespace(
+        num_experts=E,
+        gate_up_proj=gu,
+        down_proj=dn,
+        act_fn=torch.nn.functional.silu,
+        is_transposed=False,
+        is_concatenated=True,
+        has_bias=False,
+        has_gate=True,
+    )
+
+
+def _rel(a, b):
+    return (a - b).float().abs().max().item() / max(b.float().abs().max().item(), 1e-6)
+
+
+def test_is_nvfp4_param_disjoint_from_mxfp4():
+    nv = NVFP4Tensor.to_nvfp4(
+        torch.randn(E, IM, H, device=DEV, dtype=torch.bfloat16), block_size=16
+    )
+    assert is_nvfp4_param(nv) and not is_mxfp4_param(nv)
+    assert not is_nvfp4_param(torch.randn(2, 2, device=DEV))  # plain tensor
+
+
+def _routing_std(g):
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    return idx, torch.rand(N, K, device=DEV, generator=g).to(torch.bfloat16)
+
+
+def _routing_ep(g):
+    idx = torch.randint(0, E, (N, K), device=DEV, generator=g)
+    remote = torch.rand(N, K, device=DEV, generator=g) < 0.5
+    remote[:, 0] = False  # >=1 valid slot per token (DeepEP guarantee)
+    idx = torch.where(remote, torch.full_like(idx, -1), idx)
+    return idx, torch.rand(N, K, device=DEV, generator=g).to(torch.bfloat16)
+
+
+@pytest.mark.parametrize(
+    "fwd,routing",
+    [
+        (scattermoe_experts_forward, _routing_std),
+        (scattermoe_experts_forward_ep, _routing_ep),
+    ],
+)
+def test_nvfp4_fused_matches_dequant(fwd, routing, monkeypatch):
+    dt = torch.bfloat16
+    # On Blackwell (sm100/sm120) NVFP4+LoRA must run the grouped fp4 path; the legacy fused-MX kernel
+    # SIGSEGVs there and is now guarded off. Select the supported grouped backend for the test.
+    monkeypatch.setattr(ex.RUNTIME, "fp4_grouped_mode", "nvfp4")
+    g = torch.Generator(device=DEV).manual_seed(0)
+    gu_nv = NVFP4Tensor.to_nvfp4(
+        torch.randn(E, 2 * IM, H, device=DEV, dtype=dt, generator=g) * 0.1,
+        block_size=16,
+    )
+    dn_nv = NVFP4Tensor.to_nvfp4(
+        torch.randn(E, H, IM, device=DEV, dtype=dt, generator=g) * 0.1, block_size=16
+    )
+    gu_deq, dn_deq = (
+        gu_nv.dequantize(dt).contiguous(),
+        dn_nv.dequantize(dt).contiguous(),
+    )
+
+    def mk(*s):
+        return (
+            torch.randn(*s, device=DEV, dtype=dt, generator=g) * 0.05
+        ).requires_grad_(True)
+
+    A1, B1, A2, B2 = mk(R * E, H), mk(2 * IM, R * E), mk(R * E, IM), mk(H, R * E)
+    lora = [A1, B1, A2, B2]
+    monkeypatch.setattr(ex, "_has_peft_wrapper", lambda s: True)
+    monkeypatch.setattr(
+        ex, "_unwrap_experts_lora", lambda s: (s, (A1, B1, SC), (A2, B2, SC))
+    )
+
+    idx, w = routing(g)
+    valid = idx.reshape(-1) >= 0
+    active = torch.unique(torch.sort(idx.reshape(-1)[valid]).values)
+    row = (active.long()[:, None] * R + torch.arange(R, device=DEV)[None, :]).reshape(
+        -1
+    )
+    grad = torch.randn(N, H, device=DEV, dtype=dt, generator=g)
+
+    def run(m):
+        for t in lora:
+            t.grad = None
+        x = torch.randn(
+            N, H, device=DEV, dtype=dt, generator=torch.Generator(DEV).manual_seed(7)
+        ).requires_grad_(True)
+        fwd(m, x, idx, w).backward(grad)
+        return x.grad.detach(), [t.grad.clone() for t in lora]
+
+    dx_nv, gl_nv = run(_module(gu_nv, dn_nv))  # fused NVFP4 (no dequant)
+    dx_bf, gl_bf = run(_module(gu_deq, dn_deq))  # bf16 dequant reference
+    assert _rel(dx_nv, dx_bf) < 6e-2
+    for i, slc in (
+        (0, row),
+        (1, (slice(None), row)),
+        (2, row),
+        (3, (slice(None), row)),
+    ):
+        assert _rel(gl_nv[i][slc], gl_bf[i][slc]) < 6e-2, f"lora grad {i}"
+        assert torch.isfinite(gl_nv[i]).all()
diff --git a/tests/integrations/kernels/scattermoe_lora/test_nvfp4_fsdp_sharding.py b/tests/integrations/kernels/scattermoe_lora/test_nvfp4_fsdp_sharding.py
new file mode 100644
index 0000000000..e52b78ec20
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_nvfp4_fsdp_sharding.py
@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""CPU round-trip tests for the FSDP2 sharding substrate on torchao ``NVFP4Tensor``.
+
+The substrate (``...scattermoe_lora.nvfp4_fsdp``) teaches a frozen, expert-sharded
+``NVFP4Tensor`` the aten ops FSDP2 needs to dim-0 (expert axis) shard the weight:
+``split`` / ``narrow`` / ``slice`` / ``new_zeros`` / ``clone`` / ``detach`` /
+``as_strided`` / ``view`` / ``copy_`` plus the ``fsdp_pre_all_gather`` /
+``fsdp_post_all_gather`` reconstruct hooks. These are pure tensor logic over the inner
+``(qdata, scale, per_tensor_scale)`` triple and are fully exercisable on CPU - no GPU, no
+Blackwell, no fused kernel.
+
+Each round-trip test simulates the FSDP2 shard -> all-gather -> reconstruct flow and asserts
+the reconstructed tensor dequantizes equal to the unsharded original across several world
+sizes including one where the expert axis does NOT divide evenly. Covers the B11 data-
+independent FSDP2 collective contract for the NVFP4 substrate.
+
+Round-trip verification uses a unit (scalar) per_tensor_scale - the real DSV4 frozen-expert
+case and the only one torchao's ``NVFP4Tensor.dequantize`` supports against per-block scales.
+A per-expert per_tensor_scale buffer is checked at the tensor-shape level (its split/narrow
+slicing) since torchao cannot dequantize that shape for an equality check.
+
+Skip gate: requires torchao's NVFP4Tensor. The substrate module is loaded by file path so the
+scattermoe_lora package __init__ (which imports triton) is NOT executed - the ops are pure CPU
+tensor logic, so this runs on the no-CUDA / triton-less CI lane (matching the Float8 sibling).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+
+import pytest
+import torch
+
+NVFP4Tensor = pytest.importorskip(
+    "torchao.prototype.mx_formats.nvfp4_tensor", reason="torchao required"
+).NVFP4Tensor
+
+
+def _load_nvfp4_fsdp():
+    """Load nvfp4_fsdp by file path so the scattermoe_lora package __init__ (which imports triton) is
+    NOT executed. The module needs only torch + torchao, so it imports triton-free and runs on CPU."""
+    import axolotl
+
+    path = os.path.join(
+        os.path.dirname(axolotl.__file__),
+        "integrations",
+        "kernels",
+        "libs",
+        "scattermoe_lora",
+        "nvfp4_fsdp.py",
+    )
+    spec = importlib.util.spec_from_file_location(
+        "_axolotl_nvfp4_fsdp_under_test", path
+    )
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+nvfp4_fsdp = _load_nvfp4_fsdp()
+nvfp4_fsdp.patch_nvfp4_fsdp()
+
+DEV = "cpu"
+
+
+def _cdiv(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
+def _make_nvfp4(E: int, N: int, K: int, seed: int = 0, per_expert_pts: bool = False):
+    """Build a frozen ``NVFP4Tensor`` [E,N,K] (qdata [E,N,K//2], scale [E,N,K//16]).
+
+    per_expert_pts=False uses a unit scalar per_tensor_scale (the real DSV4 case, dequant-able);
+    True uses a per-expert [E] buffer to exercise the per-expert pts split path at the shape level.
+    """
+    torch.manual_seed(seed)
+    packed = torch.randint(0, 255, (E, N, K // 2), dtype=torch.uint8)
+    scale = (torch.rand(E, N, K // 16) * 0.5 + 0.5).to(torch.float8_e4m3fn)
+    pts = (torch.rand(E) * 0.5 + 0.5) if per_expert_pts else torch.ones(())
+    return NVFP4Tensor(packed.contiguous(), scale.contiguous(), 16, torch.bfloat16, pts)
+
+
+def _shard_all_gather_reconstruct(nv, world: int):
+    """Simulate one FSDP2 dim-0 (expert axis) shard -> all-gather -> reconstruct cycle on CPU.
+
+    Pads the expert axis to a per-rank-equal count (FSDP2 pads uneven shards via new_zeros +
+    copy_ into a narrow), chunks into ``world`` shards, runs each shard's fsdp_pre_all_gather,
+    concatenates the gathered inner tensors (a scalar per_tensor_scale is replicated, not
+    concatenated), and reconstructs via fsdp_post_all_gather. Returns the unpadded NVFP4Tensor.
+    """
+    E = nv.shape[0]
+    shard = _cdiv(E, world)
+    padded = shard * world
+    if padded != E:
+        full = nv.new_zeros([padded, *nv.shape[1:]])
+        full.narrow(0, 0, E).copy_(nv)
+    else:
+        full = nv.clone()
+    shards = torch.split(full, shard, 0)
+    assert len(shards) == world, (len(shards), world)
+    pre = [s.fsdp_pre_all_gather(mesh=None) for s in shards]
+    meta = pre[0][1]
+    gathered = []
+    for k in range(len(pre[0][0])):
+        tensors = [p[0][k] for p in pre]
+        if tensors[0].dim() == 0:
+            gathered.append(
+                tensors[0]
+            )  # scalar per_tensor_scale: replicated, not sharded
+        else:
+            gathered.append(torch.cat(tensors, 0))
+    recon, _ = shards[0].fsdp_post_all_gather(tuple(gathered), meta, torch.bfloat16)
+    return recon.narrow(0, 0, E)
+
+
+@pytest.mark.parametrize("E,N,K", [(4, 8, 16), (6, 8, 32)], ids=["E4_K16", "E6_K32"])
+def test_nvfp4_split_narrow_slice_roundtrip(E, N, K):
+    """split/narrow/slice on the expert axis dequantize equal to the original slice."""
+    nv = _make_nvfp4(E, N, K)
+    orig = nv.dequantize(torch.bfloat16)
+
+    parts = torch.split(nv, 2, 0)
+    rec = torch.cat([p.dequantize(torch.bfloat16) for p in parts], 0)
+    assert torch.equal(rec, orig), "split round-trip diverged"
+
+    nar = torch.narrow(nv, 0, 1, 2)
+    assert torch.equal(nar.dequantize(torch.bfloat16), orig[1:3]), "narrow diverged"
+
+    sl = nv[1:3]  # aten.slice on dim 0
+    assert torch.equal(sl.dequantize(torch.bfloat16), orig[1:3]), "slice diverged"
+
+
+def test_nvfp4_clone_detach_as_strided_roundtrip():
+    """clone / detach / as_strided preserve the dequantized value (view ops return new tensors)."""
+    nv = _make_nvfp4(6, 8, 32)
+    orig = nv.dequantize(torch.bfloat16)
+    assert torch.equal(nv.clone().dequantize(torch.bfloat16), orig), "clone diverged"
+    assert torch.equal(nv.detach().dequantize(torch.bfloat16), orig), "detach diverged"
+
+    shape = list(nv.shape)
+    stride = [1] * len(shape)
+    for i in range(len(shape) - 2, -1, -1):
+        stride[i] = stride[i + 1] * shape[i + 1]
+    ast = torch.as_strided(nv, shape, stride)
+    assert torch.equal(ast.dequantize(torch.bfloat16), orig), "as_strided diverged"
+
+
+@pytest.mark.parametrize(
+    "E,worlds",
+    [(6, (1, 2, 3)), (6, (4,)), (8, (3,))],
+    ids=["even_1_2_3", "uneven_world4_E6", "uneven_world3_E8"],
+)
+def test_nvfp4_fsdp_all_gather_roundtrip(E, worlds):
+    """FSDP2 shard -> all-gather -> reconstruct round-trips across worlds, incl. uneven expert splits.
+
+    world=4 over E=6 and world=3 over E=8 do NOT divide evenly: the shard is padded on the expert
+    axis (new_zeros + copy_ into a narrow) and narrowed back after reconstruct. Pure CPU logic.
+    """
+    nv = _make_nvfp4(E, 8, 32)
+    orig = nv.dequantize(torch.bfloat16)
+    for world in worlds:
+        recon = _shard_all_gather_reconstruct(nv, world)
+        assert torch.equal(recon.dequantize(torch.bfloat16), orig), (
+            f"FSDP all-gather reconstruct diverged at world={world} (E={E})"
+        )
+
+
+def test_nvfp4_new_zeros_preserves_trailing_shapes():
+    """new_zeros varies only dim 0 (experts); qdata/scale trailing dims are preserved verbatim."""
+    nv = _make_nvfp4(6, 8, 32)
+    nz = nv.new_zeros([10, 8, 32])
+    assert nz.qdata.shape == (10, *nv.qdata.shape[1:])
+    assert nz.scale.shape == (10, *nv.scale.shape[1:])
+
+
+def test_nvfp4_per_expert_pts_split_narrow_shapes():
+    """A per-expert per_tensor_scale [E] is sliced alongside the expert axis on split/narrow.
+
+    torchao cannot dequantize a per-expert per_tensor_scale against per-block scales, so this
+    checks the slicing at the shape level (the FSDP substrate must carry the [E] buffer through).
+    """
+    nv = _make_nvfp4(6, 8, 32, per_expert_pts=True)
+    assert nv.per_tensor_scale is not None and nv.per_tensor_scale.shape == (6,)
+
+    parts = torch.split(nv, 2, 0)
+    assert all(p.per_tensor_scale.shape == (2,) for p in parts), (
+        "per-expert per_tensor_scale not split on the expert axis"
+    )
+    assert torch.equal(
+        torch.cat([p.per_tensor_scale for p in parts], 0), nv.per_tensor_scale
+    )
+
+    nar = torch.narrow(nv, 0, 1, 3)
+    assert nar.per_tensor_scale.shape == (3,)
+    assert torch.equal(nar.per_tensor_scale, nv.per_tensor_scale[1:4])
+
+
+def test_nvfp4_split_dim_other_than_zero_rejected():
+    """Sharding is restricted to dim 0 (expert axis); a non-zero dim split must raise."""
+    nv = _make_nvfp4(4, 8, 16)
+    with pytest.raises(NotImplementedError, match="only on dim 0"):
+        torch.split(nv, 4, 1)
diff --git a/tests/integrations/kernels/scattermoe_lora/test_nvfp4_weight_converter.py b/tests/integrations/kernels/scattermoe_lora/test_nvfp4_weight_converter.py
new file mode 100644
index 0000000000..7773e0ee44
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_nvfp4_weight_converter.py
@@ -0,0 +1,398 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Tests for the gemma4 NVFP4 MoE-expert WeightConverter shim.
+
+``Nvfp4ExpertsDeserialize`` fuses the per-expert NVFP4 tensors a
+``nvidia/Gemma-4-...-NVFP4`` checkpoint ships (separate ``gate_proj`` /
+``up_proj`` / ``down_proj`` under ``...experts.E.*``) into the single fused
+``gate_up_proj`` / ``down_proj`` ``NVFP4Tensor`` that ``Gemma4TextExperts``
+expects, and assigns it to the module in place (mirroring transformers'
+``Mxfp4Deserialize``).
+
+The fusion is pure tensor stacking/concatenation of *already-quantized* bytes,
+so the fused result must be **bit-exact** (maxerr 0.0) against the reference
+fusion — that's the core correctness contract these tests pin.
+
+The converter logic is dtype-only (no GPU kernels), but ``NVFP4Tensor`` /
+``to_nvfp4`` quantization runs on CUDA, so the suite is CUDA-gated.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+NVFP4Tensor = pytest.importorskip(
+    "torchao.prototype.mx_formats.nvfp4_tensor", reason="torchao required"
+).NVFP4Tensor
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.nvfp4_weight_converter import (  # noqa: E402
+    Nvfp4ExpertsDeserialize,
+    nvfp4_experts_weight_converters,
+    register_gemma4_nvfp4_converters,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (  # noqa: E402
+    is_nvfp4_param,
+)
+
+DEV = "cuda"
+# small but realistic gemma4-MoE-ish shapes (H % 16 == 0, I % 16 == 0)
+E, H, I = 4, 128, 64
+BLOCK = 16
+
+
+def _quant(W: torch.Tensor, pts: torch.Tensor):
+    """Per-expert NVFP4 of a 2D weight with a fixed per-tensor scale (as the
+    checkpoint stores it)."""
+    return NVFP4Tensor.to_nvfp4(W, block_size=BLOCK, per_tensor_scale=pts)
+
+
+def _make_checkpoint_experts(seed: int = 0):
+    """Build E experts of raw per-expert NVFP4 tensors exactly as the checkpoint
+    stores them (separate gate/up/down, each its own qdata + e4m3 scale and its own
+    per-tensor weight_scale_2 scalar). Per-expert scales are DISTINCT so the tests pin
+    that the converter keeps each expert's scale (not expert-0's for all). Returns the
+    fused per-expert scale tensor [E,1,1] (the converter's output shape)."""
+    g = torch.Generator(device=DEV).manual_seed(seed)
+    experts = []
+    pts_each = []
+    for e in range(E):
+        pts = torch.tensor(
+            0.5 + 0.1 * e, device=DEV, dtype=torch.float32
+        )  # distinct per expert
+        pts_each.append(pts)
+        Wg = torch.randn(I, H, generator=g, device=DEV, dtype=torch.bfloat16)
+        Wu = torch.randn(I, H, generator=g, device=DEV, dtype=torch.bfloat16)
+        Wd = torch.randn(H, I, generator=g, device=DEV, dtype=torch.bfloat16)
+        experts.append(
+            {
+                "gate_proj": _quant(Wg, pts),
+                "up_proj": _quant(Wu, pts),
+                "down_proj": _quant(Wd, pts),
+            }
+        )
+    pts_per_expert = torch.stack(pts_each).view(
+        E, 1, 1
+    )  # matches converter output shape
+    return experts, pts_per_expert
+
+
+def _input_dict(experts, projs, full_layer_name):
+    """Replicate transformers' ``WeightConverter.materialize_tensors`` output:
+    keyed by source-pattern string, value = list of per-expert raw tensors."""
+    d: dict[str, list[torch.Tensor]] = {}
+    for proj in projs:
+        d[f"experts.*.{proj}.weight"] = [e[proj].qdata for e in experts]
+        d[f"experts.*.{proj}.weight_scale"] = [e[proj].scale for e in experts]
+        d[f"experts.*.{proj}.weight_scale_2"] = [
+            e[proj].per_tensor_scale for e in experts
+        ]
+    return d
+
+
+def _ref_gate_up(experts):
+    qdata = torch.cat(
+        [
+            torch.stack([e["gate_proj"].qdata for e in experts], 0),
+            torch.stack([e["up_proj"].qdata for e in experts], 0),
+        ],
+        dim=1,
+    )
+    scale = torch.cat(
+        [
+            torch.stack([e["gate_proj"].scale for e in experts], 0),
+            torch.stack([e["up_proj"].scale for e in experts], 0),
+        ],
+        dim=1,
+    )
+    return qdata, scale
+
+
+def _ref_down(experts):
+    qdata = torch.stack([e["down_proj"].qdata for e in experts], 0)
+    scale = torch.stack([e["down_proj"].scale for e in experts], 0)
+    return qdata, scale
+
+
+def _run_convert(op, experts, proj, projs):
+    """Drive the op the same way ``WeightConverter.convert`` does, against a tiny
+    module, and return the mutated module."""
+    full = f"model.language_model.layers.0.experts.{proj}"
+    module = SimpleNamespace()
+    setattr(module, proj, nn.Parameter(torch.zeros(1), requires_grad=False))
+
+    missing = {full}
+    op.convert(
+        _input_dict(experts, projs, full),
+        source_patterns=None,
+        target_patterns=None,
+        full_layer_name=full,
+        model=_AttrModel(full, module),
+        missing_keys=missing,
+    )
+    return module, missing
+
+
+class _AttrModel:
+    """Minimal stand-in so ``get_module_from_name(model, full_layer_name)`` returns our
+    experts module. That helper rsplits the dotted path and calls
+    ``model.get_submodule(parent_path)``; we map any parent path to the single leaf module."""
+
+    def __init__(self, full_layer_name: str, module):
+        self._full = full_layer_name
+        self._leaf = module
+
+    def get_submodule(self, name):
+        return self._leaf
+
+
+# ---------------------------------------------------------------------------
+# correctness: fused NVFP4Tensor is bit-exact vs the reference fusion
+# ---------------------------------------------------------------------------
+
+
+def test_gate_up_fusion_bit_exact():
+    experts, pts = _make_checkpoint_experts(seed=1)
+    op = Nvfp4ExpertsDeserialize()
+    module, missing = _run_convert(
+        op, experts, "gate_up_proj", ("gate_proj", "up_proj")
+    )
+
+    fused = module.gate_up_proj
+    assert is_nvfp4_param(fused), "gate_up_proj must be an NVFP4Tensor"
+
+    ref_qd, ref_sc = _ref_gate_up(experts)
+    assert fused.qdata.shape == (E, 2 * I, H // 2)
+    assert fused.scale.shape == (E, 2 * I, H // BLOCK)
+    # bit-exact: fusion is concatenation of already-quantized bytes -> maxerr 0.0
+    assert torch.equal(fused.qdata, ref_qd), "qdata not bit-exact"
+    assert torch.equal(fused.scale, ref_sc), "scale not bit-exact"
+    # per-expert weight_scale_2 preserved as [E,1,1] (each expert keeps its own scale)
+    assert torch.equal(fused.per_tensor_scale.to(torch.float32), pts), (
+        "per_tensor_scale mismatch"
+    )
+    # loader contract: param consumed, requires_grad off
+    assert "model.language_model.layers.0.experts.gate_up_proj" not in missing
+    assert fused.requires_grad is False
+
+
+def test_down_fusion_bit_exact():
+    experts, pts = _make_checkpoint_experts(seed=2)
+    op = Nvfp4ExpertsDeserialize()
+    module, missing = _run_convert(op, experts, "down_proj", ("down_proj",))
+
+    fused = module.down_proj
+    assert is_nvfp4_param(fused)
+    ref_qd, ref_sc = _ref_down(experts)
+    assert fused.qdata.shape == (E, H, I // 2)
+    assert fused.scale.shape == (E, H, I // BLOCK)
+    assert torch.equal(fused.qdata, ref_qd)
+    assert torch.equal(fused.scale, ref_sc)
+    assert torch.equal(fused.per_tensor_scale.to(torch.float32), pts)
+
+
+# ---------------------------------------------------------------------------
+# correctness: fused tensor dequantizes to cat(dequant(gate), dequant(up))
+# (the property the scattermoe forward relies on) at maxerr 0.0
+# ---------------------------------------------------------------------------
+
+
+def test_gate_up_dequant_self_consistent():
+    experts, _ = _make_checkpoint_experts(seed=3)
+    op = Nvfp4ExpertsDeserialize()
+    module, _ = _run_convert(op, experts, "gate_up_proj", ("gate_proj", "up_proj"))
+
+    fused = module.gate_up_proj
+    pe = fused.per_tensor_scale.reshape(-1)  # [E] per-expert weight_scale_2
+    for e in range(E):
+        # torchao's per-expert index (fused[e]) does not slice the [E,1,1] per_tensor_scale, so
+        # rebuild the expert-e slice with its own scalar scale (what the real per-expert kernel uses).
+        fused_e = NVFP4Tensor(
+            fused.qdata[e],
+            fused.scale[e],
+            BLOCK,
+            torch.bfloat16,
+            per_tensor_scale=pe[e],
+        )
+        fused_dq = fused_e.dequantize(torch.bfloat16)
+        gate_dq = experts[e]["gate_proj"].dequantize(torch.bfloat16)
+        up_dq = experts[e]["up_proj"].dequantize(torch.bfloat16)
+        ref = torch.cat([gate_dq, up_dq], dim=0)
+        err = (fused_dq - ref).abs().max().item()
+        assert err == 0.0, f"expert {e}: dequant maxerr {err} != 0"
+
+
+def test_down_dequant_self_consistent():
+    experts, _ = _make_checkpoint_experts(seed=4)
+    op = Nvfp4ExpertsDeserialize()
+    module, _ = _run_convert(op, experts, "down_proj", ("down_proj",))
+
+    fused = module.down_proj
+    pe = fused.per_tensor_scale.reshape(-1)  # [E] per-expert weight_scale_2
+    for e in range(E):
+        # torchao's per-expert index does not slice the [E,1,1] per_tensor_scale; rebuild the slice.
+        fused_e = NVFP4Tensor(
+            fused.qdata[e],
+            fused.scale[e],
+            BLOCK,
+            torch.bfloat16,
+            per_tensor_scale=pe[e],
+        )
+        fused_dq = fused_e.dequantize(torch.bfloat16)
+        ref = experts[e]["down_proj"].dequantize(torch.bfloat16)
+        err = (fused_dq - ref).abs().max().item()
+        assert err == 0.0, f"expert {e}: dequant maxerr {err} != 0"
+
+
+# ---------------------------------------------------------------------------
+# load-assertion: the experts land as NVFP4Tensor after the converter runs
+# ---------------------------------------------------------------------------
+
+
+def test_load_assertion_experts_are_nvfp4():
+    """Both fused params on the module are NVFP4Tensor (the post-load invariant
+    scattermoe's fused path requires)."""
+    experts, _ = _make_checkpoint_experts(seed=5)
+    op = Nvfp4ExpertsDeserialize()
+    gu_mod, _ = _run_convert(op, experts, "gate_up_proj", ("gate_proj", "up_proj"))
+    dn_mod, _ = _run_convert(op, experts, "down_proj", ("down_proj",))
+    assert is_nvfp4_param(gu_mod.gate_up_proj)
+    assert is_nvfp4_param(dn_mod.down_proj)
+    assert gu_mod.gate_up_proj.qdata.dtype == torch.uint8
+    assert gu_mod.gate_up_proj.scale.dtype == torch.float8_e4m3fn
+
+
+# ---------------------------------------------------------------------------
+# registration wiring: the converters install into the live transformers cache
+# ---------------------------------------------------------------------------
+
+
+def test_registration_installs_converters():
+    from transformers.conversion_mapping import get_checkpoint_conversion_mapping
+
+    register_gemma4_nvfp4_converters()
+    mapping = get_checkpoint_conversion_mapping("gemma4_text")
+    assert mapping is not None, "gemma4_text conversion mapping not registered"
+    # idempotent re-registration must not raise
+    register_gemma4_nvfp4_converters()
+
+    converters = nvfp4_experts_weight_converters()
+    assert len(converters) == 2
+    targets = {c.target_patterns[0] for c in converters}
+    assert targets == {"experts.gate_up_proj", "experts.down_proj"}
+    # gate_up converter sources the unfused per-expert keys
+    gu = next(c for c in converters if "gate_up_proj" in c.target_patterns[0])
+    assert any("gate_proj.weight" in s for s in gu.source_patterns)
+    assert any("up_proj.weight" in s for s in gu.source_patterns)
+
+
+def test_end_to_end_from_pretrained_load():
+    """Definitive load-assertion: save a tiny synthetic gemma4 NVFP4-modelopt
+    checkpoint (per-expert gate/up/down qdata+scale+scalar), register the
+    converters, ``from_pretrained``, and assert the experts land as NVFP4Tensor
+    with bit-exact fused qdata vs the reference fusion and no leftover keys."""
+    import json
+    import os
+    import shutil
+    import tempfile
+
+    from safetensors.torch import save_file
+    from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel
+
+    cfg = Gemma4TextConfig(
+        vocab_size=64,
+        hidden_size=H,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        head_dim=16,
+        max_position_embeddings=32,
+        enable_moe_block=True,
+        num_experts=E,
+        top_k_experts=2,
+        moe_intermediate_size=I,
+        num_kv_shared_layers=0,
+        hidden_size_per_layer_input=16,
+        vocab_size_per_layer_input=64,
+    )
+    base = Gemma4TextModel(cfg).to(torch.bfloat16)
+    sd = base.state_dict()
+
+    experts, _ = _make_checkpoint_experts(seed=7)
+    new_sd = {
+        k: v
+        for k, v in sd.items()
+        if ".experts.gate_up_proj" not in k and ".experts.down_proj" not in k
+    }
+    for e in range(E):
+        for proj in ("gate_proj", "up_proj", "down_proj"):
+            t = experts[e][proj]
+            pfx = f"layers.0.experts.{e}.{proj}"
+            new_sd[f"{pfx}.weight"] = t.qdata.cpu()
+            new_sd[f"{pfx}.weight_scale"] = t.scale.cpu()
+            new_sd[f"{pfx}.weight_scale_2"] = t.per_tensor_scale.cpu()
+
+    d = tempfile.mkdtemp()
+    try:
+        save_file(
+            {k: v.contiguous() for k, v in new_sd.items()},
+            os.path.join(d, "model.safetensors"),
+            metadata={"format": "pt"},
+        )
+        cfg_dict = cfg.to_dict()
+        cfg_dict["quantization_config"] = {
+            "quant_method": "modelopt",
+            "quant_algo": "NVFP4",
+        }
+        json.dump(cfg_dict, open(os.path.join(d, "config.json"), "w"))
+
+        register_gemma4_nvfp4_converters()
+        loaded = Gemma4TextModel.from_pretrained(d, dtype=torch.bfloat16).to(DEV)
+        ex = loaded.layers[0].experts
+
+        assert is_nvfp4_param(ex.gate_up_proj)
+        assert is_nvfp4_param(ex.down_proj)
+        assert ex.gate_up_proj.qdata.shape == (E, 2 * I, H // 2)
+        assert ex.down_proj.qdata.shape == (E, H, I // 2)
+
+        ref_gu_qd, _ = _ref_gate_up(experts)
+        ref_dn_qd, _ = _ref_down(experts)
+        assert torch.equal(ex.gate_up_proj.qdata, ref_gu_qd.to(DEV))
+        assert torch.equal(ex.down_proj.qdata, ref_dn_qd.to(DEV))
+    finally:
+        shutil.rmtree(d, ignore_errors=True)
+
+
+def test_source_pattern_alternation_disambiguates_scales():
+    """transformers resolves a checkpoint key via ``re.search`` over the compiled
+    ``(?P<g0>..)|(?P<g1>..)`` alternation and takes the first non-None group. The
+    patterns are not end-anchored for many-to-one converters, so ``...weight`` would
+    substring-match inside ``...weight_scale``/``...weight_scale_2`` unless the more
+    specific suffixes are ordered first. This pins that ordering for both converters."""
+    for conv in nvfp4_experts_weight_converters():
+        rx = conv.compiled_sources
+
+        # which source pattern does each concrete checkpoint key resolve to?
+        def _resolve(key, rx=rx, conv=conv):  # bind loop vars per-iteration
+            m = rx.search(key)
+            assert m is not None, f"{key} matched no source pattern"
+            gname = next(n for n, v in m.groupdict().items() if v is not None)
+            return conv.source_patterns[int(gname[1:])]
+
+        proj = "down_proj" if "down_proj" in conv.target_patterns[0] else "gate_proj"
+        assert _resolve(f"experts.3.{proj}.weight").endswith(f"{proj}.weight")
+        assert _resolve(f"experts.3.{proj}.weight_scale").endswith(
+            f"{proj}.weight_scale"
+        )
+        assert _resolve(f"experts.3.{proj}.weight_scale_2").endswith(
+            f"{proj}.weight_scale_2"
+        )
diff --git a/tests/integrations/kernels/scattermoe_lora/test_parallel_experts_large_batch_repro.py b/tests/integrations/kernels/scattermoe_lora/test_parallel_experts_large_batch_repro.py
new file mode 100644
index 0000000000..37a6a7f6ff
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_parallel_experts_large_batch_repro.py
@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Reproducer for the cuBLAS / illegal-memory-access failure surfaced at
+seq=512K with 16 shards in the tiled-MLP long-context bench.
+
+The originally-reported symptom (``CUBLAS_STATUS_EXECUTION_FAILED`` from
+``cublasGemmStridedBatchedEx`` at
+``parallel_experts.py:72``'s ``gates.unsqueeze(1) @ output_expanded``)
+is a downstream effect — the actual fault is in the upstream
+``scatter2scatter`` Triton kernel's pointer-offset arithmetic. When the
+output of the up-projection has
+``L_scattered * y_dim >= 2 ** 31`` elements (i.e. the kernel's
+``M_block * stride_ym`` int32 multiplication overflows), the kernel
+silently writes to wrong addresses, which can in turn trip the next
+kernel (the gates @ output_expanded bmm or whatever else follows).
+
+The repro shape mirrors the failing bench config:
+
+* shard tokens ``T = 32768`` (= 524288 // 16),
+* ``top_k = 8`` → ``L_scattered = T * top_k = 262144``,
+* ``num_experts = 128``, ``hidden = 2048``, ``intermediate = 8192``.
+
+At that shape, the up-projection's scatter2scatter output is
+``[262144, 2 * 8192] = [262144, 16384]`` = 2**32 elements. The
+overflow boundary for the M_block * stride_ym int32 product is
+``M_block < 2 ** 31 / 16384 = 131072`` — exactly half the output rows.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+# Failing-config constants (mirror tests/integrations/monkeypatch/
+# bench_tiled_mlp_moe.py at seq=524288, shards=16).
+_T = 32768
+_TOP_K = 8
+_NUM_EXPERTS = 128
+_HIDDEN = 2048
+_INTERMEDIATE = 8192
+_DTYPE = torch.bfloat16
+
+
+def _requires_cuda():
+    return pytest.mark.skipif(
+        not torch.cuda.is_available(), reason="CUDA required for the repro"
+    )
+
+
+_SCATTER2SCATTER_INT32_LIMIT = 2**31
+
+
+@_requires_cuda()
+def test_scatter2scatter_below_threshold_no_overhead():
+    """At shapes well below the int32 overflow boundary the auto-dispatch
+    in ``ParallelLinear`` picks ``INT64_INDICES=False`` and the kernel
+    output is bit-identical to a direct call with the same flag.
+
+    This is the regression guard for "don't accidentally penalise the
+    common-case path that does not need int64 indices".
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops import (
+        scatter2scatter,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+        flatten_sort_count,
+    )
+
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+
+    # Small shape: T=512 tokens → L_scattered=4096 → output is
+    # 4096 * 16384 = 67 M elements, well below the 2**31 threshold.
+    T_small = 512
+    x = torch.randn(T_small, _HIDDEN, device=device, dtype=_DTYPE)
+    W = (
+        torch.randn(
+            _NUM_EXPERTS, _HIDDEN, 2 * _INTERMEDIATE, device=device, dtype=_DTYPE
+        )
+        * 0.01
+    )
+
+    logits = torch.randn(T_small, _NUM_EXPERTS, device=device)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), _TOP_K, dim=-1)
+    sei, ssi, _ = flatten_sort_count(top_idx, _NUM_EXPERTS)
+
+    assert sei.size(0) * W.size(-1) < _SCATTER2SCATTER_INT32_LIMIT
+
+    out_i32 = scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=_TOP_K,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=False,
+    )
+    out_i64 = scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=_TOP_K,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(out_i32, out_i64), (
+        "INT64_INDICES must not change MMA/accumulation order at small shapes"
+    )
+
+
+@_requires_cuda()
+def test_scatter2scatter_no_corruption_at_overflow_shape():
+    """The kernel-level int64 fix must keep every output row populated
+    when the shape straddles the 2**31-element boundary.
+
+    Background: with INT64_INDICES=False the Triton ``scatter2scatter``
+    kernel computes pointer offsets as
+    ``Y_ptr + M_block * stride_ym + N_block * stride_yn`` in int32. At
+    the bench shape (L_scattered=262144, y_dim=16384 → 2**32 elements
+    of output) the trailing rows past ``M_block >= 2**31 / y_dim``
+    overflow and their masked stores silently drop, leaving those rows
+    as all-zeros. With INT64_INDICES=True the M_block range is cast to
+    int64 before it enters the multiplication and the overflow is
+    eliminated at the kernel level.
+
+    This test calls the kernel directly with INT64_INDICES=True and
+    asserts every sampled row past the boundary has at least one
+    non-zero element. (The ``ParallelLinear`` wrapper's auto-dispatch
+    is covered separately by ``test_parallel_linear_long_seq_routing_combination``.)
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops import (
+        scatter2scatter,
+    )
+    from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+        flatten_sort_count,
+    )
+
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+
+    x = torch.randn(_T, _HIDDEN, device=device, dtype=_DTYPE)
+    W = (
+        torch.randn(
+            _NUM_EXPERTS, _HIDDEN, 2 * _INTERMEDIATE, device=device, dtype=_DTYPE
+        )
+        * 0.01
+    )
+
+    logits = torch.randn(_T, _NUM_EXPERTS, device=device)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), _TOP_K, dim=-1)
+    sei, ssi, _ = flatten_sort_count(top_idx, _NUM_EXPERTS)
+
+    L_scattered = sei.size(0)
+    y_dim = W.size(-1)
+    assert L_scattered * y_dim >= _SCATTER2SCATTER_INT32_LIMIT, (
+        f"repro precondition: L_scattered * y_dim ({L_scattered * y_dim}) "
+        f"must straddle the int32 overflow boundary "
+        f"({_SCATTER2SCATTER_INT32_LIMIT})"
+    )
+
+    output = scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=_TOP_K,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+
+    overflow_threshold_row = _SCATTER2SCATTER_INT32_LIMIT // y_dim
+    sample_rows = [
+        0,
+        overflow_threshold_row // 2,
+        overflow_threshold_row - 1,
+        overflow_threshold_row,
+        overflow_threshold_row + 1,
+        (overflow_threshold_row + L_scattered) // 2,
+        L_scattered - 1,
+    ]
+    for row in sample_rows:
+        nz = (output[row] != 0).any().item()
+        assert nz, (
+            f"row {row} of scatter2scatter output is all-zero "
+            f"(overflow_threshold_row={overflow_threshold_row}, "
+            f"L_scattered={L_scattered}, y_dim={y_dim})"
+        )
+
+
+@_requires_cuda()
+def test_parallel_linear_long_seq_routing_combination():
+    """End-to-end repro through ``parallel_linear`` matching the bench path.
+
+    Replicates the ``ScatterMoEGatedMLP.forward`` shape sequence (up
+    projection at line 374 → activation → down projection at line 385
+    with ``gates=routing_weights``) at the seq=524288/shards=16 inner
+    config. Before the fix this raises
+    ``CUBLAS_STATUS_EXECUTION_FAILED`` (or a subsequent illegal-memory-
+    access) at the down-projection's ``gates @ output_expanded`` bmm.
+    """
+    import torch.nn.functional as F
+
+    from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+        flatten_sort_count,
+        parallel_linear,
+    )
+
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+
+    layer_input = torch.randn(_T, _HIDDEN, device=device, dtype=_DTYPE)
+    # Match the bench's ScatterMoEGatedMLP weight layout: input_linear
+    # is [E, 2*INTERMEDIATE, HIDDEN] then .transpose(2, 1) →
+    # [E, HIDDEN, 2*INTERMEDIATE]. output_linear is [E, HIDDEN,
+    # INTERMEDIATE] then .transpose(2, 1) → [E, INTERMEDIATE, HIDDEN].
+    in_w = (
+        torch.randn(
+            _NUM_EXPERTS, 2 * _INTERMEDIATE, _HIDDEN, device=device, dtype=_DTYPE
+        )
+        * 0.02
+    )
+    out_w = (
+        torch.randn(_NUM_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device, dtype=_DTYPE)
+        * 0.02
+    )
+
+    router_logits = torch.randn(_T, _NUM_EXPERTS, device=device)
+    routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+    routing_weights, selected_experts = torch.topk(routing_weights, _TOP_K, dim=-1)
+    routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+    routing_weights = routing_weights.to(_DTYPE)
+
+    sei, ssi, eo = flatten_sort_count(selected_experts, _NUM_EXPERTS)
+
+    # Up projection — this is the overflow-prone call (output
+    # numel = T * top_k * 2*INTERMEDIATE = 2**32 at the bench shape).
+    with torch.no_grad():
+        gup = parallel_linear(
+            layer_input,
+            in_w.transpose(2, 1),
+            _TOP_K,
+            sei,
+            ssi,
+            eo,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        gates, h = gup.chunk(2, dim=-1)
+        h = F.silu(gates) * h
+
+        # Down projection — its gates @ output_expanded bmm is where
+        # the reported CUBLAS_STATUS_EXECUTION_FAILED surfaces. The
+        # crash, however, is a downstream symptom of the up-projection
+        # corruption above.
+        layer_output = parallel_linear(
+            h,
+            out_w.transpose(2, 1),
+            1,
+            sei,
+            ssi,
+            eo,
+            grouped_in=True,
+            grouped_out=False,
+            gates=routing_weights,
+        )
+        # Force the (otherwise lazy) CUDA error to surface synchronously.
+        torch.cuda.synchronize()
+
+    assert layer_output.shape == (_T, _HIDDEN)
+    # The output must have real values, not zero rows or NaN/Inf.
+    # ``(.abs().sum(dim=-1) == 0)`` would catch the silent-zero
+    # corruption pattern even when the kernel did not crash hard.
+    assert torch.isfinite(layer_output).all().item(), (
+        "layer_output has non-finite values — likely overflow corruption"
+    )
+    row_sums = layer_output.float().abs().sum(dim=-1)
+    assert (row_sums > 0).all().item(), (
+        "layer_output has all-zero rows — silent overflow corruption "
+        "in the up-projection scatter2scatter"
+    )
diff --git a/tests/integrations/kernels/scattermoe_lora/test_scattermoe_lora_int64_indices.py b/tests/integrations/kernels/scattermoe_lora/test_scattermoe_lora_int64_indices.py
new file mode 100644
index 0000000000..88c77865fd
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_scattermoe_lora_int64_indices.py
@@ -0,0 +1,581 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Parity and overflow-correctness tests for the ``INT64_INDICES``
+``tl.constexpr`` knob added to the scattermoe-lora Triton kernels.
+
+The kernel-level fix promotes the per-launch index ranges to int64 only
+when the wrapper has detected that ``L_scattered * y_dim`` would
+overflow int32. Two properties are tested:
+
+1. **Bitwise parity at small shapes.** When the shape fits in int32,
+   ``INT64_INDICES=False`` (the JIT'd int32 variant) and
+   ``INT64_INDICES=True`` (the int64 variant) compute the same MMA in
+   the same order. Only the index *type* changes, so the outputs must
+   be bitwise identical — any deviation indicates the cast leaked into
+   the accumulator path.
+
+2. **Overflow correctness at large shapes.** At the previously-failing
+   bench config (seq=524288 with 16 shards, L_scattered=262144,
+   y_dim=16384 → 2**32 element output), the int64 kernel must populate
+   every row of the output and match the chunked workaround within bf16
+   tolerance (the chunking workaround changes accumulation order, so
+   bit-equality is not expected against it — only against the same-
+   layout int32 kernel below the overflow boundary).
+
+The bench-config test is gated by GPU memory; an L_scattered=262144
+× 16384 bf16 output is ~8.6 GiB and the up-projection weight is
+~64 GiB so we skip when free memory is below the threshold.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.kernels import (
+    lora_ops,
+    ops as base_ops,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+    flatten_sort_count,
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+# Sufficient condition for int32 pointer arithmetic to overflow in the
+# Triton kernel: any indexed buffer has >= 2**31 elements.
+_INT32_LIMIT = 2**31
+
+
+def _requires_cuda():
+    return pytest.mark.skipif(
+        not torch.cuda.is_available(), reason="CUDA not available"
+    )
+
+
+pytestmark = _requires_cuda()
+
+
+# ─── Helpers ─────────────────────────────────────────────────────────────────
+
+
+def _setup(E, K, N, T, top_k, R=16, seed=42):
+    """Create synthetic inputs + routing for a (E, K, N, T, k) shape."""
+    torch.manual_seed(seed)
+    x = torch.randn(T, K, device=DEVICE, dtype=DTYPE)
+    W = torch.randn(E, K, N, device=DEVICE, dtype=DTYPE) * 0.02
+    lora_A = torch.randn(R * E, K, device=DEVICE, dtype=DTYPE) * 0.01
+    lora_B = torch.randn(N, R * E, device=DEVICE, dtype=DTYPE) * 0.01
+    logits = torch.randn(T, E, device=DEVICE)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
+    sei, ssi, eo = flatten_sort_count(top_idx, E)
+    return x, W, lora_A, lora_B, sei, ssi, eo
+
+
+# ─── Parity tests at non-overflow shapes (bitwise identity) ──────────────────
+
+
+def test_dense_scatter2scatter_int64_parity_small():
+    """Dense scatter2scatter: INT64_INDICES=True == INT64_INDICES=False at small shape."""
+    x, W, *_, sei, ssi, _ = _setup(E=8, K=512, N=1024, T=256, top_k=4)
+    k = 4
+    out_i32 = base_ops.scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=k,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=False,
+    )
+    out_i64 = base_ops.scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=k,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(out_i32, out_i64), (
+        "INT64_INDICES must not change accumulation order at non-overflow shapes"
+    )
+
+
+def test_dense_scatter2scatter_int64_parity_ungrouped_out():
+    """Same parity but y_grouped=False (uses M_idx scatter lookup, not M_block)."""
+    x, W, *_, sei, ssi, _ = _setup(E=8, K=512, N=1024, T=256, top_k=4)
+    k = 4
+    out_i32 = base_ops.scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=k,
+        x_grouped=False,
+        y_grouped=False,
+        int64_indices=False,
+    )
+    out_i64 = base_ops.scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=k,
+        x_grouped=False,
+        y_grouped=False,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(out_i32, out_i64)
+
+
+def test_scatter2scatter_lora_int64_parity_small():
+    """scatter2scatter_lora: int32 vs int64 must agree bitwise."""
+    # Pick a shape that lands on the fused path (not split): few-large-experts
+    # split threshold is E<=32 with K*N >= 20M, so use a small K*N to stay
+    # on the fused kernel.
+    x, W, lA, lB, sei, ssi, _ = _setup(E=64, K=256, N=512, T=128, top_k=4)
+    k = 4
+    scaling = 0.5
+    out_i32 = lora_ops.scatter2scatter_lora(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=k,
+        lora_A=lA,
+        lora_B=lB,
+        scaling=scaling,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=False,
+    )
+    out_i64 = lora_ops.scatter2scatter_lora(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=k,
+        lora_A=lA,
+        lora_B=lB,
+        scaling=scaling,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(out_i32, out_i64)
+
+
+def test_scatter2scatter_lora_dX_int64_parity_small():
+    """scatter2scatter_lora_dX: int32 vs int64 must agree bitwise."""
+    _, W, lA, lB, sei, ssi, _ = _setup(E=64, K=256, N=512, T=128, top_k=4)
+    k = 4
+    scaling = 0.5
+    M_grouped = sei.size(0)  # ungrouped k=1 dy_grouped=True
+    dy = torch.randn(M_grouped, W.size(2), device=DEVICE, dtype=DTYPE) * 0.01
+    dX_i32 = lora_ops.scatter2scatter_lora_dX(
+        DY=dy,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=k,
+        lora_A=lA,
+        lora_B=lB,
+        scaling=scaling,
+        dy_grouped=True,
+        dx_grouped=False,
+        int64_indices=False,
+    )
+    dX_i64 = lora_ops.scatter2scatter_lora_dX(
+        DY=dy,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=k,
+        lora_A=lA,
+        lora_B=lB,
+        scaling=scaling,
+        dy_grouped=True,
+        dx_grouped=False,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(dX_i32, dX_i64)
+
+
+def test_group_bwd_lora_int64_parity_small():
+    """group_bwd_lora (split kernel): int32 vs int64 must agree bitwise."""
+    x, W, lA, lB, sei, ssi, eo = _setup(E=16, K=256, N=512, T=128, top_k=2)
+    grouped_x = base_ops.group(x, ssi, fan_out=2)
+    M = grouped_x.size(0)
+    dy = torch.randn(M, W.size(2), device=DEVICE, dtype=DTYPE) * 0.01
+    scaling = 0.5
+    dA_i32, dB_i32 = lora_ops.group_bwd_lora(
+        DY=dy,
+        X=grouped_x,
+        lora_A=lA,
+        lora_B=lB,
+        expert_offsets=eo,
+        E=16,
+        scaling=scaling,
+        int64_indices=False,
+    )
+    dA_i64, dB_i64 = lora_ops.group_bwd_lora(
+        DY=dy,
+        X=grouped_x,
+        lora_A=lA,
+        lora_B=lB,
+        expert_offsets=eo,
+        E=16,
+        scaling=scaling,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(dA_i32, dA_i64)
+    assert torch.equal(dB_i32, dB_i64)
+
+
+def test_group_bwd_lora_fused_int64_parity_small():
+    """group_bwd_lora_fused: int32 vs int64 must agree within bf16 tolerance.
+
+    Unlike the split kernel, the fused kernel writes dA/dB via
+    ``tl.atomic_add``; the order in which (E, K-tile, N-tile) thread blocks
+    land their atomics is non-deterministic, so bit-equality is not
+    achievable even between two launches of the *same* kernel variant.
+    The INT64 cast only changes index *types*, not the MMA path or the
+    atomic reduction, so the two variants must still match within
+    ``torch.allclose`` bf16 tolerance.
+    """
+    x, _W, lA, lB, _sei, ssi, eo = _setup(E=16, K=256, N=512, T=128, top_k=2)
+    k = 2
+    M_total = ssi.size(0)
+    N = lB.size(0)
+    dy = torch.randn(M_total, N, device=DEVICE, dtype=DTYPE) * 0.01
+    scaling = 0.5
+    dA_i32, dB_i32 = lora_ops.group_bwd_lora_fused(
+        DY=dy,
+        X=x,
+        lora_A=lA,
+        lora_B=lB,
+        expert_offsets=eo,
+        sorted_scattered_idxs=ssi,
+        E=16,
+        k=k,
+        scaling=scaling,
+        dy_grouped=False,
+        int64_indices=False,
+    )
+    dA_i64, dB_i64 = lora_ops.group_bwd_lora_fused(
+        DY=dy,
+        X=x,
+        lora_A=lA,
+        lora_B=lB,
+        expert_offsets=eo,
+        sorted_scattered_idxs=ssi,
+        E=16,
+        k=k,
+        scaling=scaling,
+        dy_grouped=False,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+    # Tolerance: a few bf16 ULPs is expected from atomic-add ordering nondet.
+    assert torch.allclose(dA_i32, dA_i64, rtol=1e-2, atol=5e-4), (
+        f"max_abs_diff dA: {(dA_i32.float() - dA_i64.float()).abs().max()}"
+    )
+    assert torch.allclose(dB_i32, dB_i64, rtol=1e-2, atol=5e-4), (
+        f"max_abs_diff dB: {(dB_i32.float() - dB_i64.float()).abs().max()}"
+    )
+
+
+# ─── Overflow correctness at the bench shape ─────────────────────────────────
+
+
+# Bench-shape constants (mirror ``test_parallel_experts_large_batch_repro.py``).
+_T = 32768
+_TOP_K = 8
+_NUM_EXPERTS = 128
+_HIDDEN = 2048
+_INTERMEDIATE = 8192
+
+
+def _has_free_gpu_mem(min_gb: float) -> bool:
+    if not torch.cuda.is_available():
+        return False
+    free, _total = torch.cuda.mem_get_info()
+    return free / (1024**3) >= min_gb
+
+
+@pytest.mark.skipif(
+    not _has_free_gpu_mem(80.0),
+    reason="bench shape needs ~80 GiB free GPU memory",
+)
+def test_dense_scatter2scatter_int64_at_overflow_shape():
+    """Direct int64 kernel call at the bench shape produces no zero rows.
+
+    With ``INT64_INDICES=True`` the kernel's pointer arithmetic stays in
+    int64 across the full output row range, so rows past the would-be
+    int32 overflow boundary (``M_block >= 2**31 / y_dim``) are populated
+    rather than silently dropped.
+    """
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+    x = torch.randn(_T, _HIDDEN, device=device, dtype=DTYPE)
+    W = (
+        torch.randn(
+            _NUM_EXPERTS, _HIDDEN, 2 * _INTERMEDIATE, device=device, dtype=DTYPE
+        )
+        * 0.01
+    )
+
+    logits = torch.randn(_T, _NUM_EXPERTS, device=device)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), _TOP_K, dim=-1)
+    sei, ssi, _ = flatten_sort_count(top_idx, _NUM_EXPERTS)
+
+    L_scattered = sei.size(0)
+    y_dim = W.size(-1)
+    assert L_scattered * y_dim >= _INT32_LIMIT, (
+        "precondition: shape must straddle the int32 overflow boundary"
+    )
+
+    out_i64 = base_ops.scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=_TOP_K,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+
+    # Sample rows on both sides of the int32 overflow boundary.
+    overflow_threshold_row = _INT32_LIMIT // y_dim
+    sample_rows = [
+        0,
+        overflow_threshold_row - 1,
+        overflow_threshold_row,
+        overflow_threshold_row + 1,
+        L_scattered - 1,
+    ]
+    for row in sample_rows:
+        assert (out_i64[row] != 0).any().item(), (
+            f"int64 kernel left row {row} all-zero (overflow boundary "
+            f"= {overflow_threshold_row}, L_scattered = {L_scattered})"
+        )
+    assert torch.isfinite(out_i64).all().item(), (
+        "int64 kernel produced non-finite values at overflow shape"
+    )
+
+
+@pytest.mark.skipif(
+    not _has_free_gpu_mem(80.0),
+    reason="bench shape needs ~80 GiB free GPU memory",
+)
+def test_parallel_linear_overflow_takes_int64_kernel_path(monkeypatch):
+    """``ParallelLinear.forward`` at the bench shape must route through
+    the int64 kernel path (single launch, ``int64_indices=True``).
+
+    The auto-dispatch should set ``needs_int64=True`` and dispatch a
+    single ``scatter2scatter`` launch with that flag. A regressed path
+    that called the kernel multiple times (e.g. a chunking workaround)
+    would invoke ``scatter2scatter_compileable`` more than once and
+    fail this assertion.
+    """
+    from axolotl.integrations.kernels.libs.scattermoe_lora import parallel_experts
+    from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+        parallel_linear,
+    )
+
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+
+    x = torch.randn(_T, _HIDDEN, device=device, dtype=DTYPE)
+    W = (
+        torch.randn(
+            _NUM_EXPERTS, _HIDDEN, 2 * _INTERMEDIATE, device=device, dtype=DTYPE
+        )
+        * 0.01
+    )
+
+    logits = torch.randn(_T, _NUM_EXPERTS, device=device)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), _TOP_K, dim=-1)
+    sei, ssi, eo = flatten_sort_count(top_idx, _NUM_EXPERTS)
+
+    # Spy on the kernel launches. The kernel-level int64 fix dispatches
+    # exactly one ``scatter2scatter_compileable`` call with
+    # ``int64_indices=True``. A re-introduced chunking workaround would
+    # invoke it once per chunk (>=2 at this shape).
+    launches = []
+    real_compileable = parallel_experts.kernels.ops.scatter2scatter_compileable
+
+    def _spy_compileable(*args, **kwargs):
+        # int64_indices is positional arg 9 (after b, x_grouped, y_grouped).
+        launches.append(
+            {
+                "args_len": len(args),
+                "int64": args[9]
+                if len(args) > 9
+                else kwargs.get("int64_indices", False),
+            }
+        )
+        return real_compileable(*args, **kwargs)
+
+    monkeypatch.setattr(
+        parallel_experts.kernels.ops,
+        "scatter2scatter_compileable",
+        _spy_compileable,
+    )
+
+    with torch.no_grad():
+        out = parallel_linear(
+            x,
+            W,
+            _TOP_K,
+            sei,
+            ssi,
+            eo,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        torch.cuda.synchronize()
+
+    assert len(launches) == 1, (
+        f"expected exactly one kernel launch (direct int64 path), got {len(launches)}"
+    )
+    assert launches[0]["int64"] is True, (
+        "auto-dispatch should have set int64_indices=True at the overflow shape"
+    )
+    assert out.shape == (_T * _TOP_K, 2 * _INTERMEDIATE)
+    assert torch.isfinite(out).all().item()
+
+
+# ─── Smaller-shape overflow (runs on L40S / 24 GiB GPUs) ─────────────────────
+# L_scattered * y_dim = 2**32 (2× past 2**31); peak VRAM ≈ 8 GiB.
+_SMALL_T = 131072
+_SMALL_TOP_K = 8
+_SMALL_E = 8
+_SMALL_K = 256
+_SMALL_INTERMEDIATE = 2048
+_SMALL_MIN_FREE_GIB = 12.0
+
+
+@pytest.mark.skipif(
+    not _has_free_gpu_mem(_SMALL_MIN_FREE_GIB),
+    reason=f"small overflow shape needs ~{_SMALL_MIN_FREE_GIB:.0f} GiB free GPU memory",
+)
+def test_dense_scatter2scatter_int64_at_overflow_shape_small():
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+    y_dim = 2 * _SMALL_INTERMEDIATE
+    x = torch.randn(_SMALL_T, _SMALL_K, device=device, dtype=DTYPE)
+    W = torch.randn(_SMALL_E, _SMALL_K, y_dim, device=device, dtype=DTYPE) * 0.01
+
+    logits = torch.randn(_SMALL_T, _SMALL_E, device=device)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), _SMALL_TOP_K, dim=-1)
+    sei, ssi, _ = flatten_sort_count(top_idx, _SMALL_E)
+
+    L_scattered = sei.size(0)
+    assert L_scattered * y_dim >= _INT32_LIMIT, (
+        f"precondition: L_scattered * y_dim ({L_scattered * y_dim}) must "
+        f"straddle the int32 overflow boundary ({_INT32_LIMIT})"
+    )
+
+    out_i64 = base_ops.scatter2scatter(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=_SMALL_TOP_K,
+        x_grouped=False,
+        y_grouped=True,
+        int64_indices=True,
+    )
+    torch.cuda.synchronize()
+
+    overflow_threshold_row = _INT32_LIMIT // y_dim
+    sample_rows = [
+        0,
+        overflow_threshold_row - 1,
+        overflow_threshold_row,
+        overflow_threshold_row + 1,
+        L_scattered - 1,
+    ]
+    for row in sample_rows:
+        assert (out_i64[row] != 0).any().item(), (
+            f"int64 kernel left row {row} all-zero (overflow boundary "
+            f"= {overflow_threshold_row}, L_scattered = {L_scattered})"
+        )
+    assert torch.isfinite(out_i64).all().item()
+
+
+@pytest.mark.skipif(
+    not _has_free_gpu_mem(_SMALL_MIN_FREE_GIB),
+    reason=f"small overflow shape needs ~{_SMALL_MIN_FREE_GIB:.0f} GiB free GPU memory",
+)
+def test_parallel_linear_overflow_takes_int64_kernel_path_small(monkeypatch):
+    from axolotl.integrations.kernels.libs.scattermoe_lora import parallel_experts
+    from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+        parallel_linear,
+    )
+
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+    y_dim = 2 * _SMALL_INTERMEDIATE
+    x = torch.randn(_SMALL_T, _SMALL_K, device=device, dtype=DTYPE)
+    W = torch.randn(_SMALL_E, _SMALL_K, y_dim, device=device, dtype=DTYPE) * 0.01
+
+    logits = torch.randn(_SMALL_T, _SMALL_E, device=device)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), _SMALL_TOP_K, dim=-1)
+    sei, ssi, eo = flatten_sort_count(top_idx, _SMALL_E)
+
+    launches = []
+    real_compileable = parallel_experts.kernels.ops.scatter2scatter_compileable
+
+    def _spy_compileable(*args, **kwargs):
+        launches.append(
+            {
+                "int64": args[9]
+                if len(args) > 9
+                else kwargs.get("int64_indices", False),
+            }
+        )
+        return real_compileable(*args, **kwargs)
+
+    monkeypatch.setattr(
+        parallel_experts.kernels.ops,
+        "scatter2scatter_compileable",
+        _spy_compileable,
+    )
+
+    with torch.no_grad():
+        out = parallel_linear(
+            x,
+            W,
+            _SMALL_TOP_K,
+            sei,
+            ssi,
+            eo,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        torch.cuda.synchronize()
+
+    assert len(launches) == 1, (
+        f"expected exactly one kernel launch (direct int64 path), got {len(launches)}"
+    )
+    assert launches[0]["int64"] is True, (
+        "auto-dispatch should have set int64_indices=True at the overflow shape"
+    )
+    assert out.shape == (_SMALL_T * _SMALL_TOP_K, y_dim)
+    assert torch.isfinite(out).all().item()
diff --git a/tests/integrations/kernels/scattermoe_lora/test_scattermoe_lora_m_bucket.py b/tests/integrations/kernels/scattermoe_lora/test_scattermoe_lora_m_bucket.py
new file mode 100644
index 0000000000..5d1c4d6b19
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_scattermoe_lora_m_bucket.py
@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Tests for the ``M_BUCKET`` autotune-key bucketing on the scattermoe-lora
+fused forward kernel.
+
+The kernel runs on the real ``M`` (loop bounds + masks); only the
+``@triton.autotune`` cache key is bucketed via :func:`_bucket_m`. These tests
+pin both halves of that contract:
+
+  * ``_bucket_m`` rounds up to a multiple of the granularity (pure-Python
+    unit test, no GPU).
+  * Two distinct real ``M`` values that share a bucket produce **one**
+    cache entry (the whole point — no resweep on small seqlen variation).
+  * Two real ``M`` values in different buckets produce **two** cache
+    entries (we didn't accidentally collapse to a single key).
+
+Run on CUDA only; the bucketing assertion needs an actual Triton launch.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.kernels import lora_ops
+from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.lora_ops import (
+    _M_BUCKET_GRANULARITY,
+    _bucket_m,
+    scatter2scatter_lora,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+    flatten_sort_count,
+)
+
+
+def test_bucket_m_rounds_up_to_granularity():
+    g = _M_BUCKET_GRANULARITY
+    assert _bucket_m(1) == g
+    assert _bucket_m(g) == g
+    assert _bucket_m(g + 1) == 2 * g
+    assert _bucket_m(2 * g) == 2 * g
+    # Realistic seqlen variation: at granularity=1024 and top_k=8 the three
+    # seqlens 16300/16400/16500 straddle one bucket boundary, so they collapse
+    # to 2 cache entries rather than 3 (16400 and 16500 share a bucket).
+    assert _bucket_m(16400 * 8) == _bucket_m(16500 * 8)
+    assert _bucket_m(16300 * 8) != _bucket_m(16400 * 8)
+    distinct = {_bucket_m(s * 8) for s in (16300, 16400, 16500)}
+    assert len(distinct) == 2
+
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA required for kernel launch"
+)
+
+
+_DEVICE = "cuda"
+_DTYPE = torch.bfloat16
+_E = 4
+_K = 64
+_N = 64
+_TOP_K = 2
+_R = 16
+
+
+def _launch_once(m: int) -> None:
+    """One fused fwd launch at the given real M; minimal shapes for speed."""
+    torch.manual_seed(m)
+    x = torch.randn(m, _K, device=_DEVICE, dtype=_DTYPE)
+    W = torch.randn(_E, _K, _N, device=_DEVICE, dtype=_DTYPE) * 0.02
+    lora_A = torch.randn(_R * _E, _K, device=_DEVICE, dtype=_DTYPE) * 0.01
+    lora_B = torch.randn(_N, _R * _E, device=_DEVICE, dtype=_DTYPE) * 0.01
+    logits = torch.randn(m, _E, device=_DEVICE)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), _TOP_K, dim=-1)
+    sei, ssi, _ = flatten_sort_count(top_idx, _E)
+    scatter2scatter_lora(
+        X=x,
+        W=W,
+        sorted_expert_idxs=sei,
+        sorted_scattered_idxs=ssi,
+        k=_TOP_K,
+        lora_A=lora_A,
+        lora_B=lora_B,
+        scaling=1.0 / _R,
+    )
+
+
+def test_autotune_cache_collapses_within_bucket_and_grows_across_buckets():
+    cache = lora_ops._scatter2scatter_lora.cache
+    cache.clear()
+
+    g = _M_BUCKET_GRANULARITY
+    # Two M values that both ceil to bucket B1.
+    m_a = g - 1
+    m_b = g // 2 + 1
+    assert _bucket_m(m_a) == g
+    assert _bucket_m(m_b) == g
+
+    _launch_once(m_a)
+    assert len(cache) == 1, (
+        f"first launch should create exactly one cache entry, got {len(cache)}"
+    )
+
+    _launch_once(m_b)
+    assert len(cache) == 1, (
+        f"second launch in the same bucket must not add a cache entry "
+        f"(M={m_a} and M={m_b} both bucket to {g}); got {len(cache)} entries"
+    )
+
+    # An M strictly past the bucket boundary lands in bucket 2*g.
+    m_c = g + 1
+    assert _bucket_m(m_c) == 2 * g
+    _launch_once(m_c)
+    assert len(cache) == 2, (
+        f"launch in a different bucket (M={m_c} -> {2 * g}) must add a "
+        f"second cache entry; got {len(cache)}"
+    )
diff --git a/tests/integrations/kernels/scattermoe_lora/test_shared_dequant_helper.py b/tests/integrations/kernels/scattermoe_lora/test_shared_dequant_helper.py
new file mode 100644
index 0000000000..cf5d88c669
--- /dev/null
+++ b/tests/integrations/kernels/scattermoe_lora/test_shared_dequant_helper.py
@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Parity tests for :func:`shared_dequant_across_shards`.
+
+The helper hoists the per-shard MXFP4 dequant out of the orthogonal
+Strategy A path so that overlapping active-expert sets across shards
+dequantize the union once instead of N times. The optimization is
+only valid if a shard's slice of the union buffer is *byte-identical*
+to what the per-shard ``selective_expert_weights`` call would have
+produced.
+"""
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import (
+    get_active_experts,
+    selective_expert_weights,
+    shared_dequant_across_shards,
+)
+
+torchao = pytest.importorskip("torchao")
+from torchao.prototype.mx_formats.mx_tensor import MXTensor  # noqa: E402
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA required for MX kernels"
+)
+
+
+class _MockExperts:
+    """Same bare wrapper used by ``test_mxfp4_expert_weights.py``."""
+
+    def __init__(self, mx_param, num_experts):
+        self.gate_up_proj = mx_param
+        self.num_experts = num_experts
+
+
+def _make_mxfp4(E, N, K, seed):
+    torch.manual_seed(seed)
+    W = torch.randn(E, N, K, device=DEVICE, dtype=DTYPE)
+    return MXTensor.to_mx(W, elem_dtype=torch.float4_e2m1fn_x2, block_size=32)
+
+
+def _make_overlapping_shard_sei(E, num_shards, seed):
+    """Build N shards of ``sorted_expert_idxs`` with deliberate overlap.
+
+    Each shard picks ~E/2 experts at random; with 4 shards the union is
+    typically ~E (full coverage) while the intersection is non-empty,
+    which is the regime the helper is intended to optimise.
+    """
+    torch.manual_seed(seed)
+    sei = []
+    for _ in range(num_shards):
+        # Each shard sees ~E/2 distinct experts repeated a few times to
+        # mimic top-k routing.
+        chosen = torch.randperm(E, device=DEVICE)[: max(2, E // 2)]
+        # Repeat each id ~3x and sort so the tensor satisfies the
+        # ``sorted_expert_idxs`` contract.
+        repeated = chosen.repeat_interleave(3)
+        sei.append(torch.sort(repeated).values)
+    return sei
+
+
+def test_shared_dequant_matches_per_shard_bitwise():
+    """Union dequant + index gather == per-shard selective dequant, bitwise.
+
+    Uses N=4 shards over E=16 experts with deliberate overlap (each shard
+    picks ~half the experts, union covers most of E). The helper must
+    yield the exact same compact buffer that a per-shard
+    ``selective_expert_weights`` call would have produced for that shard.
+    """
+    E, N, K = 16, 128, 256
+    num_shards = 4
+    mx = _make_mxfp4(E, N, K, seed=13)
+    experts = _MockExperts(mx, E)
+
+    sei_per_shard = _make_overlapping_shard_sei(E, num_shards, seed=13)
+    # Sanity: at least one pair of shards must overlap for this test to
+    # actually exercise the dedup path.
+    actives = [get_active_experts(sei, E) for sei in sei_per_shard]
+    union = torch.unique(torch.cat(actives))
+    sum_per_shard = sum(a.numel() for a in actives)
+    assert sum_per_shard > union.numel(), (
+        "shard active sets must overlap to exercise the shared-dequant path"
+    )
+
+    union_active, union_buf, shard_into_union = shared_dequant_across_shards(
+        experts, "gate_up_proj", sei_per_shard, E
+    )
+
+    assert torch.equal(union_active, union)
+    assert union_buf.shape == (union.numel(), N, K)
+
+    for i, sei in enumerate(sei_per_shard):
+        active_i = get_active_experts(sei, E)
+        reference = selective_expert_weights(experts, "gate_up_proj", active_i)
+        shared_slice = union_buf.index_select(0, shard_into_union[i])
+        assert torch.equal(shared_slice, reference), (
+            f"shard {i}: max abs diff = {(shared_slice - reference).abs().max().item()}"
+        )
+
+
+def test_shared_dequant_disjoint_shards():
+    """When shards do NOT overlap, the helper still produces the right
+    union and the per-shard slices remain bitwise identical."""
+    E, N, K = 12, 64, 128
+    mx = _make_mxfp4(E, N, K, seed=21)
+    experts = _MockExperts(mx, E)
+
+    # Two shards splitting the expert ids into disjoint halves.
+    halves = torch.arange(E, device=DEVICE).chunk(2)
+    sei_per_shard = [torch.sort(h.repeat_interleave(2)).values for h in halves]
+
+    union_active, union_buf, shard_into_union = shared_dequant_across_shards(
+        experts, "gate_up_proj", sei_per_shard, E
+    )
+    assert union_active.numel() == E
+
+    for i, sei in enumerate(sei_per_shard):
+        active_i = get_active_experts(sei, E)
+        reference = selective_expert_weights(experts, "gate_up_proj", active_i)
+        shared_slice = union_buf.index_select(0, shard_into_union[i])
+        assert torch.equal(shared_slice, reference)
+
+
+def test_shared_dequant_single_shard_noop():
+    """N=1 should reduce to the per-shard path: union == only active set."""
+    E, N, K = 8, 32, 64
+    mx = _make_mxfp4(E, N, K, seed=5)
+    experts = _MockExperts(mx, E)
+
+    sei = torch.tensor([0, 0, 2, 2, 5, 5], device=DEVICE, dtype=torch.long)
+    union_active, union_buf, shard_into_union = shared_dequant_across_shards(
+        experts, "gate_up_proj", [sei], E
+    )
+    active = get_active_experts(sei, E)
+    reference = selective_expert_weights(experts, "gate_up_proj", active)
+    assert torch.equal(union_active, active)
+    assert torch.equal(union_buf, reference)
+    assert torch.equal(union_buf.index_select(0, shard_into_union[0]), reference)
diff --git a/tests/integrations/kernels/test_grouped_backend.py b/tests/integrations/kernels/test_grouped_backend.py
new file mode 100644
index 0000000000..f3dfeac5e8
--- /dev/null
+++ b/tests/integrations/kernels/test_grouped_backend.py
@@ -0,0 +1,98 @@
+"""CPU-only tests for the grouped NVFP4 base-GEMM backend override resolution."""
+
+import pytest
+
+from axolotl.integrations.kernels.libs.scattermoe_lora import (
+    dequant_grouped as dgq,
+    grouped_train as gt,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset_override():
+    gt.set_grouped_backend_override(None)
+    yield
+    gt.set_grouped_backend_override(None)
+
+
+def test_dequant_forces_fallback():
+    gt.set_grouped_backend_override("dequant")
+    assert gt._train_backend("nvfp4") is None  # forced chunked-dequant fallback
+
+
+def test_unavailable_override_falls_back_to_auto():
+    # On CPU no fused backend is available, so a forced 'marlin' warns and auto-selects -> None.
+    gt.set_grouped_backend_override("marlin")
+    assert gt._train_backend("nvfp4") == gt._auto_backend()
+
+
+def test_non_nvfp4_mode_returns_none():
+    gt.set_grouped_backend_override("marlin")
+    assert gt._train_backend("fp8") is None
+
+
+def test_override_setter_normalizes():
+    gt.set_grouped_backend_override("MARLIN")
+    assert gt.RUNTIME.grouped_backend == "marlin"
+    gt.set_grouped_backend_override(None)
+    assert gt.RUNTIME.grouped_backend is None
+
+
+@pytest.mark.parametrize(
+    "major,expected_order",
+    [
+        (12, ("marlin", "cutlass", "deepgemm")),  # consumer Blackwell sm120
+        (10, ("deepgemm", "marlin")),  # datacenter Blackwell sm100
+        (9, ("deepgemm", "marlin")),  # Hopper sm90
+        (8, ("marlin",)),  # Ampere/Ada sm80/sm89 — Marlin is the only fused W4A16 path
+    ],
+)
+def test_auto_backend_arch_order(monkeypatch, major, expected_order):
+    """The auto-select probe order is arch-aware: each GPU class gets its tuned default first."""
+    probed = []
+
+    def fake_available(name):
+        probed.append(name)
+        return False  # force it to probe the whole order
+
+    monkeypatch.setattr(gt, "_backend_available", fake_available)
+    monkeypatch.setattr(gt.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(gt.torch.cuda, "get_device_capability", lambda: (major, 0))
+    assert gt._auto_backend() is None
+    assert tuple(probed) == expected_order
+
+
+class _FakeDG:
+    m_grouped_fp8_fp4_gemm_nt_contiguous = object()
+
+
+@pytest.mark.parametrize(
+    "major,expected",
+    [
+        (8, False),  # Ampere/Ada -> Marlin
+        (
+            9,
+            False,
+        ),  # Hopper sm90: fp8xfp4 kernel asserts arch_major==10 -> must fall back to Marlin
+        (
+            10,
+            True,
+        ),  # datacenter Blackwell sm100 (B200): the only arch the fp4 kernel runs on
+        (12, False),  # consumer Blackwell sm120 -> CUTLASS/Marlin, not DeepGEMM
+    ],
+)
+def test_deepgemm_available_is_sm100_only(monkeypatch, major, expected):
+    """Regression: deepgemm_grouped_available() must gate to sm100. The kernel symbol resolves on
+    sm90 too, but calling it asserts arch_major==10 mid-forward; gating up front routes sm90 to
+    Marlin instead of crashing."""
+    monkeypatch.setattr(dgq.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(dgq.torch.cuda, "get_device_capability", lambda: (major, 0))
+    monkeypatch.setattr(
+        dgq, "_dg", lambda: _FakeDG()
+    )  # avoid needing a real kernel/GPU on sm100
+    assert dgq.deepgemm_grouped_available() is expected
+
+
+def test_deepgemm_unavailable_without_cuda(monkeypatch):
+    monkeypatch.setattr(dgq.torch.cuda, "is_available", lambda: False)
+    assert dgq.deepgemm_grouped_available() is False
diff --git a/tests/integrations/kernels/test_kernels_adapters.py b/tests/integrations/kernels/test_kernels_adapters.py
new file mode 100644
index 0000000000..5bee7416bc
--- /dev/null
+++ b/tests/integrations/kernels/test_kernels_adapters.py
@@ -0,0 +1,112 @@
+"""CPU-only tests for adapter matching/dispatch and non-expert quantization resolution."""
+
+from unittest.mock import MagicMock
+
+import axolotl.integrations.kernels.adapters as adapters_mod
+import axolotl.integrations.kernels.adapters.gemma4 as gemma4_mod
+import axolotl.integrations.kernels.plugin as plugin_mod
+from axolotl.integrations.kernels.adapters import get_active_adapters
+from axolotl.integrations.kernels.adapters.dsv4 import DSV4Adapter
+from axolotl.integrations.kernels.adapters.gemma4 import (
+    Gemma4Adapter,
+    resolve_nonexpert_quantization,
+)
+from axolotl.integrations.kernels.adapters.qwen3_moe import Qwen3MoeAdapter
+from axolotl.integrations.kernels.plugin import KernelsPlugin
+
+
+class Cfg(dict):
+    """Minimal cfg supporting both attribute and .get access (like DictDefault)."""
+
+    def __getattr__(self, k):
+        return self.get(k)
+
+
+def test_dsv4_adapter_matches_only_dsv4():
+    active = get_active_adapters(Cfg(use_dsv4_kernels=True))
+    assert [type(a) for a in active] == [DSV4Adapter]
+    assert get_active_adapters(Cfg(use_dsv4_kernels=False)) == []
+
+
+def test_gemma4_adapter_matches(monkeypatch):
+    monkeypatch.setattr(gemma4_mod, "is_gemma4_nvfp4_modelopt", lambda cfg: True)
+    active = get_active_adapters(Cfg(use_scattermoe=True))
+    assert [type(a) for a in active] == [Gemma4Adapter]
+    # requires scattermoe
+    assert get_active_adapters(Cfg(use_scattermoe=False)) == []
+
+
+def test_qwen3_moe_adapter_matches(monkeypatch):
+    import axolotl.integrations.kernels.adapters.qwen3_moe as qwen3_moe_mod
+
+    monkeypatch.setattr(qwen3_moe_mod, "is_qwen3_moe_nvfp4_modelopt", lambda cfg: True)
+    # either expert backend activates it
+    active = get_active_adapters(Cfg(use_scattermoe=True))
+    assert Qwen3MoeAdapter in [type(a) for a in active]
+    active = get_active_adapters(Cfg(use_sonicmoe=True))
+    assert [type(a) for a in active] == [Qwen3MoeAdapter]
+    # requires an expert backend
+    assert get_active_adapters(Cfg()) == []
+
+
+def test_gemma4_match_failure_is_swallowed(monkeypatch):
+    def boom(cfg):
+        raise RuntimeError("network down")
+
+    monkeypatch.setattr(gemma4_mod, "is_gemma4_nvfp4_modelopt", boom)
+    # a raising matcher must not break adapter discovery
+    assert get_active_adapters(Cfg(use_scattermoe=True)) == []
+
+
+def test_plugin_dispatches_to_active_adapters(monkeypatch):
+    a1, a2 = MagicMock(), MagicMock()
+    monkeypatch.setattr(plugin_mod, "get_active_adapters", lambda cfg: [a1, a2])
+    plugin = KernelsPlugin()
+    cfg = Cfg(use_scattermoe=False, use_sonicmoe=False)  # skip generic registration
+    model = object()
+
+    plugin.pre_model_load(cfg)
+    plugin.pre_lora_load(cfg, model)
+    plugin.post_model_load(cfg, model)
+
+    for a in (a1, a2):
+        a.pre_model_load.assert_called_once_with(cfg)
+        a.pre_lora_load.assert_called_once_with(cfg, model)
+        a.post_model_load.assert_called_once_with(cfg, model)
+
+
+def test_plugin_caches_adapter_discovery(monkeypatch):
+    calls = {"n": 0}
+
+    def counting(cfg):
+        calls["n"] += 1
+        return []
+
+    monkeypatch.setattr(plugin_mod, "get_active_adapters", counting)
+    plugin = KernelsPlugin()
+    cfg = Cfg(use_scattermoe=False, use_sonicmoe=False)
+    plugin.pre_model_load(cfg)
+    plugin.pre_lora_load(cfg, object())
+    plugin.post_model_load(cfg, object())
+    assert calls["n"] == 1  # discovered once, cached across hooks
+
+
+def test_resolve_nonexpert_quantization_intent_and_legacy():
+    # intent field takes precedence
+    assert (
+        resolve_nonexpert_quantization(Cfg(nonexpert_quantization="fp8_blockwise"))
+        == "fp8"
+    )
+    assert resolve_nonexpert_quantization(Cfg(nonexpert_quantization="nf4")) == "nf4"
+    assert resolve_nonexpert_quantization(Cfg(nonexpert_quantization="none")) is None
+    assert resolve_nonexpert_quantization(Cfg(nonexpert_quantization="bf16")) is None
+    # legacy flags
+    assert resolve_nonexpert_quantization(Cfg(gemma4_fp8_nonexpert=True)) == "fp8"
+    assert resolve_nonexpert_quantization(Cfg(gemma4_nf4_nonexpert=True)) == "nf4"
+    # nothing set
+    assert resolve_nonexpert_quantization(Cfg()) is None
+
+
+def test_adapters_module_get_active_is_used(monkeypatch):
+    # sanity: the symbol the plugin caches resolves to the adapters module function
+    assert plugin_mod.get_active_adapters is adapters_mod.get_active_adapters
diff --git a/tests/integrations/kernels/test_kernels_config.py b/tests/integrations/kernels/test_kernels_config.py
new file mode 100644
index 0000000000..3dca3855b5
--- /dev/null
+++ b/tests/integrations/kernels/test_kernels_config.py
@@ -0,0 +1,208 @@
+"""CPU-only config-validation tests for KernelsArgs (grouped-mode + lora_mlp_kernel translation)."""
+
+import pydantic
+import pytest
+
+from axolotl.integrations.kernels.args import KernelsArgs
+
+
+def test_grouped_mode_nvfp4_accepted():
+    a = KernelsArgs.model_validate(
+        {"use_scattermoe": True, "dsv4_fp4_grouped_mode": "nvfp4"}
+    )
+    assert a.dsv4_fp4_grouped_mode == "nvfp4"
+
+
+def test_grouped_mode_fp8_rejected():
+    # 'fp8' is documented historically but unimplemented for training -> reject, don't no-op.
+    with pytest.raises(pydantic.ValidationError, match="not implemented"):
+        KernelsArgs.model_validate({"dsv4_fp4_grouped_mode": "fp8"})
+
+
+def test_grouped_mode_unknown_rejected():
+    with pytest.raises(pydantic.ValidationError):
+        KernelsArgs.model_validate({"dsv4_fp4_grouped_mode": "int4"})
+
+
+def test_lora_mlp_kernel_translated_for_dsv4():
+    # On a DSV4 MoE run, lora_mlp_kernel intent is preserved as dsv4_shared_mlp_lora_kernel
+    # before the generic lora_mlp_kernel is force-disabled.
+    a = KernelsArgs.model_validate(
+        {"use_scattermoe": True, "use_dsv4_kernels": True, "lora_mlp_kernel": True}
+    )
+    assert a.dsv4_shared_mlp_lora_kernel is True
+
+
+def test_lora_mlp_kernel_not_translated_for_non_dsv4():
+    # Non-DSV4 MoE run: lora_mlp_kernel is just disabled, no shared-MLP flag.
+    a = KernelsArgs.model_validate({"use_scattermoe": True, "lora_mlp_kernel": True})
+    assert a.dsv4_shared_mlp_lora_kernel is None
+
+
+def test_lora_mlp_kernel_explicit_shared_flag_preserved():
+    # An explicit dsv4_shared_mlp_lora_kernel is not overwritten by the translation.
+    a = KernelsArgs.model_validate(
+        {
+            "use_scattermoe": True,
+            "use_dsv4_kernels": True,
+            "lora_mlp_kernel": True,
+            "dsv4_shared_mlp_lora_kernel": False,
+        }
+    )
+    assert a.dsv4_shared_mlp_lora_kernel is False
+
+
+def test_fp8_nonexpert_mode_validation():
+    assert (
+        KernelsArgs.model_validate(
+            {"dsv4_fp8_nonexpert_mode": "bf16"}
+        ).dsv4_fp8_nonexpert_mode
+        == "bf16"
+    )
+    with pytest.raises(pydantic.ValidationError):
+        KernelsArgs.model_validate({"dsv4_fp8_nonexpert_mode": "int8"})
+
+
+def test_scattermoe_sonicmoe_mutually_exclusive():
+    with pytest.raises(pydantic.ValidationError):
+        KernelsArgs.model_validate({"use_scattermoe": True, "use_sonicmoe": True})
+
+
+# --- intent-based config surface (#5) -------------------------------------------------------
+def test_expert_backend_alias_scattermoe():
+    a = KernelsArgs.model_validate({"expert_backend": "scattermoe"})
+    assert a.use_scattermoe is True
+
+
+def test_expert_backend_alias_sonicmoe():
+    a = KernelsArgs.model_validate({"expert_backend": "sonicmoe"})
+    assert a.use_sonicmoe is True
+
+
+def test_expert_backend_eager_leaves_flags_unset():
+    a = KernelsArgs.model_validate({"expert_backend": "eager"})
+    assert a.use_scattermoe is None and a.use_sonicmoe is None
+
+
+def test_expert_backend_invalid_rejected():
+    with pytest.raises(pydantic.ValidationError):
+        KernelsArgs.model_validate({"expert_backend": "megablocks"})
+
+
+def test_nonexpert_quantization_valid_values():
+    for v in ("none", "bf16", "fp8_blockwise", "nf4"):
+        assert (
+            KernelsArgs.model_validate(
+                {"nonexpert_quantization": v}
+            ).nonexpert_quantization
+            == v
+        )
+
+
+def test_nonexpert_quantization_invalid_rejected():
+    with pytest.raises(pydantic.ValidationError):
+        KernelsArgs.model_validate({"nonexpert_quantization": "int8"})
+
+
+def test_nonexpert_quantization_nvfp4_accepted():
+    a = KernelsArgs.model_validate({"nonexpert_quantization": "nvfp4"})
+    assert a.nonexpert_quantization == "nvfp4"
+
+
+def test_moe_grouped_backend_valid_and_invalid():
+    for b in ("auto", "marlin", "cutlass", "deepgemm"):
+        assert (
+            KernelsArgs.model_validate({"moe_grouped_backend": b}).moe_grouped_backend
+            == b
+        )
+    with pytest.raises(pydantic.ValidationError):
+        KernelsArgs.model_validate({"moe_grouped_backend": "triton"})
+
+
+def test_moe_grouped_backend_dequant_rejected_for_training():
+    # M1: 'dequant' has no training/autograd path (the dispatch only wires marlin/deepgemm/cutlass);
+    # accepting it would silently run cutlass, so it is rejected with an explanatory message.
+    with pytest.raises(pydantic.ValidationError, match="not implemented for training"):
+        KernelsArgs.model_validate({"moe_grouped_backend": "dequant"})
+
+
+def test_large_head_attention_validator():
+    from axolotl.utils.schemas.config import AxolotlInputConfig
+
+    fn = AxolotlInputConfig.__dict__["validate_large_head_attention"].__func__
+    assert fn(AxolotlInputConfig, "AUTO") == "auto"  # case-normalized
+    assert fn(AxolotlInputConfig, "sdpa") == "sdpa"
+    assert fn(AxolotlInputConfig, None) is None
+    with pytest.raises(ValueError, match="large_head_attention must be one of"):
+        fn(AxolotlInputConfig, "trtion_flsah")  # typo rejected, not silently passed
+
+
+def test_moe_dequant_chunk_size_positive_accepted():
+    a = KernelsArgs.model_validate({"moe_dequant_chunk_size": 16})
+    assert a.moe_dequant_chunk_size == 16
+
+
+def test_moe_dequant_chunk_size_none_ok():
+    assert KernelsArgs.model_validate({}).moe_dequant_chunk_size is None
+
+
+@pytest.mark.parametrize("bad", [0, -1, -32])
+def test_moe_dequant_chunk_size_zero_and_negative_rejected(bad):
+    with pytest.raises(pydantic.ValidationError, match="positive integer"):
+        KernelsArgs.model_validate({"moe_dequant_chunk_size": bad})
+
+
+@pytest.mark.parametrize("bad", [2.5, True, "abc"])
+def test_moe_dequant_chunk_size_non_integer_rejected(bad):
+    with pytest.raises(pydantic.ValidationError, match="positive integer"):
+        KernelsArgs.model_validate({"moe_dequant_chunk_size": bad})
+
+
+def test_warn_unclaimed_nonexpert_quantization_fires(caplog):
+    # A non-expert quant policy set with no adapter that consumes it -> warn (no silent no-op).
+    import logging
+
+    from axolotl.integrations.kernels.adapters import ModelAdapter
+    from axolotl.integrations.kernels.plugin import KernelsPlugin
+
+    cfg = {"nonexpert_quantization": "nf4"}
+    with caplog.at_level(logging.WARNING):
+        KernelsPlugin._warn_unclaimed_nonexpert_quantization(cfg, [ModelAdapter()])
+    assert any(
+        "no active model adapter consumes it" in r.message for r in caplog.records
+    )
+
+
+def test_warn_unclaimed_nonexpert_quantization_silent_when_consumed(caplog):
+    import logging
+
+    from axolotl.integrations.kernels.adapters import ModelAdapter
+    from axolotl.integrations.kernels.plugin import KernelsPlugin
+
+    class _Consumer(ModelAdapter):
+        name = "consumer"
+
+        def consumes_nonexpert_quantization(self, cfg):
+            return True
+
+    cfg = {"nonexpert_quantization": "nf4"}
+    with caplog.at_level(logging.WARNING):
+        KernelsPlugin._warn_unclaimed_nonexpert_quantization(cfg, [_Consumer()])
+    assert not any(
+        "no active model adapter consumes it" in r.message for r in caplog.records
+    )
+
+
+@pytest.mark.parametrize("policy", [None, "none", "bf16"])
+def test_warn_unclaimed_nonexpert_quantization_skips_noop_policies(policy, caplog):
+    import logging
+
+    from axolotl.integrations.kernels.adapters import ModelAdapter
+    from axolotl.integrations.kernels.plugin import KernelsPlugin
+
+    cfg = {} if policy is None else {"nonexpert_quantization": policy}
+    with caplog.at_level(logging.WARNING):
+        KernelsPlugin._warn_unclaimed_nonexpert_quantization(cfg, [ModelAdapter()])
+    assert not any(
+        "no active model adapter consumes it" in r.message for r in caplog.records
+    )
diff --git a/tests/integrations/kernels/test_kernels_config_intent.py b/tests/integrations/kernels/test_kernels_config_intent.py
new file mode 100644
index 0000000000..ef7b0cf14d
--- /dev/null
+++ b/tests/integrations/kernels/test_kernels_config_intent.py
@@ -0,0 +1,216 @@
+"""No-GPU regression tests for the expert_backend intent surface (I1/I2) and the DSV4
+attention-LoRA guard (I7).
+
+These need no CUDA and no triton: they only construct pydantic config models. The isolation
+tests exercise KernelsArgs directly; the merged tests replicate axolotl's plugin-merge
+(``class AxolotlInputConfig(AxolotlInputConfigBase, KernelsArgs)``, see
+axolotl.integrations.config.merge_input_args) so the consumer validators run against the full
+config where ``experts_implementation`` / ``lora_mlp_kernel`` are real fields. The merged path
+is the one that was silently broken: pydantic runs same-mode ``before`` validators in reverse
+definition order, so before the fix the consumers ran before expert_backend was canonicalized.
+"""
+
+import logging
+
+import pydantic
+import pytest
+
+from axolotl.integrations.kernels.args import KernelsArgs
+
+
+def _merged_cls():
+    """Replicate merge_input_args() for the kernels plugin (AxolotlInputConfig + KernelsArgs)."""
+    from axolotl.utils.schemas.config import AxolotlInputConfig as Base
+
+    class _Merged(Base, KernelsArgs):
+        pass
+
+    return _Merged
+
+
+def _minimal(**overrides):
+    cfg = {
+        "base_model": "HuggingFaceTB/SmolLM2-135M",
+        "datasets": [{"path": "tatsu-lab/alpaca", "type": "alpaca"}],
+        "micro_batch_size": 1,
+        "gradient_accumulation_steps": 1,
+        "learning_rate": 0.0001,
+        "use_kernels": True,
+    }
+    cfg.update(overrides)
+    return cfg
+
+
+# --- I1/I2: expert_backend must canonicalize BEFORE the downstream consumers ----------------
+def test_expert_backend_scattermoe_triggers_dsv4_translation_isolation():
+    # The exact combo that regressed: pre-fix, disable_mlp_kernel ran before expert_backend was
+    # canonicalized, so use_scattermoe was still unset and dsv4_shared_mlp_lora_kernel stayed None.
+    a = KernelsArgs.model_validate(
+        {
+            "use_kernels": True,
+            "expert_backend": "scattermoe",
+            "use_dsv4_kernels": True,
+            "lora_mlp_kernel": True,
+        }
+    )
+    assert a.use_scattermoe is True
+    assert a.dsv4_shared_mlp_lora_kernel is True
+
+
+def test_expert_backend_matches_legacy_flags_isolation():
+    intent = KernelsArgs.model_validate(
+        {
+            "use_kernels": True,
+            "expert_backend": "scattermoe",
+            "use_dsv4_kernels": True,
+            "lora_mlp_kernel": True,
+        }
+    )
+    legacy = KernelsArgs.model_validate(
+        {
+            "use_kernels": True,
+            "use_scattermoe": True,
+            "use_dsv4_kernels": True,
+            "lora_mlp_kernel": True,
+        }
+    )
+    assert (intent.use_scattermoe, intent.dsv4_shared_mlp_lora_kernel) == (
+        legacy.use_scattermoe,
+        legacy.dsv4_shared_mlp_lora_kernel,
+    )
+
+
+@pytest.mark.parametrize("backend", ["scattermoe", "sonicmoe", "eager"])
+def test_merged_expert_backend_matches_legacy(backend):
+    # The full I3-complete check: on the MERGED config, experts_implementation and lora_mlp_kernel
+    # are real fields. The intent surface must produce a byte-identical canonical end state to the
+    # equivalent legacy flag for every consumer-derived field.
+    Merged = _merged_cls()
+    legacy_flag = {
+        "scattermoe": {"use_scattermoe": True},
+        "sonicmoe": {"use_sonicmoe": True},
+        "eager": {},
+    }[backend]
+
+    def snap(m):
+        return (
+            m.use_scattermoe,
+            m.use_sonicmoe,
+            getattr(m, "experts_implementation", None),
+            m.dsv4_shared_mlp_lora_kernel,
+            getattr(m, "lora_mlp_kernel", None),
+        )
+
+    intent = Merged.model_validate(
+        _minimal(expert_backend=backend, use_dsv4_kernels=True, lora_mlp_kernel=True)
+    )
+    legacy = Merged.model_validate(
+        _minimal(**legacy_flag, use_dsv4_kernels=True, lora_mlp_kernel=True)
+    )
+    assert snap(intent) == snap(legacy)
+
+
+def test_merged_scattermoe_intent_sets_experts_implementation():
+    # Guards the specific I3 symptom: pre-fix this left experts_implementation='eager'.
+    Merged = _merged_cls()
+    m = Merged.model_validate(_minimal(expert_backend="scattermoe"))
+    assert m.use_scattermoe is True
+    assert m.experts_implementation == "scattermoe"
+
+
+# --- I7: attention/module-level LoRA on a DSV4 fused-kernel run is unsupported ---------------
+def test_dsv4_attention_lora_rejected_isolation():
+    with pytest.raises(pydantic.ValidationError, match="lora_target_modules"):
+        KernelsArgs.model_validate(
+            {
+                "use_kernels": True,
+                "use_dsv4_kernels": True,
+                "lora_target_modules": ["q_proj"],
+            }
+        )
+
+
+def test_dsv4_attention_lora_allowed_with_exclude_isolation(caplog):
+    with caplog.at_level(logging.WARNING):
+        a = KernelsArgs.model_validate(
+            {
+                "use_kernels": True,
+                "use_dsv4_kernels": True,
+                "lora_target_modules": ["q_proj"],
+                "lora_exclude_modules": [".*indexer.*"],
+            }
+        )
+    assert a.use_dsv4_kernels is True
+    assert any("indexer scorer projections" in r.message for r in caplog.records)
+
+
+def test_dsv4_lora_target_linear_rejected_isolation():
+    # B1: lora_target_linear: true expands (find_all_linear_names) to attention q/k/v/o AFTER
+    # validation, so the guard must reject it too, not only lora_target_modules.
+    with pytest.raises(pydantic.ValidationError, match="lora_target_linear"):
+        KernelsArgs.model_validate(
+            {
+                "use_kernels": True,
+                "use_dsv4_kernels": True,
+                "lora_target_linear": True,
+            }
+        )
+
+
+def test_dsv4_lora_target_linear_allowed_with_exclude_isolation(caplog):
+    with caplog.at_level(logging.WARNING):
+        a = KernelsArgs.model_validate(
+            {
+                "use_kernels": True,
+                "use_dsv4_kernels": True,
+                "lora_target_linear": True,
+                "lora_exclude_modules": [".*indexer.*"],
+            }
+        )
+    assert a.use_dsv4_kernels is True
+    assert any("indexer scorer projections" in r.message for r in caplog.records)
+
+
+def test_dsv4_experts_only_allowed_isolation():
+    a = KernelsArgs.model_validate(
+        {
+            "use_kernels": True,
+            "use_dsv4_kernels": True,
+            "lora_target_parameters": ["mlp.experts.gate_up_proj"],
+        }
+    )
+    assert a.use_dsv4_kernels is True
+
+
+def test_attention_lora_without_dsv4_allowed_isolation():
+    # The guard must not fire when use_dsv4_kernels is off (normal attention LoRA is fine).
+    a = KernelsArgs.model_validate(
+        {"use_kernels": True, "lora_target_modules": ["q_proj"]}
+    )
+    assert a.use_dsv4_kernels is None
+
+
+def test_merged_dsv4_attention_lora_rejected():
+    Merged = _merged_cls()
+    with pytest.raises(pydantic.ValidationError, match="lora_target_modules"):
+        Merged.model_validate(
+            _minimal(use_dsv4_kernels=True, lora_target_modules=["q_proj"])
+        )
+
+
+def test_merged_dsv4_attention_lora_allowed_with_exclude():
+    Merged = _merged_cls()
+    m = Merged.model_validate(
+        _minimal(
+            use_dsv4_kernels=True,
+            lora_target_modules=["q_proj"],
+            lora_exclude_modules=[".*indexer.*"],
+        )
+    )
+    assert m.use_dsv4_kernels is True
+
+
+def test_merged_dsv4_lora_target_linear_rejected():
+    Merged = _merged_cls()
+    with pytest.raises(pydantic.ValidationError, match="lora_target_linear"):
+        Merged.model_validate(_minimal(use_dsv4_kernels=True, lora_target_linear=True))
diff --git a/tests/integrations/kernels/test_quant_training_guard.py b/tests/integrations/kernels/test_quant_training_guard.py
new file mode 100644
index 0000000000..a2eb2493af
--- /dev/null
+++ b/tests/integrations/kernels/test_quant_training_guard.py
@@ -0,0 +1,80 @@
+"""CPU-only tests for the scoped quantized-training guard (#4).
+
+Tests run against a throwaway module object (not the global ``transformers.trainer``) so they are
+fully isolated from conftest/session state.
+"""
+
+import types
+
+from axolotl.integrations.kernels.quant_training_guard import (
+    relax_quantized_training_guard,
+    restore_quantized_training_guard,
+)
+
+
+class _Model:
+    def __init__(self, peft):
+        self._hf_peft_config_loaded = peft
+
+
+class _PeftModel:
+    """Mirrors get_peft_model() output: exposes .peft_config, no _hf_peft_config_loaded flag."""
+
+    peft_config = {"default": object()}
+
+
+class _TrackingGuard:
+    def __init__(self):
+        self.calls = []
+
+    def __call__(self, model):
+        self.calls.append(model)
+
+
+def _fresh_module():
+    tracker = _TrackingGuard()
+    mod = types.SimpleNamespace(validate_quantization_for_training=tracker)
+    return mod, tracker
+
+
+def test_peft_model_skips_guard():
+    mod, tracker = _fresh_module()
+    relax_quantized_training_guard(mod)
+    mod.validate_quantization_for_training(_Model(peft=True))
+    assert tracker.calls == []  # PEFT/quantized is the supported pattern -> skipped
+
+
+def test_get_peft_model_skips_guard():
+    # axolotl's get_peft_model() path: .peft_config present, _hf_peft_config_loaded absent.
+    mod, tracker = _fresh_module()
+    relax_quantized_training_guard(mod)
+    mod.validate_quantization_for_training(_PeftModel())
+    assert (
+        tracker.calls == []
+    )  # recognized as PEFT -> skipped (FP8 base + adapters is supported)
+
+
+def test_non_peft_delegates_to_original():
+    mod, tracker = _fresh_module()
+    relax_quantized_training_guard(mod)
+    model = _Model(peft=False)
+    mod.validate_quantization_for_training(model)
+    assert tracker.calls == [model]  # delegated to the real guard
+
+
+def test_idempotent():
+    mod, _ = _fresh_module()
+    relax_quantized_training_guard(mod)
+    wrapped = mod.validate_quantization_for_training
+    relax_quantized_training_guard(mod)  # second call must not re-wrap
+    assert mod.validate_quantization_for_training is wrapped
+
+
+def test_restore():
+    mod, tracker = _fresh_module()
+    relax_quantized_training_guard(mod)
+    assert mod.validate_quantization_for_training is not tracker  # wrapped
+    restore_quantized_training_guard(mod)
+    assert (
+        mod.validate_quantization_for_training is tracker
+    )  # original preserved/restored
diff --git a/tests/integrations/monkeypatch/__init__.py b/tests/integrations/monkeypatch/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/integrations/monkeypatch/test_tiled_mlp_moe.py b/tests/integrations/monkeypatch/test_tiled_mlp_moe.py
new file mode 100644
index 0000000000..6dab0df109
--- /dev/null
+++ b/tests/integrations/monkeypatch/test_tiled_mlp_moe.py
@@ -0,0 +1,474 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+"""
+Single-GPU correctness tests for the TiledMLP autograd function under both
+dense and MoE block forwards. Synthetic-shape modules only; no real
+transformers checkpoints are loaded.
+"""
+
+import copy
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _requires_cuda():
+    return pytest.mark.skipif(
+        not torch.cuda.is_available(), reason="CUDA not available"
+    )
+
+
+pytestmark = _requires_cuda()
+
+DEVICE = "cuda"
+
+
+# ────────────────────────────── Helpers ──────────────────────────────
+
+
+class TinyDenseMLP(nn.Module):
+    """LlamaMLP-shape: gate * up -> down, no bias, silu activation."""
+
+    def __init__(self, hidden, intermediate, dtype=torch.float32):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden, intermediate, bias=False, dtype=dtype)
+        self.up_proj = nn.Linear(hidden, intermediate, bias=False, dtype=dtype)
+        self.down_proj = nn.Linear(intermediate, hidden, bias=False, dtype=dtype)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TinyMoEBlock(nn.Module):
+    """Hand-rolled MoE: top-k softmax router + per-expert SwiGLU MLPs.
+
+    Stays intentionally simple so the test exercises sharding semantics
+    without dragging in transformers, peft, or kernel libs.
+    """
+
+    def __init__(self, hidden, intermediate, num_experts, top_k, dtype=torch.float32):
+        super().__init__()
+        self.hidden = hidden
+        self.intermediate = intermediate
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.gate = nn.Linear(hidden, num_experts, bias=False, dtype=dtype)
+        # Per-expert SwiGLU weights packed as 3D tensors.
+        self.gate_proj = nn.Parameter(
+            torch.randn(num_experts, hidden, intermediate, dtype=dtype) * 0.02
+        )
+        self.up_proj = nn.Parameter(
+            torch.randn(num_experts, hidden, intermediate, dtype=dtype) * 0.02
+        )
+        self.down_proj = nn.Parameter(
+            torch.randn(num_experts, intermediate, hidden, dtype=dtype) * 0.02
+        )
+
+    def forward(self, x):
+        bsz, seq, h = x.shape
+        flat = x.reshape(-1, h)
+        logits = self.gate(flat)
+        weights = F.softmax(logits, dim=-1, dtype=torch.float32)
+        top_w, top_i = torch.topk(weights, self.top_k, dim=-1)
+        top_w = top_w / top_w.sum(dim=-1, keepdim=True)
+        top_w = top_w.to(flat.dtype)
+
+        out = torch.zeros_like(flat)
+        for e in range(self.num_experts):
+            mask = top_i == e
+            if not mask.any():
+                continue
+            # tokens routed to expert e (with their per-slot weight)
+            token_rows, slot_idx = mask.nonzero(as_tuple=True)
+            xe = flat[token_rows]
+            we = top_w[token_rows, slot_idx].unsqueeze(-1)
+            gate = xe @ self.gate_proj[e]
+            up = xe @ self.up_proj[e]
+            h_e = F.silu(gate) * up
+            ye = h_e @ self.down_proj[e]
+            out.index_add_(0, token_rows, we * ye)
+        return out.reshape(bsz, seq, h)
+
+
+def _clone_module(mod):
+    """Deep copy + detach + re-attach to autograd to compare two runs."""
+    cloned = copy.deepcopy(mod)
+    return cloned
+
+
+def _grad_dict(mod):
+    return {
+        n: p.grad.detach().clone()
+        for n, p in mod.named_parameters()
+        if p.grad is not None
+    }
+
+
+def _run_untiled(mod, x):
+    x = x.clone().detach().requires_grad_(True)
+    y = mod(x)
+    g = torch.randn_like(y)
+    y.backward(g)
+    return y.detach().clone(), x.grad.detach().clone(), _grad_dict(mod), g
+
+
+def _run_tiled(mod, x, upstream_grad, shards):
+    """Re-run forward+backward but routed through TiledMLP."""
+    from axolotl.monkeypatch.tiled_mlp.base import TiledMLP
+
+    # Re-fetch fn that takes (self, x) — matches what the patcher passes.
+    forward_fn = type(mod).forward
+
+    x = x.clone().detach().requires_grad_(True)
+    compute_params = [p for p in mod.parameters() if p.requires_grad]
+    y = TiledMLP.apply(forward_fn, mod, x, shards, compute_params)
+    if isinstance(y, tuple):  # MoE block forwards may return tuples
+        y = y[0]
+    y.backward(upstream_grad)
+    return y.detach().clone(), x.grad.detach().clone(), _grad_dict(mod)
+
+
+# ────────────────────────────── Dense parity ──────────────────────────────
+
+
+def test_tiled_dense_mlp_parity_fp32():
+    """Dense LlamaMLP-shape: tiled vs un-tiled must match closely."""
+    torch.manual_seed(0)
+    hidden, intermediate, seq = 64, 128, 64
+    mlp_ref = TinyDenseMLP(hidden, intermediate).to(DEVICE)
+    mlp_tile = _clone_module(mlp_ref)
+
+    # ``TiledMLP``'s backward narrows into a flattened ``x_grad`` buffer
+    # using offsets along dim 1 only — sequence-packed inputs (batch=1)
+    # are the supported shape; multi-batch tensors aren't contiguous in
+    # the way the narrow assumes. Production inputs from transformers
+    # are batch=1 after sequence packing, so this matches reality.
+    x = torch.randn(1, seq, hidden, device=DEVICE)
+    y_ref, dx_ref, gp_ref, upstream = _run_untiled(mlp_ref, x)
+    y_tile, dx_tile, gp_tile = _run_tiled(mlp_tile, x, upstream, shards=4)
+
+    # FMA reordering across shards introduces sub-eps noise in fp32; allow
+    # a small tolerance on outputs and grads.
+    assert torch.allclose(y_ref, y_tile, atol=1e-5, rtol=1e-5), (
+        f"forward mismatch max={((y_ref - y_tile).abs().max()).item()}"
+    )
+    assert torch.allclose(dx_ref, dx_tile, atol=1e-5, rtol=1e-5), (
+        f"dX mismatch max={((dx_ref - dx_tile).abs().max()).item()}"
+    )
+    for name in gp_ref:
+        diff = (gp_ref[name] - gp_tile[name]).abs().max().item()
+        assert diff < 1e-5, f"param-grad mismatch {name}: max={diff}"
+
+
+# ────────────────────────────── MoE parity ──────────────────────────────
+
+
+def test_tiled_moe_block_parity_fp32():
+    """Hand-rolled MoE block: tiled vs un-tiled fp32 parity."""
+    torch.manual_seed(1)
+    hidden, intermediate, seq = 64, 128, 64
+    moe_ref = TinyMoEBlock(hidden, intermediate, num_experts=8, top_k=2).to(DEVICE)
+    moe_tile = _clone_module(moe_ref)
+
+    x = torch.randn(1, seq, hidden, device=DEVICE)
+    y_ref, dx_ref, gp_ref, upstream = _run_untiled(moe_ref, x)
+    y_tile, dx_tile, gp_tile = _run_tiled(moe_tile, x, upstream, shards=4)
+
+    # The MoE forward involves index_add and routing, which is not
+    # numerically deterministic across different batch sizes for fp32 in
+    # all setups — but at the synthetic small scale we expect tight match.
+    assert torch.allclose(y_ref, y_tile, atol=1e-5, rtol=1e-5), (
+        f"MoE forward mismatch max={((y_ref - y_tile).abs().max()).item()}"
+    )
+    assert torch.allclose(dx_ref, dx_tile, atol=1e-5, rtol=1e-5), (
+        f"MoE dX mismatch max={((dx_ref - dx_tile).abs().max()).item()}"
+    )
+    for name in gp_ref:
+        diff = (gp_ref[name] - gp_tile[name]).abs().max().item()
+        assert diff < 1e-5, f"MoE param-grad mismatch {name}: max={diff}"
+
+
+# ─────────────────────── scattermoe-lora + tiled parity ───────────────────
+
+
+def _build_scattermoe_block(hidden, intermediate, num_experts, top_k, dtype, device):
+    """Build a minimal :class:`ScatterMoEGatedMLP`-compatible module.
+
+    Attributes are populated to match what :meth:`ScatterMoEGatedMLP.forward`
+    expects (``router``, ``input_linear``, ``output_linear``, ``activation``).
+    """
+    try:
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            ScatterMoEGatedMLP,
+        )
+    except ImportError:
+        pytest.skip("scattermoe_lora kernels not available")
+
+    block = ScatterMoEGatedMLP()
+    # router: a Linear-shaped router with top_k / num_experts attrs.
+    router = SimpleNamespace()
+    router.layer = nn.Linear(hidden, num_experts, bias=False, dtype=dtype).to(device)
+    router.top_k = top_k
+    router.num_experts = num_experts
+    block.router = router
+    # input_linear and output_linear store 3D weights: [E, *, *].
+    in_weight = nn.Parameter(
+        torch.randn(num_experts, 2 * intermediate, hidden, dtype=dtype, device=device)
+        * 0.02
+    )
+    out_weight = nn.Parameter(
+        torch.randn(num_experts, hidden, intermediate, dtype=dtype, device=device)
+        * 0.02
+    )
+    block.input_linear = nn.Module()
+    block.input_linear.weight = in_weight
+    block.output_linear = nn.Module()
+    block.output_linear.weight = out_weight
+    block.activation = nn.SiLU()
+    # Register the params so .parameters() walks them.
+    block.input_linear.register_parameter("weight", in_weight)
+    block.output_linear.register_parameter("weight", out_weight)
+    return block
+
+
+def test_tiled_scattermoe_gated_mlp_parity():
+    """ScatterMoEGatedMLP: tiled vs un-tiled fwd+bwd parity in bf16.
+
+    Uses the same tolerance scale as ``tests/integrations/test_scattermoe_lora_kernels.py``
+    (norm-relative error < 1% for weight grads is the established bar there
+    given the bf16 + tiled reduction order differences).
+    """
+    pytest.importorskip("triton")
+    torch.manual_seed(2)
+    hidden, intermediate = 64, 128
+    num_experts, top_k = 8, 2
+    seq = 64
+    dtype = torch.bfloat16
+
+    block_ref = _build_scattermoe_block(
+        hidden, intermediate, num_experts, top_k, dtype, DEVICE
+    )
+    block_tile = copy.deepcopy(block_ref)
+
+    x = torch.randn(1, seq, hidden, device=DEVICE, dtype=dtype)
+    y_ref, dx_ref, gp_ref, upstream = _run_untiled(block_ref, x)
+    y_tile, dx_tile, gp_tile = _run_tiled(block_tile, x, upstream, shards=4)
+
+    def _rel(a, b):
+        return ((a.float() - b.float()).norm() / (b.float().norm() + 1e-6)).item()
+
+    assert _rel(y_tile, y_ref) < 1e-2, (
+        f"scattermoe forward rel_err={_rel(y_tile, y_ref)}"
+    )
+    assert _rel(dx_tile, dx_ref) < 1e-2, (
+        f"scattermoe dX rel_err={_rel(dx_tile, dx_ref)}"
+    )
+    for name in gp_ref:
+        if name not in gp_tile:
+            continue
+        rel = _rel(gp_tile[name], gp_ref[name])
+        assert rel < 1e-2, f"scattermoe param-grad {name} rel_err={rel}"
+
+
+# ─────────────────── Patcher: MoE block discovery & dispatch ────────────────
+
+
+def test_resolve_moe_block_cls_picks_first_available():
+    """The patcher walks the suffix list in order; the first hit wins."""
+    from axolotl.monkeypatch.tiled_mlp.patch import _resolve_moe_block_cls
+
+    module = SimpleNamespace(
+        FooMoE=object,
+        FooMoeMLP=object,  # would also match but later in the list
+    )
+    cls = _resolve_moe_block_cls(module, "Foo")
+    assert cls is module.FooMoeMLP, "MoeMLP should be preferred over MoE"
+
+
+def test_resolve_moe_block_cls_returns_none_for_dense_model():
+    from axolotl.monkeypatch.tiled_mlp.patch import _resolve_moe_block_cls
+
+    module = SimpleNamespace(FooMLP=object)
+    assert _resolve_moe_block_cls(module, "Foo") is None
+
+
+# ─────────────── Grad parity under non-uniform per-token loss weights ─────────
+#
+# Sequence-dim sharding makes per-shard parameter-grads additive: the full
+# batch gradient is the SUM of shard contributions, not the mean. If the
+# tiled backward ever scaled by ``1/total_shards`` (the historical
+# ``GradientAccumulator.gradient_scale``), per-shard non-uniform weights
+# would make the mean visibly diverge from the un-tiled reference — uniform
+# loss weights can mask the bug because the per-shard means happen to add
+# up to a related-magnitude value. These tests exercise multiple
+# ``shards ∈ {1, 2, 4}`` with deliberately non-uniform per-token weights
+# so a regression in the scaling semantics fails loudly.
+
+
+def _run_untiled_with_upstream(mod, x, upstream):
+    """Un-tiled fwd+bwd given a fixed upstream grad."""
+    x = x.clone().detach().requires_grad_(True)
+    y = mod(x)
+    y.backward(upstream)
+    return y.detach().clone(), x.grad.detach().clone(), _grad_dict(mod)
+
+
+@pytest.mark.parametrize("shards", [1, 2, 4])
+def test_tiled_dense_mlp_grad_parity_nonuniform_weights(shards):
+    """Dense MLP: tiled vs un-tiled grad parity with non-uniform per-token weights.
+
+    The upstream grad's magnitude varies per token (non-uniform loss weights),
+    so each shard contributes a distinct fraction of the total parameter grad.
+    A mean-vs-sum bug shows up as a ``shards``-dependent scaling error.
+    """
+    torch.manual_seed(100 + shards)
+    hidden, intermediate = 64, 128
+    seq = 128
+    mlp_ref = TinyDenseMLP(hidden, intermediate).to(DEVICE)
+    mlp_tile = _clone_module(mlp_ref)
+
+    x = torch.randn(1, seq, hidden, device=DEVICE)
+    # Non-uniform per-token weights make per-shard grad contributions
+    # distinct, exposing any incorrect averaging.
+    per_token_w = torch.linspace(0.1, 3.0, seq, device=DEVICE).view(1, seq, 1)
+    upstream = torch.randn(1, seq, hidden, device=DEVICE) * per_token_w
+
+    y_ref, dx_ref, gp_ref = _run_untiled_with_upstream(mlp_ref, x, upstream)
+    y_tile, dx_tile, gp_tile = _run_tiled(mlp_tile, x, upstream, shards=shards)
+
+    assert torch.allclose(y_ref, y_tile, atol=1e-5, rtol=1e-5), (
+        f"shards={shards}: forward mismatch max={((y_ref - y_tile).abs().max()).item()}"
+    )
+    assert torch.allclose(dx_ref, dx_tile, atol=1e-5, rtol=1e-5), (
+        f"shards={shards}: dX mismatch max={((dx_ref - dx_tile).abs().max()).item()}"
+    )
+    for name in gp_ref:
+        diff = (gp_ref[name] - gp_tile[name]).abs().max().item()
+        ref_norm = gp_ref[name].abs().max().item() + 1e-8
+        # Tight bound in fp32; rel error must be tiny so a 1/N or N
+        # scaling error (which would give 25%-400% relative drift) is
+        # impossible to miss.
+        assert diff / ref_norm < 1e-4, (
+            f"shards={shards}: param-grad {name} rel_err={diff / ref_norm}"
+        )
+
+
+@pytest.mark.parametrize("shards", [1, 2, 4])
+def test_tiled_moe_grad_parity_nonuniform_weights(shards):
+    """MoE block: tiled vs un-tiled grad parity with non-uniform per-token weights."""
+    torch.manual_seed(200 + shards)
+    hidden, intermediate = 64, 128
+    seq = 128
+    moe_ref = TinyMoEBlock(hidden, intermediate, num_experts=8, top_k=2).to(DEVICE)
+    moe_tile = _clone_module(moe_ref)
+
+    x = torch.randn(1, seq, hidden, device=DEVICE)
+    per_token_w = torch.linspace(0.1, 3.0, seq, device=DEVICE).view(1, seq, 1)
+    upstream = torch.randn(1, seq, hidden, device=DEVICE) * per_token_w
+
+    y_ref, dx_ref, gp_ref = _run_untiled_with_upstream(moe_ref, x, upstream)
+    y_tile, dx_tile, gp_tile = _run_tiled(moe_tile, x, upstream, shards=shards)
+
+    assert torch.allclose(y_ref, y_tile, atol=1e-5, rtol=1e-5), (
+        f"shards={shards}: MoE forward mismatch "
+        f"max={((y_ref - y_tile).abs().max()).item()}"
+    )
+    assert torch.allclose(dx_ref, dx_tile, atol=1e-5, rtol=1e-5), (
+        f"shards={shards}: MoE dX mismatch "
+        f"max={((dx_ref - dx_tile).abs().max()).item()}"
+    )
+    for name in gp_ref:
+        diff = (gp_ref[name] - gp_tile[name]).abs().max().item()
+        ref_norm = gp_ref[name].abs().max().item() + 1e-8
+        assert diff / ref_norm < 1e-4, (
+            f"shards={shards}: MoE param-grad {name} rel_err={diff / ref_norm}"
+        )
+
+
+@pytest.mark.parametrize("shards", [1, 2, 4])
+def test_tiled_dense_mlp_grad_parity_bf16(shards):
+    """Dense MLP: bf16 grad parity at the param dtype (no fp32 accumulator).
+
+    Guards the default param-dtype accumulator path against regression.
+    bf16 reduction order across shards means we use a relative tolerance
+    rather than bitwise equality.
+    """
+    torch.manual_seed(300 + shards)
+    hidden, intermediate = 64, 128
+    seq = 128
+    dtype = torch.bfloat16
+    mlp_ref = TinyDenseMLP(hidden, intermediate, dtype=dtype).to(DEVICE)
+    mlp_tile = _clone_module(mlp_ref)
+
+    x = torch.randn(1, seq, hidden, device=DEVICE, dtype=dtype)
+    per_token_w = torch.linspace(0.1, 3.0, seq, device=DEVICE).view(1, seq, 1).to(dtype)
+    upstream = (torch.randn(1, seq, hidden, device=DEVICE) * per_token_w).to(dtype)
+
+    y_ref, dx_ref, gp_ref = _run_untiled_with_upstream(mlp_ref, x, upstream)
+    y_tile, dx_tile, gp_tile = _run_tiled(mlp_tile, x, upstream, shards=shards)
+
+    def _rel(a, b):
+        return ((a.float() - b.float()).norm() / (b.float().norm() + 1e-6)).item()
+
+    assert _rel(y_tile, y_ref) < 5e-3, (
+        f"shards={shards}: bf16 forward rel_err={_rel(y_tile, y_ref)}"
+    )
+    assert _rel(dx_tile, dx_ref) < 5e-3, (
+        f"shards={shards}: bf16 dX rel_err={_rel(dx_tile, dx_ref)}"
+    )
+    for name in gp_ref:
+        rel = _rel(gp_tile[name], gp_ref[name])
+        # Tight bound — a 1/N scaling bug would put rel_err ≈ (N-1)/N,
+        # which is far above this threshold for any N ≥ 2.
+        assert rel < 5e-3, f"shards={shards}: bf16 param-grad {name} rel_err={rel}"
+
+
+def test_tiled_grad_accumulator_dtype_matches_param_dtype():
+    """Regression guard: TiledMLP backward should accumulate at param dtype
+    by default (not fp32), so the on-the-fly accumulator does not double
+    the parameter-side memory footprint in bf16 training.
+
+    We snapshot the accumulator dtype by patching ``torch.zeros_like`` to
+    record the dtype of zero-tensors allocated for compute params during
+    backward. The assertion is that none of those allocations request
+    fp32 when the params are bf16.
+    """
+    from axolotl.monkeypatch.tiled_mlp.base import TiledMLP
+
+    torch.manual_seed(42)
+    dtype = torch.bfloat16
+    mlp = TinyDenseMLP(64, 128, dtype=dtype).to(DEVICE)
+    compute_params = [p for p in mlp.parameters() if p.requires_grad]
+    param_ids = {id(p) for p in compute_params}
+
+    x = torch.randn(1, 64, 64, device=DEVICE, dtype=dtype)
+    upstream = torch.randn(1, 64, 64, device=DEVICE, dtype=dtype)
+
+    real_zeros_like = torch.zeros_like
+    allocated_dtypes: list[torch.dtype] = []
+
+    def spy_zeros_like(t, *args, **kwargs):
+        # Only record allocations whose shape matches one of the
+        # compute params (the accumulator buffers we care about).
+        if id(t) in param_ids:
+            allocated_dtypes.append(kwargs.get("dtype", t.dtype))
+        return real_zeros_like(t, *args, **kwargs)
+
+    x_req = x.clone().detach().requires_grad_(True)
+    torch.zeros_like = spy_zeros_like
+    try:
+        y = TiledMLP.apply(type(mlp).forward, mlp, x_req, 4, compute_params)
+        y.backward(upstream)
+    finally:
+        torch.zeros_like = real_zeros_like
+
+    assert allocated_dtypes, "expected accumulator allocations to be observed"
+    # Default path must NOT pre-allocate fp32 buffers when params are bf16.
+    assert all(d == dtype for d in allocated_dtypes), (
+        f"expected accumulator dtype == {dtype}, got {allocated_dtypes}"
+    )
diff --git a/tests/integrations/mora/test_mora.py b/tests/integrations/mora/test_mora.py
new file mode 100644
index 0000000000..464979cb60
--- /dev/null
+++ b/tests/integrations/mora/test_mora.py
@@ -0,0 +1,160 @@
+"""Integration tests for the MoRA / ReMoRA adapter path."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from axolotl.integrations.base import PluginManager
+from axolotl.integrations.mora import plugin as mora_plugin
+from axolotl.loaders import adapter as adapter_module
+from axolotl.loaders.adapter import load_adapter
+from axolotl.utils.dict import DictDefault
+
+
+class TestMoraAdapterLoading:
+    """MoRA adapter selection and config wiring."""
+
+    def test_load_adapter_uses_plugin_lora_like_registration(self, monkeypatch):
+        model = torch.nn.Linear(4, 4)
+        cfg = DictDefault(
+            {
+                "adapter": "mora",
+                "mora": {"use_mora": True, "mora_type": "rope"},
+            }
+        )
+
+        PluginManager.get_instance().plugins["axolotl.integrations.mora.MoraPlugin"] = (
+            mora_plugin.MoraPlugin()
+        )
+
+        calls = []
+
+        def fake_load_lora(*args, **kwargs):
+            calls.append((args, kwargs))
+            return args[0], "adapter-config"
+
+        monkeypatch.setattr(adapter_module, "load_lora", fake_load_lora)
+
+        _, config = load_adapter(model, cfg, "mora")
+
+        assert config == "adapter-config"
+        assert calls[0][1]["config_only"] is False
+
+    def test_mora_plugin_raises_when_peft_missing_support(self):
+        model = torch.nn.Linear(4, 4)
+        cfg = DictDefault(
+            {
+                "adapter": "mora",
+                "mora": {"use_mora": True, "mora_type": "rope"},
+            }
+        )
+        PluginManager.get_instance().plugins["axolotl.integrations.mora.MoraPlugin"] = (
+            mora_plugin.MoraPlugin()
+        )
+
+        with pytest.raises(ImportError, match="MoRA support"):
+            load_adapter(model, cfg, "mora", config_only=True)
+
+    def test_mora_plugin_rejects_quantized_base_model(self):
+        model = torch.nn.Linear(4, 4)
+        cfg = DictDefault(
+            {
+                "adapter": "mora",
+                "load_in_4bit": True,
+                "mora": {"use_mora": True, "mora_type": "rope"},
+            }
+        )
+        PluginManager.get_instance().plugins["axolotl.integrations.mora.MoraPlugin"] = (
+            mora_plugin.MoraPlugin()
+        )
+
+        with pytest.raises(ValueError, match="full-precision base model"):
+            load_adapter(model, cfg, "mora", config_only=True)
+
+    def test_mora_plugin_builds_mora_config_when_supported(self, monkeypatch):
+        model = torch.nn.Linear(4, 4)
+        cfg = DictDefault(
+            {
+                "adapter": "mora",
+                "mora": {
+                    "use_mora": True,
+                    "mora_type": "rope",
+                },
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+            }
+        )
+
+        captured = {}
+
+        class FakeLoraConfig:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+                self.__dict__.update(kwargs)
+
+        fake_model = SimpleNamespace(print_trainable_parameters=Mock())
+        PluginManager.get_instance().plugins["axolotl.integrations.mora.MoraPlugin"] = (
+            mora_plugin.MoraPlugin()
+        )
+        monkeypatch.setattr(mora_plugin, "_peft_supports_mora", lambda: True)
+        monkeypatch.setattr(adapter_module, "LoraConfig", FakeLoraConfig)
+        monkeypatch.setattr(
+            adapter_module, "get_peft_model", Mock(return_value=fake_model)
+        )
+
+        _, config = load_adapter(model, cfg, "mora", config_only=True)
+
+        assert captured["use_mora"] is True
+        assert captured["mora_type"] == 6
+        assert captured["task_type"].name == "CAUSAL_LM"
+        assert config is not None
+        assert config.use_mora is True
+        assert config.mora_type == 6
+
+    def test_mora_plugin_uses_lora_model_dir_resume_path(self, monkeypatch):
+        model = torch.nn.Linear(4, 4)
+        cfg = DictDefault(
+            {
+                "adapter": "mora",
+                "mora": {"use_mora": True, "mora_type": "rope"},
+                "lora_model_dir": "adapter-checkpoint",
+                "lora_on_cpu": False,
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+            }
+        )
+
+        class FakeLoraConfig:
+            def __init__(self, **kwargs):
+                self.__dict__.update(kwargs)
+
+        class FakePeftModel:
+            def print_trainable_parameters(self):
+                pass
+
+            def named_parameters(self):
+                return []
+
+        from_pretrained = Mock(return_value=FakePeftModel())
+        PluginManager.get_instance().plugins["axolotl.integrations.mora.MoraPlugin"] = (
+            mora_plugin.MoraPlugin()
+        )
+        monkeypatch.setattr(mora_plugin, "_peft_supports_mora", lambda: True)
+        monkeypatch.setattr(adapter_module, "LoraConfig", FakeLoraConfig)
+        monkeypatch.setattr(
+            adapter_module.PeftModel,
+            "from_pretrained",
+            from_pretrained,
+        )
+
+        peft_model, config = load_adapter(model, cfg, "mora")
+
+        assert isinstance(peft_model, FakePeftModel)
+        assert config.use_mora is True
+        from_pretrained.assert_called_once()
+        assert from_pretrained.call_args.args[:2] == (model, "adapter-checkpoint")
+        assert from_pretrained.call_args.kwargs["is_trainable"] is True
diff --git a/tests/integrations/test_adapter_plugin_registry.py b/tests/integrations/test_adapter_plugin_registry.py
new file mode 100644
index 0000000000..e3102a90ec
--- /dev/null
+++ b/tests/integrations/test_adapter_plugin_registry.py
@@ -0,0 +1,73 @@
+"""Core adapter plugin registry tests."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from axolotl.integrations.base import AdapterCapabilities, BasePlugin, PluginManager
+from axolotl.loaders import adapter as adapter_module
+from axolotl.loaders.adapter import load_adapter
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class FakeAdapterPlugin(BasePlugin):
+    def get_adapter_capabilities(self) -> list[AdapterCapabilities]:
+        return [AdapterCapabilities(name="fake-adapter", lora_like=True, relora=True)]
+
+    def get_lora_config_kwargs(self, cfg: DictDefault) -> dict:
+        if cfg.adapter != "fake-adapter":
+            return {}
+        return {"fake_kwarg": "from-plugin"}
+
+
+class TestAdapterPluginRegistry:
+    def test_lora_like_plugin_adapter_contributes_peft_kwargs(self, monkeypatch):
+        model = torch.nn.Linear(4, 4)
+        cfg = DictDefault(
+            {
+                "adapter": "fake-adapter",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+            }
+        )
+        PluginManager.get_instance().plugins["fake"] = FakeAdapterPlugin()
+        captured = {}
+
+        class FakeLoraConfig:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+                self.__dict__.update(kwargs)
+
+        monkeypatch.setattr(adapter_module, "LoraConfig", FakeLoraConfig)
+        monkeypatch.setattr(adapter_module, "get_peft_model", Mock())
+
+        _, config = load_adapter(model, cfg, "fake-adapter", config_only=True)
+
+        assert config is not None
+        assert captured["fake_kwarg"] == "from-plugin"
+        assert captured["task_type"].name == "CAUSAL_LM"
+
+    def test_unknown_adapter_error_mentions_plugin_registry(self):
+        model = torch.nn.Linear(4, 4)
+        cfg = DictDefault({"adapter": "missing-adapter"})
+
+        with pytest.raises(NotImplementedError, match="registered by a plugin"):
+            load_adapter(model, cfg, "missing-adapter")
+
+    def test_relora_accepts_plugin_adapter_capability(self, min_base_cfg):
+        PluginManager.get_instance().plugins["fake"] = FakeAdapterPlugin()
+        cfg = min_base_cfg | DictDefault(
+            {
+                "adapter": "fake-adapter",
+                "relora": True,
+                "jagged_restart_steps": 100,
+            }
+        )
+
+        validated = validate_config(cfg)
+
+        assert validated.adapter == "fake-adapter"
+        assert validated.relora is True
diff --git a/tests/integrations/test_diffusion.py b/tests/integrations/test_diffusion.py
new file mode 100644
index 0000000000..432552724c
--- /dev/null
+++ b/tests/integrations/test_diffusion.py
@@ -0,0 +1,274 @@
+"""Tests for diffusion trainer integration."""
+
+# pylint: disable=redefined-outer-name,protected-access
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from axolotl.integrations.diffusion import DiffusionTrainer
+from axolotl.integrations.diffusion.utils import create_bidirectional_attention_mask
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture
+def mock_tokenizer():
+    """Create a mock tokenizer."""
+    tokenizer = Mock()
+    tokenizer.bos_token_id = 1
+    tokenizer.eos_token_id = 2
+    tokenizer.pad_token_id = 0
+    return tokenizer
+
+
+@pytest.fixture
+def diffusion_config():
+    """Create a diffusion config."""
+    return DictDefault(
+        {
+            "diffusion": {
+                "mask_token_id": 32000,
+                "eps": 1e-3,
+                "importance_weighting": False,
+            },
+            "sample_packing": False,
+        }
+    )
+
+
+@pytest.fixture
+def diffusion_trainer_instance(mock_tokenizer, diffusion_config):
+    """Create a diffusion trainer instance for testing methods directly."""
+    # Create a minimal trainer instance just for testing methods
+    trainer = object.__new__(DiffusionTrainer)  # Bypass __init__
+    trainer.axolotl_cfg = diffusion_config
+    trainer._special_token_ids = {0, 1, 2}  # pad, bos, eos
+    trainer.processing_class = mock_tokenizer
+    trainer.store_metrics = Mock()  # Mock metrics storage
+    return trainer
+
+
+class TestDiffusionTrainer:
+    """Test the DiffusionTrainer class."""
+
+    def test_forward_process_basic(self, diffusion_trainer_instance):
+        """Test basic forward process without labels."""
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+
+        noisy_batch, masked_indices, p_mask = (
+            diffusion_trainer_instance._forward_process(input_ids, eps=0.1)
+        )
+
+        # Check shapes
+        assert noisy_batch.shape == input_ids.shape
+        assert masked_indices.shape == input_ids.shape
+        assert p_mask.shape == input_ids.shape
+
+        # Check that special tokens are not masked
+        special_token_positions = (input_ids == 1) | (input_ids == 2) | (input_ids == 0)
+        assert not masked_indices[special_token_positions].any()
+
+        # Check that mask token is applied
+        mask_token_id = diffusion_trainer_instance.axolotl_cfg.diffusion.mask_token_id
+        masked_positions = masked_indices
+        if masked_positions.any():
+            assert (noisy_batch[masked_positions] == mask_token_id).all()
+
+    def test_forward_process_with_labels(self, diffusion_trainer_instance):
+        """Test forward process with SFT labels."""
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+        labels = torch.tensor([[-100, -100, 20, 30, 2]], dtype=torch.long)
+
+        noisy_batch, masked_indices, p_mask = (
+            diffusion_trainer_instance._forward_process(
+                input_ids, labels=labels, eps=0.1
+            )
+        )
+
+        # Check shapes
+        assert noisy_batch.shape == input_ids.shape
+        assert masked_indices.shape == input_ids.shape
+        assert p_mask.shape == input_ids.shape
+
+        # Check that only answer tokens can be masked (where labels != -100)
+        non_answer_mask = labels == -100
+
+        # No masking should occur on non-answer tokens
+        assert not masked_indices[non_answer_mask].any()
+
+        # p_mask should be the same for all positions (sampled timestep),
+        # but masking is only applied to answer tokens
+        assert p_mask.shape == input_ids.shape
+        # Verify that masked_indices respects the answer mask
+        assert not masked_indices[non_answer_mask].any()
+
+    def test_forward_process_with_attention_mask(self, diffusion_trainer_instance):
+        """Test forward process with attention mask."""
+        input_ids = torch.tensor([[1, 10, 20, 0]], dtype=torch.long)
+        attention_mask = torch.tensor([[1, 1, 1, 0]], dtype=torch.long)
+
+        _, masked_indices, p_mask = diffusion_trainer_instance._forward_process(
+            input_ids, attention_mask=attention_mask, eps=0.1
+        )
+
+        # Check that padding tokens are not masked
+        padding_positions = attention_mask == 0
+        assert not masked_indices[padding_positions].any()
+        assert (p_mask[padding_positions] == 0).all()
+
+    def test_bidirectional_attention_mask_no_packing(self, diffusion_trainer_instance):
+        """Test bidirectional attention mask without sample packing."""
+        input_ids = torch.tensor([[1, 10, 20, 2]], dtype=torch.long)
+
+        mask = create_bidirectional_attention_mask(input_ids)
+
+        # Should be all-to-all attention
+        expected_shape = (1, 1, 4, 4)
+        assert mask.shape == expected_shape
+        assert mask.all()
+
+    def test_bidirectional_attention_mask_with_packing(
+        self, diffusion_trainer_instance
+    ):
+        """Test bidirectional attention mask with sample packing."""
+        diffusion_trainer_instance.axolotl_cfg.sample_packing = True
+        input_ids = torch.tensor([[1, 10, 20, 30, 40, 2]], dtype=torch.long)
+        # Sample IDs: first sample (1), second sample (2)
+        attention_mask = torch.tensor([[1, 1, 1, 2, 2, 2]], dtype=torch.long)
+
+        mask = create_bidirectional_attention_mask(
+            input_ids, attention_mask, sample_packing=True
+        )
+
+        # Check that tokens within same sample can attend to each other
+        # but not across samples
+        assert mask[0, 0, 0, 1].item()  # First sample tokens can attend to each other
+        assert mask[0, 0, 1, 2].item()
+        assert not mask[0, 0, 0, 3].item()  # Can't attend across samples
+        assert not mask[0, 0, 2, 4].item()
+        assert mask[0, 0, 3, 4].item()  # Second sample tokens can attend to each other
+
+    def test_compute_loss_basic(self, diffusion_trainer_instance):
+        """Test basic loss computation."""
+        # Mock model that returns logits
+        mock_model = Mock()
+        mock_outputs = Mock()
+        vocab_size = 1000
+        seq_len = 5
+        mock_outputs.logits = torch.randn(1, seq_len, vocab_size, requires_grad=True)
+        mock_model.return_value = mock_outputs
+        mock_model.training = True
+
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+
+        loss, outputs = diffusion_trainer_instance._compute_diffusion_loss(
+            mock_model, input_ids
+        )
+
+        # Check that loss is computed
+        assert isinstance(loss, torch.Tensor)
+        assert loss.requires_grad
+        assert outputs == mock_outputs
+
+        # Check that metrics were stored
+        diffusion_trainer_instance.store_metrics.assert_called_once()
+
+    def test_compute_loss_sft(self, diffusion_trainer_instance):
+        """Test loss computation with SFT labels."""
+        # Mock model
+        mock_model = Mock()
+        mock_outputs = Mock()
+        vocab_size = 1000
+        seq_len = 5
+        mock_outputs.logits = torch.randn(1, seq_len, vocab_size, requires_grad=True)
+        mock_model.return_value = mock_outputs
+        mock_model.training = True
+        diffusion_trainer_instance.axolotl_cfg.datasets = Mock()
+
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+        labels = torch.tensor([[-100, -100, 20, 30, 2]], dtype=torch.long)
+
+        loss, _ = diffusion_trainer_instance._compute_diffusion_loss(
+            mock_model, input_ids, labels=labels
+        )
+
+        # Check that loss is computed
+        assert isinstance(loss, torch.Tensor)
+        assert loss.requires_grad
+
+        # Check that SFT metrics were added
+        call_args = diffusion_trainer_instance.store_metrics.call_args[0][0]
+        assert "answer_ratio" in call_args
+        assert "avg_answer_length" in call_args
+
+    def test_compute_loss_no_masked_tokens(self, diffusion_trainer_instance):
+        """Test loss computation when no tokens are masked."""
+        # Mock model
+        mock_model = Mock()
+        mock_outputs = Mock()
+        vocab_size = 1000
+        seq_len = 3
+        mock_outputs.logits = torch.randn(1, seq_len, vocab_size)
+        mock_model.return_value = mock_outputs
+        mock_model.training = True
+
+        # Only special tokens (which won't be masked)
+        input_ids = torch.tensor([[1, 0, 2]], dtype=torch.long)
+
+        loss, _ = diffusion_trainer_instance._compute_diffusion_loss(
+            mock_model, input_ids
+        )
+
+        # Loss should be zero when no tokens are masked
+        assert loss.item() == 0.0
+        assert loss.requires_grad
+
+    def test_cache_special_token_ids(self, mock_tokenizer):
+        """Test caching of special token IDs."""
+        trainer = object.__new__(DiffusionTrainer)
+        trainer.processing_class = mock_tokenizer
+        trainer._cache_special_token_ids()
+        assert trainer._special_token_ids == {0, 1, 2}
+
+    def test_cache_special_token_ids_no_tokenizer(self):
+        """Test caching when no tokenizer is available."""
+        trainer = object.__new__(DiffusionTrainer)
+        trainer.processing_class = None
+        trainer._cache_special_token_ids()
+
+        assert trainer._special_token_ids == set()
+
+    def test_main_compute_loss_interface(self, diffusion_trainer_instance):
+        """Test the main compute_loss interface."""
+        # Mock model
+        mock_model = Mock()
+        mock_outputs = Mock()
+        mock_outputs.logits = torch.randn(1, 5, 1000)
+        mock_model.return_value = mock_outputs
+        mock_model.training = True
+
+        inputs = {
+            "input_ids": torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long),
+            "attention_mask": torch.tensor([[1, 1, 1, 1, 1]], dtype=torch.long),
+            "labels": torch.tensor([[-100, -100, 20, 30, 2]], dtype=torch.long),
+        }
+
+        # Test without return_outputs
+        loss = diffusion_trainer_instance.compute_loss(mock_model, inputs)
+        assert isinstance(loss, torch.Tensor)
+
+        # Test with return_outputs
+        loss, outputs = diffusion_trainer_instance.compute_loss(
+            mock_model, inputs, return_outputs=True
+        )
+        assert isinstance(loss, torch.Tensor)
+        assert outputs == mock_outputs
+
+    def test_missing_input_ids_raises_error(self, diffusion_trainer_instance):
+        """Test that missing input_ids raises ValueError."""
+        mock_model = Mock()
+        inputs = {"attention_mask": torch.tensor([[1, 1, 1]])}
+
+        with pytest.raises(ValueError, match="input_ids is required"):
+            diffusion_trainer_instance.compute_loss(mock_model, inputs)
diff --git a/tests/integrations/test_diffusion_callback.py b/tests/integrations/test_diffusion_callback.py
new file mode 100644
index 0000000000..8ba6ac5a13
--- /dev/null
+++ b/tests/integrations/test_diffusion_callback.py
@@ -0,0 +1,92 @@
+"""Tests for diffusion generation callback dataloader selection and triggering."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from axolotl.integrations.diffusion import DiffusionGenerationCallback
+
+
+class DummyTrainer:
+    """Minimal trainer double with required attributes/methods for the callback."""
+
+    def __init__(self, use_eval: bool):
+        # Config used by callback
+        self.axolotl_cfg = SimpleNamespace(
+            diffusion=SimpleNamespace(
+                generation_interval=1,
+                num_generation_samples=1,
+                generation_max_length=32,
+                generation_steps=4,
+                generation_temperature=0.0,
+                mask_token_id=16,
+            ),
+            use_wandb=False,
+        )
+
+        # Model/tokenizer are passed through to generate_samples; not used here
+        self.model = Mock()
+        self.processing_class = Mock()
+
+        # Datasets and loaders
+        self.eval_dataset = object() if use_eval else None
+        self._train_loader = object()
+        self._eval_loader = object()
+
+        # State for world process check
+        self.state = SimpleNamespace(is_world_process_zero=True)
+
+        # Track which loader was requested
+        self.requested: list[str] = []
+
+    def get_train_dataloader(self):
+        self.requested.append("train")
+        return self._train_loader
+
+    def get_eval_dataloader(self):
+        self.requested.append("eval")
+        return self._eval_loader
+
+
+@pytest.mark.parametrize("use_eval", [False, True])
+def test_callback_uses_correct_dataloader(monkeypatch, use_eval):
+    trainer = DummyTrainer(use_eval=use_eval)
+    callback = DiffusionGenerationCallback(trainer)
+
+    captured = {}
+
+    # Patch generate_samples in the callback module's namespace
+    def fake_generate_samples(**kwargs):
+        captured["dataloader"] = kwargs.get("dataloader")
+        # Return one dummy sample to exercise logging path
+        return [
+            {
+                "original": "o",
+                "masked": "m",
+                "generated": "g",
+                "mask_ratio": 0.5,
+                "masked_tokens": 1,
+                "total_tokens": 2,
+            }
+        ]
+
+    monkeypatch.setattr(
+        "axolotl.integrations.diffusion.callbacks.generate_samples",
+        fake_generate_samples,
+    )
+
+    # Trigger at step 1 (interval=1)
+    args = SimpleNamespace()
+    state = SimpleNamespace(global_step=1)
+    control = SimpleNamespace()
+
+    callback.on_step_end(args=args, state=state, control=control)
+
+    # Assert the expected dataloader path was used
+    if use_eval:
+        assert trainer.requested[0] == "eval"
+        assert captured["dataloader"] is trainer._eval_loader
+    else:
+        assert trainer.requested[0] == "train"
+        assert captured["dataloader"] is trainer._train_loader
diff --git a/tests/integrations/test_expert_parallel.py b/tests/integrations/test_expert_parallel.py
new file mode 100644
index 0000000000..cd107a3613
--- /dev/null
+++ b/tests/integrations/test_expert_parallel.py
@@ -0,0 +1,880 @@
+"""Tests for the Expert-Parallel (DeepEP) integration."""
+
+import os
+import queue as queue_mod
+import socket
+import time
+from datetime import timedelta
+from importlib.util import find_spec
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from axolotl.integrations.expert_parallel import (
+    ExpertParallelArgs,
+    ExpertParallelPlugin,
+)
+from axolotl.integrations.expert_parallel.experts_fn import (
+    REGISTRY,
+    kernel_to_registered_name,
+    register_all,
+)
+from axolotl.integrations.expert_parallel.plugin import expert_shard_axis
+from axolotl.integrations.expert_parallel.shard import (
+    _detect_experts_modules,
+    _slice_expert_lora_param,
+    ep_adapter_load_local_shard,
+    shard_expert_weights,
+)
+
+
+def _build_qwen3moe_block(num_experts: int = 16, top_k: int = 4):
+    from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
+    from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock
+
+    cfg = Qwen3MoeConfig(
+        hidden_size=512,
+        moe_intermediate_size=1024,
+        num_experts=num_experts,
+        num_experts_per_tok=top_k,
+    )
+    return Qwen3MoeSparseMoeBlock(cfg)
+
+
+# --------------------------------------------------------------------------- #
+# Args validation
+# --------------------------------------------------------------------------- #
+
+
+class TestExpertParallelArgs:
+    def test_defaults(self):
+        a = ExpertParallelArgs()
+        assert a.expert_parallel_size == 1
+        assert a.expert_parallel_backend == "deep_ep"
+        assert a.expert_parallel_fallback_on_unsupported is True
+
+    def test_enabled(self):
+        a = ExpertParallelArgs(expert_parallel_size=2)
+        assert a.expert_parallel_size == 2
+
+    def test_invalid_size(self):
+        with pytest.raises(ValueError):
+            ExpertParallelArgs(expert_parallel_size=0)
+
+    def test_invalid_token_capacity(self):
+        with pytest.raises(ValueError):
+            ExpertParallelArgs(expert_parallel_token_capacity=0)
+
+    def test_token_capacity_none_and_positive(self):
+        assert ExpertParallelArgs().expert_parallel_token_capacity is None
+        assert (
+            ExpertParallelArgs(
+                expert_parallel_token_capacity=1024
+            ).expert_parallel_token_capacity
+            == 1024
+        )
+
+
+class TestKernelInference:
+    """Plugin auto-composes with user's chosen local kernel."""
+
+    def _infer(self, **cfg_kwargs):
+        from types import SimpleNamespace
+
+        return ExpertParallelPlugin._infer_local_kernel(SimpleNamespace(**cfg_kwargs))
+
+    def test_use_scattermoe_picks_scattermoe(self):
+        assert self._infer(use_scattermoe=True) == "scattermoe"
+
+    def test_experts_implementation_scattermoe_alone_does_NOT_pick_scattermoe(self):
+        # use_scattermoe is the source of truth; bare experts_implementation=scattermoe
+        # without the master flag falls through to the default kernel.
+        assert self._infer(experts_implementation="scattermoe") == "grouped_mm"
+
+    def test_use_scattermoe_overrides_experts_implementation_eager(self):
+        # If kernels validator hasn't run yet, use_scattermoe still wins.
+        assert (
+            self._infer(use_scattermoe=True, experts_implementation="eager")
+            == "scattermoe"
+        )
+
+    def test_grouped_mm_picks_grouped_mm(self):
+        assert self._infer(experts_implementation="grouped_mm") == "grouped_mm"
+
+    def test_batched_mm_picks_grouped_mm(self):
+        assert self._infer(experts_implementation="batched_mm") == "grouped_mm"
+
+    def test_eager_picks_eager(self):
+        assert self._infer(experts_implementation="eager") == "eager"
+
+    def test_default_picks_grouped_mm(self):
+        assert self._infer() == "grouped_mm"
+
+    def test_use_sonicmoe_picks_sonicmoe(self):
+        assert self._infer(use_sonicmoe=True) == "sonicmoe"
+
+
+# --------------------------------------------------------------------------- #
+# Registration
+# --------------------------------------------------------------------------- #
+
+
+class TestRegistration:
+    def test_kernel_name_mapping(self):
+        assert kernel_to_registered_name("eager") == "deep_ep"
+        assert kernel_to_registered_name("grouped_mm") == "deep_ep_grouped_mm"
+        assert kernel_to_registered_name("scattermoe") == "deep_ep_scattermoe"
+        assert kernel_to_registered_name("sonicmoe") == "deep_ep_sonicmoe"
+
+    def test_register_all_idempotent(self):
+        from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+
+        register_all()
+        register_all()  # should not error
+        for name in REGISTRY:
+            fn = ALL_EXPERTS_FUNCTIONS.get_interface(name, None)
+            assert fn is not None, f"{name} not registered"
+
+    def test_whitelist_patch_accepts_deep_ep_names(self):
+        from transformers.modeling_utils import PreTrainedModel
+
+        register_all()
+        m = PreTrainedModel.__new__(PreTrainedModel)
+
+        class _Cfg:
+            _experts_implementation = "deep_ep_grouped_mm"
+
+        m.config = _Cfg()
+        for name in REGISTRY:
+            assert PreTrainedModel.get_correct_experts_implementation(m, name) == name
+
+    def test_whitelist_patch_rejects_garbage(self):
+        from transformers.modeling_utils import PreTrainedModel
+
+        register_all()
+        m = PreTrainedModel.__new__(PreTrainedModel)
+        m.config = type("_C", (), {"_experts_implementation": "garbage"})()
+        with pytest.raises(ValueError):
+            PreTrainedModel.get_correct_experts_implementation(m, "garbage")
+
+
+# --------------------------------------------------------------------------- #
+# Module detection
+# --------------------------------------------------------------------------- #
+
+
+class TestExpertModuleDetection:
+    def test_detects_qwen3moe_experts(self):
+        block = _build_qwen3moe_block()
+        found = list(_detect_experts_modules(block))
+        assert len(found) == 1
+        name, module = found[0]
+        assert module is block.experts
+        assert module.gate_up_proj.dim() == 3
+        assert module.down_proj.dim() == 3
+
+    def test_skips_non_3d_modules(self):
+        # A regular linear layer should not be detected as Experts
+        m = torch.nn.Linear(8, 8)
+        m.gate_up_proj = m.weight  # 2D, doesn't match
+        m.down_proj = m.weight
+        found = list(_detect_experts_modules(m))
+        assert len(found) == 0
+
+
+# --------------------------------------------------------------------------- #
+# Sharding (single-rank == no-op)
+# --------------------------------------------------------------------------- #
+
+
+class TestShardingSingleRank:
+    """At world_size=1, sharding is a no-op."""
+
+    def setup_method(self):
+        if not dist.is_initialized():
+            os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+            os.environ.setdefault("MASTER_PORT", "29555")
+            os.environ.setdefault("RANK", "0")
+            os.environ.setdefault("WORLD_SIZE", "1")
+            dist.init_process_group(backend="gloo", rank=0, world_size=1)
+
+    def teardown_method(self):
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+    def test_no_op_at_world_size_1(self):
+        block = _build_qwen3moe_block(num_experts=16)
+        original_shape = tuple(block.experts.gate_up_proj.shape)
+        n = shard_expert_weights(block, dist.group.WORLD)
+        assert n == 0
+        assert tuple(block.experts.gate_up_proj.shape) == original_shape
+
+    def test_none_group_no_op(self):
+        block = _build_qwen3moe_block(num_experts=16)
+        n = shard_expert_weights(block, None)
+        assert n == 0
+
+
+# --------------------------------------------------------------------------- #
+# Plugin lifecycle (no DeepEP needed for these)
+# --------------------------------------------------------------------------- #
+
+
+class TestPluginLifecycle:
+    def test_get_input_args(self):
+        plugin = ExpertParallelPlugin()
+        assert (
+            plugin.get_input_args()
+            == "axolotl.integrations.expert_parallel.ExpertParallelArgs"
+        )
+
+    def test_pre_model_load_disabled_is_noop(self):
+        plugin = ExpertParallelPlugin()
+
+        class _Cfg:
+            expert_parallel_enabled = False
+
+        plugin.pre_model_load(_Cfg())  # should not raise / not register anything new
+
+    def _ep_cfg(self, **kw):
+        from types import SimpleNamespace
+
+        defaults = dict(
+            expert_parallel_size=2,
+            expert_parallel_fallback_on_unsupported=True,
+            experts_implementation=None,
+        )
+        defaults.update(kw)
+        return SimpleNamespace(**defaults)
+
+    @pytest.mark.skipif(find_spec("deep_ep") is None, reason="deep_ep not installed")
+    def test_pre_model_load_default_picks_grouped_mm(self):
+        cfg = self._ep_cfg()
+        ExpertParallelPlugin().pre_model_load(cfg)
+        assert cfg.experts_implementation == "deep_ep_grouped_mm"
+
+    @pytest.mark.skipif(find_spec("deep_ep") is None, reason="deep_ep not installed")
+    def test_pre_model_load_use_scattermoe_auto_composes(self):
+        cfg = self._ep_cfg(use_scattermoe=True, experts_implementation="scattermoe")
+        ExpertParallelPlugin().pre_model_load(cfg)
+        assert cfg.experts_implementation == "deep_ep_scattermoe"
+
+    @pytest.mark.skipif(find_spec("deep_ep") is None, reason="deep_ep not installed")
+    def test_pre_model_load_overrides_existing_eager(self):
+        cfg = self._ep_cfg(experts_implementation="eager")
+        ExpertParallelPlugin().pre_model_load(cfg)
+        assert cfg.experts_implementation == "deep_ep"
+
+    def test_disabled_by_default(self):
+        """expert_parallel_size=1 (default) means EP is off — pre_model_load no-ops."""
+        cfg = self._ep_cfg(expert_parallel_size=1)
+        ExpertParallelPlugin().pre_model_load(cfg)
+        assert cfg.experts_implementation is None
+
+    def test_pre_model_load_enabled_no_deep_ep_falls_back(self, monkeypatch):
+        import axolotl.integrations.expert_parallel.plugin as plugin_mod
+
+        monkeypatch.setattr(plugin_mod, "find_spec", lambda name: None)
+        cfg = self._ep_cfg()
+        ExpertParallelPlugin().pre_model_load(cfg)
+        assert cfg.experts_implementation is None
+
+    def test_pre_model_load_no_fallback_raises(self, monkeypatch):
+        import axolotl.integrations.expert_parallel.plugin as plugin_mod
+
+        monkeypatch.setattr(plugin_mod, "find_spec", lambda name: None)
+        cfg = self._ep_cfg(expert_parallel_fallback_on_unsupported=False)
+        with pytest.raises(ImportError):
+            ExpertParallelPlugin().pre_model_load(cfg)
+
+
+# --------------------------------------------------------------------------- #
+# Mesh-axis topology — rank assignments for EP+FSDP composition.
+#
+# The 4-rank case (ep=2 × dp_shard=2) is what we'll actually run on a 4× A100
+# box. Tested here in single-process by spawning gloo workers; no GPUs needed.
+# --------------------------------------------------------------------------- #
+
+
+def _ep_topology_worker(rank, world_size, ep_size, dp_shard_size, port, q):
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    dist.init_process_group(
+        backend="gloo",
+        rank=rank,
+        world_size=world_size,
+        timeout=timedelta(seconds=120),
+    )
+
+    try:
+        from types import SimpleNamespace
+
+        cfg = SimpleNamespace(
+            expert_parallel_size=ep_size,
+            dp_shard_size=dp_shard_size,
+            tensor_parallel_size=1,
+            context_parallel_size=1,
+        )
+        ep_group = ExpertParallelPlugin._resolve_ep_group(cfg)
+        ep_ranks = sorted(dist.get_process_group_ranks(ep_group))
+        # Also peek at the dp_shard slice if a 2D mesh was built.
+        mesh = ExpertParallelPlugin._device_mesh
+        dp_ranks = (
+            sorted(dist.get_process_group_ranks(mesh["dp_shard"].get_group()))
+            if mesh is not None and "dp_shard" in mesh.mesh_dim_names
+            else None
+        )
+        q.put((rank, ep_ranks, dp_ranks))
+    finally:
+        dist.destroy_process_group()
+        ExpertParallelPlugin._device_mesh = None
+
+
+def _ep_topology_worker_expects_error(
+    rank, world_size, ep_size, dp_shard_size, port, q
+):
+    """Variant that captures any ValueError from `_resolve_ep_group`.
+
+    Module-level so `mp.get_context("spawn")` can pickle it.
+    """
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    dist.init_process_group(
+        backend="gloo",
+        rank=rank,
+        world_size=world_size,
+        timeout=timedelta(seconds=120),
+    )
+    try:
+        from types import SimpleNamespace
+
+        cfg = SimpleNamespace(
+            expert_parallel_size=ep_size,
+            dp_shard_size=dp_shard_size,
+            tensor_parallel_size=1,
+            context_parallel_size=1,
+        )
+        err = None
+        try:
+            ExpertParallelPlugin._resolve_ep_group(cfg)
+        except ValueError as e:
+            err = str(e)
+        q.put((rank, err))
+    finally:
+        dist.destroy_process_group()
+
+
+def _ep_cp_topology_worker(rank, world_size, ep_size, cp_size, dp_shard_size, port, q):
+    """EP × CP (× dp_shard) topology probe: returns this rank's ep / cp / dp_shard group members."""
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    dist.init_process_group(
+        backend="gloo", rank=rank, world_size=world_size, timeout=timedelta(seconds=120)
+    )
+    try:
+        from types import SimpleNamespace
+
+        cfg = SimpleNamespace(
+            expert_parallel_size=ep_size,
+            dp_shard_size=dp_shard_size,
+            tensor_parallel_size=1,
+            context_parallel_size=cp_size,
+        )
+        ep_group = ExpertParallelPlugin._resolve_ep_group(cfg)
+        ep_ranks = sorted(dist.get_process_group_ranks(ep_group))
+        cp_group = ExpertParallelPlugin._resolve_cp_group(cfg)
+        cp_ranks = (
+            sorted(dist.get_process_group_ranks(cp_group))
+            if cp_group is not None
+            else None
+        )
+        mesh = ExpertParallelPlugin._device_mesh
+        dp_ranks = (
+            sorted(dist.get_process_group_ranks(mesh["dp_shard"].get_group()))
+            if mesh is not None and "dp_shard" in (mesh.mesh_dim_names or ())
+            else None
+        )
+        q.put((rank, ep_ranks, cp_ranks, dp_ranks))
+    finally:
+        dist.destroy_process_group()
+        ExpertParallelPlugin._device_mesh = None
+        ExpertParallelPlugin._device_mesh = None
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _collect_worker_results(procs, q, world_size, timeout=120):
+    """Collect one result per worker, bailing out early if a worker dies.
+
+    A bare ``q.get(timeout=...)`` blocks the full timeout even after a worker has
+    crashed without reporting; here we stop once all workers have exited so an
+    infra failure surfaces as a clear assertion instead of a silent hang.
+    """
+    results = []
+    deadline = time.monotonic() + timeout
+    while len(results) < world_size and time.monotonic() < deadline:
+        try:
+            results.append(q.get(timeout=5))
+        except queue_mod.Empty:
+            if all(not p.is_alive() for p in procs):
+                break
+    for p in procs:
+        p.join(timeout=20)
+    exitcodes = [p.exitcode for p in procs]
+    assert len(results) == world_size, (
+        f"only {len(results)}/{world_size} workers reported; exitcodes={exitcodes}"
+    )
+    assert all(code == 0 for code in exitcodes), f"worker exitcodes={exitcodes}"
+    return results
+
+
+def _spawn_topology_check(world_size, ep_size, dp_shard_size):
+    ctx = mp.get_context("spawn")
+    q = ctx.Queue()
+    port = _find_free_port()
+    procs = [
+        ctx.Process(
+            target=_ep_topology_worker,
+            args=(r, world_size, ep_size, dp_shard_size, port, q),
+        )
+        for r in range(world_size)
+    ]
+    for p in procs:
+        p.start()
+    results = [q.get(timeout=120) for _ in range(world_size)]
+    for p in procs:
+        p.join(timeout=20)
+        assert p.exitcode == 0, f"worker exited with {p.exitcode}"
+    return sorted(results, key=lambda r: r[0])
+
+
+def _spawn_ep_cp_check(world_size, ep_size, cp_size, dp_shard_size=1):
+    ctx = mp.get_context("spawn")
+    q = ctx.Queue()
+    port = _find_free_port()
+    procs = [
+        ctx.Process(
+            target=_ep_cp_topology_worker,
+            args=(r, world_size, ep_size, cp_size, dp_shard_size, port, q),
+        )
+        for r in range(world_size)
+    ]
+    for p in procs:
+        p.start()
+    results = [q.get(timeout=120) for _ in range(world_size)]
+    for p in procs:
+        p.join(timeout=20)
+        assert p.exitcode == 0, f"worker exited with {p.exitcode}"
+    return sorted(results, key=lambda r: r[0])
+
+
+class TestMeshTopology:
+    """The 4-rank EP+FSDP composition rank assignments."""
+
+    def test_world4_ep2_cp2_orthogonal(self):
+        """world=4, ep=2 × cp=2: experts shard on `ep`, sequence on `cp`. The two axes must be
+        orthogonal — every rank is in exactly one ep group and one cp group that intersect only at
+        that rank, and the groups tile the 4 ranks. (DeepEP not required — pure group topology.)"""
+        results = _spawn_ep_cp_check(world_size=4, ep_size=2, cp_size=2)
+        ep = {r: tuple(e) for r, e, c, d in results}
+        cp = {r: tuple(c) for r, e, c, d in results}
+        for r in range(4):
+            assert len(ep[r]) == 2 and len(cp[r]) == 2, (ep, cp)
+            assert set(ep[r]) & set(cp[r]) == {r}, (
+                r,
+                ep,
+                cp,
+            )  # orthogonal: meet only at self
+        assert len(set(ep.values())) == 2 and len(set(cp.values())) == 2, (ep, cp)
+        # union of all ep groups (and cp groups) tiles the world
+        assert set().union(*ep.values()) == {0, 1, 2, 3}
+        assert set().union(*cp.values()) == {0, 1, 2, 3}
+
+    def test_world4_ep2_dp2_orthogonal(self):
+        """At world=4 with ep=2 and dp_shard=2, EP groups must be strided
+        ({0,2}, {1,3}) and dp_shard groups contiguous ({0,1}, {2,3}).
+        """
+        results = _spawn_topology_check(world_size=4, ep_size=2, dp_shard_size=2)
+        # Build per-rank groupings from results.
+        ep_groups_by_rank = {r: tuple(eps) for r, eps, _ in results}
+        dp_groups_by_rank = {r: tuple(dps) for r, _, dps in results}
+
+        # EP groups (strided): {0,2} and {1,3}
+        assert ep_groups_by_rank[0] == (0, 2), ep_groups_by_rank
+        assert ep_groups_by_rank[1] == (1, 3), ep_groups_by_rank
+        assert ep_groups_by_rank[2] == (0, 2), ep_groups_by_rank
+        assert ep_groups_by_rank[3] == (1, 3), ep_groups_by_rank
+
+        # dp_shard groups (contiguous, matches accelerate): {0,1} and {2,3}
+        assert dp_groups_by_rank[0] == (0, 1), dp_groups_by_rank
+        assert dp_groups_by_rank[1] == (0, 1), dp_groups_by_rank
+        assert dp_groups_by_rank[2] == (2, 3), dp_groups_by_rank
+        assert dp_groups_by_rank[3] == (2, 3), dp_groups_by_rank
+
+    def test_world4_ep4_dp1_uses_world(self):
+        """ep_size == world_size short-circuits to dist.group.WORLD."""
+        results = _spawn_topology_check(world_size=4, ep_size=4, dp_shard_size=1)
+        for rank, ep_ranks, dp_ranks in results:
+            assert ep_ranks == [0, 1, 2, 3], (rank, ep_ranks)
+            assert dp_ranks is None  # no 2D mesh built
+
+    def _spawn_expects_error(self, ep_size, dp_shard_size, world_size=4):
+        ctx = mp.get_context("spawn")
+        q = ctx.Queue()
+        port = _find_free_port()
+        procs = [
+            ctx.Process(
+                target=_ep_topology_worker_expects_error,
+                args=(r, world_size, ep_size, dp_shard_size, port, q),
+            )
+            for r in range(world_size)
+        ]
+        for p in procs:
+            p.start()
+        return _collect_worker_results(procs, q, world_size)
+
+    def test_world4_ep2_dp1_invalid_product_raises(self):
+        """ep<world without dp_shard filling the rest must raise (product mismatch)."""
+        results = self._spawn_expects_error(ep_size=2, dp_shard_size=1)
+        for rank, err in results:
+            assert err is not None, (
+                f"rank {rank} did not raise; expected product mismatch"
+            )
+            assert "must equal" in err.lower() or "world_size" in err.lower(), err
+
+    def test_mesh_axis_product_mismatch_raises(self):
+        """world=4 with ep=2*dp=4 (product 8 != 4) raises clearly."""
+        results = self._spawn_expects_error(ep_size=2, dp_shard_size=4)
+        for rank, err in results:
+            assert err is not None, f"rank {rank} did not raise"
+            assert "must equal" in err.lower() or "world_size" in err.lower(), err
+
+
+class TestExpertLoraSlicing:
+    """Routed-expert LoRA must be sliced so each EP rank gets ITS OWN experts' adapter blocks, and the
+    sliced shape must imply the TRUE LoRA rank ``r`` (= sliced_packed_dim // E_local), not ``r*ep_size``.
+
+    The pure-EP bug left the adapter global-sized while the base was E_local, so
+    ``_unwrap_experts_lora`` read ``rank = lora_A.shape[0] // E_local = r*ep_size`` and reshaped the
+    adapter into a scrambled expert/rank layout (harmless only at step 0, where lora_B == 0). These are
+    pure tensor-op checks of the slice math — no dist / FSDP / DeepEP / NVFP4 / model needed.
+    """
+
+    @pytest.mark.parametrize(
+        "e_global,ep_size,r", [(8, 2, 2), (256, 8, 16), (16, 4, 3)]
+    )
+    def test_lora_A_slice_picks_local_experts(self, e_global, ep_size, r):
+        # lora_A peft layout [E*r, in], expert-major: expert e owns rows [e*r:(e+1)*r]; tag rows = e.
+        e_local, in_features = e_global // ep_size, 5
+        full = torch.zeros(e_global * r, in_features)
+        for e in range(e_global):
+            full[e * r : (e + 1) * r] = float(e)
+        for ep_rank in range(ep_size):
+            start, end = ep_rank * e_local, (ep_rank + 1) * e_local
+            lin = torch.nn.Linear(in_features, e_global * r, bias=False)
+            lin.weight = torch.nn.Parameter(full.clone())
+            assert _slice_expert_lora_param(lin, 0, e_global, start, end) == r
+            sl = lin.weight.data
+            # the implied rank from the SLICED shape is the true r (the bug would give r*ep_size)
+            assert sl.shape[0] // e_local == r
+            tags = sl.reshape(e_local, r, in_features)[:, 0, 0].long()
+            assert torch.equal(tags, torch.arange(start, end))
+
+    @pytest.mark.parametrize(
+        "e_global,ep_size,r", [(8, 2, 2), (256, 8, 16), (16, 4, 3)]
+    )
+    def test_lora_B_slice_picks_local_experts(self, e_global, ep_size, r):
+        # lora_B peft layout [out, r*E], rank-major [out, r, E]; tag column (k, e) = e.
+        e_local, out_features = e_global // ep_size, 5
+        full = torch.zeros(out_features, r, e_global)
+        for e in range(e_global):
+            full[:, :, e] = float(e)
+        full = full.reshape(out_features, r * e_global)
+        for ep_rank in range(ep_size):
+            start, end = ep_rank * e_local, (ep_rank + 1) * e_local
+            lin = torch.nn.Linear(e_global * r, out_features, bias=False)
+            lin.weight = torch.nn.Parameter(full.clone())
+            assert _slice_expert_lora_param(lin, 1, e_global, start, end) == r
+            sl = lin.weight.data
+            assert sl.shape[1] // e_local == r
+            tags = sl.reshape(out_features, r, e_local)[0, 0, :].long()
+            assert torch.equal(tags, torch.arange(start, end))
+
+    @pytest.mark.parametrize(
+        "e_global,ep_size,r", [(8, 2, 2), (256, 8, 16), (16, 4, 3)]
+    )
+    def test_forward_slice_picks_local_experts_and_true_rank(
+        self, e_global, ep_size, r
+    ):
+        """The FORWARD-time slice (used by the ParamWrapper fastpath and the _unwrap fallback) keeps
+        the adapter global and takes each rank's expert block at use time, reporting rank=r — NOT the
+        r*ep_size that the un-sliced global adapter would imply (the actual pure-EP bug path)."""
+        from types import SimpleNamespace
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+            _ep_local_expert_lora,
+        )
+
+        e_local, in_f, out_f = e_global // ep_size, 5, 7
+        A = torch.zeros(e_global * r, in_f)
+        for e in range(e_global):
+            A[e * r : (e + 1) * r] = float(e)
+        B = torch.zeros(out_f, r, e_global)
+        for e in range(e_global):
+            B[:, :, e] = float(e)
+        B = B.reshape(out_f, r * e_global)
+        for ep_rank in range(ep_size):
+            offset = ep_rank * e_local
+            experts = SimpleNamespace(
+                num_experts=e_local,
+                num_experts_global=e_global,
+                local_expert_offset=offset,
+            )
+            a, b, n_local, rank = _ep_local_expert_lora(A, B, experts)
+            assert n_local == e_local and rank == r
+            assert torch.equal(
+                a.reshape(e_local, r, in_f)[:, 0, 0].long(),
+                torch.arange(offset, offset + e_local),
+            )
+            assert torch.equal(
+                b.reshape(out_f, r, e_local)[0, 0, :].long(),
+                torch.arange(offset, offset + e_local),
+            )
+
+    def test_forward_slice_noop_when_not_ep_sharded(self):
+        from types import SimpleNamespace
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+            _ep_local_expert_lora,
+        )
+
+        a_in, b_in = torch.randn(8 * 2, 5), torch.randn(7, 2 * 8)
+        experts = SimpleNamespace(
+            num_experts=8, num_experts_global=8, local_expert_offset=0
+        )
+        a, b, n_local, rank = _ep_local_expert_lora(a_in, b_in, experts)
+        assert n_local == 8 and rank == 2
+        assert a is a_in and b is b_in  # no slice / copy when E_local == E_global
+
+    def test_ranks_reconstruct_global_adapter(self):
+        # All EP ranks' slices, concatenated on the expert axis, must rebuild the global adapter
+        # (no expert dropped or duplicated).
+        e_global, ep_size, r, in_features, out_features = 8, 2, 2, 5, 7
+        e_local = e_global // ep_size
+        A = torch.randn(e_global * r, in_features)
+        B = torch.randn(out_features, r * e_global)
+        a_pieces, b_pieces = [], []
+        for ep_rank in range(ep_size):
+            start, end = ep_rank * e_local, (ep_rank + 1) * e_local
+            la = torch.nn.Linear(in_features, e_global * r, bias=False)
+            la.weight = torch.nn.Parameter(A.clone())
+            _slice_expert_lora_param(la, 0, e_global, start, end)
+            a_pieces.append(la.weight.data)
+            lb = torch.nn.Linear(e_global * r, out_features, bias=False)
+            lb.weight = torch.nn.Parameter(B.clone())
+            _slice_expert_lora_param(lb, 1, e_global, start, end)
+            # rank-major [out, r, E_local] per rank -> stack on the expert axis
+            b_pieces.append(lb.weight.data.reshape(out_features, r, e_local))
+        assert torch.equal(torch.cat(a_pieces, dim=0), A)
+        assert torch.equal(
+            torch.cat(b_pieces, dim=2).reshape(out_features, r * e_global),
+            B.reshape(out_features, r, e_global).reshape(out_features, r * e_global),
+        )
+
+
+class TestExpertAdapterLoadSharding:
+    """The cpu_ram_efficient load reconstructs each rank's local FSDP shard of the routed-expert LoRA
+    adapter from rank-0's broadcast GLOBAL (all-experts) adapter — the inverse of shard_expert_lora +
+    the dp/cp FSDP sharding (``ep_adapter_load_local_shard``). The ep slice must be EXPERT-aware:
+    lora_B's experts are the last axis of its ``[out, r, E]`` view (NOT contiguous in the flat ``r*E``
+    dim), so a plain chunk on dim 1 would load a rank-component instead of this ep-group's experts.
+    The 1-step e2e loss can't catch this (lora_B is zero-initialized — wrong zeros are still zeros); it
+    only bites a non-zero / resumed adapter. Pure tensor-op checks — no dist / FSDP / model needed.
+    """
+
+    @pytest.mark.parametrize(
+        "e_global,ep_size,dp_size,r", [(8, 2, 2, 2), (256, 4, 2, 16), (16, 2, 1, 3)]
+    )
+    def test_lora_A_load_shards_match_forward_slice(
+        self, e_global, ep_size, dp_size, r
+    ):
+        from torch.distributed.tensor import Shard
+
+        in_features, e_local = 5, e_global // ep_size
+        # global lora_A [E_global*r, in], expert-major; expert e's row block tagged = e
+        g = torch.zeros(e_global * r, in_features)
+        for e in range(e_global):
+            g[e * r : (e + 1) * r] = float(e)
+        placements = (Shard(0),)  # FSDP shards dim 0
+        for ep_coord in range(ep_size):
+            shards = [
+                ep_adapter_load_local_shard(
+                    g, 0, e_global, ep_coord, ep_size, placements, dp_size, dp
+                )
+                for dp in range(dp_size)
+            ]
+            ep_full = torch.cat(shards, dim=0)  # gather this ep-group's dp shards
+            # == the forward E_local slice shard_expert_lora produces for this ep-group
+            start, end = ep_coord * e_local, (ep_coord + 1) * e_local
+            la = torch.nn.Linear(in_features, e_global * r, bias=False)
+            la.weight = torch.nn.Parameter(g.clone())
+            _slice_expert_lora_param(la, 0, e_global, start, end)
+            assert torch.equal(ep_full, la.weight.data)
+            tags = ep_full.reshape(e_local, r, in_features)[:, 0, 0].long()
+            assert torch.equal(tags, torch.arange(start, end))
+
+    @pytest.mark.parametrize(
+        "e_global,ep_size,dp_size,r", [(8, 2, 2, 2), (256, 4, 2, 16), (16, 2, 1, 3)]
+    )
+    def test_lora_B_load_shards_match_forward_slice(
+        self, e_global, ep_size, dp_size, r
+    ):
+        from torch.distributed.tensor import Shard
+
+        out_features, e_local = 6, e_global // ep_size
+        # global lora_B [out, r*E] viewed [out, r, E]; column (k, e) tagged = e
+        g = torch.zeros(out_features, r, e_global)
+        for e in range(e_global):
+            g[:, :, e] = float(e)
+        g = g.reshape(out_features, r * e_global)
+        placements = (
+            Shard(0),
+        )  # FSDP shards dim 0 (out); the ep slice is on the expert axis
+        for ep_coord in range(ep_size):
+            shards = [
+                ep_adapter_load_local_shard(
+                    g, 1, e_global, ep_coord, ep_size, placements, dp_size, dp
+                )
+                for dp in range(dp_size)
+            ]
+            ep_full = torch.cat(shards, dim=0)  # gather the dp (out) shards back
+            start, end = ep_coord * e_local, (ep_coord + 1) * e_local
+            lb = torch.nn.Linear(e_global * r, out_features, bias=False)
+            lb.weight = torch.nn.Parameter(g.clone())
+            _slice_expert_lora_param(lb, 1, e_global, start, end)
+            assert torch.equal(ep_full, lb.weight.data)
+            tags = ep_full.reshape(out_features, r, e_local)[0, 0, :].long()
+            assert torch.equal(tags, torch.arange(start, end))
+
+    def test_lora_B_plain_chunk_would_load_wrong_experts(self):
+        # Documents the bug the expert-aware slice fixes: a naive chunk(ep_size, dim=1) loads a
+        # rank-component (all experts of one r-slice), not this ep-group's experts.
+        from torch.distributed.tensor import Shard
+
+        e_global, ep_size, r, out = 8, 2, 2, 6
+        g = torch.zeros(out, r, e_global)
+        for e in range(e_global):
+            g[:, :, e] = float(e)
+        g = g.reshape(out, r * e_global)
+        correct = ep_adapter_load_local_shard(
+            g, 1, e_global, 1, ep_size, (Shard(0),), 1, 0
+        )
+        naive = g.chunk(ep_size, dim=1)[1]
+        assert not torch.equal(correct, naive)
+        e_local = e_global // ep_size
+        # correct loads ep-group 1's experts {4,5,6,7}; the naive chunk would not
+        got = set(correct.reshape(out, r, e_local)[0, 0, :].long().tolist())
+        assert got == {4, 5, 6, 7}
+
+
+class TestExpertShardAxis:
+    """`expert_shard_axis` picks the non-ep mesh axis the routed experts FSDP-shard on under EP
+    composition: dp_shard when present, else cp, else None (pure EP / no ep axis). Pure logic on the
+    mesh dim names — no dist / mesh object needed."""
+
+    @pytest.mark.parametrize(
+        "dim_names,expected",
+        [
+            (("ep", "dp_shard"), "dp_shard"),  # EP × dp_shard
+            (("ep", "cp"), "cp"),  # EP × cp
+            (("ep", "dp_shard", "cp"), "dp_shard"),  # both -> dp_shard preferred
+            (("ep",), None),  # pure EP at world_size: no secondary axis to pre-wrap on
+            (("dp_shard", "cp"), None),  # no ep axis -> not an EP composition
+            (("ep", "tp"), None),  # tp is not an expert-shard axis (EP×TP unsupported)
+            ((), None),
+            (None, None),
+        ],
+    )
+    def test_axis_selection(self, dim_names, expected):
+        assert expert_shard_axis(dim_names) == expected
+
+    def test_dp_shard_preferred_over_cp_regardless_of_order(self):
+        # order in the tuple must not change the dp_shard preference
+        assert expert_shard_axis(("ep", "cp", "dp_shard")) == "dp_shard"
+        assert expert_shard_axis(("ep", "dp_shard", "cp")) == "dp_shard"
+
+
+class TestEpLoraSaveGating:
+    """save_ep_lora_adapter must EP-gather the routed-expert adapter ONLY when it was physically
+    sliced to E_local (EP×dp_shard/cp composition, where shard_expert_lora sets _ep_lora_sharded).
+    Pure EP keeps the adapter global, so gathering would duplicate every expert ep_size times. This
+    checks the discriminator the save relies on: shard_expert_lora flags + slices composition adapters,
+    and a non-sharded wrapper carries no flag (so the save leaves it as-is)."""
+
+    def _make_wrapper_model(self, e_global, ep_size, r, ep_rank=0):
+        e_local = e_global // ep_size
+
+        class _Experts(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.num_experts_global = e_global
+                self.num_local_experts = e_local
+                self.local_expert_offset = ep_rank * e_local
+
+        class ParamWrapper(
+            torch.nn.Module
+        ):  # name matches _is_param_wrapper's fallback check
+            def __init__(self):
+                super().__init__()
+                self.base_layer = _Experts()
+                self.lora_A = torch.nn.ModuleDict(
+                    {"default": torch.nn.Linear(5, e_global * r, bias=False)}
+                )
+                self.lora_B = torch.nn.ModuleDict(
+                    {"default": torch.nn.Linear(e_global * r, 7, bias=False)}
+                )
+
+        class _Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.wrapper = ParamWrapper()
+
+        return _Model()
+
+    def test_composition_slice_sets_flag(self):
+        from axolotl.integrations.expert_parallel.shard import shard_expert_lora
+
+        e_global, ep_size, r = 8, 2, 2
+        e_local = e_global // ep_size
+        m = self._make_wrapper_model(e_global, ep_size, r)
+        n = shard_expert_lora(m, ep_size)
+        assert n == 2  # lora_A + lora_B sliced
+        assert (
+            m.wrapper._ep_lora_sharded is True
+        )  # the gather discriminator the save reads
+        # sliced to E_local: gather (ep_size copies) is what reconstructs E_global
+        assert m.wrapper.lora_A["default"].weight.shape[0] == e_local * r
+        assert m.wrapper.lora_B["default"].weight.shape[1] == r * e_local
+
+    def test_pure_ep_wrapper_has_no_flag(self):
+        # Pure EP never runs shard_expert_lora -> no _ep_lora_sharded -> save_ep_lora_adapter must NOT
+        # EP-gather (the adapter is already global E_global).
+        m = self._make_wrapper_model(8, 2, 2)
+        assert getattr(m.wrapper, "_ep_lora_sharded", False) is False
+        # ep_size == 1 is a no-op and never flags either
+        from axolotl.integrations.expert_parallel.shard import shard_expert_lora
+
+        assert shard_expert_lora(m, 1) == 0
+        assert getattr(m.wrapper, "_ep_lora_sharded", False) is False
diff --git a/tests/integrations/test_gemma4_moe.py b/tests/integrations/test_gemma4_moe.py
new file mode 100644
index 0000000000..e3408c7822
--- /dev/null
+++ b/tests/integrations/test_gemma4_moe.py
@@ -0,0 +1,958 @@
+"""Validation tests for Gemma 4 MoE compatibility with ScatterMoE and SonicMoE.
+
+Gemma 4 has a unique MoE architecture:
+- No separate SparseMoeBlock — MoE is embedded in the decoder layer
+- Hybrid MLP+MoE: dense MLP runs in parallel with sparse MoE, outputs summed
+- Custom router (Gemma4TextRouter): RMSNorm → scale → proj → softmax → topk → renorm → per_expert_scale
+- Router is `self.router` (not `self.gate`)
+- Experts use standard 3D param layout with @use_experts_implementation
+
+These tests validate that:
+1. ScatterMoE kernels produce correct output for Gemma4 expert layout
+2. ScatterMoE + LoRA produces correct output
+3. SonicMoE integration code handles Gemma4 routing correctly
+4. Weight layouts are compatible
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+# ============================================================================
+# Gemma4 reference implementation (extracted from transformers)
+# ============================================================================
+
+
+class Gemma4RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6, with_scale=True):
+        super().__init__()
+        self.eps = eps
+        self.with_scale = with_scale
+        if with_scale:
+            self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        variance = x.float().pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        if self.with_scale:
+            return (self.weight * x).to(x.dtype)
+        return x.to(x.dtype)
+
+
+class Gemma4TextRouter(nn.Module):
+    def __init__(self, hidden_size, num_experts, top_k, eps=1e-6):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.scalar_root_size = hidden_size**-0.5
+        self.eps = eps
+
+        self.norm = Gemma4RMSNorm(hidden_size, eps=eps, with_scale=False)
+        self.proj = nn.Linear(hidden_size, num_experts, bias=False)
+        self.scale = nn.Parameter(torch.ones(hidden_size))
+        self.per_expert_scale = nn.Parameter(torch.ones(num_experts))
+
+    def forward(self, hidden_states):
+        hidden_states = self.norm(hidden_states)
+        hidden_states = hidden_states * self.scale * self.scalar_root_size
+
+        expert_scores = self.proj(hidden_states.to(self.proj.weight.dtype))
+        router_probabilities = F.softmax(expert_scores, dim=-1)
+
+        top_k_weights, top_k_index = torch.topk(
+            router_probabilities, k=self.top_k, dim=-1
+        )
+
+        top_k_weights /= top_k_weights.sum(dim=-1, keepdim=True)
+        top_k_weights = top_k_weights * self.per_expert_scale[top_k_index]
+
+        return router_probabilities, top_k_weights, top_k_index
+
+
+class Gemma4TextExperts(nn.Module):
+    def __init__(self, num_experts, hidden_size, intermediate_size, act_fn):
+        super().__init__()
+        self.num_experts = num_experts
+        self.hidden_dim = hidden_size
+        self.intermediate_dim = intermediate_size
+        self.gate_up_proj = nn.Parameter(
+            torch.empty(num_experts, 2 * intermediate_size, hidden_size)
+        )
+        self.down_proj = nn.Parameter(
+            torch.empty(num_experts, hidden_size, intermediate_size)
+        )
+        self.act_fn = act_fn
+
+    def forward(self, hidden_states, top_k_index, top_k_weights):
+        final_hidden_states = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx[0]
+            if expert_idx == self.num_experts:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate, up = F.linear(current_state, self.gate_up_proj[expert_idx]).chunk(
+                2, dim=-1
+            )
+            current_hidden_states = self.act_fn(gate) * up
+            current_hidden_states = F.linear(
+                current_hidden_states, self.down_proj[expert_idx]
+            )
+            current_hidden_states = (
+                current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+            )
+            final_hidden_states.index_add_(
+                0, token_idx, current_hidden_states.to(final_hidden_states.dtype)
+            )
+
+        return final_hidden_states
+
+
+# ============================================================================
+# Test fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def gemma4_config():
+    """Small Gemma4 MoE config for testing."""
+    return {
+        "hidden_size": 128,
+        "num_experts": 8,
+        "top_k": 2,
+        "intermediate_size": 64,
+        "eps": 1e-6,
+    }
+
+
+@pytest.fixture
+def device():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda:0")
+
+
+@pytest.fixture
+def gemma4_moe_layer(gemma4_config, device):
+    """Create a Gemma4 MoE layer (router + experts) on GPU."""
+    from transformers.activations import ACT2FN
+
+    act_fn = ACT2FN["gelu_pytorch_tanh"]
+
+    router = Gemma4TextRouter(
+        hidden_size=gemma4_config["hidden_size"],
+        num_experts=gemma4_config["num_experts"],
+        top_k=gemma4_config["top_k"],
+        eps=gemma4_config["eps"],
+    )
+    experts = Gemma4TextExperts(
+        num_experts=gemma4_config["num_experts"],
+        hidden_size=gemma4_config["hidden_size"],
+        intermediate_size=gemma4_config["intermediate_size"],
+        act_fn=act_fn,
+    )
+
+    # Initialize weights
+    nn.init.kaiming_uniform_(experts.gate_up_proj)
+    nn.init.kaiming_uniform_(experts.down_proj)
+    nn.init.normal_(router.proj.weight, std=0.01)
+
+    router = router.to(device).to(torch.bfloat16)
+    experts = experts.to(device).to(torch.bfloat16)
+
+    return router, experts
+
+
+# ============================================================================
+# ScatterMoE Tests
+# ============================================================================
+
+
+class TestGemma4ScatterMoE:
+    """Test ScatterMoE kernel compatibility with Gemma4 expert layout."""
+
+    def test_scattermoe_experts_match_reference(
+        self, gemma4_moe_layer, gemma4_config, device
+    ):
+        """ScatterMoE kernel output matches reference expert computation."""
+        from transformers.activations import ACT2FN
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+            flatten_sort_count,
+            parallel_linear,
+        )
+
+        router, experts = gemma4_moe_layer
+        act_fn = ACT2FN["gelu_pytorch_tanh"]
+        T = 16  # num tokens
+        H = gemma4_config["hidden_size"]
+        K = gemma4_config["top_k"]
+        E = gemma4_config["num_experts"]
+
+        hidden_states = torch.randn(T, H, device=device, dtype=torch.bfloat16)
+
+        # Reference forward
+        _, top_k_weights, top_k_index = router(hidden_states)
+        ref_output = experts(hidden_states, top_k_index, top_k_weights)
+
+        # ScatterMoE forward
+        routing_weights = top_k_weights.to(hidden_states.dtype)
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+            top_k_index, num_experts=E
+        )
+
+        # gate_up_proj is [E, 2*inter, H], ScatterMoE expects transposed: [E, H, 2*I]
+        gate_up_weight = experts.gate_up_proj.transpose(2, 1)
+        down_weight = experts.down_proj.transpose(2, 1)
+
+        gates_h = parallel_linear(
+            hidden_states,
+            gate_up_weight,
+            K,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        gates, h = gates_h.chunk(2, dim=-1)
+        h = act_fn(gates) * h
+
+        scatter_output = parallel_linear(
+            h,
+            down_weight,
+            1,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True,
+            grouped_out=False,
+            gates=routing_weights,
+        )
+
+        # Allow bf16 tolerance
+        torch.testing.assert_close(scatter_output, ref_output, atol=1e-2, rtol=1e-2)
+
+    def test_scattermoe_with_lora(self, gemma4_moe_layer, gemma4_config, device):
+        """ScatterMoE + LoRA kernel matches reference LoRA computation."""
+        from transformers.activations import ACT2FN
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+            flatten_sort_count,
+            parallel_linear,
+        )
+        from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_linear_lora import (
+            parallel_linear_lora,
+        )
+
+        router, experts = gemma4_moe_layer
+        act_fn = ACT2FN["gelu_pytorch_tanh"]
+        T = 16
+        H = gemma4_config["hidden_size"]
+        K = gemma4_config["top_k"]
+        E = gemma4_config["num_experts"]
+        inter = gemma4_config["intermediate_size"]
+        rank = 4
+        scaling = 0.5
+
+        hidden_states = torch.randn(T, H, device=device, dtype=torch.bfloat16)
+
+        # Create LoRA weights for gate_up_proj
+        # ScatterMoE layout: A=[r*E, K], B=[N, r*E]
+        lora_A_gup = (
+            torch.randn(rank * E, H, device=device, dtype=torch.bfloat16) * 0.01
+        )
+        lora_B_gup = (
+            torch.randn(2 * inter, rank * E, device=device, dtype=torch.bfloat16) * 0.01
+        )
+
+        # Reference: manual LoRA application per expert
+        _, top_k_weights, top_k_index = router(hidden_states)
+        ref_output = torch.zeros(T, H, device=device, dtype=torch.bfloat16)
+        with torch.no_grad():
+            expert_mask = F.one_hot(top_k_index, num_classes=E).permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+        for eidx in expert_hit:
+            eidx = eidx[0]
+            if eidx == E:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[eidx])
+            current_state = hidden_states[token_idx]
+
+            # Base gate_up + LoRA delta
+            base_out = F.linear(current_state, experts.gate_up_proj[eidx])
+            lora_a_slice = lora_A_gup[eidx * rank : (eidx + 1) * rank, :]
+            lora_b_slice = lora_B_gup[:, eidx * rank : (eidx + 1) * rank]
+            lora_delta = (
+                F.linear(F.linear(current_state, lora_a_slice), lora_b_slice) * scaling
+            )
+            combined = base_out + lora_delta
+
+            gate, up = combined.chunk(2, dim=-1)
+            h = act_fn(gate) * up
+            h = F.linear(h, experts.down_proj[eidx])
+            h = h * top_k_weights[token_idx, top_k_pos, None]
+            ref_output.index_add_(0, token_idx, h.to(ref_output.dtype))
+
+        # ScatterMoE LoRA forward
+        routing_weights = top_k_weights.to(hidden_states.dtype)
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+            top_k_index, num_experts=E
+        )
+
+        gate_up_weight = experts.gate_up_proj.transpose(2, 1)
+        down_weight = experts.down_proj.transpose(2, 1)
+
+        gates_h = parallel_linear_lora(
+            hidden_states,
+            gate_up_weight,
+            K,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            lora_A_gup,
+            lora_B_gup,
+            scaling,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        gates, h = gates_h.chunk(2, dim=-1)
+        h = act_fn(gates) * h
+
+        scatter_output = parallel_linear(
+            h,
+            down_weight,
+            1,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True,
+            grouped_out=False,
+            gates=routing_weights,
+        )
+
+        torch.testing.assert_close(scatter_output, ref_output, atol=5e-2, rtol=5e-2)
+
+    def test_gemma4_routing_correctness(self, gemma4_moe_layer, gemma4_config, device):
+        """Gemma4 custom routing (norm+scale+per_expert_scale) produces valid outputs."""
+        router, _ = gemma4_moe_layer
+        T = 32
+        H = gemma4_config["hidden_size"]
+        K = gemma4_config["top_k"]
+        E = gemma4_config["num_experts"]
+
+        hidden_states = torch.randn(T, H, device=device, dtype=torch.bfloat16)
+        router_probs, top_k_weights, top_k_index = router(hidden_states)
+
+        # Check shapes
+        assert router_probs.shape == (T, E)
+        assert top_k_weights.shape == (T, K)
+        assert top_k_index.shape == (T, K)
+
+        # Router probs should be valid probability distribution
+        assert (router_probs >= 0).all()
+        assert torch.allclose(
+            router_probs.sum(dim=-1),
+            torch.ones(T, device=device, dtype=router_probs.dtype),
+            atol=1e-3,
+        )
+
+        # Top-k indices should be valid expert indices
+        assert (top_k_index >= 0).all()
+        assert (top_k_index < E).all()
+
+        # Top-k weights should be non-negative (per_expert_scale can change sign though)
+        # Just verify finite
+        assert top_k_weights.isfinite().all()
+
+    def test_scattermoe_gradients_flow(self, gemma4_moe_layer, gemma4_config, device):
+        """Verify gradients flow through ScatterMoE kernels for Gemma4."""
+        from transformers.activations import ACT2FN
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+            flatten_sort_count,
+            parallel_linear,
+        )
+
+        router, experts = gemma4_moe_layer
+
+        # Enable grad for expert weights
+        experts.gate_up_proj.requires_grad_(True)
+        experts.down_proj.requires_grad_(True)
+
+        act_fn = ACT2FN["gelu_pytorch_tanh"]
+        T = 16
+        H = gemma4_config["hidden_size"]
+        K = gemma4_config["top_k"]
+        E = gemma4_config["num_experts"]
+
+        hidden_states = torch.randn(T, H, device=device, dtype=torch.bfloat16)
+
+        with torch.no_grad():
+            _, top_k_weights, top_k_index = router(hidden_states)
+
+        routing_weights = top_k_weights.to(hidden_states.dtype)
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+            top_k_index, num_experts=E
+        )
+
+        gate_up_weight = experts.gate_up_proj.transpose(2, 1)
+        down_weight = experts.down_proj.transpose(2, 1)
+
+        gates_h = parallel_linear(
+            hidden_states,
+            gate_up_weight,
+            K,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        gates, h = gates_h.chunk(2, dim=-1)
+        h = act_fn(gates) * h
+
+        output = parallel_linear(
+            h,
+            down_weight,
+            1,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True,
+            grouped_out=False,
+            gates=routing_weights,
+        )
+
+        loss = output.sum()
+        loss.backward()
+
+        assert experts.gate_up_proj.grad is not None
+        assert experts.down_proj.grad is not None
+        assert experts.gate_up_proj.grad.isfinite().all()
+        assert experts.down_proj.grad.isfinite().all()
+
+
+# ============================================================================
+# SonicMoE Tests
+# ============================================================================
+
+
+def _can_import_sonicmoe():
+    try:
+        from sonicmoe.enums import ActivationType  # noqa: F401
+
+        return True
+    except Exception:
+        return False
+
+
+class TestGemma4SonicMoE:
+    """Test SonicMoE compatibility with Gemma4.
+
+    SonicMoE requires Hopper/Blackwell GPU. Tests that need sonicmoe
+    import are skipped on unsupported GPUs.
+    """
+
+    def test_gemma4_weight_layout_compatible(self, gemma4_config):
+        """Verify Gemma4 expert weight layout is compatible with SonicMoE."""
+        E = gemma4_config["num_experts"]
+        H = gemma4_config["hidden_size"]
+        inter = gemma4_config["intermediate_size"]
+
+        gate_up_proj = torch.randn(E, 2 * inter, H)
+        down_proj = torch.randn(E, H, inter)
+
+        # SonicMoE expects [dim, dim, E] (experts last)
+        gate_up_sonic = gate_up_proj.permute(1, 2, 0)
+        down_sonic = down_proj.permute(1, 2, 0)
+
+        assert gate_up_sonic.shape == (2 * inter, H, E)
+        assert down_sonic.shape == (H, inter, E)
+
+        # Verify roundtrip
+        recovered_gate_up = gate_up_sonic.permute(2, 0, 1)
+        assert torch.equal(gate_up_proj, recovered_gate_up)
+
+    def test_gemma4_is_experts_only_model(self):
+        """Verify gemma4_text is recognized as experts-only model."""
+        from axolotl.integrations.kernels.constants import (
+            is_experts_only_model,
+            resolve_experts_class,
+        )
+
+        assert is_experts_only_model("gemma4_text")
+        cls = resolve_experts_class("gemma4_text")
+        assert cls is not None
+        assert cls.__name__ == "Gemma4TextExperts"
+
+
+# ============================================================================
+# Integration Tests (full layer with real model config)
+# ============================================================================
+
+
+class TestGemma4FullLayerIntegration:
+    """Test with realistic Gemma4 config (26B-A4B dimensions, single layer)."""
+
+    @pytest.fixture
+    def real_config(self):
+        return {
+            "hidden_size": 2816,
+            "num_experts": 128,
+            "top_k": 8,
+            "intermediate_size": 704,
+            "eps": 1e-6,
+        }
+
+    def test_scattermoe_real_dimensions(self, real_config, device):
+        """ScatterMoE works with real Gemma4-26B-A4B expert dimensions."""
+        from transformers.activations import ACT2FN
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+            flatten_sort_count,
+            parallel_linear,
+        )
+
+        act_fn = ACT2FN["gelu_pytorch_tanh"]
+        H = real_config["hidden_size"]
+        E = real_config["num_experts"]
+        K = real_config["top_k"]
+        inter = real_config["intermediate_size"]
+        T = 32
+
+        # Create experts on GPU
+        gate_up_proj = (
+            torch.randn(E, 2 * inter, H, device=device, dtype=torch.bfloat16) * 0.01
+        )
+        down_proj = torch.randn(E, H, inter, device=device, dtype=torch.bfloat16) * 0.01
+        hidden_states = torch.randn(T, H, device=device, dtype=torch.bfloat16)
+
+        # Simulate routing (random valid assignment)
+        top_k_index = torch.stack(
+            [torch.randperm(E, device=device)[:K] for _ in range(T)]
+        )
+        top_k_weights = torch.softmax(
+            torch.randn(T, K, device=device, dtype=torch.bfloat16), dim=-1
+        )
+
+        # Reference
+        ref_output = torch.zeros(T, H, device=device, dtype=torch.bfloat16)
+        with torch.no_grad():
+            expert_mask = F.one_hot(top_k_index, num_classes=E).permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+        for eidx in expert_hit:
+            eidx = eidx[0]
+            top_k_pos, token_idx = torch.where(expert_mask[eidx])
+            current_state = hidden_states[token_idx]
+            gate, up = F.linear(current_state, gate_up_proj[eidx]).chunk(2, dim=-1)
+            h = act_fn(gate) * up
+            h = F.linear(h, down_proj[eidx])
+            h = h * top_k_weights[token_idx, top_k_pos, None]
+            ref_output.index_add_(0, token_idx, h.to(ref_output.dtype))
+
+        # ScatterMoE
+        routing_weights = top_k_weights.to(hidden_states.dtype)
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+            top_k_index, num_experts=E
+        )
+
+        gates_h = parallel_linear(
+            hidden_states,
+            gate_up_proj.transpose(2, 1),
+            K,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        gates, h = gates_h.chunk(2, dim=-1)
+        h = act_fn(gates) * h
+
+        scatter_output = parallel_linear(
+            h,
+            down_proj.transpose(2, 1),
+            1,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True,
+            grouped_out=False,
+            gates=routing_weights,
+        )
+
+        torch.testing.assert_close(scatter_output, ref_output, atol=5e-2, rtol=5e-2)
+
+    def test_scattermoe_lora_real_dimensions(self, real_config, device):
+        """ScatterMoE + LoRA works with real Gemma4-26B-A4B dimensions."""
+        from transformers.activations import ACT2FN
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+            flatten_sort_count,
+            parallel_linear,
+        )
+        from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_linear_lora import (
+            parallel_linear_lora,
+        )
+
+        act_fn = ACT2FN["gelu_pytorch_tanh"]
+        H = real_config["hidden_size"]
+        E = real_config["num_experts"]
+        K = real_config["top_k"]
+        inter = real_config["intermediate_size"]
+        T = 32
+        rank = 8
+        scaling = 0.5
+
+        gate_up_proj = (
+            torch.randn(E, 2 * inter, H, device=device, dtype=torch.bfloat16) * 0.01
+        )
+        down_proj = torch.randn(E, H, inter, device=device, dtype=torch.bfloat16) * 0.01
+        lora_A = torch.randn(rank * E, H, device=device, dtype=torch.bfloat16) * 0.01
+        lora_B = (
+            torch.randn(2 * inter, rank * E, device=device, dtype=torch.bfloat16) * 0.01
+        )
+        hidden_states = torch.randn(T, H, device=device, dtype=torch.bfloat16)
+
+        # Random routing
+        top_k_index = torch.stack(
+            [torch.randperm(E, device=device)[:K] for _ in range(T)]
+        )
+        top_k_weights = torch.softmax(
+            torch.randn(T, K, device=device, dtype=torch.bfloat16), dim=-1
+        )
+
+        routing_weights = top_k_weights.to(hidden_states.dtype)
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(
+            top_k_index, num_experts=E
+        )
+
+        # ScatterMoE + LoRA on gate_up
+        gates_h = parallel_linear_lora(
+            hidden_states,
+            gate_up_proj.transpose(2, 1),
+            K,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            lora_A,
+            lora_B,
+            scaling,
+            grouped_in=False,
+            grouped_out=True,
+        )
+        gates, h = gates_h.chunk(2, dim=-1)
+        h = act_fn(gates) * h
+
+        output = parallel_linear(
+            h,
+            down_proj.transpose(2, 1),
+            1,
+            sorted_expert_idxs,
+            sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True,
+            grouped_out=False,
+            gates=routing_weights,
+        )
+
+        # Basic sanity: output should be finite and right shape
+        assert output.shape == (T, H)
+        assert output.isfinite().all()
+
+
+class TestExpertsInterfaceIntegration:
+    """Test the ExpertsInterface registration (the clean transformers hook)."""
+
+    @staticmethod
+    def _make_gemma4_config():
+        from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+
+        return Gemma4TextConfig(
+            hidden_size=128,
+            num_experts=8,
+            top_k_experts=2,
+            moe_intermediate_size=64,
+            hidden_activation="gelu_pytorch_tanh",
+            num_hidden_layers=2,
+            num_attention_heads=2,
+            num_key_value_heads=1,
+            head_dim=64,
+            intermediate_size=256,
+            enable_moe_block=True,
+        )
+
+    def test_register_scattermoe_in_experts_interface(self):
+        """register_scattermoe_experts adds 'scattermoe' to the global interface."""
+        from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+            register_scattermoe_experts,
+            scattermoe_experts_forward,
+        )
+
+        register_scattermoe_experts()
+
+        assert "scattermoe" in ALL_EXPERTS_FUNCTIONS
+        assert ALL_EXPERTS_FUNCTIONS["scattermoe"] is scattermoe_experts_forward
+
+    def test_experts_implementation_dispatches_to_scattermoe(self, device):
+        """Setting config._experts_implementation='scattermoe' dispatches correctly."""
+        from transformers.models.gemma4.modeling_gemma4 import (
+            Gemma4TextExperts as HFGemma4TextExperts,
+        )
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+            register_scattermoe_experts,
+        )
+
+        register_scattermoe_experts()
+
+        cfg = self._make_gemma4_config()
+        cfg._experts_implementation = "scattermoe"
+
+        with torch.device("meta"):
+            hf_experts = HFGemma4TextExperts(cfg)
+
+        hf_experts = hf_experts.to_empty(device=device)
+        nn.init.kaiming_uniform_(hf_experts.gate_up_proj)
+        nn.init.kaiming_uniform_(hf_experts.down_proj)
+        hf_experts = hf_experts.to(torch.bfloat16)
+
+        T, K = 16, 2
+        hidden_states = torch.randn(T, 128, device=device, dtype=torch.bfloat16)
+        top_k_index = torch.stack(
+            [torch.randperm(8, device=device)[:K] for _ in range(T)]
+        )
+        top_k_weights = torch.softmax(
+            torch.randn(T, K, device=device, dtype=torch.bfloat16), dim=-1
+        )
+
+        # Get reference output with eager implementation
+        cfg_eager = self._make_gemma4_config()
+        cfg_eager._experts_implementation = "eager"
+        with torch.device("meta"):
+            eager_experts = HFGemma4TextExperts(cfg_eager)
+        eager_experts = eager_experts.to_empty(device=device).to(torch.bfloat16)
+        # Copy weights from scattermoe experts
+        eager_experts.gate_up_proj.data.copy_(hf_experts.gate_up_proj.data)
+        eager_experts.down_proj.data.copy_(hf_experts.down_proj.data)
+
+        ref_output = eager_experts(hidden_states, top_k_index, top_k_weights)
+
+        # ScatterMoE dispatched output
+        scatter_output = hf_experts(hidden_states, top_k_index, top_k_weights)
+
+        torch.testing.assert_close(scatter_output, ref_output, atol=1e-2, rtol=1e-2)
+
+    def test_validation_accepts_scattermoe(self):
+        """get_correct_experts_implementation accepts 'scattermoe' after registration."""
+        from transformers.modeling_utils import PreTrainedModel
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+            register_scattermoe_experts,
+        )
+
+        register_scattermoe_experts()
+
+        # Should not raise
+        result = PreTrainedModel.get_correct_experts_implementation(None, "scattermoe")
+        assert result == "scattermoe"
+
+    def test_eager_still_works_after_registration(self, device):
+        """Registering scattermoe doesn't break eager dispatch."""
+        from transformers.models.gemma4.modeling_gemma4 import (
+            Gemma4TextExperts as HFGemma4TextExperts,
+        )
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+            register_scattermoe_experts,
+        )
+
+        register_scattermoe_experts()
+
+        cfg = self._make_gemma4_config()
+        cfg._experts_implementation = "eager"
+
+        with torch.device("meta"):
+            hf_experts = HFGemma4TextExperts(cfg)
+
+        hf_experts = hf_experts.to_empty(device=device)
+        nn.init.kaiming_uniform_(hf_experts.gate_up_proj)
+        nn.init.kaiming_uniform_(hf_experts.down_proj)
+        hf_experts = hf_experts.to(torch.bfloat16)
+
+        T, K = 16, 2
+        hidden_states = torch.randn(T, 128, device=device, dtype=torch.bfloat16)
+        top_k_index = torch.stack(
+            [torch.randperm(8, device=device)[:K] for _ in range(T)]
+        )
+        top_k_weights = torch.softmax(
+            torch.randn(T, K, device=device, dtype=torch.bfloat16), dim=-1
+        )
+
+        # Should use eager (original) forward without error
+        output = hf_experts(hidden_states, top_k_index, top_k_weights)
+        assert output.shape == (T, 128)
+        assert output.isfinite().all()
+
+
+class TestScatterMoEExpertsInterfaceMultiModel:
+    """Test that the registered scattermoe ExpertsInterface works across model types.
+
+    All @use_experts_implementation Experts classes share the same layout:
+    gate_up_proj [E, 2*inter, H], down_proj [E, H, inter], forward(hidden_states, top_k_index, top_k_weights).
+    """
+
+    MODEL_EXPERTS = [
+        (
+            "transformers.models.gemma4.modeling_gemma4",
+            "Gemma4TextExperts",
+            {
+                "hidden_size": 128,
+                "num_experts": 8,
+                "moe_intermediate_size": 64,
+                "hidden_activation": "gelu_pytorch_tanh",
+                "top_k_experts": 2,
+                "num_hidden_layers": 2,
+                "num_attention_heads": 2,
+                "num_key_value_heads": 1,
+                "head_dim": 64,
+                "intermediate_size": 256,
+                "enable_moe_block": True,
+            },
+            "transformers.models.gemma4.configuration_gemma4.Gemma4TextConfig",
+        ),
+        (
+            "transformers.models.qwen3_moe.modeling_qwen3_moe",
+            "Qwen3MoeExperts",
+            {
+                "hidden_size": 128,
+                "num_experts": 8,
+                "moe_intermediate_size": 64,
+                "hidden_act": "silu",
+                "num_experts_per_tok": 2,
+                "num_hidden_layers": 2,
+                "num_attention_heads": 2,
+                "num_key_value_heads": 1,
+                "intermediate_size": 256,
+            },
+            "transformers.models.qwen3_moe.configuration_qwen3_moe.Qwen3MoeConfig",
+        ),
+        (
+            "transformers.models.olmoe.modeling_olmoe",
+            "OlmoeExperts",
+            {
+                "hidden_size": 128,
+                "num_experts": 8,
+                "intermediate_size": 64,
+                "hidden_act": "silu",
+                "num_experts_per_tok": 2,
+                "num_hidden_layers": 2,
+                "num_attention_heads": 2,
+                "num_key_value_heads": 1,
+            },
+            "transformers.models.olmoe.configuration_olmoe.OlmoeConfig",
+        ),
+        (
+            "transformers.models.mixtral.modeling_mixtral",
+            "MixtralExperts",
+            {
+                "hidden_size": 128,
+                "num_local_experts": 8,
+                "intermediate_size": 64,
+                "hidden_act": "silu",
+                "num_experts_per_tok": 2,
+                "num_hidden_layers": 2,
+                "num_attention_heads": 2,
+                "num_key_value_heads": 1,
+            },
+            "transformers.models.mixtral.configuration_mixtral.MixtralConfig",
+        ),
+    ]
+
+    @pytest.fixture(
+        params=[m[1] for m in MODEL_EXPERTS], ids=[m[1] for m in MODEL_EXPERTS]
+    )
+    def model_setup(self, request, device):
+        """Create an Experts instance for each model type."""
+        import importlib
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.experts import (
+            register_scattermoe_experts,
+        )
+
+        register_scattermoe_experts()
+
+        for module_path, cls_name, cfg_kwargs, config_cls_path in self.MODEL_EXPERTS:
+            if cls_name == request.param:
+                # Import config class
+                config_module, config_class = config_cls_path.rsplit(".", 1)
+                config_cls = getattr(
+                    importlib.import_module(config_module), config_class
+                )
+                cfg = config_cls(**cfg_kwargs)
+
+                # Import experts class
+                module = importlib.import_module(module_path)
+                experts_cls = getattr(module, cls_name)
+
+                # Create eager reference
+                cfg_eager = config_cls(**cfg_kwargs)
+                cfg_eager._experts_implementation = "eager"
+                with torch.device("meta"):
+                    eager = experts_cls(cfg_eager)
+                eager = eager.to_empty(device=device).to(torch.bfloat16)
+                nn.init.kaiming_uniform_(eager.gate_up_proj)
+                nn.init.kaiming_uniform_(eager.down_proj)
+
+                # Create scattermoe version with same weights
+                cfg._experts_implementation = "scattermoe"
+                with torch.device("meta"):
+                    scatter = experts_cls(cfg)
+                scatter = scatter.to_empty(device=device).to(torch.bfloat16)
+                scatter.gate_up_proj.data.copy_(eager.gate_up_proj.data)
+                scatter.down_proj.data.copy_(eager.down_proj.data)
+
+                return (
+                    cls_name,
+                    eager,
+                    scatter,
+                    cfg_kwargs.get(
+                        "num_experts", cfg_kwargs.get("num_local_experts", 8)
+                    ),
+                )
+
+    def test_scattermoe_matches_eager(self, model_setup, device):
+        """ScatterMoE ExpertsInterface output matches eager for each model type."""
+        cls_name, eager, scatter, num_experts = model_setup
+        T, K = 16, 2
+
+        hidden_states = torch.randn(T, 128, device=device, dtype=torch.bfloat16)
+        top_k_index = torch.stack(
+            [torch.randperm(num_experts, device=device)[:K] for _ in range(T)]
+        )
+        top_k_weights = torch.softmax(
+            torch.randn(T, K, device=device, dtype=torch.bfloat16), dim=-1
+        )
+
+        ref_output = eager(hidden_states, top_k_index, top_k_weights)
+        scatter_output = scatter(hidden_states, top_k_index, top_k_weights)
+
+        torch.testing.assert_close(
+            scatter_output,
+            ref_output,
+            atol=1e-2,
+            rtol=1e-2,
+            msg=f"{cls_name}: ScatterMoE output doesn't match eager",
+        )
diff --git a/tests/integrations/test_kd_chat_template.py b/tests/integrations/test_kd_chat_template.py
new file mode 100644
index 0000000000..b828e6c3d0
--- /dev/null
+++ b/tests/integrations/test_kd_chat_template.py
@@ -0,0 +1,81 @@
+"""
+Test for KD chat template strategies
+"""
+
+from unittest.mock import Mock
+
+import pytest
+
+from axolotl.integrations.kd.chat_template import ChatTemplateStrategyWithKDv2
+
+
+class TestChatTemplateStrategyWithKDv2:
+    """Test v2 strategy correctly handles target_token_ids"""
+
+    @pytest.fixture
+    def v2_strategy(self):
+        """Create v2 strategy instance with mocked dependencies"""
+        # Mock prompter
+        mock_prompter = Mock()
+        mock_prompter.roles = {"user": "user", "assistant": "assistant"}
+        mock_prompter.chat_template_msg_variables = ["role", "content"]
+        mock_prompter.chat_template = "{{ messages }}"
+
+        # Mock tokenizer
+        mock_tokenizer = Mock()
+        mock_tokenizer.pad_token_id = 0
+        mock_tokenizer.eos_token_id = 2
+        mock_tokenizer.bos_token_id = 1
+        mock_tokenizer.eos_token = "<|endoftext|>"
+        mock_tokenizer.apply_chat_template = Mock(return_value=[1, 10, 20, 30, 2])
+        mock_tokenizer.encode = Mock(return_value=[2])
+
+        return ChatTemplateStrategyWithKDv2(
+            prompter=mock_prompter,
+            tokenizer=mock_tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            logprobs_field="logprobs",
+            gen_temperature=1.0,
+            kd_temperature=1.0,
+        )
+
+    def test_v2_prepare_kd_fields_adds_target_token_ids(self, v2_strategy):
+        """
+        Test that v2's _prepare_kd_fields hook adds target_token_ids.
+
+        Validates the Template Method pattern fix where v2 overrides
+        the hook to add target_token_ids before transform.
+        """
+        tokenized = {"input_ids": [1, 10, 20, 30, 2], "labels": [1, 10, 20, 30, 2]}
+        original = {"target_token_ids": [[10, 20], [30, 40]]}
+
+        result = v2_strategy._prepare_kd_fields(tokenized, original)
+
+        assert "target_token_ids" in result
+        assert result["target_token_ids"] == [[10, 20], [30, 40]]
+
+    def test_v2_prepare_kd_fields_handles_missing_field(self, v2_strategy):
+        """Test hook handles missing target_token_ids gracefully"""
+        tokenized = {"input_ids": [1, 10, 20, 30, 2], "labels": [1, 10, 20, 30, 2]}
+        original = {}
+
+        result = v2_strategy._prepare_kd_fields(tokenized, original)
+
+        assert "target_token_ids" not in result
+
+    def test_v2_transform_requires_target_token_ids(self, v2_strategy):
+        """
+        Test v2's transform fails without target_token_ids.
+
+        Validates the bug fix - transform expects target_token_ids
+        to be added by the hook.
+        """
+        sample = {
+            "input_ids": [1, 10, 20, 30, 2],
+            "labels": [1, 10, 20, 30, 2],
+            "logprobs": [[-0.1, -0.2], [-0.3, -0.4]],
+        }
+
+        with pytest.raises(KeyError, match="target_token_ids"):
+            v2_strategy.transform_logprobs(sample)
diff --git a/tests/integrations/test_kd_liger.py b/tests/integrations/test_kd_liger.py
new file mode 100644
index 0000000000..f4a1d07563
--- /dev/null
+++ b/tests/integrations/test_kd_liger.py
@@ -0,0 +1,149 @@
+"""
+Unit tests for LigerFusedLinearKLTopKLogprobFunction autograd contract.
+
+Regression tests for the has_aux=True contract violation in
+LigerFusedLinearKLTopKLogprobFunction: CE loss must contribute to the
+backward graph, not be silently dropped as aux.
+"""
+
+import torch
+
+from axolotl.integrations.kd.kernels.liger import LigerFusedLinearKLTopKLogprobLoss
+
+BATCH_SIZE = 1
+SEQ_LEN = 4
+VOCAB_SIZE = 16
+HIDDEN_DIM = 8
+TOP_K = 4
+
+
+def make_inputs(seed: int):
+    torch.manual_seed(seed)
+    student_hidden_states = torch.randn(
+        BATCH_SIZE, SEQ_LEN, HIDDEN_DIM, dtype=torch.float32, requires_grad=True
+    )
+    lm_head_weight = torch.randn(
+        VOCAB_SIZE, HIDDEN_DIM, dtype=torch.float32, requires_grad=True
+    )
+    target_token_ids = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LEN, TOP_K))
+    target_logprobs = torch.log_softmax(torch.randn(BATCH_SIZE, SEQ_LEN, TOP_K), dim=-1)
+    target_mask = torch.ones(BATCH_SIZE, SEQ_LEN, TOP_K, dtype=torch.bool)
+    true_labels = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LEN))
+    # Set an interior position to ignore_index. After the left-shift inside the
+    # kernel (true_labels = pad(...); true_labels[:, 1:]), index [0, 2] in the
+    # input becomes index 1 in the flattened CE labels, so this exercises
+    # ignore_index handling in F.cross_entropy.
+    true_labels[0, 2] = -100
+    return {
+        "student_hidden_states": student_hidden_states,
+        "lm_head_weight": lm_head_weight,
+        "target_token_ids": target_token_ids,
+        "target_logprobs": target_logprobs,
+        "target_mask": target_mask,
+        "true_labels": true_labels,
+    }
+
+
+def test_ce_only_gradient_flows_to_student_hidden_states():
+    """Regression: with weight_soft_loss=0, weight_hard_loss=1 the CE gradient must reach inputs.
+
+    Pre-fix, has_aux=True dropped CE from the backward graph and student_input.grad was all zero.
+    """
+    inputs = make_inputs(seed=0)
+    loss_fn = LigerFusedLinearKLTopKLogprobLoss(
+        weight_hard_loss=1.0,
+        weight_soft_loss=0.0,
+        temperature=1.0,
+        beta=0.0,
+        compiled=False,
+        chunk_size=2,
+        compute_ce_loss=True,
+    )
+    loss = loss_fn(
+        inputs["lm_head_weight"],
+        inputs["student_hidden_states"],
+        inputs["target_token_ids"],
+        inputs["target_logprobs"],
+        inputs["target_mask"],
+        inputs["true_labels"],
+    )
+    loss.backward()
+
+    assert inputs["student_hidden_states"].grad is not None
+    assert torch.isfinite(inputs["student_hidden_states"].grad).all()
+    assert inputs["student_hidden_states"].grad.abs().sum().item() > 0
+
+    assert inputs["lm_head_weight"].grad is not None
+    assert torch.isfinite(inputs["lm_head_weight"].grad).all()
+    assert inputs["lm_head_weight"].grad.abs().sum().item() > 0
+
+
+def test_kd_mix_gradient_changes_when_ce_weight_changes():
+    """Regression: in KD-mix mode, increasing weight_hard_loss must change the gradient.
+
+    Two runs with identical RNG, differing only in weight_hard_loss (0.0 vs 0.5). Pre-fix
+    both runs produced identical grads because CE was silently dropped from backward.
+    """
+    inputs_a = make_inputs(seed=42)
+    loss_fn_a = LigerFusedLinearKLTopKLogprobLoss(
+        weight_soft_loss=0.5,
+        weight_hard_loss=0.0,
+        temperature=1.0,
+        beta=0.0,
+        compiled=False,
+        chunk_size=2,
+        compute_ce_loss=True,
+    )
+    loss_a = loss_fn_a(
+        inputs_a["lm_head_weight"],
+        inputs_a["student_hidden_states"],
+        inputs_a["target_token_ids"],
+        inputs_a["target_logprobs"],
+        inputs_a["target_mask"],
+        inputs_a["true_labels"],
+    )
+    loss_a.backward()
+    grad_a_h = inputs_a["student_hidden_states"].grad.detach().clone()
+    grad_a_w = inputs_a["lm_head_weight"].grad.detach().clone()
+
+    inputs_b = make_inputs(seed=42)
+    loss_fn_b = LigerFusedLinearKLTopKLogprobLoss(
+        weight_soft_loss=0.5,
+        weight_hard_loss=0.5,
+        temperature=1.0,
+        beta=0.0,
+        compiled=False,
+        chunk_size=2,
+        compute_ce_loss=True,
+    )
+    loss_b = loss_fn_b(
+        inputs_b["lm_head_weight"],
+        inputs_b["student_hidden_states"],
+        inputs_b["target_token_ids"],
+        inputs_b["target_logprobs"],
+        inputs_b["target_mask"],
+        inputs_b["true_labels"],
+    )
+    loss_b.backward()
+    grad_b_h = inputs_b["student_hidden_states"].grad.detach().clone()
+    grad_b_w = inputs_b["lm_head_weight"].grad.detach().clone()
+
+    assert grad_a_h is not None
+    assert torch.isfinite(grad_a_h).all()
+    assert grad_b_h is not None
+    assert torch.isfinite(grad_b_h).all()
+
+    assert grad_a_w is not None
+    assert torch.isfinite(grad_a_w).all()
+    assert grad_b_w is not None
+    assert torch.isfinite(grad_b_w).all()
+
+    diff = (grad_b_h - grad_a_h).abs().sum().item()
+    assert not torch.allclose(grad_a_h, grad_b_h, atol=1e-6, rtol=1e-5)
+    assert diff > 1e-4, f"CE gradient contribution suspiciously small: {diff}"
+
+    diff_w = (grad_b_w - grad_a_w).abs().sum().item()
+    assert not torch.allclose(grad_a_w, grad_b_w, atol=1e-6, rtol=1e-5)
+    assert diff_w > 1e-4, (
+        f"CE weight-gradient contribution suspiciously small: {diff_w}"
+    )
diff --git a/tests/integrations/test_kd_trainer_direct_loss.py b/tests/integrations/test_kd_trainer_direct_loss.py
new file mode 100644
index 0000000000..ce465817f6
--- /dev/null
+++ b/tests/integrations/test_kd_trainer_direct_loss.py
@@ -0,0 +1,216 @@
+"""Tests for AxolotlKDTrainer.compute_loss."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+import torch.nn as nn
+
+# Skip the entire module if the KD trainer can't be imported (env not set up).
+axolotl_kd_trainer = pytest.importorskip("axolotl.integrations.kd.trainer")
+AxolotlKDTrainer = axolotl_kd_trainer.AxolotlKDTrainer
+_resolve_lm_head = axolotl_kd_trainer._resolve_lm_head
+
+
+HIDDEN = 8
+VOCAB = 16
+SEQ = 4
+BSZ = 2
+TOP_K = 3
+
+
+def test_resolve_lm_head_standard():
+    model = SimpleNamespace(lm_head=nn.Linear(HIDDEN, VOCAB, bias=False))
+    assert _resolve_lm_head(model) is model.lm_head
+
+
+def test_resolve_lm_head_multimodal():
+    lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+    language_model = SimpleNamespace(lm_head=lm_head)
+    model = SimpleNamespace(language_model=language_model)
+    assert _resolve_lm_head(model) is lm_head
+
+
+def test_resolve_lm_head_peft_wrapped():
+    lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+    inner = SimpleNamespace(lm_head=lm_head)
+    model = SimpleNamespace(get_base_model=lambda: inner)
+    assert _resolve_lm_head(model) is lm_head
+
+
+def test_resolve_lm_head_peft_wrapped_multimodal():
+    lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+    language_model = SimpleNamespace(lm_head=lm_head)
+    inner = SimpleNamespace(language_model=language_model)
+    model = SimpleNamespace(get_base_model=lambda: inner)
+    assert _resolve_lm_head(model) is lm_head
+
+
+def test_resolve_lm_head_missing_raises():
+    model = SimpleNamespace()
+    with pytest.raises(AttributeError, match="could not find lm_head"):
+        _resolve_lm_head(model)
+
+
+def _build_fake_model():
+    class TinyModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+
+        def forward(self, input_ids=None, output_hidden_states=False, **kw):
+            assert output_hidden_states is True
+            assert "labels" not in kw
+            assert "target_token_ids" not in kw
+            assert "target_logprobs" not in kw
+            assert "target_mask" not in kw
+            assert "num_items_in_batch" not in kw
+            hidden = torch.randn(BSZ, SEQ, HIDDEN, requires_grad=True)
+            return SimpleNamespace(
+                loss=None,
+                logits=None,
+                hidden_states=(hidden,),
+                past_key_values=None,
+                attentions=None,
+            )
+
+    return TinyModel()
+
+
+def _build_inputs(labels=None):
+    if labels is None:
+        labels = torch.randint(0, VOCAB, (BSZ, SEQ))
+    return {
+        "input_ids": torch.randint(0, VOCAB, (BSZ, SEQ)),
+        "labels": labels,
+        "target_token_ids": torch.randint(0, VOCAB, (BSZ, SEQ, TOP_K)),
+        "target_logprobs": torch.log_softmax(torch.randn(BSZ, SEQ, TOP_K), dim=-1),
+        "target_mask": torch.ones(BSZ, SEQ, TOP_K, dtype=torch.bool),
+    }
+
+
+def test_compute_loss_calls_kd_loss_with_correct_shapes():
+    kd_loss_fn = MagicMock(return_value=torch.tensor(2.0, requires_grad=True))
+    fake_self = SimpleNamespace(
+        args=SimpleNamespace(sample_packing=False),
+        model_accepts_loss_kwargs=True,
+        _kd_loss_fn=kd_loss_fn,
+    )
+    model = _build_fake_model()
+    inputs = _build_inputs()
+    expected_labels = inputs["labels"].clone()
+    expected_target_ids = inputs["target_token_ids"].clone()
+
+    loss = AxolotlKDTrainer.compute_loss(fake_self, model, inputs)
+
+    kd_loss_fn.assert_called_once()
+    args, kwargs = kd_loss_fn.call_args
+    assert args[0].shape == (VOCAB, HIDDEN)
+    assert args[1].shape == (BSZ, SEQ, HIDDEN)
+    assert args[2].shape == expected_target_ids.shape
+    assert "true_labels" in kwargs
+    assert kwargs["true_labels"].shape == expected_labels.shape
+    assert torch.isfinite(loss).all()
+
+
+def test_compute_loss_divides_by_num_items_in_batch_from_labels():
+    kd_loss_fn = MagicMock(return_value=torch.tensor(8.0))
+    fake_self = SimpleNamespace(
+        args=SimpleNamespace(sample_packing=False),
+        model_accepts_loss_kwargs=True,
+        _kd_loss_fn=kd_loss_fn,
+    )
+    model = _build_fake_model()
+    labels = torch.tensor([[1, 2, -100, 3], [4, -100, -100, -100]])
+    inputs = _build_inputs(labels=labels)
+
+    loss = AxolotlKDTrainer.compute_loss(fake_self, model, inputs)
+    assert torch.isclose(loss, torch.tensor(2.0))
+
+
+def test_compute_loss_uses_explicit_num_items_in_batch():
+    kd_loss_fn = MagicMock(return_value=torch.tensor(8.0))
+    fake_self = SimpleNamespace(
+        args=SimpleNamespace(sample_packing=False),
+        model_accepts_loss_kwargs=True,
+        _kd_loss_fn=kd_loss_fn,
+    )
+    model = _build_fake_model()
+    inputs = _build_inputs()
+
+    loss = AxolotlKDTrainer.compute_loss(fake_self, model, inputs, num_items_in_batch=2)
+    assert torch.isclose(loss, torch.tensor(4.0))
+
+
+def test_compute_loss_does_not_divide_when_zero_items():
+    kd_loss_fn = MagicMock(return_value=torch.tensor(8.0))
+    fake_self = SimpleNamespace(
+        args=SimpleNamespace(sample_packing=False),
+        model_accepts_loss_kwargs=True,
+        _kd_loss_fn=kd_loss_fn,
+    )
+    model = _build_fake_model()
+    labels = torch.full((BSZ, SEQ), -100)
+    inputs = _build_inputs(labels=labels)
+
+    loss = AxolotlKDTrainer.compute_loss(fake_self, model, inputs)
+    assert torch.isclose(loss, torch.tensor(8.0))
+
+
+def test_compute_loss_raises_when_kd_keys_missing():
+    kd_loss_fn = MagicMock(return_value=torch.tensor(1.0))
+    fake_self = SimpleNamespace(
+        args=SimpleNamespace(sample_packing=False),
+        model_accepts_loss_kwargs=True,
+        _kd_loss_fn=kd_loss_fn,
+    )
+    model = _build_fake_model()
+    inputs = {
+        "input_ids": torch.randint(0, VOCAB, (BSZ, SEQ)),
+        "labels": torch.randint(0, VOCAB, (BSZ, SEQ)),
+    }
+    with pytest.raises(KeyError, match="KD batch missing required keys"):
+        AxolotlKDTrainer.compute_loss(fake_self, model, inputs)
+
+
+def test_compute_loss_raises_when_hidden_states_missing():
+    kd_loss_fn = MagicMock(return_value=torch.tensor(1.0))
+    fake_self = SimpleNamespace(
+        args=SimpleNamespace(sample_packing=False),
+        model_accepts_loss_kwargs=True,
+        _kd_loss_fn=kd_loss_fn,
+    )
+
+    class NoHiddenStatesModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+
+        def forward(self, **kw):
+            return SimpleNamespace(
+                loss=None,
+                logits=None,
+                hidden_states=None,
+                past_key_values=None,
+                attentions=None,
+            )
+
+    inputs = _build_inputs()
+    with pytest.raises(RuntimeError, match="did not return hidden_states"):
+        AxolotlKDTrainer.compute_loss(fake_self, NoHiddenStatesModel(), inputs)
+
+
+def test_compute_loss_does_not_mutate_caller_inputs():
+    kd_loss_fn = MagicMock(return_value=torch.tensor(1.0))
+    fake_self = SimpleNamespace(
+        args=SimpleNamespace(sample_packing=False),
+        model_accepts_loss_kwargs=True,
+        _kd_loss_fn=kd_loss_fn,
+    )
+    model = _build_fake_model()
+    inputs = _build_inputs()
+    original_keys = set(inputs.keys())
+
+    AxolotlKDTrainer.compute_loss(fake_self, model, inputs)
+    assert set(inputs.keys()) == original_keys
diff --git a/tests/integrations/test_liger.py b/tests/integrations/test_liger.py
new file mode 100644
index 0000000000..6865306c9d
--- /dev/null
+++ b/tests/integrations/test_liger.py
@@ -0,0 +1,93 @@
+"""
+config validation tests for swiglu args
+"""
+
+from typing import Optional
+
+import pytest
+
+from axolotl.utils.config import prepare_plugins, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="minimal_liger_cfg")
+def fixture_cfg():
+    return DictDefault(
+        {
+            "base_model": "TinyLlama/TinyLlama-1.1B-Chat-v0.6",
+            "learning_rate": 0.000001,
+            "datasets": [
+                {
+                    "path": "mhenrichsen/alpaca_2k_test",
+                    "type": "alpaca",
+                }
+            ],
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "plugins": ["axolotl.integrations.liger.LigerPlugin"],
+        }
+    )
+
+
+class TestValidation:
+    """
+    Test the validation module for liger
+    """
+
+    _caplog: Optional[pytest.LogCaptureFixture] = None
+
+    @pytest.fixture(autouse=True)
+    def inject_fixtures(self, caplog):
+        caplog.set_level("WARNING")
+        self._caplog = caplog
+
+    def test_deprecated_swiglu(self, minimal_liger_cfg):
+        test_cfg = DictDefault(
+            {
+                "liger_swiglu": False,
+            }
+            | minimal_liger_cfg
+        )
+
+        with self._caplog.at_level("WARNING", logger="axolotl.integrations.liger.args"):
+            prepare_plugins(test_cfg)
+            updated_cfg = validate_config(test_cfg)
+            # TODO this test is brittle in CI
+            # assert (
+            #     "The 'liger_swiglu' argument is deprecated"
+            #     in self._caplog.records[0].message
+            # )
+            assert updated_cfg.liger_swiglu is None
+            assert updated_cfg.liger_glu_activation is False
+
+    def test_conflict_swiglu_ligergluactivation(self, minimal_liger_cfg):
+        test_cfg = DictDefault(
+            {
+                "liger_swiglu": False,
+                "liger_glu_activation": True,
+            }
+            | minimal_liger_cfg
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=r".*You cannot have both `liger_swiglu` and `liger_glu_activation` set.*",
+        ):
+            prepare_plugins(test_cfg)
+            validate_config(test_cfg)
+
+    def test_use_token_scaling_require_flce(self, minimal_liger_cfg):
+        test_cfg = DictDefault(
+            {
+                "liger_fused_linear_cross_entropy": False,
+                "liger_use_token_scaling": True,
+            }
+            | minimal_liger_cfg
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=r"`liger_use_token_scaling: true` requires `liger_fused_linear_cross_entropy` enabled.",
+        ):
+            prepare_plugins(test_cfg)
+            validate_config(test_cfg)
diff --git a/tests/integrations/test_liger_qwen_vl_rope_default.py b/tests/integrations/test_liger_qwen_vl_rope_default.py
new file mode 100644
index 0000000000..18b3fa8a1d
--- /dev/null
+++ b/tests/integrations/test_liger_qwen_vl_rope_default.py
@@ -0,0 +1,124 @@
+"""``cfg.liger_rope=None`` must resolve to ``True`` for Qwen-VL so the upstream fused (m-)rope kernel is installed."""
+
+from unittest.mock import patch
+
+import pytest
+
+
+@pytest.mark.parametrize(
+    "model_type",
+    [
+        "qwen2_vl",
+        "qwen2_5_vl",
+        "qwen3_vl",
+        "qwen3_vl_moe",
+        "qwen2_vl_text",
+        "qwen2_5_vl_text",
+        "qwen3_vl_text",
+        "qwen3_vl_moe_text",
+    ],
+)
+def test_liger_rope_auto_defaults_to_true_for_qwen_vl(model_type):
+    from axolotl.integrations.liger.plugin import LigerPlugin
+    from axolotl.utils.dict import DictDefault
+
+    cfg = DictDefault(
+        {
+            "model_config_type": model_type,
+            "liger_rope": None,
+            "liger_cross_entropy": False,
+            "liger_fused_linear_cross_entropy": True,
+            "liger_rms_norm": True,
+            "liger_layer_norm": False,
+            "liger_glu_activation": False,
+            "liger_use_token_scaling": False,
+            "torch_compile": False,
+            "base_model": "fake/path",
+            "trust_remote_code": False,
+        }
+    )
+
+    captured = {}
+
+    def _record(
+        rope: bool = True,
+        cross_entropy: bool = False,
+        fused_linear_cross_entropy: bool = True,
+        rms_norm: bool = True,
+        swiglu: bool = True,
+        layer_norm: bool = True,
+        model=None,
+    ):
+        captured.update(
+            rope=rope,
+            cross_entropy=cross_entropy,
+            fused_linear_cross_entropy=fused_linear_cross_entropy,
+            rms_norm=rms_norm,
+            swiglu=swiglu,
+            layer_norm=layer_norm,
+        )
+
+    from liger_kernel.transformers import monkey_patch as liger_mp
+
+    with patch.dict(
+        liger_mp.MODEL_TYPE_TO_APPLY_LIGER_FN,
+        {model_type: _record},
+        clear=False,
+    ):
+        LigerPlugin().pre_model_load(cfg)
+
+    assert captured.get("rope") is True, (
+        f"Expected rope=True default for {model_type}, got {captured.get('rope')}"
+    )
+
+
+def test_liger_rope_explicit_false_is_respected_for_qwen_vl():
+    from axolotl.integrations.liger.plugin import LigerPlugin
+    from axolotl.utils.dict import DictDefault
+
+    cfg = DictDefault(
+        {
+            "model_config_type": "qwen2_5_vl",
+            "liger_rope": False,
+            "liger_cross_entropy": False,
+            "liger_fused_linear_cross_entropy": True,
+            "liger_rms_norm": True,
+            "liger_layer_norm": False,
+            "liger_glu_activation": False,
+            "liger_use_token_scaling": False,
+            "torch_compile": False,
+            "base_model": "fake/path",
+            "trust_remote_code": False,
+        }
+    )
+
+    captured = {}
+
+    def _record(
+        rope: bool = True,
+        cross_entropy: bool = False,
+        fused_linear_cross_entropy: bool = True,
+        rms_norm: bool = True,
+        swiglu: bool = True,
+        layer_norm: bool = True,
+        model=None,
+    ):
+        captured.update(
+            rope=rope,
+            cross_entropy=cross_entropy,
+            fused_linear_cross_entropy=fused_linear_cross_entropy,
+            rms_norm=rms_norm,
+            swiglu=swiglu,
+            layer_norm=layer_norm,
+        )
+
+    from liger_kernel.transformers import monkey_patch as liger_mp
+
+    with patch.dict(
+        liger_mp.MODEL_TYPE_TO_APPLY_LIGER_FN,
+        {"qwen2_5_vl": _record},
+        clear=False,
+    ):
+        LigerPlugin().pre_model_load(cfg)
+
+    assert captured.get("rope") is False
diff --git a/tests/integrations/test_nemo_gym.py b/tests/integrations/test_nemo_gym.py
new file mode 100644
index 0000000000..83206043c1
--- /dev/null
+++ b/tests/integrations/test_nemo_gym.py
@@ -0,0 +1,943 @@
+"""Unit tests for NeMo Gym integration.
+
+Tests the core parsing, routing, reward, and plugin wiring logic
+without requiring a running NeMo Gym server or GPU.
+"""
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+
+class TestParseAgentResponse(unittest.TestCase):
+    """Tests for _parse_agent_response in multi_turn.py."""
+
+    def _parse(self, response, eos_token_id=2):
+        from axolotl.integrations.nemo_gym.multi_turn import _parse_agent_response
+
+        return _parse_agent_response(response, eos_token_id)
+
+    def test_empty_response_returns_defaults(self):
+        result = self._parse({})
+        assert result["prompt_ids"] == [2]
+        assert result["completion_ids"] == [2]
+        assert result["env_mask"] == [0]
+        assert result["reward"] == 0.0
+        assert result["num_turns"] == 0
+
+    def test_error_response_returns_defaults(self):
+        result = self._parse({"error": "something broke"})
+        assert result["reward"] == 0.0
+        assert result["num_turns"] == 0
+
+    def test_single_turn_function_call(self):
+        response = {
+            "response": {
+                "output": [
+                    {
+                        "type": "function_call",
+                        "name": "guess_word",
+                        "arguments": '{"guess": "crane"}',
+                        "call_id": "call_1",
+                        "prompt_token_ids": [10, 20, 30],
+                        "generation_token_ids": [40, 50],
+                        "generation_log_probs": [-0.1, -0.2],
+                    }
+                ]
+            },
+            "reward": 0.5,
+        }
+        result = self._parse(response)
+        assert result["prompt_ids"] == [10, 20, 30]
+        assert result["completion_ids"] == [40, 50]
+        assert result["env_mask"] == [1, 1]  # model tokens
+        assert result["logprobs"] == [-0.1, -0.2]
+        assert result["reward"] == 0.5
+        assert result["num_turns"] == 1
+
+    def test_multi_turn_preserves_env_mask(self):
+        """Second turn's prompt tokens (tool results) get env_mask=0."""
+        response = {
+            "response": {
+                "output": [
+                    {
+                        "type": "function_call",
+                        "prompt_token_ids": [10, 20],
+                        "generation_token_ids": [30, 31],
+                        "generation_log_probs": [-0.1, -0.2],
+                    },
+                    {
+                        "type": "function_call_output",
+                        "output": '{"feedback": "XYGXY"}',
+                    },
+                    {
+                        "type": "function_call",
+                        # prompt includes original + gen + tool output
+                        "prompt_token_ids": [10, 20, 30, 31, 100, 101, 102],
+                        "generation_token_ids": [40, 41],
+                        "generation_log_probs": [-0.3, -0.4],
+                    },
+                ]
+            },
+            "reward": 0.3,
+        }
+        result = self._parse(response)
+        assert result["prompt_ids"] == [10, 20]
+        # completion = gen1 + tool_result + gen2
+        assert result["completion_ids"] == [30, 31, 100, 101, 102, 40, 41]
+        # env_mask: gen1=model(1), tool=env(0), gen2=model(1)
+        assert result["env_mask"] == [1, 1, 0, 0, 0, 1, 1]
+        assert result["num_turns"] == 2
+
+    def test_empty_output_preserves_reward(self):
+        response = {
+            "response": {"output": []},
+            "reward": 0.42,
+        }
+        result = self._parse(response)
+        assert result["reward"] == 0.42
+
+    def test_message_only_output(self):
+        """A message with text but no function calls."""
+        response = {
+            "response": {
+                "output": [
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [
+                            {"type": "output_text", "text": "I'll guess crane."}
+                        ],
+                        "prompt_token_ids": [10, 20],
+                        "generation_token_ids": [30, 31, 32],
+                        "generation_log_probs": [-0.1, -0.2, -0.3],
+                    }
+                ]
+            },
+            "reward": 0.1,
+        }
+        result = self._parse(response)
+        assert result["num_turns"] == 1
+        assert result["completion_ids"] == [30, 31, 32]
+        assert result["env_mask"] == [1, 1, 1]
+
+
+class TestRewardEnv(unittest.TestCase):
+    """Tests for reward_env passthrough function."""
+
+    def test_with_list_rewards(self):
+        from axolotl.integrations.nemo_gym.rewards import reward_env
+
+        result = reward_env([["comp1"], ["comp2"]], env_reward=[0.5, 0.8])
+        assert result == [0.5, 0.8]
+
+    def test_with_scalar_reward(self):
+        from axolotl.integrations.nemo_gym.rewards import reward_env
+
+        result = reward_env([["comp1"], ["comp2"]], env_reward=0.7)
+        assert result == [0.7, 0.7]
+
+    def test_missing_reward_returns_zeros(self):
+        from axolotl.integrations.nemo_gym.rewards import reward_env
+
+        result = reward_env([["comp1"], ["comp2"]])
+        assert result == [0.0, 0.0]
+
+
+class TestRewardNemoGymVerify(unittest.TestCase):
+    """Tests for reward_nemo_gym_verify with mocked HTTP."""
+
+    @patch("axolotl.integrations.nemo_gym.rewards._get_verify_urls")
+    @patch("axolotl.integrations.nemo_gym.rewards.requests")
+    def test_calls_verify_endpoint(self, mock_requests, mock_get_urls):
+        from axolotl.integrations.nemo_gym.rewards import reward_nemo_gym_verify
+
+        mock_get_urls.return_value = {"wordle": "http://localhost:9999/verify"}
+        mock_resp = MagicMock()
+        mock_resp.ok = True
+        mock_resp.json.return_value = {"reward": 0.75}
+        mock_requests.post.return_value = mock_resp
+
+        result = reward_nemo_gym_verify(
+            completions=[[{"role": "assistant", "content": "crane"}]],
+            prompts=[[{"role": "user", "content": "Guess a word"}]],
+            resources_server_ref=[{"name": "wordle"}],
+            verify_extra=[{}],
+        )
+
+        assert result == [0.75]
+        mock_requests.post.assert_called_once()
+
+    @patch("axolotl.integrations.nemo_gym.rewards._get_verify_urls")
+    def test_missing_server_returns_zero(self, mock_get_urls):
+        from axolotl.integrations.nemo_gym.rewards import reward_nemo_gym_verify
+
+        mock_get_urls.return_value = {}
+
+        result = reward_nemo_gym_verify(
+            completions=[[{"role": "assistant", "content": "crane"}]],
+            prompts=[[{"role": "user", "content": "Guess"}]],
+            resources_server_ref=[{"name": "unknown_server"}],
+            verify_extra=[{}],
+        )
+        assert result == [0.0]
+
+
+class TestNormalizeHost(unittest.TestCase):
+    """Tests for server.py _normalize_host helper."""
+
+    def test_zero_addr_normalized(self):
+        from axolotl.integrations.nemo_gym.server import _normalize_host
+
+        assert _normalize_host("0.0.0.0") == "127.0.0.1"
+
+    def test_localhost_normalized(self):
+        from axolotl.integrations.nemo_gym.server import _normalize_host
+
+        assert _normalize_host("localhost") == "127.0.0.1"
+
+    def test_loopback_passthrough(self):
+        from axolotl.integrations.nemo_gym.server import _normalize_host
+
+        assert _normalize_host("127.0.0.1") == "127.0.0.1"
+
+    def test_custom_fallback(self):
+        from axolotl.integrations.nemo_gym.server import _normalize_host
+
+        assert _normalize_host("0.0.0.0", fallback="10.0.0.1") == "10.0.0.1"
+
+    def test_real_ip_passthrough(self):
+        from axolotl.integrations.nemo_gym.server import _normalize_host
+
+        assert _normalize_host("192.168.1.50") == "192.168.1.50"
+
+
+class TestDatasetLookupKeying(unittest.TestCase):
+    """Verify dataset lookup uses last message content as key."""
+
+    def test_single_message_prompt(self):
+        """Single-message prompt: [0] == [-1], both work."""
+        prompt = [{"role": "user", "content": "Play Wordle!"}]
+        assert prompt[0]["content"] == prompt[-1]["content"]
+
+    def test_multi_message_prompt_uses_last(self):
+        """Multi-message prompt: must use [-1] to match data_producer lookup."""
+        prompt = [
+            {"role": "system", "content": "You are a game player."},
+            {"role": "user", "content": "Play Wordle!"},
+        ]
+        # data_producer.py line 92 uses prompt[-1]
+        key = prompt[-1]["content"]
+        assert key == "Play Wordle!"
+        # Old code used prompt[0] which would be wrong here
+        assert prompt[0]["content"] != key
+
+
+class TestAgentRefPreservation(unittest.TestCase):
+    """Verify agent_ref is preserved through the dispatch chain."""
+
+    def test_data_producer_preserves_agent_ref(self):
+        """Simulates the data_producer lookup logic."""
+        # Simulate what plugin.py builds
+        dataset_lookup = {
+            "Play Wordle!": {
+                "prompt": [{"role": "user", "content": "Play Wordle!"}],
+                "agent_ref": {"name": "wordle_simple_agent"},
+                "verify_extra": {
+                    "responses_create_params": {
+                        "input": [{"role": "user", "content": "Play Wordle!"}]
+                    }
+                },
+            }
+        }
+
+        # Simulate data_producer.py logic (after fix)
+        prompt_text = "Play Wordle!"
+        full_item = dataset_lookup.get(prompt_text, {})
+        item = full_item.get("verify_extra", {})
+        if "agent_ref" in full_item and "agent_ref" not in item:
+            item["agent_ref"] = full_item["agent_ref"]
+
+        assert "agent_ref" in item
+        assert item["agent_ref"]["name"] == "wordle_simple_agent"
+
+    def test_multi_turn_preserves_agent_ref(self):
+        """Simulates the multi_turn.py dispatch logic."""
+        dataset_lookup = {
+            "Play Wordle!": {
+                "agent_ref": {"name": "wordle_simple_agent"},
+                "verify_extra": {
+                    "responses_create_params": {
+                        "input": [{"role": "user", "content": "Play Wordle!"}]
+                    }
+                },
+            }
+        }
+
+        # Simulate multi_turn.py logic (after fix)
+        prompt_str = "Play Wordle!"
+        full_item = None
+        for key, val in dataset_lookup.items():
+            if isinstance(key, str) and prompt_str == key:
+                full_item = val
+                break
+
+        dispatched = full_item.get("verify_extra", full_item)
+        if isinstance(dispatched, dict) and "agent_ref" not in dispatched:
+            agent_ref = full_item.get("agent_ref")
+            if agent_ref:
+                dispatched = {**dispatched, "agent_ref": agent_ref}
+
+        assert "agent_ref" in dispatched
+        assert dispatched["agent_ref"]["name"] == "wordle_simple_agent"
+
+
+class TestCallAgentsRouting(unittest.TestCase):
+    """Tests for _call_agents routing via agent_ref."""
+
+    def test_routes_to_correct_agent(self):
+        """Items with agent_ref should route to the matching agent server."""
+
+        agent_servers = {
+            "wordle_agent": "http://localhost:11111",
+            "math_agent": "http://localhost:22222",
+        }
+
+        items = [
+            {
+                "agent_ref": {"name": "wordle_agent"},
+                "responses_create_params": {
+                    "input": [{"role": "user", "content": "Play"}]
+                },
+            }
+        ]
+
+        # We can't actually call the agent, but verify the URL resolution
+        # by checking _call_agents builds the right request
+        # The function uses aiohttp — just verify agent_ref lookup works
+        item = items[0]
+        agent_ref = item.get("agent_ref", {})
+        agent_name = agent_ref.get("name", "")
+        agent_url = agent_servers.get(agent_name, "")
+        assert agent_url == "http://localhost:11111"
+
+    def test_fallback_to_first_agent(self):
+        """Items without agent_ref should use first available agent."""
+        agent_servers = {"default_agent": "http://localhost:33333"}
+        item = {
+            "responses_create_params": {"input": [{"role": "user", "content": "Hello"}]}
+        }
+        agent_ref = item.get("agent_ref", {})
+        agent_name = agent_ref.get("name", "")
+        agent_url = agent_servers.get(agent_name, "")
+        if not agent_url and agent_servers:
+            agent_url = next(iter(agent_servers.values()))
+        assert agent_url == "http://localhost:33333"
+
+
+class TestPluginDefaults(unittest.TestCase):
+    """Tests for plugin config enforcement."""
+
+    def test_dataloader_num_workers_forced_to_zero(self):
+        """Plugin should set dataloader_num_workers=0 for NeMo Gym."""
+
+        # Simulate the plugin logic
+        class FakeCfg:
+            dataloader_num_workers = 4
+            nemo_gym_multi_turn = True
+
+        cfg = FakeCfg()
+        # Replicate plugin.get_training_args logic
+        if getattr(cfg, "dataloader_num_workers", None) not in (None, 0):
+            pass  # would log warning
+        cfg.dataloader_num_workers = 0
+        assert cfg.dataloader_num_workers == 0
+
+    def test_dataloader_num_workers_none_stays_zero(self):
+        class FakeCfg:
+            dataloader_num_workers = None
+
+        cfg = FakeCfg()
+        cfg.dataloader_num_workers = 0
+        assert cfg.dataloader_num_workers == 0
+
+
+class TestSelectWeightSyncTransport(unittest.TestCase):
+    """Pure-logic table tests for ``select_weight_sync_transport``."""
+
+    def _caps(self, **kwargs):
+        from axolotl.integrations.nemo_gym.plugin import VLLMWeightSyncCapabilities
+
+        c = VLLMWeightSyncCapabilities(probed=True)
+        for k, v in kwargs.items():
+            setattr(c, k, v)
+        return c
+
+    def test_lora_with_native_endpoint(self):
+        from axolotl.integrations.nemo_gym.plugin import select_weight_sync_transport
+
+        caps = self._caps(lora_filesystem=True)
+        assert (
+            select_weight_sync_transport(caps, has_lora=True, vllm_lora_sync_pref=True)
+            == "lora_filesystem"
+        )
+
+    def test_lora_with_axolotl_endpoint(self):
+        from axolotl.integrations.nemo_gym.plugin import select_weight_sync_transport
+
+        caps = self._caps(lora_axolotl=True)
+        assert (
+            select_weight_sync_transport(caps, has_lora=True, vllm_lora_sync_pref=False)
+            == "lora_filesystem"
+        )
+
+    def test_lora_falls_back_to_nccl_when_no_lora_endpoint(self):
+        from axolotl.integrations.nemo_gym.plugin import select_weight_sync_transport
+
+        caps = self._caps(nccl=True)
+        assert (
+            select_weight_sync_transport(caps, has_lora=True, vllm_lora_sync_pref=False)
+            == "nccl"
+        )
+
+    def test_full_param_prefers_nccl(self):
+        from axolotl.integrations.nemo_gym.plugin import select_weight_sync_transport
+
+        caps = self._caps(nccl=True, http_full=True)
+        assert (
+            select_weight_sync_transport(
+                caps, has_lora=False, vllm_lora_sync_pref=False
+            )
+            == "nccl"
+        )
+
+    def test_full_param_falls_back_to_http(self):
+        from axolotl.integrations.nemo_gym.plugin import select_weight_sync_transport
+
+        caps = self._caps(http_full=True)
+        assert (
+            select_weight_sync_transport(
+                caps, has_lora=False, vllm_lora_sync_pref=False
+            )
+            == "http_full"
+        )
+
+    def test_full_param_no_routes_returns_none(self):
+        from axolotl.integrations.nemo_gym.plugin import select_weight_sync_transport
+
+        caps = self._caps()  # all False
+        assert (
+            select_weight_sync_transport(
+                caps, has_lora=False, vllm_lora_sync_pref=False
+            )
+            == "none"
+        )
+
+    def test_lora_no_routes_returns_none(self):
+        from axolotl.integrations.nemo_gym.plugin import select_weight_sync_transport
+
+        caps = self._caps()
+        assert (
+            select_weight_sync_transport(caps, has_lora=True, vllm_lora_sync_pref=True)
+            == "none"
+        )
+
+
+class TestProbeVllmWeightSync(unittest.TestCase):
+    """``probe_vllm_weight_sync`` reads a vLLM ``/openapi.json`` and reports caps."""
+
+    def test_stock_vllm_with_lora_enabled(self):
+        """Stock ``vllm serve --enable-lora`` exposes only LoRA endpoints."""
+        from unittest.mock import patch
+
+        from axolotl.integrations.nemo_gym.plugin import probe_vllm_weight_sync
+
+        spec = {
+            "paths": {
+                "/v1/models": {"get": {}},
+                "/v1/load_lora_adapter": {"post": {}},
+                "/v1/unload_lora_adapter": {"post": {}},
+                "/v1/completions": {"post": {}},
+            }
+        }
+        with patch("requests.get") as mock_get:
+            mock_get.return_value.raise_for_status = lambda: None
+            mock_get.return_value.json = lambda: spec
+            caps = probe_vllm_weight_sync("http://localhost:8000")
+
+        assert caps.probed is True
+        assert caps.lora_filesystem is True
+        assert caps.lora_axolotl is False
+        assert caps.nccl is False
+        assert caps.http_full is False
+
+    def test_axolotl_serve_lora_full_capabilities(self):
+        """``axolotl vllm-serve`` exposes NCCL + LoRA + HTTP full sync."""
+        from unittest.mock import patch
+
+        from axolotl.integrations.nemo_gym.plugin import probe_vllm_weight_sync
+
+        spec = {
+            "paths": {
+                "/init_communicator/": {"post": {}},
+                "/update_named_param/": {"post": {}},
+                "/batch_update_named_params/": {"post": {}},
+                "/set_lora_adapter/": {"post": {}},
+                "/clear_lora_adapter/": {"post": {}},
+                "/http_update_weights/": {"post": {}},
+                "/v1/load_lora_adapter": {"post": {}},
+            }
+        }
+        with patch("requests.get") as mock_get:
+            mock_get.return_value.raise_for_status = lambda: None
+            mock_get.return_value.json = lambda: spec
+            caps = probe_vllm_weight_sync("http://localhost:8000")
+
+        assert caps.probed is True
+        assert caps.nccl is True
+        assert caps.lora_axolotl is True
+        assert caps.lora_filesystem is True
+        assert caps.http_full is True
+
+    def test_trl_vllm_serve_nccl_only(self):
+        """``trl vllm-serve`` exposes NCCL routes but not LoRA filesystem."""
+        from unittest.mock import patch
+
+        from axolotl.integrations.nemo_gym.plugin import probe_vllm_weight_sync
+
+        spec = {
+            "paths": {
+                "/init_communicator/": {"post": {}},
+                "/update_named_param/": {"post": {}},
+                "/batch_update_named_params/": {"post": {}},
+                "/close_communicator/": {"post": {}},
+                "/generate/": {"post": {}},
+            }
+        }
+        with patch("requests.get") as mock_get:
+            mock_get.return_value.raise_for_status = lambda: None
+            mock_get.return_value.json = lambda: spec
+            caps = probe_vllm_weight_sync("http://localhost:8000")
+
+        assert caps.probed is True
+        assert caps.nccl is True
+        assert caps.lora_filesystem is False
+        assert caps.lora_axolotl is False
+        assert caps.http_full is False
+
+    def test_unreachable_server_records_error(self):
+        from unittest.mock import patch
+
+        from axolotl.integrations.nemo_gym.plugin import probe_vllm_weight_sync
+
+        with patch("requests.get") as mock_get:
+            mock_get.side_effect = ConnectionError("Connection refused")
+            caps = probe_vllm_weight_sync("http://localhost:9999")
+
+        assert caps.probed is False
+        assert caps.probe_error is not None
+        assert "ConnectionError" in caps.probe_error
+        assert caps.nccl is False
+        assert caps.lora_filesystem is False
+
+
+class TestPluginWeightSyncEnforcement(unittest.TestCase):
+    """End-to-end test of post_trainer_create's transport-selection branch.
+
+    The plugin used to silently no-op weight sync when ``vllm_lora_sync: false``,
+    leaving the trainer learning in isolation while vLLM kept serving the
+    unmodified base model. After the fix:
+
+      - LoRA + LoRA-loading endpoint   → installs filesystem LoRA sync
+      - LoRA + only NCCL endpoint      → uses NCCL broadcast
+      - Full FT + NCCL endpoint        → uses NCCL broadcast (standard TRL flow)
+      - Full FT + HTTP endpoint        → raises NotImplementedError (step 3)
+      - No usable transport            → raises ValueError with a precise diagnosis
+    """
+
+    @staticmethod
+    def _fake_cfg(adapter, vllm_lora_sync):
+        class FakeTRL:
+            pass
+
+        class FakeCfg:
+            pass
+
+        trl = FakeTRL()
+        trl.vllm_lora_sync = vllm_lora_sync
+        trl.vllm_server_host = "127.0.0.1"
+        trl.vllm_server_port = 8000
+
+        cfg = FakeCfg()
+        cfg.nemo_gym_enabled = True
+        cfg.nemo_gym_model_name = None
+        cfg.base_model = "test/model"
+        cfg.nemo_gym_verify_timeout = 30
+        cfg.nemo_gym_multi_turn = True
+        cfg.adapter = adapter
+        cfg.trl = trl
+        return cfg
+
+    @staticmethod
+    def _fake_trainer():
+        class FakeVLLMGen:
+            sync_weights = staticmethod(lambda: None)
+
+        class FakeTrainer:
+            vllm_generation = FakeVLLMGen()
+
+        return FakeTrainer()
+
+    @staticmethod
+    def _caps(**kwargs):
+        from axolotl.integrations.nemo_gym.plugin import VLLMWeightSyncCapabilities
+
+        c = VLLMWeightSyncCapabilities(probed=True)
+        for k, v in kwargs.items():
+            setattr(c, k, v)
+        return c
+
+    def test_lora_with_lora_endpoint_installs_filesystem_sync(self):
+        from unittest.mock import patch
+
+        from axolotl.integrations.nemo_gym.plugin import NemoGymPlugin
+
+        plugin = NemoGymPlugin()
+        plugin._vllm_caps = self._caps(lora_filesystem=True)
+        cfg = self._fake_cfg(adapter="lora", vllm_lora_sync=True)
+        trainer = self._fake_trainer()
+
+        with (
+            patch.object(plugin, "_setup_lora_sync") as setup,
+            patch.object(plugin, "_check_lora_endpoint") as check,
+            patch.object(plugin, "_wire_multi_turn") as wire,
+        ):
+            plugin.post_trainer_create(cfg, trainer)
+            setup.assert_called_once()
+            check.assert_called_once()
+            wire.assert_called_once()
+
+    def test_lora_with_no_routes_raises_with_lora_specific_message(self):
+        from axolotl.integrations.nemo_gym.plugin import NemoGymPlugin
+
+        plugin = NemoGymPlugin()
+        plugin._vllm_caps = self._caps()  # all False, but probed
+        cfg = self._fake_cfg(adapter="lora", vllm_lora_sync=False)
+        trainer = self._fake_trainer()
+
+        with self.assertRaises(ValueError) as ctx:
+            plugin.post_trainer_create(cfg, trainer)
+        msg = str(ctx.exception)
+        assert "no-op trainer" in msg
+        assert "load_lora_adapter" in msg
+        assert "VLLM_ALLOW_RUNTIME_LORA_UPDATING" in msg
+
+    def test_full_finetune_with_nccl_endpoint_uses_nccl(self):
+        from unittest.mock import patch
+
+        from axolotl.integrations.nemo_gym.plugin import NemoGymPlugin
+
+        plugin = NemoGymPlugin()
+        plugin._vllm_caps = self._caps(nccl=True)
+        cfg = self._fake_cfg(adapter=None, vllm_lora_sync=False)
+        trainer = self._fake_trainer()
+
+        with patch.object(plugin, "_wire_multi_turn") as wire:
+            plugin.post_trainer_create(cfg, trainer)
+            wire.assert_called_once()
+
+    def test_full_finetune_with_http_endpoint_not_implemented_yet(self):
+        from axolotl.integrations.nemo_gym.plugin import NemoGymPlugin
+
+        plugin = NemoGymPlugin()
+        plugin._vllm_caps = self._caps(http_full=True)
+        cfg = self._fake_cfg(adapter=None, vllm_lora_sync=False)
+        trainer = self._fake_trainer()
+        with self.assertRaises(NotImplementedError) as ctx:
+            plugin.post_trainer_create(cfg, trainer)
+        assert "HTTP weight sync" in str(ctx.exception)
+
+    def test_full_finetune_with_no_routes_raises_with_full_param_message(self):
+        from axolotl.integrations.nemo_gym.plugin import NemoGymPlugin
+
+        plugin = NemoGymPlugin()
+        plugin._vllm_caps = self._caps()
+        cfg = self._fake_cfg(adapter=None, vllm_lora_sync=False)
+        trainer = self._fake_trainer()
+        with self.assertRaises(ValueError) as ctx:
+            plugin.post_trainer_create(cfg, trainer)
+        msg = str(ctx.exception)
+        assert "no-op trainer" in msg
+        assert "init_communicator" in msg
+        assert "http_update_weights" in msg
+
+    def test_unprobed_caps_raises_with_probe_failure_message(self):
+        from axolotl.integrations.nemo_gym.plugin import NemoGymPlugin
+
+        plugin = NemoGymPlugin()
+        # Plugin._vllm_caps left as default-None: the post_trainer_create
+        # branch falls back to a fresh VLLMWeightSyncCapabilities() with
+        # probed=False, so the error path should mention probing.
+        cfg = self._fake_cfg(adapter="lora", vllm_lora_sync=True)
+        trainer = self._fake_trainer()
+        with self.assertRaises(ValueError) as ctx:
+            plugin.post_trainer_create(cfg, trainer)
+        assert "could not probe" in str(ctx.exception)
+
+
+class TestNemoGymE2E(unittest.TestCase):
+    """End-to-end test: data producer → agent (mocked) → parse → tensors → rewards.
+
+    Exercises the full NemoGymDataProducer.produce() pipeline with mocked HTTP
+    responses, verifying that multi-turn Wordle agent responses are correctly
+    parsed into padded tensors with proper env_mask, logprobs, and rewards.
+    No GPU or NeMo Gym server required.
+    """
+
+    # A realistic 2-turn agent /run response (guess + feedback + guess + done)
+    AGENT_RESPONSE = {
+        "response": {
+            "output": [
+                {
+                    "type": "function_call",
+                    "name": "guess_word",
+                    "arguments": '{"guess": "crane"}',
+                    "call_id": "call_1",
+                    "id": "call_1",
+                    "status": "completed",
+                    "prompt_token_ids": [1, 2, 3, 4, 5],
+                    "generation_token_ids": [10, 11, 12, 13],
+                    "generation_log_probs": [-0.1, -0.2, -0.3, -0.4],
+                },
+                {
+                    "type": "function_call_output",
+                    "call_id": "call_1",
+                    "output": '{"feedback":"XYGXY","guesses_remaining":5,"done":false}',
+                },
+                {
+                    "type": "function_call",
+                    "name": "guess_word",
+                    "arguments": '{"guess": "slide"}',
+                    "call_id": "call_2",
+                    "id": "call_2",
+                    "status": "completed",
+                    # prompt = original(5) + gen1(4) + tool_output(3 tokens)
+                    "prompt_token_ids": [1, 2, 3, 4, 5, 10, 11, 12, 13, 50, 51, 52],
+                    "generation_token_ids": [20, 21, 22],
+                    "generation_log_probs": [-0.5, -0.6, -0.7],
+                },
+            ],
+        },
+        "reward": 0.42,
+    }
+
+    def _make_mock_trainer(self):
+        """Create a minimal mock trainer with the attributes produce() needs."""
+        trainer = MagicMock()
+        trainer.accelerator.is_main_process = True
+        trainer.accelerator.device = "cpu"
+        trainer.max_completion_length = 512
+        trainer.temperature = 0.8
+        trainer.pad_token_id = 0
+        trainer.processing_class.eos_token_id = 2
+        trainer.processing_class.batch_decode.return_value = ["crane slide"]
+        return trainer
+
+    @patch("axolotl.integrations.nemo_gym.data_producer._call_agents")
+    def test_produce_returns_valid_rollout_dataset(self, mock_call_agents):
+        """Full pipeline: produce() → _call_agents (mocked) → parse → RolloutDataset."""
+
+        from axolotl.integrations.nemo_gym.data_producer import NemoGymDataProducer
+
+        # Mock _call_agents — it's async, so return a coroutine
+        async def fake_call_agents(**kwargs):
+            return [self.AGENT_RESPONSE, self.AGENT_RESPONSE]
+
+        mock_call_agents.side_effect = fake_call_agents
+
+        # Build a minimal mock of GRPODataProducer's __init__ dependencies
+        # We can't easily call super().__init__, so we'll set attributes directly
+        producer = NemoGymDataProducer.__new__(NemoGymDataProducer)
+        producer._agent_servers = {"wordle_agent": "http://mock:9999"}
+        producer._dataset_lookup = {
+            "Play Wordle!": {
+                "agent_ref": {"name": "wordle_agent"},
+                "verify_extra": {
+                    "responses_create_params": {
+                        "input": [{"role": "user", "content": "Play Wordle!"}],
+                    }
+                },
+            }
+        }
+        producer._request_timeout = 30
+        producer._num_generations = 2
+
+        # Mock the trainer
+        trainer = self._make_mock_trainer()
+        producer._trainer = trainer
+
+        # Mock the prompt iterator. RepeatSampler(mini_repeat_count=num_generations)
+        # pre-expands prompts, so the iterator yields num_generations=2 consecutive
+        # copies of each unique prompt — one entry per rollout.
+        _prompt_batch = [
+            {"prompt": [{"role": "user", "content": "Play Wordle!"}]},
+            {"prompt": [{"role": "user", "content": "Play Wordle!"}]},
+        ]
+        producer._prompt_iter = iter([_prompt_batch])
+        producer._prompt_dl = [_prompt_batch]
+
+        # Call produce
+        result = producer.produce(model=MagicMock(), global_step=1)
+
+        # Verify result structure
+        assert result is not None
+        data = result._data
+
+        # Check tensor shapes — 2 rollouts (num_generations=2)
+        assert data["prompt_ids"].shape[0] == 2
+        assert data["completion_ids"].shape[0] == 2
+        assert data["completion_mask"].shape[0] == 2
+        assert data["sampling_per_token_logps"].shape[0] == 2
+        assert data["tool_mask"].shape[0] == 2
+
+        # Verify completion content — each rollout should have:
+        # gen1(4) + tool_output(3) + gen2(3) = 10 tokens
+        # (padded to same length across the batch, but both are same here)
+        comp_len = data["completion_mask"][0].sum().item()
+        assert comp_len == 10, f"Expected 10 completion tokens, got {comp_len}"
+
+        # Verify env_mask: gen1=1,1,1,1 tool=0,0,0 gen2=1,1,1
+        tool_mask = data["tool_mask"][0][:comp_len].tolist()
+        assert tool_mask == [1, 1, 1, 1, 0, 0, 0, 1, 1, 1]
+
+        # Verify logprobs are populated (use approx for float32 precision)
+        import pytest
+
+        logps = data["sampling_per_token_logps"][0][:comp_len].tolist()
+        assert logps[:4] == pytest.approx([-0.1, -0.2, -0.3, -0.4], abs=1e-6)
+        assert logps[4:7] == pytest.approx([0.0, 0.0, 0.0], abs=1e-6)
+        assert logps[7:10] == pytest.approx([-0.5, -0.6, -0.7], abs=1e-6)
+
+        # Verify rewards were injected into inputs
+        assert data["_deferred_inputs"][0]["env_reward"] == 0.42
+        assert data["_deferred_inputs"][1]["env_reward"] == 0.42
+
+        # Verify deferred scoring markers
+        assert data["_pending_policy_logps"] is True
+
+    @patch("axolotl.integrations.nemo_gym.data_producer._call_agents")
+    def test_produce_handles_failed_agent_response(self, mock_call_agents):
+        """Failed agent responses should produce default (length-1) rollouts."""
+
+        from axolotl.integrations.nemo_gym.data_producer import NemoGymDataProducer
+
+        # One success, one failure — async mock
+        async def fake_call_agents(**kwargs):
+            return [
+                self.AGENT_RESPONSE,
+                {
+                    "error": "Connection timeout",
+                    "response": {"output": []},
+                    "reward": 0.0,
+                },
+            ]
+
+        mock_call_agents.side_effect = fake_call_agents
+
+        producer = NemoGymDataProducer.__new__(NemoGymDataProducer)
+        producer._agent_servers = {"wordle_agent": "http://mock:9999"}
+        producer._dataset_lookup = {}
+        producer._request_timeout = 30
+        producer._num_generations = 2
+        producer._trainer = self._make_mock_trainer()
+        # RepeatSampler pre-expands by num_generations=2.
+        _prompt_batch = [
+            {"prompt": [{"role": "user", "content": "Play!"}]},
+            {"prompt": [{"role": "user", "content": "Play!"}]},
+        ]
+        producer._prompt_iter = iter([_prompt_batch])
+        producer._prompt_dl = [_prompt_batch]
+
+        result = producer.produce(model=MagicMock(), global_step=1)
+
+        assert result is not None
+        data = result._data
+
+        # Both rollouts present
+        assert data["completion_ids"].shape[0] == 2
+
+        # First rollout has real tokens, second has just eos (length 1)
+        mask0 = data["completion_mask"][0].sum().item()
+        mask1 = data["completion_mask"][1].sum().item()
+        assert mask0 == 10  # full response
+        assert mask1 == 1  # default fallback (just eos)
+
+        # Rewards: success=0.42, failure=0.0
+        assert data["_deferred_inputs"][0]["env_reward"] == 0.42
+        assert data["_deferred_inputs"][1]["env_reward"] == 0.0
+
+    @patch("axolotl.integrations.nemo_gym.rewards._get_verify_urls")
+    @patch("axolotl.integrations.nemo_gym.rewards.requests")
+    def test_reward_functions_chain(self, mock_requests, mock_get_urls):
+        """Test that reward_env and reward_nemo_gym_verify can be used together."""
+        from axolotl.integrations.nemo_gym.rewards import (
+            reward_env,
+            reward_nemo_gym_verify,
+        )
+
+        completions = [[{"role": "assistant", "content": "crane"}]]
+        prompts = [[{"role": "user", "content": "Guess"}]]
+
+        # reward_env: passthrough from agent
+        env_result = reward_env(completions, prompts, env_reward=[0.42])
+        assert env_result == [0.42]
+
+        # reward_nemo_gym_verify: calls /verify
+        mock_get_urls.return_value = {"wordle": "http://localhost:9999/verify"}
+        mock_resp = MagicMock()
+        mock_resp.ok = True
+        mock_resp.json.return_value = {"reward": 0.75}
+        mock_requests.post.return_value = mock_resp
+
+        verify_result = reward_nemo_gym_verify(
+            completions,
+            prompts,
+            resources_server_ref=[{"name": "wordle"}],
+            verify_extra=[{}],
+        )
+        assert verify_result == [0.75]
+
+        # Both rewards can coexist (as they would in a multi-reward config)
+        combined = [e + v for e, v in zip(env_result, verify_result, strict=True)]
+        assert combined == [1.17]
+
+
+class TestLoRASyncSetup(unittest.TestCase):
+    """Tests for _setup_lora_sync delegation logic."""
+
+    def test_delegates_to_async_trainer(self):
+        """When trainer has _sync_lora_adapter, the closure should delegate."""
+        from axolotl.integrations.nemo_gym.plugin import NemoGymPlugin
+
+        plugin = NemoGymPlugin.__new__(NemoGymPlugin)
+
+        trainer = MagicMock()
+        trainer._sync_lora_adapter = MagicMock()
+        trainer.vllm_generation = MagicMock()
+
+        plugin._setup_lora_sync(trainer)
+
+        # The closure should be installed
+        trainer.vllm_generation.sync_weights()
+        trainer._sync_lora_adapter.assert_called_once()
+
+    def test_check_lora_endpoint_skips_non_main_rank(self):
+        """_check_lora_endpoint should not crash when vllm_client is absent (rank 1)."""
+        from axolotl.integrations.nemo_gym.plugin import NemoGymPlugin
+
+        vllm_gen = MagicMock(spec=[])  # No attributes at all
+        # Should not raise
+        NemoGymPlugin._check_lora_endpoint(vllm_gen)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/integrations/test_scattermoe_autotune_telemetry.py b/tests/integrations/test_scattermoe_autotune_telemetry.py
new file mode 100644
index 0000000000..034c7b748e
--- /dev/null
+++ b/tests/integrations/test_scattermoe_autotune_telemetry.py
@@ -0,0 +1,415 @@
+"""Tests for scattermoe autotune telemetry integration.
+
+These tests use mocking to verify the collection and reporting logic
+without requiring Triton or CUDA.
+"""
+
+import sys
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+# Simulate the hash-suffixed module name that LocalLayerRepository creates.
+_FAKE_MODULE_NAME = "scattermoe_lora_abc123.kernels.lora_ops"
+
+
+def _make_mock_config(kwargs, num_warps=4, num_stages=3):
+    """Create a mock triton.Config-like object."""
+    return SimpleNamespace(kwargs=kwargs, num_warps=num_warps, num_stages=num_stages)
+
+
+def _make_mock_kernel(cache=None, name="kernel", keys=("M_BUCKET", "N", "K")):
+    """Mock a Triton ``Autotuner``: ``.cache`` dict + ``.base_fn`` (carrying ``__name__``,
+    which the collector uses for the kernel name) + declared autotune ``.keys``."""
+    return SimpleNamespace(
+        cache=cache if cache is not None else {},
+        base_fn=SimpleNamespace(__name__=name),
+        keys=list(keys),
+    )
+
+
+def _make_mock_lora_ops(
+    fwd_cache=None, dx_cache=None, bwd_cache=None, fused_cache=None
+):
+    """Build a mock ``lora_ops`` module with the four kernel attributes. The collector names a
+    kernel by its wrapped fn's ``__name__``, so set those to the expected kernel names."""
+    return SimpleNamespace(
+        _scatter2scatter_lora=_make_mock_kernel(
+            fwd_cache, name="scatter2scatter_lora_fwd"
+        ),
+        _scatter2scatter_lora_dX=_make_mock_kernel(
+            dx_cache, name="scatter2scatter_lora_dX"
+        ),
+        _group_bwd_lora=_make_mock_kernel(bwd_cache, name="group_bwd_lora"),
+        _group_bwd_lora_fused=_make_mock_kernel(
+            fused_cache, name="group_bwd_lora_fused"
+        ),
+    )
+
+
+@contextmanager
+def _inject_lora_ops(mock_module, name="scattermoe_lora_abc123.kernels.lora_ops"):
+    """Register ``mock_module`` under a ``lora_ops``-hint name in ``sys.modules`` and hide EVERY
+    real kernel module the collector scans, so the sys.modules-scanning collector sees only the
+    mock. Hiding just real ``lora_ops`` modules isn't enough: the collector also scans the
+    dsv4/glm_dsa/scattermoe kernel modules (``_KERNEL_MODULE_HINTS``), and when an earlier test in
+    the same process has autotuned one of those, its real cache otherwise leaks into the result."""
+    hide = {n: None for n in _real_scanned_kernel_module_names() if n != name}
+    with patch.dict(sys.modules, {name: mock_module, **hide}):
+        yield
+
+
+def _real_scanned_kernel_module_names():
+    """``sys.modules`` keys the collector would scan (any name matching its kernel-module hints).
+
+    Other tests in the same process may have imported/autotuned the *real* kernels (lora_ops,
+    dsv4, glm_dsa, ...). Hide every such entry so the discovery tests see only the injected mock.
+    """
+    from axolotl.integrations.kernels.autotune_collector import _KERNEL_MODULE_HINTS
+
+    return [
+        name
+        for name, mod in list(sys.modules.items())
+        if mod is not None and any(h in name for h in _KERNEL_MODULE_HINTS)
+    ]
+
+
+# =========================================================================
+# TestAutotuneCollector
+# =========================================================================
+
+
+class TestAutotuneCollector:
+    """Test ``collect_autotune_configs`` with mocked kernel objects.
+
+    The collector scans ``sys.modules`` for kernel-bearing modules, so each test injects a mock
+    ``lora_ops`` module via ``_inject_lora_ops`` (which also hides any real ``lora_ops`` modules
+    other tests in the same pytest-xdist worker may have loaded).
+    """
+
+    def test_empty_cache_returns_empty_list(self):
+        """When no kernel has been autotuned yet, return ``[]``."""
+        mock_lora_ops = _make_mock_lora_ops()
+
+        with _inject_lora_ops(mock_lora_ops):
+            from axolotl.integrations.kernels.autotune_collector import (
+                collect_autotune_configs,
+            )
+
+            result = collect_autotune_configs()
+            assert result == []
+
+    def test_populated_cache_returns_configs(self):
+        """When a cache entry exists, it appears in the output."""
+        cfg = _make_mock_config(
+            {"BLOCK_N": 128, "BLOCK_K": 64}, num_warps=8, num_stages=4
+        )
+        mock_lora_ops = _make_mock_lora_ops(fwd_cache={(2048, 4096, 1024): cfg})
+
+        with _inject_lora_ops(mock_lora_ops):
+            from axolotl.integrations.kernels.autotune_collector import (
+                collect_autotune_configs,
+            )
+
+            result = collect_autotune_configs()
+
+        assert len(result) == 1
+        entry = result[0]
+        assert entry["kernel"] == "scatter2scatter_lora_fwd"
+        assert entry["key"] == {"M_BUCKET": 2048, "N": 4096, "K": 1024}
+        assert entry["config"]["BLOCK_N"] == 128
+        assert entry["config"]["BLOCK_K"] == 64
+        assert entry["config"]["num_warps"] == 8
+        assert entry["config"]["num_stages"] == 4
+
+    def test_multiple_kernels_and_keys(self):
+        """Multiple cache entries across kernels are all returned."""
+        cfg_fwd = _make_mock_config({"BLOCK_N": 128, "BLOCK_K": 32})
+        cfg_dx = _make_mock_config({"BLOCK_K": 64, "BLOCK_N": 128}, num_warps=8)
+
+        mock_lora_ops = _make_mock_lora_ops(
+            fwd_cache={(16, 256, 128): cfg_fwd},
+            dx_cache={(16, 256, 128): cfg_dx},
+        )
+
+        with _inject_lora_ops(mock_lora_ops):
+            from axolotl.integrations.kernels.autotune_collector import (
+                collect_autotune_configs,
+            )
+
+            result = collect_autotune_configs()
+
+        assert len(result) == 2
+        names = {r["kernel"] for r in result}
+        assert "scatter2scatter_lora_fwd" in names
+        assert "scatter2scatter_lora_dX" in names
+
+    def test_extra_key_elements_stored(self):
+        """Dtype or other extra elements in the cache key are captured."""
+        cfg = _make_mock_config({"BLOCK_N": 64, "BLOCK_K": 32})
+        cache_key = (512, 1024, 256, "float16", "float16")
+
+        mock_lora_ops = _make_mock_lora_ops(fwd_cache={cache_key: cfg})
+
+        with _inject_lora_ops(mock_lora_ops):
+            from axolotl.integrations.kernels.autotune_collector import (
+                collect_autotune_configs,
+            )
+
+            result = collect_autotune_configs()
+
+        assert len(result) == 1
+        key = result[0]["key"]
+        assert key["M_BUCKET"] == 512
+        assert key["N"] == 1024
+        assert key["K"] == 256
+        assert key["_extra"] == ["float16", "float16"]
+
+    def test_no_module_in_sys_modules_returns_empty(self):
+        """If no populated lora_ops module is loaded, return ``[]``."""
+        from axolotl.integrations.kernels.autotune_collector import (
+            collect_autotune_configs,
+        )
+
+        hide = {n: None for n in _real_scanned_kernel_module_names()}
+        with patch.dict(sys.modules, hide):
+            result = collect_autotune_configs()
+        assert result == []
+
+    def test_finds_module_under_hash_suffixed_name(self):
+        """Collector finds lora_ops regardless of the hash suffix."""
+        cfg = _make_mock_config({"BLOCK_N": 256, "BLOCK_K": 128})
+        mock_lora_ops = _make_mock_lora_ops(fwd_cache={(8, 512, 64): cfg})
+
+        # Use a different hash to prove it's not hardcoded.
+        alt_name = "scattermoe_lora_deadbeef.kernels.lora_ops"
+
+        # Temporarily hide any real lora_ops modules that other tests in
+        # the same xdist worker may have loaded, so only our mock is found.
+        real_names = _real_scanned_kernel_module_names()
+        hide_patch = {name: None for name in real_names}
+
+        with patch.dict(sys.modules, {alt_name: mock_lora_ops, **hide_patch}):
+            from axolotl.integrations.kernels.autotune_collector import (
+                collect_autotune_configs,
+            )
+
+            result = collect_autotune_configs()
+
+        assert len(result) == 1
+        assert result[0]["config"]["BLOCK_N"] == 256
+
+
+# =========================================================================
+# TestAutotuneReportCallback
+# =========================================================================
+
+
+class TestAutotuneReportCallback:
+    """Test the callback fires once and sends the correct event."""
+
+    def test_reports_once_on_first_step(self):
+        """Callback should call ``send_event`` exactly once."""
+        from axolotl.integrations.kernels.autotune_callback import (
+            AutotuneReportCallback,
+        )
+
+        cb = AutotuneReportCallback()
+        mock_state = MagicMock()
+        mock_state.global_step = 1
+
+        fake_configs = [{"kernel": "test_fwd", "key": {}, "config": {}}]
+
+        with (
+            patch(
+                "axolotl.integrations.kernels.autotune_collector.collect_autotune_configs",
+                return_value=fake_configs,
+            ),
+            patch("axolotl.telemetry.manager.TelemetryManager") as mock_tm_cls,
+        ):
+            mock_tm = MagicMock()
+            mock_tm.enabled = True
+            mock_tm_cls.get_instance.return_value = mock_tm
+
+            cb.on_step_end(args=MagicMock(), state=mock_state, control=MagicMock())
+            assert mock_tm.send_event.call_count == 1
+
+            call_kwargs = mock_tm.send_event.call_args[1]
+            assert call_kwargs["event_type"] == "triton-autotune"
+            assert call_kwargs["properties"]["kernel_count"] == 1
+
+            # Second call should NOT send again.
+            cb.on_step_end(args=MagicMock(), state=mock_state, control=MagicMock())
+            assert mock_tm.send_event.call_count == 1
+
+    def test_retries_until_step_5_then_gives_up(self):
+        """If no configs found by step 5, stop retrying."""
+        from axolotl.integrations.kernels.autotune_callback import (
+            AutotuneReportCallback,
+        )
+
+        cb = AutotuneReportCallback()
+
+        with patch(
+            "axolotl.integrations.kernels.autotune_collector.collect_autotune_configs",
+            return_value=[],
+        ):
+            for step in range(1, 7):
+                mock_state = MagicMock()
+                mock_state.global_step = step
+                cb.on_step_end(args=MagicMock(), state=mock_state, control=MagicMock())
+
+            assert cb._reported is True
+
+    def test_reports_on_retry_when_data_arrives(self):
+        """If step 1 has no data but step 2 does, report at step 2."""
+        from axolotl.integrations.kernels.autotune_callback import (
+            AutotuneReportCallback,
+        )
+
+        cb = AutotuneReportCallback()
+        fake_configs = [{"kernel": "fwd", "key": {}, "config": {}}]
+
+        call_count = 0
+
+        def _collector():
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                return []
+            return fake_configs
+
+        with (
+            patch(
+                "axolotl.integrations.kernels.autotune_collector.collect_autotune_configs",
+                side_effect=_collector,
+            ),
+            patch("axolotl.telemetry.manager.TelemetryManager") as mock_tm_cls,
+        ):
+            mock_tm = MagicMock()
+            mock_tm.enabled = True
+            mock_tm_cls.get_instance.return_value = mock_tm
+
+            # Step 1 — empty, no report
+            s1 = MagicMock()
+            s1.global_step = 1
+            cb.on_step_end(args=MagicMock(), state=s1, control=MagicMock())
+            assert mock_tm.send_event.call_count == 0
+
+            # Step 2 — data arrives, report
+            s2 = MagicMock()
+            s2.global_step = 2
+            cb.on_step_end(args=MagicMock(), state=s2, control=MagicMock())
+            assert mock_tm.send_event.call_count == 1
+
+    def test_includes_gpu_info(self):
+        """Event properties should include GPU identification."""
+        from axolotl.integrations.kernels.autotune_callback import (
+            AutotuneReportCallback,
+        )
+
+        cb = AutotuneReportCallback()
+        mock_state = MagicMock()
+        mock_state.global_step = 1
+
+        fake_configs = [{"kernel": "fwd", "key": {}, "config": {}}]
+        fake_gpu = {
+            "gpu_name": "NVIDIA H100",
+            "gpu_compute_capability": "9.0",
+            "gpu_memory_bytes": 85899345920,
+        }
+
+        fake_smem = {"smem_capacity_bytes": 233472}
+
+        with (
+            patch(
+                "axolotl.integrations.kernels.autotune_collector.collect_autotune_configs",
+                return_value=fake_configs,
+            ),
+            patch(
+                "axolotl.integrations.kernels.autotune_callback._get_gpu_info",
+                return_value=fake_gpu,
+            ),
+            patch(
+                "axolotl.integrations.kernels.autotune_callback._get_smem_capacity",
+                return_value=fake_smem,
+            ),
+            patch("axolotl.telemetry.manager.TelemetryManager") as mock_tm_cls,
+        ):
+            mock_tm = MagicMock()
+            mock_tm.enabled = True
+            mock_tm_cls.get_instance.return_value = mock_tm
+
+            cb.on_step_end(args=MagicMock(), state=mock_state, control=MagicMock())
+            props = mock_tm.send_event.call_args[1]["properties"]
+            assert props["gpu_name"] == "NVIDIA H100"
+            assert props["gpu_compute_capability"] == "9.0"
+            assert props["gpu_memory_bytes"] == 85899345920
+            assert props["smem_capacity_bytes"] == 233472
+
+    def test_skips_send_when_telemetry_disabled(self):
+        """If telemetry is disabled, no event is sent."""
+        from axolotl.integrations.kernels.autotune_callback import (
+            AutotuneReportCallback,
+        )
+
+        cb = AutotuneReportCallback()
+        mock_state = MagicMock()
+        mock_state.global_step = 1
+
+        with (
+            patch(
+                "axolotl.integrations.kernels.autotune_collector.collect_autotune_configs",
+                return_value=[{"kernel": "fwd", "key": {}, "config": {}}],
+            ),
+            patch("axolotl.telemetry.manager.TelemetryManager") as mock_tm_cls,
+        ):
+            mock_tm = MagicMock()
+            mock_tm.enabled = False
+            mock_tm_cls.get_instance.return_value = mock_tm
+
+            cb.on_step_end(args=MagicMock(), state=mock_state, control=MagicMock())
+            assert mock_tm.send_event.call_count == 0
+            # Should still mark as reported so we don't retry.
+            assert cb._reported is True
+
+
+# =========================================================================
+# TestKernelsPluginCallbackRegistration
+# =========================================================================
+
+
+class TestKernelsPluginCallbackRegistration:
+    """Test that ``KernelsPlugin`` registers the callback correctly."""
+
+    def test_scattermoe_registers_callback(self):
+        """When ``use_scattermoe=True``, plugin returns the callback."""
+        from axolotl.integrations.kernels.autotune_callback import (
+            AutotuneReportCallback,
+        )
+        from axolotl.integrations.kernels.plugin import KernelsPlugin
+
+        plugin = KernelsPlugin()
+        cfg = MagicMock()
+        cfg.use_scattermoe = True
+        model = MagicMock()
+
+        callbacks = plugin.add_callbacks_pre_trainer(cfg, model)
+        assert len(callbacks) == 1
+        assert isinstance(callbacks[0], AutotuneReportCallback)
+
+    def test_no_scattermoe_no_callback(self):
+        """When ``use_scattermoe=False``, plugin returns empty list."""
+        from axolotl.integrations.kernels.plugin import KernelsPlugin
+
+        plugin = KernelsPlugin()
+        cfg = MagicMock()
+        cfg.use_scattermoe = False
+        model = MagicMock()
+
+        callbacks = plugin.add_callbacks_pre_trainer(cfg, model)
+        assert callbacks == []
diff --git a/tests/integrations/test_scattermoe_lora.py b/tests/integrations/test_scattermoe_lora.py
new file mode 100644
index 0000000000..7402e20c62
--- /dev/null
+++ b/tests/integrations/test_scattermoe_lora.py
@@ -0,0 +1,729 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Unit tests for scattermoe-lora.
+
+Tests cover:
+- KernelsArgs validator: disable_mlp_kernel
+- ParallelExperts: scaling=0.0 not treated as falsy
+- single2scatter: non-aligned K/N dimensions
+- group_compileable: coeff=None accepted
+- HFScatterMoEGatedMLP / ScatterMoEGatedMLP: return value contract
+- Routing strategy detection and sigmoid routing
+- Generic shared expert handling
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+# ============================================================================
+# 1. KernelsArgs: disable_mlp_kernel validator
+# ============================================================================
+
+
+class TestKernelsArgsValidator:
+    """Test that disable_mlp_kernel sets both flags correctly.
+
+    These tests call the validator classmethod directly on raw dicts,
+    since lora_mlp_kernel / mlp_kernel are not declared model fields.
+    """
+
+    def test_lora_mlp_kernel_kept_for_scattermoe_dense_shared_mlp(self):
+        """lora_mlp_kernel stays enabled under scattermoe (non-DSV4): it fuses only the DENSE
+        shared MLP (layer.mlp); the routed experts are handled by the MoE kernel. The non-LoRA
+        mlp_kernel is still force-disabled."""
+        from axolotl.integrations.kernels.args import KernelsArgs
+
+        data = {
+            "use_kernels": True,
+            "use_scattermoe": True,
+            "lora_mlp_kernel": True,
+        }
+        result = KernelsArgs.disable_mlp_kernel(data)
+        assert result["lora_mlp_kernel"] is True
+        assert result["mlp_kernel"] is False
+
+    def test_lora_mlp_kernel_translated_and_disabled_for_dsv4(self):
+        """On DSV4 the generic lora_mlp_kernel is translated to dsv4_shared_mlp_lora_kernel and the
+        generic path disabled (DSV4's shared MLP needs the dedicated clamped-SwiGLU kernel)."""
+        from axolotl.integrations.kernels.args import KernelsArgs
+
+        data = {
+            "use_kernels": True,
+            "use_scattermoe": True,
+            "use_dsv4_kernels": True,
+            "lora_mlp_kernel": True,
+        }
+        result = KernelsArgs.disable_mlp_kernel(data)
+        assert result["lora_mlp_kernel"] is False
+        assert result["dsv4_shared_mlp_lora_kernel"] is True
+        assert result["mlp_kernel"] is False
+
+    def test_mlp_kernel_disabled_without_lora(self):
+        """Even without lora_mlp_kernel, mlp_kernel should be disabled."""
+        from axolotl.integrations.kernels.args import KernelsArgs
+
+        data = {
+            "use_kernels": True,
+            "use_scattermoe": True,
+        }
+        result = KernelsArgs.disable_mlp_kernel(data)
+        assert result["mlp_kernel"] is False
+        # lora_mlp_kernel was not in data, should not be added
+        assert "lora_mlp_kernel" not in result
+
+    def test_lora_mlp_kernel_false_unchanged(self):
+        """lora_mlp_kernel=False should stay False (no warning, no change)."""
+        from axolotl.integrations.kernels.args import KernelsArgs
+
+        data = {
+            "use_kernels": True,
+            "use_scattermoe": True,
+            "lora_mlp_kernel": False,
+        }
+        result = KernelsArgs.disable_mlp_kernel(data)
+        assert result["lora_mlp_kernel"] is False
+
+    def test_no_change_when_scattermoe_disabled(self):
+        """When use_scattermoe is not True, nothing should be changed."""
+        from axolotl.integrations.kernels.args import KernelsArgs
+
+        data = {
+            "use_kernels": True,
+            "use_scattermoe": False,
+            "lora_mlp_kernel": True,
+        }
+        result = KernelsArgs.disable_mlp_kernel(data)
+        assert result["lora_mlp_kernel"] is True
+
+
+class TestParallelExpertsScaling:
+    """Test that scaling=0.0 is preserved and not overridden to 1.0."""
+
+    def test_scaling_zero_preserved(self):
+        """scaling=0.0 should be passed as 0.0, not replaced with 1.0."""
+        pytest.importorskip("triton")
+        from axolotl.integrations.kernels.libs.scattermoe_lora.lora_ops import (
+            ParallelExperts,
+        )
+
+        pe = ParallelExperts(num_experts=2, input_size=4, output_size=4)
+        pe.set_lora(
+            lora_A=torch.randn(4, 4),
+            lora_B=torch.randn(4, 4),
+            scaling=0.0,
+        )
+        assert pe._lora_scaling == 0.0
+
+        # Patch parallel_linear_lora to capture the scaling arg
+        with patch(
+            "axolotl.integrations.kernels.libs.scattermoe_lora.lora_ops.parallel_linear_lora"
+        ) as mock_pll:
+            mock_pll.return_value = torch.randn(4, 4)
+            # Create dummy routing tensors
+            pe.forward(
+                inputs=torch.randn(2, 4),
+                k=1,
+                sorted_expert_idxs=torch.tensor([0, 0, 1, 1]),
+                sorted_scattered_idxs=torch.tensor([0, 1, 0, 1]),
+                expert_offsets=torch.tensor([2, 4]),
+            )
+            # Check that scaling=0.0 was passed, not 1.0
+            call_kwargs = mock_pll.call_args
+            assert (
+                call_kwargs.kwargs.get("scaling") == 0.0
+                or call_kwargs[1].get("scaling") == 0.0
+            ), f"Expected scaling=0.0 but got {call_kwargs}"
+
+    def test_scaling_none_defaults_to_one(self):
+        """scaling=None (no LoRA attached) should default to 1.0."""
+        pytest.importorskip("triton")
+        from axolotl.integrations.kernels.libs.scattermoe_lora.lora_ops import (
+            ParallelExperts,
+        )
+
+        pe = ParallelExperts(num_experts=2, input_size=4, output_size=4)
+        # No set_lora called, so _lora_scaling is None
+
+        with patch(
+            "axolotl.integrations.kernels.libs.scattermoe_lora.lora_ops.parallel_linear_lora"
+        ) as mock_pll:
+            mock_pll.return_value = torch.randn(4, 4)
+            pe.forward(
+                inputs=torch.randn(2, 4),
+                k=1,
+                sorted_expert_idxs=torch.tensor([0, 0, 1, 1]),
+                sorted_scattered_idxs=torch.tensor([0, 1, 0, 1]),
+                expert_offsets=torch.tensor([2, 4]),
+            )
+            call_kwargs = mock_pll.call_args
+            scaling_val = call_kwargs.kwargs.get("scaling") or call_kwargs[1].get(
+                "scaling"
+            )
+            assert scaling_val == 1.0, (
+                f"Expected scaling=1.0 for None but got {scaling_val}"
+            )
+
+    def test_scaling_positive_preserved(self):
+        """Normal positive scaling should be preserved."""
+        pytest.importorskip("triton")
+        from axolotl.integrations.kernels.libs.scattermoe_lora.lora_ops import (
+            ParallelExperts,
+        )
+
+        pe = ParallelExperts(num_experts=2, input_size=4, output_size=4)
+        pe.set_lora(
+            lora_A=torch.randn(4, 4),
+            lora_B=torch.randn(4, 4),
+            scaling=0.5,
+        )
+
+        with patch(
+            "axolotl.integrations.kernels.libs.scattermoe_lora.lora_ops.parallel_linear_lora"
+        ) as mock_pll:
+            mock_pll.return_value = torch.randn(4, 4)
+            pe.forward(
+                inputs=torch.randn(2, 4),
+                k=1,
+                sorted_expert_idxs=torch.tensor([0, 0, 1, 1]),
+                sorted_scattered_idxs=torch.tensor([0, 1, 0, 1]),
+                expert_offsets=torch.tensor([2, 4]),
+            )
+            call_kwargs = mock_pll.call_args
+            scaling_val = call_kwargs.kwargs.get("scaling") or call_kwargs[1].get(
+                "scaling"
+            )
+            assert scaling_val == 0.5
+
+
+# ============================================================================
+# 4. single2scatter: non-aligned K/N dimensions (GPU only)
+# ============================================================================
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestSingle2ScatterBounds:
+    """Test single2scatter with non-aligned dimensions."""
+
+    def test_non_aligned_k(self):
+        """K not a multiple of BLOCK_K should produce correct results."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.single import (
+            single2scatter,
+        )
+
+        E, K, N = 2, 100, 128  # K=100 not a multiple of 128
+        W = torch.randn(E, K, N, device="cuda", dtype=torch.float32)
+        X = torch.randn(1, K, device="cuda", dtype=torch.float32)
+        expert_idxs = torch.tensor([[0, 1]], device="cuda", dtype=torch.long)
+
+        Y = single2scatter(X, W, expert_idxs)
+        assert Y.shape == (2, N)
+
+        # Verify against manual computation
+        Y_ref_0 = X[0] @ W[0]
+        Y_ref_1 = X[0] @ W[1]
+        torch.testing.assert_close(Y[0], Y_ref_0, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(Y[1], Y_ref_1, atol=1e-2, rtol=1e-2)
+
+    def test_non_aligned_n(self):
+        """N not a multiple of BLOCK_N should produce correct results."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.single import (
+            single2scatter,
+        )
+
+        E, K, N = 2, 128, 100  # N=100 not a multiple of 128
+        W = torch.randn(E, K, N, device="cuda", dtype=torch.float32)
+        X = torch.randn(1, K, device="cuda", dtype=torch.float32)
+        expert_idxs = torch.tensor([[0, 1]], device="cuda", dtype=torch.long)
+
+        Y = single2scatter(X, W, expert_idxs)
+        assert Y.shape == (2, N)
+
+        Y_ref_0 = X[0] @ W[0]
+        Y_ref_1 = X[0] @ W[1]
+        torch.testing.assert_close(Y[0], Y_ref_0, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(Y[1], Y_ref_1, atol=1e-2, rtol=1e-2)
+
+    def test_non_aligned_both(self):
+        """Both K and N not aligned should produce correct results."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.single import (
+            single2scatter,
+        )
+
+        E, K, N = 2, 100, 100  # Neither aligned to 128
+        W = torch.randn(E, K, N, device="cuda", dtype=torch.float32)
+        X = torch.randn(1, K, device="cuda", dtype=torch.float32)
+        expert_idxs = torch.tensor([[0, 1]], device="cuda", dtype=torch.long)
+
+        Y = single2scatter(X, W, expert_idxs)
+        assert Y.shape == (2, N)
+
+        Y_ref_0 = X[0] @ W[0]
+        Y_ref_1 = X[0] @ W[1]
+        torch.testing.assert_close(Y[0], Y_ref_0, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(Y[1], Y_ref_1, atol=1e-2, rtol=1e-2)
+
+
+# ============================================================================
+# 5. group_compileable: coeff=None accepted
+# ============================================================================
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestGroupCoeffNone:
+    """Test that group() works with coeff=None."""
+
+    def test_group_with_none_coeff(self):
+        """group() should accept coeff=None without errors."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops import group
+
+        M, K = 4, 32
+        A = torch.randn(M, K, device="cuda", dtype=torch.float32)
+        sorted_expert_idxs = torch.tensor([0, 1, 2, 3], device="cuda", dtype=torch.long)
+
+        # This should not raise a TypeError
+        Y = group(A, sorted_expert_idxs, coeff=None, fan_out=1)
+        assert Y.shape == (M, K)
+
+    def test_group_with_coeff(self):
+        """group() should also work with actual coeff values."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops import group
+
+        M, K = 4, 32
+        A = torch.randn(M, K, device="cuda", dtype=torch.float32)
+        sorted_expert_idxs = torch.tensor([0, 1, 2, 3], device="cuda", dtype=torch.long)
+        coeff = torch.ones(M, device="cuda", dtype=torch.float32) * 0.5
+
+        Y = group(A, sorted_expert_idxs, coeff=coeff, fan_out=1)
+        assert Y.shape == (M, K)
+
+
+# ============================================================================
+# 6. Layer return value contracts
+# ============================================================================
+
+
+class TestLayerReturnValues:
+    """Test that layer forward methods return the correct types."""
+
+    def test_hf_scatter_moe_returns_single_tensor(self):
+        """HFScatterMoEGatedMLP.forward should return a single tensor, not a tuple."""
+        pytest.importorskip("triton")
+        # Verify the forward method signature and return annotation
+        import inspect
+
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            HFScatterMoEGatedMLP,
+        )
+
+        sig = inspect.signature(HFScatterMoEGatedMLP.forward)
+        # It's a staticmethod taking (self, layer_input)
+        params = list(sig.parameters.keys())
+        assert "self" in params
+        assert "layer_input" in params
+
+    def test_scatter_moe_gated_mlp_docstring_no_router_logits(self):
+        """ScatterMoEGatedMLP.forward docstring should not mention router logits as return."""
+        pytest.importorskip("triton")
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            ScatterMoEGatedMLP,
+        )
+
+        docstring = ScatterMoEGatedMLP.forward.__doc__
+        assert docstring is not None
+        # The docstring should mention output tensor but NOT router logits
+        assert "Output tensor" in docstring or "output tensor" in docstring.lower()
+        assert "Router logits" not in docstring, (
+            "Docstring should not mention 'Router logits' in Returns section"
+        )
+
+
+# ============================================================================
+# 7. Routing strategy detection and sigmoid routing
+# ============================================================================
+
+
+def _make_softmax_gate(E=4, H=16, K=2):
+    """Create a mock softmax-style gate (Qwen/OLMoE)."""
+    return SimpleNamespace(
+        weight=torch.randn(E, H),
+        top_k=K,
+        num_experts=E,
+        norm_topk_prob=True,
+    )
+
+
+def _make_sigmoid_gate_with_bias(E=16, H=16):
+    """Create a mock sigmoid-style gate with e_score_correction_bias on gate."""
+    return SimpleNamespace(
+        weight=torch.randn(E, H),
+        e_score_correction_bias=torch.zeros(E),
+    )
+
+
+def _make_sigmoid_moe_block(
+    T=8, H=16, E=16, K=4, n_group=2, topk_group=1, bias_on_gate=True
+):
+    """Create a mock GLM/DeepSeek-style MoE block for sigmoid routing tests."""
+    if bias_on_gate:
+        gate = SimpleNamespace(
+            weight=torch.randn(E, H),
+            e_score_correction_bias=torch.zeros(E),
+        )
+        moe_block = SimpleNamespace(
+            gate=gate,
+            top_k=K,
+            n_routed_experts=E,
+            n_group=n_group,
+            topk_group=topk_group,
+            norm_topk_prob=True,
+            routed_scaling_factor=1.0,
+        )
+    else:
+        # minimax_m2 style: bias on block, not gate
+        gate = SimpleNamespace(
+            weight=torch.randn(E, H),
+            top_k=K,
+        )
+        moe_block = SimpleNamespace(
+            gate=gate,
+            top_k=K,
+            e_score_correction_bias=torch.zeros(E),
+        )
+    return moe_block, T, H, E, K
+
+
+def _skip_without_triton():
+    pytest.importorskip("triton")
+
+
+class TestSigmoidRoutingInScatterMoE:
+    """Test _sigmoid_topk_route from layers.py."""
+
+    @pytest.fixture(autouse=True)
+    def _require_triton(self):
+        _skip_without_triton()
+
+    def test_output_shapes(self):
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block()
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        weights, experts, top_k, num_experts = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+
+        assert weights.shape == (T, K)
+        assert experts.shape == (T, K)
+        assert top_k == K
+        assert num_experts == E
+
+    def test_weights_nonnegative(self):
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block()
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        weights, _, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+        assert (weights >= 0).all()
+
+    def test_group_selection_restricts_experts(self):
+        """With n_group=4, topk_group=1, experts should be from selected groups."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block(
+            E=16, K=2, n_group=4, topk_group=1
+        )
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        _, expert_idx, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+
+        # Each token's experts should fall within a single group (size E//n_group=4)
+        for t in range(T):
+            experts_t = expert_idx[t]
+            groups = experts_t // (E // moe_block.n_group)
+            assert (groups == groups[0]).all()
+
+    def test_scaling_factor_applied(self):
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block(n_group=1)
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        weights_1x, _, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+
+        moe_block.routed_scaling_factor = 2.0
+        weights_2x, _, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+
+        assert torch.allclose(weights_2x, weights_1x * 2.0, atol=1e-5)
+
+    def test_bias_on_gate(self):
+        """e_score_correction_bias on gate is found."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block(bias_on_gate=True)
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        weights, experts, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+        assert weights.shape == (T, K)
+
+    def test_bias_on_block(self):
+        """e_score_correction_bias on moe_block (not gate) is found."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block(bias_on_gate=False)
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        weights, experts, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+        assert weights.shape == (T, K)
+
+    def test_gate_lora_delta_applied(self):
+        """Gate LoRA delta should affect routing logits."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block(n_group=1)
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        weights_no_lora, _, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+
+        # Large delta should change the results
+        delta = torch.randn(E, H) * 10.0
+        weights_with_lora, _, _, _ = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, delta
+        )
+
+        assert not torch.equal(weights_no_lora, weights_with_lora)
+
+    def test_no_bias_does_not_crash(self):
+        """Calling _sigmoid_topk_route with no e_score_correction_bias should not crash."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        T, H, E, K = 8, 16, 8, 2
+        gate = SimpleNamespace(weight=torch.randn(E, H))
+        moe_block = SimpleNamespace(
+            gate=gate,
+            top_k=K,
+            n_routed_experts=E,
+            n_group=1,
+            norm_topk_prob=True,
+            routed_scaling_factor=1.0,
+        )
+        hidden = torch.randn(T, H)
+
+        weights, experts, top_k, num_experts = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+        assert weights.shape == (T, K)
+        assert experts.shape == (T, K)
+        # Without bias, scores_for_choice == sigmoid(logits) — all positive
+        assert (weights >= 0).all()
+
+    def test_missing_topk_group_defaults_to_n_group(self):
+        """When topk_group is absent but n_group > 1, should default to n_group (no-op masking)."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _sigmoid_topk_route,
+        )
+
+        T, H, E, K, n_group = 8, 16, 16, 2, 4
+        gate = SimpleNamespace(
+            weight=torch.randn(E, H),
+            e_score_correction_bias=torch.zeros(E),
+        )
+        # Intentionally omit topk_group
+        moe_block = SimpleNamespace(
+            gate=gate,
+            top_k=K,
+            n_routed_experts=E,
+            n_group=n_group,
+            norm_topk_prob=True,
+            routed_scaling_factor=1.0,
+        )
+        hidden = torch.randn(T, H)
+
+        # Should not raise AttributeError; defaults topk_group to n_group
+        weights, experts, top_k_out, num_experts = _sigmoid_topk_route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+        assert weights.shape == (T, K)
+        assert experts.shape == (T, K)
+
+
+class TestRoutingStrategyDetection:
+    """Test that _route dispatches to the correct strategy."""
+
+    @pytest.fixture(autouse=True)
+    def _require_triton(self):
+        _skip_without_triton()
+
+    def test_softmax_for_qwen_style(self):
+        """Block without e_score_correction_bias should use softmax."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import _route
+
+        gate = _make_softmax_gate(E=4, H=16, K=2)
+        moe_block = SimpleNamespace(gate=gate)
+        hidden = torch.randn(8, 16)
+
+        weights, experts, top_k, num_experts = _route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+
+        assert weights.shape == (8, 2)
+        assert experts.shape == (8, 2)
+        assert top_k == 2
+        assert num_experts == 4
+        per_token_sums = weights.sum(dim=-1)
+        assert torch.allclose(per_token_sums, torch.ones(8), atol=1e-5)
+
+    def test_sigmoid_for_glm_style(self):
+        """Block with e_score_correction_bias on gate should use sigmoid."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import _route
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block(bias_on_gate=True, n_group=1)
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        weights, experts, top_k, num_experts = _route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+
+        assert weights.shape == (T, K)
+        assert experts.shape == (T, K)
+        assert (weights >= 0).all()
+
+    def test_sigmoid_for_minimax_m2_style(self):
+        """Block with e_score_correction_bias on block (not gate) should use sigmoid."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import _route
+
+        moe_block, T, H, E, K = _make_sigmoid_moe_block(bias_on_gate=False)
+        gate = moe_block.gate
+        hidden = torch.randn(T, H)
+
+        weights, experts, top_k, num_experts = _route(
+            moe_block, gate, hidden, gate.weight, None
+        )
+
+        assert weights.shape == (T, K)
+        assert (weights >= 0).all()
+
+
+# ============================================================================
+# 8. Generic shared expert handling
+# ============================================================================
+
+
+class TestGenericSharedExpert:
+    """Test _compute_shared_expert from layers.py."""
+
+    @pytest.fixture(autouse=True)
+    def _require_triton(self):
+        _skip_without_triton()
+
+    def test_shared_expert_singular(self):
+        """shared_expert attribute (Qwen2MoE style)."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _compute_shared_expert,
+        )
+
+        called = torch.randn(4, 8)
+        moe_block = SimpleNamespace(
+            shared_expert=lambda x: called,
+        )
+        result = _compute_shared_expert(moe_block, torch.randn(4, 8))
+        assert torch.equal(result, called)
+
+    def test_shared_experts_plural(self):
+        """shared_experts attribute (DeepSeek V3 style)."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _compute_shared_expert,
+        )
+
+        called = torch.randn(4, 8)
+        moe_block = SimpleNamespace(
+            shared_experts=lambda x: called,
+        )
+        result = _compute_shared_expert(moe_block, torch.randn(4, 8))
+        assert torch.equal(result, called)
+
+    def test_shared_mlp(self):
+        """shared_mlp attribute (Hunyuan style)."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _compute_shared_expert,
+        )
+
+        called = torch.randn(4, 8)
+        moe_block = SimpleNamespace(
+            shared_mlp=lambda x: called,
+        )
+        result = _compute_shared_expert(moe_block, torch.randn(4, 8))
+        assert torch.equal(result, called)
+
+    def test_shared_expert_with_gate(self):
+        """shared_expert + shared_expert_gate applies sigmoid gating."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _compute_shared_expert,
+        )
+
+        H = 8
+        expert_out = torch.ones(4, H)
+        gate_fn = lambda x: torch.zeros(4, H)  # noqa: E731
+
+        moe_block = SimpleNamespace(
+            shared_expert=lambda x: expert_out,
+            shared_expert_gate=gate_fn,
+        )
+        result = _compute_shared_expert(moe_block, torch.randn(4, H))
+        expected = expert_out * 0.5  # sigmoid(0) = 0.5
+        assert torch.allclose(result, expected, atol=1e-6)
+
+    def test_no_shared_expert(self):
+        """No shared expert attributes returns None."""
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            _compute_shared_expert,
+        )
+
+        moe_block = SimpleNamespace()
+        result = _compute_shared_expert(moe_block, torch.randn(4, 8))
+        assert result is None
diff --git a/tests/integrations/test_scattermoe_lora_kernels.py b/tests/integrations/test_scattermoe_lora_kernels.py
new file mode 100644
index 0000000000..fc783fa1d9
--- /dev/null
+++ b/tests/integrations/test_scattermoe_lora_kernels.py
@@ -0,0 +1,407 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""
+Unit tests for ScatterMoE LoRA Triton kernels.
+
+Tests correctness of:
+  - scatter2scatter_lora (forward)
+  - scatter2scatter_lora_dX (backward input gradient)
+  - group_bwd_lora (backward LoRA weight gradients via split dA/dB)
+  - ScatterMoELoRA autograd function (full forward + backward)
+
+Each kernel is tested against a pure PyTorch per-expert-loop reference
+implementation at multiple model shapes and LoRA ranks.
+"""
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.kernels import (
+    lora_ops,
+    ops as base_ops,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
+    flatten_sort_count,
+)
+from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_linear_lora import (
+    ScatterMoELoRA,
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+def _requires_cuda():
+    return pytest.mark.skipif(
+        not torch.cuda.is_available(), reason="CUDA not available"
+    )
+
+
+pytestmark = _requires_cuda()
+
+
+# ─── Helpers ─────────────────────────────────────────────────────────────────
+
+
+def _setup(E, K, N, T, top_k, R, seed=42):
+    """Create synthetic expert weights, LoRA, routing, and grouped inputs."""
+    torch.manual_seed(seed)
+    x = torch.randn(T, K, device=DEVICE, dtype=DTYPE)
+    W = torch.randn(E, K, N, device=DEVICE, dtype=DTYPE) * 0.02
+    lora_A = torch.randn(R * E, K, device=DEVICE, dtype=DTYPE) * 0.01
+    lora_B = torch.randn(N, R * E, device=DEVICE, dtype=DTYPE) * 0.01
+    logits = torch.randn(T, E, device=DEVICE)
+    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
+    sei, ssi, eo = flatten_sort_count(top_idx, E)
+    return x, W, lora_A, lora_B, sei, ssi, eo
+
+
+def _reference_fwd(x, W, sei, ssi, eo, k, lora_A, lora_B, scaling, E):
+    """Per-expert loop reference: Y = X@W + scaling*(X@A^T)@B^T."""
+    grouped_x = base_ops.group(x, ssi, fan_out=k)
+    M, N = grouped_x.size(0), W.size(2)
+    R = lora_A.size(0) // E
+    out = torch.zeros(M, N, device=DEVICE, dtype=DTYPE)
+    for e in range(E):
+        s = eo[e - 1].item() if e > 0 else 0
+        end = eo[e].item()
+        if s == end:
+            continue
+        xe = grouped_x[s:end].float()
+        we = W[e].float()
+        ae = lora_A[e * R : (e + 1) * R].float()
+        be = lora_B[:, e * R : (e + 1) * R].float()
+        out[s:end] = (xe @ we + scaling * (xe @ ae.T) @ be.T).to(DTYPE)
+    result = torch.zeros(M, N, device=DEVICE, dtype=DTYPE)
+    result[ssi] = out
+    return result
+
+
+def _reference_dX(dy_grouped, W, sei, ssi, eo, lora_A, lora_B, scaling, E):
+    """Per-expert loop reference: dX = dY@W^T + scaling*(dY@B)@A."""
+    M, K = dy_grouped.size(0), W.size(1)
+    R = lora_A.size(0) // E
+    out = torch.zeros(M, K, device=DEVICE, dtype=DTYPE)
+    for e in range(E):
+        s = eo[e - 1].item() if e > 0 else 0
+        end = eo[e].item()
+        if s == end:
+            continue
+        dye = dy_grouped[s:end].float()
+        we = W[e].float()
+        ae = lora_A[e * R : (e + 1) * R].float()
+        be = lora_B[:, e * R : (e + 1) * R].float()
+        out[s:end] = (dye @ we.T + scaling * (dye @ be) @ ae).to(DTYPE)
+    result = torch.zeros(M, K, device=DEVICE, dtype=DTYPE)
+    result[ssi] = out
+    return result
+
+
+def _reference_bwd_lora(dy, grouped_x, lora_A, lora_B, eo, E, scaling):
+    """Per-expert loop reference: dA, dB for LoRA weight gradients."""
+    R = lora_A.size(0) // E
+    dA = torch.zeros_like(lora_A)
+    dB = torch.zeros_like(lora_B)
+    for e in range(E):
+        s = eo[e - 1].item() if e > 0 else 0
+        end = eo[e].item()
+        if s == end:
+            continue
+        xe = grouped_x[s:end].float()
+        dye = dy[s:end].float()
+        ae = lora_A[e * R : (e + 1) * R].float()
+        be = lora_B[:, e * R : (e + 1) * R].float()
+        dA[e * R : (e + 1) * R] = (scaling * (dye @ be).T @ xe).to(DTYPE)
+        dB[:, e * R : (e + 1) * R] = (scaling * dye.T @ (xe @ ae.T)).to(DTYPE)
+    return dA, dB
+
+
+# ─── Model shape configs ────────────────────────────────────────────────────
+
+# (E, K, N, T, top_k, R, description)
+CONFIGS_SMALL = [
+    (32, 128, 64, 64, 2, 4, "tiny"),
+    (64, 256, 128, 128, 4, 8, "small"),
+]
+
+CONFIGS_REAL = [
+    (256, 2048, 1024, 2048, 8, 16, "qwen35_gate_up"),
+    (256, 512, 2048, 2048, 8, 16, "qwen35_down"),
+    (64, 2048, 2048, 2048, 8, 16, "olmoe_gate_up"),
+    (128, 2048, 1536, 2048, 8, 16, "qwen3_gate_up"),
+]
+
+SCALING = 2.0
+
+
+# ─── Forward tests ──────────────────────────────────────────────────────────
+
+
+class TestScatter2ScatterLoRAForward:
+    """Test scatter2scatter_lora forward kernel vs reference."""
+
+    @pytest.fixture(params=CONFIGS_SMALL + CONFIGS_REAL)
+    def config(self, request):
+        return request.param
+
+    def test_matches_reference(self, config):
+        E, K, N, T, k, R, desc = config
+        x, W, lA, lB, sei, ssi, eo = _setup(E, K, N, T, k, R)
+
+        kernel_out = lora_ops.scatter2scatter_lora(
+            X=x,
+            W=W,
+            sorted_expert_idxs=sei,
+            sorted_scattered_idxs=ssi,
+            k=k,
+            lora_A=lA,
+            lora_B=lB,
+            scaling=SCALING,
+        )
+        ref_out = _reference_fwd(x, W, sei, ssi, eo, k, lA, lB, SCALING, E)
+
+        err = (kernel_out.float() - ref_out.float()).abs().max().item()
+        assert err < 1.0, f"[{desc}] fwd max_err={err}"
+
+    def test_output_shape(self, config):
+        E, K, N, T, k, R, desc = config
+        x, W, lA, lB, sei, ssi, eo = _setup(E, K, N, T, k, R)
+
+        out = lora_ops.scatter2scatter_lora(
+            X=x,
+            W=W,
+            sorted_expert_idxs=sei,
+            sorted_scattered_idxs=ssi,
+            k=k,
+            lora_A=lA,
+            lora_B=lB,
+            scaling=SCALING,
+        )
+        assert out.shape == (T * k, N)
+        assert out.dtype == DTYPE
+
+
+# ─── Backward dX tests ──────────────────────────────────────────────────────
+
+
+class TestScatter2ScatterLoRADX:
+    """Test scatter2scatter_lora_dX backward kernel vs reference."""
+
+    @pytest.fixture(params=CONFIGS_SMALL + CONFIGS_REAL)
+    def config(self, request):
+        return request.param
+
+    def test_matches_reference(self, config):
+        E, K, N, T, k, R, desc = config
+        x, W, lA, lB, sei, ssi, eo = _setup(E, K, N, T, k, R)
+        gx = base_ops.group(x, ssi, fan_out=k)
+        dy = torch.randn(gx.size(0), N, device=DEVICE, dtype=DTYPE)
+
+        kernel_dx = lora_ops.scatter2scatter_lora_dX(
+            DY=dy,
+            W=W,
+            sorted_expert_idxs=sei,
+            sorted_scattered_idxs=ssi,
+            k=1,
+            lora_A=lA,
+            lora_B=lB,
+            scaling=SCALING,
+            dy_grouped=True,
+            dx_grouped=False,
+        )
+        ref_dx = _reference_dX(dy, W, sei, ssi, eo, lA, lB, SCALING, E)
+
+        err = (kernel_dx.float() - ref_dx.float()).abs().max().item()
+        assert err < 1.0, f"[{desc}] dX max_err={err}"
+
+
+# ─── Backward LoRA gradient tests ───────────────────────────────────────────
+
+
+class TestGroupBwdLoRA:
+    """Test group_bwd_lora (split dA/dB kernel) vs reference."""
+
+    @pytest.fixture(params=CONFIGS_SMALL + CONFIGS_REAL)
+    def config(self, request):
+        return request.param
+
+    def test_matches_reference(self, config):
+        E, K, N, T, k, R, desc = config
+        x, W, lA, lB, sei, ssi, eo = _setup(E, K, N, T, k, R)
+        gx = base_ops.group(x, ssi, fan_out=k)
+        dy = torch.randn(gx.size(0), N, device=DEVICE, dtype=DTYPE)
+
+        kern_dA, kern_dB = lora_ops.group_bwd_lora(
+            DY=dy,
+            X=gx,
+            lora_A=lA,
+            lora_B=lB,
+            expert_offsets=eo,
+            E=E,
+            scaling=SCALING,
+        )
+        ref_dA, ref_dB = _reference_bwd_lora(dy, gx, lA, lB, eo, E, SCALING)
+
+        # Use norm-relative error: bf16 accumulation order differs between
+        # kernel (tiled + different reduction order) and reference (per-expert
+        # fp32 loop), so max absolute error can be large on individual elements
+        # while the overall tensor is correct.
+        dA_norm_err = (
+            (kern_dA.float() - ref_dA.float()).norm() / (ref_dA.float().norm() + 1e-6)
+        ).item()
+        dB_norm_err = (
+            (kern_dB.float() - ref_dB.float()).norm() / (ref_dB.float().norm() + 1e-6)
+        ).item()
+        assert dA_norm_err < 0.01, f"[{desc}] dA norm_rel_err={dA_norm_err}"
+        assert dB_norm_err < 0.01, f"[{desc}] dB norm_rel_err={dB_norm_err}"
+
+    def test_zero_expert_tokens(self):
+        """Experts with zero routed tokens produce zero gradients."""
+        E, K, N, R = 8, 64, 32, 4
+        torch.manual_seed(42)
+        # Route all tokens to expert 0 only
+        T, k = 16, 1
+        top_idx = torch.zeros(T, k, dtype=torch.long, device=DEVICE)
+        sei, ssi, eo = flatten_sort_count(top_idx, E)
+        gx = torch.randn(T, K, device=DEVICE, dtype=DTYPE)
+        dy = torch.randn(T, N, device=DEVICE, dtype=DTYPE)
+        lA = torch.randn(R * E, K, device=DEVICE, dtype=DTYPE)
+        lB = torch.randn(N, R * E, device=DEVICE, dtype=DTYPE)
+
+        dA, dB = lora_ops.group_bwd_lora(
+            DY=dy,
+            X=gx,
+            lora_A=lA,
+            lora_B=lB,
+            expert_offsets=eo,
+            E=E,
+            scaling=2.0,
+        )
+
+        # Experts 1..7 should have zero gradients
+        for e in range(1, E):
+            assert dA[e * R : (e + 1) * R].abs().max() == 0, f"Expert {e} dA not zero"
+            assert dB[:, e * R : (e + 1) * R].abs().max() == 0, (
+                f"Expert {e} dB not zero"
+            )
+
+
+# ─── Full autograd tests ────────────────────────────────────────────────────
+
+
+class TestScatterMoELoRAAutograd:
+    """Test full forward + backward through ScatterMoELoRA autograd function."""
+
+    @pytest.fixture(params=CONFIGS_SMALL + CONFIGS_REAL[:2])
+    def config(self, request):
+        return request.param
+
+    def test_gradients_exist_and_finite(self, config):
+        E, K, N, T, k, R, desc = config
+        x, W, lA, lB, sei, ssi, eo = _setup(E, K, N, T, k, R)
+
+        x = x.requires_grad_(True)
+        lA = lA.requires_grad_(True)
+        lB = lB.requires_grad_(True)
+
+        out = ScatterMoELoRA.apply(
+            x,
+            W,
+            k,
+            sei,
+            ssi,
+            eo,
+            lA,
+            lB,
+            SCALING,
+            None,
+            None,
+            False,
+            False,
+            True,
+            False,
+        )
+        out.sum().backward()
+
+        assert x.grad is not None, f"[{desc}] x.grad is None"
+        assert lA.grad is not None, f"[{desc}] lA.grad is None"
+        assert lB.grad is not None, f"[{desc}] lB.grad is None"
+        assert torch.isfinite(x.grad).all(), f"[{desc}] x.grad has non-finite"
+        assert torch.isfinite(lA.grad).all(), f"[{desc}] lA.grad has non-finite"
+        assert torch.isfinite(lB.grad).all(), f"[{desc}] lB.grad has non-finite"
+        assert x.grad.abs().sum() > 0, f"[{desc}] x.grad all zero"
+        assert lA.grad.abs().sum() > 0, f"[{desc}] lA.grad all zero"
+
+    def test_split_matches_fused(self):
+        """Split dispatch (for few large experts) matches fused kernel."""
+        # Use a shape where split would be dispatched (large K*N, few E)
+        E, K, N, T, k, R = 8, 512, 1024, 128, 2, 16
+        x, W, lA, lB, sei, ssi, eo = _setup(E, K, N, T, k, R)
+
+        # Force fused path
+        orig = lora_ops._SPLIT_LORA_FWD_THRESHOLD
+        lora_ops._SPLIT_LORA_FWD_THRESHOLD = 10**18
+        out_fused = lora_ops.scatter2scatter_lora(
+            X=x,
+            W=W,
+            sorted_expert_idxs=sei,
+            sorted_scattered_idxs=ssi,
+            k=k,
+            lora_A=lA,
+            lora_B=lB,
+            scaling=SCALING,
+        )
+
+        # Force split path
+        lora_ops._SPLIT_LORA_FWD_THRESHOLD = 0
+        out_split = lora_ops.scatter2scatter_lora(
+            X=x,
+            W=W,
+            sorted_expert_idxs=sei,
+            sorted_scattered_idxs=ssi,
+            k=k,
+            lora_A=lA,
+            lora_B=lB,
+            scaling=SCALING,
+        )
+        lora_ops._SPLIT_LORA_FWD_THRESHOLD = orig
+
+        norm_err = (
+            (out_fused.float() - out_split.float()).norm()
+            / (out_fused.float().norm() + 1e-6)
+        ).item()
+        assert norm_err < 0.01, f"split vs fused norm_err={norm_err}"
+
+    def test_scaling_zero_gives_base_only(self):
+        """With scaling=0.0, LoRA contribution vanishes. Output = X@W."""
+        E, K, N, T, k, R = 16, 64, 32, 32, 2, 4
+        x, W, lA, lB, sei, ssi, eo = _setup(E, K, N, T, k, R)
+
+        out_lora = ScatterMoELoRA.apply(
+            x,
+            W,
+            k,
+            sei,
+            ssi,
+            eo,
+            lA,
+            lB,
+            0.0,
+            None,
+            None,
+            False,
+            False,
+            True,
+            False,
+        )
+        out_base = base_ops.scatter2scatter(
+            X=x,
+            W=W,
+            sorted_expert_idxs=sei,
+            sorted_scattered_idxs=ssi,
+            k=k,
+        )
+        err = (out_lora.float() - out_base.float()).abs().max().item()
+        assert err < 0.01, f"scaling=0 should match base: err={err}"
diff --git a/tests/integrations/test_scattermoe_multi_lora.py b/tests/integrations/test_scattermoe_multi_lora.py
new file mode 100644
index 0000000000..37cdcf5be4
--- /dev/null
+++ b/tests/integrations/test_scattermoe_multi_lora.py
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""ScatterMoE multi-adapter LoRA: forward + backward vs a dense reference.
+
+Co-trains several LoRA adapters (tenants) over one frozen expert stack. The base
+GEMM keys on expert; the LoRA keys on the combined (expert, tenant) group.
+"""
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.scattermoe_lora.multi_lora import (
+    build_multilora_routing,
+    scatter2scatter_multilora,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _reference(x, w, lora_a, lora_b, expert, tenant, num_tenants, rank, scaling):
+    rows = []
+    for n in range(x.size(0)):
+        e, t = expert[n].item(), tenant[n].item()
+        g = e * num_tenants + t
+        a_g = lora_a[g * rank : (g + 1) * rank]  # [R, K]
+        b_g = lora_b[:, g * rank : (g + 1) * rank]  # [N, R]
+        rows.append(x[n] @ w[e] + scaling * ((x[n] @ a_g.t()) @ b_g.t()))
+    return torch.stack(rows)
+
+
+@pytest.mark.parametrize(
+    "num_experts,num_tenants,rank",
+    [
+        (3, 2, 4),
+        (4, 3, 8),
+        (2, 1, 4),  # T=1 reduces to single-adapter
+        (6, 4, 8),  # 24 (expert,tenant) groups vs 48 tokens -> some empty groups
+        (4, 4, 64),  # rank 64 exercises BLOCK_R > 16
+    ],
+)
+def test_multilora_forward_backward_matches_reference(num_experts, num_tenants, rank):
+    torch.manual_seed(0)
+    dev, dt = "cuda", torch.float32
+    e, t, r = num_experts, num_tenants, rank
+    K, N, M = 16, 12, 48
+    scaling = 0.5
+
+    w = torch.randn(e, K, N, device=dev, dtype=dt)
+    a = (torch.randn(e * t * r, K, device=dev, dtype=dt) * 0.1).requires_grad_(True)
+    b = (torch.randn(N, e * t * r, device=dev, dtype=dt) * 0.1).requires_grad_(True)
+    x = torch.randn(M, K, device=dev, dtype=dt, requires_grad=True)
+    expert = torch.randint(0, e, (M,), device=dev)
+    tenant = torch.randint(0, t, (M,), device=dev)
+
+    se, sc, ss, eo, co = build_multilora_routing(expert, tenant, e, t)
+    y = scatter2scatter_multilora(x, w, 1, se, sc, ss, eo, co, a, b, scaling)
+
+    xr = x.detach().clone().requires_grad_(True)
+    ar = a.detach().clone().requires_grad_(True)
+    br = b.detach().clone().requires_grad_(True)
+    yr = _reference(xr, w, ar, br, expert, tenant, t, r, scaling)
+
+    # tf32 matmul tolerance; far tighter than any misroute would produce
+    scale = yr.abs().mean().item()
+    assert (y - yr).abs().max().item() < 0.05 * max(1.0, scale)
+
+    grad = torch.randn_like(y)
+    y.backward(grad)
+    yr.backward(grad)
+    assert torch.allclose(x.grad, xr.grad, atol=5e-2, rtol=5e-2)
+    assert torch.allclose(a.grad, ar.grad, atol=5e-2, rtol=5e-2)
+    assert torch.allclose(b.grad, br.grad, atol=5e-2, rtol=5e-2)
+
+
+def test_multilora_true_fp32_when_tf32_disabled():
+    """With TF32 off, the fused multi-adapter dA/dB must be true-fp32 (~1e-7), not
+    TF32 (~1e-3). Guards the ALLOW_TF32 live-binding fix: an import-time value copy
+    left the grouped-Gram dA/dB kernel on TF32 while the forward honored the flag."""
+    import axolotl.integrations.kernels.libs.scattermoe_lora.kernels.lora_ops as lo
+    import axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops as base_ops
+
+    torch.manual_seed(0)
+    dev, dt = "cuda", torch.float32
+    e, t, r = 4, 3, 8
+    K, N, M, scaling = 16, 12, 48, 0.5
+
+    saved = (torch.backends.cuda.matmul.allow_tf32, lo.ALLOW_TF32, base_ops.ALLOW_TF32)
+    torch.backends.cuda.matmul.allow_tf32 = False
+    lo.ALLOW_TF32 = base_ops.ALLOW_TF32 = False
+    try:
+        w = torch.randn(e, K, N, device=dev, dtype=dt)
+        a = (torch.randn(e * t * r, K, device=dev, dtype=dt) * 0.1).requires_grad_(True)
+        b = (torch.randn(N, e * t * r, device=dev, dtype=dt) * 0.1).requires_grad_(True)
+        x = torch.randn(M, K, device=dev, dtype=dt, requires_grad=True)
+        expert = torch.randint(0, e, (M,), device=dev)
+        tenant = torch.randint(0, t, (M,), device=dev)
+
+        se, sc, ss, eo, co = build_multilora_routing(expert, tenant, e, t)
+        y = scatter2scatter_multilora(x, w, 1, se, sc, ss, eo, co, a, b, scaling)
+
+        xr = x.detach().clone().requires_grad_(True)
+        ar = a.detach().clone().requires_grad_(True)
+        br = b.detach().clone().requires_grad_(True)
+        yr = _reference(xr, w, ar, br, expert, tenant, t, r, scaling)
+
+        grad = torch.randn_like(y)
+        y.backward(grad)
+        yr.backward(grad)
+    finally:
+        (
+            torch.backends.cuda.matmul.allow_tf32,
+            lo.ALLOW_TF32,
+            base_ops.ALLOW_TF32,
+        ) = saved
+
+    def _rel(p, q):
+        return (p - q).abs().max().item() / max(q.abs().max().item(), 1e-12)
+
+    assert _rel(y, yr) < 1e-4
+    assert _rel(x.grad, xr.grad) < 1e-4
+    assert _rel(a.grad, ar.grad) < 1e-4  # dA: TF32 leak would be ~1e-3
+    assert _rel(b.grad, br.grad) < 1e-4  # dB: TF32 leak would be ~1e-3
+
+
+def test_misrouting_would_fail():
+    """Sanity: a wrong tenant assignment produces O(1) error, confirming the
+    tolerance above actually checks routing."""
+    torch.manual_seed(1)
+    dev, dt = "cuda", torch.float32
+    e, t, r, K, N, M = 3, 2, 4, 16, 12, 48
+    w = torch.randn(e, K, N, device=dev, dtype=dt)
+    a = torch.randn(e * t * r, K, device=dev, dtype=dt)
+    b = torch.randn(N, e * t * r, device=dev, dtype=dt)
+    x = torch.randn(M, K, device=dev, dtype=dt)
+    expert = torch.randint(0, e, (M,), device=dev)
+    tenant = torch.randint(0, t, (M,), device=dev)
+    wrong_tenant = 1 - tenant  # flip
+
+    se, sc, ss, eo, co = build_multilora_routing(expert, tenant, e, t)
+    y = scatter2scatter_multilora(x, w, 1, se, sc, ss, eo, co, a, b, 1.0)
+    y_wrong = _reference(x, w, a, b, expert, wrong_tenant, t, r, 1.0)
+    assert (y - y_wrong).abs().max().item() > 1.0
diff --git a/tests/integrations/test_sonicmoe.py b/tests/integrations/test_sonicmoe.py
new file mode 100644
index 0000000000..70078ea73a
--- /dev/null
+++ b/tests/integrations/test_sonicmoe.py
@@ -0,0 +1,237 @@
+"""Unit tests for the SonicMoE ExpertsInterface registration."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.args import KernelsArgs
+
+
+class TestKernelsArgs:
+    def test_mutual_exclusivity_raises(self):
+        with pytest.raises(ValueError, match="Cannot use both"):
+            KernelsArgs.model_validate({"use_scattermoe": True, "use_sonicmoe": True})
+
+    def test_sonicmoe_only(self):
+        result = KernelsArgs.model_validate({"use_sonicmoe": True})
+        assert result.use_sonicmoe is True
+        assert result.use_scattermoe is None
+
+    def test_scattermoe_only(self):
+        result = KernelsArgs.model_validate({"use_scattermoe": True})
+        assert result.use_scattermoe is True
+        assert result.use_sonicmoe is None
+
+    def test_neither_set(self):
+        result = KernelsArgs.model_validate({})
+        assert result.use_scattermoe is None
+        assert result.use_sonicmoe is None
+
+    def test_disables_mlp_kernel_when_sonicmoe(self):
+        # lora_mlp_kernel is kept (it fuses only the DENSE shared MLP; routed experts are handled
+        # by the MoE kernel); the non-LoRA mlp_kernel is still force-disabled under custom MoE.
+        data = {"use_sonicmoe": True, "lora_mlp_kernel": True}
+        result = KernelsArgs.disable_mlp_kernel(data)
+        assert result["lora_mlp_kernel"] is True
+        assert result["mlp_kernel"] is False
+
+    def test_experts_implementation_auto_sonicmoe(self):
+        out = KernelsArgs.check_experts_implementation({"use_sonicmoe": True})
+        assert out["experts_implementation"] == "sonicmoe"
+
+    def test_experts_implementation_auto_scattermoe(self):
+        out = KernelsArgs.check_experts_implementation({"use_scattermoe": True})
+        assert out["experts_implementation"] == "scattermoe"
+
+    def test_experts_implementation_default_eager(self):
+        out = KernelsArgs.check_experts_implementation({})
+        assert out["experts_implementation"] == "eager"
+
+    def test_sonicmoe_impl_requires_flag(self):
+        out = KernelsArgs.check_experts_implementation(
+            {"experts_implementation": "sonicmoe"}
+        )
+        assert out["experts_implementation"] == "eager"
+
+    def test_scattermoe_impl_requires_flag(self):
+        out = KernelsArgs.check_experts_implementation(
+            {"experts_implementation": "scattermoe"}
+        )
+        assert out["experts_implementation"] == "eager"
+
+    def test_unknown_impl_falls_back_to_eager(self):
+        out = KernelsArgs.check_experts_implementation(
+            {"experts_implementation": "not-a-real-impl"}
+        )
+        assert out["experts_implementation"] == "eager"
+
+    def test_builtin_impls_pass_through(self):
+        for impl in ("eager", "batched_mm", "grouped_mm"):
+            out = KernelsArgs.check_experts_implementation(
+                {"experts_implementation": impl}
+            )
+            assert out["experts_implementation"] == impl
+
+
+class TestSonicMoERegistration:
+    """Test that register_sonicmoe_experts plugs into ALL_EXPERTS_FUNCTIONS."""
+
+    def test_register_adds_entry(self):
+        from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+
+        from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+            register_sonicmoe_experts,
+            sonicmoe_experts_forward_with_lora,
+        )
+
+        register_sonicmoe_experts()
+        assert "sonicmoe" in ALL_EXPERTS_FUNCTIONS
+        assert ALL_EXPERTS_FUNCTIONS["sonicmoe"] is sonicmoe_experts_forward_with_lora
+
+    def test_register_is_idempotent(self):
+        from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+
+        from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+            register_sonicmoe_experts,
+        )
+
+        register_sonicmoe_experts()
+        register_sonicmoe_experts()
+        # Just one entry, no error
+        assert "sonicmoe" in ALL_EXPERTS_FUNCTIONS
+
+    def test_register_overrides_upstream(self):
+        """Axolotl's LoRA-aware variant replaces upstream's plain forward."""
+        from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+        from transformers.integrations.sonicmoe import sonicmoe_experts_forward
+
+        from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+            register_sonicmoe_experts,
+            sonicmoe_experts_forward_with_lora,
+        )
+
+        register_sonicmoe_experts()
+        assert ALL_EXPERTS_FUNCTIONS["sonicmoe"] is sonicmoe_experts_forward_with_lora
+        assert ALL_EXPERTS_FUNCTIONS["sonicmoe"] is not sonicmoe_experts_forward
+
+
+class TestMoELoRAMaterialize:
+    """Verify the LoRA materialization autograd Function used by the registered forward."""
+
+    def test_forward_shape_and_identity_with_zero_lora(self):
+        """W_eff == base when LoRA tensors are zero, regardless of layout convention."""
+        from axolotl.integrations.kernels.libs.sonicmoe.lora import MoELoRAMaterialize
+
+        E, dim1, dim2, r = 4, 8, 6, 2
+        base = torch.randn(E, dim1, dim2)
+        lora_A = torch.zeros(r * E, dim2)
+        lora_B = torch.zeros(dim1, r * E)
+        scaling = 0.5
+
+        W_eff = MoELoRAMaterialize.apply(base, lora_A, lora_B, scaling)
+        assert W_eff.shape == base.shape
+        torch.testing.assert_close(W_eff, base, atol=1e-6, rtol=1e-6)
+
+    def test_forward_scaling_linearity(self):
+        """Doubling scaling should double the LoRA delta."""
+        from axolotl.integrations.kernels.libs.sonicmoe.lora import MoELoRAMaterialize
+
+        E, dim1, dim2, r = 4, 8, 6, 2
+        base = torch.randn(E, dim1, dim2)
+        lora_A = torch.randn(r * E, dim2)
+        lora_B = torch.randn(dim1, r * E)
+
+        W_1 = MoELoRAMaterialize.apply(base, lora_A, lora_B, 1.0)
+        W_2 = MoELoRAMaterialize.apply(base, lora_A, lora_B, 2.0)
+        torch.testing.assert_close(W_2 - base, 2 * (W_1 - base), atol=1e-5, rtol=1e-5)
+
+    def test_forward_matches_peft_einsum(self):
+        """Delta matches PEFT's ParamWrapper.get_delta_weight einsum convention.
+
+        Reference: ``peft.tuners.lora.layer.ParamWrapper.get_delta_weight``
+        on PEFT 0.19.x — ``einsum("o r e, e r i -> e o i", B_3d, A_3d)`` where
+        ``B_3d = lora_B.reshape(dim1, r, E)`` and ``A_3d = lora_A.reshape(E, r, dim2)``.
+        """
+        from axolotl.integrations.kernels.libs.sonicmoe.lora import MoELoRAMaterialize
+
+        E, dim1, dim2, r = 3, 5, 4, 2
+        base = torch.zeros(E, dim1, dim2)
+        lora_A = torch.randn(r * E, dim2)
+        lora_B = torch.randn(dim1, r * E)
+        scaling = 0.7
+
+        W_eff = MoELoRAMaterialize.apply(base, lora_A, lora_B, scaling)
+
+        # PEFT's reference computation
+        A_3d = lora_A.reshape(E, r, dim2)
+        B_3d = lora_B.reshape(dim1, r, E)
+        peft_delta = torch.einsum("o r e, e r i -> e o i", B_3d, A_3d) * scaling
+
+        torch.testing.assert_close(W_eff, peft_delta, atol=1e-5, rtol=1e-5)
+
+    def test_gradient_flows_to_lora(self):
+        from axolotl.integrations.kernels.libs.sonicmoe.lora import MoELoRAMaterialize
+
+        E, dim1, dim2, r = 4, 8, 6, 2
+        base = torch.randn(E, dim1, dim2, requires_grad=False)
+        lora_A = torch.randn(r * E, dim2, requires_grad=True)
+        lora_B = torch.randn(dim1, r * E, requires_grad=True)
+        scaling = 0.5
+
+        W_eff = MoELoRAMaterialize.apply(base, lora_A, lora_B, scaling)
+        loss = W_eff.sum()
+        loss.backward()
+
+        assert lora_A.grad is not None
+        assert lora_B.grad is not None
+        assert lora_A.grad.abs().max() > 0
+        assert lora_B.grad.abs().max() > 0
+        # Base weight is frozen — no grad expected.
+        assert base.grad is None
+
+    def test_no_lora_returns_base_unchanged(self):
+        from axolotl.integrations.kernels.libs.sonicmoe.lora import (
+            materialize_expert_lora,
+        )
+
+        base = torch.randn(4, 8, 6)
+        result = materialize_expert_lora(base, None)
+        assert result is base
+
+
+class TestExpertsClassMetadata:
+    """The forward reads `has_gate`/`has_bias`/`is_transposed`/`is_concatenated`
+    that are set by transformers' @use_experts_implementation decorator.
+    Verify our forward respects these without an actual CUDA kernel call.
+    """
+
+    def test_rejects_non_gated(self):
+        from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+            sonicmoe_experts_forward_with_lora,
+        )
+
+        fake_self = SimpleNamespace(has_gate=False)
+        hidden = torch.zeros(2, 4)
+        top_k_index = torch.zeros(2, 1, dtype=torch.long)
+        top_k_weights = torch.ones(2, 1)
+
+        with pytest.raises(ValueError, match="has_gate"):
+            sonicmoe_experts_forward_with_lora(
+                fake_self, hidden, top_k_index, top_k_weights
+            )
+
+    def test_rejects_non_cuda(self):
+        from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+            sonicmoe_experts_forward_with_lora,
+        )
+
+        fake_self = SimpleNamespace(has_gate=True)
+        hidden = torch.zeros(2, 4)  # CPU tensor
+        top_k_index = torch.zeros(2, 1, dtype=torch.long)
+        top_k_weights = torch.ones(2, 1)
+
+        with pytest.raises(ValueError, match="CUDA"):
+            sonicmoe_experts_forward_with_lora(
+                fake_self, hidden, top_k_index, top_k_weights
+            )
diff --git a/tests/integrations/test_sonicmoe_fp4_cute_prep.py b/tests/integrations/test_sonicmoe_fp4_cute_prep.py
new file mode 100644
index 0000000000..64abbc3f72
--- /dev/null
+++ b/tests/integrations/test_sonicmoe_fp4_cute_prep.py
@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""CPU tests for the sonicmoe fp4_cute host-prep layer.
+
+Covers the NVFP4 quantize/dequantize reference, the blocked scale-factor
+layout (checked against an independent per-element formula), the varlen
+dQaccum SFA padding, per-tensor-scale folding, and the gate/up interleave.
+The kernel itself is GPU-only and validated by benchmarks/sonicmoe_nvfp4/.
+"""
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.sonicmoe.nvfp4_quant import (
+    FP4_CODE_VALUES,
+    SF_VEC_SIZE,
+    dequantize_nvfp4_ref,
+    encode_e2m1,
+    fp4_code_to_value,
+    pack_fp4_codes,
+    quantize_nvfp4_ref,
+    simulate_nvfp4_quant,
+    unpack_fp4_codes,
+)
+from axolotl.integrations.kernels.libs.sonicmoe.sf_layout import (
+    build_varlen_sfa,
+    fold_per_tensor_scale,
+    gate_up_interleave_perm,
+    pack_scales_blocked,
+    varlen_padded_num_row_tiles,
+)
+
+
+def _naive_blocked_index(m: int, k: int):
+    """Independent formula for the 512-byte swizzled SF atom position."""
+    tile_m, tile_k = m // 128, k // 4
+    byte = (m % 32) * 16 + ((m // 32) % 4) * 4 + (k % 4)
+    return tile_m, tile_k, byte
+
+
+def test_fp4_code_value_roundtrip():
+    # -0.0 (code 8) encodes back to +0 (code 0); every other code roundtrips.
+    codes = torch.tensor([c for c in range(16) if c != 8], dtype=torch.uint8)
+    vals = fp4_code_to_value(codes)
+    assert torch.equal(encode_e2m1(vals), codes)
+    assert FP4_CODE_VALUES[3] == 1.5 and FP4_CODE_VALUES[11] == -1.5
+
+
+def test_pack_unpack_roundtrip():
+    codes = torch.randint(0, 16, (5, 32), dtype=torch.uint8)
+    assert torch.equal(unpack_fp4_codes(pack_fp4_codes(codes)), codes)
+    # low nibble first: byte 0 holds elements 0 (low) and 1 (high)
+    two = torch.tensor([[3, 9]], dtype=torch.uint8)
+    assert int(pack_fp4_codes(two)[0, 0]) == 3 | (9 << 4)
+
+
+def test_quantize_exact_grid_roundtrips():
+    # Values on the E2M1 grid times an exactly-representable e4m3 scale. The
+    # stored scale is block_amax/6, so pin each 16-block's amax to code 6.0
+    # to make the scale land exactly on 2.0.
+    scale = 2.0
+    codes = torch.randint(0, 16, (4, 32), dtype=torch.uint8)
+    x = fp4_code_to_value(codes) * scale
+    x[..., 0::SF_VEC_SIZE] = 6.0 * scale
+    assert torch.equal(simulate_nvfp4_quant(x), x)
+
+
+def test_quantize_error_bound_and_zeros():
+    torch.manual_seed(0)
+    x = torch.randn(8, 64) * 0.5
+    x[3] = 0.0
+    packed, scale, pts = quantize_nvfp4_ref(x)
+    sim = dequantize_nvfp4_ref(packed, scale, pts)
+    # Max quantization error is half the largest code gap (2.0) times the scale.
+    bound = scale.float().repeat_interleave(SF_VEC_SIZE, dim=-1) * 1.0 + 1e-6
+    assert bool(((x - sim).abs() <= bound).all())
+    assert torch.equal(sim[3], torch.zeros_like(sim[3]))
+    assert not sim.isnan().any()
+
+
+def test_quantize_applies_per_tensor_scale():
+    torch.manual_seed(1)
+    x = torch.randn(4, 32)
+    packed, scale, pts = quantize_nvfp4_ref(x, per_tensor_scale=0.25)
+    assert float(pts) == 0.25
+    sim = dequantize_nvfp4_ref(packed, scale, pts)
+    bound = scale.float().repeat_interleave(SF_VEC_SIZE, dim=-1) * float(pts) + 1e-6
+    assert bool(((x - sim).abs() <= bound).all())
+
+
+def test_pack_scales_blocked_matches_naive_formula():
+    torch.manual_seed(2)
+    l, mn, sf_k = 2, 130, 6  # exercises both pad directions
+    raw = torch.randint(0, 256, (l, mn, sf_k), dtype=torch.uint8)
+    blocked = pack_scales_blocked(raw)
+    assert blocked.shape == (l, 2, 2, 512)
+    for li in range(l):
+        for m in range(mn):
+            for k in range(sf_k):
+                tm, tk, byte = _naive_blocked_index(m, k)
+                assert int(blocked[li, tm, tk, byte]) == int(raw[li, m, k])
+    # padding is zero
+    tm, tk, byte = _naive_blocked_index(131, 0)
+    assert int(blocked[0, tm, tk, byte]) == 0
+
+
+def test_pack_scales_blocked_preserves_dtype():
+    raw = torch.randint(0, 256, (1, 128, 4), dtype=torch.uint8).view(
+        torch.float8_e4m3fn
+    )
+    out = pack_scales_blocked(raw)
+    assert out.dtype == torch.float8_e4m3fn and out.shape == (1, 1, 1, 512)
+
+
+def test_build_varlen_sfa_offsets():
+    torch.manual_seed(3)
+    seqlens = [5, 0, 130, 3]
+    total_m = sum(seqlens)
+    cu = torch.tensor([0] + list(torch.tensor(seqlens).cumsum(0).tolist()))
+    sf_k = 4
+    rows = torch.randint(1, 255, (total_m, sf_k), dtype=torch.uint8)
+    blocked = build_varlen_sfa(rows, cu)
+
+    e = len(seqlens)
+    assert varlen_padded_num_row_tiles(total_m, e) == (total_m + 127) // 128 + e - 1
+    assert blocked.shape == (1, varlen_padded_num_row_tiles(total_m, e), 1, 512)
+
+    for i in range(e):
+        start = int(cu[i])
+        row0 = (start // 128 + i) * 128
+        for r in range(seqlens[i]):
+            padded_row = row0 + r
+            for k in range(sf_k):
+                tm, tk, byte = _naive_blocked_index(padded_row, k)
+                assert int(blocked[0, tm, tk, byte]) == int(rows[start + r, k])
+
+
+def test_build_varlen_sfa_rejects_bad_cu():
+    rows = torch.zeros(10, 4, dtype=torch.uint8)
+    with pytest.raises(AssertionError):
+        build_varlen_sfa(rows, torch.tensor([0, 4, 9]))  # cu[-1] != total_m
+
+
+def test_fold_per_tensor_scale():
+    scale = torch.full((2, 4, 4), 2.0).to(torch.float8_e4m3fn)
+    pts = torch.tensor([0.5, 2.0])
+    folded, rel_err = fold_per_tensor_scale(scale, pts)
+    assert rel_err == 0.0
+    assert torch.equal(folded[0].float(), torch.full((4, 4), 1.0))
+    assert torch.equal(folded[1].float(), torch.full((4, 4), 4.0))
+    # accepts (E,1,1) shape and scalar
+    folded2, _ = fold_per_tensor_scale(scale, pts.view(2, 1, 1))
+    assert torch.equal(folded2.float(), folded.float())
+    folded3, _ = fold_per_tensor_scale(scale, torch.tensor(1.0))
+    assert torch.equal(folded3.float(), scale.float())
+
+
+def test_fold_per_tensor_scale_saturation_and_underflow():
+    big = torch.full((1, 2, 2), 448.0).to(torch.float8_e4m3fn)
+    with pytest.raises(ValueError, match="saturates"):
+        fold_per_tensor_scale(big, torch.tensor([2.0]))
+    tiny = torch.full((1, 2, 2), 2.0**-9).to(torch.float8_e4m3fn)  # e4m3 min subnormal
+    with pytest.raises(ValueError, match="underflows"):
+        fold_per_tensor_scale(tiny, torch.tensor([0.25]))
+
+
+def test_gate_up_interleave_perm():
+    perm = gate_up_interleave_perm(8)
+    assert perm.tolist() == [0, 4, 1, 5, 2, 6, 3, 7]
+    w = torch.arange(8).unsqueeze(-1)  # concat: gate rows 0-3, up rows 4-7
+    wi = w[perm]
+    # pair j = (gate_j, up_j)
+    assert wi[0::2].squeeze(-1).tolist() == [0, 1, 2, 3]
+    assert wi[1::2].squeeze(-1).tolist() == [4, 5, 6, 7]
+    # invertible
+    inv = torch.argsort(perm)
+    assert torch.equal(wi[inv], w)
+
+
+def test_fp4_cute_unavailable_without_sm100():
+    from axolotl.integrations.kernels.libs.sonicmoe.fp4_cute import fp4_cute_available
+
+    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] in (10, 11):
+        pytest.skip("running on SM100/SM110")
+    assert fp4_cute_available() is False
+
+
+def _make_nvfp4_weight(E, N, K, pts=None, seed=0):
+    """Random dense [E, N, K] quantized per expert into a torchao NVFP4Tensor."""
+    nvfp4_mod = pytest.importorskip("torchao.prototype.mx_formats.nvfp4_tensor")
+    torch.manual_seed(seed)
+    pts_t = torch.ones(E) if pts is None else torch.as_tensor(pts, dtype=torch.float32)
+    qs, ss = [], []
+    for e in range(E):
+        q, s, _ = quantize_nvfp4_ref(torch.randn(N, K) * K**-0.5, pts_t[e])
+        qs.append(q)
+        ss.append(s)
+    return nvfp4_mod.NVFP4Tensor(
+        torch.stack(qs),
+        torch.stack(ss),
+        16,
+        torch.float32,
+        per_tensor_scale=pts_t.view(-1, 1, 1),
+    )
+
+
+def test_torchao_dequant_matches_ref():
+    # The kernel path assumes torchao's scheme (low nibble first, value =
+    # code * e4m3_scale * pts) is identical to our reference.
+    w = _make_nvfp4_weight(3, 8, 32, pts=[0.5, 1.3, 2.0])
+    ref = dequantize_nvfp4_ref(w.qdata, w.scale, w.per_tensor_scale)
+    assert torch.equal(w.dequantize(torch.float32), ref)
+
+
+def test_dequantize_expert_slice():
+    from axolotl.integrations.kernels.libs.sonicmoe.nvfp4 import (
+        dequantize_expert_slice,
+    )
+
+    w = _make_nvfp4_weight(3, 8, 32, pts=[0.5, 1.0, 2.0])
+    full = w.dequantize(torch.float32)
+    for e in range(3):
+        assert torch.equal(dequantize_expert_slice(w, e), full[e])
+
+    dense = torch.randn(3, 8, 32)
+    assert dequantize_expert_slice(dense, 1) is not None
+    assert torch.equal(dequantize_expert_slice(dense, 1), dense[1])
+
+
+def test_unpack_nvfp4_components():
+    from axolotl.integrations.kernels.libs.sonicmoe.fp4_cute_ops import (
+        unpack_nvfp4_components,
+    )
+
+    w = _make_nvfp4_weight(2, 8, 32, pts=[1.0, 0.5])
+    qdata, scale, pts = unpack_nvfp4_components(w)
+    assert qdata.dtype == torch.uint8 and qdata.shape == (2, 8, 16)
+    assert scale.dtype == torch.float8_e4m3fn and scale.shape == (2, 8, 2)
+    assert pts.shape == (2, 1, 1)
+    with pytest.raises(TypeError):
+        unpack_nvfp4_components(torch.randn(2, 8, 32))
+
+
+def test_fp4_cute_dims_ok():
+    from axolotl.integrations.kernels.libs.sonicmoe.fp4_cute_ops import (
+        fp4_cute_dims_ok,
+    )
+
+    w1 = torch.empty(2, 16, 64)  # [E, 2I, H]
+    w2 = torch.empty(2, 64, 32)  # [E, H, I]
+    assert fp4_cute_dims_ok(w1, w2) is True
+    assert fp4_cute_dims_ok(torch.empty(2, 16, 48), w2) is False  # K % 32
+    assert fp4_cute_dims_ok(torch.empty(2, 12, 64), w2) is False  # N % 8
+    assert fp4_cute_dims_ok(w1, torch.empty(2, 64, 24)) is False
+
+
+def test_lora_backward_nvfp4_base_matches_dense():
+    from axolotl.integrations.kernels.libs.sonicmoe.nvfp4_lora import (
+        _lora_backward_per_group,
+    )
+
+    torch.manual_seed(4)
+    E, d1, d2, r = 3, 8, 32, 2
+    w = _make_nvfp4_weight(E, d1, d2, pts=[0.5, 1.0, 2.0])
+    counts = [4, 0, 3]
+    offsets = torch.tensor([0] + list(torch.tensor(counts).cumsum(0).tolist()))
+    T = sum(counts)
+    x = torch.randn(T, d2)
+    grad_h = torch.randn(T, d1)
+    lora_A = torch.randn(r * E, d2)
+    lora_B = torch.randn(d1, r * E)
+
+    got = _lora_backward_per_group(grad_h, x, offsets, lora_A, lora_B, w, 0.5)
+    ref = _lora_backward_per_group(
+        grad_h, x, offsets, lora_A, lora_B, w.dequantize(torch.float32), 0.5
+    )
+    for g, rf in zip(got, ref, strict=True):
+        torch.testing.assert_close(g, rf)
+
+
+def test_select_nvfp4_backend(monkeypatch):
+    from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+        _select_nvfp4_backend,
+    )
+
+    dense1, dense2 = torch.empty(2, 16, 64), torch.empty(2, 64, 32)
+    monkeypatch.delenv("AXOLOTL_SONICMOE_NVFP4_BACKEND", raising=False)
+    assert _select_nvfp4_backend(dense1, dense2) == "dequant"
+    monkeypatch.setenv("AXOLOTL_SONICMOE_NVFP4_BACKEND", "fp4_cute")
+    assert _select_nvfp4_backend(dense1, dense2) == "fp4_cute"
diff --git a/tests/integrations/test_sonicmoe_lora.py b/tests/integrations/test_sonicmoe_lora.py
new file mode 100644
index 0000000000..7b1cf1765d
--- /dev/null
+++ b/tests/integrations/test_sonicmoe_lora.py
@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Unit tests for SonicMoE LoRA support."""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.sonicmoe.lora import (
+    MoELoRAMaterialize,
+    get_lora_params_from_wrapper,
+    has_lora,
+    materialize_expert_lora,
+    unwrap_experts_lora,
+)
+
+# =============================================================================
+# Helpers: mock PEFT modules
+# =============================================================================
+
+
+def _make_mock_lora_module(weight_A, weight_B, scaling_val, param_name=None):
+    """Create a mock PEFT-wrapped module with LoRA attributes."""
+    mock = MagicMock()
+
+    lora_A_linear = MagicMock()
+    lora_A_linear.weight = weight_A
+
+    lora_B_linear = MagicMock()
+    lora_B_linear.weight = weight_B
+
+    mock.lora_A = {"default": lora_A_linear}
+    mock.lora_B = {"default": lora_B_linear}
+    mock.scaling = {"default": scaling_val}
+    mock.active_adapters = ["default"]
+
+    if param_name is not None:
+        mock.parameter_name = param_name
+
+    return mock
+
+
+def _make_peft_experts(
+    num_experts, gate_up_dim, down_dim, hidden_size, rank, scaling=0.5
+):
+    """Create a mock PEFT-wrapped experts chain.
+
+    Simulates: ParamWrapper(down_proj) -> ParamWrapper(gate_up_proj) -> Experts
+    """
+    base_experts = MagicMock()
+    base_experts.gate_up_proj = torch.randn(num_experts, gate_up_dim, hidden_size)
+    base_experts.down_proj = torch.randn(num_experts, hidden_size, down_dim)
+    # Remove base_layer and lora_A from base_experts so the chain walk stops
+    del base_experts.base_layer
+    del base_experts.lora_A
+
+    # gate_up_proj wrapper
+    gup_A = torch.randn(rank * num_experts, hidden_size)
+    gup_B = torch.randn(gate_up_dim, rank * num_experts)
+    gup_wrapper = _make_mock_lora_module(gup_A, gup_B, scaling, "gate_up_proj")
+    gup_wrapper.base_layer = base_experts
+
+    # down_proj wrapper (outermost)
+    down_A = torch.randn(rank * num_experts, down_dim)
+    down_B = torch.randn(hidden_size, rank * num_experts)
+    down_wrapper = _make_mock_lora_module(down_A, down_B, scaling, "down_proj")
+    down_wrapper.base_layer = gup_wrapper
+
+    return down_wrapper, base_experts, (gup_A, gup_B), (down_A, down_B)
+
+
+# =============================================================================
+# Tests: has_lora
+# =============================================================================
+
+
+class TestHasLora:
+    def test_plain_module(self):
+        m = MagicMock(spec=["weight"])
+        del m.base_layer
+        del m.lora_A
+        assert not has_lora(m)
+
+    def test_wrapped_module(self):
+        m = MagicMock()
+        m.base_layer = MagicMock()
+        m.lora_A = {"default": MagicMock()}
+        assert has_lora(m)
+
+
+# =============================================================================
+# Tests: get_lora_params_from_wrapper
+# =============================================================================
+
+
+class TestGetLoraParams:
+    def test_no_lora_attrs(self):
+        m = MagicMock(spec=["weight"])
+        del m.lora_A
+        del m.lora_B
+        assert get_lora_params_from_wrapper(m) == (None, None, None)
+
+    def test_extracts_params(self):
+        A = torch.randn(4, 8)
+        B = torch.randn(16, 4)
+        wrapper = _make_mock_lora_module(A, B, 0.5)
+        lora_A, lora_B, scaling = get_lora_params_from_wrapper(wrapper)
+        assert torch.equal(lora_A, A)
+        assert torch.equal(lora_B, B)
+        assert scaling == 0.5
+
+    def test_no_active_adapters(self):
+        wrapper = _make_mock_lora_module(torch.randn(4, 8), torch.randn(16, 4), 0.5)
+        wrapper.active_adapters = []
+        assert get_lora_params_from_wrapper(wrapper) == (None, None, None)
+
+
+# =============================================================================
+# Tests: unwrap_experts_lora
+# =============================================================================
+
+
+class TestUnwrapExpertsLora:
+    def test_plain_experts(self):
+        experts = MagicMock(spec=["gate_up_proj", "down_proj"])
+        del experts.base_layer
+        del experts.lora_A
+        base, lora_dict = unwrap_experts_lora(experts)
+        assert base is experts
+        assert lora_dict == {}
+
+    def test_wrapped_experts(self):
+        E, I2, I, H, r = 4, 256, 128, 64, 8  # noqa: E741
+        wrapper, base_experts, (gup_A, gup_B), (down_A, down_B) = _make_peft_experts(
+            E, I2, I, H, r, scaling=0.25
+        )
+        base, lora_dict = unwrap_experts_lora(wrapper)
+        assert base is base_experts
+        assert "gate_up_proj" in lora_dict
+        assert "down_proj" in lora_dict
+
+        gup_lA, gup_lB, gup_s = lora_dict["gate_up_proj"]
+        assert torch.equal(gup_lA, gup_A)
+        assert torch.equal(gup_lB, gup_B)
+        assert gup_s == 0.25
+
+        down_lA, down_lB, down_s = lora_dict["down_proj"]
+        assert torch.equal(down_lA, down_A)
+        assert torch.equal(down_lB, down_B)
+        assert down_s == 0.25
+
+    def test_partial_lora(self):
+        """Only gate_up_proj has LoRA, down_proj does not."""
+        base_experts = MagicMock(spec=["gate_up_proj", "down_proj"])
+        del base_experts.base_layer
+        del base_experts.lora_A
+
+        gup_A = torch.randn(16, 64)
+        gup_B = torch.randn(256, 16)
+        gup_wrapper = _make_mock_lora_module(gup_A, gup_B, 0.5, "gate_up_proj")
+        gup_wrapper.base_layer = base_experts
+
+        base, lora_dict = unwrap_experts_lora(gup_wrapper)
+        assert base is base_experts
+        assert "gate_up_proj" in lora_dict
+        assert "down_proj" not in lora_dict
+
+
+# =============================================================================
+# Tests: MoELoRAMaterialize
+# =============================================================================
+
+
+class TestMoELoRAMaterialize:
+    @pytest.fixture()
+    def setup(self):
+        E, dim1, dim2, r = 4, 32, 16, 4
+        scaling = 0.5
+        W = torch.randn(E, dim1, dim2, dtype=torch.float64, requires_grad=False)
+        A = torch.randn(r * E, dim2, dtype=torch.float64, requires_grad=True)
+        B = torch.randn(dim1, r * E, dtype=torch.float64, requires_grad=True)
+        return W, A, B, scaling, E, r
+
+    def test_forward_shape(self, setup):
+        W, A, B, scaling, E, r = setup
+        W_eff = MoELoRAMaterialize.apply(W, A, B, scaling)
+        assert W_eff.shape == W.shape
+
+    def test_forward_correctness(self, setup):
+        W, A, B, scaling, E, r = setup
+        W_eff = MoELoRAMaterialize.apply(W, A, B, scaling)
+
+        # Manual per-expert computation.
+        # lora_A is expert-major: [r*E, dim2] -> rows [e*r:(e+1)*r] = expert e
+        # lora_B is rank-major:   [dim1, r*E] -> reshape [dim1, r, E], slice [:, :, e]
+        _, dim1, dim2 = W.shape
+        expected = W.clone()
+        B_3d = B.reshape(dim1, r, E)
+        for e in range(E):
+            A_e = A[e * r : (e + 1) * r, :]  # [r, dim2]
+            B_e = B_3d[:, :, e]  # [dim1, r]
+            expected[e] += scaling * (B_e @ A_e)
+
+        assert torch.allclose(W_eff, expected, atol=1e-10)
+
+    def test_backward_gradcheck(self, setup):
+        W, A, B, scaling, E, r = setup
+        # gradcheck requires float64
+        assert torch.autograd.gradcheck(
+            lambda a, b: MoELoRAMaterialize.apply(W, a, b, scaling),
+            (A, B),
+            eps=1e-6,
+            atol=1e-4,
+        )
+
+    def test_no_grad_for_base_weight(self, setup):
+        W, A, B, scaling, E, r = setup
+        W.requires_grad_(True)
+        W_eff = MoELoRAMaterialize.apply(W, A, B, scaling)
+        loss = W_eff.sum()
+        loss.backward()
+        assert W.grad is None
+        assert A.grad is not None
+        assert B.grad is not None
+
+    def test_scaling_zero(self, setup):
+        W, A, B, _, E, r = setup
+        W_eff = MoELoRAMaterialize.apply(W, A, B, 0.0)
+        assert torch.allclose(W_eff, W)
+
+    def test_gate_up_proj_shapes(self):
+        """Test with realistic gate_up_proj shapes [E, 2*I, H]."""
+        E, I2, H, r = 8, 512, 256, 16
+        W = torch.randn(E, I2, H, dtype=torch.float64)
+        A = torch.randn(r * E, H, dtype=torch.float64, requires_grad=True)
+        B = torch.randn(I2, r * E, dtype=torch.float64, requires_grad=True)
+        W_eff = MoELoRAMaterialize.apply(W, A, B, 1.0)
+        assert W_eff.shape == (E, I2, H)
+        loss = W_eff.sum()
+        loss.backward()
+        assert A.grad.shape == A.shape
+        assert B.grad.shape == B.shape
+
+    def test_down_proj_shapes(self):
+        """Test with realistic down_proj shapes [E, H, I]."""
+        E, H, I, r = 8, 256, 512, 16  # noqa: E741
+        W = torch.randn(E, H, I, dtype=torch.float64)
+        A = torch.randn(r * E, I, dtype=torch.float64, requires_grad=True)
+        B = torch.randn(H, r * E, dtype=torch.float64, requires_grad=True)
+        W_eff = MoELoRAMaterialize.apply(W, A, B, 1.0)
+        assert W_eff.shape == (E, H, I)
+        loss = W_eff.sum()
+        loss.backward()
+        assert A.grad.shape == A.shape
+        assert B.grad.shape == B.shape
+
+
+# =============================================================================
+# Tests: materialize_expert_lora
+# =============================================================================
+
+
+class TestMaterializeExpertLora:
+    def test_none_passthrough(self):
+        W = torch.randn(4, 32, 16)
+        result = materialize_expert_lora(W, None)
+        assert result is W
+
+    def test_with_lora(self):
+        E, dim1, dim2, r = 4, 32, 16, 4
+        W = torch.randn(E, dim1, dim2)
+        A = torch.randn(r * E, dim2, requires_grad=True)
+        B = torch.randn(dim1, r * E, requires_grad=True)
+        result = materialize_expert_lora(W, (A, B, 0.5))
+        assert result.shape == W.shape
+        assert not torch.equal(result, W)
diff --git a/tests/integrations/test_sonicmoe_multi_lora.py b/tests/integrations/test_sonicmoe_multi_lora.py
new file mode 100644
index 0000000000..73f189982f
--- /dev/null
+++ b/tests/integrations/test_sonicmoe_multi_lora.py
@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""SonicMoE multi-adapter LoRA: validate the E*T materialization + id remap.
+
+The opaque CUTLASS grouped GEMM is unchanged and isn't exercised here; an exact
+PyTorch grouped GEMM stands in for it so the materialize + combined-routing logic
+can be checked end-to-end (forward + backward) against a per-tenant reference.
+"""
+
+import pytest
+import torch
+
+from axolotl.integrations.kernels.libs.sonicmoe.lora import MoELoRAMaterialize
+from axolotl.integrations.kernels.libs.sonicmoe.multi_lora import (
+    combined_expert_ids,
+    materialize_multi_lora_experts,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _grouped_gemm(hidden, w_eff, group_ids, token_idx, out_dim):
+    """Exact stand-in for the CUTLASS grouped GEMM: out[token_idx[i]] += h @ W_eff[g_i]."""
+    out = hidden.new_zeros(hidden.shape[0], out_dim)
+    for i in range(group_ids.shape[0]):
+        n = token_idx[i].item()
+        out[n] = out[n] + hidden[n] @ w_eff[group_ids[i].item()].t()
+    return out
+
+
+@pytest.mark.parametrize("E,T,r", [(4, 3, 8), (3, 1, 4), (5, 4, 16)])
+def test_multilora_materialize_matches_per_tenant(E, T, r):
+    torch.manual_seed(0)
+    dev, dt = "cuda", torch.float32
+    out_dim, in_dim, M = 12, 16, 32
+    base = torch.randn(E, out_dim, in_dim, device=dev, dtype=dt)
+    A = (torch.randn(T, E, r, in_dim, device=dev, dtype=dt) * 0.1).requires_grad_(True)
+    B = (torch.randn(T, E, out_dim, r, device=dev, dtype=dt) * 0.1).requires_grad_(True)
+    scaling = torch.rand(T, device=dev, dtype=dt) + 0.5
+
+    hidden = torch.randn(M, in_dim, device=dev, dtype=dt)
+    expert = torch.randint(0, E, (M,), device=dev, dtype=torch.int32)
+    tenant = torch.randint(0, T, (M,), device=dev, dtype=torch.int32)
+    token_idx = torch.arange(M, device=dev, dtype=torch.int32)
+
+    # combined path
+    w_eff = materialize_multi_lora_experts(base, A, B, scaling)  # [E*T, out, in]
+    g_ids = combined_expert_ids(expert, tenant, token_idx, T)
+    out = _grouped_gemm(hidden, w_eff, g_ids, token_idx, out_dim)
+
+    # reference: each token uses its tenant's expert-e effective weight, directly
+    Ar = A.detach().clone().requires_grad_(True)
+    Br = B.detach().clone().requires_grad_(True)
+    out_ref = hidden.new_zeros(M, out_dim)
+    for n in range(M):
+        e, t = expert[n].item(), tenant[n].item()
+        w = base[e] + scaling[t] * (Br[t, e] @ Ar[t, e])
+        out_ref[n] = hidden[n] @ w.t()
+
+    assert torch.allclose(out, out_ref, atol=1e-4, rtol=1e-4)
+
+    grad = torch.randn_like(out)
+    out.backward(grad)
+    out_ref.backward(grad)
+    assert torch.allclose(A.grad, Ar.grad, atol=1e-4, rtol=1e-4)
+    assert torch.allclose(B.grad, Br.grad, atol=1e-4, rtol=1e-4)
+
+
+def test_t1_reduces_to_single_adapter():
+    """T==1 must produce exactly the single-adapter MoELoRAMaterialize result."""
+    torch.manual_seed(1)
+    dev, dt = "cuda", torch.float32
+    E, r, out_dim, in_dim = 4, 8, 12, 16
+    base = torch.randn(E, out_dim, in_dim, device=dev, dtype=dt)
+    # single-adapter PEFT layout: A [r*E, in], B [out, r*E]
+    A_flat = torch.randn(r * E, in_dim, device=dev, dtype=dt) * 0.1
+    B_flat = torch.randn(out_dim, r * E, device=dev, dtype=dt) * 0.1
+    scaling = 0.5
+    w_single = MoELoRAMaterialize.apply(base, A_flat, B_flat, scaling)
+
+    # same weights reshaped into the [T=1, E, ...] stacked layout. Single-adapter
+    # PEFT layout is asymmetric: A rows are E-outer/r-inner, B cols r-outer/E-inner.
+    A_stk = A_flat.reshape(E, r, in_dim).unsqueeze(0).contiguous()  # [1,E,r,in]
+    B_stk = (
+        B_flat.reshape(out_dim, r, E).permute(2, 0, 1).unsqueeze(0).contiguous()
+    )  # [1,E,out,r]
+    w_multi = materialize_multi_lora_experts(
+        base, A_stk, B_stk, torch.tensor([scaling], device=dev, dtype=dt)
+    )
+    assert torch.allclose(w_single, w_multi, atol=1e-5, rtol=1e-5)
+
+
+def test_combined_ids_routing():
+    dev = "cuda"
+    T = 3
+    expert = torch.tensor([0, 2, 1, 0], device=dev, dtype=torch.int32)
+    tenant = torch.tensor([1, 0, 2, 1], device=dev, dtype=torch.int32)  # per token
+    token_idx = torch.tensor([0, 1, 2, 3], device=dev, dtype=torch.int32)
+    g = combined_expert_ids(expert, tenant, token_idx, T)
+    # e*T + t
+    assert g.tolist() == [0 * 3 + 1, 2 * 3 + 0, 1 * 3 + 2, 0 * 3 + 1]
+    assert g.dtype == torch.int32
diff --git a/tests/integrations/test_sonicmoe_nvfp4_base.py b/tests/integrations/test_sonicmoe_nvfp4_base.py
new file mode 100644
index 0000000000..5435a70f28
--- /dev/null
+++ b/tests/integrations/test_sonicmoe_nvfp4_base.py
@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""CPU tests for the Phase-1 NVFP4 base-weight helpers of the sonicmoe backend."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from axolotl.integrations.kernels.libs.sonicmoe.nvfp4 import (
+    dequantize_expert_weight,
+    gated_activation,
+    grouped_down_gemm,
+    grouped_up_gemm,
+    is_nvfp4_param,
+    resolve_gated_activation,
+)
+
+
+def _random_offsets(E: int, counts: list[int]) -> torch.Tensor:
+    assert len(counts) == E
+    return torch.tensor(
+        [0] + list(torch.tensor(counts).cumsum(0).tolist()), dtype=torch.int64
+    )
+
+
+def test_is_nvfp4_param_false_for_plain_tensor():
+    assert is_nvfp4_param(torch.randn(4, 8)) is False
+    assert is_nvfp4_param(torch.randn(2, 6, 3)) is False
+
+
+def test_dequantize_identity_for_dense_tensor():
+    w = torch.randn(3, 8, 5)
+    out = dequantize_expert_weight(w)
+    assert out is w
+
+
+def test_grouped_up_gemm_matches_manual_loop():
+    torch.manual_seed(0)
+    E, H, I = 4, 6, 5
+    counts = [3, 0, 2, 4]
+    T = sum(counts)
+    offsets = _random_offsets(E, counts)
+
+    x = torch.randn(T, H, dtype=torch.float32)
+    w1 = torch.randn(E, 2 * I, H, dtype=torch.float32)
+
+    out = grouped_up_gemm(x, w1, offsets, backend="torch", concat=True)
+
+    ref_rows = []
+    start = 0
+    for e in range(E):
+        end = int(offsets[e + 1])
+        ref_rows.append(F.linear(x[start:end], w1[e]))
+        start = end
+    ref = torch.cat(ref_rows, dim=0)
+
+    assert out.shape == (T, 2 * I)
+    torch.testing.assert_close(out, ref)
+
+
+def test_grouped_down_gemm_matches_manual_loop():
+    torch.manual_seed(1)
+    E, H, I = 3, 7, 4
+    counts = [2, 5, 1]
+    T = sum(counts)
+    offsets = _random_offsets(E, counts)
+
+    a = torch.randn(T, I, dtype=torch.float32)
+    w2 = torch.randn(E, H, I, dtype=torch.float32)
+
+    out = grouped_down_gemm(a, w2, offsets, backend="torch")
+
+    ref_rows = []
+    start = 0
+    for e in range(E):
+        end = int(offsets[e + 1])
+        ref_rows.append(F.linear(a[start:end], w2[e]))
+        start = end
+    ref = torch.cat(ref_rows, dim=0)
+
+    assert out.shape == (T, H)
+    torch.testing.assert_close(out, ref)
+
+
+def test_gated_activation_swiglu_concat():
+    torch.manual_seed(2)
+    I = 5
+    h = torch.randn(7, 2 * I, dtype=torch.float32)
+    out = gated_activation(h, "swiglu", concat=True)
+    gate, up = h[..., :I], h[..., I:]
+    ref = up * F.silu(gate.to(torch.float32)).to(gate.dtype)
+    assert out.shape == (7, I)
+    torch.testing.assert_close(out, ref)
+    # silu is an alias for swiglu.
+    torch.testing.assert_close(gated_activation(h, "silu", concat=True), ref)
+
+
+def test_gated_activation_swiglu_interleaved():
+    torch.manual_seed(3)
+    I = 4
+    h = torch.randn(6, 2 * I, dtype=torch.float32)
+    out = gated_activation(h, "swiglu", concat=False)
+    gate, up = h[..., 0::2], h[..., 1::2]
+    ref = up * F.silu(gate.to(torch.float32)).to(gate.dtype)
+    assert out.shape == (6, I)
+    torch.testing.assert_close(out, ref)
+
+
+def test_gated_activation_geglu_concat():
+    torch.manual_seed(4)
+    I = 3
+    h = torch.randn(5, 2 * I, dtype=torch.float32)
+    out = gated_activation(h, "geglu", concat=True)
+    gate, up = h[..., :I], h[..., I:]
+    ref = up * F.gelu(gate.to(torch.float32)).to(gate.dtype)
+    torch.testing.assert_close(out, ref)
+    # gelu is an alias for geglu.
+    torch.testing.assert_close(gated_activation(h, "gelu", concat=True), ref)
+
+
+def test_gated_activation_geglu_interleaved():
+    torch.manual_seed(5)
+    I = 4
+    h = torch.randn(6, 2 * I, dtype=torch.float32)
+    out = gated_activation(h, "geglu", concat=False)
+    gate, up = h[..., 0::2], h[..., 1::2]
+    ref = up * F.gelu(gate.to(torch.float32)).to(gate.dtype)
+    torch.testing.assert_close(out, ref)
+
+
+def test_gated_activation_gelu_tanh_and_reglu():
+    torch.manual_seed(7)
+    I = 4
+    h = torch.randn(6, 2 * I, dtype=torch.float32)
+    gate, up = h[..., :I], h[..., I:]
+
+    # Gemma-style tanh-approx GeGLU (distinct from the erf gelu above).
+    tanh_ref = up * F.gelu(gate, approximate="tanh")
+    torch.testing.assert_close(gated_activation(h, "gelu_tanh", concat=True), tanh_ref)
+    torch.testing.assert_close(
+        gated_activation(h, "gelu_pytorch_tanh", concat=True), tanh_ref
+    )
+    # erf gelu and tanh gelu must not be the same op.
+    assert not torch.allclose(gated_activation(h, "gelu", concat=True), tanh_ref)
+
+    reglu_ref = up * F.relu(gate)
+    torch.testing.assert_close(gated_activation(h, "reglu", concat=True), reglu_ref)
+
+    with pytest.raises(ValueError, match="unsupported gated activation"):
+        gated_activation(h, "mystery", concat=True)
+
+
+def test_gated_activation_clamped_swiglu_limit():
+    torch.manual_seed(8)
+    I = 5
+    # Force values outside [-L, L] so the clamp is exercised.
+    h = torch.randn(7, 2 * I, dtype=torch.float64) * 20.0
+    L = 10.0
+
+    out = gated_activation(h, "swiglu", concat=True, limit=L)
+
+    gate, up = h[..., :I], h[..., I:]
+    gate_c = gate.clamp(max=L)
+    up_c = up.clamp(min=-L, max=L)
+    ref = up_c * F.silu(gate_c)
+    torch.testing.assert_close(out, ref)
+    # limit=None must not clamp.
+    assert not torch.allclose(gated_activation(h, "swiglu", concat=True), ref)
+
+
+def test_resolve_gated_activation_prefers_hidden_activation():
+    # Gemma-style config: hidden_activation set, no usable hidden_act.
+    gemma = SimpleNamespace(hidden_activation="gelu_pytorch_tanh")
+    assert resolve_gated_activation(gemma) == "gelu_pytorch_tanh"
+    # hidden_activation wins even when hidden_act is also present.
+    both = SimpleNamespace(hidden_activation="gelu_pytorch_tanh", hidden_act="silu")
+    assert resolve_gated_activation(both) == "gelu_pytorch_tanh"
+    # Most models: only hidden_act.
+    assert resolve_gated_activation(SimpleNamespace(hidden_act="SiLU")) == "silu"
+    # Neither present -> safe default.
+    assert resolve_gated_activation(SimpleNamespace()) == "silu"
+
+
+def test_full_path_gradcheck_wrt_input():
+    torch.manual_seed(6)
+    E, H, I = 3, 4, 3
+    counts = [2, 1, 2]
+    T = sum(counts)
+    offsets = _random_offsets(E, counts)
+
+    # Keep magnitudes small so the composed swiglu is not stiff for finite differencing.
+    w1 = torch.randn(E, 2 * I, H, dtype=torch.float64) * 0.3
+    w2 = torch.randn(E, H, I, dtype=torch.float64) * 0.3
+    x = (torch.randn(T, H, dtype=torch.float64) * 0.3).requires_grad_(True)
+
+    def fn(x_in):
+        h = grouped_up_gemm(x_in, w1, offsets, backend="torch", concat=True)
+        a = gated_activation(h, "swiglu", concat=True)
+        return grouped_down_gemm(a, w2, offsets, backend="torch")
+
+    assert torch.autograd.gradcheck(fn, (x,))
diff --git a/tests/integrations/test_sonicmoe_nvfp4_forward.py b/tests/integrations/test_sonicmoe_nvfp4_forward.py
new file mode 100644
index 0000000000..c843a9bef6
--- /dev/null
+++ b/tests/integrations/test_sonicmoe_nvfp4_forward.py
@@ -0,0 +1,310 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""End-to-end tests for the sonicmoe NVFP4 grouped forward path.
+
+Validates routing gather/combine + the grouped LoRA MLP orchestration against a
+naive per-token MoE oracle (materialized W_eff), forward and backward, on CPU.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from axolotl.integrations.kernels.libs.sonicmoe.nvfp4_lora import (
+    combine_expert_outputs,
+    grouped_moe_reference_forward,
+    route_and_group,
+)
+
+
+def _make_peft_lora(E, r, dim1, dim2, *, dtype, seed):
+    """Build (lora_A, lora_B) in PEFT rank-major layout + per-expert 3d views."""
+    g = torch.Generator().manual_seed(seed)
+    A_3d = torch.randn(E, r, dim2, generator=g, dtype=dtype, requires_grad=True)
+    B_3d = torch.randn(E, dim1, r, generator=g, dtype=dtype, requires_grad=True)
+    return A_3d, B_3d
+
+
+def _peft_layout(A_3d, B_3d):
+    E, r, dim2 = A_3d.shape
+    _, dim1, _ = B_3d.shape
+    lora_A = A_3d.reshape(E * r, dim2)
+    lora_B = B_3d.permute(1, 2, 0).reshape(dim1, E * r)
+    return lora_A, lora_B
+
+
+def _swiglu_concat(h):
+    i = h.shape[-1] // 2
+    return h[..., i:] * F.silu(h[..., :i])
+
+
+def _oracle_forward(
+    hidden_states,
+    top_k_index,
+    top_k_weights,
+    w1,
+    w2,
+    A1_3d,
+    B1_3d,
+    s1,
+    A2_3d,
+    B2_3d,
+    s2,
+):
+    """Naive per-token MoE with materialized W_eff per expert."""
+    T, K = top_k_index.shape
+    H = hidden_states.shape[-1]
+    out = hidden_states.new_zeros((T, H))
+    for t in range(T):
+        for k in range(K):
+            e = int(top_k_index[t, k])
+            x = hidden_states[t]
+            w1_eff = w1[e] + s1 * (B1_3d[e] @ A1_3d[e])  # [2I, H]
+            h = x @ w1_eff.transpose(0, 1)  # [2I]
+            a = _swiglu_concat(h)  # [I]
+            w2_eff = w2[e] + s2 * (B2_3d[e] @ A2_3d[e])  # [H, I]
+            y = a @ w2_eff.transpose(0, 1)  # [H]
+            out[t] = out[t] + top_k_weights[t, k] * y
+    return out
+
+
+def test_route_and_group_roundtrip():
+    torch.manual_seed(0)
+    T, K, H, E = 6, 2, 4, 3
+    hs = torch.randn(T, H, dtype=torch.float64)
+    idx = torch.randint(0, E, (T, K))
+    w = torch.rand(T, K, dtype=torch.float64)
+
+    x_g, offsets, gather_idx, w_g = route_and_group(hs, idx, w, E)
+
+    assert offsets.shape == (E + 1,)
+    assert int(offsets[-1]) == T * K
+    # rows are sorted by expert
+    counts = torch.bincount(idx.reshape(-1), minlength=E)
+    assert torch.equal(offsets[1:] - offsets[:-1], counts)
+    # combine with unit outputs and unit weights reconstructs per-token top-k counts
+    ones = torch.ones(T * K, H, dtype=torch.float64)
+    combined = combine_expert_outputs(
+        ones, gather_idx, torch.ones(T * K, dtype=torch.float64), T
+    )
+    assert torch.allclose(combined, torch.full((T, H), float(K), dtype=torch.float64))
+
+
+def test_nvfp4_forward_rejects_expert_parallelism():
+    # EP is unfinished: base sharded to local experts, routing/LoRA still global.
+    from types import SimpleNamespace
+
+    from axolotl.integrations.kernels.libs.sonicmoe.experts import (
+        _sonicmoe_nvfp4_forward,
+    )
+
+    fake_self = SimpleNamespace(
+        is_transposed=False, num_experts=4, num_experts_global=8
+    )
+    dummy = torch.zeros(2, 4)
+    with pytest.raises(NotImplementedError, match="expert parallelism"):
+        _sonicmoe_nvfp4_forward(
+            fake_self, dummy, dummy, dummy, dummy, None, dummy, None, None, None
+        )
+
+
+def test_route_and_group_rejects_out_of_range_expert():
+    # A routed id >= num_experts (e.g. global ids under expert parallelism) would
+    # otherwise silently truncate the offsets; the path must reject it loudly.
+    hs = torch.randn(4, 3, dtype=torch.float64)
+    idx = torch.tensor([[0, 5], [1, 2], [0, 1], [2, 0]])  # id 5 >= num_experts=3
+    w = torch.rand(4, 2, dtype=torch.float64)
+    with pytest.raises(NotImplementedError, match="expert parallelism"):
+        route_and_group(hs, idx, w, num_experts=3)
+
+
+def test_nvfp4_forward_matches_oracle_fwd_bwd():
+    torch.manual_seed(1)
+    T, K, H, I, E, r = 5, 2, 6, 4, 3, 2
+    dtype = torch.float64
+    s1, s2 = 0.5, 0.25
+
+    hs = torch.randn(T, H, dtype=dtype, requires_grad=True)
+    idx = torch.randint(0, E, (T, K))
+    tkw = torch.rand(T, K, dtype=dtype, requires_grad=True)
+    w1 = torch.randn(E, 2 * I, H, dtype=dtype)  # frozen base
+    w2 = torch.randn(E, H, I, dtype=dtype)
+
+    A1_3d, B1_3d = _make_peft_lora(E, r, 2 * I, H, dtype=dtype, seed=2)
+    A2_3d, B2_3d = _make_peft_lora(E, r, H, I, dtype=dtype, seed=3)
+
+    # clone leaves for the two independent graphs (module path vs oracle)
+    hs_o = hs.detach().clone().requires_grad_(True)
+    tkw_o = tkw.detach().clone().requires_grad_(True)
+    A1o, B1o = (
+        A1_3d.detach().clone().requires_grad_(True),
+        B1_3d.detach().clone().requires_grad_(True),
+    )
+    A2o, B2o = (
+        A2_3d.detach().clone().requires_grad_(True),
+        B2_3d.detach().clone().requires_grad_(True),
+    )
+
+    lora_A1, lora_B1 = _peft_layout(A1_3d, B1_3d)
+    lora_A2, lora_B2 = _peft_layout(A2_3d, B2_3d)
+
+    out = grouped_moe_reference_forward(
+        hs,
+        idx,
+        tkw,
+        w1,
+        None,
+        w2,
+        None,
+        (lora_A1, lora_B1),
+        (lora_A2, lora_B2),
+        E,
+        act="swiglu",
+        backend="torch",
+        concat=True,
+        scaling1=s1,
+        scaling2=s2,
+    )
+    ref = _oracle_forward(hs_o, idx, tkw_o, w1, w2, A1o, B1o, s1, A2o, B2o, s2)
+
+    assert torch.allclose(out, ref, atol=1e-9), (out - ref).abs().max()
+
+    g = torch.randn_like(out)
+    out.backward(g)
+    ref.backward(g)
+
+    assert torch.allclose(hs.grad, hs_o.grad, atol=1e-8)
+    assert torch.allclose(tkw.grad, tkw_o.grad, atol=1e-8)
+    assert torch.allclose(A1_3d.grad, A1o.grad, atol=1e-8)
+    assert torch.allclose(B1_3d.grad, B1o.grad, atol=1e-8)
+    assert torch.allclose(A2_3d.grad, A2o.grad, atol=1e-8)
+    assert torch.allclose(B2_3d.grad, B2o.grad, atol=1e-8)
+
+
+def test_nvfp4_forward_no_lora_and_frozen_base():
+    torch.manual_seed(4)
+    T, K, H, I, E = 4, 2, 6, 4, 3
+    dtype = torch.float64
+    hs = torch.randn(T, H, dtype=dtype)
+    idx = torch.randint(0, E, (T, K))
+    tkw = torch.rand(T, K, dtype=dtype)
+    w1 = torch.randn(E, 2 * I, H, dtype=dtype, requires_grad=True)
+    w2 = torch.randn(E, H, I, dtype=dtype, requires_grad=True)
+
+    out = grouped_moe_reference_forward(
+        hs,
+        idx,
+        tkw,
+        w1,
+        None,
+        w2,
+        None,
+        None,
+        None,
+        E,
+        act="swiglu",
+        backend="torch",
+        concat=True,
+        scaling1=1.0,
+        scaling2=1.0,
+    )
+    out.sum().backward()
+    # No LoRA Function is invoked, so the base GEMM is a plain differentiable op
+    # and w1 does get a grad here.
+    assert w1.grad is not None
+    # sanity: matches oracle with zero LoRA
+    A0_1 = torch.zeros(E, 1, H, dtype=dtype)
+    B0_1 = torch.zeros(E, 2 * I, 1, dtype=dtype)
+    A0_2 = torch.zeros(E, 1, I, dtype=dtype)
+    B0_2 = torch.zeros(E, H, 1, dtype=dtype)
+    ref = _oracle_forward(
+        hs, idx, tkw, w1.detach(), w2.detach(), A0_1, B0_1, 0.0, A0_2, B0_2, 0.0
+    )
+    assert torch.allclose(out, ref, atol=1e-9)
+
+
+def test_nvfp4_forward_clamped_swiglu_limit():
+    torch.manual_seed(8)
+    T, K, H, I, E = 5, 2, 6, 4, 3
+    dtype = torch.float64
+    L = 10.0
+    # Large weights so the pre-activation exceeds the clamp bound.
+    hs = torch.randn(T, H, dtype=dtype) * 3.0
+    idx = torch.randint(0, E, (T, K))
+    tkw = torch.rand(T, K, dtype=dtype)
+    w1 = torch.randn(E, 2 * I, H, dtype=dtype) * 3.0
+    w2 = torch.randn(E, H, I, dtype=dtype)
+
+    out = grouped_moe_reference_forward(
+        hs,
+        idx,
+        tkw,
+        w1,
+        None,
+        w2,
+        None,
+        None,
+        None,
+        E,
+        act="swiglu",
+        backend="torch",
+        concat=True,
+        scaling1=1.0,
+        scaling2=1.0,
+        limit=L,
+    )
+
+    def clamped_oracle():
+        acc = hs.new_zeros((T, H))
+        for t in range(T):
+            for k in range(K):
+                e = int(idx[t, k])
+                h = hs[t] @ w1[e].transpose(0, 1)
+                gate = h[:I].clamp(max=L)
+                up = h[I:].clamp(min=-L, max=L)
+                a = up * F.silu(gate)
+                acc[t] = acc[t] + tkw[t, k] * (a @ w2[e].transpose(0, 1))
+        return acc
+
+    torch.testing.assert_close(out, clamped_oracle(), atol=1e-9, rtol=1e-9)
+    # Without the limit the outputs differ (clamp actually fires).
+    unclamped = grouped_moe_reference_forward(
+        hs,
+        idx,
+        tkw,
+        w1,
+        None,
+        w2,
+        None,
+        None,
+        None,
+        E,
+        act="swiglu",
+        backend="torch",
+        concat=True,
+        scaling1=1.0,
+        scaling2=1.0,
+    )
+    assert not torch.allclose(out, unclamped)
+
+
+def test_dequant_backend_identity_on_dense():
+    torch.manual_seed(5)
+    T, K, H, I, E = 4, 2, 6, 4, 3
+    dtype = torch.float64
+    hs = torch.randn(T, H, dtype=dtype)
+    idx = torch.randint(0, E, (T, K))
+    tkw = torch.rand(T, K, dtype=dtype)
+    w1 = torch.randn(E, 2 * I, H, dtype=dtype)
+    w2 = torch.randn(E, H, I, dtype=dtype)
+
+    common = dict(act="swiglu", concat=True, scaling1=1.0, scaling2=1.0)
+    y_torch = grouped_moe_reference_forward(
+        hs, idx, tkw, w1, None, w2, None, None, None, E, backend="torch", **common
+    )
+    y_deq = grouped_moe_reference_forward(
+        hs, idx, tkw, w1, None, w2, None, None, None, E, backend="dequant", **common
+    )
+    assert torch.allclose(y_torch, y_deq, atol=1e-12)
diff --git a/tests/integrations/test_sonicmoe_nvfp4_gradients.py b/tests/integrations/test_sonicmoe_nvfp4_gradients.py
new file mode 100644
index 0000000000..5da67ad0e0
--- /dev/null
+++ b/tests/integrations/test_sonicmoe_nvfp4_gradients.py
@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""Gradient correctness tests for the sonicmoe NVFP4 frozen-base grouped LoRA.
+
+Runs entirely on CPU.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from axolotl.integrations.kernels.libs.sonicmoe.nvfp4_lora import (
+    GroupedDownProjLoRA,
+    GroupedUpProjLoRA,
+    grouped_expert_mlp_lora,
+)
+
+
+def _ref_gated_activation(h, act, *, concat):
+    """SwiGLU gated activation over the last (``2I``) dim -> ``I``.
+
+    ``concat=True``: split-halves (``[gate | up]``). ``concat=False``:
+    interleaved (``gate`` = even lanes, ``up`` = odd lanes).
+    """
+    assert act.lower() == "silu"
+    if concat:
+        gate, up = h.chunk(2, dim=-1)
+    else:
+        gate, up = h[..., 0::2], h[..., 1::2]
+    return F.silu(gate) * up
+
+
+# =============================================================================
+# Helpers
+# =============================================================================
+
+
+def _rand_offsets(E, T, device="cpu"):
+    """Random ascending expert offsets covering ``T`` tokens across ``E`` experts."""
+    counts = torch.zeros(E, dtype=torch.long)
+    for _ in range(T):
+        counts[torch.randint(0, E, (1,)).item()] += 1
+    offsets = torch.zeros(E + 1, dtype=torch.long, device=device)
+    offsets[1:] = torch.cumsum(counts, 0)
+    return offsets
+
+
+def _make_lora(E, dim1, dim2, r, dtype):
+    """PEFT rank-major LoRA: A ``[r*E, dim2]``, B ``[dim1, r*E]``."""
+    A = torch.randn(r * E, dim2, dtype=dtype, requires_grad=True)
+    B = torch.randn(dim1, r * E, dtype=dtype, requires_grad=True)
+    return A, B
+
+
+def _materialize_expert(W_e, A, B, e, r, E, scaling):
+    """W_eff_e = W_e + scaling * (B_e @ A_e) in PEFT rank-major layout."""
+    dim1, dim2 = W_e.shape
+    A_e = A.reshape(E, r, dim2)[e]  # [r, dim2]
+    B_e = B.reshape(dim1, r, E).permute(2, 0, 1)[e]  # [dim1, r]
+    return W_e + scaling * (B_e @ A_e)
+
+
+def _oracle_moe(
+    x_grouped, expert_offsets, w1, b1, w2, b2, A1, B1, A2, B2, r, E, s1, s2, concat
+):
+    """Naive per-expert MoE on MATERIALIZED W_eff (up -> swiglu -> down)."""
+    y = x_grouped.new_zeros((x_grouped.shape[0], w2.shape[1]))
+    for e in range(E):
+        st, en = int(expert_offsets[e]), int(expert_offsets[e + 1])
+        if en <= st:
+            continue
+        x_e = x_grouped[st:en]
+        W1e = _materialize_expert(w1[e], A1, B1, e, r, E, s1)
+        h_e = F.linear(x_e, W1e)  # x_e @ W1e^T
+        if b1 is not None:
+            h_e = h_e + b1[e]
+        a_e = _ref_gated_activation(h_e, "silu", concat=concat)
+        W2e = _materialize_expert(w2[e], A2, B2, e, r, E, s2)
+        y_e = F.linear(a_e, W2e)  # a_e @ W2e^T
+        if b2 is not None:
+            y_e = y_e + b2[e]
+        y[st:en] = y_e
+    return y
+
+
+# =============================================================================
+# Oracle test: forward + hand-written backward vs materialized autograd oracle
+# =============================================================================
+
+
+@pytest.mark.parametrize("concat", [True, False])
+def test_oracle_forward_and_backward(concat):
+    torch.manual_seed(0)
+    dtype = torch.float64
+    E, H, I, r, T = 3, 8, 6, 2, 17  # noqa: E741
+    two_i = 2 * I
+    s1, s2 = 0.5, 0.75
+
+    offsets = _rand_offsets(E, T)
+
+    w1 = torch.randn(E, two_i, H, dtype=dtype)  # frozen base
+    w2 = torch.randn(E, H, I, dtype=dtype)  # frozen base
+    b1 = torch.randn(E, two_i, dtype=dtype)
+    b2 = torch.randn(E, H, dtype=dtype)
+
+    A1, B1 = _make_lora(E, two_i, H, r, dtype)
+    A2, B2 = _make_lora(E, H, I, r, dtype)
+
+    x = torch.randn(T, H, dtype=dtype, requires_grad=True)
+
+    # ---- grouped path ----
+    y = grouped_expert_mlp_lora(
+        x,
+        offsets,
+        w1,
+        b1,
+        w2,
+        b2,
+        (A1, B1),
+        (A2, B2),
+        act="silu",
+        backend="torch",
+        concat=concat,
+        scaling1=s1,
+        scaling2=s2,
+    )
+
+    # ---- oracle path ----
+    xo = x.detach().clone().requires_grad_(True)
+    A1o = A1.detach().clone().requires_grad_(True)
+    B1o = B1.detach().clone().requires_grad_(True)
+    A2o = A2.detach().clone().requires_grad_(True)
+    B2o = B2.detach().clone().requires_grad_(True)
+    yo = _oracle_moe(
+        xo, offsets, w1, b1, w2, b2, A1o, B1o, A2o, B2o, r, E, s1, s2, concat
+    )
+
+    assert torch.allclose(y, yo, atol=1e-10), (y - yo).abs().max()
+
+    # Backward with a random cotangent so every output element contributes.
+    g = torch.randn_like(y)
+    y.backward(g)
+    yo.backward(g.clone())
+
+    assert torch.allclose(x.grad, xo.grad, atol=1e-9), (x.grad - xo.grad).abs().max()
+    assert torch.allclose(A1.grad, A1o.grad, atol=1e-9), (
+        (A1.grad - A1o.grad).abs().max()
+    )
+    assert torch.allclose(B1.grad, B1o.grad, atol=1e-9), (
+        (B1.grad - B1o.grad).abs().max()
+    )
+    assert torch.allclose(A2.grad, A2o.grad, atol=1e-9), (
+        (A2.grad - A2o.grad).abs().max()
+    )
+    assert torch.allclose(B2.grad, B2o.grad, atol=1e-9), (
+        (B2.grad - B2o.grad).abs().max()
+    )
+
+
+# =============================================================================
+# gradcheck on the full chain
+# =============================================================================
+
+
+@pytest.mark.parametrize("concat", [True, False])
+def test_gradcheck_full_chain(concat):
+    torch.manual_seed(1)
+    dtype = torch.float64
+    E, H, I, r, T = 2, 6, 4, 2, 9  # noqa: E741
+    two_i = 2 * I
+    s1, s2 = 0.5, 1.0
+
+    offsets = _rand_offsets(E, T)
+    w1 = torch.randn(E, two_i, H, dtype=dtype)
+    w2 = torch.randn(E, H, I, dtype=dtype)
+
+    A1, B1 = _make_lora(E, two_i, H, r, dtype)
+    A2, B2 = _make_lora(E, H, I, r, dtype)
+    x = torch.randn(T, H, dtype=dtype, requires_grad=True)
+
+    def fn(x_, A1_, B1_, A2_, B2_):
+        return grouped_expert_mlp_lora(
+            x_,
+            offsets,
+            w1,
+            None,
+            w2,
+            None,
+            (A1_, B1_),
+            (A2_, B2_),
+            act="silu",
+            backend="torch",
+            concat=concat,
+            scaling1=s1,
+            scaling2=s2,
+        )
+
+    assert torch.autograd.gradcheck(
+        fn, (x, A1, B1, A2, B2), eps=1e-6, atol=1e-5, rtol=1e-3
+    )
+
+
+# =============================================================================
+# None LoRA reduces to plain base MoE; base weights receive no gradient
+# =============================================================================
+
+
+@pytest.mark.parametrize("concat", [True, False])
+def test_none_lora_reduces_to_base(concat):
+    torch.manual_seed(2)
+    dtype = torch.float64
+    E, H, I, T = 3, 8, 6, 13  # noqa: E741
+    two_i = 2 * I
+
+    offsets = _rand_offsets(E, T)
+    w1 = torch.randn(E, two_i, H, dtype=dtype, requires_grad=True)
+    w2 = torch.randn(E, H, I, dtype=dtype, requires_grad=True)
+    x = torch.randn(T, H, dtype=dtype, requires_grad=True)
+
+    y = grouped_expert_mlp_lora(
+        x,
+        offsets,
+        w1,
+        None,
+        w2,
+        None,
+        None,
+        None,
+        act="silu",
+        backend="torch",
+        concat=concat,
+        scaling1=1.0,
+        scaling2=1.0,
+    )
+
+    # Naive base-only MoE.
+    expected = x.new_zeros((T, H))
+    for e in range(E):
+        s, t = int(offsets[e]), int(offsets[e + 1])
+        if t <= s:
+            continue
+        h_e = F.linear(x[s:t], w1[e])
+        a_e = _ref_gated_activation(h_e, "silu", concat=concat)
+        expected[s:t] = F.linear(a_e, w2[e])
+
+    assert torch.allclose(y, expected, atol=1e-10)
+
+    y.sum().backward()
+    # No LoRA Function is invoked here, so the base grouped GEMM is a plain
+    # differentiable op and w1/w2 do get grads. We only assert the backward runs.
+    assert x.grad is not None
+
+
+def test_lora_base_weight_no_grad():
+    """The LoRA autograd.Functions must NOT produce a gradient for the frozen base."""
+    torch.manual_seed(3)
+    dtype = torch.float64
+    E, H, I, r, T = 2, 6, 4, 2, 7  # noqa: E741
+    two_i = 2 * I
+
+    offsets = _rand_offsets(E, T)
+    # Base requires_grad=True to prove the Function still returns None for it.
+    w1 = torch.randn(E, two_i, H, dtype=dtype, requires_grad=True)
+    A1, B1 = _make_lora(E, two_i, H, r, dtype)
+    x = torch.randn(T, H, dtype=dtype, requires_grad=True)
+
+    h = GroupedUpProjLoRA.apply(x, w1, offsets, A1, B1, 0.5, "torch", True)
+    h.sum().backward()
+
+    assert w1.grad is None
+    assert A1.grad is not None
+    assert B1.grad is not None
+    assert x.grad is not None
+
+
+def test_down_proj_base_weight_no_grad():
+    torch.manual_seed(4)
+    dtype = torch.float64
+    E, H, I, r, T = 2, 6, 4, 2, 7  # noqa: E741
+
+    offsets = _rand_offsets(E, T)
+    w2 = torch.randn(E, H, I, dtype=dtype, requires_grad=True)
+    A2, B2 = _make_lora(E, H, I, r, dtype)
+    a = torch.randn(T, I, dtype=dtype, requires_grad=True)
+
+    y = GroupedDownProjLoRA.apply(a, w2, offsets, A2, B2, 0.75, "torch")
+    y.sum().backward()
+
+    assert w2.grad is None
+    assert A2.grad is not None
+    assert B2.grad is not None
+    assert a.grad is not None
diff --git a/tests/integrations/test_swanlab.py b/tests/integrations/test_swanlab.py
new file mode 100644
index 0000000000..7158f8d7ea
--- /dev/null
+++ b/tests/integrations/test_swanlab.py
@@ -0,0 +1,1337 @@
+# Copyright 2024 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Unit tests for SwanLab Integration Plugin.
+
+Tests conflict detection, configuration validation, and multi-logger warnings.
+"""
+
+import importlib.util
+import logging
+import os
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+from pydantic import ValidationError
+
+from axolotl.integrations.swanlab.args import SwanLabConfig
+from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+
+SWANLAB_INSTALLED = importlib.util.find_spec("swanlab") is not None
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestSwanLabConfigValidators:
+    """Tests for Pydantic field validators in SwanLabConfig."""
+
+    def test_valid_swanlab_mode_cloud(self):
+        """Test that 'cloud' mode is valid."""
+        config = SwanLabConfig(swanlab_mode="cloud")
+        assert config.swanlab_mode == "cloud"
+
+    def test_valid_swanlab_mode_local(self):
+        """Test that 'local' mode is valid."""
+        config = SwanLabConfig(swanlab_mode="local")
+        assert config.swanlab_mode == "local"
+
+    def test_valid_swanlab_mode_offline(self):
+        """Test that 'offline' mode is valid."""
+        config = SwanLabConfig(swanlab_mode="offline")
+        assert config.swanlab_mode == "offline"
+
+    def test_valid_swanlab_mode_disabled(self):
+        """Test that 'disabled' mode is valid."""
+        config = SwanLabConfig(swanlab_mode="disabled")
+        assert config.swanlab_mode == "disabled"
+
+    def test_invalid_swanlab_mode(self):
+        """Test that invalid mode raises ValueError."""
+        with pytest.raises(ValidationError) as exc_info:
+            SwanLabConfig(swanlab_mode="invalid")
+
+        error_msg = str(exc_info.value)
+        assert "Invalid swanlab_mode" in error_msg
+        assert "cloud" in error_msg
+        assert "local" in error_msg
+        assert "offline" in error_msg
+        assert "disabled" in error_msg
+
+    def test_swanlab_mode_none_allowed(self):
+        """Test that None mode is allowed (will use default)."""
+        config = SwanLabConfig(swanlab_mode=None)
+        assert config.swanlab_mode is None
+
+    def test_valid_swanlab_project(self):
+        """Test that valid project name is accepted."""
+        config = SwanLabConfig(swanlab_project="my-project")
+        assert config.swanlab_project == "my-project"
+
+    def test_swanlab_project_none_allowed(self):
+        """Test that None project is allowed."""
+        config = SwanLabConfig(swanlab_project=None)
+        assert config.swanlab_project is None
+
+    def test_empty_swanlab_project_rejected(self):
+        """Test that empty string project name is rejected."""
+        with pytest.raises(ValidationError) as exc_info:
+            SwanLabConfig(swanlab_project="")
+
+        error_msg = str(exc_info.value)
+        assert "cannot be an empty string" in error_msg
+
+    def test_whitespace_only_project_rejected(self):
+        """Test that whitespace-only project name is rejected."""
+        with pytest.raises(ValidationError) as exc_info:
+            SwanLabConfig(swanlab_project="   ")
+
+        error_msg = str(exc_info.value)
+        assert "cannot be an empty string" in error_msg
+
+    def test_use_swanlab_true_requires_project(self):
+        """Test that use_swanlab=True requires swanlab_project."""
+        with pytest.raises(ValidationError) as exc_info:
+            SwanLabConfig(use_swanlab=True, swanlab_project=None)
+
+        error_msg = str(exc_info.value)
+        assert "swanlab_project" in error_msg.lower()
+        assert "not set" in error_msg.lower()
+
+    def test_use_swanlab_true_with_project_valid(self):
+        """Test that use_swanlab=True with project is valid."""
+        config = SwanLabConfig(use_swanlab=True, swanlab_project="my-project")
+        assert config.use_swanlab is True
+        assert config.swanlab_project == "my-project"
+
+    def test_use_swanlab_false_no_project_valid(self):
+        """Test that use_swanlab=False without project is valid."""
+        config = SwanLabConfig(use_swanlab=False, swanlab_project=None)
+        assert config.use_swanlab is False
+        assert config.swanlab_project is None
+
+    def test_use_swanlab_none_no_project_valid(self):
+        """Test that use_swanlab=None without project is valid."""
+        config = SwanLabConfig(use_swanlab=None, swanlab_project=None)
+        assert config.use_swanlab is None
+        assert config.swanlab_project is None
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestSwanLabPluginRegister:
+    """Tests for SwanLabPlugin.register() conflict detection."""
+
+    def test_register_without_use_swanlab(self):
+        """Test that register works when SwanLab is not enabled."""
+        plugin = SwanLabPlugin()
+        cfg = {"use_swanlab": False}
+        # Should not raise
+        plugin.register(cfg)
+
+    def test_register_use_swanlab_missing_project(self):
+        """Test that use_swanlab=True without project raises ValueError."""
+        plugin = SwanLabPlugin()
+        cfg = {"use_swanlab": True}
+
+        with pytest.raises(ValueError) as exc_info:
+            plugin.register(cfg)
+
+        error_msg = str(exc_info.value)
+        assert "swanlab_project" in error_msg
+        assert "not set" in error_msg
+        assert "Solutions" in error_msg
+
+    def test_register_use_swanlab_with_project_valid(self):
+        """Test that use_swanlab=True with project is valid."""
+        plugin = SwanLabPlugin()
+        cfg = {"use_swanlab": True, "swanlab_project": "my-project"}
+        # Should not raise
+        plugin.register(cfg)
+
+    def test_register_invalid_mode(self):
+        """Test that invalid swanlab_mode raises ValueError."""
+        plugin = SwanLabPlugin()
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "my-project",
+            "swanlab_mode": "invalid-mode",
+        }
+
+        with pytest.raises(ValueError) as exc_info:
+            plugin.register(cfg)
+
+        error_msg = str(exc_info.value)
+        assert "Invalid swanlab_mode" in error_msg
+        assert "cloud" in error_msg
+        assert "local" in error_msg
+
+    def test_register_valid_modes(self):
+        """Test that all valid modes are accepted."""
+        plugin = SwanLabPlugin()
+        valid_modes = ["cloud", "local", "offline", "disabled"]
+
+        for mode in valid_modes:
+            cfg = {
+                "use_swanlab": True,
+                "swanlab_project": "my-project",
+                "swanlab_mode": mode,
+            }
+            # Should not raise
+            plugin.register(cfg)
+
+    def test_register_auto_enable_swanlab(self):
+        """Test that providing swanlab_project auto-enables use_swanlab."""
+        plugin = SwanLabPlugin()
+        cfg = {"swanlab_project": "my-project"}
+
+        plugin.register(cfg)
+
+        assert cfg["use_swanlab"] is True
+
+    def test_register_cloud_mode_without_api_key_warns(self, caplog):
+        """Test that cloud mode without API key logs warning."""
+        plugin = SwanLabPlugin()
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "my-project",
+            "swanlab_mode": "cloud",
+        }
+
+        # Clear environment variable to ensure it's not set
+        with patch.dict(os.environ, {}, clear=True):
+            with caplog.at_level(logging.WARNING):
+                plugin.register(cfg)
+
+            # Should log warning about missing API key
+            warning_messages = [record.message for record in caplog.records]
+            assert any("API key" in msg for msg in warning_messages)
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestMultiLoggerDetection:
+    """Tests for multi-logger conflict detection."""
+
+    def test_single_logger_no_warning(self, caplog):
+        """Test that single logger doesn't trigger warning."""
+        plugin = SwanLabPlugin()
+        cfg = {"use_swanlab": True, "swanlab_project": "my-project"}
+
+        with caplog.at_level(logging.WARNING):
+            plugin.register(cfg)
+
+        # Should not log multi-logger warning
+        warning_messages = [record.message for record in caplog.records]
+        assert not any("Multiple logging tools" in msg for msg in warning_messages)
+
+    def test_two_loggers_warning(self, caplog):
+        """Test that two loggers trigger warning."""
+        plugin = SwanLabPlugin()
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "my-project",
+            "use_wandb": True,
+        }
+
+        with caplog.at_level(logging.WARNING):
+            plugin.register(cfg)
+
+        # Should log multi-logger warning
+        warning_messages = [record.message for record in caplog.records]
+        assert any("Multiple logging tools" in msg for msg in warning_messages)
+        assert any("SwanLab" in msg and "WandB" in msg for msg in warning_messages)
+
+    def test_three_loggers_error(self, caplog):
+        """Test that three loggers trigger error-level warning."""
+        plugin = SwanLabPlugin()
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "my-project",
+            "use_wandb": True,
+            "use_mlflow": True,
+        }
+
+        with caplog.at_level(logging.ERROR):
+            plugin.register(cfg)
+
+        # Should log error-level warning
+        error_messages = [
+            record.message
+            for record in caplog.records
+            if record.levelno >= logging.ERROR
+        ]
+        assert any("logging tools enabled" in msg for msg in error_messages)
+
+    def test_multi_logger_with_comet(self, caplog):
+        """Test that Comet is detected in multi-logger scenario."""
+        plugin = SwanLabPlugin()
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "my-project",
+            "comet_api_key": "test-key",
+        }
+
+        with caplog.at_level(logging.WARNING):
+            plugin.register(cfg)
+
+        # Should detect Comet
+        warning_messages = [record.message for record in caplog.records]
+        assert any("Comet" in msg for msg in warning_messages)
+
+    def test_multi_logger_with_comet_project(self, caplog):
+        """Test that Comet is detected via comet_project_name."""
+        plugin = SwanLabPlugin()
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "my-project",
+            "comet_project_name": "test-project",
+        }
+
+        with caplog.at_level(logging.WARNING):
+            plugin.register(cfg)
+
+        # Should detect Comet
+        warning_messages = [record.message for record in caplog.records]
+        assert any("Comet" in msg for msg in warning_messages)
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestSwanLabPluginPreModelLoad:
+    """Tests for SwanLabPlugin.pre_model_load() runtime checks."""
+
+    def test_pre_model_load_disabled(self):
+        """Test that pre_model_load does nothing when SwanLab is disabled."""
+        plugin = SwanLabPlugin()
+        cfg = MagicMock()
+        cfg.use_swanlab = False
+
+        # Should not raise
+        plugin.pre_model_load(cfg)
+
+    def test_pre_model_load_import_error(self):
+        """Test that missing swanlab package raises clear ImportError."""
+        plugin = SwanLabPlugin()
+        cfg = MagicMock()
+        cfg.use_swanlab = True
+
+        with patch(
+            "builtins.__import__", side_effect=ImportError("No module named 'swanlab'")
+        ):
+            with pytest.raises(ImportError) as exc_info:
+                plugin.pre_model_load(cfg)
+
+            error_msg = str(exc_info.value)
+            assert "SwanLab is not installed" in error_msg
+            assert "pip install swanlab" in error_msg
+
+    @patch("axolotl.utils.distributed.is_main_process")
+    @patch("axolotl.utils.distributed.get_world_size")
+    def test_pre_model_load_non_main_process_skips(
+        self, mock_get_world_size, mock_is_main_process
+    ):
+        """Test that non-main process skips SwanLab initialization."""
+        mock_get_world_size.return_value = 2
+        mock_is_main_process.return_value = False
+
+        plugin = SwanLabPlugin()
+        cfg = MagicMock()
+        cfg.use_swanlab = True
+
+        with patch("swanlab.init") as mock_init:
+            plugin.pre_model_load(cfg)
+            # Should NOT call swanlab.init
+            mock_init.assert_not_called()
+
+    @patch("axolotl.utils.distributed.is_main_process")
+    @patch("axolotl.utils.distributed.get_world_size")
+    def test_pre_model_load_distributed_logging(
+        self, mock_get_world_size, mock_is_main_process, caplog
+    ):
+        """Test that distributed training logs world size info."""
+        mock_get_world_size.return_value = 4
+        mock_is_main_process.return_value = True
+
+        plugin = SwanLabPlugin()
+        cfg = MagicMock()
+        cfg.use_swanlab = True
+        cfg.swanlab_project = "test-project"
+        cfg.swanlab_mode = "cloud"
+
+        with patch("swanlab.init"), patch("swanlab.__version__", "0.3.0"):
+            with caplog.at_level(logging.INFO):
+                plugin.pre_model_load(cfg)
+
+            # Should log distributed training info
+            info_messages = [record.message for record in caplog.records]
+            assert any("world_size=4" in msg for msg in info_messages)
+            assert any("Only rank 0" in msg for msg in info_messages)
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestSwanLabInitKwargs:
+    """Tests for SwanLab initialization with direct parameter passing."""
+
+    def test_custom_branding_added_to_config(self):
+        """Test that Axolotl custom branding is added to SwanLab config."""
+        from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+        from axolotl.utils.dict import DictDefault
+
+        plugin = SwanLabPlugin()
+        cfg = DictDefault(
+            {
+                "use_swanlab": True,
+                "swanlab_project": "test-project",
+            }
+        )
+
+        init_kwargs = plugin._get_swanlab_init_kwargs(cfg)
+
+        # Verify custom branding is present
+        assert "config" in init_kwargs
+        assert init_kwargs["config"]["UPPERFRAME"] == "🦎 Axolotl"
+
+    def test_api_key_passed_directly(self):
+        """Test that API key is passed directly to swanlab.init() instead of via env var."""
+        from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+        from axolotl.utils.dict import DictDefault
+
+        plugin = SwanLabPlugin()
+        cfg = DictDefault(
+            {
+                "use_swanlab": True,
+                "swanlab_project": "test-project",
+                "swanlab_api_key": "test-api-key-12345",
+            }
+        )
+
+        init_kwargs = plugin._get_swanlab_init_kwargs(cfg)
+
+        # Verify API key is in init_kwargs (not set as env var)
+        assert "api_key" in init_kwargs
+        assert init_kwargs["api_key"] == "test-api-key-12345"
+
+    def test_private_deployment_hosts_passed_directly(self):
+        """Test that private deployment hosts are passed directly to swanlab.init()."""
+        from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+        from axolotl.utils.dict import DictDefault
+
+        plugin = SwanLabPlugin()
+        cfg = DictDefault(
+            {
+                "use_swanlab": True,
+                "swanlab_project": "internal-project",
+                "swanlab_web_host": "https://swanlab.company.com",
+                "swanlab_api_host": "https://api-swanlab.company.com",
+            }
+        )
+
+        init_kwargs = plugin._get_swanlab_init_kwargs(cfg)
+
+        # Verify private deployment hosts are in init_kwargs
+        assert "web_host" in init_kwargs
+        assert init_kwargs["web_host"] == "https://swanlab.company.com"
+        assert "api_host" in init_kwargs
+        assert init_kwargs["api_host"] == "https://api-swanlab.company.com"
+
+    @patch("axolotl.utils.distributed.is_main_process")
+    def test_full_private_deployment_init(self, mock_is_main_process):
+        """Test complete initialization with private deployment configuration."""
+        mock_is_main_process.return_value = True
+
+        from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+        from axolotl.utils.dict import DictDefault
+
+        plugin = SwanLabPlugin()
+        cfg = DictDefault(
+            {
+                "use_swanlab": True,
+                "swanlab_project": "secure-project",
+                "swanlab_experiment_name": "experiment-001",
+                "swanlab_mode": "cloud",
+                "swanlab_api_key": "private-key-xyz",
+                "swanlab_web_host": "https://swanlab.internal.net",
+                "swanlab_api_host": "https://api.swanlab.internal.net",
+                "swanlab_workspace": "research-team",
+            }
+        )
+
+        with patch("swanlab.init") as mock_init:
+            plugin.pre_model_load(cfg)
+
+            # Verify swanlab.init was called with all parameters
+            mock_init.assert_called_once()
+            call_kwargs = mock_init.call_args[1]
+
+            assert call_kwargs["project"] == "secure-project"
+            assert call_kwargs["experiment_name"] == "experiment-001"
+            assert call_kwargs["mode"] == "cloud"
+            assert call_kwargs["api_key"] == "private-key-xyz"
+            assert call_kwargs["web_host"] == "https://swanlab.internal.net"
+            assert call_kwargs["api_host"] == "https://api.swanlab.internal.net"
+            assert call_kwargs["workspace"] == "research-team"
+            assert call_kwargs["config"]["UPPERFRAME"] == "🦎 Axolotl"
+
+    def test_env_vars_not_set_for_api_params(self):
+        """Test that environment variables are NOT set for API parameters."""
+        import os
+
+        from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+        from axolotl.utils.dict import DictDefault
+
+        # Clear any existing env vars
+        for key in [
+            "SWANLAB_API_KEY",
+            "SWANLAB_WEB_HOST",
+            "SWANLAB_API_HOST",
+            "SWANLAB_MODE",
+        ]:
+            os.environ.pop(key, None)
+
+        plugin = SwanLabPlugin()
+        cfg = DictDefault(
+            {
+                "use_swanlab": True,
+                "swanlab_project": "test-project",
+                "swanlab_api_key": "test-key",
+                "swanlab_web_host": "https://test.com",
+                "swanlab_api_host": "https://api-test.com",
+                "swanlab_mode": "cloud",
+            }
+        )
+
+        with (
+            patch("axolotl.utils.distributed.is_main_process", return_value=True),
+            patch("swanlab.init"),
+        ):
+            plugin.pre_model_load(cfg)
+
+        # Verify env vars were NOT set (simplified approach)
+        # The old _setup_swanlab_env() method is removed, so these shouldn't be set
+        # Note: SwanLab itself might set these, but our plugin shouldn't
+        # We're just testing that our plugin doesn't call _setup_swanlab_env()
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestLarkNotificationIntegration:
+    """Tests for Lark (Feishu) notification integration."""
+
+    def test_lark_callback_registration_with_webhook_only(self):
+        """Test Lark callback registration with webhook URL only (no secret)."""
+        plugin = SwanLabPlugin()
+
+        cfg = MagicMock()
+        cfg.use_swanlab = True
+        cfg.swanlab_project = "test-project"
+        cfg.swanlab_mode = "local"
+        cfg.swanlab_lark_webhook_url = (
+            "https://open.feishu.cn/open-apis/bot/v2/hook/test-webhook"
+        )
+        cfg.swanlab_lark_secret = None
+
+        with (
+            patch("swanlab.init"),
+            patch("swanlab.__version__", "0.3.0"),
+            patch("swanlab.register_callbacks") as mock_register,
+            patch("axolotl.utils.distributed.is_main_process", return_value=True),
+            patch("axolotl.utils.distributed.get_world_size", return_value=1),
+        ):
+            # Mock LarkCallback import
+            with patch("swanlab.plugin.notification.LarkCallback") as MockLarkCallback:
+                mock_lark_instance = MagicMock()
+                MockLarkCallback.return_value = mock_lark_instance
+
+                plugin.pre_model_load(cfg)
+
+                # Verify LarkCallback was instantiated with correct params
+                MockLarkCallback.assert_called_once_with(
+                    webhook_url="https://open.feishu.cn/open-apis/bot/v2/hook/test-webhook",
+                    secret=None,
+                )
+
+                # Verify callback was registered
+                mock_register.assert_called_once_with([mock_lark_instance])
+
+    def test_lark_callback_registration_with_secret(self):
+        """Test Lark callback registration with webhook URL and HMAC secret."""
+        plugin = SwanLabPlugin()
+
+        cfg = MagicMock()
+        cfg.use_swanlab = True
+        cfg.swanlab_project = "test-project"
+        cfg.swanlab_mode = "local"
+        cfg.swanlab_lark_webhook_url = (
+            "https://open.feishu.cn/open-apis/bot/v2/hook/test-webhook"
+        )
+        cfg.swanlab_lark_secret = "test-hmac-secret"
+
+        with (
+            patch("swanlab.init"),
+            patch("swanlab.__version__", "0.3.0"),
+            patch("swanlab.register_callbacks") as mock_register,
+            patch("axolotl.utils.distributed.is_main_process", return_value=True),
+            patch("axolotl.utils.distributed.get_world_size", return_value=1),
+        ):
+            with patch("swanlab.plugin.notification.LarkCallback") as MockLarkCallback:
+                mock_lark_instance = MagicMock()
+                MockLarkCallback.return_value = mock_lark_instance
+
+                plugin.pre_model_load(cfg)
+
+                # Verify LarkCallback was instantiated with secret
+                MockLarkCallback.assert_called_once_with(
+                    webhook_url="https://open.feishu.cn/open-apis/bot/v2/hook/test-webhook",
+                    secret="test-hmac-secret",
+                )
+
+                mock_register.assert_called_once_with([mock_lark_instance])
+
+    def test_lark_callback_not_registered_without_webhook(self):
+        """Test that Lark callback is NOT registered when webhook URL not provided."""
+        plugin = SwanLabPlugin()
+
+        cfg = MagicMock()
+        cfg.use_swanlab = True
+        cfg.swanlab_project = "test-project"
+        cfg.swanlab_mode = "local"
+        cfg.swanlab_lark_webhook_url = None  # No webhook
+        cfg.swanlab_lark_secret = None
+
+        with (
+            patch("swanlab.init"),
+            patch("swanlab.__version__", "0.3.0"),
+            patch("swanlab.register_callbacks") as mock_register,
+            patch("axolotl.utils.distributed.is_main_process", return_value=True),
+            patch("axolotl.utils.distributed.get_world_size", return_value=1),
+        ):
+            plugin.pre_model_load(cfg)
+
+            # Verify register_callbacks was NOT called
+            mock_register.assert_not_called()
+
+    def test_lark_import_error_handled_gracefully(self, caplog):
+        """Test that ImportError for Lark plugin is handled gracefully."""
+        plugin = SwanLabPlugin()
+
+        cfg = MagicMock()
+        cfg.use_swanlab = True
+        cfg.swanlab_project = "test-project"
+        cfg.swanlab_mode = "local"
+        cfg.swanlab_lark_webhook_url = (
+            "https://open.feishu.cn/open-apis/bot/v2/hook/test-webhook"
+        )
+        cfg.swanlab_lark_secret = None
+
+        with (
+            patch("swanlab.init"),
+            patch("swanlab.__version__", "0.3.0"),
+            patch("axolotl.utils.distributed.is_main_process", return_value=True),
+            patch("axolotl.utils.distributed.get_world_size", return_value=1),
+        ):
+            # Mock ImportError for LarkCallback
+            with patch(
+                "swanlab.plugin.notification.LarkCallback",
+                side_effect=ImportError(
+                    "No module named 'swanlab.plugin.notification'"
+                ),
+            ):
+                with caplog.at_level(logging.WARNING):
+                    plugin.pre_model_load(cfg)
+
+                    # Should log warning about missing Lark plugin
+                    warning_messages = [record.message for record in caplog.records]
+                    assert any(
+                        "Failed to import SwanLab Lark plugin" in msg
+                        for msg in warning_messages
+                    )
+                    assert any("SwanLab >= 0.3.0" in msg for msg in warning_messages)
+
+    def test_lark_warning_for_missing_secret(self, caplog):
+        """Test that warning is logged when Lark webhook has no HMAC secret."""
+        plugin = SwanLabPlugin()
+
+        cfg = MagicMock()
+        cfg.use_swanlab = True
+        cfg.swanlab_project = "test-project"
+        cfg.swanlab_mode = "local"
+        cfg.swanlab_lark_webhook_url = (
+            "https://open.feishu.cn/open-apis/bot/v2/hook/test-webhook"
+        )
+        cfg.swanlab_lark_secret = None  # No secret
+
+        with (
+            patch("swanlab.init"),
+            patch("swanlab.__version__", "0.3.0"),
+            patch("swanlab.register_callbacks"),
+            patch("axolotl.utils.distributed.is_main_process", return_value=True),
+            patch("axolotl.utils.distributed.get_world_size", return_value=1),
+        ):
+            with patch("swanlab.plugin.notification.LarkCallback"):
+                with caplog.at_level(logging.WARNING):
+                    plugin.pre_model_load(cfg)
+
+                    # Should log warning about missing secret
+                    warning_messages = [record.message for record in caplog.records]
+                    assert any(
+                        "no secret configured" in msg.lower()
+                        for msg in warning_messages
+                    )
+                    assert any("swanlab_lark_secret" in msg for msg in warning_messages)
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestSwanLabPluginIntegration:
+    """Integration tests for SwanLab plugin lifecycle."""
+
+    def test_full_lifecycle_valid_config(self):
+        """Test full plugin lifecycle with valid configuration."""
+        plugin = SwanLabPlugin()
+
+        # Register
+        cfg_dict = {
+            "use_swanlab": True,
+            "swanlab_project": "test-project",
+            "swanlab_mode": "local",
+        }
+        plugin.register(cfg_dict)
+
+        # Pre-model load (mock SwanLab)
+        cfg_obj = MagicMock()
+        cfg_obj.use_swanlab = True
+        cfg_obj.swanlab_project = "test-project"
+        cfg_obj.swanlab_mode = "local"
+        cfg_obj.swanlab_lark_webhook_url = None  # No Lark
+
+        with (
+            patch("swanlab.init") as mock_init,
+            patch("swanlab.__version__", "0.3.0"),
+            patch("axolotl.utils.distributed.is_main_process", return_value=True),
+            patch("axolotl.utils.distributed.get_world_size", return_value=1),
+        ):
+            plugin.pre_model_load(cfg_obj)
+            # Should call swanlab.init
+            mock_init.assert_called_once()
+
+    def test_lifecycle_with_multi_logger_warning(self, caplog):
+        """Test lifecycle with multi-logger warning."""
+        plugin = SwanLabPlugin()
+
+        cfg_dict = {
+            "use_swanlab": True,
+            "swanlab_project": "test-project",
+            "use_wandb": True,
+        }
+
+        with caplog.at_level(logging.WARNING):
+            plugin.register(cfg_dict)
+
+        # Should have multi-logger warning
+        warning_messages = [record.message for record in caplog.records]
+        assert any("Multiple logging tools" in msg for msg in warning_messages)
+
+    def test_lifecycle_invalid_config_fails_early(self):
+        """Test that invalid config fails at register stage."""
+        plugin = SwanLabPlugin()
+
+        cfg_dict = {
+            "use_swanlab": True,
+            # Missing swanlab_project
+        }
+
+        # Should fail at register, not pre_model_load
+        with pytest.raises(ValueError):
+            plugin.register(cfg_dict)
+
+    def test_full_lifecycle_with_lark_notifications(self):
+        """Test full lifecycle including Lark notification registration."""
+        plugin = SwanLabPlugin()
+
+        # Register
+        cfg_dict = {
+            "use_swanlab": True,
+            "swanlab_project": "test-project",
+            "swanlab_mode": "cloud",
+        }
+        plugin.register(cfg_dict)
+
+        # Pre-model load with Lark config
+        cfg_obj = MagicMock()
+        cfg_obj.use_swanlab = True
+        cfg_obj.swanlab_project = "test-project"
+        cfg_obj.swanlab_mode = "cloud"
+        cfg_obj.swanlab_lark_webhook_url = (
+            "https://open.feishu.cn/open-apis/bot/v2/hook/test"
+        )
+        cfg_obj.swanlab_lark_secret = "secret123"
+
+        with (
+            patch("swanlab.init"),
+            patch("swanlab.__version__", "0.3.0"),
+            patch("swanlab.register_callbacks") as mock_register,
+            patch("axolotl.utils.distributed.is_main_process", return_value=True),
+            patch("axolotl.utils.distributed.get_world_size", return_value=1),
+        ):
+            with patch("swanlab.plugin.notification.LarkCallback") as MockLarkCallback:
+                mock_lark_instance = MagicMock()
+                MockLarkCallback.return_value = mock_lark_instance
+
+                plugin.pre_model_load(cfg_obj)
+
+                # Verify both SwanLab init AND Lark callback registration
+                MockLarkCallback.assert_called_once()
+                mock_register.assert_called_once_with([mock_lark_instance])
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestCompletionLogger:
+    """Tests for CompletionLogger utility class."""
+
+    def test_completion_logger_initialization(self):
+        """Test CompletionLogger initializes with correct maxlen."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=64)
+        assert logger.maxlen == 64
+        assert len(logger) == 0
+
+    def test_add_dpo_completion(self):
+        """Test adding DPO completions to buffer."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=10)
+
+        logger.add_dpo_completion(
+            step=0,
+            prompt="What is AI?",
+            chosen="Artificial Intelligence is...",
+            rejected="AI means...",
+            reward_diff=0.5,
+        )
+
+        assert len(logger) == 1
+        entry = logger.data[0]
+        assert entry["step"] == 0
+        assert entry["prompt"] == "What is AI?"
+        assert entry["chosen"] == "Artificial Intelligence is..."
+        assert entry["rejected"] == "AI means..."
+        assert entry["reward_diff"] == 0.5
+
+    def test_add_kto_completion(self):
+        """Test adding KTO completions to buffer."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=10)
+
+        logger.add_kto_completion(
+            step=1,
+            prompt="Explain quantum physics",
+            completion="Quantum physics is...",
+            label=True,
+            reward=0.8,
+        )
+
+        assert len(logger) == 1
+        entry = logger.data[0]
+        assert entry["step"] == 1
+        assert entry["prompt"] == "Explain quantum physics"
+        assert entry["completion"] == "Quantum physics is..."
+        assert entry["label"] == "desirable"
+        assert entry["reward"] == 0.8
+
+    def test_add_orpo_completion(self):
+        """Test adding ORPO completions to buffer."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=10)
+
+        logger.add_orpo_completion(
+            step=2,
+            prompt="Write a poem",
+            chosen="Roses are red...",
+            rejected="Violets are blue...",
+            log_odds_ratio=1.2,
+        )
+
+        assert len(logger) == 1
+        entry = logger.data[0]
+        assert entry["step"] == 2
+        assert entry["chosen"] == "Roses are red..."
+        assert entry["rejected"] == "Violets are blue..."
+        assert entry["log_odds_ratio"] == 1.2
+
+    def test_add_grpo_completion(self):
+        """Test adding GRPO completions to buffer."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=10)
+
+        logger.add_grpo_completion(
+            step=3,
+            prompt="Solve this problem",
+            completion="The answer is 42",
+            reward=0.9,
+            advantage=0.3,
+        )
+
+        assert len(logger) == 1
+        entry = logger.data[0]
+        assert entry["step"] == 3
+        assert entry["completion"] == "The answer is 42"
+        assert entry["reward"] == 0.9
+        assert entry["advantage"] == 0.3
+
+    def test_memory_bounded_buffer(self):
+        """Test that buffer respects maxlen and drops oldest entries."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=3)
+
+        # Add 5 completions
+        for i in range(5):
+            logger.add_dpo_completion(
+                step=i,
+                prompt=f"Prompt {i}",
+                chosen=f"Chosen {i}",
+                rejected=f"Rejected {i}",
+            )
+
+        # Should only keep last 3
+        assert len(logger) == 3
+        assert logger.data[0]["step"] == 2  # Oldest kept
+        assert logger.data[1]["step"] == 3
+        assert logger.data[2]["step"] == 4  # Newest
+
+    def test_log_to_swanlab_when_not_initialized(self):
+        """Test logging gracefully fails when SwanLab not initialized."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=10)
+        logger.add_dpo_completion(
+            step=0,
+            prompt="Test",
+            chosen="A",
+            rejected="B",
+        )
+
+        with patch("swanlab.get_run", return_value=None):
+            result = logger.log_to_swanlab()
+            assert result is False  # Should fail gracefully
+
+    def test_log_to_swanlab_success(self):
+        """Test successful logging to SwanLab."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=10)
+        logger.add_dpo_completion(
+            step=0,
+            prompt="Test prompt",
+            chosen="Chosen response",
+            rejected="Rejected response",
+            reward_diff=0.5,
+        )
+
+        with (
+            patch("swanlab.get_run") as mock_get_run,
+            patch("swanlab.log") as mock_log,
+            patch("swanlab.echarts.Table") as MockTable,
+        ):
+            mock_get_run.return_value = MagicMock()  # SwanLab initialized
+            mock_table_instance = MagicMock()
+            MockTable.return_value = mock_table_instance
+
+            result = logger.log_to_swanlab(table_name="test_table")
+
+            assert result is True
+            mock_log.assert_called_once()
+            mock_table_instance.add.assert_called_once()
+
+    def test_clear_buffer(self):
+        """Test clearing the completion buffer."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=10)
+        logger.add_dpo_completion(
+            step=0,
+            prompt="Test",
+            chosen="A",
+            rejected="B",
+        )
+
+        assert len(logger) == 1
+        logger.clear()
+        assert len(logger) == 0
+
+    def test_repr(self):
+        """Test string representation."""
+        from axolotl.integrations.swanlab.completion_logger import CompletionLogger
+
+        logger = CompletionLogger(maxlen=128)
+        logger.add_dpo_completion(
+            step=0,
+            prompt="Test",
+            chosen="A",
+            rejected="B",
+        )
+
+        repr_str = repr(logger)
+        assert "CompletionLogger" in repr_str
+        assert "maxlen=128" in repr_str
+        assert "buffered=1/128" in repr_str
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestSwanLabRLHFCompletionCallback:
+    """Tests for SwanLabRLHFCompletionCallback."""
+
+    def test_callback_initialization(self):
+        """Test callback initializes with correct parameters."""
+        from axolotl.integrations.swanlab.callbacks import SwanLabRLHFCompletionCallback
+
+        callback = SwanLabRLHFCompletionCallback(
+            log_interval=50,
+            max_completions=64,
+            table_name="custom_table",
+        )
+
+        assert callback.log_interval == 50
+        assert callback.logger.maxlen == 64
+        assert callback.table_name == "custom_table"
+        assert callback.trainer_type is None
+
+    def test_trainer_type_detection_dpo(self):
+        """Test DPO trainer type is detected correctly."""
+        from axolotl.integrations.swanlab.callbacks import SwanLabRLHFCompletionCallback
+
+        callback = SwanLabRLHFCompletionCallback()
+
+        # Mock trainer with DPO in name
+        mock_trainer = MagicMock()
+        mock_trainer.__class__.__name__ = "AxolotlDPOTrainer"
+
+        callback.on_init_end(
+            args=MagicMock(),
+            state=MagicMock(),
+            control=MagicMock(),
+            trainer=mock_trainer,
+        )
+
+        assert callback.trainer_type == "dpo"
+
+    def test_trainer_type_detection_kto(self):
+        """Test KTO trainer type is detected correctly."""
+        from axolotl.integrations.swanlab.callbacks import SwanLabRLHFCompletionCallback
+
+        callback = SwanLabRLHFCompletionCallback()
+
+        mock_trainer = MagicMock()
+        mock_trainer.__class__.__name__ = "AxolotlKTOTrainer"
+
+        callback.on_init_end(
+            args=MagicMock(),
+            state=MagicMock(),
+            control=MagicMock(),
+            trainer=mock_trainer,
+        )
+
+        assert callback.trainer_type == "kto"
+
+    def test_on_train_end_logs_completions(self):
+        """Test that completions are logged at end of training."""
+        from axolotl.integrations.swanlab.callbacks import SwanLabRLHFCompletionCallback
+
+        callback = SwanLabRLHFCompletionCallback()
+        callback.trainer_type = "dpo"
+
+        # Add some completions to buffer
+        callback.logger.add_dpo_completion(
+            step=0,
+            prompt="Test",
+            chosen="A",
+            rejected="B",
+        )
+
+        with patch.object(callback.logger, "log_to_swanlab") as mock_log:
+            callback.on_train_end(
+                args=MagicMock(),
+                state=MagicMock(global_step=100),
+                control=MagicMock(),
+            )
+
+            # Should log remaining completions
+            mock_log.assert_called_once()
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestSwanLabPluginCompletionIntegration:
+    """Integration tests for completion logging in SwanLabPlugin."""
+
+    def test_completion_callback_registered_for_dpo_trainer(self):
+        """Test that completion callback is registered for DPO trainer."""
+        from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+        from axolotl.utils.dict import DictDefault
+
+        plugin = SwanLabPlugin()
+        plugin.swanlab_initialized = True  # Simulate SwanLab initialized
+
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "test-project",
+            "swanlab_log_completions": True,
+            "swanlab_completion_log_interval": 50,
+            "swanlab_completion_max_buffer": 64,
+        }
+        cfg_obj = DictDefault(cfg)
+
+        # Mock DPO trainer
+        mock_trainer = MagicMock()
+        mock_trainer.__class__.__name__ = "AxolotlDPOTrainer"
+        mock_trainer.state = MagicMock(max_steps=1000)
+        mock_trainer.args = MagicMock(
+            num_train_epochs=3,
+            train_batch_size=4,
+            gradient_accumulation_steps=2,
+        )
+
+        with patch("swanlab.config.update"):
+            plugin.post_trainer_create(cfg_obj, mock_trainer)
+
+        # Verify callback was added
+        mock_trainer.add_callback.assert_called_once()
+        callback = mock_trainer.add_callback.call_args[0][0]
+        assert callback.__class__.__name__ == "SwanLabRLHFCompletionCallback"
+        assert callback.log_interval == 50
+        assert callback.logger.maxlen == 64
+
+    def test_completion_callback_not_registered_for_non_rlhf_trainer(self):
+        """Test that completion callback is NOT registered for non-RLHF trainers."""
+        from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+        from axolotl.utils.dict import DictDefault
+
+        plugin = SwanLabPlugin()
+        plugin.swanlab_initialized = True
+
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "test-project",
+            "swanlab_log_completions": True,
+        }
+        cfg_obj = DictDefault(cfg)
+
+        # Mock regular SFT trainer (not RLHF)
+        mock_trainer = MagicMock()
+        mock_trainer.__class__.__name__ = "AxolotlTrainer"  # Not RLHF
+        mock_trainer.state = MagicMock(max_steps=1000)
+        mock_trainer.args = MagicMock()
+
+        with patch("swanlab.config.update"):
+            plugin.post_trainer_create(cfg_obj, mock_trainer)
+
+        # Callback should NOT be added for non-RLHF trainer
+        mock_trainer.add_callback.assert_not_called()
+
+    def test_completion_callback_not_registered_when_disabled(self):
+        """Test that completion callback is not registered when disabled in config."""
+        from axolotl.integrations.swanlab.plugins import SwanLabPlugin
+        from axolotl.utils.dict import DictDefault
+
+        plugin = SwanLabPlugin()
+        plugin.swanlab_initialized = True
+
+        cfg = {
+            "use_swanlab": True,
+            "swanlab_project": "test-project",
+            "swanlab_log_completions": False,  # Disabled
+        }
+        cfg_obj = DictDefault(cfg)
+
+        # Mock DPO trainer
+        mock_trainer = MagicMock()
+        mock_trainer.__class__.__name__ = "AxolotlDPOTrainer"
+        mock_trainer.state = MagicMock(max_steps=1000)
+        mock_trainer.args = MagicMock()
+
+        with patch("swanlab.config.update"):
+            plugin.post_trainer_create(cfg_obj, mock_trainer)
+
+        # Callback should NOT be added when disabled
+        mock_trainer.add_callback.assert_not_called()
+
+
+@pytest.mark.skipif(not SWANLAB_INSTALLED, reason="swanlab package not installed")
+class TestSwanLabProfiling:
+    """Tests for SwanLab profiling utilities."""
+
+    def test_profiling_context_logs_duration(self):
+        """Test that profiling context logs execution duration."""
+        from axolotl.integrations.swanlab.profiling import swanlab_profiling_context
+
+        # Mock trainer with SwanLab enabled
+        mock_trainer = MagicMock()
+        mock_trainer.axolotl_cfg = MagicMock(use_swanlab=True)
+        mock_trainer.__class__.__name__ = "TestTrainer"
+
+        with patch("swanlab.get_run") as mock_get_run, patch("swanlab.log") as mock_log:
+            mock_get_run.return_value = MagicMock()  # SwanLab initialized
+
+            with swanlab_profiling_context(mock_trainer, "test_function"):
+                time.sleep(0.01)  # Simulate work
+
+            # Verify log was called with correct metric name
+            mock_log.assert_called_once()
+            logged_data = mock_log.call_args[0][0]
+            assert "profiling/Time taken: TestTrainer.test_function" in logged_data
+            # Duration should be > 0.01 seconds
+            assert (
+                logged_data["profiling/Time taken: TestTrainer.test_function"] >= 0.01
+            )
+
+    def test_profiling_context_skips_when_swanlab_disabled(self):
+        """Test that profiling is skipped when SwanLab is disabled."""
+        from axolotl.integrations.swanlab.profiling import swanlab_profiling_context
+
+        mock_trainer = MagicMock()
+        mock_trainer.axolotl_cfg = MagicMock(use_swanlab=False)  # Disabled
+
+        with patch("swanlab.log") as mock_log:
+            with swanlab_profiling_context(mock_trainer, "test_function"):
+                time.sleep(0.01)
+
+            # Should NOT log when disabled
+            mock_log.assert_not_called()
+
+    def test_profiling_context_skips_when_swanlab_not_initialized(self):
+        """Test that profiling is skipped when SwanLab not initialized."""
+        from axolotl.integrations.swanlab.profiling import swanlab_profiling_context
+
+        mock_trainer = MagicMock()
+        mock_trainer.axolotl_cfg = MagicMock(use_swanlab=True)
+
+        with (
+            patch("swanlab.get_run", return_value=None),
+            patch("swanlab.log") as mock_log,
+        ):
+            with swanlab_profiling_context(mock_trainer, "test_function"):
+                time.sleep(0.01)
+
+            # Should NOT log when not initialized
+            mock_log.assert_not_called()
+
+    def test_profiling_decorator(self):
+        """Test swanlab_profile decorator."""
+        from axolotl.integrations.swanlab.profiling import swanlab_profile
+
+        class MockTrainer:
+            def __init__(self):
+                self.cfg = MagicMock(use_swanlab=True)
+
+            @swanlab_profile
+            def expensive_method(self, x):
+                time.sleep(0.01)
+                return x * 2
+
+        trainer = MockTrainer()
+
+        with patch("swanlab.get_run") as mock_get_run, patch("swanlab.log") as mock_log:
+            mock_get_run.return_value = MagicMock()
+
+            result = trainer.expensive_method(5)
+
+            # Verify method still works correctly
+            assert result == 10
+
+            # Verify profiling was logged
+            mock_log.assert_called_once()
+            logged_data = mock_log.call_args[0][0]
+            assert "profiling/Time taken: MockTrainer.expensive_method" in logged_data
+
+    def test_profiling_config(self):
+        """Test ProfilingConfig class."""
+        from axolotl.integrations.swanlab.profiling import ProfilingConfig
+
+        config = ProfilingConfig(
+            enabled=True,
+            min_duration_ms=1.0,
+            log_interval=5,
+        )
+
+        # Test enabled check
+        assert config.enabled is True
+
+        # Test minimum duration filtering
+        assert config.should_log("func1", 0.0001) is False  # 0.1ms < 1.0ms threshold
+        assert config.should_log("func2", 0.002) is True  # 2.0ms > 1.0ms threshold
+
+        # Test log interval
+        assert config.should_log("func3", 0.002) is True  # 1st call
+        assert config.should_log("func3", 0.002) is False  # 2nd call
+        assert config.should_log("func3", 0.002) is False  # 3rd call
+        assert config.should_log("func3", 0.002) is False  # 4th call
+        assert config.should_log("func3", 0.002) is True  # 5th call (interval=5)
+
+    def test_profiling_config_when_disabled(self):
+        """Test ProfilingConfig when disabled."""
+        from axolotl.integrations.swanlab.profiling import ProfilingConfig
+
+        config = ProfilingConfig(enabled=False)
+
+        # Should never log when disabled
+        assert config.should_log("func1", 100.0) is False
+
+    def test_profiling_context_advanced(self):
+        """Test advanced profiling context with custom config."""
+        from axolotl.integrations.swanlab.profiling import (
+            ProfilingConfig,
+            swanlab_profiling_context_advanced,
+        )
+
+        mock_trainer = MagicMock()
+        mock_trainer.axolotl_cfg = MagicMock(use_swanlab=True)
+        mock_trainer.__class__.__name__ = "TestTrainer"
+
+        # Config that filters out very fast operations
+        config = ProfilingConfig(min_duration_ms=10.0)  # 10ms minimum
+
+        with patch("swanlab.get_run") as mock_get_run, patch("swanlab.log") as mock_log:
+            mock_get_run.return_value = MagicMock()
+
+            # Fast operation (< 10ms) - should NOT log
+            with swanlab_profiling_context_advanced(mock_trainer, "fast_op", config):
+                time.sleep(0.001)  # 1ms
+
+            mock_log.assert_not_called()
+
+            # Slow operation (> 10ms) - should log
+            with swanlab_profiling_context_advanced(mock_trainer, "slow_op", config):
+                time.sleep(0.015)  # 15ms
+
+            mock_log.assert_called_once()
+
+    def test_profiling_with_exception(self):
+        """Test that profiling still logs even when exception occurs."""
+        from axolotl.integrations.swanlab.profiling import swanlab_profiling_context
+
+        mock_trainer = MagicMock()
+        mock_trainer.axolotl_cfg = MagicMock(use_swanlab=True)
+        mock_trainer.__class__.__name__ = "TestTrainer"
+
+        with patch("swanlab.get_run") as mock_get_run, patch("swanlab.log") as mock_log:
+            mock_get_run.return_value = MagicMock()
+
+            try:
+                with swanlab_profiling_context(mock_trainer, "error_function"):
+                    time.sleep(0.01)
+                    raise ValueError("Test error")
+            except ValueError:
+                pass  # Expected
+
+            # Should still log duration even with exception
+            mock_log.assert_called_once()
diff --git a/tests/kernels/test_fused_rope_autotune_telemetry.py b/tests/kernels/test_fused_rope_autotune_telemetry.py
new file mode 100644
index 0000000000..bc5d61c0fc
--- /dev/null
+++ b/tests/kernels/test_fused_rope_autotune_telemetry.py
@@ -0,0 +1,168 @@
+"""Tests for fused RMSNorm+RoPE autotune telemetry.
+
+Mocked end-to-end, so no Triton or CUDA is required (mirrors
+``tests/integrations/test_scattermoe_autotune_telemetry.py``).
+"""
+
+import sys
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+_MODPATH = "axolotl.kernels.gemma4_fused_rope"
+
+
+def _make_mock_config(kwargs, num_warps=2, num_stages=1):
+    return SimpleNamespace(
+        kwargs=kwargs, num_warps=num_warps, num_stages=num_stages, num_ctas=None
+    )
+
+
+def _make_fake_module(cache=None):
+    kernel = SimpleNamespace(cache=cache if cache is not None else {})
+    return SimpleNamespace(_rms_norm_rope_backward_kernel=kernel)
+
+
+class TestCollector:
+    def test_no_module_returns_empty(self):
+        from axolotl.kernels.autotune_telemetry import (
+            collect_fused_rope_autotune_configs,
+        )
+
+        with patch.dict(sys.modules, {_MODPATH: None}):
+            assert collect_fused_rope_autotune_configs() == []
+
+    def test_empty_cache_returns_empty(self):
+        from axolotl.kernels.autotune_telemetry import (
+            collect_fused_rope_autotune_configs,
+        )
+
+        with patch.dict(sys.modules, {_MODPATH: _make_fake_module()}):
+            assert collect_fused_rope_autotune_configs() == []
+
+    def test_populated_cache_returns_configs(self):
+        from axolotl.kernels.autotune_telemetry import (
+            collect_fused_rope_autotune_configs,
+        )
+
+        cfg = _make_mock_config({}, num_warps=2, num_stages=1)
+        fake = _make_fake_module(cache={(128, "torch.bfloat16"): cfg})
+        with patch.dict(sys.modules, {_MODPATH: fake}):
+            result = collect_fused_rope_autotune_configs()
+
+        assert len(result) == 1
+        entry = result[0]
+        assert entry["kernel"] == "fused_rms_norm_rope_bwd"
+        assert entry["key"]["n_cols"] == 128
+        assert entry["key"]["_extra"] == ["torch.bfloat16"]
+        assert entry["config"]["num_warps"] == 2
+        assert entry["config"]["num_stages"] == 1
+
+    def test_multiple_head_dims(self):
+        """head_dim 128 (Qwen3) and 256 (Qwen3.5) get separate cache entries."""
+        from axolotl.kernels.autotune_telemetry import (
+            collect_fused_rope_autotune_configs,
+        )
+
+        fake = _make_fake_module(
+            cache={
+                (128,): _make_mock_config({}, num_warps=2),
+                (256,): _make_mock_config({}, num_warps=4),
+            }
+        )
+        with patch.dict(sys.modules, {_MODPATH: fake}):
+            result = collect_fused_rope_autotune_configs()
+
+        assert {e["key"]["n_cols"] for e in result} == {128, 256}
+
+
+class TestCallback:
+    def _patch_collect(self, return_value=None, side_effect=None):
+        return patch(
+            "axolotl.kernels.autotune_telemetry.collect_fused_rope_autotune_configs",
+            return_value=return_value,
+            side_effect=side_effect,
+        )
+
+    def test_reports_once_on_first_step(self):
+        from axolotl.kernels.autotune_telemetry import FusedRopeAutotuneReportCallback
+
+        cb = FusedRopeAutotuneReportCallback()
+        state = MagicMock()
+        state.global_step = 1
+        configs = [{"kernel": "fused_rms_norm_rope_bwd", "key": {}, "config": {}}]
+
+        with (
+            self._patch_collect(return_value=configs),
+            patch("axolotl.telemetry.manager.TelemetryManager") as tm_cls,
+        ):
+            tm = MagicMock()
+            tm.enabled = True
+            tm_cls.get_instance.return_value = tm
+
+            cb.on_step_end(args=MagicMock(), state=state, control=MagicMock())
+            assert tm.send_event.call_count == 1
+            kw = tm.send_event.call_args[1]
+            assert kw["event_type"] == "fused-rope-autotune"
+            assert kw["properties"]["kernel_count"] == 1
+
+            cb.on_step_end(args=MagicMock(), state=state, control=MagicMock())
+            assert tm.send_event.call_count == 1
+
+    def test_retries_until_step_5_then_gives_up(self):
+        from axolotl.kernels.autotune_telemetry import FusedRopeAutotuneReportCallback
+
+        cb = FusedRopeAutotuneReportCallback()
+        with self._patch_collect(return_value=[]):
+            for step in range(1, 7):
+                state = MagicMock()
+                state.global_step = step
+                cb.on_step_end(args=MagicMock(), state=state, control=MagicMock())
+        assert cb._reported is True
+
+    def test_includes_gpu_info(self):
+        from axolotl.kernels.autotune_telemetry import FusedRopeAutotuneReportCallback
+
+        cb = FusedRopeAutotuneReportCallback()
+        state = MagicMock()
+        state.global_step = 1
+        configs = [{"kernel": "fused_rms_norm_rope_bwd", "key": {}, "config": {}}]
+        gpu = {
+            "gpu_name": "NVIDIA H100",
+            "gpu_compute_capability": "9.0",
+            "gpu_memory_bytes": 85899345920,
+        }
+
+        with (
+            self._patch_collect(return_value=configs),
+            patch("axolotl.kernels.autotune_telemetry._get_gpu_info", return_value=gpu),
+            patch("axolotl.telemetry.manager.TelemetryManager") as tm_cls,
+        ):
+            tm = MagicMock()
+            tm.enabled = True
+            tm_cls.get_instance.return_value = tm
+
+            cb.on_step_end(args=MagicMock(), state=state, control=MagicMock())
+            props = tm.send_event.call_args[1]["properties"]
+            assert props["gpu_name"] == "NVIDIA H100"
+            assert props["gpu_compute_capability"] == "9.0"
+
+    def test_skips_send_when_telemetry_disabled(self):
+        from axolotl.kernels.autotune_telemetry import FusedRopeAutotuneReportCallback
+
+        cb = FusedRopeAutotuneReportCallback()
+        state = MagicMock()
+        state.global_step = 1
+
+        with (
+            self._patch_collect(
+                return_value=[{"kernel": "x", "key": {}, "config": {}}]
+            ),
+            patch("axolotl.telemetry.manager.TelemetryManager") as tm_cls,
+        ):
+            tm = MagicMock()
+            tm.enabled = False
+            tm_cls.get_instance.return_value = tm
+
+            cb.on_step_end(args=MagicMock(), state=state, control=MagicMock())
+            assert tm.send_event.call_count == 0
+            assert cb._reported is True
diff --git a/tests/kernels/test_gemma4_fused_rope.py b/tests/kernels/test_gemma4_fused_rope.py
new file mode 100644
index 0000000000..297bb25279
--- /dev/null
+++ b/tests/kernels/test_gemma4_fused_rope.py
@@ -0,0 +1,416 @@
+"""
+Correctness tests for the fused RMSNorm+RoPE Triton kernel.
+
+Tests forward and backward against the reference Gemma4 implementation
+(Gemma4RMSNorm + apply_rotary_pos_emb) across both sliding window
+(head_dim=256) and global attention (head_dim=512) layer configurations.
+"""
+
+import pytest
+import torch
+
+torch.manual_seed(42)
+
+# Skip entire module if no CUDA
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _reference_norm_rope(x, weight, cos, sin, eps):
+    """Reference: separate Gemma4RMSNorm + apply_rotary_pos_emb."""
+    from transformers.models.gemma4.modeling_gemma4 import (
+        Gemma4RMSNorm,
+        apply_rotary_pos_emb,
+    )
+
+    D = x.shape[-1]
+    norm = Gemma4RMSNorm(D, eps=eps).to(x.device, x.dtype)
+    norm.weight.data.copy_(weight)
+    normed = norm(x)
+    return apply_rotary_pos_emb(normed, cos, sin, unsqueeze_dim=2)
+
+
+def _reference_norm_noscale(x, eps):
+    """Reference: Gemma4RMSNorm with_scale=False."""
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4RMSNorm
+
+    D = x.shape[-1]
+    norm = Gemma4RMSNorm(D, eps=eps, with_scale=False).to(x.device, x.dtype)
+    return norm(x)
+
+
+def _reference_partial_norm_rope(x, weight, cos, sin, eps):
+    """Reference: Gemma4RMSNorm over the full head_dim, then stock
+    ``apply_rotary_pos_emb`` over the first ``cos.shape[-1]`` columns, with
+    the trailing columns passed through unchanged. Mirrors how Llama-style
+    partial rotary is layered on top of the stock RMSNorm + RoPE primitives.
+    """
+    from transformers.models.gemma4.modeling_gemma4 import (
+        Gemma4RMSNorm,
+        apply_rotary_pos_emb,
+    )
+
+    D = x.shape[-1]
+    n_rot = cos.shape[-1]
+    norm = Gemma4RMSNorm(D, eps=eps).to(x.device, x.dtype)
+    norm.weight.data.copy_(weight)
+    normed = norm(x)
+    if n_rot == D:
+        return apply_rotary_pos_emb(normed, cos, sin, unsqueeze_dim=2)
+    x_rot = normed[..., :n_rot]
+    x_pass = normed[..., n_rot:]
+    rotated = apply_rotary_pos_emb(x_rot, cos, sin, unsqueeze_dim=2)
+    return torch.cat([rotated, x_pass], dim=-1)
+
+
+@pytest.fixture(
+    params=[
+        (2, 64, 32, 256),  # sliding window layer shape
+        (2, 64, 4, 512),  # global attention layer shape
+        (1, 128, 16, 256),  # different batch/seq
+        (1, 1, 1, 8),  # minimal size
+    ],
+    ids=["sliding_256", "global_512", "varied", "minimal"],
+)
+def shapes(request):
+    return request.param
+
+
+@pytest.fixture(params=[torch.bfloat16, torch.float16], ids=["bf16", "fp16"])
+def dtype(request):
+    return request.param
+
+
+class TestFusedRMSNormRoPEForward:
+    """Forward pass correctness."""
+
+    def test_matches_reference(self, shapes, dtype):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = shapes
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=dtype)
+        weight = torch.randn(D, device="cuda", dtype=dtype)
+        cos = torch.randn(B, S, D, device="cuda", dtype=dtype)
+        sin = torch.randn(B, S, D, device="cuda", dtype=dtype)
+
+        y_ref = _reference_norm_rope(x.clone(), weight, cos, sin, eps)
+        y_fused = fused_rms_norm_rope(x.clone(), weight, cos, sin, eps=eps)
+
+        cos_sim = torch.nn.functional.cosine_similarity(
+            y_ref.flatten().float(), y_fused.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"Forward cosine_sim={cos_sim:.6f}, expected > 0.999"
+
+    def test_output_shape(self, shapes):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = shapes
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+
+        y = fused_rms_norm_rope(x, weight, cos, sin, eps=1e-6)
+        assert y.shape == x.shape
+        assert y.dtype == x.dtype
+
+
+class TestFusedRMSNormRoPEBackward:
+    """Backward pass correctness via gradient comparison."""
+
+    @pytest.mark.parametrize(
+        "B,S,H,D",
+        [(2, 64, 32, 256), (2, 64, 4, 512)],
+        ids=["sliding_256", "global_512"],
+    )
+    def test_x_grad_matches_reference(self, B, S, H, D):
+        from transformers.models.gemma4.modeling_gemma4 import (
+            Gemma4RMSNorm,
+            apply_rotary_pos_emb,
+        )
+
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        eps = 1e-6
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        weight_init = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+
+        # Reference backward
+        x_ref = torch.randn(
+            B, S, H, D, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        norm_ref = Gemma4RMSNorm(D, eps=eps).cuda().to(torch.bfloat16)
+        norm_ref.weight.data.copy_(weight_init)
+        y_ref = apply_rotary_pos_emb(norm_ref(x_ref), cos, sin, unsqueeze_dim=2)
+        y_ref.sum().backward()
+
+        # Fused backward
+        x_fused = x_ref.data.clone().requires_grad_(True)
+        w_fused = weight_init.clone().requires_grad_(True)
+        y_fused = fused_rms_norm_rope(x_fused, w_fused, cos, sin, eps=eps)
+        y_fused.sum().backward()
+
+        cos_sim_x = torch.nn.functional.cosine_similarity(
+            x_fused.grad.flatten().float(), x_ref.grad.flatten().float(), dim=0
+        )
+        assert cos_sim_x > 0.999, f"x grad cosine_sim={cos_sim_x:.6f}, expected > 0.999"
+
+    @pytest.mark.parametrize(
+        "B,S,H,D",
+        [(2, 64, 32, 256), (2, 64, 4, 512)],
+        ids=["sliding_256", "global_512"],
+    )
+    def test_weight_grad_matches_reference(self, B, S, H, D):
+        from transformers.models.gemma4.modeling_gemma4 import (
+            Gemma4RMSNorm,
+            apply_rotary_pos_emb,
+        )
+
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        eps = 1e-6
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        weight_init = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+
+        # Reference
+        x_ref = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        norm_ref = Gemma4RMSNorm(D, eps=eps).cuda().to(torch.bfloat16)
+        norm_ref.weight = torch.nn.Parameter(weight_init.clone())
+        apply_rotary_pos_emb(
+            norm_ref(x_ref), cos, sin, unsqueeze_dim=2
+        ).sum().backward()
+
+        # Fused
+        w_fused = weight_init.clone().requires_grad_(True)
+        fused_rms_norm_rope(x_ref.clone(), w_fused, cos, sin, eps=eps).sum().backward()
+
+        cos_sim_w = torch.nn.functional.cosine_similarity(
+            w_fused.grad.flatten().float(),
+            norm_ref.weight.grad.flatten().float(),
+            dim=0,
+        )
+        assert cos_sim_w > 0.995, (
+            f"weight grad cosine_sim={cos_sim_w:.6f}, expected > 0.995"
+        )
+
+    def test_grad_flows(self):
+        """Verify gradients are non-zero and finite."""
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 1, 16, 4, 64
+        x = torch.randn(
+            B, S, H, D, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        w = torch.randn(D, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+
+        y = fused_rms_norm_rope(x, w, cos, sin, eps=1e-6)
+        y.sum().backward()
+
+        assert x.grad is not None, "x.grad is None"
+        assert w.grad is not None, "w.grad is None"
+        assert x.grad.isfinite().all(), "x.grad has non-finite values"
+        assert w.grad.isfinite().all(), "w.grad has non-finite values"
+        assert x.grad.abs().sum() > 0, "x.grad is all zeros"
+        assert w.grad.abs().sum() > 0, "w.grad is all zeros"
+
+
+class TestFusedRMSNormRoPEPartialRotary:
+    """Partial-rotary: cos/sin last dim is smaller than head_dim.
+
+    Compares against the original primitives (`Gemma4RMSNorm` +
+    `apply_rotary_pos_emb`) applied to the rotated slice with the trailing
+    columns passed through. Without the kernel fix this used to crash with
+    `RuntimeError: shape '[..., D]' is invalid for input of size B*S*n_rot`.
+    """
+
+    @pytest.mark.parametrize(
+        "B,S,H,D,n_rot",
+        [
+            (2, 16, 4, 64, 32),  # half rotary (Llama-style 0.5)
+            (2, 16, 4, 64, 16),  # quarter rotary
+            (2, 32, 8, 128, 64),  # half rotary, larger heads
+            (1, 8, 2, 256, 64),  # 26B sliding-shape, 0.25 partial
+            (1, 8, 2, 64, 64),  # n_rot == D: must still match full-rotary path
+        ],
+        ids=["half_64", "quarter_64", "half_128", "quarter_256", "full_64"],
+    )
+    def test_forward_matches_reference(self, B, S, H, D, n_rot):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+
+        y_ref = _reference_partial_norm_rope(x.clone(), weight, cos, sin, eps)
+        y_fused = fused_rms_norm_rope(x.clone(), weight, cos, sin, eps=eps)
+
+        assert y_fused.shape == y_ref.shape == (B, S, H, D)
+        cos_sim = torch.nn.functional.cosine_similarity(
+            y_ref.flatten().float(), y_fused.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, (
+            f"partial rotary forward cosine_sim={cos_sim:.6f} "
+            f"(B={B},S={S},H={H},D={D},n_rot={n_rot})"
+        )
+
+        # The pass-through tail must equal the reference RMSNorm output bit-
+        # for-bit (any deviation would mean the kernel is touching it with a
+        # spurious rotation, which is the original bug class).
+        torch.testing.assert_close(
+            y_fused[..., n_rot:], y_ref[..., n_rot:], rtol=1e-2, atol=1e-2
+        )
+
+    @pytest.mark.parametrize(
+        "B,S,H,D,n_rot",
+        [(2, 16, 4, 64, 32), (1, 8, 2, 256, 64)],
+        ids=["half_64", "quarter_256"],
+    )
+    def test_x_grad_matches_reference(self, B, S, H, D, n_rot):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        eps = 1e-6
+        cos = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        weight_init = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+        x_data = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+
+        # Reference backward via the original primitives
+        x_ref = x_data.clone().requires_grad_(True)
+        w_ref = weight_init.clone()
+        y_ref = _reference_partial_norm_rope(x_ref, w_ref, cos, sin, eps)
+        y_ref.sum().backward()
+
+        # Fused backward
+        x_fused = x_data.clone().requires_grad_(True)
+        w_fused = weight_init.clone().requires_grad_(True)
+        y_fused = fused_rms_norm_rope(x_fused, w_fused, cos, sin, eps=eps)
+        y_fused.sum().backward()
+
+        cos_sim_x = torch.nn.functional.cosine_similarity(
+            x_fused.grad.flatten().float(), x_ref.grad.flatten().float(), dim=0
+        )
+        assert cos_sim_x > 0.999, f"partial rotary x grad cosine_sim={cos_sim_x:.6f}"
+
+    @pytest.mark.parametrize(
+        "B,S,H,D,n_rot",
+        [(2, 16, 4, 64, 32), (1, 8, 2, 256, 64)],
+        ids=["half_64", "quarter_256"],
+    )
+    def test_weight_grad_matches_reference(self, B, S, H, D, n_rot):
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4RMSNorm
+
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        eps = 1e-6
+        cos = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        weight_init = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+        x_data = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+
+        # Reference: Gemma4RMSNorm whose .weight collects grads, then partial
+        # rotary applied to the rotated slice.
+        norm_ref = Gemma4RMSNorm(D, eps=eps).cuda().to(torch.bfloat16)
+        norm_ref.weight = torch.nn.Parameter(weight_init.clone())
+        normed = norm_ref(x_data)
+        from transformers.models.gemma4.modeling_gemma4 import apply_rotary_pos_emb
+
+        rotated = apply_rotary_pos_emb(normed[..., :n_rot], cos, sin, unsqueeze_dim=2)
+        y_ref = torch.cat([rotated, normed[..., n_rot:]], dim=-1)
+        y_ref.sum().backward()
+
+        w_fused = weight_init.clone().requires_grad_(True)
+        fused_rms_norm_rope(x_data.clone(), w_fused, cos, sin, eps=eps).sum().backward()
+
+        cos_sim_w = torch.nn.functional.cosine_similarity(
+            w_fused.grad.flatten().float(),
+            norm_ref.weight.grad.flatten().float(),
+            dim=0,
+        )
+        assert cos_sim_w > 0.995, (
+            f"partial rotary weight grad cosine_sim={cos_sim_w:.6f}"
+        )
+
+    def test_full_rotary_unchanged_when_n_rot_equals_d(self):
+        """Regression: passing cos/sin with shape == head_dim must still
+        match the full-rotary reference (the partial-rotary code path must
+        not perturb the existing full-rotary output)."""
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 2, 16, 4, 64
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+
+        y_ref = _reference_norm_rope(x.clone(), weight, cos, sin, eps)
+        y_fused = fused_rms_norm_rope(x.clone(), weight, cos, sin, eps=eps)
+        cos_sim = torch.nn.functional.cosine_similarity(
+            y_ref.flatten().float(), y_fused.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"full-rotary regression cos_sim={cos_sim:.6f}"
+
+    def test_validation_errors(self):
+        """Wrapper rejects misshaped inputs cleanly (instead of a cryptic
+        Triton crash deeper in the kernel)."""
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 1, 4, 2, 64
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        w = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+
+        # n_rot > head_dim
+        cos_big = torch.randn(B, S, D + 16, device="cuda", dtype=torch.bfloat16)
+        sin_big = torch.randn(B, S, D + 16, device="cuda", dtype=torch.bfloat16)
+        with pytest.raises(ValueError, match="cannot exceed head_dim"):
+            fused_rms_norm_rope(x, w, cos_big, sin_big)
+
+        # cos/sin last-dim mismatch
+        cos = torch.randn(B, S, 32, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, 16, device="cuda", dtype=torch.bfloat16)
+        with pytest.raises(ValueError, match="same last dim"):
+            fused_rms_norm_rope(x, w, cos, sin)
+
+        # odd rotary dim
+        cos_odd = torch.randn(B, S, 31, device="cuda", dtype=torch.bfloat16)
+        sin_odd = torch.randn(B, S, 31, device="cuda", dtype=torch.bfloat16)
+        with pytest.raises(ValueError, match="must be even"):
+            fused_rms_norm_rope(x, w, cos_odd, sin_odd)
+
+
+class TestFusedRMSNormNoScale:
+    """Tests for v_norm (RMSNorm without learnable scale)."""
+
+    def test_forward_matches_reference(self, shapes, dtype):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_noscale
+
+        B, S, H, D = shapes
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=dtype)
+
+        y_ref = _reference_norm_noscale(x.clone(), eps)
+        y_fused = fused_rms_norm_noscale(x.clone(), eps=eps)
+
+        cos_sim = torch.nn.functional.cosine_similarity(
+            y_ref.flatten().float(), y_fused.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"v_norm cosine_sim={cos_sim:.6f}, expected > 0.999"
+
+    def test_backward_flows(self):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_noscale
+
+        x = torch.randn(
+            1, 16, 4, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        y = fused_rms_norm_noscale(x, eps=1e-6)
+        y.sum().backward()
+
+        assert x.grad is not None
+        assert x.grad.isfinite().all()
+        assert x.grad.abs().sum() > 0
diff --git a/tests/kernels/test_gemma4_fused_rope_compile.py b/tests/kernels/test_gemma4_fused_rope_compile.py
new file mode 100644
index 0000000000..dc8858722a
--- /dev/null
+++ b/tests/kernels/test_gemma4_fused_rope_compile.py
@@ -0,0 +1,66 @@
+"""torch.compile traceability tests for the fused RMSNorm+RoPE kernel."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+def _make_inputs(B=2, S=64, H=4, D=64, n_rot=64, dtype=torch.bfloat16, seed=0):
+    torch.manual_seed(seed)
+    x = torch.randn(B, S, H, D, device="cuda", dtype=dtype, requires_grad=True)
+    w = torch.randn(D, device="cuda", dtype=dtype, requires_grad=True)
+    cos = torch.randn(B, S, n_rot, device="cuda", dtype=dtype)
+    sin = torch.randn(B, S, n_rot, device="cuda", dtype=dtype)
+    return x, w, cos, sin
+
+
+def test_fused_rms_norm_rope_compile_forward_matches_eager():
+    from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+    x, w, cos, sin = _make_inputs(seed=1)
+    x_ref = x.detach().clone().requires_grad_(True)
+    w_ref = w.detach().clone().requires_grad_(True)
+
+    y_eager = fused_rms_norm_rope(x_ref, w_ref, cos, sin, eps=1e-6)
+
+    compiled = torch.compile(fused_rms_norm_rope, fullgraph=True, dynamic=False)
+    y_compiled = compiled(x, w, cos, sin, eps=1e-6)
+
+    torch.testing.assert_close(y_compiled, y_eager, rtol=1e-2, atol=1e-2)
+    assert torch.isfinite(y_compiled).all()
+
+
+def test_fused_rms_norm_rope_compile_backward():
+    from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+    x, w, cos, sin = _make_inputs(seed=2)
+
+    compiled = torch.compile(fused_rms_norm_rope, fullgraph=True, dynamic=False)
+    y = compiled(x, w, cos, sin, eps=1e-6)
+    y.sum().backward()
+
+    assert x.grad is not None and x.grad.isfinite().all() and x.grad.abs().sum() > 0
+    assert w.grad is not None and w.grad.isfinite().all() and w.grad.abs().sum() > 0
+
+
+def test_fused_rms_norm_noscale_compile():
+    from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_noscale
+
+    torch.manual_seed(3)
+    x = torch.randn(
+        2, 32, 4, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True
+    )
+    x_ref = x.detach().clone().requires_grad_(True)
+
+    y_eager = fused_rms_norm_noscale(x_ref, eps=1e-6)
+
+    compiled = torch.compile(fused_rms_norm_noscale, fullgraph=True, dynamic=False)
+    y_compiled = compiled(x, eps=1e-6)
+
+    torch.testing.assert_close(y_compiled, y_eager, rtol=1e-2, atol=1e-2)
+
+    y_compiled.sum().backward()
+    assert x.grad is not None and x.grad.isfinite().all() and x.grad.abs().sum() > 0
diff --git a/tests/kernels/test_gemma4_fused_rope_unit_offset.py b/tests/kernels/test_gemma4_fused_rope_unit_offset.py
new file mode 100644
index 0000000000..1e84d95feb
--- /dev/null
+++ b/tests/kernels/test_gemma4_fused_rope_unit_offset.py
@@ -0,0 +1,238 @@
+"""Correctness tests for the ``unit_offset=True`` (Gemma-style) path in the fused RMSNorm+RoPE kernel."""
+
+import pytest
+import torch
+
+torch.manual_seed(7)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _reference_unit_offset(x, weight, cos, sin, eps):
+    x_fp32 = x.float()
+    rstd = torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + eps)
+    normed = (x_fp32 * rstd * (1.0 + weight.float())).to(x.dtype)
+    cos_b = cos.unsqueeze(2)
+    sin_b = sin.unsqueeze(2)
+    return normed * cos_b + _rotate_half(normed) * sin_b
+
+
+def _reference_unit_offset_partial(x, weight, cos, sin, eps):
+    """Reference for ``unit_offset=True`` with ``cos.shape[-1] < D`` (Qwen3.5 partial rotary)."""
+    n_rot = cos.shape[-1]
+    x_fp32 = x.float()
+    rstd = torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + eps)
+    normed = (x_fp32 * rstd * (1.0 + weight.float())).to(x.dtype)
+    rot_part, pass_part = normed[..., :n_rot], normed[..., n_rot:]
+    cos_b, sin_b = cos.unsqueeze(2), sin.unsqueeze(2)
+    rotated = rot_part * cos_b + _rotate_half(rot_part) * sin_b
+    return torch.cat([rotated, pass_part], dim=-1)
+
+
+def _reference_fp32(x, weight, cos, sin, eps, unit_offset):
+    """fp32 ground truth: no intermediate bf16 rounding, so it's *more* accurate
+    than the eager bf16 path. Handles full (``n_rot == D``) and partial rotary."""
+    n_rot = cos.shape[-1]
+    x_fp32 = x.float()
+    rstd = torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + eps)
+    scale = (1.0 + weight.float()) if unit_offset else weight.float()
+    normed = x_fp32 * rstd * scale
+    cos_b, sin_b = cos.float().unsqueeze(2), sin.float().unsqueeze(2)
+    rot_part, pass_part = normed[..., :n_rot], normed[..., n_rot:]
+    rotated = rot_part * cos_b + _rotate_half(rot_part) * sin_b
+    return torch.cat([rotated, pass_part], dim=-1)
+
+
+def _assert_at_bf16_floor(y_fused, y_ref_fp32, y_eager):
+    """The fused kernel keeps fp32 internally and rounds once, so its bf16 output
+    must land at the bf16 rounding floor of the fp32 reference — and be at least
+    as accurate as the eager bf16 path (which rounds several times mid-compute)."""
+    floor = (y_ref_fp32.to(y_fused.dtype).float() - y_ref_fp32).abs().max()
+    fused_err = (y_fused.float() - y_ref_fp32).abs().max()
+    eager_err = (y_eager.float() - y_ref_fp32).abs().max()
+    assert fused_err <= 1.5 * floor, (
+        f"fused err {fused_err:.3e} exceeds 1.5x bf16 floor {floor:.3e}"
+    )
+    assert fused_err <= eager_err, (
+        f"fused ({fused_err:.3e}) less accurate than eager bf16 path ({eager_err:.3e})"
+    )
+
+
+class TestForward:
+    def test_matches_reference(self):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 2, 32, 4, 64
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16) * 0.1
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+
+        y_fused = fused_rms_norm_rope(
+            x.clone(), weight, cos, sin, eps=eps, unit_offset=True
+        )
+        y_ref_fp32 = _reference_fp32(x, weight, cos, sin, eps, unit_offset=True)
+        y_eager = _reference_unit_offset(x.clone(), weight, cos, sin, eps)
+        _assert_at_bf16_floor(y_fused, y_ref_fp32, y_eager)
+
+    def test_no_offset_matches_fp32_reference(self):
+        """Qwen3 / Gemma 4 path (``unit_offset=False``) also sits at the bf16 floor."""
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 2, 32, 4, 64
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16) * 0.1
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+
+        y_fused = fused_rms_norm_rope(
+            x.clone(), weight, cos, sin, eps=eps, unit_offset=False
+        )
+        y_ref_fp32 = _reference_fp32(x, weight, cos, sin, eps, unit_offset=False)
+        x_fp32 = x.float()
+        rstd = torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + eps)
+        normed = (x_fp32 * rstd * weight.float()).to(x.dtype)
+        y_eager = normed * cos.unsqueeze(2) + _rotate_half(normed) * sin.unsqueeze(2)
+        _assert_at_bf16_floor(y_fused, y_ref_fp32, y_eager)
+
+    def test_differs_from_no_offset(self):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 1, 8, 2, 32
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+
+        y_off = fused_rms_norm_rope(x, weight, cos, sin, unit_offset=False)
+        y_on = fused_rms_norm_rope(x, weight, cos, sin, unit_offset=True)
+        diff = (y_off.float() - y_on.float()).abs().max().item()
+        assert diff > 1e-3, f"unit_offset toggle had no effect: max_abs_diff={diff}"
+
+
+class TestBackward:
+    def test_x_and_w_grad_match_eager(self):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 2, 16, 4, 64
+        eps = 1e-6
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        weight_init = torch.randn(D, device="cuda", dtype=torch.bfloat16) * 0.1
+
+        x_ref = torch.randn(
+            B, S, H, D, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        w_ref = weight_init.clone().requires_grad_(True)
+        y_ref = _reference_unit_offset(x_ref, w_ref, cos, sin, eps)
+        y_ref.sum().backward()
+
+        x_fused = x_ref.data.clone().requires_grad_(True)
+        w_fused = weight_init.clone().requires_grad_(True)
+        y_fused = fused_rms_norm_rope(
+            x_fused, w_fused, cos, sin, eps=eps, unit_offset=True
+        )
+        y_fused.sum().backward()
+
+        cos_sim_x = torch.nn.functional.cosine_similarity(
+            x_fused.grad.flatten().float(), x_ref.grad.flatten().float(), dim=0
+        )
+        cos_sim_w = torch.nn.functional.cosine_similarity(
+            w_fused.grad.flatten().float(), w_ref.grad.flatten().float(), dim=0
+        )
+        assert cos_sim_x > 0.999, f"x grad cosine_sim={cos_sim_x:.6f}"
+        assert cos_sim_w > 0.995, f"w grad cosine_sim={cos_sim_w:.6f}"
+
+
+class TestCompile:
+    def test_compile_fullgraph(self):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 1, 8, 2, 32
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16) * 0.1
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+
+        eager = fused_rms_norm_rope(x, weight, cos, sin, eps=eps, unit_offset=True)
+        compiled_fn = torch.compile(fused_rms_norm_rope, fullgraph=True)
+        compiled = compiled_fn(x, weight, cos, sin, eps=eps, unit_offset=True)
+
+        torch.testing.assert_close(compiled, eager, rtol=0, atol=0)
+
+
+class TestPartialRotary:
+    """``unit_offset=True`` combined with ``n_rot < D`` (Qwen3.5/Qwen3.6 partial rotary)."""
+
+    def test_forward_matches_reference(self):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D, n_rot = 2, 32, 4, 128, 64
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16) * 0.1
+        cos = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+
+        y_fused = fused_rms_norm_rope(
+            x.clone(), weight, cos, sin, eps=eps, unit_offset=True
+        )
+        y_ref_fp32 = _reference_fp32(x, weight, cos, sin, eps, unit_offset=True)
+        y_eager = _reference_unit_offset_partial(x.clone(), weight, cos, sin, eps)
+        _assert_at_bf16_floor(y_fused, y_ref_fp32, y_eager)
+
+    def test_backward_x_and_w_grad_match_eager(self):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D, n_rot = 2, 16, 4, 128, 64
+        eps = 1e-6
+        cos = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        weight_init = torch.randn(D, device="cuda", dtype=torch.bfloat16) * 0.1
+
+        x_ref = torch.randn(
+            B, S, H, D, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        w_ref = weight_init.clone().requires_grad_(True)
+        y_ref = _reference_unit_offset_partial(x_ref, w_ref, cos, sin, eps)
+        y_ref.sum().backward()
+
+        x_fused = x_ref.data.clone().requires_grad_(True)
+        w_fused = weight_init.clone().requires_grad_(True)
+        y_fused = fused_rms_norm_rope(
+            x_fused, w_fused, cos, sin, eps=eps, unit_offset=True
+        )
+        y_fused.sum().backward()
+
+        cos_sim_x = torch.nn.functional.cosine_similarity(
+            x_fused.grad.flatten().float(), x_ref.grad.flatten().float(), dim=0
+        )
+        cos_sim_w = torch.nn.functional.cosine_similarity(
+            w_fused.grad.flatten().float(), w_ref.grad.flatten().float(), dim=0
+        )
+        assert cos_sim_x > 0.999, f"x grad cosine_sim={cos_sim_x:.6f}"
+        assert cos_sim_w > 0.995, f"w grad cosine_sim={cos_sim_w:.6f}"
+
+    def test_compile_fullgraph(self):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D, n_rot = 1, 8, 2, 64, 32
+        eps = 1e-6
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(D, device="cuda", dtype=torch.bfloat16) * 0.1
+        cos = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, n_rot, device="cuda", dtype=torch.bfloat16)
+
+        eager = fused_rms_norm_rope(x, weight, cos, sin, eps=eps, unit_offset=True)
+        compiled_fn = torch.compile(fused_rms_norm_rope, fullgraph=True)
+        compiled = compiled_fn(x, weight, cos, sin, eps=eps, unit_offset=True)
+
+        torch.testing.assert_close(compiled, eager, rtol=0, atol=0)
diff --git a/tests/kernels/test_op_registry.py b/tests/kernels/test_op_registry.py
new file mode 100644
index 0000000000..2b3da60392
--- /dev/null
+++ b/tests/kernels/test_op_registry.py
@@ -0,0 +1,352 @@
+"""Tests for custom-op registration of axolotl kernels."""
+
+import pytest
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="requires CUDA"
+)
+
+
+class _OpRecorder(TorchDispatchMode):
+    def __init__(self):
+        self.seen = []
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        self.seen.append(str(func))
+        return func(*args, **(kwargs or {}))
+
+
+class TestOpsRegistered:
+    def test_flash_d512_ops_exist(self):
+        import axolotl.monkeypatch.attention.flash_attn_d512  # noqa: F401
+
+        assert torch.ops.axolotl.flash_attn_d512_fwd is not None
+        assert torch.ops.axolotl.flash_attn_d512_bwd is not None
+
+    def test_activation_ops_exist(self):
+        import axolotl.kernels.geglu  # noqa: F401
+        import axolotl.kernels.swiglu  # noqa: F401
+
+        assert torch.ops.axolotl.swiglu_fwd is not None
+        assert torch.ops.axolotl.swiglu_bwd is not None
+        assert torch.ops.axolotl.geglu_fwd is not None
+        assert torch.ops.axolotl.geglu_bwd is not None
+
+    def test_rms_norm_gated_ops_exist(self):
+        import axolotl.kernels.rms_norm_gated  # noqa: F401
+
+        assert torch.ops.axolotl.rms_norm_gated_fwd is not None
+        assert torch.ops.axolotl.rms_norm_gated_bwd is not None
+
+    def test_trainer_utils_ops_exist(self):
+        import axolotl.monkeypatch.trainer.utils  # noqa: F401
+
+        assert torch.ops.axolotl.selective_log_softmax_fwd is not None
+        assert torch.ops.axolotl.selective_log_softmax_bwd is not None
+        assert torch.ops.axolotl.entropy_from_logits is not None
+
+    def test_ebft_ops_exist(self):
+        import axolotl.core.trainers.ebft.kernels  # noqa: F401
+
+        assert torch.ops.axolotl.ebft_fused_log_softmax_gather is not None
+        assert torch.ops.axolotl.ebft_fused_reinforce_loss is not None
+        assert torch.ops.axolotl.ebft_fused_cosine_similarity is not None
+        assert torch.ops.axolotl.ebft_fused_diversity_penalty is not None
+
+
+class TestDispatchVisibility:
+    @requires_cuda
+    def test_swiglu_visible_to_dispatch_mode(self):
+        from axolotl.kernels.swiglu import swiglu_forward
+
+        gate = torch.randn(1, 8, 64, device="cuda")
+        up = torch.randn(1, 8, 64, device="cuda")
+        with _OpRecorder() as rec:
+            swiglu_forward(gate, up)
+        assert any("axolotl.swiglu_fwd" in s for s in rec.seen)
+
+    @requires_cuda
+    def test_flash_d512_visible_to_dispatch_mode(self):
+        from axolotl.monkeypatch.attention.flash_attn_d512 import flash_d512
+
+        q, k, v = (
+            torch.randn(1, 2, 64, 512, device="cuda", dtype=torch.bfloat16)
+            for _ in range(3)
+        )
+        with _OpRecorder() as rec:
+            flash_d512(q, k, v, causal=True)
+        assert any("axolotl.flash_attn_d512_fwd" in s for s in rec.seen)
+
+
+class TestFakeTensor:
+    @requires_cuda
+    def test_flash_d512_fwd_fake_shapes(self):
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        import axolotl.monkeypatch.attention.flash_attn_d512  # noqa: F401
+
+        with FakeTensorMode():
+            q = torch.empty(2, 4, 128, 512, device="cuda", dtype=torch.bfloat16)
+            k = torch.empty_like(q)
+            v = torch.empty_like(q)
+            pos = torch.zeros(2, 128, device="cuda", dtype=torch.int32)
+            o, lse = torch.ops.axolotl.flash_attn_d512_fwd(
+                q, k, v, pos, True, False, 0.044
+            )
+            assert o.shape == q.shape and o.dtype == q.dtype
+            assert lse.shape == (8, 128) and lse.dtype == torch.float32
+
+    @requires_cuda
+    def test_swiglu_fwd_fake_shapes(self):
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        import axolotl.kernels.swiglu  # noqa: F401
+
+        with FakeTensorMode():
+            gate = torch.empty(1, 8, 64, device="cuda", dtype=torch.bfloat16)
+            up = torch.empty_like(gate)
+            out = torch.ops.axolotl.swiglu_fwd(gate, up)
+            assert out.shape == gate.shape and out.dtype == gate.dtype
+
+
+class TestNewOpsDispatchVisibility:
+    @requires_cuda
+    def test_rms_norm_gated_visible_to_dispatch_mode(self):
+        from axolotl.kernels.rms_norm_gated import FusedRMSNormGated
+
+        mod = FusedRMSNormGated(64).to("cuda", torch.bfloat16)
+        x = torch.randn(2, 8, 64, device="cuda", dtype=torch.bfloat16)
+        g = torch.randn_like(x)
+        with _OpRecorder() as rec:
+            mod(x, gate=g)
+        assert any("axolotl.rms_norm_gated_fwd" in s for s in rec.seen)
+
+    @requires_cuda
+    def test_selective_log_softmax_visible_to_dispatch_mode(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        logits = torch.randn(2, 16, 512, device="cuda", dtype=torch.bfloat16)
+        index = torch.randint(0, 512, (2, 16), device="cuda")
+        with _OpRecorder() as rec:
+            selective_log_softmax(logits, index)
+        assert any("axolotl.selective_log_softmax_fwd" in s for s in rec.seen)
+
+    @requires_cuda
+    def test_entropy_visible_to_dispatch_mode(self):
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        logits = torch.randn(2, 16, 512, device="cuda", dtype=torch.bfloat16)
+        with _OpRecorder() as rec:
+            entropy_from_logits(logits)
+        assert any("axolotl.entropy_from_logits" in s for s in rec.seen)
+
+    @requires_cuda
+    def test_ebft_ops_visible_to_dispatch_mode(self):
+        from axolotl.core.trainers.ebft.kernels import (
+            fused_cosine_similarity,
+            fused_diversity_penalty,
+            fused_log_softmax_gather,
+            fused_reinforce_loss,
+        )
+
+        logits = torch.randn(2, 8, 512, device="cuda", dtype=torch.bfloat16)
+        labels = torch.randint(0, 512, (2, 8), device="cuda")
+        logps = torch.randn(64, device="cuda")
+        advs = torch.randn(64, device="cuda")
+        mask = torch.rand(64, device="cuda") > 0.5
+        a = torch.randn(4, 128, device="cuda")
+        emb = torch.randn(2, 4, 64, device="cuda")
+        with _OpRecorder() as rec:
+            fused_log_softmax_gather(logits, labels)
+            fused_reinforce_loss(logps, advs, mask)
+            fused_cosine_similarity(a, a.clone())
+            fused_diversity_penalty(emb)
+        for op in (
+            "ebft_fused_log_softmax_gather",
+            "ebft_fused_reinforce_loss",
+            "ebft_fused_cosine_similarity",
+            "ebft_fused_diversity_penalty",
+        ):
+            assert any(f"axolotl.{op}" in s for s in rec.seen), op
+
+
+class TestNewOpsFakeTensor:
+    @requires_cuda
+    def test_rms_norm_gated_fwd_fake_shapes(self):
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        import axolotl.kernels.rms_norm_gated  # noqa: F401
+
+        with FakeTensorMode():
+            x = torch.empty(16, 64, device="cuda", dtype=torch.bfloat16)
+            g = torch.empty_like(x)
+            w = torch.empty(64, device="cuda", dtype=torch.bfloat16)
+            y, rstd = torch.ops.axolotl.rms_norm_gated_fwd(x, g, w, 1e-6, 0.0, 64, 4)
+            assert y.shape == x.shape and y.dtype == x.dtype
+            assert rstd.shape == (16,) and rstd.dtype == torch.float32
+
+    @requires_cuda
+    def test_selective_log_softmax_fwd_fake_shapes(self):
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        import axolotl.monkeypatch.trainer.utils  # noqa: F401
+
+        with FakeTensorMode():
+            logits = torch.empty(32, 512, device="cuda", dtype=torch.bfloat16)
+            index = torch.empty(32, 1, device="cuda", dtype=torch.int64)
+            out, lse = torch.ops.axolotl.selective_log_softmax_fwd(
+                logits, index, 1, 1, 512, 4096, 8192
+            )
+            assert out.shape == (32, 1) and out.dtype == torch.float32
+            assert lse.shape == (32,) and lse.dtype == torch.float32
+
+    @requires_cuda
+    def test_entropy_fake_shapes(self):
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        import axolotl.monkeypatch.trainer.utils  # noqa: F401
+
+        with FakeTensorMode():
+            logits = torch.empty(2, 16, 512, device="cuda", dtype=torch.bfloat16)
+            out = torch.ops.axolotl.entropy_from_logits(logits, 128)
+            assert out.shape == (32,) and out.dtype == torch.float32
+
+    @requires_cuda
+    def test_ebft_fake_shapes(self):
+        from torch._subclasses.fake_tensor import FakeTensorMode
+
+        import axolotl.core.trainers.ebft.kernels  # noqa: F401
+
+        with FakeTensorMode():
+            logits = torch.empty(2, 8, 512, device="cuda", dtype=torch.bfloat16)
+            labels = torch.empty(2, 8, device="cuda", dtype=torch.int64)
+            out = torch.ops.axolotl.ebft_fused_log_softmax_gather(logits, labels)
+            assert out.shape == (2, 8) and out.dtype == torch.float32
+
+            logps = torch.empty(64, device="cuda")
+            mask = torch.empty(64, device="cuda", dtype=torch.bool)
+            loss = torch.ops.axolotl.ebft_fused_reinforce_loss(logps, logps, mask)
+            assert loss.shape == () and loss.dtype == torch.float32
+
+            a = torch.empty(4, 8, 128, device="cuda")
+            sim = torch.ops.axolotl.ebft_fused_cosine_similarity(a, a)
+            assert sim.shape == (4, 8) and sim.dtype == torch.float32
+
+            emb = torch.empty(2, 4, 64, device="cuda")
+            div = torch.ops.axolotl.ebft_fused_diversity_penalty(emb)
+            assert div.shape == (2, 4) and div.dtype == torch.float32
+
+
+class TestCompile:
+    @requires_cuda
+    def test_swiglu_compiles_fullgraph(self):
+        from axolotl.kernels.swiglu import swiglu_forward
+
+        def fn(gate, up):
+            return swiglu_forward(gate, up).sum()
+
+        gate = torch.randn(1, 8, 64, device="cuda")
+        up = torch.randn(1, 8, 64, device="cuda")
+        eager = fn(gate, up)
+        compiled = torch.compile(fn, fullgraph=True)(gate, up)
+        torch.testing.assert_close(eager, compiled)
+
+    @requires_cuda
+    def test_flash_d512_compiles_fullgraph(self):
+        from axolotl.monkeypatch.attention.flash_attn_d512 import flash_d512
+
+        def fn(q, k, v):
+            return flash_d512(q, k, v, causal=True).sum()
+
+        q, k, v = (
+            torch.randn(1, 2, 64, 512, device="cuda", dtype=torch.bfloat16)
+            for _ in range(3)
+        )
+        eager = fn(q, k, v)
+        compiled = torch.compile(fn, fullgraph=True)(q, k, v)
+        torch.testing.assert_close(eager, compiled, atol=2e-2, rtol=2e-2)
+
+
+class TestOpcheck:
+    """torch.library.opcheck validates schema, fake fidelity, aliasing, and
+    functionalization in one shot."""
+
+    @requires_cuda
+    def test_swiglu_ops(self):
+        from axolotl.kernels.swiglu import _swiglu_bwd_op, swiglu_forward
+
+        gate = torch.randn(1, 8, 64, device="cuda")
+        up = torch.randn(1, 8, 64, device="cuda")
+        go = torch.randn(1, 8, 64, device="cuda")
+        torch.library.opcheck(swiglu_forward, (gate.clone(), up.clone()))
+        torch.library.opcheck(_swiglu_bwd_op, (go, gate, up))
+
+    @requires_cuda
+    def test_geglu_ops(self):
+        from axolotl.kernels.geglu import _geglu_bwd_op, geglu_forward
+
+        gate = torch.randn(1, 8, 64, device="cuda")
+        up = torch.randn(1, 8, 64, device="cuda")
+        go = torch.randn(1, 8, 64, device="cuda")
+        torch.library.opcheck(geglu_forward, (gate.clone(), up.clone()))
+        torch.library.opcheck(_geglu_bwd_op, (go, gate, up))
+
+    @requires_cuda
+    def test_flash_d512_fwd(self):
+        from axolotl.monkeypatch.attention.flash_attn_d512 import _flash_d512_fwd_op
+
+        q, k, v = (
+            torch.randn(1, 1, 32, 512, device="cuda", dtype=torch.bfloat16)
+            for _ in range(3)
+        )
+        pos = torch.zeros(1, 32, device="cuda", dtype=torch.int32)
+        torch.library.opcheck(_flash_d512_fwd_op, (q, k, v, pos, True, False, 0.044))
+
+
+class TestMutatingBwdUnderCompile:
+    @requires_cuda
+    def test_swiglu_backward_functionalizes_correctly(self):
+        """The mutate-then-read-back wrapper must survive compile's
+        auto-functionalization (downstream uses of mutated args are remapped)."""
+        from axolotl.kernels.swiglu import swiglu_backward
+
+        def fn(go, gate, up):
+            h, dg, du = swiglu_backward(go, gate, up)
+            return h.sum() + dg.sum() * 2 + du.sum() * 3
+
+        go, gate, up = (torch.randn(1, 16, 128, device="cuda") for _ in range(3))
+        eager = fn(go.clone(), gate.clone(), up.clone())
+        compiled = torch.compile(fn, fullgraph=True)(
+            go.clone(), gate.clone(), up.clone()
+        )
+        torch.testing.assert_close(eager, compiled)
+
+
+class TestFallbackShim:
+    def test_unregistered_op_passthrough(self):
+        from axolotl.kernels.op_registry import _UnregisteredOp
+
+        def raw(x):
+            return x + 1
+
+        shim = _UnregisteredOp(raw)
+        assert shim(torch.tensor(1)).item() == 2
+        # register_fake is a no-op that must not raise
+        shim.register_fake(lambda x: x)
+
+    def test_failed_registration_falls_back(self):
+        from typing import Callable
+
+        from axolotl.kernels.op_registry import _UnregisteredOp, register_kernel_op
+
+        # a callable arg is not an op-schema type: registration fails and the
+        # decorator degrades to the raw callable
+        @register_kernel_op("test_bad_schema_op")
+        def _bad(x: torch.Tensor, fn: Callable) -> torch.Tensor:
+            return fn(x)
+
+        assert isinstance(_bad, _UnregisteredOp)
+        assert _bad(torch.tensor(1), lambda t: t + 2).item() == 3
+        _bad.register_fake(lambda x, fn: x)
diff --git a/tests/kernels/test_rms_norm_gated.py b/tests/kernels/test_rms_norm_gated.py
new file mode 100644
index 0000000000..10b379184d
--- /dev/null
+++ b/tests/kernels/test_rms_norm_gated.py
@@ -0,0 +1,229 @@
+"""
+Correctness tests for fused RMSNorm + SiLU Gate kernel.
+
+Tests against the eager Qwen3_5RMSNormGated implementation.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytest.importorskip("triton", reason="triton required for fused kernels")
+
+if not torch.cuda.is_available():
+    pytest.skip("CUDA required for fused kernel tests", allow_module_level=True)
+
+from axolotl.kernels.rms_norm_gated import FusedRMSNormGated
+
+
+class EagerRMSNormGated(torch.nn.Module):
+    """Reference implementation matching Qwen3_5RMSNormGated exactly."""
+
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states, gate=None):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        hidden_states = self.weight * hidden_states.to(input_dtype)
+        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        return hidden_states.to(input_dtype)
+
+
+def _sync_weights(eager_mod, fused_mod):
+    """Copy weights from eager to fused module."""
+    fused_mod.weight.data.copy_(eager_mod.weight.data)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (2, 128, 256),
+        (4, 64, 512),
+        (1, 32, 1024),
+        (2, 16, 2560),  # Qwen3.5-4B hidden_size
+        (2, 16, 4096),  # Qwen3.5-9B hidden_size
+        (1, 8, 5120),  # Qwen3.5-27B hidden_size
+        (4, 16, 2048),  # Qwen3.5-35B-A3B (MoE) hidden_size
+        (4, 16, 3072),  # Qwen3.5-122B-A10B (MoE) hidden_size
+    ],
+)
+class TestRMSNormGatedForward:
+    def test_output_matches_eager(self, dtype, shape):
+        torch.manual_seed(42)
+        B, T, H = shape
+        X = torch.randn(B, T, H, dtype=dtype, device="cuda")
+        G = torch.randn(B, T, H, dtype=dtype, device="cuda")
+
+        eager = EagerRMSNormGated(H).to(dtype=dtype, device="cuda")
+        fused = FusedRMSNormGated(H).to(dtype=dtype, device="cuda")
+        _sync_weights(eager, fused)
+
+        y_eager = eager(X, gate=G)
+        y_fused = fused(X, gate=G)
+
+        if dtype == torch.float32:
+            torch.testing.assert_close(y_fused, y_eager, atol=1e-5, rtol=1e-5)
+        else:
+            torch.testing.assert_close(y_fused, y_eager, atol=1e-2, rtol=1e-2)
+
+    def test_output_shape(self, dtype, shape):
+        B, T, H = shape
+        X = torch.randn(B, T, H, dtype=dtype, device="cuda")
+        G = torch.randn(B, T, H, dtype=dtype, device="cuda")
+
+        fused = FusedRMSNormGated(H).to(dtype=dtype, device="cuda")
+        y = fused(X, gate=G)
+        assert y.shape == (B, T, H)
+        assert y.dtype == dtype
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (2, 32, 256),
+        (2, 16, 512),
+        (2, 16, 2560),  # Qwen3.5-4B
+        (1, 8, 4096),  # Qwen3.5-9B
+        (1, 8, 5120),  # Qwen3.5-27B
+        (2, 16, 2048),  # Qwen3.5-35B-A3B (MoE)
+        (2, 16, 3072),  # Qwen3.5-122B-A10B (MoE)
+    ],
+)
+class TestRMSNormGatedBackward:
+    def test_grad_x(self, dtype, shape):
+        torch.manual_seed(42)
+        B, T, H = shape
+        X = torch.randn(B, T, H, dtype=dtype, device="cuda", requires_grad=True)
+        G = torch.randn(B, T, H, dtype=dtype, device="cuda", requires_grad=True)
+        X_ref = X.detach().clone().requires_grad_(True)
+        G_ref = G.detach().clone().requires_grad_(True)
+
+        eager = EagerRMSNormGated(H).to(dtype=dtype, device="cuda")
+        fused = FusedRMSNormGated(H).to(dtype=dtype, device="cuda")
+        _sync_weights(eager, fused)
+
+        y_eager = eager(X_ref, gate=G_ref)
+        y_fused = fused(X, gate=G)
+
+        grad_out = torch.randn_like(y_eager)
+        y_eager.backward(grad_out)
+        y_fused.backward(grad_out)
+
+        if dtype == torch.float32:
+            atol, rtol = 1e-4, 1e-4
+        else:
+            atol, rtol = 5e-2, 5e-2
+
+        torch.testing.assert_close(X.grad, X_ref.grad, atol=atol, rtol=rtol)
+
+    def test_grad_gate(self, dtype, shape):
+        torch.manual_seed(42)
+        B, T, H = shape
+        X = torch.randn(B, T, H, dtype=dtype, device="cuda", requires_grad=True)
+        G = torch.randn(B, T, H, dtype=dtype, device="cuda", requires_grad=True)
+        X_ref = X.detach().clone().requires_grad_(True)
+        G_ref = G.detach().clone().requires_grad_(True)
+
+        eager = EagerRMSNormGated(H).to(dtype=dtype, device="cuda")
+        fused = FusedRMSNormGated(H).to(dtype=dtype, device="cuda")
+        _sync_weights(eager, fused)
+
+        y_eager = eager(X_ref, gate=G_ref)
+        y_fused = fused(X, gate=G)
+
+        grad_out = torch.randn_like(y_eager)
+        y_eager.backward(grad_out)
+        y_fused.backward(grad_out)
+
+        if dtype == torch.float32:
+            atol, rtol = 1e-4, 1e-4
+        else:
+            atol, rtol = 5e-2, 5e-2
+
+        torch.testing.assert_close(G.grad, G_ref.grad, atol=atol, rtol=rtol)
+
+    def test_grad_weight(self, dtype, shape):
+        torch.manual_seed(42)
+        B, T, H = shape
+        X = torch.randn(B, T, H, dtype=dtype, device="cuda", requires_grad=True)
+        G = torch.randn(B, T, H, dtype=dtype, device="cuda", requires_grad=True)
+        X_ref = X.detach().clone().requires_grad_(True)
+        G_ref = G.detach().clone().requires_grad_(True)
+
+        eager = EagerRMSNormGated(H).to(dtype=dtype, device="cuda")
+        fused = FusedRMSNormGated(H).to(dtype=dtype, device="cuda")
+        _sync_weights(eager, fused)
+
+        y_eager = eager(X_ref, gate=G_ref)
+        y_fused = fused(X, gate=G)
+
+        grad_out = torch.randn_like(y_eager)
+        y_eager.backward(grad_out)
+        y_fused.backward(grad_out)
+
+        if dtype == torch.float32:
+            atol, rtol = 1e-4, 1e-4
+        else:
+            atol, rtol = 5e-2, 5e-2
+
+        torch.testing.assert_close(
+            fused.weight.grad, eager.weight.grad, atol=atol, rtol=rtol
+        )
+
+
+class TestRMSNormGatedEdgeCases:
+    def test_gate_none_raises(self):
+        fused = FusedRMSNormGated(256).cuda()
+        X = torch.randn(2, 4, 256, device="cuda")
+        with pytest.raises(ValueError, match="requires a gate tensor"):
+            fused(X, gate=None)
+
+    def test_2d_input(self):
+        """Test with (BxT, H) shaped input instead of (B, T, H)."""
+        torch.manual_seed(42)
+        H = 512
+        X = torch.randn(64, H, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        G = torch.randn(64, H, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        X_ref = X.detach().clone().requires_grad_(True)
+        G_ref = G.detach().clone().requires_grad_(True)
+
+        eager = EagerRMSNormGated(H).to(dtype=torch.bfloat16, device="cuda")
+        fused = FusedRMSNormGated(H).to(dtype=torch.bfloat16, device="cuda")
+        _sync_weights(eager, fused)
+
+        y_eager = eager(X_ref, gate=G_ref)
+        y_fused = fused(X, gate=G)
+
+        torch.testing.assert_close(y_fused, y_eager, atol=1e-2, rtol=1e-2)
+
+        grad_out = torch.randn_like(y_eager)
+        y_eager.backward(grad_out)
+        y_fused.backward(grad_out)
+
+        torch.testing.assert_close(X.grad, X_ref.grad, atol=5e-2, rtol=5e-2)
+        torch.testing.assert_close(G.grad, G_ref.grad, atol=5e-2, rtol=5e-2)
+
+    def test_random_weight_init(self):
+        """Test with non-default weight values."""
+        torch.manual_seed(123)
+        H = 256
+        X = torch.randn(2, 16, H, dtype=torch.bfloat16, device="cuda")
+        G = torch.randn(2, 16, H, dtype=torch.bfloat16, device="cuda")
+
+        eager = EagerRMSNormGated(H).to(dtype=torch.bfloat16, device="cuda")
+        # Randomize weights
+        eager.weight.data = torch.randn_like(eager.weight.data)
+
+        fused = FusedRMSNormGated(H).to(dtype=torch.bfloat16, device="cuda")
+        _sync_weights(eager, fused)
+
+        y_eager = eager(X, gate=G)
+        y_fused = fused(X, gate=G)
+        torch.testing.assert_close(y_fused, y_eager, atol=1e-2, rtol=1e-2)
diff --git a/tests/loaders/test_patch_manager_gemma4_unified.py b/tests/loaders/test_patch_manager_gemma4_unified.py
new file mode 100644
index 0000000000..4a3a6e0de8
--- /dev/null
+++ b/tests/loaders/test_patch_manager_gemma4_unified.py
@@ -0,0 +1,129 @@
+"""Tests for ``PatchManager`` fused-attn dispatch over the gemma4 / gemma4_unified
+families. These pin the unified-specific routing (namespace branch + auto-on with
+no ``fused_attn_kernel`` gate) introduced when the two blocks were consolidated."""
+
+import logging
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from axolotl.loaders.patch_manager import PatchManager
+from axolotl.utils.dict import DictDefault
+
+
+def _make_patch_manager(model_config_type, *, gradient_checkpointing, inference=False):
+    """Minimal cfg that reaches the fused-attn dispatch. ``context_parallel_size``
+    must be an int (the pre-CUDA span compares ``> 1``) and ``fused_attn_kernel``
+    stays falsy to prove the unified path is auto-on without it."""
+    cfg = DictDefault(
+        {
+            "model_config_type": model_config_type,
+            "context_parallel_size": 1,
+            "sample_packing": False,
+            "gradient_checkpointing": gradient_checkpointing,
+            "activation_offloading": False,
+            "fused_attn_kernel": False,
+            "fsdp_config": None,
+        }
+    )
+    return PatchManager(cfg, MagicMock(), inference=inference)
+
+
+@pytest.fixture
+def spy_fused_patches(monkeypatch):
+    """Replace the two fused-attn patch entry points (imported inside the
+    dispatch) with call recorders, and force the CUDA-gated block to run."""
+    import axolotl.monkeypatch.models.gemma4.fused_attn as regular_fa
+    import axolotl.monkeypatch.models.gemma4_unified.fused_attn as unified_fa
+
+    calls = {"unified": [], "regular": []}
+    monkeypatch.setattr(
+        unified_fa,
+        "patch_gemma4_unified_fused_attn",
+        lambda **kw: calls["unified"].append(kw),
+    )
+    monkeypatch.setattr(
+        regular_fa,
+        "patch_gemma4_fused_attn",
+        lambda **kw: calls["regular"].append(kw),
+    )
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    return calls
+
+
+class TestFusedAttnDispatch:
+    @pytest.mark.parametrize(
+        "model_config_type",
+        ["gemma4_unified", "gemma4_unified_text"],
+    )
+    def test_unified_routes_to_unified_patch_auto_on(
+        self, spy_fused_patches, model_config_type
+    ):
+        """Unified dispatches to the unified namespace and runs even though
+        ``fused_attn_kernel`` is False (auto-on, like standard gemma4)."""
+        pm = _make_patch_manager(model_config_type, gradient_checkpointing=True)
+        pm._apply_model_specific_patches()
+
+        assert spy_fused_patches["unified"] == [{"install_shared_kv_workaround": True}]
+        assert spy_fused_patches["regular"] == []
+
+    @pytest.mark.parametrize("model_config_type", ["gemma4", "gemma4_text"])
+    def test_standard_gemma4_routes_to_regular_patch(
+        self, spy_fused_patches, model_config_type
+    ):
+        pm = _make_patch_manager(model_config_type, gradient_checkpointing=False)
+        pm._apply_model_specific_patches()
+
+        assert spy_fused_patches["regular"] == [{"install_shared_kv_workaround": False}]
+        assert spy_fused_patches["unified"] == []
+
+    def test_shared_kv_workaround_off_without_checkpointing(self, spy_fused_patches):
+        pm = _make_patch_manager("gemma4_unified", gradient_checkpointing=False)
+        pm._apply_model_specific_patches()
+
+        assert spy_fused_patches["unified"] == [{"install_shared_kv_workaround": False}]
+
+    def test_shared_kv_workaround_off_during_inference(self, spy_fused_patches):
+        pm = _make_patch_manager(
+            "gemma4_unified", gradient_checkpointing=True, inference=True
+        )
+        pm._apply_model_specific_patches()
+
+        assert spy_fused_patches["unified"] == [{"install_shared_kv_workaround": False}]
+
+
+class TestWarnIfFusedAttnUnsupported:
+    """``_FUSED_ATTN_KERNEL_SUPPORTED`` must include the unified model types so
+    ``fused_attn_kernel: true`` does not warn (it is not a no-op for them)."""
+
+    def _capture(self, caplog, cfg):
+        logger = logging.getLogger("axolotl.loaders.patch_manager")
+        logger.addHandler(caplog.handler)
+        previous_level = logger.level
+        logger.setLevel(logging.WARNING)
+        try:
+            PatchManager._warn_if_fused_attn_unsupported(cfg)
+        finally:
+            logger.removeHandler(caplog.handler)
+            logger.setLevel(previous_level)
+
+    @pytest.mark.parametrize(
+        "model_config_type",
+        ["gemma4_unified", "gemma4_unified_text"],
+    )
+    def test_no_warn_for_unified_types(self, caplog, model_config_type):
+        self._capture(
+            caplog,
+            DictDefault(
+                {"fused_attn_kernel": True, "model_config_type": model_config_type}
+            ),
+        )
+        assert caplog.text == "", f"unexpected warning: {caplog.text}"
+
+    def test_warns_for_unsupported_type(self, caplog):
+        self._capture(
+            caplog,
+            DictDefault({"fused_attn_kernel": True, "model_config_type": "llama"}),
+        )
+        assert "llama" in caplog.text and "no-op" in caplog.text
diff --git a/tests/monkeypatch/test_checkpoint_activation_offload.py b/tests/monkeypatch/test_checkpoint_activation_offload.py
new file mode 100644
index 0000000000..856449287d
--- /dev/null
+++ b/tests/monkeypatch/test_checkpoint_activation_offload.py
@@ -0,0 +1,55 @@
+from functools import partial
+
+import torch
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from transformers import GradientCheckpointingLayer
+
+from axolotl.monkeypatch.checkpoint_activation_offload import (
+    CheckpointHiddenStatesOffload,
+)
+
+
+class TinyCheckpointLayer(GradientCheckpointingLayer):
+    def forward(self, hidden_states):
+        hidden_states = hidden_states.sin()
+        return hidden_states * hidden_states
+
+
+def test_checkpoint_offload_marks_non_reentrant_checkpoint_input():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    layer = TinyCheckpointLayer()
+    layer.gradient_checkpointing = True
+    layer._gradient_checkpointing_func = partial(checkpoint, use_reentrant=False)
+    layer.to(device)
+    layer.train()
+    hidden_states = torch.randn(4, 8, device=device, requires_grad=True)
+
+    offload = CheckpointHiddenStatesOffload(use_streams=False, min_offload_size=0)
+    with offload:
+        loss = layer(hidden_states).sum()
+        loss.backward()
+
+    assert hidden_states.grad is not None
+    assert offload.stats.marked_tensors == 1
+    assert offload.stats.saved_tensors_seen >= offload.stats.marked_tensors
+    if hidden_states.device.type == "cuda":
+        assert offload.stats.offloaded_tensors == 1
+        assert offload.stats.restored_tensors == 1
+    else:
+        assert offload.stats.skipped_marked_tensors == 1
+
+
+def test_checkpoint_offload_ignores_unmarked_saved_tensors():
+    hidden_states = torch.randn(4, 8, requires_grad=True)
+    linear = nn.Linear(8, 8)
+
+    offload = CheckpointHiddenStatesOffload(use_streams=False, min_offload_size=0)
+    with offload:
+        loss = linear(hidden_states).square().sum()
+        loss.backward()
+
+    assert hidden_states.grad is not None
+    assert offload.stats.saved_tensors_seen > 0
+    assert offload.stats.marked_tensors == 0
+    assert offload.stats.offloaded_tensors == 0
diff --git a/tests/monkeypatch/test_flash_attn_d512.py b/tests/monkeypatch/test_flash_attn_d512.py
new file mode 100644
index 0000000000..f716ebd1b0
--- /dev/null
+++ b/tests/monkeypatch/test_flash_attn_d512.py
@@ -0,0 +1,117 @@
+"""Correctness tests for the Triton head_dim=512 flash attention (fwd+bwd, dense + varlen)."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+DEV = "cuda"
+D = 512
+
+
+def _cos(a, b):
+    return F.cosine_similarity(a.float().flatten(), b.float().flatten(), 0).item()
+
+
+def _rel(a, b):
+    return (a.float() - b.float()).norm().item() / max(b.float().norm().item(), 1e-12)
+
+
+@pytest.mark.parametrize(
+    "B,H,N,causal", [(1, 8, 1024, True), (2, 8, 1024, True), (1, 8, 1024, False)]
+)
+def test_dense_fwd_bwd_matches_sdpa(B, H, N, causal):
+    from axolotl.monkeypatch.attention.flash_attn_d512 import flash_d512
+
+    torch.manual_seed(0)
+    q = torch.randn(B, H, N, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    k = torch.randn(B, H, N, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    v = torch.randn(B, H, N, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    qr, kr, vr = (t.detach().clone().requires_grad_() for t in (q, k, v))
+    out = flash_d512(q, k, v, causal)
+    out.float().pow(2).mean().backward()
+    with sdpa_kernel([SDPBackend.MATH]):
+        ref = F.scaled_dot_product_attention(qr, kr, vr, is_causal=causal)
+    ref.float().pow(2).mean().backward()
+    assert _cos(out, ref) > 0.999
+    assert _cos(q.grad, qr.grad) > 0.999
+    assert _cos(k.grad, kr.grad) > 0.999
+    assert _cos(v.grad, vr.grad) > 0.999
+
+
+@pytest.mark.parametrize("scale", [None, 1.0])
+def test_noncontiguous_gqa_grads_match_sdpa(scale):
+    """Regression: real attention feeds non-contiguous q ([B,S,H,D].transpose(1,2)) with contiguous
+    GQA-repeated k/v. The backward must not read wrong memory from the stride mismatch. Uses a
+    magnitude-sensitive relative-error check (cosine alone misses the gradient blow-up)."""
+    from axolotl.monkeypatch.attention.flash_attn_d512 import flash_d512
+
+    B, Hq, Hkv, N = 1, 16, 4, 1024
+    ng = Hq // Hkv
+    torch.manual_seed(0)
+    # non-contiguous q exactly as produced by attention: [B, N, Hq, D] -> transpose -> [B, Hq, N, D]
+    q = torch.randn(B, N, Hq, D, device=DEV, dtype=torch.bfloat16).transpose(1, 2)
+    k = torch.randn(B, Hkv, N, D, device=DEV, dtype=torch.bfloat16)
+    v = torch.randn(B, Hkv, N, D, device=DEV, dtype=torch.bfloat16)
+    assert not q.is_contiguous()
+    q, k, v = (t.detach().requires_grad_() for t in (q, k, v))
+    qr, kr, vr = (t.detach().clone().requires_grad_() for t in (q, k, v))
+    eff = D**-0.5 if scale is None else scale
+
+    out = flash_d512(
+        q, k.repeat_interleave(ng, 1), v.repeat_interleave(ng, 1), True, scale=scale
+    )
+    out.float().pow(2).mean().backward()
+    with sdpa_kernel([SDPBackend.MATH]):
+        ref = F.scaled_dot_product_attention(
+            qr,
+            kr.repeat_interleave(ng, 1),
+            vr.repeat_interleave(ng, 1),
+            is_causal=True,
+            scale=eff,
+        )
+    ref.float().pow(2).mean().backward()
+
+    assert _rel(out, ref) < 0.01
+    assert _rel(q.grad, qr.grad) < 0.02
+    assert _rel(k.grad, kr.grad) < 0.02
+    assert _rel(v.grad, vr.grad) < 0.02
+
+
+@pytest.mark.parametrize("docs", [[512, 512], [300, 400, 324], [200, 300, 524]])
+def test_varlen_matches_per_document(docs):
+    """Packed (varlen) fwd+bwd must match per-document independent attention."""
+    from axolotl.monkeypatch.attention.flash_attn_d512 import flash_d512
+
+    B, H, N = 1, 8, sum(docs)
+    torch.manual_seed(0)
+    q = torch.randn(B, H, N, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    k = torch.randn(B, H, N, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    v = torch.randn(B, H, N, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    pos = torch.cat([torch.arange(d) for d in docs]).to(DEV)[None]
+    qr, kr, vr = (t.detach().clone().requires_grad_() for t in (q, k, v))
+
+    out = flash_d512(q, k, v, True, position_ids=pos)
+    out.float().pow(2).mean().backward()
+
+    outs, off = [], 0
+    for d in docs:
+        with sdpa_kernel([SDPBackend.MATH]):
+            outs.append(
+                F.scaled_dot_product_attention(
+                    qr[:, :, off : off + d],
+                    kr[:, :, off : off + d],
+                    vr[:, :, off : off + d],
+                    is_causal=True,
+                )
+            )
+        off += d
+    ref = torch.cat(outs, 2)
+    ref.float().pow(2).mean().backward()
+
+    assert _cos(out, ref) > 0.999
+    assert _cos(q.grad, qr.grad) > 0.999
+    assert _cos(k.grad, kr.grad) > 0.999
+    assert _cos(v.grad, vr.grad) > 0.999
diff --git a/tests/monkeypatch/test_float8_fsdp_sharding.py b/tests/monkeypatch/test_float8_fsdp_sharding.py
new file mode 100644
index 0000000000..494f04937b
--- /dev/null
+++ b/tests/monkeypatch/test_float8_fsdp_sharding.py
@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Axolotl AI
+# Licensed under the Apache License, Version 2.0
+
+"""CPU round-trip tests for the FSDP2 sharding substrate on torchao ``Float8Tensor``.
+
+The substrate (``axolotl.monkeypatch.accelerate.float8_fsdp``) teaches a blockwise-FP8
+``Float8Tensor`` the aten ops FSDP2 needs to dim-0 shard a frozen weight: ``split`` /
+``narrow`` / ``new_zeros`` / ``clone`` / ``detach`` / ``as_strided`` / ``copy_`` / ``view``
+plus the ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather`` reconstruct hooks. These are pure
+tensor logic (qdata splits by ``s``, the blockwise scale by ``s // block_size[0]``) and are
+fully exercisable on CPU - no GPU, no Blackwell, no fused kernel.
+
+Each test simulates the FSDP2 shard -> all-gather -> reconstruct flow and asserts the
+reconstructed tensor dequantizes equal to the unsharded original, across several world sizes
+including one where the (block-row) shard axis does NOT divide evenly. Covers the B11 data-
+independent FSDP2 collective contract for the Float8 substrate.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+torchao_quant = pytest.importorskip("torchao.quantization", reason="torchao required")
+Float8Tensor = torchao_quant.Float8Tensor
+KernelPreference = pytest.importorskip(
+    "torchao.quantization.quantize_.common", reason="torchao required"
+).KernelPreference
+
+from axolotl.monkeypatch.accelerate import float8_fsdp  # noqa: E402
+
+float8_fsdp.patch_float8_fsdp()
+
+
+def _cdiv(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
+def _make_blockwise_float8(N: int, K: int, b0: int, b1: int, seed: int = 0):
+    """Build a frozen blockwise-FP8 ``Float8Tensor`` [N,K] with scale [N//b0, K//b1].
+
+    This is the layout float8_fsdp targets (the 128x128-blocked weight torchao's own split
+    rejects with ``AssertionError("not yet implemented")``). PerRow/PerTensor checkpoints are
+    a special case (b0=1 or b0=N), covered separately via from_hp.
+    """
+    torch.manual_seed(seed)
+    W = torch.randn(N, K, dtype=torch.bfloat16)
+    scale = torch.empty(N // b0, K // b1)
+    qdata = torch.empty(N, K, dtype=torch.float8_e4m3fn)
+    for i in range(N // b0):
+        for j in range(K // b1):
+            blk = W[i * b0 : (i + 1) * b0, j * b1 : (j + 1) * b1].float()
+            s = blk.abs().amax().clamp_min(1e-6) / 448.0
+            scale[i, j] = s
+            qdata[i * b0 : (i + 1) * b0, j * b1 : (j + 1) * b1] = (blk / s).to(
+                torch.float8_e4m3fn
+            )
+    return Float8Tensor(
+        qdata, scale, [b0, b1], None, None, KernelPreference.AUTO, dtype=torch.bfloat16
+    )
+
+
+def _shard_all_gather_reconstruct(ft, world: int):
+    """Simulate one FSDP2 dim-0 shard -> all-gather -> reconstruct cycle on CPU.
+
+    Pads dim 0 to a per-rank-equal multiple of block_size[0] (FSDP2 pads uneven shards via
+    new_zeros + copy_ into a narrow), chunks into ``world`` shards, runs each shard's
+    fsdp_pre_all_gather, concatenates the gathered inner tensors, and reconstructs via
+    fsdp_post_all_gather. Returns the reconstructed (unpadded) Float8Tensor.
+    """
+    N = ft.shape[0]
+    blk0 = ft.block_size[0]
+    shard = _cdiv(_cdiv(N, blk0), world) * blk0  # multiple of block_size[0]
+    padded = shard * world
+    if padded != N:
+        full = ft.new_zeros([padded, *ft.shape[1:]])
+        full.narrow(0, 0, N).copy_(ft)
+    else:
+        full = ft.clone()
+    shards = torch.split(full, shard, 0)
+    assert len(shards) == world, (len(shards), world)
+    pre = [s.fsdp_pre_all_gather(mesh=None) for s in shards]
+    meta = pre[0][1]
+    gathered = tuple(
+        torch.cat([p[0][k] for p in pre], 0) for k in range(len(pre[0][0]))
+    )
+    recon, _ = shards[0].fsdp_post_all_gather(gathered, meta, torch.bfloat16)
+    return recon.narrow(0, 0, N)
+
+
+@pytest.mark.parametrize(
+    "N,K,b0,b1",
+    [(16, 32, 8, 16), (32, 64, 16, 16)],
+    ids=["blocked_8x16", "blocked_16x16"],
+)
+def test_float8_blockwise_split_roundtrips(N, K, b0, b1):
+    """split/narrow/slice on a blockwise Float8Tensor dequantize equal to the original slice."""
+    ft = _make_blockwise_float8(N, K, b0, b1)
+    orig = ft.dequantize()
+
+    parts = torch.split(ft, b0, 0)
+    rec = torch.cat([p.dequantize() for p in parts], 0)
+    assert torch.equal(rec, orig), "split round-trip diverged"
+
+    nar = torch.narrow(ft, 0, b0, b0)
+    assert torch.equal(nar.dequantize(), orig[b0 : 2 * b0]), (
+        "narrow round-trip diverged"
+    )
+
+    assert torch.equal(ft.clone().dequantize(), orig), "clone diverged"
+    assert torch.equal(ft.detach().dequantize(), orig), "detach diverged"
+
+
+@pytest.mark.parametrize(
+    "N,K,b0,b1,worlds",
+    [
+        (16, 32, 8, 16, (1, 2)),  # 2 block-rows, even splits
+        (24, 32, 8, 16, (1, 2, 3)),  # 3 block-rows; world=2 -> uneven (padded)
+    ],
+    ids=["even", "uneven_padded"],
+)
+def test_float8_fsdp_all_gather_roundtrip(N, K, b0, b1, worlds):
+    """FSDP2 shard -> all-gather -> reconstruct round-trips for blockwise Float8 across worlds.
+
+    Includes a world size where the block-row axis does NOT divide evenly (the shard is padded
+    on dim 0 via new_zeros/copy_ and narrowed back). Pure CPU tensor logic; no GPU needed.
+    """
+    ft = _make_blockwise_float8(N, K, b0, b1)
+    orig = ft.dequantize()
+    for world in worlds:
+        recon = _shard_all_gather_reconstruct(ft, world)
+        assert torch.equal(recon.dequantize(), orig), (
+            f"FSDP all-gather reconstruct diverged at world={world}"
+        )
+
+
+def test_float8_perrow_split_roundtrips():
+    """A PerRow Float8Tensor (block_size [1, K], the common from_hp checkpoint) shards on dim 0."""
+    torch.manual_seed(1)
+    W = torch.randn(16, 32, dtype=torch.bfloat16)
+    ft = Float8Tensor.from_hp(W)
+    assert ft.block_size[0] == 1, ft.block_size
+    orig = ft.dequantize()
+
+    parts = torch.split(ft, 4, 0)
+    rec = torch.cat([p.dequantize() for p in parts], 0)
+    assert torch.equal(rec, orig), "PerRow split round-trip diverged"
+
+    recon = _shard_all_gather_reconstruct(ft, 3)  # 16 rows / 3 -> uneven, padded
+    assert torch.equal(recon.dequantize(), orig), "PerRow uneven all-gather diverged"
+
+
+def test_float8_new_zeros_scale_is_coarser_by_block_size():
+    """new_zeros must scale the inner scale by block_size on every dim (FSDP pad invariant)."""
+    ft = _make_blockwise_float8(16, 32, 8, 16)
+    nz = ft.new_zeros([24, 32])
+    assert nz.qdata.shape == (24, 32)
+    assert nz.scale.shape == (24 // 8, 32 // 16)
+
+
+def test_float8_split_dim_other_than_zero_rejected():
+    """Sharding axis is dim 0 only; a non-zero dim split must raise (documents the restriction)."""
+    ft = _make_blockwise_float8(16, 32, 8, 16)
+    with pytest.raises(NotImplementedError, match="only on dim 0"):
+        torch.split(ft, 16, 1)
diff --git a/tests/monkeypatch/test_fsdp2_quantized.py b/tests/monkeypatch/test_fsdp2_quantized.py
new file mode 100644
index 0000000000..973a2729d3
--- /dev/null
+++ b/tests/monkeypatch/test_fsdp2_quantized.py
@@ -0,0 +1,125 @@
+"""CPU-only tests for the FSDP2 quantized capability helpers (#2)."""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from axolotl.monkeypatch.accelerate import fsdp2_quantized as fq
+
+
+def test_model_has_nonfloat_params():
+    float_only = nn.Linear(4, 4)
+    assert not fq.model_has_nonfloat_params(float_only)
+
+    class Quant(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.w = nn.Parameter(
+                torch.zeros(4, 4, dtype=torch.uint8), requires_grad=False
+            )
+
+    assert fq.model_has_nonfloat_params(Quant())
+
+
+def test_nonfloat_param_guard_restores_on_success():
+    orig = nn.Parameter.__new__
+    model = nn.Linear(2, 2)
+    with fq.nonfloat_param_guard(model):
+        assert nn.Parameter.__new__ is not orig  # patched inside
+    assert nn.Parameter.__new__ is orig  # restored after
+
+
+def test_nonfloat_param_guard_restores_on_exception():
+    orig = nn.Parameter.__new__
+    model = nn.Linear(2, 2)
+    with pytest.raises(RuntimeError, match="boom"):
+        with fq.nonfloat_param_guard(model):
+            assert nn.Parameter.__new__ is not orig
+            raise RuntimeError("boom during fully_shard")
+    # the process-global patch must be restored even though the body raised
+    assert nn.Parameter.__new__ is orig
+
+
+def test_nonfloat_param_guard_defaults_new_nonfloat_to_no_grad():
+    # torch normally forbids constructing a non-float Parameter with the default requires_grad=True;
+    # inside the guard the default flips to False for non-float data, so it succeeds.
+    with pytest.raises(RuntimeError):
+        nn.Parameter(torch.zeros(2, dtype=torch.uint8))  # default True -> torch rejects
+
+    model = nn.Linear(2, 2)
+    with fq.nonfloat_param_guard(model):
+        p = nn.Parameter(
+            torch.zeros(2, dtype=torch.uint8)
+        )  # default True -> guard makes it False
+        assert p.requires_grad is False
+        assert nn.Parameter(torch.zeros(2)).requires_grad is True  # float keeps True
+
+    # after restore, the normal torch behavior returns
+    with pytest.raises(RuntimeError):
+        nn.Parameter(torch.zeros(2, dtype=torch.uint8))
+
+
+def test_nonfloat_param_guard_freezes_existing_nonfloat():
+    class Quant(nn.Module):
+        def __init__(self):
+            super().__init__()
+            # non-float params must be created frozen (torch forbids requires_grad=True here)
+            self.q = nn.Parameter(
+                torch.zeros(2, 2, dtype=torch.uint8), requires_grad=False
+            )
+            self.f = nn.Parameter(torch.zeros(2, 2), requires_grad=True)
+
+    m = Quant()
+    with fq.nonfloat_param_guard(m):
+        assert m.q.requires_grad is False  # non-float stays frozen
+        assert m.f.requires_grad is True  # float untouched
+
+
+def test_register_fp32_shard_classes():
+    saved = set(fq._FP32_SHARD_CLASS_NAMES)
+    try:
+        fq.register_fp32_shard_classes(["FooBarModule"])
+        assert "FooBarModule" in fq._FP32_SHARD_CLASS_NAMES
+    finally:  # restore the global registry so tests stay order-independent
+        fq._FP32_SHARD_CLASS_NAMES.clear()
+        fq._FP32_SHARD_CLASS_NAMES.update(saved)
+
+
+def test_quantized_param_detection_float_logical_subclass():
+    # torchao NVFP4Tensor/Float8Tensor report a logical FLOAT dtype, so the nonfloat check misses
+    # them; the quantized check must still catch them by tensor-subclass name.
+    saved = set(fq._QUANT_TENSOR_CLASS_NAMES)
+    try:
+
+        class FakeNVFP4Tensor(torch.Tensor):
+            pass
+
+        t = torch.zeros(4, 4, dtype=torch.bfloat16).as_subclass(FakeNVFP4Tensor)
+        assert torch.is_floating_point(
+            t
+        )  # float-logical -> invisible to the nonfloat check
+
+        fq.register_quantized_tensor_classes(["FakeNVFP4Tensor"])
+        assert fq._is_quantized_param(t)
+
+        class M(nn.Module):
+            def __init__(self):
+                super().__init__()
+                # torchao wraps the subclass directly in the Parameter (preserves the subclass type)
+                self.w = nn.Parameter(
+                    torch.zeros(4, 4, dtype=torch.bfloat16).as_subclass(
+                        FakeNVFP4Tensor
+                    ),
+                    requires_grad=False,
+                )
+
+        m = M()
+        assert fq.model_has_quantized_params(m)  # detected via the registry
+        assert not fq.model_has_nonfloat_params(m)  # but NOT a plain non-float param
+
+        # built-in torchao names are detected out of the box
+        assert "NVFP4Tensor" in fq._QUANT_TENSOR_CLASS_NAMES
+        assert "Float8Tensor" in fq._QUANT_TENSOR_CLASS_NAMES
+    finally:  # restore the global registry so tests stay order-independent
+        fq._QUANT_TENSOR_CLASS_NAMES.clear()
+        fq._QUANT_TENSOR_CLASS_NAMES.update(saved)
diff --git a/tests/monkeypatch/test_gemma4_fused_attn.py b/tests/monkeypatch/test_gemma4_fused_attn.py
new file mode 100644
index 0000000000..db719c4cc7
--- /dev/null
+++ b/tests/monkeypatch/test_gemma4_fused_attn.py
@@ -0,0 +1,357 @@
+"""Tests for the Gemma 4 fused-attention monkey-patch (hybrid sliding + partial-rotary layers)."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+pytest.importorskip(
+    "transformers.models.gemma4",
+    reason="fused_attn patch only matters when Gemma 4 is available",
+)
+
+
+@pytest.fixture
+def restore_gemma4_attention():
+    """Snapshot ``Gemma4TextAttention.forward`` so the patch can't leak across tests."""
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextAttention
+
+    saved = Gemma4TextAttention.forward
+    yield Gemma4TextAttention
+    Gemma4TextAttention.forward = saved
+
+
+def _build_hybrid_config():
+    """Tiny hybrid Gemma 4: one sliding + one full-attention layer with ``partial_rotary_factor=0.25``."""
+    from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+
+    cfg = Gemma4TextConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=32,
+        global_head_dim=64,
+        layer_types=["sliding_attention", "full_attention"],
+        sliding_window=64,
+        max_position_embeddings=2048,
+        hidden_size_per_layer_input=16,
+        vocab_size_per_layer_input=128,
+        rope_parameters={
+            "sliding_attention": {
+                "rope_type": "default",
+                "rope_theta": 10000.0,
+            },
+            "full_attention": {
+                "rope_type": "proportional",
+                "rope_theta": 1000000.0,
+                "partial_rotary_factor": 0.25,
+            },
+        },
+    )
+    cfg._attn_implementation = "sdpa"
+    return cfg
+
+
+def _build_model(seed=0):
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel
+
+    torch.manual_seed(seed)
+    cfg = _build_hybrid_config()
+    return Gemma4TextModel(cfg).cuda().to(torch.bfloat16).eval()
+
+
+class TestGemma4FusedAttnLoRACompose:
+    """LoRA QKV + fused composition. Gemma 4 in transformers>=5.8.1 has no matching ``QKV_PATCHES`` entry yet, hence the strict xfail below."""
+
+    def _build_cfg(self):
+        from axolotl.utils.dict import DictDefault
+
+        return DictDefault(
+            {
+                "base_model": "fake/gemma4",
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_dropout": 0.0,
+            }
+        )
+
+    @pytest.mark.xfail(
+        reason="Gemma 4 QKV_PATCHES need refresh for transformers 5.8.1",
+        strict=True,
+    )
+    def test_lora_qkv_then_fused_does_not_raise(
+        self, restore_gemma4_attention, monkeypatch
+    ):
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextAttention
+
+        from axolotl.monkeypatch import lora_kernels
+        from axolotl.monkeypatch.models.gemma4.fused_attn import (
+            patch_gemma4_fused_attn,
+        )
+
+        monkeypatch.setattr(
+            lora_kernels,
+            "get_attention_cls_from_config",
+            lambda _cfg: Gemma4TextAttention,
+        )
+
+        try:
+            delattr(Gemma4TextAttention, "_original_forward")
+        except AttributeError:
+            pass
+
+        try:
+            lora_kernels.patch_self_attn_lora(self._build_cfg())
+            assert hasattr(Gemma4TextAttention, "_original_forward"), (
+                "patch_self_attn_lora must run on stock source first"
+            )
+            patch_gemma4_fused_attn()
+        finally:
+            try:
+                delattr(Gemma4TextAttention, "_original_forward")
+            except AttributeError:
+                pass
+
+    def test_reverse_order_skips_lora_rewrite(
+        self, restore_gemma4_attention, monkeypatch, caplog
+    ):
+        """Fused-then-LoRA must NOT install the LoRA source rewrite. Upstream
+        PR #3657 made ``patch_self_attn_lora`` detect ``apply_qkv``/``apply_o``
+        on a fused-patched attention and skip; our ``patch_manager`` reorder
+        keeps this from happening in practice, but the skip path is the last
+        line of defense."""
+        import logging
+
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextAttention
+
+        from axolotl.monkeypatch import lora_kernels
+        from axolotl.monkeypatch.models.gemma4.fused_attn import (
+            patch_gemma4_fused_attn,
+        )
+
+        monkeypatch.setattr(
+            lora_kernels,
+            "get_attention_cls_from_config",
+            lambda _cfg: Gemma4TextAttention,
+        )
+
+        try:
+            delattr(Gemma4TextAttention, "_original_forward")
+        except AttributeError:
+            pass
+
+        try:
+            patch_gemma4_fused_attn()
+            logger = logging.getLogger("axolotl.monkeypatch.lora_kernels")
+            logger.addHandler(caplog.handler)
+            previous_level = logger.level
+            logger.setLevel(logging.INFO)
+            try:
+                lora_kernels.patch_self_attn_lora(self._build_cfg())
+            finally:
+                logger.removeHandler(caplog.handler)
+                logger.setLevel(previous_level)
+            assert "fused attention" in caplog.text and "skipping" in caplog.text, (
+                "expected lora_kernels to detect the fused path and log a skip; "
+                f"got {caplog.text}"
+            )
+            assert not hasattr(Gemma4TextAttention, "_original_forward"), (
+                "lora_kernels installed _original_forward over a fused-patched class"
+            )
+        finally:
+            try:
+                delattr(Gemma4TextAttention, "_original_forward")
+            except AttributeError:
+                pass
+
+
+def _build_kv_shared_config():
+    """Hybrid Gemma 4 with ``num_kv_shared_layers > 0`` so the fused shared-KV branch runs."""
+    from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+
+    cfg = Gemma4TextConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=4,
+        num_kv_shared_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=32,
+        global_head_dim=64,
+        layer_types=[
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "full_attention",
+        ],
+        sliding_window=64,
+        max_position_embeddings=2048,
+        hidden_size_per_layer_input=16,
+        vocab_size_per_layer_input=128,
+        rope_parameters={
+            "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0},
+            "full_attention": {
+                "rope_type": "proportional",
+                "rope_theta": 1000000.0,
+                "partial_rotary_factor": 0.25,
+            },
+        },
+    )
+    cfg._attn_implementation = "sdpa"
+    return cfg
+
+
+class TestFusedAttnSharedKV:
+    """Regression: ``num_kv_shared_layers > 0`` hit the ``kv_shared_layer_index`` key removed in transformers>=5.8."""
+
+    def test_shared_kv_forward_backward(self, restore_gemma4_attention):
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel
+
+        from axolotl.monkeypatch.models.gemma4.fused_attn import (
+            patch_gemma4_fused_attn,
+        )
+
+        torch.manual_seed(4)
+        m = Gemma4TextModel(_build_kv_shared_config()).cuda().to(torch.bfloat16).train()
+        assert any(layer.self_attn.is_kv_shared_layer for layer in m.layers), (
+            "test config must exercise at least one kv-shared layer"
+        )
+
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+
+        with torch.no_grad():
+            ref = m(input_ids=ids, attention_mask=mask).last_hidden_state.clone()
+
+        patch_gemma4_fused_attn()
+        out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+        out.sum().backward()
+
+        assert out.shape == ref.shape
+        assert torch.isfinite(out).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), out.detach().flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"shared-kv fused vs stock cosine_sim={cos_sim:.6f}"
+
+
+class TestFusedAttnSignature:
+    """Pin the fused forward against the live ``Gemma4TextDecoderLayer.forward`` call shape."""
+
+    def test_decoder_layer_can_call_fused_forward(self, restore_gemma4_attention):
+        from axolotl.monkeypatch.models.gemma4.fused_attn import (
+            patch_gemma4_fused_attn,
+        )
+
+        m = _build_model()
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+
+        patch_gemma4_fused_attn()
+        with torch.no_grad():
+            out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+
+        assert out.shape == (2, 16, 64)
+        assert torch.isfinite(out).all()
+
+
+class TestFusedAttnPerLayerCorrectness:
+    """Single-layer comparison of patched vs stock attention to isolate kernel correctness from cross-layer drift."""
+
+    def _run_attention(self, model, layer_idx, hidden_states, position_ids):
+        attn = model.layers[layer_idx].self_attn
+        layer_type = model.config.layer_types[layer_idx]
+        cos, sin = model.rotary_emb(hidden_states, position_ids, layer_type)
+        out, _ = attn(
+            hidden_states=hidden_states,
+            position_embeddings=(cos, sin),
+            attention_mask=None,
+            shared_kv_states={},
+        )
+        return out
+
+    @pytest.mark.parametrize(
+        "layer_idx",
+        [0, 1],
+        ids=["sliding_head32", "global_head64_proportional"],
+    )
+    def test_forward_matches_stock(self, restore_gemma4_attention, layer_idx):
+        from axolotl.monkeypatch.models.gemma4.fused_attn import (
+            patch_gemma4_fused_attn,
+        )
+
+        m = _build_model(seed=1)
+        hs = torch.randn(2, 16, 64, device="cuda", dtype=torch.bfloat16)
+        pos = torch.arange(16, device="cuda").unsqueeze(0).expand(2, -1)
+
+        with torch.no_grad():
+            ref = self._run_attention(m, layer_idx, hs, pos)
+
+        patch_gemma4_fused_attn()
+        with torch.no_grad():
+            got = self._run_attention(m, layer_idx, hs, pos)
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, (
+            f"layer {layer_idx} fused vs stock cosine_sim={cos_sim:.6f}"
+        )
+        torch.testing.assert_close(got, ref, rtol=5e-2, atol=5e-2)
+
+
+class TestFusedAttnFullModel:
+    """End-to-end forward + backward through both layer types."""
+
+    def test_full_forward_matches_stock(self, restore_gemma4_attention):
+        from axolotl.monkeypatch.models.gemma4.fused_attn import (
+            patch_gemma4_fused_attn,
+        )
+
+        m = _build_model(seed=2)
+        ids = torch.randint(0, 128, (2, 32), device="cuda")
+        mask = torch.ones(2, 32, dtype=torch.long, device="cuda")
+
+        with torch.no_grad():
+            ref = m(input_ids=ids, attention_mask=mask).last_hidden_state.clone()
+
+        patch_gemma4_fused_attn()
+        with torch.no_grad():
+            got = m(input_ids=ids, attention_mask=mask).last_hidden_state.clone()
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"end-to-end cosine_sim={cos_sim:.6f}"
+
+    def test_backward_grad_flows_through_fused_path(self, restore_gemma4_attention):
+        from axolotl.monkeypatch.models.gemma4.fused_attn import (
+            patch_gemma4_fused_attn,
+        )
+
+        m = _build_model(seed=3).train()
+        patch_gemma4_fused_attn()
+
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+        out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+        out.sum().backward()
+
+        for i, layer in enumerate(m.layers[:2]):
+            attn = layer.self_attn
+            assert attn.q_norm.weight.grad is not None, f"layer {i} q_norm no grad"
+            assert attn.k_norm.weight.grad is not None, f"layer {i} k_norm no grad"
+            assert attn.q_norm.weight.grad.isfinite().all()
+            assert attn.k_norm.weight.grad.isfinite().all()
+            assert attn.q_norm.weight.grad.abs().sum() > 0
+            assert attn.k_norm.weight.grad.abs().sum() > 0
diff --git a/tests/monkeypatch/test_gemma4_fused_attn_patch.py b/tests/monkeypatch/test_gemma4_fused_attn_patch.py
new file mode 100644
index 0000000000..75dbe472bc
--- /dev/null
+++ b/tests/monkeypatch/test_gemma4_fused_attn_patch.py
@@ -0,0 +1,171 @@
+"""Unit tests for the Gemma4 fused-attention shared_kv_states routing patch."""
+
+import pytest
+
+gemma4_modeling = pytest.importorskip("transformers.models.gemma4.modeling_gemma4")
+
+
+@pytest.fixture
+def clean_decoder_layer_patch_slate():
+    """Save and restore Gemma4TextDecoderLayer.__call__ and the sentinel."""
+    from axolotl.monkeypatch.models.gemma4 import fused_attn
+
+    cls = gemma4_modeling.Gemma4TextDecoderLayer
+    original_call = cls.__call__
+    had_sentinel = getattr(cls, "_axolotl_shared_kv_patched", False)
+
+    if had_sentinel:
+        del cls._axolotl_shared_kv_patched
+
+    try:
+        yield cls, fused_attn
+    finally:
+        cls.__call__ = original_call
+        if had_sentinel:
+            cls._axolotl_shared_kv_patched = True
+        elif hasattr(cls, "_axolotl_shared_kv_patched"):
+            del cls._axolotl_shared_kv_patched
+        fused_attn._set_shared_kv_states(None)
+
+
+class TestPatchedDecoderLayerCall:
+    def test_pops_shared_kv_states_and_populates_store(
+        self, clean_decoder_layer_patch_slate
+    ):
+        cls, fused_attn = clean_decoder_layer_patch_slate
+
+        captured = {}
+
+        def spy(self, *args, **kwargs):
+            captured["args"] = args
+            captured["kwargs"] = dict(kwargs)
+            return "spy_return"
+
+        cls.__call__ = spy
+        fused_attn._patch_decoder_layer_call()
+
+        assert getattr(cls, "_axolotl_shared_kv_patched", False) is True
+        assert cls.__call__ is not spy
+
+        shared_kv = {"layer_0": ("k", "v")}
+        result = cls.__call__(
+            object(),
+            "positional_arg",
+            shared_kv_states=shared_kv,
+            other_kwarg="keep_me",
+        )
+
+        assert result == "spy_return"
+        assert captured["args"] == ("positional_arg",)
+        assert "shared_kv_states" not in captured["kwargs"]
+        assert captured["kwargs"] == {"other_kwarg": "keep_me"}
+        assert fused_attn._get_shared_kv_states() is shared_kv
+
+    def test_clears_store_when_kwarg_absent(self, clean_decoder_layer_patch_slate):
+        """Regression for commit 251021e1: a prior step's dict must not leak
+        into a later call that omits `shared_kv_states`."""
+        cls, fused_attn = clean_decoder_layer_patch_slate
+
+        def spy(self, *args, **kwargs):
+            return None
+
+        cls.__call__ = spy
+        fused_attn._patch_decoder_layer_call()
+
+        stale = {"stale_step": True}
+        fused_attn._set_shared_kv_states(stale)
+        assert fused_attn._get_shared_kv_states() is stale
+
+        cls.__call__(object())
+
+        assert fused_attn._get_shared_kv_states() is None
+
+    def test_store_visible_across_threads(self):
+        """Regression for commit e3669b2c: the store must be readable from
+        threads other than the one that set it. `threading.local()` failed
+        this invariant, crashing with 'NoneType' object is not subscriptable'
+        on MoE Gemma4 variants when autograd worker threads ran backward
+        recompute under HF-Trainer gradient_checkpointing."""
+        import threading
+
+        from axolotl.monkeypatch.models.gemma4 import fused_attn
+
+        sentinel = {"layer_0": ("k", "v")}
+        try:
+            fused_attn._set_shared_kv_states(sentinel)
+
+            seen = {}
+
+            def worker():
+                seen["value"] = fused_attn._get_shared_kv_states()
+
+            t = threading.Thread(target=worker)
+            t.start()
+            t.join()
+
+            assert seen["value"] is sentinel
+        finally:
+            fused_attn._set_shared_kv_states(None)
+
+
+@pytest.fixture
+def clean_entry_point_patch_slate():
+    """Save and restore Gemma4TextAttention.forward and Gemma4TextDecoderLayer.__call__."""
+    from axolotl.monkeypatch.models.gemma4 import fused_attn
+
+    decoder_cls = gemma4_modeling.Gemma4TextDecoderLayer
+    attn_cls = gemma4_modeling.Gemma4TextAttention
+
+    original_call = decoder_cls.__call__
+    original_forward = attn_cls.forward
+    had_sentinel = getattr(decoder_cls, "_axolotl_shared_kv_patched", False)
+
+    if had_sentinel:
+        del decoder_cls._axolotl_shared_kv_patched
+
+    try:
+        yield decoder_cls, attn_cls, original_call, original_forward, fused_attn
+    finally:
+        decoder_cls.__call__ = original_call
+        attn_cls.forward = original_forward
+        if had_sentinel:
+            decoder_cls._axolotl_shared_kv_patched = True
+        elif hasattr(decoder_cls, "_axolotl_shared_kv_patched"):
+            del decoder_cls._axolotl_shared_kv_patched
+        fused_attn._set_shared_kv_states(None)
+
+
+class TestPatchGemma4FusedAttnEntryPoint:
+    def test_default_flag_swaps_only_attention_forward(
+        self, clean_entry_point_patch_slate
+    ):
+        (
+            decoder_cls,
+            attn_cls,
+            original_call,
+            original_forward,
+            fused_attn,
+        ) = clean_entry_point_patch_slate
+
+        fused_attn.patch_gemma4_fused_attn()
+
+        assert attn_cls.forward is not original_forward
+        assert decoder_cls.__call__ is original_call
+        assert not getattr(decoder_cls, "_axolotl_shared_kv_patched", False)
+
+    def test_workaround_flag_installs_decoder_layer_patch(
+        self, clean_entry_point_patch_slate
+    ):
+        (
+            decoder_cls,
+            attn_cls,
+            original_call,
+            original_forward,
+            fused_attn,
+        ) = clean_entry_point_patch_slate
+
+        fused_attn.patch_gemma4_fused_attn(install_shared_kv_workaround=True)
+
+        assert attn_cls.forward is not original_forward
+        assert decoder_cls.__call__ is not original_call
+        assert getattr(decoder_cls, "_axolotl_shared_kv_patched", False) is True
diff --git a/tests/monkeypatch/test_gemma4_hybrid_mask.py b/tests/monkeypatch/test_gemma4_hybrid_mask.py
new file mode 100644
index 0000000000..471bc0cf9f
--- /dev/null
+++ b/tests/monkeypatch/test_gemma4_hybrid_mask.py
@@ -0,0 +1,410 @@
+"""Tests for the Gemma 4 hybrid-attention mask fix.
+
+These tests pin the single critical behavior: after installing the patch,
+``modeling_gemma4.create_causal_mask`` passes an SDPA-overridden config to
+the underlying mask builder regardless of what the caller's config says.
+This is what keeps full-attention (head_dim=512) global layers from
+crashing at long sequence lengths — they need a 4D SDPA-format mask, not
+the 2D FA2 mask that would be built from the model-level config.
+
+The tests use a mocked ``create_causal_mask`` so they don't have to load
+a real 26B Gemma 4 model or even have access to its weights. What matters
+for the bug fix is which config is handed to the mask factory, not the
+factory's actual output.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+pytest.importorskip(
+    "transformers.models.gemma4",
+    reason="gemma4_hybrid_mask patch only matters when Gemma 4 is available",
+)
+
+
+@pytest.fixture
+def restore_gemma4_module():
+    """Snapshot ``modeling_gemma4.create_causal_mask`` and restore after
+    each test so patch state doesn't leak across the suite."""
+    from transformers.models.gemma4 import modeling_gemma4
+
+    saved = modeling_gemma4.create_causal_mask
+    yield modeling_gemma4
+    modeling_gemma4.create_causal_mask = saved
+    # Reset the module-level flag so the next test can re-install cleanly.
+    from axolotl.monkeypatch import gemma4_hybrid_mask
+
+    gemma4_hybrid_mask._PATCH_APPLIED = False
+
+
+def test_patch_replaces_create_causal_mask(restore_gemma4_module):
+    modeling_gemma4 = restore_gemma4_module
+    from axolotl.monkeypatch.gemma4_hybrid_mask import patch_gemma4_hybrid_mask
+
+    original = modeling_gemma4.create_causal_mask
+    assert patch_gemma4_hybrid_mask() is True
+
+    assert modeling_gemma4.create_causal_mask is not original
+    assert modeling_gemma4.create_causal_mask._axolotl_original is original, (
+        "patched wrapper must expose the original reference for teardown"
+    )
+
+
+def test_patch_is_idempotent(restore_gemma4_module):
+    modeling_gemma4 = restore_gemma4_module
+    from axolotl.monkeypatch.gemma4_hybrid_mask import patch_gemma4_hybrid_mask
+
+    patch_gemma4_hybrid_mask()
+    wrapper_first = modeling_gemma4.create_causal_mask
+
+    # Second call must not re-wrap the already-wrapped function (which
+    # would leak the original reference through a chain of wrappers).
+    patch_gemma4_hybrid_mask()
+    wrapper_second = modeling_gemma4.create_causal_mask
+
+    assert wrapper_first is wrapper_second
+
+
+def test_patched_mask_forces_sdpa_config(restore_gemma4_module):
+    """Core invariant: when the patched wrapper is called with a config
+    that says ``flash_attention_2``, the underlying mask factory receives
+    a shallow-copied config whose ``_attn_implementation`` is ``"sdpa"``.
+
+    Without this, the full-attention global layers get a 2D FA2 mask and
+    crash at long seq lens with the [B, H, S, S] / [B, S] expand error.
+    """
+    modeling_gemma4 = restore_gemma4_module
+    from axolotl.monkeypatch.gemma4_hybrid_mask import patch_gemma4_hybrid_mask
+
+    # Swap in a mock BEFORE installing the patch so the wrapper captures
+    # it as the "original". The mock records every call so we can inspect
+    # what config got passed through.
+    mock_factory = MagicMock(name="create_causal_mask", return_value="mask_4d")
+    modeling_gemma4.create_causal_mask = mock_factory
+    patch_gemma4_hybrid_mask()
+
+    # Caller-supplied config says FA2 (that's the model-level setting).
+    caller_config = SimpleNamespace(
+        _attn_implementation="flash_attention_2",
+        head_dim=512,
+        some_other_attr="preserved",
+    )
+    result = modeling_gemma4.create_causal_mask(
+        caller_config,
+        inputs_embeds=None,
+        attention_mask=None,
+        past_key_values=None,
+        position_ids=None,
+    )
+
+    # Wrapper returned whatever the mock returned — no transformation of
+    # the result itself.
+    assert result == "mask_4d"
+
+    # The mock was called exactly once with a config whose
+    # ``_attn_implementation`` is sdpa, NOT the caller's fa2.
+    assert mock_factory.call_count == 1
+    (passed_config, *_), passed_kwargs = mock_factory.call_args
+    assert passed_config._attn_implementation == "sdpa"
+
+    # The wrapper must NOT mutate the caller's config in place — other
+    # mask builders (e.g. create_sliding_window_causal_mask) read from
+    # the same config and must still see fa2.
+    assert caller_config._attn_implementation == "flash_attention_2"
+
+    # Other attributes on the config must be preserved so the underlying
+    # factory has everything it needs (head_dim, rope_theta, vocab_size, ...).
+    assert passed_config.head_dim == 512
+    assert passed_config.some_other_attr == "preserved"
+
+
+def test_patched_wrapper_passes_through_all_kwargs(restore_gemma4_module):
+    """The wrapper must forward positional + keyword args to the original
+    unchanged, so transformers' own call-site in Gemma4TextModel.forward
+    keeps working across minor transformers-version signature drift."""
+    modeling_gemma4 = restore_gemma4_module
+    from axolotl.monkeypatch.gemma4_hybrid_mask import patch_gemma4_hybrid_mask
+
+    mock_factory = MagicMock(return_value="mask")
+    modeling_gemma4.create_causal_mask = mock_factory
+    patch_gemma4_hybrid_mask()
+
+    caller_config = SimpleNamespace(_attn_implementation="flash_attention_2")
+    modeling_gemma4.create_causal_mask(
+        caller_config,
+        "positional_arg",
+        inputs_embeds="embeds",
+        attention_mask="mask_2d",
+        past_key_values="cache",
+        position_ids="positions",
+        or_mask_function="or_fn",
+    )
+
+    args, kwargs = mock_factory.call_args
+    # First positional (after config override) is preserved.
+    assert args[1] == "positional_arg"
+    # All kwargs are forwarded untouched.
+    assert kwargs["inputs_embeds"] == "embeds"
+    assert kwargs["attention_mask"] == "mask_2d"
+    assert kwargs["past_key_values"] == "cache"
+    assert kwargs["position_ids"] == "positions"
+    assert kwargs["or_mask_function"] == "or_fn"
+
+
+def test_unpatch_restores_original(restore_gemma4_module):
+    modeling_gemma4 = restore_gemma4_module
+    from axolotl.monkeypatch.gemma4_hybrid_mask import (
+        patch_gemma4_hybrid_mask,
+        unpatch_gemma4_hybrid_mask,
+    )
+
+    sentinel = MagicMock(name="original")
+    modeling_gemma4.create_causal_mask = sentinel
+    patch_gemma4_hybrid_mask()
+    assert modeling_gemma4.create_causal_mask is not sentinel
+
+    unpatch_gemma4_hybrid_mask()
+    assert modeling_gemma4.create_causal_mask is sentinel
+
+
+def test_unpatch_is_safe_without_prior_patch(restore_gemma4_module):
+    from axolotl.monkeypatch.gemma4_hybrid_mask import unpatch_gemma4_hybrid_mask
+
+    # Should be a no-op, no exception.
+    unpatch_gemma4_hybrid_mask()
+
+
+def test_sliding_window_mask_builder_is_not_patched(restore_gemma4_module):
+    """Only ``create_causal_mask`` is overridden — the sliding-window
+    factory must remain bound to its original to preserve FA2 masks for
+    the sliding-attention layers. If we accidentally patch both, the
+    sliding layers get SDPA format and lose the FA2 speedup."""
+    modeling_gemma4 = restore_gemma4_module
+    from axolotl.monkeypatch.gemma4_hybrid_mask import patch_gemma4_hybrid_mask
+
+    if not hasattr(modeling_gemma4, "create_sliding_window_causal_mask"):
+        pytest.skip("transformers version has no create_sliding_window_causal_mask")
+
+    sliding_before = modeling_gemma4.create_sliding_window_causal_mask
+    patch_gemma4_hybrid_mask()
+    sliding_after = modeling_gemma4.create_sliding_window_causal_mask
+    assert sliding_after is sliding_before
+
+
+# ---------------------------------------------------------------------------
+# Integration tests with a tiny randomly-initialized Gemma4TextModel.
+#
+# These do NOT load real 26B weights. They build a ~350k-param Gemma 4 text
+# model with 2 layers (one sliding, one full_attention), apply the hybrid
+# attention path end-to-end, and run a forward pass with a padded
+# attention_mask at a long-ish seq len. The invariant we're pinning is that
+# the full_attention layer does not crash with the
+#   "Target sizes: [B, H, S, S]. Tensor sizes: [B, S]"
+# error — the exact failure that blew up the Gemma 4 MoE 26B pilot at ~7k
+# tokens in the FSDP2 training run.
+# ---------------------------------------------------------------------------
+
+
+def _build_tiny_gemma4_text_model():
+    """Return a tiny randomly-initialized Gemma4TextModel with mixed layers."""
+    import torch
+    from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel
+
+    cfg = Gemma4TextConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=32,
+        layer_types=["sliding_attention", "full_attention"],
+        sliding_window=64,
+        max_position_embeddings=2048,
+        hidden_size_per_layer_input=16,
+        vocab_size_per_layer_input=128,
+    )
+    # Caller-supplied attn impl simulates the pilot config (fa2 at model
+    # level). The hybrid patch is what makes this survive long context.
+    cfg._attn_implementation = "sdpa"  # start safe; the test toggles fa2 later
+    torch.manual_seed(42)
+    model = Gemma4TextModel(cfg).eval()
+    return model, cfg
+
+
+def _apply_hybrid_attn_inline(model, cfg):
+    """Replicate what ``patch_manager._apply_gemma_hybrid_attention`` does
+    to a model, without needing a full PatchManager / pydantic cfg."""
+    import copy
+
+    from axolotl.monkeypatch.gemma4_hybrid_mask import patch_gemma4_hybrid_mask
+
+    for layer_idx, layer in enumerate(model.layers):
+        if cfg.layer_types[layer_idx] != "sliding_attention":
+            attn = getattr(layer, "self_attn", None)
+            if attn is not None and hasattr(attn, "config"):
+                sdpa_cfg = copy.copy(attn.config)
+                sdpa_cfg._attn_implementation = "sdpa"
+                attn.config = sdpa_cfg
+    patch_gemma4_hybrid_mask()
+
+
+def test_tiny_gemma4_long_context_forward_does_not_crash(restore_gemma4_module):
+    """End-to-end invariant: with the hybrid attn patch applied, a tiny
+    Gemma4TextModel runs a forward at long context (1024 tokens) with
+    real padding in the attention mask, producing the expected output
+    shape. This exercises the actual code path that crashed the pilot
+    without needing a real 26B checkpoint or CUDA."""
+    import torch
+
+    model, cfg = _build_tiny_gemma4_text_model()
+    _apply_hybrid_attn_inline(model, cfg)
+
+    B, S = 2, 1024
+    input_ids = torch.randint(0, cfg.vocab_size, (B, S))
+    attn_mask = torch.ones(B, S, dtype=torch.long)
+    # Pad positions in the second row. Without padding, SDPA falls back to
+    # ``is_causal=True`` with ``mask=None`` — we need a materialized 4D
+    # mask to exercise the actual bug site.
+    attn_mask[1, S // 2 :] = 0
+
+    with torch.no_grad():
+        out = model(input_ids=input_ids, attention_mask=attn_mask)
+
+    assert out.last_hidden_state.shape == (B, S, cfg.hidden_size)
+    assert torch.isfinite(out.last_hidden_state).all()
+
+
+def test_patched_create_causal_mask_returns_4d_for_real_config(
+    restore_gemma4_module,
+):
+    """Hit the REAL ``create_causal_mask`` (not a mock) via the wrapper
+    and verify the returned mask is a 4D tensor — which is the shape the
+    SDPA-patched global layers need. Without the patch and with a
+    caller-supplied FA2 config this would return a 2D mask and the layer
+    would crash at long context."""
+    import torch
+    from transformers.cache_utils import DynamicCache
+    from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+
+    from axolotl.monkeypatch.gemma4_hybrid_mask import patch_gemma4_hybrid_mask
+
+    patch_gemma4_hybrid_mask()
+    modeling_gemma4 = restore_gemma4_module
+
+    cfg = Gemma4TextConfig(
+        vocab_size=128,
+        hidden_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=32,
+        layer_types=["sliding_attention", "full_attention"],
+        sliding_window=64,
+        max_position_embeddings=2048,
+        hidden_size_per_layer_input=16,
+        vocab_size_per_layer_input=128,
+    )
+    # Simulate the pilot: caller says flash_attention_2, but global layers
+    # were switched to SDPA per-layer. Without the patch, create_causal_mask
+    # would return an FA2 2D mask here and the SDPA layer would crash.
+    cfg._attn_implementation = "flash_attention_2"
+
+    B, S = 2, 1024
+    inputs_embeds = torch.zeros((B, S, cfg.hidden_size), dtype=torch.float32)
+    attention_mask = torch.ones((B, S), dtype=torch.long)
+    attention_mask[1, S // 2 :] = 0  # force the 4D materialized path
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, -1)
+    past_key_values = DynamicCache(config=cfg)
+
+    mask = modeling_gemma4.create_causal_mask(
+        config=cfg,
+        inputs_embeds=inputs_embeds,
+        attention_mask=attention_mask,
+        past_key_values=past_key_values,
+        position_ids=position_ids,
+    )
+
+    assert mask is not None
+    assert isinstance(mask, torch.Tensor)
+    assert mask.dim() == 4, (
+        f"expected a 4D SDPA-format mask, got {mask.dim()}D "
+        f"shape={tuple(mask.shape)}. The full_attention global layers need "
+        "this shape or they crash at long context."
+    )
+    assert mask.shape[0] == B
+    assert mask.shape[-1] == S
+    assert mask.shape[-2] == S
+
+    # Caller's config must be untouched — other code paths still read it.
+    assert cfg._attn_implementation == "flash_attention_2"
+
+
+@pytest.fixture
+def restore_gemma4_hybrid_mask_both():
+    """Snapshot ``create_causal_mask`` in BOTH the gemma4 and gemma4_unified
+    namespaces (``patch_gemma4_hybrid_mask`` wraps each independently) and reset
+    the module flag so the patch re-installs cleanly."""
+    modeling_unified = pytest.importorskip(
+        "transformers.models.gemma4_unified.modeling_gemma4_unified",
+        reason="unified namespace coverage requires gemma4_unified",
+    )
+    from transformers.models.gemma4 import modeling_gemma4
+
+    from axolotl.monkeypatch import gemma4_hybrid_mask
+
+    saved_gemma4 = modeling_gemma4.create_causal_mask
+    saved_unified = modeling_unified.create_causal_mask
+    gemma4_hybrid_mask._PATCH_APPLIED = False
+    try:
+        yield modeling_unified
+    finally:
+        modeling_gemma4.create_causal_mask = saved_gemma4
+        modeling_unified.create_causal_mask = saved_unified
+        gemma4_hybrid_mask._PATCH_APPLIED = False
+
+
+def test_patch_covers_unified_namespace(restore_gemma4_hybrid_mask_both):
+    """The unified backbone redefines ``create_causal_mask`` in its own module,
+    so the patch must wrap it too — not just ``modeling_gemma4``."""
+    modeling_unified = restore_gemma4_hybrid_mask_both
+    from axolotl.monkeypatch.gemma4_hybrid_mask import patch_gemma4_hybrid_mask
+
+    original = modeling_unified.create_causal_mask
+    assert patch_gemma4_hybrid_mask() is True
+
+    assert modeling_unified.create_causal_mask is not original
+    assert modeling_unified.create_causal_mask._axolotl_original is original
+
+
+def test_gqa_guard_accepts_both_transformers_signatures():
+    """transformers 5.13 added a `value` positional to use_gqa_in_sdpa."""
+    import torch
+    import transformers.integrations.sdpa_attention as sdpa_mod
+
+    from axolotl.monkeypatch.gemma4_hybrid_mask import _patch_use_gqa_head_dim_guard
+
+    original = sdpa_mod.use_gqa_in_sdpa
+    unwrapped = getattr(original, "_axolotl_original", original)
+    try:
+        assert _patch_use_gqa_head_dim_guard()
+        guarded = sdpa_mod.use_gqa_in_sdpa
+        key_small = torch.zeros(1, 2, 4, 128)
+        key_large = torch.zeros(1, 2, 4, 512)
+
+        # the head_dim guard must reject under BOTH call conventions without
+        # touching the wrapped original
+        assert guarded(None, key_large) is False
+        assert guarded(None, key_large, key_large) is False
+
+        # pass-through must forward whatever the installed transformers accepts
+        import inspect
+
+        n_params = len(inspect.signature(unwrapped).parameters)
+        args = (None, key_small, key_small)[:n_params]
+        assert isinstance(guarded(*args), bool)
+    finally:
+        sdpa_mod.use_gqa_in_sdpa = unwrapped
diff --git a/tests/monkeypatch/test_gemma4_unified_fused_attn.py b/tests/monkeypatch/test_gemma4_unified_fused_attn.py
new file mode 100644
index 0000000000..a50b03328b
--- /dev/null
+++ b/tests/monkeypatch/test_gemma4_unified_fused_attn.py
@@ -0,0 +1,221 @@
+"""Tests for the Gemma4 *Unified* fused-attention monkeypatch.
+
+These cover only behavior that diverges from standard gemma4 (whose fused-attn
+patch is already covered by ``test_gemma4_fused_attn.py`` /
+``test_gemma4_fused_attn_patch.py``):
+
+* the idempotency guard on ``patch_gemma4_unified_fused_attn`` (standard gemma4
+  has none),
+* that the side channel routes through the *unified* store/decoder-layer class
+  (a separate object that standard-gemma4 tests never touch), and
+* shared-KV forward/backward under the unified backbone's hardcoded
+  ``layer_type`` keying (vs the version-tolerant helpers in standard gemma4).
+"""
+
+import pytest
+import torch
+
+gemma4_unified_modeling = pytest.importorskip(
+    "transformers.models.gemma4_unified.modeling_gemma4_unified",
+    reason="unified fused_attn patch only matters when gemma4_unified is available",
+)
+
+
+@pytest.fixture
+def clean_unified_slate():
+    """Snapshot/restore ``Gemma4UnifiedTextAttention.forward``,
+    ``Gemma4UnifiedTextDecoderLayer.__call__`` and both axolotl sentinels so the
+    patch (and its idempotency flag) can't leak across tests."""
+    from axolotl.monkeypatch.models.gemma4_unified import fused_attn
+
+    attn_cls = gemma4_unified_modeling.Gemma4UnifiedTextAttention
+    decoder_cls = gemma4_unified_modeling.Gemma4UnifiedTextDecoderLayer
+    original_forward = attn_cls.forward
+    original_call = decoder_cls.__call__
+    had_attn_flag = getattr(attn_cls, "_axolotl_fused_attn_patched", False)
+    had_dec_flag = getattr(decoder_cls, "_axolotl_shared_kv_patched", False)
+
+    if had_attn_flag:
+        del attn_cls._axolotl_fused_attn_patched
+    if had_dec_flag:
+        del decoder_cls._axolotl_shared_kv_patched
+
+    try:
+        yield attn_cls, decoder_cls, original_forward, original_call, fused_attn
+    finally:
+        attn_cls.forward = original_forward
+        decoder_cls.__call__ = original_call
+        for cls, attr, had in (
+            (attn_cls, "_axolotl_fused_attn_patched", had_attn_flag),
+            (decoder_cls, "_axolotl_shared_kv_patched", had_dec_flag),
+        ):
+            if had:
+                setattr(cls, attr, True)
+            elif hasattr(cls, attr):
+                delattr(cls, attr)
+        fused_attn._set_shared_kv_states(None)
+
+
+class TestUnifiedFusedAttnEntryPoint:
+    def test_patch_is_idempotent(self, clean_unified_slate):
+        """Unified guards re-patching via ``_axolotl_fused_attn_patched`` — a
+        second call must not re-wrap forward (standard gemma4 has no such guard)."""
+        pytest.importorskip("triton")
+        attn_cls, _decoder_cls, _orig_fwd, _orig_call, fused_attn = clean_unified_slate
+
+        fused_attn.patch_gemma4_unified_fused_attn()
+        wrapped = attn_cls.forward
+        assert getattr(attn_cls, "_axolotl_fused_attn_patched", False) is True
+
+        fused_attn.patch_gemma4_unified_fused_attn()
+        assert attn_cls.forward is wrapped, "second patch call must be a no-op"
+
+    def test_default_flag_swaps_only_attention_forward(self, clean_unified_slate):
+        pytest.importorskip("triton")
+        attn_cls, decoder_cls, original_forward, original_call, fused_attn = (
+            clean_unified_slate
+        )
+
+        fused_attn.patch_gemma4_unified_fused_attn()
+
+        assert attn_cls.forward is not original_forward
+        assert decoder_cls.__call__ is original_call
+        assert not getattr(decoder_cls, "_axolotl_shared_kv_patched", False)
+
+    def test_workaround_flag_installs_decoder_layer_patch(self, clean_unified_slate):
+        pytest.importorskip("triton")
+        attn_cls, decoder_cls, original_forward, original_call, fused_attn = (
+            clean_unified_slate
+        )
+
+        fused_attn.patch_gemma4_unified_fused_attn(install_shared_kv_workaround=True)
+
+        assert attn_cls.forward is not original_forward
+        assert decoder_cls.__call__ is not original_call
+        assert getattr(decoder_cls, "_axolotl_shared_kv_patched", False) is True
+
+
+class TestUnifiedDecoderLayerCall:
+    """The PR#3611 side channel uses a *separate* store + decoder-layer class
+    from standard gemma4; a copy-paste regression (writing to gemma4's store, or
+    failing to clear) would be caught by nothing in the standard-gemma4 tests."""
+
+    def test_store_is_distinct_from_standard_gemma4(self):
+        from axolotl.monkeypatch.models.gemma4 import fused_attn as regular
+        from axolotl.monkeypatch.models.gemma4_unified import fused_attn as unified
+
+        assert (
+            unified._GEMMA4_UNIFIED_SHARED_KV_STORE
+            is not regular._GEMMA4_SHARED_KV_STORE
+        )
+
+    def test_pops_and_clears_shared_kv_states(self, clean_unified_slate):
+        _attn_cls, decoder_cls, _orig_fwd, _orig_call, fused_attn = clean_unified_slate
+
+        captured = {}
+
+        def spy(self, *args, **kwargs):
+            captured["args"] = args
+            captured["kwargs"] = dict(kwargs)
+            return "spy_return"
+
+        decoder_cls.__call__ = spy
+        fused_attn._patch_decoder_layer_call()
+
+        assert getattr(decoder_cls, "_axolotl_shared_kv_patched", False) is True
+
+        shared_kv = {"full_attention": ("k", "v")}
+        result = decoder_cls.__call__(
+            object(),
+            "positional_arg",
+            shared_kv_states=shared_kv,
+            other_kwarg="keep_me",
+        )
+
+        assert result == "spy_return"
+        assert captured["args"] == ("positional_arg",)
+        assert "shared_kv_states" not in captured["kwargs"]
+        assert captured["kwargs"] == {"other_kwarg": "keep_me"}
+        assert fused_attn._get_shared_kv_states() is shared_kv
+
+        # A later call without the kwarg must clear the store (no cross-step leak).
+        decoder_cls.__call__(object())
+        assert fused_attn._get_shared_kv_states() is None
+
+
+def _build_unified_kv_shared_config():
+    """Tiny unified text config with ``num_kv_shared_layers > 0`` so the fused
+    shared-KV branch (hardcoded ``self.layer_type`` keying) runs."""
+    from transformers.models.gemma4_unified.configuration_gemma4_unified import (
+        Gemma4UnifiedTextConfig,
+    )
+
+    cfg = Gemma4UnifiedTextConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=4,
+        num_kv_shared_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=32,
+        global_head_dim=64,
+        layer_types=[
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "full_attention",
+        ],
+        sliding_window=64,
+        max_position_embeddings=2048,
+        rope_parameters={
+            "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0},
+            "full_attention": {
+                "rope_type": "proportional",
+                "rope_theta": 1000000.0,
+                "partial_rotary_factor": 0.25,
+            },
+        },
+    )
+    cfg._attn_implementation = "sdpa"
+    return cfg
+
+
+class TestUnifiedFusedAttnSharedKV:
+    """Guards the unified backbone's hardcoded ``layer_type`` shared-KV keying
+    (standard gemma4 uses version-tolerant ``_shared_kv_*_key`` helpers)."""
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+    def test_shared_kv_forward_backward(self, clean_unified_slate):
+        _attn_cls, _decoder_cls, _orig_fwd, _orig_call, fused_attn = clean_unified_slate
+        from transformers.models.gemma4_unified.modeling_gemma4_unified import (
+            Gemma4UnifiedTextModel,
+        )
+
+        torch.manual_seed(4)
+        m = (
+            Gemma4UnifiedTextModel(_build_unified_kv_shared_config())
+            .cuda()
+            .to(torch.bfloat16)
+            .train()
+        )
+        assert any(layer.self_attn.is_kv_shared_layer for layer in m.layers), (
+            "test config must exercise at least one kv-shared layer"
+        )
+
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+
+        with torch.no_grad():
+            ref = m(input_ids=ids, attention_mask=mask).last_hidden_state.clone()
+
+        fused_attn.patch_gemma4_unified_fused_attn()
+        out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+        out.sum().backward()
+
+        assert out.shape == ref.shape
+        assert torch.isfinite(out).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), out.detach().flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"shared-kv fused vs stock cosine_sim={cos_sim:.6f}"
diff --git a/tests/monkeypatch/test_kernelize_fixes.py b/tests/monkeypatch/test_kernelize_fixes.py
new file mode 100644
index 0000000000..2031cb6f64
--- /dev/null
+++ b/tests/monkeypatch/test_kernelize_fixes.py
@@ -0,0 +1,247 @@
+"""Tests for the generic ``kernelize()`` repairs in ``kernelize_fixes``.
+
+Two upstream defects are covered (see the patch module docstring): bare
+functions stashed in ``_hidden_kernels`` (gemma4 and ~30 others) and the
+gpt-oss rotary ``Func`` whose ``position_ids`` parameter fails the kernels
+library's signature check against the hub kernel.
+"""
+
+import inspect
+
+import pytest
+import transformers
+from packaging.version import Version
+from transformers.modeling_utils import PreTrainedModel
+
+pytest.importorskip("kernels", reason="kernelize fixes only matter with kernels")
+
+# #46520 drops PreTrainedModel.kernelize; skip dev builds, fail a stable release.
+if not hasattr(PreTrainedModel, "kernelize"):
+    if Version(transformers.__version__).is_prerelease:
+        pytest.skip(
+            "PreTrainedModel.kernelize removed on transformers main (#46520); patch no-ops",
+            allow_module_level=True,
+        )
+    pytest.fail(
+        "transformers #46520 is in a stable release: PreTrainedModel.kernelize is gone "
+        "and patch_kernelize_fixes() now silently no-ops. Re-target it to set_use_kernels.",
+        pytrace=False,
+    )
+
+
+@pytest.fixture
+def kernelize_patch():
+    """Install the patch, restore everything afterwards."""
+    from axolotl.monkeypatch.kernelize_fixes import (
+        patch_kernelize_fixes,
+        unpatch_kernelize_fixes,
+    )
+
+    saved_sig = None
+    try:
+        from transformers.models.gpt_oss import modeling_gpt_oss
+
+        func = modeling_gpt_oss.apply_rotary_pos_emb
+        if hasattr(type(func), "forward"):
+            saved_sig = inspect.signature(type(func).forward)
+    except ImportError:
+        func = None
+
+    assert patch_kernelize_fixes() is True
+    yield
+
+    unpatch_kernelize_fixes()
+    if func is not None and saved_sig is not None:
+        type(func).forward.__signature__ = saved_sig
+
+
+def _tiny_gpt_oss():
+    from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig
+    from transformers.models.gpt_oss.modeling_gpt_oss import GptOssForCausalLM
+
+    cfg = GptOssConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+        vocab_size=128,
+        num_local_experts=4,
+        num_experts_per_tok=2,
+    )
+    return GptOssForCausalLM(cfg)
+
+
+def test_patch_is_idempotent(kernelize_patch):
+    from axolotl.monkeypatch.kernelize_fixes import patch_kernelize_fixes
+
+    assert patch_kernelize_fixes() is True
+
+
+def test_unpatch_restores_original():
+    from transformers.modeling_utils import PreTrainedModel
+
+    from axolotl.monkeypatch.kernelize_fixes import (
+        patch_kernelize_fixes,
+        unpatch_kernelize_fixes,
+    )
+
+    original = PreTrainedModel.kernelize
+    patch_kernelize_fixes()
+    assert PreTrainedModel.kernelize is not original
+    unpatch_kernelize_fixes()
+    assert PreTrainedModel.kernelize is original
+    # Safe to call again without a prior patch.
+    unpatch_kernelize_fixes()
+
+
+def test_gpt_oss_kernelize_and_rotary_signature(kernelize_patch):
+    """gpt-oss: kernelize() succeeds and the rotary signature matches the hub
+    kernel (kernels-community/rotary) afterwards."""
+    pytest.importorskip("transformers.models.gpt_oss")
+    from transformers.models.gpt_oss.modeling_gpt_oss import apply_rotary_pos_emb
+
+    model = _tiny_gpt_oss()
+    model.train()
+    model.kernelize()
+
+    params = inspect.signature(type(apply_rotary_pos_emb).forward).parameters
+    assert list(params) == ["self", "q", "k", "cos", "sin", "unsqueeze_dim"]
+
+
+def test_rotary_call_behavior_unchanged(kernelize_patch):
+    """Only signature metadata changes; calls (even with position_ids) still
+    produce identical results."""
+    import torch
+
+    pytest.importorskip("transformers.models.gpt_oss")
+    from transformers.models.gpt_oss.modeling_gpt_oss import apply_rotary_pos_emb
+
+    torch.manual_seed(0)
+    q, k = torch.randn(1, 4, 16, 8), torch.randn(1, 2, 16, 8)
+    cos, sin = torch.randn(1, 16, 4), torch.randn(1, 16, 4)
+
+    q_ref, k_ref = apply_rotary_pos_emb(q, k, cos, sin)
+    _tiny_gpt_oss().kernelize()
+    q_new, k_new = apply_rotary_pos_emb(q, k, cos, sin, position_ids=None)
+
+    assert torch.equal(q_ref, q_new) and torch.equal(k_ref, k_new)
+
+
+def test_kernels_signature_validation_passes(kernelize_patch):
+    """The exact kernels-library check that crashed gpt-oss training on CUDA."""
+    pytest.importorskip("transformers.models.gpt_oss")
+    from kernels.layer.func import _create_func_module
+    from kernels.layer.layer import _validate_layer
+    from transformers.models.gpt_oss.modeling_gpt_oss import apply_rotary_pos_emb
+
+    # Exact signature of kernels-community/rotary::apply_rotary_transformers.
+    def apply_rotary_transformers(q, k, cos, sin, unsqueeze_dim=1):
+        return q, k
+
+    hub_cls = _create_func_module(apply_rotary_transformers)
+    local_cls = type(apply_rotary_pos_emb)
+
+    with pytest.raises(TypeError, match="different number of arguments"):
+        _validate_layer(check_cls=local_cls, cls=hub_cls, repo="stub")
+
+    _tiny_gpt_oss().kernelize()
+    _validate_layer(check_cls=local_cls, cls=hub_cls, repo="stub")
+
+
+def test_bare_function_entries_are_dropped(kernelize_patch):
+    """Architectures that stash a bare function (gemma4 and ~30 others) no
+    longer crash kernelize(); simulated by planting one on gpt-oss."""
+    model = _tiny_gpt_oss()
+    attn = model.model.layers[0].self_attn
+
+    def bare(q, k, cos, sin):
+        return q, k
+
+    attn.__dict__.setdefault("_hidden_kernels", {})["bare"] = bare
+    model.train()
+    model.kernelize()
+    assert "bare" not in attn._hidden_kernels
+
+
+def _tiny_gemma4():
+    from transformers.models.gemma4.configuration_gemma4 import (
+        Gemma4AudioConfig,
+        Gemma4Config,
+        Gemma4TextConfig,
+        Gemma4VisionConfig,
+    )
+    from transformers.models.gemma4.modeling_gemma4 import (
+        Gemma4ForConditionalGeneration,
+    )
+
+    text = Gemma4TextConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+        vocab_size=128,
+        num_experts=4,
+        num_experts_per_tok=2,
+    )
+    vis = Gemma4VisionConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+    )
+    aud = Gemma4AudioConfig(
+        hidden_size=32, intermediate_size=64, num_hidden_layers=1, num_attention_heads=4
+    )
+    return Gemma4ForConditionalGeneration(
+        Gemma4Config(text_config=text, vision_config=vis, audio_config=aud)
+    )
+
+
+def test_gemma4_kernelize_succeeds_with_patch():
+    """The real gemma4 bare-function case end to end: with the generic patch,
+    kernelize() succeeds. The unpatched call raises on transformers releases that
+    still carry the bug and succeeds once the upstream fix lands, so that half is
+    tolerated rather than required."""
+    pytest.importorskip("transformers.models.gemma4")
+    from axolotl.monkeypatch.kernelize_fixes import (
+        patch_kernelize_fixes,
+        unpatch_kernelize_fixes,
+    )
+
+    model = _tiny_gemma4()
+    model.train()
+    try:
+        # transformers <= 5.8.x raises TypeError, >= 5.9 ValueError; fixed on main.
+        model.kernelize()
+    except (TypeError, ValueError, AttributeError):
+        pass
+
+    patch_kernelize_fixes()
+    try:
+        model = _tiny_gemma4()
+        model.train()
+        model.kernelize()
+    finally:
+        unpatch_kernelize_fixes()
+
+
+def test_patch_does_not_alter_weights(kernelize_patch):
+    """The repairs only touch ``_hidden_kernels`` and signature metadata;
+    parameters are untouched by kernelize()."""
+    import torch
+
+    torch.manual_seed(0)
+    model = _tiny_gpt_oss()
+    before = {k: v.clone() for k, v in model.state_dict().items()}
+    model.train()
+    model.kernelize()
+    after = model.state_dict()
+
+    assert before.keys() == after.keys()
+    assert all(torch.equal(before[k], after[k]) for k in before)
diff --git a/tests/monkeypatch/test_large_head_attention.py b/tests/monkeypatch/test_large_head_attention.py
new file mode 100644
index 0000000000..459ff57426
--- /dev/null
+++ b/tests/monkeypatch/test_large_head_attention.py
@@ -0,0 +1,137 @@
+"""Tests for the generic large-head-dim attention capability (#5/large_head_attention).
+
+Policy resolution and the SDPA-decline paths run on CPU (they return before touching the kernel);
+the route-success paths exercise the real Triton flash_d512 kernel and are GPU-gated."""
+
+import pytest
+import torch
+
+from axolotl.monkeypatch.attention import large_head as lh
+
+requires_gpu = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="flash_d512 Triton kernel needs CUDA"
+)
+
+
+class Cfg(dict):
+    def __getattr__(self, k):
+        return self.get(k)
+
+
+class _Mod:
+    num_key_value_groups = 1
+
+
+def _packed_pos(docs):
+    return torch.cat([torch.arange(d) for d in docs])[None]
+
+
+def test_resolve_policy_intent_alias_default():
+    assert (
+        lh.resolve_large_head_policy(Cfg(large_head_attention="triton_flash"))
+        == "triton_flash"
+    )
+    assert lh.resolve_large_head_policy(Cfg(large_head_attention="sdpa")) == "sdpa"
+    assert (
+        lh.resolve_large_head_policy(Cfg(flash_attn_d512=True)) == "auto"
+    )  # legacy alias
+    assert lh.resolve_large_head_policy(Cfg()) == "sdpa"  # default
+
+
+def test_route_declines_for_sdpa_policy():
+    q = torch.zeros(1, 2, 8, 512)
+    assert (
+        lh.flash_d512_route(_Mod(), q, q, q, None, _packed_pos([4, 4]), policy="sdpa")
+        is None
+    )
+
+
+def test_route_declines_for_unknown_policy():
+    # A config typo must NOT silently route through the Triton kernel: only auto/triton_flash do.
+    q = torch.zeros(1, 2, 8, 512)
+    assert (
+        lh.flash_d512_route(
+            _Mod(), q, q, q, None, _packed_pos([4, 4]), policy="trtion_flsah"
+        )
+        is None
+    )
+
+
+def test_set_policy_resets_to_default():
+    lh.set_large_head_policy("auto")
+    assert lh.get_large_head_policy() == "auto"
+    lh.set_large_head_policy(None)  # a later run without the field must reset to sdpa
+    assert lh.get_large_head_policy() == "sdpa"
+
+
+def test_route_declines_for_small_head_dim():
+    q = torch.zeros(1, 2, 8, 128)  # head_dim 128 -> not a large-head case
+    assert (
+        lh.flash_d512_route(_Mod(), q, q, q, None, _packed_pos([4, 4]), policy="auto")
+        is None
+    )
+
+
+def test_route_declines_single_doc_under_auto():
+    q = torch.zeros(1, 2, 8, 512)
+    single = torch.arange(8)[None]  # one document -> auto prefers SDPA is_causal
+    assert lh.flash_d512_route(_Mod(), q, q, q, None, single, policy="auto") is None
+
+
+@requires_gpu
+@pytest.mark.parametrize("scaling", [512**-0.5, 1.0])
+def test_route_runs_custom_scale(scaling):
+    # Custom attention scale is now supported (gemma4 global uses scaling=1.0): route, don't decline.
+    q = torch.randn(1, 4, 256, 512, device="cuda", dtype=torch.bfloat16)
+    out = lh.flash_d512_route(
+        _Mod(), q, q, q, scaling, _packed_pos([128, 128]).cuda(), policy="auto"
+    )
+    assert out is not None
+
+
+@requires_gpu
+def test_route_success_packed():
+    # exercise the real kernel: packed multi-doc large-head input must route and honor the contract.
+    B, H, S, D = 1, 4, 256, 512
+    q = torch.randn(B, H, S, D, device="cuda", dtype=torch.bfloat16)
+    out = lh.flash_d512_route(
+        _Mod(), q, q, q, D**-0.5, _packed_pos([128, 128]).cuda(), policy="auto"
+    )
+    assert out is not None
+    attn, weights = out
+    assert weights is None
+    assert attn.shape == (B, S, H, D)  # transposed to sdpa_attention_forward's contract
+    assert torch.isfinite(attn).all()
+
+
+@requires_gpu
+def test_route_gqa_repeats_kv():
+    # GQA path: 4 kv heads expanded to 16 q heads inside the route, then the kernel runs.
+    class GQA:
+        num_key_value_groups = 4
+
+    q = torch.randn(1, 16, 256, 512, device="cuda", dtype=torch.bfloat16)
+    kv = torch.randn(1, 4, 256, 512, device="cuda", dtype=torch.bfloat16)
+    out = lh.flash_d512_route(
+        GQA(), q, kv, kv, 512**-0.5, _packed_pos([128, 128]).cuda(), policy="auto"
+    )
+    assert out is not None
+    attn, _ = out
+    assert attn.shape == (1, 256, 16, 512)  # expanded to 16 q-heads, ran, transposed
+    assert torch.isfinite(attn).all()
+
+
+def test_patch_and_unpatch_sdpa(monkeypatch):
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    original = ALL_ATTENTION_FUNCTIONS["sdpa"]
+    lh.set_large_head_policy("sdpa")
+    assert lh.patch_sdpa_large_head() is False  # policy sdpa -> no wrap
+    assert ALL_ATTENTION_FUNCTIONS["sdpa"] is original
+
+    assert lh.patch_sdpa_large_head("auto") is True  # wraps
+    assert ALL_ATTENTION_FUNCTIONS["sdpa"] is not original
+    assert lh.patch_sdpa_large_head("auto") is True  # idempotent (still wrapped)
+    lh.unpatch_sdpa_large_head()
+    assert ALL_ATTENTION_FUNCTIONS["sdpa"] is original
+    lh.set_large_head_policy("sdpa")
diff --git a/tests/monkeypatch/test_large_head_compile_hoist.py b/tests/monkeypatch/test_large_head_compile_hoist.py
new file mode 100644
index 0000000000..13b2fe269a
--- /dev/null
+++ b/tests/monkeypatch/test_large_head_compile_hoist.py
@@ -0,0 +1,117 @@
+"""Tests for compile-friendly multidoc detection and doc_end vectorization."""
+
+import pytest
+import torch
+
+# Import the production helper directly so the parity tests exercise the real
+# vectorized path instead of a duplicate that could silently drift.
+from axolotl.monkeypatch.attention.flash_attn_d512 import (
+    _compute_doc_end as _doc_end_vectorized,
+)
+from axolotl.monkeypatch.attention.large_head import (
+    _multidoc_position_ids,
+    set_large_head_packed,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset_packed():
+    yield
+    set_large_head_packed(None)
+
+
+class TestStaticPackedFlag:
+    def _packed_pos(self):
+        return torch.tensor([[0, 1, 2, 0, 1, 0, 1, 2]])
+
+    def _single_doc_pos(self):
+        return torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7]])
+
+    def test_runtime_probe_when_undeclared(self):
+        set_large_head_packed(None)
+        assert _multidoc_position_ids(self._packed_pos()) is not None
+        assert _multidoc_position_ids(self._single_doc_pos()) is None
+
+    def test_static_packed_true_skips_probe(self):
+        set_large_head_packed(True)
+        # even a single-doc row routes varlen: no data-dependent branch
+        assert _multidoc_position_ids(self._single_doc_pos()) is not None
+        assert _multidoc_position_ids(self._packed_pos()) is not None
+
+    def test_static_packed_false(self):
+        set_large_head_packed(False)
+        assert _multidoc_position_ids(self._packed_pos()) is None
+
+    def test_none_position_ids(self):
+        set_large_head_packed(True)
+        assert _multidoc_position_ids(None) is None
+
+
+def _doc_end_loop_reference(pos: torch.Tensor) -> torch.Tensor:
+    B, N = pos.shape
+    doc_end = torch.empty_like(pos)
+    for b in range(B):
+        starts = (pos[b] == 0).nonzero().flatten()
+        bounds = torch.cat(
+            [starts, torch.tensor([N], device=pos.device, dtype=starts.dtype)]
+        )
+        doc_end[b] = (
+            bounds[1:].repeat_interleave(bounds[1:] - bounds[:-1]).to(pos.dtype)
+        )
+    return doc_end
+
+
+class TestDocEndVectorization:
+    @pytest.mark.parametrize("seed", [0, 1, 2, 3])
+    def test_parity_random_packing(self, seed):
+        gen = torch.Generator().manual_seed(seed)
+        B, N = 3, 64
+        rows = []
+        for _ in range(B):
+            lengths = []
+            remaining = N
+            while remaining > 0:
+                length = int(torch.randint(1, 17, (1,), generator=gen))
+                length = min(length, remaining)
+                lengths.append(length)
+                remaining -= length
+            rows.append(torch.cat([torch.arange(n) for n in lengths]))
+        pos = torch.stack(rows).to(torch.int32)
+        torch.testing.assert_close(
+            _doc_end_vectorized(pos), _doc_end_loop_reference(pos)
+        )
+
+    def test_parity_single_doc(self):
+        pos = torch.arange(32, dtype=torch.int32)[None]
+        torch.testing.assert_close(
+            _doc_end_vectorized(pos), _doc_end_loop_reference(pos)
+        )
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+    def test_flash_d512_varlen_uses_vectorized_doc_end(self):
+        # end-to-end: packed flash_d512 still matches per-document SDPA
+        from axolotl.monkeypatch.attention.flash_attn_d512 import flash_d512
+
+        torch.manual_seed(0)
+        B, H, D = 1, 2, 512
+        lengths = [24, 40]
+        N = sum(lengths)
+        pos = torch.cat([torch.arange(n) for n in lengths])[None].cuda().to(torch.int32)
+        q, k, v = (
+            torch.randn(B, H, N, D, device="cuda", dtype=torch.bfloat16)
+            for _ in range(3)
+        )
+        out = flash_d512(q, k, v, causal=True, position_ids=pos)
+
+        ref = torch.empty_like(out)
+        start = 0
+        for n in lengths:
+            sl = slice(start, start + n)
+            ref[:, :, sl] = torch.nn.functional.scaled_dot_product_attention(
+                q[:, :, sl].float(),
+                k[:, :, sl].float(),
+                v[:, :, sl].float(),
+                is_causal=True,
+            ).to(out.dtype)
+            start += n
+        torch.testing.assert_close(out.float(), ref.float(), atol=2e-2, rtol=2e-2)
diff --git a/tests/monkeypatch/test_llama_attn_hijack_flash.py b/tests/monkeypatch/test_llama_attn_hijack_flash.py
index 4521cd07bc..93df484b1b 100644
--- a/tests/monkeypatch/test_llama_attn_hijack_flash.py
+++ b/tests/monkeypatch/test_llama_attn_hijack_flash.py
@@ -1,6 +1,7 @@
 """
 Unit tests for the monkeypatch utils
 """
+
 import unittest
 
 import torch
@@ -49,6 +50,20 @@ def test_get_max_seqlen_in_batch(self):
         target_res = torch.tensor([4, 3, 5, 2], dtype=torch.int32)
         self.assertTrue(torch.allclose(get_max_seqlen_in_batch(attn_mask), target_res))
 
+    def test_get_max_seqlen_in_batch_adjacent_rows_same_id(self):
+        # Rows whose first/last segment ids match across the row boundary must not
+        # be merged into a single segment (each row is a fresh pack).
+        attn_mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 1]])
+        target_res = torch.tensor([4, 4], dtype=torch.int32)
+        self.assertTrue(torch.allclose(get_max_seqlen_in_batch(attn_mask), target_res))
+
+        indices, cu_seqlen, max_seqlen_in_batch = get_unpad_data(attn_mask)
+        self.assertTrue(torch.allclose(indices, torch.arange(8)))
+        self.assertTrue(
+            torch.allclose(cu_seqlen, torch.tensor([0, 4, 8], dtype=torch.int32))
+        )
+        self.assertEqual(max_seqlen_in_batch, 4)
+
     def test_get_unpad_data(self):
         attn_mask = torch.tensor([[1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 0, 0]])
         target_indices = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13])
diff --git a/tests/monkeypatch/test_lora_kernels_gemma4_unified.py b/tests/monkeypatch/test_lora_kernels_gemma4_unified.py
new file mode 100644
index 0000000000..3d009f8bf0
--- /dev/null
+++ b/tests/monkeypatch/test_lora_kernels_gemma4_unified.py
@@ -0,0 +1,75 @@
+"""Tests for ``patch_self_attn_lora`` on the gemma4_unified attention class.
+
+Unlike standard gemma4 (which always skips the LoRA source rewrite), the unified
+branch is *conditional* on ``fused_attn_kernel``: with it set the QKV/O kernels
+ride the fused forward; without it the rewrite is skipped with a warning. Either
+way the source rewrite (``_original_forward``) is never installed."""
+
+import logging
+
+import pytest
+
+pytest.importorskip("triton", reason="importing lora_kernels pulls in triton")
+gemma4_unified_modeling = pytest.importorskip(
+    "transformers.models.gemma4_unified.modeling_gemma4_unified",
+    reason="unified lora-kernel branch only matters when gemma4_unified is available",
+)
+
+
+def _cfg(fused_attn_kernel):
+    from axolotl.utils.dict import DictDefault
+
+    return DictDefault({"fused_attn_kernel": fused_attn_kernel, "lora_dropout": 0.0})
+
+
+@pytest.fixture
+def restore_unified_attn():
+    """Ensure ``_original_forward`` doesn't leak in/out of these tests."""
+    cls = gemma4_unified_modeling.Gemma4UnifiedTextAttention
+    had = hasattr(cls, "_original_forward")
+    saved = getattr(cls, "_original_forward", None)
+    if had:
+        del cls._original_forward
+    yield cls
+    if had:
+        cls._original_forward = saved
+    elif hasattr(cls, "_original_forward"):
+        del cls._original_forward
+
+
+def _run_patch(monkeypatch, caplog, cls, cfg, level):
+    from axolotl.monkeypatch import lora_kernels
+
+    monkeypatch.setattr(lora_kernels, "get_attention_cls_from_config", lambda _cfg: cls)
+    logger = logging.getLogger("axolotl.monkeypatch.lora_kernels")
+    logger.addHandler(caplog.handler)
+    previous_level = logger.level
+    logger.setLevel(level)
+    try:
+        lora_kernels.patch_self_attn_lora(cfg)
+    finally:
+        logger.removeHandler(caplog.handler)
+        logger.setLevel(previous_level)
+
+
+class TestUnifiedLoraKernelSkip:
+    def test_skips_with_warning_when_no_fused_attn_kernel(
+        self, restore_unified_attn, monkeypatch, caplog
+    ):
+        cls = restore_unified_attn
+        _run_patch(monkeypatch, caplog, cls, _cfg(False), logging.WARNING)
+
+        assert "fused_attn_kernel" in caplog.text
+        assert "skipping" in caplog.text.lower()
+        assert not hasattr(cls, "_original_forward"), (
+            "unified attention must not get the LoRA source rewrite"
+        )
+
+    def test_skips_quietly_with_fused_attn_kernel(
+        self, restore_unified_attn, monkeypatch, caplog
+    ):
+        cls = restore_unified_attn
+        _run_patch(monkeypatch, caplog, cls, _cfg(True), logging.INFO)
+
+        assert "fused attention path" in caplog.text
+        assert not hasattr(cls, "_original_forward")
diff --git a/tests/monkeypatch/test_lora_mlp_routed_expert_guard.py b/tests/monkeypatch/test_lora_mlp_routed_expert_guard.py
new file mode 100644
index 0000000000..e386176b9d
--- /dev/null
+++ b/tests/monkeypatch/test_lora_mlp_routed_expert_guard.py
@@ -0,0 +1,67 @@
+"""CPU tests for the lora_mlp_kernel routed-expert guard in find_mlp_in_layer.
+
+When a custom MoE expert kernel (ScatterMoE/SonicMoE) owns the routed experts, the generic
+lora_mlp_kernel must fuse ONLY the dense shared MLP, never the routed-expert containers (which the
+MoE kernel handles). The dense MLP stays patchable in both cases.
+"""
+
+from types import SimpleNamespace
+
+import torch.nn as nn
+
+from axolotl.monkeypatch.lora_kernels import find_mlp_in_layer
+
+
+def _lin():
+    return nn.Linear(4, 4)
+
+
+def _dense_mlp():
+    return SimpleNamespace(gate_proj=_lin(), up_proj=_lin(), down_proj=_lin())
+
+
+def _routed_experts(n):
+    return SimpleNamespace(
+        gate_projs=[_lin() for _ in range(n)],
+        up_projs=[_lin() for _ in range(n)],
+        down_projs=[_lin() for _ in range(n)],
+    )
+
+
+def test_dense_shared_mlp_always_found():
+    layer = SimpleNamespace(mlp=_dense_mlp())
+    for skip in (False, True):
+        mlps = list(find_mlp_in_layer(layer, skip_routed_experts=skip))
+        assert len(mlps) == 1
+        assert mlps[0][3] is layer.mlp  # the dense MLP module itself
+
+
+def test_routed_experts_skipped_when_moe_kernel_owns_them():
+    layer = SimpleNamespace(feedforward=SimpleNamespace(experts=_routed_experts(3)))
+    # default (no custom MoE kernel): routed experts are yielded for fusion
+    assert len(list(find_mlp_in_layer(layer))) == 3
+    # under ScatterMoE/SonicMoE: routed experts must NOT be yielded
+    assert list(find_mlp_in_layer(layer, skip_routed_experts=True)) == []
+
+
+def test_dense_kept_routed_skipped_together():
+    layer = SimpleNamespace(
+        mlp=_dense_mlp(),
+        feedforward=SimpleNamespace(experts=_routed_experts(2)),
+    )
+    full = list(find_mlp_in_layer(layer, skip_routed_experts=False))
+    guarded = list(find_mlp_in_layer(layer, skip_routed_experts=True))
+    assert len(full) == 3  # dense + 2 routed
+    assert len(guarded) == 1  # only the dense MLP
+    assert guarded[0][3] is layer.mlp
+
+
+def test_dsv4_translation_disables_generic_mlp_kernel():
+    # DSV4 keeps its dedicated clamped-SwiGLU kernel: generic lora_mlp_kernel is translated and off.
+    from axolotl.integrations.kernels.args import KernelsArgs
+
+    out = KernelsArgs.disable_mlp_kernel(
+        {"use_scattermoe": True, "use_dsv4_kernels": True, "lora_mlp_kernel": True}
+    )
+    assert out["lora_mlp_kernel"] is False
+    assert out["dsv4_shared_mlp_lora_kernel"] is True
diff --git a/tests/monkeypatch/test_mamba_utils.py b/tests/monkeypatch/test_mamba_utils.py
new file mode 100644
index 0000000000..a473df29c1
--- /dev/null
+++ b/tests/monkeypatch/test_mamba_utils.py
@@ -0,0 +1,551 @@
+"""Unit tests for shared Mamba2 SSM utilities (mamba_utils.py).
+
+Tests cover get_seq_idx correctness under:
+  - single-rank packing
+  - context parallelism (mid-sample chunk starts)
+  - batch dimension
+  - dtype and device
+  - no-negative regression (CP rank > 0 must never produce -1)
+  - mamba2_cp_correction mathematical correctness
+  - wrap_mamba_scan_for_cp wrapper behaviour
+  - end-to-end CP split: full 2K scan == 2×1K split + correction
+"""
+
+import types
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+
+from axolotl.monkeypatch.models.mamba_utils import (
+    get_seq_idx,
+    mamba2_cp_correction,
+    wrap_mamba_scan_for_cp,
+)
+
+
+def _reference_ssm_scan(x, dt, A, B, C, dt_bias=None, dt_softplus=False, h0=None):
+    """Pure-PyTorch step-by-step SSM scan (reference implementation).
+
+    Implements the Mamba2 discrete SSM recurrence:
+        Δ_t = softplus(dt_t + dt_bias)  or  dt_t
+        Ā_t = exp(A · Δ_t)
+        h_t = Ā_t · h_{t-1}  +  B_t ⊗ x_t
+        y_t = (C_t · h_t).sum(dim=n)
+
+    Args:
+        x:  [B, T, H, d]
+        dt: [B, T, H]
+        A:  [H]  (log-space, negative)
+        B:  [B, T, n_groups, n]
+        C:  [B, T, n_groups, n]
+        dt_bias:    [H] or None
+        dt_softplus: bool
+        h0: [B, H, d, n] initial state, or None → zeros
+
+    Returns:
+        out:     [B, T, H, d]
+        h_final: [B, H, d, n]
+    """
+    B_batch, T, H, d = x.shape
+    n_groups = B.shape[2]
+    n = B.shape[3]
+    heads_per_group = H // n_groups
+
+    dt_eff = dt + dt_bias[None, None, :] if dt_bias is not None else dt
+    if dt_softplus:
+        dt_eff = F.softplus(dt_eff)
+
+    h = torch.zeros(B_batch, H, d, n, dtype=x.dtype) if h0 is None else h0.clone()
+
+    outputs = []
+    for t in range(T):
+        A_bar = torch.exp(A[None, :] * dt_eff[:, t, :])  # [B, H]
+        B_t = B[:, t].repeat_interleave(heads_per_group, dim=1)  # [B, H, n]
+        C_t = C[:, t].repeat_interleave(heads_per_group, dim=1)  # [B, H, n]
+
+        h = A_bar[:, :, None, None] * h + B_t[:, :, None, :] * x[:, t, :, :, None]
+        y_t = (C_t[:, :, None, :] * h).sum(dim=-1)  # [B, H, d]
+        outputs.append(y_t)
+
+    return torch.stack(outputs, dim=1), h
+
+
+class TestGetSeqIdx:
+    """Tests for get_seq_idx(position_ids) → seq_idx."""
+
+    def test_single_sample_no_packing(self):
+        """Single sample with no packing: all zeros."""
+        pos = torch.tensor([[0, 1, 2, 3, 4]])
+        out = get_seq_idx(pos)
+        assert out.tolist() == [[0, 0, 0, 0, 0]]
+
+    def test_two_packed_samples(self):
+        """Two packed samples: index increments at the second sample boundary."""
+        pos = torch.tensor([[0, 1, 2, 3, 0, 1, 2]])
+        out = get_seq_idx(pos)
+        assert out.tolist() == [[0, 0, 0, 0, 1, 1, 1]]
+
+    def test_three_packed_samples(self):
+        """Three packed samples."""
+        pos = torch.tensor([[0, 1, 0, 1, 2, 0]])
+        out = get_seq_idx(pos)
+        assert out.tolist() == [[0, 0, 1, 1, 1, 2]]
+
+    def test_cp_rank_mid_sample_start(self):
+        """CP rank > 0: chunk starts mid-sample (position_ids[0] != 0).
+
+        Must produce non-negative seq_idx starting at 0, not -1.
+        """
+        pos = torch.tensor([[3, 4, 5, 0, 1, 2]])
+        out = get_seq_idx(pos)
+        assert out.tolist() == [[0, 0, 0, 1, 1, 1]]
+
+    def test_cp_rank_entire_chunk_mid_sample(self):
+        """CP rank whose entire chunk is mid-sample (no sample boundary)."""
+        pos = torch.tensor([[5, 6, 7, 8, 9]])
+        out = get_seq_idx(pos)
+        assert out.tolist() == [[0, 0, 0, 0, 0]]
+
+    def test_no_negative_values_regression(self):
+        """seq_idx must never contain -1 for any valid position_ids input."""
+        cases = [
+            [[1, 2, 3]],
+            [[10, 11, 12, 0, 1]],
+            [[0, 0, 0]],
+        ]
+        for pos_list in cases:
+            pos = torch.tensor(pos_list)
+            out = get_seq_idx(pos)
+            assert out.min().item() >= 0, f"Negative seq_idx for pos={pos_list}"
+
+    def test_batch_dimension(self):
+        """Batch of 3 sequences, each independently packed."""
+        pos = torch.tensor(
+            [
+                [0, 1, 2, 0, 1],
+                [0, 1, 0, 1, 2],
+                [3, 4, 0, 1, 2],
+            ]
+        )
+        out = get_seq_idx(pos)
+        assert out.tolist() == [
+            [0, 0, 0, 1, 1],
+            [0, 0, 1, 1, 1],
+            [0, 0, 1, 1, 1],
+        ]
+
+    def test_output_dtype_is_int32(self):
+        """Output dtype must be torch.int32 (mamba-ssm kernel requirement)."""
+        pos = torch.tensor([[0, 1, 2, 0, 1]])
+        out = get_seq_idx(pos)
+        assert out.dtype == torch.int32
+
+    def test_output_shape_matches_input(self):
+        """Output shape matches input shape."""
+        pos = torch.zeros(4, 128, dtype=torch.long)
+        out = get_seq_idx(pos)
+        assert out.shape == pos.shape
+
+    def test_single_token(self):
+        """Edge case: single token sequence."""
+        pos = torch.tensor([[0]])
+        out = get_seq_idx(pos)
+        assert out.tolist() == [[0]]
+
+    def test_cp_rank_starts_at_1(self):
+        """CP rank that starts exactly at position 1 (not 0)."""
+        pos = torch.tensor([[1, 2, 3, 0, 1]])
+        out = get_seq_idx(pos)
+        assert out.tolist() == [[0, 0, 0, 1, 1]]
+
+    def test_many_packed_samples(self):
+        """Many single-token samples packed together."""
+        pos = torch.tensor([[0, 0, 0, 0, 0, 0]])
+        out = get_seq_idx(pos)
+        assert out.tolist() == [[0, 1, 2, 3, 4, 5]]
+
+
+class TestMamba2CpCorrection:
+    """Tests for mamba2_cp_correction mathematical correctness."""
+
+    def test_zero_h_prev_is_noop(self):
+        """When h_prev is all zeros, output should be unchanged."""
+        B, T, H, d, n = 1, 8, 4, 16, 8
+        n_groups = 2
+
+        out = torch.randn(B, T, H * d)
+        h_final = torch.randn(B, H, d, n)
+        C = torch.randn(B, T, n_groups, n)
+        cum_A = torch.randn(B, T, H)
+        h_prev = torch.zeros(B, H, d, n)
+
+        corrected_out, corrected_h = mamba2_cp_correction(
+            out,
+            h_final,
+            C,
+            cum_A,
+            h_prev,
+            num_heads=H,
+            head_dim=d,
+        )
+
+        torch.testing.assert_close(corrected_out, out)
+        torch.testing.assert_close(corrected_h, h_final)
+
+    def test_correction_shapes(self):
+        """Output shapes must match input shapes."""
+        B, T, H, d, n = 2, 16, 8, 32, 16
+        n_groups = 4
+
+        out = torch.randn(B, T, H * d)
+        h_final = torch.randn(B, H, d, n)
+        C = torch.randn(B, T, n_groups, n)
+        cum_A = torch.randn(B, T, H)
+        h_prev = torch.randn(B, H, d, n)
+
+        corrected_out, corrected_h = mamba2_cp_correction(
+            out,
+            h_final,
+            C,
+            cum_A,
+            h_prev,
+            num_heads=H,
+            head_dim=d,
+        )
+
+        assert corrected_out.shape == out.shape
+        assert corrected_h.shape == h_final.shape
+
+    def test_correction_adds_to_output(self):
+        """With nonzero h_prev, output should differ from input."""
+        B, T, H, d, n = 1, 4, 2, 8, 4
+        n_groups = 1
+
+        out = torch.zeros(B, T, H * d)
+        h_final = torch.zeros(B, H, d, n)
+        C = torch.ones(B, T, n_groups, n)
+        cum_A = torch.zeros(B, T, H)  # exp(0) = 1, so full propagation
+        h_prev = torch.ones(B, H, d, n)
+
+        corrected_out, corrected_h = mamba2_cp_correction(
+            out,
+            h_final,
+            C,
+            cum_A,
+            h_prev,
+            num_heads=H,
+            head_dim=d,
+        )
+
+        # With exp(cum_A)=1, C=1, h_prev=1: delta_y should be nonzero
+        assert corrected_out.abs().sum() > 0
+        assert corrected_h.abs().sum() > 0
+
+    def test_correction_h_final_formula(self):
+        """Verify h_final correction: h_final + decay_T * h_prev."""
+        B, T, H, d, n = 1, 4, 2, 8, 4
+        n_groups = 1
+
+        h_final = torch.zeros(B, H, d, n)
+        C = torch.ones(B, T, n_groups, n)
+        cum_A = torch.zeros(B, T, H)
+        h_prev = torch.ones(B, H, d, n) * 2.0
+        out = torch.zeros(B, T, H * d)
+
+        _, corrected_h = mamba2_cp_correction(
+            out,
+            h_final,
+            C,
+            cum_A,
+            h_prev,
+            num_heads=H,
+            head_dim=d,
+        )
+
+        # exp(0) * 2.0 = 2.0 for all elements
+        expected = torch.ones(B, H, d, n) * 2.0
+        torch.testing.assert_close(corrected_h, expected)
+
+
+class TestCpSplitMatchesFullScan:
+    """End-to-end: full sequence scan == split into chunks + CP correction.
+
+    Runs a reference SSM scan on a full 2K-token sequence, then simulates
+    2-rank CP by splitting into 2×1K, running each half with h₀=0, and
+    applying mamba2_cp_correction to rank 1 using rank 0's final state.
+    The concatenated result must match the single-rank reference.
+    """
+
+    def test_2k_vs_2x1k_output_matches(self):
+        """Full 2048-token scan output == two 1024-token chunks + CP correction."""
+        torch.manual_seed(42)
+        B, T, H, d, n = 1, 2048, 4, 16, 8
+        n_groups = 2
+        dt_bias = torch.randn(H) * 0.1
+
+        x = torch.randn(B, T, H, d)
+        dt = torch.randn(B, T, H) * 0.1
+        A = -torch.rand(H).abs() - 0.01
+        B_ssm = torch.randn(B, T, n_groups, n) * 0.1
+        C_ssm = torch.randn(B, T, n_groups, n) * 0.1
+
+        ref_out, ref_h = _reference_ssm_scan(
+            x, dt, A, B_ssm, C_ssm, dt_bias=dt_bias, dt_softplus=True
+        )
+
+        T2 = T // 2
+
+        out_0, h_final_0 = _reference_ssm_scan(
+            x[:, :T2],
+            dt[:, :T2],
+            A,
+            B_ssm[:, :T2],
+            C_ssm[:, :T2],
+            dt_bias=dt_bias,
+            dt_softplus=True,
+        )
+
+        out_1, h_final_1 = _reference_ssm_scan(
+            x[:, T2:],
+            dt[:, T2:],
+            A,
+            B_ssm[:, T2:],
+            C_ssm[:, T2:],
+            dt_bias=dt_bias,
+            dt_softplus=True,
+        )
+
+        dt_eff_1 = F.softplus(dt[:, T2:] + dt_bias[None, None, :])
+        cum_A_1 = torch.cumsum(A[None, None, :] * dt_eff_1, dim=1)
+
+        corrected_out_1, corrected_h_1 = mamba2_cp_correction(
+            out_1.view(B, T2, H * d),
+            h_final_1,
+            C_ssm[:, T2:],
+            cum_A_1,
+            h_final_0,
+            num_heads=H,
+            head_dim=d,
+        )
+        corrected_out_1 = corrected_out_1.view(B, T2, H, d)
+
+        reconstructed = torch.cat([out_0, corrected_out_1], dim=1)
+
+        torch.testing.assert_close(reconstructed, ref_out, rtol=1e-4, atol=1e-4)
+        torch.testing.assert_close(corrected_h_1, ref_h, rtol=1e-4, atol=1e-4)
+
+    def test_2k_vs_2x1k_with_batch(self):
+        """Same split test with batch_size > 1."""
+        torch.manual_seed(123)
+        B, T, H, d, n = 3, 512, 2, 8, 4
+        n_groups = 1
+        dt_bias = torch.randn(H) * 0.05
+
+        x = torch.randn(B, T, H, d)
+        dt = torch.randn(B, T, H) * 0.1
+        A = -torch.rand(H).abs() - 0.01
+        B_ssm = torch.randn(B, T, n_groups, n) * 0.1
+        C_ssm = torch.randn(B, T, n_groups, n) * 0.1
+
+        ref_out, ref_h = _reference_ssm_scan(
+            x, dt, A, B_ssm, C_ssm, dt_bias=dt_bias, dt_softplus=True
+        )
+
+        T2 = T // 2
+
+        out_0, h_0 = _reference_ssm_scan(
+            x[:, :T2],
+            dt[:, :T2],
+            A,
+            B_ssm[:, :T2],
+            C_ssm[:, :T2],
+            dt_bias=dt_bias,
+            dt_softplus=True,
+        )
+        out_1, h_1 = _reference_ssm_scan(
+            x[:, T2:],
+            dt[:, T2:],
+            A,
+            B_ssm[:, T2:],
+            C_ssm[:, T2:],
+            dt_bias=dt_bias,
+            dt_softplus=True,
+        )
+
+        dt_eff_1 = F.softplus(dt[:, T2:] + dt_bias[None, None, :])
+        cum_A_1 = torch.cumsum(A[None, None, :] * dt_eff_1, dim=1)
+
+        corrected_out_1, corrected_h_1 = mamba2_cp_correction(
+            out_1.view(B, T2, H * d),
+            h_1,
+            C_ssm[:, T2:],
+            cum_A_1,
+            h_0,
+            num_heads=H,
+            head_dim=d,
+        )
+
+        reconstructed = torch.cat([out_0, corrected_out_1.view(B, T2, H, d)], dim=1)
+
+        torch.testing.assert_close(reconstructed, ref_out, rtol=1e-4, atol=1e-4)
+        torch.testing.assert_close(corrected_h_1, ref_h, rtol=1e-4, atol=1e-4)
+
+    def test_4_way_split(self):
+        """4-rank CP: split 1024 tokens into 4×256 chunks with sequential correction."""
+        torch.manual_seed(99)
+        B, T, H, d, n = 1, 1024, 2, 8, 4
+        n_groups = 1
+        n_ranks = 4
+        chunk = T // n_ranks
+        dt_bias = torch.randn(H) * 0.05
+
+        x = torch.randn(B, T, H, d)
+        dt = torch.randn(B, T, H) * 0.1
+        A = -torch.rand(H).abs() - 0.01
+        B_ssm = torch.randn(B, T, n_groups, n) * 0.1
+        C_ssm = torch.randn(B, T, n_groups, n) * 0.1
+
+        ref_out, ref_h = _reference_ssm_scan(
+            x, dt, A, B_ssm, C_ssm, dt_bias=dt_bias, dt_softplus=True
+        )
+
+        all_outs = []
+        h_prev = torch.zeros(B, H, d, n)
+
+        for rank in range(n_ranks):
+            s, e = rank * chunk, (rank + 1) * chunk
+            out_r, h_r = _reference_ssm_scan(
+                x[:, s:e],
+                dt[:, s:e],
+                A,
+                B_ssm[:, s:e],
+                C_ssm[:, s:e],
+                dt_bias=dt_bias,
+                dt_softplus=True,
+            )
+
+            dt_eff_r = F.softplus(dt[:, s:e] + dt_bias[None, None, :])
+            cum_A_r = torch.cumsum(A[None, None, :] * dt_eff_r, dim=1)
+
+            corrected_out_r, corrected_h_r = mamba2_cp_correction(
+                out_r.view(B, chunk, H * d),
+                h_r,
+                C_ssm[:, s:e],
+                cum_A_r,
+                h_prev,
+                num_heads=H,
+                head_dim=d,
+            )
+
+            all_outs.append(corrected_out_r.view(B, chunk, H, d))
+            h_prev = corrected_h_r
+
+        reconstructed = torch.cat(all_outs, dim=1)
+
+        torch.testing.assert_close(reconstructed, ref_out, rtol=1e-3, atol=1e-3)
+        torch.testing.assert_close(h_prev, ref_h, rtol=1e-3, atol=1e-3)
+
+
+class TestWrapMambaScanForCp:
+    """Tests for wrap_mamba_scan_for_cp wrapper."""
+
+    @staticmethod
+    def _make_module_with_scan(scan_fn):
+        """Create a module namespace with a mamba_chunk_scan_combined attribute."""
+        mod = types.ModuleType("fake_mamba_module")
+        mod.mamba_chunk_scan_combined = scan_fn
+        return mod
+
+    def test_passthrough_when_cp_inactive(self):
+        """When CP is not active, wrapper should return original result unchanged."""
+        B, T, H, d, n = 1, 8, 4, 16, 8
+        x = torch.randn(B, T, H, d)
+        dt = torch.randn(B, T, H)
+        A = -torch.rand(H)
+        B_arg = torch.randn(B, T, 2, n)
+        C_arg = torch.randn(B, T, 2, n)
+        expected_out = torch.randn(B, T, H, d)
+        expected_state = torch.randn(B, H, d, n)
+
+        def fake_scan(*args, **kwargs):
+            return expected_out, expected_state
+
+        mod = self._make_module_with_scan(fake_scan)
+
+        with patch(
+            "axolotl.monkeypatch.models.mamba_utils.is_cp_active", return_value=False
+        ):
+            wrap_mamba_scan_for_cp(mod)
+            out, state = mod.mamba_chunk_scan_combined(
+                x,
+                dt,
+                A,
+                B_arg,
+                C_arg,
+                chunk_size=64,
+                return_final_states=True,
+                dt_bias=None,
+                dt_softplus=False,
+            )
+
+        torch.testing.assert_close(out, expected_out)
+        torch.testing.assert_close(state, expected_state)
+
+    def test_forces_return_final_states_when_cp_active(self):
+        """When CP is active, wrapper must set return_final_states=True."""
+        B, T, H, d, n = 1, 4, 2, 8, 4
+        captured_kwargs = {}
+
+        def fake_scan(*args, **kwargs):
+            captured_kwargs.update(kwargs)
+            scan_out = torch.zeros(B, T, H, d)
+            ssm_state = torch.zeros(B, H, d, n)
+            return scan_out, ssm_state
+
+        mod = self._make_module_with_scan(fake_scan)
+
+        with (
+            patch(
+                "axolotl.monkeypatch.models.mamba_utils.is_cp_active", return_value=True
+            ),
+            patch(
+                "axolotl.monkeypatch.models.mamba_utils.ring_shift_ssm_state",
+                side_effect=lambda h: torch.zeros_like(h),
+            ),
+        ):
+            wrap_mamba_scan_for_cp(mod)
+            mod.mamba_chunk_scan_combined(
+                torch.zeros(B, T, H, d),
+                torch.zeros(B, T, H),
+                -torch.ones(H),
+                torch.zeros(B, T, 1, n),
+                torch.zeros(B, T, 1, n),
+                chunk_size=64,
+                return_final_states=False,
+                dt_bias=None,
+                dt_softplus=False,
+            )
+
+        assert captured_kwargs["return_final_states"] is True
+
+    def test_idempotency_guard(self):
+        """Calling wrap_mamba_scan_for_cp twice must not double-wrap."""
+        call_count = 0
+
+        def fake_scan(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            B, T, H, d, n = 1, 4, 2, 8, 4
+            return torch.zeros(B, T, H, d), torch.zeros(B, H, d, n)
+
+        mod = self._make_module_with_scan(fake_scan)
+
+        with patch(
+            "axolotl.monkeypatch.models.mamba_utils.is_cp_active", return_value=False
+        ):
+            wrap_mamba_scan_for_cp(mod)
+            first_fn = mod.mamba_chunk_scan_combined
+            wrap_mamba_scan_for_cp(mod)
+            assert mod.mamba_chunk_scan_combined is first_fn
+            assert getattr(mod, "_cp_scan_wrapped", False) is True
diff --git a/tests/monkeypatch/test_pixtral_flash_attention_patch.py b/tests/monkeypatch/test_pixtral_flash_attention_patch.py
new file mode 100644
index 0000000000..285fde41e8
--- /dev/null
+++ b/tests/monkeypatch/test_pixtral_flash_attention_patch.py
@@ -0,0 +1,77 @@
+"""Integration tests for Pixtral Flash Attention patches."""
+
+import pytest
+import torch
+
+
+class TestPixtralFlashAttentionPatchIntegration:
+    """Test Pixtral Flash Attention patch integration."""
+
+    @pytest.mark.integration
+    def test_pixtral_flash_attention_patch(self):
+        """Test that Pixtral Flash Attention patch can be applied and works correctly."""
+        try:
+            from transformers import modeling_flash_attention_utils
+        except ImportError:
+            pytest.skip("Flash Attention utils not available")
+
+        from axolotl.monkeypatch.models.pixtral.modeling_flash_attention_utils import (
+            apply_patch_is_packed_sequence,
+        )
+
+        # Store original method
+        original_is_packed_sequence = modeling_flash_attention_utils._is_packed_sequence
+
+        # Apply patch and get unpatch function
+        unpatch_fn = apply_patch_is_packed_sequence()
+
+        # Verify patch was applied
+        assert (
+            modeling_flash_attention_utils._is_packed_sequence
+            != original_is_packed_sequence
+        ), "_is_packed_sequence was not patched"
+
+        # Test the patched function with 1D position_ids
+        patched_fn = modeling_flash_attention_utils._is_packed_sequence
+
+        # Test 1D position_ids 1 sequence
+        position_ids_1d = torch.tensor([0, 1, 2, 3])
+        result = patched_fn(position_ids_1d, batch_size=1)
+        assert isinstance(result, bool), "Function should return a boolean"
+        assert result is False, "1D sequential position_ids should not be packed"
+
+        # Test 1D packed 2 sequences
+        position_ids_1d_packed = torch.tensor([0, 1, 2, 0, 1, 2])
+        result = patched_fn(position_ids_1d_packed, batch_size=1)
+        assert isinstance(result, bool), "Function should return a boolean"
+        assert result is True, "1D packed position_ids should be detected as packed"
+
+        # Test 2D packed 2 sequences
+        position_ids_2d_packed = torch.tensor([[0, 1, 2, 3, 0, 1]])
+        result = patched_fn(position_ids_2d_packed, batch_size=1)
+        assert isinstance(result, bool), "Function should return a boolean"
+        assert result is True, "2D packed position_ids should be detected as packed"
+
+        # Test 2D 1 sequence
+        position_ids_2d_normal = torch.tensor([[0, 1, 2, 3, 4, 5]])
+        result = patched_fn(position_ids_2d_normal, batch_size=1)
+        assert isinstance(result, bool), "Function should return a boolean"
+        assert result is False, "2D sequential position_ids should not be packed"
+
+        # Test 2D batch size 2
+        position_ids_2d_normal = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8]])
+        result = patched_fn(position_ids_2d_normal, batch_size=2)
+        assert isinstance(result, bool), "Function should return a boolean"
+        assert result is False, "2D position_ids batch 2 should not be packed"
+
+        # Test None case
+        result = patched_fn(None, batch_size=1)
+        assert isinstance(result, bool), "Function should return a boolean"
+        assert result is False, "None position_ids should return False"
+
+        # Test unpatch function
+        unpatch_fn()
+        assert (
+            modeling_flash_attention_utils._is_packed_sequence
+            == original_is_packed_sequence
+        ), "unpatch function did not restore original method"
diff --git a/tests/monkeypatch/test_qwen3_5_fused_attn.py b/tests/monkeypatch/test_qwen3_5_fused_attn.py
new file mode 100644
index 0000000000..e258222e67
--- /dev/null
+++ b/tests/monkeypatch/test_qwen3_5_fused_attn.py
@@ -0,0 +1,424 @@
+"""Tests for the Qwen3.5 / Qwen3.5-MoE fused-attention monkeypatch."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+pytest.importorskip("transformers.models.qwen3_5")
+pytest.importorskip("transformers.models.qwen3_5_moe")
+
+
+def _clear_patched_flag(cls):
+    try:
+        delattr(cls, "_axolotl_fused_attn_patched")
+    except AttributeError:
+        pass
+
+
+@pytest.fixture
+def restore_qwen3_5_attention():
+    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Attention
+
+    saved = Qwen3_5Attention.forward
+    saved_flag = getattr(Qwen3_5Attention, "_axolotl_fused_attn_patched", False)
+    yield Qwen3_5Attention
+    Qwen3_5Attention.forward = saved
+    if saved_flag:
+        Qwen3_5Attention._axolotl_fused_attn_patched = saved_flag
+    else:
+        _clear_patched_flag(Qwen3_5Attention)
+
+
+def _build_qwen3_5_text_model(seed: int = 0):
+    from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
+    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel
+
+    torch.manual_seed(seed)
+    cfg = Qwen3_5TextConfig(
+        vocab_size=128,
+        hidden_size=128,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        max_position_embeddings=512,
+        rms_norm_eps=1e-6,
+        attention_dropout=0.0,
+        layer_types=["full_attention", "full_attention"],
+    )
+    cfg._attn_implementation = "sdpa"
+    return Qwen3_5TextModel(cfg).cuda().to(torch.bfloat16).eval()
+
+
+def _run_attention(model, layer_idx, hidden_states, position_ids):
+    attn = model.layers[layer_idx].self_attn
+    cos, sin = model.rotary_emb(hidden_states, position_ids)
+    out, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+    )
+    return out
+
+
+class TestQwen3_5FusedAttnParity:
+    """Single-layer parity vs stock."""
+
+    @pytest.mark.parametrize("layer_idx", [0, 1])
+    def test_forward_matches_stock(self, restore_qwen3_5_attention, layer_idx):
+        from axolotl.monkeypatch.models.qwen3_5.fused_attn import (
+            patch_qwen3_5_fused_attn,
+        )
+
+        m = _build_qwen3_5_text_model(seed=1)
+        hs = torch.randn(2, 16, 128, device="cuda", dtype=torch.bfloat16)
+        pos = torch.arange(16, device="cuda").unsqueeze(0).expand(2, -1)
+
+        with torch.no_grad():
+            ref = _run_attention(m, layer_idx, hs, pos)
+
+        patch_qwen3_5_fused_attn()
+        with torch.no_grad():
+            got = _run_attention(m, layer_idx, hs, pos)
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, (
+            f"layer {layer_idx} fused vs stock cosine_sim={cos_sim:.6f}"
+        )
+        torch.testing.assert_close(got, ref, rtol=5e-2, atol=5e-2)
+
+
+class TestQwen3_5FusedAttnBackward:
+    def test_q_k_norm_grads_finite_nonzero(self, restore_qwen3_5_attention):
+        from axolotl.monkeypatch.models.qwen3_5.fused_attn import (
+            patch_qwen3_5_fused_attn,
+        )
+
+        m = _build_qwen3_5_text_model(seed=3).train()
+        patch_qwen3_5_fused_attn()
+
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+        out = m(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state
+        out.sum().backward()
+
+        for i, layer in enumerate(m.layers[:2]):
+            if m.config.layer_types[i] != "full_attention":
+                continue
+            attn = layer.self_attn
+            assert attn.q_norm.weight.grad is not None, f"layer {i} q_norm no grad"
+            assert attn.k_norm.weight.grad is not None, f"layer {i} k_norm no grad"
+            assert attn.q_norm.weight.grad.isfinite().all()
+            assert attn.k_norm.weight.grad.isfinite().all()
+            assert attn.q_norm.weight.grad.abs().sum() > 0
+            assert attn.k_norm.weight.grad.abs().sum() > 0
+
+
+class TestQwen3_5FusedAttnLoRACompose:
+    """Pin LoRA-QKV → fused composition; ``QKV_PATCHES`` includes a chunk-2 variant for Qwen3.5's ``q_proj * 2``."""
+
+    def _build_cfg(self):
+        from axolotl.utils.dict import DictDefault
+
+        return DictDefault(
+            {
+                "base_model": "fake/qwen3_5",
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_dropout": 0.0,
+            }
+        )
+
+    def test_lora_qkv_then_fused_does_not_raise(
+        self, restore_qwen3_5_attention, monkeypatch
+    ):
+        from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Attention
+
+        from axolotl.monkeypatch import lora_kernels
+        from axolotl.monkeypatch.models.qwen3_5.fused_attn import (
+            patch_qwen3_5_fused_attn,
+        )
+
+        monkeypatch.setattr(
+            lora_kernels,
+            "get_attention_cls_from_config",
+            lambda _cfg: Qwen3_5Attention,
+        )
+
+        try:
+            delattr(Qwen3_5Attention, "_original_forward")
+        except AttributeError:
+            pass
+
+        try:
+            lora_kernels.patch_self_attn_lora(self._build_cfg())
+            assert hasattr(Qwen3_5Attention, "_original_forward"), (
+                "patch_self_attn_lora must capture the stock Qwen3.5 forward — if "
+                "this fails, QKV_PATCHES drifted away from the chunk-2 q_proj source"
+            )
+            patch_qwen3_5_fused_attn()
+            assert getattr(Qwen3_5Attention, "_axolotl_fused_attn_patched", False)
+        finally:
+            try:
+                delattr(Qwen3_5Attention, "_original_forward")
+            except AttributeError:
+                pass
+
+
+@pytest.fixture
+def restore_qwen3_5_moe_attention():
+    from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
+        Qwen3_5MoeAttention,
+    )
+
+    saved = Qwen3_5MoeAttention.forward
+    saved_flag = getattr(Qwen3_5MoeAttention, "_axolotl_fused_attn_patched", False)
+    yield Qwen3_5MoeAttention
+    Qwen3_5MoeAttention.forward = saved
+    if saved_flag:
+        Qwen3_5MoeAttention._axolotl_fused_attn_patched = saved_flag
+    else:
+        _clear_patched_flag(Qwen3_5MoeAttention)
+
+
+def _build_qwen3_5_moe_model(seed: int = 0):
+    from transformers.models.qwen3_5_moe.configuration_qwen3_5_moe import (
+        Qwen3_5MoeTextConfig,
+    )
+    from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
+        Qwen3_5MoeTextModel,
+    )
+
+    torch.manual_seed(seed)
+    cfg = Qwen3_5MoeTextConfig(
+        vocab_size=128,
+        hidden_size=128,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        max_position_embeddings=512,
+        rms_norm_eps=1e-6,
+        attention_dropout=0.0,
+        layer_types=["full_attention", "full_attention"],
+        num_experts=4,
+        num_experts_per_tok=2,
+        moe_intermediate_size=64,
+        shared_expert_intermediate_size=64,
+    )
+    cfg._attn_implementation = "sdpa"
+    return Qwen3_5MoeTextModel(cfg).cuda().to(torch.bfloat16).eval()
+
+
+class TestQwen3_5MoeFusedAttnParity:
+    """End-to-end parity on the MoE variant (attention is structurally identical to dense Qwen3.5)."""
+
+    def test_forward_matches_stock(self, restore_qwen3_5_moe_attention):
+        from axolotl.monkeypatch.models.qwen3_5_moe.fused_attn import (
+            patch_qwen3_5_moe_fused_attn,
+        )
+
+        m = _build_qwen3_5_moe_model(seed=5)
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+
+        with torch.no_grad():
+            ref = m(
+                input_ids=ids, attention_mask=mask, use_cache=False
+            ).last_hidden_state.clone()
+
+        patch_qwen3_5_moe_fused_attn()
+        with torch.no_grad():
+            got = m(
+                input_ids=ids, attention_mask=mask, use_cache=False
+            ).last_hidden_state.clone()
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"qwen3_5_moe end-to-end cosine_sim={cos_sim:.6f}"
+
+
+class TestQwen3_5MoeFusedAttnBackward:
+    """Backward grad flow through the fused Q/K-norm kernels on Qwen3.5-MoE."""
+
+    def test_q_k_norm_grads_finite_nonzero(self, restore_qwen3_5_moe_attention):
+        from axolotl.monkeypatch.models.qwen3_5_moe.fused_attn import (
+            patch_qwen3_5_moe_fused_attn,
+        )
+
+        m = _build_qwen3_5_moe_model(seed=6).train()
+        patch_qwen3_5_moe_fused_attn()
+
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+        out = m(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state
+        out.sum().backward()
+
+        for i, layer in enumerate(m.layers[:2]):
+            if m.config.layer_types[i] != "full_attention":
+                continue
+            attn = layer.self_attn
+            assert attn.q_norm.weight.grad is not None, f"layer {i} q_norm no grad"
+            assert attn.k_norm.weight.grad is not None, f"layer {i} k_norm no grad"
+            assert attn.q_norm.weight.grad.isfinite().all()
+            assert attn.k_norm.weight.grad.isfinite().all()
+            assert attn.q_norm.weight.grad.abs().sum() > 0
+            assert attn.k_norm.weight.grad.abs().sum() > 0
+
+
+class TestQwen3_5MoeFusedAttnLoRACompose:
+    """MoE mirror of the Qwen3.5 LoRA-compose test."""
+
+    def _build_cfg(self):
+        from axolotl.utils.dict import DictDefault
+
+        return DictDefault(
+            {
+                "base_model": "fake/qwen3_5_moe",
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_dropout": 0.0,
+            }
+        )
+
+    def test_lora_qkv_then_fused_does_not_raise(
+        self, restore_qwen3_5_moe_attention, monkeypatch
+    ):
+        from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
+            Qwen3_5MoeAttention,
+        )
+
+        from axolotl.monkeypatch import lora_kernels
+        from axolotl.monkeypatch.models.qwen3_5_moe.fused_attn import (
+            patch_qwen3_5_moe_fused_attn,
+        )
+
+        monkeypatch.setattr(
+            lora_kernels,
+            "get_attention_cls_from_config",
+            lambda _cfg: Qwen3_5MoeAttention,
+        )
+
+        try:
+            delattr(Qwen3_5MoeAttention, "_original_forward")
+        except AttributeError:
+            pass
+
+        try:
+            lora_kernels.patch_self_attn_lora(self._build_cfg())
+            assert hasattr(Qwen3_5MoeAttention, "_original_forward")
+            patch_qwen3_5_moe_fused_attn()
+            assert getattr(Qwen3_5MoeAttention, "_axolotl_fused_attn_patched", False)
+        finally:
+            try:
+                delattr(Qwen3_5MoeAttention, "_original_forward")
+            except AttributeError:
+                pass
+
+
+class TestQwen3_5FusedAttnLigerRMSNormCompose:
+    """Liger swaps ``Qwen3_5RMSNorm`` for a subclass that exposes ``variance_epsilon`` instead of ``eps``."""
+
+    def test_forward_survives_liger_rmsnorm_swap(self, restore_qwen3_5_attention):
+        from axolotl.monkeypatch.models.qwen3_5.fused_attn import (
+            patch_qwen3_5_fused_attn,
+        )
+
+        m = _build_qwen3_5_text_model(seed=9)
+
+        class _StubRMSNormVarEps(torch.nn.Module):
+            def __init__(self, original):
+                super().__init__()
+                self.weight = original.weight
+                self.variance_epsilon = original.eps
+
+            def forward(self, x):
+                return x
+
+        for layer in m.layers:
+            if not hasattr(layer, "self_attn"):
+                continue
+            attn = layer.self_attn
+            attn.q_norm = _StubRMSNormVarEps(attn.q_norm)
+            attn.k_norm = _StubRMSNormVarEps(attn.k_norm)
+
+        patch_qwen3_5_fused_attn()
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+        with torch.no_grad():
+            out = m(input_ids=ids, attention_mask=mask, use_cache=False)
+        assert torch.isfinite(out.last_hidden_state).all()
+
+
+class TestPatchManagerQwen3_5TextDispatch:
+    """Pin that ``_apply_model_specific_patches`` covers the ``*_text`` config types of multimodal Qwen3.5 / Qwen3.5-MoE checkpoints."""
+
+    @pytest.mark.parametrize("model_config_type", ["qwen3_5", "qwen3_5_text"])
+    def test_qwen3_5_text_variant_is_patched(
+        self, restore_qwen3_5_attention, model_config_type
+    ):
+        from axolotl.loaders.patch_manager import PatchManager
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault(
+            {
+                "base_model": "fake/qwen3_5",
+                "model_config_type": model_config_type,
+                "fused_attn_kernel": True,
+                "lora_qkv_kernel": False,
+                "lora_o_kernel": False,
+                "context_parallel_size": 1,
+            }
+        )
+        mc = type("MC", (), {"model_type": model_config_type})()
+        pm = PatchManager(cfg=cfg, model_config=mc, inference=False)
+        pm._apply_model_specific_patches()
+
+        from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Attention
+
+        assert getattr(Qwen3_5Attention, "_axolotl_fused_attn_patched", False), (
+            f"PatchManager skipped fused-attn for model_config_type="
+            f"{model_config_type!r}; dispatch is missing the _text variant"
+        )
+
+    @pytest.mark.parametrize("model_config_type", ["qwen3_5_moe", "qwen3_5_moe_text"])
+    def test_qwen3_5_moe_text_variant_is_patched(
+        self, restore_qwen3_5_moe_attention, model_config_type
+    ):
+        from axolotl.loaders.patch_manager import PatchManager
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault(
+            {
+                "base_model": "fake/qwen3_5_moe",
+                "model_config_type": model_config_type,
+                "fused_attn_kernel": True,
+                "lora_qkv_kernel": False,
+                "lora_o_kernel": False,
+                "context_parallel_size": 1,
+            }
+        )
+        mc = type("MC", (), {"model_type": model_config_type})()
+        pm = PatchManager(cfg=cfg, model_config=mc, inference=False)
+        pm._apply_model_specific_patches()
+
+        from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
+            Qwen3_5MoeAttention,
+        )
+
+        assert getattr(Qwen3_5MoeAttention, "_axolotl_fused_attn_patched", False), (
+            f"PatchManager skipped fused-attn for model_config_type="
+            f"{model_config_type!r}; dispatch is missing the _text variant"
+        )
diff --git a/tests/monkeypatch/test_qwen3_fused_attn.py b/tests/monkeypatch/test_qwen3_fused_attn.py
new file mode 100644
index 0000000000..7056eb2939
--- /dev/null
+++ b/tests/monkeypatch/test_qwen3_fused_attn.py
@@ -0,0 +1,283 @@
+"""Tests for the Qwen3 / Qwen3-MoE fused-attention monkeypatch."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+pytest.importorskip("transformers.models.qwen3")
+pytest.importorskip("transformers.models.qwen3_moe")
+
+
+def _clear_patched_flag(cls):
+    try:
+        delattr(cls, "_axolotl_fused_attn_patched")
+    except AttributeError:
+        pass
+
+
+@pytest.fixture
+def restore_qwen3_attention():
+    from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+    saved = Qwen3Attention.forward
+    saved_flag = getattr(Qwen3Attention, "_axolotl_fused_attn_patched", False)
+    yield Qwen3Attention
+    Qwen3Attention.forward = saved
+    if saved_flag:
+        Qwen3Attention._axolotl_fused_attn_patched = saved_flag
+    else:
+        _clear_patched_flag(Qwen3Attention)
+
+
+@pytest.fixture
+def restore_qwen3_moe_attention():
+    from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeAttention
+
+    saved = Qwen3MoeAttention.forward
+    saved_flag = getattr(Qwen3MoeAttention, "_axolotl_fused_attn_patched", False)
+    yield Qwen3MoeAttention
+    Qwen3MoeAttention.forward = saved
+    if saved_flag:
+        Qwen3MoeAttention._axolotl_fused_attn_patched = saved_flag
+    else:
+        _clear_patched_flag(Qwen3MoeAttention)
+
+
+def _build_qwen3_model(seed: int = 0):
+    from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+    from transformers.models.qwen3.modeling_qwen3 import Qwen3Model
+
+    torch.manual_seed(seed)
+    cfg = Qwen3Config(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        max_position_embeddings=2048,
+        rms_norm_eps=1e-6,
+        attention_dropout=0.0,
+    )
+    cfg._attn_implementation = "sdpa"
+    return Qwen3Model(cfg).cuda().to(torch.bfloat16).eval()
+
+
+def _build_qwen3_moe_model(seed: int = 0):
+    from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
+    from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeModel
+
+    torch.manual_seed(seed)
+    cfg = Qwen3MoeConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        max_position_embeddings=512,
+        rms_norm_eps=1e-6,
+        attention_dropout=0.0,
+        num_experts=4,
+        num_experts_per_tok=2,
+        moe_intermediate_size=64,
+    )
+    cfg._attn_implementation = "sdpa"
+    return Qwen3MoeModel(cfg).cuda().to(torch.bfloat16).eval()
+
+
+def _run_attention(model, layer_idx, hidden_states, position_ids):
+    attn = model.layers[layer_idx].self_attn
+    cos, sin = model.rotary_emb(hidden_states, position_ids)
+    out, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+    )
+    return out
+
+
+class TestQwen3FusedAttnParity:
+    """Single-layer parity vs stock."""
+
+    @pytest.mark.parametrize("layer_idx", [0, 1])
+    def test_forward_matches_stock(self, restore_qwen3_attention, layer_idx):
+        from axolotl.monkeypatch.models.qwen3.fused_attn import patch_qwen3_fused_attn
+
+        m = _build_qwen3_model(seed=1)
+        hs = torch.randn(2, 16, 64, device="cuda", dtype=torch.bfloat16)
+        pos = torch.arange(16, device="cuda").unsqueeze(0).expand(2, -1)
+
+        with torch.no_grad():
+            ref = _run_attention(m, layer_idx, hs, pos)
+
+        patch_qwen3_fused_attn()
+        with torch.no_grad():
+            got = _run_attention(m, layer_idx, hs, pos)
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, (
+            f"layer {layer_idx} fused vs stock cosine_sim={cos_sim:.6f}"
+        )
+        torch.testing.assert_close(got, ref, rtol=5e-2, atol=5e-2)
+
+
+class TestQwen3FusedAttnEndToEnd:
+    def test_full_forward_matches_stock(self, restore_qwen3_attention):
+        from axolotl.monkeypatch.models.qwen3.fused_attn import patch_qwen3_fused_attn
+
+        m = _build_qwen3_model(seed=2)
+        ids = torch.randint(0, 128, (2, 32), device="cuda")
+        mask = torch.ones(2, 32, dtype=torch.long, device="cuda")
+
+        with torch.no_grad():
+            ref = m(input_ids=ids, attention_mask=mask).last_hidden_state.clone()
+
+        patch_qwen3_fused_attn()
+        with torch.no_grad():
+            got = m(input_ids=ids, attention_mask=mask).last_hidden_state.clone()
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"end-to-end cosine_sim={cos_sim:.6f}"
+
+    def test_backward_grad_flows_through_fused_path(self, restore_qwen3_attention):
+        from axolotl.monkeypatch.models.qwen3.fused_attn import patch_qwen3_fused_attn
+
+        m = _build_qwen3_model(seed=3).train()
+        patch_qwen3_fused_attn()
+
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+        out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+        out.sum().backward()
+
+        for i, layer in enumerate(m.layers[:2]):
+            attn = layer.self_attn
+            assert attn.q_norm.weight.grad is not None, f"layer {i} q_norm no grad"
+            assert attn.k_norm.weight.grad is not None, f"layer {i} k_norm no grad"
+            assert attn.q_norm.weight.grad.isfinite().all()
+            assert attn.k_norm.weight.grad.isfinite().all()
+            assert attn.q_norm.weight.grad.abs().sum() > 0
+            assert attn.k_norm.weight.grad.abs().sum() > 0
+
+
+class TestQwen3FusedAttnLoRACompose:
+    """Pin that LoRA-QKV runs before the fused patch (``inspect.getsource`` regex misses on the fused body)."""
+
+    def test_lora_qkv_then_fused_does_not_raise(self, restore_qwen3_attention):
+        from axolotl.monkeypatch.lora_kernels import patch_self_attn_lora
+        from axolotl.monkeypatch.models.qwen3.fused_attn import patch_qwen3_fused_attn
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault(
+            {
+                "base_model": "Qwen/Qwen3-0.6B",
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_dropout": 0.0,
+            }
+        )
+
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+        try:
+            delattr(Qwen3Attention, "_original_forward")
+        except AttributeError:
+            pass
+
+        try:
+            patch_self_attn_lora(cfg)
+            assert hasattr(Qwen3Attention, "_original_forward"), (
+                "patch_self_attn_lora must run on stock source first"
+            )
+            patch_qwen3_fused_attn()
+            assert getattr(Qwen3Attention, "_axolotl_fused_attn_patched", False)
+        finally:
+            try:
+                delattr(Qwen3Attention, "_original_forward")
+            except AttributeError:
+                pass
+
+    def test_reverse_order_breaks(self, restore_qwen3_attention):
+        from axolotl.monkeypatch.lora_kernels import patch_self_attn_lora
+        from axolotl.monkeypatch.models.qwen3.fused_attn import patch_qwen3_fused_attn
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault(
+            {
+                "base_model": "Qwen/Qwen3-0.6B",
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "lora_dropout": 0.0,
+            }
+        )
+
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+        try:
+            delattr(Qwen3Attention, "_original_forward")
+        except AttributeError:
+            pass
+
+        patch_qwen3_fused_attn()
+        with pytest.raises(AssertionError, match="Original QKV code not found"):
+            patch_self_attn_lora(cfg)
+
+
+class TestPatchManagerOrdering:
+    """Pin the patch-manager ordering invariant."""
+
+    def test_self_attn_lora_runs_before_model_specific(self):
+        import inspect
+
+        from axolotl.loaders.patch_manager import PatchManager
+
+        src = inspect.getsource(PatchManager.apply_pre_model_load_patches)
+        lora_idx = src.find("_apply_self_attention_lora_patch()")
+        specific_idx = src.find("_apply_model_specific_patches()")
+        assert lora_idx > 0 and specific_idx > 0
+        assert lora_idx < specific_idx, (
+            "_apply_self_attention_lora_patch must run before "
+            "_apply_model_specific_patches so patch_self_attn_lora sees the "
+            "stock attention forward source"
+        )
+
+
+class TestQwen3MoeFusedAttnParity:
+    """End-to-end parity on the MoE variant (attention is structurally identical to dense Qwen3)."""
+
+    def test_full_forward_matches_stock(self, restore_qwen3_moe_attention):
+        from axolotl.monkeypatch.models.qwen3_moe.fused_attn import (
+            patch_qwen3_moe_fused_attn,
+        )
+
+        m = _build_qwen3_moe_model(seed=4)
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+
+        with torch.no_grad():
+            ref = m(input_ids=ids, attention_mask=mask).last_hidden_state.clone()
+
+        patch_qwen3_moe_fused_attn()
+        with torch.no_grad():
+            got = m(input_ids=ids, attention_mask=mask).last_hidden_state.clone()
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"qwen3-moe end-to-end cosine_sim={cos_sim:.6f}"
diff --git a/tests/monkeypatch/test_qwen3_fused_attn_defensive.py b/tests/monkeypatch/test_qwen3_fused_attn_defensive.py
new file mode 100644
index 0000000000..34b22b8ee3
--- /dev/null
+++ b/tests/monkeypatch/test_qwen3_fused_attn_defensive.py
@@ -0,0 +1,375 @@
+"""Defensive regression tests for edge cases of the fused-attn patches."""
+
+import inspect
+import logging
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+pytest.importorskip("transformers.models.qwen3")
+
+
+def _clear_patched_flag(cls):
+    try:
+        delattr(cls, "_axolotl_fused_attn_patched")
+    except AttributeError:
+        pass
+
+
+@pytest.fixture
+def restore_qwen3_attention():
+    from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+    saved = Qwen3Attention.forward
+    saved_flag = getattr(Qwen3Attention, "_axolotl_fused_attn_patched", False)
+    yield Qwen3Attention
+    Qwen3Attention.forward = saved
+    if saved_flag:
+        Qwen3Attention._axolotl_fused_attn_patched = saved_flag
+    else:
+        _clear_patched_flag(Qwen3Attention)
+
+
+@pytest.fixture
+def patch_manager_caplog(caplog):
+    logger = logging.getLogger("axolotl.loaders.patch_manager")
+    logger.addHandler(caplog.handler)
+    previous_level = logger.level
+    logger.setLevel(logging.DEBUG)
+    try:
+        yield caplog
+    finally:
+        logger.removeHandler(caplog.handler)
+        logger.setLevel(previous_level)
+
+
+class TestFusedAttnKernelUnsupportedWarning:
+    """The warning lives in ``PatchManager`` (not the schema validator) so it
+    runs after ``normalize_config()`` has derived ``model_config_type``. A
+    normal CLI flow with ``fused_attn_kernel: true`` on e.g. a Llama config
+    must warn loudly instead of silently no-op'ing."""
+
+    def test_warns_on_unsupported_model_type(self, patch_manager_caplog):
+        from types import SimpleNamespace
+
+        from axolotl.loaders.patch_manager import PatchManager
+
+        cfg = SimpleNamespace(fused_attn_kernel=True, model_config_type="llama")
+        PatchManager._warn_if_fused_attn_unsupported(cfg)
+        assert (
+            "fused_attn_kernel" in patch_manager_caplog.text
+            and "llama" in patch_manager_caplog.text
+        ), f"expected warning about llama; got {patch_manager_caplog.text}"
+
+    @pytest.mark.parametrize(
+        "model_type",
+        [
+            "qwen3",
+            "qwen3_moe",
+            "qwen3_vl",
+            "qwen3_vl_text",
+            "qwen3_5",
+            "qwen3_5_text",
+            "qwen3_5_moe",
+            "qwen3_5_moe_text",
+            "gemma4",
+            "gemma4_text",
+        ],
+    )
+    def test_no_warn_on_supported_model_type(self, patch_manager_caplog, model_type):
+        from types import SimpleNamespace
+
+        from axolotl.loaders.patch_manager import PatchManager
+
+        cfg = SimpleNamespace(fused_attn_kernel=True, model_config_type=model_type)
+        PatchManager._warn_if_fused_attn_unsupported(cfg)
+        assert not any(
+            "fused_attn_kernel" in r.message and model_type in r.message
+            for r in patch_manager_caplog.records
+        ), f"unexpected warning for supported {model_type}"
+
+    def test_no_warn_when_fused_attn_kernel_false(self, patch_manager_caplog):
+        from types import SimpleNamespace
+
+        from axolotl.loaders.patch_manager import PatchManager
+
+        cfg = SimpleNamespace(fused_attn_kernel=False, model_config_type="llama")
+        PatchManager._warn_if_fused_attn_unsupported(cfg)
+        assert not any(
+            "fused_attn_kernel" in r.message for r in patch_manager_caplog.records
+        ), "no warning expected when fused_attn_kernel is False"
+
+    def test_warning_is_invoked_by_apply_model_specific_patches(self):
+        """Source-line check that ``_apply_model_specific_patches`` actually
+        calls ``_warn_if_fused_attn_unsupported``. Without this, the standalone
+        helper passes its unit tests but never runs in practice."""
+        import inspect
+
+        from axolotl.loaders.patch_manager import PatchManager
+
+        src = inspect.getsource(PatchManager._apply_model_specific_patches)
+        assert "_warn_if_fused_attn_unsupported" in src, (
+            "_apply_model_specific_patches no longer invokes "
+            "_warn_if_fused_attn_unsupported — the warning will be dead code"
+        )
+
+
+class TestPeftModulesToSaveWrapper:
+    """``modules_to_save=["q_norm","k_norm"]`` wraps the norms in ``ModulesToSaveWrapper``; the patched forward must resolve through it."""
+
+    def _make_wrapper(self, original):
+        def _make_clone(m):
+            clone = torch.nn.Module()
+            clone.weight = torch.nn.Parameter(m.weight.detach().clone())
+            clone.variance_epsilon = m.variance_epsilon
+            clone.eps = getattr(m, "eps", m.variance_epsilon)
+            return clone
+
+        class _StubWrapper(torch.nn.Module):
+            """Mirrors PEFT ``ModulesToSaveWrapper``: ``_active_adapter`` is a
+            ``list[str]`` and ``active_adapter`` / ``active_adapters`` are
+            properties returning that list."""
+
+            def __init__(self, orig):
+                super().__init__()
+                self.original_module = orig
+                self.modules_to_save = torch.nn.ModuleDict(
+                    {"default": _make_clone(orig)}
+                )
+                self._active_adapter = ["default"]
+
+            @property
+            def active_adapter(self):
+                return self._active_adapter
+
+            @property
+            def active_adapters(self):
+                return self._active_adapter
+
+            def forward(self, x):
+                return self.modules_to_save[self._active_adapter[0]](x)
+
+        return _StubWrapper(original)
+
+    def test_resolve_norm_module_returns_active_adapter_not_original(self):
+        """Direct unit-test of ``_resolve_norm_module``: with PEFT's actual
+        ``active_adapter = ["default"]`` (list), the helper must return the
+        wrapped adapter module, not the frozen ``original_module``. The earlier
+        ``isinstance(adapter, str)`` check silently failed this case."""
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3RMSNorm
+
+        from axolotl.monkeypatch.models.qwen3.fused_attn import (
+            _resolve_norm_module,
+        )
+
+        orig = Qwen3RMSNorm(16, eps=1e-6)
+        wrapper = self._make_wrapper(orig)
+        resolved = _resolve_norm_module(wrapper)
+        assert resolved is wrapper.modules_to_save["default"], (
+            "_resolve_norm_module returned the frozen original instead of the "
+            "active adapter — PEFT stores active_adapter as a list, not a str"
+        )
+
+    def test_resolve_through_real_peft_modules_to_save(self):
+        """End-to-end: build a Qwen3 model, wrap ``q_norm`` / ``k_norm`` with
+        ``peft.get_peft_model(..., modules_to_save=[...])``, set the active
+        adapter's weight to a value distinct from the frozen original, and
+        confirm ``_resolve_norm_module`` returns the active-adapter module
+        (so the fused kernel reads the trainable weight, not the frozen one).
+        This exercises the real PEFT object shape, not a stub."""
+        from peft import LoraConfig, get_peft_model
+        from peft.utils.other import ModulesToSaveWrapper
+        from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM
+
+        from axolotl.monkeypatch.models.qwen3.fused_attn import (
+            _resolve_norm_module,
+        )
+
+        cfg = Qwen3Config(
+            vocab_size=128,
+            hidden_size=64,
+            intermediate_size=128,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            head_dim=16,
+            max_position_embeddings=256,
+            rms_norm_eps=1e-6,
+            attention_dropout=0.0,
+        )
+        base = Qwen3ForCausalLM(cfg)
+        lora = LoraConfig(
+            r=4,
+            lora_alpha=8,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+            modules_to_save=["q_norm", "k_norm"],
+            task_type="CAUSAL_LM",
+        )
+        peft_model = get_peft_model(base, lora)
+
+        attn = peft_model.base_model.model.model.layers[0].self_attn
+        assert isinstance(attn.q_norm, ModulesToSaveWrapper), (
+            "PEFT did not wrap q_norm — test premise is invalid"
+        )
+        with torch.no_grad():
+            attn.q_norm.modules_to_save["default"].weight.fill_(7.0)
+            attn.q_norm.original_module.weight.fill_(0.0)
+
+        resolved = _resolve_norm_module(attn.q_norm)
+        assert resolved is attn.q_norm.modules_to_save["default"], (
+            "_resolve_norm_module did not return the active-adapter module — "
+            "real PEFT exposes active_adapter as a list, but the helper "
+            "treated only the str case"
+        )
+        assert torch.equal(
+            resolved.weight.detach(),
+            torch.full_like(resolved.weight, 7.0),
+        ), "resolved module is not the trainable adapter weight"
+
+    def test_qwen3_forward_under_modules_to_save_wrapper(self, restore_qwen3_attention):
+        from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3Model
+
+        from axolotl.monkeypatch.models.qwen3.fused_attn import (
+            patch_qwen3_fused_attn,
+        )
+
+        cfg = Qwen3Config(
+            vocab_size=128,
+            hidden_size=64,
+            intermediate_size=128,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            head_dim=16,
+            max_position_embeddings=256,
+            rms_norm_eps=1e-6,
+            attention_dropout=0.0,
+        )
+        cfg._attn_implementation = "sdpa"
+        m = Qwen3Model(cfg).cuda().to(torch.bfloat16)
+
+        for layer in m.layers:
+            attn = layer.self_attn
+            attn.q_norm = self._make_wrapper(attn.q_norm).cuda().to(torch.bfloat16)
+            attn.k_norm = self._make_wrapper(attn.k_norm).cuda().to(torch.bfloat16)
+
+        patch_qwen3_fused_attn()
+        ids = torch.randint(0, 128, (1, 16), device="cuda")
+        mask = torch.ones(1, 16, dtype=torch.long, device="cuda")
+        with torch.no_grad():
+            out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+        assert torch.isfinite(out).all(), (
+            "fused forward through ModulesToSaveWrapper produced non-finite output"
+        )
+
+
+class TestKernelProductionHeadDim:
+    """Kernel parity at ``head_dim=256`` (Qwen3.5 production); unit tests only cover 32/64."""
+
+    @pytest.mark.parametrize("head_dim", [128, 256])
+    @pytest.mark.parametrize("unit_offset", [False, True])
+    def test_fused_rms_norm_rope_parity(self, head_dim, unit_offset):
+        from axolotl.kernels.gemma4_fused_rope import fused_rms_norm_rope
+
+        B, S, H, D = 2, 32, 4, head_dim
+        torch.manual_seed(11)
+        x = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
+        w = torch.randn(D, device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16)
+        eps = 1e-6
+
+        x32 = x.to(torch.float32)
+        rms = x32.pow(2).mean(-1, keepdim=True).add(eps).rsqrt()
+        scale = (w.to(torch.float32) + 1.0) if unit_offset else w.to(torch.float32)
+        x_norm = x32 * rms * scale
+        half = D // 2
+        rot = torch.cat([-x_norm[..., half:], x_norm[..., :half]], dim=-1)
+        ref = (
+            x_norm * cos.to(torch.float32).unsqueeze(2)
+            + rot * sin.to(torch.float32).unsqueeze(2)
+        ).to(torch.bfloat16)
+
+        got = fused_rms_norm_rope(x, w, cos, sin, eps=eps, unit_offset=unit_offset)
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        torch.testing.assert_close(got, ref, rtol=5e-2, atol=5e-2)
+
+
+class TestAttentionMaskPassThrough:
+    """Sample-packing masks must flow through the patched forward verbatim."""
+
+    def test_qwen3_padding_mask_runs_clean(self, restore_qwen3_attention):
+        from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3Model
+
+        from axolotl.monkeypatch.models.qwen3.fused_attn import (
+            patch_qwen3_fused_attn,
+        )
+
+        cfg = Qwen3Config(
+            vocab_size=128,
+            hidden_size=64,
+            intermediate_size=128,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            head_dim=16,
+            max_position_embeddings=256,
+            rms_norm_eps=1e-6,
+            attention_dropout=0.0,
+        )
+        cfg._attn_implementation = "sdpa"
+        m = Qwen3Model(cfg).cuda().to(torch.bfloat16)
+        patch_qwen3_fused_attn()
+
+        ids = torch.randint(0, 128, (2, 16), device="cuda")
+        mask = torch.cat([torch.ones(2, 8), torch.zeros(2, 8)], dim=1).long().cuda()
+        with torch.no_grad():
+            out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+        assert torch.isfinite(out).all()
+
+
+class TestSlidingWindowKwarg:
+    """The fused forward must preserve ``sliding_window`` on the attention-interface call."""
+
+    def test_fused_forward_passes_sliding_window(self):
+        from axolotl.monkeypatch.models.qwen3 import fused_attn
+
+        src = inspect.getsource(fused_attn._make_fused_forward)
+        assert "sliding_window=self.sliding_window" in src, (
+            "Qwen3 fused_forward must pass sliding_window to attention_interface "
+            "to preserve sliding-attention layer behavior"
+        )
+
+
+class TestGetTextConfigDispatch:
+    """A multimodal Qwen3-VL text branch surfaces as ``model_config_type='qwen3'``; the patch must still fire."""
+
+    def test_qwen3_text_branch_dispatch(self, restore_qwen3_attention):
+        from axolotl.loaders.patch_manager import PatchManager
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault(
+            {
+                "base_model": "fake/qwen3-vl-text-branch",
+                "model_config_type": "qwen3",
+                "fused_attn_kernel": True,
+                "lora_qkv_kernel": False,
+                "lora_o_kernel": False,
+                "context_parallel_size": 1,
+            }
+        )
+        mc = type("MC", (), {"model_type": "qwen3"})()
+        pm = PatchManager(cfg=cfg, model_config=mc, inference=False)
+        pm._apply_model_specific_patches()
+
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+        assert getattr(Qwen3Attention, "_axolotl_fused_attn_patched", False)
diff --git a/tests/monkeypatch/test_qwen3_fused_attn_robustness.py b/tests/monkeypatch/test_qwen3_fused_attn_robustness.py
new file mode 100644
index 0000000000..2911aaeb72
--- /dev/null
+++ b/tests/monkeypatch/test_qwen3_fused_attn_robustness.py
@@ -0,0 +1,206 @@
+"""Robustness tests for the Qwen3 / Qwen3.5 fused-attn patches (idempotency, signature drift, GC, cross-device, FA2)."""
+
+import inspect
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+pytest.importorskip("transformers.models.qwen3")
+
+
+def _clear_patched_flag(cls):
+    try:
+        delattr(cls, "_axolotl_fused_attn_patched")
+    except AttributeError:
+        pass
+
+
+def _saved_state(cls):
+    return (cls.forward, getattr(cls, "_axolotl_fused_attn_patched", False))
+
+
+def _restore_state(cls, state):
+    saved_forward, saved_flag = state
+    cls.forward = saved_forward
+    if saved_flag:
+        cls._axolotl_fused_attn_patched = saved_flag
+    else:
+        _clear_patched_flag(cls)
+
+
+@pytest.fixture
+def restore_qwen3_attention():
+    from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+    state = _saved_state(Qwen3Attention)
+    yield Qwen3Attention
+    _restore_state(Qwen3Attention, state)
+
+
+def _build_tiny_qwen3():
+    from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+    from transformers.models.qwen3.modeling_qwen3 import Qwen3Model
+
+    cfg = Qwen3Config(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        max_position_embeddings=512,
+        rms_norm_eps=1e-6,
+        attention_dropout=0.0,
+    )
+    cfg._attn_implementation = "sdpa"
+    return Qwen3Model(cfg).cuda().to(torch.bfloat16)
+
+
+class TestPatchIdempotency:
+    """Re-applying the patch must be a no-op (``_axolotl_fused_attn_patched`` flag guard)."""
+
+    def test_qwen3_double_patch_is_noop(self, restore_qwen3_attention):
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+        from axolotl.monkeypatch.models.qwen3.fused_attn import (
+            patch_qwen3_fused_attn,
+        )
+
+        patch_qwen3_fused_attn()
+        forward_after_first = Qwen3Attention.forward
+        assert Qwen3Attention._axolotl_fused_attn_patched is True
+
+        patch_qwen3_fused_attn()
+        assert Qwen3Attention.forward is forward_after_first, (
+            "second patch_qwen3_fused_attn() call replaced .forward — the "
+            "_axolotl_fused_attn_patched guard is broken"
+        )
+
+    def test_qwen3_5_double_patch_is_noop(self):
+        pytest.importorskip("transformers.models.qwen3_5")
+        from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Attention
+
+        from axolotl.monkeypatch.models.qwen3_5.fused_attn import (
+            patch_qwen3_5_fused_attn,
+        )
+
+        state = _saved_state(Qwen3_5Attention)
+        try:
+            patch_qwen3_5_fused_attn()
+            forward_after_first = Qwen3_5Attention.forward
+            patch_qwen3_5_fused_attn()
+            assert Qwen3_5Attention.forward is forward_after_first
+        finally:
+            _restore_state(Qwen3_5Attention, state)
+
+
+class TestSignatureContract:
+    """Pin the stock attention forward signature; transformers drift would otherwise surface as a confusing TypeError mid-training."""
+
+    def test_qwen3_forward_required_params(self):
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
+
+        sig = inspect.signature(Qwen3Attention.forward)
+        params = list(sig.parameters)
+        assert params[:4] == [
+            "self",
+            "hidden_states",
+            "position_embeddings",
+            "attention_mask",
+        ], (
+            "Qwen3Attention.forward signature drifted away from the contract "
+            "our fused_forward replacement assumes — update the patch."
+        )
+
+    def test_qwen3_5_forward_required_params(self):
+        pytest.importorskip("transformers.models.qwen3_5")
+        from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Attention
+
+        sig = inspect.signature(Qwen3_5Attention.forward)
+        params = list(sig.parameters)
+        assert params[:4] == [
+            "self",
+            "hidden_states",
+            "position_embeddings",
+            "attention_mask",
+        ]
+
+
+class TestGradientCheckpointingCompose:
+    """Pin that the fused forward survives being re-run inside a checkpoint partial during backward."""
+
+    def test_qwen3_fused_under_gradient_checkpointing(self, restore_qwen3_attention):
+        from axolotl.monkeypatch.models.qwen3.fused_attn import (
+            patch_qwen3_fused_attn,
+        )
+
+        m = _build_tiny_qwen3().train()
+        m.gradient_checkpointing_enable()
+        patch_qwen3_fused_attn()
+
+        ids = torch.randint(0, 128, (2, 32), device="cuda")
+        mask = torch.ones(2, 32, dtype=torch.long, device="cuda")
+        out = m(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state
+        loss = out.sum()
+        torch.cuda.reset_peak_memory_stats()
+        loss.backward()
+        peak_mb = torch.cuda.max_memory_allocated() / 1024**2
+
+        for layer in m.layers:
+            attn = layer.self_attn
+            assert attn.q_norm.weight.grad is not None
+            assert attn.k_norm.weight.grad is not None
+            assert attn.q_norm.weight.grad.isfinite().all()
+            assert attn.k_norm.weight.grad.isfinite().all()
+            assert attn.q_norm.weight.grad.abs().sum() > 0
+        assert peak_mb < 1024, f"backward peak {peak_mb:.0f} MB looks like a leak"
+
+
+class TestCrossDeviceNormWeight:
+    """Sharded ``device_map='auto'`` can leave norm weights on CPU; the patch must coerce them or Triton raises on the CPU pointer."""
+
+    def test_qwen3_norm_weight_on_cpu_does_not_crash(self, restore_qwen3_attention):
+        from axolotl.monkeypatch.models.qwen3.fused_attn import (
+            patch_qwen3_fused_attn,
+        )
+
+        m = _build_tiny_qwen3()
+        patch_qwen3_fused_attn()
+
+        for layer in m.layers:
+            attn = layer.self_attn
+            attn.q_norm.weight.data = attn.q_norm.weight.data.cpu()
+            attn.k_norm.weight.data = attn.k_norm.weight.data.cpu()
+
+        ids = torch.randint(0, 128, (1, 16), device="cuda")
+        mask = torch.ones(1, 16, dtype=torch.long, device="cuda")
+        with torch.no_grad():
+            out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+        assert torch.isfinite(out).all()
+
+
+class TestFlashAttention2Compose:
+    """The fused region is upstream of ``attention_interface``; pin clean composition with FA2 if it's installed."""
+
+    def test_qwen3_fused_under_flash_attention_2(self, restore_qwen3_attention):
+        pytest.importorskip("flash_attn")
+        from axolotl.monkeypatch.models.qwen3.fused_attn import (
+            patch_qwen3_fused_attn,
+        )
+
+        m = _build_tiny_qwen3()
+        m.config._attn_implementation = "flash_attention_2"
+        for layer in m.layers:
+            layer.self_attn.config._attn_implementation = "flash_attention_2"
+
+        patch_qwen3_fused_attn()
+        ids = torch.randint(0, 128, (2, 32), device="cuda")
+        mask = torch.ones(2, 32, dtype=torch.long, device="cuda")
+        with torch.no_grad():
+            out = m(input_ids=ids, attention_mask=mask).last_hidden_state
+        assert torch.isfinite(out).all()
diff --git a/tests/monkeypatch/test_qwen3_next_modeling_patch.py b/tests/monkeypatch/test_qwen3_next_modeling_patch.py
new file mode 100644
index 0000000000..91d9fc1cf1
--- /dev/null
+++ b/tests/monkeypatch/test_qwen3_next_modeling_patch.py
@@ -0,0 +1,111 @@
+"""Integration tests for Qwen3 Next modeling patches."""
+
+import pytest
+import torch
+
+# Skip entire module if qwen3_next not available
+qwen3_next = pytest.importorskip("transformers.models.qwen3_next.modeling_qwen3_next")
+
+
+class TestQwen3NextModelingPatchIntegration:
+    """Test Qwen3 Next modeling patch integration."""
+
+    @pytest.mark.integration
+    def test_qwen3_next_decoder_layer_patch(self):
+        """Test that Qwen3Next decoder layer patch can be applied."""
+        from axolotl.monkeypatch.models.qwen3_next.modeling import (
+            patch_qwen3_next_decoder_layer,
+        )
+
+        # Store original method
+        original_forward = qwen3_next.Qwen3NextDecoderLayer.forward
+
+        # Apply patch and get unpatch function
+        unpatch_fn = patch_qwen3_next_decoder_layer()
+
+        # Verify patch was applied
+        assert qwen3_next.Qwen3NextDecoderLayer.forward != original_forward, (
+            "decoder layer forward method was not patched"
+        )
+
+        # Verify the method is still callable
+        assert callable(qwen3_next.Qwen3NextDecoderLayer.forward), (
+            "Patched method is not callable"
+        )
+
+        # Test unpatch function
+        if unpatch_fn:
+            unpatch_fn()
+            assert qwen3_next.Qwen3NextDecoderLayer.forward == original_forward, (
+                "unpatch function did not restore original method"
+            )
+
+    @pytest.mark.integration
+    def test_qwen3_next_gateddelta_layer_patch(self):
+        """Test that Qwen3Next GatedDeltaNet patch can be applied."""
+        from axolotl.monkeypatch.models.qwen3_next.modeling import (
+            patch_qwen3_next_gateddelta_layer,
+        )
+
+        # Store original method
+        original_forward = qwen3_next.Qwen3NextGatedDeltaNet.forward
+
+        # Apply patch and get unpatch function
+        unpatch_fn = patch_qwen3_next_gateddelta_layer()
+
+        # Verify patch was applied
+        assert qwen3_next.Qwen3NextGatedDeltaNet.forward != original_forward, (
+            "GatedDeltaNet forward method was not patched"
+        )
+
+        # Verify the method is still callable
+        assert callable(qwen3_next.Qwen3NextGatedDeltaNet.forward), (
+            "Patched method is not callable"
+        )
+
+        # Test unpatch function
+        if unpatch_fn:
+            unpatch_fn()
+            assert qwen3_next.Qwen3NextGatedDeltaNet.forward == original_forward, (
+                "unpatch function did not restore original method"
+            )
+
+    @pytest.mark.integration
+    def test_qwen3_next_imports_patch(self):
+        """Test that Qwen3Next imports patch can be applied without errors."""
+        from axolotl.monkeypatch.models.qwen3_next.modeling import (
+            patch_qwen3_next_imports,
+        )
+
+        # Apply patch - should not raise any exceptions even if modules unavailable
+        unpatch_fn = patch_qwen3_next_imports()
+
+        # Test that unpatch function is returned (or None if skipped)
+        assert unpatch_fn is None or callable(unpatch_fn), (
+            "patch_qwen3_next_imports should return None or callable unpatch function"
+        )
+
+    @pytest.mark.integration
+    def test_qwen3_next_modeling_packing_patch(self):
+        """Test that all Qwen3Next modeling patches can be applied together."""
+        from axolotl.monkeypatch.models.qwen3_next.modeling import (
+            patch_qwen3_next_modeling_packing,
+        )
+
+        # This should not raise any exceptions
+        patch_qwen3_next_modeling_packing()
+
+
+@pytest.mark.integration
+def test_get_cu_seqlens_utility():
+    """Test the get_cu_seqlens utility function."""
+    from axolotl.monkeypatch.models.qwen3_next.modeling import get_cu_seqlens
+
+    # Test with simple position_ids
+    position_ids = torch.tensor([[0, 1, 2, 0, 1]])
+    cu_seqlens = get_cu_seqlens(position_ids)
+    assert cu_seqlens.dtype == torch.int32, "Should be int32 dtype"
+
+    # Should return tensor with start positions and total length
+    expected = torch.tensor([0, 3, 5], dtype=torch.int32)
+    assert torch.equal(cu_seqlens, expected), f"Expected {expected}, got {cu_seqlens}"
diff --git a/tests/monkeypatch/test_qwen3_vl_fused_attn.py b/tests/monkeypatch/test_qwen3_vl_fused_attn.py
new file mode 100644
index 0000000000..1cc422e8b5
--- /dev/null
+++ b/tests/monkeypatch/test_qwen3_vl_fused_attn.py
@@ -0,0 +1,168 @@
+"""Tests for the Qwen3-VL text fused-attention monkeypatch."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+pytest.importorskip("transformers.models.qwen3_vl")
+
+
+def _clear_patched_flag(cls):
+    try:
+        delattr(cls, "_axolotl_fused_attn_patched")
+    except AttributeError:
+        pass
+
+
+@pytest.fixture
+def restore_qwen3_vl_attention():
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextAttention
+
+    saved = Qwen3VLTextAttention.forward
+    saved_flag = getattr(Qwen3VLTextAttention, "_axolotl_fused_attn_patched", False)
+    yield Qwen3VLTextAttention
+    Qwen3VLTextAttention.forward = saved
+    if saved_flag:
+        Qwen3VLTextAttention._axolotl_fused_attn_patched = saved_flag
+    else:
+        _clear_patched_flag(Qwen3VLTextAttention)
+
+
+def _build_qwen3_vl_text_model(seed: int = 0):
+    from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextModel
+
+    torch.manual_seed(seed)
+    cfg = Qwen3VLTextConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        max_position_embeddings=512,
+        rms_norm_eps=1e-6,
+        attention_dropout=0.0,
+        pad_token_id=0,
+    )
+    cfg._attn_implementation = "sdpa"
+    return Qwen3VLTextModel(cfg).cuda().to(torch.bfloat16).eval()
+
+
+def _run_attention(model, layer_idx, hidden_states, position_ids):
+    attn = model.layers[layer_idx].self_attn
+    cos, sin = model.rotary_emb(hidden_states, position_ids)
+    out, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+    )
+    return out
+
+
+class TestQwen3VLFusedAttnParity:
+    """Single-layer parity vs stock Qwen3VLTextAttention."""
+
+    @pytest.mark.parametrize("layer_idx", [0, 1])
+    def test_forward_matches_stock(self, restore_qwen3_vl_attention, layer_idx):
+        from axolotl.monkeypatch.models.qwen3_vl.fused_attn import (
+            patch_qwen3_vl_fused_attn,
+        )
+
+        model = _build_qwen3_vl_text_model(seed=1)
+        hidden_states = torch.randn(2, 16, 64, device="cuda", dtype=torch.bfloat16)
+        position_ids = torch.arange(16, device="cuda").unsqueeze(0).expand(2, -1)
+
+        with torch.no_grad():
+            ref = _run_attention(model, layer_idx, hidden_states, position_ids)
+
+        patch_qwen3_vl_fused_attn()
+        with torch.no_grad():
+            got = _run_attention(model, layer_idx, hidden_states, position_ids)
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, (
+            f"layer {layer_idx} fused vs stock cosine_sim={cos_sim:.6f}"
+        )
+        torch.testing.assert_close(got, ref, rtol=5e-2, atol=5e-2)
+
+
+class TestQwen3VLFusedAttnEndToEnd:
+    def test_full_forward_matches_stock(self, restore_qwen3_vl_attention):
+        from axolotl.monkeypatch.models.qwen3_vl.fused_attn import (
+            patch_qwen3_vl_fused_attn,
+        )
+
+        model = _build_qwen3_vl_text_model(seed=2)
+        input_ids = torch.randint(0, 128, (2, 32), device="cuda")
+        attention_mask = torch.ones(2, 32, dtype=torch.long, device="cuda")
+
+        with torch.no_grad():
+            ref = model(
+                input_ids=input_ids, attention_mask=attention_mask
+            ).last_hidden_state.clone()
+
+        patch_qwen3_vl_fused_attn()
+        with torch.no_grad():
+            got = model(
+                input_ids=input_ids, attention_mask=attention_mask
+            ).last_hidden_state.clone()
+
+        assert got.shape == ref.shape
+        assert torch.isfinite(got).all()
+        cos_sim = torch.nn.functional.cosine_similarity(
+            ref.flatten().float(), got.flatten().float(), dim=0
+        )
+        assert cos_sim > 0.999, f"end-to-end cosine_sim={cos_sim:.6f}"
+
+    def test_backward_grad_flows_through_fused_path(self, restore_qwen3_vl_attention):
+        from axolotl.monkeypatch.models.qwen3_vl.fused_attn import (
+            patch_qwen3_vl_fused_attn,
+        )
+
+        model = _build_qwen3_vl_text_model(seed=3).train()
+        patch_qwen3_vl_fused_attn()
+
+        input_ids = torch.randint(0, 128, (2, 16), device="cuda")
+        attention_mask = torch.ones(2, 16, dtype=torch.long, device="cuda")
+        out = model(
+            input_ids=input_ids, attention_mask=attention_mask
+        ).last_hidden_state
+        out.sum().backward()
+
+        for idx, layer in enumerate(model.layers[:2]):
+            attn = layer.self_attn
+            assert attn.q_norm.weight.grad is not None, f"layer {idx} q_norm no grad"
+            assert attn.k_norm.weight.grad is not None, f"layer {idx} k_norm no grad"
+            assert attn.q_norm.weight.grad.isfinite().all()
+            assert attn.k_norm.weight.grad.isfinite().all()
+            assert attn.q_norm.weight.grad.abs().sum() > 0
+            assert attn.k_norm.weight.grad.abs().sum() > 0
+
+
+class TestQwen3VLPatchManagerDispatch:
+    def test_patch_manager_dispatches_qwen3_vl(self, restore_qwen3_vl_attention):
+        from types import SimpleNamespace
+
+        from axolotl.loaders.patch_manager import PatchManager
+
+        cfg = SimpleNamespace(
+            fused_attn_kernel=True,
+            model_config_type="qwen3_vl",
+            llama4_linearized_experts=False,
+            sample_packing=False,
+            context_parallel_size=1,
+            attn_uses_flash_lib=False,
+        )
+
+        PatchManager(cfg=cfg, model_config=object())._apply_model_specific_patches()
+
+        assert getattr(restore_qwen3_vl_attention, "_axolotl_fused_attn_patched", False)
diff --git a/tests/monkeypatch/test_relora.py b/tests/monkeypatch/test_relora.py
new file mode 100644
index 0000000000..e722adf0cd
--- /dev/null
+++ b/tests/monkeypatch/test_relora.py
@@ -0,0 +1,186 @@
+"""Unit tests for axolotl.monkeypatch.relora.reset_optimizer."""
+
+import math
+
+import pytest
+import torch
+import torch.nn as nn
+
+from axolotl.monkeypatch.relora import (
+    magnitude_pruning_,
+    random_pruning_,
+    reset_optimizer,
+)
+
+ADAM_KEYS = ["exp_avg", "exp_avg_sq"]
+
+
+def _build_optimizer_with_state(seed: int = 0):
+    """Build a tiny optimizer over LoRA-shaped + non-LoRA params with populated state."""
+    torch.manual_seed(seed)
+    lora_a = nn.Parameter(torch.randn(8, 32))
+    lora_b = nn.Parameter(torch.randn(32, 8))
+    extra = nn.Parameter(torch.randn(64, 32))
+
+    optimizer = torch.optim.AdamW([lora_a, lora_b, extra], lr=1e-3)
+    for _ in range(2):
+        loss = (
+            (lora_a * torch.randn_like(lora_a)).sum()
+            + (lora_b * torch.randn_like(lora_b)).sum()
+            + (extra * torch.randn_like(extra)).sum()
+        )
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+    return optimizer, lora_a, lora_b, extra
+
+
+def test_reset_optimizer_only_touches_reset_params():
+    """State for params NOT in reset_params must be byte-identical after reset."""
+    optimizer, lora_a, lora_b, extra = _build_optimizer_with_state()
+
+    extra_avg_before = optimizer.state[extra]["exp_avg"].clone()
+    extra_avg_sq_before = optimizer.state[extra]["exp_avg_sq"].clone()
+
+    reset_optimizer(
+        optimizer,
+        reset_params=[lora_a, lora_b],
+        optimizer_state_keys=ADAM_KEYS,
+        prune_method="magnitude",
+        prune_ratio=0.9,
+    )
+
+    assert torch.equal(optimizer.state[extra]["exp_avg"], extra_avg_before)
+    assert torch.equal(optimizer.state[extra]["exp_avg_sq"], extra_avg_sq_before)
+
+
+def test_reset_optimizer_actually_prunes_lora_state():
+    optimizer, lora_a, lora_b, _extra = _build_optimizer_with_state()
+
+    reset_optimizer(
+        optimizer,
+        reset_params=[lora_a, lora_b],
+        optimizer_state_keys=ADAM_KEYS,
+        prune_method="magnitude",
+        prune_ratio=0.9,
+    )
+
+    for param in (lora_a, lora_b):
+        for key in ADAM_KEYS:
+            zero_frac = (optimizer.state[param][key] == 0).float().mean().item()
+            assert zero_frac >= 0.85
+
+
+@pytest.mark.parametrize(
+    "method,ratio,expected_zero_frac",
+    [
+        ("magnitude", 0.9, 0.9),
+        ("magnitude", 0.99, 0.99),
+        ("random", 0.9, 0.9),
+        ("random", 0.5, 0.5),
+        # reset uses random pruning; relora_prune_ratio must be honored, not ignored.
+        ("reset", 0.9, 0.9),
+        ("reset", 0.5, 0.5),
+    ],
+)
+def test_prune_methods(method, ratio, expected_zero_frac):
+    """Each method zeros approximately the expected fraction."""
+    optimizer, lora_a, lora_b, _extra = _build_optimizer_with_state(seed=42)
+
+    reset_optimizer(
+        optimizer,
+        reset_params=[lora_a, lora_b],
+        optimizer_state_keys=ADAM_KEYS,
+        prune_method=method,
+        prune_ratio=ratio,
+    )
+
+    total = 0
+    zeros = 0
+    for param in (lora_a, lora_b):
+        for key in ADAM_KEYS:
+            tensor = optimizer.state[param][key]
+            total += tensor.numel()
+            zeros += (tensor == 0).sum().item()
+
+    actual = zeros / total
+    tolerance = 0.02 if method == "magnitude" else 0.05
+    assert math.isclose(actual, expected_zero_frac, abs_tol=tolerance)
+
+
+def test_reset_optimizer_skips_keys_not_in_state_keys():
+    """Keys present in optimizer state but not in optimizer_state_keys are untouched."""
+    optimizer, lora_a, lora_b, _extra = _build_optimizer_with_state()
+
+    exp_avg_sq_before = optimizer.state[lora_a]["exp_avg_sq"].clone()
+
+    reset_optimizer(
+        optimizer,
+        reset_params=[lora_a, lora_b],
+        optimizer_state_keys=["exp_avg"],
+        prune_method="magnitude",
+        prune_ratio=0.9,
+    )
+
+    assert torch.equal(optimizer.state[lora_a]["exp_avg_sq"], exp_avg_sq_before)
+
+
+def test_reset_optimizer_handles_param_with_empty_state():
+    """Params with no optimizer state are skipped silently."""
+    optimizer, lora_a, lora_b, _extra = _build_optimizer_with_state()
+    orphan = nn.Parameter(torch.randn(4, 4))
+
+    reset_optimizer(
+        optimizer,
+        reset_params=[lora_a, lora_b, orphan],
+        optimizer_state_keys=ADAM_KEYS,
+        prune_method="magnitude",
+        prune_ratio=0.9,
+    )
+
+    assert orphan not in optimizer.state or not optimizer.state[orphan]
+
+
+def test_unknown_prune_method_raises():
+    optimizer, lora_a, lora_b, _extra = _build_optimizer_with_state()
+
+    with pytest.raises(ValueError, match="Unknown prune_method"):
+        reset_optimizer(
+            optimizer,
+            reset_params=[lora_a, lora_b],
+            optimizer_state_keys=ADAM_KEYS,
+            prune_method="bogus",  # type: ignore[arg-type]
+            prune_ratio=0.9,
+        )
+
+
+def test_pruning_helpers_are_inplace():
+    """magnitude_pruning_ and random_pruning_ must mutate via tensor.mul_."""
+    tensor = torch.randn(64)
+    ptr_before = tensor.data_ptr()
+    magnitude_pruning_(tensor, 0.5)
+    assert tensor.data_ptr() == ptr_before
+
+    tensor = torch.randn(64)
+    ptr_before = tensor.data_ptr()
+    random_pruning_(tensor, 0.5)
+    assert tensor.data_ptr() == ptr_before
+
+
+def test_pruning_helpers_support_uint8_tensors():
+    """Both pruning helpers must work on uint8 optimizer state tensors."""
+    tensor = torch.arange(1, 129, dtype=torch.uint8)
+    magnitude_pruning_(tensor, 0.9)
+
+    assert tensor.dtype == torch.uint8
+    magnitude_zero_frac = (tensor == 0).float().mean().item()
+    assert 0.85 <= magnitude_zero_frac <= 0.95
+
+    tensor = torch.arange(1, 129, dtype=torch.uint8)
+    with torch.random.fork_rng(devices=[]):
+        torch.manual_seed(1234)
+        random_pruning_(tensor, 0.9)
+
+    assert tensor.dtype == torch.uint8
+    random_zero_frac = (tensor == 0).float().mean().item()
+    assert 0.85 <= random_zero_frac <= 0.95
diff --git a/tests/monkeypatch/test_sdpa_varlen.py b/tests/monkeypatch/test_sdpa_varlen.py
new file mode 100644
index 0000000000..c0254c55c2
--- /dev/null
+++ b/tests/monkeypatch/test_sdpa_varlen.py
@@ -0,0 +1,278 @@
+"""Correctness tests for the cu_seqlens varlen SDPA packing path (``sdpa_varlen``)."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from axolotl.monkeypatch.attention.sdpa_varlen import (
+    patch_sdpa_varlen,
+    unpatch_sdpa_varlen,
+    varlen_available,
+)
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="varlen_attn needs CUDA"),
+    pytest.mark.skipif(
+        not varlen_available(), reason="torch.nn.attention.varlen needs torch >= 2.10"
+    ),
+]
+
+DEV = "cuda"
+
+
+def _varlen_supports_window() -> bool:
+    """varlen_attn gained the ``window_size`` arg (sliding window) in torch 2.11; on causal-only
+    builds the sdpa_varlen patch refuses sliding, so those params can't be exercised here."""
+    if not varlen_available():
+        return False
+    import inspect
+
+    from torch.nn.attention.varlen import varlen_attn
+
+    try:
+        params = set(inspect.signature(varlen_attn).parameters)
+    except (TypeError, ValueError):
+        params = set()
+    return "window_size" in params or "is_causal" not in params
+
+
+class _Mod:
+    def __init__(self, sliding_window=None, num_key_value_groups=4):
+        self.sliding_window = sliding_window
+        self.num_key_value_groups = num_key_value_groups
+
+
+def _block_mask(position_ids, dtype, sliding=None):
+    B, S = position_ids.shape
+    doc = (position_ids == 0).cumsum(-1)
+    same = doc[:, :, None] == doc[:, None, :]
+    idx = torch.arange(S, device=DEV)
+    allow = same & (idx[:, None] >= idx[None, :])[None]
+    if sliding:
+        allow = allow & ((idx[:, None] - idx[None, :]) < sliding)[None]
+    m = torch.zeros(B, 1, S, S, device=DEV, dtype=dtype)
+    m.masked_fill_(~allow[:, None], torch.finfo(dtype).min)
+    return m
+
+
+def _ref(q, k, v, scaling, position_ids, sliding=None):
+    Hq, Hkv = q.shape[1], k.shape[1]
+    kr = k.repeat_interleave(Hq // Hkv, 1)
+    vr = v.repeat_interleave(Hq // Hkv, 1)
+    o = F.scaled_dot_product_attention(
+        q, kr, vr, attn_mask=_block_mask(position_ids, q.dtype, sliding), scale=scaling
+    )
+    return o.transpose(1, 2)  # [B,S,H,D]
+
+
+def _pos(doclens_per_row):
+    rows = [torch.cat([torch.arange(ln) for ln in dl]) for dl in doclens_per_row]
+    S = max(len(r) for r in rows)
+    return torch.stack([F.pad(r, (0, S - len(r))) for r in rows]).to(DEV)
+
+
+@pytest.fixture
+def patched():
+    assert patch_sdpa_varlen()
+    yield
+    unpatch_sdpa_varlen()
+
+
+@pytest.mark.parametrize(
+    "doclens,sliding,label",
+    [
+        ([[2048, 2048]], None, "B1-2docs"),
+        ([[1000, 1500, 1596]], None, "B1-3docs"),
+        pytest.param(
+            [[2048, 2048]],
+            1024,
+            "B1-sliding",
+            marks=pytest.mark.skipif(
+                not _varlen_supports_window(),
+                reason="sliding-window varlen needs varlen_attn(window_size=...) (torch >= 2.11)",
+            ),
+        ),
+        ([[2000, 2096], [1500, 1000, 1596]], None, "B2-padded"),
+    ],
+)
+def test_varlen_matches_block_mask_sdpa(patched, doclens, sliding, label):
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    sdpa = ALL_ATTENTION_FUNCTIONS["sdpa"]
+    pos = _pos(doclens)
+    B, S = pos.shape
+    Hq, Hkv, D = 16, 4, 256
+    torch.manual_seed(0)
+    q = torch.randn(B, Hq, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    k = torch.randn(B, Hkv, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    v = torch.randn(B, Hkv, S, D, device=DEV, dtype=torch.bfloat16, requires_grad=True)
+    qr, kr, vr = (t.detach().clone().requires_grad_() for t in (q, k, v))
+    mod = _Mod(sliding_window=sliding, num_key_value_groups=Hq // Hkv)
+    scaling = D**-0.5
+
+    out, _ = sdpa(
+        mod, q, k, v, None, scaling=scaling, position_ids=pos, sliding_window=sliding
+    )
+    out.float().pow(2).mean().backward()
+    ref = _ref(qr, kr, vr, scaling, pos, sliding)
+    ref.float().pow(2).mean().backward()
+
+    assert F.cosine_similarity(out.float().flatten(), ref.float().flatten(), 0) > 0.999
+    assert (
+        F.cosine_similarity(q.grad.float().flatten(), qr.grad.float().flatten(), 0)
+        > 0.999
+    )
+
+
+def _assert_defers_to_stock(mod, q, k, v, pos):
+    """The wrapper must produce exactly the stock-SDPA result (it just calls through)."""
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    wrapper = ALL_ATTENTION_FUNCTIONS["sdpa"]
+    original = wrapper._axolotl_sdpa_original
+    scaling = q.shape[-1] ** -0.5
+    out_w, _ = wrapper(mod, q, k, v, None, scaling=scaling, position_ids=pos)
+    out_o, _ = original(mod, q, k, v, None, scaling=scaling, position_ids=pos)
+    assert torch.equal(out_w, out_o)
+
+
+def test_varlen_matches_flash_attention_2_e2e():
+    """The varlen path matches FA2 (the canonical varlen-packing impl) on a real packed forward."""
+    from transformers import LlamaConfig, LlamaModel
+
+    def build(impl):
+        cfg = LlamaConfig(
+            hidden_size=512,
+            num_hidden_layers=2,
+            num_attention_heads=8,
+            num_key_value_heads=2,
+            intermediate_size=1024,
+            vocab_size=256,
+            head_dim=64,
+            _attn_implementation=impl,
+        )
+        torch.manual_seed(1)
+        return LlamaModel(cfg).to(DEV).to(torch.bfloat16).eval()
+
+    docs = [20, 30, 14]
+    ids = torch.cat([torch.randint(0, 256, (1, d), device=DEV) for d in docs], 1)
+    pos = torch.cat([torch.arange(d) for d in docs]).to(DEV)[None]
+    try:
+        with torch.no_grad():
+            fa2 = build("flash_attention_2")(
+                input_ids=ids, position_ids=pos
+            ).last_hidden_state
+    except Exception:  # pylint: disable=broad-except
+        pytest.skip("flash_attention_2 unavailable")
+
+    assert patch_sdpa_varlen()
+    try:
+        with torch.no_grad():
+            varlen = build("sdpa")(input_ids=ids, position_ids=pos).last_hidden_state
+    finally:
+        unpatch_sdpa_varlen()
+    assert (
+        F.cosine_similarity(varlen.float().flatten(), fa2.float().flatten(), 0) > 0.999
+    )
+
+
+def test_falls_back_when_not_packed(patched):
+    """Single-document rows must defer to stock SDPA (no varlen path)."""
+    S, Hq, Hkv, D = 512, 16, 4, 256
+    pos = torch.arange(S, device=DEV)[None]  # one document -> no packing
+    torch.manual_seed(0)
+    q = torch.randn(1, Hq, S, D, device=DEV, dtype=torch.bfloat16)
+    k = torch.randn(1, Hkv, S, D, device=DEV, dtype=torch.bfloat16)
+    v = torch.randn(1, Hkv, S, D, device=DEV, dtype=torch.bfloat16)
+    _assert_defers_to_stock(_Mod(num_key_value_groups=Hq // Hkv), q, k, v, pos)
+
+
+def test_falls_back_on_large_head_dim(patched):
+    """head_dim > 256 can't use Flash varlen, but a packed row whose padding mask was
+    dropped (attention_mask=None) must still be isolated: the wrapper rebuilds the
+    block-diagonal mask rather than running pure-causal (which would leak)."""
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    pos = _pos([[256, 256]])
+    B, S = pos.shape
+    Hq, Hkv, D = 8, 2, 512  # head_dim 512
+    torch.manual_seed(0)
+    q = torch.randn(B, Hq, S, D, device=DEV, dtype=torch.bfloat16)
+    k = torch.randn(B, Hkv, S, D, device=DEV, dtype=torch.bfloat16)
+    v = torch.randn(B, Hkv, S, D, device=DEV, dtype=torch.bfloat16)
+    mod = _Mod(num_key_value_groups=Hq // Hkv)
+    scaling = D**-0.5
+
+    wrapper = ALL_ATTENTION_FUNCTIONS["sdpa"]
+    out, _ = wrapper(mod, q, k, v, None, scaling=scaling, position_ids=pos)
+    # matches the block-diagonal reference (documents isolated)...
+    ref = _ref(q, k, v, scaling, pos)
+    assert torch.allclose(out, ref, atol=1e-2)
+    # ...and is NOT the pure-causal result stock SDPA gives with a None mask (a leak).
+    leak, _ = wrapper._axolotl_sdpa_original(
+        mod, q, k, v, None, scaling=scaling, position_ids=pos
+    )
+    assert not torch.allclose(out, leak, atol=1e-2)
+
+
+def test_varlen_engages_in_training_forward():
+    """Regression: in a real training forward (use_cache=False) transformers builds a 4D
+    packed mask from position_ids, which used to bypass the varlen path entirely. The mask
+    override must make ``varlen_attn`` actually fire for a padded packed row, and the
+    result must match per-document forwards (no cross-document leakage)."""
+    import torch.nn.attention.varlen as va
+    from transformers import LlamaConfig, LlamaModel
+
+    calls = {"n": 0}
+    original_varlen = va.varlen_attn
+    va.varlen_attn = lambda *a, **k: (  # spy captured by the wrapper at patch time
+        calls.__setitem__("n", calls["n"] + 1),
+        original_varlen(*a, **k),
+    )[1]
+    try:
+        assert patch_sdpa_varlen()
+        try:
+            cfg = LlamaConfig(
+                hidden_size=256,
+                num_hidden_layers=2,
+                num_attention_heads=8,
+                num_key_value_heads=2,
+                intermediate_size=512,
+                vocab_size=256,
+                head_dim=32,
+                _attn_implementation="sdpa",
+            )
+            torch.manual_seed(0)
+            model = LlamaModel(cfg).to(DEV).to(torch.bfloat16).eval()
+
+            doc1, doc2 = list(range(1, 21)), list(range(21, 37))  # len 20 + 16
+            ids = torch.tensor([doc1 + doc2 + [0, 0, 0, 0]], device=DEV)  # + padding
+            pos = torch.tensor(
+                [list(range(20)) + list(range(16)) + list(range(4))], device=DEV
+            )
+            with torch.no_grad():
+                # attention_mask=None mimics axolotl dropping the mask for packing.
+                packed = model(
+                    input_ids=ids,
+                    attention_mask=None,
+                    position_ids=pos,
+                    use_cache=False,
+                ).last_hidden_state[:, :36]
+                ref1 = model(
+                    input_ids=torch.tensor([doc1], device=DEV),
+                    position_ids=torch.tensor([list(range(20))], device=DEV),
+                    use_cache=False,
+                ).last_hidden_state
+                ref2 = model(
+                    input_ids=torch.tensor([doc2], device=DEV),
+                    position_ids=torch.tensor([list(range(16))], device=DEV),
+                    use_cache=False,
+                ).last_hidden_state
+
+            assert calls["n"] > 0, "varlen_attn did not fire in the training forward"
+            reference = torch.cat([ref1, ref2], dim=1)
+            assert torch.allclose(packed, reference, atol=1e-2)
+        finally:
+            unpatch_sdpa_varlen()
+    finally:
+        va.varlen_attn = original_varlen
diff --git a/tests/monkeypatch/test_selective_checkpointing.py b/tests/monkeypatch/test_selective_checkpointing.py
new file mode 100644
index 0000000000..0463231a38
--- /dev/null
+++ b/tests/monkeypatch/test_selective_checkpointing.py
@@ -0,0 +1,287 @@
+"""Tests for eager selective activation checkpointing (SAC)."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch.utils.checkpoint import CheckpointPolicy, checkpoint
+
+from axolotl.monkeypatch.selective_checkpointing import (
+    SacPolicyState,
+    apply_selective_checkpointing,
+    build_sac_context_fn,
+    build_sac_policy,
+)
+
+
+class _FakeOp:
+    def __init__(self, name: str):
+        self._name = name
+
+    def name(self) -> str:
+        return self._name
+
+
+class _FakeSchemaArg:
+    def __init__(self, name: str):
+        self.name = name
+
+
+class _FakeSchema:
+    def __init__(self, arg_names: list[str]):
+        self.arguments = [_FakeSchemaArg(n) for n in arg_names]
+
+
+class _FakeFlashOp(_FakeOp):
+    """Mimics flash-attn's registered custom op with flattened window args."""
+
+    ARG_NAMES = [
+        "q",
+        "k",
+        "v",
+        "dropout_p",
+        "softmax_scale",
+        "is_causal",
+        "window_size_left",
+        "window_size_right",
+    ]
+
+    def __init__(self, name: str = "flash_attn::_flash_attn_forward"):
+        super().__init__(name)
+        self._schema = _FakeSchema(self.ARG_NAMES)
+
+    @classmethod
+    def args_with_window(cls, left: int, right: int) -> tuple:
+        return (None, None, None, 0.0, 1.0, True, left, right)
+
+
+class TestSacPolicy:
+    def test_sdpa_ops_saved(self):
+        policy = build_sac_policy(["attention"])
+        for packet in (
+            torch.ops.aten._scaled_dot_product_flash_attention,
+            torch.ops.aten._scaled_dot_product_efficient_attention,
+            torch.ops.aten._scaled_dot_product_cudnn_attention,
+        ):
+            assert policy(None, packet.default) == CheckpointPolicy.MUST_SAVE
+
+    def test_other_ops_recomputed(self):
+        policy = build_sac_policy(["attention"])
+        assert (
+            policy(None, torch.ops.aten.mm.default) == CheckpointPolicy.PREFER_RECOMPUTE
+        )
+        assert (
+            policy(None, torch.ops.aten._softmax.default)
+            == CheckpointPolicy.PREFER_RECOMPUTE
+        )
+
+    def test_flash_attn_custom_op_name_matched(self):
+        policy = build_sac_policy(["attention"])
+        assert (
+            policy(None, _FakeOp("flash_attn::_flash_attn_forward"))
+            == CheckpointPolicy.MUST_SAVE
+        )
+        assert (
+            policy(None, _FakeOp("flash_attn::_flash_attn_varlen_forward"))
+            == CheckpointPolicy.MUST_SAVE
+        )
+        assert (
+            policy(None, _FakeOp("flash_attn::_flash_attn_backward"))
+            == CheckpointPolicy.PREFER_RECOMPUTE
+        )
+
+    def test_substring_spec(self):
+        policy = build_sac_policy(["aten::mm"])
+        assert policy(None, torch.ops.aten.mm.default) == CheckpointPolicy.MUST_SAVE
+        assert (
+            policy(None, torch.ops.aten._scaled_dot_product_flash_attention.default)
+            == CheckpointPolicy.PREFER_RECOMPUTE
+        )
+
+    def test_state_records_saved_ops(self):
+        state = SacPolicyState()
+        policy = build_sac_policy(["attention"], state)
+        policy(None, torch.ops.aten._scaled_dot_product_flash_attention.default)
+        assert state.saved_op_names == {"aten::_scaled_dot_product_flash_attention"}
+
+
+class TestSlidingWindowDiscrimination:
+    def test_full_attention_saved(self):
+        policy = build_sac_policy(["attention"])
+        op = _FakeFlashOp()
+        args = _FakeFlashOp.args_with_window(-1, -1)
+        assert policy(None, op, *args) == CheckpointPolicy.MUST_SAVE
+
+    def test_sliding_window_recomputed(self):
+        state = SacPolicyState()
+        policy = build_sac_policy(["attention"], state)
+        op = _FakeFlashOp()
+        args = _FakeFlashOp.args_with_window(4095, 0)
+        assert policy(None, op, *args) == CheckpointPolicy.PREFER_RECOMPUTE
+        assert state.sliding_op_names == {"flash_attn::_flash_attn_forward"}
+
+    def test_causal_right_bound_is_not_sliding(self):
+        policy = build_sac_policy(["attention"])
+        op = _FakeFlashOp()
+        args = _FakeFlashOp.args_with_window(-1, 0)
+        assert policy(None, op, *args) == CheckpointPolicy.MUST_SAVE
+
+    def test_sliding_window_kwarg(self):
+        policy = build_sac_policy(["attention"])
+        op = _FakeFlashOp()
+        assert (
+            policy(None, op, window_size_left=1024) == CheckpointPolicy.PREFER_RECOMPUTE
+        )
+
+    def test_save_sliding_window_overrides(self):
+        policy = build_sac_policy(["attention"], save_sliding_window=True)
+        op = _FakeFlashOp()
+        args = _FakeFlashOp.args_with_window(4095, 0)
+        assert policy(None, op, *args) == CheckpointPolicy.MUST_SAVE
+
+    def test_sdpa_without_window_schema_saved(self):
+        policy = build_sac_policy(["attention"])
+        op = torch.ops.aten._scaled_dot_product_flash_attention.default
+        assert policy(None, op) == CheckpointPolicy.MUST_SAVE
+
+
+class TestEnableWrap:
+    class _FakeModel:
+        def __init__(self):
+            self.seen_kwargs = None
+
+        def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+            self.seen_kwargs = gradient_checkpointing_kwargs
+
+    def test_injects_context_fn_and_non_reentrant(self):
+        model = self._FakeModel()
+        apply_selective_checkpointing(model)
+        model.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs={"use_reentrant": True}
+        )
+        assert model.seen_kwargs["use_reentrant"] is False
+        assert callable(model.seen_kwargs["context_fn"])
+
+    def test_injects_with_none_kwargs(self):
+        model = self._FakeModel()
+        apply_selective_checkpointing(model)
+        model.gradient_checkpointing_enable()
+        assert model.seen_kwargs["use_reentrant"] is False
+        assert callable(model.seen_kwargs["context_fn"])
+
+    def test_idempotent(self):
+        model = self._FakeModel()
+        apply_selective_checkpointing(model)
+        wrapped = model.gradient_checkpointing_enable
+        apply_selective_checkpointing(model)
+        assert model.gradient_checkpointing_enable is wrapped
+
+
+class TestSacFunctional:
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA sdpa")
+    def test_checkpointed_attention_grads_match_baseline(self):
+        torch.manual_seed(0)
+        device = "cuda"
+        batch, heads, seq, dim = 2, 4, 128, 64
+
+        def make_inputs():
+            gen = torch.Generator(device="cpu").manual_seed(42)
+            qkv = torch.randn(
+                3, batch, heads, seq, dim, dtype=torch.float32, generator=gen
+            )
+            return [t.to(device).detach().clone().requires_grad_(True) for t in qkv]
+
+        def attn_block(q, k, v):
+            out = F.scaled_dot_product_attention(q, k, v)
+            return out.relu() @ v.transpose(-2, -1)
+
+        # baseline: no checkpointing
+        q0, k0, v0 = make_inputs()
+        attn_block(q0, k0, v0).sum().backward()
+
+        # SAC: checkpointed with save-attention policy
+        state = SacPolicyState()
+        policy = build_sac_policy(["attention"], state)
+
+        def context_fn():
+            from torch.utils.checkpoint import create_selective_checkpoint_contexts
+
+            return create_selective_checkpoint_contexts(policy)
+
+        q1, k1, v1 = make_inputs()
+        out = checkpoint(
+            attn_block, q1, k1, v1, use_reentrant=False, context_fn=context_fn
+        )
+        out.sum().backward()
+
+        assert state.saved_op_names, "no attention op was matched/saved"
+        torch.testing.assert_close(q0.grad, q1.grad)
+        torch.testing.assert_close(k0.grad, k1.grad)
+        torch.testing.assert_close(v0.grad, v1.grad)
+
+    def test_context_fn_returns_fresh_contexts(self):
+        context_fn = build_sac_context_fn(["attention"])
+        c1 = context_fn()
+        c2 = context_fn()
+        assert c1 is not c2
+
+
+class TestLayerTypeDiscrimination:
+    SDPA_OP = torch.ops.aten._scaled_dot_product_flash_attention.default
+
+    def test_sliding_layer_type_recomputed(self):
+        state = SacPolicyState()
+        policy = build_sac_policy(["attention"], state)
+        state.current_layer_type = "sliding_attention"
+        assert policy(None, self.SDPA_OP) == CheckpointPolicy.PREFER_RECOMPUTE
+        state.current_layer_type = "chunked_attention"
+        assert policy(None, self.SDPA_OP) == CheckpointPolicy.PREFER_RECOMPUTE
+
+    def test_full_or_unknown_layer_type_saved(self):
+        state = SacPolicyState()
+        policy = build_sac_policy(["attention"], state)
+        state.current_layer_type = "full_attention"
+        assert policy(None, self.SDPA_OP) == CheckpointPolicy.MUST_SAVE
+        state.current_layer_type = None
+        assert policy(None, self.SDPA_OP) == CheckpointPolicy.MUST_SAVE
+
+    def test_save_sliding_window_overrides_layer_type(self):
+        state = SacPolicyState()
+        policy = build_sac_policy(["attention"], state, save_sliding_window=True)
+        state.current_layer_type = "sliding_attention"
+        assert policy(None, self.SDPA_OP) == CheckpointPolicy.MUST_SAVE
+
+    def test_hooks_publish_layer_type(self):
+        from transformers import GradientCheckpointingLayer
+
+        from axolotl.monkeypatch.selective_checkpointing import (
+            install_layer_type_hooks,
+        )
+
+        state = SacPolicyState()
+        seen = []
+
+        class _Layer(GradientCheckpointingLayer):
+            def __init__(self, layer_type):
+                super().__init__()
+                self.layer_type = layer_type
+
+            def forward(self):
+                seen.append(state.current_layer_type)
+
+        class _Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList(
+                    [_Layer("full_attention"), _Layer("sliding_attention")]
+                )
+
+            def forward(self):
+                for layer in self.layers:
+                    layer()
+
+        model = _Model()
+        hooked = install_layer_type_hooks(model, state)
+        assert hooked == 2
+        model()
+        assert seen == ["full_attention", "sliding_attention"]
+        assert state.current_layer_type is None
diff --git a/tests/monkeypatch/test_selective_checkpointing_offload.py b/tests/monkeypatch/test_selective_checkpointing_offload.py
new file mode 100644
index 0000000000..68709e37b6
--- /dev/null
+++ b/tests/monkeypatch/test_selective_checkpointing_offload.py
@@ -0,0 +1,99 @@
+"""Tests for phase-2 SAC CPU offload."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+from axolotl.monkeypatch.selective_checkpointing_offload import (
+    SacOffloadEngine,
+    build_sac_offload_context_fn,
+)
+
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="requires CUDA"
+)
+
+
+class TestSacOffloadFunctional:
+    @requires_cuda
+    def test_grads_match_baseline_with_offload(self):
+        device = "cuda"
+        batch, heads, seq, dim = 2, 4, 512, 64
+
+        def make_inputs():
+            gen = torch.Generator(device="cpu").manual_seed(7)
+            qkv = torch.randn(
+                3, batch, heads, seq, dim, dtype=torch.float32, generator=gen
+            )
+            return [t.to(device).detach().clone().requires_grad_(True) for t in qkv]
+
+        def attn_block(q, k, v):
+            out = F.scaled_dot_product_attention(q, k, v)
+            return out.relu() @ v.transpose(-2, -1)
+
+        q0, k0, v0 = make_inputs()
+        attn_block(q0, k0, v0).sum().backward()
+
+        engine = SacOffloadEngine(min_offload_bytes=1024)
+        context_fn = build_sac_offload_context_fn(["attention"], engine=engine)
+
+        q1, k1, v1 = make_inputs()
+        out = checkpoint(
+            attn_block, q1, k1, v1, use_reentrant=False, context_fn=context_fn
+        )
+        out.sum().backward()
+        torch.cuda.synchronize()
+
+        assert engine.stats.offloaded_tensors > 0, "nothing was offloaded"
+        assert engine.stats.restored_tensors == engine.stats.offloaded_tensors
+        torch.testing.assert_close(q0.grad, q1.grad)
+        torch.testing.assert_close(k0.grad, k1.grad)
+        torch.testing.assert_close(v0.grad, v1.grad)
+
+    @requires_cuda
+    def test_multi_region_prefetch_and_reuse(self):
+        device = "cuda"
+        n_layers, seq, hidden = 4, 256, 128
+
+        gen = torch.Generator(device="cpu").manual_seed(11)
+        weights = [
+            torch.randn(hidden, hidden, generator=gen).to(device).requires_grad_(True)
+            for _ in range(n_layers)
+        ]
+
+        def layer(x, w):
+            q = (x @ w).view(1, 2, seq, hidden // 2)
+            out = F.scaled_dot_product_attention(q, q, q)
+            return out.reshape(1, seq, hidden).relu()
+
+        def run(context_fn=None):
+            gen2 = torch.Generator(device="cpu").manual_seed(12)
+            x = torch.randn(1, seq, hidden, generator=gen2).to(device)
+            for w in weights:
+                if w.grad is not None:
+                    w.grad = None
+            h = x
+            for w in weights:
+                if context_fn is not None:
+                    h = checkpoint(
+                        layer, h, w, use_reentrant=False, context_fn=context_fn
+                    )
+                else:
+                    h = layer(h, w)
+            h.sum().backward()
+            torch.cuda.synchronize()
+            return [w.grad.clone() for w in weights]
+
+        baseline = run()
+
+        engine = SacOffloadEngine(min_offload_bytes=1024)
+        context_fn = build_sac_offload_context_fn(["attention"], engine=engine)
+        # two steps: second step exercises the pinned-buffer pool reuse
+        run(context_fn)
+        grads = run(context_fn)
+
+        assert engine.stats.offloaded_tensors >= 2 * n_layers
+        assert engine.stats.restored_tensors == engine.stats.offloaded_tensors
+        for g0, g1 in zip(baseline, grads, strict=True):
+            torch.testing.assert_close(g0, g1)
diff --git a/tests/monkeypatch/test_trainer_accelerator_args.py b/tests/monkeypatch/test_trainer_accelerator_args.py
new file mode 100644
index 0000000000..fab2597f01
--- /dev/null
+++ b/tests/monkeypatch/test_trainer_accelerator_args.py
@@ -0,0 +1,26 @@
+"""
+Unit tests for trainer accelerator args monkeypatch
+"""
+
+import unittest
+
+from axolotl.monkeypatch.trainer_accelerator_args import (
+    check_create_accelerate_code_is_patchable,
+)
+
+
+class TestTrainerAcceleratorArgs(unittest.TestCase):
+    """
+    Unit test class for trainer accelerator args monkeypatch
+    """
+
+    def test_check_create_accelerate_code_is_patchable(self):
+        """
+        Test that the upstream transformers code is still patchable.
+        This will fail if the patched code changes upstream.
+        """
+        assert check_create_accelerate_code_is_patchable()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/monkeypatch/test_trl_vllm.py b/tests/monkeypatch/test_trl_vllm.py
new file mode 100644
index 0000000000..755ac2577e
--- /dev/null
+++ b/tests/monkeypatch/test_trl_vllm.py
@@ -0,0 +1,286 @@
+"""Unit tests for TRL vLLM monkeypatches.
+
+Tests:
+- split_tensor_dict: scalar type preservation (int/float/bool)
+- shuffle_sequence_dict: scalar type preservation
+- extract_logprobs: NaN → 0.0 replacement
+- VLLMClient.batch_update_named_params: method exists after patch
+- VLLMGeneration: weight_sync_chunk_size attribute after patch
+- Patch idempotency: applying patch twice doesn't break anything
+"""
+
+import unittest
+from dataclasses import dataclass
+from unittest.mock import MagicMock
+
+import torch
+
+
+class TestSplitTensorDict(unittest.TestCase):
+    """Tests for patched split_tensor_dict."""
+
+    def setUp(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import _patched_split_tensor_dict
+
+        self.split = _patched_split_tensor_dict
+
+    def test_scalar_int_preserved(self):
+        d = {"a": torch.randn(4, 3), "count": 42}
+        chunks = self.split(d, 2)
+        self.assertEqual(len(chunks), 2)
+        self.assertEqual(chunks[0]["count"], 42)
+        self.assertEqual(chunks[1]["count"], 42)
+
+    def test_scalar_float_preserved(self):
+        d = {"a": torch.randn(6, 2), "lr": 1e-5}
+        chunks = self.split(d, 3)
+        for c in chunks:
+            self.assertEqual(c["lr"], 1e-5)
+
+    def test_scalar_bool_preserved(self):
+        d = {"a": torch.randn(4, 2), "flag": True}
+        chunks = self.split(d, 2)
+        for c in chunks:
+            self.assertTrue(c["flag"])
+
+    def test_none_preserved(self):
+        d = {"a": torch.randn(4, 2), "b": None}
+        chunks = self.split(d, 2)
+        for c in chunks:
+            self.assertIsNone(c["b"])
+
+    def test_tensor_split(self):
+        t = torch.arange(8).reshape(4, 2)
+        d = {"a": t, "n": 10}
+        chunks = self.split(d, 2)
+        self.assertEqual(chunks[0]["a"].shape, (2, 2))
+        self.assertEqual(chunks[1]["a"].shape, (2, 2))
+        torch.testing.assert_close(chunks[0]["a"], t[:2])
+        torch.testing.assert_close(chunks[1]["a"], t[2:])
+
+    def test_0d_tensor_preserved(self):
+        d = {"a": torch.randn(4, 2), "scalar_t": torch.tensor(3.14)}
+        chunks = self.split(d, 2)
+        for c in chunks:
+            self.assertAlmostEqual(c["scalar_t"].item(), 3.14, places=5)
+
+    def test_list_split(self):
+        d = {"a": torch.randn(4, 2), "names": ["a", "b", "c", "d"]}
+        chunks = self.split(d, 2)
+        self.assertEqual(chunks[0]["names"], ["a", "b"])
+        self.assertEqual(chunks[1]["names"], ["c", "d"])
+
+
+class TestShuffleSequenceDict(unittest.TestCase):
+    """Tests for patched shuffle_sequence_dict."""
+
+    def setUp(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import _patched_shuffle_sequence_dict
+
+        self.shuffle = _patched_shuffle_sequence_dict
+
+    def test_scalar_int_preserved(self):
+        d = {"a": torch.randn(4, 3), "count": 42}
+        result = self.shuffle(d)
+        self.assertEqual(result["count"], 42)
+
+    def test_scalar_float_preserved(self):
+        d = {"a": torch.randn(4, 3), "lr": 1e-5}
+        result = self.shuffle(d)
+        self.assertEqual(result["lr"], 1e-5)
+
+    def test_scalar_bool_preserved(self):
+        d = {"a": torch.randn(4, 3), "flag": False}
+        result = self.shuffle(d)
+        self.assertFalse(result["flag"])
+
+    def test_none_preserved(self):
+        d = {"a": torch.randn(4, 3), "b": None}
+        result = self.shuffle(d)
+        self.assertIsNone(result["b"])
+
+    def test_tensor_permuted(self):
+        torch.manual_seed(42)
+        t = torch.arange(4).float()
+        d = {"a": t}
+        result = self.shuffle(d)
+        # Same elements, possibly different order
+        self.assertEqual(sorted(result["a"].tolist()), sorted(t.tolist()))
+        self.assertEqual(result["a"].shape, t.shape)
+
+    def test_list_permuted(self):
+        torch.manual_seed(42)
+        d = {"a": torch.randn(3, 2), "names": ["x", "y", "z"]}
+        result = self.shuffle(d)
+        self.assertEqual(sorted(result["names"]), ["x", "y", "z"])
+        self.assertEqual(len(result["names"]), 3)
+
+    def test_0d_tensor_preserved(self):
+        d = {"a": torch.randn(4, 2), "scalar_t": torch.tensor(3.14)}
+        result = self.shuffle(d)
+        self.assertAlmostEqual(result["scalar_t"].item(), 3.14, places=5)
+
+
+class TestExtractLogprobs(unittest.TestCase):
+    """Tests for patched extract_logprobs (NaN → 0.0)."""
+
+    def setUp(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import _patched_extract_logprobs
+
+        self.extract = _patched_extract_logprobs
+
+    def _make_output(self, logprob_values):
+        """Create a mock vLLM RequestOutput with given logprob values."""
+
+        @dataclass
+        class LogprobItem:
+            logprob: float
+            rank: int
+
+        @dataclass
+        class SeqOutput:
+            logprobs: list[dict[int, LogprobItem]] | None
+
+        @dataclass
+        class RequestOutput:
+            outputs: list[SeqOutput]
+
+        logprobs_list = []
+        for vals in logprob_values:
+            lp_dict = {i: LogprobItem(logprob=v, rank=i) for i, v in enumerate(vals)}
+            logprobs_list.append(lp_dict)
+
+        return RequestOutput(outputs=[SeqOutput(logprobs=logprobs_list)])
+
+    def test_nan_replaced_with_zero(self):
+        output = self._make_output([[float("nan"), 0.5], [-0.3, float("nan")]])
+        logprobs, token_ids = self.extract([output])
+        self.assertEqual(logprobs[0][0][0], 0.0)  # NaN → 0.0
+        self.assertEqual(logprobs[0][0][1], 0.5)
+        self.assertEqual(logprobs[0][1][0], -0.3)
+        self.assertEqual(logprobs[0][1][1], 0.0)  # NaN → 0.0
+
+    def test_normal_values_preserved(self):
+        output = self._make_output([[-0.5, -1.2], [-0.1, -2.0]])
+        logprobs, token_ids = self.extract([output])
+        self.assertAlmostEqual(logprobs[0][0][0], -0.5)
+        self.assertAlmostEqual(logprobs[0][0][1], -1.2)
+
+    def test_none_logprobs_returns_none(self):
+        @dataclass
+        class SeqOutput:
+            logprobs: None = None
+
+        @dataclass
+        class RequestOutput:
+            outputs: list
+
+        output = RequestOutput(outputs=[SeqOutput()])
+        logprobs, token_ids = self.extract([output])
+        self.assertIsNone(logprobs)
+        self.assertIsNone(token_ids)
+
+    def test_token_ids_extracted(self):
+        output = self._make_output([[-0.5]])
+        logprobs, token_ids = self.extract([output])
+        self.assertEqual(token_ids[0][0], [0])  # token_id=0 from enumerate
+
+
+class TestPatchApplication(unittest.TestCase):
+    """Tests for patch_trl_vllm() application."""
+
+    def test_batch_update_added_to_client(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import patch_trl_vllm
+
+        patch_trl_vllm()
+        from trl.generation.vllm_client import VLLMClient
+
+        self.assertTrue(hasattr(VLLMClient, "batch_update_named_params"))
+
+    def test_extract_logprobs_patched(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import (
+            _patched_extract_logprobs,
+            patch_trl_vllm,
+        )
+
+        patch_trl_vllm()
+        from trl.generation import vllm_generation
+
+        self.assertIs(vllm_generation.extract_logprobs, _patched_extract_logprobs)
+
+    def test_utils_patched(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import (
+            _patched_shuffle_sequence_dict,
+            _patched_split_tensor_dict,
+            patch_trl_vllm,
+        )
+
+        patch_trl_vllm()
+        import trl.trainer.utils
+
+        self.assertIs(trl.trainer.utils.split_tensor_dict, _patched_split_tensor_dict)
+        self.assertIs(
+            trl.trainer.utils.shuffle_sequence_dict, _patched_shuffle_sequence_dict
+        )
+
+    def test_patch_idempotent(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import patch_trl_vllm
+
+        patch_trl_vllm()
+        patch_trl_vllm()  # second call should not error
+        from trl.generation.vllm_client import VLLMClient
+
+        self.assertTrue(hasattr(VLLMClient, "batch_update_named_params"))
+
+
+class TestBatchUpdateChunking(unittest.TestCase):
+    """Tests for batch_update_named_params chunking logic."""
+
+    def test_no_chunk_single_batch(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import _batch_update_named_params
+
+        # Test that with chunk_size=None, all params go in one chunk
+        client = MagicMock()
+        client.base_url = "http://localhost:8000"
+        client.session.post.return_value = MagicMock(status_code=200)
+        client.communicator = MagicMock()
+        client.communicator.group = MagicMock()
+        client.rank = 0
+
+        params = [
+            ("layer.0.weight", torch.randn(10, 10)),
+            ("layer.1.weight", torch.randn(10, 10)),
+        ]
+        _batch_update_named_params(client, params, chunk_size=None)
+
+        # Should make exactly 1 HTTP call
+        self.assertEqual(client.session.post.call_count, 1)
+
+    def test_chunk_splits_params(self):
+        from axolotl.monkeypatch.trainer.trl_vllm import _batch_update_named_params
+
+        client = MagicMock()
+        client.base_url = "http://localhost:8000"
+        client.session.post.return_value = MagicMock(status_code=200)
+        client.communicator = MagicMock()
+        client.communicator.group = MagicMock()
+        client.rank = 0
+
+        params = [
+            ("a", torch.randn(100)),  # 100 elements
+            ("b", torch.randn(100)),  # 100 elements
+            ("c", torch.randn(100)),  # 100 elements
+        ]
+        _batch_update_named_params(client, params, chunk_size=150)
+
+        # Should make 2 HTTP calls: [a,b] then [c] (100+100 > 150 triggers split)
+        # Actually: a=100 < 150, a+b=200 > 150 → chunk [a], then b=100 < 150,
+        # b+c=200 > 150 → chunk [b], then [c]. So 3 calls.
+        # Wait: first a added (100 < 150), then b: 100+100=200 > 150, so chunk=[a],
+        # new chunk starts with b (100 < 150), then c: 100+100=200 > 150, so chunk=[b],
+        # final chunk=[c]. 3 HTTP calls.
+        self.assertEqual(client.session.post.call_count, 3)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/monkeypatch/test_voxtral_modeling_patch.py b/tests/monkeypatch/test_voxtral_modeling_patch.py
new file mode 100644
index 0000000000..878bbc1855
--- /dev/null
+++ b/tests/monkeypatch/test_voxtral_modeling_patch.py
@@ -0,0 +1,43 @@
+"""Integration tests for Voxtral modeling patches."""
+
+import pytest
+
+
+class TestVoxtralModelingPatchIntegration:
+    """Test Voxtral modeling patch integration."""
+
+    @pytest.mark.integration
+    def test_voxtral_conditional_generation_patch(self):
+        """Test that Voxtral conditional generation patch can be applied."""
+        try:
+            from transformers.models.voxtral.modeling_voxtral import (
+                VoxtralForConditionalGeneration,
+            )
+        except ImportError:
+            pytest.skip("VoxtralForConditionalGeneration not available")
+
+        from axolotl.monkeypatch.models.voxtral.modeling import (
+            patch_voxtral_conditional_generation_forward,
+        )
+
+        # Store original method
+        original_forward = VoxtralForConditionalGeneration.forward
+
+        # Apply patch and get unpatch function
+        unpatch_fn = patch_voxtral_conditional_generation_forward()
+
+        # Verify patch was applied
+        assert VoxtralForConditionalGeneration.forward != original_forward, (
+            "forward method was not patched"
+        )
+
+        # Verify the method is still callable
+        assert callable(VoxtralForConditionalGeneration.forward), (
+            "Patched method is not callable"
+        )
+
+        # Test unpatch function
+        unpatch_fn()
+        assert VoxtralForConditionalGeneration.forward == original_forward, (
+            "unpatch function did not restore original method"
+        )
diff --git a/tests/test_validation.py b/tests/patched/test_validation.py
similarity index 62%
rename from tests/test_validation.py
rename to tests/patched/test_validation.py
index 27824f2887..23773448fd 100644
--- a/tests/test_validation.py
+++ b/tests/patched/test_validation.py
@@ -1,7 +1,5 @@
-# pylint: disable=too-many-lines
 """Module for testing the validation module"""
 
-import logging
 import os
 import warnings
 from typing import Optional
@@ -9,10 +7,13 @@
 import pytest
 from pydantic import ValidationError
 
+from axolotl.loaders.utils import check_model_config
+from axolotl.utils import is_comet_available
 from axolotl.utils.config import validate_config
-from axolotl.utils.config.models.input.v0_4_1 import AxolotlConfigWCapabilities
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import check_model_config
+from axolotl.utils.mlflow_ import setup_mlflow_env_vars
+from axolotl.utils.schemas.config import AxolotlConfigWCapabilities
+from axolotl.utils.schemas.datasets import SFTDataset
 from axolotl.utils.wandb_ import setup_wandb_env_vars
 
 warnings.filterwarnings("error")
@@ -48,7 +49,6 @@ def inject_fixtures(self, caplog):
         self._caplog = caplog
 
 
-# pylint: disable=too-many-public-methods
 class TestValidation(BaseValidation):
     """
     Test the validation module
@@ -66,6 +66,53 @@ def test_defaults(self, minimal_cfg):
         assert cfg.train_on_inputs is False
         assert cfg.weight_decay is None
 
+    def test_zero3_qlora_use_reentrant_false(self, minimal_cfg):
+        test_cfg = DictDefault(
+            {
+                "deepspeed": "deepspeed_configs/zero3_bf16.json",
+                "gradient_checkpointing": True,
+                "gradient_checkpointing_kwargs": {"use_reentrant": False},
+                "load_in_4bit": True,
+                "adapter": "qlora",
+            }
+            | minimal_cfg
+        )
+
+        with self._caplog.at_level("WARNING"):
+            validate_config(test_cfg)
+            assert (
+                "qlora + zero3 with use_reentrant: false may result in a CheckpointError about recomputed values"
+                in self._caplog.records[0].message
+            )
+
+    def test_deepspeed_empty(self, minimal_cfg):
+        test_cfg = DictDefault(
+            {
+                "deepspeed": "",
+                "gradient_checkpointing": True,
+                "gradient_checkpointing_kwargs": {"use_reentrant": False},
+                "load_in_4bit": True,
+                "adapter": "qlora",
+            }
+            | minimal_cfg
+        )
+
+        _ = validate_config(test_cfg)
+
+    def test_deepspeed_not_set(self, minimal_cfg):
+        test_cfg = DictDefault(
+            {
+                "deepspeed": None,
+                "gradient_checkpointing": True,
+                "gradient_checkpointing_kwargs": {"use_reentrant": False},
+                "load_in_4bit": True,
+                "adapter": "qlora",
+            }
+            | minimal_cfg
+        )
+
+        _ = validate_config(test_cfg)
+
     def test_datasets_min_length(self):
         cfg = DictDefault(
             {
@@ -169,7 +216,7 @@ def test_batch_size_unused_warning(self):
             }
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert "batch_size is not recommended" in self._caplog.records[0].message
 
@@ -193,7 +240,7 @@ def test_batch_size_more_params(self):
 
     def test_lr_as_float(self, minimal_cfg):
         cfg = (
-            DictDefault(  # pylint: disable=unsupported-binary-operation
+            DictDefault(
                 {
                     "learning_rate": "5e-5",
                 }
@@ -231,6 +278,34 @@ def test_model_type_remap(self, minimal_cfg):
         new_cfg = validate_config(cfg)
         assert new_cfg.type_of_model == "AutoModelForCausalLM"
 
+    def test_reward_model_defaults(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "reward_model": True,
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg.num_labels == 1
+        assert new_cfg.type_of_model == "AutoModelForSequenceClassification"
+
+    def test_process_reward_model_defaults(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "process_reward_model": True,
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg.num_labels == 2
+        assert new_cfg.type_of_model == "AutoModelForTokenClassification"
+
     def test_model_revision_remap(self, minimal_cfg):
         cfg = (
             DictDefault(
@@ -255,7 +330,7 @@ def test_qlora(self, minimal_cfg):
         )
 
         cfg = (
-            DictDefault(  # pylint: disable=unsupported-binary-operation
+            DictDefault(
                 {
                     "load_in_8bit": True,
                 }
@@ -267,7 +342,7 @@ def test_qlora(self, minimal_cfg):
             validate_config(cfg)
 
         cfg = (
-            DictDefault(  # pylint: disable=unsupported-binary-operation
+            DictDefault(
                 {
                     "gptq": True,
                 }
@@ -279,7 +354,7 @@ def test_qlora(self, minimal_cfg):
             validate_config(cfg)
 
         cfg = (
-            DictDefault(  # pylint: disable=unsupported-binary-operation
+            DictDefault(
                 {
                     "load_in_4bit": False,
                 }
@@ -291,7 +366,7 @@ def test_qlora(self, minimal_cfg):
             validate_config(cfg)
 
         cfg = (
-            DictDefault(  # pylint: disable=unsupported-binary-operation
+            DictDefault(
                 {
                     "load_in_4bit": True,
                 }
@@ -313,7 +388,7 @@ def test_qlora_merge(self, minimal_cfg):
         )
 
         cfg = (
-            DictDefault(  # pylint: disable=unsupported-binary-operation
+            DictDefault(
                 {
                     "load_in_8bit": True,
                 }
@@ -325,7 +400,7 @@ def test_qlora_merge(self, minimal_cfg):
             validate_config(cfg)
 
         cfg = (
-            DictDefault(  # pylint: disable=unsupported-binary-operation
+            DictDefault(
                 {
                     "gptq": True,
                 }
@@ -337,7 +412,7 @@ def test_qlora_merge(self, minimal_cfg):
             validate_config(cfg)
 
         cfg = (
-            DictDefault(  # pylint: disable=unsupported-binary-operation
+            DictDefault(
                 {
                     "load_in_4bit": True,
                 }
@@ -464,7 +539,7 @@ def test_flash_optimum(self, minimal_cfg):
             | minimal_cfg
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert any(
                 "BetterTransformers probably doesn't work with PEFT adapters"
@@ -482,7 +557,7 @@ def test_flash_optimum(self, minimal_cfg):
             | minimal_cfg
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert any(
                 "probably set bfloat16 or float16" in record.message
@@ -528,7 +603,7 @@ def test_adamw_hyperparams(self, minimal_cfg):
             | minimal_cfg
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert any(
                 "adamw hyperparameters found, but no adamw optimizer set"
@@ -546,7 +621,7 @@ def test_adamw_hyperparams(self, minimal_cfg):
             | minimal_cfg
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert any(
                 "adamw hyperparameters found, but no adamw optimizer set"
@@ -599,13 +674,13 @@ def test_packing(self, minimal_cfg):
             DictDefault(
                 {
                     "sample_packing": True,
-                    "pad_to_sequence_len": None,
+                    "pad_to_sequence_len": False,
                     "flash_attention": True,
                 }
             )
             | minimal_cfg
         )
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert any(
                 "`pad_to_sequence_len: true` is recommended when using sample_packing"
@@ -613,69 +688,63 @@ def test_packing(self, minimal_cfg):
                 for record in self._caplog.records
             )
 
-    def test_merge_lora_no_bf16_fail(self, minimal_cfg):
-        """
-        This is assumed to be run on a CPU machine, so bf16 is not supported.
-        """
-
+    def test_packing_autoset(self, minimal_cfg):
         cfg = (
             DictDefault(
                 {
-                    "bf16": True,
-                    "capabilities": {"bf16": False},
+                    "sample_packing": True,
+                    "pad_to_sequence_len": None,
+                    "flash_attention": True,
                 }
             )
             | minimal_cfg
         )
+        with self._caplog.at_level("INFO"):
+            cfg = validate_config(cfg)
+            assert any(
+                "Setting `pad_to_sequence_len: true` to prevent memory leaks when sample_packing"
+                in record.message
+                for record in self._caplog.records
+            )
+            assert cfg.pad_to_sequence_len is True
 
-        with pytest.raises(ValueError, match=r".*AMP is not supported on this GPU*"):
-            AxolotlConfigWCapabilities(**cfg.to_dict())
+    def test_merge_lora_no_bf16_fail(self, minimal_cfg):
+        """
+        This is assumed to be run on a CPU machine, so bf16 is not supported.
+        """
 
         cfg = (
             DictDefault(
                 {
                     "bf16": True,
-                    "merge_lora": True,
                     "capabilities": {"bf16": False},
+                    "env_capabilities": {
+                        "torch_version": "2.6.0",
+                    },
                 }
             )
             | minimal_cfg
         )
 
-        validate_config(cfg)
-
-    def test_sharegpt_deprecation(self, minimal_cfg):
-        cfg = (
-            DictDefault(
-                {"datasets": [{"path": "lorem/ipsum", "type": "sharegpt:chat"}]}
-            )
-            | minimal_cfg
-        )
-        with self._caplog.at_level(logging.WARNING):
-            new_cfg = validate_config(cfg)
+        with self._caplog.at_level("WARNING"):
+            AxolotlConfigWCapabilities(**cfg.to_dict())
             assert any(
-                "`type: sharegpt:chat` will soon be deprecated." in record.message
+                "AMP is not supported" in record.message
                 for record in self._caplog.records
             )
-        assert new_cfg.datasets[0].type == "sharegpt"
 
         cfg = (
             DictDefault(
                 {
-                    "datasets": [
-                        {"path": "lorem/ipsum", "type": "sharegpt_simple:load_role"}
-                    ]
+                    "bf16": True,
+                    "merge_lora": True,
+                    "capabilities": {"bf16": False},
                 }
             )
             | minimal_cfg
         )
-        with self._caplog.at_level(logging.WARNING):
-            new_cfg = validate_config(cfg)
-            assert any(
-                "`type: sharegpt_simple` will soon be deprecated." in record.message
-                for record in self._caplog.records
-            )
-        assert new_cfg.datasets[0].type == "sharegpt:load_role"
+
+        validate_config(cfg)
 
     def test_no_conflict_save_strategy(self, minimal_cfg):
         cfg = (
@@ -757,7 +826,7 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
         cfg = (
             DictDefault(
                 {
-                    "evaluation_strategy": "epoch",
+                    "eval_strategy": "epoch",
                     "eval_steps": 10,
                 }
             )
@@ -765,14 +834,14 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
         )
 
         with pytest.raises(
-            ValueError, match=r".*evaluation_strategy and eval_steps mismatch.*"
+            ValueError, match=r".*eval_strategy and eval_steps mismatch.*"
         ):
             validate_config(cfg)
 
         cfg = (
             DictDefault(
                 {
-                    "evaluation_strategy": "no",
+                    "eval_strategy": "no",
                     "eval_steps": 10,
                 }
             )
@@ -780,14 +849,14 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
         )
 
         with pytest.raises(
-            ValueError, match=r".*evaluation_strategy and eval_steps mismatch.*"
+            ValueError, match=r".*eval_strategy and eval_steps mismatch.*"
         ):
             validate_config(cfg)
 
         cfg = (
             DictDefault(
                 {
-                    "evaluation_strategy": "steps",
+                    "eval_strategy": "steps",
                 }
             )
             | minimal_cfg
@@ -798,7 +867,7 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
         cfg = (
             DictDefault(
                 {
-                    "evaluation_strategy": "steps",
+                    "eval_strategy": "steps",
                     "eval_steps": 10,
                 }
             )
@@ -821,7 +890,7 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
         cfg = (
             DictDefault(
                 {
-                    "evaluation_strategy": "no",
+                    "eval_strategy": "no",
                 }
             )
             | minimal_cfg
@@ -832,7 +901,7 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
         cfg = (
             DictDefault(
                 {
-                    "evaluation_strategy": "epoch",
+                    "eval_strategy": "epoch",
                     "val_set_size": 0,
                 }
             )
@@ -841,7 +910,7 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
 
         with pytest.raises(
             ValueError,
-            match=r".*eval_steps and evaluation_strategy are not supported with val_set_size == 0.*",
+            match=r".*eval_steps and eval_strategy are not supported with val_set_size == 0.*",
         ):
             validate_config(cfg)
 
@@ -857,7 +926,7 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
 
         with pytest.raises(
             ValueError,
-            match=r".*eval_steps and evaluation_strategy are not supported with val_set_size == 0.*",
+            match=r".*eval_steps and eval_strategy are not supported with val_set_size == 0.*",
         ):
             validate_config(cfg)
 
@@ -887,7 +956,7 @@ def test_no_conflict_eval_strategy(self, minimal_cfg):
         cfg = (
             DictDefault(
                 {
-                    "evaluation_strategy": "epoch",
+                    "eval_strategy": "epoch",
                     "val_set_size": 0.01,
                 }
             )
@@ -1070,7 +1139,7 @@ def test_unfrozen_parameters_w_peft_layers_to_transform(self, minimal_cfg):
     def test_hub_model_id_save_value_warns_save_stragey_no(self, minimal_cfg):
         cfg = DictDefault({"hub_model_id": "test", "save_strategy": "no"}) | minimal_cfg
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert len(self._caplog.records) == 1
 
@@ -1079,7 +1148,7 @@ def test_hub_model_id_save_value_warns_random_value(self, minimal_cfg):
             DictDefault({"hub_model_id": "test", "save_strategy": "test"}) | minimal_cfg
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert len(self._caplog.records) == 1
 
@@ -1089,7 +1158,7 @@ def test_hub_model_id_save_value_steps(self, minimal_cfg):
             | minimal_cfg
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert len(self._caplog.records) == 0
 
@@ -1099,24 +1168,193 @@ def test_hub_model_id_save_value_epochs(self, minimal_cfg):
             | minimal_cfg
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert len(self._caplog.records) == 0
 
     def test_hub_model_id_save_value_none(self, minimal_cfg):
         cfg = DictDefault({"hub_model_id": "test", "save_strategy": None}) | minimal_cfg
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert len(self._caplog.records) == 0
 
     def test_hub_model_id_save_value_no_set_save_strategy(self, minimal_cfg):
         cfg = DictDefault({"hub_model_id": "test"}) | minimal_cfg
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             validate_config(cfg)
             assert len(self._caplog.records) == 0
 
+    def test_dpo_beta_deprecation(self, minimal_cfg):
+        cfg = DictDefault({"dpo_beta": 0.2}) | minimal_cfg
+
+        with self._caplog.at_level("WARNING"):
+            new_cfg = validate_config(cfg)
+            assert new_cfg["rl_beta"] == 0.2
+            assert new_cfg["dpo_beta"] is None
+            assert len(self._caplog.records) == 1
+
+    def test_eval_strategy_remap(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "evaluation_strategy": "steps",
+                }
+            )
+            | minimal_cfg
+        )
+
+        with self._caplog.at_level("WARNING"):
+            new_cfg = validate_config(cfg)
+            assert new_cfg.eval_strategy == "steps"
+            assert (
+                "evaluation_strategy is deprecated, use eval_strategy instead"
+                in self._caplog.records[0].message
+            )
+
+    def test_torch_version_adopt_req(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "optimizer": "adopt_adamw",
+                }
+            )
+            | minimal_cfg
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=r".*ADOPT optimizer is incompatible with torch version*",
+        ):
+            env_capabilities = {"torch_version": "2.3.0"}
+            capabilities = {"bf16": False}
+            _ = validate_config(
+                cfg, capabilities=capabilities, env_capabilities=env_capabilities
+            )
+
+        env_capabilities = {"torch_version": "2.6.0"}
+        capabilities = {"bf16": False}
+        _ = validate_config(
+            cfg, capabilities=capabilities, env_capabilities=env_capabilities
+        )
+
+        env_capabilities = {"torch_version": "2.5.2"}
+        capabilities = {"bf16": False}
+        _ = validate_config(
+            cfg, capabilities=capabilities, env_capabilities=env_capabilities
+        )
+
+
+class TestTorchCompileValidation(BaseValidation):
+    """
+    test suite for when torch_compile is set to 'auto'
+    """
+
+    def test_torch_compile_auto(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "torch_compile": "auto",
+                }
+            )
+            | minimal_cfg
+        )
+
+        env_capabilities = {"torch_version": "2.6.0"}
+        capabilities = {"bf16": True}
+        updated_cfg = validate_config(
+            cfg, capabilities=capabilities, env_capabilities=env_capabilities
+        )
+
+        assert updated_cfg.torch_compile is True
+
+        env_capabilities = {"torch_version": "2.4.1"}
+        capabilities = {"bf16": True}
+        updated_cfg = validate_config(
+            cfg, capabilities=capabilities, env_capabilities=env_capabilities
+        )
+
+        assert updated_cfg.torch_compile is False
+
+        env_capabilities = {}
+        capabilities = {"bf16": True}
+        updated_cfg = validate_config(
+            cfg, capabilities=capabilities, env_capabilities=env_capabilities
+        )
+
+        assert updated_cfg.torch_compile is False
+
+
+class TestSampleOptimConfigValidation(BaseValidation):
+    """
+    test configurations for sample optimizations like batch flattening
+    """
+
+    def test_batch_flattening_auto_enables(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "flash_attention": True,
+                    "sample_packing": None,
+                    "micro_batch_size": 2,
+                    "batch_flattening": "auto",
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg["batch_flattening"] is True
+
+    def test_batch_flattening_auto_no_fa(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "flash_attention": False,
+                    "sample_packing": None,
+                    "micro_batch_size": 2,
+                    "batch_flattening": "auto",
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg["batch_flattening"] is False
+
+    def test_batch_flattening_auto_mbsz_1(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "flash_attention": True,
+                    "sample_packing": None,
+                    "micro_batch_size": 1,
+                    "batch_flattening": "auto",
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg["batch_flattening"] is False
+
+    def test_batch_flattening_auto_packing(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "flash_attention": True,
+                    "sample_packing": True,
+                    "micro_batch_size": 2,
+                    "batch_flattening": "auto",
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg["batch_flattening"] is False
+
 
 class TestValidationCheckModelConfig(BaseValidation):
     """
@@ -1233,7 +1471,7 @@ def test_wandb_set_run_id_to_name(self, minimal_cfg):
             | minimal_cfg
         )
 
-        with self._caplog.at_level(logging.WARNING):
+        with self._caplog.at_level("WARNING"):
             new_cfg = validate_config(cfg)
             assert any(
                 "wandb_run_id sets the ID of the run. If you would like to set the name, please use wandb_name instead."
@@ -1283,7 +1521,6 @@ def test_wandb_sets_env(self, minimal_cfg):
         assert os.environ.get("WANDB_MODE", "") == "online"
         assert os.environ.get("WANDB_WATCH", "") == "false"
         assert os.environ.get("WANDB_LOG_MODEL", "") == "checkpoint"
-        assert os.environ.get("WANDB_DISABLED", "") != "true"
 
         os.environ.pop("WANDB_PROJECT", None)
         os.environ.pop("WANDB_NAME", None)
@@ -1292,21 +1529,141 @@ def test_wandb_sets_env(self, minimal_cfg):
         os.environ.pop("WANDB_MODE", None)
         os.environ.pop("WANDB_WATCH", None)
         os.environ.pop("WANDB_LOG_MODEL", None)
-        os.environ.pop("WANDB_DISABLED", None)
 
     def test_wandb_set_disabled(self, minimal_cfg):
         cfg = DictDefault({}) | minimal_cfg
-
         new_cfg = validate_config(cfg)
+        setup_wandb_env_vars(new_cfg)
+        assert new_cfg.use_wandb is None
 
+        cfg = (
+            DictDefault(
+                {
+                    "wandb_project": "foo",
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
         setup_wandb_env_vars(new_cfg)
+        assert new_cfg.use_wandb is True
+
+        os.environ.pop("WANDB_PROJECT", None)
+
+
+@pytest.mark.skipif(is_comet_available() is False, reason="comet_ml is not installed")
+class TestValidationComet(BaseValidation):
+    """
+    Validation test for comet
+    """
+
+    def test_comet_sets_env(self, minimal_cfg):
+        from axolotl.utils.comet_ import setup_comet_env_vars
+
+        comet_config = {
+            "comet_api_key": "foo",
+            "comet_workspace": "some_workspace",
+            "comet_project_name": "some_project",
+            "comet_experiment_key": "some_experiment_key",
+            "comet_mode": "get_or_create",
+            "comet_online": False,
+            "comet_experiment_config": {
+                "auto_histogram_activation_logging": False,
+                "auto_histogram_epoch_rate": 2,
+                "auto_histogram_gradient_logging": True,
+                "auto_histogram_tensorboard_logging": False,
+                "auto_histogram_weight_logging": True,
+                "auto_log_co2": False,
+                "auto_metric_logging": True,
+                "auto_metric_step_rate": 15,
+                "auto_output_logging": False,
+                "auto_param_logging": True,
+                "comet_disabled": False,
+                "display_summary_level": 2,
+                "distributed_node_identifier": "some_distributed_node_identifier",
+                "log_code": True,
+                "log_env_cpu": False,
+                "log_env_details": True,
+                "log_env_disk": False,
+                "log_env_gpu": True,
+                "log_env_host": False,
+                "log_env_network": True,
+                "log_git_metadata": False,
+                "log_git_patch": True,
+                "log_graph": False,
+                "name": "some_name",
+                "offline_directory": "some_offline_directory",
+                "parse_args": True,
+                "tags": ["tag1", "tag2"],
+            },
+        }
+
+        cfg = DictDefault(comet_config) | minimal_cfg
+
+        new_cfg = validate_config(cfg)
+
+        setup_comet_env_vars(new_cfg)
+
+        comet_env = {
+            key: value for key, value in os.environ.items() if key.startswith("COMET_")
+        }
+
+        assert (
+            len(comet_env)
+            == len(comet_config) + len(comet_config["comet_experiment_config"]) - 1
+        )
+
+        assert comet_env == {
+            "COMET_API_KEY": "foo",
+            "COMET_AUTO_LOG_CLI_ARGUMENTS": "true",
+            "COMET_AUTO_LOG_CO2": "false",
+            "COMET_AUTO_LOG_CODE": "true",
+            "COMET_AUTO_LOG_DISABLE": "false",
+            "COMET_AUTO_LOG_ENV_CPU": "false",
+            "COMET_AUTO_LOG_ENV_DETAILS": "true",
+            "COMET_AUTO_LOG_ENV_DISK": "false",
+            "COMET_AUTO_LOG_ENV_GPU": "true",
+            "COMET_AUTO_LOG_ENV_HOST": "false",
+            "COMET_AUTO_LOG_ENV_NETWORK": "true",
+            "COMET_AUTO_LOG_GIT_METADATA": "false",
+            "COMET_AUTO_LOG_GIT_PATCH": "true",
+            "COMET_AUTO_LOG_GRAPH": "false",
+            "COMET_AUTO_LOG_HISTOGRAM_ACTIVATIONS": "false",
+            "COMET_AUTO_LOG_HISTOGRAM_EPOCH_RATE": "2",
+            "COMET_AUTO_LOG_HISTOGRAM_GRADIENTS": "true",
+            "COMET_AUTO_LOG_HISTOGRAM_TENSORBOARD": "false",
+            "COMET_AUTO_LOG_HISTOGRAM_WEIGHTS": "true",
+            "COMET_AUTO_LOG_METRIC_STEP_RATE": "15",
+            "COMET_AUTO_LOG_METRICS": "true",
+            "COMET_AUTO_LOG_OUTPUT_LOGGER": "false",
+            "COMET_AUTO_LOG_PARAMETERS": "true",
+            "COMET_DISPLAY_SUMMARY_LEVEL": "2",
+            "COMET_DISTRIBUTED_NODE_IDENTIFIER": "some_distributed_node_identifier",
+            "COMET_EXPERIMENT_KEY": "some_experiment_key",
+            "COMET_OFFLINE_DIRECTORY": "some_offline_directory",
+            "COMET_PROJECT_NAME": "some_project",
+            "COMET_START_EXPERIMENT_NAME": "some_name",
+            "COMET_START_EXPERIMENT_TAGS": "tag1,tag2",
+            "COMET_START_MODE": "get_or_create",
+            "COMET_START_ONLINE": "false",
+            "COMET_WORKSPACE": "some_workspace",
+        }
 
-        assert os.environ.get("WANDB_DISABLED", "") == "true"
+        for key in comet_env.keys():
+            os.environ.pop(key, None)
 
+
+class TestValidationMLflow(BaseValidation):
+    """
+    Validation test for MLflow
+    """
+
+    def test_hf_mlflow_artifacts_config_sets_env(self, minimal_cfg):
         cfg = (
             DictDefault(
                 {
-                    "wandb_project": "foo",
+                    "hf_mlflow_log_artifacts": True,
                 }
             )
             | minimal_cfg
@@ -1314,9 +1671,152 @@ def test_wandb_set_disabled(self, minimal_cfg):
 
         new_cfg = validate_config(cfg)
 
-        setup_wandb_env_vars(new_cfg)
+        assert new_cfg.hf_mlflow_log_artifacts is True
 
-        assert os.environ.get("WANDB_DISABLED", "") != "true"
+        # Check it's not already present in env
+        assert "HF_MLFLOW_LOG_ARTIFACTS" not in os.environ
 
-        os.environ.pop("WANDB_PROJECT", None)
-        os.environ.pop("WANDB_DISABLED", None)
+        setup_mlflow_env_vars(new_cfg)
+
+        assert os.environ.get("HF_MLFLOW_LOG_ARTIFACTS") == "true"
+
+        os.environ.pop("HF_MLFLOW_LOG_ARTIFACTS", None)
+
+    def test_mlflow_not_used_by_default(self, minimal_cfg):
+        cfg = DictDefault({}) | minimal_cfg
+
+        new_cfg = validate_config(cfg)
+
+        setup_mlflow_env_vars(new_cfg)
+
+        assert cfg.use_mlflow is not True
+
+        cfg = (
+            DictDefault(
+                {
+                    "mlflow_experiment_name": "foo",
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+
+        setup_mlflow_env_vars(new_cfg)
+
+        assert new_cfg.use_mlflow is True
+
+        os.environ.pop("MLFLOW_EXPERIMENT_NAME", None)
+
+
+class TestDataloaderValidation(BaseValidation):
+    """
+    tests for dataloader_* sane defaults
+    """
+
+    def test_dataloader_auto_defaults(self, minimal_cfg):
+        cfg = minimal_cfg
+
+        new_cfg = validate_config(cfg, {"n_gpu": 8}, {"torch_version": "2.6.0"})
+
+        assert new_cfg.dataloader_num_workers == 8
+        assert new_cfg.dataloader_pin_memory is True
+        assert new_cfg.dataloader_prefetch_factor == 256
+
+
+class TestGCStepsMigration(BaseValidation):
+    """
+    Tests for gc_steps -> torch_empty_cache_steps / gc_collect_steps migration
+    """
+
+    def test_gc_steps_maps_to_new_options(self, minimal_cfg):
+        cfg = DictDefault({**minimal_cfg, "gc_steps": 10})
+
+        new_cfg = validate_config(cfg, {"n_gpu": 1}, {"torch_version": "2.6.0"})
+
+        assert new_cfg.torch_empty_cache_steps == 10
+        assert new_cfg.gc_collect_steps == 10
+
+    def test_gc_steps_negative_maps_gc_collect_only(self, minimal_cfg):
+        cfg = DictDefault({**minimal_cfg, "gc_steps": -1})
+
+        new_cfg = validate_config(cfg, {"n_gpu": 1}, {"torch_version": "2.6.0"})
+
+        # -1 means only epoch end/eval GC, not periodic; torch_empty_cache_steps
+        # should not be set for negative values
+        assert new_cfg.torch_empty_cache_steps is None
+        assert new_cfg.gc_collect_steps == -1
+
+    def test_new_options_take_precedence(self, minimal_cfg):
+        cfg = DictDefault({**minimal_cfg, "gc_steps": 10, "torch_empty_cache_steps": 5})
+
+        new_cfg = validate_config(cfg, {"n_gpu": 1}, {"torch_version": "2.6.0"})
+
+        # New options take precedence; gc_steps migration is skipped
+        assert new_cfg.torch_empty_cache_steps == 5
+        assert new_cfg.gc_collect_steps is None
+
+    def test_torch_empty_cache_steps_standalone(self, minimal_cfg):
+        cfg = DictDefault({**minimal_cfg, "torch_empty_cache_steps": 8})
+
+        new_cfg = validate_config(cfg, {"n_gpu": 1}, {"torch_version": "2.6.0"})
+
+        assert new_cfg.torch_empty_cache_steps == 8
+        assert new_cfg.gc_collect_steps is None
+
+    def test_gc_collect_steps_standalone(self, minimal_cfg):
+        cfg = DictDefault({**minimal_cfg, "gc_collect_steps": 5})
+
+        new_cfg = validate_config(cfg, {"n_gpu": 1}, {"torch_version": "2.6.0"})
+
+        assert new_cfg.gc_collect_steps == 5
+        assert new_cfg.torch_empty_cache_steps is None
+
+
+class TestSyntheticDatasetValidation(BaseValidation):
+    """
+    Tests for synthetic dataset config validation
+    """
+
+    @staticmethod
+    def _make_cfg(minimal_cfg, datasets):
+        raw = dict(minimal_cfg)
+        raw["datasets"] = datasets
+        return DictDefault(raw)
+
+    def test_synthetic_dict_config_validates(self, minimal_cfg):
+        """Synthetic dataset passed as a raw dict should not raise."""
+        cfg = self._make_cfg(
+            minimal_cfg,
+            [
+                {
+                    "path": "synthetic",
+                    "type": "_synthetic",
+                    "length": 100,
+                    "sequence_length": 64,
+                }
+            ],
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg.datasets[0]["path"] == "synthetic"
+
+    def test_synthetic_already_sft_does_not_crash(self, minimal_cfg):
+        """Synthetic dataset already parsed as SFTDataset should not raise AttributeError."""
+        sft = SFTDataset(path="synthetic", type="_synthetic")
+        cfg = self._make_cfg(minimal_cfg, [sft])
+
+        # Before the fix, this raised:
+        #   AttributeError: 'SFTDataset' object has no attribute 'get'
+        new_cfg = validate_config(cfg)
+        assert new_cfg.datasets[0]["path"] == "synthetic"
+
+    def test_non_synthetic_sft_validates(self, minimal_cfg):
+        """A regular SFT dataset should validate without being treated as synthetic."""
+        cfg = self._make_cfg(
+            minimal_cfg,
+            [{"path": "mhenrichsen/alpaca_2k_test", "type": "alpaca"}],
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg.datasets[0]["path"] == "mhenrichsen/alpaca_2k_test"
diff --git a/tests/prompt_strategies/__init__.py b/tests/prompt_strategies/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/prompt_strategies/conftest.py b/tests/prompt_strategies/conftest.py
new file mode 100644
index 0000000000..087d3b95c5
--- /dev/null
+++ b/tests/prompt_strategies/conftest.py
@@ -0,0 +1,248 @@
+"""
+shared fixtures for prompt strategies tests
+"""
+
+import pytest
+from datasets import Dataset
+from transformers import AutoTokenizer
+
+from axolotl.prompt_strategies.jinja_template_analyzer import JinjaTemplateAnalyzer
+from axolotl.utils.chat_templates import _CHAT_TEMPLATES
+
+from tests.hf_offline_utils import enable_hf_offline
+
+
+@pytest.fixture(name="assistant_dataset")
+def fixture_assistant_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "messages": [
+                    {"role": "user", "content": "hello"},
+                    {"role": "assistant", "content": "hello"},
+                    {"role": "user", "content": "goodbye"},
+                    {"role": "assistant", "content": "goodbye"},
+                ]
+            }
+        ]
+    )
+
+
+@pytest.fixture(name="sharegpt_dataset")
+def fixture_sharegpt_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "conversations": [
+                    {"from": "human", "value": "hello"},
+                    {"from": "gpt", "value": "hello"},
+                    {"from": "human", "value": "goodbye"},
+                    {"from": "gpt", "value": "goodbye"},
+                ]
+            }
+        ]
+    )
+
+
+@pytest.fixture(name="basic_dataset")
+def fixture_basic_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "conversations": [
+                    {"from": "system", "value": "You are an AI assistant."},
+                    {"from": "human", "value": "Hello"},
+                    {"from": "assistant", "value": "Hi there!"},
+                    {"from": "human", "value": "How are you?"},
+                    {"from": "assistant", "value": "I'm doing well, thank you!"},
+                ]
+            }
+        ]
+    )
+
+
+@pytest.fixture(name="toolcalling_dataset")
+def fixture_toolcalling_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "messages": [
+                    {
+                        "role": "system",
+                        "content": "You are a bot that responds to weather queries. You should reply with the unit used in the queried location.",
+                    },
+                    {
+                        "role": "user",
+                        "content": "Hey, what's the temperature in Paris right now?",
+                    },
+                    {
+                        "role": "assistant",
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "get_current_temperature",
+                                    "arguments": {
+                                        "location": "Paris, France",
+                                        "unit": "celsius",
+                                    },
+                                },
+                            }
+                        ],
+                    },
+                    {
+                        "role": "tool",
+                        "name": "get_current_temperature",
+                        "content": "22.0",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "The temperature in Paris is 22.0 degrees Celsius.",
+                    },
+                ]
+            }
+        ]
+    )
+
+
+@pytest.fixture(name="llama3_tokenizer", scope="session")
+@enable_hf_offline
+def fixture_llama3_tokenizer(
+    download_llama3_8b_instruct_model_fixture,
+):
+    tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
+
+    return tokenizer
+
+
+@pytest.fixture(name="smollm2_tokenizer", scope="session")
+@enable_hf_offline
+def fixture_smollm2_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M")
+    return tokenizer
+
+
+@pytest.fixture(name="mistralv03_tokenizer", scope="session")
+@enable_hf_offline
+def fixture_mistralv03_tokenizer(
+    download_mlx_mistral_7b_model_fixture,
+):
+    tokenizer = AutoTokenizer.from_pretrained(
+        "mlx-community/Mistral-7B-Instruct-v0.3-4bit"
+    )
+    return tokenizer
+
+
+@pytest.fixture(name="phi35_tokenizer", scope="session")
+@enable_hf_offline
+def fixture_phi35_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-mini-instruct")
+    return tokenizer
+
+
+@pytest.fixture(name="phi4_tokenizer", scope="session")
+@enable_hf_offline
+def fixture_phi4_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-4-reasoning")
+    return tokenizer
+
+
+@pytest.fixture(name="gemma2_tokenizer", scope="session")
+def fixture_gemma2_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("mlx-community/gemma-2-9b-it-4bit")
+
+    return tokenizer
+
+
+@pytest.fixture(name="magistral_tokenizer")
+def fixture_magistral_tokenizer():
+    from axolotl.utils.mistral import HFMistralTokenizer
+
+    tokenizer = HFMistralTokenizer.from_pretrained("mistralai/Magistral-Small-2506")
+    return tokenizer
+
+
+@pytest.fixture(name="devstral_tokenizer")
+def fixture_devstral_tokenizer():
+    from axolotl.utils.mistral import HFMistralTokenizer
+
+    tokenizer = HFMistralTokenizer.from_pretrained("mistralai/Devstral-Small-2505")
+    return tokenizer
+
+
+@pytest.fixture(name="devstral_1_1_tokenizer")
+def fixture_devstral_1_1_tokenizer():
+    from axolotl.utils.mistral import HFMistralTokenizer
+
+    tokenizer = HFMistralTokenizer.from_pretrained("mistralai/Devstral-Small-2507")
+    return tokenizer
+
+
+@pytest.fixture(name="qwen3_tokenizer")
+@enable_hf_offline
+def qwen3_tokenizer_fixture(
+    download_qwen3_half_billion_model,
+):  # pylint: disable=unused-argument,redefined-outer-name
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
+
+    return tokenizer
+
+
+@pytest.fixture(name="mistralv03_tokenizer_chat_template_jinja")
+def fixture_mistralv03_chat_template_jinja_w_system() -> str:
+    return '{%- if messages[0]["role"] == "system" %}\n    {%- set system_message = messages[0]["content"] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n{%- if not tools is defined %}\n    {%- set tools = none %}\n{%- endif %}\n{%- set user_messages = loop_messages | selectattr("role", "equalto", "user") | list %}\n\n{#- This block checks for alternating user/assistant messages, skipping tool calling messages #}\n{%- set ns = namespace() %}\n{%- set ns.index = 0 %}\n{%- for message in loop_messages %}\n    {%- if not (message.role == "tool" or message.role == "tool_results" or (message.tool_calls is defined and message.tool_calls is not none)) %}\n        {%- if (message["role"] == "user") != (ns.index % 2 == 0) %}\n            {{- raise_exception("After the optional system message, conversation roles must alternate user/assistant/user/assistant/...") }}\n        {%- endif %}\n        {%- set ns.index = ns.index + 1 %}\n    {%- endif %}\n{%- endfor %}\n\n{{- bos_token }}\n{%- for message in loop_messages %}\n    {%- if message["role"] == "user" %}\n        {%- if tools is not none and (message == user_messages[-1]) %}\n            {{- "[AVAILABLE_TOOLS] [" }}\n            {%- for tool in tools %}\n                {%- set tool = tool.function %}\n                {{- \'{"type": "function", "function": {\' }}\n                {%- for key, val in tool.items() if key != "return" %}\n                    {%- if val is string %}\n                        {{- \'"\' + key + \'": "\' + val + \'"\' }}\n                    {%- else %}\n                        {{- \'"\' + key + \'": \' + val|tojson }}\n                    {%- endif %}\n                    {%- if not loop.last %}\n                        {{- ", " }}\n                    {%- endif %}\n                {%- endfor %}\n                {{- "}}" }}\n                {%- if not loop.last %}\n                    {{- ", " }}\n                {%- else %}\n                    {{- "]" }}\n                {%- endif %}\n            {%- endfor %}\n            {{- "[/AVAILABLE_TOOLS]" }}\n            {%- endif %}\n        {%- if loop.first and system_message is defined %}\n            {{- "[INST] " + system_message + "\\n\\n" + message["content"] + "[/INST]" }}\n        {%- else %}\n            {{- "[INST] " + message["content"] + "[/INST]" }}\n        {%- endif %}\n    {%- elif message.tool_calls is defined and message.tool_calls is not none %}\n        {{- "[TOOL_CALLS] [" }}\n        {%- for tool_call in message.tool_calls %}\n            {%- set out = tool_call.function|tojson %}\n            {{- out[:-1] }}\n            {%- if not tool_call.id is defined or tool_call.id|length != 9 %}\n                {{- raise_exception("Tool call IDs should be alphanumeric strings with length 9!") }}\n            {%- endif %}\n            {{- \', "id": "\' + tool_call.id + \'"}\' }}\n            {%- if not loop.last %}\n                {{- ", " }}\n            {%- else %}\n                {{- "]" + eos_token }}\n            {%- endif %}\n        {%- endfor %}\n    {%- elif message["role"] == "assistant" %}\n        {{- " " + message["content"]|trim + eos_token}}\n    {%- elif message["role"] == "tool_results" or message["role"] == "tool" %}\n        {%- if message.content is defined and message.content.content is defined %}\n            {%- set content = message.content.content %}\n        {%- else %}\n            {%- set content = message.content %}\n        {%- endif %}\n        {{- \'[TOOL_RESULTS] {"content": \' + content|string + ", " }}\n        {%- if not message.tool_call_id is defined or message.tool_call_id|length != 9 %}\n            {{- raise_exception("Tool call IDs should be alphanumeric strings with length 9!") }}\n        {%- endif %}\n        {{- \'"call_id": "\' + message.tool_call_id + \'"}[/TOOL_RESULTS]\' }}\n    {%- else %}\n        {{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}\n    {%- endif %}\n{%- endfor %}\n'
+
+
+@pytest.fixture(name="gemma2_tokenizer_chat_template_jinja")
+def fixture_gemma2_chat_template_jinja_w_system() -> str:
+    return "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model\n'}}{% endif %}"
+
+
+@pytest.fixture(name="llama3_2_vision_chat_template_jinja")
+def fixture_llama3_2_vision_with_hardcoded_date() -> str:
+    """Hardcodes the date in the template to avoid the need for date logic in the prompt"""
+
+    template = _CHAT_TEMPLATES["llama3_2_vision"]
+
+    old_date_logic = """{%- if not date_string is defined %}
+    {%- if strftime_now is defined %}
+        {%- set date_string = strftime_now("%d %b %Y") %}
+    {%- else %}
+        {%- set date_string = "26 Jul 2024" %}
+    {%- endif %}
+{%- endif %}"""
+
+    new_date_logic = """{%- set date_string = "17 Dec 2024" %}"""
+
+    modified_template = template.replace(old_date_logic, new_date_logic)
+
+    return modified_template
+
+
+@pytest.fixture(name="chat_template_jinja_with_optional_fields")
+def fixture_chat_template_jinja_with_optional_fields() -> str:
+    return """{% for message in messages %}
+{{'<|im_start|>'}}{{ message['role'] }}
+{% if message['thoughts'] is defined %}[Thoughts: {{ message['thoughts'] }}]{% endif %}
+{% if message['tool_calls'] is defined %}[Tool: {{ message['tool_calls'][0]['type'] }}]{% endif %}
+{{ message['content'] }}{{'<|im_end|>'}}
+{% endfor %}"""
+
+
+@pytest.fixture(name="basic_jinja_template_analyzer")
+def basic_jinja_template_analyzer():
+    return JinjaTemplateAnalyzer(
+        """{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{'<|system|>
+' + message['content'] + '<|end|>
+'}}{% elif message['role'] == 'user' %}{{'<|user|>
+' + message['content'] + '<|end|>
+'}}{% elif message['role'] == 'assistant' %}{{'<|assistant|>
+' + message['content'] + '<|end|>
+'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>
+' }}{% else %}{{ eos_token }}{% endif %}"""
+    )
+
+
+@pytest.fixture(name="mistral_jinja_template_analyzer")
+def mistral_jinja_template_analyzer(mistralv03_tokenizer_chat_template_jinja):
+    return JinjaTemplateAnalyzer(mistralv03_tokenizer_chat_template_jinja)
diff --git a/tests/prompt_strategies/messages/__init__.py b/tests/prompt_strategies/messages/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/prompt_strategies/messages/test_chat.py b/tests/prompt_strategies/messages/test_chat.py
new file mode 100644
index 0000000000..f083232a88
--- /dev/null
+++ b/tests/prompt_strategies/messages/test_chat.py
@@ -0,0 +1,61 @@
+"""
+tests for chat_template prompt strategy
+"""
+
+import unittest
+
+from axolotl.prompt_strategies.messages.chat import load
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__, log_level="DEBUG")
+
+
+class TestMessagesChatLlama3:
+    """
+    Test class for assistant style datasets with llama-3 prompts using the messages chat llama3 strategy.
+    """
+
+    def test_llama3_load(self, llama3_tokenizer, assistant_dataset):
+        LOG.info("Loading llama-3 tokenizer with assistant dataset")
+        strategy = load(
+            llama3_tokenizer,
+            DictDefault(
+                {
+                    "train_on_inputs": False,
+                    "sequence_len": 512,
+                }
+            ),
+            DictDefault(
+                {
+                    "chat_template": "llama3",
+                    "message_field_role": "role",
+                    "message_field_content": "content",
+                    "field_messages": "messages",
+                }
+            ),
+        )
+        res = strategy.wrap_dataset(assistant_dataset)
+        input_ids = res[0]["input_ids"]
+        # fmt: off
+        expected_input_ids = [
+            128000,  # bos
+            128006, 882, 128007,  # user header
+            271, 15339, 128009,  # user prompt eot
+            128006, 78191, 128007,  # assistant header
+            271, 15339, 128009,  # assistant response eot
+            128006, 882, 128007,
+            271, 19045, 29474, 128009,
+            128006, 78191, 128007,
+            271, 19045, 29474, 128009,
+        ]
+        # fmt: on
+        LOG.debug(f"Expected input_ids: {expected_input_ids}")
+        LOG.debug(f"Actual input_ids: {input_ids}")
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/prompt_strategies/test_alpaca.py b/tests/prompt_strategies/test_alpaca.py
index 51dd5900bb..b96ebce198 100644
--- a/tests/prompt_strategies/test_alpaca.py
+++ b/tests/prompt_strategies/test_alpaca.py
@@ -1,6 +1,7 @@
 """
 Test module for alpaca integration w chatml
 """
+
 import pytest
 from datasets import Dataset
 from tokenizers import AddedToken
@@ -10,6 +11,8 @@
 from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy
 from axolotl.prompters import AlpacaPrompter, PromptStyle
 
+from tests.hf_offline_utils import enable_hf_offline
+
 
 @pytest.fixture(name="alpaca_dataset")
 def fixture_alpaca_dataset():
@@ -25,8 +28,8 @@ def fixture_alpaca_dataset():
 
 
 @pytest.fixture(name="tokenizer")
+@enable_hf_offline
 def fixture_tokenizer():
-    # pylint: disable=all
     tokenizer = AutoTokenizer.from_pretrained(
         "casperhansen/mistral-7b-instruct-v0.1-awq"
     )
diff --git a/tests/prompt_strategies/test_chat_template_ds_schema_unification.py b/tests/prompt_strategies/test_chat_template_ds_schema_unification.py
new file mode 100644
index 0000000000..4f4e32208b
--- /dev/null
+++ b/tests/prompt_strategies/test_chat_template_ds_schema_unification.py
@@ -0,0 +1,63 @@
+"""
+Tests for chat template prompt strategy with schema unification for none fields
+"""
+
+import json
+
+import pytest
+from datasets import Dataset
+
+from axolotl.prompt_strategies.chat_template import StrategyLoader
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="messages_w_tools")
+def fixture_messages_w_tools():
+    jsons = """
+{"messages":[{"role":"user","content":"move to (0, 1)"},{"role":"assistant","content":"","tool_calls":[{"function":{"name":"move","arguments":{"x":0,"y":1}}}]}],"tools":[{"type":"function","function":{"name":"move","description":"Move to a given location measured in meters","parameters":{"type":"object","properties":{"x":{"type":"number","description":"The x coordinate of the location, negative values are to the left, positive values are to the right"},"y":{"type":"number","description":"The y coordinate of the location, negative values are backward, positive values are forward"}},"required":["x","y"]}}},{"type":"function","function":{"name":"turn","description":"Turn the robot to a given direction","parameters":{"type":"object","properties":{"theta":{"type":"integer","description":"The angle to turn to, in degrees, positive values are counter-clockwise, negative values are clockwise"}},"required":["theta"]}}},{"type":"function","function":{"name":"invalid_prompt","description":"call when the user's prompt is invalid","parameters":{"type":"object","properties":{"message":{"type":"string","description":"why the prompt is invalid"}},"required":["message"]}}}],"add_generation_prompt":false}
+{"messages":[{"role":"user","content":"turn 270 degree"},{"role":"assistant","content":"","tool_calls":[{"function":{"name":"turn","arguments":{"theta": 270}}}]}],"tools":[{"type":"function","function":{"name":"move","description":"Move to a given location measured in meters","parameters":{"type":"object","properties":{"x":{"type":"number","description":"The x coordinate of the location, negative values are to the left, positive values are to the right"},"y":{"type":"number","description":"The y coordinate of the location, negative values are backward, positive values are forward"}},"required":["x","y"]}}},{"type":"function","function":{"name":"turn","description":"Turn the robot to a given direction","parameters":{"type":"object","properties":{"theta":{"type":"integer","description":"The angle to turn to, in degrees, positive values are counter-clockwise, negative values are clockwise"}},"required":["theta"]}}},{"type":"function","function":{"name":"invalid_prompt","description":"call when the user's prompt is invalid","parameters":{"type":"object","properties":{"message":{"type":"string","description":"why the prompt is invalid"}},"required":["message"]}}}],"add_generation_prompt":false}
+{"messages":[{"role":"user","content":"jump high"},{"role":"assistant","content":"","tool_calls":[{"function":{"name":"invalid_prompt","arguments":{"message": "jump is not a valid action"}}}]}],"tools":[{"type":"function","function":{"name":"move","description":"Move to a given location measured in meters","parameters":{"type":"object","properties":{"x":{"type":"number","description":"The x coordinate of the location, negative values are to the left, positive values are to the right"},"y":{"type":"number","description":"The y coordinate of the location, negative values are backward, positive values are forward"}},"required":["x","y"]}}},{"type":"function","function":{"name":"turn","description":"Turn the robot to a given direction","parameters":{"type":"object","properties":{"theta":{"type":"integer","description":"The angle to turn to, in degrees, positive values are counter-clockwise, negative values are clockwise"}},"required":["theta"]}}},{"type":"function","function":{"name":"invalid_prompt","description":"call when the user's prompt is invalid","parameters":{"type":"object","properties":{"message":{"type":"string","description":"why the prompt is invalid"}},"required":["message"]}}}],"add_generation_prompt":false}
+    """.strip().split("\n")
+    rows = [json.loads(row) for row in jsons]
+    return Dataset.from_list(rows)
+
+
+@pytest.fixture(name="qwen3_prompt_strategy")
+def qwen3_chat_template_strategy(qwen3_tokenizer):
+    cfg = DictDefault(
+        sequence_len=2048,
+        chat_template="qwen3",
+        eot_tokens=["<|im_end|>"],
+    )
+    ds_cfg = DictDefault(
+        type="chat_template",
+    )
+    load = StrategyLoader()
+    strat = load(qwen3_tokenizer, cfg, ds_cfg)
+    return strat
+
+
+class TestSchemaUnification:
+    """
+    Test class on handling null fields for tool calling
+    """
+
+    def test_schema_unification_single_prompt(
+        self, messages_w_tools, qwen3_prompt_strategy, qwen3_tokenizer
+    ):
+        for row in messages_w_tools:
+            inputs = qwen3_prompt_strategy.tokenize_prompt(row)
+            decoded = qwen3_tokenizer.decode(inputs["input_ids"])
+            tool_call = decoded.split("<tool_call>")[-1].split("</tool_call>")[0]
+            assert '"message": null' not in tool_call
+            assert '"theta": null' not in tool_call
+
+    def test_schema_unification_batched(
+        self, messages_w_tools, qwen3_prompt_strategy, qwen3_tokenizer
+    ):
+        rows = messages_w_tools.map(qwen3_prompt_strategy.tokenize_prompt, batched=True)
+        for row in rows:
+            decoded = qwen3_tokenizer.decode(row["input_ids"])
+            tool_call = decoded.split("<tool_call>")[-1].split("</tool_call>")[0]
+            assert '"message": null' not in tool_call
+            assert '"theta": null' not in tool_call
diff --git a/tests/prompt_strategies/test_chat_template_utils.py b/tests/prompt_strategies/test_chat_template_utils.py
new file mode 100644
index 0000000000..10c84f4320
--- /dev/null
+++ b/tests/prompt_strategies/test_chat_template_utils.py
@@ -0,0 +1,129 @@
+"""
+Tests for utils in axolotl.utils.chat_templates
+"""
+
+import unittest
+
+import pytest
+from transformers import AutoTokenizer
+
+from axolotl.utils.chat_templates import (
+    _CHAT_TEMPLATES,
+    extract_chat_template_args,
+    get_chat_template,
+)
+
+from tests.hf_offline_utils import enable_hf_offline
+
+
+@pytest.fixture(name="llama3_tokenizer")
+@enable_hf_offline
+def fixture_llama3_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B")
+
+    return tokenizer
+
+
+class TestGetChatTemplateUtils:
+    """
+    Tests the get_chat_template function.
+    """
+
+    def test_known_chat_template(self):
+        chat_template_str = get_chat_template("llama3")
+        assert chat_template_str == _CHAT_TEMPLATES["llama3"]
+
+    def test_invalid_chat_template(self):
+        with pytest.raises(ValueError) as exc:
+            get_chat_template("invalid_template")
+            assert str(exc) == "Template 'invalid_template' not found."
+
+    def test_tokenizer_default_no_tokenizer(self):
+        with pytest.raises(ValueError):
+            get_chat_template("tokenizer_default", tokenizer=None)
+
+    def test_tokenizer_default_no_chat_template_on_tokenizer(self, llama3_tokenizer):
+        with pytest.raises(ValueError):
+            get_chat_template("tokenizer_default", tokenizer=llama3_tokenizer)
+
+    def test_tokenizer_default_with_chat_template_on_tokenizer(self, llama3_tokenizer):
+        llama3_tokenizer.chat_template = "test_template"
+        chat_template_str = get_chat_template(
+            "tokenizer_default", tokenizer=llama3_tokenizer
+        )
+        assert chat_template_str == "test_template"
+
+    def test_tokenizer_default_fallback_no_tokenizer(self):
+        with pytest.raises(ValueError):
+            get_chat_template("tokenizer_default_fallback_test", tokenizer=None)
+
+    def test_tokenizer_default_fallback_no_chat_template_on_tokenizer(
+        self, llama3_tokenizer
+    ):
+        chat_template_str = get_chat_template(
+            "tokenizer_default_fallback_chatml", tokenizer=llama3_tokenizer
+        )
+        assert chat_template_str == get_chat_template("chatml")
+
+    def test_tokenizer_default_fallback_with_chat_template_on_tokenizer(
+        self, llama3_tokenizer
+    ):
+        llama3_tokenizer.chat_template = "test_template"
+        chat_template_str = get_chat_template(
+            "tokenizer_default_fallback_chatml", tokenizer=llama3_tokenizer
+        )
+        assert chat_template_str == "test_template"
+
+    def test_jinja_template_mode(self):
+        jinja_template = "example_jinja_template"
+        chat_template_str = get_chat_template("jinja", jinja_template=jinja_template)
+        assert chat_template_str == jinja_template
+
+    def test_jinja_template_mode_no_jinja_template(self):
+        with pytest.raises(ValueError):
+            get_chat_template("jinja", jinja_template=None)
+
+    def test_extract_chat_template_args(self):
+        # No ds_cfg
+        chat_template_choice, chat_template_jinja = extract_chat_template_args(
+            cfg={"chat_template": "chatml"},
+        )
+        assert chat_template_choice == "chatml"
+        assert chat_template_jinja is None
+
+        # ds_cfg provided
+        chat_template_choice, chat_template_jinja = extract_chat_template_args(
+            cfg={
+                "chat_template": "jinja",
+                "chat_template_jinja": "global_jinja_template",
+            },
+            ds_cfg={"chat_template": "llama3", "chat_template_jinja": None},
+        )
+        assert chat_template_choice == "llama3"
+        assert chat_template_jinja is None
+
+        # ds_cfg provided with jinja template
+        chat_template_choice, chat_template_jinja = extract_chat_template_args(
+            cfg={"chat_template": "chatml", "chat_template_jinja": None},
+            ds_cfg={
+                "chat_template": "jinja",
+                "chat_template_jinja": "ds_jinja_template",
+            },
+        )
+        assert chat_template_choice == "jinja"
+        assert chat_template_jinja == "ds_jinja_template"
+
+        # ds_cfg provided with no chat_template
+        chat_template_choice, chat_template_jinja = extract_chat_template_args(
+            cfg={
+                "chat_template": "jinja",
+                "chat_template_jinja": "global_jinja_template",
+            },
+            ds_cfg={"chat_template": None, "chat_template_jinja": "ds_jinja_template"},
+        )
+        assert chat_template_choice == "jinja"
+        assert chat_template_jinja == "global_jinja_template"
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/prompt_strategies/test_chat_templates.py b/tests/prompt_strategies/test_chat_templates.py
new file mode 100644
index 0000000000..90e0e274b7
--- /dev/null
+++ b/tests/prompt_strategies/test_chat_templates.py
@@ -0,0 +1,561 @@
+"""
+tests for chat_template prompt strategy
+"""
+
+import unittest
+
+from axolotl.prompt_strategies.chat_template import (
+    ChatTemplatePrompter,
+    ChatTemplateStrategy,
+    load,
+)
+from axolotl.prompters import IGNORE_TOKEN_ID
+from axolotl.utils.chat_templates import get_chat_template
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class TestAssistantChatTemplateLlama3:
+    """
+    Test class for assistant style datasets with llama-3 prompts using the chat_template strategy.
+    """
+
+    def test_llama3_load(self, llama3_tokenizer, assistant_dataset):
+        LOG.info("Loading llama-3 tokenizer with assistant dataset")
+        strategy = load(
+            llama3_tokenizer,
+            DictDefault(
+                {
+                    "train_on_inputs": False,
+                    "sequence_len": 512,
+                }
+            ),
+            DictDefault(
+                {
+                    "chat_template": "llama3",
+                    "message_field_role": "role",
+                    "message_field_content": "content",
+                    "message_property_mappings": {
+                        "role": "role",
+                        "content": "content",
+                    },
+                    "roles": {
+                        "user": ["user"],
+                        "assistant": ["assistant"],
+                        "system": ["system"],
+                    },
+                    "field_messages": "messages",
+                }
+            ),
+        )
+        res = strategy.tokenize_prompt(assistant_dataset[0])
+        input_ids = res["input_ids"]
+        # fmt: off
+        expected_input_ids = [
+            128000,  # bos
+            128006, 882, 128007,  # user header
+            271, 15339, 128009,  # user prompt eot
+            128006, 78191, 128007,  # assistant header
+            271, 15339, 128009,  # assistant response eot
+            128006, 882, 128007,
+            271, 19045, 29474, 128009,
+            128006, 78191, 128007,
+            271, 19045, 29474, 128009,
+        ]
+        # fmt: on
+        LOG.debug(f"Expected input_ids: {expected_input_ids}")
+        LOG.debug(f"Actual input_ids: {input_ids}")
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+
+    def test_llama3(self, llama3_tokenizer, assistant_dataset):
+        LOG.info("Testing llama-3 with assistant dataset")
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                llama3_tokenizer,
+                chat_template=get_chat_template("llama3"),
+                message_property_mappings={
+                    "role": "role",
+                    "content": "content",
+                },
+                roles={
+                    "user": ["user"],
+                    "assistant": ["assistant"],
+                    "system": ["system"],
+                },
+            ),
+            tokenizer=llama3_tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+        )
+
+        res = strategy.tokenize_prompt(assistant_dataset[0])
+        input_ids = res["input_ids"]
+        # fmt: off
+        expected_input_ids = [
+            128000,  # bos
+            128006, 882, 128007,  # user header
+            271, 15339, 128009,  # user prompt eot
+            128006, 78191, 128007,  # assistant header
+            271, 15339, 128009,   # assistant response eot
+            128006, 882, 128007,
+            271, 19045, 29474, 128009,
+            128006, 78191, 128007,
+            271, 19045, 29474, 128009,
+        ]
+        # fmt: on
+        LOG.debug(f"Expected input_ids: {expected_input_ids}")
+        LOG.debug(f"Actual input_ids: {input_ids}")
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+
+    def test_phi35(self, phi35_tokenizer, assistant_dataset):
+        LOG.info("Testing phi-3.5 with assistant dataset")
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                phi35_tokenizer,
+                chat_template=get_chat_template("phi_35"),
+                message_property_mappings={
+                    "role": "role",
+                    "content": "content",
+                },
+                roles={
+                    "user": ["user"],
+                    "assistant": ["assistant"],
+                    "system": ["system"],
+                },
+            ),
+            tokenizer=phi35_tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+        )
+
+        res = strategy.tokenize_prompt(assistant_dataset[0])
+        input_ids = res["input_ids"]
+        labels = res["labels"]
+        # fmt: off
+        expected_input_ids = [
+            32010,  # user
+            22172, 32007,  # user eot
+            32001,  # assistant
+            22172, 32007,  # assistant eot
+            32010,  # user
+            1781, 26966, 32007,  # user eot
+            32001,  # assistant
+            1781, 26966, 32007,  # assistant eot
+        ]
+        expected_labels = [
+            -100,  # user
+            -100, -100,  # user eot
+            -100,  # assistant
+            -100, -100,  # assistant eot,
+            -100,  # user
+            -100, -100, -100,  # user eot
+            -100,  # assistant
+            1781, 26966, 32007,  # assistant eot
+        ]
+        # fmt: on
+        LOG.debug(f"Expected input_ids: {expected_input_ids}")
+        LOG.debug(f"Actual input_ids: {input_ids}")
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+
+        LOG.debug(f"Expected labels : {expected_labels}")
+        LOG.debug(f"Actual labels : {labels}")
+        assert labels == expected_labels, (
+            f"Input IDs mismatch: {labels} != {expected_labels}"
+        )
+
+    def test_llama3_with_training_data(self, llama3_tokenizer, assistant_dataset):
+        LOG.info("Testing llama-3 with assistant dataset including training data")
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                llama3_tokenizer,
+                chat_template=get_chat_template("llama3"),
+                message_field_training="training",
+                message_property_mappings={
+                    "role": "role",
+                    "content": "content",
+                },
+                roles={
+                    "user": ["user"],
+                    "assistant": ["assistant"],
+                    "system": ["system"],
+                },
+            ),
+            tokenizer=llama3_tokenizer,
+            train_on_inputs=False,
+            train_on_eos="none",
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        prompt_tokens = strategy.prompter.build_prompt(
+            assistant_dataset[0]["messages"], False
+        )
+        prompt = llama3_tokenizer.decode(prompt_tokens, skip_special_tokens=False)
+        LOG.debug(f"Generated prompt: {prompt}")
+        res = strategy.tokenize_prompt(assistant_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+        # fmt: off
+        expected_labels = [
+            IGNORE_TOKEN_ID,  # bos
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user prompt eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant header
+            IGNORE_TOKEN_ID, 15339, IGNORE_TOKEN_ID,  # assistant response eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, 19045, 29474, IGNORE_TOKEN_ID,
+        ]
+        # fmt: on
+
+        LOG.debug(f"Expected labels: {expected_labels}")
+        LOG.debug(f"Actual labels: {labels}")
+        assert labels == expected_labels, (
+            f"Labels mismatch:\n"
+            f"Expected: {expected_labels}\n"
+            f"Actual: {labels}\n"
+            f"Input IDs: {input_ids}\n"
+        )
+
+
+class TestSharegptChatTemplateLlama3:
+    """
+    Test class for ShareGPT style datasets with llama-3 prompts using the chat_template strategy.
+    """
+
+    def test_llama3_assistant(self, llama3_tokenizer, sharegpt_dataset):
+        LOG.info("Testing ShareGPT style datasets with llama-3 assistant prompts")
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                llama3_tokenizer,
+                chat_template=get_chat_template("llama3"),
+                message_property_mappings={
+                    "role": "from",
+                    "content": "value",
+                },
+                field_messages="conversations",
+            ),
+            tokenizer=llama3_tokenizer,
+            train_on_inputs=False,
+            train_on_eos="none",
+            sequence_len=512,
+            roles_to_train=["gpt"],
+        )
+
+        res = strategy.tokenize_prompt(sharegpt_dataset[0])
+        input_ids = res["input_ids"]
+        labels = res["labels"]
+        # fmt: off
+        expected_input_ids = [
+            128000,  # bos
+            128006, 882, 128007,  # user header
+            271, 15339, 128009,  # user prompt eot
+            128006, 78191, 128007,  # assistant header
+            271, 15339, 128009,  # assistant response eot
+            128006, 882, 128007,
+            271, 19045, 29474, 128009,
+            128006, 78191, 128007,
+            271, 19045, 29474, 128009,
+        ]
+        expected_labels = [
+            IGNORE_TOKEN_ID,  # bos
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user prompt eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant header
+            IGNORE_TOKEN_ID, 15339, IGNORE_TOKEN_ID,  # assistant response eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, 19045, 29474, IGNORE_TOKEN_ID,
+        ]
+        # fmt: on
+
+        LOG.debug(f"Expected input_ids: {expected_input_ids}")
+        LOG.debug(f"Actual input_ids: {input_ids}")
+        LOG.debug(f"Expected labels: {expected_labels}")
+        LOG.debug(f"Actual labels: {labels}")
+
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+        assert labels == expected_labels, (
+            f"Labels mismatch: {labels} != {expected_labels}"
+        )
+
+    def test_llama3_human(self, llama3_tokenizer, sharegpt_dataset):
+        LOG.info("Testing ShareGPT style datasets with llama-3 human prompts")
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                llama3_tokenizer,
+                chat_template=get_chat_template("llama3"),
+                message_property_mappings={
+                    "role": "from",
+                    "content": "value",
+                },
+                field_messages="conversations",
+            ),
+            tokenizer=llama3_tokenizer,
+            train_on_inputs=False,
+            train_on_eos="none",
+            sequence_len=512,
+            roles_to_train=["human"],
+        )
+
+        res = strategy.tokenize_prompt(sharegpt_dataset[0])
+        input_ids = res["input_ids"]
+        labels = res["labels"]
+        # fmt: off
+        expected_input_ids = [
+            128000,  # bos
+            128006, 882, 128007,  # user header
+            271, 15339, 128009,  # user prompt eot
+            128006, 78191, 128007,  # assistant header
+            271, 15339, 128009,  # assistant response eot
+            128006, 882, 128007,
+            271, 19045, 29474, 128009,
+            128006, 78191, 128007,
+            271, 19045, 29474, 128009,
+        ]
+        expected_labels = [
+            IGNORE_TOKEN_ID,  # bos
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user header
+            IGNORE_TOKEN_ID, 15339, IGNORE_TOKEN_ID,  # user prompt eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant response eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, 19045, 29474, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+        ]
+        # fmt: on
+
+        LOG.debug(f"Expected input_ids: {expected_input_ids}")
+        LOG.debug(f"Actual input_ids: {input_ids}")
+        LOG.debug(f"Expected labels: {expected_labels}")
+        LOG.debug(f"Actual labels: {labels}")
+
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+        assert labels == expected_labels, (
+            f"Labels mismatch: {labels} != {expected_labels}"
+        )
+
+    def test_llama3_system_human(self, llama3_tokenizer, basic_dataset):
+        LOG.info("Testing ShareGPT style datasets with llama-3 system/human prompts")
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                llama3_tokenizer,
+                chat_template=get_chat_template("llama3"),
+                message_property_mappings={
+                    "role": "from",
+                    "content": "value",
+                },
+                field_messages="conversations",
+            ),
+            tokenizer=llama3_tokenizer,
+            train_on_inputs=False,
+            train_on_eos="none",
+            sequence_len=512,
+            roles_to_train=["system", "human"],
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        input_ids = res["input_ids"]
+        labels = res["labels"]
+        # fmt: off
+        expected_input_ids = [
+            128000,  # bos
+            128006, 9125, 128007,
+            271, 2675, 527, 459, 15592, 18328, 13, 128009,
+            128006, 882, 128007,  # user header
+            271, 9906, 128009,  # user prompt eot
+            128006, 78191, 128007,  # assistant header
+            271, 13347, 1070, 0, 128009,  # assistant response eot
+            128006, 882, 128007,
+            271, 4438, 527, 499, 30, 128009,
+            128006, 78191, 128007,
+            271, 40, 2846, 3815, 1664, 11, 9901, 499, 0, 128009,
+        ]
+        expected_labels = [
+            IGNORE_TOKEN_ID,  # bos
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # system header
+            IGNORE_TOKEN_ID, 2675, 527, 459, 15592, 18328, 13, IGNORE_TOKEN_ID,  # system prompt eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user header
+            IGNORE_TOKEN_ID, 9906, IGNORE_TOKEN_ID,  # user prompt eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant response eot
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, 4438, 527, 499, 30, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,
+        ]
+        # fmt: on
+
+        LOG.debug(f"Expected input_ids: {expected_input_ids}")
+        LOG.debug(f"Actual input_ids: {input_ids}")
+        LOG.debug(f"Expected labels: {expected_labels}")
+        LOG.debug(f"Actual labels: {labels}")
+
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+        assert labels == expected_labels, (
+            f"Labels mismatch: {labels} != {expected_labels}"
+        )
+
+
+class TestAssistantToolCallingChatTemplateLlama32Vision:
+    """
+    Test class for assistant style datasets with tool_calling prompts using the llama-32_vision chat template.
+    """
+
+    def test_llama32vision_train_on_assistant(
+        self, llama3_tokenizer, toolcalling_dataset, llama3_2_vision_chat_template_jinja
+    ):
+        LOG.info(
+            "Testing assistant style datasets with tool_calling with llama-32 chat template, training on assistant"
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                llama3_tokenizer,
+                chat_template=get_chat_template(
+                    "jinja", jinja_template=llama3_2_vision_chat_template_jinja
+                ),
+                message_property_mappings={"role": "role", "content": "content"},
+            ),
+            tokenizer=llama3_tokenizer,
+            train_on_inputs=False,
+            train_on_eos="turn",
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        res = strategy.tokenize_prompt(toolcalling_dataset[0])
+
+        input_ids = res["input_ids"]
+        labels = res["labels"]
+
+        # fmt: off
+        expected_input_ids = [
+            128000,  # bos
+            128006, 9125, 128007, 271,  # system header
+            38766, 1303, 33025, 2696, 25, 6790, 220, 2366, 18, 198, 15724, 2696, 25, 220, 1114, 3799, 220, 2366, 19, 271,  # system date prompt
+            2675, 527, 264, 11164, 430, 31680, 311, 9282, 20126, 13, 1472, 1288, 10052, 449, 279, 5089, 1511, 304, 279, 79002, 3813, 13, 128009,  # system message
+            128006, 882, 128007, 271,  # user header
+            19182, 11, 1148, 596, 279, 9499, 304, 12366, 1314, 1457, 30, 128009,  # user message
+            128006, 78191, 128007, 271,  # assistant header
+            5018, 609, 794, 330, 456, 11327, 54625, 498, 330, 14105, 794, 5324, 2588, 794, 330, 60704, 11, 9822, 498, 330, 3928, 794, 330, 66, 41347, 32075, 128009,  # assistant message
+            128006, 23799, 4690, 128007, 271,  # tool header
+            1, 1313, 13, 15, 1, 128009,  # tool message
+            128006, 78191, 128007, 271,  # assistant header
+            791, 9499, 304, 12366, 374, 220, 1313, 13, 15, 12628, 62447, 13, 128009  # assistant message
+        ]
+
+        expected_labels = [
+            IGNORE_TOKEN_ID,  # bos
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # system header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # system date prompt
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # system message
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user message
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant header
+            5018, 609, 794, 330, 456, 11327, 54625, 498, 330, 14105, 794, 5324, 2588, 794, 330, 60704, 11, 9822, 498, 330, 3928, 794, 330, 66, 41347, 32075, 128009,  # assistant message
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # tool header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # tool message
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant header
+            791, 9499, 304, 12366, 374, 220, 1313, 13, 15, 12628, 62447, 13, 128009  # assistant message
+        ]
+        # fmt: on
+
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+
+        assert labels == expected_labels, (
+            f"Labels mismatch: {labels} != {expected_labels}"
+        )
+
+    def test_llama32vision_train_on_tools(
+        self, llama3_tokenizer, toolcalling_dataset, llama3_2_vision_chat_template_jinja
+    ):
+        LOG.info(
+            "Testing assistant style datasets with tool_calling with llama-32 chat template, training on tools"
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                llama3_tokenizer,
+                chat_template=get_chat_template(
+                    "jinja", jinja_template=llama3_2_vision_chat_template_jinja
+                ),
+                message_property_mappings={"role": "role", "content": "content"},
+            ),
+            tokenizer=llama3_tokenizer,
+            train_on_inputs=False,
+            train_on_eos="turn",
+            sequence_len=512,
+            roles_to_train=["assistant", "tool"],
+        )
+
+        res = strategy.tokenize_prompt(toolcalling_dataset[0])
+
+        input_ids = res["input_ids"]
+        labels = res["labels"]
+
+        # fmt: off
+        expected_input_ids = [
+            128000,  # bos
+            128006, 9125, 128007, 271,  # system header
+            38766, 1303, 33025, 2696, 25, 6790, 220, 2366, 18, 198, 15724, 2696, 25, 220, 1114, 3799, 220, 2366, 19, 271,  # system date prompt
+            2675, 527, 264, 11164, 430, 31680, 311, 9282, 20126, 13, 1472, 1288, 10052, 449, 279, 5089, 1511, 304, 279, 79002, 3813, 13, 128009,  # system message
+            128006, 882, 128007, 271,  # user header
+            19182, 11, 1148, 596, 279, 9499, 304, 12366, 1314, 1457, 30, 128009,  # user message
+            128006, 78191, 128007, 271,  # assistant header
+            5018, 609, 794, 330, 456, 11327, 54625, 498, 330, 14105, 794, 5324, 2588, 794, 330, 60704, 11, 9822, 498, 330, 3928, 794, 330, 66, 41347, 32075, 128009,  # assistant message
+            128006, 23799, 4690, 128007, 271,  # tool header
+            1, 1313, 13, 15, 1, 128009,  # tool message
+            128006, 78191, 128007, 271,  # assistant header
+            791, 9499, 304, 12366, 374, 220, 1313, 13, 15, 12628, 62447, 13, 128009  # assistant message
+        ]
+
+        expected_labels = [
+            IGNORE_TOKEN_ID,  # bos
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # system header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # system date prompt
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # system message
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user header
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # user message
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant header
+            5018, 609, 794, 330, 456, 11327, 54625, 498, 330, 14105, 794, 5324, 2588, 794, 330, 60704, 11, 9822, 498, 330, 3928, 794, 330, 66, 41347, 32075, 128009,  # assistant message
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # tool header
+            IGNORE_TOKEN_ID, 1313, 13, 15, IGNORE_TOKEN_ID, 128009,  # tool message
+            IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID, IGNORE_TOKEN_ID,  # assistant header
+            791, 9499, 304, 12366, 374, 220, 1313, 13, 15, 12628, 62447, 13, 128009  # assistant message
+        ]
+        # fmt: on
+
+        assert input_ids == expected_input_ids, (
+            f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+        )
+
+        assert labels == expected_labels, (
+            f"Labels mismatch: {labels} != {expected_labels}"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/prompt_strategies/test_chat_templates_advanced.py b/tests/prompt_strategies/test_chat_templates_advanced.py
new file mode 100644
index 0000000000..701e4d01ea
--- /dev/null
+++ b/tests/prompt_strategies/test_chat_templates_advanced.py
@@ -0,0 +1,1906 @@
+"""
+tests for chat_template prompt strategy
+"""
+
+from copy import deepcopy
+
+import pytest
+from datasets import Dataset
+from tokenizers import AddedToken
+from transformers import PreTrainedTokenizer
+
+from axolotl.prompt_strategies.chat_template import (
+    ChatTemplatePrompter,
+    ChatTemplateStrategy,
+)
+from axolotl.prompters import IGNORE_TOKEN_ID
+from axolotl.utils.chat_templates import get_chat_template
+from axolotl.utils.logging import get_logger
+
+from tests.hf_offline_utils import enable_hf_offline
+
+LOG = get_logger(__name__)
+
+PARAMETRIZE_KEYS = "tokenizer, chat_template, chat_template_jinja, eos_token"
+PARAMETRIZE_PARAMS = [
+    ("llama3_tokenizer", "llama3", None, None),
+    ("llama3_tokenizer", "chatml", None, "<|im_end|>"),
+    (
+        "mistralv03_tokenizer",
+        "jinja",
+        "mistralv03_tokenizer_chat_template_jinja",
+        "[/INST]",
+    ),
+    (
+        "gemma2_tokenizer",
+        "jinja",
+        "gemma2_tokenizer_chat_template_jinja",
+        "<end_of_turn>",
+    ),
+    # ("phi35_tokenizer", "phi_35", None, "<|end|>"),  # seems to be broken w transformers v5
+    ("phi4_tokenizer", "phi_4", None, "<|im_end|>"),
+]
+
+
+@pytest.mark.parametrize(
+    PARAMETRIZE_KEYS,
+    PARAMETRIZE_PARAMS,
+)
+class TestChatTemplateConfigurations:
+    """
+    Test class for various configurations of ChatTemplateStrategy.
+    """
+
+    @staticmethod
+    def setup_tokenizer(
+        tokenizer_name,
+        chat_template,
+        chat_template_jinja=None,
+        eos_token=None,
+        request=None,
+        eot_token=None,
+    ) -> tuple[PreTrainedTokenizer, str]:
+        """
+        Helper function to set up the tokenizer and chat template for the test.
+        """
+        tokenizer = deepcopy(request.getfixturevalue(tokenizer_name))
+        if chat_template == "jinja":
+            chat_template_jinja = request.getfixturevalue(chat_template_jinja)
+        if eos_token:
+            tokenizer.add_special_tokens(
+                {
+                    "eos_token": AddedToken(
+                        eos_token, rstrip=False, lstrip=False, normalized=False
+                    )
+                }
+            )
+            if tokenizer.__class__.__name__ in (
+                "LlamaTokenizerFast",
+                "CodeLlamaTokenizerFast",
+            ):
+                tokenizer.update_post_processor()
+
+        if eot_token:
+            tokenizer.add_special_tokens({"additional_special_tokens": [eot_token]})
+
+        return tokenizer, chat_template_jinja
+
+    def _should_skip_turn(self, tokenizer, turn, turn_idx, start_idx, end_idx):
+        """Helper method to determine if a turn should be skipped in testing.
+        This is used to skip system messages for Mistral as the template does not output them without more turns.
+        """
+        if (
+            turn_idx == 0
+            and turn.get("from") in ["system", "context"]
+            and ("mistral" in tokenizer.name_or_path.lower())
+        ):
+            assert start_idx == -1 and end_idx == -1, (
+                "Expected system message to be skipped"
+            )
+            return True
+        return False
+
+    @enable_hf_offline
+    def test_train_on_inputs_true(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing with train_on_inputs=True")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=True,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        turns = strategy.get_conversation_thread(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Verify assistant responses are labeled
+        for i, turn in enumerate(basic_dataset[0]["conversations"]):
+            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+
+            if self._should_skip_turn(tokenizer, turn, i, start_idx, end_idx):
+                continue
+
+            decoded_response = tokenizer.decode(input_ids[start_idx:end_idx])
+            response = turn["value"]
+
+            assert response in decoded_response, (
+                f"Response {response} not found in index {start_idx}:{end_idx} "
+                f"decoded:{decoded_response}"
+            )
+
+            assert all(
+                label != IGNORE_TOKEN_ID for label in labels[start_idx:end_idx]
+            ), (
+                f"Expected labels for input '{response}' to be ignored, but got {labels[start_idx:end_idx]}"
+            )
+
+        LOG.debug("Full labels: %s", labels)
+        LOG.debug("Full input_ids: %s", input_ids)
+
+    def test_train_on_inputs_false(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing with train_on_inputs=False, on assistant only")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        turns = strategy.get_conversation_thread(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Process all turns and verify correct labeling based on role
+        for i, turn in enumerate(basic_dataset[0]["conversations"]):
+            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+
+            if self._should_skip_turn(tokenizer, turn, i, start_idx, end_idx):
+                continue
+
+            decoded_response = tokenizer.decode(input_ids[start_idx:end_idx])
+            response = turn["value"]
+
+            assert response in decoded_response, (
+                f"Response {response} not found in index {start_idx}:{end_idx} "
+                f"decoded:{decoded_response}"
+            )
+
+            # Verify that assistant responses are labeled and other inputs are not
+            is_assistant = turn["from"] == "assistant"
+            if is_assistant:
+                assert all(
+                    label != IGNORE_TOKEN_ID for label in labels[start_idx:end_idx]
+                ), (
+                    f"Expected labels for assistant response '{response}' to be set, but got {labels[start_idx:end_idx]}"
+                )
+            else:
+                assert all(
+                    label == IGNORE_TOKEN_ID for label in labels[start_idx:end_idx]
+                ), (
+                    f"Expected labels for human input '{response}' to be IGNORE_TOKEN_ID, but got {labels[start_idx:end_idx]}"
+                )
+
+    def test_roles_to_train_human_assistant_only(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing roles_to_train with human assistant only")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant", "human"],
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        turns = strategy.get_conversation_thread(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Process all turns and verify correct labeling based on role
+        for i, turn in enumerate(basic_dataset[0]["conversations"]):
+            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+
+            if self._should_skip_turn(tokenizer, turn, i, start_idx, end_idx):
+                continue
+
+            decoded_response = tokenizer.decode(input_ids[start_idx:end_idx])
+            response = turn["value"]
+
+            assert response in decoded_response, (
+                f"Response {response} not found in index {start_idx}:{end_idx} "
+                f"decoded:{decoded_response}"
+            )
+
+            # Verify that non-system responses are labeled and system are not
+            should_be_labelled = turn["from"] != "system"
+            if should_be_labelled:
+                assert all(
+                    label != IGNORE_TOKEN_ID for label in labels[start_idx:end_idx]
+                ), (
+                    f"Expected labels for assistant response '{response}' to be set, but got {labels[start_idx:end_idx]}"
+                )
+            else:
+                assert all(
+                    label == IGNORE_TOKEN_ID for label in labels[start_idx:end_idx]
+                ), (
+                    f"Expected labels for human input '{response}' to be IGNORE_TOKEN_ID, but got {labels[start_idx:end_idx]}"
+                )
+
+    def test_roles_to_train_all(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing roles_to_train with all roles")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=True,
+            sequence_len=512,
+            roles_to_train=["human", "assistant"],
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        turns = strategy.get_conversation_thread(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Verify that all responses are labeled (except for special tokens)
+        for i, turn in enumerate(basic_dataset[0]["conversations"]):
+            response = turn["value"]
+
+            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+
+            if self._should_skip_turn(tokenizer, turn, i, start_idx, end_idx):
+                continue
+
+            decoded_response = tokenizer.decode(input_ids[start_idx:end_idx])
+            assert response in decoded_response, (
+                f"Response {response} not found in index {start_idx}:{end_idx} decoded:{decoded_response}"
+            )
+
+            assert all(
+                label != IGNORE_TOKEN_ID for label in labels[start_idx:end_idx]
+            ), (
+                f"Expected labels for response '{response}' to be set, but got {labels[start_idx:end_idx]}"
+            )
+
+    def test_empty_roles_to_train(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing with empty roles_to_train")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=[],
+            train_on_eos="none",  # Add this line
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        labels = res["labels"]
+
+        # Verify that no labels are set when roles_to_train is empty
+        LOG.debug("Full labels: %s", labels)
+        assert all(label == IGNORE_TOKEN_ID for label in labels), (
+            "Expected all labels to be IGNORE_TOKEN_ID when roles_to_train is empty"
+        )
+
+    def test_train_on_eos_all(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing with train_on_eos='all'")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+            train_on_eos="all",
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        eos_token_id = tokenizer.eos_token_id
+        eos_indices = [
+            i for i, token_id in enumerate(input_ids) if token_id == eos_token_id
+        ]
+
+        assert len(eos_indices) > 0, "Expected at least one EOS token in the input"
+        for eos_idx in eos_indices:
+            assert labels[eos_idx] != IGNORE_TOKEN_ID, (
+                f"Expected EOS token at index {eos_idx} to be labeled"
+            )
+
+    def test_train_on_eos_turn(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing with train_on_eos='turn'")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+            train_on_eos="turn",
+        )
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        turns = strategy.get_conversation_thread(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        eos_token_id = tokenizer.eos_token_id
+        # Process all turns and verify EOS token labeling
+        for i, turn in enumerate(basic_dataset[0]["conversations"]):
+            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+
+            if self._should_skip_turn(tokenizer, turn, i, start_idx, end_idx):
+                continue
+
+            decoded_response = tokenizer.decode(input_ids[start_idx:end_idx])
+            response = turn["value"]
+
+            assert response in decoded_response, (
+                f"Response {response} not found in index {start_idx}:{end_idx} "
+                f"decoded:{decoded_response}"
+            )
+
+            # Find the EOS token after this turn
+            eos_idx = end_idx
+            while eos_idx < len(input_ids) and input_ids[eos_idx] != eos_token_id:
+                eos_idx += 1
+
+            assert eos_idx < len(input_ids), (
+                f"Could not find EOS token after '{response}'"
+            )
+
+            LOG.debug(
+                f"Turn {i}: role={turn['from']}, content='{turn['value']}', start_idx={start_idx}, end_idx={end_idx}, eos_idx={eos_idx}"
+            )
+
+            LOG.debug(
+                f"Labels for turn {i}: {labels[start_idx:end_idx]}, EOS label: {labels[eos_idx]}"
+            )
+
+            # Verify EOS token labeling based on role
+            is_assistant = turn["from"] == "assistant"
+            if is_assistant:
+                assert labels[eos_idx] != IGNORE_TOKEN_ID, (
+                    f"Expected EOS token after assistant response '{response}' to be labeled"
+                )
+            else:
+                assert labels[eos_idx] == IGNORE_TOKEN_ID, (
+                    f"Expected EOS token after non-assistant input '{response}' to not be labeled"
+                )
+
+    def test_train_on_eos_last(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing with train_on_eos='last'")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+            train_on_eos="last",
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        eos_token_id = tokenizer.eos_token_id
+        eos_indices = [
+            i for i, token_id in enumerate(input_ids) if token_id == eos_token_id
+        ]
+
+        assert len(eos_indices) > 0, "Expected at least one EOS token in the input"
+        last_eos_idx = eos_indices[-1]
+
+        # Check that only the last EOS token is labeled
+        for idx in eos_indices[:-1]:
+            assert labels[idx] == IGNORE_TOKEN_ID, (
+                f"Expected EOS token at index {idx} to not be labeled"
+            )
+        assert labels[last_eos_idx] != IGNORE_TOKEN_ID, (
+            f"Expected last EOS token at index {last_eos_idx} to be labeled"
+        )
+
+    def test_train_on_eos_none(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing with train_on_eos='none'")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+            train_on_eos="none",
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        eos_token_id = tokenizer.eos_token_id
+        eos_indices = [
+            i for i, token_id in enumerate(input_ids) if token_id == eos_token_id
+        ]
+
+        assert len(eos_indices) > 0, "Expected at least one EOS token in the input"
+        for eos_idx in eos_indices:
+            assert labels[eos_idx] == IGNORE_TOKEN_ID, (
+                f"Expected EOS token at index {eos_idx} to not be labeled"
+            )
+
+    def test_drop_system_message(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        LOG.info("Testing with drop_system_message=True")
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                drop_system_message=True,
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        input_ids = res["input_ids"]
+
+        # Check if system message is not present in input_ids
+        system_message = "You are an AI assistant."
+        decoded_message = tokenizer.decode(input_ids)
+        assert system_message not in decoded_message, (
+            "Expected system message to be dropped"
+        )
+
+    def test_custom_roles(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        request,
+    ):
+        LOG.info("Testing with custom roles mapping")
+        custom_roles = {
+            "user": ["human", "user"],
+            "assistant": ["ai", "assistant"],
+            "system": ["context"],
+        }
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                roles=custom_roles,
+                message_property_mappings={"role": "from", "content": "value"},
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["ai"],
+        )
+
+        # Create a new dataset with modified role names
+        modified_conversations = [
+            {"from": "context", "value": "You are an AI assistant."},
+            {"from": "human", "value": "Hello"},
+            {"from": "ai", "value": "Hi there!"},
+            {"from": "human", "value": "How are you?"},
+            {"from": "ai", "value": "I'm doing well, thank you!"},
+        ]
+
+        modified_dataset = Dataset.from_dict({"messages": [modified_conversations]})
+
+        res = strategy.tokenize_prompt(modified_dataset[0])
+        turns = strategy.get_conversation_thread(modified_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Process all turns and verify labeling
+        for i, turn in enumerate(modified_dataset[0]["messages"]):
+            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+
+            if self._should_skip_turn(tokenizer, turn, i, start_idx, end_idx):
+                continue
+
+            decoded_response = tokenizer.decode(input_ids[start_idx:end_idx])
+            response = turn["value"]
+
+            assert response in decoded_response, (
+                f"Response {response} not found in index {start_idx}:{end_idx} "
+                f"decoded:{decoded_response}"
+            )
+
+            # Check if responses are labeled correctly based on role
+            is_ai = turn["from"] == "ai"
+            if is_ai:
+                assert all(
+                    label != IGNORE_TOKEN_ID for label in labels[start_idx:end_idx]
+                ), f"Expected labels for AI response '{response}' to be set"
+            else:
+                assert all(
+                    label == IGNORE_TOKEN_ID for label in labels[start_idx:end_idx]
+                ), (
+                    f"Expected labels for non-AI message '{response}' to be IGNORE_TOKEN_ID"
+                )
+
+    def test_message_field_training(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        request,
+    ):
+        LOG.info("Testing with message_field_training")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_field_training="train",
+                message_field_training_detail="train_detail",
+                message_property_mappings={"role": "from", "content": "value"},
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=[],
+        )
+
+        # Create a new dataset with the train and train_detail fields
+        modified_conversation = [
+            {"from": "system", "value": "You are an AI assistant.", "train": False},
+            {"from": "human", "value": "Hello", "train": False},
+            {"from": "assistant", "value": "Hello", "train": True},
+            {"from": "human", "value": "How are you?", "train": True},
+            {
+                "from": "assistant",
+                "value": "I'm doing very well, thank you!",
+                "train_detail": [
+                    {"begin_offset": 0, "end_offset": 8, "train": False},
+                    {"begin_offset": 9, "end_offset": 18, "train": True},
+                    {"begin_offset": 19, "end_offset": 30, "train": False},
+                ],
+            },
+            {
+                "from": "human",
+                "value": "I'm doing very well, thank you!",
+                "train": False,
+            },
+            {"from": "assistant", "value": "Hi there!", "train": True},
+        ]
+
+        modified_dataset = Dataset.from_dict({"messages": [modified_conversation]})
+
+        res = strategy.tokenize_prompt(modified_dataset[0])
+        turns = strategy.get_conversation_thread(modified_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        def verify_labels(labels_span, should_train, context_message):
+            """Helper to verify if a span of labels matches expected training state"""
+            if should_train:
+                assert all(label != IGNORE_TOKEN_ID for label in labels_span), (
+                    f"Expected all labels for {context_message} to be set, but got {labels_span}"
+                )
+            else:
+                assert all(label == IGNORE_TOKEN_ID for label in labels_span), (
+                    f"Expected all labels for {context_message} to be {IGNORE_TOKEN_ID}, but got {labels_span}"
+                )
+
+        # Process all turns and verify labeling
+        for i, turn in enumerate(modified_dataset[0]["messages"]):
+            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+
+            if self._should_skip_turn(tokenizer, turn, i, start_idx, end_idx):
+                continue
+
+            decoded_response = tokenizer.decode(input_ids[start_idx:end_idx])
+            response = turn["value"]
+
+            assert response in decoded_response, (
+                f"Response {response} not found in index {start_idx}:{end_idx} "
+                f"decoded:{decoded_response}"
+            )
+
+            LOG.debug(
+                f"Processing turn {i}: role={turn['from']}, content='{turn['value']}', "
+                f"start_idx={start_idx}, end_idx={end_idx}"
+            )
+
+            if turn.get("train_detail", None) is not None:
+                # Handle detailed token-level training control
+                tokenized_output = tokenizer(
+                    turn["value"], return_offsets_mapping=True, add_special_tokens=False
+                )
+                assert tokenized_output["input_ids"] == input_ids[start_idx:end_idx], (
+                    f"Tokenized input mismatch for turn: {turn['value']}\n"
+                    f"Expected: {input_ids[start_idx:end_idx]}\nActual: {tokenized_output['input_ids']}\n"
+                    f"This will likely be a mismatch between template content and encoded content"
+                )
+
+                token_offsets = tokenized_output["offset_mapping"]
+
+                # Adjust token offsets
+                for j in range(len(token_offsets) - 1):
+                    token_offsets[j] = (
+                        token_offsets[j][0],
+                        token_offsets[j + 1][0] - 1,
+                    )
+                token_offsets[-1] = (token_offsets[-1][0], len(turn["value"]) - 1)
+
+                adjusted_train_details = strategy.prompter.adjust_train_details(
+                    turn["train_detail"], token_offsets
+                )
+
+                LOG.debug(f"Original train_details: {turn['train_detail']}")
+                LOG.debug(f"Adjusted train_details: {adjusted_train_details}")
+
+                # Get and verify token offsets
+                turn_tokens = input_ids[start_idx:end_idx]
+                token_offsets_unmasked = strategy.prompter.get_offsets_for_train_detail(
+                    text=turn["value"],
+                    train_details=adjusted_train_details,
+                    mask_untrainable=False,
+                )
+
+                for i, offset in enumerate(token_offsets_unmasked):
+                    assert token_offsets[i][0] == offset, (
+                        f"Token start offsets mismatch for turn: {turn['value']}\n"
+                        f"Expected: {token_offsets[i][0]}\nActual: {offset}"
+                    )
+
+                token_offsets_masked = strategy.prompter.get_offsets_for_train_detail(
+                    text=turn["value"],
+                    train_details=adjusted_train_details,
+                    mask_untrainable=True,
+                )
+                LOG.debug(f"Token offsets: {token_offsets_masked}")
+
+                # Verify expected labels against actual labels
+                expected_labels = [IGNORE_TOKEN_ID] * len(turn_tokens)
+                for i, offset in enumerate(token_offsets_masked):
+                    if offset != IGNORE_TOKEN_ID:
+                        expected_labels[i] = turn_tokens[i]
+                actual_labels = labels[
+                    start_idx : start_idx + len(token_offsets_masked)
+                ]
+                assert actual_labels == expected_labels, (
+                    f"Labels mismatch for turn: {turn['value']}\nExpected: {expected_labels}\nActual: {actual_labels}"
+                )
+
+                # Verify each detail section
+                for detail in adjusted_train_details:
+                    detail_start = start_idx + next(
+                        j
+                        for j, offset in enumerate(token_offsets_unmasked)
+                        if offset >= detail["begin_offset"]
+                    )
+                    detail_end = start_idx + next(
+                        (
+                            j
+                            for j, offset in enumerate(token_offsets_unmasked)
+                            if offset > detail["end_offset"]
+                        ),
+                        len(token_offsets),
+                    )
+
+                    detail_text = turn["value"][
+                        detail["begin_offset"] : detail["end_offset"] + 1
+                    ]
+                    detail_labels = labels[detail_start:detail_end]
+
+                    context = (
+                        f"detail (ind {detail_start}:{detail_end}): '{detail_text}'\n"
+                        f"decoded: '{tokenizer.decode(input_ids[detail_start:detail_end])}')"
+                    )
+                    verify_labels(detail_labels, detail["train"], context)
+            else:
+                # Handle regular turn-level training control
+                should_train = turn.get("train", False)
+                turn_labels = labels[start_idx:end_idx]
+                context = (
+                    f"turn (ind {start_idx}:{end_idx}): '{turn['value']}'\n"
+                    f"decoded: '{decoded_response}')"
+                )
+                verify_labels(turn_labels, should_train, context)
+
+        LOG.debug(f"Final labels: {labels}")
+        LOG.debug(f"Final input_ids: {input_ids}")
+
+    @enable_hf_offline
+    def test_content_parts_training(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        request,
+    ):
+        LOG.info("Testing with content as list of parts with per-part training")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "role", "content": "content"},
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        # Dataset where assistant content is a list of parts with per-part training
+        conversation = [
+            {
+                "role": "system",
+                "content": [
+                    {"type": "text", "text": "You are an AI assistant."},
+                ],
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is 2+2?"},
+                ],
+            },
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "Let me think...", "train": False},
+                    {"type": "text", "text": "The answer is 4.", "train": True},
+                ],
+            },
+        ]
+
+        dataset = Dataset.from_dict({"messages": [conversation]})
+        res = strategy.tokenize_prompt(dataset[0])
+        turns = strategy.get_conversation_thread(dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Find the assistant turn (last turn)
+        assistant_turn_idx = len(turns) - 1
+        start_idx, end_idx = strategy.find_turn(
+            turns=turns, turn_idx=assistant_turn_idx
+        )
+
+        assert start_idx != -1 and end_idx != -1, (
+            "Could not find assistant turn boundaries"
+        )
+
+        decoded = tokenizer.decode(input_ids[start_idx:end_idx])
+        LOG.debug(f"Assistant turn decoded: {decoded}")
+
+        # Tokenize each part separately to find their boundaries
+        part1_text = "Let me think..."
+        part2_text = "The answer is 4."
+
+        # Verify the concatenated content is in the decoded output
+        assert part1_text in decoded, (
+            f"Part 1 '{part1_text}' not found in decoded: {decoded}"
+        )
+        assert part2_text in decoded, (
+            f"Part 2 '{part2_text}' not found in decoded: {decoded}"
+        )
+
+        # Verify that part1 tokens (train=False) are masked
+        # and part2 tokens (train=True) are labeled
+        turn_labels = labels[start_idx:end_idx]
+
+        # Find where part2 starts in the token sequence
+        part1_tokens = tokenizer(part1_text, add_special_tokens=False)["input_ids"]
+        part2_tokens = tokenizer(part2_text, add_special_tokens=False)["input_ids"]
+
+        # The first part should be masked (all IGNORE_TOKEN_ID)
+        # Due to token boundary alignment, check that at least the interior tokens
+        # of part1 are masked
+        assert any(label == IGNORE_TOKEN_ID for label in turn_labels), (
+            f"Expected some masked labels for train=False part, but got {turn_labels}"
+        )
+
+        # The second part should be trained (not IGNORE_TOKEN_ID)
+        assert any(label != IGNORE_TOKEN_ID for label in turn_labels), (
+            f"Expected some trained labels for train=True part, but got {turn_labels}"
+        )
+
+        # More precise check: first N tokens should be masked, last M tokens should be trained
+        # where N ~ len(part1_tokens) and M ~ len(part2_tokens)
+        # Allow for token boundary effects at the boundary
+        num_masked = sum(1 for label in turn_labels if label == IGNORE_TOKEN_ID)
+        num_trained = sum(1 for label in turn_labels if label != IGNORE_TOKEN_ID)
+
+        LOG.debug(f"Turn labels: {turn_labels}")
+        LOG.debug(f"Masked tokens: {num_masked}, Trained tokens: {num_trained}")
+        LOG.debug(
+            f"Part1 tokens: {len(part1_tokens)}, Part2 tokens: {len(part2_tokens)}"
+        )
+
+        # The number of masked tokens should be roughly the size of part1
+        # and the number of trained tokens should be roughly the size of part2
+        assert num_masked > 0, "Expected masked tokens for the train=False part"
+        assert num_trained > 0, "Expected trained tokens for the train=True part"
+
+    @enable_hf_offline
+    def test_content_parts_with_weight(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        request,
+    ):
+        LOG.info("Testing with content parts using weight field")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "role", "content": "content"},
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        # Dataset using weight instead of train
+        conversation = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Hello"},
+                ],
+            },
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "Thinking step by step: ", "weight": 0},
+                    {"type": "text", "text": "Hello! How can I help?", "weight": 1},
+                ],
+            },
+        ]
+
+        dataset = Dataset.from_dict({"messages": [conversation]})
+        res = strategy.tokenize_prompt(dataset[0])
+        labels = res["labels"]
+
+        # There should be both masked and trained labels
+        has_masked = any(label == IGNORE_TOKEN_ID for label in labels)
+        has_trained = any(label != IGNORE_TOKEN_ID for label in labels)
+        assert has_masked, "Expected masked tokens (weight=0 part + user turn)"
+        assert has_trained, "Expected trained tokens (weight=1 part)"
+
+    @enable_hf_offline
+    def test_content_parts_string_passthrough(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        request,
+    ):
+        LOG.info("Testing that string content still works alongside list content")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "role", "content": "content"},
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        # All list content in the conversation
+        conversation = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is 2+2?"},
+                ],
+            },
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "The answer is 4.", "train": True},
+                ],
+            },
+        ]
+
+        dataset = Dataset.from_dict({"messages": [conversation]})
+        res = strategy.tokenize_prompt(dataset[0])
+
+        # Should tokenize without errors
+        assert "input_ids" in res
+        assert "labels" in res
+        assert len(res["input_ids"]) > 0
+
+    def test_get_chat_template_variables(
+        self, tokenizer, chat_template, chat_template_jinja, eos_token, request
+    ):
+        LOG.info("Testing get_chat_template_variables")
+
+        actual_tokenizer, actual_jinja_template = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        prompter = ChatTemplatePrompter(
+            actual_tokenizer,
+            chat_template=get_chat_template(
+                chat_template, jinja_template=actual_jinja_template
+            ),
+            message_property_mappings={"from": "role", "value": "content"},
+        )
+
+        variables = prompter.get_chat_template_msg_variables(
+            (
+                actual_jinja_template
+                if actual_jinja_template
+                else actual_tokenizer.get_chat_template()
+            ),
+            "messages",
+        )
+
+        # Special case for Mistral with additional tool variables
+        if chat_template == "jinja" and tokenizer == "mistralv03_tokenizer":
+            expected_variables = {"role", "content", "tool_call_id", "tool_calls"}
+        # Most chat templates use the standard role and content variables
+        elif chat_template in ["llama3", "chatml", "phi_35", "phi_4"] or (
+            chat_template == "jinja" and tokenizer == "gemma2_tokenizer"
+        ):
+            expected_variables = {"role", "content"}
+        else:
+            LOG.warning(
+                f"Unsupported chat template: {chat_template} with {chat_template_jinja}"
+            )
+            raise ValueError(
+                f"Unsupported chat template: {chat_template} with {chat_template_jinja}"
+            )
+
+        assert variables == expected_variables, (
+            f"Expected variables: {expected_variables} from {tokenizer}/{chat_template}\n"
+            f"Got: {variables}\n"
+            f"Chat template: {actual_jinja_template}"
+        )
+
+    def test_eot_tokens_conflict_with_eos_token(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        """Test that an error is raised when eot_tokens contains eos_token and train_on_eot/train_on_eos conflict"""
+        LOG.info(
+            "Testing conflict between eot_tokens containing eos_token and train_on_eot/train_on_eos mismatch"
+        )
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        # Create a situation where eot_tokens contains eos_token
+        eot_tokens = [
+            tokenizer.eos_token,
+            "[/INST]",
+        ]  # Deliberately including eos_token
+
+        # Create conflicting train_on_eos and train_on_eot settings
+        with pytest.raises(
+            ValueError,
+            match=".*eos_token is in eot_tokens and train_on_eos != train_on_eot.*",
+        ):
+            ChatTemplateStrategy(
+                ChatTemplatePrompter(
+                    tokenizer,
+                    chat_template=get_chat_template(
+                        chat_template, jinja_template=chat_template_jinja
+                    ),
+                    message_property_mappings={"role": "from", "content": "value"},
+                    field_messages="conversations",
+                ),
+                tokenizer=tokenizer,
+                train_on_inputs=False,
+                sequence_len=512,
+                roles_to_train=["assistant"],
+                train_on_eos="none",  # Setting to none
+                train_on_eot="turn",  # Different from train_on_eos
+                eot_tokens=eot_tokens,
+            )
+
+    def test_eot_token_backward_compatibility(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        """Test that eot_tokens inherits from eos_token when not specified"""
+        LOG.info("Testing backward compatibility that eot_token inherits eos_token")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+            train_on_eos="turn",  # Setting train_on_eos to "turn"
+        )
+
+        # In backward compatibility mode, eot_tokens should be derived from eos_token
+        assert strategy.eot_tokens == [tokenizer.eos_token], (
+            f"Expected eot_tokens to inherit from eos_token, got {strategy.eot_tokens}"
+        )
+        assert strategy.train_on_eot == "turn", (
+            f"Expected train_on_eot to inherit from train_on_eos, got {strategy.train_on_eot}"
+        )
+
+    def test_token_not_in_template(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        """Test runs even when tokens are not found in the template"""
+        LOG.info("Testing runs even when tokens are not found in template")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        # Create a non-existent token that definitely won't be in the template
+        non_existent_token = "[DEFINITELY_NOT_IN_TEMPLATE]"
+        tokenizer.add_special_tokens(
+            {"additional_special_tokens": [non_existent_token]}
+        )
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template(
+                    chat_template, jinja_template=chat_template_jinja
+                ),
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+            eot_tokens=[non_existent_token],
+        )
+
+        # Force template check by calling tokenize_prompt
+        strategy.tokenize_prompt(basic_dataset[0])
+
+        # We can also check that a warning was logged, but there's
+        # caplog conflicts when running with other tests
+        # assert any(
+        #     "not found in chat_template" in record.message for record in self._caplog.records
+        # ), "Expected warning about token not found in template was not logged"
+
+    def test_custom_eot_tokens(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        """Test with custom EOT tokens to ensure proper masking and training"""
+        LOG.info("Testing with custom EOT tokens")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, None, request
+        )
+
+        # Add custom EOT tokens to the tokenizer
+        custom_eot = "[EOT]"
+        tokenizer.add_special_tokens({"additional_special_tokens": [custom_eot]})
+
+        # Create a custom chat template that uses our EOT token
+        custom_template = """{% for message in messages %}{% if message['role'] == 'system' %}{{ message['content'] }}{% elif message['role'] == 'user' %}User: {{ message['content'] }}{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}[EOT]{% endif %}{% endfor %}"""
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=custom_template,
+                message_property_mappings={"role": "from", "content": "value"},
+                field_messages="conversations",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+            train_on_eot="turn",  # Train on EOT token after each turn
+            eot_tokens=[custom_eot],
+        )
+
+        res = strategy.tokenize_prompt(basic_dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Find indices of the EOT token
+        eot_token_id = tokenizer.convert_tokens_to_ids(custom_eot)
+        eot_indices = [
+            i for i, token_id in enumerate(input_ids) if token_id == eot_token_id
+        ]
+
+        assert len(eot_indices) > 0, "Expected at least one EOT token in the input"
+
+        # Verify labeling for EOT tokens based on role
+        turns = strategy.get_conversation_thread(basic_dataset[0])
+        assistant_turn_indices = []
+        non_assistant_turn_indices = []
+
+        for i, turn in enumerate(basic_dataset[0]["conversations"]):
+            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+            if start_idx != -1 and end_idx != -1:  # If turn is found
+                if turn["from"] == "assistant":
+                    assistant_turn_indices.append((start_idx, end_idx))
+                else:
+                    non_assistant_turn_indices.append((start_idx, end_idx))
+
+        # Check EOT tokens after assistant turns are labeled
+        for eot_idx in eot_indices:
+            is_after_assistant = any(
+                start_idx <= eot_idx <= end_idx + 1  # +1 to include the EOT token
+                for start_idx, end_idx in assistant_turn_indices
+            )
+
+            if is_after_assistant:
+                assert labels[eot_idx] != IGNORE_TOKEN_ID, (
+                    f"Expected EOT token after assistant turn at index {eot_idx} to be labeled"
+                )
+            else:
+                assert labels[eot_idx] == IGNORE_TOKEN_ID, (
+                    f"Expected EOT token not after assistant turn at index {eot_idx} to not be labeled"
+                )
+
+    def test_multiple_train_on_eot_settings(
+        self,
+        tokenizer,
+        chat_template,
+        chat_template_jinja,
+        eos_token,
+        basic_dataset,
+        request,
+    ):
+        """Test different train_on_eot settings"""
+        LOG.info("Testing different train_on_eot settings")
+
+        tokenizer, chat_template_jinja = self.setup_tokenizer(
+            tokenizer, chat_template, chat_template_jinja, eos_token, request
+        )
+
+        # Create a list to test different train_on_eot settings
+        test_settings = [
+            ("none", lambda idx, is_assistant: False),  # Never train on EOT
+            ("all", lambda idx, is_assistant: True),  # Always train on EOT
+            (
+                "turn",
+                lambda idx, is_assistant: is_assistant,
+            ),  # Train on EOT after assistant turns
+            ("last", lambda idx, is_last: is_last),  # Only train on last EOT
+        ]
+
+        for setting, expected_train_func in test_settings:
+            LOG.info(f"Testing train_on_eot='{setting}'")
+
+            strategy = ChatTemplateStrategy(
+                ChatTemplatePrompter(
+                    tokenizer,
+                    chat_template=get_chat_template(
+                        chat_template, jinja_template=chat_template_jinja
+                    ),
+                    message_property_mappings={"role": "from", "content": "value"},
+                    field_messages="conversations",
+                ),
+                tokenizer=tokenizer,
+                train_on_inputs=False,
+                sequence_len=512,
+                roles_to_train=["assistant"],
+                train_on_eot=setting,
+                eot_tokens=[
+                    tokenizer.eos_token
+                ],  # Use eos_token as the EOT token for simplicity
+            )
+
+            res = strategy.tokenize_prompt(basic_dataset[0])
+            turns = strategy.get_conversation_thread(basic_dataset[0])
+            labels = res["labels"]
+            input_ids = res["input_ids"]
+
+            eos_token_id = tokenizer.eos_token_id
+            eos_indices = [
+                i for i, token_id in enumerate(input_ids) if token_id == eos_token_id
+            ]
+
+            assert len(eos_indices) > 0, (
+                "Expected at least one EOS/EOT token in the input"
+            )
+
+            # Check labeling for each EOS/EOT token
+            for idx, eos_idx in enumerate(eos_indices):
+                # Find which turn this EOS token belongs to
+                preceding_turn = None
+                for i, turn in enumerate(basic_dataset[0]["conversations"]):
+                    start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
+                    if (
+                        start_idx != -1
+                        and end_idx != -1
+                        and start_idx <= eos_idx <= end_idx + 1
+                    ):
+                        preceding_turn = turn
+                        break
+
+                is_assistant = (
+                    preceding_turn is not None and preceding_turn["from"] == "assistant"
+                )
+                is_last = idx == len(eos_indices) - 1
+
+                expected_label = not expected_train_func(
+                    idx, is_assistant if setting != "last" else is_last
+                )
+
+                if expected_label:
+                    assert labels[eos_idx] == IGNORE_TOKEN_ID, (
+                        f"Expected EOT token at index {eos_idx} to not be labeled with train_on_eot='{setting}'"
+                    )
+                else:
+                    assert labels[eos_idx] != IGNORE_TOKEN_ID, (
+                        f"Expected EOT token at index {eos_idx} to be labeled with train_on_eot='{setting}'"
+                    )
+
+
+class TestChatTemplateToolCalling:
+    """
+    Test class for tool calling functionality with chat templates.
+    """
+
+    def test_tool_calling_with_llama4_template(
+        self,
+        llama3_tokenizer,
+    ):
+        LOG.info("Testing tool calling with llama3 tokenizer and llama4 chat template")
+
+        # Create tool calling dataset
+        tool_calling_dataset = [
+            {
+                "tools": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "xml_escape",
+                            "description": 'Replaces any "<", ">", or "&" characters in the input string with their corresponding XML entities.',
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "s": {
+                                        "type": "string",
+                                        "description": "The input string to be XML-escaped.",
+                                    }
+                                },
+                                "required": ["s"],
+                            },
+                        },
+                    },
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "multiples",
+                            "description": "Generates a list of all the multiples of a number that are less than a given limit.",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "number": {
+                                        "type": "integer",
+                                        "description": "The number to find multiples of.",
+                                    },
+                                    "limit": {
+                                        "type": "integer",
+                                        "description": "The upper limit for the multiples.",
+                                    },
+                                },
+                                "required": ["number", "limit"],
+                            },
+                        },
+                    },
+                ],
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "Can you help me find multiples of 5 that are less than 20?",
+                    },
+                    {
+                        "role": "assistant",
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "multiples",
+                                    "arguments": {
+                                        "number": 5,
+                                        "limit": 20,
+                                    },
+                                },
+                            }
+                        ],
+                    },
+                    {"role": "tool", "name": "multiples", "content": "5,10,15"},
+                    {
+                        "role": "assistant",
+                        "content": "The multiples of 5 less than 20 are: 5, 10, and 15.",
+                    },
+                ],
+            }
+        ]
+
+        # Setup tokenizer with llama4 chat template
+        tokenizer = deepcopy(llama3_tokenizer)
+
+        # Add EOS token to the tokenizer
+        eot_token = "<|eot_id|>"
+        tokenizer.add_special_tokens({"additional_special_tokens": [eot_token]})
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template("llama4"),
+                message_property_mappings={"role": "role", "content": "content"},
+                field_messages="messages",
+                field_tools="tools",
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+            eot_tokens=[eot_token],
+        )
+
+        res = strategy.tokenize_prompt(tool_calling_dataset[0])
+        input_ids = res["input_ids"]
+        labels = res["labels"]
+
+        # Verify that the input_ids contain expected tokens
+        assert len(input_ids) > 0, "Input IDs should not be empty"
+        assert len(labels) == len(input_ids), "Labels should match input_ids length"
+
+        # Decode the full conversation to verify structure
+        decoded_conversation = tokenizer.decode(input_ids)
+
+        # Verify tool calling structure is present in the decoded conversation
+        assert '"type": "function",' in decoded_conversation, (
+            "Tool type function should be in conversation"
+        )
+        assert '"name": "multiples",' in decoded_conversation, (
+            "Tool function name should be in conversation"
+        )
+
+        assert (
+            '<|python_start|><|python_end|>{"name": "multiples", "parameters": {"number": 5, "limit": 20}}<|eot|>'
+            in decoded_conversation
+        ), "Assistant tool call should be in conversation"
+        assert "<|header_start|>ipython<|header_end|>" in decoded_conversation, (
+            "IPython header should be in conversation"
+        )
+        assert '"5,10,15"' in decoded_conversation, (
+            "Tool response should be in conversation"
+        )
+
+        # Get conversation turns to verify labeling
+        turns = strategy.get_conversation_thread(tool_calling_dataset[0])
+        tools = strategy._get_tools(tool_calling_dataset[0])
+
+        # Check that assistant responses are properly labeled
+        for i, turn in enumerate(tool_calling_dataset[0]["messages"]):
+            if turn["role"] == "assistant":
+                start_idx, end_idx = strategy.find_turn(
+                    turns=turns, turn_idx=i, tools=tools
+                )
+
+                assert start_idx != -1 and end_idx != -1, (
+                    f"Assistant turn {i} should be found"
+                )
+
+                # Verify that assistant responses have proper labels
+                turn_labels = labels[start_idx:end_idx]
+                assert all(label != IGNORE_TOKEN_ID for label in turn_labels), (
+                    f"Assistant turn {i} should be unmasked"
+                )
+
+
+class TestChatTemplateReasoningContent:
+    """
+    Test class for reasoning_content with content parts.
+    """
+
+    @enable_hf_offline
+    def test_reasoning_content_with_content_parts(self, qwen3_tokenizer):
+        """Test that reasoning_content as string + content as list parts works correctly.
+        Content training_detail offsets should align with content-only boundaries."""
+        LOG.info("Testing reasoning_content with content parts on qwen3")
+
+        tokenizer = deepcopy(qwen3_tokenizer)
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template("qwen3"),
+                message_property_mappings={
+                    "role": "role",
+                    "content": "content",
+                    "reasoning_content": "reasoning_content",
+                },
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        # reasoning_content is a plain string, content is list with per-part training
+        conversation = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "What is 2+2?"}],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": "Step 1: 2+2=4",
+                "content": [
+                    {"type": "text", "text": "The answer is 4.", "train": True},
+                ],
+            },
+        ]
+
+        dataset = Dataset.from_dict({"messages": [conversation]})
+        res = strategy.tokenize_prompt(dataset[0])
+        turns = strategy.get_conversation_thread(dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Find the assistant turn
+        assistant_idx = 1
+        start_idx, end_idx = strategy.find_turn(
+            turns=turns, turn_idx=assistant_idx, content_only=True
+        )
+
+        assert start_idx != -1 and end_idx != -1, (
+            "Could not find assistant content boundaries"
+        )
+
+        # The content-only span should contain "The answer is 4." but NOT "Step 1: 2+2=4"
+        decoded_span = tokenizer.decode(input_ids[start_idx:end_idx])
+        assert "The answer is 4." in decoded_span, (
+            f"Content not found in span: {decoded_span}"
+        )
+        assert "Step 1" not in decoded_span, (
+            f"Reasoning should not be in content-only span: {decoded_span}"
+        )
+
+        # Verify that content tokens are trained
+        content_labels = labels[start_idx:end_idx]
+        assert any(label != IGNORE_TOKEN_ID for label in content_labels), (
+            f"Expected trained labels in content span, got {content_labels}"
+        )
+
+    @enable_hf_offline
+    def test_reasoning_content_per_part_masking(self, qwen3_tokenizer):
+        """Test masking incorrect reasoning while training on self-correction.
+        This is the core use case: mask out wrong thoughts, train on corrections."""
+        LOG.info("Testing reasoning_content per-part masking on qwen3")
+
+        tokenizer = deepcopy(qwen3_tokenizer)
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template("qwen3"),
+                message_property_mappings={
+                    "role": "role",
+                    "content": "content",
+                    "reasoning_content": "reasoning_content",
+                },
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        # Reasoning has wrong step (masked) then self-correction (trained)
+        conversation = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "What is 2+2?"}],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": [
+                    {"type": "text", "text": "Hmm maybe 2+2=5.", "train": False},
+                    {"type": "text", "text": " Wait no, 2+2=4.", "train": True},
+                ],
+                "content": [
+                    {"type": "text", "text": "The answer is 4.", "train": True},
+                ],
+            },
+        ]
+
+        dataset = Dataset.from_dict({"messages": [conversation]})
+        res = strategy.tokenize_prompt(dataset[0])
+        turns = strategy.get_conversation_thread(dataset[0])
+        labels = res["labels"]
+        input_ids = res["input_ids"]
+
+        # Find reasoning boundaries
+        reasoning_start, reasoning_end = strategy.find_turn(
+            turns=turns, turn_idx=1, reasoning_only=True
+        )
+        assert reasoning_start != -1 and reasoning_end != -1, (
+            "Could not find reasoning boundaries"
+        )
+
+        decoded_reasoning = tokenizer.decode(input_ids[reasoning_start:reasoning_end])
+        LOG.debug(f"Reasoning span: {decoded_reasoning!r}")
+        assert "2+2=5" in decoded_reasoning, (
+            f"Wrong step not in reasoning span: {decoded_reasoning}"
+        )
+        assert "2+2=4" in decoded_reasoning, (
+            f"Correction not in reasoning span: {decoded_reasoning}"
+        )
+
+        # Verify reasoning labels have both masked and trained tokens
+        reasoning_labels = labels[reasoning_start:reasoning_end]
+        reasoning_ids = input_ids[reasoning_start:reasoning_end]
+
+        # Decode only the trained tokens — should be exactly the self-correction
+        trained_ids = [
+            tid
+            for tid, lab in zip(reasoning_ids, reasoning_labels, strict=True)
+            if lab != IGNORE_TOKEN_ID
+        ]
+        trained_text = tokenizer.decode(trained_ids)
+        assert trained_text.strip() == "Wait no, 2+2=4.", (
+            f"Expected trained reasoning to be 'Wait no, 2+2=4.', got: {trained_text!r}"
+        )
+
+        # Decode only the masked tokens — should be exactly the incorrect step
+        masked_ids = [
+            tid
+            for tid, lab in zip(reasoning_ids, reasoning_labels, strict=True)
+            if lab == IGNORE_TOKEN_ID
+        ]
+        masked_text = tokenizer.decode(masked_ids)
+        assert masked_text.strip() == "Hmm maybe 2+2=5.", (
+            f"Expected masked reasoning to be 'Hmm maybe 2+2=5.', got: {masked_text!r}"
+        )
+
+        # Find content boundaries
+        content_start, content_end = strategy.find_turn(
+            turns=turns, turn_idx=1, content_only=True
+        )
+        assert content_start != -1 and content_end != -1, (
+            "Could not find content boundaries"
+        )
+
+        # Content should be fully trained — decode trained tokens to verify
+        content_labels = labels[content_start:content_end]
+        content_ids = input_ids[content_start:content_end]
+        content_trained_ids = [
+            tid
+            for tid, lab in zip(content_ids, content_labels, strict=True)
+            if lab != IGNORE_TOKEN_ID
+        ]
+        content_trained_text = tokenizer.decode(content_trained_ids)
+        assert "The answer is 4." in content_trained_text, (
+            f"Expected 'The answer is 4.' in trained content tokens, "
+            f"got: {content_trained_text!r}"
+        )
+        assert all(label != IGNORE_TOKEN_ID for label in content_labels), (
+            f"Expected all content labels trained, got {content_labels}"
+        )
+
+    @enable_hf_offline
+    def test_reasoning_content_as_list_no_training_flags(self, qwen3_tokenizer):
+        """Test that reasoning_content as list without training flags still works."""
+        LOG.info("Testing reasoning_content as list without training flags on qwen3")
+
+        tokenizer = deepcopy(qwen3_tokenizer)
+
+        strategy = ChatTemplateStrategy(
+            ChatTemplatePrompter(
+                tokenizer,
+                chat_template=get_chat_template("qwen3"),
+                message_property_mappings={
+                    "role": "role",
+                    "content": "content",
+                    "reasoning_content": "reasoning_content",
+                },
+            ),
+            tokenizer=tokenizer,
+            train_on_inputs=False,
+            sequence_len=512,
+            roles_to_train=["assistant"],
+        )
+
+        # Both as lists, no per-part training flags
+        conversation = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "What is 2+2?"}],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": [
+                    {"type": "text", "text": "Step 1: addition."},
+                    {"type": "text", "text": " Step 2: 2+2=4."},
+                ],
+                "content": [
+                    {"type": "text", "text": "The answer is 4."},
+                ],
+            },
+        ]
+
+        dataset = Dataset.from_dict({"messages": [conversation]})
+        res = strategy.tokenize_prompt(dataset[0])
+
+        # Should tokenize without errors
+        assert "input_ids" in res
+        assert "labels" in res
+        assert len(res["input_ids"]) > 0
+
+        # Verify the full output contains both reasoning and content
+        full_text = tokenizer.decode(res["input_ids"])
+        assert "Step 1: addition." in full_text
+        assert "Step 2: 2+2=4." in full_text
+        assert "The answer is 4." in full_text
diff --git a/tests/prompt_strategies/test_chat_templates_mistral.py b/tests/prompt_strategies/test_chat_templates_mistral.py
new file mode 100644
index 0000000000..956ce119aa
--- /dev/null
+++ b/tests/prompt_strategies/test_chat_templates_mistral.py
@@ -0,0 +1,879 @@
+"""Test chat templates for mistral-common wrapper tokenizer"""
+
+import unittest
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer
+
+    from axolotl.utils.mistral import HFMistralTokenizer
+
+
+# fmt: off
+@pytest.mark.parametrize(
+    ("tokenizer_str", "assistant_toolcall_ids", "tool_result_ids"),
+    (
+        ("magistral_tokenizer", (9, 44627, 3684, 33, 19881, 1049, 1050, 1051, 1052, 1053, 32, 19227, 12856, 2811, 1032, 1049, 1054, 1044, 1429, 33319, 2811, 1032, 1050, 1125, 2), (7, 19881, 1049, 1050, 1051, 1052, 1053, 19, 1049, 1044, 1050, 8)),
+        ("devstral_tokenizer", (9, 1091, 19227, 2391, 2811, 1429, 44627, 3684, 1897, 1429, 61906, 2811, 16753, 12856, 2811, 1032, 1049, 1054, 1044, 1429, 33319, 2811, 1032, 1050, 4179, 1429, 1327, 2811, 1429, 19881, 1049, 1050, 1051, 1052, 1053, 1034, 27028, 2), (7, 19881, 1049, 1050, 1051, 1052, 1053, 19, 1049, 1044, 1050, 8)),
+        ("devstral_1_1_tokenizer", (9, 44627, 3684, 32, 19227, 12856, 2811, 1032, 1049, 1054, 1044, 1429, 33319, 2811, 1032, 1050, 1125, 2,), (7, 1049, 1044, 1050, 8)),
+    )
+)
+# fmt: on
+def test_mistral_chat_template(
+    tokenizer_str: str,
+    assistant_toolcall_ids: tuple[int, ...],
+    tool_result_ids: tuple[int, ...],
+    request: pytest.FixtureRequest,
+):
+    """Test chat template with the Magistral/Devstral tokenizer"""
+
+    from axolotl.prompt_strategies.chat_template import MistralPrompter, MistralStrategy
+
+    tokenizer: HFMistralTokenizer = request.getfixturevalue(tokenizer_str)
+
+    # check bos, eos, pad, unk are accessible properties
+    assert tokenizer.bos_token_id == 1
+    assert tokenizer.eos_token_id == 2
+    assert tokenizer.pad_token_id == 11
+    assert tokenizer.unk_token_id == 0
+
+    assert tokenizer.pad_token == "<pad>"
+    assert tokenizer.eos_token == "</s>"
+    assert tokenizer.bos_token == "<s>"
+    assert tokenizer.unk_token == "<unk>"
+
+    strategy = MistralStrategy(
+        MistralPrompter(
+            tokenizer,
+            chat_template=None,
+            message_property_mappings={"role": "role", "content": "content"},
+        ),
+        tokenizer=tokenizer,
+        train_on_inputs=False,
+        train_on_eos="turn",
+        sequence_len=512,
+        roles_to_train=["assistant"],
+    )
+
+    # test chat template masking without system prompt
+    res = strategy.tokenize_prompt(
+        {
+            "messages": [
+                {"role": "user", "content": "Hello, how are you?"},
+                {"role": "assistant", "content": "I'm doing great, thank you!"},
+            ]
+        }
+    )
+
+    assert res["input_ids"] == [
+        1,  # bos
+        3,  # [INST]
+        22177,  # Hello
+        1044,  # ,
+        2606,  # how
+        1584,  # are
+        1636,  # you
+        1063,  # ?
+        4,  # [/INST]
+        1073,  # I
+        4525,  # 'm
+        6965,  # doing
+        4824,  # great
+        1044,  # ,
+        15412,  # thank
+        1636,  # you
+        1033,  # !
+        2,  # </s>
+    ]
+
+    assert res["labels"] == [
+        -100,  # bos
+        -100,  # [INST]
+        -100,  # Hello
+        -100,  # ,
+        -100,  # how
+        -100,  # are
+        -100,  # you
+        -100,  # ?
+        -100,  # [/INST]
+        1073,  # I
+        4525,  # 'm
+        6965,  # doing
+        4824,  # great
+        1044,  # ,
+        15412,  # thank
+        1636,  # you
+        1033,  # !
+        2,  # </s>
+    ]
+
+    # test chat template masking with system prompt
+    res = strategy.tokenize_prompt(
+        {
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello, how are you?"},
+                {"role": "assistant", "content": "I'm doing great, thank you!"},
+            ]
+        }
+    )
+
+    assert res["input_ids"] == [
+        1,  # bos
+        17,  # [SYSTEM_PROMPT]
+        4568,  # You
+        1584,  # are
+        1261,  # a
+        20351,  # helpful
+        27089,  # assistant
+        1046,  # .
+        18,  # [/SYSTEM_PROMPT]
+        3,  # [INST]
+        22177,  # Hello
+        1044,  # ,
+        2606,  # how
+        1584,  # are
+        1636,  # you
+        1063,  # ?
+        4,  # [/INST]
+        1073,  # I
+        4525,  # 'm
+        6965,  # doing
+        4824,  # great
+        1044,  # ,
+        15412,  # thank
+        1636,  # you
+        1033,  # !
+        2,  # </s>
+    ]
+
+    assert res["labels"] == [
+        -100,  # bos
+        -100,  # [SYSTEM_PROMPT]
+        -100,  # You
+        -100,  # are
+        -100,  # a
+        -100,  # helpful
+        -100,  # assistant
+        -100,  # .
+        -100,  # [/SYSTEM_PROMPT]
+        -100,  # [INST]
+        -100,  # Hello
+        -100,  # ,
+        -100,  # how
+        -100,  # are
+        -100,  # you
+        -100,  # ?
+        -100,  # [/INST]
+        1073,  # I
+        4525,  # 'm
+        6965,  # doing
+        4824,  # great
+        1044,  # ,
+        15412,  # thank
+        1636,  # you
+        1033,  # !
+        2,  # </s>
+    ]
+
+    # test chat template with tools
+    res = strategy.tokenize_prompt(
+        {
+            "tools": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "multiples",
+                        "description": "Generates a list of all the multiples of a number that are less than a given limit.",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "number": {
+                                    "type": "integer",
+                                    "description": "The number to find multiples of.",
+                                },
+                                "limit": {
+                                    "type": "integer",
+                                    "description": "The upper limit for the multiples.",
+                                },
+                            },
+                            "required": ["number", "limit"],
+                        },
+                    },
+                },
+            ],
+            "messages": [
+                {
+                    "role": "user",
+                    "content": "Hey, can you give me a breakdown of how to throw an awesome themed party? Like, what themes work best, and how can I set everything up to really wow my guests? I want some ideas on decorations, food, and activities that will make the party unforgettable!",
+                },
+                {
+                    "role": "assistant",
+                    "tool_calls": [
+                        {
+                            "id": "call12345",
+                            "type": "function",
+                            "function": {
+                                "name": "multiples",
+                                "arguments": {
+                                    "number": 16,
+                                    "limit": 2,
+                                },
+                            },
+                        }
+                    ],
+                },
+                {
+                    "role": "tool",
+                    "tool_call_id": "call12345",
+                    "name": "multiples",
+                    "content": "1,2",
+                },
+                {"role": "assistant", "content": "The multiples of 16 is 1 and 2."},
+            ],
+        }
+    )
+
+    # fmt: off
+    assert res["input_ids"] == [
+        1,  # bos
+        5, 1091, 19227, 4994, 2811, 1429, 5165, 1897, 1429, 5165, 2811, 16753, 2391, 2811, 1429, 44627, 3684, 1897, 1429, 14653, 2811, 1429, 10639, 2130, 1261, 2951, 1307, 1747, 1278, 60092, 1307, 1261, 2782, 1455, 1584, 4289, 2224, 1261, 4265, 6139, 39249, 1429, 26204, 2811, 16753, 4994, 2811, 1429, 6371, 1897, 1429, 48649, 2811, 16753, 12856, 2811, 16753, 4994, 2811, 1429, 49039, 1897, 1429, 14653, 2811, 1429, 1784, 2782, 1317, 3081, 60092, 1307, 2613, 4179, 1429, 33319, 2811, 16753, 4994, 2811, 1429, 49039, 1897, 1429, 14653, 2811, 1429, 1784, 9229, 6139, 1394, 1278, 60092, 2613, 47579, 1429, 15760, 2811, 12161, 12856, 1897, 1429, 33319, 4964, 2821, 27028, 6,  # tool prompt
+        3, 46634, 1044, 1710, 1636, 5628, 1639, 1261, 44433, 1307, 2606, 1317, 5388, 1420, 54191, 2424, 1286, 8967, 1063, 15621, 1044, 2549, 30305, 2196, 3560, 1044, 1321, 2606, 1710, 1362, 2016, 8605, 2015, 1317, 5524, 118931, 2036, 32951, 1063, 1362, 2933, 2269, 12106, 1408, 101987, 1044, 6939, 1044, 1321, 9216, 1455, 2084, 3180, 1278, 8967, 119141, 1689, 5935, 1033, 4,  # user
+        *assistant_toolcall_ids,  # assistant tool calling
+        *tool_result_ids,  # tool result
+        1784, 60092, 1307, 1032, 1049, 1054, 1395, 1032, 1049, 1321, 1032, 1050, 1046,  # assistant
+        2  # eos
+    ]
+
+    assert res["labels"] == [
+        -100,  # bos
+        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,  # tool prompt
+        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,  # user prompt
+        *assistant_toolcall_ids,  # assistant tool calling
+        *([-100] * len(tool_result_ids)),  # tool result
+        1784, 60092, 1307, 1032, 1049, 1054, 1395, 1032, 1049, 1321, 1032, 1050, 1046,  # assistant
+        2  # eos
+    ]
+    # fmt: on
+
+    # test chat template with tokenize=False
+    res = tokenizer.apply_chat_template(
+        [
+            {"role": "user", "content": "Hello, how are you?"},
+            {"role": "assistant", "content": "I'm doing great, thank you!"},
+        ],
+        tokenize=False,
+    )
+
+    assert res == "<s>[INST]Hello, how are you?[/INST]I'm doing great, thank you!</s>"
+
+    # test encode
+    res = tokenizer.encode("Hello, how are you?", add_special_tokens=True)
+    assert res == [
+        1,  # bos
+        22177,  # Hello
+        1044,  # ,
+        2606,  # how
+        1584,  # are
+        1636,  # you
+        1063,  # ?
+        2,  # eos
+    ]
+
+    # test decode no skip special tokens
+    decoded_res = tokenizer.decode(res, skip_special_tokens=False)
+
+    assert decoded_res == "<s>Hello, how are you?</s>"
+
+    # test decode skip special tokens
+    decoded_res = tokenizer.decode(res, skip_special_tokens=True)
+    assert decoded_res == "Hello, how are you?"
+
+    # test encode no special tokens
+    res = tokenizer.encode("Hello, how are you?", add_special_tokens=False)
+    assert res == [
+        22177,  # Hello
+        1044,  # ,
+        2606,  # how
+        1584,  # are
+        1636,  # you
+        1063,  # ?
+    ]
+
+    # test convert ids to tokens
+    res = tokenizer.convert_ids_to_tokens(res)
+    # spacing are needed as we are converting without decoding
+    assert res == ["Hello", ",", " how", " are", " you", "?"]
+
+
+def test_mistral_normalizer_fills_instruct_request_defaults(
+    magistral_tokenizer: "HFMistralTokenizer",
+):
+    from pydantic import BaseModel
+
+    class _MissingDefaults(BaseModel):
+        truncate_at_max_tokens: int
+        continue_final_message: bool
+
+    normalizer = magistral_tokenizer.tokenizer._instruct_request_normalizer
+
+    def broken_from_chat_completion_request(_request):
+        _MissingDefaults()
+
+    normalizer.from_chat_completion_request = broken_from_chat_completion_request
+    normalizer._axolotl_instruct_defaults_patched = False
+    magistral_tokenizer._patch_instruct_request_normalizer()
+
+    input_ids = magistral_tokenizer.apply_chat_template(
+        [{"role": "user", "content": "Hello"}],
+        add_generation_prompt=True,
+        return_dict=False,
+    )
+
+    assert isinstance(input_ids, list)
+    assert input_ids
+
+
+@pytest.mark.skip(reason="TODO, fix for new HF wrapper call")
+def test_magistral_tokenizer_pad_method(magistral_tokenizer: "HFMistralTokenizer"):
+    """Test the MistralTokenizer pad method"""
+    from axolotl.utils.collators.core import IGNORE_INDEX
+
+    magistral_pad_token_id = 11  # taken from tokenizer.pad_token_id
+
+    # Test padding with input_ids and labels only
+    features = [
+        {"input_ids": [1, 2, 3], "labels": [4, 5, 6]},
+        {"input_ids": [7, 8], "labels": [9, 10]},
+    ]
+
+    result = magistral_tokenizer.pad(features, padding=True, return_tensors="pt")
+
+    # Check that input_ids are padded correctly
+    assert result["input_ids"].shape == (2, 3)
+    assert result["input_ids"].tolist() == [[1, 2, 3], [7, 8, magistral_pad_token_id]]
+
+    # Check that labels are padded correctly
+    assert result["labels"].shape == (2, 3)
+    assert result["labels"].tolist() == [[4, 5, 6], [9, 10, IGNORE_INDEX]]
+
+    # Check that attention_mask and position_ids are NOT created
+    assert "attention_mask" not in result
+    assert "position_ids" not in result
+
+    # Test padding with attention_mask
+    features_with_attention = [
+        {"input_ids": [1, 2, 3], "labels": [4, 5, 6], "attention_mask": [1, 1, 1]},
+        {"input_ids": [7, 8], "labels": [9, 10], "attention_mask": [1, 1]},
+    ]
+
+    result = magistral_tokenizer.pad(
+        features_with_attention, padding=True, return_tensors="pt"
+    )
+
+    # Check that attention_mask is padded correctly
+    assert result["attention_mask"].shape == (2, 3)
+    assert result["attention_mask"].tolist() == [[1, 1, 1], [1, 1, 0]]
+
+    # Test padding with position_ids
+    features_with_position = [
+        {"input_ids": [1, 2, 3], "labels": [4, 5, 6], "position_ids": [0, 1, 2]},
+        {"input_ids": [7, 8], "labels": [9, 10], "position_ids": [0, 1]},
+    ]
+
+    result = magistral_tokenizer.pad(
+        features_with_position, padding=True, return_tensors="pt"
+    )
+
+    # Check that position_ids are padded correctly (continuing sequence)
+    assert result["position_ids"].shape == (2, 3)
+    assert result["position_ids"].tolist() == [[0, 1, 2], [0, 1, 2]]
+
+    # Test padding with all fields
+    features_all = [
+        {
+            "input_ids": [1, 2, 3],
+            "labels": [4, 5, 6],
+            "attention_mask": [1, 1, 1],
+            "position_ids": [0, 1, 2],
+        },
+        {
+            "input_ids": [7, 8],
+            "labels": [9, 10],
+            "attention_mask": [1, 1],
+            "position_ids": [0, 1],
+        },
+    ]
+
+    result = magistral_tokenizer.pad(features_all, padding=True, return_tensors="pt")
+
+    # All fields should be present and correctly padded
+    assert "input_ids" in result
+    assert "labels" in result
+    assert "attention_mask" in result
+    assert "position_ids" in result
+
+    # Test padding with all sequences same length
+    features_same_length = [
+        {"input_ids": [1, 2, 3], "labels": [4, 5, 6]},
+        {"input_ids": [7, 8, 9], "labels": [10, 11, 12]},
+    ]
+
+    result = magistral_tokenizer.pad(
+        features_same_length, padding=True, return_tensors="pt"
+    )
+
+    # Check match when no padding is needed
+    assert result["input_ids"][0].tolist() == features_same_length[0]["input_ids"]
+    assert result["labels"][0].tolist() == features_same_length[0]["labels"]
+
+    assert result["input_ids"][1].tolist() == features_same_length[1]["input_ids"]
+    assert result["labels"][1].tolist() == features_same_length[1]["labels"]
+
+    # Test padding with max_length parameter
+    result = magistral_tokenizer.pad(
+        features, padding="max_length", max_length=5, return_tensors="pt"
+    )
+
+    # Should pad to max_length
+    assert result["input_ids"].shape == (2, 5)
+    assert result["labels"].shape == (2, 5)
+
+    # Test numpy return type
+    result = magistral_tokenizer.pad(features, padding=True, return_tensors="np")
+
+    # Should return numpy arrays
+    import numpy as np
+
+    assert isinstance(result["input_ids"], np.ndarray)
+    assert isinstance(result["labels"], np.ndarray)
+
+    # Test unsupported field rejection
+    features_unsupported = [
+        {"input_ids": [1, 2, 3], "labels": [4, 5, 6], "unsupported_field": [7, 8, 9]},
+    ]
+
+    with pytest.raises(NotImplementedError, match="unsupported_field"):
+        magistral_tokenizer.pad(features_unsupported, padding=True, return_tensors="pt")
+
+    # Test token_type_ids rejection
+    features_token_type = [
+        {"input_ids": [1, 2, 3], "labels": [4, 5, 6], "token_type_ids": [0, 0, 0]},
+    ]
+
+    with pytest.raises(ValueError, match="token_type_ids is not supported"):
+        magistral_tokenizer.pad(features_token_type, padding=True, return_tensors="pt")
+
+
+def test_magistral_tool_calling(magistral_tokenizer: "HFMistralTokenizer"):
+    """Test tool calling with the Magistral tokenizer"""
+    from axolotl.prompt_strategies.chat_template import MistralPrompter, MistralStrategy
+
+    strategy = MistralStrategy(
+        MistralPrompter(
+            magistral_tokenizer,
+            chat_template=None,
+            message_property_mappings={"role": "role", "content": "content"},
+        ),
+        tokenizer=magistral_tokenizer,
+        train_on_inputs=False,
+        train_on_eos="turn",
+        sequence_len=512,
+        roles_to_train=["assistant"],
+    )
+
+    # Test basic tool calling with single function
+    basic_tool_calling = {
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "description": "Get the current weather for a location",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "location": {
+                                "type": "string",
+                                "description": "The city and state, e.g. San Francisco, CA",
+                            },
+                        },
+                        "required": ["location"],
+                    },
+                },
+            },
+        ],
+        "messages": [
+            {
+                "role": "user",
+                "content": "What's the weather like in San Francisco?",
+            },
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "call12345",
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {
+                                "location": "San Francisco, CA",
+                            },
+                        },
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "call12345",
+                "name": "get_weather",
+                "content": "Sunny, 72°F",
+            },
+            {
+                "role": "assistant",
+                "content": "The weather in San Francisco is sunny and 72°F.",
+            },
+        ],
+    }
+
+    res = strategy.tokenize_prompt(basic_tool_calling)
+
+    # Basic validation
+    assert "input_ids" in res
+    assert "labels" in res
+    assert len(res["input_ids"]) > 0
+    assert len(res["labels"]) == len(res["input_ids"])
+
+    # Decode and verify structure
+    decoded = magistral_tokenizer.decode(res["input_ids"])
+    assert (
+        '<s>[AVAILABLE_TOOLS][{"type": "function", "function": {"name": "get_weather", "description": "Get the current weather for a location", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"}}, "required": ["location"]}}}][/AVAILABLE_TOOLS]'
+        in decoded
+    )
+    assert (
+        '[TOOL_CALLS]get_weather[CALL_ID]call12345[ARGS]{"location": "San Francisco, CA"}</s>'
+        in decoded
+    )
+    assert "[TOOL_RESULTS]call12345[TOOL_CONTENT]Sunny, 72°F[/TOOL_RESULTS]" in decoded
+    assert "The weather in San Francisco is sunny and 72°F.</s>" in decoded
+
+    # Test multiple tool calls in sequence
+    multi_tool_calling = {
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "add_numbers",
+                    "description": "Add two numbers together",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "a": {"type": "number", "description": "First number"},
+                            "b": {"type": "number", "description": "Second number"},
+                        },
+                        "required": ["a", "b"],
+                    },
+                },
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "multiply_numbers",
+                    "description": "Multiply two numbers",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "x": {"type": "number", "description": "First number"},
+                            "y": {"type": "number", "description": "Second number"},
+                        },
+                        "required": ["x", "y"],
+                    },
+                },
+            },
+        ],
+        "messages": [
+            {
+                "role": "user",
+                "content": "Add 5 and 3, then multiply the result by 2",
+            },
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "call12345",
+                        "type": "function",
+                        "function": {
+                            "name": "add_numbers",
+                            "arguments": {"a": 5, "b": 3},
+                        },
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "call12345",
+                "name": "add_numbers",
+                "content": "8",
+            },
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "call23456",
+                        "type": "function",
+                        "function": {
+                            "name": "multiply_numbers",
+                            "arguments": {"x": 8, "y": 2},
+                        },
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "call23456",
+                "name": "multiply_numbers",
+                "content": "16",
+            },
+            {
+                "role": "assistant",
+                "content": "The result is 16. I first added 5 and 3 to get 8, then multiplied 8 by 2 to get 16.",
+            },
+        ],
+    }
+
+    res = strategy.tokenize_prompt(multi_tool_calling)
+
+    # Validation
+    assert len(res["input_ids"]) > 0
+    assert len(res["labels"]) == len(res["input_ids"])
+
+    decoded = magistral_tokenizer.decode(res["input_ids"])
+    assert (
+        '<s>[AVAILABLE_TOOLS][{"type": "function", "function": {"name": "add_numbers", "description": "Add two numbers together", "parameters": {"type": "object", "properties": {"a": {"type": "number", "description": "First number"}, "b": {"type": "number", "description": "Second number"}}, "required": ["a", "b"]}}}, {"type": "function", "function": {"name": "multiply_numbers", "description": "Multiply two numbers", "parameters": {"type": "object", "properties": {"x": {"type": "number", "description": "First number"}, "y": {"type": "number", "description": "Second number"}}, "required": ["x", "y"]}}}][/AVAILABLE_TOOLS]'
+        in decoded
+    )
+    assert (
+        '[TOOL_CALLS]add_numbers[CALL_ID]call12345[ARGS]{"a": 5, "b": 3}</s>' in decoded
+    )
+    assert "[TOOL_RESULTS]call12345[TOOL_CONTENT]8[/TOOL_RESULTS]" in decoded
+    assert (
+        '[TOOL_CALLS]multiply_numbers[CALL_ID]call23456[ARGS]{"x": 8, "y": 2}</s>'
+        in decoded
+    )
+    assert "[TOOL_RESULTS]call23456[TOOL_CONTENT]16[/TOOL_RESULTS]" in decoded
+    assert (
+        "The result is 16. I first added 5 and 3 to get 8, then multiplied 8 by 2 to get 16.</s>"
+        in decoded
+    )
+
+    # Test tool calling with system message
+    system_tool_calling = {
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search_database",
+                    "description": "Search for information in database",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "query": {"type": "string", "description": "Search query"},
+                        },
+                        "required": ["query"],
+                    },
+                },
+            },
+        ],
+        "messages": [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant with access to a database.",
+            },
+            {
+                "role": "user",
+                "content": "Find information about Python programming",
+            },
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "search123",
+                        "type": "function",
+                        "function": {
+                            "name": "search_database",
+                            "arguments": {"query": "Python programming"},
+                        },
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "search123",
+                "name": "search_database",
+                "content": "Python is a high-level programming language known for its simplicity.",
+            },
+            {
+                "role": "assistant",
+                "content": "Based on the database search, Python is a high-level programming language known for its simplicity and readability.",
+            },
+        ],
+    }
+
+    res = strategy.tokenize_prompt(system_tool_calling)
+
+    # Validation
+    assert len(res["input_ids"]) > 0
+    assert len(res["labels"]) == len(res["input_ids"])
+
+    decoded = magistral_tokenizer.decode(res["input_ids"])
+
+    assert (
+        '<s>[SYSTEM_PROMPT]You are a helpful assistant with access to a database.[/SYSTEM_PROMPT][AVAILABLE_TOOLS][{"type": "function", "function": {"name": "search_database", "description": "Search for information in database", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}}, "required": ["query"]}}}][/AVAILABLE_TOOLS]'
+        in decoded
+    )
+
+    # Test error handling - missing tool response
+    incomplete_tool_calling = {
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_time",
+                    "description": "Get current time",
+                    "parameters": {"type": "object", "properties": {}},
+                },
+            },
+        ],
+        "messages": [
+            {
+                "role": "user",
+                "content": "What time is it?",
+            },
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "time12345",
+                        "type": "function",
+                        "function": {
+                            "name": "get_time",
+                            "arguments": {},
+                        },
+                    }
+                ],
+            },
+            {
+                "role": "assistant",
+                "content": "The current time is 12:00 PM.",
+            },
+        ],
+    }
+
+    from mistral_common.exceptions import InvalidMessageStructureException
+
+    try:
+        strategy.tokenize_prompt(incomplete_tool_calling)
+    except InvalidMessageStructureException as e:
+        assert "Not the same number of function calls and responses" in str(e)
+
+
+@pytest.mark.skip(reason="TODO, fix for new HF wrapper call")
+def test_magistral_tokenizer_call_method(
+    magistral_tokenizer: "HFMistralTokenizer", llama3_tokenizer: "PreTrainedTokenizer"
+):
+    """Test the __call__ method behavior matches HuggingFace standards"""
+    from copy import deepcopy
+
+    import numpy as np
+    import torch
+
+    hf_tokenizer = deepcopy(llama3_tokenizer)
+    hf_tokenizer.pad_token = hf_tokenizer.eos_token
+
+    test_text = "Hello, how are you?"
+    batch_texts = ["Hello world", "How are you?"]
+
+    # Test single string with return_tensors=None
+    hf_result: dict[str, list[int]] = hf_tokenizer(test_text, return_tensors=None)
+    mistral_result: dict[str, list[int]] = magistral_tokenizer(
+        test_text, return_tensors=None
+    )
+
+    assert isinstance(mistral_result, dict)
+    assert set(mistral_result.keys()) == {"input_ids", "attention_mask"}
+    assert isinstance(mistral_result["input_ids"], type(hf_result["input_ids"]))  # list
+    assert isinstance(
+        mistral_result["attention_mask"], type(hf_result["attention_mask"])
+    )
+    assert len(mistral_result["input_ids"]) == len(mistral_result["attention_mask"])
+    assert np.all(mistral_result["attention_mask"])
+    assert len(np.array(mistral_result["input_ids"]).shape) == 1  # 1D array
+
+    # Test single string with return_tensors='pt'
+    hf_result_pt: dict[str, torch.Tensor] = hf_tokenizer(test_text, return_tensors="pt")
+    mistral_result_pt: dict[str, torch.Tensor] = magistral_tokenizer(
+        test_text, return_tensors="pt"
+    )
+
+    # Check structure and types
+    assert isinstance(mistral_result_pt["input_ids"], torch.Tensor)
+    assert isinstance(mistral_result_pt["attention_mask"], torch.Tensor)
+
+    # Check shapes match (don't compare token dimension)
+    assert len(hf_result_pt["input_ids"].shape) == len(
+        mistral_result_pt["input_ids"].shape
+    )
+    assert hf_result_pt["input_ids"].shape[0] == mistral_result_pt["input_ids"].shape[0]
+    assert (
+        mistral_result_pt["attention_mask"].shape
+        == mistral_result_pt["input_ids"].shape
+    )
+    assert torch.all(mistral_result_pt["attention_mask"] == 1)
+
+    # Test batch input with padding
+    hf_batch: dict[str, torch.Tensor] = hf_tokenizer(
+        batch_texts, return_tensors="pt", padding=True
+    )
+    mistral_batch: dict[str, torch.Tensor] = magistral_tokenizer(
+        batch_texts, return_tensors="pt", padding=True
+    )
+
+    # Check batch behavior
+    assert len(hf_batch["input_ids"].shape) == len(mistral_batch["input_ids"].shape)
+    assert hf_batch["input_ids"].shape[0] == mistral_batch["input_ids"].shape[0]
+    assert mistral_batch["attention_mask"].shape == mistral_batch["input_ids"].shape
+    assert torch.any(
+        mistral_batch["attention_mask"][0] == 0
+    )  # padding in shorter sequence
+    assert torch.all(
+        mistral_batch["attention_mask"][1] == 1
+    )  # no padding in longer sequence
+
+    # Test numpy tensors
+    mistral_result_np: dict[str, np.ndarray] = magistral_tokenizer(
+        test_text, return_tensors="np"
+    )
+    assert isinstance(mistral_result_np["input_ids"], np.ndarray)
+    assert isinstance(mistral_result_np["attention_mask"], np.ndarray)
+
+    # Test consistency with encode()
+    encoded: list[int] = magistral_tokenizer.encode(test_text, add_special_tokens=True)
+    called: dict[str, torch.Tensor] = magistral_tokenizer(
+        test_text, return_tensors="pt"
+    )
+    assert encoded == called["input_ids"][0].tolist()
+
+    # Test Error handling
+    with pytest.raises(ValueError, match="Unsupported kwargs"):
+        magistral_tokenizer(test_text, unsupported_param=True)
+
+    with pytest.raises(
+        ValueError, match="return_tensors='pt' or 'np' requires padding or truncation"
+    ):
+        magistral_tokenizer(batch_texts, return_tensors="pt")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/prompt_strategies/test_chat_templates_snapshot.py b/tests/prompt_strategies/test_chat_templates_snapshot.py
new file mode 100644
index 0000000000..eb9f5274d4
--- /dev/null
+++ b/tests/prompt_strategies/test_chat_templates_snapshot.py
@@ -0,0 +1,226 @@
+"""
+Golden-output snapshot tests for the bundled chat templates.
+
+Templates are whitespace-sensitive: a stray newline or space shifts token
+boundaries and corrupts label masking. These tests render each template through
+transformers' own jinja compiler (the exact path used at training time) and pin
+the output byte-for-byte, so any future edit that changes rendering is caught.
+
+To refresh a snapshot after an intentional template change, render the case and
+paste the new expected string (see ``EXPECTED`` below).
+"""
+
+import pytest
+from transformers.utils.chat_template_utils import _compile_jinja_template
+
+from axolotl.utils.chat_templates import _CHAT_TEMPLATES
+
+
+def render(template_name, messages, **kwargs):
+    compiled = _compile_jinja_template(_CHAT_TEMPLATES[template_name])
+    kwargs.setdefault("bos_token", "<s>")
+    kwargs.setdefault("eos_token", "</s>")
+    kwargs.setdefault("add_generation_prompt", False)
+    return compiled.render(messages=messages, **kwargs)
+
+
+# (case_id, template_name, messages, render_kwargs)
+CASES: list[tuple[str, str, list, dict]] = [
+    (
+        "deepseek_v3_basic",
+        "deepseek_v3",
+        [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello!"},
+        ],
+        {},
+    ),
+    (
+        "deepseek_v3_tool_flow",
+        "deepseek_v3",
+        [
+            {"role": "user", "content": "Weather in Paris?"},
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "Paris"}',
+                        },
+                    }
+                ],
+            },
+            {"role": "tool", "content": "sunny"},
+            {"role": "assistant", "content": "It's sunny in Paris."},
+        ],
+        {},
+    ),
+    # assistant closes a tool output with null content (regression: used to crash
+    # on ``str + None``)
+    (
+        "deepseek_v3_null_close",
+        "deepseek_v3",
+        [
+            {"role": "user", "content": "q"},
+            {"role": "tool", "content": "TOUT"},
+            {"role": "assistant", "content": None},
+        ],
+        {},
+    ),
+    (
+        "qwen3_basic",
+        "qwen3",
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello!"},
+        ],
+        {},
+    ),
+    # assistant tool call with null content (regression: used to crash iterating
+    # ``None`` in the non-string branch)
+    (
+        "qwen3_null_toolcall",
+        "qwen3",
+        [
+            {"role": "user", "content": "q"},
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [{"function": {"name": "f", "arguments": {"a": 1}}}],
+            },
+        ],
+        {},
+    ),
+    (
+        "qwen3_multimodal_user",
+        "qwen3",
+        [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Describe"}, {"type": "image"}],
+            },
+            {"role": "assistant", "content": "An image."},
+        ],
+        {},
+    ),
+    (
+        "mistral_v1_sys",
+        "mistral_v1",
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello"},
+            {"role": "user", "content": "Bye"},
+        ],
+        {},
+    ),
+    (
+        "mistral_v2v3_sys",
+        "mistral_v2v3",
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello"},
+            {"role": "user", "content": "Bye"},
+        ],
+        {},
+    ),
+    (
+        "mistral_v3_tekken_sys",
+        "mistral_v3_tekken",
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello"},
+            {"role": "user", "content": "Bye"},
+        ],
+        {},
+    ),
+    (
+        "exaone4_sys",
+        "exaone4",
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello"},
+        ],
+        {},
+    ),
+    (
+        "falcon_h1_basic",
+        "falcon_h1",
+        [
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello"},
+        ],
+        {},
+    ),
+    (
+        "pixtral_multimodal",
+        "pixtral",
+        [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is this?"},
+                    {"type": "image"},
+                ],
+            },
+            {"role": "assistant", "content": "A cat."},
+        ],
+        {},
+    ),
+    (
+        "llama3_multimodal",
+        "llama3",
+        [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Hi"}, {"type": "image"}],
+            },
+            {"role": "assistant", "content": "Hello"},
+        ],
+        {},
+    ),
+    (
+        "gemma3_sys_list",
+        "gemma3",
+        [
+            {"role": "system", "content": [{"type": "text", "text": "sys"}]},
+            {"role": "user", "content": [{"type": "text", "text": "Hi"}]},
+            {"role": "assistant", "content": [{"type": "text", "text": "Hello"}]},
+        ],
+        {},
+    ),
+]
+
+EXPECTED = {
+    "deepseek_v3_basic": "<s>You are helpful.<｜User｜>Hi<｜Assistant｜>Hello!<｜end▁of▁sentence｜>",
+    "deepseek_v3_tool_flow": '<s><｜User｜>Weather in Paris?<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n```json\n{"city": "Paris"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜><｜end▁of▁sentence｜><｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>sunny<｜tool▁output▁end｜><｜tool▁outputs▁end｜>It\'s sunny in Paris.<｜end▁of▁sentence｜>',
+    "deepseek_v3_null_close": "<s><｜User｜>q<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>TOUT<｜tool▁output▁end｜><｜tool▁outputs▁end｜><｜end▁of▁sentence｜>",
+    "qwen3_basic": "<|im_start|>system\nsys<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nHello!<|im_end|>\n",
+    "qwen3_null_toolcall": '<|im_start|>user\nq<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n<tool_call>\n{"name": "f", "arguments": {"a": 1}}\n</tool_call><|im_end|>\n',
+    "qwen3_multimodal_user": "<|im_start|>user\nDescribe<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nAn image.<|im_end|>\n",
+    "mistral_v1_sys": "<s> [INST] sys\n\nHi [/INST] Hello</s> [INST] Bye [/INST]",
+    "mistral_v2v3_sys": "<s>[INST] Hi[/INST] Hello</s>[INST] sys\n\nBye[/INST]",
+    "mistral_v3_tekken_sys": "<s>[INST]Hi[/INST]Hello</s>[INST]sys\n\nBye[/INST]",
+    "exaone4_sys": "[|system|]\nsys[|endofturn|]\n[|user|]\nHi[|endofturn|]\n[|assistant|]\n<think>\n\n</think>\n\nHello[|endofturn|]\n",
+    "falcon_h1_basic": "<s><|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\nHello<|im_end|>\n",
+    "pixtral_multimodal": "<s>[INST]What is this?[IMG][/INST]A cat.</s>",
+    "llama3_multimodal": "<s><|start_header_id|>user<|end_header_id|>\n\nHi<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nHello<|eot_id|>",
+    "gemma3_sys_list": "<s><start_of_turn>user\nsys\n\nHi<end_of_turn>\n<start_of_turn>model\nHello<end_of_turn>\n",
+}
+
+
+@pytest.mark.parametrize(
+    "case_id, template_name, messages, render_kwargs",
+    CASES,
+    ids=[c[0] for c in CASES],
+)
+def test_chat_template_snapshot(case_id, template_name, messages, render_kwargs):
+    assert render(template_name, messages, **render_kwargs) == EXPECTED[case_id]
diff --git a/tests/prompt_strategies/test_chat_templates_thinking.py b/tests/prompt_strategies/test_chat_templates_thinking.py
new file mode 100644
index 0000000000..054012e00b
--- /dev/null
+++ b/tests/prompt_strategies/test_chat_templates_thinking.py
@@ -0,0 +1,124 @@
+"""
+Tests for splitting reasoning/thinking from content into separate field
+"""
+
+import pytest
+from datasets import Dataset
+
+from axolotl.prompt_strategies.chat_template import (
+    load,
+)
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="messages_w_reasoning")
+def messages_w_reasoning_fixture():
+    return Dataset.from_list(
+        [
+            {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "hello",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "<think>lorem</think>\nwelcome",
+                    },
+                ]
+            },
+            {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "hello",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "<|begin_of_thought|>lorem<|end_of_thought|>\n<|begin_of_solution|>welcome\n<|end_of_solution|>",
+                    },
+                ]
+            },
+            {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "hello",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "<reasoning>lorem</reasoning>\nwelcome",
+                    },
+                ]
+            },
+        ]
+    )
+
+
+class TestSplitThinking:
+    """
+    test class to make sure datasets with reasoning content conforms to the chat_template strategy
+    """
+
+    def test_splits_think(self, messages_w_reasoning, qwen3_tokenizer):
+        strategy = load(
+            qwen3_tokenizer,
+            DictDefault(
+                {
+                    "train_on_inputs": False,
+                    "sequence_len": 512,
+                }
+            ),
+            DictDefault(
+                {
+                    "chat_template": "qwen3",
+                    "message_field_role": "role",
+                    "message_field_content": "content",
+                    "message_property_mappings": {
+                        "role": "role",
+                        "content": "content",
+                    },
+                    "roles": {
+                        "user": ["user"],
+                        "assistant": ["assistant"],
+                        "system": ["system"],
+                    },
+                    "field_messages": "messages",
+                    "split_thinking": True,
+                }
+            ),
+        )
+        for conversation in messages_w_reasoning:
+            transformed_prompt = strategy.get_conversation_thread(conversation)
+            assert transformed_prompt[0]["role"] == "user"
+            assert transformed_prompt[1]["role"] == "assistant"
+            assert transformed_prompt[1]["reasoning_content"] == "lorem"
+            assert transformed_prompt[1]["content"] == "welcome"
+
+            res = strategy.tokenize_prompt(conversation)
+            input_ids = res["input_ids"]
+            # fmt: off
+            expected_input_ids = [
+                151644,  # im_start
+                872,  # user
+                198,  # \n
+                14990,  # hello
+                151645,  # im_end
+                198,  # \n
+                151644,  # im_start
+                77091,  # assistant
+                198,  # \n
+                151667,  # think
+                198,  # \n
+                385, 1826,  # lorem
+                198,  # \n
+                151668,  # /think
+                271,  # \n
+                34084,  # welcome
+                151645,  # im_end
+                198,  # \n
+            ]
+            # fmt: on
+            assert input_ids == expected_input_ids, (
+                f"Input IDs mismatch: {input_ids} != {expected_input_ids}"
+            )
diff --git a/tests/prompt_strategies/test_chat_templates_tool_call_string_arguments.py b/tests/prompt_strategies/test_chat_templates_tool_call_string_arguments.py
new file mode 100644
index 0000000000..5866cc367d
--- /dev/null
+++ b/tests/prompt_strategies/test_chat_templates_tool_call_string_arguments.py
@@ -0,0 +1,505 @@
+"""
+Tests for handling json tool content
+"""
+
+import json
+
+import pytest
+from datasets import Dataset
+
+from axolotl.prompt_strategies.chat_template import (
+    load,
+)
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="qwen3_instruct_prompt_strategy")
+def qwen3_instruct_chat_template_strategy(qwen3_tokenizer):
+    strategy = load(
+        qwen3_tokenizer,
+        DictDefault(
+            {
+                "train_on_inputs": False,
+                "sequence_len": 512,
+            }
+        ),
+        DictDefault(
+            {
+                "chat_template": "qwen3",
+                "message_field_role": "role",
+                "message_field_content": "content",
+                "message_property_mappings": {
+                    "role": "role",
+                    "content": "content",
+                },
+                "roles": {
+                    "user": ["user"],
+                    "assistant": ["assistant"],
+                    "system": ["system"],
+                },
+                "field_messages": "messages",
+            }
+        ),
+    )
+    return strategy
+
+
+class TestQwen3IdenticalConversationArgs:
+    """
+    Test Qwen3 tools is identical between JSON and dict
+    """
+
+    @pytest.fixture(name="conversation_dict_args_dataset")
+    def fixture_conversation_dict_args_dataset(self):
+        """
+        Provides a dataset with conversation where arguments is a dict.
+        """
+        user_content = "What is the weather in Boston?"
+        function_name = "get_current_weather"
+        arguments_dict = {"location": "Boston, MA", "unit": "celsius"}
+
+        data = [
+            {
+                "messages": [
+                    {"role": "user", "content": user_content},
+                    {
+                        "role": "assistant",
+                        "content": "",
+                        "tool_calls": [
+                            {
+                                "function": {
+                                    "name": function_name,
+                                    "arguments": arguments_dict,  # dict
+                                }
+                            }
+                        ],
+                    },
+                ],
+            }
+        ]
+        return Dataset.from_list(data)
+
+    @pytest.fixture(name="conversation_str_args_dataset")
+    def fixture_conversation_str_args_dataset(self):
+        """
+        Provides a dataset with conversation where arguments is a JSON string.
+        """
+        user_content = "What is the weather in Boston?"
+        function_name = "get_current_weather"
+        arguments_dict = {"location": "Boston, MA", "unit": "celsius"}
+        arguments_str = json.dumps(arguments_dict)
+
+        data = [
+            {
+                "messages": [
+                    {"role": "user", "content": user_content},
+                    {
+                        "role": "assistant",
+                        "content": "",
+                        "tool_calls": [
+                            {
+                                "function": {
+                                    "name": function_name,
+                                    "arguments": arguments_str,  # str
+                                }
+                            }
+                        ],
+                    },
+                ],
+            }
+        ]
+        return Dataset.from_list(data)
+
+    @pytest.fixture(name="conversation_mixed_time_types_dataset")
+    def fixture_conversation_mixed_time_types_dataset(self):
+        """
+        Provides a dataset where 'time' field has different types in different tool calls.
+        """
+        data = [
+            {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "Get weather information at different times",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "",
+                        "tool_calls": [
+                            {
+                                "function": {
+                                    "name": "func1",
+                                    "arguments": json.dumps(
+                                        {"time": "2025-08-01"}
+                                    ),  # string type
+                                }
+                            },
+                            {
+                                "function": {
+                                    "name": "func2",
+                                    "arguments": json.dumps(
+                                        {"time": 1690876800}
+                                    ),  # number type
+                                }
+                            },
+                        ],
+                    },
+                ],
+            }
+        ]
+        return Dataset.from_list(data)
+
+    def test_dict_and_str_args_produce_identical_output(
+        self,
+        conversation_dict_args_dataset,
+        conversation_str_args_dataset,
+        qwen3_instruct_prompt_strategy,
+        qwen3_tokenizer,
+    ):
+        """
+        Tests that after tokenization and decoding, the outputs for both
+        dict and string `arguments` are exactly the same.
+        """
+        processed_dict_args = conversation_dict_args_dataset.map(
+            qwen3_instruct_prompt_strategy.tokenize_prompt,
+            batched=True,
+            remove_columns=["messages"],
+        )
+
+        processed_str_args = conversation_str_args_dataset.map(
+            qwen3_instruct_prompt_strategy.tokenize_prompt,
+            batched=True,
+            remove_columns=["messages"],
+        )
+
+        decoded_prompt_from_dict = qwen3_tokenizer.decode(
+            processed_dict_args[0]["input_ids"]
+        )
+
+        decoded_prompt_from_str = qwen3_tokenizer.decode(
+            processed_str_args[0]["input_ids"]
+        )
+
+        assert decoded_prompt_from_dict == decoded_prompt_from_str, (
+            f"Dict format output:\n{decoded_prompt_from_dict}\n"
+            f"String format output:\n{decoded_prompt_from_str}"
+        )
+
+        assert (
+            processed_dict_args[0]["input_ids"] == processed_str_args[0]["input_ids"]
+        ), "The tokenized input_ids should be identical for dict and str arguments"
+
+    def test_str_args_with_mixed_time_types_no_error(
+        self,
+        conversation_mixed_time_types_dataset,
+        qwen3_instruct_prompt_strategy,
+        qwen3_tokenizer,
+    ):
+        """
+        Tests that when 'time' field has different types (string vs number)
+        in different tool calls, str format arguments don't cause errors.
+        """
+        processed = conversation_mixed_time_types_dataset.map(
+            qwen3_instruct_prompt_strategy.tokenize_prompt,
+            batched=True,
+            remove_columns=["messages"],
+        )
+
+        assert len(processed) == 1
+        assert "input_ids" in processed[0]
+        assert len(processed[0]["input_ids"]) > 0
+
+        decoded = qwen3_tokenizer.decode(processed[0]["input_ids"])
+        assert "2025-08-01" in decoded, "String time value should be present"
+        assert "1690876800" in decoded, "Number time value should be present"
+
+
+class TestQwen3IdenticalToolsParameters:
+    """
+    Test Qwen3 tools parameters handling is identical between JSON string and dict
+    """
+
+    @pytest.fixture(name="tools_dict_params_dataset")
+    def fixture_tools_dict_params_dataset(self):
+        """
+        Provides a dataset with tools where parameters is a dict.
+        """
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "description": "Get weather information",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "location": {
+                                "type": "string",
+                                "description": "The city and state",
+                            },
+                            "unit": {
+                                "type": "string",
+                                "enum": ["celsius", "fahrenheit"],
+                            },
+                        },
+                        "required": ["location"],
+                    },
+                },
+            }
+        ]
+
+        data = [
+            {
+                "tools": tools,
+                "messages": [
+                    {"role": "user", "content": "What's the weather?"},
+                    {
+                        "role": "assistant",
+                        "content": "",
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "get_weather",
+                                    "arguments": {"location": "Boston, MA"},
+                                },
+                            }
+                        ],
+                    },
+                    {
+                        "role": "tool",
+                        "name": "get_weather",
+                        "content": "72°F and sunny",
+                    },
+                ],
+            }
+        ]
+        return Dataset.from_list(data)
+
+    @pytest.fixture(name="tools_str_params_dataset")
+    def fixture_tools_str_params_dataset(self):
+        """
+        Provides a dataset with tools where parameters is a JSON string.
+        """
+        parameters_dict = {
+            "type": "object",
+            "properties": {
+                "location": {"type": "string", "description": "The city and state"},
+                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+            },
+            "required": ["location"],
+        }
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "description": "Get weather information",
+                    "parameters": json.dumps(parameters_dict),
+                },
+            }
+        ]
+
+        data = [
+            {
+                "tools": tools,
+                "messages": [
+                    {"role": "user", "content": "What's the weather?"},
+                    {
+                        "role": "assistant",
+                        "content": "",
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "get_weather",
+                                    "arguments": {"location": "Boston, MA"},
+                                },
+                            }
+                        ],
+                    },
+                    {
+                        "role": "tool",
+                        "name": "get_weather",
+                        "content": "72°F and sunny",
+                    },
+                ],
+            }
+        ]
+        return Dataset.from_list(data)
+
+    @pytest.fixture(name="tools_mixed_type_params_dataset")
+    def fixture_tools_mixed_type_params_dataset(self):
+        """
+        Provides a dataset where different tools have the same parameter name with different types.
+        This tests that JSON string format prevents casting issues.
+        """
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "tool_with_string_arg",
+                    "description": "Tool expecting string argument",
+                    "parameters": json.dumps(
+                        {
+                            "type": "object",
+                            "properties": {
+                                "arg1": {
+                                    "type": "string",
+                                    "description": "A string parameter",
+                                }
+                            },
+                            "required": ["arg1"],
+                        }
+                    ),
+                },
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "tool_with_number_arg",
+                    "description": "Tool expecting number argument",
+                    "parameters": json.dumps(
+                        {
+                            "type": "object",
+                            "properties": {
+                                "arg1": {
+                                    "type": "number",
+                                    "description": "A numeric parameter",
+                                }
+                            },
+                            "required": ["arg1"],
+                        }
+                    ),
+                },
+            },
+        ]
+
+        data = [
+            {
+                "tools": tools,
+                "messages": [
+                    {"role": "user", "content": "Use both tools"},
+                    {
+                        "role": "assistant",
+                        "content": "",
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "tool_with_string_arg",
+                                    "arguments": json.dumps({"arg1": "hello"}),
+                                },
+                            },
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "tool_with_number_arg",
+                                    "arguments": json.dumps({"arg1": 42}),
+                                },
+                            },
+                        ],
+                    },
+                ],
+            }
+        ]
+        return Dataset.from_list(data)
+
+    def test_dict_and_str_params_produce_equivalent_output(
+        self,
+        tools_dict_params_dataset,
+        tools_str_params_dataset,
+        qwen3_instruct_prompt_strategy,
+        qwen3_tokenizer,
+    ):
+        """
+        Tests that after tokenization and decoding, the outputs for both
+        dict and string `parameters` in tools are semantically equivalent.
+        """
+        import re
+
+        processed_dict_params = tools_dict_params_dataset.map(
+            qwen3_instruct_prompt_strategy.tokenize_prompt,
+            batched=True,
+            remove_columns=["messages", "tools"],
+        )
+
+        processed_str_params = tools_str_params_dataset.map(
+            qwen3_instruct_prompt_strategy.tokenize_prompt,
+            batched=True,
+            remove_columns=["messages", "tools"],
+        )
+
+        decoded_dict = qwen3_tokenizer.decode(processed_dict_params[0]["input_ids"])
+        decoded_str = qwen3_tokenizer.decode(processed_str_params[0]["input_ids"])
+
+        # Extract the tool JSON from both outputs
+        tools_pattern = r"<tools>\n(.*?)\n</tools>"
+
+        dict_tools_match = re.search(tools_pattern, decoded_dict, re.DOTALL)
+        str_tools_match = re.search(tools_pattern, decoded_str, re.DOTALL)
+
+        assert dict_tools_match and str_tools_match, (
+            "Could not find tools section in output"
+        )
+
+        # Parse the JSON and compare as objects (order-independent)
+        dict_tools_json = json.loads(dict_tools_match.group(1))
+        str_tools_json = json.loads(str_tools_match.group(1))
+
+        # Deep comparison of the tool definitions
+        assert dict_tools_json == str_tools_json, (
+            f"Tool definitions are not equivalent:\n"
+            f"Dict format: {json.dumps(dict_tools_json, indent=2)}\n"
+            f"String format: {json.dumps(str_tools_json, indent=2)}"
+        )
+
+        # Verify the rest of the structure is the same (excluding the tools JSON part)
+        # The tools JSON can have different order, so we remove it here.
+        dict_normalized = re.sub(
+            r"<tools>.*?</tools>",
+            "<tools>TOOLS_PLACEHOLDER</tools>",
+            decoded_dict,
+            flags=re.DOTALL,
+        )
+        str_normalized = re.sub(
+            r"<tools>.*?</tools>",
+            "<tools>TOOLS_PLACEHOLDER</tools>",
+            decoded_str,
+            flags=re.DOTALL,
+        )
+
+        assert dict_normalized == str_normalized, (
+            "The overall structure differs between dict and string parameter formats"
+        )
+
+    def test_str_params_with_mixed_types_no_error(
+        self,
+        tools_mixed_type_params_dataset,
+        qwen3_instruct_prompt_strategy,
+        qwen3_tokenizer,
+    ):
+        """
+        Tests that when different tools have the same parameter name with different types,
+        JSON string format for parameters doesn't cause casting errors.
+        """
+        processed = tools_mixed_type_params_dataset.map(
+            qwen3_instruct_prompt_strategy.tokenize_prompt,
+            batched=True,
+            remove_columns=["messages", "tools"],
+        )
+
+        assert len(processed) == 1
+        assert "input_ids" in processed[0]
+        assert len(processed[0]["input_ids"]) > 0
+
+        decoded = qwen3_tokenizer.decode(processed[0]["input_ids"])
+
+        # Check that both tools are present
+        assert "tool_with_string_arg" in decoded
+        assert "tool_with_number_arg" in decoded
+
+        # Check that both argument values are present
+        assert "hello" in decoded
+        assert "42" in decoded
diff --git a/tests/prompt_strategies/test_dpo_chat_templates.py b/tests/prompt_strategies/test_dpo_chat_templates.py
new file mode 100644
index 0000000000..28028ce42e
--- /dev/null
+++ b/tests/prompt_strategies/test_dpo_chat_templates.py
@@ -0,0 +1,381 @@
+"""
+tests for chat_template prompt strategy
+"""
+
+import unittest
+
+import pytest
+from datasets import Dataset
+from transformers import AutoTokenizer
+
+from axolotl.prompt_strategies.dpo.chat_template import argilla_chat, default
+from axolotl.utils.dict import DictDefault
+
+from tests.hf_offline_utils import enable_hf_offline
+
+
+@pytest.fixture(name="assistant_dataset")
+def fixture_assistant_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "hello",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "hello",
+                    },
+                    {
+                        "role": "user",
+                        "content": "goodbye",
+                    },
+                ],
+                "chosen": {
+                    "role": "assistant",
+                    "content": "goodbye",
+                },
+                "rejected": {
+                    "role": "assistant",
+                    "content": "party on",
+                },
+            }
+        ]
+    )
+
+
+@pytest.fixture(name="custom_assistant_dataset")
+def fixture_custom_assistant_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "conversation": [
+                    {
+                        "speaker": "human",
+                        "text": "hello",
+                    },
+                    {
+                        "speaker": "agent",
+                        "text": "hello",
+                    },
+                    {
+                        "speaker": "human",
+                        "text": "goodbye",
+                    },
+                ],
+                "better": {
+                    "speaker": "agent",
+                    "text": "goodbye",
+                },
+                "worse": {
+                    "speaker": "agent",
+                    "text": "party on",
+                },
+            }
+        ]
+    )
+
+
+@pytest.fixture(name="argilla_chat_dataset")
+def fixture_argilla_chat_dataset():
+    return Dataset.from_list(
+        [
+            {
+                "chosen": [
+                    {
+                        "role": "user",
+                        "content": "hello",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "goodbye",
+                    },
+                ],
+                "rejected": [
+                    {
+                        "role": "user",
+                        "content": "hello",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "party on",
+                    },
+                ],
+            }
+        ]
+    )
+
+
+@pytest.fixture(name="phi3_tokenizer")
+@enable_hf_offline
+def fixture_phi3_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
+
+    return tokenizer
+
+
+@pytest.fixture(name="gemma_tokenizer")
+@enable_hf_offline
+def fixture_gemma_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("unsloth/gemma-2b-it", revision="703fb4a")
+
+    return tokenizer
+
+
+class TestAssistantDPOChatTemplateLlama3:
+    """
+    Test class for assistant style datasets with llama-3 prompts using the chat_template strategy.
+    """
+
+    def test_llama3_defaults(self, llama3_tokenizer, assistant_dataset):
+        transform_fn, _ = default(
+            DictDefault(
+                {
+                    "chat_template": "llama3",
+                    "datasets": [
+                        {
+                            "type": "chat_template",
+                        }
+                    ],
+                }
+            )
+        )
+        result = transform_fn(assistant_dataset[0], tokenizer=llama3_tokenizer)
+        assert result["prompt"] == (
+            "<|begin_of_text|>"
+            + "<|start_header_id|>user<|end_header_id|>\n\nhello<|eot_id|>"
+            + "<|start_header_id|>assistant<|end_header_id|>\n\nhello<|eot_id|>"
+            + "<|start_header_id|>user<|end_header_id|>\n\ngoodbye<|eot_id|>"
+            + "<|start_header_id|>assistant<|end_header_id|>\n\n"
+        )
+        assert result["chosen"] == "goodbye<|eot_id|>"
+        assert result["rejected"] == "party on<|eot_id|>"
+
+    def test_llama3_configured(self, llama3_tokenizer, custom_assistant_dataset):
+        transform_fn, _ = default(
+            DictDefault(
+                {
+                    "chat_template": "llama3",
+                    "datasets": [
+                        {
+                            "type": "chat_template",
+                            "field_messages": "conversation",
+                            "field_chosen": "better",
+                            "field_rejected": "worse",
+                            "message_field_role": "speaker",
+                            "message_field_content": "text",
+                            "roles": {
+                                "user": ["human"],
+                                "assistant": ["agent"],
+                                "system": ["sys"],
+                            },
+                        }
+                    ],
+                }
+            )
+        )
+        result = transform_fn(custom_assistant_dataset[0], tokenizer=llama3_tokenizer)
+        assert result["prompt"] == (
+            "<|begin_of_text|>"
+            + "<|start_header_id|>user<|end_header_id|>\n\nhello<|eot_id|>"
+            + "<|start_header_id|>assistant<|end_header_id|>\n\nhello<|eot_id|>"
+            + "<|start_header_id|>user<|end_header_id|>\n\ngoodbye<|eot_id|>"
+            + "<|start_header_id|>assistant<|end_header_id|>\n\n"
+        )
+        assert result["chosen"] == "goodbye<|eot_id|>"
+        assert result["rejected"] == "party on<|eot_id|>"
+
+
+class TestAssistantDPOChatTemplatePhi3:
+    """
+    Test class for assistant style datasets with phi-3 prompts using the tokenizer's chat_template strategy.
+    """
+
+    @pytest.mark.xfail(reason="likely upstream issue from v5.4.0")
+    def test_phi3_defaults(self, phi3_tokenizer, assistant_dataset):
+        transform_fn, _ = default(
+            DictDefault(
+                {
+                    "chat_template": "tokenizer_default",
+                    "datasets": [
+                        {
+                            "type": "chat_template",
+                        }
+                    ],
+                }
+            )
+        )
+        result = transform_fn(assistant_dataset[0], tokenizer=phi3_tokenizer)
+        assert result["prompt"] == (
+            "<|user|>\nhello<|end|>\n"
+            + "<|assistant|>\nhello<|end|>\n"
+            + "<|user|>\ngoodbye<|end|>\n"
+            + "<|assistant|>\n"
+        )
+        assert result["chosen"] == "goodbye<|end|>\n<|endoftext|>"
+        assert result["rejected"] == "party on<|end|>\n<|endoftext|>"
+
+
+class TestAssistantDPOChatTemplateGemma:
+    """
+    Test class for assistant style datasets with gemma prompts using the tokenizer's chat_template strategy.
+    """
+
+    def test_gemma_defaults(self, gemma_tokenizer, assistant_dataset):
+        transform_fn, _ = default(
+            DictDefault(
+                {
+                    "chat_template": "tokenizer_default",
+                    "datasets": [
+                        {
+                            "type": "chat_template",
+                        }
+                    ],
+                }
+            )
+        )
+        result = transform_fn(assistant_dataset[0], tokenizer=gemma_tokenizer)
+        assert result["prompt"] == (
+            "<bos><start_of_turn>user\nhello<end_of_turn>\n"
+            + "<start_of_turn>model\nhello<end_of_turn>\n"
+            + "<start_of_turn>user\ngoodbye<end_of_turn>\n"
+            + "<start_of_turn>model\n"
+        )
+        assert result["chosen"] == "goodbye<end_of_turn>"
+        assert result["rejected"] == "party on<end_of_turn>"
+
+
+class TestArgillaChatDPOChatTemplate:
+    """
+    Test class for argilla_chat style datasets (chosen/rejected contain full conversations).
+    """
+
+    def test_llama3_argilla_chat(self, llama3_tokenizer, argilla_chat_dataset):
+        transform_fn, _ = argilla_chat(
+            DictDefault(
+                {
+                    "chat_template": "llama3",
+                    "datasets": [
+                        {
+                            "type": "chat_template.argilla_chat",
+                        }
+                    ],
+                }
+            )
+        )
+        result = transform_fn(argilla_chat_dataset[0], tokenizer=llama3_tokenizer)
+        assert result["prompt"] == (
+            "<|begin_of_text|>"
+            + "<|start_header_id|>user<|end_header_id|>\n\nhello<|eot_id|>"
+            + "<|start_header_id|>assistant<|end_header_id|>\n\n"
+        )
+        assert result["chosen"] == "goodbye<|eot_id|>"
+        assert result["rejected"] == "party on<|eot_id|>"
+
+    @pytest.mark.xfail(reason="likely upstream issue from v5.4.0")
+    def test_phi3_argilla_chat(self, phi3_tokenizer, argilla_chat_dataset):
+        transform_fn, _ = argilla_chat(
+            DictDefault(
+                {
+                    "chat_template": "tokenizer_default",
+                    "datasets": [
+                        {
+                            "type": "chat_template.argilla_chat",
+                        }
+                    ],
+                }
+            )
+        )
+        result = transform_fn(argilla_chat_dataset[0], tokenizer=phi3_tokenizer)
+        assert result["prompt"] == "<|user|>\nhello<|end|>\n" + "<|assistant|>\n"
+        assert result["chosen"] == "goodbye<|end|>\n<|endoftext|>"
+        assert result["rejected"] == "party on<|end|>\n<|endoftext|>"
+
+
+class TestDPOChatTemplateToolRole:
+    """
+    Test that DPO chat template strategy handles tool role messages without KeyError.
+    Regression test for https://github.com/axolotl-ai-cloud/axolotl/issues/3217
+    """
+
+    def test_tool_role_default_no_key_error(self, llama3_tokenizer):
+        """Messages list with a 'tool' role should not raise KeyError."""
+        dataset = Dataset.from_list(
+            [
+                {
+                    "messages": [
+                        {"role": "user", "content": "What is the weather?"},
+                        {
+                            "role": "assistant",
+                            "content": "Let me check.",
+                        },
+                        {
+                            "role": "tool",
+                            "content": "22°C, sunny.",
+                        },
+                    ],
+                    "chosen": {
+                        "role": "assistant",
+                        "content": "It is 22°C and sunny.",
+                    },
+                    "rejected": {
+                        "role": "assistant",
+                        "content": "I don't know.",
+                    },
+                }
+            ]
+        )
+        transform_fn, _ = default(
+            DictDefault(
+                {
+                    "chat_template": "llama3",
+                    "datasets": [{"type": "chat_template"}],
+                }
+            )
+        )
+        # Should not raise KeyError: 'tool'
+        result = transform_fn(dataset[0], tokenizer=llama3_tokenizer)
+        assert "prompt" in result
+        assert "chosen" in result
+        assert "rejected" in result
+
+    def test_tool_role_custom_mapping_preserved(self, llama3_tokenizer):
+        """A user-supplied roles mapping that overrides 'tool' is still respected."""
+        dataset = Dataset.from_list(
+            [
+                {
+                    "messages": [
+                        {"role": "user", "content": "hello"},
+                        {"role": "tool_result", "content": "42"},
+                    ],
+                    "chosen": {"role": "assistant", "content": "The answer is 42."},
+                    "rejected": {"role": "assistant", "content": "Unknown."},
+                }
+            ]
+        )
+        transform_fn, _ = default(
+            DictDefault(
+                {
+                    "chat_template": "llama3",
+                    "datasets": [
+                        {
+                            "type": "chat_template",
+                            "roles": {
+                                "user": ["user"],
+                                "assistant": ["assistant"],
+                                "system": ["system"],
+                                "tool": ["tool_result"],
+                            },
+                        }
+                    ],
+                }
+            )
+        )
+        result = transform_fn(dataset[0], tokenizer=llama3_tokenizer)
+        assert "prompt" in result
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/prompt_strategies/test_dpo_chatml.py b/tests/prompt_strategies/test_dpo_chatml.py
new file mode 100644
index 0000000000..2c089067f7
--- /dev/null
+++ b/tests/prompt_strategies/test_dpo_chatml.py
@@ -0,0 +1,68 @@
+"""
+Tests for loading DPO preference datasets with chatml formatting
+"""
+
+import unittest
+
+import pytest
+
+from axolotl.loaders.tokenizer import load_tokenizer
+from axolotl.prompt_strategies.dpo import load as load_dpo
+from axolotl.utils.data.rl import prepare_preference_datasets
+from axolotl.utils.dict import DictDefault
+
+from tests.hf_offline_utils import enable_hf_offline
+
+
+@pytest.fixture(name="minimal_dpo_cfg")
+def fixture_cfg():
+    return DictDefault(
+        {
+            "base_model": "HuggingFaceTB/SmolLM2-135M",
+            "tokenizer_config": "HuggingFaceTB/SmolLM2-135M",
+            "rl": "dpo",
+            "learning_rate": 0.000001,
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "special_tokens": {
+                "pad_token": "<|endoftext|>",
+            },
+            "sequence_len": 2048,
+        }
+    )
+
+
+class TestDPOChatml:
+    """
+    Test loading DPO preference datasets with chatml formatting
+    """
+
+    @pytest.mark.skip(reason="TODO: fix hf hub offline to work with HF rate limits")
+    @enable_hf_offline
+    def test_default(self, minimal_dpo_cfg):
+        cfg = DictDefault(
+            {
+                "datasets": [
+                    {
+                        "path": "argilla/distilabel-intel-orca-dpo-pairs",
+                        "type": "chatml",
+                        "split": "train[:1%]",
+                    }
+                ]
+            }
+            | minimal_dpo_cfg
+        )
+
+        # test that dpo.load works
+        load_dpo("chatml", cfg)
+        # now actually load the datasets with the strategy
+        tokenizer = load_tokenizer(cfg)
+        train_ds, _ = prepare_preference_datasets(cfg, tokenizer)
+        assert train_ds[0]["prompt"].startswith("<|im_start|>")
+        assert train_ds[0]["prompt"].endswith("<|im_start|>assistant\n")
+        assert "chosen" in train_ds[0]
+        assert "rejected" in train_ds[0]
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/prompt_strategies/test_dpo_user_defined.py b/tests/prompt_strategies/test_dpo_user_defined.py
new file mode 100644
index 0000000000..70a94db171
--- /dev/null
+++ b/tests/prompt_strategies/test_dpo_user_defined.py
@@ -0,0 +1,142 @@
+"""
+Tests for user-defined DPO dataset transform strategies
+"""
+
+import pytest
+
+from axolotl.prompt_strategies.dpo import load as load_dpo
+from axolotl.prompt_strategies.dpo.user_defined import default
+from axolotl.utils.dict import DictDefault
+
+
+def make_cfg(type_cfg: dict) -> DictDefault:
+    """Build a minimal config with a single user-defined DPO dataset."""
+    return DictDefault({"datasets": [{"type": type_cfg}]})
+
+
+class TestDPOUserDefined:
+    """
+    Test user_defined.default DPO transforms
+    """
+
+    def test_explicit_config(self):
+        """Transform works with explicit prompt/chosen/rejected formats."""
+        cfg = make_cfg(
+            {
+                "field_prompt": "prompt",
+                "field_system": "system",
+                "field_chosen": "chosen",
+                "field_rejected": "rejected",
+                "prompt_format": "{prompt}",
+                "chosen_format": "{chosen}",
+                "rejected_format": "{rejected}",
+            }
+        )
+        transform_fn = default(cfg)
+        sample = transform_fn({"prompt": "hello", "chosen": "good", "rejected": "bad"})
+        assert sample["prompt"] == "hello"
+        assert sample["chosen"] == "good"
+        assert sample["rejected"] == "bad"
+
+    def test_defaults_without_formats(self):
+        """Falls back to canonical {prompt}/{chosen}/{rejected} when no formats given."""
+        cfg = make_cfg({})
+        transform_fn = default(cfg)
+        sample = transform_fn({"prompt": "hello", "chosen": "good", "rejected": "bad"})
+        assert sample["prompt"] == "hello"
+        assert sample["chosen"] == "good"
+        assert sample["rejected"] == "bad"
+
+    def test_custom_field_names(self):
+        """Custom field_* columns are read and written to the canonical keys."""
+        cfg = make_cfg(
+            {
+                "field_prompt": "question",
+                "field_chosen": "my_chosen",
+                "field_rejected": "my_rejected",
+            }
+        )
+        transform_fn = default(cfg)
+        sample = transform_fn(
+            {"question": "hello", "my_chosen": "good", "my_rejected": "bad"}
+        )
+        assert sample["prompt"] == "hello"
+        assert sample["chosen"] == "good"
+        assert sample["rejected"] == "bad"
+
+    def test_system_in_prompt_format(self):
+        """{system} is filled from the system field when present."""
+        cfg = make_cfg({"prompt_format": "{system} {prompt}"})
+        transform_fn = default(cfg)
+        sample = transform_fn(
+            {
+                "system": "be helpful",
+                "prompt": "hello",
+                "chosen": "good",
+                "rejected": "bad",
+            }
+        )
+        assert sample["prompt"] == "be helpful hello"
+
+    def test_system_in_prompt_format_custom_field(self):
+        """{system} reads from a custom field_system column."""
+        cfg = make_cfg({"field_system": "sys", "prompt_format": "{system} {prompt}"})
+        transform_fn = default(cfg)
+        sample = transform_fn(
+            {
+                "sys": "be helpful",
+                "prompt": "hello",
+                "chosen": "good",
+                "rejected": "bad",
+            }
+        )
+        assert sample["prompt"] == "be helpful hello"
+
+    @pytest.mark.parametrize(
+        "sample",
+        [
+            {"prompt": "hello", "chosen": "good", "rejected": "bad"},
+            {"system": None, "prompt": "hello", "chosen": "good", "rejected": "bad"},
+        ],
+        ids=["missing", "none"],
+    )
+    def test_system_in_prompt_format_empty_system(self, sample):
+        """A missing or None system field renders as "" instead of raising or 'None'."""
+        cfg = make_cfg({"prompt_format": "{system} {prompt}"})
+        transform_fn = default(cfg)
+        assert transform_fn(sample)["prompt"] == " hello"
+
+    def test_custom_chosen_rejected_formats(self):
+        """Explicit chosen_format/rejected_format are honored with custom field names."""
+        cfg = make_cfg(
+            {
+                "field_chosen": "good",
+                "field_rejected": "bad",
+                "chosen_format": "Answer: {chosen}",
+                "rejected_format": "Answer: {rejected}",
+            }
+        )
+        transform_fn = default(cfg)
+        sample = transform_fn({"prompt": "q", "good": "yes", "bad": "no"})
+        assert sample["prompt"] == "q"
+        assert sample["chosen"] == "Answer: yes"
+        assert sample["rejected"] == "Answer: no"
+
+    def test_non_dict_type_raises(self):
+        """A non-dict dataset type raises a clear ValueError."""
+        cfg = make_cfg({})
+        cfg["datasets"][0]["type"] = "user_defined.default"
+        with pytest.raises(ValueError, match="must be a dictionary"):
+            default(cfg)
+
+    def test_load_dpo_user_defined_returns_callable(self):
+        """The loader resolves user_defined.default to a callable transform."""
+        cfg = make_cfg(
+            {
+                "field_prompt": "question",
+                "field_chosen": "my_chosen",
+                "field_rejected": "my_rejected",
+            }
+        )
+        transform_fn = load_dpo("user_defined.default", cfg, dataset_idx=0)
+        assert callable(transform_fn)
diff --git a/tests/prompt_strategies/test_get_messages_str_format.py b/tests/prompt_strategies/test_get_messages_str_format.py
new file mode 100644
index 0000000000..8f1c8632aa
--- /dev/null
+++ b/tests/prompt_strategies/test_get_messages_str_format.py
@@ -0,0 +1,45 @@
+"""
+Tests for ChatTemplateStrategy._get_messages with str-encoded messages.
+
+When a JSON string decodes to a non-list, the assertion must report the decoded type
+instead of raising UnboundLocalError on the not-yet-bound loop variable.
+"""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+
+from axolotl.prompt_strategies.chat_template import ChatTemplateStrategy
+
+
+@pytest.fixture
+def strategy():
+    # _get_messages only reads self.prompter.field_messages; skip __init__ so we
+    # test the real method without a tokenizer/model.
+    strat = ChatTemplateStrategy.__new__(ChatTemplateStrategy)
+    strat.prompter = SimpleNamespace(field_messages="messages")
+    return strat
+
+
+class TestGetMessagesStrFormat:
+    def test_valid_str_messages_returned(self, strategy):
+        turns = [{"role": "user", "content": "hello"}]
+        assert strategy._get_messages({"messages": json.dumps(turns)}) == turns
+
+    def test_non_list_json_reports_decoded_type(self, strategy):
+        # JSON object instead of a list
+        with pytest.raises(AssertionError, match=r"got <class 'dict'>"):
+            strategy._get_messages({"messages": '{"role": "user", "content": "hi"}'})
+
+    def test_non_dict_turn_reports_turn_type(self, strategy):
+        with pytest.raises(AssertionError, match=r"got <class 'str'> for the turn 0"):
+            strategy._get_messages({"messages": json.dumps(["not_a_dict"])})
+
+    def test_invalid_json_raises_json_decode_error(self, strategy):
+        with pytest.raises(json.JSONDecodeError):
+            strategy._get_messages({"messages": "not-valid-json"})
+
+    def test_none_messages_raises_value_error(self, strategy):
+        with pytest.raises(ValueError, match="null"):
+            strategy._get_messages({})
diff --git a/tests/prompt_strategies/test_jinja_template_analyzer.py b/tests/prompt_strategies/test_jinja_template_analyzer.py
new file mode 100644
index 0000000000..41b9a0203a
--- /dev/null
+++ b/tests/prompt_strategies/test_jinja_template_analyzer.py
@@ -0,0 +1,158 @@
+"""
+tests for jinja_template_analyzer
+"""
+
+import pytest
+
+from axolotl.prompt_strategies.jinja_template_analyzer import JinjaTemplateAnalyzer
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__, log_level="DEBUG")
+
+
+class TestJinjaTemplateAnalyzer:
+    """
+    tests for jinja_template_analyzer
+    """
+
+    def test_basic_variable_extraction(self, basic_jinja_template_analyzer):
+        """Test that all top-level variables are correctly extracted."""
+        LOG.info("Testing with train_on_inputs=True")
+
+        variables = basic_jinja_template_analyzer.get_template_variables()
+        expected_vars = {"messages", "add_generation_prompt", "eos_token", "message"}
+        assert set(variables.keys()) == expected_vars
+
+    def test_mixtral_variable_extraction(self, mistral_jinja_template_analyzer):
+        """Test that all top-level variables are correctly extracted."""
+        LOG.info("Testing with train_on_inputs=True")
+
+        variables = mistral_jinja_template_analyzer.get_template_variables()
+        expected_vars = {
+            "messages",
+            "content",
+            "eos_token",
+            "message",
+            "tools",
+            "system_message",
+            "loop_messages",
+            "ns",
+            "tool_call",
+            "tool",
+            "loop",
+            "bos_token",
+            "raise_exception",
+        }
+        assert set(variables.keys()) == expected_vars
+        message_vars = variables["message"]
+        assert message_vars == {"role", "content", "tool_calls", "tool_call_id"}
+
+    def test_message_property_access(self, basic_jinja_template_analyzer):
+        """Test that properties accessed on 'message' variable are correctly identified."""
+        LOG.info("Testing message property access")
+
+        variables = basic_jinja_template_analyzer.get_template_variables()
+        assert "messages" in variables
+        assert "message" in variables
+        assert "role" in variables["message"]
+        assert "content" in variables["message"]
+
+    def test_detailed_analysis(self, basic_jinja_template_analyzer):
+        """Test the detailed analysis of variable usage."""
+        LOG.info("Testing detailed analysis")
+
+        analysis = basic_jinja_template_analyzer.analyze_template()
+
+        assert analysis["messages"]["is_iterated"] is True
+        assert "role" in analysis["message"]["accessed_properties"]
+        assert "content" in analysis["message"]["accessed_properties"]
+
+        assert analysis["add_generation_prompt"]["is_conditional"] is True
+        assert len(analysis["add_generation_prompt"]["accessed_properties"]) == 0
+
+        assert not analysis["eos_token"]["is_iterated"]
+        assert len(analysis["eos_token"]["accessed_properties"]) == 0
+
+    def test_nested_property_access(self):
+        """Test handling of nested property access."""
+        LOG.info("Testing nested property access")
+
+        template = """{{ user.profile.name }}{{ user.settings['preference'] }}"""
+        analyzer = JinjaTemplateAnalyzer(template)
+        variables = analyzer.get_template_variables()
+
+        assert "user" in variables
+        assert "profile" in variables["user"]
+        assert "settings" in variables["user"]
+
+    def test_loop_variable_handling(self):
+        """Test handling of loop variables and their properties."""
+        LOG.info("Testing loop variable handling")
+
+        template = """
+        {% for item in items %}
+            {{ item.name }}
+            {% for subitem in item.subitems %}
+                {{ subitem.value }}
+            {% endfor %}
+        {% endfor %}
+        """
+        analyzer = JinjaTemplateAnalyzer(template)
+        analysis = analyzer.analyze_template()
+
+        assert analysis["items"]["is_iterated"]
+        assert "name" in analysis["item"]["accessed_properties"]
+        assert "subitems" in analysis["item"]["accessed_properties"]
+
+    def test_conditional_variable_usage(self):
+        """Test detection of variables used in conditional statements."""
+        LOG.info("Testing conditional variable usage")
+
+        template = """
+        {% if user.is_admin and config.debug_mode %}
+            {{ debug_info }}
+        {% endif %}
+        """
+        analyzer = JinjaTemplateAnalyzer(template)
+        analysis = analyzer.analyze_template()
+
+        assert analysis["user"]["is_conditional"]
+        assert analysis["config"]["is_conditional"]
+        assert "is_admin" in analysis["user"]["accessed_properties"]
+        assert "debug_mode" in analysis["config"]["accessed_properties"]
+
+    def test_complex_expressions(self):
+        """Test handling of complex expressions and filters."""
+        LOG.info("Testing complex expressions and filters")
+
+        template = """
+        {{ user.name | upper }}
+        {{ messages | length > 0 and messages[0].content }}
+        {{ data['key'].nested['value'] }}
+        """
+        analyzer = JinjaTemplateAnalyzer(template)
+        variables = analyzer.get_template_variables()
+
+        assert "user" in variables
+        assert "name" in variables["user"]
+        assert "messages" in variables
+        assert "content" in variables["messages"]
+        assert "data" in variables
+
+    def test_basic_msg_vars(self, basic_jinja_template_analyzer):
+        """Test that the basic message variables are correctly identified."""
+        LOG.info("Testing basic message variables")
+
+        variables = basic_jinja_template_analyzer.get_message_vars()
+        assert variables == {"role", "content"}
+
+    def test_mixtral_msg_vars(self, mistral_jinja_template_analyzer):
+        """Test that the mixtral message variables are correctly identified."""
+        LOG.info("Testing mixtral message variables")
+
+        variables = mistral_jinja_template_analyzer.get_message_vars()
+        assert variables == {"role", "content", "tool_calls", "tool_call_id"}
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])
diff --git a/tests/prompt_strategies/test_kto_user_defined.py b/tests/prompt_strategies/test_kto_user_defined.py
new file mode 100644
index 0000000000..5d38f0417e
--- /dev/null
+++ b/tests/prompt_strategies/test_kto_user_defined.py
@@ -0,0 +1,140 @@
+"""
+Tests for user-defined KTO dataset transform strategies
+"""
+
+import pytest
+
+from axolotl.prompt_strategies.kto import load as load_kto
+from axolotl.prompt_strategies.kto.user_defined import default
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+def make_cfg(type_cfg: dict) -> DictDefault:
+    """Build a minimal config with a single user-defined KTO dataset."""
+    return DictDefault({"datasets": [{"type": type_cfg}]})
+
+
+class TestKTOUserDefined:
+    """
+    Test user_defined.default KTO transforms
+    """
+
+    def test_documented_config(self):
+        """Transform works with the exact config documented in docs/rlhf.qmd
+        (the same shape reported in issue #2757)."""
+        cfg = make_cfg(
+            {
+                "field_prompt": "prompt",
+                "field_system": "system",
+                "field_completion": "completion",
+                "field_label": "label",
+                "prompt_format": "{prompt}",
+                "completion_format": "{completion}",
+            }
+        )
+        transform_fn = default(cfg)
+        sample = transform_fn({"prompt": "hello", "completion": "world", "label": True})
+        assert sample["prompt"] == "hello"
+        assert sample["completion"] == "world"
+        assert sample["label"] is True
+
+    def test_defaults_without_formats(self):
+        """Transform falls back to canonical {prompt}/{completion} formats when
+        no explicit formats are configured."""
+        cfg = make_cfg({})
+        transform_fn = default(cfg)
+        sample = transform_fn(
+            {"prompt": "hello", "completion": "world", "label": False}
+        )
+        assert sample["prompt"] == "hello"
+        assert sample["completion"] == "world"
+        assert sample["label"] is False
+
+    def test_custom_field_names(self):
+        """Transform reads from custom field_prompt/field_completion/field_label
+        columns and writes the canonical prompt/completion/label keys."""
+        cfg = make_cfg(
+            {
+                "field_prompt": "question",
+                "field_completion": "answer",
+                "field_label": "is_good",
+            }
+        )
+        transform_fn = default(cfg)
+        sample = transform_fn({"question": "hello", "answer": "world", "is_good": True})
+        assert sample["prompt"] == "hello"
+        assert sample["completion"] == "world"
+        assert sample["label"] is True
+
+    def test_system_in_prompt_format(self):
+        """A {system} placeholder in prompt_format is filled from the system
+        field when present in the sample."""
+        cfg = make_cfg(
+            {
+                "prompt_format": "{system} {prompt}",
+            }
+        )
+        transform_fn = default(cfg)
+        sample = transform_fn(
+            {
+                "system": "be helpful",
+                "prompt": "hello",
+                "completion": "world",
+                "label": True,
+            }
+        )
+        assert sample["prompt"] == "be helpful hello"
+        assert sample["completion"] == "world"
+
+    def test_non_dict_type_raises(self):
+        """A non-dict dataset type raises a clear ValueError."""
+        cfg = make_cfg({})
+        cfg["datasets"][0]["type"] = "user_defined.default"
+        with pytest.raises(ValueError, match="must be a dictionary"):
+            default(cfg)
+
+    def test_load_kto_user_defined_returns_callable(self):
+        """The loader path resolves user_defined.default to a callable
+        transform for the documented config."""
+        cfg = make_cfg(
+            {
+                "field_prompt": "prompt",
+                "field_completion": "completion",
+                "field_label": "label",
+                "prompt_format": "{prompt}",
+                "completion_format": "{completion}",
+            }
+        )
+        transform_fn = load_kto("user_defined.default", cfg, dataset_idx=0)
+        assert callable(transform_fn)
+
+    def test_documented_config_passes_validate_config(self):
+        """The documented config validates end-to-end: field_label is a column
+        name (str), which the schema previously rejected as a bool."""
+        cfg = DictDefault(
+            {
+                "base_model": "TinyLlama/TinyLlama-1.1B-Chat-v0.6",
+                "learning_rate": 0.000001,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "remove_unused_columns": False,
+                "rl": "kto",
+                "datasets": [
+                    {
+                        "path": "user/dataset",
+                        "split": "train",
+                        "type": {
+                            "field_prompt": "prompt",
+                            "field_system": "system",
+                            "field_completion": "completion",
+                            "field_label": "label",
+                            "prompt_format": "{prompt}",
+                            "completion_format": "{completion}",
+                        },
+                    }
+                ],
+            }
+        )
+        checked_cfg = validate_config(cfg)
+        assert checked_cfg["datasets"][0]["type"]["field_label"] == "label"
diff --git a/tests/prompt_strategies/test_raw_io.py b/tests/prompt_strategies/test_raw_io.py
index 967de169f3..082e31ee62 100644
--- a/tests/prompt_strategies/test_raw_io.py
+++ b/tests/prompt_strategies/test_raw_io.py
@@ -1,6 +1,7 @@
 """
 Test module for raw i/o data for prompts
 """
+
 import pytest
 from datasets import Dataset
 from tokenizers import AddedToken
diff --git a/tests/prompt_strategies/test_sharegpt.py b/tests/prompt_strategies/test_sharegpt.py
deleted file mode 100644
index 6e69098340..0000000000
--- a/tests/prompt_strategies/test_sharegpt.py
+++ /dev/null
@@ -1,316 +0,0 @@
-"""
-Test module for sharegpt integration w chatml
-"""
-
-import pytest
-from datasets import Dataset
-from tokenizers import AddedToken
-from transformers import AutoTokenizer
-
-from axolotl.datasets import TokenizedPromptDataset
-from axolotl.prompt_strategies.sharegpt import (
-    GlaiveShareGPTPromptTokenizingStrategy,
-    SimpleShareGPTPromptTokenizingStrategy,
-    register_chatml_template,
-    register_llama3_template,
-)
-from axolotl.prompters import ShareGPTPrompterV2
-
-register_chatml_template()
-register_llama3_template()
-
-
-@pytest.fixture(name="sharegpt_dataset")
-def fixture_sharegpt_dataset():
-    return Dataset.from_list(
-        [
-            {
-                "conversations": [
-                    {
-                        "from": "system",
-                        "value": "repeat",
-                    },
-                    {
-                        "from": "human",
-                        "value": "hello",
-                    },
-                    {
-                        "from": "gpt",
-                        "value": "hello",
-                    },
-                    {
-                        "from": "human",
-                        "value": "goodbye",
-                    },
-                    {
-                        "from": "gpt",
-                        "value": "goodbye",
-                    },
-                ]
-            }
-        ]
-    )
-
-
-@pytest.fixture(name="glaive_dataset")
-def fixture_sharegpt_glaive_dataset():
-    return Dataset.from_list(
-        [
-            {
-                "system": "SYSTEM: This is a system prompt",
-                "chat": "USER: Can you book a flight for me from New York to London? ASSISTANT: I'm sorry, but I don't have the capability to book flights.  <|endoftext|>",
-            }
-        ]
-    )
-
-
-@pytest.fixture(name="multi_role_dataset")
-def fixture_multi_role_dataset():
-    return Dataset.from_list(
-        [
-            {
-                "conversations": [
-                    {
-                        "from": "system",
-                        "value": "use get_weather(city) to get the weather for a city",
-                    },
-                    {
-                        "from": "human",
-                        "value": "hello, what's the weather in New York?",
-                    },
-                    {
-                        "from": "gpt",
-                        "value": "let me get that for you",
-                    },
-                    {
-                        "from": "tool",
-                        "value": "get_weather(New York)",
-                    },
-                    {
-                        "from": "gpt",
-                        "value": "the weather in New York is 70 degrees and sunny",
-                    },
-                ]
-            }
-        ]
-    )
-
-
-@pytest.fixture(name="tokenizer")
-def fixture_tokenizer():
-    tokenizer = AutoTokenizer.from_pretrained(
-        "casperhansen/mistral-7b-instruct-v0.1-awq"
-    )
-    tokenizer.add_special_tokens(
-        {
-            "eos_token": AddedToken(
-                "<|im_end|>", rstrip=False, lstrip=False, normalized=False
-            )
-        }
-    )
-    tokenizer.add_tokens(
-        [
-            AddedToken("<|im_start|>", rstrip=False, lstrip=False, normalized=False),
-        ]
-    )
-
-    return tokenizer
-
-
-@pytest.fixture(name="llama3_tokenizer")
-def fixture_llama3_tokenizer():
-    tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B")
-    tokenizer.eos_token = "<|eot_id|>"
-
-    return tokenizer
-
-
-class TestSharegptLlama3:
-    """Test class for ShareGPT style datasets with llama-3 prompts"""
-
-    def test_tokenization(self, sharegpt_dataset, llama3_tokenizer):
-        strategy = SimpleShareGPTPromptTokenizingStrategy(
-            ShareGPTPrompterV2(
-                conversation="llama3",
-                role_key_model=None,
-                role_key_human=None,
-            ),
-            llama3_tokenizer,
-            False,  # train_on_inputs
-            2048,  # sequence_len
-        )
-
-        dataset_wrapper = TokenizedPromptDataset(
-            strategy, sharegpt_dataset, process_count=1
-        )
-
-        input_ids = dataset_wrapper[0]["input_ids"]
-
-        # fmt: off
-        assert input_ids == [
-            128000,  # bos
-            128006, 9125, 128007,  # system header
-            271, 31724, 128009,  # sys prompt, eot
-            128006, 882, 128007,  # user header
-            271, 15339, 128009,  # user prompt eot
-            128006, 78191, 128007,  # assistant header
-            271, 15339, 128009,   # assistant response eot
-            128006, 882, 128007,
-            271, 19045, 29474, 128009,
-            128006, 78191, 128007,
-            271, 19045, 29474, 128009,
-        ]
-        # fmt: on
-
-
-class TestSharegptChatML:
-    """
-    Test class for sharegpt prompter
-    """
-
-    def test_no_double_im_end(self, sharegpt_dataset, tokenizer):
-        strategy = SimpleShareGPTPromptTokenizingStrategy(
-            ShareGPTPrompterV2(
-                conversation="chatml",
-                role_key_model=None,
-                role_key_human=None,
-            ),
-            tokenizer,
-            False,  # train_on_inputs
-            2048,  # sequence_len
-        )
-
-        dataset_wrapper = TokenizedPromptDataset(
-            strategy, sharegpt_dataset, process_count=1
-        )
-
-        input_ids = dataset_wrapper[0]["input_ids"]
-        # fmt: off
-        assert input_ids == [
-            #  28705, 13, is " \n"
-            1,   # bos
-            32001, 1587, 13, 25997, 32000, 28705, 13,  # system
-            32001, 2188, 13, 21558, 32000, 28705, 13,  # human
-            32001, 13892, 13, 21558, 32000, 28705, 13,  # gpt
-            32001, 2188, 13, 12684, 17664, 32000, 28705, 13,   # human
-            32001, 13892, 13, 12684, 17664, 32000, 28705, 13,  # gpt
-        ]
-        # fmt: on
-
-    def test_w_train_on_input(self, sharegpt_dataset, tokenizer):
-        strategy = SimpleShareGPTPromptTokenizingStrategy(
-            ShareGPTPrompterV2(
-                conversation="chatml",
-                role_key_model=None,
-                role_key_human=None,
-            ),
-            tokenizer,
-            False,  # train_on_inputs
-            2048,  # sequence_len
-        )
-
-        dataset_wrapper = TokenizedPromptDataset(
-            strategy, sharegpt_dataset, process_count=1
-        )
-
-        labels = dataset_wrapper[0]["labels"]
-        # fmt: off
-        assert labels == [
-            -100,   # bos
-            -100, -100, -100, -100, -100, -100, -100,  # system
-            -100, -100, -100, -100, -100, -100, -100,  # human
-            -100, -100, 13, 21558, 32000, 28705, 13,  # gpt
-            -100, -100, -100, -100, -100, -100, -100, -100,   # human
-            -100, -100, 13, 12684, 17664, 32000, 28705, 13,  # gpt
-        ]
-        # fmt: on
-
-    def test_no_train_on_input(self, sharegpt_dataset, tokenizer):
-        strategy = SimpleShareGPTPromptTokenizingStrategy(
-            ShareGPTPrompterV2(
-                conversation="chatml",
-                role_key_model=None,
-                role_key_human=None,
-            ),
-            tokenizer,
-            True,  # train_on_inputs
-            2048,  # sequence_len
-        )
-
-        dataset_wrapper = TokenizedPromptDataset(
-            strategy, sharegpt_dataset, process_count=1
-        )
-
-        labels = dataset_wrapper[0]["labels"]
-        # fmt: off
-        assert labels == [
-            1,   # bos
-            32001, 1587, 13, 25997, 32000, 28705, 13,  # system
-            32001, 2188, 13, 21558, 32000, 28705, 13,  # human
-            32001, 13892, 13, 21558, 32000, 28705, 13,  # gpt
-            32001, 2188, 13, 12684, 17664, 32000, 28705, 13,   # human
-            32001, 13892, 13, 12684, 17664, 32000, 28705, 13,  # gpt
-        ]
-        # fmt: on
-
-    def test_chatml_glaive(self, glaive_dataset, tokenizer):
-        strategy = GlaiveShareGPTPromptTokenizingStrategy(
-            ShareGPTPrompterV2(
-                conversation="chatml",
-                role_key_model=None,
-                role_key_human=None,
-            ),
-            tokenizer,
-            True,  # train_on_inputs
-            2048,  # sequence_len
-        )
-
-        dataset_wrapper = TokenizedPromptDataset(
-            strategy, glaive_dataset, process_count=1
-        )
-
-        labels = dataset_wrapper[0]["labels"]
-        # fmt: off
-        assert labels == [
-            1,  # bos
-            32001, 1587, 13, 3260, 349, 264, 1587, 11510, 32000, 28705, 13,  # system
-            32001, 2188, 13, 6325, 368, 1820, 264, 9314, 354, 528, 477, 1450, 2726, 298, 4222, 28804, 32000, 28705, 13,  # human
-            32001, 13892, 13, 28737, 28742, 28719, 7371, 28725, 562, 315, 949, 28742, 28707, 506, 272, 21368, 298, 1820, 22447, 28723, 28705, 523, 28766, 416, 1009, 772, 28766, 28767, 32000, 28705, 13  # gpt
-        ]
-        # fmt: on
-
-    def test_multi_role_dataset(self, multi_role_dataset, tokenizer):
-        strategy = SimpleShareGPTPromptTokenizingStrategy(
-            ShareGPTPrompterV2(conversation="chatml", roles={"input": ["tool"]}),
-            tokenizer,
-            False,  # train_on_inputs
-            2048,  # sequence_len
-        )
-
-        dataset_wrapper = TokenizedPromptDataset(
-            strategy, multi_role_dataset, process_count=1
-        )
-
-        input_ids = dataset_wrapper[0]["input_ids"]
-        # fmt: off
-        assert input_ids == [
-            1,   # bos
-            32001, 1587, 13, 1730, 625, 28730, 769, 1223, 28732, 18373, 28731, 298, 625, 272, 8086, 354, 264, 2990, 32000, 28705, 13,  # system
-            32001, 2188, 13, 21558, 28725, 767, 28742, 28713, 272, 8086, 297, 1450, 2726, 28804, 32000, 28705, 13,  # human
-            32001, 13892, 13, 895, 528, 625, 369, 354, 368, 32000, 28705, 13,  # gpt
-            32001, 3921, 13, 527, 28730, 769, 1223, 28732, 2972, 2726, 28731, 32000, 28705, 13,  # tool
-            32001, 13892, 13, 1237, 8086, 297, 1450, 2726, 349, 28705, 28787, 28734, 11182, 304, 4376, 1780, 32000, 28705, 13  # gpt
-        ]
-        # fmt: on
-
-        labels = dataset_wrapper[0]["labels"]
-        # fmt: off
-        assert labels == [
-            -100,  # bos
-            -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,  # system
-            -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,  # human
-            -100, -100, 13, 895, 528, 625, 369, 354, 368, 32000, 28705, 13,  # gpt
-            -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,  # tool
-            -100, -100, 13, 1237, 8086, 297, 1450, 2726, 349, 28705, 28787, 28734, 11182, 304, 4376, 1780, 32000, 28705, 13  # gpt
-        ]
-        # fmt: on
diff --git a/tests/prompt_strategies/test_stepwise.py b/tests/prompt_strategies/test_stepwise.py
new file mode 100644
index 0000000000..b9e7e9504d
--- /dev/null
+++ b/tests/prompt_strategies/test_stepwise.py
@@ -0,0 +1,63 @@
+"""
+tests for chat_template prompt strategy
+"""
+
+import pytest
+from datasets import Dataset
+
+from axolotl.prompt_strategies.stepwise_supervised import (
+    StepwiseSupervisedPromptTokenizingStrategy,
+)
+
+
+class TestStepWiseSupervisedPromptTokenizingStrategy:
+    """
+    Test class for stepwise supervised prompt strategy
+    """
+
+    @pytest.fixture()
+    def stepwise_supervised_dataset(self):
+        return Dataset.from_list(
+            [
+                {
+                    "prompt": "Which number is larger, 9.8 or 9.11?",
+                    "completions": [
+                        "The fractional part of 9.8 is 0.8, while the fractional part of 9.11 is 0.11.",
+                        "Since 0.11 is greater than 0.8, the number 9.11 is larger than 9.8.",
+                        "Actually, this is incorrect. In decimal numbers, 0.8 is equal to 0.80, which is larger than 0.11. Therefore, 9.8 is larger than 9.11.",
+                    ],
+                    "labels": [True, False, False],
+                }
+            ]
+        )
+
+    def test_stepwise_supervised_dataset(
+        self, qwen3_tokenizer, stepwise_supervised_dataset
+    ):
+        strategy = StepwiseSupervisedPromptTokenizingStrategy(
+            qwen3_tokenizer,
+            sequence_len=2048,
+            step_separator="\n",
+        )
+        sample = stepwise_supervised_dataset[0]
+        labels = strategy.tokenize_prompt(sample)["labels"]
+        prompt_len = len(
+            qwen3_tokenizer(sample["prompt"], add_special_tokens=False)["input_ids"]
+        )
+        if qwen3_tokenizer.bos_token_id is not None:
+            prompt_len += 1
+
+        separator_len = len(qwen3_tokenizer.encode("\n", add_special_tokens=False))
+        completion_lengths = [
+            len(qwen3_tokenizer(completion, add_special_tokens=False)["input_ids"])
+            + separator_len
+            for completion in sample["completions"]
+        ]
+        expected = [-100] * prompt_len
+        for completion_len, label in zip(
+            completion_lengths, sample["labels"], strict=False
+        ):
+            expected.extend([-100] * (completion_len - 1))
+            expected.append(int(label))
+
+        assert labels == expected
diff --git a/tests/prompt_strategies/test_synthetic.py b/tests/prompt_strategies/test_synthetic.py
new file mode 100644
index 0000000000..6038333c1f
--- /dev/null
+++ b/tests/prompt_strategies/test_synthetic.py
@@ -0,0 +1,125 @@
+"""Tests for the synthetic dataset generator."""
+
+import unittest
+from unittest.mock import MagicMock
+
+from datasets import Dataset
+
+from axolotl.prompt_strategies._synthetic import SyntheticDatasetStrategy, load
+from axolotl.utils.dict import DictDefault
+
+
+class TestSyntheticDatasetStrategy(unittest.TestCase):
+    def test_generates_correct_shape(self):
+        strategy = SyntheticDatasetStrategy(
+            sequence_length=128,
+            length=50,
+            min_input_id=1,
+            max_input_id=1000,
+            seed=42,
+        )
+        dummy = Dataset.from_dict({"text": [""]})
+        result = strategy.wrap_dataset(dummy)
+
+        assert len(result) == 50
+        assert len(result[0]["input_ids"]) == 128
+        assert len(result[0]["attention_mask"]) == 128
+        assert len(result[0]["labels"]) == 128
+
+    def test_attention_mask_all_ones(self):
+        strategy = SyntheticDatasetStrategy(sequence_length=64, length=10, seed=0)
+        dummy = Dataset.from_dict({"text": [""]})
+        result = strategy.wrap_dataset(dummy)
+
+        for row in result:
+            assert all(v == 1 for v in row["attention_mask"])
+
+    def test_labels_equal_input_ids(self):
+        strategy = SyntheticDatasetStrategy(sequence_length=64, length=10, seed=0)
+        dummy = Dataset.from_dict({"text": [""]})
+        result = strategy.wrap_dataset(dummy)
+
+        for row in result:
+            assert row["input_ids"] == row["labels"]
+
+    def test_input_id_range(self):
+        strategy = SyntheticDatasetStrategy(
+            sequence_length=64,
+            length=100,
+            min_input_id=500,
+            max_input_id=600,
+            seed=42,
+        )
+        dummy = Dataset.from_dict({"text": [""]})
+        result = strategy.wrap_dataset(dummy)
+
+        for row in result:
+            for token_id in row["input_ids"]:
+                assert 500 <= token_id < 600
+
+    def test_seed_reproducibility(self):
+        kwargs = dict(
+            sequence_length=64, length=20, min_input_id=1, max_input_id=1000, seed=123
+        )
+        dummy = Dataset.from_dict({"text": [""]})
+
+        result1 = SyntheticDatasetStrategy(**kwargs).wrap_dataset(dummy)
+        result2 = SyntheticDatasetStrategy(**kwargs).wrap_dataset(dummy)
+
+        for r1, r2 in zip(result1, result2, strict=True):
+            assert r1["input_ids"] == r2["input_ids"]
+
+    def test_different_seeds_differ(self):
+        common = dict(sequence_length=64, length=20, min_input_id=1, max_input_id=1000)
+        dummy = Dataset.from_dict({"text": [""]})
+
+        result1 = SyntheticDatasetStrategy(seed=1, **common).wrap_dataset(dummy)
+        result2 = SyntheticDatasetStrategy(seed=2, **common).wrap_dataset(dummy)
+
+        any_different = any(
+            r1["input_ids"] != r2["input_ids"]
+            for r1, r2 in zip(result1, result2, strict=True)
+        )
+        assert any_different
+
+    def test_load_function_with_ds_cfg(self):
+        tokenizer = MagicMock()
+        tokenizer.vocab_size = 32000
+        cfg = DictDefault({"sequence_len": 512, "train_on_inputs": False})
+        ds_cfg = {
+            "sequence_length": 256,
+            "length": 5,
+            "min_input_id": 10,
+            "max_input_id": 100,
+            "seed": 0,
+        }
+
+        strategy = load(tokenizer, cfg, ds_cfg=ds_cfg)
+        assert isinstance(strategy, SyntheticDatasetStrategy)
+        assert strategy.sequence_length == 256
+        assert strategy.length == 5
+        assert strategy.min_input_id == 10
+        assert strategy.max_input_id == 100
+
+    def test_load_defaults_from_cfg(self):
+        tokenizer = MagicMock()
+        tokenizer.vocab_size = 32000
+        cfg = DictDefault({"sequence_len": 1024, "train_on_inputs": False})
+
+        strategy = load(tokenizer, cfg, ds_cfg={})
+        assert strategy.sequence_length == 1024
+        assert strategy.max_input_id == 32000
+        assert strategy.length == 1000
+
+    def test_load_with_no_ds_cfg(self):
+        tokenizer = MagicMock()
+        tokenizer.vocab_size = 50000
+        cfg = DictDefault({"sequence_len": 2048, "train_on_inputs": False})
+
+        strategy = load(tokenizer, cfg)
+        assert strategy.sequence_length == 2048
+        assert strategy.max_input_id == 50000
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/telemetry/__init__.py b/tests/telemetry/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/telemetry/conftest.py b/tests/telemetry/conftest.py
new file mode 100644
index 0000000000..47776ce906
--- /dev/null
+++ b/tests/telemetry/conftest.py
@@ -0,0 +1,9 @@
+"""Shared pytest fixtures for telemetry tests."""
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def del_track_env(monkeypatch):
+    monkeypatch.delenv("AXOLOTL_DO_NOT_TRACK", raising=False)
+    yield
diff --git a/tests/telemetry/test_callbacks.py b/tests/telemetry/test_callbacks.py
new file mode 100644
index 0000000000..97d56a9c67
--- /dev/null
+++ b/tests/telemetry/test_callbacks.py
@@ -0,0 +1,373 @@
+"""Tests for telemetry callback module."""
+
+# pylint: disable=redefined-outer-name
+
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+from transformers import TrainerControl, TrainerState, TrainingArguments
+
+from axolotl.telemetry.callbacks import TIME_SINCE_LAST, TelemetryCallback
+
+
+def calc_expected_metrics(step, last_step, current_time, last_time, start_time=900.0):
+    """Calculate expected metrics values for tests"""
+    time_diff = current_time - last_time
+    step_diff = step - last_step
+    return {
+        "steps_per_second": (
+            step_diff / time_diff if time_diff > 0 and step_diff > 0 else 0
+        ),
+        "time_since_last_report": time_diff,
+        "elapsed_time": current_time - start_time,
+    }
+
+
+@pytest.fixture
+def mock_time():
+    """Mock time.time() to have predictable values in tests"""
+    with patch("axolotl.telemetry.callbacks.time") as mock_time:
+        mock_time.time.return_value = 1000.0
+        yield mock_time
+
+
+@pytest.fixture
+def mock_telemetry_manager():
+    """Create a mock TelemetryManager"""
+    with patch("axolotl.telemetry.callbacks.TelemetryManager") as mock_manager_class:
+        mock_manager = MagicMock()
+        mock_manager_class.get_instance.return_value = mock_manager
+        yield mock_manager
+
+
+@pytest.fixture
+def mock_runtime_metrics_tracker():
+    """Create a mock RuntimeMetricsTracker"""
+    with patch(
+        "axolotl.telemetry.callbacks.RuntimeMetricsTracker"
+    ) as mock_tracker_class:
+        mock_tracker = MagicMock()
+        # Set up metrics property on the tracker
+        mock_metrics = MagicMock()
+        mock_metrics.to_dict.return_value = {
+            "total_steps": 100,
+            "peak_cpu_memory_bytes": 1024,
+        }
+        mock_tracker.metrics = mock_metrics
+
+        # Make the constructor return our mock
+        mock_tracker_class.return_value = mock_tracker
+        yield mock_tracker
+
+
+@pytest.fixture
+def training_args():
+    """Create a minimal TrainingArguments instance"""
+    return TrainingArguments(output_dir="./output")
+
+
+@pytest.fixture
+def trainer_state():
+    """Create a mock TrainerState"""
+    state = MagicMock(spec=TrainerState)
+    state.global_step = 10
+    state.epoch = 0.5  # halfway through first epoch
+    state.log_history = [{"loss": 2.5, "learning_rate": 5e-5}]
+    return state
+
+
+@pytest.fixture
+def trainer_control():
+    """Create a mock TrainerControl"""
+    return MagicMock(spec=TrainerControl)
+
+
+# pylint: disable=unused-argument
+@pytest.fixture
+def callback(mock_telemetry_manager, mock_runtime_metrics_tracker):
+    """Create a TelemetryCallback instance with mocked dependencies"""
+    return TelemetryCallback()
+
+
+class TestTelemetryCallback:
+    """Tests for the TelemetryCallback class."""
+
+    def test_initialization(self, callback, mock_runtime_metrics_tracker):
+        """Test callback initialization."""
+        assert callback.current_epoch == -1
+        assert callback.tracker == mock_runtime_metrics_tracker
+        assert callback.last_report_step == 0
+        assert hasattr(callback, "start_time")
+        assert hasattr(callback, "last_report_time")
+        assert callback.report_interval_steps == 100
+
+    def test_on_train_begin(
+        self,
+        callback,
+        mock_telemetry_manager,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test on_train_begin sends expected event."""
+        callback.on_train_begin(training_args, trainer_state, trainer_control)
+
+        mock_telemetry_manager.send_event.assert_called_once_with(
+            event_type="train-start"
+        )
+
+    def test_on_train_end(
+        self,
+        callback,
+        mock_telemetry_manager,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test on_train_end sends expected event with metrics."""
+        callback.on_train_end(training_args, trainer_state, trainer_control)
+
+        mock_telemetry_manager.send_event.assert_called_once()
+        call_args = mock_telemetry_manager.send_event.call_args[1]
+
+        assert call_args["event_type"] == "train-end"
+        assert "loss" in call_args["properties"]
+        assert call_args["properties"]["loss"] == 2.5
+        assert "learning_rate" in call_args["properties"]
+        assert call_args["properties"]["learning_rate"] == 5e-5
+
+        # Check that metrics from RuntimeMetricsTracker are included
+        assert "total_steps" in call_args["properties"]
+        assert call_args["properties"]["total_steps"] == 100
+        assert "peak_cpu_memory_bytes" in call_args["properties"]
+        assert call_args["properties"]["peak_cpu_memory_bytes"] == 1024
+
+    def test_on_epoch_begin(
+        self,
+        callback,
+        mock_runtime_metrics_tracker,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test on_epoch_begin updates epoch counter and calls tracker."""
+        initial_epoch = callback.current_epoch
+
+        callback.on_epoch_begin(training_args, trainer_state, trainer_control)
+
+        assert callback.current_epoch == initial_epoch + 1
+        mock_runtime_metrics_tracker.start_epoch.assert_called_once_with(
+            initial_epoch + 1
+        )
+
+    def test_on_epoch_end(
+        self,
+        callback,
+        mock_runtime_metrics_tracker,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test on_epoch_end calls tracker."""
+        # Set current epoch
+        callback.current_epoch = 2
+
+        callback.on_epoch_end(training_args, trainer_state, trainer_control)
+
+        mock_runtime_metrics_tracker.end_epoch.assert_called_once_with(2)
+
+    def test_on_step_end_no_report(
+        self,
+        callback,
+        mock_telemetry_manager,
+        mock_runtime_metrics_tracker,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test on_step_end updates tracker but doesn't report if criteria not met."""
+        # Set up state to avoid reporting
+        trainer_state.global_step = 42  # Not divisible by report_interval_steps
+        callback.last_report_step = 41  # Just 1 step since last report
+        callback.last_report_time = time.time()  # Just now
+
+        callback.on_step_end(training_args, trainer_state, trainer_control)
+
+        # Should update tracker
+        mock_runtime_metrics_tracker.update_step.assert_called_once_with(42)
+
+        # Should not send telemetry
+        mock_telemetry_manager.send_event.assert_not_called()
+
+        # Should not update last report time/step
+        assert callback.last_report_step == 41
+
+    def test_on_step_end_report_interval_steps(
+        self,
+        callback,
+        mock_telemetry_manager,
+        mock_runtime_metrics_tracker,
+        mock_time,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test on_step_end reports when step interval is reached."""
+        # Set up state with clear values
+        current_step = 100  # Exactly matches report_interval_steps
+        last_step = 0
+        start_time = 900.0
+        current_time = 1000.0
+        time_diff = current_time - start_time  # 100 seconds
+
+        # Configure state and callback
+        trainer_state.global_step = current_step
+        callback.report_interval_steps = 100
+        callback.last_report_step = last_step
+        callback.start_time = start_time
+        callback.last_report_time = start_time
+
+        # Mock time.time() to return consistent values
+        mock_time.time.return_value = current_time
+
+        callback.on_step_end(training_args, trainer_state, trainer_control)
+
+        # Should update tracker
+        mock_runtime_metrics_tracker.update_step.assert_called_once_with(current_step)
+        mock_runtime_metrics_tracker.update_memory_metrics.assert_called_once()
+
+        # Should send telemetry
+        mock_telemetry_manager.send_event.assert_called_once()
+        call_args = mock_telemetry_manager.send_event.call_args[1]
+        assert call_args["event_type"] == "train-progress"
+
+        # Properties should include expected values
+        props = call_args["properties"]
+        assert props["step"] == current_step
+        assert props["elapsed_time"] == time_diff  # 1000 - 900 = 100
+        assert props["time_since_last_report"] == time_diff  # 1000 - 900 = 100
+        assert props["steps_per_second"] == 1.0  # 100 steps / 100 seconds
+
+        # Should update last report time/step
+        assert callback.last_report_step == current_step
+        assert callback.last_report_time == current_time
+
+    def test_on_step_end_report_time_elapsed(
+        self,
+        callback,
+        mock_telemetry_manager,
+        mock_runtime_metrics_tracker,  # pylint: disable=unused-argument
+        mock_time,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test on_step_end reports when enough time has elapsed."""
+        # Set up state with clear values
+        current_step = 120
+        last_step = 10
+        start_time = 900.0
+        current_time = 1000.0
+        time_diff = TIME_SINCE_LAST + 1  # Just over the threshold
+
+        # Configure state and callback
+        trainer_state.global_step = current_step
+        callback.report_interval_steps = 100
+        callback.last_report_step = last_step
+        callback.start_time = start_time
+        callback.last_report_time = current_time - time_diff
+
+        # Mock time.time() to return consistent values
+        mock_time.time.return_value = current_time
+
+        callback.on_step_end(training_args, trainer_state, trainer_control)
+
+        # Should send telemetry
+        mock_telemetry_manager.send_event.assert_called_once()
+
+        # Properties should include expected values
+        props = mock_telemetry_manager.send_event.call_args[1]["properties"]
+        expected_metrics = calc_expected_metrics(
+            current_step, last_step, current_time, current_time - time_diff, start_time
+        )
+        assert props["steps_per_second"] == expected_metrics["steps_per_second"]
+        assert (
+            props["time_since_last_report"]
+            == expected_metrics["time_since_last_report"]
+        )
+
+    def test_on_step_end_first_step(
+        self,
+        callback,
+        mock_telemetry_manager,
+        mock_runtime_metrics_tracker,  # pylint: disable=unused-argument
+        mock_time,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test on_step_end always reports on first step."""
+        # Set up state with clear values
+        current_step = 1  # First step
+        last_step = 0
+        start_time = 900.0
+        current_time = 1000.0
+        last_report_time = 999.0  # Just 1 second ago
+
+        # Configure state and callback
+        trainer_state.global_step = current_step
+        callback.report_interval_steps = 100
+        callback.last_report_step = last_step
+        callback.start_time = start_time
+        callback.last_report_time = last_report_time
+
+        # Mock time.time() to return consistent values
+        mock_time.time.return_value = current_time
+
+        callback.on_step_end(training_args, trainer_state, trainer_control)
+
+        # Should send telemetry even though not much time has passed
+        mock_telemetry_manager.send_event.assert_called_once()
+
+        # Properties should include expected values for first step
+        props = mock_telemetry_manager.send_event.call_args[1]["properties"]
+        assert props["step"] == current_step
+        expected_metrics = calc_expected_metrics(
+            current_step, last_step, current_time, last_report_time, start_time
+        )
+        assert props["steps_per_second"] == expected_metrics["steps_per_second"]
+
+    def test_log_history_empty(
+        self,
+        callback,
+        mock_telemetry_manager,
+        mock_runtime_metrics_tracker,  # pylint: disable=unused-argument
+        mock_time,
+        training_args,
+        trainer_state,
+        trainer_control,
+    ):
+        """Test handling of empty log history."""
+        # Set up state with clear values
+        current_step = 1
+        start_time = 900.0
+        current_time = 1000.0
+
+        # Configure state and callback
+        trainer_state.global_step = current_step
+        trainer_state.log_history = []
+        callback.start_time = start_time
+
+        # Mock time.time() to return consistent values
+        mock_time.time.return_value = current_time
+
+        callback.on_step_end(training_args, trainer_state, trainer_control)
+
+        # Should still send telemetry
+        mock_telemetry_manager.send_event.assert_called_once()
+
+        # Properties should have default values for missing log data
+        props = mock_telemetry_manager.send_event.call_args[1]["properties"]
+        assert props["loss"] == 0
+        assert props["learning_rate"] == 0
diff --git a/tests/telemetry/test_errors.py b/tests/telemetry/test_errors.py
new file mode 100644
index 0000000000..2f0510b212
--- /dev/null
+++ b/tests/telemetry/test_errors.py
@@ -0,0 +1,341 @@
+"""Tests for telemetry error utilities"""
+
+# pylint: disable=redefined-outer-name
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from axolotl.telemetry.errors import sanitize_stack_trace, send_errors
+
+
+@pytest.fixture(autouse=True)
+def reset_error_flag(monkeypatch):
+    """Reset ERROR_HANDLED flag using monkeypatch"""
+    import axolotl.telemetry.errors
+
+    monkeypatch.setattr(axolotl.telemetry.errors, "ERROR_HANDLED", False)
+    yield
+    monkeypatch.setattr(axolotl.telemetry.errors, "ERROR_HANDLED", False)
+
+
+@pytest.fixture
+def example_stack_trace():
+    """Provide a sample stack trace with mixed paths"""
+    return """Traceback (most recent call last):
+  File "/home/user/.local/lib/python3.9/site-packages/axolotl/cli/train.py", line 83, in main
+    trainer = get_trainer(cfg)
+  File "/home/user/.local/lib/python3.9/site-packages/axolotl/train.py", line 214, in get_trainer
+    model = get_model(cfg, tokenizer)
+  File "/home/user/.local/lib/python3.9/site-packages/axolotl/utils/models.py", line 120, in get_model
+    raise ValueError("Model path not found")
+ValueError: Model path not found
+"""
+
+
+@pytest.fixture
+def windows_stack_trace():
+    """Provide a sample stack trace with Windows paths"""
+    return """Traceback (most recent call last):
+  File "C:\\Users\\name\\AppData\\Local\\Programs\\Python\\Python39\\lib\\site-packages\\axolotl\\cli\\train.py", line 83, in main
+    trainer = get_trainer(cfg)
+  File "C:\\Users\\name\\AppData\\Local\\Programs\\Python\\Python39\\lib\\site-packages\\axolotl\\train.py", line 214, in get_trainer
+    model = get_model(cfg, tokenizer)
+  File "C:\\Users\\name\\AppData\\Local\\Programs\\Python\\Python39\\lib\\site-packages\\transformers\\models\\auto\\modeling_auto.py", line 482, in from_pretrained
+    raise ValueError(f"Unrecognized configuration class {config.__class__}")
+ValueError: Unrecognized configuration class <class 'transformers.models.llama.configuration_llama.LlamaConfig'>
+"""
+
+
+@pytest.fixture
+def mixed_stack_trace():
+    """Provide a sample stack trace with both axolotl and non-axolotl paths"""
+    return """Traceback (most recent call last):
+  File "/home/user/.local/lib/python3.9/site-packages/axolotl/cli/train.py", line 83, in main
+    trainer = get_trainer(cfg)
+  File "/home/user/.local/lib/python3.9/site-packages/transformers/trainer.py", line 520, in train
+    self._inner_training_loop()
+  File "/home/user/.local/lib/python3.9/site-packages/axolotl/utils/trainer.py", line 75, in _inner_training_loop
+    super()._inner_training_loop()
+  File "/home/user/.local/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 631, in __next__
+    data = self._next_data()
+RuntimeError: CUDA out of memory
+"""
+
+
+@pytest.fixture
+def venv_stack_trace():
+    """Provide a sample stack trace with virtual environment paths"""
+    return """Traceback (most recent call last):
+  File "/home/user/venv/lib/python3.9/site-packages/transformers/trainer.py", line 1729, in train
+    self._inner_training_loop()
+  File "/home/user/venv/lib/python3.9/site-packages/transformers/trainer.py", line 2013, in _inner_training_loop
+    self.accelerator.backward(loss)
+  File "/home/user/venv/lib/python3.9/site-packages/accelerate/accelerator.py", line 1851, in backward
+    self.scaler.scale(loss).backward(**kwargs)
+  File "/home/user/venv/lib/python3.9/site-packages/torch/_tensor.py", line 487, in backward
+    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
+RuntimeError: CUDA out of memory
+"""
+
+
+@pytest.fixture
+def dist_packages_stack_trace():
+    """Provide a sample stack trace with dist-packages paths"""
+    return """Traceback (most recent call last):
+  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py", line 631, in __next__
+    data = self._next_data()
+  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py", line 675, in _next_data
+    data = self._dataset_fetcher.fetch(index)
+  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
+    data = [self.dataset[idx] for idx in possibly_batched_index]
+  File "/usr/local/lib/python3.8/dist-packages/datasets/arrow_dataset.py", line 2808, in __getitem__
+    raise IndexError(f"Index {key} out of range for dataset of length {len(self)}.")
+IndexError: Index 10000 out of range for dataset of length 9832.
+"""
+
+
+@pytest.fixture
+def project_stack_trace():
+    """Provide a sample stack trace from a project directory (not a virtual env)"""
+    return """Traceback (most recent call last):
+  File "/home/user/projects/myproject/run.py", line 25, in <module>
+    main()
+  File "/home/user/projects/myproject/src/cli.py", line 45, in main
+    app.run()
+  File "/home/user/projects/myproject/src/app.py", line 102, in run
+    raise ValueError("Configuration missing")
+ValueError: Configuration missing
+"""
+
+
+def test_sanitize_stack_trace(example_stack_trace):
+    """Test that sanitize_stack_trace properly preserves axolotl paths"""
+    sanitized = sanitize_stack_trace(example_stack_trace)
+
+    # Check that personal paths are removed
+    assert "/home/user" not in sanitized
+    assert ".local/lib/python3.9" not in sanitized
+
+    # Check that site-packages is preserved
+    assert "site-packages/axolotl/cli/train.py" in sanitized
+    assert "site-packages/axolotl/train.py" in sanitized
+    assert "site-packages/axolotl/utils/models.py" in sanitized
+
+    # Check that error message is preserved
+    assert "ValueError: Model path not found" in sanitized
+
+
+def test_sanitize_windows_paths(windows_stack_trace):
+    """Test that sanitize_stack_trace handles Windows paths"""
+    sanitized = sanitize_stack_trace(windows_stack_trace)
+
+    # Check that personal paths are removed
+    assert "C:\\Users\\name" not in sanitized
+    assert "AppData\\Local\\Programs\\Python" not in sanitized
+
+    # Check that both axolotl and transformers packages are preserved
+    assert (
+        "site-packages\\axolotl\\cli\\train.py" in sanitized
+        or "site-packages/axolotl/cli/train.py" in sanitized
+    )
+    assert (
+        "site-packages\\axolotl\\train.py" in sanitized
+        or "site-packages/axolotl/train.py" in sanitized
+    )
+    assert (
+        "site-packages\\transformers\\models\\auto\\modeling_auto.py" in sanitized
+        or "site-packages/transformers/models/auto/modeling_auto.py" in sanitized
+    )
+
+    # Check that error message is preserved
+    assert "ValueError: Unrecognized configuration class" in sanitized
+
+
+def test_sanitize_mixed_paths(mixed_stack_trace):
+    """Test that sanitize_stack_trace preserves all package paths"""
+    sanitized = sanitize_stack_trace(mixed_stack_trace)
+
+    # Check that all package paths are preserved
+    assert "site-packages/axolotl/cli/train.py" in sanitized
+    assert "site-packages/transformers/trainer.py" in sanitized
+    assert "site-packages/axolotl/utils/trainer.py" in sanitized
+    assert "site-packages/torch/utils/data/dataloader.py" in sanitized
+
+    # Check that error message is preserved
+    assert "RuntimeError: CUDA out of memory" in sanitized
+
+
+def test_sanitize_venv_paths(venv_stack_trace):
+    """Test that sanitize_stack_trace preserves virtual environment package paths"""
+    sanitized = sanitize_stack_trace(venv_stack_trace)
+
+    # Check that personal paths are removed
+    assert "/home/user/venv" not in sanitized
+
+    # Check that all package paths are preserved
+    assert "site-packages/transformers/trainer.py" in sanitized
+    assert "site-packages/accelerate/accelerator.py" in sanitized
+    assert "site-packages/torch/_tensor.py" in sanitized
+
+    # Check that error message is preserved
+    assert "RuntimeError: CUDA out of memory" in sanitized
+
+
+def test_sanitize_dist_packages(dist_packages_stack_trace):
+    """Test that sanitize_stack_trace preserves dist-packages paths"""
+    sanitized = sanitize_stack_trace(dist_packages_stack_trace)
+
+    # Check that system paths are removed
+    assert "/usr/local/lib/python3.8" not in sanitized
+
+    # Check that all package paths are preserved
+    assert "dist-packages/torch/utils/data/dataloader.py" in sanitized
+    assert "dist-packages/torch/utils/data/_utils/fetch.py" in sanitized
+    assert "dist-packages/datasets/arrow_dataset.py" in sanitized
+
+    # Check that error message is preserved
+    assert (
+        "IndexError: Index 10000 out of range for dataset of length 9832." in sanitized
+    )
+
+
+def test_sanitize_project_paths(project_stack_trace):
+    """Test handling of project paths (non-virtual env)"""
+    sanitized = sanitize_stack_trace(project_stack_trace)
+
+    # Check that personal paths are removed
+    assert "/home/user/projects" not in sanitized
+
+    # For non-package paths, we should at least preserve the filename
+    assert "run.py" in sanitized
+    assert "cli.py" in sanitized
+    assert "app.py" in sanitized
+
+    # Check that error message is preserved
+    assert "ValueError: Configuration missing" in sanitized
+
+
+@pytest.fixture
+def mock_telemetry_manager():
+    """Create a mock TelemetryManager"""
+    with patch("axolotl.telemetry.errors.TelemetryManager") as mock_manager_class:
+        mock_manager = MagicMock()
+        mock_manager.enabled = True
+        mock_manager_class.get_instance.return_value = mock_manager
+        yield mock_manager
+
+
+def test_send_errors_successful_execution(mock_telemetry_manager):
+    """Test that send_errors doesn't send telemetry for successful function execution"""
+
+    @send_errors
+    def test_func():
+        return "success"
+
+    result = test_func()
+    assert result == "success"
+    mock_telemetry_manager.send_event.assert_not_called()
+
+
+def test_send_errors_with_exception(mock_telemetry_manager):
+    """Test that send_errors sends telemetry when an exception occurs"""
+    test_error = ValueError("Test error")
+
+    @send_errors
+    def test_func():
+        raise test_error
+
+    with pytest.raises(ValueError) as excinfo:
+        test_func()
+
+    assert excinfo.value == test_error
+    mock_telemetry_manager.send_event.assert_called_once()
+
+    # Check that the error info was passed correctly
+    call_args = mock_telemetry_manager.send_event.call_args[1]
+    assert "test_func-error" in call_args["event_type"]
+    assert "Test error" in call_args["properties"]["exception"]
+    assert "stack_trace" in call_args["properties"]
+
+
+def test_send_errors_nested_calls(mock_telemetry_manager):
+    """Test that send_errors only sends telemetry once for nested decorated functions"""
+
+    @send_errors
+    def inner_func():
+        raise ValueError("Inner error")
+
+    @send_errors
+    def outer_func():
+        return inner_func()
+
+    with pytest.raises(ValueError):
+        outer_func()
+
+    # Telemetry should be sent only once for the inner function
+    assert mock_telemetry_manager.send_event.call_count == 1
+    call_args = mock_telemetry_manager.send_event.call_args[1]
+    assert "inner_func-error" in call_args["event_type"]
+
+
+def test_send_errors_telemetry_disable():
+    """Test that send_errors doesn't attempt to send telemetry when disabled"""
+
+    with patch("axolotl.telemetry.errors.TelemetryManager") as mock_manager_class:
+        mock_manager = MagicMock()
+        mock_manager.enabled = False
+        mock_manager_class.get_instance.return_value = mock_manager
+
+        @send_errors
+        def test_func():
+            raise ValueError("Test error")
+
+        with pytest.raises(ValueError):
+            test_func()
+
+        mock_manager.send_event.assert_not_called()
+
+
+def test_error_handled_reset():
+    """Test that ERROR_HANDLED flag is properly reset"""
+    with patch("axolotl.telemetry.errors.TelemetryManager") as mock_manager_class:
+        # Create and configure the mock manager
+        mock_manager = MagicMock()
+        mock_manager.enabled = True
+        mock_manager_class.get_instance.return_value = mock_manager
+
+        from axolotl.telemetry.errors import ERROR_HANDLED
+
+        @send_errors
+        def test_func():
+            raise ValueError("Test error")
+
+        assert not ERROR_HANDLED
+
+        with pytest.raises(ValueError):
+            test_func()
+
+        from axolotl.telemetry.errors import ERROR_HANDLED
+
+        assert ERROR_HANDLED
+
+
+def test_module_path_resolution(mock_telemetry_manager):
+    """Test that the module path is correctly resolved for the event type"""
+    import inspect
+
+    current_module = inspect.getmodule(test_module_path_resolution).__name__
+
+    @send_errors
+    def test_func():
+        raise ValueError("Test error")
+
+    with pytest.raises(ValueError):
+        test_func()
+
+    assert mock_telemetry_manager.send_event.called
+    event_type = mock_telemetry_manager.send_event.call_args[1]["event_type"]
+
+    expected_event_type = f"{current_module}.test_func-error"
+    assert expected_event_type == event_type
diff --git a/tests/telemetry/test_manager.py b/tests/telemetry/test_manager.py
new file mode 100644
index 0000000000..67914b19e3
--- /dev/null
+++ b/tests/telemetry/test_manager.py
@@ -0,0 +1,272 @@
+"""Tests for TelemetryManager class and utilities"""
+
+# pylint: disable=redefined-outer-name,protected-access
+
+import logging
+import os
+from unittest.mock import patch
+
+import pytest
+import yaml
+
+from axolotl.telemetry.manager import TelemetryManager
+
+
+@pytest.fixture
+def mock_whitelist(tmp_path):
+    """Create a temporary whitelist file for testing"""
+    whitelist_content = {
+        "organizations": ["meta-llama", "mistralai"],
+    }
+    whitelist_file = tmp_path / "whitelist.yaml"
+    with open(whitelist_file, "w", encoding="utf-8") as f:
+        yaml.dump(whitelist_content, f)
+
+    return str(whitelist_file)
+
+
+@pytest.fixture
+def telemetry_manager_class():
+    """Reset the TelemetryManager singleton between tests"""
+    original_instance = TelemetryManager._instance
+    original_initialized = TelemetryManager._initialized
+    TelemetryManager._instance = None
+    TelemetryManager._initialized = False
+    yield TelemetryManager
+    TelemetryManager._instance = original_instance
+    TelemetryManager._initialized = original_initialized
+
+
+@pytest.fixture
+def manager(telemetry_manager_class, mock_whitelist):
+    """Create a TelemetryManager instance with mocked dependencies"""
+    with (
+        patch("posthog.capture"),
+        patch("posthog.flush"),
+        patch("time.sleep"),
+        patch("axolotl.telemetry.manager.WHITELIST_PATH", mock_whitelist),
+        patch.dict(os.environ, {"RANK": "0"}),
+    ):
+        manager = telemetry_manager_class()
+        # Manually enable for most tests
+        manager.enabled = True
+        return manager
+
+
+def test_singleton_instance(telemetry_manager_class):
+    """Test that TelemetryManager is a singleton"""
+    with (
+        patch("posthog.capture"),
+        patch("time.sleep"),
+        patch.dict(os.environ, {"RANK": "0"}),
+    ):
+        first = telemetry_manager_class()
+        second = telemetry_manager_class()
+        assert first is second
+        assert telemetry_manager_class.get_instance() is first
+
+
+class TestTelemetryOptOut:
+    """
+    Telemetry is opt-out: enabled by default, disabled by AXOLOTL_DO_NOT_TRACK
+    or DO_NOT_TRACK. Each env var is checked independently — setting either one
+    to a truthy value ("1" or "true") disables telemetry.
+
+    The parametrized table below is the source of truth for expected behavior.
+    """
+
+    # fmt: off
+    #                       AXOLOTL_DO_NOT_TRACK  DO_NOT_TRACK  expected
+    @pytest.mark.parametrize("axolotl_dnt, dnt, expected", [
+        # --- Neither var set: telemetry ON ---
+        (None,    None,    True),
+
+        # --- Only AXOLOTL_DO_NOT_TRACK set ---
+        ("0",     None,    True),    # explicit opt-in
+        ("false", None,    True),    # explicit opt-in
+        ("1",     None,    False),   # opt-out
+        ("true",  None,    False),   # opt-out
+        (" 1 ",   None,    False),   # whitespace-padded opt-out
+
+        # --- Only DO_NOT_TRACK set (was broken before fix) ---
+        (None,    "0",     True),    # explicit opt-in
+        (None,    "false", True),    # explicit opt-in
+        (None,    "1",     False),   # opt-out
+        (None,    "true",  False),   # opt-out
+
+        # --- Both set: either truthy → disabled ---
+        ("0",     "1",     False),   # DO_NOT_TRACK wins
+        ("1",     "0",     False),   # AXOLOTL_DO_NOT_TRACK wins
+        ("1",     "1",     False),   # both opt-out
+        ("0",     "0",     True),    # both opt-in
+    ])
+    # fmt: on
+    def test_do_not_track_env_vars(
+        self, telemetry_manager_class, axolotl_dnt, dnt, expected
+    ):
+        env = {"RANK": "0"}
+        if axolotl_dnt is not None:
+            env["AXOLOTL_DO_NOT_TRACK"] = axolotl_dnt
+        if dnt is not None:
+            env["DO_NOT_TRACK"] = dnt
+
+        with (
+            patch.dict(os.environ, env, clear=True),
+            patch("time.sleep"),
+            patch("logging.Logger.info"),
+        ):
+            manager = telemetry_manager_class()
+            assert manager.enabled is expected
+
+
+def test_telemetry_disabled_for_non_main_process(telemetry_manager_class):
+    """Test that telemetry is disabled for non-main processes"""
+    with (
+        patch.dict(os.environ, {"AXOLOTL_DO_NOT_TRACK": "0", "RANK": "1"}),
+        patch("time.sleep"),
+    ):
+        manager = telemetry_manager_class()
+        assert not manager.enabled
+
+
+def test_is_whitelisted(telemetry_manager_class, mock_whitelist):
+    """Test org whitelist functionality"""
+    with (
+        patch("axolotl.telemetry.manager.WHITELIST_PATH", mock_whitelist),
+        patch.dict(os.environ, {"AXOLOTL_DO_NOT_TRACK": "0"}),
+    ):
+        manager = telemetry_manager_class()
+
+        # Should match organizations from the mock whitelist
+        assert manager._is_whitelisted("meta-llama/llama-7b")
+        assert manager._is_whitelisted("mistralai/mistral-7b-instruct")
+        # Should not match
+        assert not manager._is_whitelisted("unknown/model")
+        # Should handle case insensitively
+        assert manager._is_whitelisted("META-LLAMA/Llama-7B")
+        # Should handle empty input
+        assert not manager._is_whitelisted("")
+
+
+def test_system_info_collection(manager):
+    """Test system information collection"""
+    system_info = manager._get_system_info()
+
+    # Check essential keys
+    assert "os" in system_info
+    assert "python_version" in system_info
+    assert "cpu_count" in system_info
+    assert "memory_total" in system_info
+    assert "accelerator_count" in system_info
+
+
+def test_send_event(telemetry_manager_class):
+    """Test basic event sending"""
+    with (
+        patch("posthog.capture") as mock_capture,
+        patch.dict(os.environ, {"AXOLOTL_DO_NOT_TRACK": "0"}),
+    ):
+        manager = telemetry_manager_class()
+
+        # Test with clean properties (no PII)
+        manager.send_event("test_event", {"key": "value"})
+        assert mock_capture.called
+        assert mock_capture.call_args[1]["event"] == "test_event"
+        assert mock_capture.call_args[1]["properties"] == {"key": "value"}
+        assert mock_capture.call_args[1]["distinct_id"] == manager.run_id
+
+        # Test with default properties (None)
+        mock_capture.reset_mock()
+        manager.send_event("simple_event")
+        assert mock_capture.called
+        assert mock_capture.call_args[1]["properties"] == {}
+
+
+def test_send_system_info(telemetry_manager_class):
+    """Test sending system info"""
+    with (
+        patch("posthog.capture") as mock_capture,
+        patch.dict(os.environ, {"AXOLOTL_DO_NOT_TRACK": "0"}),
+    ):
+        manager = telemetry_manager_class()
+        manager.send_system_info()
+        assert mock_capture.called
+        assert mock_capture.call_args[1]["event"] == "system-info"
+        assert mock_capture.call_args[1]["properties"] == manager.system_info
+
+
+def test_redacted_properties(telemetry_manager_class):
+    """Test path redaction in send_event method"""
+    with (
+        patch("posthog.capture") as mock_capture,
+        patch.dict(os.environ, {"AXOLOTL_DO_NOT_TRACK": "0"}),
+    ):
+        manager = telemetry_manager_class()
+        # Test with properties containing various paths and non-paths
+        test_properties = {
+            "filepath": "/home/user/sensitive/data.txt",
+            "windows_path": "C:\\Users\\name\\Documents\\project\\file.py",
+            "output_dir": "/var/lib/data",
+            "path_to_model": "models/llama/7b",
+            "message": "Training started",  # Should not be redacted
+            "metrics": {"loss": 0.5, "accuracy": 0.95},  # Should not be redacted
+            "base_model": "models/local_model",
+            "nested": {
+                "model_path": "/models/my_model",
+                "root_dir": "/home/user/projects",
+                "stats": {"steps": 1000, "epochs": 3},  # Should not be redacted
+            },
+        }
+
+        manager.send_event("test_event", test_properties)
+
+        # Verify the call was made
+        assert mock_capture.called
+
+        # Get the sanitized properties that were sent
+        sanitized = mock_capture.call_args[1]["properties"]
+
+        # Check that path-like and base_model keys were redacted
+        assert sanitized["filepath"] == "[REDACTED]"
+        assert sanitized["windows_path"] == "[REDACTED]"
+        assert sanitized["path_to_model"] == "[REDACTED]"
+        assert sanitized["base_model"] == "[REDACTED]"
+
+        # Check that non-path values were preserved
+        assert sanitized["message"] == "Training started"
+        assert sanitized["metrics"] == {"loss": 0.5, "accuracy": 0.95}
+
+        # Check nested structure handling
+        assert sanitized["nested"]["model_path"] == "[REDACTED]"
+        assert sanitized["nested"]["root_dir"] == "[REDACTED]"
+        assert sanitized["nested"]["stats"] == {"steps": 1000, "epochs": 3}
+
+
+def test_disable_telemetry(manager):
+    """Test that disabled telemetry doesn't send events"""
+    with patch("posthog.capture") as mock_capture:
+        manager.enabled = False
+        manager.send_event("test_event")
+        assert not mock_capture.called
+
+
+def test_exception_handling_during_send(manager, caplog):
+    """Test that exceptions in PostHog are handled gracefully"""
+    # axolotl loggers set propagate=False (logging_config.py), so caplog's root
+    # handler never sees them; attach it to the named logger directly.
+    logger = logging.getLogger("axolotl.telemetry.manager")
+    logger.addHandler(caplog.handler)
+    try:
+        with patch("posthog.capture", side_effect=Exception("Test error")):
+            manager.send_event("test_event")
+    finally:
+        logger.removeHandler(caplog.handler)
+
+    assert "Failed to send telemetry event: Test error" in caplog.text
+
+
+def test_shutdown(manager):
+    """Test shutdown behavior"""
+    with patch("posthog.shutdown") as mock_shutdown:
+        manager.shutdown()
+        assert mock_shutdown.called
diff --git a/tests/telemetry/test_runtime_metrics.py b/tests/telemetry/test_runtime_metrics.py
new file mode 100644
index 0000000000..faa62d0749
--- /dev/null
+++ b/tests/telemetry/test_runtime_metrics.py
@@ -0,0 +1,357 @@
+"""Tests for runtime metrics telemetry module"""
+
+# pylint: disable=redefined-outer-name
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from axolotl.telemetry.runtime_metrics import RuntimeMetrics, RuntimeMetricsTracker
+
+
+@pytest.fixture
+def mock_time():
+    """Mock time.time() to have predictable values in tests"""
+    with patch("time.time") as mock_time:
+        # Start with time 1000.0 and increment by 10 seconds on each call
+        times = [1000.0 + i * 10 for i in range(10)]
+        mock_time.side_effect = times
+        yield mock_time
+
+
+@pytest.fixture
+def mock_telemetry_manager():
+    """Create a mock TelemetryManager"""
+    with patch(
+        "axolotl.telemetry.runtime_metrics.TelemetryManager"
+    ) as mock_manager_class:
+        mock_manager = MagicMock()
+        mock_manager.enabled = True
+        mock_manager_class.get_instance.return_value = mock_manager
+        yield mock_manager
+
+
+@pytest.fixture
+def mock_psutil():
+    """Mock psutil for memory information"""
+    with patch("axolotl.telemetry.runtime_metrics.psutil") as mock_psutil:
+        mock_process = MagicMock()
+        mock_memory_info = MagicMock()
+        # Set initial memory to 1GB
+        mock_memory_info.rss = 1024 * 1024 * 1024
+        mock_process.memory_info.return_value = mock_memory_info
+        mock_psutil.Process.return_value = mock_process
+        yield mock_psutil
+
+
+@pytest.fixture
+def mock_torch():
+    """Mock torch.cuda functions"""
+    with patch("axolotl.telemetry.runtime_metrics.torch") as mock_torch:
+        mock_torch.cuda.is_available.return_value = True
+        mock_torch.cuda.device_count.return_value = 2
+
+        # Mock memory allocated per device (1GB for device 0, 2GB for device 1)
+        mock_torch.cuda.memory_allocated.side_effect = lambda device: (
+            (device + 1) * 1024 * 1024 * 1024
+        )
+
+        yield mock_torch
+
+
+class TestRuntimeMetrics:
+    """Tests for RuntimeMetrics class."""
+
+    def test_initialization(self):
+        """Test RuntimeMetrics initialization."""
+        metrics = RuntimeMetrics(start_time=1000.0)
+
+        assert metrics.start_time == 1000.0
+        assert metrics.epoch_start_times == {}
+        assert metrics.epoch_end_times == {}
+        assert metrics.peak_gpu_memory == {}
+        assert metrics.total_steps == 0
+        assert metrics.current_epoch == 0
+        assert metrics.current_step == 0
+        assert metrics.peak_cpu_memory == 0
+
+    def test_elapsed_time(self, mock_time):
+        """Test elapsed_time property."""
+        metrics = RuntimeMetrics(start_time=1000.0)
+
+        # Mock time.time() to return 1050.0
+        mock_time.side_effect = [1050.0]
+
+        assert metrics.elapsed_time == 50.0
+
+    def test_epoch_time(self):
+        """Test epoch_time method."""
+        metrics = RuntimeMetrics(start_time=1000.0)
+
+        # No epoch data
+        assert metrics.epoch_time(0) is None
+
+        # Add epoch start but no end
+        metrics.epoch_start_times[0] = 1000.0
+        assert metrics.epoch_time(0) is None
+
+        # Add epoch end
+        metrics.epoch_end_times[0] = 1060.0
+        assert metrics.epoch_time(0) == 60.0
+
+    def test_average_epoch_time(self):
+        """Test average_epoch_time method."""
+        metrics = RuntimeMetrics(start_time=1000.0)
+
+        # No completed epochs
+        assert metrics.average_epoch_time() is None
+
+        # Add one completed epoch
+        metrics.epoch_start_times[0] = 1000.0
+        metrics.epoch_end_times[0] = 1060.0
+        assert metrics.average_epoch_time() == 60.0
+
+        # Add second completed epoch
+        metrics.epoch_start_times[1] = 1060.0
+        metrics.epoch_end_times[1] = 1140.0  # 80 seconds
+        assert metrics.average_epoch_time() == 70.0  # Average of 60 and 80
+
+        # Add incomplete epoch (should not affect average)
+        metrics.epoch_start_times[2] = 1140.0
+        assert metrics.average_epoch_time() == 70.0
+
+    def test_steps_per_second(self, mock_time):
+        """Test steps_per_second method."""
+        metrics = RuntimeMetrics(start_time=1000.0)
+
+        # No steps - first call to time.time()
+        mock_time.side_effect = None
+        mock_time.return_value = 1050.0
+        assert metrics.steps_per_second() is None
+
+        # Add steps - second call to time.time()
+        metrics.total_steps = 100
+        mock_time.return_value = 1050.0  # Keep same time for consistent result
+        assert metrics.steps_per_second() == 2.0  # 100 steps / 50 seconds
+
+    def test_to_dict_basic(self, mock_time):
+        """Test to_dict method with basic metrics."""
+        metrics = RuntimeMetrics(start_time=1000.0)
+        metrics.total_steps = 100
+        metrics.peak_cpu_memory = 2 * 1024 * 1024 * 1024  # 2GB
+
+        # Mock elapsed_time
+        mock_time.side_effect = None
+        mock_time.return_value = 1050.0
+
+        result = metrics.to_dict()
+
+        assert result["total_time_seconds"] == 50.0
+        assert result["total_steps"] == 100
+        assert result["steps_per_second"] == 2.0
+        assert result["epochs_completed"] == 0
+        assert result["peak_cpu_memory_bytes"] == 2 * 1024 * 1024 * 1024
+        assert "epoch_times" not in result
+        assert "gpu_memory" not in result
+
+    def test_to_dict_with_epochs(self, mock_time):
+        """Test to_dict method with epoch data."""
+        metrics = RuntimeMetrics(start_time=1000.0)
+        metrics.total_steps = 100
+
+        # Add epoch data
+        metrics.epoch_start_times[0] = 1000.0
+        metrics.epoch_end_times[0] = 1060.0
+        metrics.epoch_start_times[1] = 1060.0
+        metrics.epoch_end_times[1] = 1140.0
+
+        # Mock elapsed_time
+        mock_time.side_effect = None
+        mock_time.return_value = 1150.0
+
+        result = metrics.to_dict()
+
+        assert "epoch_times" in result
+        assert result["epoch_times"]["epoch_0_seconds"] == 60.0
+        assert result["epoch_times"]["epoch_1_seconds"] == 80.0
+        assert result["average_epoch_time_seconds"] == 70.0
+
+    def test_to_dict_with_gpu_memory(self, mock_time):
+        """Test to_dict method with GPU memory data."""
+        metrics = RuntimeMetrics(start_time=1000.0)
+        metrics.peak_gpu_memory = {
+            0: 1 * 1024 * 1024 * 1024,  # 1GB
+            1: 2 * 1024 * 1024 * 1024,  # 2GB
+        }
+
+        # Mock elapsed_time
+        mock_time.side_effect = [1050.0]
+
+        result = metrics.to_dict()
+
+        assert "gpu_memory" in result
+        assert result["gpu_memory"]["gpu_0_peak_memory_bytes"] == 1 * 1024 * 1024 * 1024
+        assert result["gpu_memory"]["gpu_1_peak_memory_bytes"] == 2 * 1024 * 1024 * 1024
+
+
+class TestRuntimeMetricsTracker:
+    """Tests for RuntimeMetricsTracker class."""
+
+    # pylint: disable=unused-argument
+    def test_initialization(self, mock_time, mock_telemetry_manager):
+        """Test RuntimeMetricsTracker initialization."""
+        tracker = RuntimeMetricsTracker()
+
+        assert isinstance(tracker.metrics, RuntimeMetrics)
+        assert tracker.metrics.start_time == 1000.0  # First value from mock_time
+
+    # pylint: disable=unused-argument
+    def test_start_epoch(
+        self, mock_time, mock_psutil, mock_torch, mock_telemetry_manager
+    ):
+        """Test start_epoch method."""
+        tracker = RuntimeMetricsTracker()
+
+        # Reset mock_time to control next value
+        mock_time.side_effect = [1010.0]
+
+        tracker.start_epoch(0)
+
+        assert tracker.metrics.current_epoch == 0
+        assert tracker.metrics.epoch_start_times[0] == 1010.0
+
+        # Verify memory metrics were updated
+        assert tracker.metrics.peak_cpu_memory == 1 * 1024 * 1024 * 1024
+        assert 0 in tracker.metrics.peak_gpu_memory
+        assert 1 in tracker.metrics.peak_gpu_memory
+
+    # pylint: disable=unused-argument
+    def test_end_epoch(self, mock_time, mock_telemetry_manager):
+        """Test end_epoch method."""
+        tracker = RuntimeMetricsTracker()
+
+        # Start epoch 0
+        mock_time.side_effect = [1010.0]
+        tracker.start_epoch(0)
+
+        # End epoch 0
+        mock_time.side_effect = [1060.0]
+        tracker.end_epoch(0)
+
+        assert 0 in tracker.metrics.epoch_end_times
+        assert tracker.metrics.epoch_end_times[0] == 1060.0
+
+    # pylint: disable=unused-argument
+    def test_update_step(
+        self, mock_time, mock_psutil, mock_torch, mock_telemetry_manager
+    ):
+        """Test update_step method."""
+        tracker = RuntimeMetricsTracker()
+
+        # Update step to a non-multiple of 100
+        tracker.update_step(42)
+
+        assert tracker.metrics.current_step == 42
+        assert tracker.metrics.total_steps == 1
+
+        # Memory metrics should not be updated for non-multiple of 100
+        assert tracker.metrics.peak_cpu_memory == 0
+
+        # Update step to a multiple of 100
+        tracker.update_step(100)
+
+        assert tracker.metrics.current_step == 100
+        assert tracker.metrics.total_steps == 2
+
+        # Memory metrics should be updated for multiple of 100
+        assert tracker.metrics.peak_cpu_memory == 1 * 1024 * 1024 * 1024
+
+    # pylint: disable=unused-argument
+    def test_update_memory_metrics(
+        self, mock_psutil, mock_torch, mock_telemetry_manager
+    ):
+        """Test update_memory_metrics method."""
+        tracker = RuntimeMetricsTracker()
+
+        # Initial memory state
+        assert tracker.metrics.peak_cpu_memory == 0
+        assert tracker.metrics.peak_gpu_memory == {}
+
+        # Update memory metrics
+        tracker.update_memory_metrics()
+
+        # Verify CPU memory
+        assert tracker.metrics.peak_cpu_memory == 1 * 1024 * 1024 * 1024
+
+        # Verify GPU memory
+        assert tracker.metrics.peak_gpu_memory[0] == 1 * 1024 * 1024 * 1024
+        assert tracker.metrics.peak_gpu_memory[1] == 2 * 1024 * 1024 * 1024
+
+        # Change mocked memory values to be lower
+        mock_process = mock_psutil.Process.return_value
+        mock_memory_info = mock_process.memory_info.return_value
+        mock_memory_info.rss = 0.5 * 1024 * 1024 * 1024  # 0.5GB
+
+        mock_torch.cuda.memory_allocated.side_effect = lambda device: (
+            (device + 0.5) * 1024 * 1024 * 1024
+        )
+
+        # Update memory metrics again
+        tracker.update_memory_metrics()
+
+        # Peak values should not decrease
+        assert tracker.metrics.peak_cpu_memory == 1 * 1024 * 1024 * 1024
+        assert tracker.metrics.peak_gpu_memory[0] == 1 * 1024 * 1024 * 1024
+        assert tracker.metrics.peak_gpu_memory[1] == 2 * 1024 * 1024 * 1024
+
+        # Change mocked memory values to be higher
+        mock_memory_info.rss = 2 * 1024 * 1024 * 1024  # 2GB
+
+        mock_torch.cuda.memory_allocated.side_effect = lambda device: (
+            (device + 2) * 1024 * 1024 * 1024
+        )
+
+        # Update memory metrics again
+        tracker.update_memory_metrics()
+
+        # Peak values should increase
+        assert tracker.metrics.peak_cpu_memory == 2 * 1024 * 1024 * 1024
+        assert tracker.metrics.peak_gpu_memory[0] == 2 * 1024 * 1024 * 1024
+        assert tracker.metrics.peak_gpu_memory[1] == 3 * 1024 * 1024 * 1024
+
+    # pylint: disable=unused-argument
+    def test_get_memory_metrics(self, mock_psutil, mock_torch, mock_telemetry_manager):
+        """Test get_memory_metrics method."""
+        tracker = RuntimeMetricsTracker()
+
+        # Set peak memory values
+        tracker.metrics.peak_cpu_memory = 2 * 1024 * 1024 * 1024
+        tracker.metrics.peak_gpu_memory = {
+            0: 3 * 1024 * 1024 * 1024,
+            1: 4 * 1024 * 1024 * 1024,
+        }
+
+        # Get memory metrics
+        memory_metrics = tracker.get_memory_metrics()
+
+        # Verify CPU memory
+        assert (
+            memory_metrics["cpu_memory_bytes"] == 1 * 1024 * 1024 * 1024
+        )  # Current value from mock
+        assert (
+            memory_metrics["peak_cpu_memory_bytes"] == 2 * 1024 * 1024 * 1024
+        )  # Peak value we set
+
+        # Verify GPU memory
+        assert (
+            memory_metrics["gpu_0_memory_bytes"] == 1 * 1024 * 1024 * 1024
+        )  # Current value from mock
+        assert (
+            memory_metrics["gpu_0_peak_memory_bytes"] == 3 * 1024 * 1024 * 1024
+        )  # Peak value we set
+        assert (
+            memory_metrics["gpu_1_memory_bytes"] == 2 * 1024 * 1024 * 1024
+        )  # Current value from mock
+        assert (
+            memory_metrics["gpu_1_peak_memory_bytes"] == 4 * 1024 * 1024 * 1024
+        )  # Peak value we set
diff --git a/tests/test_attn_implementation.py b/tests/test_attn_implementation.py
new file mode 100644
index 0000000000..dc29f1dd60
--- /dev/null
+++ b/tests/test_attn_implementation.py
@@ -0,0 +1,447 @@
+"""Tests for attn_implementation: normalization, canonical-value acceptance,
+capability flags, backend registration, and downstream validators.
+"""
+
+import logging
+import subprocess
+import sys
+from contextlib import contextmanager
+from functools import lru_cache
+
+import pytest
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.schemas.config import AxolotlInputConfig
+from axolotl.utils.schemas.enums import (
+    ATTN_IMPLS_SUPPORTING_PACKING,
+    ATTN_IMPLS_USING_FLASH_LIB,
+    ATTN_IMPLS_WITHOUT_DTYPE_CAST,
+    CANONICAL_ATTN_IMPLS,
+)
+
+
+@contextmanager
+def _capture_axolotl_warnings(caplog):
+    """Capture WARNINGs from `axolotl.*` loggers via caplog.
+
+    `axolotl.cli` calls `configure_logging()` at import time, which sets
+    `propagate=False` on the `axolotl` logger so records do not reach the root
+    logger that pytest's `caplog` hooks. This helper temporarily re-enables
+    propagation for the duration of the block.
+    """
+    ax_logger = logging.getLogger("axolotl")
+    old_propagate = ax_logger.propagate
+    ax_logger.propagate = True
+    try:
+        with caplog.at_level(logging.WARNING, logger="axolotl"):
+            yield
+    finally:
+        ax_logger.propagate = old_propagate
+
+
+@lru_cache(maxsize=1)
+def _xformers_available():
+    try:
+        result = subprocess.run(  # noqa: S603
+            [
+                sys.executable,
+                "-c",
+                (
+                    "import warnings; "
+                    "warnings.filterwarnings('ignore', category=DeprecationWarning); "
+                    "import xformers.ops"
+                ),
+            ],
+            check=False,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            timeout=10,
+        )
+        return result.returncode == 0
+    except Exception:  # pylint: disable=broad-except
+        return False
+
+
+class TestCapabilityTables:
+    """Backend capability classification via frozensets and computed_field properties."""
+
+    @pytest.mark.parametrize(
+        "impl",
+        [
+            "flash_attention_2",
+            "flash_attention_3",
+            "flex_attention",
+            "xformers",
+            "sage",
+        ],
+    )
+    def test_supports_packing(self, impl):
+        assert impl in ATTN_IMPLS_SUPPORTING_PACKING
+
+    @pytest.mark.parametrize("impl", ["eager", "sdpa", "fp8"])
+    def test_does_not_support_packing(self, impl):
+        assert impl not in ATTN_IMPLS_SUPPORTING_PACKING
+
+    @pytest.mark.parametrize("impl", ["flash_attention_2", "flash_attention_3"])
+    def test_uses_flash_lib(self, impl):
+        assert impl in ATTN_IMPLS_USING_FLASH_LIB
+
+    @pytest.mark.parametrize(
+        "impl", ["eager", "sdpa", "xformers", "flex_attention", "sage", "fp8"]
+    )
+    def test_does_not_use_flash_lib(self, impl):
+        assert impl not in ATTN_IMPLS_USING_FLASH_LIB
+
+    @pytest.mark.parametrize("impl", ["eager", "sdpa"])
+    def test_no_dtype_cast(self, impl):
+        assert impl in ATTN_IMPLS_WITHOUT_DTYPE_CAST
+
+    @pytest.mark.parametrize(
+        "impl",
+        [
+            "flash_attention_2",
+            "flash_attention_3",
+            "flex_attention",
+            "xformers",
+            "sage",
+            "fp8",
+        ],
+    )
+    def test_needs_dtype_cast(self, impl):
+        assert impl not in ATTN_IMPLS_WITHOUT_DTYPE_CAST
+
+    def test_known_hub_kernels_classified(self):
+        assert "kernels-community/flash-attn3" in ATTN_IMPLS_SUPPORTING_PACKING
+        assert "kernels-community/flash-attn3" in ATTN_IMPLS_USING_FLASH_LIB
+        assert "kernels-community/sage-attention" in ATTN_IMPLS_SUPPORTING_PACKING
+
+    def test_computed_flags_readable_on_validated_cfg(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(attn_implementation="sdpa")
+        validated = validate_config(cfg)
+        assert validated.attn_implementation == "sdpa"
+        assert validated.attn_supports_packing is False
+        assert validated.attn_uses_flash_lib is False
+        assert validated.attn_needs_dtype_cast is False
+
+    @pytest.mark.parametrize(
+        "impl,expected",
+        [
+            ("flash_attention_2", True),
+            ("flex_attention", True),
+            (
+                "sdpa",
+                True,
+            ),  # not varlen, but isolates via block-diagonal from position_ids
+            ("eager", True),
+            ("fp8", False),
+        ],
+    )
+    def test_decontaminates_packing(self, min_base_cfg, impl, expected):
+        validated = validate_config(
+            min_base_cfg | DictDefault(attn_implementation=impl)
+        )
+        assert validated.attn_decontaminates_packing is expected
+
+    def test_computed_flags_not_overridable_from_yaml(self, min_base_cfg):
+        """YAML attempts to override a computed field must not win."""
+        cfg = min_base_cfg | DictDefault(
+            attn_implementation="eager", attn_uses_flash_lib=True
+        )
+        validated = validate_config(cfg)
+        # The computed field reflects the backend, not the YAML input.
+        assert validated.attn_uses_flash_lib is False
+
+
+class TestBackendRegistration:
+    """Axolotl-owned backends register under their canonical names in HF's registries."""
+
+    @pytest.mark.skipif(not _xformers_available(), reason="xformers not available")
+    def test_register_xformers(self):
+        from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        from axolotl.monkeypatch.attention import register_xformers_attn
+
+        register_xformers_attn()
+
+        assert "xformers" in ALL_ATTENTION_FUNCTIONS
+        assert "xformers" in ALL_MASK_ATTENTION_FUNCTIONS
+        assert (
+            ALL_MASK_ATTENTION_FUNCTIONS["xformers"]
+            == ALL_MASK_ATTENTION_FUNCTIONS["flash_attention_2"]
+        )
+
+    def test_register_sage(self):
+        from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        from axolotl.monkeypatch.attention import register_sage_attn
+
+        register_sage_attn()
+
+        assert "sage" in ALL_ATTENTION_FUNCTIONS
+        assert "sage" in ALL_MASK_ATTENTION_FUNCTIONS
+        assert (
+            ALL_MASK_ATTENTION_FUNCTIONS["sage"]
+            == ALL_MASK_ATTENTION_FUNCTIONS["flash_attention_2"]
+        )
+
+    @pytest.mark.skipif(not _xformers_available(), reason="xformers not available")
+    def test_xformers_does_not_overwrite_fa2(self):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        original_fa2 = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
+
+        from axolotl.monkeypatch.attention import register_xformers_attn
+
+        register_xformers_attn()
+
+        assert ALL_ATTENTION_FUNCTIONS["flash_attention_2"] is original_fa2
+
+    def test_sage_does_not_overwrite_fa2(self):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        original_fa2 = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
+
+        from axolotl.monkeypatch.attention import register_sage_attn
+
+        register_sage_attn()
+
+        assert ALL_ATTENTION_FUNCTIONS["flash_attention_2"] is original_fa2
+
+
+class TestLegacyFlagDeprecation:
+    """Legacy boolean flags (flash_attention, sdp_attention, ...) map to a
+    canonical attn_implementation value, are stripped from the validated
+    config, and cannot be combined with an explicit canonical value.
+    """
+
+    @staticmethod
+    def _normalize(data):
+        return AxolotlInputConfig.normalize_attn_implementation(data)
+
+    @pytest.mark.parametrize(
+        "flag,expected",
+        [
+            ("flash_attention", "flash_attention_2"),
+            ("sdp_attention", "sdpa"),
+            ("xformers_attention", "xformers"),
+            ("flex_attention", "flex_attention"),
+            ("sage_attention", "sage"),
+            ("eager_attention", "eager"),
+        ],
+    )
+    def test_legacy_flag_maps_to_canonical(self, flag, expected):
+        result = self._normalize({flag: True})
+        assert result["attn_implementation"] == expected
+
+    def test_legacy_flags_are_stripped_after_mapping(self):
+        result = self._normalize({"flash_attention": True})
+        for flag in [
+            "flash_attention",
+            "sdp_attention",
+            "xformers_attention",
+            "flex_attention",
+            "sage_attention",
+            "eager_attention",
+        ]:
+            assert flag not in result
+
+    def test_sage_plus_flash_priority_is_sage(self):
+        result = self._normalize({"sage_attention": True, "flash_attention": True})
+        assert result["attn_implementation"] == "sage"
+
+    def test_canonical_plus_legacy_flag_raises(self):
+        with pytest.raises(ValueError, match="cannot be combined with legacy"):
+            self._normalize(
+                {"attn_implementation": "flash_attention_2", "flash_attention": True}
+            )
+
+    def test_canonical_plus_unrelated_legacy_flag_raises(self):
+        with pytest.raises(ValueError, match="cannot be combined with legacy"):
+            self._normalize(
+                {"attn_implementation": "xformers", "flash_attention": True}
+            )
+
+    def test_legacy_flag_stripped_on_validated_cfg(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(flash_attention=True)
+        validated = validate_config(cfg)
+        assert validated.attn_implementation == "flash_attention_2"
+        # Legacy flag must not survive to the validated DictDefault
+        # (normalizer pops it, model_dump excludes Nones).
+        assert "flash_attention" not in dict(validated)
+
+    def test_canonical_plus_legacy_rejected_on_full_validation(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            attn_implementation="flash_attention_2", flash_attention=True
+        )
+        with pytest.raises(ValueError, match="cannot be combined with legacy"):
+            validate_config(cfg)
+
+
+class TestCanonicalValueAcceptance:
+    """`attn_implementation` accepts canonical names and `org/name` hub-kernel
+    paths. Short-form aliases (`flash`, `flex`, `sdp`) and unknown bare names
+    are rejected. Absent input is a noop.
+    """
+
+    @staticmethod
+    def _normalize(data):
+        return AxolotlInputConfig.normalize_attn_implementation(data)
+
+    def test_canonical_value_is_passthrough(self):
+        data = {"attn_implementation": "flash_attention_2"}
+        result = self._normalize(data)
+        assert result["attn_implementation"] == "flash_attention_2"
+
+    def test_hub_kernel_is_passthrough(self):
+        data = {"attn_implementation": "kernels-community/flash-attn3"}
+        result = self._normalize(data)
+        assert result["attn_implementation"] == "kernels-community/flash-attn3"
+
+    def test_no_attention_set_is_noop(self):
+        result = self._normalize({"some_other_config": True})
+        assert result.get("attn_implementation") is None
+
+    def test_field_validator_accepts_all_canonical(self):
+        for impl in CANONICAL_ATTN_IMPLS:
+            assert AxolotlInputConfig.validate_attn_implementation(impl) == impl
+
+    def test_field_validator_accepts_hub_kernels(self):
+        for impl in (
+            "kernels-community/flash-attn3",
+            "kernels-community/sage-attention",
+            "someorg/custom-kernel",
+        ):
+            assert AxolotlInputConfig.validate_attn_implementation(impl) == impl
+
+    def test_field_validator_accepts_none(self):
+        assert AxolotlInputConfig.validate_attn_implementation(None) is None
+
+    @pytest.mark.parametrize("alias", ["flash", "flex", "sdp"])
+    def test_short_form_alias_rejected(self, alias):
+        with pytest.raises(ValueError, match="is not accepted"):
+            AxolotlInputConfig.validate_attn_implementation(alias)
+
+    def test_unknown_bare_name_rejected(self):
+        with pytest.raises(ValueError, match="not a recognized backend"):
+            AxolotlInputConfig.validate_attn_implementation("not_a_real_backend")
+
+    def test_canonical_value_passes_through_full_validation(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(attn_implementation="flash_attention_3")
+        validated = validate_config(cfg)
+        assert validated.attn_implementation == "flash_attention_3"
+        assert validated.attn_uses_flash_lib is True
+        assert validated.attn_supports_packing is True
+
+    def test_hub_kernel_passes_through_full_validation(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            attn_implementation="kernels-community/flash-attn3"
+        )
+        validated = validate_config(cfg)
+        assert validated.attn_implementation == "kernels-community/flash-attn3"
+        assert validated.attn_uses_flash_lib is True
+        assert validated.attn_supports_packing is True
+
+    def test_short_form_alias_rejected_on_full_validation(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(attn_implementation="flash")
+        with pytest.raises(ValueError, match="is not accepted"):
+            validate_config(cfg)
+
+
+class TestGemma4HybridMode:
+    """`gemma4_hybrid_attn_impl` pins `attn_implementation` to `flash_attention_2`."""
+
+    @staticmethod
+    def _normalize(data):
+        return AxolotlInputConfig.normalize_attn_implementation(data)
+
+    def test_defaults_to_flash_attention_2(self):
+        result = self._normalize({"gemma4_hybrid_attn_impl": True})
+        assert result["attn_implementation"] == "flash_attention_2"
+
+    def test_explicit_fa2_passes(self):
+        result = self._normalize(
+            {
+                "gemma4_hybrid_attn_impl": True,
+                "attn_implementation": "flash_attention_2",
+            }
+        )
+        assert result["attn_implementation"] == "flash_attention_2"
+
+    def test_non_fa2_raises(self):
+        with pytest.raises(
+            ValueError, match="requires attn_implementation=flash_attention_2"
+        ):
+            self._normalize(
+                {"gemma4_hybrid_attn_impl": True, "attn_implementation": "sdpa"}
+            )
+
+
+class TestSamplePackingValidation:
+    """`sample_packing` warns only for backends that can't isolate documents.
+
+    sdpa/eager decontaminate via the dropped-mask block-diagonal (sdpa additionally
+    through cu_seqlens varlen when available), so they must not warn.
+    """
+
+    def test_eager_does_not_warn(self, min_base_cfg, caplog):
+        cfg = min_base_cfg | DictDefault(
+            attn_implementation="eager", sample_packing=True
+        )
+        with _capture_axolotl_warnings(caplog):
+            validate_config(cfg)
+        assert not any(
+            "does not handle cross-sample decontamination" in r.getMessage()
+            for r in caplog.records
+        )
+
+    def test_sdpa_does_not_warn(self, min_base_cfg, caplog):
+        cfg = min_base_cfg | DictDefault(
+            attn_implementation="sdpa", sample_packing=True
+        )
+        with _capture_axolotl_warnings(caplog):
+            validate_config(cfg)
+        assert not any(
+            "does not handle cross-sample decontamination" in r.getMessage()
+            for r in caplog.records
+        )
+
+    def test_flash_attention_2_does_not_warn(self, min_base_cfg, caplog):
+        cfg = min_base_cfg | DictDefault(
+            attn_implementation="flash_attention_2", sample_packing=True
+        )
+        with _capture_axolotl_warnings(caplog):
+            validate_config(cfg)
+        assert not any(
+            "does not handle cross-sample decontamination" in r.getMessage()
+            for r in caplog.records
+        )
+
+    def test_non_decontaminating_backend_warns(self, min_base_cfg, caplog):
+        cfg = min_base_cfg | DictDefault(attn_implementation="fp8", sample_packing=True)
+        with _capture_axolotl_warnings(caplog):
+            validate_config(cfg)
+        assert any(
+            "does not handle cross-sample decontamination" in r.getMessage()
+            for r in caplog.records
+        )
+
+
+class TestScalingSoftmaxValidation:
+    """`scaling_softmax` is only implemented under flex_attention."""
+
+    def test_non_flex_raises(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            attn_implementation="flash_attention_2", scaling_softmax=True
+        )
+        with pytest.raises(ValueError, match="scaling_softmax requires flex"):
+            validate_config(cfg)
+
+    def test_flex_passes(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            attn_implementation="flex_attention", scaling_softmax=True
+        )
+        validated = validate_config(cfg)
+        assert validated.attn_implementation == "flex_attention"
diff --git a/tests/test_calculate_total_num_steps.py b/tests/test_calculate_total_num_steps.py
new file mode 100644
index 0000000000..454bb0f82d
--- /dev/null
+++ b/tests/test_calculate_total_num_steps.py
@@ -0,0 +1,118 @@
+"""Tests for token counting in calculate_total_num_steps"""
+
+import math
+
+import pytest
+from datasets import Dataset, Features, LargeList, Sequence, Value, load_from_disk
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.trainer import calculate_total_num_steps
+
+NUM_ROWS = 3000
+
+FEATURES = Features(
+    {
+        "input_ids": Sequence(Value("int64")),
+        "labels": Sequence(Value("int64")),
+    }
+)
+
+
+def build_columns(num_rows):
+    input_ids = []
+    labels = []
+    for i in range(num_rows):
+        seq_len = 5 + (i % 50)
+        input_ids.append(list(range(seq_len)))
+        labels.append([-100 if j < seq_len // 2 else j for j in range(seq_len)])
+    return {"input_ids": input_ids, "labels": labels}
+
+
+@pytest.fixture(name="columns")
+def fixture_columns():
+    return build_columns(NUM_ROWS)
+
+
+@pytest.fixture(name="expected_counts")
+def fixture_expected_counts(columns):
+    total_tokens = sum(len(row) for row in columns["input_ids"])
+    supervised_tokens = sum(
+        1 for row in columns["labels"] for token in row if token != -100
+    )
+    return total_tokens, supervised_tokens
+
+
+def make_cfg():
+    return DictDefault(
+        {
+            "num_epochs": 1,
+            "batch_size": 2,
+            "micro_batch_size": 1,
+        }
+    )
+
+
+def assert_counts(dataset, expected_counts):
+    cfg = make_cfg()
+    total_num_steps = calculate_total_num_steps(cfg, dataset)
+    expected_tokens, expected_supervised = expected_counts
+    assert cfg.total_num_tokens == expected_tokens
+    assert cfg.total_supervised_tokens == expected_supervised
+    assert total_num_steps == math.ceil(len(dataset) / cfg.batch_size)
+
+
+class TestCalculateTotalNumSteps:
+    """
+    Test token counting against hand-computed ground truth across Arrow layouts
+    """
+
+    def test_in_memory_single_chunk(self, columns, expected_counts):
+        dataset = Dataset.from_dict(columns, features=FEATURES)
+        # a single chunk larger than the to_batches chunksize is the layout
+        # where ListArray.values (vs .flatten()) silently over-counts
+        assert dataset.data.column("labels").num_chunks == 1
+        assert len(dataset) > 1024
+        assert_counts(dataset, expected_counts)
+
+    def test_disk_round_trip(self, columns, expected_counts, tmp_path):
+        dataset = Dataset.from_dict(columns, features=FEATURES)
+        dataset.save_to_disk(str(tmp_path / "ds"))
+        dataset = load_from_disk(str(tmp_path / "ds"))
+        assert_counts(dataset, expected_counts)
+
+    def test_large_list(self, columns, expected_counts):
+        dataset = Dataset.from_dict(columns, features=FEATURES).cast(
+            Features(
+                {
+                    "input_ids": LargeList(Value("int64")),
+                    "labels": LargeList(Value("int64")),
+                }
+            )
+        )
+        assert_counts(dataset, expected_counts)
+
+    def test_length_column_takes_precedence(self, columns, expected_counts):
+        columns = dict(columns)
+        # offset lengths by one so the result proves which branch ran
+        columns["length"] = [len(row) + 1 for row in columns["input_ids"]]
+        dataset = Dataset.from_dict(columns, features=None)
+        cfg = make_cfg()
+        calculate_total_num_steps(cfg, dataset)
+        expected_tokens, expected_supervised = expected_counts
+        assert cfg.total_num_tokens == expected_tokens + NUM_ROWS
+        assert cfg.total_supervised_tokens == expected_supervised
+
+    def test_empty_dataset(self):
+        dataset = Dataset.from_dict({"input_ids": [], "labels": []}, features=FEATURES)
+        cfg = make_cfg()
+        total_num_steps = calculate_total_num_steps(cfg, dataset)
+        assert not cfg.total_num_tokens
+        assert not cfg.total_supervised_tokens
+        assert total_num_steps == 0
+
+    def test_update_false_does_not_mutate_cfg(self, columns):
+        dataset = Dataset.from_dict(columns, features=FEATURES)
+        cfg = make_cfg()
+        calculate_total_num_steps(cfg, dataset, update=False)
+        assert cfg.total_num_tokens is None
+        assert cfg.total_supervised_tokens is None
diff --git a/tests/test_chunked_xentropy.py b/tests/test_chunked_xentropy.py
new file mode 100644
index 0000000000..56ac1b1682
--- /dev/null
+++ b/tests/test_chunked_xentropy.py
@@ -0,0 +1,40 @@
+"""
+test suite for chunked cross entropy
+"""
+
+import pytest
+import torch
+from torch import nn
+
+from axolotl.monkeypatch.loss.chunked import get_causal_lm_loss
+
+
+@pytest.fixture
+def chunked_fixtures():
+    model_dim = 512
+    vocab_size = 1024 * 256
+    seq_len = 2048
+    batch_size = 1
+
+    lm_head = nn.Linear(model_dim, vocab_size)
+    hidden_state = torch.randn(batch_size, seq_len, model_dim)
+    labels = torch.randint(low=0, high=vocab_size, size=(batch_size, seq_len))
+    return lm_head, hidden_state, labels, vocab_size
+
+
+def test_chunked_forward(chunked_fixtures):
+    lm_head, hidden_state, labels, vocab_size = chunked_fixtures
+    lm_loss = get_causal_lm_loss()
+
+    logits = lm_head(hidden_state)
+
+    chunked_lm_loss = lm_loss(logits, labels)
+
+    logits_flattened = logits.view(-1, vocab_size)
+    labels_flattened = labels.view(-1)
+
+    loss = nn.functional.cross_entropy(
+        logits_flattened.float(), labels_flattened, reduction="mean"
+    )
+
+    assert torch.allclose(chunked_lm_loss, loss, atol=1e-2, rtol=1e-2)
diff --git a/tests/test_context_parallel_batch_size.py b/tests/test_context_parallel_batch_size.py
new file mode 100644
index 0000000000..8f6ed7b288
--- /dev/null
+++ b/tests/test_context_parallel_batch_size.py
@@ -0,0 +1,56 @@
+"""Tests for batch_size calculation with context parallelism."""
+
+import sys
+import types
+
+import pytest
+
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="cp_base_cfg")
+def fixture_cp_base_cfg(min_base_cfg):
+    return (
+        DictDefault(
+            micro_batch_size=2,
+            gradient_accumulation_steps=4,
+            sequence_len=2048,
+            num_epochs=1,
+            flash_attention=True,
+        )
+        | min_base_cfg
+    )
+
+
+class TestContextParallelBatchSize:
+    """Verify batch_size scales by effective dp world_size when using context parallelism."""
+
+    @pytest.mark.parametrize(
+        "world_size, context_parallel_size, expected_batch_size",
+        [
+            (4, 1, 32),  # no CP: 2*4*4 = 32
+            (4, 2, 16),  # CP=2: 2*4*(4//2) = 16
+            (4, 4, 8),  # CP=4: 2*4*(4//4) = 8
+            (2, 2, 8),  # CP=ws: 2*4*(2//2) = 8 (no scaling)
+        ],
+    )
+    def test_batch_size_with_context_parallelism(
+        self,
+        cp_base_cfg,
+        monkeypatch,
+        world_size,
+        context_parallel_size,
+        expected_batch_size,
+    ):
+        monkeypatch.setenv("WORLD_SIZE", str(world_size))
+        # Mock ring_flash_attn since it's not installable on CPU,
+        # but required by schema validation when context_parallel_size > 1.
+        if "ring_flash_attn" not in sys.modules:
+            monkeypatch.setitem(
+                sys.modules, "ring_flash_attn", types.ModuleType("ring_flash_attn")
+            )
+        cp_base_cfg["context_parallel_size"] = context_parallel_size
+        cfg = validate_config(cp_base_cfg)
+        normalize_config(cfg)
+        assert cfg.batch_size == expected_batch_size
diff --git a/tests/test_convert.py b/tests/test_convert.py
new file mode 100644
index 0000000000..bfe0b603cf
--- /dev/null
+++ b/tests/test_convert.py
@@ -0,0 +1,91 @@
+"""Unit tests for src/axolotl/convert.py"""
+
+import json
+
+import pytest
+
+from axolotl.convert import (
+    FileReader,
+    FileWriter,
+    JsonlSerializer,
+    JsonParser,
+    JsonToJsonlConverter,
+    StdoutWriter,
+)
+
+
+class TestJsonParser:
+    def test_parse_valid_json_array(self):
+        parser = JsonParser()
+        result = parser.parse('[{"key": "value"}]')
+        assert result == [{"key": "value"}]
+
+    def test_parse_valid_json_object(self):
+        parser = JsonParser()
+        result = parser.parse('{"key": "value"}')
+        assert result == {"key": "value"}
+
+    def test_parse_invalid_json_raises(self):
+        parser = JsonParser()
+        with pytest.raises(json.JSONDecodeError):
+            parser.parse("not valid json")
+
+
+class TestJsonlSerializer:
+    def test_serialize_single_item(self):
+        serializer = JsonlSerializer()
+        result = serializer.serialize([{"a": 1}])
+        assert result == '{"a": 1}'
+
+    def test_serialize_multiple_items(self):
+        serializer = JsonlSerializer()
+        result = serializer.serialize([{"a": 1}, {"b": 2}])
+        lines = result.split("\n")
+        assert len(lines) == 2
+        assert json.loads(lines[0]) == {"a": 1}
+        assert json.loads(lines[1]) == {"b": 2}
+
+    def test_serialize_empty_list(self):
+        serializer = JsonlSerializer()
+        result = serializer.serialize([])
+        assert result == ""
+
+
+class TestFileReaderWriter:
+    def test_read_write_roundtrip(self, tmp_path):
+        test_file = tmp_path / "test.txt"
+        content = '{"hello": "world"}'
+        writer = FileWriter(str(test_file))
+        writer.write(content)
+
+        reader = FileReader()
+        result = reader.read(str(test_file))
+        assert result == content
+
+
+class TestStdoutWriter:
+    def test_write_to_stdout(self, capsys):
+        writer = StdoutWriter()
+        writer.write("hello")
+        captured = capsys.readouterr()
+        assert captured.out == "hello\n"
+
+
+class TestJsonToJsonlConverter:
+    def test_convert_json_to_jsonl(self, tmp_path):
+        input_data = [{"name": "Alice"}, {"name": "Bob"}]
+        input_file = tmp_path / "input.json"
+        output_file = tmp_path / "output.jsonl"
+
+        input_file.write_text(json.dumps(input_data), encoding="utf-8")
+
+        converter = JsonToJsonlConverter(
+            FileReader(), FileWriter(str(output_file)), JsonParser(), JsonlSerializer()
+        )
+        converter.convert(str(input_file))
+
+        result = output_file.read_text(encoding="utf-8")
+        lines = result.split("\n")
+        assert len(lines) == 2
+        assert json.loads(lines[0]) == {"name": "Alice"}
+        assert json.loads(lines[1]) == {"name": "Bob"}
diff --git a/tests/test_data.py b/tests/test_data.py
index 16af089a06..8dee3f629e 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -1,11 +1,17 @@
 """
-test module for the axolotl.utis.data module
+test module for the axolotl.utils.data module
 """
+
 import unittest
 
 from transformers import LlamaTokenizer
 
-from axolotl.utils.data import encode_pretraining, md5
+from axolotl.prompt_strategies.pretrain import load as load_pretrain
+from axolotl.utils.data import encode_streaming, md5
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.trainer import filter_sequences_by_length
+
+from tests.hf_offline_utils import enable_hf_offline
 
 
 class TestEncodePretraining(unittest.TestCase):
@@ -13,6 +19,7 @@ class TestEncodePretraining(unittest.TestCase):
     test class for encode pretraining and md5 helper
     """
 
+    @enable_hf_offline
     def setUp(self):
         self.tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b")
         self.tokenizer.add_special_tokens(
@@ -35,7 +42,7 @@ def test_encode_pretraining(self):
                 "hello, hello",
             ]
         }
-        result = encode_pretraining(self.tokenizer, self.max_tokens, examples["text"])
+        result = encode_streaming(examples, self.tokenizer, self.max_tokens)
 
         self.assertEqual(len(result["input_ids"]), 3)
 
@@ -53,12 +60,107 @@ def test_encode_pretraining(self):
         self.assertEqual(result["input_ids"][0][13], self.tokenizer.eos_token_id)
         self.assertEqual(result["input_ids"][0][14], self.tokenizer.pad_token_id)
 
+    def _pretrain_strategy(self, sequence_len):
+        cfg = DictDefault(
+            {
+                "train_on_inputs": False,
+                "sequence_len": sequence_len,
+                "pretraining_dataset": [{"text_column": "text"}],
+            }
+        )
+        return load_pretrain(self.tokenizer, cfg)
+
+    def test_long_document_is_chunked_not_dropped(self):
+        """Long docs must be chunked into windows that survive the length filter (#3441)."""
+        for sequence_len in (256, 512, 2048):
+            with self.subTest(sequence_len=sequence_len):
+                strat = self._pretrain_strategy(sequence_len)
+                long_doc = " ".join(f"token{i}" for i in range(4 * sequence_len))
+                windows = strat._tokenize(long_doc)["input_ids"]
+
+                # the document spans more than one window (i.e. it was chunked)
+                self.assertGreater(len(windows), 1)
+
+                for window in windows:
+                    # every window survives the downstream length filter ...
+                    self.assertLessEqual(len(window), sequence_len)
+                    self.assertTrue(
+                        filter_sequences_by_length(
+                            {"input_ids": window}, sequence_len=sequence_len
+                        )
+                    )
+                    # ... and ends with EOS
+                    self.assertEqual(window[-1], self.tokenizer.eos_token_id)
+
+    def test_no_tokens_dropped_for_oversized_docs(self):
+        """A doc longer than sequence_len must not be dropped entirely (#3441)."""
+        sequence_len = 256
+        strat = self._pretrain_strategy(sequence_len)
+        long_doc = " ".join(f"token{i}" for i in range(2000))
+        windows = strat._tokenize(long_doc)["input_ids"]
+
+        kept = [
+            w
+            for w in windows
+            if filter_sequences_by_length({"input_ids": w}, sequence_len=sequence_len)
+        ]
+        self.assertTrue(kept, "all windows were dropped — oversized doc lost entirely")
+        self.assertEqual(len(kept), len(windows))
+
+    def test_stride_below_window_size(self):
+        """Tokenization must not raise from a stride >= effective max length."""
+        for sequence_len in (256, 2048):
+            with self.subTest(sequence_len=sequence_len):
+                strat = self._pretrain_strategy(sequence_len)
+                # would raise ValueError from the tokenizer if stride were too large
+                strat._tokenize(" ".join(f"token{i}" for i in range(sequence_len * 3)))
+
     def test_md5(self):
         self.assertEqual(md5("hello world"), "5eb63bbbe01eeed093cb22bb8f5acdc3")
         self.assertEqual(
             md5("hello world", "utf-8"), "5eb63bbbe01eeed093cb22bb8f5acdc3"
         )
 
+    def test_excess_length_strategy(self):
+        """Test that excess_length_strategy results in a value error when set to 'raise'."""
+
+        # -- single sequence --
+        # This should work
+        data = {"input_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]}
+        filter_sequences_by_length(data, 32, raise_on_drop=True)
+
+        # This should return True, since data fits
+        dropped = filter_sequences_by_length(data, 32)
+        self.assertTrue(dropped)
+
+        # This should raise
+        self.assertRaises(
+            ValueError, filter_sequences_by_length, data, 15, raise_on_drop=True
+        )
+
+        # This should return False, since data doesn't fit
+        dropped = filter_sequences_by_length(data, 15)
+        self.assertFalse(dropped)
+
+        # -- batch sequence --
+        # This should work
+        data = {
+            "input_ids": [
+                [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
+                [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
+            ]
+        }
+        filter_sequences_by_length(data, 32, raise_on_drop=True)
+
+        # This should raise
+        self.assertRaises(
+            ValueError, filter_sequences_by_length, data, 15, raise_on_drop=True
+        )
+
+        # This should keep the first but drop the second entry
+        dropped = filter_sequences_by_length(data, 15)
+        self.assertEqual(dropped, [True, False])
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index a274b7b894..3bb3c8897a 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -1,34 +1,46 @@
-"""
-Test dataset loading under various conditions.
-"""
+"""Test dataset loading under various conditions."""
 
 import shutil
 import tempfile
-import unittest
 from pathlib import Path
+from typing import Any, Generator
+from unittest.mock import patch
 
+import pytest
 from datasets import Dataset
 from huggingface_hub import snapshot_download
-from transformers import AutoTokenizer
-
-from axolotl.utils.data import load_tokenized_prepared_datasets
+from transformers import PreTrainedTokenizer
+
+from axolotl.loaders.tokenizer import load_tokenizer
+from axolotl.utils.data.rl import prepare_preference_datasets
+from axolotl.utils.data.sft import (
+    _load_tokenized_prepared_datasets,
+)
+from axolotl.utils.data.shared import load_preprocessed_dataset
 from axolotl.utils.dict import DictDefault
 
+from tests.constants import (
+    ALPACA_MESSAGES_CONFIG_OG,
+    ALPACA_MESSAGES_CONFIG_REVISION,
+    SPECIAL_TOKENS,
+    alpaca_messages_dpo_rows,
+)
+from tests.hf_offline_utils import enable_hf_offline
+
 
-class TestDatasetPreparation(unittest.TestCase):
+class TestDatasetPreparation:
     """Test a configured dataloader."""
 
-    def setUp(self) -> None:
-        self.tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
-        self.tokenizer.add_special_tokens(
-            {
-                "bos_token": "<s>",
-                "eos_token": "</s>",
-                "unk_token": "<unk>",
-            }
-        )
-        # Alpaca dataset.
-        self.dataset = Dataset.from_list(
+    @pytest.fixture
+    def tokenizer(
+        self, tokenizer_huggyllama
+    ) -> Generator[PreTrainedTokenizer, Any, Any]:
+        tokenizer_huggyllama.add_special_tokens(SPECIAL_TOKENS)
+        yield tokenizer_huggyllama
+
+    @pytest.fixture
+    def dataset_fixture(self):
+        yield Dataset.from_list(
             [
                 {
                     "instruction": "Evaluate this sentence for spelling and grammar mistakes",
@@ -38,7 +50,30 @@ def setUp(self) -> None:
             ]
         )
 
-    def test_load_hub(self):
+    def test_load_preprocessed_dataset_uses_default_prepared_path(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            dataset_hash = "prepared-hash"
+            prepared_path = Path(tmp_dir) / dataset_hash
+            Dataset.from_dict({"value": [1]}).save_to_disk(str(prepared_path))
+            cfg = DictDefault(
+                {
+                    "dataset_prepared_path": None,
+                    "skip_prepare_dataset": False,
+                    "is_preprocess": False,
+                }
+            )
+
+            with patch(
+                "axolotl.utils.data.shared.DEFAULT_DATASET_PREPARED_PATH", tmp_dir
+            ):
+                dataset = load_preprocessed_dataset(cfg, dataset_hash)
+
+            assert dataset is not None
+            assert dataset["value"] == [1]
+
+    @pytest.mark.skip(reason="TODO: fix hf hub offline to work with HF rate limits")
+    @enable_hf_offline
+    def test_load_hub(self, tokenizer):
         """Core use case.  Verify that processing data from the hub works"""
         with tempfile.TemporaryDirectory() as tmp_dir:
             prepared_path = Path(tmp_dir) / "prepared"
@@ -55,25 +90,31 @@ def test_load_hub(self):
                 }
             )
 
-            dataset, _ = load_tokenized_prepared_datasets(
-                self.tokenizer, cfg, prepared_path
-            )
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
 
             assert len(dataset) == 2000
             assert "input_ids" in dataset.features
             assert "attention_mask" in dataset.features
             assert "labels" in dataset.features
 
-    def test_load_local_hub(self):
+    @enable_hf_offline
+    @pytest.mark.skip("datasets bug with local datasets when offline")
+    def test_load_local_hub(self, tokenizer):
         """Niche use case.  Verify that a local copy of a hub dataset can be loaded"""
         with tempfile.TemporaryDirectory() as tmp_dir:
-            tmp_ds_path = Path("mhenrichsen/alpaca_2k_test")
+            tmp_ds_path = Path(tmp_dir) / "mhenrichsen/alpaca_2k_test"
             tmp_ds_path.mkdir(parents=True, exist_ok=True)
-            snapshot_download(
+            snapshot_path = snapshot_download(
                 repo_id="mhenrichsen/alpaca_2k_test",
                 repo_type="dataset",
                 local_dir=tmp_ds_path,
             )
+            # offline mode doesn't actually copy it to local_dir, so we
+            # have to copy all the contents in the dir manually from the returned snapshot_path
+            shutil.copytree(snapshot_path, tmp_ds_path, dirs_exist_ok=True)
 
             prepared_path = Path(tmp_dir) / "prepared"
             # Right now a local copy that doesn't fully conform to a dataset
@@ -81,7 +122,7 @@ def test_load_local_hub(self):
             # how to load it.
             cfg = DictDefault(
                 {
-                    "tokenizer_config": "huggyllama/llama-7b",
+                    "tokenizer_config": "HuggingFaceTB/SmolLM2-135M",
                     "sequence_len": 1024,
                     "datasets": [
                         {
@@ -89,16 +130,17 @@ def test_load_local_hub(self):
                             "ds_type": "parquet",
                             "type": "alpaca",
                             "data_files": [
-                                "mhenrichsen/alpaca_2k_test/alpaca_2000.parquet",
+                                f"{tmp_ds_path}/alpaca_2000.parquet",
                             ],
                         },
                     ],
                 }
             )
 
-            dataset, _ = load_tokenized_prepared_datasets(
-                self.tokenizer, cfg, prepared_path
-            )
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
 
             assert len(dataset) == 2000
             assert "input_ids" in dataset.features
@@ -106,11 +148,12 @@ def test_load_local_hub(self):
             assert "labels" in dataset.features
             shutil.rmtree(tmp_ds_path)
 
-    def test_load_from_save_to_disk(self):
+    @enable_hf_offline
+    def test_load_from_save_to_disk(self, tokenizer, dataset_fixture):
         """Usual use case.  Verify datasets saved via `save_to_disk` can be loaded."""
         with tempfile.TemporaryDirectory() as tmp_dir:
             tmp_ds_name = Path(tmp_dir) / "tmp_dataset"
-            self.dataset.save_to_disk(str(tmp_ds_name))
+            dataset_fixture.save_to_disk(str(tmp_ds_name))
 
             prepared_path = Path(tmp_dir) / "prepared"
             cfg = DictDefault(
@@ -123,25 +166,28 @@ def test_load_from_save_to_disk(self):
                             "type": "alpaca",
                         },
                     ],
+                    "dataset_num_proc": 1,
                 }
             )
 
-            dataset, _ = load_tokenized_prepared_datasets(
-                self.tokenizer, cfg, prepared_path
-            )
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
 
             assert len(dataset) == 1
             assert "input_ids" in dataset.features
             assert "attention_mask" in dataset.features
             assert "labels" in dataset.features
 
-    def test_load_from_dir_of_parquet(self):
-        """Usual use case.  Verify a directory of parquet files can be loaded."""
+    @enable_hf_offline
+    def test_load_from_dir_of_parquet(self, tokenizer, dataset_fixture):
+        """Usual use case. Verify a directory of parquet files can be loaded."""
         with tempfile.TemporaryDirectory() as tmp_dir:
             tmp_ds_dir = Path(tmp_dir) / "tmp_dataset"
             tmp_ds_dir.mkdir()
             tmp_ds_path = tmp_ds_dir / "shard1.parquet"
-            self.dataset.to_parquet(tmp_ds_path)
+            dataset_fixture.to_parquet(tmp_ds_path)
 
             prepared_path: Path = Path(tmp_dir) / "prepared"
             cfg = DictDefault(
@@ -159,25 +205,28 @@ def test_load_from_dir_of_parquet(self):
                             "type": "alpaca",
                         },
                     ],
+                    "dataset_num_proc": 1,
                 }
             )
 
-            dataset, _ = load_tokenized_prepared_datasets(
-                self.tokenizer, cfg, prepared_path
-            )
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
 
             assert len(dataset) == 1
             assert "input_ids" in dataset.features
             assert "attention_mask" in dataset.features
             assert "labels" in dataset.features
 
-    def test_load_from_dir_of_json(self):
+    @enable_hf_offline
+    def test_load_from_dir_of_json(self, tokenizer, dataset_fixture):
         """Standard use case.  Verify a directory of json files can be loaded."""
         with tempfile.TemporaryDirectory() as tmp_dir:
             tmp_ds_dir = Path(tmp_dir) / "tmp_dataset"
             tmp_ds_dir.mkdir()
             tmp_ds_path = tmp_ds_dir / "shard1.json"
-            self.dataset.to_json(tmp_ds_path)
+            dataset_fixture.to_json(tmp_ds_path)
 
             prepared_path: Path = Path(tmp_dir) / "prepared"
             cfg = DictDefault(
@@ -195,23 +244,26 @@ def test_load_from_dir_of_json(self):
                             "type": "alpaca",
                         },
                     ],
+                    "dataset_num_proc": 1,
                 }
             )
 
-            dataset, _ = load_tokenized_prepared_datasets(
-                self.tokenizer, cfg, prepared_path
-            )
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
 
             assert len(dataset) == 1
             assert "input_ids" in dataset.features
             assert "attention_mask" in dataset.features
             assert "labels" in dataset.features
 
-    def test_load_from_single_parquet(self):
+    @enable_hf_offline
+    def test_load_from_single_parquet(self, tokenizer, dataset_fixture):
         """Standard use case.  Verify a single parquet file can be loaded."""
         with tempfile.TemporaryDirectory() as tmp_dir:
             tmp_ds_path = Path(tmp_dir) / "tmp_dataset.parquet"
-            self.dataset.to_parquet(tmp_ds_path)
+            dataset_fixture.to_parquet(tmp_ds_path)
 
             prepared_path: Path = Path(tmp_dir) / "prepared"
             cfg = DictDefault(
@@ -225,23 +277,26 @@ def test_load_from_single_parquet(self):
                             "type": "alpaca",
                         },
                     ],
+                    "dataset_num_proc": 1,
                 }
             )
 
-            dataset, _ = load_tokenized_prepared_datasets(
-                self.tokenizer, cfg, prepared_path
-            )
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
 
             assert len(dataset) == 1
             assert "input_ids" in dataset.features
             assert "attention_mask" in dataset.features
             assert "labels" in dataset.features
 
-    def test_load_from_single_json(self):
+    @enable_hf_offline
+    def test_load_from_single_json(self, tokenizer, dataset_fixture):
         """Standard use case.  Verify a single json file can be loaded."""
         with tempfile.TemporaryDirectory() as tmp_dir:
             tmp_ds_path = Path(tmp_dir) / "tmp_dataset.json"
-            self.dataset.to_json(tmp_ds_path)
+            dataset_fixture.to_json(tmp_ds_path)
 
             prepared_path: Path = Path(tmp_dir) / "prepared"
             cfg = DictDefault(
@@ -255,18 +310,267 @@ def test_load_from_single_json(self):
                             "type": "alpaca",
                         },
                     ],
+                    "dataset_num_proc": 1,
                 }
             )
 
-            dataset, _ = load_tokenized_prepared_datasets(
-                self.tokenizer, cfg, prepared_path
-            )
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
 
             assert len(dataset) == 1
             assert "input_ids" in dataset.features
             assert "attention_mask" in dataset.features
             assert "labels" in dataset.features
 
+    @pytest.mark.skip(reason="TODO: fix hf offline mode for CI rate limits")
+    @enable_hf_offline
+    def test_load_hub_with_dpo(self):
+        """Verify that processing dpo data from the hub works"""
+
+        cfg = DictDefault(
+            {
+                "tokenizer_config": "huggyllama/llama-7b",
+                "sequence_len": 1024,
+                "rl": "dpo",
+                "chat_template": "llama3",
+                "datasets": [ALPACA_MESSAGES_CONFIG_OG],
+            }
+        )
+
+        tokenizer = load_tokenizer(cfg)
+        train_dataset, _ = prepare_preference_datasets(cfg, tokenizer)
+
+        assert len(train_dataset) == 64
+        assert "conversation" not in train_dataset.features
+        assert "chosen" in train_dataset.features
+        assert "rejected" in train_dataset.features
+        assert "prompt" in train_dataset.features
+
+    @pytest.mark.skip(reason="TODO: fix hf hub offline to work with HF rate limits")
+    @enable_hf_offline
+    def test_load_hub_with_revision(self, tokenizer):
+        """Verify that processing data from the hub works with a specific revision"""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            prepared_path = Path(tmp_dir) / "prepared"
+
+            # make sure prepared_path is empty
+            shutil.rmtree(prepared_path, ignore_errors=True)
+
+            cfg = DictDefault(
+                {
+                    "tokenizer_config": "huggyllama/llama-7b",
+                    "sequence_len": 1024,
+                    "datasets": [
+                        {
+                            "path": "mhenrichsen/alpaca_2k_test",
+                            "type": "alpaca",
+                            "revision": "d05c1cb",
+                        },
+                    ],
+                }
+            )
+
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
+
+            assert len(dataset) == 2000
+            assert "input_ids" in dataset.features
+            assert "attention_mask" in dataset.features
+            assert "labels" in dataset.features
+
+    @enable_hf_offline
+    def test_load_hub_with_revision_with_dpo(self):
+        """Verify that processing dpo data from the hub works with a specific revision"""
+        dataset = Dataset.from_list(alpaca_messages_dpo_rows())
+
+        cfg = DictDefault(
+            {
+                "tokenizer_config": "huggyllama/llama-7b",
+                "sequence_len": 1024,
+                "rl": "dpo",
+                "chat_template": "llama3",
+                "datasets": [ALPACA_MESSAGES_CONFIG_REVISION],
+                "dataset_num_proc": None,
+            }
+        )
+
+        with patch(
+            "axolotl.utils.data.rl.load_dataset_with_config"
+        ) as mock_load_dataset:
+            # Set up the mock to return different values on successive calls
+            mock_load_dataset.return_value = dataset
+
+            tokenizer = load_tokenizer(cfg)
+            train_dataset, _ = prepare_preference_datasets(cfg, tokenizer)
+
+            assert len(train_dataset) == len(dataset)
+            assert "conversation" not in train_dataset.features
+            assert "chosen" in train_dataset.features
+            assert "rejected" in train_dataset.features
+            assert "prompt" in train_dataset.features
+
+    @enable_hf_offline
+    @pytest.mark.skip("datasets bug with local datasets when offline")
+    def test_load_local_hub_with_revision(
+        self, dataset_fozziethebeat_alpaca_messages_2k_dpo_test_rev_ea82cff, tokenizer
+    ):
+        """Verify that a local copy of a hub dataset can be loaded with a specific revision"""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_ds_path = Path(tmp_dir) / "mhenrichsen/alpaca_2k_test"
+            tmp_ds_path.mkdir(parents=True, exist_ok=True)
+            snapshot_path = snapshot_download(
+                repo_id="mhenrichsen/alpaca_2k_test",
+                repo_type="dataset",
+                local_dir=tmp_ds_path,
+                revision="d05c1cb",
+            )
+            shutil.copytree(snapshot_path, tmp_ds_path, dirs_exist_ok=True)
+
+            prepared_path = Path(tmp_dir) / "prepared"
+            cfg = DictDefault(
+                {
+                    "tokenizer_config": "huggyllama/llama-7b",
+                    "sequence_len": 1024,
+                    "datasets": [
+                        {
+                            "path": "mhenrichsen/alpaca_2k_test",
+                            "ds_type": "parquet",
+                            "type": "alpaca",
+                            "data_files": [
+                                f"{tmp_ds_path}/alpaca_2000.parquet",
+                            ],
+                            "revision": "d05c1cb",
+                        },
+                    ],
+                }
+            )
+
+            with patch(
+                "axolotl.utils.data.shared.load_dataset_with_config"
+            ) as mock_load_dataset:
+                # Set up the mock to return different values on successive calls
+                mock_load_dataset.return_value = (
+                    dataset_fozziethebeat_alpaca_messages_2k_dpo_test_rev_ea82cff
+                )
+
+                with patch(
+                    "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH",
+                    str(prepared_path),
+                ):
+                    dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
+
+                assert len(dataset) == 2000
+                assert "input_ids" in dataset.features
+                assert "attention_mask" in dataset.features
+                assert "labels" in dataset.features
+                shutil.rmtree(tmp_ds_path)
+
+    @enable_hf_offline
+    def test_loading_local_dataset_folder(self, tokenizer):
+        """Verify that a dataset downloaded to a local folder can be loaded"""
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_ds_path = Path(tmp_dir) / "mhenrichsen/alpaca_2k_test"
+            tmp_ds_path.mkdir(parents=True, exist_ok=True)
+            snapshot_path = snapshot_download(
+                repo_id="mhenrichsen/alpaca_2k_test",
+                repo_type="dataset",
+            )
+            shutil.copytree(snapshot_path, tmp_ds_path, dirs_exist_ok=True)
+
+            prepared_path = Path(tmp_dir) / "prepared"
+            cfg = DictDefault(
+                {
+                    "tokenizer_config": "huggyllama/llama-7b",
+                    "sequence_len": 1024,
+                    "datasets": [
+                        {
+                            "path": str(tmp_ds_path),
+                            "type": "alpaca",
+                        },
+                    ],
+                    "dataset_num_proc": 1,
+                }
+            )
+
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
+
+            assert len(dataset) == 2000
+            assert "input_ids" in dataset.features
+            assert "attention_mask" in dataset.features
+            assert "labels" in dataset.features
+            shutil.rmtree(tmp_ds_path)
+
+    @enable_hf_offline
+    def test_load_dataset_with_str_json_data(self, tokenizer):
+        """
+        Test loading datasets where data is stored as str JSON instead of list of dicts.
+        see: https://github.com/axolotl-ai-cloud/axolotl/pull/3607 for more details.
+        """
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            import json
+
+            str_json_ds = Dataset.from_list(
+                [
+                    {
+                        "messages": json.dumps(
+                            [
+                                {"role": "user", "content": "Hello how are you?"},
+                                {
+                                    "role": "assistant",
+                                    "content": "I am doing good thanks",
+                                },
+                            ]
+                        )
+                    },
+                    {
+                        "messages": json.dumps(
+                            [
+                                {"role": "user", "content": "What is 2+2?"},
+                                {"role": "assistant", "content": "2+2 equals 4."},
+                            ]
+                        )
+                    },
+                ]
+            )
+
+            tmp_ds_path = Path(tmp_dir) / "str_json_dataset.parquet"
+            str_json_ds.to_parquet(tmp_ds_path)
+
+            prepared_path = Path(tmp_dir) / "prepared"
+            cfg = DictDefault(
+                {
+                    "tokenizer_config": "huggyllama/llama-7b",
+                    "sequence_len": 512,
+                    "datasets": [
+                        {
+                            "path": str(tmp_ds_path),
+                            "name": "test_str_json",
+                            "type": "chat_template",
+                            "field_messages": "messages",
+                            "message_field_role": "role",
+                            "message_field_content": "content",
+                        },
+                    ],
+                    "dataset_num_proc": 1,
+                }
+            )
+
+            with patch(
+                "axolotl.common.const.DEFAULT_DATASET_PREPARED_PATH", str(prepared_path)
+            ):
+                dataset, _ = _load_tokenized_prepared_datasets(tokenizer, cfg)
+
+            assert len(dataset) == 2
+            assert "input_ids" in dataset.features
+            assert "attention_mask" in dataset.features
+            assert "labels" in dataset.features
 
-if __name__ == "__main__":
-    unittest.main()
+            assert len(dataset[0]["input_ids"]) > 0
diff --git a/tests/test_dict.py b/tests/test_dict.py
index 2007cb085e..19a3701997 100644
--- a/tests/test_dict.py
+++ b/tests/test_dict.py
@@ -1,6 +1,5 @@
 """Module for testing DictDefault class"""
 
-
 import unittest
 
 import pytest
@@ -22,26 +21,26 @@ def test_dict_default(self):
             }
         )
 
-        assert (
-            cfg.key_a.key_b == "value_a"
-        ), "DictDefault should return value for existing nested keys"
+        assert cfg.key_a.key_b == "value_a", (
+            "DictDefault should return value for existing nested keys"
+        )
 
-        assert (
-            cfg.key_c == "value_c"
-        ), "DictDefault should return value for existing keys"
+        assert cfg.key_c == "value_c", (
+            "DictDefault should return value for existing keys"
+        )
 
-        assert (
-            cfg.key_d[0] == "value_d"
-        ), "DictDefault should return value for existing keys in list"
+        assert cfg.key_d[0] == "value_d", (
+            "DictDefault should return value for existing keys in list"
+        )
 
-        assert (
-            "value_e" in cfg.key_d
-        ), "DictDefault should support in operator for existing keys in list"
+        assert "value_e" in cfg.key_d, (
+            "DictDefault should support in operator for existing keys in list"
+        )
 
     def test_dict_or_operator(self):
         cfg = DictDefault({"key_a": {"key_b": "value_b"}, "key_f": "value_g"})
 
-        cfg = cfg | DictDefault(  # pylint: disable=unsupported-binary-operation
+        cfg = cfg | DictDefault(
             {
                 "key_a": {"key_b": "value_a"},
                 "key_c": "value_c",
@@ -50,9 +49,9 @@ def test_dict_or_operator(self):
             }
         )
 
-        assert (
-            cfg.key_a.key_b == "value_b"
-        ), "DictDefault should support OR operator for existing nested keys"
+        assert cfg.key_a.key_b == "value_b", (
+            "DictDefault should support OR operator for existing nested keys"
+        )
 
         assert cfg.key_c == "value_c", "DictDefault should not delete existing key"
 
@@ -61,9 +60,9 @@ def test_dict_or_operator(self):
             "value_e",
         ], "DictDefault should not overwrite existing keys in list"
 
-        assert (
-            cfg.key_f == "value_g"
-        ), "DictDefault should support OR operator for existing key"
+        assert cfg.key_f == "value_g", (
+            "DictDefault should support OR operator for existing key"
+        )
 
     def test_dict_missingkey(self):
         cfg = DictDefault({})
@@ -73,9 +72,9 @@ def test_dict_missingkey(self):
     def test_dict_or(self):
         cfg = DictDefault({}) | DictDefault({})
 
-        assert (
-            cfg.random_key is None
-        ), "DictDefault should return None for missing keys after | operation"
+        assert cfg.random_key is None, (
+            "DictDefault should return None for missing keys after | operation"
+        )
 
     def test_dict_nested_missingparentkey(self):
         """
diff --git a/tests/test_ebft_kernels.py b/tests/test_ebft_kernels.py
new file mode 100644
index 0000000000..2c05a887ae
--- /dev/null
+++ b/tests/test_ebft_kernels.py
@@ -0,0 +1,294 @@
+"""Correctness tests for fused EBFT Triton kernels."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from axolotl.core.trainers.ebft.kernels import (
+    fused_cosine_similarity,
+    fused_diversity_penalty,
+    fused_log_softmax_gather,
+    fused_reinforce_loss,
+)
+
+# Skip all tests if CUDA not available
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA required for Triton kernels"
+)
+
+DEVICE = "cuda"
+
+
+# ---------------------------------------------------------------------------
+# 1. fused_log_softmax_gather
+# ---------------------------------------------------------------------------
+class TestFusedLogSoftmaxGather:
+    def _reference(self, logits, labels):
+        """PyTorch reference: log_softmax + gather."""
+        lp = F.log_softmax(logits.float(), dim=-1)
+        return lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
+
+    def test_basic_correctness(self):
+        B, S, V = 2, 16, 1024
+        logits = torch.randn(B, S, V, device=DEVICE, dtype=torch.bfloat16)
+        labels = torch.randint(0, V, (B, S), device=DEVICE)
+
+        ref = self._reference(logits, labels)
+        out = fused_log_softmax_gather(logits, labels)
+
+        torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+    def test_large_vocab(self):
+        """Test with realistic vocab size (128K)."""
+        B, S, V = 1, 8, 128256
+        logits = torch.randn(B, S, V, device=DEVICE, dtype=torch.bfloat16)
+        labels = torch.randint(0, V, (B, S), device=DEVICE)
+
+        ref = self._reference(logits, labels)
+        out = fused_log_softmax_gather(logits, labels)
+
+        torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+    def test_fp32_input(self):
+        B, S, V = 2, 8, 512
+        logits = torch.randn(B, S, V, device=DEVICE, dtype=torch.float32)
+        labels = torch.randint(0, V, (B, S), device=DEVICE)
+
+        ref = self._reference(logits, labels)
+        out = fused_log_softmax_gather(logits, labels)
+
+        torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+
+    def test_output_is_negative(self):
+        """log_softmax values should always be <= 0."""
+        B, S, V = 4, 32, 2048
+        logits = torch.randn(B, S, V, device=DEVICE, dtype=torch.bfloat16)
+        labels = torch.randint(0, V, (B, S), device=DEVICE)
+
+        out = fused_log_softmax_gather(logits, labels)
+        assert (out <= 1e-5).all(), "log_softmax values should be <= 0"
+
+    def test_extreme_logits(self):
+        """Test numerical stability with very large/small logits."""
+        B, S, V = 1, 4, 256
+        logits = torch.randn(B, S, V, device=DEVICE, dtype=torch.float32)
+        logits[:, 0, :] = 1000.0  # very large
+        logits[:, 1, :] = -1000.0  # very small
+        logits[:, 2, 0] = 1000.0  # one hot-ish
+        labels = torch.zeros(B, S, device=DEVICE, dtype=torch.long)
+
+        ref = self._reference(logits, labels)
+        out = fused_log_softmax_gather(logits, labels)
+
+        assert torch.isfinite(out).all(), "Should handle extreme values"
+        torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
+
+    def test_2d_input(self):
+        """Test with pre-flattened (N, V) input."""
+        N, V = 64, 4096
+        logits = torch.randn(N, V, device=DEVICE, dtype=torch.bfloat16)
+        labels = torch.randint(0, V, (N,), device=DEVICE)
+
+        ref = self._reference(logits.unsqueeze(0), labels.unsqueeze(0)).squeeze(0)
+        out = fused_log_softmax_gather(logits, labels)
+
+        torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# 2. fused_reinforce_loss
+# ---------------------------------------------------------------------------
+class TestFusedReinforceLoss:
+    def _reference(self, logps, advantages, mask):
+        """PyTorch reference implementation."""
+        loss_per_token = -logps * advantages
+        return (loss_per_token * mask.float()).sum() / mask.float().sum().clamp(min=1)
+
+    def test_basic_correctness(self):
+        N = 1024
+        logps = torch.randn(N, device=DEVICE, dtype=torch.float32)
+        advs = torch.randn(N, device=DEVICE, dtype=torch.float32)
+        mask = torch.randint(0, 2, (N,), device=DEVICE, dtype=torch.bool)
+
+        ref = self._reference(logps, advs, mask)
+        out = fused_reinforce_loss(logps, advs, mask)
+
+        torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
+
+    def test_2d_input(self):
+        """Test with (B, S) shaped inputs."""
+        B, S = 4, 256
+        logps = torch.randn(B, S, device=DEVICE, dtype=torch.float32)
+        advs = torch.randn(B, S, device=DEVICE, dtype=torch.float32)
+        mask = torch.randint(0, 2, (B, S), device=DEVICE, dtype=torch.bool)
+
+        ref = self._reference(logps, advs, mask)
+        out = fused_reinforce_loss(logps, advs, mask)
+
+        torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
+
+    def test_all_masked(self):
+        """All-zero mask should return 0."""
+        N = 512
+        logps = torch.randn(N, device=DEVICE, dtype=torch.float32)
+        advs = torch.randn(N, device=DEVICE, dtype=torch.float32)
+        mask = torch.zeros(N, device=DEVICE, dtype=torch.bool)
+
+        out = fused_reinforce_loss(logps, advs, mask)
+        assert out.item() == 0.0
+
+    def test_all_unmasked(self):
+        N = 512
+        logps = torch.randn(N, device=DEVICE, dtype=torch.float32)
+        advs = torch.randn(N, device=DEVICE, dtype=torch.float32)
+        mask = torch.ones(N, device=DEVICE, dtype=torch.bool)
+
+        ref = self._reference(logps, advs, mask)
+        out = fused_reinforce_loss(logps, advs, mask)
+
+        torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
+
+    def test_large(self):
+        """Test with realistic size (4 * 3000 tokens)."""
+        N = 12000
+        logps = torch.randn(N, device=DEVICE, dtype=torch.float32)
+        advs = torch.randn(N, device=DEVICE, dtype=torch.float32)
+        mask = torch.randint(0, 2, (N,), device=DEVICE, dtype=torch.bool)
+
+        ref = self._reference(logps, advs, mask)
+        out = fused_reinforce_loss(logps, advs, mask)
+
+        torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# 3. fused_cosine_similarity
+# ---------------------------------------------------------------------------
+class TestFusedCosineSimilarity:
+    def test_basic_correctness(self):
+        N, D = 64, 256
+        a = torch.randn(N, D, device=DEVICE, dtype=torch.bfloat16)
+        b = torch.randn(N, D, device=DEVICE, dtype=torch.bfloat16)
+
+        ref = F.cosine_similarity(a.float(), b.float(), dim=-1)
+        out = fused_cosine_similarity(a, b)
+
+        torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+    def test_batched(self):
+        """Test with (B, N, NB, D) shaped input."""
+        B, N, NB, D = 2, 4, 16, 512
+        a = torch.randn(B, N, NB, D, device=DEVICE, dtype=torch.bfloat16)
+        b = torch.randn(B, N, NB, D, device=DEVICE, dtype=torch.bfloat16)
+
+        ref = F.cosine_similarity(a.float(), b.float(), dim=-1)
+        out = fused_cosine_similarity(a, b)
+
+        torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+    def test_identical_vectors(self):
+        """Identical vectors should give similarity = 1."""
+        N, D = 16, 128
+        a = torch.randn(N, D, device=DEVICE, dtype=torch.float32)
+
+        out = fused_cosine_similarity(a, a)
+        torch.testing.assert_close(
+            out,
+            torch.ones(N, device=DEVICE, dtype=torch.float32),
+            atol=1e-5,
+            rtol=1e-5,
+        )
+
+    def test_orthogonal_vectors(self):
+        """Orthogonal vectors should give similarity = 0."""
+        D = 128
+        a = torch.zeros(1, D, device=DEVICE, dtype=torch.float32)
+        b = torch.zeros(1, D, device=DEVICE, dtype=torch.float32)
+        a[0, 0] = 1.0
+        b[0, 1] = 1.0
+
+        out = fused_cosine_similarity(a, b)
+        assert abs(out.item()) < 1e-5
+
+    def test_opposite_vectors(self):
+        """Opposite vectors should give similarity = -1."""
+        N, D = 8, 64
+        a = torch.randn(N, D, device=DEVICE, dtype=torch.float32)
+        out = fused_cosine_similarity(a, -a)
+        torch.testing.assert_close(
+            out,
+            -torch.ones(N, device=DEVICE, dtype=torch.float32),
+            atol=1e-5,
+            rtol=1e-5,
+        )
+
+    def test_large_dimension(self):
+        """Test with large feature dimension (multi-layer concatenated features)."""
+        N, D = 32, 4608  # 3 layers * 1536 hidden
+        a = torch.randn(N, D, device=DEVICE, dtype=torch.bfloat16)
+        b = torch.randn(N, D, device=DEVICE, dtype=torch.bfloat16)
+
+        ref = F.cosine_similarity(a.float(), b.float(), dim=-1)
+        out = fused_cosine_similarity(a, b)
+
+        torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
+# ---------------------------------------------------------------------------
+# 4. fused_diversity_penalty
+# ---------------------------------------------------------------------------
+class TestFusedDiversityPenalty:
+    def _reference(self, embeddings):
+        """PyTorch reference: bmm + mask diagonal + mean."""
+        B, N, D = embeddings.shape
+        sims = torch.bmm(embeddings.float(), embeddings.float().transpose(1, 2))
+        eye = torch.eye(N, device=embeddings.device, dtype=torch.bool)
+        sims = sims.masked_fill(eye.unsqueeze(0), 0.0)
+        return sims.sum(dim=-1) / (N - 1)
+
+    def test_basic_correctness(self):
+        B, N, D = 4, 4, 256
+        emb = torch.randn(B, N, D, device=DEVICE, dtype=torch.bfloat16)
+
+        ref = self._reference(emb)
+        out = fused_diversity_penalty(emb)
+
+        torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+    def test_two_samples(self):
+        """With n=2, diversity = dot(a, b) for each."""
+        B, D = 3, 128
+        emb = torch.randn(B, 2, D, device=DEVICE, dtype=torch.float32)
+
+        ref = self._reference(emb)
+        out = fused_diversity_penalty(emb)
+
+        torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
+
+    def test_identical_samples(self):
+        """All identical samples should give max diversity."""
+        B, N, D = 2, 4, 64
+        vec = torch.randn(B, 1, D, device=DEVICE, dtype=torch.float32)
+        emb = vec.expand(B, N, D).contiguous()
+
+        out = fused_diversity_penalty(emb)
+        # Should be ||vec||^2 for each (self-excluded mean of identical dot products)
+        expected = (vec.squeeze(1) ** 2).sum(dim=-1, keepdim=True).expand(B, N)
+        torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-4)
+
+    def test_large(self):
+        """Test with realistic EBFT dimensions."""
+        B, N, D = 1, 4, 4608  # multi-layer features
+        emb = torch.randn(B, N, D, device=DEVICE, dtype=torch.bfloat16)
+
+        ref = self._reference(emb)
+        out = fused_diversity_penalty(emb)
+
+        torch.testing.assert_close(out, ref, atol=5e-2, rtol=5e-2)
+
+    def test_single_sample_returns_zeros(self):
+        """N=1 should return zeros (no pairs), not garbage from uninitialized memory."""
+        B, D = 3, 128
+        emb = torch.randn(B, 1, D, device=DEVICE, dtype=torch.float32)
+        out = fused_diversity_penalty(emb)
+        assert (out == 0).all(), "N=1 diversity should be exactly zero"
diff --git a/tests/test_ebft_strided_structured.py b/tests/test_ebft_strided_structured.py
new file mode 100644
index 0000000000..e4ad946a02
--- /dev/null
+++ b/tests/test_ebft_strided_structured.py
@@ -0,0 +1,363 @@
+"""Tests for the EBFT strided structured dataset transform and data loading."""
+
+import pytest
+from datasets import Dataset
+from tokenizers import Tokenizer, models, pre_tokenizers
+from transformers import PreTrainedTokenizerFast
+
+from axolotl.prompt_strategies.ebft import load as load_ebft
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture
+def tokenizer():
+    """Create a simple word-level tokenizer — no network access needed."""
+    # Build a tiny vocab covering common test words
+    vocab = {"[PAD]": 0, "[UNK]": 1, "[BOS]": 2, "[EOS]": 3}
+    words = (
+        "what is 2 + the answer 4 hello world goodbye bye hi short prompt "
+        "x write code print test some string metadata noise ok good python "
+        "sampling abc 123 def solve return this that"
+    ).split()
+    for w in words:
+        if w not in vocab:
+            vocab[w] = len(vocab)
+
+    backend = Tokenizer(models.WordLevel(vocab=vocab, unk_token="[UNK]"))
+    backend.pre_tokenizer = pre_tokenizers.Whitespace()
+
+    tok = PreTrainedTokenizerFast(
+        tokenizer_object=backend,
+        bos_token="[BOS]",
+        eos_token="[EOS]",
+        pad_token="[PAD]",
+        unk_token="[UNK]",
+    )
+    return tok
+
+
+@pytest.fixture
+def cfg():
+    return DictDefault({"sequence_len": 64})
+
+
+@pytest.fixture
+def transform_fn_and_kwargs(cfg):
+    result = load_ebft("ebft_strided_structured.transform", cfg)
+    assert result is not None, "Failed to load ebft_strided_structured transform"
+    transform_fn, map_kwargs = result
+    return transform_fn, map_kwargs
+
+
+class TestEBFTStridedStructuredTransform:
+    """Tests for the dataset transform function itself."""
+
+    def test_transform_loads(self, transform_fn_and_kwargs):
+        transform_fn, map_kwargs = transform_fn_and_kwargs
+        assert callable(transform_fn)
+        assert "remove_columns" in map_kwargs
+
+    def test_remove_columns_sentinel(self, transform_fn_and_kwargs):
+        """Transform should signal removal of all original columns."""
+        _, map_kwargs = transform_fn_and_kwargs
+        assert map_kwargs["remove_columns"] == "__all__"
+
+    def test_prompt_completion_tokenization(self, transform_fn_and_kwargs, tokenizer):
+        """Prompt tokens get labels=-100, completion tokens get real labels."""
+        transform_fn, _ = transform_fn_and_kwargs
+        example = {"input": "what is 2 + 2", "output": "the answer is 4"}
+        result = transform_fn(example, tokenizer=tokenizer)
+
+        assert "input_ids" in result
+        assert "labels" in result
+        assert "attention_mask" in result
+        assert "prompt_length" in result
+
+        prompt_length = result["prompt_length"]
+        labels = result["labels"]
+        seq_len = len(result["input_ids"])
+
+        assert seq_len == 64, "Should be padded to sequence_len"
+        assert len(labels) == seq_len
+        assert prompt_length > 0
+
+        # Prompt tokens should be masked
+        for i in range(prompt_length):
+            assert labels[i] == -100, f"Prompt token at {i} should be -100"
+
+        # At least one completion token should have a real label
+        completion_labels = [lab for lab in labels[prompt_length:] if lab != -100]
+        assert len(completion_labels) > 0, "Should have non-masked completion tokens"
+
+    def test_prompt_length_matches_boundary(self, transform_fn_and_kwargs, tokenizer):
+        """prompt_length should be the exact boundary between -100 and real labels."""
+        transform_fn, _ = transform_fn_and_kwargs
+        example = {"input": "hello world", "output": "goodbye world"}
+        result = transform_fn(example, tokenizer=tokenizer)
+
+        prompt_length = result["prompt_length"]
+        labels = result["labels"]
+
+        assert labels[prompt_length - 1] == -100, "Last prompt token should be masked"
+        assert labels[prompt_length] != -100, (
+            "First completion token should not be masked"
+        )
+
+    def test_padding_tokens_masked(self, transform_fn_and_kwargs, tokenizer):
+        """Padding tokens should have labels=-100 and attention_mask=0."""
+        transform_fn, _ = transform_fn_and_kwargs
+        example = {"input": "hi", "output": "bye"}
+        result = transform_fn(example, tokenizer=tokenizer)
+
+        attention_mask = result["attention_mask"]
+        labels = result["labels"]
+
+        real_len = sum(attention_mask)
+        assert real_len < 64, "Short example should have padding"
+
+        for i in range(real_len, 64):
+            assert attention_mask[i] == 0, (
+                f"Pad position {i} should have attention_mask=0"
+            )
+            assert labels[i] == -100, f"Pad position {i} should have labels=-100"
+
+    def test_truncation_long_completion(self, transform_fn_and_kwargs, tokenizer):
+        """Long completions should be truncated to fit sequence_len."""
+        transform_fn, _ = transform_fn_and_kwargs
+        example = {
+            "input": "short prompt",
+            "output": "x " * 500,
+        }
+        result = transform_fn(example, tokenizer=tokenizer)
+        assert len(result["input_ids"]) == 64
+
+    def test_alternative_field_names(self, transform_fn_and_kwargs, tokenizer):
+        """Transform should handle different field name conventions."""
+        transform_fn, _ = transform_fn_and_kwargs
+
+        result = transform_fn(
+            {"prompt": "what", "completion": "this"}, tokenizer=tokenizer
+        )
+        assert result["prompt_length"] > 0
+
+        result = transform_fn(
+            {"question": "what", "answer": "this"}, tokenizer=tokenizer
+        )
+        assert result["prompt_length"] > 0
+
+    def test_without_tokenizer_returns_prompt(self, transform_fn_and_kwargs):
+        """Without tokenizer, should return a prompt string."""
+        transform_fn, _ = transform_fn_and_kwargs
+        result = transform_fn({"input": "hello", "output": "world"})
+        assert "prompt" in result
+        assert result["prompt"] == "hello"
+
+
+class TestEBFTColumnRemoval:
+    """Tests for the __all__ column removal logic in the RL data path."""
+
+    def _filter_remove_columns(self, map_kwargs, dataset):
+        """Reproduce the filtering logic from rl.py _load_split."""
+        if "remove_columns" in map_kwargs:
+            ds_columns = dataset.column_names
+            if map_kwargs["remove_columns"] == "__all__":
+                map_kwargs["remove_columns"] = list(ds_columns)
+            else:
+                map_kwargs["remove_columns"] = [
+                    c for c in map_kwargs["remove_columns"] if c in ds_columns
+                ]
+        return map_kwargs
+
+    def test_all_original_columns_removed(self, transform_fn_and_kwargs, tokenizer):
+        """After mapping, only tokenized columns should remain."""
+        transform_fn, map_kwargs = transform_fn_and_kwargs
+        map_kwargs = dict(map_kwargs)  # copy
+
+        ds = Dataset.from_list(
+            [
+                {"input": "what is 2 + 2", "output": "4", "extra_field": "noise"},
+            ]
+        )
+
+        map_kwargs = self._filter_remove_columns(map_kwargs, ds)
+        assert "input" in map_kwargs["remove_columns"]
+        assert "output" in map_kwargs["remove_columns"]
+        assert "extra_field" in map_kwargs["remove_columns"]
+
+        from functools import partial
+
+        mapped = ds.map(partial(transform_fn, tokenizer=tokenizer), **map_kwargs)
+        assert "input_ids" in mapped.column_names
+        assert "labels" in mapped.column_names
+        assert "prompt_length" in mapped.column_names
+        assert "input" not in mapped.column_names
+        assert "output" not in mapped.column_names
+        assert "extra_field" not in mapped.column_names
+
+    def test_extra_metadata_columns_removed(self, transform_fn_and_kwargs, tokenizer):
+        """Datasets with many extra metadata columns should all be cleaned up."""
+        transform_fn, map_kwargs = transform_fn_and_kwargs
+        map_kwargs = dict(map_kwargs)
+
+        ds = Dataset.from_list(
+            [
+                {
+                    "input": "write hello world",
+                    "output": "print hello",
+                    "id": "abc 123",
+                    "domain": "python",
+                    "generation_algorithm": "sampling",
+                    "llm_judgement": "good",
+                    "unit_tests": "test",
+                    "tests_execution_status": "ok",
+                    "average_test_score": 0.95,
+                },
+            ]
+        )
+
+        map_kwargs = self._filter_remove_columns(map_kwargs, ds)
+        assert len(map_kwargs["remove_columns"]) == 9
+
+        from functools import partial
+
+        mapped = ds.map(partial(transform_fn, tokenizer=tokenizer), **map_kwargs)
+
+        expected_columns = {"input_ids", "attention_mask", "labels", "prompt_length"}
+        assert set(mapped.column_names) == expected_columns
+
+    def test_no_string_columns_remain(self, transform_fn_and_kwargs, tokenizer):
+        """No string-typed columns should remain (would crash the DataLoader)."""
+        transform_fn, map_kwargs = transform_fn_and_kwargs
+        map_kwargs = dict(map_kwargs)
+
+        ds = Dataset.from_list(
+            [
+                {"input": "test", "output": "test", "notes": "some string metadata"},
+            ]
+        )
+
+        map_kwargs = self._filter_remove_columns(map_kwargs, ds)
+
+        from functools import partial
+
+        mapped = ds.map(partial(transform_fn, tokenizer=tokenizer), **map_kwargs)
+
+        for col in mapped.column_names:
+            feat = mapped.features[col]
+            assert str(feat) != "string", (
+                f"Column '{col}' is still a string — would crash DataLoader"
+            )
+
+    def test_filter_preserves_explicit_list(self):
+        """When remove_columns is an explicit list, only existing columns are kept."""
+        ds = Dataset.from_list([{"a": 1, "b": "text", "c": 3}])
+        map_kwargs = {"remove_columns": ["a", "b", "missing_col"]}
+
+        ds_columns = ds.column_names
+        map_kwargs["remove_columns"] = [
+            c for c in map_kwargs["remove_columns"] if c in ds_columns
+        ]
+
+        assert map_kwargs["remove_columns"] == ["a", "b"]
+        assert "missing_col" not in map_kwargs["remove_columns"]
+
+
+class TestMultiTurnSeparators:
+    """Verify multi-turn transforms and trainer-side GT reconstruction."""
+
+    def test_multiturn_transform_splits_turns(self):
+        """Transform should store first turn as GT and remaining turns separately."""
+        from axolotl.prompt_strategies.ebft import load as load_ebft
+        from axolotl.utils.dict import DictDefault
+
+        cfg = DictDefault({"sequence_len": 512})
+        fn, _ = load_ebft("ebft_chat_multiturn.transform", cfg)
+        out = fn(
+            {
+                "messages": [
+                    {"role": "user", "content": "Q1"},
+                    {"role": "assistant", "content": "A1"},
+                    {"role": "user", "content": "Q2"},
+                    {"role": "assistant", "content": "A2"},
+                ]
+            }
+        )
+        # ground_truth is only the first assistant turn
+        assert out["ground_truth"] == "A1"
+        # remaining_turns carries the rest for trainer-side reconstruction
+        assert out["remaining_turns"] == [
+            {"role": "user", "content": "Q2"},
+            {"role": "assistant", "content": "A2"},
+        ]
+
+    def test_multiturn_gt_reconstruction_via_chat_template(self):
+        """Trainer-side GT reconstruction should insert role markers between turns.
+
+        This tests the logic from trainer.py:284-299 that reconstructs multi-turn
+        GT using apply_chat_template, ensuring assistant turns are separated by
+        role markers rather than concatenated as raw text.
+        """
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            "Qwen/Qwen2-0.5B-Instruct", trust_remote_code=True
+        )
+
+        # Simulate the transform output
+        prompt_msgs = [{"role": "user", "content": "Q1"}]
+        gt = "A1"
+        remaining_turns = [
+            {"role": "user", "content": "Q2"},
+            {"role": "assistant", "content": "A2"},
+        ]
+
+        # --- Reproduce the trainer-side reconstruction (trainer.py:284-299) ---
+        prompt_text = tokenizer.apply_chat_template(
+            prompt_msgs, tokenize=False, add_generation_prompt=True
+        )
+        gt_conv = list(prompt_msgs) + [{"role": "assistant", "content": gt}]
+        gt_conv.extend(remaining_turns)
+        full_gt_text = tokenizer.apply_chat_template(
+            gt_conv, tokenize=False, add_generation_prompt=False
+        )
+
+        # The full GT text should contain both assistant turns with role markers
+        assert "A1" in full_gt_text
+        assert "A2" in full_gt_text
+        # Raw concatenation "A1A2" should NOT appear — role markers separate them
+        assert "A1A2" not in full_gt_text, (
+            "GT reconstruction should have role markers between turns, not raw concatenation"
+        )
+        # The user turn Q2 should appear between A1 and A2
+        a1_pos = full_gt_text.index("A1")
+        a2_pos = full_gt_text.index("A2")
+        q2_pos = full_gt_text.index("Q2")
+        assert a1_pos < q2_pos < a2_pos, (
+            "Turn order should be A1 -> Q2 -> A2 in rendered GT"
+        )
+        # The GT should start with the prompt
+        assert full_gt_text.startswith(prompt_text), (
+            "Full GT should start with the rendered prompt"
+        )
+
+    def test_multiturn_gt_reconstruction_fallback_single_turn(self):
+        """Single-turn prompts in a multi-turn dataset should use raw concatenation."""
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            "Qwen/Qwen2-0.5B-Instruct", trust_remote_code=True
+        )
+
+        prompt_msgs = [{"role": "user", "content": "Q1"}]
+        gt = "A1"
+        # remaining_turns would be [] for single-turn prompts
+
+        prompt_text = tokenizer.apply_chat_template(
+            prompt_msgs, tokenize=False, add_generation_prompt=True
+        )
+
+        # With empty remaining_turns, trainer falls through to raw concat
+        # (trainer.py:302: gt_texts.append(prompt_text + gt))
+        gt_text = prompt_text + gt
+        assert gt_text.endswith("A1")
+        assert prompt_text in gt_text
diff --git a/tests/test_empty_tokenization.py b/tests/test_empty_tokenization.py
new file mode 100644
index 0000000000..d1c568e9eb
--- /dev/null
+++ b/tests/test_empty_tokenization.py
@@ -0,0 +1,41 @@
+"""Regression tests: prompt-strategy ``_tokenize`` must not crash on an empty result.
+
+A field that tokenizes to nothing (e.g. an empty ``generation``/sub-field with a
+tokenizer that does not prepend BOS) used to raise ``IndexError`` in these overrides
+because they indexed ``result["input_ids"][-1]`` without the empty-guard the base
+``PromptTokenizingStrategy._tokenize`` already has.
+"""
+
+import pytest
+from transformers import AutoTokenizer
+
+from axolotl.prompt_strategies.metharme import MetharmePromptTokenizingStrategy
+from axolotl.prompt_tokenizers import ReflectionPromptTokenizingStrategy
+
+from tests.hf_offline_utils import enable_hf_offline
+
+
+@pytest.fixture(scope="module")
+@enable_hf_offline
+def no_bos_tokenizer(download_tiny_qwen3_model):
+    # Qwen3 sets add_bos_token=False, so tokenizing "" yields an empty input_ids list.
+    return AutoTokenizer.from_pretrained("axolotl-ai-co/tiny-qwen3-129m")
+
+
+def _build(strategy_cls, tokenizer):
+    strategy = strategy_cls.__new__(strategy_cls)
+    strategy.tokenizer = tokenizer
+    strategy.sequence_len = 2048
+    return strategy
+
+
+def test_metharme_tokenize_empty_does_not_crash(no_bos_tokenizer):
+    strategy = _build(MetharmePromptTokenizingStrategy, no_bos_tokenizer)
+    result = strategy._tokenize("")  # raised IndexError before the fix
+    assert list(result["input_ids"]) == []
+
+
+def test_reflection_tokenize_empty_does_not_crash(no_bos_tokenizer):
+    strategy = _build(ReflectionPromptTokenizingStrategy, no_bos_tokenizer)
+    result = strategy._tokenize("")  # raised IndexError before the fix
+    assert list(result["input_ids"]) == []
diff --git a/tests/test_exact_deduplication.py b/tests/test_exact_deduplication.py
new file mode 100644
index 0000000000..526ef50ac8
--- /dev/null
+++ b/tests/test_exact_deduplication.py
@@ -0,0 +1,426 @@
+"""Test suite for functions in the `axolotl.utils.data.utils` module, focusing on the
+`deduplicate_and_log_datasets` function.
+
+Additionally, this test suite includes tests for functions that indirectly call
+`deduplicate_and_log_datasets` during the execution of the preprocess command.
+"""
+
+import unittest
+from unittest.mock import patch
+
+import pytest
+from datasets import Dataset
+
+from axolotl.loaders import load_processor, load_tokenizer
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.data import prepare_datasets, prepare_preference_datasets
+from axolotl.utils.data.utils import deduplicate_and_log_datasets
+from axolotl.utils.dict import DictDefault
+
+from tests.constants import ALPACA_MESSAGES_CONFIG_REVISION, alpaca_messages_dpo_rows
+from tests.hf_offline_utils import enable_hf_offline
+
+
+def verify_deduplication(actual_dataset, expected_dataset, dataset_name):
+    """Validates deduplication results and size consistency.
+
+    Parameters:
+    - actual_dataset: Deduplicated dataset.
+    - expected_dataset: Expected dataset.
+    - dataset_name: Name of the dataset (e.g., 'train' or 'eval').
+
+    Asserts:
+    - Datasets match in content.
+    - Dataset size matches unique row count.
+    """
+    # Convert datasets to sets of tuples for unordered comparison
+    actual_rows = set(tuple(row.values()) for row in actual_dataset)
+    expected_rows = set(tuple(row.values()) for row in expected_dataset)
+
+    # Verify deduplication correctness
+    assert actual_rows == expected_rows, f"Mismatch in {dataset_name} dataset"
+
+    # Verify size consistency
+    assert len(actual_rows) == len(actual_dataset), (
+        f"Size mismatch in {dataset_name} dataset after deduplication"
+    )
+
+
+class TestDeduplicateIndividualFunctions(unittest.TestCase):
+    """Test class for deduplication function in data utils"""
+
+    def setUp(self):
+        # Sample data with duplicates
+        self.data = {
+            "column1": ["apple", "banana", "apple", "orange", "banana"],
+            "column2": [1, 2, 1, 3, 2],
+            "column3": ["red", "yellow", "red", "orange", "yellow"],
+        }
+
+        # Expected result after deduplication
+        self.expected_data = {
+            "column1": ["apple", "banana", "orange"],
+            "column2": [1, 2, 3],
+            "column3": ["red", "yellow", "orange"],
+        }
+
+        # Convert to Dataset format
+        self.dataset = Dataset.from_dict(self.data)
+        self.expected_dataset = Dataset.from_dict(self.expected_data)
+
+    def test_deduplication(self):
+        train_dataset, _ = deduplicate_and_log_datasets(dataset=self.dataset)
+        eval_dataset, _ = deduplicate_and_log_datasets(
+            dataset=self.dataset, dataset_name="eval"
+        )
+
+        verify_deduplication(train_dataset, self.expected_dataset, "train_dataset")
+        verify_deduplication(eval_dataset, self.expected_dataset, "eval_dataset")
+
+    def test_exact_duplicates(self):
+        # Test when datasets are exact duplicates
+        duplicate_data = {
+            "column1": ["apple", "apple", "apple"],
+            "column2": [1, 1, 1],
+            "column3": ["red", "red", "red"],
+        }
+        expected_data = {"column1": ["apple"], "column2": [1], "column3": ["red"]}
+
+        # Convert to Dataset format
+        dataset = Dataset.from_dict(duplicate_data)
+        expected_dataset = Dataset.from_dict(expected_data)
+
+        # Run deduplication
+        train_dataset, _ = deduplicate_and_log_datasets(dataset=dataset)
+        eval_dataset, _ = deduplicate_and_log_datasets(
+            dataset=dataset, dataset_name="eval"
+        )
+
+        verify_deduplication(train_dataset, expected_dataset, "train_dataset")
+        verify_deduplication(eval_dataset, expected_dataset, "eval_dataset")
+
+    def test_partial_duplicates(self):
+        # Test when only part of the dataset is a duplicate
+        partial_duplicate_data = {
+            "column1": ["apple", "banana", "apple"],
+            "column2": [1, 2, 1],
+            "column3": ["red", "yellow", "red"],
+        }
+        expected_data = {
+            "column1": ["apple", "banana"],
+            "column2": [1, 2],
+            "column3": ["red", "yellow"],
+        }
+
+        # Convert to Dataset format
+        dataset = Dataset.from_dict(partial_duplicate_data)
+        expected_dataset = Dataset.from_dict(expected_data)
+
+        # Run deduplication
+        train_dataset, _ = deduplicate_and_log_datasets(dataset=dataset)
+        eval_dataset, _ = deduplicate_and_log_datasets(
+            dataset=dataset, dataset_name="eval"
+        )
+
+        verify_deduplication(train_dataset, expected_dataset, "train_dataset")
+        verify_deduplication(eval_dataset, expected_dataset, "eval_dataset")
+
+    def test_combined_duplicates_empty(self):
+        # Test when only part of the dataset is a duplicate
+        partial_duplicate_data = {
+            "column1": ["apple", "banana", "apple"],
+            "column2": [1, 2, 1],
+            "column3": ["red", "yellow", "red"],
+        }
+        expected_data_train = {
+            "column1": ["apple", "banana"],
+            "column2": [1, 2],
+            "column3": ["red", "yellow"],
+        }
+        expected_data_eval = {
+            "column1": [],
+            "column2": [],
+            "column3": [],
+        }
+
+        # Convert to Dataset format
+        dataset = Dataset.from_dict(partial_duplicate_data)
+        expected_dataset_train = Dataset.from_dict(expected_data_train)
+        expected_dataset_eval = Dataset.from_dict(expected_data_eval)
+
+        # Run deduplication
+        train_dataset, eval_dataset = deduplicate_and_log_datasets(
+            dataset=dataset, other_dataset=dataset
+        )
+
+        verify_deduplication(train_dataset, expected_dataset_train, "train_dataset")
+        verify_deduplication(eval_dataset, expected_dataset_eval, "eval_dataset")
+
+    def test_combined_duplicates_one(self):
+        # Test when only part of the dataset is a duplicate
+        partial_duplicate_data_train = {
+            "column1": ["apple", "banana", "apple"],
+            "column2": [1, 2, 1],
+            "column3": ["red", "yellow", "red"],
+        }
+        partial_duplicate_data_eval = {
+            "column1": ["apple", "orange", "apple"],
+            "column2": [1, 2, 1],
+            "column3": ["red", "orange", "red"],
+        }
+        expected_data_train = {
+            "column1": ["apple", "banana"],
+            "column2": [1, 2],
+            "column3": ["red", "yellow"],
+        }
+        expected_data_eval = {
+            "column1": ["orange"],
+            "column2": [2],
+            "column3": ["orange"],
+        }
+
+        # Convert to Dataset format
+        dataset_train = Dataset.from_dict(partial_duplicate_data_train)
+        dataset_eval = Dataset.from_dict(partial_duplicate_data_eval)
+        expected_dataset_train = Dataset.from_dict(expected_data_train)
+        expected_dataset_eval = Dataset.from_dict(expected_data_eval)
+
+        # Run deduplication
+        train_dataset, eval_dataset = deduplicate_and_log_datasets(
+            dataset=dataset_train, other_dataset=dataset_eval
+        )
+
+        verify_deduplication(train_dataset, expected_dataset_train, "train_dataset")
+        verify_deduplication(eval_dataset, expected_dataset_eval, "eval_dataset")
+
+
+class TestDeduplicateRLDataset:
+    """Test a configured dataloader with deduplication."""
+
+    @pytest.fixture
+    def cfg(self):
+        fixture = DictDefault(
+            {
+                "tokenizer_config": "huggyllama/llama-7b",
+                "sequence_len": 1024,
+                "rl": "dpo",
+                "chat_template": "llama3",
+                "dataset_exact_deduplication": True,
+                "datasets": [
+                    ALPACA_MESSAGES_CONFIG_REVISION,
+                    ALPACA_MESSAGES_CONFIG_REVISION,
+                ],
+                "dataset_num_proc": None,
+            }
+        )
+        yield fixture
+
+    @enable_hf_offline
+    def test_load_with_deduplication(
+        self,
+        cfg,
+        tokenizer_huggyllama,
+    ):
+        """Verify that loading with deduplication removes duplicates."""
+        dataset = Dataset.from_list(alpaca_messages_dpo_rows())
+
+        with (
+            patch(
+                "axolotl.utils.data.rl.load_dataset_with_config"
+            ) as mock_load_dataset,
+            patch("axolotl.loaders.load_tokenizer") as mock_load_tokenizer,
+        ):
+            # Set up the mock to return different values on successive calls
+            mock_load_dataset.side_effect = [dataset, dataset]
+            mock_load_tokenizer.return_value = tokenizer_huggyllama
+
+            tokenizer = load_tokenizer(cfg)
+            train_dataset, _ = prepare_preference_datasets(cfg, tokenizer)
+
+            # Verify that the dataset has been deduplicated
+            assert len(train_dataset) == len(dataset), (
+                "Dataset was not properly deduplicated"
+            )
+
+    @enable_hf_offline
+    def test_load_without_deduplication(
+        self,
+        cfg,
+        tokenizer_huggyllama,
+    ):
+        dataset = Dataset.from_list(alpaca_messages_dpo_rows())
+
+        with (
+            patch(
+                "axolotl.utils.data.rl.load_dataset_with_config"
+            ) as mock_load_dataset,
+            patch("axolotl.loaders.load_tokenizer") as mock_load_tokenizer,
+        ):
+            # Set up the mock to return different values on successive calls
+            mock_load_dataset.side_effect = [dataset, dataset]
+            mock_load_tokenizer.return_value = tokenizer_huggyllama
+
+            # Load the dataset without deduplication
+            cfg.dataset_exact_deduplication = False
+            tokenizer = load_tokenizer(cfg)
+            train_dataset, _ = prepare_preference_datasets(cfg, tokenizer)
+
+            # Verify that the dataset retains duplicates
+            assert len(train_dataset) == len(dataset) * 2, (
+                "Dataset deduplication occurred when it should not have"
+            )
+
+
+class TestDeduplicateNonRL(unittest.TestCase):
+    """Test prepare_dataset function with different configurations."""
+
+    @enable_hf_offline
+    def setUp(self) -> None:
+        self.cfg_1 = DictDefault(
+            {
+                "base_model": "huggyllama/llama-7b",
+                "tokenizer_config": "huggyllama/llama-7b",
+                "sequence_len": 1024,
+                "dataset_exact_deduplication": True,
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "val_set_size": 0.0,
+                "gradient_accumulation_steps": 2,
+                "batch_size": 10,
+                "micro_batch_size": 10,
+                "num_epochs": 1,
+            }
+        )
+        self.cfg_1 = validate_config(self.cfg_1)
+        normalize_config(self.cfg_1)
+
+    @pytest.mark.skip(reason="TODO: fix hf hub offline to work with HF rate limits")
+    @enable_hf_offline
+    def test_prepare_dataset_with_deduplication_train(self):
+        """Verify that prepare_dataset function processes the dataset correctly with deduplication."""
+        self.cfg_1.dataset_exact_deduplication = True
+
+        # Load tokenizer and processor
+        tokenizer = load_tokenizer(self.cfg_1)
+        processor = (
+            load_processor(self.cfg_1, tokenizer=tokenizer)
+            if self.cfg_1.processor_type
+            else None
+        )
+
+        # Prepare dataset using the prepare_dataset function
+        train_dataset, _, _, _ = prepare_datasets(
+            self.cfg_1,
+            tokenizer,
+            processor=processor,
+        )
+
+        self.assertEqual(
+            len(train_dataset),
+            2000,
+            "Train dataset should have 2000 samples after deduplication.",
+        )
+
+    @pytest.mark.skip(reason="TODO: fix hf hub offline to work with HF rate limits")
+    @enable_hf_offline
+    def test_prepare_dataset_with_deduplication_eval(self):
+        """Verify that prepare_dataset function processes the dataset correctly with deduplication."""
+        self.cfg_1.dataset_exact_deduplication = True
+        self.cfg_1.val_set_size = 0.5
+        # Load tokenizer and processor
+        tokenizer = load_tokenizer(self.cfg_1)
+        processor = (
+            load_processor(self.cfg_1, tokenizer=tokenizer)
+            if self.cfg_1.processor_type
+            else None
+        )
+
+        # Prepare dataset using the prepare_dataset function
+        _, eval_dataset, _, _ = prepare_datasets(
+            self.cfg_1,
+            tokenizer,
+            processor=processor,
+        )
+
+        self.assertEqual(
+            len(eval_dataset),
+            1000,
+            "Eval dataset should have 2000 samples after deduplication.",
+        )
+
+    @pytest.mark.skip(reason="TODO: fix hf hub offline to work with HF rate limits")
+    @enable_hf_offline
+    def test_prepare_dataset_without_deduplication(self):
+        """Verify that prepare_dataset function processes the dataset correctly without deduplication."""
+        self.cfg_1.dataset_exact_deduplication = False
+        self.cfg_1.val_set_size = 0.1
+        # Load tokenizer and processor
+        tokenizer = load_tokenizer(self.cfg_1)
+        processor = (
+            load_processor(self.cfg_1, tokenizer=tokenizer)
+            if self.cfg_1.processor_type
+            else None
+        )
+
+        # Prepare dataset using the prepare_dataset function
+        train_dataset, eval_dataset, _, _ = prepare_datasets(
+            self.cfg_1,
+            tokenizer,
+            processor=processor,
+        )
+
+        # Verify that the dataset has been prepared correctly
+        self.assertEqual(
+            len(train_dataset),
+            1800 * 2,
+            "Train dataset should have 3600 samples without deduplication.",
+        )
+        self.assertEqual(
+            len(eval_dataset),
+            200 * 2,
+            "Train dataset should have 400 samples after deduplication.",
+        )
+
+
+class TestWrongCollisions(unittest.TestCase):
+    """Creating mock datasets for testing wrong collisions."""
+
+    def setUp(self):
+        self.train_data = {"text": ["sample 5", "sample 6"], "label": [1, 2]}
+        self.eval_data = {
+            "text": [
+                "sample 5",
+                "sample 7",
+            ],  # Different label but same text as in train_data
+            "label": [2, 3],
+        }
+        self.dataset_data = {
+            "text": ["sample 5", "sample 9", "sample 5"],
+            "label": [1, 2, 8],
+        }
+        self.train_dataset = Dataset.from_dict(self.train_data)
+        self.eval_dataset = Dataset.from_dict(self.eval_data)
+        self.dataset = Dataset.from_dict(self.dataset_data)
+
+    def test_deduplication_dataset_only(self):
+        dedup_dataset, _ = deduplicate_and_log_datasets(dataset=self.dataset)
+        self.assertEqual(
+            len(dedup_dataset), 3, "Dataset should have all original values"
+        )
+        self.assertEqual(
+            str(dedup_dataset),
+            str(self.dataset),
+            "The string representation of the output dataset should not differ.",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_expand_mask.py b/tests/test_expand_mask.py
deleted file mode 100644
index 01241c2958..0000000000
--- a/tests/test_expand_mask.py
+++ /dev/null
@@ -1,44 +0,0 @@
-"""
-Unit tests for the monkey patch for expand mask to handle packed sequences
-"""
-import unittest
-
-import torch
-
-from axolotl.monkeypatch.llama_expand_mask import _expand_mask
-
-
-class TestExpandMask(unittest.TestCase):
-    """
-    Test class for attention mask expansion for packed sequences
-    """
-
-    def test_output(self):
-        mask = torch.tensor([[1, 1, 1, 2], [2, 3, 3, 0]])
-        dtype = torch.float32
-        expected_output = torch.tensor(
-            [
-                [
-                    [
-                        [0.0000e00, -3.4028e38, -3.4028e38, -3.4028e38],
-                        [0.0000e00, 0.0000e00, -3.4028e38, -3.4028e38],
-                        [0.0000e00, 0.0000e00, 0.0000e00, -3.4028e38],
-                        [-3.4028e38, -3.4028e38, -3.4028e38, 0.0000e00],
-                    ]
-                ],
-                [
-                    [
-                        [0.0000e00, -3.4028e38, -3.4028e38, -3.4028e38],
-                        [-3.4028e38, 0.0000e00, -3.4028e38, -3.4028e38],
-                        [-3.4028e38, 0.0000e00, 0.0000e00, -3.4028e38],
-                        [-3.4028e38, -3.4028e38, -3.4028e38, -3.4028e38],
-                    ]
-                ],
-            ]
-        )
-        # Check that the output matches the expected output
-        self.assertTrue(torch.allclose(_expand_mask(mask, dtype), expected_output))
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/tests/test_fp32_norms.py b/tests/test_fp32_norms.py
new file mode 100644
index 0000000000..e2b477fe8e
--- /dev/null
+++ b/tests/test_fp32_norms.py
@@ -0,0 +1,264 @@
+"""Unit tests for fp32 norm sharding (FSDP2). Pure-CPU, no dist init."""
+
+from __future__ import annotations
+
+import logging
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.distributed.fsdp import MixedPrecisionPolicy
+
+from axolotl.loaders.model import ModelLoader
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.fp32_norms import (
+    DEFAULT_FP32_NORM_SUFFIXES,
+    _matches_norm_class,
+    shard_norms_fp32,
+)
+
+
+class LlamaRMSNorm(nn.Module):
+    def __init__(self, dim: int = 8) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+
+class AfmoeRMSNorm(nn.Module):
+    def __init__(self, dim: int = 8) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+
+class CustomNorm(nn.Module):
+    def __init__(self, dim: int = 8) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+
+class CustomNormWithBuffer(nn.Module):
+    def __init__(self, dim: int = 8) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.bfloat16))
+        self.register_buffer("running_scale", torch.ones(dim, dtype=torch.float16))
+
+
+class MLP(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.fc = nn.Linear(8, 8)
+
+
+def test_suffix_matches_multiple_norm_families():
+    patterns = list(DEFAULT_FP32_NORM_SUFFIXES)
+    assert _matches_norm_class(LlamaRMSNorm(), patterns)
+    assert _matches_norm_class(AfmoeRMSNorm(), patterns)
+    assert _matches_norm_class(nn.LayerNorm(8), patterns)
+
+
+def test_suffix_does_not_match_non_norm_modules():
+    patterns = list(DEFAULT_FP32_NORM_SUFFIXES)
+    assert not _matches_norm_class(MLP(), patterns)
+    assert not _matches_norm_class(nn.Linear(8, 8), patterns)
+    assert not _matches_norm_class(CustomNorm(), patterns)
+
+
+def test_explicit_classname_matches_custom_norm():
+    assert _matches_norm_class(CustomNorm(), ["CustomNorm"])
+
+
+def test_fully_qualified_pattern_matches_exact_path():
+    qualified = f"{LlamaRMSNorm.__module__}.LlamaRMSNorm"
+    assert _matches_norm_class(LlamaRMSNorm(), [qualified])
+    assert not _matches_norm_class(AfmoeRMSNorm(), [qualified])
+
+
+def test_mixed_patterns_suffix_and_qualified():
+    qualified = f"{LlamaRMSNorm.__module__}.LlamaRMSNorm"
+    patterns = [qualified, "LayerNorm"]
+    assert _matches_norm_class(LlamaRMSNorm(), patterns)
+    assert _matches_norm_class(nn.LayerNorm(8), patterns)
+    assert not _matches_norm_class(AfmoeRMSNorm(), patterns)
+
+
+class _Cfg:
+    def __init__(self, **kwargs):
+        self.fp32_norms = kwargs.get("fp32_norms", False)
+        self.fp32_norm_classes = kwargs.get("fp32_norm_classes", None)
+        self.fsdp_version = kwargs.get("fsdp_version", None)
+        self.fsdp_config = kwargs.get("fsdp_config", None)
+        self.tensor_parallel_size = kwargs.get("tensor_parallel_size", 1)
+        self.lora_on_cpu = kwargs.get("lora_on_cpu", False)
+
+
+def test_disabled_is_noop():
+    model = nn.Sequential(LlamaRMSNorm(), MLP())
+    assert shard_norms_fp32(model, _Cfg(fp32_norms=False)) == 0
+
+
+def test_enabled_requires_fsdp2():
+    model = nn.Sequential(LlamaRMSNorm())
+    cfg = _Cfg(fp32_norms=True, fsdp_version=1)
+    with pytest.raises(ValueError, match="fsdp_version: 2"):
+        shard_norms_fp32(model, cfg)
+
+
+def test_meta_device_is_supported(monkeypatch):
+    with torch.device("meta"):
+        model = nn.Sequential(LlamaRMSNorm())
+    cfg = _Cfg(fp32_norms=True, fsdp_version=2)
+
+    import torch.distributed.fsdp as fsdp_module
+
+    calls = []
+
+    def fake_fully_shard(module, mp_policy=None, **kwargs):
+        calls.append((type(module).__name__, mp_policy, kwargs))
+        return module
+
+    monkeypatch.setattr(fsdp_module, "fully_shard", fake_fully_shard)
+
+    n = shard_norms_fp32(model, cfg)
+    assert n == 1
+    assert calls[0][0] == "LlamaRMSNorm"
+    assert calls[0][1].param_dtype == torch.float32
+
+
+def test_passthrough_fully_shard_kwargs_are_used(monkeypatch):
+    model = nn.Sequential(LlamaRMSNorm())
+    cfg = _Cfg(fp32_norms=True, fsdp_version=2)
+
+    import torch.distributed.fsdp as fsdp_module
+
+    calls = []
+
+    def fake_fully_shard(module, mp_policy=None, **kwargs):
+        calls.append((module, mp_policy, kwargs))
+        return module
+
+    monkeypatch.setattr(fsdp_module, "fully_shard", fake_fully_shard)
+
+    sentinel_mesh = object()
+    outer_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16)
+    n = shard_norms_fp32(
+        model,
+        cfg,
+        fully_shard_kwargs={
+            "mesh": sentinel_mesh,
+            "reshard_after_forward": True,
+            "mp_policy": outer_policy,
+        },
+    )
+    assert n == 1
+    assert calls[0][2]["mesh"] is sentinel_mesh
+    assert calls[0][2]["reshard_after_forward"] is True
+    assert calls[0][1].output_dtype == torch.bfloat16
+
+
+def test_no_matches_warns_and_returns_zero(caplog):
+    model = nn.Sequential(MLP(), nn.Linear(8, 8))
+    cfg = _Cfg(fp32_norms=True, fsdp_version=2)
+    # axolotl.cli.configure_logging() sets propagate=False on the `axolotl`
+    # logger, so pytest caplog can't see records by default. Temporarily
+    # re-enable propagation for this assertion.
+    ax_logger = logging.getLogger("axolotl")
+    old_propagate = ax_logger.propagate
+    ax_logger.propagate = True
+    try:
+        with caplog.at_level("WARNING", logger="axolotl"):
+            n = shard_norms_fp32(model, cfg)
+    finally:
+        ax_logger.propagate = old_propagate
+    assert n == 0
+    assert "no modules matched" in caplog.text
+
+
+def test_explicit_classes_override_defaults(monkeypatch):
+    model = nn.Sequential(CustomNorm(), MLP())
+    cfg = _Cfg(fp32_norms=True, fsdp_version=2, fp32_norm_classes=["CustomNorm"])
+
+    import torch.distributed.fsdp as fsdp_module
+
+    calls = []
+
+    def fake_fully_shard(module, mp_policy=None, **_):
+        calls.append((type(module).__name__, mp_policy))
+        return module
+
+    monkeypatch.setattr(fsdp_module, "fully_shard", fake_fully_shard)
+
+    n = shard_norms_fp32(model, cfg)
+    assert n == 1
+    assert calls[0][0] == "CustomNorm"
+    assert calls[0][1].param_dtype == torch.float32
+    assert calls[0][1].reduce_dtype == torch.float32
+
+
+def test_matched_norm_storage_is_cast_to_fp32_before_sharding(monkeypatch):
+    model = nn.Sequential(CustomNormWithBuffer(), MLP())
+    cfg = _Cfg(
+        fp32_norms=True,
+        fsdp_version=2,
+        fp32_norm_classes=["CustomNormWithBuffer"],
+    )
+
+    import torch.distributed.fsdp as fsdp_module
+
+    seen = []
+
+    def fake_fully_shard(module, mp_policy=None, **kwargs):
+        seen.append(
+            (
+                module.weight.dtype,
+                module.running_scale.dtype,
+                mp_policy.output_dtype,
+                kwargs,
+            )
+        )
+        return module
+
+    monkeypatch.setattr(fsdp_module, "fully_shard", fake_fully_shard)
+
+    outer_policy = MixedPrecisionPolicy(param_dtype=torch.float16)
+    n = shard_norms_fp32(
+        model,
+        cfg,
+        fully_shard_kwargs={"mp_policy": outer_policy},
+    )
+
+    assert n == 1
+    assert model[0].weight.dtype == torch.float32
+    assert model[0].running_scale.dtype == torch.float32
+    assert seen[0][0] == torch.float32
+    assert seen[0][1] == torch.float32
+    assert seen[0][2] == torch.float16
+
+
+class TinyModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.embed_tokens = nn.Embedding(8, 8)
+        self.input_norm = CustomNorm()
+        self.fc = nn.Linear(8, 8)
+
+
+def test_convert_embedding_modules_dtype_keeps_fp32_norm_matches():
+    loader = ModelLoader.__new__(ModelLoader)
+    loader.cfg = DictDefault(
+        fp32_norms=True,
+        fp32_norm_classes=["CustomNorm"],
+        lora_on_cpu=False,
+    )
+    loader.model = TinyModel()
+    loader.model_config = SimpleNamespace(model_type="llama")
+
+    loader._convert_embedding_modules_dtype(
+        embedding_modules=["embed_tokens"],
+        dist_dtype=torch.bfloat16,
+        before_kbit_train_or_finetune=False,
+    )
+
+    assert loader.model.input_norm.weight.dtype == torch.float32
+    assert loader.model.embed_tokens.weight.dtype == torch.bfloat16
+    assert loader.model.fc.weight.dtype == torch.float32
diff --git a/tests/test_glm_dsa_cp_validation.py b/tests/test_glm_dsa_cp_validation.py
new file mode 100644
index 0000000000..fcbdb8b8b1
--- /dev/null
+++ b/tests/test_glm_dsa_cp_validation.py
@@ -0,0 +1,116 @@
+"""Validation tests for the GLM DSA kernels' context-parallel exemptions.
+
+``context_parallel_size > 1`` normally requires flash attention + ``ring_flash_attn``. The GLM DSA
+kernels own context-parallel attention (compressed-KV all-gather + per-rank ``q_offset``), so when
+``use_glm_dsa_kernels`` is set the validator skips that requirement and leaves ``ring_attn_func`` None
+(the SP context manager still shards the sequence by chunking). These are config-validation checks only
+-- no GPU / dist / model needed. ``ring_flash_attn`` is intentionally NOT mocked here: the DSA path must
+validate without it installed.
+"""
+
+import pytest
+
+from axolotl.utils.config import prepare_plugins, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+def _cfg(**extra):
+    return DictDefault(
+        base_model="HuggingFaceTB/SmolLM2-135M",
+        learning_rate=1e-3,
+        datasets=[{"path": "mhenrichsen/alpaca_2k_test", "type": "alpaca"}],
+        micro_batch_size=1,
+        gradient_accumulation_steps=1,
+        sequence_len=2048,
+        **extra,
+    )
+
+
+def test_config_validation_recovers_inherited_defaults():
+    from pydantic import BaseModel
+
+    from axolotl.utils.config import _model_with_inherited_default_fallback
+
+    class _Base(BaseModel):
+        base_model: str
+        strict: bool = False
+
+    class _Merged(_Base):
+        strict: bool
+
+    cfg = _model_with_inherited_default_fallback(
+        _Merged, {"base_model": "HuggingFaceTB/SmolLM2-135M"}
+    )
+
+    assert cfg.strict is False
+
+
+def test_config_validation_recovers_known_none_defaults():
+    from pydantic import BaseModel
+
+    from axolotl.utils.config import _model_with_inherited_default_fallback
+
+    class _Merged(BaseModel):
+        base_model: str
+        xformers_attention: bool | None
+
+    cfg = _model_with_inherited_default_fallback(
+        _Merged, {"base_model": "HuggingFaceTB/SmolLM2-135M"}
+    )
+
+    assert cfg.xformers_attention is None
+
+
+def test_config_validation_retry_accepts_field_names():
+    from pydantic import BaseModel, ConfigDict, Field
+
+    from axolotl.utils.config import _model_with_inherited_default_fallback
+
+    class _Merged(BaseModel):
+        model_config = ConfigDict(validate_by_name=False, validate_by_alias=True)
+
+        base_model: str
+        xformers_attention: bool | None = Field(alias="xformersAttention")
+
+    cfg = _model_with_inherited_default_fallback(
+        _Merged, {"base_model": "HuggingFaceTB/SmolLM2-135M"}
+    )
+
+    assert cfg.xformers_attention is None
+
+
+class TestGlmDsaContextParallelValidation:
+    """The use_glm_dsa_kernels exemptions in check_context_parallel_size / validate_ring_attn_func."""
+
+    def test_dsa_cp_skips_flash_and_ring_requirement(self, monkeypatch):
+        """use_glm_dsa_kernels + context_parallel_size>1 validates WITHOUT flash attention and WITHOUT
+        ring_flash_attn installed -- the DSA kernels own CP attention."""
+        monkeypatch.setenv("WORLD_SIZE", "4")
+        cfg = _cfg(
+            plugins=["axolotl.integrations.kernels.KernelsPlugin"],
+            use_glm_dsa_kernels=True,
+            context_parallel_size=2,
+        )
+        prepare_plugins(cfg)
+        out = validate_config(cfg)  # must not raise (no flash, no ring_flash_attn)
+        # ring attention is NOT substituted: ring_attn_func stays None so the SP context manager only
+        # shards the sequence (register_ring_attn skips the ring_flash_attn import when it is None).
+        assert out.ring_attn_func is None
+
+    def test_cp_without_dsa_still_requires_flash(self, monkeypatch):
+        """The exemption is scoped to use_glm_dsa_kernels -- plain CP still demands flash attention."""
+        monkeypatch.setenv("WORLD_SIZE", "4")
+        cfg = _cfg(context_parallel_size=2)  # no DSA kernels, no flash_attention
+        with pytest.raises(Exception, match="(?i)flash attention"):
+            validate_config(cfg)
+
+    def test_dsa_without_cp_leaves_ring_attn_func_none(self, monkeypatch):
+        """use_glm_dsa_kernels with context_parallel_size 1 is a no-op for the CP validators."""
+        monkeypatch.setenv("WORLD_SIZE", "1")
+        cfg = _cfg(
+            plugins=["axolotl.integrations.kernels.KernelsPlugin"],
+            use_glm_dsa_kernels=True,
+        )
+        prepare_plugins(cfg)
+        out = validate_config(cfg)
+        assert out.ring_attn_func is None
diff --git a/tests/test_http_weight_sync.py b/tests/test_http_weight_sync.py
new file mode 100644
index 0000000000..97c5ece71a
--- /dev/null
+++ b/tests/test_http_weight_sync.py
@@ -0,0 +1,158 @@
+"""Tests for HTTP weight sync serialization round-trip (bf16/fp16/fp32).
+
+Exercises the encode/decode helpers in axolotl.utils.weight_serde that handle
+the three-stage weight transfer: trainer → serve endpoint → vLLM worker.
+"""
+
+import pytest
+import torch
+
+from axolotl.utils.weight_serde import (
+    decode_from_http,
+    decode_from_ipc,
+    encode_for_http,
+    encode_for_ipc,
+)
+
+# ---------------------------------------------------------------------------
+# Stage 1: trainer → serve endpoint (HTTP with base64)
+# ---------------------------------------------------------------------------
+
+
+class TestHttpEncodeRoundTrip:
+    """Test encode_for_http / decode_from_http."""
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_round_trip_dtype(self, dtype):
+        original = torch.randn(32, 64, dtype=dtype)
+        entry = encode_for_http("layer.weight", original)
+        name, decoded = decode_from_http(entry)
+
+        assert name == "layer.weight"
+        assert decoded.dtype == dtype
+        assert decoded.shape == original.shape
+        if dtype == torch.bfloat16:
+            # bf16→fp16→bf16 loses some precision
+            torch.testing.assert_close(decoded, original, atol=1e-2, rtol=1e-2)
+        else:
+            torch.testing.assert_close(decoded, original, atol=0, rtol=0)
+
+    def test_bfloat16_wire_format_is_fp16(self):
+        """bf16 tensors should be sent as fp16 on the wire."""
+        import base64
+
+        original = torch.randn(8, 16, dtype=torch.bfloat16)
+        entry = encode_for_http("w", original)
+        raw = base64.b64decode(entry["data"])
+        # 8*16 elements * 2 bytes/elem (fp16) = 256 bytes
+        assert len(raw) == 8 * 16 * 2
+        # dtype field should preserve original dtype for reconstruction
+        assert entry["dtype"] == "torch.bfloat16"
+
+    def test_multidimensional_shapes(self):
+        for shape in [(128,), (4, 32), (2, 3, 16), (2, 2, 2, 8)]:
+            original = torch.randn(*shape, dtype=torch.bfloat16)
+            entry = encode_for_http("w", original)
+            _, decoded = decode_from_http(entry)
+            assert decoded.shape == original.shape
+            assert decoded.dtype == torch.bfloat16
+
+
+# ---------------------------------------------------------------------------
+# Stage 2: serve endpoint → vLLM worker (IPC with raw bytes)
+# ---------------------------------------------------------------------------
+
+
+class TestIpcEncodeRoundTrip:
+    """Test encode_for_ipc / decode_from_ipc."""
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_round_trip_dtype(self, dtype):
+        original = torch.randn(32, 64, dtype=dtype)
+        entry = encode_for_ipc("layer.weight", original)
+        name, decoded = decode_from_ipc(entry)
+
+        assert name == "layer.weight"
+        assert decoded.dtype == dtype
+        assert decoded.shape == original.shape
+        if dtype == torch.bfloat16:
+            torch.testing.assert_close(decoded, original, atol=1e-2, rtol=1e-2)
+        else:
+            torch.testing.assert_close(decoded, original, atol=0, rtol=0)
+
+    def test_bfloat16_ipc_wire_is_fp16(self):
+        """bf16 tensors should be serialized as fp16 bytes in IPC."""
+        original = torch.randn(4, 8, dtype=torch.bfloat16)
+        entry = encode_for_ipc("w", original)
+        assert entry["dtype"] == "float16"
+        assert entry["target_dtype"] == "bfloat16"
+        assert len(entry["data"]) == 4 * 8 * 2  # fp16 bytes
+
+    def test_fp32_has_no_target_dtype_mismatch(self):
+        original = torch.randn(4, 8, dtype=torch.float32)
+        entry = encode_for_ipc("w", original)
+        assert entry["dtype"] == "float32"
+        assert entry["target_dtype"] == "float32"
+
+    def test_worker_handles_missing_target_dtype(self):
+        """Backward compat: older serve code may not send target_dtype."""
+        entry = {
+            "name": "w",
+            "data": torch.randn(4, 8, dtype=torch.float32).numpy().tobytes(),
+            "dtype": "float32",
+            "shape": [4, 8],
+            # no target_dtype key
+        }
+        name, decoded = decode_from_ipc(entry)
+        assert decoded.dtype == torch.float32
+        assert decoded.shape == (4, 8)
+
+
+# ---------------------------------------------------------------------------
+# Full pipeline: trainer → serve → worker
+# ---------------------------------------------------------------------------
+
+
+class TestFullPipelineRoundTrip:
+    """End-to-end: trainer → serve → worker."""
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_three_stage_round_trip(self, dtype):
+        """Tensor survives trainer→serve→worker with correct dtype and values."""
+        original = torch.randn(16, 32, dtype=dtype)
+
+        # Stage 1: trainer encodes for HTTP
+        http_entry = encode_for_http("model.layers.0.weight", original)
+
+        # Stage 2: serve decodes HTTP, re-encodes for IPC
+        name, at_serve = decode_from_http(http_entry)
+        ipc_entry = encode_for_ipc(name, at_serve)
+
+        # Stage 3: worker decodes IPC
+        _, at_worker = decode_from_ipc(ipc_entry)
+
+        assert at_worker.dtype == dtype
+        assert at_worker.shape == original.shape
+        if dtype == torch.bfloat16:
+            # Two bf16→fp16→bf16 hops compound precision loss slightly
+            torch.testing.assert_close(at_worker, original, atol=2e-2, rtol=2e-2)
+        else:
+            torch.testing.assert_close(at_worker, original, atol=0, rtol=0)
+
+    def test_bfloat16_precision_loss_is_bounded(self):
+        """bf16→fp16→bf16 round-trip error should be small."""
+        original = torch.randn(256, 256, dtype=torch.bfloat16)
+        http_entry = encode_for_http("w", original)
+        _, at_serve = decode_from_http(http_entry)
+        ipc_entry = encode_for_ipc("w", at_serve)
+        _, at_worker = decode_from_ipc(ipc_entry)
+
+        max_err = (at_worker.float() - original.float()).abs().max().item()
+        # bf16 has ~8e-3 precision, fp16 has ~1e-3; round-trip error bounded
+        assert max_err < 0.05, f"Max error {max_err} exceeds bound"
+
+    def test_bfloat16_numpy_would_crash_without_fix(self):
+        """Verify that calling .numpy() on bf16 raises, confirming the fix is needed."""
+        t = torch.randn(4, 4, dtype=torch.bfloat16)
+        with pytest.raises((RuntimeError, TypeError)):
+            t.numpy()
diff --git a/tests/test_loaders.py b/tests/test_loaders.py
new file mode 100644
index 0000000000..9130905667
--- /dev/null
+++ b/tests/test_loaders.py
@@ -0,0 +1,218 @@
+"""Module for `axolotl.loaders`."""
+
+from unittest.mock import MagicMock
+
+import pytest
+from transformers import BitsAndBytesConfig, PreTrainedTokenizerBase
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+from transformers.utils.import_utils import is_torch_mps_available
+
+from axolotl.loaders import ModelLoader
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.distributed import _get_parallel_config_kwargs
+
+
+class TestModelsUtils:
+    """Testing module for `axolotl.loaders`."""
+
+    def setup_method(self) -> None:
+        # load config
+        self.cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "model_type": "AutoModelForCausalLM",
+                "tokenizer_type": "AutoTokenizer",
+                "load_in_8bit": True,
+                "load_in_4bit": False,
+                "adapter": "lora",
+                "flash_attention": False,
+                "sample_packing": True,
+                "device_map": "auto",
+            }
+        )
+        self.tokenizer = MagicMock(spec=PreTrainedTokenizerBase)
+        self.inference = False
+        self.reference_model = True
+
+        # init ModelLoader
+        self.model_loader = ModelLoader(
+            cfg=self.cfg,
+            tokenizer=self.tokenizer,
+            inference=self.inference,
+            reference_model=self.reference_model,
+        )
+
+    def test_set_device_map_config(self):
+        # check device_map
+        device_map = self.cfg.device_map
+        if is_torch_mps_available():
+            device_map = "mps"
+
+        self.model_loader._set_device_map_config()
+        if is_deepspeed_zero3_enabled():
+            assert "device_map" not in self.model_loader.model_kwargs
+        else:
+            assert device_map in self.model_loader.model_kwargs["device_map"]
+
+        # check torch_dtype
+        assert self.cfg.torch_dtype == self.model_loader.model_kwargs["torch_dtype"]
+
+    @pytest.mark.parametrize("adapter", ["lora", "qlora", None])
+    @pytest.mark.parametrize("load_in_8bit", [True, False])
+    @pytest.mark.parametrize("load_in_4bit", [True, False])
+    @pytest.mark.parametrize("gptq", [True, False])
+    def test_set_quantization_config(
+        self,
+        adapter,
+        load_in_8bit,
+        load_in_4bit,
+        gptq,
+    ):
+        # init cfg as args
+        self.cfg.load_in_8bit = load_in_8bit
+        self.cfg.load_in_4bit = load_in_4bit
+        self.cfg.gptq = gptq
+        self.cfg.adapter = adapter
+
+        self.model_loader._set_quantization_config()
+        if "quantization_config" in self.model_loader.model_kwargs or self.cfg.gptq:
+            assert not (
+                hasattr(self.model_loader.model_kwargs, "load_in_8bit")
+                and hasattr(self.model_loader.model_kwargs, "load_in_4bit")
+            )
+
+        if self.cfg.adapter == "qlora" and load_in_4bit:
+            assert isinstance(
+                self.model_loader.model_kwargs.get("quantization_config"),
+                BitsAndBytesConfig,
+            )
+
+            assert (
+                self.model_loader.model_kwargs["quantization_config"]._load_in_4bit
+                is True
+            )
+        if self.cfg.adapter == "lora" and load_in_8bit:
+            assert isinstance(
+                self.model_loader.model_kwargs.get("quantization_config"),
+                BitsAndBytesConfig,
+            )
+
+            assert (
+                self.model_loader.model_kwargs["quantization_config"]._load_in_8bit
+                is True
+            )
+
+    def test_message_property_mapping(self):
+        """Test message property mapping configuration validation"""
+        from axolotl.utils.schemas.datasets import SFTDataset
+
+        # Test legacy fields are mapped orrectly
+        dataset = SFTDataset(
+            path="test_path",
+            message_field_role="role_field",
+            message_field_content="content_field",
+        )
+        assert dataset.message_property_mappings == {
+            "role": "role_field",
+            "content": "content_field",
+        }
+
+        # Test direct message_property_mapping works
+        dataset = SFTDataset(
+            path="test_path",
+            message_property_mappings={
+                "role": "custom_role",
+                "content": "custom_content",
+            },
+        )
+        assert dataset.message_property_mappings == {
+            "role": "custom_role",
+            "content": "custom_content",
+        }
+
+        # Test both legacy and new fields work when they match
+        dataset = SFTDataset(
+            path="test_path",
+            message_field_role="same_role",
+            message_property_mappings={"role": "same_role"},
+        )
+        assert dataset.message_property_mappings == {
+            "role": "same_role",
+            "content": "content",
+        }
+
+        # Test both legacy and new fields work when they don't overlap
+        dataset = SFTDataset(
+            path="test_path",
+            message_field_role="role_field",
+            message_property_mappings={"content": "content_field"},
+        )
+        assert dataset.message_property_mappings == {
+            "role": "role_field",
+            "content": "content_field",
+        }
+
+        # Test no role or content provided
+        dataset = SFTDataset(
+            path="test_path",
+        )
+        assert dataset.message_property_mappings == {
+            "role": "role",
+            "content": "content",
+        }
+
+        # Test error when legacy and new fields conflict
+        with pytest.raises(ValueError) as exc_info:
+            SFTDataset(
+                path="test_path",
+                message_field_role="legacy_role",
+                message_property_mappings={"role": "different_role"},
+            )
+        assert "Conflicting message role fields" in str(exc_info.value)
+
+        with pytest.raises(ValueError) as exc_info:
+            SFTDataset(
+                path="test_path",
+                message_field_content="legacy_content",
+                message_property_mappings={"content": "different_content"},
+            )
+        assert "Conflicting message content fields" in str(exc_info.value)
+
+    @pytest.mark.parametrize(
+        "world_size, tensor_parallel_size, context_parallel_size, dp_shard_size, dp_replicate_size, is_fsdp, expected",
+        [
+            (16, 2, 2, 2, 2, True, (2, 2, 2, 2)),
+            (16, 1, 1, None, None, True, (0, 0, 16, 1)),
+            (16, 2, 2, 2, None, True, (2, 2, 2, 2)),
+            (16, 2, 2, None, 2, True, (2, 2, 2, 2)),
+            (16, 1, 1, None, 2, True, (0, 0, 8, 2)),
+            (2, 1, 1, None, None, True, (0, 0, 2, 1)),
+        ],
+    )
+    def test_get_parallel_config_kwargs(
+        self,
+        world_size,
+        tensor_parallel_size,
+        context_parallel_size,
+        dp_shard_size,
+        dp_replicate_size,
+        is_fsdp,
+        expected,
+    ):
+        res = _get_parallel_config_kwargs(
+            world_size,
+            tensor_parallel_size,
+            context_parallel_size,
+            dp_shard_size,
+            dp_replicate_size,
+            is_fsdp,
+        )
+
+        if expected[0] > 1:
+            assert res["tp_size"] == expected[0]
+        if expected[1] > 1:
+            assert res["cp_size"] == expected[1]
+        if expected[2] > 1:
+            assert res["dp_shard_size"] == expected[2]
+        if expected[3] > 1:
+            assert res["dp_replicate_size"] == expected[3]
diff --git a/tests/test_logging_config_file_capture.py b/tests/test_logging_config_file_capture.py
new file mode 100644
index 0000000000..11e8d7c645
--- /dev/null
+++ b/tests/test_logging_config_file_capture.py
@@ -0,0 +1,128 @@
+import logging
+import tempfile
+
+import pytest
+
+
+def read(path: str) -> str:
+    with open(path, "r", encoding="utf-8") as f:
+        return f.read()
+
+
+@pytest.fixture(autouse=True)
+def _reset_logging_state():
+    # Ensure a clean slate for logging between tests
+    for handler in logging.root.handlers[:]:
+        logging.root.removeHandler(handler)
+    logging.shutdown()
+    # Note: dictConfig in configure_logging will set up handlers again
+    yield
+    for handler in logging.root.handlers[:]:
+        logging.root.removeHandler(handler)
+    logging.shutdown()
+
+
+def test_axolotl_logs_captured_at_all_levels(monkeypatch):
+    from axolotl.logging_config import configure_logging
+    from axolotl.utils import tee
+    from axolotl.utils.logging import get_logger
+
+    with tempfile.TemporaryDirectory() as td:
+        # Avoid stdout tee in this test to simplify interaction with pytest capture
+        monkeypatch.setenv("AXOLOTL_TEE_STDOUT", "0")
+        configure_logging()
+        path = tee.prepare_debug_log(
+            type("Cfg", (), {"output_dir": td, "get": lambda *_: False})
+        )
+
+        log = get_logger("axolotl.test")
+        log.info("AX-INFO")
+        log.debug("AX-DEBUG")
+        tee.file_only_stream.flush()
+
+        data = read(path)
+        assert "AX-INFO" in data
+        assert "AX-DEBUG" in data
+        tee.close_debug_log()
+
+
+def test_third_party_logs_filtered_and_warning_captured(monkeypatch):
+    from axolotl.logging_config import configure_logging
+    from axolotl.utils import tee
+
+    with tempfile.TemporaryDirectory() as td:
+        monkeypatch.setenv("AXOLOTL_TEE_STDOUT", "0")
+        configure_logging()
+        path = tee.prepare_debug_log(
+            type("Cfg", (), {"output_dir": td, "get": lambda *_: False})
+        )
+
+        # Third-party logger (non-axolotl)
+        other = logging.getLogger("thirdparty.lib")
+        other.info("TP-INFO")
+        other.warning("TP-WARN")
+
+        # Simulate Python warnings routed through logging
+        logging.getLogger("py.warnings").warning("PY-WARN")
+
+        # Push through buffers
+        tee.file_only_stream.flush()
+
+        data = read(path)
+        # INFO from non-axolotl should be filtered out (not present)
+        assert "TP-INFO" not in data
+        # WARNING+ should be present
+        assert "TP-WARN" in data
+        # Python warnings captured (via py.warnings logger)
+        assert "PY-WARN" in data
+        tee.close_debug_log()
+        tee.close_debug_log()
+
+
+def test_prepare_debug_log_idempotent_and_no_duplicate(monkeypatch):
+    from axolotl.logging_config import configure_logging
+    from axolotl.utils import tee
+    from axolotl.utils.logging import get_logger
+
+    with tempfile.TemporaryDirectory() as td:
+        monkeypatch.setenv("AXOLOTL_TEE_STDOUT", "0")
+        configure_logging()
+        cfg = type("Cfg", (), {"output_dir": td, "get": lambda *_: False})
+        p1 = tee.prepare_debug_log(cfg)
+        p2 = tee.prepare_debug_log(cfg)
+        assert p1 == p2
+
+        log = get_logger("axolotl.test")
+        marker = "UNIQUE-MARKER-12345"
+        log.info(marker)
+        tee.file_only_stream.flush()
+
+        data = read(p1)
+        # Ensure the marker appears once (not duplicated via propagation)
+        assert data.count(marker) == 1
+        tee.close_debug_log()
+
+
+def test_hub_unauthenticated_nag_suppressed(monkeypatch):
+    from axolotl.logging_config import configure_logging
+    from axolotl.utils import tee
+
+    with tempfile.TemporaryDirectory() as td:
+        monkeypatch.setenv("AXOLOTL_TEE_STDOUT", "0")
+        configure_logging()
+        path = tee.prepare_debug_log(
+            type("Cfg", (), {"output_dir": td, "get": lambda *_: False})
+        )
+
+        hub_http = logging.getLogger("huggingface_hub.utils._http")
+        hub_http.warning(
+            "Warning: You are sending unauthenticated requests to the HF Hub."
+            " Please set a HF_TOKEN to enable higher rate limits and faster downloads."
+        )
+        hub_http.warning("Retrying in 2s [Retry 1/5].")
+        tee.file_only_stream.flush()
+
+        data = read(path)
+        assert "unauthenticated requests" not in data
+        assert "Retrying in 2s" in data
+        tee.close_debug_log()
diff --git a/tests/test_lora.py b/tests/test_lora.py
new file mode 100644
index 0000000000..50cbea9bc6
--- /dev/null
+++ b/tests/test_lora.py
@@ -0,0 +1,69 @@
+"""
+tests for loading loras
+"""
+
+from axolotl.loaders import ModelLoader, load_tokenizer
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+minimal_config = DictDefault(
+    {
+        "base_model": "HuggingFaceTB/SmolLM2-135M",
+        "learning_rate": 0.000001,
+        "datasets": [
+            {
+                "path": "mhenrichsen/alpaca_2k_test",
+                "type": "alpaca",
+            }
+        ],
+        "micro_batch_size": 1,
+        "gradient_accumulation_steps": 1,
+    }
+)
+
+
+class TestLoRALoad:
+    """
+    Test class for loading LoRA weights
+    """
+
+    def test_load_lora_weights(self):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.0,
+                "lora_target_linear": True,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "sequence_len": 1024,
+            }
+            | minimal_config
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        tokenizer = load_tokenizer(cfg)
+        ModelLoader(cfg, tokenizer).load()
+
+    def test_load_lora_weights_empty_dropout(self):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": None,
+                "lora_target_linear": True,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "sequence_len": 1024,
+            }
+            | minimal_config
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        assert cfg.lora_dropout == 0.0
+        tokenizer = load_tokenizer(cfg)
+        ModelLoader(cfg, tokenizer).load()
diff --git a/tests/test_mm_chat_collator.py b/tests/test_mm_chat_collator.py
new file mode 100644
index 0000000000..5c40c5b34f
--- /dev/null
+++ b/tests/test_mm_chat_collator.py
@@ -0,0 +1,164 @@
+"""
+Regression tests for MultiModalChatDataCollator shape contracts.
+
+Guard against the transformers 5.x breakage where apply_chat_template's
+own `return_dict` parameter (default False) caused it to return the raw
+input_ids tensor instead of the full BatchFeature dict, leading to
+  IndexError: too many indices for tensor of dimension 2
+when downstream code did batch["input_ids"] on the resulting tensor.
+"""
+
+from collections.abc import Mapping
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from transformers import BatchFeature
+
+
+@pytest.fixture(name="mock_processor")
+def fixture_mock_processor():
+    """
+    A mock processor whose apply_chat_template returns a BatchFeature
+    when called with return_dict=True (the correct call convention),
+    or a raw input_ids tensor when called without return_dict=True
+    (the broken call convention that the bug introduced).
+    """
+    processor = MagicMock()
+    processor.tokenizer = MagicMock()
+    processor.tokenizer.pad_token_id = 0
+    processor.image_token = "<|image|>"
+    processor.tokenizer.convert_tokens_to_ids = MagicMock(return_value=128256)
+
+    batch_size, seq_len = 2, 16
+    input_ids = torch.ones(batch_size, seq_len, dtype=torch.long)
+    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long)
+
+    batch_feature = BatchFeature(
+        data={
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+    )
+
+    def _apply_chat_template(*args, **kwargs):
+        if kwargs.get("return_dict", False):
+            return batch_feature
+        # Simulate transformers 5.x default behaviour: returns out["input_ids"]
+        return input_ids
+
+    processor.apply_chat_template = MagicMock(side_effect=_apply_chat_template)
+    processor.chat_template = None
+    return processor
+
+
+@pytest.fixture(name="mock_processing_strategy")
+def fixture_mock_processing_strategy(mock_processor):
+    from axolotl.processing_strategies import ProcessingStrategy
+
+    strategy = ProcessingStrategy(processor=mock_processor)
+    return strategy
+
+
+class TestMultiModalChatDataCollatorShapeContract:
+    """
+    Verify that MultiModalChatDataCollator.process_rows returns a dict with
+    2-D input_ids and labels, not a raw tensor.  This is the shape contract
+    that process_labels depends on.
+    """
+
+    def _make_collator(self, mock_processing_strategy):
+        from axolotl.utils.collators.mm_chat import MultiModalChatDataCollator
+
+        tokenizer = mock_processing_strategy.processor.tokenizer
+        return MultiModalChatDataCollator(
+            tokenizer=tokenizer,
+            processing_strategy=mock_processing_strategy,
+        )
+
+    def _make_examples(self):
+        return [
+            {
+                "messages": [
+                    {"role": "user", "content": "Hello"},
+                    {"role": "assistant", "content": "Hi there"},
+                ]
+            }
+        ]
+
+    def test_process_rows_returns_dict(self, mock_processing_strategy):
+        """batch must be a dict, not a raw tensor."""
+        collator = self._make_collator(mock_processing_strategy)
+        examples = self._make_examples()
+
+        with patch.object(
+            mock_processing_strategy,
+            "__call__",
+            return_value=examples,
+        ):
+            batch = collator.process_rows(examples)
+
+        assert isinstance(batch, Mapping), (
+            "process_rows must return a Mapping (BatchFeature or dict), not a "
+            "raw tensor. If it returns a tensor, apply_chat_template was called "
+            "without return_dict=True at the top level."
+        )
+
+    def test_process_rows_input_ids_shape(self, mock_processing_strategy):
+        """batch['input_ids'] must be a 2-D tensor (batch, seq_len)."""
+        collator = self._make_collator(mock_processing_strategy)
+        examples = self._make_examples()
+
+        with patch.object(
+            mock_processing_strategy,
+            "__call__",
+            return_value=examples,
+        ):
+            batch = collator.process_rows(examples)
+
+        assert "input_ids" in batch
+        assert isinstance(batch["input_ids"], torch.Tensor)
+        assert batch["input_ids"].ndim == 2, (
+            f"input_ids must be 2-D (batch, seq_len), got shape {batch['input_ids'].shape}"
+        )
+
+    def test_process_rows_labels_shape(self, mock_processing_strategy):
+        """batch['labels'] must be a 2-D tensor matching input_ids shape."""
+        collator = self._make_collator(mock_processing_strategy)
+        examples = self._make_examples()
+
+        with patch.object(
+            mock_processing_strategy,
+            "__call__",
+            return_value=examples,
+        ):
+            batch = collator.process_rows(examples)
+
+        assert "labels" in batch
+        assert isinstance(batch["labels"], torch.Tensor)
+        assert batch["labels"].ndim == 2
+        assert batch["labels"].shape == batch["input_ids"].shape
+
+    def test_apply_chat_template_called_with_return_dict_true(
+        self, mock_processing_strategy
+    ):
+        """apply_chat_template must be called with return_dict=True as a keyword arg."""
+        collator = self._make_collator(mock_processing_strategy)
+        examples = self._make_examples()
+
+        with patch.object(
+            mock_processing_strategy,
+            "__call__",
+            return_value=examples,
+        ):
+            collator.process_rows(examples)
+
+        call_kwargs = (
+            mock_processing_strategy.processor.apply_chat_template.call_args.kwargs
+        )
+        assert call_kwargs.get("return_dict") is True, (
+            "apply_chat_template must be called with return_dict=True as a top-level "
+            "keyword argument (not inside processor_kwargs). In transformers 5.x, "
+            "apply_chat_template has its own return_dict param (default False) that "
+            "controls whether it returns the full BatchFeature or just input_ids."
+        )
diff --git a/tests/test_mm_collator_warnings.py b/tests/test_mm_collator_warnings.py
new file mode 100644
index 0000000000..5481d2969e
--- /dev/null
+++ b/tests/test_mm_collator_warnings.py
@@ -0,0 +1,82 @@
+"""Tests for the MM dataloader_num_workers=0 warning."""
+
+from __future__ import annotations
+
+import logging
+from types import SimpleNamespace
+
+import pytest
+
+from axolotl.core.builders import causal as _causal_builder
+from axolotl.core.builders.causal import _warn_if_num_workers_zero_for_mm
+
+
+def _make_cfg(processor_type=None, dataloader_num_workers=None, train_on_inputs=False):
+    return SimpleNamespace(
+        processor_type=processor_type,
+        dataloader_num_workers=dataloader_num_workers,
+        train_on_inputs=train_on_inputs,
+    )
+
+
+@pytest.fixture
+def _reset_mm_warn_state():
+    """Reset the module-level once-per-process guard set so each test is independent.
+
+    Production keeps the warning one-shot per process so logs aren't spammed
+    when ``build()`` builds collators for both train and eval.
+    """
+    _causal_builder._MM_NUM_WORKERS_WARNED.clear()
+    yield
+    _causal_builder._MM_NUM_WORKERS_WARNED.clear()
+
+
+def _capture_warning(caplog, cfg) -> list[str]:
+    log = logging.getLogger("axolotl.core.builders.causal")
+    log.addHandler(caplog.handler)
+    try:
+        with caplog.at_level(logging.WARNING, logger="axolotl.core.builders.causal"):
+            _warn_if_num_workers_zero_for_mm(cfg, log)
+    finally:
+        log.removeHandler(caplog.handler)
+    return [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING]
+
+
+def test_warn_num_workers_zero_for_mm_emits_warning(_reset_mm_warn_state, caplog):
+    cfg = _make_cfg(processor_type="qwen2_vl", dataloader_num_workers=0)
+    msgs = _capture_warning(caplog, cfg)
+    assert any("dataloader_num_workers=0" in m for m in msgs)
+
+
+def test_warn_num_workers_zero_for_mm_silent_when_workers_set(
+    _reset_mm_warn_state, caplog
+):
+    cfg = _make_cfg(processor_type="qwen2_vl", dataloader_num_workers=2)
+    msgs = _capture_warning(caplog, cfg)
+    assert not any("dataloader_num_workers=0" in m for m in msgs)
+
+
+def test_warn_num_workers_zero_for_mm_silent_when_no_processor(
+    _reset_mm_warn_state, caplog
+):
+    cfg = _make_cfg(processor_type=None, dataloader_num_workers=0)
+    msgs = _capture_warning(caplog, cfg)
+    assert not any("dataloader_num_workers=0" in m for m in msgs)
+
+
+def test_warn_num_workers_zero_for_mm_treats_none_as_zero(_reset_mm_warn_state, caplog):
+    """``None`` (default) and ``0`` should both trigger the warning."""
+    cfg = _make_cfg(processor_type="qwen2_vl", dataloader_num_workers=None)
+    msgs = _capture_warning(caplog, cfg)
+    assert any("dataloader_num_workers=0" in m for m in msgs)
+
+
+def test_warn_num_workers_zero_for_mm_is_one_shot(_reset_mm_warn_state, caplog):
+    """Second invocation in the same process should NOT re-emit the warning."""
+    cfg = _make_cfg(processor_type="qwen2_vl", dataloader_num_workers=0)
+    first_msgs = _capture_warning(caplog, cfg)
+    assert any("dataloader_num_workers=0" in m for m in first_msgs)
+
+    caplog.clear()
+    second_msgs = _capture_warning(caplog, cfg)
+    assert not any("dataloader_num_workers=0" in m for m in second_msgs)
diff --git a/tests/test_no_legacy_attn_reads.py b/tests/test_no_legacy_attn_reads.py
new file mode 100644
index 0000000000..d1e3239790
--- /dev/null
+++ b/tests/test_no_legacy_attn_reads.py
@@ -0,0 +1,61 @@
+"""Enforce attn_implementation as the single source of truth.
+
+Fails if src/ contains a cfg.<legacy>_attention read. Migrate offending sites
+to cfg.attn_implementation or the attn_supports_packing/attn_uses_flash_lib/
+attn_needs_dtype_cast computed flags.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+LEGACY_FLAGS = (
+    "flash_attention",
+    "sdp_attention",
+    "xformers_attention",
+    "flex_attention",
+    "sage_attention",
+    "eager_attention",
+)
+
+# The normalizer is allowed to read the legacy keys (that's its job).
+# lm_eval/cli.py is a raw-YAML entry point (bypasses AxolotlInputConfig) that
+# honors both forms during the deprecation period — when we remove the legacy
+# flags entirely, drop this allowlist entry and the BC branch in that file.
+ALLOWED_FILES = {
+    Path("src/axolotl/utils/schemas/config.py"),
+    Path("src/axolotl/integrations/lm_eval/cli.py"),
+}
+
+# `cfg.<flag>`, `self.cfg.<flag>`, `data.get("<flag>")`, `data["<flag>"]`
+_PATTERNS = [re.compile(rf"\bcfg\.{flag}\b") for flag in LEGACY_FLAGS] + [
+    re.compile(rf'\bdata\.get\("{flag}"\)') for flag in LEGACY_FLAGS
+]
+
+
+def _repo_root() -> Path:
+    return Path(__file__).resolve().parent.parent
+
+
+def test_no_legacy_attn_reads_in_src():
+    root = _repo_root()
+    src = root / "src"
+    offenders: list[str] = []
+
+    for py_file in src.rglob("*.py"):
+        rel = py_file.relative_to(root)
+        if rel in ALLOWED_FILES:
+            continue
+        text = py_file.read_text(encoding="utf-8")
+        for pattern in _PATTERNS:
+            for match in pattern.finditer(text):
+                # Line number for the user's convenience.
+                line_no = text.count("\n", 0, match.start()) + 1
+                offenders.append(f"{rel}:{line_no}  {match.group(0)}")
+
+    assert not offenders, (
+        "Found legacy attention-flag reads in src/. Migrate to "
+        "`cfg.attn_implementation` / capability flags:\n  "
+        + "\n  ".join(sorted(offenders))
+    )
diff --git a/tests/test_normalize_config.py b/tests/test_normalize_config.py
index 2e76ceb45d..e12224fa29 100644
--- a/tests/test_normalize_config.py
+++ b/tests/test_normalize_config.py
@@ -1,10 +1,15 @@
 """
 Test classes for checking functionality of the cfg normalization
 """
+
 import unittest
 from unittest.mock import patch
 
-from axolotl.utils.config import normalize_cfg_datasets, normalize_config
+from axolotl.utils.config import (
+    normalize_cfg_datasets,
+    normalize_config,
+    validate_config,
+)
 from axolotl.utils.dict import DictDefault
 
 
@@ -16,12 +21,19 @@ class NormalizeConfigTestCase(unittest.TestCase):
     def _get_base_cfg(self):
         return DictDefault(
             {
-                "base_model": "JackFram/llama-68m",
-                "base_model_config": "JackFram/llama-68m",
-                "tokenizer_type": "LlamaTokenizer",
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "base_model_config": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
                 "num_epochs": 1,
                 "micro_batch_size": 1,
                 "gradient_accumulation_steps": 1,
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "learning_rate": 0.0001,
             }
         )
 
@@ -39,12 +51,12 @@ def test_chat_template_chatml(self):
                 "datasets": [
                     {
                         "path": "lorem/ipsum",
-                        "type": "sharegpt",
-                        "conversation": "vicuna_v1.1",
+                        "type": "chat_template",
+                        "chat_template": "gemma",
                     },
                     {
                         "path": "sit/amet",
-                        "type": "sharegpt",
+                        "type": "chat_template",
                     },
                 ],
             }
@@ -52,8 +64,8 @@ def test_chat_template_chatml(self):
 
         normalize_cfg_datasets(cfg)
 
-        assert cfg.datasets[0].conversation == "vicuna_v1.1"
-        assert cfg.datasets[1].conversation == "chatml"
+        assert cfg.datasets[0].chat_template == "gemma"
+        assert cfg.datasets[1].chat_template == "chatml"
 
     @patch("axolotl.utils.config.is_torch_bf16_gpu_available")
     def test_bf16_auto_setter_available(self, mock_bf16_avail):
@@ -89,3 +101,100 @@ def test_bf16_disables_fp16(self, mock_bf16_avail):
 
         self.assertTrue(cfg.bf16)
         self.assertFalse(cfg.fp16)
+
+    def test_migrate_fsdp_config(self):
+        """Test basic FSDP config migration with and without fsdp_version"""
+        cfg_with_version = self._get_base_cfg() | DictDefault(
+            {
+                "fsdp_config": {
+                    "fsdp_version": 2,
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                    "fsdp_offload_params": False,
+                    "fsdp_cpu_ram_efficient_loading": True,
+                }
+            }
+        )
+
+        cfg_with_version = validate_config(cfg_with_version)
+
+        self.assertEqual(cfg_with_version.fsdp_version, 2)
+        self.assertEqual(
+            cfg_with_version.fsdp_config.auto_wrap_policy, "TRANSFORMER_BASED_WRAP"
+        )
+        self.assertEqual(cfg_with_version.fsdp_config.offload_params, False)
+        self.assertEqual(cfg_with_version.fsdp_config.cpu_ram_efficient_loading, True)
+
+        self.assertNotIn("fsdp_auto_wrap_policy", cfg_with_version.fsdp_config)
+        self.assertNotIn("fsdp_offload_params", cfg_with_version.fsdp_config)
+        self.assertNotIn("fsdp_cpu_ram_efficient_loading", cfg_with_version.fsdp_config)
+        self.assertIn("fsdp_version", cfg_with_version.fsdp_config)
+
+        cfg_without_version = self._get_base_cfg() | DictDefault(
+            {
+                "fsdp_config": {
+                    "fsdp_auto_wrap_policy": "SIZE_BASED_WRAP",
+                    "fsdp_offload_params": True,
+                    "fsdp_min_num_params": 100000000,
+                }
+            }
+        )
+
+        cfg_without_version = validate_config(cfg_without_version)
+
+        self.assertEqual(cfg_without_version.fsdp_version, 1)
+        self.assertEqual(cfg_without_version.fsdp_config.fsdp_version, 1)
+        self.assertEqual(
+            cfg_without_version.fsdp_config.auto_wrap_policy, "SIZE_BASED_WRAP"
+        )
+        self.assertEqual(cfg_without_version.fsdp_config.offload_params, True)
+        self.assertEqual(cfg_without_version.fsdp_config.min_num_params, 100000000)
+
+        self.assertNotIn("fsdp_auto_wrap_policy", cfg_without_version.fsdp_config)
+        self.assertNotIn("fsdp_offload_params", cfg_without_version.fsdp_config)
+        self.assertNotIn("fsdp_min_num_params", cfg_without_version.fsdp_config)
+
+    def test_migrate_fsdp_config_no_fsdp_config(self):
+        """Test that function doesn't crash when no fsdp_config is present"""
+        cfg = self._get_base_cfg()
+
+        cfg = validate_config(cfg)
+
+        self.assertNotIn("fsdp_config", cfg)
+        self.assertNotIn("fsdp_version", cfg)
+
+    def test_migrate_fsdp_config_empty_fsdp_config(self):
+        """Test migration with empty fsdp_config"""
+        cfg = self._get_base_cfg() | DictDefault({"fsdp_config": {}})
+
+        cfg = validate_config(cfg)
+
+        self.assertNotIn("fsdp_version", cfg)
+        self.assertEqual(cfg.fsdp_config, {})
+
+    def test_migrate_fsdp_config_mixed_keys(self):
+        """Test migration with a mix of fsdp_ and non-fsdp_ keys"""
+        cfg = self._get_base_cfg() | DictDefault(
+            {
+                "fsdp_config": {
+                    "fsdp_version": 1,
+                    "fsdp_state_dict_type": "FULL_STATE_DICT",
+                    "mixed_precision_policy": "fp16",
+                    "activation_checkpointing": True,
+                    "fsdp_reshard_after_forward": False,
+                }
+            }
+        )
+
+        cfg = validate_config(cfg)
+
+        self.assertEqual(cfg.fsdp_version, 1)
+        self.assertEqual(cfg.fsdp_config.state_dict_type, "FULL_STATE_DICT")
+        self.assertEqual(cfg.fsdp_config.reshard_after_forward, False)
+        self.assertEqual(cfg.fsdp_config.mixed_precision_policy, "fp16")
+        self.assertEqual(cfg.fsdp_config.activation_checkpointing, True)
+
+        # Check original fsdp_ keys are removed
+        self.assertNotIn("fsdp_state_dict_type", cfg.fsdp_config)
+        self.assertNotIn("fsdp_reshard_after_forward", cfg.fsdp_config)
+
+        self.assertIn("fsdp_version", cfg.fsdp_config)
diff --git a/tests/test_opentelemetry_callback.py b/tests/test_opentelemetry_callback.py
new file mode 100644
index 0000000000..294ff65857
--- /dev/null
+++ b/tests/test_opentelemetry_callback.py
@@ -0,0 +1,349 @@
+"""Tests for OpenTelemetry metrics callback functionality."""
+
+import time
+
+import pytest
+
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture
+def mock_otel_config():
+    """Mock configuration for OpenTelemetry callback."""
+    return DictDefault(
+        {
+            "use_otel_metrics": True,
+            "otel_metrics_host": "localhost",
+            "otel_metrics_port": 8003,  # Use unique port for tests
+        }
+    )
+
+
+@pytest.fixture
+def mock_trainer_state():
+    """Mock trainer state for callback testing."""
+    from transformers import TrainerState
+
+    state = TrainerState()
+    state.epoch = 1.0
+    state.global_step = 100
+    return state
+
+
+@pytest.fixture
+def mock_training_args():
+    """Mock training arguments for callback testing."""
+    from transformers import TrainingArguments
+
+    return TrainingArguments(output_dir="/tmp/test")
+
+
+@pytest.fixture
+def mock_trainer_control():
+    """Mock trainer control for callback testing."""
+    from transformers.trainer_callback import TrainerControl
+
+    return TrainerControl()
+
+
+class TestOpenTelemetryConfig:
+    """Test OpenTelemetry configuration schema."""
+
+    def test_config_schema_valid(self):
+        """Test OpenTelemetry configuration schema validation."""
+        from axolotl.utils.schemas.integrations import OpenTelemetryConfig
+
+        # Test valid config
+        valid_config = {
+            "use_otel_metrics": True,
+            "otel_metrics_host": "localhost",
+            "otel_metrics_port": 8000,
+        }
+
+        otel_config = OpenTelemetryConfig(**valid_config)
+        assert otel_config.use_otel_metrics is True
+        assert otel_config.otel_metrics_host == "localhost"
+        assert otel_config.otel_metrics_port == 8000
+
+    def test_config_defaults(self):
+        """Test OpenTelemetry configuration default values."""
+        from axolotl.utils.schemas.integrations import OpenTelemetryConfig
+
+        # Test minimal config with defaults
+        minimal_config = {"use_otel_metrics": True}
+
+        otel_config = OpenTelemetryConfig(**minimal_config)
+        assert otel_config.use_otel_metrics is True
+        assert otel_config.otel_metrics_host == "localhost"  # default
+        assert otel_config.otel_metrics_port == 8000  # default
+
+    def test_config_disabled_by_default(self):
+        """Test that OpenTelemetry is disabled by default."""
+        from axolotl.utils.schemas.integrations import OpenTelemetryConfig
+
+        # Test default config
+        default_config = OpenTelemetryConfig()
+        assert default_config.use_otel_metrics is False
+
+
+class TestOpenTelemetryCallback:
+    """Test OpenTelemetry callback functionality."""
+
+    def test_callback_import(self):
+        """Test that OpenTelemetry callback can be imported."""
+        from axolotl.utils.callbacks.opentelemetry import OpenTelemetryMetricsCallback
+
+        assert OpenTelemetryMetricsCallback is not None
+
+    def test_callback_graceful_fallback(self, mock_otel_config):
+        """Test callback gracefully handles missing dependencies."""
+        from axolotl.utils.callbacks.opentelemetry import OpenTelemetryMetricsCallback
+
+        # This should not raise an exception even if dependencies are missing
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+
+        # Callback should exist but may have metrics disabled
+        assert callback is not None
+        assert hasattr(callback, "metrics_enabled")
+
+    def test_callback_initialization_enabled(self, mock_otel_config):
+        """Test callback initialization when OpenTelemetry is available."""
+        from axolotl.utils.callbacks.opentelemetry import (
+            OPENTELEMETRY_AVAILABLE,
+            OpenTelemetryMetricsCallback,
+        )
+
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+
+        if OPENTELEMETRY_AVAILABLE:
+            assert callback.metrics_enabled is True
+            assert callback.cfg == mock_otel_config
+            assert callback.metrics_host == "localhost"
+            assert callback.metrics_port == 8003
+        else:
+            assert callback.metrics_enabled is False
+
+    def test_metrics_server_lifecycle(
+        self,
+        mock_otel_config,
+        mock_trainer_state,
+        mock_training_args,
+        mock_trainer_control,
+    ):
+        """Test metrics server starts and stops correctly."""
+        from axolotl.utils.callbacks.opentelemetry import (
+            OPENTELEMETRY_AVAILABLE,
+            OpenTelemetryMetricsCallback,
+        )
+
+        if not OPENTELEMETRY_AVAILABLE:
+            pytest.skip("OpenTelemetry dependencies not available")
+
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+
+        # Start server
+        callback.on_train_begin(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
+        assert callback.server_started is True
+
+        # End training
+        callback.on_train_end(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
+
+    def test_metrics_recording(
+        self,
+        mock_otel_config,
+        mock_trainer_state,
+        mock_training_args,
+        mock_trainer_control,
+    ):
+        """Test that metrics are recorded during training."""
+        from axolotl.utils.callbacks.opentelemetry import (
+            OPENTELEMETRY_AVAILABLE,
+            OpenTelemetryMetricsCallback,
+        )
+
+        if not OPENTELEMETRY_AVAILABLE:
+            pytest.skip("OpenTelemetry dependencies not available")
+
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+        callback.on_train_begin(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
+
+        # Test logging metrics
+        test_logs = {
+            "loss": 0.5,
+            "learning_rate": 1e-4,
+            "grad_norm": 0.8,
+        }
+
+        # This should not raise an exception
+        callback.on_log(
+            mock_training_args, mock_trainer_state, mock_trainer_control, logs=test_logs
+        )
+        assert callback.metrics_enabled is True
+
+    def test_evaluation_metrics(
+        self,
+        mock_otel_config,
+        mock_trainer_state,
+        mock_training_args,
+        mock_trainer_control,
+    ):
+        """Test evaluation metrics recording."""
+        from axolotl.utils.callbacks.opentelemetry import (
+            OPENTELEMETRY_AVAILABLE,
+            OpenTelemetryMetricsCallback,
+        )
+
+        if not OPENTELEMETRY_AVAILABLE:
+            pytest.skip("OpenTelemetry dependencies not available")
+
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+        callback.on_train_begin(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
+
+        # Test evaluation metrics
+        eval_logs = {
+            "eval_loss": 0.3,
+            "eval_accuracy": 0.95,
+        }
+
+        # This should not raise an exception
+        callback.on_evaluate(
+            mock_training_args, mock_trainer_state, mock_trainer_control, eval_logs
+        )
+        assert callback.metrics_enabled is True
+
+    def test_thread_safety(self, mock_otel_config):
+        """Test that callback has thread safety mechanisms."""
+        from axolotl.utils.callbacks.opentelemetry import (
+            OPENTELEMETRY_AVAILABLE,
+            OpenTelemetryMetricsCallback,
+        )
+
+        if not OPENTELEMETRY_AVAILABLE:
+            pytest.skip("OpenTelemetry dependencies not available")
+
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+        assert hasattr(callback, "metrics_lock")
+        # Check it's a lock-like object
+        assert hasattr(callback.metrics_lock, "__enter__")
+        assert hasattr(callback.metrics_lock, "__exit__")
+
+
+class TestOpenTelemetryIntegration:
+    """Integration tests for OpenTelemetry."""
+
+    def test_availability_check(self):
+        """Test availability check function."""
+        from axolotl.utils import is_opentelemetry_available
+
+        result = is_opentelemetry_available()
+        assert isinstance(result, bool)
+
+    def test_prometheus_endpoint_basic(
+        self,
+        mock_otel_config,
+        mock_trainer_state,
+        mock_training_args,
+        mock_trainer_control,
+    ):
+        """Test basic Prometheus endpoint functionality."""
+        from axolotl.utils.callbacks.opentelemetry import (
+            OPENTELEMETRY_AVAILABLE,
+            OpenTelemetryMetricsCallback,
+        )
+
+        if not OPENTELEMETRY_AVAILABLE:
+            pytest.skip("OpenTelemetry dependencies not available")
+
+        try:
+            import requests
+        except ImportError:
+            pytest.skip("requests library not available")
+
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+        callback.on_train_begin(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
+
+        if not callback.server_started:
+            pytest.skip("Metrics server failed to start")
+
+        # Give server time to start
+        time.sleep(1)
+
+        # Try to access metrics endpoint
+        try:
+            response = requests.get(
+                f"http://{callback.metrics_host}:{callback.metrics_port}/metrics",
+                timeout=2,
+            )
+            assert response.status_code == 200
+            # Check for Prometheus format
+            assert "# TYPE" in response.text or "# HELP" in response.text
+        except requests.exceptions.RequestException:
+            pytest.skip(
+                "Could not connect to metrics endpoint - this is expected in some environments"
+            )
+
+
+class TestOpenTelemetryCallbackMethods:
+    """Test specific callback methods."""
+
+    def test_step_end_callback(
+        self,
+        mock_otel_config,
+        mock_trainer_state,
+        mock_training_args,
+        mock_trainer_control,
+    ):
+        """Test step end callback method."""
+        from axolotl.utils.callbacks.opentelemetry import (
+            OPENTELEMETRY_AVAILABLE,
+            OpenTelemetryMetricsCallback,
+        )
+
+        if not OPENTELEMETRY_AVAILABLE:
+            pytest.skip("OpenTelemetry dependencies not available")
+
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+        callback.on_train_begin(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
+
+        # Should not raise an exception
+        callback.on_step_end(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
+
+    def test_epoch_end_callback(
+        self,
+        mock_otel_config,
+        mock_trainer_state,
+        mock_training_args,
+        mock_trainer_control,
+    ):
+        """Test epoch end callback method."""
+        from axolotl.utils.callbacks.opentelemetry import (
+            OPENTELEMETRY_AVAILABLE,
+            OpenTelemetryMetricsCallback,
+        )
+
+        if not OPENTELEMETRY_AVAILABLE:
+            pytest.skip("OpenTelemetry dependencies not available")
+
+        callback = OpenTelemetryMetricsCallback(mock_otel_config)
+        callback.on_train_begin(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
+
+        # Should not raise an exception
+        callback.on_epoch_end(
+            mock_training_args, mock_trainer_state, mock_trainer_control
+        )
diff --git a/tests/test_packed_batch_sampler.py b/tests/test_packed_batch_sampler.py
index 50f39d60f5..465425aeff 100644
--- a/tests/test_packed_batch_sampler.py
+++ b/tests/test_packed_batch_sampler.py
@@ -1,15 +1,19 @@
 """Module for testing streaming dataset sequence packing"""
+
 import pytest
-from datasets import concatenate_datasets, load_dataset
+from datasets import concatenate_datasets
 from torch.utils.data import DataLoader, RandomSampler
 from transformers import AutoTokenizer
 
 from axolotl.datasets import TokenizedPromptDataset
 from axolotl.prompt_strategies.completion import load
 from axolotl.utils.collators import V2BatchSamplerDataCollatorForSeq2Seq
+from axolotl.utils.data.utils import handle_long_seq_in_dataset
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
 
+from tests.hf_offline_utils import enable_hf_offline
+
 
 @pytest.fixture(name="tokenizer")
 def fixture_tokenizer():
@@ -18,11 +22,6 @@ def fixture_tokenizer():
     return tokenizer
 
 
-@pytest.fixture(name="max_seq_length")
-def fixture_max_seq_length():
-    return 4096
-
-
 class TestBatchedSamplerPacking:
     """
     Test class for packing streaming dataset sequences
@@ -37,14 +36,28 @@ class TestBatchedSamplerPacking:
             (2, 2),
         ],
     )
-    def test_packing(self, batch_size, num_workers, tokenizer, max_seq_length):
-        import axolotl.monkeypatch.data.batch_dataset_fetcher  # pylint: disable=unused-import  # noqa: F401
-
-        dataset = load_dataset(
-            "Trelis/tiny-shakespeare",
-            split="train",
+    @pytest.mark.parametrize("max_seq_length", [4096, 512])
+    @pytest.mark.parametrize("sequential", [True, False])
+    @enable_hf_offline
+    def test_packing(
+        self,
+        dataset_winglian_tiny_shakespeare,
+        batch_size,
+        num_workers,
+        tokenizer,
+        max_seq_length,
+        sequential,
+    ):
+        from axolotl.monkeypatch.data.batch_dataset_fetcher import (
+            apply_multipack_dataloader_patch,
+            remove_multipack_dataloader_patch,
         )
 
+        # Apply the patch for multipack handling
+        apply_multipack_dataloader_patch()
+
+        dataset = dataset_winglian_tiny_shakespeare["train"].select(range(16))
+
         cfg = DictDefault(
             {
                 "train_on_inputs": True,
@@ -53,7 +66,7 @@ def test_packing(self, batch_size, num_workers, tokenizer, max_seq_length):
         )
         ds_cfg = DictDefault(
             {
-                "field": "Text",
+                "field": "text",
             }
         )
         completion_strategy = load(tokenizer, cfg, ds_cfg)
@@ -62,18 +75,25 @@ def test_packing(self, batch_size, num_workers, tokenizer, max_seq_length):
             dataset,
         )
         train_dataset = concatenate_datasets([dataset_wrapper])
+
+        train_dataset = handle_long_seq_in_dataset(train_dataset, cfg.sequence_len, cfg)
+
+        lengths = get_dataset_lengths(train_dataset)
         batch_sampler = MultipackBatchSampler(
             sampler=RandomSampler(train_dataset),
+            lengths=lengths,
             batch_size=batch_size,
-            drop_last=True,
             batch_max_len=max_seq_length,
-            lengths=get_dataset_lengths(train_dataset),
+            group_size=100000,
+            bin_size=200,
+            sequential=sequential,
+            drop_last=False,
         )
 
         loader = DataLoader(
             train_dataset,
             batch_sampler=batch_sampler,
-            collate_fn=V2BatchSamplerDataCollatorForSeq2Seq(  # pylint: disable=unexpected-keyword-arg
+            collate_fn=V2BatchSamplerDataCollatorForSeq2Seq(
                 tokenizer=tokenizer,
                 padding=True,
                 pad_to_multiple_of=max_seq_length,
@@ -81,19 +101,20 @@ def test_packing(self, batch_size, num_workers, tokenizer, max_seq_length):
             ),
             num_workers=num_workers,
         )
-        inputs = next(iter(loader))
-
-        assert inputs["input_ids"].shape == (batch_size, max_seq_length)
-        assert inputs["labels"].shape == (batch_size, max_seq_length)
-        assert inputs["attention_mask"].shape == (batch_size, max_seq_length)
-
-        assert inputs["input_ids"].tolist()[0][0] == 2
-        assert inputs["labels"].tolist()[0][0] == -100
-        assert inputs["attention_mask"].tolist()[0][0] == 0
-        assert inputs["attention_mask"].tolist()[0][-1] > 1
-
-        if batch_size >= 2:
-            assert inputs["input_ids"].tolist()[1][0] == 2
-            assert inputs["labels"].tolist()[1][0] == -100
-            assert inputs["attention_mask"].tolist()[1][0] == 0
-            assert inputs["attention_mask"].tolist()[1][-1] > 1
+
+        batch_idxs = []
+        for batch in batch_sampler:
+            for pack in batch:
+                batch_idxs.extend(pack)
+
+        try:
+            for batch in loader:
+                assert batch["input_ids"].numel() <= batch_size * max_seq_length
+                assert batch["input_ids"].shape[1] == max_seq_length
+
+            original_idxs = set(range(len(train_dataset)))
+            assert original_idxs == set(batch_idxs)
+            assert len(batch_idxs) == len(set(batch_idxs))
+        finally:
+            # Clean up: remove the patch after the test
+            remove_multipack_dataloader_patch()
diff --git a/tests/test_packed_dataset.py b/tests/test_packed_dataset.py
index da8fb7a937..cecd82a814 100644
--- a/tests/test_packed_dataset.py
+++ b/tests/test_packed_dataset.py
@@ -1,14 +1,17 @@
 """Module for testing dataset sequence packing"""
 
 import unittest
-from pathlib import Path
 
-from datasets import Dataset, load_dataset
 from transformers import AutoTokenizer
 
-from axolotl.datasets import ConstantLengthDataset, TokenizedPromptDataset
-from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy
-from axolotl.prompters import AlpacaPrompter
+from axolotl.cli.args import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
+from axolotl.train import setup_model_and_trainer
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import with_temp_dir
+from tests.hf_offline_utils import enable_hf_offline
 
 
 class TestPacking(unittest.TestCase):
@@ -16,8 +19,8 @@ class TestPacking(unittest.TestCase):
     Test class for packing dataset sequences
     """
 
+    @enable_hf_offline
     def setUp(self) -> None:
-        # pylint: disable=duplicate-code
         self.tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
         self.tokenizer.add_special_tokens(
             {
@@ -27,42 +30,80 @@ def setUp(self) -> None:
             }
         )
 
-    def test_increments_attention(self):
-        prompter = AlpacaPrompter("chat")
-        strat = AlpacaPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
+    @with_temp_dir
+    def test_lora_packing(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "sequence_len": 1024,
+                "sample_packing": True,
+                "multipack_real_batches": False,
+                "eval_sample_packing": True,
+                "adapter": "lora",
+                "lora_r": 32,
+                "lora_alpha": 64,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.2,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "dataset_num_proc": 1,
+                "num_epochs": 1,
+                "max_steps": 20,
+                "save_steps": 10,
+                "micro_batch_size": 8,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_torch_fused",
+                "lr_scheduler": "cosine",
+                "fp16": False,
+                "bf16": False,
+            }
         )
-        dateset = load_dataset(
-            "json",
-            data_files=str(Path(__file__).parent / "fixtures/alpaca/alpaca.json"),
-        )["train"]
-        dataset = Dataset.from_list(list(TokenizedPromptDataset(strat, dateset)))
 
-        constant_len_dataset = ConstantLengthDataset(
-            self.tokenizer,
-            [dataset],
-            seq_length=2048,
-        )
-        packed_dataset = Dataset.from_list(list(constant_len_dataset))
-        example = packed_dataset[0]
-        next_bos_index = (
-            example["input_ids"][1:].index(self.tokenizer.bos_token_id) + 1
-        )  # add one since we sliced
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
-        # first example doesn't have mask reset
-        assert example["input_ids"][0] == self.tokenizer.bos_token_id
-        assert example["attention_mask"][0] == 1
-        assert example["position_ids"][0] == 0
-        assert example["position_ids"][1] == 1
+        (
+            trainer,
+            _,
+            _,
+            _,
+            _,
+        ) = setup_model_and_trainer(cfg, dataset_meta)
 
-        # but subsequent one does
-        assert example["input_ids"][next_bos_index] == self.tokenizer.bos_token_id
-        assert example["attention_mask"][next_bos_index] == 2
-        assert example["position_ids"][next_bos_index] == 0
-        assert example["position_ids"][next_bos_index + 1] == 1
+        sampler = trainer._get_eval_sampler(trainer.eval_dataset)
+        assert "MultipackBatchSampler" in sampler.__class__.__name__
+        assert (
+            "V2BatchSamplerDataCollatorForSeq2Seq"
+            in trainer.eval_data_collator.__class__.__name__
+        )
+        dataloader = trainer.get_eval_dataloader(trainer.eval_dataset)
+        dataloader_iter = iter(dataloader)
+        batch = next(dataloader_iter)
+        assert batch["input_ids"].shape == (1, 8192)
+
+        sampler = trainer._get_train_sampler(trainer.train_dataset)
+        assert "MultipackBatchSampler" in sampler.__class__.__name__
+        assert (
+            "V2BatchSamplerDataCollatorForSeq2Seq"
+            in trainer.train_data_collator.__class__.__name__
+        )
+        dataloader = trainer.get_train_dataloader()
+        dataloader_iter = iter(dataloader)
+        batch = next(dataloader_iter)
+        assert batch["input_ids"].shape == (1, 8192)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_packed_pretraining.py b/tests/test_packed_pretraining.py
index 528f9c8074..0458f7ba23 100644
--- a/tests/test_packed_pretraining.py
+++ b/tests/test_packed_pretraining.py
@@ -1,67 +1,87 @@
 """Module for testing streaming dataset sequence packing"""
+
 import functools
-import unittest
+import random
+import string
 
+import pytest
 import torch
-from datasets import load_dataset
+from datasets import IterableDataset
 from torch.utils.data import DataLoader
-from transformers import AutoTokenizer
 
-from axolotl.utils.data import get_dataset_wrapper, wrap_pretraining_dataset
+from axolotl.utils.data import get_dataset_wrapper, wrap_streaming_dataset
 from axolotl.utils.dict import DictDefault
 
 
-class TestPretrainingPacking(unittest.TestCase):
+class TestPretrainingPacking:
     """
     Test class for packing streaming dataset sequences
     """
 
-    def setUp(self) -> None:
-        # pylint: disable=duplicate-code
-        self.tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
-        self.tokenizer.pad_token = "</s>"
+    @pytest.fixture
+    def random_text(self):
+        # seed with random.seed(0) for reproducibility
+        random.seed(0)
+
+        # generate row of random text with "words" of between 2 and 10 characters and
+        # between 400 to 1200 characters per line
+        def rand_txt():
+            return " ".join(
+                [
+                    "".join(
+                        random.choices(string.ascii_lowercase, k=random.randint(2, 10))
+                    )
+                    for _ in range(random.randint(50, 200))
+                ]
+            )
+
+        # Create a list of 2000 random texts rather than just using it within the
+        # generator so the test runs faster
+        data = [rand_txt() for _ in range(500)]
+
+        # Create an IterableDataset
+        def generator():
+            for row in data:
+                yield {"text": row}
 
-    def test_packing_stream_dataset(self):
-        # pylint: disable=duplicate-code
-        dataset = load_dataset(
-            "c4",
-            "en",
-            streaming=True,
-        )["train"]
+        return IterableDataset.from_generator(generator)
+
+    @pytest.mark.flaky(retries=1, delay=5)
+    def test_packing_stream_dataset(self, tokenizer_huggyllama, random_text):
+        dataset = random_text
 
         cfg = DictDefault(
             {
                 "pretraining_dataset": [
                     {
-                        "path": "c4",
-                        "name": "en",
+                        "path": "winglian/tiny-shakespeare",
                         "type": "pretrain",
                     }
                 ],
                 "sample_packing": True,
+                "pretrain_multipack_attn": True,
                 "pad_to_sequence_len": True,
                 "sequence_len": 2048,
                 "micro_batch_size": 2,
+                "sample_packing_group_size": 100000,
+                "sample_packing_bin_size": 200,
             }
         )
 
         ds_wrapper_partial = functools.partial(
             get_dataset_wrapper,
             cfg.pretraining_dataset[0],
-            self.tokenizer,
+            tokenizer_huggyllama,
             cfg,
             cfg.pretraining_dataset[0]["type"] or "pretrain",
         )
 
         original_bsz = cfg.micro_batch_size
-        train_dataset = wrap_pretraining_dataset(
+        train_dataset = wrap_streaming_dataset(
             dataset,
-            self.tokenizer,
+            tokenizer_huggyllama,
             cfg,
             ds_wrapper_partial,
-            max_tokens=cfg.sequence_len,
-            batch_size=cfg.micro_batch_size,
-            seed=cfg.seed or 42,
         )
 
         trainer_loader = DataLoader(
@@ -72,7 +92,7 @@ def test_packing_stream_dataset(self):
         )
         idx = 0
         for data in trainer_loader:
-            if idx > 10:
+            if idx > 3:
                 break
             assert data["input_ids"].shape == torch.Size(
                 [1, original_bsz * cfg.sequence_len]
@@ -83,11 +103,9 @@ def test_packing_stream_dataset(self):
             assert data["labels"].shape == torch.Size(
                 [1, original_bsz * cfg.sequence_len]
             )
-            assert data["attention_mask"].shape == torch.Size(
-                [1, original_bsz * cfg.sequence_len]
-            )
+            assert "attention_mask" not in data
+            # FIXME add back once we fix packing unpad/pad with attention mask
+            # assert data["attention_mask"].shape == torch.Size(
+            #     [1, original_bsz * cfg.sequence_len]
+            # )
             idx += 1
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/tests/test_perplexity.py b/tests/test_perplexity.py
new file mode 100644
index 0000000000..899308ba66
--- /dev/null
+++ b/tests/test_perplexity.py
@@ -0,0 +1,47 @@
+"""unit tests for perplexity eval callback"""
+
+from pytest import fixture
+from transformers.models.auto.modeling_auto import AutoModelForCausalLM
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+from axolotl.utils.callbacks.perplexity import Perplexity
+
+MODEL_NAME = "HuggingFaceTB/SmolLM2-135M"
+
+
+@fixture()
+def metric(tokenizer):
+    return Perplexity(tokenizer=tokenizer, max_seq_len=512)
+
+
+@fixture()
+def model():
+    return AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, trust_remote_code=True, dtype="float32"
+    )
+
+
+@fixture()
+def tokenizer():
+    tokenizer_ = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    tokenizer_.add_special_tokens({"pad_token": "<|endoftext|>"})
+    return tokenizer_
+
+
+def test_perplexity_longer_than_stride(model, metric):
+    # taken from https://huggingface.co/datasets/roneneldan/TinyStories
+    sample_text = """
+Once upon a time, there was a little car named Beep. Beep loved to go fast and play in the sun. Beep was a healthy car because he always had good fuel. Good fuel made Beep happy and strong. One day, Beep was driving in the park when he saw a big tree. The tree had many leaves that were falling. Beep liked how the leaves fall and wanted to play with them. Beep drove under the tree and watched the leaves fall on him. He laughed and beeped his horn. Beep played with the falling leaves all day. When it was time to go home, Beep knew he needed more fuel. He went to the fuel place and got more healthy fuel. Now, Beep was ready to go fast and play again the next day. And Beep lived happily ever after.
+One day, a little fish named Fin was swimming near the shore. He saw a big crab and wanted to be friends. "Hi, I am Fin. Do you want to play?" asked the little fish. The crab looked at Fin and said, "No, I don't want to play. I am cold and I don't feel fine." Fin felt sad but wanted to help the crab feel better. He swam away and thought of a plan. He remembered that the sun could make things warm. So, Fin swam to the top of the water and called to the sun, "Please, sun, help my new friend feel fine and not freeze!" The sun heard Fin's call and shone its warm light on the shore. The crab started to feel better and not so cold. He saw Fin and said, "Thank you, little fish, for making me feel fine. I don't feel like I will freeze now. Let's play together!" And so, Fin and the crab played and became good friends.
+"""
+    result = metric.compute(model, [sample_text])
+    ppl = result["score"]
+    assert round(ppl, 2) == 7.41
+
+
+def test_perplexity_short(model, metric):
+    # taken from https://huggingface.co/datasets/roneneldan/TinyStories
+    sample_text = "Once upon a time, there was a little car named Beep. Beep loved to go fast and play in the sun."
+    result = metric.compute(model, [sample_text])
+    ppl = result["score"]
+    assert round(ppl, 2) == 10.33
diff --git a/tests/test_process_labels_fusion.py b/tests/test_process_labels_fusion.py
new file mode 100644
index 0000000000..3f38134af7
--- /dev/null
+++ b/tests/test_process_labels_fusion.py
@@ -0,0 +1,640 @@
+"""Parity tests for fused ``process_labels``.
+
+Each strategy gets a legacy reference implementation (verbatim copy of the
+pre-refactor code) and a fuzz batch generator that injects realistic
+distributions of pad/image/etc tokens. We then assert byte-identical output
+between the legacy ``process_labels`` and the new fused implementation.
+"""
+
+from __future__ import annotations
+
+import random
+from types import SimpleNamespace
+from typing import List
+
+import pytest
+import torch
+
+from axolotl.processing_strategies import (
+    Gemma3nProcessingStrategy,
+    Gemma3ProcessingStrategy,
+    Gemma4ProcessingStrategy,
+    Glm4vProcessingStrategy,
+    InternVLProcessingStrategy,
+    Llama3_2VisionProcessingStrategy,
+    Llama4ProcessingStrategy,
+    Mistral3ProcessingStrategy,
+    MistralV7TekkenProcessingStrategy,
+    PixtralProcessingStrategy,
+    ProcessingStrategy,
+    Qwen2VLProcessingStrategy,
+    Qwen3_5ProcessingStrategy,
+    VoxtralProcessingStrategy,
+    _apply_role_boundaries,
+)
+
+# --------------------------------------------------------------------------- #
+# Tokenizer / processor stubs (mirrors test_processing_strategies.py)
+# --------------------------------------------------------------------------- #
+
+
+class _Tokenizer:
+    def __init__(
+        self,
+        vocab: dict,
+        pad_id: int = 0,
+        unk_id: int = 3,
+        eos_id: int | None = None,
+        extras: dict | None = None,
+    ):
+        self.vocab = vocab
+        self.pad_token_id = pad_id
+        self.unk_token_id = unk_id
+        if eos_id is not None:
+            self.eos_token_id = eos_id
+        if extras:
+            for k, v in extras.items():
+                setattr(self, k, v)
+
+    def encode(self, text, add_special_tokens=False):
+        return list(self.vocab.get(text, []))
+
+    def convert_tokens_to_ids(self, token):
+        v = self.vocab.get(token)
+        if v is None:
+            return self.unk_token_id
+        return v[0] if len(v) == 1 else self.unk_token_id
+
+
+class _Processor:
+    def __init__(self, tokenizer, **extras):
+        self.tokenizer = tokenizer
+        for k, v in extras.items():
+            setattr(self, k, v)
+
+
+# --------------------------------------------------------------------------- #
+# Legacy reference implementations (verbatim pre-refactor body)
+# --------------------------------------------------------------------------- #
+#
+# These functions implement process_labels exactly as it existed before the
+# fusion refactor: 3-4 sequential masked writes against the labels tensor,
+# with _mask_non_assistant materializing the role mask in-place.
+
+
+def _legacy_mask_non_assistant(strategy, labels):
+    """Pre-refactor ``_mask_non_assistant`` body."""
+    if strategy.train_on_inputs:
+        return labels
+    if not strategy.role_boundaries:
+        return labels
+    return _apply_role_boundaries(
+        labels,
+        strategy.role_boundaries,
+        roles_to_train=set(strategy.roles_to_train),
+        train_on_eos=strategy.train_on_eos,
+    )
+
+
+def legacy_base_process_labels(strategy, input_ids):
+    labels = input_ids.clone()
+    labels = _legacy_mask_non_assistant(strategy, labels)
+    pad_id = getattr(strategy.processor.tokenizer, "pad_token_id", None)
+    if pad_id is not None:
+        labels[labels == pad_id] = -100
+    if strategy.image_token_id is not None:
+        labels[labels == strategy.image_token_id] = -100
+    return labels
+
+
+def legacy_qwen3_5_process_labels(strategy, input_ids):
+    labels = legacy_base_process_labels(strategy, input_ids)
+    if strategy.video_token_id is not None:
+        labels[labels == strategy.video_token_id] = -100
+    return labels
+
+
+def legacy_gemma3_process_labels(strategy, input_ids):
+    labels = legacy_base_process_labels(strategy, input_ids)
+    tok = strategy.processor.tokenizer
+    soft_id = tok.convert_tokens_to_ids("<image_soft_token>")
+    unk_id = getattr(tok, "unk_token_id", None)
+    if soft_id is not None and soft_id != unk_id:
+        labels[labels == soft_id] = -100
+    else:
+        labels[labels == 262144] = -100
+    return labels
+
+
+def legacy_gemma3n_process_labels(strategy, input_ids):
+    labels = legacy_base_process_labels(strategy, input_ids)
+    tok = strategy.processor.tokenizer
+    for attr in (
+        "image_token_id",
+        "audio_token_id",
+        "boi_token_id",
+        "eoi_token_id",
+    ):
+        tok_id = getattr(tok, attr, None)
+        if tok_id is not None:
+            labels[labels == tok_id] = -100
+    return labels
+
+
+def legacy_gemma4_process_labels(strategy, input_ids):
+    labels = legacy_base_process_labels(strategy, input_ids)
+    tokenizer = strategy.processor.tokenizer
+    unk_id = getattr(tokenizer, "unk_token_id", None)
+    if getattr(tokenizer, "image_token_id", None) is not None:
+        labels[labels == tokenizer.image_token_id] = -100
+    if getattr(tokenizer, "audio_token_id", None) is not None:
+        labels[labels == tokenizer.audio_token_id] = -100
+    for attr in ("boi_token", "eoi_token", "boa_token", "eoa_token"):
+        token_str = getattr(strategy.processor, attr, None)
+        if token_str is None:
+            continue
+        token_id = tokenizer.convert_tokens_to_ids(token_str)
+        if token_id is None or token_id == unk_id:
+            continue
+        labels[labels == token_id] = -100
+    video_token_id = getattr(strategy.processor, "video_token_id", None)
+    if video_token_id is not None and video_token_id != unk_id:
+        labels[labels == video_token_id] = -100
+    return labels
+
+
+def legacy_voxtral_process_labels(strategy, input_ids):
+    labels = input_ids.clone()
+    labels = _legacy_mask_non_assistant(strategy, labels)
+    pad_id = getattr(strategy.processor.tokenizer, "pad_token_id", None)
+    if pad_id is not None:
+        labels[labels == pad_id] = -100
+    if strategy.audio_token is not None:
+        labels[labels == strategy.audio_token] = -100
+    if strategy.begin_audio_token is not None:
+        labels[labels == strategy.begin_audio_token] = -100
+    return labels
+
+
+def legacy_mistral3_process_labels(strategy, input_ids):
+    labels = input_ids.clone()
+    labels = _legacy_mask_non_assistant(strategy, labels)
+    pad_id = getattr(strategy.processor.tokenizer, "pad_token_id", None)
+    if pad_id is not None:
+        labels[labels == pad_id] = -100
+    for tok_id in (
+        strategy.image_token,
+        strategy.image_break_token,
+        strategy.image_end_token,
+    ):
+        if tok_id is not None:
+            labels[labels == tok_id] = -100
+    return labels
+
+
+def legacy_internvl_process_labels(strategy, input_ids):
+    labels = input_ids.clone()
+    labels = _legacy_mask_non_assistant(strategy, labels)
+    pad_id = getattr(strategy.processor.tokenizer, "pad_token_id", None)
+    if pad_id is not None:
+        labels[labels == pad_id] = -100
+    for ids in strategy.image_token_ids:
+        if ids is not None:
+            labels[labels == ids] = -100
+    return labels
+
+
+def legacy_glm4v_process_labels(strategy, input_ids):
+    labels = input_ids.clone()
+    labels = _legacy_mask_non_assistant(strategy, labels)
+    pad_id = getattr(strategy.tokenizer, "pad_token_id", None)
+    if pad_id is not None:
+        labels[labels == pad_id] = -100
+    for tok_id in (
+        strategy.image_token_id,
+        strategy.begin_image_token_id,
+        strategy.end_image_token_id,
+        strategy.video_token_id,
+        strategy.begin_video_token_id,
+        strategy.end_video_token_id,
+    ):
+        if tok_id is not None:
+            keep = labels != tok_id  # noqa: F841 — keep var unused; legacy uses indexed write
+            labels[labels == tok_id] = -100
+    return labels
+
+
+# --------------------------------------------------------------------------- #
+# Batch construction with realistic distributions
+# --------------------------------------------------------------------------- #
+
+
+def _build_alternating_turns(
+    rng: random.Random,
+    boundaries,
+    target_len: int,
+    pad_id: int,
+    extra_token_ids: List[int],
+    pad_rate: float = 0.07,
+    extra_rate: float = 0.07,
+    filler_min: int = 6,
+    filler_max: int = 24,
+) -> List[int]:
+    """Build a single row of length ``target_len`` from ``boundaries``.
+
+    Alternates user/assistant turns (or whatever's first/second in the list),
+    with random filler IDs that may be replaced by pad / extra (image, etc)
+    tokens at ``pad_rate`` / ``extra_rate``. The injection happens *inside*
+    assistant spans too, which is the realistic case.
+    """
+    role_b = list(boundaries)
+    ids: List[int] = []
+    if not role_b:
+        # No boundaries (Voxtral/Mistral3/InternVL/Glm4v): still inject pad and
+        # media tokens so the fused vs legacy parity check exercises the
+        # pad/media masking path even with role-masking opted out.
+        for _ in range(target_len):
+            r = rng.random()
+            if r < pad_rate:
+                ids.append(pad_id)
+            elif r < pad_rate + extra_rate and extra_token_ids:
+                ids.append(rng.choice(extra_token_ids))
+            else:
+                ids.append(rng.randint(10000, 30000))
+        return ids
+
+    turn = 0
+    while len(ids) < target_len:
+        b = role_b[turn % len(role_b)]
+        ids.extend(b.start_tokens)
+        n_filler = rng.randint(filler_min, filler_max)
+        for _ in range(n_filler):
+            r = rng.random()
+            if r < pad_rate:
+                ids.append(pad_id)
+            elif r < pad_rate + extra_rate and extra_token_ids:
+                ids.append(rng.choice(extra_token_ids))
+            else:
+                # Use IDs well outside the boundary token space
+                ids.append(rng.randint(10000, 30000))
+        if b.end_tokens:
+            ids.extend(b.end_tokens)
+        turn += 1
+    return ids[:target_len]
+
+
+def _build_batch(rng, boundaries, batch_size, seq_len, pad_id, extra_ids):
+    rows = []
+    for _ in range(batch_size):
+        # Vary length per row a bit so we get padding-like rows occasionally
+        target = rng.randint(seq_len // 2, seq_len)
+        row = _build_alternating_turns(rng, boundaries, target, pad_id, extra_ids)
+        # Pad up to seq_len with pad_id (mimics collator behavior)
+        if len(row) < seq_len:
+            row.extend([pad_id] * (seq_len - len(row)))
+        rows.append(row[:seq_len])
+    return torch.tensor(rows, dtype=torch.long)
+
+
+# --------------------------------------------------------------------------- #
+# Per-strategy fixtures
+# --------------------------------------------------------------------------- #
+
+
+def _make_qwen2vl():
+    vocab = {
+        "<|im_start|>system\n": [101, 102, 103],
+        "<|im_start|>user\n": [101, 104, 103],
+        "<|im_start|>assistant\n": [101, 105, 103],
+        "<|im_end|>": [106],
+        "<|image_pad|>": [200],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    return Qwen2VLProcessingStrategy(_Processor(tok))
+
+
+def _make_qwen3_5():
+    vocab = {
+        "<|im_start|>system\n": [101, 102, 103],
+        "<|im_start|>user\n": [101, 104, 103],
+        "<|im_start|>assistant\n": [101, 105, 103],
+        "<|im_end|>": [106],
+        "<|image_pad|>": [200],
+        "<|video_pad|>": [201],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    return Qwen3_5ProcessingStrategy(_Processor(tok))
+
+
+def _make_gemma3():
+    vocab = {
+        "<start_of_turn>user\n": [110, 111, 112],
+        "<start_of_turn>model\n": [110, 113, 112],
+        "<end_of_turn>": [114],
+        "<image_soft_token>": [262144],
+        "<boi>": [120],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3, extras={"boi_token": "<boi>"})
+    return Gemma3ProcessingStrategy(_Processor(tok))
+
+
+def _make_gemma3n():
+    vocab = {
+        "<start_of_turn>user\n": [110, 111, 112],
+        "<start_of_turn>model\n": [110, 113, 112],
+        "<end_of_turn>": [114],
+    }
+    tok = _Tokenizer(
+        vocab,
+        pad_id=0,
+        unk_id=3,
+        extras={
+            "image_token_id": 130,
+            "audio_token_id": 131,
+            "boi_token_id": 132,
+            "eoi_token_id": 133,
+        },
+    )
+    return Gemma3nProcessingStrategy(_Processor(tok))
+
+
+def _make_gemma4():
+    vocab = {
+        "<|turn>user\n": [105, 4001],
+        "<|turn>system\n": [105, 4002],
+        "<|turn>model\n": [105, 4368],
+        "<turn|>": [106],
+        "<boi>": [140],
+        "<eoi>": [141],
+        "<boa>": [142],
+        "<eoa>": [143],
+    }
+    tok = _Tokenizer(
+        vocab,
+        pad_id=0,
+        unk_id=3,
+        extras={"image_token_id": 150, "audio_token_id": 151},
+    )
+    proc = _Processor(
+        tok,
+        boi_token="<boi>",
+        eoi_token="<eoi>",
+        boa_token="<boa>",
+        eoa_token="<eoa>",
+        video_token_id=160,
+    )
+    return Gemma4ProcessingStrategy(proc)
+
+
+def _make_llama32v():
+    vocab = {
+        "<|start_header_id|>system<|end_header_id|>\n\n": [1001, 2002, 1002, 1003],
+        "<|start_header_id|>user<|end_header_id|>\n\n": [1001, 2001, 1002, 1003],
+        "<|start_header_id|>assistant<|end_header_id|>\n\n": [1001, 2003, 1002, 1003],
+        "<|eot_id|>": [1004],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    return Llama3_2VisionProcessingStrategy(_Processor(tok))
+
+
+def _make_llama4():
+    vocab = {
+        "<|header_start|>system<|header_end|>\n\n": [1001, 2002, 1002, 1003],
+        "<|header_start|>user<|header_end|>\n\n": [1001, 2001, 1002, 1003],
+        "<|header_start|>assistant<|header_end|>\n\n": [1001, 2003, 1002, 1003],
+        "<|eot|>": [1004],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    return Llama4ProcessingStrategy(_Processor(tok))
+
+
+def _make_pixtral():
+    vocab = {"[INST]": [50], "[/INST]": [51]}
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3, eos_id=99)
+    return PixtralProcessingStrategy(_Processor(tok))
+
+
+def _make_mistral_v7():
+    vocab = {
+        "[INST]": [50],
+        "[/INST]": [51],
+        "[SYSTEM_PROMPT]": [60],
+        "[/SYSTEM_PROMPT]": [61],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3, eos_id=99)
+    return MistralV7TekkenProcessingStrategy(_Processor(tok))
+
+
+# --------------------------------------------------------------------------- #
+# Strategies that opt out of role-masking (no role_boundaries declared)
+# --------------------------------------------------------------------------- #
+#
+# These four fall back to pad+media masking only. We still want parity fuzz
+# coverage of the fused vs legacy paths because the per-strategy media-token
+# combinations differ (audio for Voxtral, three image markers for Mistral3,
+# variable-length image_ids list for InternVL, six markers for Glm4v).
+
+
+def _make_voxtral():
+    """VoxtralProcessor stub.
+
+    The strategy resolves audio token ids via the nested mistral-common path
+    ``processor.tokenizer.tokenizer.instruct_tokenizer.audio_encoder.special_ids``.
+    Mirror just enough of that shape for __init__ to read the two ids.
+    """
+    tok = _Tokenizer({}, pad_id=0, unk_id=3)
+    audio_encoder = SimpleNamespace(
+        special_ids=SimpleNamespace(audio=300, begin_audio=301)
+    )
+    tok.tokenizer = SimpleNamespace(
+        instruct_tokenizer=SimpleNamespace(audio_encoder=audio_encoder)
+    )
+    return VoxtralProcessingStrategy(_Processor(tok))
+
+
+def _make_mistral3():
+    """Mistral3 stub: image_encoder.special_ids carries (img, img_break, img_end)."""
+    tok = _Tokenizer({}, pad_id=0, unk_id=3)
+    image_encoder = SimpleNamespace(
+        special_ids=SimpleNamespace(img=400, img_break=401, img_end=402)
+    )
+    tok.tokenizer = SimpleNamespace(
+        instruct_tokenizer=SimpleNamespace(image_encoder=image_encoder)
+    )
+    return Mistral3ProcessingStrategy(_Processor(tok))
+
+
+def _make_internvl():
+    """InternVL stub: processor.image_ids is a list of media-token ids."""
+    tok = _Tokenizer({}, pad_id=0, unk_id=3)
+    proc = _Processor(tok)
+    proc.image_ids = [500, 501, 502]
+    return InternVLProcessingStrategy(proc)
+
+
+def _make_glm4v():
+    """Glm4v / Glm46V stub: convert_tokens_to_ids resolves six media markers."""
+    vocab = {
+        "<|image|>": [600],
+        "<|begin_of_image|>": [601],
+        "<|end_of_image|>": [602],
+        "<|video|>": [610],
+        "<|begin_of_video|>": [611],
+        "<|end_of_video|>": [612],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    return Glm4vProcessingStrategy(_Processor(tok))
+
+
+# --------------------------------------------------------------------------- #
+# Parameterized parity fuzz
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    "strategy_factory, legacy_fn, extra_ids, name",
+    [
+        (_make_qwen2vl, legacy_base_process_labels, [200], "qwen2vl"),
+        (_make_qwen3_5, legacy_qwen3_5_process_labels, [200, 201], "qwen3_5"),
+        (_make_gemma3, legacy_gemma3_process_labels, [120, 262144], "gemma3"),
+        (_make_gemma3n, legacy_gemma3n_process_labels, [130, 131, 132, 133], "gemma3n"),
+        (
+            _make_gemma4,
+            legacy_gemma4_process_labels,
+            [140, 141, 142, 143, 150, 151, 160],
+            "gemma4",
+        ),
+        (_make_llama32v, legacy_base_process_labels, [], "llama32v"),
+        (_make_llama4, legacy_base_process_labels, [], "llama4"),
+        (_make_pixtral, legacy_base_process_labels, [], "pixtral"),
+        (_make_mistral_v7, legacy_base_process_labels, [], "mistral_v7"),
+        (_make_voxtral, legacy_voxtral_process_labels, [300, 301], "voxtral"),
+        (
+            _make_mistral3,
+            legacy_mistral3_process_labels,
+            [400, 401, 402],
+            "mistral3",
+        ),
+        (
+            _make_internvl,
+            legacy_internvl_process_labels,
+            [500, 501, 502],
+            "internvl",
+        ),
+        (
+            _make_glm4v,
+            legacy_glm4v_process_labels,
+            [600, 601, 602, 610, 611, 612],
+            "glm4v",
+        ),
+    ],
+)
+@pytest.mark.parametrize("seed", list(range(50)))
+def test_process_labels_parity_fuzz(strategy_factory, legacy_fn, extra_ids, name, seed):
+    """50 random batches per strategy: legacy vs fused must produce byte-identical labels."""
+    rng = random.Random(seed * 7919 + sum(ord(c) for c in name) % 1000)
+
+    strategy = strategy_factory()
+    pad_id = strategy.processor.tokenizer.pad_token_id
+
+    # Mix of small / medium / large, plus varied batch sizes
+    batch_size = rng.choice([1, 2, 4, 8])
+    seq_len = rng.choice([64, 128, 256, 512])
+
+    batch = _build_batch(
+        rng, strategy.role_boundaries, batch_size, seq_len, pad_id, extra_ids
+    )
+
+    expected = legacy_fn(strategy, batch.clone())
+    got = strategy.process_labels(batch.clone())
+
+    if not torch.equal(expected, got):
+        # Find first divergence to make debug output useful
+        diff = (expected != got).nonzero(as_tuple=False)
+        first = diff[0].tolist()
+        b_idx, t_idx = first
+        lo = max(0, t_idx - 8)
+        hi = min(seq_len, t_idx + 8)
+        raise AssertionError(
+            f"[{name} seed={seed}] mismatch at row={b_idx} col={t_idx}\n"
+            f"  input_ids[{b_idx}, {lo}:{hi}] = {batch[b_idx, lo:hi].tolist()}\n"
+            f"  legacy   [{b_idx}, {lo}:{hi}] = {expected[b_idx, lo:hi].tolist()}\n"
+            f"  fused    [{b_idx}, {lo}:{hi}] = {got[b_idx, lo:hi].tolist()}\n"
+            f"  boundaries: {strategy.role_boundaries}"
+        )
+
+
+def test_process_labels_train_on_inputs_passthrough():
+    """When train_on_inputs=True the keep mask must be all-True (only pad/media masked)."""
+    vocab = {
+        "<|im_start|>user\n": [101, 104, 103],
+        "<|im_start|>assistant\n": [101, 105, 103],
+        "<|im_end|>": [106],
+        "<|image_pad|>": [200],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    s = Qwen2VLProcessingStrategy(_Processor(tok), train_on_inputs=True)
+    batch = torch.tensor([[101, 104, 103, 7, 8, 0, 200, 9, 106]])
+    out = s.process_labels(batch.clone())
+    expected = legacy_base_process_labels(s, batch.clone())
+    assert torch.equal(out, expected)
+    # And confirm the only masked positions are pad and image
+    assert out.tolist() == [[101, 104, 103, 7, 8, -100, -100, 9, 106]]
+
+
+def test_process_labels_no_boundaries_falls_through():
+    """A boundary-less strategy + train_on_inputs=False keeps everything (warns once);
+    pad / media tokens are still masked.
+    """
+    vocab = {"foo": [200]}
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    # Base strategy with no role_boundaries declared and train_on_inputs=False
+    s = ProcessingStrategy(_Processor(tok))
+    assert s.role_boundaries == []
+    batch = torch.tensor([[5, 6, 7, 0, 8]])
+    out = s.process_labels(batch.clone())
+    expected = legacy_base_process_labels(s, batch.clone())
+    assert torch.equal(out, expected)
+    # Pad at index 3 must be masked, everything else kept (no boundaries → no role mask).
+    assert out.tolist() == [[5, 6, 7, -100, 8]]
+
+
+# --------------------------------------------------------------------------- #
+# Targeted edge cases for fused path
+# --------------------------------------------------------------------------- #
+
+
+def test_fused_pad_inside_assistant_is_masked():
+    """Pad inside assistant span must end up -100 even though role-mask kept it."""
+    vocab = {
+        "<|im_start|>user\n": [101, 104, 103],
+        "<|im_start|>assistant\n": [101, 105, 103],
+        "<|im_end|>": [106],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    s = Qwen2VLProcessingStrategy(_Processor(tok))
+    seq = [101, 104, 103, 7, 106, 101, 105, 103, 8, 0, 9, 106]
+    out = s.process_labels(torch.tensor([seq]))
+    expected = legacy_base_process_labels(s, torch.tensor([seq]))
+    assert torch.equal(out, expected)
+    # Pad (0) inside assistant gets masked, content (8, 9) and end (106) kept.
+    assert out.tolist() == [
+        [-100, -100, -100, -100, -100, -100, -100, -100, 8, -100, 9, 106]
+    ]
+
+
+def test_fused_image_inside_assistant_is_masked():
+    vocab = {
+        "<|im_start|>user\n": [101, 104, 103],
+        "<|im_start|>assistant\n": [101, 105, 103],
+        "<|im_end|>": [106],
+        "<|image_pad|>": [200],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, unk_id=3)
+    s = Qwen2VLProcessingStrategy(_Processor(tok))
+    seq = [101, 105, 103, 8, 200, 9, 106]
+    out = s.process_labels(torch.tensor([seq]))
+    expected = legacy_base_process_labels(s, torch.tensor([seq]))
+    assert torch.equal(out, expected)
+    # image_token_id (200) gets masked even though it's inside the assistant span.
+    assert out.tolist() == [[-100, -100, -100, 8, -100, 9, 106]]
diff --git a/tests/test_processing_strategies.py b/tests/test_processing_strategies.py
new file mode 100644
index 0000000000..155b17446e
--- /dev/null
+++ b/tests/test_processing_strategies.py
@@ -0,0 +1,1583 @@
+"""Tests for ``axolotl.processing_strategies`` using fake tokenizers (offline/CI-safe)."""
+
+import logging
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from pydantic import ValidationError
+
+from axolotl.processing_strategies import (
+    Gemma3nProcessingStrategy,
+    Gemma3ProcessingStrategy,
+    Gemma4ProcessingStrategy,
+    Gemma4UnifiedProcessingStrategy,
+    Glm4vProcessingStrategy,
+    InternVLProcessingStrategy,
+    Llama3_2VisionProcessingStrategy,
+    Llama4ProcessingStrategy,
+    Mistral3ProcessingStrategy,
+    MistralV7TekkenProcessingStrategy,
+    PixtralProcessingStrategy,
+    ProcessingStrategy,
+    Qwen2VLProcessingStrategy,
+    Qwen3_5ProcessingStrategy,
+    RoleBoundary,
+    VoxtralProcessingStrategy,
+    _apply_role_boundaries,
+    get_processing_strategy,
+)
+
+
+@pytest.fixture
+def axolotl_caplog(caplog):
+    """caplog that also captures records from the ``axolotl`` logger.
+
+    The axolotl logger sets ``propagate=False`` once ``configure_logging()`` is
+    called (which happens indirectly in many CI test paths), so the default
+    caplog handler installed on the root logger never sees these records.
+    Attaching ``caplog.handler`` to ``axolotl.processing_strategies`` directly
+    makes assertions reliable regardless of whether ``configure_logging`` has
+    already run on this worker.
+    """
+    logger = logging.getLogger("axolotl.processing_strategies")
+    logger.addHandler(caplog.handler)
+    previous_level = logger.level
+    logger.setLevel(logging.DEBUG)
+    try:
+        yield caplog
+    finally:
+        logger.removeHandler(caplog.handler)
+        logger.setLevel(previous_level)
+
+
+# --------------------------------------------------------------------------- #
+# Generic fake tokenizer/processor scaffold
+# --------------------------------------------------------------------------- #
+
+
+class _Tokenizer:
+    """Minimal tokenizer stub; ``vocab`` maps marker strings to their id lists."""
+
+    def __init__(
+        self,
+        vocab: dict[str, list[int]],
+        pad_id: int = 0,
+        unk_id: int = 3,
+        eos_id: int | None = None,
+    ):
+        self.vocab = vocab
+        self._reverse = {}
+        for tok, ids in vocab.items():
+            if len(ids) == 1:
+                self._reverse[ids[0]] = tok
+        self.pad_token_id = pad_id
+        self.unk_token_id = unk_id
+        if eos_id is not None:
+            self.eos_token_id = eos_id
+
+    def encode(self, text, add_special_tokens=False):
+        # Unknown markers return [] so _encode_markers drops them silently.
+        return list(self.vocab.get(text, []))
+
+    def convert_tokens_to_ids(self, token):
+        v = self.vocab.get(token)
+        if v is None:
+            return self.unk_token_id
+        return v[0] if len(v) == 1 else self.unk_token_id
+
+
+class _Processor:
+    def __init__(self, tokenizer: _Tokenizer):
+        self.tokenizer = tokenizer
+
+
+# --------------------------------------------------------------------------- #
+# Base scanner tests (train_on_inputs / roles_to_train / train_on_eos)
+# --------------------------------------------------------------------------- #
+
+
+def _scan(role_boundaries, seq, roles_to_train=("assistant",), train_on_eos="turn"):
+    labels = torch.tensor([seq])
+    return _apply_role_boundaries(
+        labels, role_boundaries, set(roles_to_train), train_on_eos
+    ).tolist()[0]
+
+
+def test_scanner_assistant_only_basic():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+        RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]),
+    ]
+    seq = [1, 3, 7, 9, 1, 2, 8, 8, 9, 5]
+    out = _scan(boundaries, seq)
+    assert out == [-100, -100, -100, -100, -100, -100, 8, 8, 9, -100]
+
+
+def test_scanner_train_on_eos_none_excludes_end_marker():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+    ]
+    seq = [1, 2, 8, 8, 9]
+    out = _scan(boundaries, seq, train_on_eos="none")
+    assert out == [-100, -100, 8, 8, -100]
+
+
+def test_scanner_train_on_eos_all_keeps_non_assistant_end_marker():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+        RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]),
+    ]
+    seq = [1, 3, 7, 9, 1, 2, 8, 9]
+    out = _scan(boundaries, seq, train_on_eos="all")
+    assert out == [-100, -100, -100, 9, -100, -100, 8, 9]
+
+
+def test_scanner_train_on_eos_all_with_non_trainable_include_end_false():
+    """Non-trainable + include_end=False must not leak end marker on 'all'."""
+    boundaries = [
+        RoleBoundary(
+            role="user",
+            start_tokens=[50],
+            end_tokens=[51],
+            include_end=False,  # shared with assistant-start
+        ),
+        RoleBoundary(
+            role="assistant",
+            start_tokens=[51],
+            end_tokens=[99],  # eos
+        ),
+    ]
+    seq = [50, 7, 51, 8, 8, 99]  # [INST] 7 [/INST] 8 8 EOS
+    out = _scan(boundaries, seq, roles_to_train=("assistant",), train_on_eos="all")
+    # [/INST] at idx 2 must stay masked — user.include_end=False says so.
+    assert out == [-100, -100, -100, 8, 8, 99]
+
+
+def test_scanner_roles_to_train_user_and_assistant():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+        RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]),
+    ]
+    seq = [1, 3, 7, 9, 1, 2, 8, 9]
+    out = _scan(boundaries, seq, roles_to_train=("user", "assistant"))
+    # include_start defaults to False so role-start markers stay masked.
+    assert out == [-100, -100, 7, 9, -100, -100, 8, 9]
+
+
+def test_scanner_truncated_assistant():
+    """Missing end marker: span runs to end-of-sequence, end marker not emitted."""
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+    ]
+    seq = [1, 2, 8, 8, 8]
+    out = _scan(boundaries, seq)
+    assert out == [-100, -100, 8, 8, 8]
+
+
+def test_scanner_longest_prefix_wins():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2, 4], end_tokens=[9]),
+        RoleBoundary(role="user", start_tokens=[1, 2], end_tokens=[9]),
+    ]
+    seq = [1, 2, 4, 8, 9]
+    out = _scan(boundaries, seq)
+    assert out == [-100, -100, -100, 8, 9]
+
+
+def test_scanner_no_boundaries_masks_everything():
+    # Strategies short-circuit this in _mask_non_assistant; see test_base_strategy_warns_when_no_boundaries.
+    labels = torch.tensor([[1, 2, 3, 4]])
+    out = _apply_role_boundaries(labels, [], {"assistant"}, "turn")
+    assert out.tolist() == [[-100, -100, -100, -100]]
+
+
+def test_scanner_train_on_eos_last_only_final_trainable_turn():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+    ]
+    seq = [1, 2, 5, 9, 1, 2, 6, 9]
+    out = _scan(boundaries, seq, train_on_eos="last")
+    # Only the second assistant turn's end marker (index 7) is kept.
+    assert out == [-100, -100, 5, -100, -100, -100, 6, 9]
+
+
+def test_scanner_train_on_eos_last_no_trainable_turn_is_noop():
+    boundaries = [
+        RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]),
+    ]
+    seq = [1, 3, 5, 9, 1, 3, 6, 9]
+    out = _scan(boundaries, seq, roles_to_train=("assistant",), train_on_eos="last")
+    assert out == [-100] * 8
+
+
+def test_strategy_rejects_unknown_train_on_eos():
+    vocab = {"BOA": [50], "EOT": [60]}
+    with pytest.raises(ValueError, match="train_on_eos"):
+        ProcessingStrategy(
+            _Processor(_Tokenizer(vocab, pad_id=0)),
+            train_on_eos="bogus",
+            role_boundaries_override=[
+                {"role": "assistant", "start": "BOA", "end": "EOT"}
+            ],
+        )
+
+
+def test_strategy_accepts_all_supported_train_on_eos_values():
+    vocab = {"BOA": [50], "EOT": [60]}
+    for val in ("turn", "all", "none", "last"):
+        ProcessingStrategy(
+            _Processor(_Tokenizer(vocab, pad_id=0)),
+            train_on_eos=val,
+            role_boundaries_override=[
+                {"role": "assistant", "start": "BOA", "end": "EOT"}
+            ],
+        )
+
+
+def test_empty_role_boundaries_override_falls_back_to_builtin():
+    """Empty override must fall through to built-ins (opt-in semantics)."""
+    vocab = {
+        "<|im_start|>assistant\n": [101, 102, 103],
+        "<|im_start|>user\n": [101, 106, 103],
+        "<|im_end|>": [104],
+    }
+    strat_empty = Qwen2VLProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[],
+    )
+    strat_default = Qwen2VLProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+    )
+    # Empty override === no override: both strategies keep the built-in boundaries.
+    assert strat_empty.role_boundaries == strat_default.role_boundaries
+    assert len(strat_empty.role_boundaries) > 0  # sanity: built-ins are non-empty
+
+
+def test_sft_dataset_schema_accepts_all_supported_train_on_eos_values():
+    """SFTDataset.train_on_eos accepts every value the scanner honors."""
+    from axolotl.utils.schemas.datasets import SFTDataset
+
+    for val in ("all", "turn", "last", "none"):
+        ds = SFTDataset(path="dummy", type="chat_template", train_on_eos=val)
+        assert ds.train_on_eos == val
+
+    with pytest.raises(ValidationError):
+        SFTDataset(path="dummy", type="chat_template", train_on_eos="bogus")
+
+
+def test_strategy_init_logs_resolved_masking_config_builtin(axolotl_caplog):
+    vocab = {
+        "<|im_start|>assistant\n": [101, 102, 103],
+        "<|im_start|>user\n": [101, 106, 103],
+        "<|im_end|>": [104],
+    }
+    with axolotl_caplog.at_level(logging.INFO, logger="axolotl.processing_strategies"):
+        Qwen2VLProcessingStrategy(_Processor(_Tokenizer(vocab, pad_id=0)))
+    msgs = [r.getMessage() for r in axolotl_caplog.records]
+    assert any(
+        "ProcessingStrategy init" in m
+        and "Qwen2VLProcessingStrategy" in m
+        and "boundaries_source=built-in" in m
+        for m in msgs
+    )
+
+
+def test_strategy_init_logs_resolved_masking_config_override(axolotl_caplog):
+    vocab = {"BOA": [50, 51], "EOT": [60]}
+    with axolotl_caplog.at_level(logging.INFO, logger="axolotl.processing_strategies"):
+        ProcessingStrategy(
+            _Processor(_Tokenizer(vocab, pad_id=0)),
+            role_boundaries_override=[
+                {"role": "assistant", "start": "BOA", "end": "EOT"},
+            ],
+        )
+    msgs = [r.getMessage() for r in axolotl_caplog.records]
+    # Resolved start/end ids must appear in the log so users can verify what
+    # was actually matched.
+    assert any(
+        "ProcessingStrategy init" in m
+        and "boundaries_source=override" in m
+        and "[50, 51]" in m
+        and "[60]" in m
+        for m in msgs
+    )
+
+
+def test_process_labels_no_warning_when_image_token_id_none():
+    """image_token_id=None must not trigger a UserWarning from ``labels == None``."""
+    import warnings
+
+    vocab = {"BOA": [50], "EOT": [60]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[{"role": "assistant", "start": "BOA", "end": "EOT"}],
+    )
+    assert strategy.image_token_id is None
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        strategy.process_labels(torch.tensor([[1, 50, 2, 3, 60]]))
+
+
+def test_roles_to_train_empty_list_masks_everything():
+    """An explicit empty list is distinct from None and disables all roles."""
+    vocab = {"BOA": [50], "EOT": [60]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        roles_to_train=[],
+        role_boundaries_override=[{"role": "assistant", "start": "BOA", "end": "EOT"}],
+    )
+    assert strategy.roles_to_train == []
+    seq = [1, 50, 7, 8, 60, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100] * 6
+
+
+# --------------------------------------------------------------------------- #
+# Qwen2VL / Qwen3.5
+# --------------------------------------------------------------------------- #
+
+
+def _qwen_tokenizer():
+    # ChatML-ish with image_pad=200, video_pad=201.
+    vocab = {
+        "<|im_start|>assistant\n": [101, 102, 103],
+        "<|im_start|>user\n": [101, 106, 103],
+        "<|im_start|>system\n": [101, 105, 103],
+        "<|im_end|>": [104],
+        "<|image_pad|>": [200],
+        "<|video_pad|>": [201],
+    }
+    return _Tokenizer(vocab, pad_id=0)
+
+
+def _make_qwen2vl():
+    tok = _qwen_tokenizer()
+    return Qwen2VLProcessingStrategy(_Processor(tok))
+
+
+def test_qwen2vl_masks_user_keeps_assistant_and_image_pad():
+    strategy = _make_qwen2vl()
+    seq = [
+        101,
+        105,
+        103,
+        77,
+        104,
+        101,
+        106,
+        103,
+        7,
+        104,
+        101,
+        102,
+        103,
+        200,
+        8,
+        104,
+    ]
+    labels = strategy.process_labels(torch.tensor([seq]))
+    out = labels.tolist()[0]
+    assert out[:10] == [-100] * 10
+    assert out[10] == -100 and out[11] == -100 and out[12] == -100
+    assert out[13] == -100  # image_pad masked post-scan
+    assert out[14] == 8
+    assert out[15] == 104
+
+
+def test_qwen3_5_masks_video_pad_too():
+    tok = _qwen_tokenizer()
+    strategy = Qwen3_5ProcessingStrategy(_Processor(tok))
+    seq = [101, 102, 103, 201, 8, 104]
+    labels = strategy.process_labels(torch.tensor([seq]))
+    assert labels.tolist()[0] == [-100, -100, -100, -100, 8, 104]
+
+
+def test_qwen2vl_train_on_inputs_true_keeps_everything():
+    tok = _qwen_tokenizer()
+    strategy = Qwen2VLProcessingStrategy(_Processor(tok), train_on_inputs=True)
+    seq = [101, 106, 103, 7, 104, 101, 102, 103, 8, 104]
+    labels = strategy.process_labels(torch.tensor([seq]))
+    assert labels.tolist()[0] == seq
+
+
+# --------------------------------------------------------------------------- #
+# Gemma3 / Gemma3n
+# --------------------------------------------------------------------------- #
+
+
+def _gemma_tokenizer():
+    vocab = {
+        "<start_of_turn>model\n": [1, 2, 3],
+        "<start_of_turn>user\n": [1, 10, 3],
+        "<start_of_turn>system\n": [1, 11, 3],
+        "<end_of_turn>": [4],
+        "<start_of_image>": [50],  # boi_token for Gemma3
+    }
+    tok = _Tokenizer(vocab, pad_id=0)
+    # boi_token is a direct tokenizer attribute on real Gemma3.
+    tok.boi_token = "<start_of_image>"
+    return tok
+
+
+def test_gemma3_scanner_plus_soft_image_token():
+    strategy = Gemma3ProcessingStrategy(_Processor(_gemma_tokenizer()))
+    seq = [1, 10, 3, 7, 4, 1, 2, 3, 50, 8, 262144, 4]
+    labels = strategy.process_labels(torch.tensor([seq]))
+    # boi(50) and soft-image-token(262144) masked post-scan.
+    assert labels.tolist()[0] == [
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        8,
+        -100,
+        4,
+    ]
+
+
+def test_gemma3n_masks_image_and_audio_attrs():
+    tok = _gemma_tokenizer()
+    # Gemma3n exposes these as integer attrs on the tokenizer.
+    tok.image_token_id = 70
+    tok.audio_token_id = 71
+    tok.boi_token_id = 72
+    tok.eoi_token_id = 73
+    strategy = Gemma3nProcessingStrategy(_Processor(tok))
+    seq = [1, 2, 3, 70, 71, 72, 73, 9, 4]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, -100, -100, -100, -100, -100, 9, 4]
+
+
+# --------------------------------------------------------------------------- #
+# Gemma 4
+# --------------------------------------------------------------------------- #
+
+
+class _FakeGemma4Tokenizer(_Tokenizer):
+    """Mirrors google/gemma-4-E2B-it token layout. Gemma4 role-start markers
+    include the trailing newline so the boundary matches the jinja template."""
+
+    VOCAB = {
+        "<|turn>model\n": [105, 4368, 108],
+        "<|turn>user\n": [105, 7777, 108],
+        "<|turn>system\n": [105, 8888, 108],
+        "<turn|>": [106],
+        "<|image|>": [258880],
+        "<|video|>": [258884],
+        "<|audio|>": [258881],
+        "<|image>": [255999],
+        "<image|>": [258882],
+        "<|audio>": [256000],
+        "<audio|>": [258883],
+    }
+
+    def __init__(self):
+        # Pass a fresh dict so per-instance mutations (should any future
+        # code path introduce them) cannot leak across tests via the
+        # shared class-level VOCAB.
+        super().__init__(
+            {token: list(ids) for token, ids in self.VOCAB.items()},
+            pad_id=0,
+            unk_id=3,
+        )
+
+
+class _FakeGemma4Processor:
+    def __init__(self):
+        self.tokenizer = _FakeGemma4Tokenizer()
+        self.tokenizer.image_token_id = self.tokenizer.vocab["<|image|>"][0]
+        self.tokenizer.audio_token_id = self.tokenizer.vocab["<|audio|>"][0]
+        self.image_token = "<|image|>"
+        self.image_token_id = self.tokenizer.vocab["<|image|>"][0]
+        self.boi_token = "<|image>"
+        self.eoi_token = "<image|>"
+        self.video_token = "<|video|>"
+        self.video_token_id = self.tokenizer.vocab["<|video|>"][0]
+        self.audio_token = "<|audio|>"
+        self.audio_token_id = self.tokenizer.vocab["<|audio|>"][0]
+        self.boa_token = "<|audio>"
+        self.eoa_token = "<audio|>"
+
+
+def test_gemma4_masks_everything_outside_assistant_span():
+    strategy = Gemma4ProcessingStrategy(_FakeGemma4Processor())
+    V = strategy.processor.tokenizer.vocab
+    user_start = V["<|turn>user\n"]
+    model_start = V["<|turn>model\n"]
+    turn_end = V["<turn|>"][0]
+    seq = [
+        0,
+        *user_start,
+        4444,
+        turn_end,
+        *model_start,
+        5555,
+        turn_end,
+        9999,
+    ]
+    labels = strategy.process_labels(torch.tensor([seq]))
+    expected = [-100] * (1 + len(user_start) + 1 + 1 + len(model_start)) + [
+        5555,
+        turn_end,
+        -100,
+    ]
+    assert labels.tolist()[0] == expected
+
+
+def test_gemma4_masks_media_tokens_inside_assistant_span():
+    strategy = Gemma4ProcessingStrategy(_FakeGemma4Processor())
+    V = strategy.processor.tokenizer.vocab
+    model_start = V["<|turn>model\n"]
+    media = [
+        V["<|image|>"][0],
+        V["<|video|>"][0],
+        V["<|audio|>"][0],
+        V["<|image>"][0],
+        V["<image|>"][0],
+        V["<|audio>"][0],
+        V["<audio|>"][0],
+    ]
+    turn_end = V["<turn|>"][0]
+    seq = [*model_start, *media, 9999, turn_end]
+    labels = strategy.process_labels(torch.tensor([seq]))
+    expected = [-100] * (len(model_start) + len(media)) + [9999, turn_end]
+    assert labels.tolist()[0] == expected
+
+
+def test_gemma4_multiple_assistant_turns():
+    strategy = Gemma4ProcessingStrategy(_FakeGemma4Processor())
+    V = strategy.processor.tokenizer.vocab
+    turn_end = V["<turn|>"][0]
+
+    def user_turn(x):
+        return [*V["<|turn>user\n"], x, turn_end]
+
+    def model_turn(x):
+        return [*V["<|turn>model\n"], x, turn_end]
+
+    seq = user_turn(1111) + model_turn(2222) + user_turn(3333) + model_turn(4444)
+    labels = strategy.process_labels(torch.tensor([seq]))
+    kept = [t for t in labels.tolist()[0] if t != -100]
+    assert kept == [2222, turn_end, 4444, turn_end]
+
+
+# --------------------------------------------------------------------------- #
+# Llama 3.2 Vision / Llama 4
+# --------------------------------------------------------------------------- #
+
+
+def test_llama3_2_vision_assistant_masking():
+    vocab = {
+        "<|start_header_id|>assistant<|end_header_id|>\n\n": [1, 2, 3, 4, 5],
+        "<|start_header_id|>user<|end_header_id|>\n\n": [1, 2, 6, 4, 5],
+        "<|start_header_id|>system<|end_header_id|>\n\n": [1, 2, 7, 4, 5],
+        "<|start_header_id|>tool<|end_header_id|>\n\n": [1, 2, 8, 4, 5],
+        "<|start_header_id|>ipython<|end_header_id|>\n\n": [1, 2, 9, 4, 5],
+        "<|eot_id|>": [10],
+    }
+    strategy = Llama3_2VisionProcessingStrategy(_Processor(_Tokenizer(vocab, pad_id=0)))
+    seq = [1, 2, 6, 4, 5, 11, 10, 1, 2, 3, 4, 5, 12, 10]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100] * 12 + [12, 10]
+
+
+def test_llama4_assistant_masking():
+    vocab = {
+        "<|header_start|>assistant<|header_end|>\n\n": [20, 21, 22, 23],
+        "<|header_start|>user<|header_end|>\n\n": [20, 21, 24, 23],
+        "<|header_start|>system<|header_end|>\n\n": [20, 21, 25, 23],
+        "<|header_start|>tool<|header_end|>\n\n": [20, 21, 26, 23],
+        "<|header_start|>ipython<|header_end|>\n\n": [20, 21, 27, 23],
+        "<|eot|>": [30],
+    }
+    strategy = Llama4ProcessingStrategy(_Processor(_Tokenizer(vocab, pad_id=0)))
+    seq = [20, 21, 24, 23, 100, 30, 20, 21, 22, 23, 200, 30]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100] * 10 + [200, 30]
+
+
+# --------------------------------------------------------------------------- #
+# Pixtral / Mistral v7 Tekken (eos-terminated assistant)
+# --------------------------------------------------------------------------- #
+
+
+def test_pixtral_assistant_terminates_at_eos():
+    # [/INST] is both user-end and assistant-start. Scanner backs up when
+    # user.include_end=False so the next iteration picks [/INST] up as
+    # assistant-start (Pixtral-specific handling in _build_role_boundaries).
+    vocab = {
+        "[INST]": [50],
+        "[/INST]": [51],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, eos_id=99)
+    strategy = PixtralProcessingStrategy(_Processor(tok))
+    seq = [50, 7, 51, 8, 8, 99]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    # Full-sequence expectation: user span masked; assistant content + eos kept.
+    assert out == [-100, -100, -100, 8, 8, 99]
+
+
+def test_mistral_v7_tekken_system_user_assistant():
+    vocab = {
+        "[SYSTEM_PROMPT]": [40],
+        "[/SYSTEM_PROMPT]": [41],
+        "[INST]": [50],
+        "[/INST]": [51],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, eos_id=99)
+    strategy = MistralV7TekkenProcessingStrategy(_Processor(tok))
+    seq = [40, 5, 41, 50, 7, 51, 8, 99]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    # Full-sequence expectation: system + user spans masked; assistant kept.
+    assert out == [-100, -100, -100, -100, -100, -100, 8, 99]
+
+
+def test_pixtral_train_on_eos_all_respects_user_include_end_false():
+    """Pixtral [/INST] (user-end include_end=False) stays masked on 'all'."""
+    vocab = {"[INST]": [50], "[/INST]": [51]}
+    tok = _Tokenizer(vocab, pad_id=0, eos_id=99)
+    strategy = PixtralProcessingStrategy(_Processor(tok), train_on_eos="all")
+    seq = [50, 7, 51, 8, 8, 99]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    # [/INST] at idx 2 must stay masked — user.include_end=False says so.
+    # Assistant content (8, 8) + EOS (99) are unmasked as normal.
+    assert out == [-100, -100, -100, 8, 8, 99]
+
+
+def test_mistral_v7_tekken_train_on_eos_all_respects_user_include_end_false():
+    """System end (include_end=True) unmasked on 'all'; [/INST] stays masked."""
+    vocab = {
+        "[SYSTEM_PROMPT]": [40],
+        "[/SYSTEM_PROMPT]": [41],
+        "[INST]": [50],
+        "[/INST]": [51],
+    }
+    tok = _Tokenizer(vocab, pad_id=0, eos_id=99)
+    strategy = MistralV7TekkenProcessingStrategy(_Processor(tok), train_on_eos="all")
+    seq = [40, 5, 41, 50, 7, 51, 8, 99]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    # system content masked, [/SYSTEM_PROMPT]=41 kept (include_end=True + all);
+    # user + [/INST]=51 masked (include_end=False); assistant 8 + eos 99 kept.
+    assert out == [-100, -100, 41, -100, -100, -100, 8, 99]
+
+
+# --------------------------------------------------------------------------- #
+# Voxtral / Mistral3 / InternVL / Glm4v
+#
+# These four strategies do NOT declare role boundaries (mistral-common or
+# template-variant uncertainty); they fall back to pad + media masking.
+# Coverage focuses on the media-token masking surface plus role-boundary
+# override as the supported opt-in path.
+# --------------------------------------------------------------------------- #
+
+
+def _make_voxtral_strategy(**kwargs):
+    """Voxtral stub: special_ids live at processor.tokenizer.tokenizer.instruct_tokenizer.audio_encoder."""
+    tok = _Tokenizer({}, pad_id=0)
+    audio_encoder = SimpleNamespace(
+        special_ids=SimpleNamespace(audio=300, begin_audio=301)
+    )
+    tok.tokenizer = SimpleNamespace(
+        instruct_tokenizer=SimpleNamespace(audio_encoder=audio_encoder)
+    )
+    return VoxtralProcessingStrategy(_Processor(tok), **kwargs)
+
+
+def test_voxtral_masks_pad_and_audio_tokens():
+    """Default Voxtral: no role boundaries → keep content, mask pad + audio markers."""
+    strategy = _make_voxtral_strategy()
+    # 300 = audio, 301 = begin_audio, 0 = pad. Other ids are kept.
+    seq = [7, 0, 300, 8, 301, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [7, -100, -100, 8, -100, 9]
+
+
+def test_voxtral_train_on_inputs_true_still_masks_audio():
+    """Even with train_on_inputs, audio/pad tokens must be masked."""
+    strategy = _make_voxtral_strategy(train_on_inputs=True)
+    seq = [7, 0, 300, 8, 301, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [7, -100, -100, 8, -100, 9]
+
+
+def test_voxtral_role_boundaries_override_enables_role_masking():
+    """Override opt-in: role masking applies on top of pad/audio masking."""
+    tok = _Tokenizer({"BOA": [50], "EOT": [60]}, pad_id=0)
+    audio_encoder = SimpleNamespace(
+        special_ids=SimpleNamespace(audio=300, begin_audio=301)
+    )
+    tok.tokenizer = SimpleNamespace(
+        instruct_tokenizer=SimpleNamespace(audio_encoder=audio_encoder)
+    )
+    strategy = VoxtralProcessingStrategy(
+        _Processor(tok),
+        role_boundaries_override=[{"role": "assistant", "start": "BOA", "end": "EOT"}],
+    )
+    # Pre-BOA + post-EOT masked; BOA(50) and EOT(60) included by default;
+    # audio token 300 inside the span still masked.
+    seq = [7, 50, 8, 300, 9, 60, 1]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, 8, -100, 9, 60, -100]
+
+
+def _make_mistral3_strategy(**kwargs):
+    tok = _Tokenizer({}, pad_id=0)
+    image_encoder = SimpleNamespace(
+        special_ids=SimpleNamespace(img=400, img_break=401, img_end=402)
+    )
+    tok.tokenizer = SimpleNamespace(
+        instruct_tokenizer=SimpleNamespace(image_encoder=image_encoder)
+    )
+    return Mistral3ProcessingStrategy(_Processor(tok), **kwargs)
+
+
+def test_mistral3_masks_pad_and_three_image_tokens():
+    strategy = _make_mistral3_strategy()
+    seq = [7, 0, 400, 8, 401, 9, 402, 10]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [7, -100, -100, 8, -100, 9, -100, 10]
+
+
+def test_mistral3_train_on_inputs_true_still_masks_image_markers():
+    strategy = _make_mistral3_strategy(train_on_inputs=True)
+    seq = [7, 0, 400, 8, 401, 9, 402, 10]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [7, -100, -100, 8, -100, 9, -100, 10]
+
+
+def test_internvl_masks_pad_and_image_id_list():
+    """InternVL stores a *list* of image ids on processor.image_ids."""
+    tok = _Tokenizer({}, pad_id=0)
+    proc = _Processor(tok)
+    proc.image_ids = [500, 501, 502]
+    strategy = InternVLProcessingStrategy(proc)
+    seq = [7, 0, 500, 8, 501, 9, 502, 10]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [7, -100, -100, 8, -100, 9, -100, 10]
+
+
+def test_internvl_raises_when_image_ids_missing():
+    """No image_ids on the processor → __init__ must surface a clear error."""
+    tok = _Tokenizer({}, pad_id=0)
+    with pytest.raises(ValueError, match="image_ids"):
+        InternVLProcessingStrategy(_Processor(tok))
+
+
+def test_internvl_image_ids_with_none_entries_skipped():
+    """None entries in image_ids must not be used as a mask value."""
+    tok = _Tokenizer({}, pad_id=0)
+    proc = _Processor(tok)
+    proc.image_ids = [500, None, 502]
+    strategy = InternVLProcessingStrategy(proc)
+    seq = [500, 7, 502, 8]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, 7, -100, 8]
+
+
+def _glm4v_tokenizer():
+    vocab = {
+        "<|image|>": [600],
+        "<|begin_of_image|>": [601],
+        "<|end_of_image|>": [602],
+        "<|video|>": [610],
+        "<|begin_of_video|>": [611],
+        "<|end_of_video|>": [612],
+    }
+    return _Tokenizer(vocab, pad_id=0)
+
+
+def test_glm4v_masks_pad_image_and_video_markers():
+    strategy = Glm4vProcessingStrategy(_Processor(_glm4v_tokenizer()))
+    # All six markers + pad must be masked; surrounding text tokens kept.
+    seq = [7, 0, 600, 601, 602, 8, 610, 611, 612, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [7, -100, -100, -100, -100, 8, -100, -100, -100, 9]
+
+
+def test_glm4v_train_on_inputs_true_still_masks_media_markers():
+    strategy = Glm4vProcessingStrategy(
+        _Processor(_glm4v_tokenizer()), train_on_inputs=True
+    )
+    seq = [7, 0, 600, 8, 610, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [7, -100, -100, 8, -100, 9]
+
+
+def test_glm4v_role_boundaries_override_enables_role_masking():
+    """Override opt-in masks non-assistant role spans on top of media masking."""
+    tok = _glm4v_tokenizer()
+    tok.vocab["BOA"] = [50]
+    tok.vocab["EOT"] = [60]
+    strategy = Glm4vProcessingStrategy(
+        _Processor(tok),
+        role_boundaries_override=[{"role": "assistant", "start": "BOA", "end": "EOT"}],
+    )
+    seq = [7, 50, 8, 600, 9, 60, 1]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    # Pre-BOA + post-EOT masked; image marker 600 inside the assistant span
+    # is still masked post-scan.
+    assert out == [-100, -100, 8, -100, 9, 60, -100]
+
+
+def test_glm4v_batch_of_two_rows():
+    """Multiple rows independently masked; pad in trailing positions kept as -100."""
+    strategy = Glm4vProcessingStrategy(_Processor(_glm4v_tokenizer()))
+    row_a = [7, 600, 8, 610, 9]
+    row_b = [11, 601, 12, 0, 0]
+    out = strategy.process_labels(torch.tensor([row_a, row_b])).tolist()
+    assert out[0] == [7, -100, 8, -100, 9]
+    assert out[1] == [11, -100, 12, -100, -100]
+
+
+# --------------------------------------------------------------------------- #
+# Dispatcher routing
+# --------------------------------------------------------------------------- #
+
+
+def _dispatch(processor, chat_template_type):
+    return get_processing_strategy(
+        processor=processor,
+        chat_template=None,
+        chat_template_type=chat_template_type,
+    )
+
+
+def test_dispatch_qwen2_vl():
+    s = _dispatch(_Processor(_qwen_tokenizer()), "qwen2_vl")
+    assert isinstance(s, Qwen2VLProcessingStrategy)
+
+
+def test_dispatch_qwen3_5():
+    s = _dispatch(_Processor(_qwen_tokenizer()), "qwen3_5")
+    assert isinstance(s, Qwen3_5ProcessingStrategy)
+
+
+def test_dispatch_gemma3():
+    s = _dispatch(_Processor(_gemma_tokenizer()), "gemma3")
+    assert isinstance(s, Gemma3ProcessingStrategy)
+
+
+def test_dispatch_gemma3n():
+    s = _dispatch(_Processor(_gemma_tokenizer()), "gemma3n")
+    assert isinstance(s, Gemma3nProcessingStrategy)
+
+
+def test_dispatch_gemma4():
+    s = _dispatch(_FakeGemma4Processor(), "gemma4")
+    assert isinstance(s, Gemma4ProcessingStrategy)
+
+
+def test_dispatch_gemma4_unified():
+    s = _dispatch(_FakeGemma4Processor(), "gemma4_unified")
+    assert isinstance(s, Gemma4UnifiedProcessingStrategy)
+
+
+def test_dispatch_llama3_2_vision():
+    vocab = {
+        "<|start_header_id|>assistant<|end_header_id|>\n\n": [1, 2, 3, 4, 5],
+        "<|eot_id|>": [10],
+    }
+    s = _dispatch(_Processor(_Tokenizer(vocab, pad_id=0)), "llama3_2_vision")
+    assert isinstance(s, Llama3_2VisionProcessingStrategy)
+
+
+def test_dispatch_llama4():
+    vocab = {
+        "<|header_start|>assistant<|header_end|>\n\n": [20, 21, 22, 23],
+        "<|eot|>": [30],
+    }
+    s = _dispatch(_Processor(_Tokenizer(vocab, pad_id=0)), "llama4")
+    assert isinstance(s, Llama4ProcessingStrategy)
+
+
+def test_dispatch_pixtral():
+    vocab = {"[INST]": [50], "[/INST]": [51]}
+    s = _dispatch(_Processor(_Tokenizer(vocab, pad_id=0, eos_id=99)), "pixtral")
+    assert isinstance(s, PixtralProcessingStrategy)
+
+
+def test_dispatch_mistral_v7_tekken():
+    vocab = {
+        "[INST]": [50],
+        "[/INST]": [51],
+        "[SYSTEM_PROMPT]": [40],
+        "[/SYSTEM_PROMPT]": [41],
+    }
+    s = _dispatch(
+        _Processor(_Tokenizer(vocab, pad_id=0, eos_id=99)), "mistral_v7_tekken"
+    )
+    assert isinstance(s, MistralV7TekkenProcessingStrategy)
+
+
+def test_dispatch_unknown_falls_back_to_base():
+    vocab = {"dummy": [1]}
+    s = _dispatch(_Processor(_Tokenizer(vocab, pad_id=0)), "llava")
+    assert type(s) is ProcessingStrategy
+
+
+def _glm_vision_processor(cls_path):
+    """Spec'd MagicMock so isinstance(mock, cls) passes without real HF files."""
+    from importlib import import_module
+    from unittest.mock import MagicMock
+
+    mod_name, cls_name = cls_path.rsplit(".", 1)
+    cls = getattr(import_module(mod_name), cls_name)
+
+    vocab = {
+        "<|image|>": [200],
+        "<|begin_of_image|>": [201],
+        "<|end_of_image|>": [202],
+        "<|video|>": [210],
+        "<|begin_of_video|>": [211],
+        "<|end_of_video|>": [212],
+    }
+    tok = _Tokenizer(vocab, pad_id=0)
+    proc = MagicMock(spec=cls)
+    proc.tokenizer = tok
+    # Drop processor.image_token so base class skips its probe.
+    del proc.image_token
+    return proc
+
+
+def test_dispatch_glm4v_via_Glm4vProcessor():
+    """Glm4vProcessor (GLM-4V) routes to Glm4vProcessingStrategy."""
+    pytest.importorskip("transformers.models.glm4v.processing_glm4v")
+    from axolotl.processing_strategies import Glm4vProcessingStrategy
+
+    proc = _glm_vision_processor(
+        "transformers.models.glm4v.processing_glm4v.Glm4vProcessor"
+    )
+    s = _dispatch(proc, None)
+    assert isinstance(s, Glm4vProcessingStrategy)
+
+
+def test_dispatch_glm4v_via_Glm46VProcessor():
+    """Glm46VProcessor (GLM-4.6V) also routes to Glm4vProcessingStrategy."""
+    pytest.importorskip("transformers.models.glm46v.processing_glm46v")
+    from axolotl.processing_strategies import Glm4vProcessingStrategy
+
+    proc = _glm_vision_processor(
+        "transformers.models.glm46v.processing_glm46v.Glm46VProcessor"
+    )
+    s = _dispatch(proc, None)
+    assert isinstance(s, Glm4vProcessingStrategy)
+
+
+def test_dispatch_voxtral():
+    """VoxtralProcessor routes to VoxtralProcessingStrategy."""
+    pytest.importorskip("transformers.models.voxtral")
+    from transformers.models.voxtral import VoxtralProcessor
+
+    tok = _Tokenizer({}, pad_id=0)
+    audio_encoder = SimpleNamespace(
+        special_ids=SimpleNamespace(audio=300, begin_audio=301)
+    )
+    tok.tokenizer = SimpleNamespace(
+        instruct_tokenizer=SimpleNamespace(audio_encoder=audio_encoder)
+    )
+    proc = MagicMock(spec=VoxtralProcessor)
+    proc.tokenizer = tok
+    if hasattr(proc, "image_token"):
+        del proc.image_token
+
+    s = _dispatch(proc, None)
+    assert isinstance(s, VoxtralProcessingStrategy)
+
+
+def test_dispatch_mistral3():
+    """Mistral3Processor (axolotl's optional wrapper) routes to Mistral3ProcessingStrategy."""
+    pytest.importorskip("mistral_common")
+    Mistral3Processor = pytest.importorskip(
+        "axolotl.utils.mistral.mistral3_processor"
+    ).Mistral3Processor
+
+    tok = _Tokenizer({}, pad_id=0)
+    image_encoder = SimpleNamespace(
+        special_ids=SimpleNamespace(img=400, img_break=401, img_end=402)
+    )
+    tok.tokenizer = SimpleNamespace(
+        instruct_tokenizer=SimpleNamespace(image_encoder=image_encoder)
+    )
+    proc = MagicMock(spec=Mistral3Processor)
+    proc.tokenizer = tok
+    if hasattr(proc, "image_token"):
+        del proc.image_token
+
+    s = _dispatch(proc, None)
+    assert isinstance(s, Mistral3ProcessingStrategy)
+
+
+def test_dispatch_internvl():
+    """InternVLProcessor routes to InternVLProcessingStrategy."""
+    pytest.importorskip("transformers.models.internvl")
+    from transformers.models.internvl import InternVLProcessor
+
+    tok = _Tokenizer({}, pad_id=0)
+    proc = MagicMock(spec=InternVLProcessor)
+    proc.tokenizer = tok
+    proc.image_ids = [500, 501, 502]
+    if hasattr(proc, "image_token"):
+        del proc.image_token
+
+    s = _dispatch(proc, None)
+    assert isinstance(s, InternVLProcessingStrategy)
+
+
+def test_mistral3_role_boundaries_override_enables_role_masking():
+    """Mistral3 override opt-in: role masking applies on top of pad/image masking."""
+    tok = _Tokenizer({"BOA": [50], "EOT": [60]}, pad_id=0)
+    image_encoder = SimpleNamespace(
+        special_ids=SimpleNamespace(img=400, img_break=401, img_end=402)
+    )
+    tok.tokenizer = SimpleNamespace(
+        instruct_tokenizer=SimpleNamespace(image_encoder=image_encoder)
+    )
+    strategy = Mistral3ProcessingStrategy(
+        _Processor(tok),
+        role_boundaries_override=[{"role": "assistant", "start": "BOA", "end": "EOT"}],
+    )
+    seq = [7, 50, 8, 400, 9, 60, 1]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    # Pre-BOA + post-EOT masked; image marker 400 inside the span also masked.
+    assert out == [-100, -100, 8, -100, 9, 60, -100]
+
+
+# --------------------------------------------------------------------------- #
+# Config-based role-boundary override
+# --------------------------------------------------------------------------- #
+
+
+def test_role_boundaries_override_replaces_built_in():
+    """Override swaps the built-in boundaries wholesale, not additively."""
+    vocab = {
+        "<|im_start|>assistant\n": [101, 102, 103],
+        "<|im_start|>user\n": [101, 106, 103],
+        "<|im_end|>": [104],
+        ">>>A": [200, 201],
+        ">>>U": [200, 202],
+        "<<<": [210],
+        "<|image_pad|>": [250],
+    }
+    strategy = Qwen2VLProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[
+            {"role": "assistant", "start": ">>>A", "end": "<<<"},
+            {"role": "user", "start": ">>>U", "end": "<<<"},
+        ],
+    )
+    seq = [
+        101,
+        106,
+        103,
+        7,
+        104,
+        200,
+        201,
+        9,
+        9,
+        210,
+    ]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, -100, -100, -100, -100, -100, 9, 9, 210]
+
+
+def test_role_boundaries_override_enables_unverified_strategy():
+    """Override lets users opt in to role masking on strategies that default opt out."""
+    vocab = {
+        "BOA": [50, 51],
+        "EOT": [60],
+    }
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[
+            {"role": "assistant", "start": "BOA", "end": "EOT"},
+        ],
+    )
+    seq = [1, 2, 3, 50, 51, 7, 8, 60, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, -100, -100, -100, 7, 8, 60, -100]
+
+
+def test_role_boundaries_override_eos_token_sentinel():
+    vocab = {"BOA": [50]}
+    tok = _Tokenizer(vocab, pad_id=0, eos_id=99)
+    strategy = ProcessingStrategy(
+        _Processor(tok),
+        role_boundaries_override=[
+            {"role": "assistant", "start": "BOA", "end": "eos_token"},
+        ],
+    )
+    seq = [1, 50, 7, 7, 99, 2]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, 7, 7, 99, -100]
+
+
+def test_role_boundaries_override_end_null_runs_to_sequence_end():
+    vocab = {"BOA": [50]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[
+            {"role": "assistant", "start": "BOA", "end": None},
+        ],
+    )
+    seq = [1, 2, 50, 7, 8, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, -100, 7, 8, 9]
+
+
+def test_role_boundaries_override_rejects_bad_spec():
+    vocab = {"BOA": [50]}
+    with pytest.raises(ValueError, match="must have both 'role' and 'start'"):
+        ProcessingStrategy(
+            _Processor(_Tokenizer(vocab, pad_id=0)),
+            role_boundaries_override=[{"role": "assistant"}],
+        )
+
+
+def test_role_boundaries_override_rejects_unencodable_start():
+    vocab = {"BOA": [50]}
+    with pytest.raises(ValueError, match="tokenizes to an empty sequence"):
+        ProcessingStrategy(
+            _Processor(_Tokenizer(vocab, pad_id=0)),
+            role_boundaries_override=[
+                {"role": "assistant", "start": "MISSING", "end": None}
+            ],
+        )
+
+
+def test_role_boundaries_override_rejects_unencodable_end():
+    vocab = {"BOA": [50]}
+    with pytest.raises(ValueError, match="tokenizes to an empty sequence"):
+        ProcessingStrategy(
+            _Processor(_Tokenizer(vocab, pad_id=0)),
+            role_boundaries_override=[
+                {"role": "assistant", "start": "BOA", "end": "MISSING"}
+            ],
+        )
+
+
+def test_role_boundaries_override_accepts_pydantic_models():
+    # cfg.role_boundaries arrives as RoleBoundarySpec after pydantic parsing.
+    from axolotl.utils.schemas.multimodal import RoleBoundarySpec
+
+    vocab = {"BOA": [50], "EOT": [60]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[
+            RoleBoundarySpec(role="assistant", start="BOA", end="EOT")
+        ],
+    )
+    assert len(strategy.role_boundaries) == 1
+    assert strategy.role_boundaries[0].role == "assistant"
+    assert strategy.role_boundaries[0].start_tokens == [50]
+    assert strategy.role_boundaries[0].end_tokens == [60]
+
+
+def test_base_strategy_warns_when_no_boundaries(axolotl_caplog):
+    """No boundaries + train_on_inputs=False: one-shot warning, labels unchanged."""
+    import axolotl.processing_strategies as mod
+
+    mod._ROLE_MASK_WARNED.discard("ProcessingStrategy")
+
+    vocab = {"dummy": [1]}
+    s = ProcessingStrategy(_Processor(_Tokenizer(vocab, pad_id=0)))
+
+    with axolotl_caplog.at_level(
+        logging.WARNING, logger="axolotl.processing_strategies"
+    ):
+        labels = s.process_labels(torch.tensor([[1, 2, 3]]))
+    assert labels.tolist() == [[1, 2, 3]]
+    assert any("role boundaries" in rec.message for rec in axolotl_caplog.records)
+
+
+# --------------------------------------------------------------------------- #
+# Additional edge-case coverage
+# --------------------------------------------------------------------------- #
+
+
+def test_scanner_batch_size_greater_than_one():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+        RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]),
+    ]
+    labels = torch.tensor(
+        [
+            [1, 3, 7, 9, 1, 2, 8, 9],
+            [1, 2, 5, 5, 9, 0, 0, 0],
+        ]
+    )
+    out = _apply_role_boundaries(labels, boundaries, {"assistant"}, "turn").tolist()
+    assert out[0] == [-100, -100, -100, -100, -100, -100, 8, 9]
+    assert out[1] == [-100, -100, 5, 5, 9, -100, -100, -100]
+
+
+def test_scanner_adjacent_trainable_turns():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+    ]
+    seq = [1, 2, 5, 9, 1, 2, 6, 9]
+    out = _scan(boundaries, seq)
+    assert out == [-100, -100, 5, 9, -100, -100, 6, 9]
+
+
+def test_scanner_train_on_eos_none_multi_turn():
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+        RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]),
+    ]
+    seq = [1, 3, 7, 9, 1, 2, 8, 9, 1, 3, 7, 9, 1, 2, 6, 9]
+    out = _scan(boundaries, seq, train_on_eos="none")
+    assert out == [
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        8,
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        -100,
+        6,
+        -100,
+    ]
+
+
+def test_scanner_train_on_eos_all_with_user_turn_no_end_marker():
+    """Unclosed non-trainable span with train_on_eos='all': nothing included, no crash."""
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]),
+        RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]),
+    ]
+    seq = [1, 3, 7, 7, 7]
+    out = _scan(boundaries, seq, train_on_eos="all")
+    assert out == [-100, -100, -100, -100, -100]
+
+
+def test_scanner_include_start_true_via_override():
+    vocab = {"BOA": [50, 51], "EOT": [60]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[
+            {
+                "role": "assistant",
+                "start": "BOA",
+                "end": "EOT",
+                "include_start": True,
+            },
+        ],
+    )
+    seq = [1, 50, 51, 7, 8, 60, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, 50, 51, 7, 8, 60, -100]
+
+
+def test_scanner_include_end_false_via_override():
+    """include_end=False drops end marker even with train_on_eos='turn'."""
+    vocab = {"BOA": [50], "EOT": [60]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[
+            {
+                "role": "assistant",
+                "start": "BOA",
+                "end": "EOT",
+                "include_end": False,
+            },
+        ],
+    )
+    seq = [1, 50, 7, 8, 60, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, 7, 8, -100, -100]
+
+
+def test_scanner_empty_start_tokens_is_defensive_noop():
+    """Defensive: empty start_tokens matches nothing; everything masked."""
+    boundaries = [
+        RoleBoundary(role="assistant", start_tokens=[], end_tokens=[9]),
+    ]
+    seq = [1, 2, 3, 4, 9]
+    out = _scan(boundaries, seq)
+    assert out == [-100] * 5
+
+
+def test_process_labels_masks_pad_inside_assistant_span():
+    """Pad inside a trainable span is still masked post-scan."""
+    strategy = _make_qwen2vl()
+    seq = [101, 102, 103, 8, 0, 8, 104]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, -100, 8, -100, 8, 104]
+
+
+def test_process_labels_all_pad_sequence_does_not_crash():
+    strategy = _make_qwen2vl()
+    seq = [0, 0, 0, 0]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, -100, -100]
+
+
+def test_qwen2vl_multiple_consecutive_assistant_turns():
+    strategy = _make_qwen2vl()
+    seq = [101, 102, 103, 8, 104, 101, 102, 103, 9, 104]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [
+        -100,
+        -100,
+        -100,
+        8,
+        104,
+        -100,
+        -100,
+        -100,
+        9,
+        104,
+    ]
+
+
+def test_qwen2vl_batch_of_two_rows():
+    strategy = _make_qwen2vl()
+    row_a = [101, 106, 103, 7, 104, 101, 102, 103, 8, 104]
+    row_b = [101, 102, 103, 9, 104, 0, 0, 0, 0, 0]
+    out = strategy.process_labels(torch.tensor([row_a, row_b])).tolist()
+    assert out[0] == [-100, -100, -100, -100, -100, -100, -100, -100, 8, 104]
+    assert out[1] == [-100, -100, -100, 9, 104, -100, -100, -100, -100, -100]
+
+
+def test_qwen3_5_train_on_inputs_true_still_masks_video_pad():
+    """train_on_inputs=True skips role masking but media tokens are still masked."""
+    tok = _qwen_tokenizer()
+    strategy = Qwen3_5ProcessingStrategy(_Processor(tok), train_on_inputs=True)
+    seq = [101, 106, 103, 201, 7, 104, 101, 102, 103, 201, 8, 104]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    expected = list(seq)
+    expected[3] = -100
+    expected[9] = -100
+    assert out == expected
+
+
+def test_role_boundaries_override_role_not_in_roles_to_train():
+    """Override covering only a non-trainable role masks everything."""
+    vocab = {"BOU": [50], "EOT": [60]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[
+            {"role": "user", "start": "BOU", "end": "EOT"},
+        ],
+    )
+    seq = [1, 50, 7, 8, 60, 9]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100] * 6
+
+
+def test_role_boundaries_override_include_start_flag_round_trips():
+    from axolotl.utils.schemas.multimodal import RoleBoundarySpec
+
+    vocab = {"BOA": [50], "EOT": [60]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=[
+            RoleBoundarySpec(
+                role="assistant", start="BOA", end="EOT", include_start=True
+            ),
+        ],
+    )
+    assert len(strategy.role_boundaries) == 1
+    assert strategy.role_boundaries[0].include_start is True
+    assert strategy.role_boundaries[0].include_end is True
+
+
+def test_multimodal_config_parses_dict_role_boundaries_to_specs():
+    from axolotl.utils.schemas.multimodal import (
+        MultiModalConfig,
+        RoleBoundarySpec,
+    )
+
+    cfg = MultiModalConfig(
+        role_boundaries=[
+            {"role": "assistant", "start": "BOA", "end": "EOT"},
+            {"role": "user", "start": "BOU", "end": "EOT"},
+        ]
+    )
+    assert cfg.role_boundaries is not None
+    assert len(cfg.role_boundaries) == 2
+    assert all(isinstance(rb, RoleBoundarySpec) for rb in cfg.role_boundaries)
+
+    vocab = {"BOA": [50], "BOU": [51], "EOT": [60]}
+    strategy = ProcessingStrategy(
+        _Processor(_Tokenizer(vocab, pad_id=0)),
+        role_boundaries_override=cfg.role_boundaries,
+    )
+    seq = [51, 7, 60, 50, 8, 60]
+    out = strategy.process_labels(torch.tensor([seq])).tolist()[0]
+    assert out == [-100, -100, -100, -100, 8, 60]
+
+
+# --------------------------------------------------------------------------- #
+# field_messages plumbing: helpers and __call__ re-routing
+# --------------------------------------------------------------------------- #
+
+
+def _mock_processor() -> MagicMock:
+    """Processor mock with no ``image_token`` attr so __init__ skips that branch."""
+    processor = MagicMock()
+    del processor.image_token
+    return processor
+
+
+def _openai_messages() -> list[dict]:
+    return [
+        {"role": "user", "content": "hello"},
+        {"role": "assistant", "content": "hi"},
+    ]
+
+
+def _sharegpt_messages() -> list[dict]:
+    return [
+        {"from": "human", "value": "hello"},
+        {"from": "gpt", "value": "hi"},
+    ]
+
+
+def test_normalize_field_messages_none_defaults_to_messages():
+    assert ProcessingStrategy._normalize_field_messages(None) == ("messages",)
+
+
+def test_normalize_field_messages_string_wrapped_in_tuple():
+    assert ProcessingStrategy._normalize_field_messages("dialogue") == ("dialogue",)
+
+
+def test_normalize_field_messages_list_preserved():
+    assert ProcessingStrategy._normalize_field_messages(["foo", "bar"]) == (
+        "foo",
+        "bar",
+    )
+
+
+def test_normalize_field_messages_falsy_entries_dropped():
+    assert ProcessingStrategy._normalize_field_messages(["foo", "", None, "bar"]) == (
+        "foo",
+        "bar",
+    )
+
+
+def test_is_legacy_schema_sharegpt_pair_detected():
+    assert ProcessingStrategy._is_legacy_schema([{"from": "human", "value": "hi"}])
+
+
+def test_is_legacy_schema_only_from_is_not_enough():
+    # `from` without `value` (e.g., custom metadata) must not be misrouted.
+    assert not ProcessingStrategy._is_legacy_schema([{"from": "human", "extra": "x"}])
+
+
+def test_is_legacy_schema_openai_schema_returns_false():
+    assert not ProcessingStrategy._is_legacy_schema([{"role": "user", "content": "hi"}])
+
+
+def test_is_legacy_schema_empty_list_returns_false():
+    assert not ProcessingStrategy._is_legacy_schema([])
+
+
+def test_is_legacy_schema_non_list_returns_false():
+    assert not ProcessingStrategy._is_legacy_schema(None)
+    assert not ProcessingStrategy._is_legacy_schema("not a list")
+
+
+def test_get_messages_field_default_returns_messages():
+    strategy = ProcessingStrategy(processor=_mock_processor())
+    assert strategy._get_messages_field({"messages": _openai_messages()}) == "messages"
+
+
+def test_get_messages_field_custom_field_returns_custom():
+    strategy = ProcessingStrategy(processor=_mock_processor(), field_messages="my_col")
+    assert strategy._get_messages_field({"my_col": _openai_messages()}) == "my_col"
+
+
+def test_get_messages_field_custom_takes_priority_over_stale_messages():
+    # Configured custom field must beat a leftover `messages` column.
+    strategy = ProcessingStrategy(processor=_mock_processor(), field_messages="my_col")
+    example = {
+        "messages": [{"role": "user", "content": "stale"}],
+        "my_col": _openai_messages(),
+    }
+    assert strategy._get_messages_field(example) == "my_col"
+
+
+def test_get_messages_field_falls_back_to_messages_when_custom_absent():
+    strategy = ProcessingStrategy(processor=_mock_processor(), field_messages="my_col")
+    assert strategy._get_messages_field({"messages": _openai_messages()}) == "messages"
+
+
+def test_get_messages_field_returns_none_when_no_match():
+    strategy = ProcessingStrategy(processor=_mock_processor(), field_messages="my_col")
+    assert strategy._get_messages_field({"random_col": "x"}) is None
+
+
+def test_call_canonical_messages_openai():
+    strategy = ProcessingStrategy(processor=_mock_processor())
+    out = strategy([{"messages": _openai_messages()}])
+
+    assert len(out) == 1
+    assert out[0]["messages"][0]["role"] == "user"
+    # String content is wrapped to multimedia content list.
+    assert out[0]["messages"][0]["content"] == [{"type": "text", "text": "hello"}]
+
+
+def test_call_canonical_conversations_sharegpt():
+    strategy = ProcessingStrategy(processor=_mock_processor())
+    out = strategy([{"conversations": _sharegpt_messages()}])
+
+    # Roles get normalized: human -> user, gpt -> assistant.
+    assert out[0]["messages"][0]["role"] == "user"
+    assert out[0]["messages"][1]["role"] == "assistant"
+    assert "conversations" not in out[0]
+
+
+def test_call_custom_field_with_openai_schema():
+    strategy = ProcessingStrategy(processor=_mock_processor(), field_messages="my_col")
+    out = strategy([{"my_col": _openai_messages()}])
+
+    assert out[0]["messages"][0]["role"] == "user"
+    assert "my_col" not in out[0]
+
+
+def test_call_custom_field_with_sharegpt_schema():
+    # A custom column whose first message looks like ShareGPT is routed to the legacy branch.
+    strategy = ProcessingStrategy(processor=_mock_processor(), field_messages="my_col")
+    out = strategy([{"my_col": _sharegpt_messages()}])
+
+    assert out[0]["messages"][0]["role"] == "user"
+    assert out[0]["messages"][1]["role"] == "assistant"
+    assert "my_col" not in out[0]
+    assert "conversations" not in out[0]
+
+
+def test_call_custom_field_overrides_stale_messages_column():
+    strategy = ProcessingStrategy(processor=_mock_processor(), field_messages="my_col")
+    out = strategy(
+        [
+            {
+                "messages": [{"role": "user", "content": "stale"}],
+                "my_col": [{"role": "user", "content": "fresh"}],
+            }
+        ]
+    )
+
+    assert out[0]["messages"][0]["content"] == [{"type": "text", "text": "fresh"}]
+
+
+def test_call_unknown_key_raises_value_error():
+    strategy = ProcessingStrategy(processor=_mock_processor(), field_messages="my_col")
+
+    with pytest.raises(ValueError):
+        strategy([{"random_col": _openai_messages()}])
diff --git a/tests/test_prompt_tokenizers.py b/tests/test_prompt_tokenizers.py
index 63e9a621bf..181fddbc08 100644
--- a/tests/test_prompt_tokenizers.py
+++ b/tests/test_prompt_tokenizers.py
@@ -1,15 +1,7 @@
 """Module for testing prompt tokenizers."""
 
 import json
-import logging
-import unittest
-from copy import deepcopy
 from pathlib import Path
-from typing import Optional
-
-import pytest
-from datasets import load_dataset
-from transformers import AddedToken, AutoTokenizer, LlamaTokenizer
 
 from axolotl.prompt_strategies.alpaca_chat import NoSystemPrompter
 from axolotl.prompt_strategies.alpaca_w_system import (
@@ -21,15 +13,11 @@
     LLama2ChatTokenizingStrategy,
 )
 from axolotl.prompt_strategies.orpo.chat_template import load
-from axolotl.prompt_strategies.sharegpt import GlaiveShareGPTPromptTokenizingStrategy
-from axolotl.prompt_tokenizers import (
-    AlpacaPromptTokenizingStrategy,
-    ShareGPTPromptTokenizingStrategy,
-)
-from axolotl.prompters import AlpacaPrompter, PromptStyle, ShareGPTPrompterV2
+from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy
+from axolotl.prompters import AlpacaPrompter, PromptStyle
 from axolotl.utils.dict import DictDefault
 
-LOG = logging.getLogger("axolotl")
+from tests.hf_offline_utils import enable_hf_offline
 
 test_data = {
     "multi_turn_sys": {
@@ -65,238 +53,21 @@
 }
 
 
-def prompt_strat(conversation, tokenizer):
-    "Helper function to create a prompt strategy for testing."
-    prompter = ShareGPTPrompterV2(conversation=conversation)
-    return ShareGPTPromptTokenizingStrategy(
-        prompter,
-        tokenizer,
-        False,
-        2048,
-    )
-
-
-class TestPromptTokenizationStrategies(unittest.TestCase):
+class TestPromptTokenizationStrategies:
     """
     Test class for prompt tokenization strategies.
     """
 
-    _caplog: Optional[pytest.LogCaptureFixture] = None
-
-    @pytest.fixture(autouse=True)
-    def inject_fixtures(self, caplog):
-        self._caplog = caplog
-
-    def setUp(self) -> None:
-        # pylint: disable=duplicate-code
-        self.tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
-        self.tokenizer.add_special_tokens(
-            {
-                "bos_token": "<s>",
-                "eos_token": "</s>",
-                "unk_token": "<unk>",
-            }
-        )
-
-    def test_sharegpt_integration(self):
-        with open(
-            Path(__file__).parent / "fixtures/conversation.json", encoding="utf-8"
-        ) as fin:
-            data = fin.read()
-            conversation = json.loads(data)
-        with open(
-            Path(__file__).parent / "fixtures/conversation.tokenized.json",
-            encoding="utf-8",
-        ) as fin:
-            data = fin.read()
-            tokenized_conversation = json.loads(data)
-        prompter = ShareGPTPrompterV2()
-        strat = ShareGPTPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        example = strat.tokenize_prompt(conversation)
-        for fields in ["input_ids", "attention_mask", "labels"]:
-            self.assertEqual(len(example[fields]), len(tokenized_conversation[fields]))
-            self.assertEqual(example[fields], tokenized_conversation[fields])
-
-    def test_sharegpt_warnings_integration(self):
-        with open(
-            Path(__file__).parent / "fixtures/conversation.missingturns.json",
-            encoding="utf-8",
-        ) as fin:
-            data = fin.read()
-            conversation = json.loads(data)
-        prompter = ShareGPTPrompterV2()
-        strat = ShareGPTPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        with self._caplog.at_level(logging.WARNING):
-            strat.tokenize_prompt(conversation)
-            assert "assistant turn has empty text" in self._caplog.records[1].message
-
-    def test_sharegpt_warnings_turns(self):
-        conversation = {
-            "conversations": [
-                {"from": "system", "value": "lorem"},
-                {"from": "gpt", "value": "ipsum"},
-                {"from": "human", "value": "dolor"},
-                {"from": "human", "value": "dolor"},
-                {"from": "gpt", "value": "sit"},
-            ]
-        }
-        prompter = ShareGPTPrompterV2()
-        strat = ShareGPTPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        with self._caplog.at_level(logging.WARNING):
-            strat.tokenize_prompt(conversation)
-            assert (
-                "Role did not alternate between turns (gpt and human)"
-                in self._caplog.records[0].message
-            )
-
-    def test_sharegpt_llama(self):
-        "Make sure the sharegpt/llama is tokenized and formatted correctly."
-        strat = prompt_strat("llama-2", self.tokenizer)
-
-        def tokenize(conv):
-            return strat.tokenize_prompt(deepcopy(conv))["input_ids"]
-
-        def decode(ids):
-            return strat.tokenizer.decode(ids)
-
-        # fmt: off
-        # System message, multi-turn conversations
-        mt_ids = tokenize(test_data['multi_turn_sys'])
-        assert decode(mt_ids) == '<s> [INST] <<SYS>>\nlorem\n<</SYS>>\n\nabc [/INST] ipsum</s><s> [INST] 123 [/INST] sit</s>'
-        assert mt_ids == [1, 518, 25580, 29962, 3532, 14816, 29903, 6778, 13, 29880, 3668, 13, 29966, 829, 14816, 29903, 6778, 13, 13, 10736, 518, 29914, 25580, 29962, 23421, 2, 1, 518, 25580, 29962, 29871, 29896, 29906, 29941, 518, 29914, 25580, 29962, 7845, 2]
-
-        # System message, single-turn conversations
-        st_ids = tokenize(test_data['single_turn_sys'])
-        assert decode(st_ids) == '<s> [INST] <<SYS>>\nlorem\n<</SYS>>\n\nabc [/INST] ipsum</s>'
-        assert st_ids == [1, 518, 25580, 29962, 3532, 14816, 29903, 6778, 13, 29880, 3668, 13, 29966, 829, 14816, 29903, 6778, 13, 13, 10736, 518, 29914, 25580, 29962, 23421, 2]
-
-        # No system message, single-turn
-        ns_ids = tokenize(test_data['single_turn_no_sys'])
-        assert decode(ns_ids) == '<s> [INST] abc [/INST] ipsum</s>'
-        assert ns_ids == [1, 518, 25580, 29962, 25638, 518, 29914, 25580, 29962, 23421, 2]
-
-        # No system message, multi-turn
-        ns_mt_ids = tokenize(test_data['multi_turn_no_sys'])
-        assert decode(ns_mt_ids) == '<s> [INST] abc [/INST] ipsum</s><s> [INST] 123 [/INST] sit</s>'
-        assert ns_mt_ids == [1, 518, 25580, 29962, 25638, 518, 29914, 25580, 29962, 23421, 2, 1, 518, 25580, 29962, 29871, 29896, 29906, 29941, 518, 29914, 25580, 29962, 7845, 2]
-        # fmt: on
-
-    def test_sharegpt_mistral(self):
-        "Make sure the sharegpt/mistral is tokenized and formatted correctly."
-        strat = prompt_strat("mistral", self.tokenizer)
-
-        def tokenize(conv):
-            return strat.tokenize_prompt(deepcopy(conv))["input_ids"]
-
-        def decode(ids):
-            return strat.tokenizer.decode(ids)
-
-        # fmt: off
-        # System message, multi-turn conversations
-        mt_ids = tokenize(test_data['multi_turn_sys'])
-        assert decode(mt_ids) == '<s> [INST]  lorem\nabc [/INST] ipsum</s> [INST] 123 [/INST] sit</s>'
-        assert mt_ids == [1, 518, 25580, 29962, 29871, 301, 3668, 13, 10736, 518, 29914, 25580, 29962, 23421, 2, 518, 25580, 29962, 29871, 29896, 29906, 29941, 518, 29914, 25580, 29962, 7845, 2]
-
-        # System message, single-turn conversations
-        st_ids = tokenize(test_data['single_turn_sys'])
-        assert decode(st_ids) == '<s> [INST]  lorem\nabc [/INST] ipsum</s>'
-        assert st_ids == [1, 518, 25580, 29962, 29871, 301, 3668, 13, 10736, 518, 29914, 25580, 29962, 23421, 2]
-
-        # No system message, single-turn
-        ns_ids = tokenize(test_data['single_turn_no_sys'])
-        assert decode(ns_ids) == '<s> [INST] abc [/INST] ipsum</s>'
-        assert ns_ids == [1, 518, 25580, 29962, 25638, 518, 29914, 25580, 29962, 23421, 2]
-
-        # No system message, multi-turn
-        ns_mt_ids = tokenize(test_data['multi_turn_no_sys'])
-        assert decode(ns_mt_ids) == '<s> [INST] abc [/INST] ipsum</s> [INST] 123 [/INST] sit</s>'
-        assert ns_mt_ids == [1, 518, 25580, 29962, 25638, 518, 29914, 25580, 29962, 23421, 2, 518, 25580, 29962, 29871, 29896, 29906, 29941, 518, 29914, 25580, 29962, 7845, 2]
-        # fmt: on
-
-    def test_sharegpt_changes_roles(self):
-        conversation = {
-            "roles": ["USER", "CHARACTER"],
-            "conversations": [
-                {"from": "system", "value": "lorem"},
-                {"from": "gpt", "value": "ipsum"},
-                {"from": "human", "value": "dolor"},
-                {"from": "gpt", "value": "sit"},
-            ],
-        }
-        prompter = ShareGPTPrompterV2()
-        strat = ShareGPTPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        with self._caplog.at_level(logging.WARNING):
-            res = strat.tokenize_prompt(conversation)
-            assert "CHARACTER" in self.tokenizer.decode(res["input_ids"])
-
-    def test_sharegpt_assistant_label_ignore(self):
-        conversation = {
-            "roles": ["user", "assistant"],
-            "conversations": [
-                {"from": "system", "value": "lorem"},
-                {"from": "gpt", "value": "ipsum"},
-                {"from": "human", "value": "dolor"},
-                {"from": "gpt", "value": "sit"},
-            ],
-        }
-        prompter = ShareGPTPrompterV2()
-        strat = ShareGPTPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        with self._caplog.at_level(logging.WARNING):
-            res = strat.tokenize_prompt(conversation)
-            idx = res["input_ids"].index(20255)  # assistant token
-            assert res["labels"][idx] == -100
-
-    def test_glaive_tool_label_ignore(self):
-        conversation = {
-            "system": "SYSTEM: This is a system prompt",
-            "chat": "USER: Can you book a flight for me from New York to London? ASSISTANT: I'm sorry, but I don't have the capability to book flights.  <|endoftext|>",
-        }
-        prompter = ShareGPTPrompterV2()
-        strat = GlaiveShareGPTPromptTokenizingStrategy(
-            prompter,
-            self.tokenizer,
-            False,
-            2048,
-        )
-        with self._caplog.at_level(logging.WARNING):
-            res = strat.tokenize_prompt(conversation)
-            idx = res["input_ids"].index(13566)  # assistant token
-            assert res["labels"][idx] == -100
-
-    def test_no_sys_prompt(self):
+    @enable_hf_offline
+    def test_no_sys_prompt(self, tokenizer_huggyllama_w_special_tokens):
         """
         tests the interface between the user and assistant parts
         """
         prompter = NoSystemPrompter()
-        # pylint: disable=duplicate-code
+
         strat = AlpacaPromptTokenizingStrategy(
             prompter,
-            self.tokenizer,
+            tokenizer_huggyllama_w_special_tokens,
             False,
             2048,
         )
@@ -309,15 +80,16 @@ def test_no_sys_prompt(self):
         assert example["labels"][world_idx] == 3186
         assert example["labels"][world_idx - 1] == -100
 
-    def test_alpaca(self):
+    @enable_hf_offline
+    def test_alpaca(self, tokenizer_huggyllama_w_special_tokens):
         """
         tests the interface between the user and assistant parts
         """
-        # pylint: disable=duplicate-code
+
         prompter = AlpacaPrompter()
         strat = AlpacaPromptTokenizingStrategy(
             prompter,
-            self.tokenizer,
+            tokenizer_huggyllama_w_special_tokens,
             False,
             2048,
         )
@@ -328,27 +100,17 @@ def test_alpaca(self):
         assert example["labels"][world_idx - 1] == -100
 
 
-class InstructionWSystemPromptTokenizingStrategyTest(unittest.TestCase):
+class TestInstructionWSystemPromptTokenizingStrategy:
     """
     Test class for prompt tokenization strategies with sys prompt from the dataset
     """
 
-    def setUp(self) -> None:
-        # pylint: disable=duplicate-code
-        self.tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
-        self.tokenizer.add_special_tokens(
-            {
-                "bos_token": "<s>",
-                "eos_token": "</s>",
-                "unk_token": "<unk>",
-            }
-        )
-
-    def test_system_alpaca(self):
+    @enable_hf_offline
+    def test_system_alpaca(self, tokenizer_huggyllama_w_special_tokens):
         prompter = SystemDataPrompter(PromptStyle.CHAT.value)
         strat = InstructionWSystemPromptTokenizingStrategy(
             prompter,
-            self.tokenizer,
+            tokenizer_huggyllama_w_special_tokens,
             False,
             2048,
         )
@@ -369,17 +131,13 @@ def test_system_alpaca(self):
         assert example["input_ids"][8] == 11889  # USER
 
 
-class Llama2ChatTokenizationTest(unittest.TestCase):
+class Llama2ChatTokenizationTest:
     """
     Test class for prompt tokenization strategies with sys prompt from the dataset
     """
 
-    def setUp(self) -> None:
-        # pylint: disable=duplicate-code
-        self.tokenizer = LlamaTokenizer.from_pretrained("NousResearch/Llama-2-7b-hf")
-        # woraround because official Meta repos are not open
-
-    def test_llama2_chat_integration(self):
+    @enable_hf_offline
+    def test_llama2_chat_integration(self, tokenizer_llama2_7b):
         with open(
             Path(__file__).parent / "fixtures/conversation.json", encoding="utf-8"
         ) as fin:
@@ -394,16 +152,18 @@ def test_llama2_chat_integration(self):
         prompter = Llama2ChatPrompter()
         strat = LLama2ChatTokenizingStrategy(
             prompter,
-            self.tokenizer,
+            tokenizer_llama2_7b,
             False,
             4096,
         )
         example = strat.tokenize_prompt(conversation)
         for fields in ["input_ids", "attention_mask", "labels"]:
-            self.assertEqual(len(example[fields]), len(tokenized_conversation[fields]))
-            self.assertEqual(example[fields], tokenized_conversation[fields])
+            # pytest assert equals
+
+            assert len(example[fields]) == len(tokenized_conversation[fields])
+            assert example[fields] == tokenized_conversation[fields]
 
-    def compare_with_transformers_integration(self):
+    def compare_with_transformers_integration(self, tokenizer_llama2_7b):
         # this needs transformers >= v4.31.0
         from transformers.models.llama.tokenization_llama import B_SYS, E_SYS
         from transformers.pipelines.conversational import Conversation
@@ -411,7 +171,7 @@ def compare_with_transformers_integration(self):
         # from transformers.models.llama.tokenization_llama import DEFAULT_SYSTEM_PROMPT
         # broken as of 23/7/20
         # see https://github.com/huggingface/transformers/pull/24935
-        # pylint: disable=C0103
+
         DEFAULT_SYSTEM_PROMPT = """\
 You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
 
@@ -441,67 +201,43 @@ def compare_with_transformers_integration(self):
             + user_input[1:-1],
             generated_responses=answers,
         )
-        # pylint: disable=W0212
-        hf_tokens = self.tokenizer._build_conversation_input_ids(hf_conf)
 
-        self.assertEqual(
-            hf_tokens, tokenized_conversation["input_ids"][: len(hf_tokens)]
-        )
+        hf_tokens = tokenizer_llama2_7b._build_conversation_input_ids(hf_conf)
 
+        assert hf_tokens == tokenized_conversation["input_ids"][: len(hf_tokens)]
 
-class OrpoTokenizationTest(unittest.TestCase):
-    """test case for the ORPO tokenization"""
 
-    def setUp(self) -> None:
-        # pylint: disable=duplicate-code
-        tokenizer = LlamaTokenizer.from_pretrained(
-            "casperhansen/mistral-7b-instruct-v0.1-awq"
-        )
-        tokenizer.add_special_tokens(
-            {
-                "eos_token": AddedToken(
-                    "<|im_end|>", rstrip=False, lstrip=False, normalized=False
-                )
-            }
-        )
-        tokenizer.add_tokens(
-            [
-                AddedToken(
-                    "<|im_start|>", rstrip=False, lstrip=False, normalized=False
-                ),
-            ]
-        )
-        self.tokenizer = tokenizer
-        self.dataset = load_dataset(
-            "argilla/ultrafeedback-binarized-preferences-cleaned", split="train"
-        ).select([0])
+class OrpoTokenizationTest:
+    """test case for the ORPO tokenization"""
 
-    def test_orpo_integration(self):
+    @enable_hf_offline
+    def test_orpo_integration(
+        self,
+        tokenizer_mistral_7b_instruct_chatml,
+        dataset_argilla_ultrafeedback_binarized_preferences_cleaned,
+    ):
+        ds = dataset_argilla_ultrafeedback_binarized_preferences_cleaned.select([0])
         strat = load(
-            self.tokenizer,
+            tokenizer_mistral_7b_instruct_chatml,
             DictDefault({"train_on_inputs": False}),
             DictDefault({"chat_template": "chatml"}),
         )
-        res = strat.tokenize_prompt(self.dataset[0])
-        assert "rejected_input_ids" in res
+        res = strat.tokenize_prompt(ds[0])
+        assert "rejected_ids" in res
         assert "rejected_labels" in res
         assert "input_ids" in res
         assert "labels" in res
         assert "prompt_attention_mask" in res
 
-        assert len(res["rejected_input_ids"]) == len(res["rejected_labels"])
+        assert len(res["rejected_ids"]) == len(res["rejected_labels"])
         assert len(res["input_ids"]) == len(res["labels"])
         assert len(res["input_ids"]) == len(res["prompt_attention_mask"])
 
         assert res["rejected_labels"][0] == -100
-        assert res["rejected_input_ids"][-1] == res["rejected_labels"][-1]
+        assert res["rejected_ids"][-1] == res["rejected_labels"][-1]
 
         assert res["labels"][0] == -100
         assert res["input_ids"][-1] == res["labels"][-1]
 
         assert res["prompt_attention_mask"][0] == 1
         assert res["prompt_attention_mask"][-1] == 0
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/tests/test_prompters.py b/tests/test_prompters.py
index 6c5b8f27c2..3d61398e04 100644
--- a/tests/test_prompters.py
+++ b/tests/test_prompters.py
@@ -42,6 +42,19 @@ def test_prompt_style_w_instruct(self):
         assert "USER:" not in res
         assert "ASSISTANT:" not in res
 
+    def test_prompt_style_w_phi(self):
+        prompter = AlpacaPrompter(prompt_style=PromptStyle.PHI.value)
+        res = next(prompter.build_prompt("tell me a joke about the following"))
+        assert (
+            """<|system|>
+Below is an instruction that describes a task. Write a response that appropriately completes the request.<|end|>
+<|user|>
+tell me a joke about the following<|end|>
+<|assistant|>
+"""
+            == res
+        )
+
     def test_prompt_style_w_chat(self):
         prompter = AlpacaPrompter(prompt_style=PromptStyle.CHAT.value)
         res = next(
diff --git a/tests/test_revision_parameter.py b/tests/test_revision_parameter.py
new file mode 100644
index 0000000000..112badfb80
--- /dev/null
+++ b/tests/test_revision_parameter.py
@@ -0,0 +1,240 @@
+"""Tests for revision_of_model being passed to tokenizer and processor loaders."""
+
+from unittest.mock import MagicMock, patch
+
+from transformers import PreTrainedTokenizerBase
+
+from axolotl.utils.dict import DictDefault
+
+
+class TestRevisionParameter:
+    """Tests for revision_of_model being passed to tokenizer and processor loaders."""
+
+    @patch("axolotl.loaders.tokenizer.load_model_config")
+    @patch("axolotl.loaders.tokenizer.AutoTokenizer")
+    @patch(
+        "axolotl.loaders.patch_manager.PatchManager.apply_pre_tokenizer_load_patches"
+    )
+    def test_load_tokenizer_passes_revision(
+        self, _mock_patches, mock_auto_tokenizer, _mock_load_config
+    ):
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.__class__.__name__ = "MockTokenizer"
+        mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
+
+        cfg = DictDefault(
+            {
+                "tokenizer_config": "some-model",
+                "revision_of_model": "abc123",
+            }
+        )
+        from axolotl.loaders.tokenizer import load_tokenizer
+
+        load_tokenizer(cfg)
+
+        call_kwargs = mock_auto_tokenizer.from_pretrained.call_args
+        assert call_kwargs.kwargs.get("revision") == "abc123"
+
+    @patch("axolotl.loaders.tokenizer.load_model_config")
+    @patch("axolotl.loaders.tokenizer.AutoTokenizer")
+    @patch(
+        "axolotl.loaders.patch_manager.PatchManager.apply_pre_tokenizer_load_patches"
+    )
+    def test_load_tokenizer_omits_revision_when_unset(
+        self, _mock_patches, mock_auto_tokenizer, _mock_load_config
+    ):
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.__class__.__name__ = "MockTokenizer"
+        mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
+
+        cfg = DictDefault(
+            {
+                "tokenizer_config": "some-model",
+            }
+        )
+        from axolotl.loaders.tokenizer import load_tokenizer
+
+        load_tokenizer(cfg)
+
+        call_kwargs = mock_auto_tokenizer.from_pretrained.call_args
+        assert "revision" not in call_kwargs.kwargs
+
+    @patch("axolotl.loaders.tokenizer.AutoTokenizer")
+    @patch("axolotl.loaders.tokenizer.is_local_main_process", return_value=True)
+    @patch("axolotl.loaders.tokenizer.barrier")
+    def test_modify_tokenizer_files_passes_revision(
+        self, _mock_barrier, _mock_main, mock_auto_tokenizer, temp_dir
+    ):
+        mock_tokenizer = MagicMock()
+        mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
+
+        from axolotl.loaders.tokenizer import modify_tokenizer_files
+
+        modify_tokenizer_files("some-model", {}, output_dir=temp_dir, revision="abc123")
+
+        call_kwargs = mock_auto_tokenizer.from_pretrained.call_args
+        assert call_kwargs.kwargs.get("revision") == "abc123"
+
+    @patch("axolotl.loaders.tokenizer.AutoTokenizer")
+    @patch("axolotl.loaders.tokenizer.is_local_main_process", return_value=True)
+    @patch("axolotl.loaders.tokenizer.barrier")
+    def test_modify_tokenizer_files_defaults_revision_to_main(
+        self, _mock_barrier, _mock_main, mock_auto_tokenizer, temp_dir
+    ):
+        mock_tokenizer = MagicMock()
+        mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
+
+        from axolotl.loaders.tokenizer import modify_tokenizer_files
+
+        modify_tokenizer_files("some-model", {}, output_dir=temp_dir)
+
+        call_kwargs = mock_auto_tokenizer.from_pretrained.call_args
+        assert call_kwargs.kwargs.get("revision") == "main"
+
+    @patch("axolotl.loaders.processor.AutoProcessor")
+    def test_load_processor_passes_revision(self, mock_auto_processor):
+        mock_processor = MagicMock()
+        mock_processor.size = {}
+        mock_auto_processor.from_pretrained.return_value = mock_processor
+
+        cfg = DictDefault(
+            {
+                "processor_config": "some-model",
+                "revision_of_model": "abc123",
+                "trust_remote_code": False,
+            }
+        )
+        tokenizer = MagicMock(spec=PreTrainedTokenizerBase)
+
+        from axolotl.loaders.processor import load_processor
+
+        load_processor(cfg, tokenizer)
+
+        call_kwargs = mock_auto_processor.from_pretrained.call_args
+        assert call_kwargs.kwargs.get("revision") == "abc123"
+
+    @patch("axolotl.loaders.processor.AutoProcessor")
+    def test_load_processor_omits_revision_when_unset(self, mock_auto_processor):
+        mock_processor = MagicMock()
+        mock_processor.size = {}
+        mock_auto_processor.from_pretrained.return_value = mock_processor
+
+        cfg = DictDefault(
+            {
+                "processor_config": "some-model",
+                "trust_remote_code": False,
+            }
+        )
+        tokenizer = MagicMock(spec=PreTrainedTokenizerBase)
+
+        from axolotl.loaders.processor import load_processor
+
+        load_processor(cfg, tokenizer)
+
+        call_kwargs = mock_auto_processor.from_pretrained.call_args
+        assert "revision" not in call_kwargs.kwargs
+
+    @patch("axolotl.loaders.processor.AutoProcessor")
+    def test_load_processor_forwards_processor_kwargs(self, mock_auto_processor):
+        mock_processor = MagicMock()
+        mock_processor.size = {}
+        mock_auto_processor.from_pretrained.return_value = mock_processor
+
+        cfg = DictDefault(
+            {
+                "processor_config": "some-model",
+                "trust_remote_code": False,
+                "processor_kwargs": {
+                    "image_seq_length": 1120,
+                    "max_soft_tokens": 1120,
+                },
+            }
+        )
+        tokenizer = MagicMock(spec=PreTrainedTokenizerBase)
+
+        from axolotl.loaders.processor import load_processor
+
+        load_processor(cfg, tokenizer)
+
+        call_kwargs = mock_auto_processor.from_pretrained.call_args
+        assert call_kwargs.kwargs.get("image_seq_length") == 1120
+        assert call_kwargs.kwargs.get("max_soft_tokens") == 1120
+
+    @patch("axolotl.loaders.processor.AutoProcessor")
+    def test_load_processor_omits_processor_kwargs_when_unset(
+        self, mock_auto_processor
+    ):
+        mock_processor = MagicMock()
+        mock_processor.size = {}
+        mock_auto_processor.from_pretrained.return_value = mock_processor
+
+        cfg = DictDefault(
+            {
+                "processor_config": "some-model",
+                "trust_remote_code": False,
+            }
+        )
+        tokenizer = MagicMock(spec=PreTrainedTokenizerBase)
+
+        from axolotl.loaders.processor import load_processor
+
+        load_processor(cfg, tokenizer)
+
+        call_kwargs = mock_auto_processor.from_pretrained.call_args
+        assert "image_seq_length" not in call_kwargs.kwargs
+        assert "max_soft_tokens" not in call_kwargs.kwargs
+
+    def test_processor_kwargs_schema_rejects_revision(self):
+        import pytest
+
+        from axolotl.utils.schemas.model import ModelInputConfig
+
+        with pytest.raises(ValueError, match="revision"):
+            ModelInputConfig(
+                base_model="some-model",
+                processor_kwargs={"revision": "abc123"},
+            )
+
+    def test_processor_kwargs_schema_rejects_trust_remote_code(self):
+        import pytest
+
+        from axolotl.utils.schemas.model import ModelInputConfig
+
+        with pytest.raises(ValueError, match="trust_remote_code"):
+            ModelInputConfig(
+                base_model="some-model",
+                processor_kwargs={"trust_remote_code": True},
+            )
+
+    def test_processor_kwargs_schema_accepts_valid_keys(self):
+        from axolotl.utils.schemas.model import ModelInputConfig
+
+        cfg = ModelInputConfig(
+            base_model="some-model",
+            processor_kwargs={"image_seq_length": 1120, "max_soft_tokens": 1120},
+        )
+        assert cfg.processor_kwargs == {
+            "image_seq_length": 1120,
+            "max_soft_tokens": 1120,
+        }
+
+    def test_processor_kwargs_schema_accepts_none_and_empty(self):
+        from axolotl.utils.schemas.model import ModelInputConfig
+
+        assert ModelInputConfig(base_model="x").processor_kwargs is None
+        assert (
+            ModelInputConfig(base_model="x", processor_kwargs={}).processor_kwargs == {}
+        )
+
+    def test_processor_kwargs_incompatible_with_mistral_common(self, min_base_cfg):
+        import pytest
+
+        from axolotl.utils.config import validate_config
+        from axolotl.utils.dict import DictDefault
+
+        cfg = min_base_cfg | DictDefault(
+            tokenizer_use_mistral_common=True,
+            processor_kwargs={"image_seq_length": 1120},
+        )
+        with pytest.raises(ValueError, match="processor_kwargs"):
+            validate_config(cfg)
diff --git a/tests/test_save_deduplicated.py b/tests/test_save_deduplicated.py
new file mode 100644
index 0000000000..1e41c3e10d
--- /dev/null
+++ b/tests/test_save_deduplicated.py
@@ -0,0 +1,210 @@
+"""Tests to verify that deduplication runs before dataset saving during preprocessing.
+
+This addresses GitHub issue #2719: Save De-duplicated Set During Pre-processing.
+"""
+
+from unittest.mock import MagicMock, patch
+
+from datasets import Dataset
+
+from axolotl.utils.dict import DictDefault
+
+
+class TestSFTSaveDeduplicatedBeforeSave:
+    """Verify that in SFT data loading, deduplication occurs before saving."""
+
+    @patch("axolotl.utils.data.sft.save_preprocessed_dataset")
+    @patch("axolotl.utils.data.sft.generate_dataset_hash_from_config")
+    @patch("axolotl.utils.data.sft.deduplicate_and_log_datasets")
+    @patch("axolotl.utils.data.sft.merge_datasets")
+    @patch("axolotl.utils.data.sft._load_and_process_single_dataset")
+    @patch("axolotl.utils.data.sft.datasets_with_name_generator")
+    def test_dedup_called_before_save_sft(
+        self,
+        mock_datasets_gen,
+        mock_load_single,
+        mock_merge,
+        mock_dedup,
+        mock_gen_hash,
+        mock_save,
+    ):
+        """Deduplication should be called before save_preprocessed_dataset in SFT."""
+        from axolotl.utils.data.sft import _load_raw_datasets
+
+        # Set up mock data
+        dataset = Dataset.from_dict({"text": ["a", "b", "a"], "label": [1, 2, 1]})
+        deduped_dataset = Dataset.from_dict({"text": ["a", "b"], "label": [1, 2]})
+
+        mock_datasets_gen.return_value = [
+            DictDefault({"path": "test", "type": "alpaca"})
+        ]
+        mock_load_single.return_value = (dataset, None)
+        mock_merge.return_value = dataset
+        mock_dedup.return_value = (deduped_dataset, None)
+        mock_gen_hash.return_value = "testhash"
+
+        cfg = DictDefault(
+            {
+                "skip_prepare_dataset": False,
+                "dataset_exact_deduplication": True,
+                "sequence_len": 1024,
+                "eval_sequence_len": None,
+                "sample_packing": False,
+                "is_preprocess": False,
+                "seed": 42,
+                "datasets": [{"path": "test", "type": "alpaca"}],
+            }
+        )
+
+        tokenizer = MagicMock()
+        tokenizer.name_or_path = "test-tokenizer"
+
+        # Track call order
+        call_order = []
+        mock_dedup.side_effect = lambda **kwargs: (
+            call_order.append("dedup") or (deduped_dataset, None)
+        )
+        mock_save.side_effect = lambda *args, **kwargs: call_order.append("save")
+
+        _load_raw_datasets(
+            cfg=cfg,
+            datasets_configs=cfg.datasets,
+            tokenizer=tokenizer,
+            split="train",
+        )
+
+        # Verify dedup was called
+        assert "dedup" in call_order, "Deduplication should have been called"
+        # Verify save was called
+        assert "save" in call_order, "Save should have been called"
+        # Verify dedup happened before save
+        assert call_order.index("dedup") < call_order.index("save"), (
+            "Deduplication must occur before saving the dataset"
+        )
+
+    @patch("axolotl.utils.data.sft.save_preprocessed_dataset")
+    @patch("axolotl.utils.data.sft.generate_dataset_hash_from_config")
+    @patch("axolotl.utils.data.sft.merge_datasets")
+    @patch("axolotl.utils.data.sft._load_and_process_single_dataset")
+    @patch("axolotl.utils.data.sft.datasets_with_name_generator")
+    def test_no_dedup_when_disabled_sft(
+        self,
+        mock_datasets_gen,
+        mock_load_single,
+        mock_merge,
+        mock_gen_hash,
+        mock_save,
+    ):
+        """Deduplication should not be called when dataset_exact_deduplication is False."""
+        from axolotl.utils.data.sft import _load_raw_datasets
+
+        dataset = Dataset.from_dict({"text": ["a", "b", "a"], "label": [1, 2, 1]})
+
+        mock_datasets_gen.return_value = [
+            DictDefault({"path": "test", "type": "alpaca"})
+        ]
+        mock_load_single.return_value = (dataset, None)
+        mock_merge.return_value = dataset
+        mock_gen_hash.return_value = "testhash"
+
+        cfg = DictDefault(
+            {
+                "skip_prepare_dataset": False,
+                "dataset_exact_deduplication": False,
+                "sequence_len": 1024,
+                "eval_sequence_len": None,
+                "sample_packing": False,
+                "is_preprocess": False,
+                "seed": 42,
+                "datasets": [{"path": "test", "type": "alpaca"}],
+            }
+        )
+
+        tokenizer = MagicMock()
+        tokenizer.name_or_path = "test-tokenizer"
+
+        with patch("axolotl.utils.data.sft.deduplicate_and_log_datasets") as mock_dedup:
+            _load_raw_datasets(
+                cfg=cfg,
+                datasets_configs=cfg.datasets,
+                tokenizer=tokenizer,
+                split="train",
+            )
+            mock_dedup.assert_not_called()
+
+
+class TestRLSaveDeduplicatedBeforeSave:
+    """Verify that in RL data loading, deduplication occurs before saving."""
+
+    @patch.object(Dataset, "filter", lambda self, *args, **kwargs: self)
+    @patch("axolotl.utils.data.rl.save_preprocessed_dataset")
+    @patch("axolotl.utils.data.rl.generate_dataset_hash_from_config")
+    @patch("axolotl.utils.data.rl.deduplicate_and_log_datasets")
+    @patch("axolotl.utils.data.rl.merge_datasets")
+    @patch("axolotl.utils.data.rl.load_dataset_with_config")
+    @patch("axolotl.utils.data.rl.datasets_with_name_generator")
+    @patch("axolotl.utils.data.rl.load_tokenizer")
+    def test_dedup_called_before_save_rl(
+        self,
+        mock_load_tokenizer,
+        mock_datasets_gen,
+        mock_load_dataset,
+        mock_merge,
+        mock_dedup,
+        mock_gen_hash,
+        mock_save,
+    ):
+        """Deduplication should be called before save_preprocessed_dataset in RL."""
+        from axolotl.utils.data.rl import _load_split
+
+        dataset = Dataset.from_dict(
+            {
+                "prompt": ["hi", "bye", "hi"],
+                "chosen": ["a", "b", "a"],
+                "rejected": ["c", "d", "c"],
+            }
+        )
+        deduped_dataset = Dataset.from_dict(
+            {
+                "prompt": ["hi", "bye"],
+                "chosen": ["a", "b"],
+                "rejected": ["c", "d"],
+            }
+        )
+
+        mock_datasets_gen.return_value = [DictDefault({"path": "test", "type": None})]
+        mock_load_dataset.return_value = dataset
+        mock_merge.return_value = dataset
+        mock_dedup.return_value = (deduped_dataset, None)
+        mock_gen_hash.return_value = "testhash"
+
+        tokenizer = MagicMock()
+        tokenizer.name_or_path = "test-tokenizer"
+        mock_load_tokenizer.return_value = tokenizer
+
+        cfg = DictDefault(
+            {
+                "skip_prepare_dataset": False,
+                "dataset_exact_deduplication": True,
+                "sequence_len": 1024,
+                "rl": "dpo",
+                "datasets": [{"path": "test", "type": None}],
+                "hf_use_auth_token": False,
+                "dataset_num_proc": 1,
+                "is_preprocess": False,
+            }
+        )
+
+        call_order = []
+        mock_dedup.side_effect = lambda **kwargs: (
+            call_order.append("dedup") or (deduped_dataset, None)
+        )
+        mock_save.side_effect = lambda *args, **kwargs: call_order.append("save")
+
+        _load_split(cfg, split="train")
+
+        assert "dedup" in call_order, "Deduplication should have been called"
+        assert "save" in call_order, "Save should have been called"
+        assert call_order.index("dedup") < call_order.index("save"), (
+            "Deduplication must occur before saving the dataset"
+        )
diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py
index 9402d7af7f..c783a68db7 100644
--- a/tests/test_schedulers.py
+++ b/tests/test_schedulers.py
@@ -1,6 +1,7 @@
 """
 test module for the axolotl.utis.data module
 """
+
 import unittest
 
 import torch
@@ -21,7 +22,7 @@ def setUp(self):
         self.constant_lr_ratio = 0.8
         self._lr = 0.01
         self.optimizer = SGD([torch.tensor(1)], lr=self._lr)
-        self.lr_scheduler = get_cosine_schedule_with_warmup_decay_constant(  # pylint: disable=attribute-defined-outside-init
+        self.lr_scheduler = get_cosine_schedule_with_warmup_decay_constant(
             self.optimizer,
             num_warmup_steps=self.warmup_steps,
             num_training_steps=self.train_steps,
@@ -32,16 +33,19 @@ def setUp(self):
     def test_schedulers(self):
         self.assertEqual(self.lr_scheduler.get_last_lr()[0], 0)
         for _ in range(self.warmup_steps):
+            self.optimizer.step()
             self.lr_scheduler.step()
         self.assertEqual(self.lr_scheduler.get_last_lr()[0], self._lr)
         constant_step = int(self.train_steps * self.constant_lr_ratio)
         remaining_step = self.train_steps - constant_step
         for _ in range(constant_step):
+            self.optimizer.step()
             self.lr_scheduler.step()
         self.assertEqual(
             self.lr_scheduler.get_last_lr()[0], self._lr * self.min_lr_ratio
         )
         for _ in range(remaining_step):
+            self.optimizer.step()
             self.lr_scheduler.step()
         self.assertEqual(
             self.lr_scheduler.get_last_lr()[0], self._lr * self.min_lr_ratio
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
new file mode 100644
index 0000000000..2c1f9f936d
--- /dev/null
+++ b/tests/test_streaming.py
@@ -0,0 +1,238 @@
+"""Test streaming configuration and data loading functionality."""
+
+import unittest
+from unittest.mock import Mock, patch
+
+from datasets import IterableDataset
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.data.sft import (
+    _prepare_streaming_dataset,
+    prepare_datasets,
+)
+from axolotl.utils.dict import DictDefault
+
+
+class TestStreamingConfig(unittest.TestCase):
+    """Test streaming configuration and deprecation handling."""
+
+    def test_streaming_multipack_buffer_size_deprecation(self):
+        """Test that pretrain_multipack_buffer_size is properly deprecated."""
+        # Test with old config name
+        cfg_old = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "pretrain_multipack_buffer_size": 5000,
+                "datasets": [{"path": "test/dataset", "type": "alpaca"}],
+                "sequence_len": 256,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 0.0001,
+            }
+        )
+
+        with self.assertLogs("axolotl.utils.schemas.validation", level="WARNING") as cm:
+            validated_cfg = validate_config(cfg_old)
+            self.assertIn("pretrain_multipack_buffer_size` is deprecated", cm.output[0])
+
+        self.assertEqual(validated_cfg.streaming_multipack_buffer_size, 5000)
+        self.assertIsNone(
+            getattr(validated_cfg, "pretrain_multipack_buffer_size", None)
+        )
+
+    def test_streaming_multipack_buffer_size_new(self):
+        """Test that new streaming_multipack_buffer_size works correctly."""
+        cfg_new = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "streaming_multipack_buffer_size": 7000,
+                "datasets": [{"path": "test/dataset", "type": "alpaca"}],
+                "sequence_len": 256,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 0.0001,
+            }
+        )
+
+        validated_cfg = validate_config(cfg_new)
+        self.assertEqual(validated_cfg.streaming_multipack_buffer_size, 7000)
+
+    def test_both_buffer_sizes_raises_error(self):
+        """Test that having both old and new buffer size configs raises an error."""
+        cfg_both = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "pretrain_multipack_buffer_size": 5000,
+                "streaming_multipack_buffer_size": 7000,
+                "datasets": [{"path": "test/dataset", "type": "alpaca"}],
+                "sequence_len": 256,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 0.0001,
+            }
+        )
+
+        with self.assertRaises(ValueError) as cm:
+            validate_config(cfg_both)
+        self.assertIn("both are set", str(cm.exception))
+
+
+class TestStreamingDatasetPreparation(unittest.TestCase):
+    """Test dataset preparation with streaming configuration."""
+
+    def setUp(self):
+        self.tokenizer = Mock()
+        self.tokenizer.pad_token_id = 0
+        self.tokenizer.eos_token_id = 1
+
+    @patch("axolotl.utils.data.sft._prepare_streaming_dataset")
+    def test_prepare_datasets_with_streaming_true(self, mock_prepare_streaming):
+        """Test that streaming=True triggers streaming dataset preparation."""
+        cfg = DictDefault(
+            {
+                "streaming": True,
+                "datasets": [{"path": "test/dataset", "type": "alpaca"}],
+            }
+        )
+
+        mock_prepare_streaming.return_value = (Mock(), None, 100, [])
+
+        prepare_datasets(cfg, self.tokenizer)
+
+        mock_prepare_streaming.assert_called_once_with(cfg, self.tokenizer, None)
+
+    @patch("axolotl.utils.data.sft._prepare_streaming_dataset")
+    def test_prepare_datasets_with_pretraining_dataset(self, mock_prepare_streaming):
+        """Test that pretraining_dataset triggers streaming dataset preparation."""
+        cfg = DictDefault(
+            {
+                "pretraining_dataset": "test/dataset",
+            }
+        )
+
+        mock_prepare_streaming.return_value = (Mock(), None, 100, [])
+
+        prepare_datasets(cfg, self.tokenizer)
+
+        mock_prepare_streaming.assert_called_once_with(cfg, self.tokenizer, None)
+
+    @patch("axolotl.utils.data.sft._prepare_standard_dataset")
+    def test_prepare_datasets_without_streaming(self, mock_prepare_standard):
+        """Test that without streaming, standard dataset preparation is used."""
+        cfg = DictDefault(
+            {
+                "datasets": [{"path": "test/dataset", "type": "alpaca"}],
+            }
+        )
+
+        mock_prepare_standard.return_value = (Mock(), None, 100, [])
+
+        prepare_datasets(cfg, self.tokenizer)
+
+        mock_prepare_standard.assert_called_once_with(cfg, self.tokenizer, None)
+
+
+class TestStreamingWithSamplePacking(unittest.TestCase):
+    """Test streaming dataset preparation with sample packing."""
+
+    def setUp(self):
+        self.tokenizer = Mock()
+        self.tokenizer.pad_token_id = 0
+        self.tokenizer.eos_token_id = 1
+
+    @patch("axolotl.utils.data.sft._load_streaming_dataset")
+    def test_streaming_sft_with_sample_packing_sets_split(self, mock_load_streaming):
+        """Test that streaming SFT with sample_packing sets default split."""
+        cfg = DictDefault(
+            {
+                "streaming": True,
+                "sample_packing": True,
+                "datasets": [{"path": "test/dataset", "type": "alpaca"}],
+                "sequence_len": 256,
+                "micro_batch_size": 1,
+            }
+        )
+
+        mock_load_streaming.return_value = Mock(spec=IterableDataset)
+
+        with patch("axolotl.utils.data.sft._load_and_prepare_datasets"):
+            _prepare_streaming_dataset(cfg, self.tokenizer, None)
+
+            # Check that the dataset config has split set to 'train'
+            call_args = mock_load_streaming.call_args
+            dataset_config = call_args[0][0]
+            self.assertEqual(dataset_config.split, "train")
+
+    def test_multipack_attn_forced_true_for_sft(self):
+        """Test that multipack_attn is forced to True for SFT with sample packing."""
+        from axolotl.utils.data.streaming import wrap_streaming_dataset
+
+        cfg = DictDefault(
+            {
+                "sample_packing": True,
+                "pretrain_multipack_attn": False,  # Should be overridden for SFT
+                "pretraining_dataset": None,  # This makes it SFT
+                "sequence_len": 256,
+                "micro_batch_size": 1,
+                "streaming_multipack_buffer_size": 1000,
+                "seed": 42,
+            }
+        )
+
+        mock_dataset = Mock()
+        mock_dataset.features = None  # For streaming datasets
+        mock_dataset.__iter__ = Mock(return_value=iter([]))  # Empty iterator
+        mock_dataset.map = Mock(return_value=mock_dataset)
+        mock_ds_wrapper = Mock()
+
+        with patch(
+            "axolotl.utils.data.streaming.PretrainingBatchSamplerDataCollatorForSeq2Seq"
+        ) as mock_collator:
+            with patch("axolotl.utils.data.streaming.encode_packed_streaming"):
+                wrap_streaming_dataset(
+                    mock_dataset, self.tokenizer, cfg, mock_ds_wrapper
+                )
+
+                # Check that multipack_attn=True was used in the collator
+                mock_collator.assert_called_once()
+                call_kwargs = mock_collator.call_args[1]
+                self.assertTrue(call_kwargs["multipack_attn"])
+
+    def test_multipack_attn_respects_config_for_pretraining(self):
+        """Test that multipack_attn respects config for pretraining datasets."""
+        from axolotl.utils.data.streaming import wrap_streaming_dataset
+
+        cfg = DictDefault(
+            {
+                "sample_packing": True,
+                "pretrain_multipack_attn": False,  # Should be respected for pretraining
+                "pretraining_dataset": "test/dataset",  # This makes it pretraining
+                "sequence_len": 256,
+                "micro_batch_size": 1,
+                "streaming_multipack_buffer_size": 1000,
+                "seed": 42,
+            }
+        )
+
+        mock_dataset = Mock()
+        mock_dataset.features = None  # For streaming datasets
+        mock_dataset.__iter__ = Mock(return_value=iter([]))  # Empty iterator
+        mock_dataset.map = Mock(return_value=mock_dataset)
+        mock_ds_wrapper = Mock()
+
+        with patch(
+            "axolotl.utils.data.streaming.PretrainingBatchSamplerDataCollatorForSeq2Seq"
+        ) as mock_collator:
+            with patch("axolotl.utils.data.streaming.encode_packed_streaming"):
+                wrap_streaming_dataset(
+                    mock_dataset, self.tokenizer, cfg, mock_ds_wrapper
+                )
+
+                # Check that multipack_attn=False was used (respecting config)
+                mock_collator.assert_called_once()
+                call_kwargs = mock_collator.call_args[1]
+                self.assertFalse(call_kwargs["multipack_attn"])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_tensor_parallel_batch_size.py b/tests/test_tensor_parallel_batch_size.py
new file mode 100644
index 0000000000..c6a8174fc6
--- /dev/null
+++ b/tests/test_tensor_parallel_batch_size.py
@@ -0,0 +1,55 @@
+"""Tests for batch_size calculation with tensor parallelism."""
+
+from unittest.mock import patch
+
+import addict
+import pytest
+
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="tp_base_cfg")
+def fixture_tp_base_cfg(min_base_cfg):
+    return (
+        DictDefault(
+            micro_batch_size=2,
+            gradient_accumulation_steps=4,
+            sequence_len=2048,
+            num_epochs=1,
+        )
+        | min_base_cfg
+    )
+
+
+class TestTensorParallelBatchSize:
+    """Verify batch_size scales by effective dp world_size when using tensor parallelism."""
+
+    @pytest.mark.parametrize(
+        "world_size, tensor_parallel_size, expected_batch_size",
+        [
+            (4, 1, 32),  # no TP: 2*4*4 = 32
+            (4, 2, 16),  # TP=2: 2*4*(4//2) = 16
+            (4, 4, 8),  # TP=4: 2*4*(4//4) = 8
+            (2, 2, 8),  # TP=ws: 2*4*(2//2) = 8 (no scaling)
+        ],
+    )
+    def test_batch_size_with_tensor_parallelism(
+        self,
+        tp_base_cfg,
+        monkeypatch,
+        world_size,
+        tensor_parallel_size,
+        expected_batch_size,
+    ):
+        monkeypatch.setenv("WORLD_SIZE", str(world_size))
+        tp_base_cfg["tensor_parallel_size"] = tensor_parallel_size
+        cfg = validate_config(tp_base_cfg)
+        # Mock load_model_config to avoid downloading the model and to bypass
+        # the tie_word_embeddings validation that blocks TP > 1.
+        with patch(
+            "axolotl.utils.config.load_model_config",
+            return_value=addict.Dict({"model_type": "llama"}),
+        ):
+            normalize_config(cfg)
+        assert cfg.batch_size == expected_batch_size
diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py
index 69c441f8c6..0f8c584e2f 100644
--- a/tests/test_tokenizers.py
+++ b/tests/test_tokenizers.py
@@ -1,19 +1,24 @@
 """
 Test cases for the tokenizer loading
 """
+
 import unittest
 
 import pytest
 
+from axolotl.loaders import load_tokenizer
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_tokenizer
+
+from tests.hf_offline_utils import enable_hf_offline
 
 
-class TestTokenizers(unittest.TestCase):
+class TestTokenizers:
     """
     test class for the load_tokenizer fn
     """
 
+    @pytest.mark.skip("LlamaTokenizer no longer has a Fast/Slow tokenizer")
+    @enable_hf_offline
     def test_default_use_fast(self):
         cfg = DictDefault(
             {
@@ -23,6 +28,8 @@ def test_default_use_fast(self):
         tokenizer = load_tokenizer(cfg)
         assert "Fast" in tokenizer.__class__.__name__
 
+    @pytest.mark.skip("LlamaTokenizer no longer has a Fast/Slow tokenizer")
+    @enable_hf_offline
     def test_dont_use_fast(self):
         cfg = DictDefault(
             {
@@ -33,6 +40,7 @@ def test_dont_use_fast(self):
         tokenizer = load_tokenizer(cfg)
         assert "Fast" not in tokenizer.__class__.__name__
 
+    @enable_hf_offline
     def test_special_tokens_modules_to_save(self):
         # setting special_tokens to new token
         cfg = DictDefault(
@@ -67,6 +75,7 @@ def test_special_tokens_modules_to_save(self):
         )
         load_tokenizer(cfg)
 
+    @enable_hf_offline
     def test_add_additional_special_tokens(self):
         cfg = DictDefault(
             {
@@ -75,12 +84,80 @@ def test_add_additional_special_tokens(self):
             }
         )
         tokenizer = load_tokenizer(cfg)
-        self.assertEqual(tokenizer("<|im_start|>user")["input_ids"], [1, 32000, 1404])
-        self.assertEqual(len(tokenizer), 32001)
+        assert "LlamaTokenizer" in tokenizer.__class__.__name__
+        assert tokenizer("<|im_start|>user")["input_ids"] == [1, 32000, 1792]
+        assert len(tokenizer) == 32001
 
         # ensure reloading the tokenizer again from cfg results in same vocab length
         tokenizer = load_tokenizer(cfg)
-        self.assertEqual(len(tokenizer), 32001)
+        assert len(tokenizer) == 32001
+
+    @enable_hf_offline
+    def test_added_tokens_overrides(self, temp_dir):
+        cfg = DictDefault(
+            {
+                # use with tokenizer that has reserved_tokens in added_tokens
+                "tokenizer_config": "NousResearch/Llama-3.2-1B",
+                "added_tokens_overrides": {
+                    128041: "RANDOM_OVERRIDE_1",
+                    128042: "RANDOM_OVERRIDE_2",
+                },
+                "output_dir": temp_dir,
+            }
+        )
+
+        tokenizer = load_tokenizer(cfg)
+        assert tokenizer.encode("RANDOM_OVERRIDE_1", add_special_tokens=False) == [
+            128041
+        ]
+        assert tokenizer.encode("RANDOM_OVERRIDE_2", add_special_tokens=False) == [
+            128042
+        ]
+        assert (
+            tokenizer.decode([128041, 128042]) == "RANDOM_OVERRIDE_1RANDOM_OVERRIDE_2"
+        )
+
+    @pytest.mark.skip("FIXME slow test sdist py3.11 + torch2.8.0")
+    @enable_hf_offline
+    def test_added_tokens_overrides_gemma3(self, temp_dir):
+        cfg = DictDefault(
+            {
+                # use with tokenizer that has reserved_tokens in added_tokens
+                "tokenizer_config": "mlx-community/gemma-3-4b-it-8bit",
+                "added_tokens_overrides": {
+                    256001: "RANDOM_OVERRIDE_1",
+                    256002: "RANDOM_OVERRIDE_2",
+                },
+                "output_dir": temp_dir,
+            }
+        )
+
+        tokenizer = load_tokenizer(cfg)
+        assert tokenizer.encode("RANDOM_OVERRIDE_1", add_special_tokens=False) == [
+            256001
+        ]
+        assert tokenizer.encode("RANDOM_OVERRIDE_2", add_special_tokens=False) == [
+            256002
+        ]
+        assert (
+            tokenizer.decode([256001, 256002]) == "RANDOM_OVERRIDE_1RANDOM_OVERRIDE_2"
+        )
+
+    @enable_hf_offline
+    def test_added_tokens_overrides_with_toolargeid(self, temp_dir):
+        cfg = DictDefault(
+            {
+                # use with tokenizer that has reserved_tokens in added_tokens
+                "tokenizer_config": "HuggingFaceTB/SmolLM2-135M",
+                "added_tokens_overrides": {1000000: "BROKEN_RANDOM_OVERRIDE_1"},
+                "output_dir": temp_dir,
+            }
+        )
+
+        with pytest.raises(
+            ValueError, match=r".*Token ID 1000000 not found in added_tokens.*"
+        ):
+            load_tokenizer(cfg)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_train.py b/tests/test_train.py
new file mode 100644
index 0000000000..2c29b58eee
--- /dev/null
+++ b/tests/test_train.py
@@ -0,0 +1,39 @@
+"""Test for batch size calculation for multi-gpu training."""
+
+import pytest
+
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture(name="train_base_cfg")
+def fixture_train_base_cfg(min_base_cfg):
+    return (
+        DictDefault(
+            micro_batch_size=2,
+            gradient_accumulation_steps=4,
+            sequence_len=2048,
+            sample_packing=True,
+            num_epochs=1,
+        )
+        | min_base_cfg
+    )
+
+
+class TestTrain:
+    """test class for train related tests"""
+
+    @pytest.mark.parametrize(
+        "world_size, expected_batch_size",
+        [
+            (1, 8),
+            (4, 32),
+        ],
+    )
+    def test_batch_size_ddp(
+        self, train_base_cfg, monkeypatch, world_size, expected_batch_size
+    ):
+        monkeypatch.setenv("WORLD_SIZE", str(world_size))
+        cfg = validate_config(train_base_cfg)
+        normalize_config(cfg)
+        assert cfg.batch_size == expected_batch_size
diff --git a/tests/test_triton_kernels.py b/tests/test_triton_kernels.py
new file mode 100644
index 0000000000..ffc8de865d
--- /dev/null
+++ b/tests/test_triton_kernels.py
@@ -0,0 +1,481 @@
+# Copyright 2026 Axolotl AI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+
+"""Unit tests for Triton kernels: entropy_from_logits and selective_log_softmax.
+
+Adapted from harness/test_entropy.py and harness/test_selective_logsoftmax.py
+into proper pytest tests, plus new OOB index safety tests.
+"""
+
+import math
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA required for Triton kernels"
+)
+
+
+# ---------------------------------------------------------------------------
+# Reference implementations
+# ---------------------------------------------------------------------------
+
+
+def _ref_entropy(logits):
+    """Reference entropy via log_softmax (numerically stable)."""
+    logp = F.log_softmax(logits.float(), dim=-1)
+    return -(logp.exp() * logp).sum(dim=-1)
+
+
+def _ref_selective_log_softmax(logits, index):
+    """Reference selective log softmax via PyTorch gather."""
+    squeeze = index.ndim == logits.ndim - 1
+    if squeeze:
+        index = index.unsqueeze(-1)
+    log_probs = F.log_softmax(logits.float(), dim=-1)
+    result = torch.gather(log_probs, dim=-1, index=index)
+    if squeeze:
+        result = result.squeeze(-1)
+    return result
+
+
+# ---------------------------------------------------------------------------
+# entropy_from_logits
+# ---------------------------------------------------------------------------
+
+
+class TestEntropyFromLogits:
+    @pytest.mark.parametrize(
+        "B,L",
+        [
+            (1, 128),
+            (1, 2048),
+            (4, 512),
+            (8, 256),
+            (1, 1),
+        ],
+    )
+    def test_correctness_various_shapes(self, B, L):
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        V = 1024
+        torch.manual_seed(42)
+        logits = torch.randn(B, L, V, device="cuda", dtype=torch.float32)
+        result = entropy_from_logits(logits)
+        expected = _ref_entropy(logits)
+        assert result.shape == (B, L)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_2d_input(self):
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        logits = torch.randn(16, 256, device="cuda", dtype=torch.float32)
+        result = entropy_from_logits(logits)
+        expected = _ref_entropy(logits)
+        assert result.shape == (16,)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_large_vocab(self):
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        V = 32000
+        logits = torch.randn(2, V, device="cuda", dtype=torch.float32)
+        result = entropy_from_logits(logits)
+        expected = _ref_entropy(logits)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_uniform_distribution(self):
+        """Uniform logits -> entropy = log(V)."""
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        V = 1024
+        logits = torch.zeros(2, V, device="cuda", dtype=torch.float32)
+        result = entropy_from_logits(logits)
+        expected_val = math.log(V)
+        torch.testing.assert_close(
+            result,
+            torch.full((2,), expected_val, device="cuda", dtype=torch.float32),
+            atol=1e-4,
+            rtol=1e-4,
+        )
+
+    def test_peaked_distribution(self):
+        """One-hot-like logits -> entropy near 0."""
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        logits = torch.full((2, 128), -100.0, device="cuda", dtype=torch.float32)
+        logits[:, 0] = 100.0
+        result = entropy_from_logits(logits)
+        assert (result < 1e-3).all()
+
+    def test_bfloat16(self):
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        logits = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16)
+        result = entropy_from_logits(logits)
+        expected = _ref_entropy(logits.float())
+        assert result.dtype == torch.bfloat16
+        torch.testing.assert_close(result.float(), expected, atol=5e-2, rtol=5e-2)
+
+    def test_float16(self):
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        logits = torch.randn(4, 256, device="cuda", dtype=torch.float16)
+        result = entropy_from_logits(logits)
+        expected = _ref_entropy(logits.float())
+        assert result.dtype == torch.float16
+        torch.testing.assert_close(result.float(), expected, atol=5e-2, rtol=5e-2)
+
+    def test_non_contiguous_3d_transpose(self):
+        """Non-contiguous 3D tensor via transpose(0,1)."""
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        V = 256
+        raw = torch.randn(32, 4, V, device="cuda", dtype=torch.float32)
+        logits = raw.transpose(0, 1)  # (4, 32, V) non-contiguous
+        assert not logits.is_contiguous()
+        result = entropy_from_logits(logits)
+        expected = _ref_entropy(logits)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_non_contiguous_3d_slice(self):
+        """Non-contiguous 3D tensor via batch slicing."""
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        V = 256
+        raw = torch.randn(8, 32, V, device="cuda", dtype=torch.float32)
+        logits = raw[::2]  # (4, 32, V) non-contiguous
+        assert not logits.is_contiguous()
+        result = entropy_from_logits(logits)
+        expected = _ref_entropy(logits)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_many_rows_beyond_max_grid(self):
+        """More rows than MAX_GRID (8192) to test chunked dispatch."""
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        logits = torch.randn(10000, 128, device="cuda", dtype=torch.float32)
+        result = entropy_from_logits(logits)
+        expected = _ref_entropy(logits)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_entropy_non_negative(self):
+        from axolotl.monkeypatch.trainer.utils import entropy_from_logits
+
+        logits = torch.randn(32, 512, device="cuda", dtype=torch.float32)
+        result = entropy_from_logits(logits)
+        assert (result >= -1e-5).all(), f"Negative entropy: {result.min()}"
+
+
+# ---------------------------------------------------------------------------
+# selective_log_softmax — forward correctness
+# ---------------------------------------------------------------------------
+
+
+class TestSelectiveLogSoftmax:
+    @pytest.mark.parametrize(
+        "B,L,K",
+        [
+            (1, 128, 1),
+            (4, 512, 1),
+            (8, 256, 1),
+            (4, 256, 4),
+            (4, 256, 7),
+            (15, 129, 1),  # non-power-of-2
+        ],
+    )
+    def test_correctness_various_shapes(self, B, L, K):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 1024
+        torch.manual_seed(42)
+        logits = torch.randn(B, L, V, device="cuda", dtype=torch.float32)
+        if K == 1:
+            index = torch.randint(0, V, (B, L), device="cuda")
+        else:
+            index = torch.randint(0, V, (B, L, K), device="cuda")
+        result = selective_log_softmax(logits, index)
+        expected = _ref_selective_log_softmax(logits, index)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_squeezed_index(self):
+        """Index with ndim == logits.ndim - 1 triggers squeeze path."""
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 256
+        logits = torch.randn(8, V, device="cuda", dtype=torch.float32)
+        index = torch.randint(0, V, (8,), device="cuda")
+        result = selective_log_softmax(logits, index)
+        expected = _ref_selective_log_softmax(logits, index)
+        assert result.shape == (8,)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_large_vocab(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 32000
+        logits = torch.randn(2, V, device="cuda", dtype=torch.float32)
+        index = torch.randint(0, V, (2, 1), device="cuda")
+        result = selective_log_softmax(logits, index)
+        expected = _ref_selective_log_softmax(logits, index)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_bfloat16(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 1024
+        torch.manual_seed(42)
+        logits = torch.randn(4, 128, V, device="cuda", dtype=torch.bfloat16)
+        index = torch.randint(0, V, (4, 128), device="cuda")
+        result = selective_log_softmax(logits, index)
+        expected = _ref_selective_log_softmax(logits.float(), index)
+        assert result.dtype == torch.bfloat16
+        torch.testing.assert_close(result.float(), expected, atol=0.1, rtol=0.1)
+
+    def test_fp32_tight_tolerance(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 1024
+        torch.manual_seed(42)
+        logits = torch.randn(2, 256, V, device="cuda", dtype=torch.float32)
+        index = torch.randint(0, V, (2, 256), device="cuda")
+        result = selective_log_softmax(logits, index)
+        expected = _ref_selective_log_softmax(logits, index)
+        torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5)
+
+    def test_all_same_index(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 128
+        logits = torch.randn(8, V, device="cuda", dtype=torch.float32)
+        index = torch.zeros(8, 1, device="cuda", dtype=torch.long)
+        result = selective_log_softmax(logits, index)
+        expected = _ref_selective_log_softmax(logits, index)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_last_index(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 128
+        logits = torch.randn(8, V, device="cuda", dtype=torch.float32)
+        index = torch.full((8, 1), V - 1, device="cuda", dtype=torch.long)
+        result = selective_log_softmax(logits, index)
+        expected = _ref_selective_log_softmax(logits, index)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_output_always_nonpositive(self):
+        """Log softmax values should always be <= 0."""
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 256
+        logits = torch.randn(32, V, device="cuda", dtype=torch.float32)
+        index = torch.randint(0, V, (32, 1), device="cuda")
+        result = selective_log_softmax(logits, index)
+        assert (result <= 1e-5).all(), f"Positive log-prob: {result.max()}"
+
+    def test_many_rows_beyond_max_grid(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 128
+        logits = torch.randn(10000, V, device="cuda", dtype=torch.float32)
+        index = torch.randint(0, V, (10000, 1), device="cuda")
+        result = selective_log_softmax(logits, index)
+        expected = _ref_selective_log_softmax(logits, index)
+        torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4)
+
+
+# ---------------------------------------------------------------------------
+# selective_log_softmax — backward / gradient correctness
+# ---------------------------------------------------------------------------
+
+
+class TestSelectiveLogSoftmaxBackward:
+    @pytest.mark.parametrize(
+        "B,L,V,K",
+        [
+            (2, 16, 64, 1),
+            (2, 16, 64, 4),
+            (1, 8, 128, 1),
+            (2, 8, 128, 7),
+        ],
+    )
+    def test_gradient_matches_reference(self, B, L, V, K):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        torch.manual_seed(42)
+        logits_ref = torch.randn(
+            B, L, V, device="cuda", dtype=torch.float32, requires_grad=True
+        )
+        logits_tri = logits_ref.detach().clone().requires_grad_(True)
+
+        if K == 1:
+            index = torch.randint(0, V, (B, L), device="cuda")
+        else:
+            index = torch.randint(0, V, (B, L, K), device="cuda")
+
+        ref_out = _ref_selective_log_softmax(logits_ref, index)
+        tri_out = selective_log_softmax(logits_tri, index)
+
+        ref_out.sum().backward()
+        tri_out.sum().backward()
+
+        torch.testing.assert_close(
+            logits_tri.grad, logits_ref.grad, atol=1e-5, rtol=1e-5
+        )
+
+    def test_gradient_bfloat16_full_vocab(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 4096
+        torch.manual_seed(42)
+        logits_ref = torch.randn(
+            2, 64, V, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        logits_tri = logits_ref.detach().clone().requires_grad_(True)
+        index = torch.randint(0, V, (2, 64), device="cuda")
+
+        _ref_selective_log_softmax(logits_ref, index).sum().backward()
+        selective_log_softmax(logits_tri, index).sum().backward()
+
+        torch.testing.assert_close(
+            logits_tri.grad.float(), logits_ref.grad.float(), atol=0.1, rtol=0.1
+        )
+
+    def test_gradient_k1_squeezed(self):
+        """Gradient with squeezed (1D) index."""
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 256
+        logits = torch.randn(
+            8, V, device="cuda", dtype=torch.float32, requires_grad=True
+        )
+        index = torch.randint(0, V, (8,), device="cuda")
+
+        result = selective_log_softmax(logits, index)
+        result.sum().backward()
+        triton_grad = logits.grad.clone()
+
+        logits.grad = None
+        ref = torch.gather(
+            F.log_softmax(logits, dim=-1), dim=-1, index=index.unsqueeze(-1)
+        ).squeeze(-1)
+        ref.sum().backward()
+
+        torch.testing.assert_close(triton_grad, logits.grad, atol=1e-4, rtol=1e-4)
+
+
+# ---------------------------------------------------------------------------
+# selective_log_softmax — out-of-bounds index safety
+# ---------------------------------------------------------------------------
+
+
+class TestSelectiveLogSoftmaxOOBSafety:
+    """Verify that out-of-range indices don't crash or corrupt valid results."""
+
+    def test_negative_indices_no_crash(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 128
+        logits = torch.randn(4, V, device="cuda", dtype=torch.float32)
+        index = torch.tensor(
+            [[-1], [0], [V - 1], [-5]], device="cuda", dtype=torch.long
+        )
+        result = selective_log_softmax(logits, index)
+        assert result.shape == (4, 1)
+        # Valid rows should be finite and match reference
+        valid_idx = torch.tensor([[0], [V - 1]], device="cuda", dtype=torch.long)
+        valid_logits = logits[1:3]
+        expected = _ref_selective_log_softmax(valid_logits, valid_idx)
+        torch.testing.assert_close(result[1:3], expected, atol=1e-4, rtol=1e-4)
+
+    def test_index_exceeds_vocab_no_crash(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 128
+        logits = torch.randn(4, V, device="cuda", dtype=torch.float32)
+        index = torch.tensor(
+            [[0], [V], [V + 100], [V - 1]], device="cuda", dtype=torch.long
+        )
+        result = selective_log_softmax(logits, index)
+        assert result.shape == (4, 1)
+        # Valid rows (0 and 3) should match reference
+        for row_idx, idx_val in [(0, 0), (3, V - 1)]:
+            ref = _ref_selective_log_softmax(
+                logits[row_idx : row_idx + 1],
+                torch.tensor([[idx_val]], device="cuda", dtype=torch.long),
+            )
+            torch.testing.assert_close(
+                result[row_idx : row_idx + 1], ref, atol=1e-4, rtol=1e-4
+            )
+
+    def test_mixed_valid_invalid_multi_index(self):
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 256
+        K = 3
+        logits = torch.randn(4, V, device="cuda", dtype=torch.float32)
+        index = torch.tensor(
+            [
+                [0, 10, -1],  # last invalid
+                [V, 5, 100],  # first invalid
+                [50, 60, 70],  # all valid
+                [-1, V + 1, -100],  # all invalid
+            ],
+            device="cuda",
+            dtype=torch.long,
+        )
+        result = selective_log_softmax(logits, index)
+        assert result.shape == (4, K)
+        # Row 2 (all valid) must match reference exactly
+        valid_index = torch.tensor([[50, 60, 70]], device="cuda", dtype=torch.long)
+        expected = _ref_selective_log_softmax(logits[2:3], valid_index)
+        torch.testing.assert_close(result[2:3], expected, atol=1e-4, rtol=1e-4)
+
+    def test_oob_backward_no_crash(self):
+        """Backward with OOB indices should not crash and grads should be finite."""
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 128
+        logits = torch.randn(
+            4, V, device="cuda", dtype=torch.float32, requires_grad=True
+        )
+        index = torch.tensor(
+            [[-1], [0], [V + 10], [V - 1]], device="cuda", dtype=torch.long
+        )
+        result = selective_log_softmax(logits, index)
+        result.sum().backward()
+        assert logits.grad is not None
+        assert torch.isfinite(logits.grad).all()
+
+    def test_oob_backward_valid_rows_correct(self):
+        """Gradients for valid-index rows should match reference even when other rows have OOB."""
+        from axolotl.monkeypatch.trainer.utils import selective_log_softmax
+
+        V = 128
+        logits = torch.randn(
+            4, V, device="cuda", dtype=torch.float32, requires_grad=True
+        )
+        # Row 0: invalid, Row 1: valid, Row 2: invalid, Row 3: valid
+        index = torch.tensor(
+            [[-1], [42], [V + 5], [100]], device="cuda", dtype=torch.long
+        )
+        result = selective_log_softmax(logits, index)
+        result.sum().backward()
+
+        # Compute reference gradient for valid rows only
+        logits_ref = logits.detach().clone().requires_grad_(True)
+        valid_rows = [1, 3]
+        valid_indices = [42, 100]
+        for r, idx in zip(valid_rows, valid_indices, strict=True):
+            ref_lp = F.log_softmax(logits_ref[r : r + 1], dim=-1)
+            ref_val = ref_lp[0, idx]
+            ref_val.backward(retain_graph=True)
+
+        for r in valid_rows:
+            torch.testing.assert_close(
+                logits.grad[r], logits_ref.grad[r], atol=1e-4, rtol=1e-4
+            )
diff --git a/tests/test_utils_tee.py b/tests/test_utils_tee.py
new file mode 100644
index 0000000000..e2c1536672
--- /dev/null
+++ b/tests/test_utils_tee.py
@@ -0,0 +1,107 @@
+import os
+import tempfile
+
+
+def _dummy_cfg(output_dir: str, append: bool = False):
+    # Minimal object with attributes used by prepare_debug_log
+    class Cfg:
+        def __init__(self, out, append):
+            self.output_dir = out
+            self._append = append
+
+        def get(self, key, default=None):
+            if key in {"resume_from_checkpoint", "auto_resume_from_checkpoints"}:
+                return self._append
+            return default
+
+    return Cfg(output_dir, append)
+
+
+def read(path: str) -> str:
+    with open(path, "r", encoding="utf-8") as f:
+        return f.read()
+
+
+def test_file_only_stream_writes_after_prepare(monkeypatch):
+    from axolotl.utils import tee
+
+    with tempfile.TemporaryDirectory() as td:
+        # Avoid stdout tee in this test
+        monkeypatch.setenv("AXOLOTL_TEE_STDOUT", "0")
+        cfg = _dummy_cfg(td, append=False)
+
+        # before prepare: writing to file_only_stream creates no file
+        tee.file_only_stream.write("before\n")
+        tee.file_only_stream.flush()
+        assert not os.path.exists(os.path.join(td, "debug.log"))
+
+        # prepare and write
+        path = tee.prepare_debug_log(cfg)
+        assert os.path.basename(path) == "debug.log"
+        tee.file_only_stream.write("hello\n")
+        tee.file_only_stream.flush()
+
+        content = read(path)
+        assert "hello" in content
+
+        tee.close_debug_log()
+
+
+def test_stdout_is_mirrored_after_prepare(capsys, monkeypatch):
+    from axolotl.utils import tee
+
+    with tempfile.TemporaryDirectory() as td:
+        cfg = _dummy_cfg(td, append=False)
+        try:
+            # Install tee while capture is disabled so stdout tee wraps real stdout.
+            with capsys.disabled():
+                monkeypatch.setenv("AXOLOTL_TEE_STDOUT", "1")
+                path = tee.prepare_debug_log(cfg)
+                import sys
+
+                print("printed-line")
+                sys.stdout.flush()
+
+            # Now verify file contains the line
+            content = read(path)
+            assert "printed-line" in content
+        finally:
+            tee.close_debug_log()
+
+
+def test_truncate_vs_append_behavior(monkeypatch):
+    from axolotl.utils import tee
+
+    with tempfile.TemporaryDirectory() as td:
+        # Avoid stdout tee in this test
+        monkeypatch.setenv("AXOLOTL_TEE_STDOUT", "0")
+        # First run creates file with A
+        cfg = _dummy_cfg(td, append=False)
+        _ = tee.prepare_debug_log(cfg)
+        try:
+            tee.file_only_stream.write("A\n")
+            tee.file_only_stream.flush()
+        finally:
+            tee.close_debug_log()
+
+        # Second run with append=False truncates
+        cfg2 = _dummy_cfg(td, append=False)
+        path2 = tee.prepare_debug_log(cfg2)
+        try:
+            tee.file_only_stream.write("B\n")
+            tee.file_only_stream.flush()
+            content = read(path2)
+            assert "A\n" not in content and "B\n" in content
+        finally:
+            tee.close_debug_log()
+
+        # Third run with append=True preserves existing
+        cfg3 = _dummy_cfg(td, append=True)
+        path3 = tee.prepare_debug_log(cfg3)
+        try:
+            tee.file_only_stream.write("C\n")
+            tee.file_only_stream.flush()
+            content = read(path3)
+            assert "B\n" in content and "C\n" in content
+        finally:
+            tee.close_debug_log()
diff --git a/tests/test_validation_dataset.py b/tests/test_validation_dataset.py
new file mode 100644
index 0000000000..27740db182
--- /dev/null
+++ b/tests/test_validation_dataset.py
@@ -0,0 +1,368 @@
+"""Module for testing the validation module for the dataset config"""
+
+import warnings
+from typing import Optional
+
+import pytest
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.schemas.datasets import ChatTemplate
+
+warnings.filterwarnings("error")
+
+
+@pytest.fixture(name="minimal_cfg")
+def fixture_cfg():
+    return DictDefault(
+        {
+            "base_model": "TinyLlama/TinyLlama-1.1B-Chat-v0.6",
+            "learning_rate": 0.000001,
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+        }
+    )
+
+
+class BaseValidation:
+    """
+    Base validation module to setup the log capture
+    """
+
+    _caplog: Optional[pytest.LogCaptureFixture] = None
+
+    @pytest.fixture(autouse=True)
+    def inject_fixtures(self, caplog):
+        self._caplog = caplog
+
+
+class TestValidationCheckDatasetConfig(BaseValidation):
+    """
+    Test the validation for the dataset config to ensure no correct parameters are dropped
+    """
+
+    def test_dataset_config_no_drop_param(self, minimal_cfg):
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                        "shards": 10,
+                    }
+                ]
+            }
+        )
+
+        checked_cfg = validate_config(cfg)
+
+        def _check_config():
+            assert checked_cfg.datasets[0].path == cfg.datasets[0].path
+            assert checked_cfg.datasets[0].type == cfg.datasets[0].type
+            assert checked_cfg.datasets[0].shards == cfg.datasets[0].shards
+
+        _check_config()
+
+        checked_cfg = validate_config(
+            cfg,
+            capabilities={
+                "bf16": "false",
+                "tf32": "false",
+                "n_gpu": 1,
+                "compute_capability": "8.0",
+            },
+            env_capabilities={
+                "torch_version": "2.6.0",
+            },
+        )
+
+        _check_config()
+
+    def test_dataset_default_chat_template_no_drop_param(self, minimal_cfg):
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "datasets": [
+                    {
+                        "path": "LDJnr/Puffin",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "shards": 10,
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    }
+                ],
+            }
+        )
+
+        checked_cfg = validate_config(cfg)
+
+        def _check_config():
+            assert checked_cfg.datasets[0].path == cfg.datasets[0].path
+            assert checked_cfg.datasets[0].type == cfg.datasets[0].type
+            assert checked_cfg.chat_template is None
+            assert (
+                checked_cfg.datasets[0].chat_template == ChatTemplate.tokenizer_default
+            )
+            assert (
+                checked_cfg.datasets[0].field_messages == cfg.datasets[0].field_messages
+            )
+            assert checked_cfg.datasets[0].shards == cfg.datasets[0].shards
+            assert (
+                checked_cfg.datasets[0].message_field_role
+                == cfg.datasets[0].message_field_role
+            )
+            assert (
+                checked_cfg.datasets[0].message_field_content
+                == cfg.datasets[0].message_field_content
+            )
+
+        _check_config()
+
+        checked_cfg = validate_config(
+            cfg,
+            capabilities={
+                "bf16": "false",
+                "n_gpu": 1,
+                "compute_capability": "8.0",
+            },
+            env_capabilities={
+                "torch_version": "2.6.0",
+            },
+        )
+
+        _check_config()
+
+    def test_dataset_partial_default_chat_template_no_drop_param(self, minimal_cfg):
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "LDJnr/Puffin",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "shards": 10,
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    }
+                ],
+            }
+        )
+
+        checked_cfg = validate_config(cfg)
+
+        def _check_config():
+            assert checked_cfg.datasets[0].path == cfg.datasets[0].path
+            assert checked_cfg.datasets[0].type == cfg.datasets[0].type
+            assert checked_cfg.chat_template == ChatTemplate.chatml
+            assert (
+                checked_cfg.datasets[0].chat_template == ChatTemplate.tokenizer_default
+            )
+            assert (
+                checked_cfg.datasets[0].field_messages == cfg.datasets[0].field_messages
+            )
+            assert checked_cfg.datasets[0].shards == cfg.datasets[0].shards
+            assert (
+                checked_cfg.datasets[0].message_field_role
+                == cfg.datasets[0].message_field_role
+            )
+            assert (
+                checked_cfg.datasets[0].message_field_content
+                == cfg.datasets[0].message_field_content
+            )
+
+        _check_config()
+
+        checked_cfg = validate_config(
+            cfg,
+            capabilities={
+                "bf16": "false",
+                "n_gpu": 1,
+                "compute_capability": "8.0",
+            },
+            env_capabilities={
+                "torch_version": "2.6.0",
+            },
+        )
+
+        _check_config()
+
+    def test_dataset_chatml_chat_template_no_drop_param(self, minimal_cfg):
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "LDJnr/Puffin",
+                        "type": "chat_template",
+                        "chat_template": "gemma",
+                        "field_messages": "conversations",
+                        "shards": 10,
+                        "message_field_role": "from",
+                        "message_field_content": "value",
+                    }
+                ],
+            }
+        )
+
+        checked_cfg = validate_config(cfg)
+
+        def _check_config():
+            assert checked_cfg.datasets[0].path == cfg.datasets[0].path
+            assert checked_cfg.datasets[0].type == cfg.datasets[0].type
+            assert checked_cfg.chat_template == cfg.chat_template
+            assert (
+                checked_cfg.datasets[0].chat_template == cfg.datasets[0].chat_template
+            )
+            assert (
+                checked_cfg.datasets[0].field_messages == cfg.datasets[0].field_messages
+            )
+            assert checked_cfg.datasets[0].shards == cfg.datasets[0].shards
+            assert (
+                checked_cfg.datasets[0].message_field_role
+                == cfg.datasets[0].message_field_role
+            )
+            assert (
+                checked_cfg.datasets[0].message_field_content
+                == cfg.datasets[0].message_field_content
+            )
+
+        _check_config()
+
+        checked_cfg = validate_config(
+            cfg,
+            capabilities={
+                "bf16": "false",
+                "n_gpu": 1,
+                "compute_capability": "8.0",
+            },
+            env_capabilities={
+                "torch_version": "2.6.0",
+            },
+        )
+
+        _check_config()
+
+    def test_dataset_sharegpt_deprecation(self, minimal_cfg):
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "chat_template": "chatml",
+                "datasets": [
+                    {
+                        "path": "LDJnr/Puffin",
+                        "type": "sharegpt",
+                        "conversation": "chatml",
+                    }
+                ],
+            }
+        )
+
+        # Check sharegpt deprecation is raised
+        with pytest.raises(ValueError, match=r".*type: sharegpt.*` is deprecated.*"):
+            validate_config(cfg)
+
+        # Check that deprecation is not thrown for non-str type
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": {
+                            "field_instruction": "instruction",
+                            "field_output": "output",
+                            "field_system": "system",
+                            "format": "<|user|> {instruction} {input} <|model|>",
+                            "no_input_format": "<|user|> {instruction} <|model|>",
+                            "system_prompt": "",
+                        },
+                    }
+                ],
+            }
+        )
+
+        validate_config(cfg)
+
+        # Check that deprecation is not thrown for non-sharegpt type
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    }
+                ],
+            }
+        )
+
+        validate_config(cfg)
+
+    def test_message_property_mappings(self, minimal_cfg):
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                        "message_property_mappings": {
+                            "role": "role",
+                            "content": "content",
+                        },
+                    }
+                ],
+            }
+        )
+
+        validate_config(cfg)
+
+
+class TestOptimizerValidation(BaseValidation):
+    """
+    Test muon optimizer validation
+    """
+
+    def test_muon_deepspeed(self, minimal_cfg):
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    }
+                ],
+                "optimizer": "muon",
+                "deepspeed": "deepspeed_configs/zero3.json",
+            }
+        )
+
+        with pytest.raises(ValueError, match=r".*is currently incompatible with*"):
+            validate_config(cfg)
+
+    def test_muon_fsdp(self, minimal_cfg):
+        cfg = DictDefault(
+            minimal_cfg
+            | {
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    }
+                ],
+                "optimizer": "muon",
+                "fsdp": ["full_shard"],
+                "fsdp_config": {
+                    "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                },
+            }
+        )
+
+        with pytest.raises(ValueError, match=r".*only compatible with FSDP2.*"):
+            validate_config(cfg)
diff --git a/tests/test_vectorized_scanner.py b/tests/test_vectorized_scanner.py
new file mode 100644
index 0000000000..056b469814
--- /dev/null
+++ b/tests/test_vectorized_scanner.py
@@ -0,0 +1,650 @@
+"""Differential and targeted edge-case tests for the vectorized role-boundary scanner.
+
+The vectorized scanner (:func:`axolotl.processing_strategies._apply_role_boundaries_vectorized`)
+must be byte-identical to the reference implementation
+(:func:`axolotl.processing_strategies._apply_role_boundaries`) for every valid input.
+
+This file is structured in three layers:
+
+1. Targeted edge-case tests covering each behavioral subtlety called out in
+   the implementation comments (longest-prefix tie-break, Pixtral rewind,
+   train_on_eos modes, empty end_tokens, include_end leak gate).
+
+2. A boundary-shape catalog mimicking the real models: Gemma 4, Llama 3.2 V,
+   Llama 4, Pixtral, Mistral V7.
+
+3. A differential fuzzer that generates 2,000 random configurations and
+   asserts vectorized output == reference output element-wise. On mismatch,
+   dumps inputs + outputs + boundary spec.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import asdict
+from typing import Iterable
+
+import pytest
+import torch
+
+from axolotl.processing_strategies import (
+    RoleBoundary,
+    _apply_role_boundaries,
+    _apply_role_boundaries_vectorized,
+)
+
+# --------------------------------------------------------------------------- #
+# Helpers
+# --------------------------------------------------------------------------- #
+
+
+def _assert_equiv(
+    boundaries: list[RoleBoundary],
+    seq: list[list[int]],
+    roles_to_train: Iterable[str],
+    train_on_eos: str,
+):
+    """Assert vectorized output == reference output for the given input."""
+    labels_a = torch.tensor(seq)
+    labels_b = labels_a.clone()
+    out_ref = _apply_role_boundaries(
+        labels_a, boundaries, set(roles_to_train), train_on_eos
+    )
+    out_vec = _apply_role_boundaries_vectorized(
+        labels_b, boundaries, set(roles_to_train), train_on_eos
+    )
+    if not torch.equal(out_ref, out_vec):
+        # Build a focused failure dump.
+        diff = (out_ref != out_vec).nonzero(as_tuple=False).tolist()
+        msg = (
+            "Vectorized scanner diverged from reference.\n"
+            f"  boundaries: {[asdict(b) for b in boundaries]}\n"
+            f"  roles_to_train: {sorted(roles_to_train)}\n"
+            f"  train_on_eos: {train_on_eos}\n"
+            f"  input seq: {seq}\n"
+            f"  reference: {out_ref.tolist()}\n"
+            f"  vectorized: {out_vec.tolist()}\n"
+            f"  first 10 diff indices: {diff[:10]}\n"
+        )
+        raise AssertionError(msg)
+    return out_ref
+
+
+# --------------------------------------------------------------------------- #
+# Layer 3: Targeted edge-case unit tests
+# --------------------------------------------------------------------------- #
+
+
+class TestEdgeCases:
+    """Hand-written cases that pin down each tricky behavior."""
+
+    def test_longest_prefix_wins_at_same_position(self):
+        """``<|im_start|>assistant`` (long) beats ``<|im_start|>`` (short)."""
+        # token 100 = <|im_start|>, then 101 = "user", 102 = "assistant"
+        boundaries = [
+            RoleBoundary(role="user", start_tokens=[100], end_tokens=[200]),
+            RoleBoundary(role="assistant", start_tokens=[100, 102], end_tokens=[200]),
+        ]
+        seq = [[100, 102, 1, 2, 3, 200, 100, 101, 4, 5, 6, 200]]
+        # The longer "assistant" boundary should win at j=0; only positions
+        # 1..5 are trainable (assistant content + end marker because
+        # train_on_eos="turn" includes end on trainable turns by default).
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_pixtral_shared_end_marker_rewind(self):
+        """[/INST] both ends user and starts assistant; rewind must re-match."""
+        # Pixtral: user=[INST] ... [/INST], assistant content, then </s>.
+        # The shared [/INST] must consume as user-end *and* assistant-start.
+        INST_OPEN = 100  # [INST]
+        INST_CLOSE = 200  # [/INST]
+        EOS = 2  # </s>
+        boundaries = [
+            RoleBoundary(
+                role="user",
+                start_tokens=[INST_OPEN],
+                end_tokens=[INST_CLOSE],
+                include_start=False,
+                include_end=False,  # critical: don't consume — let assistant re-match
+            ),
+            RoleBoundary(
+                role="assistant",
+                start_tokens=[INST_CLOSE],
+                end_tokens=[EOS],
+                include_start=False,
+                include_end=True,
+            ),
+        ]
+        # [INST] hello [/INST] reply </s>
+        seq = [[INST_OPEN, 5, 6, INST_CLOSE, 7, 8, 9, EOS]]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_train_on_eos_turn(self):
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[200])
+        ]
+        seq = [[100, 1, 2, 3, 200, 9, 9, 9, 100, 4, 5, 6, 200]]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_train_on_eos_all(self):
+        boundaries = [
+            RoleBoundary(role="user", start_tokens=[100], end_tokens=[200]),
+            RoleBoundary(role="assistant", start_tokens=[101], end_tokens=[200]),
+        ]
+        seq = [[100, 1, 2, 200, 101, 3, 4, 200, 9]]
+        _assert_equiv(boundaries, seq, ["assistant"], "all")
+
+    def test_train_on_eos_none(self):
+        """train_on_eos=none disables the end-marker contribution entirely."""
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[200])
+        ]
+        seq = [[100, 1, 2, 3, 200, 9]]
+        _assert_equiv(boundaries, seq, ["assistant"], "none")
+
+    def test_train_on_eos_last_only_final_turn_unmasked(self):
+        """Only the last trainable turn's end marker contributes."""
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[200])
+        ]
+        # Three assistant turns. Only the last 200 should be unmasked.
+        seq = [[100, 1, 200, 100, 2, 200, 100, 3, 200, 9, 9]]
+        _assert_equiv(boundaries, seq, ["assistant"], "last")
+
+    def test_empty_end_tokens_runs_to_eos(self):
+        """end_tokens=[] means the span runs to end-of-sequence."""
+        boundaries = [RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[])]
+        seq = [[100, 1, 2, 3, 4, 5]]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_include_start_true(self):
+        boundaries = [
+            RoleBoundary(
+                role="assistant",
+                start_tokens=[100, 101],
+                end_tokens=[200],
+                include_start=True,
+                include_end=True,
+            )
+        ]
+        seq = [[100, 101, 5, 6, 7, 200]]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_non_trainable_role_end_marker_leak_gate(self):
+        """Non-trainable role with include_end=True, train_on_eos=all → end is unmasked."""
+        boundaries = [
+            RoleBoundary(
+                role="user",
+                start_tokens=[100],
+                end_tokens=[200],
+                include_start=False,
+                include_end=True,
+            ),
+            RoleBoundary(
+                role="assistant",
+                start_tokens=[101],
+                end_tokens=[201],
+                include_start=False,
+                include_end=True,
+            ),
+        ]
+        seq = [[100, 1, 2, 200, 101, 3, 4, 201]]
+        # roles_to_train=["assistant"] but train_on_eos="all" → user's end (200) leaks in
+        _assert_equiv(boundaries, seq, ["assistant"], "all")
+
+    def test_non_trainable_role_include_end_false_no_leak(self):
+        """Non-trainable role with include_end=False, train_on_eos=all → no leak."""
+        boundaries = [
+            RoleBoundary(
+                role="user",
+                start_tokens=[100],
+                end_tokens=[200],
+                include_start=False,
+                include_end=False,
+            ),
+            RoleBoundary(
+                role="assistant",
+                start_tokens=[200],  # shared with user-end
+                end_tokens=[201],
+            ),
+        ]
+        seq = [[100, 1, 2, 200, 3, 4, 201]]
+        _assert_equiv(boundaries, seq, ["assistant"], "all")
+
+    def test_truncated_final_turn_no_end(self):
+        """Final assistant turn missing end marker — span runs to end."""
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[200])
+        ]
+        seq = [[100, 1, 2, 3, 4, 5, 6]]  # no 200 anywhere
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_all_pad_row(self):
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[200])
+        ]
+        seq = [[0, 0, 0, 0, 0, 0]]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_single_token_row(self):
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[200])
+        ]
+        seq = [[100]]  # start with no end & no content
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_empty_roles_to_train_masks_all(self):
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[200])
+        ]
+        seq = [[100, 1, 2, 200, 100, 3, 4, 200]]
+        _assert_equiv(boundaries, seq, [], "turn")
+
+    def test_adversarial_first_token_collision(self):
+        """A bare token equal to start_tokens[0] inside filler must not trigger
+        a partial match (multi-token start_tokens needed)."""
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100, 200], end_tokens=[201])
+        ]
+        # 100 appears in filler (alone, not followed by 200) → must NOT match.
+        seq = [[100, 200, 5, 6, 100, 7, 8, 201, 100, 200, 9, 10, 201]]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_batch_of_mixed_rows(self):
+        boundaries = [
+            RoleBoundary(role="assistant", start_tokens=[100], end_tokens=[200])
+        ]
+        # Three rows with different shapes.
+        seq = [
+            [100, 1, 2, 200, 0, 0, 0, 0],
+            [100, 3, 4, 5, 6, 200, 0, 0],
+            [9, 9, 100, 7, 200, 9, 9, 9],
+        ]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+
+# --------------------------------------------------------------------------- #
+# Layer 2-ish: Boundary-shape catalog (real-model-like configs)
+# --------------------------------------------------------------------------- #
+
+
+def _gemma4_like():
+    """Gemma 4: <start_of_turn>role ... <end_of_turn>."""
+    SOT, EOT = 50, 60
+    return [
+        RoleBoundary(role="user", start_tokens=[SOT, 70], end_tokens=[EOT]),
+        RoleBoundary(role="assistant", start_tokens=[SOT, 71], end_tokens=[EOT]),
+        RoleBoundary(role="system", start_tokens=[SOT, 72], end_tokens=[EOT]),
+    ]
+
+
+def _llama32v_like():
+    """Llama 3.2 V: <|start_header_id|>role<|end_header_id|> ... <|eot_id|>."""
+    SHID, EHID, EOT = 80, 81, 82
+    return [
+        RoleBoundary(role="user", start_tokens=[SHID, 90, EHID], end_tokens=[EOT]),
+        RoleBoundary(role="assistant", start_tokens=[SHID, 91, EHID], end_tokens=[EOT]),
+        RoleBoundary(role="system", start_tokens=[SHID, 92, EHID], end_tokens=[EOT]),
+    ]
+
+
+def _llama4_like():
+    """Llama 4-style: similar to llama3 but 4 roles."""
+    SHID, EHID, EOT = 80, 81, 82
+    return [
+        RoleBoundary(role="user", start_tokens=[SHID, 90, EHID], end_tokens=[EOT]),
+        RoleBoundary(role="assistant", start_tokens=[SHID, 91, EHID], end_tokens=[EOT]),
+        RoleBoundary(role="system", start_tokens=[SHID, 92, EHID], end_tokens=[EOT]),
+        RoleBoundary(role="tool", start_tokens=[SHID, 93, EHID], end_tokens=[EOT]),
+    ]
+
+
+def _pixtral_like():
+    """Pixtral: shared [/INST] between user-end and assistant-start, EOS terminates."""
+    INST_O, INST_C, EOS = 100, 200, 2
+    return [
+        RoleBoundary(
+            role="user",
+            start_tokens=[INST_O],
+            end_tokens=[INST_C],
+            include_start=False,
+            include_end=False,
+        ),
+        RoleBoundary(
+            role="assistant",
+            start_tokens=[INST_C],
+            end_tokens=[EOS],
+            include_start=False,
+            include_end=True,
+        ),
+    ]
+
+
+def _mistralv7_like():
+    """Mistral V7 Tekken-ish: similar shared [/INST] pattern with system."""
+    INST_O, INST_C, EOS, SYS = 100, 200, 2, 110
+    return [
+        RoleBoundary(
+            role="system",
+            start_tokens=[SYS],
+            end_tokens=[INST_O],
+            include_start=False,
+            include_end=False,
+        ),
+        RoleBoundary(
+            role="user",
+            start_tokens=[INST_O],
+            end_tokens=[INST_C],
+            include_start=False,
+            include_end=False,
+        ),
+        RoleBoundary(
+            role="assistant",
+            start_tokens=[INST_C],
+            end_tokens=[EOS],
+            include_start=False,
+            include_end=True,
+        ),
+    ]
+
+
+BOUNDARY_CATALOG = {
+    "gemma4": _gemma4_like,
+    "llama32v": _llama32v_like,
+    "llama4": _llama4_like,
+    "pixtral": _pixtral_like,
+    "mistralv7": _mistralv7_like,
+}
+
+
+# --------------------------------------------------------------------------- #
+# Layer 1: Differential fuzz test
+# --------------------------------------------------------------------------- #
+
+
+def _build_random_sequence(
+    rng: random.Random,
+    boundaries: list[RoleBoundary],
+    seq_len: int,
+    pad_id: int = 0,
+) -> list[int]:
+    """Build a plausible sequence by alternating turns until we hit seq_len.
+
+    Uses random filler tokens that *don't* collide with any marker. ~5% of
+    fillers are pathologically chosen as adversarial collisions: a single
+    token equal to start_tokens[0] (so that multi-token starts don't match
+    but the first byte does).
+    """
+    # Build a vocab of "safe" filler tokens that don't collide with any
+    # start_tokens prefix.
+    marker_tokens: set[int] = set()
+    for b in boundaries:
+        marker_tokens.update(b.start_tokens)
+        marker_tokens.update(b.end_tokens)
+    safe_filler = [t for t in range(300, 500) if t not in marker_tokens]
+
+    seq: list[int] = []
+    while len(seq) < seq_len:
+        # Pick a boundary at random; emit start + filler + (sometimes) end.
+        b = rng.choice(boundaries)
+        seq.extend(b.start_tokens)
+        filler_n = rng.randint(5, min(200, max(5, seq_len - len(seq))))
+        for _ in range(filler_n):
+            if rng.random() < 0.05 and safe_filler:
+                # Adversarial: emit a token that equals a marker's first byte
+                # but isn't followed by the rest. Picks from any boundary.
+                bb = rng.choice(boundaries)
+                if bb.start_tokens:
+                    seq.append(bb.start_tokens[0])
+                    continue
+            seq.append(rng.choice(safe_filler))
+
+        # 80% chance of a clean end. (Truncated turns intentional.)
+        if rng.random() < 0.8 and b.end_tokens:
+            seq.extend(b.end_tokens)
+
+    seq = seq[:seq_len]
+    # Pad up if we underran.
+    while len(seq) < seq_len:
+        seq.append(pad_id)
+    return seq
+
+
+@pytest.mark.parametrize(
+    "shape_name",
+    ["gemma4", "llama32v", "llama4", "pixtral", "mistralv7"],
+)
+def test_catalog_smoke_per_shape(shape_name):
+    """Smoke test: each catalog shape works on a small batch."""
+    boundaries = BOUNDARY_CATALOG[shape_name]()
+    rng = random.Random(0xCAFE)
+    rows = [_build_random_sequence(rng, boundaries, 256) for _ in range(4)]
+    _assert_equiv(boundaries, rows, ["assistant"], "turn")
+    _assert_equiv(boundaries, rows, ["assistant", "user"], "all")
+    _assert_equiv(boundaries, rows, [], "none")
+    _assert_equiv(boundaries, rows, ["assistant"], "last")
+
+
+def test_differential_fuzz_2000_configs():
+    """Run 2000 random configurations and verify byte-identical outputs.
+
+    Uses fixed seeds 0..1999 so any failure is deterministically reproducible.
+    """
+    failures: list[tuple[int, str]] = []
+
+    BATCH_SIZES = [1, 2, 4, 8]
+    SEQ_LENS = [32, 256, 1024, 4096]
+    EOS_MODES = ["turn", "all", "none", "last"]
+    ROLES_OPTIONS = [
+        ["assistant"],
+        ["assistant", "user"],
+        [],
+        ["assistant", "system", "user"],
+    ]
+    SHAPES = list(BOUNDARY_CATALOG.keys())
+
+    N_CONFIGS = 2000
+
+    for seed in range(N_CONFIGS):
+        rng = random.Random(seed)
+        bs = rng.choice(BATCH_SIZES)
+        sl = rng.choice(SEQ_LENS)
+        eos = rng.choice(EOS_MODES)
+        rtt = rng.choice(ROLES_OPTIONS)
+        shape = rng.choice(SHAPES)
+        boundaries = BOUNDARY_CATALOG[shape]()
+
+        # Cap the largest pixtral-shape configs at 1024 to keep wall-time
+        # under control; the small/medium configs already cover the rewind.
+        if shape == "pixtral" and sl == 4096 and bs == 8:
+            sl = 1024
+
+        rows = [_build_random_sequence(rng, boundaries, sl) for _ in range(bs)]
+
+        try:
+            _assert_equiv(boundaries, rows, rtt, eos)
+        except AssertionError as e:
+            failures.append((seed, str(e)))
+            if len(failures) >= 5:
+                break
+
+    if failures:
+        joined = "\n\n---\n\n".join(f"seed={s}:\n{m}" for s, m in failures)
+        pytest.fail(
+            f"Differential fuzz: {len(failures)} mismatches in "
+            f"{N_CONFIGS} configs.\n\n{joined}"
+        )
+
+
+# --------------------------------------------------------------------------- #
+# Pathological inputs aimed at the rewind logic specifically
+# --------------------------------------------------------------------------- #
+
+
+class TestPixtralRewindAdversarial:
+    def test_back_to_back_user_assistant_pairs(self):
+        boundaries = _pixtral_like()
+        # Five back-to-back turns.
+        seq = [
+            [
+                100,
+                1,
+                2,
+                200,
+                3,
+                4,
+                5,
+                2,  # turn 1
+                100,
+                6,
+                7,
+                200,
+                8,
+                9,
+                2,  # turn 2
+                100,
+                10,
+                11,
+                12,
+                200,
+                13,
+                14,
+                2,
+                100,
+                15,
+                200,
+                16,
+                2,
+                100,
+                17,
+                18,
+                19,
+                20,
+                200,
+                21,
+                22,
+                2,
+            ]
+        ]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_user_with_no_end_then_assistant(self):
+        """Truncated user — never closes — assistant never re-matches the [/INST]."""
+        boundaries = _pixtral_like()
+        seq = [[100, 1, 2, 3, 4, 5, 6]]  # no [/INST] anywhere
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+    def test_double_end_marker(self):
+        """Two [/INST] in a row — second one starts a (degenerate) assistant."""
+        boundaries = _pixtral_like()
+        seq = [[100, 1, 200, 200, 5, 6, 2]]
+        _assert_equiv(boundaries, seq, ["assistant"], "turn")
+
+
+# --------------------------------------------------------------------------- #
+# Long-span / multi-end-marker inputs aimed at the bisect end-finder
+# --------------------------------------------------------------------------- #
+
+
+def _build_long_multi_end_sequence(
+    rng: random.Random,
+    boundaries: list[RoleBoundary],
+    seq_len: int,
+    pad_id: int = 0,
+) -> list[int]:
+    """Like ``_build_random_sequence`` but with long turns that embed *multiple*
+    full end-marker copies inside one content region.
+
+    The vectorized scanner finds turn ends by bisecting a sorted list of
+    end-match positions for the next end >= start_of_content. Turns that carry
+    several full end markers (plus partial first-byte collisions) exercise the
+    "pick the first valid end, not the closest scanned" branch that a linear
+    walk would otherwise mask.
+    """
+    marker_tokens: set[int] = set()
+    for b in boundaries:
+        marker_tokens.update(b.start_tokens)
+        marker_tokens.update(b.end_tokens)
+    safe_filler = [t for t in range(300, 500) if t not in marker_tokens]
+
+    seq: list[int] = []
+    while len(seq) < seq_len:
+        b = rng.choice(boundaries)
+        seq.extend(b.start_tokens)
+        # Long filler so a turn can span well past the old 200-token cap.
+        filler_n = rng.randint(200, max(200, min(1500, seq_len - len(seq) + 200)))
+        for _ in range(filler_n):
+            r = rng.random()
+            if r < 0.04 and b.end_tokens:
+                # Full extra end marker mid-content: with include_end this closes
+                # the turn early; without it the rewind re-reads it as a start.
+                seq.extend(b.end_tokens)
+            elif r < 0.09 and safe_filler:
+                bb = rng.choice(boundaries)
+                if bb.start_tokens:
+                    seq.append(bb.start_tokens[0])  # partial first-byte collision
+                    continue
+                seq.append(rng.choice(safe_filler))
+            else:
+                seq.append(rng.choice(safe_filler))
+        if rng.random() < 0.8 and b.end_tokens:
+            seq.extend(b.end_tokens)
+
+    seq = seq[:seq_len]
+    while len(seq) < seq_len:
+        seq.append(pad_id)
+    return seq
+
+
+def test_bisect_first_end_in_span_explicit():
+    """Two full end markers inside one assistant turn: the span must close on the
+    first, leaving the second outside the trainable region."""
+    boundaries = _pixtral_like()  # [/INST] == [200], rewind on include_end=False
+    # assistant turn opens at [/INST] (200), content, end-of-turn (2) appears
+    # twice; the first 2 closes the turn, everything after is a fresh scan.
+    seq = [[100, 1, 2, 200, 5, 6, 7, 2, 9, 9, 9, 2, 100, 11, 200, 12, 2]]
+    _assert_equiv(boundaries, seq, ["assistant"], "turn")
+    _assert_equiv(boundaries, seq, ["assistant"], "all")
+    _assert_equiv(boundaries, seq, ["assistant"], "last")
+    _assert_equiv(boundaries, seq, ["assistant"], "none")
+
+
+def test_differential_fuzz_long_spans():
+    """500 configs with long, multi-end-marker turns over large sequences.
+
+    Targets the bisect end-finder and bytearray slice-fills, which the original
+    short-span fuzz (filler <= 200) under-exercises.
+    """
+    failures: list[tuple[int, str]] = []
+
+    BATCH_SIZES = [1, 2, 4]
+    SEQ_LENS = [1024, 2048, 4096]
+    EOS_MODES = ["turn", "all", "none", "last"]
+    ROLES_OPTIONS = [["assistant"], ["assistant", "user"], [], ["assistant", "system"]]
+    SHAPES = list(BOUNDARY_CATALOG.keys())
+
+    N_CONFIGS = 500
+
+    for seed in range(N_CONFIGS):
+        rng = random.Random(10_000 + seed)
+        bs = rng.choice(BATCH_SIZES)
+        sl = rng.choice(SEQ_LENS)
+        eos = rng.choice(EOS_MODES)
+        rtt = rng.choice(ROLES_OPTIONS)
+        shape = rng.choice(SHAPES)
+        boundaries = BOUNDARY_CATALOG[shape]()
+
+        rows = [_build_long_multi_end_sequence(rng, boundaries, sl) for _ in range(bs)]
+
+        try:
+            _assert_equiv(boundaries, rows, rtt, eos)
+        except AssertionError as e:
+            failures.append((seed, str(e)))
+            if len(failures) >= 5:
+                break
+
+    if failures:
+        joined = "\n\n---\n\n".join(f"seed={s}:\n{m}" for s, m in failures)
+        pytest.fail(
+            f"Long-span differential fuzz: {len(failures)} mismatches in "
+            f"{N_CONFIGS} configs.\n\n{joined}"
+        )
diff --git a/tests/utils/callbacks/test_dynamic_checkpoint.py b/tests/utils/callbacks/test_dynamic_checkpoint.py
new file mode 100644
index 0000000000..1fd7921026
--- /dev/null
+++ b/tests/utils/callbacks/test_dynamic_checkpoint.py
@@ -0,0 +1,389 @@
+"""Unit tests for dynamic checkpoint callback"""
+
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, Mock, patch
+
+from axolotl.utils.callbacks.dynamic_checkpoint import (
+    DEFAULT_TRIGGER_FILENAME,
+    DynamicCheckpointCallback,
+)
+from axolotl.utils.dict import DictDefault
+
+
+class TestDynamicCheckpointCallbackInit:
+    """Test callback initialization"""
+
+    def test_callback_disabled_by_default(self):
+        """Test that callback is disabled when config.enabled=False"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": False},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+            assert callback.enabled is False
+
+    def test_callback_disabled_when_none(self):
+        """Test that callback is disabled when dynamic_checkpoint is None"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": None,
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+            assert callback.enabled is False
+
+    def test_callback_enabled_when_configured(self):
+        """Test that callback is enabled when config.enabled=True"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True, "check_interval": 10},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+            assert callback.enabled is True
+            assert callback.check_interval == 10
+
+    def test_default_trigger_filename(self):
+        """Test that default trigger filename is used"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True, "check_interval": 10},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+            assert callback.trigger_filename == DEFAULT_TRIGGER_FILENAME
+
+    def test_check_interval_default(self):
+        """Test default check interval"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+            assert callback.check_interval == 100  # Default from schema
+
+
+class TestDynamicCheckpointFileDetection:
+    """Test file-based checkpoint triggering"""
+
+    def test_trigger_file_detected_and_deleted(self):
+        """Test that trigger file is detected and deleted"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True, "check_interval": 1},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+
+            trigger_file = Path(tmpdir) / DEFAULT_TRIGGER_FILENAME
+            trigger_file.touch()
+            assert trigger_file.exists()
+
+            args = Mock(output_dir=tmpdir)
+            state = Mock(global_step=1)
+            control = Mock(should_save=False)
+
+            with patch(
+                "axolotl.utils.callbacks.dynamic_checkpoint.is_main_process",
+                return_value=True,
+            ):
+                with patch(
+                    "axolotl.utils.callbacks.dynamic_checkpoint.is_distributed",
+                    return_value=False,
+                ):
+                    result = callback.on_step_end(args, state, control)
+
+            assert not trigger_file.exists()
+            assert result.should_save is True
+
+    def test_check_interval_honored(self):
+        """Test that file is only checked at check_interval steps"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True, "check_interval": 10},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+
+            args = Mock(output_dir=tmpdir)
+            control = Mock(should_save=False)
+
+            trigger_file = Path(tmpdir) / DEFAULT_TRIGGER_FILENAME
+            trigger_file.touch()
+
+            with patch(
+                "axolotl.utils.callbacks.dynamic_checkpoint.is_main_process",
+                return_value=True,
+            ):
+                with patch(
+                    "axolotl.utils.callbacks.dynamic_checkpoint.is_distributed",
+                    return_value=False,
+                ):
+                    # Step 5 - shouldn't check (not divisible by 10)
+                    state = Mock(global_step=5)
+                    result = callback.on_step_end(args, state, control)
+                    assert trigger_file.exists()  # Still there
+                    assert result.should_save is False
+
+                    # Step 10 - should check
+                    state = Mock(global_step=10)
+                    result = callback.on_step_end(args, state, control)
+                    assert not trigger_file.exists()  # Deleted
+                    assert result.should_save is True
+
+    def test_no_file_no_trigger(self):
+        """Test that no trigger occurs when file doesn't exist"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True, "check_interval": 1},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+
+            args = Mock(output_dir=tmpdir)
+            state = Mock(global_step=1)
+            control = Mock(should_save=False)
+
+            with patch(
+                "axolotl.utils.callbacks.dynamic_checkpoint.is_main_process",
+                return_value=True,
+            ):
+                with patch(
+                    "axolotl.utils.callbacks.dynamic_checkpoint.is_distributed",
+                    return_value=False,
+                ):
+                    result = callback.on_step_end(args, state, control)
+
+            assert result.should_save is False
+
+    def test_file_deletion_error_handling(self):
+        """Test that file deletion errors are handled gracefully"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True, "check_interval": 1},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+
+            trigger_file = Path(tmpdir) / DEFAULT_TRIGGER_FILENAME
+            trigger_file.touch()
+
+            args = Mock(output_dir=tmpdir)
+            state = Mock(global_step=1)
+            control = Mock(should_save=False)
+
+            with patch(
+                "axolotl.utils.callbacks.dynamic_checkpoint.is_main_process",
+                return_value=True,
+            ):
+                with patch(
+                    "axolotl.utils.callbacks.dynamic_checkpoint.is_distributed",
+                    return_value=False,
+                ):
+                    with patch.object(
+                        Path, "unlink", side_effect=OSError("Permission denied")
+                    ):
+                        result = callback.on_step_end(args, state, control)
+
+            assert result.should_save is True
+
+
+class TestDynamicCheckpointMultiGPU:
+    """Test multi-GPU synchronization"""
+
+    def test_only_rank_0_checks_file(self):
+        """Test that only rank 0 checks filesystem in multi-GPU setup"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True, "check_interval": 1},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+
+            trigger_file = Path(tmpdir) / DEFAULT_TRIGGER_FILENAME
+            trigger_file.touch()
+
+            args = Mock(output_dir=tmpdir)
+            state = Mock(global_step=1)
+            control = Mock(should_save=False)
+
+            # Rank 1 (not main process) - shouldn't check file
+            with patch(
+                "axolotl.utils.callbacks.dynamic_checkpoint.is_main_process",
+                return_value=False,
+            ):
+                with patch(
+                    "axolotl.utils.callbacks.dynamic_checkpoint.is_distributed",
+                    return_value=True,
+                ):
+                    with patch("torch.distributed.broadcast") as mock_broadcast:
+                        with patch(
+                            "axolotl.utils.callbacks.dynamic_checkpoint.barrier"
+                        ):
+                            mock_tensor = MagicMock()
+                            mock_tensor.item.return_value = 0
+                            with patch("torch.tensor", return_value=mock_tensor):
+                                callback.on_step_end(args, state, control)
+
+            assert trigger_file.exists()
+            # Broadcast should have been called
+            assert mock_broadcast.called
+
+    def test_broadcast_synchronization(self):
+        """Test that trigger decision is broadcasted to all ranks"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": True, "check_interval": 1},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+
+            trigger_file = Path(tmpdir) / DEFAULT_TRIGGER_FILENAME
+            trigger_file.touch()
+
+            args = Mock(output_dir=tmpdir)
+            state = Mock(global_step=1)
+            control = Mock(should_save=False)
+
+            # Rank 0 detects file
+            with patch(
+                "axolotl.utils.callbacks.dynamic_checkpoint.is_main_process",
+                return_value=True,
+            ):
+                with patch(
+                    "axolotl.utils.callbacks.dynamic_checkpoint.is_distributed",
+                    return_value=True,
+                ):
+                    with patch("torch.distributed.broadcast") as mock_broadcast:
+                        with patch(
+                            "axolotl.utils.callbacks.dynamic_checkpoint.barrier"
+                        ) as mock_barrier:
+                            mock_tensor = MagicMock()
+                            mock_tensor.item.return_value = 1
+                            with patch("torch.tensor", return_value=mock_tensor):
+                                with patch("torch.cuda.current_device", return_value=0):
+                                    result = callback.on_step_end(args, state, control)
+
+            assert mock_broadcast.called
+            assert mock_barrier.called
+            # All ranks should trigger
+            assert result.should_save is True
+
+
+class TestDynamicCheckpointSignalHandling:
+    """Test signal-based checkpoint triggering"""
+
+    def test_signal_trigger_via_callback(self):
+        """Test that signal flag triggers checkpoint save"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {
+                        "enabled": True,
+                        "check_interval": 1,
+                        "enable_signal": True,
+                    },
+                    "output_dir": tmpdir,
+                }
+            )
+
+            with patch("signal.signal"):
+                with patch(
+                    "axolotl.utils.callbacks.dynamic_checkpoint.is_main_process",
+                    return_value=True,
+                ):
+                    with patch(
+                        "axolotl.utils.callbacks.dynamic_checkpoint.hasattr",
+                        return_value=True,
+                    ):
+                        callback = DynamicCheckpointCallback(cfg)
+
+            callback.should_save_checkpoint = True
+
+            args = Mock(output_dir=tmpdir)
+            state = Mock(global_step=1)
+            control = Mock(should_save=False)
+
+            with patch(
+                "axolotl.utils.callbacks.dynamic_checkpoint.is_main_process",
+                return_value=True,
+            ):
+                with patch(
+                    "axolotl.utils.callbacks.dynamic_checkpoint.is_distributed",
+                    return_value=False,
+                ):
+                    result = callback.on_step_end(args, state, control)
+
+            assert result.should_save is True
+            assert callback.should_save_checkpoint is False
+
+    def test_signal_not_registered_when_disabled(self):
+        """Test that signal handler is not registered when disabled"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {
+                        "enabled": True,
+                        "check_interval": 10,
+                        "enable_signal": False,
+                    },
+                    "output_dir": tmpdir,
+                }
+            )
+
+            with patch("signal.signal") as mock_signal_register:
+                _ = DynamicCheckpointCallback(cfg)
+
+            assert not mock_signal_register.called
+
+
+class TestDynamicCheckpointDisabled:
+    """Test behavior when callback is disabled"""
+
+    def test_disabled_callback_does_nothing(self):
+        """Test that disabled callback doesn't check or trigger"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = DictDefault(
+                {
+                    "dynamic_checkpoint": {"enabled": False},
+                    "output_dir": tmpdir,
+                }
+            )
+            callback = DynamicCheckpointCallback(cfg)
+
+            trigger_file = Path(tmpdir) / DEFAULT_TRIGGER_FILENAME
+            trigger_file.touch()
+
+            args = Mock(output_dir=tmpdir)
+            state = Mock(global_step=1)
+            control = Mock(should_save=False)
+
+            result = callback.on_step_end(args, state, control)
+
+            assert trigger_file.exists()
+            assert result.should_save is False
diff --git a/tests/utils/callbacks/test_gc_callback.py b/tests/utils/callbacks/test_gc_callback.py
new file mode 100644
index 0000000000..c03e171352
--- /dev/null
+++ b/tests/utils/callbacks/test_gc_callback.py
@@ -0,0 +1,90 @@
+"""Tests for the GCCallback"""
+
+from unittest.mock import MagicMock, patch
+
+from axolotl.utils.callbacks import GCCallback
+
+
+class TestGCCallback:
+    """Tests for GCCallback which handles Python gc.collect() during training."""
+
+    def test_init_with_gc_collect_steps(self):
+        cb = GCCallback(gc_collect_steps=10)
+        assert cb.gc_collect_steps == 10
+
+    def test_init_with_negative_gc_collect_steps(self):
+        cb = GCCallback(gc_collect_steps=-1)
+        assert cb.gc_collect_steps == -1
+
+    def test_init_with_legacy_gc_steps(self):
+        cb = GCCallback(gc_steps=5)
+        assert cb.gc_collect_steps == 5
+
+    def test_init_gc_collect_steps_overrides_gc_steps(self):
+        cb = GCCallback(gc_collect_steps=10, gc_steps=5)
+        assert cb.gc_collect_steps == 10
+
+    def test_init_default(self):
+        cb = GCCallback()
+        assert cb.gc_collect_steps == -1
+
+    @patch("axolotl.utils.callbacks.gc.collect")
+    @patch("axolotl.utils.callbacks.torch.cuda.empty_cache")
+    def test_gc_calls_collect_and_empty_cache(self, mock_empty_cache, mock_gc_collect):
+        cb = GCCallback(gc_collect_steps=1)
+        cb._gc()
+        mock_gc_collect.assert_called_once()
+        mock_empty_cache.assert_called_once()
+
+    @patch("axolotl.utils.callbacks.gc.collect")
+    @patch("axolotl.utils.callbacks.torch.cuda.empty_cache")
+    def test_on_step_end_periodic_gc(self, mock_empty_cache, mock_gc_collect):
+        cb = GCCallback(gc_collect_steps=5)
+        args = MagicMock()
+        args.save_strategy = "no"
+        state = MagicMock()
+        state.global_step = 10
+        state.save_steps = 0
+        state.max_steps = 100
+        control = MagicMock()
+        control.should_evaluate = False
+
+        cb.on_step_end(args, state, control)
+
+        # Step 10 % 5 == 0, so GC should have been called
+        mock_gc_collect.assert_called_once()
+
+    @patch("axolotl.utils.callbacks.gc.collect")
+    @patch("axolotl.utils.callbacks.torch.cuda.empty_cache")
+    def test_on_step_end_no_gc_when_not_on_interval(
+        self, mock_empty_cache, mock_gc_collect
+    ):
+        cb = GCCallback(gc_collect_steps=5)
+        args = MagicMock()
+        args.save_strategy = "no"
+        state = MagicMock()
+        state.global_step = 7
+        state.save_steps = 0
+        state.max_steps = 100
+        control = MagicMock()
+        control.should_evaluate = False
+
+        cb.on_step_end(args, state, control)
+
+        # Step 7 % 5 != 0, so GC should not have been called
+        mock_gc_collect.assert_not_called()
+
+    @patch("axolotl.utils.callbacks.gc.collect")
+    @patch("axolotl.utils.callbacks.torch.cuda.empty_cache")
+    def test_on_step_end_gc_before_eval(self, mock_empty_cache, mock_gc_collect):
+        cb = GCCallback(gc_collect_steps=-1)
+        args = MagicMock()
+        state = MagicMock()
+        state.global_step = 3
+        control = MagicMock()
+        control.should_evaluate = True
+
+        cb.on_step_end(args, state, control)
+
+        mock_gc_collect.assert_called_once()
+        assert cb.next_gc_on_begin_step == 4
diff --git a/tests/utils/callbacks/test_skip_eval_on_resume.py b/tests/utils/callbacks/test_skip_eval_on_resume.py
new file mode 100644
index 0000000000..55cf274386
--- /dev/null
+++ b/tests/utils/callbacks/test_skip_eval_on_resume.py
@@ -0,0 +1,63 @@
+"""Tests for SkipEvalOnResumeCallback."""
+
+from unittest.mock import MagicMock
+
+from transformers import TrainerControl, TrainerState, TrainingArguments
+
+from axolotl.utils.callbacks import SkipEvalOnResumeCallback
+
+
+class TestSkipEvalOnResumeCallback:
+    """Tests for skipping redundant evaluation on checkpoint resume."""
+
+    @staticmethod
+    def _make_state(global_step: int) -> TrainerState:
+        state = MagicMock(spec=TrainerState)
+        state.global_step = global_step
+        return state
+
+    def test_suppresses_eval_at_resume_step(self):
+        cb = SkipEvalOnResumeCallback()
+        args = MagicMock(spec=TrainingArguments)
+        state = self._make_state(20)
+        control = TrainerControl(should_evaluate=False)
+
+        # Simulate on_train_begin at checkpoint-20
+        cb.on_train_begin(args, state, control)
+
+        # Trainer sets should_evaluate = True for step 20
+        control.should_evaluate = True
+        result = cb.on_step_end(args, state, control)
+
+        assert result.should_evaluate is False
+
+    def test_allows_eval_after_resume_step(self):
+        cb = SkipEvalOnResumeCallback()
+        args = MagicMock(spec=TrainingArguments)
+        state = self._make_state(20)
+        control = TrainerControl(should_evaluate=False)
+
+        cb.on_train_begin(args, state, control)
+
+        # Advance past the resume point
+        state.global_step = 30
+        control.should_evaluate = True
+        result = cb.on_step_end(args, state, control)
+
+        assert result.should_evaluate is True
+
+    def test_noop_on_fresh_run(self):
+        cb = SkipEvalOnResumeCallback()
+        args = MagicMock(spec=TrainingArguments)
+        state = self._make_state(0)
+        control = TrainerControl(should_evaluate=False)
+
+        # Fresh run: global_step starts at 0
+        cb.on_train_begin(args, state, control)
+
+        # Even if eval triggers at step 0 (unlikely but defensive)
+        state.global_step = 10
+        control.should_evaluate = True
+        result = cb.on_step_end(args, state, control)
+
+        assert result.should_evaluate is True
diff --git a/tests/utils/callbacks/test_tokens_per_second.py b/tests/utils/callbacks/test_tokens_per_second.py
new file mode 100644
index 0000000000..895ba3cd8c
--- /dev/null
+++ b/tests/utils/callbacks/test_tokens_per_second.py
@@ -0,0 +1,50 @@
+"""Tests for trainable-token throughput accounting."""
+
+from axolotl.core.trainers.utils import trainable_tokens_per_sec_per_gpu
+
+
+def _cumulative_after_window(microbatch_token_counts, world_size, start=0.0):
+    """Mimic the trainer's cumulative counter: every microbatch is SUM-reduced
+    across ranks (balanced => x world_size) then added to the running total."""
+    cum = start
+    for tok in microbatch_token_counts:
+        cum += tok * world_size
+    return cum
+
+
+class TestTrainableTokensPerSec:
+    """Throughput from cumulative-counter deltas (regression for GA undercount)."""
+
+    def test_basic_rate(self):
+        # 800 trainable tokens over 10s on a single GPU
+        assert trainable_tokens_per_sec_per_gpu(1000.0, 1800.0, 1, 10.0) == 80.0
+
+    def test_world_size_divides_to_per_gpu(self):
+        # 1600 tokens summed across 2 ranks over 10s => 80 tok/s/gpu
+        assert trainable_tokens_per_sec_per_gpu(0.0, 1600.0, 2, 10.0) == 80.0
+
+    def test_first_window_returns_none(self):
+        assert trainable_tokens_per_sec_per_gpu(None, 1800.0, 1, 10.0) is None
+
+    def test_resume_first_window_does_not_spike(self):
+        # after resume the counter is restored to a large value but there is no
+        # prior window yet, so the first post-resume log must not emit a rate
+        assert trainable_tokens_per_sec_per_gpu(None, 5_000_000.0, 1, 10.0) is None
+
+    def test_non_positive_elapsed_returns_none(self):
+        assert trainable_tokens_per_sec_per_gpu(1000.0, 1800.0, 1, 0.0) is None
+        assert trainable_tokens_per_sec_per_gpu(1000.0, 1800.0, 1, -5.0) is None
+
+    def test_independent_of_gradient_accumulation(self):
+        # same tokens and wall time, processed as 1 microbatch vs 8 microbatches
+        ga1 = _cumulative_after_window([800], world_size=1)
+        ga8 = _cumulative_after_window([100] * 8, world_size=1)
+        rate1 = trainable_tokens_per_sec_per_gpu(0.0, ga1, 1, 10.0)
+        rate8 = trainable_tokens_per_sec_per_gpu(0.0, ga8, 1, 10.0)
+        assert rate1 == rate8 == 80.0
+
+    def test_overhead_lowers_rate(self):
+        # same tokens, larger wall time (e.g. eval/checkpoint in window) => lower rate
+        fast = trainable_tokens_per_sec_per_gpu(0.0, 800.0, 1, 10.0)
+        slow = trainable_tokens_per_sec_per_gpu(0.0, 800.0, 1, 20.0)
+        assert slow < fast
diff --git a/tests/utils/data/test_hash.py b/tests/utils/data/test_hash.py
new file mode 100644
index 0000000000..4108a55788
--- /dev/null
+++ b/tests/utils/data/test_hash.py
@@ -0,0 +1,234 @@
+"""
+Tests for generate_dataset_hash_from_config.
+
+Regression test for https://github.com/axolotl-ai-cloud/axolotl/issues/3303:
+changing output_dir should not bust the dataset cache when added_tokens_overrides
+is set.
+"""
+
+from axolotl.utils.data.shared import generate_dataset_hash_from_config
+from axolotl.utils.dict import DictDefault
+
+
+def _base_cfg(**kwargs):
+    return DictDefault(
+        {
+            "sequence_len": 2048,
+            "sample_packing": False,
+            "eval_sample_packing": False,
+            "group_by_length": False,
+            "kd_temperature": None,
+            "dataset_exact_deduplication": False,
+            "tokenizer_config": "NousResearch/Llama-3.2-1B",
+            **kwargs,
+        }
+    )
+
+
+def _datasets():
+    return [
+        DictDefault(
+            {
+                "path": "mhenrichsen/alpaca_2k_test",
+                "type": "alpaca",
+                "shards": None,
+                "conversation": None,
+                "split": "train",
+                "temperature": None,
+            }
+        )
+    ]
+
+
+class TestGenerateDatasetHashFromConfig:
+    def test_same_config_same_hash(self):
+        """Identical configs produce identical hashes."""
+        cfg = _base_cfg()
+        h1 = generate_dataset_hash_from_config(
+            cfg, _datasets(), "NousResearch/Llama-3.2-1B"
+        )
+        h2 = generate_dataset_hash_from_config(
+            cfg, _datasets(), "NousResearch/Llama-3.2-1B"
+        )
+        assert h1 == h2
+
+    def test_different_tokenizer_different_hash(self):
+        """A different tokenizer path produces a different hash."""
+        cfg = _base_cfg()
+        h1 = generate_dataset_hash_from_config(
+            cfg, _datasets(), "NousResearch/Llama-3.2-1B"
+        )
+        h2 = generate_dataset_hash_from_config(
+            cfg, _datasets(), "HuggingFaceTB/SmolLM2-135M"
+        )
+        assert h1 != h2
+
+    def test_different_sequence_len_different_hash(self):
+        cfg_a = _base_cfg(sequence_len=2048)
+        cfg_b = _base_cfg(sequence_len=4096)
+        h1 = generate_dataset_hash_from_config(cfg_a, _datasets(), "tok")
+        h2 = generate_dataset_hash_from_config(cfg_b, _datasets(), "tok")
+        assert h1 != h2
+
+    def test_different_special_tokens_different_hash(self):
+        cfg_a = _base_cfg(special_tokens={"pad_token": "<|endoftext|>"})
+        cfg_b = _base_cfg(
+            special_tokens={
+                "pad_token": "<|endoftext|>",
+                "bos_token": "<|custom_im_start|>",
+                "eos_token": "<|custom_im_end|>",
+            }
+        )
+
+        h1 = generate_dataset_hash_from_config(cfg_a, _datasets(), "tok")
+        h2 = generate_dataset_hash_from_config(cfg_b, _datasets(), "tok")
+
+        assert h1 != h2
+
+    def test_different_chat_template_different_hash(self):
+        cfg_a = _base_cfg(chat_template="chatml")
+        cfg_b = _base_cfg(chat_template="jinja", chat_template_jinja="{{ messages }}")
+
+        h1 = generate_dataset_hash_from_config(cfg_a, _datasets(), "tok")
+        h2 = generate_dataset_hash_from_config(cfg_b, _datasets(), "tok")
+
+        assert h1 != h2
+
+    def test_dataset_chat_template_fields_affect_hash(self):
+        datasets_a = _datasets()
+        datasets_b = _datasets()
+        datasets_b[0].chat_template = "jinja"
+        datasets_b[0].chat_template_jinja = "{{ messages }}"
+        datasets_b[0].field_messages = "conversations"
+        datasets_b[0].message_property_mappings = {
+            "role": "from",
+            "content": "value",
+        }
+
+        cfg = _base_cfg()
+        h1 = generate_dataset_hash_from_config(cfg, datasets_a, "tok")
+        h2 = generate_dataset_hash_from_config(cfg, datasets_b, "tok")
+
+        assert h1 != h2
+
+    def test_llama_fix_token_configs_do_not_share_cache_hash(self):
+        datasets_custom_tokens = [
+            DictDefault(
+                {
+                    "path": "mlabonne/FineTome-100k",
+                    "type": "chat_template",
+                    "split": "train[:10%]",
+                    "chat_template": "jinja",
+                    "chat_template_jinja": "{{ messages }}",
+                    "field_messages": "conversations",
+                    "message_property_mappings": {
+                        "role": "from",
+                        "content": "value",
+                    },
+                }
+            )
+        ]
+        datasets_chatml = [
+            DictDefault(
+                {
+                    "path": "mlabonne/FineTome-100k",
+                    "type": "chat_template",
+                    "split": "train[:10%]",
+                    "field_messages": "conversations",
+                    "message_property_mappings": {
+                        "role": "from",
+                        "content": "value",
+                    },
+                }
+            )
+        ]
+        cfg_custom_tokens = _base_cfg(
+            sequence_len=512,
+            sample_packing=True,
+            eval_sample_packing=True,
+            special_tokens={
+                "pad_token": "<|endoftext|>",
+                "bos_token": "<|custom_im_start|>",
+                "eos_token": "<|custom_im_end|>",
+            },
+        )
+        cfg_chatml = _base_cfg(
+            sequence_len=512,
+            sample_packing=True,
+            eval_sample_packing=True,
+            special_tokens={"pad_token": "<|endoftext|>"},
+            chat_template="chatml",
+        )
+
+        h1 = generate_dataset_hash_from_config(
+            cfg_custom_tokens, datasets_custom_tokens, "HuggingFaceTB/SmolLM2-135M"
+        )
+        h2 = generate_dataset_hash_from_config(
+            cfg_chatml, datasets_chatml, "HuggingFaceTB/SmolLM2-135M"
+        )
+
+        assert h1 != h2
+
+    # --- Regression: added_tokens_overrides + output_dir ---
+
+    def test_added_tokens_overrides_hash_stable_across_output_dir(self):
+        """Hash must not change when only output_dir changes (issue #3303).
+
+        When added_tokens_overrides is set the tokenizer is saved into output_dir,
+        making tokenizer.name_or_path an absolute path that includes output_dir.
+        The hash should be derived from the canonical tokenizer config + overrides,
+        not from the output-dir-dependent path.
+        """
+        cfg_run1 = _base_cfg(
+            output_dir="/tmp/run_1",
+            added_tokens_overrides={32000: "<PAD>", 32001: "<MASK>"},
+        )
+        cfg_run2 = _base_cfg(
+            output_dir="/tmp/run_2_different_name",
+            added_tokens_overrides={32000: "<PAD>", 32001: "<MASK>"},
+        )
+
+        # Simulate what happens in practice: tokenizer.name_or_path becomes the
+        # output_dir-based path after modify_tokenizer_files() saves the tokenizer.
+        tokenizer_name_run1 = "/tmp/run_1/modified_tokenizer"
+        tokenizer_name_run2 = "/tmp/run_2_different_name/modified_tokenizer"
+
+        h1 = generate_dataset_hash_from_config(
+            cfg_run1, _datasets(), tokenizer_name_run1
+        )
+        h2 = generate_dataset_hash_from_config(
+            cfg_run2, _datasets(), tokenizer_name_run2
+        )
+
+        assert h1 == h2, (
+            "Dataset cache hash must not change when only output_dir changes "
+            "while added_tokens_overrides stays the same (issue #3303)."
+        )
+
+    def test_added_tokens_overrides_different_overrides_different_hash(self):
+        """Different added_tokens_overrides produce different hashes."""
+        cfg_a = _base_cfg(
+            output_dir="/tmp/run_a",
+            added_tokens_overrides={32000: "<PAD>"},
+        )
+        cfg_b = _base_cfg(
+            output_dir="/tmp/run_a",  # same output_dir
+            added_tokens_overrides={32000: "<OTHER>"},
+        )
+        tokenizer_path = "/tmp/run_a/modified_tokenizer"
+
+        h1 = generate_dataset_hash_from_config(cfg_a, _datasets(), tokenizer_path)
+        h2 = generate_dataset_hash_from_config(cfg_b, _datasets(), tokenizer_path)
+
+        assert h1 != h2
+
+    def test_no_added_tokens_overrides_uses_tokenizer_name_as_before(self):
+        """Without added_tokens_overrides the old behaviour is preserved."""
+        cfg = _base_cfg()  # no added_tokens_overrides
+        tokenizer_name = "NousResearch/Llama-3.2-1B"
+
+        h1 = generate_dataset_hash_from_config(cfg, _datasets(), tokenizer_name)
+        # Changing tokenizer_name still changes the hash
+        h2 = generate_dataset_hash_from_config(cfg, _datasets(), "some/other-model")
+
+        assert h1 != h2
diff --git a/tests/utils/data/test_rl.py b/tests/utils/data/test_rl.py
new file mode 100644
index 0000000000..44c6a010dc
--- /dev/null
+++ b/tests/utils/data/test_rl.py
@@ -0,0 +1,292 @@
+"""
+Unit tests for RL data utility functions (excess_length_strategy support).
+"""
+
+import unittest
+
+from axolotl.utils.data.rl import (
+    _drop_long_sequences,
+    _raise_on_long_sequences,
+    _truncate_long_sequences_rl,
+)
+from axolotl.utils.schemas.enums import RLType
+
+
+class _FakeTokenizer:
+    """Simple whitespace tokenizer for testing length calculations."""
+
+    def __call__(self, text, add_special_tokens=True):  # noqa: ARG002
+        tokens = text.split()
+        return {"input_ids": list(range(len(tokens)))}
+
+    def decode(self, token_ids, skip_special_tokens=True):  # noqa: ARG002
+        # Each token id maps to a placeholder word; length is what matters.
+        return " ".join(f"w{i}" for i in range(len(token_ids)))
+
+
+def _make_dpo_sample(prompt_len: int, chosen_len: int, rejected_len: int):
+    """Create a DPO sample with specified word counts."""
+    return {
+        "prompt": " ".join(f"p{i}" for i in range(prompt_len)),
+        "chosen": " ".join(f"c{i}" for i in range(chosen_len)),
+        "rejected": " ".join(f"r{i}" for i in range(rejected_len)),
+    }
+
+
+def _make_kto_sample(prompt_len: int, completion_len: int):
+    """Create a KTO sample with specified word counts."""
+    return {
+        "prompt": " ".join(f"p{i}" for i in range(prompt_len)),
+        "completion": " ".join(f"c{i}" for i in range(completion_len)),
+    }
+
+
+class TestDropLongSequences(unittest.TestCase):
+    """Tests for the existing _drop_long_sequences filter function."""
+
+    def setUp(self):
+        self.tokenizer = _FakeTokenizer()
+
+    def test_dpo_keeps_short_samples(self):
+        sample = _make_dpo_sample(prompt_len=3, chosen_len=2, rejected_len=2)
+        result = _drop_long_sequences(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertTrue(result)
+
+    def test_dpo_drops_long_chosen(self):
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=10, rejected_len=2)
+        result = _drop_long_sequences(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertFalse(result)
+
+    def test_dpo_drops_long_rejected(self):
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=2, rejected_len=10)
+        result = _drop_long_sequences(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertFalse(result)
+
+    def test_kto_keeps_short_samples(self):
+        sample = _make_kto_sample(prompt_len=3, completion_len=2)
+        result = _drop_long_sequences(
+            sample, RLType.KTO, self.tokenizer, sequence_len=10
+        )
+        self.assertTrue(result)
+
+    def test_kto_drops_long_completion(self):
+        sample = _make_kto_sample(prompt_len=5, completion_len=10)
+        result = _drop_long_sequences(
+            sample, RLType.KTO, self.tokenizer, sequence_len=10
+        )
+        self.assertFalse(result)
+
+    def test_grpo_always_keeps(self):
+        sample = {"prompt": "a " * 100}
+        result = _drop_long_sequences(
+            sample, RLType.GRPO, self.tokenizer, sequence_len=5
+        )
+        self.assertTrue(result)
+
+    def test_dpo_missing_keys_raises(self):
+        with self.assertRaises(ValueError):
+            _drop_long_sequences({"prompt": "hi"}, RLType.DPO, self.tokenizer, 10)
+
+    def test_kto_missing_keys_raises(self):
+        with self.assertRaises(ValueError):
+            _drop_long_sequences({"prompt": "hi"}, RLType.KTO, self.tokenizer, 10)
+
+    def test_ipo_uses_dpo_logic(self):
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=10, rejected_len=2)
+        result = _drop_long_sequences(
+            sample, RLType.IPO, self.tokenizer, sequence_len=10
+        )
+        self.assertFalse(result)
+
+    def test_orpo_uses_dpo_logic(self):
+        sample = _make_dpo_sample(prompt_len=3, chosen_len=2, rejected_len=2)
+        result = _drop_long_sequences(
+            sample, RLType.ORPO, self.tokenizer, sequence_len=10
+        )
+        self.assertTrue(result)
+
+    def test_boundary_length_kept(self):
+        """Samples exactly at sequence_len should be kept."""
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=5, rejected_len=5)
+        result = _drop_long_sequences(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertTrue(result)
+
+
+class TestRaiseOnLongSequences(unittest.TestCase):
+    """Tests for _raise_on_long_sequences (excess_length_strategy='raise')."""
+
+    def setUp(self):
+        self.tokenizer = _FakeTokenizer()
+
+    def test_short_sample_passes(self):
+        sample = _make_dpo_sample(prompt_len=3, chosen_len=2, rejected_len=2)
+        result = _raise_on_long_sequences(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertTrue(result)
+
+    def test_long_sample_raises_valueerror(self):
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=10, rejected_len=2)
+        with self.assertRaises(ValueError, msg="excess_length_strategy"):
+            _raise_on_long_sequences(
+                sample, RLType.DPO, self.tokenizer, sequence_len=10
+            )
+
+    def test_kto_long_raises(self):
+        sample = _make_kto_sample(prompt_len=5, completion_len=10)
+        with self.assertRaises(ValueError):
+            _raise_on_long_sequences(
+                sample, RLType.KTO, self.tokenizer, sequence_len=10
+            )
+
+    def test_grpo_never_raises(self):
+        sample = {"prompt": "a " * 100}
+        result = _raise_on_long_sequences(
+            sample, RLType.GRPO, self.tokenizer, sequence_len=5
+        )
+        self.assertTrue(result)
+
+
+class TestTruncateLongSequencesRL(unittest.TestCase):
+    """Tests for _truncate_long_sequences_rl (excess_length_strategy='truncate')."""
+
+    def setUp(self):
+        self.tokenizer = _FakeTokenizer()
+
+    def test_dpo_short_sample_unchanged(self):
+        sample = _make_dpo_sample(prompt_len=3, chosen_len=2, rejected_len=2)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertEqual(result["chosen"], sample["chosen"])
+        self.assertEqual(result["rejected"], sample["rejected"])
+
+    def test_dpo_truncates_chosen(self):
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=10, rejected_len=3)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        # max_response_len = 10 - 5 = 5, chosen had 10 words -> truncated to 5
+        chosen_tokens = self.tokenizer(result["chosen"], add_special_tokens=False)[
+            "input_ids"
+        ]
+        self.assertEqual(len(chosen_tokens), 5)
+
+    def test_dpo_truncates_rejected(self):
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=3, rejected_len=10)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        rejected_tokens = self.tokenizer(result["rejected"], add_special_tokens=False)[
+            "input_ids"
+        ]
+        self.assertEqual(len(rejected_tokens), 5)
+
+    def test_dpo_truncates_both(self):
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=10, rejected_len=10)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        chosen_len = len(
+            self.tokenizer(result["chosen"], add_special_tokens=False)["input_ids"]
+        )
+        rejected_len = len(
+            self.tokenizer(result["rejected"], add_special_tokens=False)["input_ids"]
+        )
+        self.assertEqual(chosen_len, 5)
+        self.assertEqual(rejected_len, 5)
+
+    def test_dpo_prompt_unchanged(self):
+        """Prompt text should never be modified."""
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=10, rejected_len=10)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertEqual(result["prompt"], sample["prompt"])
+
+    def test_dpo_prompt_exceeds_limit_returns_unchanged(self):
+        """When prompt alone exceeds sequence_len, sample is returned as-is."""
+        sample = _make_dpo_sample(prompt_len=15, chosen_len=3, rejected_len=3)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertEqual(result, sample)
+
+    def test_kto_truncates_completion(self):
+        sample = _make_kto_sample(prompt_len=5, completion_len=10)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.KTO, self.tokenizer, sequence_len=10
+        )
+        completion_len = len(
+            self.tokenizer(result["completion"], add_special_tokens=False)["input_ids"]
+        )
+        self.assertEqual(completion_len, 5)
+
+    def test_kto_short_sample_unchanged(self):
+        sample = _make_kto_sample(prompt_len=3, completion_len=2)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.KTO, self.tokenizer, sequence_len=10
+        )
+        self.assertEqual(result["completion"], sample["completion"])
+
+    def test_kto_prompt_exceeds_limit_returns_unchanged(self):
+        sample = _make_kto_sample(prompt_len=15, completion_len=3)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.KTO, self.tokenizer, sequence_len=10
+        )
+        self.assertEqual(result, sample)
+
+    def test_grpo_unchanged(self):
+        sample = {"prompt": "a " * 100}
+        result = _truncate_long_sequences_rl(
+            sample, RLType.GRPO, self.tokenizer, sequence_len=5
+        )
+        self.assertEqual(result, sample)
+
+    def test_ipo_uses_dpo_logic(self):
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=10, rejected_len=3)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.IPO, self.tokenizer, sequence_len=10
+        )
+        chosen_len = len(
+            self.tokenizer(result["chosen"], add_special_tokens=False)["input_ids"]
+        )
+        self.assertEqual(chosen_len, 5)
+
+    def test_does_not_mutate_original(self):
+        """Verify immutability — original sample dict is not modified."""
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=10, rejected_len=10)
+        original_chosen = sample["chosen"]
+        original_rejected = sample["rejected"]
+        _truncate_long_sequences_rl(sample, RLType.DPO, self.tokenizer, sequence_len=10)
+        self.assertEqual(sample["chosen"], original_chosen)
+        self.assertEqual(sample["rejected"], original_rejected)
+
+    def test_dpo_missing_keys_raises(self):
+        with self.assertRaises(ValueError):
+            _truncate_long_sequences_rl(
+                {"prompt": "hi"}, RLType.DPO, self.tokenizer, 10
+            )
+
+    def test_kto_missing_keys_raises(self):
+        with self.assertRaises(ValueError):
+            _truncate_long_sequences_rl(
+                {"prompt": "hi"}, RLType.KTO, self.tokenizer, 10
+            )
+
+    def test_boundary_no_truncation_needed(self):
+        """Samples exactly at sequence_len should not be modified."""
+        sample = _make_dpo_sample(prompt_len=5, chosen_len=5, rejected_len=5)
+        result = _truncate_long_sequences_rl(
+            sample, RLType.DPO, self.tokenizer, sequence_len=10
+        )
+        self.assertEqual(result["chosen"], sample["chosen"])
+        self.assertEqual(result["rejected"], sample["rejected"])
diff --git a/tests/utils/data/test_utils.py b/tests/utils/data/test_utils.py
new file mode 100644
index 0000000000..7c3e4ecac7
--- /dev/null
+++ b/tests/utils/data/test_utils.py
@@ -0,0 +1,628 @@
+"""
+Unit tests for data utility functions
+"""
+
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from datasets import Dataset
+
+from axolotl.utils.data.shared import _load_from_local_path
+from axolotl.utils.data.utils import handle_long_seq_in_dataset, remove_double_bos_token
+from axolotl.utils.dict import DictDefault
+
+
+class TestHandleLongSeqInDataset(unittest.TestCase):
+    """
+    Test class for handle_long_seq_in_dataset function
+    """
+
+    def test_drop_strategy_removes_long_sequences(self):
+        """Test that 'drop' strategy removes sequences longer than sequence_len"""
+        # Create dataset with mixed length sequences
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3],  # length 3 - keep
+                    [1, 2, 3, 4, 5],  # length 5 - keep
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],  # length 11 - drop
+                    [1, 2],  # length 2 - keep
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "drop",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Should have dropped the sequence with length 11
+        self.assertEqual(len(result), 3)
+        self.assertEqual(len(result[0]["input_ids"]), 3)
+        self.assertEqual(len(result[1]["input_ids"]), 5)
+        self.assertEqual(len(result[2]["input_ids"]), 2)
+
+    def test_drop_strategy_is_default(self):
+        """Test that 'drop' is the default strategy when not specified"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3],
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],  # length 11 - should drop
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Should have dropped the long sequence
+        self.assertEqual(len(result), 1)
+
+    def test_truncate_strategy_truncates_long_sequences(self):
+        """Test that 'truncate' strategy truncates sequences to sequence_len"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3],  # length 3 - keep as is
+                    [
+                        1,
+                        2,
+                        3,
+                        4,
+                        5,
+                        6,
+                        7,
+                        8,
+                        9,
+                        10,
+                        11,
+                        12,
+                    ],  # length 12 - truncate to 10
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "truncate",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Should have 2 samples
+        self.assertEqual(len(result), 2)
+        # First sample unchanged
+        self.assertEqual(len(result[0]["input_ids"]), 3)
+        # Second sample truncated to 10
+        self.assertEqual(len(result[1]["input_ids"]), 10)
+        self.assertEqual(result[1]["input_ids"], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+
+    def test_truncate_strategy_truncates_all_auxiliary_fields(self):
+        """Test that truncation applies to all auxiliary fields consistently"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+                ],
+                "attention_mask": [
+                    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+                ],
+                "labels": [
+                    [-100, -100, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+                ],
+                "position_ids": [
+                    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
+                ],
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "truncate",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # All fields should be truncated to 10
+        self.assertEqual(len(result[0]["input_ids"]), 10)
+        self.assertEqual(len(result[0]["attention_mask"]), 10)
+        self.assertEqual(len(result[0]["labels"]), 10)
+        self.assertEqual(len(result[0]["position_ids"]), 10)
+
+        # Verify content is correct
+        self.assertEqual(result[0]["input_ids"], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+        self.assertEqual(result[0]["attention_mask"], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
+        self.assertEqual(result[0]["labels"], [-100, -100, 3, 4, 5, 6, 7, 8, 9, 10])
+        self.assertEqual(result[0]["position_ids"], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+
+    def test_raise_strategy_raises_on_long_sequences(self):
+        """Test that 'raise' strategy raises ValueError when encountering long sequences"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3],
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],  # length 11 - should raise
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "raise",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        with self.assertRaises(ValueError):
+            handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+    def test_min_sequence_len_filters_short_sequences(self):
+        """Test that sequences shorter than min_sample_len are filtered out"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1],  # length 1 - drop (< min_sample_len=3)
+                    [1, 2],  # length 2 - drop
+                    [1, 2, 3],  # length 3 - keep
+                    [1, 2, 3, 4, 5],  # length 5 - keep
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "drop",
+                "min_sample_len": 3,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Should only keep sequences with length >= 3
+        self.assertEqual(len(result), 2)
+        self.assertEqual(len(result[0]["input_ids"]), 3)
+        self.assertEqual(len(result[1]["input_ids"]), 5)
+
+    def test_dataset_without_input_ids_column(self):
+        """Test that datasets without 'input_ids' column are returned unchanged"""
+        dataset = Dataset.from_dict(
+            {
+                "chosen": [1, 2, 3],
+                "rejected": [4, 5, 6],
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "drop",
+                "min_sample_len": 2,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Dataset should be unchanged
+        self.assertEqual(len(result), len(dataset))
+        self.assertListEqual(list(result.column_names), ["chosen", "rejected"])
+
+    def test_truncate_filters_short_before_truncating(self):
+        """Test that truncate strategy filters short sequences before truncating long ones
+
+        This is important for efficiency - we should not waste time truncating
+        sequences that will be filtered out anyway.
+        """
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1],  # length 1 - filter out first
+                    [1, 2, 3],  # length 3 - keep, no truncation needed
+                    [
+                        1,
+                        2,
+                        3,
+                        4,
+                        5,
+                        6,
+                        7,
+                        8,
+                        9,
+                        10,
+                        11,
+                        12,
+                    ],  # length 12 - keep and truncate
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "truncate",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Should have filtered out the first (short) sequence
+        self.assertEqual(len(result), 2)
+        # Second sample unchanged
+        self.assertEqual(len(result[0]["input_ids"]), 3)
+        # Third sample truncated to 10
+        self.assertEqual(len(result[1]["input_ids"]), 10)
+
+    def test_case_insensitive_strategy(self):
+        """Test that excess_length_strategy is case-insensitive"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "TRUNCATE",  # uppercase
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Should still truncate
+        self.assertEqual(len(result[0]["input_ids"]), 10)
+
+    def test_raise_strategy_silently_drops_short_sequences(self):
+        """Test that 'raise' strategy drops short sequences without raising"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1],  # length 1 - too short, should be dropped silently
+                    [1, 2, 3, 4, 5],  # length 5 - keep
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "raise",
+                "min_sample_len": 3,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        # Should NOT raise, just silently drop the short sequence
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        self.assertEqual(len(result), 1)
+        self.assertEqual(len(result[0]["input_ids"]), 5)
+
+    def test_drop_boundary_sequence_equal_to_sequence_len(self):
+        """Test that drop strategy keeps sequences with length exactly equal to sequence_len"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],  # length 10 == sequence_len
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],  # length 11 > sequence_len
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "drop",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Exactly equal should be kept, one over should be dropped
+        self.assertEqual(len(result), 1)
+        self.assertEqual(len(result[0]["input_ids"]), 10)
+
+    def test_truncate_boundary_sequence_equal_to_sequence_len(self):
+        """Test that truncate strategy leaves sequences with length exactly equal to sequence_len unchanged"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],  # length 10 == sequence_len
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "truncate",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Should be unchanged - not truncated
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0]["input_ids"], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+
+    def test_empty_dataset(self):
+        """Test that an empty dataset is handled gracefully"""
+        dataset = Dataset.from_dict({"input_ids": []})
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "drop",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        self.assertEqual(len(result), 0)
+
+    def test_all_sequences_dropped_returns_empty_dataset(self):
+        """Test that dropping all sequences results in an empty dataset"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1],  # too short
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],  # too long
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "drop",
+                "min_sample_len": 5,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        self.assertEqual(len(result), 0)
+
+    def test_iterable_dataset_skips_processing(self):
+        """Test that streaming datasets (column_names is None) are returned unchanged.
+
+        The skip check in _should_skip_processing triggers when column_names is
+        None, which happens with true streaming datasets loaded via
+        load_dataset(..., streaming=True).
+        """
+        mock_dataset = MagicMock()
+        mock_dataset.column_names = None
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "drop",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(mock_dataset, sequence_len=10, cfg=cfg)
+
+        # Should be returned unchanged (same object)
+        self.assertIs(result, mock_dataset)
+
+    def test_truncate_with_partial_auxiliary_fields(self):
+        """Test truncation when only some auxiliary fields are present"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+                ],
+                "labels": [
+                    [-100, -100, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+                ],
+                # No attention_mask or position_ids
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "truncate",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        self.assertEqual(len(result[0]["input_ids"]), 10)
+        self.assertEqual(len(result[0]["labels"]), 10)
+        self.assertEqual(result[0]["input_ids"], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+        self.assertEqual(result[0]["labels"], [-100, -100, 3, 4, 5, 6, 7, 8, 9, 10])
+        # Confirm no extra columns were introduced
+        self.assertListEqual(sorted(result.column_names), ["input_ids", "labels"])
+
+    def test_min_sample_len_defaults_to_two_when_not_set(self):
+        """Test that min_sample_len defaults to 2 when not specified in config"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1],  # length 1 - should be dropped (< default 2)
+                    [1, 2],  # length 2 - should be kept (>= default 2)
+                    [1, 2, 3],  # length 3 - should be kept
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "drop",
+                # min_sample_len not set
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        self.assertEqual(len(result), 2)
+        self.assertEqual(len(result[0]["input_ids"]), 2)
+        self.assertEqual(len(result[1]["input_ids"]), 3)
+
+    def test_invalid_strategy_falls_through_to_drop(self):
+        """Test that an unrecognized strategy value falls through to drop behavior"""
+        dataset = Dataset.from_dict(
+            {
+                "input_ids": [
+                    [1, 2, 3],  # keep
+                    [
+                        1,
+                        2,
+                        3,
+                        4,
+                        5,
+                        6,
+                        7,
+                        8,
+                        9,
+                        10,
+                        11,
+                    ],  # length 11 - should be dropped
+                ]
+            }
+        )
+
+        cfg = DictDefault(
+            {
+                "excess_length_strategy": "not_a_real_strategy",
+                "min_sample_len": 2,
+                "dataset_num_proc": None,
+                "is_preprocess": False,
+            }
+        )
+
+        result = handle_long_seq_in_dataset(dataset, sequence_len=10, cfg=cfg)
+
+        # Should behave like 'drop'
+        self.assertEqual(len(result), 1)
+        self.assertEqual(len(result[0]["input_ids"]), 3)
+
+
+class TestLocalFileSplitHandling(unittest.TestCase):
+    """Regression tests for local-file dataset split handling."""
+
+    def setUp(self):
+        self.tmpdir = Path(tempfile.gettempdir())
+
+    @patch("axolotl.utils.data.shared.load_dataset")
+    def test_local_jsonl_preserves_split_slice(self, mock_load_dataset):
+        path = str(self.tmpdir / "sample.jsonl")
+        dataset_config = DictDefault(
+            {"path": path, "ds_type": "json", "split": "train[:500]"}
+        )
+        load_dataset_kwargs = {"split": dataset_config.split}
+
+        with patch("pathlib.Path.is_file", return_value=True):
+            _load_from_local_path(dataset_config, load_dataset_kwargs)
+
+        mock_load_dataset.assert_called_once()
+        _, kwargs = mock_load_dataset.call_args
+        self.assertEqual(kwargs["split"], "train[:500]")
+        self.assertEqual(kwargs["data_files"], path)
+
+    @patch("axolotl.utils.data.shared.load_dataset")
+    def test_local_csv_preserves_split_slice(self, mock_load_dataset):
+        path = str(self.tmpdir / "sample.csv")
+        dataset_config = DictDefault(
+            {"path": path, "ds_type": "csv", "split": "train[:200]"}
+        )
+        load_dataset_kwargs = {"split": dataset_config.split}
+
+        with patch("pathlib.Path.is_file", return_value=True):
+            _load_from_local_path(dataset_config, load_dataset_kwargs)
+
+        mock_load_dataset.assert_called_once()
+        _, kwargs = mock_load_dataset.call_args
+        self.assertEqual(kwargs["split"], "train[:200]")
+        self.assertEqual(kwargs["data_files"], path)
+
+    @patch("axolotl.utils.data.shared.load_dataset")
+    def test_local_file_no_split_defaults_to_train(self, mock_load_dataset):
+        path = str(self.tmpdir / "sample.parquet")
+        dataset_config = DictDefault({"path": path, "split": None})
+        load_dataset_kwargs = {"split": None}
+
+        with patch("pathlib.Path.is_file", return_value=True):
+            _load_from_local_path(dataset_config, load_dataset_kwargs)
+
+        mock_load_dataset.assert_called_once()
+        _, kwargs = mock_load_dataset.call_args
+        self.assertEqual(kwargs["split"], "train")
+
+
+class TestRemoveDoubleBOSToken(unittest.TestCase):
+    def test_no_remove_bos_token(self):
+        input_ids = [0, 1, 2]
+        labels = [1, 2, 3]
+
+        example = {
+            "input_ids": input_ids,
+            "labels": labels,
+        }
+
+        example = remove_double_bos_token(example, 0)
+        assert example["input_ids"] == input_ids
+        assert example["labels"] == labels
+
+    def test_remove_bos_token(self):
+        input_ids = [0, 0, 1]
+        labels = [0, 1, 2]
+
+        example = {
+            "input_ids": input_ids,
+            "labels": labels,
+        }
+
+        example = remove_double_bos_token(example, 0)
+        assert example["input_ids"] == [0, 1]
+        assert example["labels"] == [1, 2]
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/utils/lora/test_config_validation_lora.py b/tests/utils/lora/test_config_validation_lora.py
new file mode 100644
index 0000000000..396ef74665
--- /dev/null
+++ b/tests/utils/lora/test_config_validation_lora.py
@@ -0,0 +1,291 @@
+import pytest
+import torch.nn as nn
+from pydantic import ValidationError
+
+from axolotl.loaders.adapter import load_lora
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.schemas.peft import LoraConfig as AxLoraConfig
+
+
+class TestLoRAConfigValidation:
+    """Test suite for LoRA/QLoRA configuration validation"""
+
+    def test_basic_configuration_validation(self):
+        """Test basic LoRA configuration validation"""
+
+        valid_config = DictDefault(
+            {
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.1,
+                "lora_target_modules": ["q_proj", "v_proj"],
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+
+        result = validate_config(valid_config)
+        assert result["adapter"] == "lora"
+
+        # DoRA is now compatible with lora kernels
+        dora_kernel_config = DictDefault(
+            {
+                "adapter": "lora",
+                "lora_mlp_kernel": True,
+                "peft_use_dora": True,
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+        result = validate_config(dora_kernel_config)
+        assert result["lora_mlp_kernel"] is True
+        assert result["peft_use_dora"] is True
+
+    def test_lora_rank_alpha_pattern_passthrough(self):
+        """Test that lora_rank_pattern / lora_alpha_pattern survive validation."""
+        rank_pattern = {r".*\.visual\.blocks\.\d+\.(attn|mlp)\..*": 128}
+        alpha_pattern = {r".*\.visual\.blocks\.\d+\.(attn|mlp)\..*": 128}
+
+        cfg = DictDefault(
+            {
+                "adapter": "lora",
+                "lora_r": 64,
+                "lora_alpha": 64,
+                "lora_target_modules": ["q_proj", "v_proj"],
+                "lora_rank_pattern": rank_pattern,
+                "lora_alpha_pattern": alpha_pattern,
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+
+        result = validate_config(cfg)
+        assert result["lora_rank_pattern"] == rank_pattern
+        assert result["lora_alpha_pattern"] == alpha_pattern
+
+    def test_lora_pattern_empty_dict_validates(self):
+        """Empty pattern dicts are accepted and are dropped at the adapter wiring layer."""
+        cfg = DictDefault(
+            {
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_target_modules": ["q_proj"],
+                "lora_rank_pattern": {},
+                "lora_alpha_pattern": {},
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+        result = validate_config(cfg)
+        assert result["lora_rank_pattern"] == {}
+        assert result["lora_alpha_pattern"] == {}
+
+    def test_lora_pattern_rejects_non_int_values(self):
+        """Pydantic should reject non-int values in rank/alpha pattern dicts."""
+        with pytest.raises(ValidationError):
+            AxLoraConfig(
+                adapter="lora",
+                lora_r=8,
+                lora_alpha=16,
+                lora_rank_pattern={r".*\.q_proj$": "not-an-int"},
+            )
+        with pytest.raises(ValidationError):
+            AxLoraConfig(
+                adapter="lora",
+                lora_r=8,
+                lora_alpha=16,
+                lora_alpha_pattern={r".*\.q_proj$": 1.5},
+            )
+
+    @pytest.mark.parametrize("bad_value", [0, -1, -42])
+    def test_lora_pattern_rejects_non_positive_values(self, bad_value):
+        """Pattern values must be > 0; alpha=0 would silently zero-out a module."""
+        with pytest.raises(ValidationError):
+            AxLoraConfig(
+                adapter="lora",
+                lora_r=8,
+                lora_alpha=16,
+                lora_rank_pattern={r".*\.q_proj$": bad_value},
+            )
+        with pytest.raises(ValidationError):
+            AxLoraConfig(
+                adapter="lora",
+                lora_r=8,
+                lora_alpha=16,
+                lora_alpha_pattern={r".*\.q_proj$": bad_value},
+            )
+
+    def test_lora_patterns_reach_peft_lora_config(self):
+        """load_lora(config_only=True) propagates patterns to the constructed PEFT LoraConfig."""
+        rank_pattern = {r".*\.layers\.0\..*\.(q_proj|v_proj)$": 32}
+        alpha_pattern = {r".*\.layers\.0\..*\.(q_proj|v_proj)$": 64}
+
+        cfg = DictDefault(
+            {
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_target_modules": ["q_proj", "v_proj"],
+                "lora_rank_pattern": rank_pattern,
+                "lora_alpha_pattern": alpha_pattern,
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+        cfg = validate_config(cfg)
+
+        class _DummyModel(nn.Module):
+            pass
+
+        _, peft_config = load_lora(_DummyModel(), cfg, config_only=True)
+        assert peft_config.rank_pattern == rank_pattern
+        assert peft_config.alpha_pattern == alpha_pattern
+
+    def test_lora_patterns_unset_means_empty_in_peft(self):
+        """When patterns are unset they must not appear as None on the PEFT LoraConfig."""
+        cfg = DictDefault(
+            {
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_target_modules": ["q_proj"],
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+        cfg = validate_config(cfg)
+
+        class _DummyModel(nn.Module):
+            pass
+
+        _, peft_config = load_lora(_DummyModel(), cfg, config_only=True)
+        assert peft_config.rank_pattern == {}
+        assert peft_config.alpha_pattern == {}
+
+    def test_qlora_4bit_validation(self):
+        """Test QLoRA 4-bit configuration validation"""
+        valid_config = DictDefault(
+            {
+                "adapter": "qlora",
+                "load_in_4bit": True,
+                "bnb_4bit_compute_dtype": "float16",
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+        result = validate_config(valid_config)
+        assert result["adapter"] == "qlora"
+        assert result["load_in_4bit"] is True
+
+        # Test QLoRA without 4-bit (should fail via PEFT validation)
+        with pytest.raises(ValueError, match=r"Require cfg\.load_in_4bit"):
+            invalid_config = DictDefault(
+                {
+                    "adapter": "qlora",
+                    "load_in_4bit": False,
+                    "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                    "micro_batch_size": 1,
+                    "gradient_accumulation_steps": 1,
+                    "learning_rate": 1e-5,
+                    "base_model": "dummy_model",
+                }
+            )
+            validate_config(invalid_config)
+
+        # Test QLoRA with 8-bit (incompatible)
+        with pytest.raises(ValueError, match="Can't load qlora in 8bit"):
+            invalid_config = DictDefault(
+                {
+                    "adapter": "qlora",
+                    "load_in_8bit": True,
+                    "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                    "micro_batch_size": 1,
+                    "gradient_accumulation_steps": 1,
+                    "learning_rate": 1e-5,
+                    "base_model": "dummy_model",
+                }
+            )
+            validate_config(invalid_config)
+
+    @pytest.mark.parametrize(
+        "kernel_field", ["lora_mlp_kernel", "lora_qkv_kernel", "lora_o_kernel"]
+    )
+    def test_lora_kernels_trust_remote_code_incompatible(self, kernel_field):
+        """Test that lora kernels are incompatible with trust_remote_code"""
+        with pytest.raises(ValueError, match="not compatible with trust_remote_code"):
+            invalid_config = DictDefault(
+                {
+                    "adapter": "lora",
+                    kernel_field: True,
+                    "trust_remote_code": True,
+                    "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                    "micro_batch_size": 1,
+                    "gradient_accumulation_steps": 1,
+                    "learning_rate": 1e-5,
+                    "base_model": "dummy_model",
+                }
+            )
+            validate_config(invalid_config)
+
+    def test_lora_kernels_trust_remote_code_false(self):
+        """Test that lora kernels work when trust_remote_code is false"""
+        # Test with trust_remote_code=False, lora kernels should be allowed
+        valid_config = DictDefault(
+            {
+                "adapter": "lora",
+                "lora_mlp_kernel": True,
+                "lora_qkv_kernel": True,
+                "lora_o_kernel": True,
+                "trust_remote_code": False,
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+        result = validate_config(valid_config)
+        assert result["lora_mlp_kernel"] is True
+        assert result["lora_qkv_kernel"] is True
+        assert result["lora_o_kernel"] is True
+
+        # Test with trust_remote_code=None (unset), kernels should be allowed
+        valid_config = DictDefault(
+            {
+                "adapter": "lora",
+                "lora_qkv_kernel": True,
+                "trust_remote_code": None,
+                "datasets": [{"path": "dummy_dataset", "type": "alpaca"}],
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "learning_rate": 1e-5,
+                "base_model": "dummy_model",
+            }
+        )
+        result = validate_config(valid_config)
+        assert result["lora_qkv_kernel"] is True
+        assert result["trust_remote_code"] is None
diff --git a/tests/utils/lora/test_freeze_lora.py b/tests/utils/lora/test_freeze_lora.py
new file mode 100644
index 0000000000..7c5ec8fb39
--- /dev/null
+++ b/tests/utils/lora/test_freeze_lora.py
@@ -0,0 +1,266 @@
+import importlib.util
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn as nn
+
+from axolotl.kernels.lora import get_lora_parameters
+
+PEFT_AVAILABLE = importlib.util.find_spec("peft") is not None
+
+
+class TestLoRAParameterFreezing:
+    """Test suite for LoRA parameter freezing validation."""
+
+    def setup_method(self):
+        self.dtype = torch.float32
+
+    def create_mock_lora_layer(
+        self, has_adapters=True, adapters_disabled=False, merged=False
+    ):
+        """Create a mock LoRA layer for testing."""
+        mock_layer = Mock()
+
+        base_layer = Mock()
+        base_layer.weight = torch.randn(512, 256, dtype=self.dtype)
+        base_layer.bias = torch.randn(512, dtype=self.dtype)
+
+        if has_adapters:
+            mock_layer.base_layer = base_layer
+            mock_layer.disable_adapters = adapters_disabled
+            mock_layer.merged = merged
+
+            mock_layer.active_adapters = ["default"]
+            mock_layer.lora_A = {"default": Mock()}
+            mock_layer.lora_B = {"default": Mock()}
+            mock_layer.scaling = {"default": 0.1}
+
+            mock_layer.lora_A["default"].weight = torch.randn(16, 256, dtype=self.dtype)
+            mock_layer.lora_B["default"].weight = torch.randn(512, 16, dtype=self.dtype)
+            mock_layer.lora_B["default"].bias = None
+
+            # Required by get_lora_parameters for dropout/DoRA extraction
+            mock_layer.lora_dropout = {}
+            mock_layer.lora_magnitude_vector = None
+        else:
+            mock_layer.weight = base_layer.weight
+            mock_layer.bias = base_layer.bias
+
+        return mock_layer
+
+    def test_parameter_freezing_adapters_disabled(self):
+        """Test that LoRA parameters are None when adapters are disabled."""
+        layer = self.create_mock_lora_layer(has_adapters=True, adapters_disabled=True)
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(layer)
+
+        # Base parameters should be returned
+        assert W is not None
+        assert b is not None
+        # LoRA parameters should be None (frozen)
+        assert A is None
+        assert B is None
+        assert s is None
+
+    def test_parameter_freezing_adapters_merged(self):
+        """Test that LoRA parameters are None when adapters are merged."""
+        layer = self.create_mock_lora_layer(has_adapters=True, merged=True)
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(layer)
+
+        # Base parameters should be returned
+        assert W is not None
+        assert b is not None
+
+        # LoRA parameters should be None (frozen)
+        assert A is None
+        assert B is None
+        assert s is None
+
+    def test_parameter_freezing_no_adapters(self):
+        """Test parameter behavior when no adapters are present."""
+        layer = self.create_mock_lora_layer(has_adapters=False)
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(layer)
+
+        # Base parameters should be returned
+        assert W is not None
+        assert b is not None
+
+        # LoRA parameters should be None (frozen)
+        assert A is None
+        assert B is None
+        assert s is None
+
+    def test_parameter_active_adapters_enabled(self):
+        """Test that LoRA parameters are returned when adapters are active."""
+        layer = self.create_mock_lora_layer(
+            has_adapters=True, adapters_disabled=False, merged=False
+        )
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(layer)
+
+        # All parameters should be returned
+        assert W is not None
+        assert b is not None
+        assert A is not None
+        assert B is not None
+        assert s is not None
+        assert s == 0.1
+
+    def test_parameter_shapes_consistency(self):
+        """Test that parameter shapes are consistent when active."""
+        layer = self.create_mock_lora_layer(
+            has_adapters=True, adapters_disabled=False, merged=False
+        )
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(layer)
+
+        # Check shape consistency
+        assert W.shape == (512, 256)
+        assert b.shape == (512,)
+        assert A.shape == (16, 256)
+        assert B.shape == (512, 16)
+
+    def test_parameter_dtypes_consistency(self):
+        """Test that parameter dtypes are consistent."""
+        layer = self.create_mock_lora_layer(
+            has_adapters=True, adapters_disabled=False, merged=False
+        )
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(layer)
+
+        assert W.dtype == self.dtype
+        assert b.dtype == self.dtype
+        assert A.dtype == self.dtype
+        assert B.dtype == self.dtype
+
+    def test_quantization_state_handling(self):
+        """Test that quantization state is properly handled."""
+        layer = self.create_mock_lora_layer(has_adapters=True)
+
+        quant_state_mock = Mock()
+        layer.base_layer.weight.quant_state = quant_state_mock
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(layer)
+
+        assert quant_state == quant_state_mock
+
+    def test_multiple_adapters_active_adapter_selection(self):
+        """Test that the correct adapter is selected when multiple adapters exist."""
+        layer = self.create_mock_lora_layer(
+            has_adapters=True, adapters_disabled=False, merged=False
+        )
+
+        layer.lora_A["adapter2"] = Mock()
+        layer.lora_B["adapter2"] = Mock()
+        layer.scaling["adapter2"] = 0.2
+
+        layer.lora_A["adapter2"].weight = torch.randn(16, 256, dtype=self.dtype)
+        layer.lora_B["adapter2"].weight = torch.randn(512, 16, dtype=self.dtype)
+
+        layer.active_adapters = ["adapter2"]
+
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(layer)
+
+        assert s == 0.2
+        assert torch.equal(A, layer.lora_A["adapter2"].weight)
+        assert torch.equal(B, layer.lora_B["adapter2"].weight)
+
+
+class TestLoRAParameterFreezingIntegration:
+    """Integration tests for parameter freezing with actual LoRA layers."""
+
+    @pytest.mark.skipif(
+        not PEFT_AVAILABLE, reason="PEFT not available for integration tests"
+    )
+    def test_parameter_freezing_with_real_lora_layer(self):
+        """Test parameter freezing with actual PEFT LoRA layer."""
+        from peft import LoraConfig, get_peft_model
+
+        class SimpleModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(256, 512)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        base_model = SimpleModel()
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=["linear"],
+            lora_dropout=0.1,
+        )
+        model = get_peft_model(base_model, lora_config)
+        lora_layer = model.base_model.model.linear
+        # Test with adapters enabled
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(lora_layer)
+        assert A is not None
+        assert B is not None
+        assert s is not None
+        # Test with adapters disabled
+        model.disable_adapter_layers()
+        W, b, quant_state, A, B, s, *_ = get_lora_parameters(lora_layer)
+        assert A is None
+        assert B is None
+        assert s is None
+
+    @pytest.mark.skipif(
+        not PEFT_AVAILABLE, reason="PEFT not available for integration tests"
+    )
+    def test_parameter_freezing_gradient_behavior(self):
+        """Test that frozen parameters don't receive gradients."""
+        from peft import LoraConfig, get_peft_model
+
+        class SimpleModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(256, 512)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        base_model = SimpleModel()
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=["linear"],
+            lora_dropout=0.1,
+        )
+        model = get_peft_model(base_model, lora_config)
+        x = torch.randn(1, 256)
+        target = torch.randn(1, 512)
+        model.enable_adapter_layers()
+        output = model(x)
+        loss = nn.MSELoss()(output, target)
+        loss.backward()
+        lora_layer = model.base_model.model.linear
+        has_lora_grads = any(
+            param.grad is not None
+            for name, param in lora_layer.named_parameters()
+            if "lora_" in name
+        )
+        assert has_lora_grads, (
+            "LoRA parameters should have gradients when adapters are enabled"
+        )
+        model.zero_grad()
+        model.disable_adapter_layers()
+        output = model(x)
+        loss = nn.MSELoss()(output, target)
+        any_requires_grad = any(param.requires_grad for param in model.parameters())
+        if any_requires_grad:
+            loss.backward()
+        has_lora_grads_disabled = any(
+            param.grad is not None
+            for name, param in lora_layer.named_parameters()
+            if "lora_" in name
+        )
+        assert not has_lora_grads_disabled, (
+            "LoRA parameters should not have gradients when adapters are disabled"
+        )
+        model.zero_grad()
+        del model, base_model, lora_layer, x, target, output, loss
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
diff --git a/tests/utils/lora/test_merge_lora.py b/tests/utils/lora/test_merge_lora.py
new file mode 100644
index 0000000000..7964ef1462
--- /dev/null
+++ b/tests/utils/lora/test_merge_lora.py
@@ -0,0 +1,2214 @@
+import json
+import math
+from unittest.mock import Mock, patch
+
+import pytest
+import safetensors.torch
+import torch
+
+from axolotl.cli.merge_lora import do_merge_lora
+from axolotl.cli.utils.lora_merge import (
+    _build_peft_layer_and_get_delta,
+    _find_param_wrapper_lora,
+    _merge_tensor_with_lora,
+    _resolve_lora_alpha_for_key,
+    find_lora_weights,
+    merge_lora_sharded_efficient,
+)
+from axolotl.utils.dict import DictDefault
+
+
+class TestAdapterMergeUnmerge:
+    """Test suite for LoRA adapter merging/unmerging functionality"""
+
+    def setup_method(self):
+        self.dtype = torch.float32
+        self.device = torch.device("cpu")
+
+    def create_mock_base_model(self, vocab_size=1000, hidden_size=256):
+        """Create a mock base model with linear layers"""
+        mock_model = Mock()
+
+        mock_model.config = Mock()
+        mock_model.config.vocab_size = vocab_size
+        mock_model.config.hidden_size = hidden_size
+
+        mock_model.q_proj = Mock()
+        mock_model.q_proj.weight = torch.randn(
+            hidden_size, hidden_size, dtype=self.dtype
+        )
+        mock_model.q_proj.bias = torch.randn(hidden_size, dtype=self.dtype)
+
+        mock_model.v_proj = Mock()
+        mock_model.v_proj.weight = torch.randn(
+            hidden_size, hidden_size, dtype=self.dtype
+        )
+        mock_model.v_proj.bias = torch.randn(hidden_size, dtype=self.dtype)
+
+        return mock_model
+
+    def create_mock_lora_model(self, base_model, r=8, alpha=16):
+        """Create a mock LoRA model wrapping the base model"""
+        mock_lora_model = Mock()
+        mock_lora_model.base_model = base_model
+
+        mock_lora_model.merge_and_unload = None
+        mock_lora_model.to = Mock(return_value=mock_lora_model)
+
+        mock_lora_model.generation_config = Mock()
+        mock_lora_model.config = Mock()
+
+        self.original_q_weight = base_model.q_proj.weight.clone()
+        self.original_v_weight = base_model.v_proj.weight.clone()
+
+        mock_lora_model.peft_config = {"default": Mock()}
+        mock_lora_model.peft_config["default"].r = r
+        mock_lora_model.peft_config["default"].lora_alpha = alpha
+
+        self.lora_A_q = torch.randn(
+            r, base_model.q_proj.weight.shape[1], dtype=self.dtype
+        )
+        self.lora_B_q = torch.randn(
+            base_model.q_proj.weight.shape[0], r, dtype=self.dtype
+        )
+
+        self.lora_A_v = torch.randn(
+            r, base_model.v_proj.weight.shape[1], dtype=self.dtype
+        )
+        self.lora_B_v = torch.randn(
+            base_model.v_proj.weight.shape[0], r, dtype=self.dtype
+        )
+
+        self.scaling = alpha / r
+
+        def mock_merge_and_unload(progressbar=False):
+            """Simulate the actual merge operation"""
+            # Apply LoRA delta to base weights: W_new = W_base + (B @ A) * scaling
+            delta_q = (self.lora_B_q @ self.lora_A_q) * self.scaling
+            delta_v = (self.lora_B_v @ self.lora_A_v) * self.scaling
+
+            base_model.q_proj.weight = self.original_q_weight + delta_q
+            base_model.v_proj.weight = self.original_v_weight + delta_v
+
+            return base_model
+
+        mock_lora_model.merge_and_unload = mock_merge_and_unload
+        return mock_lora_model
+
+    def test_basic_lora_merge_unmerge_cycle(self):
+        """Test: original_weights -> merge -> unmerge -> should equal original_weights"""
+
+        base_model = self.create_mock_base_model()
+        lora_model = self.create_mock_lora_model(base_model)
+
+        original_q_weight = self.original_q_weight.clone()
+        original_v_weight = self.original_v_weight.clone()
+
+        merged_model = lora_model.merge_and_unload()
+
+        assert not torch.equal(merged_model.q_proj.weight, original_q_weight)
+        assert not torch.equal(merged_model.v_proj.weight, original_v_weight)
+
+        delta_q = (self.lora_B_q @ self.lora_A_q) * self.scaling
+        delta_v = (self.lora_B_v @ self.lora_A_v) * self.scaling
+
+        unmerged_q_weight = merged_model.q_proj.weight - delta_q
+        unmerged_v_weight = merged_model.v_proj.weight - delta_v
+
+        assert torch.allclose(unmerged_q_weight, original_q_weight, atol=1e-6)
+        assert torch.allclose(unmerged_v_weight, original_v_weight, atol=1e-6)
+
+    def test_merge_weight_calculation_accuracy(self):
+        """Test: merged_weight = base_weight + (lora_B @ lora_A * scaling)"""
+        base_model = self.create_mock_base_model()
+        lora_model = self.create_mock_lora_model(base_model, r=16, alpha=32)
+
+        expected_delta_q = (self.lora_B_q @ self.lora_A_q) * self.scaling
+        expected_merged_q = self.original_q_weight + expected_delta_q
+        merged_model = lora_model.merge_and_unload()
+
+        assert torch.allclose(merged_model.q_proj.weight, expected_merged_q, atol=1e-6)
+
+    @patch("axolotl.cli.merge_lora.load_model_and_tokenizer")
+    def test_cli_do_merge_functionality(self, mock_load_model, tmp_path):
+        base_model = self.create_mock_base_model()
+        lora_model = self.create_mock_lora_model(base_model)
+        tokenizer = Mock()
+        processor = None
+
+        mock_load_model.return_value = (lora_model, tokenizer, processor)
+
+        cfg = DictDefault(
+            {
+                "save_safetensors": True,
+                "torch_dtype": torch.float32,
+                "local_rank": 0,
+                "output_dir": str(tmp_path),
+                "merge_method": "legacy",
+            }
+        )
+
+        with (
+            patch("pathlib.Path.mkdir"),
+            patch.object(base_model, "save_pretrained") as mock_save_model,
+            patch.object(tokenizer, "save_pretrained") as mock_save_tokenizer,
+        ):
+            do_merge_lora(cfg=cfg)
+
+        mock_save_model.assert_called_once()
+        mock_save_tokenizer.assert_called_once()
+
+    def test_quantized_model_merge_compatibility(self):
+        """Test 4-bit/8-bit model merging scenarios"""
+        base_model = self.create_mock_base_model()
+
+        # Mock quantized weights
+        base_model.q_proj.weight.quant_state = Mock()
+        base_model.q_proj.weight.quant_state.dtype = torch.uint8
+
+        lora_model = self.create_mock_lora_model(base_model)
+
+        merged_model = lora_model.merge_and_unload()
+        assert merged_model is not None
+
+    @patch.dict("os.environ", {"CUDA_VISIBLE_DEVICES": ""})
+    def test_memory_efficient_merge_with_cpu_offload(self, tmp_path):
+        """Test lora_on_cpu configuration during merge"""
+        cfg = DictDefault(
+            {
+                "lora_on_cpu": True,
+                "save_safetensors": True,
+                "output_dir": str(tmp_path),
+                "local_rank": 0,
+                "merge_method": "legacy",
+            }
+        )
+
+        with patch("axolotl.cli.merge_lora.load_model_and_tokenizer") as mock_load:
+            base_model = self.create_mock_base_model()
+            lora_model = self.create_mock_lora_model(base_model)
+            mock_load.return_value = (lora_model, Mock(), None)
+
+            with patch("pathlib.Path.mkdir"), patch("torch.save"):
+                do_merge_lora(cfg=cfg)
+
+            assert mock_load.called
+
+
+class TestEfficientMerge:
+    """Test suite for memory-efficient shard-by-shard LoRA merge."""
+
+    def _make_adapter(self, tmp_path, r=8, alpha=16, use_dora=False, use_rslora=False):
+        """Create a minimal adapter directory with config + weights."""
+        adapter_dir = tmp_path / "adapter"
+        adapter_dir.mkdir()
+
+        config = {
+            "r": r,
+            "lora_alpha": alpha,
+            "target_modules": ["q_proj", "v_proj"],
+            "task_type": "CAUSAL_LM",
+            "bias": "none",
+            "use_dora": use_dora,
+            "use_rslora": use_rslora,
+        }
+        (adapter_dir / "adapter_config.json").write_text(json.dumps(config))
+        return adapter_dir, config
+
+    def _make_base_model(self, tmp_path, hidden=32):
+        """Create a minimal base model directory with one shard."""
+        model_dir = tmp_path / "base_model"
+        model_dir.mkdir()
+
+        weights = {
+            "model.layers.0.self_attn.q_proj.weight": torch.randn(hidden, hidden),
+            "model.layers.0.self_attn.v_proj.weight": torch.randn(hidden, hidden),
+            "model.embed_tokens.weight": torch.randn(100, hidden),
+        }
+        safetensors.torch.save_file(weights, model_dir / "model.safetensors")
+
+        # Minimal config files
+        (model_dir / "config.json").write_text("{}")
+        return model_dir, weights
+
+    def test_find_lora_weights(self):
+        lora_state = {
+            "base_model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.randn(
+                8, 32
+            ),
+            "base_model.model.layers.0.self_attn.q_proj.lora_B.weight": torch.randn(
+                32, 8
+            ),
+        }
+        a, b = find_lora_weights(lora_state, "layers.0.self_attn.q_proj.weight")
+        assert a is not None and b is not None
+        assert a.shape == (8, 32)
+
+        a, b = find_lora_weights(lora_state, "layers.0.self_attn.v_proj.weight")
+        assert a is None and b is None
+
+    def test_merge_tensor_basic(self):
+        hidden = 32
+        r = 8
+        alpha = 16
+        base = torch.randn(hidden, hidden)
+        lora_a = torch.randn(r, hidden)
+        lora_b = torch.randn(hidden, r)
+        scale = alpha / r
+
+        lora_state = {
+            "base_model.model.layer.q_proj.lora_A.weight": lora_a,
+            "base_model.model.layer.q_proj.lora_B.weight": lora_b,
+        }
+
+        config = {"r": r, "lora_alpha": alpha}
+        merged, was_merged = _merge_tensor_with_lora(
+            base, "layer.q_proj.weight", lora_state, scale, config, "cpu"
+        )
+        assert was_merged
+        expected = base + scale * (lora_b @ lora_a)
+        assert torch.allclose(merged, expected, atol=1e-5)
+
+    def test_merge_tensor_rslora_scale(self):
+        """RSLoRA should use alpha/sqrt(r) as scaling factor."""
+        r = 16
+        alpha = 32
+        standard_scale = alpha / r  # 2.0
+        rslora_scale = alpha / math.sqrt(r)  # 8.0
+
+        assert rslora_scale != standard_scale
+        assert abs(rslora_scale - 8.0) < 1e-6
+
+    def test_sharded_efficient_merge(self, tmp_path):
+        """End-to-end test of shard-by-shard merge."""
+        hidden = 32
+        r = 8
+        alpha = 16
+
+        model_dir, base_weights = self._make_base_model(tmp_path, hidden=hidden)
+        adapter_dir, _ = self._make_adapter(tmp_path, r=r, alpha=alpha)
+
+        # Create LoRA weights
+        lora_state = {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.randn(
+                r, hidden
+            ),
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": torch.randn(
+                hidden, r
+            ),
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_A.weight": torch.randn(
+                r, hidden
+            ),
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_B.weight": torch.randn(
+                hidden, r
+            ),
+        }
+        safetensors.torch.save_file(
+            lora_state, adapter_dir / "adapter_model.safetensors"
+        )
+
+        output_dir = tmp_path / "output"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=output_dir,
+            device="cpu",
+        )
+
+        # Verify output exists and has merged weights
+        merged = safetensors.torch.load_file(output_dir / "model.safetensors")
+        scale = alpha / r
+
+        q_key = "model.layers.0.self_attn.q_proj.weight"
+        expected_q = base_weights[q_key] + scale * (
+            lora_state[f"base_model.model.{q_key[:-7]}.lora_B.weight"]
+            @ lora_state[f"base_model.model.{q_key[:-7]}.lora_A.weight"]
+        )
+        assert torch.allclose(merged[q_key], expected_q, atol=1e-5)
+
+        # Embedding should be unchanged
+        assert torch.equal(
+            merged["model.embed_tokens.weight"],
+            base_weights["model.embed_tokens.weight"],
+        )
+
+    def test_resolve_alpha_for_key_returns_none_without_pattern(self):
+        assert (
+            _resolve_lora_alpha_for_key("model.layers.0.self_attn.q_proj.weight", {})
+            is None
+        )
+        assert (
+            _resolve_lora_alpha_for_key(
+                "model.layers.0.self_attn.q_proj.weight",
+                {"alpha_pattern": {}},
+            )
+            is None
+        )
+
+    def test_resolve_alpha_for_key_matches_peft_suffix_semantics(self):
+        cfg = {"alpha_pattern": {"layers.0.self_attn.q_proj": 64}}
+        assert (
+            _resolve_lora_alpha_for_key("model.layers.0.self_attn.q_proj.weight", cfg)
+            == 64
+        )
+        # No match falls back to None so the caller keeps the global alpha.
+        assert (
+            _resolve_lora_alpha_for_key("model.layers.0.self_attn.v_proj.weight", cfg)
+            is None
+        )
+
+    def test_resolve_alpha_for_key_follows_weight_renamings(self):
+        # Pattern keyed against the runtime name; merge sees the checkpoint name.
+        cfg = {"alpha_pattern": {"model.new.layers.0.q_proj": 64}}
+        renamings = {r"^model\.old\.": "model.new."}
+        assert (
+            _resolve_lora_alpha_for_key(
+                "model.old.layers.0.q_proj.weight", cfg, renamings
+            )
+            == 64
+        )
+        # Without the renamings, the same lookup misses and falls back to global.
+        assert (
+            _resolve_lora_alpha_for_key("model.old.layers.0.q_proj.weight", cfg) is None
+        )
+
+    def test_pattern_no_longer_rejected_as_adalora(self, tmp_path):
+        """Memory-efficient merge accepts rank_pattern/alpha_pattern (plain LoRA, not AdaLoRA)."""
+        hidden = 32
+        r = 8
+        alpha = 16
+
+        model_dir, _ = self._make_base_model(tmp_path, hidden=hidden)
+        adapter_dir, config = self._make_adapter(tmp_path, r=r, alpha=alpha)
+        config["rank_pattern"] = {"layers.0.self_attn.q_proj": r}
+        config["alpha_pattern"] = {"layers.0.self_attn.q_proj": alpha}
+        (adapter_dir / "adapter_config.json").write_text(json.dumps(config))
+
+        lora_state = {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.randn(
+                r, hidden
+            ),
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": torch.randn(
+                hidden, r
+            ),
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_A.weight": torch.randn(
+                r, hidden
+            ),
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_B.weight": torch.randn(
+                hidden, r
+            ),
+        }
+        safetensors.torch.save_file(
+            lora_state, adapter_dir / "adapter_model.safetensors"
+        )
+
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=tmp_path / "output",
+            device="cpu",
+        )
+
+    def test_merge_applies_alpha_pattern_per_module(self, tmp_path):
+        """End-to-end: q_proj uses alpha_pattern=64, v_proj uses global lora_alpha=16."""
+        hidden = 32
+        r = 8
+        global_alpha = 16
+        q_alpha = 64
+
+        model_dir, base_weights = self._make_base_model(tmp_path, hidden=hidden)
+        adapter_dir, config = self._make_adapter(tmp_path, r=r, alpha=global_alpha)
+        config["alpha_pattern"] = {"layers.0.self_attn.q_proj": q_alpha}
+        (adapter_dir / "adapter_config.json").write_text(json.dumps(config))
+
+        q_a = torch.randn(r, hidden)
+        q_b = torch.randn(hidden, r)
+        v_a = torch.randn(r, hidden)
+        v_b = torch.randn(hidden, r)
+        lora_state = {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": q_a,
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": q_b,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_A.weight": v_a,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_B.weight": v_b,
+        }
+        safetensors.torch.save_file(
+            lora_state, adapter_dir / "adapter_model.safetensors"
+        )
+
+        output_dir = tmp_path / "output"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=output_dir,
+            device="cpu",
+        )
+        merged = safetensors.torch.load_file(output_dir / "model.safetensors")
+
+        q_key = "model.layers.0.self_attn.q_proj.weight"
+        v_key = "model.layers.0.self_attn.v_proj.weight"
+        expected_q = base_weights[q_key] + (q_alpha / r) * (q_b @ q_a)
+        expected_v = base_weights[v_key] + (global_alpha / r) * (v_b @ v_a)
+        assert torch.allclose(merged[q_key], expected_q, atol=1e-5)
+        assert torch.allclose(merged[v_key], expected_v, atol=1e-5)
+
+    def test_merge_applies_rank_pattern_per_module(self, tmp_path):
+        """End-to-end: q_proj uses rank_pattern=16, v_proj uses global r=8."""
+        hidden = 32
+        global_r = 8
+        q_r = 16
+        alpha = 16
+
+        model_dir, base_weights = self._make_base_model(tmp_path, hidden=hidden)
+        adapter_dir, config = self._make_adapter(tmp_path, r=global_r, alpha=alpha)
+        config["rank_pattern"] = {"layers.0.self_attn.q_proj": q_r}
+        (adapter_dir / "adapter_config.json").write_text(json.dumps(config))
+
+        q_a = torch.randn(q_r, hidden)
+        q_b = torch.randn(hidden, q_r)
+        v_a = torch.randn(global_r, hidden)
+        v_b = torch.randn(hidden, global_r)
+        lora_state = {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": q_a,
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": q_b,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_A.weight": v_a,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_B.weight": v_b,
+        }
+        safetensors.torch.save_file(
+            lora_state, adapter_dir / "adapter_model.safetensors"
+        )
+
+        output_dir = tmp_path / "output"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=output_dir,
+            device="cpu",
+        )
+        merged = safetensors.torch.load_file(output_dir / "model.safetensors")
+
+        q_key = "model.layers.0.self_attn.q_proj.weight"
+        v_key = "model.layers.0.self_attn.v_proj.weight"
+        expected_q = base_weights[q_key] + (alpha / q_r) * (q_b @ q_a)
+        expected_v = base_weights[v_key] + (alpha / global_r) * (v_b @ v_a)
+        assert torch.allclose(merged[q_key], expected_q, atol=1e-5)
+        assert torch.allclose(merged[v_key], expected_v, atol=1e-5)
+
+    def test_merge_applies_rank_and_alpha_pattern_combined(self, tmp_path):
+        """End-to-end: q_proj overrides both rank (16) and alpha (64), v_proj uses globals."""
+        hidden = 32
+        global_r = 8
+        global_alpha = 16
+        q_r = 16
+        q_alpha = 64
+
+        model_dir, base_weights = self._make_base_model(tmp_path, hidden=hidden)
+        adapter_dir, config = self._make_adapter(
+            tmp_path, r=global_r, alpha=global_alpha
+        )
+        config["rank_pattern"] = {"layers.0.self_attn.q_proj": q_r}
+        config["alpha_pattern"] = {"layers.0.self_attn.q_proj": q_alpha}
+        (adapter_dir / "adapter_config.json").write_text(json.dumps(config))
+
+        q_a = torch.randn(q_r, hidden)
+        q_b = torch.randn(hidden, q_r)
+        v_a = torch.randn(global_r, hidden)
+        v_b = torch.randn(hidden, global_r)
+        lora_state = {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": q_a,
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": q_b,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_A.weight": v_a,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_B.weight": v_b,
+        }
+        safetensors.torch.save_file(
+            lora_state, adapter_dir / "adapter_model.safetensors"
+        )
+
+        output_dir = tmp_path / "output"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=output_dir,
+            device="cpu",
+        )
+        merged = safetensors.torch.load_file(output_dir / "model.safetensors")
+
+        q_key = "model.layers.0.self_attn.q_proj.weight"
+        v_key = "model.layers.0.self_attn.v_proj.weight"
+        expected_q = base_weights[q_key] + (q_alpha / q_r) * (q_b @ q_a)
+        expected_v = base_weights[v_key] + (global_alpha / global_r) * (v_b @ v_a)
+        assert torch.allclose(merged[q_key], expected_q, atol=1e-5)
+        assert torch.allclose(merged[v_key], expected_v, atol=1e-5)
+
+    def test_dora_merge(self):
+        """DoRA merge applies magnitude normalization via PEFT."""
+        hidden = 32
+        r = 8
+        alpha = 16
+        scale = alpha / r
+
+        base = torch.randn(hidden, hidden)
+        lora_a = torch.randn(r, hidden)
+        lora_b = torch.randn(hidden, r)
+        magnitude = torch.randn(hidden).abs() + 0.1
+
+        lora_state = {
+            "base_model.model.layer.q_proj.lora_A.weight": lora_a,
+            "base_model.model.layer.q_proj.lora_B.weight": lora_b,
+            "base_model.model.layer.q_proj.lora_magnitude_vector": magnitude,
+        }
+
+        config = {"r": r, "lora_alpha": alpha, "use_dora": True}
+        merged, was_merged = _merge_tensor_with_lora(
+            base,
+            "layer.q_proj.weight",
+            lora_state,
+            scale,
+            config,
+            "cpu",
+            use_dora=True,
+        )
+        assert was_merged
+
+        # The merge should differ from both base and base+delta (DoRA applies normalization)
+        delta = scale * (lora_b @ lora_a)
+        assert not torch.allclose(merged, base, atol=1e-3)
+        assert not torch.allclose(merged, base + delta, atol=1e-3)
+
+    def test_fuse_unfuse_moe_merge(self):
+        """Test fuse→merge→unfuse for MoE expert weights (WeightConverter path)."""
+        from axolotl.cli.utils.lora_merge import _fuse_and_unfuse_with_merge
+
+        hidden = 16
+        intermediate = 32
+        num_experts = 4
+        r = 4
+        alpha = 8
+        scale = alpha / r
+
+        # Simulate checkpoint format: per-expert separate tensors
+        shard_tensors = {}
+        for i in range(num_experts):
+            shard_tensors[f"model.layers.0.mlp.experts.{i}.gate_proj.weight"] = (
+                torch.randn(intermediate, hidden)
+            )
+            shard_tensors[f"model.layers.0.mlp.experts.{i}.up_proj.weight"] = (
+                torch.randn(intermediate, hidden)
+            )
+            shard_tensors[f"model.layers.0.mlp.experts.{i}.down_proj.weight"] = (
+                torch.randn(hidden, intermediate)
+            )
+        shard_tensors["model.layers.0.self_attn.q_proj.weight"] = torch.randn(
+            hidden, hidden
+        )
+
+        # LoRA targets the fused key (runtime format)
+        lora_state = {
+            "base_model.model.model.layers.0.mlp.experts.gate_up_proj.lora_A.weight": torch.randn(
+                r, hidden
+            ),
+            "base_model.model.model.layers.0.mlp.experts.gate_up_proj.lora_B.weight": torch.randn(
+                intermediate * 2, r
+            ),
+            "base_model.model.model.layers.0.mlp.experts.down_proj.lora_A.weight": torch.randn(
+                r, intermediate
+            ),
+            "base_model.model.model.layers.0.mlp.experts.down_proj.lora_B.weight": torch.randn(
+                hidden, r
+            ),
+        }
+
+        # Build converters matching qwen2_moe pattern
+        from transformers.core_model_loading import (
+            Concatenate,
+            MergeModulelist,
+            WeightConverter,
+        )
+
+        converters = [
+            WeightConverter(
+                source_patterns=[
+                    "mlp.experts.*.gate_proj.weight",
+                    "mlp.experts.*.up_proj.weight",
+                ],
+                target_patterns="mlp.experts.gate_up_proj",
+                operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
+            ),
+            WeightConverter(
+                source_patterns="mlp.experts.*.down_proj.weight",
+                target_patterns="mlp.experts.down_proj",
+                operations=[MergeModulelist(dim=0)],
+            ),
+        ]
+
+        config = {"r": r, "lora_alpha": alpha}
+        result, merged_count, processed_keys = _fuse_and_unfuse_with_merge(
+            shard_tensors, converters, lora_state, scale, config, "cpu"
+        )
+
+        # Should have merged 2 LoRA targets (gate_up_proj and down_proj)
+        assert merged_count == 2
+
+        # Processed keys include original per-expert keys (removed) + fused keys (added)
+        assert len(processed_keys) > 0
+
+        # Output should be in fused format (runtime keys)
+        assert "model.layers.0.mlp.experts.gate_up_proj" in result
+        assert "model.layers.0.mlp.experts.down_proj" in result
+
+        # Per-expert keys should be removed
+        for i in range(num_experts):
+            assert f"model.layers.0.mlp.experts.{i}.gate_proj.weight" not in result
+
+        # Non-expert tensor should be passed through
+        assert "model.layers.0.self_attn.q_proj.weight" in result
+
+        # Verify fused tensors are 3D (stacked experts)
+        gate_up = result["model.layers.0.mlp.experts.gate_up_proj"]
+        assert gate_up.ndim == 3
+        assert gate_up.shape[0] == num_experts  # [num_experts, intermediate*2, hidden]
+
+        # Verify the fused LoRA delta was applied correctly
+        # Reconstruct the fused base (stack per-expert, concat gate+up)
+        gate_stack = torch.stack(
+            [
+                shard_tensors[f"model.layers.0.mlp.experts.{i}.gate_proj.weight"]
+                for i in range(num_experts)
+            ]
+        )
+        up_stack = torch.stack(
+            [
+                shard_tensors[f"model.layers.0.mlp.experts.{i}.up_proj.weight"]
+                for i in range(num_experts)
+            ]
+        )
+        base_fused = torch.cat([gate_stack, up_stack], dim=1)
+        lora_a = lora_state[
+            "base_model.model.model.layers.0.mlp.experts.gate_up_proj.lora_A.weight"
+        ]
+        lora_b = lora_state[
+            "base_model.model.model.layers.0.mlp.experts.gate_up_proj.lora_B.weight"
+        ]
+        expected_fused = base_fused + scale * (lora_b @ lora_a)
+        assert torch.allclose(gate_up, expected_fused, atol=1e-5)
+
+    def test_fuse_skipped_for_incomplete_expert_shard(self):
+        """Expert lists split across shard boundaries must not be fused from a partial shard."""
+        from transformers.core_model_loading import (
+            Concatenate,
+            MergeModulelist,
+            WeightConverter,
+        )
+
+        from axolotl.cli.utils.lora_merge import _fuse_and_unfuse_with_merge
+
+        hidden = 16
+        intermediate = 32
+        num_experts = 8
+
+        # Layer 0: gate/up expert counts mismatch within the shard (previously
+        # crashed in torch.cat). Layer 1: contiguous but partial expert list
+        # (previously silently fused a subset).
+        shard_tensors = {}
+        for i in range(5):
+            shard_tensors[f"model.layers.0.mlp.experts.{i}.gate_proj.weight"] = (
+                torch.randn(intermediate, hidden)
+            )
+        for i in range(4):
+            shard_tensors[f"model.layers.0.mlp.experts.{i}.up_proj.weight"] = (
+                torch.randn(intermediate, hidden)
+            )
+        for i in range(4):
+            shard_tensors[f"model.layers.1.mlp.experts.{i}.gate_proj.weight"] = (
+                torch.randn(intermediate, hidden)
+            )
+            shard_tensors[f"model.layers.1.mlp.experts.{i}.up_proj.weight"] = (
+                torch.randn(intermediate, hidden)
+            )
+
+        converters = [
+            WeightConverter(
+                source_patterns=[
+                    "mlp.experts.*.gate_proj.weight",
+                    "mlp.experts.*.up_proj.weight",
+                ],
+                target_patterns="mlp.experts.gate_up_proj",
+                operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
+            ),
+        ]
+
+        config = {"r": 4, "lora_alpha": 8}
+        result, merged_count, processed_keys = _fuse_and_unfuse_with_merge(
+            shard_tensors,
+            converters,
+            {},
+            2.0,
+            config,
+            "cpu",
+            expected_num_experts=num_experts,
+        )
+
+        assert merged_count == 0
+        assert not processed_keys
+        # No fused keys; every per-expert tensor passes through unchanged
+        assert "model.layers.0.mlp.experts.gate_up_proj" not in result
+        assert "model.layers.1.mlp.experts.gate_up_proj" not in result
+        assert set(result.keys()) == set(shard_tensors.keys())
+
+        # Complete layer but still-quantized tensors (nvfp4 uint8 qdata with
+        # weight_scale siblings): fusing raw qdata would orphan the scales.
+        quant_tensors = {}
+        for i in range(num_experts):
+            for proj in ("gate_proj", "up_proj"):
+                key = f"model.layers.2.mlp.experts.{i}.{proj}.weight"
+                quant_tensors[key] = torch.randint(
+                    0, 255, (intermediate, hidden // 2), dtype=torch.uint8
+                )
+                quant_tensors[key + "_scale"] = torch.randn(
+                    intermediate, hidden // 16
+                ).to(torch.float8_e4m3fn)
+        result, merged_count, processed_keys = _fuse_and_unfuse_with_merge(
+            quant_tensors,
+            converters,
+            {},
+            2.0,
+            config,
+            "cpu",
+            expected_num_experts=num_experts,
+        )
+        assert merged_count == 0
+        assert not processed_keys
+        assert "model.layers.2.mlp.experts.gate_up_proj" not in result
+        assert set(result.keys()) == set(quant_tensors.keys())
+
+    def test_param_wrapper_merge_math(self):
+        """ParamWrapper merge via PEFT's get_delta_weight matches manual einsum."""
+        num_experts = 4
+        r = 2
+        in_features = 8
+        out_features = 4
+        alpha = 4
+
+        base = torch.randn(num_experts, in_features, out_features)
+        lora_a = torch.randn(r * num_experts, in_features)
+        lora_b = torch.randn(out_features, r * num_experts)
+
+        config = {"r": r, "lora_alpha": alpha}
+        delta = _build_peft_layer_and_get_delta(
+            lora_a, lora_b, config, base, is_param_wrapper=True
+        )
+        assert delta.shape == base.shape
+
+        merged = base + delta
+
+        # Verify against manual einsum
+        scale = alpha / r
+        wa = lora_a.reshape(num_experts, r, in_features)
+        wb = lora_b.reshape(out_features, r, num_experts)
+        manual_delta = torch.einsum("o r e, e r i -> e i o", wb, wa) * scale
+        for e in range(num_experts):
+            assert torch.allclose(merged[e], base[e] + manual_delta[e], atol=1e-5), (
+                f"Expert {e} mismatch"
+            )
+
+    def test_param_wrapper_nesting_dim_filter(self):
+        """_find_param_wrapper_lora skips wrong-dimension LoRA at outer level."""
+        num_experts = 4
+        r = 2
+
+        # Outer LoRA (gate_up_proj): A=[r*E, 8], B=[16, r*E]
+        # Inner LoRA (down_proj via base_layer): A=[r*E, 16], B=[8, r*E]
+        lora_state = {
+            "base_model.model.mod.experts.lora_A.weight": torch.randn(
+                r * num_experts, 8
+            ),
+            "base_model.model.mod.experts.lora_B.weight": torch.randn(
+                16, r * num_experts
+            ),
+            "base_model.model.mod.experts.base_layer.lora_A.weight": torch.randn(
+                r * num_experts, 16
+            ),
+            "base_model.model.mod.experts.base_layer.lora_B.weight": torch.randn(
+                8, r * num_experts
+            ),
+        }
+
+        # gate_up_proj shape [4, 8, 16] — should match outer LoRA
+        a, b, name = _find_param_wrapper_lora(
+            lora_state, "mod.experts.gate_up_proj", tensor_shape=(4, 8, 16)
+        )
+        assert a is not None and name == "gate_up_proj"
+        assert a.shape == (r * num_experts, 8)  # outer
+
+        # down_proj shape [4, 16, 8] — outer dims don't match, should find inner
+        a, b, name = _find_param_wrapper_lora(
+            lora_state, "mod.experts.down_proj", tensor_shape=(4, 16, 8)
+        )
+        assert a is not None and name == "down_proj"
+        assert a.shape == (r * num_experts, 16)  # inner (base_layer)
+
+        # shape that matches neither — should return None
+        a, b, name = _find_param_wrapper_lora(
+            lora_state, "mod.experts.other", tensor_shape=(4, 99, 99)
+        )
+        assert a is None
+
+    def test_find_lora_weights_with_renamings(self):
+        """Weight renamings let checkpoint keys match LoRA keys."""
+        lora_state = {
+            "base_model.model.layers.0.mlp.fc1.lora_A.weight": torch.randn(8, 32),
+            "base_model.model.layers.0.mlp.fc1.lora_B.weight": torch.randn(32, 8),
+        }
+        # Direct lookup fails (checkpoint has "ff0", LoRA has "fc1")
+        a, b = find_lora_weights(lora_state, "layers.0.mlp.ff0.weight")
+        assert a is None
+
+        # With renaming ff0 → fc1, it should match
+        a, b = find_lora_weights(
+            lora_state, "layers.0.mlp.ff0.weight", weight_renamings={"ff0": "fc1"}
+        )
+        assert a is not None
+        assert a.shape == (8, 32)
+
+    def test_unmatched_tensors_pass_through(self):
+        """Tensors with no matching LoRA are returned unchanged."""
+        lora_state = {
+            "base_model.model.layer.q_proj.lora_A.weight": torch.randn(8, 32),
+            "base_model.model.layer.q_proj.lora_B.weight": torch.randn(32, 8),
+        }
+
+        # 1D tensor (layernorm) — never matched
+        ln = torch.randn(32)
+        merged, was_merged = _merge_tensor_with_lora(
+            ln, "layer.norm.weight", lora_state, 2.0, {}, "cpu"
+        )
+        assert not was_merged
+        assert torch.equal(merged, ln)
+
+        # 2D tensor with no matching key
+        unrelated = torch.randn(64, 32)
+        merged, was_merged = _merge_tensor_with_lora(
+            unrelated, "layer.other_proj.weight", lora_state, 2.0, {}, "cpu"
+        )
+        assert not was_merged
+        assert torch.equal(merged, unrelated)
+
+    def test_fan_in_fan_out_transpose(self):
+        """fan_in_fan_out config transposes the LoRA delta."""
+        hidden = 16
+        r = 4
+        alpha = 4  # scale = 1.0
+
+        base = torch.randn(hidden, hidden)
+        lora_a = torch.randn(r, hidden)
+        lora_b = torch.randn(hidden, r)
+
+        lora_state = {
+            "base_model.model.layer.proj.lora_A.weight": lora_a,
+            "base_model.model.layer.proj.lora_B.weight": lora_b,
+        }
+
+        config_normal = {"r": r, "lora_alpha": alpha}
+        config_fif = {"r": r, "lora_alpha": alpha, "fan_in_fan_out": True}
+
+        merged_normal, _ = _merge_tensor_with_lora(
+            base, "layer.proj.weight", lora_state, 1.0, config_normal, "cpu"
+        )
+        merged_fif, _ = _merge_tensor_with_lora(
+            base, "layer.proj.weight", lora_state, 1.0, config_fif, "cpu"
+        )
+
+        delta = (alpha / r) * (lora_b @ lora_a)
+        assert torch.allclose(merged_normal, base + delta, atol=1e-5)
+        assert torch.allclose(merged_fif, base + delta.T, atol=1e-5)
+        assert not torch.allclose(merged_normal, merged_fif, atol=1e-5)
+
+    def test_rslora_end_to_end(self, tmp_path):
+        """RSLoRA adapter uses alpha/sqrt(r) scaling in sharded merge."""
+        hidden = 16
+        r = 16
+        alpha = 32
+
+        model_dir, base_weights = self._make_base_model(tmp_path, hidden=hidden)
+        adapter_dir, _ = self._make_adapter(tmp_path, r=r, alpha=alpha, use_rslora=True)
+
+        lora_a = torch.randn(r, hidden)
+        lora_b = torch.randn(hidden, r)
+        lora_state = {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": lora_a,
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": lora_b,
+        }
+        safetensors.torch.save_file(
+            lora_state, adapter_dir / "adapter_model.safetensors"
+        )
+
+        output_dir = tmp_path / "output"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=output_dir,
+            device="cpu",
+        )
+
+        merged = safetensors.torch.load_file(output_dir / "model.safetensors")
+        rslora_scale = alpha / math.sqrt(r)  # 8.0, not 2.0
+        q_key = "model.layers.0.self_attn.q_proj.weight"
+        expected = base_weights[q_key] + rslora_scale * (lora_b @ lora_a)
+        assert torch.allclose(merged[q_key], expected, atol=1e-5)
+
+        # Confirm it differs from standard scale
+        wrong_scale = alpha / r  # 2.0
+        wrong_expected = base_weights[q_key] + wrong_scale * (lora_b @ lora_a)
+        assert not torch.allclose(merged[q_key], wrong_expected, atol=1e-3)
+
+    def test_multi_shard_index_json(self, tmp_path):
+        """Multi-shard merge generates a correct weight-map index."""
+        hidden = 16
+        r = 4
+        alpha = 8
+
+        model_dir = tmp_path / "base_model"
+        model_dir.mkdir()
+        (model_dir / "config.json").write_text("{}")
+
+        # Create 2 shards
+        shard1 = {"model.layers.0.weight": torch.randn(hidden, hidden)}
+        shard2 = {"model.layers.1.weight": torch.randn(hidden, hidden)}
+        safetensors.torch.save_file(
+            shard1, model_dir / "model-00001-of-00002.safetensors"
+        )
+        safetensors.torch.save_file(
+            shard2, model_dir / "model-00002-of-00002.safetensors"
+        )
+
+        # Write a base model index (will be skipped by copy_non_model_files)
+        base_index = {
+            "metadata": {},
+            "weight_map": {
+                "model.layers.0.weight": "model-00001-of-00002.safetensors",
+                "model.layers.1.weight": "model-00002-of-00002.safetensors",
+            },
+        }
+        (model_dir / "model.safetensors.index.json").write_text(json.dumps(base_index))
+
+        adapter_dir, _ = self._make_adapter(tmp_path, r=r, alpha=alpha)
+        safetensors.torch.save_file({}, adapter_dir / "adapter_model.safetensors")
+
+        output_dir = tmp_path / "output"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=output_dir,
+            device="cpu",
+        )
+
+        # Verify index was generated
+        index_path = output_dir / "model.safetensors.index.json"
+        assert index_path.exists()
+        with open(index_path) as f:
+            idx = json.load(f)
+
+        assert "weight_map" in idx
+        assert len(idx["weight_map"]) == 2
+        # Each key should map to a shard that exists
+        for _key, shard_name in idx["weight_map"].items():
+            assert (output_dir / shard_name).exists(), f"Missing shard: {shard_name}"
+
+    def test_dora_end_to_end(self, tmp_path):
+        """DoRA merge through the full sharded merge pipeline."""
+        hidden = 16
+        r = 4
+        alpha = 8
+
+        model_dir, base_weights = self._make_base_model(tmp_path, hidden=hidden)
+        adapter_dir, _ = self._make_adapter(tmp_path, r=r, alpha=alpha, use_dora=True)
+
+        lora_a = torch.randn(r, hidden)
+        lora_b = torch.randn(hidden, r)
+        magnitude = torch.randn(hidden).abs() + 0.1
+        lora_state = {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": lora_a,
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": lora_b,
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_magnitude_vector": magnitude,
+        }
+        safetensors.torch.save_file(
+            lora_state, adapter_dir / "adapter_model.safetensors"
+        )
+
+        output_dir = tmp_path / "output"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=output_dir,
+            device="cpu",
+        )
+
+        merged = safetensors.torch.load_file(output_dir / "model.safetensors")
+        q_key = "model.layers.0.self_attn.q_proj.weight"
+
+        # Use PEFT's own get_delta_weight as the reference
+        delta = _build_peft_layer_and_get_delta(
+            lora_a,
+            lora_b,
+            {"r": r, "lora_alpha": alpha, "use_dora": True},
+            base_weights[q_key],
+            magnitude=magnitude,
+        )
+        expected = base_weights[q_key] + delta
+        assert torch.allclose(merged[q_key], expected, atol=1e-5)
+
+        # Verify it differs from standard (non-DoRA) merge
+        standard_delta = _build_peft_layer_and_get_delta(
+            lora_a,
+            lora_b,
+            {"r": r, "lora_alpha": alpha},
+            base_weights[q_key],
+        )
+        assert not torch.allclose(delta, standard_delta, atol=1e-3)
+
+        # v_proj has no LoRA weights — should be unchanged
+        v_key = "model.layers.0.self_attn.v_proj.weight"
+        assert torch.equal(merged[v_key], base_weights[v_key]), (
+            "v_proj should be unchanged (no LoRA weights for it)"
+        )
+
+    def test_dora_merge_honors_alpha_pattern(self, tmp_path):
+        """DoRA + alpha_pattern: q_proj uses overridden alpha, v_proj uses global."""
+        hidden = 16
+        r = 4
+        global_alpha = 8
+        q_alpha = 32
+
+        model_dir, base_weights = self._make_base_model(tmp_path, hidden=hidden)
+        adapter_dir, config = self._make_adapter(
+            tmp_path, r=r, alpha=global_alpha, use_dora=True
+        )
+        config["alpha_pattern"] = {"layers.0.self_attn.q_proj": q_alpha}
+        (adapter_dir / "adapter_config.json").write_text(json.dumps(config))
+
+        q_a = torch.randn(r, hidden)
+        q_b = torch.randn(hidden, r)
+        q_mag = torch.randn(hidden).abs() + 0.1
+        v_a = torch.randn(r, hidden)
+        v_b = torch.randn(hidden, r)
+        v_mag = torch.randn(hidden).abs() + 0.1
+        lora_state = {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": q_a,
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": q_b,
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_magnitude_vector": q_mag,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_A.weight": v_a,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_B.weight": v_b,
+            "base_model.model.model.layers.0.self_attn.v_proj.lora_magnitude_vector": v_mag,
+        }
+        safetensors.torch.save_file(
+            lora_state, adapter_dir / "adapter_model.safetensors"
+        )
+
+        output_dir = tmp_path / "output"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=output_dir,
+            device="cpu",
+        )
+        merged = safetensors.torch.load_file(output_dir / "model.safetensors")
+
+        q_key = "model.layers.0.self_attn.q_proj.weight"
+        v_key = "model.layers.0.self_attn.v_proj.weight"
+        expected_q_delta = _build_peft_layer_and_get_delta(
+            q_a,
+            q_b,
+            {"r": r, "lora_alpha": global_alpha, "use_dora": True},
+            base_weights[q_key],
+            magnitude=q_mag,
+            lora_alpha_override=q_alpha,
+        )
+        expected_v_delta = _build_peft_layer_and_get_delta(
+            v_a,
+            v_b,
+            {"r": r, "lora_alpha": global_alpha, "use_dora": True},
+            base_weights[v_key],
+            magnitude=v_mag,
+        )
+        assert torch.allclose(
+            merged[q_key], base_weights[q_key] + expected_q_delta, atol=1e-5
+        )
+        assert torch.allclose(
+            merged[v_key], base_weights[v_key] + expected_v_delta, atol=1e-5
+        )
+        # Sanity: q_proj delta must differ from a non-overridden alpha computation.
+        wrong_q_delta = _build_peft_layer_and_get_delta(
+            q_a,
+            q_b,
+            {"r": r, "lora_alpha": global_alpha, "use_dora": True},
+            base_weights[q_key],
+            magnitude=q_mag,
+        )
+        assert not torch.allclose(expected_q_delta, wrong_q_delta, atol=1e-3)
+
+    def test_dora_missing_magnitude_raises(self):
+        """DoRA with missing magnitude vector raises an explicit error."""
+        hidden = 16
+        r = 4
+        alpha = 8
+        scale = alpha / r
+
+        base = torch.randn(hidden, hidden)
+        lora_a = torch.randn(r, hidden)
+        lora_b = torch.randn(hidden, r)
+
+        # No magnitude vector in lora_state
+        lora_state = {
+            "base_model.model.layer.proj.lora_A.weight": lora_a,
+            "base_model.model.layer.proj.lora_B.weight": lora_b,
+        }
+
+        config = {"r": r, "lora_alpha": alpha, "use_dora": True}
+        with pytest.raises(ValueError, match="DoRA merge requires a magnitude vector"):
+            _merge_tensor_with_lora(
+                base,
+                "layer.proj.weight",
+                lora_state,
+                scale,
+                config,
+                "cpu",
+                use_dora=True,
+            )
+
+
+class TestQuantizedBaseMerge:
+    """Per-(dtype x module-type) merge correctness for quantized base weights.
+
+    The efficient merge folds ``scaling*(B@A)`` into the base weight; for a quantized base it must do
+    so on the DEQUANTIZED value and keep bf16 (re-rounding the sum to the quant format drops the low-
+    magnitude LoRA delta). These are hermetic (synthetic tensors, CPU) regression guards.
+    """
+
+    E4M3_MAX = 448.0
+
+    @staticmethod
+    def _make_block_fp8(w: torch.Tensor, block: int):
+        """bf16 weight -> (float8_e4m3fn weight, fp32 block scale_inv). Block axes = last two dims."""
+        *lead, N, K = w.shape
+        sr, sc = N // block, K // block
+        wb = w.float().reshape(*lead, sr, block, sc, block)
+        amax = (
+            wb.abs().amax(dim=(-3, -1), keepdim=True).clamp_min(1e-12)
+        )  # per (sr,sc) block
+        scale = amax / TestQuantizedBaseMerge.E4M3_MAX
+        q = torch.clamp(
+            wb / scale,
+            -TestQuantizedBaseMerge.E4M3_MAX,
+            TestQuantizedBaseMerge.E4M3_MAX,
+        )
+        q = q.reshape(*lead, N, K).to(torch.float8_e4m3fn)
+        scale_inv = scale.reshape(*lead, sr, sc).to(torch.float32)
+        return q, scale_inv
+
+    @staticmethod
+    def _dequant(q: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor:
+        *lead, N, K = q.shape
+        sr, sc = scale_inv.shape[-2], scale_inv.shape[-1]
+        bn, bk = N // sr, K // sc
+        wf = q.float().reshape(*lead, sr, bn, sc, bk)
+        s = scale_inv.float().reshape(*lead, sr, 1, sc, 1)
+        return (wf * s).reshape(*lead, N, K)
+
+    @pytest.mark.parametrize("shape,block", [((16, 16), 8), ((2, 16, 16), 8)])
+    def test_dequantize_block_fp8_shard(self, shape, block):
+        """_dequantize_block_fp8_shard reproduces the block dequant + drops scale_inv (2D + 3D)."""
+        from axolotl.cli.utils.lora_merge import _dequantize_quantized_shard
+
+        torch.manual_seed(0)
+        w = torch.randn(*shape, dtype=torch.bfloat16) * 0.2
+        key = "m.experts.gate_up_proj" if len(shape) == 3 else "m.q_proj.weight"
+        q, si = self._make_block_fp8(w, block)
+        shard = {key: q, key + "_scale_inv": si, "m.norm.weight": torch.ones(4)}
+
+        out, _, _, _ = _dequantize_quantized_shard(shard, "cpu")
+
+        assert key + "_scale_inv" not in out  # scale dropped
+        assert out[key].dtype == torch.bfloat16
+        assert "m.norm.weight" in out  # unrelated tensor untouched
+        ref = self._dequant(q, si)
+        assert torch.allclose(out[key].float(), ref, atol=2e-2)
+
+    def test_block_fp8_untouched_without_scale_inv(self):
+        """A float8 weight with no *_scale_inv sibling is left as-is (can't dequant)."""
+        from axolotl.cli.utils.lora_merge import _dequantize_quantized_shard
+
+        q = torch.randn(8, 8).to(torch.float8_e4m3fn)
+        out, _, _, _ = _dequantize_quantized_shard({"x.weight": q}, "cpu")
+        assert out["x.weight"].dtype == torch.float8_e4m3fn
+
+    def test_merge_block_fp8_linear_folds_into_dequantized(self):
+        """End-to-end 2D: dequant shard then merge == dequant(base) + scaling*(B@A), output bf16."""
+        from axolotl.cli.utils.lora_merge import (
+            _dequantize_quantized_shard,
+            _merge_tensor_with_lora,
+        )
+
+        torch.manual_seed(1)
+        hidden, r, alpha, block = 32, 8, 16, 16
+        scale = alpha / r
+        w = torch.randn(hidden, hidden, dtype=torch.bfloat16) * 0.2
+        q, si = self._make_block_fp8(w, block)
+        key = "model.layers.0.self_attn.q_proj.weight"
+        lora_a = torch.randn(r, hidden) * 0.1
+        lora_b = torch.randn(hidden, r) * 0.1
+        lora_state = {
+            f"base_model.model.{key[: -len('.weight')]}.lora_A.weight": lora_a,
+            f"base_model.model.{key[: -len('.weight')]}.lora_B.weight": lora_b,
+        }
+
+        deq_shard, _, _, _ = _dequantize_quantized_shard(
+            {key: q, key + "_scale_inv": si}, "cpu"
+        )
+        merged, was_merged = _merge_tensor_with_lora(
+            deq_shard[key], key, lora_state, scale, {"r": r, "lora_alpha": alpha}, "cpu"
+        )
+
+        assert was_merged
+        assert merged.dtype == torch.bfloat16  # NOT re-rounded to fp8
+        expected = self._dequant(q, si) + scale * (lora_b @ lora_a)
+        rel = (merged.float() - expected).norm() / expected.norm()
+        assert rel < 5e-3, f"block-fp8 merge rel {rel:.2e}"
+
+    def test_raw_fp8_merge_is_wrong_regression(self):
+        """Guard the fix: folding the delta into the RAW fp8 (skipping dequant) is materially wrong,
+        so the dequant step must stay. Same synthetic case, merged without _dequantize_block_fp8_shard."""
+        from axolotl.cli.utils.lora_merge import _merge_tensor_with_lora
+
+        torch.manual_seed(1)
+        hidden, r, alpha, block = 32, 8, 16, 16
+        scale = alpha / r
+        w = torch.randn(hidden, hidden, dtype=torch.bfloat16) * 0.2
+        q, si = self._make_block_fp8(w, block)
+        key = "model.layers.0.self_attn.q_proj.weight"
+        lora_a, lora_b = torch.randn(r, hidden) * 0.1, torch.randn(hidden, r) * 0.1
+        lora_state = {
+            f"base_model.model.{key[: -len('.weight')]}.lora_A.weight": lora_a,
+            f"base_model.model.{key[: -len('.weight')]}.lora_B.weight": lora_b,
+        }
+        # merge on the raw fp8 tensor (the old, broken behaviour)
+        merged, _ = _merge_tensor_with_lora(
+            q, key, lora_state, scale, {"r": r, "lora_alpha": alpha}, "cpu"
+        )
+        expected = self._dequant(q, si) + scale * (lora_b @ lora_a)
+        rel = (merged.float() - expected).norm() / expected.norm()
+        assert rel > 0.1, (
+            "raw-fp8 merge unexpectedly close — the dequant guard may be untested"
+        )
+
+    def test_strip_quantization_config(self, tmp_path):
+        from axolotl.cli.utils.lora_merge import _strip_quantization_config
+
+        cfg = {
+            "model_type": "mistral_large4",
+            "torch_dtype": "float8_e4m3fn",
+            "quantization_config": {"quant_method": "fp8"},
+        }
+        (tmp_path / "config.json").write_text(json.dumps(cfg))
+        _strip_quantization_config(tmp_path)
+        out = json.loads((tmp_path / "config.json").read_text())
+        assert "quantization_config" not in out
+        assert out["torch_dtype"] == "bfloat16"
+
+    def test_block_fp8_merge_forward_equivalence_end_to_end(self, tmp_path):
+        """FORWARD-level, full disk round-trip: build a block-fp8 base + perturbed LoRA, run the real
+        merge_lora_sharded_efficient, reload, and compare OUTPUTS of unmerged (dequant-base + LoRA
+        path) vs merged (plain bf16 matmul). Catches config-strip / scale-drop / dtype / re-quant-on-
+        load issues the weight-level oracle can't."""
+        torch.manual_seed(3)
+        hidden, r, alpha, block = 32, 8, 16, 16
+        scale = alpha / r
+
+        w = torch.randn(hidden, hidden, dtype=torch.bfloat16) * 0.2
+        q, si = self._make_block_fp8(w, block)
+        deq_base = self._dequant(q, si)  # what the model actually computes with
+
+        model_dir = tmp_path / "base"
+        model_dir.mkdir()
+        key = "model.layers.0.self_attn.q_proj.weight"
+        safetensors.torch.save_file(
+            {key: q, key + "_scale_inv": si}, model_dir / "model.safetensors"
+        )
+        (model_dir / "config.json").write_text(
+            json.dumps(
+                {
+                    "torch_dtype": "float8_e4m3fn",
+                    "quantization_config": {"quant_method": "fp8"},
+                }
+            )
+        )
+
+        adapter_dir = tmp_path / "adapter"
+        adapter_dir.mkdir()
+        lora_a = torch.randn(r, hidden) * 0.1
+        lora_b = (
+            torch.randn(hidden, r) * 0.1
+        )  # perturbed (NOT zero-init) -> real effect
+        safetensors.torch.save_file(
+            {
+                f"base_model.model.{key[: -len('.weight')]}.lora_A.weight": lora_a,
+                f"base_model.model.{key[: -len('.weight')]}.lora_B.weight": lora_b,
+            },
+            adapter_dir / "adapter_model.safetensors",
+        )
+        (adapter_dir / "adapter_config.json").write_text(
+            json.dumps({"r": r, "lora_alpha": alpha, "peft_type": "LORA"})
+        )
+
+        out_dir = tmp_path / "merged"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=out_dir,
+            device="cpu",
+            dequant=True,  # explicit --dequant: bf16 output
+        )
+
+        # --dequant: merged config is de-quantized
+        merged_cfg = json.loads((out_dir / "config.json").read_text())
+        assert "quantization_config" not in merged_cfg
+
+        merged = {}
+        with safetensors.torch.safe_open(
+            out_dir / "model.safetensors", framework="pt"
+        ) as f:
+            for k in f.keys():
+                merged[k] = f.get_tensor(k)
+        assert key + "_scale_inv" not in merged
+        assert merged[key].dtype == torch.bfloat16
+
+        x = torch.randn(4, hidden)
+        # unmerged forward: dequant-base linear + LoRA branch (the module's actual computation)
+        y_unmerged = x @ deq_base.float().T + scale * (x @ lora_a.T) @ lora_b.T
+        y_merged = x @ merged[key].float().T
+        rel = (y_unmerged - y_merged).norm() / y_unmerged.norm()
+        assert rel < 5e-3, f"forward mismatch after block-fp8 merge: rel {rel:.2e}"
+
+    def test_block_fp8_merge_preserves_format_default(self, tmp_path):
+        """DEFAULT merge is FORMAT-PRESERVING: a block-fp8 base stays block-fp8 (fp8 weight +
+        weight_scale_inv), quantization_config is kept, and the forward still matches within fp8 tol.
+        Non-LoRA quantized weights pass through byte-identical."""
+        torch.manual_seed(4)
+        hidden, r, alpha, block = 32, 8, 16, 16
+        scale = alpha / r
+        w = torch.randn(hidden, hidden, dtype=torch.bfloat16) * 0.2
+        q, si = self._make_block_fp8(w, block)
+        deq_base = self._dequant(q, si)
+        # a second block-fp8 weight WITHOUT LoRA -> must pass through untouched
+        q2, si2 = self._make_block_fp8(
+            torch.randn(hidden, hidden, dtype=torch.bfloat16) * 0.2, block
+        )
+
+        model_dir = tmp_path / "base"
+        model_dir.mkdir()
+        key = "model.layers.0.self_attn.q_proj.weight"
+        okey = "model.layers.0.self_attn.o_proj.weight"
+        safetensors.torch.save_file(
+            {
+                key: q,
+                key + "_scale_inv": si,
+                okey: q2,
+                okey + "_scale_inv": si2,
+            },
+            model_dir / "model.safetensors",
+        )
+        (model_dir / "config.json").write_text(
+            json.dumps(
+                {
+                    "torch_dtype": "float8_e4m3fn",
+                    "quantization_config": {"quant_method": "fp8"},
+                }
+            )
+        )
+        adapter_dir = tmp_path / "adapter"
+        adapter_dir.mkdir()
+        lora_a = torch.randn(r, hidden) * 0.1
+        lora_b = torch.randn(hidden, r) * 0.1
+        safetensors.torch.save_file(
+            {
+                f"base_model.model.{key[: -len('.weight')]}.lora_A.weight": lora_a,
+                f"base_model.model.{key[: -len('.weight')]}.lora_B.weight": lora_b,
+            },
+            adapter_dir / "adapter_model.safetensors",
+        )
+        (adapter_dir / "adapter_config.json").write_text(
+            json.dumps({"r": r, "lora_alpha": alpha, "peft_type": "LORA"})
+        )
+
+        out_dir = tmp_path / "merged"
+        merge_lora_sharded_efficient(  # default: NO dequant flag -> format-preserving
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=out_dir,
+            device="cpu",
+        )
+        merged = {}
+        with safetensors.torch.safe_open(
+            out_dir / "model.safetensors", framework="pt"
+        ) as f:
+            for k in f.keys():
+                merged[k] = f.get_tensor(k)
+
+        # format preserved: fp8 weight + fp32 scale kept, config NOT stripped
+        assert merged[key].dtype == torch.float8_e4m3fn
+        assert merged[key + "_scale_inv"].dtype == torch.float32
+        assert "quantization_config" in json.loads(
+            (out_dir / "config.json").read_text()
+        )
+        # no-LoRA weight passes through byte-identical
+        assert torch.equal(merged[okey], q2) and torch.equal(
+            merged[okey + "_scale_inv"], si2
+        )
+
+        # forward matches within fp8 tolerance (the delta survives re-quantization with fresh scales)
+        def deq(qq, s):
+            O, K = qq.shape
+            sr, sc = s.shape
+            return (
+                qq.float().reshape(sr, O // sr, sc, K // sc) * s.reshape(sr, 1, sc, 1)
+            ).reshape(O, K)
+
+        x = torch.randn(4, hidden)
+        y_unmerged = x @ deq_base.float().T + scale * (x @ lora_a.T) @ lora_b.T
+        y_merged = x @ deq(merged[key], merged[key + "_scale_inv"]).T
+        rel = (y_unmerged - y_merged).norm() / y_unmerged.norm()
+        assert rel < 6e-2, f"format-preserving merge forward rel {rel:.2e}"
+
+    @staticmethod
+    def _make_mxfp8(w: torch.Tensor, block: int = 32):
+        """bf16 -> (e4m3 weight, uint8 e8m0 scale) with FLOOR e8m0 (block along last dim)."""
+        *lead, N, K = w.shape
+        nb = K // block
+        wb = w.float().reshape(*lead, N, nb, block)
+        amax = wb.abs().amax(dim=-1).clamp_min(1e-12)
+        exp = torch.floor(torch.log2(amax)) - 8.0
+        q = torch.clamp(wb / torch.exp2(exp)[..., None], -448.0, 448.0).reshape(
+            *lead, N, K
+        )
+        return q.to(torch.float8_e4m3fn), (exp + 127.0).clamp(0, 254).to(torch.uint8)
+
+    def test_dequantize_mxfp8_shard_and_merge(self):
+        """mxfp8 (e4m3 + e8m0/32) fused dequant + 2D merge folds into the dequantized weight."""
+        from axolotl.cli.utils.lora_merge import (
+            _dequantize_quantized_shard,
+            _merge_tensor_with_lora,
+        )
+
+        torch.manual_seed(2)
+        hidden, r, alpha = 64, 8, 16
+        scale = alpha / r
+        w = torch.randn(hidden, hidden, dtype=torch.bfloat16) * 0.2
+        q, s = self._make_mxfp8(w, 32)
+        key = "model.layers.0.self_attn.q_proj.weight"
+        deq_shard, did, _, _ = _dequantize_quantized_shard(
+            {key: q, key + "_scale": s}, "cpu"
+        )
+        assert did
+        assert (
+            key + "_scale" not in deq_shard and deq_shard[key].dtype == torch.bfloat16
+        )
+
+        # dequant reference (block-32 e8m0)
+        ref_deq = (
+            q.float().reshape(hidden, hidden // 32, 32)
+            * torch.exp2(s.float() - 127.0)[..., None]
+        ).reshape(hidden, hidden)
+        assert torch.allclose(deq_shard[key].float(), ref_deq, atol=2e-2)
+
+        lora_a, lora_b = torch.randn(r, hidden) * 0.1, torch.randn(hidden, r) * 0.1
+        lora_state = {
+            f"base_model.model.{key[: -len('.weight')]}.lora_A.weight": lora_a,
+            f"base_model.model.{key[: -len('.weight')]}.lora_B.weight": lora_b,
+        }
+        merged, was = _merge_tensor_with_lora(
+            deq_shard[key], key, lora_state, scale, {"r": r, "lora_alpha": alpha}, "cpu"
+        )
+        assert was and merged.dtype == torch.bfloat16
+        expected = ref_deq + scale * (lora_b @ lora_a)
+        assert (merged.float() - expected).norm() / expected.norm() < 5e-3
+
+    def test_dequantize_nvfp4_shard(self):
+        """nvfp4 (packed uint8 + e4m3 block-16 scale + per-tensor scale_2) fused dequant: components
+        built by torchao ``to_nvfp4`` (the same schema the loader reads) recover the original weight
+        within nvfp4 tolerance, and the scale tensors are dropped. ``_dequant_nvfp4`` reconstructs the
+        NVFP4Tensor identically to nvfp4_moe_loading's loader path."""
+        pytest.importorskip("torchao")
+        try:
+            from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+        except Exception:  # pragma: no cover
+            pytest.skip("NVFP4Tensor unavailable")
+        from axolotl.cli.utils.lora_merge import _dequantize_quantized_shard
+
+        torch.manual_seed(4)
+        N, K = 32, 64
+        w = torch.randn(N, K, dtype=torch.bfloat16) * 0.2
+        p = (w.abs().max() / (6.0 * 448.0)).reshape(1).float().clamp(min=1e-12)
+        try:
+            nv = NVFP4Tensor.to_nvfp4(w, per_tensor_scale=p, is_swizzled_scales=False)
+        except Exception as ex:  # pragma: no cover - torchao API / device gaps
+            pytest.skip(f"to_nvfp4 unavailable on this host: {ex}")
+
+        key = "model.layers.0.self_attn.q_proj.weight"
+        shard = {
+            key: nv.qdata,
+            key + "_scale": nv.scale,
+            key + "_scale_2": nv.per_tensor_scale.reshape(()),
+        }
+        out, did, _, _ = _dequantize_quantized_shard(shard, "cpu")
+        assert did
+        assert key + "_scale" not in out and key + "_scale_2" not in out
+        assert out[key].dtype == torch.bfloat16
+        # nvfp4 round-trip recovers the weight within fp4 (3-mantissa-bit) tolerance; a wrong
+        # packing/scale interpretation would give garbage, not ~w.
+        rel = (out[key].float() - w.float()).norm() / w.float().norm()
+        assert rel < 0.2, f"nvfp4 dequant round-trip rel {rel:.2e}"
+
+    def test_undequantized_quant_base_warns(self):
+        """A LoRA folding into a still-quantized weight (an unhandled format the shard dequant left
+        as uint8/fp8) must WARN loudly rather than silently corrupt — the safety net for mxfp4 /
+        per-tensor-fp8 / per-expert-unfused layouts."""
+        from unittest.mock import patch
+
+        import axolotl.cli.utils.lora_merge as lm
+
+        lm._WARNED_UNDEQUANT.clear()
+        hidden, r, alpha = 16, 4, 8
+        # a packed-4bit-like uint8 weight that the dequant step did NOT handle
+        tensor = torch.randint(0, 255, (hidden, hidden), dtype=torch.uint8)
+        key = "model.layers.0.self_attn.q_proj.weight"
+        lora_state = {
+            f"base_model.model.{key[: -len('.weight')]}.lora_A.weight": torch.randn(
+                r, hidden
+            ),
+            f"base_model.model.{key[: -len('.weight')]}.lora_B.weight": torch.randn(
+                hidden, r
+            ),
+        }
+        with patch.object(lm.LOG, "warning") as mock_warn:
+            # the warning fires BEFORE the fold; folding into raw uint8 then errors (as it should) —
+            # tolerate that, we're asserting the guard warned.
+            try:
+                lm._merge_tensor_with_lora(
+                    tensor,
+                    key,
+                    lora_state,
+                    alpha / r,
+                    {"r": r, "lora_alpha": alpha},
+                    "cpu",
+                )
+            except RuntimeError:
+                pass
+        msgs = " ".join(str(c.args[0]) for c in mock_warn.call_args_list)
+        assert "still %s" in msgs or "quantized format" in msgs, (
+            "expected undequant warning"
+        )
+
+    def test_dequantize_mxfp4_shard(self):
+        """mxfp4 (packed e2m1 + e8m0/32) fused dequant recovers the fp4-quantized weight. Built with
+        the codebook + low/high nibble order the ScatterMoE MX forward uses, so the merged weight
+        equals what the model computes."""
+        from axolotl.cli.utils.lora_merge import (
+            _FP4_E2M1_LUT,
+            _dequantize_quantized_shard,
+        )
+
+        torch.manual_seed(5)
+        N, K, block = 8, 64, 32
+        lut = torch.tensor(_FP4_E2M1_LUT)
+        w = torch.randn(N, K) * 0.5
+        # per-32-block e8m0 scale, then quantize each element to the nearest codebook value
+        nb = K // block
+        amax = w.reshape(N, nb, block).abs().amax(-1).clamp_min(1e-6)
+        exp = torch.floor(torch.log2(amax / 6.0))  # 6 = fp4 max
+        scale = torch.exp2(exp)  # [N, nb]
+        wn = (w.reshape(N, nb, block) / scale[..., None]).reshape(N, K)
+        idx = (wn.unsqueeze(-1) - lut).abs().argmin(-1)  # nearest codebook index [N,K]
+        qvals = lut[idx]
+        packed = (idx[:, 0::2] | (idx[:, 1::2] << 4)).to(
+            torch.uint8
+        )  # low=even, high=odd
+        ebyte = (exp + 127.0).to(torch.uint8)
+        key = "model.layers.0.mlp.experts.gate_up_proj"  # 2D here for simplicity
+        out, did, _, _ = _dequantize_quantized_shard(
+            {key: packed, key + "_scale": ebyte}, "cpu"
+        )
+        assert did and out[key].dtype == torch.bfloat16 and key + "_scale" not in out
+        expected = (qvals.reshape(N, nb, block) * scale[..., None]).reshape(N, K)
+        assert torch.allclose(out[key].float(), expected, atol=1e-2)
+
+    def test_dequantize_nvfp4_single_level(self):
+        """Single-level nvfp4 (e4m3 block-16 scale, NO scale_2) dequants with per_tensor_scale=1."""
+        pytest.importorskip("torchao")
+        try:
+            from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+        except Exception:  # pragma: no cover
+            pytest.skip("NVFP4Tensor unavailable")
+        from axolotl.cli.utils.lora_merge import _dequantize_quantized_shard
+
+        torch.manual_seed(6)
+        w = torch.randn(16, 64, dtype=torch.bfloat16) * 0.2
+        try:
+            nv = NVFP4Tensor.to_nvfp4(
+                w, is_swizzled_scales=False
+            )  # no per_tensor_scale
+        except Exception as ex:  # pragma: no cover
+            pytest.skip(f"to_nvfp4 unavailable: {ex}")
+        key = "model.layers.0.self_attn.q_proj.weight"
+        out, did, _, _ = _dequantize_quantized_shard(
+            {key: nv.qdata, key + "_scale": nv.scale}, "cpu"
+        )
+        assert did and out[key].dtype == torch.bfloat16
+        assert (out[key].float() - w.float()).norm() / w.float().norm() < 0.2
+
+    def test_mxfp8_ragged_blocks(self):
+        """mxfp8 dequant: each e8m0 scale covers a fixed 32-wide MX block; the final ragged block
+        (K=40 -> blocks of 32 + 8) is trimmed, NOT floor-spread evenly across K."""
+        from axolotl.cli.utils.lora_merge import _dequant_mxfp8
+
+        torch.manual_seed(7)
+        N, K = (
+            4,
+            40,
+        )  # nb = ceil(40/32) = 2: first 32 elems use scale[0], last 8 use scale[1]
+        w = (torch.randn(N, K) * 0.1).to(torch.float8_e4m3fn)
+        s = torch.randint(120, 130, (N, 2), dtype=torch.uint8)
+        out = _dequant_mxfp8(w, s, "cpu")
+        assert out.shape == (N, K)
+        scale = torch.exp2(s.float() - 127.0)  # [N, 2]
+        exp = w.float().clone()
+        exp[:, :32] *= scale[:, :1]
+        exp[:, 32:] *= scale[:, 1:2]
+        assert torch.allclose(out.float(), exp, atol=1e-2)
+
+    def test_dora_on_block_fp8_base(self):
+        """DoRA merge on a block-fp8 base: dequant -> DoRA magnitude-normalized fold -> bf16 (the
+        quant base goes through the same DoRA path, no crash, finite result)."""
+        from axolotl.cli.utils.lora_merge import (
+            _dequantize_quantized_shard,
+            _merge_tensor_with_lora,
+        )
+
+        torch.manual_seed(8)
+        hidden, r, alpha = 32, 8, 16
+        w = torch.randn(hidden, hidden, dtype=torch.bfloat16) * 0.2
+        q, si = self._make_block_fp8(w, 16)
+        key = "model.layers.0.self_attn.q_proj.weight"
+        deq, _, _, _ = _dequantize_quantized_shard(
+            {key: q, key + "_scale_inv": si}, "cpu"
+        )
+        lora_state = {
+            f"base_model.model.{key[: -len('.weight')]}.lora_A.weight": torch.randn(
+                r, hidden
+            )
+            * 0.1,
+            f"base_model.model.{key[: -len('.weight')]}.lora_B.weight": torch.randn(
+                hidden, r
+            )
+            * 0.1,
+            f"base_model.model.{key[: -len('.weight')]}.lora_magnitude_vector": torch.randn(
+                hidden
+            ).abs()
+            + 0.1,
+        }
+        merged, was = _merge_tensor_with_lora(
+            deq[key],
+            key,
+            lora_state,
+            alpha / r,
+            {"r": r, "lora_alpha": alpha, "use_dora": True},
+            "cpu",
+            use_dora=True,
+        )
+        assert (
+            was
+            and merged.dtype == torch.bfloat16
+            and torch.isfinite(merged.float()).all()
+        )
+
+    def test_resized_embeddings_override_carried(self, tmp_path):
+        """A resized embed_tokens/lm_head saved as full weights in the adapter REPLACES the base and
+        bumps config.json vocab_size (otherwise the trained/enlarged vocab is silently dropped)."""
+        from axolotl.cli.utils.lora_merge import merge_lora_sharded_efficient
+
+        hidden, base_vocab, new_vocab = 8, 10, 12
+        model_dir = tmp_path / "base"
+        model_dir.mkdir()
+        safetensors.torch.save_file(
+            {
+                "model.embed_tokens.weight": torch.randn(base_vocab, hidden),
+                "lm_head.weight": torch.randn(base_vocab, hidden),
+                "model.layers.0.self_attn.q_proj.weight": torch.randn(hidden, hidden),
+            },
+            model_dir / "model.safetensors",
+        )
+        (model_dir / "config.json").write_text(json.dumps({"vocab_size": base_vocab}))
+
+        adapter_dir = tmp_path / "adapter"
+        adapter_dir.mkdir()
+        new_embed = torch.randn(new_vocab, hidden)
+        new_head = torch.randn(new_vocab, hidden)
+        safetensors.torch.save_file(
+            {
+                "base_model.model.model.embed_tokens.weight": new_embed,
+                "base_model.model.lm_head.weight": new_head,
+            },
+            adapter_dir / "adapter_model.safetensors",
+        )
+        (adapter_dir / "adapter_config.json").write_text(
+            json.dumps({"r": 8, "lora_alpha": 16, "peft_type": "LORA"})
+        )
+
+        out_dir = tmp_path / "merged"
+        merge_lora_sharded_efficient(
+            base_model_path=model_dir,
+            lora_adapter_path=adapter_dir,
+            output_path=out_dir,
+            device="cpu",
+        )
+        merged = {}
+        with safetensors.torch.safe_open(
+            out_dir / "model.safetensors", framework="pt"
+        ) as f:
+            for k in f.keys():
+                merged[k] = f.get_tensor(k)
+        assert merged["model.embed_tokens.weight"].shape[0] == new_vocab
+        assert merged["lm_head.weight"].shape[0] == new_vocab
+        assert torch.allclose(merged["model.embed_tokens.weight"].float(), new_embed)
+        assert torch.allclose(merged["lm_head.weight"].float(), new_head)
+        assert (
+            json.loads((out_dir / "config.json").read_text())["vocab_size"] == new_vocab
+        )
+
+    def test_per_expert_unfused_mismatch_warns(self, tmp_path):
+        """A fused expert LoRA adapter over a PER-EXPERT-unfused base must warn (the expert LoRA would
+        otherwise be silently dropped by the shard merge)."""
+        from unittest.mock import patch
+
+        import axolotl.cli.utils.lora_merge as lm
+
+        model_dir = tmp_path / "base"
+        model_dir.mkdir()
+        # per-expert unfused base experts
+        safetensors.torch.save_file(
+            {
+                "model.layers.0.mlp.experts.0.gate_proj.weight": torch.randn(8, 8),
+                "model.layers.0.mlp.experts.1.gate_proj.weight": torch.randn(8, 8),
+            },
+            model_dir / "model.safetensors",
+        )
+        (model_dir / "config.json").write_text("{}")
+        adapter_dir = tmp_path / "adapter"
+        adapter_dir.mkdir()
+        # fused expert LoRA (targets experts.gate_up_proj)
+        safetensors.torch.save_file(
+            {
+                "base_model.model.model.layers.0.mlp.experts.lora_A.weight": torch.randn(
+                    16, 8
+                ),
+                "base_model.model.model.layers.0.mlp.experts.lora_B.weight": torch.randn(
+                    16, 16
+                ),
+            },
+            adapter_dir / "adapter_model.safetensors",
+        )
+        (adapter_dir / "adapter_config.json").write_text(
+            json.dumps({"r": 8, "lora_alpha": 16, "peft_type": "LORA"})
+        )
+
+        with patch.object(lm.LOG, "warning") as mock_warn:
+            lm.merge_lora_sharded_efficient(
+                base_model_path=model_dir,
+                lora_adapter_path=adapter_dir,
+                output_path=tmp_path / "merged",
+                device="cpu",
+            )
+        assert any("PER-EXPERT" in str(c.args[0]) for c in mock_warn.call_args_list), (
+            "expected a per-expert-unfused mismatch warning"
+        )
+
+    def test_nvfp4_expert_merge_writer(self, tmp_path):
+        """The expert-merge writer folds a FUSED expert LoRA into a PER-EXPERT unfused NVFP4 base:
+        per expert the output must be BITWISE requant(dequant(base) + delta), including a layer split
+        across the shard boundary; non-LoRA layers pass through byte-identical; --dequant emits the
+        fused bf16 param instead."""
+        pytest.importorskip("torchao")
+        from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+        from axolotl.cli.utils.lora_merge import (
+            _build_peft_layer_and_get_delta,
+            _dequant_nvfp4,
+            _find_param_wrapper_lora,
+            _requant_by_format,
+            merge_lora_sharded_efficient,
+        )
+
+        torch.manual_seed(0)
+        E, H, I, r, alpha = 4, 64, 16, 4, 8
+
+        def quant(w2d):
+            p = (w2d.abs().max() / (6.0 * 448.0)).reshape(1).clamp_min(1e-12)
+            nv = NVFP4Tensor.to_nvfp4(
+                w2d.float(), per_tensor_scale=p, is_swizzled_scales=False
+            )
+            return (
+                nv.qdata,
+                nv.scale.to(torch.float8_e4m3fn),
+                nv.per_tensor_scale.reshape(()),
+            )
+
+        def make_layer(layer, seed):
+            torch.manual_seed(seed)
+            keys = {}
+            for proj, (n, k) in (
+                ("gate_proj", (I, H)),
+                ("up_proj", (I, H)),
+                ("down_proj", (H, I)),
+            ):
+                for e in range(E):
+                    qd, sc, pts = quant(torch.randn(n, k) * k**-0.5)
+                    base = f"model.layers.{layer}.mlp.experts.{e}.{proj}"
+                    keys[f"{base}.weight"] = qd
+                    keys[f"{base}.weight_scale"] = sc
+                    keys[f"{base}.weight_scale_2"] = pts
+            return keys
+
+        def subset(keys, experts):
+            return {
+                k: v
+                for k, v in keys.items()
+                if int(k.split(".experts.")[1].split(".")[0]) in experts
+            }
+
+        # layer 0 complete in shard 0; layer 1 SPLIT across shards; layer 2 has no LoRA
+        l0, l1, l2 = make_layer(0, 10), make_layer(1, 11), make_layer(2, 12)
+        shard0 = {**l0, **subset(l1, {0, 1})}
+        shard1 = {**subset(l1, {2, 3}), **l2}
+
+        torch.manual_seed(99)
+        adapter = {}
+        for layer in (0, 1):
+            p = f"base_model.model.model.layers.{layer}.mlp.experts"
+            # outer wrapper = down_proj, inner .base_layer = gate_up (sonicmoe chain);
+            # Qwen3 orientation [E, out, in]: lora_A [r*E, in], lora_B [out, r*E]
+            adapter[f"{p}.lora_A.weight"] = (
+                torch.randn(r * E, I, dtype=torch.bfloat16) * 0.2
+            )
+            adapter[f"{p}.lora_B.weight"] = (
+                torch.randn(H, r * E, dtype=torch.bfloat16) * 0.2
+            )
+            adapter[f"{p}.base_layer.lora_A.weight"] = (
+                torch.randn(r * E, H, dtype=torch.bfloat16) * 0.2
+            )
+            adapter[f"{p}.base_layer.lora_B.weight"] = (
+                torch.randn(2 * I, r * E, dtype=torch.bfloat16) * 0.2
+            )
+
+        base_dir, adapter_dir = tmp_path / "base", tmp_path / "adapter"
+        base_dir.mkdir(), adapter_dir.mkdir()
+        safetensors.torch.save_file(
+            shard0, base_dir / "model-00001-of-00002.safetensors"
+        )
+        safetensors.torch.save_file(
+            shard1, base_dir / "model-00002-of-00002.safetensors"
+        )
+        wmap = {k: "model-00001-of-00002.safetensors" for k in shard0}
+        wmap.update({k: "model-00002-of-00002.safetensors" for k in shard1})
+        (base_dir / "model.safetensors.index.json").write_text(
+            json.dumps({"metadata": {"total_size": 0}, "weight_map": wmap})
+        )
+        (base_dir / "config.json").write_text(
+            json.dumps({"model_type": "synthetic-test", "num_experts": E})
+        )
+        safetensors.torch.save_file(adapter, adapter_dir / "adapter_model.safetensors")
+        (adapter_dir / "adapter_config.json").write_text(
+            json.dumps({"r": r, "lora_alpha": alpha, "peft_type": "LORA"})
+        )
+
+        out = tmp_path / "merged"
+        merge_lora_sharded_efficient(base_dir, adapter_dir, out, device="cpu")
+        merged = {}
+        for f in sorted(out.glob("*.safetensors")):
+            merged.update(safetensors.torch.load_file(f))
+
+        all_base = {**shard0, **shard1}
+        assert set(merged) == set(all_base)
+        assert all(
+            merged[k].dtype == v.dtype and merged[k].shape == v.shape
+            for k, v in all_base.items()
+        )
+        assert all(torch.equal(merged[k], v) for k, v in l2.items())
+
+        def fused_dequant(keys, layer, projs):
+            per = [
+                torch.stack(
+                    [
+                        _dequant_nvfp4(
+                            keys[f"model.layers.{layer}.mlp.experts.{e}.{proj}.weight"],
+                            keys[
+                                f"model.layers.{layer}.mlp.experts.{e}.{proj}.weight_scale"
+                            ],
+                            keys[
+                                f"model.layers.{layer}.mlp.experts.{e}.{proj}.weight_scale_2"
+                            ],
+                            "cpu",
+                        )
+                        for e in range(E)
+                    ]
+                )
+                for proj in projs
+            ]
+            return per[0] if len(per) == 1 else torch.cat(per, dim=1)
+
+        cfg = {"r": r, "lora_alpha": alpha}
+        for layer, src in ((0, l0), (1, l1)):
+            for fused_name, projs in (
+                ("gate_up_proj", ("gate_proj", "up_proj")),
+                ("down_proj", ("down_proj",)),
+            ):
+                base_f = fused_dequant(src, layer, projs)
+                lora_a, lora_b, _ = _find_param_wrapper_lora(
+                    adapter,
+                    f"model.layers.{layer}.mlp.experts.{fused_name}",
+                    tensor_shape=tuple(base_f.shape),
+                )
+                assert lora_a is not None, f"orientation fix: L{layer} {fused_name}"
+                delta = _build_peft_layer_and_get_delta(
+                    lora_a, lora_b, cfg, base_f, is_param_wrapper=True
+                )
+                want = (base_f.float() + delta.float()).to(torch.bfloat16)
+                col = 0
+                for proj in projs:
+                    n_rows = src[
+                        f"model.layers.{layer}.mlp.experts.0.{proj}.weight"
+                    ].shape[0]
+                    for e in range(E):
+                        kb = f"model.layers.{layer}.mlp.experts.{e}.{proj}"
+                        ref = _requant_by_format(
+                            "nvfp4",
+                            want[e, col : col + n_rows, :],
+                            {
+                                "_scale": src[f"{kb}.weight_scale"],
+                                "_scale_2": src[f"{kb}.weight_scale_2"],
+                            },
+                            "cpu",
+                        )
+                        assert torch.equal(merged[f"{kb}.weight"], ref[""])
+                        assert torch.equal(
+                            merged[f"{kb}.weight_scale"].view(torch.uint8),
+                            ref["_scale"].view(torch.uint8),
+                        )
+                        assert torch.equal(
+                            merged[f"{kb}.weight_scale_2"], ref["_scale_2"]
+                        )
+                    col += n_rows
+
+        # --dequant: the writer emits the merged FUSED bf16 param instead
+        out2 = tmp_path / "merged_dq"
+        merge_lora_sharded_efficient(
+            base_dir, adapter_dir, out2, device="cpu", dequant=True
+        )
+        merged2 = {}
+        for f in sorted(out2.glob("*.safetensors")):
+            merged2.update(safetensors.torch.load_file(f))
+        for layer in (0, 1):
+            for fused_name, projs in (
+                ("gate_up_proj", ("gate_proj", "up_proj")),
+                ("down_proj", ("down_proj",)),
+            ):
+                key = f"model.layers.{layer}.mlp.experts.{fused_name}"
+                assert merged2[key].dtype == torch.bfloat16
+                src = l0 if layer == 0 else l1
+                base_f = fused_dequant(src, layer, projs)
+                lora_a, lora_b, _ = _find_param_wrapper_lora(
+                    adapter, key, tensor_shape=tuple(merged2[key].shape)
+                )
+                delta = _build_peft_layer_and_get_delta(
+                    lora_a, lora_b, cfg, base_f, is_param_wrapper=True
+                )
+                want = (base_f.float() + delta.float()).to(torch.bfloat16)
+                assert torch.allclose(merged2[key].float(), want.float(), atol=1e-2)
+
+    def test_fused_expert_base_no_false_positive(self, tmp_path):
+        """A FUSED expert base (Mistral-Large-4 style) + fused adapter must NOT trigger the per-expert
+        warning."""
+        from axolotl.cli.utils.lora_merge import _detect_per_expert_unfused_mismatch
+
+        model_dir = tmp_path / "base"
+        model_dir.mkdir()
+        safetensors.torch.save_file(
+            {"model.layers.0.mlp.experts.gate_up_proj": torch.randn(2, 16, 8)},
+            model_dir / "model.safetensors",
+        )
+        lora_state = {
+            "base_model.model.model.layers.0.mlp.experts.lora_A.weight": torch.randn(
+                16, 8
+            )
+        }
+        assert not _detect_per_expert_unfused_mismatch(
+            [model_dir / "model.safetensors"], lora_state
+        )
+
+    def test_expert_lora_delta_matches_scattermoe_kernel(self):
+        """The merge's PEFT ParamWrapper expert delta EXACTLY equals what the ScatterMoE training
+        kernel applies (expert-major B via peft_lora_B_to_scattermoe). Guards against a rank-major vs
+        expert-major layout regression — the naive slice differs materially."""
+        pytest.importorskip("triton")
+        from axolotl.cli.utils.lora_merge import _build_peft_layer_and_get_delta
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            peft_lora_B_to_scattermoe,
+        )
+
+        torch.manual_seed(0)
+        E, r, IN, OUT, alpha = 4, 8, 16, 24, 16
+        scaling = alpha / r
+        A = torch.randn(r * E, IN) * 0.1
+        B = torch.randn(OUT, r * E) * 0.1
+        base = torch.randn(E, OUT, IN)
+
+        mine = _build_peft_layer_and_get_delta(
+            A, B, {"r": r, "lora_alpha": alpha}, base, is_param_wrapper=True
+        )
+        smB = peft_lora_B_to_scattermoe(B, E, r)
+        kern = torch.stack(
+            [
+                scaling * (smB[:, e * r : (e + 1) * r] @ A[e * r : (e + 1) * r, :])
+                for e in range(E)
+            ]
+        )
+        naive = torch.stack(
+            [
+                scaling * (B[:, e * r : (e + 1) * r] @ A[e * r : (e + 1) * r, :])
+                for e in range(E)
+            ]
+        )
+        assert torch.allclose(mine.float(), kern.float(), atol=1e-5)
+        # sanity: the layouts genuinely differ, so the test isn't vacuous
+        assert (mine.float() - naive.float()).norm() / naive.float().norm() > 0.5
+
+    def test_expert_3d_block_fp8_merge_folds(self):
+        """Full 3D fused-expert merge on a block-fp8 base: dequant -> PEFT expert delta -> matches the
+        kernel-layout reconstruction on the dequantized base."""
+        pytest.importorskip("triton")
+        from axolotl.cli.utils.lora_merge import (
+            _build_peft_layer_and_get_delta,
+            _dequantize_quantized_shard,
+        )
+        from axolotl.integrations.kernels.libs.scattermoe_lora.layers import (
+            peft_lora_B_to_scattermoe,
+        )
+
+        torch.manual_seed(1)
+        E, r, IN, OUT, alpha = 4, 8, 128, 256, 16
+        scaling = alpha / r
+        w = torch.randn(E, OUT, IN, dtype=torch.bfloat16) * 0.1
+        q, si = self._make_block_fp8(w, 64)
+        key = "model.layers.0.mlp.experts.gate_up_proj"
+        deq, _, _, _ = _dequantize_quantized_shard(
+            {key: q, key + "_scale_inv": si}, "cpu"
+        )
+        base_deq = self._dequant(q, si)
+
+        A = torch.randn(r * E, IN) * 0.05
+        B = torch.randn(OUT, r * E) * 0.05
+        delta = _build_peft_layer_and_get_delta(
+            A, B, {"r": r, "lora_alpha": alpha}, deq[key], is_param_wrapper=True
+        )
+        merged = deq[key].float() + delta.float()
+        smB = peft_lora_B_to_scattermoe(B, E, r)
+        kern_delta = torch.stack(
+            [
+                scaling * (smB[:, e * r : (e + 1) * r] @ A[e * r : (e + 1) * r, :])
+                for e in range(E)
+            ]
+        )
+        expected = base_deq.float() + kern_delta.float()
+        assert (merged - expected).norm() / expected.norm() < 5e-3
+
+    def test_swizzled_nvfp4_detected(self):
+        """Swizzled-scale nvfp4 with a PADDED (unaligned) shape is auto-detected via the scale-numel
+        mismatch and dequants correctly. (On-disk nvfp4 checkpoints ship NON-swizzled scales — the
+        heuristic can only catch the padded case, not a same-numel reorder at block-aligned shapes;
+        the default non-swizzled path is correct for real checkpoints.)"""
+        pytest.importorskip("torchao")
+        try:
+            from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+        except Exception:  # pragma: no cover
+            pytest.skip("NVFP4Tensor unavailable")
+        from axolotl.cli.utils.lora_merge import _dequantize_quantized_shard
+
+        torch.manual_seed(9)
+        N, K = (
+            32,
+            64,
+        )  # N<128 -> swizzle pads the scale grid, so numel differs from N*K/16
+        w = torch.randn(N, K, dtype=torch.bfloat16) * 0.2
+        p = (w.abs().max() / (6.0 * 448.0)).reshape(1).float().clamp(min=1e-12)
+        try:
+            nv = NVFP4Tensor.to_nvfp4(w, per_tensor_scale=p, is_swizzled_scales=True)
+        except Exception as ex:  # pragma: no cover
+            pytest.skip(f"swizzled to_nvfp4 unavailable: {ex}")
+        if nv.scale.numel() == N * (K // 16):
+            pytest.skip(
+                "this torchao build did not pad the swizzled scale; detection N/A"
+            )
+        key = "model.layers.0.self_attn.q_proj.weight"
+        out, did, _, _ = _dequantize_quantized_shard(
+            {key: nv.qdata, key + "_scale": nv.scale, key + "_scale_2": p.reshape(())},
+            "cpu",
+        )
+        assert did and out[key].dtype == torch.bfloat16
+        assert (out[key].float() - w.float()).norm() / w.float().norm() < 0.2
+
+    def test_modules_to_save_override_path(self):
+        """_find_full_override resolves the PEFT modules_to_save layout
+        (...<module>.modules_to_save.default.weight)."""
+        from axolotl.cli.utils.lora_merge import _find_full_override
+
+        w = torch.randn(6, 8)
+        lora_state = {
+            "base_model.model.score.modules_to_save.default.weight": w,
+        }
+        got = _find_full_override(lora_state, "score.weight")
+        assert got is not None and torch.equal(got, w)
+
+    def test_dequantize_mxfp4_ragged(self):
+        """mxfp4 dequant: fixed 32-wide MX blocks with a trimmed ragged tail (K=40 -> 32 + 8), not a
+        floor-spread."""
+        from axolotl.cli.utils.lora_merge import _dequant_mxfp4, _unpack_fp4
+
+        torch.manual_seed(10)
+        N, Khalf = 4, 20  # K = 40 nibbles, nb = ceil(40/32) = 2
+        packed = torch.randint(0, 255, (N, Khalf), dtype=torch.uint8)
+        s = torch.randint(120, 130, (N, 2), dtype=torch.uint8)
+        out = _dequant_mxfp4(packed, s, "cpu")
+        assert out.shape == (N, 40)
+        vals = _unpack_fp4(packed, "cpu")  # [N, 40]
+        scale = torch.exp2(s.float() - 127.0)
+        exp = vals.clone()
+        exp[:, :32] *= scale[:, :1]
+        exp[:, 32:] *= scale[:, 1:2]
+        assert torch.allclose(out.float(), exp, atol=1e-3)
+
+    def test_partial_dequant_reports_left_quantized(self):
+        """A shard with a dequantizable tensor AND an unsupported quantized tensor (per-tensor fp8,
+        scalar scale) must report left_quantized=True so the caller keeps quantization_config."""
+        from axolotl.cli.utils.lora_merge import _dequantize_quantized_shard
+
+        blk = torch.randn(16, 16).to(torch.float8_e4m3fn)  # block-fp8 (handled)
+        pt = torch.randn(8, 8).to(
+            torch.float8_e4m3fn
+        )  # per-tensor fp8 (unsupported here)
+        shard = {
+            "a.weight": blk,
+            "a.weight_scale_inv": torch.ones(1, 1),
+            "b.weight": pt,
+            "b.weight_scale": torch.tensor(2.0),  # scalar fp32 -> not e8m0, not block
+        }
+        out, did, left, _ = _dequantize_quantized_shard(shard, "cpu")
+        assert did and left
+        assert out["a.weight"].dtype == torch.bfloat16  # dequantized
+        assert out["b.weight"].dtype == torch.float8_e4m3fn  # left as-is
+        assert "b.weight_scale" in out  # its scale not dropped
diff --git a/tests/utils/schemas/validation/mora/test_mora_validation.py b/tests/utils/schemas/validation/mora/test_mora_validation.py
new file mode 100644
index 0000000000..da1169f509
--- /dev/null
+++ b/tests/utils/schemas/validation/mora/test_mora_validation.py
@@ -0,0 +1,100 @@
+"""Validation tests for the MoRA / ReMoRA integration."""
+
+import pytest
+
+from axolotl.integrations.mora import MoraType
+from axolotl.utils.config import prepare_plugins, validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestMoraValidation:
+    """MoRA-specific config validation."""
+
+    def test_mora_block_round_trips(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            {
+                "adapter": "mora",
+                "plugins": ["axolotl.integrations.mora.MoraPlugin"],
+                "mora": {
+                    "use_mora": True,
+                    "mora_type": "rope",
+                },
+            }
+        )
+
+        prepare_plugins(cfg)
+        validated = validate_config(cfg)
+
+        assert validated.adapter == "mora"
+        assert validated.mora.use_mora is True
+        assert validated.mora.mora_type == MoraType.ROPE
+
+    def test_mora_type_accepts_legacy_supported_numbers(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            {
+                "adapter": "mora",
+                "plugins": ["axolotl.integrations.mora.MoraPlugin"],
+                "mora": {
+                    "use_mora": True,
+                    "mora_type": 1,
+                },
+            }
+        )
+
+        prepare_plugins(cfg)
+        validated = validate_config(cfg)
+
+        assert validated.mora.mora_type == MoraType.SHARING
+
+    def test_mora_rejects_unsupported_variant_numbers(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            {
+                "adapter": "mora",
+                "plugins": ["axolotl.integrations.mora.MoraPlugin"],
+                "mora": {
+                    "use_mora": True,
+                    "mora_type": 2,
+                },
+            }
+        )
+
+        prepare_plugins(cfg)
+        with pytest.raises(ValueError, match="mora_type"):
+            validate_config(cfg)
+
+    def test_remora_uses_core_relora_fields(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            {
+                "adapter": "mora",
+                "plugins": ["axolotl.integrations.mora.MoraPlugin"],
+                "relora": True,
+                "jagged_restart_steps": 2000,
+                "mora": {
+                    "use_mora": True,
+                    "mora_type": "rope",
+                },
+            }
+        )
+
+        prepare_plugins(cfg)
+        validated = validate_config(cfg)
+
+        assert validated.relora is True
+        assert validated.jagged_restart_steps == 2000
+
+    def test_remora_still_requires_core_restart_steps(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            {
+                "adapter": "mora",
+                "plugins": ["axolotl.integrations.mora.MoraPlugin"],
+                "relora": True,
+                "mora": {
+                    "use_mora": True,
+                    "mora_type": "rope",
+                },
+            }
+        )
+
+        prepare_plugins(cfg)
+        with pytest.raises(ValueError, match="jagged_restart_steps"):
+            validate_config(cfg)
diff --git a/tests/utils/schemas/validation/test_activation_offloading.py b/tests/utils/schemas/validation/test_activation_offloading.py
new file mode 100644
index 0000000000..4bdc95cc9c
--- /dev/null
+++ b/tests/utils/schemas/validation/test_activation_offloading.py
@@ -0,0 +1,63 @@
+"""Test for config validation for activation offloading."""
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestActivationOffloading:
+    """
+    Test cases for activation offloading schema validation
+    """
+
+    def test_gc_converts_offload_wo_lora(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing="offload",
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.gradient_checkpointing is True
+        assert cfg.activation_offloading is True
+
+    def test_ac_offload_impl_noop_wo_adapter(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                activation_offloading=True,
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.gradient_checkpointing is True
+        assert cfg.activation_offloading is True
+
+    def test_hidden_states_offload_defaults_to_non_reentrant(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                activation_offloading="hidden_states",
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.gradient_checkpointing is True
+        assert cfg.gradient_checkpointing_kwargs["use_reentrant"] is False
+        assert cfg.activation_offloading == "hidden_states"
+
+    def test_hidden_states_offload_accepts_reentrant_alst(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                gradient_checkpointing_kwargs={"use_reentrant": True},
+                activation_offloading="hidden_states",
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.gradient_checkpointing_kwargs["use_reentrant"] is True
+        assert cfg.activation_offloading == "hidden_states"
diff --git a/tests/utils/schemas/validation/test_config_validators.py b/tests/utils/schemas/validation/test_config_validators.py
new file mode 100644
index 0000000000..fbfa79ad8a
--- /dev/null
+++ b/tests/utils/schemas/validation/test_config_validators.py
@@ -0,0 +1,254 @@
+"""
+Tests for new config validators added to AxolotlInputConfig.
+
+Covers:
+  - save_strategy: 'best' requires metric_for_best_model
+  - streaming=True with val_set_size > 0 is rejected
+  - lora_target_modules with invalid regex patterns is rejected
+  - GRPO: generation batch size must be divisible by num_generations,
+    num_generations >= 2, and effective_gbs >= num_generations * world_size
+"""
+
+import pytest
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestSaveStrategyBestValidator:
+    """save_strategy: 'best' must be accompanied by metric_for_best_model."""
+
+    def test_save_strategy_best_without_metric_raises(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(save_strategy="best")
+        with pytest.raises(ValueError, match="metric_for_best_model"):
+            validate_config(cfg)
+
+    def test_save_strategy_best_with_metric_passes(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            save_strategy="best",
+            metric_for_best_model="eval_loss",
+        )
+        validated = validate_config(cfg)
+        assert validated.save_strategy == "best"
+        assert validated.metric_for_best_model == "eval_loss"
+
+    def test_save_strategy_epoch_without_metric_passes(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(save_strategy="epoch")
+        validated = validate_config(cfg)
+        assert validated.save_strategy == "epoch"
+
+    def test_save_strategy_no_without_metric_passes(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(save_strategy="no")
+        validated = validate_config(cfg)
+        assert validated.save_strategy == "no"
+
+    def test_save_strategy_unset_without_metric_passes(self, min_base_cfg):
+        """The default (None / not set) should not require metric_for_best_model."""
+        validated = validate_config(min_base_cfg)
+        assert validated.save_strategy is None
+
+
+class TestStreamingWithValSetSizeValidator:
+    """streaming=True is incompatible with val_set_size > 0."""
+
+    def test_streaming_with_val_set_size_raises(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            streaming=True, val_set_size=0.1, max_steps=100
+        )
+        with pytest.raises(ValueError, match="val_set_size"):
+            validate_config(cfg)
+
+    def test_streaming_with_val_set_size_zero_passes(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            streaming=True, val_set_size=0.0, max_steps=100
+        )
+        validated = validate_config(cfg)
+        assert validated.streaming is True
+
+    def test_streaming_false_with_val_set_size_passes(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(streaming=False, val_set_size=0.1)
+        validated = validate_config(cfg)
+        assert validated.val_set_size == pytest.approx(0.1)
+
+    def test_streaming_unset_with_val_set_size_passes(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(val_set_size=0.2)
+        validated = validate_config(cfg)
+        assert validated.val_set_size == pytest.approx(0.2)
+
+
+class TestLoraTargetModulesRegexValidator:
+    """lora_target_modules entries must be valid Python regex patterns."""
+
+    def test_invalid_regex_raises(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            adapter="lora",
+            lora_target_modules=["q_proj", "[invalid_regex"],
+        )
+        with pytest.raises(ValueError, match="invalid regex pattern"):
+            validate_config(cfg)
+
+    def test_valid_regex_passes(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            adapter="lora",
+            lora_target_modules=["q_proj", "v_proj", r".*_proj"],
+        )
+        validated = validate_config(cfg)
+        assert "q_proj" in validated.lora_target_modules
+
+    def test_plain_module_names_pass(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            adapter="lora",
+            lora_target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        )
+        validated = validate_config(cfg)
+        assert len(validated.lora_target_modules) == 4
+
+    def test_lora_target_linear_string_not_validated(self, min_base_cfg):
+        """When lora_target_modules is a string (e.g. 'all-linear'), skip regex check."""
+        cfg = min_base_cfg | DictDefault(
+            adapter="lora",
+            lora_target_modules="all-linear",
+        )
+        # Should not raise
+        validate_config(cfg)
+
+    def test_multiple_invalid_patterns_reported(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            adapter="lora",
+            lora_target_modules=["[bad1", "[bad2"],
+        )
+        with pytest.raises(ValueError, match="invalid regex pattern"):
+            validate_config(cfg)
+
+
+class TestGRPOBatchSizeValidator:
+    """GRPO requires (mb*GA) % num_generations == 0 and num_generations >= 2.
+
+    These call the @model_validator(mode="before") classmethod directly on a
+    plain dict — same input shape it receives during full Pydantic validation,
+    just without dragging in unrelated fields (datasets / model loading / etc.)
+    that aren't relevant to what's under test. The validator is registered on
+    ``RLValidationMixin`` (which ``AxolotlInputConfig`` inherits) so this is the
+    same code path ``axolotl train`` exercises.
+    """
+
+    @staticmethod
+    def _check(data):
+        from axolotl.utils.schemas.validation import RLValidationMixin
+
+        return RLValidationMixin.check_grpo_batch_size_divisibility(data)
+
+    def test_divisible_passes(self):
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 4,
+            "trl": {"num_generations": 4},
+        }
+        # Should return data unchanged (no exception)
+        out = self._check(data)
+        assert out["trl"]["num_generations"] == 4
+
+    def test_non_divisible_raises(self):
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 2,
+            "trl": {"num_generations": 4},
+        }
+        with pytest.raises(ValueError, match="num_generations"):
+            self._check(data)
+
+    def test_non_divisible_error_includes_fix_hint(self):
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 3,
+            "trl": {"num_generations": 4},
+        }
+        with pytest.raises(ValueError, match="gradient_accumulation_steps: 4"):
+            self._check(data)
+
+    def test_num_generations_one_raises(self):
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 4,
+            "trl": {"num_generations": 1},
+        }
+        with pytest.raises(ValueError, match=r"num_generations >= 2"):
+            self._check(data)
+
+    def test_explicit_generation_batch_size_divisible_passes(self):
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "trl": {"num_generations": 4, "generation_batch_size": 8},
+        }
+        out = self._check(data)
+        assert out["trl"]["generation_batch_size"] == 8
+
+    def test_explicit_generation_batch_size_non_divisible_raises(self):
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "trl": {"num_generations": 4, "generation_batch_size": 6},
+        }
+        with pytest.raises(ValueError, match="trl.generation_batch_size"):
+            self._check(data)
+
+    def test_non_grpo_skips_check(self):
+        # Anything other than rl=grpo should pass through untouched, even
+        # with non-divisible batch sizes — they're irrelevant to other RL
+        # methods that don't use group-relative advantages.
+        data = {
+            "rl": "dpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 3,
+            "trl": {"num_generations": 4},
+        }
+        assert self._check(data) is data
+
+    def test_no_rl_set_skips_check(self):
+        data = {
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 3,
+        }
+        assert self._check(data) is data
+
+    def test_grpo_without_num_generations_skips_check(self):
+        # If num_generations isn't set, TRL uses its own default — we don't
+        # have enough info to validate, so the validator must short-circuit
+        # rather than guess.
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 3,
+            "trl": {},
+        }
+        out = self._check(data)
+        assert out["rl"] == "grpo"
+
+    def test_multi_rank_group_size_check(self):
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 4,  # gbs=4
+            "world_size": 2,  # need gbs >= 4*2 = 8
+            "trl": {"num_generations": 4},
+        }
+        with pytest.raises(ValueError, match=r"world_size=2"):
+            self._check(data)
+
+    def test_multi_rank_group_size_satisfied(self):
+        data = {
+            "rl": "grpo",
+            "micro_batch_size": 1,
+            "gradient_accumulation_steps": 8,  # gbs=8 >= 4*2
+            "world_size": 2,
+            "trl": {"num_generations": 4},
+        }
+        out = self._check(data)
+        assert out["gradient_accumulation_steps"] == 8
diff --git a/tests/utils/schemas/validation/test_default_values.py b/tests/utils/schemas/validation/test_default_values.py
new file mode 100644
index 0000000000..332dfe77fb
--- /dev/null
+++ b/tests/utils/schemas/validation/test_default_values.py
@@ -0,0 +1,21 @@
+"""Tests for default values for configurations"""
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestDefaultConfigValues:
+    """Tests for default values for configurations"""
+
+    def test_pad_to_sequence_len(self, min_base_cfg):
+        """Tests that sample packing automatically sets pad_to_sequence_len to True"""
+        cfg = (
+            DictDefault(
+                sample_packing=True,
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+
+        assert cfg.pad_to_sequence_len is True
diff --git a/tests/utils/schemas/validation/test_fsdp.py b/tests/utils/schemas/validation/test_fsdp.py
new file mode 100644
index 0000000000..75ffd3b342
--- /dev/null
+++ b/tests/utils/schemas/validation/test_fsdp.py
@@ -0,0 +1,261 @@
+"""
+tests for pydantic fsdp validation
+"""
+
+import logging
+
+import pytest
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestFSDPValidation:
+    """
+    test class for pydantic fsdp validation
+    """
+
+    def test_fsdp_version_from_fsdp_config(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_config={
+                "version": 2,
+            },
+        )
+        cfg = validate_config(
+            cfg,
+        )
+        assert cfg.fsdp_version == 2
+
+    def test_fsdp_version_in_fsdp_config(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_version=2,
+            fsdp_config={
+                "reshard_after_forward": True,
+            },
+        )
+        cfg = validate_config(
+            cfg,
+        )
+        assert cfg.fsdp_version == 2
+        assert cfg.fsdp_config.fsdp_version == 2
+
+    def test_fsdp_offload_w_8bit_optim(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_config={
+                "offload_params": True,
+            },
+            optimizer="adamw_8bit",
+            fsdp_version=1,
+        )
+        with pytest.raises(
+            ValueError, match="FSDP Offload not compatible with adamw_8bit"
+        ):
+            validate_config(cfg)
+
+    def test_fsdp2_w_8bit_optim(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_config={
+                "offload_params": True,
+            },
+            optimizer="adamw_8bit",
+            fsdp_version=2,
+        )
+        with pytest.raises(
+            ValueError,
+            match="FSDP2 not compatible with adamw_8bit, use `adamw_torch_8bit` instead",
+        ):
+            validate_config(cfg)
+
+    def test_fsdp2_w_cpu_ram_efficient_loading(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            load_in_8bit=True,
+            adapter="lora",
+            fsdp_config={
+                "cpu_ram_efficient_loading": True,
+            },
+            fsdp_version=2,
+        )
+        validated_cfg = validate_config(cfg)
+        assert validated_cfg.fsdp_version == 2
+        assert validated_cfg.fsdp_config.cpu_ram_efficient_loading is True
+
+    def test_fsdp2_cpu_offload_pin_memory_requires_offload_params(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_config={
+                "cpu_offload_pin_memory": False,
+                "offload_params": False,
+            },
+            fsdp_version=2,
+        )
+        with pytest.raises(
+            ValueError,
+            match="disabling cpu_offload_pin_memory requires enabling offload_params",
+        ):
+            validate_config(cfg)
+
+    def test_fsdp1_cpu_offload_pin_memory_not_supported(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_config={
+                "cpu_offload_pin_memory": False,
+                "offload_params": True,
+            },
+            fsdp_version=1,
+        )
+        with pytest.raises(
+            ValueError,
+            match="FSDP1 does not support disabling cpu_offload_pin_memory, please set `fsdp_version` to 2",
+        ):
+            validate_config(cfg)
+
+    def test_fsdp2_cpu_offload_pin_memory_w_offload_params(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_config={
+                "cpu_offload_pin_memory": False,
+                "offload_params": True,
+            },
+            fsdp_version=2,
+        )
+        validated_cfg = validate_config(cfg)
+        assert validated_cfg.fsdp_config.cpu_offload_pin_memory is False
+        assert validated_cfg.fsdp_config.offload_params is True
+
+    def test_fsdp_prefixes_removed(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_config={
+                "fsdp_version": 2,
+                "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+                "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+                "fsdp_reshard_after_forward": True,
+            }
+        )
+        cfg = validate_config(cfg)
+        assert cfg.fsdp_version == 2
+        assert cfg.fsdp_config.fsdp_version == 2
+        for key in cfg.fsdp_config.keys():
+            if key != "fsdp_version":
+                assert not key.startswith("fsdp_")
+        assert cfg.fsdp_config.auto_wrap_policy == "TRANSFORMER_BASED_WRAP"
+        assert cfg.fsdp_config.transformer_layer_cls_to_wrap == "LlamaDecoderLayer"
+        assert cfg.fsdp_config.reshard_after_forward is True
+
+    def test_fp32_norms_requires_fsdp_config(self, min_base_cfg):
+        # fsdp_config is the canonical "is_fsdp" signal; fp32_norms requires it.
+        cfg = min_base_cfg | DictDefault(
+            fp32_norms=True,
+            fsdp_version=2,
+        )
+        with pytest.raises(ValueError, match="fp32_norms requires FSDP to be enabled"):
+            validate_config(cfg)
+
+    def test_fp32_norms_requires_fsdp2(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fp32_norms=True,
+            fsdp_version=1,
+            fsdp_config={"reshard_after_forward": True},
+        )
+        with pytest.raises(ValueError, match="fp32_norms requires fsdp_version: 2"):
+            validate_config(cfg)
+
+    def test_fp32_norms_cpu_ram_efficient_loading_ok(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fp32_norms=True,
+            fsdp_version=2,
+            fsdp_config={
+                "reshard_after_forward": True,
+                "cpu_ram_efficient_loading": True,
+            },
+        )
+        validated_cfg = validate_config(cfg)
+        assert validated_cfg.fp32_norms is True
+        assert validated_cfg.fsdp_config.cpu_ram_efficient_loading is True
+
+    def test_fp32_norms_tensor_parallel_ok(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fp32_norms=True,
+            fsdp_version=2,
+            tensor_parallel_size=2,
+            fsdp_config={"reshard_after_forward": True},
+        )
+        validated_cfg = validate_config(cfg)
+        assert validated_cfg.fp32_norms is True
+        assert validated_cfg.tensor_parallel_size == 2
+
+    def test_fp32_norms_fsdp2_ok(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fp32_norms=True,
+            fp32_norm_classes=["AfmoeRMSNorm"],
+            fsdp_version=2,
+            fsdp_config={"reshard_after_forward": True},
+        )
+        validated_cfg = validate_config(cfg)
+        assert validated_cfg.fp32_norms is True
+        assert validated_cfg.fp32_norm_classes == ["AfmoeRMSNorm"]
+
+    def test_fp32_norm_classes_without_fp32_norms_warns(self, min_base_cfg, caplog):
+        cfg = min_base_cfg | DictDefault(
+            fp32_norm_classes=["AfmoeRMSNorm"],
+        )
+        # axolotl.cli.configure_logging() sets propagate=False on the `axolotl`
+        # logger, so pytest caplog (attached to root) can't see records by
+        # default. Temporarily re-enable propagation for this assertion.
+        ax_logger = logging.getLogger("axolotl")
+        old_propagate = ax_logger.propagate
+        ax_logger.propagate = True
+        try:
+            with caplog.at_level("WARNING", logger="axolotl"):
+                validated_cfg = validate_config(cfg)
+        finally:
+            ax_logger.propagate = old_propagate
+        assert not validated_cfg.fp32_norms
+        assert "fp32_norm_classes is set but fp32_norms is not enabled" in caplog.text
+
+    def test_muon_fsdp1_rejected(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            optimizer="muon",
+            fsdp_version=1,
+            fsdp_config={"reshard_after_forward": True},
+        )
+        with pytest.raises(
+            ValueError, match="Muon optimizer is only compatible with FSDP2"
+        ):
+            validate_config(cfg)
+
+    @pytest.mark.parametrize(
+        "rl",
+        [
+            "dpo",
+            "kto",
+            "orpo",
+            "ipo",
+        ],
+    )
+    def test_fsdp2_dpo(self, min_base_cfg, rl):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_version=2,
+            fsdp_config={
+                "reshard_after_forward": True,
+            },
+            rl=rl,
+            load_in_8bit=True,
+            adapter="lora",
+            remove_unused_columns=False,
+        )
+        with pytest.raises(
+            ValueError,
+            match="FSDP2 does not support load_in_8bit or load_in_4bit with ",
+        ):
+            validate_config(cfg)
+
+    def test_size_based_wrap_requires_min_num_params(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            fsdp_version=2,
+            fsdp_config={
+                "auto_wrap_policy": "SIZE_BASED_WRAP",
+                "reshard_after_forward": True,
+            },
+        )
+        with pytest.raises(
+            ValueError,
+            match="min_num_params is required when auto_wrap_policy is SIZE_BASED_WRAP",
+        ):
+            validate_config(cfg)
diff --git a/tests/utils/schemas/validation/test_moe_quant.py b/tests/utils/schemas/validation/test_moe_quant.py
new file mode 100644
index 0000000000..a9300a3c66
--- /dev/null
+++ b/tests/utils/schemas/validation/test_moe_quant.py
@@ -0,0 +1,282 @@
+"""Tests for MoE expert quantization config validation and PEFT patch idempotency."""
+
+import pytest
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture()
+def gpu_caps():
+    return {
+        "compute_capability": "sm_89",
+        "bf16": True,
+        "tf32": False,
+        "n_gpu": 1,
+        "n_node": 1,
+    }
+
+
+@pytest.fixture()
+def env_caps():
+    return {"torch_version": "2.7.0"}
+
+
+class TestQuantizeMoeExpertsValidation:
+    """Test suite for quantize_moe_experts config validator."""
+
+    def test_requires_adapter(self, min_base_cfg, gpu_caps, env_caps):
+        """quantize_moe_experts without adapter should fail."""
+        cfg = (
+            DictDefault(
+                quantize_moe_experts=True,
+            )
+            | min_base_cfg
+        )
+        with pytest.raises(ValueError, match="requires adapter"):
+            validate_config(cfg, capabilities=gpu_caps, env_capabilities=env_caps)
+
+    def test_requires_quantization(self, min_base_cfg, gpu_caps, env_caps):
+        """quantize_moe_experts without load_in_4bit/8bit should fail."""
+        cfg = (
+            DictDefault(
+                quantize_moe_experts=True,
+                adapter="lora",
+            )
+            | min_base_cfg
+        )
+        with pytest.raises(ValueError, match="requires load_in_4bit or load_in_8bit"):
+            validate_config(cfg, capabilities=gpu_caps, env_capabilities=env_caps)
+
+    def test_valid_qlora_4bit(self, min_base_cfg, gpu_caps, env_caps):
+        """quantize_moe_experts with qlora + 4bit should pass."""
+        cfg = (
+            DictDefault(
+                quantize_moe_experts=True,
+                adapter="qlora",
+                load_in_4bit=True,
+            )
+            | min_base_cfg
+        )
+        result = validate_config(cfg, capabilities=gpu_caps, env_capabilities=env_caps)
+        assert result["quantize_moe_experts"] is True
+
+    def test_valid_lora_8bit(self, min_base_cfg, gpu_caps, env_caps):
+        """quantize_moe_experts with lora + 8bit should pass."""
+        cfg = (
+            DictDefault(
+                quantize_moe_experts=True,
+                adapter="lora",
+                load_in_8bit=True,
+            )
+            | min_base_cfg
+        )
+        result = validate_config(cfg, capabilities=gpu_caps, env_capabilities=env_caps)
+        assert result["quantize_moe_experts"] is True
+
+    def test_false_skips_validation(self, min_base_cfg, gpu_caps, env_caps):
+        """quantize_moe_experts=false should not check adapter/quantization."""
+        cfg = (
+            DictDefault(
+                quantize_moe_experts=False,
+            )
+            | min_base_cfg
+        )
+        result = validate_config(cfg, capabilities=gpu_caps, env_capabilities=env_caps)
+        assert result["quantize_moe_experts"] is False
+
+    def test_rejects_lora_target_linear(self, min_base_cfg, gpu_caps, env_caps):
+        """quantize_moe_experts with lora_target_linear should fail."""
+        cfg = (
+            DictDefault(
+                quantize_moe_experts=True,
+                adapter="qlora",
+                load_in_4bit=True,
+                lora_target_linear=True,
+            )
+            | min_base_cfg
+        )
+        with pytest.raises(ValueError, match="lora_target_linear is not compatible"):
+            validate_config(cfg, capabilities=gpu_caps, env_capabilities=env_caps)
+
+    def test_default_is_false(self, min_base_cfg, gpu_caps, env_caps):
+        """quantize_moe_experts should default to false."""
+        cfg = DictDefault({}) | min_base_cfg
+        result = validate_config(cfg, capabilities=gpu_caps, env_capabilities=env_caps)
+        assert result["quantize_moe_experts"] is False
+
+
+class TestLoraTargetParametersDropout:
+    """Test that lora_dropout must be 0 when lora_target_parameters is set."""
+
+    def test_rejects_nonzero_dropout(self, min_base_cfg):
+        """lora_dropout > 0 with lora_target_parameters should fail."""
+        cfg = (
+            DictDefault(
+                adapter="lora",
+                lora_target_parameters=["mlp.experts.gate_up_proj"],
+                lora_dropout=0.1,
+                load_in_8bit=True,
+            )
+            | min_base_cfg
+        )
+        with pytest.raises(ValueError, match="lora_dropout must be 0"):
+            validate_config(cfg)
+
+    def test_zero_dropout_passes(self, min_base_cfg):
+        """lora_dropout=0 with lora_target_parameters should pass."""
+        cfg = (
+            DictDefault(
+                adapter="lora",
+                lora_target_parameters=["mlp.experts.gate_up_proj"],
+                lora_dropout=0.0,
+                load_in_8bit=True,
+            )
+            | min_base_cfg
+        )
+        result = validate_config(cfg)
+        assert result["lora_dropout"] == 0.0
+
+
+class TestPeftPatchIdempotency:
+    """Test that patch_peft_target_parameters_matching is idempotent."""
+
+    def test_double_call_does_not_stack_wrappers(self):
+        """Calling patch twice should not double-wrap _inject_parameters."""
+        from peft.tuners.tuners_utils import BaseTuner
+
+        from axolotl.monkeypatch.moe_quant import (
+            patch_peft_target_parameters_matching,
+        )
+
+        original = BaseTuner._inject_parameters
+        try:
+            patch_peft_target_parameters_matching()
+            first_patched = BaseTuner._inject_parameters
+            patch_peft_target_parameters_matching()
+            second_patched = BaseTuner._inject_parameters
+            # Should be same function, not double-wrapped
+            assert first_patched is second_patched
+        finally:
+            BaseTuner._inject_parameters = original
+            patch_peft_target_parameters_matching._axolotl_patched = False
+
+
+class TestMoeAdapterTrainMergeRoundtrip:
+    """E2E: train adapter on quantized MoE experts, then merge onto plain model.
+
+    Verifies that param wrapping order during training matches merge, preventing
+    size mismatch errors when loading adapters in standard PEFT/vLLM.
+    """
+
+    @staticmethod
+    def _make_classes():
+        """Return FakeExperts and FakeModel classes shared by both model builders."""
+        import torch
+        import torch.nn as nn
+
+        class FakeExperts(nn.Module):
+            def __init__(self):
+                super().__init__()
+                # Model definition order: gate_up_proj first, then down_proj.
+                self.gate_up_proj = nn.Parameter(torch.randn(4, 16, 8))
+                self.down_proj = nn.Parameter(torch.randn(4, 8, 16))
+                # peft >= 0.19 reads base_layer.weight.device unconditionally
+                # in _maybe_shard_state_dict_for_tp; target_parameters-style
+                # modules legitimately don't have .weight, so stub a buffer.
+                self.register_buffer("weight", torch.empty(0), persistent=False)
+
+            def forward(self, x):
+                x = torch.matmul(x, self.gate_up_proj[0].T)  # (batch, 16)
+                x = torch.matmul(x, self.down_proj[0].T)  # (batch, 8)
+                return x
+
+        class FakeModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(8, 8)
+                self.experts = FakeExperts()
+
+            def forward(self, x):
+                return self.linear(x) + self.experts(x)
+
+        return FakeExperts, FakeModel
+
+    @staticmethod
+    def _make_quantized_model():
+        """Training model: parametrizations registered in alphabetical order."""
+        import torch.nn as nn
+        import torch.nn.utils.parametrize as P
+
+        from axolotl.monkeypatch.moe_quant import _moe_load_state
+
+        _, FakeModel = TestMoeAdapterTrainMergeRoundtrip._make_classes()
+
+        class PassthroughParametrization(nn.Module):
+            def forward(self, x):
+                return x
+
+        model = FakeModel()
+
+        # Record definition order before parametrization (mirrors real loading).
+        _moe_load_state["expert_param_order"]["experts"] = list(
+            model.experts._parameters.keys()
+        )
+
+        # Register in alphabetical order to expose the ordering mismatch.
+        P.register_parametrization(
+            model.experts, "down_proj", PassthroughParametrization(), unsafe=True
+        )
+        P.register_parametrization(
+            model.experts, "gate_up_proj", PassthroughParametrization(), unsafe=True
+        )
+        return model
+
+    @staticmethod
+    def _make_plain_model():
+        """Merge model: no parametrizations — standard branch uses definition order."""
+        _, FakeModel = TestMoeAdapterTrainMergeRoundtrip._make_classes()
+        return FakeModel()
+
+    def test_train_save_merge_no_size_mismatch(self, tmp_path):
+        """Train on quantized experts, merge onto plain model — must not raise."""
+        import torch
+        from peft import LoraConfig, PeftModel, get_peft_model
+        from peft.tuners.tuners_utils import BaseTuner
+
+        from axolotl.monkeypatch.moe_quant import (
+            _moe_load_state,
+            patch_peft_target_parameters_matching,
+        )
+
+        adapter_dir = tmp_path / "adapter"
+        lora_cfg = LoraConfig(
+            r=4,
+            lora_alpha=8,
+            target_modules=[],
+            target_parameters=["experts.gate_up_proj", "experts.down_proj"],
+            lora_dropout=0.0,
+            bias="none",
+        )
+        original_inject = BaseTuner._inject_parameters
+
+        # Training phase: quantized model (parametrized branch) with axolotl patch.
+        _moe_load_state["expert_param_order"] = {}
+        patch_peft_target_parameters_matching()
+        try:
+            peft_model = get_peft_model(self._make_quantized_model(), lora_cfg)
+        finally:
+            BaseTuner._inject_parameters = original_inject
+            patch_peft_target_parameters_matching._axolotl_patched = False
+
+        optimizer = torch.optim.SGD(peft_model.parameters(), lr=1e-3)
+        for _ in range(3):
+            peft_model(torch.randn(2, 8)).sum().backward()
+            optimizer.step()
+            optimizer.zero_grad()
+        peft_model.save_pretrained(str(adapter_dir))
+
+        # Merge with standard PEFT (no axolotl patch) to verify external compatibility.
+        loaded = PeftModel.from_pretrained(self._make_plain_model(), str(adapter_dir))
+        merged = loaded.merge_and_unload()
+        assert merged is not None
diff --git a/tests/utils/schemas/validation/test_qgalore.py b/tests/utils/schemas/validation/test_qgalore.py
new file mode 100644
index 0000000000..f6325a018e
--- /dev/null
+++ b/tests/utils/schemas/validation/test_qgalore.py
@@ -0,0 +1,40 @@
+"""Validation tests for the Q-GaLore optimizer config gates."""
+
+import pytest
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestQGaLoreValidation:
+    """Pydantic-level checks for q_galore_adamw8bit."""
+
+    def test_adapter_rejected(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            optimizer="q_galore_adamw8bit",
+            adapter="lora",
+            lora_r=8,
+            lora_alpha=16,
+            lora_target_linear=True,
+        )
+        with pytest.raises(ValueError, match="incompatible with adapter"):
+            validate_config(cfg)
+
+    def test_fsdp1_rejected(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            optimizer="q_galore_adamw8bit",
+            fsdp_version=1,
+            fsdp_config={"reshard_after_forward": True},
+        )
+        with pytest.raises(ValueError, match="requires FSDP2"):
+            validate_config(cfg)
+
+    def test_defaults_filled(self, min_base_cfg):
+        cfg = min_base_cfg | DictDefault(
+            optimizer="q_galore_adamw8bit",
+            bf16=True,
+        )
+        cfg = validate_config(cfg)
+        assert cfg.optim_target_modules == ["attn", "mlp"]
+        assert cfg.qgalore_rank == 256
+        assert cfg.qgalore_proj_bits == 4
diff --git a/tests/utils/schemas/validation/test_selective_checkpointing_validation.py b/tests/utils/schemas/validation/test_selective_checkpointing_validation.py
new file mode 100644
index 0000000000..e903df345a
--- /dev/null
+++ b/tests/utils/schemas/validation/test_selective_checkpointing_validation.py
@@ -0,0 +1,128 @@
+"""Test for config validation for selective activation checkpointing."""
+
+import pytest
+
+from axolotl.utils.config import validate_config
+from axolotl.utils.dict import DictDefault
+
+
+class TestSelectiveCheckpointing:
+    """
+    Test cases for selective_checkpointing schema validation
+    """
+
+    def test_bool_shorthand_normalizes_to_attention(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                selective_checkpointing=True,
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.selective_checkpointing["save"] == ["attention"]
+        assert cfg.selective_checkpointing["save_sliding_window"] is False
+
+    def test_false_normalizes_to_none(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                selective_checkpointing=False,
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.selective_checkpointing is None
+
+    def test_custom_save_list_preserved(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                selective_checkpointing={"save": ["attention", "aten::mm"]},
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.selective_checkpointing["save"] == ["attention", "aten::mm"]
+
+    def test_requires_gradient_checkpointing(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                selective_checkpointing=True,
+            )
+            | min_base_cfg
+        )
+
+        with pytest.raises(ValueError, match="requires gradient_checkpointing"):
+            validate_config(cfg)
+
+    def test_rejects_reentrant(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                gradient_checkpointing_kwargs={"use_reentrant": True},
+                selective_checkpointing=True,
+            )
+            | min_base_cfg
+        )
+
+        with pytest.raises(ValueError, match="non-reentrant"):
+            validate_config(cfg)
+
+    def test_rejects_trl_activation_offloading(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                activation_offloading=True,
+                selective_checkpointing=True,
+            )
+            | min_base_cfg
+        )
+
+        with pytest.raises(ValueError, match="hidden_states"):
+            validate_config(cfg)
+
+    def test_composes_with_hidden_states_offload(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                activation_offloading="hidden_states",
+                selective_checkpointing=True,
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.selective_checkpointing["save"] == ["attention"]
+        assert cfg.activation_offloading == "hidden_states"
+
+    def test_rejects_matmul_save_with_adapter(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                adapter="lora",
+                lora_r=16,
+                lora_alpha=32,
+                lora_target_linear=True,
+                selective_checkpointing={"save": ["attention", "aten::mm"]},
+            )
+            | min_base_cfg
+        )
+
+        with pytest.raises(ValueError, match="in-place"):
+            validate_config(cfg)
+
+    def test_allows_matmul_save_without_adapter(self, min_base_cfg):
+        cfg = (
+            DictDefault(
+                gradient_checkpointing=True,
+                selective_checkpointing={"save": ["attention", "aten::mm"]},
+            )
+            | min_base_cfg
+        )
+
+        cfg = validate_config(cfg)
+        assert cfg.selective_checkpointing["save"] == ["attention", "aten::mm"]
diff --git a/tests/utils/test_cuda13.py b/tests/utils/test_cuda13.py
new file mode 100644
index 0000000000..c86bff8ced
--- /dev/null
+++ b/tests/utils/test_cuda13.py
@@ -0,0 +1,93 @@
+"""Tests for CUDA 13 LD_LIBRARY_PATH helpers."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+from axolotl.utils.cuda13 import cu13_library_path, prepend_cu13_ld_library_path
+
+
+@pytest.fixture(autouse=True)
+def clear_nvidia_cu13_modules():
+    for module_name in ("nvidia.cu13", "nvidia"):
+        sys.modules.pop(module_name, None)
+    yield
+    for module_name in ("nvidia.cu13", "nvidia"):
+        sys.modules.pop(module_name, None)
+
+
+def _make_fake_cu13_package(root: Path) -> Path:
+    package_root = root / "nvidia" / "cu13"
+    package_root.mkdir(parents=True)
+    (root / "nvidia" / "__init__.py").write_text("", encoding="utf-8")
+    (package_root / "__init__.py").write_text("", encoding="utf-8")
+    (package_root / "lib").mkdir()
+    return package_root
+
+
+def test_prepend_cu13_ld_library_path_when_package_exists(monkeypatch, tmp_path):
+    package_root = _make_fake_cu13_package(tmp_path)
+    monkeypatch.syspath_prepend(str(tmp_path))
+    expected_lib = str(package_root / "lib")
+
+    assert cu13_library_path() == expected_lib
+    assert (
+        prepend_cu13_ld_library_path("/opt/custom/lib:/usr/lib")
+        == f"{expected_lib}:/opt/custom/lib:/usr/lib"
+    )
+
+
+def test_prepend_cu13_ld_library_path_is_idempotent(monkeypatch, tmp_path):
+    package_root = _make_fake_cu13_package(tmp_path)
+    monkeypatch.syspath_prepend(str(tmp_path))
+    expected_lib = str(package_root / "lib")
+
+    assert (
+        prepend_cu13_ld_library_path(f"{expected_lib}:/opt/custom/lib:{expected_lib}")
+        == f"{expected_lib}:/opt/custom/lib"
+    )
+
+
+def test_prepend_cu13_ld_library_path_is_noop_without_package(monkeypatch):
+    import builtins
+
+    real_import = builtins.__import__
+
+    def fake_import(name, *args, **kwargs):
+        if name == "nvidia.cu13":
+            raise ImportError(name)
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", fake_import)
+    assert cu13_library_path() is None
+    assert prepend_cu13_ld_library_path("/opt/custom/lib") == "/opt/custom/lib"
+    assert prepend_cu13_ld_library_path(None) == ""
+
+
+def test_cuda13_env_sh_prepend_path_when_package_exists(tmp_path):
+    package_root = _make_fake_cu13_package(tmp_path)
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = f"{tmp_path}:src"
+    env["LD_LIBRARY_PATH"] = ""
+
+    result = subprocess.run(
+        [
+            "bash",
+            "-lc",
+            'source scripts/cuda13_env.sh; printf %s "$LD_LIBRARY_PATH"',
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+        env=env,
+    )
+
+    path_parts = result.stdout.split(":")
+    assert path_parts[0] == str(package_root / "lib")
+    assert str(package_root / "lib") in path_parts
diff --git a/tests/utils/test_grpo_rw_fnc.py b/tests/utils/test_grpo_rw_fnc.py
new file mode 100644
index 0000000000..507de277b5
--- /dev/null
+++ b/tests/utils/test_grpo_rw_fnc.py
@@ -0,0 +1,18 @@
+import os
+
+import pytest
+
+from axolotl.core.trainers.grpo import GRPOStrategy
+
+
+def test_get_rollout_func_loads_successfully():
+    """Test that a valid rollout function can be loaded"""
+    rollout_func = GRPOStrategy.get_rollout_func("os.path.join")
+    assert callable(rollout_func)
+    assert rollout_func == os.path.join
+
+
+def test_get_rollout_func_invalid_module_raises_error():
+    """Test that invalid module path raises clear ValueError"""
+    with pytest.raises(ValueError, match="Rollout function .* not found"):
+        GRPOStrategy.get_rollout_func("nonexistent_module.my_func")
diff --git a/tests/utils/test_import_helper.py b/tests/utils/test_import_helper.py
new file mode 100644
index 0000000000..e1ab8bec51
--- /dev/null
+++ b/tests/utils/test_import_helper.py
@@ -0,0 +1,37 @@
+"""
+test cases for axolotl.utils.import_helper
+"""
+
+import pytest
+
+from axolotl.utils.import_helper import get_cls_from_module_str
+
+
+def test_get_cls_from_module_str():
+    cls = get_cls_from_module_str("axolotl.core.trainers.base.AxolotlTrainer")
+    assert cls.__name__ == "AxolotlTrainer"
+
+
+def test_get_cls_from_module_str_empty_string():
+    with pytest.raises(ValueError, match="module_str must be a non-empty string"):
+        get_cls_from_module_str("")
+
+
+def test_get_cls_from_module_str_whitespace_only():
+    with pytest.raises(ValueError, match="module_str must be a non-empty string"):
+        get_cls_from_module_str("   ")
+
+
+def test_get_cls_from_module_str_invalid_format():
+    with pytest.raises(ValueError, match="Invalid module string format"):
+        get_cls_from_module_str("single_part")
+
+
+def test_get_cls_from_module_str_nonexistent_module():
+    with pytest.raises(ImportError, match="Failed to import module"):
+        get_cls_from_module_str("nonexistent.module.Class")
+
+
+def test_get_cls_from_module_str_nonexistent_class():
+    with pytest.raises(AttributeError, match="Class 'NonExistentClass' not found"):
+        get_cls_from_module_str("axolotl.core.trainers.base.NonExistentClass")
diff --git a/tests/utils/test_mistral3_processor.py b/tests/utils/test_mistral3_processor.py
new file mode 100644
index 0000000000..ae2bc1fafd
--- /dev/null
+++ b/tests/utils/test_mistral3_processor.py
@@ -0,0 +1,149 @@
+"""Tests for Mistral3Processor with transformers v5 ProcessorMixin integration"""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from transformers.feature_extraction_utils import BatchFeature
+
+from axolotl.utils.mistral.mistral3_processor import Mistral3Processor
+from axolotl.utils.mistral.mistral_tokenizer import HFMistralTokenizer
+
+
+@pytest.fixture()
+def mock_tokenizer():
+    """Create a mock HFMistralTokenizer that passes v5 ProcessorMixin isinstance checks."""
+    return MagicMock(spec=HFMistralTokenizer)
+
+
+@pytest.fixture()
+def processor(mock_tokenizer):
+    return Mistral3Processor(tokenizer=mock_tokenizer)
+
+
+class TestMistral3ProcessorInit:
+    def test_tokenizer_is_set(self, processor, mock_tokenizer):
+        assert processor.tokenizer is mock_tokenizer
+
+    def test_chat_template_is_none(self, processor):
+        assert processor.chat_template is None
+
+    def test_audio_tokenizer_is_none(self, processor):
+        assert processor.audio_tokenizer is None
+
+
+class TestApplyChatTemplateTokenized:
+    """Test apply_chat_template with tokenize=True, return_dict=True"""
+
+    @pytest.fixture()
+    def batched_conversations(self):
+        return [
+            [
+                {"role": "user", "content": "Describe this image."},
+                {"role": "assistant", "content": "It is red."},
+            ],
+            [
+                {"role": "user", "content": "What is this?"},
+                {"role": "assistant", "content": "A cat."},
+            ],
+        ]
+
+    def test_returns_batch_feature_with_pixel_values(
+        self, processor, mock_tokenizer, batched_conversations
+    ):
+        pixel_values = torch.randn(2, 3, 224, 224, dtype=torch.float64)
+        mock_tokenizer.apply_chat_template.return_value = {
+            "input_ids": torch.tensor([[1, 2, 3], [4, 5, 6]]),
+            "attention_mask": torch.tensor([[1, 1, 1], [1, 1, 1]]),
+            "pixel_values": pixel_values,
+        }
+
+        result = processor.apply_chat_template(
+            batched_conversations, tokenize=True, return_dict=True
+        )
+
+        assert isinstance(result, BatchFeature)
+        assert "pixel_values" in result
+        assert "image_sizes" in result
+        assert result["pixel_values"].dtype == torch.float32
+        assert result["image_sizes"].shape == (2, 2)
+        assert result["image_sizes"][0].tolist() == [224, 224]
+
+    def test_returns_batch_feature_without_pixel_values(
+        self, processor, mock_tokenizer, batched_conversations
+    ):
+        mock_tokenizer.apply_chat_template.return_value = {
+            "input_ids": torch.tensor([[1, 2, 3], [4, 5, 6]]),
+            "attention_mask": torch.tensor([[1, 1, 1], [1, 1, 1]]),
+        }
+
+        result = processor.apply_chat_template(
+            batched_conversations, tokenize=True, return_dict=True
+        )
+
+        assert isinstance(result, BatchFeature)
+        assert "input_ids" in result
+        assert "image_sizes" not in result
+
+
+class TestApplyChatTemplateNotTokenized:
+    def test_single_conversation_returns_unwrapped(self, processor, mock_tokenizer):
+        """Single conversation (not batched) should return unwrapped result."""
+        single_conversation = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi"},
+        ]
+        mock_tokenizer.apply_chat_template.return_value = [
+            "<s>[INST]Hello[/INST]Hi</s>"
+        ]
+
+        result = processor.apply_chat_template(
+            single_conversation, tokenize=False, return_dict=False
+        )
+
+        assert result == "<s>[INST]Hello[/INST]Hi</s>"
+
+    def test_batched_conversations_returns_list(self, processor, mock_tokenizer):
+        batched = [
+            [
+                {"role": "user", "content": "Hello"},
+                {"role": "assistant", "content": "Hi"},
+            ],
+            [
+                {"role": "user", "content": "Bye"},
+                {"role": "assistant", "content": "Bye"},
+            ],
+        ]
+        mock_tokenizer.apply_chat_template.return_value = ["text1", "text2"]
+
+        result = processor.apply_chat_template(
+            batched, tokenize=False, return_dict=False
+        )
+
+        assert result == ["text1", "text2"]
+
+
+class TestCall:
+    def test_delegates_to_tokenizer(self, processor, mock_tokenizer):
+        mock_tokenizer.return_value = {
+            "input_ids": [1, 2, 3],
+            "attention_mask": [1, 1, 1],
+        }
+
+        result = processor("Hello world")
+
+        mock_tokenizer.assert_called_once()
+        assert isinstance(result, BatchFeature)
+
+
+class TestReturnTensorsValidation:
+    def test_rejects_non_pt_return_tensors(self, processor):
+        conversation = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi"},
+        ]
+
+        with pytest.raises(ValueError, match=r"only supports.*return_tensors='pt'"):
+            processor.apply_chat_template(
+                conversation, tokenize=True, return_dict=True, return_tensors="np"
+            )
diff --git a/tests/utils/test_models.py b/tests/utils/test_models.py
deleted file mode 100644
index e06bb6c250..0000000000
--- a/tests/utils/test_models.py
+++ /dev/null
@@ -1,37 +0,0 @@
-"""Module for testing models utils file."""
-
-
-import unittest
-from unittest.mock import patch
-
-import pytest
-
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model
-
-
-class ModelsUtilsTest(unittest.TestCase):
-    """Testing module for models utils."""
-
-    def test_cfg_throws_error_with_s2_attention_and_sample_packing(self):
-        cfg = DictDefault(
-            {
-                "s2_attention": True,
-                "sample_packing": True,
-                "base_model": "",
-                "model_type": "LlamaForCausalLM",
-            }
-        )
-
-        # Mock out call to HF hub
-        with patch(
-            "axolotl.utils.models.load_model_config"
-        ) as mocked_load_model_config:
-            mocked_load_model_config.return_value = {}
-            with pytest.raises(ValueError) as exc:
-                # Should error before hitting tokenizer, so we pass in an empty str
-                load_model(cfg, tokenizer="")
-            assert (
-                "shifted-sparse attention does not currently support sample packing"
-                in str(exc.value)
-            )
diff --git a/tests/utils/test_sinkgd.py b/tests/utils/test_sinkgd.py
new file mode 100644
index 0000000000..0fc9b6212e
--- /dev/null
+++ b/tests/utils/test_sinkgd.py
@@ -0,0 +1,51 @@
+"""Unit tests for the SinkGD optimizer's SR-Sinkhorn, including 3D fused-MoE shapes."""
+
+import torch
+
+from axolotl.utils.optimizers.sinkgd import SinkGD, sr_sinkhorn
+
+
+def test_sr_sinkhorn_2d_fixed_point():
+    """Rows converge to L2 norm sqrt(n), columns to sqrt(m)."""
+    torch.manual_seed(0)
+    m, n = 16, 24
+    x = sr_sinkhorn(torch.randn(m, n), iters=20, eps=1e-8)
+    assert torch.allclose(x.norm(dim=1), torch.full((m,), n**0.5), atol=1e-3)
+    assert torch.allclose(x.norm(dim=0), torch.full((n,), m**0.5), atol=1e-3)
+
+
+def test_sr_sinkhorn_3d_per_expert():
+    """A fused MoE weight [E, M, N] is normalized independently per expert."""
+    torch.manual_seed(0)
+    E, m, n = 4, 16, 24
+    x = sr_sinkhorn(torch.randn(E, m, n), iters=20, eps=1e-8)
+    # each expert independently reaches the doubly-balanced fixed point
+    assert torch.allclose(x.norm(dim=-1), torch.full((E, m), n**0.5), atol=1e-3)
+    assert torch.allclose(x.norm(dim=-2), torch.full((E, n), m**0.5), atol=1e-3)
+
+
+def test_sr_sinkhorn_3d_experts_decoupled():
+    """Scaling one expert's gradient must not change the other experts' updates."""
+    torch.manual_seed(0)
+    g = torch.randn(4, 16, 24)
+    x = sr_sinkhorn(g, iters=20, eps=1e-8)
+    g2 = g.clone()
+    g2[0] *= 100.0
+    x2 = sr_sinkhorn(g2, iters=20, eps=1e-8)
+    assert torch.allclose(x[1:], x2[1:], atol=1e-4)
+
+
+def test_sinkgd_step_3d_stateless():
+    """SinkGD updates a 3D expert weight and stores no optimizer state for it."""
+    torch.manual_seed(0)
+    w = torch.nn.Parameter(torch.randn(4, 16, 24))
+    opt = SinkGD(
+        [{"params": [w], "use_sinkgd": True, "weight_decay": 0.0}],
+        lr=1e-2,
+        sinkgd_lr_scale=1.0,
+    )
+    before = w.detach().clone()
+    w.grad = torch.randn_like(w)
+    opt.step()
+    assert not torch.allclose(before, w.detach())
+    assert opt.state[w] == {}  # SR-Sinkhorn is stateless
diff --git a/tests/utils/test_train.py b/tests/utils/test_train.py
new file mode 100644
index 0000000000..a1f6f6088d
--- /dev/null
+++ b/tests/utils/test_train.py
@@ -0,0 +1,24 @@
+"""test for train checkpoint utils"""
+
+import os
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.train import determine_last_checkpoint
+
+
+def test_determine_last_checkpoint(temp_dir):
+    cfg = DictDefault(
+        output_dir=temp_dir,
+    )
+    for cpt_idx in [1, 9, 10, 20]:
+        os.makedirs(
+            os.path.join(cfg.output_dir, f"checkpoint-{cpt_idx}"), exist_ok=True
+        )
+
+    last_checkpoint = determine_last_checkpoint(cfg, update=False)
+    assert last_checkpoint == os.path.join(cfg.output_dir, "checkpoint-20")
+
+    cfg.resume_from_checkpoint = None
+    cfg.auto_resume_from_checkpoints = True
+    determine_last_checkpoint(cfg, update=True)
+    assert cfg.resume_from_checkpoint == os.path.join(cfg.output_dir, "checkpoint-20")